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
@@ -7,6 +7,7 @@ use std::sync::Arc;
use std::sync::atomic::AtomicU64;
use std::sync::atomic::Ordering;
pub(crate) use codex_code_mode_protocol::FunctionCallOutputContentItem as OutputItem;
use opentelemetry::context::FutureExt;
use serde_json::Value as JsonValue;
use tokio::sync::Mutex;
@@ -17,11 +18,9 @@ pub(crate) use self::types::CellEvent;
pub(crate) use self::types::CellId;
pub(crate) use self::types::CreateCellRequest;
pub(crate) use self::types::Error;
pub(crate) use self::types::ImageDetail;
pub(crate) use self::types::NestedToolCall;
pub(crate) use self::types::Observation;
pub(crate) use self::types::ObserveMode;
pub(crate) use self::types::OutputItem;
pub(crate) use self::types::SessionRuntimeDelegate;
pub(crate) use self::types::ToolDefinition;
pub(crate) use self::types::ToolKind;
@@ -286,19 +285,20 @@ impl<D: SessionRuntimeDelegate> CellHost for RuntimeCellHost<D> {
async fn commit_completion(
&self,
stored_value_writes: HashMap<String, Arc<JsonValue>>,
event: CellEvent,
event: impl FnOnce() -> CellEvent + Send,
cell_state: Arc<CellState>,
) -> CompletionCommit {
let cancellation_token = cell_state.cancellation_token();
let mut stored_values = tokio::select! {
biased;
_ = cancellation_token.cancelled() => {
return CompletionCommit::Rejected(event);
return CompletionCommit::Rejected;
}
stored_values = self.inner.stored_values.lock() => stored_values,
};
cell_state.commit_completion(event, || {
stored_values.extend(stored_value_writes);
drop(stored_values);
})
}
@@ -6,12 +6,14 @@ use std::task::Poll;
use std::task::Waker;
use std::time::Duration;
use codex_code_mode_protocol::CodeModeOutputTruncation;
use pretty_assertions::assert_eq;
use serde_json::Value as JsonValue;
use tokio_util::sync::CancellationToken;
use super::*;
use crate::cell_actor::CompletionCommit;
use crate::output::OutputCollector;
struct RecordingDelegate;
@@ -80,9 +82,9 @@ async fn reports_cell_actor_panics_to_the_owner() {
assert_eq!(
started.initial_event().await,
Ok(CellEvent::Completed {
content_items: vec![OutputItem::Text {
content_items: OutputCollector::Client(vec![OutputItem::InputText {
text: "done".to_string(),
}],
}]),
result_buffered_for: Duration::ZERO,
error_text: None,
})
@@ -104,21 +106,13 @@ async fn termination_rejects_a_waiting_store_commit_before_the_next_cell_can_loa
inner: Arc::clone(&runtime.inner),
execution_context: opentelemetry::Context::new(),
};
let completion = CellEvent::Completed {
content_items: vec![OutputItem::Text {
text: "uncommitted output".to_string(),
}],
result_buffered_for: Duration::ZERO,
error_text: None,
};
let stored_values = runtime.inner.stored_values.lock().await;
let commit = host.commit_completion(
HashMap::from([(
"candidate".to_string(),
Arc::new(JsonValue::String("lost".to_string())),
)]),
completion.clone(),
|| panic!("terminated cell must not construct completion output"),
Arc::clone(&cell_state),
);
tokio::pin!(commit);
@@ -128,9 +122,9 @@ async fn termination_rejects_a_waiting_store_commit_before_the_next_cell_can_loa
let termination = cell_state.request_termination();
drop(stored_values);
assert_eq!(commit.await, CompletionCommit::Rejected(completion));
assert_eq!(commit.await, CompletionCommit::Rejected);
let terminated = CellEvent::Terminated {
content_items: Vec::new(),
content_items: OutputCollector::Client(Vec::new()),
};
assert_eq!(
cell_state.finish_termination(terminated.clone()),
@@ -149,6 +143,7 @@ async fn termination_rejects_a_waiting_store_commit_before_the_next_cell_can_loa
let reader = runtime
.execute(
CreateCellRequest {
output_truncation: CodeModeOutputTruncation::Client,
tool_call_id: "reader".to_string(),
enabled_tools: Vec::new(),
source: r#"text(String(load("candidate")));"#.to_string(),
@@ -161,9 +156,9 @@ async fn termination_rejects_a_waiting_store_commit_before_the_next_cell_can_loa
assert_eq!(
reader.initial_event().await,
Ok(CellEvent::Completed {
content_items: vec![OutputItem::Text {
content_items: OutputCollector::Client(vec![OutputItem::InputText {
text: "undefined".to_string(),
}],
}]),
result_buffered_for: Duration::ZERO,
error_text: None,
})
@@ -173,6 +168,7 @@ async fn termination_rejects_a_waiting_store_commit_before_the_next_cell_can_loa
fn execute_request(source: &str) -> CreateCellRequest {
CreateCellRequest {
output_truncation: CodeModeOutputTruncation::Client,
tool_call_id: "call-1".to_string(),
enabled_tools: Vec::new(),
source: source.to_string(),
@@ -255,7 +251,7 @@ async fn drop_terminates_cells_when_the_registry_is_locked() {
assert_eq!(
started.initial_event().await,
Ok(CellEvent::Yielded {
content_items: Vec::new(),
content_items: OutputCollector::Client(Vec::new()),
})
);
@@ -5,6 +5,8 @@ use std::time::Duration;
use serde_json::Value as JsonValue;
use tokio_util::sync::CancellationToken;
use crate::output::OutputCollector;
/// Identifies one execution cell within a session runtime.
#[derive(Clone, Debug, Eq, Hash, PartialEq)]
pub(crate) struct CellId(String);
@@ -54,54 +56,32 @@ impl From<ObserveMode> for Observation {
}
/// An observable cell lifecycle event.
#[derive(Clone, Debug, PartialEq)]
#[derive(Clone, Debug)]
#[cfg_attr(test, derive(PartialEq))]
pub(crate) enum CellEvent {
Yielded {
content_items: Vec<OutputItem>,
content_items: OutputCollector,
},
Pending {
content_items: Vec<OutputItem>,
content_items: OutputCollector,
pending_tool_call_ids: Vec<String>,
},
Completed {
content_items: Vec<OutputItem>,
content_items: OutputCollector,
error_text: Option<String>,
/// Time the completed result spent waiting for an observer.
result_buffered_for: Duration,
},
Terminated {
content_items: Vec<OutputItem>,
content_items: OutputCollector,
},
}
/// Output emitted by a cell since its preceding observation.
#[derive(Clone, Debug, PartialEq)]
pub(crate) enum OutputItem {
Text {
text: String,
},
Image {
image_url: String,
detail: Option<ImageDetail>,
},
Audio {
audio_url: String,
},
}
/// Requested image fidelity for an output image.
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub(crate) enum ImageDetail {
Auto,
Low,
High,
Original,
}
/// Transport-neutral input for creating a cell.
///
/// The owning session assigns the cell ID when it admits the request.
pub(crate) struct CreateCellRequest {
pub(crate) output_truncation: codex_code_mode_protocol::CodeModeOutputTruncation,
pub(crate) tool_call_id: String,
pub(crate) enabled_tools: Vec<ToolDefinition>,
pub(crate) source: String,