diff --git a/codex-rs/code-mode/src/cell_actor/mod.rs b/codex-rs/code-mode/src/cell_actor/mod.rs index 9889d32ce9..255f1b82e3 100644 --- a/codex-rs/code-mode/src/cell_actor/mod.rs +++ b/codex-rs/code-mode/src/cell_actor/mod.rs @@ -177,6 +177,13 @@ async fn run_cell( let _ = response_tx.send(Ok(event)); break; } + if observer + .as_ref() + .is_some_and(|observer| observer.response_tx.is_closed()) + { + observer = None; + yield_timer = None; + } if observer.is_some() || termination.is_some() { let _ = response_tx.send(Err(CellError::Busy)); continue; diff --git a/codex-rs/code-mode/src/cell_actor/tests.rs b/codex-rs/code-mode/src/cell_actor/tests.rs index 63a3b8e7f1..5ad820e4cc 100644 --- a/codex-rs/code-mode/src/cell_actor/tests.rs +++ b/codex-rs/code-mode/src/cell_actor/tests.rs @@ -9,6 +9,7 @@ use codex_code_mode_protocol::ExecuteRequest; use codex_code_mode_protocol::FunctionCallOutputContentItem; use pretty_assertions::assert_eq; use serde_json::Value as JsonValue; +use tokio::sync::Semaphore; use tokio::sync::mpsc; use tokio::sync::oneshot; use tokio_util::sync::CancellationToken; @@ -18,6 +19,17 @@ use crate::session_runtime::OutputItem; struct TestHost; +#[derive(Default)] +struct RecordingHost { + committed: AtomicBool, + closed: AtomicBool, +} + +struct BlockingCommitHost { + commit_started_tx: mpsc::UnboundedSender<()>, + commit_release: Semaphore, +} + impl CellHost for TestHost { async fn invoke_tool( &self, @@ -41,6 +53,78 @@ impl CellHost for TestHost { async fn closed(&self) {} } +impl CellHost for RecordingHost { + async fn invoke_tool( + &self, + _invocation: CellToolCall, + _cancellation_token: CancellationToken, + ) -> Result { + Err("unexpected tool call".to_string()) + } + + async fn notify( + &self, + _call_id: String, + _text: String, + _cancellation_token: CancellationToken, + ) -> Result<(), String> { + Ok(()) + } + + async fn commit_stored_values(&self, _stored_value_writes: HashMap) { + self.committed.store(true, Ordering::Release); + } + + async fn closed(&self) { + self.closed.store(true, Ordering::Release); + } +} + +impl BlockingCommitHost { + fn new() -> (Arc, mpsc::UnboundedReceiver<()>) { + let (commit_started_tx, commit_started_rx) = mpsc::unbounded_channel(); + ( + Arc::new(Self { + commit_started_tx, + commit_release: Semaphore::new(/*permits*/ 0), + }), + commit_started_rx, + ) + } +} + +impl CellHost for BlockingCommitHost { + async fn invoke_tool( + &self, + _invocation: CellToolCall, + _cancellation_token: CancellationToken, + ) -> Result { + Err("unexpected tool call".to_string()) + } + + async fn notify( + &self, + _call_id: String, + _text: String, + _cancellation_token: CancellationToken, + ) -> Result<(), String> { + Ok(()) + } + + async fn commit_stored_values(&self, _stored_value_writes: HashMap) { + self.commit_started_tx + .send(()) + .expect("test did not receive commit start"); + self.commit_release + .acquire() + .await + .expect("test did not release commit") + .forget(); + } + + async fn closed(&self) {} +} + struct CellActorHarness { event_tx: mpsc::UnboundedSender, command_tx: mpsc::UnboundedSender, @@ -52,6 +136,13 @@ struct CellActorHarness { } fn spawn_cell_actor_harness(initial_observe_mode: ObserveMode) -> CellActorHarness { + spawn_cell_actor_harness_with_host(initial_observe_mode, Arc::new(TestHost)) +} + +fn spawn_cell_actor_harness_with_host( + initial_observe_mode: ObserveMode, + host: Arc, +) -> CellActorHarness { let (event_tx, event_rx) = mpsc::unbounded_channel(); let (command_tx, command_rx) = mpsc::unbounded_channel(); let (initial_event_tx, initial_event_rx) = oneshot::channel(); @@ -70,10 +161,10 @@ fn spawn_cell_actor_harness(initial_observe_mode: ObserveMode) -> CellActorHarne ) .unwrap(); let session_shutdown_token = CancellationToken::new(); - let cancellation_token = CancellationToken::new(); + let cancellation_token = session_shutdown_token.child_token(); let handle = CellHandle::new(command_tx.clone()); let task = tokio::spawn(run_cell( - Arc::new(TestHost), + host, CellContext { runtime_tx, runtime_control_tx, @@ -153,6 +244,165 @@ async fn queued_termination_preempts_unobserved_runtime_completion() { harness.task.await.unwrap(); } +#[tokio::test] +async fn termination_preempts_result_without_committing_stored_values() { + let host = Arc::new(RecordingHost::default()); + let harness = spawn_cell_actor_harness_with_host( + ObserveMode::YieldAfter(Duration::from_secs(/*secs*/ 60)), + Arc::clone(&host), + ); + harness + .event_tx + .send(RuntimeEvent::Result { + stored_value_writes: HashMap::from([("key".to_string(), JsonValue::Bool(true))]), + error_text: None, + }) + .unwrap(); + let termination = harness.handle.terminate(); + + let terminated = Ok(CellEvent::Terminated { + content_items: Vec::new(), + }); + assert_eq!(termination.await, terminated.clone()); + assert_eq!(harness.initial_event_rx.await.unwrap(), terminated); + harness.task.await.unwrap(); + assert!(!host.committed.load(Ordering::Acquire)); + assert!(host.closed.load(Ordering::Acquire)); +} + +#[tokio::test] +async fn observer_receives_a_buffered_completion_after_commit_finishes() { + let (host, mut commit_started_rx) = BlockingCommitHost::new(); + let harness = spawn_cell_actor_harness_with_host( + ObserveMode::YieldAfter(Duration::from_secs(/*secs*/ 60)), + Arc::clone(&host), + ); + harness.event_tx.send(RuntimeEvent::Started).unwrap(); + harness.event_tx.send(RuntimeEvent::YieldRequested).unwrap(); + assert_eq!( + harness.initial_event_rx.await.unwrap(), + Ok(CellEvent::Yielded { + content_items: Vec::new(), + }) + ); + + harness + .event_tx + .send(RuntimeEvent::Result { + stored_value_writes: HashMap::new(), + error_text: None, + }) + .unwrap(); + assert_eq!(commit_started_rx.recv().await, Some(())); + + let completion = harness + .handle + .observe(ObserveMode::YieldAfter(Duration::ZERO)); + host.commit_release.add_permits(/*n*/ 1); + + assert_eq!( + completion.await, + Ok(CellEvent::Completed { + content_items: Vec::new(), + error_text: None, + }) + ); + harness.task.await.unwrap(); +} + +#[tokio::test] +async fn dropped_observer_does_not_block_the_next_observer() { + let harness = spawn_cell_actor_harness(ObserveMode::YieldAfter(Duration::from_secs( + /*secs*/ 60, + ))); + harness.event_tx.send(RuntimeEvent::YieldRequested).unwrap(); + assert_eq!( + harness.initial_event_rx.await.unwrap(), + Ok(CellEvent::Yielded { + content_items: Vec::new(), + }) + ); + + let abandoned_observer = harness + .handle + .observe(ObserveMode::YieldAfter(Duration::ZERO)); + drop(abandoned_observer); + let next_observer = harness + .handle + .observe(ObserveMode::YieldAfter(Duration::ZERO)); + harness.event_tx.send(RuntimeEvent::YieldRequested).unwrap(); + + assert_eq!( + next_observer.await, + Ok(CellEvent::Yielded { + content_items: Vec::new(), + }) + ); + + 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 session_shutdown_terminates_the_active_observer() { + let harness = spawn_cell_actor_harness(ObserveMode::YieldAfter(Duration::from_secs( + /*secs*/ 60, + ))); + harness.session_shutdown_token.cancel(); + drop(harness.event_tx); + + assert_eq!( + harness.initial_event_rx.await.unwrap(), + Ok(CellEvent::Terminated { + content_items: Vec::new(), + }) + ); + harness.task.await.unwrap(); +} + +#[tokio::test] +async fn unexpected_runtime_close_reports_a_completed_error_and_closes_the_actor() { + let host = Arc::new(RecordingHost::default()); + let harness = spawn_cell_actor_harness_with_host( + ObserveMode::YieldAfter(Duration::from_secs(/*secs*/ 60)), + Arc::clone(&host), + ); + drop(harness.event_tx); + + assert_eq!( + harness.initial_event_rx.await.unwrap(), + Ok(CellEvent::Completed { + content_items: Vec::new(), + error_text: Some("exec runtime ended unexpectedly".to_string()), + }) + ); + harness.task.await.unwrap(); + assert!(host.closed.load(Ordering::Acquire)); +} + +#[tokio::test] +async fn termination_wins_when_the_runtime_channel_closes() { + let harness = spawn_cell_actor_harness(ObserveMode::YieldAfter(Duration::from_secs( + /*secs*/ 60, + ))); + let termination = harness.handle.terminate(); + drop(harness.event_tx); + + let terminated = Ok(CellEvent::Terminated { + content_items: Vec::new(), + }); + assert_eq!(termination.await, terminated.clone()); + assert_eq!(harness.initial_event_rx.await.unwrap(), terminated); + harness.task.await.unwrap(); +} + #[tokio::test] async fn session_shutdown_preempts_continuous_command_traffic() { let mut harness = spawn_cell_actor_harness(ObserveMode::YieldAfter(Duration::from_secs( diff --git a/codex-rs/code-mode/src/service_contract_tests.rs b/codex-rs/code-mode/src/service_contract_tests.rs index 9484310323..fb3338dfad 100644 --- a/codex-rs/code-mode/src/service_contract_tests.rs +++ b/codex-rs/code-mode/src/service_contract_tests.rs @@ -16,6 +16,7 @@ use crate::ToolDefinition; #[derive(Debug, PartialEq)] enum DelegateEvent { NotificationStarted, + NotificationFinished, NotificationCancelled, ToolStarted, ToolCancelled, @@ -34,6 +35,11 @@ struct HeldNotificationDelegate { notification_release: Notify, } +struct ReleasableNotificationDelegate { + events_tx: mpsc::UnboundedSender, + notification_release: Notify, +} + impl HeldNotificationDelegate { fn new() -> (Arc, mpsc::UnboundedReceiver) { let (events_tx, events_rx) = mpsc::unbounded_channel(); @@ -51,6 +57,23 @@ impl HeldNotificationDelegate { } } +impl ReleasableNotificationDelegate { + 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(); + } +} + impl CodeModeSessionDelegate for HeldNotificationDelegate { fn invoke_tool<'a>( &'a self, @@ -58,7 +81,9 @@ impl CodeModeSessionDelegate for HeldNotificationDelegate { 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()) }) } @@ -86,6 +111,47 @@ impl CodeModeSessionDelegate for HeldNotificationDelegate { } } +impl CodeModeSessionDelegate for ReleasableNotificationDelegate { + fn invoke_tool<'a>( + &'a self, + _invocation: CodeModeNestedToolCall, + cancellation_token: CancellationToken, + ) -> ToolInvocationFuture<'a> { + Box::pin(async move { + cancellation_token.cancelled().await; + Err("cancelled".to_string()) + }) + } + + fn notify<'a>( + &'a self, + _call_id: String, + _cell_id: CellId, + _text: String, + cancellation_token: CancellationToken, + ) -> NotificationFuture<'a> { + Box::pin(async move { + let _ = self.events_tx.send(DelegateEvent::NotificationStarted); + tokio::select! { + _ = self.notification_release.notified() => { + let _ = self.events_tx.send(DelegateEvent::NotificationFinished); + Ok(()) + } + _ = cancellation_token.cancelled() => { + let _ = self.events_tx.send(DelegateEvent::NotificationCancelled); + Err("cancelled".to_string()) + } + } + }) + } + + 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(); @@ -392,6 +458,159 @@ async fn shutdown_cancels_notifications_while_natural_completion_is_draining() { ); } +#[tokio::test] +async fn natural_completion_waits_for_notifications_before_responding() { + let (delegate, mut events_rx) = ReleasableNotificationDelegate::new(); + let service = CodeModeService::with_delegate(delegate.clone()); + let cell = service + .execute(execute_request(r#"notify("pending");"#)) + .await + .unwrap(); + let mut initial_response = Box::pin(cell.initial_response()); + + assert_eq!( + next_event(&mut events_rx).await, + DelegateEvent::NotificationStarted + ); + assert!( + tokio::time::timeout(Duration::from_millis(/*millis*/ 100), &mut initial_response) + .await + .is_err() + ); + + delegate.release_notification(); + + assert_eq!( + initial_response.await.unwrap(), + RuntimeResponse::Result { + cell_id: cell_id("1"), + content_items: Vec::new(), + error_text: None, + } + ); + assert_eq!( + next_event(&mut events_rx).await, + DelegateEvent::NotificationFinished + ); + assert_eq!( + next_event(&mut events_rx).await, + DelegateEvent::CellClosed(cell_id("1")) + ); +} + +#[tokio::test] +async fn termination_cancels_pending_tools_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#"await tools.block({});"#.to_string(), + ..execute_request("") + }) + .await + .unwrap(); + + assert_eq!(next_event(&mut events_rx).await, DelegateEvent::ToolStarted); + assert_eq!( + cell.initial_response().await.unwrap(), + RuntimeResponse::Yielded { + cell_id: cell_id("1"), + content_items: Vec::new(), + } + ); + + assert_eq!( + service.terminate(cell_id("1")).await.unwrap(), + WaitOutcome::LiveCell(RuntimeResponse::Terminated { + cell_id: cell_id("1"), + 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")) + ); +} + +#[tokio::test] +async fn shutdown_cancels_pending_tools_before_returning() { + let (delegate, mut events_rx) = BlockingDelegate::new(); + let service = Arc::new(CodeModeService::with_delegate(delegate.clone())); + service + .execute(ExecuteRequest { + enabled_tools: vec![blocking_tool()], + source: r#"await tools.block({});"#.to_string(), + ..execute_request("") + }) + .await + .unwrap(); + + 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!( + next_event(&mut events_rx).await, + DelegateEvent::CellClosed(cell_id("1")) + ); +} + +#[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())); + service + .execute(ExecuteRequest { + enabled_tools: vec![blocking_tool()], + source: r#"notify("pending"); 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(); @@ -423,6 +642,9 @@ async fn repeated_termination_is_rejected_while_callback_cleanup_is_pending() { 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(); @@ -437,6 +659,7 @@ async fn repeated_termination_is_rejected_while_callback_cleanup_is_pending() { content_items: Vec::new(), }) ); + assert_eq!(shutdown.await.unwrap(), Ok(())); assert_eq!( next_event(&mut events_rx).await, DelegateEvent::CellClosed(cell_id("1")) @@ -490,6 +713,54 @@ async fn second_observer_is_rejected_without_displacing_the_first() { ); } +#[tokio::test] +async fn dropping_a_wait_allows_a_later_wait_to_observe_the_cell() { + let service = CodeModeService::new(); + let cell = service + .execute(execute_request( + "yield_control(); await new Promise(() => {});", + )) + .await + .unwrap(); + + assert_eq!( + cell.initial_response().await.unwrap(), + RuntimeResponse::Yielded { + cell_id: cell_id("1"), + content_items: Vec::new(), + } + ); + + let abandoned_wait = service + .begin_wait(WaitRequest { + cell_id: cell_id("1"), + yield_time_ms: 60_000, + }) + .await; + drop(abandoned_wait); + + assert_eq!( + service + .wait(WaitRequest { + cell_id: cell_id("1"), + yield_time_ms: 0, + }) + .await + .unwrap(), + WaitOutcome::LiveCell(RuntimeResponse::Yielded { + cell_id: cell_id("1"), + content_items: Vec::new(), + }) + ); + assert_eq!( + service.terminate(cell_id("1")).await.unwrap(), + WaitOutcome::LiveCell(RuntimeResponse::Terminated { + cell_id: cell_id("1"), + content_items: Vec::new(), + }) + ); +} + #[tokio::test] async fn natural_completion_cleans_up_callbacks_before_responding() { let (delegate, mut events_rx) = BlockingDelegate::new();