diff --git a/backend/internal/service/grok_free_quota_gate.go b/backend/internal/service/grok_free_quota_gate.go index 8930db114..fe1b974bb 100644 --- a/backend/internal/service/grok_free_quota_gate.go +++ b/backend/internal/service/grok_free_quota_gate.go @@ -155,12 +155,17 @@ func filterGrokFreeQuotaAccountsCore( } if cached, ok := cache.Load(account.ID); ok { entry, valid := cached.(grokFreeQuotaGateCacheEntry) - age := now.Sub(entry.checkedAt) - if valid && settings.cacheTTL > 0 && age >= 0 && age < settings.cacheTTL { - if entry.known { - tokensByID[account.ID] = entry.tokens + if valid { + age := now.Sub(entry.checkedAt) + // cacheTTL == 0 means "no expiry" for known entries (still fail-open + // on first miss; refresh is only scheduled when missing/stale). + fresh := settings.cacheTTL <= 0 || (age >= 0 && age < settings.cacheTTL) + if fresh { + if entry.known { + tokensByID[account.ID] = entry.tokens + } + continue } - continue } } // Miss / stale: fail open on this request; refresh asynchronously. @@ -228,10 +233,10 @@ func scheduleGrokFreeQuotaStatsRefresh( statsByID, err := queryGrokFreeQuotaWindowStats(context.Background(), usageLogRepo, toFetch, now.Add(-window)) if err != nil { grokFreeQuotaGateQueryFailureTotal.Add(1) - if cacheTTL > 0 { - for _, accountID := range toFetch { - cache.Store(accountID, grokFreeQuotaGateCacheEntry{checkedAt: now}) - } + // Store a negative entry so subsequent hot-path calls do not thrash. + // known=false → still fail open until a successful refresh lands. + for _, accountID := range toFetch { + cache.Store(accountID, grokFreeQuotaGateCacheEntry{checkedAt: now}) } slog.Warn("grok_free_quota_soft_gate_stats_failed", "account_count", len(toFetch), diff --git a/backend/internal/service/grok_free_quota_gate_test.go b/backend/internal/service/grok_free_quota_gate_test.go index b1f8d604f..3b3eae769 100644 --- a/backend/internal/service/grok_free_quota_gate_test.go +++ b/backend/internal/service/grok_free_quota_gate_test.go @@ -157,15 +157,23 @@ func TestFilterGrokFreeQuotaAccountsRecoversAfterRollingUsageFalls(t *testing.T) require.Empty(t, scheduler.filterGrokFreeQuotaAccounts(context.Background(), accounts), "fresh cache keeps the soft-gate hold") // Expire entry → miss fails open and schedules refresh with recovered usage. + // Clear in-flight markers so a refresh is allowed after we force-expire the entry. + if root, ok := freeQuotaRefreshInFlight.Load(&scheduler.grokFreeQuotaGateCache); ok { + if m, ok := root.(*sync.Map); ok { + m.Delete(int64(1)) + } + } + callsBeforeExpire := repo.calls scheduler.grokFreeQuotaGateCache.Store(int64(1), grokFreeQuotaGateCacheEntry{ - tokens: 490_000, checkedAt: time.Now().Add(-time.Minute), known: true, + tokens: 490_000, checkedAt: time.Now().Add(-2 * time.Minute), known: true, // TTL=60s → stale }) + // Hot path fail-open while refresh is in flight. require.Equal(t, []int64{1}, accountIDs(scheduler.filterGrokFreeQuotaAccounts(context.Background(), accounts))) require.Eventually(t, func() bool { filtered := scheduler.filterGrokFreeQuotaAccounts(context.Background(), accounts) - return len(filtered) == 1 && filtered[0].ID == 1 + return len(filtered) == 1 && filtered[0].ID == 1 && + repo.calls > callsBeforeExpire }, 2*time.Second, 10*time.Millisecond) - require.GreaterOrEqual(t, repo.calls, 2) } func TestResolveGrokFreeQuotaGateSettingsDefaultsToNinetyFivePercent(t *testing.T) { 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 5d0525b12..dc998ec5d 100644 --- a/backend/internal/service/openai_gateway_grok_search_billing_test.go +++ b/backend/internal/service/openai_gateway_grok_search_billing_test.go @@ -8,6 +8,7 @@ import ( "io" "net/http" "net/http/httptest" + "sync" "testing" "time" @@ -87,13 +88,14 @@ func TestForwardGrokResponses_PropagatesSearchCountFromSSE(t *testing.T) { } func TestGetSchedulableAccount_AppliesGrokFreeSoftGate(t *testing.T) { - // Sticky/non-list path must not return over-gate free OAuth accounts. + // Sticky/non-list path must not return over-gate free OAuth accounts once cache is warm. + // First sticky hit fail-opens and schedules async refresh; subsequent hits use the cache. cfg := &config.Config{} cfg.Gateway.Grok.FreeQuotaSoftGateEnabled = true cfg.Gateway.Grok.FreeQuotaTokenLimit = 500_000 cfg.Gateway.Grok.FreeQuotaSoftGatePercent = 95 cfg.Gateway.Grok.FreeQuotaWindowHours = 24 - cfg.Gateway.Grok.FreeQuotaStatsCacheSeconds = 0 + cfg.Gateway.Grok.FreeQuotaStatsCacheSeconds = 60 account := healthyGrokOAuthGatewayTestAccount(8801, "tok") account.Credentials["subscription_tier"] = "free" @@ -111,15 +113,26 @@ func TestGetSchedulableAccount_AppliesGrokFreeSoftGate(t *testing.T) { gatewayGrokFreeQuotaGateCache.Delete(key) return true }) + if root, ok := freeQuotaRefreshInFlight.Load(&gatewayGrokFreeQuotaGateCache); ok { + if m, ok := root.(*sync.Map); ok { + m.Delete(account.ID) + } + } svc := &GatewayService{ cfg: cfg, accountRepo: repo, usageLogRepo: usageRepo, } + // Miss: fail open + schedule refresh. got, err := svc.getSchedulableAccount(context.Background(), account.ID) require.NoError(t, err) - require.Nil(t, got, "over free soft-gate sticky hit must miss") + require.NotNil(t, got, "first sticky hit fail-opens while free-gate stats refresh") + + require.Eventually(t, func() bool { + got, err := svc.getSchedulableAccount(context.Background(), account.ID) + return err == nil && got == nil + }, 2*time.Second, 10*time.Millisecond, "over free soft-gate sticky hit must miss after cache warm") } func TestOpenAIGetSchedulableAccount_AppliesGrokFreeSoftGate(t *testing.T) { @@ -129,7 +142,7 @@ func TestOpenAIGetSchedulableAccount_AppliesGrokFreeSoftGate(t *testing.T) { cfg.Gateway.Grok.FreeQuotaTokenLimit = 500_000 cfg.Gateway.Grok.FreeQuotaSoftGatePercent = 95 cfg.Gateway.Grok.FreeQuotaWindowHours = 24 - cfg.Gateway.Grok.FreeQuotaStatsCacheSeconds = 0 + cfg.Gateway.Grok.FreeQuotaStatsCacheSeconds = 60 account := healthyGrokOAuthGatewayTestAccount(8802, "tok") account.Credentials["subscription_tier"] = "free" @@ -146,6 +159,11 @@ func TestOpenAIGetSchedulableAccount_AppliesGrokFreeSoftGate(t *testing.T) { openaiGrokFreeQuotaGateCache.Delete(key) return true }) + if root, ok := freeQuotaRefreshInFlight.Load(&openaiGrokFreeQuotaGateCache); ok { + if m, ok := root.(*sync.Map); ok { + m.Delete(account.ID) + } + } svc := &OpenAIGatewayService{ cfg: cfg, accountRepo: repo, @@ -154,7 +172,12 @@ func TestOpenAIGetSchedulableAccount_AppliesGrokFreeSoftGate(t *testing.T) { got, err := svc.getSchedulableAccount(context.Background(), account.ID) require.NoError(t, err) - require.Nil(t, got, "OpenAI legacy sticky must apply free soft-gate") + require.NotNil(t, got, "first sticky hit fail-opens while free-gate stats refresh") + + require.Eventually(t, func() bool { + got, err := svc.getSchedulableAccount(context.Background(), account.ID) + return err == nil && got == nil + }, 2*time.Second, 10*time.Millisecond, "OpenAI legacy sticky must apply free soft-gate after cache warm") } func TestCountGrokNativeSearchCallsFromJSON_MessagesStyleBody(t *testing.T) {