From 3a415e6d04272440c2f886c6c20ab0651b03f30e Mon Sep 17 00:00:00 2001 From: IanShaw027 Date: Fri, 7 Aug 2026 18:23:28 +0800 Subject: [PATCH] =?UTF-8?q?fix(grok):=20=E4=BA=94=E8=BD=AE=E8=AF=84?= =?UTF-8?q?=E5=AE=A1=20=E2=80=94=20messages=20SearchCount=20=E4=B8=8E=20Op?= =?UTF-8?q?enAI=20free-gate?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - /v1/messages Grok：SSE/buffered 路径累计 SearchCount - OpenAI 遗留 sticky/list 应用 free soft-gate（advanced 关闭时） - 单测：OpenAI getSchedulableAccount free-gate + messages 计数器 --- .../internal/service/grok_free_quota_gate.go | 4 +- ...openai_gateway_grok_search_billing_test.go | 41 +++++++++++++++++++ .../service/openai_gateway_messages.go | 25 +++++++++-- .../service/openai_gateway_scheduling.go | 27 +++++++++++- 4 files changed, 91 insertions(+), 6 deletions(-) diff --git a/backend/internal/service/grok_free_quota_gate.go b/backend/internal/service/grok_free_quota_gate.go index 424f12e53..4d48de899 100644 --- a/backend/internal/service/grok_free_quota_gate.go +++ b/backend/internal/service/grok_free_quota_gate.go @@ -113,8 +113,10 @@ func (s *GatewayService) filterGrokFreeQuotaAccountsForGateway(ctx context.Conte return filterGrokFreeQuotaAccountsCore(ctx, s.cfg, s.usageLogRepo, &gatewayGrokFreeQuotaGateCache, accounts) } -// Shared cache for Gateway path (OpenAI scheduler keeps its own map). +// Shared caches for non-advanced-scheduler selection paths. +// Advanced scheduler keeps per-instance sync.Map on defaultOpenAIAccountScheduler. var gatewayGrokFreeQuotaGateCache sync.Map +var openaiGrokFreeQuotaGateCache sync.Map func filterGrokFreeQuotaAccountsCore( ctx context.Context, diff --git a/backend/internal/service/openai_gateway_grok_search_billing_test.go b/backend/internal/service/openai_gateway_grok_search_billing_test.go index 5ae0ee21b..ed8c1fedb 100644 --- a/backend/internal/service/openai_gateway_grok_search_billing_test.go +++ b/backend/internal/service/openai_gateway_grok_search_billing_test.go @@ -121,3 +121,44 @@ func TestGetSchedulableAccount_AppliesGrokFreeSoftGate(t *testing.T) { require.NoError(t, err) require.Nil(t, got, "over free soft-gate sticky hit must miss") } + +func TestOpenAIGetSchedulableAccount_AppliesGrokFreeSoftGate(t *testing.T) { + // Legacy OpenAI-compatible sticky (advanced scheduler off) must free-gate Grok. + 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(8802, "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}, + }} + openaiGrokFreeQuotaGateCache.Range(func(key, _ any) bool { + openaiGrokFreeQuotaGateCache.Delete(key) + return true + }) + svc := &OpenAIGatewayService{ + cfg: cfg, + accountRepo: repo, + usageLogRepo: usageRepo, + } + + got, err := svc.getSchedulableAccount(context.Background(), account.ID) + require.NoError(t, err) + require.Nil(t, got, "OpenAI legacy sticky must apply free soft-gate") +} + +func TestCountGrokNativeSearchCallsFromJSON_MessagesStyleBody(t *testing.T) { + // Proves the same counter used by Anthropic-buffered Grok /v1/messages path. + body := []byte(`{"id":"r1","output":[{"type":"web_search_call","id":"ws1"},{"type":"message","role":"assistant"}]}`) + require.Equal(t, 1, countGrokNativeSearchCallsFromJSONBytes(body)) +} diff --git a/backend/internal/service/openai_gateway_messages.go b/backend/internal/service/openai_gateway_messages.go index 066fc82ce..f5af55559 100644 --- a/backend/internal/service/openai_gateway_messages.go +++ b/backend/internal/service/openai_gateway_messages.go @@ -601,7 +601,7 @@ func (s *OpenAIGatewayService) handleAnthropicBufferedStreamingResponse( c.Header("Content-Type", "application/json; charset=utf-8") c.JSON(http.StatusOK, anthropicResp) - return &OpenAIForwardResult{ + result := &OpenAIForwardResult{ RequestID: requestID, ResponseID: finalResponse.ID, Usage: usage, @@ -610,7 +610,16 @@ func (s *OpenAIGatewayService) handleAnthropicBufferedStreamingResponse( UpstreamModel: upstreamModel, Stream: false, Duration: time.Since(startTime), - }, nil + } + // Grok /v1/messages uses Responses upstream; count native search for surcharge. + 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 } func isOpenAICompatResponsesTerminalEvent(eventType string) bool { @@ -830,6 +839,9 @@ func (s *OpenAIGatewayService) handleAnthropicStreamingResponse( clientOutputStarted := false var streamFailoverErr error var streamNonFailoverErr error + searchCount := 0 + streamSearchSeen := make(map[string]struct{}) + countSearch := account != nil && account.IsGrok() scanner := s.newUpstreamSSEScanner(resp.Body) @@ -849,7 +861,7 @@ func (s *OpenAIGatewayService) handleAnthropicStreamingResponse( // resultWithUsage builds the final result snapshot. resultWithUsage := func() *OpenAIForwardResult { - return &OpenAIForwardResult{ + out := &OpenAIForwardResult{ RequestID: requestID, ResponseID: responseID, Usage: usage, @@ -861,6 +873,10 @@ func (s *OpenAIGatewayService) handleAnthropicStreamingResponse( FirstTokenMs: firstTokenMs, ClientDisconnect: clientDisconnected, } + if searchCount > 0 { + out.SearchCount = searchCount + } + return out } // processDataLine handles a single "data: ..." SSE line from upstream. @@ -870,6 +886,9 @@ func (s *OpenAIGatewayService) handleAnthropicStreamingResponse( 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_scheduling.go b/backend/internal/service/openai_gateway_scheduling.go index c23223271..4e7606947 100644 --- a/backend/internal/service/openai_gateway_scheduling.go +++ b/backend/internal/service/openai_gateway_scheduling.go @@ -1261,7 +1261,11 @@ func (s *OpenAIGatewayService) listSchedulableAccounts(ctx context.Context, grou if err != nil { return accounts, err } - return s.filterOpenAIAccountsBySchedulingThreshold(ctx, accounts), nil + accounts = s.filterOpenAIAccountsBySchedulingThreshold(ctx, accounts) + if platform == PlatformGrok { + accounts = s.filterGrokFreeQuotaAccountsForOpenAI(ctx, accounts) + } + return accounts, nil } var accounts []Account var err error @@ -1275,7 +1279,11 @@ func (s *OpenAIGatewayService) listSchedulableAccounts(ctx context.Context, grou if err != nil { return nil, fmt.Errorf("query accounts failed: %w", err) } - return s.filterOpenAIAccountsBySchedulingThreshold(ctx, accounts), nil + accounts = s.filterOpenAIAccountsBySchedulingThreshold(ctx, accounts) + if platform == PlatformGrok { + accounts = s.filterGrokFreeQuotaAccountsForOpenAI(ctx, accounts) + } + return accounts, nil } func (s *OpenAIGatewayService) tryAcquireAccountSlot(ctx context.Context, accountID int64, maxConcurrency int) (*AcquireResult, error) { @@ -1401,9 +1409,24 @@ func (s *OpenAIGatewayService) getSchedulableAccount(ctx context.Context, accoun if s.isOpenAIAccountBlockedBySchedulingThreshold(ctx, account) { return nil, nil } + // Legacy sticky (advanced scheduler off) must still free-gate Grok OAuth. + if account.IsGrok() { + if gated := s.filterGrokFreeQuotaAccountsForOpenAI(ctx, []Account{*account}); len(gated) == 0 { + return nil, nil + } + } return account, nil } +// filterGrokFreeQuotaAccountsForOpenAI applies the same local free soft-gate as +// GatewayService / advanced scheduler, for OpenAI-compatible legacy selection. +func (s *OpenAIGatewayService) filterGrokFreeQuotaAccountsForOpenAI(ctx context.Context, accounts []Account) []Account { + if s == nil { + return accounts + } + return filterGrokFreeQuotaAccountsCore(ctx, s.cfg, s.usageLogRepo, &openaiGrokFreeQuotaGateCache, accounts) +} + func (s *OpenAIGatewayService) filterOpenAIAccountsBySchedulingThreshold(ctx context.Context, accounts []Account) []Account { if len(accounts) == 0 { return accounts