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-host/src/grpc/mod.rs

@@ -234,6 +234,7 @@ impl GrpcCodeModeHost {
         let request = WaitRequest {
             cell_id: CellId::new(request.cell_id),
             yield_time_ms: request.yield_time_ms,
+            max_output_tokens: None,
         };
         let outcome = tokio::select! {
             biased;

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

@@ -255,6 +255,7 @@ async fn session_closure_cancels_pending_termination() {
     let termination = host.terminate(Request::new(proto::TerminateRequest {
         session_id,
         cell_id,
+        max_output_tokens: None,
     }));
     tokio::pin!(termination);
     assert!(termination.as_mut().now_or_never().is_none());

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

@@ -14,6 +14,7 @@ pub(super) async fn open_session(
     let mut stream = host
         .open_session(Request::new(proto::OpenSessionRequest {
             cell_execution_limits: None,
+            output_truncation: None,
         }))
         .await
         .expect("open code-mode session")
@@ -120,6 +121,7 @@ async fn execute_stream_starts_immediately_and_wait_preserves_missing_cells() {
             cell_id: cell_id.clone(),
             wait_id: "wait-1".to_string(),
             yield_time_ms: 60_000,
+            max_output_tokens: None,
         }))
         .await
         .expect("wait for completion")
@@ -163,6 +165,7 @@ async fn execute_stream_starts_immediately_and_wait_preserves_missing_cells() {
             cell_id: "missing-cell".to_string(),
             wait_id: "wait-missing".to_string(),
             yield_time_ms: 1,
+            max_output_tokens: None,
         }))
         .await
         .expect("missing cells remain successful wait outcomes")
@@ -355,6 +358,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,
         }))
         .await
         .unwrap_err();
@@ -363,6 +367,7 @@ async fn cancellation_before_wait_admission_is_preserved() {
     host.terminate(Request::new(proto::TerminateRequest {
         session_id,
         cell_id,
+        max_output_tokens: None,
     }))
     .await
     .unwrap();

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

@@ -124,6 +124,7 @@ impl GrpcHostState {
             .try_send(Ok(proto::SessionEvent {
                 event: Some(proto::session_event::Event::Opened(proto::SessionOpened {
                     session_id: id.to_string(),
+                    output_truncation: None,
                 })),
             }))
             .map_err(|_| Status::internal("failed to publish the opened code-mode session"))?;
@@ -263,7 +264,7 @@ impl GrpcSession {
             _ = self.closed.cancelled() => {
                 Err(Status::cancelled("code-mode session is closed"))
             }
-            result = self.runtime.terminate(cell_id) => {
+            result = self.runtime.terminate(cell_id.into()) => {
                 result.map_err(Status::failed_precondition)
             }
         }

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

@@ -127,6 +127,7 @@ async fn handshake_and_multiple_session_lifecycles_are_ordered() {
                 request: HostRequest::OpenSession {
                     session_id: session_id(id),
                     cell_execution_limits: None,
+                    output_truncation: None,
                 },
             })
             .await
@@ -198,6 +199,7 @@ async fn disconnect_cancels_a_backpressured_host_writer() {
             request: HostRequest::OpenSession {
                 session_id: session_id("backpressured-session"),
                 cell_execution_limits: None,
+                output_truncation: None,
             },
         })
         .await
@@ -339,6 +341,7 @@ async fn incompatible_or_invalid_handshake_is_rejected() {
             request: HostRequest::OpenSession {
                 session_id: session_id("session-1"),
                 cell_execution_limits: None,
+                output_truncation: None,
             },
         })
         .await
@@ -405,6 +408,7 @@ async fn session_id_cannot_be_reused_after_shutdown() {
             HostRequest::OpenSession {
                 session_id: id.clone(),
                 cell_execution_limits: None,
+                output_truncation: None,
             },
         ),
         (
@@ -433,6 +437,7 @@ async fn session_id_cannot_be_reused_after_shutdown() {
             request: HostRequest::OpenSession {
                 session_id: id,
                 cell_execution_limits: None,
+                output_truncation: None,
             },
         })
         .await
