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-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");