From 358eb9db237b27dcbf39f25325253e3dc349ea89 Mon Sep 17 00:00:00 2001 From: Sayan Sisodiya Date: Thu, 9 Jul 2026 19:43:00 -0700 Subject: [PATCH] exec-server: support pending environment readiness --- codex-rs/core/tests/suite/remote_env.rs | 81 +++++---- codex-rs/exec-server/src/client.rs | 9 +- codex-rs/exec-server/src/client_api.rs | 5 +- codex-rs/exec-server/src/client_transport.rs | 1 + codex-rs/exec-server/src/environment.rs | 160 ++++++++++++++---- codex-rs/exec-server/src/environment_toml.rs | 5 +- codex-rs/exec-server/src/exec_server_url.rs | 105 ++++++++++++ .../exec-server/src/exec_server_url_tests.rs | 65 +++++++ codex-rs/exec-server/src/lib.rs | 1 + .../src/remote_file_system_path_uri_tests.rs | 3 +- .../exec-server/tests/pending_environment.rs | 102 +++++++++++ 11 files changed, 465 insertions(+), 72 deletions(-) create mode 100644 codex-rs/exec-server/src/exec_server_url.rs create mode 100644 codex-rs/exec-server/src/exec_server_url_tests.rs create mode 100644 codex-rs/exec-server/tests/pending_environment.rs diff --git a/codex-rs/core/tests/suite/remote_env.rs b/codex-rs/core/tests/suite/remote_env.rs index 5f8c503496..544121d24d 100644 --- a/codex-rs/core/tests/suite/remote_env.rs +++ b/codex-rs/core/tests/suite/remote_env.rs @@ -27,7 +27,9 @@ use codex_protocol::protocol::ReviewDecision; use codex_protocol::protocol::RolloutItem; use codex_protocol::protocol::RolloutLine; use codex_protocol::protocol::SandboxPolicy; +use codex_protocol::protocol::ThreadSettingsOverrides; use codex_protocol::protocol::TurnEnvironmentSelection; +use codex_protocol::protocol::TurnEnvironmentSelections; use codex_protocol::request_permissions::PermissionGrantScope; use codex_protocol::request_permissions::RequestPermissionProfile; use codex_protocol::request_permissions::RequestPermissionsResponse; @@ -530,25 +532,25 @@ async fn deferred_executor_updates_context_and_tools_after_startup() -> Result<( ], ) .await; - let mut builder = test_codex() - .with_exec_server_url(format!("ws://{}", listener.local_addr()?)) - .with_config(|config| { - config.project_doc_max_bytes = 0; - config.use_experimental_unified_exec_tool = true; - config.permissions.approval_policy = Constrained::allow_any(AskForApproval::OnRequest); - config.approvals_reviewer = ApprovalsReviewer::User; - assert!(config.features.enable(Feature::DeferredExecutor).is_ok()); - assert!(config.features.enable(Feature::UnifiedExec).is_ok()); - assert!( - config - .features - .enable(Feature::RequestPermissionsTool) - .is_ok() - ); - }); + let mut builder = test_codex().with_config(|config| { + config.project_doc_max_bytes = 0; + config.use_experimental_unified_exec_tool = true; + config.permissions.approval_policy = Constrained::allow_any(AskForApproval::OnRequest); + config.approvals_reviewer = ApprovalsReviewer::User; + assert!(config.features.enable(Feature::DeferredExecutor).is_ok()); + assert!(config.features.enable(Feature::UnifiedExec).is_ok()); + assert!( + config + .features + .enable(Feature::RequestPermissionsTool) + .is_ok() + ); + }); let test = timeout(Duration::from_secs(5), builder.build(&server)) .await .context("thread startup should not wait for the remote environment")??; + let environment_manager = test.thread_manager.environment_manager(); + environment_manager.register_pending_environment(REMOTE_ENVIRONMENT_ID.to_string())?; test.codex .submit(Op::UserInput { @@ -559,11 +561,24 @@ async fn deferred_executor_updates_context_and_tools_after_startup() -> Result<( final_output_json_schema: None, responsesapi_client_metadata: None, additional_context: Default::default(), - thread_settings: Default::default(), + thread_settings: ThreadSettingsOverrides { + environments: Some(TurnEnvironmentSelections::new( + test.config.cwd.clone(), + vec![TurnEnvironmentSelection { + environment_id: REMOTE_ENVIRONMENT_ID.to_string(), + cwd: PathUri::from_abs_path(&test.config.cwd), + }], + )), + ..Default::default() + }, }) .await?; wait_for_response_request_count(&response_mock, /*expected_count*/ 1).await; assert_eq!(response_mock.requests().len(), 1); + environment_manager.set_environment_exec_server_url( + REMOTE_ENVIRONMENT_ID, + format!("ws://{}", listener.local_addr()?), + )?; serve_environment_info(listener).await; let event = wait_for_event(&test.codex, |event| { matches!( @@ -740,7 +755,6 @@ fn agents_md_occurrences(request: &ResponsesRequest, contents: &str) -> usize { #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn deferred_executor_wait_reports_startup_failure() -> Result<()> { - let listener = TcpListener::bind("127.0.0.1:0").await?; let server = start_mock_server().await; let wait_call_id = "wait-for-failure"; let response_mock = mount_sse_sequence( @@ -766,16 +780,16 @@ async fn deferred_executor_wait_reports_startup_failure() -> Result<()> { ], ) .await; - let mut builder = test_codex() - .with_exec_server_url(format!("ws://{}", listener.local_addr()?)) - .with_config(|config| { - config.use_experimental_unified_exec_tool = true; - assert!(config.features.enable(Feature::DeferredExecutor).is_ok()); - assert!(config.features.enable(Feature::UnifiedExec).is_ok()); - }); + let mut builder = test_codex().with_config(|config| { + config.use_experimental_unified_exec_tool = true; + assert!(config.features.enable(Feature::DeferredExecutor).is_ok()); + assert!(config.features.enable(Feature::UnifiedExec).is_ok()); + }); let test = timeout(Duration::from_secs(5), builder.build(&server)) .await .context("thread startup should not wait for the remote environment")??; + let environment_manager = test.thread_manager.environment_manager(); + environment_manager.register_pending_environment(REMOTE_ENVIRONMENT_ID.to_string())?; test.codex .submit(Op::UserInput { @@ -786,15 +800,22 @@ async fn deferred_executor_wait_reports_startup_failure() -> Result<()> { final_output_json_schema: None, responsesapi_client_metadata: None, additional_context: Default::default(), - thread_settings: Default::default(), + thread_settings: ThreadSettingsOverrides { + environments: Some(TurnEnvironmentSelections::new( + test.config.cwd.clone(), + vec![TurnEnvironmentSelection { + environment_id: REMOTE_ENVIRONMENT_ID.to_string(), + cwd: PathUri::from_abs_path(&test.config.cwd), + }], + )), + ..Default::default() + }, }) .await?; wait_for_response_request_count(&response_mock, /*expected_count*/ 1).await; assert_eq!(response_mock.requests().len(), 1); - let (stream, _) = timeout(Duration::from_secs(5), listener.accept()) - .await - .context("exec-server connection should arrive")??; - drop(stream); + environment_manager + .fail_pending_environment(REMOTE_ENVIRONMENT_ID, "CCA provisioning failed".to_string())?; wait_for_event(&test.codex, |event| { matches!(event, EventMsg::TurnComplete(_)) }) diff --git a/codex-rs/exec-server/src/client.rs b/codex-rs/exec-server/src/client.rs index 73e49b2f54..cf90830483 100644 --- a/codex-rs/exec-server/src/client.rs +++ b/codex-rs/exec-server/src/client.rs @@ -1430,6 +1430,7 @@ mod tests { use crate::client_api::StdioExecServerCommand; use crate::client_api::StdioExecServerConnectArgs; use crate::connection::JsonRpcConnection; + use crate::exec_server_url::ExecServerUrl; use crate::process::ExecProcessEvent; use crate::protocol::EXEC_CLOSED_METHOD; use crate::protocol::EXEC_EXITED_METHOD; @@ -2113,7 +2114,7 @@ mod tests { }); let client = LazyRemoteExecServerClient::new(ExecServerTransportParams::WebSocketUrl { - websocket_url, + websocket_url: ExecServerUrl::ready(websocket_url), connect_timeout: Duration::from_secs(1), initialize_timeout: Duration::from_secs(1), }); @@ -2220,7 +2221,7 @@ mod tests { }); let client = LazyRemoteExecServerClient::new(ExecServerTransportParams::WebSocketUrl { - websocket_url, + websocket_url: ExecServerUrl::ready(websocket_url), connect_timeout: Duration::from_secs(1), initialize_timeout: Duration::from_secs(1), }); @@ -2355,7 +2356,7 @@ mod tests { .expect("client should close after the test"); }); let client = LazyRemoteExecServerClient::new(ExecServerTransportParams::WebSocketUrl { - websocket_url, + websocket_url: ExecServerUrl::ready(websocket_url), connect_timeout: Duration::from_secs(1), initialize_timeout: Duration::from_secs(1), }); @@ -2464,7 +2465,7 @@ mod tests { .expect("client should close after the test"); }); let client = LazyRemoteExecServerClient::new(ExecServerTransportParams::WebSocketUrl { - websocket_url, + websocket_url: ExecServerUrl::ready(websocket_url), connect_timeout: Duration::from_secs(1), initialize_timeout: Duration::from_secs(1), }); diff --git a/codex-rs/exec-server/src/client_api.rs b/codex-rs/exec-server/src/client_api.rs index 0c6a3576d0..f7d0f87e0d 100644 --- a/codex-rs/exec-server/src/client_api.rs +++ b/codex-rs/exec-server/src/client_api.rs @@ -11,6 +11,7 @@ use crate::HttpRequestResponse; use crate::HttpResponseBodyStream; use crate::NoiseChannelIdentity; use crate::NoiseChannelPublicKey; +use crate::exec_server_url::ExecServerUrl; pub(crate) const DEFAULT_REMOTE_EXEC_SERVER_CONNECT_TIMEOUT: Duration = Duration::from_secs(10); pub(crate) const DEFAULT_REMOTE_EXEC_SERVER_INITIALIZE_TIMEOUT: Duration = Duration::from_secs(10); @@ -92,7 +93,7 @@ pub(crate) struct StdioExecServerCommand { #[derive(Clone)] pub(crate) enum ExecServerTransportParams { WebSocketUrl { - websocket_url: String, + websocket_url: ExecServerUrl, connect_timeout: Duration, initialize_timeout: Duration, }, @@ -136,7 +137,7 @@ impl std::fmt::Debug for ExecServerTransportParams { } impl ExecServerTransportParams { - pub(crate) fn websocket_url(websocket_url: String, connect_timeout: Duration) -> Self { + pub(crate) fn websocket_url(websocket_url: ExecServerUrl, connect_timeout: Duration) -> Self { Self::WebSocketUrl { websocket_url, connect_timeout, diff --git a/codex-rs/exec-server/src/client_transport.rs b/codex-rs/exec-server/src/client_transport.rs index 1ccfd69152..86f91cec38 100644 --- a/codex-rs/exec-server/src/client_transport.rs +++ b/codex-rs/exec-server/src/client_transport.rs @@ -99,6 +99,7 @@ impl ExecServerClient { connect_timeout, initialize_timeout, } => { + let websocket_url = websocket_url.resolve().await?; Self::connect_websocket(RemoteExecServerConnectArgs { websocket_url, client_name: ENVIRONMENT_CLIENT_NAME.to_string(), diff --git a/codex-rs/exec-server/src/environment.rs b/codex-rs/exec-server/src/environment.rs index d070be6786..3cfd3a2b37 100644 --- a/codex-rs/exec-server/src/environment.rs +++ b/codex-rs/exec-server/src/environment.rs @@ -19,6 +19,7 @@ use crate::environment_provider::EnvironmentProvider; use crate::environment_provider::EnvironmentProviderSnapshot; use crate::environment_provider::normalize_exec_server_url; use crate::environment_toml::environment_provider_from_codex_home; +use crate::exec_server_url::ExecServerUrl; use crate::local_file_system::LocalFileSystem; use crate::local_process::LocalProcess; use crate::process::ExecBackend; @@ -295,34 +296,80 @@ impl EnvironmentManager { exec_server_url: String, connect_timeout: Option, ) -> Result<(), ExecServerError> { - if environment_id.is_empty() { - return Err(ExecServerError::Protocol( - "environment id cannot be empty".to_string(), - )); - } - let (exec_server_url, disabled) = normalize_exec_server_url(Some(exec_server_url)); - if disabled { - return Err(ExecServerError::Protocol( - "remote environment cannot use disabled exec-server url".to_string(), - )); - } - let Some(exec_server_url) = exec_server_url else { - return Err(ExecServerError::Protocol( - "remote environment requires an exec-server url".to_string(), - )); - }; + validate_environment_id(&environment_id)?; + let exec_server_url = validate_remote_exec_server_url(exec_server_url)?; let environment = Arc::new(Environment::remote_with_transport( ExecServerTransportParams::websocket_url( - exec_server_url, + ExecServerUrl::ready(exec_server_url), connect_timeout.unwrap_or(DEFAULT_REMOTE_EXEC_SERVER_CONNECT_TIMEOUT), ), self.local_runtime_paths.clone(), )); environment.start_connecting(); - self.environments + self.replace_environment(environment_id, environment); + Ok(()) + } + + /// Registers a remote environment before its stable exec-server URL is available. + /// Requires an active Tokio runtime to start background connection work. + pub fn register_pending_environment( + &self, + environment_id: String, + ) -> Result<(), ExecServerError> { + validate_environment_id(&environment_id)?; + let environment = Arc::new(Environment::remote_with_transport( + ExecServerTransportParams::websocket_url( + ExecServerUrl::pending(), + DEFAULT_REMOTE_EXEC_SERVER_CONNECT_TIMEOUT, + ), + self.local_runtime_paths.clone(), + )); + let mut environments = self + .environments .write() - .unwrap_or_else(std::sync::PoisonError::into_inner) - .insert(environment_id, environment); + .unwrap_or_else(std::sync::PoisonError::into_inner); + if environments.contains_key(&environment_id) { + return Err(ExecServerError::Protocol(format!( + "environment id `{environment_id}` is already registered" + ))); + } + environment.start_connecting(); + environments.insert(environment_id, environment); + Ok(()) + } + + /// Supplies the stable exec-server URL for a pending environment. + pub fn set_environment_exec_server_url( + &self, + environment_id: &str, + exec_server_url: String, + ) -> Result<(), ExecServerError> { + let exec_server_url = validate_remote_exec_server_url(exec_server_url)?; + let environments = self + .environments + .read() + .unwrap_or_else(std::sync::PoisonError::into_inner); + let environment = environments.get(environment_id).ok_or_else(|| { + ExecServerError::Protocol(format!("unknown environment id `{environment_id}`")) + })?; + environment.set_exec_server_url(exec_server_url) + } + + /// Fails and removes a pending environment so the same id may be registered again. + pub fn fail_pending_environment( + &self, + environment_id: &str, + message: String, + ) -> Result<(), ExecServerError> { + let mut environments = self + .environments + .write() + .unwrap_or_else(std::sync::PoisonError::into_inner); + let environment = environments.get(environment_id).ok_or_else(|| { + ExecServerError::Protocol(format!("unknown environment id `{environment_id}`")) + })?; + environment.fail_pending_exec_server_url(message)?; + environments.remove(environment_id); Ok(()) } @@ -336,11 +383,7 @@ impl EnvironmentManager { environment_id: String, provider: Arc, ) -> Result<(), ExecServerError> { - if environment_id.is_empty() { - return Err(ExecServerError::Protocol( - "environment id cannot be empty".to_string(), - )); - } + validate_environment_id(&environment_id)?; let identity = NoiseChannelIdentity::generate().map_err(|error| { ExecServerError::Protocol(format!( "failed to generate Noise harness identity: {error}" @@ -351,12 +394,43 @@ impl EnvironmentManager { self.local_runtime_paths.clone(), )); environment.start_connecting(); - self.environments - .write() - .unwrap_or_else(std::sync::PoisonError::into_inner) - .insert(environment_id, environment); + self.replace_environment(environment_id, environment); Ok(()) } + + fn replace_environment(&self, environment_id: String, environment: Arc) { + let previous = self + .environments + .write() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .insert(environment_id.clone(), environment); + if let Some(previous) = previous { + let _ = previous.fail_pending_exec_server_url(format!( + "environment `{environment_id}` was replaced" + )); + } + } +} + +fn validate_environment_id(environment_id: &str) -> Result<(), ExecServerError> { + if environment_id.is_empty() { + return Err(ExecServerError::Protocol( + "environment id cannot be empty".to_string(), + )); + } + Ok(()) +} + +fn validate_remote_exec_server_url(exec_server_url: String) -> Result { + let (exec_server_url, disabled) = normalize_exec_server_url(Some(exec_server_url)); + if disabled { + return Err(ExecServerError::Protocol( + "remote environment cannot use disabled exec-server url".to_string(), + )); + } + exec_server_url.ok_or_else(|| { + ExecServerError::Protocol("remote environment requires an exec-server url".to_string()) + }) } fn noise_environment_config_from_env() @@ -412,7 +486,7 @@ fn optional_environment_value(name: &str) -> Option { /// paths used by filesystem helpers. #[derive(Clone)] pub struct Environment { - exec_server_url: Option, + exec_server_url: Option, remote_client: Option, // Dropping the environment stops unfinished background startup work. startup_task: Arc>>>, @@ -503,7 +577,7 @@ impl Environment { ) -> Self { Self::remote_with_transport( ExecServerTransportParams::websocket_url( - exec_server_url, + ExecServerUrl::ready(exec_server_url), DEFAULT_REMOTE_EXEC_SERVER_CONNECT_TIMEOUT, ), local_runtime_paths, @@ -542,9 +616,29 @@ impl Environment { self.remote_client.is_some() } - /// Returns the remote exec-server URL when this environment is remote. + /// Returns the stable exec-server URL once available for a URL-backed remote environment. pub fn exec_server_url(&self) -> Option<&str> { - self.exec_server_url.as_deref() + self.exec_server_url + .as_ref() + .and_then(ExecServerUrl::current) + } + + fn set_exec_server_url(&self, exec_server_url: String) -> Result<(), ExecServerError> { + self.exec_server_url + .as_ref() + .ok_or_else(|| { + ExecServerError::Protocol("environment does not use an exec-server URL".to_string()) + })? + .set_ready(exec_server_url) + } + + fn fail_pending_exec_server_url(&self, message: String) -> Result<(), ExecServerError> { + self.exec_server_url + .as_ref() + .ok_or_else(|| { + ExecServerError::Protocol("environment does not use an exec-server URL".to_string()) + })? + .set_failed(message) } pub fn local_runtime_paths(&self) -> Option<&ExecServerRuntimePaths> { diff --git a/codex-rs/exec-server/src/environment_toml.rs b/codex-rs/exec-server/src/environment_toml.rs index bfc25873ba..cfd2a9b81e 100644 --- a/codex-rs/exec-server/src/environment_toml.rs +++ b/codex-rs/exec-server/src/environment_toml.rs @@ -19,6 +19,7 @@ use crate::client_api::StdioExecServerCommand; use crate::environment::LOCAL_ENVIRONMENT_ID; use crate::environment_provider::EnvironmentDefault; use crate::environment_provider::EnvironmentProviderSnapshot; +use crate::exec_server_url::ExecServerUrl; const ENVIRONMENTS_TOML_FILE: &str = "environments.toml"; const MAX_ENVIRONMENT_ID_LEN: usize = 64; @@ -153,7 +154,7 @@ fn parse_environment_toml( (Some(url), None) => { let url = validate_websocket_url(url)?; ExecServerTransportParams::WebSocketUrl { - websocket_url: url, + websocket_url: ExecServerUrl::ready(url), connect_timeout, initialize_timeout, } @@ -633,7 +634,7 @@ mod tests { else { panic!("expected websocket transport"); }; - assert_eq!(websocket_url, "ws://127.0.0.1:8765"); + assert_eq!(websocket_url.current(), Some("ws://127.0.0.1:8765")); assert_eq!(*connect_timeout, Duration::from_secs(12)); assert_eq!(*initialize_timeout, Duration::from_secs(34)); diff --git a/codex-rs/exec-server/src/exec_server_url.rs b/codex-rs/exec-server/src/exec_server_url.rs new file mode 100644 index 0000000000..dfc6bb24b3 --- /dev/null +++ b/codex-rs/exec-server/src/exec_server_url.rs @@ -0,0 +1,105 @@ +use std::sync::Arc; +use std::sync::OnceLock; + +use tokio::sync::watch; + +use crate::ExecServerError; + +type ExecServerUrlResult = Result; + +struct ExecServerUrlInner { + result: OnceLock, + result_available: watch::Sender, +} + +/// The stable WebSocket URL for one exec-server environment. +/// +/// A URL may be available when the environment is registered or supplied later by the owner that +/// provisions the environment. Every consumer observes the same one-shot ready or failed result. +#[derive(Clone)] +pub(crate) struct ExecServerUrl { + inner: Arc, +} + +impl ExecServerUrl { + pub(crate) fn ready(url: String) -> Self { + let (result_available, _receiver) = watch::channel(true); + Self { + inner: Arc::new(ExecServerUrlInner { + result: OnceLock::from(Ok(url)), + result_available, + }), + } + } + + pub(crate) fn pending() -> Self { + let (result_available, _receiver) = watch::channel(false); + Self { + inner: Arc::new(ExecServerUrlInner { + result: OnceLock::new(), + result_available, + }), + } + } + + pub(crate) async fn resolve(&self) -> Result { + let mut result_available = self.inner.result_available.subscribe(); + while !*result_available.borrow_and_update() { + if result_available.changed().await.is_err() { + return Err(ExecServerError::Disconnected( + "environment URL registration ended before completion".to_string(), + )); + } + } + self.inner + .result + .get() + .ok_or_else(|| { + ExecServerError::Protocol( + "exec-server URL was marked available without a result".to_string(), + ) + })? + .clone() + .map_err(|message| { + ExecServerError::Disconnected(format!("environment unavailable: {message}")) + }) + } + + pub(crate) fn current(&self) -> Option<&str> { + match self.inner.result.get()? { + Ok(url) => Some(url.as_str()), + Err(_) => None, + } + } + + pub(crate) fn set_ready(&self, url: String) -> Result<(), ExecServerError> { + self.complete(Ok(url)) + } + + pub(crate) fn set_failed(&self, message: String) -> Result<(), ExecServerError> { + self.complete(Err(message)) + } + + fn complete(&self, result: ExecServerUrlResult) -> Result<(), ExecServerError> { + self.inner + .result + .set(result) + .map_err(|_| ExecServerError::Protocol("exec-server URL is not pending".to_string()))?; + self.inner.result_available.send_replace(true); + Ok(()) + } +} + +impl std::fmt::Debug for ExecServerUrl { + fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self.inner.result.get() { + Some(Ok(url)) => std::fmt::Debug::fmt(url, formatter), + Some(Err(_)) => formatter.write_str(""), + None => formatter.write_str(""), + } + } +} + +#[cfg(test)] +#[path = "exec_server_url_tests.rs"] +mod tests; diff --git a/codex-rs/exec-server/src/exec_server_url_tests.rs b/codex-rs/exec-server/src/exec_server_url_tests.rs new file mode 100644 index 0000000000..882f488794 --- /dev/null +++ b/codex-rs/exec-server/src/exec_server_url_tests.rs @@ -0,0 +1,65 @@ +use pretty_assertions::assert_eq; + +use super::ExecServerUrl; + +#[tokio::test] +async fn ready_url_resolves_immediately() { + let url = ExecServerUrl::ready("ws://127.0.0.1:1234".to_string()); + + assert_eq!( + url.resolve().await.expect("ready URL"), + "ws://127.0.0.1:1234" + ); + assert_eq!(url.current(), Some("ws://127.0.0.1:1234")); +} + +#[tokio::test] +async fn pending_url_resolves_once_and_caches_the_result() { + let url = ExecServerUrl::pending(); + let waiter = tokio::spawn({ + let url = url.clone(); + async move { url.resolve().await } + }); + tokio::task::yield_now().await; + assert!(!waiter.is_finished()); + + url.set_ready("ws://127.0.0.1:5678".to_string()) + .expect("complete URL"); + + assert_eq!( + waiter.await.expect("waiter task").expect("resolved URL"), + "ws://127.0.0.1:5678" + ); + assert_eq!( + url.resolve().await.expect("cached URL"), + "ws://127.0.0.1:5678" + ); + assert_eq!(url.current(), Some("ws://127.0.0.1:5678")); + assert_eq!( + url.set_ready("ws://127.0.0.1:9999".to_string()) + .expect_err("URL completion is one-shot") + .to_string(), + "exec-server protocol error: exec-server URL is not pending" + ); +} + +#[tokio::test] +async fn pending_url_failure_is_shared() { + let url = ExecServerUrl::pending(); + + url.set_failed("provisioning failed".to_string()) + .expect("fail URL"); + + assert_eq!( + url.resolve().await.expect_err("failed URL").to_string(), + "environment unavailable: provisioning failed" + ); + assert_eq!( + url.resolve() + .await + .expect_err("cached failed URL") + .to_string(), + "environment unavailable: provisioning failed" + ); + assert_eq!(url.current(), None); +} diff --git a/codex-rs/exec-server/src/lib.rs b/codex-rs/exec-server/src/lib.rs index c2c8cc5194..98445da0af 100644 --- a/codex-rs/exec-server/src/lib.rs +++ b/codex-rs/exec-server/src/lib.rs @@ -6,6 +6,7 @@ mod environment; mod environment_provider; mod environment_registry; mod environment_toml; +mod exec_server_url; mod file_read; mod fs_helper; mod fs_helper_main; diff --git a/codex-rs/exec-server/src/remote_file_system_path_uri_tests.rs b/codex-rs/exec-server/src/remote_file_system_path_uri_tests.rs index 562f5d5643..e071f7df58 100644 --- a/codex-rs/exec-server/src/remote_file_system_path_uri_tests.rs +++ b/codex-rs/exec-server/src/remote_file_system_path_uri_tests.rs @@ -25,6 +25,7 @@ use tokio_tungstenite::tungstenite::Message; use super::*; use crate::client_api::DEFAULT_REMOTE_EXEC_SERVER_CONNECT_TIMEOUT; use crate::client_api::ExecServerTransportParams; +use crate::exec_server_url::ExecServerUrl; use crate::protocol::FS_READ_FILE_METHOD; use crate::protocol::FsReadFileParams; use crate::protocol::FsReadFileResponse; @@ -38,7 +39,7 @@ async fn remote_file_system_sends_path_and_sandbox_cwd_uris_without_native_conve record_read_file_params(/*expected_requests*/ 2).await; let file_system = RemoteFileSystem::new(LazyRemoteExecServerClient::new( ExecServerTransportParams::websocket_url( - websocket_url, + ExecServerUrl::ready(websocket_url), DEFAULT_REMOTE_EXEC_SERVER_CONNECT_TIMEOUT, ), )); diff --git a/codex-rs/exec-server/tests/pending_environment.rs b/codex-rs/exec-server/tests/pending_environment.rs new file mode 100644 index 0000000000..85dae6308c --- /dev/null +++ b/codex-rs/exec-server/tests/pending_environment.rs @@ -0,0 +1,102 @@ +mod common; + +use std::sync::Arc; + +use codex_exec_server::EnvironmentManager; +use common::exec_server::exec_server; +use pretty_assertions::assert_eq; + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn pending_environment_connects_after_url_is_supplied() -> anyhow::Result<()> { + let mut server = exec_server().await?; + let manager = EnvironmentManager::without_environments(); + manager.register_pending_environment("tools".to_string())?; + let environment = manager + .get_environment("tools") + .expect("pending environment"); + let readiness = tokio::spawn({ + let environment = Arc::clone(&environment); + async move { environment.wait_until_ready().await } + }); + tokio::task::yield_now().await; + + assert!(!readiness.is_finished()); + assert_eq!(environment.exec_server_url(), None); + + manager.set_environment_exec_server_url("tools", server.websocket_url().to_string())?; + readiness.await??; + + assert_eq!(environment.exec_server_url(), Some(server.websocket_url())); + environment.info().await?; + server.shutdown().await?; + Ok(()) +} + +#[tokio::test] +async fn failing_pending_environment_wakes_waiters_and_allows_reregistration() -> anyhow::Result<()> +{ + let manager = EnvironmentManager::without_environments(); + manager.register_pending_environment("tools".to_string())?; + assert_eq!( + manager + .register_pending_environment("tools".to_string()) + .expect_err("duplicate pending registration should fail") + .to_string(), + "exec-server protocol error: environment id `tools` is already registered" + ); + let environment = manager + .get_environment("tools") + .expect("pending environment"); + let readiness = tokio::spawn(async move { environment.wait_until_ready().await }); + tokio::task::yield_now().await; + + manager.fail_pending_environment("tools", "CCA provisioning failed".to_string())?; + + assert_eq!( + readiness + .await? + .expect_err("pending environment should fail") + .to_string(), + "exec-server connection attempt failed: environment unavailable: CCA provisioning failed" + ); + assert!(manager.get_environment("tools").is_none()); + manager.register_pending_environment("tools".to_string())?; + manager.fail_pending_environment("tools", "test cleanup".to_string())?; + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn replacing_pending_environment_fails_its_existing_waiters() -> anyhow::Result<()> { + let mut server = exec_server().await?; + let manager = EnvironmentManager::without_environments(); + manager.register_pending_environment("tools".to_string())?; + let pending_environment = manager + .get_environment("tools") + .expect("pending environment"); + let readiness = tokio::spawn({ + let pending_environment = Arc::clone(&pending_environment); + async move { pending_environment.wait_until_ready().await } + }); + tokio::task::yield_now().await; + + manager.upsert_environment( + "tools".to_string(), + server.websocket_url().to_string(), + /*connect_timeout*/ None, + )?; + let replacement_environment = manager + .get_environment("tools") + .expect("replacement environment"); + replacement_environment.wait_until_ready().await?; + + assert_eq!( + readiness + .await? + .expect_err("replaced pending environment should fail") + .to_string(), + "exec-server connection attempt failed: environment unavailable: environment `tools` was replaced" + ); + assert!(!Arc::ptr_eq(&pending_environment, &replacement_environment)); + server.shutdown().await?; + Ok(()) +}