diff --git a/codex-rs/exec-server/src/client.rs b/codex-rs/exec-server/src/client.rs index ae77dc7c4f..5dcaafa130 100644 --- a/codex-rs/exec-server/src/client.rs +++ b/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>; +type ConnectionAttempt = OnceCell; + #[derive(Clone)] pub(crate) struct LazyRemoteExecServerClient { transport_params: ExecServerTransportParams, - client: Arc>>, - connect_lock: Arc, + // Saves the first startup result so callers share it and failures remain final. + startup: Arc, + // The latest successful client, replaced whenever reconnecting succeeds. + current_client: Arc>>, + reconnect: Arc>>>, } 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 { 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 { + // 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 { + // 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 { @@ -282,11 +360,60 @@ impl LazyRemoteExecServerClient { } fn cached_client(&self) -> Option { - 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), + #[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); diff --git a/codex-rs/exec-server/src/environment.rs b/codex-rs/exec-server/src/environment.rs index 6159fb2dd2..76f5d7b2c9 100644 --- a/codex-rs/exec-server/src/environment.rs +++ b/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 { #[derive(Clone)] pub struct Environment { exec_server_url: Option, - remote_transport: Option, - info_provider: Arc, + remote_client: Option, exec_backend: Arc, filesystem: Arc, http_client: Arc, local_runtime_paths: Option, } -/// Provides environment metadata from either a local environment or a remote exec-server. -trait EnvironmentInfoProvider: Send + Sync { - fn info(&self) -> BoxFuture<'_, Result>; -} - -struct LocalEnvironmentInfoProvider; - -impl EnvironmentInfoProvider for LocalEnvironmentInfoProvider { - fn info(&self) -> BoxFuture<'_, Result> { - 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> { - 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 = Arc::new(RemoteProcess::new(client.clone())); let filesystem: Arc = 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 { - 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 { diff --git a/codex-rs/exec-server/tests/relay.rs b/codex-rs/exec-server/tests/relay.rs index 114fe122a2..cbb0259ec7 100644 --- a/codex-rs/exec-server/tests/relay.rs +++ b/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, + attempt_tx: mpsc::UnboundedSender<()>, } impl NoiseRendezvousConnectProvider for FailingNoiseConnectProvider { @@ -81,7 +80,7 @@ impl NoiseRendezvousConnectProvider for FailingNoiseConnectProvider { &self, _: NoiseChannelPublicKey, ) -> BoxFuture<'_, Result> { - 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(()) }