fix(grok): add guarded client tool cache routing

2026-10-26 16:03:35 +00:00 · 2026-07-19 14:41:38 +08:00
parent d4b9797ff7
commit 5af38ccf41
2 changed files with 226 additions and 21 deletions
@@ -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"} {
@@ -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"