diff --git a/codex-rs/code-mode/src/cell_actor/callbacks.rs b/codex-rs/code-mode/src/cell_actor/callbacks.rs index 9db66ed3a7..49b54f42b6 100644 --- a/codex-rs/code-mode/src/cell_actor/callbacks.rs +++ b/codex-rs/code-mode/src/cell_actor/callbacks.rs @@ -22,8 +22,14 @@ pub(super) fn spawn_notification( cancellation_token: CancellationToken, ) { tasks.spawn(async move { - if let Err(err) = host.notify(call_id, text, cancellation_token).await { - warn!("failed to deliver code mode notification: {err}"); + tokio::select! { + biased; + _ = cancellation_token.cancelled() => {} + result = host.notify(call_id, text, cancellation_token.clone()) => { + if let Err(err) = result { + warn!("failed to deliver code mode notification: {err}"); + } + } } }); } @@ -37,9 +43,15 @@ pub(super) fn spawn_tool( ) { tasks.spawn(async move { let id = invocation.id.clone(); - let command = match host.invoke_tool(invocation, cancellation_token).await { - Ok(result) => RuntimeCommand::ToolResponse { id, result }, - Err(error_text) => RuntimeCommand::ToolError { id, error_text }, + let command = tokio::select! { + biased; + _ = cancellation_token.cancelled() => { + return; + } + result = host.invoke_tool(invocation, cancellation_token.clone()) => match result { + Ok(result) => RuntimeCommand::ToolResponse { id, result }, + Err(error_text) => RuntimeCommand::ToolError { id, error_text }, + } }; let _ = runtime_tx.send(command); }); diff --git a/codex-rs/code-mode/src/cell_actor/types.rs b/codex-rs/code-mode/src/cell_actor/types.rs index 7db870cd7e..a3a1b6769b 100644 --- a/codex-rs/code-mode/src/cell_actor/types.rs +++ b/codex-rs/code-mode/src/cell_actor/types.rs @@ -34,8 +34,10 @@ pub(crate) struct CellToolCall { /// Connects a cell actor to session-owned callbacks and lifecycle state. /// -/// Implementations must honor callback cancellation and must not return from -/// `closed` until the session can no longer route requests to the cell. +/// Implementations should forward callback cancellation to downstream work. +/// The actor stops awaiting callbacks once cancellation begins. Implementations +/// must not return from `closed` until the session can no longer route requests +/// to the cell. pub(crate) trait CellHost: Send + Sync + 'static { fn invoke_tool( &self, diff --git a/codex-rs/code-mode/src/service_contract_tests.rs b/codex-rs/code-mode/src/service_contract_tests.rs index fb3338dfad..adbe6b6d25 100644 --- a/codex-rs/code-mode/src/service_contract_tests.rs +++ b/codex-rs/code-mode/src/service_contract_tests.rs @@ -1,6 +1,4 @@ use std::sync::Arc; -use std::sync::atomic::AtomicBool; -use std::sync::atomic::Ordering; use std::time::Duration; use codex_protocol::ToolName; @@ -17,22 +15,17 @@ use crate::ToolDefinition; enum DelegateEvent { NotificationStarted, NotificationFinished, - NotificationCancelled, ToolStarted, - ToolCancelled, CellClosed(CellId), } struct BlockingDelegate { events_tx: mpsc::UnboundedSender, - notification_finished: AtomicBool, - tool_finished: AtomicBool, tool_release: Notify, } -struct HeldNotificationDelegate { +struct NeverResolvingNotificationDelegate { events_tx: mpsc::UnboundedSender, - notification_release: Notify, } struct ReleasableNotificationDelegate { @@ -40,20 +33,14 @@ struct ReleasableNotificationDelegate { notification_release: Notify, } -impl HeldNotificationDelegate { +struct NeverResolvingToolDelegate { + events_tx: mpsc::UnboundedSender, +} + +impl NeverResolvingNotificationDelegate { fn new() -> (Arc, mpsc::UnboundedReceiver) { let (events_tx, events_rx) = mpsc::unbounded_channel(); - ( - Arc::new(Self { - events_tx, - notification_release: Notify::new(), - }), - events_rx, - ) - } - - fn release_notification(&self) { - self.notification_release.notify_one(); + (Arc::new(Self { events_tx }), events_rx) } } @@ -74,18 +61,20 @@ impl ReleasableNotificationDelegate { } } -impl CodeModeSessionDelegate for HeldNotificationDelegate { +impl NeverResolvingToolDelegate { + fn new() -> (Arc, mpsc::UnboundedReceiver) { + let (events_tx, events_rx) = mpsc::unbounded_channel(); + (Arc::new(Self { events_tx }), events_rx) + } +} + +impl CodeModeSessionDelegate for NeverResolvingNotificationDelegate { fn invoke_tool<'a>( &'a self, _invocation: CodeModeNestedToolCall, - cancellation_token: CancellationToken, + _cancellation_token: CancellationToken, ) -> ToolInvocationFuture<'a> { - Box::pin(async move { - let _ = self.events_tx.send(DelegateEvent::ToolStarted); - cancellation_token.cancelled().await; - let _ = self.events_tx.send(DelegateEvent::ToolCancelled); - Err("cancelled".to_string()) - }) + Box::pin(async { Err("unexpected tool call".to_string()) }) } fn notify<'a>( @@ -93,14 +82,11 @@ impl CodeModeSessionDelegate for HeldNotificationDelegate { _call_id: String, _cell_id: CellId, _text: String, - cancellation_token: CancellationToken, + _cancellation_token: CancellationToken, ) -> NotificationFuture<'a> { Box::pin(async move { let _ = self.events_tx.send(DelegateEvent::NotificationStarted); - cancellation_token.cancelled().await; - let _ = self.events_tx.send(DelegateEvent::NotificationCancelled); - self.notification_release.notified().await; - Ok(()) + std::future::pending().await }) } @@ -138,7 +124,6 @@ impl CodeModeSessionDelegate for ReleasableNotificationDelegate { Ok(()) } _ = cancellation_token.cancelled() => { - let _ = self.events_tx.send(DelegateEvent::NotificationCancelled); Err("cancelled".to_string()) } } @@ -152,14 +137,41 @@ impl CodeModeSessionDelegate for ReleasableNotificationDelegate { } } +impl CodeModeSessionDelegate for NeverResolvingToolDelegate { + fn invoke_tool<'a>( + &'a self, + _invocation: CodeModeNestedToolCall, + _cancellation_token: CancellationToken, + ) -> ToolInvocationFuture<'a> { + Box::pin(async move { + let _ = self.events_tx.send(DelegateEvent::ToolStarted); + std::future::pending().await + }) + } + + fn notify<'a>( + &'a self, + _call_id: String, + _cell_id: CellId, + _text: String, + _cancellation_token: CancellationToken, + ) -> NotificationFuture<'a> { + Box::pin(async { Ok(()) }) + } + + fn cell_closed(&self, cell_id: &CellId) { + let _ = self + .events_tx + .send(DelegateEvent::CellClosed(cell_id.clone())); + } +} + impl BlockingDelegate { fn new() -> (Arc, mpsc::UnboundedReceiver) { let (events_tx, events_rx) = mpsc::unbounded_channel(); ( Arc::new(Self { events_tx, - notification_finished: AtomicBool::new(false), - tool_finished: AtomicBool::new(false), tool_release: Notify::new(), }), events_rx, @@ -181,12 +193,9 @@ impl CodeModeSessionDelegate for BlockingDelegate { let _ = self.events_tx.send(DelegateEvent::ToolStarted); tokio::select! { _ = self.tool_release.notified() => { - self.tool_finished.store(true, Ordering::Release); Ok(serde_json::Value::Null) } _ = cancellation_token.cancelled() => { - self.tool_finished.store(true, Ordering::Release); - let _ = self.events_tx.send(DelegateEvent::ToolCancelled); Err("cancelled".to_string()) } } @@ -203,8 +212,6 @@ impl CodeModeSessionDelegate for BlockingDelegate { Box::pin(async move { let _ = self.events_tx.send(DelegateEvent::NotificationStarted); cancellation_token.cancelled().await; - self.notification_finished.store(true, Ordering::Release); - let _ = self.events_tx.send(DelegateEvent::NotificationCancelled); Err("cancelled".to_string()) }) } @@ -389,7 +396,7 @@ async fn observed_natural_completion_wins_over_termination() { } #[tokio::test] -async fn termination_cancels_pending_callbacks_before_responding() { +async fn termination_discards_pending_callbacks_before_responding() { let (delegate, mut events_rx) = BlockingDelegate::new(); let service = CodeModeService::with_delegate(delegate.clone()); let cell = service @@ -417,11 +424,6 @@ async fn termination_cancels_pending_callbacks_before_responding() { content_items: Vec::new(), }) ); - assert!(delegate.notification_finished.load(Ordering::Acquire)); - assert_eq!( - next_event(&mut events_rx).await, - DelegateEvent::NotificationCancelled - ); assert_eq!( next_event(&mut events_rx).await, DelegateEvent::CellClosed(cell_id("1")) @@ -429,9 +431,9 @@ async fn termination_cancels_pending_callbacks_before_responding() { } #[tokio::test] -async fn shutdown_cancels_notifications_while_natural_completion_is_draining() { - let (delegate, mut events_rx) = HeldNotificationDelegate::new(); - let service = Arc::new(CodeModeService::with_delegate(delegate.clone())); +async fn shutdown_does_not_await_notifications_during_natural_completion() { + let (delegate, mut events_rx) = NeverResolvingNotificationDelegate::new(); + let service = Arc::new(CodeModeService::with_delegate(delegate)); service .execute(execute_request(r#"notify("pending");"#)) .await @@ -443,15 +445,12 @@ async fn shutdown_cancels_notifications_while_natural_completion_is_draining() { ); let shutdown_service = Arc::clone(&service); - let shutdown = tokio::spawn(async move { shutdown_service.shutdown().await }); - - assert_eq!( - next_event(&mut events_rx).await, - DelegateEvent::NotificationCancelled - ); - delegate.release_notification(); - - assert_eq!(shutdown.await.unwrap(), Ok(())); + tokio::time::timeout(Duration::from_millis(/*millis*/ 100), async move { + shutdown_service.shutdown().await + }) + .await + .expect("shutdown should not await a non-cooperative notification") + .unwrap(); assert_eq!( next_event(&mut events_rx).await, DelegateEvent::CellClosed(cell_id("1")) @@ -499,7 +498,7 @@ async fn natural_completion_waits_for_notifications_before_responding() { } #[tokio::test] -async fn termination_cancels_pending_tools_before_responding() { +async fn termination_discards_pending_tools_before_responding() { let (delegate, mut events_rx) = BlockingDelegate::new(); let service = CodeModeService::with_delegate(delegate.clone()); let cell = service @@ -527,11 +526,6 @@ async fn termination_cancels_pending_tools_before_responding() { content_items: Vec::new(), }) ); - assert!(delegate.tool_finished.load(Ordering::Acquire)); - assert_eq!( - next_event(&mut events_rx).await, - DelegateEvent::ToolCancelled - ); assert_eq!( next_event(&mut events_rx).await, DelegateEvent::CellClosed(cell_id("1")) @@ -539,7 +533,7 @@ async fn termination_cancels_pending_tools_before_responding() { } #[tokio::test] -async fn shutdown_cancels_pending_tools_before_returning() { +async fn shutdown_discards_pending_tools_before_returning() { let (delegate, mut events_rx) = BlockingDelegate::new(); let service = Arc::new(CodeModeService::with_delegate(delegate.clone())); service @@ -554,13 +548,7 @@ async fn shutdown_cancels_pending_tools_before_returning() { assert_eq!(next_event(&mut events_rx).await, DelegateEvent::ToolStarted); let shutdown_service = Arc::clone(&service); - let shutdown = tokio::spawn(async move { shutdown_service.shutdown().await }); - - assert_eq!( - next_event(&mut events_rx).await, - DelegateEvent::ToolCancelled - ); - assert_eq!(shutdown.await.unwrap(), Ok(())); + assert_eq!(shutdown_service.shutdown().await, Ok(())); assert_eq!( next_event(&mut events_rx).await, DelegateEvent::CellClosed(cell_id("1")) @@ -568,98 +556,27 @@ async fn shutdown_cancels_pending_tools_before_returning() { } #[tokio::test] -async fn shutdown_cancels_mixed_callbacks_while_natural_completion_is_draining() { - let (delegate, mut events_rx) = HeldNotificationDelegate::new(); - let service = Arc::new(CodeModeService::with_delegate(delegate.clone())); +async fn shutdown_completes_when_a_nested_tool_ignores_cancellation() { + let (delegate, mut events_rx) = NeverResolvingToolDelegate::new(); + let service = Arc::new(CodeModeService::with_delegate(delegate)); service .execute(ExecuteRequest { enabled_tools: vec![blocking_tool()], - source: r#"notify("pending"); tools.block({});"#.to_string(), + source: r#"await tools.block({});"#.to_string(), ..execute_request("") }) .await .unwrap(); - assert_eq!( - next_event(&mut events_rx).await, - DelegateEvent::NotificationStarted - ); assert_eq!(next_event(&mut events_rx).await, DelegateEvent::ToolStarted); let shutdown_service = Arc::clone(&service); - let shutdown = tokio::spawn(async move { shutdown_service.shutdown().await }); - - let first_cancelled = next_event(&mut events_rx).await; - let second_cancelled = next_event(&mut events_rx).await; - assert!(matches!( - (&first_cancelled, &second_cancelled), - ( - &DelegateEvent::NotificationCancelled, - &DelegateEvent::ToolCancelled, - ) | ( - &DelegateEvent::ToolCancelled, - &DelegateEvent::NotificationCancelled, - ) - )); - - delegate.release_notification(); - - assert_eq!(shutdown.await.unwrap(), Ok(())); - assert_eq!( - next_event(&mut events_rx).await, - DelegateEvent::CellClosed(cell_id("1")) - ); -} - -#[tokio::test] -async fn repeated_termination_is_rejected_while_callback_cleanup_is_pending() { - let (delegate, mut events_rx) = HeldNotificationDelegate::new(); - let service = Arc::new(CodeModeService::with_delegate(delegate.clone())); - let cell = service - .execute(execute_request( - r#"notify("pending"); await new Promise(() => {});"#, - )) - .await - .unwrap(); - - assert_eq!( - next_event(&mut events_rx).await, - DelegateEvent::NotificationStarted - ); - assert_eq!( - cell.initial_response().await.unwrap(), - RuntimeResponse::Yielded { - cell_id: cell_id("1"), - content_items: Vec::new(), - } - ); - - let terminating_service = Arc::clone(&service); - let first_termination = - tokio::spawn(async move { terminating_service.terminate(cell_id("1")).await }); - assert_eq!( - next_event(&mut events_rx).await, - DelegateEvent::NotificationCancelled - ); - - let shutdown_service = Arc::clone(&service); - let shutdown = tokio::spawn(async move { shutdown_service.shutdown().await }); - - let repeated_termination = service.terminate(cell_id("1")).await; - delegate.release_notification(); - - assert_eq!( - repeated_termination.unwrap_err(), - "exec cell 1 is already terminating" - ); - assert_eq!( - first_termination.await.unwrap().unwrap(), - WaitOutcome::LiveCell(RuntimeResponse::Terminated { - cell_id: cell_id("1"), - content_items: Vec::new(), - }) - ); - assert_eq!(shutdown.await.unwrap(), Ok(())); + tokio::time::timeout(Duration::from_millis(/*millis*/ 100), async move { + shutdown_service.shutdown().await + }) + .await + .expect("shutdown should not await a non-cooperative nested tool") + .unwrap(); assert_eq!( next_event(&mut events_rx).await, DelegateEvent::CellClosed(cell_id("1")) @@ -760,39 +677,3 @@ async fn dropping_a_wait_allows_a_later_wait_to_observe_the_cell() { }) ); } - -#[tokio::test] -async fn natural_completion_cleans_up_callbacks_before_responding() { - let (delegate, mut events_rx) = BlockingDelegate::new(); - let service = CodeModeService::with_delegate(delegate.clone()); - let cell = service - .execute(ExecuteRequest { - enabled_tools: vec![blocking_tool()], - source: r#"tools.block({}); text("done");"#.to_string(), - yield_time_ms: Some(60_000), - ..execute_request("") - }) - .await - .unwrap(); - - assert_eq!(next_event(&mut events_rx).await, DelegateEvent::ToolStarted); - assert_eq!( - cell.initial_response().await.unwrap(), - RuntimeResponse::Result { - cell_id: cell_id("1"), - content_items: vec![FunctionCallOutputContentItem::InputText { - text: "done".to_string(), - }], - error_text: None, - } - ); - assert!(delegate.tool_finished.load(Ordering::Acquire)); - assert_eq!( - next_event(&mut events_rx).await, - DelegateEvent::ToolCancelled - ); - assert_eq!( - next_event(&mut events_rx).await, - DelegateEvent::CellClosed(cell_id("1")) - ); -} diff --git a/codex-rs/code-mode/src/session_runtime/tests.rs b/codex-rs/code-mode/src/session_runtime/tests.rs index 998df92723..dd0021c825 100644 --- a/codex-rs/code-mode/src/session_runtime/tests.rs +++ b/codex-rs/code-mode/src/session_runtime/tests.rs @@ -20,7 +20,6 @@ enum FirstCloseOutcome { #[derive(Debug, PartialEq)] enum NotificationEvent { Started, - Cancelled, Closed(CellId), } @@ -113,9 +112,7 @@ impl SessionRuntimeDelegate for BlockingNotificationDelegate { .send(NotificationEvent::Started) .map_err(|_| "test did not receive notification start".to_string())?; cancellation_token.cancelled().await; - self.events_tx - .send(NotificationEvent::Cancelled) - .map_err(|_| "test did not receive notification cancellation".to_string()) + Ok(()) } async fn cell_closed(&self, cell_id: &CellId) -> Result<(), String> { @@ -506,12 +503,6 @@ async fn drop_cancels_notifications_during_natural_completion_when_registry_is_l drop(runtime); drop(cells); - assert_eq!( - tokio::time::timeout(Duration::from_secs(/*secs*/ 1), events_rx.recv()) - .await - .unwrap(), - Some(NotificationEvent::Cancelled) - ); assert_eq!( tokio::time::timeout(Duration::from_secs(/*secs*/ 1), events_rx.recv()) .await diff --git a/codex-rs/code-mode/src/session_runtime/types.rs b/codex-rs/code-mode/src/session_runtime/types.rs index fb86f9c228..79ae8fac00 100644 --- a/codex-rs/code-mode/src/session_runtime/types.rs +++ b/codex-rs/code-mode/src/session_runtime/types.rs @@ -112,8 +112,10 @@ pub struct NestedToolCall { /// Host callbacks used by cells owned by a [`super::SessionRuntime`]. /// -/// Implementations must honor cancellation tokens and must not return from -/// `cell_closed` until downstream routing can no longer target the cell. +/// Implementations should forward callback cancellation tokens to downstream +/// work. The runtime stops awaiting callbacks once cancellation begins. +/// Implementations must not return from `cell_closed` until downstream routing +/// can no longer target the cell. pub trait SessionRuntimeDelegate: Send + Sync + 'static { fn invoke_tool( &self,