openai/codex
1
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-10-29 19:08:43 +00:00
Code Issues Packages Projects Releases Wiki Activity

Enable server-side output truncation for Code Mode sessions (#57870)

Browse Source 
## Why

Code Mode sessions always used client-side output handling, leaving server truncation settings and observation budgets unused.

## What changed

- Honor each session's `output_truncation` setting in the runtime and both host transports, and advertise the bounded-output capability.
- Apply `max_output_tokens` independently to execute, wait, and terminate responses in server mode. Keep client mode as the default when the setting is omitted.
- Use bounded output collection and rendering in server mode, preserving head and tail output with omission notices. Include script errors in the output budget while retaining failure status and a separate diagnostic capped at 64 KiB.
- Estimate audio tokens without the Tokio-dependent cache when rendering bounded output.

## Testing

Add runtime and host regression tests for session mode selection, per-request budgets, output overflow, oversized images, zero-token audio output, diagnostic preservation, and termination under output backpressure.

GitOrigin-RevId: 3d9661edcdd2dde31526bb92eb83fa0614f0fa88
This commit is contained in:
Sean Huang
2026-10-29 01:04:48 +00:00
committed by copyberry
parent 4e2dfd2230
commit edd1ae37ec
22 changed files with 1148 additions and 166 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/code-mode-runtime/Cargo.toml
+1
View File
@@ -19,6 +19,7 @@ codex-plugin-host = { workspace = true }
codex-protocol = { workspace = true }
codex-utils-audio = { workspace = true }
codex-utils-output-truncation = { workspace = true }
codex-utils-string = { workspace = true }
deno_core_icudata = { workspace = true }
futures = { workspace = true, features = ["executor", "std"] }
opentelemetry = { workspace = true }
codex-rs/code-mode-runtime/src/cell_actor/mod.rs
+5 -2
View File
@@ -474,8 +474,11 @@ async fn run_cell<H: CellHost>(
host.closed().await;
}
/// Construct completion only after it has won the terminal decision.
fn take_completed_event(output: &mut OutputCollector, error_text: Option<String>) -> CellEvent {
/// Construct completion only after the terminal outcome is settled.
pub(crate) fn take_completed_event(
output: &mut OutputCollector,
error_text: Option<String>,
) -> CellEvent {
let error_text = error_text.map(|error_text| output.collect_error(error_text));
CellEvent::Completed {
content_items: output.take(),
codex-rs/code-mode-runtime/src/cell_actor/tests.rs
+60
View File
@@ -178,6 +178,66 @@ fn spawn_cell_actor_harness_with_host_and_failure_handler<H: CellHost>(
}
}
#[tokio::test]
async fn termination_drains_a_send_already_blocked_by_backpressure() {
let harness = spawn_cell_actor_harness_with_host_and_failure_handler(
ObserveMode::YieldAfter(Duration::from_secs(60)),
Arc::new(TestHost::default()),
CodeModeOutputTruncation::Server,
/*task_failure_handler*/ None,
);
let (termination, capacity) = {
let RuntimeEventSender::Bounded(tx) = &harness.event_tx else {
panic!("server runtime must use a bounded sender");
};
// This test owns the V8-side events. The current-thread actor cannot run
// before our first await; fill its actual FIFO and register a sender.
for _ in 0..tx.max_capacity() {
tx.try_send(RuntimeEvent::ContentItem(OutputItem::InputText {
text: "queued".to_string(),
}))
.unwrap();
}
let mut send = std::pin::pin!(tx.send(RuntimeEvent::ContentItem(OutputItem::InputText {
text: "blocked".to_string(),
})));
assert!(futures::poll!(send.as_mut()).is_pending());
// terminate() synchronously cancels. No event was consumed between the
// Pending poll and cancellation; run_cell must drain while terminating.
let termination = harness.handle.terminate();
tokio::time::timeout(Duration::from_secs(1), send)
.await
.expect("termination must release the registered sender")
.expect("actor must drain the blocked item");
(termination, tx.max_capacity())
};
drop(harness.event_tx);
let event = tokio::time::timeout(Duration::from_secs(1), termination)
.await
.expect("termination must finish after draining")
.expect("termination event");
let mut expected = OutputCollector::new(CodeModeOutputTruncation::Server);
for _ in 0..capacity {
expected.push(OutputItem::InputText {
text: "queued".to_string(),
});
}
expected.push(OutputItem::InputText {
text: "blocked".to_string(),
});
assert_eq!(
event,
CellEvent::Terminated {
content_items: expected,
}
);
assert_eq!(harness.initial_event_rx.await.unwrap(), Ok(event));
tokio::time::timeout(Duration::from_secs(1), harness.task)
.await
.expect("actor must close")
.unwrap();
}
#[tokio::test]
async fn unexpected_runtime_thread_exit_is_reported_to_the_session_owner() {
let (failure_tx, mut failure_rx) = mpsc::unbounded_channel();
codex-rs/code-mode-runtime/src/output/mod.rs
+14 -4
View File
@@ -6,14 +6,19 @@
use codex_code_mode_protocol::CodeModeOutputTruncation;
use codex_code_mode_protocol::FunctionCallOutputContentItem as OutputItem;
use codex_utils_output_truncation::OutputBoundary;
use codex_utils_string::truncate_middle_with_marker;
mod output_collection;
mod output_render;
pub(crate) use output_collection::BoundedOutput;
#[allow(unused_imports, reason = "Renderer is wired in the next stack layer")]
#[cfg(test)]
pub(crate) use output_collection::MAX_OUTPUT_ITEMS;
pub(crate) use output_render::render_output;
/// Bound the buffered server diagnostic independently of model-visible output.
pub(crate) const MAX_SERVER_DIAGNOSTIC_BYTES: usize = 64 * 1024;
/// The actor transfers the whole collector at each observation, so a failed
/// synchronous send can restore the original state, including any omissions.
#[derive(Clone, Debug)]
@@ -39,8 +44,8 @@ impl OutputCollector {
}
}
/// Server errors share the observation's budget; preserve their presence
/// for completion status. Client errors retain their original formatting.
/// Server errors share the observation's output budget; their diagnostic
/// has an independent bound. Client errors retain their original formatting.
pub(crate) fn collect_error(&mut self, error_text: String) -> String {
match self {
Self::Client(_) => error_text,
@@ -48,7 +53,12 @@ impl OutputCollector {
output
.get_or_insert_with(Box::default)
.push_error(&error_text);
String::new()
// Reuse small errors without retaining oversized backing storage.
if error_text.capacity() <= MAX_SERVER_DIAGNOSTIC_BYTES {
error_text
} else {
truncate_middle_with_marker(&error_text, MAX_SERVER_DIAGNOSTIC_BYTES, "…")
}
}
}
}
codex-rs/code-mode-runtime/src/output/output_render.rs
+3 -7
View File
@@ -6,7 +6,7 @@ use codex_code_mode_protocol::ImageDetail;
use codex_protocol::models::FunctionCallOutputContentItem;
use codex_protocol::models::ImageDetail as ProtocolImageDetail;
use codex_protocol::models::ImageReference;
use codex_utils_audio::estimate_audio_token_count;
use codex_utils_audio::estimate_audio_token_count_without_cache;
use codex_utils_output_truncation::OutputBoundary;
use codex_utils_output_truncation::formatted_output_at_token_boundaries;
use codex_utils_output_truncation::mixed_output_at_token_boundaries;
@@ -30,10 +30,6 @@ const MAX_POSITIONED_NOTICES: usize = 4;
/// budget: Core's warning describes the retained text alone. The shared
/// formatter moves collector boundaries along with their source content.
/// If losses are too scattered for individual notices, report their totals once.
#[cfg_attr(
not(test),
allow(dead_code, reason = "Renderer is wired in the next stack layer")
)]
pub(crate) fn render_output(
output: CollectedOutput,
max_output_tokens: Option<usize>,
@@ -72,14 +68,14 @@ pub(crate) fn render_output(
mixed_output_at_token_boundaries(
content,
max_tokens,
estimate_audio_token_count,
estimate_audio_token_count_without_cache,
&mut boundaries,
)
} else {
formatted_output_at_token_boundaries(
content,
max_tokens,
estimate_audio_token_count,
estimate_audio_token_count_without_cache,
&mut boundaries,
)
};
codex-rs/code-mode-runtime/src/output/output_render_tests.rs
+2 -2
View File
@@ -31,8 +31,8 @@ fn missing_budget_uses_the_protocol_default() {
}
}
#[test]
fn typed_media_round_trips_with_the_audio_estimate() {
#[tokio::test]
async fn typed_media_round_trips_with_the_audio_estimate() {
let mut images: Vec<_> = [
None,
Some(ImageDetail::Auto),
codex-rs/code-mode-runtime/src/service.rs
+141 -55
View File
@@ -7,12 +7,12 @@ use codex_code_mode_protocol::CodeModeOutputTruncation;
use codex_code_mode_protocol::CodeModeSession;
use codex_code_mode_protocol::CodeModeSessionCellExecutionLimits;
use codex_code_mode_protocol::CodeModeSessionDelegate;
use codex_code_mode_protocol::CodeModeSessionOptions;
use codex_code_mode_protocol::CodeModeSessionResultFuture;
use codex_code_mode_protocol::CodeModeToolKind;
use codex_code_mode_protocol::DEFAULT_EXEC_YIELD_TIME_MS;
use codex_code_mode_protocol::ExecuteRequest;
use codex_code_mode_protocol::ExecuteToPendingOutcome;
#[cfg(test)]
use codex_code_mode_protocol::FunctionCallOutputContentItem;
use codex_code_mode_protocol::RuntimeResponse;
use codex_code_mode_protocol::StartedCell;
@@ -25,6 +25,9 @@ use serde_json::Value as JsonValue;
use tokio::sync::oneshot;
use tokio_util::sync::CancellationToken;
use crate::TaskFailureHandler;
use crate::cell_actor::take_completed_event;
use crate::output::OutputCollector;
use crate::session_runtime as runtime;
use crate::session_runtime::SessionRuntime;
@@ -33,7 +36,7 @@ const MIN_YIELD_TIME_FOR_GRACE: Duration = Duration::from_secs(10);
pub struct InProcessCodeModeSession {
runtime: SessionRuntime,
cell_execution_limits: CodeModeSessionCellExecutionLimits,
options: CodeModeSessionOptions,
}
impl InProcessCodeModeSession {
@@ -42,24 +45,29 @@ impl InProcessCodeModeSession {
}
pub fn with_limits(cell_execution_limits: CodeModeSessionCellExecutionLimits) -> Self {
Self {
runtime: SessionRuntime::new(),
cell_execution_limits: CodeModeSessionCellExecutionLimits {
max_heap_size_bytes: None,
..cell_execution_limits
Self::with_options(
CodeModeSessionOptions {
cell_execution_limits,
..CodeModeSessionOptions::default()
},
}
/*task_failure_handler*/ None,
)
}
pub fn with_task_failure_handler(
task_failure_handler: Arc<dyn Fn(String) + Send + Sync>,
cell_execution_limits: CodeModeSessionCellExecutionLimits,
/// Select output handling at creation. A session's mode must remain fixed
/// for all of its cells and observations.
pub fn with_options(
options: CodeModeSessionOptions,
task_failure_handler: Option<TaskFailureHandler>,
) -> Self {
Self {
runtime: SessionRuntime::new_with_task_failure_handler(Some(task_failure_handler)),
cell_execution_limits: CodeModeSessionCellExecutionLimits {
max_heap_size_bytes: None,
..cell_execution_limits
runtime: SessionRuntime::new_with_task_failure_handler(task_failure_handler),
options: CodeModeSessionOptions {
cell_execution_limits: CodeModeSessionCellExecutionLimits {
max_heap_size_bytes: None,
..options.cell_execution_limits
},
..options
},
}
}
@@ -71,10 +79,12 @@ impl InProcessCodeModeSession {
preempt: Option<CancellationToken>,
) -> Result<StartedCell, String> {
let yield_time_ms = request.yield_time_ms.unwrap_or(DEFAULT_EXEC_YIELD_TIME_MS);
let max_output_tokens = request.max_output_tokens;
let output_truncation = self.options.output_truncation;
let started = self
.runtime
.execute(
runtime_request(request),
runtime_request(request, self.options.output_truncation),
runtime::ObserveMode::YieldAfter(self.resolve_yield_timeout(yield_time_ms))
.with_yield_signal(preempt.unwrap_or_default()),
Arc::new(ProtocolDelegate { delegate }),
@@ -89,7 +99,14 @@ impl InProcessCodeModeSession {
.initial_event()
.await
.map_err(|error| error.to_string())
.and_then(|event| runtime_response(&response_cell_id, event));
.and_then(|event| {
runtime_response(
&response_cell_id,
event,
output_truncation,
max_output_tokens,
)
});
let _ = response_tx.send(response);
});
Ok(StartedCell::from_result_receiver(cell_id, response_rx))
@@ -100,10 +117,11 @@ impl InProcessCodeModeSession {
request: ExecuteRequest,
delegate: Arc<dyn CodeModeSessionDelegate>,
) -> Result<ExecuteToPendingOutcome, String> {
let max_output_tokens = request.max_output_tokens;
let started = self
.runtime
.execute(
runtime_request(request),
runtime_request(request, self.options.output_truncation),
runtime::ObserveMode::PendingFrontier,
Arc::new(ProtocolDelegate { delegate }),
)
@@ -114,7 +132,12 @@ impl InProcessCodeModeSession {
.initial_event()
.await
.map_err(|error| error.to_string())?;
pending_outcome(&cell_id, event)
pending_outcome(
&cell_id,
event,
self.options.output_truncation,
max_output_tokens,
)
}
pub async fn wait(
@@ -131,10 +154,11 @@ impl InProcessCodeModeSession {
preempt: Option<CancellationToken>,
) -> CodeModeSessionResultFuture<'static, WaitOutcome> {
let WaitRequest {
max_output_tokens: _,
cell_id,
yield_time_ms,
max_output_tokens,
} = request;
let output_truncation = self.options.output_truncation;
let runtime_cell_id = runtime_cell_id(&cell_id);
match self
.runtime
@@ -147,26 +171,47 @@ impl InProcessCodeModeSession {
{
Ok(pending_event) => Box::pin(async move {
match pending_event.event().await {
Ok(event) => Ok(WaitOutcome::LiveCell(runtime_response(&cell_id, event)?)),
Ok(event) => Ok(WaitOutcome::LiveCell(runtime_response(
&cell_id,
event,
output_truncation,
max_output_tokens,
)?)),
Err(runtime::Error::MissingCell(_) | runtime::Error::ClosedCell(_)) => {
Ok(WaitOutcome::MissingCell(missing_cell_response(cell_id)))
Ok(WaitOutcome::MissingCell(missing_cell_response(
cell_id,
output_truncation,
max_output_tokens,
)?))
}
Err(error) => Err(error.to_string()),
}
}),
Err(runtime::Error::MissingCell(_) | runtime::Error::ClosedCell(_)) => {
missing_wait(cell_id)
missing_wait(cell_id, output_truncation, max_output_tokens)
}
Err(error) => Box::pin(async move { Err(error.to_string()) }),
}
}
pub async fn terminate(&self, request: TerminateRequest) -> Result<WaitOutcome, String> {
let cell_id = request.cell_id;
let TerminateRequest {
cell_id,
max_output_tokens,
} = request;
match self.runtime.terminate(&runtime_cell_id(&cell_id)).await {
Ok(event) => Ok(WaitOutcome::LiveCell(runtime_response(&cell_id, event)?)),
Ok(event) => Ok(WaitOutcome::LiveCell(runtime_response(
&cell_id,
event,
self.options.output_truncation,
max_output_tokens,
)?)),
Err(runtime::Error::MissingCell(_) | runtime::Error::ClosedCell(_)) => {
Ok(WaitOutcome::MissingCell(missing_cell_response(cell_id)))
Ok(WaitOutcome::MissingCell(missing_cell_response(
cell_id,
self.options.output_truncation,
max_output_tokens,
)?))
}
Err(error) => Err(error.to_string()),
}
@@ -186,11 +231,18 @@ impl InProcessCodeModeSession {
.await
{
Ok(event) => Ok(WaitToPendingOutcome::LiveCell(pending_outcome(
&cell_id, event,
&cell_id,
event,
self.options.output_truncation,
/*max_output_tokens*/ None,
)?)),
Err(runtime::Error::MissingCell(_) | runtime::Error::ClosedCell(_)) => Ok(
WaitToPendingOutcome::MissingCell(missing_cell_response(cell_id)),
),
Err(runtime::Error::MissingCell(_) | runtime::Error::ClosedCell(_)) => {
Ok(WaitToPendingOutcome::MissingCell(missing_cell_response(
cell_id,
self.options.output_truncation,
/*max_output_tokens*/ None,
)?))
}
Err(error) => Err(error.to_string()),
}
}
@@ -210,7 +262,8 @@ impl InProcessCodeModeSession {
yield_time
};
self.cell_execution_limits
self.options
.cell_execution_limits
.max_yield_time_ms
.map(Duration::from_millis)
.map_or(timeout, |limit| timeout.min(limit))
@@ -224,6 +277,10 @@ impl Default for InProcessCodeModeSession {
}
impl CodeModeSession for InProcessCodeModeSession {
fn output_truncation(&self) -> CodeModeOutputTruncation {
self.options.output_truncation
}
fn execute<'a>(
&'a self,
request: ExecuteRequest,
@@ -318,9 +375,12 @@ impl runtime::SessionRuntimeDelegate for ProtocolDelegate {
}
}
fn runtime_request(request: ExecuteRequest) -> runtime::CreateCellRequest {
fn runtime_request(
request: ExecuteRequest,
output_truncation: CodeModeOutputTruncation,
) -> runtime::CreateCellRequest {
runtime::CreateCellRequest {
output_truncation: CodeModeOutputTruncation::Client,
output_truncation,
tool_call_id: request.tool_call_id,
enabled_tools: request
.enabled_tools
@@ -353,6 +413,8 @@ fn protocol_cell_id(cell_id: &runtime::CellId) -> CellId {
fn pending_outcome(
cell_id: &CellId,
event: runtime::CellEvent,
output_truncation: CodeModeOutputTruncation,
max_output_tokens: Option<usize>,
) -> Result<ExecuteToPendingOutcome, String> {
match event {
runtime::CellEvent::Pending {
@@ -360,13 +422,14 @@ fn pending_outcome(
pending_tool_call_ids,
} => Ok(ExecuteToPendingOutcome::Pending {
cell_id: cell_id.clone(),
content_items: content_items
.into_collected_output()
.into_untruncated_items(),
content_items: rendered_items(content_items, output_truncation, max_output_tokens),
pending_tool_call_ids,
}),
event => Ok(ExecuteToPendingOutcome::Completed(runtime_response(
cell_id, event,
cell_id,
event,
output_truncation,
max_output_tokens,
)?)),
}
}
@@ -374,13 +437,13 @@ fn pending_outcome(
fn runtime_response(
cell_id: &CellId,
event: runtime::CellEvent,
output_truncation: CodeModeOutputTruncation,
max_output_tokens: Option<usize>,
) -> Result<RuntimeResponse, String> {
match event {
runtime::CellEvent::Yielded { content_items } => Ok(RuntimeResponse::Yielded {
cell_id: cell_id.clone(),
content_items: content_items
.into_collected_output()
.into_untruncated_items(),
content_items: rendered_items(content_items, output_truncation, max_output_tokens),
code_mode_host_duration: None,
}),
runtime::CellEvent::Completed {
@@ -389,18 +452,14 @@ fn runtime_response(
result_buffered_for,
} => Ok(RuntimeResponse::Result {
cell_id: cell_id.clone(),
content_items: content_items
.into_collected_output()
.into_untruncated_items(),
content_items: rendered_items(content_items, output_truncation, max_output_tokens),
error_text,
result_buffered_for,
code_mode_host_duration: None,
}),
runtime::CellEvent::Terminated { content_items } => Ok(RuntimeResponse::Terminated {
cell_id: cell_id.clone(),
content_items: content_items
.into_collected_output()
.into_untruncated_items(),
content_items: rendered_items(content_items, output_truncation, max_output_tokens),
code_mode_host_duration: None,
}),
runtime::CellEvent::Pending { .. } => {
@@ -409,18 +468,41 @@ fn runtime_response(
}
}
fn missing_cell_response(cell_id: CellId) -> RuntimeResponse {
RuntimeResponse::Result {
error_text: Some(format!("exec cell {cell_id} not found")),
result_buffered_for: Duration::ZERO,
cell_id,
content_items: Vec::new(),
code_mode_host_duration: None,
fn rendered_items(
items: OutputCollector,
mode: CodeModeOutputTruncation,
max_output_tokens: Option<usize>,
) -> Vec<FunctionCallOutputContentItem> {
let items = items.into_collected_output();
match mode {
CodeModeOutputTruncation::Client => items.into_untruncated_items(),
CodeModeOutputTruncation::Server => crate::output::render_output(items, max_output_tokens),
}
}
fn missing_wait(cell_id: CellId) -> CodeModeSessionResultFuture<'static, WaitOutcome> {
Box::pin(async move { Ok(WaitOutcome::MissingCell(missing_cell_response(cell_id))) })
fn missing_cell_response(
cell_id: CellId,
mode: CodeModeOutputTruncation,
max_output_tokens: Option<usize>,
) -> Result<RuntimeResponse, String> {
let error_text = format!("exec cell {cell_id} not found");
let mut output = OutputCollector::new(mode);
let event = take_completed_event(&mut output, Some(error_text));
runtime_response(&cell_id, event, mode, max_output_tokens)
}
fn missing_wait(
cell_id: CellId,
mode: CodeModeOutputTruncation,
max_output_tokens: Option<usize>,
) -> CodeModeSessionResultFuture<'static, WaitOutcome> {
Box::pin(async move {
Ok(WaitOutcome::MissingCell(missing_cell_response(
cell_id,
mode,
max_output_tokens,
)?))
})
}
#[cfg(test)]
@@ -430,3 +512,7 @@ mod tests;
#[cfg(test)]
#[path = "service_contract_tests.rs"]
mod contract_tests;
#[cfg(test)]
#[path = "service_output_collection_tests.rs"]
mod output_collection_tests;
codex-rs/code-mode-runtime/src/service_audio_tests.rs
+46 -1
View File
@@ -1,4 +1,4 @@
//! Covers per-item short-WAV omissions without disrupting other Code Mode outputs.
//! Covers short-WAV handling and Server audio output alongside other Code Mode outputs.
use super::cell_id;
use super::execute;
@@ -11,6 +11,9 @@ use crate::WaitOutcome;
use crate::WaitRequest;
use base64::Engine;
use base64::engine::general_purpose::STANDARD as BASE64_STANDARD;
use codex_code_mode_protocol::CodeModeOutputTruncation;
use codex_code_mode_protocol::CodeModeSessionOptions;
use codex_code_mode_protocol::host::BOUNDED_OUTPUT_MAX_BYTES;
use pretty_assertions::assert_eq;
const SHORT_AUDIO_OMISSION_TEXT: &str =
@@ -168,3 +171,45 @@ async fn audio_helper_bounds_wav_data_and_preserves_outputs_after_yield() {
})
);
}
#[tokio::test]
async fn server_omits_audio_above_the_text_allowance_at_zero_tokens() {
let session = InProcessCodeModeSession::with_options(
CodeModeSessionOptions {
output_truncation: CodeModeOutputTruncation::Server,
..CodeModeSessionOptions::default()
},
/*task_failure_handler*/ None,
);
// Each clip exceeds the text allowance, but the clips fit the separate
// audio allowance and reach the zero-token formatter.
let wav = pcm16_wav(/*sample_rate*/ 24_000, BOUNDED_OUTPUT_MAX_BYTES / 2);
let payload = BASE64_STANDARD.encode(wav);
let source = format!(
"const clip = 'data:audio/wav;base64,{payload}'; for (let i = 0; i < 6; i++) {{ audio(clip); text(''); }}"
);
let response = execute(
&session,
ExecuteRequest {
yield_time_ms: Some(60_000),
max_output_tokens: Some(0),
..execute_request(&source)
},
)
.await;
let RuntimeResponse::Result {
content_items,
error_text,
..
} = response
else {
panic!("expected completed cell");
};
assert_eq!(error_text, None);
assert_eq!(
content_items,
[FunctionCallOutputContentItem::InputText {
text: "[omitted 6 audio items ...]".to_owned(),
}]
);
}
codex-rs/code-mode-runtime/src/service_output_collection_tests.rs
+340
View File
@@ -0,0 +1,340 @@
//! Real-runtime regressions for bounded output, observation and termination.
use std::sync::Arc;
use std::time::Duration;
use codex_code_mode_protocol::ExecuteRequest;
use codex_code_mode_protocol::FunctionCallOutputContentItem;
use codex_code_mode_protocol::NoopCodeModeSessionDelegate;
use codex_code_mode_protocol::RuntimeResponse;
use codex_code_mode_protocol::WaitOutcome;
use codex_code_mode_protocol::WaitRequest;
use pretty_assertions::assert_eq;
use super::InProcessCodeModeSession;
use crate::output::MAX_OUTPUT_ITEMS;
use crate::output::MAX_SERVER_DIAGNOSTIC_BYTES;
use codex_code_mode_protocol::CodeModeOutputTruncation;
use codex_code_mode_protocol::CodeModeSessionOptions;
use codex_code_mode_protocol::host::BOUNDED_OUTPUT_MAX_BYTES as MAX_OUTPUT_BYTES;
use codex_code_mode_protocol::host::BOUNDED_OUTPUT_MAX_IMAGE_BYTES as MAX_IMAGE_BYTES;
fn server_session() -> InProcessCodeModeSession {
InProcessCodeModeSession::with_options(
CodeModeSessionOptions {
output_truncation: CodeModeOutputTruncation::Server,
..CodeModeSessionOptions::default()
},
/*task_failure_handler*/ None,
)
}
async fn start(session: &InProcessCodeModeSession, source: &str) -> RuntimeResponse {
session
.execute(
ExecuteRequest {
tool_call_id: "bounded-output".to_owned(),
enabled_tools: Vec::new(),
source: source.to_owned(),
yield_time_ms: Some(60_000),
// These tests isolate collector overflow from the per-observation
// formatter: request enough room to display every retained item.
max_output_tokens: Some(MAX_OUTPUT_BYTES),
},
Arc::new(NoopCodeModeSessionDelegate),
/*preempt*/ None,
)
.await
.unwrap()
.initial_response()
.await
.unwrap()
}
#[tokio::test]
async fn bounded_output_keeps_head_tail_without_stopping_the_cell() {
let session = server_session();
let emitted = MAX_OUTPUT_ITEMS + 100;
let response = start(
&session,
&format!("text('begin'); for (let i = 0; i < {emitted}; i++) text(i); text('end');"),
)
.await;
let RuntimeResponse::Result {
content_items,
error_text,
..
} = response
else {
panic!("expected completed cell");
};
assert_eq!(error_text, None);
assert_eq!(content_items.len(), MAX_OUTPUT_ITEMS);
assert_eq!(
content_items.first(),
Some(&FunctionCallOutputContentItem::InputText {
text: "begin".to_owned()
})
);
assert_eq!(
content_items.last(),
Some(&FunctionCallOutputContentItem::InputText {
text: "end".to_owned()
})
);
assert!(matches!(
&content_items[MAX_OUTPUT_ITEMS / 2],
FunctionCallOutputContentItem::InputText { text }
if text.starts_with("[Code Mode session output limit reached: omitted ")
));
}
#[tokio::test]
async fn huge_output_is_bounded_and_yield_drain_starts_a_new_budget() {
let session = server_session();
// Continue immediately after yielding to exercise the FIFO boundary as
// well as discarding the already-delivered bounded snapshot.
let response = start(
&session,
r#"text('😀'.repeat(1000000)); yield_control(); text('after yield');"#,
)
.await;
let RuntimeResponse::Yielded {
cell_id,
content_items,
..
} = response
else {
panic!("expected yielded cell");
};
assert_eq!(
content_items,
vec![
FunctionCallOutputContentItem::InputText {
text: "😀".repeat(MAX_OUTPUT_BYTES / 8)
},
FunctionCallOutputContentItem::InputText {
text: format!(
"[Code Mode session output limit reached: omitted {} text bytes; discarded 0 text segments, 0 images, 0 audio items]",
4_000_000 - MAX_OUTPUT_BYTES
)
},
FunctionCallOutputContentItem::InputText {
text: "😀".repeat(MAX_OUTPUT_BYTES / 8)
},
]
);
let response = session
.wait(
WaitRequest {
cell_id,
yield_time_ms: 60_000,
max_output_tokens: None,
},
/*preempt*/ None,
)
.await
.unwrap();
let WaitOutcome::LiveCell(RuntimeResponse::Result {
content_items,
error_text,
..
}) = response
else {
panic!("expected completed cell");
};
assert_eq!(
(content_items, error_text),
(
vec![FunctionCallOutputContentItem::InputText {
text: "after yield".to_owned()
}],
None
)
);
}
#[tokio::test]
async fn a_single_oversized_image_is_omitted_without_losing_surrounding_text() {
let session = server_session();
let response = start(
&session,
&format!("text('before'); image('data:image/png;base64,' + 'A'.repeat({MAX_IMAGE_BYTES})); text('after');"),
)
.await;
let RuntimeResponse::Result {
content_items,
error_text,
..
} = response
else {
panic!("expected completed cell");
};
assert_eq!(
(content_items, error_text),
(
vec![
FunctionCallOutputContentItem::InputText {
text: "before".to_owned()
},
FunctionCallOutputContentItem::InputText {
text: "[Code Mode session output limit reached: omitted 0 text bytes; discarded 0 text segments, 1 images, 0 audio items]".to_owned()
},
FunctionCallOutputContentItem::InputText {
text: "after".to_owned()
}
],
None
)
);
}
#[tokio::test]
async fn termination_releases_a_flooding_bounded_runtime() {
let session = server_session();
let emitted = MAX_OUTPUT_ITEMS + 100;
let response = start(
&session,
&format!(
"for (let i = 0; ; i++) {{
text('output');
if (i === {emitted}) yield_control();
}}"
),
)
.await;
let RuntimeResponse::Yielded {
cell_id,
content_items,
..
} = response
else {
panic!("expected yielded cell");
};
assert_eq!(content_items.len(), MAX_OUTPUT_ITEMS);
assert!(matches!(
&content_items[MAX_OUTPUT_ITEMS / 2],
FunctionCallOutputContentItem::InputText { text }
if text.starts_with("[Code Mode session output limit reached: omitted ")
));
let response = tokio::time::timeout(Duration::from_secs(5), session.terminate(cell_id.into()))
.await
.expect("termination must unblock output")
.unwrap();
assert!(matches!(
response,
WaitOutcome::LiveCell(RuntimeResponse::Terminated { .. })
));
session.shutdown().await.unwrap();
}
#[tokio::test]
async fn terminal_error_shares_the_collector_limit_and_remains_a_failure() {
let session = server_session();
let response = start(
&session,
"text('before'); throw 'error-start:' + 'x'.repeat(2 * 1024 * 1024) + ':error-end';",
)
.await;
let RuntimeResponse::Result {
content_items,
error_text,
..
} = response
else {
panic!("expected failed cell");
};
let diagnostic = error_text.expect("overflowed diagnostics still report failure");
// The model-output budget is generous, but the diagnostic has its own limit.
assert!(diagnostic.len() <= MAX_SERVER_DIAGNOSTIC_BYTES);
assert!(diagnostic.starts_with("error-start:"));
assert!(diagnostic.ends_with(":error-end"));
let notice = content_items
.iter()
.find_map(|item| match item {
FunctionCallOutputContentItem::InputText { text }
if text.starts_with("[Code Mode session output limit reached: omitted ") =>
{
Some(text)
}
_ => None,
})
.expect("terminal overflow notice");
let total_text_bytes: usize = content_items
.iter()
.map(|item| match item {
FunctionCallOutputContentItem::InputText { text } => text.len(),
_ => 0,
})
.sum();
assert!(total_text_bytes - notice.len() <= MAX_OUTPUT_BYTES);
assert!(matches!(content_items.last(),
Some(FunctionCallOutputContentItem::InputText { text }) if text.ends_with(":error-end")
));
session.shutdown().await.unwrap();
}
#[tokio::test]
async fn wait_preserves_diagnostics_across_output_budgets() {
let source = "yield_control(); text('before'); throw 'start:' + '🎷'.repeat(4096) + ':end';";
let raw_error = format!("start:{}:end", "🎷".repeat(4096));
for mode in [
CodeModeOutputTruncation::Client,
CodeModeOutputTruncation::Server,
] {
let session = InProcessCodeModeSession::with_options(
CodeModeSessionOptions {
output_truncation: mode,
..CodeModeSessionOptions::default()
},
/*task_failure_handler*/ None,
);
// The error fits the diagnostic limit but exceeds both output budgets.
for max_output_tokens in [0, 16] {
let response = start(&session, source).await;
let RuntimeResponse::Yielded { cell_id, .. } = response else {
panic!("expected yield before error");
};
let WaitOutcome::LiveCell(RuntimeResponse::Result {
content_items,
error_text,
..
}) = session
.wait(
WaitRequest {
cell_id,
yield_time_ms: 60_000,
max_output_tokens: Some(max_output_tokens),
},
/*preempt*/ None,
)
.await
.unwrap()
else {
panic!("expected completed cell");
};
let diagnostic = error_text.expect("errors remain failures at all budgets");
assert_eq!(diagnostic, raw_error);
match mode {
CodeModeOutputTruncation::Client => {
assert_eq!(
content_items,
vec![FunctionCallOutputContentItem::InputText {
text: "before".to_owned()
}]
);
}
CodeModeOutputTruncation::Server => {
if max_output_tokens != 0 {
assert!(matches!(
content_items.first(),
Some(FunctionCallOutputContentItem::InputText { text })
if text.starts_with("Warning: truncated output")
));
}
}
}
}
session.shutdown().await.unwrap();
}
}
codex-rs/code-mode-runtime/src/service_tests.rs
+1 -1
View File
@@ -999,7 +999,7 @@ store("healthy", {[keys[0]]: "null"});
vec![FunctionCallOutputContentItem::InputText {
text: expected_text.to_string()
}],
expected_error.map(str::to_string)
expected_error.map(str::to_string),
)
);
}
codex-rs/code-mode-runtime/src/session_runtime/mod.rs
+1
View File
@@ -54,6 +54,7 @@ struct Inner {
}
impl SessionRuntime {
#[cfg(test)]
pub(crate) fn new() -> Self {
Self::new_with_task_failure_handler(/*task_failure_handler*/ None)
}
codex-rs/code-mode-runtime/src/session_runtime/types.rs
+2
View File
@@ -68,6 +68,8 @@ pub(crate) enum CellEvent {
},
Completed {
content_items: OutputCollector,
/// None is success. Server diagnostics have already been bounded
/// independently of the observing request's output budget.
error_text: Option<String>,
/// Time the completed result spent waiting for an observer.
result_buffered_for: Duration,