agent/codex
1
0
Fork 0
mirror of https://github.com/pchuan98/codex.git synced 2026-07-01 00:31:56 +08:00
Code Issues Packages Projects Releases Wiki Activity

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

Browse Source 
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
This commit is contained in:
Channing Conger
2026-06-16 13:34:16 -07:00
committed by GitHub
Unverified
parent 4b7351700f
commit e93516e259
3 changed files with 903 additions and 137 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/code-mode/src/service.rs
+260 -131
View File
@@ -86,6 +86,7 @@ struct CellHandle {
control_tx: mpsc::UnboundedSender<CellControlCommand>,
runtime_tx: std::sync::mpsc::Sender<RuntimeCommand>,
cancellation_token: CancellationToken,
termination_requested: Arc<AtomicBool>,
}
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<WaitOutcome, String> {
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<WaitOutcome, String> {
@@ -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<RuntimeResponse>,
response_tx: oneshot::Sender<Result<RuntimeResponse, String>>,
},
PollToPending {
response_tx: oneshot::Sender<ExecuteToPendingOutcome>,
response_tx: oneshot::Sender<Result<ExecuteToPendingOutcome, String>>,
},
Terminate {
response_tx: oneshot::Sender<RuntimeResponse>,
response_tx: oneshot::Sender<Result<RuntimeResponse, String>>,
},
}
enum CellResponseSender {
Runtime(oneshot::Sender<RuntimeResponse>),
ExecuteToPending(oneshot::Sender<ExecuteToPendingOutcome>),
Runtime(oneshot::Sender<Result<RuntimeResponse, String>>),
ExecuteToPending(oneshot::Sender<Result<ExecuteToPendingOutcome, String>>),
}
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<Result<RuntimeResponse, String>>,
) -> 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<CellResponseSender>,
termination_response_tx: Option<oneshot::Sender<Result<RuntimeResponse, String>>>,
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<PendingResult> = 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<std::pin::Pin<Box<tokio::time::Sleep>>> = 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::<Option<RuntimeEvent>>().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;
codex-rs/code-mode/src/service_contract_tests.rs
+634
View File
@@ -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"))
);
}
codex-rs/core/tests/suite/code_mode.rs
+9 -6
View File
@@ -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",