fix(grok): 五轮评审 — messages SearchCount 与 OpenAI free-gate

- /v1/messages Grok:SSE/buffered 路径累计 SearchCount
- OpenAI 遗留 sticky/list 应用 free soft-gate(advanced 关闭时)
- 单测:OpenAI getSchedulableAccount free-gate + messages 计数器
This commit is contained in:
IanShaw027 committed 2026-08-07 18:23:28 +08:00
1 parent 3ae94df72b
commit 3a415e6d04
4 files changed
+91 -6

No files matched your search

@@ -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,
@@ -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))
}
@@ -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 {
@@ -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