From 442d204af4c83df052a9f0f41b6a34bbff8bbd6d Mon Sep 17 00:00:00 2001 From: Channing Conger Date: Sat, 20 Jun 2026 18:07:21 +0000 Subject: [PATCH] code-mode: make session shutdown authoritative --- .../code-mode/src/cell_actor/callbacks.rs | 22 +- codex-rs/code-mode/src/cell_actor/types.rs | 9 +- .../code-mode/src/service_contract_tests.rs | 203 ++++++++++-------- codex-rs/code-mode/src/session_runtime/mod.rs | 49 ++--- .../code-mode/src/session_runtime/tests.rs | 77 +++++++ 5 files changed, 235 insertions(+), 125 deletions(-) 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 a593f5e579..a14fea4d68 100644 --- a/codex-rs/code-mode/src/cell_actor/types.rs +++ b/codex-rs/code-mode/src/cell_actor/types.rs @@ -93,16 +93,13 @@ impl CellHandle { pub(crate) fn terminate(&self) -> CellEventFuture { self.state.request_termination() } - - pub(crate) fn shutdown(&self) { - self.state.cancellation_token().cancel(); - } } /// The single linearization point for a cell's terminal outcome. /// -/// Callback cancellation tokens are children of the cell token, so a terminal -/// decision cancels runtime work and its callbacks together. +/// The cancellation token is a child of the owning session token. Callback +/// tokens are children of this token, so cancellation flows strictly from the +/// session to the cell and then to its callbacks. /// /// The mutex is held only for synchronous phase transitions and terminal /// delivery. Runtime execution, observation waits, and callbacks never run diff --git a/codex-rs/code-mode/src/service_contract_tests.rs b/codex-rs/code-mode/src/service_contract_tests.rs index 9484310323..9770dcf7fe 100644 --- a/codex-rs/code-mode/src/service_contract_tests.rs +++ b/codex-rs/code-mode/src/service_contract_tests.rs @@ -16,50 +16,84 @@ use crate::ToolDefinition; #[derive(Debug, PartialEq)] enum DelegateEvent { NotificationStarted, - NotificationCancelled, ToolStarted, - ToolCancelled, CellClosed(CellId), } struct BlockingDelegate { events_tx: mpsc::UnboundedSender, - notification_finished: AtomicBool, - tool_finished: AtomicBool, + tool_future_dropped: AtomicBool, tool_release: Notify, } -struct HeldNotificationDelegate { - events_tx: mpsc::UnboundedSender, - notification_release: Notify, +struct DropFlag<'a>(&'a AtomicBool); + +impl Drop for DropFlag<'_> { + fn drop(&mut self) { + self.0.store(true, Ordering::Release); + } } -impl HeldNotificationDelegate { +struct NeverResolvingNotificationDelegate { + events_tx: mpsc::UnboundedSender, +} + +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) } } -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 { Err("unexpected tool call".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); + std::future::pending().await + }) + } + + fn cell_closed(&self, cell_id: &CellId) { + let _ = self + .events_tx + .send(DelegateEvent::CellClosed(cell_id.clone())); + } +} + +impl CodeModeSessionDelegate for NeverResolvingToolDelegate { + 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()) + let _ = self.events_tx.send(DelegateEvent::ToolStarted); + std::future::pending().await }) } @@ -68,15 +102,9 @@ 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(()) - }) + Box::pin(async { Ok(()) }) } fn cell_closed(&self, cell_id: &CellId) { @@ -92,8 +120,7 @@ impl BlockingDelegate { ( Arc::new(Self { events_tx, - notification_finished: AtomicBool::new(false), - tool_finished: AtomicBool::new(false), + tool_future_dropped: AtomicBool::new(false), tool_release: Notify::new(), }), events_rx, @@ -111,16 +138,15 @@ impl CodeModeSessionDelegate for BlockingDelegate { _invocation: CodeModeNestedToolCall, cancellation_token: CancellationToken, ) -> ToolInvocationFuture<'a> { + let drop_flag = DropFlag(&self.tool_future_dropped); Box::pin(async move { + let _drop_flag = drop_flag; 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()) } } @@ -137,8 +163,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()) }) } @@ -289,7 +313,10 @@ async fn observed_natural_completion_wins_over_termination() { tokio::time::timeout(Duration::from_secs(1), async { loop { let response = service - .execute(execute_request(r#"text(String(load("finished")));"#)) + .execute(ExecuteRequest { + yield_time_ms: Some(60_000), + ..execute_request(r#"text(String(load("finished")));"#) + }) .await .unwrap() .initial_response() @@ -323,7 +350,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 @@ -351,11 +378,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")) @@ -363,9 +385,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 @@ -377,15 +399,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")) @@ -393,9 +412,33 @@ async fn shutdown_cancels_notifications_while_natural_completion_is_draining() { } #[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())); +async fn shutdown_does_not_await_a_non_cooperative_nested_tool() { + let (delegate, mut events_rx) = NeverResolvingToolDelegate::new(); + let service = CodeModeService::with_delegate(delegate); + let _started = 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); + tokio::time::timeout(Duration::from_millis(/*millis*/ 100), service.shutdown()) + .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")) + ); +} + +#[tokio::test] +async fn termination_does_not_await_a_non_cooperative_notification() { + let (delegate, mut events_rx) = NeverResolvingNotificationDelegate::new(); + let service = CodeModeService::with_delegate(delegate); let cell = service .execute(execute_request( r#"notify("pending"); await new Promise(() => {});"#, @@ -415,23 +458,14 @@ async fn repeated_termination_is_rejected_while_callback_cleanup_is_pending() { } ); - 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 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(), + tokio::time::timeout( + Duration::from_millis(/*millis*/ 100), + service.terminate(cell_id("1")), + ) + .await + .expect("termination should not await a non-cooperative notification") + .unwrap(), WaitOutcome::LiveCell(RuntimeResponse::Terminated { cell_id: cell_id("1"), content_items: Vec::new(), @@ -497,7 +531,12 @@ async fn natural_completion_cleans_up_callbacks_before_responding() { let cell = service .execute(ExecuteRequest { enabled_tools: vec![blocking_tool()], - source: r#"tools.block({}); text("done");"#.to_string(), + source: concat!( + "tools.block({});", + "await new Promise(resolve => setTimeout(resolve, 100));", + "text('done');", + ) + .to_string(), yield_time_ms: Some(60_000), ..execute_request("") }) @@ -515,13 +554,5 @@ async fn natural_completion_cleans_up_callbacks_before_responding() { 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")) - ); + assert!(delegate.tool_future_dropped.load(Ordering::Acquire)); } diff --git a/codex-rs/code-mode/src/session_runtime/mod.rs b/codex-rs/code-mode/src/session_runtime/mod.rs index 362dccda74..71cf3d6137 100644 --- a/codex-rs/code-mode/src/session_runtime/mod.rs +++ b/codex-rs/code-mode/src/session_runtime/mod.rs @@ -4,13 +4,13 @@ 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; use serde_json::Value as JsonValue; use tokio::sync::Mutex; use tokio_util::sync::CancellationToken; +use tokio_util::task::TaskTracker; pub(crate) use self::types::CellEvent; pub(crate) use self::types::CellId; @@ -43,8 +43,9 @@ pub(crate) struct SessionRuntime { struct Inner { stored_values: Mutex>, cells: Mutex>, + cell_tasks: TaskTracker, + shutdown_token: CancellationToken, delegate: Arc, - shutting_down: AtomicBool, next_cell_id: AtomicU64, } @@ -54,15 +55,16 @@ impl SessionRuntime { inner: Arc::new(Inner { stored_values: Mutex::new(HashMap::new()), cells: Mutex::new(HashMap::new()), + cell_tasks: TaskTracker::new(), + shutdown_token: CancellationToken::new(), delegate, - shutting_down: AtomicBool::new(false), next_cell_id: AtomicU64::new(1), }), } } pub(crate) fn is_alive(&self) -> bool { - !self.inner.shutting_down.load(Ordering::Acquire) + !self.inner.shutdown_token.is_cancelled() } pub(crate) async fn execute( @@ -70,6 +72,9 @@ impl SessionRuntime { request: CreateCellRequest, initial_observe_mode: ObserveMode, ) -> Result { + if self.inner.shutdown_token.is_cancelled() { + return Err(Error::ShuttingDown); + } let cell_id = self.allocate_cell_id(); let initial_event = self .start_cell(cell_id.clone(), request, initial_observe_mode) @@ -122,21 +127,13 @@ impl SessionRuntime { } pub(crate) async fn shutdown(&self) -> Result<(), Error> { - self.inner.shutting_down.store(true, Ordering::Release); - let handles = self - .inner - .cells - .lock() - .await - .values() - .cloned() - .collect::>(); - for handle in handles { - handle.shutdown(); - } - while !self.inner.cells.lock().await.is_empty() { - tokio::task::yield_now().await; - } + self.begin_shutdown(); + // Taking the registry lock ensures every cell that passed the shutdown + // check has registered its actor with the tracker before we wait. + let cells = self.inner.cells.lock().await; + self.inner.cell_tasks.close(); + drop(cells); + self.inner.cell_tasks.wait().await; Ok(()) } @@ -161,13 +158,13 @@ impl SessionRuntime { inner: Arc::clone(&self.inner), }); let mut cells = self.inner.cells.lock().await; - if self.inner.shutting_down.load(Ordering::Acquire) { + if self.inner.shutdown_token.is_cancelled() { return Err(Error::ShuttingDown); } if cells.contains_key(&cell_id) { return Err(Error::DuplicateCell(cell_id)); } - let cell_state = Arc::new(CellState::new(CancellationToken::new())); + let cell_state = Arc::new(CellState::new(self.inner.shutdown_token.child_token())); let (handle, initial_event, task) = CellActor::prepare( request, stored_values, @@ -177,18 +174,14 @@ impl SessionRuntime { ) .map_err(Error::Runtime)?; cells.insert(cell_id.clone(), handle); + self.inner.cell_tasks.spawn(task); drop(cells); - tokio::spawn(task); Ok(map_actor_event(cell_id, initial_event)) } fn begin_shutdown(&self) { - self.inner.shutting_down.store(true, Ordering::Release); - if let Ok(cells) = self.inner.cells.try_lock() { - for handle in cells.values() { - handle.shutdown(); - } - } + self.inner.shutdown_token.cancel(); + self.inner.cell_tasks.close(); } } diff --git a/codex-rs/code-mode/src/session_runtime/tests.rs b/codex-rs/code-mode/src/session_runtime/tests.rs index b4457b3f7b..a1dc18c0e8 100644 --- a/codex-rs/code-mode/src/session_runtime/tests.rs +++ b/codex-rs/code-mode/src/session_runtime/tests.rs @@ -108,3 +108,80 @@ async fn termination_rejects_a_waiting_store_commit_before_the_next_cell_can_loa ); runtime.shutdown().await.unwrap(); } + +fn execute_request(source: &str) -> CreateCellRequest { + CreateCellRequest { + tool_call_id: "call-1".to_string(), + enabled_tools: Vec::new(), + source: source.to_string(), + } +} + +#[tokio::test] +#[expect( + clippy::await_holding_invalid_type, + reason = "test holds the registry lock to force admission ahead of shutdown" +)] +async fn shutdown_rejects_cell_admission_queued_before_the_registry_lock() { + let runtime = Arc::new(SessionRuntime::new(Arc::new(RecordingDelegate))); + let cells = runtime.inner.cells.lock().await; + + let execution = runtime.execute( + execute_request("while (true) {}"), + ObserveMode::YieldAfter(Duration::from_millis(/*millis*/ 1)), + ); + tokio::pin!(execution); + std::future::poll_fn(|context| match execution.as_mut().poll(context) { + Poll::Pending => Poll::Ready(()), + Poll::Ready(Ok(_)) => panic!("execution completed before the registry lock was released"), + Poll::Ready(Err(error)) => { + panic!("execution failed before the registry lock was released: {error}") + } + }) + .await; + + let shutdown = runtime.shutdown(); + tokio::pin!(shutdown); + std::future::poll_fn(|context| match shutdown.as_mut().poll(context) { + Poll::Pending => Poll::Ready(()), + Poll::Ready(Ok(())) => panic!("shutdown completed before acquiring the registry lock"), + Poll::Ready(Err(error)) => { + panic!("shutdown failed before acquiring the registry lock: {error}") + } + }) + .await; + + assert!(!runtime.is_alive()); + drop(cells); + assert!(matches!(execution.await, Err(Error::ShuttingDown))); + assert_eq!(shutdown.await, Ok(())); +} + +#[tokio::test] +async fn drop_terminates_cells_when_the_registry_is_locked() { + let runtime = SessionRuntime::new(Arc::new(RecordingDelegate)); + let started = runtime + .execute( + execute_request("while (true) {}"), + ObserveMode::YieldAfter(Duration::from_millis(/*millis*/ 1)), + ) + .await + .unwrap(); + assert_eq!(started.cell_id, CellId::new("1")); + assert_eq!( + started.initial_event().await, + Ok(CellEvent::Yielded { + content_items: Vec::new(), + }) + ); + + let inner = Arc::clone(&runtime.inner); + let cells = inner.cells.lock().await; + drop(runtime); + drop(cells); + + tokio::time::timeout(Duration::from_secs(/*secs*/ 1), inner.cell_tasks.wait()) + .await + .unwrap(); + assert!(inner.cell_tasks.is_empty()); +}