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