diff --git a/codex-rs/core/src/auth.rs b/codex-rs/core/src/auth.rs index a42ddc1503..1f82caafd0 100644 --- a/codex-rs/core/src/auth.rs +++ b/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 { + 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( diff --git a/codex-rs/core/src/client.rs b/codex-rs/core/src/client.rs index b0bd129531..415f4b4294 100644 --- a/codex-rs/core/src/client.rs +++ b/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>, auth: &Option, ) -> 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,