diff --git a/changelog/unreleased/runtime-quota-recovery.md b/changelog/unreleased/runtime-quota-recovery.md new file mode 100644 index 0000000..482fbbe --- /dev/null +++ b/changelog/unreleased/runtime-quota-recovery.md @@ -0,0 +1,11 @@ +### English + +- Auto-refresh stale exhausted quota: the maintenance loop now re-fetches quota every 5 minutes, but only for accounts whose cached snapshot still reports `Exceeded`. A snapshot can go stale when upstream credits are replenished out-of-band (billing reset, add-on purchase) or when a flaky quota endpoint leaves an old `Exceeded=true` in place, so the account stayed out of routing until an operator manually refreshed it. Healthy accounts are never touched by this sweep. +- Quota error instead of "no accounts available": when every account that can serve a model is out of quota, the request now returns `429 insufficient_quota` rather than a generic "no accounts available", so the client gets the real cause. +- Back off quota probing for genuinely exhausted accounts: the maintenance sweep still starts at 5 minutes, but each account that keeps reporting `Exceeded` now steps through 5m → 5m → 5m → 15m → 30m → 1h → 3h → 6h → 12h → 24h instead of being re-fetched every five minutes forever. Provider balances reset on a daily cadence, so a permanently empty account no longer costs ~288 pointless quota calls a day, settling at roughly one probe per day. The ladder resets to the fast end as soon as a snapshot reports a non-zero balance, and a failed probe keeps its level instead of escalating a transient outage into an hours-long wait. + +### 中文 + +- 自动刷新 stale 的已耗尽配额：维护循环现在每 5 分钟重新拉取一次配额，但仅针对缓存快照仍为 `Exceeded` 的账号。快照可能因上游额度被外部恢复（计费重置、加购附加包）或配额接口抖动留下旧的 `Exceeded=true` 而变 stale，导致账号一直被路由排除，直到运维手动刷新才恢复。健康账号不会被此扫描触碰。 +- 配额耗尽返回配额错误而非「无可用账户」：当能服务某模型的所有账号都配额耗尽时，请求现在返回 `429 insufficient_quota`，而不是笼统的「无可用账户」，让客户端拿到真实原因。 +- 为真正耗尽的账号增加配额探测退避：维护扫描仍从 5 分钟起步，但连续报告 `Exceeded` 的账号会依次退避到 5m → 5m → 5m → 15m → 30m → 1h → 3h → 6h → 12h → 24h，而不是永远每 5 分钟重拉一次。供应商额度一般按天重置，长期为空的账号不会再每天产生约 288 次无效配额调用，最终稳定在每天约 1 次探测。一旦快照报告非零余额，退避立即回到最快档；探测失败则保持当前档位，不会把瞬时故障升级成数小时的等待。 diff --git a/internal/executor/chat.go b/internal/executor/chat.go index cadf021..778b9fb 100644 --- a/internal/executor/chat.go +++ b/internal/executor/chat.go @@ -311,6 +311,21 @@ func (e ChatExecutor) pick(requestID, prefer, providerFilter, regionFilter, publ Cooldown: 5 * time.Second, RetryAfter: 5 * time.Second, Failover: true, }, nil) } + // No eligible account is ready right now. Before blaming the model or + // reporting "no accounts available", check whether accounts that could + // serve this model are blocked only by an exhausted quota. A stale or + // genuinely exhausted snapshot leaves no routing hint, so without this + // the client sees a generic "no accounts available" instead of a quota + // error that explains the real cause. + if _, ok := e.Pool.QuotaBlockedForModel(query); ok { + retryAfter := NextLocalMidnightCooldown() + return Item{}, NewExecutionError(Classified{ + Kind: accounts.KindQuota, Status: 429, Code: "insufficient_quota", + Type: "insufficient_quota", Failover: false, + Cooldown: retryAfter, RetryAfter: retryAfter, + Message: quotaBlockedMessage(publicModel), + }, nil) + } if publicModel != "" && publicModel != "auto" { unfiltered := query unfiltered.PublicModel = "" @@ -412,6 +427,13 @@ func logTime(value time.Time) string { return value.UTC().Format(time.RFC3339) } +func quotaBlockedMessage(publicModel string) string { + if strings.TrimSpace(publicModel) != "" { + return fmt.Sprintf("all accounts that can serve %s are out of quota", publicModel) + } + return "all accounts are out of quota" +} + func coolingPickError(item Item, publicModel string, retryAfter time.Duration) error { failover := true kind := accounts.KindRateLimit diff --git a/internal/executor/pool.go b/internal/executor/pool.go index 47ef8e0..40f609e 100644 --- a/internal/executor/pool.go +++ b/internal/executor/pool.go @@ -327,6 +327,14 @@ func routeBaseMatches(item Item, q RouteQuery) bool { if item.Quota != nil && item.Quota.Exceeded && !itemServesFreeModel(item, q.PublicModel) { return false } + return routeReadyAndScoped(item, q) +} + +// routeReadyAndScoped is routeBaseMatches without the quota gate, so callers +// can pair it with the quota check to tell "blocked only by quota" from +// "blocked by provider/region/readiness/exclusion" without duplicating the +// scoping conditions. +func routeReadyAndScoped(item Item, q RouteQuery) bool { if !itemReady(item) { return false } @@ -348,6 +356,16 @@ func routeBaseMatches(item Item, q RouteQuery) bool { return true } +// routeBlockedByQuota reports whether an item would match the route except +// that its quota is exceeded. Free-model accounts are never quota-blocked: +// they spend no credits, so an exhausted account can still serve them. +func routeBlockedByQuota(item Item, q RouteQuery) bool { + if item.Quota == nil || !item.Quota.Exceeded || itemServesFreeModel(item, q.PublicModel) { + return false + } + return routeReadyAndScoped(item, q) && itemCouldServeModel(item, q.PublicModel) +} + func routeMatches(item Item, q RouteQuery) bool { return routeBaseMatches(item, q) && itemHasModel(item, q.PublicModel) } @@ -480,6 +498,24 @@ func (p *Pool) LenRoute(q RouteQuery) int { return count } +// QuotaBlockedForModel reports whether at least one account could serve the +// model but is kept out of routing solely because its quota is exceeded. The +// caller uses this to surface a quota error instead of a generic "no accounts +// available" when every model-serving account is out of credits. +func (p *Pool) QuotaBlockedForModel(q RouteQuery) (Item, bool) { + if p == nil { + return Item{}, false + } + p.mu.Lock() + defer p.mu.Unlock() + for _, item := range p.items { + if routeBlockedByQuota(item, q) { + return item, true + } + } + return Item{}, false +} + func (p *Pool) First() (Item, bool) { if p == nil || len(p.items) == 0 { return Item{}, false diff --git a/internal/executor/quota_blocked_test.go b/internal/executor/quota_blocked_test.go new file mode 100644 index 0000000..eb03b4e --- /dev/null +++ b/internal/executor/quota_blocked_test.go @@ -0,0 +1,69 @@ +package executor + +import "testing" + +// QuotaBlockedForModel must surface an account that could serve the requested +// model but is kept out of routing solely because its quota is exhausted. This +// is how the request path returns a quota error instead of a generic "no +// accounts available" when every model-serving account is out of credits. +func TestQuotaBlockedForModel(t *testing.T) { + pool := NewPool(nil, nil) + pool.Upsert(Item{ + ID: "acc-exhausted", + Provider: "qoder", + Region: "global", + Models: []string{"ultimate"}, + Quota: &QuotaSnapshot{Exceeded: true}, + }) + + if _, ok := pool.QuotaBlockedForModel(RouteQuery{PublicModel: "ultimate"}); !ok { + t.Fatal("exhausted account serving the model was not reported as quota-blocked") + } +} + +// A free model spends no credits, so an exhausted account is never quota-blocked +// for it: the account can still serve the route. +func TestQuotaBlockedForModelFreeModelNotBlocked(t *testing.T) { + pool := NewPool(nil, nil) + pool.Upsert(Item{ + ID: "acc-exhausted", + Provider: "qoder", + Region: "global", + Models: []string{"qwen3.8-flash", "ultimate"}, + FreeModels: []string{"qwen3.8-flash"}, + Quota: &QuotaSnapshot{Exceeded: true}, + }) + + if _, ok := pool.QuotaBlockedForModel(RouteQuery{PublicModel: "qwen3.8-flash"}); ok { + t.Fatal("free model was reported as quota-blocked; an exhausted account can still serve it") + } +} + +// An account that does not serve the requested model, is not exhausted, or is +// excluded from the attempt must not be reported as quota-blocked. +func TestQuotaBlockedForModelNotBlocked(t *testing.T) { + pool := NewPool(nil, nil) + pool.Upsert(Item{ + ID: "acc-ok", + Provider: "qoder", + Region: "global", + Models: []string{"ultimate"}, + Quota: &QuotaSnapshot{Exceeded: false}, + }) + pool.Upsert(Item{ + ID: "acc-other-model", + Provider: "qoder", + Region: "global", + Models: []string{"other"}, + Quota: &QuotaSnapshot{Exceeded: true}, + }) + + if _, ok := pool.QuotaBlockedForModel(RouteQuery{PublicModel: "ultimate"}); ok { + t.Fatal("no account serving 'ultimate' is quota-blocked, but one was reported") + } + + excluded := RouteQuery{PublicModel: "ultimate", Excluded: map[string]struct{}{"acc-ok": {}}} + if _, ok := pool.QuotaBlockedForModel(excluded); ok { + t.Fatal("excluded accounts must not be reported as quota-blocked") + } +} diff --git a/internal/runtime/maintenance_quota_test.go b/internal/runtime/maintenance_quota_test.go new file mode 100644 index 0000000..3f11554 --- /dev/null +++ b/internal/runtime/maintenance_quota_test.go @@ -0,0 +1,136 @@ +package runtime + +import ( + "testing" + "time" +) + +func TestQuotaSnapshotStale(t *testing.T) { + now := time.Date(2026, 10, 16, 12, 0, 0, 0, time.UTC) + + tests := []struct { + name string + q *QuotaSnapshot + delay time.Duration + want bool + }{ + { + name: "nil snapshot is not stale (caller skips)", + q: nil, + delay: 5 * time.Minute, + want: false, + }, + { + name: "empty FetchedAt is stale", + q: &QuotaSnapshot{Exceeded: true}, + delay: 5 * time.Minute, + want: true, + }, + { + name: "fetched under the delay is fresh", + q: &QuotaSnapshot{Exceeded: true, FetchedAt: now.Add(-1 * time.Minute).Format(time.RFC3339)}, + delay: 5 * time.Minute, + want: false, + }, + { + name: "fetched past the delay is stale", + q: &QuotaSnapshot{Exceeded: true, FetchedAt: now.Add(-(5*time.Minute + time.Minute)).Format(time.RFC3339)}, + delay: 5 * time.Minute, + want: true, + }, + { + name: "unparseable FetchedAt is stale", + q: &QuotaSnapshot{Exceeded: true, FetchedAt: "yesterday"}, + delay: 5 * time.Minute, + want: true, + }, + { + name: "staleness ignores the Exceeded flag (caller gates on it)", + q: &QuotaSnapshot{Exceeded: false, FetchedAt: now.Add(-2 * time.Hour).Format(time.RFC3339)}, + delay: 5 * time.Minute, + want: true, + }, + { + name: "a long delay holds a two-hour-old snapshot as fresh", + q: &QuotaSnapshot{Exceeded: true, FetchedAt: now.Add(-2 * time.Hour).Format(time.RFC3339)}, + delay: 3 * time.Hour, + want: false, + }, + { + name: "a long delay still expires an old snapshot", + q: &QuotaSnapshot{Exceeded: true, FetchedAt: now.Add(-7 * time.Hour).Format(time.RFC3339)}, + delay: 6 * time.Hour, + want: true, + }, + } + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + if got := quotaSnapshotStale(tc.q, now, tc.delay); got != tc.want { + t.Fatalf("quotaSnapshotStale = %v, want %v", got, tc.want) + } + }) + } +} + +func TestQuotaProbeDelay(t *testing.T) { + tests := []struct { + level int + want time.Duration + }{ + {level: -1, want: 5 * time.Minute}, + {level: 0, want: 5 * time.Minute}, + {level: 2, want: 5 * time.Minute}, + {level: 3, want: 15 * time.Minute}, + {level: 4, want: 30 * time.Minute}, + {level: 5, want: time.Hour}, + {level: 6, want: 3 * time.Hour}, + {level: 7, want: 6 * time.Hour}, + {level: 8, want: 12 * time.Hour}, + {level: 9, want: 24 * time.Hour}, + {level: 100, want: 24 * time.Hour}, + } + for _, tc := range tests { + if got := quotaProbeDelay(tc.level); got != tc.want { + t.Fatalf("quotaProbeDelay(%d) = %v, want %v", tc.level, got, tc.want) + } + } +} + +func TestQuotaProbeDelayIsMonotonic(t *testing.T) { + // The ladder must never shrink: an account that keeps failing should not + // suddenly be probed more often. + previous := time.Duration(0) + for level := 0; level < 12; level++ { + delay := quotaProbeDelay(level) + if delay < previous { + t.Fatalf("quotaProbeDelay(%d) = %v shrank below previous %v", level, delay, previous) + } + previous = delay + } +} + +func TestQuotaProbeLevelLifecycle(t *testing.T) { + manager := &Manager{} + if got := manager.quotaProbeLevel("acc"); got != 0 { + t.Fatalf("initial level = %d, want 0", got) + } + manager.bumpQuotaProbeLevel("acc") + manager.bumpQuotaProbeLevel("acc") + if got := manager.quotaProbeLevel("acc"); got != 2 { + t.Fatalf("after two bumps level = %d, want 2", got) + } + manager.resetQuotaProbeLevel("acc") + if got := manager.quotaProbeLevel("acc"); got != 0 { + t.Fatalf("after reset level = %d, want 0", got) + } + // prune drops accounts no longer seen and keeps the ones still in the pool. + manager.bumpQuotaProbeLevel("kept") + manager.bumpQuotaProbeLevel("gone") + manager.pruneQuotaProbeLevels(map[string]struct{}{"kept": {}}) + if _, ok := manager.quotaProbeLevels["gone"]; ok { + t.Fatalf("pruned account still present") + } + if got := manager.quotaProbeLevel("kept"); got != 1 { + t.Fatalf("kept account level = %d, want 1", got) + } +} diff --git a/internal/runtime/manager.go b/internal/runtime/manager.go index 1f8f33b..5437c75 100644 --- a/internal/runtime/manager.go +++ b/internal/runtime/manager.go @@ -118,6 +118,15 @@ type Manager struct { resources atomic.Pointer[ResourceSnapshot] resMu sync.Mutex // serializes sampleResources cpuPrev map[int]procTime // pid -> last jiffies sample for CPU% deltas + + // quotaProbeLevels tracks, per account ID, how many consecutive stale-quota + // sweeps have re-fetched a snapshot that still reported Exceeded. It drives + // the per-account backoff in refreshStaleQuota so a genuinely exhausted + // account is probed on a daily-ish cadence instead of every five minutes. + // Only mutated by the maintenance loop; guarded by quotaProbeMu because + // account removal can race with the sweep. + quotaProbeMu sync.Mutex + quotaProbeLevels map[string]int } func NewManager(config ManagerConfig, store AccountStore, starter ProcessStarter) *Manager { @@ -155,6 +164,7 @@ func NewManager(config ManagerConfig, store AccountStore, starter ProcessStarter persistedVersions: map[string]uint64{}, persistCloseCh: make(chan struct{}), cpuPrev: map[int]procTime{}, + quotaProbeLevels: map[string]int{}, } manager.persistCond = sync.NewCond(&manager.persistMu) manager.persistDone.Add(1) diff --git a/internal/runtime/manager_maintenance.go b/internal/runtime/manager_maintenance.go index 487a180..8f9898d 100644 --- a/internal/runtime/manager_maintenance.go +++ b/internal/runtime/manager_maintenance.go @@ -286,6 +286,9 @@ func (manager *Manager) RunMaintenanceLoop(stop <-chan struct{}) { ticker := time.NewTicker(30 * time.Second) defer ticker.Stop() lastKeepaliveDay := "" + // Seed the throttle so the first pass does not duplicate the boot-time + // RefreshAll: the next stale-quota sweep runs 5 minutes after startup. + lastStaleQuota := time.Now() for { if ctx.Err() != nil { return @@ -300,6 +303,10 @@ func (manager *Manager) RunMaintenanceLoop(stop <-chan struct{}) { lastKeepaliveDay = day } manager.sampleResources() + if now.Sub(lastStaleQuota) >= staleQuotaInterval { + manager.refreshStaleQuota(ctx) + lastStaleQuota = now + } select { case <-ctx.Done(): return @@ -307,3 +314,168 @@ func (manager *Manager) RunMaintenanceLoop(stop <-chan struct{}) { } } } + +// staleQuotaInterval is how often the maintenance loop considers exhausted +// accounts. It is only the sweep cadence: the per-account backoff below +// decides whether a given account is actually re-fetched on that sweep. +const staleQuotaInterval = 5 * time.Minute + +// quotaProbeBackoff is the per-account wait between consecutive quota probes +// while an account keeps reporting Exceeded, indexed by the number of +// consecutive exhausted probes already spent. The first few probes stay at +// 5 minutes so a transient false Exceeded or an out-of-band top-up recovers +// quickly. A genuinely exhausted account then backs off to hours, because +// provider balances reset on a daily cadence and polling it every 5 minutes +// would otherwise burn hundreds of pointless quota calls a day. +var quotaProbeBackoff = []time.Duration{ + 5 * time.Minute, // level 0: first probe + 5 * time.Minute, // level 1 + 5 * time.Minute, // level 2 + 15 * time.Minute, // level 3 + 30 * time.Minute, // level 4 + 1 * time.Hour, // level 5 + 3 * time.Hour, // level 6 + 6 * time.Hour, // level 7 + 12 * time.Hour, // level 8 + 24 * time.Hour, // level 9+ (cap: balances reset daily) +} + +// quotaProbeDelay returns the wait before the next probe at the given +// consecutive-exhausted level, clamping out-of-range levels to the ends. +func quotaProbeDelay(level int) time.Duration { + if level < 0 { + level = 0 + } + if level >= len(quotaProbeBackoff) { + return quotaProbeBackoff[len(quotaProbeBackoff)-1] + } + return quotaProbeBackoff[level] +} + +// quotaSnapshotStale reports whether an Exceeded snapshot is old enough to +// warrant a refresh at the given per-account delay. A zero/empty or +// unparseable FetchedAt means the snapshot predates timestamp tracking, so +// treat it as stale. +func quotaSnapshotStale(q *QuotaSnapshot, now time.Time, delay time.Duration) bool { + if q == nil { + return false + } + fetchedAt := strings.TrimSpace(q.FetchedAt) + if fetchedAt == "" { + return true + } + parsed, err := time.Parse(time.RFC3339, fetchedAt) + if err != nil { + return true + } + return now.Sub(parsed) >= delay +} + +// refreshStaleQuota re-fetches quota for accounts whose cached snapshot still +// reports Exceeded. A snapshot can go stale when upstream credits are +// replenished out-of-band (billing reset, add-on purchase) or when a flaky +// quota endpoint left an old Exceeded=true in place. Without this, the only +// way to clear a stale exhausted flag is a manual console refresh, so routing +// keeps skipping the account until an operator intervenes. Healthy accounts +// are never touched here; only Exceeded and stale snapshots are refreshed. +// +// Each account carries its own consecutive-exhausted level so an account that +// stays empty for days is probed less and less often instead of every five +// minutes forever. The level resets as soon as a snapshot reports a non-zero +// balance (or the account leaves the pool), so the next genuine exhaustion +// starts the fast ladder over. +func (manager *Manager) refreshStaleQuota(ctx context.Context) { + if manager == nil || manager.pool == nil { + return + } + now := time.Now() + seen := make(map[string]struct{}) + for _, item := range manager.pool.Items() { + if ctx.Err() != nil { + return + } + seen[item.ID] = struct{}{} + if item.RuntimeState == "dead" { + continue + } + if item.Quota == nil || !item.Quota.Exceeded { + // Recovered (or never exhausted): the next exhaustion is a new + // episode and starts from the 5-minute probes again. + manager.resetQuotaProbeLevel(item.ID) + continue + } + level := manager.quotaProbeLevel(item.ID) + if !quotaSnapshotStale(item.Quota, now, quotaProbeDelay(level)) { + continue + } + before := strings.TrimSpace(item.Quota.FetchedAt) + refreshCtx, cancel := context.WithTimeout(ctx, 10*time.Second) + manager.refreshQuotaOnly(refreshCtx, item) + cancel() + updated, ok := manager.pool.ByID(item.ID) + fetched := ok && updated.Quota != nil && strings.TrimSpace(updated.Quota.FetchedAt) != before + switch { + case !fetched: + // Probe failed (upstream/network). Keep the level so the next + // sweep retries at the same cadence instead of escalating a + // transient outage into an hours-long wait. + case updated.Quota.Exceeded: + manager.bumpQuotaProbeLevel(item.ID) + default: + manager.resetQuotaProbeLevel(item.ID) + } + } + manager.pruneQuotaProbeLevels(seen) +} + +// refreshQuotaOnly fetches quota for one account without a health probe or +// model discovery. It resolves the in-process provider prober when present +// (WorkBuddy/Trae/Devin) and otherwise falls back to the Qoder worker URL, +// mirroring refreshCheckinQuota's dispatch. +func (manager *Manager) refreshQuotaOnly(ctx context.Context, item Item) { + if adapter, ok := manager.providers.Get(item.Provider); ok && adapter.Prober != nil { + manager.fetchProviderQuota(ctx, item.ID, adapter.Prober) + return + } + if strings.TrimSpace(item.URL) != "" { + _ = manager.fetchQuota(ctx, item.ID, item.URL, true) + } +} + +// quotaProbeLevel reads the consecutive-exhausted probe level for an account. +func (manager *Manager) quotaProbeLevel(accountID string) int { + manager.quotaProbeMu.Lock() + defer manager.quotaProbeMu.Unlock() + return manager.quotaProbeLevels[accountID] +} + +// bumpQuotaProbeLevel advances the backoff ladder after a probe that still +// reported Exceeded. +func (manager *Manager) bumpQuotaProbeLevel(accountID string) { + manager.quotaProbeMu.Lock() + defer manager.quotaProbeMu.Unlock() + if manager.quotaProbeLevels == nil { + manager.quotaProbeLevels = map[string]int{} + } + manager.quotaProbeLevels[accountID]++ +} + +// resetQuotaProbeLevel clears the ladder once an account is no longer +// exhausted, so the next exhaustion starts with the fast 5-minute probes. +func (manager *Manager) resetQuotaProbeLevel(accountID string) { + manager.quotaProbeMu.Lock() + defer manager.quotaProbeMu.Unlock() + delete(manager.quotaProbeLevels, accountID) +} + +// pruneQuotaProbeLevels drops ladder entries for accounts no longer in the +// pool so the map stays bounded by the current account count. +func (manager *Manager) pruneQuotaProbeLevels(seen map[string]struct{}) { + manager.quotaProbeMu.Lock() + defer manager.quotaProbeMu.Unlock() + for id := range manager.quotaProbeLevels { + if _, ok := seen[id]; !ok { + delete(manager.quotaProbeLevels, id) + } + } +} diff --git a/internal/runtime/quota_backoff_test.go b/internal/runtime/quota_backoff_test.go new file mode 100644 index 0000000..097f41f --- /dev/null +++ b/internal/runtime/quota_backoff_test.go @@ -0,0 +1,160 @@ +package runtime + +import ( + "context" + "path/filepath" + "testing" + "time" + + "github.com/caigee-cmd/cli2api/internal/accounts" + "github.com/caigee-cmd/cli2api/internal/executor" + "github.com/caigee-cmd/cli2api/internal/providers" + sqlstore "github.com/caigee-cmd/cli2api/internal/store" +) + +// scriptedProber returns a fixed quota snapshot on every Quota() call and +// counts how many times it was asked. The FetchedAt it reports is stamped at +// call time so the merged snapshot looks freshly fetched, exactly like a real +// successful probe. +type scriptedProber struct { + exceeded bool + calls int + err error +} + +func (p *scriptedProber) Probe(ctx context.Context, accountID string) (providers.AccountHealth, error) { + return providers.AccountHealth{Ready: true, Hot: true}, nil +} + +func (p *scriptedProber) Quota(ctx context.Context, accountID string) (*providers.QuotaInfo, error) { + p.calls++ + if p.err != nil { + return nil, p.err + } + used, remain := 100.0, 900.0 + if p.exceeded { + used, remain = 100.0, 0.0 + } + return &providers.QuotaInfo{ + Used: used, Total: 100, Remaining: remain, Percentage: used, + Unit: "credits", Exceeded: p.exceeded, + FetchedAt: time.Now().UTC().Format(time.RFC3339), + }, nil +} + +func newQuotaBackoffTestManager(t *testing.T, prober *scriptedProber) (*Manager, string) { + t.Helper() + ctx := context.Background() + store, err := sqlstore.OpenStore(filepath.Join(t.TempDir(), "qoder.db")) + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { store.Close() }) + account, err := store.Create(ctx, accounts.CreateAccount{ + Name: "probe", Provider: "workbuddy", Region: "cn", Enabled: true, + }) + if err != nil { + t.Fatal(err) + } + registry := providers.NewRegistry() + registry.Register(providers.Adapter{ID: "workbuddy", Prober: prober}) + manager := NewManager(ManagerConfig{DataDir: t.TempDir()}, store, &noopStarter{}) + manager.SetProviders(registry) + manager.Pool().Upsert(executor.Item{ID: account.ID, Provider: "workbuddy", Runtime: "in_process"}) + return manager, account.ID +} + +// agedExceeded returns an Exceeded snapshot whose FetchedAt is age old, so the +// test can drive the sweep's staleness check without sleeping. +func agedExceeded(age time.Duration) *accounts.QuotaSnapshot { + return &accounts.QuotaSnapshot{ + Used: 100, Total: 100, Remaining: 0, Percentage: 100, + Unit: "credits", Exceeded: true, + FetchedAt: time.Now().UTC().Add(-age).Format(time.RFC3339), + } +} + +func TestRefreshStaleQuotaBacksOffExhaustedAccount(t *testing.T) { + prober := &scriptedProber{exceeded: true} + manager, id := newQuotaBackoffTestManager(t, prober) + ctx := context.Background() + + // Levels 0..2 all run at a 5-minute delay. A 10-minute-old snapshot is + // therefore probed on each of the first three sweeps, and each successful + // but still-exhausted probe bumps the level. + for i := 0; i < 3; i++ { + manager.Pool().MergeQuota(id, agedExceeded(10*time.Minute)) + manager.refreshStaleQuota(ctx) + if prober.calls != i+1 { + t.Fatalf("sweep %d: calls = %d, want %d", i, prober.calls, i+1) + } + if got := manager.quotaProbeLevel(id); got != i+1 { + t.Fatalf("sweep %d: level = %d, want %d", i, got, i+1) + } + } + + // Level 3 raises the delay to 15 minutes: a 10-minute-old snapshot must + // now be skipped instead of hammered. + manager.Pool().MergeQuota(id, agedExceeded(10*time.Minute)) + manager.refreshStaleQuota(ctx) + if prober.calls != 3 { + t.Fatalf("fresh-enough snapshot was probed: calls = %d, want 3", prober.calls) + } + if got := manager.quotaProbeLevel(id); got != 3 { + t.Fatalf("level changed while skipping: %d, want 3", got) + } + + // Once the snapshot is older than the level-3 delay it is probed again. + manager.Pool().MergeQuota(id, agedExceeded(20*time.Minute)) + manager.refreshStaleQuota(ctx) + if prober.calls != 4 { + t.Fatalf("stale snapshot was not probed: calls = %d, want 4", prober.calls) + } + if got := manager.quotaProbeLevel(id); got != 4 { + t.Fatalf("level = %d, want 4", got) + } + + // Level 4 is a 30-minute delay; a 20-minute-old snapshot is skipped. + manager.Pool().MergeQuota(id, agedExceeded(20*time.Minute)) + manager.refreshStaleQuota(ctx) + if prober.calls != 4 { + t.Fatalf("level-4 skip probed anyway: calls = %d, want 4", prober.calls) + } + + // A recovered balance resets the ladder to the fast end. + prober.exceeded = false + manager.Pool().MergeQuota(id, agedExceeded(31*time.Minute)) + manager.refreshStaleQuota(ctx) + if prober.calls != 5 { + t.Fatalf("recovery probe did not run: calls = %d, want 5", prober.calls) + } + if got := manager.quotaProbeLevel(id); got != 0 { + t.Fatalf("level after recovery = %d, want 0", got) + } +} + +func TestRefreshStaleQuotaKeepsLevelWhenProbeFails(t *testing.T) { + prober := &scriptedProber{exceeded: true, err: context.DeadlineExceeded} + manager, id := newQuotaBackoffTestManager(t, prober) + ctx := context.Background() + + manager.Pool().MergeQuota(id, agedExceeded(10*time.Minute)) + manager.refreshStaleQuota(ctx) + // A failed probe never merges a new snapshot, so FetchedAt is unchanged and + // the level must not climb: a transient outage should retry at the same + // cadence rather than escalate to an hours-long wait. + if got := manager.quotaProbeLevel(id); got != 0 { + t.Fatalf("level after failed probe = %d, want 0", got) + } + // The stale snapshot is retried on the next sweep at the same fast cadence. + manager.refreshStaleQuota(ctx) + if prober.calls != 2 { + t.Fatalf("failed probe was not retried: calls = %d, want 2", prober.calls) + } +} + +type noopStarter struct{} + +func (n *noopStarter) Start(context.Context, Account, string, int) (ManagedProcess, error) { + return nil, nil +}