From edd1ae37ec82650e8d7a97a7e84db39fd5722441 Mon Sep 17 00:00:00 2001 From: Sean Huang Date: Thu, 29 Oct 2026 00:53:16 +0000 Subject: [PATCH] Enable server-side output truncation for Code Mode sessions (#57870) ## 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 --- codex-rs/Cargo.lock | 1 + .../code-mode-host/src/grpc/conversions.rs | 16 + codex-rs/code-mode-host/src/grpc/mod.rs | 39 +- .../code-mode-host/src/grpc/service_tests.rs | 182 +++++++++- codex-rs/code-mode-host/src/grpc/session.rs | 30 +- codex-rs/code-mode-host/src/host_tests.rs | 276 +++++++++++--- codex-rs/code-mode-host/src/lib.rs | 42 ++- codex-rs/code-mode-host/tests/grpc_timing.rs | 9 +- codex-rs/code-mode-protocol/src/runtime.rs | 4 + codex-rs/code-mode-runtime/Cargo.toml | 1 + .../code-mode-runtime/src/cell_actor/mod.rs | 7 +- .../code-mode-runtime/src/cell_actor/tests.rs | 60 ++++ codex-rs/code-mode-runtime/src/output/mod.rs | 18 +- .../src/output/output_render.rs | 10 +- .../src/output/output_render_tests.rs | 4 +- codex-rs/code-mode-runtime/src/service.rs | 196 +++++++--- .../src/service_audio_tests.rs | 47 ++- .../src/service_output_collection_tests.rs | 340 ++++++++++++++++++ .../code-mode-runtime/src/service_tests.rs | 2 +- .../src/session_runtime/mod.rs | 1 + .../src/session_runtime/types.rs | 2 + codex-rs/utils/audio/src/lib.rs | 27 +- 22 files changed, 1148 insertions(+), 166 deletions(-) create mode 100644 codex-rs/code-mode-runtime/src/service_output_collection_tests.rs diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index a2ec85317a..f457de4b64 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -2805,6 +2805,7 @@ dependencies = [ "codex-protocol", "codex-utils-audio", "codex-utils-output-truncation", + "codex-utils-string", "deno_core_icudata", "futures", "opentelemetry", diff --git a/codex-rs/code-mode-host/src/grpc/conversions.rs b/codex-rs/code-mode-host/src/grpc/conversions.rs index 6dbcb6fd9d..276eee6dd9 100644 --- a/codex-rs/code-mode-host/src/grpc/conversions.rs +++ b/codex-rs/code-mode-host/src/grpc/conversions.rs @@ -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, + output_truncation: CodeModeOutputTruncation, +) -> Result, 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 { validation::identifier(&request.tool_call_id, "tool call ID")?; Ok(ExecuteRequest { diff --git a/codex-rs/code-mode-host/src/grpc/mod.rs b/codex-rs/code-mode-host/src/grpc/mod.rs index 3901da32c0..7a4a0c0d27 100644 --- a/codex-rs/code-mode-host/src/grpc/mod.rs +++ b/codex-rs/code-mode-host/src/grpc/mod.rs @@ -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>, 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())) } diff --git a/codex-rs/code-mode-host/src/grpc/service_tests.rs b/codex-rs/code-mode-host/src/grpc/service_tests.rs index c847145221..d3cb8a70c4 100644 --- a/codex-rs/code-mode-host/src/grpc/service_tests.rs +++ b/codex-rs/code-mode-host/src/grpc/service_tests.rs @@ -11,10 +11,18 @@ use tonic::Request; pub(super) async fn open_session( host: &GrpcCodeModeHost, ) -> (String, GrpcStream) { + 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, +) -> (proto::SessionOpened, GrpcStream) { 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); +} diff --git a/codex-rs/code-mode-host/src/grpc/session.rs b/codex-rs/code-mode-host/src/grpc/session.rs index f95e008b21..c1988ce8a3 100644 --- a/codex-rs/code-mode-host/src/grpc/session.rs +++ b/codex-rs/code-mode-host/src/grpc/session.rs @@ -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, - limits: CodeModeSessionCellExecutionLimits, + options: CodeModeSessionOptions, ) -> Result, 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, - limits: CodeModeSessionCellExecutionLimits, + options: CodeModeSessionOptions, ) -> Arc { Arc::new_cyclic(|weak: &Weak| { 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 { + pub(super) async fn terminate(&self, request: TerminateRequest) -> Result { 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; }); } } diff --git a/codex-rs/code-mode-host/src/host_tests.rs b/codex-rs/code-mode-host/src/host_tests.rs index 368a7ee6cb..097fa58be9 100644 --- a/codex-rs/code-mode-host/src/host_tests.rs +++ b/codex-rs/code-mode-host/src/host_tests.rs @@ -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) -> 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::().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::().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::().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::().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); diff --git a/codex-rs/code-mode-host/src/lib.rs b/codex-rs/code-mode-host/src/lib.rs index b534418123..aa6d99f3cd 100644 --- a/codex-rs/code-mode-host/src/lib.rs +++ b/codex-rs/code-mode-host/src/lib.rs @@ -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( [ 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(()) diff --git a/codex-rs/code-mode-host/tests/grpc_timing.rs b/codex-rs/code-mode-host/tests/grpc_timing.rs index 835635bf6c..ffe59d784f 100644 --- a/codex-rs/code-mode-host/tests/grpc_timing.rs +++ b/codex-rs/code-mode-host/tests/grpc_timing.rs @@ -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)?; diff --git a/codex-rs/code-mode-protocol/src/runtime.rs b/codex-rs/code-mode-protocol/src/runtime.rs index 5996ac41b5..9ed15f0d0b 100644 --- a/codex-rs/code-mode-protocol/src/runtime.rs +++ b/codex-rs/code-mode-protocol/src/runtime.rs @@ -144,6 +144,10 @@ pub enum RuntimeResponse { Result { cell_id: CellId, content_items: Vec, + /// 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, /// Time between completion and retrieval when the result was buffered. #[serde(default, skip_serializing_if = "Duration::is_zero")] diff --git a/codex-rs/code-mode-runtime/Cargo.toml b/codex-rs/code-mode-runtime/Cargo.toml index e55086f3c6..10603151ba 100644 --- a/codex-rs/code-mode-runtime/Cargo.toml +++ b/codex-rs/code-mode-runtime/Cargo.toml @@ -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 } diff --git a/codex-rs/code-mode-runtime/src/cell_actor/mod.rs b/codex-rs/code-mode-runtime/src/cell_actor/mod.rs index f3d96328e7..fe1c69d501 100644 --- a/codex-rs/code-mode-runtime/src/cell_actor/mod.rs +++ b/codex-rs/code-mode-runtime/src/cell_actor/mod.rs @@ -474,8 +474,11 @@ async fn run_cell( host.closed().await; } -/// Construct completion only after it has won the terminal decision. -fn take_completed_event(output: &mut OutputCollector, error_text: Option) -> CellEvent { +/// Construct completion only after the terminal outcome is settled. +pub(crate) fn take_completed_event( + output: &mut OutputCollector, + error_text: Option, +) -> CellEvent { let error_text = error_text.map(|error_text| output.collect_error(error_text)); CellEvent::Completed { content_items: output.take(), diff --git a/codex-rs/code-mode-runtime/src/cell_actor/tests.rs b/codex-rs/code-mode-runtime/src/cell_actor/tests.rs index 5a47b176eb..1ae1beaff8 100644 --- a/codex-rs/code-mode-runtime/src/cell_actor/tests.rs +++ b/codex-rs/code-mode-runtime/src/cell_actor/tests.rs @@ -178,6 +178,66 @@ fn spawn_cell_actor_harness_with_host_and_failure_handler( } } +#[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(); diff --git a/codex-rs/code-mode-runtime/src/output/mod.rs b/codex-rs/code-mode-runtime/src/output/mod.rs index 8989881357..08aefc3489 100644 --- a/codex-rs/code-mode-runtime/src/output/mod.rs +++ b/codex-rs/code-mode-runtime/src/output/mod.rs @@ -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, "…") + } } } } diff --git a/codex-rs/code-mode-runtime/src/output/output_render.rs b/codex-rs/code-mode-runtime/src/output/output_render.rs index 1d989f0bdc..0c506f522d 100644 --- a/codex-rs/code-mode-runtime/src/output/output_render.rs +++ b/codex-rs/code-mode-runtime/src/output/output_render.rs @@ -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, @@ -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, ) }; diff --git a/codex-rs/code-mode-runtime/src/output/output_render_tests.rs b/codex-rs/code-mode-runtime/src/output/output_render_tests.rs index fe8b6a78de..8c8db71e1f 100644 --- a/codex-rs/code-mode-runtime/src/output/output_render_tests.rs +++ b/codex-rs/code-mode-runtime/src/output/output_render_tests.rs @@ -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), diff --git a/codex-rs/code-mode-runtime/src/service.rs b/codex-rs/code-mode-runtime/src/service.rs index 925bf8ccf8..ad009994c1 100644 --- a/codex-rs/code-mode-runtime/src/service.rs +++ b/codex-rs/code-mode-runtime/src/service.rs @@ -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, - 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, ) -> 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, ) -> Result { 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, ) -> Result { + 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, ) -> 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 { - 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, ) -> Result { 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, ) -> Result { 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, +) -> Vec { + 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, +) -> Result { + 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, +) -> 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; diff --git a/codex-rs/code-mode-runtime/src/service_audio_tests.rs b/codex-rs/code-mode-runtime/src/service_audio_tests.rs index 313656782a..f8f4fba85a 100644 --- a/codex-rs/code-mode-runtime/src/service_audio_tests.rs +++ b/codex-rs/code-mode-runtime/src/service_audio_tests.rs @@ -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(), + }] + ); +} diff --git a/codex-rs/code-mode-runtime/src/service_output_collection_tests.rs b/codex-rs/code-mode-runtime/src/service_output_collection_tests.rs new file mode 100644 index 0000000000..7ba1dfa03d --- /dev/null +++ b/codex-rs/code-mode-runtime/src/service_output_collection_tests.rs @@ -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(); + } +} diff --git a/codex-rs/code-mode-runtime/src/service_tests.rs b/codex-rs/code-mode-runtime/src/service_tests.rs index e239875c3b..8a14b2421f 100644 --- a/codex-rs/code-mode-runtime/src/service_tests.rs +++ b/codex-rs/code-mode-runtime/src/service_tests.rs @@ -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), ) ); } diff --git a/codex-rs/code-mode-runtime/src/session_runtime/mod.rs b/codex-rs/code-mode-runtime/src/session_runtime/mod.rs index ca9bd034e5..06224d5a72 100644 --- a/codex-rs/code-mode-runtime/src/session_runtime/mod.rs +++ b/codex-rs/code-mode-runtime/src/session_runtime/mod.rs @@ -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) } diff --git a/codex-rs/code-mode-runtime/src/session_runtime/types.rs b/codex-rs/code-mode-runtime/src/session_runtime/types.rs index 18189aa0a1..770f43fafd 100644 --- a/codex-rs/code-mode-runtime/src/session_runtime/types.rs +++ b/codex-rs/code-mode-runtime/src/session_runtime/types.rs @@ -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, /// Time the completed result spent waiting for an observer. result_buffered_for: Duration, diff --git a/codex-rs/utils/audio/src/lib.rs b/codex-rs/utils/audio/src/lib.rs index 0f1f53b8bb..e93c6e831d 100644 --- a/codex-rs/utils/audio/src/lib.rs +++ b/codex-rs/utils/audio/src/lib.rs @@ -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 {