Wire bounded output collection into the code mode runtime (#57850)

## Why

Bounding collected output also requires limiting the raw event backlog between
V8 and the cell actor. Previously, runtime events and cell output accumulated
in an unbounded channel and vector.

## What changed

- Select the event transport and `OutputCollector` from the cell's
  `output_truncation` mode. Server mode uses an ordered, bounded channel with
  backpressure; client mode retains untruncated output.
- Transfer the collector with each observation and restore it on failed
  delivery, preserving retained media and omission metadata.
- Include server-mode script errors in the observation's output budget. Build
  completion output only after completion wins against termination, so a
  rejected error cannot evict output returned by termination.
- Keep the public service on client-side truncation.

## Testing

Add regression coverage for channel backpressure, event ordering, receiver
closure, failed observation delivery, termination during completion commit,
and oversized errors across UTF-8 and omission boundaries.

GitOrigin-RevId: 7604b60416e33c636ca41138be65dc0fd452774a
This commit is contained in:
Sean Huang
2026-10-28 22:33:13 +00:00
committed by copyberry
parent 6d97d6ecb4
commit 92e802fb61
21 changed files with 686 additions and 355 deletions
@@ -5,8 +5,9 @@ use std::sync::atomic::Ordering;
use std::sync::mpsc as std_mpsc;
use std::time::Duration;
use codex_code_mode_protocol::CodeModeOutputTruncation;
use codex_code_mode_protocol::ExecuteRequest;
use codex_code_mode_protocol::FunctionCallOutputContentItem;
use codex_code_mode_protocol::host::BOUNDED_OUTPUT_MAX_IMAGE_BYTES;
use pretty_assertions::assert_eq;
use serde_json::Value as JsonValue;
use tokio::sync::mpsc;
@@ -14,9 +15,14 @@ use tokio::sync::oneshot;
use tokio_util::sync::CancellationToken;
use super::*;
use crate::runtime::RuntimeEventSender;
use crate::session_runtime::OutputItem;
struct TestHost;
#[derive(Default)]
struct TestHost {
commit_started: CancellationToken,
resume_commit: Option<CancellationToken>,
}
#[derive(Default)]
struct RecordingHost {
@@ -44,9 +50,13 @@ impl CellHost for TestHost {
async fn commit_completion(
&self,
_stored_value_writes: HashMap<String, Arc<JsonValue>>,
event: CellEvent,
event: impl FnOnce() -> CellEvent + Send,
cell_state: Arc<CellState>,
) -> CompletionCommit {
self.commit_started.cancel();
if let Some(resume) = &self.resume_commit {
resume.cancelled().await;
}
cell_state.commit_completion(event, || {})
}
@@ -75,7 +85,7 @@ impl CellHost for RecordingHost {
async fn commit_completion(
&self,
_stored_value_writes: HashMap<String, Arc<JsonValue>>,
event: CellEvent,
event: impl FnOnce() -> CellEvent + Send,
cell_state: Arc<CellState>,
) -> CompletionCommit {
cell_state.commit_completion(event, || {})
@@ -85,7 +95,7 @@ impl CellHost for RecordingHost {
}
struct CellActorHarness {
event_tx: mpsc::UnboundedSender<RuntimeEvent>,
event_tx: RuntimeEventSender,
handle: CellHandle,
yield_signal: CancellationToken,
initial_event_rx: oneshot::Receiver<Result<CellEvent, CellError>>,
@@ -95,7 +105,7 @@ struct CellActorHarness {
}
fn spawn_cell_actor_harness(initial_observe_mode: ObserveMode) -> CellActorHarness {
spawn_cell_actor_harness_with_host(initial_observe_mode, Arc::new(TestHost))
spawn_cell_actor_harness_with_host(initial_observe_mode, Arc::new(TestHost::default()))
}
fn spawn_cell_actor_harness_with_host<H: CellHost>(
@@ -105,6 +115,7 @@ fn spawn_cell_actor_harness_with_host<H: CellHost>(
spawn_cell_actor_harness_with_host_and_failure_handler(
initial_observe_mode,
host,
CodeModeOutputTruncation::Client,
/*task_failure_handler*/ None,
)
}
@@ -112,9 +123,11 @@ fn spawn_cell_actor_harness_with_host<H: CellHost>(
fn spawn_cell_actor_harness_with_host_and_failure_handler<H: CellHost>(
initial_observe_mode: ObserveMode,
host: Arc<H>,
output_truncation: CodeModeOutputTruncation,
task_failure_handler: Option<TaskFailureHandler>,
) -> CellActorHarness {
let (event_tx, event_rx) = mpsc::unbounded_channel();
let (event_tx, event_rx) = runtime_event_channel(output_truncation);
let output = OutputCollector::new(output_truncation);
let (command_tx, command_rx) = mpsc::unbounded_channel();
let (initial_event_tx, initial_event_rx) = oneshot::channel();
let (runtime_event_tx, runtime_event_rx) = mpsc::unbounded_channel();
@@ -143,6 +156,7 @@ fn spawn_cell_actor_harness_with_host_and_failure_handler<H: CellHost>(
runtime_control_tx,
runtime_terminate_handle,
cell_state,
output,
},
event_rx,
command_rx,
@@ -169,7 +183,8 @@ async fn unexpected_runtime_thread_exit_is_reported_to_the_session_owner() {
let (failure_tx, mut failure_rx) = mpsc::unbounded_channel();
let harness = spawn_cell_actor_harness_with_host_and_failure_handler(
ObserveMode::YieldAfter(Duration::from_secs(60)),
Arc::new(TestHost),
Arc::new(TestHost::default()),
CodeModeOutputTruncation::Client,
Some(Arc::new(move |reason| {
let _ = failure_tx.send(reason);
})),
@@ -193,19 +208,74 @@ async fn unexpected_runtime_thread_exit_is_reported_to_the_session_owner() {
harness.task.await.expect("cell task");
}
#[tokio::test]
async fn termination_during_completion_commit_preserves_output_and_discards_losing_error() {
for mode in [
CodeModeOutputTruncation::Client,
CodeModeOutputTruncation::Server,
] {
// The actor has received Result but cannot commit until the host releases
// the session store. Termination must decide which output is published.
let resume_commit = CancellationToken::new();
let host = Arc::new(TestHost {
resume_commit: Some(resume_commit.clone()),
..TestHost::default()
});
let harness = spawn_cell_actor_harness_with_host_and_failure_handler(
ObserveMode::YieldAfter(Duration::from_secs(60)),
Arc::clone(&host),
mode,
/*task_failure_handler*/ None,
);
// Simulate V8's dedicated thread, where server-mode sends can block
// while the event queue is full.
let event_tx = harness.event_tx.clone();
tokio::task::spawn_blocking(move || {
event_tx
.blocking_send(RuntimeEvent::ContentItem(OutputItem::InputText {
text: "output that must survive termination".to_string(),
}))
.expect("output");
event_tx
.blocking_send(RuntimeEvent::Result {
stored_value_writes: HashMap::new(),
error_text: Some("uncommitted error".repeat(128 * 1024)),
})
.expect("completion");
})
.await
.expect("runtime producer");
tokio::time::timeout(Duration::from_secs(5), host.commit_started.cancelled())
.await
.expect("commit reached");
let termination = harness.handle.terminate();
resume_commit.cancel();
let mut expected_output = OutputCollector::new(mode);
expected_output.push(OutputItem::InputText {
text: "output that must survive termination".to_string(),
});
let expected = Ok(CellEvent::Terminated {
content_items: expected_output,
});
assert_eq!(termination.await, expected, "{mode:?}");
assert_eq!(harness.initial_event_rx.await.expect("observer"), expected);
harness.task.await.expect("cell task");
}
}
#[tokio::test]
async fn runtime_thread_panic_remains_a_cell_error_without_owner_supervision() {
let harness = spawn_cell_actor_harness(ObserveMode::YieldAfter(Duration::from_secs(60)));
harness
.event_tx
.send(RuntimeEvent::ThreadPanicked)
.blocking_send(RuntimeEvent::ThreadPanicked)
.expect("runtime panic event");
drop(harness.event_tx);
assert_eq!(
harness.initial_event_rx.await.expect("initial event"),
Ok(CellEvent::Completed {
content_items: Vec::new(),
content_items: OutputCollector::Client(Vec::new()),
result_buffered_for: Duration::ZERO,
error_text: Some("exec runtime ended unexpectedly".to_string()),
})
@@ -226,20 +296,21 @@ async fn wait_for_notification(host: &RecordingHost) {
#[tokio::test]
async fn yield_timer_preempts_buffered_runtime_output() {
let harness = spawn_cell_actor_harness(ObserveMode::YieldAfter(Duration::ZERO));
harness.event_tx.send(RuntimeEvent::Started).unwrap();
harness
.event_tx
.send(RuntimeEvent::ContentItem(
FunctionCallOutputContentItem::InputText {
text: "queued output".to_string(),
},
))
.blocking_send(RuntimeEvent::Started)
.unwrap();
harness
.event_tx
.blocking_send(RuntimeEvent::ContentItem(OutputItem::InputText {
text: "queued output".to_string(),
}))
.unwrap();
assert_eq!(
harness.initial_event_rx.await.unwrap(),
Ok(CellEvent::Yielded {
content_items: Vec::new(),
content_items: OutputCollector::Client(Vec::new()),
})
);
@@ -248,9 +319,9 @@ async fn yield_timer_preempts_buffered_runtime_output() {
assert_eq!(
termination.await,
Ok(CellEvent::Terminated {
content_items: vec![OutputItem::Text {
content_items: OutputCollector::Client(vec![OutputItem::InputText {
text: "queued output".to_string(),
}],
}]),
})
);
harness.task.await.unwrap();
@@ -259,19 +330,23 @@ async fn yield_timer_preempts_buffered_runtime_output() {
#[tokio::test]
async fn queued_termination_preempts_unobserved_runtime_completion() {
let harness = spawn_cell_actor_harness(ObserveMode::YieldAfter(Duration::ZERO));
harness.event_tx.send(RuntimeEvent::Started).unwrap();
harness.event_tx.send(RuntimeEvent::YieldRequested).unwrap();
harness
.event_tx
.send(RuntimeEvent::ContentItem(
FunctionCallOutputContentItem::InputText {
text: "queued output".to_string(),
},
))
.blocking_send(RuntimeEvent::Started)
.unwrap();
harness
.event_tx
.send(RuntimeEvent::Result {
.blocking_send(RuntimeEvent::YieldRequested)
.unwrap();
harness
.event_tx
.blocking_send(RuntimeEvent::ContentItem(OutputItem::InputText {
text: "queued output".to_string(),
}))
.unwrap();
harness
.event_tx
.blocking_send(RuntimeEvent::Result {
stored_value_writes: HashMap::new(),
error_text: None,
})
@@ -280,9 +355,9 @@ async fn queued_termination_preempts_unobserved_runtime_completion() {
let termination = harness.handle.terminate();
let terminated = Ok(CellEvent::Terminated {
content_items: vec![OutputItem::Text {
content_items: OutputCollector::Client(vec![OutputItem::InputText {
text: "queued output".to_string(),
}],
}]),
});
assert_eq!(termination.await, terminated.clone());
assert_eq!(harness.initial_event_rx.await.unwrap(), terminated);
@@ -296,7 +371,10 @@ async fn observation_dropped_before_dequeue_does_not_consume_output() {
ObserveMode::YieldAfter(Duration::from_secs(60)),
Arc::clone(&host),
);
harness.event_tx.send(RuntimeEvent::YieldRequested).unwrap();
harness
.event_tx
.blocking_send(RuntimeEvent::YieldRequested)
.unwrap();
assert!(harness.initial_event_rx.await.unwrap().is_ok());
drop(
@@ -306,16 +384,17 @@ async fn observation_dropped_before_dequeue_does_not_consume_output() {
);
harness
.event_tx
.send(RuntimeEvent::ContentItem(
FunctionCallOutputContentItem::InputText {
text: "survives pre-dequeue cancellation".to_string(),
},
))
.blocking_send(RuntimeEvent::ContentItem(OutputItem::InputText {
text: "survives pre-dequeue cancellation".to_string(),
}))
.unwrap();
harness.event_tx.send(RuntimeEvent::YieldRequested).unwrap();
harness
.event_tx
.send(RuntimeEvent::Notify {
.blocking_send(RuntimeEvent::YieldRequested)
.unwrap();
harness
.event_tx
.blocking_send(RuntimeEvent::Notify {
call_id: "after-dropped-command".to_string(),
text: "barrier".to_string(),
})
@@ -328,9 +407,9 @@ async fn observation_dropped_before_dequeue_does_not_consume_output() {
.observe(ObserveMode::YieldAfter(Duration::ZERO))
.await,
Ok(CellEvent::Yielded {
content_items: vec![OutputItem::Text {
content_items: OutputCollector::Client(vec![OutputItem::InputText {
text: "survives pre-dequeue cancellation".to_string(),
}],
}]),
})
);
@@ -339,7 +418,7 @@ async fn observation_dropped_before_dequeue_does_not_consume_output() {
assert_eq!(
termination.await,
Ok(CellEvent::Terminated {
content_items: Vec::new(),
content_items: OutputCollector::Client(Vec::new()),
})
);
harness.task.await.unwrap();
@@ -352,7 +431,10 @@ async fn dropped_yield_observer_preserves_output_for_the_next_observation() {
ObserveMode::YieldAfter(Duration::from_secs(60)),
Arc::clone(&host),
);
harness.event_tx.send(RuntimeEvent::YieldRequested).unwrap();
harness
.event_tx
.blocking_send(RuntimeEvent::YieldRequested)
.unwrap();
assert!(harness.initial_event_rx.await.unwrap().is_ok());
let dropped_observation = harness
@@ -368,16 +450,17 @@ async fn dropped_yield_observer_preserves_output_for_the_next_observation() {
drop(dropped_observation);
harness
.event_tx
.send(RuntimeEvent::ContentItem(
FunctionCallOutputContentItem::InputText {
text: "survives active cancellation".to_string(),
},
))
.blocking_send(RuntimeEvent::ContentItem(OutputItem::InputText {
text: "survives active cancellation".to_string(),
}))
.unwrap();
harness.event_tx.send(RuntimeEvent::YieldRequested).unwrap();
harness
.event_tx
.send(RuntimeEvent::Notify {
.blocking_send(RuntimeEvent::YieldRequested)
.unwrap();
harness
.event_tx
.blocking_send(RuntimeEvent::Notify {
call_id: "after-dropped-observer".to_string(),
text: "barrier".to_string(),
})
@@ -390,9 +473,9 @@ async fn dropped_yield_observer_preserves_output_for_the_next_observation() {
.observe(ObserveMode::YieldAfter(Duration::ZERO))
.await,
Ok(CellEvent::Yielded {
content_items: vec![OutputItem::Text {
content_items: OutputCollector::Client(vec![OutputItem::InputText {
text: "survives active cancellation".to_string(),
}],
}]),
})
);
@@ -401,7 +484,7 @@ async fn dropped_yield_observer_preserves_output_for_the_next_observation() {
assert_eq!(
termination.await,
Ok(CellEvent::Terminated {
content_items: Vec::new(),
content_items: OutputCollector::Client(Vec::new()),
})
);
harness.task.await.unwrap();
@@ -414,7 +497,10 @@ async fn dropped_pending_observer_preserves_the_frontier_for_the_next_observatio
ObserveMode::YieldAfter(Duration::from_secs(60)),
Arc::clone(&host),
);
harness.event_tx.send(RuntimeEvent::YieldRequested).unwrap();
harness
.event_tx
.blocking_send(RuntimeEvent::YieldRequested)
.unwrap();
assert!(harness.initial_event_rx.await.unwrap().is_ok());
let dropped_observation = harness.handle.observe(ObserveMode::PendingFrontier);
@@ -423,9 +509,16 @@ async fn dropped_pending_observer_preserves_the_frontier_for_the_next_observatio
Err(CellError::Busy)
);
drop(dropped_observation);
let item = OutputItem::InputText {
text: "before frontier".to_owned(),
};
harness
.event_tx
.send(RuntimeEvent::ToolCall {
.blocking_send(RuntimeEvent::ContentItem(item.clone()))
.unwrap();
harness
.event_tx
.blocking_send(RuntimeEvent::ToolCall {
id: "tool-1".to_string(),
name: codex_protocol::ToolName {
name: "echo".to_string(),
@@ -435,10 +528,13 @@ async fn dropped_pending_observer_preserves_the_frontier_for_the_next_observatio
input: Some(serde_json::json!({})),
})
.unwrap();
harness.event_tx.send(RuntimeEvent::Pending).unwrap();
harness
.event_tx
.send(RuntimeEvent::Notify {
.blocking_send(RuntimeEvent::Pending)
.unwrap();
harness
.event_tx
.blocking_send(RuntimeEvent::Notify {
call_id: "after-dropped-pending".to_string(),
text: "barrier".to_string(),
})
@@ -448,7 +544,7 @@ async fn dropped_pending_observer_preserves_the_frontier_for_the_next_observatio
assert_eq!(
harness.handle.observe(ObserveMode::PendingFrontier).await,
Ok(CellEvent::Pending {
content_items: Vec::new(),
content_items: OutputCollector::Client(vec![item]),
pending_tool_call_ids: vec!["tool-1".to_string()],
})
);
@@ -462,7 +558,7 @@ async fn dropped_pending_observer_preserves_the_frontier_for_the_next_observatio
assert_eq!(
termination.await,
Ok(CellEvent::Terminated {
content_items: Vec::new(),
content_items: OutputCollector::Client(Vec::new()),
})
);
harness.task.await.unwrap();
@@ -472,12 +568,12 @@ async fn dropped_pending_observer_preserves_the_frontier_for_the_next_observatio
async fn only_the_first_termination_claims_a_buffered_completion() {
let cell_state = CellState::new(CancellationToken::new());
let completion = CellEvent::Completed {
content_items: Vec::new(),
content_items: OutputCollector::Client(Vec::new()),
result_buffered_for: Duration::ZERO,
error_text: None,
};
assert_eq!(
cell_state.commit_completion(completion.clone(), || {}),
cell_state.commit_completion(|| completion.clone(), || {}),
CompletionCommit::Committed
);
assert!(matches!(
@@ -493,7 +589,7 @@ async fn only_the_first_termination_claims_a_buffered_completion() {
assert_eq!(first_termination.await, Ok(completion.clone()));
assert_eq!(
cell_state.finish_termination(CellEvent::Terminated {
content_items: Vec::new(),
content_items: OutputCollector::Client(Vec::new()),
}),
Some(completion)
);
@@ -504,20 +600,17 @@ async fn termination_claim_prevents_stored_value_commit() {
let cell_state = CellState::new(CancellationToken::new());
let termination = cell_state.request_termination();
let mut commit_ran = false;
let completion = CellEvent::Completed {
content_items: Vec::new(),
result_buffered_for: Duration::ZERO,
error_text: None,
};
assert_eq!(
cell_state.commit_completion(completion.clone(), || commit_ran = true),
CompletionCommit::Rejected(completion)
cell_state.commit_completion(
|| panic!("terminated cell must not construct completion output"),
|| commit_ran = true
),
CompletionCommit::Rejected
);
assert!(!commit_ran);
let terminated = CellEvent::Terminated {
content_items: Vec::new(),
content_items: OutputCollector::Client(Vec::new()),
};
assert_eq!(
cell_state.finish_termination(terminated.clone()),
@@ -530,14 +623,14 @@ async fn termination_claim_prevents_stored_value_commit() {
async fn failed_completion_delivery_rebuffers_the_event() {
let cell_state = CellState::new(CancellationToken::new());
let completion = |result_buffered_for| CellEvent::Completed {
content_items: vec![OutputItem::Text {
content_items: OutputCollector::Client(vec![OutputItem::InputText {
text: "done".to_string(),
}],
}]),
result_buffered_for,
error_text: None,
};
assert_eq!(
cell_state.commit_completion(completion(Duration::ZERO), || {}),
cell_state.commit_completion(|| completion(Duration::ZERO), || {}),
CompletionCommit::Committed
);
let (response_tx, response_rx) = oneshot::channel();
@@ -565,3 +658,49 @@ async fn failed_completion_delivery_rebuffers_the_event() {
));
assert_eq!(response_rx.try_recv(), Ok(Ok(completion(2 * interval))));
}
#[test]
fn a_failed_yield_restores_image_allocation_and_omission_order() {
let mut output = OutputCollector::new(CodeModeOutputTruncation::Server);
let image_url = "i".repeat(BOUNDED_OUTPUT_MAX_IMAGE_BYTES);
let pointer = image_url.as_ptr();
output.push(OutputItem::InputImage {
image_url,
detail: None,
});
output.push(OutputItem::InputImage {
image_url: "excess".to_string(),
detail: None,
});
let (tx, rx) = oneshot::channel();
drop(rx);
send_yield(&mut output, tx);
output.push(OutputItem::InputText {
text: "after".to_string(),
});
let (tx, mut rx) = oneshot::channel();
send_yield(&mut output, tx);
let CellEvent::Yielded { content_items } = rx.try_recv().unwrap().unwrap() else {
panic!("expected a yield");
};
let items = crate::output::render_output(
content_items.into_collected_output(),
/*max_output_tokens*/ None,
);
let [
OutputItem::InputImage { image_url, .. },
OutputItem::InputText { text: notice },
OutputItem::InputText { text },
] = items.as_slice()
else {
panic!("expected image, omission notice, and subsequent text");
};
assert_eq!(image_url.as_ptr(), pointer);
assert!(notice.contains("1 images"), "{notice}");
assert_eq!(text, "after");
assert_eq!(
output.take(),
OutputCollector::new(CodeModeOutputTruncation::Server)
);
}