From a7f756819d5d74c8eea18979f1acb33ef2c795e5 Mon Sep 17 00:00:00 2001 From: Eric Traut Date: Tue, 10 Mar 2026 15:25:40 -0600 Subject: [PATCH] Mitigate token refresh storms --- .../app-server/src/codex_message_processor.rs | 38 +++++--- codex-rs/app-server/tests/suite/auth.rs | 90 ++++++++++++++++++ codex-rs/core/src/auth.rs | 91 ++++++++++++++++++- codex-rs/core/tests/suite/auth_refresh.rs | 66 ++++++++++++++ 4 files changed, 269 insertions(+), 16 deletions(-) diff --git a/codex-rs/app-server/src/codex_message_processor.rs b/codex-rs/app-server/src/codex_message_processor.rs index c269fc73de..fcd8c0eb44 100644 --- a/codex-rs/app-server/src/codex_message_processor.rs +++ b/codex-rs/app-server/src/codex_message_processor.rs @@ -1325,20 +1325,25 @@ impl CodexMessageProcessor { } } - async fn refresh_token_if_requested(&self, do_refresh: bool) { + async fn refresh_token_if_requested(&self, do_refresh: bool) -> bool { if self.auth_manager.is_external_auth_active() { - return; + return false; } if do_refresh && let Err(err) = self.auth_manager.refresh_token().await { - tracing::warn!("failed to refresh token while getting account: {err}"); + let failed_reason = err.failed_reason(); + if failed_reason.is_none() { + tracing::warn!("failed to refresh token while getting account: {err}"); + } + return failed_reason.is_some(); } + false } async fn get_auth_status(&self, request_id: ConnectionRequestId, params: GetAuthStatusParams) { let include_token = params.include_token.unwrap_or(false); let do_refresh = params.refresh_token.unwrap_or(false); - self.refresh_token_if_requested(do_refresh).await; + let refresh_failed_permanently = self.refresh_token_if_requested(do_refresh).await; // Determine whether auth is required based on the active model provider. // If a custom provider is configured with `requires_openai_auth == false`, @@ -1352,18 +1357,25 @@ impl CodexMessageProcessor { requires_openai_auth: Some(false), } } else { + let refresh_failure = self.auth_manager.refresh_failure(); match self.auth_manager.auth().await { Some(auth) => { let auth_mode = auth.api_auth_mode(); - let (reported_auth_method, token_opt) = match auth.get_token() { - Ok(token) if !token.is_empty() => { - let tok = if include_token { Some(token) } else { None }; - (Some(auth_mode), tok) - } - Ok(_) => (None, None), - Err(err) => { - tracing::warn!("failed to get token for auth status: {err}"); - (None, None) + let (reported_auth_method, token_opt) = if include_token + && (refresh_failure.is_some() || refresh_failed_permanently) + { + (Some(auth_mode), None) + } else { + match auth.get_token() { + Ok(token) if !token.is_empty() => { + let tok = if include_token { Some(token) } else { None }; + (Some(auth_mode), tok) + } + Ok(_) => (None, None), + Err(err) => { + tracing::warn!("failed to get token for auth status: {err}"); + (None, None) + } } }; GetAuthStatusResponse { diff --git a/codex-rs/app-server/tests/suite/auth.rs b/codex-rs/app-server/tests/suite/auth.rs index 68d0bcd95d..0366834401 100644 --- a/codex-rs/app-server/tests/suite/auth.rs +++ b/codex-rs/app-server/tests/suite/auth.rs @@ -1,6 +1,8 @@ use anyhow::Result; +use app_test_support::ChatGptAuthFixture; use app_test_support::McpProcess; use app_test_support::to_response; +use app_test_support::write_chatgpt_auth; use codex_app_server_protocol::AuthMode; use codex_app_server_protocol::GetAuthStatusParams; use codex_app_server_protocol::GetAuthStatusResponse; @@ -8,10 +10,17 @@ use codex_app_server_protocol::JSONRPCError; use codex_app_server_protocol::JSONRPCResponse; use codex_app_server_protocol::LoginAccountResponse; use codex_app_server_protocol::RequestId; +use codex_core::auth::AuthCredentialsStoreMode; +use codex_core::auth::REFRESH_TOKEN_URL_OVERRIDE_ENV_VAR; use pretty_assertions::assert_eq; use std::path::Path; use tempfile::TempDir; use tokio::time::timeout; +use wiremock::Mock; +use wiremock::MockServer; +use wiremock::ResponseTemplate; +use wiremock::matchers::method; +use wiremock::matchers::path; const DEFAULT_READ_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(10); @@ -207,6 +216,87 @@ async fn get_auth_status_with_api_key_no_include_token() -> Result<()> { Ok(()) } +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn get_auth_status_omits_token_after_permanent_refresh_failure() -> Result<()> { + let codex_home = TempDir::new()?; + create_config_toml(codex_home.path())?; + write_chatgpt_auth( + codex_home.path(), + ChatGptAuthFixture::new("stale-access-token") + .refresh_token("stale-refresh-token") + .account_id("acct_123") + .email("user@example.com") + .plan_type("pro"), + AuthCredentialsStoreMode::File, + )?; + + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/oauth/token")) + .respond_with(ResponseTemplate::new(401).set_body_json(serde_json::json!({ + "error": { + "code": "refresh_token_reused" + } + }))) + .expect(1) + .mount(&server) + .await; + + let refresh_url = format!("{}/oauth/token", server.uri()); + let mut mcp = McpProcess::new_with_env( + codex_home.path(), + &[ + ("OPENAI_API_KEY", None), + ( + REFRESH_TOKEN_URL_OVERRIDE_ENV_VAR, + Some(refresh_url.as_str()), + ), + ], + ) + .await?; + timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; + + let request_id = mcp + .send_get_auth_status_request(GetAuthStatusParams { + include_token: Some(true), + refresh_token: Some(true), + }) + .await?; + + let resp: JSONRPCResponse = timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_response_message(RequestId::Integer(request_id)), + ) + .await??; + let status: GetAuthStatusResponse = to_response(resp)?; + assert_eq!( + status, + GetAuthStatusResponse { + auth_method: Some(AuthMode::Chatgpt), + auth_token: None, + requires_openai_auth: Some(true), + } + ); + + let second_request_id = mcp + .send_get_auth_status_request(GetAuthStatusParams { + include_token: Some(true), + refresh_token: Some(true), + }) + .await?; + + let second_resp: JSONRPCResponse = timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_response_message(RequestId::Integer(second_request_id)), + ) + .await??; + let second_status: GetAuthStatusResponse = to_response(second_resp)?; + assert_eq!(second_status, status); + + server.verify().await; + Ok(()) +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn login_api_key_rejected_when_forced_chatgpt() -> Result<()> { let codex_home = TempDir::new()?; diff --git a/codex-rs/core/src/auth.rs b/codex-rs/core/src/auth.rs index 9f13cdf2b5..bb6888e180 100644 --- a/codex-rs/core/src/auth.rs +++ b/codex-rs/core/src/auth.rs @@ -815,6 +815,13 @@ struct CachedAuth { auth: Option, /// Callback used to refresh external auth by asking the parent app for new tokens. external_refresher: Option>, + poisoned_managed_auth: Option, +} + +#[derive(Clone, Debug)] +struct PoisonedManagedAuth { + auth_dot_json: AuthDotJson, + error: RefreshTokenFailedError, } impl Debug for CachedAuth { @@ -828,6 +835,13 @@ impl Debug for CachedAuth { "external_refresher", &self.external_refresher.as_ref().map(|_| "present"), ) + .field( + "poisoned_managed_auth", + &self + .poisoned_managed_auth + .as_ref() + .map(|poisoned| poisoned.error.reason), + ) .finish() } } @@ -998,6 +1012,7 @@ impl AuthManager { inner: RwLock::new(CachedAuth { auth: managed_auth, external_refresher: None, + poisoned_managed_auth: None, }), enable_codex_api_key_env, auth_credentials_store_mode, @@ -1010,6 +1025,7 @@ impl AuthManager { let cached = CachedAuth { auth: Some(auth), external_refresher: None, + poisoned_managed_auth: None, }; Arc::new(Self { @@ -1029,6 +1045,7 @@ impl AuthManager { let cached = CachedAuth { auth: Some(auth), external_refresher: None, + poisoned_managed_auth: None, }; Arc::new(Self { codex_home, @@ -1044,6 +1061,11 @@ impl AuthManager { self.inner.read().ok().and_then(|c| c.auth.clone()) } + pub fn refresh_failure(&self) -> Option { + let auth = self.auth_cached()?; + self.refresh_failure_for_auth(&auth) + } + /// Current cached auth (clone). May be `None` if not logged in or load failed. /// Refreshes cached ChatGPT tokens if they are stale before returning. pub async fn auth(&self) -> Option { @@ -1118,6 +1140,40 @@ impl AuthManager { } } + fn refresh_failure_for_auth(&self, auth: &CodexAuth) -> Option { + let auth_dot_json = auth.get_current_auth_json()?; + self.inner + .read() + .ok() + .and_then(|cached| cached.poisoned_managed_auth.clone()) + .filter(|poisoned| poisoned.auth_dot_json == auth_dot_json) + .map(|poisoned| poisoned.error) + } + + fn poison_managed_auth_if_unchanged( + &self, + attempted_auth: &CodexAuth, + error: &RefreshTokenFailedError, + ) { + let Some(attempted_auth_dot_json) = attempted_auth.get_current_auth_json() else { + return; + }; + + if let Ok(mut guard) = self.inner.write() { + let current_auth_matches = guard + .auth + .as_ref() + .and_then(CodexAuth::get_current_auth_json) + .is_some_and(|current| current == attempted_auth_dot_json); + if current_auth_matches { + guard.poisoned_managed_auth = Some(PoisonedManagedAuth { + auth_dot_json: attempted_auth_dot_json, + error: error.clone(), + }); + } + } + } + fn load_auth_from_storage(&self) -> Option { load_auth( &self.codex_home, @@ -1132,6 +1188,19 @@ impl AuthManager { if let Ok(mut guard) = self.inner.write() { let previous = guard.auth.as_ref(); let changed = !AuthManager::auths_equal(previous, new_auth.as_ref()); + let poisoned_auth_still_matches = guard + .poisoned_managed_auth + .as_ref() + .and_then(|poisoned| { + new_auth + .as_ref() + .and_then(CodexAuth::get_current_auth_json) + .map(|current| current == poisoned.auth_dot_json) + }) + .unwrap_or(false); + if !poisoned_auth_still_matches { + guard.poisoned_managed_auth = None; + } tracing::info!("Reloaded auth, changed: {changed}"); guard.auth = new_auth; changed @@ -1197,6 +1266,11 @@ impl AuthManager { /// token is the same as the cached, then ask the token authority to refresh. pub async fn refresh_token(&self) -> Result<(), RefreshTokenError> { let auth_before_reload = self.auth_cached(); + if let Some(auth_before_reload) = auth_before_reload.as_ref() + && let Some(error) = self.refresh_failure_for_auth(auth_before_reload) + { + return Err(RefreshTokenError::Permanent(error)); + } let expected_account_id = auth_before_reload .as_ref() .and_then(CodexAuth::get_account_id); @@ -1227,7 +1301,12 @@ impl AuthManager { Some(auth) => auth, None => return Ok(()), }; - match auth { + if let Some(error) = self.refresh_failure_for_auth(&auth) { + return Err(RefreshTokenError::Permanent(error)); + } + + let attempted_auth = auth.clone(); + let result = match auth { CodexAuth::ChatgptAuthTokens(_) => { self.refresh_external_auth(ExternalAuthRefreshReason::Unauthorized) .await @@ -1239,11 +1318,14 @@ impl AuthManager { )) })?; self.refresh_and_persist_chatgpt_token(&chatgpt_auth, token_data.refresh_token) - .await?; - Ok(()) + .await } CodexAuth::ApiKey(_) => Ok(()), + }; + if let Err(RefreshTokenError::Permanent(error)) = &result { + self.poison_managed_auth_if_unchanged(&attempted_auth, error); } + result } /// Log out by deleting the on‑disk auth.json (if present). Returns Ok(true) @@ -1266,6 +1348,9 @@ impl AuthManager { } async fn refresh_if_stale(&self, auth: &CodexAuth) -> Result { + if let Some(error) = self.refresh_failure_for_auth(auth) { + return Err(RefreshTokenError::Permanent(error)); + } let chatgpt_auth = match auth { CodexAuth::Chatgpt(chatgpt_auth) => chatgpt_auth, _ => return Ok(false), diff --git a/codex-rs/core/tests/suite/auth_refresh.rs b/codex-rs/core/tests/suite/auth_refresh.rs index f5b13f0918..2e7fee8570 100644 --- a/codex-rs/core/tests/suite/auth_refresh.rs +++ b/codex-rs/core/tests/suite/auth_refresh.rs @@ -433,6 +433,72 @@ async fn refresh_token_returns_permanent_error_for_expired_refresh_token() -> Re Ok(()) } +#[serial_test::serial(auth_refresh)] +#[tokio::test] +async fn refresh_token_does_not_retry_after_permanent_failure() -> Result<()> { + skip_if_no_network!(Ok(())); + + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/oauth/token")) + .respond_with(ResponseTemplate::new(401).set_body_json(json!({ + "error": { + "code": "refresh_token_reused" + } + }))) + .expect(1) + .mount(&server) + .await; + + let ctx = RefreshTokenTestContext::new(&server)?; + let initial_last_refresh = Utc::now() - Duration::days(1); + let initial_tokens = build_tokens(INITIAL_ACCESS_TOKEN, INITIAL_REFRESH_TOKEN); + let initial_auth = AuthDotJson { + auth_mode: Some(AuthMode::Chatgpt), + openai_api_key: None, + tokens: Some(initial_tokens.clone()), + last_refresh: Some(initial_last_refresh), + }; + ctx.write_auth(&initial_auth)?; + + let first_err = ctx + .auth_manager + .refresh_token() + .await + .err() + .context("first refresh should fail")?; + assert_eq!( + first_err.failed_reason(), + Some(RefreshTokenFailedReason::Exhausted) + ); + + let second_err = ctx + .auth_manager + .refresh_token() + .await + .err() + .context("second refresh should fail without retrying")?; + assert_eq!( + second_err.failed_reason(), + Some(RefreshTokenFailedReason::Exhausted) + ); + + let stored = ctx.load_auth()?; + assert_eq!(stored, initial_auth); + let cached_auth = ctx + .auth_manager + .auth() + .await + .context("auth should remain cached")?; + let cached = cached_auth + .get_token_data() + .context("token data should remain cached")?; + assert_eq!(cached, initial_tokens); + + server.verify().await; + Ok(()) +} + #[serial_test::serial(auth_refresh)] #[tokio::test] async fn refresh_token_returns_transient_error_on_server_failure() -> Result<()> {