code-mode: clean up terminal cell dispatch gates

codex-rs/code-mode-protocol/src/session.rs

@@ -61,6 +61,12 @@ pub trait CodeModeSessionDelegate: Send + Sync {
         text: String,
         cancellation_token: CancellationToken,
     ) -> NotificationFuture<'a>;
+
+    /// Reports that a cell has reached its terminal state and will issue no more callbacks.
+    ///
+    /// Implementations must keep this non-blocking. The session does not wait for an
+    /// acknowledgement or retry delivery after a transport failure.
+    fn cell_closed(&self, cell_id: &CellId);
 }
 
 /// A stateful code-mode session owned by one Codex thread.

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

@@ -291,7 +291,7 @@ async fn run_cell<H: CellHost>(
                         content_items: std::mem::take(&mut content_items),
                         error_text: Some("exec runtime ended unexpectedly".to_string()),
                     };
-                    let rejected_event = match host
+                    let (completion_committed, rejected_event) = match host
                         .commit_completion(
                             HashMap::new(),
                             event,
@@ -300,14 +300,18 @@ async fn run_cell<H: CellHost>(
                         )
                         .await
                     {
-                        CompletionCommit::Committed => None,
-                        CompletionCommit::Rejected(event) => Some(event),
+                        CompletionCommit::Committed => (true, None),
+                        CompletionCommit::Rejected(event) => (false, Some(event)),
                     };
-                    match cell_state.deliver_completion(
+                    let completion_delivery = cell_state.deliver_completion(
                         observer
                             .take()
                             .map(|observer| (observer.mode, observer.response_tx)),
-                    ) {
+                    );
+                    if completion_committed {
+                        host.terminal();
+                    }
+                    match completion_delivery {
                         CompletionDelivery::Delivered => break,
                         CompletionDelivery::Buffered => {}
                         CompletionDelivery::Rejected(response_tx) => {
@@ -438,7 +442,7 @@ async fn run_cell<H: CellHost>(
                             content_items: std::mem::take(&mut content_items),
                             error_text,
                         };
-                        let rejected_event = match host
+                        let (completion_committed, rejected_event) = match host
                             .commit_completion(
                                 stored_value_writes,
                                 event,
@@ -447,14 +451,18 @@ async fn run_cell<H: CellHost>(
                             )
                             .await
                         {
-                            CompletionCommit::Committed => None,
-                            CompletionCommit::Rejected(event) => Some(event),
+                            CompletionCommit::Committed => (true, None),
+                            CompletionCommit::Rejected(event) => (false, Some(event)),
                         };
-                        match cell_state.deliver_completion(
+                        let completion_delivery = cell_state.deliver_completion(
                             observer
                                 .take()
                                 .map(|observer| (observer.mode, observer.response_tx)),
-                        ) {
+                        );
+                        if completion_committed {
+                            host.terminal();
+                        }
+                        match completion_delivery {
                             CompletionDelivery::Delivered => break,
                             CompletionDelivery::Buffered => {}
                             CompletionDelivery::Rejected(response_tx) => {

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

@@ -52,6 +52,8 @@ impl CellHost for TestHost {
         cell_state.commit_completion(event, pending_initial_yield_items, || {})
     }
 
+    fn terminal(&self) {}
+
     async fn closed(&self) {}
 }
 
@@ -86,6 +88,8 @@ impl CellHost for RecordingHost {
         })
     }
 
+    fn terminal(&self) {}
+
     async fn closed(&self) {}
 }
 

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

@@ -60,6 +60,9 @@ pub(crate) trait CellHost: Send + Sync + 'static {
         cell_state: Arc<CellState>,
     ) -> impl Future<Output = CompletionCommit> + Send;
 
+    /// Reports that execution reached a terminal state, even when its result remains buffered.
+    fn terminal(&self);
+
     fn closed(&self) -> impl Future<Output = ()> + Send;
 }
 

codex-rs/code-mode/src/service.rs

@@ -49,6 +49,8 @@ impl CodeModeSessionDelegate for NoopCodeModeSessionDelegate {
     ) -> NotificationFuture<'a> {
         Box::pin(async { Ok(()) })
     }
+
+    fn cell_closed(&self, _cell_id: &CellId) {}
 }
 
 #[derive(Default)]
@@ -251,6 +253,10 @@ impl runtime::SessionRuntimeDelegate for ProtocolDelegate {
             )
             .await
     }
+
+    fn cell_closed(&self, cell_id: &runtime::CellId) {
+        self.delegate.cell_closed(&protocol_cell_id(cell_id));
+    }
 }
 
 fn runtime_request(request: CreateCellRequest) -> runtime::CreateCellRequest {

codex-rs/code-mode/src/service_contract_tests.rs

@@ -19,6 +19,11 @@ enum DelegateEvent {
     NotificationStarted,
     NotificationFinished,
     ToolStarted,
+    CellClosed(CellId),
+}
+
+fn record_cell_closed(events_tx: &mpsc::UnboundedSender<DelegateEvent>, cell_id: &CellId) {
+    let _ = events_tx.send(DelegateEvent::CellClosed(cell_id.clone()));
 }
 
 struct BlockingDelegate {
@@ -100,6 +105,10 @@ impl CodeModeSessionDelegate for NeverResolvingNotificationDelegate {
             std::future::pending().await
         })
     }
