diff --git a/backend/src/upstream/accounts/pool.rs b/backend/src/upstream/accounts/pool.rs index 50dc49b52..4b410d63f 100644 --- a/backend/src/upstream/accounts/pool.rs +++ b/backend/src/upstream/accounts/pool.rs @@ -331,6 +331,7 @@ impl AccountPool { let Some(account) = self.accounts.get_mut(account_id) else { return false; }; + let status = status_after_quota_limit(status, account.quota_limit_reached); account.status = status; if status != AccountStatus::Active { self.slots.remove(account_id); @@ -353,6 +354,7 @@ impl AccountPool { .unwrap_or(cooldown_until); account.quota_limit_reached = true; account.quota_cooldown_until = Some(final_cooldown_until); + account.status = status_after_quota_limit(account.status, true); account.window_reset_at = account .window_reset_at .filter(|existing| *existing > final_cooldown_until) @@ -1376,6 +1378,13 @@ fn refresh_quota_window(account: &mut Account, now: DateTime) { } } +fn status_after_quota_limit(status: AccountStatus, quota_limit_reached: bool) -> AccountStatus { + match (status, quota_limit_reached) { + (AccountStatus::Active, true) => AccountStatus::QuotaExhausted, + (status, _) => status, + } +} + fn reset_window_counters(account: &mut Account) { account.window_request_count = 0; account.window_input_tokens = 0; diff --git a/backend/src/upstream/accounts/store/mod.rs b/backend/src/upstream/accounts/store/mod.rs index 064078bfa..7952db6ca 100644 --- a/backend/src/upstream/accounts/store/mod.rs +++ b/backend/src/upstream/accounts/store/mod.rs @@ -540,6 +540,7 @@ impl SqliteAccountStore { .bind(optional_update_value(&update.plan_type)) .bind(update.status.is_some()) .bind(status) + .bind(status) .bind(&now) .bind(account_id) .execute(&self.pool) @@ -703,6 +704,9 @@ impl SqliteAccountStore { let cloudflare_cooldown_until = account.cloudflare_cooldown_until.map(|dt| dt.to_rfc3339()); let mut tx = self.pool.begin().await?; let result = sqlx::query(SYNC_RUNTIME_ACCOUNT_STATE_SQL) + .bind(quota_limit_reached) + .bind(&now) + .bind(quota_limit_reached) .bind(status_to_db(account.status)) .bind(quota_limit_reached) .bind(&now) @@ -773,12 +777,24 @@ impl SqliteAccountStore { status: AccountStatus, ) -> SqliteAccountStoreResult { let now = Utc::now().to_rfc3339(); - let result = sqlx::query("update accounts set status = ?, updated_at = ? where id = ?") - .bind(status_to_db(status)) - .bind(&now) - .bind(account_id) - .execute(&self.pool) - .await?; + let status = status_to_db(status); + let result = sqlx::query( + r" +update accounts +set + status = case + when ? = 'active' and quota_limit_reached = 1 then 'quota_exhausted' + else ? + end, + updated_at = ? +where id = ?", + ) + .bind(status) + .bind(status) + .bind(&now) + .bind(account_id) + .execute(&self.pool) + .await?; Ok(result.rows_affected() > 0) } diff --git a/backend/src/upstream/accounts/store/queries.rs b/backend/src/upstream/accounts/store/queries.rs index d0016089c..6569fb94e 100644 --- a/backend/src/upstream/accounts/store/queries.rs +++ b/backend/src/upstream/accounts/store/queries.rs @@ -165,7 +165,13 @@ set chatgpt_user_id = case when ? then ? else chatgpt_user_id end, label = case when ? then ? else label end, plan_type = case when ? then ? else plan_type end, - status = case when ? then ? else status end, + status = case + when ? then case + when ? = 'active' and quota_limit_reached = 1 then 'quota_exhausted' + else ? + end + else status + end, updated_at = ? where id = ?"; @@ -413,6 +419,10 @@ on conflict(account_id) do update set pub(super) const MARK_QUOTA_LIMITED_UNTIL_SQL: &str = r" update accounts set + status = case + when status in ('disabled', 'banned', 'expired') then status + else 'quota_exhausted' + end, quota_limit_reached = 1, quota_verify_required = 0, quota_cooldown_until = ?, @@ -424,6 +434,12 @@ update accounts set status = case when status in ('disabled', 'banned') then status + when ( + case + when ? = 0 and quota_cooldown_until is not null and quota_cooldown_until > ? then quota_limit_reached + else ? + end + ) = 1 then 'quota_exhausted' else ? end, quota_limit_reached = case @@ -496,6 +512,7 @@ set next_refresh_at = ?, status = case when status in ('disabled', 'banned') then status + when quota_limit_reached = 1 then 'quota_exhausted' else ? end, updated_at = ? @@ -513,6 +530,7 @@ set next_refresh_at = ?, status = case when status in ('disabled', 'banned') then status + when quota_limit_reached = 1 then 'quota_exhausted' else ? end, updated_at = ? diff --git a/backend/tests/upstream/accounts/account_pool/quota.rs b/backend/tests/upstream/accounts/account_pool/quota.rs index cb56d3694..1003f39ac 100644 --- a/backend/tests/upstream/accounts/account_pool/quota.rs +++ b/backend/tests/upstream/accounts/account_pool/quota.rs @@ -39,6 +39,22 @@ fn account_pool_should_mark_quota_state_as_quota_exhausted_status() { assert!(!restored.quota_limit_reached); } +#[test] +fn account_pool_should_mark_quota_cooldown_as_quota_exhausted_status() { + let now = fixed_time(); + let mut pool = AccountPool::default(); + pool.insert(crate::support::accounts::test_account( + "limited", + AccountStatus::Active, + )); + + pool.mark_quota_limited_until("limited", now + Duration::seconds(60)); + + let limited = pool.get("limited").unwrap(); + assert_eq!(limited.status, AccountStatus::QuotaExhausted); + assert!(limited.quota_limit_reached); +} + #[test] fn account_pool_should_reuse_quota_limited_accounts_after_cooldown() { let now = fixed_time(); diff --git a/backend/tests/upstream/accounts/account_repository.rs b/backend/tests/upstream/accounts/account_repository.rs index aad362e0a..b3e209a39 100644 --- a/backend/tests/upstream/accounts/account_repository.rs +++ b/backend/tests/upstream/accounts/account_repository.rs @@ -562,6 +562,34 @@ async fn account_repository_should_mark_quota_exhausted_for_any_exhausted_window ); } +#[tokio::test] +async fn account_repository_should_mark_quota_cooldown_as_quota_exhausted_status() { + let (pool, _dir) = sqlite_account_store_parts("accounts.sqlite", 48).await; + let repo = SqliteAccountStore::new(pool.clone()); + seed_repo_account(&pool, "acct_a", "2026-06-11T00:00:00Z").await; + let future_cooldown = Utc::now() + Duration::hours(1); + + repo.mark_quota_limited_until("acct_a", future_cooldown) + .await + .unwrap(); + let stored: (String, i64, Option) = sqlx::query_as( + "select status, quota_limit_reached, quota_cooldown_until from accounts where id = ?", + ) + .bind("acct_a") + .fetch_one(&pool) + .await + .unwrap(); + + assert_eq!( + (stored.0, stored.1, stored.2), + ( + String::from("quota_exhausted"), + 1, + Some(future_cooldown.to_rfc3339()) + ) + ); +} + #[tokio::test] async fn account_repository_runtime_state_sync_should_not_clear_newer_future_quota_cooldown() { let (pool, _dir) = sqlite_account_store_parts("accounts.sqlite", 42).await; @@ -569,7 +597,7 @@ async fn account_repository_runtime_state_sync_should_not_clear_newer_future_quo seed_repo_account(&pool, "acct_a", "2026-06-11T00:00:00Z").await; let future_cooldown = Utc::now() + Duration::hours(1); sqlx::query( - "update accounts set quota_limit_reached = 1, quota_verify_required = 0, quota_cooldown_until = ? where id = ?", + "update accounts set status = 'quota_exhausted', quota_limit_reached = 1, quota_verify_required = 0, quota_cooldown_until = ? where id = ?", ) .bind(future_cooldown.to_rfc3339()) .bind("acct_a") @@ -582,6 +610,7 @@ async fn account_repository_runtime_state_sync_should_not_clear_newer_future_quo .into_iter() .find(|account| account.id == "acct_a") .unwrap(); + stale_account.status = AccountStatus::Active; stale_account.quota_limit_reached = false; stale_account.quota_verify_required = true; stale_account.quota_cooldown_until = None; @@ -589,17 +618,18 @@ async fn account_repository_runtime_state_sync_should_not_clear_newer_future_quo repo.sync_runtime_account_state(&stale_account, false) .await .unwrap(); - let stored: (i64, i64, Option) = sqlx::query_as( - "select quota_limit_reached, quota_verify_required, quota_cooldown_until from accounts where id = ?", + let stored: (String, i64, i64, Option) = sqlx::query_as( + "select status, quota_limit_reached, quota_verify_required, quota_cooldown_until from accounts where id = ?", ) .bind("acct_a") .fetch_one(&pool) .await .unwrap(); - assert_eq!(stored.0, 1); - assert_eq!(stored.1, 0); - assert_eq!(stored.2, Some(future_cooldown.to_rfc3339())); + assert_eq!(stored.0, "quota_exhausted"); + assert_eq!(stored.1, 1); + assert_eq!(stored.2, 0); + assert_eq!(stored.3, Some(future_cooldown.to_rfc3339())); } #[tokio::test] diff --git a/backend/tests/upstream/accounts/token_refresh/success.rs b/backend/tests/upstream/accounts/token_refresh/success.rs index d88b7fd19..7be369f22 100644 --- a/backend/tests/upstream/accounts/token_refresh/success.rs +++ b/backend/tests/upstream/accounts/token_refresh/success.rs @@ -147,6 +147,88 @@ async fn token_refresh_task_should_refresh_quota_exhausted_account_without_clear assert!(stored.next_refresh_at.is_some()); } +#[tokio::test] +async fn token_refresh_task_should_not_clear_quota_limit_when_expired_token_refreshes() { + let dir = tempfile::tempdir().expect("temp dir"); + let db = dir.path().join("token-refresh-expired-quota.sqlite"); + let pool = connect_sqlite(&format!("sqlite://{}", db.display())) + .await + .expect("sqlite pool"); + let store = SqliteAccountStore::new(pool.clone()); + let now = Utc.with_ymd_and_hms(2026, 6, 19, 11, 30, 0).unwrap(); + let cooldown_until = now + Duration::hours(1); + let new_expires_at = now + Duration::hours(2); + let new_access_token = test_jwt(new_expires_at.timestamp()); + store + .insert(NewAccount { + id: "acct-expired-quota".to_string(), + email: Some("expired-quota@example.com".to_string()), + account_id: Some("chatgpt-expired-quota".to_string()), + user_id: Some("user-expired-quota".to_string()), + label: None, + plan_type: Some("plus".to_string()), + access_token: SecretString::new( + test_jwt((now - Duration::seconds(30)).timestamp()).into(), + ), + refresh_token: Some(SecretString::new( + "refresh-expired-quota".to_string().into(), + )), + added_at: None, + access_token_expires_at: Some(now - Duration::seconds(30)), + status: AccountStatus::Expired, + }) + .await + .expect("account should be inserted"); + sqlx::query( + "update accounts set quota_limit_reached = 1, quota_cooldown_until = ? where id = ?", + ) + .bind(cooldown_until.to_rfc3339()) + .bind("acct-expired-quota") + .execute(&pool) + .await + .unwrap(); + let refresher = StaticTokenRefresher { + response: Arc::new(Mutex::new(Ok(TokenPair { + access_token: new_access_token.clone(), + refresh_token: None, + }))), + }; + let task = codex_proxy_rs::upstream::accounts::token_refresh::RuntimeTokenRefreshService::new( + store.clone(), + RefreshPolicy { + refresh_margin_seconds: 300, + refresh_concurrency: 1, + }, + refresher, + ); + + let summary = task + .refresh_due_accounts_once_at(now) + .await + .expect("expired quota-limited account should refresh token"); + let stored = store + .get("acct-expired-quota") + .await + .expect("account should load") + .expect("account should exist"); + let quota_state: (i64, Option) = sqlx::query_as( + "select quota_limit_reached, quota_cooldown_until from accounts where id = ?", + ) + .bind("acct-expired-quota") + .fetch_one(&pool) + .await + .unwrap(); + + assert_eq!(summary.refreshed, 1); + assert_eq!(stored.status, AccountStatus::QuotaExhausted); + assert_eq!( + stored.access_token.expose_secret(), + new_access_token.as_str() + ); + assert_eq!(quota_state.0, 1); + assert_eq!(quota_state.1, Some(cooldown_until.to_rfc3339())); +} + #[tokio::test] async fn token_refresh_task_should_persist_success_after_transient_transport_retry() { let dir = tempfile::tempdir().expect("temp dir");