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;