code-mode: extend test coverage to lock in cell lifecycle (#28468)

This PR establishes the intended behavior as an executable contract before a refactor of the cell runtime begins. It also fixes cases where a second observer or termination request could replace an existing response channel and leave the original caller unresolved. ### Behavior codified - A cell can yield output and subsequently resume to completion. - A caller can run a cell until it has no immediately runnable work, receive its accumulated output and outstanding tool-call IDs, and then resume the same cell when the awaited work is available. - Each cell admits one active observer: - a second observer receives an explicit busy error - the existing observer remains registered and is not displaced - A natural result (conclusion of the js module) that has already reached the cell controller wins over a later termination request. - Otherwise, termination preempts execution and resolves both: - the active observer, if present - the caller requesting termination - Repeated termination requests are rejected while termination is already in progress. - Terminal responses are sent only after outstanding callback work has been handled: - natural completion drains notifications and cancels outstanding tool calls - termination cancels and drains both notification and tool callbacks. - Cell removal and cell_closed notification happen after callback cleanup
2026-07-01 00:31:56 +08:00 · 2026-06-16 13:34:16 -07:00
parent 4b7351700f
commit e93516e259
3 changed files with 903 additions and 137 deletions
@@ -0,0 +1,634 @@
+use std::collections::HashMap;
+use std::sync::Arc;
+use std::sync::atomic::AtomicBool;
+use std::sync::atomic::AtomicU64;
+use std::sync::atomic::Ordering;
+use std::time::Duration;
+
+use codex_protocol::ToolName;
+use pretty_assertions::assert_eq;
+use tokio::sync::Notify;
+use tokio::sync::mpsc;
+use tokio::sync::oneshot;
+use tokio_util::sync::CancellationToken;
+
+use super::*;
+use crate::CodeModeToolKind;
+use crate::ToolDefinition;
+
+#[derive(Debug, PartialEq)]
+enum DelegateEvent {
+    NotificationStarted,
+    NotificationCancelled,
+    ToolStarted,
+    ToolCancelled,
+    CellClosed(CellId),
+}
+
+struct BlockingDelegate {
+    events_tx: mpsc::UnboundedSender<DelegateEvent>,
+    notification_finished: AtomicBool,
+    tool_finished: AtomicBool,
+}
+
+struct HeldNotificationDelegate {
+    events_tx: mpsc::UnboundedSender<DelegateEvent>,
+    notification_release: Notify,
+}
+
+impl HeldNotificationDelegate {
+    fn new() -> (Arc<Self>, mpsc::UnboundedReceiver<DelegateEvent>) {
+        let (events_tx, events_rx) = mpsc::unbounded_channel();
+        (
+            Arc::new(Self {
+                events_tx,
+                notification_release: Notify::new(),
+            }),
+            events_rx,
+        )
+    }
+
+    fn release_notification(&self) {
+        self.notification_release.notify_one();
+    }
+}
+
+impl CodeModeSessionDelegate for HeldNotificationDelegate {
+    fn invoke_tool<'a>(
+        &'a self,
+        _invocation: CodeModeNestedToolCall,
+        cancellation_token: CancellationToken,
+    ) -> ToolInvocationFuture<'a> {
+        Box::pin(async move {
+            cancellation_token.cancelled().await;
+            Err("cancelled".to_string())
+        })
+    }
+
+    fn notify<'a>(
+        &'a self,
+        _call_id: String,
+        _cell_id: CellId,
+        _text: String,
+        cancellation_token: CancellationToken,
+    ) -> NotificationFuture<'a> {
+        Box::pin(async move {
+            let _ = self.events_tx.send(DelegateEvent::NotificationStarted);
+            cancellation_token.cancelled().await;
+            let _ = self.events_tx.send(DelegateEvent::NotificationCancelled);
+            self.notification_release.notified().await;
+            Ok(())
+        })
+    }
+
+    fn cell_closed(&self, cell_id: &CellId) {
+        let _ = self
+            .events_tx
+            .send(DelegateEvent::CellClosed(cell_id.clone()));
+    }
+}
+
+struct CellControlHarness {
+    event_tx: mpsc::UnboundedSender<RuntimeEvent>,
+    control_tx: mpsc::UnboundedSender<CellControlCommand>,
+    initial_response_rx: oneshot::Receiver<Result<RuntimeResponse, String>>,
+    task: tokio::task::JoinHandle<()>,
+    _runtime_event_rx: mpsc::UnboundedReceiver<RuntimeEvent>,
+}
+
+fn spawn_cell_control_harness(
+    initial_yield_time_ms: Option<u64>,
+    delegate: Arc<dyn CodeModeSessionDelegate>,
+) -> CellControlHarness {
+    let (event_tx, event_rx) = mpsc::unbounded_channel();
+    let (control_tx, control_rx) = mpsc::unbounded_channel();
+    let (initial_response_tx, initial_response_rx) = oneshot::channel();
+    let (runtime_event_tx, runtime_event_rx) = mpsc::unbounded_channel();
+    let (runtime_tx, runtime_control_tx, runtime_terminate_handle) = spawn_runtime(
+        HashMap::new(),
+        execute_request("await new Promise(() => {});"),
+        runtime_event_tx,
+        PendingRuntimeMode::Continue,
+    )
+    .unwrap();
+    let inner = Arc::new(Inner {
+        stored_values: Mutex::new(HashMap::new()),
+        cells: Mutex::new(HashMap::new()),
+        delegate,
+        shutting_down: AtomicBool::new(false),
+        next_cell_id: AtomicU64::new(1),
+    });
+    let task = tokio::spawn(run_cell_control(
+        inner,
+        CellControlContext {
+            cell_id: cell_id("1"),
+            runtime_tx,
+            runtime_control_tx,
+            pending_mode: PendingRuntimeMode::Continue,
+            runtime_terminate_handle,
+            cancellation_token: CancellationToken::new(),
+        },
+        event_rx,
+        control_rx,
+        CellResponseSender::Runtime(initial_response_tx),
+        initial_yield_time_ms,
+    ));
+
+    CellControlHarness {
+        event_tx,
+        control_tx,
+        initial_response_rx,
+        task,
+        _runtime_event_rx: runtime_event_rx,
+    }
+}
+
+impl BlockingDelegate {
+    fn new() -> (Arc<Self>, mpsc::UnboundedReceiver<DelegateEvent>) {
+        let (events_tx, events_rx) = mpsc::unbounded_channel();
+        (
+            Arc::new(Self {
+                events_tx,
+                notification_finished: AtomicBool::new(false),
+                tool_finished: AtomicBool::new(false),
+            }),
+            events_rx,
+        )
+    }
+}
+
+impl CodeModeSessionDelegate for BlockingDelegate {
+    fn invoke_tool<'a>(
+        &'a self,
+        _invocation: CodeModeNestedToolCall,
+        cancellation_token: CancellationToken,
+    ) -> ToolInvocationFuture<'a> {
+        Box::pin(async move {
+            let _ = self.events_tx.send(DelegateEvent::ToolStarted);
+            cancellation_token.cancelled().await;
+            self.tool_finished.store(true, Ordering::Release);
+            let _ = self.events_tx.send(DelegateEvent::ToolCancelled);
+            Err("cancelled".to_string())
+        })
+    }
+
+    fn notify<'a>(
+        &'a self,
+        _call_id: String,
+        _cell_id: CellId,
+        _text: String,
+        cancellation_token: CancellationToken,
+    ) -> NotificationFuture<'a> {
+        Box::pin(async move {
+            let _ = self.events_tx.send(DelegateEvent::NotificationStarted);
+            cancellation_token.cancelled().await;
+            self.notification_finished.store(true, Ordering::Release);
+            let _ = self.events_tx.send(DelegateEvent::NotificationCancelled);
+            Err("cancelled".to_string())
+        })
+    }
+
+    fn cell_closed(&self, cell_id: &CellId) {
+        let _ = self
+            .events_tx
+            .send(DelegateEvent::CellClosed(cell_id.clone()));
+    }
+}
+
+fn cell_id(value: &str) -> CellId {
+    CellId::new(value.to_string())
+}
+
+fn execute_request(source: &str) -> ExecuteRequest {
+    ExecuteRequest {
+        tool_call_id: "call-1".to_string(),
+        enabled_tools: Vec::new(),
+        source: source.to_string(),
+        yield_time_ms: Some(1),
+        max_output_tokens: None,
+    }
+}
+
+fn blocking_tool() -> ToolDefinition {
+    ToolDefinition {
+        name: "block".to_string(),
+        tool_name: ToolName::plain("block"),
+        description: String::new(),
+        kind: CodeModeToolKind::Function,
+        input_schema: None,
+        output_schema: None,
+    }
+}
+
+async fn next_event(events_rx: &mut mpsc::UnboundedReceiver<DelegateEvent>) -> DelegateEvent {
+    tokio::time::timeout(Duration::from_secs(2), events_rx.recv())
+        .await
+        .expect("delegate event timeout")
+        .expect("delegate event channel closed")
+}
+
+#[tokio::test]
+async fn yield_timer_preempts_buffered_runtime_output() {
+    let harness = spawn_cell_control_harness(
+        Some(/*initial_yield_time_ms*/ 0),
+        Arc::new(NoopCodeModeSessionDelegate),
+    );
+    harness.event_tx.send(RuntimeEvent::Started).unwrap();
+    harness
+        .event_tx
+        .send(RuntimeEvent::ContentItem(
+            FunctionCallOutputContentItem::InputText {
+                text: "queued output".to_string(),
+            },
+        ))
+        .unwrap();
+
+    assert_eq!(
+        harness.initial_response_rx.await.unwrap(),
+        Ok(RuntimeResponse::Yielded {
+            cell_id: cell_id("1"),
+            content_items: Vec::new(),
+        })
+    );
+
+    let (termination_tx, termination_rx) = oneshot::channel();
+    harness
+        .control_tx
+        .send(CellControlCommand::Terminate {
+            response_tx: termination_tx,
+        })
+        .unwrap();
+    drop(harness.event_tx);
+    assert_eq!(
+        termination_rx.await.unwrap(),
+        Ok(RuntimeResponse::Terminated {
+            cell_id: cell_id("1"),
+            content_items: vec![FunctionCallOutputContentItem::InputText {
+                text: "queued output".to_string(),
+            }],
+        })
+    );
+    harness.task.await.unwrap();
+}
+
+#[tokio::test]
+async fn queued_termination_preempts_unobserved_runtime_completion() {
+    let harness = spawn_cell_control_harness(
+        Some(/*initial_yield_time_ms*/ 60_000),
+        Arc::new(NoopCodeModeSessionDelegate),
+    );
+    harness
+        .event_tx
+        .send(RuntimeEvent::Result {
+            stored_value_writes: HashMap::new(),
+            error_text: None,
+        })
+        .unwrap();
+    let (termination_tx, termination_rx) = oneshot::channel();
+    harness
+        .control_tx
+        .send(CellControlCommand::Terminate {
+            response_tx: termination_tx,
+        })
+        .unwrap();
+
+    let terminated = Ok(RuntimeResponse::Terminated {
+        cell_id: cell_id("1"),
+        content_items: Vec::new(),
+    });
+    assert_eq!(termination_rx.await.unwrap(), terminated.clone());
+    assert_eq!(harness.initial_response_rx.await.unwrap(), terminated);
+    harness.task.await.unwrap();
+}
+
+#[tokio::test]
+async fn yields_and_resumes() {
+    let service = CodeModeService::new();
+    let cell = service
+        .execute(execute_request(
+            r#"text("before"); yield_control(); text("after");"#,
+        ))
+        .await
+        .unwrap();
+
+    assert_eq!(
+        cell.initial_response().await.unwrap(),
+        RuntimeResponse::Yielded {
+            cell_id: cell_id("1"),
+            content_items: vec![FunctionCallOutputContentItem::InputText {
+                text: "before".to_string(),
+            }],
+        }
+    );
+    assert_eq!(
+        service
+            .wait(WaitRequest {
+                cell_id: cell_id("1"),
+                yield_time_ms: 1,
+            })
+            .await
+            .unwrap(),
+        WaitOutcome::LiveCell(RuntimeResponse::Result {
+            cell_id: cell_id("1"),
+            content_items: vec![FunctionCallOutputContentItem::InputText {
+                text: "after".to_string(),
+            }],
+            error_text: None,
+        })
+    );
+}
+
+#[tokio::test]
+async fn returns_and_resumes_from_the_pending_frontier() {
+    let service = CodeModeService::new();
+
+    assert_eq!(
+        service
+            .execute_to_pending(execute_request(
+                r#"
+await new Promise((resolve) => setTimeout(resolve, 60_000));
+text("after");
+"#,
+            ))
+            .await
+            .unwrap(),
+        ExecuteToPendingOutcome::Pending {
+            cell_id: cell_id("1"),
+            content_items: Vec::new(),
+            pending_tool_call_ids: Vec::new(),
+        }
+    );
+
+    service
+        .inner
+        .cells
+        .lock()
+        .await
+        .get(&cell_id("1"))
+        .unwrap()
+        .runtime_tx
+        .send(RuntimeCommand::TimeoutFired { id: 1 })
+        .unwrap();
+
+    assert_eq!(
+        service
+            .wait_to_pending(WaitToPendingRequest {
+                cell_id: cell_id("1"),
+            })
+            .await
+            .unwrap(),
+        WaitToPendingOutcome::LiveCell(ExecuteToPendingOutcome::Completed(
+            RuntimeResponse::Result {
+                cell_id: cell_id("1"),
+                content_items: vec![FunctionCallOutputContentItem::InputText {
+                    text: "after".to_string(),
+                }],
+                error_text: None,
+            }
+        ))
+    );
+}
+
+#[tokio::test]
+async fn observed_natural_completion_wins_over_termination() {
+    let (delegate, mut events_rx) = BlockingDelegate::new();
+    let harness =
+        spawn_cell_control_harness(Some(/*initial_yield_time_ms*/ 60_000), delegate.clone());
+    harness.event_tx.send(RuntimeEvent::YieldRequested).unwrap();
+
+    assert_eq!(
+        harness.initial_response_rx.await.unwrap(),
+        Ok(RuntimeResponse::Yielded {
+            cell_id: cell_id("1"),
+            content_items: Vec::new(),
+        })
+    );
+    harness
+        .event_tx
+        .send(RuntimeEvent::ContentItem(
+            FunctionCallOutputContentItem::InputText {
+                text: "done".to_string(),
+            },
+        ))
+        .unwrap();
+    harness
+        .event_tx
+        .send(RuntimeEvent::Result {
+            stored_value_writes: HashMap::new(),
+            error_text: None,
+        })
+        .unwrap();
+    harness
+        .event_tx
+        .send(RuntimeEvent::Notify {
+            call_id: "notify-1".to_string(),
+            text: "completion observed".to_string(),
+        })
+        .unwrap();
+    assert_eq!(
+        next_event(&mut events_rx).await,
+        DelegateEvent::NotificationStarted
+    );
+
+    let (termination_tx, termination_rx) = oneshot::channel();
+    harness
+        .control_tx
+        .send(CellControlCommand::Terminate {
+            response_tx: termination_tx,
+        })
+        .unwrap();
+    assert_eq!(
+        termination_rx.await.unwrap(),
+        Ok(RuntimeResponse::Result {
+            cell_id: cell_id("1"),
+            content_items: vec![FunctionCallOutputContentItem::InputText {
+                text: "done".to_string(),
+            }],
+            error_text: None,
+        })
+    );
+    harness.task.await.unwrap();
+    assert!(delegate.notification_finished.load(Ordering::Acquire));
+    assert_eq!(
+        next_event(&mut events_rx).await,
+        DelegateEvent::NotificationCancelled
+    );
+    assert_eq!(
+        next_event(&mut events_rx).await,
+        DelegateEvent::CellClosed(cell_id("1"))
+    );
+}
+
+#[tokio::test]
+async fn termination_cancels_pending_callbacks_before_responding() {
+    let (delegate, mut events_rx) = BlockingDelegate::new();
+    let service = CodeModeService::with_delegate(delegate.clone());
+    let cell = service
+        .execute(execute_request(
+            r#"notify("pending"); await new Promise(() => {});"#,
+        ))
+        .await
+        .unwrap();
+
+    assert_eq!(
+        next_event(&mut events_rx).await,
+        DelegateEvent::NotificationStarted
+    );
+    assert_eq!(
+        cell.initial_response().await.unwrap(),
+        RuntimeResponse::Yielded {
+            cell_id: cell_id("1"),
+            content_items: Vec::new(),
+        }
+    );
+    assert_eq!(
+        service.terminate(cell_id("1")).await.unwrap(),
+        WaitOutcome::LiveCell(RuntimeResponse::Terminated {
+            cell_id: cell_id("1"),
+            content_items: Vec::new(),
+        })
+    );
+    assert!(delegate.notification_finished.load(Ordering::Acquire));
+    assert_eq!(
+        next_event(&mut events_rx).await,
+        DelegateEvent::NotificationCancelled
+    );
+    assert_eq!(
+        next_event(&mut events_rx).await,
+        DelegateEvent::CellClosed(cell_id("1"))
+    );
+}
+
+#[tokio::test]
+async fn repeated_termination_is_rejected_while_callback_cleanup_is_pending() {
+    let (delegate, mut events_rx) = HeldNotificationDelegate::new();
+    let service = Arc::new(CodeModeService::with_delegate(delegate.clone()));
+    let cell = service
+        .execute(execute_request(
+            r#"notify("pending"); await new Promise(() => {});"#,
+        ))
+        .await
+        .unwrap();
+
+    assert_eq!(
+        next_event(&mut events_rx).await,
+        DelegateEvent::NotificationStarted
+    );
+    assert_eq!(
+        cell.initial_response().await.unwrap(),
+        RuntimeResponse::Yielded {
+            cell_id: cell_id("1"),
+            content_items: Vec::new(),
+        }
+    );
+
+    let terminating_service = Arc::clone(&service);
+    let first_termination =
+        tokio::spawn(async move { terminating_service.terminate(cell_id("1")).await });
+    assert_eq!(
+        next_event(&mut events_rx).await,
+        DelegateEvent::NotificationCancelled
+    );
+
+    let repeated_termination = service.terminate(cell_id("1")).await;
+    delegate.release_notification();
+
+    assert_eq!(
+        repeated_termination.unwrap_err(),
+        "exec cell 1 is already terminating"
+    );
+    assert_eq!(
+        first_termination.await.unwrap().unwrap(),
+        WaitOutcome::LiveCell(RuntimeResponse::Terminated {
+            cell_id: cell_id("1"),
+            content_items: Vec::new(),
+        })
+    );
+    assert_eq!(
+        next_event(&mut events_rx).await,
+        DelegateEvent::CellClosed(cell_id("1"))
+    );
+}
+
+#[tokio::test]
+async fn second_observer_is_rejected_without_displacing_the_first() {
+    let service = CodeModeService::new();
+    let cell = service
+        .execute(execute_request("await new Promise(() => {});"))
+        .await
+        .unwrap();
+
+    assert_eq!(
+        cell.initial_response().await.unwrap(),
+        RuntimeResponse::Yielded {
+            cell_id: cell_id("1"),
+            content_items: Vec::new(),
+        }
+    );
+
+    let first_observer = service
+        .begin_wait(WaitRequest {
+            cell_id: cell_id("1"),
+            yield_time_ms: 60_000,
+        })
+        .await;
+    assert_eq!(
+        service
+            .wait(WaitRequest {
+                cell_id: cell_id("1"),
+                yield_time_ms: 60_000,
+            })
+            .await
+            .unwrap_err(),
+        "exec cell 1 already has an active observer"
+    );
+
+    let terminated = RuntimeResponse::Terminated {
+        cell_id: cell_id("1"),
+        content_items: Vec::new(),
+    };
+    assert_eq!(
+        service.terminate(cell_id("1")).await.unwrap(),
+        WaitOutcome::LiveCell(terminated.clone())
+    );
+    assert_eq!(
+        first_observer.await.unwrap(),
+        WaitOutcome::LiveCell(terminated)
+    );
+}
+
+#[tokio::test]
+async fn natural_completion_cleans_up_callbacks_before_responding() {
+    let (delegate, mut events_rx) = BlockingDelegate::new();
+    let service = CodeModeService::with_delegate(delegate.clone());
+    let cell = service
+        .execute(ExecuteRequest {
+            enabled_tools: vec![blocking_tool()],
+            source: r#"tools.block({}); text("done");"#.to_string(),
+            yield_time_ms: Some(60_000),
+            ..execute_request("")
+        })
+        .await
+        .unwrap();
+
+    assert_eq!(next_event(&mut events_rx).await, DelegateEvent::ToolStarted);
+    assert_eq!(
+        cell.initial_response().await.unwrap(),
+        RuntimeResponse::Result {
+            cell_id: cell_id("1"),
+            content_items: vec![FunctionCallOutputContentItem::InputText {
+                text: "done".to_string(),
+            }],
+            error_text: None,
+        }
+    );
+    assert!(delegate.tool_finished.load(Ordering::Acquire));
+    assert_eq!(
+        next_event(&mut events_rx).await,
+        DelegateEvent::ToolCancelled
+    );
+    assert_eq!(
+        next_event(&mut events_rx).await,
+        DelegateEvent::CellClosed(cell_id("1"))
+    );
+}