Add hidden grace to code mode yields (#2054)

## Why

Code-mode `exec` and `wait` can use the same yield duration as an
awaited nested tool call. Because the outer timer starts first, it can
expire just before the nested result arrives, forcing an otherwise
unnecessary continuation and extra `wait` call.

## What changed

- Add a one-second internal grace period to code-mode `exec` and `wait`
yield timers only when the configured yield is at least 10 seconds.
- Keep the model-visible yield values, defaults, and accepted parameters
unchanged.
- Preserve exact timing for explicit short yields, including the
immediate-yield `0` case.
- Cover the 10-second threshold and nested-tool behavior through both
`exec` and `wait` with virtual-time tests instead of slow wall-clock
integration tests.

## User impact

Nested calls using the 10-second default (or a longer configured yield)
get one extra second to return through the current code-mode call
instead of needlessly yielding the outer cell. Sub-10-second yields
retain their existing timing.

## Test plan

- `just test -p codex-code-mode`

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),