@@ -556,7 +561,7 @@ async fn execute_request_id_remains_active_until_initial_response() {
     state
         .session(&session_id)
         .expect("session")
-        .terminate(cell_id.into())
+        .terminate(codex_code_mode_protocol::CellId::from(cell_id).into())
         .await
         .expect("terminate cell");
     state.disconnect().await;

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

@@ -382,6 +382,7 @@ impl HostState {
             HostRequest::OpenSession {
                 session_id,
                 cell_execution_limits,
+                ..
             } => {
                 let result = CodeModeSessionCellExecutionLimits::try_from(
                     cell_execution_limits.unwrap_or_default(),
@@ -452,6 +453,7 @@ impl HostState {
             HostRequest::Wait {
                 session_id,
                 request,
+                ..
             } => {
                 let result = match self.session(&session_id) {
                     Ok(session) => {
@@ -476,15 +478,20 @@ impl HostState {
             HostRequest::Terminate {
                 session_id,
                 cell_id,
+                ..
             } => {
                 let result = match self.session(&session_id) {
-                    Ok(session) => session.terminate(cell_id.into()).await.and_then(|outcome| {
-                        let outcome = outcome.with_code_mode_host_duration(received_at.elapsed());
-                        Ok(HostResponse::WaitCompleted {
-                            outcome: WireWaitOutcome::try_from(outcome)
-                                .map_err(|error| error.to_string())?,
-                        })
-                    }),
+                    Ok(session) => session
+                        .terminate(codex_code_mode_protocol::CellId::from(cell_id).into())
+                        .await
+                        .and_then(|outcome| {
+                            let outcome =
+                                outcome.with_code_mode_host_duration(received_at.elapsed());
+                            Ok(HostResponse::WaitCompleted {
+                                outcome: WireWaitOutcome::try_from(outcome)
+                                    .map_err(|error| error.to_string())?,
+                            })
+                        }),
                     Err(err) => Err(err),
                 };
                 self.respond(request_id, result);

codex-rs/code-mode-host/tests/grpc.rs

@@ -327,6 +327,7 @@ async fn grpc_early_yield_preserves_cell_and_later_waits() -> Result<()> {
         WaitRequest {
             cell_id: cell_id.clone(),
             yield_time_ms: 60_000,
+            max_output_tokens: None,
         },
         Some(wait_signal.clone()),
     );
@@ -347,6 +348,7 @@ async fn grpc_early_yield_preserves_cell_and_later_waits() -> Result<()> {
             WaitRequest {
                 cell_id: cell_id.clone(),
                 yield_time_ms: 60_000,
+                max_output_tokens: None,
             },
             /*preempt*/ None,
         ),
@@ -397,13 +399,17 @@ async fn shutdown_immediately_rejects_new_operations() -> Result<()> {
                 WaitRequest {
                     cell_id: cell_id("missing"),
                     yield_time_ms: 1,
+                    max_output_tokens: None,
                 },
                 /*preempt*/ None,
             )
             .await,
         Err(expected.clone())
     );
-    assert_eq!(session.terminate(cell_id("missing")).await, Err(expected));
+    assert_eq!(
+        session.terminate(cell_id("missing").into()).await,
+        Err(expected)
+    );
 
     shutdown.await.map_err(anyhow::Error::msg)?;
     Ok(())
@@ -452,6 +458,7 @@ async fn cancelling_execution_before_admission_keeps_the_session_usable() -> Res
                     WaitRequest {
                         cell_id: abandoned_cell.clone(),
                         yield_time_ms: 1,
+                        max_output_tokens: None,
                     },
                     /*preempt*/ None,
                 )
@@ -512,6 +519,7 @@ async fn dropping_a_started_cell_off_runtime_terminates_its_buffered_remote_exec
                     WaitRequest {
                         cell_id: abandoned_cell.clone(),
                         yield_time_ms: 1,
+                        max_output_tokens: None,
                     },
                     /*preempt*/ None,
                 )
@@ -540,6 +548,7 @@ async fn dropping_a_started_cell_off_runtime_terminates_its_buffered_remote_exec
                     WaitRequest {
                         cell_id: abandoned_cell.clone(),
                         yield_time_ms: 1,
+                        max_output_tokens: None,
                     },
                     /*preempt*/ None,
                 )
@@ -751,6 +760,7 @@ async fn concurrent_wait_rejects_without_displacing_the_active_observer() -> Res
         WaitRequest {
             cell_id: running_cell.clone(),
             yield_time_ms: 100,
+            max_output_tokens: None,
         },
     )
     .await?;
@@ -762,6 +772,7 @@ async fn concurrent_wait_rejects_without_displacing_the_active_observer() -> Res
                 WaitRequest {
                     cell_id: running_cell.clone(),
                     yield_time_ms: 60_000,
+                    max_output_tokens: None,
                 },
                 /*preempt*/ None,
             ),
@@ -785,7 +796,7 @@ async fn concurrent_wait_rejects_without_displacing_the_active_observer() -> Res
         })
     );
     let actual = session
-        .terminate(running_cell.clone())
+        .terminate(running_cell.clone().into())
         .await
         .map_err(anyhow::Error::msg)?;
     assert_eq!(
@@ -838,6 +849,7 @@ async fn dropping_a_wait_retires_its_observer_before_the_next_wait() -> Result<(
         WaitRequest {
             cell_id: running_cell.clone(),
             yield_time_ms: 60_000,
+            max_output_tokens: None,
         },
     )
     .await?;
@@ -850,6 +862,7 @@ async fn dropping_a_wait_retires_its_observer_before_the_next_wait() -> Result<(
             WaitRequest {
                 cell_id: running_cell.clone(),
                 yield_time_ms: 1,
+                max_output_tokens: None,
             },
             /*preempt*/ None,
         ),
@@ -866,7 +879,7 @@ async fn dropping_a_wait_retires_its_observer_before_the_next_wait() -> Result<(
         })
     );
     let actual = session
