From 2dc0791c7d2cb4decda6b8ebd8b54d38d7411f64 Mon Sep 17 00:00:00 2001 From: Channing Conger Date: Fri, 9 Oct 2026 23:12:33 +0000 Subject: [PATCH] Add opt-in gRPC over stdio for the code-mode host (#52723) ## What changed - Add `grpc+stdio://` as a host transport and the disabled-by-default `code_mode_host_grpc` feature flag to select it. - Share a lazy HTTP/2 channel and host process across code-mode sessions while keeping session state isolated. - Use HTTP/2 keepalives to detect unresponsive hosts, kill and reap disconnected children, and spawn a fresh host on reconnect. Exit the host when its stdio connection closes. ## Testing Add coverage for shared-host session isolation, host exit on stdin EOF, and recovery after an unresponsive host is reaped. Exercise nested tool calls with both legacy stdio and gRPC stdio, asserting that the feature flag selects the expected transport. GitOrigin-RevId: 31389d8540de236b8e0571d3c24e29e117b75c4a --- codex-rs/Cargo.lock | 3 + codex-rs/Cargo.toml | 1 + codex-rs/code-mode-host/src/grpc_stdio.rs | 88 +++++++++ codex-rs/code-mode-host/src/lib.rs | 1 + codex-rs/code-mode-host/src/main.rs | 2 +- codex-rs/code-mode-host/src/transport.rs | 7 +- .../code-mode-host/src/transport_tests.rs | 4 + codex-rs/core/Cargo.toml | 3 + codex-rs/core/config.schema.json | 6 + codex-rs/core/src/code_mode_host.rs | 169 ++++++++++++++++++ codex-rs/core/src/code_mode_host_tests.rs | 148 +++++++++++++++ codex-rs/core/src/lib.rs | 1 + codex-rs/core/src/thread_manager.rs | 30 +++- codex-rs/core/tests/common/code_mode_host.rs | 59 ++++++ codex-rs/core/tests/common/lib.rs | 2 + codex-rs/core/tests/suite/code_mode.rs | 52 ++++-- codex-rs/features/src/lib.rs | 8 + 17 files changed, 566 insertions(+), 18 deletions(-) create mode 100644 codex-rs/code-mode-host/src/grpc_stdio.rs create mode 100644 codex-rs/core/src/code_mode_host.rs create mode 100644 codex-rs/core/src/code_mode_host_tests.rs create mode 100644 codex-rs/core/tests/common/code_mode_host.rs diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index 0e2fd03801..530060483e 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -3020,6 +3020,7 @@ dependencies = [ "eventsource-stream", "futures", "http 1.4.0", + "hyper-util", "iana-time-zone", "image", "indexmap 2.14.0", @@ -3050,6 +3051,8 @@ dependencies = [ "tokio-util", "toml 0.9.11+spec-1.1.0", "toml_edit 0.24.0+spec-1.1.0", + "tonic", + "tower", "tracing", "tracing-opentelemetry", "tracing-subscriber", diff --git a/codex-rs/Cargo.toml b/codex-rs/Cargo.toml index 0e77acfa9a..1dcbaf88c0 100644 --- a/codex-rs/Cargo.toml +++ b/codex-rs/Cargo.toml @@ -386,6 +386,7 @@ globset = "0.4" hmac = "0.12.1" http = "1.3.1" httpdate = "1.0.3" +hyper-util = { version = "0.1.20", features = ["tokio"] } iana-time-zone = "0.1.64" icu_decimal = "2.1" icu_locale_core = "2.1" diff --git a/codex-rs/code-mode-host/src/grpc_stdio.rs b/codex-rs/code-mode-host/src/grpc_stdio.rs new file mode 100644 index 0000000000..c05bf8e135 --- /dev/null +++ b/codex-rs/code-mode-host/src/grpc_stdio.rs @@ -0,0 +1,88 @@ +//! Serves one HTTP/2 gRPC connection on stdin/stdout. Stdout contains only protocol +//! bytes, and closing the connection stops the host rather than accepting another client. + +use std::io; +use std::pin::Pin; +use std::task::Context; +use std::task::Poll; + +use anyhow::Context as _; +use anyhow::Result; +use codex_code_mode_protocol::grpc::code_mode_host_server::CodeModeHostServer; +use codex_code_mode_protocol::host::MAX_FRAME_BYTES; +use futures::StreamExt; +use tokio::io::AsyncRead; +use tokio::io::AsyncWrite; +use tokio::io::ReadBuf; +use tokio_util::sync::CancellationToken; +use tonic::transport::Server; +use tonic::transport::server::Connected; + +use crate::GrpcCodeModeHost; + +pub(crate) async fn run() -> Result<()> { + let closed = CancellationToken::new(); + let connection = StdioConnection { + io: tokio::io::join(tokio::io::stdin(), tokio::io::stdout()), + closed: closed.clone(), + }; + // Keep the incoming stream open while the single connection is being served. + let incoming = futures::stream::once(async { Ok::<_, io::Error>(connection) }) + .chain(futures::stream::pending()); + let server = Server::builder() + .add_service( + CodeModeHostServer::new(GrpcCodeModeHost::new()) + .max_decoding_message_size(MAX_FRAME_BYTES) + .max_encoding_message_size(MAX_FRAME_BYTES), + ) + .serve_with_incoming(incoming); + tokio::select! { + result = server => result.context("code-mode gRPC stdio connection failed"), + _ = closed.cancelled() => Ok(()), + } +} + +struct StdioConnection { + io: tokio::io::Join, + closed: CancellationToken, +} + +impl Connected for StdioConnection { + type ConnectInfo = (); + + fn connect_info(&self) -> Self::ConnectInfo {} +} + +impl AsyncRead for StdioConnection { + fn poll_read( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + buf: &mut ReadBuf<'_>, + ) -> Poll> { + Pin::new(&mut self.io).poll_read(cx, buf) + } +} + +impl AsyncWrite for StdioConnection { + fn poll_write( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + buf: &[u8], + ) -> Poll> { + Pin::new(&mut self.io).poll_write(cx, buf) + } + + fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + Pin::new(&mut self.io).poll_flush(cx) + } + + fn poll_shutdown(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + Pin::new(&mut self.io).poll_shutdown(cx) + } +} + +impl Drop for StdioConnection { + fn drop(&mut self) { + self.closed.cancel(); + } +} diff --git a/codex-rs/code-mode-host/src/lib.rs b/codex-rs/code-mode-host/src/lib.rs index 37687e660c..7c73ae23df 100644 --- a/codex-rs/code-mode-host/src/lib.rs +++ b/codex-rs/code-mode-host/src/lib.rs @@ -48,6 +48,7 @@ pub use self::transport::DEFAULT_LISTEN_URL; mod delegate; mod grpc; +mod grpc_stdio; mod grpc_transport; mod peer; mod transport; diff --git a/codex-rs/code-mode-host/src/main.rs b/codex-rs/code-mode-host/src/main.rs index fb205fc1bf..cc797ffa47 100644 --- a/codex-rs/code-mode-host/src/main.rs +++ b/codex-rs/code-mode-host/src/main.rs @@ -18,7 +18,7 @@ const OTEL_SHUTDOWN_TIMEOUT: Duration = Duration::from_secs(/*secs*/ 5); #[derive(Debug, Parser)] struct Cli { - /// Transport endpoint: `stdio`, `stdio://`, or `grpc://IP:PORT`. + /// Transport endpoint: `stdio`, `stdio://`, `grpc+stdio://`, or `grpc://IP:PORT`. #[arg( long, value_name = "URL", diff --git a/codex-rs/code-mode-host/src/transport.rs b/codex-rs/code-mode-host/src/transport.rs index 2c56855726..095d0d78e0 100644 --- a/codex-rs/code-mode-host/src/transport.rs +++ b/codex-rs/code-mode-host/src/transport.rs @@ -11,17 +11,22 @@ pub const DEFAULT_LISTEN_URL: &str = "stdio"; #[derive(Debug, Clone, Eq, PartialEq)] enum ListenTransport { Stdio, + GrpcStdio, Grpc(SocketAddr), } pub(crate) async fn run_transport(listen_url: &str) -> Result<()> { match parse_listen_url(listen_url)? { ListenTransport::Stdio => crate::run_stdio().await, + ListenTransport::GrpcStdio => crate::grpc_stdio::run().await, ListenTransport::Grpc(bind_address) => grpc_transport::run_tcp_listener(bind_address).await, } } fn parse_listen_url(listen_url: &str) -> Result { + if listen_url == "grpc+stdio://" { + return Ok(ListenTransport::GrpcStdio); + } if matches!(listen_url, "stdio" | "stdio://") { return Ok(ListenTransport::Stdio); } @@ -36,7 +41,7 @@ fn parse_listen_url(listen_url: &str) -> Result { } anyhow::bail!( - "unsupported --listen URL `{listen_url}`; expected `grpc://IP:PORT`, `stdio`, or `stdio://`" + "unsupported --listen URL `{listen_url}`; expected `grpc://IP:PORT`, `grpc+stdio://`, `stdio`, or `stdio://`" ); } diff --git a/codex-rs/code-mode-host/src/transport_tests.rs b/codex-rs/code-mode-host/src/transport_tests.rs index d09728fda6..3f9348e69e 100644 --- a/codex-rs/code-mode-host/src/transport_tests.rs +++ b/codex-rs/code-mode-host/src/transport_tests.rs @@ -41,6 +41,10 @@ fn parse_listen_url_accepts_stdio_transports() { parse_listen_url("stdio://").expect("stdio URL should parse"), ListenTransport::Stdio ); + assert_eq!( + parse_listen_url("grpc+stdio://").expect("gRPC stdio URL should parse"), + ListenTransport::GrpcStdio + ); } #[test] diff --git a/codex-rs/core/Cargo.toml b/codex-rs/core/Cargo.toml index a8ad7d678a..39c860803a 100644 --- a/codex-rs/core/Cargo.toml +++ b/codex-rs/core/Cargo.toml @@ -92,6 +92,7 @@ dunce = { workspace = true } eventsource-stream = { workspace = true } futures = { workspace = true } http = { workspace = true } +hyper-util = { workspace = true } iana-time-zone = { workspace = true } image = { workspace = true, features = ["jpeg", "png", "webp"] } indexmap = { workspace = true } @@ -123,6 +124,8 @@ tokio-util = { workspace = true, features = ["rt"] } tokio-tungstenite = { workspace = true } toml = { workspace = true } toml_edit = { workspace = true } +tonic = { workspace = true } +tower = { version = "0.5.3", features = ["util"] } tracing = { workspace = true, features = ["log"] } url = { workspace = true } uuid = { workspace = true, features = ["serde", "v4", "v5", "v7"] } diff --git a/codex-rs/core/config.schema.json b/codex-rs/core/config.schema.json index dbf2606dd5..bc1fd50407 100644 --- a/codex-rs/core/config.schema.json +++ b/codex-rs/core/config.schema.json @@ -803,6 +803,9 @@ "code_mode_host": { "$ref": "#/definitions/FeatureToml_for_CodeModeHostConfigToml" }, + "code_mode_host_grpc": { + "type": "boolean" + }, "code_mode_interrupt": { "type": "boolean" }, @@ -7100,6 +7103,9 @@ "code_mode_host": { "$ref": "#/definitions/FeatureToml_for_CodeModeHostConfigToml" }, + "code_mode_host_grpc": { + "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 new file mode 100644 index 0000000000..de328875cb --- /dev/null +++ b/codex-rs/core/src/code_mode_host.rs @@ -0,0 +1,169 @@ +//! Owns the local gRPC code-mode client and host process. Sessions share a lazy +//! HTTP/2 channel; reconnecting spawns a fresh host, and dropping its I/O kills and +//! reaps the child. Session recovery and stale cell IDs use the shared gRPC client. + +use std::io; +use std::path::Path; +use std::path::PathBuf; +use std::pin::Pin; +use std::process::Stdio; +use std::sync::OnceLock; +use std::task::Context; +use std::task::Poll; +use std::time::Duration; + +use codex_code_mode::CodeModeSessionCellExecutionLimits; +use codex_code_mode::CodeModeSessionProvider; +use codex_code_mode::CodeModeSessionProviderFuture; +use codex_code_mode::GrpcCodeModeSessionProvider; +use codex_code_mode::ProcessOwnedCodeModeSessionProvider; +use codex_protocol::shell_environment::scrub_non_inheritable_env_vars; +use codex_utils_process::background_command; +use hyper_util::rt::TokioIo; +use tokio::io::AsyncBufReadExt; +use tokio::io::AsyncRead; +use tokio::io::AsyncWrite; +use tokio::io::BufReader; +use tokio::io::ReadBuf; +use tokio::process::Child; +use tokio::process::ChildStdin; +use tokio::process::ChildStdout; +use tokio::process::Command; +use tonic::transport::Endpoint; +use tower::service_fn; + +pub(crate) struct ProcessOwnedGrpcCodeModeSessionProvider { + host_program: PathBuf, + provider: OnceLock, +} + +impl ProcessOwnedGrpcCodeModeSessionProvider { + pub(crate) fn with_host_program(host_program: PathBuf) -> Self { + Self { + host_program, + provider: OnceLock::new(), + } + } +} + +impl CodeModeSessionProvider for ProcessOwnedGrpcCodeModeSessionProvider { + fn availability(&self) -> Result<(), String> { + ProcessOwnedCodeModeSessionProvider::with_host_program(self.host_program.clone()) + .availability() + } + + fn create_session(&self) -> CodeModeSessionProviderFuture<'_> { + self.create_session_with_limits(CodeModeSessionCellExecutionLimits::default()) + } + + fn create_session_with_limits<'a>( + &'a self, + limits: CodeModeSessionCellExecutionLimits, + ) -> CodeModeSessionProviderFuture<'a> { + Box::pin(async move { + self.availability()?; + let provider = self.provider.get_or_init(|| { + let host_program = self.host_program.clone(); + // The URI is an HTTP/2 authority only; the connector exclusively uses pipes. + // A live but unresponsive child can leave both pipes open. Let HTTP/2 + // detect that and drop the I/O, which reaps the child on reconnect. + let channel = Endpoint::from_static("http://code-mode-host") + .http2_keep_alive_interval(Duration::from_secs(/*secs*/ 30)) + .keep_alive_timeout(Duration::from_secs(/*secs*/ 10)) + .keep_alive_while_idle(/*enabled*/ true) + .connect_with_connector_lazy(service_fn(move |_| { + let host_program = host_program.clone(); + async move { HostConnection::spawn(&host_program).map(TokioIo::new) } + })); + GrpcCodeModeSessionProvider::with_channel(channel) + }); + provider.create_session_with_limits(limits).await + }) + } +} + +struct HostConnection { + io: tokio::io::Join, + child: Option, +} + +impl HostConnection { + fn spawn(host_program: &Path) -> io::Result { + let mut command = Command::from(background_command(host_program)); + #[cfg(unix)] + command.process_group(/*pgroup*/ 0); + command + .args(["--listen", "grpc+stdio://"]) + .stdin(Stdio::piped()) + .stdout(Stdio::piped()) + .stderr(Stdio::piped()) + .kill_on_drop(/*kill_on_drop*/ true); + scrub_non_inheritable_env_vars(command.as_std_mut()); + let mut child = command.spawn()?; + let stdin = child + .stdin + .take() + .ok_or_else(|| io::Error::other("code-mode host has no stdin"))?; + let stdout = child + .stdout + .take() + .ok_or_else(|| io::Error::other("code-mode host has no stdout"))?; + if let Some(stderr) = child.stderr.take() { + tokio::spawn(async move { + let mut lines = BufReader::new(stderr).lines(); + while let Ok(Some(line)) = lines.next_line().await { + tracing::debug!("code-mode host stderr: {line}"); + } + }); + } + Ok(Self { + io: tokio::io::join(stdout, stdin), + child: Some(child), + }) + } +} + +impl AsyncRead for HostConnection { + fn poll_read( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + buf: &mut ReadBuf<'_>, + ) -> Poll> { + Pin::new(&mut self.io).poll_read(cx, buf) + } +} + +impl AsyncWrite for HostConnection { + fn poll_write( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + buf: &[u8], + ) -> Poll> { + Pin::new(&mut self.io).poll_write(cx, buf) + } + + fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + Pin::new(&mut self.io).poll_flush(cx) + } + + fn poll_shutdown(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + Pin::new(&mut self.io).poll_shutdown(cx) + } +} + +impl Drop for HostConnection { + fn drop(&mut self) { + if let Some(mut child) = self.child.take() { + let _ = child.start_kill(); + if let Ok(runtime) = tokio::runtime::Handle::try_current() { + runtime.spawn(async move { + let _ = child.wait().await; + }); + } + } + } +} + +#[cfg(test)] +#[path = "code_mode_host_tests.rs"] +mod tests; diff --git a/codex-rs/core/src/code_mode_host_tests.rs b/codex-rs/core/src/code_mode_host_tests.rs new file mode 100644 index 0000000000..5f0d1892fc --- /dev/null +++ b/codex-rs/core/src/code_mode_host_tests.rs @@ -0,0 +1,148 @@ +//! Exercises the process-owned gRPC transport against the standalone host. + +use std::sync::Arc; +use std::time::Duration; + +use anyhow::Context; +use anyhow::Result; +use codex_code_mode::CodeModeSession; +use codex_code_mode::CodeModeSessionProvider; +use codex_code_mode::ExecuteRequest; +use codex_code_mode::FunctionCallOutputContentItem; +use codex_code_mode::NoopCodeModeSessionDelegate; +use codex_code_mode::RuntimeResponse; +#[cfg(unix)] +use core_test_support::code_mode_host::CodeModeHostRecorder; +use pretty_assertions::assert_eq; +use tokio::time::timeout; + +use super::HostConnection; +use super::ProcessOwnedGrpcCodeModeSessionProvider; + +const TEST_TIMEOUT: Duration = Duration::from_secs(/*secs*/ 20); + +async fn execute( + session: &Arc, + source: &str, +) -> Result> { + let response = timeout(TEST_TIMEOUT, async { + session + .execute( + ExecuteRequest { + tool_call_id: "call-1".to_string(), + enabled_tools: Vec::new(), + source: source.to_string(), + yield_time_ms: Some(/*value*/ 5_000), + max_output_tokens: Some(/*value*/ 1_000), + }, + Arc::new(NoopCodeModeSessionDelegate), + /*preempt*/ None, + ) + .await? + .initial_response() + .await + }) + .await + .context("timed out executing a stdio gRPC cell")? + .map_err(anyhow::Error::msg)?; + match response { + RuntimeResponse::Result { + content_items, + error_text: None, + .. + } => Ok(content_items), + other => anyhow::bail!("unexpected code-mode response: {other:?}"), + } +} + +#[tokio::test] +async fn grpc_stdio_sessions_share_a_host_and_isolate_state() -> Result<()> { + #[cfg(unix)] + let recorder = CodeModeHostRecorder::new()?; + #[cfg(unix)] + let host_program = recorder.program(); + #[cfg(not(unix))] + let host_program = codex_utils_cargo_bin::cargo_bin("codex-code-mode-host")?; + let provider = ProcessOwnedGrpcCodeModeSessionProvider::with_host_program(host_program); + let (first, second) = tokio::try_join!(provider.create_session(), provider.create_session()) + .map_err(anyhow::Error::msg)?; + assert_eq!( + execute(&first, "store(\"stored\", 42); text(load(\"stored\"));").await?, + vec![FunctionCallOutputContentItem::InputText { + text: "42".to_string() + }] + ); + assert_eq!( + execute(&second, "text(typeof load(\"stored\"));").await?, + vec![FunctionCallOutputContentItem::InputText { + text: "undefined".to_string() + }] + ); + first.shutdown().await.map_err(anyhow::Error::msg)?; + assert_eq!( + execute(&second, "text(\"still open\");").await?, + vec![FunctionCallOutputContentItem::InputText { + text: "still open".to_string() + }] + ); + second.shutdown().await.map_err(anyhow::Error::msg)?; + #[cfg(unix)] + assert_eq!( + recorder.invocations()?, + vec!["--listen grpc+stdio://".to_string()], + "both sessions must share one gRPC host process" + ); + Ok(()) +} + +#[tokio::test] +async fn grpc_stdio_host_exits_when_stdin_closes() -> Result<()> { + let mut connection = + HostConnection::spawn(&codex_utils_cargo_bin::cargo_bin("codex-code-mode-host")?)?; + let mut child = connection.child.take().context("host process")?; + drop(connection); + let status = timeout(TEST_TIMEOUT, child.wait()) + .await + .context("host did not exit on stdin EOF")??; + assert!(status.success()); + Ok(()) +} + +#[cfg(unix)] +#[tokio::test] +async fn hung_grpc_stdio_host_is_reaped_and_the_session_recovers() -> Result<()> { + let recorder = CodeModeHostRecorder::new()?; + let provider = ProcessOwnedGrpcCodeModeSessionProvider::with_host_program(recorder.program()); + let session = provider + .create_session() + .await + .map_err(anyhow::Error::msg)?; + let pid = recorder.pid()?; + // Suspended processes keep their pipes open: EOF cannot detect this failure. + if unsafe { libc::kill(pid as libc::pid_t, libc::SIGSTOP) } != 0 { + return Err(std::io::Error::last_os_error().into()); + } + timeout(Duration::from_secs(/*secs*/ 55), async { + while unsafe { + libc::kill(pid as libc::pid_t, /*sig*/ 0) + } == 0 + { + tokio::time::sleep(Duration::from_millis(/*millis*/ 10)).await; + } + }) + .await + .context("unresponsive host was not killed and reaped")?; + assert_eq!( + execute(&session, "text(\"recovered\");").await?, + vec![FunctionCallOutputContentItem::InputText { + text: "recovered".to_string() + }] + ); + assert_eq!( + recorder.invocations()?, + vec!["--listen grpc+stdio://".to_string(); 2], + "the session must reconnect to a fresh gRPC host" + ); + session.shutdown().await.map_err(anyhow::Error::msg)?; + Ok(()) +} diff --git a/codex-rs/core/src/lib.rs b/codex-rs/core/src/lib.rs index e59f841ef1..d4ec1918bf 100644 --- a/codex-rs/core/src/lib.rs +++ b/codex-rs/core/src/lib.rs @@ -79,6 +79,7 @@ mod agent_communication; mod agent_message_board; pub use agent_message_board::install_agent_message_board; mod attestation; +mod code_mode_host; mod codex_delegate; mod command_canonicalization; pub mod config; diff --git a/codex-rs/core/src/thread_manager.rs b/codex-rs/core/src/thread_manager.rs index 6685663eb6..b6931ad555 100644 --- a/codex-rs/core/src/thread_manager.rs +++ b/codex-rs/core/src/thread_manager.rs @@ -61,6 +61,7 @@ use codex_history::HistoryInitialization; use codex_history::InitialHistory; use codex_history::ResumedHistory; use codex_history::RolloutItem; +use codex_install_context::InstallContext; use codex_login::AuthManager; use codex_login::CodexAuth; use codex_login::default_client::CODEX_INTERNAL_ORIGINATOR_OVERRIDE_ENV_VAR; @@ -574,7 +575,10 @@ impl ThreadManager { if config.features.enabled(Feature::CodeModeHost) || config.code_mode.disable_in_process_fallback { - Arc::new(ProcessOwnedCodeModeSessionProvider::default()) + Self::process_code_mode_session_provider( + config, + InstallContext::current().code_mode_host_program(), + ) } else { Arc::new(DisabledCodeModeSessionProvider) }; @@ -660,17 +664,33 @@ impl ThreadManager { pub(crate) fn with_code_mode_host_program_for_tests( mut self, host_program: PathBuf, - _config: &Config, + config: &Config, ) -> Self { let Some(state) = Arc::get_mut(&mut self.state) else { unreachable!("new thread manager state should not be shared"); }; - state.code_mode_session_provider = Arc::new( - ProcessOwnedCodeModeSessionProvider::with_host_program(host_program), - ); + state.code_mode_session_provider = + Self::process_code_mode_session_provider(config, host_program); self } + fn process_code_mode_session_provider( + config: &Config, + host_program: PathBuf, + ) -> 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, + )) + } + } + /// Construct with a dummy AuthManager containing the provided CodexAuth. /// Used for integration tests: should not be used by ordinary business logic. pub(crate) fn with_models_provider_for_tests( diff --git a/codex-rs/core/tests/common/code_mode_host.rs b/codex-rs/core/tests/common/code_mode_host.rs new file mode 100644 index 0000000000..47aaf0e670 --- /dev/null +++ b/codex-rs/core/tests/common/code_mode_host.rs @@ -0,0 +1,59 @@ +//! Records launches of the real code-mode host through a Unix wrapper. Each log +//! line represents one process spawn and its arguments, independently of sessions. + +use std::fs; +use std::path::PathBuf; + +use anyhow::Context; +use anyhow::Result; +use tempfile::TempDir; + +/// Observes host process launches without replacing either transport implementation. +pub struct CodeModeHostRecorder { + _directory: TempDir, + program: PathBuf, + invocations: PathBuf, + pid: PathBuf, +} + +impl CodeModeHostRecorder { + pub fn new() -> Result { + let directory = tempfile::tempdir()?; + let program = directory.path().join("code mode host wrapper"); + let invocations = directory.path().join("host invocations"); + let pid = directory.path().join("host pid"); + fs::write(&invocations, "")?; + let host_program = codex_utils_cargo_bin::cargo_bin("codex-code-mode-host")?; + let quoted_host = + shlex::try_quote(host_program.to_str().context("host path is not UTF-8")?)?; + let quoted_log = shlex::try_quote(invocations.to_str().context("log path is not UTF-8")?)?; + let quoted_pid = shlex::try_quote(pid.to_str().context("PID path is not UTF-8")?)?; + codex_utils_cargo_bin::write_executable( + &program, + &format!( + "#!/bin/sh\nprintf '%s\\n' \"$*\" >> {quoted_log}\nprintf '%s\\n' \"$$\" > {quoted_pid}\nexec {quoted_host} \"$@\"\n" + ), + )?; + Ok(Self { + _directory: directory, + program, + invocations, + pid, + }) + } + + pub fn program(&self) -> PathBuf { + self.program.clone() + } + + pub fn invocations(&self) -> Result> { + Ok(fs::read_to_string(&self.invocations)? + .lines() + .map(str::to_string) + .collect()) + } + + pub fn pid(&self) -> Result { + Ok(fs::read_to_string(&self.pid)?.trim().parse()?) + } +} diff --git a/codex-rs/core/tests/common/lib.rs b/codex-rs/core/tests/common/lib.rs index 2fe2b7de16..77758d0408 100644 --- a/codex-rs/core/tests/common/lib.rs +++ b/codex-rs/core/tests/common/lib.rs @@ -25,6 +25,8 @@ use std::path::Path; use std::path::PathBuf; pub mod apps_test_server; +#[cfg(unix)] +pub mod code_mode_host; pub mod context_snapshot; pub mod exec_server; pub mod hooks; diff --git a/codex-rs/core/tests/suite/code_mode.rs b/codex-rs/core/tests/suite/code_mode.rs index f27581cd41..3bb15e2e63 100644 --- a/codex-rs/core/tests/suite/code_mode.rs +++ b/codex-rs/core/tests/suite/code_mode.rs @@ -77,6 +77,8 @@ use core_test_support::apps_test_server::DIRECT_CALENDAR_APP_ONLY_TOOL; use core_test_support::apps_test_server::recorded_apps_tool_calls; use core_test_support::apps_test_server::search_capable_apps_builder; use core_test_support::assert_regex_match; +#[cfg(unix)] +use core_test_support::code_mode_host::CodeModeHostRecorder; use core_test_support::responses; use core_test_support::responses::ResponseMock; use core_test_support::responses::ResponsesRequest; @@ -3958,23 +3960,34 @@ text(JSON.stringify({{ } #[cfg_attr(windows, ignore = "no exec_command on Windows")] +#[test_case(false; "legacy_stdio")] +#[test_case(true; "grpc_stdio")] #[tokio::test(flavor = "multi_thread", worker_threads = 2)] -async fn code_mode_only_can_call_nested_tools() -> Result<()> { +async fn code_mode_only_can_call_nested_tools(grpc_stdio: bool) -> Result<()> { skip_if_no_network!(Ok(())); + #[cfg(unix)] + let recorder = CodeModeHostRecorder::new()?; let server = responses::start_mock_server().await; + let command = match core_test_support::test_target_os() { + core_test_support::TestTargetOs::Windows => { + "[Console]::Write('code_mode_only_nested_tool_marker')" + } + core_test_support::TestTargetOs::Linux | core_test_support::TestTargetOs::MacOs => { + "printf code_mode_only_nested_tool_marker" + } + }; + let source = format!( + r#" +const output = await tools.exec_command({{ cmd: {command:?} }}); +text(output.output); +"#, + ); responses::mount_sse_once( &server, sse(vec![ ev_response_created("resp-1"), - ev_custom_tool_call( - "call-1", - "exec", - r#" -const output = await tools.exec_command({ cmd: "printf code_mode_only_nested_tool_marker" }); -text(output.output); -"#, - ), + ev_custom_tool_call("call-1", "exec", &source), ev_completed("resp-1"), ]), ) @@ -3988,10 +4001,17 @@ text(output.output); ) .await; - let mut builder = test_codex().with_config(|config| { + let mut builder = test_codex().with_config(move |config| { let _ = config.features.enable(Feature::CodeModeOnly); + if grpc_stdio { + let _ = config.features.enable(Feature::CodeModeHostGrpc); + } }); - let test = builder.build(&server).await?; + #[cfg(unix)] + { + builder = builder.with_code_mode_host_program(recorder.program()); + } + let test = builder.build_with_auto_env(&server).await?; test.submit_turn("use exec to run nested tool in code mode only") .await?; @@ -4003,6 +4023,16 @@ text(output.output); "code_mode_only nested tool call failed unexpectedly: {output}" ); assert_eq!(output, "code_mode_only_nested_tool_marker"); + #[cfg(unix)] + assert_eq!( + recorder.invocations()?, + vec![if grpc_stdio { + "--listen grpc+stdio://".to_string() + } else { + String::new() + }], + "the feature flag must select the expected host transport" + ); assert!( request .custom_tool_call_output("call-1") diff --git a/codex-rs/features/src/lib.rs b/codex-rs/features/src/lib.rs index 08451de27e..0b8386febd 100644 --- a/codex-rs/features/src/lib.rs +++ b/codex-rs/features/src/lib.rs @@ -133,6 +133,8 @@ pub enum Feature { CodeModeBufferedExec, /// Run JavaScript code mode in the standalone host process. CodeModeHost, + /// Use gRPC over stdio for the process-owned code-mode host. + CodeModeHostGrpc, /// Establish the code-mode host connection during session startup. CodeModePrewarm, /// Terminate active code mode cells when their turn is interrupted. @@ -1162,6 +1164,12 @@ pub const FEATURES: &[FeatureSpec] = &[ stage: Stage::Stable, default_enabled: true, }, + FeatureSpec { + id: Feature::CodeModeHostGrpc, + key: "code_mode_host_grpc", + stage: Stage::UnderDevelopment, + default_enabled: false, + }, FeatureSpec { id: Feature::CodeModePrewarm, key: "code_mode_prewarm",