From 9b8c2bb50b4baee873ef1baa41a7136cc3e2c04a Mon Sep 17 00:00:00 2001 From: jif Date: Fri, 30 Oct 2026 14:42:33 +0000 Subject: [PATCH] Fix remote executor authentication timeout recovery (#58355) ## Why Registration can time out while resolving credentials, before any request is sent. Treating this as a transport timeout prevents a safe retry. Credential resolution timeouts can also incorrectly quarantine credentials when no OAuth refresh exchange has started. ## What changed - Classify registration deadlines reached before the first send as `RequestNotSent`, allowing the existing registration retry loop to recover. Preserve transport timeout handling once a request has been sent. - Track unsettled ChatGPT refreshes in `AuthManager` from the start of the OAuth exchange through credential persistence. Cancellation while waiting for the refresh guard or reading credentials does not mark them as unsettled. - Have the executor auth provider withhold credentials with an unsettled refresh. Preserve uncertainty after a cancelled exchange even if a later refresh fails, while allowing replacement credentials and successfully completed refreshes. ## Testing Add regression coverage for registration and connection deadlines across auth stages, retrying registration after an unsent credential timeout, recovery after slow replacement PAT metadata, and cancellation before and during OAuth refresh. GitOrigin-RevId: f4d7ec2db7f79b8d313f1ee52dbd441094856669 --- codex-rs/cli/tests/exec_server.rs | 25 +++- .../exec-server-cli/src/exec_server_auth.rs | 44 +++---- codex-rs/exec-server/src/remote.rs | 26 ++-- codex-rs/exec-server/src/remote/auth.rs | 101 +++++++++------- codex-rs/exec-server/src/remote/auth_tests.rs | 73 +++++++----- .../src/remote/registration_retry_tests.rs | 54 +++++---- codex-rs/login/src/auth/auth_tests.rs | 112 ++++++++++++++++++ codex-rs/login/src/auth/manager.rs | 108 +++++++++++------ 8 files changed, 375 insertions(+), 168 deletions(-) diff --git a/codex-rs/cli/tests/exec_server.rs b/codex-rs/cli/tests/exec_server.rs index 4dd347cd68..689fff9f1f 100644 --- a/codex-rs/cli/tests/exec_server.rs +++ b/codex-rs/cli/tests/exec_server.rs @@ -213,6 +213,7 @@ enum PatCredentialChange { Logout, Unchanged, UnavailableReplacement, + SlowReplacement, } async fn assert_pat_registry_reload(change: PatCredentialChange) -> Result<()> { @@ -271,9 +272,9 @@ async fn assert_pat_registry_reload(change: PatCredentialChange) -> Result<()> { PatCredentialChange::Logout | PatCredentialChange::Unchanged ) { let account = match change { - PatCredentialChange::Rotate | PatCredentialChange::UnavailableReplacement => { - "account-pat-test" - } + PatCredentialChange::Rotate + | PatCredentialChange::UnavailableReplacement + | PatCredentialChange::SlowReplacement => "account-pat-test", PatCredentialChange::SwitchAccount => "another-account", PatCredentialChange::Logout | PatCredentialChange::Unchanged => unreachable!(), }; @@ -297,6 +298,13 @@ async fn assert_pat_registry_reload(change: PatCredentialChange) -> Result<()> { "chatgpt_plan_type": "enterprise", "chatgpt_account_is_fedramp": false, })); + // Exceed registration's 10-second deadline while the detached auth + // resolution can still finish within its own 30-second budget. + let response = if matches!(change, PatCredentialChange::SlowReplacement) { + response.set_delay(Duration::from_secs(11)) + } else { + response + }; Mock::given(method("GET")) .and(path("/v1/user-auth-credential/whoami")) .and(header( @@ -312,7 +320,8 @@ async fn assert_pat_registry_reload(change: PatCredentialChange) -> Result<()> { let expected_registry_requests = match change { PatCredentialChange::Rotate | PatCredentialChange::Unchanged - | PatCredentialChange::UnavailableReplacement => 2, + | PatCredentialChange::UnavailableReplacement + | PatCredentialChange::SlowReplacement => 2, PatCredentialChange::SwitchAccount | PatCredentialChange::Logout => 1, }; let second_bearer = if matches!(change, PatCredentialChange::Unchanged) { @@ -331,7 +340,8 @@ async fn assert_pat_registry_reload(change: PatCredentialChange) -> Result<()> { let result = match change { PatCredentialChange::Rotate | PatCredentialChange::SwitchAccount - | PatCredentialChange::UnavailableReplacement => std::fs::write( + | PatCredentialChange::UnavailableReplacement + | PatCredentialChange::SlowReplacement => std::fs::write( &auth_path, r#"{"personal_access_token":"at-synthetic-replacement-pat"}"#, ), @@ -436,6 +446,11 @@ async fn remote_exec_server_pat_recovers_after_replacement_metadata_outage() -> assert_pat_registry_reload(PatCredentialChange::UnavailableReplacement).await } +#[tokio::test] +async fn remote_exec_server_pat_recovers_after_replacement_metadata_timeout() -> Result<()> { + assert_pat_registry_reload(PatCredentialChange::SlowReplacement).await +} + #[tokio::test] async fn remote_exec_server_gracefully_stops_when_parent_stdin_closes() -> Result<()> { const ENVIRONMENT_ID: &str = "environment-parent-lifetime"; diff --git a/codex-rs/exec-server-cli/src/exec_server_auth.rs b/codex-rs/exec-server-cli/src/exec_server_auth.rs index 75beb1d6d5..8fea10fc64 100644 --- a/codex-rs/exec-server-cli/src/exec_server_auth.rs +++ b/codex-rs/exec-server-cli/src/exec_server_auth.rs @@ -3,9 +3,9 @@ //! ChatGPT credentials may rotate while an executor is idle. Reload them for //! registry requests while retaining the identity that registered the executor. //! Failed credential reads block headers until revalidation, without counting as logout. +//! Quarantine begins with the OAuth exchange, after credential reads and refresh-guard waits. use std::sync::Arc; -use std::sync::Mutex; use std::sync::atomic::AtomicBool; use std::sync::atomic::Ordering; use std::time::Duration; @@ -55,7 +55,6 @@ pub(super) fn chatgpt_auth_provider( fedramp_account: expected_auth.is_fedramp_account(), unidentified_headers, resolution_lock: Arc::new(Semaphore::new(/*permits*/ 1)), - uncertain_refresh: Arc::new(Mutex::new(None)), headers_valid: Arc::new(AtomicBool::new(true)), })) } @@ -70,7 +69,6 @@ struct ReloadingChatgptAuthProvider { fedramp_account: bool, unidentified_headers: Option, resolution_lock: Arc, - uncertain_refresh: Arc>>, headers_valid: Arc, } @@ -93,13 +91,13 @@ impl ReloadingChatgptAuthProvider { "executor credential has an empty bearer".to_string(), )); } - let uncertain = self.uncertain_refresh.lock().map_err(|_| { - AuthError::Transient("executor authentication is unavailable".to_string()) - })?; - if uncertain - .as_ref() - .is_some_and(|previous| auth.get_token_data().as_ref().ok() == Some(previous)) - { + let uncertain = self + .auth_manager + .is_chatgpt_refresh_unsettled(&auth) + .map_err(|_| { + AuthError::Transient("executor authentication is unavailable".to_string()) + })?; + if uncertain { return Err(AuthError::Transient( "executor refresh outcome is unknown; replace or renew stored credentials" .to_string(), @@ -171,7 +169,6 @@ impl ReloadingChatgptAuthProvider { // that generation until another login/refresh supplies different credentials. tokio::spawn(async move { let _guard = guard; - let mut attempted_refresh = None; let result = tokio::time::timeout(Duration::from_secs(30), async { provider .auth_manager @@ -182,12 +179,10 @@ impl ReloadingChatgptAuthProvider { "executor credential storage or hydration is unavailable".to_string(), ) })?; - let auth = provider.auth_manager.auth_cached(); - let current = provider.checked_headers(auth.clone())?; + let current = provider.checked_headers(provider.auth_manager.auth_cached())?; if provider.unidentified_headers.is_some() { return Ok(current); } - attempted_refresh = auth.as_ref().and_then(|auth| auth.get_token_data().ok()); if let Some(rejected) = rejected_headers { if !current .iter() @@ -198,8 +193,6 @@ impl ReloadingChatgptAuthProvider { } else { provider.auth_manager.auth().await; } - // A later storage timeout is not an interrupted OAuth exchange. - attempted_refresh = None; // auth() deliberately falls back to a snapshot on refresh errors. // Re-read the selected source before dispatch: logout, replacement, // or failed hydration must never authorize that fallback snapshot. @@ -217,19 +210,14 @@ impl ReloadingChatgptAuthProvider { .await; let result = match result { Ok(result) => result, - Err(_) => { - if let Some(attempted) = attempted_refresh - && let Ok(mut uncertain) = provider.uncertain_refresh.lock() - { - *uncertain = Some(attempted); - } - Err(AuthError::Transient( - "executor credential resolution timed out; retry after repairing credentials" - .to_string(), - )) - } + Err(_) => Err(AuthError::Transient( + "executor credential resolution timed out; retry after repairing credentials" + .to_string(), + )), }; - provider.headers_valid.store(result.is_ok(), Ordering::Release); + provider + .headers_valid + .store(result.is_ok(), Ordering::Release); result }) .await diff --git a/codex-rs/exec-server/src/remote.rs b/codex-rs/exec-server/src/remote.rs index 118cc81aa5..cf88006ff0 100644 --- a/codex-rs/exec-server/src/remote.rs +++ b/codex-rs/exec-server/src/remote.rs @@ -180,16 +180,12 @@ impl EnvironmentRegistryClient { security_profile: NOISE_RELAY_SECURITY_PROFILE.to_string(), executor_public_key: executor_public_key.clone(), }; - let (response, executor_credential) = timeout_at( - deadline, - self.send_authenticated(Request::new(Method::POST, url).with_json(&body)), - ) - .await - .unwrap_or_else(|_| { - Err(ExecServerError::EnvironmentRegistryRequest( - codex_http_client::RouteAwareRequestError::Timeout, - )) - })?; + let (response, executor_credential) = self + .send_authenticated( + Request::new(Method::POST, url).with_json(&body), + Some(deadline), + ) + .await?; let status = response.status(); // Read diagnostics within the same attempt budget, preserving a known error status. let response: EnvironmentRegistryRegistrationResponse = timeout_at(deadline, async { @@ -264,7 +260,10 @@ impl EnvironmentRegistryClient { let response: EnvironmentRegistryConnectResponse = timeout_at(tokio::time::Instant::now() + self.connect_timeout, async { let (response, _) = self - .send_authenticated(Request::new(Method::POST, url).with_json(&body)) + .send_authenticated( + Request::new(Method::POST, url).with_json(&body), + /*deadline*/ None, + ) .await?; self.parse_json_response(response).await }) @@ -421,7 +420,10 @@ impl HarnessKeyValidator for RegistryHarnessKeyValidator { })?, ); } - let (response, executor_credential) = self.client.send_authenticated(request).await?; + let (response, executor_credential) = self + .client + .send_authenticated(request, /*deadline*/ None) + .await?; let status = response.status(); if !status.is_success() { // The request contains the short-lived authorization. Do not include diff --git a/codex-rs/exec-server/src/remote/auth.rs b/codex-rs/exec-server/src/remote/auth.rs index 58bbd006af..e78de93991 100644 --- a/codex-rs/exec-server/src/remote/auth.rs +++ b/codex-rs/exec-server/src/remote/auth.rs @@ -1,10 +1,13 @@ //! Sends registry requests with current credentials and at most one retry after a 401. //! Transport failures and other HTTP errors are never replayed here. +//! A deadline before the first transport send remains a retryable unsent request. use codex_http_client::HttpResponse; use codex_http_client::Request; use http::HeaderValue; use http::StatusCode; +use tokio::time::Instant; +use tokio::time::timeout_at; use super::EnvironmentRegistryClient; use super::RegistryAuthMode; @@ -24,52 +27,70 @@ impl EnvironmentRegistryClient { pub(super) async fn send_authenticated( &self, request: Request, + deadline: Option, ) -> Result<(HttpResponse, Option), ExecServerError> { let request = request .into_prepared() .map_err(ExecServerError::EnvironmentRegistryConfig)?; - let mut retried = false; - loop { - let (authenticated, credential) = - if self.auth_mode == RegistryAuthMode::ExecutorCredential { - let (headers, credential) = self.resolve_auth().await?; - let mut authenticated = request.clone(); - authenticated.headers.extend(headers); - (authenticated, credential) - } else { - ( - self.auth_provider - .apply_auth(request.clone()) - .await - .map_err(registry_auth_error)?, - None, - ) - }; - let prepared = authenticated - .prepare_body_for_send() - .map_err(ExecServerError::EnvironmentRegistryConfig)?; - let rejected_headers = prepared.headers.clone(); - let mut outgoing = self - .http_for_auth() - .request(authenticated.method, authenticated.url) - .headers(prepared.headers) - .headers(current_trace_context_headers()) - .body(prepared.body.unwrap_or_default()); - if let Some(timeout) = authenticated.timeout { - outgoing = outgoing.timeout(timeout); + let mut sent = false; + let send = async { + let mut retried = false; + loop { + let (authenticated, credential) = + if self.auth_mode == RegistryAuthMode::ExecutorCredential { + let (headers, credential) = self.resolve_auth().await?; + let mut authenticated = request.clone(); + authenticated.headers.extend(headers); + (authenticated, credential) + } else { + ( + self.auth_provider + .apply_auth(request.clone()) + .await + .map_err(registry_auth_error)?, + None, + ) + }; + let prepared = authenticated + .prepare_body_for_send() + .map_err(ExecServerError::EnvironmentRegistryConfig)?; + let rejected_headers = prepared.headers.clone(); + let mut outgoing = self + .http_for_auth() + .request(authenticated.method, authenticated.url) + .headers(prepared.headers) + .headers(current_trace_context_headers()) + .body(prepared.body.unwrap_or_default()); + if let Some(timeout) = authenticated.timeout { + outgoing = outgoing.timeout(timeout); + } + sent = true; + let response = outgoing.send().await?; + if !retried + && response.status() == StatusCode::UNAUTHORIZED + && self + .auth_provider + .refresh_after_unauthorized(rejected_headers) + .await + { + retried = true; + continue; + } + return Ok((response, credential)); } - let response = outgoing.send().await?; - if !retried - && response.status() == StatusCode::UNAUTHORIZED - && self - .auth_provider - .refresh_after_unauthorized(rejected_headers) - .await - { - retried = true; - continue; + }; + if let Some(deadline) = deadline { + match timeout_at(deadline, send).await { + Ok(result) => result, + Err(_) if !sent => Err(ExecServerError::RequestNotSent( + "environment registry credential resolution timed out".to_string(), + )), + Err(_) => Err(ExecServerError::EnvironmentRegistryRequest( + codex_http_client::RouteAwareRequestError::Timeout, + )), } - return Ok((response, credential)); + } else { + send.await } } } diff --git a/codex-rs/exec-server/src/remote/auth_tests.rs b/codex-rs/exec-server/src/remote/auth_tests.rs index c8d0dd210c..d5d2750b51 100644 --- a/codex-rs/exec-server/src/remote/auth_tests.rs +++ b/codex-rs/exec-server/src/remote/auth_tests.rs @@ -84,32 +84,48 @@ impl AuthProvider for RotatingAuth { } #[tokio::test] -async fn connect_deadline_covers_initial_auth_refresh_and_retry() -> anyhow::Result<()> { - for stage in [AuthStage::Initial, AuthStage::Refresh, AuthStage::Retry] { - let server = MockServer::start().await; - Mock::given(wiremock::matchers::any()) - .respond_with(ResponseTemplate::new(401)) - .mount(&server) - .await; - let provider = Arc::new(RotatingAuth { - stall: Some(stage), - ..Default::default() - }); - let mut client = EnvironmentRegistryClient::new(server.uri(), provider.clone())?; - client.connect_timeout = Duration::from_secs(60); - let key = NoiseChannelIdentity::generate()?.public_key(); - let task = tokio::spawn(async move { client.connect_environment("env", key).await }); - provider.reached.notified().await; - // Advance only after real socket I/O reached the selected auth stage. - tokio::time::pause(); - tokio::time::advance(Duration::from_secs(61)).await; - let result = tokio::time::timeout(Duration::from_secs(1), task) - .await - .expect("connect must finish at its original deadline")?; - assert!( - matches!(result, Err(ExecServerError::EnvironmentRegistryRequest(error)) if error.is_timeout()) - ); - tokio::time::resume(); +async fn registry_deadline_covers_initial_auth_refresh_and_retry() -> anyhow::Result<()> { + for operation in ["connect", "register"] { + for stage in [AuthStage::Initial, AuthStage::Refresh, AuthStage::Retry] { + let server = MockServer::start().await; + Mock::given(wiremock::matchers::any()) + .respond_with(ResponseTemplate::new(401)) + .mount(&server) + .await; + let provider = Arc::new(RotatingAuth { + stall: Some(stage), + ..Default::default() + }); + let mut client = EnvironmentRegistryClient::new(server.uri(), provider.clone())?; + client.connect_timeout = Duration::from_secs(60); + let key = NoiseChannelIdentity::generate()?.public_key(); + let task = tokio::spawn(async move { + if operation == "register" { + client.register_environment("env", &key).await.map(|_| ()) + } else { + client.connect_environment("env", key).await.map(|_| ()) + } + }); + provider.reached.notified().await; + // Advance only after real socket I/O reached the selected auth stage. + tokio::time::pause(); + tokio::time::advance(Duration::from_secs(61)).await; + let result = tokio::time::timeout(Duration::from_secs(1), task) + .await + .expect("registry operation must finish at its original deadline")?; + if operation == "register" && stage == AuthStage::Initial { + assert!(matches!(result, Err(ExecServerError::RequestNotSent(_)))); + } else { + assert!( + matches!(result, Err(ExecServerError::EnvironmentRegistryRequest(error)) if error.is_timeout()) + ); + } + assert_eq!( + server.received_requests().await.unwrap().len(), + if stage == AuthStage::Initial { 0 } else { 1 } + ); + tokio::time::resume(); + } } Ok(()) } @@ -192,7 +208,10 @@ async fn only_401_retries_and_only_once() -> anyhow::Result<()> { let provider = Arc::new(RotatingAuth::default()); let client = EnvironmentRegistryClient::new(server.uri(), provider.clone())?; let (response, _) = client - .send_authenticated(Request::new(Method::POST, server.uri())) + .send_authenticated( + Request::new(Method::POST, server.uri()), + /*deadline*/ None, + ) .await?; assert_eq!(response.status().as_u16(), status); assert_eq!(provider.refreshes.load(Ordering::SeqCst), refreshes); diff --git a/codex-rs/exec-server/src/remote/registration_retry_tests.rs b/codex-rs/exec-server/src/remote/registration_retry_tests.rs index 6db2608f1c..d74bab6f4d 100644 --- a/codex-rs/exec-server/src/remote/registration_retry_tests.rs +++ b/codex-rs/exec-server/src/remote/registration_retry_tests.rs @@ -41,6 +41,7 @@ const SUCCESS_BODY: &str = r#"{"environment_id":"registration-retry-test","url": struct RegistrationAuthProvider { calls: AtomicUsize, fail_first: bool, + stall_first: bool, } impl AuthProvider for RegistrationAuthProvider { @@ -51,6 +52,9 @@ impl AuthProvider for RegistrationAuthProvider { fn apply_auth_headers(&self, headers: HeaderMap) -> codex_api::AuthHeadersFuture<'_> { if self.calls.fetch_add(1, Ordering::Relaxed) == 0 && self.fail_first { return Box::pin(async { + if self.stall_first { + std::future::pending::<()>().await; + } Err(codex_api::AuthError::Transient( "metadata unavailable".to_string(), )) @@ -62,30 +66,34 @@ impl AuthProvider for RegistrationAuthProvider { #[tokio::test] async fn transient_credential_failure_retries_before_sending_registration() -> anyhow::Result<()> { - let server = wiremock::MockServer::start().await; - wiremock::Mock::given(wiremock::matchers::any()) - .respond_with( - wiremock::ResponseTemplate::new(200).set_body_raw(SUCCESS_BODY, "application/json"), + for stall_first in [false, true] { + let server = wiremock::MockServer::start().await; + wiremock::Mock::given(wiremock::matchers::any()) + .respond_with( + wiremock::ResponseTemplate::new(200).set_body_raw(SUCCESS_BODY, "application/json"), + ) + .expect(1) + .mount(&server) + .await; + let provider = Arc::new(RegistrationAuthProvider { + fail_first: true, + stall_first, + ..Default::default() + }); + let mut client = EnvironmentRegistryClient::new(server.uri(), provider.clone())?; + client.connect_timeout = Duration::from_secs(1); + let response = timeout( + Duration::from_secs(10), + client.register_environment_with_retry( + ENVIRONMENT_ID, + &NoiseChannelIdentity::generate()?.public_key(), + ), ) - .expect(1) - .mount(&server) - .await; - let provider = Arc::new(RegistrationAuthProvider { - fail_first: true, - ..Default::default() - }); - let client = EnvironmentRegistryClient::new(server.uri(), provider.clone())?; - let response = timeout( - Duration::from_secs(10), - client.register_environment_with_retry( - ENVIRONMENT_ID, - &NoiseChannelIdentity::generate()?.public_key(), - ), - ) - .await??; - assert_eq!(response.environment_id, ENVIRONMENT_ID); - assert_eq!(provider.calls.load(Ordering::Relaxed), 2); - server.verify().await; + .await??; + assert_eq!(response.environment_id, ENVIRONMENT_ID); + assert_eq!(provider.calls.load(Ordering::Relaxed), 2); + server.verify().await; + } Ok(()) } diff --git a/codex-rs/login/src/auth/auth_tests.rs b/codex-rs/login/src/auth/auth_tests.rs index d1c9206275..bd811d6d56 100644 --- a/codex-rs/login/src/auth/auth_tests.rs +++ b/codex-rs/login/src/auth/auth_tests.rs @@ -59,6 +59,118 @@ fn header_auth_exposes_a_valid_chatgpt_account_id() { } } +#[tokio::test] +async fn cancelled_refresh_marks_only_started_exchanges_as_unsettled() -> anyhow::Result<()> { + let home = tempdir()?; + write_auth_file( + AuthFileParams { + openai_api_key: None, + chatgpt_plan_type: Some("pro".into()), + chatgpt_account_id: Some(WORKSPACE_ID_ALLOWED.into()), + }, + home.path(), + )?; + let server = MockServer::start().await; + let storage = FileAuthStorage::new(home.path().to_path_buf()); + let mut stored = storage.load()?.context("stored credentials")?; + stored + .tokens + .as_mut() + .context("stored tokens")? + .oauth_provenance = Some(AccountOAuthProvenance { + issuer: server.uri(), + client_id: "codex-client".into(), + }); + storage.save(&stored)?; + let manager = Arc::new( + AuthManager::new_from_auth_config( + build_config( + home.path(), + /*forced_login_method*/ None, + /*forced_chatgpt_workspace_id*/ None, + ) + .await, + /*enable_codex_api_key_env*/ false, + /*allowed_auth_modes*/ None, + ) + .await, + ); + let initial = manager.auth_cached().context("initial auth")?; + + // Poll directly while the guard is held, then cancel before any OAuth request. + let guard = manager.refresh_lock.acquire().await?; + let mut waiting = Box::pin(manager.refresh_token_if_unchanged(&initial)); + assert!( + waiting + .as_mut() + .poll(&mut std::task::Context::from_waker( + std::task::Waker::noop(), + )) + .is_pending() + ); + drop(waiting); + assert!(!manager.is_chatgpt_refresh_unsettled(&initial)?); + assert!(server.received_requests().await.unwrap().is_empty()); + drop(guard); + + let received = Arc::new(tokio::sync::Notify::new()); + let notify = Arc::clone(&received); + Mock::given(path("/oauth/token")) + .respond_with(move |_: &wiremock::Request| { + notify.notify_one(); + ResponseTemplate::new(200) + .set_delay(Duration::from_secs(60)) + .set_body_json(json!({"access_token": "late", "refresh_token": "late-refresh"})) + }) + .expect(1) + .mount(&server) + .await; + let refreshing = Arc::clone(&manager); + let task = tokio::spawn(async move { refreshing.refresh_token_from_authority().await }); + received.notified().await; + assert!(manager.is_chatgpt_refresh_unsettled(&initial)?); + task.abort(); + assert!(task.await.unwrap_err().is_cancelled()); + assert!(manager.is_chatgpt_refresh_unsettled(&initial)?); + + // A later known failure cannot resolve an earlier cancelled exchange. + server.reset().await; + Mock::given(path("/oauth/token")) + .respond_with(ResponseTemplate::new(500)) + .expect(1) + .mount(&server) + .await; + assert!(manager.refresh_token_from_authority().await.is_err()); + assert!(manager.is_chatgpt_refresh_unsettled(&initial)?); + + // A replacement is usable, and a completed refresh settles its own generation. + server.reset().await; + Mock::given(path("/oauth/token")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "access_token": "completed", "refresh_token": "completed-refresh", + }))) + .expect(1) + .mount(&server) + .await; + let tokens = stored.tokens.as_mut().context("replacement tokens")?; + tokens.access_token = "replacement".into(); + tokens.refresh_token = "replacement-refresh".into(); + storage.save(&stored)?; + manager.reload().await; + let replacement = manager.auth_cached().context("replacement auth")?; + assert!(!manager.is_chatgpt_refresh_unsettled(&replacement)?); + manager.refresh_token_from_authority().await?; + assert!(!manager.is_chatgpt_refresh_unsettled(&replacement)?); + assert_eq!( + manager + .auth_cached() + .context("refreshed auth")? + .get_token()?, + "completed" + ); + Ok(()) +} + #[tokio::test] async fn refresh_without_id_token() -> anyhow::Result<()> { let codex_home = tempdir()?; diff --git a/codex-rs/login/src/auth/manager.rs b/codex-rs/login/src/auth/manager.rs index f561c2dd9b..9b51c70be1 100644 --- a/codex-rs/login/src/auth/manager.rs +++ b/codex-rs/login/src/auth/manager.rs @@ -2214,6 +2214,8 @@ pub struct AuthManager { chatgpt_base_url: Option, agent_identity_authapi_base_url: Option, refresh_lock: Arc, + // Writers hold refresh_lock; cancellation retains the attempted generation. + unsettled_chatgpt_refresh: Mutex>, agent_identity_lock: Semaphore, agent_identity_bootstrap_cooldown: Mutex, external_auth: RwLock>>, @@ -2377,6 +2379,7 @@ impl AuthManager { chatgpt_base_url, agent_identity_authapi_base_url, refresh_lock: Arc::new(Semaphore::new(/*permits*/ 1)), + unsettled_chatgpt_refresh: Mutex::new(None), agent_identity_lock: Semaphore::new(/*permits*/ 1), agent_identity_bootstrap_cooldown: Mutex::default(), external_auth: RwLock::new(None), @@ -2414,6 +2417,7 @@ impl AuthManager { chatgpt_base_url: None, agent_identity_authapi_base_url: default_agent_identity_authapi_base_url(), refresh_lock: Arc::new(Semaphore::new(/*permits*/ 1)), + unsettled_chatgpt_refresh: Mutex::new(None), agent_identity_lock: Semaphore::new(/*permits*/ 1), agent_identity_bootstrap_cooldown: Mutex::default(), external_auth: RwLock::new(None), @@ -2445,6 +2449,7 @@ impl AuthManager { chatgpt_base_url: None, agent_identity_authapi_base_url: default_agent_identity_authapi_base_url(), refresh_lock: Arc::new(Semaphore::new(/*permits*/ 1)), + unsettled_chatgpt_refresh: Mutex::new(None), agent_identity_lock: Semaphore::new(/*permits*/ 1), agent_identity_bootstrap_cooldown: Mutex::default(), external_auth: RwLock::new(None), @@ -2484,6 +2489,7 @@ impl AuthManager { .to_string(), ), refresh_lock: Arc::new(Semaphore::new(/*permits*/ 1)), + unsettled_chatgpt_refresh: Mutex::new(None), agent_identity_lock: Semaphore::new(/*permits*/ 1), agent_identity_bootstrap_cooldown: Mutex::default(), external_auth: RwLock::new(None), @@ -2522,6 +2528,7 @@ impl AuthManager { chatgpt_base_url: None, agent_identity_authapi_base_url: default_agent_identity_authapi_base_url(), refresh_lock: Arc::new(Semaphore::new(/*permits*/ 1)), + unsettled_chatgpt_refresh: Mutex::new(None), agent_identity_lock: Semaphore::new(/*permits*/ 1), agent_identity_bootstrap_cooldown: Mutex::default(), external_auth: RwLock::new(Some(Arc::new(BearerTokenRefresher::new(config, env)))), @@ -2572,6 +2579,19 @@ impl AuthManager { }) } + /// Whether this credential's refresh exchange or persistence started but did not finish. + /// Cancellation retains the record so callers can withhold an uncertain generation. + /// Waiting for the refresh guard or reloading storage does not mark a credential. + pub fn is_chatgpt_refresh_unsettled(&self, auth: &CodexAuth) -> std::io::Result { + let unsettled = self + .unsettled_chatgpt_refresh + .lock() + .map_err(|_| std::io::Error::other("Refresh progress is unavailable."))?; + Ok(unsettled + .as_ref() + .is_some_and(|tokens| auth.get_token_data().as_ref().ok() == Some(tokens))) + } + /// Current cached auth (clone). May be `None` if not logged in or load failed. /// For managed ChatGPT auth that needs a proactive refresh, first performs /// a guarded reload and then refreshes only if the on-disk auth is unchanged. @@ -3597,42 +3617,64 @@ impl AuthManager { .as_ref() .map(|provenance| provenance.client_id.clone()) .unwrap_or_else(oauth_client_id); - let refresh_response = request_chatgpt_token_refresh( - tokens.refresh_token.clone(), - &endpoint, - &client_id, - auth.client(), - reason, - ) - .await?; + let previous_unsettled = self + .unsettled_chatgpt_refresh + .lock() + .map_err(|_| std::io::Error::other("Refresh progress is unavailable."))? + .replace(tokens.clone()); + let result: Result<(), RefreshTokenError> = async { + let refresh_response = request_chatgpt_token_refresh( + tokens.refresh_token.clone(), + &endpoint, + &client_id, + auth.client(), + reason, + ) + .await?; - // An explicit override may send refresh to another authority. Retain provenance - // only when the endpoint used by this request matches the recorded issuer. - let provenance = tokens.oauth_provenance.clone().filter(|provenance| { - endpoint == format!("{}/oauth/token", provenance.issuer.trim_end_matches('/')) - }); + // An explicit override may send refresh to another authority. Retain provenance + // only when the endpoint used by this request matches the recorded issuer. + let provenance = tokens.oauth_provenance.clone().filter(|provenance| { + endpoint == format!("{}/oauth/token", provenance.issuer.trim_end_matches('/')) + }); - let storage = Arc::clone(auth.storage()); - let originator = codex_otel::auth_storage::AuthStorageOriginator::current(); - tokio::task::spawn_blocking(move || { - originator.sync_scope(|| { - persist_tokens( - &storage, - &tokens, - provenance, - refresh_response.id_token, - refresh_response.access_token, - refresh_response.refresh_token, - ) + let storage = Arc::clone(auth.storage()); + let originator = codex_otel::auth_storage::AuthStorageOriginator::current(); + tokio::task::spawn_blocking(move || { + originator.sync_scope(|| { + persist_tokens( + &storage, + &tokens, + provenance, + refresh_response.id_token, + refresh_response.access_token, + refresh_response.refresh_token, + ) + }) }) - }) - .await - .map_err(|_| { - RefreshTokenError::Transient(std::io::Error::other( - "credential persistence task failed", - )) - })? - .map_err(RefreshTokenError::from)?; + .await + .map_err(|_| { + RefreshTokenError::Transient(std::io::Error::other( + "credential persistence task failed", + )) + })? + .map_err(RefreshTokenError::from)?; + Ok(()) + } + .await; + // Dropping this future while exchange or persistence is pending leaves + // the attempted generation recorded. A known failure cannot settle an + // earlier cancelled attempt; successful persistence can. + *self + .unsettled_chatgpt_refresh + .lock() + .map_err(|_| std::io::Error::other("Refresh progress is unavailable."))? = + if result.is_ok() { + None + } else { + previous_unsettled + }; + result?; self.reload().await; Ok(())