diff --git a/codex-rs/core/src/client.rs b/codex-rs/core/src/client.rs index af5e3007b4..d8aec17be9 100644 --- a/codex-rs/core/src/client.rs +++ b/codex-rs/core/src/client.rs @@ -31,9 +31,7 @@ use std::sync::atomic::AtomicBool; use std::sync::atomic::Ordering; use crate::api_bridge::CoreAuthProvider; -use crate::api_bridge::auth_provider_from_auth; use crate::api_bridge::map_api_error; -use crate::auth::UnauthorizedRecovery; use crate::auth_env_telemetry::AuthEnvTelemetry; use crate::auth_env_telemetry::collect_auth_env_telemetry; use codex_api::CompactClient as ApiCompactClient; @@ -104,6 +102,12 @@ use crate::error::Result; use crate::flags::CODEX_RS_SSE_FIXTURE; use crate::model_provider_info::ModelProviderInfo; use crate::model_provider_info::WireApi; +use crate::request_auth::RequestUnauthorizedRecovery; +use crate::request_auth::ResolvedRequestAuth; +use crate::request_auth::UnauthorizedRecoveryError; +use crate::request_auth::UnauthorizedRecoveryExecution; +use crate::request_auth::UnauthorizedRecoveryOutcome; +use crate::request_auth::resolve_request_auth; use crate::response_debug_context::extract_response_debug_context; use crate::response_debug_context::extract_response_debug_context_from_api_error; use crate::response_debug_context::telemetry_api_error_message; @@ -144,16 +148,6 @@ struct ModelClientState { cached_websocket_session: StdMutex, } -/// Resolved API client setup for a single request attempt. -/// -/// Keeping this as a single bundle ensures prewarm and normal request paths -/// share the same auth/provider setup flow. -struct CurrentClientSetup { - auth: Option, - api_provider: codex_api::Provider, - api_auth: CoreAuthProvider, -} - #[derive(Clone, Copy)] struct RequestRouteTelemetry { endpoint: &'static str, @@ -523,21 +517,8 @@ impl ModelClient { /// /// This centralizes setup used by both prewarm and normal request paths so they stay in /// lockstep when auth/provider resolution changes. - async fn current_client_setup(&self) -> Result { - let auth = match self.state.auth_manager.as_ref() { - Some(manager) => manager.auth().await, - None => None, - }; - let api_provider = self - .state - .provider - .to_api_provider(auth.as_ref().map(CodexAuth::auth_mode))?; - let api_auth = auth_provider_from_auth(auth.clone(), &self.state.provider)?; - Ok(CurrentClientSetup { - auth, - api_provider, - api_auth, - }) + async fn current_client_setup(&self) -> Result { + resolve_request_auth(self.state.auth_manager.as_ref(), &self.state.provider).await } /// Opens a websocket connection using the same header and telemetry wiring as normal turns. @@ -1016,11 +997,11 @@ impl ModelClientSession { return Ok(stream); } - let auth_manager = self.client.state.auth_manager.clone(); - let mut auth_recovery = auth_manager - .as_ref() - .map(super::auth::AuthManager::unauthorized_recovery); + let mut unauthorized_recovery = + RequestUnauthorizedRecovery::new(self.client.state.auth_manager.as_ref()); let mut pending_retry = PendingUnauthorizedRetry::default(); + // Only loop after a successful auth-recovery step. Each retry must rebuild + // the client and request headers before issuing the same streaming request again. loop { let client_setup = self.client.current_client_setup().await?; let transport = ReqwestTransport::new(build_reqwest_client()); @@ -1065,7 +1046,7 @@ impl ModelClientSession { pending_retry = PendingUnauthorizedRetry::from_recovery( handle_unauthorized( unauthorized_transport, - &mut auth_recovery, + &mut unauthorized_recovery, session_telemetry, ) .await?, @@ -1104,12 +1085,11 @@ impl ModelClientSession { warmup: bool, request_trace: Option, ) -> Result { - let auth_manager = self.client.state.auth_manager.clone(); - - let mut auth_recovery = auth_manager - .as_ref() - .map(super::auth::AuthManager::unauthorized_recovery); + let mut unauthorized_recovery = + RequestUnauthorizedRecovery::new(self.client.state.auth_manager.as_ref()); let mut pending_retry = PendingUnauthorizedRetry::default(); + // Only loop after a successful auth-recovery step. WebSocket auth is attached + // during connect, so a recovered token requires a fresh connection attempt. loop { let client_setup = self.client.current_client_setup().await?; let request_auth_context = AuthRequestTelemetryContext::new( @@ -1162,14 +1142,13 @@ impl ModelClientSession { Err(ApiError::Transport( unauthorized_transport @ TransportError::Http { status, .. }, )) if status == StatusCode::UNAUTHORIZED => { - pending_retry = PendingUnauthorizedRetry::from_recovery( - handle_unauthorized( - unauthorized_transport, - &mut auth_recovery, - session_telemetry, - ) - .await?, - ); + let recovery = handle_unauthorized( + unauthorized_transport, + &mut unauthorized_recovery, + session_telemetry, + ) + .await?; + pending_retry = PendingUnauthorizedRetry::from_recovery(recovery); continue; } Err(err) => return Err(map_api_error(err)), @@ -1484,16 +1463,6 @@ where (ResponseStream { rx_event }, rx_last_response) } -/// Handles a 401 response by optionally refreshing ChatGPT tokens once. -/// -/// When refresh succeeds, the caller should retry the API call; otherwise -/// the mapped `CodexErr` is returned to the caller. -#[derive(Clone, Copy, Debug)] -struct UnauthorizedRecoveryExecution { - mode: &'static str, - phase: &'static str, -} - #[derive(Clone, Copy, Debug, Default)] struct PendingUnauthorizedRetry { retry_after_unauthorized: bool, @@ -1551,45 +1520,71 @@ struct WebsocketConnectParams<'a> { request_route_telemetry: RequestRouteTelemetry, } +/// Handles a `401 Unauthorized` from the transport used by the request loops above. +/// +/// The helper centralizes three coupled concerns: +/// - ask `RequestUnauthorizedRecovery` whether another recovery step is available +/// - record the matching telemetry / feedback tags for the outcome of that step +/// - return the recovery execution to the caller so it can rebuild auth state and retry, +/// or map the failure into the `CodexErr` that should terminate the loop async fn handle_unauthorized( transport: TransportError, - auth_recovery: &mut Option, + unauthorized_recovery: &mut RequestUnauthorizedRecovery, session_telemetry: &SessionTelemetry, ) -> Result { let debug = extract_response_debug_context(&transport); - if let Some(recovery) = auth_recovery - && recovery.has_next() - { - let mode = recovery.mode_name(); - let phase = recovery.step_name(); - return match recovery.next().await { - Ok(step_result) => { + match unauthorized_recovery.next().await { + Ok(UnauthorizedRecoveryOutcome::Recovered(recovery)) => { + session_telemetry.record_auth_recovery( + recovery.mode, + recovery.phase, + "recovery_succeeded", + debug.request_id.as_deref(), + debug.cf_ray.as_deref(), + debug.auth_error.as_deref(), + debug.auth_error_code.as_deref(), + /*recovery_reason*/ None, + recovery.auth_state_changed, + ); + emit_feedback_auth_recovery_tags( + recovery.mode, + recovery.phase, + "recovery_succeeded", + debug.request_id.as_deref(), + debug.cf_ray.as_deref(), + debug.auth_error.as_deref(), + debug.auth_error_code.as_deref(), + ); + Ok(recovery) + } + Ok(UnauthorizedRecoveryOutcome::Unavailable(unavailable)) => { + session_telemetry.record_auth_recovery( + unavailable.mode, + unavailable.phase, + "recovery_not_run", + debug.request_id.as_deref(), + debug.cf_ray.as_deref(), + debug.auth_error.as_deref(), + debug.auth_error_code.as_deref(), + Some(unavailable.reason), + /*auth_state_changed*/ None, + ); + emit_feedback_auth_recovery_tags( + unavailable.mode, + unavailable.phase, + "recovery_not_run", + debug.request_id.as_deref(), + debug.cf_ray.as_deref(), + debug.auth_error.as_deref(), + debug.auth_error_code.as_deref(), + ); + Err(map_api_error(ApiError::Transport(transport))) + } + Err(UnauthorizedRecoveryError::Chatgpt { execution, error }) => match error { + RefreshTokenError::Permanent(failed) => { session_telemetry.record_auth_recovery( - mode, - phase, - "recovery_succeeded", - debug.request_id.as_deref(), - debug.cf_ray.as_deref(), - debug.auth_error.as_deref(), - debug.auth_error_code.as_deref(), - /*recovery_reason*/ None, - step_result.auth_state_changed(), - ); - emit_feedback_auth_recovery_tags( - mode, - phase, - "recovery_succeeded", - debug.request_id.as_deref(), - debug.cf_ray.as_deref(), - debug.auth_error.as_deref(), - debug.auth_error_code.as_deref(), - ); - Ok(UnauthorizedRecoveryExecution { mode, phase }) - } - Err(RefreshTokenError::Permanent(failed)) => { - session_telemetry.record_auth_recovery( - mode, - phase, + execution.mode, + execution.phase, "recovery_failed_permanent", debug.request_id.as_deref(), debug.cf_ray.as_deref(), @@ -1599,8 +1594,8 @@ async fn handle_unauthorized( /*auth_state_changed*/ None, ); emit_feedback_auth_recovery_tags( - mode, - phase, + execution.mode, + execution.phase, "recovery_failed_permanent", debug.request_id.as_deref(), debug.cf_ray.as_deref(), @@ -1609,10 +1604,10 @@ async fn handle_unauthorized( ); Err(CodexErr::RefreshTokenFailed(failed)) } - Err(RefreshTokenError::Transient(other)) => { + RefreshTokenError::Transient(other) => { session_telemetry.record_auth_recovery( - mode, - phase, + execution.mode, + execution.phase, "recovery_failed_transient", debug.request_id.as_deref(), debug.cf_ray.as_deref(), @@ -1622,8 +1617,8 @@ async fn handle_unauthorized( /*auth_state_changed*/ None, ); emit_feedback_auth_recovery_tags( - mode, - phase, + execution.mode, + execution.phase, "recovery_failed_transient", debug.request_id.as_deref(), debug.cf_ray.as_deref(), @@ -1632,39 +1627,8 @@ async fn handle_unauthorized( ); Err(CodexErr::Io(other)) } - }; + }, } - - let (mode, phase, recovery_reason) = match auth_recovery.as_ref() { - Some(recovery) => ( - recovery.mode_name(), - recovery.step_name(), - Some(recovery.unavailable_reason()), - ), - None => ("none", "none", Some("auth_manager_missing")), - }; - session_telemetry.record_auth_recovery( - mode, - phase, - "recovery_not_run", - debug.request_id.as_deref(), - debug.cf_ray.as_deref(), - debug.auth_error.as_deref(), - debug.auth_error_code.as_deref(), - recovery_reason, - /*auth_state_changed*/ None, - ); - emit_feedback_auth_recovery_tags( - mode, - phase, - "recovery_not_run", - debug.request_id.as_deref(), - debug.cf_ray.as_deref(), - debug.auth_error.as_deref(), - debug.auth_error_code.as_deref(), - ); - - Err(map_api_error(ApiError::Transport(transport))) } fn api_error_http_status(error: &ApiError) -> Option { diff --git a/codex-rs/core/src/client_tests.rs b/codex-rs/core/src/client_tests.rs index b7d8075f0d..fb50ab7ff6 100644 --- a/codex-rs/core/src/client_tests.rs +++ b/codex-rs/core/src/client_tests.rs @@ -110,6 +110,7 @@ fn auth_request_telemetry_context_tracks_attached_auth_and_retry_phase() { PendingUnauthorizedRetry::from_recovery(UnauthorizedRecoveryExecution { mode: "managed", phase: "refresh_token", + auth_state_changed: None, }), ); diff --git a/codex-rs/core/src/lib.rs b/codex-rs/core/src/lib.rs index 5276b09de3..4219b1a51d 100644 --- a/codex-rs/core/src/lib.rs +++ b/codex-rs/core/src/lib.rs @@ -52,6 +52,7 @@ pub mod models_manager; mod network_policy_decision; pub mod network_proxy_loader; mod original_image_detail; +mod request_auth; pub use mcp_connection_manager::MCP_SANDBOX_STATE_CAPABILITY; pub use mcp_connection_manager::MCP_SANDBOX_STATE_METHOD; pub use mcp_connection_manager::SandboxState; diff --git a/codex-rs/core/src/models_manager/manager.rs b/codex-rs/core/src/models_manager/manager.rs index 29a1a85767..9497276a93 100644 --- a/codex-rs/core/src/models_manager/manager.rs +++ b/codex-rs/core/src/models_manager/manager.rs @@ -1,9 +1,7 @@ use super::cache::ModelsCacheManager; -use crate::api_bridge::auth_provider_from_auth; use crate::api_bridge::map_api_error; use crate::auth::AuthManager; use crate::auth::AuthMode; -use crate::auth::CodexAuth; use crate::auth_env_telemetry::AuthEnvTelemetry; use crate::auth_env_telemetry::collect_auth_env_telemetry; use crate::config::Config; @@ -14,6 +12,9 @@ use crate::model_provider_info::ModelProviderInfo; use crate::models_manager::collaboration_mode_presets::CollaborationModesConfig; use crate::models_manager::collaboration_mode_presets::builtin_collaboration_mode_presets; use crate::models_manager::model_info; +use crate::request_auth::RequestUnauthorizedRecovery; +use crate::request_auth::UnauthorizedRecoveryOutcome; +use crate::request_auth::resolve_request_auth; use crate::response_debug_context::extract_response_debug_context; use crate::response_debug_context::telemetry_transport_error_message; use crate::util::FeedbackRequestTags; @@ -22,6 +23,7 @@ use codex_api::ModelsClient; use codex_api::RequestTelemetry; use codex_api::ReqwestTransport; use codex_api::TransportError; +use codex_api::error::ApiError; use codex_otel::TelemetryAuthMode; use codex_protocol::config_types::CollaborationModeMask; use codex_protocol::openai_models::ModelInfo; @@ -431,39 +433,60 @@ impl ModelsManager { async fn fetch_and_update_models(&self) -> CoreResult<()> { let _timer = codex_otel::start_global_timer("codex.remote_models.fetch_update.duration_ms", &[]); - let auth = self.auth_manager.auth().await; - let auth_mode = auth.as_ref().map(CodexAuth::auth_mode); - let api_provider = self.provider.to_api_provider(auth_mode)?; - let api_auth = auth_provider_from_auth(auth.clone(), &self.provider)?; let auth_env = collect_auth_env_telemetry( &self.provider, self.auth_manager.codex_api_key_env_enabled(), ); - let transport = ReqwestTransport::new(build_reqwest_client()); - let request_telemetry: Arc = Arc::new(ModelsRequestTelemetry { - auth_mode: auth_mode.map(|mode| TelemetryAuthMode::from(mode).to_string()), - auth_header_attached: api_auth.auth_header_attached(), - auth_header_name: api_auth.auth_header_name(), - auth_env, - }); - let client = ModelsClient::new(transport, api_provider, api_auth) - .with_telemetry(Some(request_telemetry)); - let client_version = crate::models_manager::client_version_to_whole(); - let (models, etag) = timeout( - MODELS_REFRESH_TIMEOUT, - client.list_models(&client_version, HeaderMap::new()), - ) - .await - .map_err(|_| CodexErr::Timeout)? - .map_err(map_api_error)?; + let mut unauthorized_recovery = RequestUnauthorizedRecovery::new(Some(&self.auth_manager)); - self.apply_remote_models(models.clone()).await; - *self.etag.write().await = etag.clone(); - self.cache_manager - .persist_cache(&models, etag, client_version) - .await; - Ok(()) + // Only loop after a successful auth-recovery step so `/models` retries with + // the same freshly resolved auth state as normal request paths. + loop { + let request_auth = + resolve_request_auth(Some(&self.auth_manager), &self.provider).await?; + let transport = ReqwestTransport::new(build_reqwest_client()); + let request_telemetry: Arc = Arc::new(ModelsRequestTelemetry { + auth_mode: request_auth + .auth_mode + .map(|mode| TelemetryAuthMode::from(mode).to_string()), + auth_header_attached: request_auth.api_auth.auth_header_attached(), + auth_header_name: request_auth.api_auth.auth_header_name(), + auth_env: auth_env.clone(), + }); + let client = + ModelsClient::new(transport, request_auth.api_provider, request_auth.api_auth) + .with_telemetry(Some(request_telemetry)); + + match timeout( + MODELS_REFRESH_TIMEOUT, + client.list_models(&client_version, HeaderMap::new()), + ) + .await + .map_err(|_| CodexErr::Timeout)? + { + Ok((models, etag)) => { + self.apply_remote_models(models.clone()).await; + *self.etag.write().await = etag.clone(); + self.cache_manager + .persist_cache(&models, etag, client_version) + .await; + return Ok(()); + } + Err(ApiError::Transport( + unauthorized_transport @ TransportError::Http { status, .. }, + )) if status == http::StatusCode::UNAUTHORIZED => { + match unauthorized_recovery.next().await { + Ok(UnauthorizedRecoveryOutcome::Recovered(_)) => continue, + Ok(UnauthorizedRecoveryOutcome::Unavailable(_)) => { + return Err(map_api_error(ApiError::Transport(unauthorized_transport))); + } + Err(error) => return Err(error.into_codex_err()), + } + } + Err(err) => return Err(map_api_error(err)), + } + } } async fn get_etag(&self) -> Option { diff --git a/codex-rs/core/src/request_auth.rs b/codex-rs/core/src/request_auth.rs new file mode 100644 index 0000000000..31f74bed7d --- /dev/null +++ b/codex-rs/core/src/request_auth.rs @@ -0,0 +1,137 @@ +use std::sync::Arc; + +use crate::api_bridge::CoreAuthProvider; +use crate::api_bridge::auth_provider_from_auth; +use crate::auth::AuthManager; +use crate::auth::AuthMode; +use crate::auth::CodexAuth; +use crate::auth::RefreshTokenError; +use crate::auth::UnauthorizedRecovery; +use crate::error::CodexErr; +use crate::error::Result; +use crate::model_provider_info::ModelProviderInfo; + +#[derive(Clone)] +pub(crate) struct ResolvedRequestAuth { + pub(crate) auth: Option, + pub(crate) auth_mode: Option, + pub(crate) api_provider: codex_api::Provider, + pub(crate) api_auth: CoreAuthProvider, +} + +pub(crate) async fn resolve_request_auth( + auth_manager: Option<&Arc>, + provider: &ModelProviderInfo, +) -> Result { + let auth = match auth_manager { + Some(manager) => manager.auth().await, + None => None, + }; + let auth_mode = auth.as_ref().map(CodexAuth::auth_mode); + let api_provider = provider.to_api_provider(auth_mode)?; + let api_auth = auth_provider_from_auth(auth.clone(), provider)?; + Ok(ResolvedRequestAuth { + auth, + auth_mode, + api_provider, + api_auth, + }) +} + +#[derive(Clone, Copy, Debug)] +pub(crate) struct UnauthorizedRecoveryExecution { + pub(crate) mode: &'static str, + pub(crate) phase: &'static str, + pub(crate) auth_state_changed: Option, +} + +#[derive(Clone, Copy, Debug)] +pub(crate) struct UnauthorizedRecoveryUnavailable { + pub(crate) mode: &'static str, + pub(crate) phase: &'static str, + pub(crate) reason: &'static str, +} + +#[derive(Debug)] +pub(crate) enum UnauthorizedRecoveryOutcome { + Recovered(UnauthorizedRecoveryExecution), + Unavailable(UnauthorizedRecoveryUnavailable), +} + +#[derive(Debug)] +pub(crate) enum UnauthorizedRecoveryError { + Chatgpt { + execution: UnauthorizedRecoveryExecution, + error: RefreshTokenError, + }, +} + +impl UnauthorizedRecoveryError { + pub(crate) fn into_codex_err(self) -> CodexErr { + match self { + Self::Chatgpt { error, .. } => match error { + RefreshTokenError::Permanent(failed) => CodexErr::RefreshTokenFailed(failed), + RefreshTokenError::Transient(error) => CodexErr::Io(error), + }, + } + } +} + +/// Tracks which auth-recovery steps are still available after a request fails with `401`. +/// +/// The request loops reuse the same instance across retries so this type can enforce the +/// recovery ordering and make sure each one-shot step only runs once. +pub(crate) struct RequestUnauthorizedRecovery { + auth_recovery: Option, +} + +impl RequestUnauthorizedRecovery { + pub(crate) fn new(auth_manager: Option<&Arc>) -> Self { + Self { + auth_recovery: auth_manager.map(AuthManager::unauthorized_recovery), + } + } + + /// Runs the next available auth-recovery step, if any. + /// + /// Provider auth refresh is attempted first because it is specific to the configured + /// model provider and should repair the currently attached bearer token before falling + /// back to ChatGPT-managed recovery. Once all recovery steps have been consumed, this + /// returns `UnauthorizedRecoveryOutcome::Unavailable` so the caller can stop retrying. + pub(crate) async fn next( + &mut self, + ) -> std::result::Result { + if let Some(recovery) = self.auth_recovery.as_mut() + && recovery.has_next() + { + let execution = UnauthorizedRecoveryExecution { + mode: recovery.mode_name(), + phase: recovery.step_name(), + auth_state_changed: None, + }; + return match recovery.next().await { + Ok(step_result) => Ok(UnauthorizedRecoveryOutcome::Recovered( + UnauthorizedRecoveryExecution { + auth_state_changed: step_result.auth_state_changed(), + ..execution + }, + )), + Err(error) => Err(UnauthorizedRecoveryError::Chatgpt { execution, error }), + }; + } + + let unavailable = match self.auth_recovery.as_ref() { + Some(recovery) => UnauthorizedRecoveryUnavailable { + mode: recovery.mode_name(), + phase: recovery.step_name(), + reason: recovery.unavailable_reason(), + }, + None => UnauthorizedRecoveryUnavailable { + mode: "none", + phase: "none", + reason: "auth_manager_missing", + }, + }; + Ok(UnauthorizedRecoveryOutcome::Unavailable(unavailable)) + } +}