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

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<SessionInner> {
     Arc::new(SessionInner {
         id: "session".to_string(),
+        output_truncation: CodeModeOutputTruncation::Client,
         client: client.clone(),
         route: None,
         runtime: tokio::runtime::Handle::current(),

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<usize>,
+) -> Result<Option<u64>, 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,

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<CodeModeOutputTruncation>,
     ) -> Result<Arc<GrpcCodeModeSession>, 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.

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<WaitOutcome, String> {
+    pub(super) async fn terminate(&self, request: TerminateRequest) -> Result<WaitOutcome, String> {
         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(

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<Option<SessionBinding>>,
+    // 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<CodeModeOutputTruncation>,
     opening_permit: Semaphore,
     shutdown_requested: CancellationToken,
     shutdown_result: Mutex<Option<ShutdownResultReceiver>>,
@@ -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
             }
         };