code-mode: block requests for closing cells

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

@@ -384,6 +384,7 @@ async fn run_cell<H: CellHost>(
             }
         }
     }
+    drop(command_rx);
     begin_termination(
         &runtime_tx,
         &runtime_control_tx,

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

@@ -41,12 +41,17 @@ pub struct SessionRuntime<D: SessionRuntimeDelegate> {
 
 struct Inner<D: SessionRuntimeDelegate> {
     stored_values: Mutex<HashMap<String, JsonValue>>,
-    cells: Mutex<HashMap<CellId, CellHandle>>,
+    cells: Mutex<HashMap<CellId, CellState>>,
     delegate: Arc<D>,
     shutting_down: AtomicBool,
     next_cell_id: AtomicU64,
 }
 
+enum CellState {
+    Live(CellHandle),
+    Closing,
+}
+
 impl<D: SessionRuntimeDelegate> SessionRuntime<D> {
     pub fn new(delegate: Arc<D>) -> Self {
         Self {
@@ -91,28 +96,25 @@ impl<D: SessionRuntimeDelegate> SessionRuntime<D> {
         cell_id: &CellId,
         mode: ObserveMode,
     ) -> Result<PendingEvent, Error> {
-        let handle = self
-            .inner
-            .cells
-            .lock()
-            .await
-            .get(cell_id)
-            .cloned()
-            .ok_or_else(|| Error::MissingCell(cell_id.clone()))?;
+        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 = handle.clone();
+        drop(cells);
         Ok(PendingEvent {
             event: map_actor_event(cell_id.clone(), handle.observe(mode)),
         })
     }
 
     pub async fn terminate(&self, cell_id: &CellId) -> Result<CellEvent, Error> {
-        let handle = self
-            .inner
-            .cells
-            .lock()
-            .await
-            .get(cell_id)
-            .cloned()
-            .ok_or_else(|| 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()
+        };
         handle
             .terminate()
             .await
@@ -127,7 +129,10 @@ impl<D: SessionRuntimeDelegate> SessionRuntime<D> {
             .lock()
             .await
             .values()
-            .cloned()
+            .filter_map(|state| match state {
+                CellState::Live(handle) => Some(handle.clone()),
+                CellState::Closing => None,
+            })
             .collect::<Vec<_>>();
         for handle in handles {
             handle.shutdown();
@@ -168,7 +173,7 @@ impl<D: SessionRuntimeDelegate> SessionRuntime<D> {
         let (handle, initial_event, task) =
             CellActor::prepare(request, stored_values, host, initial_observe_mode)
                 .map_err(Error::Runtime)?;
-        cells.insert(cell_id.clone(), handle);
+        cells.insert(cell_id.clone(), CellState::Live(handle));
         drop(cells);
         tokio::spawn(task);
         Ok(map_actor_event(cell_id, initial_event))
@@ -177,8 +182,10 @@ impl<D: SessionRuntimeDelegate> SessionRuntime<D> {
     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();
+            for state in cells.values() {
+                if let CellState::Live(handle) = state {
+                    handle.shutdown();
+                }
             }
         }
     }
@@ -260,6 +267,11 @@ impl<D: SessionRuntimeDelegate> CellHost for RuntimeCellHost<D> {
     }
 
     async fn closed(&self) {
+        self.inner
+            .cells
+            .lock()
+            .await
+            .insert(self.cell_id.clone(), CellState::Closing);
         if let Err(err) = self.inner.delegate.cell_closed(&self.cell_id).await {
             warn!(
                 "failed to close code mode delegate state for cell {}: {err}",
@@ -281,3 +293,7 @@ fn actor_error(cell_id: &CellId, error: CellError) -> Error {
         CellError::Closed => Error::MissingCell(cell_id.clone()),
     }
 }
+
+#[cfg(test)]
+#[path = "tests.rs"]
+mod tests;

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

@@ -0,0 +1,106 @@
+use std::sync::Arc;
+use std::time::Duration;
+
+use pretty_assertions::assert_eq;
+use serde_json::Value as JsonValue;
+use tokio::sync::Notify;
+use tokio_util::sync::CancellationToken;
+
+use super::*;
+
+struct BlockingCloseDelegate {
+    close_started: Notify,
+    close_release: Notify,
+}
+
+impl BlockingCloseDelegate {
+    fn release_close(&self) {
+        self.close_release.notify_one();
+    }
+}
+
+impl SessionRuntimeDelegate for BlockingCloseDelegate {
+    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> {
+        Ok(())
+    }
+
+    async fn cell_closed(&self, _cell_id: &CellId) -> Result<(), String> {
+        self.close_started.notify_one();
+        self.close_release.notified().await;
+        Ok(())
+    }
+}
+
+fn execute_request() -> ExecuteRequest {
+    ExecuteRequest {
+        tool_call_id: "call-1".to_string(),
+        enabled_tools: Vec::new(),
+        source: r#"text("done");"#.to_string(),
+    }
+}
+
+#[tokio::test]
+async fn closing_cells_reject_requests_while_delegate_cleanup_runs() {
+    let delegate = Arc::new(BlockingCloseDelegate {
+        close_started: Notify::new(),
+        close_release: Notify::new(),
+    });
+    let runtime = Arc::new(SessionRuntime::new(Arc::clone(&delegate)));
+    let close_started = delegate.close_started.notified();
+    let started = runtime
+        .execute(
+            execute_request(),
+            ObserveMode::YieldAfter(Duration::from_secs(/*secs*/ 60)),
+        )
+        .await
+        .unwrap();
+    let cell_id = started.cell_id.clone();
+
+    assert_eq!(
+        started.initial_event().await.unwrap(),
+        CellEvent::Completed {
+            content_items: vec![OutputItem::Text {
+                text: "done".to_string(),
+            }],
+            error_text: None,
+        }
+    );
+    close_started.await;
+
+    assert_eq!(
+        runtime
+            .observe(&cell_id, ObserveMode::PendingFrontier)
+            .await,
+        Err(Error::MissingCell(cell_id.clone()))
+    );
+    assert_eq!(
+        runtime.terminate(&cell_id).await,
+        Err(Error::MissingCell(cell_id.clone()))
+    );
+
+    let shutdown_runtime = Arc::clone(&runtime);
+    let mut shutdown = tokio::spawn(async move { shutdown_runtime.shutdown().await });
+    assert!(
+        tokio::time::timeout(Duration::from_millis(/*millis*/ 100), &mut shutdown)
+            .await
+            .is_err()
+    );
+
+    delegate.release_close();
+
+    assert_eq!(shutdown.await.unwrap(), Ok(()));
+}