diff --git a/backend/src/bootstrap/services.rs b/backend/src/bootstrap/services.rs index 0bb0784a3..da751cdae 100644 --- a/backend/src/bootstrap/services.rs +++ b/backend/src/bootstrap/services.rs @@ -321,7 +321,8 @@ impl Services { refresh_policy.clone(), token_client.clone(), ) - .with_refresh_lease_store(stores.refresh_leases.clone()), + .with_refresh_lease_store(stores.refresh_leases.clone()) + .with_account_pool_sync(account_pool.clone()), ); let upstream_client: Arc = codex.clone(); let store: Arc = Arc::new(stores.model_snapshots.clone()); diff --git a/backend/src/fleet/refresh/service.rs b/backend/src/fleet/refresh/service.rs index 20a0ff4fd..3183ceaf6 100644 --- a/backend/src/fleet/refresh/service.rs +++ b/backend/src/fleet/refresh/service.rs @@ -22,6 +22,7 @@ use uuid::Uuid; use super::TokenRefresher; use crate::fleet::{ account::{Account, AccountStatus}, + pool::AccountPoolService, store::{ AccountClaimsUpdate, AccountStore, PgAccountStore, PgAccountStoreError, StoredAccount, }, @@ -48,6 +49,7 @@ where C: TokenRefresher, { store: PgAccountStore, + account_pool: Option>, scheduler: Arc>, policy: RuntimeRefreshPolicy, refresh_leases: Option, @@ -193,6 +195,7 @@ where fn clone(&self) -> Self { Self { store: self.store.clone(), + account_pool: self.account_pool.clone(), scheduler: self.scheduler.clone(), policy: self.policy.clone(), refresh_leases: self.refresh_leases.clone(), @@ -215,6 +218,7 @@ where let policy = policy.into(); Self { store, + account_pool: None, scheduler: Arc::new(RefreshScheduler::new(policy.clone(), client)), policy, refresh_leases: None, @@ -233,6 +237,12 @@ where self } + /// 刷新后同步持久化账号到运行时池,避免旧 token 快照覆盖新状态。 + pub fn with_account_pool_sync(mut self, account_pool: Arc) -> Self { + self.account_pool = Some(account_pool); + self + } + /// 更新刷新策略。 pub fn update_policy(&self, policy: RefreshPolicy) { self.policy.update(policy); @@ -587,6 +597,7 @@ where if attempt.refresh_token.is_none() { persist_status(&self.store, &attempt.id, AccountStatus::Expired).await?; persist_next_refresh_at(&self.store, &attempt.id, None).await?; + self.sync_account_pool_after_update(&attempt.id).await; return Ok(TokenRefreshOutcome::StatusUpdated); } @@ -601,6 +612,7 @@ where persist_status(&self.store, &attempt.id, updated.status).await?; persist_next_refresh_at(&self.store, &attempt.id, None).await?; + self.sync_account_pool_after_update(&attempt.id).await; return Ok(TokenRefreshOutcome::StatusUpdated); } Ok(mut updated) => { @@ -609,6 +621,7 @@ where updated.next_refresh_at = self.computed_refresh_at(&updated); let refresh_token_rotated = updated.refresh_token != attempt.refresh_token; persist_token_update(&self.store, &updated).await?; + self.sync_account_pool_after_update(&updated.id).await; info!( account_id = %updated.id, attempt = attempt_index + 1, @@ -643,6 +656,19 @@ where Ok(TokenRefreshOutcome::Skipped) } + async fn sync_account_pool_after_update(&self, account_id: &str) { + let Some(account_pool) = self.account_pool.as_ref() else { + return; + }; + if let Err(error) = account_pool.sync_account_from_store(account_id).await { + warn!( + account_id, + error = %error, + "Failed to sync runtime account pool after token refresh" + ); + } + } + async fn prepare_refresh_attempt( &self, account: &Account, diff --git a/backend/src/fleet/store/queries.rs b/backend/src/fleet/store/queries.rs index 9c5e34e26..b1a2dd135 100644 --- a/backend/src/fleet/store/queries.rs +++ b/backend/src/fleet/store/queries.rs @@ -174,6 +174,7 @@ update accounts set status = case when status in ('disabled', 'banned') then status + when $4 = 'expired' and access_token_expires_at > $2 then status when ( case when not $1 and quota_cooldown_until is not null and quota_cooldown_until > $2 then quota_limit_reached diff --git a/backend/tests/fleet/refresh/mod.rs b/backend/tests/fleet/refresh/mod.rs index adfb64779..806287e0d 100644 --- a/backend/tests/fleet/refresh/mod.rs +++ b/backend/tests/fleet/refresh/mod.rs @@ -9,12 +9,17 @@ use std::{ use async_trait::async_trait; use chrono::{Duration, TimeZone, Utc}; -use codex_proxy_rs::fleet::account::AccountStatus; use codex_proxy_rs::fleet::refresh::{RedisRefreshLeaseStore, RefreshPolicy}; use codex_proxy_rs::fleet::refresh::{RefreshFailure, TokenPair, TokenRefresher}; use codex_proxy_rs::fleet::store::{ AccountClaimsUpdate, NewAccount, PgAccountStore, StoredAccount, }; +use codex_proxy_rs::fleet::{ + account::AccountStatus, + pool::{AccountPoolOptions, AccountPoolService}, + store::AccountStore, +}; +use codex_proxy_rs::telemetry::account_usage::store::PgAccountUsageStore; use secrecy::{ExposeSecret, SecretString}; use serde_json::json; use tokio::{ diff --git a/backend/tests/fleet/refresh/success.rs b/backend/tests/fleet/refresh/success.rs index c99b2d5e5..3b240c518 100644 --- a/backend/tests/fleet/refresh/success.rs +++ b/backend/tests/fleet/refresh/success.rs @@ -1,5 +1,85 @@ use super::*; +#[tokio::test] +async fn token_refresh_task_should_sync_refreshed_token_into_runtime_account_pool() { + let (pool, _guard) = init_test_db("token-refresh-runtime-pool-sync").await; + let store = PgAccountStore::new(pool.clone()); + let now = Utc.with_ymd_and_hms(2026, 6, 18, 7, 0, 0).unwrap(); + let new_expires_at = now + Duration::hours(1); + let new_access_token = test_jwt(new_expires_at.timestamp()); + store + .insert(NewAccount { + id: "acct-refresh-runtime-pool".to_string(), + email: Some("runtime-pool@example.com".to_string()), + account_id: Some("chatgpt-runtime-pool".to_string()), + user_id: Some("user-runtime-pool".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-runtime-pool".to_string().into())), + added_at: None, + access_token_expires_at: Some(now + Duration::seconds(30)), + status: AccountStatus::Active, + }) + .await + .expect("account should be inserted"); + let account_pool = Arc::new(AccountPoolService::new( + Arc::new(store.clone()) as Arc, + Arc::new(PgAccountUsageStore::new(pool)), + AccountPoolOptions { + max_concurrent_per_account: 1, + ..AccountPoolOptions::default() + }, + 0, + )); + account_pool + .restore_from_store() + .await + .expect("runtime account pool should restore"); + let refresher = StaticTokenRefresher { + response: Arc::new(Mutex::new(Ok(TokenPair { + access_token: new_access_token.clone(), + refresh_token: None, + }))), + }; + let task = codex_proxy_rs::fleet::refresh::TokenRefreshService::new( + store.clone(), + RefreshPolicy { + refresh_margin_seconds: 300, + refresh_concurrency: 1, + }, + refresher, + ) + .with_account_pool_sync(Arc::clone(&account_pool)); + + task.schedule_account_timers_once_at(now) + .await + .expect("refresh timer should be scheduled"); + let stored = wait_for_account(&store, "acct-refresh-runtime-pool", |account| { + account.access_token.expose_secret() == new_access_token.as_str() + }) + .await; + let capacity = timeout(StdDuration::from_secs(2), async { + loop { + let capacity = account_pool + .capacity_summary(now + Duration::minutes(1)) + .await; + if capacity.available_slots == 1 { + return capacity; + } + sleep(StdDuration::from_millis(20)).await; + } + }) + .await + .expect("refreshed account should remain available in the runtime pool"); + task.shutdown().await; + + assert_eq!(stored.status, AccountStatus::Active); + assert_eq!(capacity.available_slots, 1); +} + #[tokio::test] async fn token_refresh_task_should_persist_refreshed_access_token_and_keep_refresh_token() { let (pool, _guard) = init_test_db("token-refresh").await; diff --git a/backend/tests/fleet/store/mod.rs b/backend/tests/fleet/store/mod.rs index 043c08002..d827a6b18 100644 --- a/backend/tests/fleet/store/mod.rs +++ b/backend/tests/fleet/store/mod.rs @@ -793,6 +793,66 @@ async fn account_store_runtime_state_sync_should_not_clear_newer_future_quota_co assert_eq!(stored.3, Some(future_cooldown)); } +#[tokio::test] +async fn account_store_runtime_state_sync_should_not_expire_a_newer_refreshed_token() { + let (pool, _dir) = account_store_parts("accounts", 49).await; + let repo = PgAccountStore::new(pool.clone()); + let stale_expires_at = pg_timestamp(Utc::now() - Duration::minutes(1)); + repo.insert(NewAccount { + id: "acct_a".to_string(), + email: Some("user@example.com".to_string()), + account_id: Some("chatgpt-account".to_string()), + user_id: Some("chatgpt-user".to_string()), + label: None, + plan_type: Some("plus".to_string()), + access_token: SecretString::new("stale-access-token".to_string().into()), + refresh_token: Some(SecretString::new("refresh-token".to_string().into())), + access_token_expires_at: Some(stale_expires_at), + status: AccountStatus::Active, + added_at: None, + }) + .await + .unwrap(); + let mut stale_account = AccountStore::list_pool_accounts(&repo) + .await + .unwrap() + .into_iter() + .find(|account| account.id == "acct_a") + .unwrap(); + let refreshed_expires_at = pg_timestamp(Utc::now() + Duration::hours(1)); + repo.update_from_claims( + "acct_a", + AccountClaimsUpdate { + email: Some("user@example.com".to_string()), + account_id: Some("chatgpt-account".to_string()), + user_id: Some("chatgpt-user".to_string()), + plan_type: Some("plus".to_string()), + access_token: SecretString::new("refreshed-access-token".to_string().into()), + refresh_token: None, + access_token_expires_at: Some(refreshed_expires_at), + next_refresh_at: Some(refreshed_expires_at - Duration::minutes(5)), + status: AccountStatus::Active, + }, + ) + .await + .unwrap(); + stale_account.status = AccountStatus::Expired; + + repo.sync_runtime_account_state(&stale_account) + .await + .unwrap(); + let stored = repo.get("acct_a").await.unwrap().unwrap(); + + assert_eq!(stored.status, AccountStatus::Active); + assert_eq!( + ( + stored.access_token.expose_secret(), + stored.access_token_expires_at + ), + ("refreshed-access-token", Some(refreshed_expires_at)) + ); +} + #[tokio::test] async fn account_store_runtime_state_sync_should_not_regress_newer_usage_window() { let (pool, _dir) = account_store_parts("accounts", 43).await;