diff --git a/codex-rs/cloud-config/src/bundle_loader.rs b/codex-rs/cloud-config/src/bundle_loader.rs index 5b027d21cd..0d7bedc91f 100644 --- a/codex-rs/cloud-config/src/bundle_loader.rs +++ b/codex-rs/cloud-config/src/bundle_loader.rs @@ -13,7 +13,7 @@ use tokio::task::JoinHandle; struct CloudConfigBundleLoaderLifetime { service: Arc>, - refresh_task: JoinHandle<()>, + refresh_task: AbortHandle, } impl Drop for CloudConfigBundleLoaderLifetime { @@ -43,7 +43,7 @@ pub fn cloud_config_bundle_loader( pub(crate) fn cloud_config_bundle_loader_for_service( service: CloudConfigBundleService, -) -> (CloudConfigBundleLoader, AbortHandle) +) -> (CloudConfigBundleLoader, JoinHandle<()>) where C: BundleClient + 'static, { @@ -56,7 +56,7 @@ where let abort_handle = refresh_task.abort_handle(); let lifetime = Arc::new(CloudConfigBundleLoaderLifetime { service, - refresh_task, + refresh_task: abort_handle, }); let snapshot_lifetime = Arc::clone(&lifetime); @@ -69,7 +69,7 @@ where let lifetime = Arc::clone(&snapshot_lifetime); async move { lifetime.service.get_latest_snapshot().await } }); - (loader, abort_handle) + (loader, refresh_task) } pub async fn cloud_config_bundle_loader_for_storage( diff --git a/codex-rs/cloud-config/src/service_tests.rs b/codex-rs/cloud-config/src/service_tests.rs index 08086fded5..b3f00454cf 100644 --- a/codex-rs/cloud-config/src/service_tests.rs +++ b/codex-rs/cloud-config/src/service_tests.rs @@ -1564,14 +1564,13 @@ async fn each_loader_owns_refresh_until_its_last_clone_is_dropped() { .expect("both loader owners should keep refreshing"); assert_eq!(fetcher.request_count.load(Ordering::SeqCst), 3); + // Aborting is asynchronous; joining waits for cancellation to finish. drop(cloned_loader); - tokio::task::yield_now().await; - assert!(original_task.is_finished()); + assert!(original_task.await.unwrap_err().is_cancelled()); assert!(!replacement_task.is_finished()); drop(replacement_loader); - tokio::task::yield_now().await; - assert!(replacement_task.is_finished()); + assert!(replacement_task.await.unwrap_err().is_cancelled()); tokio::time::advance(CLOUD_CONFIG_BUNDLE_CACHE_REFRESH_INTERVAL + Duration::from_secs(1)).await; tokio::task::yield_now().await; assert_eq!(replacement_fetcher.request_count.load(Ordering::SeqCst), 2);