diff --git a/codex-rs/core/src/client.rs b/codex-rs/core/src/client.rs index 39513b3a27..f51f2ced1e 100644 --- a/codex-rs/core/src/client.rs +++ b/codex-rs/core/src/client.rs @@ -2053,7 +2053,7 @@ async fn handle_unauthorized( provider: &SharedModelProvider, ) -> Result { let debug = extract_response_debug_context(&transport); - if is_chatgpt_ip_workspace_restricted_unauthorized(&transport) + if is_chatgpt_ip_workspace_restricted_unauthorized(&transport, &debug) && let Some(recovery) = auth_recovery && recovery.current_auth_uses_codex_backend() { @@ -2210,13 +2210,19 @@ async fn handle_unauthorized( Err(provider.map_api_error(ApiError::Transport(transport))) } -fn is_chatgpt_ip_workspace_restricted_unauthorized(transport: &TransportError) -> bool { +fn is_chatgpt_ip_workspace_restricted_unauthorized( + transport: &TransportError, + debug: &codex_response_debug_context::ResponseDebugContext, +) -> bool { let TransportError::Http { status, body, .. } = transport else { return false; }; if *status != StatusCode::UNAUTHORIZED { return false; } + if debug.auth_error_code.as_deref() == Some(CHATGPT_IP_WORKSPACE_RESTRICTED_ERROR_CODE) { + return true; + } body.as_deref() .is_some_and(|body| body_has_error_code(body, CHATGPT_IP_WORKSPACE_RESTRICTED_ERROR_CODE)) } diff --git a/codex-rs/core/src/client_tests.rs b/codex-rs/core/src/client_tests.rs index 50df8c8d98..4cdac2b7de 100644 --- a/codex-rs/core/src/client_tests.rs +++ b/codex-rs/core/src/client_tests.rs @@ -13,6 +13,7 @@ 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; @@ -22,9 +23,14 @@ use codex_login::AuthDotJson; 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; @@ -161,14 +167,30 @@ 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: Default::default(), - access_token: "access-token".to_string(), - refresh_token: "refresh-token".to_string(), - account_id: Some("account_id".to_string()), + 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, @@ -178,14 +200,62 @@ fn test_auth_dot_json() -> AuthDotJson { } fn unauthorized_transport(body: &str) -> TransportError { + unauthorized_transport_with_headers(body, None) +} + +fn unauthorized_transport_with_headers( + body: &str, + headers: Option, +) -> TransportError { TransportError::Http { status: http::StatusCode::UNAUTHORIZED, url: Some("https://chatgpt.com/backend-api/codex/responses".to_string()), - headers: None, + headers, body: Some(body.to_string()), } } +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", + None, + )) + }) + } +} + #[derive(Default)] struct TagCollectorVisitor { tags: BTreeMap, @@ -667,6 +737,108 @@ async fn workspace_restricted_401_does_not_clear_api_key_auth() { ); } +#[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)); + 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(), + ) + .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_replacement_auth() -> anyhow::Result<()> { + let codex_home = TempDir::new()?; + save_auth( + codex_home.path(), + &test_auth_dot_json_with_account("replacement_account"), + 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(), + ) + .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("replacement auth should not be removed"); + assert_eq!( + persisted_auth + .tokens + .and_then(|tokens| tokens.account_id) + .as_deref(), + Some("replacement_account") + ); + + Ok(()) +} + +#[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(), + ) + .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 generic_or_malformed_401_uses_existing_recovery_path() { for body in [ @@ -698,11 +870,32 @@ fn workspace_restricted_error_code_parser_accepts_supported_shapes() { r#"{"code":"chatgpt_ip_workspace_restricted"}"#, ] { assert!(super::is_chatgpt_ip_workspace_restricted_unauthorized( - &unauthorized_transport(body) + &unauthorized_transport(body), + &Default::default() )); } } +#[test] +fn workspace_restricted_error_code_parser_accepts_x_error_json_header() { + let mut headers = http::HeaderMap::new(); + headers.insert( + "x-error-json", + http::HeaderValue::from_static( + "eyJlcnJvciI6eyJjb2RlIjoiY2hhdGdwdF9pcF93b3Jrc3BhY2VfcmVzdHJpY3RlZCJ9fQ==", + ), + ); + let transport = unauthorized_transport_with_headers( + r#"{"error":{"code":"different_error_code"}}"#, + Some(headers), + ); + let debug = codex_response_debug_context::extract_response_debug_context(&transport); + + assert!(super::is_chatgpt_ip_workspace_restricted_unauthorized( + &transport, &debug + )); +} + fn model_client_with_counting_attestation( include_attestation: bool, ) -> (ModelClient, Arc) { diff --git a/codex-rs/login/src/auth/auth_tests.rs b/codex-rs/login/src/auth/auth_tests.rs index bcce562e69..fb17d245a9 100644 --- a/codex-rs/login/src/auth/auth_tests.rs +++ b/codex-rs/login/src/auth/auth_tests.rs @@ -940,6 +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, mode: UnauthorizedRecoveryMode::Managed, }; assert_eq!(managed.mode_name(), "managed"); @@ -949,6 +950,7 @@ async fn unauthorized_recovery_reports_mode_and_step_names() { manager, step: UnauthorizedRecoveryStep::ExternalRefresh, expected_account_id: None, + rejected_auth: None, mode: UnauthorizedRecoveryMode::External, }; assert_eq!(external.mode_name(), "external"); diff --git a/codex-rs/login/src/auth/manager.rs b/codex-rs/login/src/auth/manager.rs index a95000702a..5b1d4c68f5 100644 --- a/codex-rs/login/src/auth/manager.rs +++ b/codex-rs/login/src/auth/manager.rs @@ -1149,6 +1149,84 @@ fn logout_all_stores( Ok(removed_ephemeral || removed_managed) } +fn logout_store_matching_rejected_auth( + codex_home: &Path, + auth_credentials_store_mode: AuthCredentialsStoreMode, + keyring_backend_kind: AuthKeyringBackendKind, + rejected_auth: &CodexAuth, +) -> std::io::Result { + let Some(auth_dot_json) = load_auth_dot_json( + codex_home, + auth_credentials_store_mode, + keyring_backend_kind, + )? + else { + return Ok(false); + }; + + if auth_dot_json_matches_rejected_auth(&auth_dot_json, rejected_auth) { + logout( + codex_home, + auth_credentials_store_mode, + keyring_backend_kind, + ) + } else { + Ok(false) + } +} + +fn auth_dot_json_matches_rejected_auth( + auth_dot_json: &AuthDotJson, + rejected_auth: &CodexAuth, +) -> bool { + if auth_dot_json.resolved_mode() != rejected_auth.api_auth_mode() { + return false; + } + + match rejected_auth { + CodexAuth::Chatgpt(_) | CodexAuth::ChatgptAuthTokens(_) => { + auth_dot_json_matches_rejected_chatgpt_auth(auth_dot_json, rejected_auth) + } + CodexAuth::AgentIdentity(rejected) => auth_dot_json + .agent_identity + .as_deref() + .and_then(|jwt| AgentIdentityAuthRecord::from_agent_identity_jwt(jwt).ok()) + .is_some_and(|record| record.account_id == rejected.account_id()), + CodexAuth::PersonalAccessToken(rejected) => { + auth_dot_json.personal_access_token.as_deref() == Some(rejected.access_token()) + } + CodexAuth::ApiKey(_) | CodexAuth::BedrockApiKey(_) => false, + } +} + +fn auth_dot_json_matches_rejected_chatgpt_auth( + auth_dot_json: &AuthDotJson, + rejected_auth: &CodexAuth, +) -> bool { + let Some(rejected_auth_dot_json) = rejected_auth.get_current_auth_json() else { + return false; + }; + let (Some(tokens), Some(rejected_tokens)) = + (&auth_dot_json.tokens, &rejected_auth_dot_json.tokens) + else { + return false; + }; + + tokens.access_token == rejected_tokens.access_token + && tokens.refresh_token == rejected_tokens.refresh_token + && auth_dot_json_chatgpt_account_id(auth_dot_json) + == auth_dot_json_chatgpt_account_id(&rejected_auth_dot_json) +} + +fn auth_dot_json_chatgpt_account_id(auth_dot_json: &AuthDotJson) -> Option<&str> { + auth_dot_json.tokens.as_ref().and_then(|tokens| { + tokens + .account_id + .as_deref() + .or(tokens.id_token.chatgpt_account_id.as_deref()) + }) +} + #[allow(clippy::too_many_arguments)] async fn load_auth( codex_home: &Path, @@ -1547,6 +1625,7 @@ pub struct UnauthorizedRecovery { manager: Arc, step: UnauthorizedRecoveryStep, expected_account_id: Option, + rejected_auth: Option, mode: UnauthorizedRecoveryMode, } @@ -1582,6 +1661,7 @@ impl UnauthorizedRecovery { manager, step, expected_account_id, + rejected_auth: cached_auth, mode, } } @@ -1668,7 +1748,7 @@ impl UnauthorizedRecovery { pub async fn force_logout_due_to_server_auth_rejection(&mut self) -> std::io::Result { self.step = UnauthorizedRecoveryStep::Done; self.manager - .force_logout_due_to_server_auth_rejection() + .force_logout_due_to_server_auth_rejection(self.rejected_auth.as_ref()) .await } @@ -2412,21 +2492,54 @@ impl AuthManager { Ok(result) } - pub async fn force_logout_due_to_server_auth_rejection(&self) -> std::io::Result { + pub async fn force_logout_due_to_server_auth_rejection( + &self, + rejected_auth: Option<&CodexAuth>, + ) -> std::io::Result { if !self.current_auth_uses_codex_backend() { return Ok(false); } - let removal_result = logout_all_stores( + let removal_result = self.logout_stores_matching_rejected_auth(rejected_auth); + let cache_changed = self.set_cached_auth(self.load_auth_from_storage().await); + let removed = removal_result?; + Ok(removed || cache_changed) + } + + fn logout_stores_matching_rejected_auth( + &self, + rejected_auth: Option<&CodexAuth>, + ) -> std::io::Result { + let Some(rejected_auth) = rejected_auth else { + return logout_all_stores( + &self.codex_home, + self.auth_credentials_store_mode, + self.keyring_backend_kind, + ); + }; + + if self.auth_credentials_store_mode == AuthCredentialsStoreMode::Ephemeral { + return logout_store_matching_rejected_auth( + &self.codex_home, + AuthCredentialsStoreMode::Ephemeral, + AuthKeyringBackendKind::default(), + rejected_auth, + ); + } + + let removed_ephemeral = logout_store_matching_rejected_auth( + &self.codex_home, + AuthCredentialsStoreMode::Ephemeral, + AuthKeyringBackendKind::default(), + rejected_auth, + )?; + let removed_managed = logout_store_matching_rejected_auth( &self.codex_home, self.auth_credentials_store_mode, self.keyring_backend_kind, - ); - let had_external_auth = self.has_external_auth(); - self.clear_external_auth(); - let cache_changed = self.set_cached_auth(None); - let removed = removal_result?; - Ok(removed || had_external_auth || cache_changed) + rejected_auth, + )?; + Ok(removed_ephemeral || removed_managed) } pub fn get_api_auth_mode(&self) -> Option {