From 09db085f9456c608a7960fdd4c5c55a34c6a84ea Mon Sep 17 00:00:00 2001 From: Sean Huang Date: Thu, 29 Oct 2026 16:49:35 +0000 Subject: [PATCH] Support negotiated server-side truncation in Code Mode clients (#58125) ## Why Host-side truncation can bound Code Mode output before it crosses the transport, including script errors that would otherwise exceed the frame limit. ## What changed - Negotiate server-side output truncation for stdio and gRPC sessions, with client-side fallback for older hosts. - Forward output budgets for execution, waiting, and termination when server-side truncation is selected. Avoid truncating host-formatted output or appending script errors a second time. - Preserve the negotiated mode across reconnects and reject hosts that cannot maintain it. ## Testing Add transport and integration coverage for output budgets, oversized script errors, reconnect compatibility, and preservation of output-limit notices with a zero-token budget. GitOrigin-RevId: 2075101e4845eca64f45491c27a3b326788645c8 --- .../tests/grpc/session_recovery_tests.rs | 139 +++++++++++++++ codex-rs/code-mode-host/tests/stdio.rs | 38 +++- .../src/grpc_session/callbacks_tests.rs | 2 + .../code-mode/src/grpc_session/conversion.rs | 14 ++ codex-rs/code-mode/src/grpc_session/mod.rs | 57 +++++- .../code-mode/src/grpc_session/operations.rs | 19 +- .../code-mode/src/grpc_session/reconnect.rs | 26 ++- codex-rs/code-mode/src/remote_session.rs | 65 +++++-- .../src/remote_session/connection.rs | 55 ++++-- .../connection/driver/commands.rs | 33 ++-- .../connection/driver/responses.rs | 1 + .../remote_session/connection/driver/types.rs | 9 +- .../remote_session/connection/driver_tests.rs | 165 ++++++++++++++---- codex-rs/core/src/code_mode_host.rs | 19 +- codex-rs/core/src/test_support.rs | 1 + .../src/tools/code_mode/execute_handler.rs | 1 + codex-rs/core/src/tools/code_mode/mod.rs | 25 ++- .../core/src/tools/code_mode/wait_handler.rs | 8 +- codex-rs/core/tests/suite/code_mode.rs | 154 +++++++++++++--- codex-rs/core/tests/suite/scenarios.rs | 44 +++-- ...cenarios__astra_code_mode_call_timing.snap | 30 +++- 21 files changed, 751 insertions(+), 154 deletions(-) diff --git a/codex-rs/code-mode-host/tests/grpc/session_recovery_tests.rs b/codex-rs/code-mode-host/tests/grpc/session_recovery_tests.rs index b11dbca224..822ebab886 100644 --- a/codex-rs/code-mode-host/tests/grpc/session_recovery_tests.rs +++ b/codex-rs/code-mode-host/tests/grpc/session_recovery_tests.rs @@ -6,9 +6,13 @@ use std::sync::atomic::AtomicUsize; use std::sync::atomic::Ordering; use anyhow::Result; +use codex_code_mode::CodeModeOutputTruncation; +use codex_code_mode::CodeModeSessionOptions; use codex_code_mode::CodeModeSessionProvider; +use codex_code_mode::FunctionCallOutputContentItem; use codex_code_mode::GrpcCodeModeSessionProvider; use codex_code_mode::NoopCodeModeSessionDelegate; +use codex_code_mode::RuntimeResponse; use codex_code_mode_host::GrpcCodeModeHost; use codex_code_mode_protocol::grpc::code_mode_host_server::CodeModeHostServer; use pretty_assertions::assert_eq; @@ -56,6 +60,10 @@ async fn only_missing_sessions_reopen_without_replaying_execution() -> Result<() .create_session() .await .map_err(anyhow::Error::msg)?; + assert_eq!( + session.output_truncation(), + CodeModeOutputTruncation::Client + ); execute( &session, request(r#"store("before", 1);"#), @@ -121,3 +129,134 @@ async fn only_missing_sessions_reopen_without_replaying_execution() -> Result<() } Ok(()) } + +#[tokio::test] +async fn grpc_session_pins_output_mode_across_recovery_and_shutdown() -> Result<()> { + use CodeModeOutputTruncation::Client; + use CodeModeOutputTruncation::Server; + use axum::body::Body; + use axum::extract::Request; + use axum::middleware::Next; + + // A Server-capable host honors either choice. A Client-only host ignores the new field. + for (initial_host, replacement_host) in [(Server, Server), (Server, Client), (Client, Server)] { + let opens = Arc::new(AtomicUsize::new(0)); + let subscriptions = Arc::new(AtomicUsize::new(0)); + let subscription_count = Arc::clone(&subscriptions); + let app = axum::Router::new() + .fallback_service(CodeModeHostServer::new(GrpcCodeModeHost::new())) + .layer(axum::middleware::from_fn( + move |mut rpc: Request, next: Next| { + let opens = Arc::clone(&opens); + let subscriptions = Arc::clone(&subscription_count); + async move { + if rpc.uri().path().ends_with("/OpenSession") { + let host = if opens.fetch_add(1, Ordering::SeqCst) == 0 { + initial_host + } else { + replacement_host + }; + if host == Client { + // No resource limits were requested, so ignoring the output mode + // leaves an empty OpenSessionRequest in one uncompressed gRPC frame. + rpc.headers_mut().remove("content-length"); + *rpc.body_mut() = Body::from(&[0_u8; 5][..]); + } + } else if rpc.uri().path().ends_with("/Execute") + && opens.load(Ordering::SeqCst) == 1 + { + return Status::not_found("session was lost").into_http::(); + } else if rpc.uri().path().ends_with("/SubscribeToToolCalls") { + subscriptions.fetch_add(1, Ordering::SeqCst); + } + next.run(rpc).await + } + }, + )); + let listener = TcpListener::bind("127.0.0.1:0").await?; + let endpoint = format!("http://{}", listener.local_addr()?); + let _server = AbortOnDropHandle::new(tokio::spawn(async move { + axum::serve(listener, app) + .await + .expect("recovery test host"); + })); + let session = GrpcCodeModeSessionProvider::new(endpoint) + .create_session_with_options(CodeModeSessionOptions { + output_truncation: Server, + ..CodeModeSessionOptions::default() + }) + .await + .map_err(anyhow::Error::msg)?; + assert_eq!(session.output_truncation(), initial_host); + + let delegate = Arc::new(NoopCodeModeSessionDelegate); + let failed = execute(&session, request("text('must not run')"), delegate.clone()) + .await + .unwrap_err(); + assert!(failed.to_string().contains("session was lost")); + let reset = execute( + &session, + request("text('must not replay')"), + delegate.clone(), + ) + .await + .unwrap_err(); + if initial_host == Server && replacement_host == Client { + assert_eq!( + reset.to_string(), + "gRPC code-mode host changed the session output truncation mode on reconnect" + ); + assert_eq!( + subscriptions.load(Ordering::SeqCst), + 1, + "the downgraded binding must not subscribe" + ); + } else { + assert!( + reset + .to_string() + .contains("Stored values and running cells were lost") + ); + let source = match initial_host { + Client => "text('x'.repeat(2 * 1024 * 1024))", + Server => "throw 'begin:' + 'X'.repeat(65 * 1024 * 1024) + ':end'", + }; + let mut observation = request(source); + observation.max_output_tokens = Some(5); + observation.yield_time_ms = Some(60_000); + let response = execute(&session, observation, delegate).await?; + match initial_host { + Client => assert_eq!( + response, + text_response( + "g2:1", + &"x".repeat(2 * 1024 * 1024), + response.code_mode_host_duration() + ) + ), + Server => { + let RuntimeResponse::Result { + cell_id, + content_items, + error_text: Some(diagnostic), + .. + } = response + else { + panic!("expected a script failure"); + }; + assert_eq!(cell_id.as_str(), "g2:1"); + assert_eq!(diagnostic.len(), 64 * 1024); + assert!(diagnostic.starts_with("begin:") && diagnostic.ends_with(":end")); + assert_eq!(content_items, [ + "Warning: truncated output (original token count: 262144)\nTotal output lines: 2\n\nScript err…262139 tokens truncated…", + "[Code Mode session output limit reached: omitted 67108888 text bytes; discarded 0 text segments, 0 images, 0 audio items]", + "XXXXXX:end", + ].map(|text| FunctionCallOutputContentItem::InputText { text: text.into() })); + } + } + } + session.shutdown().await.map_err(anyhow::Error::msg)?; + assert_eq!(session.output_truncation(), initial_host); + } + Ok(()) +} diff --git a/codex-rs/code-mode-host/tests/stdio.rs b/codex-rs/code-mode-host/tests/stdio.rs index de75d914ed..41151ee68a 100644 --- a/codex-rs/code-mode-host/tests/stdio.rs +++ b/codex-rs/code-mode-host/tests/stdio.rs @@ -10,9 +10,11 @@ use std::os::unix::fs::PermissionsExt; use codex_code_mode::CellId; use codex_code_mode::CodeModeNestedToolCall; +use codex_code_mode::CodeModeOutputTruncation; use codex_code_mode::CodeModeSession; use codex_code_mode::CodeModeSessionCellExecutionLimits; use codex_code_mode::CodeModeSessionDelegate; +use codex_code_mode::CodeModeSessionOptions; use codex_code_mode::CodeModeSessionProvider; use codex_code_mode::CodeModeToolKind; use codex_code_mode::ExecuteRequest; @@ -908,9 +910,14 @@ async fn oversized_initial_response_does_not_close_the_shared_host() { "unexpected error: {error}" ); + // The default session still returns raw output after the oversized frame failed. let actual = execute( &session, - execute_request(r#"text("still alive");"#), + ExecuteRequest { + max_output_tokens: Some(5), + yield_time_ms: Some(60_000), + ..execute_request("text('x'.repeat(2 * 1024 * 1024));") + }, Arc::new(RecordingDelegate::default()), ) .await; @@ -920,13 +927,40 @@ async fn oversized_initial_response_does_not_close_the_shared_host() { code_mode_host_duration: actual.code_mode_host_duration(), cell_id: cell_id("2"), content_items: vec![FunctionCallOutputContentItem::InputText { - text: "still alive".to_string(), + text: "x".repeat(2 * 1024 * 1024), }], result_buffered_for: Duration::ZERO, error_text: None, } ); session.shutdown().await.expect("shutdown remote session"); + + // With Server truncation negotiated, a script error larger than the frame limit still returns. + let session = provider + .create_session_with_options(CodeModeSessionOptions { + output_truncation: CodeModeOutputTruncation::Server, + ..Default::default() + }) + .await + .expect("open server-truncated session"); + let mut request = execute_request("throw 'begin:' + 'X'.repeat(65 * 1024 * 1024) + ':end';"); + request.max_output_tokens = Some(5); + request.yield_time_ms = Some(60_000); + let response = execute(&session, request, Arc::new(RecordingDelegate::default())).await; + let RuntimeResponse::Result { + content_items, + error_text: Some(error), + .. + } = response + else { + panic!("expected completed error within the frame limit"); + }; + assert!(error.starts_with("begin:") && error.ends_with(":end")); + assert!(error.len() > 20 && error.len() <= 64 * 1024); + assert!( + matches!(content_items.last(), Some(FunctionCallOutputContentItem::InputText { text }) if text.ends_with(":end")) + ); + session.shutdown().await.expect("shutdown remote session"); } #[cfg(unix)] diff --git a/codex-rs/code-mode/src/grpc_session/callbacks_tests.rs b/codex-rs/code-mode/src/grpc_session/callbacks_tests.rs index b057e3f7e8..d32bc1efcd 100644 --- a/codex-rs/code-mode/src/grpc_session/callbacks_tests.rs +++ b/codex-rs/code-mode/src/grpc_session/callbacks_tests.rs @@ -6,6 +6,7 @@ use std::sync::Mutex; use std::sync::atomic::AtomicBool; use std::time::Duration; +use codex_code_mode_protocol::CodeModeOutputTruncation; use http_body_util::Full; use pretty_assertions::assert_eq; use tokio::sync::Notify; @@ -48,6 +49,7 @@ async fn missing_invocation_completion_racing_cancellation_preserves_session() { fn session(client: GrpcClient) -> Arc { Arc::new(SessionInner { id: "session".to_string(), + output_truncation: CodeModeOutputTruncation::Client, client: client.clone(), route: None, runtime: tokio::runtime::Handle::current(), diff --git a/codex-rs/code-mode/src/grpc_session/conversion.rs b/codex-rs/code-mode/src/grpc_session/conversion.rs index d78f3f571a..4120a8b1df 100644 --- a/codex-rs/code-mode/src/grpc_session/conversion.rs +++ b/codex-rs/code-mode/src/grpc_session/conversion.rs @@ -2,6 +2,7 @@ use std::time::Duration; use codex_code_mode_protocol::CellId; use codex_code_mode_protocol::CodeModeNestedToolCall; +use codex_code_mode_protocol::CodeModeOutputTruncation; use codex_code_mode_protocol::CodeModeToolKind; use codex_code_mode_protocol::ExecuteRequest; use codex_code_mode_protocol::FunctionCallOutputContentItem; @@ -12,6 +13,19 @@ use codex_code_mode_protocol::WaitOutcome; use codex_code_mode_protocol::grpc; use codex_protocol::ToolName; +pub(super) fn output_budget( + mode: CodeModeOutputTruncation, + max_output_tokens: Option, +) -> Result, String> { + match mode { + CodeModeOutputTruncation::Client => Ok(None), + CodeModeOutputTruncation::Server => max_output_tokens + .map(u64::try_from) + .transpose() + .map_err(|error| format!("invalid code-mode output token limit: {error}")), + } +} + pub(super) fn execute_request( session_id: &str, execution_id: String, diff --git a/codex-rs/code-mode/src/grpc_session/mod.rs b/codex-rs/code-mode/src/grpc_session/mod.rs index 602f2fe2d0..e805e24666 100644 --- a/codex-rs/code-mode/src/grpc_session/mod.rs +++ b/codex-rs/code-mode/src/grpc_session/mod.rs @@ -10,9 +10,11 @@ use std::sync::atomic::Ordering; use std::time::Duration; use codex_code_mode_protocol::CellId; +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::CodeModeSessionProvider; use codex_code_mode_protocol::CodeModeSessionProviderFuture; use codex_code_mode_protocol::CodeModeSessionResultFuture; @@ -129,9 +131,12 @@ impl GrpcCodeModeSessionProvider { #[tracing::instrument(name = "code_mode.grpc.open_binding", level = "info", skip_all)] async fn open_binding( &self, - limits: CodeModeSessionCellExecutionLimits, + options: &CodeModeSessionOptions, + pinned_mode: Option, ) -> Result, String> { let mut client = deadline::startup("transport connection", self.transport.client()).await?; + let requested_mode = pinned_mode.unwrap_or(options.output_truncation); + let limits = &options.cell_execution_limits; let limits = grpc::SessionCellExecutionLimits { max_yield_time_ms: limits.max_yield_time_ms, max_heap_size_bytes: limits @@ -149,7 +154,12 @@ impl GrpcCodeModeSessionProvider { loop { let mut request = tonic::Request::new(grpc::OpenSessionRequest { cell_execution_limits, - output_truncation: None, + output_truncation: match requested_mode { + CodeModeOutputTruncation::Client => None, + CodeModeOutputTruncation::Server => { + Some(grpc::OutputTruncation::Server.into()) + } + }, }); inject_span_traceparent(&mut request, &open_session_span); let result = async { @@ -197,9 +207,33 @@ impl GrpcCodeModeSessionProvider { return Err("gRPC code-mode session lease omitted its opening event".to_string()); }; validate_identifier(&opened.session_id, "session ID")?; + let host_mode = + grpc::OutputTruncation::try_from(opened.output_truncation.unwrap_or_default()) + .map_err(|_| "gRPC code-mode host returned an unknown output truncation mode")?; + let output_truncation = match host_mode { + grpc::OutputTruncation::Unspecified | grpc::OutputTruncation::Client => { + CodeModeOutputTruncation::Client + } + grpc::OutputTruncation::Server => CodeModeOutputTruncation::Server, + }; + if output_truncation == CodeModeOutputTruncation::Server + && requested_mode == CodeModeOutputTruncation::Client + { + return Err(format!( + "gRPC code-mode host selected {output_truncation:?} output truncation, requested {requested_mode:?}" + )); + } + if pinned_mode.is_some_and(|previous| previous != output_truncation) { + // Dropping the opening lease releases this incompatible session before subscribing. + return Err( + "gRPC code-mode host changed the session output truncation mode on reconnect" + .to_string(), + ); + } let inner = Arc::new(SessionInner { id: opened.session_id, + output_truncation, client, route, runtime: tokio::runtime::Handle::current(), @@ -259,9 +293,19 @@ impl CodeModeSessionProvider for GrpcCodeModeSessionProvider { fn create_session_with_limits<'a>( &'a self, limits: CodeModeSessionCellExecutionLimits, + ) -> CodeModeSessionProviderFuture<'a> { + self.create_session_with_options(CodeModeSessionOptions { + cell_execution_limits: limits, + ..CodeModeSessionOptions::default() + }) + } + + fn create_session_with_options<'a>( + &'a self, + options: CodeModeSessionOptions, ) -> CodeModeSessionProviderFuture<'a> { Box::pin(async move { - let session = Arc::new(reconnect::ReconnectableSession::new(self.clone(), limits)); + let session = Arc::new(reconnect::ReconnectableSession::new(self.clone(), options)); session.initialize().await?; Ok(session as _) }) @@ -290,6 +334,10 @@ impl Drop for OpeningSession { } impl CodeModeSession for GrpcCodeModeSession { + fn output_truncation(&self) -> CodeModeOutputTruncation { + self.inner.output_truncation + } + fn execute<'a>( &'a self, request: ExecuteRequest, @@ -311,7 +359,7 @@ impl CodeModeSession for GrpcCodeModeSession { &'a self, request: TerminateRequest, ) -> CodeModeSessionResultFuture<'a, WaitOutcome> { - Box::pin(self.inner.terminate(request.cell_id)) + Box::pin(self.inner.terminate(request)) } fn shutdown<'a>(&'a self) -> CodeModeSessionResultFuture<'a, ()> { @@ -327,6 +375,7 @@ impl Drop for GrpcCodeModeSession { pub(super) struct SessionInner { pub(super) id: String, + output_truncation: CodeModeOutputTruncation, pub(super) client: GrpcClient, // Affinity belongs to the session: parents and children may share a transport // while being admitted to different hosts. The token is opaque to Codex. diff --git a/codex-rs/code-mode/src/grpc_session/operations.rs b/codex-rs/code-mode/src/grpc_session/operations.rs index bd02e961d3..877590198e 100644 --- a/codex-rs/code-mode/src/grpc_session/operations.rs +++ b/codex-rs/code-mode/src/grpc_session/operations.rs @@ -10,6 +10,7 @@ use codex_code_mode_protocol::DEFAULT_EXEC_YIELD_TIME_MS; use codex_code_mode_protocol::ExecuteRequest; use codex_code_mode_protocol::RuntimeResponse; use codex_code_mode_protocol::StartedCell; +use codex_code_mode_protocol::TerminateRequest; use codex_code_mode_protocol::WaitOutcome; use codex_code_mode_protocol::WaitRequest; use codex_code_mode_protocol::grpc; @@ -58,7 +59,7 @@ impl Drop for ExecutionOwnership { } let session = Arc::clone(&self.session); self.session.runtime.spawn(async move { - if let Err(error) = session.terminate(cell_id).await + if let Err(error) = session.terminate(cell_id.into()).await && !session.stopped.is_cancelled() { debug!("abandoned code-mode execution termination raced closure: {error}"); @@ -313,7 +314,10 @@ impl SessionInner { cell_id: expected_cell_id.as_str().to_string(), wait_id: wait_id.clone(), yield_time_ms: request.yield_time_ms, - max_output_tokens: None, + max_output_tokens: conversion::output_budget( + self.output_truncation, + request.max_output_tokens, + )?, }; let mut client = self.client(); let watcher = preempt.map(|signal| { @@ -359,7 +363,7 @@ impl SessionInner { })) } - pub(super) async fn terminate(&self, cell_id: CellId) -> Result { + pub(super) async fn terminate(&self, request: TerminateRequest) -> Result { self.require_open()?; let mut client = self.client(); let response = deadline::session_request( @@ -368,13 +372,16 @@ impl SessionInner { Duration::ZERO, client.terminate(self.request(grpc::TerminateRequest { session_id: self.id.clone(), - cell_id: cell_id.as_str().to_string(), - max_output_tokens: None, + cell_id: request.cell_id.as_str().to_string(), + max_output_tokens: conversion::output_budget( + self.output_truncation, + request.max_output_tokens, + )?, })), ) .await?; let outcome = conversion::wait_outcome(response)?; - self.validate_wait_cell(&cell_id, outcome).await + self.validate_wait_cell(&request.cell_id, outcome).await } async fn validate_wait_cell( diff --git a/codex-rs/code-mode/src/grpc_session/reconnect.rs b/codex-rs/code-mode/src/grpc_session/reconnect.rs index 117865866e..9b39a445f4 100644 --- a/codex-rs/code-mode/src/grpc_session/reconnect.rs +++ b/codex-rs/code-mode/src/grpc_session/reconnect.rs @@ -1,11 +1,13 @@ use std::sync::Arc; use std::sync::Mutex; +use std::sync::OnceLock; use std::sync::PoisonError; use std::sync::atomic::Ordering; +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::ExecuteRequest; use codex_code_mode_protocol::StartedCell; @@ -30,8 +32,11 @@ pub(super) struct ReconnectableSession { struct ReconnectInner { provider: GrpcCodeModeSessionProvider, - limits: CodeModeSessionCellExecutionLimits, + options: CodeModeSessionOptions, binding: Mutex>, + // The negotiated mode belongs to the logical session, including after + // shutdown releases the binding. Never pin a mode from an unpublished open. + output_truncation: OnceLock, opening_permit: Semaphore, shutdown_requested: CancellationToken, shutdown_result: Mutex>, @@ -46,13 +51,14 @@ struct SessionBinding { impl ReconnectableSession { pub(super) fn new( provider: GrpcCodeModeSessionProvider, - limits: CodeModeSessionCellExecutionLimits, + options: CodeModeSessionOptions, ) -> Self { Self { inner: Arc::new(ReconnectInner { provider, - limits, + options, binding: Mutex::new(None), + output_truncation: OnceLock::new(), opening_permit: Semaphore::new(/*permits*/ 1), shutdown_requested: CancellationToken::new(), shutdown_result: Mutex::new(None), @@ -66,6 +72,14 @@ impl ReconnectableSession { } impl CodeModeSession for ReconnectableSession { + fn output_truncation(&self) -> CodeModeOutputTruncation { + self.inner + .output_truncation + .get() + .copied() + .unwrap_or_default() + } + fn execute<'a>( &'a self, request: ExecuteRequest, @@ -184,7 +198,7 @@ impl ReconnectInner { _ = self.shutdown_requested.cancelled() => { return Err(SHUTDOWN_ERROR.to_string()); } - session = self.provider.open_binding(self.limits.clone()) => session?, + session = self.provider.open_binding(&self.options, self.output_truncation.get().copied()) => session?, }; let binding = SessionBinding { session, @@ -196,6 +210,8 @@ impl ReconnectInner { false } else { *current = Some(binding.clone()); + self.output_truncation + .get_or_init(|| binding.session.output_truncation()); true } }; diff --git a/codex-rs/code-mode/src/remote_session.rs b/codex-rs/code-mode/src/remote_session.rs index 4c4c3c2428..55aa16c601 100644 --- a/codex-rs/code-mode/src/remote_session.rs +++ b/codex-rs/code-mode/src/remote_session.rs @@ -2,13 +2,16 @@ use std::io; use std::path::PathBuf; use std::sync::Arc; use std::sync::Mutex as StdMutex; +use std::sync::OnceLock; use std::sync::atomic::AtomicBool; use std::sync::atomic::AtomicU64; use std::sync::atomic::Ordering; +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::CodeModeSessionProvider; use codex_code_mode_protocol::CodeModeSessionProviderFuture; use codex_code_mode_protocol::CodeModeSessionResultFuture; @@ -83,7 +86,17 @@ impl CodeModeSessionProvider for ProcessOwnedCodeModeSessionProvider { &'a self, limits: CodeModeSessionCellExecutionLimits, ) -> CodeModeSessionProviderFuture<'a> { - Box::pin(create_host_session(self.process_host(), limits)) + self.create_session_with_options(CodeModeSessionOptions { + cell_execution_limits: limits, + ..CodeModeSessionOptions::default() + }) + } + + fn create_session_with_options<'a>( + &'a self, + options: CodeModeSessionOptions, + ) -> CodeModeSessionProviderFuture<'a> { + Box::pin(create_host_session(self.process_host(), options)) } } @@ -106,9 +119,9 @@ impl CodeModeSessionProvider for DisabledCodeModeSessionProvider { async fn create_host_session( host: Arc, - limits: CodeModeSessionCellExecutionLimits, + options: CodeModeSessionOptions, ) -> Result, String> { - let session = ProcessOwnedCodeModeSession::with_host(host, limits); + let session = ProcessOwnedCodeModeSession::with_host(host, options); session.connection().await?; Ok(Arc::new(session)) } @@ -218,7 +231,9 @@ struct SessionBinding { struct SessionInner { host: Arc, - limits: CodeModeSessionCellExecutionLimits, + options: CodeModeSessionOptions, + // The first successful open fixes the output policy for all generations of this session. + output_truncation: OnceLock, state: StdMutex, shutdown_requested: AtomicBool, shutdown_result: StdMutex>, @@ -236,15 +251,16 @@ impl ProcessOwnedCodeModeSession { Arc::new(OwnedCodeModeHost::new( InstallContext::current().code_mode_host_program(), )), - CodeModeSessionCellExecutionLimits::default(), + CodeModeSessionOptions::default(), ) } - fn with_host(host: Arc, limits: CodeModeSessionCellExecutionLimits) -> Self { + fn with_host(host: Arc, options: CodeModeSessionOptions) -> Self { Self { inner: Arc::new(SessionInner { host, - limits, + options, + output_truncation: OnceLock::new(), state: StdMutex::new(SessionState::New), shutdown_requested: AtomicBool::new(false), shutdown_result: StdMutex::new(None), @@ -257,9 +273,13 @@ impl ProcessOwnedCodeModeSession { self.inner.connection().await } + fn server_output_budget(&self, budget: Option) -> Option { + budget.filter(|_| self.output_truncation() == CodeModeOutputTruncation::Server) + } + pub async fn execute( &self, - request: ExecuteRequest, + mut request: ExecuteRequest, delegate: Arc, preempt: Option, ) -> Result { @@ -268,6 +288,10 @@ impl ProcessOwnedCodeModeSession { delegate, trace: codex_otel::current_span_w3c_trace_context(), }); + // V1 encodes the exec budget as i32. Its maximum already exceeds the host collector limit. + request.max_output_tokens = self + .server_output_budget(request.max_output_tokens) + .map(|tokens| tokens.min(i32::MAX as usize)); binding .connection .execute(binding.remote, request, delegate, preempt) @@ -276,22 +300,21 @@ impl ProcessOwnedCodeModeSession { pub async fn wait( &self, - request: WaitRequest, + mut request: WaitRequest, preempt: Option, ) -> Result { let binding = self.connection().await?; + request.max_output_tokens = self.server_output_budget(request.max_output_tokens); binding .connection .wait(binding.remote, request, preempt) .await } - pub async fn terminate(&self, request: TerminateRequest) -> Result { + pub async fn terminate(&self, mut request: TerminateRequest) -> Result { let binding = self.connection().await?; - binding - .connection - .terminate(binding.remote, request.cell_id) - .await + request.max_output_tokens = self.server_output_budget(request.max_output_tokens); + binding.connection.terminate(binding.remote, request).await } pub async fn shutdown(&self) -> Result<(), String> { @@ -352,7 +375,11 @@ impl SessionInner { let result = match self.host.connection().await { Ok(connection) => { let cleanup = connection - .open_session(remote.clone(), self.limits.clone()) + .open_session( + remote.clone(), + self.options.clone(), + &self.output_truncation, + ) .await; cleanup.map(|cleanup| SessionBinding { connection, @@ -529,6 +556,14 @@ impl Default for ProcessOwnedCodeModeSession { } impl CodeModeSession for ProcessOwnedCodeModeSession { + fn output_truncation(&self) -> CodeModeOutputTruncation { + self.inner + .output_truncation + .get() + .copied() + .unwrap_or_default() + } + fn execute<'a>( &'a self, request: ExecuteRequest, diff --git a/codex-rs/code-mode/src/remote_session/connection.rs b/codex-rs/code-mode/src/remote_session/connection.rs index ed5a22074b..504c0fa80a 100644 --- a/codex-rs/code-mode/src/remote_session/connection.rs +++ b/codex-rs/code-mode/src/remote_session/connection.rs @@ -5,17 +5,21 @@ use std::path::Path; use std::path::PathBuf; use std::process::Stdio; use std::sync::Arc; +use std::sync::OnceLock; use std::sync::atomic::AtomicBool; use std::sync::atomic::Ordering; use std::time::Duration; -use codex_code_mode_protocol::CellId; +use codex_code_mode_protocol::CodeModeOutputTruncation; use codex_code_mode_protocol::CodeModeSessionCellExecutionLimits; use codex_code_mode_protocol::CodeModeSessionDelegate; +use codex_code_mode_protocol::CodeModeSessionOptions; use codex_code_mode_protocol::ExecuteRequest; use codex_code_mode_protocol::StartedCell; +use codex_code_mode_protocol::TerminateRequest; use codex_code_mode_protocol::WaitOutcome; 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::ClientHello; @@ -209,9 +213,14 @@ impl Connection { .map_err(|error| error.to_string())?; let yield_capability = Capability::new(YIELD_OBSERVATION_CAPABILITY).map_err(|error| error.to_string())?; - let optional_capabilities = - CapabilitySet::try_new([session_limits_capability, yield_capability]) - .map_err(|error| error.to_string())?; + let output_capability = + Capability::new(BOUNDED_OUTPUT_CAPABILITY).map_err(|error| error.to_string())?; + let optional_capabilities = CapabilitySet::try_new([ + session_limits_capability, + yield_capability, + output_capability, + ]) + .map_err(|error| error.to_string())?; let hello = ClientHello::new( SupportedProtocolVersions::try_new([ProtocolVersion::V1]) .map_err(|err| err.to_string())?, @@ -326,27 +335,42 @@ impl Connection { self.alive.load(Ordering::Acquire) } + fn supports(&self, capability: &str) -> bool { + self.capabilities + .iter() + .any(|value| value.as_str() == capability) + } + pub(super) async fn open_session( &self, session: RemoteSession, - limits: CodeModeSessionCellExecutionLimits, + mut options: CodeModeSessionOptions, + pinned: &OnceLock, ) -> Result { - if limits != CodeModeSessionCellExecutionLimits::default() - && !self - .capabilities - .iter() - .any(|capability| capability.as_str() == SESSION_RESOURCE_LIMITS_CAPABILITY) + if options.cell_execution_limits != CodeModeSessionCellExecutionLimits::default() + && !self.supports(SESSION_RESOURCE_LIMITS_CAPABILITY) { return Err(format!( "code-mode host does not support session resource limits: missing `{SESSION_RESOURCE_LIMITS_CAPABILITY}` capability" )); } + let pinned_mode = pinned.get().copied(); + let output_truncation = if self.supports(BOUNDED_OUTPUT_CAPABILITY) { + pinned_mode.unwrap_or(options.output_truncation) + } else if pinned_mode == Some(CodeModeOutputTruncation::Server) { + return Err( + "code-mode host cannot preserve server output truncation after reconnect".into(), + ); + } else { + CodeModeOutputTruncation::Client + }; + options.output_truncation = output_truncation; let cleanup = SessionCleanup::new(); let cancellation = CallerCancellation::new(); let (response_tx, response_rx) = oneshot::channel(); self.send(DriverCommand::OpenSession { session, - limits, + options, cleanup: cleanup.clone(), caller_cancellation: cancellation.token(), response_tx, @@ -355,6 +379,7 @@ impl Connection { let result = self.receive(response_rx).await; cancellation.disarm(); result?; + pinned.get_or_init(|| output_truncation); Ok(cleanup) } @@ -424,9 +449,7 @@ impl Connection { &self, signal: Option, ) -> Option { - self.capabilities - .iter() - .any(|capability| capability.as_str() == YIELD_OBSERVATION_CAPABILITY) + self.supports(YIELD_OBSERVATION_CAPABILITY) .then_some(signal) .flatten() } @@ -434,13 +457,13 @@ impl Connection { pub(super) async fn terminate( &self, session: RemoteSession, - cell_id: CellId, + request: TerminateRequest, ) -> Result { let (response_tx, response_rx) = oneshot::channel(); self.with_transport_deadline(Duration::ZERO, "terminate", async { self.send(DriverCommand::Terminate { session, - cell_id, + request, response_tx, }) .await?; diff --git a/codex-rs/code-mode/src/remote_session/connection/driver/commands.rs b/codex-rs/code-mode/src/remote_session/connection/driver/commands.rs index 3c4e050418..bc76df8fdb 100644 --- a/codex-rs/code-mode/src/remote_session/connection/driver/commands.rs +++ b/codex-rs/code-mode/src/remote_session/connection/driver/commands.rs @@ -1,14 +1,16 @@ use std::sync::Arc; -use codex_code_mode_protocol::CellId; -use codex_code_mode_protocol::CodeModeSessionCellExecutionLimits; +use codex_code_mode_protocol::CodeModeOutputTruncation; use codex_code_mode_protocol::CodeModeSessionDelegate; +use codex_code_mode_protocol::CodeModeSessionOptions; use codex_code_mode_protocol::ExecuteRequest; +use codex_code_mode_protocol::TerminateRequest; use codex_code_mode_protocol::WaitOutcome; use codex_code_mode_protocol::WaitRequest; use codex_code_mode_protocol::host::ClientToHost; use codex_code_mode_protocol::host::EncodedFrame; use codex_code_mode_protocol::host::HostRequest; +use codex_code_mode_protocol::host::WireOutputTruncation; use codex_code_mode_protocol::host::WireSessionCellExecutionLimits; use codex_code_mode_protocol::host::WireWaitRequest; use tokio::sync::oneshot; @@ -30,11 +32,11 @@ impl ConnectionDriver { match command { DriverCommand::OpenSession { session, - limits, + options, cleanup, caller_cancellation, response_tx, - } => self.open_session(session, limits, cleanup, caller_cancellation, response_tx), + } => self.open_session(session, options, cleanup, caller_cancellation, response_tx), DriverCommand::Execute { session, request, @@ -65,9 +67,9 @@ impl ConnectionDriver { ), DriverCommand::Terminate { session, - cell_id, + request, response_tx, - } => self.terminate(session, cell_id, response_tx), + } => self.terminate(session, request, response_tx), DriverCommand::ShutdownSession { session, response_tx, @@ -78,7 +80,7 @@ impl ConnectionDriver { fn open_session( &mut self, session: RemoteSession, - limits: CodeModeSessionCellExecutionLimits, + options: CodeModeSessionOptions, cleanup: super::cleanup::SessionCleanup, caller_cancellation: CancellationToken, response_tx: oneshot::Sender>, @@ -90,7 +92,7 @@ impl ConnectionDriver { ))); return true; } - let limits = match WireSessionCellExecutionLimits::try_from(limits) { + let limits = match WireSessionCellExecutionLimits::try_from(options.cell_execution_limits) { Ok(limits) => limits, Err(error) => { let _ = response_tx.send(Err(format!( @@ -112,7 +114,8 @@ impl ConnectionDriver { session_id: session.id.clone(), cell_execution_limits: (limits != WireSessionCellExecutionLimits::default()) .then_some(limits), - output_truncation: None, + output_truncation: (options.output_truncation == CodeModeOutputTruncation::Server) + .then_some(WireOutputTruncation::Server), }, }; let frame = match EncodedFrame::encode(&message) { @@ -213,6 +216,7 @@ impl ConnectionDriver { let _ = response_tx.send(Err(err)); return true; } + let max_output_tokens = request.max_output_tokens; let request = match remote_wait_request(&session, request) { Ok(request) => request, Err(err) => { @@ -224,6 +228,7 @@ impl ConnectionDriver { self.requests.push_deferred_wait(DeferredWait { session, request, + max_output_tokens, caller_cancellation, yield_signal, response_tx, @@ -233,6 +238,7 @@ impl ConnectionDriver { self.start_wait( session, request, + max_output_tokens, caller_cancellation, yield_signal, response_tx, @@ -243,6 +249,7 @@ impl ConnectionDriver { &mut self, session: RemoteSession, request: WireWaitRequest, + max_output_tokens: Option, caller_cancellation: CancellationToken, yield_signal: Option, response_tx: oneshot::Sender>, @@ -252,7 +259,7 @@ impl ConnectionDriver { HostRequest::Wait { session_id: session.id.clone(), request, - max_output_tokens: None, + max_output_tokens, }, PendingRequest::Wait { session, @@ -267,14 +274,14 @@ impl ConnectionDriver { fn terminate( &mut self, session: RemoteSession, - cell_id: CellId, + request: TerminateRequest, response_tx: oneshot::Sender>, ) -> bool { if let Err(err) = self.sessions.require_ready(&session) { let _ = response_tx.send(Err(err)); return true; } - let cell_id = match remote_cell_id(&session, &cell_id) { + let cell_id = match remote_cell_id(&session, &request.cell_id) { Ok(cell_id) => cell_id, Err(err) => { let _ = response_tx.send(Err(err)); @@ -286,7 +293,7 @@ impl ConnectionDriver { HostRequest::Terminate { session_id: session.id.clone(), cell_id, - max_output_tokens: None, + max_output_tokens: request.max_output_tokens, }, PendingRequest::Terminate { session, diff --git a/codex-rs/code-mode/src/remote_session/connection/driver/responses.rs b/codex-rs/code-mode/src/remote_session/connection/driver/responses.rs index f3bd836c12..7f20fb6e23 100644 --- a/codex-rs/code-mode/src/remote_session/connection/driver/responses.rs +++ b/codex-rs/code-mode/src/remote_session/connection/driver/responses.rs @@ -41,6 +41,7 @@ impl ConnectionDriver { if !self.start_wait( wait.session, wait.request, + wait.max_output_tokens, wait.caller_cancellation, wait.yield_signal, wait.response_tx, diff --git a/codex-rs/code-mode/src/remote_session/connection/driver/types.rs b/codex-rs/code-mode/src/remote_session/connection/driver/types.rs index 8aed64908e..b1bc8862c1 100644 --- a/codex-rs/code-mode/src/remote_session/connection/driver/types.rs +++ b/codex-rs/code-mode/src/remote_session/connection/driver/types.rs @@ -1,11 +1,11 @@ use std::sync::Arc; -use codex_code_mode_protocol::CellId; -use codex_code_mode_protocol::CodeModeSessionCellExecutionLimits; use codex_code_mode_protocol::CodeModeSessionDelegate; +use codex_code_mode_protocol::CodeModeSessionOptions; use codex_code_mode_protocol::ExecuteRequest; use codex_code_mode_protocol::RuntimeResponse; use codex_code_mode_protocol::StartedCell; +use codex_code_mode_protocol::TerminateRequest; use codex_code_mode_protocol::WaitOutcome; use codex_code_mode_protocol::WaitRequest; use codex_code_mode_protocol::host::DelegateRequestId; @@ -31,7 +31,7 @@ pub(in crate::remote_session) struct RemoteSession { pub(in crate::remote_session::connection) enum DriverCommand { OpenSession { session: RemoteSession, - limits: CodeModeSessionCellExecutionLimits, + options: CodeModeSessionOptions, cleanup: SessionCleanup, caller_cancellation: CancellationToken, response_tx: oneshot::Sender>, @@ -53,7 +53,7 @@ pub(in crate::remote_session::connection) enum DriverCommand { }, Terminate { session: RemoteSession, - cell_id: CellId, + request: TerminateRequest, response_tx: oneshot::Sender>, }, ShutdownSession { @@ -200,6 +200,7 @@ pub(super) enum PendingRequest { pub(super) struct DeferredWait { pub(super) session: RemoteSession, pub(super) request: WireWaitRequest, + pub(super) max_output_tokens: Option, pub(super) caller_cancellation: CancellationToken, pub(super) yield_signal: Option, pub(super) response_tx: oneshot::Sender>, diff --git a/codex-rs/code-mode/src/remote_session/connection/driver_tests.rs b/codex-rs/code-mode/src/remote_session/connection/driver_tests.rs index 99e65368d8..7c89483d97 100644 --- a/codex-rs/code-mode/src/remote_session/connection/driver_tests.rs +++ b/codex-rs/code-mode/src/remote_session/connection/driver_tests.rs @@ -7,12 +7,18 @@ use std::time::Duration; use codex_code_mode_protocol::CellId; use codex_code_mode_protocol::CodeModeNestedToolCall; +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::ExecuteRequest; use codex_code_mode_protocol::NotificationFuture; +use codex_code_mode_protocol::TerminateRequest; use codex_code_mode_protocol::ToolInvocationFuture; 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; use codex_code_mode_protocol::host::DelegateRequest; @@ -31,6 +37,7 @@ use codex_code_mode_protocol::host::WireSessionCellExecutionLimits; use codex_code_mode_protocol::host::WireWaitOutcome; use codex_protocol::ToolName; use pretty_assertions::assert_eq; +use serde_json::json; use tokio::sync::mpsc; use tokio::sync::oneshot; use tokio_util::sync::CancellationToken; @@ -43,6 +50,8 @@ use super::DriverEvent; use super::DriverLifecycle; use super::RemoteSession; use super::SessionCleanup; +use crate::remote_session::OwnedCodeModeHost; +use crate::remote_session::ProcessOwnedCodeModeSession; struct DriverHarness { command_tx: mpsc::Sender, @@ -93,26 +102,15 @@ impl DriverHarness { self.command_tx .send(DriverCommand::OpenSession { session: session.clone(), - - limits: Default::default(), + options: CodeModeSessionOptions::default(), cleanup: cleanup.clone(), caller_cancellation: CancellationToken::new(), response_tx, }) .await .expect("open command"); - self.outgoing_rx.recv().await.expect("open frame"); - self.event_tx - .send(DriverEvent::HostMessage(HostToClient::Response { - id: RequestId::new(/*value*/ 1), - result: WireResult::Ok { - value: HostResponse::SessionReady { - session_id: session.id, - }, - }, - })) - .await - .expect("open response"); + self.accept_open(session.id, CodeModeOutputTruncation::Client) + .await; response_rx .await .expect("open reply") @@ -120,6 +118,30 @@ impl DriverHarness { cleanup } + async fn accept_open(&mut self, session_id: SessionId, mode: CodeModeOutputTruncation) { + let mut expected = json!({ "method": "session/open", "sessionId": session_id }); + if mode == CodeModeOutputTruncation::Server { + expected["outputTruncation"] = json!("server"); + } + assert_eq!(self.next_request().await, expected); + self.event_tx + .send(DriverEvent::HostMessage(HostToClient::Response { + id: RequestId::new(/*value*/ 1), + result: WireResult::Ok { + value: HostResponse::SessionReady { session_id }, + }, + })) + .await + .expect("open response"); + } + + async fn next_request(&mut self) -> serde_json::Value { + let frame = self.outgoing_rx.recv().await.expect("request"); + let mut message = + EncodedFrame::decode_framed::(&frame.into_framed_bytes()).unwrap(); + message["request"].take() + } + async fn start_cell( &mut self, session: RemoteSession, @@ -167,6 +189,17 @@ impl DriverHarness { delivered.started } + fn connection(&self, capabilities: CapabilitySet) -> Connection { + Connection { + command_tx: self.command_tx.clone(), + execute_claim_tx: self.execute_claim_tx.clone(), + alive: Arc::clone(&self.alive), + failure: Arc::clone(&self.failure), + cancellation: self.cancellation.clone(), + capabilities, + } + } + async fn start_tool_delegate(&self, session: &RemoteSession, id: DelegateRequestId) { self.event_tx .send(DriverEvent::HostMessage(HostToClient::DelegateRequest { @@ -220,14 +253,7 @@ async fn old_host_does_not_receive_yield_frames() { })) .await .unwrap(); - let connection = Connection { - command_tx: harness.command_tx.clone(), - execute_claim_tx: harness.execute_claim_tx.clone(), - alive: Arc::clone(&harness.alive), - failure: Arc::clone(&harness.failure), - cancellation: harness.cancellation.clone(), - capabilities: CapabilitySet::empty(), - }; + let connection = harness.connection(CapabilitySet::empty()); let signal = CancellationToken::new(); signal.cancel(); let wait = tokio::spawn(async move { @@ -287,8 +313,10 @@ async fn open_session_includes_nondefault_cell_execution_limits() { .command_tx .send(DriverCommand::OpenSession { session: session.clone(), - - limits, + options: CodeModeSessionOptions { + cell_execution_limits: limits, + ..CodeModeSessionOptions::default() + }, cleanup: SessionCleanup::new(), caller_cancellation: CancellationToken::new(), response_tx, @@ -318,6 +346,80 @@ async fn open_session_includes_nondefault_cell_execution_limits() { ); } +#[tokio::test] +async fn stdio_logical_session_pins_output_mode_across_reconnect() { + use CodeModeOutputTruncation::Client; + use CodeModeOutputTruncation::Server; + + let legacy = CapabilitySet::empty(); + let bounded = + CapabilitySet::try_new([Capability::new(BOUNDED_OUTPUT_CAPABILITY).unwrap()]).unwrap(); + for (requested, hosts, expected_mode) in [ + (Client, [&bounded, &bounded].as_slice(), Client), + (Server, [&legacy, &bounded].as_slice(), Client), + (Server, [&bounded, &bounded, &legacy].as_slice(), Server), + ] { + let host = Arc::new(OwnedCodeModeHost::new("unused".into())); + let session = ProcessOwnedCodeModeSession::with_host( + Arc::clone(&host), + CodeModeSessionOptions { + output_truncation: requested, + ..CodeModeSessionOptions::default() + }, + ); + + for (index, &capabilities) in hosts.iter().enumerate() { + let mut harness = DriverHarness::start(); + *host.connection.lock().unwrap() = + Some(Arc::new(harness.connection(capabilities.clone()))); + if expected_mode == Server && capabilities == &legacy { + tokio::select! { + result = session.connection() => assert!(matches!(result, Err(error) if error.contains("cannot preserve"))), + _ = harness.outgoing_rx.recv() => panic!("opened a server session on an incompatible host"), + } + assert_eq!(session.output_truncation(), Server); + break; + } + + let session_id = SessionId::new(format!("session-{}", index + 1)).unwrap(); + tokio::try_join!(session.connection(), async { + harness.accept_open(session_id, expected_mode).await; + Ok(()) + }) + .unwrap(); + assert_eq!(session.output_truncation(), expected_mode); + + if capabilities == &legacy { + let cell_id = CellId::new("1".into()); + let wait = WaitRequest { + cell_id: cell_id.clone(), + yield_time_ms: 1, + max_output_tokens: Some(5), + }; + let terminate = TerminateRequest { + cell_id, + max_output_tokens: Some(5), + }; + let legacy_host = async { + for method in ["session/wait", "session/terminate"] { + let request = harness.next_request().await; + assert_eq!(request["method"], method); + // A null value would also be rejected by a legacy host. + assert!(request.get("maxOutputTokens").is_none()); + } + harness.cancellation.cancel(); + }; + let (wait, terminate, ()) = tokio::join!(biased; + session.wait(wait, /*preempt*/ None), session.terminate(terminate), legacy_host + ); + assert!(wait.is_err() && terminate.is_err()); + } + harness.cancellation.cancel(); + (&mut harness.driver_task).await.unwrap(); + } + } +} + #[derive(Default)] struct RecordingDelegate { closed_cells: StdMutex>, @@ -626,8 +728,7 @@ async fn dropped_open_waiter_shuts_down_committed_session() { .command_tx .send(DriverCommand::OpenSession { session: session.clone(), - - limits: Default::default(), + options: CodeModeSessionOptions::default(), cleanup, caller_cancellation: CancellationToken::new(), response_tx: open_tx, @@ -857,7 +958,7 @@ async fn terminate_closes_cell_without_waiting_for_delegate_cleanup() { .command_tx .send(DriverCommand::Terminate { session: session.clone(), - cell_id: CellId::new("1".to_string()), + request: CellId::new("1".to_string()).into(), response_tx, }) .await @@ -1330,7 +1431,7 @@ async fn mismatched_terminate_response_fails_connection() { .command_tx .send(DriverCommand::Terminate { session, - cell_id: CellId::new("1".to_string()), + request: CellId::new("1".to_string()).into(), response_tx, }) .await @@ -1498,7 +1599,7 @@ async fn queued_remote_termination_times_out_and_invalidates_the_connection() { }; let response = tokio::spawn(async move { let result = connection - .terminate(session, CellId::new("1".to_string())) + .terminate(session, CellId::new("1".to_string()).into()) .await; (connection, result) }); @@ -1539,7 +1640,7 @@ async fn cancelled_wait_is_retired_before_next_wait_is_sent() { request: WaitRequest { cell_id: CellId::new("1".to_string()), yield_time_ms: 60_000, - max_output_tokens: None, + max_output_tokens: Some(19), }, caller_cancellation: first_cancellation.clone(), yield_signal: None, @@ -1559,7 +1660,7 @@ async fn cancelled_wait_is_retired_before_next_wait_is_sent() { request: WaitRequest { cell_id: CellId::new("1".to_string()), yield_time_ms: 1, - max_output_tokens: None, + max_output_tokens: Some(37), }, caller_cancellation: CancellationToken::new(), yield_signal: None, @@ -1587,7 +1688,7 @@ async fn cancelled_wait_is_retired_before_next_wait_is_sent() { })) .await .expect("cancelled wait response"); - harness.outgoing_rx.recv().await.expect("second wait frame"); + assert_eq!(harness.next_request().await["maxOutputTokens"], 37); harness .event_tx .send(DriverEvent::HostMessage(HostToClient::Response { diff --git a/codex-rs/core/src/code_mode_host.rs b/codex-rs/core/src/code_mode_host.rs index ffad6e2da3..4d6eea1a5b 100644 --- a/codex-rs/core/src/code_mode_host.rs +++ b/codex-rs/core/src/code_mode_host.rs @@ -19,6 +19,7 @@ use std::time::Duration; use codex_code_mode::CodeModeProcessMetricsSample; use codex_code_mode::CodeModeSessionCellExecutionLimits; +use codex_code_mode::CodeModeSessionOptions; use codex_code_mode::CodeModeSessionProvider; use codex_code_mode::CodeModeSessionProviderFuture; use codex_code_mode::CodeModeSessionResultFuture; @@ -43,14 +44,14 @@ use tonic::transport::Endpoint; use tower::ServiceExt; use tower::service_fn; -pub(crate) struct ProcessOwnedGrpcCodeModeSessionProvider { +pub struct ProcessOwnedGrpcCodeModeSessionProvider { host_program: PathBuf, provider: OnceLock, current: Arc>>, } impl ProcessOwnedGrpcCodeModeSessionProvider { - pub(crate) fn with_host_program(host_program: PathBuf) -> Self { + pub fn with_host_program(host_program: PathBuf) -> Self { Self { host_program, provider: OnceLock::new(), @@ -86,12 +87,22 @@ impl CodeModeSessionProvider for ProcessOwnedGrpcCodeModeSessionProvider { } fn create_session(&self) -> CodeModeSessionProviderFuture<'_> { - self.create_session_with_limits(CodeModeSessionCellExecutionLimits::default()) + self.create_session_with_options(CodeModeSessionOptions::default()) } fn create_session_with_limits<'a>( &'a self, limits: CodeModeSessionCellExecutionLimits, + ) -> CodeModeSessionProviderFuture<'a> { + self.create_session_with_options(CodeModeSessionOptions { + cell_execution_limits: limits, + ..CodeModeSessionOptions::default() + }) + } + + fn create_session_with_options<'a>( + &'a self, + options: CodeModeSessionOptions, ) -> CodeModeSessionProviderFuture<'a> { Box::pin(async move { self.availability()?; @@ -105,7 +116,7 @@ impl CodeModeSessionProvider for ProcessOwnedGrpcCodeModeSessionProvider { async move { transport.send(request).await } })) }); - provider.create_session_with_limits(limits).await + provider.create_session_with_options(options).await }) } } diff --git a/codex-rs/core/src/test_support.rs b/codex-rs/core/src/test_support.rs index b9f94b7eb1..c14f568d00 100644 --- a/codex-rs/core/src/test_support.rs +++ b/codex-rs/core/src/test_support.rs @@ -32,6 +32,7 @@ use codex_protocol::protocol::SessionSource; use once_cell::sync::Lazy; use crate::ThreadManager; +pub use crate::code_mode_host::ProcessOwnedGrpcCodeModeSessionProvider; use crate::config::Config; use crate::responses_metadata::CodexResponsesMetadata; use crate::responses_metadata::CodexResponsesRequestKind; diff --git a/codex-rs/core/src/tools/code_mode/execute_handler.rs b/codex-rs/core/src/tools/code_mode/execute_handler.rs index 317e029a46..756c8057d8 100644 --- a/codex-rs/core/src/tools/code_mode/execute_handler.rs +++ b/codex-rs/core/src/tools/code_mode/execute_handler.rs @@ -165,6 +165,7 @@ impl CodeModeExecuteHandler { &step_context.settings.model_info, response, args.max_output_tokens, + exec.session.services.code_mode_service.output_truncation(), wall_time, exec.turn.config.code_mode.experimental_show_cell_overhead, )) diff --git a/codex-rs/core/src/tools/code_mode/mod.rs b/codex-rs/core/src/tools/code_mode/mod.rs index 936d1bc266..a8c9095480 100644 --- a/codex-rs/core/src/tools/code_mode/mod.rs +++ b/codex-rs/core/src/tools/code_mode/mod.rs @@ -15,6 +15,7 @@ use std::time::Duration; use codex_code_mode::CellId; use codex_code_mode::CodeModeNestedToolCall; +use codex_code_mode::CodeModeOutputTruncation; use codex_code_mode::CodeModeSession; use codex_code_mode::CodeModeSessionProvider; use codex_code_mode::CodeModeToolKind; @@ -126,6 +127,13 @@ impl CodeModeService { Arc::clone(&self.session_provider) } + fn output_truncation(&self) -> CodeModeOutputTruncation { + self.session + .get() + .map(|session| session.output_truncation()) + .unwrap_or_default() + } + pub(crate) async fn execute( &self, mut request: codex_code_mode::ExecuteRequest, @@ -155,9 +163,9 @@ impl CodeModeService { pub(crate) async fn terminate( &self, - cell_id: CellId, + request: codex_code_mode::TerminateRequest, ) -> Result { - self.session().await?.terminate(cell_id.into()).await + self.session().await?.terminate(request).await } pub(crate) async fn interrupt_active_cells(&self) { @@ -263,6 +271,7 @@ fn handle_runtime_response( model_info: &codex_protocol::openai_models::ModelInfo, response: RuntimeResponse, max_output_tokens: Option, + output_truncation: CodeModeOutputTruncation, wall_time: Duration, experimental_show_cell_overhead: bool, ) -> CodeModeToolOutput { @@ -284,12 +293,14 @@ fn handle_runtime_response( let mut content_items = into_function_call_output_content_items(content_items); sanitize_image_detail_items(supports_original, &mut content_items); let success = error_text.is_none(); - if let Some(error_text) = error_text { - content_items.push(FunctionCallOutputContentItem::InputText { - text: format!("Script error:\n{error_text}"), - }); + if output_truncation == CodeModeOutputTruncation::Client { + if let Some(error_text) = error_text { + content_items.push(FunctionCallOutputContentItem::InputText { + text: format!("Script error:\n{error_text}"), + }); + } + content_items = truncate_code_mode_result(content_items, max_output_tokens); } - content_items = truncate_code_mode_result(content_items, max_output_tokens); CodeModeToolOutput::new( FunctionToolOutput::from_content(content_items, Some(success)), script_status, diff --git a/codex-rs/core/src/tools/code_mode/wait_handler.rs b/codex-rs/core/src/tools/code_mode/wait_handler.rs index 1623d4e427..1d3bc5643f 100644 --- a/codex-rs/core/src/tools/code_mode/wait_handler.rs +++ b/codex-rs/core/src/tools/code_mode/wait_handler.rs @@ -130,7 +130,10 @@ impl CodeModeWaitHandler { exec.session .services .code_mode_service - .terminate(cell_id.clone()) + .terminate(codex_code_mode::TerminateRequest { + cell_id: cell_id.clone(), + max_output_tokens: args.max_tokens, + }) .await } else { exec.session @@ -140,7 +143,7 @@ impl CodeModeWaitHandler { codex_code_mode::WaitRequest { cell_id: cell_id.clone(), yield_time_ms: args.yield_time_ms, - max_output_tokens: None, + max_output_tokens: args.max_tokens, }, step_context.preempt.clone(), ) @@ -208,6 +211,7 @@ impl CodeModeWaitHandler { &step_context.settings.model_info, wait_response.into(), args.max_tokens, + exec.session.services.code_mode_service.output_truncation(), wall_time, exec.turn.config.code_mode.experimental_show_cell_overhead, ))) diff --git a/codex-rs/core/tests/suite/code_mode.rs b/codex-rs/core/tests/suite/code_mode.rs index 3532ab19b1..42991e04e3 100644 --- a/codex-rs/core/tests/suite/code_mode.rs +++ b/codex-rs/core/tests/suite/code_mode.rs @@ -12,6 +12,11 @@ use anyhow::Result; use base64::Engine; use base64::engine::general_purpose::STANDARD as BASE64_STANDARD; use codex_analytics::AnalyticsEventsClient; +use codex_code_mode::CodeModeOutputTruncation; +use codex_code_mode::CodeModeSessionOptions; +use codex_code_mode::CodeModeSessionProvider; +use codex_code_mode::CodeModeSessionProviderFuture; +use codex_code_mode::ProcessOwnedCodeModeSessionProvider; use codex_config::types::McpServerConfig; use codex_config::types::McpServerTransportConfig; use codex_core::StartThreadOptions; @@ -20,6 +25,7 @@ use codex_core::config::Config; use codex_core::config::Constrained; use codex_core::config::CurrentTimeReminderConfig; use codex_core::context::NodeReplReviewEvidence; +use codex_core::test_support::ProcessOwnedGrpcCodeModeSessionProvider; use codex_extension_api::ExtensionData; use codex_extension_api::ExtensionFuture; use codex_extension_api::ExtensionRegistryBuilder; @@ -308,6 +314,74 @@ async fn run_code_mode_turn_with_builder( Ok((test, second_mock)) } +// Request server truncation from the real host and verify what it negotiated. +struct ServerOutputProvider(Arc); + +impl CodeModeSessionProvider for ServerOutputProvider { + fn availability(&self) -> std::result::Result<(), String> { + self.0.availability() + } + + fn create_session(&self) -> CodeModeSessionProviderFuture<'_> { + Box::pin(async { + let session = self + .0 + .create_session_with_options(CodeModeSessionOptions { + output_truncation: CodeModeOutputTruncation::Server, + ..CodeModeSessionOptions::default() + }) + .await?; + assert_eq!( + session.output_truncation(), + CodeModeOutputTruncation::Server + ); + Ok(session) + }) + } +} + +pub(super) enum OutputHost { + DefaultClient, + ServerStdio, + ServerGrpc, +} + +impl OutputHost { + pub(super) fn test_codex(self) -> TestCodexBuilder { + let builder = test_codex() + .with_model("test-gpt-5.1-codex") + .with_config(|config| { + config + .features + .enable(Feature::CodeMode) + .expect("enable code mode"); + }); + let host = || { + codex_utils_cargo_bin::cargo_bin("codex-code-mode-host") + .expect("locate the code-mode host") + }; + let provider: Arc = match self { + Self::DefaultClient => return builder, + Self::ServerStdio => Arc::new(ProcessOwnedCodeModeSessionProvider::with_host_program( + host(), + )), + Self::ServerGrpc => Arc::new( + ProcessOwnedGrpcCodeModeSessionProvider::with_host_program(host()), + ), + }; + builder + .with_config(|config| { + config + .features + .disable(Feature::CodeModeHost) + .expect("use the injected provider"); + }) + .with_thread_manager(move |manager| { + manager.with_code_mode_session_provider(Arc::new(ServerOutputProvider(provider))) + }) + } +} + async fn run_unavailable_code_mode_turn( server: &MockServer, test: &TestCodex, @@ -4570,12 +4644,15 @@ text(`Variable truncated: ${resultVariableWasTruncated ? "True" : "False"}. Vari // The outer directive limits output after JavaScript emits it; it does not // limit `result.output` returned by the nested command. #[cfg_attr(windows, ignore = "no exec_command on Windows")] +#[test_case(OutputHost::DefaultClient; "client_stdio")] +#[test_case(OutputHost::ServerStdio; "server_stdio")] +#[test_case(OutputHost::ServerGrpc; "server_grpc")] #[tokio::test(flavor = "multi_thread", worker_threads = 2)] -async fn code_mode_exec_outer_limit_truncates_emitted_output() -> Result<()> { +async fn code_mode_exec_outer_limit_truncates_emitted_output(host: OutputHost) -> Result<()> { skip_if_no_network!(Ok(())); let server = responses::start_mock_server().await; - let (_test, second_mock) = run_code_mode_turn( + let (_test, second_mock) = run_code_mode_turn_with_builder( &server, "use exec_command from code mode", r#"// @exec: {"max_output_tokens": 5} @@ -4584,6 +4661,7 @@ const result = await tools.exec_command({ }); text(result.output); "#, + host.test_codex(), ) .await?; @@ -5546,23 +5624,49 @@ yield_control(); Ok(()) } +// The cell emits its termination output before starting the nested command. Keep +// the first turn's dispatch worker alive until that command has been dispatched. +#[derive(Default)] +struct TerminationOutputReady(tokio::sync::Notify); + +impl ToolLifecycleContributor for TerminationOutputReady { + fn on_tool_finish<'a>(&'a self, input: ToolFinishInput<'a>) -> ToolLifecycleFuture<'a> { + Box::pin(async move { + if input.tool_name.name == "exec_command" { + self.0.notify_one(); + } else if input.tool_name.name == "exec" && input.call_id == "call-1" { + assert_eq!(input.outcome, ToolCallOutcome::Completed { success: true }); + tokio::time::timeout(Duration::from_secs(10), self.0.notified()) + .await + .expect("the yielded cell should reach its nested command"); + } + }) + } +} + #[cfg_attr(windows, ignore = "no exec_command on Windows")] +#[test_case(OutputHost::DefaultClient; "client_stdio")] +#[test_case(OutputHost::ServerStdio; "server_stdio")] +#[test_case(OutputHost::ServerGrpc; "server_grpc")] #[tokio::test(flavor = "multi_thread", worker_threads = 2)] -async fn code_mode_wait_can_terminate_and_continue() -> Result<()> { +async fn code_mode_wait_can_terminate_and_continue(host: OutputHost) -> Result<()> { skip_if_no_network!(Ok(())); let server = responses::start_mock_server().await; - let mut builder = test_codex().with_config(move |config| { - let _ = config.features.enable(Feature::CodeMode); - }); + let mut extensions = ExtensionRegistryBuilder::::new(); + extensions.tool_lifecycle_contributor(Arc::new(TerminationOutputReady::default())); + let mut builder = host + .test_codex() + .with_extensions(Arc::new(extensions.build())); let test = builder.build(&server).await?; let termination_gate = test.workspace_path("code-mode-terminate.ready"); let termination_wait = wait_for_file_source(&termination_gate)?; let code = format!( - r#" + r#"// @exec: {{"max_output_tokens": 100}} text("phase 1"); yield_control(); +text("0123456789012345678901234567890123456789"); {termination_wait} text("phase 2"); "# @@ -5604,6 +5708,7 @@ text("phase 2"); &serde_json::to_string(&serde_json::json!({ "cell_id": cell_id.clone(), "terminate": true, + "max_tokens": 5, }))?, ), ev_completed("resp-3"), @@ -5623,7 +5728,7 @@ text("phase 2"); let second_request = second_completion.single_request(); let second_items = function_tool_output_items(&second_request, "call-2"); - assert_eq!(second_items.len(), 1); + assert_eq!(second_items.len(), 2); assert_regex_match( concat!( r"(?s)\A", @@ -5631,6 +5736,10 @@ text("phase 2"); ), text_item(&second_items, /*index*/ 0), ); + assert_eq!( + text_item(&second_items, /*index*/ 1), + "Warning: truncated output (original token count: 10)\nTotal output lines: 1\n\n0123456789…5 tokens truncated…0123456789" + ); responses::mount_sse_once( &server, @@ -6224,19 +6333,19 @@ async fn code_mode_interrupt_terminates_active_cells_and_nested_tools() -> Resul Ok(()) } -#[cfg_attr(windows, ignore = "no exec_command on Windows")] +#[test_case(OutputHost::DefaultClient; "client_stdio")] +#[test_case(OutputHost::ServerStdio; "server_stdio")] +#[test_case(OutputHost::ServerGrpc; "server_grpc")] #[tokio::test(flavor = "multi_thread", worker_threads = 2)] -async fn code_mode_wait_uses_its_own_max_tokens_budget() -> Result<()> { +async fn code_mode_wait_uses_its_own_max_tokens_budget(host: OutputHost) -> Result<()> { skip_if_no_network!(Ok(())); let server = responses::start_mock_server().await; - let mut builder = test_codex().with_config(move |config| { - let _ = config.features.enable(Feature::CodeMode); - }); - let test = builder.build(&server).await?; + let mut builder = host.test_codex(); + let test = builder.build_with_auto_env(&server).await?; - // An explicit yield separates the output batches without a nested shell call. - let code = r#"// @exec: {"max_output_tokens": 100} + // Both transports accept exec budgets above i32::MAX. The wait uses its own budget. + let code = r#"// @exec: {"max_output_tokens": 2147483648} text("phase 1"); yield_control(); text("token one token two token three token four token five token six token seven"); @@ -6306,15 +6415,10 @@ text("token one token two token three token four token five token six token seve ), text_item(&second_items, /*index*/ 0), ); - let expected_pattern = r#"(?sx) -\A -Warning:\ truncated\ output\ \(original\ token\ count:\ \d+\)\n -Total\ output\ lines:\ 1\n -\n -.*…\d+\ tokens\ truncated….* -\z -"#; - assert_regex_match(expected_pattern, text_item(&second_items, /*index*/ 1)); + assert_eq!( + text_item(&second_items, /*index*/ 1), + "Warning: truncated output (original token count: 19)\nTotal output lines: 1\n\ntoken one to…13 tokens truncated… token seven" + ); Ok(()) } diff --git a/codex-rs/core/tests/suite/scenarios.rs b/codex-rs/core/tests/suite/scenarios.rs index b5410da20f..0bd1df327a 100644 --- a/codex-rs/core/tests/suite/scenarios.rs +++ b/codex-rs/core/tests/suite/scenarios.rs @@ -2198,25 +2198,18 @@ text(`MCP: ${ping.structuredContent?.echo ?? "missing"}`);"#, async fn astra_reads_code_mode_call_timing() -> Result<()> { skip_if_no_network!(Ok(())); let server = start_mock_server().await; - let test = test_codex() + let test = super::code_mode::OutputHost::ServerGrpc + .test_codex() .with_model("gpt-6-astra") .with_config(|config| { configure_scenario_catalog(config); // Use the selected cwd as the workspace root on local and remote executors. config.workspace_roots = vec![config.cwd.clone()]; config.code_mode.experimental_show_cell_overhead = true; - config - .features - .enable(Feature::CodeMode) - .expect("enable code mode"); config .features .enable(Feature::CodeModeOnly) .expect("enable code-mode-only tools"); - config - .features - .enable(Feature::CodeModeHost) - .expect("enable the code-mode host"); }) .build_with_auto_env(&server) .await?; @@ -2225,7 +2218,13 @@ async fn astra_reads_code_mode_call_timing() -> Result<()> { vec![ sse(vec![ ev_response_created("exec-response"), - ev_custom_tool_call("exec-call", "exec", "text('ready');"), + ev_custom_tool_call( + "exec-call", + "exec", + r#"// @exec: {"max_output_tokens": 0} +text('ready'); +text('x'.repeat(2 * 1024 * 1024));"#, + ), ev_completed("exec-response"), ]), sse(vec![ @@ -2239,18 +2238,35 @@ async fn astra_reads_code_mode_call_timing() -> Result<()> { ev_completed("wait-response"), ]), sse(vec![ - ev_assistant_message("final", "The call completed; the missing-cell wait failed."), + ev_response_created("overflow-response"), + ev_custom_tool_call( + "overflow-call", + "exec", + r#"// @exec: {"max_output_tokens": 32} +text('before'); +text('x'.repeat(2 * 1024 * 1024)); +throw 'boom';"#, + ), + ev_completed("overflow-response"), + ]), + sse(vec![ + ev_assistant_message( + "final", + "The first call completed. The wait and overflow script failed.", + ), ev_completed("final-response"), ]), ], ) .await; - test.submit_turn("Run a code cell, then inspect its timing and a failed wait.") - .await?; + test.submit_turn( + "Run a code cell, inspect its timing and a failed wait, then check the diagnostic script.", + ) + .await?; insta::assert_snapshot!( "astra_code_mode_call_timing", context_snapshot::format_request_history_snapshot( - "Astra receives host and handler timings on completed and failed code-mode calls.", + "Astra reads code-mode timings and the collector notice even at zero tokens, then a failed wait and a truncated script error.", &mock.requests(), &ContextSnapshotOptions::default(), ) diff --git a/codex-rs/core/tests/suite/snapshots/all__suite__scenarios__astra_code_mode_call_timing.snap b/codex-rs/core/tests/suite/snapshots/all__suite__scenarios__astra_code_mode_call_timing.snap index 061450da7a..5fdeedba16 100644 --- a/codex-rs/core/tests/suite/snapshots/all__suite__scenarios__astra_code_mode_call_timing.snap +++ b/codex-rs/core/tests/suite/snapshots/all__suite__scenarios__astra_code_mode_call_timing.snap @@ -1,8 +1,8 @@ --- source: core/tests/suite/scenarios.rs -expression: "context_snapshot::format_request_history_snapshot(\"Astra receives host and handler timings on completed and failed code-mode calls.\",\n&mock.requests(), &ContextSnapshotOptions::default(),)" +expression: "context_snapshot::format_request_history_snapshot(\"Astra reads code-mode timings and the collector notice even at zero tokens, then a failed wait and a truncated script error.\",\n&mock.requests(), &ContextSnapshotOptions::default(),)" --- -Scenario: Astra receives host and handler timings on completed and failed code-mode calls. +Scenario: Astra reads code-mode timings and the collector notice even at zero tokens, then a failed wait and a truncated script error. ## Window 1 -- request 1 (turn) -- @@ -89,14 +89,19 @@ Scenario: Astra receives host and handler timings on completed and failed code-m <...e="unrestricted" /> [hash=d40b2f45cb028888] 06:message/user: - Run a code cell, then inspect its timing and a failed wait. + Run a code cell, inspect its timing and a failed wait, then check the diagnostic script. 07:configuration_update:reasoning={"effort":"low"} -- request 2 (turn) -- -08:custom_tool_call/exec:text('ready'); +08:custom_tool_call/exec:// @exec: {"max_output_tokens": 0} + text('ready'); + text('x'.repeat(2 * 1024 * 1024)); 09:custom_tool_call_output:Script completed Wall time seconds (code-mode seconds; overhead seconds) Output: - | ready + | Warning: truncated output (original token count: 262145) + Total output lines: 2 + + …262145 tokens truncated… | [Code Mode session output limit reached: omitted 1048581 text bytes; discarded 0 text segments, 0 images, 0 audio items] -- request 3 (turn) -- 10:function_call/functions.wait:{"cell_id":"missing"} 11:function_call_output:Script failed @@ -104,3 +109,18 @@ Scenario: Astra receives host and handler timings on completed and failed code-m Output: | Script error: exec cell missing not found +-- request 4 (turn) -- +12:custom_tool_call/exec:// @exec: {"max_output_tokens": 32} + text('before'); + text('x'.repeat(2 * 1024 * 1024)); + throw 'boom'; +13:custom_tool_call_output:Script failed + Wall time seconds (code-mode seconds; overhead seconds) + Output: + | Warning: truncated output (original token count: 262145) + Total output lines: 4 + + before + xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx…262113 tokens truncated… | [Code Mode session output limit reached: omitted 1048600 text bytes; discarded 0 text segments, 0 images, 0 audio items] | xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx + Script error: + boom