exec-server: start remote environments on registration

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!(

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<String>,
@@ -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<DetectedShell> 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();