code-mode: clean up terminal cell dispatch gates

codex-rs/core/src/tools/code_mode/delegate.rs

@@ -1,4 +1,5 @@
 use std::collections::HashMap;
+use std::collections::hash_map::Entry;
 use std::sync::Arc;
 use std::sync::Mutex;
 
@@ -24,7 +25,12 @@ use crate::tools::parallel::ToolCallRuntime;
 pub(super) struct CodeModeDispatchBroker {
     dispatch_tx: async_channel::Sender<DispatchMessage>,
     dispatch_rx: async_channel::Receiver<DispatchMessage>,
-    dispatch_gates: Arc<Mutex<HashMap<CellId, watch::Sender<bool>>>>,
+    dispatch_gates: Arc<Mutex<HashMap<CellId, DispatchGate>>>,
+}
+
+enum DispatchGate {
+    Waiting(watch::Sender<bool>),
+    ClosedBeforeReady,
 }
 
 impl CodeModeDispatchBroker {
@@ -38,11 +44,35 @@ impl CodeModeDispatchBroker {
     }
 
     pub(super) fn mark_cell_ready_for_dispatch(&self, cell_id: &CellId) {
-        dispatch_gate(&self.dispatch_gates, cell_id).send_replace(true);
+        let mut dispatch_gates = match self.dispatch_gates.lock() {
+            Ok(dispatch_gates) => dispatch_gates,
+            Err(poisoned) => poisoned.into_inner(),
+        };
+        match dispatch_gates.entry(cell_id.clone()) {
+            Entry::Occupied(mut entry) => match entry.get_mut() {
+                DispatchGate::Waiting(ready_tx) => {
+                    ready_tx.send_replace(true);
+                }
+                DispatchGate::ClosedBeforeReady => {
+                    entry.remove();
+                }
+            },
+            Entry::Vacant(entry) => {
+                entry.insert(DispatchGate::Waiting(watch::channel(true).0));
+            }
+        }
     }
 
     pub(super) fn close_cell(&self, cell_id: &CellId) {
-        remove_dispatch_gate(&self.dispatch_gates, cell_id);
+        close_dispatch_gate(&self.dispatch_gates, cell_id);
+    }
+
+    pub(super) fn forget_cell(&self, cell_id: &CellId) {
+        let mut dispatch_gates = match self.dispatch_gates.lock() {
+            Ok(dispatch_gates) => dispatch_gates,
+            Err(poisoned) => poisoned.into_inner(),
+        };
+        dispatch_gates.remove(cell_id);
     }
 
     pub(super) fn start_turn_worker(
@@ -87,7 +117,7 @@ impl CodeModeDispatchBroker {
                         {
                             host.notify(call_id, cell_id, text).await
                         } else {
-                            remove_dispatch_gate(&dispatch_gates, &cell_id);
+                            close_dispatch_gate(&dispatch_gates, &cell_id);
                             Err("code mode notification cancelled".to_string())
                         };
                         let _ = response_tx.send(response);
@@ -105,7 +135,7 @@ impl CodeModeDispatchBroker {
                         )
                         .await
                         {
-                            remove_dispatch_gate(&dispatch_gates, &cell_id);
+                            close_dispatch_gate(&dispatch_gates, &cell_id);
                             continue;
                         }
                         let host = Arc::clone(&host);
@@ -129,40 +159,53 @@ impl CodeModeDispatchBroker {
     }
 }
 
-fn dispatch_gate(
-    dispatch_gates: &Mutex<HashMap<CellId, watch::Sender<bool>>>,
-    cell_id: &CellId,
-) -> watch::Sender<bool> {
+fn close_dispatch_gate(dispatch_gates: &Mutex<HashMap<CellId, DispatchGate>>, cell_id: &CellId) {
     let mut dispatch_gates = match dispatch_gates.lock() {
         Ok(dispatch_gates) => dispatch_gates,
         Err(poisoned) => poisoned.into_inner(),
     };
-    dispatch_gates
-        .entry(cell_id.clone())
-        .or_insert_with(|| watch::channel(false).0)
-        .clone()
-}
-
-fn remove_dispatch_gate(
-    dispatch_gates: &Mutex<HashMap<CellId, watch::Sender<bool>>>,
-    cell_id: &CellId,
-) {
-    let mut dispatch_gates = match dispatch_gates.lock() {
-        Ok(dispatch_gates) => dispatch_gates,
-        Err(poisoned) => poisoned.into_inner(),
-    };
-    dispatch_gates.remove(cell_id);
+    match dispatch_gates.entry(cell_id.clone()) {
+        Entry::Occupied(mut entry) => match entry.get_mut() {
+            DispatchGate::Waiting(ready_tx) if *ready_tx.borrow() => {
+                entry.remove();
+            }
+            DispatchGate::Waiting(_) => {
+                entry.insert(DispatchGate::ClosedBeforeReady);
+            }
+            DispatchGate::ClosedBeforeReady => {}
+        },
+        Entry::Vacant(entry) => {
+            entry.insert(DispatchGate::ClosedBeforeReady);
+        }
+    }
 }
 
 async fn wait_until_cell_ready_for_dispatch(
-    dispatch_gates: &Mutex<HashMap<CellId, watch::Sender<bool>>>,
+    dispatch_gates: &Mutex<HashMap<CellId, DispatchGate>>,
     cell_id: &CellId,
     cancellation_token: &CancellationToken,
 ) -> bool {
     if cancellation_token.is_cancelled() {
         return false;
     }
-    let mut ready_rx = dispatch_gate(dispatch_gates, cell_id).subscribe();
+    let mut ready_rx = {
+        let mut dispatch_gates = match dispatch_gates.lock() {
+            Ok(dispatch_gates) => dispatch_gates,
+            Err(poisoned) => poisoned.into_inner(),
+        };
+        match dispatch_gates.entry(cell_id.clone()) {
+            Entry::Occupied(entry) => match entry.get() {
+                DispatchGate::Waiting(ready_tx) => ready_tx.subscribe(),
+                DispatchGate::ClosedBeforeReady => return false,
+            },
+            Entry::Vacant(entry) => {
+                let ready_tx = watch::channel(false).0;
+                let ready_rx = ready_tx.subscribe();
+                entry.insert(DispatchGate::Waiting(ready_tx));
+                ready_rx
+            }
+        }
+    };
     loop {
         if *ready_rx.borrow_and_update() {
             return true;
@@ -238,6 +281,10 @@ impl CodeModeSessionDelegate for CodeModeDispatchBroker {
             }
         })
     }
+
+    fn cell_closed(&self, cell_id: &CellId) {
+        self.close_cell(cell_id);
+    }
 }
 
 enum DispatchMessage {
@@ -307,3 +354,7 @@ impl CoreTurnHost {
             })
     }
 }
