code-mode: make session shutdown cancellation hierarchical

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

@@ -63,8 +63,8 @@ impl CellActor {
             event_tx,
             PendingRuntimeMode::PauseUntilResumed,
         )?;
-        let cancellation_token = CancellationToken::new();
-        let handle = CellHandle::new(command_tx, cancellation_token.clone());
+        let cancellation_token = session_shutdown_token.child_token();
+        let handle = CellHandle::new(command_tx);
         let task = run_cell(
             host,
             CellContext {

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

@@ -69,9 +69,9 @@ fn spawn_cell_actor_harness(initial_observe_mode: ObserveMode) -> CellActorHarne
         PendingRuntimeMode::PauseUntilResumed,
     )
     .unwrap();
-    let cancellation_token = CancellationToken::new();
     let session_shutdown_token = CancellationToken::new();
-    let handle = CellHandle::new(command_tx.clone(), cancellation_token.clone());
+    let cancellation_token = CancellationToken::new();
+    let handle = CellHandle::new(command_tx.clone());
     let task = tokio::spawn(run_cell(
         Arc::new(TestHost),
         CellContext {

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

@@ -61,18 +61,13 @@ pub(crate) trait CellHost: Send + Sync + 'static {
 #[derive(Clone)]
 pub(crate) struct CellHandle {
     command_tx: mpsc::UnboundedSender<CellCommand>,
-    cancellation_token: CancellationToken,
     termination_requested: Arc<AtomicBool>,
 }
 
 impl CellHandle {
-    pub(super) fn new(
-        command_tx: mpsc::UnboundedSender<CellCommand>,
-        cancellation_token: CancellationToken,
-    ) -> Self {
+    pub(super) fn new(command_tx: mpsc::UnboundedSender<CellCommand>) -> Self {
         Self {
             command_tx,
-            cancellation_token,
             termination_requested: Arc::new(AtomicBool::new(false)),
         }
     }
@@ -110,14 +105,6 @@ impl CellHandle {
         }
         response_event(response_rx)
     }
-
-    pub(crate) fn shutdown(&self) {
-        self.termination_requested.store(true, Ordering::Relaxed);
-        self.cancellation_token.cancel();
-        let _ = self
-            .command_tx
-            .send(CellCommand::Terminate { response_tx: None });
-    }
 }
 
 pub(super) enum CellCommand {

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

@@ -128,20 +128,6 @@ impl<D: SessionRuntimeDelegate> SessionRuntime<D> {
     pub async fn shutdown(&self) -> Result<(), Error> {
         let mut cell_count = self.inner.cell_count_tx.subscribe();
         self.begin_shutdown();
-        let handles = self
-            .inner
-            .cells
-            .lock()
-            .await
-            .values()
-            .filter_map(|state| match state {
-                CellState::Live(handle) => Some(handle.clone()),
-                CellState::Closing => None,
-            })
-            .collect::<Vec<_>>();
-        for handle in handles {
-            handle.shutdown();
-        }
         while *cell_count.borrow_and_update() != 0 {
             if cell_count.changed().await.is_err() {
                 break;
@@ -194,14 +180,6 @@ impl<D: SessionRuntimeDelegate> SessionRuntime<D> {
 
     fn begin_shutdown(&self) {
         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 {
-                    handle.shutdown();
-                }
-            }
-        }
     }
 }
 

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

@@ -16,12 +16,23 @@ enum FirstCloseOutcome {
     Fails,
 }
 
+#[derive(Debug, PartialEq)]
+enum NotificationEvent {
+    Started,
+    Cancelled,
+    Closed(CellId),
+}
+
 struct BlockingCloseDelegate {
     close_started_tx: mpsc::UnboundedSender<CellId>,
     close_release: Semaphore,
     fail_first_close: AtomicBool,
 }
 
+struct BlockingNotificationDelegate {
+    events_tx: mpsc::UnboundedSender<NotificationEvent>,
+}
+
 impl BlockingCloseDelegate {
     fn new(first_close_outcome: FirstCloseOutcome) -> (Arc<Self>, mpsc::UnboundedReceiver<CellId>) {
         let (close_started_tx, close_started_rx) = mpsc::unbounded_channel();
@@ -39,6 +50,13 @@ impl BlockingCloseDelegate {
     }
 }
 
+impl BlockingNotificationDelegate {
+    fn new() -> (Arc<Self>, mpsc::UnboundedReceiver<NotificationEvent>) {
+        let (events_tx, events_rx) = mpsc::unbounded_channel();
+        (Arc::new(Self { events_tx }), events_rx)
+    }
+}
+
 impl SessionRuntimeDelegate for BlockingCloseDelegate {
     async fn invoke_tool(
         &self,
@@ -74,6 +92,38 @@ impl SessionRuntimeDelegate for BlockingCloseDelegate {
     }
 }
 
+impl SessionRuntimeDelegate for BlockingNotificationDelegate {
+    async fn invoke_tool(
+        &self,
+        _invocation: NestedToolCall,
+        _cancellation_token: CancellationToken,
+    ) -> Result<JsonValue, String> {
+        Ok(JsonValue::Null)
+    }
+
+    async fn notify(
+        &self,
+        _call_id: String,
+        _cell_id: CellId,
+        _text: String,
+        cancellation_token: CancellationToken,
+    ) -> Result<(), String> {
+        self.events_tx
+            .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())
+    }
+
+    async fn cell_closed(&self, cell_id: &CellId) -> Result<(), String> {
+        self.events_tx
+            .send(NotificationEvent::Closed(cell_id.clone()))
+            .map_err(|_| "test did not receive cell close".to_string())
+    }
+}
+
 fn execute_request(source: &str) -> ExecuteRequest {
     ExecuteRequest {
         tool_call_id: "call-1".to_string(),
@@ -326,3 +376,48 @@ async fn drop_terminates_cells_when_the_registry_is_locked() {
         .unwrap();
     assert_eq!(*cell_count.borrow_and_update(), 0);
 }
+
+#[tokio::test]
+async fn drop_cancels_notifications_during_natural_completion_when_registry_is_locked() {
+    let (delegate, mut events_rx) = BlockingNotificationDelegate::new();
+    let runtime = SessionRuntime::new(Arc::clone(&delegate));
+    let started = runtime
+        .execute(
+            execute_request(r#"notify("pending");"#),
+            ObserveMode::YieldAfter(Duration::from_secs(/*secs*/ 60)),
+        )
+        .await
+        .unwrap();
+    let cell_id = started.cell_id;
+
+    assert_eq!(
+        tokio::time::timeout(Duration::from_secs(/*secs*/ 1), events_rx.recv())
+            .await
+            .unwrap(),
+        Some(NotificationEvent::Started)
+    );
+
+    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), events_rx.recv())
+            .await
+            .unwrap(),
+        Some(NotificationEvent::Cancelled)
+    );
+    assert_eq!(
+        tokio::time::timeout(Duration::from_secs(/*secs*/ 1), events_rx.recv())
+            .await
+            .unwrap(),
+        Some(NotificationEvent::Closed(cell_id))
+    );
+    tokio::time::timeout(Duration::from_secs(/*secs*/ 1), cell_count.changed())
+        .await
+        .unwrap()
+        .unwrap();
+    assert_eq!(*cell_count.borrow_and_update(), 0);
+}