openai/codex
1
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-11-14 21:40:39 +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/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;