diff --git a/.env.example b/.env.example index 6e8911fa0..d2094cb35 100644 --- a/.env.example +++ b/.env.example @@ -50,6 +50,13 @@ CODEX_PORT=8080 # Codex 上游传输：默认 http；可设 ws 强制使用 WebSocket；auto 在 HTTP 入站下仍走 http # CODEX_UPSTREAM_TRANSPORT=http +# WS 逐帧用量快照落库的节流间隔（秒），默认 30。内存快照每帧实时更新； +# 跨阈值、5pp 变化和 100% 仍会强制落库。大号池可按数据库承压情况调大。 +# CODEX_WS_USAGE_FRAME_PERSIST_SECONDS=30 + +# WS 逐帧用量摄入开关，默认 true；设为 false 回退到 WHAM 探针与 HTTP 响应头 +# CODEX_WS_USAGE_FRAME_INGEST=true + # 旧版 WebSocket 开关（兼容保留；推荐改用 CODEX_UPSTREAM_TRANSPORT=ws） # USE_WEBSOCKET=true diff --git a/CHANGELOG.md b/CHANGELOG.md index 9ec96a37c..42a8e04a5 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,5 +1,23 @@ # Changelog +## Unreleased + +### Features + +- **Codex WebSocket quota frames update account usage in real time (#800).** On the persistent WebSocket upstream, quota data arrives in per-turn `codex.rate_limits` and `codex.response.metadata` frames instead of HTTP headers, so account usage only advanced on WHAM probes and a burst could exhaust a 5h window before auto-pause saw it. Every Codex relay path now maps those frames onto the existing header sync: Responses streaming and non-streaming, Chat Completions, downstream WebSocket, and Claude Messages. Connection-level frames that arrive 1–2 seconds after a terminal are captured by a non-blocking observer instead of being dropped. Partial frames never clear a missing window, metered limits (Spark/reserve) cannot overwrite the main 5h/7d windows, a workspace hard stop is preserved because the frame has no `rate_limit_reached_type`, and a late frame from an earlier concurrent turn cannot roll usage back within the same reset window. Relay-style accounts are skipped because forwarded frames describe the relay's own upstream account. A frame without `reset_at` after the stored reset has passed starts the new window instead of inheriting the expired timestamp, and connection-level frames also apply the post-turn credits snapshot and refresh the scheduler so queued dispatches wake when an account recovers. In-memory snapshots update on every frame; database writes are throttled to 30 seconds on a per-account coalescing worker, so a slow or locked database cannot stall the response stream, with forced writes at threshold crossings, 5pp changes and 100% (`CODEX_WS_USAGE_FRAME_PERSIST_SECONDS`). `CODEX_WS_USAGE_FRAME_INGEST=false` turns frame ingestion off and falls back to the WHAM probe and HTTP headers. `/api/admin/accounts/live` now returns runtime usage, status, credits and the last applied frame, and the account page merges it in its 1-second poll. The merge keeps unchanged row identities, and missing-usage and stale-stats retries no longer restart on every poll; a hung live request is aborted after 5 seconds so the next tick can retry. + +### Fixes + +- **Account live polling refreshes the existing plan and credit badges.** The live endpoint now includes the runtime plan and credit snapshot using the account-list field contract. Missing fields preserve older-backend values, hidden balances clear stale values, and unchanged rows keep their identity; Antigravity rows retain their own projection. Persistent-WebSocket coverage checks post-terminal plan and balance changes through the authenticated live endpoint. + +- **Quota-frame reset timestamps require strict integers.** Fractional, exponent-form and overflowing JSON numbers are rejected instead of being truncated or wrapped into reset deadlines; valid usage fields remain usable with the existing missing-reset fallback. + +- **Post-terminal quota frames survive observation contention without stalling the WebSocket reader.** Active and idle frames share an ordered runtime path; a bounded per-account FIFO retains partial updates and requests an authoritative probe on overflow. Credits, plan, quota and cooldown writes are coalesced off the reader, including exhausted-window and credits-only frames. Old observations cannot roll credits or frame freshness backward, 7d-only writes retain their reset, and delayed cooldown publication reads current state so it cannot overwrite a newer authentication error or resurrect a cleared cooldown. Failed cooldown writes retry on later observations without another failure count; repeated real Responses 429s still fence older in-flight successes. A persistent-WebSocket regression exercises the production idle observer through scheduling and the authenticated live endpoint, with no subsequent business request or probe. + +- **Quota observations cannot shorten an active account cooldown.** The 5h and 7d usage paths keep a longer existing cooldown and its reason, including an authoritative Responses 429 or authentication error. Later resets extend the cooldown without replacing that reason or counting another failure. A real Responses 429 can upgrade a local usage cooldown to `responses_rate_limited` without shortening it, so `ignore_usage_limit_status` cannot bypass an actual rejection; a stronger authentication failure retains priority. Usage snapshots still update when the cooldown is unchanged, and repeated observations do not write the cooldown again. + +- **Account page copy matches the actual quota and cleanup behavior.** The auto-pause threshold hint now says blank or 0 inherits the group/global threshold and that the switch above disables that window, instead of claiming 0 disables it. The rate-limited cleanup dialog and result toast say accounts move to the recycle bin and that auto-paused accounts appear in the cleanup when their runtime status is rate-limited or `usage_exhausted`, instead of saying status flags are cleared. + ## v3.1.1 - 2026-10-17 ### Features diff --git a/admin/account_live.go b/admin/account_live.go index 4adca4d4d..1e7ebaa81 100644 --- a/admin/account_live.go +++ b/admin/account_live.go @@ -2,11 +2,14 @@ package admin import ( "net/http" + "time" "github.com/gin-gonic/gin" ) type accountLiveItem struct { + // A nil embedded projection omits Codex credit fields for other planes. + *AccountLiveCredits ActiveRequests int64 `json:"active_requests"` OccupiedRequests int64 `json:"occupied_requests"` // DynamicConcurrencyLimit is the admission cap the scheduler currently @@ -15,11 +18,43 @@ type accountLiveItem struct { // BaseConcurrencyEffective is the configured cap (account override, // then group, then the global default) before any runtime reduction. BaseConcurrencyEffective int64 `json:"base_concurrency_effective"` + + // Runtime usage/status projection so the account page can reflect WS frame + // observations without waiting for a list reload. Optional pointers keep the + // backward-compat contract: an older backend omits the key (frontend keeps + // the list value); a pointer that is nil serializes as null (frontend clears + // the now-invalid window). + Status *string `json:"status,omitempty"` + UsingCredits *bool `json:"using_credits,omitempty"` + PlanType *string `json:"plan_type,omitempty"` + UsagePercent5h *float64 `json:"usage_percent_5h"` + Reset5hAt *string `json:"reset_5h_at"` + UsageUpdatedAt5h *string `json:"codex_5h_usage_updated_at"` + UsagePercent7d *float64 `json:"usage_percent_7d"` + Reset7dAt *string `json:"reset_7d_at"` + UsageUpdatedAt *string `json:"codex_usage_updated_at"` + // Frame freshness proves the WS quota ingestion path is working for this + // account; omitted until the first frame is applied. + UsageFrameAt *string `json:"codex_usage_frame_at,omitempty"` + UsageFrameKind *string `json:"codex_usage_frame_kind,omitempty"` +} + +// AccountLiveCredits uses the existing account-list credit field contract. +type AccountLiveCredits struct { + Valid bool `json:"credits_valid"` + Balance *string `json:"credits_balance"` + HasCredits *bool `json:"credits_has_credits"` + Unlimited *bool `json:"credits_unlimited"` + OverageLimitReached *bool `json:"credits_overage_limit_reached"` + SpendControlReached *bool `json:"credits_spend_control_reached"` + RateLimitReachedType string `json:"credits_rate_limit_reached_type"` } // GetAccountLiveState returns request-local runtime counters for the visible // account page. Scheduler counters use in-memory atomics; concurrency caps are -// the scheduler's maintained values and are not recomputed here. +// the scheduler's maintained values and are not recomputed here. The usage +// projection is a single read lock per account and never rebuilds a snapshot, +// so it stays cheap at the one-second cadence even for a full page. func (h *Handler) GetAccountLiveState(c *gin.Context) { ids, err := parseAccountListIDs(c.Query("ids")) if err != nil { @@ -37,15 +72,60 @@ func (h *Handler) GetAccountLiveState(c *gin.Context) { if account == nil { continue } - live[id] = accountLiveItem{ + item := accountLiveItem{ ActiveRequests: account.GetActiveRequests(), OccupiedRequests: account.GetOccupiedRequests(), DynamicConcurrencyLimit: account.GetDynamicConcurrencyLimit(), BaseConcurrencyEffective: account.GetBaseConcurrencyEffective(), } + // Antigravity rows use a persisted quota snapshot; runtimeStatusLocked + // does not model their display status, so leave the row's fields alone. + if !account.IsAntigravityAPI() { + usage := account.GetAccountLiveUsageSnapshot() + item.PlanType = &usage.PlanType + item.AccountLiveCredits = &AccountLiveCredits{Valid: usage.CreditsValid} + if usage.CreditsValid { + item.Balance = usage.Credits.Balance + item.HasCredits = &usage.Credits.HasCredits + item.Unlimited = &usage.Credits.Unlimited + item.OverageLimitReached = &usage.Credits.OverageLimitReached + item.SpendControlReached = usage.Credits.SpendControlReached + item.RateLimitReachedType = usage.Credits.RateLimitReachedType + } + status := usage.Status + item.Status = &status + usingCredits := usage.UsingCredits + item.UsingCredits = &usingCredits + if usage.UsagePercent5hValid { + pct := usage.UsagePercent5h + item.UsagePercent5h = &pct + } + item.Reset5hAt = optionalLiveTime(usage.Reset5hAt) + item.UsageUpdatedAt5h = optionalLiveTime(usage.UsageUpdatedAt5h) + if usage.UsagePercent7dValid { + pct := usage.UsagePercent7d + item.UsagePercent7d = &pct + } + item.Reset7dAt = optionalLiveTime(usage.Reset7dAt) + item.UsageUpdatedAt = optionalLiveTime(usage.UsageUpdatedAt) + if !usage.UsageFrameLastAt.IsZero() { + item.UsageFrameAt = optionalLiveTime(usage.UsageFrameLastAt) + kind := usage.UsageFrameLastKind + item.UsageFrameKind = &kind + } + } + live[id] = item } c.JSON(http.StatusOK, gin.H{ "accounts": live, "session_slot_buffer_enabled": h.store.SessionSlotBufferEnabled(), }) } + +func optionalLiveTime(t time.Time) *string { + if t.IsZero() { + return nil + } + formatted := t.Format(time.RFC3339) + return &formatted +} diff --git a/admin/account_live_test.go b/admin/account_live_test.go index 550668a19..38b302e76 100644 --- a/admin/account_live_test.go +++ b/admin/account_live_test.go @@ -5,6 +5,7 @@ import ( "net/http" "net/http/httptest" "testing" + "time" "github.com/codex2api/auth" "github.com/gin-gonic/gin" @@ -58,3 +59,128 @@ func TestGetAccountLiveStateReturnsVisibleInflightCounts(t *testing.T) { t.Fatal("missing account unexpectedly returned") } } + +func TestGetAccountLiveStateReturnsRuntimeUsage(t *testing.T) { + gin.SetMode(gin.TestMode) + store := auth.NewStore(nil, nil, nil) + now := time.Now() + account := &auth.Account{ + DBID: 7, + AccessToken: "token", + PlanType: "plus", + Status: auth.StatusReady, + HealthTier: auth.HealthTierHealthy, + } + account.UsagePercent5h = 42 + account.UsagePercent5hValid = true + account.Reset5hAt = now.Add(time.Hour) + account.UsageUpdatedAt5h = now + account.UsagePercent7d = 12 + account.UsagePercent7dValid = true + account.Reset7dAt = now.Add(24 * time.Hour) + account.UsageUpdatedAt = now + store.AddAccount(account) + store.ObserveUsageFramePlan(account, "pro") + balance := "12.5" + account.ApplySparseCreditsHeaders(true, false, &balance, nil, "") + handler := &Handler{store: store} + + recorder := httptest.NewRecorder() + ctx, _ := gin.CreateTestContext(recorder) + ctx.Request = httptest.NewRequest(http.MethodGet, "/api/admin/accounts/live?ids=7", nil) + handler.GetAccountLiveState(ctx) + + if recorder.Code != http.StatusOK { + t.Fatalf("status = %d, body=%s", recorder.Code, recorder.Body.String()) + } + var response struct { + Accounts map[string]accountLiveItem `json:"accounts"` + } + if err := json.Unmarshal(recorder.Body.Bytes(), &response); err != nil { + t.Fatal(err) + } + item := response.Accounts["7"] + if item.UsagePercent5h == nil || *item.UsagePercent5h != 42 { + t.Fatalf("usage_percent_5h = %v, want 42", item.UsagePercent5h) + } + if item.UsagePercent7d == nil || *item.UsagePercent7d != 12 { + t.Fatalf("usage_percent_7d = %v, want 12", item.UsagePercent7d) + } + if item.Reset5hAt == nil || item.Reset7dAt == nil { + t.Fatalf("reset times missing: %+v", item) + } + if item.Status == nil || *item.Status == "" { + t.Fatalf("status missing: %+v", item) + } + if item.UsingCredits == nil { + t.Fatalf("using_credits missing: %+v", item) + } + var raw struct { + Accounts map[string]map[string]json.RawMessage `json:"accounts"` + } + if err := json.Unmarshal(recorder.Body.Bytes(), &raw); err != nil { + t.Fatal(err) + } + for key, want := range map[string]string{"plan_type": `"pro"`, "credits_valid": "true", "credits_balance": `"12.5"`, "credits_has_credits": "true", "credits_unlimited": "false", "credits_spend_control_reached": "null"} { + if got := string(raw.Accounts["7"][key]); got != want { + t.Errorf("live %s=%s, want %s", key, got, want) + } + } +} + +func TestGetAccountLiveStateSerializesNullWindowsAndSkipsAntigravityFields(t *testing.T) { + gin.SetMode(gin.TestMode) + store := auth.NewStore(nil, nil, nil) + codex := &auth.Account{ + DBID: 11, + AccessToken: "token", + Status: auth.StatusReady, + HealthTier: auth.HealthTierHealthy, + } + antigravity := &auth.Account{ + DBID: 12, + AccessToken: "token", + UpstreamType: auth.UpstreamAntigravity, + } + store.AddAccount(codex) + store.AddAccount(antigravity) + codex.MarkUsageFrameObserved("codex.rate_limits", time.Now()) + handler := &Handler{store: store} + + recorder := httptest.NewRecorder() + ctx, _ := gin.CreateTestContext(recorder) + ctx.Request = httptest.NewRequest(http.MethodGet, "/api/admin/accounts/live?ids=11,12", nil) + handler.GetAccountLiveState(ctx) + + if recorder.Code != http.StatusOK { + t.Fatalf("status = %d, body=%s", recorder.Code, recorder.Body.String()) + } + var raw struct { + Accounts map[string]map[string]json.RawMessage `json:"accounts"` + } + if err := json.Unmarshal(recorder.Body.Bytes(), &raw); err != nil { + t.Fatal(err) + } + codexItem := raw.Accounts["11"] + if got := string(codexItem["usage_percent_5h"]); got != "null" { + t.Fatalf("codex usage_percent_5h = %s, want null", got) + } + if got := string(codexItem["status"]); got == "" || got == "null" { + t.Fatalf("codex status missing: %s", got) + } + if got := string(codexItem["codex_usage_frame_kind"]); got != `"codex.rate_limits"` { + t.Fatalf("codex frame kind = %s, want codex.rate_limits", got) + } + if got := string(codexItem["codex_usage_frame_at"]); got == "" || got == "null" { + t.Fatalf("codex frame time missing: %s", got) + } + antigravityItem := raw.Accounts["12"] + for _, key := range []string{"status", "using_credits", "plan_type", "credits_valid", "credits_balance", "codex_usage_frame_at", "codex_usage_frame_kind"} { + if value, ok := antigravityItem[key]; ok { + t.Fatalf("antigravity item unexpectedly carries %s=%s", key, value) + } + } + if got := string(antigravityItem["usage_percent_5h"]); got != "null" { + t.Fatalf("antigravity usage_percent_5h = %s, want null", got) + } +} diff --git a/auth/oauth_refresh_lock.go b/auth/oauth_refresh_lock.go index 02d720462..971f1b111 100644 --- a/auth/oauth_refresh_lock.go +++ b/auth/oauth_refresh_lock.go @@ -453,11 +453,8 @@ func (s *Store) finishReloadedOAuthRefresh(ctx context.Context, acc *Account) { } } if expiredCooldown { - s.deleteCachedAccountCooldown(dbID) - if s.db != nil { - _ = s.db.ClearCooldown(ctx, dbID) - } + s.persistAccountCooldownState(acc, true, false) } else if !activeCooldown && s.db != nil { - _ = s.db.ClearError(ctx, dbID) + s.persistAccountCooldownState(acc, true, true) } } diff --git a/auth/premium_rate_limit.go b/auth/premium_rate_limit.go index 70885b7be..df16daf2d 100644 --- a/auth/premium_rate_limit.go +++ b/auth/premium_rate_limit.go @@ -133,31 +133,44 @@ func (a *Account) GetUsageSnapshot5h() (pct float64, resetAt time.Time, ok bool) // PersistUsageSnapshot5hOnly 持久化仅包含 5h 数据的用量快照。 func (s *Store) PersistUsageSnapshot5hOnly(acc *Account) { + s.PublishUsageSnapshot5hOnly(acc, time.Now(), true) +} + +// PublishUsageSnapshot5hOnly updates the in-memory 5h snapshot and scheduler. +// persist=false skips the DB write for the throttled WS frame path. +func (s *Store) PublishUsageSnapshot5hOnly(acc *Account, observedAt time.Time, persist bool) bool { if acc == nil || s == nil { - return + return false } pct5h, reset5hAt, ok := acc.GetUsageSnapshot5h() if !ok { - return + return false } - updatedAt := time.Now() + if observedAt.IsZero() { + observedAt = time.Now() + } acc.mu.Lock() - acc.UsageUpdatedAt5h = updatedAt + acc.UsageUpdatedAt5h = observedAt acc.mu.Unlock() s.fastSchedulerUpdate(acc) - if s.db == nil { - return + if !persist || s.db == nil { + return false } ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second) defer cancel() - if err := s.db.UpdateUsageSnapshot5h(ctx, acc.DBID, pct5h, reset5hAt, updatedAt); err != nil { + err := s.withUsageSnapshotWrite(acc, func() error { + return s.db.UpdateUsageSnapshot5h(ctx, acc.DBID, pct5h, reset5hAt, observedAt) + }) + if err != nil { log.Printf("[账号 %d] 持久化 5h 用量快照失败: %v", acc.DBID, err) + return false } + return true } // Persist5hWindowActivated 把「已为哪个 Reset5hAt 发过开窗请求」写入 credentials，重启后不重复打。 @@ -215,7 +228,6 @@ func (s *Store) ClearAbsentUsageSnapshot5hAt(acc *Account, observedAt time.Time) acc.usageObservedAt = observedAt had5h := acc.UsagePercent5hValid cleared5hCooldown := s != nil && acc.Status == StatusCooldown && acc.CooldownReason == premium5hCooldownReason - priorCooldownUntil := acc.CooldownUtil if !had5h && !cleared5hCooldown { acc.mu.Unlock() return false @@ -245,32 +257,22 @@ func (s *Store) ClearAbsentUsageSnapshot5hAt(acc *Account, observedAt time.Time) return true } s.fastSchedulerUpdate(acc) - if cleared5hCooldown { - s.deleteCachedAccountCooldown(acc.DBID) - // A newer 401/generic 429 can land after the in-memory transition but - // before the cache delete. Restore that newer cooldown if present. - acc.mu.RLock() - status := acc.Status - reason := acc.CooldownReason - until := acc.CooldownUtil - acc.mu.RUnlock() - if status == StatusCooldown && reason != "" { - s.setCachedAccountCooldown(acc.DBID, reason, until) - } - } if s.db == nil { + if cleared5hCooldown { + s.persistAccountCooldownState(acc, true, false) + } return true } ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second) defer cancel() - if err := s.db.ClearUsageSnapshot5h(ctx, acc.DBID); err != nil { + if err := s.withUsageSnapshotWrite(acc, func() error { + return s.db.ClearUsageSnapshot5h(ctx, acc.DBID) + }); err != nil { log.Printf("[账号 %d] 清除 5h 用量快照失败: %v", acc.DBID, err) } if cleared5hCooldown { - if _, err := s.db.ClearCooldownIfReasonAndUntil(ctx, acc.DBID, premium5hCooldownReason, priorCooldownUntil); err != nil { - log.Printf("[账号 %d] 清除 premium 5h 冷却状态失败: %v", acc.DBID, err) - } + s.persistAccountCooldownState(acc, true, false) } return true } @@ -289,7 +291,7 @@ func (s *Store) MarkPremium5hRateLimitedAt(acc *Account, resetAt, observedAt tim return acc.ApplyUsageObservation(observedAt, func() { // observedAt orders competing upstream observations; stateAt reflects // when the state is actually applied after any preceding DB/identity work. - s.markPremium5hRateLimited(acc, resetAt, time.Now(), premium5hCooldownReason) + s.markPremium5hRateLimited(acc, resetAt, time.Now(), premium5hCooldownReason, true) }) } @@ -302,11 +304,11 @@ func (s *Store) MarkResponsesPremium5hRateLimited(acc *Account, resetAt time.Tim } now := time.Now() _ = acc.ApplyUsageObservation(now, func() { - s.markPremium5hRateLimited(acc, resetAt, time.Now(), ResponsesRateLimitedCooldownReason) + s.markPremium5hRateLimited(acc, resetAt, time.Now(), ResponsesRateLimitedCooldownReason, true) }) } -func (s *Store) markPremium5hRateLimited(acc *Account, resetAt, observedAt time.Time, cooldownReason string) { +func (s *Store) markPremium5hRateLimited(acc *Account, resetAt, observedAt time.Time, cooldownReason string, persist bool) { now := observedAt if now.IsZero() { now = time.Now() @@ -316,33 +318,44 @@ func (s *Store) markPremium5hRateLimited(acc *Account, resetAt, observedAt time. } acc.mu.Lock() + snapshotChanged := !acc.UsagePercent5hValid || acc.UsagePercent5h != 100 || !acc.Reset5hAt.Equal(resetAt) acc.UsagePercent5h = 100 acc.UsagePercent5hValid = true acc.Reset5hAt = resetAt acc.UsageUpdatedAt5h = now - acc.LastRateLimitedAt = now - acc.Status = StatusCooldown - acc.CooldownUtil = resetAt - acc.CooldownReason = cooldownReason - if acc.HealthTier != HealthTierBanned { - acc.HealthTier = HealthTierRisky + until, reason, cooldownChanged := acc.usageCooldownTargetLocked(now, resetAt, cooldownReason) + // A repeated real 429 is new rejection evidence even when its deadline is + // unchanged. An older in-flight success must not clear that evidence. + if cooldownReason == ResponsesRateLimitedCooldownReason { + acc.LastRateLimitedAt = now + } + if cooldownChanged { + acc.LastRateLimitedAt = now + acc.Status = StatusCooldown + acc.CooldownUtil = until + acc.CooldownReason = reason + acc.usageCooldownPersistPending = true + if acc.HealthTier != HealthTierBanned { + acc.HealthTier = HealthTierRisky + } } acc.recomputeSchedulerLocked(s.maxConcurrency.Load()) acc.mu.Unlock() s.fastSchedulerUpdate(acc) - s.setCachedAccountCooldown(acc.DBID, cooldownReason, resetAt) + if persist { + s.persistUsageCooldown(acc) + } - if s.db == nil { + if !persist || s.db == nil || (!cooldownChanged && !snapshotChanged) { return } ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second) defer cancel() - if err := s.db.SetCooldown(ctx, acc.DBID, cooldownReason, resetAt); err != nil { - log.Printf("[账号 %d] 持久化 premium 5h 限流冷却状态失败: %v", acc.DBID, err) - } - if err := s.db.UpdateUsageSnapshot5h(ctx, acc.DBID, 100, resetAt, now); err != nil { + if err := s.withUsageSnapshotWrite(acc, func() error { + return s.db.UpdateUsageSnapshot5h(ctx, acc.DBID, 100, resetAt, now) + }); err != nil { log.Printf("[账号 %d] 持久化 premium 5h 限流快照失败: %v", acc.DBID, err) } } diff --git a/auth/scheduler_outbox_consumer.go b/auth/scheduler_outbox_consumer.go index c3e82d412..db35722d4 100644 --- a/auth/scheduler_outbox_consumer.go +++ b/auth/scheduler_outbox_consumer.go @@ -528,15 +528,25 @@ func (s *Store) applyPersistentAccountSnapshot(dst, src *Account, enabled bool) dst.CooldownUtil = src.CooldownUtil dst.CooldownReason = src.CooldownReason dst.ErrorMsg = src.ErrorMsg - dst.UsagePercent7d = src.UsagePercent7d - dst.UsagePercent7dValid = src.UsagePercent7dValid - dst.Reset7dAt = src.Reset7dAt - dst.Window7dSeconds = src.Window7dSeconds - dst.UsagePercent5h = src.UsagePercent5h - dst.UsagePercent5hValid = src.UsagePercent5hValid - dst.Reset5hAt = src.Reset5hAt - dst.UsageUpdatedAt = src.UsageUpdatedAt - dst.UsageUpdatedAt5h = src.UsageUpdatedAt5h + // A WS frame can update the runtime snapshot between two persisted + // observations; the outbox row then still carries the older DB value. Do + // not let that reload roll usage back unless the row is newer, or the + // credential identity changed and usage belongs to a new account. + useRuntime5h := identityChanged || !src.UsageUpdatedAt5h.Before(dst.UsageUpdatedAt5h) + if useRuntime5h { + dst.UsagePercent5h = src.UsagePercent5h + dst.UsagePercent5hValid = src.UsagePercent5hValid + dst.Reset5hAt = src.Reset5hAt + dst.UsageUpdatedAt5h = src.UsageUpdatedAt5h + } + useRuntime7d := identityChanged || !src.UsageUpdatedAt.Before(dst.UsageUpdatedAt) + if useRuntime7d { + dst.UsagePercent7d = src.UsagePercent7d + dst.UsagePercent7dValid = src.UsagePercent7dValid + dst.Reset7dAt = src.Reset7dAt + dst.Window7dSeconds = src.Window7dSeconds + dst.UsageUpdatedAt = src.UsageUpdatedAt + } if src.usageObservedAt.After(dst.usageObservedAt) { dst.usageObservedAt = src.usageObservedAt } diff --git a/auth/scheduler_outbox_consumer_test.go b/auth/scheduler_outbox_consumer_test.go index 882bb0f9d..80342e5cf 100644 --- a/auth/scheduler_outbox_consumer_test.go +++ b/auth/scheduler_outbox_consumer_test.go @@ -228,6 +228,52 @@ func TestApplyPersistentAccountSnapshotPreservesRuntimeState(t *testing.T) { } } +func TestApplyPersistentAccountSnapshotKeepsNewerRuntimeUsage(t *testing.T) { + store := newIndexedRoutingTestStore(nil) + now := time.Now().Add(-time.Minute) + dst := newFastSchedulerTestAccount(1, HealthTierHealthy, 100, 1) + dst.UsagePercent5h = 90 + dst.UsagePercent5hValid = true + dst.Reset5hAt = now.Add(time.Hour) + dst.UsageUpdatedAt5h = now + dst.UsagePercent7d = 80 + dst.UsagePercent7dValid = true + dst.Reset7dAt = now.Add(24 * time.Hour) + dst.UsageUpdatedAt = now + + older := newFastSchedulerTestAccount(1, HealthTierHealthy, 100, 1) + older.CredentialGeneration = dst.CredentialGeneration + older.UsagePercent5h = 10 + older.UsagePercent5hValid = true + older.Reset5hAt = now.Add(-time.Hour) + older.UsageUpdatedAt5h = now.Add(-time.Minute) + older.UsagePercent7d = 20 + older.UsagePercent7dValid = true + older.Reset7dAt = now.Add(-time.Hour) + older.UsageUpdatedAt = now.Add(-time.Minute) + + store.applyPersistentAccountSnapshot(dst, older, true) + if dst.UsagePercent5h != 90 || dst.UsagePercent7d != 80 { + t.Fatalf("older outbox row rolled usage back: 5h=%v 7d=%v", dst.UsagePercent5h, dst.UsagePercent7d) + } + + newer := newFastSchedulerTestAccount(1, HealthTierHealthy, 100, 1) + newer.CredentialGeneration = dst.CredentialGeneration + newer.UsagePercent5h = 95 + newer.UsagePercent5hValid = true + newer.Reset5hAt = now.Add(2 * time.Hour) + newer.UsageUpdatedAt5h = now.Add(time.Minute) + newer.UsagePercent7d = 85 + newer.UsagePercent7dValid = true + newer.Reset7dAt = now.Add(48 * time.Hour) + newer.UsageUpdatedAt = now.Add(time.Minute) + + store.applyPersistentAccountSnapshot(dst, newer, true) + if dst.UsagePercent5h != 95 || dst.UsagePercent7d != 85 { + t.Fatalf("newer outbox row was ignored: 5h=%v 7d=%v", dst.UsagePercent5h, dst.UsagePercent7d) + } +} + func TestReloadDispatchAccountsByIDsAppliesBatchProjection(t *testing.T) { ctx := context.Background() db, err := database.New("sqlite", filepath.Join(t.TempDir(), "scheduler-batch-reload.db")) diff --git a/auth/store.go b/auth/store.go index 88b650d12..4e1d840b0 100644 --- a/auth/store.go +++ b/auth/store.go @@ -163,7 +163,9 @@ type Account struct { UpstreamRequestIDHeader string mu sync.RWMutex usageSyncMu sync.Mutex - modelCatalogMu sync.Mutex + // Snapshot writes serialize per account without blocking other accounts or frame ingestion. + usageSnapshotPersistMu sync.Mutex + modelCatalogMu sync.Mutex // grokRuntimeFactsMu serializes inference-response observations for this // account. The sink performs generation-fenced database writes before it // publishes any hard gate or routing invalidation back to memory. @@ -279,11 +281,13 @@ type Account struct { // 客户端压缩阈值。仅内存保留：值只能在收到响应后才知道，落库也救不了首个请求。 grokContextWindow int64 // grokFreeQuota 是免费额度耗尽 429 解析出的权威用量快照，随 credentials 落库。 - grokFreeQuota *GrokFreeQuotaSnapshot - Status AccountStatus - CooldownUtil time.Time - CooldownReason string // rate_limited / rate_limited_5h / responses_rate_limited / unauthorized / 空 - ErrorMsg string + grokFreeQuota *GrokFreeQuotaSnapshot + Status AccountStatus + CooldownUtil time.Time + CooldownReason string // rate_limited / rate_limited_5h / responses_rate_limited / unauthorized / 空 + usageCooldownPersistMu sync.Mutex + usageCooldownPersistPending bool + ErrorMsg string // 用量进度（从 Codex 响应头被动解析） UsagePercent7d float64 // 7d 窗口使用率 0-100+ @@ -298,6 +302,22 @@ type Account struct { Reset5hAt time.Time // 5h 窗口重置时间 UsageUpdatedAt time.Time // 7d 用量快照刷新时间 UsageUpdatedAt5h time.Time // 5h 用量快照刷新时间 + // usageFramePersisted* 记录 WS 逐帧观测最近一次落库的时间与百分比。 + // 内存快照每帧都实时更新；落库只为重启恢复，大号池下按时间/变化幅度节流， + // 避免把每个上游响应都变成一次数据库写入。 + usageFramePersistedAt time.Time + usageFramePersisted5h float64 + usageFramePersisted7d float64 + usageFramePersistPending bool + usageFramePlanPending bool + usageFrameObservationMu sync.Mutex + usageFramePending []usageFrameObservation + usageFrameObservationRunning bool + usageFrameObservationOverflow bool + // usageFrameLastAt/Kind 记录最近一次被应用的 WS 配额帧，仅供管理端 + // 验证帧摄入是否工作，不参与调度。 + usageFrameLastAt time.Time + usageFrameLastKind string // activated5hResetAt 是已经为其发送过「开窗」最小 /responses 的那个 Reset5hAt。 // 每个观测到的 5h 窗口最多激活一次（issue #581）。 activated5hResetAt time.Time @@ -334,6 +354,7 @@ type Account struct { // 业务标志变化时落库，余额逐请求递减不能每次都写。 creditsPersistedKey string creditsPersistedFlagsKey string + creditsPersistMu sync.Mutex // resetCreditsProbedAt 记录最近一次成功 wham 用量探针的时间。 // 「主动重置次数」只能通过 wham 探针刷新（普通 /responses 流量不携带该字段）， // 因此用它独立判断重置次数是否过期，避免活跃账号因用量快照一直被流量刷新而长期不探针。 @@ -2345,32 +2366,113 @@ func (a *Account) GetUsagePercent7d() (float64, bool) { return a.UsagePercent7d, a.UsagePercent7dValid } +// cooldownMatchesReset reports whether an existing cooldown already ends at +// the observed window reset. Repeated observations of the same exhausted +// window (several in-flight turns, or a WHAM probe while the cooldown is +// active) must not issue another DB write or inflate FailureStreak. +func cooldownMatchesReset(cooldownUntil, resetAt time.Time) bool { + if resetAt.IsZero() { + return true + } + diff := cooldownUntil.Sub(resetAt) + if diff < 0 { + diff = -diff + } + return diff <= time.Minute +} + +// usageCooldownTargetLocked keeps passive quota observations from weakening an +// active cooldown. Preserve every existing reason rather than maintaining an +// incomplete list of authoritative errors. A real Responses 429 also upgrades a +// local usage cooldown so IgnoreUsageLimitStatus cannot bypass that rejection. +// Authentication failures retain the higher priority used by the runtime cache. +// The caller must hold a.mu. +func (a *Account) usageCooldownTargetLocked(now, resetAt time.Time, reason string) (time.Time, string, bool) { + if a.Status == StatusCooldown && now.Before(a.CooldownUtil) { + if reason == ResponsesRateLimitedCooldownReason && a.usageWindowCooldownLocked() { + if resetAt.Before(a.CooldownUtil) { + resetAt = a.CooldownUtil + } + return resetAt, reason, true + } + if !resetAt.After(a.CooldownUtil) || + (a.CooldownReason == reason && cooldownMatchesReset(a.CooldownUtil, resetAt)) { + return a.CooldownUtil, a.CooldownReason, false + } + if reason != ResponsesRateLimitedCooldownReason || a.CooldownReason == "unauthorized" { + reason = a.CooldownReason + } + } + return resetAt, reason, true +} + // MarkUsage7dRateLimited marks an account as rate-limited when its active 7d // usage window is exhausted. A future reset time is preferred; missing reset // metadata falls back to a full 7d cooldown, while stale reset times are ignored. func (s *Store) MarkUsage7dRateLimited(acc *Account) bool { - if s == nil || acc == nil || acc.IsBanned() { + return s.markUsage7dRateLimited(acc, true) +} + +func (s *Store) markUsage7dRateLimited(acc *Account, persist bool) bool { + if s == nil || acc == nil { return false } - if acc.SkipsUsageWindowLimits() { + acc.mu.Lock() + if acc.healthTierLocked() == HealthTierBanned || acc.skipsUsageWindowLimitsLocked() { + acc.mu.Unlock() return false } - pct, ok := acc.GetUsagePercent7d() - if !ok || pct < 100 { + if !acc.UsagePercent7dValid || acc.UsagePercent7d < 100 { + acc.mu.Unlock() return false } - duration := 7 * 24 * time.Hour - if resetAt := acc.GetReset7dAt(); !resetAt.IsZero() { - untilReset := time.Until(resetAt) - if untilReset <= 0 { - return false + now := time.Now() + resetAt := acc.Reset7dAt + if resetAt.IsZero() { + resetAt = now.Add(7 * 24 * time.Hour) + } else if !resetAt.After(now) { + acc.mu.Unlock() + return false + } + + activeCooldown := acc.Status == StatusCooldown && now.Before(acc.CooldownUtil) + // With no reset metadata, an active cooldown already fences this account. + // Do not slide the fallback window forward on every observation. + if activeCooldown && acc.Reset7dAt.IsZero() { + acc.mu.Unlock() + if persist { + s.persistUsageCooldown(acc) } - duration = untilReset + return true + } + until, reason, changed := acc.usageCooldownTargetLocked(now, resetAt, "rate_limited") + if !changed { + acc.mu.Unlock() + if persist { + s.persistUsageCooldown(acc) + } + return true } - s.MarkCooldown(acc, duration, "rate_limited") + if activeCooldown { + acc.CooldownUtil = until + } else { + acc.LastRateLimitedAt = now + acc.LastFailureAt = now + acc.FailureStreak++ + acc.SuccessStreak = 0 + acc.setCooldownUntilLocked(until, reason) + } + acc.recomputeSchedulerLocked(s.maxConcurrency.Load()) + acc.usageCooldownPersistPending = true + acc.mu.Unlock() + + s.fastSchedulerUpdate(acc) + if persist { + s.persistUsageCooldown(acc) + } return true } @@ -2418,6 +2520,33 @@ func (a *Account) ApplyUsageObservation(observedAt time.Time, apply func()) bool return true } +// TryApplyUsageObservation is the non-blocking variant used by connection-level +// WS frame ingress. A false result lets ApplyUsageFrameObservation queue the +// observation instead of blocking the WebSocket read pump. +func (a *Account) TryApplyUsageObservation(observedAt time.Time, apply func()) bool { + if a == nil || apply == nil { + return false + } + if !a.usageSyncMu.TryLock() { + return false + } + defer a.usageSyncMu.Unlock() + if observedAt.IsZero() { + observedAt = time.Now() + } + + a.mu.Lock() + if observedAt.Before(a.usageObservedAt) { + a.mu.Unlock() + return false + } + a.usageObservedAt = observedAt + a.mu.Unlock() + + apply() + return true +} + // GetUsagePercent5h 获取 5h 用量百分比 func (a *Account) GetUsagePercent5h() (float64, bool) { a.mu.RLock() @@ -2425,6 +2554,152 @@ func (a *Account) GetUsagePercent5h() (float64, bool) { return a.UsagePercent5h, a.UsagePercent5hValid } +// UsageFramePersistInterval is the default maximum cadence for persisting WS +// frame observations. In-memory usage is updated on every frame; persistence +// exists only so a restart does not begin from a badly stale snapshot. +const UsageFramePersistInterval = 30 * time.Second + +// usageFramePersistMinDelta forces a write when the reported percentage has +// moved this much since the last persisted observation, so a burst cannot hide +// behind the interval. Threshold crossings and 100% always force a write. +const usageFramePersistMinDelta = 5.0 + +// ShouldPersistUsageFrame decides whether a WS usage-frame observation needs a +// database write. The caller always applies the in-memory snapshot first; this +// only throttles persistence for large pools. The throttle state advances only +// via MarkUsageFramePersisted after a successful write, so a failed write is +// retried on the next frame instead of being suppressed for the interval. +func (a *Account) ShouldPersistUsageFrame(now time.Time, pct5h float64, has5h bool, pct7d float64, has7d bool, interval time.Duration) bool { + if a == nil { + return false + } + a.mu.RLock() + defer a.mu.RUnlock() + + if interval <= 0 { + interval = UsageFramePersistInterval + } + force := false + if has5h { + if pct5h >= 100 && a.usageFramePersisted5h < 100 { + force = true + } + if a.effectiveAutoPause5h > 0 { + threshold := a.effectiveAutoPause5h * 100 + if pct5h >= threshold && a.usageFramePersisted5h < threshold { + force = true + } + } + if math.Abs(pct5h-a.usageFramePersisted5h) >= usageFramePersistMinDelta { + force = true + } + } + if has7d { + if pct7d >= 100 && a.usageFramePersisted7d < 100 { + force = true + } + if a.effectiveAutoPause7d > 0 { + threshold := a.effectiveAutoPause7d * 100 + if pct7d >= threshold && a.usageFramePersisted7d < threshold { + force = true + } + } + if math.Abs(pct7d-a.usageFramePersisted7d) >= usageFramePersistMinDelta { + force = true + } + } + if !force && !a.usageFramePersistedAt.IsZero() && now.Sub(a.usageFramePersistedAt) < interval { + return false + } + return true +} + +// MarkUsageFramePersisted records a successful frame snapshot write. +func (a *Account) MarkUsageFramePersisted(now time.Time, pct5h float64, has5h bool, pct7d float64, has7d bool) { + if a == nil { + return + } + a.mu.Lock() + defer a.mu.Unlock() + a.usageFramePersistedAt = now + if has5h { + a.usageFramePersisted5h = pct5h + } + if has7d { + a.usageFramePersisted7d = pct7d + } +} + +// MarkUsageFrameObserved records the last WS quota frame ingested for the +// account. The admin live projection exposes it so operators can tell +// frame-driven updates apart from WHAM probes; it does not affect scheduling. +func (a *Account) MarkUsageFrameObserved(kind string, at time.Time) { + if a == nil { + return + } + if at.IsZero() { + at = time.Now() + } + a.mu.Lock() + defer a.mu.Unlock() + if at.Before(a.usageFrameLastAt) { + return + } + a.usageFrameLastAt = at + a.usageFrameLastKind = strings.TrimSpace(kind) +} + +// AccountLiveUsageSnapshot is a single-lock projection for the admin live poll. +// It intentionally avoids GetSchedulerDebugSnapshot and other recomputing +// getters, which are too expensive for a one-second cadence over a large page. +type AccountLiveUsageSnapshot struct { + Status string + UsingCredits bool + PlanType string + CreditsValid bool + Credits CreditBalanceSnapshot + UsagePercent5h float64 + UsagePercent5hValid bool + Reset5hAt time.Time + UsageUpdatedAt5h time.Time + UsagePercent7d float64 + UsagePercent7dValid bool + Reset7dAt time.Time + UsageUpdatedAt time.Time + // UsageFrameLastAt/Kind record the last WS quota frame ingested for this + // account, as opposed to snapshots that only advance on WHAM probes. + UsageFrameLastAt time.Time + UsageFrameLastKind string +} + +// GetAccountLiveUsageSnapshot returns the runtime usage/status projection used +// by the account list's live poll. +func (a *Account) GetAccountLiveUsageSnapshot() AccountLiveUsageSnapshot { + if a == nil { + return AccountLiveUsageSnapshot{} + } + a.mu.RLock() + defer a.mu.RUnlock() + now := time.Now() + return AccountLiveUsageSnapshot{ + Status: a.runtimeStatusLocked(now), + UsingCredits: a.usingCreditsLocked(now), + PlanType: a.PlanType, + CreditsValid: a.CreditsValid, + Credits: a.creditSnapshotLocked(), + UsagePercent5h: a.UsagePercent5h, + UsagePercent5hValid: a.UsagePercent5hValid, + Reset5hAt: a.Reset5hAt, + UsageUpdatedAt5h: a.UsageUpdatedAt5h, + UsagePercent7d: a.UsagePercent7d, + UsagePercent7dValid: a.UsagePercent7dValid, + Reset7dAt: a.Reset7dAt, + UsageUpdatedAt: a.UsageUpdatedAt, + UsageFrameLastAt: a.usageFrameLastAt, + UsageFrameLastKind: a.usageFrameLastKind, + } +} + // Mark5hWindowActivated 记录已经为哪个 Reset5hAt 发送过开窗请求。 func (a *Account) Mark5hWindowActivated(resetAt time.Time) { if a == nil || resetAt.IsZero() { @@ -2570,6 +2845,18 @@ func (a *Account) SetRateLimitReachedType(kind string) { a.CreditsRateLimitReachedType = strings.TrimSpace(kind) } +// GetRateLimitReachedType returns the last workspace hard-stop marker. WS +// frames do not carry the per-response field, so the frame path preserves this +// value until WHAM or an explicit header observation clears it. +func (a *Account) GetRateLimitReachedType() string { + if a == nil { + return "" + } + a.mu.RLock() + defer a.mu.RUnlock() + return a.CreditsRateLimitReachedType +} + // ApplySparseCreditsHeaders 应用普通 Codex 响应头携带的 sparse credits 观测信息。 // 余额 / 超额头缺席时保留旧值；rate_limit_reached_type 缺席即「未触达」，必须覆盖。 func (a *Account) ApplySparseCreditsHeaders(hasCredits, unlimited bool, balance *string, overageReached *bool, rateLimitReachedType string) { @@ -2740,6 +3027,8 @@ func (s *Store) persistCreditSnapshot(acc *Account, flagsOnly bool) { if s == nil || s.db == nil || acc == nil { return } + acc.creditsPersistMu.Lock() + defer acc.creditsPersistMu.Unlock() acc.mu.Lock() snap := acc.creditSnapshotLocked() fullKey := creditBalanceKey(snap) @@ -3595,6 +3884,7 @@ type Store struct { armedBoundaryAt atomic.Int64 lazyRefreshInFlight sync.Map stopCh chan struct{} + usageFramePersist *usageFramePersistQueue stopOnce sync.Once backgroundCtx context.Context backgroundCancel context.CancelFunc @@ -4249,6 +4539,7 @@ func NewStore(db *database.DB, tc cache.TokenCache, settings *database.SystemSet backgroundRefreshWakeCh: make(chan struct{}, 1), boundaryProbeWakeCh: make(chan struct{}, 1), stopCh: make(chan struct{}), + usageFramePersist: newUsageFramePersistQueue(), backgroundCtx: backgroundCtx, backgroundCancel: backgroundCancel, schedulerMetrics: newSchedulerRuntimeMetrics(), @@ -4260,6 +4551,7 @@ func NewStore(db *database.DB, tc cache.TokenCache, settings *database.SystemSet oauthRefreshLocks: make(map[string]*oauthRefreshLocalLock), } s.maxConcurrency.Store(int64(settings.MaxConcurrency)) + s.usageFramePersist.start(s) s.testConcurrency.Store(int64(settings.TestConcurrency)) s.codexRequestCompression.Store(settings.CodexRequestCompression) s.availability.Store(newAvailabilityHub()) @@ -4465,6 +4757,14 @@ func (s *Store) fastSchedulerUpdate(acc *Account) { s.notifySchedulerAccountAvailability(acc, true) } +// RefreshUsageFrameScheduler applies an in-memory usage snapshot to the fast +// scheduler and wakes queued dispatchers without touching the database. The +// connection-level WS frame path updates the snapshot on the read pump, so the +// scheduler is refreshed separately from the persistence-throttled path. +func (s *Store) RefreshUsageFrameScheduler(acc *Account) { + s.fastSchedulerUpdate(acc) +} + func (s *Store) fastSchedulerRemove(dbID int64) { if s == nil || dbID == 0 { return @@ -10066,16 +10366,7 @@ func (s *Store) markCooldownUntil(acc *Account, until time.Time, reason string, if updateScheduler { s.fastSchedulerUpdate(acc) } - s.setCachedAccountCooldown(acc.DBID, reason, until) - - if s.db == nil { - return - } - ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second) - defer cancel() - if err := s.db.SetCooldown(ctx, acc.DBID, reason, until); err != nil { - log.Printf("[账号 %d] 持久化冷却状态失败: %v", acc.DBID, err) - } + s.persistAccountCooldownState(acc, true, false) } // markCooldown 根据 exactDuration 选择自适应或精确时长并应用账号冷却。 @@ -10128,23 +10419,7 @@ func (s *Store) markCooldown(acc *Account, duration time.Duration, reason string acc.setCooldownUntilLocked(until, reason) acc.mu.Unlock() s.fastSchedulerUpdate(acc) - s.setCachedAccountCooldown(acc.DBID, reason, until) - - if s.db == nil { - return - } - - ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second) - defer cancel() - var err error - if errorMsg != "" { - err = s.db.SetCooldownWithError(ctx, acc.DBID, reason, until, errorMsg) - } else { - err = s.db.SetCooldown(ctx, acc.DBID, reason, until) - } - if err != nil { - log.Printf("[账号 %d] 持久化冷却状态失败: %v", acc.DBID, err) - } + s.persistAccountCooldownState(acc, true, false) } func (s *Store) MarkModelCooldown(acc *Account, model string, duration time.Duration, reason string) ModelCooldown { @@ -10548,17 +10823,7 @@ func (s *Store) MarkError(acc *Account, errorMsg string) { acc.recomputeSchedulerLocked(s.maxConcurrency.Load()) acc.mu.Unlock() s.fastSchedulerUpdate(acc) - s.deleteCachedAccountCooldown(acc.DBID) - - if s.db == nil { - return - } - - ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second) - defer cancel() - if err := s.db.SetError(ctx, acc.DBID, errorMsg); err != nil { - log.Printf("[账号 %d] 持久化错误状态失败: %v", acc.DBID, err) - } + s.persistAccountCooldownState(acc, true, false) } // ClearCooldown 清除账号冷却状态，并同步清理数据库 @@ -10594,17 +10859,7 @@ func (s *Store) ClearCooldown(acc *Account) { acc.recomputeSchedulerLocked(s.maxConcurrency.Load()) acc.mu.Unlock() s.fastSchedulerUpdate(acc) - s.deleteCachedAccountCooldown(acc.DBID) - - if s.db == nil { - return - } - - ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second) - defer cancel() - if err := s.db.ClearError(ctx, acc.DBID); err != nil { - log.Printf("[账号 %d] 清理账号状态失败: %v", acc.DBID, err) - } + s.persistAccountCooldownState(acc, true, true) } // ReleaseUsageWindowCooldownForCredits 在积分门打开后释放由本地用量窗口判罚产生的 cooldown。 @@ -10649,8 +10904,6 @@ func (s *Store) ClearUsageLimitCooldownSince(acc *Account, observedAt time.Time) acc.mu.Unlock() return false } - reason := acc.CooldownReason - until := acc.CooldownUtil acc.Status = StatusReady acc.CooldownUtil = time.Time{} acc.CooldownReason = "" @@ -10663,24 +10916,7 @@ func (s *Store) ClearUsageLimitCooldownSince(acc *Account, observedAt time.Time) acc.mu.Unlock() s.fastSchedulerUpdate(acc) - s.deleteCachedAccountCooldown(acc.DBID) - acc.mu.RLock() - status := acc.Status - currentReason := acc.CooldownReason - currentUntil := acc.CooldownUtil - acc.mu.RUnlock() - if status == StatusCooldown && currentReason != "" { - s.setCachedAccountCooldown(acc.DBID, currentReason, currentUntil) - } - if s.db == nil { - return true - } - - ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second) - defer cancel() - if _, err := s.db.ClearCooldownIfReasonAndUntil(ctx, acc.DBID, reason, until); err != nil { - log.Printf("[账号 %d] 清理过期用量冷却状态失败: %v", acc.DBID, err) - } + s.persistAccountCooldownState(acc, true, false) return true } @@ -10732,14 +10968,7 @@ func (s *Store) confirmResponsesAvailable(acc *Account, requestStartedAt time.Ti acc.mu.Unlock() s.fastSchedulerUpdate(acc) - s.deleteCachedAccountCooldown(acc.DBID) - if s.db != nil { - ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second) - defer cancel() - if err := s.db.ClearCooldown(ctx, acc.DBID); err != nil { - log.Printf("[账号 %d] Responses 成功后清理用量限流冷却失败: %v", acc.DBID, err) - } - } + s.persistAccountCooldownState(acc, true, false) return true } @@ -10784,17 +11013,7 @@ func (s *Store) RecordManualTestSuccess(acc *Account, latency time.Duration) { acc.recomputeSchedulerLocked(s.maxConcurrency.Load()) acc.mu.Unlock() s.fastSchedulerUpdate(acc) - s.deleteCachedAccountCooldown(acc.DBID) - - if s.db == nil { - return - } - - ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second) - defer cancel() - if err := s.db.ClearError(ctx, acc.DBID); err != nil { - log.Printf("[账号 %d] 清理账号测试成功状态失败: %v", acc.DBID, err) - } + s.persistAccountCooldownState(acc, true, true) } // ReportRequestSuccess 记录一次成功请求，用于动态调度评分 @@ -10890,16 +11109,25 @@ func (s *Store) ReportRequestFailure(acc *Account, kind string, latency time.Dur // PersistUsageSnapshot 持久化账号用量快照（7d + 5h） func (s *Store) PersistUsageSnapshot(acc *Account, pct7d float64) { - if acc == nil { - return + s.PublishUsageSnapshot(acc, pct7d, time.Now(), true) +} + +// PublishUsageSnapshot updates the in-memory usage snapshot and the scheduler. +// persist=false skips the database write, used by the WS frame path where the +// in-memory value must change immediately but persistence is throttled. +func (s *Store) PublishUsageSnapshot(acc *Account, pct7d float64, observedAt time.Time, persist bool) bool { + if acc == nil || s == nil { + return false + } + if observedAt.IsZero() { + observedAt = time.Now() } - now := time.Now() - acc.SetUsageSnapshot(pct7d, now) + acc.SetUsageSnapshot(pct7d, observedAt) s.fastSchedulerUpdate(acc) - if s.db == nil { - return + if !persist || s.db == nil { + return false } // 如果有 5h 数据，使用完整存储 @@ -10908,17 +11136,34 @@ func (s *Store) PersistUsageSnapshot(acc *Account, pct7d float64) { reset7dAt := acc.GetReset7dAt() ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second) defer cancel() - if err := s.db.UpdateUsageSnapshotFull(ctx, acc.DBID, pct7d, reset7dAt, pct5h, reset5hAt, now, acc.GetUsageUpdatedAt5h()); err != nil { + err := s.withUsageSnapshotWrite(acc, func() error { + return s.db.UpdateUsageSnapshotFull(ctx, acc.DBID, pct7d, reset7dAt, pct5h, reset5hAt, observedAt, acc.GetUsageUpdatedAt5h()) + }) + if err != nil { log.Printf("[账号 %d] 持久化用量快照失败: %v", acc.DBID, err) + return false } - return + return true } ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second) defer cancel() - if err := s.db.UpdateUsageSnapshot(ctx, acc.DBID, pct7d, now); err != nil { + err := s.withUsageSnapshotWrite(acc, func() error { + return s.db.UpdateUsageSnapshot(ctx, acc.DBID, pct7d, observedAt) + }) + if err != nil { log.Printf("[账号 %d] 持久化用量快照失败: %v", acc.DBID, err) + return false } + return true +} + +// withUsageSnapshotWrite serializes credential snapshot writes with the async +// frame-persist worker. +func (s *Store) withUsageSnapshotWrite(acc *Account, fn func() error) error { + acc.usageSnapshotPersistMu.Lock() + defer acc.usageSnapshotPersistMu.Unlock() + return fn() } // UpdateAccountSubscriptionExpiresAt persists the latest subscription expiration observed from upstream. @@ -11026,6 +11271,15 @@ func (s *Store) ClearStaleSubscriptionExpiresAt(acc *Account) bool { // UpdateAccountPlanType persists the latest Codex plan type observed from upstream headers. func (s *Store) UpdateAccountPlanType(acc *Account, planType string) bool { + return s.updateAccountPlanType(acc, planType, true) +} + +// ObserveUsageFramePlan updates runtime routing without writing on a WS reader. +func (s *Store) ObserveUsageFramePlan(acc *Account, planType string) bool { + return s.updateAccountPlanType(acc, planType, false) +} + +func (s *Store) updateAccountPlanType(acc *Account, planType string, persist bool) bool { if s == nil || acc == nil { return false } @@ -11038,6 +11292,7 @@ func (s *Store) UpdateAccountPlanType(acc *Account, planType string) bool { changed := acc.PlanType != plan if changed { acc.PlanType = plan + acc.usageFramePlanPending = true acc.recomputeSchedulerLocked(s.maxConcurrency.Load()) } acc.mu.Unlock() @@ -11046,13 +11301,11 @@ func (s *Store) UpdateAccountPlanType(acc *Account, planType string) bool { s.fastSchedulerUpdate(acc) } - if s.db == nil || !changed { + if s.db == nil || !persist { return changed } - ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second) - defer cancel() - if err := s.db.UpdateCredentials(ctx, acc.DBID, map[string]interface{}{"plan_type": plan}); err != nil { + if err := s.withUsageSnapshotWrite(acc, func() error { return s.persistUsageFramePlan(acc) }); err != nil { log.Printf("[账号 %d] 持久化 plan_type 失败: %v", acc.DBID, err) } return changed @@ -11706,12 +11959,7 @@ func (s *Store) parallelRecoveryProbe(ctx context.Context) { } account.mu.Unlock() // 清理数据库冷却状态 - s.deleteCachedAccountCooldown(account.DBID) - if s.db != nil { - clearCtx, clearCancel := context.WithTimeout(ctx, 2*time.Second) - _ = s.db.ClearCooldown(clearCtx, account.DBID) - clearCancel() - } + s.persistAccountCooldownState(account, true, false) } } }(acc) diff --git a/auth/usage_cooldown_persist.go b/auth/usage_cooldown_persist.go new file mode 100644 index 000000000..102987d1a --- /dev/null +++ b/auth/usage_cooldown_persist.go @@ -0,0 +1,68 @@ +package auth + +import ( + "context" + "log" + "time" +) + +// persistUsageCooldown retries a failed usage-owned cooldown write without +// treating the retry as a new failure or changing the runtime deadline. +func (s *Store) persistUsageCooldown(acc *Account) { + s.persistAccountCooldownState(acc, false, false) +} + +// All native account cooldown writers/clearers publish through this critical +// section. Read current state only after acquiring the publication lock: a +// delayed frame, 401, or clear must never publish a captured older state. +// Never hold this lock while acquiring acc.usageSnapshotPersistMu. +func (s *Store) persistAccountCooldownState(acc *Account, force, clearError bool) { + if s == nil || acc == nil { + return + } + acc.usageCooldownPersistMu.Lock() + defer acc.usageCooldownPersistMu.Unlock() + acc.mu.RLock() + pending := acc.usageCooldownPersistPending + status, errorMsg := acc.Status, acc.ErrorMsg + reason, until := acc.CooldownReason, acc.CooldownUtil + acc.mu.RUnlock() + if !force && (!pending || status != StatusCooldown) { + return + } + if status == StatusCooldown { + s.setCachedAccountCooldown(acc.DBID, reason, until) + } else { + s.deleteCachedAccountCooldown(acc.DBID) + } + ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second) + defer cancel() + var err error + if s.db != nil { + switch status { + case StatusCooldown: + if errorMsg != "" { + err = s.db.SetCooldownWithError(ctx, acc.DBID, reason, until, errorMsg) + } else { + err = s.db.SetCooldown(ctx, acc.DBID, reason, until) + } + case StatusError: + err = s.db.SetError(ctx, acc.DBID, errorMsg) + default: + if clearError { + err = s.db.ClearError(ctx, acc.DBID) + } else { + err = s.db.ClearCooldown(ctx, acc.DBID) + } + } + } + if err != nil { + log.Printf("[账号 %d] 持久化用量冷却失败: %v", acc.DBID, err) + return + } + acc.mu.Lock() + if acc.Status == status && acc.CooldownReason == reason && acc.CooldownUtil.Equal(until) { + acc.usageCooldownPersistPending = false + } + acc.mu.Unlock() +} diff --git a/auth/usage_cooldown_test.go b/auth/usage_cooldown_test.go new file mode 100644 index 000000000..e51fb8f6e --- /dev/null +++ b/auth/usage_cooldown_test.go @@ -0,0 +1,323 @@ +package auth + +import ( + "context" + "database/sql" + "path/filepath" + "sync" + "testing" + "time" + + "github.com/codex2api/cache" + "github.com/codex2api/database" +) + +func TestUsageCooldownNeverShortensActiveCooldown(t *testing.T) { + ctx := context.Background() + path := filepath.Join(t.TempDir(), "usage-cooldown.db") + db, err := database.New("sqlite", path) + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { _ = db.Close() }) + conn, err := sql.Open("sqlite", path) + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { _ = conn.Close() }) + // Count actual cooldown writes, independently of snapshot writes. + for _, query := range []string{ + `CREATE TABLE cooldown_writes (account_id INTEGER)`, + `CREATE TRIGGER count_cooldown_writes AFTER UPDATE OF cooldown_until ON accounts BEGIN INSERT INTO cooldown_writes VALUES (NEW.id); END`, + } { + if _, err := conn.ExecContext(ctx, query); err != nil { + t.Fatal(err) + } + } + tokenCache := cache.NewMemory(4) + t.Cleanup(func() { _ = tokenCache.Close() }) + store := NewStore(db, tokenCache, &database.SystemSettings{MaxConcurrency: 2}) + t.Cleanup(store.Stop) + + for _, source := range []string{"5h", "7d", "responses"} { + t.Run(source, func(t *testing.T) { + defaultReason := premium5hCooldownReason + if source == "7d" { + defaultReason = "rate_limited" + } else if source == "responses" { + defaultReason = ResponsesRateLimitedCooldownReason + } + for _, tc := range []struct { + name string + prior time.Duration + reset time.Duration + reason string + wantChange bool + }{ + {"shorter authoritative", 6 * time.Hour, time.Hour, ResponsesRateLimitedCooldownReason, false}, + {"equal authoritative", time.Hour, time.Hour, ResponsesRateLimitedCooldownReason, false}, + {"extend authoritative", time.Hour, 2 * time.Hour, ResponsesRateLimitedCooldownReason, true}, + {"extend other reason", time.Hour, 2 * time.Hour, "upstream_error", true}, + {"shorter unauthorized", 6 * time.Hour, time.Hour, "unauthorized", false}, + {"extend unauthorized", time.Hour, 2 * time.Hour, "unauthorized", true}, + {"shorter usage", 6 * time.Hour, time.Hour, premium5hCooldownReason, false}, + {"equal usage", time.Hour, time.Hour, premium5hCooldownReason, false}, + {"shorter 7d usage", 6 * time.Hour, time.Hour, "rate_limited_7d", false}, + {"equal 7d usage", time.Hour, time.Hour, "rate_limited_7d", false}, + {"extend usage", time.Hour, 2 * time.Hour, premium5hCooldownReason, true}, + {"new cooldown", 0, 2 * time.Hour, "", true}, + {"expired cooldown", -time.Hour, 2 * time.Hour, ResponsesRateLimitedCooldownReason, true}, + } { + t.Run(tc.name, func(t *testing.T) { + now := time.Now().UTC().Truncate(time.Second) + resetAt := now.Add(tc.reset) + id, err := db.InsertAccountWithCredentials(ctx, "usage-cooldown", map[string]interface{}{"access_token": "test-token", "plan_type": "plus"}, "") + if err != nil { + t.Fatal(err) + } + acc := newPremium5hTestAccount("plus", resetAt) + acc.DBID = id + acc.FailureStreak = 3 + acc.UsagePercent5h = 40 + acc.UsagePercent7d, acc.UsagePercent7dValid, acc.Reset7dAt = 100, true, resetAt + wantUntil, wantReason := resetAt, defaultReason + wantWrites := 1 + if tc.prior != 0 { + acc.Status, acc.CooldownUtil, acc.CooldownReason = StatusCooldown, now.Add(tc.prior), tc.reason + if source == "responses" && tc.reason == "rate_limited_7d" { + acc.Reset7dAt = acc.CooldownUtil + } + if err := db.SetCooldown(ctx, id, tc.reason, acc.CooldownUtil); err != nil { + t.Fatal(err) + } + wantWrites++ + if tc.prior > 0 { + store.setCachedAccountCooldown(id, tc.reason, acc.CooldownUtil) + if !tc.wantChange { + wantUntil, wantReason = acc.CooldownUtil, tc.reason + if source == "responses" && (tc.reason == premium5hCooldownReason || tc.reason == "rate_limited_7d") { + wantReason = ResponsesRateLimitedCooldownReason + } else { + wantWrites-- + } + } else if source != "responses" || tc.reason == "unauthorized" { + wantReason = tc.reason + } + } + } + switch source { + case "5h": + store.MarkPremium5hRateLimited(acc, resetAt) + case "7d": + if !store.MarkUsage7dRateLimited(acc) { + t.Fatal("exhausted 7d window was not recognized") + } + case "responses": + store.MarkResponsesPremium5hRateLimited(acc, resetAt) + } + if acc.Status != StatusCooldown || !acc.CooldownUtil.Equal(wantUntil) || acc.CooldownReason != wantReason { + t.Errorf("runtime cooldown = (%v, %v, %q), want cooldown/%v/%q", acc.Status, acc.CooldownUtil, acc.CooldownReason, wantUntil, wantReason) + } + wantStreak := 3 + if source == "7d" && tc.prior <= 0 { + wantStreak++ // Preserve the existing new-cooldown behavior. + } + if acc.FailureStreak != wantStreak { + t.Errorf("FailureStreak = %d, want %d", acc.FailureStreak, wantStreak) + } + row, err := db.GetAccountByID(ctx, id) + if err != nil { + t.Fatal(err) + } + if !row.CooldownUntil.Valid || !row.CooldownUntil.Time.Equal(wantUntil) || row.CooldownReason != wantReason { + t.Errorf("persisted cooldown = (%v, %q), want %v/%q", row.CooldownUntil, row.CooldownReason, wantUntil, wantReason) + } + record, ok := store.getCachedAccountCooldown(id) + if !ok || !record.ResetAt.Equal(wantUntil) || record.Reason != wantReason { + t.Errorf("cached cooldown = %+v, present=%v, want %v/%q", record, ok, wantUntil, wantReason) + } + var writes int + if err := conn.QueryRowContext(ctx, `SELECT COUNT(*) FROM cooldown_writes WHERE account_id = ?`, id).Scan(&writes); err != nil { + t.Fatal(err) + } + if writes != wantWrites { + t.Errorf("cooldown writes = %d, want %d", writes, wantWrites) + } + if source != "7d" { + if pct, reset, ok := acc.GetUsageSnapshot5h(); !ok || pct != 100 || !reset.Equal(resetAt) { + t.Errorf("5h snapshot = (%v, %v, %v), want 100/%v/true", pct, reset, ok, resetAt) + } + if row.GetCredential("codex_5h_used_percent") != "100" { + t.Error("5h snapshot was not persisted independently of the cooldown") + } + } + // The same observation must not keep writing the cooldown. + switch source { + case "5h": + store.MarkPremium5hRateLimited(acc, resetAt) + case "7d": + store.MarkUsage7dRateLimited(acc) + case "responses": + store.MarkResponsesPremium5hRateLimited(acc, resetAt) + } + if err := conn.QueryRowContext(ctx, `SELECT COUNT(*) FROM cooldown_writes WHERE account_id = ?`, id).Scan(&writes); err != nil { + t.Fatal(err) + } + if writes != wantWrites || acc.FailureStreak != wantStreak { + t.Errorf("repeated observation: writes=%d streak=%d, want %d/%d", writes, acc.FailureStreak, wantWrites, wantStreak) + } + }) + } + }) + } +} + +func TestUsageCooldownRetriesFailedPersistence(t *testing.T) { + ctx := context.Background() + path := filepath.Join(t.TempDir(), "cooldown-retry.db") + db, err := database.New("sqlite", path) + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { _ = db.Close() }) + conn, err := sql.Open("sqlite", path) + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { _ = conn.Close() }) + store := NewStore(db, nil, &database.SystemSettings{MaxConcurrency: 1}) + t.Cleanup(store.Stop) + for _, source := range []string{"5h", "7d", "responses"} { + t.Run(source, func(t *testing.T) { + id, err := db.InsertAccountWithCredentials(ctx, "cooldown-retry", map[string]interface{}{"plan_type": "plus"}, "") + if err != nil { + t.Fatal(err) + } + reset := time.Now().Add(time.Hour).UTC().Truncate(time.Second) + acc := newPremium5hTestAccount("plus", reset) + acc.DBID = id + acc.UsagePercent7d, acc.UsagePercent7dValid, acc.Reset7dAt = 100, true, reset + mark := func() { + switch source { + case "5h": + store.MarkPremium5hRateLimited(acc, reset) + case "7d": + store.MarkUsage7dRateLimited(acc) + case "responses": + store.MarkResponsesPremium5hRateLimited(acc, reset) + } + } + if _, err := conn.Exec(`CREATE TRIGGER fail_cooldown BEFORE UPDATE OF cooldown_until ON accounts BEGIN SELECT RAISE(FAIL, 'synthetic transient failure'); END`); err != nil { + t.Fatal(err) + } + mark() + streak := acc.FailureStreak + row, err := db.GetAccountByID(ctx, id) + if err != nil { + t.Fatal(err) + } + if row.CooldownUntil.Valid { + t.Fatal("fault injection did not reject the cooldown write") + } + if _, err := conn.Exec(`DROP TRIGGER fail_cooldown`); err != nil { + t.Fatal(err) + } + mark() + row, err = db.GetAccountByID(ctx, id) + if err != nil { + t.Fatal(err) + } + if !row.CooldownUntil.Valid || !row.CooldownUntil.Time.Equal(reset) || row.CooldownReason != acc.GetCooldownReason() { + t.Fatalf("cooldown was not retried after recovery: %+v", row.CooldownUntil) + } + if acc.FailureStreak != streak { + t.Fatal("persistence retry counted another failure") + } + }) + } +} + +func TestUsageCooldownDelayedPublisherRespectsCurrentState(t *testing.T) { + for _, action := range []string{"unauthorized", "clear", "error"} { + t.Run(action, func(t *testing.T) { + store, db, acc := newUsageFramePersistenceFixture(t) + reset := time.Now().Add(time.Hour).UTC().Truncate(time.Second) + acc.SetUsageSnapshot5h(100, reset) + store.MarkUsageFrameCooldowns(acc, time.Now()) + // Pause the frame publisher before it can read/publish its state. + acc.usageCooldownPersistMu.Lock() + var release sync.Once + unlock := func() { release.Do(acc.usageCooldownPersistMu.Unlock) } + t.Cleanup(unlock) + oldDone := make(chan struct{}) + go func() { store.persistUsageCooldown(acc); close(oldDone) }() + newDone := make(chan struct{}) + go func() { + defer close(newDone) + switch action { + case "unauthorized": + store.MarkCooldownWithErrorExactDuration(acc, 6*time.Hour, "unauthorized", "synthetic auth error") + case "clear": + store.ClearCooldown(acc) + case "error": + store.MarkError(acc, "synthetic permanent error") + } + }() + deadline := time.Now().Add(3 * time.Second) + changed := false + for time.Now().Before(deadline) { + acc.mu.RLock() + changed = acc.CooldownReason != premium5hCooldownReason + acc.mu.RUnlock() + if changed { + break + } + time.Sleep(time.Millisecond) + } + unlock() + for _, done := range []chan struct{}{oldDone, newDone} { + select { + case <-done: + case <-time.After(3 * time.Second): + t.Fatal("cooldown publication deadlocked") + } + } + if !changed { + t.Fatal("new runtime state did not precede delayed publication") + } + row, err := db.GetAccountByID(context.Background(), acc.DBID) + if err != nil { + t.Fatal(err) + } + record, cached := store.getCachedAccountCooldown(acc.DBID) + switch action { + case "unauthorized": + if row.CooldownReason != "unauthorized" || !row.CooldownUntil.Valid || !row.CooldownUntil.Time.Equal(acc.CooldownUtil) || !cached || record.Reason != "unauthorized" || !record.ResetAt.Equal(row.CooldownUntil.Time) { + t.Fatalf("delayed usage publisher overwrote newer authentication cooldown: row=%v/%v runtime=%v cache=%+v present=%v", row.CooldownReason, row.CooldownUntil, acc.CooldownUtil, record, cached) + } + case "clear", "error": + if row.CooldownUntil.Valid || row.CooldownReason != "" || cached { + t.Fatal("delayed usage publisher resurrected cleared cooldown") + } + if action == "error" && row.Status != "error" { + t.Fatal("delayed usage publisher cleared newer error") + } + } + }) + } +} + +func TestUsageCooldownRepeatedResponses429FencesInflightSuccess(t *testing.T) { + store := NewStore(nil, nil, nil) + t.Cleanup(store.Stop) + reset := time.Now().Add(time.Hour) + acc := newPremium5hTestAccount("plus", reset) + store.SetIgnoreUsageLimitStatus(true) + store.AddAccount(acc) + store.MarkResponsesPremium5hRateLimited(acc, reset) + started := acc.LastRateLimitedAt.Add(time.Second) + store.markPremium5hRateLimited(acc, reset, started.Add(time.Second), ResponsesRateLimitedCooldownReason, true) + if store.ConfirmResponsesAvailableSince(acc, started) { + t.Fatal("an in-flight success cleared newer repeated Responses 429 evidence") + } +} diff --git a/auth/usage_frame_observation.go b/auth/usage_frame_observation.go new file mode 100644 index 000000000..9e95c3ca4 --- /dev/null +++ b/auth/usage_frame_observation.go @@ -0,0 +1,103 @@ +package auth + +import "time" + +type usageFrameObservation struct { + at time.Time + apply func() +} + +const maxPendingUsageFrames = 64 + +// ApplyUsageFrameObservation keeps the sole WS reader non-blocking even when a +// WHAM observation holds usageSyncMu during a write. A bounded FIFO preserves +// partial-frame semantics and original timestamps; callbacks must not be merged +// as whole frames. On overflow retain the newest bounded tail and request the +// existing authoritative probe. Queued callbacks never perform database I/O. +func (s *Store) ApplyUsageFrameObservation(acc *Account, at time.Time, apply func()) { + if acc == nil || apply == nil { + return + } + if at.IsZero() { + at = time.Now() + } + guarded := func() { + if s != nil { + select { + case <-s.stopCh: + return + default: + } + } + apply() + } + acc.usageFrameObservationMu.Lock() + // Do not let a newly arriving partial frame jump past a queued frame and + // make that frame appear stale before its worker acquires usageSyncMu. + if !acc.usageFrameObservationRunning && len(acc.usageFramePending) == 0 && acc.TryApplyUsageObservation(at, guarded) { + acc.usageFrameObservationMu.Unlock() + return + } + acc.mu.RLock() + stale := at.Before(acc.usageObservedAt) + acc.mu.RUnlock() + if stale { + acc.usageFrameObservationMu.Unlock() + return + } + overflow := len(acc.usageFramePending) == maxPendingUsageFrames + wakeProbe := overflow && !acc.usageFrameObservationOverflow + if overflow { + copy(acc.usageFramePending, acc.usageFramePending[1:]) + acc.usageFramePending = acc.usageFramePending[:maxPendingUsageFrames-1] + acc.usageFrameObservationOverflow = true + } + acc.usageFramePending = append(acc.usageFramePending, usageFrameObservation{at: at, apply: guarded}) + if acc.usageFrameObservationRunning { + acc.usageFrameObservationMu.Unlock() + if wakeProbe && s != nil { + s.TriggerUsageProbeForAccountAsync(acc) + } + return + } + acc.usageFrameObservationRunning = true + acc.usageFrameObservationMu.Unlock() + go func() { + for { + acc.usageFrameObservationMu.Lock() + if len(acc.usageFramePending) == 0 { + acc.usageFrameObservationRunning = false + acc.usageFrameObservationOverflow = false + acc.usageFramePending = nil + acc.usageFrameObservationMu.Unlock() + return + } + pending := acc.usageFramePending[0] + acc.usageFramePending[0] = usageFrameObservation{} + acc.usageFramePending = acc.usageFramePending[1:] + acc.usageFrameObservationMu.Unlock() + acc.ApplyUsageObservation(pending.at, pending.apply) + } + }() +} + +// MarkUsageFrameCooldowns is called inside the frame observation ordering gate. +// Only runtime state changes here; the existing worker owns DB/cache writes. +func (s *Store) MarkUsageFrameCooldowns(acc *Account, at time.Time) (bool, bool) { + if s == nil || acc == nil { + return false, false + } + acc.mu.RLock() + isolated := acc.Status == StatusError || acc.healthTierLocked() == HealthTierBanned + acc.mu.RUnlock() + if isolated { + return false, false + } + limited7d := s.markUsage7dRateLimited(acc, false) + pct, reset, valid := acc.GetUsageSnapshot5h() + limited5h := valid && pct >= 100 && acc.IsPremium5hPlan() && !acc.SkipsUsageWindowLimits() + if limited5h { + s.markPremium5hRateLimited(acc, reset, at, premium5hCooldownReason, false) + } + return limited7d, limited5h +} diff --git a/auth/usage_frame_observation_test.go b/auth/usage_frame_observation_test.go new file mode 100644 index 000000000..6819bf451 --- /dev/null +++ b/auth/usage_frame_observation_test.go @@ -0,0 +1,129 @@ +package auth + +import ( + "context" + "reflect" + "sync" + "testing" + "time" +) + +func waitUsageFrameDrain(t *testing.T, acc *Account) { + t.Helper() + deadline := time.Now().Add(3 * time.Second) + for time.Now().Before(deadline) { + acc.usageFrameObservationMu.Lock() + running := acc.usageFrameObservationRunning + acc.usageFrameObservationMu.Unlock() + if !running { + return + } + time.Sleep(time.Millisecond) + } + t.Fatal("frame observation worker did not drain") +} + +func TestUsageFrameObservationPreservesPartialFIFO(t *testing.T) { + acc := &Account{} + var store *Store + at := time.Now() + var applied []int + acc.usageSyncMu.Lock() + var release sync.Once + unlock := func() { release.Do(acc.usageSyncMu.Unlock) } + t.Cleanup(unlock) + for i := 0; i < 3; i++ { + store.ApplyUsageFrameObservation(acc, at.Add(time.Duration(i)*time.Second), func() { + applied = append(applied, i) + switch i { + case 0: + acc.SetUsageSnapshot5hAt(42, at.Add(time.Hour), at) + case 1: + acc.SetUsageSnapshot(65, at.Add(time.Second)) + case 2: + acc.ApplySparseCreditsHeaders(true, false, nil, nil, "") + } + }) + } + unlock() + // A frame arriving as the worker acquires the observation gate must not + // bypass the backlog and make older partial windows stale. + store.ApplyUsageFrameObservation(acc, at.Add(3*time.Second), func() { applied = append(applied, 3) }) + waitUsageFrameDrain(t, acc) + if !reflect.DeepEqual(applied, []int{0, 1, 2, 3}) { + t.Fatalf("partial observations applied out of order: %v", applied) + } + if pct, _, valid := acc.GetUsageSnapshot5h(); !valid || pct != 42 { + t.Fatalf("5h partial lost: %v/%v", pct, valid) + } + if pct, valid := acc.GetUsagePercent7d(); !valid || pct != 65 { + t.Fatalf("7d partial lost: %v/%v", pct, valid) + } + if !acc.CreditsValid || !acc.CreditsHasCredits || !acc.UsageUpdatedAt5h.Equal(at) || !acc.usageObservedAt.Equal(at.Add(3*time.Second)) { + t.Fatal("credits or original observation timestamp lost") + } +} + +func TestUsageFrameObservationOverflowRequestsProbe(t *testing.T) { + store := NewStore(nil, nil, nil) + t.Cleanup(store.Stop) + acc := &Account{} + probed := make(chan struct{}, 2) + probeRelease := make(chan struct{}) + defer close(probeRelease) + store.SetUsageProbeFunc(func(ctx context.Context, _ *Account) error { + probed <- struct{}{} + select { + case <-probeRelease: + case <-ctx.Done(): + } + return nil + }) + acc.usageSyncMu.Lock() + var release sync.Once + unlock := func() { release.Do(acc.usageSyncMu.Unlock) } + t.Cleanup(unlock) + last := -1 + at := time.Now() + for i := 0; i < 2*maxPendingUsageFrames; i++ { + store.ApplyUsageFrameObservation(acc, at.Add(time.Duration(i)*time.Millisecond), func() { last = i }) + } + acc.usageFrameObservationMu.Lock() + pending := len(acc.usageFramePending) + acc.usageFrameObservationMu.Unlock() + if pending > maxPendingUsageFrames { + t.Fatalf("unbounded frame queue: %d", pending) + } + select { + case <-probed: + case <-time.After(3 * time.Second): + t.Fatal("overflow did not trigger authoritative probe") + } + unlock() + waitUsageFrameDrain(t, acc) + if last != 2*maxPendingUsageFrames-1 { + t.Fatalf("newest tail lost: %d", last) + } + select { + case <-probed: + t.Fatal("same overflow burst triggered duplicate probe") + default: + } +} + +func TestUsageFrameObservationDoesNotApplyAfterStop(t *testing.T) { + store := NewStore(nil, nil, nil) + acc := &Account{} + acc.usageSyncMu.Lock() + var release sync.Once + unlock := func() { release.Do(acc.usageSyncMu.Unlock) } + t.Cleanup(unlock) + applied := false + store.ApplyUsageFrameObservation(acc, time.Now(), func() { applied = true }) + store.Stop() + unlock() + waitUsageFrameDrain(t, acc) + if applied { + t.Fatal("queued observation mutated runtime after Stop") + } +} diff --git a/auth/usage_frame_persist.go b/auth/usage_frame_persist.go new file mode 100644 index 000000000..1f69bd8ce --- /dev/null +++ b/auth/usage_frame_persist.go @@ -0,0 +1,241 @@ +package auth + +import ( + "context" + "log" + "sync" + "time" +) + +// usageFramePersistJob is a snapshot of the fields needed to persist one +// frame-driven usage observation. Values are captured at enqueue time so the +// write is coherent even if newer frames arrive while it is queued. +type usageFramePersistJob struct { + account *Account + dbID int64 + observedAt time.Time + pct5h float64 + reset5hAt time.Time + has5h bool + updatedAt5h time.Time + pct7d float64 + reset7dAt time.Time + has7d bool + updatedAt7d time.Time + persistUsage bool + planObserved bool +} + +// usageFramePersistQueue coalesces frame snapshots per account and writes them +// on a worker goroutine. The upstream read path only captures and enqueues, so +// a slow or locked database cannot stall a streaming response. +type usageFramePersistQueue struct { + mu sync.Mutex + pending map[int64]usageFramePersistJob + wake chan struct{} +} + +func newUsageFramePersistQueue() *usageFramePersistQueue { + return &usageFramePersistQueue{ + pending: make(map[int64]usageFramePersistJob), + wake: make(chan struct{}, 1), + } +} + +func (q *usageFramePersistQueue) start(s *Store) { + if q == nil || s == nil || s.db == nil { + return + } + s.wg.Add(1) + go q.run(s) +} + +func (q *usageFramePersistQueue) run(s *Store) { + defer s.wg.Done() + for { + select { + case <-s.stopCh: + // Usage snapshots are a cache that WHAM can rebuild; drop pending + // jobs instead of delaying shutdown on a blocked database. + return + case <-q.wake: + for { + select { + case <-s.stopCh: + return + default: + } + job, ok := q.take() + if !ok { + break + } + q.persist(s, job) + } + } + } +} + +func (q *usageFramePersistQueue) enqueue(job usageFramePersistJob) { + q.mu.Lock() + if prior, ok := q.pending[job.dbID]; ok { + if job.observedAt.Before(prior.observedAt) { + job, prior = prior, job + } + job.persistUsage = job.persistUsage || prior.persistUsage + job.planObserved = job.planObserved || prior.planObserved + } + q.pending[job.dbID] = job + q.mu.Unlock() + select { + case q.wake <- struct{}{}: + default: + } +} + +func (q *usageFramePersistQueue) take() (usageFramePersistJob, bool) { + q.mu.Lock() + defer q.mu.Unlock() + for id, job := range q.pending { + delete(q.pending, id) + return job, true + } + return usageFramePersistJob{}, false +} + +func (q *usageFramePersistQueue) persist(s *Store, job usageFramePersistJob) { + if s == nil || s.db == nil || job.dbID == 0 { + return + } + acc := s.FindByID(job.dbID) + if acc == nil || acc != job.account { + return + } + // Serialize with the synchronous snapshot writers and re-check inside the + // lock: a WHAM probe or a newer frame may apply while this job waits, and + // its write must win. Holding usageSyncMu here would stall the upstream + // read path, so use this account's snapshot-write lock instead. + acc.usageSnapshotPersistMu.Lock() + if s.FindByID(job.dbID) != acc { + acc.usageSnapshotPersistMu.Unlock() + return + } + acc.mu.RLock() + superseded := acc.usageObservedAt.After(job.observedAt) + acc.mu.RUnlock() + ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second) + defer cancel() + var err error + // Supersession fences the captured quota values, not the pending side + // effects below, which always read current runtime state under their locks. + if job.persistUsage && !superseded { + switch { + case job.has7d && job.has5h: + err = s.db.UpdateUsageSnapshotFull(ctx, job.dbID, job.pct7d, job.reset7dAt, job.pct5h, job.reset5hAt, job.updatedAt7d, job.updatedAt5h) + case job.has7d: + err = s.db.UpdateCredentials(ctx, job.dbID, map[string]interface{}{ + "codex_7d_used_percent": job.pct7d, + "codex_7d_reset_at": job.reset7dAt.Format(time.RFC3339), + "codex_usage_updated_at": job.updatedAt7d.Format(time.RFC3339), + }) + case job.has5h: + err = s.db.UpdateUsageSnapshot5h(ctx, job.dbID, job.pct5h, job.reset5hAt, job.updatedAt5h) + default: + // Credits-only observations need no quota write. + } + if err != nil { + // Leave the throttle marker untouched; the next frame observation sees + // that no successful write happened and enqueues again. + log.Printf("[账号 %d] 异步持久化用量快照失败: %v", job.dbID, err) + } else { + acc.MarkUsageFramePersisted(job.observedAt, job.pct5h, job.has5h, job.pct7d, job.has7d) + acc.mu.Lock() + if !acc.usageObservedAt.After(job.observedAt) { + acc.usageFramePersistPending = false + } + acc.mu.Unlock() + } + } + if err := s.persistUsageFramePlan(acc); err != nil { + log.Printf("[账号 %d] 异步持久化套餐失败: %v", acc.DBID, err) + } + acc.usageSnapshotPersistMu.Unlock() + // Credit/cooldown publication and subscription hooks own their synchronization. + acc.mu.RLock() + creditsValid := acc.CreditsValid + cooldownPending := acc.usageCooldownPersistPending && acc.Status == StatusCooldown + acc.mu.RUnlock() + if creditsValid { + s.persistCreditSnapshot(acc, true) + } + if cooldownPending { + s.persistUsageCooldown(acc) + } + if job.planObserved { + s.ClearStaleSubscriptionExpiresAt(acc) + } +} + +// EnqueueUsageFramePersist schedules a frame snapshot write off the upstream +// read path. The latest snapshot per account wins while a write is queued. +func (s *Store) EnqueueUsageFramePersist(acc *Account, observedAt time.Time) { + s.EnqueueUsageFrameEffects(acc, observedAt, true, false) +} + +// EnqueueUsageFrameEffects coalesces quota, credits, plan and cooldown writes. +// The quota dirty bit survives a newer credits-only/throttled observation so a +// superseded job cannot silently lose an already requested snapshot write. +func (s *Store) EnqueueUsageFrameEffects(acc *Account, observedAt time.Time, persistUsage, planObserved bool) { + if s == nil || acc == nil || s.db == nil || s.usageFramePersist == nil || acc.DBID == 0 { + return + } + select { + case <-s.stopCh: + return + default: + } + if observedAt.IsZero() { + observedAt = time.Now() + } + acc.mu.Lock() + if persistUsage && (acc.UsagePercent5hValid || acc.UsagePercent7dValid) { + acc.usageFramePersistPending = true + } + job := usageFramePersistJob{ + account: acc, + dbID: acc.DBID, + observedAt: observedAt, + pct5h: acc.UsagePercent5h, + reset5hAt: acc.Reset5hAt, + has5h: acc.UsagePercent5hValid, + updatedAt5h: acc.UsageUpdatedAt5h, + pct7d: acc.UsagePercent7d, + reset7dAt: acc.Reset7dAt, + has7d: acc.UsagePercent7dValid, + updatedAt7d: acc.UsageUpdatedAt, + persistUsage: acc.usageFramePersistPending, + planObserved: planObserved, + } + acc.mu.Unlock() + s.usageFramePersist.enqueue(job) +} + +// Called with acc.usageSnapshotPersistMu held, including synchronous plan writes. +func (s *Store) persistUsageFramePlan(acc *Account) error { + acc.mu.RLock() + pending, plan := acc.usageFramePlanPending, acc.PlanType + acc.mu.RUnlock() + if !pending { + return nil + } + ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second) + defer cancel() + if err := s.db.UpdateCredentials(ctx, acc.DBID, map[string]interface{}{"plan_type": plan}); err != nil { + return err + } + acc.mu.Lock() + if acc.PlanType == plan { + acc.usageFramePlanPending = false + } + acc.mu.Unlock() + return nil +} diff --git a/auth/usage_frame_persist_test.go b/auth/usage_frame_persist_test.go new file mode 100644 index 000000000..fe2203f63 --- /dev/null +++ b/auth/usage_frame_persist_test.go @@ -0,0 +1,211 @@ +package auth + +import ( + "context" + "database/sql" + "path/filepath" + "sync" + "testing" + "time" + + "github.com/codex2api/cache" + "github.com/codex2api/database" + "github.com/tidwall/gjson" +) + +func TestUsageSnapshotWritesAreAccountScoped(t *testing.T) { + store := &Store{} + first, second := &Account{}, &Account{} + entered, release, done := make(chan struct{}), make(chan struct{}), make(chan struct{}) + var once sync.Once + unblock := func() { once.Do(func() { close(release) }) } + t.Cleanup(unblock) + go func() { + _ = store.withUsageSnapshotWrite(first, func() error { + close(entered) + <-release + return nil + }) + close(done) + }() + <-entered + otherDone := make(chan struct{}) + go func() { + _ = store.withUsageSnapshotWrite(second, func() error { return nil }) + close(otherDone) + }() + select { + case <-otherDone: + case <-time.After(time.Second): + t.Error("a blocked account stalled another account's snapshot writer") + } + unblock() + <-done + <-otherDone +} + +func TestUsageFramePersistReleasesSnapshotLockBeforeCredits(t *testing.T) { + store, db, acc := newUsageFramePersistenceFixture(t) + store.ObserveUsageFramePlan(acc, "pro") + acc.ApplySparseCreditsHeaders(true, false, nil, nil, "") + acc.creditsPersistMu.Lock() + var once sync.Once + unblock := func() { once.Do(acc.creditsPersistMu.Unlock) } + t.Cleanup(unblock) + done := make(chan struct{}) + go func() { + newUsageFramePersistQueue().persist(store, usageFramePersistJob{account: acc, dbID: acc.DBID, observedAt: time.Now()}) + close(done) + }() + deadline := time.Now().Add(2 * time.Second) + released := false + for time.Now().Before(deadline) { + row, err := db.GetAccountByID(context.Background(), acc.DBID) + if err == nil && row.GetCredential("plan_type") == "pro" && acc.usageSnapshotPersistMu.TryLock() { + acc.usageSnapshotPersistMu.Unlock() + released = true + break + } + time.Sleep(time.Millisecond) + } + unblock() + <-done + if !released { + t.Fatal("snapshot lock remained held while credit persistence was blocked") + } +} + +// TestUsageFramePersistSkipsSupersededJob verifies the worker re-checks the +// observation order under the snapshot-write lock: a queued older frame must +// not overwrite a newer WHAM/frame observation that was applied while the job +// waited for the database. +func TestUsageFramePersistSkipsSupersededJob(t *testing.T) { + ctx := context.Background() + dbPath := filepath.Join(t.TempDir(), "superseded.db") + db, err := database.New("sqlite", dbPath) + if err != nil { + t.Fatalf("database.New: %v", err) + } + t.Cleanup(func() { _ = db.Close() }) + id, err := db.InsertAccountWithCredentials(ctx, "superseded", map[string]any{ + "access_token": "at-test", "plan_type": "plus", + }, "") + if err != nil { + t.Fatalf("InsertAccountWithCredentials: %v", err) + } + store := NewStore(db, nil, &database.SystemSettings{MaxConcurrency: 2}) + t.Cleanup(store.Stop) + account := &Account{DBID: id, AccessToken: "at-test", PlanType: "plus", Status: StatusReady} + store.AddAccount(account) + + t0 := time.Now().Add(-time.Minute) + account.SetUsageSnapshot(10, t0) + + // Block the worker inside the write lock so the newer observation is + // applied before its superseded check runs. + account.usageSnapshotPersistMu.Lock() + store.EnqueueUsageFramePersist(account, t0) + t1 := t0.Add(30 * time.Second) + if !account.ApplyUsageObservation(t1, func() { account.SetUsageSnapshot(80, t1) }) { + t.Fatal("newer observation was rejected") + } + account.usageSnapshotPersistMu.Unlock() + + time.Sleep(200 * time.Millisecond) + conn, err := sql.Open("sqlite", dbPath) + if err != nil { + t.Fatalf("sql.Open: %v", err) + } + defer func() { _ = conn.Close() }() + var raw string + if err := conn.QueryRow(`SELECT credentials FROM accounts WHERE id = $1`, id).Scan(&raw); err != nil { + t.Fatalf("read credentials: %v", err) + } + if got := gjson.Get(raw, "codex_7d_used_percent").Float(); got == 10 { + t.Fatalf("superseded job wrote the older snapshot: %v", got) + } +} + +func newUsageFramePersistenceFixture(t *testing.T) (*Store, *database.DB, *Account) { + t.Helper() + db, err := database.New("sqlite", filepath.Join(t.TempDir(), "frame-effects.db")) + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { _ = db.Close() }) + id, err := db.InsertAccountWithCredentials(context.Background(), "synthetic-frame", map[string]any{"plan_type": "plus"}, "") + if err != nil { + t.Fatal(err) + } + tc := cache.NewMemory(1) + t.Cleanup(func() { _ = tc.Close() }) + store := NewStore(db, tc, &database.SystemSettings{MaxConcurrency: 1}) + t.Cleanup(store.Stop) + acc := &Account{DBID: id, PlanType: "plus", Status: StatusReady} + store.AddAccount(acc) + return store, db, acc +} + +func TestUsageFramePersistSupersededQuotaStillFlushesCurrentEffects(t *testing.T) { + store, db, acc := newUsageFramePersistenceFixture(t) + at := time.Now().UTC().Truncate(time.Second) + reset := at.Add(time.Hour) + acc.SetUsageSnapshot5hAt(100, reset, at) + store.MarkUsageFrameCooldowns(acc, at) + store.ObserveUsageFramePlan(acc, "pro") + acc.ApplySparseCreditsHeaders(true, false, nil, nil, "") + // A newer non-frame observation supersedes the quota job, but does not + // replace its pending credits, plan or cooldown persistence request. + acc.ApplyUsageObservation(at.Add(time.Second), func() { acc.SetUsageSnapshot5hAt(40, reset, at.Add(time.Second)) }) + queue := newUsageFramePersistQueue() + queue.persist(store, usageFramePersistJob{account: acc, dbID: acc.DBID, observedAt: at, has5h: true, pct5h: 100, reset5hAt: reset, persistUsage: true}) + row, err := db.GetAccountByID(context.Background(), acc.DBID) + if err != nil { + t.Fatal(err) + } + if row.GetCredential("codex_5h_used_percent") != "" { + t.Fatal("superseded quota was persisted") + } + if row.GetCredential("plan_type") != "pro" || row.GetCredential("codex_credits") == "" { + t.Fatal("superseded quota job discarded current plan/credits effects") + } + if row.CooldownReason != premium5hCooldownReason || !row.CooldownUntil.Valid || !row.CooldownUntil.Time.Equal(reset) { + t.Fatal("superseded quota job discarded pending cooldown") + } +} + +func TestUsageFramePersist7dOnlyPreservesReset(t *testing.T) { + store, db, acc := newUsageFramePersistenceFixture(t) + at := time.Now().UTC().Truncate(time.Second) + reset := at.Add(7 * 24 * time.Hour) + queue := newUsageFramePersistQueue() + queue.persist(store, usageFramePersistJob{account: acc, dbID: acc.DBID, observedAt: at.Add(time.Second), has7d: true, pct7d: 42, reset7dAt: reset, updatedAt7d: at, persistUsage: true}) + row, err := db.GetAccountByID(context.Background(), acc.DBID) + if err != nil { + t.Fatal(err) + } + if row.GetCredential("codex_7d_used_percent") != "42" || row.GetCredential("codex_7d_reset_at") != reset.Format(time.RFC3339) { + t.Fatal("7d-only frame did not persist its reset") + } + if row.GetCredential("codex_5h_used_percent") != "" { + t.Fatal("7d-only frame fabricated a 5h window") + } + if row.GetCredential("codex_usage_updated_at") != at.Format(time.RFC3339) { + t.Fatal("side-effect observation changed the persisted quota timestamp") + } +} + +func TestUsageFramePersistRejectsReplacedAccount(t *testing.T) { + store, db, acc := newUsageFramePersistenceFixture(t) + job := usageFramePersistJob{account: acc, dbID: acc.DBID, observedAt: time.Now(), has5h: true, pct5h: 99, persistUsage: true} + store.RemoveAccount(acc.DBID) + store.AddAccount(&Account{DBID: acc.DBID, PlanType: "plus", Status: StatusReady}) + newUsageFramePersistQueue().persist(store, job) + row, err := db.GetAccountByID(context.Background(), acc.DBID) + if err != nil { + t.Fatal(err) + } + if row.GetCredential("codex_5h_used_percent") != "" { + t.Fatal("old account job wrote onto its replacement") + } +} diff --git a/auth/usage_frame_test.go b/auth/usage_frame_test.go new file mode 100644 index 000000000..d0e1c4b11 --- /dev/null +++ b/auth/usage_frame_test.go @@ -0,0 +1,117 @@ +package auth + +import ( + "testing" + "time" +) + +func TestShouldPersistUsageFrameThrottlesAndForces(t *testing.T) { + acc := newQuotaAutoPauseTestAccount() + now := time.Unix(1_700_000_000, 0) + interval := time.Minute + + if !acc.ShouldPersistUsageFrame(now, 10, true, 5, true, interval) { + t.Fatal("first observation must persist") + } + acc.MarkUsageFramePersisted(now, 10, true, 5, true) + if acc.ShouldPersistUsageFrame(now.Add(time.Second), 10.5, true, 5, true, interval) { + t.Fatal("small delta inside the interval must be throttled") + } + if !acc.ShouldPersistUsageFrame(now.Add(time.Second), 16, true, 5, true, interval) { + t.Fatal("a 5pp delta must force persistence") + } + acc.MarkUsageFramePersisted(now.Add(time.Second), 16, true, 5, true) + if acc.ShouldPersistUsageFrame(now.Add(2*time.Second), 16.5, true, 5, true, interval) { + t.Fatal("small delta after a forced write must be throttled again") + } + if !acc.ShouldPersistUsageFrame(now.Add(2*time.Minute), 16.5, true, 5, true, interval) { + t.Fatal("elapsed interval must persist") + } +} + +func TestShouldPersistUsageFrameRetriesUntilMarked(t *testing.T) { + acc := newQuotaAutoPauseTestAccount() + now := time.Unix(1_700_000_000, 0) + + if !acc.ShouldPersistUsageFrame(now, 10, true, 5, true, time.Minute) { + t.Fatal("first observation must persist") + } + // Simulate a failed write: without MarkUsageFramePersisted the decision + // must be retried instead of being throttled for the interval. + if !acc.ShouldPersistUsageFrame(now.Add(time.Second), 10, true, 5, true, time.Minute) { + t.Fatal("failed write must be retried on the next observation") + } +} + +func TestShouldPersistUsageFrameForcesThresholdCrossing(t *testing.T) { + acc := newQuotaAutoPauseTestAccount() + now := time.Unix(1_700_000_000, 0) + acc.effectiveAutoPause5h = 0.5 + + if !acc.ShouldPersistUsageFrame(now, 48, true, 0, false, time.Hour) { + t.Fatal("first observation must persist") + } + acc.MarkUsageFramePersisted(now, 48, true, 0, false) + if !acc.ShouldPersistUsageFrame(now.Add(time.Second), 50, true, 0, false, time.Hour) { + t.Fatal("crossing the auto-pause threshold must force persistence") + } +} + +func TestGetAccountLiveUsageSnapshotProjectsStatusAndUsage(t *testing.T) { + acc := newQuotaAutoPauseTestAccount() + acc.AutoPause5hThreshold = 0.9 + acc.UsagePercent5h = 95 + acc.UsagePercent5hValid = true + acc.Reset5hAt = time.Now().Add(time.Hour) + acc.UsagePercent7d = 31 + acc.UsagePercent7dValid = true + acc.UsageUpdatedAt = time.Now() + setAutoPauseThresholdsWithGuard(acc, 5) + + snapshot := acc.GetAccountLiveUsageSnapshot() + if snapshot.Status != "quota_paused" { + t.Fatalf("status = %q, want quota_paused", snapshot.Status) + } + if !snapshot.UsagePercent5hValid || snapshot.UsagePercent5h != 95 { + t.Fatalf("5h snapshot = %+v", snapshot) + } + if !snapshot.UsagePercent7dValid || snapshot.UsagePercent7d != 31 { + t.Fatalf("7d snapshot = %+v", snapshot) + } +} + +func TestTryApplyUsageObservationSkipsWhileLocked(t *testing.T) { + acc := &Account{} + acc.usageSyncMu.Lock() + applied := false + start := time.Now() + if acc.TryApplyUsageObservation(time.Now(), func() { applied = true }) { + t.Fatal("TryApply succeeded while the observation lock was held") + } + if elapsed := time.Since(start); elapsed > 100*time.Millisecond { + t.Fatalf("TryApply blocked for %v, want a non-blocking skip", elapsed) + } + acc.usageSyncMu.Unlock() + if applied { + t.Fatal("skipped observation still ran its callback") + } + + if !acc.TryApplyUsageObservation(time.Now(), func() { applied = true }) || !applied { + t.Fatal("TryApply did not run after the lock was released") + } +} + +func TestTryApplyUsageObservationRejectsOlderTimestamp(t *testing.T) { + acc := &Account{} + now := time.Now() + applied := 0 + if !acc.TryApplyUsageObservation(now, func() { applied++ }) { + t.Fatal("first observation was rejected") + } + if acc.TryApplyUsageObservation(now.Add(-time.Minute), func() { applied++ }) { + t.Fatal("older observation was accepted") + } + if applied != 1 { + t.Fatalf("callbacks = %d, want 1", applied) + } +} diff --git a/docs/API.md b/docs/API.md index b35264076..d57bf1e1d 100644 --- a/docs/API.md +++ b/docs/API.md @@ -774,6 +774,70 @@ Codex 的流式 remote compact v2（`POST /v1/responses`，`stream:true`，`inpu | billed_5h | number/null | 过去 5 小时窗口内的累计计费金额（USD） | | billed_7d | number/null | 过去 7 天窗口内的累计计费金额（USD） | +#### GET /api/admin/accounts/live + +返回当前页账号的运行时并发计数与用量投影，供账号页约每秒一次的轻量轮询。只读取内存状态，不查数据库、不重建调度快照。 + +**请求参数:** + +| 参数 | 类型 | 说明 | +| ---- | ------ | ---------------------------------- | +| ids | string | 逗号分隔的账号 ID，最多 500 个 | + +**响应:** + +```json +{ + "accounts": { + "1": { + "active_requests": 1, + "occupied_requests": 1, + "dynamic_concurrency_limit": 2, + "base_concurrency_effective": 2, + "status": "quota_paused", + "using_credits": false, + "plan_type": "pro", + "credits_valid": true, + "credits_balance": "12.5", + "credits_has_credits": true, + "credits_unlimited": false, + "credits_overage_limit_reached": false, + "credits_spend_control_reached": null, + "credits_rate_limit_reached_type": "", + "usage_percent_5h": 91.2, + "reset_5h_at": "2026-10-18T18:00:00+08:00", + "codex_5h_usage_updated_at": "2026-10-18T15:10:00+08:00", + "usage_percent_7d": 42, + "reset_7d_at": "2026-10-25T07:00:00+08:00", + "codex_usage_updated_at": "2026-10-18T15:10:00+08:00", + "codex_usage_frame_at": "2026-10-18T15:10:00+08:00", + "codex_usage_frame_kind": "codex.rate_limits" + } + }, + "session_slot_buffer_enabled": false +} +``` + +**字段说明:** + +| 字段 | 类型 | 说明 | +| ---- | ---- | ---- | +| status / using_credits | string / bool | 运行时状态；旧后端或 Antigravity 行省略该 key，前端保留列表值 | +| plan_type | string | 最新运行时套餐，沿用账号列表字段；旧后端或 Antigravity 行省略时保留列表值 | +| credits_valid / credits_balance | bool / string/null | 当前积分快照是否已观测、余额；未知或隐藏的余额为 `null`，清除旧显示值 | +| credits_has_credits / credits_unlimited / credits_overage_limit_reached | bool/null | 沿用账号列表的积分业务标志；未观测快照时为 `null` | +| credits_spend_control_reached / credits_rate_limit_reached_type | bool/null / string | 支出限制与工作区硬停止；未知支出限制为 `null`，无硬停止为 `""` | +| usage_percent_5h / usage_percent_7d | number/null | 运行时用量；显式 `null` 表示窗口已失效，前端清除列表值 | +| reset_5h_at / reset_7d_at | string/null | 窗口重置时间（RFC3339） | +| codex_5h_usage_updated_at / codex_usage_updated_at | string/null | 最近一次用量刷新时间；WS 每轮配额帧也会刷新 | +| codex_usage_frame_at / codex_usage_frame_kind | string | 最近一次被应用的 WS 配额帧时间与事件类型；首帧前省略，可用于确认帧摄入是否生效 | +| dynamic_concurrency_limit / base_concurrency_effective | integer | 调度器当前生效的并发上限与配置值 | +| session_slot_buffer_enabled | bool | 全局会话槽缓冲开关 | + +缺失的 key 表示旧后端，前端保留列表值；用量字段显式 `null` 表示窗口失效。套餐与 credits 的 live 字段只投影到使用 Codex 运行时字段的行，Antigravity 行省略这些字段，不改变原列表展示。 + +终态后到达的连接级配额帧也会更新上述字段。只有通过观测排序门的帧才推进帧时间、credits 和用量；遇到并发观测锁争用时按原始时间排队应用，不等待数据库。用量、credits、套餐与冷却的数据库写入在后台合并执行，因此 live 结果可先于持久化结果可见。 + #### PATCH /api/admin/accounts/:id/scheduler 更新账号调度配置。 diff --git a/docs/CONFIGURATION.md b/docs/CONFIGURATION.md index 7f07ee404..49c175fff 100644 --- a/docs/CONFIGURATION.md +++ b/docs/CONFIGURATION.md @@ -87,6 +87,8 @@ Codex2API 采用三层配置架构： | `USE_WEBSOCKET` | 否 | `false` | 旧版开关；未设置 `CODEX_UPSTREAM_TRANSPORT` 时，`true` 等价于 `CODEX_UPSTREAM_TRANSPORT=ws` | | `CODEX_TRANSPORT_MODE` | 否 | `standard` | Codex HTTP transport：默认标准 Go TLS；`utls_chrome` 可回滚旧 Chrome uTLS 行为 | | `CODEX_WS_SEND_USER_AGENT` | 否 | `true` | WS 握手是否发送 Codex `User-Agent`/`Version`；设为 `false` 可关闭 | +| `CODEX_WS_USAGE_FRAME_PERSIST_SECONDS` | 否 | `30` | WS 逐帧用量快照落库的最小间隔（秒），正整数，重启生效。内存快照每帧更新；落库由每账号合并的异步 worker 执行，不阻塞上游响应流；跨自动暂停阈值、5pp 变化或 100% 时仍强制落库。大号池可按数据库承压调大 | +| `CODEX_WS_USAGE_FRAME_INGEST` | 否 | `true` | 是否消费 WS 上游每轮及终态后连接级 `codex.rate_limits` / `codex.response.metadata` 配额帧；设为 `false` 时回退到 WHAM 探针与 HTTP 响应头，重启生效 | | `CODEX_SESSION_AFFINITY_TTL` | 否 | `1h` | Codex 会话到账号/代理的黏性 TTL，支持 `1h`、`90m` 或秒数 | | `CODEX_COMPACTION_AFFINITY_TTL` | 否 | `168h` | 加密压缩状态的来源亲和 TTL。缓存仅保存密文的 SHA-256 摘要、来源账号和兼容域；已知状态不会跨 Codex 官方、不同 Responses 中转或 Grok 上游流转 | | `CODEX_FINGERPRINT_DEBUG` | 否 | `false` | 输出脱敏指纹策略诊断日志，不记录 token | @@ -98,7 +100,11 @@ Codex2API 采用三层配置架构： | `DOWNSTREAM_HTTP_KEEPALIVE_INTERVAL` | 否 | `30s` | 下游 HTTP/SSE 保活周期，使用 Go duration；`0` 关闭。流式端点从首个心跳起建立 SSE 200，发送注释或 Messages ping；非流式端点发送 HTTP 102 | | `DOWNSTREAM_WS_KEEPALIVE_INTERVAL` | 否 | `45s` | 下游 WebSocket Ping 周期，使用 Go duration；`0` 关闭。覆盖 Responses、Realtime 与 Live Sideband | +配额帧读路径只更新内存和调度器；用量、credits、套餐及冷却的写库共用后台合并路径。并发观测锁争用时，每账号最多保留 64 个待应用 partial 帧，按原始时间顺序处理；溢出保留最新的有界尾部，并通过已有探针补齐权威快照。停机丢弃未完成的缓存写，重启后由探针重建。冷却写失败会在后续观测重试，不重复计入失败次数。 + > `CODEX_UPSTREAM_TRANSPORT` 只控制 HTTP 入站请求转发到 Codex 上游时使用 `http` 还是 `ws`。客户端侧 WebSocket 入口独立可用：使用 `GET ws:///v1/responses` 建连，首帧发送 `response.create` JSON，服务端会通过 Codex 上游 WS 返回 Responses 事件帧。 +> +> WS 传输下，服务端会消费上游每轮下发的 `codex.rate_limits` 帧，以及响应终态后 1–2 秒补发的连接级用量帧，实时更新账号的 5h/7d 用量快照；自动暂停阈值不再需要等待 WHAM 探针周期即可生效。这些帧随对话本身下发，不产生额外上游请求。 ### 数据库配置 diff --git a/frontend/src/hooks/useAccountLiveState.ts b/frontend/src/hooks/useAccountLiveState.ts index b8b2bf825..bae255a7e 100644 --- a/frontend/src/hooks/useAccountLiveState.ts +++ b/frontend/src/hooks/useAccountLiveState.ts @@ -1,53 +1,9 @@ import { useEffect, useMemo } from 'react' import { api } from '../api' import type { AccountLiveStateResponse } from '../types' +import { mergeAccountLiveState } from '../lib/accountLiveState' -// Merge a live poll response into an account list. Rows whose live concurrency -// counters did not change keep their object identity, and a no-op poll returns the -// original array, so the 1s polling cadence cannot defeat row-level memoization. -export function mergeAccountLiveState( - current: T[], - response: AccountLiveStateResponse, -): T[] { - let changed = false - const next = current.map((account) => { - const live = response.accounts[String(account.id)] - const activeRequests = live?.active_requests ?? 0 - const occupiedRequests = live?.occupied_requests ?? activeRequests - const slotBufferEnabled = response.session_slot_buffer_enabled === true - // Concurrency caps are optional so an older backend keeps the list values. - const dynamicLimit = typeof live?.dynamic_concurrency_limit === 'number' - ? live.dynamic_concurrency_limit - : account.dynamic_concurrency_limit - const baseConcurrency = typeof live?.base_concurrency_effective === 'number' - ? live.base_concurrency_effective - : account.base_concurrency_effective - if ( - (account.active_requests ?? 0) === activeRequests && - (account.occupied_requests ?? account.active_requests ?? 0) === occupiedRequests && - (account.session_slot_buffer_enabled ?? false) === slotBufferEnabled && - account.dynamic_concurrency_limit === dynamicLimit && - account.base_concurrency_effective === baseConcurrency - ) return account - changed = true - return { - ...account, - active_requests: activeRequests, - occupied_requests: occupiedRequests, - session_slot_buffer_enabled: slotBufferEnabled, - dynamic_concurrency_limit: dynamicLimit, - base_concurrency_effective: baseConcurrency, - } - }) - return changed ? next : current -} +export { mergeAccountLiveState } export function useAccountLiveState( ids: number[], @@ -69,15 +25,21 @@ export function useAccountLiveState( return } controller?.abort() - controller = new AbortController() + const current = new AbortController() + controller = current + // A hung fetch must not stop the poll forever; abort it and let the next + // tick retry. The list reload remains the fallback for stale data. + const timeout = window.setTimeout(() => current.abort(), 5000) try { const response = await api.getAccountLiveState( idsKey.split(',').map(Number), - controller.signal, + current.signal, ) - if (!stopped && !controller.signal.aborted) apply(response) + if (!stopped && !current.signal.aborted) apply(response) } catch { // The next tick retries; live counters must never disrupt the page. + } finally { + window.clearTimeout(timeout) } if (!stopped) timer = window.setTimeout(poll, 1000) } diff --git a/frontend/src/lib/accountLiveState.test.mjs b/frontend/src/lib/accountLiveState.test.mjs new file mode 100644 index 000000000..f4c3ad8c0 --- /dev/null +++ b/frontend/src/lib/accountLiveState.test.mjs @@ -0,0 +1,165 @@ +import assert from "node:assert/strict"; +import test from "node:test"; +import { mergeAccountLiveState } from "./accountLiveState.ts"; + +const baseRow = { + id: 1, + status: "active", + using_credits: false, + plan_type: "plus", + credits_valid: true, + credits_balance: "10", + credits_has_credits: true, + usage_percent_5h: 40, + reset_5h_at: "2026-10-18T17:00:00+08:00", + codex_5h_usage_updated_at: "2026-10-18T13:00:00+08:00", + usage_percent_7d: 31, + reset_7d_at: "2026-10-25T07:00:00+08:00", + codex_usage_updated_at: "2026-10-18T13:00:00+08:00", + active_requests: 1, + occupied_requests: 1, + dynamic_concurrency_limit: 5, + base_concurrency_effective: 5, +}; + +function live(overrides = {}) { + return { + session_slot_buffer_enabled: false, + accounts: { + "1": { + active_requests: 1, + occupied_requests: 1, + dynamic_concurrency_limit: 5, + base_concurrency_effective: 5, + status: "active", + using_credits: false, + plan_type: baseRow.plan_type, + credits_valid: baseRow.credits_valid, + credits_balance: baseRow.credits_balance, + credits_has_credits: baseRow.credits_has_credits, + usage_percent_5h: 40, + reset_5h_at: baseRow.reset_5h_at, + codex_5h_usage_updated_at: baseRow.codex_5h_usage_updated_at, + usage_percent_7d: baseRow.usage_percent_7d, + reset_7d_at: baseRow.reset_7d_at, + codex_usage_updated_at: baseRow.codex_usage_updated_at, + ...overrides, + }, + }, + }; +} + +test("no live entry keeps the original array", () => { + const rows = [baseRow]; + assert.equal( + mergeAccountLiveState(rows, { accounts: {}, session_slot_buffer_enabled: false }), + rows, + ); +}); + +test("unchanged live fields keep the same row object", () => { + const rows = [baseRow]; + const merged = mergeAccountLiveState(rows, live()); + assert.equal(merged, rows); + assert.equal(merged[0], baseRow); +}); + +test("usage change only replaces the affected row", () => { + const other = { ...baseRow, id: 2 }; + const rows = [baseRow, other]; + const merged = mergeAccountLiveState(rows, live({ usage_percent_5h: 91, status: "quota_paused" })); + assert.notEqual(merged, rows); + assert.notEqual(merged[0], baseRow); + assert.equal(merged[0].usage_percent_5h, 91); + assert.equal(merged[0].status, "quota_paused"); + assert.equal(merged[1], other); +}); + +test("older backend without usage keys preserves the list values", () => { + const rows = [baseRow]; + const merged = mergeAccountLiveState(rows, { + session_slot_buffer_enabled: false, + accounts: { + "1": { + active_requests: 2, + occupied_requests: 2, + dynamic_concurrency_limit: 4, + }, + }, + }); + assert.equal(merged[0].usage_percent_5h, 40); + assert.equal(merged[0].reset_5h_at, baseRow.reset_5h_at); + assert.equal(merged[0].status, "active"); + assert.equal(merged[0].active_requests, 2); +}); + +test("explicit null clears an invalid window", () => { + const rows = [baseRow]; + const merged = mergeAccountLiveState(rows, live({ + usage_percent_5h: null, + reset_5h_at: null, + codex_5h_usage_updated_at: null, + })); + assert.equal(merged[0].usage_percent_5h, null); + assert.equal(merged[0].reset_5h_at, undefined); + assert.equal(merged[0].codex_5h_usage_updated_at, undefined); +}); + +test("explicit null clears an invalid 7d window", () => { + const rows = [baseRow]; + const merged = mergeAccountLiveState(rows, live({ + usage_percent_7d: null, + reset_7d_at: null, + codex_usage_updated_at: null, + })); + assert.equal(merged[0].usage_percent_7d, null); + assert.equal(merged[0].reset_7d_at, undefined); + assert.equal(merged[0].codex_usage_updated_at, undefined); +}); + +test("missing status and credits fields keep the list values", () => { + const rows = [baseRow]; + const payload = live(); + delete payload.accounts["1"].status; + delete payload.accounts["1"].using_credits; + const merged = mergeAccountLiveState(rows, payload); + assert.equal(merged[0].status, "active"); + assert.equal(merged[0].using_credits, false); +}); + +test("mixed rows only replace changed entries", () => { + const other = { ...baseRow, id: 2, usage_percent_7d: 20 }; + const rows = [baseRow, other]; + const payload = live({ usage_percent_7d: 35 }); + payload.accounts["2"] = { ...payload.accounts["1"], usage_percent_7d: 20 }; + const merged = mergeAccountLiveState(rows, payload); + assert.notEqual(merged[0], baseRow); + assert.equal(merged[1], other); +}); + +test("session slot buffer change is merged into rows", () => { + const rows = [baseRow]; + const payload = live(); + payload.session_slot_buffer_enabled = true; + const merged = mergeAccountLiveState(rows, payload); + assert.equal(merged[0].session_slot_buffer_enabled, true); +}); + +test("plan and credit changes reach the existing account row", () => { + const other = { ...baseRow, id: 2 }; + const rows = [baseRow, other]; + const merged = mergeAccountLiveState(rows, live({ plan_type: "pro", credits_balance: "0", credits_has_credits: false })); + assert.equal(merged[0].plan_type, "pro"); + assert.equal(merged[0].credits_balance, "0"); + assert.equal(merged[0].credits_has_credits, false); + assert.equal(merged[1], other); +}); + +test("hidden credit balance clears the prior value and absent fields preserve it", () => { + const merged = mergeAccountLiveState([baseRow], live({ credits_balance: null })); + assert.equal(merged[0].credits_balance, null); + const payload = live(); + delete payload.accounts["1"].credits_balance; + delete payload.accounts["1"].plan_type; + assert.equal(mergeAccountLiveState([baseRow], payload)[0], baseRow); +}); diff --git a/frontend/src/lib/accountLiveState.ts b/frontend/src/lib/accountLiveState.ts new file mode 100644 index 000000000..636199cff --- /dev/null +++ b/frontend/src/lib/accountLiveState.ts @@ -0,0 +1,128 @@ +import type { AccountLiveStateResponse } from '../types' + +// Account rows carry a large optional usage/status projection from the list +// API. The live poll refreshes the same fields in place so the card's bars, +// status badge and countdown react to per-turn WS usage observations without a +// list reload. Missing keys mean "older backend: keep the list value"; an +// explicit null means "the window is no longer valid: clear it". +export type AccountLiveMergeable = { + id: number + active_requests?: number + occupied_requests?: number + session_slot_buffer_enabled?: boolean + dynamic_concurrency_limit?: number + base_concurrency_effective?: number + status?: string + using_credits?: boolean + plan_type?: string + credits_valid?: boolean + credits_balance?: string | null + credits_has_credits?: boolean | null + credits_unlimited?: boolean | null + credits_overage_limit_reached?: boolean | null + credits_spend_control_reached?: boolean | null + credits_rate_limit_reached_type?: string | null + usage_percent_5h?: number | null + reset_5h_at?: string | null + codex_5h_usage_updated_at?: string | null + usage_percent_7d?: number | null + reset_7d_at?: string | null + codex_usage_updated_at?: string | null +} + +export function mergeAccountLiveState( + current: T[], + response: AccountLiveStateResponse, +): T[] { + let changed = false + const next = current.map((account) => { + const live = response.accounts[String(account.id)] + if (!live) return account + + const activeRequests = live.active_requests ?? 0 + const occupiedRequests = live.occupied_requests ?? activeRequests + const slotBufferEnabled = response.session_slot_buffer_enabled === true + // Concurrency caps are optional so an older backend keeps the list values. + const dynamicLimit = typeof live.dynamic_concurrency_limit === 'number' + ? live.dynamic_concurrency_limit + : account.dynamic_concurrency_limit + const baseConcurrency = typeof live.base_concurrency_effective === 'number' + ? live.base_concurrency_effective + : account.base_concurrency_effective + // Usage/status fields follow the same contract: undefined keeps the row, + // null clears it. The list API omits some fields for Antigravity rows. + const status = live.status === undefined ? account.status : live.status + const usingCredits = live.using_credits === undefined + ? account.using_credits + : live.using_credits + const usage5h = live.usage_percent_5h === undefined + ? account.usage_percent_5h + : live.usage_percent_5h + const reset5h = live.reset_5h_at === undefined + ? account.reset_5h_at + : (live.reset_5h_at ?? undefined) + const updated5h = live.codex_5h_usage_updated_at === undefined + ? account.codex_5h_usage_updated_at + : (live.codex_5h_usage_updated_at ?? undefined) + const usage7d = live.usage_percent_7d === undefined + ? account.usage_percent_7d + : live.usage_percent_7d + const reset7d = live.reset_7d_at === undefined + ? account.reset_7d_at + : (live.reset_7d_at ?? undefined) + const updated7d = live.codex_usage_updated_at === undefined + ? account.codex_usage_updated_at + : (live.codex_usage_updated_at ?? undefined) + // Existing plan/credit badges use the list API's field names and null + // semantics. A hidden balance must clear an older visible value. + const details = { + plan_type: live.plan_type === undefined ? account.plan_type : live.plan_type, + credits_valid: live.credits_valid === undefined ? account.credits_valid : live.credits_valid, + credits_balance: live.credits_balance === undefined ? account.credits_balance : live.credits_balance, + credits_has_credits: live.credits_has_credits === undefined ? account.credits_has_credits : live.credits_has_credits, + credits_unlimited: live.credits_unlimited === undefined ? account.credits_unlimited : live.credits_unlimited, + credits_overage_limit_reached: live.credits_overage_limit_reached === undefined ? account.credits_overage_limit_reached : live.credits_overage_limit_reached, + credits_spend_control_reached: live.credits_spend_control_reached === undefined ? account.credits_spend_control_reached : live.credits_spend_control_reached, + credits_rate_limit_reached_type: live.credits_rate_limit_reached_type === undefined ? account.credits_rate_limit_reached_type : live.credits_rate_limit_reached_type, + } + const detailsUnchanged = (Object.keys(details) as (keyof typeof details)[]) + .every((key) => account[key] === details[key]) + + if ( + detailsUnchanged && + (account.active_requests ?? 0) === activeRequests && + (account.occupied_requests ?? account.active_requests ?? 0) === occupiedRequests && + (account.session_slot_buffer_enabled ?? false) === slotBufferEnabled && + account.dynamic_concurrency_limit === dynamicLimit && + account.base_concurrency_effective === baseConcurrency && + account.status === status && + account.using_credits === usingCredits && + account.usage_percent_5h === usage5h && + account.reset_5h_at === reset5h && + account.codex_5h_usage_updated_at === updated5h && + account.usage_percent_7d === usage7d && + account.reset_7d_at === reset7d && + account.codex_usage_updated_at === updated7d + ) return account + + changed = true + return { + ...account, + ...details, + active_requests: activeRequests, + occupied_requests: occupiedRequests, + session_slot_buffer_enabled: slotBufferEnabled, + dynamic_concurrency_limit: dynamicLimit, + base_concurrency_effective: baseConcurrency, + status, + using_credits: usingCredits, + usage_percent_5h: usage5h, + reset_5h_at: reset5h, + codex_5h_usage_updated_at: updated5h, + usage_percent_7d: usage7d, + reset_7d_at: reset7d, + codex_usage_updated_at: updated7d, + } + }) + return changed ? next : current +} diff --git a/frontend/src/locales/en.json b/frontend/src/locales/en.json index b78fc4747..3f9bad69c 100644 --- a/frontend/src/locales/en.json +++ b/frontend/src/locales/en.json @@ -1543,15 +1543,15 @@ "cleanBannedTitle": "Clear Banned Accounts", "cleanBannedDesc": "This clears all unauthorized status flags so those accounts can rejoin scheduling and testing.", "cleanRateLimitedTitle": "Clear Rate-Limited Accounts", - "cleanRateLimitedDesc": "This clears all rate_limited status flags so those accounts can return to the schedulable pool.", + "cleanRateLimitedDesc": "This moves every rate-limited account in the pool to the recycle bin (restorable there), regardless of the current filter or page. Auto-paused accounts are included when their runtime status is rate-limited or usage_exhausted.", "cleanConfirm": "Confirm Cleanup", "cleanBannedProgressTitle": "Clearing banned accounts", "cleanBannedSuccess": "Banned account flags cleared", "cleanBannedSuccessCount": "Cleared {{count}} banned accounts", "cleanBannedFailed": "Failed to clear banned accounts: {{error}}", "cleanRateLimitedProgressTitle": "Clearing rate-limited accounts", - "cleanRateLimitedSuccess": "Rate-limited account flags cleared", - "cleanRateLimitedSuccessCount": "Cleared {{count}} rate-limited accounts", + "cleanRateLimitedSuccess": "Rate-limited accounts moved to the recycle bin", + "cleanRateLimitedSuccessCount": "Moved {{count}} rate-limited accounts to the recycle bin", "cleanRateLimitedFailed": "Failed to clear rate-limited accounts: {{error}}", "cleanErrorTitle": "Clear Error Accounts", "cleanErrorDesc": "This removes all accounts in error state (e.g. expired refresh tokens) from the pool.", @@ -1631,7 +1631,7 @@ "autoPause5hThreshold": "5h Usage Threshold (%)", "autoPause7dThreshold": "7d Usage Threshold (%)", "autoPauseThresholdPlaceholder": "Blank or 0", - "autoPauseThresholdHint": "Leave blank or enter 0 to disable this window threshold. Reaching it only skips scheduling.", + "autoPauseThresholdHint": "Leave blank or enter 0 to inherit the group/global threshold. To disable this window's auto-pause entirely, turn on the matching switch above. Reaching the threshold only skips scheduling.", "autoPauseThresholdRange": "Enter a percentage from 0 to 100.", "ignoreUsageLimitStatus": "Ignore Usage-Limit Status", "ignoreUsageLimitStatusHint": "When enabled, a bound active Codex turn carrying turn state may continue after 5h/7d usage reaches 100%. Fresh turns remain limited, and an actual Responses usage_limit_reached stops continuation. Manual disablement, auth failures, and configured auto-pause thresholds still apply.", diff --git a/frontend/src/locales/zh.json b/frontend/src/locales/zh.json index 812f0d4ce..4a7aa5c07 100644 --- a/frontend/src/locales/zh.json +++ b/frontend/src/locales/zh.json @@ -1541,15 +1541,15 @@ "cleanBannedTitle": "清理封禁账号", "cleanBannedDesc": "会清除所有 unauthorized 状态标记，让这些账号重新参与调度与测试。", "cleanRateLimitedTitle": "清理限流账号", - "cleanRateLimitedDesc": "会清除所有 rate_limited 状态标记，让这些账号重新回到可调度状态。", + "cleanRateLimitedDesc": "会把全池处于限流状态的账号移入回收站（可在回收站恢复），不受当前筛选或分页影响；自动暂停的账号如果同时处于限流或 usage_exhausted 状态，也会被纳入清理范围。", "cleanConfirm": "确认清理", "cleanBannedProgressTitle": "正在清理封禁账号", "cleanBannedSuccess": "封禁账号状态已清理", "cleanBannedSuccessCount": "已清理 {{count}} 个封禁账号", "cleanBannedFailed": "清理封禁账号失败：{{error}}", "cleanRateLimitedProgressTitle": "正在清理限流账号", - "cleanRateLimitedSuccess": "限流账号状态已清理", - "cleanRateLimitedSuccessCount": "已清理 {{count}} 个限流账号", + "cleanRateLimitedSuccess": "限流账号已移入回收站", + "cleanRateLimitedSuccessCount": "已将 {{count}} 个限流账号移入回收站", "cleanRateLimitedFailed": "清理限流账号失败：{{error}}", "cleanErrorTitle": "清理错误账号", "cleanErrorDesc": "会移除所有 error 状态的账号（如 refresh token 失效），从号池中彻底清除。", @@ -1629,7 +1629,7 @@ "autoPause5hThreshold": "5h 用量阈值(%)", "autoPause7dThreshold": "7d 用量阈值(%)", "autoPauseThresholdPlaceholder": "留空或 0", - "autoPauseThresholdHint": "留空或填 0 表示不启用该窗口阈值；达到阈值后只在调度时跳过账号。", + "autoPauseThresholdHint": "留空或填 0 表示继承分组/全局阈值；要完全停用该窗口的自动暂停，请打开上方对应的禁用开关。达到阈值后只在调度时跳过账号。", "autoPauseThresholdRange": "请输入 0 到 100 之间的百分比。", "ignoreUsageLimitStatus": "忽略用量限流状态", "ignoreUsageLimitStatusHint": "开启后，仅允许携带 turn state 且已绑定账号的活跃 Codex 轮次在 5h/7d 用量达到 100% 后继续；新轮次仍会限流，Responses 实际返回 usage_limit_reached 时立即停止续传。不会绕过手动禁用、鉴权失败或自定义自动暂停阈值。", diff --git a/frontend/src/pages/Accounts.tsx b/frontend/src/pages/Accounts.tsx index 2d3a8edcb..cebcb6e7d 100644 --- a/frontend/src/pages/Accounts.tsx +++ b/frontend/src/pages/Accounts.tsx @@ -3056,40 +3056,91 @@ export default function Accounts() { ); // Codex 视图的统计卡/额度分布/列表/批量操作一律排除 Grok 账号 // （Grok 账号由顶部切换后的 Grok 页单独统计与管理）。 - const allAccounts = useMemo( - () => - data.accounts.map((account) => { - const stats = accountPageStats[String(account.id)]; - if (!stats) return account; - // 行自身已带的字段优先(如 refreshAccountRow 拉回的完整详情比 - // 本页 stats 快照更新),page-stats 只补基础行缺失的部分。 - // 官方结算例外：点进官方统计刷新后必须以最新快照为准， - // 否则行上的旧额度会挡住这次同步时间点。 - const merged = { ...account }; - if (merged.billed_5h == null && stats.billed_5h != null) merged.billed_5h = stats.billed_5h; - if (merged.billed_7d == null && stats.billed_7d != null) merged.billed_7d = stats.billed_7d; - const officialUsd = stats.official_usd ?? stats.official_usd_7d; - if (officialUsd != null) { - merged.official_usd = officialUsd; - merged.official_usd_7d = officialUsd; - } else if (stats.official_usage_synced) { - merged.official_usd = undefined; - merged.official_usd_7d = undefined; - } - if (stats.official_usage_synced != null) { - merged.official_usage_synced = stats.official_usage_synced; - } - if (!merged.usage_5h_detail && stats.usage_5h_detail) merged.usage_5h_detail = stats.usage_5h_detail; - if (!merged.usage_7d_detail && stats.usage_7d_detail) merged.usage_7d_detail = stats.usage_7d_detail; - if (!merged.usage_today_detail && stats.usage_today_detail) merged.usage_today_detail = stats.usage_today_detail; - return merged; - }), - [data.accounts, accountPageStats], - ); + const allAccounts = useMemo(() => { + const sameOptional = (a: unknown, b: unknown) => + (a ?? undefined) === (b ?? undefined); + return data.accounts.map((account) => { + const stats = accountPageStats[String(account.id)]; + if (!stats) return account; + // 行自身已带的字段优先(如 refreshAccountRow 拉回的完整详情比 + // 本页 stats 快照更新),page-stats 只补基础行缺失的部分。 + // 官方结算例外：点进官方统计刷新后必须以最新快照为准， + // 否则行上的旧额度会挡住这次同步时间点。 + const billed5h = + account.billed_5h == null && stats.billed_5h != null + ? stats.billed_5h + : account.billed_5h; + const billed7d = + account.billed_7d == null && stats.billed_7d != null + ? stats.billed_7d + : account.billed_7d; + const officialUsd = stats.official_usd ?? stats.official_usd_7d; + let nextOfficialUsd = account.official_usd; + let nextOfficialUsd7d = account.official_usd_7d; + if (officialUsd != null) { + nextOfficialUsd = officialUsd; + nextOfficialUsd7d = officialUsd; + } else if (stats.official_usage_synced) { + nextOfficialUsd = undefined; + nextOfficialUsd7d = undefined; + } + const officialSynced = + stats.official_usage_synced != null + ? stats.official_usage_synced + : account.official_usage_synced; + const usage5hDetail = + !account.usage_5h_detail && stats.usage_5h_detail + ? stats.usage_5h_detail + : account.usage_5h_detail; + const usage7dDetail = + !account.usage_7d_detail && stats.usage_7d_detail + ? stats.usage_7d_detail + : account.usage_7d_detail; + const usageTodayDetail = + !account.usage_today_detail && stats.usage_today_detail + ? stats.usage_today_detail + : account.usage_today_detail; + // page-stats 对本页每个账号都有条目；合并结果没变时保留行对象， + // 否则 1 秒 live 更新会让整页行级 memo 每秒全部失效。 + if ( + sameOptional(billed5h, account.billed_5h) && + sameOptional(billed7d, account.billed_7d) && + sameOptional(nextOfficialUsd, account.official_usd) && + sameOptional(nextOfficialUsd7d, account.official_usd_7d) && + sameOptional(officialSynced, account.official_usage_synced) && + sameOptional(usage5hDetail, account.usage_5h_detail) && + sameOptional(usage7dDetail, account.usage_7d_detail) && + sameOptional(usageTodayDetail, account.usage_today_detail) + ) { + return account; + } + const merged = { ...account }; + merged.billed_5h = billed5h; + merged.billed_7d = billed7d; + merged.official_usd = nextOfficialUsd; + merged.official_usd_7d = nextOfficialUsd7d; + merged.official_usage_synced = officialSynced; + merged.usage_5h_detail = usage5hDetail; + merged.usage_7d_detail = usage7dDetail; + merged.usage_today_detail = usageTodayDetail; + return merged; + }); + }, [data.accounts, accountPageStats]); const accounts = useMemo( () => allAccounts.filter((account) => !account.grok_api), [allAccounts], ); + // 只随“缺用量、需要补探”的账号集合变化，而不是随 1 秒 live 合并产生的 + // 数组身份变化；否则下面的补探 effect 每秒被清理重建，定时器永远不触发。 + const usageReloadIdsKey = useMemo( + () => + accounts + .filter(needsUsageReload) + .map((account) => account.id) + .sort((a, b) => a - b) + .join(","), + [accounts], + ); const accountPageIDsKey = useMemo( () => accounts.map((account) => account.id).join(","), [accounts], @@ -3194,6 +3245,8 @@ export default function Accounts() { // 测试连接后需要强制刷新用量的账号 id：即使其用量数据已存在（如已显示 100%）， // 也要在后台探针跑完后重新拉取，确保进度条更新为最新值。 const forceUsageReloadRef = useRef>(new Set()); + // 强制刷新是命令式信号，ref 本身不触发 effect；用计数器明确唤醒补探调度。 + const [usageReloadTick, setUsageReloadTick] = useState(0); useEffect(() => { persistAnalysisVisibility(showAnalysisCharts); @@ -3235,12 +3288,14 @@ export default function Accounts() { useEffect(() => { if (providerView !== "codex") return; - const missingUsageIds = accounts - .filter(needsUsageReload) - .map((account) => account.id); + const missingUsageIds = usageReloadIdsKey + ? usageReloadIdsKey.split(",").map(Number) + : []; // 测试连接后被标记强制刷新的账号也要参与重拉，即使其用量数据已存在。 // 仅保留仍在当前列表中的 id，避免泄漏。 - const accountIdSet = new Set(accounts.map((account) => account.id)); + const accountIdSet = new Set( + accountPageIDsKey ? accountPageIDsKey.split(",").map(Number) : [], + ); const forceIds = Array.from(forceUsageReloadRef.current).filter((id) => accountIdSet.has(id), ); @@ -3267,23 +3322,28 @@ export default function Accounts() { usageReloadAttemptsRef.current.get(id) ?? 0, ); } - for (const id of retryIds) { - usageReloadAttemptsRef.current.set( - id, - (usageReloadAttemptsRef.current.get(id) ?? 0) + 1, - ); - } - // 强制刷新的账号已安排本轮重拉，移除标记，避免无谓地反复重拉到上限。 - // 若重拉后数据仍未更新且该账号确实缺数据，会由 needsUsageReload 接管继续重试。 - for (const id of forceIds) { - forceUsageReloadRef.current.delete(id); - } - // 指数退避:2.5s/5s/10s/20s/40s/80s。固定 2.5s 会让列表加载后的十几秒里 // 反复整树重渲染;退避既压低churn,又把慢探针的覆盖窗口拉长到 ~2.5 分钟。 // 测试连接触发的强制刷新保持首档延迟,用户在等即时反馈。 const delay = forceIds.length > 0 ? 2500 : 2500 * 2 ** Math.min(minAttempt, 5); + const fire = () => { + // 只在真正发起重拉时消耗重试次数；被 live 更新取消的定时器不算一次。 + for (const id of retryIds) { + usageReloadAttemptsRef.current.set( + id, + (usageReloadAttemptsRef.current.get(id) ?? 0) + 1, + ); + } + // 强制刷新的账号已安排本轮重拉，移除标记，避免无谓地反复重拉到上限。 + // 若重拉后数据仍未更新且该账号确实缺数据，会由 needsUsageReload 接管继续重试。 + for (const id of forceIds) { + forceUsageReloadRef.current.delete(id); + } + void reloadSilently(); + // 集合不变时也要把下一档退避排上；否则一次尝试后重试链会断。 + setUsageReloadTick((tick) => tick + 1); + }; let visibilityHandler: (() => void) | null = null; const timer = window.setTimeout(() => { // 页面在后台时挂起本轮重拉,回到前台再补,不在不可见标签页里空转打请求。 @@ -3292,12 +3352,12 @@ export default function Accounts() { if (document.hidden || visibilityHandler === null) return; document.removeEventListener("visibilitychange", visibilityHandler); visibilityHandler = null; - void reloadSilently(); + fire(); }; document.addEventListener("visibilitychange", visibilityHandler); return; } - void reloadSilently(); + fire(); }, delay); return () => { @@ -3307,13 +3367,24 @@ export default function Accounts() { visibilityHandler = null; } }; - }, [accounts, reloadSilently, providerView]); + // 不依赖 accounts：1 秒 live 合并会不断产生新的数组身份，但只有 + // “需要补探的 id 集合 / 页面 id 集合 / 强制刷新信号”变化时才应重建定时器。 + }, [ + usageReloadIdsKey, + usageReloadTick, + accountPageIDsKey, + reloadSilently, + providerView, + ]); // stats_state 补拉:统计缓存分两层(请求数缓存→列表快照)且都是"先返回旧值 // 后台重建",从 stale 转 ready 需要连续两三次轮询。上面的静默刷新退避后间隔 // 最长 80s,不补拉的话「统计刷新中…」会常驻几分钟。这里在非 ready 时用 3s // 短间隔追到 ready 为止;带次数上限,防后端统计查询持续失败时退化成常驻轮询。 const statsStaleRetriesRef = useRef(0); + // 自增触发重试链，不依赖 data 对象身份：1 秒 live 合并会不断替换 data， + // 但 statsState 可能连续多次都是 stale，字符串不变时也必须继续补拉。 + const [statsRetryTick, setStatsRetryTick] = useState(0); useEffect(() => { if (loading) return undefined; if (data.statsState === "ready") { @@ -3325,12 +3396,11 @@ export default function Accounts() { const timer = window.setTimeout(() => { if (document.hidden) return; statsStaleRetriesRef.current += 1; + setStatsRetryTick((tick) => tick + 1); void reloadSilently(); }, 3000); return () => window.clearTimeout(timer); - // 依赖 data 本身(而非 data.statsState):连续两次都返回 stale 时字符串不变, - // 只有对象身份变化才能把补拉定时器重新拉起来。 - }, [loading, data, reloadSilently]); + }, [loading, data.statsState, data.disabledSorts, statsRetryTick, reloadSilently]); const accountSummary = { totalAccounts: data.summary?.total ?? data.total, @@ -9230,6 +9300,7 @@ export default function Accounts() { // 标记该账号强制刷新用量，配合后台探针确保进度条更新为最新值。 forceUsageReloadRef.current.add(testingAccount.id); usageReloadAttemptsRef.current.delete(testingAccount.id); + setUsageReloadTick((tick) => tick + 1); void reloadSilently(); }} onClose={() => setTestingAccount(null)} @@ -9245,6 +9316,7 @@ export default function Accounts() { forceUsageReloadRef.current.add(detectorAccount.id); usageReloadAttemptsRef.current.delete(detectorAccount.id); setDetectorAccount(null); + setUsageReloadTick((tick) => tick + 1); void reloadSilently(); }} /> diff --git a/frontend/src/types.ts b/frontend/src/types.ts index d62a497f4..05fabaf0c 100644 --- a/frontend/src/types.ts +++ b/frontend/src/types.ts @@ -628,6 +628,26 @@ export interface AccountLiveStateResponse { // 调度器当前实际执行的并发上限与配置值；旧后端不返回时保留列表里的值。 dynamic_concurrency_limit?: number base_concurrency_effective?: number + // 运行时用量投影：旧后端缺 key 时保留列表值；显式 null 表示窗口已失效。 + status?: AccountStatus + using_credits?: boolean + plan_type?: string + credits_valid?: boolean + credits_balance?: string | null + credits_has_credits?: boolean | null + credits_unlimited?: boolean | null + credits_overage_limit_reached?: boolean | null + credits_spend_control_reached?: boolean | null + credits_rate_limit_reached_type?: string | null + usage_percent_5h?: number | null + reset_5h_at?: ISODateString | null + codex_5h_usage_updated_at?: ISODateString | null + usage_percent_7d?: number | null + reset_7d_at?: ISODateString | null + codex_usage_updated_at?: ISODateString | null + // 最近一次被应用的 WS 配额帧；旧后端不返回，仅用于观测帧摄入。 + codex_usage_frame_at?: ISODateString + codex_usage_frame_kind?: string }> session_slot_buffer_enabled: boolean } diff --git a/main.go b/main.go index 998042a78..f5ea77342 100644 --- a/main.go +++ b/main.go @@ -40,6 +40,12 @@ func migrateOnlyEnabled() bool { return value == "1" || strings.EqualFold(value, "true") } +func configureCodexIdleUsageObserver(store *auth.Store) { + wsrelay.SetIdleMetadataObserver(func(account *auth.Account, eventType string, payload []byte) { + proxy.ApplyCodexIdleUsageFrame(store, account, eventType, payload) + }) +} + // main 加载配置、初始化存储与路由，并启动 Codex2API HTTP 服务。 func main() { log.SetFlags(log.LstdFlags | log.Lshortfile) @@ -445,6 +451,20 @@ func main() { // 注册 WebSocket 执行函数（避免 proxy ↔ wsrelay 循环依赖） proxy.WebsocketExecuteFunc = wsrelay.ExecuteRequestWebsocket + // 连接级（终态后补发）的用量帧也交给 proxy 更新内存快照；observer 必须非阻塞。 + configureCodexIdleUsageObserver(store) + // WS 逐帧用量落库节流；大号池可按数据库承压情况调大，0/非法值保持默认 30s。 + if raw := strings.TrimSpace(os.Getenv("CODEX_WS_USAGE_FRAME_PERSIST_SECONDS")); raw != "" { + if seconds, err := strconv.Atoi(raw); err == nil && seconds > 0 { + proxy.SetCodexUsageFramePersistInterval(time.Duration(seconds) * time.Second) + } + } + // 帧摄入的部署级逃生阀；false 时回退到 WHAM 探针与 HTTP 响应头。 + if raw := strings.TrimSpace(os.Getenv("CODEX_WS_USAGE_FRAME_INGEST")); raw != "" { + if enabled, err := strconv.ParseBool(raw); err == nil { + proxy.SetCodexUsageFrameIngestEnabled(enabled) + } + } // 注册 Agent Identity task 确保函数（proxy 无 Store 引用，启动时注入） proxy.EnsureCodexAgentIdentityTaskFunc = store.EnsureCodexAgentIdentityTask diff --git a/main_ws_usage_ingest_test.go b/main_ws_usage_ingest_test.go new file mode 100644 index 000000000..4bbee3492 --- /dev/null +++ b/main_ws_usage_ingest_test.go @@ -0,0 +1,311 @@ +package main + +import ( + "context" + "encoding/json" + "fmt" + "net/http" + "net/http/httptest" + "path/filepath" + "strings" + "sync" + "sync/atomic" + "testing" + "time" + + "github.com/codex2api/admin" + "github.com/codex2api/auth" + "github.com/codex2api/cache" + "github.com/codex2api/database" + "github.com/codex2api/proxy" + "github.com/codex2api/proxy/wsrelay" + "github.com/gin-gonic/gin" + "github.com/gorilla/websocket" +) + +// Exercise the production observer wiring, not a test replacement observer: +// Manager/read pump -> idle quota ingest -> scheduler -> authenticated live API. +// No background probes run, and no second business request rescues the last frame. +func TestMainPersistentWebSocketIngestsLastPostTerminalUsageFrame(t *testing.T) { + for _, test := range []struct { + name string + eventType string + contended bool + disabled bool + }{ + {name: "rate limits", eventType: "codex.rate_limits"}, + {name: "response metadata", eventType: "codex.response.metadata"}, + {name: "last frame waits for observation", eventType: "codex.rate_limits", contended: true}, + {name: "kill switch", eventType: "codex.rate_limits", disabled: true}, + } { + t.Run(test.name, func(t *testing.T) { + gin.SetMode(gin.TestMode) + db, err := database.New("sqlite", filepath.Join(t.TempDir(), "usage.db")) + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { _ = db.Close() }) + tc := cache.NewMemory(1) + t.Cleanup(func() { _ = tc.Close() }) + store := auth.NewStore(db, tc, &database.SystemSettings{MaxConcurrency: 1}) + t.Cleanup(store.Stop) + id, err := db.InsertATAccount(context.Background(), "synthetic-ws-usage", "synthetic-access-token", "") + if err != nil { + t.Fatal(err) + } + account := &auth.Account{DBID: id, AccessToken: "synthetic-access-token", PlanType: "plus", Status: auth.StatusReady, HealthTier: auth.HealthTierHealthy} + store.AddAccount(account) + store.SetGlobalAutoPauseThresholds(.95, 0) + if !account.IsAvailable() || store.AvailableCount() != 1 { + t.Fatal("synthetic account must be dispatchable before its last frame") + } + if selected := store.Next(); selected != account { + t.Fatal("scheduler did not dispatch the initial synthetic account") + } else { + store.Release(selected) + } + configureCodexIdleUsageObserver(store) + t.Cleanup(func() { wsrelay.SetIdleMetadataObserver(nil) }) + proxy.SetCodexUsageFrameIngestEnabled(!test.disabled) + t.Cleanup(func() { proxy.SetCodexUsageFrameIngestEnabled(true) }) + + const adminKey = "synthetic-ws-usage-admin" + adminHandler := admin.NewHandler(store, db, tc, proxy.NewRateLimiter(0), adminKey) + router := gin.New() + adminHandler.RegisterRoutes(router) + liveServer := httptest.NewServer(router) + t.Cleanup(liveServer.Close) + client := &http.Client{Timeout: time.Second} + readLive := func() auth.AccountLiveUsageSnapshot { + t.Helper() + request, err := http.NewRequest(http.MethodGet, fmt.Sprintf("%s/api/admin/accounts/live?ids=%d", liveServer.URL, id), nil) + if err != nil { + t.Fatal(err) + } + request.Header.Set("X-Admin-Key", adminKey) + response, err := client.Do(request) + if err != nil { + t.Fatal(err) + } + defer response.Body.Close() + var body struct { + Accounts map[int64]struct { + Status string `json:"status"` + Percent *float64 `json:"usage_percent_5h"` + Reset string `json:"reset_5h_at"` + Updated string `json:"codex_5h_usage_updated_at"` + FrameAt string `json:"codex_usage_frame_at"` + FrameKind string `json:"codex_usage_frame_kind"` + Plan string `json:"plan_type"` + CreditsValid bool `json:"credits_valid"` + CreditsBalance *string `json:"credits_balance"` + HasCredits bool `json:"credits_has_credits"` + } `json:"accounts"` + } + if response.StatusCode != http.StatusOK { + t.Fatalf("authenticated live status = %d", response.StatusCode) + } + if err := json.NewDecoder(response.Body).Decode(&body); err != nil { + t.Fatal(err) + } + item, ok := body.Accounts[id] + if !ok { + t.Fatal("live endpoint omitted the synthetic account") + } + result := auth.AccountLiveUsageSnapshot{Status: item.Status, UsageFrameLastKind: item.FrameKind, PlanType: item.Plan, CreditsValid: item.CreditsValid, Credits: auth.CreditBalanceSnapshot{Balance: item.CreditsBalance, HasCredits: item.HasCredits}} + if item.Percent != nil { + result.UsagePercent5h, result.UsagePercent5hValid = *item.Percent, true + } + for _, field := range []struct { + text string + target *time.Time + }{{item.Reset, &result.Reset5hAt}, {item.Updated, &result.UsageUpdatedAt5h}, {item.FrameAt, &result.UsageFrameLastAt}} { + if field.text != "" { + *field.target, err = time.Parse(time.RFC3339, field.text) + if err != nil { + t.Fatal(err) + } + } + } + return result + } + if before := readLive(); before.Status != "active" || before.UsagePercent5hValid || !before.UsageFrameLastAt.IsZero() { + t.Fatalf("initial live projection = %+v", before) + } + + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + t.Cleanup(cancel) + handshake := make(chan struct{}) + sendLastFrame := make(chan struct{}) + processedLastFrame := make(chan struct{}, 1) + serverErrors := make(chan error, 1) + reportServerError := func(err error) { + select { + case serverErrors <- err: + default: + } + } + var handshakeOnce sync.Once + var serverHandlers sync.WaitGroup + var requests atomic.Int32 + reset := time.Now().Add(2 * time.Hour).Truncate(time.Second) + payload := fmt.Sprintf(`{"type":"codex.rate_limits","plan_type":"pro","credits":{"has_credits":true,"unlimited":false,"balance":"3.5"},"rate_limits":{"primary":{"used_percent":96,"window_minutes":300,"reset_at":%d}}}`, reset.Unix()) + if test.eventType == "codex.response.metadata" { + payload = fmt.Sprintf(`{"type":"codex.response.metadata","headers":{"x-codex-plan-type":"pro","x-codex-credits-has-credits":"true","x-codex-credits-unlimited":"false","x-codex-credits-balance":"3.5","x-codex-primary-used-percent":"96","x-codex-primary-window-minutes":"300","x-codex-primary-reset-at":"%d"}}`, reset.Unix()) + } + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, request *http.Request) { + serverHandlers.Add(1) + defer serverHandlers.Done() + conn, err := (&websocket.Upgrader{}).Upgrade(w, request, nil) + if err != nil { + reportServerError(err) + return + } + defer conn.Close() + handshakeOnce.Do(func() { close(handshake) }) + if _, _, err := conn.ReadMessage(); err != nil { + reportServerError(err) + return + } + requests.Add(1) + if err := conn.WriteMessage(websocket.TextMessage, []byte(`{"type":"response.completed","response":{"id":"resp_synthetic","status":"completed"}}`)); err != nil { + reportServerError(err) + return + } + select { + case <-sendLastFrame: + case <-ctx.Done(): + return + } + conn.SetPongHandler(func(data string) error { + if data == "last-frame-consumed" { + select { + case processedLastFrame <- struct{}{}: + default: + } + } + return nil + }) + if err := conn.WriteMessage(websocket.TextMessage, []byte(payload)); err != nil { + reportServerError(err) + return + } + // A control-only barrier proves the sole reader passed the final data + // frame, even while another observer owns the usage lock. No later + // business frame or request is sent to trigger ingestion/recovery. + if err := conn.WriteControl(websocket.PingMessage, []byte("last-frame-consumed"), time.Now().Add(time.Second)); err != nil { + reportServerError(err) + return + } + for { + if _, _, err := conn.ReadMessage(); err != nil { + return // Manager shutdown closes the held-open connection. + } + requests.Add(1) + } + })) + manager := wsrelay.NewManager() + t.Cleanup(func() { + manager.Stop() + cancel() + upstream.Close() + serverHandlers.Wait() + }) + wc, pending, err := manager.AcquireConnection(ctx, account, "ws"+strings.TrimPrefix(upstream.URL, "http"), "synthetic-session", http.Header{}, "") + if err != nil { + t.Fatal(err) + } + waitMainUsageSignal(t, ctx, handshake, serverErrors) + if err := wc.WriteMessage(websocket.TextMessage, []byte(`{"type":"response.create","response":{"model":"synthetic-model"}}`)); err != nil { + t.Fatal(err) + } + terminal := make(chan error, 1) + terminalDone := make(chan struct{}) + go func() { + defer close(terminalDone) + _, data, err := wc.ReadMessage() + if err == nil && !strings.Contains(string(data), `"type":"response.completed"`) { + err = fmt.Errorf("unexpected terminal: %s", data) + } + terminal <- err + }() + t.Cleanup(func() { manager.Stop(); <-terminalDone }) + select { + case err := <-terminal: + if err != nil { + t.Fatal(err) + } + case <-ctx.Done(): + t.Fatal("terminal response was not delivered") + } + // Match the executor's terminal cleanup before the server sends its + // connection-level frame; the retained socket has no pending request. + session, ok := manager.GetSession(id, wc.URL, "synthetic-session", "") + if !ok { + t.Fatal("manager lost the persistent session") + } + session.RemovePendingRequest(pending.RequestID) + manager.ReleaseConnection(wc) + + if test.contended { + locked, release, finished := make(chan struct{}), make(chan struct{}), make(chan struct{}) + var once sync.Once + t.Cleanup(func() { once.Do(func() { close(release) }); <-finished }) + go func() { + defer close(finished) + account.ApplyUsageObservation(time.Now(), func() { close(locked); <-release }) + }() + waitMainUsageSignal(t, ctx, locked, serverErrors) + close(sendLastFrame) + waitMainUsageSignal(t, ctx, processedLastFrame, serverErrors) + once.Do(func() { close(release) }) + waitMainUsageSignal(t, ctx, finished, serverErrors) + } else { + close(sendLastFrame) + waitMainUsageSignal(t, ctx, processedLastFrame, serverErrors) + } + + if test.disabled { + if live := readLive(); live.Status != "active" || live.UsagePercent5hValid || !live.UsageFrameLastAt.IsZero() || !account.IsAvailable() { + t.Fatalf("disabled ingestion changed state: %+v", live) + } + } else { + ticker := time.NewTicker(time.Millisecond) + defer ticker.Stop() + for { + live := readLive() + if live.Status == "quota_paused" && live.UsagePercent5hValid && live.UsagePercent5h == 96 && live.Reset5hAt.Equal(reset) && !live.UsageUpdatedAt5h.IsZero() && !live.UsageFrameLastAt.IsZero() && live.UsageFrameLastKind == test.eventType && live.PlanType == "pro" && live.CreditsValid && live.Credits.HasCredits && live.Credits.Balance != nil && *live.Credits.Balance == "3.5" && !account.IsAvailable() && store.AvailableCount() == 0 { + break + } + select { + case <-ticker.C: + case err := <-serverErrors: + t.Fatal(err) + case <-ctx.Done(): + t.Fatalf("final idle frame never reached scheduler/live projection: %+v; available=%v", live, account.IsAvailable()) + } + } + } + if !test.disabled { + if selected := store.Next(); selected != nil { + store.Release(selected) + t.Fatal("scheduler dispatched an account paused by its final idle frame") + } + } + if requests.Load() != 1 || session.PendingCount() != 0 || !wc.IsConnected() || manager.ConnectionCount() != 1 { + t.Fatalf("last frame required another request or destroyed its persistent connection: requests=%d pending=%d connected=%v pooled=%d", requests.Load(), session.PendingCount(), wc.IsConnected(), manager.ConnectionCount()) + } + }) + } +} + +func waitMainUsageSignal(t *testing.T, ctx context.Context, signal <-chan struct{}, errors <-chan error) { + t.Helper() + select { + case <-signal: + case err := <-errors: + t.Fatal(err) + case <-ctx.Done(): + t.Fatal("persistent WS fixture synchronization timed out") + } +} diff --git a/proxy/codex_usage_frames.go b/proxy/codex_usage_frames.go new file mode 100644 index 000000000..a5bfe0f81 --- /dev/null +++ b/proxy/codex_usage_frames.go @@ -0,0 +1,396 @@ +package proxy + +import ( + "net/http" + "strconv" + "strings" + "sync/atomic" + "time" + + "github.com/codex2api/auth" + "github.com/tidwall/gjson" +) + +const ( + codexRateLimitsFrameType = "codex.rate_limits" + codexResponseMetadataFrameType = "codex.response.metadata" + responseMetadataFrameType = "response.metadata" + // Reset offsets are later converted to time.Duration in the shared parser. + maxCodexFrameResetSeconds = int64((1<<63 - 1) / time.Second) +) + +// codexUsageFramePersistIntervalNanos bounds how often a WS frame observation +// is written to the database. The in-memory snapshot is updated on every frame; +// persistence only matters for restart recovery, so a burst across a large pool +// must not turn every turn into a write. Operators can tune the floor with +// CODEX_WS_USAGE_FRAME_PERSIST_SECONDS (default 30; threshold crossings, 5pp +// deltas and 100% still force a write regardless of the interval). +var codexUsageFramePersistIntervalNanos = func() *atomic.Int64 { + value := &atomic.Int64{} + value.Store(int64(auth.UsageFramePersistInterval)) + return value +}() + +// codexUsageFrameIngestEnabled is the deployment-level kill switch for WS +// quota ingestion. When disabled, usage falls back to WHAM probes and HTTP +// response headers without rebuilding the image. +var codexUsageFrameIngestEnabled = func() *atomic.Bool { + value := &atomic.Bool{} + value.Store(true) + return value +}() + +// SetCodexUsageFrameIngestEnabled overrides the WS frame ingestion switch. +// Call before serving traffic. +func SetCodexUsageFrameIngestEnabled(enabled bool) { + codexUsageFrameIngestEnabled.Store(enabled) +} + +// SetCodexUsageFramePersistInterval overrides the frame persistence cadence. +// Values <= 0 restore the default. Call before serving traffic. +func SetCodexUsageFramePersistInterval(interval time.Duration) { + if interval <= 0 { + interval = auth.UsageFramePersistInterval + } + codexUsageFramePersistIntervalNanos.Store(int64(interval)) +} + +func codexUsageFramePersistInterval() time.Duration { + return time.Duration(codexUsageFramePersistIntervalNanos.Load()) +} + +// ApplyCodexUsageEventFrame ingests the preflight metadata frames that the +// upstream Codex WebSocket sends in place of response headers. On the WS +// transport there are no x-codex-* response headers per turn, so without this +// path the account usage snapshot only advances on WHAM probes. Frames are +// already on the wire (the gateway forwards them downstream), so normal ingest +// adds no upstream requests or request headers. Observation queue overflow can +// request the existing authoritative probe to recover dropped partial frames. +func ApplyCodexUsageEventFrame(store *auth.Store, account *auth.Account, eventType string, frame gjson.Result, observedAt time.Time) { + // Relay accounts can pass through their backend's quota frames. Those + // describe the relay's own upstream account, not this account's Codex + // windows, so they must not drive local usage or cooldowns. + if account == nil || account.IsRelayStyle() || !codexUsageFrameIngestEnabled.Load() { + return + } + switch strings.ToLower(strings.TrimSpace(eventType)) { + case codexRateLimitsFrameType: + applyCodexRateLimitsFrame(store, account, frame, observedAt) + case codexResponseMetadataFrameType, responseMetadataFrameType: + applyCodexResponseMetadataFrame(store, account, frame, observedAt) + } +} + +// applyCodexRateLimitsFrame maps the codex.rate_limits event onto the same +// synthetic headers the HTTP transport exposes. Only the default "codex" limit +// is applied to the account's 5h/7d snapshots; metered limits (for example a +// Spark or reserve bucket) must not overwrite the main windows. +func applyCodexRateLimitsFrame(store *auth.Store, account *auth.Account, frame gjson.Result, observedAt time.Time) { + headers := codexRateLimitsHeaders(frame, observedAt) + if len(headers) == 0 { + return + } + applyCodexUsageHeaders(store, account, headers, observedAt, codexRateLimitsFrameType) +} + +func codexRateLimitsHeaders(frame gjson.Result, observedAt time.Time) http.Header { + if name := firstCodexFrameString(frame, "metered_limit_name", "limit_id", "limit_name"); name != "" && !isDefaultCodexLimitName(name) { + return nil + } + headers := make(http.Header) + if plan, ok := codexFrameText(frame.Get("plan_type").String(), 32); ok { + headers.Set("x-codex-plan-type", plan) + } + setCodexFrameWindow(headers, "x-codex-primary-", frame.Get("rate_limits.primary"), observedAt) + setCodexFrameWindow(headers, "x-codex-secondary-", frame.Get("rate_limits.secondary"), observedAt) + // The frame carries a compact credits snapshot. Only adopt it when both + // booleans are present, matching the sparse-header contract: a missing flag + // must not overwrite a richer WHAM observation with false. + if credits := frame.Get("credits"); credits.IsObject() { + hasCredits := credits.Get("has_credits") + unlimited := credits.Get("unlimited") + if hasCredits.Exists() && unlimited.Exists() { + headers.Set("x-codex-credits-has-credits", hasCredits.String()) + headers.Set("x-codex-credits-unlimited", unlimited.String()) + if balance, ok := codexFrameText(credits.Get("balance").String(), 128); ok { + headers.Set("x-codex-credits-balance", balance) + } + } + } + return headers +} + +// codexFrameText validates free-text values copied from an upstream frame. +// These values can reach plain-text logs or persistence, so reject control +// characters that could forge a line. +func codexFrameText(value string, maxLen int) (string, bool) { + value = strings.TrimSpace(value) + if value == "" || len(value) > maxLen { + return "", false + } + for _, char := range value { + if char < 0x20 || char == 0x7f { + return "", false + } + } + return value, true +} + +// applyCodexResponseMetadataFrame adopts the quota headers embedded in a +// codex.response.metadata frame. The same frame also carries protocol routing +// data (x-codex-turn-state) and model catalog hints (x-models-etag); only the +// explicit quota-header allowlist below is copied, so routing or model metadata +// never reaches the usage sync or its persistence path. +func applyCodexResponseMetadataFrame(store *auth.Store, account *auth.Account, frame gjson.Result, observedAt time.Time) { + headers := codexResponseMetadataHeaders(frame, observedAt) + if len(headers) == 0 { + return + } + applyCodexUsageHeaders(store, account, headers, observedAt, codexResponseMetadataFrameType) +} + +func codexResponseMetadataHeaders(frame gjson.Result, observedAt time.Time) http.Header { + headers := make(http.Header) + // Absolute reset values win over a relative header carrying the same + // window, regardless of the JSON key order in the frame. + absoluteResets := make(map[string]string, 2) + frame.Get("headers").ForEach(func(key, value gjson.Result) bool { + name := strings.ToLower(strings.TrimSpace(key.String())) + if !isCodexUsageFrameHeader(name) { + return true + } + // The official client accepts header values that are arrays and reads the + // first element (json_value_as_string). Mirror that so a wrapped value + // does not silently drop the whole window. + if value.IsArray() { + items := value.Array() + if len(items) == 0 { + return true + } + value = items[0] + } + switch value.Type { + case gjson.String, gjson.Number, gjson.True, gjson.False: + default: + return true + } + // Convert the absolute reset form to the relative form using the frame's + // observation time. The header parser otherwise resolves it against the + // wall clock, which is identical in production but not in deterministic + // tests and can drift if a frame is replayed later. + if strings.HasSuffix(name, "-reset-at") { + epoch := codexFrameInt64(value) + if epoch > 0 { + seconds := epoch - observedAt.Unix() + if seconds > maxCodexFrameResetSeconds { + return true + } + if seconds < 0 { + seconds = 0 + } + prefix := strings.TrimSuffix(name, "reset-at") + absoluteResets[prefix+"reset-after-seconds"] = strconv.FormatInt(seconds, 10) + } + return true + } + if strings.HasSuffix(name, "-reset-after-seconds") { + seconds, ok := codexFrameInt64Value(value) + if !ok || seconds > maxCodexFrameResetSeconds { + return true + } + headers.Set(name, strconv.FormatInt(max(seconds, 0), 10)) + return true + } + raw, ok := codexFrameText(value.String(), 256) + if !ok { + return true + } + headers.Set(name, raw) + return true + }) + for name, value := range absoluteResets { + headers.Set(name, value) + } + return headers +} + +// isCodexUsageFrameHeader lists the headers consumed by the synthetic HTTP +// response. The metadata frame also carries routing and model fields in the +// same object; an explicit allowlist keeps those out of the usage path even if +// the sync later starts persisting more header fields. +func isCodexUsageFrameHeader(name string) bool { + switch name { + case "x-codex-plan-type", + "x-codex-primary-used-percent", "x-codex-primary-window-minutes", + "x-codex-primary-reset-at", "x-codex-primary-reset-after-seconds", + "x-codex-secondary-used-percent", "x-codex-secondary-window-minutes", + "x-codex-secondary-reset-at", "x-codex-secondary-reset-after-seconds", + "x-codex-credits-has-credits", "x-codex-credits-unlimited", + "x-codex-credits-balance", "x-codex-credits-overage-reached", + "x-codex-rate-limit-reached-type": + return true + } + return false +} + +func applyCodexUsageHeaders(store *auth.Store, account *auth.Account, headers http.Header, observedAt time.Time, kind string) { + if len(headers) == 0 { + return + } + if observedAt.IsZero() { + observedAt = time.Now() + } + // WS frames do not carry the per-response rate_limit_reached_type field, and + // the shared sync treats an absent value as "not reached" (correct for HTTP + // response headers). Preserve the last workspace hard stop instead of + // clearing it; only WHAM or an explicit metadata header may clear it. + resp := &http.Response{Header: headers} + // A per-turn frame may omit one window; never let it erase a valid snapshot. + // WHAM probes remain the authority for window presence/absence. The persist + // decision runs inside ApplyUsageObservation so a stale/out-of-order frame + // cannot advance the throttling state without being applied. + store.ApplyUsageFrameObservation(account, observedAt, func() { + if !codexUsageFrameIngestEnabled.Load() { + return + } + result := syncCodexUsageState(store, account, resp, observedAt, codexUsageSyncOptions{ + persistUsageSnapshot: true, + clearAbsent5h: false, + asyncPersist: true, + preserveAbsentHardStop: true, + observationLocked: true, + shouldPersistUsage: func(pct5h float64, has5h bool, pct7d float64, has7d bool) bool { + return account.ShouldPersistUsageFrame(observedAt, pct5h, has5h, pct7d, has7d, codexUsageFramePersistInterval()) + }, + }) + if result.Applied { + account.MarkUsageFrameObserved(kind, observedAt) + } + }) +} + +func setCodexFrameWindow(headers http.Header, prefix string, window gjson.Result, observedAt time.Time) { + if !window.IsObject() { + return + } + used := window.Get("used_percent") + minutes := window.Get("window_minutes") + if !used.Exists() || !minutes.Exists() { + return + } + headers.Set(prefix+"used-percent", used.String()) + headers.Set(prefix+"window-minutes", minutes.String()) + if reset := codexFrameResetSeconds(window, observedAt); reset != "" { + headers.Set(prefix+"reset-after-seconds", reset) + } +} + +// codexFrameResetSeconds accepts both relative and absolute reset shapes so the +// parser tolerates upstream schema drift. The official frame uses reset_at +// (absolute Unix seconds); older/observed fixtures used resets_in_seconds or +// reset_after_seconds. +func codexFrameResetSeconds(window gjson.Result, observedAt time.Time) string { + relativeSeen := false + for _, key := range []string{"resets_in_seconds", "reset_after_seconds"} { + if value := window.Get(key); value.Exists() { + if n, ok := codexFrameInt64Value(value); ok && n <= maxCodexFrameResetSeconds { + relativeSeen = true + // Zero/negative means "resets now", matching the absolute form + // where a past epoch is clamped to zero. A positive relative + // value is authoritative; a zero yields to a future absolute + // reset below so the real window end is not discarded. + if n > 0 { + return strconv.FormatInt(n, 10) + } + } + } + } + for _, key := range []string{"reset_at", "resets_at"} { + if value := window.Get(key); value.Exists() { + epoch, ok := codexFrameInt64Value(value) + if !ok || epoch <= 0 { + continue + } + seconds := epoch - observedAt.Unix() + if seconds > maxCodexFrameResetSeconds { + continue + } + if seconds < 0 { + seconds = 0 + } + return strconv.FormatInt(seconds, 10) + } + } + if relativeSeen { + return "0" + } + return "" +} + +func codexFrameInt64Value(value gjson.Result) (int64, bool) { + // Parse the original textual integer, not gjson's permissive Int conversion: + // fractional/exponent numbers and overflow must not become reset deadlines. + raw := value.String() + if value.Type == gjson.Number { + raw = value.Raw + } + n, err := strconv.ParseInt(strings.TrimSpace(raw), 10, 64) + if err != nil { + return 0, false + } + return n, true +} + +func codexFrameInt64(value gjson.Result) int64 { + n, _ := codexFrameInt64Value(value) + return n +} + +func firstCodexFrameString(frame gjson.Result, paths ...string) string { + for _, path := range paths { + if value := strings.TrimSpace(frame.Get(path).String()); value != "" { + return value + } + } + return "" +} + +var codexLimitNameSeparators = strings.NewReplacer("_", "-", " ", "-") + +func isDefaultCodexLimitName(name string) bool { + return codexLimitNameSeparators.Replace(strings.ToLower(strings.TrimSpace(name))) == "codex" +} + +// ApplyCodexIdleUsageFrame is the observer for connection-level metadata frames +// that arrive after a response terminal (the wsrelay read pump drops those +// instead of routing them to a request). Both frame paths share the same ordered +// runtime update and coalescing persistence. A busy observation lock queues the +// partial frames in order instead of blocking the read pump or discarding the +// final update. +func ApplyCodexIdleUsageFrame(store *auth.Store, account *auth.Account, eventType string, payload []byte) { + if account == nil || account.IsRelayStyle() || len(payload) == 0 || !codexUsageFrameIngestEnabled.Load() { + return + } + normalizedType := strings.ToLower(strings.TrimSpace(eventType)) + switch normalizedType { + case codexRateLimitsFrameType, codexResponseMetadataFrameType, responseMetadataFrameType: + default: + return + } + // Only quota-bearing frames are parsed. The read pump also forwards + // responsesapi.* telemetry frames here; skipping them before gjson keeps + // the sole raw reader cheap. + observedAt := time.Now() + frame := gjson.ParseBytes(payload) + var headers http.Header + switch normalizedType { + case codexRateLimitsFrameType: + headers = codexRateLimitsHeaders(frame, observedAt) + default: + headers = codexResponseMetadataHeaders(frame, observedAt) + } + if len(headers) == 0 { + return + } + applyCodexUsageHeaders(store, account, headers, observedAt, normalizedType) +} diff --git a/proxy/codex_usage_frames_integration_test.go b/proxy/codex_usage_frames_integration_test.go new file mode 100644 index 000000000..338d124b4 --- /dev/null +++ b/proxy/codex_usage_frames_integration_test.go @@ -0,0 +1,157 @@ +package proxy + +import ( + "bytes" + "net/http" + "net/http/httptest" + "strconv" + "testing" + "time" + + "github.com/codex2api/auth" + "github.com/codex2api/config" + "github.com/codex2api/database" + "github.com/gin-gonic/gin" +) + +func TestResponsesStreamAppliesCodexRateLimitsFrameToAccount(t *testing.T) { + gin.SetMode(gin.TestMode) + previousResin := resinCfg.Load() + t.Cleanup(func() { resinCfg.Store(previousResin) }) + + now := time.Now() + events := [][]string{{ + `{"type":"codex.rate_limits","plan_type":"plus","rate_limits":{` + + `"primary":{"used_percent":42,"window_minutes":300,"reset_at":` + strconv.FormatInt(now.Add(2*time.Hour).Unix(), 10) + `},` + + `"secondary":{"used_percent":8,"window_minutes":10080,"reset_at":` + strconv.FormatInt(now.Add(48*time.Hour).Unix(), 10) + `}}}`, + `{"type":"response.created","response":{"id":"resp_frame_1"}}`, + `{"type":"response.output_text.delta","delta":"hello"}`, + `{"type":"response.completed","response":{"id":"resp_frame_1","status":"completed","usage":{"input_tokens":1,"output_tokens":1,"total_tokens":2}}}`, + }} + upstream, _ := newAttemptSequenceSSEServer(t, events) + SetResinConfig(&ResinConfig{BaseURL: upstream.URL, PlatformName: "test"}) + + store := auth.NewStore(nil, nil, &database.SystemSettings{ + MaxConcurrency: 1, + MaxRetries: 0, + MaxRateLimitRetries: 0, + }) + t.Cleanup(store.Stop) + account := &auth.Account{ + DBID: 1, + AccessToken: "at-1", + AccountID: "acct-1", + Models: []string{"gpt-5.6-sol"}, + PlanType: "plus", + Status: auth.StatusReady, + HealthTier: auth.HealthTierHealthy, + } + store.AddAccount(account) + handler := NewHandler(store, nil, &config.Config{AllowAnonymousV1: true}, nil) + + recorder := httptest.NewRecorder() + ctx, _ := gin.CreateTestContext(recorder) + ctx.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", bytes.NewBufferString( + `{"model":"gpt-5.6-sol","input":"hello","stream":true}`, + )) + ctx.Request.Header.Set("Content-Type", "application/json") + handler.Responses(ctx) + + if recorder.Code != http.StatusOK { + t.Fatalf("status = %d, body=%s", recorder.Code, recorder.Body.String()) + } + pct5h, ok := account.GetUsagePercent5h() + if !ok || pct5h != 42 { + t.Fatalf("5h usage = %v valid=%v, want 42 from the rate_limits frame", pct5h, ok) + } + pct7d, ok := account.GetUsagePercent7d() + if !ok || pct7d != 8 { + t.Fatalf("7d usage = %v valid=%v, want 8 from the rate_limits frame", pct7d, ok) + } +} + +func TestResponsesNonStreamAppliesCodexRateLimitsFrameToAccount(t *testing.T) { + gin.SetMode(gin.TestMode) + previousResin := resinCfg.Load() + t.Cleanup(func() { resinCfg.Store(previousResin) }) + + now := time.Now() + events := [][]string{{ + `{"type":"codex.rate_limits","plan_type":"plus","rate_limits":{` + + `"primary":{"used_percent":42,"window_minutes":300,"reset_at":` + strconv.FormatInt(now.Add(2*time.Hour).Unix(), 10) + `}}}`, + `{"type":"response.created","response":{"id":"resp_nonstream_1"}}`, + `{"type":"response.output_text.delta","delta":"hello"}`, + `{"type":"response.completed","response":{"id":"resp_nonstream_1","status":"completed","usage":{"input_tokens":1,"output_tokens":1,"total_tokens":2}}}`, + }} + upstream, _ := newAttemptSequenceSSEServer(t, events) + SetResinConfig(&ResinConfig{BaseURL: upstream.URL, PlatformName: "test"}) + + store := auth.NewStore(nil, nil, &database.SystemSettings{MaxConcurrency: 1}) + t.Cleanup(store.Stop) + account := &auth.Account{ + DBID: 1, AccessToken: "at-1", AccountID: "acct-1", + Models: []string{"gpt-5.6-sol"}, PlanType: "plus", + Status: auth.StatusReady, HealthTier: auth.HealthTierHealthy, + } + store.AddAccount(account) + handler := NewHandler(store, nil, &config.Config{AllowAnonymousV1: true}, nil) + + recorder := httptest.NewRecorder() + ctx, _ := gin.CreateTestContext(recorder) + ctx.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", bytes.NewBufferString( + `{"model":"gpt-5.6-sol","input":"hello","stream":false}`, + )) + ctx.Request.Header.Set("Content-Type", "application/json") + handler.Responses(ctx) + + if recorder.Code != http.StatusOK { + t.Fatalf("status = %d, body=%s", recorder.Code, recorder.Body.String()) + } + if pct, ok := account.GetUsagePercent5h(); !ok || pct != 42 { + t.Fatalf("5h usage = %v valid=%v, want 42 from the non-streaming frame", pct, ok) + } +} + +func TestChatCompletionsStreamAppliesCodexRateLimitsFrameToAccount(t *testing.T) { + handler, _ := newChatStreamTerminalTestHandler(t, []string{ + `{"type":"codex.rate_limits","plan_type":"plus","rate_limits":{"primary":{"used_percent":42,"window_minutes":300}}}`, + `{"type":"response.created","response":{"id":"resp_chat_frame"}}`, + `{"type":"response.output_text.delta","delta":"hello"}`, + `{"type":"response.completed","response":{"id":"resp_chat_frame","status":"completed","usage":{"input_tokens":1,"output_tokens":1,"total_tokens":2}}}`, + }) + + recorder := invokeChatCompletionsStream(t, handler) + if recorder.Code != http.StatusOK { + t.Fatalf("status = %d, body=%s", recorder.Code, recorder.Body.String()) + } + account := handler.store.FindByID(1) + if account == nil { + t.Fatal("test account missing") + } + if pct, ok := account.GetUsagePercent5h(); !ok || pct != 42 { + t.Fatalf("5h usage = %v valid=%v, want 42 from the chat frame", pct, ok) + } +} + +func TestMessagesStreamAppliesCodexRateLimitsFrameToAccount(t *testing.T) { + handler, _ := newAnthropicStreamFailureTestHandler(t, func(_ int32, w http.ResponseWriter) { + writeCodexSSE(w, + `{"type":"codex.rate_limits","plan_type":"plus","rate_limits":{"primary":{"used_percent":42,"window_minutes":300}}}`, + `{"type":"response.created","response":{"id":"resp_messages_frame"}}`, + `{"type":"response.output_text.delta","delta":"hello"}`, + `{"type":"response.completed","response":{"id":"resp_messages_frame","status":"completed","usage":{"input_tokens":1,"output_tokens":1,"total_tokens":2}}}`, + ) + }) + + recorder := invokeAnthropicMessagesStream(t, handler) + if recorder.Code != http.StatusOK { + t.Fatalf("status = %d, body=%s", recorder.Code, recorder.Body.String()) + } + account := handler.store.FindByID(1) + if account == nil { + t.Fatal("test account missing") + } + if pct, ok := account.GetUsagePercent5h(); !ok || pct != 42 { + t.Fatalf("5h usage = %v valid=%v, want 42 from the messages frame", pct, ok) + } +} diff --git a/proxy/codex_usage_frames_test.go b/proxy/codex_usage_frames_test.go new file mode 100644 index 000000000..05e2c0241 --- /dev/null +++ b/proxy/codex_usage_frames_test.go @@ -0,0 +1,1002 @@ +package proxy + +import ( + "context" + "database/sql" + "fmt" + "net/http" + "path/filepath" + "strconv" + "testing" + "time" + + "github.com/codex2api/auth" + "github.com/codex2api/database" + "github.com/tidwall/gjson" +) + +func newUsageFrameTestAccount() *auth.Account { + return &auth.Account{ + DBID: 1, + PlanType: "plus", + Status: auth.StatusReady, + HealthTier: auth.HealthTierHealthy, + } +} + +func TestApplyCodexRateLimitsFrameUpdatesWindows(t *testing.T) { + acc := newUsageFrameTestAccount() + observedAt := time.Unix(1_700_000_000, 0) + frame := gjson.Parse(`{ + "type":"codex.rate_limits", + "plan_type":"plus", + "rate_limits":{ + "primary":{"used_percent":42,"window_minutes":300,"reset_at":1700000900}, + "secondary":{"used_percent":8,"window_minutes":10080,"reset_at":1700604800} + } + }`) + + ApplyCodexUsageEventFrame(nil, acc, "codex.rate_limits", frame, observedAt) + + pct5h, ok := acc.GetUsagePercent5h() + if !ok || pct5h != 42 { + t.Fatalf("5h usage = %v, valid=%v, want 42", pct5h, ok) + } + if want := observedAt.Add(15 * time.Minute); !acc.GetReset5hAt().Equal(want) { + t.Fatalf("Reset5hAt = %v, want %v", acc.GetReset5hAt(), want) + } + pct7d, ok := acc.GetUsagePercent7d() + if !ok || pct7d != 8 { + t.Fatalf("7d usage = %v, valid=%v, want 8", pct7d, ok) + } +} + +func TestApplyCodexRateLimitsFrameKeepsAbsent5hWindow(t *testing.T) { + acc := newUsageFrameTestAccount() + observedAt := time.Unix(1_700_000_000, 0) + acc.SetUsageSnapshot5hAt(77, observedAt.Add(time.Hour), observedAt) + frame := gjson.Parse(`{ + "type":"codex.rate_limits", + "rate_limits":{ + "primary":null, + "secondary":{"used_percent":12,"window_minutes":10080,"reset_at":1700604800} + } + }`) + + ApplyCodexUsageEventFrame(nil, acc, "codex.rate_limits", frame, observedAt.Add(time.Minute)) + + pct5h, ok := acc.GetUsagePercent5h() + if !ok || pct5h != 77 { + t.Fatalf("partial frame cleared 5h snapshot: usage=%v valid=%v", pct5h, ok) + } + pct7d, ok := acc.GetUsagePercent7d() + if !ok || pct7d != 12 { + t.Fatalf("partial frame dropped its present 7d window: usage=%v valid=%v", pct7d, ok) + } +} + +func TestApplyCodexRateLimitsFrameIgnoresMeteredLimits(t *testing.T) { + acc := newUsageFrameTestAccount() + observedAt := time.Unix(1_700_000_000, 0) + acc.SetUsageSnapshot5hAt(55, observedAt.Add(time.Hour), observedAt) + frame := gjson.Parse(`{ + "type":"codex.rate_limits", + "metered_limit_name":"codex_spark", + "rate_limits":{ + "primary":{"used_percent":99,"window_minutes":300,"reset_at":1700000900} + } + }`) + + ApplyCodexUsageEventFrame(nil, acc, "codex.rate_limits", frame, observedAt.Add(time.Minute)) + + pct5h, ok := acc.GetUsagePercent5h() + if !ok || pct5h != 55 { + t.Fatalf("metered limit frame overwrote main 5h snapshot: usage=%v valid=%v", pct5h, ok) + } +} + +func TestApplyCodexResponseMetadataFrameAdoptsUsageHeaders(t *testing.T) { + acc := newUsageFrameTestAccount() + observedAt := time.Unix(1_700_000_000, 0) + frame := gjson.Parse(`{ + "type":"codex.response.metadata", + "headers":{ + "x-codex-primary-used-percent":"33", + "x-codex-primary-window-minutes":"300", + "x-codex-primary-reset-at":"1700000900", + "x-codex-turn-state":"ts-must-not-be-used", + "x-codex-safety-buffering-enabled":"true", + "x-models-etag":"etag-must-not-be-used", + "authorization":"Bearer must-not-be-used", + "set-cookie":"must-not-be-used" + } + }`) + + headers := codexResponseMetadataHeaders(frame, observedAt) + if got := headers.Get("x-codex-primary-used-percent"); got != "33" { + t.Fatalf("quota header = %q, want 33", got) + } + for _, name := range []string{ + "x-codex-turn-state", + "x-codex-safety-buffering-enabled", + "x-models-etag", + "authorization", + "set-cookie", + } { + if got := headers.Get(name); got != "" { + t.Fatalf("%s leaked into synthetic usage headers: %q", name, got) + } + } + + ApplyCodexUsageEventFrame(nil, acc, "codex.response.metadata", frame, observedAt) + + pct5h, ok := acc.GetUsagePercent5h() + if !ok || pct5h != 33 { + t.Fatalf("5h usage = %v valid=%v, want 33", pct5h, ok) + } + if want := observedAt.Add(15 * time.Minute); !acc.GetReset5hAt().Equal(want) { + t.Fatalf("Reset5hAt = %v, want %v", acc.GetReset5hAt(), want) + } +} + +func TestSyncCodexUsageStateMarks100PercentOnlyWhenPersisting(t *testing.T) { + store := auth.NewStore(nil, nil, &database.SystemSettings{MaxConcurrency: 1}) + t.Cleanup(store.Stop) + observedAt := time.Now() + + newFullResponse := func() *http.Response { + resp := &http.Response{Header: http.Header{}} + resp.Header.Set("x-codex-primary-used-percent", "100") + resp.Header.Set("x-codex-primary-window-minutes", "300") + resp.Header.Set("x-codex-primary-reset-after-seconds", "3600") + resp.Header.Set("x-codex-secondary-used-percent", "100") + resp.Header.Set("x-codex-secondary-window-minutes", "10080") + resp.Header.Set("x-codex-secondary-reset-after-seconds", "86400") + return resp + } + + throttled := newUsageFrameTestAccount() + store.AddAccount(throttled) + result := syncCodexUsageState(store, throttled, newFullResponse(), observedAt, codexUsageSyncOptions{ + persistUsageSnapshot: true, + shouldPersistUsage: func(float64, bool, float64, bool) bool { return false }, + }) + if result.Usage7dRateLimited || result.Premium5hRateLimited { + t.Fatalf("throttled frame marked cooldowns: 7d=%v 5h=%v", result.Usage7dRateLimited, result.Premium5hRateLimited) + } + if throttled.Status == auth.StatusCooldown { + t.Fatal("throttled 100% frame put the account into cooldown") + } + + persisted := newUsageFrameTestAccount() + persisted.DBID = 2 + store.AddAccount(persisted) + result = syncCodexUsageState(store, persisted, newFullResponse(), observedAt, codexUsageSyncOptions{ + persistUsageSnapshot: true, + shouldPersistUsage: func(float64, bool, float64, bool) bool { return true }, + }) + if !result.Usage7dRateLimited || !result.Premium5hRateLimited { + t.Fatalf("persisted 100%% marks = 7d:%v 5h:%v, want both", result.Usage7dRateLimited, result.Premium5hRateLimited) + } +} + +func TestCodexUsageFramesPreserveLongerResponsesCooldown(t *testing.T) { + for _, idle := range []bool{false, true} { + for _, kind := range []string{"codex.rate_limits", "codex.response.metadata"} { + t.Run(fmt.Sprintf("idle=%v/%s", idle, kind), func(t *testing.T) { + store := auth.NewStore(nil, nil, &database.SystemSettings{MaxConcurrency: 1}) + t.Cleanup(store.Stop) + acc := newUsageFrameTestAccount() + until := time.Now().Add(6 * time.Hour) + acc.Status, acc.CooldownUtil, acc.CooldownReason = auth.StatusCooldown, until, auth.ResponsesRateLimitedCooldownReason + acc.FailureStreak = 4 + store.AddAccount(acc) + payload := `{"type":"codex.rate_limits","rate_limits":{"primary":{"used_percent":100,"window_minutes":300,"reset_after_seconds":3600},"secondary":{"used_percent":100,"window_minutes":10080,"reset_after_seconds":7200}}}` + if kind == "codex.response.metadata" { + payload = `{"type":"codex.response.metadata","headers":{"x-codex-primary-used-percent":"100","x-codex-primary-window-minutes":"300","x-codex-primary-reset-after-seconds":"3600","x-codex-secondary-used-percent":"100","x-codex-secondary-window-minutes":"10080","x-codex-secondary-reset-after-seconds":"7200"}}` + } + if idle { + ApplyCodexIdleUsageFrame(store, acc, kind, []byte(payload)) + } else { + ApplyCodexUsageEventFrame(store, acc, kind, gjson.Parse(payload), time.Now()) + } + if acc.Status != auth.StatusCooldown || !acc.CooldownUtil.Equal(until) || acc.GetCooldownReason() != auth.ResponsesRateLimitedCooldownReason || acc.FailureStreak != 4 { + t.Fatalf("frame weakened the Responses cooldown: status=%v until=%v reason=%q streak=%d", acc.Status, acc.CooldownUtil, acc.GetCooldownReason(), acc.FailureStreak) + } + if pct, ok := acc.GetUsagePercent5h(); !ok || pct != 100 { + t.Fatalf("5h usage = %v valid=%v, want 100", pct, ok) + } + if pct, ok := acc.GetUsagePercent7d(); !ok || pct != 100 { + t.Fatalf("7d usage = %v valid=%v, want 100", pct, ok) + } + }) + } + } +} + +func TestApplyCodexRateLimitsFrameProPrimaryWeekly(t *testing.T) { + acc := newUsageFrameTestAccount() + acc.PlanType = "pro" + now := time.Unix(1_700_000_000, 0) + acc.SetUsageSnapshot5hAt(12, now.Add(time.Hour), now.Add(-time.Minute)) + // Shape captured from a live CLIProxyAPI fixture: a Pro plan can carry the + // weekly window in primary and leave secondary null. + frame := gjson.Parse(`{ + "type":"codex.rate_limits", + "plan_type":"pro", + "rate_limits":{ + "allowed":true, + "limit_reached":false, + "primary":{"used_percent":81,"window_minutes":10080,"reset_after_seconds":137160,"reset_at":1700137160}, + "secondary":null + }, + "credits":{"has_credits":false,"unlimited":false,"balance":"0"} + }`) + + ApplyCodexUsageEventFrame(nil, acc, "codex.rate_limits", frame, now) + + if pct, ok := acc.GetUsagePercent7d(); !ok || pct != 81 { + t.Fatalf("7d usage = %v valid=%v, want 81", pct, ok) + } + if pct, ok := acc.GetUsagePercent5h(); !ok || pct != 12 { + t.Fatalf("primary weekly frame cleared 5h snapshot: usage=%v valid=%v", pct, ok) + } + if got, want := acc.GetReset7dAt(), now.Add(137160*time.Second); !got.Equal(want) { + t.Fatalf("Reset7dAt = %v, want %v", got, want) + } +} + +func TestParseCodexWindowUsageRejectsOutOfRangePercent(t *testing.T) { + for _, used := range []string{"101", "-1"} { + if got := parseCodexWindowUsage(used, "300", "60"); got.valid { + t.Fatalf("used_percent %s accepted: %+v", used, got) + } + } + if got := parseCodexWindowUsage("100", "300", "60"); !got.valid { + t.Fatal("used_percent 100 rejected") + } +} + +func TestCodexUsageFrameIngestKillSwitch(t *testing.T) { + SetCodexUsageFrameIngestEnabled(false) + t.Cleanup(func() { SetCodexUsageFrameIngestEnabled(true) }) + + acc := newUsageFrameTestAccount() + frame := gjson.Parse(`{"rate_limits":{"primary":{"used_percent":42,"window_minutes":300}}}`) + ApplyCodexUsageEventFrame(nil, acc, "codex.rate_limits", frame, time.Now()) + ApplyCodexIdleUsageFrame(nil, acc, "codex.rate_limits", []byte(`{"rate_limits":{"primary":{"used_percent":42,"window_minutes":300}}}`)) + + if pct, ok := acc.GetUsagePercent5h(); ok { + t.Fatalf("disabled ingestion updated usage to %v", pct) + } +} + +func TestApplyCodexUsageEventFrameRecordsFrameFreshness(t *testing.T) { + acc := newUsageFrameTestAccount() + observedAt := time.Unix(1_700_000_000, 0) + frame := gjson.Parse(`{"rate_limits":{"primary":{"used_percent":42,"window_minutes":300}}}`) + + ApplyCodexUsageEventFrame(nil, acc, "codex.rate_limits", frame, observedAt) + + snapshot := acc.GetAccountLiveUsageSnapshot() + if !snapshot.UsageFrameLastAt.Equal(observedAt) { + t.Fatalf("frame time = %v, want %v", snapshot.UsageFrameLastAt, observedAt) + } + if snapshot.UsageFrameLastKind != "codex.rate_limits" { + t.Fatalf("frame kind = %q, want codex.rate_limits", snapshot.UsageFrameLastKind) + } +} + +func TestApplyCodexUsageEventFrameIgnoresUnrelatedEvents(t *testing.T) { + acc := newUsageFrameTestAccount() + observedAt := time.Unix(1_700_000_000, 0) + acc.SetUsageSnapshot5hAt(21, observedAt.Add(time.Hour), observedAt) + frame := gjson.Parse(`{"type":"response.created","response":{"id":"resp-1"}}`) + + ApplyCodexUsageEventFrame(nil, acc, "response.created", frame, observedAt) + + pct5h, ok := acc.GetUsagePercent5h() + if !ok || pct5h != 21 { + t.Fatalf("unrelated event changed usage: usage=%v valid=%v", pct5h, ok) + } +} + +func TestApplyCodexRateLimitsFrameFreePlanHasOnlyLongWindow(t *testing.T) { + acc := newUsageFrameTestAccount() + acc.PlanType = "free" + observedAt := time.Unix(1_700_000_000, 0) + frame := gjson.Parse(`{ + "type":"codex.rate_limits", + "plan_type":"free", + "rate_limits":{ + "primary":{"used_percent":64,"window_minutes":10080,"reset_at":1700604800}, + "secondary":null + } + }`) + + ApplyCodexUsageEventFrame(nil, acc, "codex.rate_limits", frame, observedAt) + + if _, ok := acc.GetUsagePercent5h(); ok { + t.Fatal("free plan frame must not fabricate a 5h window") + } + pct7d, ok := acc.GetUsagePercent7d() + if !ok || pct7d != 64 { + t.Fatalf("7d usage = %v valid=%v, want 64", pct7d, ok) + } +} + +func TestApplyCodexRateLimitsFrame5hOnlyKeeps7dSnapshot(t *testing.T) { + acc := newUsageFrameTestAccount() + observedAt := time.Unix(1_700_000_000, 0) + acc.SetUsagePercent7d(31) + frame := gjson.Parse(`{ + "type":"codex.rate_limits", + "rate_limits":{ + "primary":{"used_percent":56,"window_minutes":300,"reset_at":1700000900}, + "secondary":null + } + }`) + + ApplyCodexUsageEventFrame(nil, acc, "codex.rate_limits", frame, observedAt) + + pct5h, ok := acc.GetUsagePercent5h() + if !ok || pct5h != 56 { + t.Fatalf("5h usage = %v valid=%v, want 56", pct5h, ok) + } + pct7d, ok := acc.GetUsagePercent7d() + if !ok || pct7d != 31 { + t.Fatalf("partial 5h frame cleared 7d snapshot: usage=%v valid=%v", pct7d, ok) + } +} + +func TestApplyCodexRateLimitsFrameTeamMonthlyWindow(t *testing.T) { + acc := newUsageFrameTestAccount() + acc.PlanType = "team" + observedAt := time.Unix(1_700_000_000, 0) + frame := gjson.Parse(`{ + "type":"codex.rate_limits", + "rate_limits":{ + "primary":{"used_percent":44,"window_minutes":43200,"reset_at":1702592000}, + "secondary":null + } + }`) + + ApplyCodexUsageEventFrame(nil, acc, "codex.rate_limits", frame, observedAt) + + if _, ok := acc.GetUsagePercent5h(); ok { + t.Fatal("monthly window must not be classified as 5h") + } + pct7d, ok := acc.GetUsagePercent7d() + if !ok || pct7d != 44 { + t.Fatalf("7d usage = %v valid=%v, want 44", pct7d, ok) + } + if got := acc.GetWindow7dSeconds(); got != 43200*60 { + t.Fatalf("Window7dSeconds = %d, want %d", got, 43200*60) + } +} + +func TestApplyCodexRateLimitsFramePreservesKnownReset(t *testing.T) { + acc := newUsageFrameTestAccount() + observedAt := time.Unix(1_700_000_000, 0) + resetAt := observedAt.Add(2 * time.Hour) + acc.SetUsageSnapshot5hAt(20, resetAt, observedAt) + frame := gjson.Parse(`{ + "type":"codex.rate_limits", + "rate_limits":{ + "primary":{"used_percent":70,"window_minutes":300}, + "secondary":null + } + }`) + + ApplyCodexUsageEventFrame(nil, acc, "codex.rate_limits", frame, observedAt.Add(time.Minute)) + + pct5h, ok := acc.GetUsagePercent5h() + if !ok || pct5h != 70 { + t.Fatalf("5h usage = %v valid=%v, want 70", pct5h, ok) + } + if !acc.GetReset5hAt().Equal(resetAt) { + t.Fatalf("Reset5hAt = %v, want preserved %v", acc.GetReset5hAt(), resetAt) + } +} + +func TestApplyCodexIdleUsageFrameWithoutStoreUpdatesRuntime(t *testing.T) { + acc := newUsageFrameTestAccount() + now := time.Now() + acc.SetUsageSnapshot5hAt(20, now.Add(time.Hour), now) + acc.SetUsagePercent7d(10) + + ApplyCodexIdleUsageFrame(nil, acc, "codex.rate_limits", []byte(`{ + "type":"codex.rate_limits", + "rate_limits":{ + "primary":{"used_percent":75,"window_minutes":300,"reset_at":1700000900}, + "secondary":null + } + }`)) + + pct5h, ok := acc.GetUsagePercent5h() + if !ok || pct5h != 75 { + t.Fatalf("idle frame 5h usage = %v valid=%v, want 75", pct5h, ok) + } + pct7d, ok := acc.GetUsagePercent7d() + if !ok || pct7d != 10 { + t.Fatalf("idle partial frame cleared 7d snapshot: usage=%v valid=%v", pct7d, ok) + } +} + +func TestParseCodexUsageHeaderObservationAcceptsResetAt(t *testing.T) { + now := time.Now() + resp := &http.Response{Header: http.Header{}} + resp.Header.Set("x-codex-primary-used-percent", "33") + resp.Header.Set("x-codex-primary-window-minutes", "300") + resp.Header.Set("x-codex-primary-reset-at", strconv.FormatInt(now.Add(15*time.Minute).Unix(), 10)) + + observation := parseCodexUsageHeaderObservation(resp) + if !observation.w5h.valid { + t.Fatal("reset-at header did not produce a valid 5h window") + } + if observation.w5h.resetSec < 14*60 || observation.w5h.resetSec > 16*60 { + t.Fatalf("resetSec = %v, want about 900", observation.w5h.resetSec) + } +} + +func TestApplyCodexRateLimitsFramePreservesWorkspaceHardStop(t *testing.T) { + acc := newUsageFrameTestAccount() + acc.SetRateLimitReachedType("workspace_member_credits_depleted") + frame := gjson.Parse(`{ + "type":"codex.rate_limits", + "plan_type":"plus", + "rate_limits":{ + "primary":{"used_percent":42,"window_minutes":300,"reset_at":1700000900}, + "secondary":{"used_percent":8,"window_minutes":10080,"reset_at":1700604800} + }, + "credits":{"has_credits":true,"unlimited":false,"balance":"10"} + }`) + + ApplyCodexUsageEventFrame(nil, acc, "codex.rate_limits", frame, time.Unix(1_700_000_000, 0)) + + if got := acc.GetRateLimitReachedType(); got != "workspace_member_credits_depleted" { + t.Fatalf("frame cleared workspace hard stop: %q", got) + } +} + +func TestApplyCodexMetadataFramePreservesWorkspaceHardStop(t *testing.T) { + acc := newUsageFrameTestAccount() + acc.SetRateLimitReachedType("workspace_member_credits_depleted") + frame := gjson.Parse(`{ + "type":"codex.response.metadata", + "headers":{"x-codex-primary-used-percent":"33","x-codex-primary-window-minutes":"300"} + }`) + + ApplyCodexUsageEventFrame(nil, acc, "codex.response.metadata", frame, time.Unix(1_700_000_000, 0)) + + if got := acc.GetRateLimitReachedType(); got != "workspace_member_credits_depleted" { + t.Fatalf("metadata frame cleared workspace hard stop: %q", got) + } +} + +func TestCodexResponseMetadataHeadersUnwrapsArrayValues(t *testing.T) { + observedAt := time.Unix(1_700_000_000, 0) + frame := gjson.Parse(`{"headers":{ + "x-codex-primary-used-percent":["42"], + "x-codex-primary-window-minutes":["300"], + "x-codex-primary-reset-at":[1700000900], + "x-codex-credits-has-credits":["true"], + "x-codex-turn-state":["ts-must-not-be-used"] + }}`) + + headers := codexResponseMetadataHeaders(frame, observedAt) + if got := headers.Get("x-codex-primary-used-percent"); got != "42" { + t.Fatalf("array header = %q, want 42", got) + } + if got := headers.Get("x-codex-primary-reset-after-seconds"); got != "900" { + t.Fatalf("array reset = %q, want 900", got) + } + if got := headers.Get("x-codex-credits-has-credits"); got != "true" { + t.Fatalf("array credits = %q, want true", got) + } + if got := headers.Get("x-codex-turn-state"); got != "" { + t.Fatalf("turn state leaked into quota headers: %q", got) + } +} + +func TestCodexFrameRejectsControlCharacters(t *testing.T) { + observedAt := time.Unix(1_700_000_000, 0) + frame := gjson.Parse(`{ + "plan_type":"plus\r\nX-Injected: 1", + "rate_limits":{"primary":{"used_percent":42,"window_minutes":300}}, + "credits":{"has_credits":true,"unlimited":false,"balance":"10\n20"} + }`) + + headers := codexRateLimitsHeaders(frame, observedAt) + if got := headers.Get("x-codex-plan-type"); got != "" { + t.Fatalf("control characters in plan type were accepted: %q", got) + } + if got := headers.Get("x-codex-credits-balance"); got != "" { + t.Fatalf("control characters in balance were accepted: %q", got) + } + if got := headers.Get("x-codex-primary-used-percent"); got != "42" { + t.Fatalf("valid window was discarded: %q", got) + } + + metadata := codexResponseMetadataHeaders(gjson.Parse(`{"headers":{ + "x-codex-plan-type":"plus\r\nX-Injected: 1" + }}`), observedAt) + if got := metadata.Get("x-codex-plan-type"); got != "" { + t.Fatalf("control characters in metadata header were accepted: %q", got) + } +} + +func TestCodexRateLimitsFrameTreatsZeroRelativeResetAsNow(t *testing.T) { + observedAt := time.Unix(1_700_000_000, 0) + frame := gjson.Parse(`{"rate_limits":{"primary":{ + "used_percent":50,"window_minutes":300,"resets_in_seconds":0 + }}}`) + + headers := codexRateLimitsHeaders(frame, observedAt) + if got := headers.Get("x-codex-primary-reset-after-seconds"); got != "0" { + t.Fatalf("relative zero reset = %q, want 0", got) + } +} + +func TestCodexFrameResetRejectsNonIntegerNumbers(t *testing.T) { + at := time.Unix(1_700_000_000, 0) + for _, raw := range []string{"1.5", "1e3", "9223372036854775808", "-9223372036854775809"} { + t.Run(raw, func(t *testing.T) { + for _, field := range []string{"reset_after_seconds", "reset_at"} { + window := gjson.Parse(`{"` + field + `":` + raw + `}`) + if got := codexFrameResetSeconds(window, at); got != "" { + t.Errorf("%s=%s accepted as reset %q", field, raw, got) + } + } + for _, field := range []string{"reset-at", "reset-after-seconds"} { + metadata := codexResponseMetadataHeaders(gjson.Parse(`{"headers":{"x-codex-primary-`+field+`":`+raw+`}}`), at) + if got := metadata.Get("x-codex-primary-reset-after-seconds"); got != "" { + t.Errorf("metadata %s=%s accepted as reset %q", field, raw, got) + } + } + }) + } +} + +func TestCodexFrameResetRejectsDurationOverflow(t *testing.T) { + at := time.Unix(1_700_000_000, 0) + for _, test := range []struct { + field, header, raw string + }{ + {"reset_after_seconds", "reset-after-seconds", "9223372037"}, + {"reset_at", "reset-at", "10923372037"}, + } { + if got := codexFrameResetSeconds(gjson.Parse(`{"`+test.field+`":`+test.raw+`}`), at); got != "" { + t.Errorf("duration-overflow %s accepted: %q", test.field, got) + } + headers := codexResponseMetadataHeaders(gjson.Parse(`{"headers":{"x-codex-primary-`+test.header+`":`+test.raw+`}}`), at) + if got := headers.Get("x-codex-primary-reset-after-seconds"); got != "" { + t.Errorf("metadata duration-overflow %s accepted: %q", test.header, got) + } + } +} + +func TestApplyCodexRateLimitsFrameIgnoresGPTCodexMeteredLimit(t *testing.T) { + acc := newUsageFrameTestAccount() + now := time.Now() + acc.SetUsageSnapshot5hAt(55, now.Add(time.Hour), now) + frame := gjson.Parse(`{ + "type":"codex.rate_limits", + "limit_name":"gpt-codex", + "rate_limits":{"primary":{"used_percent":99,"window_minutes":300,"reset_at":1700000900}} + }`) + + ApplyCodexUsageEventFrame(nil, acc, "codex.rate_limits", frame, now) + + pct, ok := acc.GetUsagePercent5h() + if !ok || pct != 55 { + t.Fatalf("metered gpt-codex frame overwrote main window: usage=%v valid=%v", pct, ok) + } +} + +func TestApplyCodexRateLimitsFrameAcceptsDefaultLimitName(t *testing.T) { + acc := newUsageFrameTestAccount() + now := time.Unix(1_700_000_000, 0) + frame := gjson.Parse(`{ + "type":"codex.rate_limits", + "limit_name":"codex", + "rate_limits":{"primary":{"used_percent":42,"window_minutes":300}} + }`) + + ApplyCodexUsageEventFrame(nil, acc, "codex.rate_limits", frame, now) + + if pct, ok := acc.GetUsagePercent5h(); !ok || pct != 42 { + t.Fatalf("default limit_name frame rejected: usage=%v valid=%v", pct, ok) + } +} + +func TestApplyCodexUsageHeaderObservationKeepsHigherUsageWithinSameWindow(t *testing.T) { + acc := newUsageFrameTestAccount() + t0 := time.Unix(1_700_000_000, 0) + first := codexUsageHeaderObservation{ + w5h: codexWindowUsage{usedPct: 50, resetSec: 3600, windowMin: 300, valid: true, hasReset: true}, + authoritative: true, + } + applyCodexUsageHeaderObservation(nil, acc, first, t0, false) + + second := first + second.w5h.usedPct = 40 + second.w5h.resetSec = 3590 // t0+10s+3590s is the same absolute reset. + applyCodexUsageHeaderObservation(nil, acc, second, t0.Add(10*time.Second), false) + if pct, _ := acc.GetUsagePercent5h(); pct != 50 { + t.Fatalf("out-of-order frame rolled 5h usage back to %v, want 50", pct) + } + + third := first + third.w5h.usedPct = 10 + third.w5h.resetSec = 18000 // next window. + applyCodexUsageHeaderObservation(nil, acc, third, t0.Add(20*time.Second), false) + if pct, _ := acc.GetUsagePercent5h(); pct != 10 { + t.Fatalf("new window usage = %v, want 10", pct) + } +} + +func TestApplyCodexUsageHeaderObservationBoundsMissingReset(t *testing.T) { + acc := newUsageFrameTestAccount() + t0 := time.Unix(1_700_000_000, 0) + + applyCodexUsageHeaderObservation(nil, acc, codexUsageHeaderObservation{ + w5h: codexWindowUsage{usedPct: 80, windowMin: 300, valid: true}, + authoritative: true, + }, t0, false) + if got, want := acc.GetReset5hAt(), t0.Add(5*time.Hour); !got.Equal(want) { + t.Fatalf("5h reset fallback = %v, want %v", got, want) + } + + applyCodexUsageHeaderObservation(nil, acc, codexUsageHeaderObservation{ + w7d: codexWindowUsage{usedPct: 90, windowMin: 10080, valid: true}, + authoritative: true, + }, t0, false) + if got, want := acc.GetReset7dAt(), t0.Add(7*24*time.Hour); !got.Equal(want) { + t.Fatalf("7d reset fallback = %v, want %v", got, want) + } +} + +func TestApplyCodexUsageHeaderObservationStaleResetStartsNewWindow(t *testing.T) { + acc := newUsageFrameTestAccount() + t0 := time.Unix(1_700_000_000, 0) + // The stored window has already reset, but the new frame omits reset_at. + acc.SetUsageSnapshot5hAt(100, t0.Add(-time.Minute), t0.Add(-time.Hour)) + acc.SetUsagePercent7d(100) + acc.SetReset7dAt(t0.Add(-time.Minute)) + + applyCodexUsageHeaderObservation(nil, acc, codexUsageHeaderObservation{ + w5h: codexWindowUsage{usedPct: 5, windowMin: 300, valid: true}, + w7d: codexWindowUsage{usedPct: 6, windowMin: 10080, valid: true}, + authoritative: true, + }, t0, false) + + if pct, ok := acc.GetUsagePercent5h(); !ok || pct != 5 { + t.Fatalf("5h usage = %v valid=%v, want 5 after the window rolled over", pct, ok) + } + if want := t0.Add(5 * time.Hour); !acc.GetReset5hAt().Equal(want) { + t.Fatalf("5h reset = %v, want %v", acc.GetReset5hAt(), want) + } + if pct, ok := acc.GetUsagePercent7d(); !ok || pct != 6 { + t.Fatalf("7d usage = %v valid=%v, want 6 after the window rolled over", pct, ok) + } + if want := t0.Add(7 * 24 * time.Hour); !acc.GetReset7dAt().Equal(want) { + t.Fatalf("7d reset = %v, want %v", acc.GetReset7dAt(), want) + } +} + +func TestCodexResetPrefersAbsoluteOverZeroRelative(t *testing.T) { + observedAt := time.Unix(1_700_000_000, 0) + frame := gjson.Parse(`{"rate_limits":{"primary":{ + "used_percent":50,"window_minutes":300, + "resets_in_seconds":0,"reset_at":1700003600 + }}}`) + + headers := codexRateLimitsHeaders(frame, observedAt) + if got := headers.Get("x-codex-primary-reset-after-seconds"); got != "3600" { + t.Fatalf("reset = %q, want the future absolute reset 3600", got) + } +} + +func TestParseCodexWindowUsageRejectsTrailingGarbage(t *testing.T) { + if got := parseCodexWindowUsage("42abc", "300", "60"); got.valid { + t.Fatalf("trailing garbage in used_percent accepted: %+v", got) + } + if got := parseCodexWindowUsage("42", "300x", "60"); got.valid { + t.Fatalf("trailing garbage in window_minutes accepted: %+v", got) + } + if got := parseCodexWindowUsage("42", "300", "60abc"); !got.valid { + t.Fatal("valid window rejected because the reset value is malformed") + } else if got.hasReset { + t.Fatalf("malformed reset accepted as a reset value: %+v", got) + } +} + +func TestCodexRateLimitsFrameUsesLimitIDWhenDisplayNameDiffers(t *testing.T) { + acc := newUsageFrameTestAccount() + now := time.Unix(1_700_000_000, 0) + frame := gjson.Parse(`{ + "type":"codex.rate_limits", + "limit_id":"codex", + "limit_name":"gpt-5.2-codex-sonic", + "rate_limits":{"primary":{"used_percent":42,"window_minutes":300}} + }`) + + ApplyCodexUsageEventFrame(nil, acc, "codex.rate_limits", frame, now) + + if pct, ok := acc.GetUsagePercent5h(); !ok || pct != 42 { + t.Fatalf("main limit frame rejected: usage=%v valid=%v", pct, ok) + } +} + +func TestApplyCodexIdleUsageFrameAppliesCredits(t *testing.T) { + acc := newUsageFrameTestAccount() + acc.ApplySparseCreditsHeaders(true, false, nil, nil, "workspace_member_credits_depleted") + + ApplyCodexIdleUsageFrame(nil, acc, "codex.rate_limits", []byte(`{ + "type":"codex.rate_limits", + "credits":{"has_credits":false,"unlimited":false,"balance":"0"} + }`)) + + credits, ok := acc.GetCreditBalance() + if !ok || credits.HasCredits { + t.Fatalf("idle frame did not clear the credits flag: %+v valid=%v", credits, ok) + } + if got := acc.GetRateLimitReachedType(); got != "workspace_member_credits_depleted" { + t.Fatalf("idle credits frame cleared the workspace hard stop: %q", got) + } +} + +func TestCodexRateLimitsFrameIgnoresSingleCreditsFlag(t *testing.T) { + observedAt := time.Unix(1_700_000_000, 0) + frame := gjson.Parse(`{"credits":{"has_credits":true}}`) + + headers := codexRateLimitsHeaders(frame, observedAt) + if got := headers.Get("x-codex-credits-has-credits"); got != "" { + t.Fatalf("single credits flag was adopted: %q", got) + } +} + +func TestCodexUsageFramesSkipRelayAccounts(t *testing.T) { + acc := newUsageFrameTestAccount() + acc.UpstreamType = auth.UpstreamOpenAIResponses + acc.BaseURL = "relay.example" + acc.APIKey = "relay-key" + frame := gjson.Parse(`{"rate_limits":{"primary":{"used_percent":42,"window_minutes":300}}}`) + + ApplyCodexUsageEventFrame(nil, acc, "codex.rate_limits", frame, time.Now()) + ApplyCodexIdleUsageFrame(nil, acc, "codex.rate_limits", []byte(`{"rate_limits":{"primary":{"used_percent":42,"window_minutes":300}}}`)) + + if pct, ok := acc.GetUsagePercent5h(); ok { + t.Fatalf("relay account ingested a forwarded frame: usage=%v", pct) + } +} + +func TestApplyCodexUsageEventFrameIgnoresOlderObservation(t *testing.T) { + acc := newUsageFrameTestAccount() + t2 := time.Unix(1_700_000_000, 0) + t1 := t2.Add(-time.Minute) + newer := gjson.Parse(`{"rate_limits":{"primary":{"used_percent":60,"window_minutes":300}}}`) + older := gjson.Parse(`{"rate_limits":{"primary":{"used_percent":30,"window_minutes":300}}}`) + + ApplyCodexUsageEventFrame(nil, acc, "codex.rate_limits", newer, t2) + ApplyCodexUsageEventFrame(nil, acc, "codex.rate_limits", older, t1) + + if pct, _ := acc.GetUsagePercent5h(); pct != 60 { + t.Fatalf("older observation rolled usage back to %v, want 60", pct) + } +} + +func TestCodexUsageFramesOrderCreditsAndFreshnessWithUsage(t *testing.T) { + for _, creditsOnly := range []bool{false, true} { + t.Run(fmt.Sprintf("creditsOnly=%v", creditsOnly), func(t *testing.T) { + acc := newUsageFrameTestAccount() + at := time.Now() + newer := gjson.Parse(`{"rate_limits":{"primary":{"used_percent":60,"window_minutes":300}},"credits":{"has_credits":false,"unlimited":false,"balance":"0"}}`) + older := `{"rate_limits":{"primary":{"used_percent":10,"window_minutes":300}},"credits":{"has_credits":true,"unlimited":false,"balance":"100"}}` + if creditsOnly { + older = `{"credits":{"has_credits":true,"unlimited":false,"balance":"100"}}` + } + ApplyCodexUsageEventFrame(nil, acc, "codex.rate_limits", newer, at) + ApplyCodexUsageEventFrame(nil, acc, "codex.rate_limits", gjson.Parse(older), at.Add(-time.Second)) + credits, valid := acc.GetCreditBalance() + snap := acc.GetAccountLiveUsageSnapshot() + if !valid || credits.HasCredits || snap.UsagePercent5h != 60 || !snap.UsageFrameLastAt.Equal(at) { + t.Fatalf("older frame applied: credits=%+v valid=%v snapshot=%+v", credits, valid, snap) + } + ApplyCodexUsageEventFrame(nil, acc, "codex.rate_limits", gjson.Parse(`{"credits":{"has_credits":true,"unlimited":false,"balance":"100"}}`), at.Add(time.Second)) + credits, valid = acc.GetCreditBalance() + if !valid || !credits.HasCredits || !acc.GetAccountLiveUsageSnapshot().UsageFrameLastAt.Equal(at.Add(time.Second)) { + t.Fatal("newer credits-only observation was not applied") + } + }) + } +} + +func TestCodexUsageFramePersistenceRespectsThrottle(t *testing.T) { + ctx := context.Background() + dbPath := filepath.Join(t.TempDir(), "usage-frame.db") + db, err := database.New("sqlite", dbPath) + if err != nil { + t.Fatalf("database.New: %v", err) + } + t.Cleanup(func() { _ = db.Close() }) + accountID, err := db.InsertAccountWithCredentials(ctx, "frame-test", map[string]any{ + "access_token": "at-test", + "plan_type": "plus", + }, "") + if err != nil { + t.Fatalf("InsertAccountWithCredentials: %v", err) + } + + store := auth.NewStore(db, nil, &database.SystemSettings{MaxConcurrency: 4}) + t.Cleanup(store.Stop) + account := newUsageFrameTestAccount() + account.DBID = accountID + account.AccessToken = "at-test" + store.AddAccount(account) + + previous := codexUsageFramePersistInterval() + SetCodexUsageFramePersistInterval(time.Hour) + t.Cleanup(func() { SetCodexUsageFramePersistInterval(previous) }) + + t0 := time.Now() + apply := func(pct float64, at time.Time) { + frame := gjson.Parse(fmt.Sprintf( + `{"type":"codex.rate_limits","rate_limits":{"primary":{"used_percent":%v,"window_minutes":300,"reset_at":%d}}}`, + pct, at.Add(3*time.Hour).Unix(), + )) + ApplyCodexUsageEventFrame(store, account, "codex.rate_limits", frame, at) + } + + apply(10, t0) + waitPersisted5h(t, dbPath, accountID, 10) + waitUsageFrameThrottled(t, account) + + apply(12, t0.Add(time.Second)) + if pct, _ := account.GetUsagePercent5h(); pct != 12 { + t.Fatalf("throttled frame did not update memory: %v", pct) + } + time.Sleep(150 * time.Millisecond) + if got := readPersisted5h(t, dbPath, accountID); got != 10 { + t.Fatalf("throttled frame wrote %v, want the previous 10", got) + } + + apply(100, t0.Add(2*time.Second)) + waitPersisted5h(t, dbPath, accountID, 100) +} + +// TestCodexUsageFramePersistDoesNotBlockOnLockedDB verifies the read path only +// enqueues the snapshot; a database write lock must not stall the caller. +func TestCodexUsageFramePersistDoesNotBlockOnLockedDB(t *testing.T) { + t.Run("usage", func(t *testing.T) { + testCodexFrameOnLockedDB(t, `{"rate_limits":{"primary":{"used_percent":42,"window_minutes":300}}}`, 42) + }) + t.Run("exhaustion_credits_and_plan", func(t *testing.T) { + testCodexFrameOnLockedDB(t, `{"plan_type":"pro","rate_limits":{"primary":{"used_percent":100,"window_minutes":300},"secondary":{"used_percent":100,"window_minutes":10080}},"credits":{"has_credits":true,"unlimited":false,"balance":"10"}}`, 100) + }) + t.Run("credits_only", func(t *testing.T) { + testCodexFrameOnLockedDB(t, `{"credits":{"has_credits":true,"unlimited":false,"balance":"10"}}`, -1) + }) +} + +func testCodexFrameOnLockedDB(t *testing.T, payload string, wantUsage float64) { + t.Helper() + ctx := context.Background() + dbPath := filepath.Join(t.TempDir(), "usage-frame-locked.db") + db, err := database.New("sqlite", dbPath) + if err != nil { + t.Fatalf("database.New: %v", err) + } + t.Cleanup(func() { _ = db.Close() }) + accountID, err := db.InsertAccountWithCredentials(ctx, "frame-locked", map[string]any{ + "access_token": "at-test", "plan_type": "plus", + }, "") + if err != nil { + t.Fatalf("InsertAccountWithCredentials: %v", err) + } + store := auth.NewStore(db, nil, &database.SystemSettings{MaxConcurrency: 4}) + t.Cleanup(store.Stop) + account := newUsageFrameTestAccount() + account.DBID = accountID + account.AccessToken = "at-test" + store.AddAccount(account) + + lockConn, err := sql.Open("sqlite", dbPath) + if err != nil { + t.Fatalf("sql.Open: %v", err) + } + t.Cleanup(func() { _ = lockConn.Close() }) + conn, err := lockConn.Conn(ctx) + if err != nil { + t.Fatalf("conn: %v", err) + } + t.Cleanup(func() { _ = conn.Close() }) + if _, err := conn.ExecContext(ctx, "BEGIN IMMEDIATE"); err != nil { + t.Fatalf("begin immediate: %v", err) + } + + frame := gjson.Parse(payload) + done := make(chan struct{}) + go func() { + ApplyCodexUsageEventFrame(store, account, "codex.rate_limits", frame, time.Now()) + close(done) + }() + returned := false + select { + case <-done: + returned = true + case <-time.After(500 * time.Millisecond): + t.Error("frame path waited on a locked database") + } + if _, err := conn.ExecContext(ctx, "ROLLBACK"); err != nil { + t.Fatalf("rollback: %v", err) + } + if !returned { + select { + case <-done: + case <-time.After(5 * time.Second): + t.Fatal("frame did not finish after the DB lock was released") + } + } + if wantUsage >= 0 { + if pct, ok := account.GetUsagePercent5h(); !ok || pct != wantUsage { + t.Fatalf("memory usage = %v valid=%v, want %v", pct, ok, wantUsage) + } + waitPersisted5h(t, dbPath, accountID, wantUsage) + } else { + deadline := time.Now().Add(3 * time.Second) + for { + row, err := db.GetAccountByID(ctx, accountID) + if err != nil { + t.Fatal(err) + } + if row.GetCredential("codex_credits") != "" { + break + } + if time.Now().After(deadline) { + t.Fatal("credits-only frame was not persisted") + } + time.Sleep(10 * time.Millisecond) + } + } +} + +func waitUsageFrameThrottled(t *testing.T, account *auth.Account) { + t.Helper() + deadline := time.Now().Add(3 * time.Second) + for account.ShouldPersistUsageFrame(time.Now(), 10, true, 0, false, time.Hour) { + if time.Now().After(deadline) { + t.Fatal("usage frame throttle did not advance after the async write") + } + time.Sleep(10 * time.Millisecond) + } +} + +func waitPersisted5h(t *testing.T, dbPath string, accountID int64, want float64) { + t.Helper() + deadline := time.Now().Add(3 * time.Second) + for { + if got := readPersisted5h(t, dbPath, accountID); got == want { + return + } + if time.Now().After(deadline) { + t.Fatalf("persisted 5h usage = %v, want %v", readPersisted5h(t, dbPath, accountID), want) + } + time.Sleep(20 * time.Millisecond) + } +} + +func readPersisted5h(t *testing.T, dbPath string, accountID int64) float64 { + t.Helper() + conn, err := sql.Open("sqlite", dbPath) + if err != nil { + t.Fatalf("sql.Open: %v", err) + } + defer func() { _ = conn.Close() }() + var raw string + if err := conn.QueryRow(`SELECT credentials FROM accounts WHERE id = $1`, accountID).Scan(&raw); err != nil { + t.Fatalf("read credentials: %v", err) + } + return gjson.Get(raw, "codex_5h_used_percent").Float() +} diff --git a/proxy/handler.go b/proxy/handler.go index a0c7acb1a..0493ff2c1 100644 --- a/proxy/handler.go +++ b/proxy/handler.go @@ -10,6 +10,7 @@ import ( "fmt" "io" "log" + "math" "math/rand" "net" "net/http" @@ -180,15 +181,22 @@ type usageLimitDetails struct { } type CodexUsageSyncResult struct { - UsagePct7d float64 - HasUsage7d bool - Usage7dRateLimited bool - UsagePct5h float64 - Reset5hAt time.Time - HasUsage5h bool - Used5hHeaders bool - Persisted5hOnly bool - Premium5hRateLimited bool + UsagePct7d float64 + HasUsage7d bool + Usage7dRateLimited bool + UsagePct5h float64 + Reset5hAt time.Time + HasUsage5h bool + Used5hHeaders bool + Persisted5hOnly bool + // Premium5hRateLimited 表示本次同步把账号标记为 premium 5h 限流。 + Premium5hRateLimited bool + // PersistUsage 表示节流策略允许本次观察落库。WS 帧路径可能只更新内存 + // 快照；100% 冷却标记只跟着允许落库的观察执行，避免每帧重复写库。 + PersistUsage bool + // Applied 表示本次观察通过了排序门并写入了账号状态；WS 帧路径用它决定 + // 是否更新“最后一帧”观测标记，避免迟到的旧帧伪装成最新帧。 + Applied bool UsageWindowLimitsIgnored bool // Cleared5h 表示本次同步因上游未返回 5h 窗口而清除了本地陈旧 5h 快照（issue #382）。 Cleared5h bool @@ -4688,6 +4696,7 @@ func (h *Handler) Responses(c *gin.Context) { } parsed := gjson.ParseBytes(data) eventType := normalizedUpstreamSSEEventType(sseEvent, data) + ApplyCodexUsageEventFrame(h.store, account, eventType, parsed, time.Now()) ttftGuard.MarkProgress(eventType) isFirstToken := isLooseFirstTokenResult(parsed) if !ttftRecorded && isFirstToken { @@ -5362,6 +5371,7 @@ func (h *Handler) Responses(c *gin.Context) { } parsed := gjson.ParseBytes(data) eventType := normalizedUpstreamSSEEventType(sseEvent, data) + ApplyCodexUsageEventFrame(h.store, account, eventType, parsed, time.Now()) // TTFT: 记录第一个实际内容事件的时间 ttftGuard.MarkProgress(eventType) @@ -5612,6 +5622,7 @@ func (h *Handler) Responses(c *gin.Context) { } parsed := gjson.ParseBytes(data) eventType := normalizedUpstreamSSEEventType(sseEvent, data) + ApplyCodexUsageEventFrame(h.store, account, eventType, parsed, time.Now()) if eventType == "error" { terminalFailurePayload = terminalUpstreamErrorPayload(data) gotTerminal = true @@ -7457,6 +7468,7 @@ func (h *Handler) ChatCompletions(c *gin.Context) { readErr = readSSEStreamWithContinuousRetryKeepalive(readCtx, resp.Body, func(sseEvent string, data []byte) bool { parsed := gjson.ParseBytes(data) eventType := normalizedUpstreamSSEEventType(sseEvent, data) + ApplyCodexUsageEventFrame(h.store, account, eventType, parsed, time.Now()) if eventType == "response.failed" { statusCode := classifyResponseFailedOutcome(data).logStatusCode var incidentID string @@ -7625,6 +7637,7 @@ func (h *Handler) ChatCompletions(c *gin.Context) { outputCollector.Add(data) parsed := gjson.ParseBytes(data) eventType := normalizedUpstreamSSEEventType(sseEvent, data) + ApplyCodexUsageEventFrame(h.store, account, eventType, parsed, time.Now()) ttftGuard.MarkProgress(eventType) if !ttftRecorded && isLooseFirstTokenResult(parsed) { firstTokenMs = int(time.Since(start).Milliseconds()) @@ -8134,6 +8147,11 @@ type codexWindowUsage struct { resetSec float64 windowMin float64 valid bool + // hasReset distinguishes an explicit reset (possibly "0", meaning the + // window resets now) from a payload that omitted reset information. A + // partial frame must not overwrite a known reset with "now" and disable + // the auto-pause fence. + hasReset bool } func parseCodexWindowUsage(usedStr, windowStr, resetStr string) codexWindowUsage { @@ -8141,15 +8159,28 @@ func parseCodexWindowUsage(usedStr, windowStr, resetStr string) codexWindowUsage return codexWindowUsage{} } var usedPct, windowMin, resetSec float64 - if _, err := fmt.Sscanf(usedStr, "%f", &usedPct); err != nil { + // Parse the whole value: fmt.Sscanf("%f") accepts trailing garbage, so a + // malformed header such as "42abc" would be treated as 42. + parsedUsed, err := strconv.ParseFloat(strings.TrimSpace(usedStr), 64) + if err != nil || math.IsNaN(parsedUsed) { return codexWindowUsage{} } - if _, err := fmt.Sscanf(windowStr, "%f", &windowMin); err != nil || windowMin <= 0 { + usedPct = parsedUsed + // Reject out-of-range percentages instead of pacing an account on a + // malformed observation; keep the previous snapshot until the next probe. + if usedPct < 0 || usedPct > 100 { return codexWindowUsage{} } - if strings.TrimSpace(resetStr) != "" { - if _, err := fmt.Sscanf(resetStr, "%f", &resetSec); err != nil { - resetSec = 0 + parsedWindow, err := strconv.ParseFloat(strings.TrimSpace(windowStr), 64) + if err != nil || math.IsNaN(parsedWindow) || math.IsInf(parsedWindow, 0) || parsedWindow <= 0 { + return codexWindowUsage{} + } + windowMin = parsedWindow + hasReset := false + if rawReset := strings.TrimSpace(resetStr); rawReset != "" { + if parsedReset, err := strconv.ParseFloat(rawReset, 64); err == nil && !math.IsNaN(parsedReset) && !math.IsInf(parsedReset, 0) { + resetSec = parsedReset + hasReset = true } } return codexWindowUsage{ @@ -8157,7 +8188,32 @@ func parseCodexWindowUsage(usedStr, windowStr, resetStr string) codexWindowUsage windowMin: windowMin, resetSec: resetSec, valid: true, + hasReset: hasReset, + } +} + +// codexResetAfterSeconds returns the reset offset for one rate-limit window. +// Upstreams have used both families: x-codex--reset-after-seconds (relative) +// and x-codex--reset-at (absolute Unix seconds, the family the official +// Codex client parses). Normalize the absolute form to a relative offset so the +// existing snapshot code keeps one representation. +func codexResetAfterSeconds(header http.Header, prefix string) string { + if raw := strings.TrimSpace(header.Get(prefix + "reset-after-seconds")); raw != "" { + return raw + } + raw := strings.TrimSpace(header.Get(prefix + "reset-at")) + if raw == "" { + return "" } + epoch, err := strconv.ParseInt(raw, 10, 64) + if err != nil || epoch <= 0 { + return "" + } + seconds := time.Until(time.Unix(epoch, 0)).Seconds() + if seconds < 0 { + seconds = 0 + } + return strconv.FormatFloat(seconds, 'f', -1, 64) } func classifyCodex429Window(resp *http.Response, now time.Time) (codexRateLimitWindow, time.Time, bool) { @@ -8168,12 +8224,12 @@ func classifyCodex429Window(resp *http.Response, now time.Time) (codexRateLimitW primary := parseCodexWindowUsage( resp.Header.Get("x-codex-primary-used-percent"), resp.Header.Get("x-codex-primary-window-minutes"), - resp.Header.Get("x-codex-primary-reset-after-seconds"), + codexResetAfterSeconds(resp.Header, "x-codex-primary-"), ) secondary := parseCodexWindowUsage( resp.Header.Get("x-codex-secondary-used-percent"), resp.Header.Get("x-codex-secondary-window-minutes"), - resp.Header.Get("x-codex-secondary-reset-after-seconds"), + codexResetAfterSeconds(resp.Header, "x-codex-secondary-"), ) var exhausted []codexWindowUsage @@ -8209,7 +8265,7 @@ func responseHasCodex5hHeaders(resp *http.Response) bool { primary := parseCodexWindowUsage( resp.Header.Get("x-codex-primary-used-percent"), resp.Header.Get("x-codex-primary-window-minutes"), - resp.Header.Get("x-codex-primary-reset-after-seconds"), + codexResetAfterSeconds(resp.Header, "x-codex-primary-"), ) if primary.valid && codexWindowType(primary.windowMin) == codexRateLimitWindow5h { return true @@ -8218,7 +8274,7 @@ func responseHasCodex5hHeaders(resp *http.Response) bool { secondary := parseCodexWindowUsage( resp.Header.Get("x-codex-secondary-used-percent"), resp.Header.Get("x-codex-secondary-window-minutes"), - resp.Header.Get("x-codex-secondary-reset-after-seconds"), + codexResetAfterSeconds(resp.Header, "x-codex-secondary-"), ) return secondary.valid && codexWindowType(secondary.windowMin) == codexRateLimitWindow5h } @@ -8682,88 +8738,178 @@ func windowMinutesToCooldown(windowMinutes float64) time.Duration { // SyncCodexUsageState 解析 Codex 响应头并完成 7d / 5h 快照持久化与 premium 5h 提前限流。 func SyncCodexUsageState(store *auth.Store, account *auth.Account, resp *http.Response) CodexUsageSyncResult { + return syncCodexUsageState(store, account, resp, time.Now(), codexUsageSyncOptions{ + persistUsageSnapshot: true, + clearAbsent5h: true, + }) +} + +type codexUsageSyncOptions struct { + // persistUsageSnapshot=false keeps the in-memory/scheduler update but skips + // the DB write. WS frames use this for throttled persistence. + persistUsageSnapshot bool + // clearAbsent5h controls whether an authoritative observation that omits + // the 5h window clears a previous 5h snapshot. Per-turn WS frames may be + // partial, so they must not erase a valid window. + clearAbsent5h bool + // shouldPersistUsage, when set, is evaluated inside ApplyUsageObservation's + // ordering gate. That gate rejects stale observers, so the throttling state + // must not advance for an observation that will not be applied. + shouldPersistUsage func(pct5h float64, has5h bool, pct7d float64, has7d bool) bool + // asyncPersist moves the frame snapshot write to the store's coalescing + // worker. The upstream read path only applies memory/scheduler state, so a + // slow or locked database cannot stall a streaming response. + asyncPersist bool + // preserveAbsentHardStop is used by partial WS quota frames, which cannot + // clear a workspace stop unless they explicitly carry that field. + preserveAbsentHardStop bool + // observationLocked is only set by the WS ingress ordering gate. + observationLocked bool +} + +func syncCodexUsageState(store *auth.Store, account *auth.Account, resp *http.Response, observedAt time.Time, opts codexUsageSyncOptions) CodexUsageSyncResult { result := CodexUsageSyncResult{} if account == nil || resp == nil { return result } - observedAt := time.Now() - if store != nil { - planHeader := resp.Header.Get("x-codex-plan-type") - store.UpdateAccountPlanType(account, planHeader) - // 权威付费 plan_type 与「订阅已过期」互斥，借每次响应校正陈旧到期时间。(issue #360) - if planHeader != "" { - store.ClearStaleSubscriptionExpiresAt(account) - } + if observedAt.IsZero() { + observedAt = time.Now() } - observation := parseCodexUsageHeaderObservation(resp) + hasCredits, unlimited, creditsObserved := parseCodexCreditsFlagHeaders(resp) rateLimitReachedType := strings.TrimSpace(resp.Header.Get("x-codex-rate-limit-reached-type")) - result.RateLimitReachedType = rateLimitReachedType - if hasCredits, unlimited, ok := parseCodexCreditsFlagHeaders(resp); ok { - result.CreditsObserved = true - var balPtr *string - if trimmed := strings.TrimSpace(resp.Header.Get("x-codex-credits-balance")); trimmed != "" { - balPtr = &trimmed - } - var overagePtr *bool - if raw := strings.TrimSpace(resp.Header.Get("x-codex-credits-overage-reached")); raw != "" { - if ov, err := strconv.ParseBool(raw); err == nil { - overagePtr = &ov + apply := func() { + if store != nil { + planHeader := resp.Header.Get("x-codex-plan-type") + if opts.asyncPersist { + store.ObserveUsageFramePlan(account, planHeader) + } else { + store.UpdateAccountPlanType(account, planHeader) + // 权威付费 plan_type 与「订阅已过期」互斥，借每次响应校正陈旧到期时间。(issue #360) + if planHeader != "" { + store.ClearStaleSubscriptionExpiresAt(account) + } } } - if store != nil { - store.PersistSparseCreditObservation(account, hasCredits, unlimited, balPtr, overagePtr, rateLimitReachedType) - } else { - account.ApplySparseCreditsHeaders(hasCredits, unlimited, balPtr, overagePtr, rateLimitReachedType) + + if opts.preserveAbsentHardStop && rateLimitReachedType == "" { + rateLimitReachedType = account.GetRateLimitReachedType() } - } else if rateLimitReachedType != "" || observation.authoritative { - // reached-type 是逐响应字段：带用量头的响应没有它就是「未触达」，要把旧的 - // workspace hard-stop 清掉，否则工作区充值后积分顶替永远回不来。 - if store != nil { - store.PersistRateLimitReachedType(account, rateLimitReachedType) - } else { - account.SetRateLimitReachedType(rateLimitReachedType) + result.RateLimitReachedType = rateLimitReachedType + if creditsObserved { + result.CreditsObserved = true + var balPtr *string + if trimmed := strings.TrimSpace(resp.Header.Get("x-codex-credits-balance")); trimmed != "" { + balPtr = &trimmed + } + var overagePtr *bool + if raw := strings.TrimSpace(resp.Header.Get("x-codex-credits-overage-reached")); raw != "" { + if ov, err := strconv.ParseBool(raw); err == nil { + overagePtr = &ov + } + } + if store != nil && !opts.asyncPersist { + store.PersistSparseCreditObservation(account, hasCredits, unlimited, balPtr, overagePtr, rateLimitReachedType) + } else { + account.ApplySparseCreditsHeaders(hasCredits, unlimited, balPtr, overagePtr, rateLimitReachedType) + } + } else if rateLimitReachedType != "" || observation.authoritative { + // reached-type 是逐响应字段：带用量头的响应没有它就是「未触达」，要把旧的 + // workspace hard-stop 清掉，否则工作区充值后积分顶替永远回不来。 + if store != nil && !opts.asyncPersist { + store.PersistRateLimitReachedType(account, rateLimitReachedType) + } else { + account.SetRateLimitReachedType(rateLimitReachedType) + } } - } - result.UsageWindowLimitsIgnored = account.SkipsUsageWindowLimits() + result.UsageWindowLimitsIgnored = account.SkipsUsageWindowLimits() - result.Used5hHeaders = observation.w5h.valid - usageApplied := false - if observation.authoritative { - usageApplied = account.ApplyUsageObservation(observedAt, func() { - parsed := applyCodexUsageHeaderObservation(store, account, observation, observedAt) + result.Used5hHeaders = observation.w5h.valid + if observation.authoritative { + parsed := applyCodexUsageHeaderObservation(store, account, observation, observedAt, opts.clearAbsent5h) result.UsagePct7d = parsed.usagePct7d result.HasUsage7d = parsed.hasUsage7d result.Cleared5h = parsed.cleared5h if store == nil { return } + persistUsage := opts.persistUsageSnapshot + if opts.shouldPersistUsage != nil { + persistUsage = opts.shouldPersistUsage( + observation.w5h.usedPct, + observation.w5h.valid, + observation.w7d.usedPct, + observation.w7d.valid, + ) + } + result.PersistUsage = persistUsage + persistedUsage := false if result.HasUsage7d { - store.PersistUsageSnapshot(account, result.UsagePct7d) - if result.UsagePct7d >= 100 { + if opts.asyncPersist { + // Memory and scheduler update stay synchronous; the DB write + // is coalesced per account off the read path. + store.PublishUsageSnapshot(account, result.UsagePct7d, observedAt, false) + } else { + persistedUsage = store.PublishUsageSnapshot(account, result.UsagePct7d, observedAt, persistUsage) + } + if result.UsagePct7d >= 100 && persistUsage && !opts.asyncPersist { result.Usage7dRateLimited = store.MarkUsage7dRateLimited(account) } } else if result.Used5hHeaders { - store.PersistUsageSnapshot5hOnly(account) - result.Persisted5hOnly = true + if opts.asyncPersist { + store.PublishUsageSnapshot5hOnly(account, observedAt, false) + } else { + persistedUsage = store.PublishUsageSnapshot5hOnly(account, observedAt, persistUsage) + } + result.Persisted5hOnly = persistUsage } - }) + if persistUsage && persistedUsage { + // Advance the frame throttle only after the write succeeded so a + // failed write is retried on the next frame instead of sleeping + // through the interval. + pct5h, has5h := account.GetUsagePercent5h() + pct7d, has7d := account.GetUsagePercent7d() + account.MarkUsageFramePersisted(observedAt, pct5h, has5h, pct7d, has7d) + } + if opts.asyncPersist && persistUsage { + result.Usage7dRateLimited, result.Premium5hRateLimited = store.MarkUsageFrameCooldowns(account, observedAt) + } + } } + applied := false + if observation.authoritative || creditsObserved || rateLimitReachedType != "" || resp.Header.Get("x-codex-plan-type") != "" { + if opts.observationLocked { + apply() + applied = true + } else { + applied = account.ApplyUsageObservation(observedAt, apply) + } + } + usageApplied := applied && observation.authoritative result.UsagePct5h, result.Reset5hAt, result.HasUsage5h = account.GetUsageSnapshot5h() if usageApplied && store != nil && result.HasUsage5h { // 被动 /responses 头刷新了 5h 窗口重置时刻：武装「到点即探」，窗口翻新即刷新进度条。 store.WakeBoundaryProbe(result.Reset5hAt) } - if usageApplied && result.Used5hHeaders && account.IsPremium5hPlan() && result.HasUsage5h && result.UsagePct5h >= 100 && !account.SkipsUsageWindowLimits() { + if !opts.asyncPersist && usageApplied && result.Used5hHeaders && account.IsPremium5hPlan() && result.HasUsage5h && result.UsagePct5h >= 100 && !account.SkipsUsageWindowLimits() { if store != nil { - result.Premium5hRateLimited = store.MarkPremium5hRateLimitedAt(account, result.Reset5hAt, observedAt) + // The per-frame path may decide not to persist this observation; a + // repeated 100% frame must not issue another cooldown write. + if result.PersistUsage { + result.Premium5hRateLimited = store.MarkPremium5hRateLimitedAt(account, result.Reset5hAt, observedAt) + } } else { result.Premium5hRateLimited = true } } + result.Applied = applied && (usageApplied || result.CreditsObserved) + if applied && opts.asyncPersist && store != nil { + store.RefreshUsageFrameScheduler(account) + store.EnqueueUsageFrameEffects(account, observedAt, result.PersistUsage, resp.Header.Get("x-codex-plan-type") != "") + } return result } @@ -8791,7 +8937,7 @@ func parseCodexUsageHeaders(resp *http.Response, account *auth.Account) (float64 } parsed := codexUsageHeaderParseResult{} account.ApplyUsageObservation(observedAt, func() { - parsed = applyCodexUsageHeaderObservation(nil, account, observation, observedAt) + parsed = applyCodexUsageHeaderObservation(nil, account, observation, observedAt, true) }) return parsed.usagePct7d, parsed.hasUsage7d } @@ -8828,10 +8974,10 @@ func parseCodexUsageHeaderObservation(resp *http.Response) codexUsageHeaderObser // 解析 primary 和 secondary 窗口 primaryUsedStr := resp.Header.Get("x-codex-primary-used-percent") primaryWindowStr := resp.Header.Get("x-codex-primary-window-minutes") - primaryResetStr := resp.Header.Get("x-codex-primary-reset-after-seconds") + primaryResetStr := codexResetAfterSeconds(resp.Header, "x-codex-primary-") secondaryUsedStr := resp.Header.Get("x-codex-secondary-used-percent") secondaryWindowStr := resp.Header.Get("x-codex-secondary-window-minutes") - secondaryResetStr := resp.Header.Get("x-codex-secondary-reset-after-seconds") + secondaryResetStr := codexResetAfterSeconds(resp.Header, "x-codex-secondary-") primary := parseCodexWindowUsage(primaryUsedStr, primaryWindowStr, primaryResetStr) secondary := parseCodexWindowUsage(secondaryUsedStr, secondaryWindowStr, secondaryResetStr) @@ -8855,26 +9001,72 @@ func parseCodexUsageHeaderObservation(resp *http.Response) codexUsageHeaderObser return out } -func applyCodexUsageHeaderObservation(store *auth.Store, account *auth.Account, observation codexUsageHeaderObservation, observedAt time.Time) codexUsageHeaderParseResult { +func applyCodexUsageHeaderObservation(store *auth.Store, account *auth.Account, observation codexUsageHeaderObservation, observedAt time.Time, clearAbsent5h bool) codexUsageHeaderParseResult { out := codexUsageHeaderParseResult{} if observation.w5h.valid { - resetAt := observedAt.Add(time.Duration(observation.w5h.resetSec) * time.Second) - account.SetUsageSnapshot5hAt(observation.w5h.usedPct, resetAt, observedAt) - } else if observation.authoritative { + existingReset := account.GetReset5hAt() + existingPct, existingValid := account.GetUsagePercent5h() + resetAt := existingReset + // Only a reset that is still in the future identifies the current + // window. If it has already passed, a frame without reset_at belongs + // to the next window and must not inherit the stale timestamp. + if observation.w5h.hasReset { + resetAt = observedAt.Add(time.Duration(observation.w5h.resetSec) * time.Second) + } else if resetAt.IsZero() || !resetAt.After(observedAt) { + // A partial frame without reset must not create a zero reset that + // never expires; bound the pause to one window length instead. + resetAt = observedAt.Add(time.Duration(observation.w5h.windowMin) * time.Minute) + } + pct := observation.w5h.usedPct + if existingValid && existingReset.After(observedAt) && sameUsageResetWindow(existingReset, resetAt) && pct < existingPct { + // Within one window usage only increases. A late frame from an + // earlier concurrent turn must not roll the snapshot back and + // temporarily un-pause the account. + pct = existingPct + } + account.SetUsageSnapshot5hAt(pct, resetAt, observedAt) + } else if observation.authoritative && clearAbsent5h { out.cleared5h = store.ClearAbsentUsageSnapshot5hAt(account, observedAt) } if observation.w7d.valid { - resetAt := observedAt.Add(time.Duration(observation.w7d.resetSec) * time.Second) - account.SetReset7dAt(resetAt) + existingReset := account.GetReset7dAt() + existingPct, existingValid := account.GetUsagePercent7d() + resetAt := existingReset + if observation.w7d.hasReset { + resetAt = observedAt.Add(time.Duration(observation.w7d.resetSec) * time.Second) + account.SetReset7dAt(resetAt) + } else if resetAt.IsZero() || !resetAt.After(observedAt) { + resetAt = observedAt.Add(time.Duration(observation.w7d.windowMin) * time.Minute) + account.SetReset7dAt(resetAt) + } + pct := observation.w7d.usedPct + if existingValid && existingReset.After(observedAt) && sameUsageResetWindow(existingReset, resetAt) && pct < existingPct { + pct = existingPct + } account.SetWindow7dSeconds(int64(observation.w7d.windowMin * 60)) - account.SetUsagePercent7d(observation.w7d.usedPct) - out.usagePct7d = observation.w7d.usedPct + account.SetUsagePercent7d(pct) + out.usagePct7d = pct out.hasUsage7d = true } return out } +// sameUsageResetWindow reports whether two observations describe the same +// rolling window. Relative reset headers are recomputed against the receipt +// time, so allow a small tolerance while still distinguishing the next +// (hours later) window. +func sameUsageResetWindow(a, b time.Time) bool { + if a.IsZero() || b.IsZero() { + return false + } + diff := a.Sub(b) + if diff < 0 { + diff = -diff + } + return diff <= time.Minute +} + // ParseCodexUsageHeaders 从响应头提取并更新账号用量信息 func ParseCodexUsageHeaders(resp *http.Response, account *auth.Account) (float64, bool) { return parseCodexUsageHeaders(resp, account) diff --git a/proxy/handler_anthropic.go b/proxy/handler_anthropic.go index f0fdde762..79bdbf9fb 100644 --- a/proxy/handler_anthropic.go +++ b/proxy/handler_anthropic.go @@ -1460,6 +1460,7 @@ func (h *Handler) Messages(c *gin.Context) { } parsed := gjson.ParseBytes(data) eventType := normalizedUpstreamSSEEventType(sseEvent, data) + ApplyCodexUsageEventFrame(h.store, account, eventType, parsed, time.Now()) // TTFT 跟踪 ttftGuard.MarkProgress(eventType) @@ -1619,6 +1620,7 @@ func (h *Handler) Messages(c *gin.Context) { readErr = readSSEStreamWithContinuousRetryKeepalive(readCtx, resp.Body, func(sseEvent string, data []byte) bool { parsed := gjson.ParseBytes(data) eventType := normalizedUpstreamSSEEventType(sseEvent, data) + ApplyCodexUsageEventFrame(h.store, account, eventType, parsed, time.Now()) if eventType == "error" { terminalFailurePayload = terminalUpstreamErrorPayload(data) gotTerminal = true diff --git a/proxy/responses_ws.go b/proxy/responses_ws.go index 66314bdf0..3edc18be7 100644 --- a/proxy/responses_ws.go +++ b/proxy/responses_ws.go @@ -1234,6 +1234,7 @@ func (h *Handler) streamResponsesWSUpstream( outputCollector.Add(data) parsed := gjson.ParseBytes(data) eventType := normalizedUpstreamSSEEventType(sseEvent, data) + ApplyCodexUsageEventFrame(h.store, account, eventType, parsed, time.Now()) eventType, data, parsed = rewriteEmptyIncompleteTerminal(emptyIncomplete, eventType, data, parsed) clientData := data if options != nil && options.transformClientEvent != nil { diff --git a/proxy/wsrelay/metadata_observer.go b/proxy/wsrelay/metadata_observer.go new file mode 100644 index 000000000..8eb0c407a --- /dev/null +++ b/proxy/wsrelay/metadata_observer.go @@ -0,0 +1,35 @@ +package wsrelay + +import ( + "sync/atomic" + + "github.com/codex2api/auth" +) + +// IdleMetadataObserver receives connection-level metadata frames that the read +// pump would otherwise drop (they arrive outside any request lease, typically +// shortly after a terminal frame). The observer must not block: it runs on the +// sole raw reader goroutine. +type IdleMetadataObserver func(account *auth.Account, eventType string, payload []byte) + +var idleMetadataObserver atomic.Pointer[IdleMetadataObserver] + +// SetIdleMetadataObserver installs the connection-level metadata observer. +// Passing nil removes it. The proxy package registers a usage-only observer so +// post-turn rate-limit frames can update the account snapshot without being +// routed to a request. +func SetIdleMetadataObserver(observer IdleMetadataObserver) { + if observer == nil { + idleMetadataObserver.Store(nil) + return + } + idleMetadataObserver.Store(&observer) +} + +func notifyIdleMetadataObserver(account *auth.Account, eventType string, payload []byte) { + observer := idleMetadataObserver.Load() + if observer == nil { + return + } + (*observer)(account, eventType, payload) +} diff --git a/proxy/wsrelay/read_pump.go b/proxy/wsrelay/read_pump.go index 489662148..f8fa50cff 100644 --- a/proxy/wsrelay/read_pump.go +++ b/proxy/wsrelay/read_pump.go @@ -203,6 +203,9 @@ func (wc *WsConnection) runReadPump() { // 只在"丢弃元数据"与"销毁连接"之间裁决，不延长业务空闲窗口。 eventType := strings.TrimSpace(gjson.GetBytes(payload, "type").String()) if isIdleDroppableMetadataFrame(eventType) { + // 连接级元数据帧不属于任何请求，但其中的 codex.rate_limits 是 + // 上一轮结束后的权威用量信号；在丢弃前交给只做内存更新的 observer。 + notifyIdleMetadataObserver(wc.account, eventType, payload) wc.endIdleSniff() wc.touchInbound() continue diff --git a/proxy/wsrelay/read_pump_test.go b/proxy/wsrelay/read_pump_test.go index 0d16f1b98..00ddbbae1 100644 --- a/proxy/wsrelay/read_pump_test.go +++ b/proxy/wsrelay/read_pump_test.go @@ -28,8 +28,18 @@ func newReadPumpTestConnection(t *testing.T, serve func(*websocket.Conn)) (*Mana return newReadPumpTestConnectionWithUpgrader(t, websocket.Upgrader{}, serve) } +func newReadPumpTestConnectionWithAccount(t *testing.T, serve func(*websocket.Conn), account *auth.Account) (*Manager, *WsConnection) { + t.Helper() + return newReadPumpTestConnectionWithUpgraderAndAccount(t, websocket.Upgrader{}, serve, account) +} + func newReadPumpTestConnectionWithUpgrader(t *testing.T, upgrader websocket.Upgrader, serve func(*websocket.Conn)) (*Manager, *WsConnection) { t.Helper() + return newReadPumpTestConnectionWithUpgraderAndAccount(t, upgrader, serve, nil) +} + +func newReadPumpTestConnectionWithUpgraderAndAccount(t *testing.T, upgrader websocket.Upgrader, serve func(*websocket.Conn), account *auth.Account) (*Manager, *WsConnection) { + t.Helper() ready := make(chan struct{}) upgrader.CheckOrigin = func(*http.Request) bool { return true } @@ -56,6 +66,7 @@ func newReadPumpTestConnectionWithUpgrader(t *testing.T, upgrader websocket.Upgr session := NewSession(1, manager) session.SetConnected(true) wc := NewWsConnection(conn, session, wsURL) + wc.account = account wc.PoolKey = key wc.onReadFailure = manager.DiscardConnection wc.installControlHandlers() @@ -1365,6 +1376,51 @@ func TestReadPumpDropsIdleMetadataFrame(t *testing.T) { } } +func TestReadPumpNotifiesIdleMetadataObserver(t *testing.T) { + got := make(chan struct { + accountID int64 + eventType string + payload string + }, 1) + SetIdleMetadataObserver(func(account *auth.Account, eventType string, payload []byte) { + var accountID int64 + if account != nil { + accountID = account.DBID + } + got <- struct { + accountID int64 + eventType string + payload string + }{accountID: accountID, eventType: eventType, payload: string(payload)} + }) + t.Cleanup(func() { SetIdleMetadataObserver(nil) }) + + holdOpen := make(chan struct{}) + t.Cleanup(func() { close(holdOpen) }) + newReadPumpTestConnectionWithAccount(t, func(conn *websocket.Conn) { + if err := conn.WriteMessage(websocket.TextMessage, []byte(`{"type":"codex.rate_limits","rate_limits":{"primary":{"used_percent":42}}}`)); err != nil { + t.Errorf("write idle metadata frame: %v", err) + return + } + <-holdOpen + }, &auth.Account{DBID: 42}) + + select { + case event := <-got: + if event.accountID != 42 { + t.Fatalf("observer account = %d, want 42", event.accountID) + } + if event.eventType != "codex.rate_limits" { + t.Fatalf("observer event type = %q", event.eventType) + } + if !strings.Contains(event.payload, `"used_percent":42`) { + t.Fatalf("observer payload = %q", event.payload) + } + case <-time.After(readPumpTestTimeout): + t.Fatal("idle metadata observer was not notified") + } +} + func TestReadPumpReusesConnectionAcrossLeases(t *testing.T) { holdOpen := make(chan struct{}) t.Cleanup(func() { close(holdOpen) })