From 026d2a80839286e4992d882d49140c780dcf7a03 Mon Sep 17 00:00:00 2001 From: superman2003 <2112076433zcr@gmail.com> Date: Sun, 19 Jul 2026 14:41:47 +0800 Subject: [PATCH] fix(grok): promote Codex Lite tools for caching --- .../internal/service/openai_gateway_grok.go | 67 ++++++++- .../service/openai_gateway_grok_test.go | 136 +++++++++++++++++- .../internal/service/openai_ws_http_bridge.go | 7 +- .../service/openai_ws_http_bridge_test.go | 72 ++++++++++ 4 files changed, 271 insertions(+), 11 deletions(-) diff --git a/backend/internal/service/openai_gateway_grok.go b/backend/internal/service/openai_gateway_grok.go index 198055f20e..dec30c3bd2 100644 --- a/backend/internal/service/openai_gateway_grok.go +++ b/backend/internal/service/openai_gateway_grok.go @@ -52,18 +52,21 @@ func (s *OpenAIGatewayService) forwardGrokResponses( if isGrokImageGenerationModel(upstreamModel) { return nil, fmt.Errorf("model %s is an image model and is not available on the Responses endpoint; use /v1/images/generations instead", upstreamModel) } - cacheIdentity := resolveGrokCacheIdentity(c, body, "", upstreamModel) patchedBody, err := patchGrokResponsesBody(body, upstreamModel) if err != nil { return nil, err } + // Derive the identity from the request xAI will actually see. This makes + // Codex Responses Lite additional_tools part of the stable tool prefix. + cacheIdentity := resolveGrokCacheIdentity(c, patchedBody, "", upstreamModel) + mixedCacheIntentBody := append([]byte(nil), patchedBody...) patchedBody, err = applyGrokResponsesCacheIdentity(patchedBody, body, cacheIdentity, account.IsGrokOAuth()) if err != nil { return nil, fmt.Errorf("apply grok prompt cache identity: %w", err) } // Free OAuth + client function tools: reuse Messages mixed-tools cache route // (append web_search/x_search so xAI does not force non-cacheable build-free). - patchedBody, err = applyGrokFreeMessagesFunctionToolCacheRoute(patchedBody, body, account, cacheIdentity) + patchedBody, err = applyGrokFreeRequestToolCacheRoute(c, patchedBody, mixedCacheIntentBody, account, cacheIdentity) if err != nil { return nil, fmt.Errorf("apply grok Free function-tool cache route: %w", err) } @@ -374,10 +377,9 @@ func deleteJSONFields(value any, fields map[string]struct{}) bool { } // additional_tools is a Codex/Responses Lite private input carrier. xAI's -// Responses schema accepts ordinary message/function-call input items but -// rejects this carrier before inference with a ModelInput deserialization -// error. Top-level supported tools remain available through the separate -// sanitizeGrokResponsesTools path. +// Responses schema rejects the carrier itself, but accepts supported tools at +// the top level. Preserve top-level order, append newly discovered tools in +// carrier order, then let sanitizeGrokResponsesTools filter unsupported types. func sanitizeGrokResponsesInput(body []byte) ([]byte, error) { if !bytes.Contains(body, []byte(`"additional_tools"`)) { return body, nil @@ -389,8 +391,36 @@ func sanitizeGrokResponsesInput(body []byte) ([]byte, error) { rawItems := input.Array() filtered := make([]json.RawMessage, 0, len(rawItems)) + topLevelTools := gjson.GetBytes(body, "tools") + mergedTools := make([]json.RawMessage, 0) + seenTools := make(map[string]struct{}) + appendTool := func(tool gjson.Result) bool { + key := grokResponsesToolDedupKey(tool) + if _, exists := seenTools[key]; exists { + return false + } + seenTools[key] = struct{}{} + mergedTools = append(mergedTools, json.RawMessage(tool.Raw)) + return true + } + if topLevelTools.IsArray() { + for _, tool := range topLevelTools.Array() { + seenTools[grokResponsesToolDedupKey(tool)] = struct{}{} + mergedTools = append(mergedTools, json.RawMessage(tool.Raw)) + } + } + + promoted := false for _, item := range rawItems { if strings.TrimSpace(item.Get("type").String()) == "additional_tools" { + tools := item.Get("tools") + if tools.IsArray() { + for _, tool := range tools.Array() { + if appendTool(tool) { + promoted = true + } + } + } continue } filtered = append(filtered, json.RawMessage(item.Raw)) @@ -402,7 +432,30 @@ func sanitizeGrokResponsesInput(body []byte) ([]byte, error) { if err != nil { return nil, err } - return sjson.SetRawBytes(body, "input", encoded) + body, err = sjson.SetRawBytes(body, "input", encoded) + if err != nil || !promoted { + return body, err + } + encodedTools, err := json.Marshal(mergedTools) + if err != nil { + return nil, err + } + return sjson.SetRawBytes(body, "tools", encodedTools) +} + +func grokResponsesToolDedupKey(tool gjson.Result) string { + toolType := strings.TrimSpace(tool.Get("type").String()) + if toolType != "" { + if name := strings.TrimSpace(tool.Get("name").String()); name != "" { + return "type:" + toolType + "\x00name:" + name + } + if toolType == "mcp" { + if label := strings.TrimSpace(tool.Get("server_label").String()); label != "" { + return "type:mcp\x00server_label:" + label + } + } + } + return "json:" + normalizeCompatSeedJSON(json.RawMessage(tool.Raw)) } // sanitizeGrokReasoningNullContent 删除 reasoning 项中的 "content": null。 diff --git a/backend/internal/service/openai_gateway_grok_test.go b/backend/internal/service/openai_gateway_grok_test.go index fd7eb09c3d..42d88875db 100644 --- a/backend/internal/service/openai_gateway_grok_test.go +++ b/backend/internal/service/openai_gateway_grok_test.go @@ -206,18 +206,27 @@ func TestPatchGrokResponsesBodyDropsToolChoiceWhenNoSupportedToolsRemain(t *test require.False(t, gjson.GetBytes(patched, "tool_choice").Exists()) } -func TestPatchGrokResponsesBodyDropsCodexAdditionalToolsInputItems(t *testing.T) { +func TestPatchGrokResponsesBodyPromotesCodexAdditionalTools(t *testing.T) { t.Parallel() body := []byte(`{ "model": "grok", + "tools": [ + {"type": "function", "name": "existing", "description": "top-level wins"}, + {"type": "web_search"} + ], + "tool_choice": "auto", "input": [ { "type": "additional_tools", "role": "developer", "tools": [ - {"type": "namespace", "name": "image_gen"}, - {"type": "function", "name": "wait"} + {"type": "function", "name": "existing", "description": "duplicate carrier definition"}, + {"type": "function", "name": "wait"}, + {"type": "web_search"}, + {"type": "shell"}, + {"type": "custom", "name": "apply_patch"}, + {"type": "namespace", "name": "collaboration"} ] }, { @@ -239,12 +248,133 @@ func TestPatchGrokResponsesBodyDropsCodexAdditionalToolsInputItems(t *testing.T) require.Equal(t, "grok-4.5", gjson.GetBytes(patched, "model").String()) require.Equal(t, 2, len(gjson.GetBytes(patched, "input").Array())) require.False(t, gjson.GetBytes(patched, `input.#(type=="additional_tools")`).Exists()) + tools := gjson.GetBytes(patched, "tools").Array() + require.Len(t, tools, 4) + require.Equal(t, "existing", tools[0].Get("name").String()) + require.Equal(t, "top-level wins", tools[0].Get("description").String()) + require.Equal(t, "web_search", tools[1].Get("type").String()) + require.Equal(t, "wait", tools[2].Get("name").String()) + require.Equal(t, "shell", tools[3].Get("type").String()) + require.False(t, gjson.GetBytes(patched, `tools.#(type=="custom")`).Exists()) + require.False(t, gjson.GetBytes(patched, `tools.#(type=="namespace")`).Exists()) + require.Equal(t, "auto", gjson.GetBytes(patched, "tool_choice").String()) require.Equal(t, "developer", gjson.GetBytes(patched, "input.0.role").String()) require.Equal(t, "system prompt", gjson.GetBytes(patched, "input.0.content.0.text").String()) require.Equal(t, "user", gjson.GetBytes(patched, "input.1.role").String()) require.Equal(t, "hello", gjson.GetBytes(patched, "input.1.content.0.text").String()) } +func TestForwardGrokResponsesCodexAdditionalToolsUsesMixedCacheIntent(t *testing.T) { + gin.SetMode(gin.TestMode) + + body := []byte(`{ + "model":"grok", + "stream":false, + "prompt_cache_key":"codex-session", + "input":[ + {"type":"additional_tools","role":"developer","tools":[ + {"type":"function","name":"lookup","description":"look up a key","parameters":{"type":"object"}}, + {"type":"function","name":"web_search","description":"search","parameters":{"type":"object"}}, + {"type":"custom","name":"apply_patch"}, + {"type":"namespace","name":"collaboration"} + ]}, + {"type":"message","role":"user","content":[{"type":"input_text","text":"hello"}]} + ] + }`) + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", bytes.NewReader(body)) + c.Request.Header.Set(grokClientToolCacheOptInHeader, "prefer-cache") + c.Set("api_key", &APIKey{ID: 4501}) + + account := healthyGrokOAuthGatewayTestAccount(4501, "access-token") + account.Credentials["subscription_tier"] = "free" + repo := &grokQuotaAccountRepo{ + mockAccountRepoForPlatform: &mockAccountRepoForPlatform{ + accountsByID: map[int64]*Account{account.ID: account}, + }, + } + upstream := &httpUpstreamRecorder{resp: &http.Response{ + StatusCode: http.StatusOK, + Header: http.Header{"Content-Type": []string{"application/json"}}, + Body: io.NopCloser(strings.NewReader(`{ + "id":"resp_codex_lite","object":"response","model":"grok-4.5","status":"completed", + "output":[],"usage":{"input_tokens":10,"output_tokens":1} + }`)), + }} + svc := &OpenAIGatewayService{ + httpUpstream: upstream, + grokTokenProvider: NewGrokTokenProvider(repo, nil), + accountRepo: repo, + } + + result, err := svc.forwardGrokResponses(context.Background(), c, account, body, "grok", false, time.Now()) + + require.NoError(t, err) + require.NotNil(t, result) + require.Equal(t, "resp_codex_lite", result.ResponseID) + require.False(t, gjson.GetBytes(upstream.lastBody, `input.#(type=="additional_tools")`).Exists()) + 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, "web_search", tools[1].Get("type").String()) + require.Equal(t, "x_search", tools[2].Get("type").String()) + require.False(t, gjson.GetBytes(upstream.lastBody, "tool_choice").Exists()) + require.False(t, gjson.GetBytes(upstream.lastBody, `tools.#(type=="custom")`).Exists()) + require.False(t, gjson.GetBytes(upstream.lastBody, `tools.#(type=="namespace")`).Exists()) + identity := gjson.GetBytes(upstream.lastBody, "prompt_cache_key").String() + require.NotEmpty(t, identity) + require.Equal(t, identity, upstream.lastReq.Header.Get(grokConversationIDHeader)) + require.Empty(t, upstream.lastReq.Header.Get(grokClientToolCacheOptInHeader)) +} + +func TestGrokResponsesCacheIdentityIncludesPromotedCodexTools(t *testing.T) { + c := newGrokCacheTestContext(4503) + lookupBody := []byte(`{"model":"grok","input":[{"type":"additional_tools","tools":[{"type":"function","name":"lookup","parameters":{"type":"object"}}]},{"type":"message","role":"user","content":"same prompt"}]}`) + readBody := []byte(`{"model":"grok","input":[{"type":"additional_tools","tools":[{"type":"function","name":"read_file","parameters":{"type":"object"}}]},{"type":"message","role":"user","content":"same prompt"}]}`) + + patchedLookup, err := patchGrokResponsesBody(lookupBody, "grok-4.5") + require.NoError(t, err) + patchedRead, err := patchGrokResponsesBody(readBody, "grok-4.5") + require.NoError(t, err) + + lookupIdentity := resolveGrokCacheIdentity(c, patchedLookup, "", "grok-4.5") + readIdentity := resolveGrokCacheIdentity(c, patchedRead, "", "grok-4.5") + require.NotEmpty(t, lookupIdentity) + require.NotEmpty(t, readIdentity) + require.NotEqual(t, lookupIdentity, readIdentity) +} + +func TestCodexUnsupportedAdditionalToolsDoNotBecomeToolFreeCacheIntent(t *testing.T) { + body := []byte(`{ + "model":"grok","tool_choice":"auto", + "input":[ + {"type":"additional_tools","role":"developer","tools":[ + {"type":"custom","name":"apply_patch"}, + {"type":"namespace","name":"collaboration"} + ]}, + {"type":"message","role":"user","content":"hello"} + ] + }`) + patched, err := patchGrokResponsesBody(body, "grok-4.5") + require.NoError(t, err) + require.False(t, gjson.GetBytes(patched, "tools").Exists()) + require.False(t, gjson.GetBytes(patched, "tool_choice").Exists()) + + mixedCacheIntent := patched + patched, err = applyGrokResponsesCacheIdentity(patched, body, "isolated-id", true) + require.NoError(t, err) + account := healthyGrokOAuthGatewayTestAccount(4502, "access-token") + account.Credentials["subscription_tier"] = "free" + patched, err = applyGrokFreeRequestToolCacheRoute(nil, patched, mixedCacheIntent, account, "isolated-id") + + require.NoError(t, err) + require.False(t, gjson.GetBytes(patched, "tools").Exists()) + require.False(t, gjson.GetBytes(patched, "tool_choice").Exists()) + require.Equal(t, "isolated-id", gjson.GetBytes(patched, "prompt_cache_key").String()) +} + func TestBuildGrokResponsesRequestUsesAccountBaseURLAndBearerToken(t *testing.T) { t.Setenv(xai.EnvAllowUnsafeURLOverrides, "true") diff --git a/backend/internal/service/openai_ws_http_bridge.go b/backend/internal/service/openai_ws_http_bridge.go index cc7b307982..dd9875d0ad 100644 --- a/backend/internal/service/openai_ws_http_bridge.go +++ b/backend/internal/service/openai_ws_http_bridge.go @@ -196,12 +196,13 @@ func (s *OpenAIGatewayService) proxyOpenAIWSHTTPBridgeTurn( releaseUpstreamCtx() return nil, err } + grokMixedCacheIntentBody := append([]byte(nil), body...) body, err = applyGrokResponsesCacheIdentity(body, grokIntentSourceBody, grokCacheIdentity, account.IsGrokOAuth()) if err != nil { releaseUpstreamCtx() return nil, fmt.Errorf("apply grok prompt cache identity: %w", err) } - body, err = applyGrokFreeMessagesFunctionToolCacheRoute(body, grokIntentSourceBody, account, grokCacheIdentity) + body, err = applyGrokFreeRequestToolCacheRoute(c, body, grokMixedCacheIntentBody, account, grokCacheIdentity) if err != nil { releaseUpstreamCtx() return nil, fmt.Errorf("apply grok Free function-tool cache route: %w", err) @@ -446,6 +447,10 @@ func resolveGrokWSCacheIdentity(c *gin.Context, account *Account, payload []byte return "", err } upstreamModel := resolveGrokWSUpstreamModel(account, body, originalModel) + body, err = patchGrokResponsesBody(body, upstreamModel) + if err != nil { + return "", err + } return resolveGrokCacheIdentity(c, body, "", upstreamModel), nil } diff --git a/backend/internal/service/openai_ws_http_bridge_test.go b/backend/internal/service/openai_ws_http_bridge_test.go index 808a63c6b1..2f77690b22 100644 --- a/backend/internal/service/openai_ws_http_bridge_test.go +++ b/backend/internal/service/openai_ws_http_bridge_test.go @@ -235,6 +235,78 @@ func TestProxyOpenAIWSHTTPBridgeTurnForGrokDefaultsEmptyModelTo45(t *testing.T) require.Len(t, events, 2) } +func TestProxyOpenAIWSHTTPBridgeTurnPromotesCodexAdditionalToolsForMixedCache(t *testing.T) { + gin.SetMode(gin.TestMode) + + upstream := &httpUpstreamRecorder{resp: &http.Response{ + StatusCode: http.StatusOK, + Header: http.Header{"Content-Type": []string{"text/event-stream"}}, + Body: io.NopCloser(strings.NewReader(strings.Join([]string{ + `data: {"type":"response.created","response":{"id":"resp_grok_codex_lite","model":"grok-4.5"}}`, + "", + `data: {"type":"response.completed","response":{"id":"resp_grok_codex_lite","model":"grok-4.5","usage":{"input_tokens":4,"output_tokens":1}}}`, + "", + }, "\n"))), + }} + svc := &OpenAIGatewayService{ + cfg: &config.Config{Gateway: config.GatewayConfig{MaxLineSize: defaultMaxLineSize}}, + httpUpstream: upstream, + } + account := &Account{ + ID: 73, + Platform: PlatformGrok, + Type: AccountTypeOAuth, + Concurrency: 1, + Credentials: map[string]any{ + "base_url": xai.DefaultCLIBaseURL, + "subscription_tier": "free", + }, + } + payload := []byte(`{ + "type":"response.create","generate":true,"model":"grok","stream":true, + "input":[ + {"type":"additional_tools","role":"developer","tools":[ + {"type":"function","name":"lookup","parameters":{"type":"object"}}, + {"type":"function","name":"web_search","parameters":{"type":"object"}}, + {"type":"custom","name":"apply_patch"}, + {"type":"namespace","name":"collaboration"} + ]}, + {"type":"message","role":"user","content":[{"type":"input_text","text":"hello"}]} + ] + }`) + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + c.Request = httptest.NewRequest(http.MethodGet, "/v1/responses", nil) + c.Request.Header.Set(grokClientToolCacheOptInHeader, "prefer-cache") + var events [][]byte + + result, err := svc.proxyOpenAIWSHTTPBridgeTurn( + context.Background(), c, account, "access-token", payload, len(payload), + "grok", "", "", "", "isolated-ws-cache-id", 1, + func(message []byte) error { + events = append(events, append([]byte(nil), message...)) + return nil + }, + ) + + require.NoError(t, err) + require.NotNil(t, result) + require.Len(t, events, 2) + require.False(t, gjson.GetBytes(upstream.lastBody, `input.#(type=="additional_tools")`).Exists()) + 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, "web_search", tools[1].Get("type").String()) + require.Equal(t, "x_search", tools[2].Get("type").String()) + require.False(t, gjson.GetBytes(upstream.lastBody, "tool_choice").Exists()) + require.False(t, gjson.GetBytes(upstream.lastBody, `tools.#(type=="custom")`).Exists()) + require.False(t, gjson.GetBytes(upstream.lastBody, `tools.#(type=="namespace")`).Exists()) + require.Equal(t, "isolated-ws-cache-id", gjson.GetBytes(upstream.lastBody, "prompt_cache_key").String()) + require.Equal(t, "isolated-ws-cache-id", upstream.lastReq.Header.Get(grokConversationIDHeader)) + require.Empty(t, upstream.lastReq.Header.Get(grokClientToolCacheOptInHeader)) +} + func TestProxyResponsesWebSocketFromClientForGrokUsesXAIHTTPBridge(t *testing.T) { gin.SetMode(gin.TestMode)