-        .terminate(running_cell.clone())
+        .terminate(running_cell.clone().into())
         .await
         .map_err(anyhow::Error::msg)?;
     assert_eq!(
@@ -1091,6 +1104,7 @@ async fn sessions_enforce_independent_yield_limits() -> Result<()> {
             WaitRequest {
                 cell_id: cell_id("1"),
                 yield_time_ms: 60_000,
+                max_output_tokens: None,
             },
             /*preempt*/ None,
         ),
@@ -1117,7 +1131,7 @@ async fn sessions_enforce_independent_yield_limits() -> Result<()> {
         text_response("1", "isolated", actual.code_mode_host_duration())
     );
     let actual = limited
-        .terminate(cell_id("1"))
+        .terminate(cell_id("1").into())
         .await
         .map_err(anyhow::Error::msg)?;
     assert_eq!(
@@ -1156,6 +1170,7 @@ async fn dropping_a_grpc_lease_retires_its_server_session() -> Result<()> {
     let mut lease = client
         .open_session(grpc::OpenSessionRequest {
             cell_execution_limits: None,
+            output_truncation: None,
         })
         .await
         .context("open raw gRPC session")?
@@ -1236,6 +1251,7 @@ async fn cached_session_recovers_after_a_remote_host_restarts() -> Result<()> {
         WaitRequest {
             cell_id: old_cell_id.clone(),
             yield_time_ms: 60_000,
+            max_output_tokens: None,
         },
     )
     .await?;
@@ -1280,7 +1296,7 @@ async fn cached_session_recovers_after_a_remote_host_restarts() -> Result<()> {
     .context("host loss did not retire the original generation's cell")?;
 
     let _replacement = HostHarness::start(&listen_url).await?;
-    assert!(session.terminate(old_cell_id.clone()).await.is_err());
+    assert!(session.terminate(old_cell_id.clone().into()).await.is_err());
 
     let mut rejected = request(r#"await tools.echo({must_not_run: true});"#);
     rejected.enabled_tools = vec![tool("echo")];
@@ -1291,6 +1307,7 @@ async fn cached_session_recovers_after_a_remote_host_restarts() -> Result<()> {
         WaitRequest {
             cell_id: old_cell_id.clone(),
             yield_time_ms: 1,
+            max_output_tokens: None,
         },
         /*preempt*/ None,
     );
@@ -1406,6 +1423,7 @@ async fn cached_session_recovers_after_a_remote_host_restarts() -> Result<()> {
             WaitRequest {
                 cell_id: replacement_cell_id.clone(),
                 yield_time_ms: 1,
+                max_output_tokens: None,
             },
             /*preempt*/ None,
         )
@@ -1420,7 +1438,7 @@ async fn cached_session_recovers_after_a_remote_host_restarts() -> Result<()> {
         })
     );
     let actual = session
-        .terminate(replacement_cell_id.clone())
+        .terminate(replacement_cell_id.clone().into())
         .await
         .map_err(anyhow::Error::msg)?;
     assert_eq!(
@@ -1437,13 +1455,14 @@ async fn cached_session_recovers_after_a_remote_host_restarts() -> Result<()> {
             WaitRequest {
                 cell_id: old_cell_id.clone(),
                 yield_time_ms: 1,
+                max_output_tokens: None,
             },
             /*preempt*/ None,
         )
         .await
         .unwrap_err();
     assert!(stale_wait.contains("stale code-mode host generation"));
-    let stale_termination = session.terminate(old_cell_id).await.unwrap_err();
+    let stale_termination = session.terminate(old_cell_id.into()).await.unwrap_err();
     assert!(stale_termination.contains("stale code-mode host generation"));
     session.shutdown().await.map_err(anyhow::Error::msg)?;
     Ok(())

codex-rs/code-mode-host/tests/grpc/admission_tests.rs

@@ -228,6 +228,7 @@ async fn routing_metadata_is_scoped_to_each_session_on_a_shared_provider() -> Re
                 WaitRequest {
                     cell_id: id.clone(),
                     yield_time_ms: 1,
+                    max_output_tokens: None,
                 },
                 /*preempt*/ None,
             )
@@ -237,7 +238,10 @@ async fn routing_metadata_is_scoped_to_each_session_on_a_shared_provider() -> Re
             outcome,
             WaitOutcome::LiveCell(RuntimeResponse::Yielded { .. })
         ));
-        session.terminate(id).await.map_err(anyhow::Error::msg)?;
+        session
+            .terminate(id.into())
+            .await
+            .map_err(anyhow::Error::msg)?;
     }
     for session in [&parent, &child] {
         execute(

codex-rs/code-mode-host/tests/grpc/image_cancellation_tests.rs

@@ -57,7 +57,7 @@ try { image(value); } catch { text(calls); }
         .await
         .context("getter did not signal startup")?;
 
-    let result = timeout(Duration::from_secs(1), session.terminate(id.clone()))
+    let result = timeout(Duration::from_secs(1), session.terminate(id.clone().into()))
         .await
         .context("workload did not stop within the cancellation deadline")?
         .map_err(anyhow::Error::msg)?;

codex-rs/code-mode-host/tests/grpc_notifications.rs

@@ -225,6 +225,7 @@ async fn completed_waits_drain_pending_notifications_before_returning() -> Resul
                 WaitRequest {
                     cell_id: cell_id("1"),
                     yield_time_ms: 5_000,
+                    max_output_tokens: None,
                 },
                 /*preempt*/ None,
             )
@@ -292,7 +293,7 @@ async fn termination_cancels_pending_notifications() -> Result<()> {
         }
     );
     let actual = session
