Add early yielding for code-mode observations (#48123)

## What changed

Add an optional `preempt` signal to `CodeModeSession::execute` and `CodeModeSession::wait` so callers can end an observation early while the cell continues running and remains available for later waits.

Support the signal in the in-process runtime and across both remote transports: `operation/yield` with the negotiated `yield-observation` capability for stdio, and `YieldObservation` for gRPC. Preserve signals received before gRPC observation registration, and retain normal timeout behavior for older hosts that lack support.

## Testing

Add runtime and transport tests covering early execute and wait yields followed by successful cell completion, yields before observation registration, and suppression of yield frames for older stdio hosts.

GitOrigin-RevId: 3f8049d8f0582be6b8f0832cc03f2d974a0b8f18

codex-rs/code-mode-host/src/grpc/mod.rs

@@ -139,7 +139,7 @@ impl GrpcCodeModeHost {
         let execution_id = request.execution_id.clone();
         let request = conversions::execute_request(request)?;
         let cell_permit = self.state.cell_permit()?;
-        session.reserve_execution(&execution_id)?;
+        let yield_signal = session.reserve_execution(&execution_id)?;
         let mut admission = ExecutionAdmission {
             session: Arc::clone(&session),
             execution_id: Some(execution_id.clone()),
@@ -148,7 +148,7 @@ impl GrpcCodeModeHost {
             _ = session.closed.cancelled() => {
                 return Err(Status::cancelled("code-mode session is closed"));
             }
-            result = session.runtime.execute(request, Arc::new(delegate::GrpcDelegate::new(Arc::downgrade(&session)))) => {
+            result = session.runtime.execute(request, Arc::new(delegate::GrpcDelegate::new(Arc::downgrade(&session))), Some(yield_signal)) => {
                 result.map_err(Status::failed_precondition)?
             }
         };
@@ -243,7 +243,7 @@ impl GrpcCodeModeHost {
             _ = session.closed.cancelled() => {
                 return Err(Status::cancelled("code-mode session is closed"));
             }
-            outcome = session.runtime.wait(request) => {
+            outcome = session.runtime.wait(request, Some(registration.yield_signal())) => {
                 outcome.map_err(Status::failed_precondition)?
             }
         };
@@ -264,6 +264,24 @@ impl GrpcCodeModeHost {
         Ok(Response::new(proto::CancelWaitResponse {}))
     }
 
+    async fn yield_observation_request(
+        &self,
+        request: proto::YieldObservationRequest,
+    ) -> Result<Response<proto::YieldObservationResponse>, Status> {
+        let _permit = self.state.control_permit()?;
+        let session = self.state.session(&request.session_id)?;
+        match request.observation {
+            Some(proto::yield_observation_request::Observation::ExecutionId(id)) => {
+                session.yield_execution(&id)?;
+            }
+            Some(proto::yield_observation_request::Observation::WaitId(id)) => {
+                session.yield_wait(&id)?;
+            }
+            None => return Err(Status::invalid_argument("missing code-mode observation ID")),
+        }
+        Ok(Response::new(proto::YieldObservationResponse {}))
+    }
+
     async fn terminate_request(
         &self,
         request: proto::TerminateRequest,
@@ -405,6 +423,17 @@ impl CodeModeHost for GrpcCodeModeHost {
         Box::pin(self.cancel_wait_request(request.into_inner()))
     }
 
+    fn yield_observation<'a, 'async_trait>(
+        &'a self,
+        request: Request<proto::YieldObservationRequest>,
+    ) -> GrpcFuture<'async_trait, proto::YieldObservationResponse>
+    where
+        'a: 'async_trait,
+        Self: 'async_trait,
+    {
+        Box::pin(self.yield_observation_request(request.into_inner()))
+    }
+
     fn terminate<'a, 'async_trait>(
         &'a self,
         request: Request<proto::TerminateRequest>,
@@ -424,7 +453,9 @@ struct ExecutionAdmission {
 
 impl ExecutionAdmission {
     fn disarm(&mut self) {
-        self.execution_id = None;
+        if let Some(execution_id) = self.execution_id.take() {
+            self.session.retire_execution_observation(&execution_id);
+        }
     }
 }
 

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

@@ -368,6 +368,28 @@ async fn cancellation_before_wait_admission_is_preserved() {
     .unwrap();
 }
 
+#[tokio::test]
+async fn yield_before_observation_registration_is_remembered() {
+    let host = GrpcCodeModeHost::new();
+    let (session_id, _events) = open_session(&host).await;
+    let session = host.state.session(&session_id).unwrap();
+    for observation in [
+        proto::yield_observation_request::Observation::ExecutionId("exec".to_string()),
+        proto::yield_observation_request::Observation::WaitId("wait".to_string()),
+    ] {
+        host.yield_observation(Request::new(proto::YieldObservationRequest {
+            session_id: session_id.clone(),
+            observation: Some(observation),
+        }))
+        .await
+        .unwrap();
+    }
+
+    assert!(session.reserve_execution("exec").unwrap().is_cancelled());
+    let wait = super::waits::WaitRegistration::new(session, "wait".to_string()).unwrap();
+    assert!(wait.yield_signal().is_cancelled());
+}
+
 #[tokio::test]
 async fn notifications_do_not_delay_cell_completion() {
     let host = GrpcCodeModeHost::new();

codex-rs/code-mode-host/src/grpc/session.rs

@@ -58,7 +58,9 @@ pub(super) struct SessionState {
     pub(super) cells: HashMap<String, ExecutionState>,
     pending_executions: HashSet<String>,
     pending_closures: HashSet<String>,
-    seen_executions: BoundedIds,
+    pub(super) seen_executions: BoundedIds,
+    pub(super) execution_yields: HashMap<String, CancellationToken>,
+    pub(super) yielded_executions: BoundedIds,
     pub(super) subscriptions: Vec<ToolSubscription>,
     pub(super) next_subscription: usize,
     pub(super) pending_invocations: HashMap<Uuid, PendingInvocation>,
@@ -66,6 +68,7 @@ pub(super) struct SessionState {
     pub(super) waits: HashMap<String, ActiveWait>,
     pub(super) seen_waits: BoundedIds,
     pub(super) cancelled_waits: BoundedIds,
+    pub(super) yielded_waits: BoundedIds,
 }
 
 pub(super) struct ExecutionState {
@@ -266,7 +269,10 @@ impl GrpcSession {
         }
     }
 
-    pub(super) fn reserve_execution(&self, execution_id: &str) -> Result<(), Status> {
+    pub(super) fn reserve_execution(
+        &self,
+        execution_id: &str,
+    ) -> Result<CancellationToken, Status> {
         validation::identifier(execution_id, "execution ID")?;
         let mut state = self.state.lock().unwrap_or_else(PoisonError::into_inner);
         if self.closed.is_cancelled() {
@@ -284,7 +290,14 @@ impl GrpcSession {
             )));
         }
         state.pending_executions.insert(execution_id.to_string());
-        Ok(())
+        let yield_signal = CancellationToken::new();
+        if state.yielded_executions.remove(execution_id) {
+            yield_signal.cancel();
+        }
+        state
+            .execution_yields
+            .insert(execution_id.to_string(), yield_signal.clone());
+        Ok(yield_signal)
     }
 
     pub(super) fn admit_execution(
@@ -323,6 +336,7 @@ impl GrpcSession {
         let cell_id = {
             let mut state = self.state.lock().unwrap_or_else(PoisonError::into_inner);
             state.pending_executions.remove(execution_id);
+            state.execution_yields.remove(execution_id);
             state
                 .cells
                 .iter()

codex-rs/code-mode-host/src/grpc/waits.rs

@@ -9,6 +9,7 @@ use super::validation;
 
 pub(super) struct ActiveWait {
     pub(super) cancellation: CancellationToken,
+    pub(super) yield_signal: CancellationToken,
     retired: CancellationToken,
 }
 
@@ -16,6 +17,7 @@ pub(super) struct WaitRegistration {
     session: Arc<GrpcSession>,
     id: String,
     cancellation: CancellationToken,
+    yield_signal: CancellationToken,
     retired: CancellationToken,
 }
 
@@ -23,6 +25,7 @@ impl WaitRegistration {
     pub(super) fn new(session: Arc<GrpcSession>, id: String) -> Result<Self, Status> {
         validation::identifier(&id, "wait ID")?;
         let cancellation = CancellationToken::new();
+        let yield_signal = CancellationToken::new();
         let retired = CancellationToken::new();
         let mut state = session.state.lock().unwrap_or_else(PoisonError::into_inner);
         if session.closed.is_cancelled() {
@@ -34,12 +37,17 @@ impl WaitRegistration {
             )));
         }
         if state.cancelled_waits.remove(&id) {
+            state.yielded_waits.remove(&id);
             return Err(Status::cancelled("code-mode wait was cancelled"));
         }
+        if state.yielded_waits.remove(&id) {
+            yield_signal.cancel();
+        }
         state.waits.insert(
             id.clone(),
             ActiveWait {
                 cancellation: cancellation.clone(),
+                yield_signal: yield_signal.clone(),
                 retired: retired.clone(),
             },
         );
@@ -48,6 +56,7 @@ impl WaitRegistration {
             session,
             id,
             cancellation,
+            yield_signal,
             retired,
         })
     }
@@ -55,6 +64,10 @@ impl WaitRegistration {
     pub(super) fn cancellation(&self) -> &CancellationToken {
         &self.cancellation
     }
+
+    pub(super) fn yield_signal(&self) -> CancellationToken {
+        self.yield_signal.clone()
+    }
 }
 
 impl Drop for WaitRegistration {
@@ -70,6 +83,36 @@ impl Drop for WaitRegistration {
 }
 
 impl GrpcSession {
+    pub(super) fn yield_execution(&self, id: &str) -> Result<(), Status> {
+        validation::identifier(id, "execution ID")?;
+        let mut state = self.state.lock().unwrap_or_else(PoisonError::into_inner);
+        if let Some(signal) = state.execution_yields.get(id) {
+            signal.cancel();
+        } else if !state.seen_executions.contains(id) {
+            state.yielded_executions.remember(id.to_string());
+        }
+        Ok(())
+    }
+
+    pub(super) fn yield_wait(&self, id: &str) -> Result<(), Status> {
+        validation::identifier(id, "wait ID")?;
+        let mut state = self.state.lock().unwrap_or_else(PoisonError::into_inner);
+        if let Some(wait) = state.waits.get(id) {
+            wait.yield_signal.cancel();
+        } else if !state.seen_waits.contains(id) {
+            state.yielded_waits.remember(id.to_string());
+        }
+        Ok(())
+    }
+
+    pub(super) fn retire_execution_observation(&self, id: &str) {
+        self.state
+            .lock()
+            .unwrap_or_else(PoisonError::into_inner)
+            .execution_yields
+            .remove(id);
+    }
+
     pub(super) async fn cancel_wait(&self, id: &str) -> Result<(), Status> {
         validation::identifier(id, "wait ID")?;
         let active = {

codex-rs/code-mode-host/src/host_tests.rs

@@ -595,6 +595,7 @@ async fn active_cell_limit_rejects_execute_without_disconnecting() {
                 request: execute_request("text(\"hello\");"),
             },
             CancellationToken::new(),
+            CancellationToken::new(),
             Instant::now(),
         )
         .await;

codex-rs/code-mode-host/src/lib.rs

@@ -29,6 +29,7 @@ use codex_code_mode_protocol::host::SESSION_RESOURCE_LIMITS_CAPABILITY;
 use codex_code_mode_protocol::host::SessionId;
 use codex_code_mode_protocol::host::SupportedProtocolVersions;
 use codex_code_mode_protocol::host::WireWaitOutcome;
+use codex_code_mode_protocol::host::YIELD_OBSERVATION_CAPABILITY;
 use codex_code_mode_runtime::InProcessCodeModeSession;
 use tokio::io::AsyncRead;
 use tokio::io::AsyncWrite;
@@ -163,6 +164,9 @@ where
                 ClientToHost::CancelRequest { id } => {
                     state.cancel_request(id);
                 }
+                ClientToHost::YieldRequest { id } => {
+                    state.yield_request(id);
+                }
                 ClientToHost::DelegateResponse { id, result } => {
                     peer.complete(id, result.into_result()).await;
                 }
@@ -253,15 +257,20 @@ async fn negotiate<R: AsyncRead + Unpin, W: AsyncWrite + Unpin>(
         return Ok(NegotiatedConnection::Rejected);
     }
 
-    let resource_limits_capability = Capability::new(SESSION_RESOURCE_LIMITS_CAPABILITY)?;
-    let resource_limits_requested = client_hello
-        .required_capabilities()
-        .contains(&resource_limits_capability)
-        || client_hello
-            .optional_capabilities()
-            .contains(&resource_limits_capability);
-    let host_capabilities =
-        CapabilitySet::try_new(resource_limits_requested.then_some(resource_limits_capability))?;
+    let host_capabilities = CapabilitySet::try_new(
+        [
+            SESSION_RESOURCE_LIMITS_CAPABILITY,
+            YIELD_OBSERVATION_CAPABILITY,
+        ]
+        .into_iter()
+        .map(Capability::new)
+        .collect::<Result<Vec<_>, _>>()?
+        .into_iter()
+        .filter(|capability| {
+            client_hello.required_capabilities().contains(capability)
+                || client_hello.optional_capabilities().contains(capability)
+        }),
+    )?;
     if let Some(capability) = client_hello
         .required_capabilities()
         .iter()
@@ -305,7 +314,7 @@ impl HostState {
         // Include host admission and scheduling before the runtime observes the request.
         let received_at = Instant::now();
         let request_kind = RequestKind::from(&request);
-        let cancellation = self
+        let (cancellation, yield_signal) = self
             .requests
             .lock()
             .unwrap_or_else(PoisonError::into_inner)
@@ -322,7 +331,7 @@ impl HostState {
         let request_task = self.request_tasks.spawn(async move {
             let _permit = permit;
             state
-                .handle_request(request_id, request, cancellation, received_at)
+                .handle_request(request_id, request, cancellation, yield_signal, received_at)
                 .await;
             state.finish_request(request_id);
         });
@@ -355,6 +364,7 @@ impl HostState {
         request_id: RequestId,
         request: HostRequest,
         cancellation: CancellationToken,
+        yield_signal: CancellationToken,
         received_at: Instant,
     ) {
         if let HostRequest::Execute { request, .. } = &request {
@@ -416,7 +426,7 @@ impl HostState {
                     session_id.clone(),
                     Arc::clone(&self.peer),
                 ));
-                let result = session.execute(request, delegate).await;
+                let result = session.execute(request, delegate, Some(yield_signal)).await;
                 match result {
                     Ok(started) => {
                         let cell_id = started.cell_id.clone();
@@ -449,7 +459,7 @@ impl HostState {
                             _ = cancellation.cancelled() => {
                                 Err("code-mode request cancelled".to_string())
                             }
-                            result = session.wait(request.into()) => result.and_then(|outcome| {
+                            result = session.wait(request.into(), Some(yield_signal)) => result.and_then(|outcome| {
                                 let outcome = outcome.with_code_mode_host_duration(received_at.elapsed());
                                 Ok(HostResponse::WaitCompleted {
                                     outcome: WireWaitOutcome::try_from(outcome)
@@ -557,6 +567,13 @@ impl HostState {
             .cancel(request_id);
     }
 
+    fn yield_request(&self, request_id: RequestId) {
+        self.requests
+            .lock()
+            .unwrap_or_else(PoisonError::into_inner)
+            .yield_observation(request_id);
+    }
+
     fn finish_request(&self, request_id: RequestId) {
         self.requests
             .lock()
@@ -623,6 +640,7 @@ impl RequestKind {
 struct ActiveRequest {
     kind: RequestKind,
     cancellation: CancellationToken,
+    yield_signal: CancellationToken,
 }
 
 #[derive(Default)]
@@ -637,19 +655,21 @@ impl RequestRegistry {
         &mut self,
         request_id: RequestId,
         kind: RequestKind,
-    ) -> Result<CancellationToken, anyhow::Error> {
+    ) -> Result<(CancellationToken, CancellationToken), anyhow::Error> {
         if self.active.contains_key(&request_id) || self.recent.contains(&request_id) {
             anyhow::bail!("duplicate code-mode request ID {request_id:?}");
         }
         let cancellation = CancellationToken::new();
+        let yield_signal = CancellationToken::new();
         self.active.insert(
             request_id,
             ActiveRequest {
                 kind,
                 cancellation: cancellation.clone(),
+                yield_signal: yield_signal.clone(),
             },
         );
-        Ok(cancellation)
+        Ok((cancellation, yield_signal))
     }
 
     fn cancel(&mut self, request_id: RequestId) {
@@ -660,6 +680,14 @@ impl RequestRegistry {
         }
     }
 
+    fn yield_observation(&self, request_id: RequestId) {
+        if let Some(request) = self.active.get(&request_id)
+            && matches!(request.kind, RequestKind::Execute | RequestKind::Wait)
+        {
+            request.yield_signal.cancel();
+        }
+    }
+
     fn finish(&mut self, request_id: RequestId) {
         if self.active.remove(&request_id).is_none() {
             return;

codex-rs/code-mode-host/tests/grpc.rs

@@ -151,7 +151,7 @@ async fn execute(
 ) -> Result<RuntimeResponse> {
     timeout(TEST_TIMEOUT, async {
         session
-            .execute(request, delegate.clone())
+            .execute(request, delegate.clone(), /*preempt*/ None)
             .await
             .map_err(anyhow::Error::msg)?
             .initial_response()
@@ -168,7 +168,7 @@ async fn start_active_wait(
 ) -> Result<tokio::task::JoinHandle<std::result::Result<WaitOutcome, String>>> {
     let (admitted_tx, admitted_rx) = tokio::sync::oneshot::channel();
     let wait = tokio::spawn(async move {
-        let mut wait = session.wait(request);
+        let mut wait = session.wait(request, /*preempt*/ None);
         match wait.as_mut().now_or_never() {
             Some(result) => result,
             None => {
@@ -285,6 +285,81 @@ async fn tcp_session_persists_values_and_forwards_tools_notifications_and_closur
     Ok(())
 }
 
+#[tokio::test]
+async fn grpc_early_yield_preserves_cell_and_later_waits() -> Result<()> {
+    let host = HostHarness::start("grpc://127.0.0.1:0").await?;
+    let session = GrpcCodeModeSessionProvider::new(host.endpoint)
+        .create_session()
+        .await
+        .map_err(anyhow::Error::msg)?;
+    let delegate = Arc::new(LargeToolResultDelegate {
+        started: Semaphore::new(0),
+        release: Semaphore::new(0),
+    });
+    let mut execute = request(r#"await tools.large({}); text("done");"#);
+    execute.enabled_tools = vec![tool("large")];
+    execute.yield_time_ms = Some(60_000);
+    let signal = CancellationToken::new();
+    signal.cancel();
+    let cell = session
+        .execute(execute, delegate.clone(), Some(signal))
+        .await
+        .map_err(anyhow::Error::msg)?;
+    let cell_id = cell.cell_id.clone();
+    let response = timeout(TEST_TIMEOUT, cell.initial_response())
+        .await?
+        .map_err(anyhow::Error::msg)?;
+    assert!(matches!(response, RuntimeResponse::Yielded { .. }));
+    timeout(TEST_TIMEOUT, delegate.started.acquire())
+        .await??
+        .forget();
+
+    let wait_signal = CancellationToken::new();
+    let mut wait = session.wait(
+        WaitRequest {
+            cell_id: cell_id.clone(),
+            yield_time_ms: 60_000,
+        },
+        Some(wait_signal.clone()),
+    );
+    assert!(wait.as_mut().now_or_never().is_none());
+    wait_signal.cancel();
+    let response = timeout(TEST_TIMEOUT, wait)
+        .await?
+        .map_err(anyhow::Error::msg)?;
+    assert!(matches!(
+        response,
+        WaitOutcome::LiveCell(RuntimeResponse::Yielded { .. })
+    ));
+
+    delegate.release.add_permits(1);
+    let response = timeout(
+        TEST_TIMEOUT,
+        session.wait(
+            WaitRequest {
+                cell_id: cell_id.clone(),
+                yield_time_ms: 60_000,
+            },
+            /*preempt*/ None,
+        ),
+    )
+    .await?
+    .map_err(anyhow::Error::msg)?;
+    assert_eq!(
+        response,
+        WaitOutcome::LiveCell(RuntimeResponse::Result {
+            code_mode_host_duration: response.code_mode_host_duration(),
+            cell_id,
+            content_items: vec![FunctionCallOutputContentItem::InputText {
+                text: "done".to_string(),
+            }],
+            error_text: None,
+        })
+    );
+    session.shutdown().await.map_err(anyhow::Error::msg)?;
+    Ok(())
+}
+
 #[tokio::test]
 async fn shutdown_immediately_rejects_new_operations() -> Result<()> {
     let host = HostHarness::start("grpc://127.0.0.1:0").await?;
@@ -300,7 +375,8 @@ async fn shutdown_immediately_rejects_new_operations() -> Result<()> {
         session
             .execute(
                 request("text('too late');"),
-                Arc::new(NoopCodeModeSessionDelegate)
+                Arc::new(NoopCodeModeSessionDelegate),
+                /*preempt*/ None,
             )
             .await
             .err(),
@@ -308,10 +384,13 @@ async fn shutdown_immediately_rejects_new_operations() -> Result<()> {
     );
     assert_eq!(
         session
-            .wait(WaitRequest {
-                cell_id: cell_id("missing"),
-                yield_time_ms: 1,
-            })
+            .wait(
+                WaitRequest {
+                    cell_id: cell_id("missing"),
+                    yield_time_ms: 1,
+                },
+                /*preempt*/ None,
+            )
             .await,
         Err(expected.clone())
     );
@@ -335,7 +414,7 @@ async fn cancelling_execution_before_admission_keeps_the_session_usable() -> Res
 
     assert!(
         session
-            .execute(pending, delegate.clone())
+            .execute(pending, delegate.clone(), /*preempt*/ None)
             .now_or_never()
             .is_none()
     );
@@ -360,10 +439,13 @@ async fn cancelling_execution_before_admission_keeps_the_session_usable() -> Res
     timeout(TEST_TIMEOUT, async {
         loop {
             match session
-                .wait(WaitRequest {
-                    cell_id: abandoned_cell.clone(),
-                    yield_time_ms: 1,
-                })
+                .wait(
+                    WaitRequest {
+                        cell_id: abandoned_cell.clone(),
+                        yield_time_ms: 1,
+                    },
+                    /*preempt*/ None,
+                )
                 .await
             {
                 Ok(WaitOutcome::MissingCell(_))
@@ -409,7 +491,7 @@ async fn dropping_a_started_cell_off_runtime_terminates_its_buffered_remote_exec
     let mut pending = request("await new Promise(() => {});");
     pending.yield_time_ms = Some(/*value*/ 1);
     let started = session
-        .execute(pending, delegate.clone())
+        .execute(pending, delegate.clone(), /*preempt*/ None)
         .await
         .map_err(anyhow::Error::msg)?;
     let abandoned_cell = started.cell_id.clone();
@@ -417,10 +499,13 @@ async fn dropping_a_started_cell_off_runtime_terminates_its_buffered_remote_exec
     timeout(TEST_TIMEOUT, async {
         loop {
             match session
-                .wait(WaitRequest {
-                    cell_id: abandoned_cell.clone(),
-                    yield_time_ms: 1,
-                })
+                .wait(
+                    WaitRequest {
+                        cell_id: abandoned_cell.clone(),
+                        yield_time_ms: 1,
+                    },
+                    /*preempt*/ None,
+                )
                 .await
             {
                 Ok(WaitOutcome::LiveCell(RuntimeResponse::Yielded { .. })) => break Ok(()),
@@ -442,10 +527,13 @@ async fn dropping_a_started_cell_off_runtime_terminates_its_buffered_remote_exec
     timeout(TEST_TIMEOUT, async {
         loop {
             match session
-                .wait(WaitRequest {
-                    cell_id: abandoned_cell.clone(),
-                    yield_time_ms: 1,
-                })
+                .wait(
+                    WaitRequest {
+                        cell_id: abandoned_cell.clone(),
+                        yield_time_ms: 1,
+                    },
+                    /*preempt*/ None,
+                )
                 .await
             {
                 Ok(WaitOutcome::MissingCell(_))
@@ -498,7 +586,7 @@ async fn dropping_an_initial_response_terminates_its_pending_remote_execution()
     let mut pending = request("await new Promise(() => {});");
     pending.yield_time_ms = Some(/*value*/ 60_000);
     let started = session
-        .execute(pending, delegate.clone())
+        .execute(pending, delegate.clone(), /*preempt*/ None)
         .await
         .map_err(anyhow::Error::msg)?;
     let abandoned_cell = started.cell_id.clone();
@@ -628,7 +716,11 @@ async fn concurrent_wait_rejects_without_displacing_the_active_observer() -> Res
     let mut pending = request("await new Promise(() => {});");
     pending.yield_time_ms = Some(/*value*/ 1);
     let started = session
-        .execute(pending, Arc::new(NoopCodeModeSessionDelegate))
+        .execute(
+            pending,
+            Arc::new(NoopCodeModeSessionDelegate),
+            /*preempt*/ None,
+        )
         .await
         .map_err(anyhow::Error::msg)?;
     let running_cell = started.cell_id.clone();
@@ -657,10 +749,13 @@ async fn concurrent_wait_rejects_without_displacing_the_active_observer() -> Res
     assert_eq!(
         timeout(
             Duration::from_secs(/*secs*/ 1),
-            session.wait(WaitRequest {
-                cell_id: running_cell.clone(),
-                yield_time_ms: 60_000,
-            }),
+            session.wait(
+                WaitRequest {
+                    cell_id: running_cell.clone(),
+                    yield_time_ms: 60_000,
+                },
+                /*preempt*/ None,
+            ),
         )
         .await
         .context("concurrent wait did not reject immediately")?
@@ -708,7 +803,11 @@ async fn dropping_a_wait_retires_its_observer_before_the_next_wait() -> Result<(
     let mut pending = request("await new Promise(() => {});");
     pending.yield_time_ms = Some(/*value*/ 1);
     let started = session
-        .execute(pending, Arc::new(NoopCodeModeSessionDelegate))
+        .execute(
+            pending,
+            Arc::new(NoopCodeModeSessionDelegate),
+            /*preempt*/ None,
+        )
         .await
         .map_err(anyhow::Error::msg)?;
     let running_cell = started.cell_id.clone();
@@ -738,10 +837,13 @@ async fn dropping_a_wait_retires_its_observer_before_the_next_wait() -> Result<(
 
     let actual = timeout(
         TEST_TIMEOUT,
-        session.wait(WaitRequest {
-            cell_id: running_cell.clone(),
-            yield_time_ms: 1,
-        }),
+        session.wait(
+            WaitRequest {
+                cell_id: running_cell.clone(),
+                yield_time_ms: 1,
+            },
+            /*preempt*/ None,
+        ),
     )
     .await
     .context("replacement wait did not observe cancellation retirement")?
@@ -836,7 +938,7 @@ async fn large_unary_tool_completion_does_not_block_an_independent_session() ->
     slow_request.enabled_tools = vec![tool("large")];
     slow_request.yield_time_ms = Some(/*value*/ 20_000);
     let slow_cell = slow_session
-        .execute(slow_request, delegate.clone())
+        .execute(slow_request, delegate.clone(), /*preempt*/ None)
         .await
         .map_err(anyhow::Error::msg)?;
     timeout(TEST_TIMEOUT, delegate.started.acquire())
@@ -906,7 +1008,7 @@ async fn single_subscription_processes_slow_and_fast_tools_concurrently() -> Res
     slow.enabled_tools = vec![tool("large")];
     slow.yield_time_ms = Some(/*value*/ 20_000);
     let slow_cell = session
-        .execute(slow, delegate.clone())
+        .execute(slow, delegate.clone(), /*preempt*/ None)
         .await
         .map_err(anyhow::Error::msg)?;
     timeout(TEST_TIMEOUT, delegate.started.acquire())
@@ -974,10 +1076,13 @@ async fn sessions_enforce_independent_yield_limits() -> Result<()> {
     );
     let actual = timeout(
         TEST_TIMEOUT,
-        limited.wait(WaitRequest {
-            cell_id: cell_id("1"),
-            yield_time_ms: 60_000,
-        }),
+        limited.wait(
+            WaitRequest {
+                cell_id: cell_id("1"),
+                yield_time_ms: 60_000,
+            },
+            /*preempt*/ None,
+        ),
     )
     .await
     .context("session yield limit did not bound an explicit wait")?
@@ -1094,7 +1199,7 @@ async fn cached_session_recovers_after_a_remote_host_restarts() -> Result<()> {
     pending.enabled_tools = vec![tool("echo")];
     pending.yield_time_ms = Some(/*value*/ 1);
     let started = session
-        .execute(pending, delegate.clone())
+        .execute(pending, delegate.clone(), /*preempt*/ None)
         .await
         .map_err(anyhow::Error::msg)?;
     let old_cell_id = started.cell_id.clone();
@@ -1227,7 +1332,7 @@ async fn cached_session_recovers_after_a_remote_host_restarts() -> Result<()> {
     let mut pending = request("await new Promise(() => {});");
     pending.yield_time_ms = Some(/*value*/ 1);
     let started = session
-        .execute(pending, delegate.clone())
+        .execute(pending, delegate.clone(), /*preempt*/ None)
         .await
         .map_err(anyhow::Error::msg)?;
     let replacement_cell_id = started.cell_id.clone();
@@ -1245,10 +1350,13 @@ async fn cached_session_recovers_after_a_remote_host_restarts() -> Result<()> {
         }
     );
     let actual = session
-        .wait(WaitRequest {
-            cell_id: replacement_cell_id.clone(),
-            yield_time_ms: 1,
-        })
+        .wait(
+            WaitRequest {
+                cell_id: replacement_cell_id.clone(),
+                yield_time_ms: 1,
+            },
+            /*preempt*/ None,
+        )
         .await
         .map_err(anyhow::Error::msg)?;
     assert_eq!(
@@ -1273,10 +1381,13 @@ async fn cached_session_recovers_after_a_remote_host_restarts() -> Result<()> {
     );
 
     let stale_wait = session
-        .wait(WaitRequest {
-            cell_id: old_cell_id.clone(),
-            yield_time_ms: 1,
-        })
+        .wait(
+            WaitRequest {
+                cell_id: old_cell_id.clone(),
+                yield_time_ms: 1,
+            },
+            /*preempt*/ None,
+        )
         .await
         .unwrap_err();
     assert!(stale_wait.contains("stale code-mode host generation"));

codex-rs/code-mode-host/tests/grpc/network_policy_tests.rs

@@ -60,7 +60,11 @@ async fn account_change_revokes_active_rpc_and_renews_the_cached_client() -> Res
     let mut pending = request("await new Promise(() => {});");
     pending.yield_time_ms = Some(60_000);
     let started = session
-        .execute(pending, Arc::new(NoopCodeModeSessionDelegate))
+        .execute(
+            pending,
+            Arc::new(NoopCodeModeSessionDelegate),
+            /*preempt*/ None,
+        )
         .await
         .map_err(anyhow::Error::msg)?;
 

codex-rs/code-mode-host/tests/grpc_notifications.rs

@@ -124,7 +124,7 @@ async fn execute(
 ) -> Result<RuntimeResponse> {
     timeout(TEST_TIMEOUT, async {
         session
-            .execute(request, delegate.clone())
+            .execute(request, delegate.clone(), /*preempt*/ None)
             .await
             .map_err(anyhow::Error::msg)?
             .initial_response()
@@ -204,7 +204,7 @@ async fn completed_waits_drain_pending_notifications_before_returning() -> Resul
         .map_err(anyhow::Error::msg)?;
     let pending = request(r#"yield_control(); notify("notice"); text("done");"#);
     let cell = session
-        .execute(pending, delegate.clone())
+        .execute(pending, delegate.clone(), /*preempt*/ None)
         .await
         .map_err(anyhow::Error::msg)?;
     let actual = cell.initial_response().await.map_err(anyhow::Error::msg)?;
@@ -220,10 +220,13 @@ async fn completed_waits_drain_pending_notifications_before_returning() -> Resul
     let waiting = Arc::clone(&session);
     let completion = tokio::spawn(async move {
         waiting
-            .wait(WaitRequest {
-                cell_id: cell_id("1"),
-                yield_time_ms: 5_000,
-            })
+            .wait(
+                WaitRequest {
+                    cell_id: cell_id("1"),
+                    yield_time_ms: 5_000,
+                },
+                /*preempt*/ None,
+            )
             .await
             .map_err(anyhow::Error::msg)
     });
@@ -269,7 +272,7 @@ async fn termination_cancels_pending_notifications() -> Result<()> {
     let mut pending = request(r#"notify("notice"); await new Promise(() => {});"#);
     pending.yield_time_ms = Some(/*value*/ 1);
     let cell = session
-        .execute(pending, delegate.clone())
+        .execute(pending, delegate.clone(), /*preempt*/ None)
         .await
         .map_err(anyhow::Error::msg)?;
 

codex-rs/code-mode-host/tests/grpc_timing.rs

@@ -140,6 +140,7 @@ async fn stdio_execution_timing_includes_javascript_but_excludes_delayed_reads()
                     max_output_tokens: Some(/*value*/ 1_000),
                 },
                 delegate.clone(),
+                /*preempt*/ None,
             )
             .await
             .map_err(anyhow::Error::msg)?;
@@ -200,6 +201,7 @@ async fn observation_timing_excludes_previous_requests_and_background_time() ->
                     max_output_tokens: Some(/*value*/ 1_000),
                 },
                 Arc::new(NoopCodeModeSessionDelegate),
+                /*preempt*/ None,
             )
             .await
             .map_err(anyhow::Error::msg)?;
@@ -223,10 +225,13 @@ async fn observation_timing_excludes_previous_requests_and_background_time() ->
             sleep(Duration::from_millis(/*millis*/ 300)).await;
             let observed = Instant::now();
             let outcome = session
-                .wait(WaitRequest {
-                    cell_id: cell_id.clone(),
-                    yield_time_ms: 50,
-                })
+                .wait(
+                    WaitRequest {
+                        cell_id: cell_id.clone(),
+                        yield_time_ms: 50,
+                    },
+                    /*preempt*/ None,
+                )
                 .await
                 .map_err(anyhow::Error::msg)?;
             let duration = outcome
@@ -267,10 +272,13 @@ async fn observation_timing_excludes_previous_requests_and_background_time() ->
 
         let observed = Instant::now();
         let missing = session
-            .wait(WaitRequest {
-                cell_id: cell_id.clone(),
-                yield_time_ms: 50,
-            })
+            .wait(
+                WaitRequest {
+                    cell_id: cell_id.clone(),
+                    yield_time_ms: 50,
+                },
+                /*preempt*/ None,
+            )
             .await
             .map_err(anyhow::Error::msg)?;
         let duration = missing

codex-rs/code-mode-host/tests/stdio.rs

@@ -180,7 +180,7 @@ async fn execute(
     delegate: Arc<dyn CodeModeSessionDelegate>,
 ) -> RuntimeResponse {
     session
-        .execute(request, delegate.clone())
+        .execute(request, delegate.clone(), /*preempt*/ None)
         .await
         .expect("start execution")
         .initial_response()
@@ -194,7 +194,7 @@ async fn execute_to_terminal(
     delegate: Arc<dyn CodeModeSessionDelegate>,
 ) -> RuntimeResponse {
     let started = session
-        .execute(request, delegate.clone())
+        .execute(request, delegate.clone(), /*preempt*/ None)
         .await
         .expect("start execution");
     let mut response = started.initial_response().await.expect("initial response");
@@ -202,10 +202,13 @@ async fn execute_to_terminal(
         match response {
             RuntimeResponse::Yielded { cell_id, .. } => {
                 response = match session
-                    .wait(WaitRequest {
-                        cell_id,
-                        yield_time_ms: 60_000,
-                    })
+                    .wait(
+                        WaitRequest {
+                            cell_id,
+                            yield_time_ms: 60_000,
+                        },
+                        /*preempt*/ None,
+                    )
                     .await
                     .expect("wait for terminal response")
                 {
@@ -228,6 +231,79 @@ async fn next_callback_event(
         .expect("callback event stream closed")
 }
 
+#[tokio::test]
+async fn interrupt_yields_observations_without_stopping_the_cell() {
+    let provider = ProcessOwnedCodeModeSessionProvider::with_host_program(
+        codex_utils_cargo_bin::cargo_bin("codex-code-mode-host").expect("host binary"),
+    );
+    let session = provider.create_session().await.expect("create session");
+    let (delegate, mut events) = CancellationDelegate::new();
+    let mut request = execute_request("await tools.block({}); text('done');");
+    request.yield_time_ms = Some(60_000);
+    request.enabled_tools = vec![ToolDefinition {
+        name: "block".to_string(),
+        tool_name: ToolName::plain("block"),
+        description: String::new(),
+        kind: CodeModeToolKind::Function,
+        input_schema: None,
+        output_schema: None,
+    }];
+    let signal = CancellationToken::new();
+    let started = session
+        .execute(request, delegate.clone(), Some(signal.clone()))
+        .await
+        .expect("start execution");
+    let cell_id = started.cell_id.clone();
+    assert_eq!(
+        next_callback_event(&mut events).await,
+        CallbackEvent::Started("block".to_string())
+    );
+    signal.cancel();
+    let initial = tokio::time::timeout(Duration::from_secs(5), started.initial_response())
+        .await
+        .expect("yield initial observation")
+        .expect("initial response");
+    assert!(matches!(initial, RuntimeResponse::Yielded { .. }));
+    let signal = CancellationToken::new();
+    signal.cancel(); // The interrupt may precede host admission.
+    let request = WaitRequest {
+        cell_id: cell_id.clone(),
+        yield_time_ms: 60_000,
+    };
+    let response =
+        tokio::time::timeout(Duration::from_secs(5), session.wait(request, Some(signal)))
+            .await
+            .expect("yield wait observation")
+            .expect("wait response");
+    assert!(matches!(
+        response,
+        WaitOutcome::LiveCell(RuntimeResponse::Yielded { .. })
+    ));
+
+    delegate.fast_tool_release.add_permits(1);
+    let response = tokio::time::timeout(
+        Duration::from_secs(5),
+        session.wait(
+            WaitRequest {
+                cell_id,
+                yield_time_ms: 60_000,
+            },
+            /*preempt*/ None,
+        ),
+    )
+    .await
+    .expect("cell completes")
+    .expect("completion response");
+    assert!(matches!(
+        response,
+        WaitOutcome::LiveCell(RuntimeResponse::Result {
+            error_text: None,
+            ..
+        })
+    ));
+    session.shutdown().await.expect("shutdown session");
+}
+
 #[tokio::test]
 async fn session_execution_limits_are_isolated_on_a_shared_process_host() {
     let provider: Arc<dyn CodeModeSessionProvider> =
@@ -269,10 +345,13 @@ async fn session_execution_limits_are_isolated_on_a_shared_process_host() {
     );
     let actual = tokio::time::timeout(
         Duration::from_secs(5),
-        limited.wait(WaitRequest {
-            cell_id: cell_id("1"),
-            yield_time_ms: 60_000,
-        }),
+        limited.wait(
+            WaitRequest {
+                cell_id: cell_id("1"),
+                yield_time_ms: 60_000,
+            },
+            /*preempt*/ None,
+        ),
     )
     .await
     .expect("session limit should bound explicit waits")
@@ -382,10 +461,13 @@ text(result.value);
         }
     );
     let actual = session
-        .wait(WaitRequest {
-            cell_id: cell_id("3"),
-            yield_time_ms: 1,
-        })
+        .wait(
+            WaitRequest {
+                cell_id: cell_id("3"),
+                yield_time_ms: 1,
+            },
+            /*preempt*/ None,
+        )
         .await
         .expect("wait for cell");
     assert_eq!(
@@ -435,7 +517,11 @@ async fn dropping_long_wait_releases_observer_before_next_wait() {
     let mut request = execute_request("await new Promise(() => {});");
     request.yield_time_ms = Some(1);
     let started = session
-        .execute(request, Arc::new(RecordingDelegate::default()))
+        .execute(
+            request,
+            Arc::new(RecordingDelegate::default()),
+            /*preempt*/ None,
+        )
         .await
         .expect("start execution");
     let running_cell_id = started.cell_id.clone();
@@ -453,10 +539,13 @@ async fn dropping_long_wait_releases_observer_before_next_wait() {
     let wait_cell_id = running_cell_id.clone();
     let first_wait = tokio::spawn(async move {
         wait_session
-            .wait(WaitRequest {
-                cell_id: wait_cell_id,
-                yield_time_ms: 60_000,
-            })
+            .wait(
+                WaitRequest {
+                    cell_id: wait_cell_id,
+                    yield_time_ms: 60_000,
+                },
+                /*preempt*/ None,
+            )
             .await
     });
     tokio::time::sleep(Duration::from_millis(100)).await;
@@ -465,10 +554,13 @@ async fn dropping_long_wait_releases_observer_before_next_wait() {
 
     let actual = tokio::time::timeout(
         Duration::from_secs(2),
-        session.wait(WaitRequest {
-            cell_id: running_cell_id.clone(),
-            yield_time_ms: 1,
-        }),
+        session.wait(
+            WaitRequest {
+                cell_id: running_cell_id.clone(),
+                yield_time_ms: 1,
+            },
+            /*preempt*/ None,
+        ),
     )
     .await
     .expect("second wait timeout")
@@ -533,7 +625,7 @@ return;
     .collect();
 
     let started = session
-        .execute(request, delegate.clone())
+        .execute(request, delegate.clone(), /*preempt*/ None)
         .await
         .expect("start execution");
     let running_cell_id = started.cell_id.clone();
@@ -553,10 +645,13 @@ return;
     let wait_cell_id = running_cell_id.clone();
     let wait_task = tokio::spawn(async move {
         wait_session
-            .wait(WaitRequest {
-                cell_id: wait_cell_id,
-                yield_time_ms: 60_000,
-            })
+            .wait(
+                WaitRequest {
+                    cell_id: wait_cell_id,
+                    yield_time_ms: 60_000,
+                },
+                /*preempt*/ None,
+            )
             .await
     });
 
@@ -621,6 +716,7 @@ async fn oversized_execute_request_does_not_close_the_shared_host() {
         .execute(
             execute_request(&"x".repeat(MAX_FRAME_BYTES)),
             Arc::new(RecordingDelegate::default()),
+            /*preempt*/ None,
         )
         .await
         .err()
@@ -760,6 +856,7 @@ async fn oversized_initial_response_does_not_close_the_shared_host() {
         .execute(
             execute_request(&format!(r#"text("x".repeat({MAX_FRAME_BYTES}));"#)),
             Arc::new(RecordingDelegate::default()),
+            /*preempt*/ None,
         )
         .await
         .expect("start oversized response");
@@ -840,7 +937,7 @@ async fn child_process_loss_cleans_up_and_rebuilds_the_shared_host() {
         output_schema: None,
     }];
     let started_a = session_a
-        .execute(request_a, delegate_a.clone())
+        .execute(request_a, delegate_a.clone(), /*preempt*/ None)
         .await
         .expect("start first cell");
     let cell_a = started_a.cell_id.clone();
@@ -864,7 +961,7 @@ async fn child_process_loss_cleans_up_and_rebuilds_the_shared_host() {
     let mut request_b = execute_request("await new Promise(() => {});");
     request_b.yield_time_ms = Some(1);
     let started_b = session_b
-        .execute(request_b, delegate_b.clone())
+        .execute(request_b, delegate_b.clone(), /*preempt*/ None)
         .await
         .expect("start second cell");
     let cell_b = started_b.cell_id.clone();
@@ -885,20 +982,26 @@ async fn child_process_loss_cleans_up_and_rebuilds_the_shared_host() {
     let wait_a_cell = cell_a.clone();
     let wait_a = tokio::spawn(async move {
         wait_a_session
-            .wait(WaitRequest {
-                cell_id: wait_a_cell,
-                yield_time_ms: 60_000,
-            })
+            .wait(
+                WaitRequest {
+                    cell_id: wait_a_cell,
+                    yield_time_ms: 60_000,
+                },
+                /*preempt*/ None,
+            )
             .await
     });
     let wait_b_session = Arc::clone(&session_b);
     let wait_b_cell = cell_b.clone();
     let wait_b = tokio::spawn(async move {
         wait_b_session
-            .wait(WaitRequest {
-                cell_id: wait_b_cell,
-                yield_time_ms: 60_000,
-            })
+            .wait(
+                WaitRequest {
+                    cell_id: wait_b_cell,
+                    yield_time_ms: 60_000,
+                },
+                /*preempt*/ None,
+            )
             .await
     });
     tokio::time::sleep(Duration::from_millis(100)).await;
@@ -975,10 +1078,13 @@ async fn child_process_loss_cleans_up_and_rebuilds_the_shared_host() {
         }
     );
     let stale_error = session_b
-        .wait(WaitRequest {
-            cell_id: cell_b.clone(),
-            yield_time_ms: 1,
-        })
+        .wait(
+            WaitRequest {
+                cell_id: cell_b.clone(),
+                yield_time_ms: 1,
+            },
+            /*preempt*/ None,
+        )
         .await
         .expect_err("stale cell should be rejected");
     assert!(stale_error.contains("stale code-mode host generation"));

codex-rs/code-mode-protocol/src/grpc/codex.code_mode.v1.proto

@@ -29,6 +29,10 @@ service CodeModeHost {
   rpc Execute(ExecuteRequest) returns (stream ExecuteEvent);
   rpc Wait(WaitRequest) returns (WaitResponse);
 
+  // Ends the current observation early without stopping its cell. The
+  // corresponding Execute or Wait still returns its ordinary outcome.
+  rpc YieldObservation(YieldObservationRequest) returns (YieldObservationResponse);
+
   // Acknowledges that a canceled wait has retired before another wait starts.
   rpc CancelWait(CancelWaitRequest) returns (CancelWaitResponse);
   rpc Terminate(TerminateRequest) returns (WaitResponse);
@@ -189,6 +193,16 @@ message CancelWaitRequest {
 
 message CancelWaitResponse {}
 
+message YieldObservationRequest {
+  string session_id = 1;
+  oneof observation {
+    string execution_id = 2;
+    string wait_id = 3;
+  }
+}
+
+message YieldObservationResponse {}
+
 message TerminateRequest {
   string session_id = 1;
   string cell_id = 2;

codex-rs/code-mode-protocol/src/host/message.rs

@@ -135,6 +135,9 @@ pub enum ClientToHost {
     Request { id: RequestId, request: HostRequest },
     #[serde(rename = "operation/cancel")]
     CancelRequest { id: RequestId },
+    /// Yield the active observation without cancelling its cell.
+    #[serde(rename = "operation/yield")]
+    YieldRequest { id: RequestId },
     #[serde(rename = "delegate/response")]
     DelegateResponse {
         id: DelegateRequestId,

codex-rs/code-mode-protocol/src/host/mod.rs

@@ -15,6 +15,9 @@ pub const MAX_PENDING_DELEGATE_CALLS: usize = 1_024;
 /// Negotiated support for cell execution resource limits on `session/open`.
 pub const SESSION_RESOURCE_LIMITS_CAPABILITY: &str = "session-cell-execution-resource-limits";
 
+/// Support for yielding an in-flight execute or wait observation.
+pub const YIELD_OBSERVATION_CAPABILITY: &str = "yield-observation";
+
 pub use codec::EncodedFrame;
 pub use codec::FramedReader;
 pub use codec::FramedWriter;

codex-rs/code-mode-protocol/src/session.rs

@@ -152,13 +152,21 @@ impl CodeModeSessionDelegate for NoopCodeModeSessionDelegate {
 /// must keep those values isolated. Implementations may execute cells
 /// in-process or remotely.
 pub trait CodeModeSession: Send + Sync {
+    /// Executes a cell, yielding its foreground observation if `preempt` is signaled.
+    /// The cell continues running after its observation yields.
     fn execute<'a>(
         &'a self,
         request: ExecuteRequest,
         delegate: Arc<dyn CodeModeSessionDelegate>,
+        preempt: Option<CancellationToken>,
     ) -> CodeModeSessionResultFuture<'a, StartedCell>;
 
-    fn wait<'a>(&'a self, request: WaitRequest) -> CodeModeSessionResultFuture<'a, WaitOutcome>;
+    /// Waits for a cell, yielding the observation if `preempt` is signaled while the cell continues.
+    fn wait<'a>(
+        &'a self,
+        request: WaitRequest,
+        preempt: Option<CancellationToken>,
+    ) -> CodeModeSessionResultFuture<'a, WaitOutcome>;
 
     fn terminate<'a>(&'a self, cell_id: CellId) -> CodeModeSessionResultFuture<'a, WaitOutcome>;
 

codex-rs/code-mode-runtime/src/cell_actor/mod.rs

@@ -38,6 +38,7 @@ use crate::runtime::RuntimeEvent;
 use crate::runtime::spawn_runtime;
 use crate::session_runtime::CellEvent;
 use crate::session_runtime::CreateCellRequest as CellRequest;
+use crate::session_runtime::Observation;
 use crate::session_runtime::ObserveMode;
 use crate::session_runtime::OutputItem;
 use crate::session_runtime::ToolName as CellToolName;
@@ -49,7 +50,7 @@ impl CellActor {
         request: CellRequest,
         stored_values: HashMap<String, Arc<JsonValue>>,
         host: Arc<H>,
-        initial_observe_mode: ObserveMode,
+        initial_observation: Observation,
         cell_state: Arc<CellState>,
         task_failure_handler: Option<TaskFailureHandler>,
     ) -> Result<
@@ -82,7 +83,7 @@ impl CellActor {
             event_rx,
             command_rx,
             Observer {
-                mode: initial_observe_mode,
+                observation: initial_observation,
                 response_tx: initial_response_tx,
             },
             task_failure_handler,
@@ -101,7 +102,7 @@ struct CellContext {
 }
 
 struct Observer {
-    mode: ObserveMode,
+    observation: Observation,
     response_tx: oneshot::Sender<Result<CellEvent, CellError>>,
 }
 
@@ -173,10 +174,11 @@ async fn run_cell<H: CellHost>(
                     None => std::future::pending::<Option<CellCommand>>().await,
                 }
             } => {
-                let Some(CellCommand::Observe { mode, response_tx }) = maybe_command else {
+                let Some(CellCommand::Observe { observation, response_tx }) = maybe_command else {
                     cancellation_token.cancel();
                     continue;
                 };
+                let mode = observation.mode;
                 if response_tx.is_closed() {
                     continue;
                 }
@@ -220,7 +222,7 @@ async fn run_cell<H: CellHost>(
                     }
                     continue;
                 }
-                observer = Some(Observer { mode, response_tx });
+                observer = Some(Observer { observation, response_tx });
                 yield_timer = observer.as_ref().and_then(observer_timer);
                 if runtime_paused && matches!(mode, ObserveMode::YieldAfter(_)) {
                     pending_frontier_ready = false;
@@ -234,10 +236,21 @@ async fn run_cell<H: CellHost>(
                 );
             }
             _ = async {
-                if let Some(yield_timer) = yield_timer.as_mut() {
-                    yield_timer.await;
-                } else {
-                    std::future::pending::<()>().await;
+                tokio::select! {
+                    _ = async {
+                        if let Some(yield_timer) = yield_timer.as_mut() {
+                            yield_timer.await;
+                        } else {
+                            std::future::pending::<()>().await;
+                        }
+                    } => {}
+                    _ = async {
+                        if let Some(observer) = observer.as_ref() {
+                            observer.observation.yield_signal.cancelled().await;
+                        } else {
+                            std::future::pending::<()>().await;
+                        }
+                    } => {}
                 }
             } => {
                 yield_timer = None;
@@ -334,7 +347,7 @@ async fn run_cell<H: CellHost>(
                     RuntimeEvent::Pending => {
                         runtime_paused = true;
                         if matches!(
-                            observer.as_ref().map(|observer| observer.mode),
+                            observer.as_ref().map(|observer| observer.observation.mode),
                             Some(ObserveMode::PendingFrontier)
                         ) {
                             yield_timer = None;
@@ -370,7 +383,7 @@ async fn run_cell<H: CellHost>(
                     RuntimeEvent::ContentItem(item) => content_items.push(output_item(item)),
                     RuntimeEvent::YieldRequested => {
                         let yield_observer = matches!(
-                            observer.as_ref().map(|observer| observer.mode),
+                            observer.as_ref().map(|observer| observer.observation.mode),
                             Some(ObserveMode::YieldAfter(_))
                         );
                         if yield_observer {
@@ -565,7 +578,7 @@ fn finish_termination(
 }
 
 fn observer_timer(observer: &Observer) -> Option<std::pin::Pin<Box<tokio::time::Sleep>>> {
-    match observer.mode {
+    match observer.observation.mode {
         ObserveMode::YieldAfter(duration) => Some(Box::pin(tokio::time::sleep(duration))),
         ObserveMode::PendingFrontier => None,
     }

codex-rs/code-mode-runtime/src/cell_actor/tests.rs

@@ -147,7 +147,7 @@ fn spawn_cell_actor_harness_with_host_and_failure_handler<H: CellHost>(
         event_rx,
         command_rx,
         Observer {
-            mode: initial_observe_mode,
+            observation: initial_observe_mode.into(),
             response_tx: initial_event_tx,
         },
         task_failure_handler,

codex-rs/code-mode-runtime/src/cell_actor/types.rs

@@ -10,6 +10,7 @@ use tokio::sync::oneshot;
 use tokio_util::sync::CancellationToken;
 
 use crate::session_runtime::CellEvent;
+use crate::session_runtime::Observation;
 use crate::session_runtime::ObserveMode;
 use crate::session_runtime::OutputItem;
 use crate::session_runtime::ToolKind;
@@ -76,14 +77,17 @@ impl CellHandle {
         Self { command_tx, state }
     }
 
-    pub(crate) fn observe(&self, mode: ObserveMode) -> CellEventFuture {
+    pub(crate) fn observe(&self, observation: impl Into<Observation>) -> CellEventFuture {
         if !self.state.accepting_observations() {
             return closed_event();
         }
         let (response_tx, response_rx) = oneshot::channel();
         if self
             .command_tx
-            .send(CellCommand::Observe { mode, response_tx })
+            .send(CellCommand::Observe {
+                observation: observation.into(),
+                response_tx,
+            })
             .is_err()
         {
             return closed_event();
@@ -426,7 +430,7 @@ fn prepend_initial_yield(
 
 pub(super) enum CellCommand {
     Observe {
-        mode: ObserveMode,
+        observation: Observation,
         response_tx: oneshot::Sender<Result<CellEvent, CellError>>,
     },
 }

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

@@ -66,13 +66,15 @@ impl InProcessCodeModeSession {
         &self,
         request: ExecuteRequest,
         delegate: Arc<dyn CodeModeSessionDelegate>,
+        preempt: Option<CancellationToken>,
     ) -> Result<StartedCell, String> {
         let yield_time_ms = request.yield_time_ms.unwrap_or(DEFAULT_EXEC_YIELD_TIME_MS);
         let started = self
             .runtime
             .execute(
                 runtime_request(request),
-                runtime::ObserveMode::YieldAfter(self.resolve_yield_timeout(yield_time_ms)),
+                runtime::ObserveMode::YieldAfter(self.resolve_yield_timeout(yield_time_ms))
+                    .with_yield_signal(preempt.unwrap_or_default()),
                 Arc::new(ProtocolDelegate { delegate }),
             )
             .await
@@ -113,13 +115,18 @@ impl InProcessCodeModeSession {
         pending_outcome(&cell_id, event)
     }
 
-    pub async fn wait(&self, request: WaitRequest) -> Result<WaitOutcome, String> {
-        self.begin_wait(request).await.await
+    pub async fn wait(
+        &self,
+        request: WaitRequest,
+        preempt: Option<CancellationToken>,
+    ) -> Result<WaitOutcome, String> {
+        self.begin_wait(request, preempt).await.await
     }
 
     async fn begin_wait(
         &self,
         request: WaitRequest,
+        preempt: Option<CancellationToken>,
     ) -> CodeModeSessionResultFuture<'static, WaitOutcome> {
         let WaitRequest {
             cell_id,
@@ -130,7 +137,8 @@ impl InProcessCodeModeSession {
             .runtime
             .begin_observe(
                 &runtime_cell_id,
-                runtime::ObserveMode::YieldAfter(self.resolve_yield_timeout(yield_time_ms)),
+                runtime::ObserveMode::YieldAfter(self.resolve_yield_timeout(yield_time_ms))
+                    .with_yield_signal(preempt.unwrap_or_default()),
             )
             .await
         {
@@ -216,12 +224,19 @@ impl CodeModeSession for InProcessCodeModeSession {
         &'a self,
         request: ExecuteRequest,
         delegate: Arc<dyn CodeModeSessionDelegate>,
+        preempt: Option<CancellationToken>,
     ) -> CodeModeSessionResultFuture<'a, StartedCell> {
-        Box::pin(InProcessCodeModeSession::execute(self, request, delegate))
+        Box::pin(InProcessCodeModeSession::execute(
+            self, request, delegate, preempt,
+        ))
     }
 
-    fn wait<'a>(&'a self, request: WaitRequest) -> CodeModeSessionResultFuture<'a, WaitOutcome> {
-        Box::pin(InProcessCodeModeSession::wait(self, request))
+    fn wait<'a>(
+        &'a self,
+        request: WaitRequest,
+        preempt: Option<CancellationToken>,
+    ) -> CodeModeSessionResultFuture<'a, WaitOutcome> {
+        Box::pin(InProcessCodeModeSession::wait(self, request, preempt))
     }
 
     fn terminate<'a>(&'a self, cell_id: CellId) -> CodeModeSessionResultFuture<'a, WaitOutcome> {

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

@@ -147,10 +147,13 @@ async fn audio_helper_bounds_wav_data_and_preserves_outputs_after_yield() {
     });
     assert_eq!(
         service
-            .wait(WaitRequest {
-                cell_id: cell_id("1"),
-                yield_time_ms: 60_000,
-            })
+            .wait(
+                WaitRequest {
+                    cell_id: cell_id("1"),
+                    yield_time_ms: 60_000,
+                },
+                /*preempt*/ None,
+            )
             .await
             .unwrap(),
         WaitOutcome::LiveCell(RuntimeResponse::Result {

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

@@ -197,7 +197,7 @@ async fn yielded_cells_retain_their_own_delegate_until_closed() {
         ..execute_request("await tools.block({}); await tools.block({});")
     };
     let a = service
-        .execute(request.clone(), delegate_a.clone())
+        .execute(request.clone(), delegate_a.clone(), /*preempt*/ None)
         .await
         .unwrap();
     assert_eq!(next_event(&mut events_a).await, DelegateEvent::ToolStarted);
@@ -206,7 +206,10 @@ async fn yielded_cells_retain_their_own_delegate_until_closed() {
         RuntimeResponse::Yielded { .. }
     ));
 
-    let b = service.execute(request, delegate_b.clone()).await.unwrap();
+    let b = service
+        .execute(request, delegate_b.clone(), /*preempt*/ None)
+        .await
+        .unwrap();
     assert_eq!(next_event(&mut events_b).await, DelegateEvent::ToolStarted);
     assert!(matches!(
         b.initial_response().await.unwrap(),
@@ -219,10 +222,13 @@ async fn yielded_cells_retain_their_own_delegate_until_closed() {
     delegate_a.release_tool();
     assert_eq!(
         service
-            .wait(WaitRequest {
-                cell_id: cell_id("1"),
-                yield_time_ms: 60_000
-            })
+            .wait(
+                WaitRequest {
+                    cell_id: cell_id("1"),
+                    yield_time_ms: 60_000,
+                },
+                /*preempt*/ None,
+            )
             .await
             .unwrap(),
         WaitOutcome::LiveCell(RuntimeResponse::Result {
@@ -267,6 +273,7 @@ async fn yields_and_resumes() {
                 ..execute_request("")
             },
             Arc::new(NoopCodeModeSessionDelegate),
+            /*preempt*/ None,
         )
         .await
         .unwrap();
@@ -283,10 +290,13 @@ async fn yields_and_resumes() {
     );
     assert_eq!(
         service
-            .wait(WaitRequest {
-                cell_id: cell_id("1"),
-                yield_time_ms: 60_000,
-            })
+            .wait(
+                WaitRequest {
+                    cell_id: cell_id("1"),
+                    yield_time_ms: 60_000,
+                },
+                /*preempt*/ None,
+            )
             .await
             .unwrap(),
         WaitOutcome::LiveCell(RuntimeResponse::Result {
@@ -359,6 +369,7 @@ async fn observed_natural_completion_wins_over_termination() {
         .execute(
             execute_request(r#"yield_control(); store("finished", true); text("done");"#),
             Arc::new(NoopCodeModeSessionDelegate),
+            /*preempt*/ None,
         )
         .await
         .unwrap();
@@ -380,6 +391,7 @@ async fn observed_natural_completion_wins_over_termination() {
                         ..execute_request(r#"text(String(load("finished")));"#)
                     },
                     Arc::new(NoopCodeModeSessionDelegate),
+                    /*preempt*/ None,
                 )
                 .await
                 .unwrap()
@@ -422,6 +434,7 @@ async fn termination_cancels_pending_callbacks_before_responding() {
         .execute(
             execute_request(r#"notify("pending"); await new Promise(() => {});"#),
             delegate.clone(),
+            /*preempt*/ None,
         )
         .await
         .unwrap();
@@ -462,7 +475,11 @@ async fn shutdown_cancels_notifications_while_natural_completion_is_draining() {
     let (delegate, mut events_rx) = HeldNotificationDelegate::new();
     let service = Arc::new(InProcessCodeModeSession::new());
     service
-        .execute(execute_request(r#"notify("pending");"#), delegate.clone())
+        .execute(
+            execute_request(r#"notify("pending");"#),
+            delegate.clone(),
+            /*preempt*/ None,
+        )
         .await
         .unwrap();
 
@@ -495,6 +512,7 @@ async fn repeated_termination_is_rejected_while_callback_cleanup_is_pending() {
         .execute(
             execute_request(r#"notify("pending"); await new Promise(() => {});"#),
             delegate.clone(),
+            /*preempt*/ None,
         )
         .await
         .unwrap();
@@ -548,6 +566,7 @@ async fn second_observer_is_rejected_without_displacing_the_first() {
         .execute(
             execute_request("await new Promise(() => {});"),
             Arc::new(NoopCodeModeSessionDelegate),
+            /*preempt*/ None,
         )
         .await
         .unwrap();
@@ -562,17 +581,23 @@ async fn second_observer_is_rejected_without_displacing_the_first() {
     );
 
     let first_observer = service
-        .begin_wait(WaitRequest {
-            cell_id: cell_id("1"),
-            yield_time_ms: 60_000,
-        })
+        .begin_wait(
+            WaitRequest {
+                cell_id: cell_id("1"),
+                yield_time_ms: 60_000,
+            },
+            /*preempt*/ None,
+        )
         .await;
     assert_eq!(
         service
-            .wait(WaitRequest {
-                cell_id: cell_id("1"),
-                yield_time_ms: 60_000,
-            })
+            .wait(
+                WaitRequest {
+                    cell_id: cell_id("1"),
+                    yield_time_ms: 60_000,
+                },
+                /*preempt*/ None,
+            )
             .await
             .unwrap_err(),
         "exec cell 1 already has an active observer"
@@ -606,6 +631,7 @@ async fn natural_completion_cleans_up_callbacks_before_responding() {
                 ..execute_request("")
             },
             delegate.clone(),
+            /*preempt*/ None,
         )
         .await
         .unwrap();

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

@@ -18,6 +18,7 @@ use crate::ExecuteRequest;
 use crate::ExecuteToPendingOutcome;
 use crate::FunctionCallOutputContentItem;
 use crate::ToolDefinition;
+use codex_code_mode_protocol::CodeModeSession;
 use codex_code_mode_protocol::CodeModeSessionCellExecutionLimits;
 use codex_code_mode_protocol::NotificationFuture;
 use codex_code_mode_protocol::ToolInvocationFuture;
@@ -57,6 +58,95 @@ 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 delegate = Arc::new(ReleasableToolDelegate::default());
+    let service: Arc<dyn CodeModeSession> = Arc::new(InProcessCodeModeSession::new());
+    let execute_signal = CancellationToken::new();
+    let started = service
+        .execute(
+            ExecuteRequest {
+                enabled_tools: vec![echo_tool()],
+                source: r#"await tools.echo({}); text("done");"#.to_string(),
+                yield_time_ms: Some(60_000),
+                ..execute_request("")
+            },
+            delegate.clone(),
+            Some(execute_signal.clone()),
+        )
+        .await
+        .unwrap();
+    let response = tokio::spawn(started.initial_response());
+    wait_until_tool_started(&delegate).await;
+    execute_signal.cancel();
+    wait_until_finished(&response).await;
+    assert_eq!(
+        response.await.unwrap().unwrap(),
+        RuntimeResponse::Yielded {
+            code_mode_host_duration: None,
+            cell_id: cell_id("1"),
+            content_items: Vec::new(),
+        }
+    );
+
+    let wait_signal = CancellationToken::new();
+    let response = tokio::spawn({
+        let service = Arc::clone(&service);
+        let wait_signal = wait_signal.clone();
+        async move {
+            service
+                .wait(
+                    WaitRequest {
+                        cell_id: cell_id("1"),
+                        yield_time_ms: 60_000,
+                    },
+                    Some(wait_signal),
+                )
+                .await
+        }
+    });
+    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 {
+            code_mode_host_duration: None,
+            cell_id: cell_id("1"),
+            content_items: Vec::new(),
+        })
+    );
+
+    delegate.release_tool();
+    let response = tokio::spawn({
+        let service = Arc::clone(&service);
+        async move {
+            service
+                .wait(
+                    WaitRequest {
+                        cell_id: cell_id("1"),
+                        yield_time_ms: 60_000,
+                    },
+                    /*preempt*/ None,
+                )
+                .await
+        }
+    });
+    wait_until_finished(&response).await;
+    assert_eq!(
+        response.await.unwrap().unwrap(),
+        WaitOutcome::LiveCell(RuntimeResponse::Result {
+            code_mode_host_duration: None,
+            cell_id: cell_id("1"),
+            content_items: vec![FunctionCallOutputContentItem::InputText {
+                text: "done".to_string(),
+            }],
+            error_text: None,
+        })
+    );
+}
+
 #[tokio::test(start_paused = true)]
 async fn execute_waits_for_nested_tool_during_yield_grace() {
     let delegate = Arc::new(ReleasableToolDelegate::default());
@@ -67,7 +157,10 @@ async fn execute_waits_for_nested_tool_during_yield_grace() {
         yield_time_ms: Some(10_000),
         ..execute_request("")
     };
-    let started = service.execute(request, delegate.clone()).await.unwrap();
+    let started = service
+        .execute(request, delegate.clone(), /*preempt*/ None)
+        .await
+        .unwrap();
     let response = tokio::spawn(started.initial_response());
     wait_until_tool_started(&delegate).await;
     tokio::time::advance(Duration::from_millis(10_500)).await;
@@ -104,6 +197,7 @@ async fn execute_and_wait_clamp_yield_grace_without_stopping_the_cell() {
                 ..execute_request("")
             },
             delegate.clone(),
+            /*preempt*/ None,
         )
         .await
         .unwrap();
@@ -124,10 +218,13 @@ async fn execute_and_wait_clamp_yield_grace_without_stopping_the_cell() {
     );
 
     let wait_response = service
-        .begin_wait(WaitRequest {
-            cell_id: cell_id("1"),
-            yield_time_ms: 10_000,
-        })
+        .begin_wait(
+            WaitRequest {
+                cell_id: cell_id("1"),
+                yield_time_ms: 10_000,
+            },
+            /*preempt*/ None,
+        )
         .await;
     let wait_response = tokio::spawn(wait_response);
     tokio::task::yield_now().await;
@@ -144,10 +241,13 @@ async fn execute_and_wait_clamp_yield_grace_without_stopping_the_cell() {
 
     delegate.release_tool();
     let completion = service
-        .begin_wait(WaitRequest {
-            cell_id: cell_id("1"),
-            yield_time_ms: 10_000,
-        })
+        .begin_wait(
+            WaitRequest {
+                cell_id: cell_id("1"),
+                yield_time_ms: 10_000,
+            },
+            /*preempt*/ None,
+        )
         .await;
     let completion = tokio::spawn(completion);
     wait_until_finished(&completion).await;
@@ -188,10 +288,13 @@ async fn wait_waits_for_nested_tool_during_yield_grace() {
         }
     );
     let response = service
-        .begin_wait(WaitRequest {
-            cell_id: cell_id("1"),
-            yield_time_ms: 10_000,
-        })
+        .begin_wait(
+            WaitRequest {
+                cell_id: cell_id("1"),
+                yield_time_ms: 10_000,
+            },
+            /*preempt*/ None,
+        )
         .await;
     let response = tokio::spawn(response);
     tokio::task::yield_now().await;
@@ -233,11 +336,15 @@ async fn zero_yield_limit_is_immediate_and_scoped_to_its_session() {
         ..execute_request("")
     };
     let zero_started = zero_session
-        .execute(request.clone(), zero_delegate.clone())
+        .execute(
+            request.clone(),
+            zero_delegate.clone(),
+            /*preempt*/ None,
+        )
         .await
         .unwrap();
     let limited_started = limited_session
-        .execute(request, limited_delegate.clone())
+        .execute(request, limited_delegate.clone(), /*preempt*/ None)
         .await
         .unwrap();
     let zero_response = tokio::spawn(zero_started.initial_response());
@@ -256,10 +363,13 @@ async fn zero_yield_limit_is_immediate_and_scoped_to_its_session() {
     );
 
     let zero_wait = zero_session
-        .begin_wait(WaitRequest {
-            cell_id: cell_id("1"),
-            yield_time_ms: 60_000,
-        })
+        .begin_wait(
+            WaitRequest {
+                cell_id: cell_id("1"),
+                yield_time_ms: 60_000,
+            },
+            /*preempt*/ None,
+        )
         .await;
     let zero_wait = tokio::spawn(zero_wait);
     wait_until_finished(&zero_wait).await;
@@ -375,7 +485,11 @@ fn echo_tool() -> ToolDefinition {
 
 async fn execute(service: &InProcessCodeModeSession, request: ExecuteRequest) -> RuntimeResponse {
     service
-        .execute(request, Arc::new(NoopCodeModeSessionDelegate))
+        .execute(
+            request,
+            Arc::new(NoopCodeModeSessionDelegate),
+            /*preempt*/ None,
+        )
         .await
         .unwrap()
         .initial_response()
@@ -530,6 +644,7 @@ async fn shutdown_interrupts_cpu_bound_cells() {
                 ..execute_request("")
             },
             Arc::new(NoopCodeModeSessionDelegate),
+            /*preempt*/ None,
         )
         .await
         .unwrap();
@@ -557,6 +672,7 @@ async fn start_cell_rejects_new_cell_after_shutdown_begins() {
         .execute(
             execute_request("text('late');"),
             Arc::new(NoopCodeModeSessionDelegate),
+            /*preempt*/ None,
         )
         .await
         .err()
@@ -1618,10 +1734,13 @@ async fn wait_reports_missing_cell_separately_from_runtime_results() {
     let service = InProcessCodeModeSession::new();
 
     let response = service
-        .wait(WaitRequest {
-            cell_id: cell_id("missing"),
-            yield_time_ms: 1,
-        })
+        .wait(
+            WaitRequest {
+                cell_id: cell_id("missing"),
+                yield_time_ms: 1,
+            },
+            /*preempt*/ None,
+        )
         .await
         .unwrap();
 

codex-rs/code-mode-runtime/src/session_runtime/mod.rs

@@ -19,6 +19,7 @@ pub(crate) use self::types::CreateCellRequest;
 pub(crate) use self::types::Error;
 pub(crate) use self::types::ImageDetail;
 pub(crate) use self::types::NestedToolCall;
+pub(crate) use self::types::Observation;
 pub(crate) use self::types::ObserveMode;
 pub(crate) use self::types::OutputItem;
 pub(crate) use self::types::SessionRuntimeDelegate;
@@ -76,7 +77,7 @@ impl SessionRuntime {
     pub(crate) async fn execute<D: SessionRuntimeDelegate>(
         &self,
         request: CreateCellRequest,
-        initial_observe_mode: ObserveMode,
+        initial_observation: impl Into<Observation> + Send,
         delegate: Arc<D>,
     ) -> Result<StartedCell, Error> {
         if self.inner.shutdown_token.is_cancelled() {
@@ -84,7 +85,12 @@ impl SessionRuntime {
         }
         let cell_id = self.allocate_cell_id()?;
         let initial_event = self
-            .start_cell(cell_id.clone(), request, initial_observe_mode, delegate)
+            .start_cell(
+                cell_id.clone(),
+                request,
+                initial_observation.into(),
+                delegate,
+            )
             .await?;
         Ok(StartedCell {
             cell_id,
@@ -103,7 +109,7 @@ impl SessionRuntime {
     pub(crate) async fn begin_observe(
         &self,
         cell_id: &CellId,
-        mode: ObserveMode,
+        observation: impl Into<Observation> + Send,
     ) -> Result<PendingEvent, Error> {
         let handle = self
             .inner
@@ -114,7 +120,7 @@ impl SessionRuntime {
             .cloned()
             .ok_or_else(|| Error::MissingCell(cell_id.clone()))?;
         Ok(PendingEvent {
-            event: map_actor_event(cell_id.clone(), handle.observe(mode)),
+            event: map_actor_event(cell_id.clone(), handle.observe(observation)),
         })
     }
 
@@ -158,7 +164,7 @@ impl SessionRuntime {
         &self,
         cell_id: CellId,
         request: CreateCellRequest,
-        initial_observe_mode: ObserveMode,
+        initial_observation: Observation,
         delegate: Arc<D>,
     ) -> Result<RuntimeEventFuture, Error> {
         let stored_values = self.inner.stored_values.lock().await.clone();
@@ -180,7 +186,7 @@ impl SessionRuntime {
             request,
             stored_values,
             host,
-            initial_observe_mode,
+            initial_observation,
             cell_state,
             self.inner.task_failure_handler.clone(),
         )

codex-rs/code-mode-runtime/src/session_runtime/types.rs

@@ -32,6 +32,27 @@ pub(crate) enum ObserveMode {
     PendingFrontier,
 }
 
+/// A signal may end an observation early without cancelling the cell.
+pub(crate) struct Observation {
+    pub(crate) mode: ObserveMode,
+    pub(crate) yield_signal: CancellationToken,
+}
+
+impl ObserveMode {
+    pub(crate) fn with_yield_signal(self, yield_signal: CancellationToken) -> Observation {
+        Observation {
+            mode: self,
+            yield_signal,
+        }
+    }
+}
+
+impl From<ObserveMode> for Observation {
+    fn from(mode: ObserveMode) -> Self {
+        mode.with_yield_signal(CancellationToken::new())
+    }
+}
+
 /// An observable cell lifecycle event.
 #[derive(Clone, Debug, PartialEq)]
 pub(crate) enum CellEvent {

codex-rs/code-mode-runtime/tests/array_sort.rs

@@ -61,6 +61,7 @@ text(JSON.stringify([3, 1, 2].sort((a, b) => a - b)));
                 max_output_tokens: None,
             },
             Arc::new(NoopCodeModeSessionDelegate),
+            /*preempt*/ None,
         )
         .await
         .expect("start code-mode cell");

codex-rs/code-mode-runtime/tests/jit.rs

@@ -22,6 +22,7 @@ async fn code_mode_runs_with_jit_disabled() {
                 max_output_tokens: None,
             },
             Arc::new(NoopCodeModeSessionDelegate),
+            /*preempt*/ None,
         )
         .await
         .expect("start code-mode cell");

codex-rs/code-mode/src/grpc_session/mod.rs

@@ -221,12 +221,17 @@ impl CodeModeSession for GrpcCodeModeSession {
         &'a self,
         request: ExecuteRequest,
         delegate: Arc<dyn CodeModeSessionDelegate>,
+        preempt: Option<CancellationToken>,
     ) -> CodeModeSessionResultFuture<'a, StartedCell> {
-        Box::pin(self.inner.execute(request, delegate))
+        Box::pin(self.inner.execute(request, delegate, preempt))
     }
 
-    fn wait<'a>(&'a self, request: WaitRequest) -> CodeModeSessionResultFuture<'a, WaitOutcome> {
-        Box::pin(self.inner.wait(request))
+    fn wait<'a>(
+        &'a self,
+        request: WaitRequest,
+        preempt: Option<CancellationToken>,
+    ) -> CodeModeSessionResultFuture<'a, WaitOutcome> {
+        Box::pin(self.inner.wait(request, preempt))
     }
 
     fn terminate<'a>(&'a self, cell_id: CellId) -> CodeModeSessionResultFuture<'a, WaitOutcome> {

codex-rs/code-mode/src/grpc_session/operations.rs

@@ -16,6 +16,8 @@ use codex_code_mode_protocol::grpc;
 use codex_protocol::protocol::W3cTraceContext;
 use tokio::sync::OwnedMutexGuard;
 use tokio::sync::oneshot;
+use tokio_util::sync::CancellationToken;
+use tokio_util::task::AbortOnDropHandle;
 use tracing::Instrument;
 use tracing::debug;
 use uuid::Uuid;
@@ -70,6 +72,7 @@ impl SessionInner {
         self: &Arc<Self>,
         request: ExecuteRequest,
         delegate: Arc<dyn CodeModeSessionDelegate>,
+        preempt: Option<CancellationToken>,
     ) -> Result<StartedCell, String> {
         self.require_open()?;
         let execution_id = Uuid::new_v4().to_string();
@@ -87,12 +90,18 @@ impl SessionInner {
             .begin_execution(&request, delegate)?;
         let ownership = ExecutionOwnership {
             session: Arc::clone(self),
-            execution_id,
+            execution_id: execution_id.clone(),
             armed: true,
         };
         let (started_tx, started_rx) = oneshot::channel();
         let inner = Arc::clone(self);
         self.stream_tasks.spawn(async move {
+            let _watcher = preempt.map(|signal| {
+                inner.watch_yield(
+                    signal,
+                    grpc::yield_observation_request::Observation::ExecutionId(execution_id),
+                )
+            });
             inner
                 .drive_execution(request, ownership, started_tx, trace)
                 .instrument(execute_span)
@@ -249,6 +258,7 @@ impl SessionInner {
     pub(super) async fn wait(
         self: &Arc<Self>,
         request: WaitRequest,
+        preempt: Option<CancellationToken>,
     ) -> Result<WaitOutcome, String> {
         self.require_open()?;
         let slot = {
@@ -297,17 +307,47 @@ impl SessionInner {
         let request = grpc::WaitRequest {
             session_id: self.id.clone(),
             cell_id: expected_cell_id.as_str().to_string(),
-            wait_id,
+            wait_id: wait_id.clone(),
             yield_time_ms: request.yield_time_ms,
         };
         let mut client = self.client();
+        let watcher = preempt.map(|signal| {
+            self.watch_yield(
+                signal,
+                grpc::yield_observation_request::Observation::WaitId(wait_id),
+            )
+        });
         let response = deadline::request(self, "wait", runtime_timeout, client.wait(request)).await;
+        drop(watcher);
         cancellation.disarm();
         self.prune_wait_slots();
         let outcome = conversion::wait_outcome(response?.into_inner())?;
         self.validate_wait_cell(&expected_cell_id, outcome).await
     }
 
+    fn watch_yield(
+        &self,
+        yield_signal: CancellationToken,
+        observation: grpc::yield_observation_request::Observation,
+    ) -> AbortOnDropHandle<()> {
+        let mut client = self.client();
+        let request = grpc::YieldObservationRequest {
+            session_id: self.id.clone(),
+            observation: Some(observation),
+        };
+        AbortOnDropHandle::new(tokio::spawn(async move {
+            yield_signal.cancelled().await;
+            if let Err(error) = deadline::startup("yield observation", async move {
+                client.yield_observation(request).await.map(|_| ())
+            })
+            .await
+            {
+                // An older host may not implement this RPC; the normal timeout still applies.
+                debug!("could not yield code-mode observation early: {error}");
+            }
+        }))
+    }
+
     pub(super) async fn terminate(&self, cell_id: CellId) -> Result<WaitOutcome, String> {
         self.require_open()?;
         let mut client = self.client();

codex-rs/code-mode/src/grpc_session/reconnect.rs

@@ -73,6 +73,7 @@ impl CodeModeSession for ReconnectableSession {
         &'a self,
         request: ExecuteRequest,
         delegate: Arc<dyn CodeModeSessionDelegate>,
+        preempt: Option<CancellationToken>,
     ) -> CodeModeSessionResultFuture<'a, StartedCell> {
         Box::pin(async move {
             let binding = self.inner.get_or_open_binding().await?;
@@ -80,19 +81,23 @@ impl CodeModeSession for ReconnectableSession {
                 delegate,
                 generation: binding.generation,
             });
-            let started = binding.session.execute(request, delegate).await?;
+            let started = binding.session.execute(request, delegate, preempt).await?;
             Ok(generation::public_started_cell(binding.generation, started))
         })
     }
 
-    fn wait<'a>(&'a self, request: WaitRequest) -> CodeModeSessionResultFuture<'a, WaitOutcome> {
+    fn wait<'a>(
+        &'a self,
+        request: WaitRequest,
+        preempt: Option<CancellationToken>,
+    ) -> CodeModeSessionResultFuture<'a, WaitOutcome> {
         Box::pin(async move {
             let binding = self.inner.get_or_open_binding().await?;
             let request = WaitRequest {
                 cell_id: generation::remote_cell_id(binding.generation, &request.cell_id)?,
                 yield_time_ms: request.yield_time_ms,
             };
-            let outcome = binding.session.wait(request).await?;
+            let outcome = binding.session.wait(request, preempt).await?;
             Ok(generation::public_wait_outcome(binding.generation, outcome))
         })
     }

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

@@ -21,6 +21,7 @@ use codex_code_mode_protocol::host::SessionId;
 use codex_install_context::InstallContext;
 use tokio::sync::Semaphore;
 use tokio::sync::watch;
+use tokio_util::sync::CancellationToken;
 
 use self::connection::Connection;
 use self::connection::ConnectionError;
@@ -258,17 +259,25 @@ impl ProcessOwnedCodeModeSession {
         &self,
         request: ExecuteRequest,
         delegate: Arc<dyn CodeModeSessionDelegate>,
+        preempt: Option<CancellationToken>,
     ) -> Result<StartedCell, String> {
         let binding = self.connection().await?;
         binding
             .connection
-            .execute(binding.remote, request, delegate)
+            .execute(binding.remote, request, delegate, preempt)
             .await
     }
 
-    pub async fn wait(&self, request: WaitRequest) -> Result<WaitOutcome, String> {
+    pub async fn wait(
+        &self,
+        request: WaitRequest,
+        preempt: Option<CancellationToken>,
+    ) -> Result<WaitOutcome, String> {
         let binding = self.connection().await?;
-        binding.connection.wait(binding.remote, request).await
+        binding
+            .connection
+            .wait(binding.remote, request, preempt)
+            .await
     }
 
     pub async fn terminate(&self, cell_id: CellId) -> Result<WaitOutcome, String> {
@@ -519,14 +528,19 @@ impl CodeModeSession for ProcessOwnedCodeModeSession {
         &'a self,
         request: ExecuteRequest,
         delegate: Arc<dyn CodeModeSessionDelegate>,
+        preempt: Option<CancellationToken>,
     ) -> CodeModeSessionResultFuture<'a, StartedCell> {
         Box::pin(ProcessOwnedCodeModeSession::execute(
-            self, request, delegate,
+            self, request, delegate, preempt,
         ))
     }
 
-    fn wait<'a>(&'a self, request: WaitRequest) -> CodeModeSessionResultFuture<'a, WaitOutcome> {
-        Box::pin(ProcessOwnedCodeModeSession::wait(self, request))
+    fn wait<'a>(
+        &'a self,
+        request: WaitRequest,
+        preempt: Option<CancellationToken>,
+    ) -> CodeModeSessionResultFuture<'a, WaitOutcome> {
+        Box::pin(ProcessOwnedCodeModeSession::wait(self, request, preempt))
     }
 
     fn terminate<'a>(&'a self, cell_id: CellId) -> CodeModeSessionResultFuture<'a, WaitOutcome> {

codex-rs/code-mode/src/remote_session/connection.rs

@@ -28,6 +28,7 @@ use codex_code_mode_protocol::host::ProtocolVersion;
 use codex_code_mode_protocol::host::RequestId;
 use codex_code_mode_protocol::host::SESSION_RESOURCE_LIMITS_CAPABILITY;
 use codex_code_mode_protocol::host::SupportedProtocolVersions;
+use codex_code_mode_protocol::host::YIELD_OBSERVATION_CAPABILITY;
 use codex_protocol::shell_environment::scrub_non_inheritable_env_vars;
 use tokio::io::AsyncBufReadExt;
 use tokio::io::BufReader;
@@ -204,8 +205,11 @@ impl Connection {
         let handshake = async {
             let session_limits_capability = Capability::new(SESSION_RESOURCE_LIMITS_CAPABILITY)
                 .map_err(|error| error.to_string())?;
-            let optional_capabilities = CapabilitySet::try_new([session_limits_capability])
-                .map_err(|error| error.to_string())?;
+            let yield_capability =
+                Capability::new(YIELD_OBSERVATION_CAPABILITY).map_err(|error| error.to_string())?;
+            let optional_capabilities =
+                CapabilitySet::try_new([session_limits_capability, yield_capability])
+                    .map_err(|error| error.to_string())?;
             let hello = ClientHello::new(
                 SupportedProtocolVersions::try_new([ProtocolVersion::V1])
                     .map_err(|err| err.to_string())?,
@@ -357,7 +361,9 @@ impl Connection {
         session: RemoteSession,
         request: ExecuteRequest,
         delegate: Arc<dyn CodeModeSessionDelegate>,
+        yield_signal: Option<CancellationToken>,
     ) -> Result<StartedCell, String> {
+        let yield_signal = self.supported_yield_signal(yield_signal);
         let cancellation = CallerCancellation::new();
         let (response_tx, response_rx) = oneshot::channel();
         self.send(DriverCommand::Execute {
@@ -365,6 +371,7 @@ impl Connection {
             request,
             delegate,
             caller_cancellation: cancellation.token(),
+            yield_signal,
             response_tx,
         })
         .await?;
@@ -386,7 +393,9 @@ impl Connection {
         &self,
         session: RemoteSession,
         request: WaitRequest,
+        yield_signal: Option<CancellationToken>,
     ) -> Result<WaitOutcome, String> {
+        let yield_signal = self.supported_yield_signal(yield_signal);
         // Account for the runtime's one-second yield grace separately from transport.
         let runtime_timeout =
             Duration::from_millis(request.yield_time_ms).saturating_add(Duration::from_secs(1));
@@ -398,6 +407,7 @@ impl Connection {
                     session,
                     request,
                     caller_cancellation: cancellation.token(),
+                    yield_signal,
                     response_tx,
                 })
                 .await?;
@@ -408,6 +418,17 @@ impl Connection {
         result
     }
 
+    fn supported_yield_signal(
+        &self,
+        signal: Option<CancellationToken>,
+    ) -> Option<CancellationToken> {
+        self.capabilities
+            .iter()
+            .any(|capability| capability.as_str() == YIELD_OBSERVATION_CAPABILITY)
+            .then_some(signal)
+            .flatten()
+    }
+
     pub(super) async fn terminate(
         &self,
         session: RemoteSession,

codex-rs/code-mode/src/remote_session/connection/driver.rs

@@ -122,6 +122,7 @@ impl ConnectionDriver {
             DriverEvent::HostMessage(message) => self.handle_host_message(message),
             DriverEvent::DelegateCompleted { id, result } => self.complete_delegate(id, result),
             DriverEvent::RequestCancelled(id) => self.cancel_request(id),
+            DriverEvent::YieldRequest(id) => self.yield_request(id),
             DriverEvent::Failed(reason) => {
                 self.fail(reason);
                 false

codex-rs/code-mode/src/remote_session/connection/driver/commands.rs

@@ -21,6 +21,7 @@ use super::types::CancellableRequest;
 use super::types::DeferredWait;
 use super::types::DeliveredExecute;
 use super::types::DriverCommand;
+use super::types::ObservationYield;
 use super::types::PendingRequest;
 use super::types::RemoteSession;
 
@@ -39,14 +40,29 @@ impl ConnectionDriver {
                 request,
                 delegate,
                 caller_cancellation,
+                yield_signal,
                 response_tx,
-            } => self.execute(session, request, delegate, caller_cancellation, response_tx),
+            } => self.execute(
+                session,
+                request,
+                delegate,
+                caller_cancellation,
+                yield_signal,
+                response_tx,
+            ),
             DriverCommand::Wait {
                 session,
                 request,
                 caller_cancellation,
+                yield_signal,
                 response_tx,
-            } => self.wait(session, request, caller_cancellation, response_tx),
+            } => self.wait(
+                session,
+                request,
+                caller_cancellation,
+                yield_signal,
+                response_tx,
+            ),
             DriverCommand::Terminate {
                 session,
                 cell_id,
@@ -127,6 +143,7 @@ impl ConnectionDriver {
         request: ExecuteRequest,
         delegate: Arc<dyn CodeModeSessionDelegate>,
         caller_cancellation: CancellationToken,
+        yield_signal: Option<CancellationToken>,
         response_tx: oneshot::Sender<Result<DeliveredExecute, String>>,
     ) -> bool {
         if let Err(err) = self.sessions.require_ready(&session) {
@@ -176,6 +193,7 @@ impl ConnectionDriver {
                 initial_response_tx,
                 initial_response_rx,
                 cancellation,
+                yield_observation: ObservationYield::new(yield_signal),
             },
             &self.event_tx,
         );
@@ -187,6 +205,7 @@ impl ConnectionDriver {
         session: RemoteSession,
         request: WaitRequest,
         caller_cancellation: CancellationToken,
+        yield_signal: Option<CancellationToken>,
         response_tx: oneshot::Sender<Result<WaitOutcome, String>>,
     ) -> bool {
         if let Err(err) = self.sessions.require_ready(&session) {
@@ -205,11 +224,18 @@ impl ConnectionDriver {
                 session,
                 request,
                 caller_cancellation,
+                yield_signal,
                 response_tx,
             });
             return true;
         }
-        self.start_wait(session, request, caller_cancellation, response_tx)
+        self.start_wait(
+            session,
+            request,
+            caller_cancellation,
+            yield_signal,
+            response_tx,
+        )
     }
 
     pub(super) fn start_wait(
@@ -217,6 +243,7 @@ impl ConnectionDriver {
         session: RemoteSession,
         request: WireWaitRequest,
         caller_cancellation: CancellationToken,
+        yield_signal: Option<CancellationToken>,
         response_tx: oneshot::Sender<Result<WaitOutcome, String>>,
     ) -> bool {
         let cell_id = request.cell_id.clone();
@@ -229,6 +256,7 @@ impl ConnectionDriver {
                 session,
                 cell_id,
                 cancellation: CancellableRequest::new(caller_cancellation),
+                yield_observation: ObservationYield::new(yield_signal),
                 response_tx,
             },
         )

codex-rs/code-mode/src/remote_session/connection/driver/request_tracker.rs

@@ -73,9 +73,12 @@ impl RequestTracker {
     pub(super) fn insert_pending(
         &mut self,
         id: RequestId,
-        pending: PendingRequest,
+        mut pending: PendingRequest,
         event_tx: &mpsc::Sender<DriverEvent>,
     ) {
+        if let Some(observation) = pending.yield_observation() {
+            observation.spawn_watcher(id, event_tx.clone());
+        }
         self.pending.insert(id, pending);
         if let Some(cancellation) = self
             .pending

codex-rs/code-mode/src/remote_session/connection/driver/responses.rs

@@ -42,6 +42,7 @@ impl ConnectionDriver {
                 wait.session,
                 wait.request,
                 wait.caller_cancellation,
+                wait.yield_signal,
                 wait.response_tx,
             ) {
                 for wait in deferred {
@@ -182,6 +183,7 @@ impl ConnectionDriver {
                 initial_response_tx,
                 initial_response_rx,
                 cancellation,
+                yield_observation,
             } => match result {
                 Ok(HostResponse::ExecutionStarted { cell_id }) => {
                     // The host owns a checked, never-reused ID sequence. Retain only live
@@ -211,6 +213,7 @@ impl ConnectionDriver {
                             generation: session.generation,
                             cell_id: remote_cell_id.clone(),
                             response_tx: initial_response_tx,
+                            _yield_observation: yield_observation,
                         },
                     );
                     let started = StartedCell::from_result_receiver(public_id, initial_response_rx);
@@ -247,6 +250,7 @@ impl ConnectionDriver {
                 session,
                 cell_id,
                 cancellation: _,
+                yield_observation: _,
                 response_tx,
             } => {
                 let result = match result {
@@ -362,6 +366,17 @@ impl ConnectionDriver {
         }
     }
 
+    pub(super) fn yield_request(&mut self, id: RequestId) -> bool {
+        let frame = match EncodedFrame::encode(&ClientToHost::YieldRequest { id }) {
+            Ok(frame) => frame,
+            Err(err) => {
+                self.fail(format!("failed to encode code-mode yield request: {err}"));
+                return false;
+            }
+        };
+        self.queue_frame(frame)
+    }
+
     fn send_cancel_request(&mut self, id: RequestId) -> bool {
         let message = ClientToHost::CancelRequest { id };
         let frame = match EncodedFrame::encode(&message) {

codex-rs/code-mode/src/remote_session/connection/driver/types.rs

@@ -18,6 +18,7 @@ use codex_code_mode_protocol::host::WireWaitRequest;
 use tokio::sync::mpsc;
 use tokio::sync::oneshot;
 use tokio_util::sync::CancellationToken;
+use tokio_util::task::AbortOnDropHandle;
 
 use super::cleanup::SessionCleanup;
 
@@ -40,12 +41,14 @@ pub(in crate::remote_session::connection) enum DriverCommand {
         delegate: Arc<dyn CodeModeSessionDelegate>,
         request: ExecuteRequest,
         caller_cancellation: CancellationToken,
+        yield_signal: Option<CancellationToken>,
         response_tx: oneshot::Sender<Result<DeliveredExecute, String>>,
     },
     Wait {
         session: RemoteSession,
         request: WaitRequest,
         caller_cancellation: CancellationToken,
+        yield_signal: Option<CancellationToken>,
         response_tx: oneshot::Sender<Result<WaitOutcome, String>>,
     },
     Terminate {
@@ -66,6 +69,7 @@ pub(in crate::remote_session::connection) enum DriverEvent {
         result: Result<DelegateResponse, String>,
     },
     RequestCancelled(RequestId),
+    YieldRequest(RequestId),
     Failed(String),
 }
 
@@ -116,10 +120,36 @@ impl Drop for CancellableRequest {
     }
 }
 
+/// Owns the watcher until an observation's final response is received.
+pub(super) struct ObservationYield {
+    signal: Option<CancellationToken>,
+    watcher: Option<AbortOnDropHandle<()>>,
+}
+
+impl ObservationYield {
+    pub(super) fn new(signal: Option<CancellationToken>) -> Self {
+        Self {
+            signal,
+            watcher: None,
+        }
+    }
+
+    pub(super) fn spawn_watcher(&mut self, id: RequestId, event_tx: mpsc::Sender<DriverEvent>) {
+        let Some(signal) = self.signal.take() else {
+            return;
+        };
+        self.watcher = Some(AbortOnDropHandle::new(tokio::spawn(async move {
+            signal.cancelled().await;
+            let _ = event_tx.send(DriverEvent::YieldRequest(id)).await;
+        })));
+    }
+}
+
 pub(super) struct InitialResponse {
     pub(super) generation: u64,
     pub(super) cell_id: WireCellId,
     pub(super) response_tx: oneshot::Sender<Result<RuntimeResponse, String>>,
+    pub(super) _yield_observation: ObservationYield,
 }
 
 pub(in crate::remote_session::connection) struct DeliveredExecute {
@@ -147,11 +177,13 @@ pub(super) enum PendingRequest {
         initial_response_tx: oneshot::Sender<Result<RuntimeResponse, String>>,
         initial_response_rx: oneshot::Receiver<Result<RuntimeResponse, String>>,
         cancellation: CancellableRequest,
+        yield_observation: ObservationYield,
     },
     Wait {
         session: RemoteSession,
         cell_id: WireCellId,
         cancellation: CancellableRequest,
+        yield_observation: ObservationYield,
         response_tx: oneshot::Sender<Result<WaitOutcome, String>>,
     },
     Terminate {
@@ -169,6 +201,7 @@ pub(super) struct DeferredWait {
     pub(super) session: RemoteSession,
     pub(super) request: WireWaitRequest,
     pub(super) caller_cancellation: CancellationToken,
+    pub(super) yield_signal: Option<CancellationToken>,
     pub(super) response_tx: oneshot::Sender<Result<WaitOutcome, String>>,
 }
 
@@ -182,6 +215,18 @@ impl PendingRequest {
         }
     }
 
+    pub(super) fn yield_observation(&mut self) -> Option<&mut ObservationYield> {
+        match self {
+            Self::Execute {
+                yield_observation, ..
+            }
+            | Self::Wait {
+                yield_observation, ..
+            } => Some(yield_observation),
+            _ => None,
+        }
+    }
+
     pub(super) fn fail(self, reason: String) {
         match self {
             Self::OpenSession { response_tx, .. } | Self::ShutdownSession { response_tx, .. } => {

codex-rs/code-mode/src/remote_session/connection/driver_tests.rs

@@ -140,6 +140,7 @@ impl DriverHarness {
                     max_output_tokens: None,
                 },
                 caller_cancellation: CancellationToken::new(),
+                yield_signal: None,
                 response_tx,
             })
             .await
@@ -192,6 +193,85 @@ impl Drop for DriverHarness {
     }
 }
 
+#[tokio::test(start_paused = true)]
+async fn old_host_does_not_receive_yield_frames() {
+    let mut harness = DriverHarness::start();
+    let session = remote_session();
+    let _cleanup = harness.open(session.clone()).await;
+    let _started = harness
+        .start_cell(
+            session.clone(),
+            /*request_id*/ 2,
+            "1",
+            Arc::new(RecordingDelegate::default()),
+        )
+        .await;
+    harness
+        .event_tx
+        .send(DriverEvent::HostMessage(HostToClient::InitialResponse {
+            id: RequestId::new(/*value*/ 2),
+            result: WireResult::Ok {
+                value: WireRuntimeResponse::Yielded {
+                    cell_id: CellId::new("1".to_string()).into(),
+                    content_items: Vec::new(),
+                    code_mode_host_duration_ns: 1,
+                },
+            },
+        }))
+        .await
+        .unwrap();
+    let connection = Connection {
+        command_tx: harness.command_tx.clone(),
+        execute_claim_tx: harness.execute_claim_tx.clone(),
+        alive: Arc::clone(&harness.alive),
+        failure: Arc::clone(&harness.failure),
+        cancellation: harness.cancellation.clone(),
+        capabilities: CapabilitySet::empty(),
+    };
+    let signal = CancellationToken::new();
+    signal.cancel();
+    let wait = tokio::spawn(async move {
+        connection
+            .wait(
+                session,
+                WaitRequest {
+                    cell_id: CellId::new("1".to_string()),
+                    yield_time_ms: 60_000,
+                },
+                Some(signal),
+            )
+            .await
+    });
+    let frame = harness.outgoing_rx.recv().await.unwrap();
+    assert!(
+        matches!(EncodedFrame::decode_framed::<ClientToHost>(&frame.into_framed_bytes()).unwrap(), ClientToHost::Request { id, request: HostRequest::Wait { .. } } if id == RequestId::new(/*value*/ 3))
+    );
+    // Let any yield watcher run while the wait is still active. A response would
+    // drop the watcher and could otherwise hide an incorrectly sent yield frame.
+    assert!(
+        tokio::time::timeout(Duration::from_millis(1), harness.outgoing_rx.recv())
+            .await
+            .is_err()
+    );
+    harness
+        .event_tx
+        .send(DriverEvent::HostMessage(HostToClient::Response {
+            id: RequestId::new(/*value*/ 3),
+            result: WireResult::Ok {
+                value: HostResponse::WaitCompleted {
+                    outcome: WireWaitOutcome::LiveCell(WireRuntimeResponse::Yielded {
+                        cell_id: CellId::new("1".to_string()).into(),
+                        content_items: Vec::new(),
+                        code_mode_host_duration_ns: 1,
+                    }),
+                },
+            },
+        }))
+        .await
+        .unwrap();
+    wait.await.unwrap().unwrap();
+}
+
 #[tokio::test]
 async fn open_session_includes_nondefault_cell_execution_limits() {
     let mut harness = DriverHarness::start();
@@ -409,6 +489,7 @@ async fn deferred_delegates_follow_cell_readiness_and_cancellation() {
                     max_output_tokens: None,
                 },
                 caller_cancellation: CancellationToken::new(),
+                yield_signal: None,
                 response_tx,
             })
             .await
@@ -597,6 +678,7 @@ async fn dropped_open_waiter_shuts_down_committed_session() {
                 max_output_tokens: None,
             },
             caller_cancellation: CancellationToken::new(),
+            yield_signal: None,
             response_tx: execute_tx,
         })
         .await
@@ -1195,6 +1277,7 @@ async fn mismatched_wait_response_fails_connection() {
                 yield_time_ms: 1,
             },
             caller_cancellation: CancellationToken::new(),
+            yield_signal: None,
             response_tx,
         })
         .await
@@ -1298,6 +1381,7 @@ async fn remote_wait_accepts_durations_longer_than_five_minutes() {
                 yield_time_ms: 300_001,
             },
             caller_cancellation: CancellationToken::new(),
+            yield_signal: None,
             response_tx,
         })
         .await
@@ -1365,6 +1449,7 @@ async fn queued_remote_wait_times_out_and_invalidates_the_connection() {
                     cell_id: CellId::new("1".to_string()),
                     yield_time_ms: 1,
                 },
+                /*yield_signal*/ None,
             )
             .await;
         (connection, result)
@@ -1451,6 +1536,7 @@ async fn cancelled_wait_is_retired_before_next_wait_is_sent() {
                 yield_time_ms: 60_000,
             },
             caller_cancellation: first_cancellation.clone(),
+            yield_signal: None,
             response_tx: first_tx,
         })
         .await
@@ -1469,6 +1555,7 @@ async fn cancelled_wait_is_retired_before_next_wait_is_sent() {
                 yield_time_ms: 1,
             },
             caller_cancellation: CancellationToken::new(),
+            yield_signal: None,
             response_tx: second_tx,
         })
         .await
@@ -1544,6 +1631,7 @@ async fn abandoned_execute_is_tracked_and_terminated_after_admission() {
                 max_output_tokens: None,
             },
             caller_cancellation: cancellation.clone(),
+            yield_signal: None,
             response_tx: execute_tx,
         })
         .await
@@ -1642,6 +1730,7 @@ async fn delivered_but_unclaimed_execute_is_terminated_when_the_caller_is_cancel
                 max_output_tokens: None,
             },
             caller_cancellation: cancellation.clone(),
+            yield_signal: None,
             response_tx: execute_tx,
         })
         .await
@@ -1806,6 +1895,7 @@ async fn connection_failure_closes_every_live_cell_once() {
                 max_output_tokens: None,
             },
             caller_cancellation: CancellationToken::new(),
+            yield_signal: None,
             response_tx: execute_tx,
         })
         .await
@@ -1917,6 +2007,7 @@ async fn aborting_driver_marks_connection_dead_and_closes_cells() {
                 yield_time_ms: 60_000,
             },
             caller_cancellation: CancellationToken::new(),
+            yield_signal: None,
             response_tx: wait_tx,
         })
         .await
@@ -1983,6 +2074,7 @@ async fn dropped_shutdown_waiter_does_not_abort_remote_cleanup() {
                 max_output_tokens: None,
             },
             caller_cancellation: CancellationToken::new(),
+            yield_signal: None,
             response_tx: execute_tx,
         })
         .await

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

@@ -85,6 +85,7 @@ async fn shutdown_before_open_does_not_spawn_the_host() {
                 max_output_tokens: None,
             },
             Arc::new(NoopCodeModeSessionDelegate),
+            /*preempt*/ None,
         )
         .await
         .err()

codex-rs/core/src/tools/code_mode/mod.rs

@@ -137,14 +137,17 @@ impl CodeModeService {
             broker: Arc::clone(&self.dispatch_broker),
             step_context,
         });
-        self.session().await?.execute(request, delegate).await
+        self.session()
+            .await?
+            .execute(request, delegate, /*preempt*/ None)
+            .await
     }
 
     pub(crate) async fn wait(
         &self,
         request: codex_code_mode::WaitRequest,
     ) -> Result<codex_code_mode::WaitOutcome, String> {
-        self.session().await?.wait(request).await
+        self.session().await?.wait(request, /*preempt*/ None).await
     }
 
     pub(crate) async fn terminate(

codex-rs/core/src/tools/tool_dispatch_trace_tests.rs

@@ -80,6 +80,7 @@ impl codex_code_mode::CodeModeSession for MissingCellCodeModeSession {
         &'a self,
         _request: codex_code_mode::ExecuteRequest,
         _delegate: Arc<dyn codex_code_mode::CodeModeSessionDelegate>,
+        _preempt: Option<CancellationToken>,
     ) -> codex_code_mode::CodeModeSessionResultFuture<'a, codex_code_mode::StartedCell> {
         Box::pin(async { Err("test session cannot execute cells".to_string()) })
     }
@@ -87,6 +88,7 @@ impl codex_code_mode::CodeModeSession for MissingCellCodeModeSession {
     fn wait<'a>(
         &'a self,
         request: codex_code_mode::WaitRequest,
+        _preempt: Option<CancellationToken>,
     ) -> codex_code_mode::CodeModeSessionResultFuture<'a, codex_code_mode::WaitOutcome> {
         self.terminate(request.cell_id)
     }