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/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,
         )