From f041b5cf0a4ad947e87ce6859e31a4fc7e2b0060 Mon Sep 17 00:00:00 2001 From: superman2003 <2112076433zcr@gmail.com> Date: Wed, 15 Jul 2026 13:05:13 +0800 Subject: [PATCH] fix(grok): cache Free Messages function tools --- .../service/openai_gateway_grok_cache.go | 121 ++++++++++- .../service/openai_gateway_grok_cache_test.go | 188 +++++++++++++++++- .../service/openai_gateway_grok_test.go | 67 +++++++ .../service/openai_gateway_messages.go | 11 +- 4 files changed, 375 insertions(+), 12 deletions(-) diff --git a/backend/internal/service/openai_gateway_grok_cache.go b/backend/internal/service/openai_gateway_grok_cache.go index 1d689bce8a..bbdd5be44a 100644 --- a/backend/internal/service/openai_gateway_grok_cache.go +++ b/backend/internal/service/openai_gateway_grok_cache.go @@ -1,6 +1,7 @@ package service import ( + "encoding/json" "fmt" "net/http" "strings" @@ -100,8 +101,8 @@ func isGrokRequestContext(c *gin.Context) bool { // Free OAuth requests without native search tools are routed by xAI to the // non-cacheable build-free model. For otherwise tool-free requests, add the // native tools with tool_choice=none: this selects the cache-capable tier -// without allowing an actual search. Any explicit client tools or tool_choice -// disable this augmentation so client function-calling semantics stay intact. +// without allowing an actual search. Explicit client tools are handled by the +// narrower Messages-only mixed-tools policy below. func applyGrokResponsesCacheIdentity(body, intentSourceBody []byte, identity string, injectFreeTierTools bool) ([]byte, error) { identity = strings.TrimSpace(identity) if identity == "" { @@ -130,6 +131,122 @@ func applyGrokResponsesCacheIdentity(body, intentSourceBody []byte, identity str return sjson.SetBytes(out, "tool_choice", grokFreeCacheDisabledToolChoice) } +// 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. +func applyGrokFreeMessagesFunctionToolCacheRoute(body, intentSourceBody []byte, account *Account, cacheIdentity string) ([]byte, error) { + if strings.TrimSpace(cacheIdentity) == "" || !isKnownGrokFreeAccount(account) { + return body, nil + } + intentTools := gjson.GetBytes(intentSourceBody, "tools") + intentToolChoice := gjson.GetBytes(intentSourceBody, "tool_choice") + if !isGrokFreeCacheFunctionToolIntent(intentTools, intentToolChoice) { + return body, nil + } + return appendMissingGrokFreeCacheNativeTools(body) +} + +func isKnownGrokFreeAccount(account *Account) bool { + if account == nil || !account.IsGrokOAuth() { + return false + } + if billing, err := grokBillingSnapshotFromExtra(account.Extra); err == nil && billing != nil { + if tier := strings.TrimSpace(billing.Plan); tier != "" { + return isGrokFreeSubscriptionTier(tier) + } + } + if snapshot, err := grokQuotaSnapshotFromExtra(account.Extra); err == nil && snapshot != nil { + if tier := strings.TrimSpace(snapshot.SubscriptionTier); tier != "" { + return isGrokFreeSubscriptionTier(tier) + } + } + return isGrokFreeSubscriptionTier(account.GetCredential("subscription_tier")) +} + +func isGrokFreeSubscriptionTier(tier string) bool { + switch strings.ToLower(strings.TrimSpace(tier)) { + case "free", "grok-free", "grok_free", "free-tier", "free_tier": + return true + default: + return false + } +} + +func isGrokFreeCacheFunctionToolIntent(tools, toolChoice gjson.Result) bool { + if !tools.IsArray() { + return false + } + items := tools.Array() + if len(items) == 0 { + return false + } + for _, tool := range items { + if !tool.IsObject() || strings.TrimSpace(tool.Get("type").String()) != "function" { + 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() { + return false + } + } + if !toolChoice.Exists() { + return true + } + return toolChoice.Type == gjson.String && strings.TrimSpace(toolChoice.String()) == "auto" +} + +func appendMissingGrokFreeCacheNativeTools(body []byte) ([]byte, error) { + tools := gjson.GetBytes(body, "tools") + if !tools.Exists() || !tools.IsArray() { + return body, nil + } + + items := tools.Array() + if len(items) == 0 { + return body, nil + } + merged := make([]json.RawMessage, 0, len(items)+2) + present := make(map[string]bool, 2) + hasFunction := false + for _, tool := range items { + toolType := strings.TrimSpace(tool.Get("type").String()) + switch toolType { + case "function": + if !tool.IsObject() || strings.TrimSpace(tool.Get("name").String()) == "" || tool.Get("function").Exists() { + return body, nil + } + hasFunction = true + case "web_search", "x_search": + // Native tools may already be present when this helper is retried. + default: + return body, nil + } + merged = append(merged, json.RawMessage(tool.Raw)) + present[toolType] = true + } + if !hasFunction { + return body, nil + } + for _, toolType := range []string{"web_search", "x_search"} { + if present[toolType] { + continue + } + raw, err := json.Marshal(map[string]string{"type": toolType}) + if err != nil { + return nil, err + } + merged = append(merged, raw) + } + encoded, err := json.Marshal(merged) + if err != nil { + return nil, err + } + return sjson.SetRawBytes(body, "tools", encoded) +} + // applyGrokCacheHeaders applies the documented Chat Completions conversation // routing header. The request is built from a fresh header map, so client // supplied x-grok headers cannot override this server-derived value. diff --git a/backend/internal/service/openai_gateway_grok_cache_test.go b/backend/internal/service/openai_gateway_grok_cache_test.go index 42abfc5800..4f811a4530 100644 --- a/backend/internal/service/openai_gateway_grok_cache_test.go +++ b/backend/internal/service/openai_gateway_grok_cache_test.go @@ -197,15 +197,161 @@ func TestApplyGrokCacheIdentityWritesResponsesBodyAndHeader(t *testing.T) { require.False(t, gjson.GetBytes(unscopedBody, "tool_choice").Exists()) } -func TestApplyGrokCacheIdentityPreservesExplicitClientToolFields(t *testing.T) { +func TestApplyGrokCacheIdentityAppendsNativeToolsToResponseFunctions(t *testing.T) { + account := healthyGrokOAuthGatewayTestAccount(901, "access-token") + account.Credentials["subscription_tier"] = " FREE " + tests := []struct { + name string + toolChoiceJSON string + wantChoice bool + }{ + {name: "missing tool choice"}, + {name: "automatic tool choice", toolChoiceJSON: `,"tool_choice":"auto"`, wantChoice: true}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + intentBody := []byte(`{"model":"grok","tools":[{"type":"function","name":"lookup","description":"look up a value","parameters":{"type":"object"}},{"type":"function","name":"save","parameters":{"type":"object"}}]` + tt.toolChoiceJSON + `}`) + body, err := applyGrokResponsesCacheIdentity(intentBody, intentBody, "isolated-id", true) + require.NoError(t, err) + body, err = applyGrokFreeMessagesFunctionToolCacheRoute(body, intentBody, account, "isolated-id") + + require.NoError(t, err) + require.Equal(t, "isolated-id", gjson.GetBytes(body, "prompt_cache_key").String()) + tools := gjson.GetBytes(body, "tools").Array() + require.Len(t, tools, 4) + require.Equal(t, "function", tools[0].Get("type").String()) + require.Equal(t, "lookup", tools[0].Get("name").String()) + require.Equal(t, "function", tools[1].Get("type").String()) + require.Equal(t, "save", 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()) + require.Equal(t, tt.wantChoice, gjson.GetBytes(body, "tool_choice").Exists()) + if tt.wantChoice { + require.Equal(t, "auto", gjson.GetBytes(body, "tool_choice").String()) + } + + second, err := applyGrokResponsesCacheIdentity(body, intentBody, "isolated-id", true) + require.NoError(t, err) + second, err = applyGrokFreeMessagesFunctionToolCacheRoute(second, intentBody, account, "isolated-id") + require.NoError(t, err) + require.JSONEq(t, string(body), string(second), "native tools must not be duplicated") + require.Len(t, gjson.GetBytes(second, "tools").Array(), 4) + }) + } +} + +func TestApplyGrokCacheIdentityRequiresPatchedFunctionTools(t *testing.T) { + account := healthyGrokOAuthGatewayTestAccount(902, "access-token") + account.Credentials["subscription_tier"] = "free" + intentBody := []byte(`{"model":"grok","tools":[{"type":"function","name":"lookup"}],"tool_choice":"auto"}`) + tests := []struct { + name string + patchedBody string + }{ + {name: "missing tools", patchedBody: `{"model":"grok-4.5"}`}, + {name: "empty tools", patchedBody: `{"model":"grok-4.5","tools":[]}`}, + {name: "native tools only", patchedBody: `{"model":"grok-4.5","tools":[{"type":"web_search"}]}`}, + {name: "unexpected patched tool", patchedBody: `{"model":"grok-4.5","tools":[{"type":"function","name":"lookup"},{"type":"mcp","name":"server"}]}`}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + beforeTools := gjson.Get(tt.patchedBody, "tools") + body, err := applyGrokResponsesCacheIdentity([]byte(tt.patchedBody), intentBody, "isolated-id", true) + require.NoError(t, err) + body, err = applyGrokFreeMessagesFunctionToolCacheRoute(body, intentBody, account, "isolated-id") + + require.NoError(t, err) + require.Equal(t, "isolated-id", gjson.GetBytes(body, "prompt_cache_key").String()) + afterTools := gjson.GetBytes(body, "tools") + require.Equal(t, beforeTools.Exists(), afterTools.Exists()) + require.Equal(t, beforeTools.Raw, afterTools.Raw) + }) + } +} + +func TestGrokFreeMessagesFunctionToolCacheRouteRequiresKnownFreeTier(t *testing.T) { + intentBody := []byte(`{"model":"grok","tools":[{"type":"function","name":"lookup"}],"tool_choice":"auto"}`) + tests := []struct { + name string + account *Account + wantMix bool + }{ + { + name: "free credential tier", + account: func() *Account { + a := healthyGrokOAuthGatewayTestAccount(910, "access-token") + a.Credentials["subscription_tier"] = "free" + return a + }(), + wantMix: true, + }, + { + name: "free billing tier", + account: func() *Account { + a := healthyGrokOAuthGatewayTestAccount(911, "access-token") + a.Extra = map[string]any{grokBillingExtraKey: map[string]any{"plan": "FREE"}} + return a + }(), + wantMix: true, + }, + { + name: "supergrok remains unchanged", + account: func() *Account { + a := healthyGrokOAuthGatewayTestAccount(912, "access-token") + a.Credentials["subscription_tier"] = "supergrok" + return a + }(), + }, + { + name: "unknown tier remains unchanged", + account: healthyGrokOAuthGatewayTestAccount(913, "access-token"), + }, + { + name: "api key remains unchanged", + account: &Account{ + ID: 914, + Platform: PlatformGrok, + Type: AccountTypeAPIKey, + }, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + body, err := applyGrokFreeMessagesFunctionToolCacheRoute(intentBody, intentBody, tt.account, "isolated-id") + + require.NoError(t, err) + tools := gjson.GetBytes(body, "tools").Array() + if tt.wantMix { + require.Len(t, tools, 3) + require.Equal(t, "web_search", tools[1].Get("type").String()) + require.Equal(t, "x_search", tools[2].Get("type").String()) + return + } + require.Len(t, tools, 1) + }) + } +} + +func TestGrokFreeMessagesFunctionToolCacheRouteRequiresIdentity(t *testing.T) { + account := healthyGrokOAuthGatewayTestAccount(915, "access-token") + account.Credentials["subscription_tier"] = "free" + body := []byte(`{"model":"grok","tools":[{"type":"function","name":"lookup"}],"tool_choice":"auto"}`) + + patched, err := applyGrokFreeMessagesFunctionToolCacheRoute(body, body, account, "") + + require.NoError(t, err) + require.JSONEq(t, string(body), string(patched)) + require.Len(t, gjson.GetBytes(patched, "tools").Array(), 1) +} + +func TestApplyGrokCacheIdentityPreservesIneligibleClientToolFields(t *testing.T) { tests := []struct { name string body string }{ - { - name: "tools only", - body: `{"model":"grok","tools":[{"type":"function","name":"lookup","parameters":{"type":"object"}}]}`, - }, { name: "empty tools array", body: `{"model":"grok","tools":[]}`, @@ -223,13 +369,41 @@ func TestApplyGrokCacheIdentityPreservesExplicitClientToolFields(t *testing.T) { body: `{"model":"grok","tool_choice":null}`, }, { - name: "both fields", + name: "native tool with auto choice", body: `{"model":"grok","tools":[{"type":"web_search"}],"tool_choice":"auto"}`, }, { - name: "unsupported tool", + name: "function with required choice", + body: `{"model":"grok","tools":[{"type":"function","name":"lookup"}],"tool_choice":"required"}`, + }, + { + name: "function with none choice", + body: `{"model":"grok","tools":[{"type":"function","name":"lookup"}],"tool_choice":"none"}`, + }, + { + name: "function with specific choice", + body: `{"model":"grok","tools":[{"type":"function","name":"lookup"}],"tool_choice":{"type":"function","name":"lookup"}}`, + }, + { + name: "function with object auto choice", + body: `{"model":"grok","tools":[{"type":"function","name":"lookup"}],"tool_choice":{"type":"auto"}}`, + }, + { + name: "function mixed with unsupported tool", + body: `{"model":"grok","tools":[{"type":"function","name":"lookup"},{"type":"namespace","name":"client_tools"}],"tool_choice":"auto"}`, + }, + { + name: "unsupported tool only", body: `{"model":"grok","tools":[{"type":"namespace","name":"client_tools"}]}`, }, + { + name: "chat completions function shape", + body: `{"model":"grok","tools":[{"type":"function","function":{"name":"lookup","parameters":{"type":"object"}}}],"tool_choice":"auto"}`, + }, + { + name: "incomplete responses function", + body: `{"model":"grok","tools":[{"type":"function","parameters":{"type":"object"}}]}`, + }, } for _, tt := range tests { diff --git a/backend/internal/service/openai_gateway_grok_test.go b/backend/internal/service/openai_gateway_grok_test.go index d93415b2e0..60997b396f 100644 --- a/backend/internal/service/openai_gateway_grok_test.go +++ b/backend/internal/service/openai_gateway_grok_test.go @@ -1458,6 +1458,73 @@ func TestForwardAsAnthropicForGrokUsesXAIResponses(t *testing.T) { require.Contains(t, recorder.Body.String(), "ok") } +func TestForwardAsAnthropicForGrokFunctionToolUsesCacheCapableMixedRoute(t *testing.T) { + gin.SetMode(gin.TestMode) + + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + body := []byte(`{ + "model":"grok","max_tokens":32,"stream":false, + "messages":[{"role":"user","content":"look up alpha"}], + "tools":[{"name":"lookup","description":"look up a key","input_schema":{"type":"object","properties":{"key":{"type":"string"}},"required":["key"]}}], + "tool_choice":{"type":"auto"} + }`) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/messages", bytes.NewReader(body)) + c.Set("api_key", &APIKey{ID: 5403}) + + account := healthyGrokOAuthGatewayTestAccount(58, "access-token") + account.Credentials["subscription_tier"] = "free" + repo := &grokQuotaAccountRepo{ + mockAccountRepoForPlatform: &mockAccountRepoForPlatform{ + accountsByID: map[int64]*Account{58: account}, + }, + } + responseBody := strings.Join([]string{ + `data: {"type":"response.completed","response":{"id":"resp_grok_function","object":"response","model":"grok-4.5","status":"completed","output":[{"type":"function_call","id":"fc_lookup","call_id":"call_lookup","name":"lookup","arguments":"{\"key\":\"alpha\"}","status":"completed"}],"usage":{"input_tokens":7000,"output_tokens":2,"total_tokens":7002,"input_tokens_details":{"cached_tokens":6144}}}}`, + "", + "data: [DONE]", + "", + }, "\n") + upstream := &httpUpstreamRecorder{resp: &http.Response{ + StatusCode: http.StatusOK, + Header: http.Header{"Content-Type": []string{"text/event-stream"}}, + Body: io.NopCloser(strings.NewReader(responseBody)), + }} + svc := &OpenAIGatewayService{ + httpUpstream: upstream, + grokTokenProvider: NewGrokTokenProvider(repo, nil), + accountRepo: repo, + } + + result, err := svc.ForwardAsAnthropic(context.Background(), c, account, body, "", "") + + require.NoError(t, err) + require.NotNil(t, result) + require.Equal(t, xai.DefaultCLIBaseURL+"/responses", upstream.lastReq.URL.String()) + identity := gjson.GetBytes(upstream.lastBody, "prompt_cache_key").String() + require.NotEmpty(t, identity) + require.Equal(t, identity, upstream.lastReq.Header.Get(grokConversationIDHeader)) + tools := gjson.GetBytes(upstream.lastBody, "tools").Array() + require.Len(t, tools, 3) + require.Equal(t, "function", tools[0].Get("type").String()) + require.Equal(t, "lookup", tools[0].Get("name").String()) + require.Equal(t, "object", tools[0].Get("parameters.type").String()) + require.Equal(t, "web_search", tools[1].Get("type").String()) + require.Equal(t, "x_search", tools[2].Get("type").String()) + require.Equal(t, "auto", gjson.GetBytes(upstream.lastBody, "tool_choice").String()) + + require.Equal(t, 7000, result.Usage.InputTokens) + require.Equal(t, 6144, result.Usage.CacheReadInputTokens) + clientBody := recorder.Body.String() + require.Equal(t, "tool_use", gjson.Get(clientBody, "content.0.type").String()) + require.Equal(t, "call_lookup", gjson.Get(clientBody, "content.0.id").String()) + require.Equal(t, "lookup", gjson.Get(clientBody, "content.0.name").String()) + require.Equal(t, "alpha", gjson.Get(clientBody, "content.0.input.key").String()) + require.Equal(t, "tool_use", gjson.Get(clientBody, "stop_reason").String()) + require.Equal(t, int64(856), gjson.Get(clientBody, "usage.input_tokens").Int()) + require.Equal(t, int64(6144), gjson.Get(clientBody, "usage.cache_read_input_tokens").Int()) +} + func TestForwardAsAnthropicForGrokStreamingPreservesCacheUsage(t *testing.T) { gin.SetMode(gin.TestMode) diff --git a/backend/internal/service/openai_gateway_messages.go b/backend/internal/service/openai_gateway_messages.go index cbfd218a22..abf80de75f 100644 --- a/backend/internal/service/openai_gateway_messages.go +++ b/backend/internal/service/openai_gateway_messages.go @@ -245,15 +245,20 @@ func (s *OpenAIGatewayService) ForwardAsAnthropic( responsesBody = updatedBody grokCacheIdentity := "" if account.Platform == PlatformGrok { - grokCacheIdentity = resolveGrokCacheIdentity(c, responsesBody, promptCacheKey, upstreamModel) - patchedBody, patchErr := patchGrokResponsesBody(responsesBody, upstreamModel) + grokIntentBody := responsesBody + grokCacheIdentity = resolveGrokCacheIdentity(c, grokIntentBody, promptCacheKey, upstreamModel) + patchedBody, patchErr := patchGrokResponsesBody(grokIntentBody, upstreamModel) if patchErr != nil { return nil, patchErr } - responsesBody, patchErr = applyGrokResponsesCacheIdentity(patchedBody, responsesBody, grokCacheIdentity, account.IsGrokOAuth()) + responsesBody, patchErr = applyGrokResponsesCacheIdentity(patchedBody, grokIntentBody, grokCacheIdentity, account.IsGrokOAuth()) if patchErr != nil { return nil, fmt.Errorf("apply grok prompt cache identity: %w", patchErr) } + responsesBody, patchErr = applyGrokFreeMessagesFunctionToolCacheRoute(responsesBody, grokIntentBody, account, grokCacheIdentity) + if patchErr != nil { + return nil, fmt.Errorf("apply grok Free function-tool cache route: %w", patchErr) + } } // 5. Get access token