From e674f19129f187ddff65dacd5b1baed9bd9f413b Mon Sep 17 00:00:00 2001 From: Luis Pater Date: Mon, 13 Jul 2026 14:58:03 +0800 Subject: [PATCH] Fix xAI allowed tool namespace choices --- internal/runtime/executor/xai_executor.go | 55 ++++++++++++------- .../runtime/executor/xai_executor_test.go | 48 +++++++++++++++- 2 files changed, 83 insertions(+), 20 deletions(-) diff --git a/internal/runtime/executor/xai_executor.go b/internal/runtime/executor/xai_executor.go index f52be84d4..45ab5422b 100644 --- a/internal/runtime/executor/xai_executor.go +++ b/internal/runtime/executor/xai_executor.go @@ -1289,32 +1289,49 @@ 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. +// normalizeXAINamespaceToolChoice qualifies namespaced function choices using +// the same names sent in the flattened tools list. xAI does not accept the +// Responses namespace field on tool choices. 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 + original := body + normalizeAtPath := func(path string) bool { + toolChoice := gjson.GetBytes(body, path) + if !toolChoice.IsObject() || toolChoice.Get("type").String() != xaiFunctionToolType { + return true + } + namespaceName := strings.TrimSpace(toolChoice.Get("namespace").String()) + toolName := strings.TrimSpace(toolChoice.Get("name").String()) + qualifiedName := qualifyXAINamespaceToolName(namespaceName, toolName) + if namespaceName == "" || qualifiedName == "" { + return true + } + updated, errSet := sjson.SetBytes(body, path+".name", qualifiedName) + if errSet != nil { + return false + } + updated, errDelete := sjson.DeleteBytes(updated, path+".namespace") + if errDelete != nil { + return false + } + body = updated + return true } - namespaceName := strings.TrimSpace(toolChoice.Get("namespace").String()) - toolName := strings.TrimSpace(toolChoice.Get("name").String()) - qualifiedName := qualifyXAINamespaceToolName(namespaceName, toolName) - if namespaceName == "" || qualifiedName == "" { - return body + + if !normalizeAtPath("tool_choice") { + return original } - updated, errSet := sjson.SetBytes(body, "tool_choice.name", qualifiedName) - if errSet != nil { - return body + tools := gjson.GetBytes(body, "tool_choice.tools") + if tools.IsArray() { + for index := range tools.Array() { + if !normalizeAtPath(fmt.Sprintf("tool_choice.tools.%d", index)) { + return original + } + } } - updated, errDelete := sjson.DeleteBytes(updated, "tool_choice.namespace") - if errDelete != nil { - return body - } - return updated + return body } func normalizeXAITool(tool gjson.Result, namespaceName string) ([]byte, bool, bool) { diff --git a/internal/runtime/executor/xai_executor_test.go b/internal/runtime/executor/xai_executor_test.go index ddea75e61..929e58c9d 100644 --- a/internal/runtime/executor/xai_executor_test.go +++ b/internal/runtime/executor/xai_executor_test.go @@ -73,7 +73,7 @@ func TestXAIExecutorExecuteShapesResponsesRequest(t *testing.T) { _, err := exec.Execute(context.Background(), auth, cliproxyexecutor.Request{ Model: "grok-4.3", - Payload: []byte(`{"model":"grok-4.3","input":[{"type":"reasoning","summary":[{"type":"summary_text","text":"test"}],"content":null,"encrypted_content":null},{"type":"reasoning","summary":[{"type":"summary_text","text":"second"}]},{"role":"user","content":"hello"}],"include":["reasoning.encrypted_content"],"reasoning":{"effort":"high"},"tools":[{"type":"tool_search"},{"type":"image_generation"},{"type":"custom","name":"apply_patch"},{"type":"custom","name":"custom_lookup"},{"type":"function","name":"lookup"},{"type":"web_search","external_web_access":true,"search_content_types":["text","image"]},{"type":"namespace","name":"codex_app","description":"Tools in the codex_app namespace.","tools":[{"type":"function","name":"automation_update"},{"type":"custom","name":"namespace_custom"},{"type":"tool_search"}]}]}`), + Payload: []byte(`{"model":"grok-4.3","input":[{"type":"reasoning","summary":[{"type":"summary_text","text":"test"}],"content":null,"encrypted_content":null},{"type":"reasoning","summary":[{"type":"summary_text","text":"second"}]},{"role":"user","content":"hello"}],"include":["reasoning.encrypted_content"],"reasoning":{"effort":"high"},"tools":[{"type":"tool_search"},{"type":"image_generation"},{"type":"custom","name":"apply_patch"},{"type":"custom","name":"custom_lookup"},{"type":"function","name":"lookup"},{"type":"web_search","external_web_access":true,"search_content_types":["text","image"]},{"type":"namespace","name":"codex_app","description":"Tools in the codex_app namespace.","tools":[{"type":"function","name":"automation_update"},{"type":"custom","name":"namespace_custom"},{"type":"tool_search"}]}],"tool_choice":{"type":"allowed_tools","tools":[{"type":"function","name":"automation_update","namespace":"codex_app"},{"type":"function","name":"lookup"},{"type":"web_search"}]}}`), }, cliproxyexecutor.Options{ SourceFormat: sdktranslator.FormatOpenAIResponse, Stream: false, @@ -168,6 +168,18 @@ func TestXAIExecutorExecuteShapesResponsesRequest(t *testing.T) { if !foundNamespaceCustom { t.Fatalf("namespace custom tool was not moved to top-level tools; body=%s", string(gotBody)) } + if got := gjson.GetBytes(gotBody, "tool_choice.tools.0.name").String(); got != "codex_app__automation_update" { + t.Fatalf("tool_choice.tools.0.name = %q, want codex_app__automation_update; body=%s", got, string(gotBody)) + } + if gjson.GetBytes(gotBody, "tool_choice.tools.0.namespace").Exists() { + t.Fatalf("tool_choice.tools.0.namespace should be removed for xAI upstream: %s", string(gotBody)) + } + if got := gjson.GetBytes(gotBody, "tool_choice.tools.1.name").String(); got != "lookup" { + t.Fatalf("tool_choice.tools.1.name = %q, want lookup; body=%s", got, string(gotBody)) + } + if got := gjson.GetBytes(gotBody, "tool_choice.tools.2.type").String(); got != "web_search" { + t.Fatalf("tool_choice.tools.2.type = %q, want web_search; body=%s", got, string(gotBody)) + } foundEncryptedReasoningInclude := false for _, include := range gjson.GetBytes(gotBody, "include").Array() { if include.String() == "reasoning.encrypted_content" { @@ -1351,6 +1363,40 @@ func TestNormalizeXAINamespaceToolChoice(t *testing.T) { } } +func TestNormalizeXAINamespaceToolChoiceAllowedTools(t *testing.T) { + body := []byte(`{ + "tool_choice":{ + "type":"allowed_tools", + "tools":[ + {"type":"function","name":"search","namespace":"mcp__exa"}, + {"type":"function","name":"collaboration__send_message","namespace":"collaboration"}, + {"type":"function","name":"lookup"}, + {"type":"web_search","namespace":"ignored"} + ] + } + }`) + out := normalizeXAINamespaceToolChoice(body) + + if got := gjson.GetBytes(out, "tool_choice.tools.0.name").String(); got != "mcp__exa__search" { + t.Fatalf("tool_choice.tools.0.name = %q, want mcp__exa__search; body=%s", got, string(out)) + } + if gjson.GetBytes(out, "tool_choice.tools.0.namespace").Exists() { + t.Fatalf("tool_choice.tools.0.namespace should be removed: %s", string(out)) + } + if got := gjson.GetBytes(out, "tool_choice.tools.1.name").String(); got != "collaboration__send_message" { + t.Fatalf("tool_choice.tools.1.name = %q, want collaboration__send_message; body=%s", got, string(out)) + } + if gjson.GetBytes(out, "tool_choice.tools.1.namespace").Exists() { + t.Fatalf("tool_choice.tools.1.namespace should be removed: %s", string(out)) + } + if got := gjson.GetBytes(out, "tool_choice.tools.2.name").String(); got != "lookup" { + t.Fatalf("tool_choice.tools.2.name = %q, want lookup; body=%s", got, string(out)) + } + if got := gjson.GetBytes(out, "tool_choice.tools.3.namespace").String(); got != "ignored" { + t.Fatalf("non-function namespace = %q, want ignored; body=%s", got, string(out)) + } +} + func TestNormalizeXAINamespaceToolChoice_PreservesOtherChoices(t *testing.T) { tests := []struct { name string