From 2490fb2b469ec4e01e8790f5207550c08315ac28 Mon Sep 17 00:00:00 2001 From: starr-openai Date: Fri, 1 May 2026 12:05:39 -0700 Subject: [PATCH 01/42] Add stdio exec-server client transport Allow exec-server clients to connect through a shell command over stdio. The connection can now retain a drop resource so the spawned child is terminated when the JSON-RPC client is dropped. Co-authored-by: Codex --- codex-rs/exec-server/src/client.rs | 75 ++++---- codex-rs/exec-server/src/client_api.rs | 16 ++ codex-rs/exec-server/src/client_transport.rs | 176 +++++++++++++++++++ codex-rs/exec-server/src/connection.rs | 12 ++ codex-rs/exec-server/src/environment.rs | 5 +- codex-rs/exec-server/src/lib.rs | 3 + codex-rs/exec-server/src/rpc.rs | 6 +- codex-rs/exec-server/src/server/processor.rs | 2 +- 8 files changed, 252 insertions(+), 43 deletions(-) create mode 100644 codex-rs/exec-server/src/client_transport.rs diff --git a/codex-rs/exec-server/src/client.rs b/codex-rs/exec-server/src/client.rs index 47359393d368..0729be151d73 100644 --- a/codex-rs/exec-server/src/client.rs +++ b/codex-rs/exec-server/src/client.rs @@ -17,13 +17,14 @@ use tokio::sync::mpsc; use tokio::sync::watch; use tokio::time::timeout; -use tokio_tungstenite::connect_async; use tracing::debug; use crate::ProcessId; use crate::client_api::ExecServerClientConnectOptions; +use crate::client_api::ExecServerTransport; use crate::client_api::HttpClient; use crate::client_api::RemoteExecServerConnectArgs; +use crate::client_api::StdioExecServerConnectArgs; use crate::connection::JsonRpcConnection; use crate::process::ExecProcessEvent; use crate::process::ExecProcessEventLog; @@ -105,6 +106,16 @@ impl From for ExecServerClientConnectOptions { } } +impl From for ExecServerClientConnectOptions { + fn from(value: StdioExecServerConnectArgs) -> Self { + Self { + client_name: value.client_name, + initialize_timeout: value.initialize_timeout, + resume_session_id: value.resume_session_id, + } + } +} + impl RemoteExecServerConnectArgs { pub fn new(websocket_url: String, client_name: String) -> Self { Self { @@ -180,29 +191,23 @@ pub struct ExecServerClient { #[derive(Clone)] pub(crate) struct LazyRemoteExecServerClient { - websocket_url: String, + transport: ExecServerTransport, client: Arc>, } impl LazyRemoteExecServerClient { - pub(crate) fn new(websocket_url: String) -> Self { + pub(crate) fn new(transport: ExecServerTransport) -> Self { Self { - websocket_url, + transport, client: Arc::new(OnceCell::new()), } } pub(crate) async fn get(&self) -> Result { self.client - .get_or_try_init(|| async { - ExecServerClient::connect_websocket(RemoteExecServerConnectArgs { - websocket_url: self.websocket_url.clone(), - client_name: "codex-environment".to_string(), - connect_timeout: Duration::from_secs(5), - initialize_timeout: Duration::from_secs(5), - resume_session_id: None, - }) - .await + .get_or_try_init(|| { + let transport = self.transport.clone(); + async move { transport.connect_for_environment().await } }) .await .cloned() @@ -269,32 +274,6 @@ pub enum ExecServerError { } impl ExecServerClient { - pub async fn connect_websocket( - args: RemoteExecServerConnectArgs, - ) -> Result { - let websocket_url = args.websocket_url.clone(); - let connect_timeout = args.connect_timeout; - let (stream, _) = timeout(connect_timeout, connect_async(websocket_url.as_str())) - .await - .map_err(|_| ExecServerError::WebSocketConnectTimeout { - url: websocket_url.clone(), - timeout: connect_timeout, - })? - .map_err(|source| ExecServerError::WebSocketConnect { - url: websocket_url.clone(), - source, - })?; - - Self::connect( - JsonRpcConnection::from_websocket( - stream, - format!("exec-server websocket {websocket_url}"), - ), - args.into(), - ) - .await - } - pub async fn initialize( &self, options: ExecServerClientConnectOptions, @@ -443,7 +422,7 @@ impl ExecServerClient { .clone() } - async fn connect( + pub(crate) async fn connect( connection: JsonRpcConnection, options: ExecServerClientConnectOptions, ) -> Result { @@ -905,6 +884,7 @@ mod tests { use super::ExecServerClient; use super::ExecServerClientConnectOptions; use crate::ProcessId; + use crate::client_api::StdioExecServerConnectArgs; use crate::connection::JsonRpcConnection; use crate::process::ExecProcessEvent; use crate::protocol::EXEC_CLOSED_METHOD; @@ -942,6 +922,21 @@ mod tests { .expect("json-rpc line should write"); } + #[cfg(not(windows))] + #[tokio::test] + async fn connect_stdio_command_initializes_json_rpc_client() { + let client = ExecServerClient::connect_stdio_command(StdioExecServerConnectArgs { + shell_command: "read _line; printf '%s\\n' '{\"id\":1,\"result\":{\"sessionId\":\"stdio-test\"}}'; read _line; sleep 60".to_string(), + client_name: "stdio-test-client".to_string(), + initialize_timeout: Duration::from_secs(1), + resume_session_id: None, + }) + .await + .expect("stdio client should connect"); + + assert_eq!(client.session_id().as_deref(), Some("stdio-test")); + } + #[tokio::test] async fn process_events_are_delivered_in_seq_order_when_notifications_are_reordered() { let (client_stdin, server_reader) = duplex(1 << 20); diff --git a/codex-rs/exec-server/src/client_api.rs b/codex-rs/exec-server/src/client_api.rs index b1761b69f11b..95ed053476d9 100644 --- a/codex-rs/exec-server/src/client_api.rs +++ b/codex-rs/exec-server/src/client_api.rs @@ -25,6 +25,22 @@ pub struct RemoteExecServerConnectArgs { pub resume_session_id: Option, } +/// Stdio connection arguments for a command-backed exec-server. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct StdioExecServerConnectArgs { + pub shell_command: String, + pub client_name: String, + pub initialize_timeout: Duration, + pub resume_session_id: Option, +} + +/// Transport used to connect to a remote exec-server environment. +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum ExecServerTransport { + WebSocketUrl(String), + StdioShellCommand(String), +} + /// Sends HTTP requests through a runtime-selected transport. /// /// This is the HTTP capability counterpart to [`crate::ExecBackend`]. Callers diff --git a/codex-rs/exec-server/src/client_transport.rs b/codex-rs/exec-server/src/client_transport.rs new file mode 100644 index 000000000000..908a9e1b0596 --- /dev/null +++ b/codex-rs/exec-server/src/client_transport.rs @@ -0,0 +1,176 @@ +use std::process::Stdio; +use std::time::Duration; + +use tokio::io::AsyncBufReadExt; +use tokio::io::BufReader; +use tokio::process::Child; +use tokio::process::Command; +use tokio::runtime::Handle; +use tokio::time::timeout; +use tokio_tungstenite::connect_async; +use tracing::debug; +use tracing::warn; + +use crate::ExecServerClient; +use crate::ExecServerError; +use crate::client_api::ExecServerTransport; +use crate::client_api::RemoteExecServerConnectArgs; +use crate::client_api::StdioExecServerConnectArgs; +use crate::connection::JsonRpcConnection; + +const ENVIRONMENT_CLIENT_NAME: &str = "codex-environment"; +const ENVIRONMENT_CONNECT_TIMEOUT: Duration = Duration::from_secs(5); +const ENVIRONMENT_INITIALIZE_TIMEOUT: Duration = Duration::from_secs(5); + +impl ExecServerTransport { + pub(crate) async fn connect_for_environment(self) -> Result { + match self { + ExecServerTransport::WebSocketUrl(websocket_url) => { + ExecServerClient::connect_websocket(RemoteExecServerConnectArgs { + websocket_url, + client_name: ENVIRONMENT_CLIENT_NAME.to_string(), + connect_timeout: ENVIRONMENT_CONNECT_TIMEOUT, + initialize_timeout: ENVIRONMENT_INITIALIZE_TIMEOUT, + resume_session_id: None, + }) + .await + } + ExecServerTransport::StdioShellCommand(shell_command) => { + ExecServerClient::connect_stdio_command(StdioExecServerConnectArgs { + shell_command, + client_name: ENVIRONMENT_CLIENT_NAME.to_string(), + initialize_timeout: ENVIRONMENT_INITIALIZE_TIMEOUT, + resume_session_id: None, + }) + .await + } + } + } +} + +impl ExecServerClient { + pub async fn connect_websocket( + args: RemoteExecServerConnectArgs, + ) -> Result { + let websocket_url = args.websocket_url.clone(); + let connect_timeout = args.connect_timeout; + let (stream, _) = timeout(connect_timeout, connect_async(websocket_url.as_str())) + .await + .map_err(|_| ExecServerError::WebSocketConnectTimeout { + url: websocket_url.clone(), + timeout: connect_timeout, + })? + .map_err(|source| ExecServerError::WebSocketConnect { + url: websocket_url.clone(), + source, + })?; + + Self::connect( + JsonRpcConnection::from_websocket( + stream, + format!("exec-server websocket {websocket_url}"), + ), + args.into(), + ) + .await + } + + pub async fn connect_stdio_command( + args: StdioExecServerConnectArgs, + ) -> Result { + let shell_command = args.shell_command.clone(); + let mut child = shell_command_process(&shell_command) + .stdin(Stdio::piped()) + .stdout(Stdio::piped()) + .stderr(Stdio::piped()) + .spawn() + .map_err(ExecServerError::Spawn)?; + + let stdin = child.stdin.take().ok_or_else(|| { + ExecServerError::Protocol("spawned exec-server command has no stdin".to_string()) + })?; + let stdout = child.stdout.take().ok_or_else(|| { + ExecServerError::Protocol("spawned exec-server command has no stdout".to_string()) + })?; + if let Some(stderr) = child.stderr.take() { + tokio::spawn(async move { + let mut lines = BufReader::new(stderr).lines(); + loop { + match lines.next_line().await { + Ok(Some(line)) => debug!("exec-server stdio stderr: {line}"), + Ok(None) => break, + Err(err) => { + warn!("failed to read exec-server stdio stderr: {err}"); + break; + } + } + } + }); + } + + Self::connect( + JsonRpcConnection::from_stdio( + stdout, + stdin, + format!("exec-server stdio command `{shell_command}`"), + ) + .with_lifetime_guard(Box::new(StdioChildGuard { child: Some(child) })), + args.into(), + ) + .await + } +} + +struct StdioChildGuard { + child: Option, +} + +impl Drop for StdioChildGuard { + fn drop(&mut self) { + let Some(child) = self.child.take() else { + return; + }; + + match Handle::try_current() { + Ok(handle) => { + let _terminate_task = handle.spawn(terminate_stdio_child(child)); + } + Err(_) => { + terminate_stdio_child_now(child); + } + } + } +} + +async fn terminate_stdio_child(mut child: Child) { + kill_stdio_child(&mut child); + if let Err(err) = child.wait().await { + debug!("failed to wait for exec-server stdio child: {err}"); + } +} + +fn terminate_stdio_child_now(mut child: Child) { + kill_stdio_child(&mut child); +} + +fn kill_stdio_child(child: &mut Child) { + if let Err(err) = child.start_kill() { + debug!("failed to terminate exec-server stdio child: {err}"); + } +} + +fn shell_command_process(shell_command: &str) -> Command { + #[cfg(windows)] + { + let mut command = Command::new("cmd"); + command.arg("/C").arg(shell_command); + command + } + + #[cfg(not(windows))] + { + let mut command = Command::new("sh"); + command.arg("-lc").arg(shell_command); + command + } +} diff --git a/codex-rs/exec-server/src/connection.rs b/codex-rs/exec-server/src/connection.rs index 71f4f31059fc..f79ac3da57d4 100644 --- a/codex-rs/exec-server/src/connection.rs +++ b/codex-rs/exec-server/src/connection.rs @@ -15,6 +15,8 @@ use tokio::io::BufWriter; pub(crate) const CHANNEL_CAPACITY: usize = 128; +pub(crate) type JsonRpcConnectionLifetimeGuard = Box; + #[derive(Debug)] pub(crate) enum JsonRpcConnectionEvent { Message(JSONRPCMessage), @@ -27,6 +29,7 @@ pub(crate) struct JsonRpcConnection { incoming_rx: mpsc::Receiver, disconnected_rx: watch::Receiver, task_handles: Vec>, + lifetime_guard: Option, } impl JsonRpcConnection { @@ -117,6 +120,7 @@ impl JsonRpcConnection { incoming_rx, disconnected_rx, task_handles: vec![reader_task, writer_task], + lifetime_guard: None, } } @@ -251,9 +255,15 @@ impl JsonRpcConnection { incoming_rx, disconnected_rx, task_handles: vec![reader_task, writer_task], + lifetime_guard: None, } } + pub(crate) fn with_lifetime_guard(mut self, guard: JsonRpcConnectionLifetimeGuard) -> Self { + self.lifetime_guard = Some(guard); + self + } + pub(crate) fn into_parts( self, ) -> ( @@ -261,12 +271,14 @@ impl JsonRpcConnection { mpsc::Receiver, watch::Receiver, Vec>, + Option, ) { ( self.outgoing_tx, self.incoming_rx, self.disconnected_rx, self.task_handles, + self.lifetime_guard, ) } } diff --git a/codex-rs/exec-server/src/environment.rs b/codex-rs/exec-server/src/environment.rs index 855989dafbc2..3764b29fe567 100644 --- a/codex-rs/exec-server/src/environment.rs +++ b/codex-rs/exec-server/src/environment.rs @@ -7,6 +7,7 @@ use crate::ExecutorFileSystem; use crate::HttpClient; use crate::client::LazyRemoteExecServerClient; use crate::client::http_client::ReqwestHttpClient; +use crate::client_api::ExecServerTransport; use crate::environment_provider::DefaultEnvironmentProvider; use crate::environment_provider::EnvironmentProvider; use crate::environment_provider::normalize_exec_server_url; @@ -274,7 +275,9 @@ impl Environment { exec_server_url: String, local_runtime_paths: Option, ) -> Self { - let client = LazyRemoteExecServerClient::new(exec_server_url.clone()); + let client = LazyRemoteExecServerClient::new(ExecServerTransport::WebSocketUrl( + exec_server_url.clone(), + )); let exec_backend: Arc = Arc::new(RemoteProcess::new(client.clone())); let filesystem: Arc = Arc::new(RemoteFileSystem::new(client.clone())); diff --git a/codex-rs/exec-server/src/lib.rs b/codex-rs/exec-server/src/lib.rs index d860d59aba98..9bec4ed1da6f 100644 --- a/codex-rs/exec-server/src/lib.rs +++ b/codex-rs/exec-server/src/lib.rs @@ -1,5 +1,6 @@ mod client; mod client_api; +mod client_transport; mod connection; mod environment; mod environment_provider; @@ -24,8 +25,10 @@ pub use client::ExecServerError; pub use client::http_client::HttpResponseBodyStream; pub use client::http_client::ReqwestHttpClient; pub use client_api::ExecServerClientConnectOptions; +pub use client_api::ExecServerTransport; pub use client_api::HttpClient; pub use client_api::RemoteExecServerConnectArgs; +pub use client_api::StdioExecServerConnectArgs; pub use codex_file_system::CopyOptions; pub use codex_file_system::CreateDirectoryOptions; pub use codex_file_system::ExecutorFileSystem; diff --git a/codex-rs/exec-server/src/rpc.rs b/codex-rs/exec-server/src/rpc.rs index 723b99f5028d..2cce5c04008a 100644 --- a/codex-rs/exec-server/src/rpc.rs +++ b/codex-rs/exec-server/src/rpc.rs @@ -23,6 +23,7 @@ use tokio::task::JoinHandle; use crate::connection::JsonRpcConnection; use crate::connection::JsonRpcConnectionEvent; +use crate::connection::JsonRpcConnectionLifetimeGuard; #[derive(Debug)] pub(crate) enum RpcCallError { @@ -229,12 +230,14 @@ pub(crate) struct RpcClient { disconnected_rx: watch::Receiver, next_request_id: AtomicI64, transport_tasks: Vec>, + _transport_lifetime_guard: Option, reader_task: JoinHandle<()>, } impl RpcClient { pub(crate) fn new(connection: JsonRpcConnection) -> (Self, mpsc::Receiver) { - let (write_tx, mut incoming_rx, disconnected_rx, transport_tasks) = connection.into_parts(); + let (write_tx, mut incoming_rx, disconnected_rx, transport_tasks, lifetime_guard) = + connection.into_parts(); let pending = Arc::new(Mutex::new(HashMap::::new())); let (event_tx, event_rx) = mpsc::channel(128); @@ -275,6 +278,7 @@ impl RpcClient { disconnected_rx, next_request_id: AtomicI64::new(1), transport_tasks, + _transport_lifetime_guard: lifetime_guard, reader_task, }, event_rx, diff --git a/codex-rs/exec-server/src/server/processor.rs b/codex-rs/exec-server/src/server/processor.rs index dc1a9b9ffe74..50907f7e4285 100644 --- a/codex-rs/exec-server/src/server/processor.rs +++ b/codex-rs/exec-server/src/server/processor.rs @@ -47,7 +47,7 @@ async fn run_connection( runtime_paths: ExecServerRuntimePaths, ) { let router = Arc::new(build_router()); - let (json_outgoing_tx, mut incoming_rx, mut disconnected_rx, connection_tasks) = + let (json_outgoing_tx, mut incoming_rx, mut disconnected_rx, connection_tasks, _lifetime_guard) = connection.into_parts(); let (outgoing_tx, mut outgoing_rx) = mpsc::channel::(CHANNEL_CAPACITY); From 8626d27c860fb369110058d295a67b653e9e1876 Mon Sep 17 00:00:00 2001 From: starr-openai Date: Fri, 1 May 2026 12:53:17 -0700 Subject: [PATCH 02/42] Make exec-server RPC client Send-safe Co-authored-by: Codex --- codex-rs/exec-server/src/connection.rs | 17 ++++++++--------- codex-rs/exec-server/src/rpc.rs | 5 +++-- 2 files changed, 11 insertions(+), 11 deletions(-) diff --git a/codex-rs/exec-server/src/connection.rs b/codex-rs/exec-server/src/connection.rs index f79ac3da57d4..d76cbcae86de 100644 --- a/codex-rs/exec-server/src/connection.rs +++ b/codex-rs/exec-server/src/connection.rs @@ -16,6 +16,13 @@ use tokio::io::BufWriter; pub(crate) const CHANNEL_CAPACITY: usize = 128; pub(crate) type JsonRpcConnectionLifetimeGuard = Box; +pub(crate) type JsonRpcConnectionParts = ( + mpsc::Sender, + mpsc::Receiver, + watch::Receiver, + Vec>, + Option, +); #[derive(Debug)] pub(crate) enum JsonRpcConnectionEvent { @@ -264,15 +271,7 @@ impl JsonRpcConnection { self } - pub(crate) fn into_parts( - self, - ) -> ( - mpsc::Sender, - mpsc::Receiver, - watch::Receiver, - Vec>, - Option, - ) { + pub(crate) fn into_parts(self) -> JsonRpcConnectionParts { ( self.outgoing_tx, self.incoming_rx, diff --git a/codex-rs/exec-server/src/rpc.rs b/codex-rs/exec-server/src/rpc.rs index 2cce5c04008a..8985849ec98f 100644 --- a/codex-rs/exec-server/src/rpc.rs +++ b/codex-rs/exec-server/src/rpc.rs @@ -2,6 +2,7 @@ use std::collections::HashMap; use std::future::Future; use std::pin::Pin; use std::sync::Arc; +use std::sync::Mutex as StdMutex; use std::sync::atomic::AtomicI64; use std::sync::atomic::Ordering; @@ -230,7 +231,7 @@ pub(crate) struct RpcClient { disconnected_rx: watch::Receiver, next_request_id: AtomicI64, transport_tasks: Vec>, - _transport_lifetime_guard: Option, + _transport_lifetime_guard: Option>, reader_task: JoinHandle<()>, } @@ -278,7 +279,7 @@ impl RpcClient { disconnected_rx, next_request_id: AtomicI64::new(1), transport_tasks, - _transport_lifetime_guard: lifetime_guard, + _transport_lifetime_guard: lifetime_guard.map(StdMutex::new), reader_task, }, event_rx, From 3d7522777cb6c023c54293a47f03d267b8601882 Mon Sep 17 00:00:00 2001 From: starr-openai Date: Fri, 1 May 2026 13:37:01 -0700 Subject: [PATCH 03/42] Remove duplicate stdio client test import Co-authored-by: Codex --- codex-rs/exec-server/src/client.rs | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/codex-rs/exec-server/src/client.rs b/codex-rs/exec-server/src/client.rs index 0729be151d73..c91902f3f9cd 100644 --- a/codex-rs/exec-server/src/client.rs +++ b/codex-rs/exec-server/src/client.rs @@ -884,7 +884,8 @@ mod tests { use super::ExecServerClient; use super::ExecServerClientConnectOptions; use crate::ProcessId; - use crate::client_api::StdioExecServerConnectArgs; + #[cfg(not(windows))] + use crate::StdioExecServerConnectArgs; use crate::connection::JsonRpcConnection; use crate::process::ExecProcessEvent; use crate::protocol::EXEC_CLOSED_METHOD; From 7834efe6529d95e7da1ea334e07a42bc6cec6b4a Mon Sep 17 00:00:00 2001 From: starr-openai Date: Mon, 4 May 2026 11:10:40 -0700 Subject: [PATCH 04/42] Clarify exec-server transport lifetime ownership Co-authored-by: Codex --- codex-rs/exec-server/src/client_transport.rs | 2 +- codex-rs/exec-server/src/connection.rs | 16 ++++++++-------- codex-rs/exec-server/src/rpc.rs | 16 ++++++++++++---- codex-rs/exec-server/src/server/processor.rs | 9 +++++++-- 4 files changed, 28 insertions(+), 15 deletions(-) diff --git a/codex-rs/exec-server/src/client_transport.rs b/codex-rs/exec-server/src/client_transport.rs index 908a9e1b0596..c49e9cdb9818 100644 --- a/codex-rs/exec-server/src/client_transport.rs +++ b/codex-rs/exec-server/src/client_transport.rs @@ -114,7 +114,7 @@ impl ExecServerClient { stdin, format!("exec-server stdio command `{shell_command}`"), ) - .with_lifetime_guard(Box::new(StdioChildGuard { child: Some(child) })), + .with_transport_lifetime(Box::new(StdioChildGuard { child: Some(child) })), args.into(), ) .await diff --git a/codex-rs/exec-server/src/connection.rs b/codex-rs/exec-server/src/connection.rs index d76cbcae86de..0ec4c6bc5f1b 100644 --- a/codex-rs/exec-server/src/connection.rs +++ b/codex-rs/exec-server/src/connection.rs @@ -15,13 +15,13 @@ use tokio::io::BufWriter; pub(crate) const CHANNEL_CAPACITY: usize = 128; -pub(crate) type JsonRpcConnectionLifetimeGuard = Box; +pub(crate) type JsonRpcTransportLifetime = Box; pub(crate) type JsonRpcConnectionParts = ( mpsc::Sender, mpsc::Receiver, watch::Receiver, Vec>, - Option, + Option, ); #[derive(Debug)] @@ -36,7 +36,7 @@ pub(crate) struct JsonRpcConnection { incoming_rx: mpsc::Receiver, disconnected_rx: watch::Receiver, task_handles: Vec>, - lifetime_guard: Option, + transport_lifetime: Option, } impl JsonRpcConnection { @@ -127,7 +127,7 @@ impl JsonRpcConnection { incoming_rx, disconnected_rx, task_handles: vec![reader_task, writer_task], - lifetime_guard: None, + transport_lifetime: None, } } @@ -262,12 +262,12 @@ impl JsonRpcConnection { incoming_rx, disconnected_rx, task_handles: vec![reader_task, writer_task], - lifetime_guard: None, + transport_lifetime: None, } } - pub(crate) fn with_lifetime_guard(mut self, guard: JsonRpcConnectionLifetimeGuard) -> Self { - self.lifetime_guard = Some(guard); + pub(crate) fn with_transport_lifetime(mut self, lifetime: JsonRpcTransportLifetime) -> Self { + self.transport_lifetime = Some(lifetime); self } @@ -277,7 +277,7 @@ impl JsonRpcConnection { self.incoming_rx, self.disconnected_rx, self.task_handles, - self.lifetime_guard, + self.transport_lifetime, ) } } diff --git a/codex-rs/exec-server/src/rpc.rs b/codex-rs/exec-server/src/rpc.rs index 8985849ec98f..d9d8fbbf726d 100644 --- a/codex-rs/exec-server/src/rpc.rs +++ b/codex-rs/exec-server/src/rpc.rs @@ -24,7 +24,7 @@ use tokio::task::JoinHandle; use crate::connection::JsonRpcConnection; use crate::connection::JsonRpcConnectionEvent; -use crate::connection::JsonRpcConnectionLifetimeGuard; +use crate::connection::JsonRpcTransportLifetime; #[derive(Debug)] pub(crate) enum RpcCallError { @@ -231,13 +231,19 @@ pub(crate) struct RpcClient { disconnected_rx: watch::Receiver, next_request_id: AtomicI64, transport_tasks: Vec>, - _transport_lifetime_guard: Option>, + _transport_lifetime: Option, reader_task: JoinHandle<()>, } +// Holds transport-owned resources, such as a stdio child process, for as long +// as the RPC client owns the underlying connection. +struct TransportLifetime { + _guard: StdMutex, +} + impl RpcClient { pub(crate) fn new(connection: JsonRpcConnection) -> (Self, mpsc::Receiver) { - let (write_tx, mut incoming_rx, disconnected_rx, transport_tasks, lifetime_guard) = + let (write_tx, mut incoming_rx, disconnected_rx, transport_tasks, transport_lifetime) = connection.into_parts(); let pending = Arc::new(Mutex::new(HashMap::::new())); let (event_tx, event_rx) = mpsc::channel(128); @@ -279,7 +285,9 @@ impl RpcClient { disconnected_rx, next_request_id: AtomicI64::new(1), transport_tasks, - _transport_lifetime_guard: lifetime_guard.map(StdMutex::new), + _transport_lifetime: transport_lifetime.map(|lifetime| TransportLifetime { + _guard: StdMutex::new(lifetime), + }), reader_task, }, event_rx, diff --git a/codex-rs/exec-server/src/server/processor.rs b/codex-rs/exec-server/src/server/processor.rs index 50907f7e4285..b7e1a03bd529 100644 --- a/codex-rs/exec-server/src/server/processor.rs +++ b/codex-rs/exec-server/src/server/processor.rs @@ -47,8 +47,13 @@ async fn run_connection( runtime_paths: ExecServerRuntimePaths, ) { let router = Arc::new(build_router()); - let (json_outgoing_tx, mut incoming_rx, mut disconnected_rx, connection_tasks, _lifetime_guard) = - connection.into_parts(); + let ( + json_outgoing_tx, + mut incoming_rx, + mut disconnected_rx, + connection_tasks, + _transport_lifetime, + ) = connection.into_parts(); let (outgoing_tx, mut outgoing_rx) = mpsc::channel::(CHANNEL_CAPACITY); let notifications = RpcNotificationSender::new(outgoing_tx.clone()); From 881e7b5ddf693a2591349a98a2a32b33930d8a97 Mon Sep 17 00:00:00 2001 From: starr-openai Date: Mon, 4 May 2026 12:03:46 -0700 Subject: [PATCH 05/42] Clean up stdio client process groups Use the existing process-group cleanup pattern for stdio command transports so wrapper shell children are terminated with the client lifetime. Add a regression test that drops the client after spawning a background shell child through the command-backed transport. Co-authored-by: Codex --- codex-rs/exec-server/src/client.rs | 76 +++++++++++++++++++ codex-rs/exec-server/src/client_transport.rs | 77 ++++++++++++++++---- 2 files changed, 140 insertions(+), 13 deletions(-) diff --git a/codex-rs/exec-server/src/client.rs b/codex-rs/exec-server/src/client.rs index c91902f3f9cd..c04c57eb25d4 100644 --- a/codex-rs/exec-server/src/client.rs +++ b/codex-rs/exec-server/src/client.rs @@ -872,6 +872,10 @@ mod tests { use codex_app_server_protocol::JSONRPCNotification; use codex_app_server_protocol::JSONRPCResponse; use pretty_assertions::assert_eq; + #[cfg(unix)] + use std::path::Path; + #[cfg(unix)] + use std::process::Command; use tokio::io::AsyncBufReadExt; use tokio::io::AsyncWrite; use tokio::io::AsyncWriteExt; @@ -879,6 +883,8 @@ mod tests { use tokio::io::duplex; use tokio::sync::mpsc; use tokio::time::Duration; + #[cfg(unix)] + use tokio::time::sleep; use tokio::time::timeout; use super::ExecServerClient; @@ -938,6 +944,76 @@ mod tests { assert_eq!(client.session_id().as_deref(), Some("stdio-test")); } + #[cfg(unix)] + #[tokio::test] + async fn dropping_stdio_client_terminates_shell_process_group() { + let tempdir = tempfile::tempdir().expect("tempdir should be created"); + let pid_file = tempdir.path().join("child.pid"); + let shell_command = format!( + "read _line; \ + (trap 'exit 0' TERM; while true; do sleep 1; done) & \ + child=$!; \ + echo \"$child\" > {}; \ + printf '%s\\n' '{{\"id\":1,\"result\":{{\"sessionId\":\"stdio-test\"}}}}'; \ + read _line; \ + wait \"$child\"", + shell_quote(pid_file.as_path()), + ); + + let client = ExecServerClient::connect_stdio_command(StdioExecServerConnectArgs { + shell_command, + client_name: "stdio-test-client".to_string(), + initialize_timeout: Duration::from_secs(1), + resume_session_id: None, + }) + .await + .expect("stdio client should connect"); + let child_pid = read_pid_file(pid_file.as_path()).await; + assert!( + process_exists(child_pid), + "wrapper child process should be running before client drop" + ); + + drop(client); + + for _ in 0..20 { + if !process_exists(child_pid) { + return; + } + sleep(Duration::from_millis(100)).await; + } + panic!("wrapper child process {child_pid} should exit after client drop"); + } + + #[cfg(unix)] + async fn read_pid_file(path: &Path) -> u32 { + for _ in 0..20 { + if let Ok(contents) = std::fs::read_to_string(path) { + return contents + .trim() + .parse() + .expect("pid file should contain a pid"); + } + sleep(Duration::from_millis(50)).await; + } + panic!("pid file {} should be written", path.display()); + } + + #[cfg(unix)] + fn process_exists(pid: u32) -> bool { + Command::new("kill") + .arg("-0") + .arg(pid.to_string()) + .status() + .is_ok_and(|status| status.success()) + } + + #[cfg(unix)] + fn shell_quote(path: &Path) -> String { + let value = path.to_string_lossy(); + format!("'{}'", value.replace('\'', "'\\''")) + } + #[tokio::test] async fn process_events_are_delivered_in_seq_order_when_notifications_are_reordered() { let (client_stdin, server_reader) = duplex(1 << 20); diff --git a/codex-rs/exec-server/src/client_transport.rs b/codex-rs/exec-server/src/client_transport.rs index c49e9cdb9818..d877aae3dfb4 100644 --- a/codex-rs/exec-server/src/client_transport.rs +++ b/codex-rs/exec-server/src/client_transport.rs @@ -1,6 +1,14 @@ use std::process::Stdio; +#[cfg(unix)] +use std::thread::sleep; +#[cfg(unix)] +use std::thread::spawn; use std::time::Duration; +#[cfg(unix)] +use codex_utils_pty::process_group::kill_process_group; +#[cfg(unix)] +use codex_utils_pty::process_group::terminate_process_group; use tokio::io::AsyncBufReadExt; use tokio::io::BufReader; use tokio::process::Child; @@ -21,6 +29,8 @@ use crate::connection::JsonRpcConnection; const ENVIRONMENT_CLIENT_NAME: &str = "codex-environment"; const ENVIRONMENT_CONNECT_TIMEOUT: Duration = Duration::from_secs(5); const ENVIRONMENT_INITIALIZE_TIMEOUT: Duration = Duration::from_secs(5); +#[cfg(unix)] +const STDIO_CHILD_TERM_GRACE_PERIOD: Duration = Duration::from_millis(500); impl ExecServerTransport { pub(crate) async fn connect_for_environment(self) -> Result { @@ -80,11 +90,13 @@ impl ExecServerClient { ) -> Result { let shell_command = args.shell_command.clone(); let mut child = shell_command_process(&shell_command) + .kill_on_drop(true) .stdin(Stdio::piped()) .stdout(Stdio::piped()) .stderr(Stdio::piped()) .spawn() .map_err(ExecServerError::Spawn)?; + let process_id = child.id(); let stdin = child.stdin.take().ok_or_else(|| { ExecServerError::Protocol("spawned exec-server command has no stdin".to_string()) @@ -114,7 +126,10 @@ impl ExecServerClient { stdin, format!("exec-server stdio command `{shell_command}`"), ) - .with_transport_lifetime(Box::new(StdioChildGuard { child: Some(child) })), + .with_transport_lifetime(Box::new(StdioChildGuard { + child: Some(child), + process_id, + })), args.into(), ) .await @@ -123,34 +138,69 @@ impl ExecServerClient { struct StdioChildGuard { child: Option, + process_id: Option, } impl Drop for StdioChildGuard { fn drop(&mut self) { - let Some(child) = self.child.take() else { + let Some(mut child) = self.child.take() else { return; }; - match Handle::try_current() { - Ok(handle) => { - let _terminate_task = handle.spawn(terminate_stdio_child(child)); - } - Err(_) => { - terminate_stdio_child_now(child); - } + terminate_stdio_child_process(self.process_id, &mut child); + + if let Ok(handle) = Handle::try_current() { + let _wait_task = handle.spawn(wait_stdio_child(child)); } } } -async fn terminate_stdio_child(mut child: Child) { - kill_stdio_child(&mut child); +async fn wait_stdio_child(mut child: Child) { if let Err(err) = child.wait().await { debug!("failed to wait for exec-server stdio child: {err}"); } } -fn terminate_stdio_child_now(mut child: Child) { - kill_stdio_child(&mut child); +#[cfg(unix)] +fn terminate_stdio_child_process(process_group_id: Option, child: &mut Child) { + let Some(process_group_id) = process_group_id else { + kill_stdio_child(child); + return; + }; + + let should_escalate = match terminate_process_group(process_group_id) { + Ok(exists) => exists, + Err(err) => { + debug!("failed to terminate exec-server stdio process group {process_group_id}: {err}"); + false + } + }; + if should_escalate { + spawn(move || { + sleep(STDIO_CHILD_TERM_GRACE_PERIOD); + if let Err(err) = kill_process_group(process_group_id) { + debug!("failed to kill exec-server stdio process group {process_group_id}: {err}"); + } + }); + } +} + +#[cfg(windows)] +fn terminate_stdio_child_process(process_id: Option, child: &mut Child) { + if let Some(process_id) = process_id { + let _ = std::process::Command::new("taskkill") + .arg("/PID") + .arg(process_id.to_string()) + .arg("/T") + .arg("/F") + .output(); + } + kill_stdio_child(child); +} + +#[cfg(not(any(unix, windows)))] +fn terminate_stdio_child_process(_process_id: Option, child: &mut Child) { + kill_stdio_child(child); } fn kill_stdio_child(child: &mut Child) { @@ -171,6 +221,7 @@ fn shell_command_process(shell_command: &str) -> Command { { let mut command = Command::new("sh"); command.arg("-lc").arg(shell_command); + command.process_group(0); command } } From 55bcc97228ac649337919961acc4b49c4ea51891 Mon Sep 17 00:00:00 2001 From: starr-openai Date: Mon, 4 May 2026 12:52:40 -0700 Subject: [PATCH 06/42] Simplify exec-server transport internals Keep environment transport connection policy on ExecServerClient instead of the transport enum, and replace the JSON-RPC connection tuple alias with named connection parts. Co-authored-by: Codex --- codex-rs/exec-server/src/client.rs | 2 +- codex-rs/exec-server/src/client_transport.rs | 19 ++++++------- codex-rs/exec-server/src/connection.rs | 29 ++++++++++---------- codex-rs/exec-server/src/rpc.rs | 8 ++++-- codex-rs/exec-server/src/server/processor.rs | 13 ++++----- 5 files changed, 37 insertions(+), 34 deletions(-) diff --git a/codex-rs/exec-server/src/client.rs b/codex-rs/exec-server/src/client.rs index c04c57eb25d4..54a67fea3cd1 100644 --- a/codex-rs/exec-server/src/client.rs +++ b/codex-rs/exec-server/src/client.rs @@ -207,7 +207,7 @@ impl LazyRemoteExecServerClient { self.client .get_or_try_init(|| { let transport = self.transport.clone(); - async move { transport.connect_for_environment().await } + async move { ExecServerClient::connect_for_environment(transport).await } }) .await .cloned() diff --git a/codex-rs/exec-server/src/client_transport.rs b/codex-rs/exec-server/src/client_transport.rs index d877aae3dfb4..00ecc2bd2e04 100644 --- a/codex-rs/exec-server/src/client_transport.rs +++ b/codex-rs/exec-server/src/client_transport.rs @@ -21,7 +21,6 @@ use tracing::warn; use crate::ExecServerClient; use crate::ExecServerError; -use crate::client_api::ExecServerTransport; use crate::client_api::RemoteExecServerConnectArgs; use crate::client_api::StdioExecServerConnectArgs; use crate::connection::JsonRpcConnection; @@ -32,11 +31,13 @@ const ENVIRONMENT_INITIALIZE_TIMEOUT: Duration = Duration::from_secs(5); #[cfg(unix)] const STDIO_CHILD_TERM_GRACE_PERIOD: Duration = Duration::from_millis(500); -impl ExecServerTransport { - pub(crate) async fn connect_for_environment(self) -> Result { - match self { - ExecServerTransport::WebSocketUrl(websocket_url) => { - ExecServerClient::connect_websocket(RemoteExecServerConnectArgs { +impl ExecServerClient { + pub(crate) async fn connect_for_environment( + transport: crate::client_api::ExecServerTransport, + ) -> Result { + match transport { + crate::client_api::ExecServerTransport::WebSocketUrl(websocket_url) => { + Self::connect_websocket(RemoteExecServerConnectArgs { websocket_url, client_name: ENVIRONMENT_CLIENT_NAME.to_string(), connect_timeout: ENVIRONMENT_CONNECT_TIMEOUT, @@ -45,8 +46,8 @@ impl ExecServerTransport { }) .await } - ExecServerTransport::StdioShellCommand(shell_command) => { - ExecServerClient::connect_stdio_command(StdioExecServerConnectArgs { + crate::client_api::ExecServerTransport::StdioShellCommand(shell_command) => { + Self::connect_stdio_command(StdioExecServerConnectArgs { shell_command, client_name: ENVIRONMENT_CLIENT_NAME.to_string(), initialize_timeout: ENVIRONMENT_INITIALIZE_TIMEOUT, @@ -56,9 +57,7 @@ impl ExecServerTransport { } } } -} -impl ExecServerClient { pub async fn connect_websocket( args: RemoteExecServerConnectArgs, ) -> Result { diff --git a/codex-rs/exec-server/src/connection.rs b/codex-rs/exec-server/src/connection.rs index 0ec4c6bc5f1b..9d06f0841d00 100644 --- a/codex-rs/exec-server/src/connection.rs +++ b/codex-rs/exec-server/src/connection.rs @@ -16,13 +16,14 @@ use tokio::io::BufWriter; pub(crate) const CHANNEL_CAPACITY: usize = 128; pub(crate) type JsonRpcTransportLifetime = Box; -pub(crate) type JsonRpcConnectionParts = ( - mpsc::Sender, - mpsc::Receiver, - watch::Receiver, - Vec>, - Option, -); + +pub(crate) struct JsonRpcConnectionParts { + pub(crate) outgoing_tx: mpsc::Sender, + pub(crate) incoming_rx: mpsc::Receiver, + pub(crate) disconnected_rx: watch::Receiver, + pub(crate) task_handles: Vec>, + pub(crate) transport_lifetime: Option, +} #[derive(Debug)] pub(crate) enum JsonRpcConnectionEvent { @@ -272,13 +273,13 @@ impl JsonRpcConnection { } pub(crate) fn into_parts(self) -> JsonRpcConnectionParts { - ( - self.outgoing_tx, - self.incoming_rx, - self.disconnected_rx, - self.task_handles, - self.transport_lifetime, - ) + JsonRpcConnectionParts { + outgoing_tx: self.outgoing_tx, + incoming_rx: self.incoming_rx, + disconnected_rx: self.disconnected_rx, + task_handles: self.task_handles, + transport_lifetime: self.transport_lifetime, + } } } diff --git a/codex-rs/exec-server/src/rpc.rs b/codex-rs/exec-server/src/rpc.rs index d9d8fbbf726d..3b155c08a2e0 100644 --- a/codex-rs/exec-server/src/rpc.rs +++ b/codex-rs/exec-server/src/rpc.rs @@ -243,8 +243,12 @@ struct TransportLifetime { impl RpcClient { pub(crate) fn new(connection: JsonRpcConnection) -> (Self, mpsc::Receiver) { - let (write_tx, mut incoming_rx, disconnected_rx, transport_tasks, transport_lifetime) = - connection.into_parts(); + let connection_parts = connection.into_parts(); + let write_tx = connection_parts.outgoing_tx; + let mut incoming_rx = connection_parts.incoming_rx; + let disconnected_rx = connection_parts.disconnected_rx; + let transport_tasks = connection_parts.task_handles; + let transport_lifetime = connection_parts.transport_lifetime; let pending = Arc::new(Mutex::new(HashMap::::new())); let (event_tx, event_rx) = mpsc::channel(128); diff --git a/codex-rs/exec-server/src/server/processor.rs b/codex-rs/exec-server/src/server/processor.rs index b7e1a03bd529..11472dc6264a 100644 --- a/codex-rs/exec-server/src/server/processor.rs +++ b/codex-rs/exec-server/src/server/processor.rs @@ -47,13 +47,12 @@ async fn run_connection( runtime_paths: ExecServerRuntimePaths, ) { let router = Arc::new(build_router()); - let ( - json_outgoing_tx, - mut incoming_rx, - mut disconnected_rx, - connection_tasks, - _transport_lifetime, - ) = connection.into_parts(); + let connection_parts = connection.into_parts(); + let json_outgoing_tx = connection_parts.outgoing_tx; + let mut incoming_rx = connection_parts.incoming_rx; + let mut disconnected_rx = connection_parts.disconnected_rx; + let connection_tasks = connection_parts.task_handles; + let _transport_lifetime = connection_parts.transport_lifetime; let (outgoing_tx, mut outgoing_rx) = mpsc::channel::(CHANNEL_CAPACITY); let notifications = RpcNotificationSender::new(outgoing_tx.clone()); From 52ca8fa8b83733df8d02d7d9133980ed6ea16fd8 Mon Sep 17 00:00:00 2001 From: starr-openai Date: Tue, 5 May 2026 10:08:36 -0700 Subject: [PATCH 07/42] Address stdio exec-server review feedback Spawn stdio exec-server commands directly from structured argv/env/cwd instead of wrapping a shell string, redact the connection label, and tie the stdio child guard to transport disconnect. Co-authored-by: Codex --- codex-rs/exec-server/src/client.rs | 137 +++++++++++++++---- codex-rs/exec-server/src/client_api.rs | 15 +- codex-rs/exec-server/src/client_transport.rs | 127 +++++------------ codex-rs/exec-server/src/connection.rs | 18 ++- codex-rs/exec-server/src/lib.rs | 1 + codex-rs/exec-server/src/rpc.rs | 19 +-- codex-rs/exec-server/src/server/processor.rs | 8 +- 7 files changed, 184 insertions(+), 141 deletions(-) diff --git a/codex-rs/exec-server/src/client.rs b/codex-rs/exec-server/src/client.rs index 54a67fea3cd1..2842838ce423 100644 --- a/codex-rs/exec-server/src/client.rs +++ b/codex-rs/exec-server/src/client.rs @@ -12,7 +12,6 @@ use futures::FutureExt; use futures::future::BoxFuture; use serde_json::Value; use tokio::sync::Mutex; -use tokio::sync::OnceCell; use tokio::sync::mpsc; use tokio::sync::watch; @@ -192,25 +191,28 @@ pub struct ExecServerClient { #[derive(Clone)] pub(crate) struct LazyRemoteExecServerClient { transport: ExecServerTransport, - client: Arc>, + client: Arc>>, } impl LazyRemoteExecServerClient { pub(crate) fn new(transport: ExecServerTransport) -> Self { Self { transport, - client: Arc::new(OnceCell::new()), + client: Arc::new(Mutex::new(None)), } } pub(crate) async fn get(&self) -> Result { - self.client - .get_or_try_init(|| { - let transport = self.transport.clone(); - async move { ExecServerClient::connect_for_environment(transport).await } - }) - .await - .cloned() + let mut client = self.client.lock().await; + if let Some(client) = client.as_ref() + && !client.is_disconnected() + { + return Ok(client.clone()); + } + + let connected = ExecServerClient::connect_for_environment(self.transport.clone()).await?; + *client = Some(connected.clone()); + Ok(connected) } } @@ -274,6 +276,10 @@ pub enum ExecServerError { } impl ExecServerClient { + fn is_disconnected(&self) -> bool { + self.inner.disconnected_error().is_some() || self.inner.client.is_disconnected() + } + pub async fn initialize( &self, options: ExecServerClientConnectOptions, @@ -872,6 +878,7 @@ mod tests { use codex_app_server_protocol::JSONRPCNotification; use codex_app_server_protocol::JSONRPCResponse; use pretty_assertions::assert_eq; + use std::collections::HashMap; #[cfg(unix)] use std::path::Path; #[cfg(unix)] @@ -890,7 +897,7 @@ mod tests { use super::ExecServerClient; use super::ExecServerClientConnectOptions; use crate::ProcessId; - #[cfg(not(windows))] + use crate::StdioExecServerCommand; use crate::StdioExecServerConnectArgs; use crate::connection::JsonRpcConnection; use crate::process::ExecProcessEvent; @@ -933,7 +940,38 @@ mod tests { #[tokio::test] async fn connect_stdio_command_initializes_json_rpc_client() { let client = ExecServerClient::connect_stdio_command(StdioExecServerConnectArgs { - shell_command: "read _line; printf '%s\\n' '{\"id\":1,\"result\":{\"sessionId\":\"stdio-test\"}}'; read _line; sleep 60".to_string(), + command: StdioExecServerCommand { + program: "sh".to_string(), + args: vec![ + "-c".to_string(), + "read _line; printf '%s\\n' '{\"id\":1,\"result\":{\"sessionId\":\"stdio-test\"}}'; read _line; sleep 60".to_string(), + ], + env: HashMap::new(), + cwd: None, + }, + client_name: "stdio-test-client".to_string(), + initialize_timeout: Duration::from_secs(1), + resume_session_id: None, + }) + .await + .expect("stdio client should connect"); + + assert_eq!(client.session_id().as_deref(), Some("stdio-test")); + } + + #[cfg(windows)] + #[tokio::test] + async fn connect_stdio_command_initializes_json_rpc_client_on_windows() { + let client = ExecServerClient::connect_stdio_command(StdioExecServerConnectArgs { + command: StdioExecServerCommand { + program: "cmd".to_string(), + args: vec![ + "/C".to_string(), + "set /p _line= & echo {\"id\":1,\"result\":{\"sessionId\":\"stdio-test\"}} & set /p _line= & ping -n 60 127.0.0.1 >nul".to_string(), + ], + env: HashMap::new(), + cwd: None, + }, client_name: "stdio-test-client".to_string(), initialize_timeout: Duration::from_secs(1), resume_session_id: None, @@ -946,43 +984,71 @@ mod tests { #[cfg(unix)] #[tokio::test] - async fn dropping_stdio_client_terminates_shell_process_group() { + async fn dropping_stdio_client_terminates_spawned_process() { let tempdir = tempfile::tempdir().expect("tempdir should be created"); - let pid_file = tempdir.path().join("child.pid"); - let shell_command = format!( + let pid_file = tempdir.path().join("server.pid"); + let stdio_script = format!( "read _line; \ - (trap 'exit 0' TERM; while true; do sleep 1; done) & \ - child=$!; \ - echo \"$child\" > {}; \ + echo \"$$\" > {}; \ printf '%s\\n' '{{\"id\":1,\"result\":{{\"sessionId\":\"stdio-test\"}}}}'; \ read _line; \ - wait \"$child\"", + sleep 60", shell_quote(pid_file.as_path()), ); let client = ExecServerClient::connect_stdio_command(StdioExecServerConnectArgs { - shell_command, + command: StdioExecServerCommand { + program: "sh".to_string(), + args: vec!["-c".to_string(), stdio_script], + env: HashMap::new(), + cwd: None, + }, client_name: "stdio-test-client".to_string(), initialize_timeout: Duration::from_secs(1), resume_session_id: None, }) .await .expect("stdio client should connect"); - let child_pid = read_pid_file(pid_file.as_path()).await; + let server_pid = read_pid_file(pid_file.as_path()).await; assert!( - process_exists(child_pid), - "wrapper child process should be running before client drop" + process_exists(server_pid), + "spawned stdio process should be running before client drop" ); drop(client); - for _ in 0..20 { - if !process_exists(child_pid) { - return; - } - sleep(Duration::from_millis(100)).await; - } - panic!("wrapper child process {child_pid} should exit after client drop"); + wait_for_process_exit(server_pid).await; + } + + #[cfg(unix)] + #[tokio::test] + async fn malformed_stdio_message_terminates_spawned_process() { + let tempdir = tempfile::tempdir().expect("tempdir should be created"); + let pid_file = tempdir.path().join("server.pid"); + let stdio_script = format!( + "read _line; \ + echo \"$$\" > {}; \ + printf '%s\\n' 'not-json'; \ + sleep 60", + shell_quote(pid_file.as_path()), + ); + + let result = ExecServerClient::connect_stdio_command(StdioExecServerConnectArgs { + command: StdioExecServerCommand { + program: "sh".to_string(), + args: vec!["-c".to_string(), stdio_script], + env: HashMap::new(), + cwd: None, + }, + client_name: "stdio-test-client".to_string(), + initialize_timeout: Duration::from_secs(1), + resume_session_id: None, + }) + .await; + assert!(result.is_err(), "malformed stdio server should not connect"); + + let server_pid = read_pid_file(pid_file.as_path()).await; + wait_for_process_exit(server_pid).await; } #[cfg(unix)] @@ -999,6 +1065,17 @@ mod tests { panic!("pid file {} should be written", path.display()); } + #[cfg(unix)] + async fn wait_for_process_exit(pid: u32) { + for _ in 0..20 { + if !process_exists(pid) { + return; + } + sleep(Duration::from_millis(100)).await; + } + panic!("process {pid} should exit"); + } + #[cfg(unix)] fn process_exists(pid: u32) -> bool { Command::new("kill") diff --git a/codex-rs/exec-server/src/client_api.rs b/codex-rs/exec-server/src/client_api.rs index 95ed053476d9..8110ee24e56a 100644 --- a/codex-rs/exec-server/src/client_api.rs +++ b/codex-rs/exec-server/src/client_api.rs @@ -1,3 +1,5 @@ +use std::collections::HashMap; +use std::path::PathBuf; use std::time::Duration; use futures::future::BoxFuture; @@ -28,17 +30,26 @@ pub struct RemoteExecServerConnectArgs { /// Stdio connection arguments for a command-backed exec-server. #[derive(Debug, Clone, PartialEq, Eq)] pub struct StdioExecServerConnectArgs { - pub shell_command: String, + pub command: StdioExecServerCommand, pub client_name: String, pub initialize_timeout: Duration, pub resume_session_id: Option, } +/// Structured process command used to start an exec-server over stdio. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct StdioExecServerCommand { + pub program: String, + pub args: Vec, + pub env: HashMap, + pub cwd: Option, +} + /// Transport used to connect to a remote exec-server environment. #[derive(Debug, Clone, PartialEq, Eq)] pub enum ExecServerTransport { WebSocketUrl(String), - StdioShellCommand(String), + StdioCommand(StdioExecServerCommand), } /// Sends HTTP requests through a runtime-selected transport. diff --git a/codex-rs/exec-server/src/client_transport.rs b/codex-rs/exec-server/src/client_transport.rs index 00ecc2bd2e04..d6d1c92deb33 100644 --- a/codex-rs/exec-server/src/client_transport.rs +++ b/codex-rs/exec-server/src/client_transport.rs @@ -1,19 +1,11 @@ use std::process::Stdio; -#[cfg(unix)] -use std::thread::sleep; -#[cfg(unix)] -use std::thread::spawn; use std::time::Duration; -#[cfg(unix)] -use codex_utils_pty::process_group::kill_process_group; -#[cfg(unix)] -use codex_utils_pty::process_group::terminate_process_group; use tokio::io::AsyncBufReadExt; use tokio::io::BufReader; use tokio::process::Child; use tokio::process::Command; -use tokio::runtime::Handle; +use tokio::sync::oneshot; use tokio::time::timeout; use tokio_tungstenite::connect_async; use tracing::debug; @@ -22,14 +14,13 @@ use tracing::warn; use crate::ExecServerClient; use crate::ExecServerError; use crate::client_api::RemoteExecServerConnectArgs; +use crate::client_api::StdioExecServerCommand; use crate::client_api::StdioExecServerConnectArgs; use crate::connection::JsonRpcConnection; const ENVIRONMENT_CLIENT_NAME: &str = "codex-environment"; const ENVIRONMENT_CONNECT_TIMEOUT: Duration = Duration::from_secs(5); const ENVIRONMENT_INITIALIZE_TIMEOUT: Duration = Duration::from_secs(5); -#[cfg(unix)] -const STDIO_CHILD_TERM_GRACE_PERIOD: Duration = Duration::from_millis(500); impl ExecServerClient { pub(crate) async fn connect_for_environment( @@ -46,9 +37,9 @@ impl ExecServerClient { }) .await } - crate::client_api::ExecServerTransport::StdioShellCommand(shell_command) => { + crate::client_api::ExecServerTransport::StdioCommand(command) => { Self::connect_stdio_command(StdioExecServerConnectArgs { - shell_command, + command, client_name: ENVIRONMENT_CLIENT_NAME.to_string(), initialize_timeout: ENVIRONMENT_INITIALIZE_TIMEOUT, resume_session_id: None, @@ -87,15 +78,13 @@ impl ExecServerClient { pub async fn connect_stdio_command( args: StdioExecServerConnectArgs, ) -> Result { - let shell_command = args.shell_command.clone(); - let mut child = shell_command_process(&shell_command) + let mut child = stdio_command_process(&args.command) .kill_on_drop(true) .stdin(Stdio::piped()) .stdout(Stdio::piped()) .stderr(Stdio::piped()) .spawn() .map_err(ExecServerError::Spawn)?; - let process_id = child.id(); let stdin = child.stdin.take().ok_or_else(|| { ExecServerError::Protocol("spawned exec-server command has no stdin".to_string()) @@ -120,15 +109,8 @@ impl ExecServerClient { } Self::connect( - JsonRpcConnection::from_stdio( - stdout, - stdin, - format!("exec-server stdio command `{shell_command}`"), - ) - .with_transport_lifetime(Box::new(StdioChildGuard { - child: Some(child), - process_id, - })), + JsonRpcConnection::from_stdio(stdout, stdin, "exec-server stdio command".to_string()) + .with_transport_lifetime(Box::new(StdioChildGuard::spawn(child))), args.into(), ) .await @@ -136,70 +118,44 @@ impl ExecServerClient { } struct StdioChildGuard { - child: Option, - process_id: Option, + shutdown_tx: Option>, } -impl Drop for StdioChildGuard { - fn drop(&mut self) { - let Some(mut child) = self.child.take() else { - return; - }; - - terminate_stdio_child_process(self.process_id, &mut child); - - if let Ok(handle) = Handle::try_current() { - let _wait_task = handle.spawn(wait_stdio_child(child)); +impl StdioChildGuard { + fn spawn(child: Child) -> Self { + let (shutdown_tx, shutdown_rx) = oneshot::channel(); + tokio::spawn(supervise_stdio_child(child, shutdown_rx)); + Self { + shutdown_tx: Some(shutdown_tx), } } } -async fn wait_stdio_child(mut child: Child) { - if let Err(err) = child.wait().await { - debug!("failed to wait for exec-server stdio child: {err}"); +impl Drop for StdioChildGuard { + fn drop(&mut self) { + if let Some(shutdown_tx) = self.shutdown_tx.take() { + let _ = shutdown_tx.send(()); + } } } -#[cfg(unix)] -fn terminate_stdio_child_process(process_group_id: Option, child: &mut Child) { - let Some(process_group_id) = process_group_id else { - kill_stdio_child(child); - return; - }; - - let should_escalate = match terminate_process_group(process_group_id) { - Ok(exists) => exists, - Err(err) => { - debug!("failed to terminate exec-server stdio process group {process_group_id}: {err}"); +async fn supervise_stdio_child(mut child: Child, shutdown_rx: oneshot::Receiver<()>) { + let shutdown_requested = tokio::select! { + result = child.wait() => { + if let Err(err) = result { + debug!("failed to wait for exec-server stdio child: {err}"); + } false } + _ = shutdown_rx => true, }; - if should_escalate { - spawn(move || { - sleep(STDIO_CHILD_TERM_GRACE_PERIOD); - if let Err(err) = kill_process_group(process_group_id) { - debug!("failed to kill exec-server stdio process group {process_group_id}: {err}"); - } - }); - } -} -#[cfg(windows)] -fn terminate_stdio_child_process(process_id: Option, child: &mut Child) { - if let Some(process_id) = process_id { - let _ = std::process::Command::new("taskkill") - .arg("/PID") - .arg(process_id.to_string()) - .arg("/T") - .arg("/F") - .output(); + if shutdown_requested { + kill_stdio_child(&mut child); + if let Err(err) = child.wait().await { + debug!("failed to wait for exec-server stdio child after shutdown: {err}"); + } } - kill_stdio_child(child); -} - -#[cfg(not(any(unix, windows)))] -fn terminate_stdio_child_process(_process_id: Option, child: &mut Child) { - kill_stdio_child(child); } fn kill_stdio_child(child: &mut Child) { @@ -208,19 +164,12 @@ fn kill_stdio_child(child: &mut Child) { } } -fn shell_command_process(shell_command: &str) -> Command { - #[cfg(windows)] - { - let mut command = Command::new("cmd"); - command.arg("/C").arg(shell_command); - command - } - - #[cfg(not(windows))] - { - let mut command = Command::new("sh"); - command.arg("-lc").arg(shell_command); - command.process_group(0); - command +fn stdio_command_process(stdio_command: &StdioExecServerCommand) -> Command { + let mut command = Command::new(&stdio_command.program); + command.args(&stdio_command.args); + command.envs(&stdio_command.env); + if let Some(cwd) = &stdio_command.cwd { + command.current_dir(cwd); } + command } diff --git a/codex-rs/exec-server/src/connection.rs b/codex-rs/exec-server/src/connection.rs index 9d06f0841d00..c0a4a42a5404 100644 --- a/codex-rs/exec-server/src/connection.rs +++ b/codex-rs/exec-server/src/connection.rs @@ -272,7 +272,23 @@ impl JsonRpcConnection { self } - pub(crate) fn into_parts(self) -> JsonRpcConnectionParts { + pub(crate) fn into_parts( + self, + ) -> ( + mpsc::Sender, + mpsc::Receiver, + watch::Receiver, + Vec>, + ) { + ( + self.outgoing_tx, + self.incoming_rx, + self.disconnected_rx, + self.task_handles, + ) + } + + pub(crate) fn into_parts_with_lifetime(self) -> JsonRpcConnectionParts { JsonRpcConnectionParts { outgoing_tx: self.outgoing_tx, incoming_rx: self.incoming_rx, diff --git a/codex-rs/exec-server/src/lib.rs b/codex-rs/exec-server/src/lib.rs index 9bec4ed1da6f..29897c528532 100644 --- a/codex-rs/exec-server/src/lib.rs +++ b/codex-rs/exec-server/src/lib.rs @@ -28,6 +28,7 @@ pub use client_api::ExecServerClientConnectOptions; pub use client_api::ExecServerTransport; pub use client_api::HttpClient; pub use client_api::RemoteExecServerConnectArgs; +pub use client_api::StdioExecServerCommand; pub use client_api::StdioExecServerConnectArgs; pub use codex_file_system::CopyOptions; pub use codex_file_system::CreateDirectoryOptions; diff --git a/codex-rs/exec-server/src/rpc.rs b/codex-rs/exec-server/src/rpc.rs index 3b155c08a2e0..49438b4fd806 100644 --- a/codex-rs/exec-server/src/rpc.rs +++ b/codex-rs/exec-server/src/rpc.rs @@ -2,7 +2,6 @@ use std::collections::HashMap; use std::future::Future; use std::pin::Pin; use std::sync::Arc; -use std::sync::Mutex as StdMutex; use std::sync::atomic::AtomicI64; use std::sync::atomic::Ordering; @@ -24,7 +23,6 @@ use tokio::task::JoinHandle; use crate::connection::JsonRpcConnection; use crate::connection::JsonRpcConnectionEvent; -use crate::connection::JsonRpcTransportLifetime; #[derive(Debug)] pub(crate) enum RpcCallError { @@ -231,19 +229,12 @@ pub(crate) struct RpcClient { disconnected_rx: watch::Receiver, next_request_id: AtomicI64, transport_tasks: Vec>, - _transport_lifetime: Option, reader_task: JoinHandle<()>, } -// Holds transport-owned resources, such as a stdio child process, for as long -// as the RPC client owns the underlying connection. -struct TransportLifetime { - _guard: StdMutex, -} - impl RpcClient { pub(crate) fn new(connection: JsonRpcConnection) -> (Self, mpsc::Receiver) { - let connection_parts = connection.into_parts(); + let connection_parts = connection.into_parts_with_lifetime(); let write_tx = connection_parts.outgoing_tx; let mut incoming_rx = connection_parts.incoming_rx; let disconnected_rx = connection_parts.disconnected_rx; @@ -254,6 +245,7 @@ impl RpcClient { let pending_for_reader = Arc::clone(&pending); let reader_task = tokio::spawn(async move { + let _transport_lifetime = transport_lifetime; while let Some(event) = incoming_rx.recv().await { match event { JsonRpcConnectionEvent::Message(message) => { @@ -289,9 +281,6 @@ impl RpcClient { disconnected_rx, next_request_id: AtomicI64::new(1), transport_tasks, - _transport_lifetime: transport_lifetime.map(|lifetime| TransportLifetime { - _guard: StdMutex::new(lifetime), - }), reader_task, }, event_rx, @@ -318,6 +307,10 @@ impl RpcClient { }) } + pub(crate) fn is_disconnected(&self) -> bool { + *self.disconnected_rx.borrow() + } + pub(crate) async fn call(&self, method: &str, params: &P) -> Result where P: Serialize, diff --git a/codex-rs/exec-server/src/server/processor.rs b/codex-rs/exec-server/src/server/processor.rs index 11472dc6264a..dc1a9b9ffe74 100644 --- a/codex-rs/exec-server/src/server/processor.rs +++ b/codex-rs/exec-server/src/server/processor.rs @@ -47,12 +47,8 @@ async fn run_connection( runtime_paths: ExecServerRuntimePaths, ) { let router = Arc::new(build_router()); - let connection_parts = connection.into_parts(); - let json_outgoing_tx = connection_parts.outgoing_tx; - let mut incoming_rx = connection_parts.incoming_rx; - let mut disconnected_rx = connection_parts.disconnected_rx; - let connection_tasks = connection_parts.task_handles; - let _transport_lifetime = connection_parts.transport_lifetime; + let (json_outgoing_tx, mut incoming_rx, mut disconnected_rx, connection_tasks) = + connection.into_parts(); let (outgoing_tx, mut outgoing_rx) = mpsc::channel::(CHANNEL_CAPACITY); let notifications = RpcNotificationSender::new(outgoing_tx.clone()); From 21faf08349ea4b7f9b5e2526bf35ff7abbf589af Mon Sep 17 00:00:00 2001 From: starr-openai Date: Tue, 5 May 2026 13:17:06 -0700 Subject: [PATCH 08/42] Simplify stdio exec-server transport ownership Co-authored-by: Codex --- codex-rs/exec-server/src/client.rs | 34 ++++--- codex-rs/exec-server/src/client_transport.rs | 52 +---------- codex-rs/exec-server/src/connection.rs | 93 ++++++++++++-------- codex-rs/exec-server/src/rpc.rs | 25 ++---- 4 files changed, 81 insertions(+), 123 deletions(-) diff --git a/codex-rs/exec-server/src/client.rs b/codex-rs/exec-server/src/client.rs index 2842838ce423..37f3e2455a67 100644 --- a/codex-rs/exec-server/src/client.rs +++ b/codex-rs/exec-server/src/client.rs @@ -12,6 +12,7 @@ use futures::FutureExt; use futures::future::BoxFuture; use serde_json::Value; use tokio::sync::Mutex; +use tokio::sync::OnceCell; use tokio::sync::mpsc; use tokio::sync::watch; @@ -152,6 +153,9 @@ pub(crate) struct Session { struct Inner { client: RpcClient, + // Keep the connection alive for any transport-specific owned state such as + // the stdio child process. RpcClient only takes the runtime channels/tasks. + _connection: JsonRpcConnection, // The remote transport delivers one shared notification stream for every // process on the connection. Keep a local process_id -> session registry so // we can turn those connection-global notifications into process wakeups @@ -191,28 +195,25 @@ pub struct ExecServerClient { #[derive(Clone)] pub(crate) struct LazyRemoteExecServerClient { transport: ExecServerTransport, - client: Arc>>, + client: Arc>, } impl LazyRemoteExecServerClient { pub(crate) fn new(transport: ExecServerTransport) -> Self { Self { transport, - client: Arc::new(Mutex::new(None)), + client: Arc::new(OnceCell::new()), } } pub(crate) async fn get(&self) -> Result { - let mut client = self.client.lock().await; - if let Some(client) = client.as_ref() - && !client.is_disconnected() - { - return Ok(client.clone()); - } - - let connected = ExecServerClient::connect_for_environment(self.transport.clone()).await?; - *client = Some(connected.clone()); - Ok(connected) + self.client + .get_or_try_init(|| { + let transport = self.transport.clone(); + async move { ExecServerClient::connect_for_environment(transport).await } + }) + .await + .cloned() } } @@ -276,10 +277,6 @@ pub enum ExecServerError { } impl ExecServerClient { - fn is_disconnected(&self) -> bool { - self.inner.disconnected_error().is_some() || self.inner.client.is_disconnected() - } - pub async fn initialize( &self, options: ExecServerClientConnectOptions, @@ -429,10 +426,10 @@ impl ExecServerClient { } pub(crate) async fn connect( - connection: JsonRpcConnection, + mut connection: JsonRpcConnection, options: ExecServerClientConnectOptions, ) -> Result { - let (rpc_client, mut events_rx) = RpcClient::new(connection); + let (rpc_client, mut events_rx) = RpcClient::new(&mut connection); let inner = Arc::new_cyclic(|weak| { let weak = weak.clone(); let reader_task = tokio::spawn(async move { @@ -467,6 +464,7 @@ impl ExecServerClient { Inner { client: rpc_client, + _connection: connection, sessions: ArcSwap::from_pointee(HashMap::new()), sessions_write_lock: Mutex::new(()), disconnected: OnceLock::new(), diff --git a/codex-rs/exec-server/src/client_transport.rs b/codex-rs/exec-server/src/client_transport.rs index d6d1c92deb33..df9d84beabb8 100644 --- a/codex-rs/exec-server/src/client_transport.rs +++ b/codex-rs/exec-server/src/client_transport.rs @@ -3,9 +3,7 @@ use std::time::Duration; use tokio::io::AsyncBufReadExt; use tokio::io::BufReader; -use tokio::process::Child; use tokio::process::Command; -use tokio::sync::oneshot; use tokio::time::timeout; use tokio_tungstenite::connect_async; use tracing::debug; @@ -79,7 +77,6 @@ impl ExecServerClient { args: StdioExecServerConnectArgs, ) -> Result { let mut child = stdio_command_process(&args.command) - .kill_on_drop(true) .stdin(Stdio::piped()) .stdout(Stdio::piped()) .stderr(Stdio::piped()) @@ -110,60 +107,13 @@ impl ExecServerClient { Self::connect( JsonRpcConnection::from_stdio(stdout, stdin, "exec-server stdio command".to_string()) - .with_transport_lifetime(Box::new(StdioChildGuard::spawn(child))), + .with_stdio_child(child), args.into(), ) .await } } -struct StdioChildGuard { - shutdown_tx: Option>, -} - -impl StdioChildGuard { - fn spawn(child: Child) -> Self { - let (shutdown_tx, shutdown_rx) = oneshot::channel(); - tokio::spawn(supervise_stdio_child(child, shutdown_rx)); - Self { - shutdown_tx: Some(shutdown_tx), - } - } -} - -impl Drop for StdioChildGuard { - fn drop(&mut self) { - if let Some(shutdown_tx) = self.shutdown_tx.take() { - let _ = shutdown_tx.send(()); - } - } -} - -async fn supervise_stdio_child(mut child: Child, shutdown_rx: oneshot::Receiver<()>) { - let shutdown_requested = tokio::select! { - result = child.wait() => { - if let Err(err) = result { - debug!("failed to wait for exec-server stdio child: {err}"); - } - false - } - _ = shutdown_rx => true, - }; - - if shutdown_requested { - kill_stdio_child(&mut child); - if let Err(err) = child.wait().await { - debug!("failed to wait for exec-server stdio child after shutdown: {err}"); - } - } -} - -fn kill_stdio_child(child: &mut Child) { - if let Err(err) = child.start_kill() { - debug!("failed to terminate exec-server stdio child: {err}"); - } -} - fn stdio_command_process(stdio_command: &StdioExecServerCommand) -> Command { let mut command = Command::new(&stdio_command.program); command.args(&stdio_command.args); diff --git a/codex-rs/exec-server/src/connection.rs b/codex-rs/exec-server/src/connection.rs index c0a4a42a5404..f7832e5a3414 100644 --- a/codex-rs/exec-server/src/connection.rs +++ b/codex-rs/exec-server/src/connection.rs @@ -3,10 +3,12 @@ use futures::SinkExt; use futures::StreamExt; use tokio::io::AsyncRead; use tokio::io::AsyncWrite; +use tokio::process::Child; use tokio::sync::mpsc; use tokio::sync::watch; use tokio_tungstenite::WebSocketStream; use tokio_tungstenite::tungstenite::Message; +use tracing::debug; use tokio::io::AsyncBufReadExt; use tokio::io::AsyncWriteExt; @@ -15,16 +17,6 @@ use tokio::io::BufWriter; pub(crate) const CHANNEL_CAPACITY: usize = 128; -pub(crate) type JsonRpcTransportLifetime = Box; - -pub(crate) struct JsonRpcConnectionParts { - pub(crate) outgoing_tx: mpsc::Sender, - pub(crate) incoming_rx: mpsc::Receiver, - pub(crate) disconnected_rx: watch::Receiver, - pub(crate) task_handles: Vec>, - pub(crate) transport_lifetime: Option, -} - #[derive(Debug)] pub(crate) enum JsonRpcConnectionEvent { Message(JSONRPCMessage), @@ -32,12 +24,24 @@ pub(crate) enum JsonRpcConnectionEvent { Disconnected { reason: Option }, } +struct StdioTransport { + child: Child, +} + +impl Drop for StdioTransport { + fn drop(&mut self) { + if let Err(err) = self.child.start_kill() { + debug!("failed to terminate exec-server stdio child: {err}"); + } + } +} + pub(crate) struct JsonRpcConnection { - outgoing_tx: mpsc::Sender, - incoming_rx: mpsc::Receiver, - disconnected_rx: watch::Receiver, + outgoing_tx: Option>, + incoming_rx: Option>, + disconnected_rx: Option>, task_handles: Vec>, - transport_lifetime: Option, + _stdio_transport: Option, } impl JsonRpcConnection { @@ -124,11 +128,11 @@ impl JsonRpcConnection { }); Self { - outgoing_tx, - incoming_rx, - disconnected_rx, + outgoing_tx: Some(outgoing_tx), + incoming_rx: Some(incoming_rx), + disconnected_rx: Some(disconnected_rx), task_handles: vec![reader_task, writer_task], - transport_lifetime: None, + _stdio_transport: None, } } @@ -259,16 +263,38 @@ impl JsonRpcConnection { }); Self { - outgoing_tx, - incoming_rx, - disconnected_rx, + outgoing_tx: Some(outgoing_tx), + incoming_rx: Some(incoming_rx), + disconnected_rx: Some(disconnected_rx), task_handles: vec![reader_task, writer_task], - transport_lifetime: None, + _stdio_transport: None, } } - pub(crate) fn with_transport_lifetime(mut self, lifetime: JsonRpcTransportLifetime) -> Self { - self.transport_lifetime = Some(lifetime); + pub(crate) fn take_client_runtime( + &mut self, + ) -> ( + mpsc::Sender, + mpsc::Receiver, + watch::Receiver, + Vec>, + ) { + ( + self.outgoing_tx + .take() + .expect("JSON-RPC client runtime already taken"), + self.incoming_rx + .take() + .expect("JSON-RPC client runtime already taken"), + self.disconnected_rx + .take() + .expect("JSON-RPC client runtime already taken"), + std::mem::take(&mut self.task_handles), + ) + } + + pub(crate) fn with_stdio_child(mut self, child: Child) -> Self { + self._stdio_transport = Some(StdioTransport { child }); self } @@ -281,22 +307,15 @@ impl JsonRpcConnection { Vec>, ) { ( - self.outgoing_tx, - self.incoming_rx, - self.disconnected_rx, + self.outgoing_tx + .expect("JSON-RPC connection parts already taken"), + self.incoming_rx + .expect("JSON-RPC connection parts already taken"), + self.disconnected_rx + .expect("JSON-RPC connection parts already taken"), self.task_handles, ) } - - pub(crate) fn into_parts_with_lifetime(self) -> JsonRpcConnectionParts { - JsonRpcConnectionParts { - outgoing_tx: self.outgoing_tx, - incoming_rx: self.incoming_rx, - disconnected_rx: self.disconnected_rx, - task_handles: self.task_handles, - transport_lifetime: self.transport_lifetime, - } - } } async fn send_disconnected( diff --git a/codex-rs/exec-server/src/rpc.rs b/codex-rs/exec-server/src/rpc.rs index 49438b4fd806..a9c77d549a3b 100644 --- a/codex-rs/exec-server/src/rpc.rs +++ b/codex-rs/exec-server/src/rpc.rs @@ -233,19 +233,16 @@ pub(crate) struct RpcClient { } impl RpcClient { - pub(crate) fn new(connection: JsonRpcConnection) -> (Self, mpsc::Receiver) { - let connection_parts = connection.into_parts_with_lifetime(); - let write_tx = connection_parts.outgoing_tx; - let mut incoming_rx = connection_parts.incoming_rx; - let disconnected_rx = connection_parts.disconnected_rx; - let transport_tasks = connection_parts.task_handles; - let transport_lifetime = connection_parts.transport_lifetime; + pub(crate) fn new( + connection: &mut JsonRpcConnection, + ) -> (Self, mpsc::Receiver) { + let (write_tx, mut incoming_rx, disconnected_rx, transport_tasks) = + connection.take_client_runtime(); let pending = Arc::new(Mutex::new(HashMap::::new())); let (event_tx, event_rx) = mpsc::channel(128); let pending_for_reader = Arc::clone(&pending); let reader_task = tokio::spawn(async move { - let _transport_lifetime = transport_lifetime; while let Some(event) = incoming_rx.recv().await { match event { JsonRpcConnectionEvent::Message(message) => { @@ -307,10 +304,6 @@ impl RpcClient { }) } - pub(crate) fn is_disconnected(&self) -> bool { - *self.disconnected_rx.borrow() - } - pub(crate) async fn call(&self, method: &str, params: &P) -> Result where P: Serialize, @@ -575,11 +568,9 @@ mod tests { async fn rpc_client_matches_out_of_order_responses_by_request_id() { let (client_stdin, server_reader) = tokio::io::duplex(4096); let (mut server_writer, client_stdout) = tokio::io::duplex(4096); - let (client, _events_rx) = RpcClient::new(JsonRpcConnection::from_stdio( - client_stdout, - client_stdin, - "test-rpc".to_string(), - )); + let mut connection = + JsonRpcConnection::from_stdio(client_stdout, client_stdin, "test-rpc".to_string()); + let (client, _events_rx) = RpcClient::new(&mut connection); let server = tokio::spawn(async move { let mut lines = BufReader::new(server_reader).lines(); From bedee6e8cfdafd9850410f9b322619f36b017a95 Mon Sep 17 00:00:00 2001 From: starr-openai Date: Tue, 5 May 2026 13:21:10 -0700 Subject: [PATCH 09/42] Clarify exec-server transport connect naming Co-authored-by: Codex --- codex-rs/exec-server/src/client.rs | 8 +++++--- codex-rs/exec-server/src/client_transport.rs | 2 +- 2 files changed, 6 insertions(+), 4 deletions(-) diff --git a/codex-rs/exec-server/src/client.rs b/codex-rs/exec-server/src/client.rs index 37f3e2455a67..df068466b82d 100644 --- a/codex-rs/exec-server/src/client.rs +++ b/codex-rs/exec-server/src/client.rs @@ -153,8 +153,8 @@ pub(crate) struct Session { struct Inner { client: RpcClient, - // Keep the connection alive for any transport-specific owned state such as - // the stdio child process. RpcClient only takes the runtime channels/tasks. + // Keep the underlying transport connection alive. RpcClient only takes the + // runtime channels/tasks it needs to drive the JSON-RPC client. _connection: JsonRpcConnection, // The remote transport delivers one shared notification stream for every // process on the connection. Keep a local process_id -> session registry so @@ -208,9 +208,11 @@ impl LazyRemoteExecServerClient { pub(crate) async fn get(&self) -> Result { self.client + // TODO: Add reconnect/disconnect handling here instead of reusing + // the first successfully initialized connection forever. .get_or_try_init(|| { let transport = self.transport.clone(); - async move { ExecServerClient::connect_for_environment(transport).await } + async move { ExecServerClient::connect_for_transport(transport).await } }) .await .cloned() diff --git a/codex-rs/exec-server/src/client_transport.rs b/codex-rs/exec-server/src/client_transport.rs index df9d84beabb8..7543d283b4ed 100644 --- a/codex-rs/exec-server/src/client_transport.rs +++ b/codex-rs/exec-server/src/client_transport.rs @@ -21,7 +21,7 @@ const ENVIRONMENT_CONNECT_TIMEOUT: Duration = Duration::from_secs(5); const ENVIRONMENT_INITIALIZE_TIMEOUT: Duration = Duration::from_secs(5); impl ExecServerClient { - pub(crate) async fn connect_for_environment( + pub(crate) async fn connect_for_transport( transport: crate::client_api::ExecServerTransport, ) -> Result { match transport { From 039011068287360d01cfc035a7851d2ac5579d44 Mon Sep 17 00:00:00 2001 From: starr-openai Date: Tue, 5 May 2026 13:23:19 -0700 Subject: [PATCH 10/42] Order exec-server transport teardown before RPC teardown Co-authored-by: Codex --- codex-rs/exec-server/src/client.rs | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/codex-rs/exec-server/src/client.rs b/codex-rs/exec-server/src/client.rs index df068466b82d..9b845d36113e 100644 --- a/codex-rs/exec-server/src/client.rs +++ b/codex-rs/exec-server/src/client.rs @@ -152,10 +152,10 @@ pub(crate) struct Session { } struct Inner { - client: RpcClient, - // Keep the underlying transport connection alive. RpcClient only takes the - // runtime channels/tasks it needs to drive the JSON-RPC client. + // Keep the underlying transport connection alive and drop it before the RPC + // client starts tearing down its channel/task handles. _connection: JsonRpcConnection, + client: RpcClient, // The remote transport delivers one shared notification stream for every // process on the connection. Keep a local process_id -> session registry so // we can turn those connection-global notifications into process wakeups @@ -465,8 +465,8 @@ impl ExecServerClient { }); Inner { - client: rpc_client, _connection: connection, + client: rpc_client, sessions: ArcSwap::from_pointee(HashMap::new()), sessions_write_lock: Mutex::new(()), disconnected: OnceLock::new(), From 1870847d44be6fa860af4c14feedf732ee966289 Mon Sep 17 00:00:00 2001 From: starr-openai Date: Tue, 5 May 2026 13:33:53 -0700 Subject: [PATCH 11/42] Name retained exec-server connection field Co-authored-by: Codex --- codex-rs/exec-server/src/client.rs | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/codex-rs/exec-server/src/client.rs b/codex-rs/exec-server/src/client.rs index 9b845d36113e..675e4e1b4d67 100644 --- a/codex-rs/exec-server/src/client.rs +++ b/codex-rs/exec-server/src/client.rs @@ -154,7 +154,8 @@ pub(crate) struct Session { struct Inner { // Keep the underlying transport connection alive and drop it before the RPC // client starts tearing down its channel/task handles. - _connection: JsonRpcConnection, + #[allow(dead_code)] + connection: JsonRpcConnection, client: RpcClient, // The remote transport delivers one shared notification stream for every // process on the connection. Keep a local process_id -> session registry so @@ -465,7 +466,7 @@ impl ExecServerClient { }); Inner { - _connection: connection, + connection, client: rpc_client, sessions: ArcSwap::from_pointee(HashMap::new()), sessions_write_lock: Mutex::new(()), From f3ba2aa4f87009222a9be3117d46659ea46937e0 Mon Sep 17 00:00:00 2001 From: starr-openai Date: Tue, 5 May 2026 13:36:02 -0700 Subject: [PATCH 12/42] Model retained JSON-RPC transport generically Co-authored-by: Codex --- codex-rs/exec-server/src/client_transport.rs | 2 +- codex-rs/exec-server/src/connection.rs | 102 +++++++++++-------- 2 files changed, 61 insertions(+), 43 deletions(-) diff --git a/codex-rs/exec-server/src/client_transport.rs b/codex-rs/exec-server/src/client_transport.rs index 7543d283b4ed..2f8f1b8a0581 100644 --- a/codex-rs/exec-server/src/client_transport.rs +++ b/codex-rs/exec-server/src/client_transport.rs @@ -107,7 +107,7 @@ impl ExecServerClient { Self::connect( JsonRpcConnection::from_stdio(stdout, stdin, "exec-server stdio command".to_string()) - .with_stdio_child(child), + .with_child_process(child), args.into(), ) .await diff --git a/codex-rs/exec-server/src/connection.rs b/codex-rs/exec-server/src/connection.rs index f7832e5a3414..07b3d7f4d592 100644 --- a/codex-rs/exec-server/src/connection.rs +++ b/codex-rs/exec-server/src/connection.rs @@ -24,24 +24,40 @@ pub(crate) enum JsonRpcConnectionEvent { Disconnected { reason: Option }, } -struct StdioTransport { - child: Child, +#[derive(Default)] +struct JsonRpcTransport { + child_process: Option, } -impl Drop for StdioTransport { +impl JsonRpcTransport { + fn with_child_process(child_process: Child) -> Self { + Self { + child_process: Some(child_process), + } + } +} + +impl Drop for JsonRpcTransport { fn drop(&mut self) { - if let Err(err) = self.child.start_kill() { + if let Some(child_process) = self.child_process.as_mut() + && let Err(err) = child_process.start_kill() + { debug!("failed to terminate exec-server stdio child: {err}"); } } } -pub(crate) struct JsonRpcConnection { - outgoing_tx: Option>, - incoming_rx: Option>, - disconnected_rx: Option>, +struct JsonRpcConnectionRuntime { + outgoing_tx: mpsc::Sender, + incoming_rx: mpsc::Receiver, + disconnected_rx: watch::Receiver, task_handles: Vec>, - _stdio_transport: Option, +} + +pub(crate) struct JsonRpcConnection { + runtime: Option, + #[allow(dead_code)] + transport: JsonRpcTransport, } impl JsonRpcConnection { @@ -128,11 +144,13 @@ impl JsonRpcConnection { }); Self { - outgoing_tx: Some(outgoing_tx), - incoming_rx: Some(incoming_rx), - disconnected_rx: Some(disconnected_rx), - task_handles: vec![reader_task, writer_task], - _stdio_transport: None, + runtime: Some(JsonRpcConnectionRuntime { + outgoing_tx, + incoming_rx, + disconnected_rx, + task_handles: vec![reader_task, writer_task], + }), + transport: JsonRpcTransport::default(), } } @@ -263,11 +281,13 @@ impl JsonRpcConnection { }); Self { - outgoing_tx: Some(outgoing_tx), - incoming_rx: Some(incoming_rx), - disconnected_rx: Some(disconnected_rx), - task_handles: vec![reader_task, writer_task], - _stdio_transport: None, + runtime: Some(JsonRpcConnectionRuntime { + outgoing_tx, + incoming_rx, + disconnected_rx, + task_handles: vec![reader_task, writer_task], + }), + transport: JsonRpcTransport::default(), } } @@ -279,22 +299,20 @@ impl JsonRpcConnection { watch::Receiver, Vec>, ) { - ( - self.outgoing_tx - .take() - .expect("JSON-RPC client runtime already taken"), - self.incoming_rx - .take() - .expect("JSON-RPC client runtime already taken"), - self.disconnected_rx - .take() - .expect("JSON-RPC client runtime already taken"), - std::mem::take(&mut self.task_handles), - ) + let JsonRpcConnectionRuntime { + outgoing_tx, + incoming_rx, + disconnected_rx, + task_handles, + } = self + .runtime + .take() + .expect("JSON-RPC client runtime already taken"); + (outgoing_tx, incoming_rx, disconnected_rx, task_handles) } - pub(crate) fn with_stdio_child(mut self, child: Child) -> Self { - self._stdio_transport = Some(StdioTransport { child }); + pub(crate) fn with_child_process(mut self, child_process: Child) -> Self { + self.transport = JsonRpcTransport::with_child_process(child_process); self } @@ -306,15 +324,15 @@ impl JsonRpcConnection { watch::Receiver, Vec>, ) { - ( - self.outgoing_tx - .expect("JSON-RPC connection parts already taken"), - self.incoming_rx - .expect("JSON-RPC connection parts already taken"), - self.disconnected_rx - .expect("JSON-RPC connection parts already taken"), - self.task_handles, - ) + let JsonRpcConnectionRuntime { + outgoing_tx, + incoming_rx, + disconnected_rx, + task_handles, + } = self + .runtime + .expect("JSON-RPC connection parts already taken"); + (outgoing_tx, incoming_rx, disconnected_rx, task_handles) } } From 9f771fb0b6297b0d7246e8b7b8fc9d97f1860235 Mon Sep 17 00:00:00 2001 From: starr-openai Date: Tue, 5 May 2026 13:38:57 -0700 Subject: [PATCH 13/42] Split JSON-RPC transport variants Co-authored-by: Codex --- codex-rs/exec-server/src/connection.rs | 35 ++++++++++++++++---------- 1 file changed, 22 insertions(+), 13 deletions(-) diff --git a/codex-rs/exec-server/src/connection.rs b/codex-rs/exec-server/src/connection.rs index 07b3d7f4d592..2dbda9e60403 100644 --- a/codex-rs/exec-server/src/connection.rs +++ b/codex-rs/exec-server/src/connection.rs @@ -24,24 +24,33 @@ pub(crate) enum JsonRpcConnectionEvent { Disconnected { reason: Option }, } -#[derive(Default)] -struct JsonRpcTransport { - child_process: Option, +enum JsonRpcTransport { + Plain, + Stdio(StdioTransport), } impl JsonRpcTransport { - fn with_child_process(child_process: Child) -> Self { - Self { - child_process: Some(child_process), - } + fn from_child_process(child_process: Child) -> Self { + Self::Stdio(StdioTransport { child_process }) } } impl Drop for JsonRpcTransport { fn drop(&mut self) { - if let Some(child_process) = self.child_process.as_mut() - && let Err(err) = child_process.start_kill() - { + match self { + Self::Plain => {} + Self::Stdio(transport) => transport.shutdown(), + } + } +} + +struct StdioTransport { + child_process: Child, +} + +impl StdioTransport { + fn shutdown(&mut self) { + if let Err(err) = self.child_process.start_kill() { debug!("failed to terminate exec-server stdio child: {err}"); } } @@ -150,7 +159,7 @@ impl JsonRpcConnection { disconnected_rx, task_handles: vec![reader_task, writer_task], }), - transport: JsonRpcTransport::default(), + transport: JsonRpcTransport::Plain, } } @@ -287,7 +296,7 @@ impl JsonRpcConnection { disconnected_rx, task_handles: vec![reader_task, writer_task], }), - transport: JsonRpcTransport::default(), + transport: JsonRpcTransport::Plain, } } @@ -312,7 +321,7 @@ impl JsonRpcConnection { } pub(crate) fn with_child_process(mut self, child_process: Child) -> Self { - self.transport = JsonRpcTransport::with_child_process(child_process); + self.transport = JsonRpcTransport::from_child_process(child_process); self } From 0190927d62014b2a119d5b4acac419364a7b3d72 Mon Sep 17 00:00:00 2001 From: starr-openai Date: Tue, 5 May 2026 13:44:58 -0700 Subject: [PATCH 14/42] Rename exec-server transport input params Co-authored-by: Codex --- codex-rs/exec-server/src/client.rs | 16 ++-- codex-rs/exec-server/src/client_api.rs | 8 +- codex-rs/exec-server/src/client_transport.rs | 10 +-- codex-rs/exec-server/src/connection.rs | 38 ++++++-- codex-rs/exec-server/src/environment.rs | 4 +- codex-rs/exec-server/src/lib.rs | 3 - codex-rs/exec-server/src/rpc.rs | 95 +++++++++++++++----- 7 files changed, 123 insertions(+), 51 deletions(-) diff --git a/codex-rs/exec-server/src/client.rs b/codex-rs/exec-server/src/client.rs index 675e4e1b4d67..8a1178c93d41 100644 --- a/codex-rs/exec-server/src/client.rs +++ b/codex-rs/exec-server/src/client.rs @@ -21,7 +21,7 @@ use tracing::debug; use crate::ProcessId; use crate::client_api::ExecServerClientConnectOptions; -use crate::client_api::ExecServerTransport; +use crate::client_api::ExecServerTransportParams; use crate::client_api::HttpClient; use crate::client_api::RemoteExecServerConnectArgs; use crate::client_api::StdioExecServerConnectArgs; @@ -195,14 +195,14 @@ pub struct ExecServerClient { #[derive(Clone)] pub(crate) struct LazyRemoteExecServerClient { - transport: ExecServerTransport, + transport_params: ExecServerTransportParams, client: Arc>, } impl LazyRemoteExecServerClient { - pub(crate) fn new(transport: ExecServerTransport) -> Self { + pub(crate) fn new(transport_params: ExecServerTransportParams) -> Self { Self { - transport, + transport_params, client: Arc::new(OnceCell::new()), } } @@ -212,8 +212,8 @@ impl LazyRemoteExecServerClient { // TODO: Add reconnect/disconnect handling here instead of reusing // the first successfully initialized connection forever. .get_or_try_init(|| { - let transport = self.transport.clone(); - async move { ExecServerClient::connect_for_transport(transport).await } + let transport_params = self.transport_params.clone(); + async move { ExecServerClient::connect_for_transport(transport_params).await } }) .await .cloned() @@ -898,8 +898,8 @@ mod tests { use super::ExecServerClient; use super::ExecServerClientConnectOptions; use crate::ProcessId; - use crate::StdioExecServerCommand; - use crate::StdioExecServerConnectArgs; + use crate::client_api::StdioExecServerCommand; + use crate::client_api::StdioExecServerConnectArgs; use crate::connection::JsonRpcConnection; use crate::process::ExecProcessEvent; use crate::protocol::EXEC_CLOSED_METHOD; diff --git a/codex-rs/exec-server/src/client_api.rs b/codex-rs/exec-server/src/client_api.rs index 8110ee24e56a..20520f002ef6 100644 --- a/codex-rs/exec-server/src/client_api.rs +++ b/codex-rs/exec-server/src/client_api.rs @@ -29,7 +29,7 @@ pub struct RemoteExecServerConnectArgs { /// Stdio connection arguments for a command-backed exec-server. #[derive(Debug, Clone, PartialEq, Eq)] -pub struct StdioExecServerConnectArgs { +pub(crate) struct StdioExecServerConnectArgs { pub command: StdioExecServerCommand, pub client_name: String, pub initialize_timeout: Duration, @@ -38,16 +38,16 @@ pub struct StdioExecServerConnectArgs { /// Structured process command used to start an exec-server over stdio. #[derive(Debug, Clone, PartialEq, Eq)] -pub struct StdioExecServerCommand { +pub(crate) struct StdioExecServerCommand { pub program: String, pub args: Vec, pub env: HashMap, pub cwd: Option, } -/// Transport used to connect to a remote exec-server environment. +/// Parameters used to connect to a remote exec-server environment. #[derive(Debug, Clone, PartialEq, Eq)] -pub enum ExecServerTransport { +pub(crate) enum ExecServerTransportParams { WebSocketUrl(String), StdioCommand(StdioExecServerCommand), } diff --git a/codex-rs/exec-server/src/client_transport.rs b/codex-rs/exec-server/src/client_transport.rs index 2f8f1b8a0581..560630e3db79 100644 --- a/codex-rs/exec-server/src/client_transport.rs +++ b/codex-rs/exec-server/src/client_transport.rs @@ -22,10 +22,10 @@ const ENVIRONMENT_INITIALIZE_TIMEOUT: Duration = Duration::from_secs(5); impl ExecServerClient { pub(crate) async fn connect_for_transport( - transport: crate::client_api::ExecServerTransport, + transport_params: crate::client_api::ExecServerTransportParams, ) -> Result { - match transport { - crate::client_api::ExecServerTransport::WebSocketUrl(websocket_url) => { + match transport_params { + crate::client_api::ExecServerTransportParams::WebSocketUrl(websocket_url) => { Self::connect_websocket(RemoteExecServerConnectArgs { websocket_url, client_name: ENVIRONMENT_CLIENT_NAME.to_string(), @@ -35,7 +35,7 @@ impl ExecServerClient { }) .await } - crate::client_api::ExecServerTransport::StdioCommand(command) => { + crate::client_api::ExecServerTransportParams::StdioCommand(command) => { Self::connect_stdio_command(StdioExecServerConnectArgs { command, client_name: ENVIRONMENT_CLIENT_NAME.to_string(), @@ -73,7 +73,7 @@ impl ExecServerClient { .await } - pub async fn connect_stdio_command( + pub(crate) async fn connect_stdio_command( args: StdioExecServerConnectArgs, ) -> Result { let mut child = stdio_command_process(&args.command) diff --git a/codex-rs/exec-server/src/connection.rs b/codex-rs/exec-server/src/connection.rs index 2dbda9e60403..05f8c94f2fb3 100644 --- a/codex-rs/exec-server/src/connection.rs +++ b/codex-rs/exec-server/src/connection.rs @@ -31,12 +31,12 @@ enum JsonRpcTransport { impl JsonRpcTransport { fn from_child_process(child_process: Child) -> Self { - Self::Stdio(StdioTransport { child_process }) + Self::Stdio(StdioTransport { + child_process: Some(child_process), + }) } -} -impl Drop for JsonRpcTransport { - fn drop(&mut self) { + fn shutdown(&mut self) { match self { Self::Plain => {} Self::Stdio(transport) => transport.shutdown(), @@ -45,14 +45,30 @@ impl Drop for JsonRpcTransport { } struct StdioTransport { - child_process: Child, + child_process: Option, } impl StdioTransport { fn shutdown(&mut self) { - if let Err(err) = self.child_process.start_kill() { + let Some(mut child_process) = self.child_process.take() else { + return; + }; + + if let Err(err) = child_process.start_kill() { debug!("failed to terminate exec-server stdio child: {err}"); } + match tokio::runtime::Handle::try_current() { + Ok(handle) => { + handle.spawn(async move { + if let Err(err) = child_process.wait().await { + debug!("failed to wait for exec-server stdio child: {err}"); + } + }); + } + Err(err) => { + debug!("failed to wait for exec-server stdio child without a Tokio runtime: {err}"); + } + } } } @@ -65,10 +81,15 @@ struct JsonRpcConnectionRuntime { pub(crate) struct JsonRpcConnection { runtime: Option, - #[allow(dead_code)] transport: JsonRpcTransport, } +impl Drop for JsonRpcConnection { + fn drop(&mut self) { + self.transport.shutdown(); + } +} + impl JsonRpcConnection { pub(crate) fn from_stdio(reader: R, writer: W, connection_label: String) -> Self where @@ -326,7 +347,7 @@ impl JsonRpcConnection { } pub(crate) fn into_parts( - self, + mut self, ) -> ( mpsc::Sender, mpsc::Receiver, @@ -340,6 +361,7 @@ impl JsonRpcConnection { task_handles, } = self .runtime + .take() .expect("JSON-RPC connection parts already taken"); (outgoing_tx, incoming_rx, disconnected_rx, task_handles) } diff --git a/codex-rs/exec-server/src/environment.rs b/codex-rs/exec-server/src/environment.rs index 3764b29fe567..395cac9e6333 100644 --- a/codex-rs/exec-server/src/environment.rs +++ b/codex-rs/exec-server/src/environment.rs @@ -7,7 +7,7 @@ use crate::ExecutorFileSystem; use crate::HttpClient; use crate::client::LazyRemoteExecServerClient; use crate::client::http_client::ReqwestHttpClient; -use crate::client_api::ExecServerTransport; +use crate::client_api::ExecServerTransportParams; use crate::environment_provider::DefaultEnvironmentProvider; use crate::environment_provider::EnvironmentProvider; use crate::environment_provider::normalize_exec_server_url; @@ -275,7 +275,7 @@ impl Environment { exec_server_url: String, local_runtime_paths: Option, ) -> Self { - let client = LazyRemoteExecServerClient::new(ExecServerTransport::WebSocketUrl( + let client = LazyRemoteExecServerClient::new(ExecServerTransportParams::WebSocketUrl( exec_server_url.clone(), )); let exec_backend: Arc = Arc::new(RemoteProcess::new(client.clone())); diff --git a/codex-rs/exec-server/src/lib.rs b/codex-rs/exec-server/src/lib.rs index 29897c528532..b36ab39d0105 100644 --- a/codex-rs/exec-server/src/lib.rs +++ b/codex-rs/exec-server/src/lib.rs @@ -25,11 +25,8 @@ pub use client::ExecServerError; pub use client::http_client::HttpResponseBodyStream; pub use client::http_client::ReqwestHttpClient; pub use client_api::ExecServerClientConnectOptions; -pub use client_api::ExecServerTransport; pub use client_api::HttpClient; pub use client_api::RemoteExecServerConnectArgs; -pub use client_api::StdioExecServerCommand; -pub use client_api::StdioExecServerConnectArgs; pub use codex_file_system::CopyOptions; pub use codex_file_system::CreateDirectoryOptions; pub use codex_file_system::ExecutorFileSystem; diff --git a/codex-rs/exec-server/src/rpc.rs b/codex-rs/exec-server/src/rpc.rs index a9c77d549a3b..8eb94450771b 100644 --- a/codex-rs/exec-server/src/rpc.rs +++ b/codex-rs/exec-server/src/rpc.rs @@ -223,10 +223,10 @@ where pub(crate) struct RpcClient { write_tx: mpsc::Sender, pending: Arc>>, - // Shared transport state from `JsonRpcConnection`. Calls use this to fail - // immediately when the socket closes, even if no JSON-RPC error response - // can be delivered for their request id. - disconnected_rx: watch::Receiver, + // This flips when either the underlying transport closes or the RPC reader + // exits, so new calls fail quickly after the connection is no longer usable. + closed_rx: watch::Receiver, + transport_disconnected_rx: watch::Receiver, next_request_id: AtomicI64, transport_tasks: Vec>, reader_task: JoinHandle<()>, @@ -236,38 +236,36 @@ impl RpcClient { pub(crate) fn new( connection: &mut JsonRpcConnection, ) -> (Self, mpsc::Receiver) { - let (write_tx, mut incoming_rx, disconnected_rx, transport_tasks) = + let (write_tx, mut incoming_rx, transport_disconnected_rx, transport_tasks) = connection.take_client_runtime(); let pending = Arc::new(Mutex::new(HashMap::::new())); let (event_tx, event_rx) = mpsc::channel(128); + let (closed_tx, closed_rx) = watch::channel(*transport_disconnected_rx.borrow()); let pending_for_reader = Arc::clone(&pending); let reader_task = tokio::spawn(async move { - while let Some(event) = incoming_rx.recv().await { + let reason = loop { + let Some(event) = incoming_rx.recv().await else { + break None; + }; + match event { JsonRpcConnectionEvent::Message(message) => { if let Err(err) = handle_server_message(&pending_for_reader, &event_tx, message).await { - let _ = err; - break; + break Some(err); } } JsonRpcConnectionEvent::MalformedMessage { reason } => { - let _ = reason; - break; - } - JsonRpcConnectionEvent::Disconnected { reason } => { - let _ = event_tx.send(RpcClientEvent::Disconnected { reason }).await; - drain_pending(&pending_for_reader).await; - return; + break Some(reason); } + JsonRpcConnectionEvent::Disconnected { reason } => break reason, } - } + }; - let _ = event_tx - .send(RpcClientEvent::Disconnected { reason: None }) - .await; + let _ = closed_tx.send(true); + let _ = event_tx.send(RpcClientEvent::Disconnected { reason }).await; drain_pending(&pending_for_reader).await; }); @@ -275,7 +273,8 @@ impl RpcClient { Self { write_tx, pending, - disconnected_rx, + closed_rx, + transport_disconnected_rx, next_request_id: AtomicI64::new(1), transport_tasks, reader_task, @@ -290,6 +289,12 @@ impl RpcClient { params: &P, ) -> Result<(), serde_json::Error> { let params = serde_json::to_value(params)?; + if *self.closed_rx.borrow() || *self.transport_disconnected_rx.borrow() { + return Err(serde_json::Error::io(std::io::Error::new( + std::io::ErrorKind::BrokenPipe, + "JSON-RPC transport closed", + ))); + } self.write_tx .send(JSONRPCMessage::Notification(JSONRPCNotification { method: method.to_string(), @@ -316,7 +321,7 @@ impl RpcClient { // Registering the pending request and checking disconnect must be // atomic with the reader's drain_pending path. Otherwise a call // can sneak in after the drain and wait forever. - if *self.disconnected_rx.borrow() { + if *self.closed_rx.borrow() || *self.transport_disconnected_rx.borrow() { return Err(RpcCallError::Closed); } pending.insert(request_id.clone(), response_tx); @@ -518,7 +523,9 @@ mod tests { use std::time::Duration; use codex_app_server_protocol::JSONRPCMessage; + use codex_app_server_protocol::JSONRPCRequest; use codex_app_server_protocol::JSONRPCResponse; + use codex_app_server_protocol::RequestId; use pretty_assertions::assert_eq; use tokio::io::AsyncBufReadExt; use tokio::io::AsyncWriteExt; @@ -526,6 +533,7 @@ mod tests { use tokio::time::timeout; use super::RpcClient; + use super::RpcClientEvent; use crate::connection::JsonRpcConnection; async fn read_jsonrpc_line(lines: &mut tokio::io::Lines>) -> JSONRPCMessage @@ -629,4 +637,49 @@ mod tests { panic!("server task failed: {err}"); } } + + #[tokio::test] + async fn rpc_client_rejects_new_calls_after_reader_protocol_error() { + let (client_stdin, _server_reader) = tokio::io::duplex(4096); + let (mut server_writer, client_stdout) = tokio::io::duplex(4096); + let mut connection = + JsonRpcConnection::from_stdio(client_stdout, client_stdin, "test-rpc".to_string()); + let (client, mut events_rx) = RpcClient::new(&mut connection); + + write_jsonrpc_line( + &mut server_writer, + JSONRPCMessage::Request(JSONRPCRequest { + id: RequestId::Integer(1), + method: "server/request".to_string(), + params: None, + trace: None, + }), + ) + .await; + + let event = timeout(Duration::from_secs(1), events_rx.recv()) + .await + .expect("timed out waiting for disconnect event"); + match event { + Some(RpcClientEvent::Disconnected { reason }) => { + assert!( + reason + .as_deref() + .is_some_and(|reason| reason.contains("unexpected JSON-RPC request")), + "unexpected disconnect reason: {reason:?}" + ); + } + event => panic!("expected disconnect event, got {event:?}"), + } + + let result = timeout( + Duration::from_secs(1), + client.call::<_, serde_json::Value>("after-close", &serde_json::json!({})), + ) + .await + .expect("timed out waiting for closed call"); + + assert!(matches!(result, Err(super::RpcCallError::Closed))); + assert_eq!(client.pending_request_count().await, 0); + } } From fb93315b4b1cfbf41e1997b316f73d904db1e9b7 Mon Sep 17 00:00:00 2001 From: starr-openai Date: Tue, 5 May 2026 15:25:21 -0700 Subject: [PATCH 15/42] Fix exec-server transport CI failures Co-authored-by: Codex --- codex-rs/exec-server/src/client.rs | 93 +++++++++++++++++++- codex-rs/exec-server/src/client_api.rs | 1 - codex-rs/exec-server/src/client_transport.rs | 31 +++---- codex-rs/exec-server/src/connection.rs | 17 ++-- codex-rs/exec-server/src/rpc.rs | 49 +++++++++++ codex-rs/tui/src/resume_picker.rs | 3 - 6 files changed, 158 insertions(+), 36 deletions(-) diff --git a/codex-rs/exec-server/src/client.rs b/codex-rs/exec-server/src/client.rs index 8a1178c93d41..346b7d841b8c 100644 --- a/codex-rs/exec-server/src/client.rs +++ b/codex-rs/exec-server/src/client.rs @@ -154,8 +154,7 @@ pub(crate) struct Session { struct Inner { // Keep the underlying transport connection alive and drop it before the RPC // client starts tearing down its channel/task handles. - #[allow(dead_code)] - connection: JsonRpcConnection, + connection: Option, client: RpcClient, // The remote transport delivers one shared notification stream for every // process on the connection. Keep a local process_id -> session registry so @@ -185,6 +184,7 @@ struct Inner { impl Drop for Inner { fn drop(&mut self) { self.reader_task.abort(); + drop(self.connection.take()); } } @@ -466,7 +466,7 @@ impl ExecServerClient { }); Inner { - connection, + connection: Some(connection), client: rpc_client, sessions: ArcSwap::from_pointee(HashMap::new()), sessions_write_lock: Mutex::new(()), @@ -890,6 +890,7 @@ mod tests { use tokio::io::BufReader; use tokio::io::duplex; use tokio::sync::mpsc; + use tokio::sync::oneshot; use tokio::time::Duration; #[cfg(unix)] use tokio::time::sleep; @@ -1235,6 +1236,92 @@ mod tests { server.await.expect("server task should finish"); } + #[tokio::test] + async fn transport_disconnect_fails_sessions_and_rejects_new_sessions() { + let (client_stdin, server_reader) = duplex(1 << 20); + let (mut server_writer, client_stdout) = duplex(1 << 20); + let (disconnect_tx, disconnect_rx) = oneshot::channel(); + let server = tokio::spawn(async move { + let mut lines = BufReader::new(server_reader).lines(); + let initialize = read_jsonrpc_line(&mut lines).await; + let request = match initialize { + JSONRPCMessage::Request(request) if request.method == INITIALIZE_METHOD => request, + other => panic!("expected initialize request, got {other:?}"), + }; + write_jsonrpc_line( + &mut server_writer, + JSONRPCMessage::Response(JSONRPCResponse { + id: request.id, + result: serde_json::to_value(InitializeResponse { + session_id: "session-1".to_string(), + }) + .expect("initialize response should serialize"), + }), + ) + .await; + + let initialized = read_jsonrpc_line(&mut lines).await; + match initialized { + JSONRPCMessage::Notification(notification) + if notification.method == INITIALIZED_METHOD => {} + other => panic!("expected initialized notification, got {other:?}"), + } + + let _ = disconnect_rx.await; + drop(server_writer); + }); + + let client = ExecServerClient::connect( + JsonRpcConnection::from_stdio( + client_stdout, + client_stdin, + "test-exec-server-client".to_string(), + ), + ExecServerClientConnectOptions::default(), + ) + .await + .expect("client should connect"); + + let process_id = ProcessId::from("disconnect"); + let session = client + .register_session(&process_id) + .await + .expect("session should register"); + let mut events = session.subscribe_events(); + + disconnect_tx.send(()).expect("disconnect should signal"); + + let event = timeout(Duration::from_secs(1), events.recv()) + .await + .expect("session failure should not time out") + .expect("session event stream should stay open"); + let ExecProcessEvent::Failed(message) = event else { + panic!("expected session failure after disconnect, got {event:?}"); + }; + assert_eq!(message, "exec-server transport disconnected"); + + let response = session + .read( + /*after_seq*/ None, /*max_bytes*/ None, /*wait_ms*/ None, + ) + .await + .expect("disconnected session read should synthesize a response"); + assert_eq!( + response.failure.as_deref(), + Some("exec-server transport disconnected") + ); + assert!(response.closed); + + let new_session = client.register_session(&ProcessId::from("new")).await; + assert!(matches!( + new_session, + Err(super::ExecServerError::Disconnected(_)) + )); + + drop(client); + server.await.expect("server task should finish"); + } + #[tokio::test] async fn wake_notifications_do_not_block_other_sessions() { let (client_stdin, server_reader) = duplex(1 << 20); diff --git a/codex-rs/exec-server/src/client_api.rs b/codex-rs/exec-server/src/client_api.rs index 20520f002ef6..9320efac304b 100644 --- a/codex-rs/exec-server/src/client_api.rs +++ b/codex-rs/exec-server/src/client_api.rs @@ -49,7 +49,6 @@ pub(crate) struct StdioExecServerCommand { #[derive(Debug, Clone, PartialEq, Eq)] pub(crate) enum ExecServerTransportParams { WebSocketUrl(String), - StdioCommand(StdioExecServerCommand), } /// Sends HTTP requests through a runtime-selected transport. diff --git a/codex-rs/exec-server/src/client_transport.rs b/codex-rs/exec-server/src/client_transport.rs index 560630e3db79..6f4492853cb9 100644 --- a/codex-rs/exec-server/src/client_transport.rs +++ b/codex-rs/exec-server/src/client_transport.rs @@ -24,27 +24,16 @@ impl ExecServerClient { pub(crate) async fn connect_for_transport( transport_params: crate::client_api::ExecServerTransportParams, ) -> Result { - match transport_params { - crate::client_api::ExecServerTransportParams::WebSocketUrl(websocket_url) => { - Self::connect_websocket(RemoteExecServerConnectArgs { - websocket_url, - client_name: ENVIRONMENT_CLIENT_NAME.to_string(), - connect_timeout: ENVIRONMENT_CONNECT_TIMEOUT, - initialize_timeout: ENVIRONMENT_INITIALIZE_TIMEOUT, - resume_session_id: None, - }) - .await - } - crate::client_api::ExecServerTransportParams::StdioCommand(command) => { - Self::connect_stdio_command(StdioExecServerConnectArgs { - command, - client_name: ENVIRONMENT_CLIENT_NAME.to_string(), - initialize_timeout: ENVIRONMENT_INITIALIZE_TIMEOUT, - resume_session_id: None, - }) - .await - } - } + let crate::client_api::ExecServerTransportParams::WebSocketUrl(websocket_url) = + transport_params; + Self::connect_websocket(RemoteExecServerConnectArgs { + websocket_url, + client_name: ENVIRONMENT_CLIENT_NAME.to_string(), + connect_timeout: ENVIRONMENT_CONNECT_TIMEOUT, + initialize_timeout: ENVIRONMENT_INITIALIZE_TIMEOUT, + resume_session_id: None, + }) + .await } pub async fn connect_websocket( diff --git a/codex-rs/exec-server/src/connection.rs b/codex-rs/exec-server/src/connection.rs index 05f8c94f2fb3..367d83f15da4 100644 --- a/codex-rs/exec-server/src/connection.rs +++ b/codex-rs/exec-server/src/connection.rs @@ -334,10 +334,7 @@ impl JsonRpcConnection { incoming_rx, disconnected_rx, task_handles, - } = self - .runtime - .take() - .expect("JSON-RPC client runtime already taken"); + } = self.take_runtime("JSON-RPC client runtime already taken"); (outgoing_tx, incoming_rx, disconnected_rx, task_handles) } @@ -359,12 +356,16 @@ impl JsonRpcConnection { incoming_rx, disconnected_rx, task_handles, - } = self - .runtime - .take() - .expect("JSON-RPC connection parts already taken"); + } = self.take_runtime("JSON-RPC connection parts already taken"); (outgoing_tx, incoming_rx, disconnected_rx, task_handles) } + + fn take_runtime(&mut self, message: &'static str) -> JsonRpcConnectionRuntime { + match self.runtime.take() { + Some(runtime) => runtime, + None => panic!("{message}"), + } + } } async fn send_disconnected( diff --git a/codex-rs/exec-server/src/rpc.rs b/codex-rs/exec-server/src/rpc.rs index 8eb94450771b..82948b920c09 100644 --- a/codex-rs/exec-server/src/rpc.rs +++ b/codex-rs/exec-server/src/rpc.rs @@ -682,4 +682,53 @@ mod tests { assert!(matches!(result, Err(super::RpcCallError::Closed))); assert_eq!(client.pending_request_count().await, 0); } + + #[tokio::test] + async fn rpc_client_drains_pending_call_on_transport_eof() { + let (client_stdin, server_reader) = tokio::io::duplex(4096); + let (server_writer, client_stdout) = tokio::io::duplex(4096); + let mut connection = + JsonRpcConnection::from_stdio(client_stdout, client_stdin, "test-rpc".to_string()); + let (client, mut events_rx) = RpcClient::new(&mut connection); + + let server = tokio::spawn(async move { + let mut lines = BufReader::new(server_reader).lines(); + let request = read_jsonrpc_line(&mut lines).await; + match request { + JSONRPCMessage::Request(request) if request.method == "will-close" => {} + other => panic!("expected will-close request, got {other:?}"), + } + drop(server_writer); + }); + + let result = timeout( + Duration::from_secs(1), + client.call::<_, serde_json::Value>("will-close", &serde_json::json!({})), + ) + .await + .expect("timed out waiting for closed call"); + assert!(matches!(result, Err(super::RpcCallError::Closed))); + + let event = timeout(Duration::from_secs(1), events_rx.recv()) + .await + .expect("timed out waiting for disconnect event"); + assert!(matches!( + event, + Some(RpcClientEvent::Disconnected { reason: None }) + )); + assert_eq!(client.pending_request_count().await, 0); + + let result = timeout( + Duration::from_secs(1), + client.call::<_, serde_json::Value>("after-close", &serde_json::json!({})), + ) + .await + .expect("timed out waiting for fast closed call"); + assert!(matches!(result, Err(super::RpcCallError::Closed))); + + let notify = client.notify("after-close", &serde_json::json!({})).await; + assert!(notify.is_err()); + + server.await.expect("server task should finish"); + } } diff --git a/codex-rs/tui/src/resume_picker.rs b/codex-rs/tui/src/resume_picker.rs index e06ccfd118b9..06ad0a61a7ce 100644 --- a/codex-rs/tui/src/resume_picker.rs +++ b/codex-rs/tui/src/resume_picker.rs @@ -5753,7 +5753,6 @@ session_picker_view = "dense" text: String::from("1. Do the thing"), }, ], - items_view: codex_app_server_protocol::TurnItemsView::Full, status: codex_app_server_protocol::TurnStatus::Completed, error: None, started_at: None, @@ -5805,7 +5804,6 @@ session_picker_view = "dense" summary: Vec::new(), content: vec![String::from("private raw chain of thought")], }], - items_view: codex_app_server_protocol::TurnItemsView::Full, status: codex_app_server_protocol::TurnStatus::Completed, error: None, started_at: None, @@ -5861,7 +5859,6 @@ session_picker_view = "dense" summary: vec![String::from("public summary")], content: vec![String::from("raw reasoning content")], }], - items_view: codex_app_server_protocol::TurnItemsView::Full, status: codex_app_server_protocol::TurnStatus::Completed, error: None, started_at: None, From bc34e376f75029751d2e7a193c8778ffe825166e Mon Sep 17 00:00:00 2001 From: starr-openai Date: Tue, 5 May 2026 15:37:07 -0700 Subject: [PATCH 16/42] Simplify exec-server disconnect plumbing Keep transport shutdown responsible for stdio child cleanup, and remove the separate disconnect watch channel from the JSON-RPC connection/runtime. The RPC client now keeps a single closed flag for rejecting calls after the ordered reader exits. Co-authored-by: Codex --- codex-rs/exec-server/src/connection.rs | 44 +++----------------- codex-rs/exec-server/src/rpc.rs | 24 +++++------ codex-rs/exec-server/src/server/processor.rs | 21 ++-------- 3 files changed, 19 insertions(+), 70 deletions(-) diff --git a/codex-rs/exec-server/src/connection.rs b/codex-rs/exec-server/src/connection.rs index 367d83f15da4..bd9c9ccce304 100644 --- a/codex-rs/exec-server/src/connection.rs +++ b/codex-rs/exec-server/src/connection.rs @@ -5,7 +5,6 @@ use tokio::io::AsyncRead; use tokio::io::AsyncWrite; use tokio::process::Child; use tokio::sync::mpsc; -use tokio::sync::watch; use tokio_tungstenite::WebSocketStream; use tokio_tungstenite::tungstenite::Message; use tracing::debug; @@ -75,7 +74,6 @@ impl StdioTransport { struct JsonRpcConnectionRuntime { outgoing_tx: mpsc::Sender, incoming_rx: mpsc::Receiver, - disconnected_rx: watch::Receiver, task_handles: Vec>, } @@ -98,11 +96,9 @@ impl JsonRpcConnection { { let (outgoing_tx, mut outgoing_rx) = mpsc::channel(CHANNEL_CAPACITY); let (incoming_tx, incoming_rx) = mpsc::channel(CHANNEL_CAPACITY); - let (disconnected_tx, disconnected_rx) = watch::channel(false); let reader_label = connection_label.clone(); let incoming_tx_for_reader = incoming_tx.clone(); - let disconnected_tx_for_reader = disconnected_tx.clone(); let reader_task = tokio::spawn(async move { let mut lines = BufReader::new(reader).lines(); loop { @@ -133,18 +129,12 @@ impl JsonRpcConnection { } } Ok(None) => { - send_disconnected( - &incoming_tx_for_reader, - &disconnected_tx_for_reader, - /*reason*/ None, - ) - .await; + send_disconnected(&incoming_tx_for_reader, /*reason*/ None).await; break; } Err(err) => { send_disconnected( &incoming_tx_for_reader, - &disconnected_tx_for_reader, Some(format!( "failed to read JSON-RPC message from {reader_label}: {err}" )), @@ -162,7 +152,6 @@ impl JsonRpcConnection { if let Err(err) = write_jsonrpc_line_message(&mut writer, &message).await { send_disconnected( &incoming_tx, - &disconnected_tx, Some(format!( "failed to write JSON-RPC message to {connection_label}: {err}" )), @@ -177,7 +166,6 @@ impl JsonRpcConnection { runtime: Some(JsonRpcConnectionRuntime { outgoing_tx, incoming_rx, - disconnected_rx, task_handles: vec![reader_task, writer_task], }), transport: JsonRpcTransport::Plain, @@ -190,12 +178,10 @@ impl JsonRpcConnection { { let (outgoing_tx, mut outgoing_rx) = mpsc::channel(CHANNEL_CAPACITY); let (incoming_tx, incoming_rx) = mpsc::channel(CHANNEL_CAPACITY); - let (disconnected_tx, disconnected_rx) = watch::channel(false); let (mut websocket_writer, mut websocket_reader) = stream.split(); let reader_label = connection_label.clone(); let incoming_tx_for_reader = incoming_tx.clone(); - let disconnected_tx_for_reader = disconnected_tx.clone(); let reader_task = tokio::spawn(async move { loop { match websocket_reader.next().await { @@ -244,12 +230,7 @@ impl JsonRpcConnection { } } Some(Ok(Message::Close(_))) => { - send_disconnected( - &incoming_tx_for_reader, - &disconnected_tx_for_reader, - /*reason*/ None, - ) - .await; + send_disconnected(&incoming_tx_for_reader, /*reason*/ None).await; break; } Some(Ok(Message::Ping(_))) | Some(Ok(Message::Pong(_))) => {} @@ -257,7 +238,6 @@ impl JsonRpcConnection { Some(Err(err)) => { send_disconnected( &incoming_tx_for_reader, - &disconnected_tx_for_reader, Some(format!( "failed to read websocket JSON-RPC message from {reader_label}: {err}" )), @@ -266,12 +246,7 @@ impl JsonRpcConnection { break; } None => { - send_disconnected( - &incoming_tx_for_reader, - &disconnected_tx_for_reader, - /*reason*/ None, - ) - .await; + send_disconnected(&incoming_tx_for_reader, /*reason*/ None).await; break; } } @@ -286,7 +261,6 @@ impl JsonRpcConnection { { send_disconnected( &incoming_tx, - &disconnected_tx, Some(format!( "failed to write websocket JSON-RPC message to {connection_label}: {err}" )), @@ -298,7 +272,6 @@ impl JsonRpcConnection { Err(err) => { send_disconnected( &incoming_tx, - &disconnected_tx, Some(format!( "failed to serialize JSON-RPC message for {connection_label}: {err}" )), @@ -314,7 +287,6 @@ impl JsonRpcConnection { runtime: Some(JsonRpcConnectionRuntime { outgoing_tx, incoming_rx, - disconnected_rx, task_handles: vec![reader_task, writer_task], }), transport: JsonRpcTransport::Plain, @@ -326,16 +298,14 @@ impl JsonRpcConnection { ) -> ( mpsc::Sender, mpsc::Receiver, - watch::Receiver, Vec>, ) { let JsonRpcConnectionRuntime { outgoing_tx, incoming_rx, - disconnected_rx, task_handles, } = self.take_runtime("JSON-RPC client runtime already taken"); - (outgoing_tx, incoming_rx, disconnected_rx, task_handles) + (outgoing_tx, incoming_rx, task_handles) } pub(crate) fn with_child_process(mut self, child_process: Child) -> Self { @@ -348,16 +318,14 @@ impl JsonRpcConnection { ) -> ( mpsc::Sender, mpsc::Receiver, - watch::Receiver, Vec>, ) { let JsonRpcConnectionRuntime { outgoing_tx, incoming_rx, - disconnected_rx, task_handles, } = self.take_runtime("JSON-RPC connection parts already taken"); - (outgoing_tx, incoming_rx, disconnected_rx, task_handles) + (outgoing_tx, incoming_rx, task_handles) } fn take_runtime(&mut self, message: &'static str) -> JsonRpcConnectionRuntime { @@ -370,10 +338,8 @@ impl JsonRpcConnection { async fn send_disconnected( incoming_tx: &mpsc::Sender, - disconnected_tx: &watch::Sender, reason: Option, ) { - let _ = disconnected_tx.send(true); let _ = incoming_tx .send(JsonRpcConnectionEvent::Disconnected { reason }) .await; diff --git a/codex-rs/exec-server/src/rpc.rs b/codex-rs/exec-server/src/rpc.rs index 82948b920c09..65cc363aa734 100644 --- a/codex-rs/exec-server/src/rpc.rs +++ b/codex-rs/exec-server/src/rpc.rs @@ -2,6 +2,7 @@ use std::collections::HashMap; use std::future::Future; use std::pin::Pin; use std::sync::Arc; +use std::sync::atomic::AtomicBool; use std::sync::atomic::AtomicI64; use std::sync::atomic::Ordering; @@ -18,7 +19,6 @@ use serde_json::Value; use tokio::sync::Mutex; use tokio::sync::mpsc; use tokio::sync::oneshot; -use tokio::sync::watch; use tokio::task::JoinHandle; use crate::connection::JsonRpcConnection; @@ -223,10 +223,9 @@ where pub(crate) struct RpcClient { write_tx: mpsc::Sender, pending: Arc>>, - // This flips when either the underlying transport closes or the RPC reader - // exits, so new calls fail quickly after the connection is no longer usable. - closed_rx: watch::Receiver, - transport_disconnected_rx: watch::Receiver, + // This flips before the ordered RPC reader drains pending requests, so new + // calls fail instead of registering work that can never complete. + closed: Arc, next_request_id: AtomicI64, transport_tasks: Vec>, reader_task: JoinHandle<()>, @@ -236,13 +235,13 @@ impl RpcClient { pub(crate) fn new( connection: &mut JsonRpcConnection, ) -> (Self, mpsc::Receiver) { - let (write_tx, mut incoming_rx, transport_disconnected_rx, transport_tasks) = - connection.take_client_runtime(); + let (write_tx, mut incoming_rx, transport_tasks) = connection.take_client_runtime(); let pending = Arc::new(Mutex::new(HashMap::::new())); let (event_tx, event_rx) = mpsc::channel(128); - let (closed_tx, closed_rx) = watch::channel(*transport_disconnected_rx.borrow()); + let closed = Arc::new(AtomicBool::new(false)); let pending_for_reader = Arc::clone(&pending); + let closed_for_reader = Arc::clone(&closed); let reader_task = tokio::spawn(async move { let reason = loop { let Some(event) = incoming_rx.recv().await else { @@ -264,7 +263,7 @@ impl RpcClient { } }; - let _ = closed_tx.send(true); + closed_for_reader.store(true, Ordering::SeqCst); let _ = event_tx.send(RpcClientEvent::Disconnected { reason }).await; drain_pending(&pending_for_reader).await; }); @@ -273,8 +272,7 @@ impl RpcClient { Self { write_tx, pending, - closed_rx, - transport_disconnected_rx, + closed, next_request_id: AtomicI64::new(1), transport_tasks, reader_task, @@ -289,7 +287,7 @@ impl RpcClient { params: &P, ) -> Result<(), serde_json::Error> { let params = serde_json::to_value(params)?; - if *self.closed_rx.borrow() || *self.transport_disconnected_rx.borrow() { + if self.closed.load(Ordering::SeqCst) { return Err(serde_json::Error::io(std::io::Error::new( std::io::ErrorKind::BrokenPipe, "JSON-RPC transport closed", @@ -321,7 +319,7 @@ impl RpcClient { // Registering the pending request and checking disconnect must be // atomic with the reader's drain_pending path. Otherwise a call // can sneak in after the drain and wait forever. - if *self.closed_rx.borrow() || *self.transport_disconnected_rx.borrow() { + if self.closed.load(Ordering::SeqCst) { return Err(RpcCallError::Closed); } pending.insert(request_id.clone(), response_tx); diff --git a/codex-rs/exec-server/src/server/processor.rs b/codex-rs/exec-server/src/server/processor.rs index dc1a9b9ffe74..132de639212f 100644 --- a/codex-rs/exec-server/src/server/processor.rs +++ b/codex-rs/exec-server/src/server/processor.rs @@ -47,8 +47,7 @@ async fn run_connection( runtime_paths: ExecServerRuntimePaths, ) { let router = Arc::new(build_router()); - let (json_outgoing_tx, mut incoming_rx, mut disconnected_rx, connection_tasks) = - connection.into_parts(); + let (json_outgoing_tx, mut incoming_rx, connection_tasks) = connection.into_parts(); let (outgoing_tx, mut outgoing_rx) = mpsc::channel::(CHANNEL_CAPACITY); let notifications = RpcNotificationSender::new(outgoing_tx.clone()); @@ -96,13 +95,7 @@ async fn run_connection( JsonRpcConnectionEvent::Message(message) => match message { codex_app_server_protocol::JSONRPCMessage::Request(request) => { if let Some(route) = router.request_route(request.method.as_str()) { - let message = tokio::select! { - message = route(Arc::clone(&handler), request) => message, - _ = disconnected_rx.changed() => { - debug!("exec-server transport disconnected while handling request"); - break; - } - }; + let message = route(Arc::clone(&handler), request).await; if let Some(message) = message && outgoing_tx.send(message).await.is_err() { @@ -131,15 +124,7 @@ async fn run_connection( ); break; }; - let result = tokio::select! { - result = route(Arc::clone(&handler), notification) => result, - _ = disconnected_rx.changed() => { - debug!( - "exec-server transport disconnected while handling notification" - ); - break; - } - }; + let result = route(Arc::clone(&handler), notification).await; if let Err(err) = result { warn!("closing exec-server connection after protocol error: {err}"); break; From 39634adbbe664f9c91c23f2effc41d6e4533c325 Mon Sep 17 00:00:00 2001 From: starr-openai Date: Tue, 5 May 2026 15:47:17 -0700 Subject: [PATCH 17/42] Remove server disconnect race test The stdio transport no longer adds a processor-side disconnect side channel, so drop the test that asserted that removed behavior. Client cleanup is covered at the RPC/client transport boundary instead. Co-authored-by: Codex --- codex-rs/exec-server/src/server/processor.rs | 238 ------------------- 1 file changed, 238 deletions(-) diff --git a/codex-rs/exec-server/src/server/processor.rs b/codex-rs/exec-server/src/server/processor.rs index 132de639212f..9aa2b715eed5 100644 --- a/codex-rs/exec-server/src/server/processor.rs +++ b/codex-rs/exec-server/src/server/processor.rs @@ -163,241 +163,3 @@ async fn run_connection( } let _ = outbound_task.await; } - -#[cfg(test)] -mod tests { - use std::collections::HashMap; - use std::sync::Arc; - use std::time::Duration; - - use codex_app_server_protocol::JSONRPCMessage; - use codex_app_server_protocol::JSONRPCNotification; - use codex_app_server_protocol::JSONRPCRequest; - use codex_app_server_protocol::JSONRPCResponse; - use codex_app_server_protocol::RequestId; - use serde::Serialize; - use serde::de::DeserializeOwned; - use tokio::io::AsyncBufReadExt; - use tokio::io::AsyncWriteExt; - use tokio::io::BufReader; - use tokio::io::DuplexStream; - use tokio::io::Lines; - use tokio::io::duplex; - use tokio::task::JoinHandle; - use tokio::time::timeout; - - use super::run_connection; - use crate::ExecServerRuntimePaths; - use crate::ProcessId; - use crate::connection::JsonRpcConnection; - use crate::protocol::EXEC_METHOD; - use crate::protocol::EXEC_READ_METHOD; - use crate::protocol::EXEC_TERMINATE_METHOD; - use crate::protocol::ExecParams; - use crate::protocol::ExecResponse; - use crate::protocol::INITIALIZE_METHOD; - use crate::protocol::INITIALIZED_METHOD; - use crate::protocol::InitializeParams; - use crate::protocol::InitializeResponse; - use crate::protocol::ReadParams; - use crate::protocol::TerminateParams; - use crate::protocol::TerminateResponse; - use crate::server::session_registry::SessionRegistry; - - #[tokio::test] - async fn transport_disconnect_detaches_session_during_in_flight_read() { - let registry = SessionRegistry::new(); - let (mut first_writer, mut first_lines, first_task) = - spawn_test_connection(Arc::clone(®istry), "first"); - - send_request( - &mut first_writer, - /*id*/ 1, - INITIALIZE_METHOD, - &InitializeParams { - client_name: "exec-server-test".to_string(), - resume_session_id: None, - }, - ) - .await; - let initialize_response: InitializeResponse = - read_response(&mut first_lines, /*expected_id*/ 1).await; - send_notification(&mut first_writer, INITIALIZED_METHOD, &()).await; - - let process_id = ProcessId::from("proc-long-poll"); - send_request( - &mut first_writer, - /*id*/ 2, - EXEC_METHOD, - &exec_params(process_id.clone()), - ) - .await; - let _: ExecResponse = read_response(&mut first_lines, /*expected_id*/ 2).await; - - send_request( - &mut first_writer, - /*id*/ 3, - EXEC_READ_METHOD, - &ReadParams { - process_id: process_id.clone(), - after_seq: None, - max_bytes: None, - wait_ms: Some(5_000), - }, - ) - .await; - drop(first_writer); - tokio::time::sleep(Duration::from_millis(25)).await; - - let (mut second_writer, mut second_lines, second_task) = - spawn_test_connection(Arc::clone(®istry), "second"); - send_request( - &mut second_writer, - /*id*/ 1, - INITIALIZE_METHOD, - &InitializeParams { - client_name: "exec-server-test".to_string(), - resume_session_id: Some(initialize_response.session_id.clone()), - }, - ) - .await; - let second_initialize_response = timeout( - Duration::from_secs(1), - read_response::(&mut second_lines, /*expected_id*/ 1), - ) - .await - .expect("resume initialize should not wait for the old read to finish"); - assert_eq!( - second_initialize_response.session_id, - initialize_response.session_id - ); - timeout(Duration::from_secs(1), first_task) - .await - .expect("first processor should exit") - .expect("first processor should join"); - send_notification(&mut second_writer, INITIALIZED_METHOD, &()).await; - - send_request( - &mut second_writer, - /*id*/ 2, - EXEC_TERMINATE_METHOD, - &TerminateParams { process_id }, - ) - .await; - let _: TerminateResponse = read_response(&mut second_lines, /*expected_id*/ 2).await; - - drop(second_writer); - drop(second_lines); - timeout(Duration::from_secs(1), second_task) - .await - .expect("second processor should exit") - .expect("second processor should join"); - } - - fn spawn_test_connection( - registry: Arc, - label: &str, - ) -> (DuplexStream, Lines>, JoinHandle<()>) { - let (client_writer, server_reader) = duplex(1 << 20); - let (server_writer, client_reader) = duplex(1 << 20); - let connection = - JsonRpcConnection::from_stdio(server_reader, server_writer, label.to_string()); - let task = tokio::spawn(run_connection(connection, registry, test_runtime_paths())); - (client_writer, BufReader::new(client_reader).lines(), task) - } - - fn test_runtime_paths() -> ExecServerRuntimePaths { - ExecServerRuntimePaths::new( - std::env::current_exe().expect("current exe"), - /*codex_linux_sandbox_exe*/ None, - ) - .expect("runtime paths") - } - - async fn send_request( - writer: &mut DuplexStream, - id: i64, - method: &str, - params: &P, - ) { - write_message( - writer, - &JSONRPCMessage::Request(JSONRPCRequest { - id: RequestId::Integer(id), - method: method.to_string(), - params: Some(serde_json::to_value(params).expect("serialize params")), - trace: None, - }), - ) - .await; - } - - async fn send_notification(writer: &mut DuplexStream, method: &str, params: &P) { - write_message( - writer, - &JSONRPCMessage::Notification(JSONRPCNotification { - method: method.to_string(), - params: Some(serde_json::to_value(params).expect("serialize params")), - }), - ) - .await; - } - - async fn write_message(writer: &mut DuplexStream, message: &JSONRPCMessage) { - let encoded = serde_json::to_vec(message).expect("serialize JSON-RPC message"); - writer.write_all(&encoded).await.expect("write request"); - writer.write_all(b"\n").await.expect("write newline"); - } - - async fn read_response( - lines: &mut Lines>, - expected_id: i64, - ) -> T { - let line = lines - .next_line() - .await - .expect("read response") - .expect("response line"); - match serde_json::from_str::(&line).expect("decode JSON-RPC response") { - JSONRPCMessage::Response(JSONRPCResponse { id, result }) => { - assert_eq!(id, RequestId::Integer(expected_id)); - serde_json::from_value(result).expect("decode response result") - } - JSONRPCMessage::Error(error) => panic!("unexpected JSON-RPC error: {error:?}"), - other => panic!("expected JSON-RPC response, got {other:?}"), - } - } - - fn exec_params(process_id: ProcessId) -> ExecParams { - let mut env = HashMap::new(); - if let Some(path) = std::env::var_os("PATH") { - env.insert("PATH".to_string(), path.to_string_lossy().into_owned()); - } - ExecParams { - process_id, - argv: sleep_then_print_argv(), - cwd: std::env::current_dir().expect("cwd"), - env_policy: None, - env, - tty: false, - pipe_stdin: false, - arg0: None, - } - } - - fn sleep_then_print_argv() -> Vec { - if cfg!(windows) { - vec![ - std::env::var("COMSPEC").unwrap_or_else(|_| "cmd.exe".to_string()), - "/C".to_string(), - "ping -n 3 127.0.0.1 >NUL && echo late".to_string(), - ] - } else { - vec![ - "/bin/sh".to_string(), - "-c".to_string(), - "sleep 1; printf late".to_string(), - ] - } - } -} From 2a1200e62b528e266eff16825a99cb2172b317a9 Mon Sep 17 00:00:00 2001 From: starr-openai Date: Tue, 5 May 2026 15:53:08 -0700 Subject: [PATCH 18/42] Simplify exec-server transport ownership Remove the Option wrapper used only to force connection drop order and call transport shutdown explicitly instead. Also drop dead-code allowances that are no longer needed. Co-authored-by: Codex --- codex-rs/exec-server/src/client.rs | 9 +++--- codex-rs/exec-server/src/client_api.rs | 1 + codex-rs/exec-server/src/client_transport.rs | 31 +++++++++++++------- codex-rs/exec-server/src/connection.rs | 27 +++++------------ codex-rs/exec-server/src/rpc.rs | 6 +--- codex-rs/exec-server/src/server/processor.rs | 4 +-- 6 files changed, 37 insertions(+), 41 deletions(-) diff --git a/codex-rs/exec-server/src/client.rs b/codex-rs/exec-server/src/client.rs index 346b7d841b8c..3fbb35cb90ec 100644 --- a/codex-rs/exec-server/src/client.rs +++ b/codex-rs/exec-server/src/client.rs @@ -152,9 +152,8 @@ pub(crate) struct Session { } struct Inner { - // Keep the underlying transport connection alive and drop it before the RPC - // client starts tearing down its channel/task handles. - connection: Option, + // Keep the underlying transport connection alive for the client lifetime. + connection: JsonRpcConnection, client: RpcClient, // The remote transport delivers one shared notification stream for every // process on the connection. Keep a local process_id -> session registry so @@ -184,7 +183,7 @@ struct Inner { impl Drop for Inner { fn drop(&mut self) { self.reader_task.abort(); - drop(self.connection.take()); + self.connection.shutdown(); } } @@ -466,7 +465,7 @@ impl ExecServerClient { }); Inner { - connection: Some(connection), + connection, client: rpc_client, sessions: ArcSwap::from_pointee(HashMap::new()), sessions_write_lock: Mutex::new(()), diff --git a/codex-rs/exec-server/src/client_api.rs b/codex-rs/exec-server/src/client_api.rs index 9320efac304b..20520f002ef6 100644 --- a/codex-rs/exec-server/src/client_api.rs +++ b/codex-rs/exec-server/src/client_api.rs @@ -49,6 +49,7 @@ pub(crate) struct StdioExecServerCommand { #[derive(Debug, Clone, PartialEq, Eq)] pub(crate) enum ExecServerTransportParams { WebSocketUrl(String), + StdioCommand(StdioExecServerCommand), } /// Sends HTTP requests through a runtime-selected transport. diff --git a/codex-rs/exec-server/src/client_transport.rs b/codex-rs/exec-server/src/client_transport.rs index 6f4492853cb9..560630e3db79 100644 --- a/codex-rs/exec-server/src/client_transport.rs +++ b/codex-rs/exec-server/src/client_transport.rs @@ -24,16 +24,27 @@ impl ExecServerClient { pub(crate) async fn connect_for_transport( transport_params: crate::client_api::ExecServerTransportParams, ) -> Result { - let crate::client_api::ExecServerTransportParams::WebSocketUrl(websocket_url) = - transport_params; - Self::connect_websocket(RemoteExecServerConnectArgs { - websocket_url, - client_name: ENVIRONMENT_CLIENT_NAME.to_string(), - connect_timeout: ENVIRONMENT_CONNECT_TIMEOUT, - initialize_timeout: ENVIRONMENT_INITIALIZE_TIMEOUT, - resume_session_id: None, - }) - .await + match transport_params { + crate::client_api::ExecServerTransportParams::WebSocketUrl(websocket_url) => { + Self::connect_websocket(RemoteExecServerConnectArgs { + websocket_url, + client_name: ENVIRONMENT_CLIENT_NAME.to_string(), + connect_timeout: ENVIRONMENT_CONNECT_TIMEOUT, + initialize_timeout: ENVIRONMENT_INITIALIZE_TIMEOUT, + resume_session_id: None, + }) + .await + } + crate::client_api::ExecServerTransportParams::StdioCommand(command) => { + Self::connect_stdio_command(StdioExecServerConnectArgs { + command, + client_name: ENVIRONMENT_CLIENT_NAME.to_string(), + initialize_timeout: ENVIRONMENT_INITIALIZE_TIMEOUT, + resume_session_id: None, + }) + .await + } + } } pub async fn connect_websocket( diff --git a/codex-rs/exec-server/src/connection.rs b/codex-rs/exec-server/src/connection.rs index bd9c9ccce304..ba1761240705 100644 --- a/codex-rs/exec-server/src/connection.rs +++ b/codex-rs/exec-server/src/connection.rs @@ -84,11 +84,15 @@ pub(crate) struct JsonRpcConnection { impl Drop for JsonRpcConnection { fn drop(&mut self) { - self.transport.shutdown(); + self.shutdown(); } } impl JsonRpcConnection { + pub(crate) fn shutdown(&mut self) { + self.transport.shutdown(); + } + pub(crate) fn from_stdio(reader: R, writer: W, connection_label: String) -> Self where R: AsyncRead + Unpin + Send + 'static, @@ -293,7 +297,7 @@ impl JsonRpcConnection { } } - pub(crate) fn take_client_runtime( + pub(crate) fn take_runtime( &mut self, ) -> ( mpsc::Sender, @@ -304,7 +308,7 @@ impl JsonRpcConnection { outgoing_tx, incoming_rx, task_handles, - } = self.take_runtime("JSON-RPC client runtime already taken"); + } = self.take_runtime_or_panic("JSON-RPC connection runtime already taken"); (outgoing_tx, incoming_rx, task_handles) } @@ -313,22 +317,7 @@ impl JsonRpcConnection { self } - pub(crate) fn into_parts( - mut self, - ) -> ( - mpsc::Sender, - mpsc::Receiver, - Vec>, - ) { - let JsonRpcConnectionRuntime { - outgoing_tx, - incoming_rx, - task_handles, - } = self.take_runtime("JSON-RPC connection parts already taken"); - (outgoing_tx, incoming_rx, task_handles) - } - - fn take_runtime(&mut self, message: &'static str) -> JsonRpcConnectionRuntime { + fn take_runtime_or_panic(&mut self, message: &'static str) -> JsonRpcConnectionRuntime { match self.runtime.take() { Some(runtime) => runtime, None => panic!("{message}"), diff --git a/codex-rs/exec-server/src/rpc.rs b/codex-rs/exec-server/src/rpc.rs index 65cc363aa734..8c7c17753670 100644 --- a/codex-rs/exec-server/src/rpc.rs +++ b/codex-rs/exec-server/src/rpc.rs @@ -58,11 +58,9 @@ pub(crate) enum RpcServerOutboundMessage { request_id: RequestId, error: JSONRPCErrorError, }, - #[allow(dead_code)] Notification(JSONRPCNotification), } -#[allow(dead_code)] #[derive(Clone)] pub(crate) struct RpcNotificationSender { outgoing_tx: mpsc::Sender, @@ -84,7 +82,6 @@ impl RpcNotificationSender { .map_err(|_| internal_error("RPC connection closed while sending response".into())) } - #[allow(dead_code)] pub(crate) async fn notify( &self, method: &str, @@ -235,7 +232,7 @@ impl RpcClient { pub(crate) fn new( connection: &mut JsonRpcConnection, ) -> (Self, mpsc::Receiver) { - let (write_tx, mut incoming_rx, transport_tasks) = connection.take_client_runtime(); + let (write_tx, mut incoming_rx, transport_tasks) = connection.take_runtime(); let pending = Arc::new(Mutex::new(HashMap::::new())); let (event_tx, event_rx) = mpsc::channel(128); let closed = Arc::new(AtomicBool::new(false)); @@ -363,7 +360,6 @@ impl RpcClient { } #[cfg(test)] - #[allow(dead_code)] pub(crate) async fn pending_request_count(&self) -> usize { self.pending.lock().await.len() } diff --git a/codex-rs/exec-server/src/server/processor.rs b/codex-rs/exec-server/src/server/processor.rs index 9aa2b715eed5..88a282e0d544 100644 --- a/codex-rs/exec-server/src/server/processor.rs +++ b/codex-rs/exec-server/src/server/processor.rs @@ -42,12 +42,12 @@ impl ConnectionProcessor { } async fn run_connection( - connection: JsonRpcConnection, + mut connection: JsonRpcConnection, session_registry: Arc, runtime_paths: ExecServerRuntimePaths, ) { let router = Arc::new(build_router()); - let (json_outgoing_tx, mut incoming_rx, connection_tasks) = connection.into_parts(); + let (json_outgoing_tx, mut incoming_rx, connection_tasks) = connection.take_runtime(); let (outgoing_tx, mut outgoing_rx) = mpsc::channel::(CHANNEL_CAPACITY); let notifications = RpcNotificationSender::new(outgoing_tx.clone()); From 94956cda6ea28ad0838c0c9cc75d138d9aa5b710 Mon Sep 17 00:00:00 2001 From: starr-openai Date: Tue, 5 May 2026 16:27:03 -0700 Subject: [PATCH 19/42] Restore exec-server processor ownership boundary Keep the server-side connection processor on the original by-value parts API, and move the compatibility needed for that shape into JsonRpcConnection. The client still borrows the connection mutably so it can keep transport ownership with ExecServerClient. Co-authored-by: Codex --- codex-rs/exec-server/src/connection.rs | 54 +++- codex-rs/exec-server/src/server/processor.rs | 261 ++++++++++++++++++- 2 files changed, 308 insertions(+), 7 deletions(-) diff --git a/codex-rs/exec-server/src/connection.rs b/codex-rs/exec-server/src/connection.rs index ba1761240705..e159dade288e 100644 --- a/codex-rs/exec-server/src/connection.rs +++ b/codex-rs/exec-server/src/connection.rs @@ -5,6 +5,7 @@ use tokio::io::AsyncRead; use tokio::io::AsyncWrite; use tokio::process::Child; use tokio::sync::mpsc; +use tokio::sync::watch; use tokio_tungstenite::WebSocketStream; use tokio_tungstenite::tungstenite::Message; use tracing::debug; @@ -74,6 +75,7 @@ impl StdioTransport { struct JsonRpcConnectionRuntime { outgoing_tx: mpsc::Sender, incoming_rx: mpsc::Receiver, + disconnected_rx: watch::Receiver, task_handles: Vec>, } @@ -100,9 +102,11 @@ impl JsonRpcConnection { { let (outgoing_tx, mut outgoing_rx) = mpsc::channel(CHANNEL_CAPACITY); let (incoming_tx, incoming_rx) = mpsc::channel(CHANNEL_CAPACITY); + let (disconnected_tx, disconnected_rx) = watch::channel(false); let reader_label = connection_label.clone(); let incoming_tx_for_reader = incoming_tx.clone(); + let disconnected_tx_for_reader = disconnected_tx.clone(); let reader_task = tokio::spawn(async move { let mut lines = BufReader::new(reader).lines(); loop { @@ -133,12 +137,18 @@ impl JsonRpcConnection { } } Ok(None) => { - send_disconnected(&incoming_tx_for_reader, /*reason*/ None).await; + send_disconnected( + &incoming_tx_for_reader, + &disconnected_tx_for_reader, + /*reason*/ None, + ) + .await; break; } Err(err) => { send_disconnected( &incoming_tx_for_reader, + &disconnected_tx_for_reader, Some(format!( "failed to read JSON-RPC message from {reader_label}: {err}" )), @@ -156,6 +166,7 @@ impl JsonRpcConnection { if let Err(err) = write_jsonrpc_line_message(&mut writer, &message).await { send_disconnected( &incoming_tx, + &disconnected_tx, Some(format!( "failed to write JSON-RPC message to {connection_label}: {err}" )), @@ -170,6 +181,7 @@ impl JsonRpcConnection { runtime: Some(JsonRpcConnectionRuntime { outgoing_tx, incoming_rx, + disconnected_rx, task_handles: vec![reader_task, writer_task], }), transport: JsonRpcTransport::Plain, @@ -182,10 +194,12 @@ impl JsonRpcConnection { { let (outgoing_tx, mut outgoing_rx) = mpsc::channel(CHANNEL_CAPACITY); let (incoming_tx, incoming_rx) = mpsc::channel(CHANNEL_CAPACITY); + let (disconnected_tx, disconnected_rx) = watch::channel(false); let (mut websocket_writer, mut websocket_reader) = stream.split(); let reader_label = connection_label.clone(); let incoming_tx_for_reader = incoming_tx.clone(); + let disconnected_tx_for_reader = disconnected_tx.clone(); let reader_task = tokio::spawn(async move { loop { match websocket_reader.next().await { @@ -234,7 +248,12 @@ impl JsonRpcConnection { } } Some(Ok(Message::Close(_))) => { - send_disconnected(&incoming_tx_for_reader, /*reason*/ None).await; + send_disconnected( + &incoming_tx_for_reader, + &disconnected_tx_for_reader, + /*reason*/ None, + ) + .await; break; } Some(Ok(Message::Ping(_))) | Some(Ok(Message::Pong(_))) => {} @@ -242,6 +261,7 @@ impl JsonRpcConnection { Some(Err(err)) => { send_disconnected( &incoming_tx_for_reader, + &disconnected_tx_for_reader, Some(format!( "failed to read websocket JSON-RPC message from {reader_label}: {err}" )), @@ -250,7 +270,12 @@ impl JsonRpcConnection { break; } None => { - send_disconnected(&incoming_tx_for_reader, /*reason*/ None).await; + send_disconnected( + &incoming_tx_for_reader, + &disconnected_tx_for_reader, + /*reason*/ None, + ) + .await; break; } } @@ -265,6 +290,7 @@ impl JsonRpcConnection { { send_disconnected( &incoming_tx, + &disconnected_tx, Some(format!( "failed to write websocket JSON-RPC message to {connection_label}: {err}" )), @@ -276,6 +302,7 @@ impl JsonRpcConnection { Err(err) => { send_disconnected( &incoming_tx, + &disconnected_tx, Some(format!( "failed to serialize JSON-RPC message for {connection_label}: {err}" )), @@ -291,6 +318,7 @@ impl JsonRpcConnection { runtime: Some(JsonRpcConnectionRuntime { outgoing_tx, incoming_rx, + disconnected_rx, task_handles: vec![reader_task, writer_task], }), transport: JsonRpcTransport::Plain, @@ -307,11 +335,29 @@ impl JsonRpcConnection { let JsonRpcConnectionRuntime { outgoing_tx, incoming_rx, + disconnected_rx: _, task_handles, } = self.take_runtime_or_panic("JSON-RPC connection runtime already taken"); (outgoing_tx, incoming_rx, task_handles) } + pub(crate) fn into_parts( + mut self, + ) -> ( + mpsc::Sender, + mpsc::Receiver, + watch::Receiver, + Vec>, + ) { + let JsonRpcConnectionRuntime { + outgoing_tx, + incoming_rx, + disconnected_rx, + task_handles, + } = self.take_runtime_or_panic("JSON-RPC connection runtime already taken"); + (outgoing_tx, incoming_rx, disconnected_rx, task_handles) + } + pub(crate) fn with_child_process(mut self, child_process: Child) -> Self { self.transport = JsonRpcTransport::from_child_process(child_process); self @@ -327,8 +373,10 @@ impl JsonRpcConnection { async fn send_disconnected( incoming_tx: &mpsc::Sender, + disconnected_tx: &watch::Sender, reason: Option, ) { + let _ = disconnected_tx.send(true); let _ = incoming_tx .send(JsonRpcConnectionEvent::Disconnected { reason }) .await; diff --git a/codex-rs/exec-server/src/server/processor.rs b/codex-rs/exec-server/src/server/processor.rs index 88a282e0d544..dc1a9b9ffe74 100644 --- a/codex-rs/exec-server/src/server/processor.rs +++ b/codex-rs/exec-server/src/server/processor.rs @@ -42,12 +42,13 @@ impl ConnectionProcessor { } async fn run_connection( - mut connection: JsonRpcConnection, + connection: JsonRpcConnection, session_registry: Arc, runtime_paths: ExecServerRuntimePaths, ) { let router = Arc::new(build_router()); - let (json_outgoing_tx, mut incoming_rx, connection_tasks) = connection.take_runtime(); + let (json_outgoing_tx, mut incoming_rx, mut disconnected_rx, connection_tasks) = + connection.into_parts(); let (outgoing_tx, mut outgoing_rx) = mpsc::channel::(CHANNEL_CAPACITY); let notifications = RpcNotificationSender::new(outgoing_tx.clone()); @@ -95,7 +96,13 @@ async fn run_connection( JsonRpcConnectionEvent::Message(message) => match message { codex_app_server_protocol::JSONRPCMessage::Request(request) => { if let Some(route) = router.request_route(request.method.as_str()) { - let message = route(Arc::clone(&handler), request).await; + let message = tokio::select! { + message = route(Arc::clone(&handler), request) => message, + _ = disconnected_rx.changed() => { + debug!("exec-server transport disconnected while handling request"); + break; + } + }; if let Some(message) = message && outgoing_tx.send(message).await.is_err() { @@ -124,7 +131,15 @@ async fn run_connection( ); break; }; - let result = route(Arc::clone(&handler), notification).await; + let result = tokio::select! { + result = route(Arc::clone(&handler), notification) => result, + _ = disconnected_rx.changed() => { + debug!( + "exec-server transport disconnected while handling notification" + ); + break; + } + }; if let Err(err) = result { warn!("closing exec-server connection after protocol error: {err}"); break; @@ -163,3 +178,241 @@ async fn run_connection( } let _ = outbound_task.await; } + +#[cfg(test)] +mod tests { + use std::collections::HashMap; + use std::sync::Arc; + use std::time::Duration; + + use codex_app_server_protocol::JSONRPCMessage; + use codex_app_server_protocol::JSONRPCNotification; + use codex_app_server_protocol::JSONRPCRequest; + use codex_app_server_protocol::JSONRPCResponse; + use codex_app_server_protocol::RequestId; + use serde::Serialize; + use serde::de::DeserializeOwned; + use tokio::io::AsyncBufReadExt; + use tokio::io::AsyncWriteExt; + use tokio::io::BufReader; + use tokio::io::DuplexStream; + use tokio::io::Lines; + use tokio::io::duplex; + use tokio::task::JoinHandle; + use tokio::time::timeout; + + use super::run_connection; + use crate::ExecServerRuntimePaths; + use crate::ProcessId; + use crate::connection::JsonRpcConnection; + use crate::protocol::EXEC_METHOD; + use crate::protocol::EXEC_READ_METHOD; + use crate::protocol::EXEC_TERMINATE_METHOD; + use crate::protocol::ExecParams; + use crate::protocol::ExecResponse; + use crate::protocol::INITIALIZE_METHOD; + use crate::protocol::INITIALIZED_METHOD; + use crate::protocol::InitializeParams; + use crate::protocol::InitializeResponse; + use crate::protocol::ReadParams; + use crate::protocol::TerminateParams; + use crate::protocol::TerminateResponse; + use crate::server::session_registry::SessionRegistry; + + #[tokio::test] + async fn transport_disconnect_detaches_session_during_in_flight_read() { + let registry = SessionRegistry::new(); + let (mut first_writer, mut first_lines, first_task) = + spawn_test_connection(Arc::clone(®istry), "first"); + + send_request( + &mut first_writer, + /*id*/ 1, + INITIALIZE_METHOD, + &InitializeParams { + client_name: "exec-server-test".to_string(), + resume_session_id: None, + }, + ) + .await; + let initialize_response: InitializeResponse = + read_response(&mut first_lines, /*expected_id*/ 1).await; + send_notification(&mut first_writer, INITIALIZED_METHOD, &()).await; + + let process_id = ProcessId::from("proc-long-poll"); + send_request( + &mut first_writer, + /*id*/ 2, + EXEC_METHOD, + &exec_params(process_id.clone()), + ) + .await; + let _: ExecResponse = read_response(&mut first_lines, /*expected_id*/ 2).await; + + send_request( + &mut first_writer, + /*id*/ 3, + EXEC_READ_METHOD, + &ReadParams { + process_id: process_id.clone(), + after_seq: None, + max_bytes: None, + wait_ms: Some(5_000), + }, + ) + .await; + drop(first_writer); + tokio::time::sleep(Duration::from_millis(25)).await; + + let (mut second_writer, mut second_lines, second_task) = + spawn_test_connection(Arc::clone(®istry), "second"); + send_request( + &mut second_writer, + /*id*/ 1, + INITIALIZE_METHOD, + &InitializeParams { + client_name: "exec-server-test".to_string(), + resume_session_id: Some(initialize_response.session_id.clone()), + }, + ) + .await; + let second_initialize_response = timeout( + Duration::from_secs(1), + read_response::(&mut second_lines, /*expected_id*/ 1), + ) + .await + .expect("resume initialize should not wait for the old read to finish"); + assert_eq!( + second_initialize_response.session_id, + initialize_response.session_id + ); + timeout(Duration::from_secs(1), first_task) + .await + .expect("first processor should exit") + .expect("first processor should join"); + send_notification(&mut second_writer, INITIALIZED_METHOD, &()).await; + + send_request( + &mut second_writer, + /*id*/ 2, + EXEC_TERMINATE_METHOD, + &TerminateParams { process_id }, + ) + .await; + let _: TerminateResponse = read_response(&mut second_lines, /*expected_id*/ 2).await; + + drop(second_writer); + drop(second_lines); + timeout(Duration::from_secs(1), second_task) + .await + .expect("second processor should exit") + .expect("second processor should join"); + } + + fn spawn_test_connection( + registry: Arc, + label: &str, + ) -> (DuplexStream, Lines>, JoinHandle<()>) { + let (client_writer, server_reader) = duplex(1 << 20); + let (server_writer, client_reader) = duplex(1 << 20); + let connection = + JsonRpcConnection::from_stdio(server_reader, server_writer, label.to_string()); + let task = tokio::spawn(run_connection(connection, registry, test_runtime_paths())); + (client_writer, BufReader::new(client_reader).lines(), task) + } + + fn test_runtime_paths() -> ExecServerRuntimePaths { + ExecServerRuntimePaths::new( + std::env::current_exe().expect("current exe"), + /*codex_linux_sandbox_exe*/ None, + ) + .expect("runtime paths") + } + + async fn send_request( + writer: &mut DuplexStream, + id: i64, + method: &str, + params: &P, + ) { + write_message( + writer, + &JSONRPCMessage::Request(JSONRPCRequest { + id: RequestId::Integer(id), + method: method.to_string(), + params: Some(serde_json::to_value(params).expect("serialize params")), + trace: None, + }), + ) + .await; + } + + async fn send_notification(writer: &mut DuplexStream, method: &str, params: &P) { + write_message( + writer, + &JSONRPCMessage::Notification(JSONRPCNotification { + method: method.to_string(), + params: Some(serde_json::to_value(params).expect("serialize params")), + }), + ) + .await; + } + + async fn write_message(writer: &mut DuplexStream, message: &JSONRPCMessage) { + let encoded = serde_json::to_vec(message).expect("serialize JSON-RPC message"); + writer.write_all(&encoded).await.expect("write request"); + writer.write_all(b"\n").await.expect("write newline"); + } + + async fn read_response( + lines: &mut Lines>, + expected_id: i64, + ) -> T { + let line = lines + .next_line() + .await + .expect("read response") + .expect("response line"); + match serde_json::from_str::(&line).expect("decode JSON-RPC response") { + JSONRPCMessage::Response(JSONRPCResponse { id, result }) => { + assert_eq!(id, RequestId::Integer(expected_id)); + serde_json::from_value(result).expect("decode response result") + } + JSONRPCMessage::Error(error) => panic!("unexpected JSON-RPC error: {error:?}"), + other => panic!("expected JSON-RPC response, got {other:?}"), + } + } + + fn exec_params(process_id: ProcessId) -> ExecParams { + let mut env = HashMap::new(); + if let Some(path) = std::env::var_os("PATH") { + env.insert("PATH".to_string(), path.to_string_lossy().into_owned()); + } + ExecParams { + process_id, + argv: sleep_then_print_argv(), + cwd: std::env::current_dir().expect("cwd"), + env_policy: None, + env, + tty: false, + pipe_stdin: false, + arg0: None, + } + } + + fn sleep_then_print_argv() -> Vec { + if cfg!(windows) { + vec![ + std::env::var("COMSPEC").unwrap_or_else(|_| "cmd.exe".to_string()), + "/C".to_string(), + "ping -n 3 127.0.0.1 >NUL && echo late".to_string(), + ] + } else { + vec![ + "/bin/sh".to_string(), + "-c".to_string(), + "sleep 1; printf late".to_string(), + ] + } + } +} From 579f4731dfd40d5a9f136550f64101a1c5bc6dc7 Mon Sep 17 00:00:00 2001 From: starr-openai Date: Tue, 5 May 2026 16:30:33 -0700 Subject: [PATCH 20/42] Simplify exec-server connection ownership Remove the runtime extraction helpers and make JsonRpcConnection ownership explicit at the destructuring sites. Let the stdio transport clean up through Drop so ExecServerClient no longer needs to call an explicit shutdown hook. Co-authored-by: Codex --- codex-rs/exec-server/src/client.rs | 8 +- codex-rs/exec-server/src/connection.rs | 87 ++++---------------- codex-rs/exec-server/src/rpc.rs | 31 ++++--- codex-rs/exec-server/src/server/processor.rs | 13 ++- 4 files changed, 49 insertions(+), 90 deletions(-) diff --git a/codex-rs/exec-server/src/client.rs b/codex-rs/exec-server/src/client.rs index 3fbb35cb90ec..d73750206910 100644 --- a/codex-rs/exec-server/src/client.rs +++ b/codex-rs/exec-server/src/client.rs @@ -152,8 +152,6 @@ pub(crate) struct Session { } struct Inner { - // Keep the underlying transport connection alive for the client lifetime. - connection: JsonRpcConnection, client: RpcClient, // The remote transport delivers one shared notification stream for every // process on the connection. Keep a local process_id -> session registry so @@ -183,7 +181,6 @@ struct Inner { impl Drop for Inner { fn drop(&mut self) { self.reader_task.abort(); - self.connection.shutdown(); } } @@ -428,10 +425,10 @@ impl ExecServerClient { } pub(crate) async fn connect( - mut connection: JsonRpcConnection, + connection: JsonRpcConnection, options: ExecServerClientConnectOptions, ) -> Result { - let (rpc_client, mut events_rx) = RpcClient::new(&mut connection); + let (rpc_client, mut events_rx) = RpcClient::new(connection); let inner = Arc::new_cyclic(|weak| { let weak = weak.clone(); let reader_task = tokio::spawn(async move { @@ -465,7 +462,6 @@ impl ExecServerClient { }); Inner { - connection, client: rpc_client, sessions: ArcSwap::from_pointee(HashMap::new()), sessions_write_lock: Mutex::new(()), diff --git a/codex-rs/exec-server/src/connection.rs b/codex-rs/exec-server/src/connection.rs index e159dade288e..bb92791138a7 100644 --- a/codex-rs/exec-server/src/connection.rs +++ b/codex-rs/exec-server/src/connection.rs @@ -24,7 +24,7 @@ pub(crate) enum JsonRpcConnectionEvent { Disconnected { reason: Option }, } -enum JsonRpcTransport { +pub(crate) enum JsonRpcTransport { Plain, Stdio(StdioTransport), } @@ -35,21 +35,14 @@ impl JsonRpcTransport { child_process: Some(child_process), }) } - - fn shutdown(&mut self) { - match self { - Self::Plain => {} - Self::Stdio(transport) => transport.shutdown(), - } - } } -struct StdioTransport { +pub(crate) struct StdioTransport { child_process: Option, } -impl StdioTransport { - fn shutdown(&mut self) { +impl Drop for StdioTransport { + fn drop(&mut self) { let Some(mut child_process) = self.child_process.take() else { return; }; @@ -72,29 +65,19 @@ impl StdioTransport { } } -struct JsonRpcConnectionRuntime { - outgoing_tx: mpsc::Sender, - incoming_rx: mpsc::Receiver, - disconnected_rx: watch::Receiver, - task_handles: Vec>, +pub(crate) struct JsonRpcConnectionRuntime { + pub(crate) outgoing_tx: mpsc::Sender, + pub(crate) incoming_rx: mpsc::Receiver, + pub(crate) disconnected_rx: watch::Receiver, + pub(crate) task_handles: Vec>, } pub(crate) struct JsonRpcConnection { - runtime: Option, - transport: JsonRpcTransport, -} - -impl Drop for JsonRpcConnection { - fn drop(&mut self) { - self.shutdown(); - } + pub(crate) runtime: JsonRpcConnectionRuntime, + pub(crate) transport: JsonRpcTransport, } impl JsonRpcConnection { - pub(crate) fn shutdown(&mut self) { - self.transport.shutdown(); - } - pub(crate) fn from_stdio(reader: R, writer: W, connection_label: String) -> Self where R: AsyncRead + Unpin + Send + 'static, @@ -178,12 +161,12 @@ impl JsonRpcConnection { }); Self { - runtime: Some(JsonRpcConnectionRuntime { + runtime: JsonRpcConnectionRuntime { outgoing_tx, incoming_rx, disconnected_rx, task_handles: vec![reader_task, writer_task], - }), + }, transport: JsonRpcTransport::Plain, } } @@ -315,60 +298,20 @@ impl JsonRpcConnection { }); Self { - runtime: Some(JsonRpcConnectionRuntime { + runtime: JsonRpcConnectionRuntime { outgoing_tx, incoming_rx, disconnected_rx, task_handles: vec![reader_task, writer_task], - }), + }, transport: JsonRpcTransport::Plain, } } - pub(crate) fn take_runtime( - &mut self, - ) -> ( - mpsc::Sender, - mpsc::Receiver, - Vec>, - ) { - let JsonRpcConnectionRuntime { - outgoing_tx, - incoming_rx, - disconnected_rx: _, - task_handles, - } = self.take_runtime_or_panic("JSON-RPC connection runtime already taken"); - (outgoing_tx, incoming_rx, task_handles) - } - - pub(crate) fn into_parts( - mut self, - ) -> ( - mpsc::Sender, - mpsc::Receiver, - watch::Receiver, - Vec>, - ) { - let JsonRpcConnectionRuntime { - outgoing_tx, - incoming_rx, - disconnected_rx, - task_handles, - } = self.take_runtime_or_panic("JSON-RPC connection runtime already taken"); - (outgoing_tx, incoming_rx, disconnected_rx, task_handles) - } - pub(crate) fn with_child_process(mut self, child_process: Child) -> Self { self.transport = JsonRpcTransport::from_child_process(child_process); self } - - fn take_runtime_or_panic(&mut self, message: &'static str) -> JsonRpcConnectionRuntime { - match self.runtime.take() { - Some(runtime) => runtime, - None => panic!("{message}"), - } - } } async fn send_disconnected( diff --git a/codex-rs/exec-server/src/rpc.rs b/codex-rs/exec-server/src/rpc.rs index 8c7c17753670..c90fef2ab9a4 100644 --- a/codex-rs/exec-server/src/rpc.rs +++ b/codex-rs/exec-server/src/rpc.rs @@ -23,6 +23,8 @@ use tokio::task::JoinHandle; use crate::connection::JsonRpcConnection; use crate::connection::JsonRpcConnectionEvent; +use crate::connection::JsonRpcConnectionRuntime; +use crate::connection::JsonRpcTransport; #[derive(Debug)] pub(crate) enum RpcCallError { @@ -225,14 +227,22 @@ pub(crate) struct RpcClient { closed: Arc, next_request_id: AtomicI64, transport_tasks: Vec>, + _transport: JsonRpcTransport, reader_task: JoinHandle<()>, } impl RpcClient { - pub(crate) fn new( - connection: &mut JsonRpcConnection, - ) -> (Self, mpsc::Receiver) { - let (write_tx, mut incoming_rx, transport_tasks) = connection.take_runtime(); + pub(crate) fn new(connection: JsonRpcConnection) -> (Self, mpsc::Receiver) { + let JsonRpcConnection { + runtime: + JsonRpcConnectionRuntime { + outgoing_tx: write_tx, + incoming_rx: mut incoming_rx, + disconnected_rx: _, + task_handles: transport_tasks, + }, + transport, + } = connection; let pending = Arc::new(Mutex::new(HashMap::::new())); let (event_tx, event_rx) = mpsc::channel(128); let closed = Arc::new(AtomicBool::new(false)); @@ -272,6 +282,7 @@ impl RpcClient { closed, next_request_id: AtomicI64::new(1), transport_tasks, + _transport: transport, reader_task, }, event_rx, @@ -570,9 +581,9 @@ mod tests { async fn rpc_client_matches_out_of_order_responses_by_request_id() { let (client_stdin, server_reader) = tokio::io::duplex(4096); let (mut server_writer, client_stdout) = tokio::io::duplex(4096); - let mut connection = + let connection = JsonRpcConnection::from_stdio(client_stdout, client_stdin, "test-rpc".to_string()); - let (client, _events_rx) = RpcClient::new(&mut connection); + let (client, _events_rx) = RpcClient::new(connection); let server = tokio::spawn(async move { let mut lines = BufReader::new(server_reader).lines(); @@ -636,9 +647,9 @@ mod tests { async fn rpc_client_rejects_new_calls_after_reader_protocol_error() { let (client_stdin, _server_reader) = tokio::io::duplex(4096); let (mut server_writer, client_stdout) = tokio::io::duplex(4096); - let mut connection = + let connection = JsonRpcConnection::from_stdio(client_stdout, client_stdin, "test-rpc".to_string()); - let (client, mut events_rx) = RpcClient::new(&mut connection); + let (client, mut events_rx) = RpcClient::new(connection); write_jsonrpc_line( &mut server_writer, @@ -681,9 +692,9 @@ mod tests { async fn rpc_client_drains_pending_call_on_transport_eof() { let (client_stdin, server_reader) = tokio::io::duplex(4096); let (server_writer, client_stdout) = tokio::io::duplex(4096); - let mut connection = + let connection = JsonRpcConnection::from_stdio(client_stdout, client_stdin, "test-rpc".to_string()); - let (client, mut events_rx) = RpcClient::new(&mut connection); + let (client, mut events_rx) = RpcClient::new(connection); let server = tokio::spawn(async move { let mut lines = BufReader::new(server_reader).lines(); diff --git a/codex-rs/exec-server/src/server/processor.rs b/codex-rs/exec-server/src/server/processor.rs index dc1a9b9ffe74..a880e35fe962 100644 --- a/codex-rs/exec-server/src/server/processor.rs +++ b/codex-rs/exec-server/src/server/processor.rs @@ -8,6 +8,7 @@ use crate::ExecServerRuntimePaths; use crate::connection::CHANNEL_CAPACITY; use crate::connection::JsonRpcConnection; use crate::connection::JsonRpcConnectionEvent; +use crate::connection::JsonRpcConnectionRuntime; use crate::rpc::RpcNotificationSender; use crate::rpc::RpcServerOutboundMessage; use crate::rpc::encode_server_message; @@ -47,8 +48,16 @@ async fn run_connection( runtime_paths: ExecServerRuntimePaths, ) { let router = Arc::new(build_router()); - let (json_outgoing_tx, mut incoming_rx, mut disconnected_rx, connection_tasks) = - connection.into_parts(); + let JsonRpcConnection { + runtime: + JsonRpcConnectionRuntime { + outgoing_tx: json_outgoing_tx, + incoming_rx: mut incoming_rx, + disconnected_rx: mut disconnected_rx, + task_handles: connection_tasks, + }, + transport: _transport, + } = connection; let (outgoing_tx, mut outgoing_rx) = mpsc::channel::(CHANNEL_CAPACITY); let notifications = RpcNotificationSender::new(outgoing_tx.clone()); From c0f9dafadb978e0a8eb6714685b4264e618c9f46 Mon Sep 17 00:00:00 2001 From: starr-openai Date: Tue, 5 May 2026 16:32:19 -0700 Subject: [PATCH 21/42] Flatten JSON-RPC connection state Drop the separate JsonRpcConnectionRuntime wrapper so JsonRpcConnection directly owns the channels, disconnect watch, transport tasks, and transport guard. This keeps the lifetime model explicit without helper extraction methods. Co-authored-by: Codex --- codex-rs/exec-server/src/connection.rs | 26 +++++++------------- codex-rs/exec-server/src/rpc.rs | 12 +++------ codex-rs/exec-server/src/server/processor.rs | 12 +++------ 3 files changed, 17 insertions(+), 33 deletions(-) diff --git a/codex-rs/exec-server/src/connection.rs b/codex-rs/exec-server/src/connection.rs index bb92791138a7..2c0d7ee6d260 100644 --- a/codex-rs/exec-server/src/connection.rs +++ b/codex-rs/exec-server/src/connection.rs @@ -65,15 +65,11 @@ impl Drop for StdioTransport { } } -pub(crate) struct JsonRpcConnectionRuntime { +pub(crate) struct JsonRpcConnection { pub(crate) outgoing_tx: mpsc::Sender, pub(crate) incoming_rx: mpsc::Receiver, pub(crate) disconnected_rx: watch::Receiver, pub(crate) task_handles: Vec>, -} - -pub(crate) struct JsonRpcConnection { - pub(crate) runtime: JsonRpcConnectionRuntime, pub(crate) transport: JsonRpcTransport, } @@ -161,12 +157,10 @@ impl JsonRpcConnection { }); Self { - runtime: JsonRpcConnectionRuntime { - outgoing_tx, - incoming_rx, - disconnected_rx, - task_handles: vec![reader_task, writer_task], - }, + outgoing_tx, + incoming_rx, + disconnected_rx, + task_handles: vec![reader_task, writer_task], transport: JsonRpcTransport::Plain, } } @@ -298,12 +292,10 @@ impl JsonRpcConnection { }); Self { - runtime: JsonRpcConnectionRuntime { - outgoing_tx, - incoming_rx, - disconnected_rx, - task_handles: vec![reader_task, writer_task], - }, + outgoing_tx, + incoming_rx, + disconnected_rx, + task_handles: vec![reader_task, writer_task], transport: JsonRpcTransport::Plain, } } diff --git a/codex-rs/exec-server/src/rpc.rs b/codex-rs/exec-server/src/rpc.rs index c90fef2ab9a4..3e99976be841 100644 --- a/codex-rs/exec-server/src/rpc.rs +++ b/codex-rs/exec-server/src/rpc.rs @@ -23,7 +23,6 @@ use tokio::task::JoinHandle; use crate::connection::JsonRpcConnection; use crate::connection::JsonRpcConnectionEvent; -use crate::connection::JsonRpcConnectionRuntime; use crate::connection::JsonRpcTransport; #[derive(Debug)] @@ -234,13 +233,10 @@ pub(crate) struct RpcClient { impl RpcClient { pub(crate) fn new(connection: JsonRpcConnection) -> (Self, mpsc::Receiver) { let JsonRpcConnection { - runtime: - JsonRpcConnectionRuntime { - outgoing_tx: write_tx, - incoming_rx: mut incoming_rx, - disconnected_rx: _, - task_handles: transport_tasks, - }, + outgoing_tx: write_tx, + incoming_rx: mut incoming_rx, + disconnected_rx: _, + task_handles: transport_tasks, transport, } = connection; let pending = Arc::new(Mutex::new(HashMap::::new())); diff --git a/codex-rs/exec-server/src/server/processor.rs b/codex-rs/exec-server/src/server/processor.rs index a880e35fe962..5382512be35a 100644 --- a/codex-rs/exec-server/src/server/processor.rs +++ b/codex-rs/exec-server/src/server/processor.rs @@ -8,7 +8,6 @@ use crate::ExecServerRuntimePaths; use crate::connection::CHANNEL_CAPACITY; use crate::connection::JsonRpcConnection; use crate::connection::JsonRpcConnectionEvent; -use crate::connection::JsonRpcConnectionRuntime; use crate::rpc::RpcNotificationSender; use crate::rpc::RpcServerOutboundMessage; use crate::rpc::encode_server_message; @@ -49,13 +48,10 @@ async fn run_connection( ) { let router = Arc::new(build_router()); let JsonRpcConnection { - runtime: - JsonRpcConnectionRuntime { - outgoing_tx: json_outgoing_tx, - incoming_rx: mut incoming_rx, - disconnected_rx: mut disconnected_rx, - task_handles: connection_tasks, - }, + outgoing_tx: json_outgoing_tx, + incoming_rx: mut incoming_rx, + disconnected_rx: mut disconnected_rx, + task_handles: connection_tasks, transport: _transport, } = connection; let (outgoing_tx, mut outgoing_rx) = From d42a1e01fa16419427b1747a379079fb899701b7 Mon Sep 17 00:00:00 2001 From: starr-openai Date: Wed, 6 May 2026 11:42:16 -0700 Subject: [PATCH 22/42] Narrow stdio client lifetime handling Keep the retained transport ownership needed for stdio child cleanup, but drop the broader AtomicBool closed-state behavior and its targeted tests from this PR. Co-authored-by: Codex --- codex-rs/exec-server/src/rpc.rs | 145 +++++--------------------------- 1 file changed, 22 insertions(+), 123 deletions(-) diff --git a/codex-rs/exec-server/src/rpc.rs b/codex-rs/exec-server/src/rpc.rs index 3e99976be841..08700343b656 100644 --- a/codex-rs/exec-server/src/rpc.rs +++ b/codex-rs/exec-server/src/rpc.rs @@ -2,7 +2,6 @@ use std::collections::HashMap; use std::future::Future; use std::pin::Pin; use std::sync::Arc; -use std::sync::atomic::AtomicBool; use std::sync::atomic::AtomicI64; use std::sync::atomic::Ordering; @@ -19,6 +18,7 @@ use serde_json::Value; use tokio::sync::Mutex; use tokio::sync::mpsc; use tokio::sync::oneshot; +use tokio::sync::watch; use tokio::task::JoinHandle; use crate::connection::JsonRpcConnection; @@ -221,9 +221,10 @@ where pub(crate) struct RpcClient { write_tx: mpsc::Sender, pending: Arc>>, - // This flips before the ordered RPC reader drains pending requests, so new - // calls fail instead of registering work that can never complete. - closed: Arc, + // Shared transport state from `JsonRpcConnection`. Calls use this to fail + // immediately when the socket closes, even if no JSON-RPC error response + // can be delivered for their request id. + disconnected_rx: watch::Receiver, next_request_id: AtomicI64, transport_tasks: Vec>, _transport: JsonRpcTransport, @@ -235,39 +236,40 @@ impl RpcClient { let JsonRpcConnection { outgoing_tx: write_tx, incoming_rx: mut incoming_rx, - disconnected_rx: _, + disconnected_rx, task_handles: transport_tasks, transport, } = connection; let pending = Arc::new(Mutex::new(HashMap::::new())); let (event_tx, event_rx) = mpsc::channel(128); - let closed = Arc::new(AtomicBool::new(false)); let pending_for_reader = Arc::clone(&pending); - let closed_for_reader = Arc::clone(&closed); let reader_task = tokio::spawn(async move { - let reason = loop { - let Some(event) = incoming_rx.recv().await else { - break None; - }; - + while let Some(event) = incoming_rx.recv().await { match event { JsonRpcConnectionEvent::Message(message) => { if let Err(err) = handle_server_message(&pending_for_reader, &event_tx, message).await { - break Some(err); + let _ = err; + break; } } JsonRpcConnectionEvent::MalformedMessage { reason } => { - break Some(reason); + let _ = reason; + break; + } + JsonRpcConnectionEvent::Disconnected { reason } => { + let _ = event_tx.send(RpcClientEvent::Disconnected { reason }).await; + drain_pending(&pending_for_reader).await; + return; } - JsonRpcConnectionEvent::Disconnected { reason } => break reason, } - }; + } - closed_for_reader.store(true, Ordering::SeqCst); - let _ = event_tx.send(RpcClientEvent::Disconnected { reason }).await; + let _ = event_tx + .send(RpcClientEvent::Disconnected { reason: None }) + .await; drain_pending(&pending_for_reader).await; }); @@ -275,7 +277,7 @@ impl RpcClient { Self { write_tx, pending, - closed, + disconnected_rx, next_request_id: AtomicI64::new(1), transport_tasks, _transport: transport, @@ -291,12 +293,6 @@ impl RpcClient { params: &P, ) -> Result<(), serde_json::Error> { let params = serde_json::to_value(params)?; - if self.closed.load(Ordering::SeqCst) { - return Err(serde_json::Error::io(std::io::Error::new( - std::io::ErrorKind::BrokenPipe, - "JSON-RPC transport closed", - ))); - } self.write_tx .send(JSONRPCMessage::Notification(JSONRPCNotification { method: method.to_string(), @@ -323,7 +319,7 @@ impl RpcClient { // Registering the pending request and checking disconnect must be // atomic with the reader's drain_pending path. Otherwise a call // can sneak in after the drain and wait forever. - if self.closed.load(Ordering::SeqCst) { + if *self.disconnected_rx.borrow() { return Err(RpcCallError::Closed); } pending.insert(request_id.clone(), response_tx); @@ -524,9 +520,7 @@ mod tests { use std::time::Duration; use codex_app_server_protocol::JSONRPCMessage; - use codex_app_server_protocol::JSONRPCRequest; use codex_app_server_protocol::JSONRPCResponse; - use codex_app_server_protocol::RequestId; use pretty_assertions::assert_eq; use tokio::io::AsyncBufReadExt; use tokio::io::AsyncWriteExt; @@ -534,7 +528,6 @@ mod tests { use tokio::time::timeout; use super::RpcClient; - use super::RpcClientEvent; use crate::connection::JsonRpcConnection; async fn read_jsonrpc_line(lines: &mut tokio::io::Lines>) -> JSONRPCMessage @@ -638,98 +631,4 @@ mod tests { panic!("server task failed: {err}"); } } - - #[tokio::test] - async fn rpc_client_rejects_new_calls_after_reader_protocol_error() { - let (client_stdin, _server_reader) = tokio::io::duplex(4096); - let (mut server_writer, client_stdout) = tokio::io::duplex(4096); - let connection = - JsonRpcConnection::from_stdio(client_stdout, client_stdin, "test-rpc".to_string()); - let (client, mut events_rx) = RpcClient::new(connection); - - write_jsonrpc_line( - &mut server_writer, - JSONRPCMessage::Request(JSONRPCRequest { - id: RequestId::Integer(1), - method: "server/request".to_string(), - params: None, - trace: None, - }), - ) - .await; - - let event = timeout(Duration::from_secs(1), events_rx.recv()) - .await - .expect("timed out waiting for disconnect event"); - match event { - Some(RpcClientEvent::Disconnected { reason }) => { - assert!( - reason - .as_deref() - .is_some_and(|reason| reason.contains("unexpected JSON-RPC request")), - "unexpected disconnect reason: {reason:?}" - ); - } - event => panic!("expected disconnect event, got {event:?}"), - } - - let result = timeout( - Duration::from_secs(1), - client.call::<_, serde_json::Value>("after-close", &serde_json::json!({})), - ) - .await - .expect("timed out waiting for closed call"); - - assert!(matches!(result, Err(super::RpcCallError::Closed))); - assert_eq!(client.pending_request_count().await, 0); - } - - #[tokio::test] - async fn rpc_client_drains_pending_call_on_transport_eof() { - let (client_stdin, server_reader) = tokio::io::duplex(4096); - let (server_writer, client_stdout) = tokio::io::duplex(4096); - let connection = - JsonRpcConnection::from_stdio(client_stdout, client_stdin, "test-rpc".to_string()); - let (client, mut events_rx) = RpcClient::new(connection); - - let server = tokio::spawn(async move { - let mut lines = BufReader::new(server_reader).lines(); - let request = read_jsonrpc_line(&mut lines).await; - match request { - JSONRPCMessage::Request(request) if request.method == "will-close" => {} - other => panic!("expected will-close request, got {other:?}"), - } - drop(server_writer); - }); - - let result = timeout( - Duration::from_secs(1), - client.call::<_, serde_json::Value>("will-close", &serde_json::json!({})), - ) - .await - .expect("timed out waiting for closed call"); - assert!(matches!(result, Err(super::RpcCallError::Closed))); - - let event = timeout(Duration::from_secs(1), events_rx.recv()) - .await - .expect("timed out waiting for disconnect event"); - assert!(matches!( - event, - Some(RpcClientEvent::Disconnected { reason: None }) - )); - assert_eq!(client.pending_request_count().await, 0); - - let result = timeout( - Duration::from_secs(1), - client.call::<_, serde_json::Value>("after-close", &serde_json::json!({})), - ) - .await - .expect("timed out waiting for fast closed call"); - assert!(matches!(result, Err(super::RpcCallError::Closed))); - - let notify = client.notify("after-close", &serde_json::json!({})).await; - assert!(notify.is_err()); - - server.await.expect("server task should finish"); - } } From 40ea3b74889f75a1c2e3dab777ca849df3e05cf0 Mon Sep 17 00:00:00 2001 From: starr-openai Date: Wed, 6 May 2026 14:29:55 -0700 Subject: [PATCH 23/42] Fix stdio transport clippy issues Keep the stack-introduced stdio transport variant explicit while avoiding dead-code and redundant-pattern lints reported by PR20664 CI. Co-authored-by: Codex --- codex-rs/exec-server/src/client_api.rs | 1 + codex-rs/exec-server/src/connection.rs | 10 ++++++---- codex-rs/exec-server/src/rpc.rs | 2 +- codex-rs/exec-server/src/server/processor.rs | 4 ++-- 4 files changed, 10 insertions(+), 7 deletions(-) diff --git a/codex-rs/exec-server/src/client_api.rs b/codex-rs/exec-server/src/client_api.rs index 20520f002ef6..8adfadd6e705 100644 --- a/codex-rs/exec-server/src/client_api.rs +++ b/codex-rs/exec-server/src/client_api.rs @@ -49,6 +49,7 @@ pub(crate) struct StdioExecServerCommand { #[derive(Debug, Clone, PartialEq, Eq)] pub(crate) enum ExecServerTransportParams { WebSocketUrl(String), + #[allow(dead_code)] StdioCommand(StdioExecServerCommand), } diff --git a/codex-rs/exec-server/src/connection.rs b/codex-rs/exec-server/src/connection.rs index 2c0d7ee6d260..e672d1ae38a3 100644 --- a/codex-rs/exec-server/src/connection.rs +++ b/codex-rs/exec-server/src/connection.rs @@ -26,14 +26,16 @@ pub(crate) enum JsonRpcConnectionEvent { pub(crate) enum JsonRpcTransport { Plain, - Stdio(StdioTransport), + Stdio { _transport: StdioTransport }, } impl JsonRpcTransport { fn from_child_process(child_process: Child) -> Self { - Self::Stdio(StdioTransport { - child_process: Some(child_process), - }) + Self::Stdio { + _transport: StdioTransport { + child_process: Some(child_process), + }, + } } } diff --git a/codex-rs/exec-server/src/rpc.rs b/codex-rs/exec-server/src/rpc.rs index 08700343b656..9ea41f3854c9 100644 --- a/codex-rs/exec-server/src/rpc.rs +++ b/codex-rs/exec-server/src/rpc.rs @@ -235,7 +235,7 @@ impl RpcClient { pub(crate) fn new(connection: JsonRpcConnection) -> (Self, mpsc::Receiver) { let JsonRpcConnection { outgoing_tx: write_tx, - incoming_rx: mut incoming_rx, + mut incoming_rx, disconnected_rx, task_handles: transport_tasks, transport, diff --git a/codex-rs/exec-server/src/server/processor.rs b/codex-rs/exec-server/src/server/processor.rs index 5382512be35a..6fc0723f0c1e 100644 --- a/codex-rs/exec-server/src/server/processor.rs +++ b/codex-rs/exec-server/src/server/processor.rs @@ -49,8 +49,8 @@ async fn run_connection( let router = Arc::new(build_router()); let JsonRpcConnection { outgoing_tx: json_outgoing_tx, - incoming_rx: mut incoming_rx, - disconnected_rx: mut disconnected_rx, + mut incoming_rx, + mut disconnected_rx, task_handles: connection_tasks, transport: _transport, } = connection; From f5d901470cf477f9871c64e4bf281aa39ac099b1 Mon Sep 17 00:00:00 2001 From: starr-openai Date: Wed, 6 May 2026 14:39:01 -0700 Subject: [PATCH 24/42] Box retained stdio transport guard Avoid the Windows clippy large-enum-variant failure while preserving the retained stdio child cleanup guard behavior. Co-authored-by: Codex --- codex-rs/exec-server/src/connection.rs | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/codex-rs/exec-server/src/connection.rs b/codex-rs/exec-server/src/connection.rs index e672d1ae38a3..bf3d28a8330e 100644 --- a/codex-rs/exec-server/src/connection.rs +++ b/codex-rs/exec-server/src/connection.rs @@ -26,15 +26,17 @@ pub(crate) enum JsonRpcConnectionEvent { pub(crate) enum JsonRpcTransport { Plain, - Stdio { _transport: StdioTransport }, + Stdio { + _transport: Box, + }, } impl JsonRpcTransport { fn from_child_process(child_process: Child) -> Self { Self::Stdio { - _transport: StdioTransport { + _transport: Box::new(StdioTransport { child_process: Some(child_process), - }, + }), } } } From 9d933226013ae425aa7c6e0b550a69fed88840a8 Mon Sep 17 00:00:00 2001 From: starr-openai Date: Wed, 6 May 2026 14:43:47 -0700 Subject: [PATCH 25/42] Apply rustfmt to stdio transport guard Match the rustfmt shape reported by the PR20664 Format / etc CI job after boxing the retained stdio transport guard. Co-authored-by: Codex --- codex-rs/exec-server/src/connection.rs | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/codex-rs/exec-server/src/connection.rs b/codex-rs/exec-server/src/connection.rs index bf3d28a8330e..f1e65e321aac 100644 --- a/codex-rs/exec-server/src/connection.rs +++ b/codex-rs/exec-server/src/connection.rs @@ -26,9 +26,7 @@ pub(crate) enum JsonRpcConnectionEvent { pub(crate) enum JsonRpcTransport { Plain, - Stdio { - _transport: Box, - }, + Stdio { _transport: Box }, } impl JsonRpcTransport { From 6e3d29412ddeee1344238f7c185071d1136da97f Mon Sep 17 00:00:00 2001 From: starr-openai Date: Wed, 6 May 2026 15:11:16 -0700 Subject: [PATCH 26/42] Fix Windows stdio test JSON quoting Escape the JSON-RPC response quotes in the cmd.exe stdio test command so Windows emits valid JSON before the client initialize timeout. Co-authored-by: Codex --- codex-rs/exec-server/src/client.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/codex-rs/exec-server/src/client.rs b/codex-rs/exec-server/src/client.rs index d73750206910..3ab1be8112a6 100644 --- a/codex-rs/exec-server/src/client.rs +++ b/codex-rs/exec-server/src/client.rs @@ -964,7 +964,7 @@ mod tests { program: "cmd".to_string(), args: vec![ "/C".to_string(), - "set /p _line= & echo {\"id\":1,\"result\":{\"sessionId\":\"stdio-test\"}} & set /p _line= & ping -n 60 127.0.0.1 >nul".to_string(), + "set /p _line= & echo {^\"id^\":1,^\"result^\":{^\"sessionId^\":^\"stdio-test^\"}} & set /p _line= & ping -n 60 127.0.0.1 >nul".to_string(), ], env: HashMap::new(), cwd: None, From f0234866f6faacad47a33ebb17d57026cdb7efeb Mon Sep 17 00:00:00 2001 From: starr-openai Date: Wed, 6 May 2026 16:02:24 -0700 Subject: [PATCH 27/42] Use PowerShell for Windows stdio test helper Avoid cmd.exe echo quoting semantics in the Windows stdio client test by reading stdin and writing the JSON-RPC initialize response from PowerShell. Co-authored-by: Codex --- codex-rs/exec-server/src/client.rs | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/codex-rs/exec-server/src/client.rs b/codex-rs/exec-server/src/client.rs index 3ab1be8112a6..55cd99b6c9c2 100644 --- a/codex-rs/exec-server/src/client.rs +++ b/codex-rs/exec-server/src/client.rs @@ -961,10 +961,11 @@ mod tests { async fn connect_stdio_command_initializes_json_rpc_client_on_windows() { let client = ExecServerClient::connect_stdio_command(StdioExecServerConnectArgs { command: StdioExecServerCommand { - program: "cmd".to_string(), + program: "powershell".to_string(), args: vec![ - "/C".to_string(), - "set /p _line= & echo {^\"id^\":1,^\"result^\":{^\"sessionId^\":^\"stdio-test^\"}} & set /p _line= & ping -n 60 127.0.0.1 >nul".to_string(), + "-NoProfile".to_string(), + "-Command".to_string(), + "$null = [Console]::In.ReadLine(); [Console]::Out.WriteLine('{\"id\":1,\"result\":{\"sessionId\":\"stdio-test\"}}'); $null = [Console]::In.ReadLine(); Start-Sleep -Seconds 60".to_string(), ], env: HashMap::new(), cwd: None, From 9012decfc39bd4b3a80774b7a8c7c5710a092ac1 Mon Sep 17 00:00:00 2001 From: starr-openai Date: Fri, 1 May 2026 12:07:47 -0700 Subject: [PATCH 28/42] Make environment providers own default selection Let environment providers return an explicit default selection and let remote environments track the underlying transport instead of treating only websocket URLs as remote. This prepares the environment layer for stdio-backed remotes without introducing config-file loading. Co-authored-by: Codex --- codex-rs/exec-server/src/environment.rs | 154 ++++++++++++++---- .../exec-server/src/environment_provider.rs | 23 +++ codex-rs/exec-server/src/lib.rs | 1 + 3 files changed, 149 insertions(+), 29 deletions(-) diff --git a/codex-rs/exec-server/src/environment.rs b/codex-rs/exec-server/src/environment.rs index 395cac9e6333..4372f69c1d15 100644 --- a/codex-rs/exec-server/src/environment.rs +++ b/codex-rs/exec-server/src/environment.rs @@ -9,6 +9,7 @@ use crate::client::LazyRemoteExecServerClient; use crate::client::http_client::ReqwestHttpClient; use crate::client_api::ExecServerTransportParams; use crate::environment_provider::DefaultEnvironmentProvider; +use crate::environment_provider::DefaultEnvironmentSelection; use crate::environment_provider::EnvironmentProvider; use crate::environment_provider::normalize_exec_server_url; use crate::local_file_system::LocalFileSystem; @@ -32,8 +33,8 @@ pub const CODEX_EXEC_SERVER_URL_ENV_VAR: &str = "CODEX_EXEC_SERVER_URL"; /// shell/filesystem tool availability. /// /// Remote environments create remote filesystem and execution backends that -/// lazy-connect to the configured exec-server on first use. The websocket is -/// not opened when the manager or environment is constructed. +/// lazy-connect to the configured exec-server on first use. The remote +/// transport is not opened when the manager or environment is constructed. #[derive(Debug)] pub struct EnvironmentManager { default_environment: Option, @@ -72,9 +73,12 @@ impl EnvironmentManager { /// Builds a test-only manager with environment access disabled. pub fn disabled_for_tests(local_runtime_paths: ExecServerRuntimePaths) -> Self { - let mut manager = Self::from_environments(HashMap::new(), local_runtime_paths); - manager.default_environment = None; - manager + Self::from_environments( + HashMap::new(), + local_runtime_paths, + DefaultEnvironmentSelection::Disabled, + ) + .expect("disabled test environment manager") } /// Builds a test-only manager from a raw exec-server URL value. @@ -99,16 +103,14 @@ impl EnvironmentManager { exec_server_url: Option, local_runtime_paths: ExecServerRuntimePaths, ) -> Self { - let environment_disabled = normalize_exec_server_url(exec_server_url.clone()).1; let provider = DefaultEnvironmentProvider::new(exec_server_url); let provider_environments = provider.environments(&local_runtime_paths); - let mut manager = Self::from_environments(provider_environments, local_runtime_paths); - if environment_disabled { - // TODO: Remove this legacy `CODEX_EXEC_SERVER_URL=none` crutch once - // environment attachment defaulting moves out of EnvironmentManager. - manager.default_environment = None; - } - manager + Self::from_environments( + provider_environments, + local_runtime_paths, + provider.default_environment_selection(), + ) + .expect("default provider should create valid environments") } /// Builds a manager from a provider-supplied startup snapshot. @@ -122,12 +124,14 @@ impl EnvironmentManager { Self::from_provider_environments( provider.get_environments(&local_runtime_paths).await?, local_runtime_paths, + provider.default_environment_selection(), ) } fn from_provider_environments( environments: HashMap, local_runtime_paths: ExecServerRuntimePaths, + default_selection: DefaultEnvironmentSelection, ) -> Result { for id in environments.keys() { if id.is_empty() { @@ -137,21 +141,35 @@ impl EnvironmentManager { } } - Ok(Self::from_environments(environments, local_runtime_paths)) + Self::from_environments(environments, local_runtime_paths, default_selection) } fn from_environments( environments: HashMap, local_runtime_paths: ExecServerRuntimePaths, - ) -> Self { + default_selection: DefaultEnvironmentSelection, + ) -> Result { // TODO: Stop deriving a default environment here once omitted // environment attachment is owned by thread/session setup. - let default_environment = if environments.contains_key(REMOTE_ENVIRONMENT_ID) { - Some(REMOTE_ENVIRONMENT_ID.to_string()) - } else if environments.contains_key(LOCAL_ENVIRONMENT_ID) { - Some(LOCAL_ENVIRONMENT_ID.to_string()) - } else { - None + let default_environment = match default_selection { + DefaultEnvironmentSelection::Derived => { + if environments.contains_key(REMOTE_ENVIRONMENT_ID) { + Some(REMOTE_ENVIRONMENT_ID.to_string()) + } else if environments.contains_key(LOCAL_ENVIRONMENT_ID) { + Some(LOCAL_ENVIRONMENT_ID.to_string()) + } else { + None + } + } + DefaultEnvironmentSelection::Environment(environment_id) => { + if !environments.contains_key(&environment_id) { + return Err(ExecServerError::Protocol(format!( + "default environment `{environment_id}` is not configured" + ))); + } + Some(environment_id) + } + DefaultEnvironmentSelection::Disabled => None, }; let local_environment = Arc::new(Environment::local(local_runtime_paths)); let environments = environments @@ -159,11 +177,11 @@ impl EnvironmentManager { .map(|(id, environment)| (id, Arc::new(environment))) .collect(); - Self { + Ok(Self { default_environment, environments, local_environment, - } + }) } /// Returns the default environment instance. @@ -196,6 +214,7 @@ impl EnvironmentManager { #[derive(Clone)] pub struct Environment { exec_server_url: Option, + remote_transport: Option, exec_backend: Arc, filesystem: Arc, http_client: Arc, @@ -207,6 +226,7 @@ impl Environment { pub fn default_for_tests() -> Self { Self { exec_server_url: None, + remote_transport: None, exec_backend: Arc::new(LocalProcess::default()), filesystem: Arc::new(LocalFileSystem::unsandboxed()), http_client: Arc::new(ReqwestHttpClient), @@ -262,6 +282,7 @@ impl Environment { pub(crate) fn local(local_runtime_paths: ExecServerRuntimePaths) -> Self { Self { exec_server_url: None, + remote_transport: None, exec_backend: Arc::new(LocalProcess::default()), filesystem: Arc::new(LocalFileSystem::with_runtime_paths( local_runtime_paths.clone(), @@ -275,15 +296,28 @@ impl Environment { exec_server_url: String, local_runtime_paths: Option, ) -> Self { - let client = LazyRemoteExecServerClient::new(ExecServerTransportParams::WebSocketUrl( - exec_server_url.clone(), - )); + Self::remote_with_transport( + ExecServerTransportParams::WebSocketUrl(exec_server_url), + local_runtime_paths, + ) + } + + fn remote_with_transport( + transport_params: ExecServerTransportParams, + local_runtime_paths: Option, + ) -> Self { + let exec_server_url = match &transport_params { + ExecServerTransportParams::WebSocketUrl(url) => Some(url.clone()), + ExecServerTransportParams::StdioCommand(_) => None, + }; + let client = LazyRemoteExecServerClient::new(transport_params.clone()); let exec_backend: Arc = Arc::new(RemoteProcess::new(client.clone())); let filesystem: Arc = Arc::new(RemoteFileSystem::new(client.clone())); Self { - exec_server_url: Some(exec_server_url), + exec_server_url, + remote_transport: Some(transport_params), exec_backend, filesystem, http_client: Arc::new(client), @@ -292,7 +326,7 @@ impl Environment { } pub fn is_remote(&self) -> bool { - self.exec_server_url.is_some() + self.remote_transport.is_some() } /// Returns the remote exec-server URL when this environment is remote. @@ -322,6 +356,7 @@ mod tests { use std::collections::HashMap; use std::sync::Arc; + use super::DefaultEnvironmentSelection; use super::Environment; use super::EnvironmentManager; use super::LOCAL_ENVIRONMENT_ID; @@ -428,7 +463,9 @@ mod tests { .expect("remote environment"), )]), test_runtime_paths(), - ); + DefaultEnvironmentSelection::Derived, + ) + .expect("environment manager"); assert_eq!( manager.default_environment_id(), @@ -449,6 +486,7 @@ mod tests { let err = EnvironmentManager::from_provider_environments( HashMap::from([("".to_string(), Environment::default_for_tests())]), test_runtime_paths(), + DefaultEnvironmentSelection::Derived, ) .expect_err("empty id should fail"); @@ -458,6 +496,64 @@ mod tests { ); } + #[tokio::test] + async fn environment_manager_uses_explicit_provider_default() { + let manager = EnvironmentManager::from_provider_environments( + HashMap::from([ + ( + LOCAL_ENVIRONMENT_ID.to_string(), + Environment::default_for_tests(), + ), + ( + "devbox".to_string(), + Environment::create_for_tests(Some("ws://127.0.0.1:8765".to_string())) + .expect("remote environment"), + ), + ]), + test_runtime_paths(), + DefaultEnvironmentSelection::Environment("devbox".to_string()), + ) + .expect("manager"); + + assert_eq!(manager.default_environment_id(), Some("devbox")); + assert!(manager.default_environment().expect("default").is_remote()); + } + + #[tokio::test] + async fn environment_manager_disables_provider_default() { + let manager = EnvironmentManager::from_provider_environments( + HashMap::from([( + LOCAL_ENVIRONMENT_ID.to_string(), + Environment::default_for_tests(), + )]), + test_runtime_paths(), + DefaultEnvironmentSelection::Disabled, + ) + .expect("manager"); + + assert_eq!(manager.default_environment_id(), None); + assert!(manager.default_environment().is_none()); + assert!(manager.get_environment(LOCAL_ENVIRONMENT_ID).is_some()); + } + + #[tokio::test] + async fn environment_manager_rejects_unknown_provider_default() { + let err = EnvironmentManager::from_provider_environments( + HashMap::from([( + LOCAL_ENVIRONMENT_ID.to_string(), + Environment::default_for_tests(), + )]), + test_runtime_paths(), + DefaultEnvironmentSelection::Environment("missing".to_string()), + ) + .expect_err("unknown default should fail"); + + assert_eq!( + err.to_string(), + "exec-server protocol error: default environment `missing` is not configured" + ); + } + #[tokio::test] async fn environment_manager_uses_provider_supplied_local_environment() { let manager = EnvironmentManager::create_for_tests( diff --git a/codex-rs/exec-server/src/environment_provider.rs b/codex-rs/exec-server/src/environment_provider.rs index 7c8db07e85e5..eebfff6f139a 100644 --- a/codex-rs/exec-server/src/environment_provider.rs +++ b/codex-rs/exec-server/src/environment_provider.rs @@ -21,6 +21,17 @@ pub trait EnvironmentProvider: Send + Sync { &self, local_runtime_paths: &ExecServerRuntimePaths, ) -> Result, ExecServerError>; + + fn default_environment_selection(&self) -> DefaultEnvironmentSelection { + DefaultEnvironmentSelection::Derived + } +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub enum DefaultEnvironmentSelection { + Derived, + Environment(String), + Disabled, } /// Default provider backed by `CODEX_EXEC_SERVER_URL`. @@ -69,6 +80,14 @@ impl EnvironmentProvider for DefaultEnvironmentProvider { ) -> Result, ExecServerError> { Ok(self.environments(local_runtime_paths)) } + + fn default_environment_selection(&self) -> DefaultEnvironmentSelection { + if normalize_exec_server_url(self.exec_server_url.clone()).1 { + DefaultEnvironmentSelection::Disabled + } else { + DefaultEnvironmentSelection::Derived + } + } } pub(crate) fn normalize_exec_server_url(exec_server_url: Option) -> (Option, bool) { @@ -135,6 +154,10 @@ mod tests { assert!(!environments[LOCAL_ENVIRONMENT_ID].is_remote()); assert!(!environments.contains_key(REMOTE_ENVIRONMENT_ID)); + assert_eq!( + provider.default_environment_selection(), + DefaultEnvironmentSelection::Disabled + ); } #[tokio::test] diff --git a/codex-rs/exec-server/src/lib.rs b/codex-rs/exec-server/src/lib.rs index b36ab39d0105..1fd1ea9c1d77 100644 --- a/codex-rs/exec-server/src/lib.rs +++ b/codex-rs/exec-server/src/lib.rs @@ -42,6 +42,7 @@ pub use environment::EnvironmentManagerArgs; pub use environment::LOCAL_ENVIRONMENT_ID; pub use environment::REMOTE_ENVIRONMENT_ID; pub use environment_provider::DefaultEnvironmentProvider; +pub use environment_provider::DefaultEnvironmentSelection; pub use environment_provider::EnvironmentProvider; pub use fs_helper::CODEX_FS_HELPER_ARG1; pub use fs_helper_main::main as run_fs_helper_main; From 941da7982ee2ba46b0f028a17c3e924893364bf5 Mon Sep 17 00:00:00 2001 From: starr-openai Date: Fri, 1 May 2026 13:07:58 -0700 Subject: [PATCH 29/42] Fix environment manager clippy lints Co-authored-by: Codex --- codex-rs/exec-server/src/environment.rs | 16 ++++++++++------ 1 file changed, 10 insertions(+), 6 deletions(-) diff --git a/codex-rs/exec-server/src/environment.rs b/codex-rs/exec-server/src/environment.rs index 4372f69c1d15..433d51e77793 100644 --- a/codex-rs/exec-server/src/environment.rs +++ b/codex-rs/exec-server/src/environment.rs @@ -73,12 +73,14 @@ impl EnvironmentManager { /// Builds a test-only manager with environment access disabled. pub fn disabled_for_tests(local_runtime_paths: ExecServerRuntimePaths) -> Self { - Self::from_environments( + match Self::from_environments( HashMap::new(), local_runtime_paths, DefaultEnvironmentSelection::Disabled, - ) - .expect("disabled test environment manager") + ) { + Ok(manager) => manager, + Err(err) => panic!("disabled test environment manager: {err}"), + } } /// Builds a test-only manager from a raw exec-server URL value. @@ -105,12 +107,14 @@ impl EnvironmentManager { ) -> Self { let provider = DefaultEnvironmentProvider::new(exec_server_url); let provider_environments = provider.environments(&local_runtime_paths); - Self::from_environments( + match Self::from_environments( provider_environments, local_runtime_paths, provider.default_environment_selection(), - ) - .expect("default provider should create valid environments") + ) { + Ok(manager) => manager, + Err(err) => panic!("default provider should create valid environments: {err}"), + } } /// Builds a manager from a provider-supplied startup snapshot. From 0422edeed3adaa804d8ddc3e2dddc635b324f3ce Mon Sep 17 00:00:00 2001 From: starr-openai Date: Mon, 4 May 2026 12:37:38 -0700 Subject: [PATCH 30/42] Simplify provider default environment selection Have providers return a concrete default environment id after constructing their environment map, using None to disable the default. This removes the DefaultEnvironmentSelection tri-state while preserving legacy derived defaults through the trait's default implementation. Co-authored-by: Codex --- codex-rs/exec-server/src/environment.rs | 60 +++++++------------ .../exec-server/src/environment_provider.rs | 38 +++++++----- codex-rs/exec-server/src/lib.rs | 1 - 3 files changed, 43 insertions(+), 56 deletions(-) diff --git a/codex-rs/exec-server/src/environment.rs b/codex-rs/exec-server/src/environment.rs index 433d51e77793..b36ed33f3f21 100644 --- a/codex-rs/exec-server/src/environment.rs +++ b/codex-rs/exec-server/src/environment.rs @@ -9,7 +9,6 @@ use crate::client::LazyRemoteExecServerClient; use crate::client::http_client::ReqwestHttpClient; use crate::client_api::ExecServerTransportParams; use crate::environment_provider::DefaultEnvironmentProvider; -use crate::environment_provider::DefaultEnvironmentSelection; use crate::environment_provider::EnvironmentProvider; use crate::environment_provider::normalize_exec_server_url; use crate::local_file_system::LocalFileSystem; @@ -76,7 +75,7 @@ impl EnvironmentManager { match Self::from_environments( HashMap::new(), local_runtime_paths, - DefaultEnvironmentSelection::Disabled, + /*default_environment*/ None, ) { Ok(manager) => manager, Err(err) => panic!("disabled test environment manager: {err}"), @@ -107,10 +106,11 @@ impl EnvironmentManager { ) -> Self { let provider = DefaultEnvironmentProvider::new(exec_server_url); let provider_environments = provider.environments(&local_runtime_paths); + let default_environment = provider.default_environment_id(&provider_environments); match Self::from_environments( provider_environments, local_runtime_paths, - provider.default_environment_selection(), + default_environment, ) { Ok(manager) => manager, Err(err) => panic!("default provider should create valid environments: {err}"), @@ -125,17 +125,15 @@ impl EnvironmentManager { where P: EnvironmentProvider + ?Sized, { - Self::from_provider_environments( - provider.get_environments(&local_runtime_paths).await?, - local_runtime_paths, - provider.default_environment_selection(), - ) + let environments = provider.get_environments(&local_runtime_paths).await?; + let default_environment = provider.default_environment_id(&environments); + Self::from_provider_environments(environments, local_runtime_paths, default_environment) } fn from_provider_environments( environments: HashMap, local_runtime_paths: ExecServerRuntimePaths, - default_selection: DefaultEnvironmentSelection, + default_environment: Option, ) -> Result { for id in environments.keys() { if id.is_empty() { @@ -145,36 +143,21 @@ impl EnvironmentManager { } } - Self::from_environments(environments, local_runtime_paths, default_selection) + Self::from_environments(environments, local_runtime_paths, default_environment) } fn from_environments( environments: HashMap, local_runtime_paths: ExecServerRuntimePaths, - default_selection: DefaultEnvironmentSelection, + default_environment: Option, ) -> Result { - // TODO: Stop deriving a default environment here once omitted - // environment attachment is owned by thread/session setup. - let default_environment = match default_selection { - DefaultEnvironmentSelection::Derived => { - if environments.contains_key(REMOTE_ENVIRONMENT_ID) { - Some(REMOTE_ENVIRONMENT_ID.to_string()) - } else if environments.contains_key(LOCAL_ENVIRONMENT_ID) { - Some(LOCAL_ENVIRONMENT_ID.to_string()) - } else { - None - } - } - DefaultEnvironmentSelection::Environment(environment_id) => { - if !environments.contains_key(&environment_id) { - return Err(ExecServerError::Protocol(format!( - "default environment `{environment_id}` is not configured" - ))); - } - Some(environment_id) - } - DefaultEnvironmentSelection::Disabled => None, - }; + if let Some(environment_id) = default_environment.as_ref() + && !environments.contains_key(environment_id) + { + return Err(ExecServerError::Protocol(format!( + "default environment `{environment_id}` is not configured" + ))); + } let local_environment = Arc::new(Environment::local(local_runtime_paths)); let environments = environments .into_iter() @@ -360,7 +343,6 @@ mod tests { use std::collections::HashMap; use std::sync::Arc; - use super::DefaultEnvironmentSelection; use super::Environment; use super::EnvironmentManager; use super::LOCAL_ENVIRONMENT_ID; @@ -467,7 +449,7 @@ mod tests { .expect("remote environment"), )]), test_runtime_paths(), - DefaultEnvironmentSelection::Derived, + Some(REMOTE_ENVIRONMENT_ID.to_string()), ) .expect("environment manager"); @@ -490,7 +472,7 @@ mod tests { let err = EnvironmentManager::from_provider_environments( HashMap::from([("".to_string(), Environment::default_for_tests())]), test_runtime_paths(), - DefaultEnvironmentSelection::Derived, + /*default_environment*/ None, ) .expect_err("empty id should fail"); @@ -515,7 +497,7 @@ mod tests { ), ]), test_runtime_paths(), - DefaultEnvironmentSelection::Environment("devbox".to_string()), + Some("devbox".to_string()), ) .expect("manager"); @@ -531,7 +513,7 @@ mod tests { Environment::default_for_tests(), )]), test_runtime_paths(), - DefaultEnvironmentSelection::Disabled, + /*default_environment*/ None, ) .expect("manager"); @@ -548,7 +530,7 @@ mod tests { Environment::default_for_tests(), )]), test_runtime_paths(), - DefaultEnvironmentSelection::Environment("missing".to_string()), + Some("missing".to_string()), ) .expect_err("unknown default should fail"); diff --git a/codex-rs/exec-server/src/environment_provider.rs b/codex-rs/exec-server/src/environment_provider.rs index eebfff6f139a..1f85780622f7 100644 --- a/codex-rs/exec-server/src/environment_provider.rs +++ b/codex-rs/exec-server/src/environment_provider.rs @@ -22,18 +22,14 @@ pub trait EnvironmentProvider: Send + Sync { local_runtime_paths: &ExecServerRuntimePaths, ) -> Result, ExecServerError>; - fn default_environment_selection(&self) -> DefaultEnvironmentSelection { - DefaultEnvironmentSelection::Derived + fn default_environment_id( + &self, + environments: &HashMap, + ) -> Option { + derived_default_environment_id(environments) } } -#[derive(Clone, Debug, PartialEq, Eq)] -pub enum DefaultEnvironmentSelection { - Derived, - Environment(String), - Disabled, -} - /// Default provider backed by `CODEX_EXEC_SERVER_URL`. #[derive(Clone, Debug)] pub struct DefaultEnvironmentProvider { @@ -81,15 +77,28 @@ impl EnvironmentProvider for DefaultEnvironmentProvider { Ok(self.environments(local_runtime_paths)) } - fn default_environment_selection(&self) -> DefaultEnvironmentSelection { + fn default_environment_id( + &self, + environments: &HashMap, + ) -> Option { if normalize_exec_server_url(self.exec_server_url.clone()).1 { - DefaultEnvironmentSelection::Disabled + None } else { - DefaultEnvironmentSelection::Derived + derived_default_environment_id(environments) } } } +fn derived_default_environment_id(environments: &HashMap) -> Option { + if environments.contains_key(REMOTE_ENVIRONMENT_ID) { + Some(REMOTE_ENVIRONMENT_ID.to_string()) + } else if environments.contains_key(LOCAL_ENVIRONMENT_ID) { + Some(LOCAL_ENVIRONMENT_ID.to_string()) + } else { + None + } +} + pub(crate) fn normalize_exec_server_url(exec_server_url: Option) -> (Option, bool) { match exec_server_url.as_deref().map(str::trim) { None | Some("") => (None, false), @@ -154,10 +163,7 @@ mod tests { assert!(!environments[LOCAL_ENVIRONMENT_ID].is_remote()); assert!(!environments.contains_key(REMOTE_ENVIRONMENT_ID)); - assert_eq!( - provider.default_environment_selection(), - DefaultEnvironmentSelection::Disabled - ); + assert_eq!(provider.default_environment_id(&environments), None); } #[tokio::test] diff --git a/codex-rs/exec-server/src/lib.rs b/codex-rs/exec-server/src/lib.rs index 1fd1ea9c1d77..b36ab39d0105 100644 --- a/codex-rs/exec-server/src/lib.rs +++ b/codex-rs/exec-server/src/lib.rs @@ -42,7 +42,6 @@ pub use environment::EnvironmentManagerArgs; pub use environment::LOCAL_ENVIRONMENT_ID; pub use environment::REMOTE_ENVIRONMENT_ID; pub use environment_provider::DefaultEnvironmentProvider; -pub use environment_provider::DefaultEnvironmentSelection; pub use environment_provider::EnvironmentProvider; pub use fs_helper::CODEX_FS_HELPER_ARG1; pub use fs_helper_main::main as run_fs_helper_main; From 6e03a7ced989b867d827f3db43288de54e1c671f Mon Sep 17 00:00:00 2001 From: starr-openai Date: Mon, 4 May 2026 12:46:29 -0700 Subject: [PATCH 31/42] Return provider environment snapshots Make environment providers return the environment map and default id together. This keeps provider-owned startup state in one boundary and removes the separate default callback over a map. Co-authored-by: Codex --- codex-rs/exec-server/src/environment.rs | 130 +++++++++--------- .../exec-server/src/environment_provider.rs | 85 ++++++------ 2 files changed, 109 insertions(+), 106 deletions(-) diff --git a/codex-rs/exec-server/src/environment.rs b/codex-rs/exec-server/src/environment.rs index b36ed33f3f21..9fec3372f4e0 100644 --- a/codex-rs/exec-server/src/environment.rs +++ b/codex-rs/exec-server/src/environment.rs @@ -10,6 +10,7 @@ use crate::client::http_client::ReqwestHttpClient; use crate::client_api::ExecServerTransportParams; use crate::environment_provider::DefaultEnvironmentProvider; use crate::environment_provider::EnvironmentProvider; +use crate::environment_provider::EnvironmentProviderSnapshot; use crate::environment_provider::normalize_exec_server_url; use crate::local_file_system::LocalFileSystem; use crate::local_process::LocalProcess; @@ -72,11 +73,11 @@ impl EnvironmentManager { /// Builds a test-only manager with environment access disabled. pub fn disabled_for_tests(local_runtime_paths: ExecServerRuntimePaths) -> Self { - match Self::from_environments( - HashMap::new(), - local_runtime_paths, - /*default_environment*/ None, - ) { + let snapshot = EnvironmentProviderSnapshot { + environments: HashMap::new(), + default_environment_id: None, + }; + match Self::from_provider_snapshot(snapshot, local_runtime_paths) { Ok(manager) => manager, Err(err) => panic!("disabled test environment manager: {err}"), } @@ -105,13 +106,8 @@ impl EnvironmentManager { local_runtime_paths: ExecServerRuntimePaths, ) -> Self { let provider = DefaultEnvironmentProvider::new(exec_server_url); - let provider_environments = provider.environments(&local_runtime_paths); - let default_environment = provider.default_environment_id(&provider_environments); - match Self::from_environments( - provider_environments, - local_runtime_paths, - default_environment, - ) { + let snapshot = provider.snapshot(&local_runtime_paths); + match Self::from_provider_snapshot(snapshot, local_runtime_paths) { Ok(manager) => manager, Err(err) => panic!("default provider should create valid environments: {err}"), } @@ -125,16 +121,20 @@ impl EnvironmentManager { where P: EnvironmentProvider + ?Sized, { - let environments = provider.get_environments(&local_runtime_paths).await?; - let default_environment = provider.default_environment_id(&environments); - Self::from_provider_environments(environments, local_runtime_paths, default_environment) + let snapshot = provider + .get_environment_snapshot(&local_runtime_paths) + .await?; + Self::from_provider_snapshot(snapshot, local_runtime_paths) } - fn from_provider_environments( - environments: HashMap, + fn from_provider_snapshot( + snapshot: EnvironmentProviderSnapshot, local_runtime_paths: ExecServerRuntimePaths, - default_environment: Option, ) -> Result { + let EnvironmentProviderSnapshot { + environments, + default_environment_id, + } = snapshot; for id in environments.keys() { if id.is_empty() { return Err(ExecServerError::Protocol( @@ -143,15 +143,7 @@ impl EnvironmentManager { } } - Self::from_environments(environments, local_runtime_paths, default_environment) - } - - fn from_environments( - environments: HashMap, - local_runtime_paths: ExecServerRuntimePaths, - default_environment: Option, - ) -> Result { - if let Some(environment_id) = default_environment.as_ref() + if let Some(environment_id) = default_environment_id.as_ref() && !environments.contains_key(environment_id) { return Err(ExecServerError::Protocol(format!( @@ -165,7 +157,7 @@ impl EnvironmentManager { .collect(); Ok(Self { - default_environment, + default_environment: default_environment_id, environments, local_environment, }) @@ -441,15 +433,17 @@ mod tests { } #[tokio::test] - async fn environment_manager_builds_from_provider_environments() { - let manager = EnvironmentManager::from_environments( - HashMap::from([( - REMOTE_ENVIRONMENT_ID.to_string(), - Environment::create_for_tests(Some("ws://127.0.0.1:8765".to_string())) - .expect("remote environment"), - )]), + async fn environment_manager_builds_from_provider_snapshot() { + let manager = EnvironmentManager::from_provider_snapshot( + EnvironmentProviderSnapshot { + environments: HashMap::from([( + REMOTE_ENVIRONMENT_ID.to_string(), + Environment::create_for_tests(Some("ws://127.0.0.1:8765".to_string())) + .expect("remote environment"), + )]), + default_environment_id: Some(REMOTE_ENVIRONMENT_ID.to_string()), + }, test_runtime_paths(), - Some(REMOTE_ENVIRONMENT_ID.to_string()), ) .expect("environment manager"); @@ -469,10 +463,12 @@ mod tests { #[tokio::test] async fn environment_manager_rejects_empty_environment_id() { - let err = EnvironmentManager::from_provider_environments( - HashMap::from([("".to_string(), Environment::default_for_tests())]), + let err = EnvironmentManager::from_provider_snapshot( + EnvironmentProviderSnapshot { + environments: HashMap::from([("".to_string(), Environment::default_for_tests())]), + default_environment_id: None, + }, test_runtime_paths(), - /*default_environment*/ None, ) .expect_err("empty id should fail"); @@ -484,20 +480,22 @@ mod tests { #[tokio::test] async fn environment_manager_uses_explicit_provider_default() { - let manager = EnvironmentManager::from_provider_environments( - HashMap::from([ - ( - LOCAL_ENVIRONMENT_ID.to_string(), - Environment::default_for_tests(), - ), - ( - "devbox".to_string(), - Environment::create_for_tests(Some("ws://127.0.0.1:8765".to_string())) - .expect("remote environment"), - ), - ]), + let manager = EnvironmentManager::from_provider_snapshot( + EnvironmentProviderSnapshot { + environments: HashMap::from([ + ( + LOCAL_ENVIRONMENT_ID.to_string(), + Environment::default_for_tests(), + ), + ( + "devbox".to_string(), + Environment::create_for_tests(Some("ws://127.0.0.1:8765".to_string())) + .expect("remote environment"), + ), + ]), + default_environment_id: Some("devbox".to_string()), + }, test_runtime_paths(), - Some("devbox".to_string()), ) .expect("manager"); @@ -507,13 +505,15 @@ mod tests { #[tokio::test] async fn environment_manager_disables_provider_default() { - let manager = EnvironmentManager::from_provider_environments( - HashMap::from([( - LOCAL_ENVIRONMENT_ID.to_string(), - Environment::default_for_tests(), - )]), + let manager = EnvironmentManager::from_provider_snapshot( + EnvironmentProviderSnapshot { + environments: HashMap::from([( + LOCAL_ENVIRONMENT_ID.to_string(), + Environment::default_for_tests(), + )]), + default_environment_id: None, + }, test_runtime_paths(), - /*default_environment*/ None, ) .expect("manager"); @@ -524,13 +524,15 @@ mod tests { #[tokio::test] async fn environment_manager_rejects_unknown_provider_default() { - let err = EnvironmentManager::from_provider_environments( - HashMap::from([( - LOCAL_ENVIRONMENT_ID.to_string(), - Environment::default_for_tests(), - )]), + let err = EnvironmentManager::from_provider_snapshot( + EnvironmentProviderSnapshot { + environments: HashMap::from([( + LOCAL_ENVIRONMENT_ID.to_string(), + Environment::default_for_tests(), + )]), + default_environment_id: Some("missing".to_string()), + }, test_runtime_paths(), - Some("missing".to_string()), ) .expect_err("unknown default should fail"); diff --git a/codex-rs/exec-server/src/environment_provider.rs b/codex-rs/exec-server/src/environment_provider.rs index 1f85780622f7..39226cd0d1e0 100644 --- a/codex-rs/exec-server/src/environment_provider.rs +++ b/codex-rs/exec-server/src/environment_provider.rs @@ -17,17 +17,15 @@ use crate::environment::REMOTE_ENVIRONMENT_ID; #[async_trait] pub trait EnvironmentProvider: Send + Sync { /// Returns the environments available for a new manager. - async fn get_environments( + async fn get_environment_snapshot( &self, local_runtime_paths: &ExecServerRuntimePaths, - ) -> Result, ExecServerError>; + ) -> Result; +} - fn default_environment_id( - &self, - environments: &HashMap, - ) -> Option { - derived_default_environment_id(environments) - } +pub struct EnvironmentProviderSnapshot { + pub environments: HashMap, + pub default_environment_id: Option, } /// Default provider backed by `CODEX_EXEC_SERVER_URL`. @@ -47,15 +45,15 @@ impl DefaultEnvironmentProvider { Self::new(std::env::var(CODEX_EXEC_SERVER_URL_ENV_VAR).ok()) } - pub(crate) fn environments( + pub(crate) fn snapshot( &self, local_runtime_paths: &ExecServerRuntimePaths, - ) -> HashMap { + ) -> EnvironmentProviderSnapshot { let mut environments = HashMap::from([( LOCAL_ENVIRONMENT_ID.to_string(), Environment::local(local_runtime_paths.clone()), )]); - let exec_server_url = normalize_exec_server_url(self.exec_server_url.clone()).0; + let (exec_server_url, disabled) = normalize_exec_server_url(self.exec_server_url.clone()); if let Some(exec_server_url) = exec_server_url { environments.insert( @@ -64,28 +62,26 @@ impl DefaultEnvironmentProvider { ); } - environments + let default_environment_id = if disabled { + None + } else { + derived_default_environment_id(&environments) + }; + + EnvironmentProviderSnapshot { + environments, + default_environment_id, + } } } #[async_trait] impl EnvironmentProvider for DefaultEnvironmentProvider { - async fn get_environments( + async fn get_environment_snapshot( &self, local_runtime_paths: &ExecServerRuntimePaths, - ) -> Result, ExecServerError> { - Ok(self.environments(local_runtime_paths)) - } - - fn default_environment_id( - &self, - environments: &HashMap, - ) -> Option { - if normalize_exec_server_url(self.exec_server_url.clone()).1 { - None - } else { - derived_default_environment_id(environments) - } + ) -> Result { + Ok(self.snapshot(local_runtime_paths)) } } @@ -126,10 +122,11 @@ mod tests { async fn default_provider_returns_local_environment_when_url_is_missing() { let provider = DefaultEnvironmentProvider::new(/*exec_server_url*/ None); let runtime_paths = test_runtime_paths(); - let environments = provider - .get_environments(&runtime_paths) + let snapshot = provider + .get_environment_snapshot(&runtime_paths) .await - .expect("environments"); + .expect("environment snapshot"); + let environments = snapshot.environments; assert!(!environments[LOCAL_ENVIRONMENT_ID].is_remote()); assert_eq!( @@ -143,10 +140,11 @@ mod tests { async fn default_provider_returns_local_environment_when_url_is_empty() { let provider = DefaultEnvironmentProvider::new(Some(String::new())); let runtime_paths = test_runtime_paths(); - let environments = provider - .get_environments(&runtime_paths) + let snapshot = provider + .get_environment_snapshot(&runtime_paths) .await - .expect("environments"); + .expect("environment snapshot"); + let environments = snapshot.environments; assert!(!environments[LOCAL_ENVIRONMENT_ID].is_remote()); assert!(!environments.contains_key(REMOTE_ENVIRONMENT_ID)); @@ -156,24 +154,26 @@ mod tests { async fn default_provider_returns_local_environment_for_none_value() { let provider = DefaultEnvironmentProvider::new(Some("none".to_string())); let runtime_paths = test_runtime_paths(); - let environments = provider - .get_environments(&runtime_paths) + let snapshot = provider + .get_environment_snapshot(&runtime_paths) .await - .expect("environments"); + .expect("environment snapshot"); + let environments = snapshot.environments; assert!(!environments[LOCAL_ENVIRONMENT_ID].is_remote()); assert!(!environments.contains_key(REMOTE_ENVIRONMENT_ID)); - assert_eq!(provider.default_environment_id(&environments), None); + assert_eq!(snapshot.default_environment_id, None); } #[tokio::test] async fn default_provider_adds_remote_environment_for_websocket_url() { let provider = DefaultEnvironmentProvider::new(Some("ws://127.0.0.1:8765".to_string())); let runtime_paths = test_runtime_paths(); - let environments = provider - .get_environments(&runtime_paths) + let snapshot = provider + .get_environment_snapshot(&runtime_paths) .await - .expect("environments"); + .expect("environment snapshot"); + let environments = snapshot.environments; assert!(!environments[LOCAL_ENVIRONMENT_ID].is_remote()); let remote_environment = &environments[REMOTE_ENVIRONMENT_ID]; @@ -188,10 +188,11 @@ mod tests { async fn default_provider_normalizes_exec_server_url() { let provider = DefaultEnvironmentProvider::new(Some(" ws://127.0.0.1:8765 ".to_string())); let runtime_paths = test_runtime_paths(); - let environments = provider - .get_environments(&runtime_paths) + let snapshot = provider + .get_environment_snapshot(&runtime_paths) .await - .expect("environments"); + .expect("environment snapshot"); + let environments = snapshot.environments; assert_eq!( environments[REMOTE_ENVIRONMENT_ID].exec_server_url(), From 4b9daddeea13eb6fda1e4305df2d77f34e042773 Mon Sep 17 00:00:00 2001 From: starr-openai Date: Mon, 4 May 2026 13:03:44 -0700 Subject: [PATCH 32/42] Split provider environments from default id Remove the EnvironmentProviderSnapshot wrapper. Providers now expose environments and the selected default id directly, while EnvironmentManager validates that the default id exists in the returned environment map. Co-authored-by: Codex --- codex-rs/exec-server/src/environment.rs | 117 ++++++++---------- .../exec-server/src/environment_provider.rs | 92 ++++++-------- 2 files changed, 92 insertions(+), 117 deletions(-) diff --git a/codex-rs/exec-server/src/environment.rs b/codex-rs/exec-server/src/environment.rs index 9fec3372f4e0..dc4bd6498901 100644 --- a/codex-rs/exec-server/src/environment.rs +++ b/codex-rs/exec-server/src/environment.rs @@ -10,7 +10,6 @@ use crate::client::http_client::ReqwestHttpClient; use crate::client_api::ExecServerTransportParams; use crate::environment_provider::DefaultEnvironmentProvider; use crate::environment_provider::EnvironmentProvider; -use crate::environment_provider::EnvironmentProviderSnapshot; use crate::environment_provider::normalize_exec_server_url; use crate::local_file_system::LocalFileSystem; use crate::local_process::LocalProcess; @@ -24,7 +23,8 @@ pub const CODEX_EXEC_SERVER_URL_ENV_VAR: &str = "CODEX_EXEC_SERVER_URL"; /// /// `EnvironmentManager` is a shared registry for concrete environments. Its /// default constructor preserves the legacy `CODEX_EXEC_SERVER_URL` behavior -/// while provider-based construction accepts a provider-supplied snapshot. +/// while provider-based construction accepts a provider-supplied environment +/// list and default id. /// /// Setting `CODEX_EXEC_SERVER_URL=none` disables environment access by leaving /// the default environment unset while still keeping an explicit local @@ -73,11 +73,7 @@ impl EnvironmentManager { /// Builds a test-only manager with environment access disabled. pub fn disabled_for_tests(local_runtime_paths: ExecServerRuntimePaths) -> Self { - let snapshot = EnvironmentProviderSnapshot { - environments: HashMap::new(), - default_environment_id: None, - }; - match Self::from_provider_snapshot(snapshot, local_runtime_paths) { + match Self::from_provider_parts(HashMap::new(), None, local_runtime_paths) { Ok(manager) => manager, Err(err) => panic!("disabled test environment manager: {err}"), } @@ -106,14 +102,17 @@ impl EnvironmentManager { local_runtime_paths: ExecServerRuntimePaths, ) -> Self { let provider = DefaultEnvironmentProvider::new(exec_server_url); - let snapshot = provider.snapshot(&local_runtime_paths); - match Self::from_provider_snapshot(snapshot, local_runtime_paths) { + match Self::from_provider_parts( + provider.environments(&local_runtime_paths), + provider.default_id(), + local_runtime_paths, + ) { Ok(manager) => manager, Err(err) => panic!("default provider should create valid environments: {err}"), } } - /// Builds a manager from a provider-supplied startup snapshot. + /// Builds a manager from a provider-supplied environment list and default. pub async fn from_provider

( provider: &P, local_runtime_paths: ExecServerRuntimePaths, @@ -121,20 +120,16 @@ impl EnvironmentManager { where P: EnvironmentProvider + ?Sized, { - let snapshot = provider - .get_environment_snapshot(&local_runtime_paths) - .await?; - Self::from_provider_snapshot(snapshot, local_runtime_paths) + let environments = provider.get_environments(&local_runtime_paths).await?; + let default_environment_id = provider.default_environment_id(); + Self::from_provider_parts(environments, default_environment_id, local_runtime_paths) } - fn from_provider_snapshot( - snapshot: EnvironmentProviderSnapshot, + fn from_provider_parts( + environments: HashMap, + default_environment_id: Option, local_runtime_paths: ExecServerRuntimePaths, ) -> Result { - let EnvironmentProviderSnapshot { - environments, - default_environment_id, - } = snapshot; for id in environments.keys() { if id.is_empty() { return Err(ExecServerError::Protocol( @@ -433,16 +428,14 @@ mod tests { } #[tokio::test] - async fn environment_manager_builds_from_provider_snapshot() { - let manager = EnvironmentManager::from_provider_snapshot( - EnvironmentProviderSnapshot { - environments: HashMap::from([( - REMOTE_ENVIRONMENT_ID.to_string(), - Environment::create_for_tests(Some("ws://127.0.0.1:8765".to_string())) - .expect("remote environment"), - )]), - default_environment_id: Some(REMOTE_ENVIRONMENT_ID.to_string()), - }, + async fn environment_manager_builds_from_provider_parts() { + let manager = EnvironmentManager::from_provider_parts( + HashMap::from([( + REMOTE_ENVIRONMENT_ID.to_string(), + Environment::create_for_tests(Some("ws://127.0.0.1:8765".to_string())) + .expect("remote environment"), + )]), + Some(REMOTE_ENVIRONMENT_ID.to_string()), test_runtime_paths(), ) .expect("environment manager"); @@ -463,11 +456,9 @@ mod tests { #[tokio::test] async fn environment_manager_rejects_empty_environment_id() { - let err = EnvironmentManager::from_provider_snapshot( - EnvironmentProviderSnapshot { - environments: HashMap::from([("".to_string(), Environment::default_for_tests())]), - default_environment_id: None, - }, + let err = EnvironmentManager::from_provider_parts( + HashMap::from([("".to_string(), Environment::default_for_tests())]), + None, test_runtime_paths(), ) .expect_err("empty id should fail"); @@ -480,21 +471,19 @@ mod tests { #[tokio::test] async fn environment_manager_uses_explicit_provider_default() { - let manager = EnvironmentManager::from_provider_snapshot( - EnvironmentProviderSnapshot { - environments: HashMap::from([ - ( - LOCAL_ENVIRONMENT_ID.to_string(), - Environment::default_for_tests(), - ), - ( - "devbox".to_string(), - Environment::create_for_tests(Some("ws://127.0.0.1:8765".to_string())) - .expect("remote environment"), - ), - ]), - default_environment_id: Some("devbox".to_string()), - }, + let manager = EnvironmentManager::from_provider_parts( + HashMap::from([ + ( + LOCAL_ENVIRONMENT_ID.to_string(), + Environment::default_for_tests(), + ), + ( + "devbox".to_string(), + Environment::create_for_tests(Some("ws://127.0.0.1:8765".to_string())) + .expect("remote environment"), + ), + ]), + Some("devbox".to_string()), test_runtime_paths(), ) .expect("manager"); @@ -505,14 +494,12 @@ mod tests { #[tokio::test] async fn environment_manager_disables_provider_default() { - let manager = EnvironmentManager::from_provider_snapshot( - EnvironmentProviderSnapshot { - environments: HashMap::from([( - LOCAL_ENVIRONMENT_ID.to_string(), - Environment::default_for_tests(), - )]), - default_environment_id: None, - }, + let manager = EnvironmentManager::from_provider_parts( + HashMap::from([( + LOCAL_ENVIRONMENT_ID.to_string(), + Environment::default_for_tests(), + )]), + None, test_runtime_paths(), ) .expect("manager"); @@ -524,14 +511,12 @@ mod tests { #[tokio::test] async fn environment_manager_rejects_unknown_provider_default() { - let err = EnvironmentManager::from_provider_snapshot( - EnvironmentProviderSnapshot { - environments: HashMap::from([( - LOCAL_ENVIRONMENT_ID.to_string(), - Environment::default_for_tests(), - )]), - default_environment_id: Some("missing".to_string()), - }, + let err = EnvironmentManager::from_provider_parts( + HashMap::from([( + LOCAL_ENVIRONMENT_ID.to_string(), + Environment::default_for_tests(), + )]), + Some("missing".to_string()), test_runtime_paths(), ) .expect_err("unknown default should fail"); diff --git a/codex-rs/exec-server/src/environment_provider.rs b/codex-rs/exec-server/src/environment_provider.rs index 39226cd0d1e0..57f4340498d8 100644 --- a/codex-rs/exec-server/src/environment_provider.rs +++ b/codex-rs/exec-server/src/environment_provider.rs @@ -11,21 +11,20 @@ use crate::environment::REMOTE_ENVIRONMENT_ID; /// Lists the concrete environments available to Codex. /// -/// Implementations should return the provider-owned startup snapshot that -/// `EnvironmentManager` will cache. Providers that want the local environment to -/// be addressable by id should include it explicitly in the returned map. +/// Implementations own both the available environment list and the default +/// environment id. Providers that want the local environment to be addressable +/// by id should include it explicitly in the returned map. #[async_trait] pub trait EnvironmentProvider: Send + Sync { /// Returns the environments available for a new manager. - async fn get_environment_snapshot( + async fn get_environments( &self, local_runtime_paths: &ExecServerRuntimePaths, - ) -> Result; -} + ) -> Result, ExecServerError>; -pub struct EnvironmentProviderSnapshot { - pub environments: HashMap, - pub default_environment_id: Option, + /// Returns the provider-selected default environment id, or `None` to + /// disable model-facing environment access. + fn default_environment_id(&self) -> Option; } /// Default provider backed by `CODEX_EXEC_SERVER_URL`. @@ -45,15 +44,15 @@ impl DefaultEnvironmentProvider { Self::new(std::env::var(CODEX_EXEC_SERVER_URL_ENV_VAR).ok()) } - pub(crate) fn snapshot( + pub(crate) fn environments( &self, local_runtime_paths: &ExecServerRuntimePaths, - ) -> EnvironmentProviderSnapshot { + ) -> HashMap { let mut environments = HashMap::from([( LOCAL_ENVIRONMENT_ID.to_string(), Environment::local(local_runtime_paths.clone()), )]); - let (exec_server_url, disabled) = normalize_exec_server_url(self.exec_server_url.clone()); + let (exec_server_url, _disabled) = normalize_exec_server_url(self.exec_server_url.clone()); if let Some(exec_server_url) = exec_server_url { environments.insert( @@ -62,36 +61,32 @@ impl DefaultEnvironmentProvider { ); } - let default_environment_id = if disabled { + environments + } + + pub(crate) fn default_id(&self) -> Option { + let (exec_server_url, disabled) = normalize_exec_server_url(self.exec_server_url.clone()); + if disabled { None + } else if exec_server_url.is_some() { + Some(REMOTE_ENVIRONMENT_ID.to_string()) } else { - derived_default_environment_id(&environments) - }; - - EnvironmentProviderSnapshot { - environments, - default_environment_id, + Some(LOCAL_ENVIRONMENT_ID.to_string()) } } } #[async_trait] impl EnvironmentProvider for DefaultEnvironmentProvider { - async fn get_environment_snapshot( + async fn get_environments( &self, local_runtime_paths: &ExecServerRuntimePaths, - ) -> Result { - Ok(self.snapshot(local_runtime_paths)) + ) -> Result, ExecServerError> { + Ok(self.environments(local_runtime_paths)) } -} -fn derived_default_environment_id(environments: &HashMap) -> Option { - if environments.contains_key(REMOTE_ENVIRONMENT_ID) { - Some(REMOTE_ENVIRONMENT_ID.to_string()) - } else if environments.contains_key(LOCAL_ENVIRONMENT_ID) { - Some(LOCAL_ENVIRONMENT_ID.to_string()) - } else { - None + fn default_environment_id(&self) -> Option { + self.default_id() } } @@ -122,11 +117,10 @@ mod tests { async fn default_provider_returns_local_environment_when_url_is_missing() { let provider = DefaultEnvironmentProvider::new(/*exec_server_url*/ None); let runtime_paths = test_runtime_paths(); - let snapshot = provider - .get_environment_snapshot(&runtime_paths) + let environments = provider + .get_environments(&runtime_paths) .await - .expect("environment snapshot"); - let environments = snapshot.environments; + .expect("environments"); assert!(!environments[LOCAL_ENVIRONMENT_ID].is_remote()); assert_eq!( @@ -140,11 +134,10 @@ mod tests { async fn default_provider_returns_local_environment_when_url_is_empty() { let provider = DefaultEnvironmentProvider::new(Some(String::new())); let runtime_paths = test_runtime_paths(); - let snapshot = provider - .get_environment_snapshot(&runtime_paths) + let environments = provider + .get_environments(&runtime_paths) .await - .expect("environment snapshot"); - let environments = snapshot.environments; + .expect("environments"); assert!(!environments[LOCAL_ENVIRONMENT_ID].is_remote()); assert!(!environments.contains_key(REMOTE_ENVIRONMENT_ID)); @@ -154,26 +147,24 @@ mod tests { async fn default_provider_returns_local_environment_for_none_value() { let provider = DefaultEnvironmentProvider::new(Some("none".to_string())); let runtime_paths = test_runtime_paths(); - let snapshot = provider - .get_environment_snapshot(&runtime_paths) + let environments = provider + .get_environments(&runtime_paths) .await - .expect("environment snapshot"); - let environments = snapshot.environments; + .expect("environments"); assert!(!environments[LOCAL_ENVIRONMENT_ID].is_remote()); assert!(!environments.contains_key(REMOTE_ENVIRONMENT_ID)); - assert_eq!(snapshot.default_environment_id, None); + assert_eq!(provider.default_environment_id(), None); } #[tokio::test] async fn default_provider_adds_remote_environment_for_websocket_url() { let provider = DefaultEnvironmentProvider::new(Some("ws://127.0.0.1:8765".to_string())); let runtime_paths = test_runtime_paths(); - let snapshot = provider - .get_environment_snapshot(&runtime_paths) + let environments = provider + .get_environments(&runtime_paths) .await - .expect("environment snapshot"); - let environments = snapshot.environments; + .expect("environments"); assert!(!environments[LOCAL_ENVIRONMENT_ID].is_remote()); let remote_environment = &environments[REMOTE_ENVIRONMENT_ID]; @@ -188,11 +179,10 @@ mod tests { async fn default_provider_normalizes_exec_server_url() { let provider = DefaultEnvironmentProvider::new(Some(" ws://127.0.0.1:8765 ".to_string())); let runtime_paths = test_runtime_paths(); - let snapshot = provider - .get_environment_snapshot(&runtime_paths) + let environments = provider + .get_environments(&runtime_paths) .await - .expect("environment snapshot"); - let environments = snapshot.environments; + .expect("environments"); assert_eq!( environments[REMOTE_ENVIRONMENT_ID].exec_server_url(), From ac065c320b2b13dcc726b30f5cf8afa6514c60ce Mon Sep 17 00:00:00 2001 From: starr-openai Date: Mon, 4 May 2026 13:12:50 -0700 Subject: [PATCH 33/42] Inline provider manager construction Remove the private from_provider_parts helper. EnvironmentManager::from_provider now performs the provider read, validation, and manager construction directly, and tests use a small provider implementation instead of bypassing that path. Co-authored-by: Codex --- codex-rs/exec-server/src/environment.rs | 111 ++++++++++++++---------- 1 file changed, 64 insertions(+), 47 deletions(-) diff --git a/codex-rs/exec-server/src/environment.rs b/codex-rs/exec-server/src/environment.rs index dc4bd6498901..3c70374bade3 100644 --- a/codex-rs/exec-server/src/environment.rs +++ b/codex-rs/exec-server/src/environment.rs @@ -73,9 +73,10 @@ impl EnvironmentManager { /// Builds a test-only manager with environment access disabled. pub fn disabled_for_tests(local_runtime_paths: ExecServerRuntimePaths) -> Self { - match Self::from_provider_parts(HashMap::new(), None, local_runtime_paths) { - Ok(manager) => manager, - Err(err) => panic!("disabled test environment manager: {err}"), + Self { + default_environment: None, + environments: HashMap::new(), + local_environment: Arc::new(Environment::local(local_runtime_paths)), } } @@ -102,11 +103,7 @@ impl EnvironmentManager { local_runtime_paths: ExecServerRuntimePaths, ) -> Self { let provider = DefaultEnvironmentProvider::new(exec_server_url); - match Self::from_provider_parts( - provider.environments(&local_runtime_paths), - provider.default_id(), - local_runtime_paths, - ) { + match Self::from_provider(&provider, local_runtime_paths).await { Ok(manager) => manager, Err(err) => panic!("default provider should create valid environments: {err}"), } @@ -122,14 +119,6 @@ impl EnvironmentManager { { let environments = provider.get_environments(&local_runtime_paths).await?; let default_environment_id = provider.default_environment_id(); - Self::from_provider_parts(environments, default_environment_id, local_runtime_paths) - } - - fn from_provider_parts( - environments: HashMap, - default_environment_id: Option, - local_runtime_paths: ExecServerRuntimePaths, - ) -> Result { for id in environments.keys() { if id.is_empty() { return Err(ExecServerError::Protocol( @@ -334,10 +323,33 @@ mod tests { use super::EnvironmentManager; use super::LOCAL_ENVIRONMENT_ID; use super::REMOTE_ENVIRONMENT_ID; + use crate::EnvironmentProvider; + use crate::ExecServerError; use crate::ExecServerRuntimePaths; use crate::ProcessId; + use async_trait::async_trait; use pretty_assertions::assert_eq; + #[derive(Clone)] + struct TestEnvironmentProvider { + environments: HashMap, + default_environment_id: Option, + } + + #[async_trait] + impl EnvironmentProvider for TestEnvironmentProvider { + async fn get_environments( + &self, + _local_runtime_paths: &ExecServerRuntimePaths, + ) -> Result, ExecServerError> { + Ok(self.environments.clone()) + } + + fn default_environment_id(&self) -> Option { + self.default_environment_id.clone() + } + } + fn test_runtime_paths() -> ExecServerRuntimePaths { ExecServerRuntimePaths::new( std::env::current_exe().expect("current exe"), @@ -428,17 +440,18 @@ mod tests { } #[tokio::test] - async fn environment_manager_builds_from_provider_parts() { - let manager = EnvironmentManager::from_provider_parts( - HashMap::from([( + async fn environment_manager_builds_from_provider() { + let provider = TestEnvironmentProvider { + environments: HashMap::from([( REMOTE_ENVIRONMENT_ID.to_string(), Environment::create_for_tests(Some("ws://127.0.0.1:8765".to_string())) .expect("remote environment"), )]), - Some(REMOTE_ENVIRONMENT_ID.to_string()), - test_runtime_paths(), - ) - .expect("environment manager"); + default_environment_id: Some(REMOTE_ENVIRONMENT_ID.to_string()), + }; + let manager = EnvironmentManager::from_provider(&provider, test_runtime_paths()) + .await + .expect("environment manager"); assert_eq!( manager.default_environment_id(), @@ -456,12 +469,13 @@ mod tests { #[tokio::test] async fn environment_manager_rejects_empty_environment_id() { - let err = EnvironmentManager::from_provider_parts( - HashMap::from([("".to_string(), Environment::default_for_tests())]), - None, - test_runtime_paths(), - ) - .expect_err("empty id should fail"); + let provider = TestEnvironmentProvider { + environments: HashMap::from([("".to_string(), Environment::default_for_tests())]), + default_environment_id: None, + }; + let err = EnvironmentManager::from_provider(&provider, test_runtime_paths()) + .await + .expect_err("empty id should fail"); assert_eq!( err.to_string(), @@ -471,8 +485,8 @@ mod tests { #[tokio::test] async fn environment_manager_uses_explicit_provider_default() { - let manager = EnvironmentManager::from_provider_parts( - HashMap::from([ + let provider = TestEnvironmentProvider { + environments: HashMap::from([ ( LOCAL_ENVIRONMENT_ID.to_string(), Environment::default_for_tests(), @@ -483,10 +497,11 @@ mod tests { .expect("remote environment"), ), ]), - Some("devbox".to_string()), - test_runtime_paths(), - ) - .expect("manager"); + default_environment_id: Some("devbox".to_string()), + }; + let manager = EnvironmentManager::from_provider(&provider, test_runtime_paths()) + .await + .expect("manager"); assert_eq!(manager.default_environment_id(), Some("devbox")); assert!(manager.default_environment().expect("default").is_remote()); @@ -494,15 +509,16 @@ mod tests { #[tokio::test] async fn environment_manager_disables_provider_default() { - let manager = EnvironmentManager::from_provider_parts( - HashMap::from([( + let provider = TestEnvironmentProvider { + environments: HashMap::from([( LOCAL_ENVIRONMENT_ID.to_string(), Environment::default_for_tests(), )]), - None, - test_runtime_paths(), - ) - .expect("manager"); + default_environment_id: None, + }; + let manager = EnvironmentManager::from_provider(&provider, test_runtime_paths()) + .await + .expect("manager"); assert_eq!(manager.default_environment_id(), None); assert!(manager.default_environment().is_none()); @@ -511,15 +527,16 @@ mod tests { #[tokio::test] async fn environment_manager_rejects_unknown_provider_default() { - let err = EnvironmentManager::from_provider_parts( - HashMap::from([( + let provider = TestEnvironmentProvider { + environments: HashMap::from([( LOCAL_ENVIRONMENT_ID.to_string(), Environment::default_for_tests(), )]), - Some("missing".to_string()), - test_runtime_paths(), - ) - .expect_err("unknown default should fail"); + default_environment_id: Some("missing".to_string()), + }; + let err = EnvironmentManager::from_provider(&provider, test_runtime_paths()) + .await + .expect_err("unknown default should fail"); assert_eq!( err.to_string(), From 16ae047c3276076c966f3cc9431af791b07e0830 Mon Sep 17 00:00:00 2001 From: starr-openai Date: Tue, 5 May 2026 13:58:49 -0700 Subject: [PATCH 34/42] Simplify environment provider defaults Co-authored-by: Codex --- codex-rs/app-server-client/src/lib.rs | 17 ++-- codex-rs/app-server/src/lib.rs | 15 ++-- codex-rs/core/src/connectors.rs | 2 +- codex-rs/core/src/environment_selection.rs | 3 +- codex-rs/core/src/prompt_debug.rs | 4 +- codex-rs/core/src/thread_manager_tests.rs | 11 +-- codex-rs/core/tests/common/test_codex.rs | 8 +- codex-rs/exec-server/src/environment.rs | 77 ++++++------------- .../exec-server/src/environment_provider.rs | 50 ++++++------ codex-rs/exec/src/lib.rs | 6 +- codex-rs/mcp-server/src/lib.rs | 15 ++-- codex-rs/thread-manager-sample/src/main.rs | 5 +- codex-rs/tui/src/lib.rs | 18 ++--- 13 files changed, 95 insertions(+), 136 deletions(-) diff --git a/codex-rs/app-server-client/src/lib.rs b/codex-rs/app-server-client/src/lib.rs index 0911cc448dd3..6e78e4e6c8fc 100644 --- a/codex-rs/app-server-client/src/lib.rs +++ b/codex-rs/app-server-client/src/lib.rs @@ -2046,17 +2046,14 @@ mod tests { #[tokio::test] async fn runtime_start_args_forward_environment_manager() { let config = Arc::new(build_test_config().await); - let environment_manager = Arc::new( - EnvironmentManager::create_for_tests( - Some("ws://127.0.0.1:8765".to_string()), - ExecServerRuntimePaths::new( - std::env::current_exe().expect("current exe"), - /*codex_linux_sandbox_exe*/ None, - ) - .expect("runtime paths"), + let environment_manager = Arc::new(EnvironmentManager::create_for_tests( + Some("ws://127.0.0.1:8765".to_string()), + ExecServerRuntimePaths::new( + std::env::current_exe().expect("current exe"), + /*codex_linux_sandbox_exe*/ None, ) - .await, - ); + .expect("runtime paths"), + )); let runtime_args = InProcessClientStartArgs { arg0_paths: Arg0DispatchPaths::default(), diff --git a/codex-rs/app-server/src/lib.rs b/codex-rs/app-server/src/lib.rs index 4013bbe76bc9..a1c23a6a1174 100644 --- a/codex-rs/app-server/src/lib.rs +++ b/codex-rs/app-server/src/lib.rs @@ -417,15 +417,12 @@ pub async fn run_main_with_transport_options( auth: AppServerWebsocketAuthSettings, runtime_options: AppServerRuntimeOptions, ) -> IoResult<()> { - let environment_manager = Arc::new( - EnvironmentManager::new(EnvironmentManagerArgs::new( - ExecServerRuntimePaths::from_optional_paths( - arg0_paths.codex_self_exe.clone(), - arg0_paths.codex_linux_sandbox_exe.clone(), - )?, - )) - .await, - ); + let environment_manager = Arc::new(EnvironmentManager::new(EnvironmentManagerArgs::new( + ExecServerRuntimePaths::from_optional_paths( + arg0_paths.codex_self_exe.clone(), + arg0_paths.codex_linux_sandbox_exe.clone(), + )?, + ))); let (transport_event_tx, mut transport_event_rx) = mpsc::channel::(CHANNEL_CAPACITY); let (outgoing_tx, mut outgoing_rx) = mpsc::channel::(CHANNEL_CAPACITY); diff --git a/codex-rs/core/src/connectors.rs b/codex-rs/core/src/connectors.rs index 9f0381f53a92..a2fae1a8edb7 100644 --- a/codex-rs/core/src/connectors.rs +++ b/codex-rs/core/src/connectors.rs @@ -201,7 +201,7 @@ pub async fn list_accessible_connectors_from_mcp_tools_with_options_and_status( config.codex_linux_sandbox_exe.clone(), )?; let environment_manager = - EnvironmentManager::new(EnvironmentManagerArgs::new(local_runtime_paths)).await; + EnvironmentManager::new(EnvironmentManagerArgs::new(local_runtime_paths)); list_accessible_connectors_from_mcp_tools_with_environment_manager( config, force_refetch, diff --git a/codex-rs/core/src/environment_selection.rs b/codex-rs/core/src/environment_selection.rs index e9d617cbf432..714592009ac9 100644 --- a/codex-rs/core/src/environment_selection.rs +++ b/codex-rs/core/src/environment_selection.rs @@ -108,8 +108,7 @@ mod tests { let manager = EnvironmentManager::create_for_tests( Some("ws://127.0.0.1:8765".to_string()), test_runtime_paths(), - ) - .await; + ); assert_eq!( default_thread_environment_selections(&manager, &cwd), diff --git a/codex-rs/core/src/prompt_debug.rs b/codex-rs/core/src/prompt_debug.rs index a40fd2794244..2925d0dd2397 100644 --- a/codex-rs/core/src/prompt_debug.rs +++ b/codex-rs/core/src/prompt_debug.rs @@ -46,7 +46,9 @@ pub async fn build_prompt_input( &config, Arc::clone(&auth_manager), SessionSource::Exec, - Arc::new(EnvironmentManager::new(EnvironmentManagerArgs::new(local_runtime_paths)).await), + Arc::new(EnvironmentManager::new(EnvironmentManagerArgs::new( + local_runtime_paths, + ))), /*analytics_events_client*/ None, state_db, thread_store, diff --git a/codex-rs/core/src/thread_manager_tests.rs b/codex-rs/core/src/thread_manager_tests.rs index 643309aac166..71c61cf15d2f 100644 --- a/codex-rs/core/src/thread_manager_tests.rs +++ b/codex-rs/core/src/thread_manager_tests.rs @@ -313,13 +313,10 @@ async fn start_thread_accepts_explicit_environment_when_default_environment_is_d /*codex_linux_sandbox_exe*/ None, ) .expect("runtime paths"); - let environment_manager = Arc::new( - codex_exec_server::EnvironmentManager::create_for_tests( - Some("none".to_string()), - runtime_paths, - ) - .await, - ); + let environment_manager = Arc::new(codex_exec_server::EnvironmentManager::create_for_tests( + Some("none".to_string()), + runtime_paths, + )); let manager = ThreadManager::with_models_provider_and_home_for_tests( CodexAuth::from_api_key("dummy"), config.model_provider.clone(), diff --git a/codex-rs/core/tests/common/test_codex.rs b/codex-rs/core/tests/common/test_codex.rs index 13dd026654c0..56faad366723 100644 --- a/codex-rs/core/tests/common/test_codex.rs +++ b/codex-rs/core/tests/common/test_codex.rs @@ -390,13 +390,11 @@ impl TestCodexBuilder { std::env::current_exe()?, /*codex_linux_sandbox_exe*/ None, )?; - let environment_manager = Arc::new( - codex_exec_server::EnvironmentManager::create_for_tests( + let environment_manager = + Arc::new(codex_exec_server::EnvironmentManager::create_for_tests( exec_server_url, local_runtime_paths, - ) - .await, - ); + )); let file_system = test_env.environment().get_filesystem(); let mut workspace_setups = vec![]; swap(&mut self.workspace_setups, &mut workspace_setups); diff --git a/codex-rs/exec-server/src/environment.rs b/codex-rs/exec-server/src/environment.rs index 3c70374bade3..7dc4ae2ae8e3 100644 --- a/codex-rs/exec-server/src/environment.rs +++ b/codex-rs/exec-server/src/environment.rs @@ -81,43 +81,43 @@ impl EnvironmentManager { } /// Builds a test-only manager from a raw exec-server URL value. - pub async fn create_for_tests( + pub fn create_for_tests( exec_server_url: Option, local_runtime_paths: ExecServerRuntimePaths, ) -> Self { - Self::from_default_provider_url(exec_server_url, local_runtime_paths).await + Self::from_default_provider_url(exec_server_url, local_runtime_paths) } /// Builds a manager from `CODEX_EXEC_SERVER_URL` and local runtime paths /// used when creating local filesystem helpers. - pub async fn new(args: EnvironmentManagerArgs) -> Self { + pub fn new(args: EnvironmentManagerArgs) -> Self { let EnvironmentManagerArgs { local_runtime_paths, } = args; let exec_server_url = std::env::var(CODEX_EXEC_SERVER_URL_ENV_VAR).ok(); - Self::from_default_provider_url(exec_server_url, local_runtime_paths).await + Self::from_default_provider_url(exec_server_url, local_runtime_paths) } - async fn from_default_provider_url( + fn from_default_provider_url( exec_server_url: Option, local_runtime_paths: ExecServerRuntimePaths, ) -> Self { let provider = DefaultEnvironmentProvider::new(exec_server_url); - match Self::from_provider(&provider, local_runtime_paths).await { + match Self::from_provider(&provider, local_runtime_paths) { Ok(manager) => manager, Err(err) => panic!("default provider should create valid environments: {err}"), } } /// Builds a manager from a provider-supplied environment list and default. - pub async fn from_provider

( + pub fn from_provider

( provider: &P, local_runtime_paths: ExecServerRuntimePaths, ) -> Result where P: EnvironmentProvider + ?Sized, { - let environments = provider.get_environments(&local_runtime_paths).await?; + let environments = provider.get_environments(&local_runtime_paths)?; let default_environment_id = provider.default_environment_id(); for id in environments.keys() { if id.is_empty() { @@ -177,7 +177,6 @@ impl EnvironmentManager { #[derive(Clone)] pub struct Environment { exec_server_url: Option, - remote_transport: Option, exec_backend: Arc, filesystem: Arc, http_client: Arc, @@ -189,7 +188,6 @@ impl Environment { pub fn default_for_tests() -> Self { Self { exec_server_url: None, - remote_transport: None, exec_backend: Arc::new(LocalProcess::default()), filesystem: Arc::new(LocalFileSystem::unsandboxed()), http_client: Arc::new(ReqwestHttpClient), @@ -245,7 +243,6 @@ impl Environment { pub(crate) fn local(local_runtime_paths: ExecServerRuntimePaths) -> Self { Self { exec_server_url: None, - remote_transport: None, exec_backend: Arc::new(LocalProcess::default()), filesystem: Arc::new(LocalFileSystem::with_runtime_paths( local_runtime_paths.clone(), @@ -259,28 +256,15 @@ impl Environment { exec_server_url: String, local_runtime_paths: Option, ) -> Self { - Self::remote_with_transport( - ExecServerTransportParams::WebSocketUrl(exec_server_url), - local_runtime_paths, - ) - } - - fn remote_with_transport( - transport_params: ExecServerTransportParams, - local_runtime_paths: Option, - ) -> Self { - let exec_server_url = match &transport_params { - ExecServerTransportParams::WebSocketUrl(url) => Some(url.clone()), - ExecServerTransportParams::StdioCommand(_) => None, - }; - let client = LazyRemoteExecServerClient::new(transport_params.clone()); + let client = LazyRemoteExecServerClient::new(ExecServerTransportParams::WebSocketUrl( + exec_server_url.clone(), + )); let exec_backend: Arc = Arc::new(RemoteProcess::new(client.clone())); let filesystem: Arc = Arc::new(RemoteFileSystem::new(client.clone())); Self { - exec_server_url, - remote_transport: Some(transport_params), + exec_server_url: Some(exec_server_url), exec_backend, filesystem, http_client: Arc::new(client), @@ -289,7 +273,7 @@ impl Environment { } pub fn is_remote(&self) -> bool { - self.remote_transport.is_some() + self.exec_server_url.is_some() } /// Returns the remote exec-server URL when this environment is remote. @@ -327,18 +311,15 @@ mod tests { use crate::ExecServerError; use crate::ExecServerRuntimePaths; use crate::ProcessId; - use async_trait::async_trait; use pretty_assertions::assert_eq; - #[derive(Clone)] struct TestEnvironmentProvider { environments: HashMap, default_environment_id: Option, } - #[async_trait] impl EnvironmentProvider for TestEnvironmentProvider { - async fn get_environments( + fn get_environments( &self, _local_runtime_paths: &ExecServerRuntimePaths, ) -> Result, ExecServerError> { @@ -370,7 +351,7 @@ mod tests { #[tokio::test] async fn environment_manager_normalizes_empty_url() { let manager = - EnvironmentManager::create_for_tests(Some(String::new()), test_runtime_paths()).await; + EnvironmentManager::create_for_tests(Some(String::new()), test_runtime_paths()); let environment = manager.default_environment().expect("default environment"); assert_eq!(manager.default_environment_id(), Some(LOCAL_ENVIRONMENT_ID)); @@ -400,8 +381,7 @@ mod tests { let manager = EnvironmentManager::create_for_tests( Some("ws://127.0.0.1:8765".to_string()), test_runtime_paths(), - ) - .await; + ); let environment = manager.default_environment().expect("default environment"); assert_eq!( @@ -450,7 +430,6 @@ mod tests { default_environment_id: Some(REMOTE_ENVIRONMENT_ID.to_string()), }; let manager = EnvironmentManager::from_provider(&provider, test_runtime_paths()) - .await .expect("environment manager"); assert_eq!( @@ -474,7 +453,6 @@ mod tests { default_environment_id: None, }; let err = EnvironmentManager::from_provider(&provider, test_runtime_paths()) - .await .expect_err("empty id should fail"); assert_eq!( @@ -499,9 +477,8 @@ mod tests { ]), default_environment_id: Some("devbox".to_string()), }; - let manager = EnvironmentManager::from_provider(&provider, test_runtime_paths()) - .await - .expect("manager"); + let manager = + EnvironmentManager::from_provider(&provider, test_runtime_paths()).expect("manager"); assert_eq!(manager.default_environment_id(), Some("devbox")); assert!(manager.default_environment().expect("default").is_remote()); @@ -516,9 +493,8 @@ mod tests { )]), default_environment_id: None, }; - let manager = EnvironmentManager::from_provider(&provider, test_runtime_paths()) - .await - .expect("manager"); + let manager = + EnvironmentManager::from_provider(&provider, test_runtime_paths()).expect("manager"); assert_eq!(manager.default_environment_id(), None); assert!(manager.default_environment().is_none()); @@ -535,7 +511,6 @@ mod tests { default_environment_id: Some("missing".to_string()), }; let err = EnvironmentManager::from_provider(&provider, test_runtime_paths()) - .await .expect_err("unknown default should fail"); assert_eq!( @@ -549,8 +524,7 @@ mod tests { let manager = EnvironmentManager::create_for_tests( /*exec_server_url*/ None, test_runtime_paths(), - ) - .await; + ); assert_eq!(manager.default_environment_id(), Some(LOCAL_ENVIRONMENT_ID)); let provider_local = manager @@ -567,8 +541,7 @@ mod tests { let manager = EnvironmentManager::create_for_tests( /*exec_server_url*/ None, runtime_paths.clone(), - ) - .await; + ); let environment = manager.default_environment().expect("default environment"); @@ -579,8 +552,7 @@ mod tests { .local_runtime_paths() .expect("local runtime paths") .clone(), - ) - .await; + ); let environment = manager.default_environment().expect("default environment"); assert_eq!(environment.local_runtime_paths(), Some(&runtime_paths)); } @@ -596,8 +568,7 @@ mod tests { #[tokio::test] async fn environment_manager_keeps_default_provider_local_lookup_when_default_disabled() { let manager = - EnvironmentManager::create_for_tests(Some("none".to_string()), test_runtime_paths()) - .await; + EnvironmentManager::create_for_tests(Some("none".to_string()), test_runtime_paths()); assert!(manager.default_environment().is_none()); assert_eq!(manager.default_environment_id(), None); diff --git a/codex-rs/exec-server/src/environment_provider.rs b/codex-rs/exec-server/src/environment_provider.rs index 57f4340498d8..677a5544804d 100644 --- a/codex-rs/exec-server/src/environment_provider.rs +++ b/codex-rs/exec-server/src/environment_provider.rs @@ -1,7 +1,5 @@ use std::collections::HashMap; -use async_trait::async_trait; - use crate::Environment; use crate::ExecServerError; use crate::ExecServerRuntimePaths; @@ -14,10 +12,9 @@ use crate::environment::REMOTE_ENVIRONMENT_ID; /// Implementations own both the available environment list and the default /// environment id. Providers that want the local environment to be addressable /// by id should include it explicitly in the returned map. -#[async_trait] pub trait EnvironmentProvider: Send + Sync { /// Returns the environments available for a new manager. - async fn get_environments( + fn get_environments( &self, local_runtime_paths: &ExecServerRuntimePaths, ) -> Result, ExecServerError>; @@ -63,22 +60,10 @@ impl DefaultEnvironmentProvider { environments } - - pub(crate) fn default_id(&self) -> Option { - let (exec_server_url, disabled) = normalize_exec_server_url(self.exec_server_url.clone()); - if disabled { - None - } else if exec_server_url.is_some() { - Some(REMOTE_ENVIRONMENT_ID.to_string()) - } else { - Some(LOCAL_ENVIRONMENT_ID.to_string()) - } - } } -#[async_trait] impl EnvironmentProvider for DefaultEnvironmentProvider { - async fn get_environments( + fn get_environments( &self, local_runtime_paths: &ExecServerRuntimePaths, ) -> Result, ExecServerError> { @@ -86,7 +71,19 @@ impl EnvironmentProvider for DefaultEnvironmentProvider { } fn default_environment_id(&self) -> Option { - self.default_id() + let (exec_server_url, disabled) = normalize_exec_server_url(self.exec_server_url.clone()); + if disabled { + return None; + } + + Some( + if exec_server_url.is_some() { + REMOTE_ENVIRONMENT_ID + } else { + LOCAL_ENVIRONMENT_ID + } + .to_string(), + ) } } @@ -119,7 +116,6 @@ mod tests { let runtime_paths = test_runtime_paths(); let environments = provider .get_environments(&runtime_paths) - .await .expect("environments"); assert!(!environments[LOCAL_ENVIRONMENT_ID].is_remote()); @@ -128,6 +124,10 @@ mod tests { Some(&runtime_paths) ); assert!(!environments.contains_key(REMOTE_ENVIRONMENT_ID)); + assert_eq!( + provider.default_environment_id(), + Some(LOCAL_ENVIRONMENT_ID.to_string()) + ); } #[tokio::test] @@ -136,11 +136,14 @@ mod tests { let runtime_paths = test_runtime_paths(); let environments = provider .get_environments(&runtime_paths) - .await .expect("environments"); assert!(!environments[LOCAL_ENVIRONMENT_ID].is_remote()); assert!(!environments.contains_key(REMOTE_ENVIRONMENT_ID)); + assert_eq!( + provider.default_environment_id(), + Some(LOCAL_ENVIRONMENT_ID.to_string()) + ); } #[tokio::test] @@ -149,7 +152,6 @@ mod tests { let runtime_paths = test_runtime_paths(); let environments = provider .get_environments(&runtime_paths) - .await .expect("environments"); assert!(!environments[LOCAL_ENVIRONMENT_ID].is_remote()); @@ -163,7 +165,6 @@ mod tests { let runtime_paths = test_runtime_paths(); let environments = provider .get_environments(&runtime_paths) - .await .expect("environments"); assert!(!environments[LOCAL_ENVIRONMENT_ID].is_remote()); @@ -173,6 +174,10 @@ mod tests { remote_environment.exec_server_url(), Some("ws://127.0.0.1:8765") ); + assert_eq!( + provider.default_environment_id(), + Some(REMOTE_ENVIRONMENT_ID.to_string()) + ); } #[tokio::test] @@ -181,7 +186,6 @@ mod tests { let runtime_paths = test_runtime_paths(); let environments = provider .get_environments(&runtime_paths) - .await .expect("environments"); assert_eq!( diff --git a/codex-rs/exec/src/lib.rs b/codex-rs/exec/src/lib.rs index d61346f1d09e..dd866a506de0 100644 --- a/codex-rs/exec/src/lib.rs +++ b/codex-rs/exec/src/lib.rs @@ -515,9 +515,9 @@ pub async fn run_main(cli: Cli, arg0_paths: Arg0DispatchPaths) -> anyhow::Result feedback: CodexFeedback::new(), log_db: None, state_db: state_db.clone(), - environment_manager: std::sync::Arc::new( - EnvironmentManager::new(EnvironmentManagerArgs::new(local_runtime_paths)).await, - ), + environment_manager: std::sync::Arc::new(EnvironmentManager::new( + EnvironmentManagerArgs::new(local_runtime_paths), + )), config_warnings, session_source: SessionSource::Exec, enable_codex_api_key_env: true, diff --git a/codex-rs/mcp-server/src/lib.rs b/codex-rs/mcp-server/src/lib.rs index ac764456f5e5..8fac7e1cdfe6 100644 --- a/codex-rs/mcp-server/src/lib.rs +++ b/codex-rs/mcp-server/src/lib.rs @@ -60,15 +60,12 @@ pub async fn run_main( arg0_paths: Arg0DispatchPaths, cli_config_overrides: CliConfigOverrides, ) -> IoResult<()> { - let environment_manager = Arc::new( - EnvironmentManager::new(EnvironmentManagerArgs::new( - ExecServerRuntimePaths::from_optional_paths( - arg0_paths.codex_self_exe.clone(), - arg0_paths.codex_linux_sandbox_exe.clone(), - )?, - )) - .await, - ); + let environment_manager = Arc::new(EnvironmentManager::new(EnvironmentManagerArgs::new( + ExecServerRuntimePaths::from_optional_paths( + arg0_paths.codex_self_exe.clone(), + arg0_paths.codex_linux_sandbox_exe.clone(), + )?, + ))); // Parse CLI overrides once and derive the base Config eagerly so later // components do not need to work with raw TOML values. let cli_kv_overrides = cli_config_overrides.parse_overrides().map_err(|e| { diff --git a/codex-rs/thread-manager-sample/src/main.rs b/codex-rs/thread-manager-sample/src/main.rs index a27c78d54358..2d3e365fbae0 100644 --- a/codex-rs/thread-manager-sample/src/main.rs +++ b/codex-rs/thread-manager-sample/src/main.rs @@ -117,8 +117,9 @@ async fn run_main(arg0_paths: Arg0DispatchPaths) -> anyhow::Result<()> { }; let thread_store = thread_store_from_config(&config, state_db.clone()); let agent_graph_store = agent_graph_store_from_state_db(state_db.clone()); - let environment_manager = - Arc::new(EnvironmentManager::new(EnvironmentManagerArgs::new(local_runtime_paths)).await); + let environment_manager = Arc::new(EnvironmentManager::new(EnvironmentManagerArgs::new( + local_runtime_paths, + ))); let thread_manager = ThreadManager::new( &config, auth_manager, diff --git a/codex-rs/tui/src/lib.rs b/codex-rs/tui/src/lib.rs index b2e92a19b485..5bcc6ffa26af 100644 --- a/codex-rs/tui/src/lib.rs +++ b/codex-rs/tui/src/lib.rs @@ -761,15 +761,12 @@ pub async fn run_main( } }; - let environment_manager = Arc::new( - EnvironmentManager::new(EnvironmentManagerArgs::new( - ExecServerRuntimePaths::from_optional_paths( - arg0_paths.codex_self_exe.clone(), - arg0_paths.codex_linux_sandbox_exe.clone(), - )?, - )) - .await, - ); + let environment_manager = Arc::new(EnvironmentManager::new(EnvironmentManagerArgs::new( + ExecServerRuntimePaths::from_optional_paths( + arg0_paths.codex_self_exe.clone(), + arg0_paths.codex_linux_sandbox_exe.clone(), + )?, + ))); let cwd = cli.cwd.clone(); let config_cwd = config_cwd_for_app_server_target(cwd.as_deref(), &app_server_target, &environment_manager)?; @@ -2141,8 +2138,7 @@ mod tests { std::env::current_exe().expect("current exe"), /*codex_linux_sandbox_exe*/ None, )?, - ) - .await; + ); let config_cwd = config_cwd_for_app_server_target(Some(remote_only_cwd), &target, &environment_manager)?; From bc0677df9848080a454188f824dd816689e23333 Mon Sep 17 00:00:00 2001 From: starr-openai Date: Wed, 6 May 2026 12:04:08 -0700 Subject: [PATCH 35/42] Represent provider defaults with snapshots Keep EnvironmentManager construction async to preserve caller behavior while moving provider-owned default selection into a single snapshot object. Co-authored-by: Codex --- codex-rs/app-server-client/src/lib.rs | 17 +- codex-rs/app-server/src/lib.rs | 15 +- codex-rs/core/src/connectors.rs | 2 +- codex-rs/core/src/environment_selection.rs | 3 +- codex-rs/core/src/prompt_debug.rs | 4 +- codex-rs/core/src/thread_manager_tests.rs | 11 +- codex-rs/core/tests/common/test_codex.rs | 8 +- codex-rs/exec-server/src/environment.rs | 173 +++++++++++------- .../exec-server/src/environment_provider.rs | 115 +++++++----- codex-rs/exec/src/lib.rs | 6 +- codex-rs/mcp-server/src/lib.rs | 15 +- codex-rs/thread-manager-sample/src/main.rs | 5 +- codex-rs/tui/src/lib.rs | 18 +- 13 files changed, 230 insertions(+), 162 deletions(-) diff --git a/codex-rs/app-server-client/src/lib.rs b/codex-rs/app-server-client/src/lib.rs index 6e78e4e6c8fc..0911cc448dd3 100644 --- a/codex-rs/app-server-client/src/lib.rs +++ b/codex-rs/app-server-client/src/lib.rs @@ -2046,14 +2046,17 @@ mod tests { #[tokio::test] async fn runtime_start_args_forward_environment_manager() { let config = Arc::new(build_test_config().await); - let environment_manager = Arc::new(EnvironmentManager::create_for_tests( - Some("ws://127.0.0.1:8765".to_string()), - ExecServerRuntimePaths::new( - std::env::current_exe().expect("current exe"), - /*codex_linux_sandbox_exe*/ None, + let environment_manager = Arc::new( + EnvironmentManager::create_for_tests( + Some("ws://127.0.0.1:8765".to_string()), + ExecServerRuntimePaths::new( + std::env::current_exe().expect("current exe"), + /*codex_linux_sandbox_exe*/ None, + ) + .expect("runtime paths"), ) - .expect("runtime paths"), - )); + .await, + ); let runtime_args = InProcessClientStartArgs { arg0_paths: Arg0DispatchPaths::default(), diff --git a/codex-rs/app-server/src/lib.rs b/codex-rs/app-server/src/lib.rs index a1c23a6a1174..4013bbe76bc9 100644 --- a/codex-rs/app-server/src/lib.rs +++ b/codex-rs/app-server/src/lib.rs @@ -417,12 +417,15 @@ pub async fn run_main_with_transport_options( auth: AppServerWebsocketAuthSettings, runtime_options: AppServerRuntimeOptions, ) -> IoResult<()> { - let environment_manager = Arc::new(EnvironmentManager::new(EnvironmentManagerArgs::new( - ExecServerRuntimePaths::from_optional_paths( - arg0_paths.codex_self_exe.clone(), - arg0_paths.codex_linux_sandbox_exe.clone(), - )?, - ))); + let environment_manager = Arc::new( + EnvironmentManager::new(EnvironmentManagerArgs::new( + ExecServerRuntimePaths::from_optional_paths( + arg0_paths.codex_self_exe.clone(), + arg0_paths.codex_linux_sandbox_exe.clone(), + )?, + )) + .await, + ); let (transport_event_tx, mut transport_event_rx) = mpsc::channel::(CHANNEL_CAPACITY); let (outgoing_tx, mut outgoing_rx) = mpsc::channel::(CHANNEL_CAPACITY); diff --git a/codex-rs/core/src/connectors.rs b/codex-rs/core/src/connectors.rs index a2fae1a8edb7..9f0381f53a92 100644 --- a/codex-rs/core/src/connectors.rs +++ b/codex-rs/core/src/connectors.rs @@ -201,7 +201,7 @@ pub async fn list_accessible_connectors_from_mcp_tools_with_options_and_status( config.codex_linux_sandbox_exe.clone(), )?; let environment_manager = - EnvironmentManager::new(EnvironmentManagerArgs::new(local_runtime_paths)); + EnvironmentManager::new(EnvironmentManagerArgs::new(local_runtime_paths)).await; list_accessible_connectors_from_mcp_tools_with_environment_manager( config, force_refetch, diff --git a/codex-rs/core/src/environment_selection.rs b/codex-rs/core/src/environment_selection.rs index 714592009ac9..e9d617cbf432 100644 --- a/codex-rs/core/src/environment_selection.rs +++ b/codex-rs/core/src/environment_selection.rs @@ -108,7 +108,8 @@ mod tests { let manager = EnvironmentManager::create_for_tests( Some("ws://127.0.0.1:8765".to_string()), test_runtime_paths(), - ); + ) + .await; assert_eq!( default_thread_environment_selections(&manager, &cwd), diff --git a/codex-rs/core/src/prompt_debug.rs b/codex-rs/core/src/prompt_debug.rs index 2925d0dd2397..a40fd2794244 100644 --- a/codex-rs/core/src/prompt_debug.rs +++ b/codex-rs/core/src/prompt_debug.rs @@ -46,9 +46,7 @@ pub async fn build_prompt_input( &config, Arc::clone(&auth_manager), SessionSource::Exec, - Arc::new(EnvironmentManager::new(EnvironmentManagerArgs::new( - local_runtime_paths, - ))), + Arc::new(EnvironmentManager::new(EnvironmentManagerArgs::new(local_runtime_paths)).await), /*analytics_events_client*/ None, state_db, thread_store, diff --git a/codex-rs/core/src/thread_manager_tests.rs b/codex-rs/core/src/thread_manager_tests.rs index 71c61cf15d2f..643309aac166 100644 --- a/codex-rs/core/src/thread_manager_tests.rs +++ b/codex-rs/core/src/thread_manager_tests.rs @@ -313,10 +313,13 @@ async fn start_thread_accepts_explicit_environment_when_default_environment_is_d /*codex_linux_sandbox_exe*/ None, ) .expect("runtime paths"); - let environment_manager = Arc::new(codex_exec_server::EnvironmentManager::create_for_tests( - Some("none".to_string()), - runtime_paths, - )); + let environment_manager = Arc::new( + codex_exec_server::EnvironmentManager::create_for_tests( + Some("none".to_string()), + runtime_paths, + ) + .await, + ); let manager = ThreadManager::with_models_provider_and_home_for_tests( CodexAuth::from_api_key("dummy"), config.model_provider.clone(), diff --git a/codex-rs/core/tests/common/test_codex.rs b/codex-rs/core/tests/common/test_codex.rs index 56faad366723..13dd026654c0 100644 --- a/codex-rs/core/tests/common/test_codex.rs +++ b/codex-rs/core/tests/common/test_codex.rs @@ -390,11 +390,13 @@ impl TestCodexBuilder { std::env::current_exe()?, /*codex_linux_sandbox_exe*/ None, )?; - let environment_manager = - Arc::new(codex_exec_server::EnvironmentManager::create_for_tests( + let environment_manager = Arc::new( + codex_exec_server::EnvironmentManager::create_for_tests( exec_server_url, local_runtime_paths, - )); + ) + .await, + ); let file_system = test_env.environment().get_filesystem(); let mut workspace_setups = vec![]; swap(&mut self.workspace_setups, &mut workspace_setups); diff --git a/codex-rs/exec-server/src/environment.rs b/codex-rs/exec-server/src/environment.rs index 7dc4ae2ae8e3..d7867c973289 100644 --- a/codex-rs/exec-server/src/environment.rs +++ b/codex-rs/exec-server/src/environment.rs @@ -9,7 +9,9 @@ use crate::client::LazyRemoteExecServerClient; use crate::client::http_client::ReqwestHttpClient; use crate::client_api::ExecServerTransportParams; use crate::environment_provider::DefaultEnvironmentProvider; +use crate::environment_provider::EnvironmentDefault; use crate::environment_provider::EnvironmentProvider; +use crate::environment_provider::EnvironmentProviderSnapshot; use crate::environment_provider::normalize_exec_server_url; use crate::local_file_system::LocalFileSystem; use crate::local_process::LocalProcess; @@ -23,8 +25,7 @@ pub const CODEX_EXEC_SERVER_URL_ENV_VAR: &str = "CODEX_EXEC_SERVER_URL"; /// /// `EnvironmentManager` is a shared registry for concrete environments. Its /// default constructor preserves the legacy `CODEX_EXEC_SERVER_URL` behavior -/// while provider-based construction accepts a provider-supplied environment -/// list and default id. +/// while provider-based construction accepts a provider-supplied snapshot. /// /// Setting `CODEX_EXEC_SERVER_URL=none` disables environment access by leaving /// the default environment unset while still keeping an explicit local @@ -81,44 +82,56 @@ impl EnvironmentManager { } /// Builds a test-only manager from a raw exec-server URL value. - pub fn create_for_tests( + pub async fn create_for_tests( exec_server_url: Option, local_runtime_paths: ExecServerRuntimePaths, ) -> Self { - Self::from_default_provider_url(exec_server_url, local_runtime_paths) + Self::from_default_provider_url(exec_server_url, local_runtime_paths).await } /// Builds a manager from `CODEX_EXEC_SERVER_URL` and local runtime paths /// used when creating local filesystem helpers. - pub fn new(args: EnvironmentManagerArgs) -> Self { + pub async fn new(args: EnvironmentManagerArgs) -> Self { let EnvironmentManagerArgs { local_runtime_paths, } = args; let exec_server_url = std::env::var(CODEX_EXEC_SERVER_URL_ENV_VAR).ok(); - Self::from_default_provider_url(exec_server_url, local_runtime_paths) + Self::from_default_provider_url(exec_server_url, local_runtime_paths).await } - fn from_default_provider_url( + async fn from_default_provider_url( exec_server_url: Option, local_runtime_paths: ExecServerRuntimePaths, ) -> Self { let provider = DefaultEnvironmentProvider::new(exec_server_url); - match Self::from_provider(&provider, local_runtime_paths) { + match Self::from_provider(&provider, local_runtime_paths).await { Ok(manager) => manager, Err(err) => panic!("default provider should create valid environments: {err}"), } } - /// Builds a manager from a provider-supplied environment list and default. - pub fn from_provider

( + /// Builds a manager from a provider-supplied startup snapshot. + pub async fn from_provider

( provider: &P, local_runtime_paths: ExecServerRuntimePaths, ) -> Result where P: EnvironmentProvider + ?Sized, { - let environments = provider.get_environments(&local_runtime_paths)?; - let default_environment_id = provider.default_environment_id(); + Self::from_provider_snapshot( + provider.snapshot(&local_runtime_paths).await?, + local_runtime_paths, + ) + } + + fn from_provider_snapshot( + snapshot: EnvironmentProviderSnapshot, + local_runtime_paths: ExecServerRuntimePaths, + ) -> Result { + let EnvironmentProviderSnapshot { + environments, + default, + } = snapshot; for id in environments.keys() { if id.is_empty() { return Err(ExecServerError::Protocol( @@ -127,13 +140,17 @@ impl EnvironmentManager { } } - if let Some(environment_id) = default_environment_id.as_ref() - && !environments.contains_key(environment_id) - { - return Err(ExecServerError::Protocol(format!( - "default environment `{environment_id}` is not configured" - ))); - } + let default_environment = match default { + EnvironmentDefault::Disabled => None, + EnvironmentDefault::EnvironmentId(environment_id) => { + if !environments.contains_key(&environment_id) { + return Err(ExecServerError::Protocol(format!( + "default environment `{environment_id}` is not configured" + ))); + } + Some(environment_id) + } + }; let local_environment = Arc::new(Environment::local(local_runtime_paths)); let environments = environments .into_iter() @@ -141,7 +158,7 @@ impl EnvironmentManager { .collect(); Ok(Self { - default_environment: default_environment_id, + default_environment, environments, local_environment, }) @@ -311,23 +328,21 @@ mod tests { use crate::ExecServerError; use crate::ExecServerRuntimePaths; use crate::ProcessId; + use crate::environment_provider::EnvironmentDefault; + use crate::environment_provider::EnvironmentProviderSnapshot; use pretty_assertions::assert_eq; struct TestEnvironmentProvider { - environments: HashMap, - default_environment_id: Option, + snapshot: EnvironmentProviderSnapshot, } + #[async_trait::async_trait] impl EnvironmentProvider for TestEnvironmentProvider { - fn get_environments( + async fn snapshot( &self, _local_runtime_paths: &ExecServerRuntimePaths, - ) -> Result, ExecServerError> { - Ok(self.environments.clone()) - } - - fn default_environment_id(&self) -> Option { - self.default_environment_id.clone() + ) -> Result { + Ok(self.snapshot.clone()) } } @@ -351,7 +366,7 @@ mod tests { #[tokio::test] async fn environment_manager_normalizes_empty_url() { let manager = - EnvironmentManager::create_for_tests(Some(String::new()), test_runtime_paths()); + EnvironmentManager::create_for_tests(Some(String::new()), test_runtime_paths()).await; let environment = manager.default_environment().expect("default environment"); assert_eq!(manager.default_environment_id(), Some(LOCAL_ENVIRONMENT_ID)); @@ -381,7 +396,8 @@ mod tests { let manager = EnvironmentManager::create_for_tests( Some("ws://127.0.0.1:8765".to_string()), test_runtime_paths(), - ); + ) + .await; let environment = manager.default_environment().expect("default environment"); assert_eq!( @@ -422,14 +438,17 @@ mod tests { #[tokio::test] async fn environment_manager_builds_from_provider() { let provider = TestEnvironmentProvider { - environments: HashMap::from([( - REMOTE_ENVIRONMENT_ID.to_string(), - Environment::create_for_tests(Some("ws://127.0.0.1:8765".to_string())) - .expect("remote environment"), - )]), - default_environment_id: Some(REMOTE_ENVIRONMENT_ID.to_string()), + snapshot: EnvironmentProviderSnapshot { + environments: HashMap::from([( + REMOTE_ENVIRONMENT_ID.to_string(), + Environment::create_for_tests(Some("ws://127.0.0.1:8765".to_string())) + .expect("remote environment"), + )]), + default: EnvironmentDefault::EnvironmentId(REMOTE_ENVIRONMENT_ID.to_string()), + }, }; let manager = EnvironmentManager::from_provider(&provider, test_runtime_paths()) + .await .expect("environment manager"); assert_eq!( @@ -449,10 +468,13 @@ mod tests { #[tokio::test] async fn environment_manager_rejects_empty_environment_id() { let provider = TestEnvironmentProvider { - environments: HashMap::from([("".to_string(), Environment::default_for_tests())]), - default_environment_id: None, + snapshot: EnvironmentProviderSnapshot { + environments: HashMap::from([("".to_string(), Environment::default_for_tests())]), + default: EnvironmentDefault::Disabled, + }, }; let err = EnvironmentManager::from_provider(&provider, test_runtime_paths()) + .await .expect_err("empty id should fail"); assert_eq!( @@ -464,21 +486,24 @@ mod tests { #[tokio::test] async fn environment_manager_uses_explicit_provider_default() { let provider = TestEnvironmentProvider { - environments: HashMap::from([ - ( - LOCAL_ENVIRONMENT_ID.to_string(), - Environment::default_for_tests(), - ), - ( - "devbox".to_string(), - Environment::create_for_tests(Some("ws://127.0.0.1:8765".to_string())) - .expect("remote environment"), - ), - ]), - default_environment_id: Some("devbox".to_string()), + snapshot: EnvironmentProviderSnapshot { + environments: HashMap::from([ + ( + LOCAL_ENVIRONMENT_ID.to_string(), + Environment::default_for_tests(), + ), + ( + "devbox".to_string(), + Environment::create_for_tests(Some("ws://127.0.0.1:8765".to_string())) + .expect("remote environment"), + ), + ]), + default: EnvironmentDefault::EnvironmentId("devbox".to_string()), + }, }; - let manager = - EnvironmentManager::from_provider(&provider, test_runtime_paths()).expect("manager"); + let manager = EnvironmentManager::from_provider(&provider, test_runtime_paths()) + .await + .expect("manager"); assert_eq!(manager.default_environment_id(), Some("devbox")); assert!(manager.default_environment().expect("default").is_remote()); @@ -487,14 +512,17 @@ mod tests { #[tokio::test] async fn environment_manager_disables_provider_default() { let provider = TestEnvironmentProvider { - environments: HashMap::from([( - LOCAL_ENVIRONMENT_ID.to_string(), - Environment::default_for_tests(), - )]), - default_environment_id: None, + snapshot: EnvironmentProviderSnapshot { + environments: HashMap::from([( + LOCAL_ENVIRONMENT_ID.to_string(), + Environment::default_for_tests(), + )]), + default: EnvironmentDefault::Disabled, + }, }; - let manager = - EnvironmentManager::from_provider(&provider, test_runtime_paths()).expect("manager"); + let manager = EnvironmentManager::from_provider(&provider, test_runtime_paths()) + .await + .expect("manager"); assert_eq!(manager.default_environment_id(), None); assert!(manager.default_environment().is_none()); @@ -504,13 +532,16 @@ mod tests { #[tokio::test] async fn environment_manager_rejects_unknown_provider_default() { let provider = TestEnvironmentProvider { - environments: HashMap::from([( - LOCAL_ENVIRONMENT_ID.to_string(), - Environment::default_for_tests(), - )]), - default_environment_id: Some("missing".to_string()), + snapshot: EnvironmentProviderSnapshot { + environments: HashMap::from([( + LOCAL_ENVIRONMENT_ID.to_string(), + Environment::default_for_tests(), + )]), + default: EnvironmentDefault::EnvironmentId("missing".to_string()), + }, }; let err = EnvironmentManager::from_provider(&provider, test_runtime_paths()) + .await .expect_err("unknown default should fail"); assert_eq!( @@ -524,7 +555,8 @@ mod tests { let manager = EnvironmentManager::create_for_tests( /*exec_server_url*/ None, test_runtime_paths(), - ); + ) + .await; assert_eq!(manager.default_environment_id(), Some(LOCAL_ENVIRONMENT_ID)); let provider_local = manager @@ -541,7 +573,8 @@ mod tests { let manager = EnvironmentManager::create_for_tests( /*exec_server_url*/ None, runtime_paths.clone(), - ); + ) + .await; let environment = manager.default_environment().expect("default environment"); @@ -552,7 +585,8 @@ mod tests { .local_runtime_paths() .expect("local runtime paths") .clone(), - ); + ) + .await; let environment = manager.default_environment().expect("default environment"); assert_eq!(environment.local_runtime_paths(), Some(&runtime_paths)); } @@ -568,7 +602,8 @@ mod tests { #[tokio::test] async fn environment_manager_keeps_default_provider_local_lookup_when_default_disabled() { let manager = - EnvironmentManager::create_for_tests(Some("none".to_string()), test_runtime_paths()); + EnvironmentManager::create_for_tests(Some("none".to_string()), test_runtime_paths()) + .await; assert!(manager.default_environment().is_none()); assert_eq!(manager.default_environment_id(), None); diff --git a/codex-rs/exec-server/src/environment_provider.rs b/codex-rs/exec-server/src/environment_provider.rs index 677a5544804d..0e4bcc519162 100644 --- a/codex-rs/exec-server/src/environment_provider.rs +++ b/codex-rs/exec-server/src/environment_provider.rs @@ -1,5 +1,7 @@ use std::collections::HashMap; +use async_trait::async_trait; + use crate::Environment; use crate::ExecServerError; use crate::ExecServerRuntimePaths; @@ -9,19 +11,29 @@ use crate::environment::REMOTE_ENVIRONMENT_ID; /// Lists the concrete environments available to Codex. /// -/// Implementations own both the available environment list and the default -/// environment id. Providers that want the local environment to be addressable -/// by id should include it explicitly in the returned map. +/// Implementations own a startup snapshot containing both the available +/// environment list and default environment selection. Providers that want the +/// local environment to be addressable by id should include it explicitly in +/// the returned map. +#[async_trait] pub trait EnvironmentProvider: Send + Sync { - /// Returns the environments available for a new manager. - fn get_environments( + /// Returns the provider-owned environment startup snapshot. + async fn snapshot( &self, local_runtime_paths: &ExecServerRuntimePaths, - ) -> Result, ExecServerError>; + ) -> Result; +} + +#[derive(Clone, Debug)] +pub struct EnvironmentProviderSnapshot { + pub environments: HashMap, + pub default: EnvironmentDefault, +} - /// Returns the provider-selected default environment id, or `None` to - /// disable model-facing environment access. - fn default_environment_id(&self) -> Option; +#[derive(Clone, Debug, PartialEq, Eq)] +pub enum EnvironmentDefault { + Disabled, + EnvironmentId(String), } /// Default provider backed by `CODEX_EXEC_SERVER_URL`. @@ -41,15 +53,15 @@ impl DefaultEnvironmentProvider { Self::new(std::env::var(CODEX_EXEC_SERVER_URL_ENV_VAR).ok()) } - pub(crate) fn environments( + pub(crate) fn snapshot_inner( &self, local_runtime_paths: &ExecServerRuntimePaths, - ) -> HashMap { + ) -> EnvironmentProviderSnapshot { let mut environments = HashMap::from([( LOCAL_ENVIRONMENT_ID.to_string(), Environment::local(local_runtime_paths.clone()), )]); - let (exec_server_url, _disabled) = normalize_exec_server_url(self.exec_server_url.clone()); + let (exec_server_url, disabled) = normalize_exec_server_url(self.exec_server_url.clone()); if let Some(exec_server_url) = exec_server_url { environments.insert( @@ -58,32 +70,28 @@ impl DefaultEnvironmentProvider { ); } - environments + let default = if disabled { + EnvironmentDefault::Disabled + } else if environments.contains_key(REMOTE_ENVIRONMENT_ID) { + EnvironmentDefault::EnvironmentId(REMOTE_ENVIRONMENT_ID.to_string()) + } else { + EnvironmentDefault::EnvironmentId(LOCAL_ENVIRONMENT_ID.to_string()) + }; + + EnvironmentProviderSnapshot { + environments, + default, + } } } +#[async_trait] impl EnvironmentProvider for DefaultEnvironmentProvider { - fn get_environments( + async fn snapshot( &self, local_runtime_paths: &ExecServerRuntimePaths, - ) -> Result, ExecServerError> { - Ok(self.environments(local_runtime_paths)) - } - - fn default_environment_id(&self) -> Option { - let (exec_server_url, disabled) = normalize_exec_server_url(self.exec_server_url.clone()); - if disabled { - return None; - } - - Some( - if exec_server_url.is_some() { - REMOTE_ENVIRONMENT_ID - } else { - LOCAL_ENVIRONMENT_ID - } - .to_string(), - ) + ) -> Result { + Ok(self.snapshot_inner(local_runtime_paths)) } } @@ -114,9 +122,11 @@ mod tests { async fn default_provider_returns_local_environment_when_url_is_missing() { let provider = DefaultEnvironmentProvider::new(/*exec_server_url*/ None); let runtime_paths = test_runtime_paths(); - let environments = provider - .get_environments(&runtime_paths) + let snapshot = provider + .snapshot(&runtime_paths) + .await .expect("environments"); + let environments = snapshot.environments; assert!(!environments[LOCAL_ENVIRONMENT_ID].is_remote()); assert_eq!( @@ -125,8 +135,8 @@ mod tests { ); assert!(!environments.contains_key(REMOTE_ENVIRONMENT_ID)); assert_eq!( - provider.default_environment_id(), - Some(LOCAL_ENVIRONMENT_ID.to_string()) + snapshot.default, + EnvironmentDefault::EnvironmentId(LOCAL_ENVIRONMENT_ID.to_string()) ); } @@ -134,15 +144,17 @@ mod tests { async fn default_provider_returns_local_environment_when_url_is_empty() { let provider = DefaultEnvironmentProvider::new(Some(String::new())); let runtime_paths = test_runtime_paths(); - let environments = provider - .get_environments(&runtime_paths) + let snapshot = provider + .snapshot(&runtime_paths) + .await .expect("environments"); + let environments = snapshot.environments; assert!(!environments[LOCAL_ENVIRONMENT_ID].is_remote()); assert!(!environments.contains_key(REMOTE_ENVIRONMENT_ID)); assert_eq!( - provider.default_environment_id(), - Some(LOCAL_ENVIRONMENT_ID.to_string()) + snapshot.default, + EnvironmentDefault::EnvironmentId(LOCAL_ENVIRONMENT_ID.to_string()) ); } @@ -150,22 +162,26 @@ mod tests { async fn default_provider_returns_local_environment_for_none_value() { let provider = DefaultEnvironmentProvider::new(Some("none".to_string())); let runtime_paths = test_runtime_paths(); - let environments = provider - .get_environments(&runtime_paths) + let snapshot = provider + .snapshot(&runtime_paths) + .await .expect("environments"); + let environments = snapshot.environments; assert!(!environments[LOCAL_ENVIRONMENT_ID].is_remote()); assert!(!environments.contains_key(REMOTE_ENVIRONMENT_ID)); - assert_eq!(provider.default_environment_id(), None); + assert_eq!(snapshot.default, EnvironmentDefault::Disabled); } #[tokio::test] async fn default_provider_adds_remote_environment_for_websocket_url() { let provider = DefaultEnvironmentProvider::new(Some("ws://127.0.0.1:8765".to_string())); let runtime_paths = test_runtime_paths(); - let environments = provider - .get_environments(&runtime_paths) + let snapshot = provider + .snapshot(&runtime_paths) + .await .expect("environments"); + let environments = snapshot.environments; assert!(!environments[LOCAL_ENVIRONMENT_ID].is_remote()); let remote_environment = &environments[REMOTE_ENVIRONMENT_ID]; @@ -175,8 +191,8 @@ mod tests { Some("ws://127.0.0.1:8765") ); assert_eq!( - provider.default_environment_id(), - Some(REMOTE_ENVIRONMENT_ID.to_string()) + snapshot.default, + EnvironmentDefault::EnvironmentId(REMOTE_ENVIRONMENT_ID.to_string()) ); } @@ -185,11 +201,12 @@ mod tests { let provider = DefaultEnvironmentProvider::new(Some(" ws://127.0.0.1:8765 ".to_string())); let runtime_paths = test_runtime_paths(); let environments = provider - .get_environments(&runtime_paths) + .snapshot(&runtime_paths) + .await .expect("environments"); assert_eq!( - environments[REMOTE_ENVIRONMENT_ID].exec_server_url(), + environments.environments[REMOTE_ENVIRONMENT_ID].exec_server_url(), Some("ws://127.0.0.1:8765") ); } diff --git a/codex-rs/exec/src/lib.rs b/codex-rs/exec/src/lib.rs index dd866a506de0..d61346f1d09e 100644 --- a/codex-rs/exec/src/lib.rs +++ b/codex-rs/exec/src/lib.rs @@ -515,9 +515,9 @@ pub async fn run_main(cli: Cli, arg0_paths: Arg0DispatchPaths) -> anyhow::Result feedback: CodexFeedback::new(), log_db: None, state_db: state_db.clone(), - environment_manager: std::sync::Arc::new(EnvironmentManager::new( - EnvironmentManagerArgs::new(local_runtime_paths), - )), + environment_manager: std::sync::Arc::new( + EnvironmentManager::new(EnvironmentManagerArgs::new(local_runtime_paths)).await, + ), config_warnings, session_source: SessionSource::Exec, enable_codex_api_key_env: true, diff --git a/codex-rs/mcp-server/src/lib.rs b/codex-rs/mcp-server/src/lib.rs index 8fac7e1cdfe6..ac764456f5e5 100644 --- a/codex-rs/mcp-server/src/lib.rs +++ b/codex-rs/mcp-server/src/lib.rs @@ -60,12 +60,15 @@ pub async fn run_main( arg0_paths: Arg0DispatchPaths, cli_config_overrides: CliConfigOverrides, ) -> IoResult<()> { - let environment_manager = Arc::new(EnvironmentManager::new(EnvironmentManagerArgs::new( - ExecServerRuntimePaths::from_optional_paths( - arg0_paths.codex_self_exe.clone(), - arg0_paths.codex_linux_sandbox_exe.clone(), - )?, - ))); + let environment_manager = Arc::new( + EnvironmentManager::new(EnvironmentManagerArgs::new( + ExecServerRuntimePaths::from_optional_paths( + arg0_paths.codex_self_exe.clone(), + arg0_paths.codex_linux_sandbox_exe.clone(), + )?, + )) + .await, + ); // Parse CLI overrides once and derive the base Config eagerly so later // components do not need to work with raw TOML values. let cli_kv_overrides = cli_config_overrides.parse_overrides().map_err(|e| { diff --git a/codex-rs/thread-manager-sample/src/main.rs b/codex-rs/thread-manager-sample/src/main.rs index 2d3e365fbae0..a27c78d54358 100644 --- a/codex-rs/thread-manager-sample/src/main.rs +++ b/codex-rs/thread-manager-sample/src/main.rs @@ -117,9 +117,8 @@ async fn run_main(arg0_paths: Arg0DispatchPaths) -> anyhow::Result<()> { }; let thread_store = thread_store_from_config(&config, state_db.clone()); let agent_graph_store = agent_graph_store_from_state_db(state_db.clone()); - let environment_manager = Arc::new(EnvironmentManager::new(EnvironmentManagerArgs::new( - local_runtime_paths, - ))); + let environment_manager = + Arc::new(EnvironmentManager::new(EnvironmentManagerArgs::new(local_runtime_paths)).await); let thread_manager = ThreadManager::new( &config, auth_manager, diff --git a/codex-rs/tui/src/lib.rs b/codex-rs/tui/src/lib.rs index 5bcc6ffa26af..b2e92a19b485 100644 --- a/codex-rs/tui/src/lib.rs +++ b/codex-rs/tui/src/lib.rs @@ -761,12 +761,15 @@ pub async fn run_main( } }; - let environment_manager = Arc::new(EnvironmentManager::new(EnvironmentManagerArgs::new( - ExecServerRuntimePaths::from_optional_paths( - arg0_paths.codex_self_exe.clone(), - arg0_paths.codex_linux_sandbox_exe.clone(), - )?, - ))); + let environment_manager = Arc::new( + EnvironmentManager::new(EnvironmentManagerArgs::new( + ExecServerRuntimePaths::from_optional_paths( + arg0_paths.codex_self_exe.clone(), + arg0_paths.codex_linux_sandbox_exe.clone(), + )?, + )) + .await, + ); let cwd = cli.cwd.clone(); let config_cwd = config_cwd_for_app_server_target(cwd.as_deref(), &app_server_target, &environment_manager)?; @@ -2138,7 +2141,8 @@ mod tests { std::env::current_exe().expect("current exe"), /*codex_linux_sandbox_exe*/ None, )?, - ); + ) + .await; let config_cwd = config_cwd_for_app_server_target(Some(remote_only_cwd), &target, &environment_manager)?; From 2a5b5ab00be3e27feded1a60b7419a74b2c10e9d Mon Sep 17 00:00:00 2001 From: starr-openai Date: Fri, 1 May 2026 12:08:20 -0700 Subject: [PATCH 36/42] Add CODEX_HOME environments TOML provider Add the environments.toml schema, parser, validation, and provider implementation for configured websocket and stdio-command environments. This keeps the provider load helper available but does not make product entrypoints use it yet. Co-authored-by: Codex --- codex-rs/Cargo.lock | 1 + codex-rs/exec-server/BUILD.bazel | 3 + codex-rs/exec-server/Cargo.toml | 1 + codex-rs/exec-server/src/environment.rs | 31 +- codex-rs/exec-server/src/environment_toml.rs | 673 +++++++++++++++++++ codex-rs/exec-server/src/lib.rs | 1 + 6 files changed, 700 insertions(+), 10 deletions(-) create mode 100644 codex-rs/exec-server/src/environment_toml.rs diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index cd3e93c9cb2c..89ea9611c046 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -2708,6 +2708,7 @@ dependencies = [ "tokio", "tokio-tungstenite", "tokio-util", + "toml 0.9.11+spec-1.1.0", "tracing", "uuid", "wiremock", diff --git a/codex-rs/exec-server/BUILD.bazel b/codex-rs/exec-server/BUILD.bazel index 57ebe041f8cb..237d3a7f2b0d 100644 --- a/codex-rs/exec-server/BUILD.bazel +++ b/codex-rs/exec-server/BUILD.bazel @@ -3,6 +3,9 @@ load("//:defs.bzl", "codex_rust_crate") codex_rust_crate( name = "exec-server", crate_name = "codex_exec_server", + deps_extra = [ + "@crates//:toml", + ], # Keep the crate's integration tests single-threaded under Bazel because # they install process-global test-binary dispatch state, and the remote # exec-server cases already rely on serialization around the full CLI path. diff --git a/codex-rs/exec-server/Cargo.toml b/codex-rs/exec-server/Cargo.toml index 1495397c7828..c466a234c1ed 100644 --- a/codex-rs/exec-server/Cargo.toml +++ b/codex-rs/exec-server/Cargo.toml @@ -28,6 +28,7 @@ serde = { workspace = true, features = ["derive"] } serde_json = { workspace = true } sha2 = { workspace = true } thiserror = { workspace = true } +toml = { workspace = true } tokio = { workspace = true, features = [ "fs", "io-std", diff --git a/codex-rs/exec-server/src/environment.rs b/codex-rs/exec-server/src/environment.rs index d7867c973289..b0d91308a74f 100644 --- a/codex-rs/exec-server/src/environment.rs +++ b/codex-rs/exec-server/src/environment.rs @@ -193,7 +193,7 @@ impl EnvironmentManager { /// paths used by filesystem helpers. #[derive(Clone)] pub struct Environment { - exec_server_url: Option, + remote_transport: Option, exec_backend: Arc, filesystem: Arc, http_client: Arc, @@ -204,7 +204,7 @@ impl Environment { /// Builds a test-only local environment without configured sandbox helper paths. pub fn default_for_tests() -> Self { Self { - exec_server_url: None, + remote_transport: None, exec_backend: Arc::new(LocalProcess::default()), filesystem: Arc::new(LocalFileSystem::unsandboxed()), http_client: Arc::new(ReqwestHttpClient), @@ -216,7 +216,7 @@ impl Environment { impl std::fmt::Debug for Environment { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { f.debug_struct("Environment") - .field("exec_server_url", &self.exec_server_url) + .field("exec_server_url", &self.exec_server_url()) .finish_non_exhaustive() } } @@ -259,7 +259,7 @@ impl Environment { pub(crate) fn local(local_runtime_paths: ExecServerRuntimePaths) -> Self { Self { - exec_server_url: None, + remote_transport: None, exec_backend: Arc::new(LocalProcess::default()), filesystem: Arc::new(LocalFileSystem::with_runtime_paths( local_runtime_paths.clone(), @@ -273,15 +273,23 @@ impl Environment { exec_server_url: String, local_runtime_paths: Option, ) -> Self { - let client = LazyRemoteExecServerClient::new(ExecServerTransportParams::WebSocketUrl( - exec_server_url.clone(), - )); + Self::remote_with_transport( + ExecServerTransportParams::WebSocketUrl(exec_server_url), + local_runtime_paths, + ) + } + + pub(crate) fn remote_with_transport( + transport_params: ExecServerTransportParams, + local_runtime_paths: Option, + ) -> Self { + let client = LazyRemoteExecServerClient::new(transport_params.clone()); let exec_backend: Arc = Arc::new(RemoteProcess::new(client.clone())); let filesystem: Arc = Arc::new(RemoteFileSystem::new(client.clone())); Self { - exec_server_url: Some(exec_server_url), + remote_transport: Some(transport_params), exec_backend, filesystem, http_client: Arc::new(client), @@ -290,12 +298,15 @@ impl Environment { } pub fn is_remote(&self) -> bool { - self.exec_server_url.is_some() + self.remote_transport.is_some() } /// Returns the remote exec-server URL when this environment is remote. pub fn exec_server_url(&self) -> Option<&str> { - self.exec_server_url.as_deref() + match self.remote_transport.as_ref() { + Some(ExecServerTransportParams::WebSocketUrl(url)) => Some(url.as_str()), + Some(ExecServerTransportParams::StdioCommand(_)) | None => None, + } } pub fn local_runtime_paths(&self) -> Option<&ExecServerRuntimePaths> { diff --git a/codex-rs/exec-server/src/environment_toml.rs b/codex-rs/exec-server/src/environment_toml.rs new file mode 100644 index 000000000000..5907b0a17226 --- /dev/null +++ b/codex-rs/exec-server/src/environment_toml.rs @@ -0,0 +1,673 @@ +use std::collections::HashMap; +use std::collections::HashSet; +use std::path::Path; +use std::path::PathBuf; + +use serde::Deserialize; +use tokio_tungstenite::tungstenite::client::IntoClientRequest; + +use crate::DefaultEnvironmentProvider; +use crate::Environment; +use crate::EnvironmentProvider; +use crate::ExecServerError; +use crate::ExecServerRuntimePaths; +use crate::client_api::ExecServerTransportParams; +use crate::client_api::StdioExecServerCommand; +use crate::environment::LOCAL_ENVIRONMENT_ID; + +const ENVIRONMENTS_TOML_FILE: &str = "environments.toml"; +const MAX_ENVIRONMENT_ID_LEN: usize = 64; + +#[derive(Deserialize, Debug, Default)] +#[serde(deny_unknown_fields)] +struct EnvironmentsToml { + default: Option, + + #[serde(default)] + environments: Vec, +} + +#[derive(Deserialize, Debug, Default, PartialEq, Eq)] +#[serde(deny_unknown_fields)] +struct EnvironmentToml { + id: String, + url: Option, + program: Option, + args: Option>, + env: Option>, + cwd: Option, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +struct TomlEnvironmentProvider { + default_environment_id: Option, + environments: HashMap, +} + +impl TomlEnvironmentProvider { + fn new(config: EnvironmentsToml) -> Result { + Self::new_with_config_dir(config, None) + } + + fn new_with_config_dir( + config: EnvironmentsToml, + config_dir: Option<&Path>, + ) -> Result { + let mut ids = HashSet::from([LOCAL_ENVIRONMENT_ID.to_string()]); + let mut environments = HashMap::with_capacity(config.environments.len()); + for item in config.environments { + let (id, transport) = parse_environment_toml(item, config_dir)?; + if !ids.insert(id.clone()) { + return Err(ExecServerError::Protocol(format!( + "environment id `{id}` is duplicated" + ))); + } + environments.insert(id, transport); + } + let default_environment_id = + normalize_default_environment_id(config.default.as_deref(), &ids)?; + Ok(Self { + default_environment_id, + environments, + }) + } +} + +impl EnvironmentProvider for TomlEnvironmentProvider { + fn get_environments( + &self, + local_runtime_paths: &ExecServerRuntimePaths, + ) -> Result, ExecServerError> { + let mut environments = HashMap::from([( + LOCAL_ENVIRONMENT_ID.to_string(), + Environment::local(local_runtime_paths.clone()), + )]); + + for (id, transport_params) in &self.environments { + environments.insert( + id.clone(), + Environment::remote_with_transport( + transport_params.clone(), + Some(local_runtime_paths.clone()), + ), + ); + } + + Ok(environments) + } + + fn default_environment_id(&self) -> Option { + self.default_environment_id.clone() + } +} + +fn parse_environment_toml( + item: EnvironmentToml, + config_dir: Option<&Path>, +) -> Result<(String, ExecServerTransportParams), ExecServerError> { + let EnvironmentToml { + id, + url, + program, + args, + env, + cwd, + } = item; + validate_environment_id(&id)?; + if program.is_none() && (args.is_some() || env.is_some() || cwd.is_some()) { + return Err(ExecServerError::Protocol(format!( + "environment `{id}` args, env, and cwd require program" + ))); + } + + let transport_params = match (url, program) { + (Some(url), None) => { + let url = validate_websocket_url(url)?; + ExecServerTransportParams::WebSocketUrl(url) + } + (None, Some(program)) => { + let program = program.trim().to_string(); + if program.is_empty() { + return Err(ExecServerError::Protocol(format!( + "environment `{id}` program cannot be empty" + ))); + } + let cwd = normalize_stdio_cwd(&id, cwd, config_dir)?; + ExecServerTransportParams::StdioCommand(StdioExecServerCommand { + program, + args: args.unwrap_or_default(), + env: env.unwrap_or_default(), + cwd, + }) + } + (None, None) | (Some(_), Some(_)) => { + return Err(ExecServerError::Protocol(format!( + "environment `{id}` must set exactly one of url or program" + ))); + } + }; + + Ok((id, transport_params)) +} + +fn normalize_stdio_cwd( + id: &str, + cwd: Option, + config_dir: Option<&Path>, +) -> Result, ExecServerError> { + let Some(cwd) = cwd else { + return Ok(None); + }; + if cwd.is_absolute() { + return Ok(Some(cwd)); + } + let Some(config_dir) = config_dir else { + return Err(ExecServerError::Protocol(format!( + "environment `{id}` cwd must be absolute" + ))); + }; + Ok(Some(config_dir.join(cwd))) +} + +pub(crate) fn environment_provider_from_codex_home( + codex_home: &Path, +) -> Result, ExecServerError> { + let path = codex_home.join(ENVIRONMENTS_TOML_FILE); + if !path.try_exists().map_err(|err| { + ExecServerError::Protocol(format!( + "failed to inspect environment config `{}`: {err}", + path.display() + )) + })? { + return Ok(Box::new(DefaultEnvironmentProvider::from_env())); + } + + let environments = load_environments_toml(&path)?; + Ok(Box::new(TomlEnvironmentProvider::new_with_config_dir( + environments, + Some(codex_home), + )?)) +} + +fn normalize_default_environment_id( + default: Option<&str>, + ids: &HashSet, +) -> Result, ExecServerError> { + let Some(default) = default.map(str::trim) else { + return Ok(Some(LOCAL_ENVIRONMENT_ID.to_string())); + }; + if default.is_empty() { + return Err(ExecServerError::Protocol( + "default environment id cannot be empty".to_string(), + )); + } + if !default.eq_ignore_ascii_case("none") && !ids.contains(default) { + return Err(ExecServerError::Protocol(format!( + "default environment `{default}` is not configured" + ))); + } + if default.eq_ignore_ascii_case("none") { + Ok(None) + } else { + Ok(Some(default.to_string())) + } +} + +fn validate_environment_id(id: &str) -> Result<(), ExecServerError> { + let trimmed_id = id.trim(); + if trimmed_id.is_empty() { + return Err(ExecServerError::Protocol( + "environment id cannot be empty".to_string(), + )); + } + if trimmed_id != id { + return Err(ExecServerError::Protocol(format!( + "environment id `{id}` must not contain surrounding whitespace" + ))); + } + if id == LOCAL_ENVIRONMENT_ID || id.eq_ignore_ascii_case("none") { + return Err(ExecServerError::Protocol(format!( + "environment id `{id}` is reserved" + ))); + } + if id.len() > MAX_ENVIRONMENT_ID_LEN { + return Err(ExecServerError::Protocol(format!( + "environment id `{id}` cannot be longer than {MAX_ENVIRONMENT_ID_LEN} characters" + ))); + } + if !id + .chars() + .all(|ch| ch.is_ascii_alphanumeric() || ch == '-' || ch == '_') + { + return Err(ExecServerError::Protocol(format!( + "environment id `{id}` must contain only ASCII letters, numbers, '-' or '_'" + ))); + } + Ok(()) +} + +fn validate_websocket_url(url: String) -> Result { + let url = url.trim(); + if url.is_empty() { + return Err(ExecServerError::Protocol( + "environment url cannot be empty".to_string(), + )); + } + if !url.starts_with("ws://") && !url.starts_with("wss://") { + return Err(ExecServerError::Protocol(format!( + "environment url `{url}` must use ws:// or wss://" + ))); + } + url.into_client_request().map_err(|err| { + ExecServerError::Protocol(format!("environment url `{url}` is invalid: {err}")) + })?; + Ok(url.to_string()) +} + +fn load_environments_toml(path: &Path) -> Result { + let contents = std::fs::read_to_string(path).map_err(|err| { + ExecServerError::Protocol(format!( + "failed to read environment config `{}`: {err}", + path.display() + )) + })?; + + toml::from_str(&contents).map_err(|err| { + ExecServerError::Protocol(format!( + "failed to parse environment config `{}`: {err}", + path.display() + )) + }) +} + +#[cfg(test)] +mod tests { + use pretty_assertions::assert_eq; + use tempfile::tempdir; + + use super::*; + + fn test_runtime_paths() -> ExecServerRuntimePaths { + ExecServerRuntimePaths::new( + std::env::current_exe().expect("current exe"), + /*codex_linux_sandbox_exe*/ None, + ) + .expect("runtime paths") + } + + #[tokio::test] + async fn toml_provider_adds_implicit_local_and_configured_environments() { + let ssh_transport = ExecServerTransportParams::StdioCommand(StdioExecServerCommand { + program: "ssh".to_string(), + args: vec![ + "dev".to_string(), + "codex exec-server --listen stdio".to_string(), + ], + env: HashMap::from([("CODEX_LOG".to_string(), "debug".to_string())]), + cwd: None, + }); + let provider = TomlEnvironmentProvider::new(EnvironmentsToml { + default: Some("ssh-dev".to_string()), + environments: vec![ + EnvironmentToml { + id: "devbox".to_string(), + url: Some(" ws://127.0.0.1:8765 ".to_string()), + ..Default::default() + }, + EnvironmentToml { + id: "ssh-dev".to_string(), + program: Some(" ssh ".to_string()), + args: Some(vec![ + "dev".to_string(), + "codex exec-server --listen stdio".to_string(), + ]), + env: Some(HashMap::from([( + "CODEX_LOG".to_string(), + "debug".to_string(), + )])), + ..Default::default() + }, + ], + }) + .expect("provider"); + let runtime_paths = test_runtime_paths(); + + let environments = provider + .get_environments(&runtime_paths) + .expect("environments"); + + assert!(!environments[LOCAL_ENVIRONMENT_ID].is_remote()); + assert_eq!( + environments["devbox"].exec_server_url(), + Some("ws://127.0.0.1:8765") + ); + assert_eq!(provider.environments["ssh-dev"], ssh_transport); + assert!(environments["ssh-dev"].is_remote()); + assert_eq!( + provider.default_environment_id(), + Some("ssh-dev".to_string()) + ); + } + + #[test] + fn toml_provider_default_omitted_selects_local() { + let provider = TomlEnvironmentProvider::new(EnvironmentsToml::default()).expect("provider"); + + assert_eq!( + provider.default_environment_id, + Some(LOCAL_ENVIRONMENT_ID.to_string()) + ); + } + + #[test] + fn toml_provider_default_none_disables_default() { + let provider = TomlEnvironmentProvider::new(EnvironmentsToml { + default: Some("none".to_string()), + environments: Vec::new(), + }) + .expect("provider"); + + assert_eq!(provider.default_environment_id, None); + } + + #[test] + fn toml_provider_rejects_invalid_environments() { + let cases = [ + ( + EnvironmentToml { + id: "local".to_string(), + url: Some("ws://127.0.0.1:8765".to_string()), + ..Default::default() + }, + "environment id `local` is reserved", + ), + ( + EnvironmentToml { + id: " devbox ".to_string(), + url: Some("ws://127.0.0.1:8765".to_string()), + ..Default::default() + }, + "environment id ` devbox ` must not contain surrounding whitespace", + ), + ( + EnvironmentToml { + id: "dev box".to_string(), + url: Some("ws://127.0.0.1:8765".to_string()), + ..Default::default() + }, + "environment id `dev box` must contain only ASCII letters, numbers, '-' or '_'", + ), + ( + EnvironmentToml { + id: "devbox".to_string(), + url: Some("http://127.0.0.1:8765".to_string()), + ..Default::default() + }, + "environment url `http://127.0.0.1:8765` must use ws:// or wss://", + ), + ( + EnvironmentToml { + id: "devbox".to_string(), + url: Some("ws://127.0.0.1:8765".to_string()), + program: Some("codex".to_string()), + ..Default::default() + }, + "environment `devbox` must set exactly one of url or program", + ), + ( + EnvironmentToml { + id: "devbox".to_string(), + program: Some(" ".to_string()), + ..Default::default() + }, + "environment `devbox` program cannot be empty", + ), + ( + EnvironmentToml { + id: "devbox".to_string(), + args: Some(Vec::new()), + ..Default::default() + }, + "environment `devbox` args, env, and cwd require program", + ), + ]; + + for (item, expected) in cases { + let err = TomlEnvironmentProvider::new(EnvironmentsToml { + default: None, + environments: vec![item], + }) + .expect_err("invalid item should fail"); + + assert_eq!( + err.to_string(), + format!("exec-server protocol error: {expected}") + ); + } + } + + #[test] + fn toml_provider_resolves_relative_stdio_cwd_from_config_dir() { + let config_dir = tempdir().expect("tempdir"); + let provider = TomlEnvironmentProvider::new_with_config_dir( + EnvironmentsToml { + default: None, + environments: vec![EnvironmentToml { + id: "ssh-dev".to_string(), + program: Some("ssh".to_string()), + cwd: Some(PathBuf::from("workspace")), + ..Default::default() + }], + }, + Some(config_dir.path()), + ) + .expect("provider"); + + assert_eq!( + provider.environments["ssh-dev"], + ExecServerTransportParams::StdioCommand(StdioExecServerCommand { + program: "ssh".to_string(), + args: Vec::new(), + env: HashMap::new(), + cwd: Some(config_dir.path().join("workspace")), + }) + ); + } + + #[test] + fn toml_provider_rejects_relative_stdio_cwd_without_config_dir() { + let err = TomlEnvironmentProvider::new(EnvironmentsToml { + default: None, + environments: vec![EnvironmentToml { + id: "ssh-dev".to_string(), + program: Some("ssh".to_string()), + cwd: Some(PathBuf::from("workspace")), + ..Default::default() + }], + }) + .expect_err("relative cwd without config dir should fail"); + + assert_eq!( + err.to_string(), + "exec-server protocol error: environment `ssh-dev` cwd must be absolute" + ); + } + + #[test] + fn toml_provider_rejects_duplicate_ids() { + let err = TomlEnvironmentProvider::new(EnvironmentsToml { + default: None, + environments: vec![ + EnvironmentToml { + id: "devbox".to_string(), + url: Some("ws://127.0.0.1:8765".to_string()), + ..Default::default() + }, + EnvironmentToml { + id: "devbox".to_string(), + program: Some("codex".to_string()), + ..Default::default() + }, + ], + }) + .expect_err("duplicate id should fail"); + + assert_eq!( + err.to_string(), + "exec-server protocol error: environment id `devbox` is duplicated" + ); + } + + #[test] + fn toml_provider_rejects_overlong_id() { + let id = "a".repeat(MAX_ENVIRONMENT_ID_LEN + 1); + let err = TomlEnvironmentProvider::new(EnvironmentsToml { + default: None, + environments: vec![EnvironmentToml { + id: id.clone(), + url: Some("ws://127.0.0.1:8765".to_string()), + ..Default::default() + }], + }) + .expect_err("overlong id should fail"); + + assert_eq!( + err.to_string(), + format!( + "exec-server protocol error: environment id `{id}` cannot be longer than {MAX_ENVIRONMENT_ID_LEN} characters" + ) + ); + } + + #[test] + fn toml_provider_rejects_unknown_default() { + let err = TomlEnvironmentProvider::new(EnvironmentsToml { + default: Some("missing".to_string()), + environments: Vec::new(), + }) + .expect_err("unknown default should fail"); + + assert_eq!( + err.to_string(), + "exec-server protocol error: default environment `missing` is not configured" + ); + } + + #[test] + fn load_environments_toml_reads_root_environment_list() { + let codex_home = tempdir().expect("tempdir"); + let path = codex_home.path().join(ENVIRONMENTS_TOML_FILE); + std::fs::write( + &path, + r#" +default = "ssh-dev" + +[[environments]] +id = "devbox" +url = "ws://127.0.0.1:4512" + +[[environments]] +id = "ssh-dev" +program = "ssh" +args = ["dev", "codex exec-server --listen stdio"] +cwd = "/tmp" +[environments.env] +CODEX_LOG = "debug" +"#, + ) + .expect("write environments.toml"); + + let environments = load_environments_toml(&path).expect("environments.toml"); + + assert_eq!(environments.default.as_deref(), Some("ssh-dev")); + assert_eq!(environments.environments.len(), 2); + assert_eq!(environments.environments[0].id, "devbox"); + assert_eq!( + environments.environments[1], + EnvironmentToml { + id: "ssh-dev".to_string(), + program: Some("ssh".to_string()), + args: Some(vec![ + "dev".to_string(), + "codex exec-server --listen stdio".to_string(), + ]), + env: Some(HashMap::from([( + "CODEX_LOG".to_string(), + "debug".to_string(), + )])), + cwd: Some(PathBuf::from("/tmp")), + ..Default::default() + } + ); + } + + #[test] + fn load_environments_toml_rejects_unknown_fields() { + let codex_home = tempdir().expect("tempdir"); + let cases = [ + ("unknown = true\n", "unknown field `unknown`"), + ( + r#" +[[environments]] +id = "devbox" +url = "ws://127.0.0.1:4512" +unknown = true +"#, + "unknown field `unknown`", + ), + ]; + + for (index, (contents, expected)) in cases.into_iter().enumerate() { + let path = codex_home.path().join(format!("environments-{index}.toml")); + std::fs::write(&path, contents).expect("write environments.toml"); + + let err = load_environments_toml(&path).expect_err("unknown field should fail"); + + assert!( + err.to_string().contains(expected), + "expected `{err}` to contain `{expected}`" + ); + } + } + + #[test] + fn toml_provider_rejects_malformed_websocket_url() { + let err = TomlEnvironmentProvider::new(EnvironmentsToml { + default: None, + environments: vec![EnvironmentToml { + id: "devbox".to_string(), + url: Some("ws://".to_string()), + ..Default::default() + }], + }) + .expect_err("malformed websocket url should fail"); + + assert!( + err.to_string() + .contains("environment url `ws://` is invalid"), + "expected malformed URL error, got `{err}`" + ); + } + + #[tokio::test] + async fn environment_provider_from_codex_home_uses_present_environments_file() { + let codex_home = tempdir().expect("tempdir"); + std::fs::write( + codex_home.path().join(ENVIRONMENTS_TOML_FILE), + r#" +default = "none" +"#, + ) + .expect("write environments.toml"); + + let provider = + environment_provider_from_codex_home(codex_home.path()).expect("environment provider"); + + let environments = provider + .get_environments(&test_runtime_paths()) + .expect("environments"); + + assert!(environments.contains_key(LOCAL_ENVIRONMENT_ID)); + assert_eq!(provider.default_environment_id(), None); + } +} diff --git a/codex-rs/exec-server/src/lib.rs b/codex-rs/exec-server/src/lib.rs index b36ab39d0105..85de8258f2dc 100644 --- a/codex-rs/exec-server/src/lib.rs +++ b/codex-rs/exec-server/src/lib.rs @@ -4,6 +4,7 @@ mod client_transport; mod connection; mod environment; mod environment_provider; +mod environment_toml; mod fs_helper; mod fs_helper_main; mod fs_sandbox; From 1607c0baade567e173c1fe3929b2d12ccca96380 Mon Sep 17 00:00:00 2001 From: starr-openai Date: Tue, 5 May 2026 15:26:29 -0700 Subject: [PATCH 37/42] Fix environments TOML lint coverage Co-authored-by: Codex --- codex-rs/exec-server/src/environment_toml.rs | 16 +++++++++++++++- 1 file changed, 15 insertions(+), 1 deletion(-) diff --git a/codex-rs/exec-server/src/environment_toml.rs b/codex-rs/exec-server/src/environment_toml.rs index 5907b0a17226..cbbd01ff53bc 100644 --- a/codex-rs/exec-server/src/environment_toml.rs +++ b/codex-rs/exec-server/src/environment_toml.rs @@ -46,7 +46,7 @@ struct TomlEnvironmentProvider { impl TomlEnvironmentProvider { fn new(config: EnvironmentsToml) -> Result { - Self::new_with_config_dir(config, None) + Self::new_with_config_dir(config, /*config_dir*/ None) } fn new_with_config_dir( @@ -670,4 +670,18 @@ default = "none" assert!(environments.contains_key(LOCAL_ENVIRONMENT_ID)); assert_eq!(provider.default_environment_id(), None); } + + #[tokio::test] + async fn environment_provider_from_codex_home_falls_back_when_file_is_missing() { + let codex_home = tempdir().expect("tempdir"); + + let provider = + environment_provider_from_codex_home(codex_home.path()).expect("environment provider"); + + let environments = provider + .get_environments(&test_runtime_paths()) + .expect("environments"); + + assert!(environments.contains_key(LOCAL_ENVIRONMENT_ID)); + } } From 5e019e4e3771c4f4580c92d0df9399978a934654 Mon Sep 17 00:00:00 2001 From: starr-openai Date: Tue, 5 May 2026 15:43:44 -0700 Subject: [PATCH 38/42] Limit TOML provider test constructor to tests Avoid keeping the test-only constructor in normal builds now that production construction uses the config-dir aware path. Co-authored-by: Codex --- codex-rs/exec-server/src/environment_toml.rs | 1 + 1 file changed, 1 insertion(+) diff --git a/codex-rs/exec-server/src/environment_toml.rs b/codex-rs/exec-server/src/environment_toml.rs index cbbd01ff53bc..49b55089674e 100644 --- a/codex-rs/exec-server/src/environment_toml.rs +++ b/codex-rs/exec-server/src/environment_toml.rs @@ -45,6 +45,7 @@ struct TomlEnvironmentProvider { } impl TomlEnvironmentProvider { + #[cfg(test)] fn new(config: EnvironmentsToml) -> Result { Self::new_with_config_dir(config, /*config_dir*/ None) } From 551d46ad1905116440932c8bbacf53a7357adebc Mon Sep 17 00:00:00 2001 From: starr-openai Date: Wed, 6 May 2026 13:23:35 -0700 Subject: [PATCH 39/42] Narrow exec server URL accessor Co-authored-by: Codex --- codex-rs/core/tests/common/test_codex.rs | 11 +++++------ codex-rs/exec-server/src/environment.rs | 2 +- codex-rs/exec-server/src/environment_toml.rs | 1 + 3 files changed, 7 insertions(+), 7 deletions(-) diff --git a/codex-rs/core/tests/common/test_codex.rs b/codex-rs/core/tests/common/test_codex.rs index 13dd026654c0..6311e8fe1d93 100644 --- a/codex-rs/core/tests/common/test_codex.rs +++ b/codex-rs/core/tests/common/test_codex.rs @@ -73,6 +73,7 @@ const SUBMIT_TURN_COMPLETE_TIMEOUT: Duration = Duration::from_secs(30); #[derive(Debug)] pub struct TestEnv { environment: codex_exec_server::Environment, + exec_server_url: Option, cwd: AbsolutePathBuf, local_cwd_temp_dir: Option>, remote_container_name: Option, @@ -86,6 +87,7 @@ impl TestEnv { codex_exec_server::Environment::create_for_tests(/*exec_server_url*/ None)?; Ok(Self { environment, + exec_server_url: None, cwd, local_cwd_temp_dir: Some(local_cwd_temp_dir), remote_container_name: None, @@ -100,10 +102,6 @@ impl TestEnv { &self.environment } - pub fn exec_server_url(&self) -> Option<&str> { - self.environment.exec_server_url() - } - fn local_cwd_temp_dir(&self) -> Option> { self.local_cwd_temp_dir.clone() } @@ -123,7 +121,7 @@ pub async fn test_env() -> Result { Some(remote_env) => { let websocket_url = remote_exec_server_url()?; let environment = - codex_exec_server::Environment::create_for_tests(Some(websocket_url))?; + codex_exec_server::Environment::create_for_tests(Some(websocket_url.clone()))?; let cwd = remote_aware_cwd_path(); environment .get_filesystem() @@ -135,6 +133,7 @@ pub async fn test_env() -> Result { .await?; Ok(TestEnv { environment, + exec_server_url: Some(websocket_url), cwd, local_cwd_temp_dir: None, remote_container_name: Some(remote_env.container_name), @@ -385,7 +384,7 @@ impl TestCodexBuilder { let exec_server_url = self .exec_server_url .clone() - .or_else(|| test_env.exec_server_url().map(str::to_owned)); + .or_else(|| test_env.exec_server_url.clone()); let local_runtime_paths = codex_exec_server::ExecServerRuntimePaths::new( std::env::current_exe()?, /*codex_linux_sandbox_exe*/ None, diff --git a/codex-rs/exec-server/src/environment.rs b/codex-rs/exec-server/src/environment.rs index b0d91308a74f..81731af2bfee 100644 --- a/codex-rs/exec-server/src/environment.rs +++ b/codex-rs/exec-server/src/environment.rs @@ -302,7 +302,7 @@ impl Environment { } /// Returns the remote exec-server URL when this environment is remote. - pub fn exec_server_url(&self) -> Option<&str> { + pub(crate) fn exec_server_url(&self) -> Option<&str> { match self.remote_transport.as_ref() { Some(ExecServerTransportParams::WebSocketUrl(url)) => Some(url.as_str()), Some(ExecServerTransportParams::StdioCommand(_)) | None => None, diff --git a/codex-rs/exec-server/src/environment_toml.rs b/codex-rs/exec-server/src/environment_toml.rs index 49b55089674e..d4cac5e85e3c 100644 --- a/codex-rs/exec-server/src/environment_toml.rs +++ b/codex-rs/exec-server/src/environment_toml.rs @@ -344,6 +344,7 @@ mod tests { ); assert_eq!(provider.environments["ssh-dev"], ssh_transport); assert!(environments["ssh-dev"].is_remote()); + assert_eq!(environments["ssh-dev"].exec_server_url(), None); assert_eq!( provider.default_environment_id(), Some("ssh-dev".to_string()) From 9878ced5cc971d2219afb04932a8007e508d04c2 Mon Sep 17 00:00:00 2001 From: starr-openai Date: Wed, 6 May 2026 13:33:13 -0700 Subject: [PATCH 40/42] Align TOML provider with snapshot trait Co-authored-by: Codex --- codex-rs/exec-server/src/environment_toml.rs | 85 ++++++++++++-------- 1 file changed, 52 insertions(+), 33 deletions(-) diff --git a/codex-rs/exec-server/src/environment_toml.rs b/codex-rs/exec-server/src/environment_toml.rs index d4cac5e85e3c..99808d7896cc 100644 --- a/codex-rs/exec-server/src/environment_toml.rs +++ b/codex-rs/exec-server/src/environment_toml.rs @@ -3,6 +3,7 @@ use std::collections::HashSet; use std::path::Path; use std::path::PathBuf; +use async_trait::async_trait; use serde::Deserialize; use tokio_tungstenite::tungstenite::client::IntoClientRequest; @@ -14,6 +15,8 @@ use crate::ExecServerRuntimePaths; use crate::client_api::ExecServerTransportParams; use crate::client_api::StdioExecServerCommand; use crate::environment::LOCAL_ENVIRONMENT_ID; +use crate::environment_provider::EnvironmentDefault; +use crate::environment_provider::EnvironmentProviderSnapshot; const ENVIRONMENTS_TOML_FILE: &str = "environments.toml"; const MAX_ENVIRONMENT_ID_LEN: usize = 64; @@ -40,7 +43,7 @@ struct EnvironmentToml { #[derive(Clone, Debug, PartialEq, Eq)] struct TomlEnvironmentProvider { - default_environment_id: Option, + default: EnvironmentDefault, environments: HashMap, } @@ -65,20 +68,20 @@ impl TomlEnvironmentProvider { } environments.insert(id, transport); } - let default_environment_id = - normalize_default_environment_id(config.default.as_deref(), &ids)?; + let default = normalize_default_environment_id(config.default.as_deref(), &ids)?; Ok(Self { - default_environment_id, + default, environments, }) } } +#[async_trait] impl EnvironmentProvider for TomlEnvironmentProvider { - fn get_environments( + async fn snapshot( &self, local_runtime_paths: &ExecServerRuntimePaths, - ) -> Result, ExecServerError> { + ) -> Result { let mut environments = HashMap::from([( LOCAL_ENVIRONMENT_ID.to_string(), Environment::local(local_runtime_paths.clone()), @@ -94,11 +97,10 @@ impl EnvironmentProvider for TomlEnvironmentProvider { ); } - Ok(environments) - } - - fn default_environment_id(&self) -> Option { - self.default_environment_id.clone() + Ok(EnvironmentProviderSnapshot { + environments, + default: self.default.clone(), + }) } } @@ -193,9 +195,11 @@ pub(crate) fn environment_provider_from_codex_home( fn normalize_default_environment_id( default: Option<&str>, ids: &HashSet, -) -> Result, ExecServerError> { +) -> Result { let Some(default) = default.map(str::trim) else { - return Ok(Some(LOCAL_ENVIRONMENT_ID.to_string())); + return Ok(EnvironmentDefault::EnvironmentId( + LOCAL_ENVIRONMENT_ID.to_string(), + )); }; if default.is_empty() { return Err(ExecServerError::Protocol( @@ -208,9 +212,9 @@ fn normalize_default_environment_id( ))); } if default.eq_ignore_ascii_case("none") { - Ok(None) + Ok(EnvironmentDefault::Disabled) } else { - Ok(Some(default.to_string())) + Ok(EnvironmentDefault::EnvironmentId(default.to_string())) } } @@ -333,9 +337,14 @@ mod tests { .expect("provider"); let runtime_paths = test_runtime_paths(); - let environments = provider - .get_environments(&runtime_paths) + let snapshot = provider + .snapshot(&runtime_paths) + .await .expect("environments"); + let EnvironmentProviderSnapshot { + environments, + default, + } = snapshot; assert!(!environments[LOCAL_ENVIRONMENT_ID].is_remote()); assert_eq!( @@ -346,30 +355,38 @@ mod tests { assert!(environments["ssh-dev"].is_remote()); assert_eq!(environments["ssh-dev"].exec_server_url(), None); assert_eq!( - provider.default_environment_id(), - Some("ssh-dev".to_string()) + default, + EnvironmentDefault::EnvironmentId("ssh-dev".to_string()) ); } - #[test] - fn toml_provider_default_omitted_selects_local() { + #[tokio::test] + async fn toml_provider_default_omitted_selects_local() { let provider = TomlEnvironmentProvider::new(EnvironmentsToml::default()).expect("provider"); + let snapshot = provider + .snapshot(&test_runtime_paths()) + .await + .expect("environments"); assert_eq!( - provider.default_environment_id, - Some(LOCAL_ENVIRONMENT_ID.to_string()) + snapshot.default, + EnvironmentDefault::EnvironmentId(LOCAL_ENVIRONMENT_ID.to_string()) ); } - #[test] - fn toml_provider_default_none_disables_default() { + #[tokio::test] + async fn toml_provider_default_none_disables_default() { let provider = TomlEnvironmentProvider::new(EnvironmentsToml { default: Some("none".to_string()), environments: Vec::new(), }) .expect("provider"); + let snapshot = provider + .snapshot(&test_runtime_paths()) + .await + .expect("environments"); - assert_eq!(provider.default_environment_id, None); + assert_eq!(snapshot.default, EnvironmentDefault::Disabled); } #[test] @@ -665,12 +682,13 @@ default = "none" let provider = environment_provider_from_codex_home(codex_home.path()).expect("environment provider"); - let environments = provider - .get_environments(&test_runtime_paths()) + let snapshot = provider + .snapshot(&test_runtime_paths()) + .await .expect("environments"); - assert!(environments.contains_key(LOCAL_ENVIRONMENT_ID)); - assert_eq!(provider.default_environment_id(), None); + assert!(snapshot.environments.contains_key(LOCAL_ENVIRONMENT_ID)); + assert_eq!(snapshot.default, EnvironmentDefault::Disabled); } #[tokio::test] @@ -680,10 +698,11 @@ default = "none" let provider = environment_provider_from_codex_home(codex_home.path()).expect("environment provider"); - let environments = provider - .get_environments(&test_runtime_paths()) + let snapshot = provider + .snapshot(&test_runtime_paths()) + .await .expect("environments"); - assert!(environments.contains_key(LOCAL_ENVIRONMENT_ID)); + assert!(snapshot.environments.contains_key(LOCAL_ENVIRONMENT_ID)); } } From fc52ee6c5e534c7e273e1a69ac383873cd95b7c0 Mon Sep 17 00:00:00 2001 From: starr-openai Date: Wed, 6 May 2026 14:34:56 -0700 Subject: [PATCH 41/42] Expose CODEX_HOME environment manager constructor Make the environments.toml provider reachable from the exec-server crate API so the provider PR passes clippy before entrypoint wiring lands in the next stack PR. Co-authored-by: Codex --- codex-rs/exec-server/src/environment.rs | 25 +++++++++++++++++++++++++ 1 file changed, 25 insertions(+) diff --git a/codex-rs/exec-server/src/environment.rs b/codex-rs/exec-server/src/environment.rs index 81731af2bfee..d3926059347c 100644 --- a/codex-rs/exec-server/src/environment.rs +++ b/codex-rs/exec-server/src/environment.rs @@ -13,6 +13,7 @@ use crate::environment_provider::EnvironmentDefault; use crate::environment_provider::EnvironmentProvider; use crate::environment_provider::EnvironmentProviderSnapshot; use crate::environment_provider::normalize_exec_server_url; +use crate::environment_toml::environment_provider_from_codex_home; use crate::local_file_system::LocalFileSystem; use crate::local_process::LocalProcess; use crate::process::ExecBackend; @@ -99,6 +100,30 @@ impl EnvironmentManager { Self::from_default_provider_url(exec_server_url, local_runtime_paths).await } + /// Builds a manager from `CODEX_HOME` and local runtime paths used when + /// creating local filesystem helpers. + /// + /// If `CODEX_HOME/environments.toml` is present, it defines the configured + /// environments. Otherwise this preserves the legacy + /// `CODEX_EXEC_SERVER_URL` behavior. Callers that ignore user config + /// should use [`Self::from_env`] instead. + pub async fn from_codex_home( + codex_home: impl AsRef, + local_runtime_paths: ExecServerRuntimePaths, + ) -> Result { + let provider = environment_provider_from_codex_home(codex_home.as_ref())?; + Self::from_provider(provider.as_ref(), local_runtime_paths).await + } + + /// Builds a manager from the legacy environment-variable provider without + /// reading user config files from `CODEX_HOME`. + pub async fn from_env( + local_runtime_paths: ExecServerRuntimePaths, + ) -> Result { + let provider = DefaultEnvironmentProvider::from_env(); + Self::from_provider(&provider, local_runtime_paths).await + } + async fn from_default_provider_url( exec_server_url: Option, local_runtime_paths: ExecServerRuntimePaths, From 4d827aa7800d323d2887c989018e0293d3441063 Mon Sep 17 00:00:00 2001 From: starr-openai Date: Wed, 6 May 2026 19:49:13 -0700 Subject: [PATCH 42/42] Add optional environment resolver hook Introduce an app-server/runtime-facing resolver interface that can be passed through to EnvironmentManager without changing current strict environment-id lookup behavior. Co-authored-by: Codex --- codex-rs/app-server-client/src/lib.rs | 1 + codex-rs/app-server/src/lib.rs | 32 ++++++++++++------- codex-rs/exec-server/src/environment.rs | 29 +++++++++++++++-- .../exec-server/src/environment_resolver.rs | 21 ++++++++++++ codex-rs/exec-server/src/lib.rs | 2 ++ 5 files changed, 71 insertions(+), 14 deletions(-) create mode 100644 codex-rs/exec-server/src/environment_resolver.rs diff --git a/codex-rs/app-server-client/src/lib.rs b/codex-rs/app-server-client/src/lib.rs index 0911cc448dd3..b9e95ef84fd1 100644 --- a/codex-rs/app-server-client/src/lib.rs +++ b/codex-rs/app-server-client/src/lib.rs @@ -50,6 +50,7 @@ pub use codex_core::StateDbHandle; use codex_core::config::Config; pub use codex_exec_server::EnvironmentManager; pub use codex_exec_server::EnvironmentManagerArgs; +pub use codex_exec_server::EnvironmentResolver; pub use codex_exec_server::ExecServerRuntimePaths; use codex_feedback::CodexFeedback; use codex_protocol::protocol::SessionSource; diff --git a/codex-rs/app-server/src/lib.rs b/codex-rs/app-server/src/lib.rs index 4013bbe76bc9..80f00f60807e 100644 --- a/codex-rs/app-server/src/lib.rs +++ b/codex-rs/app-server/src/lib.rs @@ -52,6 +52,7 @@ use codex_core::check_execpolicy_for_warnings; use codex_core::config::find_codex_home; use codex_core::init_state_db_from_config; use codex_exec_server::EnvironmentManager; +use codex_exec_server::EnvironmentResolver; use codex_exec_server::ExecServerRuntimePaths; use codex_feedback::CodexFeedback; use codex_protocol::protocol::SessionSource; @@ -371,15 +372,19 @@ pub enum PluginStartupTasks { Skip, } -#[derive(Debug, Clone, Copy, PartialEq, Eq)] +#[derive(Debug, Clone)] pub struct AppServerRuntimeOptions { pub plugin_startup_tasks: PluginStartupTasks, + /// Optional resolver passed through to the environment manager. App-server + /// still uses strict environment-id lookup today. + pub environment_resolver: Option>, } impl Default for AppServerRuntimeOptions { fn default() -> Self { Self { plugin_startup_tasks: PluginStartupTasks::Start, + environment_resolver: None, } } } @@ -417,15 +422,20 @@ pub async fn run_main_with_transport_options( auth: AppServerWebsocketAuthSettings, runtime_options: AppServerRuntimeOptions, ) -> IoResult<()> { - let environment_manager = Arc::new( - EnvironmentManager::new(EnvironmentManagerArgs::new( - ExecServerRuntimePaths::from_optional_paths( - arg0_paths.codex_self_exe.clone(), - arg0_paths.codex_linux_sandbox_exe.clone(), - )?, - )) - .await, - ); + let AppServerRuntimeOptions { + plugin_startup_tasks, + environment_resolver, + } = runtime_options; + let mut environment_manager_args = + EnvironmentManagerArgs::new(ExecServerRuntimePaths::from_optional_paths( + arg0_paths.codex_self_exe.clone(), + arg0_paths.codex_linux_sandbox_exe.clone(), + )?); + if let Some(environment_resolver) = environment_resolver { + environment_manager_args = + environment_manager_args.with_environment_resolver(environment_resolver); + } + let environment_manager = Arc::new(EnvironmentManager::new(environment_manager_args).await); let (transport_event_tx, mut transport_event_rx) = mpsc::channel::(CHANNEL_CAPACITY); let (outgoing_tx, mut outgoing_rx) = mpsc::channel::(CHANNEL_CAPACITY); @@ -765,7 +775,7 @@ pub async fn run_main_with_transport_options( auth_manager, rpc_transport: analytics_rpc_transport(&transport), remote_control_handle: Some(remote_control_handle.clone()), - plugin_startup_tasks: runtime_options.plugin_startup_tasks, + plugin_startup_tasks, })); let mut thread_created_rx = processor.thread_created_receiver(); let mut running_turn_count_rx = processor.subscribe_running_assistant_turn_count(); diff --git a/codex-rs/exec-server/src/environment.rs b/codex-rs/exec-server/src/environment.rs index d3926059347c..53b2c95f17c3 100644 --- a/codex-rs/exec-server/src/environment.rs +++ b/codex-rs/exec-server/src/environment.rs @@ -13,6 +13,7 @@ use crate::environment_provider::EnvironmentDefault; use crate::environment_provider::EnvironmentProvider; use crate::environment_provider::EnvironmentProviderSnapshot; use crate::environment_provider::normalize_exec_server_url; +use crate::environment_resolver::EnvironmentResolver; use crate::environment_toml::environment_provider_from_codex_home; use crate::local_file_system::LocalFileSystem; use crate::local_process::LocalProcess; @@ -42,6 +43,7 @@ pub struct EnvironmentManager { default_environment: Option, environments: HashMap>, local_environment: Arc, + environment_resolver: Option>, } pub const LOCAL_ENVIRONMENT_ID: &str = "local"; @@ -50,14 +52,26 @@ pub const REMOTE_ENVIRONMENT_ID: &str = "remote"; #[derive(Clone, Debug)] pub struct EnvironmentManagerArgs { pub local_runtime_paths: ExecServerRuntimePaths, + /// Optional resolver supplied by embedding runtimes. It is stored only; + /// environment lookup remains strict until policy is wired explicitly. + environment_resolver: Option>, } impl EnvironmentManagerArgs { pub fn new(local_runtime_paths: ExecServerRuntimePaths) -> Self { Self { local_runtime_paths, + environment_resolver: None, } } + + pub fn with_environment_resolver( + mut self, + environment_resolver: Arc, + ) -> Self { + self.environment_resolver = Some(environment_resolver); + self + } } impl EnvironmentManager { @@ -70,6 +84,7 @@ impl EnvironmentManager { Arc::new(Environment::default_for_tests()), )]), local_environment: Arc::new(Environment::default_for_tests()), + environment_resolver: None, } } @@ -79,6 +94,7 @@ impl EnvironmentManager { default_environment: None, environments: HashMap::new(), local_environment: Arc::new(Environment::local(local_runtime_paths)), + environment_resolver: None, } } @@ -87,7 +103,7 @@ impl EnvironmentManager { exec_server_url: Option, local_runtime_paths: ExecServerRuntimePaths, ) -> Self { - Self::from_default_provider_url(exec_server_url, local_runtime_paths).await + Self::from_default_provider_url(exec_server_url, local_runtime_paths, None).await } /// Builds a manager from `CODEX_EXEC_SERVER_URL` and local runtime paths @@ -95,9 +111,11 @@ impl EnvironmentManager { pub async fn new(args: EnvironmentManagerArgs) -> Self { let EnvironmentManagerArgs { local_runtime_paths, + environment_resolver, } = args; let exec_server_url = std::env::var(CODEX_EXEC_SERVER_URL_ENV_VAR).ok(); - Self::from_default_provider_url(exec_server_url, local_runtime_paths).await + Self::from_default_provider_url(exec_server_url, local_runtime_paths, environment_resolver) + .await } /// Builds a manager from `CODEX_HOME` and local runtime paths used when @@ -127,10 +145,14 @@ impl EnvironmentManager { async fn from_default_provider_url( exec_server_url: Option, local_runtime_paths: ExecServerRuntimePaths, + environment_resolver: Option>, ) -> Self { let provider = DefaultEnvironmentProvider::new(exec_server_url); match Self::from_provider(&provider, local_runtime_paths).await { - Ok(manager) => manager, + Ok(mut manager) => { + manager.environment_resolver = environment_resolver; + manager + } Err(err) => panic!("default provider should create valid environments: {err}"), } } @@ -186,6 +208,7 @@ impl EnvironmentManager { default_environment, environments, local_environment, + environment_resolver: None, }) } diff --git a/codex-rs/exec-server/src/environment_resolver.rs b/codex-rs/exec-server/src/environment_resolver.rs new file mode 100644 index 000000000000..5efa203c3b75 --- /dev/null +++ b/codex-rs/exec-server/src/environment_resolver.rs @@ -0,0 +1,21 @@ +use std::fmt::Debug; + +use async_trait::async_trait; + +use crate::Environment; +use crate::ExecServerError; +use crate::ExecServerRuntimePaths; + +/// Resolves environment ids that are not already present in the manager snapshot. +/// +/// This is an optional extension point for embedders. `EnvironmentManager` +/// stores the resolver, but `get_environment` remains a strict snapshot lookup +/// until resolution policy is wired explicitly. +#[async_trait] +pub trait EnvironmentResolver: Send + Sync + Debug { + async fn resolve_environment( + &self, + environment_id: &str, + local_runtime_paths: &ExecServerRuntimePaths, + ) -> Result, ExecServerError>; +} diff --git a/codex-rs/exec-server/src/lib.rs b/codex-rs/exec-server/src/lib.rs index 85de8258f2dc..4bdb1cb5101c 100644 --- a/codex-rs/exec-server/src/lib.rs +++ b/codex-rs/exec-server/src/lib.rs @@ -4,6 +4,7 @@ mod client_transport; mod connection; mod environment; mod environment_provider; +mod environment_resolver; mod environment_toml; mod fs_helper; mod fs_helper_main; @@ -44,6 +45,7 @@ pub use environment::LOCAL_ENVIRONMENT_ID; pub use environment::REMOTE_ENVIRONMENT_ID; pub use environment_provider::DefaultEnvironmentProvider; pub use environment_provider::EnvironmentProvider; +pub use environment_resolver::EnvironmentResolver; pub use fs_helper::CODEX_FS_HELPER_ARG1; pub use fs_helper_main::main as run_fs_helper_main; pub use local_file_system::LOCAL_FS;