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",