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/Cargo.lock
Generated
+1
View File
@@ -2805,6 +2805,7 @@ dependencies = [
"codex-protocol",
"codex-utils-audio",
"codex-utils-output-truncation",
"codex-utils-string",
"deno_core_icudata",
"futures",
"opentelemetry",
codex-rs/code-mode-host/src/grpc/conversions.rs
+16
View File
@@ -1,3 +1,4 @@
use codex_code_mode_protocol::CodeModeOutputTruncation;
use codex_code_mode_protocol::CodeModeSessionCellExecutionLimits;
use codex_code_mode_protocol::CodeModeToolKind;
use codex_code_mode_protocol::ExecuteRequest;
@@ -29,6 +30,21 @@ pub(super) fn session_limits(
})
}
pub(super) fn observation_budget(
max_output_tokens: Option<u64>,
output_truncation: CodeModeOutputTruncation,
) -> Result<Option<usize>, Status> {
if output_truncation == CodeModeOutputTruncation::Client {
return Ok(None);
}
max_output_tokens
.map(|value| {
usize::try_from(value)
.map_err(|_| Status::invalid_argument("output budget is too large"))
})
.transpose()
}
pub(super) fn execute_request(request: proto::ExecuteRequest) -> Result<ExecuteRequest, Status> {
validation::identifier(&request.tool_call_id, "tool call ID")?;
Ok(ExecuteRequest {
codex-rs/code-mode-host/src/grpc/mod.rs
+36 -3
View File
@@ -12,6 +12,10 @@ use std::sync::Arc;
use std::time::Instant;
use codex_code_mode_protocol::CellId;
use codex_code_mode_protocol::CodeModeOutputTruncation;
use codex_code_mode_protocol::CodeModeSession;
use codex_code_mode_protocol::CodeModeSessionOptions;
use codex_code_mode_protocol::TerminateRequest;
use codex_code_mode_protocol::WaitRequest;
use codex_code_mode_protocol::grpc as proto;
use codex_code_mode_protocol::grpc::code_mode_host_server::CodeModeHost;
@@ -75,7 +79,23 @@ impl GrpcCodeModeHost {
) -> Result<Response<GrpcStream<proto::SessionEvent>>, Status> {
let _permit = self.state.request_permit()?;
let limits = conversions::session_limits(request.cell_execution_limits)?;
Ok(Response::new(self.state.open_session(limits)?))
let output_truncation = match request
.output_truncation
.map(proto::OutputTruncation::try_from)
.transpose()
.map_err(|_| Status::invalid_argument("unknown code-mode output truncation mode"))?
{
None | Some(proto::OutputTruncation::Unspecified | proto::OutputTruncation::Client) => {
CodeModeOutputTruncation::Client
}
Some(proto::OutputTruncation::Server) => CodeModeOutputTruncation::Server,
};
Ok(Response::new(self.state.open_session(
CodeModeSessionOptions {
cell_execution_limits: limits,
output_truncation,
},
)?))
}
#[tracing::instrument(
@@ -242,11 +262,15 @@ impl GrpcCodeModeHost {
validation::identifier(&request.cell_id, "cell ID")?;
validation::identifier(&request.wait_id, "wait ID")?;
let _permit = self.state.request_permit()?;
let max_output_tokens = conversions::observation_budget(
request.max_output_tokens,
session.runtime.output_truncation(),
)?;
let registration = WaitRegistration::new(Arc::clone(&session), request.wait_id)?;
let request = WaitRequest {
cell_id: CellId::new(request.cell_id),
yield_time_ms: request.yield_time_ms,
max_output_tokens: None,
max_output_tokens,
};
let outcome = tokio::select! {
biased;
@@ -301,7 +325,16 @@ impl GrpcCodeModeHost {
let session = self.state.session(&request.session_id)?;
validation::identifier(&request.cell_id, "cell ID")?;
let _permit = self.state.request_permit()?;
let outcome = session.terminate(CellId::new(request.cell_id)).await?;
let max_output_tokens = conversions::observation_budget(
request.max_output_tokens,
session.runtime.output_truncation(),
)?;
let outcome = session
.terminate(TerminateRequest {
cell_id: CellId::new(request.cell_id),
max_output_tokens,
})
.await?;
let outcome = outcome.with_code_mode_host_duration(received_at.elapsed());
conversions::wait_response(outcome).map_err(|error| Status::internal(error.to_string()))
}
codex-rs/code-mode-host/src/grpc/service_tests.rs
+178 -4
View File
@@ -11,10 +11,18 @@ use tonic::Request;
pub(super) async fn open_session(
host: &GrpcCodeModeHost,
) -> (String, GrpcStream<proto::SessionEvent>) {
let (opened, stream) = open_session_with_truncation(host, /*output_truncation*/ None).await;
(opened.session_id, stream)
}
async fn open_session_with_truncation(
host: &GrpcCodeModeHost,
output_truncation: Option<i32>,
) -> (proto::SessionOpened, GrpcStream<proto::SessionEvent>) {
let mut stream = host
.open_session(Request::new(proto::OpenSessionRequest {
cell_execution_limits: None,
output_truncation: None,
output_truncation,
}))
.await
.expect("open code-mode session")
@@ -27,7 +35,7 @@ pub(super) async fn open_session(
let Some(proto::session_event::Event::Opened(opened)) = event.event else {
panic!("expected the first session event to open its lease");
};
(opened.session_id, stream)
(opened, stream)
}
pub(super) fn execute_request(
@@ -80,6 +88,18 @@ pub(super) async fn execute_events(
(started.cell_id, stream)
}
fn only_text(items: &[proto::ContentItem]) -> &str {
let [
proto::ContentItem {
item: Some(proto::content_item::Item::Text(text)),
},
] = items
else {
panic!("expected one text item, got {items:?}");
};
&text.text
}
#[tokio::test]
async fn execute_stream_starts_immediately_and_wait_preserves_missing_cells() {
let host = GrpcCodeModeHost::new();
@@ -358,7 +378,7 @@ async fn cancellation_before_wait_admission_is_preserved() {
cell_id: cell_id.clone(),
wait_id: "pre-cancelled".to_string(),
yield_time_ms: 60_000,
max_output_tokens: None,
max_output_tokens: Some(0),
}))
.await
.unwrap_err();
@@ -367,7 +387,7 @@ async fn cancellation_before_wait_admission_is_preserved() {
host.terminate(Request::new(proto::TerminateRequest {
session_id,
cell_id,
max_output_tokens: None,
max_output_tokens: Some(0),
}))
.await
.unwrap();
@@ -440,3 +460,157 @@ async fn notifications_do_not_delay_cell_completion() {
.await
.expect("legacy notification acknowledgments remain accepted");
}
#[tokio::test]
async fn output_truncation_is_a_session_setting_and_omitting_it_keeps_legacy_output() {
for (output_truncation, bounded_output) in [
(None, false),
(Some(proto::OutputTruncation::Unspecified.into()), false),
(Some(proto::OutputTruncation::Client.into()), false),
(Some(proto::OutputTruncation::Server.into()), true),
] {
let host = GrpcCodeModeHost::new();
let (opened, _lease) = open_session_with_truncation(&host, output_truncation).await;
assert_eq!(
opened.output_truncation,
bounded_output.then_some(proto::OutputTruncation::Server.into()),
);
let mut request = execute_request(
&opened.session_id,
"execution",
"text('x'.repeat(2 * 1024 * 1024));",
);
request.yield_time_ms = Some(60_000);
let (_, mut execution) = execute_events(&host, request).await;
let outcome = execution
.next()
.await
.expect("outcome event")
.expect("outcome");
let Some(proto::execute_event::Event::Outcome(outcome)) = outcome.event else {
panic!("expected execution outcome");
};
assert!(matches!(
outcome.outcome,
Some(proto::execution_outcome::Outcome::Completed(_))
));
assert_eq!(
outcome.content_items.iter().any(
|item| matches!(&item.item, Some(proto::content_item::Item::Text(text))
if text.text.starts_with("[Code Mode session output limit reached: omitted "))
),
bounded_output
);
if !bounded_output {
assert_eq!(
only_text(&outcome.content_items),
"x".repeat(2 * 1024 * 1024)
);
}
}
}
#[tokio::test]
async fn server_output_budgets_follow_exec_wait_and_terminate_over_grpc() {
for (execution_id, after_yield) in [
("failed", "throw '0123456789'.repeat(8);"),
(
"terminated",
"text('0123456789'.repeat(8)); notify('ready to terminate'); await new Promise(() => {});",
),
] {
let host = GrpcCodeModeHost::new();
let (opened, mut lease) =
open_session_with_truncation(&host, Some(proto::OutputTruncation::Server.into())).await;
let source = format!("text('0123456789'.repeat(4)); yield_control(); {after_yield}");
let mut request = execute_request(&opened.session_id, execution_id, &source);
request.yield_time_ms = Some(60_000);
request.max_output_tokens = Some(5);
let (cell_id, mut execution) = execute_events(&host, request).await;
let Some(proto::execute_event::Event::Outcome(first)) =
execution.next().await.unwrap().unwrap().event
else {
panic!("expected execution outcome");
};
assert_eq!(first.cell_id, cell_id);
assert!(matches!(
first.outcome,
Some(proto::execution_outcome::Outcome::Yielded(_))
));
assert_eq!(
only_text(&first.content_items),
"Warning: truncated output (original token count: 10)\nTotal output lines: 1\n\n0123456789…5 tokens truncated…0123456789"
);
let response = if execution_id == "failed" {
host.wait(Request::new(proto::WaitRequest {
session_id: opened.session_id,
cell_id: cell_id.clone(),
wait_id: "wait-after-yield".to_string(),
yield_time_ms: 60_000,
max_output_tokens: Some(7),
}))
.await
.expect("script failures are completed results, not RPC failures")
.into_inner()
} else {
// The event is ordered after text in the runtime FIFO, so the
// termination observes it without a sleep or racing the cell.
assert!(matches!(
lease.next().await.unwrap().unwrap().event,
Some(proto::session_event::Event::Notification(
proto::Notification { text, .. }
)) if text == "ready to terminate"
));
host.terminate(Request::new(proto::TerminateRequest {
session_id: opened.session_id,
cell_id: cell_id.clone(),
max_output_tokens: Some(7),
}))
.await
.expect("terminate with an independent output budget")
.into_inner()
};
let Some(proto::wait_response::State::LiveCell(last)) = response.state else {
panic!("expected final observation from the live cell");
};
let expected_text = if execution_id == "failed" {
let Some(proto::execution_outcome::Outcome::Completed(proto::ExecutionCompleted {
error_text: Some(diagnostic),
})) = &last.outcome
else {
panic!(
"expected completed failure with a diagnostic, got {:?}",
last.outcome
);
};
assert_eq!(diagnostic, &"0123456789".repeat(8));
"Warning: truncated output (original token count: 24)\nTotal output lines: 2\n\nScript error:\n…17 tokens truncated…67890123456789"
} else {
assert_eq!(
last.outcome,
Some(proto::execution_outcome::Outcome::Terminated(
proto::ExecutionTerminated {}
))
);
"Warning: truncated output (original token count: 20)\nTotal output lines: 1\n\n01234567890123…13 tokens truncated…67890123456789"
};
assert_eq!(last.cell_id, cell_id);
assert_eq!(only_text(&last.content_items), expected_text);
}
}
#[tokio::test]
async fn unknown_output_truncation_is_rejected_before_opening_a_session() {
let host = GrpcCodeModeHost::new();
let result = host
.open_session(Request::new(proto::OpenSessionRequest {
cell_execution_limits: None,
output_truncation: Some(42),
}))
.await;
let Err(status) = result else {
panic!("unknown truncation mode should be rejected");
};
assert_eq!(status.code(), Code::InvalidArgument);
}
codex-rs/code-mode-host/src/grpc/session.rs
+19 -11
View File
@@ -10,7 +10,10 @@ use std::sync::PoisonError;
use std::sync::Weak;
use codex_code_mode_protocol::CellId;
use codex_code_mode_protocol::CodeModeSessionCellExecutionLimits;
use codex_code_mode_protocol::CodeModeOutputTruncation;
use codex_code_mode_protocol::CodeModeSession;
use codex_code_mode_protocol::CodeModeSessionOptions;
use codex_code_mode_protocol::TerminateRequest;
use codex_code_mode_protocol::WaitOutcome;
use codex_code_mode_protocol::grpc as proto;
use codex_code_mode_protocol::host::MAX_PENDING_DELEGATE_CALLS;
@@ -107,7 +110,7 @@ impl GrpcHostState {
pub(super) fn open_session(
self: &Arc<Self>,
limits: CodeModeSessionCellExecutionLimits,
options: CodeModeSessionOptions,
) -> Result<GrpcStream<proto::SessionEvent>, Status> {
let id = Uuid::new_v4();
let (events, receiver) = mpsc::channel(OUTGOING_CHANNEL_CAPACITY);
@@ -118,13 +121,18 @@ impl GrpcHostState {
event_sender,
closed,
Arc::clone(&self.delegate_permits),
limits,
options,
);
events
.try_send(Ok(proto::SessionEvent {
event: Some(proto::session_event::Event::Opened(proto::SessionOpened {
session_id: id.to_string(),
output_truncation: None,
output_truncation: match session.runtime.output_truncation() {
CodeModeOutputTruncation::Server => {
Some(proto::OutputTruncation::Server.into())
}
CodeModeOutputTruncation::Client => None,
},
})),
}))
.map_err(|_| Status::internal("failed to publish the opened code-mode session"))?;
@@ -213,7 +221,7 @@ impl GrpcSession {
events: EventSender,
closed: CancellationToken,
delegate_permits: Arc<Semaphore>,
limits: CodeModeSessionCellExecutionLimits,
options: CodeModeSessionOptions,
) -> Arc<Self> {
Arc::new_cyclic(|weak: &Weak<Self>| {
let failure_session = weak.clone();
@@ -225,9 +233,9 @@ impl GrpcSession {
});
Self {
id,
runtime: Arc::new(InProcessCodeModeSession::with_task_failure_handler(
failure_handler,
limits,
runtime: Arc::new(InProcessCodeModeSession::with_options(
options,
Some(failure_handler),
)),
closed,
state: Mutex::new(SessionState::default()),
@@ -258,13 +266,13 @@ impl GrpcSession {
result
}
pub(super) async fn terminate(&self, cell_id: CellId) -> Result<WaitOutcome, Status> {
pub(super) async fn terminate(&self, request: TerminateRequest) -> Result<WaitOutcome, Status> {
tokio::select! {
biased;
_ = self.closed.cancelled() => {
Err(Status::cancelled("code-mode session is closed"))
}
result = self.runtime.terminate(cell_id.into()) => {
result = self.runtime.terminate(request) => {
result.map_err(Status::failed_precondition)
}
}
@@ -347,7 +355,7 @@ impl GrpcSession {
if let Some(cell_id) = cell_id {
let session = Arc::clone(self);
tokio::spawn(async move {
let _ = session.terminate(CellId::new(cell_id)).await;
let _ = session.terminate(CellId::new(cell_id).into()).await;
});
}
}
codex-rs/code-mode-host/src/host_tests.rs
+226 -50
View File
@@ -9,7 +9,9 @@ use std::task::Poll;
use std::time::Duration;
use std::time::Instant;
use codex_code_mode_protocol::CodeModeSessionCellExecutionLimits;
use codex_code_mode_protocol::CodeModeSessionOptions;
use codex_code_mode_protocol::TerminateRequest;
use codex_code_mode_protocol::host::BOUNDED_OUTPUT_CAPABILITY;
use codex_code_mode_protocol::host::Capability;
use codex_code_mode_protocol::host::CapabilitySet;
use codex_code_mode_protocol::host::ClientHello;
@@ -27,9 +29,15 @@ use codex_code_mode_protocol::host::RequestId;
use codex_code_mode_protocol::host::SESSION_RESOURCE_LIMITS_CAPABILITY;
use codex_code_mode_protocol::host::SessionId;
use codex_code_mode_protocol::host::SupportedProtocolVersions;
use codex_code_mode_protocol::host::WireContentItem;
use codex_code_mode_protocol::host::WireExecuteRequest;
use codex_code_mode_protocol::host::WireOutputTruncation;
use codex_code_mode_protocol::host::WireResult;
use codex_code_mode_protocol::host::WireRuntimeResponse;
use codex_code_mode_protocol::host::WireWaitOutcome;
use codex_code_mode_protocol::host::WireWaitRequest;
use pretty_assertions::assert_eq;
use tokio::io::AsyncRead;
use tokio::io::AsyncWrite;
use tokio::sync::Semaphore;
use tokio::sync::mpsc;
@@ -96,6 +104,20 @@ fn execute_request(source: &str) -> WireExecuteRequest {
}
}
async fn next_response(reader: &mut FramedReader<impl AsyncRead + Unpin>) -> HostToClient {
loop {
match reader
.read()
.await
.expect("host response")
.expect("host connected")
{
HostToClient::CellClosed { .. } => {}
message => return message,
}
}
}
#[tokio::test]
async fn handshake_and_multiple_session_lifecycles_are_ordered() {
let (host_stream, client_stream) = tokio::io::duplex(/*max_buf_size*/ 4096);
@@ -258,52 +280,209 @@ impl AsyncWrite for BlockingWriter {
}
#[tokio::test]
async fn session_resource_limits_are_negotiated_when_optional_or_required() {
let resource_limits_capability =
Capability::new(SESSION_RESOURCE_LIMITS_CAPABILITY).expect("resource limits capability");
let resource_limits =
CapabilitySet::try_new([resource_limits_capability.clone()]).expect("host capabilities");
for (required, optional) in [
(
CapabilitySet::empty(),
CapabilitySet::try_new([resource_limits_capability]).expect("optional capabilities"),
),
(resource_limits.clone(), CapabilitySet::empty()),
async fn supported_capabilities_are_negotiated_when_optional_or_required() {
for capability in [
SESSION_RESOURCE_LIMITS_CAPABILITY,
BOUNDED_OUTPUT_CAPABILITY,
] {
let (host_stream, client_stream) = tokio::io::duplex(/*max_buf_size*/ 1024);
let (host_reader, host_writer) = tokio::io::split(host_stream);
let (client_reader, client_writer) = tokio::io::split(client_stream);
let host = tokio::spawn(run(host_reader, host_writer));
let mut reader = FramedReader::new(client_reader);
let mut writer = FramedWriter::new(client_writer);
let capability = Capability::new(capability).expect("supported capability");
let host_capabilities =
CapabilitySet::try_new([capability.clone()]).expect("host capabilities");
writer
.write(&ClientToHost::ClientHello(
ClientHello::new(
SupportedProtocolVersions::try_new([ProtocolVersion::V1])
.expect("supported versions"),
required,
optional,
)
.expect("client hello"),
))
.await
.expect("write hello");
assert_eq!(
reader.read::<HostToClient>().await.expect("host hello"),
Some(HostToClient::HostHello(HostHello::new(
ProtocolVersion::V1,
resource_limits.clone(),
)))
);
for (required, optional) in [
(
CapabilitySet::empty(),
CapabilitySet::try_new([capability]).expect("optional capabilities"),
),
(host_capabilities.clone(), CapabilitySet::empty()),
] {
let (host_stream, client_stream) = tokio::io::duplex(/*max_buf_size*/ 1024);
let (host_reader, host_writer) = tokio::io::split(host_stream);
let (client_reader, client_writer) = tokio::io::split(client_stream);
let host = tokio::spawn(run(host_reader, host_writer));
let mut reader = FramedReader::new(client_reader);
let mut writer = FramedWriter::new(client_writer);
drop(writer);
drop(reader);
host.await.expect("host task").expect("host connection");
writer
.write(&ClientToHost::ClientHello(
ClientHello::new(
SupportedProtocolVersions::try_new([ProtocolVersion::V1])
.expect("supported versions"),
required,
optional,
)
.expect("client hello"),
))
.await
.expect("write hello");
assert_eq!(
reader.read::<HostToClient>().await.expect("host hello"),
Some(HostToClient::HostHello(HostHello::new(
ProtocolVersion::V1,
host_capabilities.clone(),
)))
);
drop(writer);
drop(reader);
host.await.expect("host task").expect("host connection");
}
}
}
#[tokio::test]
async fn output_selection_is_per_session() {
let (host_stream, client_stream) = tokio::io::duplex(/*max_buf_size*/ 4096);
let (host_reader, host_writer) = tokio::io::split(host_stream);
let (client_reader, client_writer) = tokio::io::split(client_stream);
let host = tokio::spawn(run(host_reader, host_writer));
let mut reader = FramedReader::new(client_reader);
let mut writer = FramedWriter::new(client_writer);
writer
.write(&client_hello([ProtocolVersion::V1], CapabilitySet::empty()))
.await
.expect("write hello");
assert_eq!(
reader.read::<HostToClient>().await.expect("host hello"),
Some(HostToClient::HostHello(HostHello::new(
ProtocolVersion::V1,
CapabilitySet::empty(),
)))
);
for (id, name, output_truncation) in [
(1, "client", None),
(2, "server", Some(WireOutputTruncation::Server)),
] {
writer
.write(&ClientToHost::Request {
id: request_id(id),
request: HostRequest::OpenSession {
session_id: session_id(name),
cell_execution_limits: None,
output_truncation,
},
})
.await
.expect("open session");
assert_eq!(
reader.read::<HostToClient>().await.expect("session ready"),
Some(HostToClient::Response {
id: request_id(id),
result: WireResult::Ok {
value: HostResponse::SessionReady {
session_id: session_id(name),
},
},
})
);
}
for (id, name, expected_text, expected_wait) in [
(
3,
"client",
"0123456789012345678901234567890123456789",
"01234567890123456789012345678901234567890123456789012345678901234567890123456789",
),
(
4,
"server",
"Warning: truncated output (original token count: 10)\nTotal output lines: 1\n\n0123456789…5 tokens truncated…0123456789",
"Warning: truncated output (original token count: 20)\nTotal output lines: 1\n\n01234567890123…13 tokens truncated…67890123456789",
),
] {
let mut request = execute_request(
"text('0123456789'.repeat(4)); yield_control(); text('0123456789'.repeat(8));",
);
request.max_output_tokens = Some(5);
writer
.write(&ClientToHost::Request {
id: request_id(id),
request: HostRequest::Execute {
session_id: session_id(name),
request,
},
})
.await
.expect("execute");
// Admission is sent before start_cell can send its initial result.
assert!(matches!(
next_response(&mut reader).await,
HostToClient::Response {
id: actual_id,
result: WireResult::Ok { value: HostResponse::ExecutionStarted { .. } },
} if actual_id == request_id(id)
));
let HostToClient::InitialResponse {
id: actual_id,
result:
WireResult::Ok {
value:
WireRuntimeResponse::Yielded {
cell_id,
content_items,
..
},
},
} = next_response(&mut reader).await
else {
panic!("expected initial yield");
};
assert_eq!(actual_id, request_id(id));
assert_eq!(
content_items,
vec![WireContentItem::InputText {
text: expected_text.to_string()
}]
);
writer
.write(&ClientToHost::Request {
id: request_id(id + 2),
request: HostRequest::Wait {
session_id: session_id(name),
request: WireWaitRequest {
cell_id,
yield_time_ms: 60_000,
},
// A supplied budget only affects Server mode and belongs
// to this observation, not the preceding exec.
max_output_tokens: Some(7),
},
})
.await
.expect("wait");
let HostToClient::Response {
id: actual_id,
result:
WireResult::Ok {
value:
HostResponse::WaitCompleted {
outcome:
WireWaitOutcome::LiveCell(WireRuntimeResponse::Result {
content_items,
error_text,
..
}),
},
},
} = next_response(&mut reader).await
else {
panic!("expected final wait");
};
assert_eq!(actual_id, request_id(id + 2));
assert_eq!(error_text, None);
assert_eq!(
content_items,
vec![WireContentItem::InputText {
text: expected_wait.to_string()
}]
);
}
drop(writer);
drop(reader);
host.await.expect("host task").expect("host connection");
}
#[tokio::test]
async fn incompatible_or_invalid_handshake_is_rejected() {
let (host_stream, client_stream) = tokio::io::duplex(/*max_buf_size*/ 1024);
@@ -521,10 +700,7 @@ async fn execute_request_id_remains_active_until_initial_response() {
});
let session_id = session_id("session-1");
state
.open_session(
session_id.clone(),
CodeModeSessionCellExecutionLimits::default(),
)
.open_session(session_id.clone(), CodeModeSessionOptions::default())
.expect("open session");
let request_id = request_id(/*value*/ 1);
@@ -561,7 +737,10 @@ async fn execute_request_id_remains_active_until_initial_response() {
state
.session(&session_id)
.expect("session")
.terminate(codex_code_mode_protocol::CellId::from(cell_id).into())
.terminate(TerminateRequest {
cell_id: cell_id.into(),
max_output_tokens: None,
})
.await
.expect("terminate cell");
state.disconnect().await;
@@ -585,10 +764,7 @@ async fn active_cell_limit_rejects_execute_without_disconnecting() {
};
let session_id = session_id("session-1");
state
.open_session(
session_id.clone(),
CodeModeSessionCellExecutionLimits::default(),
)
.open_session(session_id.clone(), CodeModeSessionOptions::default())
.expect("open session");
let request_id = request_id(/*value*/ 1);
codex-rs/code-mode-host/src/lib.rs
+32 -10
View File
@@ -11,7 +11,12 @@ use std::time::Instant;
use anyhow::Context;
use anyhow::Result;
use codex_code_mode_protocol::CodeModeOutputTruncation;
use codex_code_mode_protocol::CodeModeSessionCellExecutionLimits;
use codex_code_mode_protocol::CodeModeSessionOptions;
use codex_code_mode_protocol::TerminateRequest;
use codex_code_mode_protocol::WaitRequest;
use codex_code_mode_protocol::host::BOUNDED_OUTPUT_CAPABILITY;
use codex_code_mode_protocol::host::Capability;
use codex_code_mode_protocol::host::CapabilitySet;
use codex_code_mode_protocol::host::ClientToHost;
@@ -278,6 +283,7 @@ async fn negotiate<R: AsyncRead + Unpin, W: AsyncWrite + Unpin>(
[
SESSION_RESOURCE_LIMITS_CAPABILITY,
YIELD_OBSERVATION_CAPABILITY,
BOUNDED_OUTPUT_CAPABILITY,
]
.into_iter()
.map(Capability::new)
@@ -398,13 +404,23 @@ impl HostState {
HostRequest::OpenSession {
session_id,
cell_execution_limits,
..
output_truncation,
} => {
let result = CodeModeSessionCellExecutionLimits::try_from(
cell_execution_limits.unwrap_or_default(),
)
.map_err(|error| format!("invalid code-mode session execution limits: {error}"))
.and_then(|limits| self.open_session(session_id.clone(), limits))
.and_then(|limits| {
self.open_session(
session_id.clone(),
CodeModeSessionOptions {
cell_execution_limits: limits,
output_truncation: output_truncation
.map(CodeModeOutputTruncation::from)
.unwrap_or_default(),
},
)
})
.map(|()| HostResponse::SessionReady { session_id });
self.respond(request_id, result);
}
@@ -469,7 +485,7 @@ impl HostState {
HostRequest::Wait {
session_id,
request,
..
max_output_tokens,
} => {
let result = match self.session(&session_id) {
Ok(session) => {
@@ -478,7 +494,10 @@ impl HostState {
_ = cancellation.cancelled() => {
Err("code-mode request cancelled".to_string())
}
result = session.wait(request.into(), Some(yield_signal)) => result.and_then(|outcome| {
result = session.wait(WaitRequest {
max_output_tokens,
..request.into()
}, Some(yield_signal)) => result.and_then(|outcome| {
let outcome = outcome.with_code_mode_host_duration(received_at.elapsed());
Ok(HostResponse::WaitCompleted {
outcome: WireWaitOutcome::try_from(outcome)
@@ -494,11 +513,14 @@ impl HostState {
HostRequest::Terminate {
session_id,
cell_id,
..
max_output_tokens,
} => {
let result = match self.session(&session_id) {
Ok(session) => session
.terminate(codex_code_mode_protocol::CellId::from(cell_id).into())
.terminate(TerminateRequest {
cell_id: cell_id.into(),
max_output_tokens,
})
.await
.and_then(|outcome| {
let outcome =
@@ -536,7 +558,7 @@ impl HostState {
fn open_session(
&self,
session_id: SessionId,
cell_execution_limits: CodeModeSessionCellExecutionLimits,
options: CodeModeSessionOptions,
) -> Result<(), String> {
let mut sessions = self.sessions.lock().unwrap_or_else(PoisonError::into_inner);
if sessions.contains_key(&session_id) {
@@ -563,9 +585,9 @@ impl HostState {
});
sessions.insert(
session_id,
Arc::new(InProcessCodeModeSession::with_task_failure_handler(
task_failure_handler,
cell_execution_limits,
Arc::new(InProcessCodeModeSession::with_options(
options,
Some(task_failure_handler),
)),
);
Ok(())
codex-rs/code-mode-host/tests/grpc_timing.rs
+4 -5
View File
@@ -165,11 +165,10 @@ async fn stdio_execution_timing_includes_javascript_but_excludes_delayed_reads()
if suffix.is_empty() {
assert_eq!(error_text, None);
} else {
assert!(
error_text
.context("expected JS failure")?
.contains("timed failure")
);
let Some(error_text) = error_text else {
anyhow::bail!("expected JS failure");
};
assert!(error_text.contains("timed failure"));
}
}
session.shutdown().await.map_err(anyhow::Error::msg)?;
codex-rs/code-mode-protocol/src/runtime.rs
+4
View File
@@ -144,6 +144,10 @@ pub enum RuntimeResponse {
Result {
cell_id: CellId,
content_items: Vec<FunctionCallOutputContentItem>,
/// None means success; Some (including an empty string) means failure.
/// Client mode supplies the original diagnostic to append before truncation.
/// Server mode includes the diagnostic in the already-budgeted content and
/// supplies an independently bounded copy here for logging; do not append it again.
error_text: Option<String>,
/// Time between completion and retrieval when the result was buffered.
#[serde(default, skip_serializing_if = "Duration::is_zero")]
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,
codex-rs/utils/audio/src/lib.rs
+16 -11
View File
@@ -147,17 +147,22 @@ fn canonical_audio_mime(mime: &str) -> Option<&'static str> {
/// Estimates audio tokens from decoded duration, falling back to the data URL size.
pub fn estimate_audio_token_count(audio_url: &str) -> usize {
let key = sha1_digest(audio_url.as_bytes());
AUDIO_TOKEN_ESTIMATE_CACHE.get_or_init(key, || {
let Some(duration_seconds) = audio_duration_seconds(audio_url) else {
return approx_token_count(audio_url);
};
let token_count = (duration_seconds * AUDIO_TOKENS_PER_SECOND).ceil();
if token_count >= usize::MAX as f64 {
usize::MAX
} else {
token_count as usize
}
})
AUDIO_TOKEN_ESTIMATE_CACHE
.get_or_init(key, || estimate_audio_token_count_without_cache(audio_url))
}
/// The same duration-based estimate without accessing the Tokio-dependent cache.
/// Suitable for bounded output rendered on a current-thread runtime.
pub fn estimate_audio_token_count_without_cache(audio_url: &str) -> usize {
let Some(duration_seconds) = audio_duration_seconds(audio_url) else {
return approx_token_count(audio_url);
};
let token_count = (duration_seconds * AUDIO_TOKENS_PER_SECOND).ceil();
if token_count >= usize::MAX as f64 {
usize::MAX
} else {
token_count as usize
}
}
fn audio_duration_seconds(audio_url: &str) -> Option<f64> {