diff --git a/codex-rs/code-mode/Cargo.toml b/codex-rs/code-mode/Cargo.toml
index 011fee6ebc..be14424181 100644
--- a/codex-rs/code-mode/Cargo.toml
+++ b/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"] }
diff --git a/codex-rs/code-mode/src/service.rs b/codex-rs/code-mode/src/service.rs
index 8c5b94bb75..c2651758fb 100644
--- a/codex-rs/code-mode/src/service.rs
+++ b/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
         {
diff --git a/codex-rs/code-mode/src/service_tests.rs b/codex-rs/code-mode/src/service_tests.rs
index 49cfc53d99..7e2313f79d 100644
--- a/codex-rs/code-mode/src/service_tests.rs
+++ b/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),