From f4f57bb882a8886b2b7e84892839ce8ed74cca60 Mon Sep 17 00:00:00 2001 From: Sayan Sisodiya Date: Wed, 17 Jun 2026 18:06:42 -0700 Subject: [PATCH] exec-server: start remote environments on registration --- codex-rs/exec-server/src/client.rs | 40 ++++++++++++------- codex-rs/exec-server/src/environment.rs | 53 ++++++++++++++++++++----- 2 files changed, 67 insertions(+), 26 deletions(-) diff --git a/codex-rs/exec-server/src/client.rs b/codex-rs/exec-server/src/client.rs index 30e0185a7a..f2f6bb050c 100644 --- a/codex-rs/exec-server/src/client.rs +++ b/codex-rs/exec-server/src/client.rs @@ -1245,6 +1245,7 @@ mod tests { use tokio::net::TcpStream; use tokio::sync::mpsc; use tokio::sync::oneshot; + use tokio::sync::watch; use tokio::time::Duration; #[cfg(unix)] use tokio::time::sleep; @@ -2047,6 +2048,7 @@ mod tests { listener.local_addr().expect("listener should have address") ); let (replacement_initialized_tx, replacement_initialized_rx) = oneshot::channel(); + let (allow_replacement_tx, allow_replacement_rx) = watch::channel(false); let server = tokio::spawn(async move { let mut first = accept_websocket(&listener).await; complete_websocket_initialize( @@ -2060,17 +2062,20 @@ mod tests { .await .expect("startup websocket should close"); - let (mut failed_reconnect, _) = listener - .accept() + let successful_reconnect = loop { + let (stream, _) = listener.accept().await.expect("reconnect should arrive"); + if *allow_replacement_rx.borrow() { + break stream; + } + let mut failed_reconnect = stream; + failed_reconnect + .write_all(b"HTTP/1.1 500 Internal Server Error\r\nContent-Length: 0\r\n\r\n") + .await + .expect("failed handshake response should write"); + }; + let mut successful_reconnect = accept_async(successful_reconnect) .await - .expect("first reconnect should arrive"); - failed_reconnect - .write_all(b"HTTP/1.1 500 Internal Server Error\r\nContent-Length: 0\r\n\r\n") - .await - .expect("failed handshake response should write"); - drop(failed_reconnect); - - let mut successful_reconnect = accept_websocket(&listener).await; + .expect("replacement websocket handshake should succeed"); complete_websocket_initialize( &mut successful_reconnect, "replacement-session", @@ -2091,11 +2096,16 @@ mod tests { }); let initial = client.get().await.expect("startup should connect"); - wait_for_disconnect(&initial).await; - assert!(matches!( - client.get().await, - Err(super::ExecServerError::ConnectionAttempt(_)) - )); + timeout(Duration::from_secs(1), async { + while !initial.is_disconnected() { + tokio::task::yield_now().await; + } + }) + .await + .expect("client should observe disconnect"); + allow_replacement_tx + .send(true) + .expect("server should allow a fresh client"); let replacement = client.get().await.expect("later reconnect should succeed"); assert_eq!( diff --git a/codex-rs/exec-server/src/environment.rs b/codex-rs/exec-server/src/environment.rs index 76f5d7b2c9..1d77b0c7ae 100644 --- a/codex-rs/exec-server/src/environment.rs +++ b/codex-rs/exec-server/src/environment.rs @@ -47,9 +47,9 @@ pub const CODEX_EXEC_SERVER_NOISE_CHATGPT_ACCOUNT_ID_ENV_VAR: &str = /// use `default_environment().is_some()` as the signal for model-facing /// shell/filesystem tool availability. /// -/// Remote environments create remote filesystem and execution backends that -/// lazy-connect to the configured exec-server on first use. The remote -/// transport is not opened when the manager or environment is constructed. +/// Remote environments begin connecting when added to the manager. Their +/// filesystem and execution backends share that startup result and reconnect +/// after later disconnects as needed. #[derive(Debug)] pub struct EnvironmentManager { default_environment: Option, @@ -221,6 +221,10 @@ impl EnvironmentManager { Some(environment_id) } }; + // The snapshot is valid; start connecting its remote environments in the background. + for environment in environment_map.values() { + environment.start_connecting(); + } Ok(Self { default_environment, environments: RwLock::new(environment_map), @@ -304,12 +308,15 @@ impl EnvironmentManager { "remote environment requires an exec-server url".to_string(), )); }; - let environment = - Environment::remote_inner(exec_server_url, self.local_runtime_paths.clone()); + let environment = Arc::new(Environment::remote_inner( + exec_server_url, + self.local_runtime_paths.clone(), + )); + environment.start_connecting(); self.environments .write() .unwrap_or_else(std::sync::PoisonError::into_inner) - .insert(environment_id, Arc::new(environment)); + .insert(environment_id, environment); Ok(()) } @@ -333,14 +340,15 @@ impl EnvironmentManager { "failed to generate Noise harness identity: {error}" )) })?; - let environment = Environment::remote_with_transport( + let environment = Arc::new(Environment::remote_with_transport( ExecServerTransportParams::NoiseRendezvous { provider, identity }, self.local_runtime_paths.clone(), - ); + )); + environment.start_connecting(); self.environments .write() .unwrap_or_else(std::sync::PoisonError::into_inner) - .insert(environment_id, Arc::new(environment)); + .insert(environment_id, environment); Ok(()) } } @@ -590,6 +598,7 @@ impl From for ShellInfo { #[cfg(test)] mod tests { use std::sync::Arc; + use std::time::Duration; use super::Environment; use super::EnvironmentManager; @@ -602,6 +611,8 @@ mod tests { use crate::environment_provider::EnvironmentProviderSnapshot; use codex_utils_path_uri::PathUri; use pretty_assertions::assert_eq; + use tokio::net::TcpListener; + use tokio::time::timeout; fn test_runtime_paths() -> ExecServerRuntimePaths { ExecServerRuntimePaths::new( @@ -615,8 +626,8 @@ mod tests { assert!(manager.try_local_environment().is_none()); } - #[test] - fn noise_environment_config_selects_remote_as_default() { + #[tokio::test] + async fn noise_environment_config_selects_remote_as_default() { let config = noise_environment_config_from_values( Some("http://registry.example/api".to_string()), Some("environment-requested".to_string()), @@ -971,6 +982,26 @@ mod tests { assert!(!Arc::ptr_eq(&first, &second)); } + #[tokio::test] + async fn environment_manager_starts_remote_environment_when_upserted() { + let listener = TcpListener::bind("127.0.0.1:0") + .await + .expect("bind websocket listener"); + let manager = EnvironmentManager::without_environments(); + + manager + .upsert_environment( + "executor-a".to_string(), + format!("ws://{}", listener.local_addr().expect("listener address")), + ) + .expect("remote environment"); + + timeout(Duration::from_secs(5), listener.accept()) + .await + .expect("environment should start connecting when registered") + .expect("accept connection"); + } + #[tokio::test] async fn environment_manager_rejects_empty_remote_environment_url() { let manager = EnvironmentManager::without_environments();