code-mode: cancel cells when session ownership drops

codex-rs/Cargo.lock

@@ -2498,6 +2498,7 @@ dependencies = [
  "deno_core_icudata",
  "pretty_assertions",
  "serde_json",
+ "thiserror 2.0.18",
  "tokio",
  "tokio-util",
  "tracing",

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]

codex-rs/code-mode/src/cell_actor/mod.rs

@@ -45,6 +45,7 @@ impl CellActor {
         stored_values: HashMap<String, JsonValue>,
         host: Arc<H>,
         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<RuntimeControlCommand>,
     runtime_terminate_handle: v8::IsolateHandle,
     cancellation_token: CancellationToken,
+    session_shutdown_token: CancellationToken,
 }
 
 struct Observer {
@@ -113,6 +116,7 @@ async fn run_cell<H: CellHost>(
         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<H: CellHost>(
     let mut yield_timer: Option<std::pin::Pin<Box<tokio::time::Sleep>>> = 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<H: CellHost>(
             }
         }
     }
+    // Reject requests that arrive while asynchronous terminal cleanup runs.
     drop(command_rx);
     begin_termination(
         &runtime_tx,

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,

codex-rs/code-mode/src/session_runtime/mod.rs

@@ -44,6 +44,7 @@ struct Inner<D: SessionRuntimeDelegate> {
     stored_values: Mutex<HashMap<String, JsonValue>>,
     cells: Mutex<HashMap<CellId, CellState>>,
     cell_count_tx: watch::Sender<usize>,
+    shutdown_token: CancellationToken,
     delegate: Arc<D>,
     shutting_down: AtomicBool,
     next_cell_id: AtomicU64,
@@ -62,6 +63,7 @@ impl<D: SessionRuntimeDelegate> SessionRuntime<D> {
                 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<D: SessionRuntimeDelegate> SessionRuntime<D> {
         cell_id: &CellId,
         mode: ObserveMode,
     ) -> Result<PendingEvent, Error> {
-        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<D: SessionRuntimeDelegate> SessionRuntime<D> {
 
     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<D: SessionRuntimeDelegate> SessionRuntime<D> {
         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<D: SessionRuntimeDelegate> SessionRuntime<D> {
 
     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 {

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);
+}

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 {}