From 9ede0f71658a32ece3bf649a602aa973d4f662b6 Mon Sep 17 00:00:00 2001 From: Nai Long 7th Date: Thu, 20 Aug 2026 12:37:39 +0800 Subject: [PATCH] fix(grok): promote tool-search discoveries into callable tools --- .../pkg/apicompat/responses_client_tools.go | 65 +++-- .../apicompat/responses_client_tools_test.go | 254 +++++++++++++++-- .../responses_tool_search_discoveries.go | 263 ++++++++++++++++++ .../openai_gateway_grok_tool_protocol.go | 3 +- .../openai_gateway_grok_tool_protocol_test.go | 71 ++++- .../internal/service/openai_ws_http_bridge.go | 65 ++++- .../service/openai_ws_http_bridge_test.go | 131 ++++++++- 7 files changed, 792 insertions(+), 60 deletions(-) create mode 100644 backend/internal/pkg/apicompat/responses_tool_search_discoveries.go diff --git a/backend/internal/pkg/apicompat/responses_client_tools.go b/backend/internal/pkg/apicompat/responses_client_tools.go index 4237a95ff0..b7ce024e8a 100644 --- a/backend/internal/pkg/apicompat/responses_client_tools.go +++ b/backend/internal/pkg/apicompat/responses_client_tools.go @@ -27,6 +27,13 @@ func AdaptResponsesClientTools(req map[string]any) (ResponsesClientToolMapping, if !ok || len(tools) == 0 { return ResponsesClientToolMapping{}, false, nil } + discovered, err := promoteResponsesToolSearchDiscoveries(req) + if err != nil { + return ResponsesClientToolMapping{}, false, err + } + if discovered { + tools, _ = req["tools"].([]any) + } adapter := ResponsesClientToolMapping{CustomTools: make(map[string]bool)} functionNames := make(map[string]bool) @@ -73,7 +80,7 @@ func AdaptResponsesClientTools(req map[string]any) (ResponsesClientToolMapping, tools, _ = req["tools"].([]any) lowered := make([]any, 0, len(tools)) - changed := flattened + changed := discovered || flattened seenSearch := false for _, raw := range tools { tool, ok := raw.(map[string]any) @@ -115,7 +122,11 @@ func AdaptResponsesClientTools(req map[string]any) (ResponsesClientToolMapping, if changed { req["tools"] = lowered } - if rewriteClientToolHistory(req["input"], &adapter) { + historyChanged, err := rewriteClientToolHistory(req["input"], &adapter) + if err != nil { + return ResponsesClientToolMapping{}, false, err + } + if historyChanged { changed = true } if rewriteClientToolChoice(req, &adapter) { @@ -138,6 +149,7 @@ func AdaptResponsesClientTools(req map[string]any) (ResponsesClientToolMapping, func AdaptResponsesClientToolsWithInheritedMapping( req map[string]any, inherited ResponsesClientToolMapping, + inheritedLoweredTools ...[]any, ) (ResponsesClientToolMapping, bool, error) { if req == nil { return ResponsesClientToolMapping{}, false, nil @@ -148,8 +160,15 @@ func AdaptResponsesClientToolsWithInheritedMapping( if len(inherited.CustomTools) == 0 && !inherited.ToolSearch && len(inherited.NamespaceTools) == 0 { return ResponsesClientToolMapping{}, false, nil } + if len(inheritedLoweredTools) > 0 && len(inheritedLoweredTools[0]) > 0 { + req["tools"] = restoreInheritedResponsesClientToolDeclarations(inheritedLoweredTools[0], inherited) + return AdaptResponsesClientTools(req) + } - changed := rewriteClientToolHistory(req["input"], &inherited) + changed, err := rewriteClientToolHistory(req["input"], &inherited) + if err != nil { + return ResponsesClientToolMapping{}, false, err + } if len(inherited.NamespaceTools) > 0 { before := changed rewriteNamespaceQualifiedCalls(req["input"], inherited.NamespaceTools) @@ -174,14 +193,16 @@ func copyClientTool(tool map[string]any) map[string]any { return copy } -func rewriteClientToolHistory(value any, adapter *ResponsesClientToolMapping) bool { +func rewriteClientToolHistory(value any, adapter *ResponsesClientToolMapping) (bool, error) { changed := false - var visit func(any) - visit = func(value any) { + var visit func(any) error + visit = func(value any) error { switch typed := value.(type) { case []any: for _, item := range typed { - visit(item) + if err := visit(item); err != nil { + return err + } } case map[string]any: typ := strings.TrimSpace(stringValue(typed["type"])) @@ -210,19 +231,30 @@ func rewriteClientToolHistory(value any, adapter *ResponsesClientToolMapping) bo } case "tool_search_output": if adapter.ToolSearch { + callID := strings.TrimSpace(stringValue(typed["call_id"])) + if callID == "" { + return fmt.Errorf("tool_search_output requires a non-empty string call_id before it can be lowered to function_call_output") + } typed["type"] = "function_call_output" dropInvalidLoweredFunctionItemID(typed) - normalizeToolSearchOutput(typed) + if err := normalizeToolSearchOutput(typed); err != nil { + return err + } changed = true } } for _, child := range typed { - visit(child) + if err := visit(child); err != nil { + return err + } } } + return nil } - visit(value) - return changed + if err := visit(value); err != nil { + return false, err + } + return changed, nil } // dropInvalidLoweredFunctionItemID removes Codex client-only item IDs such as @@ -261,7 +293,7 @@ func normalizeClientToolOutput(item map[string]any) { // output field directly; newer Codex clients return discovered definitions in // a top-level tools field. Codex treats that field's value as the tool output, // so serialize the value directly rather than wrapping it in another object. -func normalizeToolSearchOutput(item map[string]any) { +func normalizeToolSearchOutput(item map[string]any) error { if output, hasOutput := item["output"]; hasOutput { switch typed := output.(type) { case string: @@ -271,23 +303,24 @@ func normalizeToolSearchOutput(item map[string]any) { default: encoded, err := json.Marshal(typed) if err != nil { - return + return fmt.Errorf("tool_search_output output cannot be encoded as function_call_output output: %w", err) } item["output"] = string(encoded) } dropToolSearchOutputPrivateFields(item) - return + return nil } tools, hasTools := item["tools"] if !hasTools { - return + return fmt.Errorf("tool_search_output requires output or tools before it can be lowered to function_call_output") } encoded, err := json.Marshal(tools) if err != nil { - return + return fmt.Errorf("tool_search_output tools cannot be encoded as function_call_output output: %w", err) } item["output"] = string(encoded) dropToolSearchOutputPrivateFields(item) + return nil } func dropToolSearchOutputPrivateFields(item map[string]any) { diff --git a/backend/internal/pkg/apicompat/responses_client_tools_test.go b/backend/internal/pkg/apicompat/responses_client_tools_test.go index 41c436374a..5b0dac4356 100644 --- a/backend/internal/pkg/apicompat/responses_client_tools_test.go +++ b/backend/internal/pkg/apicompat/responses_client_tools_test.go @@ -73,24 +73,42 @@ func TestAdaptResponsesClientTools_LowersDiscoveredToolSearchOutput(t *testing.T {"type":"tool_search_call","id":"tsc_client","call_id":"call_search","arguments":{"query":"codex app"},"execution":"client","status":"completed"}, {"type":"tool_search_output","id":"tso_client","call_id":"call_search","execution":"client","status":"completed","tools":[ {"type":"namespace","name":"codex_app","tools":[{"type":"function","name":"load_workspace_dependencies","description":"Load workspace dependencies","parameters":{"type":"object","properties":{},"additionalProperties":false}}]}, - {"type":"namespace","name":"multi_agent_v1","tools":[{"type":"function","name":"spawn_agent","description":"Spawn an agent","parameters":{"type":"object","properties":{"message":{"type":"string"}},"required":["message"],"additionalProperties":false}}]} + {"type":"namespace","name":"multi_agent_v1","tools":[ + {"type":"function","name":"spawn_agent","description":"Spawn an agent","parameters":{"type":"object","properties":{"message":{"type":"string"}},"required":["message"],"additionalProperties":false}}, + {"type":"function","name":"wait_agent","description":"Wait for agents","parameters":{"type":"object","properties":{},"additionalProperties":false}} + ]} ]} ] }` - adapt := func() map[string]any { + type adaptedRequest struct { + req map[string]any + mapping ResponsesClientToolMapping + } + adapt := func() adaptedRequest { var req map[string]any require.NoError(t, json.Unmarshal([]byte(requestJSON), &req)) - _, changed, err := AdaptResponsesClientTools(req) + mapping, changed, err := AdaptResponsesClientTools(req) require.NoError(t, err) require.True(t, changed) - return req + return adaptedRequest{req: req, mapping: mapping} } first := adapt() second := adapt() - firstInput := requireResponsesClientToolValue[[]any](t, first["input"]) - secondInput := requireResponsesClientToolValue[[]any](t, second["input"]) + firstInput := requireResponsesClientToolValue[[]any](t, first.req["input"]) + secondInput := requireResponsesClientToolValue[[]any](t, second.req["input"]) + + tools := requireResponsesClientToolValue[[]any](t, first.req["tools"]) + require.Len(t, tools, 4) + require.Equal(t, []string{ + "tool_search", + "codex_app__load_workspace_dependencies", + "multi_agent_v1__spawn_agent", + "multi_agent_v1__wait_agent", + }, responsesClientToolNames(t, tools)) + require.Equal(t, ResponsesNamespaceName{Namespace: "multi_agent_v1", Name: "spawn_agent"}, first.mapping.NamespaceTools["multi_agent_v1__spawn_agent"]) + require.Equal(t, ResponsesNamespaceName{Namespace: "multi_agent_v1", Name: "wait_agent"}, first.mapping.NamespaceTools["multi_agent_v1__wait_agent"]) call := requireResponsesClientToolValue[map[string]any](t, firstInput[0]) require.Equal(t, "function_call", call["type"]) @@ -107,10 +125,141 @@ func TestAdaptResponsesClientTools_LowersDiscoveredToolSearchOutput(t *testing.T outputText := requireResponsesClientToolValue[string](t, output["output"]) require.JSONEq(t, `[ {"type":"namespace","name":"codex_app","tools":[{"type":"function","name":"load_workspace_dependencies","description":"Load workspace dependencies","parameters":{"type":"object","properties":{},"additionalProperties":false}}]}, - {"type":"namespace","name":"multi_agent_v1","tools":[{"type":"function","name":"spawn_agent","description":"Spawn an agent","parameters":{"type":"object","properties":{"message":{"type":"string"}},"required":["message"],"additionalProperties":false}}]} + {"type":"namespace","name":"multi_agent_v1","tools":[ + {"type":"function","name":"spawn_agent","description":"Spawn an agent","parameters":{"type":"object","properties":{"message":{"type":"string"}},"required":["message"],"additionalProperties":false}}, + {"type":"function","name":"wait_agent","description":"Wait for agents","parameters":{"type":"object","properties":{},"additionalProperties":false}} + ]} ]`, outputText) secondOutput := requireResponsesClientToolValue[map[string]any](t, secondInput[1]) require.Equal(t, outputText, secondOutput["output"], "tool discovery output encoding must be deterministic") + + restored, changed, err := RestoreResponsesClientToolPayload( + []byte(`{"output":[{"type":"function_call","name":"multi_agent_v1__spawn_agent","call_id":"call_spawn","arguments":"{\"message\":\"work\"}"}]}`), + first.mapping, + ) + require.NoError(t, err) + require.True(t, changed) + require.JSONEq(t, `{"output":[{"type":"function_call","name":"spawn_agent","namespace":"multi_agent_v1","call_id":"call_spawn","arguments":"{\"message\":\"work\"}"}]}`, string(restored)) +} + +func TestAdaptResponsesClientTools_PromotesDirectDiscoveryAndDeduplicatesIdenticalDeclarations(t *testing.T) { + direct := map[string]any{ + "type": "function", "name": "inspect_result", "description": "Inspect a result", + "parameters": map[string]any{"type": "object", "properties": map[string]any{}}, + } + custom := map[string]any{ + "type": "custom", "name": "run_script", "description": "Run a script", + "format": map[string]any{"type": "grammar"}, + } + namespace := map[string]any{ + "type": "namespace", "name": "multi_agent_v1", "tools": []any{map[string]any{ + "type": "function", "name": "spawn_agent", "parameters": map[string]any{"type": "object"}, + }}, + } + req := map[string]any{ + "tools": []any{ + map[string]any{"type": "function", "name": "static_first", "parameters": map[string]any{"type": "object"}}, + map[string]any{"type": "tool_search"}, + }, + "input": []any{ + map[string]any{"type": "tool_search_output", "status": "completed", "call_id": "search_1", "tools": []any{direct, custom, namespace}}, + map[string]any{"type": "tool_search_output", "status": "completed", "call_id": "search_2", "tools": []any{copyClientTool(direct), copyClientTool(custom), copyClientTool(namespace)}}, + }, + } + + mapping, changed, err := AdaptResponsesClientTools(req) + require.NoError(t, err) + require.True(t, changed) + require.True(t, mapping.CustomTools["run_script"]) + require.Equal(t, ResponsesNamespaceName{Namespace: "multi_agent_v1", Name: "spawn_agent"}, mapping.NamespaceTools["multi_agent_v1__spawn_agent"]) + tools := requireResponsesClientToolValue[[]any](t, req["tools"]) + require.Equal(t, []string{"static_first", "tool_search", "inspect_result", "run_script", "multi_agent_v1__spawn_agent"}, responsesClientToolNames(t, tools)) + customTool := requireResponsesClientToolValue[map[string]any](t, tools[3]) + require.Equal(t, "function", customTool["type"]) + require.NotContains(t, customTool, "format") + for _, raw := range requireResponsesClientToolValue[[]any](t, req["input"]) { + item := requireResponsesClientToolValue[map[string]any](t, raw) + require.Equal(t, "function_call_output", item["type"]) + require.NotContains(t, item, "tools") + require.NotContains(t, item, "status") + } +} + +func TestAdaptResponsesClientTools_RejectsDiscoveredSchemaAndNamespaceCollisions(t *testing.T) { + tests := []struct { + name string + staticTools []any + discovered []any + }{ + { + name: "direct schema collision", + staticTools: []any{map[string]any{ + "type": "function", "name": "inspect", "parameters": map[string]any{"type": "object"}, + }}, + discovered: []any{map[string]any{ + "type": "function", "name": "inspect", "parameters": map[string]any{"type": "string"}, + }}, + }, + { + name: "namespace schema collision", + staticTools: []any{map[string]any{ + "type": "namespace", "name": "multi_agent_v1", "tools": []any{map[string]any{ + "type": "function", "name": "spawn_agent", "parameters": map[string]any{"type": "object"}, + }}, + }}, + discovered: []any{map[string]any{ + "type": "namespace", "name": "multi_agent_v1", "tools": []any{map[string]any{ + "type": "function", "name": "spawn_agent", "parameters": map[string]any{"type": "string"}, + }}, + }}, + }, + { + name: "flattened namespace collision", + staticTools: []any{map[string]any{"type": "function", "name": "multi_agent_v1__spawn_agent"}}, + discovered: []any{map[string]any{ + "type": "namespace", "name": "multi_agent_v1", "tools": []any{map[string]any{"type": "function", "name": "spawn_agent"}}, + }}, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + req := map[string]any{ + "tools": append(tt.staticTools, map[string]any{"type": "tool_search"}), + "input": []any{map[string]any{ + "type": "tool_search_output", "status": "completed", "tools": tt.discovered, + }}, + } + _, _, err := AdaptResponsesClientTools(req) + require.ErrorContains(t, err, "conflicts") + }) + } +} + +func TestAdaptResponsesClientTools_DoesNotPromoteUnusableDiscoveries(t *testing.T) { + req := map[string]any{ + "tools": []any{map[string]any{"type": "tool_search"}}, + "input": []any{ + map[string]any{"type": "tool_search_output", "call_id": "search_in_progress", "status": "in_progress", "tools": []any{map[string]any{"type": "function", "name": "not_ready"}}}, + map[string]any{"type": "tool_search_output", "call_id": "search_malformed", "status": "completed", "tools": []any{map[string]any{"type": "function"}}}, + }, + } + + _, changed, err := AdaptResponsesClientTools(req) + require.NoError(t, err) + require.True(t, changed, "the static tool_search declaration is still lowered") + tools := requireResponsesClientToolValue[[]any](t, req["tools"]) + require.Equal(t, []string{"tool_search"}, responsesClientToolNames(t, tools)) +} + +func responsesClientToolNames(t *testing.T, tools []any) []string { + t.Helper() + names := make([]string, 0, len(tools)) + for _, raw := range tools { + tool := requireResponsesClientToolValue[map[string]any](t, raw) + names = append(names, requireResponsesClientToolValue[string](t, tool["name"])) + } + return names } func TestAdaptResponsesClientTools_ToolSearchOutputEdgeCases(t *testing.T) { @@ -122,12 +271,13 @@ func TestAdaptResponsesClientTools_ToolSearchOutputEdgeCases(t *testing.T) { wantOutputExists bool wantPrivateKeys []string wantExactOutput bool + wantErr bool }{ { - name: "absent tools and output remains visibly malformed", + name: "absent tools and output is rejected", item: map[string]any{"type": "tool_search_output", "call_id": "call_empty", "status": "completed"}, wantOutputExists: false, - wantPrivateKeys: []string{"status"}, + wantErr: true, }, { name: "preexisting string output wins", @@ -150,14 +300,13 @@ func TestAdaptResponsesClientTools_ToolSearchOutputEdgeCases(t *testing.T) { wantExactOutput: true, }, { - name: "unencodable preexisting output remains visibly malformed", + name: "unencodable preexisting output is rejected", item: map[string]any{ "type": "tool_search_output", "call_id": "call_bad_output", "output": unencodableOutput, "tools": []any{map[string]any{"type": "function", "name": "retained"}}, "status": "completed", "execution": "client", }, - wantOutput: unencodableOutput, - wantOutputExists: true, - wantPrivateKeys: []string{"tools", "status", "execution"}, + wantOutput: unencodableOutput, + wantErr: true, }, { name: "empty tools array is a valid empty output", @@ -180,12 +329,18 @@ func TestAdaptResponsesClientTools_ToolSearchOutputEdgeCases(t *testing.T) { wantExactOutput: true, }, { - name: "unencodable tools remains visibly malformed", + name: "unencodable tools is rejected", item: map[string]any{ "type": "tool_search_output", "call_id": "call_unencodable", "tools": make(chan struct{}), "status": "completed", }, - wantOutputExists: false, - wantPrivateKeys: []string{"tools", "status"}, + wantErr: true, + }, + { + name: "missing call id is rejected", + item: map[string]any{ + "type": "tool_search_output", "tools": []any{}, "status": "completed", + }, + wantErr: true, }, } @@ -196,6 +351,10 @@ func TestAdaptResponsesClientTools_ToolSearchOutputEdgeCases(t *testing.T) { "input": []any{tt.item}, } _, changed, err := AdaptResponsesClientTools(req) + if tt.wantErr { + require.Error(t, err) + return + } require.NoError(t, err) require.True(t, changed) input := requireResponsesClientToolValue[[]any](t, req["input"]) @@ -277,6 +436,47 @@ func TestAdaptResponsesClientToolsWithInheritedMapping_LowersFollowupHistoryWith require.JSONEq(t, `[{"text":"ok","type":"input_text"}]`, requireResponsesClientToolValue[string](t, output["output"])) } +func TestAdaptResponsesClientToolsWithInheritedMapping_PromotesOmittedToolsDiscoveryIntoEffectiveDeclarations(t *testing.T) { + req := map[string]any{ + "input": []any{map[string]any{ + "type": "tool_search_output", "call_id": "call_search", "status": "completed", "execution": "client", + "tools": []any{map[string]any{ + "type": "namespace", "name": "multi_agent_v1", "tools": []any{map[string]any{ + "type": "function", "name": "spawn_agent", "parameters": map[string]any{"type": "object"}, + }}, + }}, + }}, + } + inherited := ResponsesClientToolMapping{ + ToolSearch: true, + NamespaceTools: map[string]ResponsesNamespaceName{ + "codex_app__read_resource": {Namespace: "codex_app", Name: "read_resource"}, + }, + } + lowered := []any{ + map[string]any{"type": "function", "name": "static_first", "parameters": map[string]any{"type": "object"}}, + map[string]any{"type": "function", "name": "tool_search", "parameters": json.RawMessage(toolSearchProxySchema)}, + map[string]any{"type": "function", "name": "codex_app__read_resource", "parameters": map[string]any{"type": "object"}}, + } + + mapping, changed, err := AdaptResponsesClientToolsWithInheritedMapping(req, inherited, lowered) + require.NoError(t, err) + require.True(t, changed) + require.True(t, mapping.ToolSearch) + require.Equal(t, ResponsesNamespaceName{Namespace: "codex_app", Name: "read_resource"}, mapping.NamespaceTools["codex_app__read_resource"]) + require.Equal(t, ResponsesNamespaceName{Namespace: "multi_agent_v1", Name: "spawn_agent"}, mapping.NamespaceTools["multi_agent_v1__spawn_agent"]) + tools := requireResponsesClientToolValue[[]any](t, req["tools"]) + require.Equal(t, []string{ + "static_first", "tool_search", "codex_app__read_resource", "multi_agent_v1__spawn_agent", + }, responsesClientToolNames(t, tools)) + output := requireResponsesClientToolValue[map[string]any](t, requireResponsesClientToolValue[[]any](t, req["input"])[0]) + require.Equal(t, "function_call_output", output["type"]) + require.IsType(t, "", output["output"]) + require.NotContains(t, output, "tools") + require.NotContains(t, output, "status") + require.NotContains(t, output, "execution") +} + func TestAdaptResponsesClientToolsWithInheritedMapping_ExplicitToolsReplaceInheritedMapping(t *testing.T) { req := map[string]any{ "tools": []any{}, @@ -298,6 +498,28 @@ func TestAdaptResponsesClientToolsWithInheritedMapping_ExplicitToolsReplaceInher require.Equal(t, "custom_tool_call", call["type"]) } +func TestAdaptResponsesClientToolsWithInheritedMapping_ExplicitToolResetDoesNotPromoteDiscovery(t *testing.T) { + for _, reset := range []any{nil, []any{}} { + req := map[string]any{ + "tools": reset, + "input": []any{map[string]any{ + "type": "tool_search_output", "call_id": "call_reset", "status": "completed", + "tools": []any{map[string]any{"type": "function", "name": "must_not_promote"}}, + }}, + } + mapping, changed, err := AdaptResponsesClientToolsWithInheritedMapping( + req, + ResponsesClientToolMapping{ToolSearch: true}, + []any{map[string]any{"type": "function", "name": "tool_search"}}, + ) + require.NoError(t, err) + require.False(t, changed) + require.Empty(t, mapping) + item := requireResponsesClientToolValue[map[string]any](t, requireResponsesClientToolValue[[]any](t, req["input"])[0]) + require.Equal(t, "tool_search_output", item["type"]) + } +} + func TestRestoreResponsesClientToolPayload_RestoresClientAndNamespaceCalls(t *testing.T) { mapping := ResponsesClientToolMapping{ CustomTools: map[string]bool{"exec": true}, ToolSearch: true, diff --git a/backend/internal/pkg/apicompat/responses_tool_search_discoveries.go b/backend/internal/pkg/apicompat/responses_tool_search_discoveries.go new file mode 100644 index 0000000000..2617d5f940 --- /dev/null +++ b/backend/internal/pkg/apicompat/responses_tool_search_discoveries.go @@ -0,0 +1,263 @@ +package apicompat + +import ( + "encoding/json" + "fmt" + "strings" +) + +type responsesDiscoveredToolIdentity struct { + typ string + name string + namespace string + encoded string + ambiguous bool +} + +// promoteResponsesToolSearchDiscoveries makes successfully discovered client +// tools callable for function-only upstreams. The original tool_search_output +// remains request history and is normalized separately; declarations are +// appended after static tools so the client's declaration order stays stable. +func promoteResponsesToolSearchDiscoveries(req map[string]any) (bool, error) { + tools, ok := req["tools"].([]any) + if !ok || len(tools) == 0 || !hasResponsesToolSearchDeclaration(tools) { + return false, nil + } + input, ok := req["input"].([]any) + if !ok || len(input) == 0 { + return false, nil + } + + known := make(map[string]responsesDiscoveredToolIdentity) + for _, raw := range tools { + registerExistingResponsesToolIdentity(known, raw) + } + + promoted := make([]any, 0) + for _, rawItem := range input { + item, ok := rawItem.(map[string]any) + if !ok || strings.TrimSpace(stringValue(item["type"])) != "tool_search_output" || !usableResponsesToolSearchOutput(item) { + continue + } + discoveries, ok := item["tools"].([]any) + if !ok || len(discoveries) == 0 { + continue + } + if _, err := json.Marshal(discoveries); err != nil { + continue + } + for _, rawDiscovery := range discoveries { + discovery, ok := rawDiscovery.(map[string]any) + if !ok { + continue + } + typ := strings.TrimSpace(stringValue(discovery["type"])) + switch typ { + case "function", "custom": + copy, identity, ok := responsesDirectToolDiscovery(discovery, typ) + if !ok { + continue + } + appendTool, err := admitResponsesDiscoveredTool(known, identity.name, identity) + if err != nil { + return false, err + } + if appendTool { + promoted = append(promoted, copy) + } + case "namespace": + copy, identities, ok := responsesNamespaceToolDiscovery(discovery) + if !ok { + continue + } + children := make([]any, 0, len(identities)) + for _, candidate := range identities { + appendTool, err := admitResponsesDiscoveredTool(known, candidate.flat, candidate.identity) + if err != nil { + return false, err + } + if appendTool { + children = append(children, candidate.child) + } + } + if len(children) > 0 { + copy["tools"] = children + delete(copy, "children") + promoted = append(promoted, copy) + } + } + } + } + if len(promoted) == 0 { + return false, nil + } + req["tools"] = append(tools, promoted...) + return true, nil +} + +func hasResponsesToolSearchDeclaration(tools []any) bool { + for _, raw := range tools { + tool, ok := raw.(map[string]any) + if ok && strings.TrimSpace(stringValue(tool["type"])) == "tool_search" { + return true + } + } + return false +} + +func usableResponsesToolSearchOutput(item map[string]any) bool { + status, present := item["status"] + if !present { + return true + } + text, ok := status.(string) + return ok && strings.TrimSpace(text) == "completed" +} + +func registerExistingResponsesToolIdentity(known map[string]responsesDiscoveredToolIdentity, raw any) { + tool, ok := raw.(map[string]any) + if !ok { + return + } + typ := strings.TrimSpace(stringValue(tool["type"])) + switch typ { + case "function", "custom": + _, identity, ok := responsesDirectToolDiscovery(tool, typ) + if ok { + registerExistingResponsesIdentity(known, identity.name, identity) + } + case "namespace": + _, identities, ok := responsesNamespaceToolDiscovery(tool) + if !ok { + return + } + for _, candidate := range identities { + registerExistingResponsesIdentity(known, candidate.flat, candidate.identity) + } + } +} + +func registerExistingResponsesIdentity(known map[string]responsesDiscoveredToolIdentity, key string, identity responsesDiscoveredToolIdentity) { + previous, exists := known[key] + if !exists { + known[key] = identity + return + } + if !sameResponsesDiscoveredTool(previous, identity) { + previous.ambiguous = true + known[key] = previous + } +} + +func admitResponsesDiscoveredTool(known map[string]responsesDiscoveredToolIdentity, key string, identity responsesDiscoveredToolIdentity) (bool, error) { + previous, exists := known[key] + if !exists { + known[key] = identity + return true, nil + } + if !previous.ambiguous && sameResponsesDiscoveredTool(previous, identity) { + return false, nil + } + return false, fmt.Errorf("discovered tool %q conflicts with an existing declaration; this upstream cannot safely disambiguate different names, namespaces, or schemas", key) +} + +func sameResponsesDiscoveredTool(left, right responsesDiscoveredToolIdentity) bool { + return left.typ == right.typ && left.name == right.name && left.namespace == right.namespace && left.encoded == right.encoded +} + +func responsesDirectToolDiscovery(tool map[string]any, typ string) (map[string]any, responsesDiscoveredToolIdentity, bool) { + name := strings.TrimSpace(stringValue(tool["name"])) + if name == "" { + return nil, responsesDiscoveredToolIdentity{}, false + } + copy := copyClientTool(tool) + copy["type"] = typ + copy["name"] = name + identityCopy := copy + if typ == "custom" { + identityCopy = copyClientTool(copy) + identityCopy["type"] = "function" + identityCopy["parameters"] = json.RawMessage(customToolInputSchema) + delete(identityCopy, "format") + } + encoded, err := json.Marshal(identityCopy) + if err != nil { + return nil, responsesDiscoveredToolIdentity{}, false + } + return copy, responsesDiscoveredToolIdentity{typ: typ, name: name, encoded: string(encoded)}, true +} + +// restoreInheritedResponsesClientToolDeclarations reverses only the declaration +// identities recorded by ResponsesClientToolMapping. It is used when a WS +// continuation omits tools but the HTTP function upstream still needs the +// effective session declarations on every request. +func restoreInheritedResponsesClientToolDeclarations(lowered []any, mapping ResponsesClientToolMapping) []any { + restored := make([]any, 0, len(lowered)) + for _, raw := range lowered { + tool, ok := raw.(map[string]any) + if !ok { + restored = append(restored, raw) + continue + } + name := strings.TrimSpace(stringValue(tool["name"])) + switch { + case mapping.ToolSearch && name == toolSearchProxyName: + restored = append(restored, map[string]any{"type": "tool_search"}) + case mapping.CustomTools[name]: + copy := copyClientTool(tool) + copy["type"] = "custom" + restored = append(restored, copy) + case mapping.NamespaceTools[name].Namespace != "": + identity := mapping.NamespaceTools[name] + child := copyClientTool(tool) + child["type"] = "function" + child["name"] = identity.Name + restored = append(restored, map[string]any{ + "type": "namespace", "name": identity.Namespace, "tools": []any{child}, + }) + default: + restored = append(restored, copyClientTool(tool)) + } + } + return restored +} + +type responsesNamespaceToolCandidate struct { + flat string + child map[string]any + identity responsesDiscoveredToolIdentity +} + +func responsesNamespaceToolDiscovery(tool map[string]any) (map[string]any, []responsesNamespaceToolCandidate, bool) { + namespace := strings.TrimSpace(stringValue(tool["name"])) + children := namespaceChildren(tool) + if namespace == "" || len(children) == 0 { + return nil, nil, false + } + copy := copyClientTool(tool) + copy["type"] = "namespace" + copy["name"] = namespace + identities := make([]responsesNamespaceToolCandidate, 0, len(children)) + for _, rawChild := range children { + child, ok := rawChild.(map[string]any) + if !ok || strings.TrimSpace(stringValue(child["type"])) != "function" { + continue + } + childCopy, direct, ok := responsesDirectToolDiscovery(child, "function") + if !ok { + continue + } + flat := flattenNamespaceToolName(namespace, direct.name) + identities = append(identities, responsesNamespaceToolCandidate{ + flat: flat, + child: childCopy, + identity: responsesDiscoveredToolIdentity{ + typ: "namespace", name: direct.name, namespace: namespace, encoded: direct.encoded, + }, + }) + } + if len(identities) == 0 { + return nil, nil, false + } + return copy, identities, true +} diff --git a/backend/internal/service/openai_gateway_grok_tool_protocol.go b/backend/internal/service/openai_gateway_grok_tool_protocol.go index c22282f8ca..9f876a1a6e 100644 --- a/backend/internal/service/openai_gateway_grok_tool_protocol.go +++ b/backend/internal/service/openai_gateway_grok_tool_protocol.go @@ -27,6 +27,7 @@ func adaptResponsesClientToolsForFunctionUpstreamWithMapping( body []byte, upstream string, inherited apicompat.ResponsesClientToolMapping, + inheritedLoweredTools ...[]any, ) ([]byte, apicompat.ResponsesClientToolMapping, error) { decoder := json.NewDecoder(bytes.NewReader(body)) decoder.UseNumber() @@ -35,7 +36,7 @@ func adaptResponsesClientToolsForFunctionUpstreamWithMapping( return body, apicompat.ResponsesClientToolMapping{}, fmt.Errorf("decode %s Responses client tools: %w", upstream, err) } - mapping, changed, err := apicompat.AdaptResponsesClientToolsWithInheritedMapping(requestBody, inherited) + mapping, changed, err := apicompat.AdaptResponsesClientToolsWithInheritedMapping(requestBody, inherited, inheritedLoweredTools...) if err != nil { return body, apicompat.ResponsesClientToolMapping{}, err } diff --git a/backend/internal/service/openai_gateway_grok_tool_protocol_test.go b/backend/internal/service/openai_gateway_grok_tool_protocol_test.go index 50efdbfce8..cc25cb1924 100644 --- a/backend/internal/service/openai_gateway_grok_tool_protocol_test.go +++ b/backend/internal/service/openai_gateway_grok_tool_protocol_test.go @@ -72,10 +72,13 @@ func TestPatchGrokResponsesBodyWithClientToolsLowersDiscoveredToolsOutput(t *tes "model":"grok-4.5", "tools":[{"type":"tool_search"}], "input":[ - {"type":"tool_search_call","id":"tsc_request_74547","call_id":"call_request_74547","arguments":{"query":"subagent"},"execution":"client","status":"completed"}, - {"type":"tool_search_output","id":"tso_request_74547","call_id":"call_request_74547","execution":"client","status":"completed","tools":[ + {"type":"tool_search_call","id":"tsc_fixture","call_id":"call_fixture","arguments":{"query":"subagent"},"execution":"client","status":"completed"}, + {"type":"tool_search_output","id":"tso_fixture","call_id":"call_fixture","execution":"client","status":"completed","tools":[ {"type":"namespace","name":"codex_app","tools":[{"type":"function","name":"load_workspace_dependencies","parameters":{"type":"object","properties":{},"additionalProperties":false}}]}, - {"type":"namespace","name":"multi_agent_v1","tools":[{"type":"function","name":"spawn_agent","parameters":{"type":"object","properties":{"message":{"type":"string"}},"required":["message"],"additionalProperties":false}}]} + {"type":"namespace","name":"multi_agent_v1","tools":[ + {"type":"function","name":"spawn_agent","parameters":{"type":"object","properties":{"message":{"type":"string"}},"required":["message"],"additionalProperties":false}}, + {"type":"function","name":"wait_agent","parameters":{"type":"object","properties":{"timeout_ms":{"type":"integer"}},"additionalProperties":false}} + ]} ]} ] }`) @@ -83,19 +86,37 @@ func TestPatchGrokResponsesBodyWithClientToolsLowersDiscoveredToolsOutput(t *tes patched, mapping, err := patchGrokResponsesBodyWithClientTools(body, "grok-4.5") require.NoError(t, err) require.True(t, mapping.ToolSearch) - require.Equal(t, "function_call", gjson.GetBytes(patched, "input.0.type").String()) - require.Equal(t, "tool_search", gjson.GetBytes(patched, "input.0.name").String()) - require.Equal(t, "function_call_output", gjson.GetBytes(patched, "input.1.type").String()) - require.Equal(t, "call_request_74547", gjson.GetBytes(patched, "input.1.call_id").String()) - require.Len(t, gjson.GetBytes(patched, "input.1").Map(), 3) - require.False(t, gjson.GetBytes(patched, "input.1.tools").Exists()) - require.False(t, gjson.GetBytes(patched, "input.1.status").Exists()) - require.False(t, gjson.GetBytes(patched, "input.1.execution").Exists()) + require.Equal(t, apicompat.ResponsesNamespaceName{Namespace: "multi_agent_v1", Name: "spawn_agent"}, mapping.NamespaceTools["multi_agent_v1__spawn_agent"]) + require.Equal(t, apicompat.ResponsesNamespaceName{Namespace: "multi_agent_v1", Name: "wait_agent"}, mapping.NamespaceTools["multi_agent_v1__wait_agent"]) output := gjson.GetBytes(patched, "input.1.output").String() require.JSONEq(t, `[ {"type":"namespace","name":"codex_app","tools":[{"type":"function","name":"load_workspace_dependencies","parameters":{"type":"object","properties":{},"additionalProperties":false}}]}, - {"type":"namespace","name":"multi_agent_v1","tools":[{"type":"function","name":"spawn_agent","parameters":{"type":"object","properties":{"message":{"type":"string"}},"required":["message"],"additionalProperties":false}}]} + {"type":"namespace","name":"multi_agent_v1","tools":[ + {"type":"function","name":"spawn_agent","parameters":{"type":"object","properties":{"message":{"type":"string"}},"required":["message"],"additionalProperties":false}}, + {"type":"function","name":"wait_agent","parameters":{"type":"object","properties":{"timeout_ms":{"type":"integer"}},"additionalProperties":false}} + ]} ]`, output) + + require.JSONEq(t, `{ + "model":"grok-4.5", + "tools":[ + {"type":"function","name":"tool_search","description":"Search and load Codex tools, plugins, connectors, and MCP namespaces for the current task.","parameters":{"type":"object","properties":{"query":{"type":"string","description":"Search query for tools or connectors to load."},"limit":{"type":"integer","description":"Maximum number of tool groups to return."}},"required":["query"]}}, + {"type":"function","name":"codex_app__load_workspace_dependencies","parameters":{"type":"object","properties":{},"additionalProperties":false}}, + {"type":"function","name":"multi_agent_v1__spawn_agent","parameters":{"type":"object","properties":{"message":{"type":"string"}},"required":["message"],"additionalProperties":false}}, + {"type":"function","name":"multi_agent_v1__wait_agent","parameters":{"type":"object","properties":{"timeout_ms":{"type":"integer"}},"additionalProperties":false}} + ], + "input":[ + {"type":"function_call","call_id":"call_fixture","name":"tool_search","arguments":"{\"query\":\"subagent\"}","status":"completed"}, + {"type":"function_call_output","call_id":"call_fixture","output":`+string(mustMarshalJSONForTest(t, output))+`} + ] + }`, string(patched)) +} + +func mustMarshalJSONForTest(t *testing.T, value string) []byte { + t.Helper() + encoded, err := json.Marshal(value) + require.NoError(t, err) + return encoded } func TestPatchGrokResponsesBodyWithClientToolsRewritesEveryToolChoice(t *testing.T) { @@ -207,6 +228,32 @@ func TestForwardGrokResponsesClientToolNameConflictReturns400(t *testing.T) { require.Empty(t, upstream.requests, "an ambiguous request must not reach xAI") } +func TestForwardGrokResponsesMalformedToolSearchOutputReturns400BeforeUpstream(t *testing.T) { + gin.SetMode(gin.TestMode) + + body := []byte(`{ + "model":"grok","stream":false, + "tools":[{"type":"tool_search"}], + "input":[{"type":"tool_search_output","status":"completed"}] + }`) + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", bytes.NewReader(body)) + upstream := &httpUpstreamRecorder{} + svc := &OpenAIGatewayService{httpUpstream: upstream} + account := grokProtocolAPIKeyAccount(7103) + + result, err := svc.forwardGrokResponses(context.Background(), c, account, body, "grok", false, time.Now()) + + require.Error(t, err) + require.Nil(t, result) + require.Equal(t, http.StatusBadRequest, recorder.Code) + require.Equal(t, "invalid_request_error", gjson.Get(recorder.Body.String(), "error.type").String()) + require.Equal(t, "tools", gjson.Get(recorder.Body.String(), "error.param").String()) + require.Contains(t, gjson.Get(recorder.Body.String(), "error.message").String(), "call_id") + require.Empty(t, upstream.requests, "malformed lowered output must not reach xAI") +} + func TestForwardGrokResponsesOAuthRestoresClientToolsNonStreaming(t *testing.T) { gin.SetMode(gin.TestMode) diff --git a/backend/internal/service/openai_ws_http_bridge.go b/backend/internal/service/openai_ws_http_bridge.go index 1fab9be1a5..d68e26c3d6 100644 --- a/backend/internal/service/openai_ws_http_bridge.go +++ b/backend/internal/service/openai_ws_http_bridge.go @@ -15,7 +15,6 @@ import ( "github.com/Wei-Shaw/sub2api/internal/pkg/apicompat" "github.com/gin-gonic/gin" "github.com/tidwall/gjson" - "github.com/tidwall/sjson" ) const ( @@ -48,6 +47,33 @@ func setOpenAIWSHTTPBridgeToolState(c *gin.Context, state openAIWSHTTPBridgeTool c.Set(openAIWSHTTPBridgeToolStateContextKey, state) } +func decodeOpenAIWSHTTPBridgeLoweredTools(raw json.RawMessage) []any { + if len(raw) == 0 { + return nil + } + var tools []any + if err := json.Unmarshal(raw, &tools); err != nil { + return nil + } + return tools +} + +func openAIWSHTTPBridgeRawField(body []byte, name string) (json.RawMessage, bool) { + var fields map[string]json.RawMessage + if err := json.Unmarshal(body, &fields); err != nil { + return nil, false + } + raw, present := fields[name] + return append(json.RawMessage(nil), raw...), present +} + +func openAIWSHTTPBridgeToolUpstreamName(account *Account) string { + if account != nil && account.Platform == PlatformGrok { + return "Grok WS HTTP bridge" + } + return "OpenAI WS HTTP bridge" +} + // ResolveOpenAIWSClientFirstMessageTimeout returns the effective client ingress deadline. func ResolveOpenAIWSClientFirstMessageTimeout(cfg *config.Config) time.Duration { seconds := config.DefaultOpenAIWSClientFirstMessageTimeoutSeconds @@ -246,26 +272,40 @@ func (s *OpenAIGatewayService) proxyOpenAIWSHTTPBridgeTurn( if err != nil { return nil, fmt.Errorf("prepare http bridge body: %w", err) } + grokIntentSourceBody := append([]byte(nil), body...) + _, grokExplicitToolsField := openAIWSHTTPBridgeRawField(grokIntentSourceBody, "tools") + grokExplicitToolIntent := account.Platform == PlatformGrok && hasGrokResponsesToolIntent(grokIntentSourceBody) var clientToolMapping apicompat.ResponsesClientToolMapping - if account.Platform == PlatformOpenAI && account.Type == AccountTypeAPIKey { + functionToolUpstream := (account.Platform == PlatformOpenAI && account.Type == AccountTypeAPIKey) || account.Platform == PlatformGrok + if functionToolUpstream { + if account.Platform == PlatformGrok { + body, err = sanitizeGrokResponsesInput(body) + if err != nil { + return nil, fmt.Errorf("sanitize Grok WS HTTP bridge input: %w", err) + } + } inheritedState, _ := openAIWSHTTPBridgeToolStateFromContext(c) - toolsPresent := gjson.GetBytes(body, "tools").Exists() + inheritedLoweredTools := decodeOpenAIWSHTTPBridgeLoweredTools(inheritedState.LoweredTools) body, clientToolMapping, err = adaptResponsesClientToolsForFunctionUpstreamWithMapping( body, - "OpenAI WS HTTP bridge", + openAIWSHTTPBridgeToolUpstreamName(account), inheritedState.ClientMapping, + inheritedLoweredTools, ) if err != nil { - return nil, fmt.Errorf("adapt OpenAI WS HTTP bridge client tools: %w", err) + return nil, fmt.Errorf("adapt %s client tools: %w", openAIWSHTTPBridgeToolUpstreamName(account), err) + } + if account.Platform == PlatformGrok && !grokExplicitToolsField && !grokExplicitToolIntent && len(inheritedLoweredTools) > 0 && hasGrokResponsesToolIntent(body) { + // This continuation omitted tools, so the pre-adapter source cannot + // represent the effective inherited declarations. Cache routing must + // see the rehydrated tool intent or it will replace client functions + // with the native-search tool-free route. Explicit current-turn tool + // intent still uses the original pre-sanitization source above. + grokIntentSourceBody = append(grokIntentSourceBody[:0], body...) } loweredTools := inheritedState.LoweredTools - if toolsPresent { - loweredTools = json.RawMessage(gjson.GetBytes(body, "tools").Raw) - } else if len(loweredTools) > 0 { - body, err = sjson.SetRawBytes(body, "tools", loweredTools) - if err != nil { - return nil, fmt.Errorf("inherit OpenAI WS HTTP bridge tools: %w", err) - } + if currentTools, present := openAIWSHTTPBridgeRawField(body, "tools"); present { + loweredTools = currentTools } setOpenAIWSHTTPBridgeToolState(c, openAIWSHTTPBridgeToolState{ ClientMapping: clientToolMapping, @@ -277,7 +317,6 @@ func (s *OpenAIGatewayService) proxyOpenAIWSHTTPBridgeTurn( var upstreamReq *http.Request if account.Platform == PlatformGrok { upstreamModel := resolveGrokWSUpstreamModel(account, body, originalModel) - grokIntentSourceBody := body body, err = patchGrokResponsesBody(body, upstreamModel) if err != nil { releaseUpstreamCtx() diff --git a/backend/internal/service/openai_ws_http_bridge_test.go b/backend/internal/service/openai_ws_http_bridge_test.go index 9617bd464d..49b4c70735 100644 --- a/backend/internal/service/openai_ws_http_bridge_test.go +++ b/backend/internal/service/openai_ws_http_bridge_test.go @@ -170,6 +170,131 @@ func TestProxyOpenAIWSHTTPBridgeTurnAPIKeyRestoresClientToolsInResponseDone(t *t require.Equal(t, "custom_tool_call", gjson.GetBytes(result.wsReplayInput[0], "type").String()) } +func TestProxyOpenAIWSHTTPBridgeTurnGrokPromotesDiscoveryAndRestoresNamespaceSSE(t *testing.T) { + gin.SetMode(gin.TestMode) + + sse := strings.Join([]string{ + `data: {"type":"response.output_item.added","sequence_number":0,"output_index":0,"item":{"type":"function_call","id":"item_spawn","call_id":"call_spawn","name":"multi_agent_v1__spawn_agent","status":"in_progress"}}`, + "", + `data: {"type":"response.function_call_arguments.done","sequence_number":1,"output_index":0,"item_id":"item_spawn","call_id":"call_spawn","name":"multi_agent_v1__spawn_agent","arguments":"{\"message\":\"work\"}"}`, + "", + `data: {"type":"response.output_item.done","sequence_number":2,"output_index":0,"item":{"type":"function_call","id":"item_spawn","call_id":"call_spawn","name":"multi_agent_v1__spawn_agent","arguments":"{\"message\":\"work\"}","status":"completed"}}`, + "", + `data: {"type":"response.completed","sequence_number":3,"response":{"id":"resp_spawn","status":"completed","output":[{"type":"function_call","id":"item_spawn","call_id":"call_spawn","name":"multi_agent_v1__spawn_agent","arguments":"{\"message\":\"work\"}","status":"completed"}],"usage":{"input_tokens":1,"output_tokens":1}}}`, + "", + }, "\n") + upstream := &httpUpstreamRecorder{resp: &http.Response{ + StatusCode: http.StatusOK, + Header: http.Header{"Content-Type": []string{"text/event-stream"}}, + Body: io.NopCloser(strings.NewReader(sse)), + }} + svc := &OpenAIGatewayService{ + cfg: &config.Config{Gateway: config.GatewayConfig{MaxLineSize: defaultMaxLineSize}}, + httpUpstream: upstream, + } + account := &Account{ + ID: 5765, Platform: PlatformGrok, Type: AccountTypeOAuth, Concurrency: 1, + Credentials: map[string]any{"base_url": xai.DefaultCLIBaseURL}, + } + payload := []byte(`{ + "type":"response.create","model":"grok-4.5","stream":true, + "tools":[{"type":"tool_search"}], + "input":[ + {"type":"tool_search_call","call_id":"call_search","arguments":{"query":"subagent"},"status":"completed"}, + {"type":"tool_search_output","call_id":"call_search","execution":"client","status":"completed","tools":[ + {"type":"namespace","name":"multi_agent_v1","tools":[{"type":"function","name":"spawn_agent","parameters":{"type":"object"}}]} + ]} + ] + }`) + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + c.Request = httptest.NewRequest(http.MethodGet, "/v1/responses", nil) + var events [][]byte + + result, err := svc.proxyOpenAIWSHTTPBridgeTurn( + context.Background(), c, account, "access-token", payload, len(payload), + "grok-4.5", "", "", "", "", 1, + func(message []byte) error { + events = append(events, append([]byte(nil), message...)) + return nil + }, + ) + + require.NoError(t, err) + require.NotNil(t, result) + require.Equal(t, "multi_agent_v1__spawn_agent", gjson.GetBytes(upstream.lastBody, "tools.1.name").String()) + require.Equal(t, "function_call_output", gjson.GetBytes(upstream.lastBody, "input.1.type").String()) + require.False(t, gjson.GetBytes(upstream.lastBody, "input.1.tools").Exists()) + require.False(t, gjson.GetBytes(upstream.lastBody, "input.1.status").Exists()) + require.False(t, gjson.GetBytes(upstream.lastBody, "input.1.execution").Exists()) + state, ok := openAIWSHTTPBridgeToolStateFromContext(c) + require.True(t, ok) + require.Equal(t, "multi_agent_v1", state.ClientMapping.NamespaceTools["multi_agent_v1__spawn_agent"].Namespace) + require.Len(t, events, 4) + require.Equal(t, "spawn_agent", gjson.GetBytes(events[0], "item.name").String()) + require.Equal(t, "multi_agent_v1", gjson.GetBytes(events[0], "item.namespace").String()) + require.Equal(t, "spawn_agent", gjson.GetBytes(events[2], "item.name").String()) + require.Equal(t, "multi_agent_v1", gjson.GetBytes(events[2], "item.namespace").String()) + require.Equal(t, "spawn_agent", gjson.GetBytes(events[3], "response.output.0.name").String()) + require.Equal(t, "multi_agent_v1", gjson.GetBytes(events[3], "response.output.0.namespace").String()) +} + +func TestProxyOpenAIWSHTTPBridgeTurnGrokInheritsToolSearchAndPromotesFollowupDiscovery(t *testing.T) { + gin.SetMode(gin.TestMode) + + firstSSE := "data: {\"type\":\"response.completed\",\"response\":{\"id\":\"resp_first\",\"output\":[],\"usage\":{\"input_tokens\":1,\"output_tokens\":1}}}\n\n" + secondSSE := "data: {\"type\":\"response.completed\",\"response\":{\"id\":\"resp_second\",\"output\":[{\"type\":\"function_call\",\"id\":\"item_spawn\",\"call_id\":\"call_spawn\",\"name\":\"multi_agent_v1__spawn_agent\",\"arguments\":\"{}\"}],\"usage\":{\"input_tokens\":1,\"output_tokens\":1}}}\n\n" + upstream := &httpUpstreamRecorder{responses: []*http.Response{ + {StatusCode: http.StatusOK, Header: http.Header{"Content-Type": []string{"text/event-stream"}}, Body: io.NopCloser(strings.NewReader(firstSSE))}, + {StatusCode: http.StatusOK, Header: http.Header{"Content-Type": []string{"text/event-stream"}}, Body: io.NopCloser(strings.NewReader(secondSSE))}, + }} + svc := &OpenAIGatewayService{ + cfg: &config.Config{Gateway: config.GatewayConfig{MaxLineSize: defaultMaxLineSize}}, + httpUpstream: upstream, + } + account := &Account{ + ID: 5766, Platform: PlatformGrok, Type: AccountTypeOAuth, Concurrency: 1, + Credentials: map[string]any{"base_url": xai.DefaultCLIBaseURL, "subscription_tier": "free"}, + } + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + c.Request = httptest.NewRequest(http.MethodGet, "/v1/responses", nil) + + first := []byte(`{"type":"response.create","model":"grok-4.5","stream":true,"tools":[{"type":"tool_search"}],"input":"discover tools"}`) + _, err := svc.proxyOpenAIWSHTTPBridgeTurn( + context.Background(), c, account, "access-token", first, len(first), + "grok-4.5", "", "", "", "grok-ws-cache", 1, func([]byte) error { return nil }, + ) + require.NoError(t, err) + + second := []byte(`{ + "type":"response.create","model":"grok-4.5","stream":true, + "input":[{"type":"tool_search_output","call_id":"call_search","status":"completed","tools":[ + {"type":"namespace","name":"multi_agent_v1","tools":[{"type":"function","name":"spawn_agent","parameters":{"type":"object"}}]} + ]}] + }`) + var events [][]byte + _, err = svc.proxyOpenAIWSHTTPBridgeTurn( + context.Background(), c, account, "access-token", second, len(second), + "grok-4.5", "", "", "", "grok-ws-cache", 2, + func(message []byte) error { + events = append(events, append([]byte(nil), message...)) + return nil + }, + ) + + require.NoError(t, err) + require.Len(t, upstream.bodies, 2) + require.Equal(t, "tool_search", gjson.GetBytes(upstream.bodies[1], "tools.0.name").String()) + require.Equal(t, "multi_agent_v1__spawn_agent", gjson.GetBytes(upstream.bodies[1], "tools.1.name").String()) + require.NotEqual(t, grokFreeCacheDisabledToolChoice, gjson.GetBytes(upstream.bodies[1], "tool_choice").String()) + require.Equal(t, "grok-ws-cache", gjson.GetBytes(upstream.bodies[1], "prompt_cache_key").String()) + require.Equal(t, "function_call_output", gjson.GetBytes(upstream.bodies[1], "input.0.type").String()) + require.Len(t, events, 1) + require.Equal(t, "spawn_agent", gjson.GetBytes(events[0], "response.output.0.name").String()) + require.Equal(t, "multi_agent_v1", gjson.GetBytes(events[0], "response.output.0.namespace").String()) +} + func TestOpenAIWSHTTPBridgeAPIKeyReusesClientToolMappingWhenFollowupOmitsTools(t *testing.T) { gin.SetMode(gin.TestMode) @@ -894,11 +1019,13 @@ func TestProxyOpenAIWSHTTPBridgeTurnPromotesCodexAdditionalToolsForMixedCache(t require.Len(t, events, 2) require.False(t, gjson.GetBytes(upstream.lastBody, `input.#(type=="additional_tools")`).Exists()) tools := gjson.GetBytes(upstream.lastBody, "tools").Array() - require.Len(t, tools, 3) + require.Len(t, tools, 4) require.Equal(t, "function", tools[0].Get("type").String()) require.Equal(t, "lookup", tools[0].Get("name").String()) require.Equal(t, "web_search", tools[1].Get("type").String()) - require.Equal(t, "x_search", tools[2].Get("type").String()) + require.Equal(t, "function", tools[2].Get("type").String()) + require.Equal(t, "apply_patch", tools[2].Get("name").String()) + require.Equal(t, "x_search", tools[3].Get("type").String()) require.False(t, gjson.GetBytes(upstream.lastBody, "tool_choice").Exists()) require.False(t, gjson.GetBytes(upstream.lastBody, `tools.#(type=="custom")`).Exists()) require.False(t, gjson.GetBytes(upstream.lastBody, `tools.#(type=="namespace")`).Exists())