+
+    fn cell_closed(&self, cell_id: &CellId) {
+        record_cell_closed(&self.events_tx, cell_id);
+    }
 }
 
 impl CodeModeSessionDelegate for ReleasableNotificationDelegate {
@@ -134,6 +143,10 @@ impl CodeModeSessionDelegate for ReleasableNotificationDelegate {
             }
         })
     }
+
+    fn cell_closed(&self, cell_id: &CellId) {
+        record_cell_closed(&self.events_tx, cell_id);
+    }
 }
 
 impl CodeModeSessionDelegate for NeverResolvingToolDelegate {
@@ -157,6 +170,10 @@ impl CodeModeSessionDelegate for NeverResolvingToolDelegate {
     ) -> NotificationFuture<'a> {
         Box::pin(async { Ok(()) })
     }
+
+    fn cell_closed(&self, cell_id: &CellId) {
+        record_cell_closed(&self.events_tx, cell_id);
+    }
 }
 
 impl BlockingDelegate {
@@ -211,6 +228,10 @@ impl CodeModeSessionDelegate for BlockingDelegate {
             Err("cancelled".to_string())
         })
     }
+
+    fn cell_closed(&self, cell_id: &CellId) {
+        record_cell_closed(&self.events_tx, cell_id);
+    }
 }
 
 fn cell_id(value: &str) -> CellId {
@@ -372,6 +393,39 @@ await tools.block({});
     );
 }
 
+#[tokio::test]
+async fn background_completion_notifies_the_delegate_without_another_observation() {
+    let (delegate, mut events_rx) = BlockingDelegate::new();
+    let service = CodeModeService::with_delegate(delegate);
+    let created_cell_id = service
+        .create_cell(execute_request(
+            r#"await new Promise(resolve => setTimeout(resolve, 100)); text("done");"#,
+        ))
+        .await
+        .unwrap();
+    assert_eq!(
+        tokio::time::timeout(
+            Duration::from_secs(2),
+            service.observe(ObserveRequest {
+                cell_id: created_cell_id.clone(),
+                yield_time_ms: 1,
+            }),
+        )
+        .await
+        .expect("initial observation should yield while the cell is still running")
+        .unwrap(),
+        CellOutcome::LiveCell(RuntimeResponse::Yielded {
+            cell_id: created_cell_id.clone(),
+            content_items: Vec::new(),
+        })
+    );
+
+    assert_eq!(
+        next_event(&mut events_rx).await,
+        DelegateEvent::CellClosed(created_cell_id)
+    );
+}
+
 #[tokio::test]
 async fn returns_and_resumes_from_the_pending_frontier() {
     let (delegate, mut events_rx) = BlockingDelegate::new();

codex-rs/code-mode/src/service_tests.rs

@@ -68,6 +68,8 @@ impl CodeModeSessionDelegate for ReleasableToolDelegate {
     ) -> NotificationFuture<'a> {
         Box::pin(async { Ok(()) })
     }
+
+    fn cell_closed(&self, _cell_id: &CellId) {}
 }
 
 #[derive(Debug, PartialEq)]
@@ -82,6 +84,9 @@ enum RecordedDelegateCall {
         text: String,
         cancellation_requested: bool,
     },
+    CellClosed {
+        cell_id: CellId,
+    },
 }
 
 struct RecordingDelegate {
@@ -150,6 +155,15 @@ impl CodeModeSessionDelegate for RecordingDelegate {
                 .expect("test must provide one result per notification")
         })
     }
+
+    fn cell_closed(&self, cell_id: &CellId) {
+        self.calls
+            .lock()
+            .unwrap()
+            .push(RecordedDelegateCall::CellClosed {
+                cell_id: cell_id.clone(),
+            });
+    }
 }
 
 fn execute_request(source: &str) -> CreateCellRequest {
@@ -1204,6 +1218,7 @@ async fn protocol_delegate_maps_callbacks_cancellation_and_errors_field_for_fiel
         .await,
         Err("notification failed".to_string())
     );
+    runtime::SessionRuntimeDelegate::cell_closed(&adapter, &runtime::CellId::new("cell-d4"));
 
     assert_eq!(
         delegate.take_calls(),
@@ -1234,6 +1249,9 @@ async fn protocol_delegate_maps_callbacks_cancellation_and_errors_field_for_fiel
                 text: "progress".to_string(),
                 cancellation_requested: true,
             },
