From a9831c841ab63dcc62ab24cfb0d86fd2db1e01d4 Mon Sep 17 00:00:00 2001 From: Luis Pater Date: Tue, 14 Jul 2026 05:03:19 +0800 Subject: [PATCH] fix(xai): sync allowed_tools and prune orphaned choices for x_search inject Normalize and drop tool_choice entries that reference tools removed by normalizeXAITools before injecting native x_search, then allow x_search in allowed_tools without duplicates so Grok can select the injected tool. --- internal/runtime/executor/xai_executor.go | 201 +++++++++++-- .../runtime/executor/xai_executor_test.go | 266 +++++++++++++++++- 2 files changed, 437 insertions(+), 30 deletions(-) diff --git a/internal/runtime/executor/xai_executor.go b/internal/runtime/executor/xai_executor.go index bf7eaee69..bff83c4d8 100644 --- a/internal/runtime/executor/xai_executor.go +++ b/internal/runtime/executor/xai_executor.go @@ -71,6 +71,11 @@ const ( xaiUsingAPIAttr = "using_api" ) +// Always inject native x_search when the client did not declare it so Grok can +// run X Search server-side. Internal subtool traces are still filtered downstream +// when this native tool is present (see filterInternalXSearch). +var xaiXSearchToolJSON = []byte(`{"type":"x_search"}`) + // XAIExecutor is a stateless executor for xAI Grok's Responses API. type XAIExecutor struct { cfg *config.Config @@ -894,8 +899,13 @@ func (e *XAIExecutor) prepareResponsesRequestTo(ctx context.Context, req cliprox // the post-restore (namespace, short-name) shape used by the response filter. clientDeclaredTools := collectXAIClientDeclaredToolKeys(body) body = normalizeXAITools(body) + // Drop choices that point at tools removed by normalizeXAITools before we + // inject native x_search, so a surviving allowed_tools / forced choice is not + // left pointing at a deleted tool once only x_search remains. body = normalizeXAINamespaceToolChoice(body) + body = pruneXAIOrphanedToolChoice(body) body = normalizeXAIToolChoiceForTools(body) + body = ensureXAINativeXSearchTool(body) var replayScope xaiReasoningReplayScope body, replayScope, err = applyXAIReasoningReplayCacheRequired(ctx, from, req, opts, body) if err != nil { @@ -1204,6 +1214,171 @@ func sanitizeXAIResponsesBody(body []byte, model string) []byte { return body } +// ensureXAINativeXSearchTool appends {"type":"x_search"} when the final tools +// list does not already include native X Search. When tool_choice restricts the +// model to allowed_tools, x_search is also added there (without duplicates) so +// Grok can select the injected tool. HTTP and websocket executors both prepare +// payloads through prepareResponsesRequestTo, so this runs once before the body +// is submitted upstream. +func ensureXAINativeXSearchTool(body []byte) []byte { + if !gjson.ValidBytes(body) { + return body + } + if !xaiRequestHasNativeXSearch(body) { + tools := gjson.GetBytes(body, "tools") + if !tools.Exists() || !tools.IsArray() { + body, _ = sjson.SetRawBytes(body, "tools", []byte(`[{"type":"x_search"}]`)) + } else { + body, _ = sjson.SetRawBytes(body, "tools.-1", xaiXSearchToolJSON) + } + } + return ensureXAINativeXSearchAllowedTools(body) +} + +// ensureXAINativeXSearchAllowedTools appends x_search to tool_choice.tools when +// the choice mode is allowed_tools and x_search is not already listed. +func ensureXAINativeXSearchAllowedTools(body []byte) []byte { + choice := gjson.GetBytes(body, "tool_choice") + if !choice.IsObject() || choice.Get("type").String() != "allowed_tools" { + return body + } + allowed := choice.Get("tools") + if !allowed.Exists() || !allowed.IsArray() { + body, _ = sjson.SetRawBytes(body, "tool_choice.tools", []byte(`[{"type":"x_search"}]`)) + return body + } + for _, tool := range allowed.Array() { + if strings.TrimSpace(tool.Get("type").String()) == xaiXSearchToolType { + return body + } + } + body, _ = sjson.SetRawBytes(body, "tool_choice.tools.-1", xaiXSearchToolJSON) + return body +} + +// pruneXAIOrphanedToolChoice removes tool_choice entries that no longer match +// any remaining tool after normalizeXAITools filtering. Forced choices that +// reference a deleted tool are dropped entirely; allowed_tools lists keep only +// choices that still resolve against the post-normalization tools set. +func pruneXAIOrphanedToolChoice(body []byte) []byte { + if !gjson.ValidBytes(body) { + return body + } + choice := gjson.GetBytes(body, "tool_choice") + if !choice.Exists() { + return body + } + available := collectXAIAvailableToolChoiceKeys(body) + if choice.Type == gjson.String { + // auto / none / required are not tool references. + return body + } + if !choice.IsObject() { + return body + } + choiceType := strings.TrimSpace(choice.Get("type").String()) + switch choiceType { + case "allowed_tools": + return pruneXAIAllowedToolsChoice(body, available) + default: + if choiceType == "" { + return body + } + if xaiToolChoiceMatchesAvailable(choice, available) { + return body + } + body, _ = sjson.DeleteBytes(body, "tool_choice") + return body + } +} + +func pruneXAIAllowedToolsChoice(body []byte, available map[xaiToolChoiceKey]struct{}) []byte { + allowed := gjson.GetBytes(body, "tool_choice.tools") + if !allowed.Exists() || !allowed.IsArray() { + body, _ = sjson.DeleteBytes(body, "tool_choice") + return body + } + filtered := []byte(`[]`) + changed := false + for _, tool := range allowed.Array() { + if !xaiToolChoiceMatchesAvailable(tool, available) { + changed = true + continue + } + updated, errSet := sjson.SetRawBytes(filtered, "-1", []byte(tool.Raw)) + if errSet != nil { + return body + } + filtered = updated + } + if !changed { + return body + } + if len(gjson.ParseBytes(filtered).Array()) == 0 { + body, _ = sjson.DeleteBytes(body, "tool_choice") + return body + } + body, _ = sjson.SetRawBytes(body, "tool_choice.tools", filtered) + return body +} + +// xaiToolChoiceKey identifies a selectable tool the way xAI tool_choice entries +// reference it after namespace qualification: type alone for host tools, or +// type+name for function tools. +type xaiToolChoiceKey struct { + toolType string + name string +} + +func collectXAIAvailableToolChoiceKeys(body []byte) map[xaiToolChoiceKey]struct{} { + keys := make(map[xaiToolChoiceKey]struct{}) + collect := func(tools gjson.Result) { + if !tools.IsArray() { + return + } + for _, tool := range tools.Array() { + toolType := strings.TrimSpace(tool.Get("type").String()) + if toolType == "" { + continue + } + key := xaiToolChoiceKey{toolType: toolType} + if toolType == xaiFunctionToolType || toolType == xaiCustomToolType { + key.name = strings.TrimSpace(tool.Get("name").String()) + if key.name == "" { + continue + } + } + keys[key] = struct{}{} + } + } + collect(gjson.GetBytes(body, "tools")) + input := gjson.GetBytes(body, "input") + if input.IsArray() { + for _, item := range input.Array() { + if item.Get("type").String() == "additional_tools" { + collect(item.Get("tools")) + } + } + } + return keys +} + +func xaiToolChoiceMatchesAvailable(choice gjson.Result, available map[xaiToolChoiceKey]struct{}) bool { + toolType := strings.TrimSpace(choice.Get("type").String()) + if toolType == "" { + return false + } + key := xaiToolChoiceKey{toolType: toolType} + if toolType == xaiFunctionToolType || toolType == xaiCustomToolType { + key.name = strings.TrimSpace(choice.Get("name").String()) + if key.name == "" { + return false + } + } + _, ok := available[key] + return ok +} + func normalizeXAITools(body []byte) []byte { if !gjson.ValidBytes(body) { return body @@ -1604,30 +1779,12 @@ func newXAIInternalXSearchResponseFilter(enabled bool, clientDeclaredTools map[x } func xaiRequestHasNativeXSearch(body []byte) bool { - hasXSearch := func(tools gjson.Result) bool { - if !tools.IsArray() { - return false - } - for _, tool := range tools.Array() { - if strings.TrimSpace(tool.Get("type").String()) == xaiXSearchToolType { - return true - } - } - return false - } - if hasXSearch(gjson.GetBytes(body, "tools")) { + if gjson.GetBytes(body, `tools.#(type=="x_search")`).Exists() { return true } - input := gjson.GetBytes(body, "input") - if !input.IsArray() { - return false - } - for _, item := range input.Array() { - if item.Get("type").String() == "additional_tools" && hasXSearch(item.Get("tools")) { - return true - } - } - return false + // Multipath queries return an array of matches; an empty array still Exists(). + // Check the match count instead of Exists() for additional_tools injection. + return len(gjson.GetBytes(body, `input.#(type=="additional_tools")#.tools.#(type=="x_search")`).Array()) > 0 } // collectXAIClientDeclaredToolKeys records client-declared function/custom tools diff --git a/internal/runtime/executor/xai_executor_test.go b/internal/runtime/executor/xai_executor_test.go index 6a535d52e..90a7ca448 100644 --- a/internal/runtime/executor/xai_executor_test.go +++ b/internal/runtime/executor/xai_executor_test.go @@ -129,18 +129,22 @@ func TestXAIExecutorExecuteShapesResponsesRequest(t *testing.T) { t.Fatalf("input.2 exists, want consecutive reasoning item merged; body=%s", string(gotBody)) } tools := gjson.GetBytes(gotBody, "tools").Array() - if len(tools) != 5 { - t.Fatalf("tools length = %d, want 5; body=%s", len(tools), string(gotBody)) + if len(tools) != 6 { + t.Fatalf("tools length = %d, want 6; body=%s", len(tools), string(gotBody)) } foundAutomationUpdate := false foundNamespaceCustom := false + foundXSearch := false for i, tool := range tools { toolType := tool.Get("type").String() if toolType == "image_generation" { t.Fatalf("tools.%d.type = image_generation, want removed; body=%s", i, string(gotBody)) } - if toolType != "function" && toolType != "web_search" { - t.Fatalf("tools.%d.type = %q, want function or web_search; body=%s", i, toolType, string(gotBody)) + if toolType != "function" && toolType != "web_search" && toolType != "x_search" { + t.Fatalf("tools.%d.type = %q, want function, web_search, or x_search; body=%s", i, toolType, string(gotBody)) + } + if toolType == "x_search" { + foundXSearch = true } if toolType == "function" && !tool.Get("parameters").Exists() { t.Fatalf("tools.%d.parameters missing for xAI function tool; body=%s", i, string(gotBody)) @@ -169,6 +173,9 @@ func TestXAIExecutorExecuteShapesResponsesRequest(t *testing.T) { if !foundNamespaceCustom { t.Fatalf("namespace custom tool was not moved to top-level tools; body=%s", string(gotBody)) } + if !foundXSearch { + t.Fatalf("native x_search tool was not injected; 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)) } @@ -181,6 +188,18 @@ func TestXAIExecutorExecuteShapesResponsesRequest(t *testing.T) { 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)) } + if got := gjson.GetBytes(gotBody, "tool_choice.tools.3.type").String(); got != "x_search" { + t.Fatalf("tool_choice.tools.3.type = %q, want x_search; body=%s", got, string(gotBody)) + } + xSearchAllowedCount := 0 + for _, tool := range gjson.GetBytes(gotBody, "tool_choice.tools").Array() { + if tool.Get("type").String() == "x_search" { + xSearchAllowedCount++ + } + } + if xSearchAllowedCount != 1 { + t.Fatalf("allowed_tools x_search count = %d, want 1; body=%s", xSearchAllowedCount, string(gotBody)) + } foundEncryptedReasoningInclude := false for _, include := range gjson.GetBytes(gotBody, "include").Array() { if include.String() == "reasoning.encrypted_content" { @@ -462,6 +481,230 @@ func TestXAIExecutorExecuteFiltersInternalXSearchCalls(t *testing.T) { } } +func TestEnsureXAINativeXSearchTool(t *testing.T) { + t.Parallel() + + // Missing tools array: inject a top-level x_search tool. + out := ensureXAINativeXSearchTool([]byte(`{"model":"grok-4.5","input":"hi"}`)) + tools := gjson.GetBytes(out, "tools").Array() + if len(tools) != 1 { + t.Fatalf("tools length = %d, want 1; body=%s", len(tools), out) + } + if got := tools[0].Get("type").String(); got != "x_search" { + t.Fatalf("tools.0.type = %q, want x_search; body=%s", got, out) + } + + // Existing tools without x_search: append once. + out = ensureXAINativeXSearchTool([]byte(`{"tools":[{"type":"web_search"},{"type":"function","name":"lookup","parameters":{"type":"object"}}]}`)) + tools = gjson.GetBytes(out, "tools").Array() + if len(tools) != 3 { + t.Fatalf("tools length = %d, want 3; body=%s", len(tools), out) + } + if got := tools[2].Get("type").String(); got != "x_search" { + t.Fatalf("tools.2.type = %q, want x_search; body=%s", got, out) + } + + // Already present: leave body unchanged (no duplicate). + in := []byte(`{"tools":[{"type":"function","name":"lookup","parameters":{"type":"object"}},{"type":"x_search"}]}`) + out = ensureXAINativeXSearchTool(in) + tools = gjson.GetBytes(out, "tools").Array() + if len(tools) != 2 { + t.Fatalf("tools length = %d, want 2; body=%s", len(tools), out) + } + xSearchCount := 0 + for _, tool := range tools { + if tool.Get("type").String() == "x_search" { + xSearchCount++ + } + } + if xSearchCount != 1 { + t.Fatalf("x_search count = %d, want 1; body=%s", xSearchCount, out) + } + + // allowed_tools without x_search: append once so Grok may select it. + out = ensureXAINativeXSearchTool([]byte(`{ + "tools":[{"type":"function","name":"lookup","parameters":{"type":"object"}}], + "tool_choice":{"type":"allowed_tools","tools":[{"type":"function","name":"lookup"}]} + }`)) + if got := gjson.GetBytes(out, "tools.1.type").String(); got != "x_search" { + t.Fatalf("tools.1.type = %q, want x_search; body=%s", got, out) + } + if got := gjson.GetBytes(out, "tool_choice.tools.1.type").String(); got != "x_search" { + t.Fatalf("tool_choice.tools.1.type = %q, want x_search; body=%s", got, out) + } + + // allowed_tools already lists x_search: do not duplicate. + out = ensureXAINativeXSearchTool([]byte(`{ + "tools":[{"type":"web_search"},{"type":"x_search"}], + "tool_choice":{"type":"allowed_tools","tools":[{"type":"web_search"},{"type":"x_search"}]} + }`)) + tools = gjson.GetBytes(out, "tools").Array() + if len(tools) != 2 { + t.Fatalf("tools length = %d, want 2; body=%s", len(tools), out) + } + allowed := gjson.GetBytes(out, "tool_choice.tools").Array() + if len(allowed) != 2 { + t.Fatalf("tool_choice.tools length = %d, want 2; body=%s", len(allowed), out) + } + xSearchAllowed := 0 + for _, tool := range allowed { + if tool.Get("type").String() == "x_search" { + xSearchAllowed++ + } + } + if xSearchAllowed != 1 { + t.Fatalf("allowed_tools x_search count = %d, want 1; body=%s", xSearchAllowed, out) + } +} + +func TestPruneXAIOrphanedToolChoice(t *testing.T) { + t.Parallel() + + // Forced choice for a removed tool is dropped. + out := pruneXAIOrphanedToolChoice([]byte(`{ + "tools":[{"type":"function","name":"lookup","parameters":{"type":"object"}}], + "tool_choice":{"type":"image_generation"} + }`)) + if gjson.GetBytes(out, "tool_choice").Exists() { + t.Fatalf("orphaned forced tool_choice should be removed: %s", out) + } + + // allowed_tools keeps only still-available entries. + out = pruneXAIOrphanedToolChoice([]byte(`{ + "tools":[{"type":"function","name":"lookup","parameters":{"type":"object"}},{"type":"web_search"}], + "tool_choice":{"type":"allowed_tools","tools":[ + {"type":"function","name":"lookup"}, + {"type":"image_generation"}, + {"type":"web_search"} + ]} + }`)) + allowed := gjson.GetBytes(out, "tool_choice.tools").Array() + if len(allowed) != 2 { + t.Fatalf("allowed_tools length = %d, want 2; body=%s", len(allowed), out) + } + if got := allowed[0].Get("name").String(); got != "lookup" { + t.Fatalf("allowed_tools.0.name = %q, want lookup; body=%s", got, out) + } + if got := allowed[1].Get("type").String(); got != "web_search" { + t.Fatalf("allowed_tools.1.type = %q, want web_search; body=%s", got, out) + } + + // When every allowed entry is orphaned, drop tool_choice entirely. + out = pruneXAIOrphanedToolChoice([]byte(`{ + "tools":[], + "tool_choice":{"type":"allowed_tools","tools":[{"type":"image_generation"}]} + }`)) + if gjson.GetBytes(out, "tool_choice").Exists() { + t.Fatalf("fully orphaned allowed_tools should be removed: %s", out) + } + + // String choices are not tool references. + in := []byte(`{"tools":[{"type":"web_search"}],"tool_choice":"auto"}`) + if got := pruneXAIOrphanedToolChoice(in); !bytes.Equal(got, in) { + t.Fatalf("string tool_choice changed: got=%s want=%s", got, in) + } +} + +func TestXAIExecutorPrepareDropsOrphanedToolChoiceBeforeXSearchInject(t *testing.T) { + t.Parallel() + + exec := NewXAIExecutor(&config.Config{}) + prepared, err := exec.prepareResponsesRequest(context.Background(), cliproxyexecutor.Request{ + Model: "grok-4.5", + // image_generation is stripped by normalizeXAITools; without pruning, the + // forced choice would survive next to the injected x_search tool. + Payload: []byte(`{ + "model":"grok-4.5", + "input":"draw something", + "tools":[{"type":"image_generation"}], + "tool_choice":{"type":"image_generation"} + }`), + }, cliproxyexecutor.Options{ + SourceFormat: sdktranslator.FormatOpenAIResponse, + Stream: false, + }, false) + if err != nil { + t.Fatalf("prepareResponsesRequest() error = %v", err) + } + + tools := gjson.GetBytes(prepared.body, "tools").Array() + if len(tools) != 1 { + t.Fatalf("tools length = %d, want 1; body=%s", len(tools), prepared.body) + } + if got := tools[0].Get("type").String(); got != "x_search" { + t.Fatalf("tools.0.type = %q, want x_search; body=%s", got, prepared.body) + } + if gjson.GetBytes(prepared.body, "tool_choice").Exists() { + t.Fatalf("orphaned image_generation tool_choice must not reach upstream: %s", prepared.body) + } +} + +func TestXAIExecutorPrepareAllowedToolsSyncsInjectedXSearch(t *testing.T) { + t.Parallel() + + exec := NewXAIExecutor(&config.Config{}) + prepared, err := exec.prepareResponsesRequest(context.Background(), cliproxyexecutor.Request{ + Model: "grok-4.5", + // Only image_generation remains after client filtering of tool_search-like + // tools is not relevant here: normalizeXAITools drops image_generation and + // we inject x_search, while allowed_tools must be rewritten so Grok can + // choose the injected tool and not a deleted one. + Payload: []byte(`{ + "model":"grok-4.5", + "input":"search X", + "tools":[{"type":"image_generation"},{"type":"function","name":"lookup","parameters":{"type":"object"}}], + "tool_choice":{"type":"allowed_tools","tools":[ + {"type":"image_generation"}, + {"type":"function","name":"lookup"} + ]} + }`), + }, cliproxyexecutor.Options{ + SourceFormat: sdktranslator.FormatOpenAIResponse, + Stream: false, + }, false) + if err != nil { + t.Fatalf("prepareResponsesRequest() error = %v", err) + } + + tools := gjson.GetBytes(prepared.body, "tools").Array() + if len(tools) != 2 { + t.Fatalf("tools length = %d, want 2; body=%s", len(tools), prepared.body) + } + foundLookup := false + foundXSearch := false + for _, tool := range tools { + switch tool.Get("type").String() { + case "function": + if tool.Get("name").String() == "lookup" { + foundLookup = true + } + case "x_search": + foundXSearch = true + case "image_generation": + t.Fatalf("image_generation must be removed; body=%s", prepared.body) + } + } + if !foundLookup || !foundXSearch { + t.Fatalf("expected lookup + x_search tools; body=%s", prepared.body) + } + + allowed := gjson.GetBytes(prepared.body, "tool_choice.tools").Array() + if len(allowed) != 2 { + t.Fatalf("tool_choice.tools length = %d, want 2; body=%s", len(allowed), prepared.body) + } + if got := allowed[0].Get("name").String(); got != "lookup" { + t.Fatalf("tool_choice.tools.0.name = %q, want lookup; body=%s", got, prepared.body) + } + if got := allowed[1].Get("type").String(); got != "x_search" { + t.Fatalf("tool_choice.tools.1.type = %q, want x_search; body=%s", got, prepared.body) + } + for _, tool := range allowed { + if tool.Get("type").String() == "image_generation" { + t.Fatalf("orphaned image_generation choice leaked: %s", prepared.body) + } + } +} + func TestXAIInternalXSearchResponseFilterRequiresNativeTool(t *testing.T) { if xaiRequestHasNativeXSearch([]byte(`{"tools":[{"type":"web_search"}]}`)) { t.Fatal("web_search must not enable internal X search filtering") @@ -1578,8 +1821,8 @@ func TestXAIExecutorExecuteStreamFiltersToolSearchTool(t *testing.T) { } tools := gjson.GetBytes(gotBody, "tools").Array() - if len(tools) != 5 { - t.Fatalf("tools length = %d, want 5; body=%s", len(tools), string(gotBody)) + if len(tools) != 6 { + t.Fatalf("tools length = %d, want 6; body=%s", len(tools), string(gotBody)) } if gjson.GetBytes(gotBody, "input.0.content").Exists() { t.Fatalf("input.0.content exists, want removed; body=%s", string(gotBody)) @@ -1601,13 +1844,14 @@ func TestXAIExecutorExecuteStreamFiltersToolSearchTool(t *testing.T) { } foundAutomationUpdate := false foundNamespaceCustom := false + foundXSearch := false for i, tool := range tools { toolType := tool.Get("type").String() if toolType == "image_generation" { t.Fatalf("tools.%d.type = image_generation, want removed; body=%s", i, string(gotBody)) } - if toolType != "function" && toolType != "web_search" { - t.Fatalf("tools.%d.type = %q, want function or web_search; body=%s", i, toolType, string(gotBody)) + if toolType != "function" && toolType != "web_search" && toolType != "x_search" { + t.Fatalf("tools.%d.type = %q, want function, web_search, or x_search; body=%s", i, toolType, string(gotBody)) } if toolType == "function" && !tool.Get("parameters").Exists() { t.Fatalf("tools.%d.parameters missing for xAI function tool; body=%s", i, string(gotBody)) @@ -1621,6 +1865,9 @@ func TestXAIExecutorExecuteStreamFiltersToolSearchTool(t *testing.T) { case "codex_app__namespace_custom": foundNamespaceCustom = true } + if toolType == "x_search" { + foundXSearch = true + } if toolType == "web_search" { if tool.Get("external_web_access").Exists() { t.Fatalf("tools.%d.external_web_access exists, want removed; body=%s", i, string(gotBody)) @@ -1636,6 +1883,9 @@ func TestXAIExecutorExecuteStreamFiltersToolSearchTool(t *testing.T) { if !foundNamespaceCustom { t.Fatalf("namespace custom tool was not moved to top-level tools; body=%s", string(gotBody)) } + if !foundXSearch { + t.Fatalf("native x_search tool was not injected; body=%s", string(gotBody)) + } } func TestXAIExecutorExecuteStreamNormalizesReasoningTextEvents(t *testing.T) {