diff --git a/codex-rs/code-mode-protocol/src/session.rs b/codex-rs/code-mode-protocol/src/session.rs index 8dffeceabd..a1b7dbffa5 100644 --- a/codex-rs/code-mode-protocol/src/session.rs +++ b/codex-rs/code-mode-protocol/src/session.rs @@ -61,6 +61,12 @@ pub trait CodeModeSessionDelegate: Send + Sync { text: String, cancellation_token: CancellationToken, ) -> NotificationFuture<'a>; + + /// Reports that a cell has reached its terminal state and will issue no more callbacks. + /// + /// Implementations must keep this non-blocking. The session does not wait for an + /// acknowledgement or retry delivery after a transport failure. + fn cell_closed(&self, cell_id: &CellId); } /// A stateful code-mode session owned by one Codex thread. diff --git a/codex-rs/code-mode/src/cell_actor/mod.rs b/codex-rs/code-mode/src/cell_actor/mod.rs index 124a190933..137affe845 100644 --- a/codex-rs/code-mode/src/cell_actor/mod.rs +++ b/codex-rs/code-mode/src/cell_actor/mod.rs @@ -291,7 +291,7 @@ async fn run_cell( content_items: std::mem::take(&mut content_items), error_text: Some("exec runtime ended unexpectedly".to_string()), }; - let rejected_event = match host + let (completion_committed, rejected_event) = match host .commit_completion( HashMap::new(), event, @@ -300,14 +300,18 @@ async fn run_cell( ) .await { - CompletionCommit::Committed => None, - CompletionCommit::Rejected(event) => Some(event), + CompletionCommit::Committed => (true, None), + CompletionCommit::Rejected(event) => (false, Some(event)), }; - match cell_state.deliver_completion( + let completion_delivery = cell_state.deliver_completion( observer .take() .map(|observer| (observer.mode, observer.response_tx)), - ) { + ); + if completion_committed { + host.terminal(); + } + match completion_delivery { CompletionDelivery::Delivered => break, CompletionDelivery::Buffered => {} CompletionDelivery::Rejected(response_tx) => { @@ -438,7 +442,7 @@ async fn run_cell( content_items: std::mem::take(&mut content_items), error_text, }; - let rejected_event = match host + let (completion_committed, rejected_event) = match host .commit_completion( stored_value_writes, event, @@ -447,14 +451,18 @@ async fn run_cell( ) .await { - CompletionCommit::Committed => None, - CompletionCommit::Rejected(event) => Some(event), + CompletionCommit::Committed => (true, None), + CompletionCommit::Rejected(event) => (false, Some(event)), }; - match cell_state.deliver_completion( + let completion_delivery = cell_state.deliver_completion( observer .take() .map(|observer| (observer.mode, observer.response_tx)), - ) { + ); + if completion_committed { + host.terminal(); + } + match completion_delivery { CompletionDelivery::Delivered => break, CompletionDelivery::Buffered => {} CompletionDelivery::Rejected(response_tx) => { diff --git a/codex-rs/code-mode/src/cell_actor/tests.rs b/codex-rs/code-mode/src/cell_actor/tests.rs index 93fc3e4771..66fcf3fbad 100644 --- a/codex-rs/code-mode/src/cell_actor/tests.rs +++ b/codex-rs/code-mode/src/cell_actor/tests.rs @@ -52,6 +52,8 @@ impl CellHost for TestHost { cell_state.commit_completion(event, pending_initial_yield_items, || {}) } + fn terminal(&self) {} + async fn closed(&self) {} } @@ -86,6 +88,8 @@ impl CellHost for RecordingHost { }) } + fn terminal(&self) {} + async fn closed(&self) {} } diff --git a/codex-rs/code-mode/src/cell_actor/types.rs b/codex-rs/code-mode/src/cell_actor/types.rs index 745ce7b23d..3cdaa809f0 100644 --- a/codex-rs/code-mode/src/cell_actor/types.rs +++ b/codex-rs/code-mode/src/cell_actor/types.rs @@ -60,6 +60,9 @@ pub(crate) trait CellHost: Send + Sync + 'static { cell_state: Arc, ) -> impl Future + Send; + /// Reports that execution reached a terminal state, even when its result remains buffered. + fn terminal(&self); + fn closed(&self) -> impl Future + Send; } diff --git a/codex-rs/code-mode/src/service.rs b/codex-rs/code-mode/src/service.rs index ca2955e85d..14f600d970 100644 --- a/codex-rs/code-mode/src/service.rs +++ b/codex-rs/code-mode/src/service.rs @@ -49,6 +49,8 @@ impl CodeModeSessionDelegate for NoopCodeModeSessionDelegate { ) -> NotificationFuture<'a> { Box::pin(async { Ok(()) }) } + + fn cell_closed(&self, _cell_id: &CellId) {} } #[derive(Default)] @@ -251,6 +253,10 @@ impl runtime::SessionRuntimeDelegate for ProtocolDelegate { ) .await } + + fn cell_closed(&self, cell_id: &runtime::CellId) { + self.delegate.cell_closed(&protocol_cell_id(cell_id)); + } } fn runtime_request(request: CreateCellRequest) -> runtime::CreateCellRequest { diff --git a/codex-rs/code-mode/src/service_contract_tests.rs b/codex-rs/code-mode/src/service_contract_tests.rs index 23338c53f9..3e8e665633 100644 --- a/codex-rs/code-mode/src/service_contract_tests.rs +++ b/codex-rs/code-mode/src/service_contract_tests.rs @@ -19,6 +19,11 @@ enum DelegateEvent { NotificationStarted, NotificationFinished, ToolStarted, + CellClosed(CellId), +} + +fn record_cell_closed(events_tx: &mpsc::UnboundedSender, cell_id: &CellId) { + let _ = events_tx.send(DelegateEvent::CellClosed(cell_id.clone())); } struct BlockingDelegate { @@ -100,6 +105,10 @@ impl CodeModeSessionDelegate for NeverResolvingNotificationDelegate { std::future::pending().await }) } + + fn cell_closed(&self, cell_id: &CellId) { + record_cell_closed(&self.events_tx, cell_id); + } } impl CodeModeSessionDelegate for ReleasableNotificationDelegate { @@ -134,6 +143,10 @@ impl CodeModeSessionDelegate for ReleasableNotificationDelegate { } }) } + + fn cell_closed(&self, cell_id: &CellId) { + record_cell_closed(&self.events_tx, cell_id); + } } impl CodeModeSessionDelegate for NeverResolvingToolDelegate { @@ -157,6 +170,10 @@ impl CodeModeSessionDelegate for NeverResolvingToolDelegate { ) -> NotificationFuture<'a> { Box::pin(async { Ok(()) }) } + + fn cell_closed(&self, cell_id: &CellId) { + record_cell_closed(&self.events_tx, cell_id); + } } impl BlockingDelegate { @@ -211,6 +228,10 @@ impl CodeModeSessionDelegate for BlockingDelegate { Err("cancelled".to_string()) }) } + + fn cell_closed(&self, cell_id: &CellId) { + record_cell_closed(&self.events_tx, cell_id); + } } fn cell_id(value: &str) -> CellId { @@ -372,6 +393,39 @@ await tools.block({}); ); } +#[tokio::test] +async fn background_completion_notifies_the_delegate_without_another_observation() { + let (delegate, mut events_rx) = BlockingDelegate::new(); + let service = CodeModeService::with_delegate(delegate); + let created_cell_id = service + .create_cell(execute_request( + r#"await new Promise(resolve => setTimeout(resolve, 100)); text("done");"#, + )) + .await + .unwrap(); + assert_eq!( + tokio::time::timeout( + Duration::from_secs(2), + service.observe(ObserveRequest { + cell_id: created_cell_id.clone(), + yield_time_ms: 1, + }), + ) + .await + .expect("initial observation should yield while the cell is still running") + .unwrap(), + CellOutcome::LiveCell(RuntimeResponse::Yielded { + cell_id: created_cell_id.clone(), + content_items: Vec::new(), + }) + ); + + assert_eq!( + next_event(&mut events_rx).await, + DelegateEvent::CellClosed(created_cell_id) + ); +} + #[tokio::test] async fn returns_and_resumes_from_the_pending_frontier() { let (delegate, mut events_rx) = BlockingDelegate::new(); diff --git a/codex-rs/code-mode/src/service_tests.rs b/codex-rs/code-mode/src/service_tests.rs index e0b036e8df..6cfe8b9107 100644 --- a/codex-rs/code-mode/src/service_tests.rs +++ b/codex-rs/code-mode/src/service_tests.rs @@ -68,6 +68,8 @@ impl CodeModeSessionDelegate for ReleasableToolDelegate { ) -> NotificationFuture<'a> { Box::pin(async { Ok(()) }) } + + fn cell_closed(&self, _cell_id: &CellId) {} } #[derive(Debug, PartialEq)] @@ -82,6 +84,9 @@ enum RecordedDelegateCall { text: String, cancellation_requested: bool, }, + CellClosed { + cell_id: CellId, + }, } struct RecordingDelegate { @@ -150,6 +155,15 @@ impl CodeModeSessionDelegate for RecordingDelegate { .expect("test must provide one result per notification") }) } + + fn cell_closed(&self, cell_id: &CellId) { + self.calls + .lock() + .unwrap() + .push(RecordedDelegateCall::CellClosed { + cell_id: cell_id.clone(), + }); + } } fn execute_request(source: &str) -> CreateCellRequest { @@ -1204,6 +1218,7 @@ async fn protocol_delegate_maps_callbacks_cancellation_and_errors_field_for_fiel .await, Err("notification failed".to_string()) ); + runtime::SessionRuntimeDelegate::cell_closed(&adapter, &runtime::CellId::new("cell-d4")); assert_eq!( delegate.take_calls(), @@ -1234,6 +1249,9 @@ async fn protocol_delegate_maps_callbacks_cancellation_and_errors_field_for_fiel text: "progress".to_string(), cancellation_requested: true, }, + RecordedDelegateCall::CellClosed { + cell_id: cell_id("cell-d4"), + }, ] ); } diff --git a/codex-rs/code-mode/src/session_runtime/mod.rs b/codex-rs/code-mode/src/session_runtime/mod.rs index 3bfe91db9d..693569da4e 100644 --- a/codex-rs/code-mode/src/session_runtime/mod.rs +++ b/codex-rs/code-mode/src/session_runtime/mod.rs @@ -4,6 +4,7 @@ use std::collections::HashMap; use std::future::Future; use std::pin::Pin; use std::sync::Arc; +use std::sync::atomic::AtomicBool; use std::sync::atomic::AtomicU64; use std::sync::atomic::Ordering; @@ -177,6 +178,7 @@ impl SessionRuntime { let host = Arc::new(RuntimeCellHost { cell_id: cell_id.clone(), inner: Arc::clone(&self.inner), + terminal_notified: AtomicBool::new(false), }); let mut cells = self.inner.cells.lock().await; if self.inner.shutdown_token.is_cancelled() { @@ -220,6 +222,7 @@ impl PendingEvent { struct RuntimeCellHost { cell_id: CellId, inner: Arc>, + terminal_notified: AtomicBool, } impl CellHost for RuntimeCellHost { @@ -275,8 +278,15 @@ impl CellHost for RuntimeCellHost { }) } + fn terminal(&self) { + if !self.terminal_notified.swap(true, Ordering::AcqRel) { + self.inner.delegate.cell_closed(&self.cell_id); + } + } + async fn closed(&self) { self.inner.cells.lock().await.remove(&self.cell_id); + self.terminal(); } } diff --git a/codex-rs/code-mode/src/session_runtime/tests.rs b/codex-rs/code-mode/src/session_runtime/tests.rs index c7dfde5a6b..26704bca75 100644 --- a/codex-rs/code-mode/src/session_runtime/tests.rs +++ b/codex-rs/code-mode/src/session_runtime/tests.rs @@ -1,6 +1,7 @@ use std::collections::HashMap; use std::future::Future; use std::sync::Arc; +use std::sync::atomic::AtomicBool; use std::task::Context; use std::task::Poll; use std::task::Waker; @@ -33,6 +34,8 @@ impl SessionRuntimeDelegate for RecordingDelegate { ) -> Result<(), String> { Ok(()) } + + fn cell_closed(&self, _cell_id: &CellId) {} } #[tokio::test] @@ -42,6 +45,7 @@ async fn termination_rejects_a_waiting_store_commit_before_the_next_cell_can_loa let host = RuntimeCellHost { cell_id: CellId::new("terminating-writer"), inner: Arc::clone(&runtime.inner), + terminal_notified: AtomicBool::new(false), }; let completion = CellEvent::Completed { content_items: vec![OutputItem::Text { diff --git a/codex-rs/code-mode/src/session_runtime/types.rs b/codex-rs/code-mode/src/session_runtime/types.rs index 3498b6fc52..a5446c4c00 100644 --- a/codex-rs/code-mode/src/session_runtime/types.rs +++ b/codex-rs/code-mode/src/session_runtime/types.rs @@ -118,7 +118,8 @@ pub struct NestedToolCall { /// Host callbacks used by cells owned by a [`super::SessionRuntime`]. /// /// Implementations should forward callback cancellation tokens to downstream -/// work. The runtime stops awaiting callbacks once cancellation begins. +/// work. The runtime stops awaiting callbacks once cancellation begins. Terminal +/// notifications must be non-blocking and are not acknowledged or retried. pub trait SessionRuntimeDelegate: Send + Sync + 'static { fn invoke_tool( &self, @@ -133,6 +134,11 @@ pub trait SessionRuntimeDelegate: Send + Sync + 'static { text: String, cancellation_token: CancellationToken, ) -> impl Future> + Send; + + /// Reports that execution is terminal and will issue no more callbacks. + /// + /// The terminal result may remain buffered for a later observation. + fn cell_closed(&self, cell_id: &CellId); } /// A failure reported by a session runtime operation. diff --git a/codex-rs/core/src/tools/code_mode/delegate.rs b/codex-rs/core/src/tools/code_mode/delegate.rs index fa6fb89300..993c30eaa3 100644 --- a/codex-rs/core/src/tools/code_mode/delegate.rs +++ b/codex-rs/core/src/tools/code_mode/delegate.rs @@ -238,6 +238,10 @@ impl CodeModeSessionDelegate for CodeModeDispatchBroker { } }) } + + fn cell_closed(&self, cell_id: &CellId) { + self.close_cell(cell_id); + } } enum DispatchMessage { @@ -307,3 +311,7 @@ impl CoreTurnHost { }) } } + +#[cfg(test)] +#[path = "delegate_tests.rs"] +mod tests; diff --git a/codex-rs/core/src/tools/code_mode/delegate_tests.rs b/codex-rs/core/src/tools/code_mode/delegate_tests.rs new file mode 100644 index 0000000000..23e153d4bf --- /dev/null +++ b/codex-rs/core/src/tools/code_mode/delegate_tests.rs @@ -0,0 +1,66 @@ +use std::sync::Arc; +use std::time::Duration; + +use codex_code_mode::CellId; +use codex_code_mode::CellOutcome; +use codex_code_mode::CodeModeService; +use codex_code_mode::CodeModeSessionDelegate; +use codex_code_mode::CreateCellRequest; +use codex_code_mode::ObserveRequest; +use codex_code_mode::RuntimeResponse; + +use super::CodeModeDispatchBroker; + +#[test] +fn terminal_notification_removes_the_cell_dispatch_gate() { + let broker = CodeModeDispatchBroker::new(); + let cell_id = CellId::new("cell-a7".to_string()); + broker.mark_cell_ready_for_dispatch(&cell_id); + assert!(broker.dispatch_gates.lock().unwrap().contains_key(&cell_id)); + + CodeModeSessionDelegate::cell_closed(&broker, &cell_id); + + assert!(!broker.dispatch_gates.lock().unwrap().contains_key(&cell_id)); +} + +#[tokio::test] +async fn background_completion_removes_the_dispatch_gate_without_another_observation() { + let broker = Arc::new(CodeModeDispatchBroker::new()); + let service = CodeModeService::with_delegate(broker.clone()); + let cell_id = service + .create_cell(CreateCellRequest { + tool_call_id: "call-1".to_string(), + enabled_tools: Vec::new(), + source: concat!( + "await new Promise(resolve => setTimeout(resolve, 100));", + "text('done');", + ) + .to_string(), + }) + .await + .unwrap(); + broker.mark_cell_ready_for_dispatch(&cell_id); + + assert_eq!( + service + .observe(ObserveRequest { + cell_id: cell_id.clone(), + yield_time_ms: 1, + }) + .await + .unwrap(), + CellOutcome::LiveCell(RuntimeResponse::Yielded { + cell_id: cell_id.clone(), + content_items: Vec::new(), + }) + ); + tokio::time::timeout(Duration::from_secs(/*secs*/ 2), async { + while broker.dispatch_gates.lock().unwrap().contains_key(&cell_id) { + tokio::task::yield_now().await; + } + }) + .await + .expect("terminal notification should remove the dispatch gate"); + + service.shutdown().await.unwrap(); +}