From 3dcf412b6514a617fa7d709a1d3ddf1265a2e1f9 Mon Sep 17 00:00:00 2001 From: superman2003 <2112076433zcr@gmail.com> Date: Sun, 19 Jul 2026 19:00:08 +0800 Subject: [PATCH] fix(grok): preserve Trae cache routing across turns --- .../openai_gateway_grok_chat_bridge.go | 65 +++++++++++-- .../openai_gateway_grok_chat_bridge_test.go | 97 +++++++++++++++++++ 2 files changed, 156 insertions(+), 6 deletions(-) diff --git a/backend/internal/service/openai_gateway_grok_chat_bridge.go b/backend/internal/service/openai_gateway_grok_chat_bridge.go index baf1d8a59..89ebb4049 100644 --- a/backend/internal/service/openai_gateway_grok_chat_bridge.go +++ b/backend/internal/service/openai_gateway_grok_chat_bridge.go @@ -22,18 +22,23 @@ const ( var grokChatResponsesBridgeTopLevelFields = map[string]struct{}{ "model": {}, "messages": {}, + "instructions": {}, "stream": {}, "stream_options": {}, "max_tokens": {}, "max_completion_tokens": {}, "temperature": {}, "top_p": {}, + "stop": {}, + "reasoning_effort": {}, "prompt_cache_key": {}, "tools": {}, "tool_choice": {}, "functions": {}, "function_call": {}, "parallel_tool_calls": {}, + "response_format": {}, + "service_tier": {}, } // grokChatResponsesBridgeEligibility deliberately accepts only request shapes @@ -46,11 +51,33 @@ func grokChatResponsesBridgeEligibility(body []byte) (bool, string) { return false, "invalid_json" } + // These fields have no effect when explicitly set to JSON null. Accepting + // that common SDK representation keeps the request on the bridge path, + // while non-null values remain unsupported because the Responses converter + // cannot preserve their Chat Completions semantics. for _, field := range []string{"stop", "reasoning_effort"} { - if _, exists := root[field]; exists { + if raw, exists := root[field]; exists && !grokChatJSONNull(raw) { return false, "unsupported_" + field } } + if raw, exists := root["instructions"]; exists { + var instructions string + if !grokChatJSONNull(raw) && json.Unmarshal(raw, &instructions) != nil { + return false, "invalid_instructions" + } + } + if raw, exists := root["response_format"]; exists { + var responseFormat map[string]json.RawMessage + if !grokChatJSONNull(raw) && (json.Unmarshal(raw, &responseFormat) != nil || responseFormat == nil) { + return false, "invalid_response_format" + } + } + if raw, exists := root["service_tier"]; exists { + var serviceTier string + if !grokChatJSONNull(raw) && json.Unmarshal(raw, &serviceTier) != nil { + return false, "invalid_service_tier" + } + } if raw, exists := root["tools"]; exists { if ok, reason := grokChatFunctionDeclarationsBridgeable(raw); !ok { return false, reason @@ -161,9 +188,16 @@ func grokChatResponsesBridgeEligibility(body []byte) (bool, string) { return false, reason } case "assistant": - if ok, reason := grokChatMessageFieldsBridgeable(message, "role", "content", "tool_calls"); !ok { + if ok, reason := grokChatMessageFieldsBridgeable(message, "role", "content", "reasoning_content", "tool_calls"); !ok { return false, reason } + reasoningContent := "" + if raw, exists := message["reasoning_content"]; exists { + if !grokChatJSONNull(raw) && json.Unmarshal(raw, &reasoningContent) != nil { + return false, "invalid_reasoning_content" + } + } + hasReasoningContent := strings.TrimSpace(reasoningContent) != "" toolCallCount := 0 if raw, exists := message["tool_calls"]; exists { var reason string @@ -174,20 +208,25 @@ func grokChatResponsesBridgeEligibility(body []byte) (bool, string) { } raw, hasContent := message["content"] if !hasContent || strings.TrimSpace(string(raw)) == "null" { - if toolCallCount == 0 { + if toolCallCount == 0 && !hasReasoningContent { return false, "non_text_message_content" } continue } var content string if json.Unmarshal(raw, &content) == nil { - if strings.TrimSpace(content) == "" && toolCallCount == 0 { + if strings.TrimSpace(content) == "" && toolCallCount == 0 && !hasReasoningContent { return false, "empty_message_content" } continue } if ok, reason := grokChatStructuredContentBridgeable(raw); !ok { - return false, reason + // The converter can still emit a standalone reasoning part when + // an otherwise empty content array accompanies reasoning_content. + // Do not broaden this exception to unsupported/malformed parts. + if !hasReasoningContent || reason != "empty_message_content" { + return false, reason + } } case "tool": if ok, reason := grokChatMessageFieldsBridgeable(message, "role", "content", "tool_call_id"); !ok { @@ -342,11 +381,17 @@ func grokChatAssistantToolCallsBridgeable(raw json.RawMessage) (int, string) { } for field := range call { switch field { - case "id", "type", "function": + case "id", "type", "function", "index": default: return 0, "unsafe_tool_call_field_" + field } } + if rawIndex, exists := call["index"]; exists { + var index *int + if json.Unmarshal(rawIndex, &index) != nil || (index != nil && *index < 0) { + return 0, "invalid_tool_call_index" + } + } var callID string if rawID, exists := call["id"]; !exists || json.Unmarshal(rawID, &callID) != nil || strings.TrimSpace(callID) == "" { return 0, "invalid_tool_call_id" @@ -419,6 +464,10 @@ func grokChatNullOrNone(raw json.RawMessage) bool { return json.Unmarshal(raw, &value) == nil && strings.EqualFold(strings.TrimSpace(value), "none") } +func grokChatJSONNull(raw json.RawMessage) bool { + return strings.TrimSpace(string(raw)) == "null" +} + func grokChatNullOrEmptyArray(raw json.RawMessage) bool { if strings.TrimSpace(string(raw)) == "null" { return true @@ -488,6 +537,10 @@ func (s *OpenAIGatewayService) forwardGrokChatCompletionsViaResponses( } responsesReq.Model = upstreamModel responsesReq.Stream = true + // Keep Chat and native Responses paths aligned for OpenAI-compatible + // service_tier aliases (for example, "fast" -> "priority"). Unknown + // values are omitted by the shared normalizer instead of reaching xAI. + normalizeResponsesRequestServiceTier(responsesReq) // These fields are useful to Codex but are not needed by the Grok CLI // protocol. Keep the bridge request as close as possible to native Grok. responsesReq.Include = nil diff --git a/backend/internal/service/openai_gateway_grok_chat_bridge_test.go b/backend/internal/service/openai_gateway_grok_chat_bridge_test.go index 6e39bae15..188e38f70 100644 --- a/backend/internal/service/openai_gateway_grok_chat_bridge_test.go +++ b/backend/internal/service/openai_gateway_grok_chat_bridge_test.go @@ -39,6 +39,31 @@ func TestGrokChatResponsesBridgeEligibility(t *testing.T) { body: `{"model":"grok","messages":[{"role":"user","content":"hi"}],"stream":true,"stream_options":{"include_usage":true},"max_completion_tokens":256,"temperature":0.2,"top_p":0.9,"prompt_cache_key":"session","tools":[],"functions":null,"tool_choice":"none"}`, want: true, }, + { + name: "Trae compatible fields", + body: `{"model":"grok","messages":[{"role":"user","content":"return json"}],"instructions":"Be concise","response_format":{"type":"json_object"},"service_tier":"fast","stop":null,"reasoning_effort":null}`, + want: true, + }, + { + name: "Trae nullable SDK fields", + body: `{"model":"grok","messages":[{"role":"user","content":"hi"}],"instructions":null,"response_format":null,"service_tier":null,"stop":null,"reasoning_effort":null}`, + want: true, + }, + { + name: "invalid instructions fall back", + body: `{"model":"grok","messages":[{"role":"user","content":"hi"}],"instructions":{"text":"invalid"}}`, + reason: "invalid_instructions", + }, + { + name: "invalid response format falls back", + body: `{"model":"grok","messages":[{"role":"user","content":"hi"}],"response_format":"json_object"}`, + reason: "invalid_response_format", + }, + { + name: "invalid service tier falls back", + body: `{"model":"grok","messages":[{"role":"user","content":"hi"}],"service_tier":1}`, + reason: "invalid_service_tier", + }, { name: "stop falls back", body: `{"model":"grok","messages":[{"role":"user","content":"hi"}],"stop":"done"}`, @@ -99,6 +124,26 @@ func TestGrokChatResponsesBridgeEligibility(t *testing.T) { body: `{"model":"grok","messages":[{"role":"assistant","content":null,"tool_calls":[{"id":"call_lookup","type":"function","function":{"name":"lookup","arguments":"{\"key\":\"alpha\"}"}}]},{"role":"tool","tool_call_id":"call_lookup","content":"{\"value\":\"ok\"}"},{"role":"user","content":"summarize"}],"tools":[{"type":"function","function":{"name":"lookup","parameters":{"type":"object"}}}],"tool_choice":"auto","parallel_tool_calls":true}`, want: true, }, + { + name: "Trae reasoning and indexed tool history bridges", + body: `{"model":"grok","messages":[{"role":"assistant","content":null,"reasoning_content":"I should call lookup","tool_calls":[{"index":0,"id":"call_lookup","type":"function","function":{"name":"lookup","arguments":"{\"key\":\"alpha\"}"}}]},{"role":"tool","tool_call_id":"call_lookup","content":"{\"value\":\"ok\"}"},{"role":"user","content":"summarize"}],"tools":[{"type":"function","function":{"name":"lookup","parameters":{"type":"object"}}}]}`, + want: true, + }, + { + name: "reasoning only assistant history bridges", + body: `{"model":"grok","messages":[{"role":"assistant","reasoning_content":"Prior reasoning"},{"role":"user","content":"continue"}]}`, + want: true, + }, + { + name: "invalid reasoning content falls back", + body: `{"model":"grok","messages":[{"role":"assistant","reasoning_content":{"text":"invalid"}},{"role":"user","content":"continue"}]}`, + reason: "invalid_reasoning_content", + }, + { + name: "negative tool call index falls back", + body: `{"model":"grok","messages":[{"role":"assistant","content":null,"tool_calls":[{"index":-1,"id":"call_lookup","type":"function","function":{"name":"lookup","arguments":"{}"}}]}]}`, + reason: "invalid_tool_call_index", + }, { name: "unknown tool type falls back", body: `{"model":"grok","messages":[{"role":"user","content":"hi"}],"tools":[{"type":"web_search","function":{"name":"lookup","parameters":{"type":"object"}}}]}`, @@ -286,6 +331,58 @@ func TestForwardGrokChatViaResponsesTraeToolHistoryKeepsCacheRoute(t *testing.T) require.Equal(t, `{"value":"ok"}`, gjson.GetBytes(upstream.lastBody, "input.3.output").String()) } +func TestForwardGrokChatViaResponsesTraeCompatibilityFieldsKeepCacheRoute(t *testing.T) { + gin.SetMode(gin.TestMode) + + firstTurnBody := []byte(`{"model":"grok","messages":[{"role":"user","content":"Find alpha"}],"instructions":"Return concise JSON","stream":false,"response_format":{"type":"json_object"},"service_tier":"fast","stop":null,"reasoning_effort":null,"tools":[{"type":"function","function":{"name":"lookup","description":"Lookup a value","parameters":{"type":"object","properties":{"key":{"type":"string"}},"required":["key"]}}}],"tool_choice":"auto","parallel_tool_calls":true}`) + body := []byte(`{"model":"grok","messages":[{"role":"user","content":"Find alpha"},{"role":"assistant","content":null,"reasoning_content":"I should use lookup","tool_calls":[{"index":0,"id":"call_lookup","type":"function","function":{"name":"lookup","arguments":"{\"key\":\"alpha\"}"}}]},{"role":"tool","tool_call_id":"call_lookup","content":"{\"value\":\"ok\"}"},{"role":"user","content":"Summarize"}],"instructions":"Return concise JSON","stream":false,"response_format":{"type":"json_object"},"service_tier":"fast","stop":null,"reasoning_effort":null,"tools":[{"type":"function","function":{"name":"lookup","description":"Lookup a value","parameters":{"type":"object","properties":{"key":{"type":"string"}},"required":["key"]}}}],"tool_choice":"auto","parallel_tool_calls":true}`) + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + c.Request = httptest.NewRequest(http.MethodPost, grokChatRawEndpoint, bytes.NewReader(body)) + c.Set("api_key", &APIKey{ID: 7161}) + + account := grokChatBridgeTestAccount(716) + account.Credentials["subscription_tier"] = "free" + repo := &grokQuotaAccountRepo{mockAccountRepoForPlatform: &mockAccountRepoForPlatform{ + accountsByID: map[int64]*Account{account.ID: account}, + }} + upstream := &httpUpstreamRecorder{resp: grokChatBridgeCompletedResponse("resp_grok_chat_trae_compat", 12288)} + svc := &OpenAIGatewayService{ + httpUpstream: upstream, + grokTokenProvider: NewGrokTokenProvider(repo, nil), + accountRepo: repo, + } + + firstTurnIdentity := resolveGrokCacheIdentity(c, firstTurnBody, "", "grok-4.5") + extendedTurnIdentity := resolveGrokCacheIdentity(c, body, "", "grok-4.5") + require.NotEmpty(t, firstTurnIdentity) + require.Equal(t, firstTurnIdentity, extendedTurnIdentity) + + result, err := svc.ForwardAsChatCompletions(context.Background(), c, account, body, "", "") + require.NoError(t, err) + require.NotNil(t, result) + require.Equal(t, xai.DefaultCLIBaseURL+"/responses", upstream.lastReq.URL.String()) + require.Equal(t, grokChatResponsesEndpoint, result.UpstreamEndpoint) + require.Equal(t, 12288, result.Usage.CacheReadInputTokens) + require.Equal(t, extendedTurnIdentity, gjson.GetBytes(upstream.lastBody, "prompt_cache_key").String()) + require.Equal(t, extendedTurnIdentity, upstream.lastReq.Header.Get(grokConversationIDHeader)) + require.Equal(t, "Return concise JSON", gjson.GetBytes(upstream.lastBody, "instructions").String()) + require.Equal(t, "json_object", gjson.GetBytes(upstream.lastBody, "text.format.type").String()) + require.Equal(t, "priority", gjson.GetBytes(upstream.lastBody, "service_tier").String()) + require.False(t, gjson.GetBytes(upstream.lastBody, "stop").Exists()) + require.False(t, gjson.GetBytes(upstream.lastBody, "reasoning").Exists()) + require.Contains(t, gjson.GetBytes(upstream.lastBody, "input.1.content.0.text").String(), "I should use lookup") + require.Equal(t, "function_call", gjson.GetBytes(upstream.lastBody, "input.2.type").String()) + require.Equal(t, "function_call_output", gjson.GetBytes(upstream.lastBody, "input.3.type").String()) + + tools := gjson.GetBytes(upstream.lastBody, "tools").Array() + require.Len(t, tools, 3) + 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.Equal(t, int64(12288), gjson.Get(recorder.Body.String(), "usage.prompt_tokens_details.cached_tokens").Int()) +} + func TestForwardGrokChatViaResponsesStreamingPropagatesCachedUsage(t *testing.T) { gin.SetMode(gin.TestMode)