From ca87ef86e101c4e0b18847116fadcff078abaeaa Mon Sep 17 00:00:00 2001
From: goodman-oai <goodman@openai.com>
Date: Thu, 22 Oct 2026 21:50:34 +0000
Subject: [PATCH] Wait for cloud config refresh cancellation in loader lifetime
 tests (#56274)

## Why

Aborting a refresh task is asynchronous, so yielding once before checking `is_finished()` does not guarantee cancellation has completed.

## What changed

Return the refresh task's `JoinHandle` from `cloud_config_bundle_loader_for_service` and retain its `AbortHandle` in the loader lifetime. Update `each_loader_owns_refresh_until_its_last_clone_is_dropped` to await each task and assert cancellation after dropping its last loader clone, while still checking that the replacement loader's task remains active.

GitOrigin-RevId: 2eb65211aefc21f9195ff766d8793f4b8f5a74c4
---
 codex-rs/cloud-config/src/bundle_loader.rs | 8 ++++----
 codex-rs/cloud-config/src/service_tests.rs | 7 +++----
 2 files changed, 7 insertions(+), 8 deletions(-)

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<C> {
     service: Arc<CloudConfigBundleService<C>>,
-    refresh_task: JoinHandle<()>,
+    refresh_task: AbortHandle,
 }
 
 impl<C> Drop for CloudConfigBundleLoaderLifetime<C> {
@@ -43,7 +43,7 @@ pub fn cloud_config_bundle_loader(
 
 pub(crate) fn cloud_config_bundle_loader_for_service<C>(
     service: CloudConfigBundleService<C>,
-) -> (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);