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(())