diff --git a/backend/internal/service/openai_gateway_chat_completions.go b/backend/internal/service/openai_gateway_chat_completions.go index 171a6c8805..a66489da00 100644 --- a/backend/internal/service/openai_gateway_chat_completions.go +++ b/backend/internal/service/openai_gateway_chat_completions.go @@ -463,6 +463,17 @@ func (s *OpenAIGatewayService) handleChatBufferedStreamingResponse( // When the terminal event has an empty output array, reconstruct from // accumulated delta events so the client receives the full content. acc.SupplementResponseOutput(finalResponse) + if account != nil && account.Platform == PlatformGrok && !hasBillableOpenAIUsage(usage) { + return nil, s.newOpenAIStreamFailoverError( + c, + account, + false, + firstNonEmpty(requestID, resp.Header.Get("xai-request-id")), + nil, + grokMissingUsageMessage, + resp.Header, + ) + } chatResp := apicompat.ResponsesToChatCompletions(finalResponse, originalModel) diff --git a/backend/internal/service/openai_gateway_chat_completions_raw.go b/backend/internal/service/openai_gateway_chat_completions_raw.go index 49cd591829..0aee0fa5f5 100644 --- a/backend/internal/service/openai_gateway_chat_completions_raw.go +++ b/backend/internal/service/openai_gateway_chat_completions_raw.go @@ -214,7 +214,7 @@ func (s *OpenAIGatewayService) forwardAsRawChatCompletions( if clientStream { result, forwardErr = s.streamRawChatCompletions(c, resp, account, originalModel, billingModel, upstreamModel, reasoningEffort, serviceTier, startTime, len(body)) } else { - result, forwardErr = s.bufferRawChatCompletions(c, resp, originalModel, billingModel, upstreamModel, reasoningEffort, serviceTier, startTime) + result, forwardErr = s.bufferRawChatCompletions(c, resp, account, originalModel, billingModel, upstreamModel, reasoningEffort, serviceTier, startTime) } if result != nil { addOpenAIUsage(&result.Usage, bridgeUsage) @@ -409,6 +409,7 @@ func extractCCStreamUsage(payload string) *OpenAIUsage { func (s *OpenAIGatewayService) bufferRawChatCompletions( c *gin.Context, resp *http.Response, + account *Account, originalModel string, billingModel string, upstreamModel string, @@ -430,6 +431,17 @@ func (s *OpenAIGatewayService) bufferRawChatCompletions( if parsedUsage, ok := extractOpenAIUsageFromJSONBytes(respBody); ok { usage = parsedUsage } + if account != nil && account.Platform == PlatformGrok && !hasBillableOpenAIUsage(usage) { + return nil, s.newOpenAIStreamFailoverError( + c, + account, + false, + firstNonEmpty(requestID, resp.Header.Get("xai-request-id")), + nil, + grokMissingUsageMessage, + resp.Header, + ) + } if s.responseHeaderFilter != nil { responseheaders.WriteFilteredHeaders(c.Writer.Header(), resp.Header, s.responseHeaderFilter) diff --git a/backend/internal/service/openai_gateway_chat_completions_raw_test.go b/backend/internal/service/openai_gateway_chat_completions_raw_test.go index 61663e636a..c19066f3fc 100644 --- a/backend/internal/service/openai_gateway_chat_completions_raw_test.go +++ b/backend/internal/service/openai_gateway_chat_completions_raw_test.go @@ -638,7 +638,7 @@ func TestBufferRawChatCompletions_RejectsOversizedResponse(t *testing.T) { svc := &OpenAIGatewayService{cfg: rawChatCompletionsTestConfig()} svc.cfg.Gateway.UpstreamResponseReadMaxBytes = 3 - result, err := svc.bufferRawChatCompletions(c, resp, "gpt-5.4", "gpt-5.4", "gpt-5.4", nil, nil, time.Now()) + result, err := svc.bufferRawChatCompletions(c, resp, rawChatCompletionsTestAccount(), "gpt-5.4", "gpt-5.4", "gpt-5.4", nil, nil, time.Now()) require.ErrorIs(t, err, ErrUpstreamResponseBodyTooLarge) require.Nil(t, result) require.Equal(t, http.StatusBadGateway, rec.Code) 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 06f434341a..abb21906aa 100644 --- a/backend/internal/service/openai_gateway_grok_chat_bridge_test.go +++ b/backend/internal/service/openai_gateway_grok_chat_bridge_test.go @@ -273,6 +273,45 @@ func TestForwardGrokChatViaResponsesNonStreamingCachesAndReturnsChat(t *testing. require.NotNil(t, repo.updates[account.ID][grokQuotaSnapshotExtraKey]) } +func TestForwardGrokChatViaResponsesNonStreamingRejectsCompletedResponseWithoutUsage(t *testing.T) { + gin.SetMode(gin.TestMode) + + body := []byte(`{"model":"grok","messages":[{"role":"user","content":"hi"}],"stream":false,"prompt_cache_key":"stable-session"}`) + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + c.Request = httptest.NewRequest(http.MethodPost, grokChatRawEndpoint, bytes.NewReader(body)) + c.Set("api_key", &APIKey{ID: 7102}) + + account := grokChatBridgeTestAccount(72) + repo := &grokQuotaAccountRepo{mockAccountRepoForPlatform: &mockAccountRepoForPlatform{ + accountsByID: map[int64]*Account{account.ID: account}, + }} + upstreamBody := strings.Join([]string{ + `data: {"type":"response.output_text.delta","sequence_number":0,"delta":"ok"}`, + "", + `data: {"type":"response.completed","sequence_number":1,"response":{"id":"resp_missing_usage","object":"response","model":"grok-4.5","status":"completed","output":[{"type":"message","id":"msg_1","role":"assistant","status":"completed","content":[{"type":"output_text","text":"ok"}]}]}}`, + "", + }, "\n") + upstream := &httpUpstreamRecorder{resp: &http.Response{ + StatusCode: http.StatusOK, + Header: http.Header{"Content-Type": []string{"text/event-stream"}}, + Body: io.NopCloser(strings.NewReader(upstreamBody)), + }} + svc := &OpenAIGatewayService{ + httpUpstream: upstream, + grokTokenProvider: NewGrokTokenProvider(repo, nil), + accountRepo: repo, + } + + result, err := svc.ForwardAsChatCompletions(context.Background(), c, account, body, "", "") + + require.Nil(t, result) + var failoverErr *UpstreamFailoverError + require.ErrorAs(t, err, &failoverErr) + require.False(t, c.Writer.Written(), "an unbillable response must not be committed to the client") + require.Empty(t, recorder.Body.String()) +} + func TestForwardGrokChatViaResponsesCodeBuddyUsesStableConversationHeader(t *testing.T) { gin.SetMode(gin.TestMode) const conversationID = "codebuddy-session-42" diff --git a/backend/internal/service/openai_gateway_grok_test.go b/backend/internal/service/openai_gateway_grok_test.go index d087bd5184..24f87fed40 100644 --- a/backend/internal/service/openai_gateway_grok_test.go +++ b/backend/internal/service/openai_gateway_grok_test.go @@ -1937,6 +1937,41 @@ func TestForwardAsChatCompletionsForGrokAPIKeyUsesConfiguredRawEndpointWithoutOA require.NotEqual(t, grokUpstreamUserAgent, upstream.lastReq.Header.Get("User-Agent")) } +func TestForwardAsChatCompletionsForGrokAPIKeyRejectsNonStreamingResponseWithoutUsage(t *testing.T) { + gin.SetMode(gin.TestMode) + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + body := []byte(`{"model":"grok","messages":[{"role":"user","content":"hi"}],"stream":false}`) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", bytes.NewReader(body)) + c.Request.Header.Set("Content-Type", "application/json") + account := &Account{ + ID: 707, + Platform: PlatformGrok, + Type: AccountTypeAPIKey, + Concurrency: 1, + Credentials: map[string]any{ + "api_key": "third-party-key", + "base_url": "https://grok.example.test/v1", + }, + } + upstream := &httpUpstreamRecorder{resp: &http.Response{ + StatusCode: http.StatusOK, + Header: http.Header{"Content-Type": []string{"application/json"}}, + Body: io.NopCloser(strings.NewReader( + `{"id":"resp_missing_usage","object":"chat.completion","model":"grok-4.5","choices":[{"index":0,"message":{"role":"assistant","content":"ok"},"finish_reason":"stop"}]}`, + )), + }} + svc := &OpenAIGatewayService{cfg: rawChatCompletionsTestConfig(), httpUpstream: upstream} + + result, err := svc.ForwardAsChatCompletions(context.Background(), c, account, body, "", "") + + require.Nil(t, result) + var failoverErr *UpstreamFailoverError + require.ErrorAs(t, err, &failoverErr) + require.False(t, c.Writer.Written(), "an unbillable response must not be committed to the client") + require.Empty(t, recorder.Body.String()) +} + func TestAccountTestServiceGrokAPIKeyUsesXAIResponses(t *testing.T) { gin.SetMode(gin.TestMode) diff --git a/backend/internal/service/openai_gateway_usage_integrity.go b/backend/internal/service/openai_gateway_usage_integrity.go new file mode 100644 index 0000000000..c1b9e1523b --- /dev/null +++ b/backend/internal/service/openai_gateway_usage_integrity.go @@ -0,0 +1,16 @@ +package service + +const grokMissingUsageMessage = "Grok upstream returned a successful response without billable usage" + +// hasBillableOpenAIUsage distinguishes a usable accounting result from the +// zero value produced when an upstream omits usage entirely. A successful Grok +// completion without any positive usage cannot be safely charged, so callers +// must fail over before committing the response to the client. +func hasBillableOpenAIUsage(usage OpenAIUsage) bool { + return usage.InputTokens > 0 || + usage.ImageInputTokens > 0 || + usage.OutputTokens > 0 || + usage.CacheCreationInputTokens > 0 || + usage.CacheReadInputTokens > 0 || + usage.ImageOutputTokens > 0 +}