Stabilize paused-time code-mode service tests (#50443)

## Why

Tokio's paused clock can auto-advance while tests await the V8 thread. The service tests need explicit control over time when checking yield limits and grace periods.

## What changed

- Hold virtual time with a blocking-task guard that is released when dropped.
- Replace bounded polling loops with tool-start notifications and direct awaits of response tasks.
- Return an error from the test helper if the cell closes before a nested tool starts.

## Testing

Add tests covering tool-start acknowledgement when a cell throws or is terminated before invoking a nested tool, including assertions on the resulting cell responses.

GitOrigin-RevId: c221c0501ec0354052102c0211cc651008babff3

codex-rs/code-mode-runtime/src/service_tests.rs

@@ -1,7 +1,5 @@
 use codex_code_mode_protocol::NoopCodeModeSessionDelegate;
 use std::sync::Arc;
-use std::sync::atomic::AtomicBool;
-use std::sync::atomic::Ordering;
 use std::time::Duration;
 
 use super::CellId;
@@ -26,6 +24,7 @@ use codex_protocol::ToolName;
 use pretty_assertions::assert_eq;
 use serde_json::Value as JsonValue;
 use tokio::sync::Notify;
+use tokio::sync::oneshot;
 use tokio_util::sync::CancellationToken;
 
 #[test]
@@ -60,6 +59,7 @@ fn resolve_yield_timeout_applies_grace_before_session_limits() {
 
 #[tokio::test(start_paused = true)]
 async fn interrupt_yields_execute_and_wait_without_stopping_the_cell() {
+    let _time_guard = hold_virtual_time();
     let delegate = Arc::new(ReleasableToolDelegate::default());
     let service: Arc<dyn CodeModeSession> = Arc::new(InProcessCodeModeSession::new());
     let execute_signal = CancellationToken::new();
@@ -77,9 +77,8 @@ async fn interrupt_yields_execute_and_wait_without_stopping_the_cell() {
         .await
         .unwrap();
     let response = tokio::spawn(started.initial_response());
-    wait_until_tool_started(&delegate).await;
+    wait_until_tool_started(&delegate).await.unwrap();
     execute_signal.cancel();
-    wait_until_finished(&response).await;
     assert_eq!(
         response.await.unwrap().unwrap(),
         RuntimeResponse::Yielded {
@@ -108,7 +107,6 @@ async fn interrupt_yields_execute_and_wait_without_stopping_the_cell() {
     tokio::task::yield_now().await;
     assert!(!response.is_finished());
     wait_signal.cancel();
-    wait_until_finished(&response).await;
     assert_eq!(
         response.await.unwrap().unwrap(),
         WaitOutcome::LiveCell(RuntimeResponse::Yielded {
@@ -133,7 +131,6 @@ async fn interrupt_yields_execute_and_wait_without_stopping_the_cell() {
                 .await
         }
     });
-    wait_until_finished(&response).await;
     assert_eq!(
         response.await.unwrap().unwrap(),
         WaitOutcome::LiveCell(RuntimeResponse::Result {
@@ -149,6 +146,7 @@ async fn interrupt_yields_execute_and_wait_without_stopping_the_cell() {
 
 #[tokio::test(start_paused = true)]
 async fn execute_waits_for_nested_tool_during_yield_grace() {
+    let _time_guard = hold_virtual_time();
     let delegate = Arc::new(ReleasableToolDelegate::default());
     let service = InProcessCodeModeSession::new();
     let request = ExecuteRequest {
@@ -162,10 +160,9 @@ async fn execute_waits_for_nested_tool_during_yield_grace() {
         .await
         .unwrap();
     let response = tokio::spawn(started.initial_response());
-    wait_until_tool_started(&delegate).await;
+    wait_until_tool_started(&delegate).await.unwrap();
     tokio::time::advance(Duration::from_millis(10_500)).await;
     delegate.release_tool();
-    wait_until_finished(&response).await;
     let response = response.await.unwrap().unwrap();
 
     assert_eq!(
@@ -183,6 +180,7 @@ async fn execute_waits_for_nested_tool_during_yield_grace() {
 
 #[tokio::test(start_paused = true)]
 async fn execute_and_wait_clamp_yield_grace_without_stopping_the_cell() {
+    let _time_guard = hold_virtual_time();
     let delegate = Arc::new(ReleasableToolDelegate::default());
     let service = InProcessCodeModeSession::with_limits(CodeModeSessionCellExecutionLimits {
         max_yield_time_ms: Some(/*value*/ 10_000),
@@ -202,12 +200,11 @@ async fn execute_and_wait_clamp_yield_grace_without_stopping_the_cell() {
         .await
         .unwrap();
     let initial_response = tokio::spawn(started.initial_response());
-    wait_until_tool_started(&delegate).await;
+    wait_until_tool_started(&delegate).await.unwrap();
 
     tokio::time::advance(Duration::from_millis(9_999)).await;
     assert!(!initial_response.is_finished());
     tokio::time::advance(Duration::from_millis(/*millis*/ 1)).await;
-    wait_until_finished(&initial_response).await;
     assert_eq!(
         initial_response.await.unwrap().unwrap(),
         RuntimeResponse::Yielded {
@@ -229,7 +226,6 @@ async fn execute_and_wait_clamp_yield_grace_without_stopping_the_cell() {
     let wait_response = tokio::spawn(wait_response);
     tokio::task::yield_now().await;
     tokio::time::advance(Duration::from_secs(/*secs*/ 10)).await;
-    wait_until_finished(&wait_response).await;
     assert_eq!(
         wait_response.await.unwrap().unwrap(),
         WaitOutcome::LiveCell(RuntimeResponse::Yielded {
@@ -250,7 +246,6 @@ async fn execute_and_wait_clamp_yield_grace_without_stopping_the_cell() {
         )
         .await;
     let completion = tokio::spawn(completion);
-    wait_until_finished(&completion).await;
     assert_eq!(
         completion.await.unwrap().unwrap(),
         WaitOutcome::LiveCell(RuntimeResponse::Result {
@@ -266,6 +261,7 @@ async fn execute_and_wait_clamp_yield_grace_without_stopping_the_cell() {
 
 #[tokio::test(start_paused = true)]
 async fn wait_waits_for_nested_tool_during_yield_grace() {
+    let _time_guard = hold_virtual_time();
     let delegate = Arc::new(ReleasableToolDelegate::default());
     let service = InProcessCodeModeSession::new();
     let initial_response = service
@@ -300,7 +296,6 @@ async fn wait_waits_for_nested_tool_during_yield_grace() {
     tokio::task::yield_now().await;
     tokio::time::advance(Duration::from_millis(10_500)).await;
     delegate.release_tool();
-    wait_until_finished(&response).await;
     let response = response.await.unwrap();
 
     assert_eq!(
@@ -318,6 +313,7 @@ async fn wait_waits_for_nested_tool_during_yield_grace() {
 
 #[tokio::test(start_paused = true)]
 async fn zero_yield_limit_is_immediate_and_scoped_to_its_session() {
+    let _time_guard = hold_virtual_time();
     let zero_delegate = Arc::new(ReleasableToolDelegate::default());
     let zero_session = InProcessCodeModeSession::with_limits(CodeModeSessionCellExecutionLimits {
         max_yield_time_ms: Some(/*value*/ 0),
@@ -349,9 +345,8 @@ async fn zero_yield_limit_is_immediate_and_scoped_to_its_session() {
         .unwrap();
     let zero_response = tokio::spawn(zero_started.initial_response());
     let limited_response = tokio::spawn(limited_started.initial_response());
-    wait_until_tool_started(&zero_delegate).await;
-    wait_until_tool_started(&limited_delegate).await;
-    wait_until_finished(&zero_response).await;
+    wait_until_tool_started(&zero_delegate).await.unwrap();
+    wait_until_tool_started(&limited_delegate).await.unwrap();
     assert!(!limited_response.is_finished());
     assert_eq!(
         zero_response.await.unwrap().unwrap(),
@@ -372,7 +367,6 @@ async fn zero_yield_limit_is_immediate_and_scoped_to_its_session() {
         )
         .await;
     let zero_wait = tokio::spawn(zero_wait);
-    wait_until_finished(&zero_wait).await;
     assert_eq!(
         zero_wait.await.unwrap().unwrap(),
         WaitOutcome::LiveCell(RuntimeResponse::Yielded {
@@ -383,7 +377,6 @@ async fn zero_yield_limit_is_immediate_and_scoped_to_its_session() {
     );
 
     tokio::time::advance(Duration::from_millis(/*millis*/ 10)).await;
-    wait_until_finished(&limited_response).await;
     assert_eq!(
         limited_response.await.unwrap().unwrap(),
         RuntimeResponse::Yielded {
@@ -397,30 +390,88 @@ async fn zero_yield_limit_is_immediate_and_scoped_to_its_session() {
     limited_session.shutdown().await.unwrap();
 }
 
-async fn wait_until_finished<T>(task: &tokio::task::JoinHandle<T>) {
-    for _ in 0..10_000 {
-        if task.is_finished() {
-            return;
+#[tokio::test(start_paused = true)]
+async fn tool_start_acknowledgement_reports_error_before_invocation() {
+    let _time_guard = hold_virtual_time();
+    let delegate = Arc::new(ReleasableToolDelegate::default());
+    let service = InProcessCodeModeSession::new();
+    let started = service
+        .execute(
+            execute_request("throw 'startup failed';"),
+            delegate.clone(),
+            /*preempt*/ None,
+        )
+        .await
+        .unwrap();
+
+    let (startup, response) = tokio::join!(
+        wait_until_tool_started(&delegate),
+        started.initial_response(),
+    );
+    assert_eq!(startup, Err("cell closed before nested tool started"));
+    assert_eq!(
+        response.unwrap(),
+        RuntimeResponse::Result {
+            code_mode_host_duration: None,
+            cell_id: cell_id("1"),
+            content_items: Vec::new(),
+            error_text: Some("startup failed".to_string()),
         }
-        tokio::task::yield_now().await;
-    }
-    panic!("code-mode response did not finish while virtual time was held in the grace period");
+    );
 }
 
-async fn wait_until_tool_started(delegate: &ReleasableToolDelegate) {
-    for _ in 0..10_000 {
-        if delegate.tool_started.load(Ordering::Acquire) {
-            return;
-        }
-        tokio::task::yield_now().await;
+#[tokio::test(start_paused = true)]
+async fn tool_start_acknowledgement_reports_termination_before_invocation() {
+    let _time_guard = hold_virtual_time();
+    let delegate = Arc::new(ReleasableToolDelegate::default());
+    let service = InProcessCodeModeSession::new();
+    let started = service
+        .execute(
+            execute_request("await new Promise(() => {});"),
+            delegate.clone(),
+            /*preempt*/ None,
+        )
+        .await
+        .unwrap();
+
+    let (startup, terminated) = tokio::join!(
+        wait_until_tool_started(&delegate),
+        service.terminate(cell_id("1")),
+    );
+    assert_eq!(startup, Err("cell closed before nested tool started"));
+    let response = RuntimeResponse::Terminated {
+        code_mode_host_duration: None,
+        cell_id: cell_id("1"),
+        content_items: Vec::new(),
+    };
+    assert_eq!(terminated.unwrap(), WaitOutcome::LiveCell(response.clone()));
+    assert_eq!(started.initial_response().await.unwrap(), response);
+}
+
+fn hold_virtual_time() -> oneshot::Sender<()> {
+    // A live blocking task prevents Tokio's paused clock from auto-advancing while
+    // we await the V8 thread. Tests advance time explicitly; dropping the guard
+    // releases the task even if the test panics or is cancelled.
+    let (guard, released) = oneshot::channel();
+    tokio::task::spawn_blocking(move || {
+        let _ = released.blocking_recv();
+    });
+    guard
+}
+
+async fn wait_until_tool_started(delegate: &ReleasableToolDelegate) -> Result<(), &'static str> {
+    tokio::select! {
+        biased;
+        _ = delegate.tool_started.notified() => Ok(()),
+        _ = delegate.cell_closed.cancelled() => Err("cell closed before nested tool started"),
     }
-    panic!("nested code-mode tool did not start");
 }
 
 #[derive(Default)]
 struct ReleasableToolDelegate {
     tool_release: Notify,
-    tool_started: AtomicBool,
+    tool_started: Notify,
+    cell_closed: CancellationToken,
 }
 
 impl ReleasableToolDelegate {
@@ -435,7 +486,7 @@ impl CodeModeSessionDelegate for ReleasableToolDelegate {
         _invocation: CodeModeNestedToolCall,
         cancellation_token: CancellationToken,
     ) -> ToolInvocationFuture<'a> {
-        self.tool_started.store(true, Ordering::Release);
+        self.tool_started.notify_one();
         Box::pin(async move {
             tokio::select! {
                 _ = self.tool_release.notified() => Ok(JsonValue::Null),
@@ -454,7 +505,9 @@ impl CodeModeSessionDelegate for ReleasableToolDelegate {
         Box::pin(async { Ok(()) })
     }
 
-    fn cell_closed(&self, _cell_id: &CellId) {}
+    fn cell_closed(&self, _cell_id: &CellId) {
+        self.cell_closed.cancel();
+    }
 }
 
 fn execute_request(source: &str) -> ExecuteRequest {