diff --git a/codex-rs/core/src/client.rs b/codex-rs/core/src/client.rs index 189816fef1..236be5c56f 100644 --- a/codex-rs/core/src/client.rs +++ b/codex-rs/core/src/client.rs @@ -276,7 +276,7 @@ struct LastResponse { #[derive(Debug, Default)] struct WebsocketSession { connection: Option, - connection_auth: Option, + connection_auth_fingerprint: Option, last_request: Option, last_response_rx: Option>, last_response_from_untraced_warmup: bool, @@ -1041,7 +1041,7 @@ impl ModelClientSession { fn reset_websocket_session(&mut self) { self.websocket_session.connection = None; - self.websocket_session.connection_auth = None; + self.websocket_session.connection_auth_fingerprint = None; self.websocket_session.last_request = None; self.websocket_session.last_response_rx = None; self.websocket_session.last_response_from_untraced_warmup = false; @@ -1194,7 +1194,7 @@ impl ModelClientSession { client_setup.api_auth.as_ref(), PendingUnauthorizedRetry::default(), ); - let connection_auth = client_setup + let connection_auth_fingerprint = client_setup .auth .as_ref() .and_then(AuthFingerprint::from_auth); @@ -1210,7 +1210,7 @@ impl ModelClientSession { ) .await?; self.websocket_session.connection = Some(connection); - self.websocket_session.connection_auth = connection_auth; + self.websocket_session.connection_auth_fingerprint = connection_auth_fingerprint; self.websocket_session .set_connection_reused(/*connection_reused*/ false); Ok(()) @@ -1250,7 +1250,7 @@ impl ModelClientSession { self.websocket_session.last_request = None; self.websocket_session.last_response_rx = None; self.websocket_session.last_response_from_untraced_warmup = false; - self.websocket_session.connection_auth = None; + self.websocket_session.connection_auth_fingerprint = None; let new_conn = match self .client .connect_websocket( @@ -1272,7 +1272,7 @@ impl ModelClientSession { } }; self.websocket_session.connection = Some(new_conn); - self.websocket_session.connection_auth = auth; + self.websocket_session.connection_auth_fingerprint = auth; self.websocket_session .set_connection_reused(/*connection_reused*/ false); } else { @@ -1389,7 +1389,7 @@ impl ModelClientSession { Arc::clone(&self.client.state.provider), Some(StreamAuthRecovery { auth_recovery, - rejected_auth: client_setup + rejected_auth_fingerprint: client_setup .auth .as_ref() .and_then(AuthFingerprint::from_auth), @@ -1407,7 +1407,7 @@ impl ModelClientSession { response_debug_context.request_id.as_deref(), /*output_items*/ &[], ); - let rejected_auth = client_setup + let rejected_auth_fingerprint = client_setup .auth .as_ref() .and_then(AuthFingerprint::from_auth); @@ -1417,7 +1417,7 @@ impl ModelClientSession { &mut auth_recovery, session_telemetry, &self.client.state.provider, - rejected_auth.as_ref(), + rejected_auth_fingerprint.as_ref(), ) .await?, ); @@ -1537,7 +1537,7 @@ impl ModelClientSession { Err(ApiError::Transport( unauthorized_transport @ TransportError::Http { status, .. }, )) if status == StatusCode::UNAUTHORIZED => { - let rejected_auth = client_setup + let rejected_auth_fingerprint = client_setup .auth .as_ref() .and_then(AuthFingerprint::from_auth); @@ -1547,7 +1547,7 @@ impl ModelClientSession { &mut auth_recovery, session_telemetry, &self.client.state.provider, - rejected_auth.as_ref(), + rejected_auth_fingerprint.as_ref(), ) .await?, ); @@ -1604,12 +1604,16 @@ impl ModelClientSession { ); err })?; - let rejected_auth = self.websocket_session.connection_auth.clone().or_else(|| { - client_setup - .auth - .as_ref() - .and_then(AuthFingerprint::from_auth) - }); + let rejected_auth_fingerprint = self + .websocket_session + .connection_auth_fingerprint + .clone() + .or_else(|| { + client_setup + .auth + .as_ref() + .and_then(AuthFingerprint::from_auth) + }); let (stream, last_request_rx) = map_response_stream( stream_result, request_session_telemetry, @@ -1617,7 +1621,7 @@ impl ModelClientSession { Arc::clone(&self.client.state.provider), Some(StreamAuthRecovery { auth_recovery, - rejected_auth, + rejected_auth_fingerprint, }), ); self.websocket_session.last_response_rx = Some(last_request_rx); @@ -1999,7 +2003,7 @@ where &transport, &mut stream_auth_recovery.auth_recovery, &session_telemetry, - stream_auth_recovery.rejected_auth.as_ref(), + stream_auth_recovery.rejected_auth_fingerprint.as_ref(), ) .await } else { @@ -2115,14 +2119,14 @@ struct WebsocketConnectParams<'a> { struct StreamAuthRecovery { auth_recovery: Option, - rejected_auth: Option, + rejected_auth_fingerprint: Option, } async fn force_logout_if_workspace_restricted_unauthorized( transport: &TransportError, auth_recovery: &mut Option, session_telemetry: &SessionTelemetry, - rejected_auth: Option<&AuthFingerprint>, + rejected_auth_fingerprint: Option<&AuthFingerprint>, ) -> Option { let debug = extract_response_debug_context(transport); if !is_chatgpt_ip_workspace_restricted_unauthorized(transport, &debug) { @@ -2144,7 +2148,7 @@ async fn force_logout_if_workspace_restricted_unauthorized( .unwrap_or(CHATGPT_IP_WORKSPACE_RESTRICTED_ERROR_CODE), ); let auth_state_changed = match recovery - .force_logout_due_to_server_auth_rejection(rejected_auth) + .force_logout_due_to_server_auth_rejection(rejected_auth_fingerprint) .await { Ok(changed) => Some(changed), @@ -2186,13 +2190,13 @@ async fn handle_unauthorized( auth_recovery: &mut Option, session_telemetry: &SessionTelemetry, provider: &SharedModelProvider, - rejected_auth: Option<&AuthFingerprint>, + rejected_auth_fingerprint: Option<&AuthFingerprint>, ) -> Result { if let Some(err) = force_logout_if_workspace_restricted_unauthorized( &transport, auth_recovery, session_telemetry, - rejected_auth, + rejected_auth_fingerprint, ) .await { diff --git a/codex-rs/core/src/client_tests.rs b/codex-rs/core/src/client_tests.rs index 571d595f3e..7045c0e5b1 100644 --- a/codex-rs/core/src/client_tests.rs +++ b/codex-rs/core/src/client_tests.rs @@ -13,25 +13,13 @@ use crate::GenerateAttestationFuture; use crate::responses_metadata::CodexResponsesMetadata; use crate::test_support::TestCodexResponsesRequestKind; use crate::test_support::responses_metadata as test_responses_metadata; -use base64::Engine; use codex_api::ApiError; use codex_api::ResponseEvent; use codex_api::TransportError; use codex_app_server_protocol::AuthMode; -use codex_config::types::AuthCredentialsStoreMode; -use codex_login::AuthDotJson; use codex_login::AuthFingerprint; -use codex_login::AuthKeyringBackendKind; use codex_login::AuthManager; use codex_login::CodexAuth; -use codex_login::ExternalAuth; -use codex_login::ExternalAuthFuture; -use codex_login::ExternalAuthRefreshContext; -use codex_login::ExternalAuthTokens; -use codex_login::TokenData; -use codex_login::load_auth_dot_json; -use codex_login::save_auth; -use codex_login::token_data::IdTokenInfo; use codex_model_provider::BearerAuthProvider; use codex_model_provider::SharedModelProvider; use codex_model_provider::create_model_provider; @@ -167,45 +155,6 @@ fn test_session_telemetry() -> SessionTelemetry { ) } -fn test_auth_dot_json() -> AuthDotJson { - test_auth_dot_json_with_account("account_id") -} - -fn test_auth_dot_json_with_account(account_id: &str) -> AuthDotJson { - test_auth_dot_json_with_account_and_tokens(account_id, "Access Token", "test") -} - -fn test_auth_dot_json_with_account_and_tokens( - account_id: &str, - access_token: &str, - refresh_token: &str, -) -> AuthDotJson { - AuthDotJson { - auth_mode: Some(AuthMode::Chatgpt), - openai_api_key: None, - tokens: Some(TokenData { - id_token: IdTokenInfo { - chatgpt_account_id: Some(account_id.to_string()), - raw_jwt: fake_jwt_with_account(account_id), - ..Default::default() - }, - access_token: access_token.to_string(), - refresh_token: refresh_token.to_string(), - account_id: Some(account_id.to_string()), - }), - last_refresh: Some(chrono::Utc::now()), - agent_identity: None, - personal_access_token: None, - bedrock_api_key: None, - } -} - -fn test_chatgpt_auth_tokens_dot_json(account_id: &str, access_token: &str) -> AuthDotJson { - let mut auth = test_auth_dot_json_with_account_and_tokens(account_id, access_token, ""); - auth.auth_mode = Some(AuthMode::ChatgptAuthTokens); - auth -} - fn unauthorized_transport(body: &str) -> TransportError { unauthorized_transport_with_headers(body, /*headers*/ None) } @@ -222,47 +171,6 @@ fn unauthorized_transport_with_headers( } } -fn fake_jwt_with_account(account_id: &str) -> String { - let header = json!({ - "alg": "none", - "typ": "JWT", - }); - let payload = json!({ - "email": "user@example.com", - "https://api.openai.com/auth": { - "chatgpt_user_id": "user-12345", - "user_id": "user-12345", - "chatgpt_account_id": account_id, - }, - }); - let b64 = |value: serde_json::Value| { - base64::engine::general_purpose::URL_SAFE_NO_PAD - .encode(serde_json::to_vec(&value).expect("serialize jwt segment")) - }; - format!("{}.{}.{}", b64(header), b64(payload), b64(json!("sig"))) -} - -struct TestExternalAuth; - -impl ExternalAuth for TestExternalAuth { - fn auth_mode(&self) -> AuthMode { - AuthMode::ChatgptAuthTokens - } - - fn refresh( - &self, - _context: ExternalAuthRefreshContext, - ) -> ExternalAuthFuture<'_, ExternalAuthTokens> { - Box::pin(async { - Ok(ExternalAuthTokens::chatgpt( - "access-token", - "account_id", - /*chatgpt_plan_type*/ None, - )) - }) - } -} - #[derive(Default)] struct TagCollectorVisitor { tags: BTreeMap, @@ -681,10 +589,11 @@ fn auth_request_telemetry_context_tracks_attached_auth_and_retry_phase() { } #[tokio::test] -async fn workspace_restricted_401_preserves_external_auth_bridge() { - let auth_manager = - AuthManager::from_auth_for_testing(CodexAuth::create_dummy_chatgpt_auth_for_testing()); - auth_manager.set_external_auth(Arc::new(TestExternalAuth)); +async fn workspace_restricted_401_forces_chatgpt_logout_without_retry() { + let rejected_auth_snapshot = CodexAuth::create_dummy_chatgpt_auth_for_testing(); + let rejected_auth_fingerprint = AuthFingerprint::from_auth(&rejected_auth_snapshot) + .expect("ChatGPT auth can be fingerprinted"); + let auth_manager = AuthManager::from_auth_for_testing(rejected_auth_snapshot); let mut recovery = Some(auth_manager.unauthorized_recovery()); let err = super::handle_unauthorized( @@ -692,97 +601,7 @@ async fn workspace_restricted_401_preserves_external_auth_bridge() { &mut recovery, &test_session_telemetry(), &test_model_provider(), - /*rejected_auth*/ None, - ) - .await - .expect_err("matching 401 should force reauth without retry"); - - assert!(matches!(err, CodexErr::RefreshTokenFailed(_))); - assert!(auth_manager.auth_cached().is_none()); - assert!(auth_manager.has_external_auth()); -} - -#[tokio::test] -async fn workspace_restricted_401_preserves_newer_same_account_auth() -> anyhow::Result<()> { - let codex_home = TempDir::new()?; - save_auth( - codex_home.path(), - &test_auth_dot_json_with_account_and_tokens("account_id", "new-access-token", "new-test"), - AuthCredentialsStoreMode::File, - AuthKeyringBackendKind::default(), - )?; - let auth_manager = AuthManager::from_auth_for_testing_with_home( - CodexAuth::create_dummy_chatgpt_auth_for_testing(), - codex_home.path().to_path_buf(), - ); - let mut recovery = Some(auth_manager.unauthorized_recovery()); - - let err = super::handle_unauthorized( - unauthorized_transport(r#"{"error":{"code":"chatgpt_ip_workspace_restricted"}}"#), - &mut recovery, - &test_session_telemetry(), - &test_model_provider(), - /*rejected_auth*/ None, - ) - .await - .expect_err("matching 401 should force reauth without retry"); - - assert!(matches!(err, CodexErr::RefreshTokenFailed(_))); - let persisted_auth = load_auth_dot_json( - codex_home.path(), - AuthCredentialsStoreMode::File, - AuthKeyringBackendKind::default(), - )? - .expect("newer auth should not be removed"); - assert_eq!( - persisted_auth - .tokens - .map(|tokens| tokens.access_token) - .as_deref(), - Some("new-access-token") - ); - - Ok(()) -} - -#[tokio::test] -async fn workspace_restricted_401_uses_request_auth_for_forced_logout() -> anyhow::Result<()> { - let codex_home = TempDir::new()?; - save_auth( - codex_home.path(), - &test_auth_dot_json_with_account_and_tokens( - "account_id", - "request-access-token", - "request-refresh-token", - ), - AuthCredentialsStoreMode::File, - AuthKeyringBackendKind::default(), - )?; - let request_auth_manager = AuthManager::shared( - codex_home.path().to_path_buf(), - /*enable_codex_api_key_env*/ false, - AuthCredentialsStoreMode::File, - /*chatgpt_base_url*/ None, - AuthKeyringBackendKind::default(), - ) - .await; - let request_auth = request_auth_manager - .auth_cached() - .expect("request auth should load"); - let request_auth_fingerprint = - AuthFingerprint::from_auth(&request_auth).expect("request auth can be fingerprinted"); - let auth_manager = AuthManager::from_auth_for_testing_with_home( - CodexAuth::create_dummy_chatgpt_auth_for_testing(), - codex_home.path().to_path_buf(), - ); - let mut recovery = Some(auth_manager.unauthorized_recovery()); - - let err = super::handle_unauthorized( - unauthorized_transport(r#"{"error":{"code":"chatgpt_ip_workspace_restricted"}}"#), - &mut recovery, - &test_session_telemetry(), - &test_model_provider(), - Some(&request_auth_fingerprint), + Some(&rejected_auth_fingerprint), ) .await .expect_err("matching 401 should force reauth without retry"); @@ -796,94 +615,16 @@ async fn workspace_restricted_401_uses_request_auth_for_forced_logout() -> anyho } other => panic!("expected RefreshTokenFailed, got {other:?}"), } - assert!( - load_auth_dot_json( - codex_home.path(), - AuthCredentialsStoreMode::File, - AuthKeyringBackendKind::default(), - )? - .is_none() - ); - assert!(!recovery.as_ref().expect("recovery").has_next()); - - Ok(()) -} - -#[tokio::test] -async fn workspace_restricted_401_clears_stale_persistent_external_auth() -> anyhow::Result<()> { - let codex_home = TempDir::new()?; - save_auth( - codex_home.path(), - &test_auth_dot_json_with_account("old_persistent_account"), - AuthCredentialsStoreMode::File, - AuthKeyringBackendKind::default(), - )?; - save_auth( - codex_home.path(), - &test_chatgpt_auth_tokens_dot_json("account_id", "external-access-token"), - AuthCredentialsStoreMode::Ephemeral, - AuthKeyringBackendKind::default(), - )?; - let auth_manager = AuthManager::shared( - codex_home.path().to_path_buf(), - /*enable_codex_api_key_env*/ false, - AuthCredentialsStoreMode::File, - /*chatgpt_base_url*/ None, - AuthKeyringBackendKind::default(), - ) - .await; - assert!( - auth_manager - .auth_cached() - .is_some_and(|auth| { matches!(auth.api_auth_mode(), AuthMode::ChatgptAuthTokens) }) - ); - let mut recovery = Some(auth_manager.unauthorized_recovery()); - - let err = super::handle_unauthorized( - unauthorized_transport(r#"{"error":{"code":"chatgpt_ip_workspace_restricted"}}"#), - &mut recovery, - &test_session_telemetry(), - &test_model_provider(), - /*rejected_auth*/ None, - ) - .await - .expect_err("matching 401 should force reauth without retry"); - - assert!(matches!(err, CodexErr::RefreshTokenFailed(_))); assert!(auth_manager.auth_cached().is_none()); - assert!( - load_auth_dot_json( - codex_home.path(), - AuthCredentialsStoreMode::Ephemeral, - AuthKeyringBackendKind::default(), - )? - .is_none() - ); - assert!( - load_auth_dot_json( - codex_home.path(), - AuthCredentialsStoreMode::File, - AuthKeyringBackendKind::default(), - )? - .is_none() - ); - - Ok(()) + assert!(!recovery.as_ref().expect("recovery").has_next()); } #[tokio::test] -async fn workspace_restricted_stream_401_forces_chatgpt_logout() -> anyhow::Result<()> { - let codex_home = TempDir::new()?; - save_auth( - codex_home.path(), - &test_auth_dot_json(), - AuthCredentialsStoreMode::File, - AuthKeyringBackendKind::default(), - )?; - let auth_manager = AuthManager::from_auth_for_testing_with_home( - CodexAuth::create_dummy_chatgpt_auth_for_testing(), - codex_home.path().to_path_buf(), - ); +async fn workspace_restricted_stream_401_forces_chatgpt_logout() { + let rejected_auth_snapshot = CodexAuth::create_dummy_chatgpt_auth_for_testing(); + let rejected_auth_fingerprint = AuthFingerprint::from_auth(&rejected_auth_snapshot) + .expect("ChatGPT auth can be fingerprinted"); + let auth_manager = AuthManager::from_auth_for_testing(rejected_auth_snapshot); let api_stream = futures::stream::iter([Err(ApiError::Transport(unauthorized_transport( r#"{"type":"error","status":401,"error":{"code":"chatgpt_ip_workspace_restricted"}}"#, )))]); @@ -895,9 +636,7 @@ async fn workspace_restricted_stream_401_forces_chatgpt_logout() -> anyhow::Resu test_model_provider(), Some(super::StreamAuthRecovery { auth_recovery: Some(auth_manager.unauthorized_recovery()), - rejected_auth: AuthFingerprint::from_auth( - &CodexAuth::create_dummy_chatgpt_auth_for_testing(), - ), + rejected_auth_fingerprint: Some(rejected_auth_fingerprint), }), ); @@ -909,42 +648,26 @@ async fn workspace_restricted_stream_401_forces_chatgpt_logout() -> anyhow::Resu assert!(matches!(err, CodexErr::RefreshTokenFailed(_))); assert!(auth_manager.auth_cached().is_none()); - assert!( - load_auth_dot_json( - codex_home.path(), - AuthCredentialsStoreMode::File, - AuthKeyringBackendKind::default(), - )? - .is_none() - ); - - Ok(()) } #[tokio::test] async fn workspace_restricted_401_ignores_non_matching_cases() { - for body in [ - r#"{"error":{"code":"token_expired"}}"#, - "not json", - r#"{"error":{"message":"missing code"}}"#, - ] { - let auth_manager = - AuthManager::from_auth_for_testing(CodexAuth::create_dummy_chatgpt_auth_for_testing()); - let mut recovery = Some(auth_manager.unauthorized_recovery()); + let auth_manager = + AuthManager::from_auth_for_testing(CodexAuth::create_dummy_chatgpt_auth_for_testing()); + let mut recovery = Some(auth_manager.unauthorized_recovery()); - let err = super::handle_unauthorized( - unauthorized_transport(body), - &mut recovery, - &test_session_telemetry(), - &test_model_provider(), - /*rejected_auth*/ None, - ) - .await - .expect_err("non-matching 401 should fall back to normal recovery"); + let err = super::handle_unauthorized( + unauthorized_transport(r#"{"error":{"code":"token_expired"}}"#), + &mut recovery, + &test_session_telemetry(), + &test_model_provider(), + /*rejected_auth_fingerprint*/ None, + ) + .await + .expect_err("generic 401 should fall back to normal recovery"); - assert!(matches!(err, CodexErr::RefreshTokenFailed(_))); - assert!(auth_manager.auth_cached().is_some()); - } + assert!(matches!(err, CodexErr::RefreshTokenFailed(_))); + assert!(auth_manager.auth_cached().is_some()); let auth_manager = AuthManager::from_auth_for_testing(CodexAuth::from_api_key("sk-test")); let mut recovery = Some(auth_manager.unauthorized_recovery()); @@ -953,7 +676,7 @@ async fn workspace_restricted_401_ignores_non_matching_cases() { &mut recovery, &test_session_telemetry(), &test_model_provider(), - /*rejected_auth*/ None, + /*rejected_auth_fingerprint*/ None, ) .await .expect_err("api key auth should not use ChatGPT forced logout"); @@ -968,14 +691,23 @@ async fn workspace_restricted_401_ignores_non_matching_cases() { #[test] fn workspace_restricted_error_code_parser_accepts_body_and_header_shapes() { - for body in [ - r#"{"error":{"code":"chatgpt_ip_workspace_restricted"}}"#, - r#"{"code":"chatgpt_ip_workspace_restricted"}"#, + for (body, expected) in [ + ( + r#"{"error":{"code":"chatgpt_ip_workspace_restricted"}}"#, + true, + ), + (r#"{"code":"chatgpt_ip_workspace_restricted"}"#, true), + (r#"{"error":{"code":"token_expired"}}"#, false), + ("not json", false), + (r#"{"error":{"message":"missing code"}}"#, false), ] { - assert!(super::is_chatgpt_ip_workspace_restricted_unauthorized( - &unauthorized_transport(body), - &Default::default() - )); + assert_eq!( + super::is_chatgpt_ip_workspace_restricted_unauthorized( + &unauthorized_transport(body), + &Default::default() + ), + expected + ); } let mut headers = http::HeaderMap::new(); diff --git a/codex-rs/login/src/auth/auth_tests.rs b/codex-rs/login/src/auth/auth_tests.rs index bb99b75ca0..ac95d22be3 100644 --- a/codex-rs/login/src/auth/auth_tests.rs +++ b/codex-rs/login/src/auth/auth_tests.rs @@ -940,7 +940,7 @@ async fn unauthorized_recovery_reports_mode_and_step_names() { manager: Arc::clone(&manager), step: UnauthorizedRecoveryStep::Reload, expected_account_id: None, - rejected_auth: None, + rejected_auth_snapshot: None, mode: UnauthorizedRecoveryMode::Managed, }; assert_eq!(managed.mode_name(), "managed"); @@ -950,7 +950,7 @@ async fn unauthorized_recovery_reports_mode_and_step_names() { manager, step: UnauthorizedRecoveryStep::ExternalRefresh, expected_account_id: None, - rejected_auth: None, + rejected_auth_snapshot: None, mode: UnauthorizedRecoveryMode::External, }; assert_eq!(external.mode_name(), "external"); @@ -1608,6 +1608,60 @@ fn auth_fingerprint_debug_redacts_auth_material() { assert!(!debug.contains("account_id")); } +fn test_auth_dot_json_with_mode_and_tokens( + auth_mode: AuthMode, + account_id: &str, + access_token: &str, + refresh_token: &str, +) -> AuthDotJson { + let id_token = fake_jwt_for_auth_file_params(&AuthFileParams { + openai_api_key: None, + chatgpt_plan_type: None, + chatgpt_account_id: Some(account_id.to_string()), + }) + .expect("fake jwt should serialize"); + AuthDotJson { + auth_mode: Some(auth_mode), + openai_api_key: None, + tokens: Some(TokenData { + id_token: IdTokenInfo { + chatgpt_account_id: Some(account_id.to_string()), + raw_jwt: id_token, + ..Default::default() + }, + access_token: access_token.to_string(), + refresh_token: refresh_token.to_string(), + account_id: Some(account_id.to_string()), + }), + last_refresh: Some(Utc::now()), + agent_identity: None, + personal_access_token: None, + bedrock_api_key: None, + } +} + +fn test_chatgpt_auth_dot_json( + account_id: &str, + access_token: &str, + refresh_token: &str, +) -> AuthDotJson { + test_auth_dot_json_with_mode_and_tokens( + AuthMode::Chatgpt, + account_id, + access_token, + refresh_token, + ) +} + +fn test_chatgpt_auth_tokens_dot_json(account_id: &str, access_token: &str) -> AuthDotJson { + test_auth_dot_json_with_mode_and_tokens( + AuthMode::ChatgptAuthTokens, + account_id, + access_token, + "", + ) +} + #[tokio::test] #[serial(codex_auth_env)] async fn forced_server_auth_rejection_does_not_reload_access_token_env() { @@ -1625,11 +1679,11 @@ async fn forced_server_auth_rejection_does_not_reload_access_token_env() { .await; let _authapi_guard = EnvVarGuard::set("CODEX_AUTHAPI_BASE_URL", &server.uri()); let _access_token_guard = EnvVarGuard::set(CODEX_ACCESS_TOKEN_ENV_VAR, "at-env-test"); - let rejected_auth = CodexAuth::create_dummy_chatgpt_auth_for_testing(); - let rejected_auth_fingerprint = - AuthFingerprint::from_auth(&rejected_auth).expect("ChatGPT auth can be fingerprinted"); + let rejected_auth_snapshot = CodexAuth::create_dummy_chatgpt_auth_for_testing(); + let rejected_auth_fingerprint = AuthFingerprint::from_auth(&rejected_auth_snapshot) + .expect("ChatGPT auth can be fingerprinted"); let manager = AuthManager::from_auth_for_testing_with_home( - rejected_auth.clone(), + rejected_auth_snapshot.clone(), codex_home.path().to_path_buf(), ); @@ -1643,6 +1697,101 @@ async fn forced_server_auth_rejection_does_not_reload_access_token_env() { server.verify().await; } +#[tokio::test] +async fn forced_server_auth_rejection_preserves_newer_same_account_auth() { + let codex_home = tempdir().unwrap(); + save_auth( + codex_home.path(), + &test_chatgpt_auth_dot_json("account_id", "new-access-token", "new-refresh-token"), + AuthCredentialsStoreMode::File, + AuthKeyringBackendKind::default(), + ) + .expect("seed newer auth"); + let rejected_auth_snapshot = CodexAuth::create_dummy_chatgpt_auth_for_testing(); + let rejected_auth_fingerprint = AuthFingerprint::from_auth(&rejected_auth_snapshot) + .expect("ChatGPT auth can be fingerprinted"); + let manager = AuthManager::from_auth_for_testing_with_home( + rejected_auth_snapshot, + codex_home.path().to_path_buf(), + ); + + manager + .force_logout_due_to_server_auth_rejection(Some(&rejected_auth_fingerprint)) + .await + .expect("forced logout should succeed"); + + let persisted_auth = load_auth_dot_json( + codex_home.path(), + AuthCredentialsStoreMode::File, + AuthKeyringBackendKind::default(), + ) + .expect("load auth") + .expect("newer auth should not be removed"); + assert_eq!( + persisted_auth + .tokens + .map(|tokens| tokens.access_token) + .as_deref(), + Some("new-access-token") + ); +} + +#[tokio::test] +async fn forced_server_auth_rejection_clears_stale_persistent_external_auth() { + let codex_home = tempdir().unwrap(); + save_auth( + codex_home.path(), + &test_chatgpt_auth_dot_json("old_persistent_account", "old-access-token", "old-refresh"), + AuthCredentialsStoreMode::File, + AuthKeyringBackendKind::default(), + ) + .expect("seed stale persistent auth"); + save_auth( + codex_home.path(), + &test_chatgpt_auth_tokens_dot_json("account_id", "external-access-token"), + AuthCredentialsStoreMode::Ephemeral, + AuthKeyringBackendKind::default(), + ) + .expect("seed ephemeral external auth"); + let manager = AuthManager::shared( + codex_home.path().to_path_buf(), + /*enable_codex_api_key_env*/ false, + AuthCredentialsStoreMode::File, + /*chatgpt_base_url*/ None, + AuthKeyringBackendKind::default(), + ) + .await; + let rejected_auth_fingerprint = + AuthFingerprint::from_auth(&manager.auth_cached().expect("external auth should load")) + .expect("ChatGPT auth can be fingerprinted"); + + let changed = manager + .force_logout_due_to_server_auth_rejection(Some(&rejected_auth_fingerprint)) + .await + .expect("forced logout should succeed"); + + assert!(changed); + assert!(manager.auth_cached().is_none()); + assert!( + load_auth_dot_json( + codex_home.path(), + AuthCredentialsStoreMode::Ephemeral, + AuthKeyringBackendKind::default(), + ) + .expect("load ephemeral auth") + .is_none() + ); + assert!( + load_auth_dot_json( + codex_home.path(), + AuthCredentialsStoreMode::File, + AuthKeyringBackendKind::default(), + ) + .expect("load file auth") + .is_none() + ); +} + #[tokio::test] #[serial(codex_auth_env)] async fn load_auth_keeps_codex_api_key_env_precedence() { diff --git a/codex-rs/login/src/auth/manager.rs b/codex-rs/login/src/auth/manager.rs index 629e7bd238..3b388b9b99 100644 --- a/codex-rs/login/src/auth/manager.rs +++ b/codex-rs/login/src/auth/manager.rs @@ -1324,17 +1324,17 @@ fn logout_all_stores( Ok(removed_ephemeral || removed_managed) } -fn logout_store_matching_rejected_auth( +fn logout_store_matching_rejected_auth_fingerprint( codex_home: &Path, auth_credentials_store_mode: AuthCredentialsStoreMode, keyring_backend_kind: AuthKeyringBackendKind, - rejected_auth: &AuthFingerprint, + rejected_auth_fingerprint: &AuthFingerprint, ) -> std::io::Result { logout_store_if( codex_home, auth_credentials_store_mode, keyring_backend_kind, - |auth_dot_json| rejected_auth.matches_auth_dot_json(auth_dot_json), + |auth_dot_json| rejected_auth_fingerprint.matches_auth_dot_json(auth_dot_json), ) } @@ -1848,7 +1848,7 @@ pub struct UnauthorizedRecovery { manager: Arc, step: UnauthorizedRecoveryStep, expected_account_id: Option, - rejected_auth: Option, + rejected_auth_snapshot: Option, mode: UnauthorizedRecoveryMode, } @@ -1884,7 +1884,7 @@ impl UnauthorizedRecovery { manager, step, expected_account_id, - rejected_auth: cached_auth, + rejected_auth_snapshot: cached_auth, mode, } } @@ -1970,16 +1970,17 @@ impl UnauthorizedRecovery { pub async fn force_logout_due_to_server_auth_rejection( &mut self, - rejected_auth: Option<&AuthFingerprint>, + rejected_auth_fingerprint: Option<&AuthFingerprint>, ) -> std::io::Result { self.step = UnauthorizedRecoveryStep::Done; - let fallback_rejected_auth = self - .rejected_auth + let fallback_rejected_auth_fingerprint = self + .rejected_auth_snapshot .as_ref() .and_then(AuthFingerprint::from_auth); - let rejected_auth = rejected_auth.or(fallback_rejected_auth.as_ref()); + let rejected_auth_fingerprint = + rejected_auth_fingerprint.or(fallback_rejected_auth_fingerprint.as_ref()); self.manager - .force_logout_due_to_server_auth_rejection(rejected_auth) + .force_logout_due_to_server_auth_rejection(rejected_auth_fingerprint) .await } @@ -2742,13 +2743,14 @@ impl AuthManager { pub async fn force_logout_due_to_server_auth_rejection( &self, - rejected_auth: Option<&AuthFingerprint>, + rejected_auth_fingerprint: Option<&AuthFingerprint>, ) -> std::io::Result { if !self.current_auth_uses_codex_backend() { return Ok(false); } - let removal_result = self.logout_stores_matching_rejected_auth(rejected_auth); + let removal_result = + self.logout_stores_matching_rejected_auth_fingerprint(rejected_auth_fingerprint); let cache_changed = self.set_cached_auth( self.load_auth_from_storage_without_codex_access_token_env() .await, @@ -2757,11 +2759,11 @@ impl AuthManager { Ok(removed || cache_changed) } - fn logout_stores_matching_rejected_auth( + fn logout_stores_matching_rejected_auth_fingerprint( &self, - rejected_auth: Option<&AuthFingerprint>, + rejected_auth_fingerprint: Option<&AuthFingerprint>, ) -> std::io::Result { - let Some(rejected_auth) = rejected_auth else { + let Some(rejected_auth_fingerprint) = rejected_auth_fingerprint else { return logout_all_stores( &self.codex_home, self.auth_credentials_store_mode, @@ -2769,12 +2771,12 @@ impl AuthManager { ); }; - if rejected_auth.is_external_chatgpt_tokens() { - let removed_ephemeral = logout_store_matching_rejected_auth( + if rejected_auth_fingerprint.is_external_chatgpt_tokens() { + let removed_ephemeral = logout_store_matching_rejected_auth_fingerprint( &self.codex_home, AuthCredentialsStoreMode::Ephemeral, AuthKeyringBackendKind::default(), - rejected_auth, + rejected_auth_fingerprint, )?; if self.auth_credentials_store_mode == AuthCredentialsStoreMode::Ephemeral { return Ok(removed_ephemeral); @@ -2800,35 +2802,35 @@ impl AuthManager { } if self.auth_credentials_store_mode == AuthCredentialsStoreMode::Ephemeral { - return logout_store_matching_rejected_auth( + return logout_store_matching_rejected_auth_fingerprint( &self.codex_home, AuthCredentialsStoreMode::Ephemeral, AuthKeyringBackendKind::default(), - rejected_auth, + rejected_auth_fingerprint, ); } - let removed_ephemeral = logout_store_matching_rejected_auth( + let removed_ephemeral = logout_store_matching_rejected_auth_fingerprint( &self.codex_home, AuthCredentialsStoreMode::Ephemeral, AuthKeyringBackendKind::default(), - rejected_auth, + rejected_auth_fingerprint, )?; let removed_file = if self.auth_credentials_store_mode == AuthCredentialsStoreMode::File { false } else { - logout_store_matching_rejected_auth( + logout_store_matching_rejected_auth_fingerprint( &self.codex_home, AuthCredentialsStoreMode::File, AuthKeyringBackendKind::default(), - rejected_auth, + rejected_auth_fingerprint, )? }; - let removed_managed = logout_store_matching_rejected_auth( + let removed_managed = logout_store_matching_rejected_auth_fingerprint( &self.codex_home, self.auth_credentials_store_mode, self.keyring_backend_kind, - rejected_auth, + rejected_auth_fingerprint, )?; Ok(removed_ephemeral || removed_file || removed_managed) }