From 0e3a3e6199a3bda641e78d9a05df58454b6aca9d Mon Sep 17 00:00:00 2001 From: huahai0202 <2023217276@mail.hfut.edu.cn> Date: Sun, 12 Jul 2026 22:37:03 +0800 Subject: [PATCH 1/4] fix(xai): preserve namespace tool routing xAI Responses does not accept namespace tool containers, so the executor flattens nested tools before forwarding them upstream. The original namespace support in 8b3670b8 moved children to the top level but discarded their namespace identity; ca67caf later retained namespace context only for targeted schema simplification. Qualify flattened names to avoid collisions, normalize namespaced function-call history, and restore name plus namespace on HTTP and WebSocket responses. This keeps Codex MCP routing intact without changing the existing xAI tool_search filtering policy. --- internal/runtime/executor/xai_executor.go | 117 ++++++++++++++++++ .../runtime/executor/xai_executor_test.go | 90 +++++++++++++- .../executor/xai_websockets_executor.go | 2 + .../executor/xai_websockets_executor_test.go | 104 ++++++++++++++++ 4 files changed, 308 insertions(+), 5 deletions(-) diff --git a/internal/runtime/executor/xai_executor.go b/internal/runtime/executor/xai_executor.go index 20faaea7..5cc56735 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) @@ -196,6 +197,7 @@ func (e *XAIExecutor) Execute(ctx context.Context, auth *cliproxyauth.Auth, req reporter.Publish(ctx, detail) } completedData := xaiPatchCompletedOutput(eventData, outputItemsByIndex, outputItemsFallback) + completedData = restoreXAINamespaceToolCalls(completedData, prepared.namespaceTools) completedData = xaiNormalizeReasoningSummaryData(completedData) cacheXAIReasoningReplayFromCompleted(ctx, prepared.replayScope, completedData) var param any @@ -671,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": @@ -680,6 +683,7 @@ func (e *XAIExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth.Auth reporter.Publish(ctx, detail) } eventData = xaiPatchCompletedOutput(eventData, outputItemsByIndex, outputItemsFallback) + eventData = restoreXAINamespaceToolCalls(eventData, prepared.namespaceTools) eventData = xaiNormalizeReasoningSummaryData(eventData) cacheXAIReasoningReplayFromCompleted(ctx, prepared.replayScope, eventData) normalizedEventName = gjson.GetBytes(eventData, "type").String() @@ -816,10 +820,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) } @@ -851,6 +861,8 @@ 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 = normalizeXAIInputNamespaceToolCalls(body) body = normalizeXAITools(body) body = normalizeXAIToolChoiceForTools(body) var replayScope xaiReasoningReplayScope @@ -878,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 @@ -1289,9 +1302,113 @@ 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 + } + if strings.HasPrefix(toolName, namespaceName) { + return toolName + } + if strings.HasSuffix(namespaceName, "__") { + return namespaceName + toolName + } + return namespaceName + "__" + toolName +} + +func collectXAINamespaceToolRefs(body []byte) map[string]xaiNamespaceToolRef { + refs := make(map[string]xaiNamespaceToolRef) + tools := gjson.GetBytes(body, "tools") + if !tools.Exists() || !tools.IsArray() { + return refs + } + 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} + } + } + return refs +} + +func normalizeXAIInputNamespaceToolCalls(body []byte) []byte { + 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) + body, _ = sjson.SetBytes(body, namePath, qualifiedName) + body, _ = sjson.DeleteBytes(body, namespacePath) + } + return body +} + +func restoreXAINamespaceToolCalls(data []byte, refs map[string]xaiNamespaceToolRef) []byte { + if len(refs) == 0 || len(data) == 0 { + 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 + } + data, _ = sjson.SetBytes(data, path+".name", ref.name) + data, _ = sjson.SetBytes(data, path+".namespace", ref.namespace) + return data +} + // 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 d0183460..5429a8a8 100644 --- a/internal/runtime/executor/xai_executor_test.go +++ b/internal/runtime/executor/xai_executor_test.go @@ -139,9 +139,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" { @@ -668,9 +668,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" { @@ -1085,7 +1085,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"`) { @@ -1115,6 +1115,48 @@ 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 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"}, + } + 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 { @@ -1184,6 +1226,44 @@ 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 TestNormalizeXAIToolChoiceForTools_DropsWhenToolsEmpty(t *testing.T) { body := []byte(`{"model":"grok-4","tools":[],"tool_choice":"auto","parallel_tool_calls":true,"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..adbe2b2e 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) @@ -654,6 +655,7 @@ func (e *XAIWebsocketsExecutor) ExecuteStream(ctx context.Context, auth *cliprox reporter.Publish(ctx, detail) } payload = xaiPatchCompletedOutput(payload, outputItemsByIndex, outputItemsFallback) + payload = restoreXAINamespaceToolCalls(payload, prepared.namespaceTools) payload = xaiNormalizeReasoningSummaryData(payload) cacheXAIReasoningReplayFromCompleted(ctx, prepared.replayScope, payload) if !warmupRequest && idMapper != nil && idMapper.state != nil && !recordedTranscript { diff --git a/internal/runtime/executor/xai_websockets_executor_test.go b/internal/runtime/executor/xai_websockets_executor_test.go index d9de647e..a62f70a0 100644 --- a/internal/runtime/executor/xai_websockets_executor_test.go +++ b/internal/runtime/executor/xai_websockets_executor_test.go @@ -121,6 +121,110 @@ 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":"use Exa", + "tools":[{ + "type":"namespace", + "name":"mcp__exa", + "tools":[{"type":"function","name":"web_search_exa","parameters":{"type":"object"}}] + }] + }`), + } + 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, "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("namespace").Exists() { + t.Fatalf("upstream tool should not contain namespace: %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) { From c4eda81b93126cc1cec000e7855507348ef0614f Mon Sep 17 00:00:00 2001 From: huahai0202 <2023217276@mail.hfut.edu.cn> Date: Sun, 12 Jul 2026 23:33:53 +0800 Subject: [PATCH 2/4] fix(xai): address namespace routing review feedback Use exact namespace separator matching, normalize replay-inserted calls, preserve payloads on JSON update failures, and remove redundant completed-response restoration. --- internal/runtime/executor/xai_executor.go | 43 +++++++++++++------ .../runtime/executor/xai_executor_test.go | 12 ++++++ .../executor/xai_websockets_executor.go | 1 - 3 files changed, 41 insertions(+), 15 deletions(-) diff --git a/internal/runtime/executor/xai_executor.go b/internal/runtime/executor/xai_executor.go index 5cc56735..c52d3f5f 100644 --- a/internal/runtime/executor/xai_executor.go +++ b/internal/runtime/executor/xai_executor.go @@ -197,7 +197,6 @@ func (e *XAIExecutor) Execute(ctx context.Context, auth *cliproxyauth.Auth, req reporter.Publish(ctx, detail) } completedData := xaiPatchCompletedOutput(eventData, outputItemsByIndex, outputItemsFallback) - completedData = restoreXAINamespaceToolCalls(completedData, prepared.namespaceTools) completedData = xaiNormalizeReasoningSummaryData(completedData) cacheXAIReasoningReplayFromCompleted(ctx, prepared.replayScope, completedData) var param any @@ -683,7 +682,6 @@ func (e *XAIExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth.Auth reporter.Publish(ctx, detail) } eventData = xaiPatchCompletedOutput(eventData, outputItemsByIndex, outputItemsFallback) - eventData = restoreXAINamespaceToolCalls(eventData, prepared.namespaceTools) eventData = xaiNormalizeReasoningSummaryData(eventData) cacheXAIReasoningReplayFromCompleted(ctx, prepared.replayScope, eventData) normalizedEventName = gjson.GetBytes(eventData, "type").String() @@ -862,7 +860,6 @@ func (e *XAIExecutor) prepareResponsesRequestTo(ctx context.Context, req cliprox body, _ = sjson.DeleteBytes(body, "safety_identifier") body, _ = sjson.DeleteBytes(body, "stream_options") namespaceTools := collectXAINamespaceToolRefs(body) - body = normalizeXAIInputNamespaceToolCalls(body) body = normalizeXAITools(body) body = normalizeXAIToolChoiceForTools(body) var replayScope xaiReasoningReplayScope @@ -870,6 +867,7 @@ func (e *XAIExecutor) prepareResponsesRequestTo(ctx context.Context, req cliprox if err != nil { return nil, err } + body = normalizeXAIInputNamespaceToolCalls(body) body = normalizeXAIInputReasoningItems(body) body = sanitizeXAIInputEncryptedContent(body) body = normalizeCodexInstructions(body) @@ -1323,13 +1321,14 @@ func qualifyXAINamespaceToolName(namespaceName, toolName string) string { if namespaceName == "" || toolName == "" || strings.HasPrefix(toolName, "mcp__") { return toolName } - if strings.HasPrefix(toolName, namespaceName) { + prefix := namespaceName + if !strings.HasSuffix(prefix, "__") { + prefix += "__" + } + if strings.HasPrefix(toolName, prefix) { return toolName } - if strings.HasSuffix(namespaceName, "__") { - return namespaceName + toolName - } - return namespaceName + "__" + toolName + return prefix + toolName } func collectXAINamespaceToolRefs(body []byte) map[string]xaiNamespaceToolRef { @@ -1359,6 +1358,9 @@ func collectXAINamespaceToolRefs(body []byte) map[string]xaiNamespaceToolRef { } func normalizeXAIInputNamespaceToolCalls(body []byte) []byte { + if !gjson.ValidBytes(body) { + return body + } input := gjson.GetBytes(body, "input") if !input.Exists() || !input.IsArray() { return body @@ -1375,14 +1377,21 @@ func normalizeXAIInputNamespaceToolCalls(body []byte) []byte { } namePath := fmt.Sprintf("input.%d.name", index) namespacePath := fmt.Sprintf("input.%d.namespace", index) - body, _ = sjson.SetBytes(body, namePath, qualifiedName) - body, _ = sjson.DeleteBytes(body, namespacePath) + 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 { + if len(refs) == 0 || len(data) == 0 || !gjson.ValidBytes(data) { return data } data = restoreXAINamespaceToolCallAtPath(data, "item", refs) @@ -1404,9 +1413,15 @@ func restoreXAINamespaceToolCallAtPath(data []byte, path string, refs map[string if !ok { return data } - data, _ = sjson.SetBytes(data, path+".name", ref.name) - data, _ = sjson.SetBytes(data, path+".namespace", ref.namespace) - 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 diff --git a/internal/runtime/executor/xai_executor_test.go b/internal/runtime/executor/xai_executor_test.go index 5429a8a8..f69cb737 100644 --- a/internal/runtime/executor/xai_executor_test.go +++ b/internal/runtime/executor/xai_executor_test.go @@ -1147,6 +1147,7 @@ func TestQualifyXAINamespaceToolNamePreservesQualifiedNames(t *testing.T) { {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) { @@ -1264,6 +1265,17 @@ func TestRestoreXAINamespaceToolCalls(t *testing.T) { } } +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) diff --git a/internal/runtime/executor/xai_websockets_executor.go b/internal/runtime/executor/xai_websockets_executor.go index adbe2b2e..0a55f1c9 100644 --- a/internal/runtime/executor/xai_websockets_executor.go +++ b/internal/runtime/executor/xai_websockets_executor.go @@ -655,7 +655,6 @@ func (e *XAIWebsocketsExecutor) ExecuteStream(ctx context.Context, auth *cliprox reporter.Publish(ctx, detail) } payload = xaiPatchCompletedOutput(payload, outputItemsByIndex, outputItemsFallback) - payload = restoreXAINamespaceToolCalls(payload, prepared.namespaceTools) payload = xaiNormalizeReasoningSummaryData(payload) cacheXAIReasoningReplayFromCompleted(ctx, prepared.replayScope, payload) if !warmupRequest && idMapper != nil && idMapper.state != nil && !recordedTranscript { From eb6d169407f29393f88690c568bbc28c03e29cf9 Mon Sep 17 00:00:00 2001 From: huahai0202 <2023217276@mail.hfut.edu.cn> Date: Sun, 12 Jul 2026 23:58:16 +0800 Subject: [PATCH 3/4] fix(xai): qualify namespaced tool choices --- internal/runtime/executor/xai_executor.go | 29 ++++++++++++++ .../runtime/executor/xai_executor_test.go | 38 +++++++++++++++++++ 2 files changed, 67 insertions(+) diff --git a/internal/runtime/executor/xai_executor.go b/internal/runtime/executor/xai_executor.go index c52d3f5f..0acc5af7 100644 --- a/internal/runtime/executor/xai_executor.go +++ b/internal/runtime/executor/xai_executor.go @@ -861,6 +861,7 @@ func (e *XAIExecutor) prepareResponsesRequestTo(ctx context.Context, req cliprox 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) @@ -1246,6 +1247,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 diff --git a/internal/runtime/executor/xai_executor_test.go b/internal/runtime/executor/xai_executor_test.go index f69cb737..112d7e8e 100644 --- a/internal/runtime/executor/xai_executor_test.go +++ b/internal/runtime/executor/xai_executor_test.go @@ -1136,6 +1136,44 @@ func TestNormalizeXAITools_QualifiesSameNamedNamespaceTools(t *testing.T) { } } +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 From 4123e275830c93c7061723d7c800ec95b7c9b551 Mon Sep 17 00:00:00 2001 From: huahai0202 <2023217276@mail.hfut.edu.cn> Date: Mon, 13 Jul 2026 10:39:39 +0800 Subject: [PATCH 4/4] fix(xai): handle additional namespace tools --- internal/runtime/executor/xai_executor.go | 112 +++++++++++++----- .../runtime/executor/xai_executor_test.go | 90 ++++++++++++++ .../executor/xai_websockets_executor_test.go | 20 ++-- 3 files changed, 183 insertions(+), 39 deletions(-) diff --git a/internal/runtime/executor/xai_executor.go b/internal/runtime/executor/xai_executor.go index 0acc5af7..0791ea6e 100644 --- a/internal/runtime/executor/xai_executor.go +++ b/internal/runtime/executor/xai_executor.go @@ -1170,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() { @@ -1186,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 { @@ -1194,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 } @@ -1203,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 { @@ -1211,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 @@ -1232,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 } @@ -1362,25 +1404,35 @@ func qualifyXAINamespaceToolName(namespaceName, toolName string) string { func collectXAINamespaceToolRefs(body []byte) map[string]xaiNamespaceToolRef { refs := make(map[string]xaiNamespaceToolRef) - tools := gjson.GetBytes(body, "tools") - if !tools.Exists() || !tools.IsArray() { - return refs - } - for _, tool := range tools.Array() { - if tool.Get("type").String() != xaiNamespaceToolType { - continue + collect := func(tools gjson.Result) { + if !tools.Exists() || !tools.IsArray() { + return } - 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 == "" { + for _, tool := range tools.Array() { + if tool.Get("type").String() != xaiNamespaceToolType { continue } - refs[qualifiedName] = xaiNamespaceToolRef{namespace: namespaceName, name: toolName} + 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 diff --git a/internal/runtime/executor/xai_executor_test.go b/internal/runtime/executor/xai_executor_test.go index 112d7e8e..6035ac35 100644 --- a/internal/runtime/executor/xai_executor_test.go +++ b/internal/runtime/executor/xai_executor_test.go @@ -166,6 +166,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{ @@ -1136,6 +1193,27 @@ func TestNormalizeXAITools_QualifiesSameNamedNamespaceTools(t *testing.T) { } } +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"}}]}], @@ -1359,6 +1437,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_test.go b/internal/runtime/executor/xai_websockets_executor_test.go index a62f70a0..9d8b900e 100644 --- a/internal/runtime/executor/xai_websockets_executor_test.go +++ b/internal/runtime/executor/xai_websockets_executor_test.go @@ -165,12 +165,14 @@ func TestXAIWebsocketsExecuteStreamRestoresNamespaceToolCalls(t *testing.T) { Model: "grok-4.3", Payload: []byte(`{ "model":"grok-4.3", - "input":"use Exa", - "tools":[{ - "type":"namespace", - "name":"mcp__exa", - "tools":[{"type":"function","name":"web_search_exa","parameters":{"type":"object"}}] - }] + "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{ @@ -187,12 +189,12 @@ func TestXAIWebsocketsExecuteStreamRestoresNamespaceToolCalls(t *testing.T) { select { case payload := <-capturedPayload: - tool := gjson.GetBytes(payload, "tools.0") + 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("namespace").Exists() { - t.Fatalf("upstream tool should not contain namespace: %s", 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")