diff --git a/codex-rs/exec-server/src/client.rs b/codex-rs/exec-server/src/client.rs index 325fb7442..88c0c8034 100644 --- a/codex-rs/exec-server/src/client.rs +++ b/codex-rs/exec-server/src/client.rs @@ -14,9 +14,10 @@ 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; +use tokio_util::task::AbortOnDropHandle; use tokio::time::timeout; use tracing::debug; @@ -227,54 +228,132 @@ 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) -> Option> { + // Stdio starts a process, so keep it lazy until the environment is used. + if matches!( + self.transport_params, + ExecServerTransportParams::StdioCommand { .. } + ) { + return None; + } + let client = self.clone(); + Some(AbortOnDropHandle::new(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_once(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 { @@ -283,11 +362,25 @@ 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_once(transport_params: ExecServerTransportParams) -> ConnectionResult { + ExecServerClient::connect_for_transport(transport_params) + .await + .map_err(Arc::new) } impl HttpClient for LazyRemoteExecServerClient { @@ -353,6 +446,8 @@ 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), } impl ExecServerClient { @@ -1112,6 +1207,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; @@ -1919,6 +2015,177 @@ mod tests { server.await.expect("server task should finish"); } + #[tokio::test] + async fn initial_connection_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 mut connection = accept_websocket(&listener).await; + complete_websocket_initialize( + &mut connection, + "startup-session", + /*expected_resume_session_id*/ None, + ) + .await; + timeout(Duration::from_secs(1), connection.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()); + let _startup_task = 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), + }); + + assert!(client.start_connecting().is_none()); + assert!(!client.startup_finished()); + 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 (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( + &mut first, + "startup-session", + /*expected_resume_session_id*/ None, + ) + .await; + first + .close(None) + .await + .expect("startup websocket should close"); + + 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("replacement websocket handshake should succeed"); + 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"); + timeout(Duration::from_secs(1), async { + while !initial.is_disconnected() { + tokio::task::yield_now().await; + } + }) + .await + .expect("client should observe disconnect"); + let failed_reconnect = match client.get().await { + Ok(_) => panic!("first lazy reconnect should fail"), + Err(error) => error, + }; + assert!(matches!( + failed_reconnect, + super::ExecServerError::ConnectionAttempt(_) + )); + allow_replacement_tx + .send(true) + .expect("server should allow a fresh client"); + 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 6159fb2dd..03d0d996c 100644 --- a/codex-rs/exec-server/src/environment.rs +++ b/codex-rs/exec-server/src/environment.rs @@ -1,10 +1,8 @@ use std::collections::HashMap; use std::sync::Arc; +use std::sync::Mutex; use std::sync::RwLock; -use futures::FutureExt; -use futures::future::BoxFuture; - use crate::ExecServerError; use crate::ExecServerRuntimePaths; use crate::ExecutorFileSystem; @@ -29,6 +27,7 @@ use crate::remote::NoiseRendezvousEnvironmentConfig; use crate::remote_file_system::RemoteFileSystem; use crate::remote_process::RemoteProcess; use codex_shell_command::shell_detect::DetectedShell; +use tokio_util::task::AbortOnDropHandle; pub const CODEX_EXEC_SERVER_URL_ENV_VAR: &str = "CODEX_EXEC_SERVER_URL"; pub const CODEX_EXEC_SERVER_NOISE_REGISTRY_URL_ENV_VAR: &str = @@ -50,9 +49,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, @@ -224,6 +223,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), @@ -307,12 +310,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(()) } @@ -336,14 +342,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(()) } } @@ -402,50 +409,22 @@ 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, + // Dropping the environment stops unfinished background startup work. + startup_task: Arc>>>, 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, + startup_task: Arc::new(Mutex::new(None)), exec_backend: Arc::new(LocalProcess::default()), filesystem: Arc::new(LocalFileSystem::unsandboxed()), http_client: Arc::new(ReqwestHttpClient), @@ -501,8 +480,8 @@ 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, + startup_task: Arc::new(Mutex::new(None)), exec_backend: Arc::new(LocalProcess::default()), filesystem: Arc::new(LocalFileSystem::with_runtime_paths( local_runtime_paths.clone(), @@ -534,15 +513,15 @@ 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()), + startup_task: Arc::new(Mutex::new(None)), exec_backend, filesystem, http_client: Arc::new(client), @@ -551,7 +530,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 +544,40 @@ 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. + /// Requires an active Tokio runtime when background startup is supported. + pub fn start_connecting(&self) { + let Some(client) = &self.remote_client else { + return; + }; + let mut startup_task = self + .startup_task + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + if startup_task.is_none() { + *startup_task = 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 { @@ -600,7 +612,9 @@ impl From for ShellInfo { #[cfg(test)] mod tests { + use std::collections::HashMap; use std::sync::Arc; + use std::time::Duration; use super::Environment; use super::EnvironmentManager; @@ -609,10 +623,14 @@ mod tests { use super::noise_environment_config_from_values; use crate::ExecServerRuntimePaths; use crate::ProcessId; + use crate::client_api::ExecServerTransportParams; + use crate::client_api::StdioExecServerCommand; use crate::environment_provider::EnvironmentDefault; 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( @@ -626,8 +644,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()), @@ -982,6 +1000,108 @@ 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_leaves_stdio_environment_lazy() { + let environment = Environment::remote_with_transport( + 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), + }, + /*local_runtime_paths*/ None, + ); + let manager = EnvironmentManager::from_snapshot( + EnvironmentProviderSnapshot { + environments: vec![("stdio".to_string(), environment)], + default: EnvironmentDefault::Disabled, + include_local: false, + }, + /*local_runtime_paths*/ None, + ) + .expect("environment manager"); + let environment = manager.get_environment("stdio").expect("stdio environment"); + + assert!(!environment.startup_finished()); + assert!(environment.wait_until_ready().await.is_err()); + assert!(environment.startup_finished()); + } + + #[tokio::test] + async fn replacing_environment_stops_its_startup_task() { + let first_listener = TcpListener::bind("127.0.0.1:0") + .await + .expect("bind first websocket listener"); + let second_listener = TcpListener::bind("127.0.0.1:0") + .await + .expect("bind second websocket listener"); + let manager = EnvironmentManager::without_environments(); + manager + .upsert_environment( + "executor-a".to_string(), + format!( + "ws://{}", + first_listener.local_addr().expect("first listener address") + ), + ) + .expect("first remote environment"); + let environment = manager + .get_environment("executor-a") + .expect("first remote environment"); + let startup_abort = environment + .startup_task + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .as_ref() + .expect("startup task") + .abort_handle(); + assert!(!startup_abort.is_finished()); + drop(environment); + + manager + .upsert_environment( + "executor-a".to_string(), + format!( + "ws://{}", + second_listener + .local_addr() + .expect("second listener address") + ), + ) + .expect("replacement remote environment"); + + timeout(Duration::from_secs(1), async { + while !startup_abort.is_finished() { + tokio::task::yield_now().await; + } + }) + .await + .expect("replacing the environment should cancel its startup task"); + } + #[tokio::test] async fn environment_manager_rejects_empty_remote_environment_url() { let manager = EnvironmentManager::without_environments(); diff --git a/codex-rs/exec-server/tests/relay.rs b/codex-rs/exec-server/tests/relay.rs index 114fe122a..5e5861841 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; @@ -15,25 +13,20 @@ use anyhow::Result; use base64::Engine as _; use base64::engine::general_purpose::STANDARD; use codex_api::AuthProvider; -use codex_exec_server::EnvironmentManager; use codex_exec_server::ExecParams; use codex_exec_server::ExecResponse; use codex_exec_server::ExecServerClient; -use codex_exec_server::ExecServerError; use codex_exec_server::ExecServerRuntimePaths; use codex_exec_server::FsReadFileParams; use codex_exec_server::NoiseChannelIdentity; use codex_exec_server::NoiseChannelPublicKey; use codex_exec_server::NoiseRendezvousConnectArgs; use codex_exec_server::NoiseRendezvousConnectBundle; -use codex_exec_server::NoiseRendezvousConnectProvider; use codex_exec_server::ProcessId; use codex_exec_server::RemoteEnvironmentConfig; use codex_utils_path_uri::PathUri; -use futures::FutureExt; use futures::SinkExt; use futures::StreamExt; -use futures::future::BoxFuture; use http::HeaderMap; use http::HeaderValue; use pretty_assertions::assert_eq; @@ -72,68 +65,10 @@ impl AuthProvider for StaticRegistryAuthProvider { } } -struct FailingNoiseConnectProvider { - attempts: Arc, -} - -impl NoiseRendezvousConnectProvider for FailingNoiseConnectProvider { - fn connect_bundle( - &self, - _: NoiseChannelPublicKey, - ) -> BoxFuture<'_, Result> { - self.attempts.fetch_add(1, Ordering::SeqCst); - async { - Err(ExecServerError::Protocol( - "test registry connect failure".to_string(), - )) - } - .boxed() - } -} - fn static_registry_auth_provider() -> codex_api::SharedAuthProvider { Arc::new(StaticRegistryAuthProvider) } -#[tokio::test] -async fn noise_environment_refreshes_bundle_for_each_connection_attempt() -> Result<()> { - let attempts = Arc::new(AtomicUsize::new(0)); - let manager = EnvironmentManager::without_environments(); - manager.upsert_noise_environment( - ENVIRONMENT_ID.to_string(), - Arc::new(FailingNoiseConnectProvider { - attempts: Arc::clone(&attempts), - }), - )?; - let backend = manager - .get_environment(ENVIRONMENT_ID) - .context("Noise environment should be materialized")? - .get_exec_backend(); - - for attempt in 1..=2 { - let result = backend - .start(ExecParams { - process_id: ProcessId::new(format!("proc-{attempt}")), - argv: vec!["true".to_string()], - cwd: PathUri::from_path(std::env::current_dir()?)?, - 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" - )); - } - - assert_eq!(attempts.load(Ordering::SeqCst), 2); - Ok(()) -} - #[tokio::test(flavor = "multi_thread", worker_threads = 4)] async fn remote_environment_routes_encrypted_exec_server_rpc() -> Result<()> { let listener = TcpListener::bind("127.0.0.1:0").await?;