From d0180b74b3eed06073a65528cbca502f0e1ca6bc Mon Sep 17 00:00:00 2001 From: 2han9wen71an <2han9wen71an@gmail.com> Date: Sat, 19 Sep 2026 18:13:33 +0800 Subject: [PATCH] fix(openai): refresh expired quota windows after reset grace (#170) Refresh non-exhausted accounts after the quota reset grace period while preserving observed usage and exhausted-window attribution. Keep periodic checks for exhausted accounts to detect early upstream resets, and cover refresh timing, throttling, and display regressions. Co-authored-by: Zyy Co-authored-by: Codex --- .../openai/src/credential/quota/mod.rs | 105 ++++++++++------ .../crates/providers/openai/tests/admin.rs | 115 ++++++++++++++++++ .../tests/credential/quota/refresh_timing.rs | 94 ++++++++++++++ docs/api.md | 6 +- 4 files changed, 282 insertions(+), 38 deletions(-) diff --git a/backend/crates/providers/openai/src/credential/quota/mod.rs b/backend/crates/providers/openai/src/credential/quota/mod.rs index b31b4d0a2..3167c797f 100644 --- a/backend/crates/providers/openai/src/credential/quota/mod.rs +++ b/backend/crates/providers/openai/src/credential/quota/mod.rs @@ -60,8 +60,8 @@ use snapshot::{ const DEFAULT_RATE_LIMIT_COOLDOWN: Duration = Duration::from_secs(60); pub(crate) const QUOTA_SCHEDULING_TTL: Duration = Duration::from_secs(10 * 60); const QUOTA_HYDRATION_FAILURE_TTL: Duration = Duration::from_secs(5); -const EXHAUSTED_QUOTA_REFRESH_RETRY_INTERVAL: Duration = Duration::from_secs(30 * 60); -const EXHAUSTED_QUOTA_RESET_GRACE: Duration = Duration::from_secs(2 * 60); +const PERIODIC_QUOTA_REFRESH_RETRY_INTERVAL: Duration = Duration::from_secs(30 * 60); +const QUOTA_RESET_GRACE: Duration = Duration::from_secs(2 * 60); /// 首次 OAuth 异步观察失败时，由既有 quota worker 兜底重试的单轮上限。 const INITIAL_QUOTA_SYNC_BATCH: usize = 100; // 5xx 上游拒绝的短退避重试预算；指数退避 1s/2s，吞掉瞬时抖动。 @@ -394,15 +394,19 @@ impl CodexQuotaSchedulingProjection { fn reserve_periodic_refreshes( &self, accounts: Vec, + observed_snapshots: &BTreeMap, now: SystemTime, ) -> Vec { let candidates = accounts .into_iter() - .filter_map(|account| quota_refresh_candidate(account, now)) + .filter_map(|account| { + let snapshot = observed_snapshots.get(account.id()); + quota_refresh_candidate(account, snapshot, now) + }) .collect::>(); let candidate_ids = candidates .iter() - .map(|account| account.id().clone()) + .map(|(account, _)| account.id().clone()) .collect::>(); let refreshed_at = Instant::now(); let mut state = self @@ -413,18 +417,12 @@ impl CodexQuotaSchedulingProjection { .last_periodic_refresh_at .retain(|account_id, _| candidate_ids.contains(account_id)); - // 正常账号只由真实请求的响应头和 `codex.rate_limits` 被动同步。 - // 已耗尽账号每 30 分钟复核，不能等待旧 reset：官方活动可能提前重置额度。 - // reset + 2 分钟可提前触发一次复核，给上游重置留出传播时间。 + // 已耗尽账号首次立即复核，之后每 30 分钟复核，以发现官方提前重置。 + // reset + 2 分钟额外触发一次复核，给上游重置留出传播时间。 + // 正常账号仅在非零用量窗口经过宽限期后参与，未更新时复用周期节流。 let mut reserved = Vec::new(); - for account in candidates { - if !periodic_quota_refresh_due( - &state, - account.id(), - account.quota().reset_at(), - now, - refreshed_at, - ) { + for (account, target_reset) in candidates { + if !periodic_quota_refresh_due(&state, account.id(), target_reset, now, refreshed_at) { continue; } state.last_periodic_refresh_at.insert( @@ -440,9 +438,35 @@ impl CodexQuotaSchedulingProjection { } } -fn quota_refresh_candidate(account: ProviderAccount, now: SystemTime) -> Option { - (eligible_periodic_quota_refresh(&account, now) && account.quota().is_exhausted()) - .then_some(account) +fn quota_refresh_candidate( + account: ProviderAccount, + snapshot: Option<&CodexAccountQuotaSnapshot>, + now: SystemTime, +) -> Option<(ProviderAccount, Option)> { + if !eligible_periodic_quota_refresh(&account, now) { + return None; + } + if account.quota().is_exhausted() { + let reset_at = account.quota().reset_at(); + return Some((account, reset_at)); + } + // 正常账号首次进入周期候选也要等待宽限期，不能由“无刷新历史”绕过。 + let snapshot = snapshot?; + let expired_window_reset = snapshot + .windows() + .iter() + .filter(|window| { + window.used_percent().is_some_and(|used| used > 0.0) || window.limit_reached() + }) + .filter_map(CodexQuotaWindow::reset_at) + .map(SystemTime::from) + .filter(|reset| { + reset + .checked_add(QUOTA_RESET_GRACE) + .is_some_and(|due_at| due_at <= now) + }) + .min(); + expired_window_reset.map(|reset_at| (account, Some(reset_at))) } fn periodic_quota_refresh_due( @@ -457,9 +481,9 @@ fn periodic_quota_refresh_due( .get(account_id) .is_none_or(|last| { monotonic_now.saturating_duration_since(last.monotonic_at) - >= EXHAUSTED_QUOTA_REFRESH_RETRY_INTERVAL + >= PERIODIC_QUOTA_REFRESH_RETRY_INTERVAL || reset_at - .and_then(|reset| reset.checked_add(EXHAUSTED_QUOTA_RESET_GRACE)) + .and_then(|reset| reset.checked_add(QUOTA_RESET_GRACE)) // 已在该边界之后复核过时回到周期重试，避免过期 reset 每轮触发。 .is_some_and(|due_at| last.wall_at < due_at && due_at <= now) }) @@ -781,8 +805,26 @@ impl CodexCredentialQuotaService { }); let mut summary = CodexQuotaSyncSummary::default(); let now = SystemTime::now(); - let initial = self.initial_quota_sync_accounts(&accounts, now).await?; - let periodic = self.scheduling.reserve_periodic_refreshes(accounts, now); + let account_ids = accounts + .iter() + .map(|account| account.id().clone()) + .collect::>(); + let observed = self.store.get_quotas(&account_ids).await?; + let observed_snapshots = observed + .iter() + .filter_map(|obs| { + quota_snapshot_from_observation(obs) + .map(|snapshot| (obs.account_id.clone(), snapshot)) + }) + .collect::>(); + let observed_ids = observed + .into_iter() + .map(|observation| observation.account_id) + .collect::>(); + let initial = Self::initial_quota_sync_accounts(&accounts, &observed_ids, now); + let periodic = + self.scheduling + .reserve_periodic_refreshes(accounts, &observed_snapshots, now); let accounts = initial .into_iter() .chain(periodic) @@ -857,28 +899,19 @@ impl CodexCredentialQuotaService { } /// `quota_observed_at` 为空代表首次异步观察尚未成功；不另建同步状态表。 - async fn initial_quota_sync_accounts( - &self, + fn initial_quota_sync_accounts( accounts: &[ProviderAccount], + observed_ids: &BTreeSet, now: SystemTime, - ) -> Result, CodexCredentialQuotaError> { - let account_ids = accounts - .iter() - .map(|account| account.id().clone()) - .collect::>(); - let observed = self.store.get_quotas(&account_ids).await?; - let observed_ids = observed - .into_iter() - .map(|observation| observation.account_id) - .collect::>(); - Ok(accounts + ) -> Vec { + accounts .iter() .filter(|account| { !observed_ids.contains(account.id()) && eligible_initial_quota_sync(account, now) }) .take(INITIAL_QUOTA_SYNC_BATCH) .cloned() - .collect()) + .collect() } /// 解析并 revision-fenced 落库单账号的 Provider quota JSON。 diff --git a/backend/crates/providers/openai/tests/admin.rs b/backend/crates/providers/openai/tests/admin.rs index b1fda9bb8..1b2b7468b 100644 --- a/backend/crates/providers/openai/tests/admin.rs +++ b/backend/crates/providers/openai/tests/admin.rs @@ -1142,6 +1142,121 @@ async fn openai_admin_keeps_confirmed_exhaustion_separate_from_raw_usage_display assert_eq!(recovered.windows[0].used_percent, Some(86.0)); } +#[tokio::test] +async fn openai_admin_preserves_expired_window_usage_and_exhaustion_attribution() { + for exhausted in [false, true] { + let store = Arc::new(MemoryAccountStore::default()); + let account_id = "acct_admin_expired_window"; + store + .seed_oauth_credential(ImportCodexOAuthCredential { + account_id: account_id.to_owned(), + name: "admin expired window".to_owned(), + secret: secret("admin-expired-window-access"), + verified_account: profile("chatgpt-admin-expired-window"), + next_refresh_at: Some(Utc::now() + chrono::Duration::minutes(30)), + enabled: true, + }) + .await; + let account = store.account(account_id).expect("stored account"); + let past_reset_at = Utc::now().timestamp() - 60; + let observed_at = SystemTime::now() - Duration::from_secs(300); + let weekly_used = if exhausted { 100.0 } else { 74.0 }; + let state = if exhausted { + QuotaState::exhausted( + QuotaEvidence::ProviderDenied, + observed_at, + Some(SystemTime::UNIX_EPOCH + Duration::from_secs(past_reset_at as u64)), + ) + } else { + QuotaState::allowed(observed_at) + }; + store + .compare_and_swap_quota(QuotaObservation { + plan_type: None, + account_id: account.id().clone(), + expected_revision: account.revision(), + quota: OpaqueProviderData::new( + json!({ + "rate_limit": { + "allowed": !exhausted, + "limit_reached": exhausted, + "primary_window": { + "used_percent": 15, + "reset_at": past_reset_at + 18_000, + "limit_window_seconds": 18_000, + }, + "secondary_window": { + "used_percent": weekly_used, + "reset_at": past_reset_at, + "limit_window_seconds": 604_800, + } + } + }) + .as_object() + .expect("quota object") + .clone(), + ), + observed_at, + state, + }) + .await + .expect("persist raw quota"); + + let config = valid_config(); + let bundle = provider_openai::initialize( + config.config.clone(), + provider_ports_with(Arc::clone(&store), Arc::new(TestOAuthPending::default())), + ) + .await + .expect("OpenAI bundle"); + let mut projected = bundle + .admin_provider() + .quota(ProviderQuotaRequest { + account_id: account.id().clone(), + refresh: false, + rolling_usage: None, + }) + .await + .expect("project quota"); + assert_eq!(projected.limit_reached, exhausted); + // 账号接口还会归一化耗尽展示；过期周窗口不能把触顶错误转移到短期窗口。 + projected.apply_limit_reached_display(); + let primary = projected + .windows + .iter() + .find(|w| w.window_seconds == Some(18_000)) + .expect("primary"); + let weekly = projected + .windows + .iter() + .find(|w| w.window_seconds == Some(604_800)) + .expect("weekly"); + assert_eq!( + (primary.used_percent, primary.limit_reached), + (Some(15.0), false) + ); + assert_eq!( + (weekly.used_percent, weekly.limit_reached), + (Some(weekly_used), exhausted) + ); + assert_eq!( + weekly.reset_at.map(|reset| reset.timestamp()), + Some(past_reset_at) + ); + let raw = store + .get_quotas(std::slice::from_ref(account.id())) + .await + .expect("raw quota") + .pop() + .expect("observation"); + assert_eq!(raw.observed_at, observed_at); + assert_eq!( + raw.quota.expose_to_provider()["rate_limit"]["secondary_window"]["used_percent"], + weekly_used + ); + } +} + #[tokio::test] async fn openai_admin_provider_rejects_unprepared_mutations_before_store_commit() { let store = Arc::new(MemoryAccountStore::default()); diff --git a/backend/crates/providers/openai/tests/credential/quota/refresh_timing.rs b/backend/crates/providers/openai/tests/credential/quota/refresh_timing.rs index 2a2129972..40f5f30a4 100644 --- a/backend/crates/providers/openai/tests/credential/quota/refresh_timing.rs +++ b/backend/crates/providers/openai/tests/credential/quota/refresh_timing.rs @@ -128,3 +128,97 @@ async fn periodic_checks_continue_when_reset_is_unknown_or_far_in_the_future() { assert_eq!(server.received_requests().await.expect("requests").len(), 1); } } + +#[tokio::test] +async fn allowed_account_with_expired_window_synchronizes_at_reset_grace() { + let store = Arc::new(MemoryAccountStore::default()); + create_account(&store, "acct_allowed_expired_window").await; + let account = store + .account("acct_allowed_expired_window") + .expect("account"); + let server = MockServer::start().await; + let service = quota_service_with_base_url(&store, reqwest::Client::new(), server.uri()); + + let short_reset = Utc::now().timestamp() - 114; + let usage = |used: u32, reset: i64| { + json!({"rate_limit": { + "allowed": true, + "primary_window": { + "used_percent": used, + "reset_at": reset, + "limit_window_seconds": 18_000, + } + }}) + }; + + mount_usage(&server, usage(100, short_reset)).await; + service + .refresh_account(account.id()) + .await + .expect("seed quota with 100% usage"); + + assert!( + Utc::now().timestamp() < short_reset + 120, + "fixture must precede grace deadline" + ); + let requests = server.received_requests().await.expect("requests").len(); + service.synchronize().await.expect("initial periodic check"); + service + .synchronize() + .await + .expect("repeat check before grace"); + assert_eq!( + server.received_requests().await.expect("requests").len(), + requests, + "正常账号首次复核也必须等待 reset 宽限期" + ); + + // 当到达 reset + 120s 宽限期后，账号虽然处于 allowed 状态，但包含已到期的非零用量窗口，被调度主动同步。 + mount_usage(&server, usage(0, short_reset + 18_000)).await; + wait_for_reset_grace(short_reset).await; + + let summary = service + .synchronize() + .await + .expect("synchronize expired window"); + assert_eq!(summary.updated, 1); + assert_eq!(server.received_requests().await.expect("requests").len(), 1); + + let snapshot = service + .read_account(account.id()) + .await + .expect("read quota") + .expect("snapshot"); + assert_eq!(snapshot.windows()[0].used_percent(), Some(0.0)); + service.synchronize().await.expect("skip fresh window"); + assert_eq!(server.received_requests().await.expect("requests").len(), 1); +} + +#[tokio::test] +async fn allowed_expired_window_refresh_is_throttled_when_observation_is_unchanged() { + let store = Arc::new(MemoryAccountStore::default()); + create_account(&store, "acct_allowed_expired_retry").await; + let account = store + .account("acct_allowed_expired_retry") + .expect("account"); + let server = MockServer::start().await; + let service = quota_service_with_base_url(&store, reqwest::Client::new(), server.uri()); + let usage = json!({"rate_limit": { + "allowed": true, + "primary_window": { + "used_percent": 74, + "reset_at": Utc::now().timestamp() - 180, + "limit_window_seconds": 18_000, + } + }}); + mount_usage(&server, usage.clone()).await; + service + .refresh_account(account.id()) + .await + .expect("seed quota"); + mount_usage(&server, usage).await; + + assert_eq!(service.synchronize().await.expect("first check").updated, 1); + service.synchronize().await.expect("throttled repeat check"); + assert_eq!(server.received_requests().await.expect("requests").len(), 1); +} diff --git a/docs/api.md b/docs/api.md index 002153258..3ea2f6c22 100644 --- a/docs/api.md +++ b/docs/api.md @@ -694,8 +694,10 @@ OAuth start 使用： - `GET /accounts/quota` 只读取最后一次落库快照；`POST /accounts/quota/refresh` 才访问上游。access token 已过期时，额度刷新要求先走 credential 刷新或重新授权，不会拿过期 token 探测额度。 - OpenAI 已耗尽账号每 30 分钟主动复核一次，也会在最早未恢复窗口的 `resetAt + 2 分钟` 到期后 - 提前复核。后台每 30 秒检查触发条件；同一重置边界复核后仍未恢复时回到 30 分钟重试， - 避免旧 reset 持续触发请求。各窗口独立确认恢复，时间到期本身不会直接解除账号耗尽。 + 提前复核；周期复核不等待旧重置时间，因此也能发现官方提前重置。正常账号有非零用量或触顶窗口时， + 在该窗口的 `resetAt + 2 分钟` 后主动复核。后台每 30 秒检查触发条件；同一重置边界复核后仍未更新时 + 回到 30 分钟重试，避免旧 reset 持续触发请求。各窗口独立确认恢复，时间到期本身不会直接解除账号 + 耗尽或将展示用量归零。 - `POST /accounts/recover` 是管理员对本地事实的强制恢复：它清除 Redis cooldown 和已保存的额度/错误， 把账号重新启用并恢复为可调度 credential；它不验证上游账号是否已经恢复，下一次真实请求仍可重新写入 失败事实。