Add grace period to code-mode yield timeouts (#33867)

## Why

Nested tool calls can finish just after a long code-mode yield deadline, causing
an `exec` or `wait` response to yield instead of returning the completed result.

## What changed

- Add a one-second grace period to `exec` and `wait` yield timeouts of at least
  ten seconds.
- Preserve the exact requested timeout for shorter yields.

## Testing

Add paused-time coverage for the timeout threshold and for nested tools that
complete during the grace period in both `exec` and `wait`.

GitOrigin-RevId: 403f587e07e2cde88c86424e46b0eff96ac8c409

codex-rs/code-mode/Cargo.toml

@@ -28,3 +28,4 @@ v8 = { workspace = true }
 
 [dev-dependencies]
 pretty_assertions = { workspace = true }
+tokio = { workspace = true, features = ["test-util"] }

codex-rs/code-mode/src/service.rs

@@ -29,6 +29,18 @@ use tokio_util::sync::CancellationToken;
 use crate::session_runtime as runtime;
 use crate::session_runtime::SessionRuntime;
 
+const YIELD_GRACE_PERIOD: Duration = Duration::from_secs(1);
+const MIN_YIELD_TIME_FOR_GRACE: Duration = Duration::from_secs(10);
+
+fn yield_timeout(yield_time_ms: u64) -> Duration {
+    let yield_time = Duration::from_millis(yield_time_ms);
+    if yield_time >= MIN_YIELD_TIME_FOR_GRACE {
+        yield_time.saturating_add(YIELD_GRACE_PERIOD)
+    } else {
+        yield_time
+    }
+}
+
 pub struct NoopCodeModeSessionDelegate;
 
 impl CodeModeSessionDelegate for NoopCodeModeSessionDelegate {
@@ -105,7 +117,7 @@ impl InProcessCodeModeSession {
             .runtime
             .execute(
                 runtime_request(request),
-                runtime::ObserveMode::YieldAfter(Duration::from_millis(yield_time_ms)),
+                runtime::ObserveMode::YieldAfter(yield_timeout(yield_time_ms)),
             )
             .await
             .map_err(|error| error.to_string())?;
@@ -160,7 +172,7 @@ impl InProcessCodeModeSession {
             .runtime
             .begin_observe(
                 &runtime_cell_id,
-                runtime::ObserveMode::YieldAfter(Duration::from_millis(yield_time_ms)),
+                runtime::ObserveMode::YieldAfter(yield_timeout(yield_time_ms)),
             )
             .await
         {

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

@@ -1,4 +1,6 @@
 use std::sync::Arc;
+use std::sync::atomic::AtomicBool;
+use std::sync::atomic::Ordering;
 use std::time::Duration;
 
 use super::CellId;
@@ -12,6 +14,7 @@ use super::WaitOutcome;
 use super::WaitRequest;
 use super::WaitToPendingOutcome;
 use super::WaitToPendingRequest;
+use super::yield_timeout;
 use crate::CodeModeToolKind;
 use crate::ExecuteRequest;
 use crate::ExecuteToPendingOutcome;
@@ -23,9 +26,117 @@ use serde_json::Value as JsonValue;
 use tokio::sync::Notify;
 use tokio_util::sync::CancellationToken;
 
+#[test]
+fn yield_timeout_adds_grace_only_at_ten_seconds() {
+    assert_eq!(
+        yield_timeout(/*yield_time_ms*/ 9_999),
+        Duration::from_millis(9_999)
+    );
+    assert_eq!(
+        yield_timeout(/*yield_time_ms*/ 10_000),
+        Duration::from_secs(11)
+    );
+}
+
+#[tokio::test(start_paused = true)]
+async fn execute_waits_for_nested_tool_during_yield_grace() {
+    let delegate = Arc::new(ReleasableToolDelegate::default());
+    let service = InProcessCodeModeSession::with_delegate(delegate.clone());
+    let request = ExecuteRequest {
+        enabled_tools: vec![echo_tool()],
+        source: r#"await tools.echo({}); text("done");"#.to_string(),
+        yield_time_ms: Some(10_000),
+        ..execute_request("")
+    };
+    let started = service.execute(request).await.unwrap();
+    let response = tokio::spawn(started.initial_response());
+    wait_until_tool_started(&delegate).await;
+    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!(
+        response,
+        RuntimeResponse::Result {
+            cell_id: cell_id("1"),
+            content_items: vec![FunctionCallOutputContentItem::InputText {
+                text: "done".to_string(),
+            }],
+            error_text: None,
+        }
+    );
+}
+
+#[tokio::test(start_paused = true)]
+async fn wait_waits_for_nested_tool_during_yield_grace() {
+    let delegate = Arc::new(ReleasableToolDelegate::default());
+    let service = InProcessCodeModeSession::with_delegate(delegate.clone());
+    let initial_response = service
+        .execute_to_pending(ExecuteRequest {
+            enabled_tools: vec![echo_tool()],
+            source: r#"await tools.echo({}); text("done");"#.to_string(),
+            ..execute_request("")
+        })
+        .await
+        .unwrap();
+    assert_eq!(
+        initial_response,
+        ExecuteToPendingOutcome::Pending {
+            cell_id: cell_id("1"),
+            content_items: Vec::new(),
+            pending_tool_call_ids: vec!["tool-1".to_string()],
+        }
+    );
+    let response = service
+        .begin_wait(WaitRequest {
+            cell_id: cell_id("1"),
+            yield_time_ms: 10_000,
+        })
+        .await;
+    let response = tokio::spawn(response);
+    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!(
+        response.unwrap(),
+        WaitOutcome::LiveCell(RuntimeResponse::Result {
+            cell_id: cell_id("1"),
+            content_items: vec![FunctionCallOutputContentItem::InputText {
+                text: "done".to_string(),
+            }],
+            error_text: None,
+        })
+    );
+}
+
+async fn wait_until_finished<T>(task: &tokio::task::JoinHandle<T>) {
+    for _ in 0..10_000 {
+        if task.is_finished() {
+            return;
+        }
+        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;
+    }
+    panic!("nested code-mode tool did not start");
+}
+
 #[derive(Default)]
 struct ReleasableToolDelegate {
     tool_release: Notify,
+    tool_started: AtomicBool,
 }
 
 impl ReleasableToolDelegate {
@@ -40,6 +151,7 @@ impl CodeModeSessionDelegate for ReleasableToolDelegate {
         _invocation: CodeModeNestedToolCall,
         cancellation_token: CancellationToken,
     ) -> ToolInvocationFuture<'a> {
+        self.tool_started.store(true, Ordering::Release);
         Box::pin(async move {
             tokio::select! {
                 _ = self.tool_release.notified() => Ok(JsonValue::Null),