diff --git a/codex-rs/code-mode-host/src/grpc/mod.rs b/codex-rs/code-mode-host/src/grpc/mod.rs index 513eb3aaab..a76f825e89 100644 --- a/codex-rs/code-mode-host/src/grpc/mod.rs +++ b/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, 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, + ) -> 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, @@ -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); + } } } diff --git a/codex-rs/code-mode-host/src/grpc/service_tests.rs b/codex-rs/code-mode-host/src/grpc/service_tests.rs index 14c0b8d488..cf30586ad2 100644 --- a/codex-rs/code-mode-host/src/grpc/service_tests.rs +++ b/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(); diff --git a/codex-rs/code-mode-host/src/grpc/session.rs b/codex-rs/code-mode-host/src/grpc/session.rs index d576e6700d..dbfc96ec95 100644 --- a/codex-rs/code-mode-host/src/grpc/session.rs +++ b/codex-rs/code-mode-host/src/grpc/session.rs @@ -58,7 +58,9 @@ pub(super) struct SessionState { pub(super) cells: HashMap, pending_executions: HashSet, pending_closures: HashSet, - seen_executions: BoundedIds, + pub(super) seen_executions: BoundedIds, + pub(super) execution_yields: HashMap, + pub(super) yielded_executions: BoundedIds, pub(super) subscriptions: Vec, pub(super) next_subscription: usize, pub(super) pending_invocations: HashMap, @@ -66,6 +68,7 @@ pub(super) struct SessionState { pub(super) waits: HashMap, 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 { 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() diff --git a/codex-rs/code-mode-host/src/grpc/waits.rs b/codex-rs/code-mode-host/src/grpc/waits.rs index 9f0b0ef477..bbff2d9a95 100644 --- a/codex-rs/code-mode-host/src/grpc/waits.rs +++ b/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, id: String, cancellation: CancellationToken, + yield_signal: CancellationToken, retired: CancellationToken, } @@ -23,6 +25,7 @@ impl WaitRegistration { pub(super) fn new(session: Arc, id: String) -> Result { 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 = { diff --git a/codex-rs/code-mode-host/src/host_tests.rs b/codex-rs/code-mode-host/src/host_tests.rs index da449a780d..0d091ac4ef 100644 --- a/codex-rs/code-mode-host/src/host_tests.rs +++ b/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; diff --git a/codex-rs/code-mode-host/src/lib.rs b/codex-rs/code-mode-host/src/lib.rs index f128fa1513..37687e660c 100644 --- a/codex-rs/code-mode-host/src/lib.rs +++ b/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( 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::, _>>()? + .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 { + ) -> 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; diff --git a/codex-rs/code-mode-host/tests/grpc.rs b/codex-rs/code-mode-host/tests/grpc.rs index 10e706fc32..bf13a4e117 100644 --- a/codex-rs/code-mode-host/tests/grpc.rs +++ b/codex-rs/code-mode-host/tests/grpc.rs @@ -151,7 +151,7 @@ async fn execute( ) -> Result { 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>> { 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")); diff --git a/codex-rs/code-mode-host/tests/grpc/network_policy_tests.rs b/codex-rs/code-mode-host/tests/grpc/network_policy_tests.rs index 5200c358f8..838fd968fa 100644 --- a/codex-rs/code-mode-host/tests/grpc/network_policy_tests.rs +++ b/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)?; diff --git a/codex-rs/code-mode-host/tests/grpc_notifications.rs b/codex-rs/code-mode-host/tests/grpc_notifications.rs index 29533cc5b9..9b9500192f 100644 --- a/codex-rs/code-mode-host/tests/grpc_notifications.rs +++ b/codex-rs/code-mode-host/tests/grpc_notifications.rs @@ -124,7 +124,7 @@ async fn execute( ) -> Result { 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)?; diff --git a/codex-rs/code-mode-host/tests/grpc_timing.rs b/codex-rs/code-mode-host/tests/grpc_timing.rs index 980e2ee697..ec067a842e 100644 --- a/codex-rs/code-mode-host/tests/grpc_timing.rs +++ b/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 diff --git a/codex-rs/code-mode-host/tests/stdio.rs b/codex-rs/code-mode-host/tests/stdio.rs index 0138569d15..e9a0fa805a 100644 --- a/codex-rs/code-mode-host/tests/stdio.rs +++ b/codex-rs/code-mode-host/tests/stdio.rs @@ -180,7 +180,7 @@ async fn execute( delegate: Arc, ) -> 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, ) -> 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 = @@ -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")); diff --git a/codex-rs/code-mode-protocol/src/grpc/codex.code_mode.v1.proto b/codex-rs/code-mode-protocol/src/grpc/codex.code_mode.v1.proto index 7bd1756043..84fb22a373 100644 --- a/codex-rs/code-mode-protocol/src/grpc/codex.code_mode.v1.proto +++ b/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; diff --git a/codex-rs/code-mode-protocol/src/host/message.rs b/codex-rs/code-mode-protocol/src/host/message.rs index bc442426a3..39dd61315a 100644 --- a/codex-rs/code-mode-protocol/src/host/message.rs +++ b/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, diff --git a/codex-rs/code-mode-protocol/src/host/mod.rs b/codex-rs/code-mode-protocol/src/host/mod.rs index 43bb711162..af3f960fdd 100644 --- a/codex-rs/code-mode-protocol/src/host/mod.rs +++ b/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; diff --git a/codex-rs/code-mode-protocol/src/session.rs b/codex-rs/code-mode-protocol/src/session.rs index 36b88b916f..1c87cb8ea1 100644 --- a/codex-rs/code-mode-protocol/src/session.rs +++ b/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, + preempt: Option, ) -> 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, + ) -> CodeModeSessionResultFuture<'a, WaitOutcome>; fn terminate<'a>(&'a self, cell_id: CellId) -> CodeModeSessionResultFuture<'a, WaitOutcome>; diff --git a/codex-rs/code-mode-runtime/src/cell_actor/mod.rs b/codex-rs/code-mode-runtime/src/cell_actor/mod.rs index d403464746..7d02cfb915 100644 --- a/codex-rs/code-mode-runtime/src/cell_actor/mod.rs +++ b/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>, host: Arc, - initial_observe_mode: ObserveMode, + initial_observation: Observation, cell_state: Arc, task_failure_handler: Option, ) -> 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>, } @@ -173,10 +174,11 @@ async fn run_cell( None => std::future::pending::>().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( } 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( ); } _ = 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( 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( 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>> { - match observer.mode { + match observer.observation.mode { ObserveMode::YieldAfter(duration) => Some(Box::pin(tokio::time::sleep(duration))), ObserveMode::PendingFrontier => None, } diff --git a/codex-rs/code-mode-runtime/src/cell_actor/tests.rs b/codex-rs/code-mode-runtime/src/cell_actor/tests.rs index ca472cb3bc..a2cf836b35 100644 --- a/codex-rs/code-mode-runtime/src/cell_actor/tests.rs +++ b/codex-rs/code-mode-runtime/src/cell_actor/tests.rs @@ -147,7 +147,7 @@ fn spawn_cell_actor_harness_with_host_and_failure_handler( event_rx, command_rx, Observer { - mode: initial_observe_mode, + observation: initial_observe_mode.into(), response_tx: initial_event_tx, }, task_failure_handler, diff --git a/codex-rs/code-mode-runtime/src/cell_actor/types.rs b/codex-rs/code-mode-runtime/src/cell_actor/types.rs index 56762f96c9..2aa343022a 100644 --- a/codex-rs/code-mode-runtime/src/cell_actor/types.rs +++ b/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) -> 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>, }, } diff --git a/codex-rs/code-mode-runtime/src/service.rs b/codex-rs/code-mode-runtime/src/service.rs index 002dd8a2b2..37fa8506c7 100644 --- a/codex-rs/code-mode-runtime/src/service.rs +++ b/codex-rs/code-mode-runtime/src/service.rs @@ -66,13 +66,15 @@ impl InProcessCodeModeSession { &self, request: ExecuteRequest, delegate: Arc, + preempt: Option, ) -> Result { 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 { - self.begin_wait(request).await.await + pub async fn wait( + &self, + request: WaitRequest, + preempt: Option, + ) -> Result { + self.begin_wait(request, preempt).await.await } async fn begin_wait( &self, request: WaitRequest, + preempt: Option, ) -> 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, + preempt: Option, ) -> 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, + ) -> CodeModeSessionResultFuture<'a, WaitOutcome> { + Box::pin(InProcessCodeModeSession::wait(self, request, preempt)) } fn terminate<'a>(&'a self, cell_id: CellId) -> CodeModeSessionResultFuture<'a, WaitOutcome> { diff --git a/codex-rs/code-mode-runtime/src/service_audio_tests.rs b/codex-rs/code-mode-runtime/src/service_audio_tests.rs index 45c1bc453a..ec4070e074 100644 --- a/codex-rs/code-mode-runtime/src/service_audio_tests.rs +++ b/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 { diff --git a/codex-rs/code-mode-runtime/src/service_contract_tests.rs b/codex-rs/code-mode-runtime/src/service_contract_tests.rs index 8f9ae2d985..1fb094cd6b 100644 --- a/codex-rs/code-mode-runtime/src/service_contract_tests.rs +++ b/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(); diff --git a/codex-rs/code-mode-runtime/src/service_tests.rs b/codex-rs/code-mode-runtime/src/service_tests.rs index 9785f34fe0..9f7356626f 100644 --- a/codex-rs/code-mode-runtime/src/service_tests.rs +++ b/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 = 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(); diff --git a/codex-rs/code-mode-runtime/src/session_runtime/mod.rs b/codex-rs/code-mode-runtime/src/session_runtime/mod.rs index 0fa6abd58c..0d2f453923 100644 --- a/codex-rs/code-mode-runtime/src/session_runtime/mod.rs +++ b/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( &self, request: CreateCellRequest, - initial_observe_mode: ObserveMode, + initial_observation: impl Into + Send, delegate: Arc, ) -> Result { 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 + Send, ) -> Result { 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, ) -> Result { 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(), ) diff --git a/codex-rs/code-mode-runtime/src/session_runtime/types.rs b/codex-rs/code-mode-runtime/src/session_runtime/types.rs index 8b54e1f32e..d1cdd79cee 100644 --- a/codex-rs/code-mode-runtime/src/session_runtime/types.rs +++ b/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 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 { diff --git a/codex-rs/code-mode-runtime/tests/array_sort.rs b/codex-rs/code-mode-runtime/tests/array_sort.rs index ca9039d6df..d1a25a9b81 100644 --- a/codex-rs/code-mode-runtime/tests/array_sort.rs +++ b/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"); diff --git a/codex-rs/code-mode-runtime/tests/jit.rs b/codex-rs/code-mode-runtime/tests/jit.rs index 2ab5be527e..44359005c4 100644 --- a/codex-rs/code-mode-runtime/tests/jit.rs +++ b/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"); diff --git a/codex-rs/code-mode/src/grpc_session/mod.rs b/codex-rs/code-mode/src/grpc_session/mod.rs index 1905434f3d..5c58592e22 100644 --- a/codex-rs/code-mode/src/grpc_session/mod.rs +++ b/codex-rs/code-mode/src/grpc_session/mod.rs @@ -221,12 +221,17 @@ impl CodeModeSession for GrpcCodeModeSession { &'a self, request: ExecuteRequest, delegate: Arc, + preempt: Option, ) -> 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, + ) -> CodeModeSessionResultFuture<'a, WaitOutcome> { + Box::pin(self.inner.wait(request, preempt)) } fn terminate<'a>(&'a self, cell_id: CellId) -> CodeModeSessionResultFuture<'a, WaitOutcome> { diff --git a/codex-rs/code-mode/src/grpc_session/operations.rs b/codex-rs/code-mode/src/grpc_session/operations.rs index 727eb6ad3a..922a036a3b 100644 --- a/codex-rs/code-mode/src/grpc_session/operations.rs +++ b/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, request: ExecuteRequest, delegate: Arc, + preempt: Option, ) -> Result { 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, request: WaitRequest, + preempt: Option, ) -> Result { 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 { self.require_open()?; let mut client = self.client(); diff --git a/codex-rs/code-mode/src/grpc_session/reconnect.rs b/codex-rs/code-mode/src/grpc_session/reconnect.rs index dee623d265..5d40f8902c 100644 --- a/codex-rs/code-mode/src/grpc_session/reconnect.rs +++ b/codex-rs/code-mode/src/grpc_session/reconnect.rs @@ -73,6 +73,7 @@ impl CodeModeSession for ReconnectableSession { &'a self, request: ExecuteRequest, delegate: Arc, + preempt: Option, ) -> 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, + ) -> 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)) }) } diff --git a/codex-rs/code-mode/src/remote_session.rs b/codex-rs/code-mode/src/remote_session.rs index 0aff7acbd0..36e4e1a74f 100644 --- a/codex-rs/code-mode/src/remote_session.rs +++ b/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, + preempt: Option, ) -> Result { 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 { + pub async fn wait( + &self, + request: WaitRequest, + preempt: Option, + ) -> Result { 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 { @@ -519,14 +528,19 @@ impl CodeModeSession for ProcessOwnedCodeModeSession { &'a self, request: ExecuteRequest, delegate: Arc, + preempt: Option, ) -> 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, + ) -> CodeModeSessionResultFuture<'a, WaitOutcome> { + Box::pin(ProcessOwnedCodeModeSession::wait(self, request, preempt)) } fn terminate<'a>(&'a self, cell_id: CellId) -> CodeModeSessionResultFuture<'a, WaitOutcome> { diff --git a/codex-rs/code-mode/src/remote_session/connection.rs b/codex-rs/code-mode/src/remote_session/connection.rs index eecf9c64ee..48345d506c 100644 --- a/codex-rs/code-mode/src/remote_session/connection.rs +++ b/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, + yield_signal: Option, ) -> Result { + 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, ) -> Result { + 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, + ) -> Option { + self.capabilities + .iter() + .any(|capability| capability.as_str() == YIELD_OBSERVATION_CAPABILITY) + .then_some(signal) + .flatten() + } + pub(super) async fn terminate( &self, session: RemoteSession, diff --git a/codex-rs/code-mode/src/remote_session/connection/driver.rs b/codex-rs/code-mode/src/remote_session/connection/driver.rs index a4c8ab8749..da30a019b8 100644 --- a/codex-rs/code-mode/src/remote_session/connection/driver.rs +++ b/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 diff --git a/codex-rs/code-mode/src/remote_session/connection/driver/commands.rs b/codex-rs/code-mode/src/remote_session/connection/driver/commands.rs index ca6b5bd235..bdbde44c1e 100644 --- a/codex-rs/code-mode/src/remote_session/connection/driver/commands.rs +++ b/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, caller_cancellation: CancellationToken, + yield_signal: Option, response_tx: oneshot::Sender>, ) -> 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, response_tx: oneshot::Sender>, ) -> 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, response_tx: oneshot::Sender>, ) -> 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, }, ) diff --git a/codex-rs/code-mode/src/remote_session/connection/driver/request_tracker.rs b/codex-rs/code-mode/src/remote_session/connection/driver/request_tracker.rs index d39981acd0..75f6724b65 100644 --- a/codex-rs/code-mode/src/remote_session/connection/driver/request_tracker.rs +++ b/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, ) { + 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 diff --git a/codex-rs/code-mode/src/remote_session/connection/driver/responses.rs b/codex-rs/code-mode/src/remote_session/connection/driver/responses.rs index aa48b7bada..6c95cc6e5d 100644 --- a/codex-rs/code-mode/src/remote_session/connection/driver/responses.rs +++ b/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) { diff --git a/codex-rs/code-mode/src/remote_session/connection/driver/types.rs b/codex-rs/code-mode/src/remote_session/connection/driver/types.rs index 19595786db..8aed64908e 100644 --- a/codex-rs/code-mode/src/remote_session/connection/driver/types.rs +++ b/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, request: ExecuteRequest, caller_cancellation: CancellationToken, + yield_signal: Option, response_tx: oneshot::Sender>, }, Wait { session: RemoteSession, request: WaitRequest, caller_cancellation: CancellationToken, + yield_signal: Option, response_tx: oneshot::Sender>, }, Terminate { @@ -66,6 +69,7 @@ pub(in crate::remote_session::connection) enum DriverEvent { result: Result, }, 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, + watcher: Option>, +} + +impl ObservationYield { + pub(super) fn new(signal: Option) -> Self { + Self { + signal, + watcher: None, + } + } + + pub(super) fn spawn_watcher(&mut self, id: RequestId, event_tx: mpsc::Sender) { + 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>, + 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>, initial_response_rx: oneshot::Receiver>, cancellation: CancellableRequest, + yield_observation: ObservationYield, }, Wait { session: RemoteSession, cell_id: WireCellId, cancellation: CancellableRequest, + yield_observation: ObservationYield, response_tx: oneshot::Sender>, }, 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, pub(super) response_tx: oneshot::Sender>, } @@ -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, .. } => { diff --git a/codex-rs/code-mode/src/remote_session/connection/driver_tests.rs b/codex-rs/code-mode/src/remote_session/connection/driver_tests.rs index 6b3428b1af..4442fa62af 100644 --- a/codex-rs/code-mode/src/remote_session/connection/driver_tests.rs +++ b/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::(&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 diff --git a/codex-rs/code-mode/src/remote_session_tests.rs b/codex-rs/code-mode/src/remote_session_tests.rs index e3f86bd2f1..4336cec795 100644 --- a/codex-rs/code-mode/src/remote_session_tests.rs +++ b/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() diff --git a/codex-rs/core/src/tools/code_mode/mod.rs b/codex-rs/core/src/tools/code_mode/mod.rs index 95133ec108..49615443e8 100644 --- a/codex-rs/core/src/tools/code_mode/mod.rs +++ b/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 { - self.session().await?.wait(request).await + self.session().await?.wait(request, /*preempt*/ None).await } pub(crate) async fn terminate( diff --git a/codex-rs/core/src/tools/tool_dispatch_trace_tests.rs b/codex-rs/core/src/tools/tool_dispatch_trace_tests.rs index 1317810361..086ed3902c 100644 --- a/codex-rs/core/src/tools/tool_dispatch_trace_tests.rs +++ b/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, + _preempt: Option, ) -> 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, ) -> codex_code_mode::CodeModeSessionResultFuture<'a, codex_code_mode::WaitOutcome> { self.terminate(request.cell_id) }