From e93516e25983e115fe39162cfa7912ca374c43eb Mon Sep 17 00:00:00 2001 From: Channing Conger Date: Tue, 16 Jun 2026 13:34:16 -0700 Subject: [PATCH] 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 --- codex-rs/code-mode/src/service.rs | 391 +++++++---- .../code-mode/src/service_contract_tests.rs | 634 ++++++++++++++++++ codex-rs/core/tests/suite/code_mode.rs | 15 +- 3 files changed, 903 insertions(+), 137 deletions(-) create mode 100644 codex-rs/code-mode/src/service_contract_tests.rs diff --git a/codex-rs/code-mode/src/service.rs b/codex-rs/code-mode/src/service.rs index 6eb084ea5..db5cb55a6 100644 --- a/codex-rs/code-mode/src/service.rs +++ b/codex-rs/code-mode/src/service.rs @@ -86,6 +86,7 @@ struct CellHandle { control_tx: mpsc::UnboundedSender, runtime_tx: std::sync::mpsc::Sender, cancellation_token: CancellationToken, + termination_requested: Arc, } struct Inner { @@ -142,7 +143,7 @@ impl CodeModeService { ) .await?; - Ok(StartedCell::new(cell_id, response_rx)) + Ok(StartedCell::from_result_receiver(cell_id, response_rx)) } pub async fn execute_to_pending( @@ -162,7 +163,7 @@ impl CodeModeService { response_rx .await - .map_err(|_| "exec runtime ended unexpectedly".to_string()) + .map_err(|_| "exec runtime ended unexpectedly".to_string())? } async fn start_cell( @@ -195,6 +196,7 @@ impl CodeModeService { control_tx, runtime_tx: runtime_tx.clone(), cancellation_token: cancellation_token.clone(), + termination_requested: Arc::new(AtomicBool::new(false)), }, ); (runtime_tx, runtime_control_tx, runtime_terminate_handle) @@ -220,13 +222,20 @@ impl CodeModeService { } pub async fn wait(&self, request: WaitRequest) -> Result { + self.begin_wait(request).await.await + } + + async fn begin_wait( + &self, + request: WaitRequest, + ) -> CodeModeSessionResultFuture<'static, WaitOutcome> { let WaitRequest { cell_id, yield_time_ms, } = request; let handle = self.inner.cells.lock().await.get(&cell_id).cloned(); let Some(handle) = handle else { - return Ok(WaitOutcome::MissingCell(missing_cell_response(cell_id))); + return missing_wait(cell_id); }; let (response_tx, response_rx) = oneshot::channel(); let control_message = CellControlCommand::Poll { @@ -234,12 +243,9 @@ impl CodeModeService { response_tx, }; if handle.control_tx.send(control_message).is_err() { - return Ok(WaitOutcome::MissingCell(missing_cell_response(cell_id))); - } - match response_rx.await { - Ok(response) => Ok(WaitOutcome::LiveCell(response)), - Err(_) => Ok(WaitOutcome::MissingCell(missing_cell_response(cell_id))), + return missing_wait(cell_id); } + wait_for_response(cell_id, response_rx) } pub async fn terminate(&self, cell_id: CellId) -> Result { @@ -247,16 +253,25 @@ impl CodeModeService { let Some(handle) = handle else { return Ok(WaitOutcome::MissingCell(missing_cell_response(cell_id))); }; + if handle + .termination_requested + .compare_exchange(false, true, Ordering::Relaxed, Ordering::Relaxed) + .is_err() + { + return Err(already_terminating_error(&cell_id)); + } let (response_tx, response_rx) = oneshot::channel(); if handle .control_tx .send(CellControlCommand::Terminate { response_tx }) .is_err() { + handle.termination_requested.store(false, Ordering::Relaxed); return Ok(WaitOutcome::MissingCell(missing_cell_response(cell_id))); } match response_rx.await { - Ok(response) => Ok(WaitOutcome::LiveCell(response)), + Ok(Ok(response)) => Ok(WaitOutcome::LiveCell(response)), + Ok(Err(error_text)) => Err(error_text), Err(_) => Ok(WaitOutcome::MissingCell(missing_cell_response(cell_id))), } } @@ -283,7 +298,8 @@ impl CodeModeService { ))); } match response_rx.await { - Ok(response) => Ok(WaitToPendingOutcome::LiveCell(response)), + Ok(Ok(response)) => Ok(WaitToPendingOutcome::LiveCell(response)), + Ok(Err(error_text)) => Err(error_text), Err(_) => Ok(WaitToPendingOutcome::MissingCell(missing_cell_response( cell_id, ))), @@ -365,19 +381,19 @@ impl CodeModeSession for CodeModeService { enum CellControlCommand { Poll { yield_time_ms: u64, - response_tx: oneshot::Sender, + response_tx: oneshot::Sender>, }, PollToPending { - response_tx: oneshot::Sender, + response_tx: oneshot::Sender>, }, Terminate { - response_tx: oneshot::Sender, + response_tx: oneshot::Sender>, }, } enum CellResponseSender { - Runtime(oneshot::Sender), - ExecuteToPending(oneshot::Sender), + Runtime(oneshot::Sender>), + ExecuteToPending(oneshot::Sender>), } struct PendingResult { @@ -402,6 +418,31 @@ fn missing_cell_response(cell_id: CellId) -> RuntimeResponse { } } +fn missing_wait(cell_id: CellId) -> CodeModeSessionResultFuture<'static, WaitOutcome> { + Box::pin(async move { Ok(WaitOutcome::MissingCell(missing_cell_response(cell_id))) }) +} + +fn wait_for_response( + cell_id: CellId, + response_rx: oneshot::Receiver>, +) -> CodeModeSessionResultFuture<'static, WaitOutcome> { + Box::pin(async move { + match response_rx.await { + Ok(Ok(response)) => Ok(WaitOutcome::LiveCell(response)), + Ok(Err(error_text)) => Err(error_text), + Err(_) => Ok(WaitOutcome::MissingCell(missing_cell_response(cell_id))), + } + }) +} + +fn busy_observer_error(cell_id: &CellId) -> String { + format!("exec cell {cell_id} already has an active observer") +} + +fn already_terminating_error(cell_id: &CellId) -> String { + format!("exec cell {cell_id} is already terminating") +} + fn pending_result_response(cell_id: &CellId, result: PendingResult) -> RuntimeResponse { RuntimeResponse::Result { cell_id: cell_id.clone(), @@ -413,14 +454,27 @@ fn pending_result_response(cell_id: &CellId, result: PendingResult) -> RuntimeRe fn send_terminal_response(response_tx: CellResponseSender, response: RuntimeResponse) { match response_tx { CellResponseSender::Runtime(response_tx) => { - let _ = response_tx.send(response); + let _ = response_tx.send(Ok(response)); } CellResponseSender::ExecuteToPending(response_tx) => { - let _ = response_tx.send(ExecuteToPendingOutcome::Completed(response)); + let _ = response_tx.send(Ok(ExecuteToPendingOutcome::Completed(response))); } } } +fn send_termination_responses( + response_tx: Option, + termination_response_tx: Option>>, + response: RuntimeResponse, +) { + if let Some(response_tx) = response_tx { + send_terminal_response(response_tx, response.clone()); + } + if let Some(termination_response_tx) = termination_response_tx { + let _ = termination_response_tx.send(Ok(response)); + } +} + fn send_or_buffer_result( cell_id: &CellId, result: PendingResult, @@ -447,10 +501,10 @@ fn send_yield_response( }; match current_response_tx { CellResponseSender::Runtime(response_tx) => { - let _ = response_tx.send(RuntimeResponse::Yielded { + let _ = response_tx.send(Ok(RuntimeResponse::Yielded { cell_id: cell_id.clone(), content_items: std::mem::take(content_items), - }); + })); } CellResponseSender::ExecuteToPending(execute_to_pending_tx) => { *response_tx = Some(CellResponseSender::ExecuteToPending(execute_to_pending_tx)); @@ -478,30 +532,134 @@ async fn run_cell_control( let mut pending_tool_call_ids = Vec::new(); let mut pending_result: Option = None; let mut response_tx = Some(initial_response_tx); + let mut termination_response_tx = None; let mut termination_requested = false; let mut runtime_closed = false; let mut yield_timer: Option>> = None; let mut notification_tasks = JoinSet::new(); + let mut tool_tasks = JoinSet::new(); loop { + let yield_deadline_elapsed = yield_timer + .as_ref() + .is_some_and(|yield_timer| yield_timer.deadline() <= tokio::time::Instant::now()); tokio::select! { + biased; + maybe_command = control_rx.recv() => { + let Some(command) = maybe_command else { + break; + }; + match command { + CellControlCommand::Poll { + yield_time_ms, + response_tx: next_response_tx, + } => { + if let Some(result) = pending_result.take() { + let _ = next_response_tx.send(Ok(pending_result_response(&cell_id, result))); + break; + } + if response_tx.is_some() || termination_response_tx.is_some() { + let _ = next_response_tx.send(Err(busy_observer_error(&cell_id))); + continue; + } + response_tx = Some(CellResponseSender::Runtime(next_response_tx)); + yield_timer = Some(Box::pin(tokio::time::sleep(Duration::from_millis(yield_time_ms)))); + resume_paused_runtime(&runtime_control_tx, pending_mode); + } + CellControlCommand::PollToPending { + response_tx: next_response_tx, + } => { + if let Some(result) = pending_result.take() { + let response = pending_result_response(&cell_id, result); + let _ = next_response_tx + .send(Ok(ExecuteToPendingOutcome::Completed(response))); + break; + } + if response_tx.is_some() || termination_response_tx.is_some() { + let _ = next_response_tx.send(Err(busy_observer_error(&cell_id))); + continue; + } + response_tx = + Some(CellResponseSender::ExecuteToPending(next_response_tx)); + yield_timer = None; + resume_paused_runtime(&runtime_control_tx, pending_mode); + } + CellControlCommand::Terminate { response_tx: next_response_tx } => { + if let Some(result) = pending_result.take() { + let _ = next_response_tx.send(Ok(pending_result_response(&cell_id, result))); + break; + } + + if termination_response_tx.is_some() { + let _ = next_response_tx.send(Err(already_terminating_error(&cell_id))); + continue; + } + + termination_response_tx = Some(next_response_tx); + termination_requested = true; + cancellation_token.cancel(); + yield_timer = None; + let _ = runtime_tx.send(RuntimeCommand::Terminate); + terminate_paused_runtime(&runtime_control_tx, pending_mode); + let _ = runtime_terminate_handle.terminate_execution(); + if runtime_closed { + finish_callbacks( + &cancellation_token, + &mut notification_tasks, + &mut tool_tasks, + CallbackCompletion::Cancel, + ).await; + let response = RuntimeResponse::Terminated { + cell_id: cell_id.clone(), + content_items: std::mem::take(&mut content_items), + }; + send_termination_responses( + response_tx.take(), + termination_response_tx.take(), + response, + ); + break; + } else { + continue; + } + } + } + } + _ = async { + if let Some(yield_timer) = yield_timer.as_mut() { + yield_timer.await; + } else { + std::future::pending::<()>().await; + } + } => { + yield_timer = None; + send_yield_response(&cell_id, &mut content_items, &mut response_tx); + } maybe_event = async { if runtime_closed { std::future::pending::>().await } else { event_rx.recv().await } - } => { + }, if !yield_deadline_elapsed => { let Some(event) = maybe_event else { runtime_closed = true; if termination_requested { - if let Some(response_tx) = response_tx.take() { - let response = RuntimeResponse::Terminated { - cell_id: cell_id.clone(), - content_items: std::mem::take(&mut content_items), - }; - send_terminal_response(response_tx, response); - } + finish_callbacks( + &cancellation_token, + &mut notification_tasks, + &mut tool_tasks, + CallbackCompletion::Cancel, + ).await; + let response = RuntimeResponse::Terminated { + cell_id: cell_id.clone(), + content_items: std::mem::take(&mut content_items), + }; + send_termination_responses( + response_tx.take(), + termination_response_tx.take(), + response, + ); break; } if pending_result.is_none() { @@ -534,13 +692,13 @@ async fn run_cell_control( Some(CellResponseSender::Runtime(runtime_response_tx)); } CellResponseSender::ExecuteToPending(response_tx) => { - let _ = response_tx.send(ExecuteToPendingOutcome::Pending { + let _ = response_tx.send(Ok(ExecuteToPendingOutcome::Pending { cell_id: cell_id.clone(), content_items: std::mem::take(&mut content_items), pending_tool_call_ids: std::mem::take( &mut pending_tool_call_ids, ), - }); + })); } } } @@ -557,20 +715,13 @@ async fn run_cell_control( let cell_id = cell_id.clone(); let cancellation_token = cancellation_token.child_token(); notification_tasks.spawn(async move { - tokio::select! { - result = delegate.notify( - call_id, - cell_id.clone(), - text, - cancellation_token.clone(), - ) => { - if let Err(err) = result { - warn!( - "failed to deliver code mode notification for cell {cell_id}: {err}" - ); - } - } - _ = cancellation_token.cancelled() => {} + if let Err(err) = delegate + .notify(call_id, cell_id.clone(), text, cancellation_token) + .await + { + warn!( + "failed to deliver code mode notification for cell {cell_id}: {err}" + ); } }); } @@ -593,11 +744,8 @@ async fn run_cell_control( let delegate = Arc::clone(&inner.delegate); let runtime_tx = runtime_tx.clone(); let cancellation_token = cancellation_token.child_token(); - tokio::spawn(async move { - let response = tokio::select! { - response = delegate.invoke_tool(tool_call, cancellation_token.clone()) => response, - _ = cancellation_token.cancelled() => return, - }; + tool_tasks.spawn(async move { + let response = delegate.invoke_tool(tool_call, cancellation_token).await; let command = match response { Ok(result) => RuntimeCommand::ToolResponse { id, result }, Err(error_text) => RuntimeCommand::ToolError { id, error_text }, @@ -611,16 +759,29 @@ async fn run_cell_control( } => { yield_timer = None; if termination_requested { - if let Some(response_tx) = response_tx.take() { - let response = RuntimeResponse::Terminated { - cell_id: cell_id.clone(), - content_items: std::mem::take(&mut content_items), - }; - send_terminal_response(response_tx, response); - } + finish_callbacks( + &cancellation_token, + &mut notification_tasks, + &mut tool_tasks, + CallbackCompletion::Cancel, + ).await; + let response = RuntimeResponse::Terminated { + cell_id: cell_id.clone(), + content_items: std::mem::take(&mut content_items), + }; + send_termination_responses( + response_tx.take(), + termination_response_tx.take(), + response, + ); break; } - drain_notification_tasks(&mut notification_tasks).await; + finish_callbacks( + &cancellation_token, + &mut notification_tasks, + &mut tool_tasks, + CallbackCompletion::DrainNotifications, + ).await; inner .stored_values .lock() @@ -648,92 +809,56 @@ async fn run_cell_control( warn!("code mode notification task failed: {err}"); } } - maybe_command = control_rx.recv() => { - let Some(command) = maybe_command else { - break; - }; - match command { - CellControlCommand::Poll { - yield_time_ms, - response_tx: next_response_tx, - } => { - if let Some(result) = pending_result.take() { - let _ = next_response_tx.send(pending_result_response(&cell_id, result)); - break; - } - response_tx = Some(CellResponseSender::Runtime(next_response_tx)); - yield_timer = Some(Box::pin(tokio::time::sleep(Duration::from_millis(yield_time_ms)))); - resume_paused_runtime(&runtime_control_tx, pending_mode); - } - CellControlCommand::PollToPending { - response_tx: next_response_tx, - } => { - if let Some(result) = pending_result.take() { - let response = pending_result_response(&cell_id, result); - let _ = next_response_tx - .send(ExecuteToPendingOutcome::Completed(response)); - break; - } - response_tx = - Some(CellResponseSender::ExecuteToPending(next_response_tx)); - yield_timer = None; - resume_paused_runtime(&runtime_control_tx, pending_mode); - } - CellControlCommand::Terminate { response_tx: next_response_tx } => { - if let Some(result) = pending_result.take() { - let _ = next_response_tx.send(pending_result_response(&cell_id, result)); - break; - } - - response_tx = Some(CellResponseSender::Runtime(next_response_tx)); - termination_requested = true; - cancellation_token.cancel(); - yield_timer = None; - let _ = runtime_tx.send(RuntimeCommand::Terminate); - terminate_paused_runtime(&runtime_control_tx, pending_mode); - let _ = runtime_terminate_handle.terminate_execution(); - if runtime_closed { - if let Some(response_tx) = response_tx.take() { - let response = RuntimeResponse::Terminated { - cell_id: cell_id.clone(), - content_items: std::mem::take(&mut content_items), - }; - send_terminal_response(response_tx, response); - } - break; - } else { - continue; - } - } + task_result = tool_tasks.join_next(), if !tool_tasks.is_empty() => { + if let Some(Err(err)) = task_result + && !err.is_cancelled() + { + warn!("code mode tool task failed: {err}"); } } - _ = async { - if let Some(yield_timer) = yield_timer.as_mut() { - yield_timer.await; - } else { - std::future::pending::<()>().await; - } - } => { - yield_timer = None; - send_yield_response(&cell_id, &mut content_items, &mut response_tx); - } } } let _ = runtime_tx.send(RuntimeCommand::Terminate); cancellation_token.cancel(); - drain_notification_tasks(&mut notification_tasks).await; + finish_callbacks( + &cancellation_token, + &mut notification_tasks, + &mut tool_tasks, + CallbackCompletion::Cancel, + ) + .await; terminate_paused_runtime(&runtime_control_tx, pending_mode); inner.cells.lock().await.remove(&cell_id); inner.delegate.cell_closed(&cell_id); } -async fn drain_notification_tasks(notification_tasks: &mut JoinSet<()>) { - while let Some(result) = notification_tasks.join_next().await { +#[derive(Clone, Copy)] +enum CallbackCompletion { + DrainNotifications, + Cancel, +} + +async fn finish_callbacks( + cancellation_token: &CancellationToken, + notification_tasks: &mut JoinSet<()>, + tool_tasks: &mut JoinSet<()>, + completion: CallbackCompletion, +) { + if matches!(completion, CallbackCompletion::Cancel) { + cancellation_token.cancel(); + } + drain_tasks(notification_tasks, "notification").await; + cancellation_token.cancel(); + drain_tasks(tool_tasks, "tool").await; +} + +async fn drain_tasks(tasks: &mut JoinSet<()>, description: &str) { + while let Some(result) = tasks.join_next().await { if let Err(err) = result && !err.is_cancelled() { - warn!("code mode notification task failed: {err}"); + warn!("code mode {description} task failed: {err}"); } } } @@ -1786,10 +1911,10 @@ image({ event_tx.send(RuntimeEvent::YieldRequested).unwrap(); assert_eq!( initial_response_rx.await.unwrap(), - RuntimeResponse::Yielded { + Ok(RuntimeResponse::Yielded { cell_id: cell_id("cell-1"), content_items: Vec::new(), - } + }) ); let (terminate_response_tx, terminate_response_rx) = oneshot::channel(); @@ -1810,12 +1935,16 @@ image({ assert_eq!( terminate_response.await, - RuntimeResponse::Terminated { + Ok(RuntimeResponse::Terminated { cell_id: cell_id("cell-1"), content_items: Vec::new(), - } + }) ); let _ = runtime_tx.send(RuntimeCommand::Terminate); } } + +#[cfg(test)] +#[path = "service_contract_tests.rs"] +mod contract_tests; diff --git a/codex-rs/code-mode/src/service_contract_tests.rs b/codex-rs/code-mode/src/service_contract_tests.rs new file mode 100644 index 000000000..752c6cc6c --- /dev/null +++ b/codex-rs/code-mode/src/service_contract_tests.rs @@ -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, + notification_finished: AtomicBool, + tool_finished: AtomicBool, +} + +struct HeldNotificationDelegate { + events_tx: mpsc::UnboundedSender, + notification_release: Notify, +} + +impl HeldNotificationDelegate { + fn new() -> (Arc, mpsc::UnboundedReceiver) { + 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, + control_tx: mpsc::UnboundedSender, + initial_response_rx: oneshot::Receiver>, + task: tokio::task::JoinHandle<()>, + _runtime_event_rx: mpsc::UnboundedReceiver, +} + +fn spawn_cell_control_harness( + initial_yield_time_ms: Option, + delegate: Arc, +) -> 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, mpsc::UnboundedReceiver) { + 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 { + 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")) + ); +} diff --git a/codex-rs/core/tests/suite/code_mode.rs b/codex-rs/core/tests/suite/code_mode.rs index 25fc3ccf2..6a4121cc5 100644 --- a/codex-rs/core/tests/suite/code_mode.rs +++ b/codex-rs/core/tests/suite/code_mode.rs @@ -1411,7 +1411,7 @@ text("phase 3"); #[cfg_attr(windows, ignore = "no exec_command on Windows")] #[tokio::test(flavor = "multi_thread", worker_threads = 2)] -async fn code_mode_yield_timeout_works_for_busy_loop() -> Result<()> { +async fn code_mode_yield_and_termination_are_not_starved_by_runtime_output() -> Result<()> { skip_if_no_network!(Ok(())); let server = responses::start_mock_server().await; @@ -1420,8 +1420,12 @@ async fn code_mode_yield_timeout_works_for_busy_loop() -> Result<()> { }); let test = builder.build(&server).await?; - let code = r#"// @exec: {"yield_time_ms": 100} -text("phase 1"); + // Exact controller arbitration is covered by deterministic code-mode contract tests. Keep + // this end-to-end load bounded while exercising a substantial runtime output backlog. + let code = r#"// @exec: {"yield_time_ms": 0, "max_output_tokens": 16} +for (let index = 0; index < 16_384; index++) { + text(`event ${index}`); +} while (true) {} "#; @@ -1451,7 +1455,7 @@ while (true) {} let first_request = first_completion.single_request(); let first_items = custom_tool_output_items(&first_request, "call-1"); - assert_eq!(first_items.len(), 2); + assert_eq!(first_items.len(), 1); assert_regex_match( concat!( r"(?s)\A", @@ -1459,7 +1463,6 @@ while (true) {} ), text_item(&first_items, /*index*/ 0), ); - assert_eq!(text_item(&first_items, /*index*/ 1), "phase 1"); let cell_id = extract_running_cell_id(text_item(&first_items, /*index*/ 0)); responses::mount_sse_once( @@ -1491,7 +1494,7 @@ while (true) {} let second_request = second_completion.single_request(); let second_items = function_tool_output_items(&second_request, "call-2"); - assert_eq!(second_items.len(), 1); + assert!(!second_items.is_empty()); assert_regex_match( concat!( r"(?s)\A",