fix(grok): preserve Trae cache routing across turns

2026-10-14 12:02:27 +08:00 · 2026-07-19 19:00:08 +08:00
parent c8ff478876
commit 3dcf412b65
2 changed files with 156 additions and 6 deletions
@@ -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
@@ -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(), "<thinking>I should use lookup</thinking>")
+	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)