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/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()