exec-server: add remote environment connection lifecycle

codex-rs/exec-server/src/client.rs

@@ -14,7 +14,7 @@ use futures::FutureExt;
 use futures::future::BoxFuture;
 use serde_json::Value;
 use tokio::sync::Mutex;
-use tokio::sync::Semaphore;
+use tokio::sync::OnceCell;
 use tokio::sync::mpsc;
 use tokio::sync::watch;
 
@@ -107,6 +107,13 @@ const CONNECT_TIMEOUT: Duration = Duration::from_secs(10);
 const INITIALIZE_TIMEOUT: Duration = Duration::from_secs(10);
 const PROCESS_EVENT_CHANNEL_CAPACITY: usize = 256;
 const PROCESS_EVENT_RETAINED_BYTES: usize = 1024 * 1024;
+const ENVIRONMENT_STARTUP_TIMEOUT: Duration = Duration::from_secs(5 * 60);
+const ENVIRONMENT_INITIAL_RETRY_DELAY: Duration = Duration::from_secs(1);
+const ENVIRONMENT_MAX_RETRY_DELAY: Duration = Duration::from_secs(30);
+
+// ThreadEnvironments::snapshot() currently waits for Environment::info(), so this
+// must land with the follow-up that stops snapshots from waiting on starting environments.
+// Otherwise an unavailable executor can delay session startup for five minutes.
 
 impl Default for ExecServerClientConnectOptions {
     fn default() -> Self {
@@ -226,54 +233,125 @@ impl Drop for ActiveProcessStart {
     }
 }
 
+type ConnectionResult = Result<ExecServerClient, Arc<ExecServerError>>;
+type ConnectionAttempt = OnceCell<ConnectionResult>;
+
 #[derive(Clone)]
 pub(crate) struct LazyRemoteExecServerClient {
     transport_params: ExecServerTransportParams,
-    client: Arc<StdMutex<Option<ExecServerClient>>>,
-    connect_lock: Arc<Semaphore>,
+    // Saves the first startup result so callers share it and failures remain final.
+    startup: Arc<ConnectionAttempt>,
+    // The latest successful client, replaced whenever reconnecting succeeds.
+    current_client: Arc<StdMutex<Option<ExecServerClient>>>,
+    reconnect: Arc<StdMutex<Option<Arc<ConnectionAttempt>>>>,
 }
 
 impl LazyRemoteExecServerClient {
     pub(crate) fn new(transport_params: ExecServerTransportParams) -> Self {
         Self {
             transport_params,
-            client: Arc::new(StdMutex::new(None)),
-            connect_lock: Arc::new(Semaphore::new(/*permits*/ 1)),
+            startup: Arc::new(ConnectionAttempt::new()),
+            current_client: Arc::new(StdMutex::new(None)),
+            reconnect: Arc::new(StdMutex::new(None)),
         }
     }
 
+    pub(crate) fn start_connecting(&self) {
+        let client = self.clone();
+        drop(tokio::spawn(async move {
+            if let Err(error) = client.wait_until_ready().await {
+                debug!(%error, "exec-server environment startup failed");
+            }
+        }));
+    }
+
+    pub(crate) fn startup_finished(&self) -> bool {
+        self.startup.get().is_some()
+    }
+
+    pub(crate) async fn wait_until_ready(&self) -> Result<(), ExecServerError> {
+        self.initial_client().await.map(drop)
+    }
+
     pub(crate) async fn get(&self) -> Result<ExecServerClient, ExecServerError> {
         if let Some(client) = self.connected_client() {
             return Ok(client);
         }
 
-        let _connect_permit = self.connect_lock.acquire().await.map_err(|_| {
-            ExecServerError::Protocol("exec-server connect lock closed".to_string())
-        })?;
-        if let Some(client) = self.connected_client() {
-            return Ok(client);
-        }
-
-        let next_client = match self.cached_client() {
-            Some(_client)
-                if matches!(
-                    &self.transport_params,
-                    ExecServerTransportParams::WebSocketUrl { .. }
-                        | ExecServerTransportParams::NoiseRendezvous { .. }
-                ) =>
-            {
-                ExecServerClient::connect_for_transport(self.transport_params.clone()).await?
+        let Some(cached_client) = self.cached_client() else {
+            let client = self.initial_client().await?;
+            if !client.is_disconnected() || !self.can_reconnect() {
+                return Ok(client);
             }
-            Some(client) => return Ok(client),
-            None => ExecServerClient::connect_for_transport(self.transport_params.clone()).await?,
+            return self.reconnect().await;
         };
 
-        let mut cached_client = self
-            .client
+        if !self.can_reconnect() {
+            return Ok(cached_client);
+        }
+
+        self.reconnect().await
+    }
+
+    async fn initial_client(&self) -> Result<ExecServerClient, ExecServerError> {
+        // The first caller starts the work; every other caller waits for that same result.
+        let result = self
+            .startup
+            .get_or_init(|| connect_with_startup_retries(self.transport_params.clone()))
+            .await;
+        match result {
+            Ok(client) => {
+                let mut current_client = self
+                    .current_client
+                    .lock()
+                    .unwrap_or_else(std::sync::PoisonError::into_inner);
+                if current_client.is_none() {
+                    *current_client = Some(client.clone());
+                }
+                Ok(client.clone())
+            }
+            Err(error) => Err(ExecServerError::ConnectionAttempt(Arc::clone(error))),
+        }
+    }
+
+    async fn reconnect(&self) -> Result<ExecServerClient, ExecServerError> {
+        // Callers handling the same outage share one reconnect attempt.
+        let attempt = {
+            let mut reconnect = self
+                .reconnect
+                .lock()
+                .unwrap_or_else(std::sync::PoisonError::into_inner);
+            if let Some(client) = self.connected_client() {
+                return Ok(client);
+            }
+            reconnect
+                .get_or_insert_with(|| Arc::new(ConnectionAttempt::new()))
+                .clone()
+        };
+        let result = attempt
+            .get_or_init(|| async {
+                let result = connect_once(self.transport_params.clone()).await;
+                if let Ok(client) = &result {
+                    *self
+                        .current_client
+                        .lock()
+                        .unwrap_or_else(std::sync::PoisonError::into_inner) = Some(client.clone());
+                }
+                result
+            })
+            .await;
+        let mut reconnect = self
+            .reconnect
             .lock()
             .unwrap_or_else(std::sync::PoisonError::into_inner);
-        *cached_client = Some(next_client.clone());
-        Ok(next_client)
+        // Forget only this completed attempt so a later operation can retry after failure.
+        if reconnect
+            .as_ref()
+            .is_some_and(|current| Arc::ptr_eq(current, &attempt))
+        {
+            *reconnect = None;
+        }
+        result.clone().map_err(ExecServerError::ConnectionAttempt)
     }
 
     fn connected_client(&self) -> Option<ExecServerClient> {
@@ -282,11 +360,60 @@ impl LazyRemoteExecServerClient {
     }
 
     fn cached_client(&self) -> Option<ExecServerClient> {
-        self.client
+        self.current_client
             .lock()
             .unwrap_or_else(std::sync::PoisonError::into_inner)
             .clone()
     }
+
+    fn can_reconnect(&self) -> bool {
+        matches!(
+            self.transport_params,
+            ExecServerTransportParams::WebSocketUrl { .. }
+                | ExecServerTransportParams::NoiseRendezvous { .. }
+        )
+    }
+}
+
+async fn connect_with_startup_retries(
+    transport_params: ExecServerTransportParams,
+) -> ConnectionResult {
+    if matches!(
+        transport_params,
+        ExecServerTransportParams::StdioCommand { .. }
+    ) {
+        return connect_once(transport_params).await;
+    }
+
+    let startup = async {
+        let mut retry_delay = ENVIRONMENT_INITIAL_RETRY_DELAY;
+        loop {
+            match ExecServerClient::connect_for_transport(transport_params.clone()).await {
+                Ok(client) => return Ok(client),
+                Err(error) => {
+                    debug!(
+                        %error,
+                        retry_in = ?retry_delay,
+                        "exec-server environment is not ready; retrying"
+                    );
+                    tokio::time::sleep(retry_delay).await;
+                    retry_delay = (retry_delay * 2).min(ENVIRONMENT_MAX_RETRY_DELAY);
+                }
+            }
+        }
+    };
+    match timeout(ENVIRONMENT_STARTUP_TIMEOUT, startup).await {
+        Ok(result) => result,
+        Err(_) => Err(Arc::new(ExecServerError::StartupTimedOut {
+            timeout: ENVIRONMENT_STARTUP_TIMEOUT,
+        })),
+    }
+}
+
+async fn connect_once(transport_params: ExecServerTransportParams) -> ConnectionResult {
+    ExecServerClient::connect_for_transport(transport_params)
+        .await
+        .map_err(Arc::new)
 }
 
 impl HttpClient for LazyRemoteExecServerClient {
@@ -352,6 +479,10 @@ pub enum ExecServerError {
     EnvironmentRegistryAuth(String),
     #[error("environment registry request failed: {0}")]
     EnvironmentRegistryRequest(#[from] reqwest::Error),
+    #[error("exec-server connection attempt failed: {0}")]
+    ConnectionAttempt(#[source] Arc<ExecServerError>),
+    #[error("exec-server did not become ready within {timeout:?}")]
+    StartupTimedOut { timeout: Duration },
 }
 
 impl ExecServerClient {
@@ -1773,6 +1904,164 @@ mod tests {
         server.await.expect("server task should finish");
     }
 
+    #[tokio::test]
+    async fn initial_connection_retries_once_and_is_shared_by_all_waiters() {
+        let listener = TcpListener::bind("127.0.0.1:0")
+            .await
+            .expect("listener should bind");
+        let websocket_url = format!(
+            "ws://{}",
+            listener.local_addr().expect("listener should have address")
+        );
+        let server = tokio::spawn(async move {
+            let (first, _) = listener
+                .accept()
+                .await
+                .expect("first connection should arrive");
+            drop(first);
+
+            let mut second = accept_websocket(&listener).await;
+            complete_websocket_initialize(
+                &mut second,
+                "startup-session",
+                /*expected_resume_session_id*/ None,
+            )
+            .await;
+            timeout(Duration::from_secs(1), second.next())
+                .await
+                .expect("client should close after the test");
+        });
+        let client = LazyRemoteExecServerClient::new(ExecServerTransportParams::WebSocketUrl {
+            websocket_url,
+            connect_timeout: Duration::from_secs(1),
+            initialize_timeout: Duration::from_secs(1),
+        });
+
+        assert!(!client.startup_finished());
+        client.start_connecting();
+        let (ready, first, second) =
+            tokio::join!(client.wait_until_ready(), client.get(), client.get());
+        ready.expect("background startup should finish");
+        let first = first.expect("first waiter should receive the client");
+        let second = second.expect("second waiter should receive the same client");
+
+        assert!(client.startup_finished());
+        assert_eq!(first.session_id().as_deref(), Some("startup-session"));
+        assert!(Arc::ptr_eq(&first.inner, &second.inner));
+
+        drop(first);
+        drop(second);
+        drop(client);
+        server.await.expect("server task should finish");
+    }
+
+    #[tokio::test]
+    async fn terminal_stdio_startup_failure_is_remembered() {
+        let client = LazyRemoteExecServerClient::new(ExecServerTransportParams::StdioCommand {
+            command: StdioExecServerCommand {
+                program: "codex-missing-exec-server-for-test".to_string(),
+                args: Vec::new(),
+                env: HashMap::new(),
+                cwd: None,
+            },
+            initialize_timeout: Duration::from_secs(1),
+        });
+
+        let first = match client.get().await {
+            Ok(_) => panic!("missing executable should fail"),
+            Err(error) => error,
+        };
+        assert!(client.startup_finished());
+        let second = match client.get().await {
+            Ok(_) => panic!("burned environment should stay failed"),
+            Err(error) => error,
+        };
+
+        let (
+            super::ExecServerError::ConnectionAttempt(first),
+            super::ExecServerError::ConnectionAttempt(second),
+        ) = (first, second)
+        else {
+            panic!("expected saved connection failures");
+        };
+        assert!(Arc::ptr_eq(&first, &second));
+    }
+
+    #[tokio::test]
+    async fn failed_reconnect_does_not_burn_environment() {
+        let listener = TcpListener::bind("127.0.0.1:0")
+            .await
+            .expect("listener should bind");
+        let websocket_url = format!(
+            "ws://{}",
+            listener.local_addr().expect("listener should have address")
+        );
+        let (replacement_initialized_tx, replacement_initialized_rx) = oneshot::channel();
+        let server = tokio::spawn(async move {
+            let mut first = accept_websocket(&listener).await;
+            complete_websocket_initialize(
+                &mut first,
+                "startup-session",
+                /*expected_resume_session_id*/ None,
+            )
+            .await;
+            first
+                .close(None)
+                .await
+                .expect("startup websocket should close");
+
+            let (mut failed_reconnect, _) = listener
+                .accept()
+                .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;
+            complete_websocket_initialize(
+                &mut successful_reconnect,
+                "replacement-session",
+                /*expected_resume_session_id*/ None,
+            )
+            .await;
+            replacement_initialized_tx
+                .send(())
+                .expect("replacement initialization should be observed");
+            timeout(Duration::from_secs(1), successful_reconnect.next())
+                .await
+                .expect("client should close after the test");
+        });
+        let client = LazyRemoteExecServerClient::new(ExecServerTransportParams::WebSocketUrl {
+            websocket_url,
+            connect_timeout: Duration::from_secs(1),
+            initialize_timeout: Duration::from_secs(1),
+        });
+
+        let initial = client.get().await.expect("startup should connect");
+        wait_for_disconnect(&initial).await;
+        assert!(matches!(
+            client.get().await,
+            Err(super::ExecServerError::ConnectionAttempt(_))
+        ));
+        let replacement = client.get().await.expect("later reconnect should succeed");
+
+        assert_eq!(
+            replacement.session_id().as_deref(),
+            Some("replacement-session")
+        );
+        replacement_initialized_rx
+            .await
+            .expect("server should observe replacement initialization");
+
+        drop(initial);
+        drop(replacement);
+        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);

codex-rs/exec-server/src/environment.rs

@@ -2,9 +2,6 @@ use std::collections::HashMap;
 use std::sync::Arc;
 use std::sync::RwLock;
 
-use futures::FutureExt;
-use futures::future::BoxFuture;
-
 use crate::ExecServerError;
 use crate::ExecServerRuntimePaths;
 use crate::ExecutorFileSystem;
@@ -402,50 +399,19 @@ fn optional_environment_value(name: &str) -> Option<String> {
 #[derive(Clone)]
 pub struct Environment {
     exec_server_url: Option<String>,
-    remote_transport: Option<ExecServerTransportParams>,
-    info_provider: Arc<dyn EnvironmentInfoProvider>,
+    remote_client: Option<LazyRemoteExecServerClient>,
     exec_backend: Arc<dyn ExecBackend>,
     filesystem: Arc<dyn ExecutorFileSystem>,
     http_client: Arc<dyn HttpClient>,
     local_runtime_paths: Option<ExecServerRuntimePaths>,
 }
 
-/// Provides environment metadata from either a local environment or a remote exec-server.
-trait EnvironmentInfoProvider: Send + Sync {
-    fn info(&self) -> BoxFuture<'_, Result<EnvironmentInfo, ExecServerError>>;
-}
-
-struct LocalEnvironmentInfoProvider;
-
-impl EnvironmentInfoProvider for LocalEnvironmentInfoProvider {
-    fn info(&self) -> BoxFuture<'_, Result<EnvironmentInfo, ExecServerError>> {
-        std::future::ready(Ok(EnvironmentInfo::local())).boxed()
-    }
-}
-
-struct RemoteEnvironmentInfoProvider {
-    client: LazyRemoteExecServerClient,
-}
-
-impl RemoteEnvironmentInfoProvider {
-    fn new(client: LazyRemoteExecServerClient) -> Self {
-        Self { client }
-    }
-}
-
-impl EnvironmentInfoProvider for RemoteEnvironmentInfoProvider {
-    fn info(&self) -> BoxFuture<'_, Result<EnvironmentInfo, ExecServerError>> {
-        async move { self.client.environment_info().await }.boxed()
-    }
-}
-
 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,
-            info_provider: Arc::new(LocalEnvironmentInfoProvider),
+            remote_client: None,
             exec_backend: Arc::new(LocalProcess::default()),
             filesystem: Arc::new(LocalFileSystem::unsandboxed()),
             http_client: Arc::new(ReqwestHttpClient),
@@ -501,8 +467,7 @@ impl Environment {
     pub(crate) fn local(local_runtime_paths: ExecServerRuntimePaths) -> Self {
         Self {
             exec_server_url: None,
-            remote_transport: None,
-            info_provider: Arc::new(LocalEnvironmentInfoProvider),
+            remote_client: None,
             exec_backend: Arc::new(LocalProcess::default()),
             filesystem: Arc::new(LocalFileSystem::with_runtime_paths(
                 local_runtime_paths.clone(),
@@ -534,15 +499,14 @@ impl Environment {
             ExecServerTransportParams::NoiseRendezvous { .. } => None,
             ExecServerTransportParams::StdioCommand { .. } => None,
         };
-        let client = LazyRemoteExecServerClient::new(remote_transport.clone());
+        let client = LazyRemoteExecServerClient::new(remote_transport);
         let exec_backend: Arc<dyn ExecBackend> = Arc::new(RemoteProcess::new(client.clone()));
         let filesystem: Arc<dyn ExecutorFileSystem> =
             Arc::new(RemoteFileSystem::new(client.clone()));
 
         Self {
             exec_server_url,
-            remote_transport: Some(remote_transport),
-            info_provider: Arc::new(RemoteEnvironmentInfoProvider::new(client.clone())),
+            remote_client: Some(client.clone()),
             exec_backend,
             filesystem,
             http_client: Arc::new(client),
@@ -551,7 +515,7 @@ impl Environment {
     }
 
     pub fn is_remote(&self) -> bool {
-        self.remote_transport.is_some()
+        self.remote_client.is_some()
     }
 
     /// Returns the remote exec-server URL when this environment is remote.
@@ -565,7 +529,32 @@ impl Environment {
 
     /// Returns environment information from the selected execution/filesystem environment.
     pub async fn info(&self) -> Result<EnvironmentInfo, ExecServerError> {
-        self.info_provider.info().await
+        match &self.remote_client {
+            Some(client) => client.environment_info().await,
+            None => Ok(EnvironmentInfo::local()),
+        }
+    }
+
+    /// Starts connecting a remote environment without waiting for it.
+    pub fn start_connecting(&self) {
+        if let Some(client) = &self.remote_client {
+            client.start_connecting();
+        }
+    }
+
+    /// Returns whether initial startup has either succeeded or permanently failed.
+    pub fn startup_finished(&self) -> bool {
+        self.remote_client
+            .as_ref()
+            .is_none_or(LazyRemoteExecServerClient::startup_finished)
+    }
+
+    /// Waits for initial startup. A failed startup is never attempted again.
+    pub async fn wait_until_ready(&self) -> Result<(), ExecServerError> {
+        match &self.remote_client {
+            Some(client) => client.wait_until_ready().await,
+            None => Ok(()),
+        }
     }
 
     pub fn get_exec_backend(&self) -> Arc<dyn ExecBackend> {

codex-rs/exec-server/tests/relay.rs

@@ -6,8 +6,6 @@ mod relay_proto;
 use std::collections::HashMap;
 use std::sync::Arc;
 use std::sync::Mutex;
-use std::sync::atomic::AtomicUsize;
-use std::sync::atomic::Ordering;
 use std::time::Duration;
 
 use anyhow::Context;
@@ -43,6 +41,7 @@ use relay_proto::relay_message_frame;
 use tempfile::TempDir;
 use tokio::net::TcpListener;
 use tokio::net::TcpStream;
+use tokio::sync::mpsc;
 use tokio::time::timeout;
 use tokio_tungstenite::WebSocketStream;
 use tokio_tungstenite::accept_async;
@@ -73,7 +72,7 @@ impl AuthProvider for StaticRegistryAuthProvider {
 }
 
 struct FailingNoiseConnectProvider {
-    attempts: Arc<AtomicUsize>,
+    attempt_tx: mpsc::UnboundedSender<()>,
 }
 
 impl NoiseRendezvousConnectProvider for FailingNoiseConnectProvider {
@@ -81,7 +80,7 @@ impl NoiseRendezvousConnectProvider for FailingNoiseConnectProvider {
         &self,
         _: NoiseChannelPublicKey,
     ) -> BoxFuture<'_, Result<NoiseRendezvousConnectBundle, ExecServerError>> {
-        self.attempts.fetch_add(1, Ordering::SeqCst);
+        let _ = self.attempt_tx.send(());
         async {
             Err(ExecServerError::Protocol(
                 "test registry connect failure".to_string(),
@@ -96,41 +95,52 @@ fn static_registry_auth_provider() -> codex_api::SharedAuthProvider {
 }
 
 #[tokio::test]
-async fn noise_environment_refreshes_bundle_for_each_connection_attempt() -> Result<()> {
-    let attempts = Arc::new(AtomicUsize::new(0));
+async fn noise_environment_refreshes_bundle_during_startup_retries() -> Result<()> {
+    let (attempt_tx, mut attempt_rx) = mpsc::unbounded_channel();
     let manager = EnvironmentManager::without_environments();
     manager.upsert_noise_environment(
         ENVIRONMENT_ID.to_string(),
-        Arc::new(FailingNoiseConnectProvider {
-            attempts: Arc::clone(&attempts),
-        }),
+        Arc::new(FailingNoiseConnectProvider { attempt_tx }),
     )?;
     let backend = manager
         .get_environment(ENVIRONMENT_ID)
         .context("Noise environment should be materialized")?
         .get_exec_backend();
+    let cwd = PathUri::from_path(std::env::current_dir()?)?;
 
-    for attempt in 1..=2 {
-        let result = backend
+    let startup = tokio::spawn(async move {
+        backend
             .start(ExecParams {
-                process_id: ProcessId::new(format!("proc-{attempt}")),
+                process_id: ProcessId::from("proc-1"),
                 argv: vec!["true".to_string()],
-                cwd: PathUri::from_path(std::env::current_dir()?)?,
+                cwd,
                 env_policy: None,
                 env: HashMap::new(),
                 tty: false,
                 pipe_stdin: false,
                 arg0: None,
             })
-            .await;
-        assert!(matches!(
-            result,
-            Err(ExecServerError::Protocol(ref message))
-                if message == "test registry connect failure"
-        ));
-    }
+            .await
+    });
 
-    assert_eq!(attempts.load(Ordering::SeqCst), 2);
+    timeout(TEST_TIMEOUT, async {
+        for _ in 0..2 {
+            attempt_rx
+                .recv()
+                .await
+                .context("connection provider should remain available")?;
+        }
+        Ok::<_, anyhow::Error>(())
+    })
+    .await
+    .context("startup should retry the Noise connection")??;
+
+    startup.abort();
+    let cancellation = match startup.await {
+        Err(error) => error,
+        Ok(_) => panic!("startup should be cancelled"),
+    };
+    assert!(cancellation.is_cancelled());
     Ok(())
 }