diff --git a/backend/crates/providers/openai/src/config.rs b/backend/crates/providers/openai/src/config.rs index e2732d06d..6fe7c8f85 100644 --- a/backend/crates/providers/openai/src/config.rs +++ b/backend/crates/providers/openai/src/config.rs @@ -217,7 +217,7 @@ pub struct CodexQuotaSettings { impl Default for CodexQuotaSettings { fn default() -> Self { Self { - refresh_interval_minutes: 5, + refresh_interval_minutes: 15, skip_exhausted: true, } } diff --git a/backend/crates/providers/openai/src/credential/quota.rs b/backend/crates/providers/openai/src/credential/quota.rs index 0e056a5a6..a02e8c586 100644 --- a/backend/crates/providers/openai/src/credential/quota.rs +++ b/backend/crates/providers/openai/src/credential/quota.rs @@ -30,15 +30,12 @@ const MAX_RATE_LIMIT_COOLDOWN: Duration = Duration::from_secs(60 * 60); const QUOTA_SCHEDULING_TTL: Duration = Duration::from_secs(10 * 60); const QUOTA_HYDRATION_FAILURE_TTL: Duration = Duration::from_secs(5); const QUOTA_PERIODIC_REFRESH_MIN_INTERVAL: Duration = Duration::from_secs(30 * 60); -const QUOTA_RESET_RECHECK_GRACE: Duration = Duration::from_secs(30); -// ready 账号主要依赖真实请求的被动观察；这里只提供低频、定额的状态兜底。 -const QUOTA_READY_REFRESH_MAX_AGE: Duration = Duration::from_secs(2 * 60 * 60); -const QUOTA_READY_REFRESH_BATCH_SIZE: usize = 8; /// OpenAI Provider 主动额度刷新的调度策略。 /// -/// 账号筛选仍由 core 的可调度状态统一处理；该策略只控制 Provider 自己的 -/// quota worker 频率，避免把 OpenAI 额度语义泄漏到公共调度层。 +/// 正常账号依赖请求响应的被动额度同步；周期 worker 仅复核已耗尽账号。 +/// 该策略只控制 Provider 自己的 quota worker 频率，避免把 OpenAI +/// 额度语义泄漏到公共调度层。 #[derive(Debug, Clone, Copy)] pub struct CodexQuotaRefreshPolicy { interval: Duration, @@ -249,8 +246,6 @@ struct CodexQuotaProjectionState { next_version: u64, entries: BTreeMap, last_periodic_refresh_at: BTreeMap, - last_exhausted_reset_recheck: BTreeMap, - ready_refresh_cursor: Option, } #[derive(Debug, Clone, Copy)] @@ -261,24 +256,6 @@ struct CodexQuotaSchedulingEntry { signals: Option, } -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -struct CodexQuotaResetMarker { - revision: CredentialRevision, - reset_at: SystemTime, -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -enum CodexQuotaRefreshKind { - Exhausted { reset_at: Option }, - Ready, -} - -#[derive(Clone)] -struct CodexQuotaRefreshCandidate { - account: ProviderAccount, - kind: CodexQuotaRefreshKind, -} - #[derive(Clone)] struct CodexQuotaHydrationTarget { account: ProviderAccount, @@ -461,16 +438,15 @@ impl CodexQuotaSchedulingProjection { fn reserve_periodic_refreshes( &self, accounts: Vec, - observations: Option<&BTreeMap>, now: SystemTime, ) -> Vec { let candidates = accounts .into_iter() - .filter_map(|account| quota_refresh_candidate(account, observations, now)) + .filter_map(|account| quota_refresh_candidate(account, now)) .collect::>(); let candidate_ids = candidates .iter() - .map(|candidate| candidate.account.id().clone()) + .map(|account| account.id().clone()) .collect::>(); let refreshed_at = Instant::now(); let mut state = self @@ -480,95 +456,27 @@ impl CodexQuotaSchedulingProjection { state .last_periodic_refresh_at .retain(|account_id, _| candidate_ids.contains(account_id)); - state - .last_exhausted_reset_recheck - .retain(|account_id, _| candidate_ids.contains(account_id)); + // 正常账号只由真实请求的响应头和 `codex.rate_limits` 被动同步。定时器只兜底 + // 已耗尽账号，且不依据 reset 时间跳过复核；每个账号仍至少间隔 30 分钟。 let mut reserved = Vec::new(); - let mut ready = Vec::new(); - for candidate in candidates { - match candidate.kind { - CodexQuotaRefreshKind::Exhausted { reset_at } => { - if reserve_exhausted_quota_refresh( - &mut state, - &candidate.account, - reset_at, - now, - refreshed_at, - ) { - reserved.push(candidate.account); - } - } - CodexQuotaRefreshKind::Ready => ready.push(candidate.account), + for account in candidates { + if !periodic_quota_refresh_due(&state, account.id(), refreshed_at) { + continue; } - } - - ready.sort_by(|left, right| left.id().cmp(right.id())); - let ready = ready - .into_iter() - .filter(|account| periodic_quota_refresh_due(&state, account.id(), refreshed_at)) - .collect::>(); - let ready_len = ready.len(); - if ready_len == 0 { - return reserved; - } - let start = state - .ready_refresh_cursor - .as_ref() - .and_then(|cursor| ready.iter().position(|account| account.id() > cursor)) - .unwrap_or(0); - let ready = ready - .into_iter() - .cycle() - .skip(start) - .take(ready_len.min(QUOTA_READY_REFRESH_BATCH_SIZE)) - .collect::>(); - if let Some(last) = ready.last() { - state.ready_refresh_cursor = Some(last.id().clone()); - } - for account in &ready { state .last_periodic_refresh_at .insert(account.id().clone(), refreshed_at); + reserved.push(account); } - reserved.extend(ready); reserved } } -fn quota_refresh_candidate( - account: ProviderAccount, - observations: Option<&BTreeMap>, - now: SystemTime, -) -> Option { - if !eligible_periodic_quota_refresh(&account, now) { - return None; - } - let observation = observations.and_then(|observations| observations.get(account.id())); - let snapshot = quota_snapshot_for_account(&account, observation); - let kind = match account.availability() { - AccountAvailability::QuotaExhausted => CodexQuotaRefreshKind::Exhausted { - reset_at: account - .cooldown_until() - .or_else(|| quota_reset_at_from_snapshot(snapshot.as_ref())), - }, - AccountAvailability::Cooldown - if snapshot - .as_ref() - .is_some_and(|snapshot| snapshot.fact().exhausted()) => - { - CodexQuotaRefreshKind::Exhausted { - reset_at: quota_reset_at_from_snapshot(snapshot.as_ref()), - } - } - AccountAvailability::Ready - if observations.is_some() && quota_observation_is_stale(&account, observation, now) => - { - CodexQuotaRefreshKind::Ready - } - _ => return None, - }; - Some(CodexQuotaRefreshCandidate { account, kind }) +fn quota_refresh_candidate(account: ProviderAccount, now: SystemTime) -> Option { + (eligible_periodic_quota_refresh(&account, now) + && account.availability() == AccountAvailability::QuotaExhausted) + .then_some(account) } fn quota_snapshot_for_account( @@ -597,72 +505,6 @@ fn quota_observation_confirms_active_exhaustion( .is_some_and(|snapshot| quota_fact_is_active_exhaustion(snapshot.fact(), now)) } -fn quota_observation_is_stale( - account: &ProviderAccount, - observation: Option<&QuotaObservation>, - now: SystemTime, -) -> bool { - let Some(observation) = observation.filter(|observation| { - observation.expected_revision == account.revision() - && quota_snapshot_from_observation(observation).is_some() - }) else { - return true; - }; - observation - .observed_at - .and_then(|observed_at| now.duration_since(observed_at).ok()) - .is_none_or(|age| age >= QUOTA_READY_REFRESH_MAX_AGE) -} - -fn reserve_exhausted_quota_refresh( - state: &mut CodexQuotaProjectionState, - account: &ProviderAccount, - reset_at: Option, - now: SystemTime, - refreshed_at: Instant, -) -> bool { - if let Some(reset_at) = reset_at { - if !quota_reset_recheck_ready(reset_at, now) { - return false; - } - let same_reset_already_checked = state - .last_exhausted_reset_recheck - .get(account.id()) - .is_some_and(|marker| { - marker.revision == account.revision() && marker.reset_at == reset_at - }); - if same_reset_already_checked - && !periodic_quota_refresh_due(state, account.id(), refreshed_at) - { - return false; - } - state - .last_periodic_refresh_at - .insert(account.id().clone(), refreshed_at); - state.last_exhausted_reset_recheck.insert( - account.id().clone(), - CodexQuotaResetMarker { - revision: account.revision(), - reset_at, - }, - ); - return true; - } - if !periodic_quota_refresh_due(state, account.id(), refreshed_at) { - return false; - } - state - .last_periodic_refresh_at - .insert(account.id().clone(), refreshed_at); - true -} - -fn quota_reset_recheck_ready(reset_at: SystemTime, now: SystemTime) -> bool { - reset_at - .checked_add(QUOTA_RESET_RECHECK_GRACE) - .is_some_and(|ready_at| ready_at <= now) -} - fn periodic_quota_refresh_due( state: &CodexQuotaProjectionState, account_id: &ProviderAccountId, @@ -774,16 +616,14 @@ impl CodexCredentialQuotaService { .collect::>(), ), Err(error) => { - // 持久观察不可用时不盲扫 ready；保留已耗尽账号的既有恢复轮询。 - tracing::warn!(error = %error, "OpenAI quota refresh skipped ready-account scan"); + // 持久观察不可用时跳过旧状态修正；仍复核当前已标记耗尽的账号。 + tracing::warn!(error = %error, "OpenAI quota refresh skipped persisted exhaustion reconciliation"); None } }; self.reconcile_persisted_exhaustion(&mut accounts, observations.as_ref(), now) .await; - let accounts = - self.scheduling - .reserve_periodic_refreshes(accounts, observations.as_ref(), now); + let accounts = self.scheduling.reserve_periodic_refreshes(accounts, now); if accounts.is_empty() { return Ok(summary); } diff --git a/backend/crates/providers/openai/src/provider.rs b/backend/crates/providers/openai/src/provider.rs index 11339b753..a5f4fd2c5 100644 --- a/backend/crates/providers/openai/src/provider.rs +++ b/backend/crates/providers/openai/src/provider.rs @@ -84,6 +84,9 @@ const MAX_COOKIE_HEADER_BYTES: usize = 16 * 1024; /// 提交边界前最多保留 64 KiB 原始上游 chunk；达到阈值后结束无感换号窗口， /// 但不会把上游数据改写成协议失败。 const MAX_STREAM_PREFETCH_BYTES: usize = 64 * 1024; +// 额度拒绝后先给上游额度结算留出时间，再以受限时长同步 usage 快照。 +const QUOTA_FAILURE_REFRESH_DELAY: Duration = Duration::from_secs(2); +const QUOTA_FAILURE_REFRESH_TIMEOUT: Duration = Duration::from_secs(5); pub const OFFICIAL_CODEX_BASE_PATH: &str = "/backend-api"; pub const OFFICIAL_CODEX_BASE_URL: &str = "https://chatgpt.com/backend-api"; @@ -1937,7 +1940,7 @@ impl ReplayBoundary { struct OpenAiFailureContext<'a> { client: &'a CodexBackendClient, selector: &'a CodexCredentialSelector, - quota: &'a CodexCredentialQuotaService, + quota: &'a Arc, response_origin: &'a Url, cyber_policy_scope: Option<&'a CodexCyberPolicyScope>, allows_account_state_mutation: bool, @@ -1952,6 +1955,10 @@ async fn apply_failure( return; } synchronize_passive_quota(context.quota, account, &failure.rate_limit_headers).await; + let needs_authoritative_quota_refresh = matches!( + failure.account_failure, + Some(CodexAccountFailure::QuotaExhausted | CodexAccountFailure::UsageLimitExhausted { .. }) + ); if failure.cyber_policy_failure { context .selector @@ -1975,6 +1982,9 @@ async fn apply_failure( ); } } + if needs_authoritative_quota_refresh { + schedule_authoritative_quota_refresh_after_failure(context.quota, account); + } if failure.capture_response_cookies && !failure.set_cookie_headers.is_empty() && let Err(error) = context @@ -1994,6 +2004,42 @@ async fn apply_failure( } } +fn schedule_authoritative_quota_refresh_after_failure( + quota: &Arc, + account: &ProviderAccount, +) { + // 限额错误可以确认账号状态，但 Responses 事件中的 used_percent 可能仍停在上一结算点。 + // usage 快照在后台补齐，不能阻塞原始失败响应。先等待 2 秒让上游结算， + // 再查询；5 秒仅限制实际查询自身。 + let quota = Arc::clone(quota); + let account_id = account.id().clone(); + drop(tokio::spawn(async move { + tokio::time::sleep(QUOTA_FAILURE_REFRESH_DELAY).await; + match tokio::time::timeout( + QUOTA_FAILURE_REFRESH_TIMEOUT, + quota.refresh_account(&account_id), + ) + .await + { + Ok(Ok(_)) => {} + Ok(Err(error)) => { + tracing::warn!( + account_id = %account_id, + error = %error, + "OpenAI authoritative quota synchronization after quota rejection failed" + ); + } + Err(_) => { + tracing::warn!( + account_id = %account_id, + timeout_ms = QUOTA_FAILURE_REFRESH_TIMEOUT.as_millis(), + "OpenAI authoritative quota synchronization after quota rejection timed out" + ); + } + } + })); +} + fn map_handshake_error(error: CodexClientError) -> MappedProviderFailure { map_client_error(error, UpstreamSendState::Ambiguous, true) } diff --git a/backend/crates/providers/openai/tests/config.rs b/backend/crates/providers/openai/tests/config.rs index 7f712d721..c26e29fd7 100644 --- a/backend/crates/providers/openai/tests/config.rs +++ b/backend/crates/providers/openai/tests/config.rs @@ -117,7 +117,7 @@ fn openai_config_defaults_to_the_provider_owned_operating_values() { 64, 8, 20_000, - 5, + 15, true, true, "app_EMoamEEZ73f0CkXaXp7hrann", diff --git a/backend/crates/providers/openai/tests/credential/quota.rs b/backend/crates/providers/openai/tests/credential/quota.rs index 02ef92b76..40fe85eba 100644 --- a/backend/crates/providers/openai/tests/credential/quota.rs +++ b/backend/crates/providers/openai/tests/credential/quota.rs @@ -755,7 +755,7 @@ async fn synchronize_without_accounts_is_a_noop_before_network_io() { } #[tokio::test] -async fn periodic_quota_synchronization_skips_fresh_ready_accounts() { +async fn periodic_quota_synchronization_does_not_scan_stale_ready_accounts() { let store = Arc::new(MemoryAccountStore::default()); create_account(&store, "acct_periodic_ready").await; let account = store @@ -776,15 +776,16 @@ async fn periodic_quota_synchronization_skips_fresh_ready_accounts() { .expect("quota object") .clone(), )), - observed_at: Some(SystemTime::now()), + // 正常账号即使缓存已旧，也只等待真实请求被动同步，不能被定时任务主动拉取。 + observed_at: Some(SystemTime::UNIX_EPOCH), }) .await - .expect("persist fresh quota"); + .expect("persist stale quota"); let summary = blocked_network_quota_service(&store) .synchronize() .await - .expect("fresh ready account must not query upstream"); + .expect("ready account must not query upstream"); assert_eq!(summary, CodexQuotaSyncSummary::default()); } @@ -830,7 +831,13 @@ async fn periodic_quota_synchronization_repairs_cooldown_with_active_exhausted_q let server = MockServer::start().await; Mock::given(method("GET")) .and(path("/api/codex/usage")) - .respond_with(ResponseTemplate::new(429)) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "rate_limit": { + "allowed": false, + "limit_reached": true, + "primary_window": {"used_percent": 100, "reset_at": reset_at} + } + }))) .mount(&server) .await; let service = quota_service_with_base_url( @@ -939,72 +946,6 @@ async fn periodic_quota_synchronization_skips_ordinary_cooldown_without_exhauste ); } -#[tokio::test] -async fn periodic_quota_synchronization_limits_stale_ready_account_scans() { - let store = Arc::new(MemoryAccountStore::default()); - for index in 0..9 { - let account_id = format!("acct_periodic_stale_ready_{index}"); - create_account(&store, &account_id).await; - let account = store.account(&account_id).expect("created account"); - store - .compare_and_swap_quota(QuotaObservation { - account_id: account.id().clone(), - expected_revision: account.revision(), - quota: Some(OpaqueProviderData::new( - json!({ - "rate_limit": { - "allowed": true, - "primary_window": {"used_percent": 10, "reset_at": 1_900_000_000} - } - }) - .as_object() - .expect("quota object") - .clone(), - )), - observed_at: Some(SystemTime::UNIX_EPOCH), - }) - .await - .expect("persist stale quota"); - } - let server = MockServer::start().await; - Mock::given(method("GET")) - .and(path("/api/codex/usage")) - .respond_with(ResponseTemplate::new(200).set_body_json(json!({ - "rate_limit": { - "allowed": false, - "primary_window": {"used_percent": 100, "reset_at": 1_900_000_000} - } - }))) - .mount(&server) - .await; - let service = quota_service_with_base_url( - &store, - reqwest::Client::builder() - .no_proxy() - .build() - .expect("client"), - server.uri(), - ); - - let summary = service.synchronize().await.expect("quota synchronization"); - - assert_eq!(summary.exhausted, 8); - assert_eq!( - store - .account("acct_periodic_stale_ready_0") - .expect("scanned account") - .availability(), - AccountAvailability::QuotaExhausted - ); - assert_eq!( - store - .account("acct_periodic_stale_ready_8") - .expect("deferred account") - .availability(), - AccountAvailability::Ready - ); -} - #[tokio::test] async fn periodic_quota_synchronization_attempts_quota_exhausted_accounts() { let store = Arc::new(MemoryAccountStore::default()); @@ -1033,7 +974,7 @@ async fn periodic_quota_synchronization_attempts_quota_exhausted_accounts() { } #[tokio::test] -async fn periodic_quota_synchronization_waits_for_quota_cooldown_deadline_with_grace() { +async fn periodic_quota_synchronization_rechecks_exhausted_accounts_without_waiting_for_reset() { let store = Arc::new(MemoryAccountStore::default()); create_account(&store, "acct_periodic_exhausted_cooldown").await; let account = store @@ -1045,38 +986,29 @@ async fn periodic_quota_synchronization_waits_for_quota_cooldown_deadline_with_g expected_revision: account.revision(), availability: AccountAvailability::QuotaExhausted, reason: Some("usage_limit_exhausted".to_owned()), - cooldown_until: SystemTime::now().checked_sub(Duration::from_secs(29)), + cooldown_until: SystemTime::now().checked_add(Duration::from_secs(24 * 60 * 60)), observed_at: SystemTime::now(), }) .await - .expect("mark account exhausted with recent cooldown"); - let service = blocked_network_quota_service(&store); + .expect("mark account exhausted with a future reset"); + let server = MockServer::start().await; + Mock::given(method("GET")) + .and(path("/api/codex/usage")) + .respond_with(ResponseTemplate::new(500)) + .expect(1) + .mount(&server) + .await; + let service = quota_service_with_base_url( + &store, + reqwest::Client::builder() + .no_proxy() + .build() + .expect("client"), + server.uri(), + ); - let waiting = service - .synchronize() - .await - .expect("quota synchronization before grace"); - - assert_eq!(waiting, CodexQuotaSyncSummary::default()); - let current = store - .account("acct_periodic_exhausted_cooldown") - .expect("current account"); - store - .apply_state_change(AccountStateChange { - account_id: current.id().clone(), - expected_revision: current.revision(), - availability: AccountAvailability::QuotaExhausted, - reason: Some("usage_limit_exhausted".to_owned()), - cooldown_until: SystemTime::now().checked_sub(Duration::from_secs(31)), - observed_at: SystemTime::now(), - }) - .await - .expect("mark account exhausted past grace"); - - let due = service - .synchronize() - .await - .expect("quota synchronization after grace"); + // reset 仅用于展示/调度状态，不能推迟额度耗尽账号的定时复核。 + let due = service.synchronize().await.expect("quota synchronization"); let repeated = service .synchronize() .await @@ -1113,91 +1045,7 @@ async fn periodic_quota_synchronization_throttles_the_same_exhausted_account() { } #[tokio::test] -async fn periodic_quota_synchronization_rechecks_an_exhausted_account_once_when_its_reset_is_due() { - let store = Arc::new(MemoryAccountStore::default()); - create_account(&store, "acct_periodic_reset_due").await; - let account = store - .account("acct_periodic_reset_due") - .expect("created account"); - store - .apply_state_change(AccountStateChange { - account_id: account.id().clone(), - expected_revision: account.revision(), - availability: AccountAvailability::QuotaExhausted, - reason: Some("quota_exhausted".to_owned()), - cooldown_until: None, - observed_at: SystemTime::now(), - }) - .await - .expect("mark account exhausted"); - store - .compare_and_swap_quota(QuotaObservation { - account_id: account.id().clone(), - expected_revision: account.revision(), - quota: Some(OpaqueProviderData::new( - json!({ - "rate_limit": { - "allowed": false, - "primary_window": {"used_percent": 100, "reset_at": 1_900_000_000} - } - }) - .as_object() - .expect("quota object") - .clone(), - )), - observed_at: Some(SystemTime::now()), - }) - .await - .expect("persist future reset"); - let server = MockServer::start().await; - Mock::given(method("GET")) - .and(path("/api/codex/usage")) - .respond_with(ResponseTemplate::new(500)) - .mount(&server) - .await; - let service = quota_service_with_base_url( - &store, - reqwest::Client::builder() - .no_proxy() - .build() - .expect("client"), - server.uri(), - ); - - let waiting = service.synchronize().await.expect("initial cooldown wait"); - assert_eq!(waiting, CodexQuotaSyncSummary::default()); - store - .compare_and_swap_quota(QuotaObservation { - account_id: account.id().clone(), - expected_revision: account.revision(), - quota: Some(OpaqueProviderData::new( - json!({ - "rate_limit": { - "allowed": false, - "primary_window": {"used_percent": 100, "reset_at": 1} - } - }) - .as_object() - .expect("quota object") - .clone(), - )), - observed_at: Some(SystemTime::now()), - }) - .await - .expect("persist due reset"); - - let reset_recheck = service.synchronize().await.expect("reset recheck"); - let repeated = service - .synchronize() - .await - .expect("throttled reset recheck"); - - assert_eq!(reset_recheck.transient, 1); - assert_eq!(repeated, CodexQuotaSyncSummary::default()); -} - -#[tokio::test] -async fn periodic_quota_synchronization_reschedules_a_confirmed_new_reset() { +async fn periodic_quota_synchronization_does_not_bypass_throttle_for_a_new_reset() { let store = Arc::new(MemoryAccountStore::default()); let account_id = "acct_periodic_new_reset"; create_account(&store, account_id).await; @@ -1241,7 +1089,10 @@ async fn periodic_quota_synchronization_reschedules_a_confirmed_new_reset() { .account(account_id) .expect("account after first recheck") .cooldown_until(); - let second = service.synchronize().await.expect("new reset recheck"); + let second = service + .synchronize() + .await + .expect("throttled reset recheck"); let repeated = service.synchronize().await.expect("same reset throttled"); let request_count = server .received_requests() @@ -1260,9 +1111,9 @@ async fn periodic_quota_synchronization_reschedules_a_confirmed_new_reset() { ( 1, Some(confirmed_reset), - 1, + 0, CodexQuotaSyncSummary::default(), - 2, + 1, ) ); } diff --git a/backend/crates/providers/openai/tests/provider/contract.rs b/backend/crates/providers/openai/tests/provider/contract.rs index 4beaaa635..dc08539f2 100644 --- a/backend/crates/providers/openai/tests/provider/contract.rs +++ b/backend/crates/providers/openai/tests/provider/contract.rs @@ -7,7 +7,8 @@ use base64::{Engine as _, engine::general_purpose::STANDARD}; use chrono::Utc; use futures::StreamExt; use gateway_core::engine::credential::{ - AccountAvailability, AccountFeedbackStats, ProviderAccountId, ProviderAccountStore as _, + AccountAvailability, AccountFeedbackStats, OpaqueProviderData, ProviderAccountId, + ProviderAccountStore as _, QuotaObservation, }; use gateway_core::engine::provider::{Provider as _, ProviderRequest}; use gateway_core::engine::{ @@ -1006,16 +1007,49 @@ async fn response_failed_before_semantic_output_is_atomic_and_persists_quota_loc } #[tokio::test] -async fn usage_limit_failure_marks_quota_exhausted_without_usage_probe() { +async fn usage_limit_failure_returns_promptly_and_refreshes_the_authoritative_quota_snapshot() { let store = Arc::new(MemoryAccountStore::default()); let account_id = "acct_usage_limit_request_path"; create_account(&store, account_id).await; + let account = store.account(account_id).expect("created account"); + store + .compare_and_swap_quota(QuotaObservation { + account_id: account.id().clone(), + expected_revision: account.revision(), + quota: Some(OpaqueProviderData::new( + json!({ + "rate_limit": { + "allowed": true, + "limit_reached": false, + "primary_window": {"used_percent": 99, "reset_at": 1_900_000_000} + } + }) + .as_object() + .expect("stale quota object") + .clone(), + )), + observed_at: Some(SystemTime::now()), + }) + .await + .expect("seed stale passive quota"); let server = MockServer::start().await; Mock::given(method("GET")) .and(path("/api/codex/usage")) .and(header("authorization", format!("Bearer at-{account_id}"))) - .respond_with(ResponseTemplate::new(500)) - .expect(0) + // 限额错误本身只能确认状态；这里必须覆盖上一帧被动观察到的 99%。 + .respond_with( + ResponseTemplate::new(200) + // 验证后台 usage 同步不能把原始的额度错误响应拖到查询完成之后。 + .set_delay(Duration::from_millis(750)) + .set_body_json(json!({ + "rate_limit": { + "allowed": false, + "limit_reached": true, + "primary_window": {"used_percent": 100, "reset_at": 1_900_000_000} + } + })), + ) + .expect(1) .mount(&server) .await; Mock::given(method("POST")) @@ -1040,13 +1074,17 @@ async fn usage_limit_failure_marks_quota_exhausted_without_usage_probe() { ) .await .expect("prepare provider stream"); - let failure = loop { - match stream.next().await { - Some(Ok(_)) => {} - Some(Err(error)) => break error, - None => panic!("usage-limit failure must surface a typed failure"), + let failure = timeout(Duration::from_millis(300), async { + loop { + match stream.next().await { + Some(Ok(_)) => {} + Some(Err(error)) => break error, + None => panic!("usage-limit failure must surface a typed failure"), + } } - }; + }) + .await + .expect("usage refresh must not block the original quota failure"); assert_eq!(failure.kind(), ProviderErrorKind::QuotaExhausted); let account = store.account(account_id).expect("usage-limit account"); @@ -1056,14 +1094,53 @@ async fn usage_limit_failure_marks_quota_exhausted_without_usage_probe() { .cooldown_until() .is_some_and(|until| until > SystemTime::now()) ); - assert!(!store.has_quota(account_id)); + tokio::time::sleep(Duration::from_millis(200)).await; let requests = server.received_requests().await.expect("received requests"); assert_eq!( requests .iter() .filter(|request| request.url.path() == "/api/codex/usage") .count(), - 0 + 0, + "usage refresh must wait for the upstream quota settlement delay" + ); + + let quota = timeout(Duration::from_secs(5), async { + loop { + let quota = store + .get_quotas(&[account.id().clone()]) + .await + .expect("read refreshed quota") + .into_iter() + .next() + .and_then(|observation| observation.quota) + .expect("authoritative quota snapshot"); + let quota = Value::Object(quota.into_inner()); + if quota + .pointer("/rate_limit/primary_window/used_percent") + .and_then(Value::as_u64) + == Some(100) + { + break quota; + } + tokio::time::sleep(Duration::from_millis(10)).await; + } + }) + .await + .expect("background usage refresh must replace the stale quota snapshot"); + assert_eq!( + quota + .pointer("/rate_limit/primary_window/used_percent") + .and_then(Value::as_u64), + Some(100) + ); + let requests = server.received_requests().await.expect("received requests"); + assert_eq!( + requests + .iter() + .filter(|request| request.url.path() == "/api/codex/usage") + .count(), + 1 ); }