-        .terminate(cell_id("1"))
+        .terminate(cell_id("1").into())
         .await
         .map_err(anyhow::Error::msg)?;
     assert_eq!(

codex-rs/code-mode-host/tests/grpc_tcp.rs

@@ -150,6 +150,7 @@ async fn tcp_listener_opens_a_grpc_session() -> Result<()> {
         let response = client
             .open_session(grpc::OpenSessionRequest {
                 cell_execution_limits: None,
+                output_truncation: None,
             })
             .await
             .context("failed to open gRPC code-mode session")?;

codex-rs/code-mode-host/tests/grpc_timing.rs

@@ -40,6 +40,7 @@ async fn execution_timing_includes_javascript_but_excludes_delayed_reads() -> Re
     let mut events = client
         .open_session(grpc::OpenSessionRequest {
             cell_execution_limits: None,
+            output_truncation: None,
         })
         .await?
         .into_inner();
@@ -229,6 +230,7 @@ async fn observation_timing_excludes_previous_requests_and_background_time() ->
                     WaitRequest {
                         cell_id: cell_id.clone(),
                         yield_time_ms: 50,
+                        max_output_tokens: None,
                     },
                     /*preempt*/ None,
                 )
@@ -254,7 +256,7 @@ async fn observation_timing_excludes_previous_requests_and_background_time() ->
 
         let observed = Instant::now();
         let terminated = session
-            .terminate(cell_id.clone())
+            .terminate(cell_id.clone().into())
             .await
             .map_err(anyhow::Error::msg)?;
         let duration = terminated
@@ -276,6 +278,7 @@ async fn observation_timing_excludes_previous_requests_and_background_time() ->
                 WaitRequest {
                     cell_id: cell_id.clone(),
                     yield_time_ms: 50,
+                    max_output_tokens: None,
                 },
                 /*preempt*/ None,
             )

codex-rs/code-mode-host/tests/settled.rs

@@ -109,6 +109,7 @@ async fn assert_partial_results(provider: &dyn CodeModeSessionProvider) {
                     WaitRequest {
                         cell_id: cell_id.clone(),
                         yield_time_ms: 10,
+                        max_output_tokens: None,
                     },
                     /*preempt*/ None,
                 )