+
+#[cfg(test)]
+#[path = "delegate_tests.rs"]
+mod tests;

codex-rs/core/src/tools/code_mode/delegate_tests.rs

@@ -0,0 +1,83 @@
+use std::sync::Arc;
+use std::time::Duration;
+
+use codex_code_mode::CellId;
+use codex_code_mode::CodeModeService;
+use codex_code_mode::CodeModeSessionDelegate;
+use codex_code_mode::CreateCellRequest;
+use codex_code_mode::ObserveOutcome;
+use codex_code_mode::ObserveRequest;
+
+use super::CodeModeDispatchBroker;
+
+#[test]
+fn terminal_notification_removes_the_cell_dispatch_gate() {
+    let broker = CodeModeDispatchBroker::new();
+    let cell_id = CellId::new("cell-a7".to_string());
+    broker.mark_cell_ready_for_dispatch(&cell_id);
+    assert!(broker.dispatch_gates.lock().unwrap().contains_key(&cell_id));
+
+    CodeModeSessionDelegate::cell_closed(&broker, &cell_id);
+
+    assert!(!broker.dispatch_gates.lock().unwrap().contains_key(&cell_id));
+}
+
+#[test]
+fn terminal_notification_before_readiness_does_not_reopen_the_dispatch_gate() {
+    let broker = CodeModeDispatchBroker::new();
+    let cell_id = CellId::new("cell-a7".to_string());
+
+    CodeModeSessionDelegate::cell_closed(&broker, &cell_id);
+    assert!(matches!(
+        broker.dispatch_gates.lock().unwrap().get(&cell_id),
+        Some(super::DispatchGate::ClosedBeforeReady)
+    ));
+
+    broker.mark_cell_ready_for_dispatch(&cell_id);
+
+    assert!(!broker.dispatch_gates.lock().unwrap().contains_key(&cell_id));
+}
+
+#[tokio::test]
+async fn background_completion_removes_the_dispatch_gate_without_another_observation() {
+    let broker = Arc::new(CodeModeDispatchBroker::new());
+    let service = CodeModeService::with_delegate(broker.clone());
+    let cell_id = service
+        .create_cell(CreateCellRequest {
+            idempotency_key: "delegate-cleanup-cell".to_string(),
+            tool_call_id: "call-1".to_string(),
+            enabled_tools: Vec::new(),
+            source: concat!(
+                "await new Promise(resolve => setTimeout(resolve, 100));",
+                "text('done');",
+            )
+            .to_string(),
+        })
+        .await
+        .unwrap();
+    broker.mark_cell_ready_for_dispatch(&cell_id);
+
+    assert_eq!(
+        service
+            .observe(ObserveRequest {
+                idempotency_key: "delegate-cleanup-observation".to_string(),
+                cell_id: cell_id.clone(),
+                yield_time_ms: 1,
+            })
+            .await
+            .unwrap(),
+        ObserveOutcome::Yielded {
+            cell_id: cell_id.clone(),
+            content_items: Vec::new(),
+        }
+    );
+    tokio::time::timeout(Duration::from_secs(/*secs*/ 2), async {
+        while broker.dispatch_gates.lock().unwrap().contains_key(&cell_id) {
+            tokio::task::yield_now().await;
+        }
+    })
+    .await
+    .expect("terminal notification should remove the dispatch gate");
+
+    service.shutdown().await.unwrap();
+}

codex-rs/core/src/tools/code_mode/mod.rs

@@ -164,7 +164,7 @@ impl CodeModeService {
     }
 
     pub(crate) fn finish_cell_dispatch(&self, cell_id: &CellId) {
-        self.dispatch_broker.close_cell(cell_id);
+        self.dispatch_broker.forget_cell(cell_id);
     }
 
     pub(crate) fn start_turn_worker(