diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index 023c5dfaa3..7c22e5e91f 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -2553,7 +2553,6 @@ dependencies = [ "core-foundation 0.9.4", "dns-lookup", "dunce", - "futures", "gethostname", "indexmap 2.14.0", "libc", diff --git a/codex-rs/cloud-config/Cargo.toml b/codex-rs/cloud-config/Cargo.toml index 363cb21b6f..d88bbc6a87 100644 --- a/codex-rs/cloud-config/Cargo.toml +++ b/codex-rs/cloud-config/Cargo.toml @@ -21,7 +21,7 @@ serde = { workspace = true, features = ["derive"] } serde_json = { workspace = true } sha2 = { workspace = true } thiserror = { workspace = true } -tokio = { workspace = true, features = ["fs", "rt", "sync", "time"] } +tokio = { workspace = true, features = ["fs", "macros", "rt", "sync", "time"] } tracing = { workspace = true } [dev-dependencies] diff --git a/codex-rs/cloud-config/src/bundle_loader.rs b/codex-rs/cloud-config/src/bundle_loader.rs index e4266c6195..853dc842db 100644 --- a/codex-rs/cloud-config/src/bundle_loader.rs +++ b/codex-rs/cloud-config/src/bundle_loader.rs @@ -1,8 +1,8 @@ use crate::backend::BackendBundleClient; +use crate::service::CLOUD_CONFIG_BUNDLE_CACHE_REFRESH_RETRY_INTERVAL; use crate::service::CLOUD_CONFIG_BUNDLE_TIMEOUT; use crate::service::CloudConfigBundleService; -use codex_config::CloudConfigBundleLoadError; -use codex_config::CloudConfigBundleLoadErrorCode; +use crate::service::StartupLoad; use codex_config::CloudConfigBundleLoader; use codex_config::types::AuthCredentialsStoreMode; use codex_login::AuthKeyringBackendKind; @@ -10,14 +10,6 @@ use codex_login::AuthManager; use codex_login::AuthRouteConfig; use std::path::PathBuf; use std::sync::Arc; -use std::sync::Mutex; -use std::sync::OnceLock; -use tokio::task::JoinHandle; - -fn refresher_task_slot() -> &'static Mutex>> { - static REFRESHER_TASK: OnceLock>>> = OnceLock::new(); - REFRESHER_TASK.get_or_init(|| Mutex::new(None)) -} pub fn cloud_config_bundle_loader( auth_manager: Arc, @@ -30,27 +22,29 @@ pub fn cloud_config_bundle_loader( codex_home, CLOUD_CONFIG_BUNDLE_TIMEOUT, ); - let refresh_service = service.clone(); - let task = tokio::spawn(async move { service.load_startup_bundle_with_timeout().await }); - let refresh_task = - tokio::spawn(async move { refresh_service.refresh_cache_in_background().await }); - let mut refresher_guard = refresher_task_slot().lock().unwrap_or_else(|err| { - tracing::warn!("cloud config bundle refresher task slot was poisoned"); - err.into_inner() - }); - if let Some(existing_task) = refresher_guard.replace(refresh_task) { - existing_task.abort(); - } - CloudConfigBundleLoader::new(async move { - task.await.map_err(|err| { - tracing::error!(error = %err, "Cloud config bundle task failed"); - CloudConfigBundleLoadError::new( - CloudConfigBundleLoadErrorCode::Internal, - /*status_code*/ None, - format!("cloud config bundle load failed: {err}"), - ) - })? - }) + let (loader, publisher) = CloudConfigBundleLoader::pending(); + // Intentionally detach the task that owns startup and refresh. Publishing + // stops the lifecycle if this loader and all clones were dropped. + drop(tokio::spawn(async move { + let result = service.load_startup_bundle().await; + let (initial_bundle, refresh_in) = match result { + Ok(StartupLoad::Inactive) => (Ok(None), None), + Ok(StartupLoad::Active(loaded)) => (Ok(loaded.bundle), Some(loaded.refresh_in)), + Err(err) => ( + Err(err), + Some(CLOUD_CONFIG_BUNDLE_CACHE_REFRESH_RETRY_INTERVAL), + ), + }; + if !publisher.publish(initial_bundle) { + return; + } + if let Some(refresh_in) = refresh_in { + service + .refresh_cache_in_background(refresh_in, publisher) + .await; + } + })); + loader } pub async fn cloud_config_bundle_loader_for_storage( diff --git a/codex-rs/cloud-config/src/cache.rs b/codex-rs/cloud-config/src/cache.rs index 887f84b8c2..3ecaf6d905 100644 --- a/codex-rs/cloud-config/src/cache.rs +++ b/codex-rs/cloud-config/src/cache.rs @@ -87,7 +87,7 @@ impl CloudConfigBundleCache { &self, chatgpt_user_id: Option<&str>, account_id: Option<&str>, - ) -> Result { + ) -> Result { let (Some(chatgpt_user_id), Some(account_id)) = (chatgpt_user_id, account_id) else { return Err(CacheLoadStatus::AuthIdentityIncomplete); }; @@ -139,14 +139,27 @@ impl CloudConfigBundleCache { let now = Utc::now(); let cache_age = now .signed_duration_since(cache_file.signed_payload.cached_at) - .to_std(); - if cache_file.signed_payload.expires_at <= now - || !matches!(cache_age, Ok(cache_age) if cache_age < CLOUD_CONFIG_BUNDLE_CACHE_TTL) + .to_std() + .map_err(|_| CacheLoadStatus::CacheExpired)?; + if cache_file.signed_payload.expires_at <= now || cache_age >= CLOUD_CONFIG_BUNDLE_CACHE_TTL { return Err(CacheLoadStatus::CacheExpired); } - Ok(cache_file.signed_payload) + let expires_in = cache_file + .signed_payload + .expires_at + .signed_duration_since(now) + .to_std() + .map_err(|_| CacheLoadStatus::CacheExpired)?; + let ttl_remaining = CLOUD_CONFIG_BUNDLE_CACHE_TTL + .checked_sub(cache_age) + .ok_or(CacheLoadStatus::CacheExpired)?; + + Ok(LoadedCloudConfigBundleCache { + signed_payload: cache_file.signed_payload, + refresh_in: expires_in.min(ttl_remaining), + }) } pub(super) fn log_load_status(&self, status: &CacheLoadStatus) { @@ -173,7 +186,7 @@ impl CloudConfigBundleCache { chatgpt_user_id: Option, account_id: Option, bundle: CloudConfigBundle, - ) -> Result<(), CloudConfigBundleCacheError> { + ) -> Result { let now = Utc::now(); let expires_at = now .checked_add_signed( @@ -206,7 +219,10 @@ impl CloudConfigBundleCache { fs::write(&self.path, serialized) .await .map_err(|_| CloudConfigBundleCacheError)?; - Ok(()) + expires_at + .signed_duration_since(Utc::now()) + .to_std() + .map_err(|_| CloudConfigBundleCacheError) } } @@ -238,6 +254,12 @@ pub(super) enum CacheLoadStatus { #[error("failed to write cloud config bundle cache")] pub(super) struct CloudConfigBundleCacheError; +#[derive(Clone, Debug, Eq, PartialEq)] +pub(super) struct LoadedCloudConfigBundleCache { + pub(super) signed_payload: CloudConfigBundleCacheSignedPayload, + pub(super) refresh_in: Duration, +} + #[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)] pub(super) struct CloudConfigBundleCacheFile { pub(super) signed_payload: CloudConfigBundleCacheSignedPayload, diff --git a/codex-rs/cloud-config/src/cache_tests.rs b/codex-rs/cloud-config/src/cache_tests.rs index af30e34a1b..2616077479 100644 --- a/codex-rs/cloud-config/src/cache_tests.rs +++ b/codex-rs/cloud-config/src/cache_tests.rs @@ -98,10 +98,12 @@ async fn save_writes_signed_payload_and_loads_for_matching_identity() { }) ); - assert_eq!( - cache.load(Some("user-12345"), Some("account-12345")).await, - Ok(cache_file.signed_payload) - ); + let loaded_cache = cache + .load(Some("user-12345"), Some("account-12345")) + .await + .expect("load cache"); + assert_eq!(loaded_cache.signed_payload, cache_file.signed_payload); + assert!(loaded_cache.refresh_in <= CLOUD_CONFIG_BUNDLE_CACHE_TTL); } #[tokio::test] @@ -209,6 +211,25 @@ async fn load_rejects_cache_older_than_ttl_even_when_expiry_is_later() { ); } +#[tokio::test] +async fn load_uses_ttl_cap_for_refresh_delay() { + let codex_home = tempdir().expect("tempdir"); + let cache = create_test_cache(codex_home.path()); + let now = Utc::now(); + let mut signed_payload = valid_signed_payload(); + signed_payload.cached_at = now - ChronoDuration::minutes(5); + signed_payload.expires_at = now + ChronoDuration::minutes(30); + write_cache_file(&cache, &signed_cache_file(signed_payload)); + + let loaded_cache = cache + .load(Some("user-12345"), Some("account-12345")) + .await + .expect("load cache"); + + assert!(loaded_cache.refresh_in <= Duration::from_secs(10 * 60)); + assert!(loaded_cache.refresh_in >= Duration::from_secs(10 * 60 - 5)); +} + #[tokio::test] async fn load_rejects_unsupported_cache_version() { let codex_home = tempdir().expect("tempdir"); diff --git a/codex-rs/cloud-config/src/service.rs b/codex-rs/cloud-config/src/service.rs index 71962c44c9..4e97497375 100644 --- a/codex-rs/cloud-config/src/service.rs +++ b/codex-rs/cloud-config/src/service.rs @@ -1,8 +1,7 @@ //! Cloud config bundle lifecycle orchestration. //! -//! Startup loads a single shared bundle from cache or backend, and a background -//! refresher keeps the cache warm for future app starts without changing the -//! already-snapshotted runtime config. +//! Startup loads a shared bundle from cache or backend. A background refresher +//! updates both the on-disk cache and the bundle used by future config loads. use crate::backend::BundleClient; use crate::backend::BundleRequestError; @@ -18,6 +17,7 @@ use codex_config::AbsolutePathBuf; use codex_config::CloudConfigBundle; use codex_config::CloudConfigBundleLoadError; use codex_config::CloudConfigBundleLoadErrorCode; +use codex_config::CloudConfigBundlePublisher; use codex_core::util::backoff; use codex_login::AuthManager; use codex_login::CodexAuth; @@ -33,7 +33,8 @@ use tokio::time::timeout; pub(crate) const CLOUD_CONFIG_BUNDLE_TIMEOUT: Duration = Duration::from_secs(15); const CLOUD_CONFIG_BUNDLE_MAX_ATTEMPTS: usize = 5; -const CLOUD_CONFIG_BUNDLE_CACHE_CHECK_INTERVAL: Duration = Duration::from_secs(60); +pub(crate) const CLOUD_CONFIG_BUNDLE_CACHE_REFRESH_RETRY_INTERVAL: Duration = + Duration::from_secs(60); const CLOUD_CONFIG_BUNDLE_CACHE_LOCK_RETRY_INTERVAL: Duration = Duration::from_secs(1); const CLOUD_CONFIG_BUNDLE_LOAD_FAILED_MESSAGE: &str = "Failed to load cloud config bundle (workspace-managed policies)."; @@ -63,9 +64,24 @@ fn optional_bundle(bundle: CloudConfigBundle) -> Option { } } -enum CachedBundleLookup { - Hit(Option), - Miss, +#[derive(Debug, Eq, PartialEq)] +pub(crate) struct LoadedBundle { + pub(crate) bundle: Option, + pub(crate) refresh_in: Duration, +} + +#[derive(Debug, Eq, PartialEq)] +pub(crate) enum StartupLoad { + /// Cloud config does not apply to the current auth, so no refresh is needed. + Inactive, + /// Cloud config applies and must refresh, even when the current bundle is empty. + Active(LoadedBundle), +} + +#[derive(Debug, Eq, PartialEq)] +enum CacheRefreshSchedule { + Stop, + ContinueAfter(Duration), } #[derive(Clone, Copy)] @@ -87,18 +103,6 @@ pub(crate) struct CloudConfigBundleService { timeout: Duration, } -impl Clone for CloudConfigBundleService { - fn clone(&self) -> Self { - Self { - auth_manager: Arc::clone(&self.auth_manager), - client: Arc::clone(&self.client), - cache: self.cache.clone(), - codex_home: self.codex_home.clone(), - timeout: self.timeout, - } - } -} - impl CloudConfigBundleService where C: BundleClient + 'static, @@ -119,32 +123,43 @@ where } } - pub(crate) async fn load_startup_bundle_with_timeout( + pub(crate) async fn load_startup_bundle( &self, - ) -> Result, CloudConfigBundleLoadError> { + ) -> Result { let _timer = codex_otel::start_global_timer("codex.cloud_config_bundle.fetch.duration_ms", &[]); let started_at = Instant::now(); - let load_result = timeout(self.timeout, self.load_startup_bundle()) - .await - .inspect_err(|_| { - let message = format!( - "Timed out waiting for cloud config bundle after {}s", + let load_result = timeout(self.timeout, async { + let Some(auth) = self.auth_manager.auth().await else { + return Ok(StartupLoad::Inactive); + }; + if !cloud_config_eligible_auth(&auth) { + return Ok(StartupLoad::Inactive); + } + + self.load_bundle(auth, BundleFetchTrigger::Startup) + .await + .map(StartupLoad::Active) + }) + .await + .inspect_err(|_| { + let message = format!( + "Timed out waiting for cloud config bundle after {}s", + self.timeout.as_secs() + ); + tracing::error!("{message}"); + emit_load_metric("startup", "error", /*bundle*/ None); + }) + .map_err(|_| { + CloudConfigBundleLoadError::new( + CloudConfigBundleLoadErrorCode::Timeout, + /*status_code*/ None, + format!( + "timed out waiting for cloud config bundle after {}s", self.timeout.as_secs() - ); - tracing::error!("{message}"); - emit_load_metric("startup", "error", /*bundle*/ None); - }) - .map_err(|_| { - CloudConfigBundleLoadError::new( - CloudConfigBundleLoadErrorCode::Timeout, - /*status_code*/ None, - format!( - "timed out waiting for cloud config bundle after {}s", - self.timeout.as_secs() - ), - ) - })?; + ), + ) + })?; let result = match load_result { Ok(result) => result, @@ -154,8 +169,11 @@ where } }; - match result.as_ref() { - Some(bundle) => { + match &result { + StartupLoad::Active(LoadedBundle { + bundle: Some(bundle), + .. + }) => { tracing::info!( elapsed_ms = started_at.elapsed().as_millis(), config_fragments = bundle.config_toml.enterprise_managed.len(), @@ -164,7 +182,7 @@ where ); emit_load_metric("startup", "success", Some(bundle)); } - None => { + StartupLoad::Inactive | StartupLoad::Active(LoadedBundle { bundle: None, .. }) => { tracing::info!( elapsed_ms = started_at.elapsed().as_millis(), "Cloud config bundle load completed (none)" @@ -176,28 +194,17 @@ where Ok(result) } - async fn load_startup_bundle( - &self, - ) -> Result, CloudConfigBundleLoadError> { - let Some(auth) = self.auth_manager.auth().await else { - return Ok(None); - }; - if !cloud_config_eligible_auth(&auth) { - return Ok(None); - } - - self.load_bundle(auth, BundleFetchTrigger::Startup).await - } - - async fn load_valid_cached_bundle(&self, auth: &CodexAuth) -> CachedBundleLookup { + async fn load_valid_cached_bundle(&self, auth: &CodexAuth) -> Option { let (chatgpt_user_id, account_id) = auth_identity(auth); match self .cache .load(chatgpt_user_id.as_deref(), account_id.as_deref()) .await { - Ok(signed_payload) => { - if let Err(err) = validate_bundle(&signed_payload.bundle, &self.codex_home) { + Ok(loaded_cache) => { + if let Err(err) = + validate_bundle(&loaded_cache.signed_payload.bundle, &self.codex_home) + { tracing::warn!( path = %self.cache.path().display(), error = %err, @@ -205,18 +212,21 @@ where ); self.cache .log_load_status(&CacheLoadStatus::CacheInvalidBundle); - CachedBundleLookup::Miss + None } else { tracing::info!( path = %self.cache.path().display(), "Using cached cloud config bundle" ); - CachedBundleLookup::Hit(optional_bundle(signed_payload.bundle)) + Some(LoadedBundle { + bundle: optional_bundle(loaded_cache.signed_payload.bundle), + refresh_in: loaded_cache.refresh_in, + }) } } Err(cache_load_status) => { self.cache.log_load_status(&cache_load_status); - CachedBundleLookup::Miss + None } } } @@ -225,19 +235,17 @@ where &self, auth: CodexAuth, trigger: BundleFetchTrigger, - ) -> Result, CloudConfigBundleLoadError> { + ) -> Result { loop { - if let CachedBundleLookup::Hit(bundle) = self.load_valid_cached_bundle(&auth).await { - return Ok(bundle); + if let Some(loaded) = self.load_valid_cached_bundle(&auth).await { + return Ok(loaded); } match self.cache.try_acquire_lock().await { Ok(CacheLockAttempt::Acquired(_cache_lock)) => { // Close the race between the cache read and lock acquisition. - if let CachedBundleLookup::Hit(bundle) = - self.load_valid_cached_bundle(&auth).await - { - return Ok(bundle); + if let Some(loaded) = self.load_valid_cached_bundle(&auth).await { + return Ok(loaded); } return self .fetch_remote_bundle_and_update_cache_with_retries(auth, trigger) @@ -247,16 +255,14 @@ where sleep(CLOUD_CONFIG_BUNDLE_CACHE_LOCK_RETRY_INTERVAL).await; } Err(err) => { - tracing::error!( + tracing::warn!( path = %self.cache.path().display(), error = %err, - "Failed to acquire cloud config bundle cache lock" + "Failed to acquire cloud config bundle cache lock; fetching without coordination" ); - return Err(CloudConfigBundleLoadError::new( - CloudConfigBundleLoadErrorCode::Internal, - /*status_code*/ None, - format!("failed to acquire cloud config bundle cache lock: {err}"), - )); + return self + .fetch_remote_bundle_and_update_cache_with_retries(auth, trigger) + .await; } } } @@ -266,7 +272,7 @@ where &self, mut auth: CodexAuth, trigger: BundleFetchTrigger, - ) -> Result, CloudConfigBundleLoadError> { + ) -> Result { let trigger_name = match trigger { BundleFetchTrigger::Startup => "startup", BundleFetchTrigger::Refresh => "refresh", @@ -350,7 +356,7 @@ where trigger: &'static str, attempt: usize, bundle: CloudConfigBundle, - ) -> Result, CloudConfigBundleLoadError> { + ) -> Result { emit_fetch_attempt_metric(trigger, attempt, "success", /*status_code*/ None); if let Err(err) = validate_bundle(&bundle, &self.codex_home) { emit_fetch_final_metric( @@ -365,16 +371,20 @@ where } let (chatgpt_user_id, account_id) = auth_identity(auth); - if let Err(err) = self + let refresh_in = match self .cache .save(chatgpt_user_id, account_id, bundle.clone()) .await { - tracing::warn!( - error = %err, - "Failed to write cloud config bundle cache" - ); - } + Ok(refresh_in) => refresh_in, + Err(err) => { + tracing::warn!( + error = %err, + "Failed to write cloud config bundle cache" + ); + CLOUD_CONFIG_BUNDLE_CACHE_REFRESH_RETRY_INTERVAL + } + }; emit_fetch_final_metric( trigger, @@ -384,7 +394,10 @@ where /*status_code*/ None, Some(&bundle), ); - Ok(optional_bundle(bundle)) + Ok(LoadedBundle { + bundle: optional_bundle(bundle), + refresh_in, + }) } async fn retry_after_request_failure( @@ -501,36 +514,54 @@ where )) } - pub(crate) async fn refresh_cache_in_background(&self) { + pub(crate) async fn refresh_cache_in_background( + &self, + mut refresh_in: Duration, + publisher: CloudConfigBundlePublisher, + ) { loop { - // Poll more frequently than the TTL so the task's start time does - // not leave an expired cache unwarmed for another full TTL. Valid - // entries only require a local file read; remote fetches remain - // bounded by the cache TTL and cross-process lock. - sleep(CLOUD_CONFIG_BUNDLE_CACHE_CHECK_INTERVAL).await; - match timeout(self.timeout, self.refresh_cache_once()).await { - Ok(true) => {} - Ok(false) => break, + tokio::select! { + biased; + _ = publisher.closed() => break, + _ = sleep(refresh_in) => {} + } + // A peer may have replaced the shared cache while we slept. + // load_bundle() rechecks it before attempting the lock, then + // returns the peer's later deadline without making a remote call. + let refresh_result = timeout(self.timeout, self.refresh_cache_once(&publisher)).await; + refresh_in = match refresh_result { + Ok(CacheRefreshSchedule::ContinueAfter(refresh_in)) => refresh_in, + Ok(CacheRefreshSchedule::Stop) => break, Err(_) => { tracing::error!( "Timed out refreshing cloud config bundle cache; keeping existing cache" ); emit_load_metric("refresh", "error", /*bundle*/ None); + CLOUD_CONFIG_BUNDLE_CACHE_REFRESH_RETRY_INTERVAL } - } + }; } } - async fn refresh_cache_once(&self) -> bool { + async fn refresh_cache_once( + &self, + publisher: &CloudConfigBundlePublisher, + ) -> CacheRefreshSchedule { let Some(auth) = self.auth_manager.auth().await else { - return false; + return CacheRefreshSchedule::Stop; }; if !cloud_config_eligible_auth(&auth) { - return false; + return CacheRefreshSchedule::Stop; } match self.load_bundle(auth, BundleFetchTrigger::Refresh).await { - Ok(bundle) => emit_load_metric("refresh", "success", bundle.as_ref()), + Ok(loaded) => { + emit_load_metric("refresh", "success", loaded.bundle.as_ref()); + if !publisher.publish(Ok(loaded.bundle)) { + return CacheRefreshSchedule::Stop; + } + CacheRefreshSchedule::ContinueAfter(loaded.refresh_in) + } Err(err) => { tracing::error!( path = %self.cache.path().display(), @@ -538,9 +569,11 @@ where "Failed to refresh cloud config bundle cache" ); emit_load_metric("refresh", "error", /*bundle*/ None); + CacheRefreshSchedule::ContinueAfter( + CLOUD_CONFIG_BUNDLE_CACHE_REFRESH_RETRY_INTERVAL, + ) } } - true } } diff --git a/codex-rs/cloud-config/src/service_tests.rs b/codex-rs/cloud-config/src/service_tests.rs index 84eb18f09f..4c9d636e38 100644 --- a/codex-rs/cloud-config/src/service_tests.rs +++ b/codex-rs/cloud-config/src/service_tests.rs @@ -16,6 +16,8 @@ use base64::engine::general_purpose::URL_SAFE_NO_PAD; use codex_backend_client::ConfigBundleResponse; use codex_backend_client::DeliveredTomlFragment; use codex_config::AbsolutePathBuf; +use codex_config::CloudConfigBundleLoader; +use codex_config::CloudConfigBundlePublisher; use codex_config::CloudConfigFragment; use codex_config::CloudConfigTomlBundle; use codex_config::CloudRequirementsFragment; @@ -207,6 +209,19 @@ fn test_bundle() -> CloudConfigBundle { } } +fn pending_loader() -> (CloudConfigBundleLoader, CloudConfigBundlePublisher) { + let (loader, publisher) = CloudConfigBundleLoader::pending(); + assert!(publisher.publish(Ok(None))); + (loader, publisher) +} + +fn expect_active(load: Result) -> LoadedBundle { + match load.expect("startup load should succeed") { + StartupLoad::Active(loaded) => loaded, + StartupLoad::Inactive => panic!("startup load should be active"), + } +} + fn test_config_fragment() -> CloudConfigFragment { CloudConfigFragment { id: "cfg_1".to_string(), @@ -441,7 +456,10 @@ async fn get_bundle_skips_non_chatgpt_auth() { CLOUD_CONFIG_BUNDLE_TIMEOUT, ); - assert_eq!(service.load_startup_bundle().await, Ok(None)); + assert_eq!( + service.load_startup_bundle().await, + Ok(StartupLoad::Inactive) + ); assert_eq!(fetcher.request_count.load(Ordering::SeqCst), 0); } @@ -456,7 +474,10 @@ async fn get_bundle_skips_individual_plan() { CLOUD_CONFIG_BUNDLE_TIMEOUT, ); - assert_eq!(service.load_startup_bundle().await, Ok(None)); + assert_eq!( + service.load_startup_bundle().await, + Ok(StartupLoad::Inactive) + ); assert_eq!(fetcher.request_count.load(Ordering::SeqCst), 0); } @@ -481,8 +502,8 @@ async fn get_bundle_allows_eligible_workspace_plans_and_writes_cache() { ); assert_eq!( - service.load_startup_bundle().await, - Ok(Some(bundle)), + expect_active(service.load_startup_bundle().await).bundle, + Some(bundle), "plan_type: {plan_type}" ); assert_eq!( @@ -533,7 +554,10 @@ async fn get_bundle_skips_team_like_usage_based_plan() { CLOUD_CONFIG_BUNDLE_TIMEOUT, ); - assert_eq!(service.load_startup_bundle().await, Ok(None)); + assert_eq!( + service.load_startup_bundle().await, + Ok(StartupLoad::Inactive) + ); assert_eq!(fetcher.request_count.load(Ordering::SeqCst), 0); } @@ -586,8 +610,8 @@ async fn get_bundle_ignores_invalid_cache_and_refetches() { ); assert_eq!( - service.load_startup_bundle().await, - Ok(Some(replacement_bundle.clone())) + expect_active(service.load_startup_bundle().await).bundle, + Some(replacement_bundle.clone()) ); assert_eq!(fetcher.request_count.load(Ordering::SeqCst), 1); assert_eq!( @@ -595,6 +619,7 @@ async fn get_bundle_ignores_invalid_cache_and_refetches() { .load(Some("user-12345"), Some("account-12345")) .await .expect("load refreshed cache") + .signed_payload .bundle, replacement_bundle ); @@ -611,7 +636,15 @@ async fn get_bundle_empty_response_is_success_and_cached() { CLOUD_CONFIG_BUNDLE_TIMEOUT, ); - assert_eq!(service.load_startup_bundle().await, Ok(None)); + let loaded = service + .load_startup_bundle() + .await + .expect("empty response should be an active startup load"); + let StartupLoad::Active(LoadedBundle { bundle, refresh_in }) = loaded else { + panic!("eligible auth should keep cloud config refresh active"); + }; + assert_eq!(bundle, None); + assert!(refresh_in > Duration::ZERO); assert_eq!(fetcher.request_count.load(Ordering::SeqCst), 1); assert!( codex_home @@ -622,7 +655,7 @@ async fn get_bundle_empty_response_is_success_and_cached() { } #[tokio::test] -async fn get_bundle_surfaces_cache_lock_failure() { +async fn get_bundle_fetches_and_caches_when_cache_lock_fails() { let codex_home = tempdir().expect("tempdir"); std::fs::create_dir( codex_home @@ -638,12 +671,16 @@ async fn get_bundle_surfaces_cache_lock_failure() { CLOUD_CONFIG_BUNDLE_TIMEOUT, ); - let err = service - .load_startup_bundle() - .await - .expect_err("cache lock failure should fail startup"); - assert_eq!(err.code(), CloudConfigBundleLoadErrorCode::Internal); - assert_eq!(fetcher.request_count.load(Ordering::SeqCst), 0); + let loaded = expect_active(service.load_startup_bundle().await); + assert_eq!(loaded.bundle, Some(test_bundle())); + assert!(loaded.refresh_in > CLOUD_CONFIG_BUNDLE_CACHE_REFRESH_RETRY_INTERVAL); + assert_eq!(fetcher.request_count.load(Ordering::SeqCst), 1); + assert!( + codex_home + .path() + .join(CLOUD_CONFIG_BUNDLE_CACHE_FILENAME) + .exists() + ); } #[tokio::test] @@ -656,7 +693,7 @@ async fn get_bundle_refetches_cache_older_than_ttl() { codex_home.path().to_path_buf(), CLOUD_CONFIG_BUNDLE_TIMEOUT, ); - let _ = prime_service.load_startup_bundle().await; + expect_active(prime_service.load_startup_bundle().await); shift_cache_timestamps( &create_test_cache(codex_home.path()), -chrono::Duration::from_std(CLOUD_CONFIG_BUNDLE_CACHE_TTL + Duration::from_secs(1)) @@ -671,7 +708,10 @@ async fn get_bundle_refetches_cache_older_than_ttl() { CLOUD_CONFIG_BUNDLE_TIMEOUT, ); - assert_eq!(service.load_startup_bundle().await, Ok(Some(bundle))); + assert_eq!( + expect_active(service.load_startup_bundle().await).bundle, + Some(bundle) + ); assert_eq!(fetcher.request_count.load(Ordering::SeqCst), 1); } @@ -685,7 +725,7 @@ async fn get_bundle_ignores_cache_for_different_auth_identity() { codex_home.path().to_path_buf(), CLOUD_CONFIG_BUNDLE_TIMEOUT, ); - let _ = prime_service.load_startup_bundle().await; + expect_active(prime_service.load_startup_bundle().await); let replacement_bundle = CloudConfigBundle { config_toml: CloudConfigTomlBundle::default(), @@ -709,8 +749,8 @@ async fn get_bundle_ignores_cache_for_different_auth_identity() { ); assert_eq!( - service.load_startup_bundle().await, - Ok(Some(replacement_bundle)) + expect_active(service.load_startup_bundle().await).bundle, + Some(replacement_bundle) ); assert_eq!(fetcher.request_count.load(Ordering::SeqCst), 1); } @@ -724,7 +764,7 @@ async fn get_bundle_times_out() { codex_home.path().to_path_buf(), CLOUD_CONFIG_BUNDLE_TIMEOUT, ); - let handle = tokio::spawn(async move { service.load_startup_bundle_with_timeout().await }); + let handle = tokio::spawn(async move { service.load_startup_bundle().await }); tokio::time::advance(CLOUD_CONFIG_BUNDLE_TIMEOUT + Duration::from_millis(1)).await; let result = handle.await.expect("cloud config bundle task"); @@ -753,7 +793,10 @@ async fn get_bundle_retries_until_success() { tokio::task::yield_now().await; tokio::time::advance(Duration::from_secs(1)).await; - assert_eq!(handle.await.expect("bundle task"), Ok(Some(test_bundle()))); + assert_eq!( + expect_active(handle.await.expect("bundle task")).bundle, + Some(test_bundle()) + ); assert_eq!(fetcher.request_count.load(Ordering::SeqCst), 2); } @@ -812,7 +855,10 @@ async fn get_bundle_recovers_after_unauthorized_reload() { CLOUD_CONFIG_BUNDLE_TIMEOUT, ); - assert_eq!(service.load_startup_bundle().await, Ok(Some(test_bundle()))); + assert_eq!( + expect_active(service.load_startup_bundle().await).bundle, + Some(test_bundle()) + ); assert_eq!(fetcher.request_count.load(Ordering::SeqCst), 2); } @@ -869,13 +915,17 @@ async fn get_bundle_recovers_after_unauthorized_reload_updates_cache_identity() CLOUD_CONFIG_BUNDLE_TIMEOUT, ); - assert_eq!(service.load_startup_bundle().await, Ok(Some(test_bundle()))); + assert_eq!( + expect_active(service.load_startup_bundle().await).bundle, + Some(test_bundle()) + ); let cache = create_test_cache(codex_home.path()); assert_eq!( cache .load(Some("user-99999"), Some("account-12345")) .await .expect("load cache") + .signed_payload .bundle, test_bundle() ); @@ -1012,7 +1062,7 @@ async fn get_bundle_does_not_use_cache_when_auth_identity_is_incomplete() { codex_home.path().to_path_buf(), CLOUD_CONFIG_BUNDLE_TIMEOUT, ); - let _ = prime_service.load_startup_bundle().await; + expect_active(prime_service.load_startup_bundle().await); let replacement_bundle = CloudConfigBundle { config_toml: CloudConfigTomlBundle::default(), @@ -1040,8 +1090,8 @@ async fn get_bundle_does_not_use_cache_when_auth_identity_is_incomplete() { ); assert_eq!( - service.load_startup_bundle().await, - Ok(Some(replacement_bundle)) + expect_active(service.load_startup_bundle().await).bundle, + Some(replacement_bundle) ); assert_eq!(fetcher.request_count.load(Ordering::SeqCst), 1); } @@ -1087,8 +1137,8 @@ async fn refresh_skips_remote_fetch_when_shared_cache_was_refreshed_recently() { CLOUD_CONFIG_BUNDLE_TIMEOUT, ); assert_eq!( - prime_service.load_startup_bundle().await, - Ok(Some(test_bundle())) + expect_active(prime_service.load_startup_bundle().await).bundle, + Some(test_bundle()) ); let fetcher = Arc::new(SequenceBundleClient::new(vec![Err(request_error())])); @@ -1099,8 +1149,93 @@ async fn refresh_skips_remote_fetch_when_shared_cache_was_refreshed_recently() { CLOUD_CONFIG_BUNDLE_TIMEOUT, ); - assert!(service.refresh_cache_once().await); + let (loader, publisher) = pending_loader(); + let CacheRefreshSchedule::ContinueAfter(refresh_in) = + service.refresh_cache_once(&publisher).await + else { + panic!("refresh should remain scheduled"); + }; + assert!(refresh_in > CLOUD_CONFIG_BUNDLE_CACHE_REFRESH_RETRY_INTERVAL); assert_eq!(fetcher.request_count.load(Ordering::SeqCst), 0); + assert_eq!(loader.get().await, Ok(Some(test_bundle()))); +} + +#[tokio::test(start_paused = true)] +async fn background_refresh_stops_when_loader_is_dropped() { + let codex_home = tempdir().expect("tempdir"); + let fetcher = Arc::new(StaticBundleClient::new(test_bundle())); + let service = CloudConfigBundleService::new( + auth_manager_with_plan("business").await, + fetcher.clone(), + codex_home.path().to_path_buf(), + CLOUD_CONFIG_BUNDLE_TIMEOUT, + ); + let (loader, publisher) = pending_loader(); + let refresh_task = tokio::spawn(async move { + service + .refresh_cache_in_background(CLOUD_CONFIG_BUNDLE_CACHE_TTL, publisher) + .await; + }); + + drop(loader); + tokio::time::timeout(Duration::from_secs(1), refresh_task) + .await + .expect("loader drop should stop refresh") + .expect("refresh task"); + assert_eq!(fetcher.request_count.load(Ordering::SeqCst), 0); +} + +#[tokio::test(start_paused = true)] +async fn refresh_failure_uses_retry_interval() { + let codex_home = tempdir().expect("tempdir"); + let fetcher = Arc::new(SequenceBundleClient::new(vec![ + Err(request_error()); + CLOUD_CONFIG_BUNDLE_MAX_ATTEMPTS + ])); + let service = CloudConfigBundleService::new( + auth_manager_with_plan("business").await, + fetcher.clone(), + codex_home.path().to_path_buf(), + CLOUD_CONFIG_BUNDLE_TIMEOUT, + ); + + let (_loader, publisher) = pending_loader(); + let refresh = tokio::spawn(async move { service.refresh_cache_once(&publisher).await }); + tokio::task::yield_now().await; + tokio::time::advance(Duration::from_secs(5)).await; + tokio::task::yield_now().await; + + assert_eq!( + refresh.await.expect("refresh task"), + CacheRefreshSchedule::ContinueAfter(CLOUD_CONFIG_BUNDLE_CACHE_REFRESH_RETRY_INTERVAL) + ); + assert_eq!( + fetcher.request_count.load(Ordering::SeqCst), + CLOUD_CONFIG_BUNDLE_MAX_ATTEMPTS + ); +} + +#[tokio::test] +async fn startup_uses_retry_interval_when_cache_write_fails() { + let codex_home = tempdir().expect("tempdir"); + let cache = create_test_cache(codex_home.path()); + std::fs::create_dir(cache.path()).expect("create directory at cache path"); + let fetcher = Arc::new(StaticBundleClient::new(test_bundle())); + let service = CloudConfigBundleService::new( + auth_manager_with_plan("business").await, + fetcher.clone(), + codex_home.path().to_path_buf(), + CLOUD_CONFIG_BUNDLE_TIMEOUT, + ); + + assert_eq!( + service.load_startup_bundle().await, + Ok(StartupLoad::Active(LoadedBundle { + bundle: Some(test_bundle()), + refresh_in: CLOUD_CONFIG_BUNDLE_CACHE_REFRESH_RETRY_INTERVAL, + })) + ); + assert_eq!(fetcher.request_count.load(Ordering::SeqCst), 1); } #[tokio::test] @@ -1125,8 +1260,13 @@ async fn refresh_fetches_remote_when_cache_timestamp_is_in_future() { CLOUD_CONFIG_BUNDLE_TIMEOUT, ); - assert!(service.refresh_cache_once().await); + let (loader, publisher) = pending_loader(); + assert!(matches!( + service.refresh_cache_once(&publisher).await, + CacheRefreshSchedule::ContinueAfter(_) + )); assert_eq!(fetcher.request_count.load(Ordering::SeqCst), 1); + assert_eq!(loader.get().await, Ok(Some(test_bundle()))); } #[tokio::test] @@ -1154,18 +1294,18 @@ async fn concurrent_startups_make_one_remote_request() { fetcher.release_request.notify_one(); assert_eq!( - first_load.await.expect("first load task"), - Ok(Some(test_bundle())) + expect_active(first_load.await.expect("first load task")).bundle, + Some(test_bundle()) ); assert_eq!( - second_load.await.expect("second load task"), - Ok(Some(test_bundle())) + expect_active(second_load.await.expect("second load task")).bundle, + Some(test_bundle()) ); assert_eq!(fetcher.request_count.load(Ordering::SeqCst), 1); } #[tokio::test] -async fn refresh_skips_remote_fetch_when_cache_lock_fails() { +async fn refresh_fetches_and_caches_when_cache_lock_fails() { let codex_home = tempdir().expect("tempdir"); let cache = create_test_cache(codex_home.path()); cache @@ -1196,8 +1336,24 @@ async fn refresh_skips_remote_fetch_when_cache_lock_fails() { CLOUD_CONFIG_BUNDLE_TIMEOUT, ); - assert!(service.refresh_cache_once().await); - assert_eq!(fetcher.request_count.load(Ordering::SeqCst), 0); + let (loader, publisher) = pending_loader(); + let CacheRefreshSchedule::ContinueAfter(refresh_in) = + service.refresh_cache_once(&publisher).await + else { + panic!("refresh should remain scheduled"); + }; + assert!(refresh_in > CLOUD_CONFIG_BUNDLE_CACHE_REFRESH_RETRY_INTERVAL); + assert_eq!(fetcher.request_count.load(Ordering::SeqCst), 1); + assert_eq!(loader.get().await, Ok(Some(test_bundle()))); + assert_eq!( + cache + .load(Some("user-12345"), Some("account-12345")) + .await + .expect("load refreshed cache") + .signed_payload + .bundle, + test_bundle() + ); } #[tokio::test] @@ -1224,20 +1380,28 @@ async fn refresh_from_remote_updates_stale_cached_bundle() { CLOUD_CONFIG_BUNDLE_TIMEOUT, ); - assert_eq!(service.load_startup_bundle().await, Ok(Some(test_bundle()))); + assert_eq!( + expect_active(service.load_startup_bundle().await).bundle, + Some(test_bundle()) + ); shift_cache_timestamps( &create_test_cache(codex_home.path()), -chrono::Duration::from_std(CLOUD_CONFIG_BUNDLE_CACHE_TTL + Duration::from_secs(1)) .expect("cache age should fit chrono duration"), ); - assert!(service.refresh_cache_once().await); + let (loader, publisher) = pending_loader(); + assert!(matches!( + service.refresh_cache_once(&publisher).await, + CacheRefreshSchedule::ContinueAfter(_) + )); let cache = create_test_cache(codex_home.path()); - let signed_payload = cache + let loaded_cache = cache .load(Some("user-12345"), Some("account-12345")) .await .expect("load cache"); - assert_eq!(signed_payload.bundle, replacement_bundle); + assert_eq!(loaded_cache.signed_payload.bundle, replacement_bundle); + assert_eq!(loader.get().await, Ok(Some(replacement_bundle))); } #[test] diff --git a/codex-rs/config/Cargo.toml b/codex-rs/config/Cargo.toml index 4fe529d743..dc7124cc78 100644 --- a/codex-rs/config/Cargo.toml +++ b/codex-rs/config/Cargo.toml @@ -25,7 +25,6 @@ codex-utils-absolute-path = { workspace = true } codex-utils-path = { workspace = true } codex-utils-path-uri = { workspace = true } dunce = { workspace = true } -futures = { workspace = true, features = ["alloc", "std"] } gethostname = { workspace = true } indexmap = { workspace = true, features = ["serde"] } multimap = { workspace = true } @@ -38,7 +37,7 @@ serde_json = { workspace = true } serde_path_to_error = { workspace = true } sha2 = { workspace = true } thiserror = { workspace = true } -tokio = { workspace = true, features = ["fs"] } +tokio = { workspace = true, features = ["fs", "sync"] } toml = { workspace = true, features = ["preserve_order"] } toml_edit = { workspace = true } tonic = { workspace = true } diff --git a/codex-rs/config/src/cloud_config_bundle.rs b/codex-rs/config/src/cloud_config_bundle.rs index cfece32a01..c02407d44c 100644 --- a/codex-rs/config/src/cloud_config_bundle.rs +++ b/codex-rs/config/src/cloud_config_bundle.rs @@ -12,14 +12,11 @@ use crate::cloud_config_layers::CloudConfigLayerError; use crate::cloud_config_layers::cloud_config_layers_from_fragments_strict; use crate::cloud_config_layers_from_fragments; use codex_utils_absolute_path::AbsolutePathBuf; -use futures::future::BoxFuture; -use futures::future::FutureExt; -use futures::future::Shared; use serde::Deserialize; use serde::Serialize; use std::fmt; -use std::future::Future; use thiserror::Error; +use tokio::sync::watch; #[derive(Clone, Debug, Default, Deserialize, Eq, PartialEq, Serialize)] pub struct CloudConfigBundle { @@ -177,23 +174,86 @@ impl CloudConfigBundleLoadError { #[derive(Clone)] pub struct CloudConfigBundleLoader { - fut: Shared, CloudConfigBundleLoadError>>>, + state: watch::Receiver, +} + +#[derive(Clone, Debug)] +enum CloudConfigBundleState { + Pending, + Ready(Result, CloudConfigBundleLoadError>), +} + +/// Publishes results to a [`CloudConfigBundleLoader`]. +#[derive(Debug)] +pub struct CloudConfigBundlePublisher { + state: watch::Sender, } impl CloudConfigBundleLoader { - pub fn new(fut: F) -> Self - where - F: Future, CloudConfigBundleLoadError>> - + Send - + 'static, - { - Self { - fut: fut.boxed().shared(), - } + /// Creates a loader that immediately returns the supplied result. + pub fn from_result( + result: Result, CloudConfigBundleLoadError>, + ) -> Self { + let (_publisher, state) = watch::channel(CloudConfigBundleState::Ready(result)); + Self { state } + } + + /// Creates a pending loader and its publisher. + pub fn pending() -> (Self, CloudConfigBundlePublisher) { + let (publisher, state) = watch::channel(CloudConfigBundleState::Pending); + ( + Self { state }, + CloudConfigBundlePublisher { state: publisher }, + ) } pub async fn get(&self) -> Result, CloudConfigBundleLoadError> { - self.fut.clone().await + let mut state = self.state.clone(); + loop { + let result = match &*state.borrow_and_update() { + CloudConfigBundleState::Pending => None, + CloudConfigBundleState::Ready(result) => Some(result.clone()), + }; + if let Some(result) = result { + return result; + } + if state.changed().await.is_err() { + return Err(CloudConfigBundleLoadError::new( + CloudConfigBundleLoadErrorCode::Internal, + /*status_code*/ None, + "cloud config bundle lifecycle ended before startup completed", + )); + } + } + } +} + +impl CloudConfigBundlePublisher { + /// Waits until every associated loader and clone has been dropped. + pub async fn closed(&self) { + self.state.closed().await; + } + + /// Publishes the latest load result. + /// + /// An error resolves a pending loader, but does not replace a result that + /// was already published. Successful results always replace the current + /// value. Returns `false` if every associated loader was dropped. + pub fn publish( + &self, + result: Result, CloudConfigBundleLoadError>, + ) -> bool { + if self.state.is_closed() { + return false; + } + self.state.send_if_modified(move |state| { + if matches!(state, CloudConfigBundleState::Ready(_)) && result.is_err() { + return false; + } + *state = CloudConfigBundleState::Ready(result); + true + }); + !self.state.is_closed() } } @@ -205,7 +265,7 @@ impl fmt::Debug for CloudConfigBundleLoader { impl Default for CloudConfigBundleLoader { fn default() -> Self { - Self::new(async { Ok(None) }) + Self::from_result(Ok(None)) } } diff --git a/codex-rs/config/src/cloud_config_bundle_tests.rs b/codex-rs/config/src/cloud_config_bundle_tests.rs index 714aca6bc6..f1b78547b3 100644 --- a/codex-rs/config/src/cloud_config_bundle_tests.rs +++ b/codex-rs/config/src/cloud_config_bundle_tests.rs @@ -4,23 +4,104 @@ use crate::ConfigRequirementsToml; use crate::compose_requirements; use codex_protocol::protocol::AskForApproval; use pretty_assertions::assert_eq; -use std::sync::Arc; -use std::sync::atomic::AtomicUsize; -use std::sync::atomic::Ordering; use tempfile::tempdir; -#[tokio::test] -async fn shared_future_runs_once() { - let counter = Arc::new(AtomicUsize::new(0)); - let counter_clone = Arc::clone(&counter); - let loader = CloudConfigBundleLoader::new(async move { - counter_clone.fetch_add(1, Ordering::SeqCst); - Ok(Some(CloudConfigBundle::default())) - }); +fn bundle_with_model(model: &str) -> CloudConfigBundle { + CloudConfigBundle { + config_toml: CloudConfigTomlBundle { + enterprise_managed: vec![CloudConfigFragment { + id: model.to_string(), + name: model.to_string(), + contents: format!("model = \"{model}\""), + }], + }, + ..Default::default() + } +} - let (first, second) = tokio::join!(loader.get(), loader.get()); - assert_eq!(first, second); - assert_eq!(counter.load(Ordering::SeqCst), 1); +#[tokio::test] +async fn pending_loader_receives_initial_result() { + let bundle = bundle_with_model("initial"); + let (loader, publisher) = CloudConfigBundleLoader::pending(); + let cloned_loader = loader.clone(); + let load_task = tokio::spawn(async move { tokio::join!(loader.get(), cloned_loader.get()) }); + tokio::task::yield_now().await; + + assert!(publisher.publish(Ok(Some(bundle.clone())))); + + assert_eq!( + tokio::time::timeout(std::time::Duration::from_secs(1), load_task) + .await + .expect("initial result should wake loaders") + .expect("loader task"), + (Ok(Some(bundle.clone())), Ok(Some(bundle))) + ); +} + +#[tokio::test] +async fn successful_result_replaces_error_and_later_error_is_ignored() { + let bundle = bundle_with_model("recovered"); + let initial_error = CloudConfigBundleLoadError::new( + CloudConfigBundleLoadErrorCode::RequestFailed, + /*status_code*/ None, + "initial load failed", + ); + let (loader, publisher) = CloudConfigBundleLoader::pending(); + + assert!(publisher.publish(Err(initial_error.clone()))); + assert_eq!(loader.get().await, Err(initial_error)); + + assert!(publisher.publish(Ok(Some(bundle.clone())))); + assert_eq!(loader.get().await, Ok(Some(bundle.clone()))); + + assert!(publisher.publish(Err(CloudConfigBundleLoadError::new( + CloudConfigBundleLoadErrorCode::RequestFailed, + /*status_code*/ None, + "refresh failed", + )))); + assert_eq!(loader.get().await, Ok(Some(bundle))); +} + +#[tokio::test] +async fn refresh_updates_only_loader_and_its_clones() { + let initial_bundle = bundle_with_model("initial"); + let refreshed_bundle = bundle_with_model("refreshed"); + let independent_loader = CloudConfigBundleLoader::from_result(Ok(None)); + let (loader, publisher) = CloudConfigBundleLoader::pending(); + let cloned_loader = loader.clone(); + + assert!(publisher.publish(Ok(Some(initial_bundle.clone())))); + assert_eq!(loader.get().await, Ok(Some(initial_bundle))); + + assert!(publisher.publish(Ok(Some(refreshed_bundle.clone())))); + + assert_eq!(loader.get().await, Ok(Some(refreshed_bundle.clone()))); + assert_eq!(cloned_loader.get().await, Ok(Some(refreshed_bundle))); + assert_eq!(independent_loader.get().await, Ok(None)); + + assert!(publisher.publish(Ok(None))); + + assert_eq!(loader.get().await, Ok(None)); + + drop(cloned_loader); + drop(loader); + assert!(!publisher.publish(Ok(Some(CloudConfigBundle::default())))); +} + +#[tokio::test] +async fn pending_loader_reports_dropped_publisher() { + let (loader, publisher) = CloudConfigBundleLoader::pending(); + + drop(publisher); + + assert_eq!( + loader.get().await, + Err(CloudConfigBundleLoadError::new( + CloudConfigBundleLoadErrorCode::Internal, + /*status_code*/ None, + "cloud config bundle lifecycle ended before startup completed", + )) + ); } #[test] diff --git a/codex-rs/config/src/lib.rs b/codex-rs/config/src/lib.rs index df3e45d76d..d46125df4f 100644 --- a/codex-rs/config/src/lib.rs +++ b/codex-rs/config/src/lib.rs @@ -38,6 +38,7 @@ pub use cloud_config_bundle::CloudConfigBundleLayers; pub use cloud_config_bundle::CloudConfigBundleLoadError; pub use cloud_config_bundle::CloudConfigBundleLoadErrorCode; pub use cloud_config_bundle::CloudConfigBundleLoader; +pub use cloud_config_bundle::CloudConfigBundlePublisher; pub use cloud_config_bundle::CloudConfigTomlBundle; pub use cloud_config_bundle::CloudRequirementsFragment; pub use cloud_config_bundle::CloudRequirementsTomlBundle; diff --git a/codex-rs/config/src/loader/tests.rs b/codex-rs/config/src/loader/tests.rs index 068e29aa6e..c156251206 100644 --- a/codex-rs/config/src/loader/tests.rs +++ b/codex-rs/config/src/loader/tests.rs @@ -1,4 +1,11 @@ use super::*; +use crate::CloudConfigBundle; +use crate::CloudConfigBundleLoader; +use crate::CloudConfigFragment; +use crate::CloudConfigTomlBundle; +use crate::CloudRequirementsFragment; +use crate::CloudRequirementsTomlBundle; +use crate::ConfigRequirementsToml; use codex_file_system::CopyOptions; use codex_file_system::CreateDirectoryOptions; use codex_file_system::ExecutorFileSystemFuture; @@ -104,6 +111,77 @@ impl ExecutorFileSystem for TestFileSystem { } } +#[tokio::test] +async fn refreshed_cloud_bundle_applies_to_future_layer_loads() { + let tmp = tempdir().expect("tempdir"); + let bundle = |suffix: &str, model: &str, approval_policy: &str| CloudConfigBundle { + config_toml: CloudConfigTomlBundle { + enterprise_managed: vec![CloudConfigFragment { + id: format!("cfg_{suffix}"), + name: format!("Config {suffix}"), + contents: format!("model = \"{model}\""), + }], + }, + requirements_toml: CloudRequirementsTomlBundle { + enterprise_managed: vec![CloudRequirementsFragment { + id: format!("req_{suffix}"), + name: format!("Requirements {suffix}"), + contents: format!("allowed_approval_policies = [\"{approval_policy}\"]"), + }], + }, + }; + let initial_bundle = bundle("initial", "initial-model", "never"); + let refreshed_bundle = bundle("refreshed", "refreshed-model", "on-request"); + let (loader, publisher) = CloudConfigBundleLoader::pending(); + assert!(publisher.publish(Ok(Some(initial_bundle)))); + let mut loader_overrides = LoaderOverrides::without_managed_config_for_tests(); + loader_overrides.managed_config_path = Some(tmp.path().join("managed_config.toml")); + loader_overrides.system_config_path = Some(tmp.path().join("system-config.toml")); + loader_overrides.system_requirements_path = Some(tmp.path().join("requirements.toml")); + let load_layers = || { + load_config_layers_state( + &TestFileSystem, + tmp.path(), + /*cwd*/ None, + &[], + ConfigLoadOptions { + loader_overrides: loader_overrides.clone(), + strict_config: false, + cloud_config_bundle: loader.clone(), + }, + &crate::NoopThreadConfigLoader, + ) + }; + + let initial_layers = load_layers().await.expect("load initial layers"); + assert_eq!( + initial_layers.effective_config(), + toml::from_str("model = \"initial-model\"").expect("expected initial config") + ); + assert_eq!( + initial_layers.requirements_toml(), + &ConfigRequirementsToml { + allowed_approval_policies: Some(vec![AskForApproval::Never]), + ..Default::default() + } + ); + + assert!(publisher.publish(Ok(Some(refreshed_bundle)))); + + let refreshed_layers = load_layers().await.expect("load refreshed layers"); + assert_eq!( + refreshed_layers.effective_config(), + toml::from_str("model = \"refreshed-model\"").expect("expected refreshed config") + ); + assert_eq!( + refreshed_layers.requirements_toml(), + &ConfigRequirementsToml { + allowed_approval_policies: Some(vec![AskForApproval::OnRequest]), + ..Default::default() + } + ); +} + #[tokio::test] async fn profile_v2_rejects_matching_legacy_profile_in_base_user_config() { let tmp = tempdir().expect("tempdir"); diff --git a/codex-rs/config/src/test_support.rs b/codex-rs/config/src/test_support.rs index 44627f3892..54b93ba97c 100644 --- a/codex-rs/config/src/test_support.rs +++ b/codex-rs/config/src/test_support.rs @@ -71,7 +71,7 @@ impl CloudConfigBundleFixture { pub fn into_loader(self) -> CloudConfigBundleLoader { let bundle = self.into_bundle(); - CloudConfigBundleLoader::new(async move { Ok(Some(bundle)) }) + CloudConfigBundleLoader::from_result(Ok(Some(bundle))) } } diff --git a/codex-rs/core/src/config/config_loader_tests.rs b/codex-rs/core/src/config/config_loader_tests.rs index 9857e08f16..c2250ca2af 100644 --- a/codex-rs/core/src/config/config_loader_tests.rs +++ b/codex-rs/core/src/config/config_loader_tests.rs @@ -2235,13 +2235,13 @@ async fn load_config_layers_fails_when_cloud_config_bundle_loader_fails() -> any Some(cwd), &[] as &[(String, TomlValue)], ConfigLoadOptions { - cloud_config_bundle: CloudConfigBundleLoader::new(async { - Err(CloudConfigBundleLoadError::new( + cloud_config_bundle: CloudConfigBundleLoader::from_result(Err( + CloudConfigBundleLoadError::new( codex_config::CloudConfigBundleLoadErrorCode::RequestFailed, /*status_code*/ None, "cloud config bundle failed", - )) - }), + ), + )), ..Default::default() }, &codex_config::NoopThreadConfigLoader,