@@ -133,6 +134,7 @@ async fn assert_partial_results(provider: &dyn CodeModeSessionProvider) {
                 WaitRequest {
                     cell_id: cell_id.clone(),
                     yield_time_ms: 5_000,
+                    max_output_tokens: None,
                 },
                 /*preempt*/ None,
             )

codex-rs/code-mode-host/tests/stdio.rs

@@ -206,6 +206,7 @@ async fn execute_to_terminal(
                         WaitRequest {
                             cell_id,
                             yield_time_ms: 60_000,
+                            max_output_tokens: None,
                         },
                         /*preempt*/ None,
                     )
@@ -270,6 +271,7 @@ async fn interrupt_yields_observations_without_stopping_the_cell() {
     let request = WaitRequest {
         cell_id: cell_id.clone(),
         yield_time_ms: 60_000,
+        max_output_tokens: None,
     };
     let response =
         tokio::time::timeout(Duration::from_secs(5), session.wait(request, Some(signal)))
@@ -288,6 +290,7 @@ async fn interrupt_yields_observations_without_stopping_the_cell() {
             WaitRequest {
                 cell_id,
                 yield_time_ms: 60_000,
+                max_output_tokens: None,
             },
             /*preempt*/ None,
         ),
@@ -350,6 +353,7 @@ async fn session_execution_limits_are_isolated_on_a_shared_process_host() {
             WaitRequest {
                 cell_id: cell_id("1"),
                 yield_time_ms: 60_000,
+                max_output_tokens: None,
             },
             /*preempt*/ None,
         ),
@@ -367,7 +371,7 @@ async fn session_execution_limits_are_isolated_on_a_shared_process_host() {
     );
 
     limited
-        .terminate(cell_id("1"))
+        .terminate(cell_id("1").into())
         .await
         .expect("terminate cell");
     limited.shutdown().await.expect("shutdown limited session");
@@ -468,6 +472,7 @@ text(result.value);
             WaitRequest {
                 cell_id: cell_id("3"),
                 yield_time_ms: 1,
+                max_output_tokens: None,
             },
             /*preempt*/ None,
         )
@@ -482,7 +487,7 @@ text(result.value);
         })
     );
     let actual = session
-        .terminate(cell_id("3"))
+        .terminate(cell_id("3").into())
         .await
         .expect("terminate cell");
     assert_eq!(
@@ -546,6 +551,7 @@ async fn dropping_long_wait_releases_observer_before_next_wait() {
                 WaitRequest {
                     cell_id: wait_cell_id,
                     yield_time_ms: 60_000,
+                    max_output_tokens: None,
                 },
                 /*preempt*/ None,
             )
@@ -561,6 +567,7 @@ async fn dropping_long_wait_releases_observer_before_next_wait() {
             WaitRequest {
                 cell_id: running_cell_id.clone(),
                 yield_time_ms: 1,
+                max_output_tokens: None,
             },
             /*preempt*/ None,
         ),
@@ -577,7 +584,7 @@ async fn dropping_long_wait_releases_observer_before_next_wait() {
         })
     );
     session
-        .terminate(running_cell_id)
+        .terminate(running_cell_id.into())
         .await
         .expect("terminate cell");
     session.shutdown().await.expect("shutdown remote session");
@@ -652,6 +659,7 @@ return;
                 WaitRequest {
                     cell_id: wait_cell_id,
                     yield_time_ms: 60_000,
+                    max_output_tokens: None,
                 },
                 /*preempt*/ None,
             )
@@ -995,6 +1003,7 @@ async fn child_process_loss_cleans_up_and_rebuilds_the_shared_host() {
                 WaitRequest {
                     cell_id: wait_a_cell,
                     yield_time_ms: 60_000,
+                    max_output_tokens: None,
                 },
                 /*preempt*/ None,
             )
@@ -1008,6 +1017,7 @@ async fn child_process_loss_cleans_up_and_rebuilds_the_shared_host() {
                 WaitRequest {
                     cell_id: wait_b_cell,
                     yield_time_ms: 60_000,
+                    max_output_tokens: None,
                 },
                 /*preempt*/ None,
             )
@@ -1092,6 +1102,7 @@ async fn child_process_loss_cleans_up_and_rebuilds_the_shared_host() {
             WaitRequest {
                 cell_id: cell_b.clone(),
                 yield_time_ms: 1,
+                max_output_tokens: None,
             },
             /*preempt*/ None,
         )