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/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<u64>,
+    output_truncation: CodeModeOutputTruncation,
+) -> Result<Option<usize>, Status> {
+    if output_truncation == CodeModeOutputTruncation::Client {
+        return Ok(None);
+    }
+    max_output_tokens
+        .map(|value| {
+            usize::try_from(value)
+                .map_err(|_| Status::invalid_argument("output budget is too large"))
+        })
+        .transpose()
+}
+
 pub(super) fn execute_request(request: proto::ExecuteRequest) -> Result<ExecuteRequest, Status> {
     validation::identifier(&request.tool_call_id, "tool call ID")?;
     Ok(ExecuteRequest {

codex-rs/code-mode-host/src/grpc/mod.rs

@@ -12,6 +12,10 @@ use std::sync::Arc;
 use std::time::Instant;
 
 use codex_code_mode_protocol::CellId;
+use codex_code_mode_protocol::CodeModeOutputTruncation;
+use codex_code_mode_protocol::CodeModeSession;
+use codex_code_mode_protocol::CodeModeSessionOptions;
+use codex_code_mode_protocol::TerminateRequest;
 use codex_code_mode_protocol::WaitRequest;
 use codex_code_mode_protocol::grpc as proto;
 use codex_code_mode_protocol::grpc::code_mode_host_server::CodeModeHost;
@@ -75,7 +79,23 @@ impl GrpcCodeModeHost {
     ) -> Result<Response<GrpcStream<proto::SessionEvent>>, Status> {
         let _permit = self.state.request_permit()?;
         let limits = conversions::session_limits(request.cell_execution_limits)?;
-        Ok(Response::new(self.state.open_session(limits)?))
+        let output_truncation = match request
+            .output_truncation
+            .map(proto::OutputTruncation::try_from)
+            .transpose()
+            .map_err(|_| Status::invalid_argument("unknown code-mode output truncation mode"))?
+        {
+            None | Some(proto::OutputTruncation::Unspecified | proto::OutputTruncation::Client) => {
+                CodeModeOutputTruncation::Client
+            }
+            Some(proto::OutputTruncation::Server) => CodeModeOutputTruncation::Server,
+        };
+        Ok(Response::new(self.state.open_session(
+            CodeModeSessionOptions {
+                cell_execution_limits: limits,
+                output_truncation,
+            },
+        )?))
     }
 
     #[tracing::instrument(
@@ -242,11 +262,15 @@ impl GrpcCodeModeHost {
         validation::identifier(&request.cell_id, "cell ID")?;
         validation::identifier(&request.wait_id, "wait ID")?;
         let _permit = self.state.request_permit()?;
+        let max_output_tokens = conversions::observation_budget(
+            request.max_output_tokens,
+            session.runtime.output_truncation(),
+        )?;
         let registration = WaitRegistration::new(Arc::clone(&session), request.wait_id)?;
         let request = WaitRequest {
             cell_id: CellId::new(request.cell_id),
             yield_time_ms: request.yield_time_ms,
-            max_output_tokens: None,
+            max_output_tokens,
         };
         let outcome = tokio::select! {
             biased;
@@ -301,7 +325,16 @@ impl GrpcCodeModeHost {
         let session = self.state.session(&request.session_id)?;
         validation::identifier(&request.cell_id, "cell ID")?;
         let _permit = self.state.request_permit()?;
-        let outcome = session.terminate(CellId::new(request.cell_id)).await?;
+        let max_output_tokens = conversions::observation_budget(
+            request.max_output_tokens,
+            session.runtime.output_truncation(),
+        )?;
+        let outcome = session
+            .terminate(TerminateRequest {
+                cell_id: CellId::new(request.cell_id),
+                max_output_tokens,
+            })
+            .await?;
         let outcome = outcome.with_code_mode_host_duration(received_at.elapsed());
         conversions::wait_response(outcome).map_err(|error| Status::internal(error.to_string()))
     }

codex-rs/code-mode-host/src/grpc/service_tests.rs

@@ -11,10 +11,18 @@ use tonic::Request;
 pub(super) async fn open_session(
     host: &GrpcCodeModeHost,
 ) -> (String, GrpcStream<proto::SessionEvent>) {
+    let (opened, stream) = open_session_with_truncation(host, /*output_truncation*/ None).await;
+    (opened.session_id, stream)
+}
+
+async fn open_session_with_truncation(
+    host: &GrpcCodeModeHost,
+    output_truncation: Option<i32>,
+) -> (proto::SessionOpened, GrpcStream<proto::SessionEvent>) {
     let mut stream = host
         .open_session(Request::new(proto::OpenSessionRequest {
             cell_execution_limits: None,
-            output_truncation: None,
+            output_truncation,
         }))
         .await
         .expect("open code-mode session")
@@ -27,7 +35,7 @@ pub(super) async fn open_session(
     let Some(proto::session_event::Event::Opened(opened)) = event.event else {
         panic!("expected the first session event to open its lease");
     };
-    (opened.session_id, stream)
+    (opened, stream)
 }
 
 pub(super) fn execute_request(
@@ -80,6 +88,18 @@ pub(super) async fn execute_events(
     (started.cell_id, stream)
 }
 
+fn only_text(items: &[proto::ContentItem]) -> &str {
+    let [
+        proto::ContentItem {
+            item: Some(proto::content_item::Item::Text(text)),
+        },
+    ] = items
+    else {
+        panic!("expected one text item, got {items:?}");
+    };
+    &text.text
+}
+
 #[tokio::test]
 async fn execute_stream_starts_immediately_and_wait_preserves_missing_cells() {
     let host = GrpcCodeModeHost::new();
@@ -358,7 +378,7 @@ async fn cancellation_before_wait_admission_is_preserved() {
             cell_id: cell_id.clone(),
             wait_id: "pre-cancelled".to_string(),
             yield_time_ms: 60_000,
-            max_output_tokens: None,
+            max_output_tokens: Some(0),
         }))
         .await
         .unwrap_err();
@@ -367,7 +387,7 @@ async fn cancellation_before_wait_admission_is_preserved() {
     host.terminate(Request::new(proto::TerminateRequest {
         session_id,
         cell_id,
-        max_output_tokens: None,
+        max_output_tokens: Some(0),
     }))
     .await
     .unwrap();
@@ -440,3 +460,157 @@ async fn notifications_do_not_delay_cell_completion() {
     .await
     .expect("legacy notification acknowledgments remain accepted");
 }
+
+#[tokio::test]
+async fn output_truncation_is_a_session_setting_and_omitting_it_keeps_legacy_output() {
+    for (output_truncation, bounded_output) in [
+        (None, false),
+        (Some(proto::OutputTruncation::Unspecified.into()), false),
+        (Some(proto::OutputTruncation::Client.into()), false),
+        (Some(proto::OutputTruncation::Server.into()), true),
+    ] {
+        let host = GrpcCodeModeHost::new();
+        let (opened, _lease) = open_session_with_truncation(&host, output_truncation).await;
+        assert_eq!(
+            opened.output_truncation,
+            bounded_output.then_some(proto::OutputTruncation::Server.into()),
+        );
+        let mut request = execute_request(
+            &opened.session_id,
+            "execution",
+            "text('x'.repeat(2 * 1024 * 1024));",
+        );
+        request.yield_time_ms = Some(60_000);
+        let (_, mut execution) = execute_events(&host, request).await;
+        let outcome = execution
+            .next()
+            .await
+            .expect("outcome event")
+            .expect("outcome");
+        let Some(proto::execute_event::Event::Outcome(outcome)) = outcome.event else {
+            panic!("expected execution outcome");
+        };
+        assert!(matches!(
+            outcome.outcome,
+            Some(proto::execution_outcome::Outcome::Completed(_))
+        ));
+        assert_eq!(
+            outcome.content_items.iter().any(
+                |item| matches!(&item.item, Some(proto::content_item::Item::Text(text))
+                    if text.text.starts_with("[Code Mode session output limit reached: omitted "))
+            ),
+            bounded_output
+        );
+        if !bounded_output {
+            assert_eq!(
+                only_text(&outcome.content_items),
+                "x".repeat(2 * 1024 * 1024)
+            );
+        }
+    }
+}
+
+#[tokio::test]
+async fn server_output_budgets_follow_exec_wait_and_terminate_over_grpc() {
+    for (execution_id, after_yield) in [
+        ("failed", "throw '0123456789'.repeat(8);"),
+        (
+            "terminated",
+            "text('0123456789'.repeat(8)); notify('ready to terminate'); await new Promise(() => {});",
+        ),
+    ] {
+        let host = GrpcCodeModeHost::new();
+        let (opened, mut lease) =
+            open_session_with_truncation(&host, Some(proto::OutputTruncation::Server.into())).await;
+        let source = format!("text('0123456789'.repeat(4)); yield_control(); {after_yield}");
+        let mut request = execute_request(&opened.session_id, execution_id, &source);
+        request.yield_time_ms = Some(60_000);
+        request.max_output_tokens = Some(5);
+        let (cell_id, mut execution) = execute_events(&host, request).await;
+        let Some(proto::execute_event::Event::Outcome(first)) =
+            execution.next().await.unwrap().unwrap().event
+        else {
+            panic!("expected execution outcome");
+        };
+        assert_eq!(first.cell_id, cell_id);
+        assert!(matches!(
+            first.outcome,
+            Some(proto::execution_outcome::Outcome::Yielded(_))
+        ));
+        assert_eq!(
+            only_text(&first.content_items),
+            "Warning: truncated output (original token count: 10)\nTotal output lines: 1\n\n0123456789…5 tokens truncated…0123456789"
+        );
+
+        let response = if execution_id == "failed" {
+            host.wait(Request::new(proto::WaitRequest {
+                session_id: opened.session_id,
+                cell_id: cell_id.clone(),
+                wait_id: "wait-after-yield".to_string(),
+                yield_time_ms: 60_000,
+                max_output_tokens: Some(7),
+            }))
+            .await
+            .expect("script failures are completed results, not RPC failures")
+            .into_inner()
+        } else {
+            // The event is ordered after text in the runtime FIFO, so the
+            // termination observes it without a sleep or racing the cell.
+            assert!(matches!(
+                lease.next().await.unwrap().unwrap().event,
+                Some(proto::session_event::Event::Notification(
+                    proto::Notification { text, .. }
+                )) if text == "ready to terminate"
+            ));
+            host.terminate(Request::new(proto::TerminateRequest {
+                session_id: opened.session_id,
+                cell_id: cell_id.clone(),
+                max_output_tokens: Some(7),
+            }))
+            .await
+            .expect("terminate with an independent output budget")
+            .into_inner()
+        };
+        let Some(proto::wait_response::State::LiveCell(last)) = response.state else {
+            panic!("expected final observation from the live cell");
+        };
+        let expected_text = if execution_id == "failed" {
+            let Some(proto::execution_outcome::Outcome::Completed(proto::ExecutionCompleted {
+                error_text: Some(diagnostic),
+            })) = &last.outcome
+            else {
+                panic!(
+                    "expected completed failure with a diagnostic, got {:?}",
+                    last.outcome
+                );
+            };
+            assert_eq!(diagnostic, &"0123456789".repeat(8));
+            "Warning: truncated output (original token count: 24)\nTotal output lines: 2\n\nScript error:\n…17 tokens truncated…67890123456789"
+        } else {
+            assert_eq!(
+                last.outcome,
+                Some(proto::execution_outcome::Outcome::Terminated(
+                    proto::ExecutionTerminated {}
+                ))
+            );
+            "Warning: truncated output (original token count: 20)\nTotal output lines: 1\n\n01234567890123…13 tokens truncated…67890123456789"
+        };
+        assert_eq!(last.cell_id, cell_id);
+        assert_eq!(only_text(&last.content_items), expected_text);
+    }
+}
+
+#[tokio::test]
+async fn unknown_output_truncation_is_rejected_before_opening_a_session() {
+    let host = GrpcCodeModeHost::new();
+    let result = host
+        .open_session(Request::new(proto::OpenSessionRequest {
+            cell_execution_limits: None,
+            output_truncation: Some(42),
+        }))
+        .await;
+    let Err(status) = result else {
+        panic!("unknown truncation mode should be rejected");
+    };
+    assert_eq!(status.code(), Code::InvalidArgument);
+}

codex-rs/code-mode-host/src/grpc/session.rs

@@ -10,7 +10,10 @@ use std::sync::PoisonError;
 use std::sync::Weak;
 
 use codex_code_mode_protocol::CellId;
-use codex_code_mode_protocol::CodeModeSessionCellExecutionLimits;
+use codex_code_mode_protocol::CodeModeOutputTruncation;
+use codex_code_mode_protocol::CodeModeSession;
+use codex_code_mode_protocol::CodeModeSessionOptions;
+use codex_code_mode_protocol::TerminateRequest;
 use codex_code_mode_protocol::WaitOutcome;
 use codex_code_mode_protocol::grpc as proto;
 use codex_code_mode_protocol::host::MAX_PENDING_DELEGATE_CALLS;
@@ -107,7 +110,7 @@ impl GrpcHostState {
 
     pub(super) fn open_session(
         self: &Arc<Self>,
-        limits: CodeModeSessionCellExecutionLimits,
+        options: CodeModeSessionOptions,
     ) -> Result<GrpcStream<proto::SessionEvent>, Status> {
         let id = Uuid::new_v4();
         let (events, receiver) = mpsc::channel(OUTGOING_CHANNEL_CAPACITY);
@@ -118,13 +121,18 @@ impl GrpcHostState {
             event_sender,
             closed,
             Arc::clone(&self.delegate_permits),
-            limits,
+            options,
         );
         events
             .try_send(Ok(proto::SessionEvent {
                 event: Some(proto::session_event::Event::Opened(proto::SessionOpened {
                     session_id: id.to_string(),
-                    output_truncation: None,
+                    output_truncation: match session.runtime.output_truncation() {
+                        CodeModeOutputTruncation::Server => {
+                            Some(proto::OutputTruncation::Server.into())
+                        }
+                        CodeModeOutputTruncation::Client => None,
+                    },
                 })),
             }))
             .map_err(|_| Status::internal("failed to publish the opened code-mode session"))?;
@@ -213,7 +221,7 @@ impl GrpcSession {
         events: EventSender,
         closed: CancellationToken,
         delegate_permits: Arc<Semaphore>,
-        limits: CodeModeSessionCellExecutionLimits,
+        options: CodeModeSessionOptions,
     ) -> Arc<Self> {
         Arc::new_cyclic(|weak: &Weak<Self>| {
             let failure_session = weak.clone();
@@ -225,9 +233,9 @@ impl GrpcSession {
             });
             Self {
                 id,
-                runtime: Arc::new(InProcessCodeModeSession::with_task_failure_handler(
-                    failure_handler,
-                    limits,
+                runtime: Arc::new(InProcessCodeModeSession::with_options(
+                    options,
+                    Some(failure_handler),
                 )),
                 closed,
                 state: Mutex::new(SessionState::default()),
@@ -258,13 +266,13 @@ impl GrpcSession {
         result
     }
 
-    pub(super) async fn terminate(&self, cell_id: CellId) -> Result<WaitOutcome, Status> {
+    pub(super) async fn terminate(&self, request: TerminateRequest) -> Result<WaitOutcome, Status> {
         tokio::select! {
             biased;
             _ = self.closed.cancelled() => {
                 Err(Status::cancelled("code-mode session is closed"))
             }
-            result = self.runtime.terminate(cell_id.into()) => {
+            result = self.runtime.terminate(request) => {
                 result.map_err(Status::failed_precondition)
             }
         }
@@ -347,7 +355,7 @@ impl GrpcSession {
         if let Some(cell_id) = cell_id {
             let session = Arc::clone(self);
             tokio::spawn(async move {
-                let _ = session.terminate(CellId::new(cell_id)).await;
+                let _ = session.terminate(CellId::new(cell_id).into()).await;
             });
         }
     }