History

Merge pull request #2896 from edlsh/fix/oauth-tool-rename-per-request-map

fix(amp): smart-mode tool name fixes + deep-mode response repair

Luis Pater committed 6 months agov6.10.6

Commit bdc42402 parents e4a93c0 80eb0376 files · +558 −85Browse files

Unified · Split

6 files
  1. internal/api/modules/amp/response_rewriter.go +50 −0
  2. internal/api/modules/amp/response_rewriter_test.go +51 −0
  3. internal/runtime/executor/claude_executor.go +87 −66
  4. internal/runtime/executor/claude_executor_test.go +141 −14
  5. sdk/api/handlers/openai/openai_responses_handlers.go +131 −4
  6. sdk/api/handlers/openai/openai_responses_handlers_stream_test.go +98 −1
internal/api/modules/amp/response_rewriter.go+50−0
@@ -123,6 +123,52 @@ func (rw *ResponseRewriter) Flush() { · 122 unmodified lines
123 123
124var modelFieldPaths = []string{"message.model", "model", "modelVersion", "response.model", "response.modelVersion"}124var modelFieldPaths = []string{"message.model", "model", "modelVersion", "response.model", "response.modelVersion"}
125 125
126// ampCanonicalToolNames maps tool names to the exact casing expected by the
127// Amp mode tool whitelist (case-sensitive match).
128var ampCanonicalToolNames = map[string]string{
129 "bash": "Bash",
130 "read": "Read",
131 "grep": "Grep",
132 "glob": "glob",
133 "task": "Task",
134 "check": "Check",
135}
136
137// normalizeAmpToolNames fixes tool_use block names to match Amp's canonical casing.
138// Some upstream models return lowercase tool names (e.g. "bash" instead of "Bash")
139// which causes Amp's case-sensitive mode whitelist to reject them.
140func normalizeAmpToolNames(data []byte) []byte {
141 // Non-streaming: content[].name in tool_use blocks
142 for index, block := range gjson.GetBytes(data, "content").Array() {
143 if block.Get("type").String() != "tool_use" {
144 continue
145 }
146 name := block.Get("name").String()
147 if canonical, ok := ampCanonicalToolNames[strings.ToLower(name)]; ok && name != canonical {
148 path := fmt.Sprintf("content.%d.name", index)
149 var err error
150 data, err = sjson.SetBytes(data, path, canonical)
151 if err != nil {
152 log.Warnf("Amp ResponseRewriter: failed to normalize tool name %q to %q: %v", name, canonical, err)
153 }
154 }
155 }
156
157 // Streaming: content_block.name in content_block_start events
158 if gjson.GetBytes(data, "content_block.type").String() == "tool_use" {
159 name := gjson.GetBytes(data, "content_block.name").String()
160 if canonical, ok := ampCanonicalToolNames[strings.ToLower(name)]; ok && name != canonical {
161 var err error
162 data, err = sjson.SetBytes(data, "content_block.name", canonical)
163 if err != nil {
164 log.Warnf("Amp ResponseRewriter: failed to normalize streaming tool name %q to %q: %v", name, canonical, err)
165 }
166 }
167 }
168
169 return data
170}
171
126// ensureAmpSignature injects empty signature fields into tool_use/thinking blocks172// ensureAmpSignature injects empty signature fields into tool_use/thinking blocks
127// in API responses so that the Amp TUI does not crash on P.signature.length.173// in API responses so that the Amp TUI does not crash on P.signature.length.
128func ensureAmpSignature(data []byte) []byte {174func ensureAmpSignature(data []byte) []byte {
@@ -179,6 +225,7 @@ func (rw *ResponseRewriter) suppressAmpThinking(data []byte) []byte { · 50 unmodified lines
179 225
180func (rw *ResponseRewriter) rewriteModelInResponse(data []byte) []byte {226func (rw *ResponseRewriter) rewriteModelInResponse(data []byte) []byte {
181 data = ensureAmpSignature(data)227 data = ensureAmpSignature(data)
228 data = normalizeAmpToolNames(data)
182 data = rw.suppressAmpThinking(data)229 data = rw.suppressAmpThinking(data)
183 if len(data) == 0 {230 if len(data) == 0 {
184 return data231 return data
@@ -278,6 +325,9 @@ func (rw *ResponseRewriter) rewriteStreamEvent(data []byte) []byte { · 93 unmodified lines
278 // Inject empty signature where needed325 // Inject empty signature where needed
279 data = ensureAmpSignature(data)326 data = ensureAmpSignature(data)
280 327
328 // Normalize tool names to canonical casing
329 data = normalizeAmpToolNames(data)
330
281 // Rewrite model name331 // Rewrite model name
282 if rw.originalModel != "" {332 if rw.originalModel != "" {
283 for _, path := range modelFieldPaths {333 for _, path := range modelFieldPaths {
internal/api/modules/amp/response_rewriter_test.go+51−0
@@ -175,6 +175,57 @@ func TestSanitizeAmpRequestBody_MixedInvalidThinkingAndToolUseSignature(t *testi · 174 unmodified lines
175 }175 }
176}176}
177 177
178func TestNormalizeAmpToolNames_NonStreaming(t *testing.T) {
179 input := []byte(`{"content":[{"type":"tool_use","id":"toolu_01","name":"bash","input":{"cmd":"ls"}},{"type":"tool_use","id":"toolu_02","name":"read","input":{"path":"/tmp"}},{"type":"text","text":"hello"}]}`)
180 result := normalizeAmpToolNames(input)
181
182 if !contains(result, []byte(`"name":"Bash"`)) {
183 t.Errorf("expected bash->Bash, got %s", string(result))
184 }
185 if !contains(result, []byte(`"name":"Read"`)) {
186 t.Errorf("expected read->Read, got %s", string(result))
187 }
188 if contains(result, []byte(`"name":"bash"`)) {
189 t.Errorf("expected lowercase bash to be replaced, got %s", string(result))
190 }
191}
192
193func TestNormalizeAmpToolNames_Streaming(t *testing.T) {
194 input := []byte(`{"type":"content_block_start","index":1,"content_block":{"type":"tool_use","name":"grep","id":"toolu_01","input":{}}}`)
195 result := normalizeAmpToolNames(input)
196
197 if !contains(result, []byte(`"name":"Grep"`)) {
198 t.Errorf("expected grep->Grep in streaming, got %s", string(result))
199 }
200}
201
202func TestNormalizeAmpToolNames_AlreadyCorrect(t *testing.T) {
203 input := []byte(`{"content":[{"type":"tool_use","id":"toolu_01","name":"Bash","input":{"cmd":"ls"}}]}`)
204 result := normalizeAmpToolNames(input)
205
206 if string(result) != string(input) {
207 t.Errorf("expected no modification for correctly-cased tool, got %s", string(result))
208 }
209}
210
211func TestNormalizeAmpToolNames_GlobPreserved(t *testing.T) {
212 input := []byte(`{"content":[{"type":"tool_use","id":"toolu_01","name":"glob","input":{"pattern":"*.go"}}]}`)
213 result := normalizeAmpToolNames(input)
214
215 if string(result) != string(input) {
216 t.Errorf("expected glob to remain lowercase, got %s", string(result))
217 }
218}
219
220func TestNormalizeAmpToolNames_UnknownToolUntouched(t *testing.T) {
221 input := []byte(`{"content":[{"type":"tool_use","id":"toolu_01","name":"edit_file","input":{"path":"/tmp/x"}}]}`)
222 result := normalizeAmpToolNames(input)
223
224 if string(result) != string(input) {
225 t.Errorf("expected no modification for unknown tool, got %s", string(result))
226 }
227}
228
178func contains(data, substr []byte) bool {229func contains(data, substr []byte) bool {
179 for i := 0; i <= len(data)-len(substr); i++ {230 for i := 0; i <= len(data)-len(substr); i++ {
180 if string(data[i:i+len(substr)]) == string(substr) {231 if string(data[i:i+len(substr)]) == string(substr) {
internal/runtime/executor/claude_executor.go+87−66
@@ -65,14 +65,13 @@ var oauthToolRenameMap = map[string]string{ · 64 unmodified lines
65 "notebookedit": "NotebookEdit",65 "notebookedit": "NotebookEdit",
66}66}
67 67
68// oauthToolRenameReverseMap is the inverse of oauthToolRenameMap for response decoding.68// The reverse map is now computed per-request in remapOAuthToolNames so that
69var oauthToolRenameReverseMap = func() map[string]string {69// only names the client actually caused us to rewrite are restored on the
70 m := make(map[string]string, len(oauthToolRenameMap))70// response. A global reverse map — as used previously — corrupted responses
71 for k, v := range oauthToolRenameMap {71// for clients that sent mixed casing (e.g. Amp CLI sends `Bash` TitleCase
72 m[v] = k72// alongside `glob` lowercase; the request flagged renames via `glob→Glob`,
73 }73// then the global reverse map incorrectly rewrote every `Bash` in the
74 return m74// response to `bash`, causing Amp to reject the tool_use as unknown).
75}()
76 75
77// oauthToolsToRemove lists tool names that must be stripped from OAuth requests76// oauthToolsToRemove lists tool names that must be stripped from OAuth requests
78// even after remapping. Currently empty — all tools are mapped instead of removed.77// even after remapping. Currently empty — all tools are mapped instead of removed.
@@ -192,15 +191,9 @@ func (e *ClaudeExecutor) Execute(ctx context.Context, auth *cliproxyauth.Auth, r · 113 unmodified lines
192 bodyForTranslation := body191 bodyForTranslation := body
193 bodyForUpstream := body192 bodyForUpstream := body
194 oauthToken := isClaudeOAuthToken(apiKey)193 oauthToken := isClaudeOAuthToken(apiKey)
195 oauthToolNamesRemapped := false194 var oauthToolNamesReverseMap map[string]string
196 if oauthToken && !auth.ToolPrefixDisabled() {
197 bodyForUpstream = applyClaudeToolPrefix(body, claudeToolPrefix)
198 }
199 // Remap third-party tool names to Claude Code equivalents and remove
200 // tools without official counterparts. This prevents Anthropic from
201 // fingerprinting the request as third-party via tool naming patterns.
202 if oauthToken {195 if oauthToken {
203 bodyForUpstream, oauthToolNamesRemapped = remapOAuthToolNames(bodyForUpstream)196 bodyForUpstream, oauthToolNamesReverseMap = prepareClaudeOAuthToolNamesForUpstream(bodyForUpstream, claudeToolPrefix, auth.ToolPrefixDisabled())
204 }197 }
205 // Enable cch signing by default for OAuth tokens (not just experimental flag).198 // Enable cch signing by default for OAuth tokens (not just experimental flag).
206 // Claude Code always computes cch; missing or invalid cch is a detectable fingerprint.199 // Claude Code always computes cch; missing or invalid cch is a detectable fingerprint.
@@ -298,13 +291,7 @@ func (e *ClaudeExecutor) Execute(ctx context.Context, auth *cliproxyauth.Auth, r · 91 unmodified lines
298 } else {291 } else {
299 reporter.Publish(ctx, helps.ParseClaudeUsage(data))292 reporter.Publish(ctx, helps.ParseClaudeUsage(data))
300 }293 }
301 if isClaudeOAuthToken(apiKey) && !auth.ToolPrefixDisabled() {294 data = restoreClaudeOAuthToolNamesFromResponse(data, claudeToolPrefix, auth.ToolPrefixDisabled(), oauthToolNamesReverseMap)
302 data = stripClaudeToolPrefixFromResponse(data, claudeToolPrefix)
303 }
304 // Reverse the OAuth tool name remap so the downstream client sees original names.
305 if isClaudeOAuthToken(apiKey) && oauthToolNamesRemapped {
306 data = reverseRemapOAuthToolNames(data)
307 }
308 var param any295 var param any
309 out := sdktranslator.TranslateNonStream(296 out := sdktranslator.TranslateNonStream(
310 ctx,297 ctx,
@@ -379,15 +366,9 @@ func (e *ClaudeExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth.A · 68 unmodified lines
379 bodyForTranslation := body366 bodyForTranslation := body
380 bodyForUpstream := body367 bodyForUpstream := body
381 oauthToken := isClaudeOAuthToken(apiKey)368 oauthToken := isClaudeOAuthToken(apiKey)
382 oauthToolNamesRemapped := false369 var oauthToolNamesReverseMap map[string]string
383 if oauthToken && !auth.ToolPrefixDisabled() {
384 bodyForUpstream = applyClaudeToolPrefix(body, claudeToolPrefix)
385 }
386 // Remap third-party tool names to Claude Code equivalents and remove
387 // tools without official counterparts. This prevents Anthropic from
388 // fingerprinting the request as third-party via tool naming patterns.
389 if oauthToken {370 if oauthToken {
390 bodyForUpstream, oauthToolNamesRemapped = remapOAuthToolNames(bodyForUpstream)371 bodyForUpstream, oauthToolNamesReverseMap = prepareClaudeOAuthToolNamesForUpstream(bodyForUpstream, claudeToolPrefix, auth.ToolPrefixDisabled())
391 }372 }
392 // Enable cch signing by default for OAuth tokens (not just experimental flag).373 // Enable cch signing by default for OAuth tokens (not just experimental flag).
393 if oauthToken || experimentalCCHSigningEnabled(e.cfg, auth) {374 if oauthToken || experimentalCCHSigningEnabled(e.cfg, auth) {
@@ -478,12 +459,7 @@ func (e *ClaudeExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth.A · 84 unmodified lines
478 if detail, ok := helps.ParseClaudeStreamUsage(line); ok {459 if detail, ok := helps.ParseClaudeStreamUsage(line); ok {
479 reporter.Publish(ctx, detail)460 reporter.Publish(ctx, detail)
480 }461 }
481 if isClaudeOAuthToken(apiKey) && !auth.ToolPrefixDisabled() {462 line = restoreClaudeOAuthToolNamesFromStreamLine(line, claudeToolPrefix, auth.ToolPrefixDisabled(), oauthToolNamesReverseMap)
482 line = stripClaudeToolPrefixFromStreamLine(line, claudeToolPrefix)
483 }
484 if isClaudeOAuthToken(apiKey) && oauthToolNamesRemapped {
485 line = reverseRemapOAuthToolNamesFromStreamLine(line)
486 }
487 // Forward the line as-is to preserve SSE format463 // Forward the line as-is to preserve SSE format
488 cloned := make([]byte, len(line)+1)464 cloned := make([]byte, len(line)+1)
489 copy(cloned, line)465 copy(cloned, line)
@@ -515,12 +491,7 @@ func (e *ClaudeExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth.A · 25 unmodified lines
515 if detail, ok := helps.ParseClaudeStreamUsage(line); ok {491 if detail, ok := helps.ParseClaudeStreamUsage(line); ok {
516 reporter.Publish(ctx, detail)492 reporter.Publish(ctx, detail)
517 }493 }
518 if isClaudeOAuthToken(apiKey) && !auth.ToolPrefixDisabled() {494 line = restoreClaudeOAuthToolNamesFromStreamLine(line, claudeToolPrefix, auth.ToolPrefixDisabled(), oauthToolNamesReverseMap)
519 line = stripClaudeToolPrefixFromStreamLine(line, claudeToolPrefix)
520 }
521 if isClaudeOAuthToken(apiKey) && oauthToolNamesRemapped {
522 line = reverseRemapOAuthToolNamesFromStreamLine(line)
523 }
524 chunks := sdktranslator.TranslateStream(495 chunks := sdktranslator.TranslateStream(
525 ctx,496 ctx,
526 to,497 to,
@@ -635,12 +606,8 @@ func (e *ClaudeExecutor) CountTokens(ctx context.Context, auth *cliproxyauth.Aut · 108 unmodified lines
635 // Extract betas from body and convert to header (for count_tokens too)606 // Extract betas from body and convert to header (for count_tokens too)
636 var extraBetas []string607 var extraBetas []string
637 extraBetas, body = extractAndRemoveBetas(body)608 extraBetas, body = extractAndRemoveBetas(body)
638 if isClaudeOAuthToken(apiKey) && !auth.ToolPrefixDisabled() {
639 body = applyClaudeToolPrefix(body, claudeToolPrefix)
640 }
641 // Remap tool names for OAuth token requests to avoid third-party fingerprinting.
642 if isClaudeOAuthToken(apiKey) {609 if isClaudeOAuthToken(apiKey) {
643 body, _ = remapOAuthToolNames(body)610 body, _ = prepareClaudeOAuthToolNamesForUpstream(body, claudeToolPrefix, auth.ToolPrefixDisabled())
644 }611 }
645 612
646 url := fmt.Sprintf("%s/v1/messages/count_tokens?beta=true", baseURL)613 url := fmt.Sprintf("%s/v1/messages/count_tokens?beta=true", baseURL)
@@ -1080,6 +1047,36 @@ func isClaudeOAuthToken(apiKey string) bool { · 433 unmodified lines
1080 return strings.Contains(apiKey, "sk-ant-oat")1047 return strings.Contains(apiKey, "sk-ant-oat")
1081}1048}
1082 1049
1050// prepareClaudeOAuthToolNamesForUpstream applies the Claude OAuth tool-name
1051// transforms in the same order across request paths. Remap runs before prefixing
1052// so any future non-empty prefix still composes correctly with the per-request
1053// reverse map.
1054func prepareClaudeOAuthToolNamesForUpstream(body []byte, prefix string, prefixDisabled bool) ([]byte, map[string]string) {
1055 body, reverseMap := remapOAuthToolNames(body)
1056 if !prefixDisabled {
1057 body = applyClaudeToolPrefix(body, prefix)
1058 }
1059 return body, reverseMap
1060}
1061
1062// restoreClaudeOAuthToolNamesFromResponse undoes the Claude OAuth tool-name
1063// transforms for non-stream responses in reverse order.
1064func restoreClaudeOAuthToolNamesFromResponse(body []byte, prefix string, prefixDisabled bool, reverseMap map[string]string) []byte {
1065 if !prefixDisabled {
1066 body = stripClaudeToolPrefixFromResponse(body, prefix)
1067 }
1068 return reverseRemapOAuthToolNames(body, reverseMap)
1069}
1070
1071// restoreClaudeOAuthToolNamesFromStreamLine undoes the Claude OAuth tool-name
1072// transforms for SSE lines in reverse order.
1073func restoreClaudeOAuthToolNamesFromStreamLine(line []byte, prefix string, prefixDisabled bool, reverseMap map[string]string) []byte {
1074 if !prefixDisabled {
1075 line = stripClaudeToolPrefixFromStreamLine(line, prefix)
1076 }
1077 return reverseRemapOAuthToolNamesFromStreamLine(line, reverseMap)
1078}
1079
1083// remapOAuthToolNames renames third-party tool names to Claude Code equivalents1080// remapOAuthToolNames renames third-party tool names to Claude Code equivalents
1084// and removes tools without an official counterpart. This prevents Anthropic from1081// and removes tools without an official counterpart. This prevents Anthropic from
1085// fingerprinting the request as a third-party client via tool naming patterns.1082// fingerprinting the request as a third-party client via tool naming patterns.
@@ -1087,8 +1084,25 @@ func isClaudeOAuthToken(apiKey string) bool { · 1 unmodified line
1087// It operates on: tools[].name, tool_choice.name, and all tool_use/tool_reference1084// It operates on: tools[].name, tool_choice.name, and all tool_use/tool_reference
1088// references in messages. Removed tools' corresponding tool_result blocks are preserved1085// references in messages. Removed tools' corresponding tool_result blocks are preserved
1089// (they just become orphaned, which is safe for Claude).1086// (they just become orphaned, which is safe for Claude).
1090func remapOAuthToolNames(body []byte) ([]byte, bool) {1087//
1091 renamed := false1088// The returned map is keyed on the upstream (TitleCase) name and maps to the
1089// client-supplied original name. Callers MUST pass this map to the reverse
1090// functions so only names the client actually caused us to rewrite are restored
1091// on the response. A global reverse map (the previous implementation) incorrectly
1092// rewrote names the client originally sent in TitleCase (e.g. Amp CLI's `Bash`)
1093// when any OTHER tool in the same request triggered a forward rename (e.g.
1094// Amp's `glob`→`Glob`), because the global reverse map contained `Bash`→`bash`
1095// regardless of what the client originally sent.
1096func remapOAuthToolNames(body []byte) ([]byte, map[string]string) {
1097 reverseMap := make(map[string]string, len(oauthToolRenameMap))
1098 recordRename := func(original, renamed string) {
1099 // Preserve the first-seen original name if the same upstream name is
1100 // produced from multiple call sites; they all map back identically.
1101 if _, exists := reverseMap[renamed]; !exists {
1102 reverseMap[renamed] = original
1103 }
1104 }
1105
1092 // 1. Rewrite tools array in a single pass (if present).1106 // 1. Rewrite tools array in a single pass (if present).
1093 // IMPORTANT: do not mutate names first and then rebuild from an older gjson1107 // IMPORTANT: do not mutate names first and then rebuild from an older gjson
1094 // snapshot. gjson results are snapshots of the original bytes; rebuilding from a1108 // snapshot. gjson results are snapshots of the original bytes; rebuilding from a
@@ -1121,7 +1135,7 @@ func remapOAuthToolNames(body []byte) ([]byte, bool) { · 26 unmodified lines
1121 updatedTool, err := sjson.Set(toolJSON, "name", newName)1135 updatedTool, err := sjson.Set(toolJSON, "name", newName)
1122 if err == nil {1136 if err == nil {
1123 toolJSON = updatedTool1137 toolJSON = updatedTool
1124 renamed = true1138 recordRename(name, newName)
1125 }1139 }
1126 }1140 }
1127 1141
@@ -1146,7 +1160,7 @@ func remapOAuthToolNames(body []byte) ([]byte, bool) { · 18 unmodified lines
1146 body, _ = sjson.DeleteBytes(body, "tool_choice")1160 body, _ = sjson.DeleteBytes(body, "tool_choice")
1147 } else if newName, ok := oauthToolRenameMap[tcName]; ok && newName != tcName {1161 } else if newName, ok := oauthToolRenameMap[tcName]; ok && newName != tcName {
1148 body, _ = sjson.SetBytes(body, "tool_choice.name", newName)1162 body, _ = sjson.SetBytes(body, "tool_choice.name", newName)
1149 renamed = true1163 recordRename(tcName, newName)
1150 }1164 }
1151 }1165 }
1152 1166
@@ -1166,14 +1180,14 @@ func remapOAuthToolNames(body []byte) ([]byte, bool) { · 13 unmodified lines
1166 if newName, ok := oauthToolRenameMap[name]; ok && newName != name {1180 if newName, ok := oauthToolRenameMap[name]; ok && newName != name {
1167 path := fmt.Sprintf("messages.%d.content.%d.name", msgIndex.Int(), contentIndex.Int())1181 path := fmt.Sprintf("messages.%d.content.%d.name", msgIndex.Int(), contentIndex.Int())
1168 body, _ = sjson.SetBytes(body, path, newName)1182 body, _ = sjson.SetBytes(body, path, newName)
1169 renamed = true1183 recordRename(name, newName)
1170 }1184 }
1171 case "tool_reference":1185 case "tool_reference":
1172 toolName := part.Get("tool_name").String()1186 toolName := part.Get("tool_name").String()
1173 if newName, ok := oauthToolRenameMap[toolName]; ok && newName != toolName {1187 if newName, ok := oauthToolRenameMap[toolName]; ok && newName != toolName {
1174 path := fmt.Sprintf("messages.%d.content.%d.tool_name", msgIndex.Int(), contentIndex.Int())1188 path := fmt.Sprintf("messages.%d.content.%d.tool_name", msgIndex.Int(), contentIndex.Int())
1175 body, _ = sjson.SetBytes(body, path, newName)1189 body, _ = sjson.SetBytes(body, path, newName)
1176 renamed = true1190 recordRename(toolName, newName)
1177 }1191 }
1178 case "tool_result":1192 case "tool_result":
1179 // Handle nested tool_reference blocks inside tool_result.content[]1193 // Handle nested tool_reference blocks inside tool_result.content[]
@@ -1187,7 +1201,7 @@ func remapOAuthToolNames(body []byte) ([]byte, bool) { · 7 unmodified lines
1187 if newName, ok := oauthToolRenameMap[nestedToolName]; ok && newName != nestedToolName {1201 if newName, ok := oauthToolRenameMap[nestedToolName]; ok && newName != nestedToolName {
1188 nestedPath := fmt.Sprintf("messages.%d.content.%d.content.%d.tool_name", msgIndex.Int(), contentIndex.Int(), nestedIndex.Int())1202 nestedPath := fmt.Sprintf("messages.%d.content.%d.content.%d.tool_name", msgIndex.Int(), contentIndex.Int(), nestedIndex.Int())
1189 body, _ = sjson.SetBytes(body, nestedPath, newName)1203 body, _ = sjson.SetBytes(body, nestedPath, newName)
1190 renamed = true1204 recordRename(nestedToolName, newName)
1191 }1205 }
1192 }1206 }
1193 return true1207 return true
@@ -1200,13 +1214,16 @@ func remapOAuthToolNames(body []byte) ([]byte, bool) { · 6 unmodified lines
1200 })1214 })
1201 }1215 }
1202 1216
1203 return body, renamed1217 return body, reverseMap
1204}1218}
1205 1219
1206// reverseRemapOAuthToolNames reverses the tool name mapping for non-stream responses.1220// reverseRemapOAuthToolNames reverses the tool name mapping for non-stream responses
1207// It maps Claude Code TitleCase names back to the original lowercase names so the1221// using the per-request map produced by remapOAuthToolNames. Names the client sent
1208// downstream client receives tool names it recognizes.1222// that were NOT forward-renamed are passed through unchanged.
1209func reverseRemapOAuthToolNames(body []byte) []byte {1223func reverseRemapOAuthToolNames(body []byte, reverseMap map[string]string) []byte {
1224 if len(reverseMap) == 0 {
1225 return body
1226 }
1210 content := gjson.GetBytes(body, "content")1227 content := gjson.GetBytes(body, "content")
1211 if !content.Exists() || !content.IsArray() {1228 if !content.Exists() || !content.IsArray() {
1212 return body1229 return body
@@ -1216,13 +1233,13 @@ func reverseRemapOAuthToolNames(body []byte) []byte { · 3 unmodified lines
1216 switch partType {1233 switch partType {
1217 case "tool_use":1234 case "tool_use":
1218 name := part.Get("name").String()1235 name := part.Get("name").String()
1219 if origName, ok := oauthToolRenameReverseMap[name]; ok {1236 if origName, ok := reverseMap[name]; ok {
1220 path := fmt.Sprintf("content.%d.name", index.Int())1237 path := fmt.Sprintf("content.%d.name", index.Int())
1221 body, _ = sjson.SetBytes(body, path, origName)1238 body, _ = sjson.SetBytes(body, path, origName)
1222 }1239 }
1223 case "tool_reference":1240 case "tool_reference":
1224 toolName := part.Get("tool_name").String()1241 toolName := part.Get("tool_name").String()
1225 if origName, ok := oauthToolRenameReverseMap[toolName]; ok {1242 if origName, ok := reverseMap[toolName]; ok {
1226 path := fmt.Sprintf("content.%d.tool_name", index.Int())1243 path := fmt.Sprintf("content.%d.tool_name", index.Int())
1227 body, _ = sjson.SetBytes(body, path, origName)1244 body, _ = sjson.SetBytes(body, path, origName)
1228 }1245 }
@@ -1232,8 +1249,12 @@ func reverseRemapOAuthToolNames(body []byte) []byte { · 3 unmodified lines
1232 return body1249 return body
1233}1250}
1234 1251
1235// reverseRemapOAuthToolNamesFromStreamLine reverses the tool name mapping for SSE stream lines.1252// reverseRemapOAuthToolNamesFromStreamLine reverses the tool name mapping for SSE
1236func reverseRemapOAuthToolNamesFromStreamLine(line []byte) []byte {1253// stream lines, using the per-request reverseMap produced by remapOAuthToolNames.
1254func reverseRemapOAuthToolNamesFromStreamLine(line []byte, reverseMap map[string]string) []byte {
1255 if len(reverseMap) == 0 {
1256 return line
1257 }
1237 payload := helps.JSONPayload(line)1258 payload := helps.JSONPayload(line)
1238 if len(payload) == 0 || !gjson.ValidBytes(payload) {1259 if len(payload) == 0 || !gjson.ValidBytes(payload) {
1239 return line1260 return line
@@ -1251,7 +1272,7 @@ func reverseRemapOAuthToolNamesFromStreamLine(line []byte) []byte { · 11 unmodified lines
1251 switch blockType {1272 switch blockType {
1252 case "tool_use":1273 case "tool_use":
1253 name := contentBlock.Get("name").String()1274 name := contentBlock.Get("name").String()
1254 if origName, ok := oauthToolRenameReverseMap[name]; ok {1275 if origName, ok := reverseMap[name]; ok {
1255 updated, err = sjson.SetBytes(payload, "content_block.name", origName)1276 updated, err = sjson.SetBytes(payload, "content_block.name", origName)
1256 if err != nil {1277 if err != nil {
1257 return line1278 return line
@@ -1261,7 +1282,7 @@ func reverseRemapOAuthToolNamesFromStreamLine(line []byte) []byte { · 3 unmodified lines
1261 }1282 }
1262 case "tool_reference":1283 case "tool_reference":
1263 toolName := contentBlock.Get("tool_name").String()1284 toolName := contentBlock.Get("tool_name").String()
1264 if origName, ok := oauthToolRenameReverseMap[toolName]; ok {1285 if origName, ok := reverseMap[toolName]; ok {
1265 updated, err = sjson.SetBytes(payload, "content_block.tool_name", origName)1286 updated, err = sjson.SetBytes(payload, "content_block.tool_name", origName)
1266 if err != nil {1287 if err != nil {
1267 return line1288 return line
internal/runtime/executor/claude_executor_test.go+141−14
@@ -2096,19 +2096,16 @@ func TestNormalizeClaudeTemperatureForThinking_AfterForcedToolChoiceKeepsOrigina · 2095 unmodified lines
2096func TestRemapOAuthToolNames_TitleCase_NoReverseNeeded(t *testing.T) {2096func TestRemapOAuthToolNames_TitleCase_NoReverseNeeded(t *testing.T) {
2097 body := []byte(`{"tools":[{"name":"Bash","description":"Run shell commands","input_schema":{"type":"object","properties":{"cmd":{"type":"string"}}}}],"messages":[{"role":"user","content":[{"type":"text","text":"hi"}]}]}`)2097 body := []byte(`{"tools":[{"name":"Bash","description":"Run shell commands","input_schema":{"type":"object","properties":{"cmd":{"type":"string"}}}}],"messages":[{"role":"user","content":[{"type":"text","text":"hi"}]}]}`)
2098 2098
2099 out, renamed := remapOAuthToolNames(body)2099 out, reverseMap := remapOAuthToolNames(body)
2100 if renamed {2100 if len(reverseMap) != 0 {
2101 t.Fatalf("renamed = true, want false")2101 t.Fatalf("reverseMap = %v, want empty", reverseMap)
2102 }2102 }
2103 if got := gjson.GetBytes(out, "tools.0.name").String(); got != "Bash" {2103 if got := gjson.GetBytes(out, "tools.0.name").String(); got != "Bash" {
2104 t.Fatalf("tools.0.name = %q, want %q", got, "Bash")2104 t.Fatalf("tools.0.name = %q, want %q", got, "Bash")
2105 }2105 }
2106 2106
2107 resp := []byte(`{"content":[{"type":"tool_use","id":"toolu_01","name":"Bash","input":{"cmd":"ls"}}]}`)2107 resp := []byte(`{"content":[{"type":"tool_use","id":"toolu_01","name":"Bash","input":{"cmd":"ls"}}]}`)
2108 reversed := resp2108 reversed := reverseRemapOAuthToolNames(resp, reverseMap)
2109 if renamed {
2110 reversed = reverseRemapOAuthToolNames(resp)
2111 }
2112 if got := gjson.GetBytes(reversed, "content.0.name").String(); got != "Bash" {2109 if got := gjson.GetBytes(reversed, "content.0.name").String(); got != "Bash" {
2113 t.Fatalf("content.0.name = %q, want %q", got, "Bash")2110 t.Fatalf("content.0.name = %q, want %q", got, "Bash")
2114 }2111 }
@@ -2117,20 +2114,150 @@ func TestRemapOAuthToolNames_TitleCase_NoReverseNeeded(t *testing.T) { · 2 unmodified lines
2117func TestRemapOAuthToolNames_Lowercase_ReverseApplied(t *testing.T) {2114func TestRemapOAuthToolNames_Lowercase_ReverseApplied(t *testing.T) {
2118 body := []byte(`{"tools":[{"name":"bash","description":"Run shell commands","input_schema":{"type":"object","properties":{"cmd":{"type":"string"}}}}],"messages":[{"role":"user","content":[{"type":"text","text":"hi"}]}]}`)2115 body := []byte(`{"tools":[{"name":"bash","description":"Run shell commands","input_schema":{"type":"object","properties":{"cmd":{"type":"string"}}}}],"messages":[{"role":"user","content":[{"type":"text","text":"hi"}]}]}`)
2119 2116
2120 out, renamed := remapOAuthToolNames(body)2117 out, reverseMap := remapOAuthToolNames(body)
2121 if !renamed {2118 if reverseMap["Bash"] != "bash" {
2122 t.Fatalf("renamed = false, want true")2119 t.Fatalf("reverseMap = %v, want entry Bash->bash", reverseMap)
2123 }2120 }
2124 if got := gjson.GetBytes(out, "tools.0.name").String(); got != "Bash" {2121 if got := gjson.GetBytes(out, "tools.0.name").String(); got != "Bash" {
2125 t.Fatalf("tools.0.name = %q, want %q", got, "Bash")2122 t.Fatalf("tools.0.name = %q, want %q", got, "Bash")
2126 }2123 }
2127 2124
2128 resp := []byte(`{"content":[{"type":"tool_use","id":"toolu_01","name":"Bash","input":{"cmd":"ls"}}]}`)2125 resp := []byte(`{"content":[{"type":"tool_use","id":"toolu_01","name":"Bash","input":{"cmd":"ls"}}]}`)
2129 reversed := resp2126 reversed := reverseRemapOAuthToolNames(resp, reverseMap)
2130 if renamed {
2131 reversed = reverseRemapOAuthToolNames(resp)
2132 }
2133 if got := gjson.GetBytes(reversed, "content.0.name").String(); got != "bash" {2127 if got := gjson.GetBytes(reversed, "content.0.name").String(); got != "bash" {
2134 t.Fatalf("content.0.name = %q, want %q", got, "bash")2128 t.Fatalf("content.0.name = %q, want %q", got, "bash")
2135 }2129 }
2136}2130}
2131
2132// TestRemapOAuthToolNames_MixedCase_OnlyRenamedToolsReversed is the regression
2133// test for a case where a single request contains both a TitleCase tool (which
2134// must pass through unchanged) and a lowercase tool that we forward-rename.
2135// Before the fix, triggering ANY forward rename caused the reverse pass to
2136// lowercase every TitleCase tool in the response using a global reverse map,
2137// corrupting tool names the client originally sent in TitleCase (notably Amp
2138// CLI's `Bash`, which its registry lookup cannot find as `bash`).
2139func TestRemapOAuthToolNames_MixedCase_OnlyRenamedToolsReversed(t *testing.T) {
2140 body := []byte(`{"tools":[` +
2141 `{"name":"Bash","input_schema":{"type":"object","properties":{"cmd":{"type":"string"}}}},` +
2142 `{"name":"glob","input_schema":{"type":"object","properties":{"filePattern":{"type":"string"}}}}` +
2143 `]}`)
2144
2145 out, reverseMap := remapOAuthToolNames(body)
2146
2147 // Forward: TitleCase `Bash` is not a forward-map key, must pass through.
2148 if got := gjson.GetBytes(out, "tools.0.name").String(); got != "Bash" {
2149 t.Fatalf("tools.0.name = %q, want %q (TitleCase tool must not be renamed)", got, "Bash")
2150 }
2151 // Forward: `glob` is a forward-map key, upstream sees `Glob`.
2152 if got := gjson.GetBytes(out, "tools.1.name").String(); got != "Glob" {
2153 t.Fatalf("tools.1.name = %q, want %q", got, "Glob")
2154 }
2155
2156 // Reverse map records ONLY the rename that happened.
2157 if len(reverseMap) != 1 || reverseMap["Glob"] != "glob" {
2158 t.Fatalf("reverseMap = %v, want {Glob:glob}", reverseMap)
2159 }
2160
2161 // Upstream responds with a `Bash` tool_use. Since we never renamed `Bash`,
2162 // reverseRemap MUST leave it alone.
2163 bashResp := []byte(`{"content":[{"type":"tool_use","id":"toolu_01","name":"Bash","input":{"cmd":"ls"}}]}`)
2164 reversed := reverseRemapOAuthToolNames(bashResp, reverseMap)
2165 if got := gjson.GetBytes(reversed, "content.0.name").String(); got != "Bash" {
2166 t.Fatalf("content.0.name = %q, want %q (Bash must be preserved; was never forward-renamed)", got, "Bash")
2167 }
2168
2169 // Upstream responds with a `Glob` tool_use. Since we renamed `glob`→`Glob`,
2170 // reverseRemap MUST restore the original `glob`.
2171 globResp := []byte(`{"content":[{"type":"tool_use","id":"toolu_02","name":"Glob","input":{"filePattern":"**/*.go"}}]}`)
2172 reversed = reverseRemapOAuthToolNames(globResp, reverseMap)
2173 if got := gjson.GetBytes(reversed, "content.0.name").String(); got != "glob" {
2174 t.Fatalf("content.0.name = %q, want %q (Glob must be restored to client's original `glob`)", got, "glob")
2175 }
2176}
2177
2178// TestReverseRemapOAuthToolNamesFromStreamLine_HonorsPerRequestMap guards the
2179// SSE streaming code path against the same mixed-case bug.
2180func TestReverseRemapOAuthToolNamesFromStreamLine_HonorsPerRequestMap(t *testing.T) {
2181 reverseMap := map[string]string{"Glob": "glob"}
2182
2183 // Bash block was never renamed, must pass through as-is.
2184 bashLine := []byte(`data: {"type":"content_block_start","index":0,"content_block":{"type":"tool_use","id":"toolu_01","name":"Bash","input":{}}}`)
2185 out := reverseRemapOAuthToolNamesFromStreamLine(bashLine, reverseMap)
2186 if !bytes.Contains(out, []byte(`"name":"Bash"`)) {
2187 t.Fatalf("Bash should be preserved, got: %s", string(out))
2188 }
2189 if bytes.Contains(out, []byte(`"name":"bash"`)) {
2190 t.Fatalf("Bash must not be lowercased, got: %s", string(out))
2191 }
2192
2193 // Glob block IS in the reverseMap, must be restored to `glob`.
2194 globLine := []byte(`data: {"type":"content_block_start","index":0,"content_block":{"type":"tool_use","id":"toolu_02","name":"Glob","input":{}}}`)
2195 out = reverseRemapOAuthToolNamesFromStreamLine(globLine, reverseMap)
2196 if !bytes.Contains(out, []byte(`"name":"glob"`)) {
2197 t.Fatalf("Glob should be restored to glob, got: %s", string(out))
2198 }
2199}
2200
2201func TestPrepareClaudeOAuthToolNamesForUpstream_MixedCaseWithPrefix(t *testing.T) {
2202 body := []byte(`{"tools":[` +
2203 `{"name":"Bash","input_schema":{"type":"object","properties":{"cmd":{"type":"string"}}}},` +
2204 `{"name":"glob","input_schema":{"type":"object","properties":{"filePattern":{"type":"string"}}}}` +
2205 `],"messages":[{"role":"assistant","content":[` +
2206 `{"type":"tool_use","id":"toolu_01","name":"Bash","input":{}},` +
2207 `{"type":"tool_use","id":"toolu_02","name":"glob","input":{}}` +
2208 `]}]}`)
2209
2210 out, reverseMap := prepareClaudeOAuthToolNamesForUpstream(body, "proxy_", false)
2211
2212 if got := gjson.GetBytes(out, "tools.0.name").String(); got != "proxy_Bash" {
2213 t.Fatalf("tools.0.name = %q, want %q", got, "proxy_Bash")
2214 }
2215 if got := gjson.GetBytes(out, "tools.1.name").String(); got != "proxy_Glob" {
2216 t.Fatalf("tools.1.name = %q, want %q", got, "proxy_Glob")
2217 }
2218 if got := gjson.GetBytes(out, "messages.0.content.0.name").String(); got != "proxy_Bash" {
2219 t.Fatalf("messages.0.content.0.name = %q, want %q", got, "proxy_Bash")
2220 }
2221 if got := gjson.GetBytes(out, "messages.0.content.1.name").String(); got != "proxy_Glob" {
2222 t.Fatalf("messages.0.content.1.name = %q, want %q", got, "proxy_Glob")
2223 }
2224 if len(reverseMap) != 1 || reverseMap["Glob"] != "glob" {
2225 t.Fatalf("reverseMap = %v, want {Glob:glob}", reverseMap)
2226 }
2227}
2228
2229func TestRestoreClaudeOAuthToolNamesFromResponse_MixedCaseWithPrefix(t *testing.T) {
2230 reverseMap := map[string]string{"Glob": "glob"}
2231 resp := []byte(`{"content":[` +
2232 `{"type":"tool_use","id":"toolu_01","name":"proxy_Bash","input":{}},` +
2233 `{"type":"tool_use","id":"toolu_02","name":"proxy_Glob","input":{}}` +
2234 `]}`)
2235
2236 out := restoreClaudeOAuthToolNamesFromResponse(resp, "proxy_", false, reverseMap)
2237
2238 if got := gjson.GetBytes(out, "content.0.name").String(); got != "Bash" {
2239 t.Fatalf("content.0.name = %q, want %q", got, "Bash")
2240 }
2241 if got := gjson.GetBytes(out, "content.1.name").String(); got != "glob" {
2242 t.Fatalf("content.1.name = %q, want %q", got, "glob")
2243 }
2244}
2245
2246func TestRestoreClaudeOAuthToolNamesFromStreamLine_MixedCaseWithPrefix(t *testing.T) {
2247 reverseMap := map[string]string{"Glob": "glob"}
2248
2249 bashLine := []byte(`data: {"type":"content_block_start","index":0,"content_block":{"type":"tool_use","id":"toolu_01","name":"proxy_Bash","input":{}}}`)
2250 out := restoreClaudeOAuthToolNamesFromStreamLine(bashLine, "proxy_", false, reverseMap)
2251 if !bytes.Contains(out, []byte(`"name":"Bash"`)) {
2252 t.Fatalf("Bash should be preserved, got: %s", string(out))
2253 }
2254 if bytes.Contains(out, []byte(`"name":"bash"`)) {
2255 t.Fatalf("Bash must not be lowercased, got: %s", string(out))
2256 }
2257
2258 globLine := []byte(`data: {"type":"content_block_start","index":0,"content_block":{"type":"tool_use","id":"toolu_02","name":"proxy_Glob","input":{}}}`)
2259 out = restoreClaudeOAuthToolNamesFromStreamLine(globLine, "proxy_", false, reverseMap)
2260 if !bytes.Contains(out, []byte(`"name":"glob"`)) {
2261 t.Fatalf("Glob should be restored to glob, got: %s", string(out))
2262 }
2263}
sdk/api/handlers/openai/openai_responses_handlers.go+131−4
@@ -13,6 +13,7 @@ import ( · 12 unmodified lines
13 "fmt"13 "fmt"
14 "io"14 "io"
15 "net/http"15 "net/http"
16 "sort"
16 17
17 "github.com/gin-gonic/gin"18 "github.com/gin-gonic/gin"
18 . "github.com/router-for-me/CLIProxyAPI/v6/internal/constant"19 . "github.com/router-for-me/CLIProxyAPI/v6/internal/constant"
@@ -45,7 +46,10 @@ func writeResponsesSSEChunk(w io.Writer, chunk []byte) { · 26 unmodified lines
45}46}
46 47
47type responsesSSEFramer struct {48type responsesSSEFramer struct {
48 pending []byte49 pending []byte
50 outputItems map[int][]byte
51 outputOrder []int
52 unindexedOutputItems [][]byte
49}53}
50 54
51func (f *responsesSSEFramer) WriteChunk(w io.Writer, chunk []byte) {55func (f *responsesSSEFramer) WriteChunk(w io.Writer, chunk []byte) {
@@ -61,7 +65,7 @@ func (f *responsesSSEFramer) WriteChunk(w io.Writer, chunk []byte) { · 9 unmodified lines
61 if frameLen == 0 {65 if frameLen == 0 {
62 break66 break
63 }67 }
64 writeResponsesSSEChunk(w, f.pending[:frameLen])68 f.writeFrame(w, f.pending[:frameLen])
65 copy(f.pending, f.pending[frameLen:])69 copy(f.pending, f.pending[frameLen:])
66 f.pending = f.pending[:len(f.pending)-frameLen]70 f.pending = f.pending[:len(f.pending)-frameLen]
67 }71 }
@@ -72,7 +76,7 @@ func (f *responsesSSEFramer) WriteChunk(w io.Writer, chunk []byte) { · 4 unmodified lines
72 if len(f.pending) == 0 || !responsesSSECanEmitWithoutDelimiter(f.pending) {76 if len(f.pending) == 0 || !responsesSSECanEmitWithoutDelimiter(f.pending) {
73 return77 return
74 }78 }
75 writeResponsesSSEChunk(w, f.pending)79 f.writeFrame(w, f.pending)
76 f.pending = f.pending[:0]80 f.pending = f.pending[:0]
77}81}
78 82
@@ -88,10 +92,133 @@ func (f *responsesSSEFramer) Flush(w io.Writer) { · 9 unmodified lines
88 f.pending = f.pending[:0]92 f.pending = f.pending[:0]
89 return93 return
90 }94 }
91 writeResponsesSSEChunk(w, f.pending)95 f.writeFrame(w, f.pending)
92 f.pending = f.pending[:0]96 f.pending = f.pending[:0]
93}97}
94 98
99func (f *responsesSSEFramer) writeFrame(w io.Writer, frame []byte) {
100 writeResponsesSSEChunk(w, f.repairFrame(frame))
101}
102
103func (f *responsesSSEFramer) repairFrame(frame []byte) []byte {
104 payload, ok := responsesSSEDataPayload(frame)
105 if !ok || len(payload) == 0 || bytes.Equal(payload, []byte("[DONE]")) || !json.Valid(payload) {
106 return frame
107 }
108
109 switch gjson.GetBytes(payload, "type").String() {
110 case "response.output_item.done":
111 f.recordOutputItem(payload)
112 case "response.completed":
113 repaired := f.repairCompletedPayload(payload)
114 if !bytes.Equal(repaired, payload) {
115 return responsesSSEFrameWithData(frame, repaired)
116 }
117 }
118 return frame
119}
120
121func responsesSSEDataPayload(frame []byte) ([]byte, bool) {
122 var payload []byte
123 found := false
124 for _, line := range bytes.Split(frame, []byte("\n")) {
125 line = bytes.TrimRight(line, "\r")
126 trimmed := bytes.TrimSpace(line)
127 if !bytes.HasPrefix(trimmed, []byte("data:")) {
128 continue
129 }
130 data := bytes.TrimSpace(trimmed[len("data:"):])
131 if found {
132 payload = append(payload, '\n')
133 }
134 payload = append(payload, data...)
135 found = true
136 }
137 return payload, found
138}
139
140func responsesSSEFrameWithData(frame, payload []byte) []byte {
141 var out bytes.Buffer
142 for _, line := range bytes.Split(frame, []byte("\n")) {
143 line = bytes.TrimRight(line, "\r")
144 trimmed := bytes.TrimSpace(line)
145 if len(trimmed) == 0 || bytes.HasPrefix(trimmed, []byte("data:")) {
146 continue
147 }
148 out.Write(line)
149 out.WriteByte('\n')
150 }
151 for _, line := range bytes.Split(payload, []byte("\n")) {
152 out.WriteString("data: ")
153 out.Write(line)
154 out.WriteByte('\n')
155 }
156 out.WriteByte('\n')
157 return out.Bytes()
158}
159
160func (f *responsesSSEFramer) recordOutputItem(payload []byte) {
161 item := gjson.GetBytes(payload, "item")
162 if !item.Exists() || !item.IsObject() || item.Get("type").String() == "" {
163 return
164 }
165
166 if outputIndex := gjson.GetBytes(payload, "output_index"); outputIndex.Exists() {
167 index := int(outputIndex.Int())
168 if f.outputItems == nil {
169 f.outputItems = make(map[int][]byte)
170 }
171 if _, exists := f.outputItems[index]; !exists {
172 f.outputOrder = append(f.outputOrder, index)
173 }
174 f.outputItems[index] = append([]byte(nil), item.Raw...)
175 return
176 }
177
178 f.unindexedOutputItems = append(f.unindexedOutputItems, append([]byte(nil), item.Raw...))
179}
180
181func (f *responsesSSEFramer) repairCompletedPayload(payload []byte) []byte {
182 if len(f.outputOrder) == 0 && len(f.unindexedOutputItems) == 0 {
183 return payload
184 }
185 output := gjson.GetBytes(payload, "response.output")
186 if output.Exists() && (!output.IsArray() || len(output.Array()) > 0) {
187 return payload
188 }
189
190 var outputJSON bytes.Buffer
191 outputJSON.WriteByte('[')
192 indexes := append([]int(nil), f.outputOrder...)
193 sort.Ints(indexes)
194 written := 0
195 for _, index := range indexes {
196 item, ok := f.outputItems[index]
197 if !ok {
198 continue
199 }
200 if written > 0 {
201 outputJSON.WriteByte(',')
202 }
203 outputJSON.Write(item)
204 written++
205 }
206 for _, item := range f.unindexedOutputItems {
207 if written > 0 {
208 outputJSON.WriteByte(',')
209 }
210 outputJSON.Write(item)
211 written++
212 }
213 outputJSON.WriteByte(']')
214
215 repaired, err := sjson.SetRawBytes(payload, "response.output", outputJSON.Bytes())
216 if err != nil {
217 return payload
218 }
219 return repaired
220}
221
95func responsesSSEFrameLen(chunk []byte) int {222func responsesSSEFrameLen(chunk []byte) int {
96 if len(chunk) == 0 {223 if len(chunk) == 0 {
97 return 0224 return 0
sdk/api/handlers/openai/openai_responses_handlers_stream_test.go+98−1
@@ -10,6 +10,7 @@ import ( · 9 unmodified lines
10 "github.com/router-for-me/CLIProxyAPI/v6/internal/interfaces"10 "github.com/router-for-me/CLIProxyAPI/v6/internal/interfaces"
11 "github.com/router-for-me/CLIProxyAPI/v6/sdk/api/handlers"11 "github.com/router-for-me/CLIProxyAPI/v6/sdk/api/handlers"
12 sdkconfig "github.com/router-for-me/CLIProxyAPI/v6/sdk/config"12 sdkconfig "github.com/router-for-me/CLIProxyAPI/v6/sdk/config"
13 "github.com/tidwall/gjson"
13)14)
14 15
15func newResponsesStreamTestHandler(t *testing.T) (*OpenAIResponsesAPIHandler, *httptest.ResponseRecorder, *gin.Context, http.Flusher) {16func newResponsesStreamTestHandler(t *testing.T) (*OpenAIResponsesAPIHandler, *httptest.ResponseRecorder, *gin.Context, http.Flusher) {
@@ -53,12 +54,108 @@ func TestForwardResponsesStreamSeparatesDataOnlySSEChunks(t *testing.T) { · 37 unmodified lines
53 t.Errorf("unexpected first event.\nGot: %q\nWant: %q", parts[0], expectedPart1)54 t.Errorf("unexpected first event.\nGot: %q\nWant: %q", parts[0], expectedPart1)
54 }55 }
55 56
56 expectedPart2 := "data: {\"type\":\"response.completed\",\"response\":{\"id\":\"resp-1\",\"output\":[]}}"57 expectedPart2 := "data: {\"type\":\"response.completed\",\"response\":{\"id\":\"resp-1\",\"output\":[{\"type\":\"function_call\",\"arguments\":\"{}\"}]}}"
57 if parts[1] != expectedPart2 {58 if parts[1] != expectedPart2 {
58 t.Errorf("unexpected second event.\nGot: %q\nWant: %q", parts[1], expectedPart2)59 t.Errorf("unexpected second event.\nGot: %q\nWant: %q", parts[1], expectedPart2)
59 }60 }
60}61}
61 62
63func TestForwardResponsesStreamRepairsEmptyCompletedOutputFromDoneItems(t *testing.T) {
64 h, recorder, c, flusher := newResponsesStreamTestHandler(t)
65
66 data := make(chan []byte, 3)
67 errs := make(chan *interfaces.ErrorMessage)
68 data <- []byte(`data: {"type":"response.output_item.done","output_index":0,"item":{"type":"reasoning","id":"rs-1","summary":[]}}`)
69 data <- []byte(`data: {"type":"response.output_item.done","output_index":1,"item":{"type":"function_call","id":"fc-1","call_id":"call-1","name":"shell","arguments":"{\"cmd\":\"pwd\"}","status":"completed"}}`)
70 data <- []byte(`data: {"type":"response.completed","response":{"id":"resp-1","output":[]}}`)
71 close(data)
72 close(errs)
73
74 h.forwardResponsesStream(c, flusher, func(error) {}, data, errs, nil)
75
76 parts := strings.Split(strings.TrimSpace(recorder.Body.String()), "\n\n")
77 if len(parts) != 3 {
78 t.Fatalf("expected 3 SSE events, got %d. Body: %q", len(parts), recorder.Body.String())
79 }
80
81 payload := strings.TrimPrefix(parts[2], "data: ")
82 output := gjson.Get(payload, "response.output")
83 if !output.IsArray() || len(output.Array()) != 2 {
84 t.Fatalf("expected repaired completed output with 2 items, got %s", output.Raw)
85 }
86 if got := gjson.Get(payload, "response.output.1.name").String(); got != "shell" {
87 t.Fatalf("expected function_call name to be preserved, got %q in %s", got, payload)
88 }
89 if got := gjson.Get(payload, "response.output.1.arguments").String(); got != `{"cmd":"pwd"}` {
90 t.Fatalf("expected function_call arguments to be preserved, got %q in %s", got, payload)
91 }
92}
93
94func TestForwardResponsesStreamRepairsMixedIndexedAndUnindexedDoneItems(t *testing.T) {
95 h, recorder, c, flusher := newResponsesStreamTestHandler(t)
96
97 data := make(chan []byte, 3)
98 errs := make(chan *interfaces.ErrorMessage)
99 data <- []byte(`data: {"type":"response.output_item.done","output_index":1,"item":{"type":"function_call","id":"fc-1","call_id":"call-1","name":"shell","arguments":"{}","status":"completed"}}`)
100 data <- []byte(`data: {"type":"response.output_item.done","item":{"type":"message","id":"msg-1","role":"assistant","content":[{"type":"output_text","text":"done"}]}}`)
101 data <- []byte(`data: {"type":"response.completed","response":{"id":"resp-1","output":[]}}`)
102 close(data)
103 close(errs)
104
105 h.forwardResponsesStream(c, flusher, func(error) {}, data, errs, nil)
106
107 parts := strings.Split(strings.TrimSpace(recorder.Body.String()), "\n\n")
108 if len(parts) != 3 {
109 t.Fatalf("expected 3 SSE events, got %d. Body: %q", len(parts), recorder.Body.String())
110 }
111
112 payload := strings.TrimPrefix(parts[2], "data: ")
113 output := gjson.Get(payload, "response.output")
114 if !output.IsArray() || len(output.Array()) != 2 {
115 t.Fatalf("expected repaired completed output with 2 items, got %s", output.Raw)
116 }
117 if got := gjson.Get(payload, "response.output.0.name").String(); got != "shell" {
118 t.Fatalf("expected indexed function_call to be preserved first, got %q in %s", got, payload)
119 }
120 if got := gjson.Get(payload, "response.output.1.id").String(); got != "msg-1" {
121 t.Fatalf("expected unindexed message to be appended, got %q in %s", got, payload)
122 }
123}
124
125func TestForwardResponsesStreamRepairsMultilineCompletedOutputAsSSEDataLines(t *testing.T) {
126 h, recorder, c, flusher := newResponsesStreamTestHandler(t)
127
128 data := make(chan []byte, 2)
129 errs := make(chan *interfaces.ErrorMessage)
130 data <- []byte(`data: {"type":"response.output_item.done","item":{"type":"function_call","arguments":"{}"}}`)
131 data <- []byte("data: {\"type\":\"response.completed\",\ndata: \"response\":{\"id\":\"resp-1\",\"output\":[]}}\n\n")
132 close(data)
133 close(errs)
134
135 h.forwardResponsesStream(c, flusher, func(error) {}, data, errs, nil)
136
137 parts := strings.Split(strings.TrimSpace(recorder.Body.String()), "\n\n")
138 if len(parts) != 2 {
139 t.Fatalf("expected 2 SSE events, got %d. Body: %q", len(parts), recorder.Body.String())
140 }
141
142 completedFrame := []byte(parts[1])
143 for _, line := range strings.Split(parts[1], "\n") {
144 if line != "" && !strings.HasPrefix(line, "data: ") {
145 t.Fatalf("expected every completed payload line to be an SSE data line, got %q in %q", line, parts[1])
146 }
147 }
148
149 payload, ok := responsesSSEDataPayload(completedFrame)
150 if !ok {
151 t.Fatalf("expected completed frame to contain data payload: %q", parts[1])
152 }
153 output := gjson.GetBytes(payload, "response.output")
154 if !output.IsArray() || len(output.Array()) != 1 {
155 t.Fatalf("expected repaired completed output with 1 item, got %s from %q", output.Raw, payload)
156 }
157}
158
62func TestForwardResponsesStreamReassemblesSplitSSEEventChunks(t *testing.T) {159func TestForwardResponsesStreamReassemblesSplitSSEEventChunks(t *testing.T) {
63 h, recorder, c, flusher := newResponsesStreamTestHandler(t)160 h, recorder, c, flusher := newResponsesStreamTestHandler(t)
64 161