From 5a1c247dd5fd682c22a34af93df697dd06300a32 Mon Sep 17 00:00:00 2001 From: smurching-oai Date: Thu, 22 Oct 2026 23:30:51 +0000 Subject: [PATCH] Decouple managed configuration loading from ChatGPT authentication (#56295) ## What changed Introduce a `BundleSource` interface for authentication, cache identity, bundle fetching, and unauthorized-request recovery. Move ChatGPT plan eligibility and credential recovery into `ChatGptSource`, while keeping validation, caching, and retry orchestration in the shared service. Support source-specific cache namespaces while preserving the existing ChatGPT cache location. Retain one authentication recovery context across all retries of a bundle fetch. ## Testing Extend the external authentication test to verify that recovery happens only once per bundle fetch, including when an intervening request failure triggers a retry. GitOrigin-RevId: f5009f72dc4479f96d4ed1da9232f250988444a5 --- codex-rs/cloud-config/src/backend.rs | 73 ++++-- codex-rs/cloud-config/src/bundle_loader.rs | 11 +- codex-rs/cloud-config/src/cache.rs | 6 + codex-rs/cloud-config/src/chatgpt.rs | 99 ++++++++ codex-rs/cloud-config/src/config_source.rs | 43 ++++ codex-rs/cloud-config/src/lib.rs | 2 + codex-rs/cloud-config/src/service.rs | 272 ++++++++------------- codex-rs/cloud-config/src/service_tests.rs | 31 ++- 8 files changed, 334 insertions(+), 203 deletions(-) create mode 100644 codex-rs/cloud-config/src/chatgpt.rs create mode 100644 codex-rs/cloud-config/src/config_source.rs diff --git a/codex-rs/cloud-config/src/backend.rs b/codex-rs/cloud-config/src/backend.rs index fc56579940..3787229b43 100644 --- a/codex-rs/cloud-config/src/backend.rs +++ b/codex-rs/cloud-config/src/backend.rs @@ -13,12 +13,13 @@ use std::future::Future; #[derive(Clone, Copy, Debug, Eq, PartialEq)] pub(crate) enum RetryableFailureKind { Request { status_code: Option }, + AuthRecovery { status_code: Option }, } impl RetryableFailureKind { pub(crate) fn status_code(self) -> Option { match self { - Self::Request { status_code } => status_code, + Self::Request { status_code } | Self::AuthRecovery { status_code } => status_code, } } } @@ -26,6 +27,10 @@ impl RetryableFailureKind { #[derive(Clone, Debug, Eq, PartialEq)] pub(crate) enum BundleRequestError { Policy(codex_http_client::NetworkPolicyDenied), + Fatal { + error: codex_config::CloudConfigBundleLoadError, + metric_reason: Option<&'static str>, + }, Retryable(RetryableFailureKind), Unauthorized { status_code: Option, @@ -33,6 +38,23 @@ pub(crate) enum BundleRequestError { }, } +impl BundleRequestError { + pub(crate) fn auth_failure( + status_code: Option, + message: impl Into, + metric_reason: Option<&'static str>, + ) -> Self { + Self::Fatal { + error: codex_config::CloudConfigBundleLoadError::new( + codex_config::CloudConfigBundleLoadErrorCode::Auth, + status_code, + message, + ), + metric_reason, + } + } +} + /// Retrieves one cloud config bundle from the backend. /// /// Implementations should return the backend-selected bundle exactly as delivered and leave @@ -66,28 +88,35 @@ impl BundleClient for BackendBundleClient { self.http_client_factory.clone(), ); - let response = client - .get_config_bundle() - .await - .inspect_err(|err| { - tracing::warn!(error = %err, "Failed to fetch cloud config bundle"); - }) - .map_err(|err| { - if let codex_backend_client::RequestError::Policy(denied) = &err { - return BundleRequestError::Policy(*denied); - } - let status_code = err.status().map(|status| status.as_u16()); - if err.is_unauthorized() { - BundleRequestError::Unauthorized { - status_code, - message: err.to_string(), - } - } else { - BundleRequestError::Retryable(RetryableFailureKind::Request { status_code }) - } - })?; + fetch_bundle(&client).await + } +} - Ok(bundle_from_response(response)) +pub(crate) async fn fetch_bundle( + client: &BackendClient, +) -> Result { + client + .get_config_bundle() + .await + .inspect_err(|err| { + tracing::warn!(error = %err, "Failed to fetch cloud config bundle"); + }) + .map(bundle_from_response) + .map_err(bundle_request_error) +} + +fn bundle_request_error(error: codex_backend_client::RequestError) -> BundleRequestError { + if let codex_backend_client::RequestError::Policy(denied) = &error { + return BundleRequestError::Policy(*denied); + } + let status_code = error.status().map(|status| status.as_u16()); + if error.is_unauthorized() { + BundleRequestError::Unauthorized { + status_code, + message: "Managed configuration sign-in is required".into(), + } + } else { + BundleRequestError::Retryable(RetryableFailureKind::Request { status_code }) } } diff --git a/codex-rs/cloud-config/src/bundle_loader.rs b/codex-rs/cloud-config/src/bundle_loader.rs index 0d7bedc91f..2016e092ab 100644 --- a/codex-rs/cloud-config/src/bundle_loader.rs +++ b/codex-rs/cloud-config/src/bundle_loader.rs @@ -1,5 +1,6 @@ use crate::backend::BackendBundleClient; -use crate::backend::BundleClient; +use crate::chatgpt::ChatGptSource; +use crate::config_source::BundleSource; use crate::service::CLOUD_CONFIG_BUNDLE_TIMEOUT; use crate::service::CloudConfigBundleService; use codex_config::CloudConfigBundleLoader; @@ -11,12 +12,12 @@ use std::sync::Arc; use tokio::task::AbortHandle; use tokio::task::JoinHandle; -struct CloudConfigBundleLoaderLifetime { +struct CloudConfigBundleLoaderLifetime { service: Arc>, refresh_task: AbortHandle, } -impl Drop for CloudConfigBundleLoaderLifetime { +impl Drop for CloudConfigBundleLoaderLifetime { fn drop(&mut self) { self.refresh_task.abort(); } @@ -45,7 +46,7 @@ pub(crate) fn cloud_config_bundle_loader_for_service( service: CloudConfigBundleService, ) -> (CloudConfigBundleLoader, JoinHandle<()>) where - C: BundleClient + 'static, + C: BundleSource + 'static, { let service = Arc::new(service); let background_service = Arc::clone(&service); @@ -102,7 +103,7 @@ pub async fn cloud_config_bundle_loader_for_storage_without_cache( async fn cloud_config_bundle_service_for_storage( auth_config: AuthConfig, enable_codex_api_key_env: bool, -) -> std::io::Result> { +) -> std::io::Result>> { let auth_manager = AuthManager::shared_from_auth_config(auth_config.clone(), enable_codex_api_key_env).await?; Ok(CloudConfigBundleService::new( diff --git a/codex-rs/cloud-config/src/cache.rs b/codex-rs/cloud-config/src/cache.rs index 94748ec052..a08573500d 100644 --- a/codex-rs/cloud-config/src/cache.rs +++ b/codex-rs/cloud-config/src/cache.rs @@ -46,6 +46,12 @@ impl CloudConfigBundleCache { } } + pub(super) fn for_source(codex_home: AbsolutePathBuf, namespace: &str) -> Self { + Self { + path: codex_home.join(format!("managed-config-{namespace}.json")), + } + } + pub(super) fn path(&self) -> &Path { &self.path } diff --git a/codex-rs/cloud-config/src/chatgpt.rs b/codex-rs/cloud-config/src/chatgpt.rs new file mode 100644 index 0000000000..4d2c813fd2 --- /dev/null +++ b/codex-rs/cloud-config/src/chatgpt.rs @@ -0,0 +1,99 @@ +//! ChatGPT eligibility, credentials, and account identity for managed configuration. + +use crate::backend::BundleClient; +use crate::backend::BundleRequestError; +use crate::backend::RetryableFailureKind; +use crate::config_source::BundleSource; +use codex_config::CloudConfigBundle; +use codex_config::CloudConfigBundleLoadError; +use codex_login::AuthManager; +use codex_login::CodexAuth; +use codex_login::RefreshTokenError; +use codex_login::UnauthorizedRecovery; +use codex_protocol::account::PlanType; +use std::sync::Arc; + +const CLOUD_CONFIG_BUNDLE_AUTH_RECOVERY_FAILED_MESSAGE: &str = concat!( + "Your authentication session could not be refreshed automatically. ", + "Please log out and sign in again." +); + +pub(crate) struct ChatGptSource { + pub(crate) auth_manager: Arc, + pub(crate) client: Arc, +} + +impl BundleSource for ChatGptSource { + type Auth = CodexAuth; + type FetchContext = UnauthorizedRecovery; + fn begin_fetch(&self) -> Self::FetchContext { + self.auth_manager.unauthorized_recovery() + } + + async fn auth(&self) -> Result, CloudConfigBundleLoadError> { + Ok(self.auth_manager.auth().await.filter(|auth| { + auth.uses_codex_backend() + && auth.account_plan_type().is_some_and(|plan| { + plan.is_business_like() + || plan.is_education_like() + || matches!( + plan, + PlanType::Enterprise | PlanType::Sci | PlanType::SciPlus + ) + }) + })) + } + + fn identity(&self, auth: &CodexAuth) -> (Option, Option) { + (auth.get_chatgpt_user_id(), auth.get_account_id()) + } + + async fn get_bundle( + &self, + auth: &mut CodexAuth, + _context: &mut UnauthorizedRecovery, + ) -> Result { + self.client.get_bundle(auth).await + } + + async fn recover_auth( + &self, + auth: &mut CodexAuth, + recovery: &mut UnauthorizedRecovery, + status_code: Option, + _recovery_attempt: usize, + ) -> Result { + if !recovery.has_next() { + return Err(BundleRequestError::auth_failure( + status_code, + CLOUD_CONFIG_BUNDLE_AUTH_RECOVERY_FAILED_MESSAGE, + Some("auth_recovery_unavailable"), + )); + } + match recovery.next().await { + Ok(_) => { + *auth = self.auth_manager.auth().await.ok_or_else(|| { + BundleRequestError::auth_failure( + status_code, + CLOUD_CONFIG_BUNDLE_AUTH_RECOVERY_FAILED_MESSAGE, + Some("auth_recovery_missing_auth"), + ) + })?; + Ok(true) + } + Err(RefreshTokenError::Transient(_)) => Err(BundleRequestError::Retryable( + RetryableFailureKind::AuthRecovery { status_code }, + )), + Err(RefreshTokenError::Permanent(failed)) => Err(BundleRequestError::auth_failure( + status_code, + failed.message, + Some("auth_recovery_unrecoverable"), + )), + Err(RefreshTokenError::Policy(error)) => Err(BundleRequestError::auth_failure( + status_code, + error.to_string(), + /*metric_reason*/ None, + )), + } + } +} diff --git a/codex-rs/cloud-config/src/config_source.rs b/codex-rs/cloud-config/src/config_source.rs new file mode 100644 index 0000000000..523d8fc90e --- /dev/null +++ b/codex-rs/cloud-config/src/config_source.rs @@ -0,0 +1,43 @@ +//! Common interface for fetching managed configuration + +use crate::backend::BundleRequestError; +use codex_config::CloudConfigBundle; +use codex_config::CloudConfigBundleLoadError; +use std::future::Future; + +/// Supplies authentication and cache identity to the managed-config service. +/// Cache lookup and fetch must use the same authorization identity. +pub(crate) trait BundleSource: Send + Sync { + type Auth: Send + Sync; + type FetchContext: Send; + /// One context per load or refresh, retained across the service's retries. + fn begin_fetch(&self) -> Self::FetchContext; + + /// `None` means this source does not apply, such as an ineligible ChatGPT plan. + /// A required source returns `Some` even when its endpoint is unauthenticated. + fn auth( + &self, + ) -> impl Future, CloudConfigBundleLoadError>> + Send; + /// Stable subject and authority identifiers; missing values disable cache reads. + fn identity(&self, auth: &Self::Auth) -> (Option, Option); + /// Cache lookups do not call this operation or require a token refresh. + fn get_bundle( + &self, + auth: &mut Self::Auth, + context: &mut Self::FetchContext, + ) -> impl Future> + Send; + + /// Recover one unauthorized request. The service retries only when this returns true. + fn recover_auth( + &self, + auth: &mut Self::Auth, + context: &mut Self::FetchContext, + status_code: Option, + recovery_attempt: usize, + ) -> impl Future> + Send; + + /// The default preserves the existing ChatGPT cache location. + fn cache_namespace(&self) -> Option { + None + } +} diff --git a/codex-rs/cloud-config/src/lib.rs b/codex-rs/cloud-config/src/lib.rs index 3bcf30f4c3..0eb289a0f2 100644 --- a/codex-rs/cloud-config/src/lib.rs +++ b/codex-rs/cloud-config/src/lib.rs @@ -6,6 +6,8 @@ mod backend; mod bundle_loader; mod cache; +mod chatgpt; +mod config_source; mod metrics; mod service; mod validation; diff --git a/codex-rs/cloud-config/src/service.rs b/codex-rs/cloud-config/src/service.rs index e478b974b2..00c67c3613 100644 --- a/codex-rs/cloud-config/src/service.rs +++ b/codex-rs/cloud-config/src/service.rs @@ -9,6 +9,8 @@ use crate::backend::BundleRequestError; use crate::backend::RetryableFailureKind; use crate::cache::CacheLoadStatus; use crate::cache::CloudConfigBundleCache; +use crate::chatgpt::ChatGptSource; +use crate::config_source::BundleSource; use crate::metrics::emit_fetch_attempt_metric; use crate::metrics::emit_fetch_final_metric; use crate::metrics::emit_load_metric; @@ -21,10 +23,6 @@ use codex_config::CloudConfigBundleLoadErrorCode; use codex_config::CloudConfigBundlePolicy; use codex_config::CloudConfigBundleSnapshot; use codex_login::AuthManager; -use codex_login::CodexAuth; -use codex_login::RefreshTokenError; -use codex_login::UnauthorizedRecovery; -use codex_protocol::account::PlanType; use std::path::PathBuf; use std::sync::Arc; use std::time::Duration; @@ -40,28 +38,6 @@ const CLOUD_CONFIG_BUNDLE_CACHE_REFRESH_INTERVAL: Duration = Duration::from_secs const CLOUD_CONFIG_BUNDLE_TIMEOUT_RETRY_INTERVAL: Duration = Duration::from_secs(5); const CLOUD_CONFIG_BUNDLE_LOAD_FAILED_MESSAGE: &str = "Failed to load cloud config bundle (workspace-managed policies)."; -const CLOUD_CONFIG_BUNDLE_AUTH_RECOVERY_FAILED_MESSAGE: &str = concat!( - "Your authentication session could not be refreshed automatically. ", - "Please log out and sign in again." -); - -fn auth_identity(auth: &CodexAuth) -> (Option, Option) { - (auth.get_chatgpt_user_id(), auth.get_account_id()) -} - -fn cloud_config_eligible_auth(auth: &CodexAuth) -> bool { - let Some(plan_type) = auth.account_plan_type() else { - return false; - }; - auth.uses_codex_backend() - && (plan_type.is_business_like() - || plan_type.is_education_like() - || matches!( - plan_type, - PlanType::Enterprise | PlanType::Sci | PlanType::SciPlus - )) -} - fn optional_bundle(bundle: CloudConfigBundle) -> Option { if bundle.is_empty() { None @@ -75,14 +51,8 @@ enum CachedBundleLookup { Miss, } -enum UnauthorizedRecoveryAction { - RetrySameAttempt, - RetryNextAttempt, -} - -pub(crate) struct CloudConfigBundleService { - auth_manager: Arc, - client: Arc, +pub(crate) struct CloudConfigBundleService { + source: S, cache: CloudConfigBundleCache, cache_enabled: bool, codex_home: AbsolutePathBuf, @@ -91,21 +61,34 @@ pub(crate) struct CloudConfigBundleService { pub(crate) policy: CloudConfigBundlePolicy, } -impl CloudConfigBundleService -where - C: BundleClient + 'static, -{ +impl CloudConfigBundleService> { pub(crate) fn new( auth_manager: Arc, client: Arc, codex_home: PathBuf, timeout: Duration, ) -> Self { + Self::from_source( + ChatGptSource { + auth_manager, + client, + }, + codex_home, + timeout, + ) + } +} + +impl CloudConfigBundleService { + pub(crate) fn from_source(source: S, codex_home: PathBuf, timeout: Duration) -> Self { let codex_home = AbsolutePathBuf::resolve_path_against_base(codex_home, "/"); + let cache = match source.cache_namespace() { + Some(namespace) => CloudConfigBundleCache::for_source(codex_home.clone(), &namespace), + None => CloudConfigBundleCache::new(codex_home.clone()), + }; Self { - auth_manager, - client, - cache: CloudConfigBundleCache::new(codex_home.clone()), + source, + cache, cache_enabled: true, codex_home, timeout, @@ -203,17 +186,14 @@ where async fn load_startup_bundle( &self, ) -> Result, CloudConfigBundleLoadError> { - let Some(auth) = self.auth_manager.auth().await else { + let Some(auth) = self.source.auth().await? else { return Ok(None); }; - if !cloud_config_eligible_auth(&auth) { - return Ok(None); - } if self.cache_enabled { // Startup prefers a valid, identity-matched cache entry. The backend is // only consulted on cache miss or invalid cache contents. - let (chatgpt_user_id, account_id) = auth_identity(&auth); + let (chatgpt_user_id, account_id) = self.source.identity(&auth); match self .load_valid_cached_bundle(chatgpt_user_id.as_deref(), account_id.as_deref()) .await @@ -260,15 +240,18 @@ where async fn fetch_remote_bundle_and_update_cache_with_retries( &self, - mut auth: CodexAuth, + mut auth: S::Auth, trigger: &'static str, ) -> Result, CloudConfigBundleLoadError> { + let mut context = self.source.begin_fetch(); let mut attempt = 1; let mut last_status_code: Option = None; - let mut auth_recovery = self.auth_manager.unauthorized_recovery(); while attempt <= CLOUD_CONFIG_BUNDLE_MAX_ATTEMPTS { - match self.client.get_bundle(&auth).await { + match self + .fetch_with_auth_recovery(&mut auth, &mut context, trigger, attempt) + .await + { Ok(bundle) => { return self .validate_and_cache_remote_bundle(&auth, trigger, attempt, bundle) @@ -291,28 +274,39 @@ where continue; } } + Err(BundleRequestError::Fatal { + error, + metric_reason, + }) => { + if let Some(reason) = metric_reason { + emit_fetch_final_metric( + trigger, + "error", + reason, + attempt, + error.status_code(), + /*bundle*/ None, + ); + } + return Err(error); + } Err(BundleRequestError::Unauthorized { status_code, message, }) => { - last_status_code = status_code; - match self - .handle_unauthorized( - &mut auth, - &mut auth_recovery, - trigger, - attempt, - status_code, - &message, - ) - .await? - { - UnauthorizedRecoveryAction::RetrySameAttempt => continue, - UnauthorizedRecoveryAction::RetryNextAttempt => { - attempt += 1; - continue; - } - } + emit_fetch_final_metric( + trigger, + "error", + "auth", + attempt, + status_code, + /*bundle*/ None, + ); + return Err(CloudConfigBundleLoadError::new( + CloudConfigBundleLoadErrorCode::Auth, + status_code, + message, + )); } } @@ -338,9 +332,34 @@ where )) } + async fn fetch_with_auth_recovery( + &self, + auth: &mut S::Auth, + context: &mut S::FetchContext, + trigger: &'static str, + attempt: usize, + ) -> Result { + let mut recovery_attempt = 0; + loop { + let result = self.source.get_bundle(auth, context).await; + let Err(BundleRequestError::Unauthorized { status_code, .. }) = &result else { + return result; + }; + emit_fetch_attempt_metric(trigger, attempt, "unauthorized", *status_code); + if !self + .source + .recover_auth(auth, context, *status_code, recovery_attempt) + .await? + { + return result; + } + recovery_attempt += 1; + } + } + async fn validate_and_cache_remote_bundle( &self, - auth: &CodexAuth, + auth: &S::Auth, trigger: &'static str, attempt: usize, bundle: CloudConfigBundle, @@ -359,7 +378,7 @@ where return Err(err); } - let (chatgpt_user_id, account_id) = auth_identity(auth); + let (chatgpt_user_id, account_id) = self.source.identity(auth); if self.cache_enabled && let Some(revision) = revision { @@ -394,7 +413,9 @@ where status: RetryableFailureKind, ) -> bool { let status_code = status.status_code(); - emit_fetch_attempt_metric(trigger, attempt, "error", status_code); + if matches!(status, RetryableFailureKind::Request { .. }) { + emit_fetch_attempt_metric(trigger, attempt, "error", status_code); + } if attempt < CLOUD_CONFIG_BUNDLE_MAX_ATTEMPTS { tracing::warn!( status = ?status, @@ -409,105 +430,6 @@ where } } - async fn handle_unauthorized( - &self, - auth: &mut CodexAuth, - auth_recovery: &mut UnauthorizedRecovery, - trigger: &'static str, - attempt: usize, - status_code: Option, - message: &str, - ) -> Result { - emit_fetch_attempt_metric(trigger, attempt, "unauthorized", status_code); - if auth_recovery.has_next() { - tracing::warn!( - attempt, - max_attempts = CLOUD_CONFIG_BUNDLE_MAX_ATTEMPTS, - "Cloud config bundle request was unauthorized; attempting auth recovery" - ); - match auth_recovery.next().await { - Ok(_) => { - let Some(refreshed_auth) = self.auth_manager.auth().await else { - tracing::error!( - "Auth recovery succeeded but no auth is available for cloud config bundle" - ); - emit_fetch_final_metric( - trigger, - "error", - "auth_recovery_missing_auth", - attempt, - status_code, - /*bundle*/ None, - ); - return Err(CloudConfigBundleLoadError::new( - CloudConfigBundleLoadErrorCode::Auth, - status_code, - CLOUD_CONFIG_BUNDLE_AUTH_RECOVERY_FAILED_MESSAGE, - )); - }; - *auth = refreshed_auth; - return Ok(UnauthorizedRecoveryAction::RetrySameAttempt); - } - Err(RefreshTokenError::Policy(error)) => { - return Err(CloudConfigBundleLoadError::new( - CloudConfigBundleLoadErrorCode::Auth, - status_code, - error.to_string(), - )); - } - Err(RefreshTokenError::Permanent(failed)) => { - tracing::warn!( - error = %failed, - "Failed to recover from unauthorized cloud config bundle request" - ); - emit_fetch_final_metric( - trigger, - "error", - "auth_recovery_unrecoverable", - attempt, - status_code, - /*bundle*/ None, - ); - return Err(CloudConfigBundleLoadError::new( - CloudConfigBundleLoadErrorCode::Auth, - status_code, - failed.message, - )); - } - Err(RefreshTokenError::Transient(recovery_err)) => { - if attempt < CLOUD_CONFIG_BUNDLE_MAX_ATTEMPTS { - tracing::warn!( - error = %recovery_err, - attempt, - max_attempts = CLOUD_CONFIG_BUNDLE_MAX_ATTEMPTS, - "Failed to recover from unauthorized cloud config bundle request; retrying" - ); - sleep(backoff(attempt as u64)).await; - } - return Ok(UnauthorizedRecoveryAction::RetryNextAttempt); - } - } - } - - tracing::warn!( - error = %message, - "Cloud config bundle request was unauthorized and no auth recovery is available" - ); - emit_fetch_final_metric( - trigger, - "error", - "auth_recovery_unavailable", - attempt, - status_code, - /*bundle*/ None, - ); - Err(CloudConfigBundleLoadError::new( - CloudConfigBundleLoadErrorCode::Auth, - status_code, - CLOUD_CONFIG_BUNDLE_AUTH_RECOVERY_FAILED_MESSAGE, - )) - } - pub(crate) async fn refresh_cache_in_background(&self) { loop { let mut refresh_interval = CLOUD_CONFIG_BUNDLE_CACHE_REFRESH_INTERVAL; @@ -542,12 +464,14 @@ where } async fn refresh_cache_once(&self) -> bool { - let Some(auth) = self.auth_manager.auth().await else { - return false; + let auth = match self.source.auth().await { + Ok(Some(auth)) => auth, + Ok(None) => return false, + Err(error) => { + self.publish_refresh_result(Err(error)).await; + return true; + } }; - if !cloud_config_eligible_auth(&auth) { - return false; - } match self .fetch_remote_bundle_and_update_cache_with_retries(auth, "refresh") diff --git a/codex-rs/cloud-config/src/service_tests.rs b/codex-rs/cloud-config/src/service_tests.rs index b3f00454cf..2a131a6dee 100644 --- a/codex-rs/cloud-config/src/service_tests.rs +++ b/codex-rs/cloud-config/src/service_tests.rs @@ -19,10 +19,12 @@ use codex_config::CloudRequirementsTomlBundle; use codex_config::types::AuthCredentialsStoreMode; use codex_core::config::ConfigBuilder; use codex_login::AuthKeyringBackendKind; +use codex_login::CodexAuth; use codex_login::auth::AgentIdentityAuth; use codex_login::auth::AgentIdentityAuthRecord; use codex_login::auth::ExternalAuth; use codex_login::auth::ExternalAuthRefreshContext; +use codex_protocol::account::PlanType; use pretty_assertions::assert_eq; use serde_json::json; use std::collections::VecDeque; @@ -1131,7 +1133,7 @@ async fn get_bundle_surfaces_auth_recovery_message() { } #[tokio::test] -async fn get_bundle_refreshes_external_auth_after_unauthorized() { +async fn external_auth_recovers_once_per_bundle_fetch() { let auth_home = tempdir().expect("tempdir"); let auth_manager = Arc::new( AuthManager::new( @@ -1175,7 +1177,7 @@ async fn get_bundle_refreshes_external_auth_after_unauthorized() { }); let codex_home = tempdir().expect("tempdir"); let service = CloudConfigBundleService::new( - auth_manager, + Arc::clone(&auth_manager), fetcher.clone(), codex_home.path().to_path_buf(), CLOUD_CONFIG_BUNDLE_TIMEOUT, @@ -1184,6 +1186,31 @@ async fn get_bundle_refreshes_external_auth_after_unauthorized() { assert_eq!(service.load_startup_bundle().await, Ok(Some(test_bundle()))); assert_eq!(fetcher.request_count.load(Ordering::SeqCst), 2); assert_eq!(external_auth.refresh_count.load(Ordering::SeqCst), 1); + let retried_fetcher = Arc::new(SequenceBundleClient::new(vec![ + Err(BundleRequestError::Unauthorized { + status_code: Some(401), + message: "rejected".into(), + }), + Err(request_error()), + Err(BundleRequestError::Unauthorized { + status_code: Some(401), + message: "still rejected".into(), + }), + Ok(test_bundle()), + ])); + let retrying = CloudConfigBundleService::new( + Arc::clone(&auth_manager), + retried_fetcher.clone(), + codex_home.path().to_path_buf(), + CLOUD_CONFIG_BUNDLE_TIMEOUT, + ) + .without_cache(); + assert_eq!( + retrying.get_latest().await.unwrap_err().code(), + CloudConfigBundleLoadErrorCode::Auth + ); + assert_eq!(retried_fetcher.request_count.load(Ordering::SeqCst), 3); + assert_eq!(external_auth.refresh_count.load(Ordering::SeqCst), 2); } #[tokio::test]