From 8423cce2d1004e80948a9e2c60ee69354c0aabc3 Mon Sep 17 00:00:00 2001 From: Luis Pater Date: Sun, 26 Jul 2026 20:41:24 +0000 Subject: [PATCH] feat(executor): add configurable injection of x_search tool for xAI requests - Introduced `InjectXSearch` in `XAIConfig` to enable automatic injection of the native `x_search` tool when not explicitly declared. - Updated `XAIExecutor` to honor the `InjectXSearch` configuration, ensuring consistent tool availability. - Enhanced configuration handling with support for dynamic diffing to track changes in `InjectXSearch`. - Added comprehensive tests to validate `InjectXSearch` behavior, including preparation and tool choice synchronization. - Updated example config and documentation to outline `InjectXSearch` usage. Closes: #4339 --- config.example.yaml | 6 ++++++ internal/config/config.go | 3 +++ internal/config/config_types.go | 6 ++++++ internal/config/xai_api_key_test.go | 20 ++++++++++++++++++++ internal/runtime/executor/xai_executor.go | 5 ++--- internal/runtime/executor/xai_executor_request.go | 15 ++++++++------- internal/runtime/executor/xai_executor_test.go | 78 ++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++---- internal/watcher/diff/config_diff.go | 3 +++ internal/watcher/diff/config_diff_test.go | 2 ++ 9 file(s) changed, 124 insertion(s)(+), 14 deletion(s)(-) diff --git a/config.example.yaml b/config.example.yaml --- a/config.example.yaml +++ b/config.example.yaml @@ -244,6 +244,12 @@ # username: "user" # credential: "secret" +# xAI provider behavior. +xai: + # When true, inject the native x_search tool when the request does not declare it. + # The injected tool is also added to tool_choice.allowed_tools when applicable. + inject-x-search: false + # When true, enable authentication for the WebSocket API (/v1/ws). ws-auth: true diff --git a/internal/config/config.go b/internal/config/config.go --- a/internal/config/config.go +++ b/internal/config/config.go @@ -112,6 +112,9 @@ // XAIKey defines xAI API key configurations using the same structure as Codex API keys. XAIKey []XAIKey `yaml:"xai-api-key" json:"xai-api-key"` + // XAI configures provider-wide xAI request behavior. + XAI XAIConfig `yaml:"xai" json:"xai"` + // Codex configures provider-wide Codex request behavior. Codex CodexConfig `yaml:"codex" json:"codex"` diff --git a/internal/config/config_types.go b/internal/config/config_types.go --- a/internal/config/config_types.go +++ b/internal/config/config_types.go @@ -118,6 +118,12 @@ BetaFeatures string `yaml:"beta-features" json:"beta-features"` } +// XAIConfig configures provider-wide xAI request behavior. +type XAIConfig struct { + // InjectXSearch injects xAI's native x_search tool when the request does not declare it. + InjectXSearch bool `yaml:"inject-x-search" json:"inject-x-search"` +} + // CodexConfig configures provider-wide Codex request behavior. type CodexConfig struct { IdentityConfuse bool `yaml:"identity-confuse" json:"identity-confuse"` diff --git a/internal/config/xai_api_key_test.go b/internal/config/xai_api_key_test.go --- a/internal/config/xai_api_key_test.go +++ b/internal/config/xai_api_key_test.go @@ -2,6 +2,26 @@ import "testing" +func TestParseConfigBytesXAIConfig(t *testing.T) { + defaultCfg, errDefault := ParseConfigBytes([]byte(`{}`)) + if errDefault != nil { + t.Fatalf("ParseConfigBytes(default) error = %v", errDefault) + } + if defaultCfg.XAI.InjectXSearch { + t.Fatal("xai.inject-x-search = true by default, want false") + } + + enabledCfg, errEnabled := ParseConfigBytes([]byte(`xai: + inject-x-search: true +`)) + if errEnabled != nil { + t.Fatalf("ParseConfigBytes(enabled) error = %v", errEnabled) + } + if !enabledCfg.XAI.InjectXSearch { + t.Fatal("xai.inject-x-search = false, want true") + } +} + func TestParseConfigBytesXAIAPIKeyMatchesCodexShape(t *testing.T) { cfg, errParse := ParseConfigBytes([]byte(`xai-api-key: - api-key: " xai-key " diff --git a/internal/runtime/executor/xai_executor.go b/internal/runtime/executor/xai_executor.go --- a/internal/runtime/executor/xai_executor.go +++ b/internal/runtime/executor/xai_executor.go @@ -52,9 +52,8 @@ 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). +// xaiXSearchToolJSON is the native X Search tool injected when enabled by config. +// Internal subtool traces are still filtered downstream when this tool is present. var xaiXSearchToolJSON = []byte(`{"type":"x_search"}`) // XAIExecutor is a stateless executor for xAI Grok's Responses API. diff --git a/internal/runtime/executor/xai_executor_request.go b/internal/runtime/executor/xai_executor_request.go --- a/internal/runtime/executor/xai_executor_request.go +++ b/internal/runtime/executor/xai_executor_request.go @@ -94,13 +94,14 @@ clientDeclaredTools := collectXAIClientDeclaredToolKeys(body) body = normalizeXAITools(body) body = promoteXAIAdditionalTools(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. + // Drop choices that point at tools removed by normalizeXAITools before any + // configured x_search injection, so no surviving choice references a deleted tool. body = normalizeXAINamespaceToolChoice(body) body = pruneXAIOrphanedToolChoice(body) body = normalizeXAIToolChoiceForTools(body) - body = ensureXAINativeXSearchTool(body) + if e.cfg != nil && e.cfg.XAI.InjectXSearch { + body = ensureXAINativeXSearchTool(body) + } var replayScope xaiReasoningReplayScope body, replayScope, err = applyXAIReasoningReplayCacheRequired(ctx, from, req, opts, body) if err != nil { @@ -542,9 +543,9 @@ // 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. +// Grok can select the injected tool. When injection is enabled, 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 diff --git a/internal/runtime/executor/xai_executor_test.go b/internal/runtime/executor/xai_executor_test.go --- a/internal/runtime/executor/xai_executor_test.go +++ b/internal/runtime/executor/xai_executor_test.go @@ -175,7 +175,7 @@ })) defer server.Close() - exec := NewXAIExecutor(&config.Config{}) + exec := NewXAIExecutor(&config.Config{XAI: config.XAIConfig{InjectXSearch: true}}) auth := &cliproxyauth.Auth{ ID: "xai-auth", Provider: "xai", @@ -656,6 +656,76 @@ } } +func TestXAIExecutorPrepareHonorsInjectXSearchConfig(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + cfg *config.Config + wantXSearch bool + }{ + {name: "default disabled", cfg: &config.Config{}, wantXSearch: false}, + {name: "explicitly enabled", cfg: &config.Config{XAI: config.XAIConfig{InjectXSearch: true}}, wantXSearch: true}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + exec := NewXAIExecutor(tt.cfg) + prepared, errPrepare := exec.prepareResponsesRequest(context.Background(), cliproxyexecutor.Request{ + Model: "grok-4.5", + Payload: []byte(`{ + "model":"grok-4.5", + "input":"search the web", + "tools":[{"type":"function","name":"web_search","parameters":{"type":"object"}}], + "tool_choice":{"type":"allowed_tools","tools":[{"type":"function","name":"web_search"}]} + }`), + }, cliproxyexecutor.Options{ + SourceFormat: sdktranslator.FormatOpenAIResponse, + Stream: false, + }, false) + if errPrepare != nil { + t.Fatalf("prepareResponsesRequest() error = %v", errPrepare) + } + + wantXSearchCount := 0 + if tt.wantXSearch { + wantXSearchCount = 1 + } + tools := gjson.GetBytes(prepared.body, "tools").Array() + if len(tools) != 1+wantXSearchCount { + t.Fatalf("tools length = %d, want %d; body=%s", len(tools), 1+wantXSearchCount, prepared.body) + } + if got := tools[0].Get("name").String(); got != "web_search" { + t.Fatalf("client web_search tool missing; body=%s", prepared.body) + } + xSearchTools := 0 + for _, tool := range tools { + if tool.Get("type").String() == "x_search" { + xSearchTools++ + } + } + if xSearchTools != wantXSearchCount { + t.Fatalf("x_search tools = %d, want %d; body=%s", xSearchTools, wantXSearchCount, prepared.body) + } + + xSearchAllowed := 0 + for _, tool := range gjson.GetBytes(prepared.body, "tool_choice.tools").Array() { + if tool.Get("type").String() == "x_search" { + xSearchAllowed++ + } + } + if xSearchAllowed != wantXSearchCount { + t.Fatalf("allowed x_search tools = %d, want %d; body=%s", xSearchAllowed, wantXSearchCount, prepared.body) + } + if prepared.filterInternalXSearch != tt.wantXSearch { + t.Fatalf("filterInternalXSearch = %t, want %t", prepared.filterInternalXSearch, tt.wantXSearch) + } + }) + } +} + func TestEnsureXAINativeXSearchTool(t *testing.T) { t.Parallel() @@ -783,7 +853,7 @@ func TestXAIExecutorPrepareDropsOrphanedToolChoiceBeforeXSearchInject(t *testing.T) { t.Parallel() - exec := NewXAIExecutor(&config.Config{}) + exec := NewXAIExecutor(&config.Config{XAI: config.XAIConfig{InjectXSearch: true}}) prepared, err := exec.prepareResponsesRequest(context.Background(), cliproxyexecutor.Request{ Model: "grok-4.5", // image_generation is stripped by normalizeXAITools; without pruning, the @@ -1053,7 +1123,7 @@ func TestXAIExecutorPrepareAllowedToolsSyncsInjectedXSearch(t *testing.T) { t.Parallel() - exec := NewXAIExecutor(&config.Config{}) + exec := NewXAIExecutor(&config.Config{XAI: config.XAIConfig{InjectXSearch: true}}) prepared, err := exec.prepareResponsesRequest(context.Background(), cliproxyexecutor.Request{ Model: "grok-4.5", // Only image_generation remains after client filtering of tool_search-like @@ -2306,7 +2376,7 @@ })) defer server.Close() - exec := NewXAIExecutor(&config.Config{}) + exec := NewXAIExecutor(&config.Config{XAI: config.XAIConfig{InjectXSearch: true}}) auth := &cliproxyauth.Auth{ Provider: "xai", Attributes: map[string]string{"base_url": server.URL}, diff --git a/internal/watcher/diff/config_diff.go b/internal/watcher/diff/config_diff.go --- a/internal/watcher/diff/config_diff.go +++ b/internal/watcher/diff/config_diff.go @@ -108,6 +108,9 @@ if oldCfg.Codex.OptimizeMultiAgentV2 != newCfg.Codex.OptimizeMultiAgentV2 { changes = append(changes, fmt.Sprintf("codex.optimize-multi-agent-v2: %t -> %t", oldCfg.Codex.OptimizeMultiAgentV2, newCfg.Codex.OptimizeMultiAgentV2)) } + if oldCfg.XAI.InjectXSearch != newCfg.XAI.InjectXSearch { + changes = append(changes, fmt.Sprintf("xai.inject-x-search: %t -> %t", oldCfg.XAI.InjectXSearch, newCfg.XAI.InjectXSearch)) + } oldLiveRelay := oldCfg.Codex.LiveMediaRelay newLiveRelay := newCfg.Codex.LiveMediaRelay if oldLiveRelay.Enabled != newLiveRelay.Enabled { diff --git a/internal/watcher/diff/config_diff_test.go b/internal/watcher/diff/config_diff_test.go --- a/internal/watcher/diff/config_diff_test.go +++ b/internal/watcher/diff/config_diff_test.go @@ -352,6 +352,7 @@ MaxRetryInterval: 3, WebsocketAuth: true, QuotaExceeded: config.QuotaExceeded{SwitchProject: true, SwitchPreviewModel: true, AntigravityCredits: true}, + XAI: config.XAIConfig{InjectXSearch: true}, ClaudeKey: []config.ClaudeKey{ {APIKey: "c1", BaseURL: "http://new", ProxyURL: "http://p", Headers: map[string]string{"H": "1"}, ExcludedModels: []string{"a"}}, {APIKey: "c2"}, @@ -395,6 +396,7 @@ expectContains(t, details, "quota-exceeded.switch-project: false -> true") expectContains(t, details, "quota-exceeded.switch-preview-model: false -> true") expectContains(t, details, "quota-exceeded.antigravity-credits: false -> true") + expectContains(t, details, "xai.inject-x-search: false -> true") expectContains(t, details, "api-keys count: 1 -> 2") expectContains(t, details, "claude-api-key count: 1 -> 2") expectContains(t, details, "codex-api-key count: 1 -> 2") -- tangled.sh