diff --git a/backend/internal/service/openai_gateway_chat_completions.go b/backend/internal/service/openai_gateway_chat_completions.go index a66489da00..ddf9a288a8 100644 --- a/backend/internal/service/openai_gateway_chat_completions.go +++ b/backend/internal/service/openai_gateway_chat_completions.go @@ -463,7 +463,7 @@ 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) { + if requiresBillableGrokChatUsage(account, originalModel, billingModel, upstreamModel, finalResponse.Model) && !hasBillableOpenAIUsage(usage) { return nil, s.newOpenAIStreamFailoverError( c, account, diff --git a/backend/internal/service/openai_gateway_chat_completions_raw.go b/backend/internal/service/openai_gateway_chat_completions_raw.go index 0aee0fa5f5..cdafc30d1b 100644 --- a/backend/internal/service/openai_gateway_chat_completions_raw.go +++ b/backend/internal/service/openai_gateway_chat_completions_raw.go @@ -431,7 +431,8 @@ func (s *OpenAIGatewayService) bufferRawChatCompletions( if parsedUsage, ok := extractOpenAIUsageFromJSONBytes(respBody); ok { usage = parsedUsage } - if account != nil && account.Platform == PlatformGrok && !hasBillableOpenAIUsage(usage) { + responseModel := gjson.GetBytes(respBody, "model").String() + if requiresBillableGrokChatUsage(account, originalModel, billingModel, upstreamModel, responseModel) && !hasBillableOpenAIUsage(usage) { return nil, s.newOpenAIStreamFailoverError( c, account, 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 c19066f3fc..0a9741889c 100644 --- a/backend/internal/service/openai_gateway_chat_completions_raw_test.go +++ b/backend/internal/service/openai_gateway_chat_completions_raw_test.go @@ -14,6 +14,7 @@ import ( "time" "github.com/Wei-Shaw/sub2api/internal/config" + "github.com/Wei-Shaw/sub2api/internal/pkg/openai_compat" "github.com/gin-gonic/gin" "github.com/stretchr/testify/require" "github.com/tidwall/gjson" @@ -122,6 +123,41 @@ func TestForwardAsRawChatCompletions_ForcesStreamUsageUpstreamAndPassesUsageDown require.Contains(t, rec.Body.String(), "data: [DONE]") } +func TestForwardAsChatCompletions_OpenAICompatibleGrokRawMissingUsageFailsBeforeWrite(t *testing.T) { + gin.SetMode(gin.TestMode) + + body := []byte(`{"model":"grok-4.5","messages":[{"role":"user","content":"hello"}],"stream":false}`) + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", bytes.NewReader(body)) + + upstream := &httpUpstreamRecorder{resp: &http.Response{ + StatusCode: http.StatusOK, + Header: http.Header{ + "Content-Type": []string{"application/json"}, + "X-Request-Id": []string{"rid-openai-compat-grok-no-usage"}, + }, + Body: io.NopCloser(strings.NewReader( + `{"id":"resp_missing_usage","object":"chat.completion","model":"grok-4.5","choices":[{"index":0,"message":{"role":"assistant","content":"hello"},"finish_reason":"stop"}]}`, + )), + }} + svc := &OpenAIGatewayService{ + cfg: rawChatCompletionsTestConfig(), + httpUpstream: upstream, + } + account := rawChatCompletionsTestAccount() + account.Name = "openai-compatible-grok" + account.Extra = map[string]any{openai_compat.ExtraKeyResponsesSupported: false} + + 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(), "unbilled Grok content must not be returned by an OpenAI-compatible account") + require.Empty(t, recorder.Body.String()) +} + func TestForwardAsRawChatCompletions_PreservesMappedGPT56MaxEffort(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 index c1b9e1523b..a70357cc6f 100644 --- a/backend/internal/service/openai_gateway_usage_integrity.go +++ b/backend/internal/service/openai_gateway_usage_integrity.go @@ -14,3 +14,19 @@ func hasBillableOpenAIUsage(usage OpenAIUsage) bool { usage.CacheReadInputTokens > 0 || usage.ImageOutputTokens > 0 } + +// requiresBillableGrokChatUsage identifies Grok traffic by both account +// platform and model identity. Grok models may be served through generic +// OpenAI-compatible accounts, so account.Platform alone is not a safe billing +// boundary. +func requiresBillableGrokChatUsage(account *Account, models ...string) bool { + if account != nil && account.Platform == PlatformGrok { + return true + } + for _, model := range models { + if platform, ok := DetectModelPlatform(model); ok && platform == PlatformGrok { + return true + } + } + return false +} diff --git a/backend/internal/service/openai_gateway_usage_integrity_test.go b/backend/internal/service/openai_gateway_usage_integrity_test.go new file mode 100644 index 0000000000..3a1dc441f9 --- /dev/null +++ b/backend/internal/service/openai_gateway_usage_integrity_test.go @@ -0,0 +1,58 @@ +//go:build unit + +package service + +import ( + "testing" + + "github.com/stretchr/testify/require" +) + +func TestRequiresBillableGrokChatUsage(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + account *Account + models []string + want bool + }{ + { + name: "grok platform", + account: &Account{Platform: PlatformGrok}, + models: []string{"alias"}, + want: true, + }, + { + name: "OpenAI-compatible requested Grok model", + account: &Account{Platform: PlatformOpenAI}, + models: []string{"grok-4.5"}, + want: true, + }, + { + name: "OpenAI-compatible mapped Grok model", + account: &Account{Platform: PlatformOpenAI}, + models: []string{"alias", "grok-4.5"}, + want: true, + }, + { + name: "xAI-qualified Grok model", + account: &Account{Platform: PlatformOpenAI}, + models: []string{"xai/grok-4.5"}, + want: true, + }, + { + name: "ordinary OpenAI model", + account: &Account{Platform: PlatformOpenAI}, + models: []string{"gpt-5.4"}, + want: false, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + require.Equal(t, tt.want, requiresBillableGrokChatUsage(tt.account, tt.models...)) + }) + } +}