diff --git a/backend/internal/service/openai_gateway_grok_chat_bridge.go b/backend/internal/service/openai_gateway_grok_chat_bridge.go index 0b8f0e2d83..baf1d8a592 100644 --- a/backend/internal/service/openai_gateway_grok_chat_bridge.go +++ b/backend/internal/service/openai_gateway_grok_chat_bridge.go @@ -33,6 +33,7 @@ var grokChatResponsesBridgeTopLevelFields = map[string]struct{}{ "tool_choice": {}, "functions": {}, "function_call": {}, + "parallel_tool_calls": {}, } // grokChatResponsesBridgeEligibility deliberately accepts only request shapes @@ -50,13 +51,22 @@ func grokChatResponsesBridgeEligibility(body []byte) (bool, string) { return false, "unsupported_" + field } } - for _, field := range []string{"tools", "functions"} { - if raw, exists := root[field]; exists && !grokChatNullOrEmptyArray(raw) { - return false, "unsupported_" + field + if raw, exists := root["tools"]; exists { + if ok, reason := grokChatFunctionDeclarationsBridgeable(raw); !ok { + return false, reason } } - if raw, exists := root["tool_choice"]; exists && !grokChatNullOrNone(raw) { - return false, "unsupported_tool_choice" + if raw, exists := root["functions"]; exists && !grokChatNullOrEmptyArray(raw) { + return false, "unsupported_functions" + } + if raw, exists := root["tool_choice"]; exists { + if ok, reason := grokChatToolChoiceBridgeable(raw); !ok { + return false, reason + } + var choice string + if json.Unmarshal(raw, &choice) == nil && choice == "required" && !grokChatHasFunctionDeclarations(root) { + return false, "required_tool_choice_without_tools" + } } if raw, exists := root["function_call"]; exists && !grokChatNullOrNone(raw) { return false, "unsupported_function_call" @@ -78,6 +88,12 @@ func grokChatResponsesBridgeEligibility(body []byte) (bool, string) { return false, "invalid_stream" } } + if raw, ok := root["parallel_tool_calls"]; ok { + var parallelToolCalls *bool + if json.Unmarshal(raw, ¶llelToolCalls) != nil || parallelToolCalls == nil { + return false, "invalid_parallel_tool_calls" + } + } if raw, ok := root["stream_options"]; ok { var options map[string]json.RawMessage if json.Unmarshal(raw, &options) != nil || options == nil { @@ -128,42 +144,238 @@ func grokChatResponsesBridgeEligibility(body []byte) (bool, string) { return false, "invalid_messages" } for _, message := range messages { - for field := range message { - if field != "role" && field != "content" { - return false, "unsafe_message_field_" + field - } - } var role string if raw, exists := message["role"]; !exists || json.Unmarshal(raw, &role) != nil { return false, "invalid_message_role" } switch role { - case "system", "user", "assistant": + case "system", "user": + if ok, reason := grokChatMessageFieldsBridgeable(message, "role", "content"); !ok { + return false, reason + } + raw, exists := message["content"] + if !exists { + return false, "non_text_message_content" + } + if ok, reason := grokChatRequiredMessageContentBridgeable(raw); !ok { + return false, reason + } + case "assistant": + if ok, reason := grokChatMessageFieldsBridgeable(message, "role", "content", "tool_calls"); !ok { + return false, reason + } + toolCallCount := 0 + if raw, exists := message["tool_calls"]; exists { + var reason string + toolCallCount, reason = grokChatAssistantToolCallsBridgeable(raw) + if reason != "" { + return false, reason + } + } + raw, hasContent := message["content"] + if !hasContent || strings.TrimSpace(string(raw)) == "null" { + if toolCallCount == 0 { + return false, "non_text_message_content" + } + continue + } + var content string + if json.Unmarshal(raw, &content) == nil { + if strings.TrimSpace(content) == "" && toolCallCount == 0 { + return false, "empty_message_content" + } + continue + } + if ok, reason := grokChatStructuredContentBridgeable(raw); !ok { + return false, reason + } + case "tool": + if ok, reason := grokChatMessageFieldsBridgeable(message, "role", "content", "tool_call_id"); !ok { + return false, reason + } + var callID string + if raw, exists := message["tool_call_id"]; !exists || json.Unmarshal(raw, &callID) != nil || strings.TrimSpace(callID) == "" { + return false, "invalid_tool_call_id" + } + var output string + if raw, exists := message["content"]; !exists || json.Unmarshal(raw, &output) != nil || output == "" { + return false, "invalid_tool_message_content" + } default: return false, "unsupported_message_role_" + role } - raw, exists := message["content"] - if !exists { - return false, "non_text_message_content" - } - var content string - if json.Unmarshal(raw, &content) == nil { - if strings.TrimSpace(content) == "" { - return false, "empty_message_content" - } - continue - } - // Structured content: only allow arrays whose parts are text or - // image_url. These are losslessly convertible to Responses input_text/ - // input_image parts, so the bridge preserves Chat Completions semantics. - if ok, reason := grokChatStructuredContentBridgeable(raw); !ok { - return false, reason - } } return true, "" } +func grokChatFunctionDeclarationsBridgeable(raw json.RawMessage) (bool, string) { + if strings.TrimSpace(string(raw)) == "null" { + return true, "" + } + var declarations []json.RawMessage + if json.Unmarshal(raw, &declarations) != nil { + return false, "invalid_tools" + } + for _, declaration := range declarations { + var tool map[string]json.RawMessage + if json.Unmarshal(declaration, &tool) != nil || tool == nil { + return false, "invalid_tool" + } + for field := range tool { + if field != "type" && field != "function" { + return false, "unsafe_tool_field_" + field + } + } + var toolType string + if rawType, exists := tool["type"]; !exists || json.Unmarshal(rawType, &toolType) != nil || toolType != "function" { + return false, "unsupported_tool_type" + } + functionRaw, exists := tool["function"] + if !exists { + return false, "invalid_tool_function" + } + + var function map[string]json.RawMessage + if json.Unmarshal(functionRaw, &function) != nil || function == nil { + return false, "invalid_tool_function" + } + for field := range function { + switch field { + case "name", "description", "parameters", "strict": + default: + return false, "unsafe_tool_function_field_" + field + } + } + var name string + if rawName, exists := function["name"]; !exists || json.Unmarshal(rawName, &name) != nil || strings.TrimSpace(name) == "" { + return false, "invalid_tool_function_name" + } + if rawDescription, exists := function["description"]; exists { + var description string + if json.Unmarshal(rawDescription, &description) != nil { + return false, "invalid_tool_function_description" + } + } + var parameters map[string]json.RawMessage + if rawParameters, exists := function["parameters"]; !exists || json.Unmarshal(rawParameters, ¶meters) != nil || parameters == nil { + return false, "invalid_tool_function_parameters" + } + if rawStrict, exists := function["strict"]; exists { + var strict bool + if json.Unmarshal(rawStrict, &strict) != nil { + return false, "invalid_tool_function_strict" + } + } + } + return true, "" +} + +func grokChatToolChoiceBridgeable(raw json.RawMessage) (bool, string) { + if strings.TrimSpace(string(raw)) == "null" { + return true, "" + } + var choice string + if json.Unmarshal(raw, &choice) != nil { + return false, "unsupported_tool_choice" + } + switch choice { + case "auto", "none", "required": + return true, "" + default: + return false, "unsupported_tool_choice" + } +} + +func grokChatHasFunctionDeclarations(root map[string]json.RawMessage) bool { + for _, field := range []string{"tools", "functions"} { + raw, exists := root[field] + if !exists { + continue + } + var declarations []json.RawMessage + if json.Unmarshal(raw, &declarations) == nil && len(declarations) > 0 { + return true + } + } + return false +} + +func grokChatMessageFieldsBridgeable(message map[string]json.RawMessage, allowedFields ...string) (bool, string) { + allowed := make(map[string]struct{}, len(allowedFields)) + for _, field := range allowedFields { + allowed[field] = struct{}{} + } + for field := range message { + if _, ok := allowed[field]; !ok { + return false, "unsafe_message_field_" + field + } + } + return true, "" +} + +func grokChatRequiredMessageContentBridgeable(raw json.RawMessage) (bool, string) { + var content string + if json.Unmarshal(raw, &content) == nil { + if strings.TrimSpace(content) == "" { + return false, "empty_message_content" + } + return true, "" + } + // Structured content: only allow arrays whose parts are text or + // image_url. These are losslessly convertible to Responses input_text/ + // input_image parts, so the bridge preserves Chat Completions semantics. + return grokChatStructuredContentBridgeable(raw) +} + +func grokChatAssistantToolCallsBridgeable(raw json.RawMessage) (int, string) { + if strings.TrimSpace(string(raw)) == "null" { + return 0, "" + } + var calls []map[string]json.RawMessage + if json.Unmarshal(raw, &calls) != nil { + return 0, "invalid_tool_calls" + } + for _, call := range calls { + if call == nil { + return 0, "invalid_tool_call" + } + for field := range call { + switch field { + case "id", "type", "function": + default: + return 0, "unsafe_tool_call_field_" + field + } + } + var callID string + if rawID, exists := call["id"]; !exists || json.Unmarshal(rawID, &callID) != nil || strings.TrimSpace(callID) == "" { + return 0, "invalid_tool_call_id" + } + var callType string + if rawType, exists := call["type"]; !exists || json.Unmarshal(rawType, &callType) != nil || callType != "function" { + return 0, "unsupported_tool_call_type" + } + var function map[string]json.RawMessage + if rawFunction, exists := call["function"]; !exists || json.Unmarshal(rawFunction, &function) != nil || function == nil { + return 0, "invalid_tool_call_function" + } + for field := range function { + if field != "name" && field != "arguments" { + return 0, "unsafe_tool_call_function_field_" + field + } + } + var name string + if rawName, exists := function["name"]; !exists || json.Unmarshal(rawName, &name) != nil || strings.TrimSpace(name) == "" { + return 0, "invalid_tool_call_function_name" + } + var arguments string + if rawArguments, exists := function["arguments"]; !exists || json.Unmarshal(rawArguments, &arguments) != nil || !json.Valid([]byte(arguments)) { + return 0, "invalid_tool_call_arguments" + } + } + return len(calls), "" +} + func grokChatStructuredContentBridgeable(raw json.RawMessage) (bool, string) { var parts []map[string]json.RawMessage if err := json.Unmarshal(raw, &parts); err != nil { @@ -199,14 +411,6 @@ func grokChatStructuredContentBridgeable(raw json.RawMessage) (bool, string) { return true, "" } -func grokChatNullOrEmptyArray(raw json.RawMessage) bool { - if strings.TrimSpace(string(raw)) == "null" { - return true - } - var values []json.RawMessage - return json.Unmarshal(raw, &values) == nil && len(values) == 0 -} - func grokChatNullOrNone(raw json.RawMessage) bool { if strings.TrimSpace(string(raw)) == "null" { return true @@ -215,14 +419,30 @@ func grokChatNullOrNone(raw json.RawMessage) bool { return json.Unmarshal(raw, &value) == nil && strings.EqualFold(strings.TrimSpace(value), "none") } -func grokChatCacheIntentBody(body []byte) ([]byte, error) { +func grokChatNullOrEmptyArray(raw json.RawMessage) bool { + if strings.TrimSpace(string(raw)) == "null" { + return true + } + var values []json.RawMessage + return json.Unmarshal(raw, &values) == nil && len(values) == 0 +} + +func grokChatResponsesCacheIntentBody(body []byte) ([]byte, error) { + // An empty Chat tools array is omitted by the Responses converter. In that + // case auto/none is also a semantic no-op and must not suppress the normal + // tool-free cache route. Non-empty converted tools are always kept intact. + if gjson.GetBytes(body, "tools").Exists() { + return append([]byte(nil), body...), nil + } + choice := gjson.GetBytes(body, "tool_choice") + if !choice.Exists() || choice.Type != gjson.String || (choice.String() != "auto" && choice.String() != "none") { + return append([]byte(nil), body...), nil + } var root map[string]json.RawMessage if err := json.Unmarshal(body, &root); err != nil { return nil, err } - for _, field := range []string{"tools", "tool_choice", "functions", "function_call"} { - delete(root, field) - } + delete(root, "tool_choice") return json.Marshal(root) } @@ -277,18 +497,25 @@ func (s *OpenAIGatewayService) forwardGrokChatCompletionsViaResponses( if err != nil { return nil, fmt.Errorf("marshal grok responses bridge request: %w", err) } + // Preserve the converted Responses intent before Grok capability + // sanitization. Cache routing must see the actual client function tools, + // not the nested Chat Completions declarations and not a tool-free copy. + intentBody, err := grokChatResponsesCacheIntentBody(responsesBody) + if err != nil { + return nil, fmt.Errorf("normalize grok responses bridge cache intent: %w", err) + } responsesBody, err = patchGrokResponsesBody(responsesBody, upstreamModel) if err != nil { return nil, fmt.Errorf("patch grok responses bridge request: %w", err) } - intentBody, err := grokChatCacheIntentBody(body) - if err != nil { - return nil, fmt.Errorf("normalize grok responses bridge tool intent: %w", err) - } responsesBody, err = applyGrokResponsesCacheIdentity(responsesBody, intentBody, cacheIdentity, true) if err != nil { return nil, fmt.Errorf("apply grok responses bridge cache identity: %w", err) } + responsesBody, err = applyGrokFreeRequestToolCacheRoute(c, responsesBody, intentBody, account, cacheIdentity) + if err != nil { + return nil, fmt.Errorf("apply grok responses bridge function-tool cache route: %w", err) + } updatedBody, policyErr := s.applyOpenAIFastPolicyToBody(ctx, account, upstreamModel, responsesBody) if policyErr != 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 86fc45e952..6e39bae15d 100644 --- a/backend/internal/service/openai_gateway_grok_chat_bridge_test.go +++ b/backend/internal/service/openai_gateway_grok_chat_bridge_test.go @@ -75,15 +75,60 @@ func TestGrokChatResponsesBridgeEligibility(t *testing.T) { reason: "empty_message_content", }, { - name: "function tools fall back", - body: `{"model":"grok","messages":[{"role":"user","content":"hi"}],"tools":[{"type":"function","function":{"name":"lookup"}}]}`, - reason: "unsupported_tools", + name: "function tools bridge", + body: `{"model":"grok","messages":[{"role":"user","content":"hi"}],"tools":[{"type":"function","function":{"name":"lookup","parameters":{"type":"object"},"strict":false}}]}`, + want: true, }, { - name: "automatic tool choice falls back", - body: `{"model":"grok","messages":[{"role":"user","content":"hi"}],"tools":[],"tool_choice":"auto"}`, + name: "legacy functions fall back", + body: `{"model":"grok","messages":[{"role":"user","content":"hi"}],"functions":[{"name":"lookup","parameters":{"type":"object"}}]}`, + reason: "unsupported_functions", + }, + { + name: "automatic tool choice bridges", + body: `{"model":"grok","messages":[{"role":"user","content":"hi"}],"tools":[],"tool_choice":"auto"}`, + want: true, + }, + { + name: "required tool choice without tools falls back", + body: `{"model":"grok","messages":[{"role":"user","content":"hi"}],"tools":[],"tool_choice":"required"}`, + reason: "required_tool_choice_without_tools", + }, + { + name: "tool history bridges", + 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: "unknown tool type falls back", + body: `{"model":"grok","messages":[{"role":"user","content":"hi"}],"tools":[{"type":"web_search","function":{"name":"lookup","parameters":{"type":"object"}}}]}`, + reason: "unsupported_tool_type", + }, + { + name: "missing tool schema falls back", + body: `{"model":"grok","messages":[{"role":"user","content":"hi"}],"tools":[{"type":"function","function":{"name":"lookup"}}]}`, + reason: "invalid_tool_function_parameters", + }, + { + name: "named tool choice falls back", + body: `{"model":"grok","messages":[{"role":"user","content":"hi"}],"tools":[{"type":"function","function":{"name":"lookup","parameters":{"type":"object"}}}],"tool_choice":{"type":"function","function":{"name":"lookup"}}}`, reason: "unsupported_tool_choice", }, + { + name: "invalid tool call arguments fall back", + body: `{"model":"grok","messages":[{"role":"assistant","content":null,"tool_calls":[{"id":"call_lookup","type":"function","function":{"name":"lookup","arguments":"{"}}]}]}`, + reason: "invalid_tool_call_arguments", + }, + { + name: "tool result without call id falls back", + body: `{"model":"grok","messages":[{"role":"tool","content":"ok"}]}`, + reason: "invalid_tool_call_id", + }, + { + name: "non boolean parallel tool calls falls back", + body: `{"model":"grok","messages":[{"role":"user","content":"hi"}],"parallel_tool_calls":"true"}`, + reason: "invalid_parallel_tool_calls", + }, { name: "reasoning effort falls back because conversion adds summary", body: `{"model":"grok","messages":[{"role":"user","content":"hi"}],"reasoning_effort":"high"}`, @@ -100,9 +145,9 @@ func TestGrokChatResponsesBridgeEligibility(t *testing.T) { reason: "empty_message_content", }, { - name: "tool history falls back", + name: "empty tool history falls back", body: `{"model":"grok","messages":[{"role":"assistant","content":"","tool_calls":[]}]}`, - reason: "unsafe_message_field_tool_calls", + reason: "empty_message_content", }, { name: "unknown field falls back", @@ -183,6 +228,64 @@ func TestForwardGrokChatViaResponsesNonStreamingCachesAndReturnsChat(t *testing. require.NotNil(t, repo.updates[account.ID][grokQuotaSnapshotExtraKey]) } +func TestForwardGrokChatViaResponsesTraeToolHistoryKeepsCacheRoute(t *testing.T) { + gin.SetMode(gin.TestMode) + + firstTurnBody := []byte(`{"model":"grok","messages":[{"role":"system","content":"Be concise"},{"role":"user","content":"Find alpha"}],"stream":false,"prompt_cache_key":"trae-session","tools":[{"type":"function","function":{"name":"lookup","description":"Lookup a value","parameters":{"type":"object","properties":{"key":{"type":"string"}},"required":["key"]},"strict":false}}],"tool_choice":"auto","parallel_tool_calls":true}`) + body := []byte(`{"model":"grok","messages":[{"role":"system","content":"Be concise"},{"role":"user","content":"Find alpha"},{"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"}],"stream":false,"prompt_cache_key":"trae-session","tools":[{"type":"function","function":{"name":"lookup","description":"Lookup a value","parameters":{"type":"object","properties":{"key":{"type":"string"}},"required":["key"]},"strict":false}}],"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.Request.Header.Set(grokClientToolCacheOptInHeader, "prefer-cache") + c.Set("api_key", &APIKey{ID: 7151}) + + account := grokChatBridgeTestAccount(715) + account.Credentials["subscription_tier"] = "free" + repo := &grokQuotaAccountRepo{mockAccountRepoForPlatform: &mockAccountRepoForPlatform{ + accountsByID: map[int64]*Account{account.ID: account}, + }} + upstream := &httpUpstreamRecorder{resp: grokChatBridgeCompletedResponse("resp_grok_chat_trae", 8192)} + 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, extendedTurnIdentity, gjson.GetBytes(upstream.lastBody, "prompt_cache_key").String()) + require.Equal(t, extendedTurnIdentity, 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, "Lookup a value", tools[0].Get("description").String()) + require.Equal(t, "string", tools[0].Get("parameters.properties.key.type").String()) + require.True(t, tools[0].Get("strict").Exists()) + require.False(t, tools[0].Get("strict").Bool()) + 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.True(t, gjson.GetBytes(upstream.lastBody, "parallel_tool_calls").Bool()) + + require.Equal(t, "function_call", gjson.GetBytes(upstream.lastBody, "input.2.type").String()) + require.Equal(t, "call_lookup", gjson.GetBytes(upstream.lastBody, "input.2.call_id").String()) + require.Equal(t, "lookup", gjson.GetBytes(upstream.lastBody, "input.2.name").String()) + require.Equal(t, `{"key":"alpha"}`, gjson.GetBytes(upstream.lastBody, "input.2.arguments").String()) + require.Equal(t, "function_call_output", gjson.GetBytes(upstream.lastBody, "input.3.type").String()) + require.Equal(t, "call_lookup", gjson.GetBytes(upstream.lastBody, "input.3.call_id").String()) + require.Equal(t, `{"value":"ok"}`, gjson.GetBytes(upstream.lastBody, "input.3.output").String()) +} + func TestForwardGrokChatViaResponsesStreamingPropagatesCachedUsage(t *testing.T) { gin.SetMode(gin.TestMode)