diff --git a/codex-rs/app-server/tests/suite/v2/account.rs b/codex-rs/app-server/tests/suite/v2/account.rs index a15b46abd9..3f14d28379 100644 --- a/codex-rs/app-server/tests/suite/v2/account.rs +++ b/codex-rs/app-server/tests/suite/v2/account.rs @@ -446,7 +446,7 @@ async fn external_auth_refreshes_on_unauthorized() -> Result<()> { responses::ev_completed("resp-turn"), ]); let unauthorized = ResponseTemplate::new(401).set_body_json(json!({ - "error": { "message": "unauthorized" } + "error": { "message": "unauthorized", "type": "token_invalidated" } })); let responses_mock = responses::mount_response_sequence( &mock_server, diff --git a/codex-rs/core/src/client.rs b/codex-rs/core/src/client.rs index ef9f482c54..f89205c528 100644 --- a/codex-rs/core/src/client.rs +++ b/codex-rs/core/src/client.rs @@ -1987,31 +1987,58 @@ async fn handle_unauthorized( ); if let Some(recovery_reason) = revocation_recovery_reason && let Some(recovery) = auth_recovery.as_ref() + && recovery.handles_invalidated_access_token_auth() { let mode = recovery.mode_name(); let phase = recovery.step_name(); - let failed = recovery.clear_invalidated_access_token_auth().await; - session_telemetry.record_auth_recovery( - mode, - phase, - "recovery_not_run", - debug.request_id.as_deref(), - debug.cf_ray.as_deref(), - debug.auth_error.as_deref(), - debug.auth_error_code.as_deref(), - Some(recovery_reason), - /*auth_state_changed*/ None, - ); - emit_feedback_auth_recovery_tags( - mode, - phase, - "recovery_not_run", - debug.request_id.as_deref(), - debug.cf_ray.as_deref(), - debug.auth_error.as_deref(), - debug.auth_error_code.as_deref(), - ); - return Err(CodexErr::RefreshTokenFailed(failed)); + return match recovery.handle_invalidated_access_token_auth().await { + Ok(step_result) => { + session_telemetry.record_auth_recovery( + mode, + phase, + "recovery_succeeded", + debug.request_id.as_deref(), + debug.cf_ray.as_deref(), + debug.auth_error.as_deref(), + debug.auth_error_code.as_deref(), + Some(recovery_reason), + step_result.auth_state_changed(), + ); + emit_feedback_auth_recovery_tags( + mode, + phase, + "recovery_succeeded", + debug.request_id.as_deref(), + debug.cf_ray.as_deref(), + debug.auth_error.as_deref(), + debug.auth_error_code.as_deref(), + ); + Ok(UnauthorizedRecoveryExecution { mode, phase }) + } + Err(failed) => { + session_telemetry.record_auth_recovery( + mode, + phase, + "recovery_failed_permanent", + debug.request_id.as_deref(), + debug.cf_ray.as_deref(), + debug.auth_error.as_deref(), + debug.auth_error_code.as_deref(), + Some(recovery_reason), + /*auth_state_changed*/ None, + ); + emit_feedback_auth_recovery_tags( + mode, + phase, + "recovery_failed_permanent", + debug.request_id.as_deref(), + debug.cf_ray.as_deref(), + debug.auth_error.as_deref(), + debug.auth_error_code.as_deref(), + ); + Err(CodexErr::RefreshTokenFailed(failed)) + } + }; } if let Some(recovery) = auth_recovery diff --git a/codex-rs/core/src/client_tests.rs b/codex-rs/core/src/client_tests.rs index dbae792cd7..7f6b2640bf 100644 --- a/codex-rs/core/src/client_tests.rs +++ b/codex-rs/core/src/client_tests.rs @@ -17,7 +17,6 @@ use codex_api::ResponseEvent; use codex_api::TransportError; use codex_app_server_protocol::AuthMode; use codex_login::AuthManager; -use codex_login::CodexAuth; use codex_model_provider::BearerAuthProvider; use codex_model_provider_info::CHATGPT_CODEX_BASE_URL; use codex_model_provider_info::ModelProviderInfo; @@ -49,6 +48,7 @@ use pretty_assertions::assert_eq; use serde_json::json; use std::collections::BTreeMap; use std::collections::VecDeque; +use std::path::Path; use std::pin::Pin; use std::sync::Arc; use std::sync::Mutex; @@ -131,6 +131,53 @@ fn test_session_telemetry() -> SessionTelemetry { ) } +fn write_managed_chatgpt_auth(codex_home: &Path, access_token: &str) { + 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.join("auth.json"), + serde_json::to_vec_pretty(&auth_json).expect("managed auth should serialize"), + ) + .expect("managed auth should persist"); +} + +async fn managed_chatgpt_auth_manager(access_token: &str) -> (TempDir, Arc) { + let codex_home = TempDir::new().expect("managed auth tempdir"); + write_managed_chatgpt_auth(codex_home.path(), access_token); + 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; + (codex_home, manager) +} + #[derive(Default)] struct TagCollectorVisitor { tags: BTreeMap, @@ -592,8 +639,7 @@ async fn non_chatgpt_codex_endpoints_omit_attestation_generation() { #[tokio::test] async fn token_invalidated_401_clears_auth_and_requires_relogin() { - let manager = - AuthManager::from_auth_for_testing(CodexAuth::create_dummy_chatgpt_auth_for_testing()); + let (_codex_home, manager) = managed_chatgpt_auth_manager("revoked-access-token").await; let mut recovery = Some(manager.unauthorized_recovery()); let x_error_json = base64::engine::general_purpose::STANDARD .encode(r#"{"error":{"code":"token_invalidated"}}"#); @@ -636,8 +682,7 @@ async fn token_invalidated_401_clears_auth_and_requires_relogin() { #[tokio::test] async fn token_invalidated_error_type_401_clears_auth_and_requires_relogin() { - let manager = - AuthManager::from_auth_for_testing(CodexAuth::create_dummy_chatgpt_auth_for_testing()); + let (_codex_home, manager) = managed_chatgpt_auth_manager("revoked-access-token").await; let mut recovery = Some(manager.unauthorized_recovery()); let err = handle_unauthorized( @@ -670,3 +715,34 @@ async fn token_invalidated_error_type_401_clears_auth_and_requires_relogin() { "reload" ); } + +#[tokio::test] +async fn token_invalidated_401_retries_when_persisted_auth_changed() { + let (codex_home, manager) = managed_chatgpt_auth_manager("revoked-access-token").await; + write_managed_chatgpt_auth(codex_home.path(), "replacement-access-token"); + let mut recovery = Some(manager.unauthorized_recovery()); + + let retry = 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_invalidated"}}"#.to_string()), + }, + &mut recovery, + &test_session_telemetry(), + ) + .await + .expect("new persisted auth should retry instead of logging out"); + + assert_eq!(retry.mode, "managed"); + assert_eq!(retry.phase, "reload"); + assert_eq!( + manager + .auth_cached() + .expect("replacement auth should remain cached") + .get_token() + .expect("replacement token should resolve"), + "replacement-access-token" + ); +} diff --git a/codex-rs/login/src/auth/auth_tests.rs b/codex-rs/login/src/auth/auth_tests.rs index 5097d0a8a9..b29be8c735 100644 --- a/codex-rs/login/src/auth/auth_tests.rs +++ b/codex-rs/login/src/auth/auth_tests.rs @@ -326,11 +326,30 @@ async fn unauthorized_recovery_reports_mode_and_step_names() { #[tokio::test] async fn invalidated_access_token_logout_clears_cached_auth() { - let manager = - AuthManager::from_auth_for_testing(CodexAuth::create_dummy_chatgpt_auth_for_testing()); + let codex_home = tempdir().unwrap(); + 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 recovery = manager.unauthorized_recovery(); - let failed = recovery.clear_invalidated_access_token_auth().await; + assert!(recovery.handles_invalidated_access_token_auth()); + let failed = recovery + .handle_invalidated_access_token_auth() + .await + .expect_err("unchanged revoked auth should force login"); assert_eq!(failed.reason, RefreshTokenFailedReason::Revoked); assert_eq!( @@ -338,6 +357,60 @@ async fn invalidated_access_token_logout_clears_cached_auth() { "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] +async fn invalidated_access_token_preserves_reloaded_auth() { + let codex_home = tempdir().unwrap(); + 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 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 step_result = recovery + .handle_invalidated_access_token_auth() + .await + .expect("new persisted auth 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] @@ -469,6 +542,7 @@ async fn unauthorized_recovery_uses_external_refresh_for_bearer_manager() { assert!(recovery.has_next()); assert_eq!(recovery.mode_name(), "external"); assert_eq!(recovery.step_name(), "external_refresh"); + assert!(!recovery.handles_invalidated_access_token_auth()); let result = recovery .next() @@ -626,7 +700,8 @@ fn write_auth_file(params: AuthFileParams, codex_home: &Path) -> std::io::Result "tokens": { "id_token": fake_jwt, "access_token": "test-access-token", - "refresh_token": "test-refresh-token" + "refresh_token": "test-refresh-token", + "account_id": params.chatgpt_account_id, }, "last_refresh": Utc::now(), }); diff --git a/codex-rs/login/src/auth/manager.rs b/codex-rs/login/src/auth/manager.rs index 95a547e9e9..ef1ef8f12e 100644 --- a/codex-rs/login/src/auth/manager.rs +++ b/codex-rs/login/src/auth/manager.rs @@ -1066,8 +1066,9 @@ enum UnauthorizedRecoveryMode { } // UnauthorizedRecovery is a state machine that handles an attempt to refresh the authentication when requests -// to API fail with a refreshable 401 status code. Callers must leave explicit access-token revocation signals -// outside this flow so those sessions can be cleared instead of refreshed. +// to API fail with a refreshable 401 status code. Managed ChatGPT access-token revocation is handled beside +// this flow so persisted credentials can be guarded before they are cleared; externally owned auth stays here +// so its owner can supply replacement credentials. // The client calls next() every time it encounters a 401 error, one time per retry. // For API key based authentication, we don't do anything and let the error bubble to the user. // @@ -1193,6 +1194,15 @@ impl UnauthorizedRecovery { } } + pub fn handles_invalidated_access_token_auth(&self) -> bool { + self.mode == UnauthorizedRecoveryMode::Managed + && self + .manager + .auth_cached() + .as_ref() + .is_some_and(|auth| matches!(auth, CodexAuth::Chatgpt(_))) + } + pub async fn next(&mut self) -> Result { if !self.has_next() { return Err(RefreshTokenError::Permanent(RefreshTokenFailedError::new( @@ -1252,14 +1262,34 @@ impl UnauthorizedRecovery { }) } - pub async fn clear_invalidated_access_token_auth(&self) -> RefreshTokenFailedError { - 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}" - ), - }; - RefreshTokenFailedError::new(RefreshTokenFailedReason::Revoked, message) + pub async fn handle_invalidated_access_token_auth( + &self, + ) -> Result { + match self + .manager + .reload_if_account_id_matches(self.expected_account_id.as_deref()) + .await + { + ReloadOutcome::ReloadedChanged => 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}" + ), + }; + Err(RefreshTokenFailedError::new( + RefreshTokenFailedReason::Revoked, + message, + )) + } + ReloadOutcome::Skipped => Err(RefreshTokenFailedError::new( + RefreshTokenFailedReason::Other, + REFRESH_TOKEN_ACCOUNT_MISMATCH_MESSAGE.to_string(), + )), + } } }