From 4eedd05bbaab8b034fbb87ea20a378f39132c10d Mon Sep 17 00:00:00 2001 From: Cooper Gamble Date: Tue, 19 May 2026 23:35:08 +0000 Subject: [PATCH] [codex-core] close revoked auth review gaps [ci changed_files] --- codex-rs/core/src/client.rs | 148 +++++-- codex-rs/core/src/client_tests.rs | 216 +++++++++- codex-rs/core/src/compact.rs | 8 + codex-rs/core/src/session/turn.rs | 10 + codex-rs/core/tests/common/test_codex.rs | 5 +- codex-rs/core/tests/suite/client.rs | 3 +- .../core/tests/suite/client_websockets.rs | 391 ++++++++++++++++++ codex-rs/login/src/auth/auth_tests.rs | 382 ++++++++++++++++- codex-rs/login/src/auth/manager.rs | 202 +++++++-- 9 files changed, 1283 insertions(+), 82 deletions(-) diff --git a/codex-rs/core/src/client.rs b/codex-rs/core/src/client.rs index f89205c528..5dc1070c4a 100644 --- a/codex-rs/core/src/client.rs +++ b/codex-rs/core/src/client.rs @@ -1283,6 +1283,7 @@ impl ModelClientSession { stream, session_telemetry.clone(), inference_trace_attempt, + /*revoked_auth_recovery*/ None, ); return Ok(stream); } @@ -1466,6 +1467,7 @@ impl ModelClientSession { stream_result, session_telemetry.clone(), inference_trace_attempt, + auth_recovery, ); self.websocket_session.last_response_rx = Some(last_request_rx); return Ok(WebsocketStreamOutcome::Stream(stream)); @@ -1526,37 +1528,52 @@ impl ModelClientSession { } let disabled_trace = InferenceTraceContext::disabled(); - match self - .stream_responses_websocket( - prompt, - model_info, - session_telemetry, - effort, - summary, - service_tier, - turn_metadata_header, - /*warmup*/ true, - current_span_w3c_trace_context(), - &disabled_trace, - ) - .await - { - Ok(WebsocketStreamOutcome::Stream(mut stream)) => { - // Wait for the v2 warmup request to complete before sending the first turn request. - while let Some(event) = stream.next().await { - match event { - Ok(ResponseEvent::Completed { .. }) => break, - Err(err) => return Err(err), - _ => {} + let mut retried_after_auth_recovery = false; + loop { + match self + .stream_responses_websocket( + prompt, + model_info, + session_telemetry, + effort, + summary, + service_tier.clone(), + turn_metadata_header, + /*warmup*/ true, + current_span_w3c_trace_context(), + &disabled_trace, + ) + .await + { + Ok(WebsocketStreamOutcome::Stream(mut stream)) => { + // Wait for the v2 warmup request to complete before sending the first turn request. + let mut retry_after_auth_recovery = false; + while let Some(event) = stream.next().await { + match event { + Ok(ResponseEvent::Completed { .. }) => return Ok(()), + Err(err) + if !retried_after_auth_recovery + && is_websocket_auth_recovery_retry(&err) => + { + retried_after_auth_recovery = true; + retry_after_auth_recovery = true; + break; + } + Err(err) => return Err(err), + _ => {} + } } + if retry_after_auth_recovery { + continue; + } + return Ok(()); } - Ok(()) + Ok(WebsocketStreamOutcome::FallbackToHttp) => { + self.try_switch_fallback_transport(session_telemetry, model_info); + return Ok(()); + } + Err(err) => return Err(err), } - Ok(WebsocketStreamOutcome::FallbackToHttp) => { - self.try_switch_fallback_transport(session_telemetry, model_info); - Ok(()) - } - Err(err) => Err(err), } } @@ -1736,11 +1753,22 @@ fn parent_thread_id_header_value(session_source: &SessionSource) -> Option bool { + matches!( + err, + CodexErr::Stream(message, _) + if message == WEBSOCKET_AUTH_RECOVERY_RETRY_REASON + ) +} fn map_response_stream( api_stream: codex_api::ResponseStream, session_telemetry: SessionTelemetry, inference_trace_attempt: InferenceTraceAttempt, + revoked_auth_recovery: Option, ) -> (ResponseStream, oneshot::Receiver) { let codex_api::ResponseStream { rx_event, @@ -1755,6 +1783,7 @@ fn map_response_stream( api_stream, session_telemetry, inference_trace_attempt, + revoked_auth_recovery, ) } @@ -1763,6 +1792,7 @@ fn map_response_events( api_stream: S, session_telemetry: SessionTelemetry, inference_trace_attempt: InferenceTraceAttempt, + mut revoked_auth_recovery: Option, ) -> (ResponseStream, oneshot::Receiver) where S: futures::Stream> @@ -1873,7 +1903,12 @@ where if let Some(upstream_request_id) = upstream_request_id { feedback_tags!(last_model_request_id = upstream_request_id); } - let mapped = map_api_error(err); + let mapped = map_response_stream_error( + err, + &mut revoked_auth_recovery, + &session_telemetry, + ) + .await; inference_trace_attempt.record_failed( &mapped, upstream_request_id, @@ -1905,6 +1940,59 @@ where ) } +async fn map_response_stream_error( + err: ApiError, + revoked_auth_recovery: &mut Option, + session_telemetry: &SessionTelemetry, +) -> CodexErr { + match err { + ApiError::Transport( + transport @ TransportError::Http { + status: StatusCode::UNAUTHORIZED, + .. + }, + ) if stream_error_has_revoked_access_token(&transport) => { + let recovery_transport = clone_http_transport_error(&transport); + match handle_unauthorized(recovery_transport, revoked_auth_recovery, session_telemetry) + .await + { + Ok(_) => CodexErr::Stream( + WEBSOCKET_AUTH_RECOVERY_RETRY_REASON.to_string(), + /*requested_delay*/ None, + ), + Err(err) => err, + } + } + err => map_api_error(err), + } +} + +fn stream_error_has_revoked_access_token(transport: &TransportError) -> bool { + let debug = extract_response_debug_context(transport); + access_token_revocation_recovery_reason( + debug.auth_error_code.as_deref(), + debug.auth_error_type.as_deref(), + ) + .is_some() +} + +fn clone_http_transport_error(transport: &TransportError) -> TransportError { + match transport { + TransportError::Http { + status, + url, + headers, + body, + } => TransportError::Http { + status: *status, + url: url.clone(), + headers: headers.clone(), + body: body.clone(), + }, + _ => unreachable!("stream auth recovery only clones HTTP transport errors"), + } +} + /// Handles a 401 response by optionally refreshing ChatGPT tokens once. /// /// When refresh succeeds, the caller should retry the API call; otherwise @@ -1986,7 +2074,7 @@ async fn handle_unauthorized( debug.auth_error_type.as_deref(), ); if let Some(recovery_reason) = revocation_recovery_reason - && let Some(recovery) = auth_recovery.as_ref() + && let Some(recovery) = auth_recovery.as_mut() && recovery.handles_invalidated_access_token_auth() { let mode = recovery.mode_name(); diff --git a/codex-rs/core/src/client_tests.rs b/codex-rs/core/src/client_tests.rs index bdafbadaf4..30815757c9 100644 --- a/codex-rs/core/src/client_tests.rs +++ b/codex-rs/core/src/client_tests.rs @@ -2,6 +2,7 @@ use super::AuthRequestTelemetryContext; use super::ModelClient; use super::PendingUnauthorizedRetry; use super::UnauthorizedRecoveryExecution; +use super::WEBSOCKET_AUTH_RECOVERY_RETRY_REASON; use super::X_CODEX_INSTALLATION_ID_HEADER; use super::X_CODEX_PARENT_THREAD_ID_HEADER; use super::X_CODEX_TURN_METADATA_HEADER; @@ -11,12 +12,18 @@ use super::handle_unauthorized; use crate::AttestationContext; use crate::AttestationProvider; use crate::GenerateAttestationFuture; +use async_trait::async_trait; use base64::Engine; use codex_api::ApiError; use codex_api::ResponseEvent; use codex_api::TransportError; use codex_app_server_protocol::AuthMode; use codex_login::AuthManager; +use codex_login::CodexAuth; +use codex_login::ExternalAuth; +use codex_login::ExternalAuthRefreshContext; +use codex_login::ExternalAuthRefreshReason; +use codex_login::ExternalAuthTokens; use codex_model_provider::BearerAuthProvider; use codex_model_provider_info::CHATGPT_CODEX_BASE_URL; use codex_model_provider_info::ModelProviderInfo; @@ -131,12 +138,12 @@ fn test_session_telemetry() -> SessionTelemetry { ) } -fn write_managed_chatgpt_auth(codex_home: &Path, access_token: &str) { +fn fake_chatgpt_jwt(signature: &str) -> String { let encode_json = |value: serde_json::Value| { base64::engine::general_purpose::URL_SAFE_NO_PAD .encode(serde_json::to_vec(&value).expect("managed auth JWT segment should serialize")) }; - let id_token = format!( + format!( "{}.{}.{}", encode_json(json!({"alg": "none", "typ": "JWT"})), encode_json(json!({ @@ -147,8 +154,12 @@ fn write_managed_chatgpt_auth(codex_home: &Path, access_token: &str) { "user_id": "user-12345" } })), - base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(b"sig"), - ); + base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(signature.as_bytes()), + ) +} + +fn write_managed_chatgpt_auth(codex_home: &Path, access_token: &str) { + let id_token = fake_chatgpt_jwt("managed-auth"); let auth_json = json!({ "tokens": { "id_token": id_token, @@ -178,6 +189,35 @@ async fn managed_chatgpt_auth_manager(access_token: &str) -> (TempDir, Arc AuthMode { + AuthMode::Chatgpt + } + + async fn refresh( + &self, + context: ExternalAuthRefreshContext, + ) -> std::io::Result { + assert_eq!( + context.reason, + ExternalAuthRefreshReason::Unauthorized, + "websocket revoked-token recovery should use unauthorized external refresh" + ); + self.refresh_count.fetch_add(1, Ordering::SeqCst); + Ok(ExternalAuthTokens::chatgpt( + self.refreshed_token.clone(), + "account_id", + Some("pro".to_string()), + )) + } +} + #[derive(Default)] struct TagCollectorVisitor { tags: BTreeMap, @@ -401,6 +441,7 @@ async fn dropped_response_stream_traces_cancelled_partial_output() -> anyhow::Re api_stream, test_session_telemetry(), attempt, + /*revoked_auth_recovery*/ None, ); let observed = stream @@ -450,6 +491,7 @@ async fn response_stream_records_last_model_feedback_ids() { api_stream, test_session_telemetry(), InferenceTraceAttempt::disabled(), + /*revoked_auth_recovery*/ None, ); while stream.next().await.is_some() {} @@ -491,6 +533,7 @@ async fn dropped_backpressured_response_stream_traces_cancelled_partial_output() api_stream, test_session_telemetry(), attempt, + /*revoked_auth_recovery*/ None, ); // Fill the mapper channel with non-terminal events, then yield one output @@ -745,6 +788,171 @@ async fn token_revoked_error_code_401_clears_auth_and_requires_relogin() { assert!(manager.auth_cached().is_none()); } +#[tokio::test] +async fn token_revoked_error_type_401_clears_auth_and_requires_relogin() { + let (_codex_home, manager) = managed_chatgpt_auth_manager("revoked-access-token").await; + let mut recovery = Some(manager.unauthorized_recovery()); + + let err = handle_unauthorized( + TransportError::Http { + status: StatusCode::UNAUTHORIZED, + url: Some("https://chatgpt.com/backend-api/codex/responses".to_string()), + headers: None, + body: Some(r#"{"error":{"type":"token_revoked"}}"#.to_string()), + }, + &mut recovery, + &test_session_telemetry(), + ) + .await + .expect_err("revoked access tokens should force relogin"); + + let CodexErr::RefreshTokenFailed(failed) = err else { + panic!("expected revoked access token to force relogin, got {err:?}"); + }; + assert_eq!(failed.reason, RefreshTokenFailedReason::Revoked); + assert_eq!( + failed.message, + "Your ChatGPT session is no longer valid. Please sign in again." + ); + assert!(manager.auth_cached().is_none()); +} + +#[tokio::test] +async fn websocket_stream_token_revoked_401_clears_auth_and_requires_relogin() { + let (_codex_home, manager) = managed_chatgpt_auth_manager("revoked-access-token").await; + let api_stream = futures::stream::iter([Err(ApiError::Transport(TransportError::Http { + status: StatusCode::UNAUTHORIZED, + url: None, + headers: None, + body: Some(r#"{"type":"error","status":401,"error":{"type":"token_revoked"}}"#.to_string()), + }))]); + + let (mut stream, _) = super::map_response_events( + /*upstream_request_id*/ None, + api_stream, + test_session_telemetry(), + InferenceTraceAttempt::disabled(), + Some(manager.unauthorized_recovery()), + ); + let err = stream + .next() + .await + .expect("revoked websocket stream should emit an error") + .expect_err("revoked websocket stream should require relogin"); + + let CodexErr::RefreshTokenFailed(failed) = err else { + panic!("expected revoked websocket stream to force relogin, got {err:?}"); + }; + assert_eq!(failed.reason, RefreshTokenFailedReason::Revoked); + assert_eq!( + failed.message, + "Your ChatGPT session is no longer valid. Please sign in again." + ); + assert!(manager.auth_cached().is_none()); +} + +#[tokio::test] +async fn websocket_stream_token_revoked_401_retries_after_loading_replacement_auth() { + let (codex_home, manager) = managed_chatgpt_auth_manager("revoked-access-token").await; + write_managed_chatgpt_auth(codex_home.path(), "replacement-access-token"); + let api_stream = futures::stream::iter([Err(ApiError::Transport(TransportError::Http { + status: StatusCode::UNAUTHORIZED, + url: None, + headers: None, + body: Some(r#"{"type":"error","status":401,"error":{"type":"token_revoked"}}"#.to_string()), + }))]); + + let (mut stream, _) = super::map_response_events( + /*upstream_request_id*/ None, + api_stream, + test_session_telemetry(), + InferenceTraceAttempt::disabled(), + Some(manager.unauthorized_recovery()), + ); + let err = stream + .next() + .await + .expect("recovered websocket stream should emit a retryable error") + .expect_err("recovered websocket stream should ask the turn loop to retry"); + + let CodexErr::Stream(message, None) = err else { + panic!( + "expected recovered websocket revocation to surface a retryable stream error, got {err:?}" + ); + }; + assert_eq!(message, WEBSOCKET_AUTH_RECOVERY_RETRY_REASON); + assert_eq!( + manager + .auth_cached() + .expect("replacement auth should remain cached") + .get_token() + .expect("replacement token should resolve"), + "replacement-access-token" + ); +} + +#[tokio::test] +async fn websocket_stream_token_revoked_401_refreshes_external_auth_before_retry() { + let codex_home = TempDir::new().expect("external auth tempdir"); + let initial_access_token = fake_chatgpt_jwt("external-initial"); + let refreshed_access_token = fake_chatgpt_jwt("external-refreshed"); + codex_login::auth::login_with_chatgpt_auth_tokens( + codex_home.path(), + &initial_access_token, + "account_id", + Some("pro"), + ) + .expect("external ChatGPT auth should seed"); + let manager = AuthManager::shared( + codex_home.path().to_path_buf(), + /*enable_codex_api_key_env*/ false, + codex_login::AuthCredentialsStoreMode::File, + /*chatgpt_base_url*/ None, + ) + .await; + let external_auth = Arc::new(RefreshingExternalChatgptAuth { + refreshed_token: refreshed_access_token.clone(), + refresh_count: AtomicUsize::new(0), + }); + manager.set_external_auth(external_auth.clone()); + + let api_stream = futures::stream::iter([Err(ApiError::Transport(TransportError::Http { + status: StatusCode::UNAUTHORIZED, + url: None, + headers: None, + body: Some(r#"{"type":"error","status":401,"error":{"type":"token_revoked"}}"#.to_string()), + }))]); + + let (mut stream, _) = super::map_response_events( + /*upstream_request_id*/ None, + api_stream, + test_session_telemetry(), + InferenceTraceAttempt::disabled(), + Some(manager.unauthorized_recovery()), + ); + let err = stream + .next() + .await + .expect("refreshed websocket stream should emit a retryable error") + .expect_err("refreshed websocket stream should ask the turn loop to retry"); + + let CodexErr::Stream(message, None) = err else { + panic!( + "expected external websocket revocation to surface a retryable stream error, got {err:?}" + ); + }; + assert_eq!(message, WEBSOCKET_AUTH_RECOVERY_RETRY_REASON); + assert_eq!(external_auth.refresh_count.load(Ordering::SeqCst), 1); + assert_eq!( + manager + .auth_cached() + .expect("refreshed external auth should remain cached") + .get_token() + .expect("refreshed token should resolve"), + refreshed_access_token + ); +} + #[tokio::test] async fn token_invalidated_401_retries_when_persisted_auth_changed() { let (codex_home, manager) = managed_chatgpt_auth_manager("revoked-access-token").await; diff --git a/codex-rs/core/src/compact.rs b/codex-rs/core/src/compact.rs index b5802268d3..b4b67bb088 100644 --- a/codex-rs/core/src/compact.rs +++ b/codex-rs/core/src/compact.rs @@ -186,6 +186,7 @@ async fn run_compact_task_inner_impl( let max_retries = turn_context.provider.info().stream_max_retries(); let mut retries = 0; + let mut websocket_auth_recovery_retries = 0; let mut client_session = sess.services.model_client.new_session(); // Reuse one client session so turn-scoped state (sticky routing, websocket incremental // request tracking) @@ -235,6 +236,13 @@ async fn run_compact_task_inner_impl( sess.send_event(&turn_context, event).await; return Err(e); } + Err(e) + if crate::client::is_websocket_auth_recovery_retry(&e) + && websocket_auth_recovery_retries == 0 => + { + websocket_auth_recovery_retries += 1; + continue; + } Err(e) => { if retries < max_retries { retries += 1; diff --git a/codex-rs/core/src/session/turn.rs b/codex-rs/core/src/session/turn.rs index 26ba845460..40b085dc01 100644 --- a/codex-rs/core/src/session/turn.rs +++ b/codex-rs/core/src/session/turn.rs @@ -943,6 +943,7 @@ async fn run_sampling_request( ) .await; let mut retries = 0; + let mut websocket_auth_recovery_retries = 0; let mut initial_input = Some(input); loop { let prompt_input = if let Some(input) = initial_input.take() { @@ -992,6 +993,15 @@ async fn run_sampling_request( return Err(err); } + // Auth recovery already loaded replacement credentials. Give that one immediate retry + // without spending the transport reconnect budget. + if crate::client::is_websocket_auth_recovery_retry(&err) + && websocket_auth_recovery_retries == 0 + { + websocket_auth_recovery_retries += 1; + continue; + } + // Use the configured provider-specific stream retry budget. let max_retries = turn_context.provider.info().stream_max_retries(); if retries >= max_retries diff --git a/codex-rs/core/tests/common/test_codex.rs b/codex-rs/core/tests/common/test_codex.rs index a398758291..8a904141ce 100644 --- a/codex-rs/core/tests/common/test_codex.rs +++ b/codex-rs/core/tests/common/test_codex.rs @@ -467,7 +467,10 @@ impl TestCodexBuilder { let installation_id = resolve_installation_id(&config.codex_home).await?; let thread_manager = ThreadManager::new( &config, - codex_core::test_support::auth_manager_from_auth(auth.clone()), + codex_core::test_support::auth_manager_from_auth_with_home( + auth.clone(), + config.codex_home.to_path_buf(), + ), SessionSource::Exec, Arc::clone(&environment_manager), empty_extension_registry(), diff --git a/codex-rs/core/tests/suite/client.rs b/codex-rs/core/tests/suite/client.rs index 45cf7acdb1..135ab780bd 100644 --- a/codex-rs/core/tests/suite/client.rs +++ b/codex-rs/core/tests/suite/client.rs @@ -1099,7 +1099,7 @@ async fn revoked_chatgpt_auth_user_turn_clears_auth_and_requests_relogin() -> an .await? .expect("managed ChatGPT auth should load"); - let mut model_provider = built_in_model_providers(/* openai_base_url */ None)["openai"].clone(); + let mut model_provider = built_in_model_providers(/*openai_base_url*/ None)["openai"].clone(); model_provider.base_url = Some(format!("{}/api/codex", server.uri())); model_provider.supports_websockets = false; let mut builder = test_codex() @@ -1120,6 +1120,7 @@ async fn revoked_chatgpt_auth_user_turn_clears_auth_and_requests_relogin() -> an }], final_output_json_schema: None, responsesapi_client_metadata: None, + thread_settings: Default::default(), }) .await?; diff --git a/codex-rs/core/tests/suite/client_websockets.rs b/codex-rs/core/tests/suite/client_websockets.rs index 74eb053325..5cc53958af 100755 --- a/codex-rs/core/tests/suite/client_websockets.rs +++ b/codex-rs/core/tests/suite/client_websockets.rs @@ -1,6 +1,7 @@ #![allow(clippy::expect_used, clippy::unwrap_used)] use codex_api::WS_REQUEST_HEADER_TRACEPARENT_CLIENT_METADATA_KEY; use codex_api::WS_REQUEST_HEADER_TRACESTATE_CLIENT_METADATA_KEY; +use codex_config::types::AuthCredentialsStoreMode; use codex_core::ModelClient; use codex_core::ModelClientSession; use codex_core::Prompt; @@ -66,6 +67,42 @@ const TEST_INSTALLATION_ID: &str = "11111111-1111-4111-8111-111111111111"; const X_CODEX_WS_STREAM_REQUEST_START_MS_CLIENT_METADATA_KEY: &str = "x-codex-ws-stream-request-start-ms"; +fn write_managed_chatgpt_auth(codex_home: &TempDir, access_token: &str) { + use base64::Engine as _; + + let encode_json = |value: serde_json::Value| { + base64::engine::general_purpose::URL_SAFE_NO_PAD + .encode(serde_json::to_vec(&value).expect("managed auth JWT segment should serialize")) + }; + let id_token = format!( + "{}.{}.{}", + encode_json(json!({"alg": "none", "typ": "JWT"})), + encode_json(json!({ + "email": "user@example.com", + "https://api.openai.com/auth": { + "chatgpt_account_id": "account_id", + "chatgpt_user_id": "user-12345", + "user_id": "user-12345" + } + })), + base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(b"sig"), + ); + let auth_json = json!({ + "tokens": { + "id_token": id_token, + "access_token": access_token, + "refresh_token": "test-refresh-token", + "account_id": "account_id" + }, + "last_refresh": chrono::Utc::now(), + }); + std::fs::write( + codex_home.path().join("auth.json"), + serde_json::to_vec_pretty(&auth_json).expect("managed auth should serialize"), + ) + .expect("managed auth should persist"); +} + fn assert_request_trace_matches(body: &serde_json::Value, expected_trace: &W3cTraceContext) { let client_metadata = body["client_metadata"] .as_object() @@ -1429,6 +1466,360 @@ async fn responses_websocket_invalid_request_error_with_status_is_forwarded() { server.shutdown().await; } +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn responses_websocket_revoked_managed_auth_clears_auth_and_requests_relogin() { + skip_if_no_network!(); + + let revoked_token_error = json!({ + "type": "error", + "status": 401, + "error": { + "type": "token_revoked", + "message": "revoked" + } + }); + let server = start_websocket_server(vec![vec![ + vec![ + ev_response_created("resp-prewarm"), + ev_completed("resp-prewarm"), + ], + vec![revoked_token_error], + ]]) + .await; + + let codex_home = Arc::new(TempDir::new().expect("managed auth tempdir")); + write_managed_chatgpt_auth(codex_home.as_ref(), "revoked-access-token"); + let auth = CodexAuth::from_auth_storage( + codex_home.path(), + AuthCredentialsStoreMode::File, + /*chatgpt_base_url*/ None, + ) + .await + .expect("managed ChatGPT auth should load") + .expect("managed ChatGPT auth should exist"); + + let mut builder = test_codex() + .with_home(codex_home.clone()) + .with_auth(auth) + .with_config(|config| { + config.model_provider.request_max_retries = Some(0); + config.model_provider.stream_max_retries = Some(0); + }); + let test = builder + .build_with_websocket_server(&server) + .await + .expect("build websocket codex"); + + let submission_id = test + .codex + .submit(Op::UserInput { + environments: None, + items: vec![UserInput::Text { + text: "hello".into(), + text_elements: Vec::new(), + }], + final_output_json_schema: None, + responsesapi_client_metadata: None, + thread_settings: Default::default(), + }) + .await + .expect("submission should emit relogin events"); + + let error_event = wait_for_event(&test.codex, |msg| matches!(msg, EventMsg::Error(_))).await; + let EventMsg::Error(error_event) = error_event else { + unreachable!(); + }; + assert!( + error_event + .message + .contains("Your ChatGPT session is no longer valid. Please sign in again."), + "unexpected error message for submission {submission_id}: {}", + error_event.message + ); + wait_for_event(&test.codex, |msg| matches!(msg, EventMsg::TurnComplete(_))).await; + + assert!( + !test.codex_home_path().join("auth.json").exists(), + "revoked managed ChatGPT auth should be removed" + ); + assert_eq!(server.single_connection().len(), 2); + + server.shutdown().await; +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn responses_websocket_revoked_managed_auth_retries_reloaded_auth() { + skip_if_no_network!(); + + let revoked_token_error = json!({ + "type": "error", + "status": 401, + "error": { + "type": "token_revoked", + "message": "revoked" + } + }); + let server = start_websocket_server(vec![ + vec![ + vec![ + ev_response_created("resp-prewarm"), + ev_completed("resp-prewarm"), + ], + vec![revoked_token_error], + ], + vec![vec![ + ev_response_created("resp-retry"), + ev_completed("resp-retry"), + ]], + ]) + .await; + + let codex_home = Arc::new(TempDir::new().expect("managed auth tempdir")); + write_managed_chatgpt_auth(codex_home.as_ref(), "revoked-access-token"); + let auth = CodexAuth::from_auth_storage( + codex_home.path(), + AuthCredentialsStoreMode::File, + /*chatgpt_base_url*/ None, + ) + .await + .expect("managed ChatGPT auth should load") + .expect("managed ChatGPT auth should exist"); + + let mut builder = test_codex() + .with_home(codex_home.clone()) + .with_auth(auth) + .with_config(|config| { + config.model_provider.request_max_retries = Some(0); + config.model_provider.stream_max_retries = Some(0); + }); + let test = builder + .build_with_websocket_server(&server) + .await + .expect("build websocket codex"); + + write_managed_chatgpt_auth(codex_home.as_ref(), "replacement-access-token"); + + test.submit_turn("hello") + .await + .expect("reloaded managed auth should retry the websocket turn"); + + let persisted_auth = CodexAuth::from_auth_storage( + codex_home.path(), + AuthCredentialsStoreMode::File, + /*chatgpt_base_url*/ None, + ) + .await + .expect("replacement managed auth should load") + .expect("replacement managed auth should persist"); + assert_eq!( + persisted_auth + .get_token() + .expect("replacement token should resolve"), + "replacement-access-token" + ); + + let connections = server.connections(); + assert_eq!(connections.len(), 2); + assert_eq!(connections[0].len(), 2); + assert_eq!(connections[1].len(), 1); + + let handshakes = server.handshakes(); + assert_eq!(handshakes.len(), 2); + assert_eq!( + handshakes[0].header("authorization").as_deref(), + Some("Bearer revoked-access-token") + ); + assert_eq!( + handshakes[1].header("authorization").as_deref(), + Some("Bearer replacement-access-token") + ); + + server.shutdown().await; +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn responses_websocket_revoked_managed_auth_prewarm_retries_reloaded_auth() { + skip_if_no_network!(); + + let revoked_token_error = json!({ + "type": "error", + "status": 401, + "error": { + "type": "token_revoked", + "message": "revoked" + } + }); + let server = start_websocket_server(vec![ + vec![vec![revoked_token_error]], + vec![ + vec![ + ev_response_created("warm-retry"), + ev_completed("warm-retry"), + ], + vec![ev_response_created("resp-1"), ev_completed("resp-1")], + ], + ]) + .await; + + let codex_home = Arc::new(TempDir::new().expect("managed auth tempdir")); + write_managed_chatgpt_auth(codex_home.as_ref(), "revoked-access-token"); + let auth = CodexAuth::from_auth_storage( + codex_home.path(), + AuthCredentialsStoreMode::File, + /*chatgpt_base_url*/ None, + ) + .await + .expect("managed ChatGPT auth should load") + .expect("managed ChatGPT auth should exist"); + write_managed_chatgpt_auth(codex_home.as_ref(), "replacement-access-token"); + + let mut builder = test_codex() + .with_home(codex_home.clone()) + .with_auth(auth) + .with_config(|config| { + config.model_provider.request_max_retries = Some(0); + config.model_provider.stream_max_retries = Some(0); + }); + let test = builder + .build_with_websocket_server(&server) + .await + .expect("build websocket codex"); + + test.submit_turn("hello") + .await + .expect("reloaded managed auth should retry startup websocket prewarm"); + + let persisted_auth = CodexAuth::from_auth_storage( + codex_home.path(), + AuthCredentialsStoreMode::File, + /*chatgpt_base_url*/ None, + ) + .await + .expect("replacement managed auth should load") + .expect("replacement managed auth should persist"); + assert_eq!( + persisted_auth + .get_token() + .expect("replacement token should resolve"), + "replacement-access-token" + ); + + let connections = server.connections(); + assert_eq!(connections.len(), 2); + assert_eq!(connections[0].len(), 1); + assert_eq!(connections[1].len(), 2); + + let handshakes = server.handshakes(); + assert_eq!(handshakes.len(), 2); + assert_eq!( + handshakes[0].header("authorization").as_deref(), + Some("Bearer revoked-access-token") + ); + assert_eq!( + handshakes[1].header("authorization").as_deref(), + Some("Bearer replacement-access-token") + ); + + server.shutdown().await; +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn responses_websocket_revoked_managed_auth_compact_retries_reloaded_auth() { + skip_if_no_network!(); + + let revoked_token_error = json!({ + "type": "error", + "status": 401, + "error": { + "type": "token_revoked", + "message": "revoked" + } + }); + let server = start_websocket_server(vec![ + vec![ + vec![ + ev_response_created("resp-prewarm"), + ev_completed("resp-prewarm"), + ], + vec![ev_response_created("resp-1"), ev_completed("resp-1")], + vec![revoked_token_error], + ], + vec![vec![ + ev_response_created("compact-retry"), + ev_assistant_message("msg-compact", "COMPACT_SUMMARY"), + ev_completed("compact-retry"), + ]], + ]) + .await; + + let codex_home = Arc::new(TempDir::new().expect("managed auth tempdir")); + write_managed_chatgpt_auth(codex_home.as_ref(), "revoked-access-token"); + let auth = CodexAuth::from_auth_storage( + codex_home.path(), + AuthCredentialsStoreMode::File, + /*chatgpt_base_url*/ None, + ) + .await + .expect("managed ChatGPT auth should load") + .expect("managed ChatGPT auth should exist"); + + let mut builder = test_codex() + .with_home(codex_home.clone()) + .with_auth(auth) + .with_config(|config| { + config.model_provider.request_max_retries = Some(0); + config.model_provider.stream_max_retries = Some(0); + }); + let test = builder + .build_with_websocket_server(&server) + .await + .expect("build websocket codex"); + + test.submit_turn("hello") + .await + .expect("initial websocket turn should complete"); + write_managed_chatgpt_auth(codex_home.as_ref(), "replacement-access-token"); + + test.codex + .submit(Op::Compact) + .await + .expect("compact submission should succeed"); + wait_for_event(&test.codex, |msg| matches!(msg, EventMsg::TurnComplete(_))).await; + + let persisted_auth = CodexAuth::from_auth_storage( + codex_home.path(), + AuthCredentialsStoreMode::File, + /*chatgpt_base_url*/ None, + ) + .await + .expect("replacement managed auth should load") + .expect("replacement managed auth should persist"); + assert_eq!( + persisted_auth + .get_token() + .expect("replacement token should resolve"), + "replacement-access-token" + ); + + let connections = server.connections(); + assert_eq!(connections.len(), 2); + assert_eq!(connections[0].len(), 3); + assert_eq!(connections[1].len(), 1); + + let handshakes = server.handshakes(); + assert_eq!(handshakes.len(), 2); + assert_eq!( + handshakes[0].header("authorization").as_deref(), + Some("Bearer revoked-access-token") + ); + assert_eq!( + handshakes[1].header("authorization").as_deref(), + Some("Bearer replacement-access-token") + ); + + server.shutdown().await; +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn responses_websocket_connection_limit_error_reconnects_and_completes() { skip_if_no_network!(); diff --git a/codex-rs/login/src/auth/auth_tests.rs b/codex-rs/login/src/auth/auth_tests.rs index bf10c2c183..42b3e7a84a 100644 --- a/codex-rs/login/src/auth/auth_tests.rs +++ b/codex-rs/login/src/auth/auth_tests.rs @@ -308,6 +308,7 @@ async fn unauthorized_recovery_reports_mode_and_step_names() { let managed = UnauthorizedRecovery { manager: Arc::clone(&manager), step: UnauthorizedRecoveryStep::Reload, + expected_auth: None, expected_account_id: None, mode: UnauthorizedRecoveryMode::Managed, }; @@ -317,6 +318,7 @@ async fn unauthorized_recovery_reports_mode_and_step_names() { let external = UnauthorizedRecovery { manager, step: UnauthorizedRecoveryStep::ExternalRefresh, + expected_auth: None, expected_account_id: None, mode: UnauthorizedRecoveryMode::External, }; @@ -325,8 +327,10 @@ async fn unauthorized_recovery_reports_mode_and_step_names() { } #[tokio::test] +#[serial(codex_auth_env)] async fn invalidated_access_token_logout_clears_cached_auth() { let codex_home = tempdir().unwrap(); + let _access_token_guard = remove_access_token_env_var(); write_auth_file( AuthFileParams { openai_api_key: None, @@ -343,7 +347,7 @@ async fn invalidated_access_token_logout_clears_cached_auth() { /*chatgpt_base_url*/ None, ) .await; - let recovery = manager.unauthorized_recovery(); + let mut recovery = manager.unauthorized_recovery(); assert!(recovery.handles_invalidated_access_token_auth()); let failed = recovery @@ -361,8 +365,10 @@ async fn invalidated_access_token_logout_clears_cached_auth() { } #[tokio::test] +#[serial(codex_auth_env)] async fn invalidated_access_token_logout_clears_cached_auth_without_account_id() { let codex_home = tempdir().unwrap(); + let _access_token_guard = remove_access_token_env_var(); write_auth_file( AuthFileParams { openai_api_key: None, @@ -379,7 +385,7 @@ async fn invalidated_access_token_logout_clears_cached_auth_without_account_id() /*chatgpt_base_url*/ None, ) .await; - let recovery = manager.unauthorized_recovery(); + let mut recovery = manager.unauthorized_recovery(); let failed = recovery .handle_invalidated_access_token_auth() @@ -396,8 +402,10 @@ async fn invalidated_access_token_logout_clears_cached_auth_without_account_id() } #[tokio::test] -async fn invalidated_access_token_preserves_reloaded_auth() { +#[serial(codex_auth_env)] +async fn invalidated_access_token_clears_cached_auth_when_persisted_auth_was_removed() { let codex_home = tempdir().unwrap(); + let _access_token_guard = remove_access_token_env_var(); write_auth_file( AuthFileParams { openai_api_key: None, @@ -414,7 +422,85 @@ async fn invalidated_access_token_preserves_reloaded_auth() { /*chatgpt_base_url*/ None, ) .await; - let recovery = manager.unauthorized_recovery(); + let mut recovery = manager.unauthorized_recovery(); + + std::fs::remove_file(codex_home.path().join("auth.json")) + .expect("auth file should be removable"); + + let failed = recovery + .handle_invalidated_access_token_auth() + .await + .expect_err("removed persisted auth should force login"); + + assert_eq!(failed.reason, RefreshTokenFailedReason::Revoked); + assert_eq!( + failed.message, + "Your ChatGPT session is no longer valid. Please sign in again." + ); + assert!(manager.auth_cached().is_none()); +} + +#[tokio::test] +#[serial(codex_auth_env)] +async fn invalidated_access_token_preserves_cached_auth_when_storage_inspection_fails() { + let codex_home = tempdir().unwrap(); + let _access_token_guard = remove_access_token_env_var(); + write_auth_file( + AuthFileParams { + openai_api_key: None, + chatgpt_plan_type: Some("pro".to_string()), + chatgpt_account_id: Some("org_mine".to_string()), + }, + codex_home.path(), + ) + .expect("failed to write auth file"); + let manager = AuthManager::shared( + codex_home.path().to_path_buf(), + /*enable_codex_api_key_env*/ false, + AuthCredentialsStoreMode::File, + /*chatgpt_base_url*/ None, + ) + .await; + let mut recovery = manager.unauthorized_recovery(); + + std::fs::write(codex_home.path().join("auth.json"), "{not-json") + .expect("auth file should be corruptible"); + + let failed = recovery + .handle_invalidated_access_token_auth() + .await + .expect_err("storage inspection failures should stop invalidation cleanup"); + + assert_eq!(failed.reason, RefreshTokenFailedReason::Revoked); + assert!(failed.message.starts_with( + "Your ChatGPT session is no longer valid. Please sign in again. Codex could not inspect saved auth:" + )); + assert!(manager.auth_cached().is_some()); + assert!(codex_home.path().join("auth.json").exists()); +} + +#[tokio::test] +#[serial(codex_auth_env)] +async fn invalidated_access_token_preserves_reloaded_auth() { + let codex_home = tempdir().unwrap(); + let _access_token_guard = remove_access_token_env_var(); + write_auth_file( + AuthFileParams { + openai_api_key: None, + chatgpt_plan_type: Some("pro".to_string()), + chatgpt_account_id: Some("org_mine".to_string()), + }, + codex_home.path(), + ) + .expect("failed to write auth file"); + let manager = AuthManager::shared( + codex_home.path().to_path_buf(), + /*enable_codex_api_key_env*/ false, + AuthCredentialsStoreMode::File, + /*chatgpt_base_url*/ None, + ) + .await; + let mut recovery = manager.unauthorized_recovery(); let mut reauthenticated = load_auth_dot_json(codex_home.path(), AuthCredentialsStoreMode::File) .expect("auth should load") @@ -448,6 +534,294 @@ async fn invalidated_access_token_preserves_reloaded_auth() { assert!(codex_home.path().join("auth.json").exists()); } +#[tokio::test] +#[serial(codex_auth_env)] +async fn invalidated_access_token_logs_out_reloaded_auth_if_retry_is_also_revoked() { + let codex_home = tempdir().unwrap(); + let _access_token_guard = remove_access_token_env_var(); + write_auth_file( + AuthFileParams { + openai_api_key: None, + chatgpt_plan_type: Some("pro".to_string()), + chatgpt_account_id: Some("org_mine".to_string()), + }, + codex_home.path(), + ) + .expect("failed to write auth file"); + let manager = AuthManager::shared( + codex_home.path().to_path_buf(), + /*enable_codex_api_key_env*/ false, + AuthCredentialsStoreMode::File, + /*chatgpt_base_url*/ None, + ) + .await; + let mut recovery = manager.unauthorized_recovery(); + + let mut reauthenticated = load_auth_dot_json(codex_home.path(), AuthCredentialsStoreMode::File) + .expect("auth should load") + .expect("auth should exist"); + reauthenticated + .tokens + .as_mut() + .expect("tokens should exist") + .access_token = "replacement-access-token".to_string(); + save_auth( + codex_home.path(), + &reauthenticated, + AuthCredentialsStoreMode::File, + ) + .expect("replacement auth should persist"); + + recovery + .handle_invalidated_access_token_auth() + .await + .expect("first revoked token should retry the replacement auth"); + + let failed = recovery + .handle_invalidated_access_token_auth() + .await + .expect_err("repeated revocation of the replacement auth should force login"); + + assert_eq!(failed.reason, RefreshTokenFailedReason::Revoked); + assert_eq!( + failed.message, + "Your ChatGPT session is no longer valid. Please sign in again." + ); + assert!(manager.auth_cached().is_none()); + assert!(!codex_home.path().join("auth.json").exists()); +} + +#[tokio::test] +#[serial(codex_auth_env)] +async fn invalidated_access_token_logs_out_auth_reloaded_by_normal_recovery() { + let codex_home = tempdir().unwrap(); + let _access_token_guard = remove_access_token_env_var(); + write_auth_file( + AuthFileParams { + openai_api_key: None, + chatgpt_plan_type: Some("pro".to_string()), + chatgpt_account_id: Some("org_mine".to_string()), + }, + codex_home.path(), + ) + .expect("failed to write auth file"); + let manager = AuthManager::shared( + codex_home.path().to_path_buf(), + /*enable_codex_api_key_env*/ false, + AuthCredentialsStoreMode::File, + /*chatgpt_base_url*/ None, + ) + .await; + let mut recovery = manager.unauthorized_recovery(); + + let mut reauthenticated = load_auth_dot_json(codex_home.path(), AuthCredentialsStoreMode::File) + .expect("auth should load") + .expect("auth should exist"); + reauthenticated + .tokens + .as_mut() + .expect("tokens should exist") + .access_token = "replacement-access-token".to_string(); + save_auth( + codex_home.path(), + &reauthenticated, + AuthCredentialsStoreMode::File, + ) + .expect("replacement auth should persist"); + + let reloaded = recovery + .next() + .await + .expect("normal unauthorized recovery should reload replacement auth"); + assert_eq!(reloaded.auth_state_changed(), Some(true)); + + let failed = recovery + .handle_invalidated_access_token_auth() + .await + .expect_err("revoking the just-reloaded auth should force login"); + + assert_eq!(failed.reason, RefreshTokenFailedReason::Revoked); + assert_eq!( + failed.message, + "Your ChatGPT session is no longer valid. Please sign in again." + ); + assert!(manager.auth_cached().is_none()); + assert!(!codex_home.path().join("auth.json").exists()); +} + +#[tokio::test] +#[serial(codex_auth_env)] +async fn invalidated_access_token_preserves_auth_refreshed_after_request_started() { + let codex_home = tempdir().unwrap(); + let _access_token_guard = remove_access_token_env_var(); + write_auth_file( + AuthFileParams { + openai_api_key: None, + chatgpt_plan_type: Some("pro".to_string()), + chatgpt_account_id: Some("org_mine".to_string()), + }, + codex_home.path(), + ) + .expect("failed to write auth file"); + let manager = AuthManager::shared( + codex_home.path().to_path_buf(), + /*enable_codex_api_key_env*/ false, + AuthCredentialsStoreMode::File, + /*chatgpt_base_url*/ None, + ) + .await; + let mut recovery = manager.unauthorized_recovery(); + + let mut refreshed_auth = load_auth_dot_json(codex_home.path(), AuthCredentialsStoreMode::File) + .expect("auth should load") + .expect("auth should exist"); + refreshed_auth + .tokens + .as_mut() + .expect("tokens should exist") + .access_token = "replacement-access-token".to_string(); + save_auth( + codex_home.path(), + &refreshed_auth, + AuthCredentialsStoreMode::File, + ) + .expect("replacement auth should persist"); + manager.reload().await; + assert_eq!( + manager + .auth_cached() + .expect("replacement auth should load into cache") + .get_token() + .expect("replacement token should resolve"), + "replacement-access-token" + ); + + let step_result = recovery + .handle_invalidated_access_token_auth() + .await + .expect("auth refreshed after request start should be retried"); + + assert_eq!(step_result.auth_state_changed(), Some(true)); + assert_eq!( + manager + .auth_cached() + .expect("replacement auth should remain cached") + .get_token() + .expect("replacement token should resolve"), + "replacement-access-token" + ); + assert!(codex_home.path().join("auth.json").exists()); +} + +#[tokio::test] +#[serial(codex_auth_env)] +async fn invalidated_access_token_preserves_reloaded_auth_from_a_different_account() { + let codex_home = tempdir().unwrap(); + let _access_token_guard = remove_access_token_env_var(); + write_auth_file( + AuthFileParams { + openai_api_key: None, + chatgpt_plan_type: Some("pro".to_string()), + chatgpt_account_id: Some("org_mine".to_string()), + }, + codex_home.path(), + ) + .expect("failed to write auth file"); + let manager = AuthManager::shared( + codex_home.path().to_path_buf(), + /*enable_codex_api_key_env*/ false, + AuthCredentialsStoreMode::File, + /*chatgpt_base_url*/ None, + ) + .await; + let mut recovery = manager.unauthorized_recovery(); + + write_auth_file( + AuthFileParams { + openai_api_key: None, + chatgpt_plan_type: Some("pro".to_string()), + chatgpt_account_id: Some("org_elsewhere".to_string()), + }, + codex_home.path(), + ) + .expect("replacement auth should persist"); + + let step_result = recovery + .handle_invalidated_access_token_auth() + .await + .expect("new persisted auth should be retried across account changes"); + + assert_eq!(step_result.auth_state_changed(), Some(true)); + assert_eq!( + manager + .auth_cached() + .expect("replacement auth should remain cached") + .get_account_id() + .as_deref(), + Some("org_elsewhere") + ); + let follow_up_step = recovery + .next() + .await + .expect("follow-up refreshable 401 should use the replacement account"); + assert_eq!(follow_up_step.auth_state_changed(), Some(false)); + assert_eq!(recovery.step_name(), "refresh_token"); + assert!(codex_home.path().join("auth.json").exists()); +} + +#[tokio::test] +#[serial(codex_auth_env)] +async fn invalidated_access_token_logout_preserves_auth_changed_before_delete() { + let codex_home = tempdir().unwrap(); + let _access_token_guard = remove_access_token_env_var(); + write_auth_file( + AuthFileParams { + openai_api_key: None, + chatgpt_plan_type: Some("pro".to_string()), + chatgpt_account_id: Some("org_mine".to_string()), + }, + codex_home.path(), + ) + .expect("failed to write auth file"); + let manager = AuthManager::shared( + codex_home.path().to_path_buf(), + /*enable_codex_api_key_env*/ false, + AuthCredentialsStoreMode::File, + /*chatgpt_base_url*/ None, + ) + .await; + + let mut reauthenticated = load_auth_dot_json(codex_home.path(), AuthCredentialsStoreMode::File) + .expect("auth should load") + .expect("auth should exist"); + reauthenticated + .tokens + .as_mut() + .expect("tokens should exist") + .access_token = "replacement-access-token".to_string(); + save_auth( + codex_home.path(), + &reauthenticated, + AuthCredentialsStoreMode::File, + ) + .expect("replacement auth should persist"); + + let outcome = manager + .logout_if_auth_snapshot_unchanged() + .await + .expect("replaced persisted auth should avoid logout"); + assert!(matches!(outcome, InvalidatedAuthLogoutOutcome::AuthChanged)); + assert_eq!( + manager + .auth_cached() + .expect("replacement auth should remain cached") + .get_token() + .expect("replacement token should resolve"), + "replacement-access-token" + ); + assert!(codex_home.path().join("auth.json").exists()); +} + #[tokio::test] #[serial(codex_auth_env)] async fn refresh_failure_is_scoped_to_the_matching_auth_snapshot() { diff --git a/codex-rs/login/src/auth/manager.rs b/codex-rs/login/src/auth/manager.rs index d8c89ec846..17615620dd 100644 --- a/codex-rs/login/src/auth/manager.rs +++ b/codex-rs/login/src/auth/manager.rs @@ -1059,6 +1059,11 @@ enum ReloadOutcome { Skipped, } +enum InvalidatedAuthLogoutOutcome { + LoggedOut, + AuthChanged, +} + #[derive(Clone, Copy, Debug, PartialEq, Eq)] enum UnauthorizedRecoveryMode { Managed, @@ -1088,6 +1093,7 @@ enum UnauthorizedRecoveryMode { pub struct UnauthorizedRecovery { manager: Arc, step: UnauthorizedRecoveryStep, + expected_auth: Option, expected_account_id: Option, mode: UnauthorizedRecoveryMode, } @@ -1123,6 +1129,7 @@ impl UnauthorizedRecovery { Self { manager, step, + expected_auth: cached_auth, expected_account_id, mode, } @@ -1219,6 +1226,7 @@ impl UnauthorizedRecovery { .await { ReloadOutcome::ReloadedChanged => { + self.update_expected_auth_from_cache(); self.step = UnauthorizedRecoveryStep::RefreshToken; return Ok(UnauthorizedRecoveryStepResult { auth_state_changed: Some(true), @@ -1241,6 +1249,7 @@ impl UnauthorizedRecovery { } UnauthorizedRecoveryStep::RefreshToken => { self.manager.refresh_token_from_authority().await?; + self.update_expected_auth_from_cache(); self.step = UnauthorizedRecoveryStep::Done; return Ok(UnauthorizedRecoveryStepResult { auth_state_changed: Some(true), @@ -1250,6 +1259,7 @@ impl UnauthorizedRecovery { self.manager .refresh_external_auth(ExternalAuthRefreshReason::Unauthorized) .await?; + self.update_expected_auth_from_cache(); self.step = UnauthorizedRecoveryStep::Done; return Ok(UnauthorizedRecoveryStepResult { auth_state_changed: Some(true), @@ -1262,47 +1272,125 @@ impl UnauthorizedRecovery { }) } - pub async fn handle_invalidated_access_token_auth( - &self, - ) -> Result { - let reload_outcome = match self.expected_account_id.as_deref() { - Some(expected_account_id) => { - self.manager - .reload_if_account_id_matches(Some(expected_account_id)) - .await - } - None => self.manager.reload_if_auth_snapshot_changed().await, - }; + fn update_expected_auth_from_cache(&mut self) { + let expected_auth = self.manager.auth_cached(); + self.expected_account_id = expected_auth.as_ref().and_then(CodexAuth::get_account_id); + self.expected_auth = expected_auth; + } - match reload_outcome { - ReloadOutcome::ReloadedChanged => { - if self.manager.auth_cached().is_none() { - return Err(RefreshTokenFailedError::new( - RefreshTokenFailedReason::Revoked, - ACCESS_TOKEN_INVALIDATED_MESSAGE.to_string(), - )); + pub async fn handle_invalidated_access_token_auth( + &mut self, + ) -> Result { + let (result, next_expected_auth) = { + let _refresh_guard = self.manager.refresh_lock.acquire().await.map_err(|_| { + RefreshTokenFailedError::new( + RefreshTokenFailedReason::Other, + REFRESH_TOKEN_UNKNOWN_MESSAGE.to_string(), + ) + })?; + let cached_auth = self.manager.auth_cached(); + if !AuthManager::auths_equal_for_refresh( + cached_auth.as_ref(), + self.expected_auth.as_ref(), + ) { + self.invalidated_auth_state_changed_result() + } else { + let reload_outcome = self + .manager + .reload_if_auth_snapshot_changed() + .await + .map_err(Self::invalidated_auth_storage_error)?; + + match reload_outcome { + ReloadOutcome::ReloadedChanged => self.invalidated_auth_state_changed_result(), + ReloadOutcome::ReloadedNoChange => { + match self.manager.logout_if_auth_snapshot_unchanged().await { + Ok(InvalidatedAuthLogoutOutcome::AuthChanged) => { + self.invalidated_auth_state_changed_result() + } + Ok(InvalidatedAuthLogoutOutcome::LoggedOut) => ( + Err(RefreshTokenFailedError::new( + RefreshTokenFailedReason::Revoked, + ACCESS_TOKEN_INVALIDATED_MESSAGE.to_string(), + )), + None, + ), + Err(err) => ( + Err(RefreshTokenFailedError::new( + RefreshTokenFailedReason::Revoked, + format!( + "{ACCESS_TOKEN_INVALIDATED_MESSAGE} Codex could not clear saved auth: {err}" + ), + )), + None, + ), + } + } + ReloadOutcome::Skipped => { + if self + .manager + .clear_cached_auth_if_storage_missing() + .await + .map_err(Self::invalidated_auth_storage_error)? + { + ( + Err(RefreshTokenFailedError::new( + RefreshTokenFailedReason::Revoked, + ACCESS_TOKEN_INVALIDATED_MESSAGE.to_string(), + )), + None, + ) + } else { + ( + Err(RefreshTokenFailedError::new( + RefreshTokenFailedReason::Other, + REFRESH_TOKEN_ACCOUNT_MISMATCH_MESSAGE.to_string(), + )), + None, + ) + } + } } - Ok(UnauthorizedRecoveryStepResult { - auth_state_changed: Some(true), - }) } - ReloadOutcome::ReloadedNoChange => { - let message = match self.manager.logout().await { - Ok(_) => ACCESS_TOKEN_INVALIDATED_MESSAGE.to_string(), - Err(err) => format!( - "{ACCESS_TOKEN_INVALIDATED_MESSAGE} Codex could not clear saved auth: {err}" - ), - }; + }; + if result.is_ok() { + self.expected_account_id = next_expected_auth + .as_ref() + .and_then(CodexAuth::get_account_id); + self.expected_auth = next_expected_auth; + } + result + } + + fn invalidated_auth_state_changed_result( + &self, + ) -> ( + Result, + Option, + ) { + let cached_auth = self.manager.auth_cached(); + if cached_auth.is_none() { + return ( Err(RefreshTokenFailedError::new( RefreshTokenFailedReason::Revoked, - message, - )) - } - ReloadOutcome::Skipped => Err(RefreshTokenFailedError::new( - RefreshTokenFailedReason::Other, - REFRESH_TOKEN_ACCOUNT_MISMATCH_MESSAGE.to_string(), - )), + ACCESS_TOKEN_INVALIDATED_MESSAGE.to_string(), + )), + None, + ); } + ( + Ok(UnauthorizedRecoveryStepResult { + auth_state_changed: Some(true), + }), + cached_auth, + ) + } + + fn invalidated_auth_storage_error(err: std::io::Error) -> RefreshTokenFailedError { + RefreshTokenFailedError::new( + RefreshTokenFailedReason::Revoked, + format!("{ACCESS_TOKEN_INVALIDATED_MESSAGE} Codex could not inspect saved auth: {err}"), + ) } } @@ -1545,18 +1633,46 @@ impl AuthManager { } } - async fn reload_if_auth_snapshot_changed(&self) -> ReloadOutcome { - let new_auth = self.load_auth_from_storage().await; + async fn reload_if_auth_snapshot_changed(&self) -> std::io::Result { + let new_auth = self.try_load_auth_from_storage().await?; let cached_before_reload = self.auth_cached(); let auth_changed = !Self::auths_equal_for_refresh(cached_before_reload.as_ref(), new_auth.as_ref()); if !auth_changed { - return ReloadOutcome::ReloadedNoChange; + return Ok(ReloadOutcome::ReloadedNoChange); } tracing::info!("Reloading auth because the persisted auth snapshot changed."); self.set_cached_auth(new_auth); - ReloadOutcome::ReloadedChanged + Ok(ReloadOutcome::ReloadedChanged) + } + + async fn clear_cached_auth_if_storage_missing(&self) -> std::io::Result { + if self.try_load_auth_from_storage().await?.is_some() { + return Ok(false); + } + + tracing::info!("Clearing cached auth because persisted auth is no longer available."); + self.set_cached_auth(/*new_auth*/ None); + Ok(true) + } + + async fn logout_if_auth_snapshot_unchanged( + &self, + ) -> std::io::Result { + let persisted_auth = self.try_load_auth_from_storage().await?; + let cached_auth = self.auth_cached(); + if !Self::auths_equal_for_refresh(cached_auth.as_ref(), persisted_auth.as_ref()) { + tracing::info!( + "Skipping auth logout because the persisted auth snapshot changed before deletion." + ); + self.set_cached_auth(persisted_auth); + return Ok(InvalidatedAuthLogoutOutcome::AuthChanged); + } + + logout_all_stores(&self.codex_home, self.auth_credentials_store_mode)?; + self.reload().await; + Ok(InvalidatedAuthLogoutOutcome::LoggedOut) } fn auths_equal_for_refresh(a: Option<&CodexAuth>, b: Option<&CodexAuth>) -> bool { @@ -1607,7 +1723,7 @@ impl AuthManager { } } - async fn load_auth_from_storage(&self) -> Option { + async fn try_load_auth_from_storage(&self) -> std::io::Result> { load_auth( &self.codex_home, self.enable_codex_api_key_env, @@ -1615,8 +1731,10 @@ impl AuthManager { self.chatgpt_base_url.as_deref(), ) .await - .ok() - .flatten() + } + + async fn load_auth_from_storage(&self) -> Option { + self.try_load_auth_from_storage().await.ok().flatten() } fn set_cached_auth(&self, new_auth: Option) -> bool {