Code review feedback

codex-rs/core/src/auth.rs

@@ -1249,6 +1249,22 @@ pub(crate) enum SyncFromStorageResult {
     Applied { changed: bool },
 }
 
+/// Per-request state used to coordinate a limited 401 recovery flow.
+///
+/// This is intentionally managed by `AuthManager` so callers don't need to
+/// understand the recovery stages.
+#[derive(Default, Debug)]
+pub(crate) struct UnauthorizedRecovery {
+    synced_from_storage: bool,
+    refreshed_token: bool,
+}
+
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub(crate) enum UnauthorizedRecoveryDecision {
+    Retry,
+    GiveUp,
+}
+
 impl AuthManager {
     /// Create a new manager loading the initial auth using the provided
     /// preferred auth method. Errors loading auth are swallowed; `auth()` will
@@ -1523,6 +1539,48 @@ impl AuthManager {
             }
         }
     }
+
+    pub(crate) async fn recover_from_unauthorized_for_request(
+        &self,
+        expected: &CodexAuth,
+        recovery: &mut UnauthorizedRecovery,
+    ) -> Result<UnauthorizedRecoveryDecision, RefreshTokenError> {
+        if recovery.refreshed_token {
+            return Ok(UnauthorizedRecoveryDecision::GiveUp);
+        }
+
+        if expected.mode != AuthMode::ChatGPT {
+            return Ok(UnauthorizedRecoveryDecision::GiveUp);
+        }
+
+        if !recovery.synced_from_storage {
+            let sync = self
+                .sync_from_storage_for_request(expected)
+                .await
+                .map_err(RefreshTokenError::Transient)?;
+            recovery.synced_from_storage = true;
+            match sync {
+                SyncFromStorageResult::Applied { changed } => {
+                    tracing::debug!(changed, "synced ChatGPT credentials from storage after 401");
+                    Ok(UnauthorizedRecoveryDecision::Retry)
+                }
+                SyncFromStorageResult::SkippedMissingIdentity => {
+                    Ok(UnauthorizedRecoveryDecision::Retry)
+                }
+                SyncFromStorageResult::LoggedOut | SyncFromStorageResult::IdentityMismatch => {
+                    Ok(UnauthorizedRecoveryDecision::GiveUp)
+                }
+            }
+        } else {
+            match self.refresh_token_for_request(expected).await? {
+                Some(_) => {
+                    recovery.refreshed_token = true;
+                    Ok(UnauthorizedRecoveryDecision::Retry)
+                }
+                None => Ok(UnauthorizedRecoveryDecision::GiveUp),
+            }
+        }
+    }
 }
 
 async fn load_stored_refresh_token_if_identity_matches(

codex-rs/core/src/client.rs

@@ -17,7 +17,6 @@ use codex_api::TransportError;
 use codex_api::common::Reasoning;
 use codex_api::create_text_param_for_request;
 use codex_api::error::ApiError;
-use codex_app_server_protocol::AuthMode;
 use codex_otel::otel_manager::OtelManager;
 use codex_protocol::ConversationId;
 use codex_protocol::config_types::ReasoningSummary as ReasoningSummaryConfig;
@@ -153,7 +152,7 @@ impl ModelClient {
         let conversation_id = self.conversation_id.to_string();
         let session_source = self.session_source.clone();
 
-        let mut recovery = UnauthorizedRecovery::default();
+        let mut recovery = crate::auth::UnauthorizedRecovery::default();
         loop {
             let auth = auth_manager.as_ref().and_then(|m| m.auth());
             let api_provider = self
@@ -242,7 +241,7 @@ impl ModelClient {
         let conversation_id = self.conversation_id.to_string();
         let session_source = self.session_source.clone();
 
-        let mut recovery = UnauthorizedRecovery::default();
+        let mut recovery = crate::auth::UnauthorizedRecovery::default();
         loop {
             let auth = auth_manager.as_ref().and_then(|m| m.auth());
             let api_provider = self
@@ -485,59 +484,29 @@ where
 /// the mapped `CodexErr` is returned to the caller.
 async fn handle_unauthorized(
     status: StatusCode,
-    recovery: &mut UnauthorizedRecovery,
+    recovery: &mut crate::auth::UnauthorizedRecovery,
     auth_manager: &Option<Arc<AuthManager>>,
     auth: &Option<crate::auth::CodexAuth>,
 ) -> Result<()> {
-    if recovery.refreshed_token {
-        return Err(map_unauthorized_status(status));
-    }
-
     if let Some(manager) = auth_manager.as_ref()
         && let Some(auth) = auth.as_ref()
-        && auth.mode == AuthMode::ChatGPT
     {
-        if !recovery.synced_from_storage {
-            let sync = manager
-                .sync_from_storage_for_request(auth)
-                .await
-                .map_err(CodexErr::Io)?;
-            recovery.synced_from_storage = true;
-            match sync {
-                crate::auth::SyncFromStorageResult::Applied { changed } => {
-                    tracing::debug!(changed, "synced ChatGPT credentials from storage after 401");
-                    Ok(())
-                }
-                crate::auth::SyncFromStorageResult::SkippedMissingIdentity => Ok(()),
-                crate::auth::SyncFromStorageResult::LoggedOut
-                | crate::auth::SyncFromStorageResult::IdentityMismatch => {
-                    Err(map_unauthorized_status(status))
-                }
-            }
-        } else {
-            match manager.refresh_token_for_request(auth).await {
-                Ok(Some(_)) => {
-                    recovery.refreshed_token = true;
-                    Ok(())
-                }
-                Ok(None) => Err(map_unauthorized_status(status)),
-                Err(RefreshTokenError::Permanent(failed)) => {
-                    Err(CodexErr::RefreshTokenFailed(failed))
-                }
-                Err(RefreshTokenError::Transient(other)) => Err(CodexErr::Io(other)),
+        match manager
+            .recover_from_unauthorized_for_request(auth, recovery)
+            .await
+        {
+            Ok(crate::auth::UnauthorizedRecoveryDecision::Retry) => Ok(()),
+            Ok(crate::auth::UnauthorizedRecoveryDecision::GiveUp) => {
+                Err(map_unauthorized_status(status))
             }
+            Err(RefreshTokenError::Permanent(failed)) => Err(CodexErr::RefreshTokenFailed(failed)),
+            Err(RefreshTokenError::Transient(other)) => Err(CodexErr::Io(other)),
         }
     } else {
         Err(map_unauthorized_status(status))
     }
 }
 
-#[derive(Default, Debug)]
-struct UnauthorizedRecovery {
-    synced_from_storage: bool,
-    refreshed_token: bool,
-}
-
 fn map_unauthorized_status(status: StatusCode) -> CodexErr {
     map_api_error(ApiError::Transport(TransportError::Http {
         status,