From f8ba23d96a15e7a9cb5e53d8f056746e7f2f4f4c Mon Sep 17 00:00:00 2001 From: Channing Conger Date: Thu, 18 Jun 2026 19:48:20 +0000 Subject: [PATCH] code-mode: cancel cells when session ownership drops --- codex-rs/Cargo.lock | 1 + codex-rs/code-mode/Cargo.toml | 1 + codex-rs/code-mode/src/cell_actor/mod.rs | 16 +++++++- codex-rs/code-mode/src/cell_actor/tests.rs | 6 ++- codex-rs/code-mode/src/session_runtime/mod.rs | 28 ++++++++----- .../code-mode/src/session_runtime/tests.rs | 40 +++++++++++++++++++ .../code-mode/src/session_runtime/types.rs | 30 ++++---------- 7 files changed, 86 insertions(+), 36 deletions(-) diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index 844a84ef73..87ef121484 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -2498,6 +2498,7 @@ dependencies = [ "deno_core_icudata", "pretty_assertions", "serde_json", + "thiserror 2.0.18", "tokio", "tokio-util", "tracing", diff --git a/codex-rs/code-mode/Cargo.toml b/codex-rs/code-mode/Cargo.toml index 09f1d7380d..0bd68b153f 100644 --- a/codex-rs/code-mode/Cargo.toml +++ b/codex-rs/code-mode/Cargo.toml @@ -23,6 +23,7 @@ serde_json = { workspace = true } tokio = { workspace = true, features = ["macros", "rt", "sync", "time"] } tokio-util = { workspace = true, features = ["rt"] } tracing = { workspace = true } +thiserror = { workspace = true } v8 = { workspace = true } [dev-dependencies] diff --git a/codex-rs/code-mode/src/cell_actor/mod.rs b/codex-rs/code-mode/src/cell_actor/mod.rs index 18be3350f4..6fcac36d62 100644 --- a/codex-rs/code-mode/src/cell_actor/mod.rs +++ b/codex-rs/code-mode/src/cell_actor/mod.rs @@ -45,6 +45,7 @@ impl CellActor { stored_values: HashMap, host: Arc, initial_observe_mode: ObserveMode, + session_shutdown_token: CancellationToken, ) -> Result< ( CellHandle, @@ -71,6 +72,7 @@ impl CellActor { runtime_control_tx, runtime_terminate_handle, cancellation_token, + session_shutdown_token, }, event_rx, command_rx, @@ -90,6 +92,7 @@ struct CellContext { runtime_control_tx: std::sync::mpsc::Sender, runtime_terminate_handle: v8::IsolateHandle, cancellation_token: CancellationToken, + session_shutdown_token: CancellationToken, } struct Observer { @@ -113,6 +116,7 @@ async fn run_cell( runtime_control_tx, runtime_terminate_handle, cancellation_token, + session_shutdown_token, } = context; let mut content_items = Vec::new(); let mut pending_tool_call_ids = Vec::new(); @@ -124,14 +128,21 @@ async fn run_cell( let mut yield_timer: Option>> = None; let mut notification_tasks = JoinSet::new(); let mut tool_tasks = JoinSet::new(); - loop { let yield_deadline_elapsed = yield_timer .as_ref() .is_some_and(|yield_timer| yield_timer.deadline() <= tokio::time::Instant::now()); tokio::select! { biased; - maybe_command = command_rx.recv() => { + maybe_command = async { + tokio::select! { + biased; + command = command_rx.recv() => command, + _ = session_shutdown_token.cancelled(), if termination.is_none() => { + Some(CellCommand::Terminate { response_tx: None }) + } + } + } => { let Some(command) = maybe_command else { if completed_event.is_some() { break; @@ -384,6 +395,7 @@ async fn run_cell( } } } + // Reject requests that arrive while asynchronous terminal cleanup runs. drop(command_rx); begin_termination( &runtime_tx, diff --git a/codex-rs/code-mode/src/cell_actor/tests.rs b/codex-rs/code-mode/src/cell_actor/tests.rs index ca01c3b142..2d37d678ba 100644 --- a/codex-rs/code-mode/src/cell_actor/tests.rs +++ b/codex-rs/code-mode/src/cell_actor/tests.rs @@ -64,14 +64,16 @@ fn spawn_cell_actor_harness(initial_observe_mode: ObserveMode) -> CellActorHarne PendingRuntimeMode::PauseUntilResumed, ) .unwrap(); - let handle = CellHandle::new(command_tx, CancellationToken::new()); + let cancellation_token = CancellationToken::new(); + let handle = CellHandle::new(command_tx, cancellation_token.clone()); let task = tokio::spawn(run_cell( Arc::new(TestHost), CellContext { runtime_tx, runtime_control_tx, runtime_terminate_handle, - cancellation_token: CancellationToken::new(), + cancellation_token, + session_shutdown_token: CancellationToken::new(), }, event_rx, command_rx, diff --git a/codex-rs/code-mode/src/session_runtime/mod.rs b/codex-rs/code-mode/src/session_runtime/mod.rs index f0ca9ea8f2..893e54ef31 100644 --- a/codex-rs/code-mode/src/session_runtime/mod.rs +++ b/codex-rs/code-mode/src/session_runtime/mod.rs @@ -44,6 +44,7 @@ struct Inner { stored_values: Mutex>, cells: Mutex>, cell_count_tx: watch::Sender, + shutdown_token: CancellationToken, delegate: Arc, shutting_down: AtomicBool, next_cell_id: AtomicU64, @@ -62,6 +63,7 @@ impl SessionRuntime { stored_values: Mutex::new(HashMap::new()), cells: Mutex::new(HashMap::new()), cell_count_tx, + shutdown_token: CancellationToken::new(), delegate, shutting_down: AtomicBool::new(false), next_cell_id: AtomicU64::new(1), @@ -100,12 +102,13 @@ impl SessionRuntime { cell_id: &CellId, mode: ObserveMode, ) -> Result { - let cells = self.inner.cells.lock().await; - let Some(CellState::Live(handle)) = cells.get(cell_id) else { - return Err(Error::MissingCell(cell_id.clone())); + let handle = { + let cells = self.inner.cells.lock().await; + let Some(CellState::Live(handle)) = cells.get(cell_id) else { + return Err(Error::MissingCell(cell_id.clone())); + }; + handle.clone() }; - let handle = handle.clone(); - drop(cells); Ok(PendingEvent { event: map_actor_event(cell_id.clone(), handle.observe(mode)), }) @@ -127,7 +130,7 @@ impl SessionRuntime { pub async fn shutdown(&self) -> Result<(), Error> { let mut cell_count = self.inner.cell_count_tx.subscribe(); - self.inner.shutting_down.store(true, Ordering::Release); + self.begin_shutdown(); let handles = self .inner .cells @@ -177,9 +180,14 @@ impl SessionRuntime { if cells.contains_key(&cell_id) { return Err(Error::DuplicateCell(cell_id)); } - let (handle, initial_event, task) = - CellActor::prepare(request, stored_values, host, initial_observe_mode) - .map_err(Error::Runtime)?; + let (handle, initial_event, task) = CellActor::prepare( + request, + stored_values, + host, + initial_observe_mode, + self.inner.shutdown_token.clone(), + ) + .map_err(Error::Runtime)?; cells.insert(cell_id.clone(), CellState::Live(handle)); self.inner.cell_count_tx.send_replace(cells.len()); drop(cells); @@ -189,6 +197,8 @@ impl SessionRuntime { fn begin_shutdown(&self) { self.inner.shutting_down.store(true, Ordering::Release); + self.inner.shutdown_token.cancel(); + // The token reaches every cell if the registry is temporarily locked. if let Ok(cells) = self.inner.cells.try_lock() { for state in cells.values() { if let CellState::Live(handle) = state { diff --git a/codex-rs/code-mode/src/session_runtime/tests.rs b/codex-rs/code-mode/src/session_runtime/tests.rs index 0020565c2d..b01c2d02a9 100644 --- a/codex-rs/code-mode/src/session_runtime/tests.rs +++ b/codex-rs/code-mode/src/session_runtime/tests.rs @@ -286,3 +286,43 @@ async fn concurrent_shutdowns_wait_for_the_same_cell_cleanup() { assert_eq!(first_shutdown.await.unwrap(), Ok(())); assert_eq!(second_shutdown.await.unwrap(), Ok(())); } + +#[tokio::test] +async fn drop_terminates_cells_when_the_registry_is_locked() { + let (delegate, mut close_started_rx) = BlockingCloseDelegate::new(FirstCloseOutcome::Succeeds); + let runtime = SessionRuntime::new(Arc::clone(&delegate)); + let started = runtime + .execute( + execute_request("while (true) {}"), + ObserveMode::YieldAfter(Duration::from_millis(/*millis*/ 1)), + ) + .await + .unwrap(); + let cell_id = started.cell_id.clone(); + assert_eq!( + started.initial_event().await.unwrap(), + CellEvent::Yielded { + content_items: Vec::new(), + } + ); + + let inner = Arc::clone(&runtime.inner); + let mut cell_count = inner.cell_count_tx.subscribe(); + let cells = inner.cells.lock().await; + drop(runtime); + drop(cells); + + assert_eq!( + tokio::time::timeout(Duration::from_secs(/*secs*/ 1), close_started_rx.recv()) + .await + .unwrap(), + Some(cell_id) + ); + + delegate.close_release.add_permits(/*n*/ 1); + tokio::time::timeout(Duration::from_secs(/*secs*/ 1), cell_count.changed()) + .await + .unwrap() + .unwrap(); + assert_eq!(*cell_count.borrow_and_update(), 0); +} diff --git a/codex-rs/code-mode/src/session_runtime/types.rs b/codex-rs/code-mode/src/session_runtime/types.rs index db3eaca227..fb86f9c228 100644 --- a/codex-rs/code-mode/src/session_runtime/types.rs +++ b/codex-rs/code-mode/src/session_runtime/types.rs @@ -133,34 +133,18 @@ pub trait SessionRuntimeDelegate: Send + Sync + 'static { } /// A failure reported by a session runtime operation. -#[derive(Clone, Debug, Eq, PartialEq)] +#[derive(Clone, Debug, Eq, PartialEq, thiserror::Error)] pub enum Error { + #[error("code mode session is shutting down")] ShuttingDown, + #[error("exec cell {0} already exists")] DuplicateCell(CellId), + #[error("exec cell {0} not found")] MissingCell(CellId), + #[error("exec cell {0} already has an active observer")] BusyObserver(CellId), + #[error("exec cell {0} is already terminating")] AlreadyTerminating(CellId), + #[error("{0}")] Runtime(String), } - -impl fmt::Display for Error { - fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { - match self { - Self::ShuttingDown => formatter.write_str("code mode session is shutting down"), - Self::DuplicateCell(cell_id) => write!(formatter, "exec cell {cell_id} already exists"), - Self::MissingCell(cell_id) => write!(formatter, "exec cell {cell_id} not found"), - Self::BusyObserver(cell_id) => { - write!( - formatter, - "exec cell {cell_id} already has an active observer" - ) - } - Self::AlreadyTerminating(cell_id) => { - write!(formatter, "exec cell {cell_id} is already terminating") - } - Self::Runtime(error_text) => formatter.write_str(error_text), - } - } -} - -impl std::error::Error for Error {}