diff --git a/internal/runtime/executor/xai_executor.go b/internal/runtime/executor/xai_executor.go index 3df00adf..f52be84d 100644 --- a/internal/runtime/executor/xai_executor.go +++ b/internal/runtime/executor/xai_executor.go @@ -188,6 +188,7 @@ func (e *XAIExecutor) Execute(ctx context.Context, auth *cliproxyauth.Auth, req continue } eventData := xaiNormalizeReasoningSummaryData(bytes.TrimSpace(line[len(xaiDataTag):])) + eventData = restoreXAINamespaceToolCalls(eventData, prepared.namespaceTools) switch gjson.GetBytes(eventData, "type").String() { case "response.output_item.done": xaiCollectOutputItemDone(eventData, outputItemsByIndex, &outputItemsFallback) @@ -672,6 +673,7 @@ func (e *XAIExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth.Auth eventDataList := xaiNormalizeReasoningSummaryDataEvents(bytes.TrimSpace(line[len(xaiDataTag):])) hasPendingEventLine := pendingEventLine != nil for i, eventData := range eventDataList { + eventData = restoreXAINamespaceToolCalls(eventData, prepared.namespaceTools) normalizedEventName := gjson.GetBytes(eventData, "type").String() switch normalizedEventName { case "response.output_item.done": @@ -817,10 +819,16 @@ type xaiPreparedRequest struct { to sdktranslator.Format originalPayload []byte body []byte + namespaceTools map[string]xaiNamespaceToolRef sessionID string replayScope xaiReasoningReplayScope } +type xaiNamespaceToolRef struct { + namespace string + name string +} + func (e *XAIExecutor) prepareResponsesRequest(ctx context.Context, req cliproxyexecutor.Request, opts cliproxyexecutor.Options, stream bool) (*xaiPreparedRequest, error) { return e.prepareResponsesRequestTo(ctx, req, opts, stream, sdktranslator.FormatCodex) } @@ -852,13 +860,16 @@ func (e *XAIExecutor) prepareResponsesRequestTo(ctx context.Context, req cliprox body, _ = sjson.DeleteBytes(body, "prompt_cache_retention") body, _ = sjson.DeleteBytes(body, "safety_identifier") body, _ = sjson.DeleteBytes(body, "stream_options") + namespaceTools := collectXAINamespaceToolRefs(body) body = normalizeXAITools(body) + body = normalizeXAINamespaceToolChoice(body) body = normalizeXAIToolChoiceForTools(body) var replayScope xaiReasoningReplayScope body, replayScope, err = applyXAIReasoningReplayCacheRequired(ctx, from, req, opts, body) if err != nil { return nil, err } + body = normalizeXAIInputNamespaceToolCalls(body) body = normalizeXAIInputReasoningItems(body) body = sanitizeXAIInputEncryptedContent(body) body = normalizeCodexInstructions(body) @@ -879,6 +890,7 @@ func (e *XAIExecutor) prepareResponsesRequestTo(ctx context.Context, req cliprox to: to, originalPayload: originalPayload, body: body, + namespaceTools: namespaceTools, sessionID: sessionID, replayScope: replayScope, }, nil @@ -1158,11 +1170,48 @@ func sanitizeXAIResponsesBody(body []byte, model string) []byte { } func normalizeXAITools(body []byte) []byte { - tools := gjson.GetBytes(body, "tools") - if !tools.Exists() || !tools.IsArray() { + if !gjson.ValidBytes(body) { return body } + original := body + normalizeAtPath := func(path string) bool { + tools := gjson.GetBytes(body, path) + if !tools.Exists() || !tools.IsArray() { + return true + } + filtered, changed, ok := normalizeXAIToolArray(tools) + if !ok { + return false + } + if !changed { + return true + } + updated, errSet := sjson.SetRawBytes(body, path, filtered) + if errSet != nil { + return false + } + body = updated + return true + } + if !normalizeAtPath("tools") { + return original + } + input := gjson.GetBytes(body, "input") + if input.Exists() && input.IsArray() { + for index, item := range input.Array() { + if item.Get("type").String() != "additional_tools" { + continue + } + if !normalizeAtPath(fmt.Sprintf("input.%d.tools", index)) { + return original + } + } + } + return body +} + +func normalizeXAIToolArray(tools gjson.Result) ([]byte, bool, bool) { changed := false filtered := []byte(`[]`) for _, tool := range tools.Array() { @@ -1174,7 +1223,7 @@ func normalizeXAITools(body []byte) []byte { for _, nestedTool := range namespaceTools.Array() { nestedRaw, nestedChanged, ok := normalizeXAITool(nestedTool, namespaceName) if !ok { - return body + return nil, false, false } changed = changed || nestedChanged if len(nestedRaw) == 0 { @@ -1182,7 +1231,7 @@ func normalizeXAITools(body []byte) []byte { } updated, errSet := sjson.SetRawBytes(filtered, "-1", nestedRaw) if errSet != nil { - return body + return nil, false, false } filtered = updated } @@ -1191,7 +1240,7 @@ func normalizeXAITools(body []byte) []byte { } raw, toolChanged, ok := normalizeXAITool(tool, "") if !ok { - return body + return nil, false, false } changed = changed || toolChanged if len(raw) == 0 { @@ -1199,18 +1248,11 @@ func normalizeXAITools(body []byte) []byte { } updated, errSet := sjson.SetRawBytes(filtered, "-1", raw) if errSet != nil { - return body + return nil, false, false } filtered = updated } - if !changed { - return body - } - updated, errSet := sjson.SetRawBytes(body, "tools", filtered) - if errSet != nil { - return body - } - return updated + return filtered, changed, true } // normalizeXAIToolChoiceForTools drops tool_choice and parallel_tool_calls @@ -1220,6 +1262,18 @@ func normalizeXAITools(body []byte) []byte { func normalizeXAIToolChoiceForTools(body []byte) []byte { tools := gjson.GetBytes(body, "tools") hasTools := tools.Exists() && tools.IsArray() && len(tools.Array()) > 0 + if !hasTools { + input := gjson.GetBytes(body, "input") + if input.Exists() && input.IsArray() { + for _, item := range input.Array() { + additionalTools := item.Get("tools") + if item.Get("type").String() == "additional_tools" && additionalTools.IsArray() && len(additionalTools.Array()) > 0 { + hasTools = true + break + } + } + } + } if hasTools { return body } @@ -1235,6 +1289,34 @@ func normalizeXAIToolChoiceForTools(body []byte) []byte { return body } +// normalizeXAINamespaceToolChoice qualifies a forced namespaced function choice +// using the same name sent in the flattened tools list. xAI does not accept the +// Responses namespace field on tool_choice. +func normalizeXAINamespaceToolChoice(body []byte) []byte { + if !gjson.ValidBytes(body) { + return body + } + toolChoice := gjson.GetBytes(body, "tool_choice") + if !toolChoice.IsObject() || toolChoice.Get("type").String() != xaiFunctionToolType { + return body + } + namespaceName := strings.TrimSpace(toolChoice.Get("namespace").String()) + toolName := strings.TrimSpace(toolChoice.Get("name").String()) + qualifiedName := qualifyXAINamespaceToolName(namespaceName, toolName) + if namespaceName == "" || qualifiedName == "" { + return body + } + updated, errSet := sjson.SetBytes(body, "tool_choice.name", qualifiedName) + if errSet != nil { + return body + } + updated, errDelete := sjson.DeleteBytes(updated, "tool_choice.namespace") + if errDelete != nil { + return body + } + return updated +} + func normalizeXAITool(tool gjson.Result, namespaceName string) ([]byte, bool, bool) { toolType := tool.Get("type").String() changed := false @@ -1289,9 +1371,140 @@ func normalizeXAITool(tool gjson.Result, namespaceName string) ([]byte, bool, bo changed = true log.Debugf("xai: simplified parameters for tool %s.%s to avoid upstream hang", namespaceName, tool.Get("name").String()) } + if toolType == xaiFunctionToolType && strings.TrimSpace(namespaceName) != "" { + qualifiedName := qualifyXAINamespaceToolName(namespaceName, tool.Get("name").String()) + if qualifiedName == "" { + return nil, false, false + } + updatedTool, errSet := sjson.SetBytes(raw, "name", qualifiedName) + if errSet != nil { + return nil, false, false + } + raw = updatedTool + changed = true + } return raw, changed, true } +func qualifyXAINamespaceToolName(namespaceName, toolName string) string { + namespaceName = strings.TrimSpace(namespaceName) + toolName = strings.TrimSpace(toolName) + if namespaceName == "" || toolName == "" || strings.HasPrefix(toolName, "mcp__") { + return toolName + } + prefix := namespaceName + if !strings.HasSuffix(prefix, "__") { + prefix += "__" + } + if strings.HasPrefix(toolName, prefix) { + return toolName + } + return prefix + toolName +} + +func collectXAINamespaceToolRefs(body []byte) map[string]xaiNamespaceToolRef { + refs := make(map[string]xaiNamespaceToolRef) + collect := func(tools gjson.Result) { + if !tools.Exists() || !tools.IsArray() { + return + } + for _, tool := range tools.Array() { + if tool.Get("type").String() != xaiNamespaceToolType { + continue + } + namespaceName := strings.TrimSpace(tool.Get("name").String()) + if namespaceName == "" { + continue + } + for _, nestedTool := range tool.Get("tools").Array() { + toolName := strings.TrimSpace(nestedTool.Get("name").String()) + qualifiedName := qualifyXAINamespaceToolName(namespaceName, toolName) + if qualifiedName == "" { + continue + } + refs[qualifiedName] = xaiNamespaceToolRef{namespace: namespaceName, name: toolName} + } + } + } + collect(gjson.GetBytes(body, "tools")) + input := gjson.GetBytes(body, "input") + if input.Exists() && input.IsArray() { + for _, item := range input.Array() { + if item.Get("type").String() == "additional_tools" { + collect(item.Get("tools")) + } + } + } + return refs +} + +func normalizeXAIInputNamespaceToolCalls(body []byte) []byte { + if !gjson.ValidBytes(body) { + return body + } + input := gjson.GetBytes(body, "input") + if !input.Exists() || !input.IsArray() { + return body + } + for index, item := range input.Array() { + if item.Get("type").String() != "function_call" { + continue + } + namespaceName := strings.TrimSpace(item.Get("namespace").String()) + toolName := strings.TrimSpace(item.Get("name").String()) + qualifiedName := qualifyXAINamespaceToolName(namespaceName, toolName) + if namespaceName == "" || qualifiedName == "" { + continue + } + namePath := fmt.Sprintf("input.%d.name", index) + namespacePath := fmt.Sprintf("input.%d.namespace", index) + updated, errSet := sjson.SetBytes(body, namePath, qualifiedName) + if errSet != nil { + continue + } + updated, errDelete := sjson.DeleteBytes(updated, namespacePath) + if errDelete != nil { + continue + } + body = updated + } + return body +} + +func restoreXAINamespaceToolCalls(data []byte, refs map[string]xaiNamespaceToolRef) []byte { + if len(refs) == 0 || len(data) == 0 || !gjson.ValidBytes(data) { + return data + } + data = restoreXAINamespaceToolCallAtPath(data, "item", refs) + output := gjson.GetBytes(data, "response.output") + if output.Exists() && output.IsArray() { + for index := range output.Array() { + data = restoreXAINamespaceToolCallAtPath(data, fmt.Sprintf("response.output.%d", index), refs) + } + } + return data +} + +func restoreXAINamespaceToolCallAtPath(data []byte, path string, refs map[string]xaiNamespaceToolRef) []byte { + if gjson.GetBytes(data, path+".type").String() != "function_call" { + return data + } + qualifiedName := strings.TrimSpace(gjson.GetBytes(data, path+".name").String()) + ref, ok := refs[qualifiedName] + if !ok { + return data + } + updated, errSet := sjson.SetBytes(data, path+".name", ref.name) + if errSet != nil { + return data + } + updated, errSet = sjson.SetBytes(updated, path+".namespace", ref.namespace) + if errSet != nil { + return data + } + return updated +} + // xaiFunctionParametersNeedSimplification reports whether a function tool is // the Codex Desktop automation tool known to hang xAI Responses streaming. func xaiFunctionParametersNeedSimplification(tool gjson.Result, namespaceName string) bool { diff --git a/internal/runtime/executor/xai_executor_test.go b/internal/runtime/executor/xai_executor_test.go index 99b4610a..ddea75e6 100644 --- a/internal/runtime/executor/xai_executor_test.go +++ b/internal/runtime/executor/xai_executor_test.go @@ -148,9 +148,9 @@ func TestXAIExecutorExecuteShapesResponsesRequest(t *testing.T) { t.Fatalf("tools.%d.name = apply_patch, want removed; body=%s", i, string(gotBody)) } switch tool.Get("name").String() { - case "automation_update": + case "codex_app__automation_update": foundAutomationUpdate = true - case "namespace_custom": + case "codex_app__namespace_custom": foundNamespaceCustom = true } if toolType == "web_search" { @@ -180,6 +180,63 @@ func TestXAIExecutorExecuteShapesResponsesRequest(t *testing.T) { } } +func TestXAIExecutorExecuteRestoresAdditionalToolsNamespaceCalls(t *testing.T) { + var gotBody []byte + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + var errRead error + gotBody, errRead = io.ReadAll(r.Body) + if errRead != nil { + t.Fatalf("read body: %v", errRead) + } + w.Header().Set("Content-Type", "text/event-stream") + _, _ = w.Write([]byte("data: {\"type\":\"response.output_item.done\",\"output_index\":0,\"item\":{\"type\":\"function_call\",\"name\":\"mcp__exa__web_search_exa\",\"call_id\":\"call_1\",\"arguments\":\"{}\"}}\n\n")) + _, _ = w.Write([]byte("data: {\"type\":\"response.completed\",\"response\":{\"id\":\"resp_1\",\"object\":\"response\",\"status\":\"completed\",\"model\":\"grok-4.3\",\"output\":[],\"usage\":{\"input_tokens\":1,\"output_tokens\":1,\"total_tokens\":2}}}\n\n")) + })) + defer server.Close() + + exec := NewXAIExecutor(&config.Config{}) + auth := &cliproxyauth.Auth{ + Provider: "xai", + Attributes: map[string]string{"base_url": server.URL}, + Metadata: map[string]any{"access_token": "xai-token"}, + } + resp, err := exec.Execute(context.Background(), auth, cliproxyexecutor.Request{ + Model: "grok-4.3", + Payload: []byte(`{ + "model":"grok-4.3", + "input":[ + {"type":"additional_tools","role":"developer","tools":[{"type":"namespace","name":"mcp__exa","tools":[{"type":"function","name":"web_search_exa","parameters":{"type":"object"}}]}]}, + {"role":"user","content":"use Exa"} + ] + }`), + }, cliproxyexecutor.Options{ + SourceFormat: sdktranslator.FormatOpenAIResponse, + ResponseFormat: sdktranslator.FormatOpenAIResponse, + Stream: false, + }) + if err != nil { + t.Fatalf("Execute() error = %v", err) + } + + tool := gjson.GetBytes(gotBody, "input.0.tools.0") + if got := tool.Get("name").String(); got != "mcp__exa__web_search_exa" { + t.Fatalf("upstream additional tool name = %q, want qualified name; body=%s", got, gotBody) + } + if got := tool.Get("type").String(); got != "function" { + t.Fatalf("upstream additional tool type = %q, want function; body=%s", got, gotBody) + } + if tool.Get("tools").Exists() { + t.Fatalf("upstream additional tool should not contain namespace children: %s", gotBody) + } + output := gjson.GetBytes(resp.Payload, "output.0") + if got := output.Get("name").String(); got != "web_search_exa" { + t.Fatalf("response output name = %q, want child name; payload=%s", got, resp.Payload) + } + if got := output.Get("namespace").String(); got != "mcp__exa" { + t.Fatalf("response output namespace = %q, want mcp__exa; payload=%s", got, resp.Payload) + } +} + func TestXAIExecutorComposerSessionIsolation(t *testing.T) { exec := NewXAIExecutor(&config.Config{}) auth := &cliproxyauth.Auth{ @@ -786,9 +843,9 @@ func TestXAIExecutorExecuteStreamFiltersToolSearchTool(t *testing.T) { t.Fatalf("tools.%d.name = apply_patch, want removed; body=%s", i, string(gotBody)) } switch tool.Get("name").String() { - case "automation_update": + case "codex_app__automation_update": foundAutomationUpdate = true - case "namespace_custom": + case "codex_app__namespace_custom": foundNamespaceCustom = true } if toolType == "web_search" { @@ -1203,7 +1260,7 @@ func TestNormalizeXAITools_SimplifiesCodexAppAutomationUpdateSchema(t *testing.T foundExec := false for _, tool := range tools.Array() { switch tool.Get("name").String() { - case "automation_update": + case "codex_app__automation_update": foundAuto = true paramsRaw := tool.Get("parameters").Raw if strings.Contains(paramsRaw, `"oneOf"`) || strings.Contains(paramsRaw, `"$defs"`) { @@ -1233,6 +1290,108 @@ func TestNormalizeXAITools_SimplifiesCodexAppAutomationUpdateSchema(t *testing.T } } +func TestNormalizeXAITools_QualifiesSameNamedNamespaceTools(t *testing.T) { + body := []byte(`{ + "tools":[ + {"type":"namespace","name":"mcp__exa","tools":[{"type":"function","name":"search","parameters":{"type":"object"}}]}, + {"type":"namespace","name":"mcp__docs","tools":[{"type":"function","name":"search","parameters":{"type":"object"}}]} + ] + }`) + out := normalizeXAITools(body) + + tools := gjson.GetBytes(out, "tools").Array() + if len(tools) != 2 { + t.Fatalf("tools length = %d, want 2; body=%s", len(tools), string(out)) + } + if got := tools[0].Get("name").String(); got != "mcp__exa__search" { + t.Fatalf("tools.0.name = %q, want mcp__exa__search; body=%s", got, string(out)) + } + if got := tools[1].Get("name").String(); got != "mcp__docs__search" { + t.Fatalf("tools.1.name = %q, want mcp__docs__search; body=%s", got, string(out)) + } +} + +func TestNormalizeXAITools_AdditionalToolsNamespace(t *testing.T) { + body := []byte(`{ + "input":[ + {"type":"additional_tools","role":"developer","tools":[{"type":"namespace","name":"mcp__exa","tools":[{"type":"function","name":"search","parameters":{"type":"object"}}]}]}, + {"role":"user","content":"hello"} + ] + }`) + out := normalizeXAITools(body) + + tools := gjson.GetBytes(out, "input.0.tools").Array() + if len(tools) != 1 { + t.Fatalf("additional tools length = %d, want 1; body=%s", len(tools), string(out)) + } + if got := tools[0].Get("name").String(); got != "mcp__exa__search" { + t.Fatalf("additional tool name = %q, want mcp__exa__search; body=%s", got, string(out)) + } + if got := tools[0].Get("type").String(); got != "function" { + t.Fatalf("additional tool type = %q, want function; body=%s", got, string(out)) + } +} + +func TestNormalizeXAINamespaceToolChoice(t *testing.T) { + body := []byte(`{ + "tools":[{"type":"namespace","name":"mcp__exa","tools":[{"type":"function","name":"search","parameters":{"type":"object"}}]}], + "tool_choice":{"type":"function","name":"search","namespace":"mcp__exa"} + }`) + out := normalizeXAITools(body) + out = normalizeXAINamespaceToolChoice(out) + + if got := gjson.GetBytes(out, "tools.0.name").String(); got != "mcp__exa__search" { + t.Fatalf("tools.0.name = %q, want mcp__exa__search; body=%s", got, string(out)) + } + if got := gjson.GetBytes(out, "tool_choice.name").String(); got != "mcp__exa__search" { + t.Fatalf("tool_choice.name = %q, want mcp__exa__search; body=%s", got, string(out)) + } + if gjson.GetBytes(out, "tool_choice.namespace").Exists() { + t.Fatalf("tool_choice.namespace should be removed for xAI upstream: %s", string(out)) + } +} + +func TestNormalizeXAINamespaceToolChoice_PreservesOtherChoices(t *testing.T) { + tests := []struct { + name string + body []byte + }{ + {name: "automatic choice", body: []byte(`{"tool_choice":"auto"}`)}, + {name: "top-level function", body: []byte(`{"tool_choice":{"type":"function","name":"search"}}`)}, + {name: "non-function choice", body: []byte(`{"tool_choice":{"type":"web_search","name":"search","namespace":"mcp__exa"}}`)}, + {name: "malformed payload", body: []byte(`{"tool_choice":{"type":"function"`)}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + if got := normalizeXAINamespaceToolChoice(tt.body); !bytes.Equal(got, tt.body) { + t.Fatalf("payload changed: got=%q want=%q", got, tt.body) + } + }) + } +} + +func TestQualifyXAINamespaceToolNamePreservesQualifiedNames(t *testing.T) { + tests := []struct { + name string + namespace string + tool string + want string + }{ + {name: "plain child", namespace: "mcp__exa", tool: "search", want: "mcp__exa__search"}, + {name: "prequalified MCP child", namespace: "mcp__exa", tool: "mcp__exa__search", want: "mcp__exa__search"}, + {name: "prequalified generic child", namespace: "collaboration", tool: "collaboration__send_message", want: "collaboration__send_message"}, + {name: "namespace with separator", namespace: "collaboration__", tool: "send_message", want: "collaboration__send_message"}, + {name: "partial prefix is not qualified", namespace: "exa", tool: "example_tool", want: "exa__example_tool"}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + if got := qualifyXAINamespaceToolName(tt.namespace, tt.tool); got != tt.want { + t.Fatalf("qualifyXAINamespaceToolName(%q, %q) = %q, want %q", tt.namespace, tt.tool, got, tt.want) + } + }) + } +} + func TestNormalizeXAITools_PreservesUnrelatedSchemas(t *testing.T) { largeParams := `{"oneOf":[{"type":"object","properties":{"mode":{"type":"string"}}}],"$defs":{"a":{"type":"string"}},"x":"` + strings.Repeat("y", 1600) + `"}` tests := []struct { @@ -1302,6 +1461,55 @@ func TestXAIFunctionParametersNeedSimplification(t *testing.T) { } } +func TestNormalizeXAIInputNamespaceToolCalls(t *testing.T) { + body := []byte(`{"input":[{"type":"function_call","name":"web_search_exa","namespace":"mcp__exa","call_id":"call_1","arguments":"{}"},{"type":"function_call","name":"plain_tool","call_id":"call_2","arguments":"{}"}]}`) + out := normalizeXAIInputNamespaceToolCalls(body) + + if got := gjson.GetBytes(out, "input.0.name").String(); got != "mcp__exa__web_search_exa" { + t.Fatalf("input.0.name = %q, want qualified namespace name; body=%s", got, string(out)) + } + if gjson.GetBytes(out, "input.0.namespace").Exists() { + t.Fatalf("input.0.namespace should be removed for xAI upstream: %s", string(out)) + } + if got := gjson.GetBytes(out, "input.1.name").String(); got != "plain_tool" { + t.Fatalf("plain function call name changed to %q", got) + } +} + +func TestRestoreXAINamespaceToolCalls(t *testing.T) { + request := []byte(`{"tools":[{"type":"namespace","name":"mcp__exa","tools":[{"type":"function","name":"web_search_exa","parameters":{"type":"object"}}]}]}`) + refs := collectXAINamespaceToolRefs(request) + + event := []byte(`{"type":"response.output_item.done","item":{"type":"function_call","name":"mcp__exa__web_search_exa","call_id":"call_1","arguments":"{}"}}`) + restoredEvent := restoreXAINamespaceToolCalls(event, refs) + if got := gjson.GetBytes(restoredEvent, "item.name").String(); got != "web_search_exa" { + t.Fatalf("item.name = %q, want child name; event=%s", got, string(restoredEvent)) + } + if got := gjson.GetBytes(restoredEvent, "item.namespace").String(); got != "mcp__exa" { + t.Fatalf("item.namespace = %q, want mcp__exa; event=%s", got, string(restoredEvent)) + } + + completed := []byte(`{"type":"response.completed","response":{"output":[{"type":"function_call","name":"mcp__exa__web_search_exa","call_id":"call_1","arguments":"{}"}]}}`) + restoredCompleted := restoreXAINamespaceToolCalls(completed, refs) + if got := gjson.GetBytes(restoredCompleted, "response.output.0.name").String(); got != "web_search_exa" { + t.Fatalf("response.output.0.name = %q, want child name; event=%s", got, string(restoredCompleted)) + } + if got := gjson.GetBytes(restoredCompleted, "response.output.0.namespace").String(); got != "mcp__exa" { + t.Fatalf("response.output.0.namespace = %q, want mcp__exa; event=%s", got, string(restoredCompleted)) + } +} + +func TestRestoreXAINamespaceToolCallsPreservesMalformedPayload(t *testing.T) { + data := []byte(`{"item":{"type":"function_call","name":"mcp__exa__web_search_exa"`) + refs := map[string]xaiNamespaceToolRef{ + "mcp__exa__web_search_exa": {namespace: "mcp__exa", name: "web_search_exa"}, + } + + if got := restoreXAINamespaceToolCalls(data, refs); !bytes.Equal(got, data) { + t.Fatalf("malformed payload changed: got=%q want=%q", got, data) + } +} + func TestNormalizeXAIToolChoiceForTools_DropsWhenToolsEmpty(t *testing.T) { body := []byte(`{"model":"grok-4","tools":[],"tool_choice":"auto","parallel_tool_calls":true,"input":"hi"}`) out := normalizeXAIToolChoiceForTools(body) @@ -1347,6 +1555,18 @@ func TestNormalizeXAIToolChoiceForTools_KeepsWhenToolsPresent(t *testing.T) { } } +func TestNormalizeXAIToolChoiceForTools_KeepsWhenAdditionalToolsPresent(t *testing.T) { + body := []byte(`{"model":"grok-4","input":[{"type":"additional_tools","tools":[{"type":"function","name":"Bash"}]}],"tool_choice":"auto","parallel_tool_calls":true}`) + out := normalizeXAIToolChoiceForTools(body) + + if got := gjson.GetBytes(out, "tool_choice").String(); got != "auto" { + t.Fatalf("tool_choice = %q, want auto: %s", got, string(out)) + } + if !gjson.GetBytes(out, "parallel_tool_calls").Bool() { + t.Fatalf("parallel_tool_calls should be kept: %s", string(out)) + } +} + func TestNormalizeXAIToolChoiceForTools_NoOpWhenBothAbsent(t *testing.T) { body := []byte(`{"model":"grok-4","input":"hi"}`) out := normalizeXAIToolChoiceForTools(body) diff --git a/internal/runtime/executor/xai_websockets_executor.go b/internal/runtime/executor/xai_websockets_executor.go index 287a7cd5..0a55f1c9 100644 --- a/internal/runtime/executor/xai_websockets_executor.go +++ b/internal/runtime/executor/xai_websockets_executor.go @@ -637,6 +637,7 @@ func (e *XAIWebsocketsExecutor) ExecuteStream(ctx context.Context, auth *cliprox } for _, payload := range xaiNormalizeReasoningSummaryDataEvents(payload) { + payload = restoreXAINamespaceToolCalls(payload, prepared.namespaceTools) eventType := gjson.GetBytes(payload, "type").String() isTerminalEvent := eventType == "response.completed" || eventType == "response.done" || eventType == "error" warmupCompletedPayload := []byte(nil) diff --git a/internal/runtime/executor/xai_websockets_executor_test.go b/internal/runtime/executor/xai_websockets_executor_test.go index d9de647e..9d8b900e 100644 --- a/internal/runtime/executor/xai_websockets_executor_test.go +++ b/internal/runtime/executor/xai_websockets_executor_test.go @@ -121,6 +121,112 @@ func TestXAIWebsocketsExecuteStreamSendsResponseCreateWithPreviousResponseID(t * } } +func TestXAIWebsocketsExecuteStreamRestoresNamespaceToolCalls(t *testing.T) { + upgrader := websocket.Upgrader{CheckOrigin: func(*http.Request) bool { return true }} + capturedPayload := make(chan []byte, 1) + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + conn, err := upgrader.Upgrade(w, r, nil) + if err != nil { + t.Errorf("upgrade websocket: %v", err) + return + } + defer func() { _ = conn.Close() }() + + _, payload, errRead := conn.ReadMessage() + if errRead != nil { + t.Errorf("read upstream websocket message: %v", errRead) + return + } + capturedPayload <- bytes.Clone(payload) + + events := [][]byte{ + []byte(`{"type":"response.output_item.done","output_index":0,"item":{"type":"function_call","name":"mcp__exa__web_search_exa","call_id":"call_1","arguments":"{}"}}`), + []byte(`{"type":"response.completed","response":{"id":"resp_1","output":[],"usage":{"input_tokens":1,"output_tokens":1,"total_tokens":2}}}`), + } + for _, event := range events { + if errWrite := conn.WriteMessage(websocket.TextMessage, event); errWrite != nil { + t.Errorf("write websocket event: %v", errWrite) + return + } + } + })) + defer server.Close() + + exec := NewXAIWebsocketsExecutor(&config.Config{}) + auth := &cliproxyauth.Auth{ + Provider: "xai", + Attributes: map[string]string{ + "base_url": server.URL, + "websockets": "true", + }, + Metadata: map[string]any{"access_token": "xai-token"}, + } + req := cliproxyexecutor.Request{ + Model: "grok-4.3", + Payload: []byte(`{ + "model":"grok-4.3", + "input":[ + {"type":"additional_tools","role":"developer","tools":[{ + "type":"namespace", + "name":"mcp__exa", + "tools":[{"type":"function","name":"web_search_exa","parameters":{"type":"object"}}] + }]}, + {"role":"user","content":"use Exa"} + ] + }`), + } + opts := cliproxyexecutor.Options{ + SourceFormat: sdktranslator.FormatOpenAIResponse, + ResponseFormat: sdktranslator.FormatOpenAIResponse, + Stream: true, + } + ctx := cliproxyexecutor.WithDownstreamWebsocket(context.Background()) + + result, err := exec.ExecuteStream(ctx, auth, req, opts) + if err != nil { + t.Fatalf("ExecuteStream() error = %v", err) + } + + select { + case payload := <-capturedPayload: + tool := gjson.GetBytes(payload, "input.0.tools.0") + if got := tool.Get("name").String(); got != "mcp__exa__web_search_exa" { + t.Fatalf("upstream tool name = %q, want qualified name; payload=%s", got, payload) + } + if tool.Get("tools").Exists() { + t.Fatalf("upstream tool should not contain namespace children: %s", payload) + } + case <-time.After(5 * time.Second): + t.Fatal("timed out waiting for upstream websocket payload") + } + + var outputItemDone, completed gjson.Result + for chunk := range result.Chunks { + if chunk.Err != nil { + t.Fatalf("stream chunk error = %v", chunk.Err) + } + payload := gjson.ParseBytes(bytes.TrimSpace(chunk.Payload)) + switch payload.Get("type").String() { + case "response.output_item.done": + outputItemDone = payload + case "response.completed": + completed = payload + } + } + + for label, item := range map[string]gjson.Result{ + "output_item.done": outputItemDone.Get("item"), + "completed": completed.Get("response.output.0"), + } { + if got := item.Get("name").String(); got != "web_search_exa" { + t.Fatalf("%s name = %q, want child name; item=%s", label, got, item.Raw) + } + if got := item.Get("namespace").String(); got != "mcp__exa" { + t.Fatalf("%s namespace = %q, want mcp__exa; item=%s", label, got, item.Raw) + } + } +} + func TestXAIWebsocketsExecuteStreamNormalizesReasoningTextEvents(t *testing.T) { upgrader := websocket.Upgrader{CheckOrigin: func(*http.Request) bool { return true }} server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {