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-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::<Body>();
+                        } 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(())
+}

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)]