Add code-mode output truncation protocol contracts (#57272)

## What changed

Add session options for client- or server-owned output truncation and optional `max_output_tokens` budgets for wait and terminate observations. Introduce `TerminateRequest` and update session implementations and callers to use it.

Extend the stdio and gRPC schemas with optional truncation settings and budgets, and define the `bounded-output-v1` capability. Existing sessions retain client-side truncation; providers reject server mode by default. This change defines the contracts without enabling server-side truncation.

## Testing

Extend protocol serialization coverage for `outputTruncation: "server"` while preserving existing wire shapes when optional fields are absent.

GitOrigin-RevId: 22145897ea16d385c5ba9e7549109f7a669fceb4

codex-rs/code-mode-protocol/src/grpc/codex.code_mode.v1.proto

@@ -45,6 +45,14 @@ service CodeModeHost {
 
 message OpenSessionRequest {
   optional SessionCellExecutionLimits cell_execution_limits = 1;
+  // Absent preserves the legacy client-side truncation behavior.
+  optional OutputTruncation output_truncation = 2;
+}
+
+enum OutputTruncation {
+  OUTPUT_TRUNCATION_UNSPECIFIED = 0;
+  OUTPUT_TRUNCATION_CLIENT = 1;
+  OUTPUT_TRUNCATION_SERVER = 2;
 }
 
 message SessionCellExecutionLimits {
@@ -64,6 +72,8 @@ message SessionEvent {
 
 message SessionOpened {
   string session_id = 1;
+  // Absent (including from an older host) means the client owns truncation.
+  optional OutputTruncation output_truncation = 2;
 }
 
 message CloseSessionRequest {
@@ -182,6 +192,8 @@ message WaitRequest {
   string cell_id = 2;
   string wait_id = 3;
   uint64 yield_time_ms = 4;
+  // Only set for a session that selected server-side output truncation.
+  optional uint64 max_output_tokens = 5;
 }
 
 message WaitResponse {
@@ -211,6 +223,8 @@ message YieldObservationResponse {}
 message TerminateRequest {
   string session_id = 1;
   string cell_id = 2;
+  // The terminating observation has its own output budget.
+  optional uint64 max_output_tokens = 3;
 }
 
 message ExecutionOutcome {

codex-rs/code-mode-protocol/src/host/host_tests.rs

@@ -27,6 +27,7 @@ use super::WireContentItem;
 use super::WireExecuteRequest;
 use super::WireImageDetail;
 use super::WireNestedToolCall;
+use super::WireOutputTruncation;
 use super::WireResult;
 use super::WireRuntimeResponse;
 use super::WireSessionCellExecutionLimits;
@@ -249,6 +250,7 @@ fn open_session_serializes_optional_cell_execution_limits() {
                 max_yield_time_ms: Some(250),
                 max_heap_size_bytes: Some(16 * 1024 * 1024),
             }),
+            output_truncation: None,
         },
         json!({
             "method": "session/open",
@@ -304,9 +306,23 @@ fn client_to_host_v1_variants_are_pinned() {
             HostRequest::OpenSession {
                 session_id: session_id(),
                 cell_execution_limits: None,
+                output_truncation: None,
             },
             json!({ "method": "session/open", "sessionId": "session-1" }),
         ),
+        (
+            request_id(/*value*/ 10),
+            HostRequest::OpenSession {
+                session_id: session_id(),
+                cell_execution_limits: None,
+                output_truncation: Some(WireOutputTruncation::Server),
+            },
+            json!({
+                "method": "session/open",
+                "sessionId": "session-1",
+                "outputTruncation": "server",
+            }),
+        ),
         (
             request_id(/*value*/ 2),
             HostRequest::Execute {
@@ -353,6 +369,7 @@ fn client_to_host_v1_variants_are_pinned() {
                     cell_id: cell_id("cell-1"),
                     yield_time_ms: 50,
                 },
+                max_output_tokens: None,
             },
             json!({
                 "method": "session/wait",
@@ -365,6 +382,7 @@ fn client_to_host_v1_variants_are_pinned() {
             HostRequest::Terminate {
                 session_id: session_id(),
                 cell_id: cell_id("cell-1"),
+                max_output_tokens: None,
             },
             json!({
                 "method": "session/terminate",

codex-rs/code-mode-protocol/src/host/message.rs

@@ -15,6 +15,7 @@ use super::SupportedProtocolVersions;
 use super::WireCellId;
 use super::WireExecuteRequest;
 use super::WireNestedToolCall;
+use super::WireOutputTruncation;
 use super::WireRuntimeResponse;
 use super::WireSessionCellExecutionLimits;
 use super::WireWaitOutcome;
@@ -186,6 +187,10 @@ pub enum HostRequest {
         session_id: SessionId,
         #[serde(default, skip_serializing_if = "Option::is_none")]
         cell_execution_limits: Option<WireSessionCellExecutionLimits>,
+        /// Sent only when output truncation support was agreed at connection hello.
+        /// Omission preserves the legacy session behavior and wire shape.
+        #[serde(default, skip_serializing_if = "Option::is_none")]
+        output_truncation: Option<WireOutputTruncation>,
     },
     #[serde(rename = "session/execute")]
     Execute {
@@ -196,11 +201,15 @@ pub enum HostRequest {
     Wait {
         session_id: SessionId,
         request: WireWaitRequest,
+        #[serde(default, skip_serializing_if = "Option::is_none")]
+        max_output_tokens: Option<usize>,
     },
     #[serde(rename = "session/terminate")]
     Terminate {
         session_id: SessionId,
         cell_id: WireCellId,
+        #[serde(default, skip_serializing_if = "Option::is_none")]
+        max_output_tokens: Option<usize>,
     },
     #[serde(rename = "session/shutdown")]
     ShutdownSession { session_id: SessionId },

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

@@ -12,6 +12,12 @@ mod types;
 /// Maximum number of unresolved delegate callbacks allowed per host connection.
 pub const MAX_PENDING_DELEGATE_CALLS: usize = 1_024;
 
+/// Stdio connection support for the optional `session/open.outputTruncation` field.
+/// Each session selects its own setting after the connection hello.
+pub const BOUNDED_OUTPUT_CAPABILITY: &str = "bounded-output-v1";
+/// Maximum retained bytes between server-truncated observations.
+pub const BOUNDED_OUTPUT_MAX_BYTES: usize = 1024 * 1024;
+
 /// Negotiated support for cell execution resource limits on `session/open`.
 pub const SESSION_RESOURCE_LIMITS_CAPABILITY: &str = "session-cell-execution-resource-limits";
 
@@ -38,6 +44,7 @@ pub use payload::WireContentItem;
 pub use payload::WireExecuteRequest;
 pub use payload::WireImageDetail;
 pub use payload::WireNestedToolCall;
+pub use payload::WireOutputTruncation;
 pub use payload::WireRuntimeResponse;
 pub use payload::WireSessionCellExecutionLimits;
 pub use payload::WireToolDefinition;

codex-rs/code-mode-protocol/src/host/payload.rs

@@ -8,6 +8,7 @@ use serde_json::Value as JsonValue;
 
 use crate::CellId;
 use crate::CodeModeNestedToolCall;
+use crate::CodeModeOutputTruncation;
 use crate::CodeModeSessionCellExecutionLimits;
 use crate::CodeModeToolKind;
 use crate::ExecuteRequest;
@@ -51,6 +52,32 @@ impl TryFrom<WireSessionCellExecutionLimits> for CodeModeSessionCellExecutionLim
     }
 }
 
+/// The output truncation setting carried by a V1 session-open request.
+#[derive(Clone, Copy, Debug, Deserialize, Eq, PartialEq, Serialize)]
+#[serde(rename_all = "camelCase")]
+pub enum WireOutputTruncation {
+    Client,
+    Server,
+}
+
+impl From<CodeModeOutputTruncation> for WireOutputTruncation {
+    fn from(value: CodeModeOutputTruncation) -> Self {
+        match value {
+            CodeModeOutputTruncation::Client => Self::Client,
+            CodeModeOutputTruncation::Server => Self::Server,
+        }
+    }
+}
+
+impl From<WireOutputTruncation> for CodeModeOutputTruncation {
+    fn from(value: WireOutputTruncation) -> Self {
+        match value {
+            WireOutputTruncation::Client => Self::Client,
+            WireOutputTruncation::Server => Self::Server,
+        }
+    }
+}
+
 /// A cell identifier with a wire representation owned by protocol V1.
 #[derive(Clone, Debug, Deserialize, Eq, Hash, PartialEq, Serialize)]
 #[serde(transparent)]
@@ -211,7 +238,8 @@ impl TryFrom<WireExecuteRequest> for ExecuteRequest {
     }
 }
 
-/// The complete wait request shape supported by protocol V1.
+/// The nested wait payload for protocol V1. Its optional budget is carried by
+/// the enclosing `HostRequest::Wait`.
 #[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)]
 #[serde(deny_unknown_fields)]
 pub struct WireWaitRequest {
@@ -233,6 +261,7 @@ impl From<WireWaitRequest> for WaitRequest {
         Self {
             cell_id: value.cell_id.into(),
             yield_time_ms: value.yield_time_ms,
+            max_output_tokens: None,
         }
     }
 }

codex-rs/code-mode-protocol/src/lib.rs

@@ -35,14 +35,17 @@ pub use runtime::ExecuteRequest;
 pub use runtime::ExecuteToPendingOutcome;
 pub use runtime::MissingCodeModeHostDuration;
 pub use runtime::RuntimeResponse;
+pub use runtime::TerminateRequest;
 pub use runtime::WaitOutcome;
 pub use runtime::WaitRequest;
 pub use runtime::WaitToPendingOutcome;
 pub use runtime::WaitToPendingRequest;
 pub use session::CellId;
+pub use session::CodeModeOutputTruncation;
 pub use session::CodeModeSession;
 pub use session::CodeModeSessionCellExecutionLimits;
 pub use session::CodeModeSessionDelegate;
+pub use session::CodeModeSessionOptions;
 pub use session::CodeModeSessionProvider;
 pub use session::CodeModeSessionProviderFuture;
 pub use session::CodeModeSessionResultFuture;

codex-rs/code-mode-protocol/src/runtime.rs

@@ -29,6 +29,26 @@ pub struct ExecuteRequest {
 pub struct WaitRequest {
     pub cell_id: CellId,
     pub yield_time_ms: u64,
+    /// Limit for this observation in Server mode; ignored in Client mode.
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub max_output_tokens: Option<usize>,
+}
+
+#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)]
+pub struct TerminateRequest {
+    pub cell_id: CellId,
+    /// Limit for the final observation in Server mode; ignored in Client mode.
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub max_output_tokens: Option<usize>,
+}
+
+impl From<CellId> for TerminateRequest {
+    fn from(cell_id: CellId) -> Self {
+        Self {
+            cell_id,
+            max_output_tokens: None,
+        }
+    }
 }
 
 #[derive(Clone, Debug, Deserialize, PartialEq, Serialize)]

codex-rs/code-mode-protocol/src/session.rs

@@ -12,6 +12,7 @@ use tokio_util::sync::CancellationToken;
 use crate::CodeModeNestedToolCall;
 use crate::ExecuteRequest;
 use crate::RuntimeResponse;
+use crate::TerminateRequest;
 use crate::WaitOutcome;
 use crate::WaitRequest;
 
@@ -30,6 +31,33 @@ pub struct CodeModeSessionCellExecutionLimits {
     pub max_heap_size_bytes: Option<usize>,
 }
 
+/// The owner of the output budget for each exec or wait observation.
+#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
+pub enum CodeModeOutputTruncation {
+    /// Return raw output; the caller applies its own budget and formatting.
+    #[default]
+    Client,
+    /// The host applies the exec, wait, or terminate observation's budget before
+    /// transmitting output (10,000 tokens by default). Clients must not truncate
+    /// it again or append the script error: status and error metadata remain separate.
+    ///
+    /// Between observations, the host retains at most
+    /// [`crate::host::BOUNDED_OUTPUT_MAX_BYTES`] of output payload. It also
+    /// bounds the storage for retained items separately, without charging it
+    /// against the payload or observation token budgets.
+    /// Ordinary text notices report overflow, even for a zero-token budget.
+    /// Media are indivisible and can be omitted if they exceed the collector.
+    /// Without overflow, formatting matches Core's client policy.
+    Server,
+}
+
+/// Settings shared by every cell in one code-mode session.
+#[derive(Clone, Debug, Default, Eq, PartialEq)]
+pub struct CodeModeSessionOptions {
+    pub cell_execution_limits: CodeModeSessionCellExecutionLimits,
+    pub output_truncation: CodeModeOutputTruncation,
+}
+
 #[derive(Clone, Debug, Deserialize, Eq, Hash, PartialEq, Serialize)]
 pub struct CellId(String);
 
@@ -152,6 +180,12 @@ impl CodeModeSessionDelegate for NoopCodeModeSessionDelegate {
 /// must keep those values isolated. Implementations may execute cells
 /// in-process or remotely.
 pub trait CodeModeSession: Send + Sync {
+    /// Identifies who applies output truncation for this session's responses.
+    /// Providers must not change modes when reconnecting to a new host.
+    fn output_truncation(&self) -> CodeModeOutputTruncation {
+        CodeModeOutputTruncation::Client
+    }
+
     /// Executes a cell, yielding its foreground observation if `preempt` is signaled.
     /// The cell continues running after its observation yields.
     fn execute<'a>(
@@ -162,13 +196,19 @@ pub trait CodeModeSession: Send + Sync {
     ) -> CodeModeSessionResultFuture<'a, StartedCell>;
 
     /// Waits for a cell, yielding the observation if `preempt` is signaled while the cell continues.
+    /// Server mode applies the request's output limit; Client mode ignores it.
     fn wait<'a>(
         &'a self,
         request: WaitRequest,
         preempt: Option<CancellationToken>,
     ) -> CodeModeSessionResultFuture<'a, WaitOutcome>;
 
-    fn terminate<'a>(&'a self, cell_id: CellId) -> CodeModeSessionResultFuture<'a, WaitOutcome>;
+    /// Terminates a cell, applying the request's limit to its final observation
+    /// in Server mode; Client mode ignores the limit.
+    fn terminate<'a>(
+        &'a self,
+        request: TerminateRequest,
+    ) -> CodeModeSessionResultFuture<'a, WaitOutcome>;
 
     fn shutdown<'a>(&'a self) -> CodeModeSessionResultFuture<'a, ()>;
 }
@@ -201,6 +241,25 @@ pub trait CodeModeSessionProvider: Send + Sync {
             })
         }
     }
+
+    /// Requests output handling for one session, without changing other sessions
+    /// created by this provider. Server mode requires host support.
+    fn create_session_with_options<'a>(
+        &'a self,
+        options: CodeModeSessionOptions,
+    ) -> CodeModeSessionProviderFuture<'a> {
+        match options.output_truncation {
+            CodeModeOutputTruncation::Client => {
+                self.create_session_with_limits(options.cell_execution_limits)
+            }
+            CodeModeOutputTruncation::Server => Box::pin(async {
+                Err(
+                    "code-mode session provider does not support server output truncation"
+                        .to_string(),
+                )
+            }),
+        }
+    }
 }
 
 #[cfg(test)]