From 22b860e80b06dff3c5e86160e5dc3199b7816693 Mon Sep 17 00:00:00 2001 From: alexsong-oai Date: Tue, 18 Aug 2026 19:59:21 +0000 Subject: [PATCH] Add provider-owned authentication recovery (#39274) ## What changed - Let model providers identify recoverable authentication errors and refresh their own authentication state before falling back to `AuthManager` recovery. - Apply provider recovery to both HTTP and WebSocket response requests, with at most one provider-owned recovery attempt per request. - Preserve existing `401 Unauthorized` handling as the default for providers that do not customize recovery. ## Testing - Add coverage for successful and failed provider recovery, including the single-attempt bound and preservation of the original unauthorized error. GitOrigin-RevId: 9a737e3be96354699586f97dd88e544a45082969 --- codex-rs/core/src/client.rs | 55 +++++++++-- codex-rs/core/src/client_tests.rs | 118 ++++++++++++++++++++++++ codex-rs/model-provider/src/lib.rs | 1 + codex-rs/model-provider/src/provider.rs | 28 ++++++ 4 files changed, 193 insertions(+), 9 deletions(-) diff --git a/codex-rs/core/src/client.rs b/codex-rs/core/src/client.rs index da3b6dad38..f26d41955e 100644 --- a/codex-rs/core/src/client.rs +++ b/codex-rs/core/src/client.rs @@ -126,6 +126,7 @@ use codex_login::auth_env_telemetry::AuthEnvTelemetry; use codex_login::auth_env_telemetry::collect_auth_env_telemetry; use codex_model_provider::AgentIdentitySessionFallback; use codex_model_provider::ProviderAuthScope; +use codex_model_provider::ProviderUnauthorizedRecovery; use codex_model_provider::SharedModelProvider; use codex_model_provider::create_model_provider; #[cfg(test)] @@ -1451,6 +1452,7 @@ impl ModelClientSession { let mut auth_recovery = auth_manager .as_ref() .map(AuthManager::unauthorized_recovery); + let mut provider_auth_recovery_attempted = false; let mut pending_retry = PendingUnauthorizedRetry::default(); loop { let client_setup = self.client.current_client_setup().await?; @@ -1520,9 +1522,13 @@ impl ModelClientSession { ); return Ok(stream); } - Err(ApiError::Transport( - unauthorized_transport @ TransportError::Http { status, .. }, - )) if status == StatusCode::UNAUTHORIZED => { + Err(ApiError::Transport(unauthorized_transport)) + if self + .client + .state + .provider + .is_recoverable_auth_error(&unauthorized_transport) => + { let response_debug_context = extract_response_debug_context(&unauthorized_transport); inference_trace_attempt.record_failed( @@ -1534,6 +1540,7 @@ impl ModelClientSession { handle_unauthorized( unauthorized_transport, &mut auth_recovery, + &mut provider_auth_recovery_attempted, session_telemetry, &self.client.state.provider, ) @@ -1584,11 +1591,13 @@ impl ModelClientSession { request_trace: Option, inference_trace: &InferenceTraceContext, ) -> Result { - let auth_manager = self.client.state.provider.auth_manager(); + let provider = Arc::clone(&self.client.state.provider); + let auth_manager = provider.auth_manager(); let mut auth_recovery = auth_manager .as_ref() .map(AuthManager::unauthorized_recovery); + let mut provider_auth_recovery_attempted = false; let mut pending_retry = PendingUnauthorizedRetry::default(); loop { let client_setup = self.client.current_client_setup().await?; @@ -1643,21 +1652,22 @@ impl ModelClientSession { { return Ok(WebsocketStreamOutcome::FallbackToHttp); } - Err(ApiError::Transport( - unauthorized_transport @ TransportError::Http { status, .. }, - )) if status == StatusCode::UNAUTHORIZED => { + Err(ApiError::Transport(unauthorized_transport)) + if provider.is_recoverable_auth_error(&unauthorized_transport) => + { pending_retry = PendingUnauthorizedRetry::from_recovery( handle_unauthorized( unauthorized_transport, &mut auth_recovery, + &mut provider_auth_recovery_attempted, session_telemetry, - &self.client.state.provider, + &provider, ) .await?, ); continue; } - Err(err) => return Err(self.client.state.provider.map_api_error(err)), + Err(err) => return Err(provider.map_api_error(err)), } let (incremental_request, previous_response_id_from_untraced_warmup) = @@ -2225,10 +2235,37 @@ struct WebsocketConnectParams<'a> { async fn handle_unauthorized( transport: TransportError, auth_recovery: &mut Option, + provider_auth_recovery_attempted: &mut bool, session_telemetry: &SessionTelemetry, provider: &SharedModelProvider, ) -> Result { let debug = extract_response_debug_context(&transport); + if !*provider_auth_recovery_attempted { + *provider_auth_recovery_attempted = true; + match provider.recover_from_unauthorized().await { + Ok(ProviderUnauthorizedRecovery::Recovered) => { + return Ok(UnauthorizedRecoveryExecution { + mode: "provider", + phase: "provider_refresh", + }); + } + Ok(ProviderUnauthorizedRecovery::NotConfigured) => {} + Err(error) => { + let original = provider.map_api_error(ApiError::Transport(transport)); + warn!( + error = %error, + original_error = %original, + "provider authentication recovery failed" + ); + return Err(if error.is_retryable() { + original + } else { + error + }); + } + } + } + if let Some(recovery) = auth_recovery && recovery.has_next() { diff --git a/codex-rs/core/src/client_tests.rs b/codex-rs/core/src/client_tests.rs index 41c7dc1ab0..58f2f5f72e 100644 --- a/codex-rs/core/src/client_tests.rs +++ b/codex-rs/core/src/client_tests.rs @@ -27,19 +27,27 @@ use codex_login::AuthManager; use codex_login::CodexAuth; use codex_login::auth::AgentIdentityAuthPolicy; use codex_model_provider::BearerAuthProvider; +use codex_model_provider::ModelProvider; +use codex_model_provider::ModelProviderFuture; +use codex_model_provider::ProviderAccountResult; +use codex_model_provider::ProviderUnauthorizedRecovery; use codex_model_provider::SharedModelProvider; use codex_model_provider::create_model_provider; use codex_model_provider_info::CHATGPT_CODEX_BASE_URL; use codex_model_provider_info::ModelProviderInfo; use codex_model_provider_info::WireApi; use codex_model_provider_info::create_oss_provider_with_base_url; +use codex_models_manager::manager::SharedModelsManager; use codex_otel::SessionTelemetry; use codex_protocol::ThreadId; use codex_protocol::auth::AuthMode; +use codex_protocol::error::CodexErr; +use codex_protocol::error::CodexErrorDetails; use codex_protocol::models::BaseInstructions; use codex_protocol::models::ContentItem; use codex_protocol::models::ResponseItem; use codex_protocol::openai_models::ModelInfo; +use codex_protocol::openai_models::ModelsResponse; use codex_protocol::openai_models::ReasoningEffort; use codex_protocol::protocol::InternalSessionSource; use codex_protocol::protocol::SessionSource; @@ -57,6 +65,7 @@ use pretty_assertions::assert_eq; use serde_json::json; use std::collections::BTreeMap; use std::collections::VecDeque; +use std::path::PathBuf; use std::pin::Pin; use std::sync::Arc; use std::sync::Mutex; @@ -659,6 +668,7 @@ async fn bedrock_unauthorized_error_uses_provider_mapping() { /*auth_manager*/ None, ); let mut auth_recovery = None; + let mut provider_auth_recovery_attempted = false; let url = "https://bedrock-mantle.us-east-2.api.aws/openai/v1/responses"; let error = super::handle_unauthorized( TransportError::Http { @@ -671,6 +681,7 @@ async fn bedrock_unauthorized_error_uses_provider_mapping() { ), }, &mut auth_recovery, + &mut provider_auth_recovery_attempted, &test_session_telemetry(), &provider, ) @@ -685,6 +696,113 @@ async fn bedrock_unauthorized_error_uses_provider_mapping() { ); } +#[derive(Debug)] +struct TestRecoveryProvider { + inner: SharedModelProvider, + should_fail: bool, + attempts: Arc, +} + +impl ModelProvider for TestRecoveryProvider { + fn info(&self) -> &ModelProviderInfo { + self.inner.info() + } + + fn auth_manager(&self) -> Option> { + None + } + + fn auth(&self) -> ModelProviderFuture<'_, Option> { + self.inner.auth() + } + + fn account_state(&self) -> ProviderAccountResult { + self.inner.account_state() + } + + fn recover_from_unauthorized( + &self, + ) -> ModelProviderFuture<'_, codex_protocol::error::Result> { + self.attempts.fetch_add(1, Ordering::Relaxed); + Box::pin(async move { + if self.should_fail { + Err(CodexErr::Io(std::io::Error::other( + "provider recovery failed", + ))) + } else { + Ok(ProviderUnauthorizedRecovery::Recovered) + } + }) + } + + fn models_manager( + &self, + codex_home: PathBuf, + config_model_catalog: Option, + ) -> SharedModelsManager { + self.inner.models_manager(codex_home, config_model_catalog) + } +} + +#[tokio::test] +async fn provider_owned_auth_recovery_is_bounded_and_preserves_unauthorized_failures() { + for should_fail in [false, true] { + let attempts = Arc::new(AtomicUsize::new(0)); + let provider: SharedModelProvider = Arc::new(TestRecoveryProvider { + inner: test_model_provider(), + should_fail, + attempts: Arc::clone(&attempts), + }); + assert!(provider.auth_manager().is_none()); + + let unauthorized = || TransportError::Http { + status: http::StatusCode::UNAUTHORIZED, + url: Some("https://example.com/v1/responses".to_string()), + headers: None, + body: Some("unauthorized".to_string()), + }; + let mut auth_recovery = None; + let mut provider_auth_recovery_attempted = false; + let telemetry = test_session_telemetry(); + let result = super::handle_unauthorized( + unauthorized(), + &mut auth_recovery, + &mut provider_auth_recovery_attempted, + &telemetry, + &provider, + ) + .await; + + let error = if should_fail { + result.expect_err("failed provider recovery should return the original error") + } else { + let recovered = result.expect("provider recovery should succeed without AuthManager"); + assert_eq!( + (recovered.mode, recovered.phase), + ("provider", "provider_refresh") + ); + super::handle_unauthorized( + unauthorized(), + &mut auth_recovery, + &mut provider_auth_recovery_attempted, + &telemetry, + &provider, + ) + .await + .expect_err("provider recovery should not run more than once") + }; + + match error.details() { + CodexErrorDetails::UnexpectedStatus(response) => { + assert_eq!(response.status, http::StatusCode::UNAUTHORIZED); + assert_eq!(response.body, "unauthorized"); + } + other => panic!("unexpected error after provider recovery: {other}"), + } + assert_eq!(attempts.load(Ordering::Relaxed), 1); + } +} + #[tokio::test] async fn dropped_backpressured_response_stream_traces_cancelled_partial_output() -> anyhow::Result<()> { diff --git a/codex-rs/model-provider/src/lib.rs b/codex-rs/model-provider/src/lib.rs index b1bd6d2162..22b5da74ed 100644 --- a/codex-rs/model-provider/src/lib.rs +++ b/codex-rs/model-provider/src/lib.rs @@ -23,6 +23,7 @@ pub use provider::ProviderAccountError; pub use provider::ProviderAccountResult; pub use provider::ProviderAccountState; pub use provider::ProviderCapabilities; +pub use provider::ProviderUnauthorizedRecovery; pub use provider::RemoteCompactionSupport; pub use provider::SharedModelProvider; pub use provider::create_model_provider; diff --git a/codex-rs/model-provider/src/provider.rs b/codex-rs/model-provider/src/provider.rs index aee4f9a2b5..f4fc6b6688 100644 --- a/codex-rs/model-provider/src/provider.rs +++ b/codex-rs/model-provider/src/provider.rs @@ -7,6 +7,7 @@ use std::sync::Arc; use codex_api::ApiError; use codex_api::Provider; use codex_api::SharedAuthProvider; +use codex_api::TransportError; use codex_api::is_azure_responses_provider; use codex_login::AuthManager; use codex_login::CodexAuth; @@ -71,6 +72,15 @@ pub struct ProviderAccountState { pub requires_openai_auth: bool, } +/// Outcome of a provider-owned attempt to recover from an authentication failure. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum ProviderUnauthorizedRecovery { + /// The provider has no provider-specific authentication recovery configured. + NotConfigured, + /// The provider recovered its authentication state and the request can be retried. + Recovered, +} + /// Error returned when a provider cannot construct its app-visible account state. #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub enum ProviderAccountError { @@ -160,6 +170,24 @@ pub trait ModelProvider: fmt::Debug + Send + Sync { /// manager throughout the codebase; that is a larger refactor than this change. fn auth_manager(&self) -> Option>; + /// Returns whether this transport failure can be recovered by provider-scoped authentication. + /// + /// The default preserves existing unauthorized-response handling. Providers with other + /// authentication failure shapes may recognize additional response or request-signing errors. + fn is_recoverable_auth_error(&self, error: &TransportError) -> bool { + matches!( + error, + TransportError::Http { status, .. } if *status == http::StatusCode::UNAUTHORIZED + ) + } + + /// Attempts provider-owned authentication recovery before using the auth manager. + fn recover_from_unauthorized( + &self, + ) -> ModelProviderFuture<'_, codex_protocol::error::Result> { + Box::pin(async { Ok(ProviderUnauthorizedRecovery::NotConfigured) }) + } + /// Returns the current provider-scoped auth value, if one is configured. fn auth(&self) -> ModelProviderFuture<'_, Option>;