diff --git a/codex-rs/code-mode/src/cell_actor/mod.rs b/codex-rs/code-mode/src/cell_actor/mod.rs index 124a190933..019e0d990e 100644 --- a/codex-rs/code-mode/src/cell_actor/mod.rs +++ b/codex-rs/code-mode/src/cell_actor/mod.rs @@ -36,6 +36,7 @@ use crate::runtime::RuntimeControlCommand; use crate::runtime::RuntimeEvent; use crate::runtime::spawn_runtime; use crate::session_runtime::CellEvent; +use crate::session_runtime::CellExecutionPolicy; use crate::session_runtime::CreateCellRequest as CellRequest; use crate::session_runtime::ObserveMode; use crate::session_runtime::OutputItem; @@ -49,6 +50,7 @@ impl CellActor { stored_values: HashMap, host: Arc, cell_state: Arc, + execution_policy: CellExecutionPolicy, ) -> Result<(CellHandle, impl Future + Send + 'static), String> { let (event_tx, event_rx) = mpsc::unbounded_channel(); let (command_tx, command_rx) = mpsc::unbounded_channel(); @@ -69,6 +71,7 @@ impl CellActor { }, event_rx, command_rx, + execution_policy, ); Ok((handle, task)) } @@ -91,6 +94,7 @@ async fn run_cell( context: CellContext, mut event_rx: mpsc::UnboundedReceiver, command_rx: mpsc::UnboundedReceiver, + execution_policy: CellExecutionPolicy, ) { let CellContext { runtime_tx, @@ -339,7 +343,12 @@ async fn run_cell( &mut content_items, &mut pending_tool_call_ids, ); - } else if observer.is_some() || has_been_observed { + } else if observer.is_some() + || matches!( + execution_policy, + CellExecutionPolicy::ContinueWhenUnblocked + ) + { pending_frontier_ready = false; pending_tool_call_ids.clear(); let _ = runtime_control_tx.send(RuntimeControlCommand::Continue); diff --git a/codex-rs/code-mode/src/cell_actor/tests.rs b/codex-rs/code-mode/src/cell_actor/tests.rs index b5bf89d753..ec5c112c2c 100644 --- a/codex-rs/code-mode/src/cell_actor/tests.rs +++ b/codex-rs/code-mode/src/cell_actor/tests.rs @@ -104,6 +104,17 @@ fn spawn_cell_actor_harness() -> CellActorHarness { } fn spawn_cell_actor_harness_with_host(host: Arc) -> CellActorHarness { + spawn_cell_actor_harness_with_policy(host, CellExecutionPolicy::ContinueWhenUnblocked) +} + +fn spawn_pausable_cell_actor_harness_with_host(host: Arc) -> CellActorHarness { + spawn_cell_actor_harness_with_policy(host, CellExecutionPolicy::PauseAtPendingFrontier) +} + +fn spawn_cell_actor_harness_with_policy( + host: Arc, + execution_policy: CellExecutionPolicy, +) -> CellActorHarness { let (event_tx, event_rx) = mpsc::unbounded_channel(); let (command_tx, command_rx) = mpsc::unbounded_channel(); let (runtime_event_tx, runtime_event_rx) = mpsc::unbounded_channel(); @@ -133,6 +144,7 @@ fn spawn_cell_actor_harness_with_host(host: Arc) -> CellActorHar }, event_rx, command_rx, + execution_policy, )); CellActorHarness { @@ -195,10 +207,45 @@ async fn completion_and_output_are_buffered_until_the_first_observation() { } #[tokio::test] -async fn pending_frontier_is_buffered_while_runtime_commands_are_queued() { +async fn continuing_harness_advances_an_unobserved_pending_frontier() { let host = Arc::new(RecordingHost::default()); let harness = spawn_cell_actor_harness_with_host(Arc::clone(&host)); harness.event_tx.send(RuntimeEvent::Pending).unwrap(); + harness + .event_tx + .send(RuntimeEvent::Notify { + call_id: "continuing-barrier".to_string(), + text: "barrier".to_string(), + }) + .unwrap(); + wait_for_notification(&host).await; + + loop { + match harness.runtime_control_rx.try_recv() { + Ok(RuntimeControlCommand::Continue) => break, + Ok(command) => panic!("expected continue, got {command:?}"), + Err(std_mpsc::TryRecvError::Empty) => tokio::task::yield_now().await, + Err(std_mpsc::TryRecvError::Disconnected) => { + panic!("runtime control channel disconnected") + } + } + } + let termination = harness.handle.terminate(); + drop(harness.event_tx); + assert_eq!( + termination.await, + Ok(CellEvent::Terminated { + content_items: Vec::new(), + }) + ); + harness.task.await.unwrap(); +} + +#[tokio::test] +async fn pending_frontier_is_buffered_while_runtime_commands_are_queued() { + let host = Arc::new(RecordingHost::default()); + let harness = spawn_pausable_cell_actor_harness_with_host(Arc::clone(&host)); + harness.event_tx.send(RuntimeEvent::Pending).unwrap(); harness .event_tx .send(RuntimeEvent::Notify { @@ -247,7 +294,7 @@ async fn pending_frontier_is_buffered_while_runtime_commands_are_queued() { #[tokio::test] async fn buffered_yield_observation_resumes_an_unobserved_pending_frontier() { let host = Arc::new(RecordingHost::default()); - let harness = spawn_cell_actor_harness_with_host(Arc::clone(&host)); + let harness = spawn_pausable_cell_actor_harness_with_host(Arc::clone(&host)); harness.event_tx.send(RuntimeEvent::YieldRequested).unwrap(); harness.event_tx.send(RuntimeEvent::Pending).unwrap(); harness @@ -295,7 +342,7 @@ async fn buffered_yield_observation_resumes_an_unobserved_pending_frontier() { } assert!(matches!( harness.runtime_control_rx.try_recv(), - Ok(RuntimeControlCommand::Continue) + Err(std_mpsc::TryRecvError::Empty) )); let termination = harness.handle.terminate(); @@ -380,7 +427,7 @@ async fn first_observation_preserves_a_yield_that_raced_with_creation() { #[tokio::test] async fn dropped_pending_observer_preserves_pre_observation_yield() { let host = Arc::new(RecordingHost::default()); - let harness = spawn_cell_actor_harness_with_host(Arc::clone(&host)); + let harness = spawn_pausable_cell_actor_harness_with_host(Arc::clone(&host)); harness .event_tx .send(RuntimeEvent::ContentItem( @@ -456,7 +503,7 @@ async fn dropped_pending_observer_preserves_pre_observation_yield() { #[tokio::test] async fn dropped_pending_observer_preserves_pre_observation_yield_at_completion() { let host = Arc::new(RecordingHost::default()); - let harness = spawn_cell_actor_harness_with_host(Arc::clone(&host)); + let harness = spawn_pausable_cell_actor_harness_with_host(Arc::clone(&host)); harness .event_tx .send(RuntimeEvent::ContentItem( @@ -785,7 +832,7 @@ async fn dropped_yield_observer_preserves_output_for_the_next_observation() { #[tokio::test] async fn dropped_pending_observer_preserves_the_frontier_for_the_next_observation() { let host = Arc::new(RecordingHost::default()); - let harness = spawn_cell_actor_harness_with_host(Arc::clone(&host)); + let harness = spawn_pausable_cell_actor_harness_with_host(Arc::clone(&host)); assert_eq!( harness .handle diff --git a/codex-rs/code-mode/src/service.rs b/codex-rs/code-mode/src/service.rs index 24baa0c457..59afe9b7b1 100644 --- a/codex-rs/code-mode/src/service.rs +++ b/codex-rs/code-mode/src/service.rs @@ -91,7 +91,10 @@ impl CodeModeService { let yield_time_ms = request.yield_time_ms.unwrap_or(DEFAULT_EXEC_YIELD_TIME_MS); let runtime_cell_id = self .runtime - .create_cell(runtime_request(request)) + .create_cell_with_execution_policy( + runtime_request(request), + runtime::CellExecutionPolicy::PauseAtPendingFrontier, + ) .await .map_err(|error| error.to_string())?; let pending_event = self diff --git a/codex-rs/code-mode/src/session_runtime/mod.rs b/codex-rs/code-mode/src/session_runtime/mod.rs index 229d93bcb9..2a5e843273 100644 --- a/codex-rs/code-mode/src/session_runtime/mod.rs +++ b/codex-rs/code-mode/src/session_runtime/mod.rs @@ -13,6 +13,7 @@ use tokio_util::sync::CancellationToken; use tokio_util::task::TaskTracker; pub(crate) use self::types::CellEvent; +pub(crate) use self::types::CellExecutionPolicy; pub(crate) use self::types::CellId; pub(crate) use self::types::CreateCellRequest; pub(crate) use self::types::Error; @@ -68,11 +69,21 @@ impl SessionRuntime { } pub(crate) async fn create_cell(&self, request: CreateCellRequest) -> Result { + self.create_cell_with_execution_policy(request, CellExecutionPolicy::ContinueWhenUnblocked) + .await + } + + pub(crate) async fn create_cell_with_execution_policy( + &self, + request: CreateCellRequest, + execution_policy: CellExecutionPolicy, + ) -> Result { if self.inner.shutdown_token.is_cancelled() { return Err(Error::ShuttingDown); } let cell_id = self.allocate_cell_id(); - self.start_cell(cell_id.clone(), request).await?; + self.start_cell(cell_id.clone(), request, execution_policy) + .await?; Ok(cell_id) } @@ -137,7 +148,12 @@ impl SessionRuntime { ) } - async fn start_cell(&self, cell_id: CellId, request: CreateCellRequest) -> Result<(), Error> { + async fn start_cell( + &self, + cell_id: CellId, + request: CreateCellRequest, + execution_policy: CellExecutionPolicy, + ) -> Result<(), Error> { let stored_values = self.inner.stored_values.lock().await.clone(); let host = Arc::new(RuntimeCellHost { cell_id: cell_id.clone(), @@ -152,7 +168,8 @@ impl SessionRuntime { } let cell_state = Arc::new(CellState::new(self.inner.shutdown_token.child_token())); let (handle, task) = - CellActor::prepare(request, stored_values, host, cell_state).map_err(Error::Runtime)?; + CellActor::prepare(request, stored_values, host, cell_state, execution_policy) + .map_err(Error::Runtime)?; cells.insert(cell_id.clone(), handle); self.inner.cell_tasks.spawn(task); drop(cells); diff --git a/codex-rs/code-mode/src/session_runtime/tests.rs b/codex-rs/code-mode/src/session_runtime/tests.rs index 13c540c786..3b32fd9533 100644 --- a/codex-rs/code-mode/src/session_runtime/tests.rs +++ b/codex-rs/code-mode/src/session_runtime/tests.rs @@ -8,6 +8,7 @@ use std::time::Duration; use pretty_assertions::assert_eq; use serde_json::Value as JsonValue; +use tokio::sync::mpsc; use tokio_util::sync::CancellationToken; use super::*; @@ -15,6 +16,10 @@ use crate::cell_actor::CompletionCommit; struct RecordingDelegate; +struct ImmediateToolDelegate { + invocations_tx: mpsc::UnboundedSender, +} + impl SessionRuntimeDelegate for RecordingDelegate { async fn invoke_tool( &self, @@ -37,6 +42,85 @@ impl SessionRuntimeDelegate for RecordingDelegate { fn cell_closed(&self, _cell_id: &CellId) {} } +impl SessionRuntimeDelegate for ImmediateToolDelegate { + async fn invoke_tool( + &self, + invocation: NestedToolCall, + _cancellation_token: CancellationToken, + ) -> Result { + let _ = self.invocations_tx.send(invocation.tool_name.name); + Ok(JsonValue::Null) + } + + async fn notify( + &self, + _call_id: String, + _cell_id: CellId, + _text: String, + _cancellation_token: CancellationToken, + ) -> Result<(), String> { + Ok(()) + } + + fn cell_closed(&self, _cell_id: &CellId) {} +} + +fn tool_definition(name: &str) -> ToolDefinition { + ToolDefinition { + name: name.to_string(), + tool_name: ToolName { + name: name.to_string(), + namespace: None, + }, + description: String::new(), + kind: ToolKind::Function, + } +} + +#[tokio::test] +async fn continuing_cell_resolves_tools_before_the_first_observation() { + let (invocations_tx, mut invocations_rx) = mpsc::unbounded_channel(); + let runtime = SessionRuntime::new(Arc::new(ImmediateToolDelegate { invocations_tx })); + let cell_id = runtime + .create_cell(CreateCellRequest { + tool_call_id: "call-1".to_string(), + enabled_tools: vec![tool_definition("first"), tool_definition("second")], + source: r#" +await tools.first({}); +await tools.second({}); +text("done"); +"# + .to_string(), + }) + .await + .unwrap(); + + assert_eq!( + tokio::time::timeout(Duration::from_secs(1), invocations_rx.recv()) + .await + .expect("first tool invocation timed out"), + Some("first".to_string()) + ); + assert_eq!( + tokio::time::timeout(Duration::from_secs(1), invocations_rx.recv()) + .await + .expect("second tool invocation timed out"), + Some("second".to_string()) + ); + assert_eq!( + runtime + .observe(&cell_id, ObserveMode::YieldAfter(Duration::from_secs(1))) + .await, + Ok(CellEvent::Completed { + content_items: vec![OutputItem::Text { + text: "done".to_string(), + }], + error_text: None, + }) + ); + runtime.shutdown().await.unwrap(); +} + #[tokio::test] async fn termination_rejects_a_waiting_store_commit_before_the_next_cell_can_load_it() { let runtime = SessionRuntime::new(Arc::new(RecordingDelegate)); diff --git a/codex-rs/code-mode/src/session_runtime/types.rs b/codex-rs/code-mode/src/session_runtime/types.rs index b04bd0f236..ea6816907f 100644 --- a/codex-rs/code-mode/src/session_runtime/types.rs +++ b/codex-rs/code-mode/src/session_runtime/types.rs @@ -32,6 +32,24 @@ pub(crate) enum ObserveMode { PendingFrontier, } +/// Controls how a cell advances when its runtime is waiting for external input. +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub(crate) enum CellExecutionPolicy { + /// Process tool and timer results even when no observation is attached. + ContinueWhenUnblocked, + /// Remain paused at a pending frontier until pending execution is advanced. + PauseAtPendingFrontier, +} + +impl From for CellExecutionPolicy { + fn from(mode: ObserveMode) -> Self { + match mode { + ObserveMode::YieldAfter(_) => Self::ContinueWhenUnblocked, + ObserveMode::PendingFrontier => Self::PauseAtPendingFrontier, + } + } +} + /// An observable cell lifecycle event. #[derive(Clone, Debug, PartialEq)] pub(crate) enum CellEvent {