address review comments

'

codex-rs/core/src/client.rs

@@ -2053,7 +2053,7 @@ async fn handle_unauthorized(
     provider: &SharedModelProvider,
 ) -> Result<UnauthorizedRecoveryExecution> {
     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))
 }

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<http::HeaderMap>,
+) -> 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<String, String>,
@@ -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<AtomicUsize>) {

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");

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<bool> {
+    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<AuthManager>,
     step: UnauthorizedRecoveryStep,
     expected_account_id: Option<String>,
+    rejected_auth: Option<CodexAuth>,
     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<bool> {
         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<bool> {
+    pub async fn force_logout_due_to_server_auth_rejection(
+        &self,
+        rejected_auth: Option<&CodexAuth>,
+    ) -> std::io::Result<bool> {
         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<bool> {
+        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<ApiAuthMode> {