diff --git a/src/admin/accounts/routes.rs b/src/admin/accounts/routes.rs index 94aa5c215..60a9cbaf2 100644 --- a/src/admin/accounts/routes.rs +++ b/src/admin/accounts/routes.rs @@ -26,7 +26,9 @@ use crate::{ }, admin::{ accounts::service::{ - AdminAccountError, AdminAccountMetadata, AdminAccountUpdate, OAuthExchangeInput, + AdminAccountError, AdminAccountHealthCheck, AdminAccountMetadata, + AdminAccountRefreshOutcome, AdminAccountRefreshResult, AdminAccountUpdate, + OAuthExchangeInput, }, monitoring::{billing, service::AdminUsageRecord}, }, @@ -405,6 +407,73 @@ enum AccountUpdateData { Account(Box), } +#[derive(Debug, Clone, Serialize)] +#[serde(rename_all = "camelCase")] +struct AccountRefreshData { + id: String, + email: Option, + previous_status: String, + result: &'static str, + error: Option, + duration_ms: i64, +} + +impl From for AccountRefreshData { + fn from(result: AdminAccountRefreshResult) -> Self { + Self { + id: result.id, + email: result.email, + previous_status: account_status_str(result.previous_status).to_string(), + result: account_refresh_outcome_str(result.outcome), + error: result.error, + duration_ms: result.duration_ms, + } + } +} + +#[derive(Debug, Clone, Deserialize, Default)] +#[serde(default)] +pub(crate) struct AccountHealthCheckRequest { + ids: Option>, + stagger_ms: Option, + concurrency: Option, +} + +#[derive(Debug, Clone, Serialize)] +#[serde(rename_all = "camelCase")] +struct AccountHealthCheckSummaryData { + total: usize, + alive: usize, + dead: usize, + skipped: usize, +} + +#[derive(Debug, Clone, Serialize)] +#[serde(rename_all = "camelCase")] +struct AccountHealthCheckData { + summary: AccountHealthCheckSummaryData, + results: Vec, +} + +impl From for AccountHealthCheckData { + fn from(result: AdminAccountHealthCheck) -> Self { + let summary = AccountHealthCheckSummaryData { + total: result.results.len(), + alive: result.alive(), + dead: result.dead(), + skipped: result.skipped(), + }; + Self { + summary, + results: result + .results + .into_iter() + .map(AccountRefreshData::from) + .collect(), + } + } +} + #[derive(Debug, Clone, Serialize)] #[serde(rename_all = "camelCase")] struct AccountImportData { @@ -541,9 +610,43 @@ pub(crate) async fn refresh_account( .refresh_account(&account_id) .await { - Ok(account) => Ok(AdminResponse::new( + Ok(result) => Ok(AdminResponse::new( StatusCode::OK, - AdminEnvelope::ok(AdminAccountData::from(account)), + AdminEnvelope::ok(AccountRefreshData::from(result)), + )), + Err(error) => Err(account_error(&error)), + } +} + +/// `POST /api/admin/accounts/health-check` +pub(crate) async fn health_check_accounts( + State(state): State, + headers: HeaderMap, + Json(payload): Json, +) -> Result { + require_admin_auth(&state, &headers).await?; + let stagger_ms = payload.stagger_ms.unwrap_or(3_000); + if !(500..=30_000).contains(&stagger_ms) && stagger_ms != 0 { + return Err(AdminError::bad_request( + "stagger_ms must be between 500 and 30000", + )); + } + if let Some(concurrency) = payload.concurrency { + if !(1..=10).contains(&concurrency) { + return Err(AdminError::bad_request( + "concurrency must be between 1 and 10", + )); + } + } + match state + .services + .admin_accounts + .health_check_accounts(payload.ids, stagger_ms, payload.concurrency.unwrap_or(2)) + .await + { + Ok(result) => Ok(AdminResponse::new( + StatusCode::OK, + AdminEnvelope::ok(AccountHealthCheckData::from(result)), )), Err(error) => Err(account_error(&error)), } @@ -1464,6 +1567,14 @@ fn account_error(error: &AdminAccountError) -> AdminError { } } +fn account_refresh_outcome_str(outcome: AdminAccountRefreshOutcome) -> &'static str { + match outcome { + AdminAccountRefreshOutcome::Alive => "alive", + AdminAccountRefreshOutcome::Dead => "dead", + AdminAccountRefreshOutcome::Skipped => "skipped", + } +} + fn account_not_found() -> AdminError { AdminError::not_found("Account not found") } diff --git a/src/admin/accounts/service/importing.rs b/src/admin/accounts/service/importing.rs index b0f9f40fc..b6d817d2d 100644 --- a/src/admin/accounts/service/importing.rs +++ b/src/admin/accounts/service/importing.rs @@ -92,8 +92,6 @@ impl AdminAccountService { .as_ref() .map(|claims| claims.expires_at) .or(access_token_expires_at); - let next_refresh_at = access_token_expires_at - .map(|expires_at| self.next_refresh_at_for_expires_at(expires_at)); let claims = resolved_tokens.claims.as_ref(); let mut plan_type = claims .and_then(|claims| claims.plan_type.clone()) @@ -145,6 +143,8 @@ impl AdminAccountService { chatgpt_user_id.as_deref(), ) .await?; + let next_refresh_at = access_token_expires_at + .map(|expires_at| self.next_refresh_at_for_expires_at(&account_id, expires_at)); if plan_type.is_none() { plan_type = supplemental.plan_type; } @@ -297,7 +297,7 @@ impl AdminAccountService { ), )) .ok_or(AdminAccountError::TokenRequired)?; - refresh_token = refreshed.refresh_token.or(Some(existing_refresh_token)); + refresh_token = refreshed.refresh_token; let claims = crate::upstream::accounts::token_refresh::manual_account_claims( &access_token, Utc::now(), @@ -337,7 +337,7 @@ impl AdminAccountService { crate::upstream::accounts::importing::normalize_bearer_token(&refreshed.access_token), )) .ok_or(AdminAccountError::TokenRequired)?; - refresh_token = refreshed.refresh_token.or(Some(existing_refresh_token)); + refresh_token = refreshed.refresh_token; let claims = crate::upstream::accounts::token_refresh::manual_account_claims( &access_token, Utc::now(), diff --git a/src/admin/accounts/service/lifecycle.rs b/src/admin/accounts/service/lifecycle.rs index 324335592..ab16d0b7a 100644 --- a/src/admin/accounts/service/lifecycle.rs +++ b/src/admin/accounts/service/lifecycle.rs @@ -1,11 +1,22 @@ -use chrono::Utc; -use secrecy::{ExposeSecret, SecretString}; +use std::{ + collections::{hash_map::DefaultHasher, HashSet}, + hash::{Hash, Hasher}, +}; -use crate::infra::json::{NumberedPage, Page}; +use chrono::Utc; +use futures::{stream, StreamExt}; +use secrecy::{ExposeSecret, SecretString}; +use tokio::time::sleep; + +use crate::{ + infra::json::{NumberedPage, Page}, + upstream::accounts::store::AccountStore, +}; use super::{ types::{ - parse_account_status, stored_to_admin_metadata, AdminAccountError, AdminAccountMetadata, + parse_account_status, stored_to_admin_metadata, AdminAccountError, AdminAccountHealthCheck, + AdminAccountMetadata, AdminAccountRefreshOutcome, AdminAccountRefreshResult, AdminAccountUpdate, BatchDeleteAccounts, ManualCreateTokens, UpdatedAccountStatus, }, AdminAccountService, @@ -97,10 +108,7 @@ impl AdminAccountService { .ok_or(AdminAccountError::TokenRequired)?; ManualCreateTokens { access_token, - refresh_token_for_new: token_pair - .refresh_token - .clone() - .or_else(|| Some(refresh_token.clone())), + refresh_token_for_new: token_pair.refresh_token.clone(), refresh_token_for_existing: token_pair.refresh_token, } } else { @@ -137,7 +145,7 @@ impl AdminAccountService { .map(|token| SecretString::new(token.into())), access_token_expires_at: Some(claims.expires_at), next_refresh_at: Some( - self.next_refresh_at_for_expires_at(claims.expires_at), + self.next_refresh_at_for_expires_at(&existing.id, claims.expires_at), ), status: crate::upstream::accounts::model::AccountStatus::Active, }, @@ -292,16 +300,34 @@ impl AdminAccountService { pub async fn refresh_account( &self, account_id: &str, - ) -> Result { + ) -> Result { let account = match self.store.get(account_id).await { Ok(Some(account)) => account, Ok(None) => return Err(AdminAccountError::NotFound), Err(_) => return Err(AdminAccountError::Inspect), }; - let Some(refresh_token) = account.refresh_token.as_ref() else { - return Err(AdminAccountError::TokenRequired); + let previous_status = account.status; + let started_at = std::time::Instant::now(); + + let skipped = |error: &'static str| AdminAccountRefreshResult { + id: account.id.clone(), + email: account.email.clone(), + previous_status, + outcome: AdminAccountRefreshOutcome::Skipped, + error: Some(error.to_string()), + duration_ms: elapsed_millis_i64(started_at), }; + let Some(refresh_token) = account.refresh_token.as_ref() else { + return Ok(skipped("no refresh token")); + }; + if account.status == crate::upstream::accounts::model::AccountStatus::Disabled { + return Ok(skipped("manually disabled")); + } + if account.status == crate::upstream::accounts::model::AccountStatus::Refreshing { + return Ok(skipped("refresh already in progress")); + } + match self .token_refresher .refresh(refresh_token.expose_secret()) @@ -334,7 +360,7 @@ impl AdminAccountService { .map(|token| SecretString::new(token.into())), access_token_expires_at: Some(claims.expires_at), next_refresh_at: Some( - self.next_refresh_at_for_expires_at(claims.expires_at), + self.next_refresh_at_for_expires_at(account_id, claims.expires_at), ), status: crate::upstream::accounts::model::AccountStatus::Active, }, @@ -345,11 +371,18 @@ impl AdminAccountService { return Err(AdminAccountError::NotFound); } self.sync_account_pool(account_id).await?; + Ok(AdminAccountRefreshResult { + id: account.id, + email: account.email, + previous_status, + outcome: AdminAccountRefreshOutcome::Alive, + error: None, + duration_ms: elapsed_millis_i64(started_at), + }) } Err(failure) => { - if let Some(status) = - manual_refresh_failure_status(&failure, account.access_token.expose_secret()) - { + if manual_refresh_failure_is_permanent(&failure) { + let status = crate::upstream::accounts::model::AccountStatus::Expired; let updated = self .store .set_status(account_id, status) @@ -373,46 +406,85 @@ impl AdminAccountService { self.sync_account_pool_best_effort(account_id, "account refresh failure") .await; } - return Err(AdminAccountError::RefreshTokenExchange(failure)); + Ok(AdminAccountRefreshResult { + id: account.id, + email: account.email, + previous_status, + outcome: AdminAccountRefreshOutcome::Dead, + error: Some(format!("token refresh exchange failed: {failure}")), + duration_ms: elapsed_millis_i64(started_at), + }) } } + } - self.store - .get(account_id) + pub async fn health_check_accounts( + &self, + ids: Option>, + stagger_ms: u64, + concurrency: usize, + ) -> Result { + let accounts = self + .store + .list_pool_accounts() .await - .map_err(|_| AdminAccountError::Inspect)? - .map(stored_to_admin_metadata) - .ok_or(AdminAccountError::NotFound) + .map_err(|_| AdminAccountError::List)?; + let ids = ids.map(|ids| ids.into_iter().collect::>()); + let candidate_ids = accounts + .into_iter() + .filter(|account| { + ids.as_ref() + .is_none_or(|ids| ids.contains(account.id.as_str())) + }) + .map(|account| account.id) + .collect::>(); + let results = stream::iter(candidate_ids.into_iter().enumerate().map( + |(index, account_id)| async move { + if index > 0 && stagger_ms > 0 { + let base_delay = stagger_ms.saturating_mul(index.min(concurrency) as u64); + let delay = stable_jittered_millis(&account_id, base_delay, 0.30); + sleep(std::time::Duration::from_millis(delay)).await; + } + self.refresh_account(&account_id).await + }, + )) + .buffer_unordered(concurrency.max(1)) + .collect::>() + .await + .into_iter() + .collect::, _>>()?; + + Ok(AdminAccountHealthCheck { results }) } } -fn manual_refresh_failure_status( +fn stable_jittered_millis(account_id: &str, base_millis: u64, variance: f64) -> u64 { + if base_millis == 0 { + return 0; + } + + let mut hasher = DefaultHasher::new(); + account_id.hash(&mut hasher); + "health-check-stagger".hash(&mut hasher); + let unit = hasher.finish() as f64 / u64::MAX as f64; + let factor = (1.0 - variance) + unit * variance * 2.0; + (base_millis as f64 * factor) + .round() + .clamp(0.0, u64::MAX as f64) as u64 +} + +fn manual_refresh_failure_is_permanent( failure: &crate::upstream::accounts::token_refresh::RefreshFailure, - current_access_token: &str, -) -> Option { +) -> bool { match failure { - crate::upstream::accounts::token_refresh::RefreshFailure::QuotaExhausted => { - Some(crate::upstream::accounts::model::AccountStatus::QuotaExhausted) - } - crate::upstream::accounts::token_refresh::RefreshFailure::Banned => { - Some(crate::upstream::accounts::model::AccountStatus::Banned) - } - crate::upstream::accounts::token_refresh::RefreshFailure::Disabled => { - Some(crate::upstream::accounts::model::AccountStatus::Disabled) - } - crate::upstream::accounts::token_refresh::RefreshFailure::InvalidGrant - if matches!( - crate::upstream::accounts::token_refresh::jwt_expiry( - current_access_token, - Utc::now() - ), - crate::upstream::accounts::token_refresh::JwtExpiry::Expired - ) => - { - Some(crate::upstream::accounts::model::AccountStatus::Expired) - } crate::upstream::accounts::token_refresh::RefreshFailure::InvalidGrant + | crate::upstream::accounts::token_refresh::RefreshFailure::Banned => true, + crate::upstream::accounts::token_refresh::RefreshFailure::QuotaExhausted | crate::upstream::accounts::token_refresh::RefreshFailure::RetryableTransport - | crate::upstream::accounts::token_refresh::RefreshFailure::Transport => None, + | crate::upstream::accounts::token_refresh::RefreshFailure::Transport => false, } } + +fn elapsed_millis_i64(started_at: std::time::Instant) -> i64 { + i64::try_from(started_at.elapsed().as_millis()).unwrap_or(i64::MAX) +} diff --git a/src/admin/accounts/service/mod.rs b/src/admin/accounts/service/mod.rs index 2dbe2a5ec..d67336213 100644 --- a/src/admin/accounts/service/mod.rs +++ b/src/admin/accounts/service/mod.rs @@ -11,7 +11,7 @@ mod types; use std::sync::Arc as StdArc; -use chrono::{DateTime, Duration}; +use chrono::DateTime; use crate::{ upstream::accounts::{ @@ -25,9 +25,9 @@ use crate::{ }; pub use types::{ - AdminAccountError, AdminAccountMetadata, AdminAccountUpdate, BatchDeleteAccounts, - ExportedAccounts, ImportedAccounts, OAuthAuthorizeResult, OAuthExchangeInput, - UpdatedAccountStatus, + AdminAccountError, AdminAccountHealthCheck, AdminAccountMetadata, AdminAccountRefreshOutcome, + AdminAccountRefreshResult, AdminAccountUpdate, BatchDeleteAccounts, ExportedAccounts, + ImportedAccounts, OAuthAuthorizeResult, OAuthExchangeInput, UpdatedAccountStatus, }; #[derive(Clone)] @@ -72,13 +72,14 @@ impl AdminAccountService { pub(crate) fn next_refresh_at_for_expires_at( &self, + account_id: &str, expires_at: DateTime, ) -> DateTime { - let margin_seconds = self - .refresh_policy - .refresh_margin_seconds() - .min(i64::MAX as u64) as i64; - expires_at - Duration::seconds(margin_seconds) + crate::upstream::accounts::token_refresh::jittered_refresh_at( + account_id, + expires_at, + self.refresh_policy.refresh_margin_seconds(), + ) } pub(crate) async fn sync_account_pool( diff --git a/src/admin/accounts/service/types.rs b/src/admin/accounts/service/types.rs index bab6e3cb7..06d5b0a12 100644 --- a/src/admin/accounts/service/types.rs +++ b/src/admin/accounts/service/types.rs @@ -117,6 +117,51 @@ pub struct UpdatedAccountStatus { pub status: AcctStatus, } +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum AdminAccountRefreshOutcome { + Alive, + Dead, + Skipped, +} + +#[derive(Debug, Clone)] +pub struct AdminAccountRefreshResult { + pub id: String, + pub email: Option, + pub previous_status: AcctStatus, + pub outcome: AdminAccountRefreshOutcome, + pub error: Option, + pub duration_ms: i64, +} + +#[derive(Debug, Clone, Default)] +pub struct AdminAccountHealthCheck { + pub results: Vec, +} + +impl AdminAccountHealthCheck { + pub fn alive(&self) -> usize { + self.results + .iter() + .filter(|result| result.outcome == AdminAccountRefreshOutcome::Alive) + .count() + } + + pub fn dead(&self) -> usize { + self.results + .iter() + .filter(|result| result.outcome == AdminAccountRefreshOutcome::Dead) + .count() + } + + pub fn skipped(&self) -> usize { + self.results + .iter() + .filter(|result| result.outcome == AdminAccountRefreshOutcome::Skipped) + .count() + } +} + #[derive(Debug, Clone, Default)] pub struct AdminAccountUpdate { pub label: Option>, @@ -280,6 +325,11 @@ pub(super) fn import_status_from_usage_error(error: &CodexClientError) -> Option } pub(super) fn parse_account_status(status: &str) -> Result { - crate::upstream::accounts::importing::parse_account_status(status) - .map_err(|_| AdminAccountError::InvalidStatus(status.trim().to_ascii_lowercase())) + match status.trim().to_ascii_lowercase().as_str() { + "active" => Ok(AcctStatus::Active), + "disabled" => Ok(AcctStatus::Disabled), + _ => Err(AdminAccountError::InvalidStatus( + status.trim().to_ascii_lowercase(), + )), + } } diff --git a/src/admin/router.rs b/src/admin/router.rs index 40e714876..62dc8a4ad 100644 --- a/src/admin/router.rs +++ b/src/admin/router.rs @@ -10,8 +10,8 @@ use crate::runtime::state::AppState; use super::{ accounts::routes::{ account_models, account_quota, accounts, batch_delete_accounts, create_account, - export_accounts, import_accounts, oauth_authorize_account, oauth_exchange_account, - refresh_account, test_account_connection, update_account, + export_accounts, health_check_accounts, import_accounts, oauth_authorize_account, + oauth_exchange_account, refresh_account, test_account_connection, update_account, }, auth::session::{login, logout, session_status}, keys::routes::{api_keys, batch_delete_api_keys, create_api_key, update_api_key}, @@ -62,6 +62,10 @@ pub fn router() -> Router { .route("/api/admin/accounts/delete", post(batch_delete_accounts)) .route("/api/admin/accounts/update", post(update_account)) .route("/api/admin/accounts/refresh", post(refresh_account)) + .route( + "/api/admin/accounts/health-check", + post(health_check_accounts), + ) .route("/api/admin/accounts/quota", get(account_quota)) .route("/api/admin/usage/records", get(usage_records)) .route( diff --git a/src/proxy/dispatch/chat.rs b/src/proxy/dispatch/chat.rs index 58ce88997..45bfdf926 100644 --- a/src/proxy/dispatch/chat.rs +++ b/src/proxy/dispatch/chat.rs @@ -34,10 +34,12 @@ use crate::{ upstream::accounts::{ model::AccountStatus, pool::{AccountAcquireRequest, RuntimeAccountPoolService}, + token_refresh::RuntimeTokenRefreshService, }, upstream::{ models::ModelCatalog, protocol::responses::{apply_response_model_options, CodexResponsesRequest}, + token_client::OpenAiTokenClient, transport::{ backend_transport_for_response_request, is_banned_upstream_error, CodexBackendClient, CodexClientError, @@ -52,6 +54,7 @@ pub struct ChatDispatchService { models: Arc, codex: Arc, usage_records: Arc, + token_refresh: Arc>, installation_id: Option, cloudflare: CloudflareRecovery, } @@ -62,6 +65,7 @@ impl ChatDispatchService { models: Arc, codex: Arc, usage_records: Arc, + token_refresh: Arc>, installation_id: Option, cloudflare: CloudflareRecovery, ) -> Self { @@ -70,11 +74,30 @@ impl ChatDispatchService { models, codex, usage_records, + token_refresh, installation_id, cloudflare, } } + fn trigger_refresh_after_auth_failure(&self, account_id: &str, status: AccountStatus) { + if status != AccountStatus::Expired { + return; + } + + let token_refresh = self.token_refresh.clone(); + let account_id = account_id.to_string(); + tokio::spawn(async move { + if let Err(error) = token_refresh.trigger_account_refresh_now(&account_id).await { + tracing::warn!( + account_id = %account_id, + error = %error, + "reactive token refresh after upstream auth failure failed" + ); + } + }); + } + /// 调度非流式 Chat Completions 请求到 Codex Responses 上游。 pub async fn complete( &self, @@ -288,6 +311,7 @@ impl ChatDispatchService { self.account_pool .set_status(&release_account_id, account_status) .await; + self.trigger_refresh_after_auth_failure(&release_account_id, account_status); excluded_account_ids.push(release_account_id); } Err(error) if is_cloudflare_challenge_upstream_error(&error) => { diff --git a/src/proxy/dispatch/responses.rs b/src/proxy/dispatch/responses.rs index c5b37c3a4..89c78b510 100644 --- a/src/proxy/dispatch/responses.rs +++ b/src/proxy/dispatch/responses.rs @@ -56,6 +56,7 @@ use crate::{ upstream::accounts::{ model::{Account, AccountStatus}, pool::{AccountAcquireRequest, RuntimeAccountPoolService}, + token_refresh::RuntimeTokenRefreshService, }, upstream::{ models::ModelCatalog, @@ -71,6 +72,7 @@ use crate::{ encode_sse_event, parse_sse_events, sse_body_has_done, SseError, DONE_SSE_FRAME, }, }, + token_client::OpenAiTokenClient, transport::{ backend_transport_for_response_request, is_banned_auth_signal, is_banned_upstream_error, CodexBackendClient, CodexBackendResponse, @@ -107,6 +109,7 @@ pub struct ResponseDispatchService { session_affinity: Arc, reasoning_replay: Arc>, usage_records: Arc, + token_refresh: Arc>, installation_id: Option, cloudflare: CloudflareRecovery, } @@ -154,6 +157,7 @@ impl ResponseDispatchService { codex: Arc, session_affinity: Arc, usage_records: Arc, + token_refresh: Arc>, installation_id: Option, cloudflare: CloudflareRecovery, ) -> Self { @@ -166,11 +170,30 @@ impl ResponseDispatchService { DEFAULT_REASONING_REPLAY_TTL_SECS, )))), usage_records, + token_refresh, installation_id, cloudflare, } } + fn trigger_refresh_after_auth_failure(&self, account_id: &str, status: AccountStatus) { + if status != AccountStatus::Expired { + return; + } + + let token_refresh = self.token_refresh.clone(); + let account_id = account_id.to_string(); + tokio::spawn(async move { + if let Err(error) = token_refresh.trigger_account_refresh_now(&account_id).await { + tracing::warn!( + account_id = %account_id, + error = %error, + "reactive token refresh after upstream auth failure failed" + ); + } + }); + } + async fn prepare_response_session( &self, request: &mut CodexResponsesRequest, @@ -711,6 +734,10 @@ impl ResponseDispatchService { self.account_pool .set_status(&release_account_id, account_status) .await; + self.trigger_refresh_after_auth_failure( + &release_account_id, + account_status, + ); excluded_account_ids.push(release_account_id); continue; } @@ -771,6 +798,7 @@ impl ResponseDispatchService { self.account_pool .set_status(&release_account_id, account_status) .await; + self.trigger_refresh_after_auth_failure(&release_account_id, account_status); excluded_account_ids.push(release_account_id); } Err(error) if is_cloudflare_challenge_upstream_error(&error) => { @@ -1338,6 +1366,10 @@ impl ResponseDispatchService { self.account_pool .set_status(&release_account_id, account_status) .await; + self.trigger_refresh_after_auth_failure( + &release_account_id, + account_status, + ); excluded_account_ids.push(release_account_id); self.account_pool.release(&account.id).await; continue; @@ -1447,6 +1479,7 @@ impl ResponseDispatchService { self.account_pool .set_status(&release_account_id, account_status) .await; + self.trigger_refresh_after_auth_failure(&release_account_id, account_status); excluded_account_ids.push(release_account_id); } Err(error) if is_cloudflare_challenge_upstream_error(&error) => { @@ -1837,6 +1870,7 @@ impl ResponseDispatchService { self.account_pool .set_status(&release_account_id, account_status) .await; + self.trigger_refresh_after_auth_failure(&release_account_id, account_status); excluded_account_ids.push(release_account_id); } Err(error) if is_cloudflare_challenge_upstream_error(&error) => { diff --git a/src/runtime/services.rs b/src/runtime/services.rs index efd21350f..9bd8db2f0 100644 --- a/src/runtime/services.rs +++ b/src/runtime/services.rs @@ -34,12 +34,13 @@ use crate::{ store::{AccountStore as AccountStoreTrait, SqliteAccountStore}, token_refresh::{ RefreshLeaseStore as SqliteRefreshLeaseStore, RefreshPolicy, RuntimeRefreshPolicy, + RuntimeTokenRefreshService, }, }, upstream::{ fingerprint::{Fingerprint, FingerprintRepository}, models::{ModelService, ModelServiceError, ModelSnapshotStore}, - token_client::{default_openai_token_client, TokenClientConfig}, + token_client::{default_openai_token_client, OpenAiTokenClient, TokenClientConfig}, transport::{ build_reqwest_client, tls::CustomCaError, CodexBackendClient, CodexModelCatalogClient, }, @@ -90,6 +91,7 @@ pub struct Services { pub usage_records: StdArc, pub usage: StdArc, pub account_pool: StdArc, + pub token_refresh: StdArc>, pub chat: StdArc, pub responses: StdArc, pub session_affinity: StdArc, @@ -204,7 +206,15 @@ impl Services { runtime_client_keys.clone(), )); let client_keys = StdArc::new(ClientKeyService::new(runtime_client_keys)); - let token_client = StdArc::new(default_openai_token_client(token_client_config(config))); + let token_client = default_openai_token_client(token_client_config(config)); + let token_refresh = StdArc::new( + RuntimeTokenRefreshService::new( + stores.accounts.clone(), + refresh_policy.clone(), + token_client.clone(), + ) + .with_refresh_lease_store(stores.refresh_leases.clone()), + ); let upstream_client: StdArc = codex.clone(); let snapshot_store: StdArc = StdArc::new( crate::upstream::models::SqliteModelSnapshotStore::new(stores.accounts.pool().clone()), @@ -222,7 +232,7 @@ impl Services { codex: codex.clone(), models: models.clone(), account_pool: account_pool.clone(), - token_refresher: token_client, + token_refresher: StdArc::new(token_client), oauth: crate::admin::accounts::service::oauth::AccountOAuthService::new( reqwest::Client::new(), config.auth.oauth_client_id.clone(), @@ -250,6 +260,7 @@ impl Services { models.clone(), codex.clone(), usage_records.clone(), + token_refresh.clone(), installation_id.clone(), cloudflare_recovery.clone(), )); @@ -259,6 +270,7 @@ impl Services { codex.clone(), session_affinity.clone(), usage_records.clone(), + token_refresh.clone(), installation_id.clone(), cloudflare_recovery, )); @@ -275,6 +287,7 @@ impl Services { usage_records, usage, account_pool, + token_refresh, chat, responses, session_affinity, diff --git a/src/runtime/tasks/coordinator.rs b/src/runtime/tasks/coordinator.rs index 4e7783f36..89feaaef6 100644 --- a/src/runtime/tasks/coordinator.rs +++ b/src/runtime/tasks/coordinator.rs @@ -4,8 +4,6 @@ use std::path::PathBuf; use tokio::task::JoinHandle; -use crate::upstream::token_client::{default_openai_token_client, TokenClientConfig}; - use super::{ cookie_cleanup::CookieCleanupTask, fingerprint_update::{ @@ -68,13 +66,7 @@ impl TaskCoordinator { if config.auth.refresh_enabled { coordinator.push( "token_refresh", - TokenRefreshTask::new( - stores.accounts.clone(), - services.refresh_policy.clone(), - default_openai_token_client(token_client_config(config)), - ) - .with_refresh_lease_store(stores.refresh_leases.clone()) - .start(), + TokenRefreshTask::from_service((*services.token_refresh).clone()).start(), ); } coordinator.push( @@ -119,13 +111,6 @@ impl TaskCoordinator { } } -fn token_client_config(config: &crate::runtime::state::RuntimeConfig) -> TokenClientConfig { - TokenClientConfig { - client_id: config.auth.oauth_client_id.clone(), - token_endpoint: config.auth.oauth_token_endpoint.clone(), - } -} - /// 后台任务关闭句柄。 pub enum SchedulerHandle { /// 通过 channel 发送关闭信号。 diff --git a/src/runtime/tasks/token_refresh.rs b/src/runtime/tasks/token_refresh.rs index 5cef19ae8..0b024b535 100644 --- a/src/runtime/tasks/token_refresh.rs +++ b/src/runtime/tasks/token_refresh.rs @@ -37,6 +37,14 @@ where } } + /// 使用已创建的运行时刷新服务构造后台任务。 + pub fn from_service(service: RuntimeTokenRefreshService) -> Self { + Self { + service, + interval_secs: DEFAULT_INTERVAL_SECS, + } + } + /// 使用刷新租约存储保护账号刷新。 pub fn with_refresh_lease_store(mut self, refresh_leases: RefreshLeaseStore) -> Self { self.service = self.service.with_refresh_lease_store(refresh_leases); diff --git a/src/upstream/accounts/store.rs b/src/upstream/accounts/store.rs index 89e53a279..ec00c0a8e 100644 --- a/src/upstream/accounts/store.rs +++ b/src/upstream/accounts/store.rs @@ -465,7 +465,10 @@ set refresh_token = ?, access_token_expires_at = ?, next_refresh_at = ?, - status = ?, + status = case + when status in ('disabled', 'banned') then status + else ? + end, updated_at = ? where id = ?"; @@ -479,7 +482,10 @@ set access_token = ?, access_token_expires_at = ?, next_refresh_at = ?, - status = ?, + status = case + when status in ('disabled', 'banned') then status + else ? + end, updated_at = ? where id = ?"; diff --git a/src/upstream/accounts/token_refresh/mod.rs b/src/upstream/accounts/token_refresh/mod.rs index 5da4331c8..13ae602d3 100644 --- a/src/upstream/accounts/token_refresh/mod.rs +++ b/src/upstream/accounts/token_refresh/mod.rs @@ -1,6 +1,10 @@ //! 刷新租约存储、JWT 解码与账号 claims 验证。 -use std::sync::{Arc, RwLock}; +use std::{ + collections::hash_map::DefaultHasher, + hash::{Hash, Hasher}, + sync::{Arc, RwLock}, +}; use async_trait::async_trait; use base64::{engine::general_purpose::URL_SAFE_NO_PAD, Engine as _}; @@ -217,6 +221,29 @@ impl From for RuntimeRefreshPolicy { } } +/// 按账号稳定扰动刷新提前量，范围与 TS 版 `jitter(margin, 0.15)` 一致。 +pub(crate) fn jittered_refresh_at( + account_id: &str, + expires_at: DateTime, + margin_seconds: u64, +) -> DateTime { + expires_at - Duration::seconds(jittered_refresh_margin_seconds(account_id, margin_seconds)) +} + +fn jittered_refresh_margin_seconds(account_id: &str, margin_seconds: u64) -> i64 { + if margin_seconds == 0 { + return 0; + } + + let mut hasher = DefaultHasher::new(); + account_id.hash(&mut hasher); + let unit = hasher.finish() as f64 / u64::MAX as f64; + let factor = 0.85 + unit * 0.30; + (margin_seconds as f64 * factor) + .round() + .clamp(0.0, i64::MAX as f64) as i64 +} + /// 触发刷新动作的原因。 #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub enum RefreshTrigger { @@ -238,9 +265,6 @@ pub enum RefreshFailure { /// 账号被上游封禁。 #[error("account is banned")] Banned, - /// 账号被显式禁用。 - #[error("account is disabled")] - Disabled, /// 刷新请求在到达服务端前失败，可安全复用当前 refresh token 重试。 #[error("refresh transport failed before server processing")] RetryableTransport, @@ -326,6 +350,7 @@ where match self.client.refresh(refresh_token).await { Ok(token_pair) => Ok(apply_token_pair(account, token_pair)), + Err(RefreshFailure::QuotaExhausted) => Err(RefreshError::Transport), Err(RefreshFailure::RetryableTransport) => Err(RefreshError::RetryableTransport), Err(RefreshFailure::Transport) => Err(RefreshError::Transport), Err(error) => Ok(apply_refresh_failure(account, error)), @@ -388,7 +413,6 @@ pub fn apply_refresh_failure(account: &Account, failure: RefreshFailure) -> Acco RefreshFailure::InvalidGrant => AccountStatus::Expired, RefreshFailure::QuotaExhausted => AccountStatus::QuotaExhausted, RefreshFailure::Banned => AccountStatus::Banned, - RefreshFailure::Disabled => AccountStatus::Disabled, RefreshFailure::RetryableTransport => AccountStatus::Active, RefreshFailure::Transport => AccountStatus::Active, }; diff --git a/src/upstream/accounts/token_refresh/runtime.rs b/src/upstream/accounts/token_refresh/runtime.rs index eb35427d2..50f5c419b 100644 --- a/src/upstream/accounts/token_refresh/runtime.rs +++ b/src/upstream/accounts/token_refresh/runtime.rs @@ -1,8 +1,9 @@ //! 运行时 token refresh 业务服务。 use std::{ - collections::{HashMap, HashSet}, + collections::{hash_map::DefaultHasher, HashMap, HashSet}, future::Future, + hash::{Hash, Hasher}, pin::Pin, sync::Arc, time::Duration, @@ -22,8 +23,9 @@ use crate::upstream::accounts::{ StoredAccount, }, token_refresh::{ - jwt_expiration, RefreshError, RefreshLeaseStore, RefreshLeaseStoreError, RefreshPolicy, - RefreshScheduler, RefreshTrigger, RuntimeRefreshPolicy, TokenRefresher, + jittered_refresh_at, jwt_expiration, RefreshError, RefreshLeaseStore, + RefreshLeaseStoreError, RefreshPolicy, RefreshScheduler, RefreshTrigger, + RuntimeRefreshPolicy, TokenRefresher, }, }; @@ -33,6 +35,8 @@ const REFRESH_RETRY_BASE_DELAY_MILLIS: u64 = 5_000; const REFRESH_RETRY_MAX_DELAY_MILLIS: u64 = 300_000; const PERMANENT_FAILURE_CONFIRMATION_THRESHOLD: usize = 2; const RECOVERY_DELAY_SECONDS: i64 = 10 * 60; +const RETRY_DELAY_JITTER: f64 = 0.30; +const RECOVERY_DELAY_JITTER: f64 = 0.20; /// 运行时 token refresh 服务。 pub struct RuntimeTokenRefreshService @@ -226,6 +230,32 @@ where self.timers.lock().await.len() } + /// 立即刷新被上游认证失败命中的账号。 + pub async fn trigger_account_refresh_now( + &self, + account_id: &str, + ) -> TokenRefreshServiceResult<()> { + let Some(account) = self.store.get(account_id).await? else { + return Err(TokenRefreshServiceError::AccountNotFound( + account_id.to_string(), + )); + }; + if account.refresh_token.is_none() + || matches!( + account.status, + AccountStatus::Disabled | AccountStatus::Banned + ) + { + return Ok(()); + } + + self.clear_scheduled_timer(account_id).await; + persist_next_refresh_at(&self.store, account_id, None).await?; + self.refresh_scheduled_account(account_id, RefreshTrigger::Unauthorized, Utc::now()) + .await + .map(|_| ()) + } + /// 在指定时间点执行一次到期刷新扫描。 pub async fn refresh_due_accounts_once_at( &self, @@ -310,8 +340,11 @@ where let expires_at = account .access_token_expires_at .or_else(|| jwt_expiration(&account.access_token))?; - let margin_seconds = self.policy.refresh_margin_seconds().min(i64::MAX as u64) as i64; - Some(expires_at - chrono::Duration::seconds(margin_seconds)) + Some(jittered_refresh_at( + &account.id, + expires_at, + self.policy.refresh_margin_seconds(), + )) } async fn persist_next_refresh_at_if_changed( @@ -521,7 +554,7 @@ where if is_permanent_refresh_status(updated.status) { permanent_failures += 1; if permanent_failures < PERMANENT_FAILURE_CONFIRMATION_THRESHOLD - && self.sleep_before_retry(attempt_index).await + && self.sleep_before_retry(&attempt.id, attempt_index).await { continue; } @@ -532,11 +565,17 @@ where return Ok(TokenRefreshOutcome::StatusUpdated); } Ok(_) => return Ok(TokenRefreshOutcome::Skipped), - Err(RefreshError::RetryableTransport) - if self.sleep_before_retry(attempt_index).await => {} + Err(RefreshError::RetryableTransport | RefreshError::Transport) + if self.sleep_before_retry(&attempt.id, attempt_index).await => {} Err(RefreshError::RetryableTransport | RefreshError::Transport) => { persist_status(&self.store, &attempt.id, AccountStatus::Active).await?; - let next_refresh_at = now + chrono::Duration::seconds(RECOVERY_DELAY_SECONDS); + let recovery_delay = stable_jittered_duration( + &attempt.id, + Duration::from_secs(RECOVERY_DELAY_SECONDS as u64), + RECOVERY_DELAY_JITTER, + "recovery", + ); + let next_refresh_at = scheduled_at_from_delay(now, recovery_delay); persist_next_refresh_at(&self.store, &attempt.id, Some(next_refresh_at)) .await?; return Ok(TokenRefreshOutcome::Failed); @@ -576,11 +615,17 @@ where Ok(Some(attempt)) } - async fn sleep_before_retry(&self, attempt_index: usize) -> bool { + async fn sleep_before_retry(&self, account_id: &str, attempt_index: usize) -> bool { let Some(delay) = self.retry_delays.get(attempt_index).copied() else { return false; }; + let delay = stable_jittered_duration( + account_id, + delay, + RETRY_DELAY_JITTER, + &format!("retry:{attempt_index}"), + ); if !delay.is_zero() { sleep(delay).await; } @@ -792,6 +837,27 @@ fn default_refresh_retry_delays() -> Vec { .collect() } +fn stable_jittered_duration( + account_id: &str, + base: Duration, + variance: f64, + salt: &str, +) -> Duration { + if base.is_zero() { + return Duration::ZERO; + } + + let mut hasher = DefaultHasher::new(); + account_id.hash(&mut hasher); + salt.hash(&mut hasher); + let unit = hasher.finish() as f64 / u64::MAX as f64; + let factor = (1.0 - variance) + unit * variance * 2.0; + let millis = (base.as_millis() as f64 * factor) + .round() + .clamp(0.0, u64::MAX as f64) as u64; + Duration::from_millis(millis) +} + fn scheduled_at_from_delay(now: DateTime, delay: Duration) -> DateTime { match chrono::Duration::from_std(delay) { Ok(delay) => now + delay, diff --git a/src/upstream/token_client.rs b/src/upstream/token_client.rs index 7bd488a62..0dd703607 100644 --- a/src/upstream/token_client.rs +++ b/src/upstream/token_client.rs @@ -77,25 +77,15 @@ fn parse_token_pair(body: &str) -> Result { }) } -fn classify_refresh_failure(status: StatusCode, body: &str) -> RefreshFailure { +fn classify_refresh_failure(_status: StatusCode, body: &str) -> RefreshFailure { let lower = body.to_ascii_lowercase(); - if lower.contains("quota") { - return RefreshFailure::QuotaExhausted; - } - if lower.contains("banned") || lower.contains("account has been deactivated") { + if lower.contains("account has been deactivated") || lower.contains("refresh_token_reused") { return RefreshFailure::Banned; } - if lower.contains("account disabled") { - return RefreshFailure::Disabled; - } - if status == StatusCode::BAD_REQUEST - || status == StatusCode::UNAUTHORIZED - || lower.contains("invalid_grant") + if lower.contains("invalid_grant") || lower.contains("invalid_token") || lower.contains("access_denied") || lower.contains("refresh_token_expired") - || lower.contains("refresh_token_reused") - || lower.contains("token_revoked") { return RefreshFailure::InvalidGrant; } diff --git a/src/upstream/transport/client.rs b/src/upstream/transport/client.rs index 699c5a8eb..0a44278de 100644 --- a/src/upstream/transport/client.rs +++ b/src/upstream/transport/client.rs @@ -146,10 +146,15 @@ pub fn is_banned_upstream_error(error: &CodexClientError) -> bool { matches!( error, CodexClientError::Upstream { status, body, .. } - if status.as_u16() == 403 && is_banned_auth_signal(body) + if status.as_u16() == 403 && !is_html_error_body(body) ) } +fn is_html_error_body(value: &str) -> bool { + let value = value.trim_start().to_ascii_lowercase(); + value.starts_with(" = Result; diff --git a/tests/admin/accounts/importing.rs b/tests/admin/accounts/importing.rs index d877db0e9..d17a97554 100644 --- a/tests/admin/accounts/importing.rs +++ b/tests/admin/accounts/importing.rs @@ -255,7 +255,8 @@ async fn admin_accounts_import_should_complete_chatgpt_account_id_from_refresh_t .await .unwrap(); let body = response_json(response).await; - let stored = SqliteAccountStore::new(pool) + let store = SqliteAccountStore::new(pool); + let stored = store .list_metadata_page(1, 10, None) .await .unwrap() @@ -285,6 +286,68 @@ async fn admin_accounts_import_should_complete_chatgpt_account_id_from_refresh_t ); } +#[tokio::test] +async fn admin_accounts_import_from_refresh_token_should_not_store_consumed_refresh_token_when_not_rotated( +) { + let server = MockServer::start().await; + let access_token = test_jwt( + "rt-import-no-rotation-account", + Some("rt-import-no-rotation-user"), + Some("rt-import-no-rotation@example.com"), + Some("plus"), + ); + Mock::given(method("POST")) + .and(path("/oauth/token")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "access_token": access_token, + }))) + .mount(&server) + .await; + let (app, _state, pool, _dir) = admin_accounts_test_app_with_oauth_token_endpoint( + "admin-accounts-import-rt-no-rotation.sqlite", + 126, + format!("{}/oauth/token", server.uri()), + ) + .await; + + let response = app + .oneshot( + Request::builder() + .method("POST") + .uri("/api/admin/accounts/import") + .header("content-type", "application/json") + .header("cookie", "cpr_admin_session=session_1") + .header("x-request-id", "req_accounts_import_rt_no_rotation") + .body(Body::from( + json!({ + "accounts": [{ "refreshToken": "rt-import-no-rotation-source" }] + }) + .to_string(), + )) + .unwrap(), + ) + .await + .unwrap(); + let body = response_json(response).await; + let store = SqliteAccountStore::new(pool); + let stored = store + .list_metadata_page(1, 10, None) + .await + .unwrap() + .items + .into_iter() + .next() + .unwrap(); + + assert_eq!(body["data"]["imported"], 1); + assert_eq!( + stored.account_id.as_deref(), + Some("rt-import-no-rotation-account") + ); + let account = store.get(&stored.id).await.unwrap().unwrap(); + assert!(account.refresh_token.is_none()); +} + #[tokio::test] async fn admin_accounts_import_should_fetch_usage_to_complete_missing_plan_and_quota() { let server = MockServer::start().await; diff --git a/tests/admin/accounts/lifecycle.rs b/tests/admin/accounts/lifecycle.rs index a69370565..7ec978a03 100644 --- a/tests/admin/accounts/lifecycle.rs +++ b/tests/admin/accounts/lifecycle.rs @@ -262,7 +262,7 @@ async fn admin_account_status_update_should_update_proxy_account_pool() { .header("content-type", "application/json") .header("cookie", "cpr_admin_session=session_1") .body(Body::from( - json!({"id": "acct_runtime_status", "status": "banned"}).to_string(), + json!({"id": "acct_runtime_status", "status": "disabled"}).to_string(), )) .unwrap(), ) @@ -282,7 +282,61 @@ async fn admin_account_status_update_should_update_proxy_account_pool() { } #[tokio::test] -async fn admin_account_refresh_should_not_mark_valid_account_banned_when_refresh_token_reused() { +async fn admin_account_status_update_should_reject_non_manual_statuses() { + let (app, _state, pool, _dir) = + admin_accounts_test_app("admin-account-status-reject-banned.sqlite", 125).await; + seed_account( + &pool, + NewAccount { + id: "acct_reject_status".to_string(), + email: Some("reject-status@example.com".to_string()), + account_id: Some("chatgpt-reject-status".to_string()), + user_id: Some("user-reject-status".to_string()), + label: None, + plan_type: None, + access_token: SecretString::new( + test_jwt( + "chatgpt-reject-status", + Some("user-reject-status"), + Some("reject-status@example.com"), + Some("free"), + ) + .into(), + ), + refresh_token: None, + access_token_expires_at: None, + status: AccountStatus::Active, + added_at: None, + }, + ) + .await; + + let response = app + .oneshot( + Request::builder() + .method("POST") + .uri("/api/admin/accounts/update") + .header("content-type", "application/json") + .header("cookie", "cpr_admin_session=session_1") + .body(Body::from( + json!({"id": "acct_reject_status", "status": "banned"}).to_string(), + )) + .unwrap(), + ) + .await + .unwrap(); + let stored = SqliteAccountStore::new(pool) + .get("acct_reject_status") + .await + .unwrap() + .unwrap(); + + assert_eq!(response.status(), StatusCode::BAD_REQUEST); + assert_eq!(stored.status, AccountStatus::Active); +} + +#[tokio::test] +async fn admin_account_refresh_should_mark_account_expired_when_refresh_token_reused() { let server = wiremock::MockServer::start().await; wiremock::Mock::given(wiremock::matchers::method("POST")) .and(wiremock::matchers::path("/oauth/token")) @@ -338,16 +392,355 @@ async fn admin_account_refresh_should_not_mark_valid_account_banned_when_refresh ) .await .unwrap(); + let status = response.status(); + let body = response_json(response).await; let stored = SqliteAccountStore::new(pool) .get("acct_refresh_rt_reused") .await .unwrap() .unwrap(); - assert_eq!(response.status(), StatusCode::BAD_REQUEST); + assert_eq!(status, StatusCode::OK); + assert_eq!(body["data"]["result"], "dead"); + assert_eq!(body["data"]["previousStatus"], "active"); + assert_eq!(stored.status, AccountStatus::Expired); +} + +#[tokio::test] +async fn admin_account_refresh_should_skip_disabled_account_without_consuming_refresh_token() { + let server = wiremock::MockServer::start().await; + wiremock::Mock::given(wiremock::matchers::method("POST")) + .and(wiremock::matchers::path("/oauth/token")) + .respond_with(wiremock::ResponseTemplate::new(200).set_body_json(json!({ + "access_token": test_jwt( + "chatgpt-disabled-refresh", + Some("user-disabled-refresh"), + Some("disabled-refresh@example.com"), + Some("plus"), + ), + "refresh_token": "refresh-disabled-rotated" + }))) + .expect(0) + .mount(&server) + .await; + let (app, _state, pool, _dir) = admin_accounts_test_app_with_oauth_token_endpoint( + "admin-account-refresh-disabled-skip.sqlite", + 124, + format!("{}/oauth/token", server.uri()), + ) + .await; + seed_account( + &pool, + NewAccount { + id: "acct_refresh_disabled".to_string(), + email: Some("disabled-refresh@example.com".to_string()), + account_id: Some("chatgpt-disabled-refresh".to_string()), + user_id: Some("user-disabled-refresh".to_string()), + label: None, + plan_type: None, + access_token: SecretString::new( + test_jwt( + "chatgpt-disabled-refresh", + Some("user-disabled-refresh"), + Some("disabled-refresh@example.com"), + Some("plus"), + ) + .into(), + ), + refresh_token: Some(SecretString::new("refresh-disabled".to_string().into())), + access_token_expires_at: None, + status: AccountStatus::Disabled, + added_at: None, + }, + ) + .await; + + let response = app + .oneshot( + Request::builder() + .method("POST") + .uri("/api/admin/accounts/refresh") + .header("content-type", "application/json") + .header("cookie", "cpr_admin_session=session_1") + .body(Body::from( + json!({"id": "acct_refresh_disabled"}).to_string(), + )) + .unwrap(), + ) + .await + .unwrap(); + let status = response.status(); + let body = response_json(response).await; + let stored = SqliteAccountStore::new(pool) + .get("acct_refresh_disabled") + .await + .unwrap() + .unwrap(); + + assert_eq!(status, StatusCode::OK); + assert_eq!(body["data"]["result"], "skipped"); + assert_eq!(body["data"]["error"], "manually disabled"); + assert_eq!(stored.status, AccountStatus::Disabled); + assert_eq!( + stored.refresh_token.unwrap().expose_secret(), + "refresh-disabled" + ); +} + +#[tokio::test] +async fn admin_account_refresh_should_skip_account_without_refresh_token() { + let (app, _state, pool, _dir) = + admin_accounts_test_app("admin-account-refresh-no-rt-skip.sqlite", 124).await; + seed_account( + &pool, + NewAccount { + id: "acct_refresh_no_rt".to_string(), + email: Some("no-rt-refresh@example.com".to_string()), + account_id: Some("chatgpt-no-rt-refresh".to_string()), + user_id: Some("user-no-rt-refresh".to_string()), + label: None, + plan_type: None, + access_token: SecretString::new( + test_jwt( + "chatgpt-no-rt-refresh", + Some("user-no-rt-refresh"), + Some("no-rt-refresh@example.com"), + Some("plus"), + ) + .into(), + ), + refresh_token: None, + access_token_expires_at: None, + status: AccountStatus::Active, + added_at: None, + }, + ) + .await; + + let response = app + .oneshot( + Request::builder() + .method("POST") + .uri("/api/admin/accounts/refresh") + .header("content-type", "application/json") + .header("cookie", "cpr_admin_session=session_1") + .body(Body::from(json!({"id": "acct_refresh_no_rt"}).to_string())) + .unwrap(), + ) + .await + .unwrap(); + let status = response.status(); + let body = response_json(response).await; + let stored = SqliteAccountStore::new(pool) + .get("acct_refresh_no_rt") + .await + .unwrap() + .unwrap(); + + assert_eq!(status, StatusCode::OK); + assert_eq!(body["data"]["result"], "skipped"); + assert_eq!(body["data"]["error"], "no refresh token"); assert_eq!(stored.status, AccountStatus::Active); } +#[tokio::test] +async fn admin_account_refresh_should_not_change_status_for_non_permanent_refresh_error() { + let server = wiremock::MockServer::start().await; + wiremock::Mock::given(wiremock::matchers::method("POST")) + .and(wiremock::matchers::path("/oauth/token")) + .respond_with(wiremock::ResponseTemplate::new(400).set_body_string("quota exceeded")) + .mount(&server) + .await; + let (app, _state, pool, _dir) = admin_accounts_test_app_with_oauth_token_endpoint( + "admin-account-refresh-quota-temporary.sqlite", + 124, + format!("{}/oauth/token", server.uri()), + ) + .await; + seed_account( + &pool, + NewAccount { + id: "acct_refresh_quota".to_string(), + email: Some("quota-refresh@example.com".to_string()), + account_id: Some("chatgpt-quota-refresh".to_string()), + user_id: Some("user-quota-refresh".to_string()), + label: None, + plan_type: None, + access_token: SecretString::new( + test_jwt( + "chatgpt-quota-refresh", + Some("user-quota-refresh"), + Some("quota-refresh@example.com"), + Some("plus"), + ) + .into(), + ), + refresh_token: Some(SecretString::new("refresh-quota".to_string().into())), + access_token_expires_at: None, + status: AccountStatus::Active, + added_at: None, + }, + ) + .await; + + let response = app + .oneshot( + Request::builder() + .method("POST") + .uri("/api/admin/accounts/refresh") + .header("content-type", "application/json") + .header("cookie", "cpr_admin_session=session_1") + .body(Body::from(json!({"id": "acct_refresh_quota"}).to_string())) + .unwrap(), + ) + .await + .unwrap(); + let status = response.status(); + let body = response_json(response).await; + let stored = SqliteAccountStore::new(pool) + .get("acct_refresh_quota") + .await + .unwrap() + .unwrap(); + + assert_eq!(status, StatusCode::OK); + assert_eq!(body["data"]["result"], "dead"); + assert_eq!(stored.status, AccountStatus::Active); +} + +#[tokio::test] +async fn admin_accounts_health_check_should_return_alive_dead_and_skipped_summary() { + let server = wiremock::MockServer::start().await; + wiremock::Mock::given(wiremock::matchers::method("POST")) + .and(wiremock::matchers::path("/oauth/token")) + .respond_with(wiremock::ResponseTemplate::new(200).set_body_json(json!({ + "access_token": test_jwt( + "chatgpt-health-alive", + Some("user-health-alive"), + Some("health-alive@example.com"), + Some("plus"), + ), + "refresh_token": "refresh-health-alive-rotated" + }))) + .expect(1) + .mount(&server) + .await; + let (app, _state, pool, _dir) = admin_accounts_test_app_with_oauth_token_endpoint( + "admin-accounts-health-check.sqlite", + 124, + format!("{}/oauth/token", server.uri()), + ) + .await; + seed_account( + &pool, + NewAccount { + id: "acct_health_alive".to_string(), + email: Some("health-alive@example.com".to_string()), + account_id: Some("chatgpt-health-alive".to_string()), + user_id: Some("user-health-alive".to_string()), + label: None, + plan_type: None, + access_token: SecretString::new( + test_jwt( + "chatgpt-health-alive", + Some("user-health-alive"), + Some("health-alive@example.com"), + Some("plus"), + ) + .into(), + ), + refresh_token: Some(SecretString::new("refresh-health-alive".to_string().into())), + access_token_expires_at: None, + status: AccountStatus::Active, + added_at: None, + }, + ) + .await; + seed_account( + &pool, + NewAccount { + id: "acct_health_no_rt".to_string(), + email: Some("health-no-rt@example.com".to_string()), + account_id: Some("chatgpt-health-no-rt".to_string()), + user_id: Some("user-health-no-rt".to_string()), + label: None, + plan_type: None, + access_token: SecretString::new( + test_jwt( + "chatgpt-health-no-rt", + Some("user-health-no-rt"), + Some("health-no-rt@example.com"), + Some("plus"), + ) + .into(), + ), + refresh_token: None, + access_token_expires_at: None, + status: AccountStatus::Active, + added_at: None, + }, + ) + .await; + seed_account( + &pool, + NewAccount { + id: "acct_health_disabled".to_string(), + email: Some("health-disabled@example.com".to_string()), + account_id: Some("chatgpt-health-disabled".to_string()), + user_id: Some("user-health-disabled".to_string()), + label: None, + plan_type: None, + access_token: SecretString::new( + test_jwt( + "chatgpt-health-disabled", + Some("user-health-disabled"), + Some("health-disabled@example.com"), + Some("plus"), + ) + .into(), + ), + refresh_token: Some(SecretString::new( + "refresh-health-disabled".to_string().into(), + )), + access_token_expires_at: None, + status: AccountStatus::Disabled, + added_at: None, + }, + ) + .await; + + let response = app + .oneshot( + Request::builder() + .method("POST") + .uri("/api/admin/accounts/health-check") + .header("content-type", "application/json") + .header("cookie", "cpr_admin_session=session_1") + .body(Body::from(json!({"stagger_ms": 0}).to_string())) + .unwrap(), + ) + .await + .unwrap(); + let status = response.status(); + let body = response_json(response).await; + let stored_alive = SqliteAccountStore::new(pool) + .get("acct_health_alive") + .await + .unwrap() + .unwrap(); + + assert_eq!(status, StatusCode::OK); + assert_eq!(body["data"]["summary"]["total"], 3); + assert_eq!(body["data"]["summary"]["alive"], 1); + assert_eq!(body["data"]["summary"]["dead"], 0); + assert_eq!(body["data"]["summary"]["skipped"], 2); + assert_eq!(stored_alive.status, AccountStatus::Active); + assert_eq!( + stored_alive.refresh_token.unwrap().expose_secret(), + "refresh-health-alive-rotated" + ); +} + #[tokio::test] async fn admin_account_create_should_derive_claims_and_store_plain_tokens() { let dir = tempfile::tempdir().unwrap(); @@ -438,6 +831,43 @@ async fn admin_account_manual_create_should_reject_missing_invalid_expired_or_un } } +#[tokio::test] +async fn admin_account_manual_create_from_refresh_token_should_not_store_consumed_refresh_token_when_not_rotated( +) { + let server = wiremock::MockServer::start().await; + let access_token = test_jwt( + "rt-create-account", + Some("rt-create-user"), + Some("rt-create@example.com"), + Some("plus"), + ); + wiremock::Mock::given(wiremock::matchers::method("POST")) + .and(wiremock::matchers::path("/oauth/token")) + .respond_with(wiremock::ResponseTemplate::new(200).set_body_json(json!({ + "access_token": access_token, + }))) + .mount(&server) + .await; + let (app, _state, pool, _dir) = admin_accounts_test_app_with_oauth_token_endpoint( + "admin-account-create-rt-no-rotation.sqlite", + 126, + format!("{}/oauth/token", server.uri()), + ) + .await; + + let response = post_admin_account(&app, json!({"refreshToken": "rt-create-source"})).await; + let status = response.status(); + let body = response_json(response).await; + let stored = SqliteAccountStore::new(pool) + .get(body["data"]["id"].as_str().unwrap()) + .await + .unwrap() + .unwrap(); + + assert_eq!(status, StatusCode::OK); + assert!(stored.refresh_token.is_none()); +} + #[tokio::test] async fn admin_account_manual_create_should_accept_current_openai_token_without_chatgpt_account_id() { diff --git a/tests/proxy/dispatch/chat_upstream/chat_routes.rs b/tests/proxy/dispatch/chat_upstream/chat_routes.rs index a72874a68..6bf6d0353 100644 --- a/tests/proxy/dispatch/chat_upstream/chat_routes.rs +++ b/tests/proxy/dispatch/chat_upstream/chat_routes.rs @@ -840,7 +840,7 @@ async fn chat_completions_should_mark_banned_after_403_and_fallback() { .and(path("/codex/responses")) .and(header("authorization", "Bearer access-primary")) .respond_with(ResponseTemplate::new(403).set_body_json(json!({ - "error": {"message": "account banned"} + "error": {"message": "request forbidden"} }))) .expect(1) .mount(&server) diff --git a/tests/proxy/dispatch/chat_upstream/responses_recovery.rs b/tests/proxy/dispatch/chat_upstream/responses_recovery.rs index a287e3c12..22010aa50 100644 --- a/tests/proxy/dispatch/chat_upstream/responses_recovery.rs +++ b/tests/proxy/dispatch/chat_upstream/responses_recovery.rs @@ -418,7 +418,7 @@ async fn responses_should_mark_banned_after_403_and_fallback() { .and(path("/codex/responses")) .and(header("authorization", "Bearer access-primary")) .respond_with(ResponseTemplate::new(403).set_body_json(json!({ - "error": {"message": "account banned"} + "error": {"message": "request forbidden"} }))) .expect(1) .mount(&server) @@ -1524,7 +1524,7 @@ async fn responses_stream_should_mark_banned_after_403_and_fallback() { .and(path("/codex/responses")) .and(header("authorization", "Bearer access-primary")) .respond_with(ResponseTemplate::new(403).set_body_json(json!({ - "error": {"message": "account banned"} + "error": {"message": "request forbidden"} }))) .expect(1) .mount(&server) diff --git a/tests/proxy/dispatch/chat_upstream/responses_websocket.rs b/tests/proxy/dispatch/chat_upstream/responses_websocket.rs index e80985201..dbe6c56c4 100644 --- a/tests/proxy/dispatch/chat_upstream/responses_websocket.rs +++ b/tests/proxy/dispatch/chat_upstream/responses_websocket.rs @@ -693,7 +693,7 @@ async fn responses_websocket_should_implicitly_resume_full_history_with_reasonin } #[tokio::test] -async fn responses_websocket_should_strip_implicit_banned_affinity_when_switching_accounts() { +async fn responses_websocket_should_strip_implicit_disabled_affinity_when_switching_accounts() { let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); let base_url = format!("http://{}", listener.local_addr().unwrap()); let upstream = tokio::spawn(async move { @@ -746,7 +746,7 @@ async fn responses_websocket_should_strip_implicit_banned_affinity_when_switchin .await .contains("\"id\":\"resp_implicit_ban_first\"")); - update_admin_account_status(&app, "acct_primary", "banned").await; + update_admin_account_status(&app, "acct_primary", "disabled").await; let second_response = app .oneshot(responses_json_request( diff --git a/tests/upstream/accounts/account_repository.rs b/tests/upstream/accounts/account_repository.rs index a6bddf62f..f388f0cc6 100644 --- a/tests/upstream/accounts/account_repository.rs +++ b/tests/upstream/accounts/account_repository.rs @@ -3,7 +3,7 @@ use codex_proxy_rs::{ infra::database::connect_sqlite, upstream::accounts::{ model::{AccountStatus, AccountUsageDelta}, - store::{AccountStore, NewAccount, SqliteAccountStore}, + store::{AccountClaimsUpdate, AccountStore, NewAccount, SqliteAccountStore}, }, }; use secrecy::{ExposeSecret, SecretString}; @@ -155,6 +155,64 @@ async fn account_repository_should_update_status_and_label_without_rewriting_tok assert!(updated_at.0.parse::>().is_ok()); } +#[tokio::test] +async fn account_repository_should_not_reactivate_disabled_or_banned_accounts_from_claims_update() { + let (pool, _dir) = sqlite_account_store_parts("accounts.sqlite", 11).await; + let repo = SqliteAccountStore::new(pool); + seed_new_account(&repo, "acct_disabled").await; + seed_new_account(&repo, "acct_banned").await; + repo.set_status("acct_disabled", AccountStatus::Disabled) + .await + .expect("disabled status should persist"); + repo.set_status("acct_banned", AccountStatus::Banned) + .await + .expect("banned status should persist"); + + for account_id in ["acct_disabled", "acct_banned"] { + assert!(repo + .update_from_claims( + account_id, + AccountClaimsUpdate { + email: Some(format!("{account_id}@example.com")), + account_id: Some(format!("chatgpt-{account_id}")), + user_id: Some(format!("user-{account_id}")), + plan_type: Some("plus".to_string()), + access_token: SecretString::new(format!("new-access-{account_id}").into()), + refresh_token: Some(SecretString::new( + format!("new-refresh-{account_id}").into(), + )), + access_token_expires_at: Some(Utc::now() + Duration::hours(1)), + next_refresh_at: Some(Utc::now() + Duration::minutes(30)), + status: AccountStatus::Active, + }, + ) + .await + .expect("claims update should persist")); + } + + let disabled = repo + .get("acct_disabled") + .await + .expect("account should load") + .expect("account should exist"); + let banned = repo + .get("acct_banned") + .await + .expect("account should load") + .expect("account should exist"); + + assert_eq!(disabled.status, AccountStatus::Disabled); + assert_eq!( + disabled.access_token.expose_secret(), + "new-access-acct_disabled" + ); + assert_eq!(banned.status, AccountStatus::Banned); + assert_eq!( + banned.access_token.expose_secret(), + "new-access-acct_banned" + ); +} + #[tokio::test] async fn account_repository_should_accumulate_usage_counters() { let (pool, _dir) = sqlite_account_store_parts("accounts.sqlite", 7).await; diff --git a/tests/upstream/accounts/token_refresh/failures.rs b/tests/upstream/accounts/token_refresh/failures.rs index 4edca0b16..4a9bedcb0 100644 --- a/tests/upstream/accounts/token_refresh/failures.rs +++ b/tests/upstream/accounts/token_refresh/failures.rs @@ -129,7 +129,7 @@ async fn token_refresh_task_should_recover_refreshing_account_after_restart() { } #[tokio::test] -async fn token_refresh_task_should_not_retry_ambiguous_transport_failure() { +async fn token_refresh_task_should_retry_transport_failure_before_recovery() { let dir = tempfile::tempdir().expect("temp dir"); let db = dir.path().join("token-refresh-transport-failure.sqlite"); let pool = connect_sqlite(&format!("sqlite://{}", db.display())) @@ -183,7 +183,7 @@ async fn token_refresh_task_should_not_retry_ambiguous_transport_failure() { let observed_statuses = observed_statuses.lock().await.clone(); assert_eq!(summary.failed, 1); - assert_eq!(observed_statuses, [AccountStatus::Refreshing]); + assert_eq!(observed_statuses, [AccountStatus::Refreshing; 5]); assert_eq!(stored.status, AccountStatus::Active); assert!(stored.next_refresh_at.is_some_and(|next| next > now)); assert_eq!( @@ -248,12 +248,19 @@ async fn token_refresh_task_should_delay_recovery_after_retry_exhaustion() { .refresh_due_accounts_once_at(now) .await .expect("first scan should summarize retry exhaustion"); + let recovery_at = store + .get("acct-refresh-delayed-recovery") + .await + .expect("account should load") + .expect("account should exist") + .next_refresh_at + .expect("retry exhaustion should persist recovery time"); let delayed = task .refresh_due_accounts_once_at(now + Duration::minutes(5)) .await .expect("recovery window should skip refresh"); let refreshed = task - .refresh_due_accounts_once_at(now + Duration::minutes(10) + Duration::seconds(1)) + .refresh_due_accounts_once_at(recovery_at + Duration::seconds(1)) .await .expect("recovery window should allow refresh"); let stored = store @@ -264,6 +271,8 @@ async fn token_refresh_task_should_delay_recovery_after_retry_exhaustion() { let observed_statuses = observed_statuses.lock().await.clone(); assert_eq!(failed.failed, 1); + assert!(recovery_at >= now + Duration::minutes(8)); + assert!(recovery_at <= now + Duration::minutes(12)); assert_eq!(delayed.skipped, 1); assert_eq!(refreshed.refreshed, 1); assert_eq!(observed_statuses, [AccountStatus::Refreshing; 6]); diff --git a/tests/upstream/accounts/token_refresh/scheduling.rs b/tests/upstream/accounts/token_refresh/scheduling.rs index cae29dcb8..6097ffdf4 100644 --- a/tests/upstream/accounts/token_refresh/scheduling.rs +++ b/tests/upstream/accounts/token_refresh/scheduling.rs @@ -183,7 +183,7 @@ async fn token_refresh_task_should_schedule_future_per_account_timer_without_ref .expect("account should be inserted"); let refresher = CountingTokenRefresher::default(); let task = codex_proxy_rs::upstream::accounts::token_refresh::RuntimeTokenRefreshService::new( - store, + store.clone(), RefreshPolicy { refresh_margin_seconds: 300, refresh_concurrency: 1, @@ -195,11 +195,150 @@ async fn token_refresh_task_should_schedule_future_per_account_timer_without_ref .schedule_account_timers_once_at(now) .await .expect("timer scheduling should succeed"); + let stored = store + .get("acct-refresh-future-timer") + .await + .expect("account should load") + .expect("account should exist"); tokio::task::yield_now().await; assert_eq!(summary.scheduled, 1); assert_eq!(task.scheduled_timer_count().await, 1); assert_eq!(refresher.calls.load(Ordering::SeqCst), 0); + let scheduled_at = stored + .next_refresh_at + .expect("future refresh timer should persist next_refresh_at"); + let exact_margin_at = now + Duration::minutes(5); + assert_ne!(scheduled_at, exact_margin_at); + assert!(scheduled_at >= now + Duration::seconds(255)); + assert!(scheduled_at <= now + Duration::seconds(345)); +} + +#[tokio::test] +async fn token_refresh_task_should_refresh_immediately_after_unauthorized_even_with_future_timer() { + let dir = tempfile::tempdir().expect("temp dir"); + let db = dir.path().join("token-refresh-unauthorized-now.sqlite"); + let pool = connect_sqlite(&format!("sqlite://{}", db.display())) + .await + .expect("sqlite pool"); + let store = SqliteAccountStore::new(pool); + let now = Utc::now(); + let old_access_token = test_jwt((now + Duration::minutes(10)).timestamp()); + let new_access_token = test_jwt((now + Duration::hours(1)).timestamp()); + store + .insert(NewAccount { + id: "acct-refresh-unauthorized-now".to_string(), + email: Some("unauthorized-now@example.com".to_string()), + account_id: Some("chatgpt-unauthorized-now".to_string()), + user_id: Some("user-unauthorized-now".to_string()), + label: None, + plan_type: Some("plus".to_string()), + access_token: SecretString::new(old_access_token.into()), + refresh_token: Some(SecretString::new( + "refresh-unauthorized-now".to_string().into(), + )), + added_at: None, + access_token_expires_at: Some(now + Duration::minutes(10)), + status: AccountStatus::Active, + }) + .await + .expect("account should be inserted"); + let refresher = NotifyingTokenRefresher { + calls: Arc::new(AtomicUsize::new(0)), + started: Arc::new(Notify::new()), + response: 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.clone(), + ) + .with_retry_delays(vec![StdDuration::ZERO; 4]); + task.schedule_account_timers_once_at(now) + .await + .expect("future timer should be scheduled"); + assert_eq!(task.scheduled_timer_count().await, 1); + assert!(store + .set_status("acct-refresh-unauthorized-now", AccountStatus::Expired) + .await + .expect("expired status should persist")); + + task.trigger_account_refresh_now("acct-refresh-unauthorized-now") + .await + .expect("unauthorized refresh should run immediately"); + let stored = store + .get("acct-refresh-unauthorized-now") + .await + .expect("account should load") + .expect("account should exist"); + + assert_eq!(refresher.calls.load(Ordering::SeqCst), 1); + assert_eq!(stored.status, AccountStatus::Active); + assert_eq!( + stored.access_token.expose_secret(), + new_access_token.as_str() + ); + assert!(stored.next_refresh_at.is_some()); + assert_eq!(task.scheduled_timer_count().await, 1); +} + +#[tokio::test] +async fn token_refresh_task_should_skip_unauthorized_refresh_without_refresh_token() { + let dir = tempfile::tempdir().expect("temp dir"); + let db = dir.path().join("token-refresh-unauthorized-no-rt.sqlite"); + let pool = connect_sqlite(&format!("sqlite://{}", db.display())) + .await + .expect("sqlite pool"); + let store = SqliteAccountStore::new(pool); + let now = Utc::now(); + let old_access_token = test_jwt((now + Duration::minutes(10)).timestamp()); + store + .insert(NewAccount { + id: "acct-refresh-unauthorized-no-rt".to_string(), + email: Some("unauthorized-no-rt@example.com".to_string()), + account_id: Some("chatgpt-unauthorized-no-rt".to_string()), + user_id: Some("user-unauthorized-no-rt".to_string()), + label: None, + plan_type: Some("plus".to_string()), + access_token: SecretString::new(old_access_token.clone().into()), + refresh_token: None, + added_at: None, + access_token_expires_at: Some(now + Duration::minutes(10)), + status: AccountStatus::Expired, + }) + .await + .expect("account should be inserted"); + let refresher = CountingTokenRefresher::default(); + let task = codex_proxy_rs::upstream::accounts::token_refresh::RuntimeTokenRefreshService::new( + store.clone(), + RefreshPolicy { + refresh_margin_seconds: 300, + refresh_concurrency: 1, + }, + refresher.clone(), + ); + + task.trigger_account_refresh_now("acct-refresh-unauthorized-no-rt") + .await + .expect("unauthorized refresh should no-op without refresh token"); + let stored = store + .get("acct-refresh-unauthorized-no-rt") + .await + .expect("account should load") + .expect("account should exist"); + + assert_eq!(refresher.calls.load(Ordering::SeqCst), 0); + assert_eq!(stored.status, AccountStatus::Expired); + assert_eq!( + stored.access_token.expose_secret(), + old_access_token.as_str() + ); } #[tokio::test] diff --git a/tests/upstream/token_client.rs b/tests/upstream/token_client.rs index fc3d41671..4026c4aa0 100644 --- a/tests/upstream/token_client.rs +++ b/tests/upstream/token_client.rs @@ -33,7 +33,7 @@ async fn openai_token_client_should_exchange_refresh_token_with_form_body() { } #[tokio::test] -async fn openai_token_client_should_not_treat_refresh_token_reuse_as_banned() { +async fn openai_token_client_should_treat_refresh_token_reuse_as_banned() { let server = MockServer::start().await; Mock::given(method("POST")) .and(path("/oauth/token")) @@ -47,17 +47,32 @@ async fn openai_token_client_should_not_treat_refresh_token_reuse_as_banned() { let failure = client.refresh("refresh-secret").await.unwrap_err(); - assert_eq!(failure, RefreshFailure::InvalidGrant); + assert_eq!(failure, RefreshFailure::Banned); } #[tokio::test] -async fn openai_token_client_should_treat_explicit_banned_signal_as_banned() { +async fn openai_token_client_should_not_treat_generic_banned_signal_as_banned() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/oauth/token")) + .respond_with(ResponseTemplate::new(403).set_body_string("account is banned")) + .mount(&server) + .await; + let client = test_token_client(&server); + + let failure = client.refresh("refresh-secret").await.unwrap_err(); + + assert_eq!(failure, RefreshFailure::Transport); +} + +#[tokio::test] +async fn openai_token_client_should_treat_deactivated_refresh_error_as_banned() { let server = MockServer::start().await; Mock::given(method("POST")) .and(path("/oauth/token")) .respond_with(ResponseTemplate::new(403).set_body_json(serde_json::json!({ "error": "access_denied", - "error_description": "account is banned" + "error_description": "account has been deactivated" }))) .mount(&server) .await; @@ -68,6 +83,51 @@ async fn openai_token_client_should_treat_explicit_banned_signal_as_banned() { assert_eq!(failure, RefreshFailure::Banned); } +#[tokio::test] +async fn openai_token_client_should_retry_non_permanent_http_refresh_errors() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/oauth/token")) + .respond_with(ResponseTemplate::new(400).set_body_string("account disabled")) + .mount(&server) + .await; + let client = test_token_client(&server); + + let failure = client.refresh("refresh-secret").await.unwrap_err(); + + assert_eq!(failure, RefreshFailure::Transport); +} + +#[tokio::test] +async fn openai_token_client_should_treat_quota_refresh_error_as_temporary() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/oauth/token")) + .respond_with(ResponseTemplate::new(400).set_body_string("quota exceeded")) + .mount(&server) + .await; + let client = test_token_client(&server); + + let failure = client.refresh("refresh-secret").await.unwrap_err(); + + assert_eq!(failure, RefreshFailure::Transport); +} + +#[tokio::test] +async fn openai_token_client_should_treat_token_revoked_without_invalid_grant_as_temporary() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/oauth/token")) + .respond_with(ResponseTemplate::new(400).set_body_string("token_revoked")) + .mount(&server) + .await; + let client = test_token_client(&server); + + let failure = client.refresh("refresh-secret").await.unwrap_err(); + + assert_eq!(failure, RefreshFailure::Transport); +} + fn test_token_client(server: &MockServer) -> OpenAiTokenClient { OpenAiTokenClient::new( reqwest::Client::builder().no_proxy().build().unwrap(), diff --git a/web/src/views/login/index.vue b/web/src/views/login/index.vue index b04ab47ba..63ba24266 100644 --- a/web/src/views/login/index.vue +++ b/web/src/views/login/index.vue @@ -1,8 +1,9 @@