diff --git a/backend/internal/service/gateway_scheduling.go b/backend/internal/service/gateway_scheduling.go index aee0ad7f75..be63f804f5 100644 --- a/backend/internal/service/gateway_scheduling.go +++ b/backend/internal/service/gateway_scheduling.go @@ -1451,6 +1451,12 @@ func (s *GatewayService) getSchedulableAccount(ctx context.Context, accountID in if s.isAccountBlockedBySchedulingThreshold(ctx, account) { return nil, nil } + // Sticky / non-list selection must honor free soft-gate (same as listSchedulableAccounts). + if account.IsGrok() { + if gated := s.filterGrokFreeQuotaAccountsForGateway(ctx, []Account{*account}); len(gated) == 0 { + return nil, nil + } + } return account, nil } diff --git a/backend/internal/service/grok_search_count.go b/backend/internal/service/grok_search_count.go index 80dc372168..70c46c835a 100644 --- a/backend/internal/service/grok_search_count.go +++ b/backend/internal/service/grok_search_count.go @@ -42,6 +42,10 @@ func countGrokNativeSearchCallsInSSEData(data []byte) int { // countGrokNativeSearchCallsInSSEDataDedup increments only unseen call ids. // Callers must reuse the same seen map for the full stream lifetime. +// +// When call_id/id is missing, a synthetic key is built from item type + name so +// item.done + response.completed for the same tool still count once (never fall +// back to raw multi-event n, which ~2× overbills). func countGrokNativeSearchCallsInSSEDataDedup(data []byte, seen map[string]struct{}) int { if seen == nil { return countGrokNativeSearchCallsInSSEData(data) @@ -50,11 +54,15 @@ func countGrokNativeSearchCallsInSSEDataDedup(data []byte, seen map[string]struc if n <= 0 { return 0 } - // No stable id: fall back to raw count once (cannot dedup across events). + // Prefer stable ids; fill gaps with synthetic keys so we never raw-add n. + if len(keys) < n { + // Rebuild keys for every item so unkeyed items still get a fingerprint. + keys = collectGrokNativeSearchCallKeys(data) + } if len(keys) == 0 { - return n + // True empty — should not happen when n>0; fail-closed to 0 extra bill. + return 0 } - // Intra-event + cross-event dedup by call_id/id. added := 0 local := make(map[string]struct{}, len(keys)) for _, k := range keys { @@ -74,6 +82,54 @@ func countGrokNativeSearchCallsInSSEDataDedup(data []byte, seen map[string]struc return added } +func collectGrokNativeSearchCallKeys(data []byte) []string { + if len(data) == 0 || !gjson.ValidBytes(data) { + return nil + } + eventType := strings.TrimSpace(gjson.GetBytes(data, "type").String()) + switch eventType { + case "response.output_item.done", "response.completed", "response.done", "": + default: + if eventType != "" { + return nil + } + } + var keys []string + consider := func(item gjson.Result) { + if !isGrokNativeSearchOutputItem(item) { + return + } + key := firstNonEmpty( + strings.TrimSpace(item.Get("call_id").String()), + strings.TrimSpace(item.Get("id").String()), + strings.TrimSpace(item.Get("item.call_id").String()), + strings.TrimSpace(item.Get("item.id").String()), + ) + if key == "" { + // Synthetic fingerprint: type + name is stable across done/completed + // for the same tool invocation when upstream omits call_id. + key = "synth:" + strings.ToLower(strings.TrimSpace(item.Get("type").String())) + + ":" + strings.ToLower(strings.TrimSpace(item.Get("name").String())) + } + keys = append(keys, key) + } + if item := gjson.GetBytes(data, "item"); item.Exists() { + consider(item) + } + gjson.GetBytes(data, "response.output").ForEach(func(_, item gjson.Result) bool { + consider(item) + return true + }) + gjson.GetBytes(data, "output").ForEach(func(_, item gjson.Result) bool { + consider(item) + return true + }) + if len(keys) == 0 && isGrokNativeSearchOutputItem(gjson.ParseBytes(data)) { + consider(gjson.ParseBytes(data)) + } + return keys +} + func countGrokNativeSearchCallsInSSEDataWithKeys(data []byte) (int, []string) { if len(data) == 0 || !gjson.ValidBytes(data) { return 0, nil diff --git a/backend/internal/service/grok_search_count_test.go b/backend/internal/service/grok_search_count_test.go index b626e29c6d..799bd051d7 100644 --- a/backend/internal/service/grok_search_count_test.go +++ b/backend/internal/service/grok_search_count_test.go @@ -43,6 +43,16 @@ func TestCountGrokNativeSearchCallsInSSEDataDedup_LiveStreamPath(t *testing.T) { require.Equal(t, 2, countGrokNativeSearchCallsInSSEData(completed)) } +func TestCountGrokNativeSearchCallsInSSEDataDedup_NoIDStillDedups(t *testing.T) { + t.Parallel() + // Upstream sometimes omits call_id/id; synthetic keys must still prevent 2×. + seen := make(map[string]struct{}) + done := []byte(`{"type":"response.output_item.done","item":{"type":"web_search_call"}}`) + completed := []byte(`{"type":"response.completed","response":{"output":[{"type":"web_search_call"}]}}`) + require.Equal(t, 1, countGrokNativeSearchCallsInSSEDataDedup(done, seen)) + require.Equal(t, 0, countGrokNativeSearchCallsInSSEDataDedup(completed, seen)) +} + func stringsJoin(lines ...string) string { out := "" for _, l := range lines { diff --git a/backend/internal/service/openai_gateway_chat_completions.go b/backend/internal/service/openai_gateway_chat_completions.go index 171a6c8805..514e68266e 100644 --- a/backend/internal/service/openai_gateway_chat_completions.go +++ b/backend/internal/service/openai_gateway_chat_completions.go @@ -476,7 +476,7 @@ func (s *OpenAIGatewayService) handleChatBufferedStreamingResponse( c.Writer.Header().Set("Content-Type", "application/json; charset=utf-8") c.JSON(http.StatusOK, chatResp) - return &OpenAIForwardResult{ + result := &OpenAIForwardResult{ RequestID: requestID, Usage: usage, Model: originalModel, @@ -484,7 +484,16 @@ func (s *OpenAIGatewayService) handleChatBufferedStreamingResponse( UpstreamModel: upstreamModel, Stream: false, Duration: time.Since(startTime), - }, nil + } + // Grok chat bridge: bill native search tools found in the terminal Responses body. + if account != nil && account.IsGrok() && finalResponse != nil { + if body, err := json.Marshal(finalResponse); err == nil { + if n := countGrokNativeSearchCallsFromJSONBytes(body); n > 0 { + result.SearchCount = n + } + } + } + return result, nil } // handleChatStreamingResponse reads Responses SSE events from upstream, @@ -517,6 +526,10 @@ func (s *OpenAIGatewayService) handleChatStreamingResponse( refusalDetector := newOpenAIChatSilentRefusalDetector(requestBodyLen) var streamFailoverErr *UpstreamFailoverError var streamNonFailoverErr error + // Grok chat bridge reuses Responses SSE; count native search tools for surcharge. + searchCount := 0 + streamSearchSeen := make(map[string]struct{}) + countSearch := account != nil && account.IsGrok() scanner := s.newUpstreamSSEScanner(resp.Body) @@ -535,7 +548,7 @@ func (s *OpenAIGatewayService) handleChatStreamingResponse( } resultWithUsage := func() *OpenAIForwardResult { - return &OpenAIForwardResult{ + out := &OpenAIForwardResult{ RequestID: requestID, Usage: usage, Model: originalModel, @@ -545,6 +558,10 @@ func (s *OpenAIGatewayService) handleChatStreamingResponse( Duration: time.Since(startTime), FirstTokenMs: firstTokenMs, } + if searchCount > 0 { + out.SearchCount = searchCount + } + return out } processDataLine := func(payload string) bool { @@ -553,6 +570,9 @@ func (s *OpenAIGatewayService) handleChatStreamingResponse( ms := int(time.Since(startTime).Milliseconds()) firstTokenMs = &ms } + if countSearch { + searchCount += countGrokNativeSearchCallsInSSEDataDedup([]byte(payload), streamSearchSeen) + } var event apicompat.ResponsesStreamEvent if err := json.Unmarshal([]byte(payload), &event); err != nil { diff --git a/backend/internal/service/openai_gateway_grok.go b/backend/internal/service/openai_gateway_grok.go index 4160819e67..13b0f121b5 100644 --- a/backend/internal/service/openai_gateway_grok.go +++ b/backend/internal/service/openai_gateway_grok.go @@ -186,6 +186,9 @@ func (s *OpenAIGatewayService) forwardGrokResponses( var usage *OpenAIUsage var firstTokenMs *int responseID := "" + searchCount := 0 + imageCount := 0 + var imageOutputSizes []string if reqStream { maxLineSize := defaultMaxLineSize if s.cfg != nil && s.cfg.Gateway.MaxLineSize > 0 { @@ -202,6 +205,9 @@ func (s *OpenAIGatewayService) forwardGrokResponses( usage = streamResult.usage firstTokenMs = streamResult.firstTokenMs responseID = strings.TrimSpace(streamResult.responseID) + searchCount = streamResult.searchCount + imageCount = streamResult.imageCount + imageOutputSizes = streamResult.imageOutputSizes } else { nonStreamResult, err := s.handleNonStreamingResponse(ctx, resp, c, account, originalModel, upstreamModel) if err != nil { @@ -209,13 +215,16 @@ func (s *OpenAIGatewayService) forwardGrokResponses( } usage = nonStreamResult.usage responseID = strings.TrimSpace(nonStreamResult.responseID) + searchCount = nonStreamResult.searchCount + imageCount = nonStreamResult.imageCount + imageOutputSizes = nonStreamResult.imageOutputSizes } if usage == nil { usage = &OpenAIUsage{} } reasoningEffort := extractOpenAIReasoningEffortFromBody(patchedBody, originalModel) - return &OpenAIForwardResult{ + result := &OpenAIForwardResult{ RequestID: firstNonEmpty(resp.Header.Get("x-request-id"), resp.Header.Get("xai-request-id")), ResponseID: responseID, Usage: *usage, @@ -227,7 +236,17 @@ func (s *OpenAIGatewayService) forwardGrokResponses( ResponseHeaders: resp.Header.Clone(), Duration: time.Since(startTime), FirstTokenMs: firstTokenMs, - }, nil + } + // Propagate search/image counters from the shared Responses handler — without + // this, stream/JSON counting runs but search_price_per_1k / image bills never apply. + if searchCount > 0 { + result.SearchCount = searchCount + } + if imageCount > 0 { + result.ImageCount = imageCount + result.ImageOutputSizes = imageOutputSizes + } + return result, nil } func isGrokInvalidEncryptedContentResponse(statusCode int, body []byte) bool { diff --git a/backend/internal/service/openai_gateway_grok_search_billing_test.go b/backend/internal/service/openai_gateway_grok_search_billing_test.go new file mode 100644 index 0000000000..5ae0ee21b4 --- /dev/null +++ b/backend/internal/service/openai_gateway_grok_search_billing_test.go @@ -0,0 +1,123 @@ +//go:build unit + +package service + +import ( + "bytes" + "context" + "io" + "net/http" + "net/http/httptest" + "testing" + "time" + + "github.com/Wei-Shaw/sub2api/internal/config" + "github.com/Wei-Shaw/sub2api/internal/pkg/usagestats" + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/require" +) + +func TestForwardGrokResponses_PropagatesSearchCountFromJSON(t *testing.T) { + gin.SetMode(gin.TestMode) + body := []byte(`{"model":"grok","input":"search something","tools":[{"type":"web_search"}],"stream":false}`) + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", bytes.NewReader(body)) + + account := healthyGrokOAuthGatewayTestAccount(9901, "access-token") + repo := &mockAccountRepoForPlatform{accountsByID: map[int64]*Account{account.ID: account}} + upstreamBody := `{ + "id":"resp_search_bill", + "object":"response", + "model":"grok-4.5", + "status":"completed", + "output":[ + {"type":"web_search_call","id":"ws1","call_id":"c1","status":"completed"}, + {"type":"x_search_call","id":"xs1","call_id":"c2"}, + {"type":"message","role":"assistant","content":[{"type":"output_text","text":"ok"}]} + ], + "usage":{"input_tokens":10,"output_tokens":5} + }` + upstream := &httpUpstreamRecorder{resp: &http.Response{ + StatusCode: http.StatusOK, + Header: http.Header{"Content-Type": []string{"application/json"}}, + Body: io.NopCloser(bytes.NewReader([]byte(upstreamBody))), + }} + svc := &OpenAIGatewayService{ + httpUpstream: upstream, + grokTokenProvider: NewGrokTokenProvider(repo, nil), + accountRepo: repo, + } + + result, err := svc.forwardGrokResponses(context.Background(), c, account, body, "grok", false, time.Now()) + require.NoError(t, err) + require.NotNil(t, result) + require.Equal(t, 2, result.SearchCount, "Grok Responses must surface search tool calls for surcharge billing") + require.Equal(t, 10, result.Usage.InputTokens) + require.Equal(t, 5, result.Usage.OutputTokens) +} + +func TestForwardGrokResponses_PropagatesSearchCountFromSSE(t *testing.T) { + gin.SetMode(gin.TestMode) + body := []byte(`{"model":"grok","input":"search","tools":[{"type":"web_search"}],"stream":true}`) + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", bytes.NewReader(body)) + + account := healthyGrokOAuthGatewayTestAccount(9902, "access-token") + repo := &mockAccountRepoForPlatform{accountsByID: map[int64]*Account{account.ID: account}} + // item.done + response.completed for same call_id must count once after wire-up. + sse := "data: {\"type\":\"response.output_item.done\",\"item\":{\"type\":\"web_search_call\",\"id\":\"ws1\",\"call_id\":\"c1\"}}\n\n" + + "data: {\"type\":\"response.completed\",\"response\":{\"id\":\"resp_s\",\"status\":\"completed\",\"output\":[{\"type\":\"web_search_call\",\"id\":\"ws1\",\"call_id\":\"c1\"}],\"usage\":{\"input_tokens\":3,\"output_tokens\":1}}}\n\n" + upstream := &httpUpstreamRecorder{resp: &http.Response{ + StatusCode: http.StatusOK, + Header: http.Header{"Content-Type": []string{"text/event-stream"}}, + Body: io.NopCloser(bytes.NewReader([]byte(sse))), + }} + svc := &OpenAIGatewayService{ + httpUpstream: upstream, + grokTokenProvider: NewGrokTokenProvider(repo, nil), + accountRepo: repo, + } + + result, err := svc.forwardGrokResponses(context.Background(), c, account, body, "grok", true, time.Now()) + require.NoError(t, err) + require.NotNil(t, result) + require.Equal(t, 1, result.SearchCount, "stream SearchCount must be wired and deduped") +} + +func TestGetSchedulableAccount_AppliesGrokFreeSoftGate(t *testing.T) { + // Sticky/non-list path must not return over-gate free OAuth accounts. + cfg := &config.Config{} + cfg.Gateway.Grok.FreeQuotaSoftGateEnabled = true + cfg.Gateway.Grok.FreeQuotaTokenLimit = 1_000_000 + cfg.Gateway.Grok.FreeQuotaSoftGatePercent = 95 + cfg.Gateway.Grok.FreeQuotaWindowHours = 24 + cfg.Gateway.Grok.FreeQuotaStatsCacheSeconds = 0 + + account := healthyGrokOAuthGatewayTestAccount(8801, "tok") + account.Credentials["subscription_tier"] = "free" + account.Status = StatusActive + account.Schedulable = true + + repo := &mockAccountRepoForPlatform{ + accountsByID: map[int64]*Account{account.ID: account}, + } + usageRepo := &grokFreeQuotaUsageRepoStub{stats: map[int64]*usagestats.AccountStats{ + account.ID: {Tokens: 999_000}, + }} + // Clear shared gateway free-gate cache so this test is deterministic. + gatewayGrokFreeQuotaGateCache.Range(func(key, _ any) bool { + gatewayGrokFreeQuotaGateCache.Delete(key) + return true + }) + svc := &GatewayService{ + cfg: cfg, + accountRepo: repo, + usageLogRepo: usageRepo, + } + + got, err := svc.getSchedulableAccount(context.Background(), account.ID) + require.NoError(t, err) + require.Nil(t, got, "over free soft-gate sticky hit must miss") +}