From bb9531a6fecb6e6d9a52bf340ebdcbdfe442a775 Mon Sep 17 00:00:00 2001 From: Delicious233 <101502465+DeliciousBuding@users.noreply.github.com> Date: Sun, 18 Oct 2026 14:08:43 +0800 Subject: [PATCH 01/24] feat(proxy): ingest Codex WS usage frames for realtime quota state The persistent WebSocket transport carries quota data in per-turn codex.rate_limits and codex.response.metadata frames instead of HTTP headers, so account usage only advanced on WHAM probes. Parse those frames on every Codex relay path (Responses, Chat Completions, downstream WebSocket) and apply them through the existing header sync path. - Map primary/secondary windows, plan type and the compact credits snapshot to the same synthetic headers used by HTTP; support the official reset_at shape and keep the legacy reset-after-seconds family. Do not clear an absent window from a partial frame. - Ignore metered-limit frames (spark/reserve) so they cannot overwrite the main 5h/7d windows. - Bound persistence writes per account (default 30s, CODEX_WS_USAGE_FRAME_PERSIST_SECONDS); in-memory snapshots still update on every frame. Threshold crossings, 5pp moves and 100% force a write. - Observe connection-level metadata frames that arrive after a terminal (previously dropped by the read pump) through a non-blocking in-memory observer. - Expose runtime usage/status in /api/admin/accounts/live and merge it into the 1s account poll so cards update without a list reload. --- .env.example | 4 + admin/account_live.go | 51 +++- admin/account_live_test.go | 55 ++++ auth/premium_rate_limit.go | 16 +- auth/store.go | 159 ++++++++++- auth/usage_frame_test.go | 64 +++++ frontend/src/hooks/useAccountLiveState.ts | 48 +--- frontend/src/lib/accountLiveState.test.mjs | 98 +++++++ frontend/src/lib/accountLiveState.ts | 104 ++++++++ frontend/src/types.ts | 9 + main.go | 8 + proxy/codex_usage_frames.go | 263 +++++++++++++++++++ proxy/codex_usage_frames_integration_test.go | 71 +++++ proxy/codex_usage_frames_test.go | 263 +++++++++++++++++++ proxy/handler.go | 109 ++++++-- proxy/responses_ws.go | 1 + proxy/wsrelay/metadata_observer.go | 35 +++ proxy/wsrelay/read_pump.go | 3 + proxy/wsrelay/read_pump_test.go | 56 ++++ 19 files changed, 1341 insertions(+), 76 deletions(-) create mode 100644 auth/usage_frame_test.go create mode 100644 frontend/src/lib/accountLiveState.test.mjs create mode 100644 frontend/src/lib/accountLiveState.ts create mode 100644 proxy/codex_usage_frames.go create mode 100644 proxy/codex_usage_frames_integration_test.go create mode 100644 proxy/codex_usage_frames_test.go create mode 100644 proxy/wsrelay/metadata_observer.go diff --git a/.env.example b/.env.example index 6e8911fa0..c962178d5 100644 --- a/.env.example +++ b/.env.example @@ -50,6 +50,10 @@ 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 + # 旧版 WebSocket 开关（兼容保留；推荐改用 CODEX_UPSTREAM_TRANSPORT=ws） # USE_WEBSOCKET=true diff --git a/admin/account_live.go b/admin/account_live.go index 4adca4d4d..4a62c984c 100644 --- a/admin/account_live.go +++ b/admin/account_live.go @@ -2,6 +2,7 @@ package admin import ( "net/http" + "time" "github.com/gin-gonic/gin" ) @@ -15,11 +16,27 @@ 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"` + 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"` } // 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 +54,45 @@ 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() + 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) + } + 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..06c10d8f6 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,57 @@ 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) + 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) + } +} diff --git a/auth/premium_rate_limit.go b/auth/premium_rate_limit.go index 70885b7be..9a4508192 100644 --- a/auth/premium_rate_limit.go +++ b/auth/premium_rate_limit.go @@ -133,6 +133,12 @@ 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) { if acc == nil || s == nil { return } @@ -142,20 +148,22 @@ func (s *Store) PersistUsageSnapshot5hOnly(acc *Account) { return } - 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 { + if !persist || s.db == nil { return } ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second) defer cancel() - if err := s.db.UpdateUsageSnapshot5h(ctx, acc.DBID, pct5h, reset5hAt, updatedAt); err != nil { + if err := s.db.UpdateUsageSnapshot5h(ctx, acc.DBID, pct5h, reset5hAt, observedAt); err != nil { log.Printf("[账号 %d] 持久化 5h 用量快照失败: %v", acc.DBID, err) } } diff --git a/auth/store.go b/auth/store.go index 88b650d12..87d3bbed0 100644 --- a/auth/store.go +++ b/auth/store.go @@ -298,6 +298,12 @@ type Account struct { Reset5hAt time.Time // 5h 窗口重置时间 UsageUpdatedAt time.Time // 7d 用量快照刷新时间 UsageUpdatedAt5h time.Time // 5h 用量快照刷新时间 + // usageFramePersisted* 记录 WS 逐帧观测最近一次落库的时间与百分比。 + // 内存快照每帧都实时更新；落库只为重启恢复，大号池下按时间/变化幅度节流， + // 避免把每个上游响应都变成一次数据库写入。 + usageFramePersistedAt time.Time + usageFramePersisted5h float64 + usageFramePersisted7d float64 // activated5hResetAt 是已经为其发送过「开窗」最小 /responses 的那个 Reset5hAt。 // 每个观测到的 5h 窗口最多激活一次（issue #581）。 activated5hResetAt time.Time @@ -2418,6 +2424,34 @@ func (a *Account) ApplyUsageObservation(observedAt time.Time, apply func()) bool return true } +// TryApplyUsageObservation is the non-blocking variant used by connection-level +// WS frame observers. It skips the observation when another goroutine is +// already applying a usage observation (for example a preflight frame doing a +// throttled DB write), so the WebSocket read pump is never stalled. +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 +2459,110 @@ 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. +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.Lock() + defer a.mu.Unlock() + + 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 + } + a.usageFramePersistedAt = now + if has5h { + a.usageFramePersisted5h = pct5h + } + if has7d { + a.usageFramePersisted7d = pct7d + } + return true +} + +// 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 + UsagePercent5h float64 + UsagePercent5hValid bool + Reset5hAt time.Time + UsageUpdatedAt5h time.Time + UsagePercent7d float64 + UsagePercent7dValid bool + Reset7dAt time.Time + UsageUpdatedAt time.Time +} + +// 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), + UsagePercent5h: a.UsagePercent5h, + UsagePercent5hValid: a.UsagePercent5hValid, + Reset5hAt: a.Reset5hAt, + UsageUpdatedAt5h: a.UsageUpdatedAt5h, + UsagePercent7d: a.UsagePercent7d, + UsagePercent7dValid: a.UsagePercent7dValid, + Reset7dAt: a.Reset7dAt, + UsageUpdatedAt: a.UsageUpdatedAt, + } +} + // Mark5hWindowActivated 记录已经为哪个 Reset5hAt 发送过开窗请求。 func (a *Account) Mark5hWindowActivated(resetAt time.Time) { if a == nil || resetAt.IsZero() { @@ -10890,15 +11028,24 @@ 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 { + 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) { + if acc == nil || s == nil { return } + if observedAt.IsZero() { + observedAt = time.Now() + } - now := time.Now() - acc.SetUsageSnapshot(pct7d, now) + acc.SetUsageSnapshot(pct7d, observedAt) s.fastSchedulerUpdate(acc) - if s.db == nil { + if !persist || s.db == nil { return } @@ -10908,7 +11055,7 @@ 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 { + if err := s.db.UpdateUsageSnapshotFull(ctx, acc.DBID, pct7d, reset7dAt, pct5h, reset5hAt, observedAt, acc.GetUsageUpdatedAt5h()); err != nil { log.Printf("[账号 %d] 持久化用量快照失败: %v", acc.DBID, err) } return @@ -10916,7 +11063,7 @@ func (s *Store) PersistUsageSnapshot(acc *Account, pct7d float64) { ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second) defer cancel() - if err := s.db.UpdateUsageSnapshot(ctx, acc.DBID, pct7d, now); err != nil { + if err := s.db.UpdateUsageSnapshot(ctx, acc.DBID, pct7d, observedAt); err != nil { log.Printf("[账号 %d] 持久化用量快照失败: %v", acc.DBID, err) } } diff --git a/auth/usage_frame_test.go b/auth/usage_frame_test.go new file mode 100644 index 000000000..ae1d1969d --- /dev/null +++ b/auth/usage_frame_test.go @@ -0,0 +1,64 @@ +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") + } + 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") + } + 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 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") + } + 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) + } +} diff --git a/frontend/src/hooks/useAccountLiveState.ts b/frontend/src/hooks/useAccountLiveState.ts index b8b2bf825..8c7c52151 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[], diff --git a/frontend/src/lib/accountLiveState.test.mjs b/frontend/src/lib/accountLiveState.test.mjs new file mode 100644 index 000000000..84da73712 --- /dev/null +++ b/frontend/src/lib/accountLiveState.test.mjs @@ -0,0 +1,98 @@ +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, + 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, + 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); +}); diff --git a/frontend/src/lib/accountLiveState.ts b/frontend/src/lib/accountLiveState.ts new file mode 100644 index 000000000..b44896447 --- /dev/null +++ b/frontend/src/lib/accountLiveState.ts @@ -0,0 +1,104 @@ +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 + 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) + + 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 && + 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, + 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/types.ts b/frontend/src/types.ts index d62a497f4..b1bf31417 100644 --- a/frontend/src/types.ts +++ b/frontend/src/types.ts @@ -628,6 +628,15 @@ export interface AccountLiveStateResponse { // 调度器当前实际执行的并发上限与配置值；旧后端不返回时保留列表里的值。 dynamic_concurrency_limit?: number base_concurrency_effective?: number + // 运行时用量投影：旧后端缺 key 时保留列表值；显式 null 表示窗口已失效。 + status?: AccountStatus + using_credits?: boolean + 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 }> session_slot_buffer_enabled: boolean } diff --git a/main.go b/main.go index 998042a78..8e2c24270 100644 --- a/main.go +++ b/main.go @@ -445,6 +445,14 @@ func main() { // 注册 WebSocket 执行函数（避免 proxy ↔ wsrelay 循环依赖） proxy.WebsocketExecuteFunc = wsrelay.ExecuteRequestWebsocket + // 连接级（终态后补发）的用量帧也交给 proxy 更新内存快照；observer 必须非阻塞。 + wsrelay.SetIdleMetadataObserver(proxy.ApplyCodexIdleUsageFrame) + // 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) + } + } // 注册 Agent Identity task 确保函数（proxy 无 Store 引用，启动时注入） proxy.EnsureCodexAgentIdentityTaskFunc = store.EnsureCodexAgentIdentityTask diff --git a/proxy/codex_usage_frames.go b/proxy/codex_usage_frames.go new file mode 100644 index 000000000..4564c119c --- /dev/null +++ b/proxy/codex_usage_frames.go @@ -0,0 +1,263 @@ +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" +) + +// 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 +}() + +// 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 this adds no +// upstream requests and no new request headers. +func ApplyCodexUsageEventFrame(store *auth.Store, account *auth.Account, eventType string, frame gjson.Result, observedAt time.Time) { + if account == nil { + 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) +} + +func codexRateLimitsHeaders(frame gjson.Result, observedAt time.Time) http.Header { + if name := firstCodexFrameString(frame, "metered_limit_name", "limit_name"); name != "" && !isDefaultCodexLimitName(name) { + return nil + } + headers := make(http.Header) + if plan := strings.TrimSpace(frame.Get("plan_type").String()); plan != "" { + 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 := strings.TrimSpace(credits.Get("balance").String()); balance != "" { + headers.Set("x-codex-credits-balance", balance) + } + } + } + return headers +} + +// applyCodexResponseMetadataFrame adopts the x-codex-* headers embedded in a +// codex.response.metadata frame. Other headers are intentionally dropped: some +// upstream metadata frames also carry cookies or request identifiers that must +// not be persisted or logged. +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) +} + +func codexResponseMetadataHeaders(frame gjson.Result, observedAt time.Time) http.Header { + headers := make(http.Header) + frame.Get("headers").ForEach(func(key, value gjson.Result) bool { + name := strings.ToLower(strings.TrimSpace(key.String())) + if !strings.HasPrefix(name, "x-codex-") { + 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 < 0 { + seconds = 0 + } + prefix := strings.TrimSuffix(name, "reset-at") + headers.Set(prefix+"reset-after-seconds", strconv.FormatInt(seconds, 10)) + } + return true + } + headers.Set(name, value.String()) + return true + }) + return headers +} + +func applyCodexUsageHeaders(store *auth.Store, account *auth.Account, headers http.Header, observedAt time.Time) { + if len(headers) == 0 { + return + } + 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. + syncCodexUsageState(store, account, resp, observedAt, codexUsageSyncOptions{ + persistUsageSnapshot: true, + clearAbsent5h: false, + shouldPersistUsage: func(pct5h float64, has5h bool, pct7d float64, has7d bool) bool { + return account.ShouldPersistUsageFrame(observedAt, pct5h, has5h, pct7d, has7d, codexUsageFramePersistInterval()) + }, + }) +} + +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 { + for _, key := range []string{"resets_in_seconds", "reset_after_seconds"} { + if value := window.Get(key); value.Exists() { + if n := codexFrameInt64(value); n > 0 { + return strconv.FormatInt(n, 10) + } + } + } + for _, key := range []string{"reset_at", "resets_at"} { + if value := window.Get(key); value.Exists() { + epoch := codexFrameInt64(value) + if epoch <= 0 { + continue + } + seconds := epoch - observedAt.Unix() + if seconds < 0 { + seconds = 0 + } + return strconv.FormatInt(seconds, 10) + } + } + return "" +} + +func codexFrameInt64(value gjson.Result) int64 { + if value.Type == gjson.Number { + return value.Int() + } + n, err := strconv.ParseInt(strings.TrimSpace(value.String()), 10, 64) + if err != nil { + return 0 + } + 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 "" +} + +func isDefaultCodexLimitName(name string) bool { + normalized := strings.NewReplacer("_", "-", " ", "-").Replace(strings.ToLower(strings.TrimSpace(name))) + return normalized == "codex" || normalized == "gpt-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). It updates only the in-memory usage +// snapshot, with a non-blocking try-lock, so the read pump never waits on a DB +// write. Persistence and the 100% cooldown markings stay on the preflight/WHAM +// paths; this path exists to shrink the window between a turn finishing and +// the scheduler seeing its consumption. +func ApplyCodexIdleUsageFrame(account *auth.Account, eventType string, payload []byte) { + if account == nil || len(payload) == 0 { + return + } + observedAt := time.Now() + frame := gjson.ParseBytes(payload) + var headers http.Header + switch strings.ToLower(strings.TrimSpace(eventType)) { + case codexRateLimitsFrameType: + headers = codexRateLimitsHeaders(frame, observedAt) + case codexResponseMetadataFrameType, responseMetadataFrameType: + headers = codexResponseMetadataHeaders(frame, observedAt) + default: + return + } + if len(headers) == 0 { + return + } + observation := parseCodexUsageHeaderObservation(&http.Response{Header: headers}) + if !observation.authoritative { + return + } + account.TryApplyUsageObservation(observedAt, func() { + // Never clear an absent 5h window from a partial connection-level frame. + applyCodexUsageHeaderObservation(nil, account, observation, observedAt, false) + }) +} diff --git a/proxy/codex_usage_frames_integration_test.go b/proxy/codex_usage_frames_integration_test.go new file mode 100644 index 000000000..69e193002 --- /dev/null +++ b/proxy/codex_usage_frames_integration_test.go @@ -0,0 +1,71 @@ +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) + } +} diff --git a/proxy/codex_usage_frames_test.go b/proxy/codex_usage_frames_test.go new file mode 100644 index 000000000..a275d50ab --- /dev/null +++ b/proxy/codex_usage_frames_test.go @@ -0,0 +1,263 @@ +package proxy + +import ( + "net/http" + "strconv" + "testing" + "time" + + "github.com/codex2api/auth" + "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) + } +} + +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", + "set-cookie":"must-not-be-used" + } + }`) + + 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 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 TestApplyCodexIdleUsageFrameUpdatesMemoryOnly(t *testing.T) { + acc := newUsageFrameTestAccount() + now := time.Now() + acc.SetUsageSnapshot5hAt(20, now.Add(time.Hour), now) + acc.SetUsagePercent7d(10) + + ApplyCodexIdleUsageFrame(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) + } +} diff --git a/proxy/handler.go b/proxy/handler.go index a0c7acb1a..8980485d6 100644 --- a/proxy/handler.go +++ b/proxy/handler.go @@ -4688,6 +4688,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 +5363,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 +5614,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 +7460,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 +7629,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 +8139,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 { @@ -8147,7 +8157,8 @@ func parseCodexWindowUsage(usedStr, windowStr, resetStr string) codexWindowUsage if _, err := fmt.Sscanf(windowStr, "%f", &windowMin); err != nil || windowMin <= 0 { return codexWindowUsage{} } - if strings.TrimSpace(resetStr) != "" { + hasReset := strings.TrimSpace(resetStr) != "" + if hasReset { if _, err := fmt.Sscanf(resetStr, "%f", &resetSec); err != nil { resetSec = 0 } @@ -8157,9 +8168,34 @@ 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 "" + } + var epoch int64 + if _, err := fmt.Sscanf(raw, "%d", &epoch); 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) { if resp == nil { return codexRateLimitWindowUnknown, time.Time{}, false @@ -8168,12 +8204,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 +8245,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 +8254,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,11 +8718,34 @@ 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 +} + +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 observedAt.IsZero() { + observedAt = time.Now() + } if store != nil { planHeader := resp.Header.Get("x-codex-plan-type") store.UpdateAccountPlanType(account, planHeader) @@ -8732,21 +8791,30 @@ func SyncCodexUsageState(store *auth.Store, account *auth.Account, resp *http.Re usageApplied := false if observation.authoritative { usageApplied = account.ApplyUsageObservation(observedAt, func() { - parsed := applyCodexUsageHeaderObservation(store, account, observation, observedAt) + 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, + ) + } if result.HasUsage7d { - store.PersistUsageSnapshot(account, result.UsagePct7d) + store.PublishUsageSnapshot(account, result.UsagePct7d, observedAt, persistUsage) if result.UsagePct7d >= 100 { result.Usage7dRateLimited = store.MarkUsage7dRateLimited(account) } } else if result.Used5hHeaders { - store.PersistUsageSnapshot5hOnly(account) - result.Persisted5hOnly = true + store.PublishUsageSnapshot5hOnly(account, observedAt, persistUsage) + result.Persisted5hOnly = persistUsage } }) } @@ -8791,7 +8859,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 +8896,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,18 +8923,23 @@ 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) + resetAt := account.GetReset5hAt() + if observation.w5h.hasReset { + resetAt = observedAt.Add(time.Duration(observation.w5h.resetSec) * time.Second) + } account.SetUsageSnapshot5hAt(observation.w5h.usedPct, resetAt, observedAt) - } else if observation.authoritative { + } 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) + if observation.w7d.hasReset { + resetAt := observedAt.Add(time.Duration(observation.w7d.resetSec) * time.Second) + account.SetReset7dAt(resetAt) + } account.SetWindow7dSeconds(int64(observation.w7d.windowMin * 60)) account.SetUsagePercent7d(observation.w7d.usedPct) out.usagePct7d = observation.w7d.usedPct 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) }) From 205a06f80aa27a763517e7eef8c9239d608bb40e Mon Sep 17 00:00:00 2001 From: Delicious233 <101502465+DeliciousBuding@users.noreply.github.com> Date: Sun, 18 Oct 2026 14:15:41 +0800 Subject: [PATCH 02/24] docs: document WS usage frame ingestion and persist interval Record CODEX_WS_USAGE_FRAME_PERSIST_SECONDS and the per-turn rate-limit frame behavior under the Codex upstream stability settings so operators know auto-pause reacts between WHAM probes without extra upstream calls. --- docs/CONFIGURATION.md | 3 +++ 1 file changed, 3 insertions(+) diff --git a/docs/CONFIGURATION.md b/docs/CONFIGURATION.md index 7f07ee404..556adf52f 100644 --- a/docs/CONFIGURATION.md +++ b/docs/CONFIGURATION.md @@ -87,6 +87,7 @@ 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 逐帧用量快照落库的最小间隔（秒），正整数，重启生效。内存快照每帧更新；跨自动暂停阈值、5pp 变化或 100% 时仍强制落库。大号池可按数据库承压调大 | | `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 | @@ -99,6 +100,8 @@ Codex2API 采用三层配置架构： | `DOWNSTREAM_WS_KEEPALIVE_INTERVAL` | 否 | `45s` | 下游 WebSocket Ping 周期，使用 Go duration；`0` 关闭。覆盖 Responses、Realtime 与 Live Sideband | > `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 探针周期即可生效。这些帧随对话本身下发，不产生额外上游请求。 ### 数据库配置 From 8cdc6631ab500340fbfb3d4d33381942d38c2903 Mon Sep 17 00:00:00 2001 From: Delicious233 <101502465+DeliciousBuding@users.noreply.github.com> Date: Sun, 18 Oct 2026 14:54:04 +0800 Subject: [PATCH 03/24] fix(proxy): harden WS usage frame ingestion - preserve a stored workspace hard stop when a frame carries no rate_limit_reached_type - ignore metered limits such as gpt-codex instead of letting them overwrite 5h/7d - unwrap array-valued metadata headers and reject control characters and out-of-range percentages - keep usage monotonic within the same reset window and bound a missing reset to the window length - advance the persistence throttle only after the DB write succeeds - dedupe repeated cooldown marks for the same reason and reset window - add CODEX_WS_USAGE_FRAME_INGEST as a deployment-level kill switch --- .env.example | 3 + auth/premium_rate_limit.go | 21 ++- auth/store.go | 99 ++++++++-- auth/usage_frame_test.go | 17 ++ docs/CONFIGURATION.md | 1 + main.go | 6 + proxy/codex_usage_frames.go | 147 ++++++++++++--- proxy/codex_usage_frames_test.go | 299 +++++++++++++++++++++++++++++++ proxy/handler.go | 99 ++++++++-- proxy/handler_anthropic.go | 2 + 10 files changed, 638 insertions(+), 56 deletions(-) diff --git a/.env.example b/.env.example index c962178d5..d2094cb35 100644 --- a/.env.example +++ b/.env.example @@ -54,6 +54,9 @@ CODEX_PORT=8080 # 跨阈值、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/auth/premium_rate_limit.go b/auth/premium_rate_limit.go index 9a4508192..c2c31f7a6 100644 --- a/auth/premium_rate_limit.go +++ b/auth/premium_rate_limit.go @@ -138,14 +138,14 @@ func (s *Store) PersistUsageSnapshot5hOnly(acc *Account) { // 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) { +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 } if observedAt.IsZero() { @@ -158,14 +158,16 @@ func (s *Store) PublishUsageSnapshot5hOnly(acc *Account, observedAt time.Time, p s.fastSchedulerUpdate(acc) if !persist || s.db == nil { - return + return false } ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second) defer cancel() if err := s.db.UpdateUsageSnapshot5h(ctx, acc.DBID, pct5h, reset5hAt, observedAt); err != nil { log.Printf("[账号 %d] 持久化 5h 用量快照失败: %v", acc.DBID, err) + return false } + return true } // Persist5hWindowActivated 把「已为哪个 Reset5hAt 发过开窗请求」写入 credentials，重启后不重复打。 @@ -323,6 +325,17 @@ func (s *Store) markPremium5hRateLimited(acc *Account, resetAt, observedAt time. resetAt = now.Add(premium5hFallbackWindow) } + // Re-observing the same exhausted window must not write the cooldown again. + acc.mu.RLock() + alreadyLimited := acc.Status == StatusCooldown && + acc.CooldownReason == cooldownReason && + now.Before(acc.CooldownUtil) && + cooldownMatchesReset(acc.CooldownUtil, resetAt) + acc.mu.RUnlock() + if alreadyLimited { + return + } + acc.mu.Lock() acc.UsagePercent5h = 100 acc.UsagePercent5hValid = true diff --git a/auth/store.go b/auth/store.go index 87d3bbed0..d58735728 100644 --- a/auth/store.go +++ b/auth/store.go @@ -304,6 +304,10 @@ type Account struct { usageFramePersistedAt time.Time usageFramePersisted5h float64 usageFramePersisted7d float64 + // usageFrameLastAt/Kind 记录最近一次被应用的 WS 配额帧，仅供管理端 + // 验证帧摄入是否工作，不参与调度。 + usageFrameLastAt time.Time + usageFrameLastKind string // activated5hResetAt 是已经为其发送过「开窗」最小 /responses 的那个 Reset5hAt。 // 每个观测到的 5h 窗口最多激活一次（issue #581）。 activated5hResetAt time.Time @@ -2351,6 +2355,21 @@ 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 +} + // 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. @@ -2367,15 +2386,27 @@ func (s *Store) MarkUsage7dRateLimited(acc *Account) bool { return false } + now := time.Now() duration := 7 * 24 * time.Hour - if resetAt := acc.GetReset7dAt(); !resetAt.IsZero() { - untilReset := time.Until(resetAt) + resetAt := acc.GetReset7dAt() + if !resetAt.IsZero() { + untilReset := resetAt.Sub(now) if untilReset <= 0 { return false } duration = untilReset } + acc.mu.RLock() + alreadyLimited := acc.Status == StatusCooldown && + (acc.CooldownReason == "rate_limited" || acc.CooldownReason == "rate_limited_7d") && + now.Before(acc.CooldownUtil) && + cooldownMatchesReset(acc.CooldownUtil, resetAt) + acc.mu.RUnlock() + if alreadyLimited { + return true + } + s.MarkCooldown(acc, duration, "rate_limited") return true } @@ -2471,13 +2502,15 @@ 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. +// 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.Lock() - defer a.mu.Unlock() + a.mu.RLock() + defer a.mu.RUnlock() if interval <= 0 { interval = UsageFramePersistInterval @@ -2514,6 +2547,16 @@ func (a *Account) ShouldPersistUsageFrame(now time.Time, pct5h float64, has5h bo 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 @@ -2521,7 +2564,22 @@ func (a *Account) ShouldPersistUsageFrame(now time.Time, pct5h float64, has5h bo if has7d { a.usageFramePersisted7d = pct7d } - return true +} + +// 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() + a.usageFrameLastAt = at + a.usageFrameLastKind = strings.TrimSpace(kind) } // AccountLiveUsageSnapshot is a single-lock projection for the admin live poll. @@ -2538,6 +2596,10 @@ type AccountLiveUsageSnapshot struct { 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 @@ -2560,6 +2622,8 @@ func (a *Account) GetAccountLiveUsageSnapshot() AccountLiveUsageSnapshot { UsagePercent7dValid: a.UsagePercent7dValid, Reset7dAt: a.Reset7dAt, UsageUpdatedAt: a.UsageUpdatedAt, + UsageFrameLastAt: a.usageFrameLastAt, + UsageFrameLastKind: a.usageFrameLastKind, } } @@ -2708,6 +2772,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) { @@ -11034,9 +11110,9 @@ func (s *Store) PersistUsageSnapshot(acc *Account, pct7d float64) { // 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) { +func (s *Store) PublishUsageSnapshot(acc *Account, pct7d float64, observedAt time.Time, persist bool) bool { if acc == nil || s == nil { - return + return false } if observedAt.IsZero() { observedAt = time.Now() @@ -11046,7 +11122,7 @@ func (s *Store) PublishUsageSnapshot(acc *Account, pct7d float64, observedAt tim s.fastSchedulerUpdate(acc) if !persist || s.db == nil { - return + return false } // 如果有 5h 数据，使用完整存储 @@ -11057,15 +11133,18 @@ func (s *Store) PublishUsageSnapshot(acc *Account, pct7d float64, observedAt tim defer cancel() if err := s.db.UpdateUsageSnapshotFull(ctx, acc.DBID, pct7d, reset7dAt, pct5h, reset5hAt, observedAt, acc.GetUsageUpdatedAt5h()); 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, observedAt); err != nil { log.Printf("[账号 %d] 持久化用量快照失败: %v", acc.DBID, err) + return false } + return true } // UpdateAccountSubscriptionExpiresAt persists the latest subscription expiration observed from upstream. diff --git a/auth/usage_frame_test.go b/auth/usage_frame_test.go index ae1d1969d..c56a2e7d3 100644 --- a/auth/usage_frame_test.go +++ b/auth/usage_frame_test.go @@ -13,12 +13,14 @@ func TestShouldPersistUsageFrameThrottlesAndForces(t *testing.T) { 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") } @@ -27,6 +29,20 @@ func TestShouldPersistUsageFrameThrottlesAndForces(t *testing.T) { } } +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) @@ -35,6 +51,7 @@ func TestShouldPersistUsageFrameForcesThresholdCrossing(t *testing.T) { 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") } diff --git a/docs/CONFIGURATION.md b/docs/CONFIGURATION.md index 556adf52f..f22a1d9d6 100644 --- a/docs/CONFIGURATION.md +++ b/docs/CONFIGURATION.md @@ -88,6 +88,7 @@ Codex2API 采用三层配置架构： | `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 逐帧用量快照落库的最小间隔（秒），正整数，重启生效。内存快照每帧更新；跨自动暂停阈值、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 | diff --git a/main.go b/main.go index 8e2c24270..3961f726b 100644 --- a/main.go +++ b/main.go @@ -453,6 +453,12 @@ func main() { 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/proxy/codex_usage_frames.go b/proxy/codex_usage_frames.go index 4564c119c..102594a4e 100644 --- a/proxy/codex_usage_frames.go +++ b/proxy/codex_usage_frames.go @@ -29,6 +29,21 @@ var codexUsageFramePersistIntervalNanos = func() *atomic.Int64 { 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) { @@ -49,7 +64,7 @@ func codexUsageFramePersistInterval() time.Duration { // already on the wire (the gateway forwards them downstream), so this adds no // upstream requests and no new request headers. func ApplyCodexUsageEventFrame(store *auth.Store, account *auth.Account, eventType string, frame gjson.Result, observedAt time.Time) { - if account == nil { + if account == nil || !codexUsageFrameIngestEnabled.Load() { return } switch strings.ToLower(strings.TrimSpace(eventType)) { @@ -69,15 +84,16 @@ func applyCodexRateLimitsFrame(store *auth.Store, account *auth.Account, frame g if len(headers) == 0 { return } + account.MarkUsageFrameObserved(codexRateLimitsFrameType, observedAt) applyCodexUsageHeaders(store, account, headers, observedAt) } func codexRateLimitsHeaders(frame gjson.Result, observedAt time.Time) http.Header { - if name := firstCodexFrameString(frame, "metered_limit_name", "limit_name"); name != "" && !isDefaultCodexLimitName(name) { + if name := firstCodexFrameString(frame, "metered_limit_name", "limit_name", "limit_id"); name != "" && !isDefaultCodexLimitName(name) { return nil } headers := make(http.Header) - if plan := strings.TrimSpace(frame.Get("plan_type").String()); plan != "" { + 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) @@ -91,7 +107,7 @@ func codexRateLimitsHeaders(frame gjson.Result, observedAt time.Time) http.Heade if hasCredits.Exists() && unlimited.Exists() { headers.Set("x-codex-credits-has-credits", hasCredits.String()) headers.Set("x-codex-credits-unlimited", unlimited.String()) - if balance := strings.TrimSpace(credits.Get("balance").String()); balance != "" { + if balance, ok := codexFrameText(credits.Get("balance").String(), 128); ok { headers.Set("x-codex-credits-balance", balance) } } @@ -99,15 +115,33 @@ func codexRateLimitsHeaders(frame gjson.Result, observedAt time.Time) http.Heade return headers } -// applyCodexResponseMetadataFrame adopts the x-codex-* headers embedded in a -// codex.response.metadata frame. Other headers are intentionally dropped: some -// upstream metadata frames also carry cookies or request identifiers that must -// not be persisted or logged. +// 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 } + account.MarkUsageFrameObserved(codexResponseMetadataFrameType, observedAt) applyCodexUsageHeaders(store, account, headers, observedAt) } @@ -115,7 +149,22 @@ func codexResponseMetadataHeaders(frame gjson.Result, observedAt time.Time) http headers := make(http.Header) frame.Get("headers").ForEach(func(key, value gjson.Result) bool { name := strings.ToLower(strings.TrimSpace(key.String())) - if !strings.HasPrefix(name, "x-codex-") { + 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 @@ -134,16 +183,48 @@ func codexResponseMetadataHeaders(frame gjson.Result, observedAt time.Time) http } return true } - headers.Set(name, value.String()) + raw, ok := codexFrameText(value.String(), 256) + if !ok { + return true + } + headers.Set(name, raw) return true }) 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) { if len(headers) == 0 { return } + // 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. + if headers.Get("x-codex-rate-limit-reached-type") == "" { + if current := account.GetRateLimitReachedType(); current != "" { + headers.Set("x-codex-rate-limit-reached-type", current) + } + } 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 @@ -181,15 +262,21 @@ func setCodexFrameWindow(headers http.Header, prefix string, window gjson.Result func codexFrameResetSeconds(window gjson.Result, observedAt time.Time) string { for _, key := range []string{"resets_in_seconds", "reset_after_seconds"} { if value := window.Get(key); value.Exists() { - if n := codexFrameInt64(value); n > 0 { + if n, ok := codexFrameInt64Value(value); ok { + // Zero/negative means "resets now", matching the absolute form + // where a past epoch is clamped to zero. Only a missing or + // unparseable field keeps the previous reset. + if n < 0 { + n = 0 + } return strconv.FormatInt(n, 10) } } } for _, key := range []string{"reset_at", "resets_at"} { if value := window.Get(key); value.Exists() { - epoch := codexFrameInt64(value) - if epoch <= 0 { + epoch, ok := codexFrameInt64Value(value) + if !ok || epoch <= 0 { continue } seconds := epoch - observedAt.Unix() @@ -202,14 +289,19 @@ func codexFrameResetSeconds(window gjson.Result, observedAt time.Time) string { return "" } -func codexFrameInt64(value gjson.Result) int64 { +func codexFrameInt64Value(value gjson.Result) (int64, bool) { if value.Type == gjson.Number { - return value.Int() + return value.Int(), true } n, err := strconv.ParseInt(strings.TrimSpace(value.String()), 10, 64) if err != nil { - return 0 + return 0, false } + return n, true +} + +func codexFrameInt64(value gjson.Result) int64 { + n, _ := codexFrameInt64Value(value) return n } @@ -222,9 +314,10 @@ func firstCodexFrameString(frame gjson.Result, paths ...string) string { return "" } +var codexLimitNameSeparators = strings.NewReplacer("_", "-", " ", "-") + func isDefaultCodexLimitName(name string) bool { - normalized := strings.NewReplacer("_", "-", " ", "-").Replace(strings.ToLower(strings.TrimSpace(name))) - return normalized == "codex" || normalized == "gpt-codex" + return codexLimitNameSeparators.Replace(strings.ToLower(strings.TrimSpace(name))) == "codex" } // ApplyCodexIdleUsageFrame is the observer for connection-level metadata frames @@ -235,19 +328,26 @@ func isDefaultCodexLimitName(name string) bool { // paths; this path exists to shrink the window between a turn finishing and // the scheduler seeing its consumption. func ApplyCodexIdleUsageFrame(account *auth.Account, eventType string, payload []byte) { - if account == nil || len(payload) == 0 { + if account == nil || 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 strings.ToLower(strings.TrimSpace(eventType)) { + switch normalizedType { case codexRateLimitsFrameType: headers = codexRateLimitsHeaders(frame, observedAt) - case codexResponseMetadataFrameType, responseMetadataFrameType: - headers = codexResponseMetadataHeaders(frame, observedAt) default: - return + headers = codexResponseMetadataHeaders(frame, observedAt) } if len(headers) == 0 { return @@ -257,6 +357,7 @@ func ApplyCodexIdleUsageFrame(account *auth.Account, eventType string, payload [ return } account.TryApplyUsageObservation(observedAt, func() { + account.MarkUsageFrameObserved(normalizedType, observedAt) // Never clear an absent 5h window from a partial connection-level frame. applyCodexUsageHeaderObservation(nil, account, observation, observedAt, false) }) diff --git a/proxy/codex_usage_frames_test.go b/proxy/codex_usage_frames_test.go index a275d50ab..0f9bba411 100644 --- a/proxy/codex_usage_frames_test.go +++ b/proxy/codex_usage_frames_test.go @@ -7,6 +7,7 @@ import ( "time" "github.com/codex2api/auth" + "github.com/codex2api/database" "github.com/tidwall/gjson" ) @@ -95,10 +96,30 @@ func TestApplyCodexResponseMetadataFrameAdoptsUsageHeaders(t *testing.T) { "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() @@ -110,6 +131,120 @@ func TestApplyCodexResponseMetadataFrameAdoptsUsageHeaders(t *testing.T) { } } +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 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(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) @@ -261,3 +396,167 @@ func TestParseCodexUsageHeaderObservationAcceptsResetAt(t *testing.T) { 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 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 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) + } +} diff --git a/proxy/handler.go b/proxy/handler.go index 8980485d6..fb2f97593 100644 --- a/proxy/handler.go +++ b/proxy/handler.go @@ -180,15 +180,19 @@ 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 UsageWindowLimitsIgnored bool // Cleared5h 表示本次同步因上游未返回 5h 窗口而清除了本地陈旧 5h 快照（issue #382）。 Cleared5h bool @@ -8154,6 +8158,11 @@ func parseCodexWindowUsage(usedStr, windowStr, resetStr string) codexWindowUsage if _, err := fmt.Sscanf(usedStr, "%f", &usedPct); err != nil { return codexWindowUsage{} } + // 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 _, err := fmt.Sscanf(windowStr, "%f", &windowMin); err != nil || windowMin <= 0 { return codexWindowUsage{} } @@ -8807,15 +8816,25 @@ func syncCodexUsageState(store *auth.Store, account *auth.Account, resp *http.Re observation.w7d.valid, ) } + result.PersistUsage = persistUsage + persistedUsage := false if result.HasUsage7d { - store.PublishUsageSnapshot(account, result.UsagePct7d, observedAt, persistUsage) - if result.UsagePct7d >= 100 { + persistedUsage = store.PublishUsageSnapshot(account, result.UsagePct7d, observedAt, persistUsage) + if result.UsagePct7d >= 100 && persistUsage { result.Usage7dRateLimited = store.MarkUsage7dRateLimited(account) } } else if result.Used5hHeaders { - store.PublishUsageSnapshot5hOnly(account, observedAt, persistUsage) + 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) + } }) } @@ -8826,7 +8845,11 @@ func syncCodexUsageState(store *auth.Store, account *auth.Account, resp *http.Re } if 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 } @@ -8926,28 +8949,66 @@ func parseCodexUsageHeaderObservation(resp *http.Response) codexUsageHeaderObser func applyCodexUsageHeaderObservation(store *auth.Store, account *auth.Account, observation codexUsageHeaderObservation, observedAt time.Time, clearAbsent5h bool) codexUsageHeaderParseResult { out := codexUsageHeaderParseResult{} if observation.w5h.valid { - resetAt := account.GetReset5hAt() + existingReset := account.GetReset5hAt() + existingPct, existingValid := account.GetUsagePercent5h() + resetAt := existingReset if observation.w5h.hasReset { resetAt = observedAt.Add(time.Duration(observation.w5h.resetSec) * time.Second) - } - account.SetUsageSnapshot5hAt(observation.w5h.usedPct, resetAt, observedAt) + } else if resetAt.IsZero() { + // 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 && 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 { + existingReset := account.GetReset7dAt() + existingPct, existingValid := account.GetUsagePercent7d() + resetAt := existingReset if observation.w7d.hasReset { - resetAt := observedAt.Add(time.Duration(observation.w7d.resetSec) * time.Second) + resetAt = observedAt.Add(time.Duration(observation.w7d.resetSec) * time.Second) + account.SetReset7dAt(resetAt) + } else if resetAt.IsZero() { + resetAt = observedAt.Add(time.Duration(observation.w7d.windowMin) * time.Minute) account.SetReset7dAt(resetAt) } + pct := observation.w7d.usedPct + if existingValid && 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 From 3f3aa7246ab3795d590228852d8ed57ddbd020fa Mon Sep 17 00:00:00 2001 From: Delicious233 <101502465+DeliciousBuding@users.noreply.github.com> Date: Sun, 18 Oct 2026 14:54:14 +0800 Subject: [PATCH 04/24] fix(admin): expose WS usage frame freshness in live state The account live endpoint now returns the time and event type of the last applied WS quota frame, omitted until the first frame, so operators can tell frame-driven usage from WHAM probe updates. Raw JSON semantics stay intact: missing keys keep list values and explicit null clears invalid windows. --- admin/account_live.go | 9 ++++++ admin/account_live_test.go | 57 ++++++++++++++++++++++++++++++++++++++ docs/API.md | 50 +++++++++++++++++++++++++++++++++ 3 files changed, 116 insertions(+) diff --git a/admin/account_live.go b/admin/account_live.go index 4a62c984c..17f6f4178 100644 --- a/admin/account_live.go +++ b/admin/account_live.go @@ -30,6 +30,10 @@ type accountLiveItem struct { 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"` } // GetAccountLiveState returns request-local runtime counters for the visible @@ -80,6 +84,11 @@ func (h *Handler) GetAccountLiveState(c *gin.Context) { } 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 } diff --git a/admin/account_live_test.go b/admin/account_live_test.go index 06c10d8f6..87e908abe 100644 --- a/admin/account_live_test.go +++ b/admin/account_live_test.go @@ -113,3 +113,60 @@ func TestGetAccountLiveStateReturnsRuntimeUsage(t *testing.T) { t.Fatalf("using_credits missing: %+v", item) } } + +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", "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/docs/API.md b/docs/API.md index b35264076..1c8e6e3aa 100644 --- a/docs/API.md +++ b/docs/API.md @@ -774,6 +774,56 @@ 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, + "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，前端保留列表值 | +| 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` 表示窗口失效。Antigravity 行不使用 Codex 用量字段。 + #### PATCH /api/admin/accounts/:id/scheduler 更新账号调度配置。 From 39bc046a18df6eec19382ba3bae3f73c52bf1dd4 Mon Sep 17 00:00:00 2001 From: Delicious233 <101502465+DeliciousBuding@users.noreply.github.com> Date: Sun, 18 Oct 2026 14:54:20 +0800 Subject: [PATCH 05/24] fix(web): keep account polling from starving usage reloads The 1-second live merge rebuilt every row and every retry effect, so the missing-usage reload and stale-stats backoff restarted before their timers could fire. Merge only changed fields, keep unchanged row identities, drive the retry chains from explicit id-set and tick dependencies, and abort a hung live request after 5 seconds so the next tick still runs. --- frontend/src/hooks/useAccountLiveState.ts | 12 +- frontend/src/lib/accountLiveState.test.mjs | 40 +++++ frontend/src/pages/Accounts.tsx | 176 +++++++++++++++------ frontend/src/types.ts | 3 + 4 files changed, 176 insertions(+), 55 deletions(-) diff --git a/frontend/src/hooks/useAccountLiveState.ts b/frontend/src/hooks/useAccountLiveState.ts index 8c7c52151..bae255a7e 100644 --- a/frontend/src/hooks/useAccountLiveState.ts +++ b/frontend/src/hooks/useAccountLiveState.ts @@ -25,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 index 84da73712..a09188cf1 100644 --- a/frontend/src/lib/accountLiveState.test.mjs +++ b/frontend/src/lib/accountLiveState.test.mjs @@ -96,3 +96,43 @@ test("explicit null clears an invalid window", () => { 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); +}); 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 b1bf31417..baaa4761c 100644 --- a/frontend/src/types.ts +++ b/frontend/src/types.ts @@ -637,6 +637,9 @@ export interface AccountLiveStateResponse { 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 } From a56c6d334fb010c2a5d896acaf788fe98eee6ab2 Mon Sep 17 00:00:00 2001 From: Delicious233 <101502465+DeliciousBuding@users.noreply.github.com> Date: Sun, 18 Oct 2026 14:54:25 +0800 Subject: [PATCH 06/24] fix(web): correct auto-pause and cleanup copy Blank or 0 on an account auto-pause threshold inherits the group/global threshold; the window switch is what disables it. The rate-limited cleanup moves accounts to the recycle bin and leaves auto-paused accounts in place. --- frontend/src/locales/en.json | 8 ++++---- frontend/src/locales/zh.json | 8 ++++---- 2 files changed, 8 insertions(+), 8 deletions(-) diff --git a/frontend/src/locales/en.json b/frontend/src/locales/en.json index b78fc4747..b53956978 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 not included and stay in the list.", "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..c50b1e9ef 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": "会把全池处于限流状态的账号移入回收站（可在回收站恢复），不受当前筛选或分页影响；自动暂停的账号不在清理范围内，会保留。", "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 时立即停止续传。不会绕过手动禁用、鉴权失败或自定义自动暂停阈值。", From 8a29f748e92cba199b98ca4da9e0c5170f5861cb Mon Sep 17 00:00:00 2001 From: Delicious233 <101502465+DeliciousBuding@users.noreply.github.com> Date: Sun, 18 Oct 2026 14:54:31 +0800 Subject: [PATCH 07/24] docs: note frame ingest kill switch and account page fixes --- CHANGELOG.md | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index 9ec96a37c..0e568e1ba 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,5 +1,15 @@ # Changelog +## Unreleased + +### Features + +- **Codex WebSocket quota frames update account usage in real time.** 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. In-memory snapshots update on every frame; database writes are throttled to 30 seconds 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 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 auto-paused accounts are left in place, instead of saying status flags are cleared. + ## v3.1.1 - 2026-10-17 ### Features From 5aadcd01392a56ac7eebcca5ece6ccde93921cd7 Mon Sep 17 00:00:00 2001 From: Delicious233 <101502465+DeliciousBuding@users.noreply.github.com> Date: Sun, 18 Oct 2026 14:56:08 +0800 Subject: [PATCH 08/24] docs: reference the WS usage frame issue --- CHANGELOG.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 0e568e1ba..a1c7be0d6 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,7 +4,7 @@ ### Features -- **Codex WebSocket quota frames update account usage in real time.** 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. In-memory snapshots update on every frame; database writes are throttled to 30 seconds 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. +- **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. In-memory snapshots update on every frame; database writes are throttled to 30 seconds 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 From 278688bfd00e9ed64fefaaa4391db6bd6a43d8ff Mon Sep 17 00:00:00 2001 From: Delicious233 <101502465+DeliciousBuding@users.noreply.github.com> Date: Sun, 18 Oct 2026 15:12:31 +0800 Subject: [PATCH 09/24] fix(proxy): start a new window when the stored reset has passed A partial frame without reset_at kept an already-expired reset and the within-window monotonic guard then pinned the old 99-100% value, so a rolled over account stayed paused until the next WHAM probe. Only apply the guard while the stored reset is still ahead; otherwise start a new window with the configured length. Also prefer a future absolute reset over a zero relative one, require full-string numeric parses, treat limit_id before a display limit_name, apply credits from connection-level frames, and advance the frame freshness marker only when the observation passed the ordering gate. --- proxy/codex_usage_frames.go | 70 +++++++++++++++++++++++++++---------- proxy/handler.go | 36 +++++++++++++------ 2 files changed, 77 insertions(+), 29 deletions(-) diff --git a/proxy/codex_usage_frames.go b/proxy/codex_usage_frames.go index 102594a4e..43d2dc5bc 100644 --- a/proxy/codex_usage_frames.go +++ b/proxy/codex_usage_frames.go @@ -84,12 +84,13 @@ func applyCodexRateLimitsFrame(store *auth.Store, account *auth.Account, frame g if len(headers) == 0 { return } - account.MarkUsageFrameObserved(codexRateLimitsFrameType, observedAt) - applyCodexUsageHeaders(store, account, headers, observedAt) + if result := applyCodexUsageHeaders(store, account, headers, observedAt); result.Applied { + account.MarkUsageFrameObserved(codexRateLimitsFrameType, observedAt) + } } func codexRateLimitsHeaders(frame gjson.Result, observedAt time.Time) http.Header { - if name := firstCodexFrameString(frame, "metered_limit_name", "limit_name", "limit_id"); name != "" && !isDefaultCodexLimitName(name) { + if name := firstCodexFrameString(frame, "metered_limit_name", "limit_id", "limit_name"); name != "" && !isDefaultCodexLimitName(name) { return nil } headers := make(http.Header) @@ -141,12 +142,16 @@ func applyCodexResponseMetadataFrame(store *auth.Store, account *auth.Account, f if len(headers) == 0 { return } - account.MarkUsageFrameObserved(codexResponseMetadataFrameType, observedAt) - applyCodexUsageHeaders(store, account, headers, observedAt) + if result := applyCodexUsageHeaders(store, account, headers, observedAt); result.Applied { + account.MarkUsageFrameObserved(codexResponseMetadataFrameType, observedAt) + } } 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) { @@ -179,7 +184,7 @@ func codexResponseMetadataHeaders(frame gjson.Result, observedAt time.Time) http seconds = 0 } prefix := strings.TrimSuffix(name, "reset-at") - headers.Set(prefix+"reset-after-seconds", strconv.FormatInt(seconds, 10)) + absoluteResets[prefix+"reset-after-seconds"] = strconv.FormatInt(seconds, 10) } return true } @@ -190,6 +195,9 @@ func codexResponseMetadataHeaders(frame gjson.Result, observedAt time.Time) http headers.Set(name, raw) return true }) + for name, value := range absoluteResets { + headers.Set(name, value) + } return headers } @@ -212,9 +220,9 @@ func isCodexUsageFrameHeader(name string) bool { return false } -func applyCodexUsageHeaders(store *auth.Store, account *auth.Account, headers http.Header, observedAt time.Time) { +func applyCodexUsageHeaders(store *auth.Store, account *auth.Account, headers http.Header, observedAt time.Time) CodexUsageSyncResult { if len(headers) == 0 { - return + return CodexUsageSyncResult{} } // 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 @@ -230,7 +238,7 @@ func applyCodexUsageHeaders(store *auth.Store, account *auth.Account, headers ht // 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. - syncCodexUsageState(store, account, resp, observedAt, codexUsageSyncOptions{ + return syncCodexUsageState(store, account, resp, observedAt, codexUsageSyncOptions{ persistUsageSnapshot: true, clearAbsent5h: false, shouldPersistUsage: func(pct5h float64, has5h bool, pct7d float64, has7d bool) bool { @@ -260,16 +268,18 @@ func setCodexFrameWindow(headers http.Header, prefix string, window gjson.Result // (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 { + relativeSeen = true // Zero/negative means "resets now", matching the absolute form - // where a past epoch is clamped to zero. Only a missing or - // unparseable field keeps the previous reset. - if n < 0 { - n = 0 + // 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) } - return strconv.FormatInt(n, 10) } } } @@ -286,6 +296,9 @@ func codexFrameResetSeconds(window gjson.Result, observedAt time.Time) string { return strconv.FormatInt(seconds, 10) } } + if relativeSeen { + return "0" + } return "" } @@ -352,13 +365,34 @@ func ApplyCodexIdleUsageFrame(account *auth.Account, eventType string, payload [ if len(headers) == 0 { return } - observation := parseCodexUsageHeaderObservation(&http.Response{Header: headers}) - if !observation.authoritative { + resp := &http.Response{Header: headers} + observation := parseCodexUsageHeaderObservation(resp) + hasCredits, unlimited, creditsObserved := parseCodexCreditsFlagHeaders(resp) + if !observation.authoritative && !creditsObserved { return } account.TryApplyUsageObservation(observedAt, func() { account.MarkUsageFrameObserved(normalizedType, observedAt) - // Never clear an absent 5h window from a partial connection-level frame. - applyCodexUsageHeaderObservation(nil, account, observation, observedAt, false) + if observation.authoritative { + // Never clear an absent 5h window from a partial connection-level frame. + applyCodexUsageHeaderObservation(nil, account, observation, observedAt, false) + } + if creditsObserved { + // Connection-level frames also carry the post-turn credits snapshot. + // Apply it so the last spent credit can flip credit-based scheduling + // before the next WHAM probe. Keep the stored hard stop: the frame has + // no rate_limit_reached_type of its own. + var balance *string + if trimmed := strings.TrimSpace(headers.Get("x-codex-credits-balance")); trimmed != "" { + balance = &trimmed + } + var overage *bool + if raw := strings.TrimSpace(headers.Get("x-codex-credits-overage-reached")); raw != "" { + if value, err := strconv.ParseBool(raw); err == nil { + overage = &value + } + } + account.ApplySparseCreditsHeaders(hasCredits, unlimited, balance, overage, account.GetRateLimitReachedType()) + } }) } diff --git a/proxy/handler.go b/proxy/handler.go index fb2f97593..08379ba35 100644 --- a/proxy/handler.go +++ b/proxy/handler.go @@ -10,6 +10,7 @@ import ( "fmt" "io" "log" + "math" "math/rand" "net" "net/http" @@ -192,7 +193,10 @@ type CodexUsageSyncResult struct { Premium5hRateLimited bool // PersistUsage 表示节流策略允许本次观察落库。WS 帧路径可能只更新内存 // 快照；100% 冷却标记只跟着允许落库的观察执行，避免每帧重复写库。 - PersistUsage bool + PersistUsage bool + // Applied 表示本次观察通过了排序门并写入了账号状态；WS 帧路径用它决定 + // 是否更新“最后一帧”观测标记，避免迟到的旧帧伪装成最新帧。 + Applied bool UsageWindowLimitsIgnored bool // Cleared5h 表示本次同步因上游未返回 5h 窗口而清除了本地陈旧 5h 快照（issue #382）。 Cleared5h bool @@ -8155,21 +8159,27 @@ 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{} } + 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 _, err := fmt.Sscanf(windowStr, "%f", &windowMin); err != nil || windowMin <= 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 := strings.TrimSpace(resetStr) != "" if hasReset { - if _, err := fmt.Sscanf(resetStr, "%f", &resetSec); err != nil { - resetSec = 0 + if parsedReset, err := strconv.ParseFloat(strings.TrimSpace(resetStr), 64); err == nil && !math.IsNaN(parsedReset) && !math.IsInf(parsedReset, 0) { + resetSec = parsedReset } } return codexWindowUsage{ @@ -8194,8 +8204,8 @@ func codexResetAfterSeconds(header http.Header, prefix string) string { if raw == "" { return "" } - var epoch int64 - if _, err := fmt.Sscanf(raw, "%d", &epoch); err != nil || epoch <= 0 { + epoch, err := strconv.ParseInt(raw, 10, 64) + if err != nil || epoch <= 0 { return "" } seconds := time.Until(time.Unix(epoch, 0)).Seconds() @@ -8855,6 +8865,7 @@ func syncCodexUsageState(store *auth.Store, account *auth.Account, resp *http.Re } } + result.Applied = usageApplied || result.CreditsObserved return result } @@ -8952,15 +8963,18 @@ func applyCodexUsageHeaderObservation(store *auth.Store, account *auth.Account, 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() { + } 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 && sameUsageResetWindow(existingReset, resetAt) && pct < existingPct { + 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. @@ -8978,12 +8992,12 @@ func applyCodexUsageHeaderObservation(store *auth.Store, account *auth.Account, if observation.w7d.hasReset { resetAt = observedAt.Add(time.Duration(observation.w7d.resetSec) * time.Second) account.SetReset7dAt(resetAt) - } else if resetAt.IsZero() { + } 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 && sameUsageResetWindow(existingReset, resetAt) && pct < existingPct { + if existingValid && existingReset.After(observedAt) && sameUsageResetWindow(existingReset, resetAt) && pct < existingPct { pct = existingPct } account.SetWindow7dSeconds(int64(observation.w7d.windowMin * 60)) From 01fb2512031f125a04f6aecad2c901b8ec861a8a Mon Sep 17 00:00:00 2001 From: Delicious233 <101502465+DeliciousBuding@users.noreply.github.com> Date: Sun, 18 Oct 2026 15:12:41 +0800 Subject: [PATCH 10/24] test(proxy): cover frame entry points, persistence, and idle credits - assert the non-streaming Responses, Chat Completions and Messages streams apply a rate_limits frame to the account - round-trip the 30-second throttle against SQLite: first frame persists, a small delta stays memory-only, 100% forces a write - cover rolled-over windows, absolute-over-zero resets, trailing-garbage numbers, limit_id ordering, partial-frame 7d updates, single credits flags and idle credit application - cover the non-blocking TryApply skip and the observation ordering gate --- auth/usage_frame_test.go | 36 ++++ proxy/codex_usage_frames_integration_test.go | 86 ++++++++ proxy/codex_usage_frames_test.go | 206 +++++++++++++++++++ 3 files changed, 328 insertions(+) diff --git a/auth/usage_frame_test.go b/auth/usage_frame_test.go index c56a2e7d3..d0e1c4b11 100644 --- a/auth/usage_frame_test.go +++ b/auth/usage_frame_test.go @@ -79,3 +79,39 @@ func TestGetAccountLiveUsageSnapshotProjectsStatusAndUsage(t *testing.T) { 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/proxy/codex_usage_frames_integration_test.go b/proxy/codex_usage_frames_integration_test.go index 69e193002..338d124b4 100644 --- a/proxy/codex_usage_frames_integration_test.go +++ b/proxy/codex_usage_frames_integration_test.go @@ -69,3 +69,89 @@ func TestResponsesStreamAppliesCodexRateLimitsFrameToAccount(t *testing.T) { 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 index 0f9bba411..14e727ca0 100644 --- a/proxy/codex_usage_frames_test.go +++ b/proxy/codex_usage_frames_test.go @@ -1,7 +1,11 @@ package proxy import ( + "context" + "database/sql" + "fmt" "net/http" + "path/filepath" "strconv" "testing" "time" @@ -65,6 +69,10 @@ func TestApplyCodexRateLimitsFrameKeepsAbsent5hWindow(t *testing.T) { 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) { @@ -514,6 +522,22 @@ func TestApplyCodexRateLimitsFrameIgnoresGPTCodexMeteredLimit(t *testing.T) { } } +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) @@ -560,3 +584,185 @@ func TestApplyCodexUsageHeaderObservationBoundsMissingReset(t *testing.T) { 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") + } +} + +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(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 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 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) + if got := readPersisted5h(t, dbPath, accountID); got != 10 { + t.Fatalf("first frame persisted %v, want 10", got) + } + + apply(12, t0.Add(time.Second)) + if pct, _ := account.GetUsagePercent5h(); pct != 12 { + t.Fatalf("throttled frame did not update memory: %v", pct) + } + 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)) + if got := readPersisted5h(t, dbPath, accountID); got != 100 { + t.Fatalf("100%% frame persisted %v, want 100", got) + } +} + +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() +} From 280471d9449a5c3d6b0c7f6479d124ddd6266d7a Mon Sep 17 00:00:00 2001 From: Delicious233 <101502465+DeliciousBuding@users.noreply.github.com> Date: Sun, 18 Oct 2026 15:12:46 +0800 Subject: [PATCH 11/24] docs: note rolled-over windows and idle frame credits --- CHANGELOG.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index a1c7be0d6..0430505c7 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,7 +4,7 @@ ### 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. In-memory snapshots update on every frame; database writes are throttled to 30 seconds 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. +- **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. 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. In-memory snapshots update on every frame; database writes are throttled to 30 seconds 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 From 2c49b8367831a7016f44e81c54890914a21f88db Mon Sep 17 00:00:00 2001 From: Delicious233 <101502465+DeliciousBuding@users.noreply.github.com> Date: Sun, 18 Oct 2026 15:16:09 +0800 Subject: [PATCH 12/24] fix(proxy): treat a malformed reset value as absent A reset header such as "60abc" previously left hasReset true with a zero offset, which fabricated an immediate reset and could skip a cooldown. Only accept a fully parsed value as a reset; malformed input now falls back to the existing or window-length reset. --- proxy/codex_usage_frames_test.go | 2 ++ proxy/handler.go | 7 ++++--- 2 files changed, 6 insertions(+), 3 deletions(-) diff --git a/proxy/codex_usage_frames_test.go b/proxy/codex_usage_frames_test.go index 14e727ca0..270b2f223 100644 --- a/proxy/codex_usage_frames_test.go +++ b/proxy/codex_usage_frames_test.go @@ -635,6 +635,8 @@ func TestParseCodexWindowUsageRejectsTrailingGarbage(t *testing.T) { } 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) } } diff --git a/proxy/handler.go b/proxy/handler.go index 08379ba35..26cd04f4c 100644 --- a/proxy/handler.go +++ b/proxy/handler.go @@ -8176,10 +8176,11 @@ func parseCodexWindowUsage(usedStr, windowStr, resetStr string) codexWindowUsage return codexWindowUsage{} } windowMin = parsedWindow - hasReset := strings.TrimSpace(resetStr) != "" - if hasReset { - if parsedReset, err := strconv.ParseFloat(strings.TrimSpace(resetStr), 64); err == nil && !math.IsNaN(parsedReset) && !math.IsInf(parsedReset, 0) { + 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{ From 07d83d5cd3c7e7785b3b0b1e56e0497a337863fc Mon Sep 17 00:00:00 2001 From: Delicious233 <101502465+DeliciousBuding@users.noreply.github.com> Date: Sun, 18 Oct 2026 15:17:07 +0800 Subject: [PATCH 13/24] fix(auth): keep newer runtime usage across outbox reloads A WS frame updates memory immediately while persistence is throttled, so an outbox row can still carry the older DB snapshot. The reload copied those usage fields unconditionally and kept the newer usageObservedAt, leaving the scheduler with lower usage than the account actually has. Copy the 5h/7d groups only when the row timestamp is not older, except on credential identity changes. --- auth/scheduler_outbox_consumer.go | 28 +++++++++++----- auth/scheduler_outbox_consumer_test.go | 46 ++++++++++++++++++++++++++ 2 files changed, 65 insertions(+), 9 deletions(-) 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")) From 169c5e9fab538a20b953e855752540ba012b24c8 Mon Sep 17 00:00:00 2001 From: Delicious233 <101502465+DeliciousBuding@users.noreply.github.com> Date: Sun, 18 Oct 2026 15:30:16 +0800 Subject: [PATCH 14/24] fix(proxy): skip forwarded quota frames for relay accounts A relay upstream can forward codex.rate_limits frames that describe the relay's own backend account, not the local relay account's Codex windows. Ingesting them could set a spurious 100% snapshot and put the relay account into a multi-day cooldown. Gate both the per-turn and connection-level entry points on IsRelayStyle, matching the rest of the relay paths. --- proxy/codex_usage_frames.go | 7 +++++-- proxy/codex_usage_frames_test.go | 15 +++++++++++++++ 2 files changed, 20 insertions(+), 2 deletions(-) diff --git a/proxy/codex_usage_frames.go b/proxy/codex_usage_frames.go index 43d2dc5bc..a4752ac81 100644 --- a/proxy/codex_usage_frames.go +++ b/proxy/codex_usage_frames.go @@ -64,7 +64,10 @@ func codexUsageFramePersistInterval() time.Duration { // already on the wire (the gateway forwards them downstream), so this adds no // upstream requests and no new request headers. func ApplyCodexUsageEventFrame(store *auth.Store, account *auth.Account, eventType string, frame gjson.Result, observedAt time.Time) { - if account == nil || !codexUsageFrameIngestEnabled.Load() { + // 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)) { @@ -341,7 +344,7 @@ func isDefaultCodexLimitName(name string) bool { // paths; this path exists to shrink the window between a turn finishing and // the scheduler seeing its consumption. func ApplyCodexIdleUsageFrame(account *auth.Account, eventType string, payload []byte) { - if account == nil || len(payload) == 0 || !codexUsageFrameIngestEnabled.Load() { + if account == nil || account.IsRelayStyle() || len(payload) == 0 || !codexUsageFrameIngestEnabled.Load() { return } normalizedType := strings.ToLower(strings.TrimSpace(eventType)) diff --git a/proxy/codex_usage_frames_test.go b/proxy/codex_usage_frames_test.go index 270b2f223..5dfbf20b5 100644 --- a/proxy/codex_usage_frames_test.go +++ b/proxy/codex_usage_frames_test.go @@ -685,6 +685,21 @@ func TestCodexRateLimitsFrameIgnoresSingleCreditsFlag(t *testing.T) { } } +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(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) From ab5c5fd1ab3f45e21e421a236ca478e31e784bec Mon Sep 17 00:00:00 2001 From: Delicious233 <101502465+DeliciousBuding@users.noreply.github.com> Date: Sun, 18 Oct 2026 15:30:17 +0800 Subject: [PATCH 15/24] fix(web): include usage_exhausted auto-pause in cleanup copy The rate-limited cleanup collector includes usage_exhausted accounts, so an auto-paused account at 100% can be cleaned. Say so instead of claiming every auto-paused account is left in place. --- CHANGELOG.md | 4 ++-- frontend/src/locales/en.json | 2 +- frontend/src/locales/zh.json | 2 +- 3 files changed, 4 insertions(+), 4 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 0430505c7..153ecbbf9 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,11 +4,11 @@ ### 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. 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. In-memory snapshots update on every frame; database writes are throttled to 30 seconds 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. +- **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. In-memory snapshots update on every frame; database writes are throttled to 30 seconds 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 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 auto-paused accounts are left in place, instead of saying status flags are cleared. +- **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 diff --git a/frontend/src/locales/en.json b/frontend/src/locales/en.json index b53956978..3f9bad69c 100644 --- a/frontend/src/locales/en.json +++ b/frontend/src/locales/en.json @@ -1543,7 +1543,7 @@ "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 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 not included and stay in the list.", + "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", diff --git a/frontend/src/locales/zh.json b/frontend/src/locales/zh.json index c50b1e9ef..4a7aa5c07 100644 --- a/frontend/src/locales/zh.json +++ b/frontend/src/locales/zh.json @@ -1541,7 +1541,7 @@ "cleanBannedTitle": "清理封禁账号", "cleanBannedDesc": "会清除所有 unauthorized 状态标记，让这些账号重新参与调度与测试。", "cleanRateLimitedTitle": "清理限流账号", - "cleanRateLimitedDesc": "会把全池处于限流状态的账号移入回收站（可在回收站恢复），不受当前筛选或分页影响；自动暂停的账号不在清理范围内，会保留。", + "cleanRateLimitedDesc": "会把全池处于限流状态的账号移入回收站（可在回收站恢复），不受当前筛选或分页影响；自动暂停的账号如果同时处于限流或 usage_exhausted 状态，也会被纳入清理范围。", "cleanConfirm": "确认清理", "cleanBannedProgressTitle": "正在清理封禁账号", "cleanBannedSuccess": "封禁账号状态已清理", From 72857c905bae3a8aed303714f3b6f7d2e2a5bf7d Mon Sep 17 00:00:00 2001 From: Delicious233 <101502465+DeliciousBuding@users.noreply.github.com> Date: Sun, 18 Oct 2026 15:47:09 +0800 Subject: [PATCH 16/24] fix(proxy): refresh the scheduler after connection-level frames An idle frame can move an account back under its threshold, but the read-pump observer had no store handle, so queued dispatchers waited for the periodic availability recheck. Pass the store into the observer, expose a non-persisting scheduler refresh, and call it after a frame is applied. --- CHANGELOG.md | 2 +- auth/store.go | 8 ++++++++ main.go | 4 +++- proxy/codex_usage_frames.go | 10 ++++++++-- proxy/codex_usage_frames_test.go | 8 ++++---- 5 files changed, 24 insertions(+), 8 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 153ecbbf9..29e6d45af 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,7 +4,7 @@ ### 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. In-memory snapshots update on every frame; database writes are throttled to 30 seconds 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. +- **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 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 diff --git a/auth/store.go b/auth/store.go index d58735728..dbb9c7fda 100644 --- a/auth/store.go +++ b/auth/store.go @@ -4679,6 +4679,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 diff --git a/main.go b/main.go index 3961f726b..d39c4ed6b 100644 --- a/main.go +++ b/main.go @@ -446,7 +446,9 @@ func main() { // 注册 WebSocket 执行函数（避免 proxy ↔ wsrelay 循环依赖） proxy.WebsocketExecuteFunc = wsrelay.ExecuteRequestWebsocket // 连接级（终态后补发）的用量帧也交给 proxy 更新内存快照；observer 必须非阻塞。 - wsrelay.SetIdleMetadataObserver(proxy.ApplyCodexIdleUsageFrame) + wsrelay.SetIdleMetadataObserver(func(account *auth.Account, eventType string, payload []byte) { + proxy.ApplyCodexIdleUsageFrame(store, account, eventType, payload) + }) // 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 { diff --git a/proxy/codex_usage_frames.go b/proxy/codex_usage_frames.go index a4752ac81..8842c31c3 100644 --- a/proxy/codex_usage_frames.go +++ b/proxy/codex_usage_frames.go @@ -343,7 +343,7 @@ func isDefaultCodexLimitName(name string) bool { // write. Persistence and the 100% cooldown markings stay on the preflight/WHAM // paths; this path exists to shrink the window between a turn finishing and // the scheduler seeing its consumption. -func ApplyCodexIdleUsageFrame(account *auth.Account, eventType string, payload []byte) { +func ApplyCodexIdleUsageFrame(store *auth.Store, account *auth.Account, eventType string, payload []byte) { if account == nil || account.IsRelayStyle() || len(payload) == 0 || !codexUsageFrameIngestEnabled.Load() { return } @@ -374,7 +374,7 @@ func ApplyCodexIdleUsageFrame(account *auth.Account, eventType string, payload [ if !observation.authoritative && !creditsObserved { return } - account.TryApplyUsageObservation(observedAt, func() { + applied := account.TryApplyUsageObservation(observedAt, func() { account.MarkUsageFrameObserved(normalizedType, observedAt) if observation.authoritative { // Never clear an absent 5h window from a partial connection-level frame. @@ -398,4 +398,10 @@ func ApplyCodexIdleUsageFrame(account *auth.Account, eventType string, payload [ account.ApplySparseCreditsHeaders(hasCredits, unlimited, balance, overage, account.GetRateLimitReachedType()) } }) + if applied { + // A frame can restore an account that crossed back under its threshold; + // queued dispatchers wait on the availability notification, so refresh + // the scheduler even though persistence stays on the preflight path. + store.RefreshUsageFrameScheduler(account) + } } diff --git a/proxy/codex_usage_frames_test.go b/proxy/codex_usage_frames_test.go index 5dfbf20b5..5b694725a 100644 --- a/proxy/codex_usage_frames_test.go +++ b/proxy/codex_usage_frames_test.go @@ -230,7 +230,7 @@ func TestCodexUsageFrameIngestKillSwitch(t *testing.T) { acc := newUsageFrameTestAccount() frame := gjson.Parse(`{"rate_limits":{"primary":{"used_percent":42,"window_minutes":300}}}`) ApplyCodexUsageEventFrame(nil, acc, "codex.rate_limits", frame, time.Now()) - ApplyCodexIdleUsageFrame(acc, "codex.rate_limits", []byte(`{"rate_limits":{"primary":{"used_percent":42,"window_minutes":300}}}`)) + 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) @@ -371,7 +371,7 @@ func TestApplyCodexIdleUsageFrameUpdatesMemoryOnly(t *testing.T) { acc.SetUsageSnapshot5hAt(20, now.Add(time.Hour), now) acc.SetUsagePercent7d(10) - ApplyCodexIdleUsageFrame(acc, "codex.rate_limits", []byte(`{ + ApplyCodexIdleUsageFrame(nil, acc, "codex.rate_limits", []byte(`{ "type":"codex.rate_limits", "rate_limits":{ "primary":{"used_percent":75,"window_minutes":300,"reset_at":1700000900}, @@ -661,7 +661,7 @@ func TestApplyCodexIdleUsageFrameAppliesCredits(t *testing.T) { acc := newUsageFrameTestAccount() acc.ApplySparseCreditsHeaders(true, false, nil, nil, "workspace_member_credits_depleted") - ApplyCodexIdleUsageFrame(acc, "codex.rate_limits", []byte(`{ + ApplyCodexIdleUsageFrame(nil, acc, "codex.rate_limits", []byte(`{ "type":"codex.rate_limits", "credits":{"has_credits":false,"unlimited":false,"balance":"0"} }`)) @@ -693,7 +693,7 @@ func TestCodexUsageFramesSkipRelayAccounts(t *testing.T) { frame := gjson.Parse(`{"rate_limits":{"primary":{"used_percent":42,"window_minutes":300}}}`) ApplyCodexUsageEventFrame(nil, acc, "codex.rate_limits", frame, time.Now()) - ApplyCodexIdleUsageFrame(acc, "codex.rate_limits", []byte(`{"rate_limits":{"primary":{"used_percent":42,"window_minutes":300}}}`)) + 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) From 3fdbd449888c163c0f12c27aee7d443564998fd2 Mon Sep 17 00:00:00 2001 From: Delicious233 <101502465+DeliciousBuding@users.noreply.github.com> Date: Sun, 18 Oct 2026 17:07:37 +0800 Subject: [PATCH 17/24] fix(proxy): persist frame usage off the read path The frame path wrote the usage snapshot synchronously inside the upstream SSE read loop, under downstreamMu, so a locked or slow database could stall a streaming response (15s with SQLite busy_timeout in a local reproduction). Keep memory and scheduler updates synchronous, coalesce the snapshot per account, and let a store worker perform the throttled write. A job is skipped when a newer WHAM/frame observation has been applied, and the throttle marker still advances only after a successful write. Pending jobs are dropped on stop; the snapshot is a cache WHAM can rebuild. --- CHANGELOG.md | 2 +- auth/store.go | 3 + auth/usage_frame_persist.go | 157 +++++++++++++++++++++++++++++++ docs/CONFIGURATION.md | 2 +- proxy/codex_usage_frames.go | 1 + proxy/codex_usage_frames_test.go | 85 ++++++++++++++++- proxy/handler.go | 24 ++++- 7 files changed, 265 insertions(+), 9 deletions(-) create mode 100644 auth/usage_frame_persist.go diff --git a/CHANGELOG.md b/CHANGELOG.md index 29e6d45af..aa4cd9112 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,7 +4,7 @@ ### 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 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. +- **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 diff --git a/auth/store.go b/auth/store.go index dbb9c7fda..b718f410c 100644 --- a/auth/store.go +++ b/auth/store.go @@ -3809,6 +3809,7 @@ type Store struct { armedBoundaryAt atomic.Int64 lazyRefreshInFlight sync.Map stopCh chan struct{} + usageFramePersist *usageFramePersistQueue stopOnce sync.Once backgroundCtx context.Context backgroundCancel context.CancelFunc @@ -4463,6 +4464,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(), @@ -4474,6 +4476,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()) diff --git a/auth/usage_frame_persist.go b/auth/usage_frame_persist.go new file mode 100644 index 000000000..64d885d39 --- /dev/null +++ b/auth/usage_frame_persist.go @@ -0,0 +1,157 @@ +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 { + dbID int64 + observedAt time.Time + pct5h float64 + reset5hAt time.Time + has5h bool + updatedAt5h time.Time + pct7d float64 + reset7dAt time.Time + has7d 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 { + job, ok := q.take() + if !ok { + break + } + q.persist(s, job) + } + } + } +} + +func (q *usageFramePersistQueue) enqueue(job usageFramePersistJob) { + q.mu.Lock() + 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 { + return + } + // A WHAM probe or a newer frame may have been applied after this job was + // queued; in that case the newer observation (and its own persistence path) + // owns the snapshot. Writing this older job would roll usage back. + acc.mu.RLock() + superseded := acc.usageObservedAt.After(job.observedAt) + acc.mu.RUnlock() + if superseded { + return + } + ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second) + defer cancel() + var err error + switch { + case job.has7d && job.has5h: + err = s.db.UpdateUsageSnapshotFull(ctx, job.dbID, job.pct7d, job.reset7dAt, job.pct5h, job.reset5hAt, job.observedAt, job.updatedAt5h) + case job.has7d: + err = s.db.UpdateUsageSnapshot(ctx, job.dbID, job.pct7d, job.observedAt) + case job.has5h: + err = s.db.UpdateUsageSnapshot5h(ctx, job.dbID, job.pct5h, job.reset5hAt, job.observedAt) + default: + return + } + 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) + return + } + acc.MarkUsageFramePersisted(job.observedAt, job.pct5h, job.has5h, job.pct7d, job.has7d) +} + +// 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) { + if s == nil || acc == nil || s.usageFramePersist == nil || acc.DBID == 0 { + return + } + select { + case <-s.stopCh: + return + default: + } + if observedAt.IsZero() { + observedAt = time.Now() + } + snap := acc.GetAccountLiveUsageSnapshot() + if !snap.UsagePercent5hValid && !snap.UsagePercent7dValid { + return + } + s.usageFramePersist.enqueue(usageFramePersistJob{ + dbID: acc.DBID, + observedAt: observedAt, + pct5h: snap.UsagePercent5h, + reset5hAt: snap.Reset5hAt, + has5h: snap.UsagePercent5hValid, + updatedAt5h: snap.UsageUpdatedAt5h, + pct7d: snap.UsagePercent7d, + reset7dAt: snap.Reset7dAt, + has7d: snap.UsagePercent7dValid, + }) +} diff --git a/docs/CONFIGURATION.md b/docs/CONFIGURATION.md index f22a1d9d6..eccdad9ad 100644 --- a/docs/CONFIGURATION.md +++ b/docs/CONFIGURATION.md @@ -87,7 +87,7 @@ 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 逐帧用量快照落库的最小间隔（秒），正整数，重启生效。内存快照每帧更新；跨自动暂停阈值、5pp 变化或 100% 时仍强制落库。大号池可按数据库承压调大 | +| `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 上游流转 | diff --git a/proxy/codex_usage_frames.go b/proxy/codex_usage_frames.go index 8842c31c3..65244a1e6 100644 --- a/proxy/codex_usage_frames.go +++ b/proxy/codex_usage_frames.go @@ -244,6 +244,7 @@ func applyCodexUsageHeaders(store *auth.Store, account *auth.Account, headers ht return syncCodexUsageState(store, account, resp, observedAt, codexUsageSyncOptions{ persistUsageSnapshot: true, clearAbsent5h: false, + asyncPersist: true, shouldPersistUsage: func(pct5h float64, has5h bool, pct7d float64, has7d bool) bool { return account.ShouldPersistUsageFrame(observedAt, pct5h, has5h, pct7d, has7d, codexUsageFramePersistInterval()) }, diff --git a/proxy/codex_usage_frames_test.go b/proxy/codex_usage_frames_test.go index 5b694725a..f481e817f 100644 --- a/proxy/codex_usage_frames_test.go +++ b/proxy/codex_usage_frames_test.go @@ -752,21 +752,96 @@ func TestCodexUsageFramePersistenceRespectsThrottle(t *testing.T) { } apply(10, t0) - if got := readPersisted5h(t, dbPath, accountID); got != 10 { - t.Fatalf("first frame persisted %v, want 10", got) - } + 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)) - if got := readPersisted5h(t, dbPath, accountID); got != 100 { - t.Fatalf("100%% frame persisted %v, want 100", got) + 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) { + 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(`{"type":"codex.rate_limits","rate_limits":{"primary":{"used_percent":42,"window_minutes":300}}}`) + start := time.Now() + ApplyCodexUsageEventFrame(store, account, "codex.rate_limits", frame, time.Now()) + if elapsed := time.Since(start); elapsed > 500*time.Millisecond { + t.Fatalf("frame path blocked %s on a locked database, want enqueue-only", elapsed) + } + if pct, ok := account.GetUsagePercent5h(); !ok || pct != 42 { + t.Fatalf("memory usage = %v valid=%v, want 42", pct, ok) + } + if _, err := conn.ExecContext(ctx, "ROLLBACK"); err != nil { + t.Fatalf("rollback: %v", err) + } + waitPersisted5h(t, dbPath, accountID, 42) +} + +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) } } diff --git a/proxy/handler.go b/proxy/handler.go index 26cd04f4c..e4cde2280 100644 --- a/proxy/handler.go +++ b/proxy/handler.go @@ -8756,6 +8756,10 @@ type codexUsageSyncOptions struct { // 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 } func syncCodexUsageState(store *auth.Store, account *auth.Account, resp *http.Response, observedAt time.Time, opts codexUsageSyncOptions) CodexUsageSyncResult { @@ -8830,12 +8834,28 @@ func syncCodexUsageState(store *auth.Store, account *auth.Account, resp *http.Re result.PersistUsage = persistUsage persistedUsage := false if result.HasUsage7d { - persistedUsage = store.PublishUsageSnapshot(account, result.UsagePct7d, observedAt, persistUsage) + 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) + if persistUsage { + store.EnqueueUsageFramePersist(account, observedAt) + } + } else { + persistedUsage = store.PublishUsageSnapshot(account, result.UsagePct7d, observedAt, persistUsage) + } if result.UsagePct7d >= 100 && persistUsage { result.Usage7dRateLimited = store.MarkUsage7dRateLimited(account) } } else if result.Used5hHeaders { - persistedUsage = store.PublishUsageSnapshot5hOnly(account, observedAt, persistUsage) + if opts.asyncPersist { + store.PublishUsageSnapshot5hOnly(account, observedAt, false) + if persistUsage { + store.EnqueueUsageFramePersist(account, observedAt) + } + } else { + persistedUsage = store.PublishUsageSnapshot5hOnly(account, observedAt, persistUsage) + } result.Persisted5hOnly = persistUsage } if persistUsage && persistedUsage { From 5dd7b5ebc0c524966193d312ab3b7312b664c51a Mon Sep 17 00:00:00 2001 From: Delicious233 <101502465+DeliciousBuding@users.noreply.github.com> Date: Sun, 18 Oct 2026 17:26:35 +0800 Subject: [PATCH 18/24] fix(auth): serialize frame snapshot writes with newer observations The async worker checked usageObservedAt, released the account lock, then wrote; a WHAM observation could apply and persist in that gap, after which the older job overwrote it. Serialize the check and write with the synchronous snapshot writers on a dedicated lock (holding usageSyncMu would reintroduce the read-path stall) and add a deterministic regression that applies a newer observation while the worker waits. --- auth/premium_rate_limit.go | 13 +++++-- auth/store.go | 30 ++++++++++++--- auth/usage_frame_persist.go | 13 +++++-- auth/usage_frame_persist_test.go | 63 ++++++++++++++++++++++++++++++++ 4 files changed, 106 insertions(+), 13 deletions(-) create mode 100644 auth/usage_frame_persist_test.go diff --git a/auth/premium_rate_limit.go b/auth/premium_rate_limit.go index c2c31f7a6..9110eefd6 100644 --- a/auth/premium_rate_limit.go +++ b/auth/premium_rate_limit.go @@ -163,7 +163,10 @@ func (s *Store) PublishUsageSnapshot5hOnly(acc *Account, observedAt time.Time, p ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second) defer cancel() - if err := s.db.UpdateUsageSnapshot5h(ctx, acc.DBID, pct5h, reset5hAt, observedAt); err != nil { + err := s.withUsageSnapshotWrite(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 } @@ -274,7 +277,9 @@ func (s *Store) ClearAbsentUsageSnapshot5hAt(acc *Account, observedAt time.Time) 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(func() error { + return s.db.ClearUsageSnapshot5h(ctx, acc.DBID) + }); err != nil { log.Printf("[账号 %d] 清除 5h 用量快照失败: %v", acc.DBID, err) } if cleared5hCooldown { @@ -363,7 +368,9 @@ func (s *Store) markPremium5hRateLimited(acc *Account, resetAt, observedAt time. 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(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/store.go b/auth/store.go index b718f410c..589c66b81 100644 --- a/auth/store.go +++ b/auth/store.go @@ -3810,10 +3810,14 @@ type Store struct { lazyRefreshInFlight sync.Map stopCh chan struct{} usageFramePersist *usageFramePersistQueue - stopOnce sync.Once - backgroundCtx context.Context - backgroundCancel context.CancelFunc - wg sync.WaitGroup + // usageSnapshotPersistMu serializes 5h/7d credential writes between the + // synchronous observation paths and the async frame-persist worker, so a + // queued older frame cannot overwrite a newer WHAM observation. + usageSnapshotPersistMu sync.Mutex + stopOnce sync.Once + backgroundCtx context.Context + backgroundCancel context.CancelFunc + wg sync.WaitGroup // 代理池 proxyPoolReloadMu sync.Mutex @@ -11142,7 +11146,10 @@ func (s *Store) PublishUsageSnapshot(acc *Account, pct7d float64, observedAt tim 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, observedAt, acc.GetUsageUpdatedAt5h()); err != nil { + err := s.withUsageSnapshotWrite(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 } @@ -11151,13 +11158,24 @@ func (s *Store) PublishUsageSnapshot(acc *Account, pct7d float64, observedAt tim ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second) defer cancel() - if err := s.db.UpdateUsageSnapshot(ctx, acc.DBID, pct7d, observedAt); err != nil { + err := s.withUsageSnapshotWrite(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(fn func() error) error { + s.usageSnapshotPersistMu.Lock() + defer s.usageSnapshotPersistMu.Unlock() + return fn() +} + // UpdateAccountSubscriptionExpiresAt persists the latest subscription expiration observed from upstream. func (s *Store) UpdateAccountSubscriptionExpiresAt(acc *Account, expiresAt time.Time) bool { if s == nil || acc == nil || expiresAt.IsZero() { diff --git a/auth/usage_frame_persist.go b/auth/usage_frame_persist.go index 64d885d39..4a17b83b4 100644 --- a/auth/usage_frame_persist.go +++ b/auth/usage_frame_persist.go @@ -94,13 +94,16 @@ func (q *usageFramePersistQueue) persist(s *Store, job usageFramePersistJob) { if acc == nil { return } - // A WHAM probe or a newer frame may have been applied after this job was - // queued; in that case the newer observation (and its own persistence path) - // owns the snapshot. Writing this older job would roll usage back. + // 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 the store's snapshot-write lock instead. + s.usageSnapshotPersistMu.Lock() acc.mu.RLock() superseded := acc.usageObservedAt.After(job.observedAt) acc.mu.RUnlock() if superseded { + s.usageSnapshotPersistMu.Unlock() return } ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second) @@ -114,8 +117,10 @@ func (q *usageFramePersistQueue) persist(s *Store, job usageFramePersistJob) { case job.has5h: err = s.db.UpdateUsageSnapshot5h(ctx, job.dbID, job.pct5h, job.reset5hAt, job.observedAt) default: + s.usageSnapshotPersistMu.Unlock() return } + s.usageSnapshotPersistMu.Unlock() if err != nil { // Leave the throttle marker untouched; the next frame observation sees // that no successful write happened and enqueues again. @@ -128,7 +133,7 @@ func (q *usageFramePersistQueue) persist(s *Store, job usageFramePersistJob) { // 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) { - if s == nil || acc == nil || s.usageFramePersist == nil || acc.DBID == 0 { + if s == nil || acc == nil || s.db == nil || s.usageFramePersist == nil || acc.DBID == 0 { return } select { diff --git a/auth/usage_frame_persist_test.go b/auth/usage_frame_persist_test.go new file mode 100644 index 000000000..cdcd2bef7 --- /dev/null +++ b/auth/usage_frame_persist_test.go @@ -0,0 +1,63 @@ +package auth + +import ( + "context" + "database/sql" + "path/filepath" + "testing" + "time" + + "github.com/codex2api/database" + "github.com/tidwall/gjson" +) + +// 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. + store.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") + } + store.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) + } +} From 301ec4964ca4da2ed9799e2a1679cf430a986d7b Mon Sep 17 00:00:00 2001 From: Delicious233 <101502465+DeliciousBuding@users.noreply.github.com> Date: Sun, 18 Oct 2026 18:10:29 +0800 Subject: [PATCH 19/24] fix(auth): never shorten an active cooldown with usage frames Keep passive 5h and 7d observations from replacing longer cooldowns or authoritative reasons. Update usage snapshots independently and avoid repeated cooldown writes and failure penalties. Preserve authentication priority and allow actual Responses rejections to upgrade local usage reasons without reducing the deadline. Co-authored-by: Codex --- CHANGELOG.md | 2 + auth/premium_rate_limit.go | 39 ++++--- auth/store.go | 88 ++++++++++++---- auth/usage_cooldown_test.go | 172 +++++++++++++++++++++++++++++++ proxy/codex_usage_frames_test.go | 34 ++++++ 5 files changed, 294 insertions(+), 41 deletions(-) create mode 100644 auth/usage_cooldown_test.go diff --git a/CHANGELOG.md b/CHANGELOG.md index aa4cd9112..53404b999 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -8,6 +8,8 @@ ### Fixes +- **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 diff --git a/auth/premium_rate_limit.go b/auth/premium_rate_limit.go index 9110eefd6..15672d1da 100644 --- a/auth/premium_rate_limit.go +++ b/auth/premium_rate_limit.go @@ -330,43 +330,40 @@ func (s *Store) markPremium5hRateLimited(acc *Account, resetAt, observedAt time. resetAt = now.Add(premium5hFallbackWindow) } - // Re-observing the same exhausted window must not write the cooldown again. - acc.mu.RLock() - alreadyLimited := acc.Status == StatusCooldown && - acc.CooldownReason == cooldownReason && - now.Before(acc.CooldownUtil) && - cooldownMatchesReset(acc.CooldownUtil, resetAt) - acc.mu.RUnlock() - if alreadyLimited { - return - } - 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) + if cooldownChanged { + acc.LastRateLimitedAt = now + acc.Status = StatusCooldown + acc.CooldownUtil = until + acc.CooldownReason = reason + 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 cooldownChanged { + s.setCachedAccountCooldown(acc.DBID, reason, until) + } - if s.db == nil { + if 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 cooldownChanged { + if err := s.db.SetCooldown(ctx, acc.DBID, reason, until); err != nil { + log.Printf("[账号 %d] 持久化 premium 5h 限流冷却状态失败: %v", acc.DBID, err) + } } if err := s.withUsageSnapshotWrite(func() error { return s.db.UpdateUsageSnapshot5h(ctx, acc.DBID, 100, resetAt, now) diff --git a/auth/store.go b/auth/store.go index 589c66b81..8b44a18f4 100644 --- a/auth/store.go +++ b/auth/store.go @@ -2370,44 +2370,92 @@ func cooldownMatchesReset(cooldownUntil, resetAt time.Time) bool { 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() { + 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 } now := time.Now() - duration := 7 * 24 * time.Hour - resetAt := acc.GetReset7dAt() - if !resetAt.IsZero() { - untilReset := resetAt.Sub(now) - if untilReset <= 0 { - return false - } - duration = untilReset + resetAt := acc.Reset7dAt + if resetAt.IsZero() { + resetAt = now.Add(7 * 24 * time.Hour) + } else if !resetAt.After(now) { + acc.mu.Unlock() + return false } - acc.mu.RLock() - alreadyLimited := acc.Status == StatusCooldown && - (acc.CooldownReason == "rate_limited" || acc.CooldownReason == "rate_limited_7d") && - now.Before(acc.CooldownUtil) && - cooldownMatchesReset(acc.CooldownUtil, resetAt) - acc.mu.RUnlock() - if alreadyLimited { + 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() return true } + until, reason, changed := acc.usageCooldownTargetLocked(now, resetAt, "rate_limited") + if !changed { + acc.mu.Unlock() + return true + } + + 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.mu.Unlock() - s.MarkCooldown(acc, duration, "rate_limited") + s.fastSchedulerUpdate(acc) + s.setCachedAccountCooldown(acc.DBID, reason, until) + if s.db != nil { + 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] 持久化 7d 限流冷却状态失败: %v", acc.DBID, err) + } + } return true } diff --git a/auth/usage_cooldown_test.go b/auth/usage_cooldown_test.go new file mode 100644 index 000000000..53fe989b4 --- /dev/null +++ b/auth/usage_cooldown_test.go @@ -0,0 +1,172 @@ +package auth + +import ( + "context" + "database/sql" + "path/filepath" + "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) + } + }) + } + }) + } +} diff --git a/proxy/codex_usage_frames_test.go b/proxy/codex_usage_frames_test.go index f481e817f..31cd8891f 100644 --- a/proxy/codex_usage_frames_test.go +++ b/proxy/codex_usage_frames_test.go @@ -180,6 +180,40 @@ func TestSyncCodexUsageStateMarks100PercentOnlyWhenPersisting(t *testing.T) { } } +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" From 38f874c1104d5da55046248952b2899ab8f97f6b Mon Sep 17 00:00:00 2001 From: Delicious233 <101502465+DeliciousBuding@users.noreply.github.com> Date: Sun, 18 Oct 2026 18:49:07 +0800 Subject: [PATCH 20/24] fix(proxy): order credit observations and retry failed cooldown writes Apply quota, credits and hard-stop observations through one ordering gate. Keep cooldown writes pending until successful so the same observation can retry after a transient database error without adding another failure. Co-authored-by: Codex --- auth/premium_rate_limit.go | 7 +-- auth/store.go | 26 +++++----- auth/usage_cooldown_persist.go | 35 +++++++++++++ auth/usage_cooldown_test.go | 65 +++++++++++++++++++++++ proxy/codex_usage_frames.go | 12 ++--- proxy/codex_usage_frames_test.go | 26 ++++++++++ proxy/handler.go | 89 ++++++++++++++++++-------------- 7 files changed, 196 insertions(+), 64 deletions(-) create mode 100644 auth/usage_cooldown_persist.go diff --git a/auth/premium_rate_limit.go b/auth/premium_rate_limit.go index 15672d1da..3d67badae 100644 --- a/auth/premium_rate_limit.go +++ b/auth/premium_rate_limit.go @@ -342,6 +342,7 @@ func (s *Store) markPremium5hRateLimited(acc *Account, resetAt, observedAt time. acc.Status = StatusCooldown acc.CooldownUtil = until acc.CooldownReason = reason + acc.usageCooldownPersistPending = true if acc.HealthTier != HealthTierBanned { acc.HealthTier = HealthTierRisky } @@ -353,6 +354,7 @@ func (s *Store) markPremium5hRateLimited(acc *Account, resetAt, observedAt time. if cooldownChanged { s.setCachedAccountCooldown(acc.DBID, reason, until) } + s.persistUsageCooldown(acc) if s.db == nil || (!cooldownChanged && !snapshotChanged) { return @@ -360,11 +362,6 @@ func (s *Store) markPremium5hRateLimited(acc *Account, resetAt, observedAt time. ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second) defer cancel() - if cooldownChanged { - if err := s.db.SetCooldown(ctx, acc.DBID, reason, until); err != nil { - log.Printf("[账号 %d] 持久化 premium 5h 限流冷却状态失败: %v", acc.DBID, err) - } - } if err := s.withUsageSnapshotWrite(func() error { return s.db.UpdateUsageSnapshot5h(ctx, acc.DBID, 100, resetAt, now) }); err != nil { diff --git a/auth/store.go b/auth/store.go index 8b44a18f4..98f634ef0 100644 --- a/auth/store.go +++ b/auth/store.go @@ -279,11 +279,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+ @@ -2427,11 +2429,13 @@ func (s *Store) MarkUsage7dRateLimited(acc *Account) bool { // Do not slide the fallback window forward on every observation. if activeCooldown && acc.Reset7dAt.IsZero() { acc.mu.Unlock() + s.persistUsageCooldown(acc) return true } until, reason, changed := acc.usageCooldownTargetLocked(now, resetAt, "rate_limited") if !changed { acc.mu.Unlock() + s.persistUsageCooldown(acc) return true } @@ -2445,17 +2449,12 @@ func (s *Store) MarkUsage7dRateLimited(acc *Account) bool { acc.setCooldownUntilLocked(until, reason) } acc.recomputeSchedulerLocked(s.maxConcurrency.Load()) + acc.usageCooldownPersistPending = true acc.mu.Unlock() s.fastSchedulerUpdate(acc) s.setCachedAccountCooldown(acc.DBID, reason, until) - if s.db != nil { - 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] 持久化 7d 限流冷却状态失败: %v", acc.DBID, err) - } - } + s.persistUsageCooldown(acc) return true } @@ -2626,6 +2625,9 @@ func (a *Account) MarkUsageFrameObserved(kind string, at time.Time) { } a.mu.Lock() defer a.mu.Unlock() + if at.Before(a.usageFrameLastAt) { + return + } a.usageFrameLastAt = at a.usageFrameLastKind = strings.TrimSpace(kind) } diff --git a/auth/usage_cooldown_persist.go b/auth/usage_cooldown_persist.go new file mode 100644 index 000000000..409e942ae --- /dev/null +++ b/auth/usage_cooldown_persist.go @@ -0,0 +1,35 @@ +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) { + if s == nil || s.db == nil || acc == nil { + return + } + acc.usageCooldownPersistMu.Lock() + defer acc.usageCooldownPersistMu.Unlock() + acc.mu.RLock() + pending := acc.usageCooldownPersistPending && acc.Status == StatusCooldown + reason, until := acc.CooldownReason, acc.CooldownUtil + acc.mu.RUnlock() + if !pending { + 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) + return + } + acc.mu.Lock() + if acc.Status == StatusCooldown && 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 index 53fe989b4..a9b43f681 100644 --- a/auth/usage_cooldown_test.go +++ b/auth/usage_cooldown_test.go @@ -170,3 +170,68 @@ func TestUsageCooldownNeverShortensActiveCooldown(t *testing.T) { }) } } + +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") + } + }) + } +} diff --git a/proxy/codex_usage_frames.go b/proxy/codex_usage_frames.go index 65244a1e6..59c10a2a2 100644 --- a/proxy/codex_usage_frames.go +++ b/proxy/codex_usage_frames.go @@ -231,20 +231,16 @@ func applyCodexUsageHeaders(store *auth.Store, account *auth.Account, headers ht // 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. - if headers.Get("x-codex-rate-limit-reached-type") == "" { - if current := account.GetRateLimitReachedType(); current != "" { - headers.Set("x-codex-rate-limit-reached-type", current) - } - } 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. return syncCodexUsageState(store, account, resp, observedAt, codexUsageSyncOptions{ - persistUsageSnapshot: true, - clearAbsent5h: false, - asyncPersist: true, + persistUsageSnapshot: true, + clearAbsent5h: false, + asyncPersist: true, + preserveAbsentHardStop: true, shouldPersistUsage: func(pct5h float64, has5h bool, pct7d float64, has7d bool) bool { return account.ShouldPersistUsageFrame(observedAt, pct5h, has5h, pct7d, has7d, codexUsageFramePersistInterval()) }, diff --git a/proxy/codex_usage_frames_test.go b/proxy/codex_usage_frames_test.go index 31cd8891f..3179916b1 100644 --- a/proxy/codex_usage_frames_test.go +++ b/proxy/codex_usage_frames_test.go @@ -749,6 +749,32 @@ func TestApplyCodexUsageEventFrameIgnoresOlderObservation(t *testing.T) { } } +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") diff --git a/proxy/handler.go b/proxy/handler.go index e4cde2280..7f71e50f9 100644 --- a/proxy/handler.go +++ b/proxy/handler.go @@ -8760,6 +8760,9 @@ type codexUsageSyncOptions struct { // 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 } func syncCodexUsageState(store *auth.Store, account *auth.Account, resp *http.Response, observedAt time.Time, opts codexUsageSyncOptions) CodexUsageSyncResult { @@ -8770,51 +8773,54 @@ func syncCodexUsageState(store *auth.Store, account *auth.Account, resp *http.Re if observedAt.IsZero() { 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) - } - } - 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") + 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 { + 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 { + 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() { + result.Used5hHeaders = observation.w5h.valid + if observation.authoritative { parsed := applyCodexUsageHeaderObservation(store, account, observation, observedAt, opts.clearAbsent5h) result.UsagePct7d = parsed.usagePct7d result.HasUsage7d = parsed.hasUsage7d @@ -8866,8 +8872,13 @@ func syncCodexUsageState(store *auth.Store, account *auth.Account, resp *http.Re pct7d, has7d := account.GetUsagePercent7d() account.MarkUsageFramePersisted(observedAt, pct5h, has5h, pct7d, has7d) } - }) + } + } + applied := false + if observation.authoritative || creditsObserved || rateLimitReachedType != "" || resp.Header.Get("x-codex-plan-type") != "" { + applied = account.ApplyUsageObservation(observedAt, apply) } + usageApplied := applied && observation.authoritative result.UsagePct5h, result.Reset5hAt, result.HasUsage5h = account.GetUsageSnapshot5h() if usageApplied && store != nil && result.HasUsage5h { @@ -8886,7 +8897,7 @@ func syncCodexUsageState(store *auth.Store, account *auth.Account, resp *http.Re } } - result.Applied = usageApplied || result.CreditsObserved + result.Applied = applied && (usageApplied || result.CreditsObserved) return result } From 376b9a0c4ca3d4f297a48ce73ca5840e6c98b056 Mon Sep 17 00:00:00 2001 From: Delicious233 <101502465+DeliciousBuding@users.noreply.github.com> Date: Sun, 18 Oct 2026 19:13:14 +0800 Subject: [PATCH 21/24] fix(proxy): retain ordered quota frames and defer reader side effects Serialize partial observations without blocking persistent WebSocket readers. Coalesce quota, credits, plan and cooldown persistence, and publish current cooldown state under one lock. Cover idle transport wiring and contention through the authenticated live endpoint. Co-authored-by: Codex --- CHANGELOG.md | 2 + auth/oauth_refresh_lock.go | 7 +- auth/premium_rate_limit.go | 39 ++-- auth/store.go | 155 ++++---------- auth/usage_cooldown_persist.go | 43 +++- auth/usage_cooldown_test.go | 86 ++++++++ auth/usage_frame_observation.go | 103 +++++++++ auth/usage_frame_observation_test.go | 129 +++++++++++ auth/usage_frame_persist.go | 167 +++++++++++---- auth/usage_frame_persist_test.go | 85 ++++++++ docs/API.md | 2 + docs/CONFIGURATION.md | 4 +- main.go | 10 +- main_ws_usage_ingest_test.go | 307 +++++++++++++++++++++++++++ proxy/codex_usage_frames.go | 86 +++----- proxy/codex_usage_frames_test.go | 62 +++++- proxy/handler.go | 42 ++-- 17 files changed, 1057 insertions(+), 272 deletions(-) create mode 100644 auth/usage_frame_observation.go create mode 100644 auth/usage_frame_observation_test.go create mode 100644 main_ws_usage_ingest_test.go diff --git a/CHANGELOG.md b/CHANGELOG.md index 53404b999..95f17beee 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -8,6 +8,8 @@ ### Fixes +- **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. 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 3d67badae..d82f574e6 100644 --- a/auth/premium_rate_limit.go +++ b/auth/premium_rate_limit.go @@ -228,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 @@ -258,20 +257,10 @@ 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 } @@ -283,9 +272,7 @@ func (s *Store) ClearAbsentUsageSnapshot5hAt(acc *Account, observedAt time.Time) 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 } @@ -304,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) }) } @@ -317,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() @@ -337,6 +324,11 @@ func (s *Store) markPremium5hRateLimited(acc *Account, resetAt, observedAt time. acc.Reset5hAt = resetAt acc.UsageUpdatedAt5h = now 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 @@ -351,12 +343,11 @@ func (s *Store) markPremium5hRateLimited(acc *Account, resetAt, observedAt time. acc.mu.Unlock() s.fastSchedulerUpdate(acc) - if cooldownChanged { - s.setCachedAccountCooldown(acc.DBID, reason, until) + if persist { + s.persistUsageCooldown(acc) } - s.persistUsageCooldown(acc) - if s.db == nil || (!cooldownChanged && !snapshotChanged) { + if !persist || s.db == nil || (!cooldownChanged && !snapshotChanged) { return } diff --git a/auth/store.go b/auth/store.go index 98f634ef0..0b1d8b296 100644 --- a/auth/store.go +++ b/auth/store.go @@ -303,9 +303,15 @@ type Account struct { // usageFramePersisted* 记录 WS 逐帧观测最近一次落库的时间与百分比。 // 内存快照每帧都实时更新；落库只为重启恢复，大号池下按时间/变化幅度节流， // 避免把每个上游响应都变成一次数据库写入。 - usageFramePersistedAt time.Time - usageFramePersisted5h float64 - usageFramePersisted7d float64 + 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 @@ -346,6 +352,7 @@ type Account struct { // 业务标志变化时落库，余额逐请求递减不能每次都写。 creditsPersistedKey string creditsPersistedFlagsKey string + creditsPersistMu sync.Mutex // resetCreditsProbedAt 记录最近一次成功 wham 用量探针的时间。 // 「主动重置次数」只能通过 wham 探针刷新（普通 /responses 流量不携带该字段）， // 因此用它独立判断重置次数是否过期，避免活跃账号因用量快照一直被流量刷新而长期不探针。 @@ -2401,6 +2408,10 @@ func (a *Account) usageCooldownTargetLocked(now, resetAt time.Time, reason strin // 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 { + return s.markUsage7dRateLimited(acc, true) +} + +func (s *Store) markUsage7dRateLimited(acc *Account, persist bool) bool { if s == nil || acc == nil { return false } @@ -2429,13 +2440,17 @@ func (s *Store) MarkUsage7dRateLimited(acc *Account) bool { // Do not slide the fallback window forward on every observation. if activeCooldown && acc.Reset7dAt.IsZero() { acc.mu.Unlock() - s.persistUsageCooldown(acc) + if persist { + s.persistUsageCooldown(acc) + } return true } until, reason, changed := acc.usageCooldownTargetLocked(now, resetAt, "rate_limited") if !changed { acc.mu.Unlock() - s.persistUsageCooldown(acc) + if persist { + s.persistUsageCooldown(acc) + } return true } @@ -2453,8 +2468,9 @@ func (s *Store) MarkUsage7dRateLimited(acc *Account) bool { acc.mu.Unlock() s.fastSchedulerUpdate(acc) - s.setCachedAccountCooldown(acc.DBID, reason, until) - s.persistUsageCooldown(acc) + if persist { + s.persistUsageCooldown(acc) + } return true } @@ -2503,9 +2519,8 @@ func (a *Account) ApplyUsageObservation(observedAt time.Time, apply func()) bool } // TryApplyUsageObservation is the non-blocking variant used by connection-level -// WS frame observers. It skips the observation when another goroutine is -// already applying a usage observation (for example a preflight frame doing a -// throttled DB write), so the WebSocket read pump is never stalled. +// 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 @@ -3004,6 +3019,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) @@ -10345,16 +10362,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 选择自适应或精确时长并应用账号冷却。 @@ -10407,23 +10415,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 { @@ -10827,17 +10819,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 清除账号冷却状态，并同步清理数据库 @@ -10873,17 +10855,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。 @@ -10928,8 +10900,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 = "" @@ -10942,24 +10912,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 } @@ -11011,14 +10964,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 } @@ -11063,17 +11009,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 记录一次成功请求，用于动态调度评分 @@ -11331,6 +11267,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 } @@ -11343,6 +11288,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() @@ -11351,13 +11297,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(func() error { return s.persistUsageFramePlan(acc) }); err != nil { log.Printf("[账号 %d] 持久化 plan_type 失败: %v", acc.DBID, err) } return changed @@ -12011,12 +11955,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 index 409e942ae..bcd41cc9e 100644 --- a/auth/usage_cooldown_persist.go +++ b/auth/usage_cooldown_persist.go @@ -9,26 +9,59 @@ import ( // 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) { - if s == nil || s.db == nil || acc == nil { + 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 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 && acc.Status == StatusCooldown + pending := acc.usageCooldownPersistPending + status, errorMsg := acc.Status, acc.ErrorMsg reason, until := acc.CooldownReason, acc.CooldownUtil acc.mu.RUnlock() - if !pending { + 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() - if err := s.db.SetCooldown(ctx, acc.DBID, reason, until); err != nil { + 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 == StatusCooldown && acc.CooldownReason == reason && acc.CooldownUtil.Equal(until) { + 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 index a9b43f681..e51fb8f6e 100644 --- a/auth/usage_cooldown_test.go +++ b/auth/usage_cooldown_test.go @@ -4,6 +4,7 @@ import ( "context" "database/sql" "path/filepath" + "sync" "testing" "time" @@ -235,3 +236,88 @@ func TestUsageCooldownRetriesFailedPersistence(t *testing.T) { }) } } + +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 index 4a17b83b4..f069d3ffa 100644 --- a/auth/usage_frame_persist.go +++ b/auth/usage_frame_persist.go @@ -11,15 +11,19 @@ import ( // 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 { - dbID int64 - observedAt time.Time - pct5h float64 - reset5hAt time.Time - has5h bool - updatedAt5h time.Time - pct7d float64 - reset7dAt time.Time - has7d bool + 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 @@ -56,6 +60,11 @@ func (q *usageFramePersistQueue) run(s *Store) { return case <-q.wake: for { + select { + case <-s.stopCh: + return + default: + } job, ok := q.take() if !ok { break @@ -68,6 +77,13 @@ func (q *usageFramePersistQueue) run(s *Store) { 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 { @@ -91,7 +107,7 @@ func (q *usageFramePersistQueue) persist(s *Store, job usageFramePersistJob) { return } acc := s.FindByID(job.dbID) - if acc == nil { + if acc == nil || acc != job.account { return } // Serialize with the synchronous snapshot writers and re-check inside the @@ -99,40 +115,74 @@ func (q *usageFramePersistQueue) persist(s *Store, job usageFramePersistJob) { // its write must win. Holding usageSyncMu here would stall the upstream // read path, so use the store's snapshot-write lock instead. s.usageSnapshotPersistMu.Lock() + defer s.usageSnapshotPersistMu.Unlock() + if s.FindByID(job.dbID) != acc { + return + } acc.mu.RLock() superseded := acc.usageObservedAt.After(job.observedAt) acc.mu.RUnlock() - if superseded { - s.usageSnapshotPersistMu.Unlock() - return - } ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second) defer cancel() var err error - switch { - case job.has7d && job.has5h: - err = s.db.UpdateUsageSnapshotFull(ctx, job.dbID, job.pct7d, job.reset7dAt, job.pct5h, job.reset5hAt, job.observedAt, job.updatedAt5h) - case job.has7d: - err = s.db.UpdateUsageSnapshot(ctx, job.dbID, job.pct7d, job.observedAt) - case job.has5h: - err = s.db.UpdateUsageSnapshot5h(ctx, job.dbID, job.pct5h, job.reset5hAt, job.observedAt) - default: - s.usageSnapshotPersistMu.Unlock() - return + // 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() + } } - s.usageSnapshotPersistMu.Unlock() - 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) - return + 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 err := s.persistUsageFramePlan(acc); err != nil { + log.Printf("[账号 %d] 异步持久化套餐失败: %v", acc.DBID, err) + } + if job.planObserved { + s.ClearStaleSubscriptionExpiresAt(acc) } - acc.MarkUsageFramePersisted(job.observedAt, job.pct5h, job.has5h, job.pct7d, job.has7d) } // 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 } @@ -144,19 +194,46 @@ func (s *Store) EnqueueUsageFramePersist(acc *Account, observedAt time.Time) { if observedAt.IsZero() { observedAt = time.Now() } - snap := acc.GetAccountLiveUsageSnapshot() - if !snap.UsagePercent5hValid && !snap.UsagePercent7dValid { - return + 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 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 } - s.usageFramePersist.enqueue(usageFramePersistJob{ - dbID: acc.DBID, - observedAt: observedAt, - pct5h: snap.UsagePercent5h, - reset5hAt: snap.Reset5hAt, - has5h: snap.UsagePercent5hValid, - updatedAt5h: snap.UsageUpdatedAt5h, - pct7d: snap.UsagePercent7d, - reset7dAt: snap.Reset7dAt, - has7d: snap.UsagePercent7dValid, - }) + acc.mu.Unlock() + return nil } diff --git a/auth/usage_frame_persist_test.go b/auth/usage_frame_persist_test.go index cdcd2bef7..d95ccea77 100644 --- a/auth/usage_frame_persist_test.go +++ b/auth/usage_frame_persist_test.go @@ -7,6 +7,7 @@ import ( "testing" "time" + "github.com/codex2api/cache" "github.com/codex2api/database" "github.com/tidwall/gjson" ) @@ -61,3 +62,87 @@ func TestUsageFramePersistSkipsSupersededJob(t *testing.T) { 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/docs/API.md b/docs/API.md index 1c8e6e3aa..858dc4a28 100644 --- a/docs/API.md +++ b/docs/API.md @@ -824,6 +824,8 @@ Codex 的流式 remote compact v2（`POST /v1/responses`，`stream:true`，`inpu 缺失的 key 表示旧后端，前端保留列表值；显式 `null` 表示窗口失效。Antigravity 行不使用 Codex 用量字段。 +终态后到达的连接级配额帧也会更新上述字段。只有通过观测排序门的帧才推进帧时间、credits 和用量；遇到并发观测锁争用时按原始时间排队应用，不等待数据库。用量、credits、套餐与冷却的数据库写入在后台合并执行，因此 live 结果可先于持久化结果可见。 + #### PATCH /api/admin/accounts/:id/scheduler 更新账号调度配置。 diff --git a/docs/CONFIGURATION.md b/docs/CONFIGURATION.md index eccdad9ad..49c175fff 100644 --- a/docs/CONFIGURATION.md +++ b/docs/CONFIGURATION.md @@ -88,7 +88,7 @@ Codex2API 采用三层配置架构： | `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_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 | @@ -100,6 +100,8 @@ 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/main.go b/main.go index d39c4ed6b..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) @@ -446,9 +452,7 @@ func main() { // 注册 WebSocket 执行函数（避免 proxy ↔ wsrelay 循环依赖） proxy.WebsocketExecuteFunc = wsrelay.ExecuteRequestWebsocket // 连接级（终态后补发）的用量帧也交给 proxy 更新内存快照；observer 必须非阻塞。 - wsrelay.SetIdleMetadataObserver(func(account *auth.Account, eventType string, payload []byte) { - proxy.ApplyCodexIdleUsageFrame(store, account, eventType, payload) - }) + 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 { diff --git a/main_ws_usage_ingest_test.go b/main_ws_usage_ingest_test.go new file mode 100644 index 000000000..1ff57850b --- /dev/null +++ b/main_ws_usage_ingest_test.go @@ -0,0 +1,307 @@ +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"` + } `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} + 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","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-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 && !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 index 59c10a2a2..24941259e 100644 --- a/proxy/codex_usage_frames.go +++ b/proxy/codex_usage_frames.go @@ -87,9 +87,7 @@ func applyCodexRateLimitsFrame(store *auth.Store, account *auth.Account, frame g if len(headers) == 0 { return } - if result := applyCodexUsageHeaders(store, account, headers, observedAt); result.Applied { - account.MarkUsageFrameObserved(codexRateLimitsFrameType, observedAt) - } + applyCodexUsageHeaders(store, account, headers, observedAt, codexRateLimitsFrameType) } func codexRateLimitsHeaders(frame gjson.Result, observedAt time.Time) http.Header { @@ -145,9 +143,7 @@ func applyCodexResponseMetadataFrame(store *auth.Store, account *auth.Account, f if len(headers) == 0 { return } - if result := applyCodexUsageHeaders(store, account, headers, observedAt); result.Applied { - account.MarkUsageFrameObserved(codexResponseMetadataFrameType, observedAt) - } + applyCodexUsageHeaders(store, account, headers, observedAt, codexResponseMetadataFrameType) } func codexResponseMetadataHeaders(frame gjson.Result, observedAt time.Time) http.Header { @@ -223,9 +219,12 @@ func isCodexUsageFrameHeader(name string) bool { return false } -func applyCodexUsageHeaders(store *auth.Store, account *auth.Account, headers http.Header, observedAt time.Time) CodexUsageSyncResult { +func applyCodexUsageHeaders(store *auth.Store, account *auth.Account, headers http.Header, observedAt time.Time, kind string) { if len(headers) == 0 { - return CodexUsageSyncResult{} + 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 @@ -236,14 +235,23 @@ func applyCodexUsageHeaders(store *auth.Store, account *auth.Account, headers ht // 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. - return syncCodexUsageState(store, account, resp, observedAt, codexUsageSyncOptions{ - persistUsageSnapshot: true, - clearAbsent5h: false, - asyncPersist: true, - preserveAbsentHardStop: true, - shouldPersistUsage: func(pct5h float64, has5h bool, pct7d float64, has7d bool) bool { - return account.ShouldPersistUsageFrame(observedAt, pct5h, has5h, pct7d, has7d, codexUsageFramePersistInterval()) - }, + 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) + } }) } @@ -335,11 +343,10 @@ func isDefaultCodexLimitName(name string) bool { // 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). It updates only the in-memory usage -// snapshot, with a non-blocking try-lock, so the read pump never waits on a DB -// write. Persistence and the 100% cooldown markings stay on the preflight/WHAM -// paths; this path exists to shrink the window between a turn finishing and -// the scheduler seeing its consumption. +// 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 @@ -365,40 +372,5 @@ func ApplyCodexIdleUsageFrame(store *auth.Store, account *auth.Account, eventTyp if len(headers) == 0 { return } - resp := &http.Response{Header: headers} - observation := parseCodexUsageHeaderObservation(resp) - hasCredits, unlimited, creditsObserved := parseCodexCreditsFlagHeaders(resp) - if !observation.authoritative && !creditsObserved { - return - } - applied := account.TryApplyUsageObservation(observedAt, func() { - account.MarkUsageFrameObserved(normalizedType, observedAt) - if observation.authoritative { - // Never clear an absent 5h window from a partial connection-level frame. - applyCodexUsageHeaderObservation(nil, account, observation, observedAt, false) - } - if creditsObserved { - // Connection-level frames also carry the post-turn credits snapshot. - // Apply it so the last spent credit can flip credit-based scheduling - // before the next WHAM probe. Keep the stored hard stop: the frame has - // no rate_limit_reached_type of its own. - var balance *string - if trimmed := strings.TrimSpace(headers.Get("x-codex-credits-balance")); trimmed != "" { - balance = &trimmed - } - var overage *bool - if raw := strings.TrimSpace(headers.Get("x-codex-credits-overage-reached")); raw != "" { - if value, err := strconv.ParseBool(raw); err == nil { - overage = &value - } - } - account.ApplySparseCreditsHeaders(hasCredits, unlimited, balance, overage, account.GetRateLimitReachedType()) - } - }) - if applied { - // A frame can restore an account that crossed back under its threshold; - // queued dispatchers wait on the availability notification, so refresh - // the scheduler even though persistence stays on the preflight path. - store.RefreshUsageFrameScheduler(account) - } + applyCodexUsageHeaders(store, account, headers, observedAt, normalizedType) } diff --git a/proxy/codex_usage_frames_test.go b/proxy/codex_usage_frames_test.go index 3179916b1..b273bc76c 100644 --- a/proxy/codex_usage_frames_test.go +++ b/proxy/codex_usage_frames_test.go @@ -831,6 +831,19 @@ func TestCodexUsageFramePersistenceRespectsThrottle(t *testing.T) { // 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) @@ -865,19 +878,50 @@ func TestCodexUsageFramePersistDoesNotBlockOnLockedDB(t *testing.T) { t.Fatalf("begin immediate: %v", err) } - frame := gjson.Parse(`{"type":"codex.rate_limits","rate_limits":{"primary":{"used_percent":42,"window_minutes":300}}}`) - start := time.Now() - ApplyCodexUsageEventFrame(store, account, "codex.rate_limits", frame, time.Now()) - if elapsed := time.Since(start); elapsed > 500*time.Millisecond { - t.Fatalf("frame path blocked %s on a locked database, want enqueue-only", elapsed) - } - if pct, ok := account.GetUsagePercent5h(); !ok || pct != 42 { - t.Fatalf("memory usage = %v valid=%v, want 42", pct, ok) + 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) } - waitPersisted5h(t, dbPath, accountID, 42) + 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) { diff --git a/proxy/handler.go b/proxy/handler.go index 7f71e50f9..0493ff2c1 100644 --- a/proxy/handler.go +++ b/proxy/handler.go @@ -8763,6 +8763,8 @@ type codexUsageSyncOptions struct { // 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 { @@ -8779,10 +8781,14 @@ func syncCodexUsageState(store *auth.Store, account *auth.Account, resp *http.Re apply := func() { 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 opts.asyncPersist { + store.ObserveUsageFramePlan(account, planHeader) + } else { + store.UpdateAccountPlanType(account, planHeader) + // 权威付费 plan_type 与「订阅已过期」互斥，借每次响应校正陈旧到期时间。(issue #360) + if planHeader != "" { + store.ClearStaleSubscriptionExpiresAt(account) + } } } @@ -8802,7 +8808,7 @@ func syncCodexUsageState(store *auth.Store, account *auth.Account, resp *http.Re overagePtr = &ov } } - if store != nil { + if store != nil && !opts.asyncPersist { store.PersistSparseCreditObservation(account, hasCredits, unlimited, balPtr, overagePtr, rateLimitReachedType) } else { account.ApplySparseCreditsHeaders(hasCredits, unlimited, balPtr, overagePtr, rateLimitReachedType) @@ -8810,7 +8816,7 @@ func syncCodexUsageState(store *auth.Store, account *auth.Account, resp *http.Re } else if rateLimitReachedType != "" || observation.authoritative { // reached-type 是逐响应字段：带用量头的响应没有它就是「未触达」，要把旧的 // workspace hard-stop 清掉，否则工作区充值后积分顶替永远回不来。 - if store != nil { + if store != nil && !opts.asyncPersist { store.PersistRateLimitReachedType(account, rateLimitReachedType) } else { account.SetRateLimitReachedType(rateLimitReachedType) @@ -8844,21 +8850,15 @@ func syncCodexUsageState(store *auth.Store, account *auth.Account, resp *http.Re // Memory and scheduler update stay synchronous; the DB write // is coalesced per account off the read path. store.PublishUsageSnapshot(account, result.UsagePct7d, observedAt, false) - if persistUsage { - store.EnqueueUsageFramePersist(account, observedAt) - } } else { persistedUsage = store.PublishUsageSnapshot(account, result.UsagePct7d, observedAt, persistUsage) } - if result.UsagePct7d >= 100 && persistUsage { + if result.UsagePct7d >= 100 && persistUsage && !opts.asyncPersist { result.Usage7dRateLimited = store.MarkUsage7dRateLimited(account) } } else if result.Used5hHeaders { if opts.asyncPersist { store.PublishUsageSnapshot5hOnly(account, observedAt, false) - if persistUsage { - store.EnqueueUsageFramePersist(account, observedAt) - } } else { persistedUsage = store.PublishUsageSnapshot5hOnly(account, observedAt, persistUsage) } @@ -8872,11 +8872,19 @@ func syncCodexUsageState(store *auth.Store, account *auth.Account, resp *http.Re 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") != "" { - applied = account.ApplyUsageObservation(observedAt, apply) + if opts.observationLocked { + apply() + applied = true + } else { + applied = account.ApplyUsageObservation(observedAt, apply) + } } usageApplied := applied && observation.authoritative @@ -8885,7 +8893,7 @@ func syncCodexUsageState(store *auth.Store, account *auth.Account, resp *http.Re // 被动 /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 { // The per-frame path may decide not to persist this observation; a // repeated 100% frame must not issue another cooldown write. @@ -8898,6 +8906,10 @@ func syncCodexUsageState(store *auth.Store, account *auth.Account, resp *http.Re } 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 } From 1c759fa3b4ef37b0329521051c2482dc3710ac38 Mon Sep 17 00:00:00 2001 From: Delicious233 <101502465+DeliciousBuding@users.noreply.github.com> Date: Sun, 18 Oct 2026 19:26:38 +0800 Subject: [PATCH 22/24] fix(proxy): reject malformed quota frame reset numbers Parse the original integer representation and bound offsets before duration conversion. Cover fractional, exponent and overflow values in compact and metadata frames without changing HTTP header semantics. Co-authored-by: Codex --- CHANGELOG.md | 2 ++ proxy/codex_usage_frames.go | 30 ++++++++++++++++++++---- proxy/codex_usage_frames_test.go | 40 +++++++++++++++++++++++++++++++- 3 files changed, 66 insertions(+), 6 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 95f17beee..512bf74c1 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -8,6 +8,8 @@ ### Fixes +- **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. diff --git a/proxy/codex_usage_frames.go b/proxy/codex_usage_frames.go index 24941259e..a5bfe0f81 100644 --- a/proxy/codex_usage_frames.go +++ b/proxy/codex_usage_frames.go @@ -15,6 +15,8 @@ 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 @@ -61,8 +63,9 @@ func codexUsageFramePersistInterval() time.Duration { // 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 this adds no -// upstream requests and no new request headers. +// 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 @@ -179,6 +182,9 @@ func codexResponseMetadataHeaders(frame gjson.Result, observedAt time.Time) http epoch := codexFrameInt64(value) if epoch > 0 { seconds := epoch - observedAt.Unix() + if seconds > maxCodexFrameResetSeconds { + return true + } if seconds < 0 { seconds = 0 } @@ -187,6 +193,14 @@ func codexResponseMetadataHeaders(frame gjson.Result, observedAt time.Time) http } 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 @@ -279,7 +293,7 @@ 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 { + 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 @@ -298,6 +312,9 @@ func codexFrameResetSeconds(window gjson.Result, observedAt time.Time) string { continue } seconds := epoch - observedAt.Unix() + if seconds > maxCodexFrameResetSeconds { + continue + } if seconds < 0 { seconds = 0 } @@ -311,10 +328,13 @@ func codexFrameResetSeconds(window gjson.Result, observedAt time.Time) string { } 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 { - return value.Int(), true + raw = value.Raw } - n, err := strconv.ParseInt(strings.TrimSpace(value.String()), 10, 64) + n, err := strconv.ParseInt(strings.TrimSpace(raw), 10, 64) if err != nil { return 0, false } diff --git a/proxy/codex_usage_frames_test.go b/proxy/codex_usage_frames_test.go index b273bc76c..05e2c0241 100644 --- a/proxy/codex_usage_frames_test.go +++ b/proxy/codex_usage_frames_test.go @@ -399,7 +399,7 @@ func TestApplyCodexRateLimitsFramePreservesKnownReset(t *testing.T) { } } -func TestApplyCodexIdleUsageFrameUpdatesMemoryOnly(t *testing.T) { +func TestApplyCodexIdleUsageFrameWithoutStoreUpdatesRuntime(t *testing.T) { acc := newUsageFrameTestAccount() now := time.Now() acc.SetUsageSnapshot5hAt(20, now.Add(time.Hour), now) @@ -538,6 +538,44 @@ func TestCodexRateLimitsFrameTreatsZeroRelativeResetAsNow(t *testing.T) { } } +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() From 85f1a3af979e9c79b924ac70766e1ee4c864eb3f Mon Sep 17 00:00:00 2001 From: Delicious233 <101502465+DeliciousBuding@users.noreply.github.com> Date: Sun, 18 Oct 2026 19:44:59 +0800 Subject: [PATCH 23/24] fix(admin): project frame plan and credit details in live polls Complete the original live-sync contract with existing account-list fields. Preserve absent keys, clear hidden balances and leave Antigravity fields untouched. Exercise plan and balance changes through the production persistent WebSocket observer and authenticated live API. Co-authored-by: Codex --- CHANGELOG.md | 2 ++ admin/account_live.go | 24 +++++++++++++++++++ admin/account_live_test.go | 16 ++++++++++++- auth/store.go | 6 +++++ docs/API.md | 14 ++++++++++- frontend/src/lib/accountLiveState.test.mjs | 27 ++++++++++++++++++++++ frontend/src/lib/accountLiveState.ts | 24 +++++++++++++++++++ frontend/src/types.ts | 8 +++++++ main_ws_usage_ingest_test.go | 24 +++++++++++-------- 9 files changed, 133 insertions(+), 12 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 512bf74c1..42a8e04a5 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -8,6 +8,8 @@ ### 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. diff --git a/admin/account_live.go b/admin/account_live.go index 17f6f4178..1e7ebaa81 100644 --- a/admin/account_live.go +++ b/admin/account_live.go @@ -8,6 +8,8 @@ import ( ) 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 @@ -24,6 +26,7 @@ type accountLiveItem struct { // 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"` @@ -36,6 +39,17 @@ type accountLiveItem struct { 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 usage @@ -68,6 +82,16 @@ func (h *Handler) GetAccountLiveState(c *gin.Context) { // 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 diff --git a/admin/account_live_test.go b/admin/account_live_test.go index 87e908abe..38b302e76 100644 --- a/admin/account_live_test.go +++ b/admin/account_live_test.go @@ -80,6 +80,9 @@ func TestGetAccountLiveStateReturnsRuntimeUsage(t *testing.T) { 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() @@ -112,6 +115,17 @@ func TestGetAccountLiveStateReturnsRuntimeUsage(t *testing.T) { 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) { @@ -161,7 +175,7 @@ func TestGetAccountLiveStateSerializesNullWindowsAndSkipsAntigravityFields(t *te t.Fatalf("codex frame time missing: %s", got) } antigravityItem := raw.Accounts["12"] - for _, key := range []string{"status", "using_credits", "codex_usage_frame_at", "codex_usage_frame_kind"} { + 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) } diff --git a/auth/store.go b/auth/store.go index 0b1d8b296..d79bddaa0 100644 --- a/auth/store.go +++ b/auth/store.go @@ -2653,6 +2653,9 @@ func (a *Account) MarkUsageFrameObserved(kind string, at time.Time) { type AccountLiveUsageSnapshot struct { Status string UsingCredits bool + PlanType string + CreditsValid bool + Credits CreditBalanceSnapshot UsagePercent5h float64 UsagePercent5hValid bool Reset5hAt time.Time @@ -2679,6 +2682,9 @@ func (a *Account) GetAccountLiveUsageSnapshot() AccountLiveUsageSnapshot { 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, diff --git a/docs/API.md b/docs/API.md index 858dc4a28..d57bf1e1d 100644 --- a/docs/API.md +++ b/docs/API.md @@ -796,6 +796,14 @@ Codex 的流式 remote compact v2（`POST /v1/responses`，`stream:true`，`inpu "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", @@ -815,6 +823,10 @@ Codex 的流式 remote compact v2（`POST /v1/responses`，`stream:true`，`inpu | 字段 | 类型 | 说明 | | ---- | ---- | ---- | | 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 每轮配额帧也会刷新 | @@ -822,7 +834,7 @@ Codex 的流式 remote compact v2（`POST /v1/responses`，`stream:true`，`inpu | dynamic_concurrency_limit / base_concurrency_effective | integer | 调度器当前生效的并发上限与配置值 | | session_slot_buffer_enabled | bool | 全局会话槽缓冲开关 | -缺失的 key 表示旧后端，前端保留列表值；显式 `null` 表示窗口失效。Antigravity 行不使用 Codex 用量字段。 +缺失的 key 表示旧后端，前端保留列表值；用量字段显式 `null` 表示窗口失效。套餐与 credits 的 live 字段只投影到使用 Codex 运行时字段的行，Antigravity 行省略这些字段，不改变原列表展示。 终态后到达的连接级配额帧也会更新上述字段。只有通过观测排序门的帧才推进帧时间、credits 和用量；遇到并发观测锁争用时按原始时间排队应用，不等待数据库。用量、credits、套餐与冷却的数据库写入在后台合并执行，因此 live 结果可先于持久化结果可见。 diff --git a/frontend/src/lib/accountLiveState.test.mjs b/frontend/src/lib/accountLiveState.test.mjs index a09188cf1..f4c3ad8c0 100644 --- a/frontend/src/lib/accountLiveState.test.mjs +++ b/frontend/src/lib/accountLiveState.test.mjs @@ -6,6 +6,10 @@ 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", @@ -29,6 +33,10 @@ function live(overrides = {}) { 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, @@ -136,3 +144,22 @@ test("session slot buffer change is merged into rows", () => { 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 index b44896447..636199cff 100644 --- a/frontend/src/lib/accountLiveState.ts +++ b/frontend/src/lib/accountLiveState.ts @@ -14,6 +14,14 @@ export type AccountLiveMergeable = { 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 @@ -65,8 +73,23 @@ export function mergeAccountLiveState( 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 && @@ -85,6 +108,7 @@ export function mergeAccountLiveState( changed = true return { ...account, + ...details, active_requests: activeRequests, occupied_requests: occupiedRequests, session_slot_buffer_enabled: slotBufferEnabled, diff --git a/frontend/src/types.ts b/frontend/src/types.ts index baaa4761c..05fabaf0c 100644 --- a/frontend/src/types.ts +++ b/frontend/src/types.ts @@ -631,6 +631,14 @@ export interface AccountLiveStateResponse { // 运行时用量投影：旧后端缺 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 diff --git a/main_ws_usage_ingest_test.go b/main_ws_usage_ingest_test.go index 1ff57850b..4bbee3492 100644 --- a/main_ws_usage_ingest_test.go +++ b/main_ws_usage_ingest_test.go @@ -90,12 +90,16 @@ func TestMainPersistentWebSocketIngestsLastPostTerminalUsageFrame(t *testing.T) 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"` + 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 { @@ -108,7 +112,7 @@ func TestMainPersistentWebSocketIngestsLastPostTerminalUsageFrame(t *testing.T) if !ok { t.Fatal("live endpoint omitted the synthetic account") } - result := auth.AccountLiveUsageSnapshot{Status: item.Status, UsageFrameLastKind: item.FrameKind} + 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 } @@ -145,9 +149,9 @@ func TestMainPersistentWebSocketIngestsLastPostTerminalUsageFrame(t *testing.T) 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","rate_limits":{"primary":{"used_percent":96,"window_minutes":300,"reset_at":%d}}}`, reset.Unix()) + 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-primary-used-percent":"96","x-codex-primary-window-minutes":"300","x-codex-primary-reset-at":"%d"}}`, reset.Unix()) + 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) @@ -270,7 +274,7 @@ func TestMainPersistentWebSocketIngestsLastPostTerminalUsageFrame(t *testing.T) 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 && !account.IsAvailable() && store.AvailableCount() == 0 { + 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 { From c07028a8539ce9ba10a9061f05610900c94bf147 Mon Sep 17 00:00:00 2001 From: Delicious233 <101502465+DeliciousBuding@users.noreply.github.com> Date: Sun, 18 Oct 2026 21:28:35 +0800 Subject: [PATCH 24/24] fix(auth): scope quota snapshot writes to each account Co-authored-by: Codex --- auth/premium_rate_limit.go | 6 +-- auth/store.go | 28 +++++++------ auth/usage_cooldown_persist.go | 2 +- auth/usage_frame_persist.go | 16 ++++---- auth/usage_frame_persist_test.go | 67 +++++++++++++++++++++++++++++++- 5 files changed, 91 insertions(+), 28 deletions(-) diff --git a/auth/premium_rate_limit.go b/auth/premium_rate_limit.go index d82f574e6..df16daf2d 100644 --- a/auth/premium_rate_limit.go +++ b/auth/premium_rate_limit.go @@ -163,7 +163,7 @@ func (s *Store) PublishUsageSnapshot5hOnly(acc *Account, observedAt time.Time, p ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second) defer cancel() - err := s.withUsageSnapshotWrite(func() error { + err := s.withUsageSnapshotWrite(acc, func() error { return s.db.UpdateUsageSnapshot5h(ctx, acc.DBID, pct5h, reset5hAt, observedAt) }) if err != nil { @@ -266,7 +266,7 @@ func (s *Store) ClearAbsentUsageSnapshot5hAt(acc *Account, observedAt time.Time) ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second) defer cancel() - if err := s.withUsageSnapshotWrite(func() error { + if err := s.withUsageSnapshotWrite(acc, func() error { return s.db.ClearUsageSnapshot5h(ctx, acc.DBID) }); err != nil { log.Printf("[账号 %d] 清除 5h 用量快照失败: %v", acc.DBID, err) @@ -353,7 +353,7 @@ func (s *Store) markPremium5hRateLimited(acc *Account, resetAt, observedAt time. ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second) defer cancel() - if err := s.withUsageSnapshotWrite(func() error { + 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/store.go b/auth/store.go index d79bddaa0..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. @@ -3883,14 +3885,10 @@ type Store struct { lazyRefreshInFlight sync.Map stopCh chan struct{} usageFramePersist *usageFramePersistQueue - // usageSnapshotPersistMu serializes 5h/7d credential writes between the - // synchronous observation paths and the async frame-persist worker, so a - // queued older frame cannot overwrite a newer WHAM observation. - usageSnapshotPersistMu sync.Mutex - stopOnce sync.Once - backgroundCtx context.Context - backgroundCancel context.CancelFunc - wg sync.WaitGroup + stopOnce sync.Once + backgroundCtx context.Context + backgroundCancel context.CancelFunc + wg sync.WaitGroup // 代理池 proxyPoolReloadMu sync.Mutex @@ -11138,7 +11136,7 @@ func (s *Store) PublishUsageSnapshot(acc *Account, pct7d float64, observedAt tim reset7dAt := acc.GetReset7dAt() ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second) defer cancel() - err := s.withUsageSnapshotWrite(func() error { + err := s.withUsageSnapshotWrite(acc, func() error { return s.db.UpdateUsageSnapshotFull(ctx, acc.DBID, pct7d, reset7dAt, pct5h, reset5hAt, observedAt, acc.GetUsageUpdatedAt5h()) }) if err != nil { @@ -11150,7 +11148,7 @@ func (s *Store) PublishUsageSnapshot(acc *Account, pct7d float64, observedAt tim ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second) defer cancel() - err := s.withUsageSnapshotWrite(func() error { + err := s.withUsageSnapshotWrite(acc, func() error { return s.db.UpdateUsageSnapshot(ctx, acc.DBID, pct7d, observedAt) }) if err != nil { @@ -11162,9 +11160,9 @@ func (s *Store) PublishUsageSnapshot(acc *Account, pct7d float64, observedAt tim // withUsageSnapshotWrite serializes credential snapshot writes with the async // frame-persist worker. -func (s *Store) withUsageSnapshotWrite(fn func() error) error { - s.usageSnapshotPersistMu.Lock() - defer s.usageSnapshotPersistMu.Unlock() +func (s *Store) withUsageSnapshotWrite(acc *Account, fn func() error) error { + acc.usageSnapshotPersistMu.Lock() + defer acc.usageSnapshotPersistMu.Unlock() return fn() } @@ -11307,7 +11305,7 @@ func (s *Store) updateAccountPlanType(acc *Account, planType string, persist boo return changed } - if err := s.withUsageSnapshotWrite(func() error { return s.persistUsageFramePlan(acc) }); 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 diff --git a/auth/usage_cooldown_persist.go b/auth/usage_cooldown_persist.go index bcd41cc9e..102987d1a 100644 --- a/auth/usage_cooldown_persist.go +++ b/auth/usage_cooldown_persist.go @@ -15,7 +15,7 @@ func (s *Store) persistUsageCooldown(acc *Account) { // 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 usageSnapshotPersistMu. +// Never hold this lock while acquiring acc.usageSnapshotPersistMu. func (s *Store) persistAccountCooldownState(acc *Account, force, clearError bool) { if s == nil || acc == nil { return diff --git a/auth/usage_frame_persist.go b/auth/usage_frame_persist.go index f069d3ffa..1f69bd8ce 100644 --- a/auth/usage_frame_persist.go +++ b/auth/usage_frame_persist.go @@ -113,10 +113,10 @@ func (q *usageFramePersistQueue) persist(s *Store, job usageFramePersistJob) { // 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 the store's snapshot-write lock instead. - s.usageSnapshotPersistMu.Lock() - defer s.usageSnapshotPersistMu.Unlock() + // 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() @@ -155,6 +155,11 @@ func (q *usageFramePersistQueue) persist(s *Store, job usageFramePersistJob) { 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 @@ -165,9 +170,6 @@ func (q *usageFramePersistQueue) persist(s *Store, job usageFramePersistJob) { if cooldownPending { s.persistUsageCooldown(acc) } - if err := s.persistUsageFramePlan(acc); err != nil { - log.Printf("[账号 %d] 异步持久化套餐失败: %v", acc.DBID, err) - } if job.planObserved { s.ClearStaleSubscriptionExpiresAt(acc) } @@ -217,7 +219,7 @@ func (s *Store) EnqueueUsageFrameEffects(acc *Account, observedAt time.Time, per s.usageFramePersist.enqueue(job) } -// Called with usageSnapshotPersistMu held, including synchronous plan writes. +// 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 diff --git a/auth/usage_frame_persist_test.go b/auth/usage_frame_persist_test.go index d95ccea77..fe2203f63 100644 --- a/auth/usage_frame_persist_test.go +++ b/auth/usage_frame_persist_test.go @@ -4,6 +4,7 @@ import ( "context" "database/sql" "path/filepath" + "sync" "testing" "time" @@ -12,6 +13,68 @@ import ( "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 @@ -40,13 +103,13 @@ func TestUsageFramePersistSkipsSupersededJob(t *testing.T) { // Block the worker inside the write lock so the newer observation is // applied before its superseded check runs. - store.usageSnapshotPersistMu.Lock() + 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") } - store.usageSnapshotPersistMu.Unlock() + account.usageSnapshotPersistMu.Unlock() time.Sleep(200 * time.Millisecond) conn, err := sql.Open("sqlite", dbPath)