+            RecordedDelegateCall::CellClosed {
+                cell_id: cell_id("cell-d4"),
+            },
         ]
     );
 }

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

@@ -4,6 +4,7 @@ use std::collections::HashMap;
 use std::future::Future;
 use std::pin::Pin;
 use std::sync::Arc;
+use std::sync::atomic::AtomicBool;
 use std::sync::atomic::AtomicU64;
 use std::sync::atomic::Ordering;
 
@@ -177,6 +178,7 @@ impl<D: SessionRuntimeDelegate> SessionRuntime<D> {
         let host = Arc::new(RuntimeCellHost {
             cell_id: cell_id.clone(),
             inner: Arc::clone(&self.inner),
+            terminal_notified: AtomicBool::new(false),
         });
         let mut cells = self.inner.cells.lock().await;
         if self.inner.shutdown_token.is_cancelled() {
@@ -220,6 +222,7 @@ impl PendingEvent {
 struct RuntimeCellHost<D: SessionRuntimeDelegate> {
     cell_id: CellId,
     inner: Arc<Inner<D>>,
+    terminal_notified: AtomicBool,
 }
 
 impl<D: SessionRuntimeDelegate> CellHost for RuntimeCellHost<D> {
@@ -275,8 +278,15 @@ impl<D: SessionRuntimeDelegate> CellHost for RuntimeCellHost<D> {
         })
     }
 
+    fn terminal(&self) {
+        if !self.terminal_notified.swap(true, Ordering::AcqRel) {
+            self.inner.delegate.cell_closed(&self.cell_id);
+        }
+    }
+
     async fn closed(&self) {
         self.inner.cells.lock().await.remove(&self.cell_id);
+        self.terminal();
     }
 }
 

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

@@ -1,6 +1,7 @@
 use std::collections::HashMap;
 use std::future::Future;
 use std::sync::Arc;
+use std::sync::atomic::AtomicBool;
 use std::task::Context;
 use std::task::Poll;
 use std::task::Waker;
@@ -33,6 +34,8 @@ impl SessionRuntimeDelegate for RecordingDelegate {
     ) -> Result<(), String> {
         Ok(())
     }
+
+    fn cell_closed(&self, _cell_id: &CellId) {}
 }
 
 #[tokio::test]
@@ -42,6 +45,7 @@ async fn termination_rejects_a_waiting_store_commit_before_the_next_cell_can_loa
     let host = RuntimeCellHost {
         cell_id: CellId::new("terminating-writer"),
         inner: Arc::clone(&runtime.inner),
+        terminal_notified: AtomicBool::new(false),
     };
     let completion = CellEvent::Completed {
         content_items: vec![OutputItem::Text {

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

@@ -118,7 +118,8 @@ pub struct NestedToolCall {
 /// Host callbacks used by cells owned by a [`super::SessionRuntime`].
 ///
 /// Implementations should forward callback cancellation tokens to downstream
-/// work. The runtime stops awaiting callbacks once cancellation begins.
+/// work. The runtime stops awaiting callbacks once cancellation begins. Terminal
+/// notifications must be non-blocking and are not acknowledged or retried.
 pub trait SessionRuntimeDelegate: Send + Sync + 'static {
     fn invoke_tool(
         &self,
@@ -133,6 +134,11 @@ pub trait SessionRuntimeDelegate: Send + Sync + 'static {
         text: String,
         cancellation_token: CancellationToken,
     ) -> impl Future<Output = Result<(), String>> + Send;
+
+    /// Reports that execution is terminal and will issue no more callbacks.
+    ///
+    /// The terminal result may remain buffered for a later observation.
+    fn cell_closed(&self, cell_id: &CellId);
 }
 
 /// A failure reported by a session runtime operation.

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

@@ -238,6 +238,10 @@ impl CodeModeSessionDelegate for CodeModeDispatchBroker {
             }
         })
     }
+
+    fn cell_closed(&self, cell_id: &CellId) {
+        self.close_cell(cell_id);
+    }
 }
 
 enum DispatchMessage {
@@ -307,3 +311,7 @@ impl CoreTurnHost {
             })
     }
 }
+
+#[cfg(test)]
+#[path = "delegate_tests.rs"]
+mod tests;

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

@@ -0,0 +1,66 @@
+use std::sync::Arc;
+use std::time::Duration;
+
+use codex_code_mode::CellId;
+use codex_code_mode::CellOutcome;
+use codex_code_mode::CodeModeService;
+use codex_code_mode::CodeModeSessionDelegate;
+use codex_code_mode::CreateCellRequest;
+use codex_code_mode::ObserveRequest;
+use codex_code_mode::RuntimeResponse;
+
+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));
+}
+
+#[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 {
+            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 {
+                cell_id: cell_id.clone(),
+                yield_time_ms: 1,
+            })
+            .await
+            .unwrap(),
+        CellOutcome::LiveCell(RuntimeResponse::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();
+}