diff --git a/codex-rs/core/config.schema.json b/codex-rs/core/config.schema.json index 285af315d6..168c26dbdc 100644 --- a/codex-rs/core/config.schema.json +++ b/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" }, diff --git a/codex-rs/core/src/code_mode_host.rs b/codex-rs/core/src/code_mode_host.rs index 4d6eea1a5b..976b637178 100644 --- a/codex-rs/core/src/code_mode_host.rs +++ b/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, +} + +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, diff --git a/codex-rs/core/src/thread_manager.rs b/codex-rs/core/src/thread_manager.rs index b56485d741..800f8fa52a 100644 --- a/codex-rs/core/src/thread_manager.rs +++ b/codex-rs/core/src/thread_manager.rs @@ -756,16 +756,25 @@ impl ThreadManager { config: &Config, host_program: PathBuf, ) -> Arc { - if config.features.enabled(Feature::CodeModeHostGrpc) { - Arc::new( + let provider: Arc = + 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 } } diff --git a/codex-rs/core/src/thread_manager_tests.rs b/codex-rs/core/src/thread_manager_tests.rs index b46fe16e2e..26d480070d 100644 --- a/codex-rs/core/src/thread_manager_tests.rs +++ b/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 = 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())) diff --git a/codex-rs/core/tests/suite/code_mode.rs b/codex-rs/core/tests/suite/code_mode.rs index 42991e04e3..e9cb691bc0 100644 --- a/codex-rs/core/tests/suite/code_mode.rs +++ b/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<()> { diff --git a/codex-rs/features/src/lib.rs b/codex-rs/features/src/lib.rs index d75146e5d3..a171b1636b 100644 --- a/codex-rs/features/src/lib.rs +++ b/codex-rs/features/src/lib.rs @@ -150,6 +150,8 @@ pub enum Feature { CodeModeHost, /// Use gRPC over stdio for the process-owned code-mode host. CodeModeHostGrpc, + /// Truncate output in the process-owned code-mode host before sending it to Core. + CodeModeHostOutputTruncation, /// Establish the code-mode host connection during session startup. CodeModePrewarm, /// Terminate active code mode cells when their turn is interrupted. @@ -1281,6 +1283,12 @@ pub const FEATURES: &[FeatureSpec] = &[ stage: Stage::UnderDevelopment, default_enabled: false, }, + FeatureSpec { + id: Feature::CodeModeHostOutputTruncation, + key: "code_mode_host_output_truncation", + stage: Stage::UnderDevelopment, + default_enabled: false, + }, FeatureSpec { id: Feature::CodeModePrewarm, key: "code_mode_prewarm",