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:
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
|
||||
|
||||
Reference in new issue
Block a user