diff --git a/backend/internal/pkg/apicompat/responses_client_tools.go b/backend/internal/pkg/apicompat/responses_client_tools.go new file mode 100644 index 000000000..947d3e8ff --- /dev/null +++ b/backend/internal/pkg/apicompat/responses_client_tools.go @@ -0,0 +1,617 @@ +package apicompat + +import ( + "bytes" + "encoding/json" + "fmt" + "strings" +) + +// ResponsesClientToolMapping records the reversible lowering applied before a +// native Responses request is sent to an upstream that only understands +// function tools. +type ResponsesClientToolMapping struct { + CustomTools map[string]bool + ToolSearch bool + NamespaceTools map[string]ResponsesNamespaceName +} + +// AdaptResponsesClientTools lowers Codex client-only tools in req to +// ordinary function tools. It mutates req and returns the mapping required to +// restore the upstream response. +func AdaptResponsesClientTools(req map[string]any) (ResponsesClientToolMapping, bool, error) { + if req == nil { + return ResponsesClientToolMapping{}, false, nil + } + tools, ok := req["tools"].([]any) + if !ok || len(tools) == 0 { + return ResponsesClientToolMapping{}, false, nil + } + + adapter := ResponsesClientToolMapping{CustomTools: make(map[string]bool)} + functionNames := make(map[string]bool) + customNames := make(map[string]bool) + for _, raw := range tools { + tool, ok := raw.(map[string]any) + if !ok { + continue + } + name := strings.TrimSpace(stringValue(tool["name"])) + switch strings.TrimSpace(stringValue(tool["type"])) { + case "function": + if name != "" { + functionNames[name] = true + } + case "custom": + if name != "" { + customNames[name] = true + } + case "tool_search": + adapter.ToolSearch = true + } + } + for name := range customNames { + if functionNames[name] { + return ResponsesClientToolMapping{}, false, fmt.Errorf("custom tool %q conflicts with a function tool of the same name; this upstream cannot disambiguate them, rename one of the tools", name) + } + } + if adapter.ToolSearch && (functionNames[toolSearchProxyName] || customNames[toolSearchProxyName]) { + return ResponsesClientToolMapping{}, false, fmt.Errorf("built-in tool_search conflicts with a declared tool named %q; this upstream cannot disambiguate them, rename the tool", toolSearchProxyName) + } + + // Namespace flattening also rewrites namespace-qualified history and choice. + names, flattened, err := FlattenResponsesNamespaces(req) + if err != nil { + return ResponsesClientToolMapping{}, false, err + } + adapter.NamespaceTools = names + if adapter.ToolSearch { + if _, exists := names[toolSearchProxyName]; exists { + return ResponsesClientToolMapping{}, false, fmt.Errorf("built-in tool_search conflicts with namespace tool flattened as %q; this upstream cannot disambiguate them, rename the tool", toolSearchProxyName) + } + } + + tools, _ = req["tools"].([]any) + lowered := make([]any, 0, len(tools)) + changed := flattened + seenSearch := false + for _, raw := range tools { + tool, ok := raw.(map[string]any) + if !ok { + lowered = append(lowered, raw) + continue + } + typ := strings.TrimSpace(stringValue(tool["type"])) + name := strings.TrimSpace(stringValue(tool["name"])) + switch typ { + case "custom": + if name == "" { + lowered = append(lowered, raw) + continue + } + copy := copyClientTool(tool) + copy["type"] = "function" + copy["parameters"] = json.RawMessage(customToolInputSchema) + delete(copy, "format") + adapter.CustomTools[name] = true + lowered = append(lowered, copy) + changed = true + case "tool_search": + if seenSearch { + changed = true + continue + } + seenSearch = true + lowered = append(lowered, map[string]any{ + "type": "function", "name": toolSearchProxyName, + "description": "Search and load Codex tools, plugins, connectors, and MCP namespaces for the current task.", + "parameters": json.RawMessage(toolSearchProxySchema), + }) + changed = true + default: + lowered = append(lowered, raw) + } + } + if changed { + req["tools"] = lowered + } + if rewriteClientToolHistory(req["input"], &adapter) { + changed = true + } + if rewriteClientToolChoice(req, &adapter) { + changed = true + } + if len(adapter.CustomTools) == 0 { + adapter.CustomTools = nil + } + if len(adapter.NamespaceTools) == 0 { + adapter.NamespaceTools = nil + } + return adapter, changed, nil +} + +func copyClientTool(tool map[string]any) map[string]any { + copy := make(map[string]any, len(tool)) + for key, value := range tool { + copy[key] = value + } + return copy +} + +func rewriteClientToolHistory(value any, adapter *ResponsesClientToolMapping) bool { + changed := false + var visit func(any) + visit = func(value any) { + switch typed := value.(type) { + case []any: + for _, item := range typed { + visit(item) + } + case map[string]any: + typ := strings.TrimSpace(stringValue(typed["type"])) + switch typ { + case "custom_tool_call": + if adapter.CustomTools[strings.TrimSpace(stringValue(typed["name"]))] { + typed["type"] = "function_call" + typed["arguments"] = customToolCallArguments(stringValue(typed["input"])) + delete(typed, "input") + changed = true + } + case "custom_tool_call_output": + typed["type"] = "function_call_output" + normalizeClientToolOutput(typed) + changed = true + case "tool_search_call": + if adapter.ToolSearch { + typed["type"] = "function_call" + typed["name"] = toolSearchProxyName + typed["arguments"] = rawObjectString(typed["arguments"]) + delete(typed, "execution") + changed = true + } + case "tool_search_output": + if adapter.ToolSearch { + typed["type"] = "function_call_output" + normalizeClientToolOutput(typed) + changed = true + } + } + for _, child := range typed { + visit(child) + } + } + } + visit(value) + return changed +} + +func normalizeClientToolOutput(item map[string]any) { + output, exists := item["output"] + if !exists { + return + } + if _, ok := output.(string); ok { + return + } + if output == nil { + item["output"] = "" + return + } + encoded, err := json.Marshal(output) + if err != nil { + item["output"] = "" + return + } + item["output"] = string(encoded) +} + +func rewriteClientToolChoice(req map[string]any, adapter *ResponsesClientToolMapping) bool { + choice, ok := req["tool_choice"].(map[string]any) + if !ok { + return false + } + typ := strings.TrimSpace(stringValue(choice["type"])) + name := strings.TrimSpace(stringValue(choice["name"])) + if typ == "custom" && adapter.CustomTools[name] { + choice["type"] = "function" + return true + } + if typ == "tool_search" && adapter.ToolSearch { + req["tool_choice"] = map[string]any{"type": "function", "name": toolSearchProxyName} + return true + } + return false +} + +func customToolCallArguments(input string) string { + encoded, _ := json.Marshal(map[string]string{"input": input}) + return string(encoded) +} + +func rawObjectString(value any) string { + if text, ok := value.(string); ok { + return text + } + encoded, err := json.Marshal(value) + if err != nil { + return "{}" + } + return string(encoded) +} + +// RestoreResponsesClientToolPayload restores client tool calls in a non-stream +// native Responses JSON payload. +func RestoreResponsesClientToolPayload(payload []byte, mapping ResponsesClientToolMapping) ([]byte, bool, error) { + if len(payload) == 0 { + return payload, false, nil + } + var value any + if err := json.Unmarshal(payload, &value); err != nil { + return payload, false, err + } + changed := restoreClientToolValue(value, &mapping) + if !changed { + if len(mapping.NamespaceTools) == 0 { + return payload, false, nil + } + return RestoreResponsesNamespaceCalls(payload, mapping.NamespaceTools) + } + var rebuilt bytes.Buffer + encoder := json.NewEncoder(&rebuilt) + encoder.SetEscapeHTML(false) + if err := encoder.Encode(value); err != nil { + return payload, false, err + } + rebuiltPayload := bytes.TrimSuffix(rebuilt.Bytes(), []byte("\n")) + if len(mapping.NamespaceTools) == 0 { + return rebuiltPayload, true, nil + } + restored, _, err := RestoreResponsesNamespaceCalls(rebuiltPayload, mapping.NamespaceTools) + if err != nil { + return payload, false, err + } + return restored, true, nil +} + +func restoreClientToolValue(value any, adapter *ResponsesClientToolMapping) bool { + changed := false + switch typed := value.(type) { + case []any: + for _, item := range typed { + changed = restoreClientToolValue(item, adapter) || changed + } + case map[string]any: + if strings.TrimSpace(stringValue(typed["type"])) == "function_call" { + name := strings.TrimSpace(stringValue(typed["name"])) + if adapter.CustomTools[name] { + typed["type"] = "custom_tool_call" + typed["input"] = extractCustomToolCallInput(rawObjectString(typed["arguments"])) + delete(typed, "arguments") + delete(typed, "namespace") + changed = true + } else if adapter.ToolSearch && name == toolSearchProxyName { + typed["type"] = "tool_search_call" + typed["execution"] = "client" + typed["arguments"] = json.RawMessage(toolSearchCallArgumentsJSON(rawObjectString(typed["arguments"]))) + delete(typed, "name") + delete(typed, "namespace") + changed = true + } + } + for _, child := range typed { + changed = restoreClientToolValue(child, adapter) || changed + } + } + return changed +} + +// ResponsesClientToolStreamRestorer restores client tool stream lifecycles. +// It is intentionally stateful because custom tools need their function +// arguments buffered until the upstream signals the call is complete. +type ResponsesClientToolStreamRestorer struct { + adapter ResponsesClientToolMapping + nextSeq int + seenSeq bool + calls map[string]*responsesClientToolStreamCall + byOutput map[int]*responsesClientToolStreamCall +} + +type responsesClientToolStreamCall struct { + kind string + name string + callID string + itemID string + outputIdx int + arguments strings.Builder +} + +func NewResponsesClientToolStreamRestorer(mapping ResponsesClientToolMapping) *ResponsesClientToolStreamRestorer { + return &ResponsesClientToolStreamRestorer{adapter: mapping, calls: make(map[string]*responsesClientToolStreamCall), byOutput: make(map[int]*responsesClientToolStreamCall)} +} + +// Restore transforms one upstream SSE event into zero or more client events. +// Returned sequence numbers are continuous even when function argument events +// are suppressed or a custom completion expands into two events. +func (r *ResponsesClientToolStreamRestorer) Restore(event ResponsesStreamEvent) []ResponsesStreamEvent { + if r == nil { + return []ResponsesStreamEvent{event} + } + if !r.seenSeq { + r.nextSeq = event.SequenceNumber + r.seenSeq = true + } + var out []ResponsesStreamEvent + emit := func(event ResponsesStreamEvent) { + event.SequenceNumber = r.nextSeq + r.nextSeq++ + out = append(out, event) + } + + switch event.Type { + case "response.output_item.added": + if call := r.recordItem(event); call != nil { + if call.kind == "custom" { + event.Item.Type = "custom_tool_call" + event.Item.Input = "" + event.Item.Arguments = "" + event.Item.Namespace = "" + } else { + event.Item.Type = "tool_search_call" + event.Item.Name = "" + event.Item.Arguments = "{}" + event.Item.Namespace = "" + } + } + emit(r.restoreNamespaceEvent(event)) + case "response.function_call_arguments.delta": + if call := r.callFor(event); call != nil { + _, _ = call.arguments.WriteString(event.Delta) + return nil + } + emit(r.restoreNamespaceEvent(event)) + case "response.function_call_arguments.done": + if call := r.callFor(event); call != nil { + if event.Arguments != "" { + call.arguments.Reset() + _, _ = call.arguments.WriteString(event.Arguments) + } + if call.kind == "custom" { + input := extractCustomToolCallInput(call.arguments.String()) + if input != "" { + emit(ResponsesStreamEvent{Type: "response.custom_tool_call_input.delta", OutputIndex: call.outputIdx, ItemID: call.itemID, Delta: input}) + } + emit(ResponsesStreamEvent{Type: "response.custom_tool_call_input.done", OutputIndex: call.outputIdx, ItemID: call.itemID, CallID: call.callID, Name: call.name, Input: input}) + } + return out + } + emit(r.restoreNamespaceEvent(event)) + case "response.output_item.done": + if call := r.recordItem(event); call != nil { + if call.kind == "custom" { + event.Item.Type = "custom_tool_call" + event.Item.Input = extractCustomToolCallInput(call.arguments.String()) + event.Item.Arguments = "" + event.Item.Namespace = "" + } else { + event.Item.Type = "tool_search_call" + event.Item.Name = "" + event.Item.Arguments = call.arguments.String() + if strings.TrimSpace(event.Item.Arguments) == "" { + event.Item.Arguments = "{}" + } + event.Item.Namespace = "" + } + delete(r.calls, call.itemID) + delete(r.calls, call.callID) + delete(r.byOutput, call.outputIdx) + } + emit(r.restoreNamespaceEvent(event)) + default: + // response.completed carries the non-stream representation. + if event.Response != nil { + restoreResponsesOutputClientTools(event.Response.Output, &r.adapter) + } + emit(r.restoreNamespaceEvent(event)) + } + return out +} + +// RestoreEvent restores one Responses SSE JSON data payload. Custom tool +// completions can expand to multiple payloads and proxy argument deltas can be +// intentionally dropped, hence the slice return value. +func (r *ResponsesClientToolStreamRestorer) RestoreEvent(payload []byte) ([][]byte, bool, error) { + if len(payload) == 0 { + return nil, false, nil + } + var wire struct { + Type string `json:"type"` + Sequence int `json:"sequence_number"` + } + if err := json.Unmarshal(payload, &wire); err != nil { + return nil, false, err + } + if wire.Type == "response.completed" || wire.Type == "response.incomplete" || wire.Type == "response.failed" { + restored, changed, err := RestoreResponsesClientToolPayload(payload, r.adapter) + if err != nil { + return nil, false, err + } + return r.resequenceRaw(restored, wire.Sequence, changed) + } + if !clientToolLifecycleEvent(wire.Type) { + return r.resequenceRaw(payload, wire.Sequence, false) + } + if !r.clientToolEventPayload(payload) { + return r.resequenceRaw(payload, wire.Sequence, false) + } + var event ResponsesStreamEvent + if err := json.Unmarshal(payload, &event); err != nil { + return nil, false, err + } + events := r.Restore(event) + if len(events) == 1 { + unchanged, err := json.Marshal(events[0]) + if err == nil && bytes.Equal(bytes.TrimSpace(unchanged), bytes.TrimSpace(payload)) { + return [][]byte{payload}, false, nil + } + } + result := make([][]byte, 0, len(events)) + for _, restored := range events { + encoded, err := json.Marshal(restored) + if err != nil { + return nil, false, err + } + result = append(result, encoded) + } + return result, true, nil +} + +func (r *ResponsesClientToolStreamRestorer) clientToolEventPayload(payload []byte) bool { + var raw struct { + ItemID string `json:"item_id"` + CallID string `json:"call_id"` + Name string `json:"name"` + OutputIndex int `json:"output_index"` + Item *struct { + Type string `json:"type"` + ID string `json:"id"` + CallID string `json:"call_id"` + Name string `json:"name"` + } `json:"item"` + } + if err := json.Unmarshal(payload, &raw); err != nil { + return false + } + if raw.Item != nil { + if raw.Item.Type != "function_call" { + return false + } + _, namespaceTool := r.adapter.NamespaceTools[raw.Item.Name] + return r.adapter.CustomTools[raw.Item.Name] || (r.adapter.ToolSearch && raw.Item.Name == toolSearchProxyName) || namespaceTool || r.calls[raw.Item.ID] != nil || r.calls[raw.Item.CallID] != nil + } + if _, namespaceTool := r.adapter.NamespaceTools[raw.Name]; namespaceTool { + return true + } + if r.calls[raw.ItemID] != nil || r.calls[raw.CallID] != nil || r.byOutput[raw.OutputIndex] != nil { + return true + } + return false +} + +func clientToolLifecycleEvent(typ string) bool { + switch typ { + case "response.output_item.added", "response.output_item.done", "response.function_call_arguments.delta", "response.function_call_arguments.done": + return true + default: + return false + } +} + +// resequenceRaw deliberately keeps opaque upstream event fields untouched. +func (r *ResponsesClientToolStreamRestorer) resequenceRaw(payload []byte, sequence int, changed bool) ([][]byte, bool, error) { + if !r.seenSeq { + r.nextSeq, r.seenSeq = sequence, true + } + if r.nextSeq == sequence && !changed { + r.nextSeq++ + return [][]byte{payload}, false, nil + } + var raw map[string]any + if err := json.Unmarshal(payload, &raw); err != nil { + return nil, false, err + } + raw["sequence_number"] = r.nextSeq + r.nextSeq++ + encoded, err := json.Marshal(raw) + if err != nil { + return nil, false, err + } + return [][]byte{encoded}, true, nil +} + +func (r *ResponsesClientToolStreamRestorer) recordItem(event ResponsesStreamEvent) *responsesClientToolStreamCall { + if event.Item == nil || event.Item.Type != "function_call" { + return nil + } + name := event.Item.Name + kind := "" + if r.adapter.CustomTools[name] { + kind = "custom" + } else if r.adapter.ToolSearch && name == toolSearchProxyName { + kind = "tool_search" + } + if kind == "" { + return nil + } + key := event.Item.ID + if key == "" { + key = event.Item.CallID + } + call := r.calls[key] + if call == nil { + call = &responsesClientToolStreamCall{kind: kind, name: name, callID: event.Item.CallID, itemID: event.Item.ID, outputIdx: event.OutputIndex} + r.calls[key] = call + if call.callID != "" { + r.calls[call.callID] = call + } + r.byOutput[call.outputIdx] = call + } + if event.Item.Arguments != "" { + call.arguments.Reset() + _, _ = call.arguments.WriteString(event.Item.Arguments) + } + return call +} + +func (r *ResponsesClientToolStreamRestorer) callFor(event ResponsesStreamEvent) *responsesClientToolStreamCall { + if call := r.calls[event.ItemID]; call != nil { + return call + } + if call := r.byOutput[event.OutputIndex]; call != nil { + return call + } + for _, call := range r.calls { + if (event.CallID != "" && call.callID == event.CallID) || (event.ItemID == "" && event.Name != "" && call.name == event.Name) { + return call + } + } + return nil +} + +func (r *ResponsesClientToolStreamRestorer) restoreNamespaceEvent(event ResponsesStreamEvent) ResponsesStreamEvent { + if len(r.adapter.NamespaceTools) == 0 { + return event + } + if event.Item != nil && event.Item.Type == "function_call" { + if name, ok := r.adapter.NamespaceTools[event.Item.Name]; ok { + event.Item.Name, event.Item.Namespace = name.Name, name.Namespace + } + } + if event.Type == "response.function_call_arguments.done" { + if name, ok := r.adapter.NamespaceTools[event.Name]; ok { + event.Name = name.Name + } + } + return event +} + +func restoreResponsesOutputClientTools(outputs []ResponsesOutput, adapter *ResponsesClientToolMapping) { + for index := range outputs { + output := &outputs[index] + if output.Type != "function_call" { + continue + } + if adapter.CustomTools[output.Name] { + output.Type = "custom_tool_call" + output.Input = extractCustomToolCallInput(output.Arguments) + output.Arguments = "" + output.Namespace = "" + } else if adapter.ToolSearch && output.Name == toolSearchProxyName { + output.Type = "tool_search_call" + output.Name = "" + output.Namespace = "" + } + if name, ok := adapter.NamespaceTools[output.Name]; ok && output.Type == "function_call" { + output.Name, output.Namespace = name.Name, name.Namespace + } + } +} diff --git a/backend/internal/pkg/apicompat/responses_client_tools_test.go b/backend/internal/pkg/apicompat/responses_client_tools_test.go new file mode 100644 index 000000000..61dc0d648 --- /dev/null +++ b/backend/internal/pkg/apicompat/responses_client_tools_test.go @@ -0,0 +1,170 @@ +package apicompat + +import ( + "encoding/json" + "testing" + + "github.com/stretchr/testify/require" + "github.com/tidwall/gjson" +) + +func TestAdaptResponsesClientTools_LowersDeclarationsHistoryChoiceAndNamespaces(t *testing.T) { + req := map[string]any{ + "tools": []any{ + map[string]any{"type": "custom", "name": "exec", "format": map[string]any{"type": "grammar"}}, + map[string]any{"type": "tool_search"}, + map[string]any{"type": "namespace", "name": "team", "tools": []any{map[string]any{"type": "function", "name": "send"}}}, + }, + "tool_choice": map[string]any{"type": "custom", "name": "exec"}, + "input": []any{ + map[string]any{"type": "custom_tool_call", "call_id": "c1", "name": "exec", "input": "dir"}, + map[string]any{"type": "custom_tool_call_output", "call_id": "c1", "output": "ok"}, + map[string]any{"type": "tool_search_call", "call_id": "s1", "arguments": map[string]any{"query": "git"}}, + map[string]any{"type": "tool_search_output", "call_id": "s1", "output": map[string]any{"groups": []string{"git"}}}, + map[string]any{"type": "function_call", "call_id": "n1", "namespace": "team", "name": "send", "arguments": "{}"}, + }, + } + + mapping, changed, err := AdaptResponsesClientTools(req) + require.NoError(t, err) + require.True(t, changed) + require.True(t, mapping.CustomTools["exec"]) + require.True(t, mapping.ToolSearch) + require.Equal(t, ResponsesNamespaceName{Namespace: "team", Name: "send"}, mapping.NamespaceTools["team__send"]) + + tools := requireResponsesClientToolValue[[]any](t, req["tools"]) + require.Len(t, tools, 3) + exec := requireResponsesClientToolValue[map[string]any](t, tools[0]) + require.Equal(t, "function", exec["type"]) + parameters := requireResponsesClientToolValue[json.RawMessage](t, exec["parameters"]) + require.JSONEq(t, customToolInputSchema, string(parameters)) + search := requireResponsesClientToolValue[map[string]any](t, tools[1]) + require.Equal(t, toolSearchProxyName, search["name"]) + namespaceTool := requireResponsesClientToolValue[map[string]any](t, tools[2]) + require.Equal(t, "team__send", namespaceTool["name"]) + + choice := requireResponsesClientToolValue[map[string]any](t, req["tool_choice"]) + require.Equal(t, "function", choice["type"]) + input := requireResponsesClientToolValue[[]any](t, req["input"]) + customCall := requireResponsesClientToolValue[map[string]any](t, input[0]) + require.Equal(t, "function_call", customCall["type"]) + require.JSONEq(t, `{"input":"dir"}`, requireResponsesClientToolValue[string](t, customCall["arguments"])) + customOutput := requireResponsesClientToolValue[map[string]any](t, input[1]) + require.Equal(t, "function_call_output", customOutput["type"]) + searchCall := requireResponsesClientToolValue[map[string]any](t, input[2]) + require.Equal(t, "function_call", searchCall["type"]) + require.Equal(t, toolSearchProxyName, searchCall["name"]) + require.JSONEq(t, `{"query":"git"}`, requireResponsesClientToolValue[string](t, searchCall["arguments"])) + searchOutput := requireResponsesClientToolValue[map[string]any](t, input[3]) + require.Equal(t, "function_call_output", searchOutput["type"]) + require.JSONEq(t, `{"groups":["git"]}`, requireResponsesClientToolValue[string](t, searchOutput["output"])) + namespaceCall := requireResponsesClientToolValue[map[string]any](t, input[4]) + require.Equal(t, "team__send", namespaceCall["name"]) +} + +func requireResponsesClientToolValue[T any](t *testing.T, value any) T { + t.Helper() + typed, ok := value.(T) + require.True(t, ok, "unexpected value type %T", value) + return typed +} + +func TestAdaptResponsesClientTools_RejectsAmbiguousNames(t *testing.T) { + cases := []map[string]any{ + {"tools": []any{map[string]any{"type": "custom", "name": "same"}, map[string]any{"type": "function", "name": "same"}}}, + {"tools": []any{map[string]any{"type": "tool_search"}, map[string]any{"type": "function", "name": "tool_search"}}}, + {"tools": []any{map[string]any{"type": "function", "name": "team__send"}, map[string]any{"type": "namespace", "name": "team", "tools": []any{map[string]any{"type": "function", "name": "send"}}}}}, + } + for _, req := range cases { + _, _, err := AdaptResponsesClientTools(req) + require.Error(t, err) + } +} + +func TestRestoreResponsesClientToolPayload_RestoresClientAndNamespaceCalls(t *testing.T) { + mapping := ResponsesClientToolMapping{ + CustomTools: map[string]bool{"exec": true}, ToolSearch: true, + NamespaceTools: map[string]ResponsesNamespaceName{"team__send": {Namespace: "team", Name: "send"}}, + } + payload := []byte(`{"id":"resp","output":[{"type":"function_call","id":"i1","call_id":"c1","name":"exec","arguments":"{\"input\":\"dir\"}","namespace":"ignore"},{"type":"function_call","id":"i2","call_id":"s1","name":"tool_search","arguments":"{\"query\":\"git\"}"},{"type":"function_call","id":"i3","call_id":"n1","name":"team__send","arguments":"{}"}]}`) + + restored, changed, err := RestoreResponsesClientToolPayload(payload, mapping) + require.NoError(t, err) + require.True(t, changed) + require.JSONEq(t, `{"id":"resp","output":[{"type":"custom_tool_call","id":"i1","call_id":"c1","name":"exec","input":"dir"},{"type":"tool_search_call","id":"i2","call_id":"s1","execution":"client","arguments":{"query":"git"}},{"type":"function_call","id":"i3","call_id":"n1","name":"send","namespace":"team","arguments":"{}"}]}`, string(restored)) +} + +func TestResponsesClientToolStreamRestorer_CustomToolBuffersWrapperAndSequences(t *testing.T) { + restorer := NewResponsesClientToolStreamRestorer(ResponsesClientToolMapping{CustomTools: map[string]bool{"exec": true}}) + added := restorer.Restore(ResponsesStreamEvent{Type: "response.output_item.added", SequenceNumber: 7, OutputIndex: 0, Item: &ResponsesOutput{Type: "function_call", ID: "i1", CallID: "c1", Name: "exec", Status: "in_progress"}}) + require.Len(t, added, 1) + require.Equal(t, 7, added[0].SequenceNumber) + require.Equal(t, "custom_tool_call", added[0].Item.Type) + require.Empty(t, restorer.Restore(ResponsesStreamEvent{Type: "response.function_call_arguments.delta", SequenceNumber: 8, ItemID: "i1", Delta: `{"input":"di`})) + done := restorer.Restore(ResponsesStreamEvent{Type: "response.function_call_arguments.done", SequenceNumber: 9, ItemID: "i1", CallID: "c1", Name: "exec", Arguments: `{"input":"dir"}`}) + require.Len(t, done, 2) + require.Equal(t, 8, done[0].SequenceNumber) + require.Equal(t, "response.custom_tool_call_input.delta", done[0].Type) + require.Equal(t, "dir", done[0].Delta) + require.Equal(t, 9, done[1].SequenceNumber) + require.Equal(t, "response.custom_tool_call_input.done", done[1].Type) + require.Equal(t, "dir", done[1].Input) + closed := restorer.Restore(ResponsesStreamEvent{Type: "response.output_item.done", SequenceNumber: 10, OutputIndex: 0, Item: &ResponsesOutput{Type: "function_call", ID: "i1", CallID: "c1", Name: "exec", Arguments: `{"input":"dir"}`, Status: "completed"}}) + require.Equal(t, 10, closed[0].SequenceNumber) + require.Equal(t, "custom_tool_call", closed[0].Item.Type) + require.Equal(t, "dir", closed[0].Item.Input) +} + +func TestResponsesClientToolStreamRestorer_ToolSearchAndFunction(t *testing.T) { + restorer := NewResponsesClientToolStreamRestorer(ResponsesClientToolMapping{ToolSearch: true}) + search := restorer.Restore(ResponsesStreamEvent{Type: "response.output_item.added", SequenceNumber: 0, OutputIndex: 0, Item: &ResponsesOutput{Type: "function_call", ID: "s1", CallID: "c1", Name: "tool_search", Status: "in_progress"}}) + require.Equal(t, "tool_search_call", search[0].Item.Type) + require.Empty(t, restorer.Restore(ResponsesStreamEvent{Type: "response.function_call_arguments.delta", SequenceNumber: 1, ItemID: "s1", Delta: `{"query":"git"}`})) + require.Empty(t, restorer.Restore(ResponsesStreamEvent{Type: "response.function_call_arguments.done", SequenceNumber: 2, ItemID: "s1", Arguments: `{"query":"git"}`})) + closed := restorer.Restore(ResponsesStreamEvent{Type: "response.output_item.done", SequenceNumber: 3, OutputIndex: 0, Item: &ResponsesOutput{Type: "function_call", ID: "s1", CallID: "c1", Name: "tool_search", Status: "completed"}}) + require.Equal(t, 1, closed[0].SequenceNumber) + require.Equal(t, "tool_search_call", closed[0].Item.Type) + require.JSONEq(t, `{"query":"git"}`, string(toolSearchCallArgumentsJSON(closed[0].Item.Arguments))) + + function := restorer.Restore(ResponsesStreamEvent{Type: "response.function_call_arguments.done", SequenceNumber: 4, ItemID: "plain", Name: "plain", Arguments: "{}"}) + require.Len(t, function, 1) + require.Equal(t, "response.function_call_arguments.done", function[0].Type) + require.Equal(t, 2, function[0].SequenceNumber) +} + +func TestResponsesClientToolStreamRestorer_RestoresNamespaceLifecycle(t *testing.T) { + restorer := NewResponsesClientToolStreamRestorer(ResponsesClientToolMapping{ + NamespaceTools: map[string]ResponsesNamespaceName{ + "browser__open": {Namespace: "browser", Name: "open"}, + }, + }) + + added, changed, err := restorer.RestoreEvent([]byte(`{"type":"response.output_item.added","sequence_number":4,"output_index":0,"item":{"type":"function_call","id":"i1","call_id":"c1","name":"browser__open","arguments":"","status":"in_progress"}}`)) + require.NoError(t, err) + require.True(t, changed) + require.Len(t, added, 1) + require.Equal(t, "open", gjson.GetBytes(added[0], "item.name").String()) + require.Equal(t, "browser", gjson.GetBytes(added[0], "item.namespace").String()) + + done, changed, err := restorer.RestoreEvent([]byte(`{"type":"response.function_call_arguments.done","sequence_number":5,"output_index":0,"item_id":"i1","name":"browser__open","arguments":"{}"}`)) + require.NoError(t, err) + require.True(t, changed) + require.Len(t, done, 1) + require.Equal(t, "open", gjson.GetBytes(done[0], "name").String()) +} + +func TestResponsesClientToolStreamRestorer_RawEventsPreserveUnknownFieldsAndOutputFallback(t *testing.T) { + restorer := NewResponsesClientToolStreamRestorer(ResponsesClientToolMapping{CustomTools: map[string]bool{"exec": true}}) + passthrough, changed, err := restorer.RestoreEvent([]byte(`{"type":"response.created","sequence_number":4,"response":{"id":"r"},"upstream_extension":{"keep":true}}`)) + require.NoError(t, err) + require.False(t, changed) + require.Len(t, passthrough, 1) + require.Contains(t, string(passthrough[0]), `"upstream_extension":{"keep":true}`) + + restorer.Restore(ResponsesStreamEvent{Type: "response.output_item.added", SequenceNumber: 5, OutputIndex: 9, Item: &ResponsesOutput{Type: "function_call", ID: "item", CallID: "call", Name: "exec"}}) + // Some upstreams omit every tool identity field on later argument chunks. + require.Empty(t, restorer.Restore(ResponsesStreamEvent{Type: "response.function_call_arguments.delta", SequenceNumber: 6, OutputIndex: 9, Delta: `{"input":"pwd"}`})) + done := restorer.Restore(ResponsesStreamEvent{Type: "response.function_call_arguments.done", SequenceNumber: 7, OutputIndex: 9}) + require.Len(t, done, 2) + require.Equal(t, "pwd", done[1].Input) +} diff --git a/backend/internal/service/openai_gateway_forward.go b/backend/internal/service/openai_gateway_forward.go index 93697be77..7a42a445e 100644 --- a/backend/internal/service/openai_gateway_forward.go +++ b/backend/internal/service/openai_gateway_forward.go @@ -19,6 +19,7 @@ import ( // Forward forwards request to OpenAI API func (s *OpenAIGatewayService) Forward(ctx context.Context, c *gin.Context, account *Account, body []byte) (*OpenAIForwardResult, error) { + clearGrokResponsesClientToolMapping(c) startTime := time.Now() // 固定渠道映射后的请求级 canonical body;账号 normalize/strip 不得改写跨 failover hint。 canonicalImageIntentBody := body diff --git a/backend/internal/service/openai_gateway_grok.go b/backend/internal/service/openai_gateway_grok.go index 1ce73dc4b..7077aca1c 100644 --- a/backend/internal/service/openai_gateway_grok.go +++ b/backend/internal/service/openai_gateway_grok.go @@ -52,10 +52,15 @@ func (s *OpenAIGatewayService) forwardGrokResponses( if isGrokImageGenerationModel(upstreamModel) { return nil, fmt.Errorf("model %s is an image model and is not available on the Responses endpoint; use /v1/images/generations instead", upstreamModel) } - patchedBody, err := patchGrokResponsesBody(body, upstreamModel) + patchedBody, clientToolMapping, err := patchGrokResponsesBodyWithClientTools(body, upstreamModel) if err != nil { + setOpsUpstreamError(c, http.StatusBadRequest, err.Error(), "") + c.JSON(http.StatusBadRequest, gin.H{"error": gin.H{ + "type": "invalid_request_error", "message": err.Error(), "param": "tools", + }}) return nil, err } + setGrokResponsesClientToolMapping(c, clientToolMapping) // OpenAI /responses/compact is not a native xAI endpoint. Convert it into a // normal Grok Responses turn that asks for a structured summary, then map the // reply back to an OpenAI compaction item on the way out. @@ -170,6 +175,13 @@ func (s *OpenAIGatewayService) forwardGrokResponses( var firstTokenMs *int responseID := "" if reqStream { + if hasGrokResponsesClientToolMapping(clientToolMapping) { + maxLineSize := defaultMaxLineSize + if s.cfg != nil && s.cfg.Gateway.MaxLineSize > 0 { + maxLineSize = s.cfg.Gateway.MaxLineSize + } + resp.Body = newGrokResponsesClientToolStreamBody(resp.Body, clientToolMapping, maxLineSize) + } streamResult, err := s.handleStreamingResponse(ctx, resp, c, account, startTime, originalModel, upstreamModel) if err != nil { return nil, err @@ -367,6 +379,29 @@ func trimGrokInvalidEncryptedContentRetryBody(body []byte) ([]byte, bool, error) } func patchGrokResponsesBody(body []byte, upstreamModel string) ([]byte, error) { + return patchGrokResponsesBodyBase(body, upstreamModel) +} + +func patchGrokResponsesBodyWithClientTools(body []byte, upstreamModel string) ([]byte, apicompat.ResponsesClientToolMapping, error) { + if !json.Valid(body) { + return nil, apicompat.ResponsesClientToolMapping{}, fmt.Errorf("invalid json request body") + } + promoted, err := sanitizeGrokResponsesInput(body) + if err != nil { + return nil, apicompat.ResponsesClientToolMapping{}, err + } + adapted, mapping, err := adaptGrokResponsesClientTools(promoted) + if err != nil { + return nil, apicompat.ResponsesClientToolMapping{}, err + } + patched, err := patchGrokResponsesBodyBase(adapted, upstreamModel) + if err != nil { + return nil, apicompat.ResponsesClientToolMapping{}, err + } + return patched, mapping, nil +} + +func patchGrokResponsesBodyBase(body []byte, upstreamModel string) ([]byte, error) { if !json.Valid(body) { return nil, fmt.Errorf("invalid json request body") } diff --git a/backend/internal/service/openai_gateway_grok_test.go b/backend/internal/service/openai_gateway_grok_test.go index 2f33fe489..922a57be0 100644 --- a/backend/internal/service/openai_gateway_grok_test.go +++ b/backend/internal/service/openai_gateway_grok_test.go @@ -162,29 +162,31 @@ func TestPatchGrokResponsesBodyDropsNestedUnsupportedFields(t *testing.T) { require.Equal(t, "kept_fn", gjson.GetBytes(patched, "tools.0.name").String()) } -func TestPatchGrokResponsesBodyDropsUnsupportedNamespaceTools(t *testing.T) { +func TestPatchGrokResponsesBodyFlattensNamespaceTools(t *testing.T) { t.Parallel() body := []byte(`{ "model": "grok", "input": "hello", "tools": [ - {"type": "namespace", "namespace": "functions", "tools": [{"type": "function", "name": "inner"}]}, + {"type": "namespace", "name": "functions", "tools": [{"type": "function", "name": "inner"}]}, {"type": "function", "name": "kept_fn", "parameters": {"type": "object"}}, {"type": "shell", "name": "kept_shell"} ], - "tool_choice": {"type": "function", "name": "kept_fn"} + "tool_choice": {"type": "function", "namespace": "functions", "name": "inner"} }`) - patched, err := patchGrokResponsesBody(body, "grok-4.3") + patched, _, err := patchGrokResponsesBodyWithClientTools(body, "grok-4.3") require.NoError(t, err) require.True(t, json.Valid(patched)) require.Equal(t, "grok-4.3", gjson.GetBytes(patched, "model").String()) - require.Len(t, gjson.GetBytes(patched, "tools").Array(), 2) + require.Len(t, gjson.GetBytes(patched, "tools").Array(), 3) require.False(t, gjson.GetBytes(patched, `tools.#(type=="namespace")`).Exists()) require.True(t, gjson.GetBytes(patched, `tools.#(type=="function")`).Exists()) require.True(t, gjson.GetBytes(patched, `tools.#(type=="shell")`).Exists()) - require.Equal(t, "kept_fn", gjson.GetBytes(patched, "tool_choice.name").String()) + require.Equal(t, "functions__inner", gjson.GetBytes(patched, "tools.0.name").String()) + require.Equal(t, "functions__inner", gjson.GetBytes(patched, "tool_choice.name").String()) + require.False(t, gjson.GetBytes(patched, "tool_choice.namespace").Exists()) } func TestPatchGrokResponsesBodyDropsToolChoiceWhenNoSupportedToolsRemain(t *testing.T) { @@ -243,19 +245,22 @@ func TestPatchGrokResponsesBodyPromotesCodexAdditionalTools(t *testing.T) { ] }`) - patched, err := patchGrokResponsesBody(body, "grok-4.5") + patched, _, err := patchGrokResponsesBodyWithClientTools(body, "grok-4.5") require.NoError(t, err) require.True(t, json.Valid(patched)) require.Equal(t, "grok-4.5", gjson.GetBytes(patched, "model").String()) require.Equal(t, 2, len(gjson.GetBytes(patched, "input").Array())) require.False(t, gjson.GetBytes(patched, `input.#(type=="additional_tools")`).Exists()) tools := gjson.GetBytes(patched, "tools").Array() - require.Len(t, tools, 4) + require.Len(t, tools, 5) require.Equal(t, "existing", tools[0].Get("name").String()) require.Equal(t, "top-level wins", tools[0].Get("description").String()) require.Equal(t, "web_search", tools[1].Get("type").String()) require.Equal(t, "wait", tools[2].Get("name").String()) require.Equal(t, "shell", tools[3].Get("type").String()) + require.Equal(t, "function", tools[4].Get("type").String()) + require.Equal(t, "apply_patch", tools[4].Get("name").String()) + require.Equal(t, "string", tools[4].Get("parameters.properties.input.type").String()) require.False(t, gjson.GetBytes(patched, `tools.#(type=="custom")`).Exists()) require.False(t, gjson.GetBytes(patched, `tools.#(type=="namespace")`).Exists()) require.Equal(t, "auto", gjson.GetBytes(patched, "tool_choice").String()) @@ -316,11 +321,13 @@ func TestForwardGrokResponsesCodexAdditionalToolsUsesMixedCacheIntent(t *testing require.Equal(t, "resp_codex_lite", result.ResponseID) 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()) diff --git a/backend/internal/service/openai_gateway_grok_tool_protocol.go b/backend/internal/service/openai_gateway_grok_tool_protocol.go new file mode 100644 index 000000000..1a8dfc35a --- /dev/null +++ b/backend/internal/service/openai_gateway_grok_tool_protocol.go @@ -0,0 +1,255 @@ +package service + +import ( + "bufio" + "bytes" + "encoding/json" + "fmt" + "io" + "strings" + + "github.com/Wei-Shaw/sub2api/internal/pkg/apicompat" + "github.com/gin-gonic/gin" + "github.com/tidwall/gjson" +) + +const grokResponsesClientToolMappingContextKey = "grok_responses_client_tool_mapping" + +func adaptGrokResponsesClientTools(body []byte) ([]byte, apicompat.ResponsesClientToolMapping, error) { + decoder := json.NewDecoder(bytes.NewReader(body)) + decoder.UseNumber() + var requestBody map[string]any + if err := decoder.Decode(&requestBody); err != nil { + return body, apicompat.ResponsesClientToolMapping{}, fmt.Errorf("decode Grok Responses client tools: %w", err) + } + + mapping, changed, err := apicompat.AdaptResponsesClientTools(requestBody) + if err != nil { + return body, apicompat.ResponsesClientToolMapping{}, err + } + if !changed { + return body, mapping, nil + } + rebuilt, err := marshalOpenAIUpstreamJSON(requestBody) + if err != nil { + return body, apicompat.ResponsesClientToolMapping{}, fmt.Errorf("encode Grok Responses client tools: %w", err) + } + return rebuilt, mapping, nil +} + +func hasGrokResponsesClientToolMapping(mapping apicompat.ResponsesClientToolMapping) bool { + return len(mapping.CustomTools) > 0 || mapping.ToolSearch || len(mapping.NamespaceTools) > 0 +} + +func setGrokResponsesClientToolMapping(c *gin.Context, mapping apicompat.ResponsesClientToolMapping) { + if c == nil { + return + } + if !hasGrokResponsesClientToolMapping(mapping) { + clearGrokResponsesClientToolMapping(c) + return + } + c.Set(grokResponsesClientToolMappingContextKey, mapping) +} + +func clearGrokResponsesClientToolMapping(c *gin.Context) { + if c == nil { + return + } + if _, exists := c.Get(grokResponsesClientToolMappingContextKey); !exists { + return + } + c.Set(grokResponsesClientToolMappingContextKey, apicompat.ResponsesClientToolMapping{}) +} + +func grokResponsesClientToolMapping(c *gin.Context) (apicompat.ResponsesClientToolMapping, bool) { + if c == nil { + return apicompat.ResponsesClientToolMapping{}, false + } + value, ok := c.Get(grokResponsesClientToolMappingContextKey) + if !ok { + return apicompat.ResponsesClientToolMapping{}, false + } + mapping, ok := value.(apicompat.ResponsesClientToolMapping) + return mapping, ok && hasGrokResponsesClientToolMapping(mapping) +} + +func restoreGrokResponsesClientToolPayload(c *gin.Context, payload []byte) ([]byte, error) { + mapping, ok := grokResponsesClientToolMapping(c) + if !ok || !bytes.Contains(payload, []byte(`"function_call"`)) || !json.Valid(payload) { + return payload, nil + } + restored, _, err := apicompat.RestoreResponsesClientToolPayload(payload, mapping) + return restored, err +} + +type grokResponsesClientToolStreamBody struct { + *io.PipeReader + source io.Closer +} + +func (b *grokResponsesClientToolStreamBody) Close() error { + readerErr := b.PipeReader.Close() + sourceErr := b.source.Close() + if readerErr != nil { + return readerErr + } + return sourceErr +} + +func newGrokResponsesClientToolStreamBody( + source io.ReadCloser, + mapping apicompat.ResponsesClientToolMapping, + maxLineSize int, +) io.ReadCloser { + reader, writer := io.Pipe() + body := &grokResponsesClientToolStreamBody{PipeReader: reader, source: source} + go transformGrokResponsesClientToolStream(source, writer, mapping, maxLineSize) + return body +} + +func transformGrokResponsesClientToolStream( + source io.ReadCloser, + destination *io.PipeWriter, + mapping apicompat.ResponsesClientToolMapping, + maxLineSize int, +) { + defer func() { _ = source.Close() }() + if maxLineSize <= 0 { + maxLineSize = defaultMaxLineSize + } + + scanner := bufio.NewScanner(source) + scanBuf := getSSEScannerBuf64K() + defer putSSEScannerBuf64K(scanBuf) + scanner.Buffer(scanBuf[:0], maxLineSize) + documents := newOpenAISSEJSONDocumentScanner(scanner) + restorer := apicompat.NewResponsesClientToolStreamRestorer(mapping) + buffered := bufio.NewWriterSize(destination, 4*1024) + pendingFields := make([]string, 0, 2) + frameHadEventField := false + frameEmitted := false + + writeLine := func(line string) error { + if _, err := buffered.WriteString(line); err != nil { + return err + } + return buffered.WriteByte('\n') + } + writePendingFields := func(payload []byte, includeNonEvent bool) error { + eventType := strings.TrimSpace(gjson.GetBytes(payload, "type").String()) + for _, field := range pendingFields { + if _, isEvent := extractOpenAISSEEventLine(field); isEvent { + if eventType != "" { + if err := writeLine("event: " + eventType); err != nil { + return err + } + } else if err := writeLine(field); err != nil { + return err + } + continue + } + if includeNonEvent { + if err := writeLine(field); err != nil { + return err + } + } + } + return nil + } + writePayloads := func(payloads [][]byte) error { + for index, payload := range payloads { + if index == 0 { + if err := writePendingFields(payload, true); err != nil { + return err + } + } else if frameHadEventField { + eventType := strings.TrimSpace(gjson.GetBytes(payload, "type").String()) + if eventType != "" { + if err := writeLine("event: " + eventType); err != nil { + return err + } + } + } + if err := writeLine("data: " + string(payload)); err != nil { + return err + } + if err := writeLine(""); err != nil { + return err + } + } + return buffered.Flush() + } + + for documents.Scan() { + line := documents.Text() + data, isData := extractOpenAISSEDataLine(line) + if isData { + payload := []byte(data) + payloads := [][]byte{payload} + if json.Valid(payload) { + var err error + payloads, _, err = restorer.RestoreEvent(payload) + if err != nil { + _ = buffered.Flush() + _ = destination.CloseWithError(fmt.Errorf("restore Grok Responses client tool event: %w", err)) + return + } + } + if err := writePayloads(payloads); err != nil { + _ = destination.CloseWithError(err) + return + } + pendingFields = pendingFields[:0] + frameHadEventField = false + frameEmitted = true + continue + } + + if line == "" { + if !frameEmitted { + for _, field := range pendingFields { + if err := writeLine(field); err != nil { + _ = destination.CloseWithError(err) + return + } + } + if len(pendingFields) > 0 { + if err := writeLine(""); err != nil { + _ = destination.CloseWithError(err) + return + } + if err := buffered.Flush(); err != nil { + _ = destination.CloseWithError(err) + return + } + } + } + pendingFields = pendingFields[:0] + frameHadEventField = false + frameEmitted = false + continue + } + + if _, isEvent := extractOpenAISSEEventLine(line); isEvent { + frameHadEventField = true + } + pendingFields = append(pendingFields, line) + } + + for _, field := range pendingFields { + if err := writeLine(field); err != nil { + _ = destination.CloseWithError(err) + return + } + } + if err := buffered.Flush(); err != nil { + _ = destination.CloseWithError(err) + return + } + if err := documents.Err(); err != nil { + _ = destination.CloseWithError(err) + return + } + _ = destination.Close() +} diff --git a/backend/internal/service/openai_gateway_grok_tool_protocol_test.go b/backend/internal/service/openai_gateway_grok_tool_protocol_test.go new file mode 100644 index 000000000..bca3d3f2b --- /dev/null +++ b/backend/internal/service/openai_gateway_grok_tool_protocol_test.go @@ -0,0 +1,518 @@ +//go:build unit + +package service + +import ( + "bufio" + "bytes" + "context" + "encoding/json" + "fmt" + "io" + "net/http" + "net/http/httptest" + "strings" + "testing" + "time" + + "github.com/Wei-Shaw/sub2api/internal/pkg/apicompat" + "github.com/Wei-Shaw/sub2api/internal/pkg/xai" + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/require" + "github.com/tidwall/gjson" +) + +func TestPatchGrokResponsesBodyWithClientToolsLowersCodexProtocol(t *testing.T) { + t.Parallel() + + body := grokClientToolProtocolRequest(false) + patched, mapping, err := patchGrokResponsesBodyWithClientTools(body, "grok-4.5") + require.NoError(t, err) + require.True(t, json.Valid(patched)) + require.True(t, mapping.CustomTools["apply_patch"]) + require.True(t, mapping.ToolSearch) + require.Equal(t, "collaboration", mapping.NamespaceTools["collaboration__send_message"].Namespace) + require.Equal(t, "send_message", mapping.NamespaceTools["collaboration__send_message"].Name) + + tools := gjson.GetBytes(patched, "tools").Array() + require.Len(t, tools, 3) + require.Equal(t, "function", tools[0].Get("type").String()) + require.Equal(t, "apply_patch", tools[0].Get("name").String()) + require.Equal(t, "string", tools[0].Get("parameters.properties.input.type").String()) + require.False(t, tools[0].Get("format").Exists()) + require.Equal(t, "function", tools[1].Get("type").String()) + require.Equal(t, "tool_search", tools[1].Get("name").String()) + require.Equal(t, "function", tools[2].Get("type").String()) + require.Equal(t, "collaboration__send_message", tools[2].Get("name").String()) + require.False(t, gjson.GetBytes(patched, `tools.#(type=="custom")`).Exists()) + require.False(t, gjson.GetBytes(patched, `tools.#(type=="namespace")`).Exists()) + require.False(t, gjson.GetBytes(patched, `tools.#(type=="tool_search")`).Exists()) + + require.Equal(t, "function", gjson.GetBytes(patched, "tool_choice.type").String()) + require.Equal(t, "apply_patch", gjson.GetBytes(patched, "tool_choice.name").String()) + require.Equal(t, "function_call", gjson.GetBytes(patched, "input.0.type").String()) + require.JSONEq(t, `{"input":"*** Begin Patch"}`, gjson.GetBytes(patched, "input.0.arguments").String()) + require.False(t, gjson.GetBytes(patched, "input.0.input").Exists()) + require.Equal(t, "function_call_output", gjson.GetBytes(patched, "input.1.type").String()) + require.Equal(t, "function_call", gjson.GetBytes(patched, "input.2.type").String()) + require.Equal(t, "tool_search", gjson.GetBytes(patched, "input.2.name").String()) + require.JSONEq(t, `{"query":"github"}`, gjson.GetBytes(patched, "input.2.arguments").String()) + require.False(t, gjson.GetBytes(patched, "input.2.execution").Exists()) + require.Equal(t, "function_call_output", gjson.GetBytes(patched, "input.3.type").String()) + require.JSONEq(t, `{"groups":["github"]}`, gjson.GetBytes(patched, "input.3.output").String()) + require.Equal(t, "function_call", gjson.GetBytes(patched, "input.4.type").String()) + require.Equal(t, "collaboration__send_message", gjson.GetBytes(patched, "input.4.name").String()) + require.False(t, gjson.GetBytes(patched, "input.4.namespace").Exists()) +} + +func TestPatchGrokResponsesBodyWithClientToolsRewritesEveryToolChoice(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + choice string + wantName string + wantType string + wantNoNS bool + }{ + { + name: "custom", + choice: `{"type":"custom","name":"apply_patch"}`, + wantName: "apply_patch", + wantType: "function", + }, + { + name: "tool search", + choice: `{"type":"tool_search"}`, + wantName: "tool_search", + wantType: "function", + }, + { + name: "namespace function", + choice: `{"type":"function","namespace":"collaboration","name":"send_message"}`, + wantName: "collaboration__send_message", + wantType: "function", + wantNoNS: true, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + body := []byte(fmt.Sprintf(`{ + "model":"grok","input":"hello", + "tools":[ + {"type":"custom","name":"apply_patch"}, + {"type":"tool_search"}, + {"type":"namespace","name":"collaboration","tools":[{"type":"function","name":"send_message","parameters":{"type":"object"}}]} + ], + "tool_choice":%s + }`, tt.choice)) + + patched, _, err := patchGrokResponsesBodyWithClientTools(body, "grok-4.5") + require.NoError(t, err) + require.Equal(t, tt.wantType, gjson.GetBytes(patched, "tool_choice.type").String()) + require.Equal(t, tt.wantName, gjson.GetBytes(patched, "tool_choice.name").String()) + if tt.wantNoNS { + require.False(t, gjson.GetBytes(patched, "tool_choice.namespace").Exists()) + } + }) + } +} + +func TestPatchGrokResponsesBodyWithClientToolsRejectsTrailingJSONDocument(t *testing.T) { + t.Parallel() + + body := []byte(`{"model":"grok","input":"hello","tools":[{"type":"custom","name":"apply_patch"}]} {"ignored":true}`) + patched, mapping, err := patchGrokResponsesBodyWithClientTools(body, "grok-4.5") + + require.Error(t, err) + require.Contains(t, strings.ToLower(err.Error()), "invalid json") + require.Nil(t, patched) + require.Empty(t, mapping.CustomTools) +} + +func TestClearGrokResponsesClientToolMappingRemovesStaleContextState(t *testing.T) { + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + setGrokResponsesClientToolMapping(c, apicompat.ResponsesClientToolMapping{ + CustomTools: map[string]bool{"stale_tool": true}, + }) + + _, seeded := grokResponsesClientToolMapping(c) + require.True(t, seeded) + clearGrokResponsesClientToolMapping(c) + _, remains := grokResponsesClientToolMapping(c) + require.False(t, remains) +} + +func TestForwardGrokResponsesClientToolNameConflictReturns400(t *testing.T) { + gin.SetMode(gin.TestMode) + + body := []byte(`{ + "model":"grok","stream":false,"input":"hello", + "tools":[ + {"type":"custom","name":"duplicate"}, + {"type":"function","name":"duplicate","parameters":{"type":"object"}} + ] + }`) + 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(7101) + + 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(), "conflicts") + require.Empty(t, upstream.requests, "an ambiguous request must not reach xAI") +} + +func TestForwardGrokResponsesOAuthRestoresClientToolsNonStreaming(t *testing.T) { + gin.SetMode(gin.TestMode) + + body := grokClientToolProtocolRequest(false) + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", bytes.NewReader(body)) + c.Request.Header.Set("Content-Type", "application/json") + c.Set("api_key", &APIKey{ID: 7102}) + + account := grokProtocolOAuthAccount(7102) + repo := &grokQuotaAccountRepo{mockAccountRepoForPlatform: &mockAccountRepoForPlatform{ + accountsByID: map[int64]*Account{account.ID: account}, + }} + upstream := &httpUpstreamRecorder{resp: &http.Response{ + StatusCode: http.StatusOK, + Header: http.Header{ + "Content-Type": []string{"application/json"}, + "Xai-Request-Id": []string{"protocol-oauth"}, + }, + Body: io.NopCloser(strings.NewReader(`{ + "id":"resp_protocol_oauth","object":"response","model":"grok-4.5","status":"completed", + "output":[ + {"type":"function_call","id":"item_custom","call_id":"call_custom","name":"apply_patch","arguments":"{\"input\":\"*** Begin Patch\"}","namespace":"must_not_leak"}, + {"type":"function_call","id":"item_search","call_id":"call_search","name":"tool_search","arguments":"{\"query\":\"github\"}"}, + {"type":"function_call","id":"item_namespace","call_id":"call_namespace","name":"collaboration__send_message","arguments":"{\"target\":\"root\"}"} + ], + "usage":{"input_tokens":9,"output_tokens":3,"total_tokens":12} + }`)), + }} + svc := &OpenAIGatewayService{ + httpUpstream: upstream, + grokTokenProvider: NewGrokTokenProvider(repo, nil), + accountRepo: repo, + } + + result, err := svc.forwardGrokResponses(context.Background(), c, account, body, "grok", false, time.Now()) + + require.NoError(t, err) + require.NotNil(t, result) + require.False(t, result.Stream) + require.Equal(t, "resp_protocol_oauth", result.ResponseID) + require.Equal(t, xai.DefaultCLIBaseURL+"/responses", upstream.lastReq.URL.String()) + require.Equal(t, "Bearer oauth-protocol-token", upstream.lastReq.Header.Get("Authorization")) + assertGrokProtocolRequestLowered(t, upstream.lastBody) + + response := recorder.Body.Bytes() + require.Equal(t, "custom_tool_call", gjson.GetBytes(response, "output.0.type").String()) + require.Equal(t, "*** Begin Patch", gjson.GetBytes(response, "output.0.input").String()) + require.False(t, gjson.GetBytes(response, "output.0.arguments").Exists()) + require.False(t, gjson.GetBytes(response, "output.0.namespace").Exists()) + require.Equal(t, "tool_search_call", gjson.GetBytes(response, "output.1.type").String()) + require.Equal(t, "client", gjson.GetBytes(response, "output.1.execution").String()) + require.Equal(t, "github", gjson.GetBytes(response, "output.1.arguments.query").String()) + require.False(t, gjson.GetBytes(response, "output.1.name").Exists()) + require.Equal(t, "function_call", gjson.GetBytes(response, "output.2.type").String()) + require.Equal(t, "collaboration", gjson.GetBytes(response, "output.2.namespace").String()) + require.Equal(t, "send_message", gjson.GetBytes(response, "output.2.name").String()) +} + +func TestForwardGrokResponsesAPIKeyRestoresClientToolsFromSSEForNonStreamingRequest(t *testing.T) { + gin.SetMode(gin.TestMode) + + body := grokClientToolProtocolRequest(false) + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", bytes.NewReader(body)) + c.Request.Header.Set("Content-Type", "application/json") + + upstream := &httpUpstreamRecorder{resp: &http.Response{ + StatusCode: http.StatusOK, + Header: http.Header{ + "Content-Type": []string{"text/event-stream"}, + "Xai-Request-Id": []string{"protocol-api-key-sse-nonstream"}, + }, + Body: io.NopCloser(strings.NewReader(grokProtocolUpstreamSSE())), + }} + svc := &OpenAIGatewayService{httpUpstream: upstream} + account := grokProtocolAPIKeyAccount(7104) + + result, err := svc.forwardGrokResponses(context.Background(), c, account, body, "grok", false, time.Now()) + + require.NoError(t, err) + require.NotNil(t, result) + require.False(t, result.Stream) + require.Equal(t, "resp_protocol_stream", result.ResponseID) + assertGrokProtocolRequestLowered(t, upstream.lastBody) + + response := recorder.Body.Bytes() + require.True(t, json.Valid(response)) + require.Equal(t, "custom_tool_call", gjson.GetBytes(response, "output.0.type").String()) + require.Equal(t, "*** Begin Patch", gjson.GetBytes(response, "output.0.input").String()) + require.Equal(t, "tool_search_call", gjson.GetBytes(response, "output.1.type").String()) + require.Equal(t, "client", gjson.GetBytes(response, "output.1.execution").String()) + require.Equal(t, "collaboration", gjson.GetBytes(response, "output.2.namespace").String()) + require.Equal(t, "send_message", gjson.GetBytes(response, "output.2.name").String()) +} + +func TestForwardGrokResponsesAPIKeyRestoresClientToolsStreaming(t *testing.T) { + gin.SetMode(gin.TestMode) + + body := grokClientToolProtocolRequest(true) + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", bytes.NewReader(body)) + c.Request.Header.Set("Content-Type", "application/json") + + upstream := &httpUpstreamRecorder{resp: &http.Response{ + StatusCode: http.StatusOK, + Header: http.Header{ + "Content-Type": []string{"text/event-stream"}, + "Xai-Request-Id": []string{"protocol-api-key"}, + }, + Body: io.NopCloser(strings.NewReader(grokProtocolUpstreamSSE())), + }} + svc := &OpenAIGatewayService{httpUpstream: upstream} + account := grokProtocolAPIKeyAccount(7103) + + result, err := svc.forwardGrokResponses(context.Background(), c, account, body, "grok", true, time.Now()) + + require.NoError(t, err) + require.NotNil(t, result) + require.True(t, result.Stream) + require.Equal(t, "resp_protocol_stream", result.ResponseID) + require.Equal(t, "https://api.x.ai/v1/responses", upstream.lastReq.URL.String()) + require.Equal(t, "Bearer xai-protocol-key", upstream.lastReq.Header.Get("Authorization")) + assertGrokProtocolRequestLowered(t, upstream.lastBody) + + frames := parseGrokProtocolSSEFrames(t, recorder.Body.String()) + require.NotEmpty(t, frames) + for index, frame := range frames { + require.Equal(t, frame.event, gjson.GetBytes(frame.data, "type").String(), "SSE event field must follow the restored data.type") + require.Equal(t, 40+index, int(gjson.GetBytes(frame.data, "sequence_number").Int()), "sequence_number must be continuous after suppressed and expanded events") + } + + created := requireGrokProtocolFrame(t, frames, "response.created", "", "") + require.True(t, gjson.GetBytes(created.data, "upstream_extension.preserved").Bool()) + customAdded := requireGrokProtocolFrame(t, frames, "response.output_item.added", "item.type", "custom_tool_call") + require.Equal(t, "apply_patch", gjson.GetBytes(customAdded.data, "item.name").String()) + customInputDelta := requireGrokProtocolFrame(t, frames, "response.custom_tool_call_input.delta", "", "") + require.Equal(t, "*** Begin Patch", gjson.GetBytes(customInputDelta.data, "delta").String()) + customInputDone := requireGrokProtocolFrame(t, frames, "response.custom_tool_call_input.done", "", "") + require.Equal(t, "*** Begin Patch", gjson.GetBytes(customInputDone.data, "input").String()) + customDone := requireGrokProtocolFrame(t, frames, "response.output_item.done", "item.type", "custom_tool_call") + require.Equal(t, "*** Begin Patch", gjson.GetBytes(customDone.data, "item.input").String()) + + namespaceAdded := requireGrokProtocolFrame(t, frames, "response.output_item.added", "item.namespace", "collaboration") + require.Equal(t, "send_message", gjson.GetBytes(namespaceAdded.data, "item.name").String()) + namespaceDone := requireGrokProtocolFrame(t, frames, "response.output_item.done", "item.namespace", "collaboration") + require.Equal(t, "send_message", gjson.GetBytes(namespaceDone.data, "item.name").String()) + namespaceArgumentsDone := requireGrokProtocolFrame(t, frames, "response.function_call_arguments.done", "name", "send_message") + require.Equal(t, "response.function_call_arguments.done", gjson.GetBytes(namespaceArgumentsDone.data, "type").String()) + require.False(t, gjson.GetBytes(namespaceArgumentsDone.data, "namespace").Exists()) + + searchAdded := requireGrokProtocolFrame(t, frames, "response.output_item.added", "item.type", "tool_search_call") + require.Equal(t, "client", gjson.GetBytes(searchAdded.data, "item.execution").String()) + searchDone := requireGrokProtocolFrame(t, frames, "response.output_item.done", "item.type", "tool_search_call") + require.Equal(t, "github", gjson.GetBytes(searchDone.data, "item.arguments.query").String()) + + for _, frame := range frames { + itemID := gjson.GetBytes(frame.data, "item_id").String() + if itemID == "item_custom" || itemID == "item_search" { + require.NotContains(t, frame.event, "function_call_arguments", "client-only proxy argument events must not leak") + } + } + completed := requireGrokProtocolFrame(t, frames, "response.completed", "", "") + require.Equal(t, "custom_tool_call", gjson.GetBytes(completed.data, "response.output.0.type").String()) + require.Equal(t, "tool_search_call", gjson.GetBytes(completed.data, "response.output.1.type").String()) + require.Equal(t, "collaboration", gjson.GetBytes(completed.data, "response.output.2.namespace").String()) +} + +func TestGrokResponsesClientToolStreamBodyFlushesFrameBeforeEOF(t *testing.T) { + sourceReader, sourceWriter := io.Pipe() + body := newGrokResponsesClientToolStreamBody(sourceReader, apicompat.ResponsesClientToolMapping{ + CustomTools: map[string]bool{"apply_patch": true}, + }, defaultMaxLineSize) + defer func() { _ = body.Close() }() + defer func() { _ = sourceWriter.Close() }() + + type readResult struct { + frame string + err error + } + read := make(chan readResult, 1) + go func() { + reader := bufio.NewReader(body) + var frame strings.Builder + for { + line, err := reader.ReadString('\n') + if err != nil { + read <- readResult{err: err} + return + } + frame.WriteString(line) + if strings.TrimSpace(line) == "" { + read <- readResult{frame: frame.String()} + return + } + } + }() + + firstFrame := "event: response.created\n" + + `data: {"type":"response.created","sequence_number":0,"response":{"id":"flush-before-eof"}}` + "\n\n" + _, err := sourceWriter.Write([]byte(firstFrame)) + require.NoError(t, err) + + select { + case result := <-read: + require.NoError(t, result.err) + require.Contains(t, result.frame, "flush-before-eof") + require.Contains(t, result.frame, "event: response.created") + case <-time.After(3 * time.Second): + t.Fatal("first transformed SSE frame was not flushed while the upstream connection remained open") + } +} + +type grokProtocolSSEFrame struct { + event string + data []byte +} + +func grokClientToolProtocolRequest(stream bool) []byte { + return []byte(fmt.Sprintf(`{ + "model":"grok","stream":%t, + "tools":[ + {"type":"custom","name":"apply_patch","description":"apply a patch","format":{"type":"grammar","syntax":"lark","definition":"start: /.+/"}}, + {"type":"tool_search"}, + {"type":"namespace","name":"collaboration","tools":[{"type":"function","name":"send_message","description":"send a message","parameters":{"type":"object","properties":{"target":{"type":"string"}}}}]} + ], + "tool_choice":{"type":"custom","name":"apply_patch"}, + "input":[ + {"type":"custom_tool_call","id":"old_custom","call_id":"old_custom_call","name":"apply_patch","input":"*** Begin Patch"}, + {"type":"custom_tool_call_output","call_id":"old_custom_call","output":"Done!"}, + {"type":"tool_search_call","id":"old_search","call_id":"old_search_call","arguments":{"query":"github"},"execution":"client"}, + {"type":"tool_search_output","call_id":"old_search_call","output":{"groups":["github"]}}, + {"type":"function_call","id":"old_namespace","call_id":"old_namespace_call","namespace":"collaboration","name":"send_message","arguments":"{\"target\":\"root\"}"}, + {"type":"function_call_output","call_id":"old_namespace_call","output":"ok"}, + {"type":"message","role":"user","content":[{"type":"input_text","text":"continue"}]} + ] + }`, stream)) +} + +func grokProtocolOAuthAccount(id int64) *Account { + return &Account{ + ID: id, Name: "grok-oauth-protocol", Platform: PlatformGrok, Type: AccountTypeOAuth, + Status: StatusActive, Schedulable: true, Concurrency: 1, + Credentials: map[string]any{ + "access_token": "oauth-protocol-token", "refresh_token": "refresh-token", + "expires_at": time.Now().Add(2 * grokTokenRefreshSkew).UTC().Format(time.RFC3339), + "base_url": xai.DefaultCLIBaseURL, "subscription_tier": "supergrok", + }, + } +} + +func grokProtocolAPIKeyAccount(id int64) *Account { + return &Account{ + ID: id, Name: "grok-api-key-protocol", Platform: PlatformGrok, Type: AccountTypeAPIKey, + Status: StatusActive, Schedulable: true, Concurrency: 1, + Credentials: map[string]any{"api_key": "xai-protocol-key", "base_url": "https://api.x.ai/v1"}, + } +} + +func assertGrokProtocolRequestLowered(t *testing.T, body []byte) { + t.Helper() + require.True(t, json.Valid(body)) + require.False(t, gjson.GetBytes(body, `tools.#(type=="custom")`).Exists()) + require.False(t, gjson.GetBytes(body, `tools.#(type=="namespace")`).Exists()) + require.False(t, gjson.GetBytes(body, `tools.#(type=="tool_search")`).Exists()) + require.True(t, gjson.GetBytes(body, `tools.#(name=="apply_patch")`).Exists()) + require.True(t, gjson.GetBytes(body, `tools.#(name=="tool_search")`).Exists()) + require.True(t, gjson.GetBytes(body, `tools.#(name=="collaboration__send_message")`).Exists()) + require.Equal(t, "function", gjson.GetBytes(body, "tool_choice.type").String()) + require.Equal(t, "apply_patch", gjson.GetBytes(body, "tool_choice.name").String()) + require.Equal(t, "function_call", gjson.GetBytes(body, "input.0.type").String()) + require.Equal(t, "function_call_output", gjson.GetBytes(body, "input.1.type").String()) + require.Equal(t, "function_call", gjson.GetBytes(body, "input.2.type").String()) + require.Equal(t, "tool_search", gjson.GetBytes(body, "input.2.name").String()) + require.Equal(t, "function_call_output", gjson.GetBytes(body, "input.3.type").String()) + require.Equal(t, "collaboration__send_message", gjson.GetBytes(body, "input.4.name").String()) + require.False(t, gjson.GetBytes(body, "input.4.namespace").Exists()) +} + +func grokProtocolUpstreamSSE() string { + events := []string{ + `{"type":"response.created","sequence_number":40,"response":{"id":"resp_protocol_stream","model":"grok-4.5"},"upstream_extension":{"preserved":true}}`, + `{"type":"response.output_item.added","sequence_number":41,"output_index":0,"item":{"type":"function_call","id":"item_custom","call_id":"call_custom","name":"apply_patch","arguments":"","status":"in_progress"}}`, + `{"type":"response.function_call_arguments.delta","sequence_number":42,"output_index":0,"item_id":"item_custom","delta":"{\"input\":\"*** Begin"}`, + `{"type":"response.function_call_arguments.delta","sequence_number":43,"output_index":0,"item_id":"item_custom","delta":" Patch\"}"}`, + `{"type":"response.function_call_arguments.done","sequence_number":44,"output_index":0,"item_id":"item_custom","call_id":"call_custom","name":"apply_patch","arguments":"{\"input\":\"*** Begin Patch\"}"}`, + `{"type":"response.output_item.done","sequence_number":45,"output_index":0,"item":{"type":"function_call","id":"item_custom","call_id":"call_custom","name":"apply_patch","arguments":"{\"input\":\"*** Begin Patch\"}","status":"completed"}}`, + `{"type":"response.output_item.added","sequence_number":46,"output_index":1,"item":{"type":"function_call","id":"item_namespace","call_id":"call_namespace","name":"collaboration__send_message","arguments":"","status":"in_progress"}}`, + `{"type":"response.function_call_arguments.done","sequence_number":47,"output_index":1,"item_id":"item_namespace","call_id":"call_namespace","name":"collaboration__send_message","arguments":"{\"target\":\"root\"}"}`, + `{"type":"response.output_item.done","sequence_number":48,"output_index":1,"item":{"type":"function_call","id":"item_namespace","call_id":"call_namespace","name":"collaboration__send_message","arguments":"{\"target\":\"root\"}","status":"completed"}}`, + `{"type":"response.output_item.added","sequence_number":49,"output_index":2,"item":{"type":"function_call","id":"item_search","call_id":"call_search","name":"tool_search","arguments":"","status":"in_progress"}}`, + `{"type":"response.function_call_arguments.delta","sequence_number":50,"output_index":2,"item_id":"item_search","delta":"{\"query\":\"github\"}"}`, + `{"type":"response.function_call_arguments.done","sequence_number":51,"output_index":2,"item_id":"item_search","call_id":"call_search","name":"tool_search","arguments":"{\"query\":\"github\"}"}`, + `{"type":"response.output_item.done","sequence_number":52,"output_index":2,"item":{"type":"function_call","id":"item_search","call_id":"call_search","name":"tool_search","arguments":"{\"query\":\"github\"}","status":"completed"}}`, + `{"type":"response.completed","sequence_number":53,"response":{"id":"resp_protocol_stream","object":"response","model":"grok-4.5","status":"completed","output":[{"type":"function_call","id":"item_custom","call_id":"call_custom","name":"apply_patch","arguments":"{\"input\":\"*** Begin Patch\"}"},{"type":"function_call","id":"item_search","call_id":"call_search","name":"tool_search","arguments":"{\"query\":\"github\"}"},{"type":"function_call","id":"item_namespace","call_id":"call_namespace","name":"collaboration__send_message","arguments":"{\"target\":\"root\"}"}],"usage":{"input_tokens":11,"output_tokens":4,"total_tokens":15}}}`, + } + var out strings.Builder + for _, event := range events { + typ := gjson.Get(event, "type").String() + fmt.Fprintf(&out, "event: %s\ndata: %s\n\n", typ, event) + } + return out.String() +} + +func parseGrokProtocolSSEFrames(t *testing.T, body string) []grokProtocolSSEFrame { + t.Helper() + var frames []grokProtocolSSEFrame + event := "" + for _, rawLine := range strings.Split(body, "\n") { + line := strings.TrimSuffix(rawLine, "\r") + if value, ok := extractOpenAISSEEventLine(line); ok { + event = strings.TrimSpace(value) + continue + } + data, ok := extractOpenAISSEDataLine(line) + if !ok || strings.TrimSpace(data) == "[DONE]" { + continue + } + require.NotEmpty(t, event, "every data frame from this upstream should retain an event field") + require.JSONEq(t, data, data) + frames = append(frames, grokProtocolSSEFrame{event: event, data: []byte(data)}) + event = "" + } + return frames +} + +func requireGrokProtocolFrame(t *testing.T, frames []grokProtocolSSEFrame, eventType, path, value string) grokProtocolSSEFrame { + t.Helper() + for _, frame := range frames { + if frame.event != eventType { + continue + } + if path == "" || gjson.GetBytes(frame.data, path).String() == value { + return frame + } + } + t.Fatalf("missing SSE frame event=%q %s=%q", eventType, path, value) + return grokProtocolSSEFrame{} +} diff --git a/backend/internal/service/openai_gateway_response_handling.go b/backend/internal/service/openai_gateway_response_handling.go index 20a230d1b..12923274c 100644 --- a/backend/internal/service/openai_gateway_response_handling.go +++ b/backend/internal/service/openai_gateway_response_handling.go @@ -477,7 +477,12 @@ func (s *OpenAIGatewayService) handleStreamingResponseWithReasoning(ctx context. line = "data: " + data eventType = strings.TrimSpace(gjson.GetBytes(dataBytes, "type").String()) } - restoredData, restoreErr := restoreOpenAIResponsesNamespacePayload(c, dataBytes) + restoredData, restoreErr := restoreGrokResponsesClientToolPayload(c, dataBytes) + if restoreErr != nil { + streamEarlyErr = fmt.Errorf("restore Grok Responses client tool response: %w", restoreErr) + return + } + restoredData, restoreErr = restoreOpenAIResponsesNamespacePayload(c, restoredData) if restoreErr != nil { streamEarlyErr = fmt.Errorf("restore OpenAI namespace response: %w", restoreErr) return @@ -1140,6 +1145,10 @@ func (s *OpenAIGatewayService) handleNonStreamingResponse(ctx context.Context, r if originalModel != mappedModel { body = s.replaceModelInResponseBody(body, mappedModel, originalModel) } + body, err = restoreGrokResponsesClientToolPayload(c, body) + if err != nil { + return nil, fmt.Errorf("restore Grok Responses client tool response: %w", err) + } body, err = restoreOpenAIResponsesNamespacePayload(c, body) if err != nil { return nil, fmt.Errorf("restore OpenAI namespace response: %w", err) @@ -1213,7 +1222,11 @@ func (s *OpenAIGatewayService) handleSSEToJSON(resp *http.Response, c *gin.Conte } // Correct tool calls in final response body = s.correctToolCallsInResponseBody(body) - restoredBody, restoreErr := restoreOpenAIResponsesNamespacePayload(c, body) + restoredBody, restoreErr := restoreGrokResponsesClientToolPayload(c, body) + if restoreErr != nil { + return nil, fmt.Errorf("restore Grok Responses client tool response: %w", restoreErr) + } + restoredBody, restoreErr = restoreOpenAIResponsesNamespacePayload(c, restoredBody) if restoreErr != nil { return nil, fmt.Errorf("restore OpenAI namespace response: %w", restoreErr) } diff --git a/backend/internal/service/openai_responses_namespace.go b/backend/internal/service/openai_responses_namespace.go index 7d71b6481..4c3f433bc 100644 --- a/backend/internal/service/openai_responses_namespace.go +++ b/backend/internal/service/openai_responses_namespace.go @@ -17,7 +17,7 @@ const openAIResponsesNamespaceNamesContextKey = "openai_responses_namespace_name // 平名无法还原会破坏客户端工具匹配,因此实际走 WSv2 分支的请求保持 namespace // 原样。透传账号先于 WSv2 分支经 HTTP 转发返回,仍需摊平。 func shouldFlattenOpenAIResponsesNamespaces(account *Account, transport OpenAIUpstreamTransport, passthroughEnabled bool) bool { - if account == nil || account.Type != AccountTypeOAuth { + if account == nil || !account.IsOpenAIOAuth() { return false } if transport == OpenAIUpstreamTransportResponsesWebsocketV2 && !passthroughEnabled { diff --git a/backend/internal/service/openai_responses_namespace_test.go b/backend/internal/service/openai_responses_namespace_test.go index 7aa139226..00b8772d1 100644 --- a/backend/internal/service/openai_responses_namespace_test.go +++ b/backend/internal/service/openai_responses_namespace_test.go @@ -7,8 +7,9 @@ import ( ) func TestShouldFlattenOpenAIResponsesNamespaces(t *testing.T) { - oauth := &Account{Type: AccountTypeOAuth} - apiKey := &Account{Type: AccountTypeAPIKey} + oauth := &Account{Platform: PlatformOpenAI, Type: AccountTypeOAuth} + apiKey := &Account{Platform: PlatformOpenAI, Type: AccountTypeAPIKey} + grokOAuth := &Account{Platform: PlatformGrok, Type: AccountTypeOAuth} tests := []struct { name string @@ -24,6 +25,7 @@ func TestShouldFlattenOpenAIResponsesNamespaces(t *testing.T) { // 透传账号先于 WSv2 分支经 HTTP 转发返回,仍需摊平。 {name: "oauth_wsv2_passthrough", account: oauth, transport: OpenAIUpstreamTransportResponsesWebsocketV2, passthroughEnabled: true, want: true}, {name: "apikey_http", account: apiKey, transport: OpenAIUpstreamTransportHTTPSSE, want: false}, + {name: "grok_oauth_http", account: grokOAuth, transport: OpenAIUpstreamTransportHTTPSSE, want: false}, {name: "nil_account", account: nil, transport: OpenAIUpstreamTransportHTTPSSE, want: false}, } for _, tt := range tests {