Add a feature flag for code-mode host output truncation (#58159)

## What changed

Add the disabled-by-default `code_mode_host_output_truncation` feature flag. When enabled, Core requests server-side output truncation for its built-in code-mode session providers over both stdio and gRPC. Injected providers retain control of their session settings, and explicit session options pass through unchanged.

## Testing

Add coverage for both transports verifying that the host's output-limit notice survives a zero-token observation budget. Extend the shared-provider test to exercise an injected provider with the new flag enabled.

GitOrigin-RevId: edbb65e930bf544e3d2f8384767037a6fb82d451

codex-rs/core/config.schema.json

@@ -924,6 +924,9 @@
             "code_mode_host_grpc": {
               "type": "boolean"
             },
+            "code_mode_host_output_truncation": {
+              "type": "boolean"
+            },
             "code_mode_interrupt": {
               "type": "boolean"
             },
@@ -7589,6 +7592,9 @@
         "code_mode_host_grpc": {
           "type": "boolean"
         },
+        "code_mode_host_output_truncation": {
+          "type": "boolean"
+        },
         "code_mode_interrupt": {
           "type": "boolean"
         },

codex-rs/core/src/code_mode_host.rs

@@ -1,4 +1,4 @@
-//! Owns the local gRPC code-mode client and host process. Sessions share a lazy
+//! Owns Core's local code-mode output rollout and gRPC host. Sessions share a lazy
 //! HTTP/2 connection. Only session admission can replace a dead host; metrics
 //! use the current generation and cannot reconnect. Dropping its I/O reaps the child.
 
@@ -17,6 +17,7 @@ use std::task::Context;
 use std::task::Poll;
 use std::time::Duration;
 
+use codex_code_mode::CodeModeOutputTruncation;
 use codex_code_mode::CodeModeProcessMetricsSample;
 use codex_code_mode::CodeModeSessionCellExecutionLimits;
 use codex_code_mode::CodeModeSessionOptions;
@@ -44,6 +45,45 @@ use tonic::transport::Endpoint;
 use tower::ServiceExt;
 use tower::service_fn;
 
+/// Only wraps Core's built-in providers; injected providers own their session settings.
+pub(crate) struct LocalServerOutputProvider {
+    pub(crate) inner: Arc<dyn CodeModeSessionProvider>,
+}
+
+impl CodeModeSessionProvider for LocalServerOutputProvider {
+    fn availability(&self) -> Result<(), String> {
+        self.inner.availability()
+    }
+
+    fn sample_process_metrics(
+        &self,
+    ) -> CodeModeSessionResultFuture<'_, CodeModeProcessMetricsSample> {
+        self.inner.sample_process_metrics()
+    }
+
+    fn create_session(&self) -> CodeModeSessionProviderFuture<'_> {
+        self.create_session_with_limits(CodeModeSessionCellExecutionLimits::default())
+    }
+
+    fn create_session_with_limits<'a>(
+        &'a self,
+        limits: CodeModeSessionCellExecutionLimits,
+    ) -> CodeModeSessionProviderFuture<'a> {
+        self.inner
+            .create_session_with_options(CodeModeSessionOptions {
+                cell_execution_limits: limits,
+                output_truncation: CodeModeOutputTruncation::Server,
+            })
+    }
+
+    fn create_session_with_options<'a>(
+        &'a self,
+        options: CodeModeSessionOptions,
+    ) -> CodeModeSessionProviderFuture<'a> {
+        self.inner.create_session_with_options(options)
+    }
+}
+
 pub struct ProcessOwnedGrpcCodeModeSessionProvider {
     host_program: PathBuf,
     provider: OnceLock<GrpcCodeModeSessionProvider>,

codex-rs/core/src/thread_manager.rs

@@ -756,16 +756,25 @@ impl ThreadManager {
         config: &Config,
         host_program: PathBuf,
     ) -> Arc<dyn CodeModeSessionProvider> {
-        if config.features.enabled(Feature::CodeModeHostGrpc) {
-            Arc::new(
+        let provider: Arc<dyn CodeModeSessionProvider> =
+            if config.features.enabled(Feature::CodeModeHostGrpc) {
+                Arc::new(
                 crate::code_mode_host::ProcessOwnedGrpcCodeModeSessionProvider::with_host_program(
                     host_program,
                 ),
             )
+            } else {
+                Arc::new(ProcessOwnedCodeModeSessionProvider::with_host_program(
+                    host_program,
+                ))
+            };
+        if config
+            .features
+            .enabled(Feature::CodeModeHostOutputTruncation)
+        {
+            Arc::new(crate::code_mode_host::LocalServerOutputProvider { inner: provider })
         } else {
-            Arc::new(ProcessOwnedCodeModeSessionProvider::with_host_program(
-                host_program,
-            ))
+            provider
         }
     }
 

codex-rs/core/src/thread_manager_tests.rs

@@ -943,6 +943,10 @@ async fn code_mode_session_provider_is_shared_across_threads() {
     config.codex_home = temp_dir.path().join("codex-home").abs();
     config.cwd = config.codex_home.abs();
     std::fs::create_dir_all(&config.codex_home).expect("create codex home");
+    config
+        .features
+        .enable(Feature::CodeModeHostOutputTruncation)
+        .expect("enable the built-in provider's server truncation feature");
 
     let provider: Arc<dyn CodeModeSessionProvider> = Arc::new(DisabledCodeModeSessionProvider);
     let manager = ThreadManager::with_models_provider_and_home_for_tests(
@@ -951,6 +955,9 @@ async fn code_mode_session_provider_is_shared_across_threads() {
         config.codex_home.to_path_buf(),
         Arc::new(codex_exec_server::EnvironmentManager::default_for_tests()),
     )
+    // Exercise the production factory with the startup flag before overriding
+    // its lazy built-in provider. No host is started by this test.
+    .with_code_mode_host_program_for_tests(temp_dir.path().join("unused-code-mode-host"), &config)
     .with_code_mode_session_provider(Arc::clone(&provider));
     let first = manager
         .start_thread(StartThreadOptions::new(config.clone()))

codex-rs/core/tests/suite/code_mode.rs

@@ -4718,6 +4718,49 @@ Error:\ boom\n
     Ok(())
 }
 
+// An overflow notice is outside the observation budget. This distinguishes
+// actual server formatting from Core silently repeating client truncation.
+#[test_case(false; "stdio")]
+#[test_case(true; "grpc")]
+#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+async fn code_mode_collector_notice_survives_the_observation_budget(
+    grpc_stdio: bool,
+) -> Result<()> {
+    skip_if_no_network!(Ok(()));
+
+    let server = responses::start_mock_server().await;
+    let (_test, follow_up) = run_code_mode_turn_with_config(
+        &server,
+        "emit enough text to overflow the collector",
+        r#"// @exec: {"max_output_tokens": 0}
+text("begin");
+text("x".repeat(2 * 1024 * 1024));
+text("end");
+"#,
+        move |config| {
+            let _ = config.features.enable(Feature::CodeModeHost);
+            let _ = config
+                .features
+                .enable(Feature::CodeModeHostOutputTruncation);
+            if grpc_stdio {
+                let _ = config.features.enable(Feature::CodeModeHostGrpc);
+            } else {
+                let _ = config.features.disable(Feature::CodeModeHostGrpc);
+            }
+        },
+    )
+    .await?;
+    let (output, success) =
+        custom_tool_output_body_and_success(&follow_up.single_request(), "call-1");
+    assert_ne!(success, Some(false), "{output}");
+    assert!(
+        output.contains("Code Mode session output limit reached")
+            && output.contains("text segments, 0 images, 0 audio items"),
+        "host omission notice must survive Core's zero-token exec budget: {output}"
+    );
+    Ok(())
+}
+
 #[cfg_attr(windows, ignore = "no exec_command on Windows")]
 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
 async fn code_mode_exec_surfaces_handler_errors_as_exceptions() -> Result<()> {