From 5af38ccf416a8bcba36174c8fe68e1afa0bb8ab0 Mon Sep 17 00:00:00 2001 From: superman2003 <2112076433zcr@gmail.com> Date: Sun, 19 Jul 2026 14:41:38 +0800 Subject: [PATCH] fix(grok): add guarded client tool cache routing --- .../service/openai_gateway_grok_cache.go | 136 +++++++++++++++--- .../service/openai_gateway_grok_cache_test.go | 111 ++++++++++++++ 2 files changed, 226 insertions(+), 21 deletions(-) diff --git a/backend/internal/service/openai_gateway_grok_cache.go b/backend/internal/service/openai_gateway_grok_cache.go index d9e34cbb58..1ec6159c1a 100644 --- a/backend/internal/service/openai_gateway_grok_cache.go +++ b/backend/internal/service/openai_gateway_grok_cache.go @@ -12,10 +12,12 @@ import ( ) const ( - grokConversationIDHeader = "X-Grok-Conv-Id" - grokFreeCacheNativeToolsJSON = `[{"type":"web_search"},{"type":"x_search"}]` - grokFreeCacheDisabledToolChoice = "none" - grokFreeRolling24hTokenLimit = int64(2_000_000) + grokConversationIDHeader = "X-Grok-Conv-Id" + grokClientToolCacheOptInHeader = "X-Sub2API-Grok-Client-Tool-Cache" + grokFreeCacheNativeToolsJSON = `[{"type":"web_search"},{"type":"x_search"}]` + grokFreeCacheDisabledToolChoice = "none" + grokClientToolCacheOptInExtraKey = "grok_client_tool_cache_enabled" + grokFreeRolling24hTokenLimit = int64(2_000_000) ) // resolveGrokCacheIdentity derives one stable, tenant-isolated routing identity @@ -122,7 +124,7 @@ func applyGrokResponsesCacheIdentity(body, intentSourceBody []byte, identity str // Inspect the pre-sanitization source. patchGrokResponsesBody may remove an // unsupported client tool and its tool_choice; that must not turn an // explicit client tool intent into an eligible native-tool request. - if gjson.GetBytes(intentSourceBody, "tools").Exists() || gjson.GetBytes(intentSourceBody, "tool_choice").Exists() { + if hasGrokResponsesToolIntent(intentSourceBody) { return out, nil } out, err = sjson.SetRawBytes(out, "tools", []byte(grokFreeCacheNativeToolsJSON)) @@ -132,12 +134,55 @@ func applyGrokResponsesCacheIdentity(body, intentSourceBody []byte, identity str return sjson.SetBytes(out, "tool_choice", grokFreeCacheDisabledToolChoice) } +func hasGrokResponsesToolIntent(body []byte) bool { + if gjson.GetBytes(body, "tools").Exists() || gjson.GetBytes(body, "tool_choice").Exists() { + return true + } + input := gjson.GetBytes(body, "input") + if !input.IsArray() { + return false + } + for _, item := range input.Array() { + if strings.TrimSpace(item.Get("type").String()) != "additional_tools" { + continue + } + tools := item.Get("tools") + if !tools.Exists() || !tools.IsArray() || len(tools.Array()) > 0 { + return true + } + } + return false +} + // applyGrokFreeMessagesFunctionToolCacheRoute enables xAI's cache-capable -// mixed-tools route only for the Anthropic Messages bridge and only when the -// selected account is known to be Free. Native tools become eligible under -// auto selection, so callers must not apply this policy to paid accounts or -// other ingress protocols implicitly. +// mixed-tools route only for known Free accounts. Requests that already expose +// search capability can use the route automatically. Pure client-tool requests +// require an account-level opt-in because adding native search changes xAI's +// automatic tool selection behavior (#4486). func applyGrokFreeMessagesFunctionToolCacheRoute(body, intentSourceBody []byte, account *Account, cacheIdentity string) ([]byte, error) { + allowPureClientTools := account != nil && account.getExtraBool(grokClientToolCacheOptInExtraKey) + return applyGrokFreeToolCacheRoute(body, intentSourceBody, account, cacheIdentity, allowPureClientTools, true) +} + +// applyGrokFreeRequestToolCacheRoute also accepts a request-scoped opt-in. The +// sub2api header is consumed locally because buildGrokResponsesRequest only +// forwards the explicitly supported OpenAI-Beta header from downstream. +func applyGrokFreeRequestToolCacheRoute(c *gin.Context, body, intentSourceBody []byte, account *Account, cacheIdentity string) ([]byte, error) { + allowPureClientTools := account != nil && account.getExtraBool(grokClientToolCacheOptInExtraKey) + if !allowPureClientTools && c != nil { + switch strings.ToLower(strings.TrimSpace(c.GetHeader(grokClientToolCacheOptInHeader))) { + case "1", "true", "yes", "on", "prefer-cache": + allowPureClientTools = true + } + } + // A function merely named web_search/x_search is still a client function. + // Native Responses search types are safe to recognize automatically, while + // converting a client function requires the same explicit opt-in as every + // other pure client-tool request (#4486). + return applyGrokFreeToolCacheRoute(body, intentSourceBody, account, cacheIdentity, allowPureClientTools, allowPureClientTools) +} + +func applyGrokFreeToolCacheRoute(body, intentSourceBody []byte, account *Account, cacheIdentity string, allowPureClientTools, allowFunctionSearch bool) ([]byte, error) { if strings.TrimSpace(cacheIdentity) == "" || !isKnownGrokFreeAccount(account) { return body, nil } @@ -146,7 +191,12 @@ func applyGrokFreeMessagesFunctionToolCacheRoute(body, intentSourceBody []byte, if !isGrokFreeCacheFunctionToolIntent(intentTools, intentToolChoice) { return body, nil } - return appendMissingGrokFreeCacheNativeTools(body) + if intentToolChoice.Type == gjson.String && strings.TrimSpace(intentToolChoice.String()) == grokFreeCacheDisabledToolChoice { + // Adding native cache-routing tools cannot change behavior when the + // client has explicitly disabled all tool execution. + return appendGrokFreeCacheNativeToolsWithPolicy(body, true, false) + } + return appendGrokFreeCacheNativeToolsWithPolicy(body, allowPureClientTools, allowFunctionSearch) } func isKnownGrokFreeAccount(account *Account) bool { @@ -231,22 +281,44 @@ func isGrokFreeCacheFunctionToolIntent(tools, toolChoice gjson.Result) bool { return false } for _, tool := range items { - if !tool.IsObject() || strings.TrimSpace(tool.Get("type").String()) != "function" { + if !tool.IsObject() { return false } - // Responses function declarations keep name at the top level. Reject - // Chat Completions' nested function shape and incomplete declarations. - if strings.TrimSpace(tool.Get("name").String()) == "" || tool.Get("function").Exists() { + toolType := strings.TrimSpace(tool.Get("type").String()) + if _, ok := grokResponsesSupportedToolTypes[toolType]; !ok { return false } + if toolType == "function" { + // Responses function declarations keep name at the top level. Reject + // Chat Completions' nested function shape and incomplete declarations. + if strings.TrimSpace(tool.Get("name").String()) == "" || tool.Get("function").Exists() { + return false + } + } } if !toolChoice.Exists() { return true } - return toolChoice.Type == gjson.String && strings.TrimSpace(toolChoice.String()) == "auto" + if toolChoice.Type != gjson.String { + return false + } + switch strings.TrimSpace(toolChoice.String()) { + case "auto", grokFreeCacheDisabledToolChoice: + return true + default: + return false + } } func appendMissingGrokFreeCacheNativeTools(body []byte) ([]byte, error) { + return appendGrokFreeCacheNativeTools(body, false) +} + +func appendGrokFreeCacheNativeTools(body []byte, allowPureClientTools bool) ([]byte, error) { + return appendGrokFreeCacheNativeToolsWithPolicy(body, allowPureClientTools, true) +} + +func appendGrokFreeCacheNativeToolsWithPolicy(body []byte, allowPureClientTools, allowFunctionSearch bool) ([]byte, error) { tools := gjson.GetBytes(body, "tools") if !tools.Exists() || !tools.IsArray() { return body, nil @@ -256,9 +328,19 @@ func appendMissingGrokFreeCacheNativeTools(body []byte) ([]byte, error) { if len(items) == 0 { return body, nil } + hasNativeSearch := false + for _, tool := range items { + switch strings.TrimSpace(tool.Get("type").String()) { + case "web_search", "x_search": + hasNativeSearch = true + } + } + if !allowPureClientTools && !allowFunctionSearch && !hasNativeSearch { + return body, nil + } merged := make([]json.RawMessage, 0, len(items)+2) present := make(map[string]bool, 2) - hasFunction := false + hasCompanionTool := false for _, tool := range items { toolType := strings.TrimSpace(tool.Get("type").String()) switch toolType { @@ -269,7 +351,7 @@ func appendMissingGrokFreeCacheNativeTools(body []byte) ([]byte, error) { } // Grok Build may declare search as function tools. Convert to native // entries so Free OAuth stays cache-capable without duplicate names. - if name == "web_search" || name == "x_search" { + if (name == "web_search" || name == "x_search") && allowFunctionSearch { if present[name] { continue } @@ -279,9 +361,17 @@ func appendMissingGrokFreeCacheNativeTools(body []byte) ([]byte, error) { } merged = append(merged, raw) present[name] = true + if allowPureClientTools { + hasCompanionTool = true + } continue } - hasFunction = true + if name == "web_search" || name == "x_search" { + // Keep the client function intact and avoid adding a same-named + // native tool unless conversion was explicitly enabled. + present[name] = true + } + hasCompanionTool = true merged = append(merged, json.RawMessage(tool.Raw)) case "web_search", "x_search": if present[toolType] { @@ -290,17 +380,21 @@ func appendMissingGrokFreeCacheNativeTools(body []byte) ([]byte, error) { merged = append(merged, json.RawMessage(tool.Raw)) present[toolType] = true default: - return body, nil + if _, ok := grokResponsesSupportedToolTypes[toolType]; !ok { + return body, nil + } + hasCompanionTool = true + merged = append(merged, json.RawMessage(tool.Raw)) } } - if !hasFunction { + if !hasCompanionTool { return body, nil } // Only complement missing native search tools when the request already contains // at least one search tool (native or function-form). Pure client function tools // (e.g. view_image) must not trigger injection to avoid biasing model tool // selection (#4486). - if !present["web_search"] && !present["x_search"] { + if !allowPureClientTools && !present["web_search"] && !present["x_search"] { return body, nil } for _, toolType := range []string{"web_search", "x_search"} { diff --git a/backend/internal/service/openai_gateway_grok_cache_test.go b/backend/internal/service/openai_gateway_grok_cache_test.go index 73b1ea5ec1..a6f7d00881 100644 --- a/backend/internal/service/openai_gateway_grok_cache_test.go +++ b/backend/internal/service/openai_gateway_grok_cache_test.go @@ -249,6 +249,117 @@ func TestApplyGrokCacheIdentityAppendsNativeToolsWhenSearchPresent(t *testing.T) require.Equal(t, "x_search", tools[2].Get("type").String()) } +func TestGrokFreeClientToolCacheAccountOptIn(t *testing.T) { + account := healthyGrokOAuthGatewayTestAccount(9011, "access-token") + account.Credentials["subscription_tier"] = "free" + account.Extra = map[string]any{grokClientToolCacheOptInExtraKey: true} + intentBody := []byte(`{"model":"grok","tools":[{"type":"function","name":"view_image","parameters":{"type":"object"}},{"type":"function","name":"read_file","parameters":{"type":"object"}}],"tool_choice":"auto"}`) + + body, err := applyGrokResponsesCacheIdentity(intentBody, intentBody, "isolated-id", true) + require.NoError(t, err) + body, err = applyGrokFreeMessagesFunctionToolCacheRoute(body, intentBody, account, "isolated-id") + require.NoError(t, err) + + tools := gjson.GetBytes(body, "tools").Array() + require.Len(t, tools, 4) + require.Equal(t, "view_image", tools[0].Get("name").String()) + require.Equal(t, "read_file", tools[1].Get("name").String()) + require.Equal(t, "web_search", tools[2].Get("type").String()) + require.Equal(t, "x_search", tools[3].Get("type").String()) +} + +func TestGrokFreeClientToolCacheRequestOptIn(t *testing.T) { + account := healthyGrokOAuthGatewayTestAccount(9014, "access-token") + account.Credentials["subscription_tier"] = "free" + c := newGrokCacheTestContext(9014) + c.Request.Header.Set(grokClientToolCacheOptInHeader, "prefer-cache") + intentBody := []byte(`{"model":"grok","tools":[{"type":"function","name":"view_image","parameters":{"type":"object"}}],"tool_choice":"auto"}`) + + body, err := applyGrokResponsesCacheIdentity(intentBody, intentBody, "isolated-id", true) + require.NoError(t, err) + body, err = applyGrokFreeRequestToolCacheRoute(c, body, intentBody, account, "isolated-id") + require.NoError(t, err) + + tools := gjson.GetBytes(body, "tools").Array() + require.Len(t, tools, 3) + require.Equal(t, "view_image", tools[0].Get("name").String()) + require.Equal(t, "web_search", tools[1].Get("type").String()) + require.Equal(t, "x_search", tools[2].Get("type").String()) +} + +func TestGrokFreeRequestClientSearchFunctionRequiresOptIn(t *testing.T) { + account := healthyGrokOAuthGatewayTestAccount(9015, "access-token") + account.Credentials["subscription_tier"] = "free" + c := newGrokCacheTestContext(9015) + body := []byte(`{"model":"grok","tools":[{"type":"function","name":"view_image","parameters":{"type":"object"}},{"type":"function","name":"web_search","parameters":{"type":"object"}}],"tool_choice":"auto"}`) + + patched, err := applyGrokFreeRequestToolCacheRoute(c, body, body, account, "isolated-id") + + require.NoError(t, err) + require.JSONEq(t, string(body), string(patched)) +} + +func TestGrokFreeRequestToolChoiceNoneUsesSafeCacheRoute(t *testing.T) { + account := healthyGrokOAuthGatewayTestAccount(9016, "access-token") + account.Credentials["subscription_tier"] = "free" + c := newGrokCacheTestContext(9016) + body := []byte(`{"model":"grok","tools":[{"type":"function","name":"view_image","parameters":{"type":"object"}}],"tool_choice":"none"}`) + + patched, err := applyGrokFreeRequestToolCacheRoute(c, body, body, account, "isolated-id") + + require.NoError(t, err) + tools := gjson.GetBytes(patched, "tools").Array() + require.Len(t, tools, 3) + require.Equal(t, "view_image", tools[0].Get("name").String()) + require.Equal(t, "web_search", tools[1].Get("type").String()) + require.Equal(t, "x_search", tools[2].Get("type").String()) + require.Equal(t, "none", gjson.GetBytes(patched, "tool_choice").String()) +} + +func TestApplyGrokCacheIdentityRecognizesResponsesLiteAdditionalTools(t *testing.T) { + intentBody := []byte(`{"model":"grok","input":[{"type":"additional_tools","tools":[{"type":"function","name":"lookup","parameters":{"type":"object"}}]},{"type":"message","role":"user","content":"hello"}]}`) + patchedBody := []byte(`{"model":"grok-4.5","tools":[{"type":"function","name":"lookup","parameters":{"type":"object"}}],"input":[{"type":"message","role":"user","content":"hello"}]}`) + + patched, err := applyGrokResponsesCacheIdentity(patchedBody, intentBody, "isolated-id", true) + + require.NoError(t, err) + tools := gjson.GetBytes(patched, "tools").Array() + require.Len(t, tools, 1) + require.Equal(t, "lookup", tools[0].Get("name").String()) + require.False(t, gjson.GetBytes(patched, "tool_choice").Exists()) + require.Equal(t, "isolated-id", gjson.GetBytes(patched, "prompt_cache_key").String()) +} + +func TestGrokFreeCacheRoutePreservesMixedSupportedToolsWithSearchIntent(t *testing.T) { + account := healthyGrokOAuthGatewayTestAccount(9012, "access-token") + account.Credentials["subscription_tier"] = "free" + intentBody := []byte(`{"model":"grok","tools":[{"type":"function","name":"view_image","parameters":{"type":"object"}},{"type":"shell"},{"type":"web_search"}],"tool_choice":"auto"}`) + + body, err := applyGrokResponsesCacheIdentity(intentBody, intentBody, "isolated-id", true) + require.NoError(t, err) + body, err = applyGrokFreeMessagesFunctionToolCacheRoute(body, intentBody, account, "isolated-id") + require.NoError(t, err) + + tools := gjson.GetBytes(body, "tools").Array() + require.Len(t, tools, 4) + require.Equal(t, "function", tools[0].Get("type").String()) + require.Equal(t, "shell", tools[1].Get("type").String()) + require.Equal(t, "web_search", tools[2].Get("type").String()) + require.Equal(t, "x_search", tools[3].Get("type").String()) +} + +func TestGrokClientToolCacheOptInDoesNotOverridePaidTier(t *testing.T) { + account := healthyGrokOAuthGatewayTestAccount(9013, "access-token") + account.Credentials["subscription_tier"] = "supergrok" + account.Extra = map[string]any{grokClientToolCacheOptInExtraKey: true} + body := []byte(`{"model":"grok","tools":[{"type":"function","name":"view_image","parameters":{"type":"object"}}],"tool_choice":"auto"}`) + + patched, err := applyGrokFreeMessagesFunctionToolCacheRoute(body, body, account, "isolated-id") + + require.NoError(t, err) + require.JSONEq(t, string(body), string(patched)) +} + func TestApplyGrokCacheIdentityRequiresPatchedFunctionTools(t *testing.T) { account := healthyGrokOAuthGatewayTestAccount(902, "access-token") account.Credentials["subscription_tier"] = "free"