From 4472698728b1a8f4bd184086c5d1d0c65f2918b3 Mon Sep 17 00:00:00 2001 From: pakrym-oai Date: Fri, 10 Jul 2026 18:03:04 +0000 Subject: [PATCH] Support pending remote environment registration (#32231) ## Why Remote environment provisioning can finish after a thread starts, before an exec-server WebSocket URL is available. ## What changed - Add `EnvironmentManager::register_pending_environment` and a one-shot `PendingEnvironmentRegistration` handle that resolves to either a validated WebSocket URL or a terminal provisioning error. - Let lazy remote exec-server clients wait for that result, while preserving reconnection behavior after a successful registration. - Keep replacement registrations isolated so completing an older handle does not resolve the current environment with the same ID. ## Testing Add coverage for successful connection and reconnection, provisioning and dropped-registration failures, invalid URLs, replacement isolation, and the deferred-executor startup flow. GitOrigin-RevId: 5c05be2b72291b77a1f71176d7075b1ad63332a5 --- .../tests/common/test_app_server.rs | 2 +- codex-rs/core/tests/common/test_codex.rs | 4 + codex-rs/core/tests/suite/remote_env.rs | 79 +++++++----- codex-rs/exec-server/src/client.rs | 1 + codex-rs/exec-server/src/client_api.rs | 8 ++ codex-rs/exec-server/src/client_transport.rs | 53 +++++--- codex-rs/exec-server/src/environment.rs | 119 +++++++++++------- .../exec-server/src/environment_provider.rs | 16 +-- codex-rs/exec-server/src/environment_toml.rs | 6 +- codex-rs/exec-server/src/lib.rs | 1 + .../exec-server/tests/pending_environment.rs | 70 +++++++++++ 11 files changed, 244 insertions(+), 115 deletions(-) create mode 100644 codex-rs/exec-server/tests/pending_environment.rs diff --git a/codex-rs/app-server/tests/common/test_app_server.rs b/codex-rs/app-server/tests/common/test_app_server.rs index 825ea46ab8..3852399999 100644 --- a/codex-rs/app-server/tests/common/test_app_server.rs +++ b/codex-rs/app-server/tests/common/test_app_server.rs @@ -1843,7 +1843,7 @@ impl TestAppServerBuilder { let mut auto_env_overrides = vec![ ( CODEX_EXEC_SERVER_URL_ENV_VAR.to_string(), - auto_env.environment().exec_server_url().map(str::to_string), + auto_env.exec_server_url().map(str::to_string), ), ( CODEX_EXEC_SERVER_NOISE_REGISTRY_URL_ENV_VAR.to_string(), diff --git a/codex-rs/core/tests/common/test_codex.rs b/codex-rs/core/tests/common/test_codex.rs index b95a003c5b..dd0b5204fa 100644 --- a/codex-rs/core/tests/common/test_codex.rs +++ b/codex-rs/core/tests/common/test_codex.rs @@ -163,6 +163,10 @@ impl TestEnv { &self.environment } + pub fn exec_server_url(&self) -> Option<&str> { + self.exec_server_url.as_deref() + } + /// Returns the environment and target-native cwd selected by the test harness. pub fn selection(&self) -> &TurnEnvironmentSelection { &self.selection diff --git a/codex-rs/core/tests/suite/remote_env.rs b/codex-rs/core/tests/suite/remote_env.rs index 5f8c503496..31c8853a35 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,26 @@ 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(); + let registration = + environment_manager.register_pending_environment(REMOTE_ENVIRONMENT_ID.to_string())?; test.codex .submit(Op::UserInput { @@ -559,11 +562,21 @@ 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); + registration.complete(Ok(format!("ws://{}", listener.local_addr()?)))?; serve_environment_info(listener).await; let event = wait_for_event(&test.codex, |event| { matches!( @@ -740,7 +753,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 +778,17 @@ 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(); + let registration = + environment_manager.register_pending_environment(REMOTE_ENVIRONMENT_ID.to_string())?; test.codex .submit(Op::UserInput { @@ -786,15 +799,21 @@ 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); + registration.complete(Err("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 ae601c5dda..b34ec7ee3d 100644 --- a/codex-rs/exec-server/src/client.rs +++ b/codex-rs/exec-server/src/client.rs @@ -427,6 +427,7 @@ impl LazyRemoteExecServerClient { matches!( self.transport_params, ExecServerTransportParams::WebSocketUrl { .. } + | ExecServerTransportParams::PendingWebSocketUrl(..) | ExecServerTransportParams::NoiseRendezvous { .. } ) } diff --git a/codex-rs/exec-server/src/client_api.rs b/codex-rs/exec-server/src/client_api.rs index 0c6a3576d0..5e72508df6 100644 --- a/codex-rs/exec-server/src/client_api.rs +++ b/codex-rs/exec-server/src/client_api.rs @@ -4,6 +4,8 @@ use std::sync::Arc; use std::time::Duration; use futures::future::BoxFuture; +use futures::future::Shared; +use tokio::sync::oneshot; use crate::ExecServerError; use crate::HttpRequestParams; @@ -88,6 +90,8 @@ pub(crate) struct StdioExecServerCommand { pub cwd: Option, } +pub(crate) type PendingExecServerUrl = Shared>>; + /// Parameters used to connect to a remote exec-server environment. #[derive(Clone)] pub(crate) enum ExecServerTransportParams { @@ -96,6 +100,7 @@ pub(crate) enum ExecServerTransportParams { connect_timeout: Duration, initialize_timeout: Duration, }, + PendingWebSocketUrl(PendingExecServerUrl), NoiseRendezvous { provider: Arc, identity: NoiseChannelIdentity, @@ -120,6 +125,9 @@ impl std::fmt::Debug for ExecServerTransportParams { .field("connect_timeout", connect_timeout) .field("initialize_timeout", initialize_timeout) .finish(), + Self::PendingWebSocketUrl(..) => { + f.debug_tuple("PendingWebSocketUrl").finish_non_exhaustive() + } Self::NoiseRendezvous { .. } => { f.debug_struct("NoiseRendezvous").finish_non_exhaustive() } diff --git a/codex-rs/exec-server/src/client_transport.rs b/codex-rs/exec-server/src/client_transport.rs index 1ccfd69152..99cb8a9aa8 100644 --- a/codex-rs/exec-server/src/client_transport.rs +++ b/codex-rs/exec-server/src/client_transport.rs @@ -19,6 +19,7 @@ use crate::ExecServerError; use crate::client_api::DEFAULT_REMOTE_EXEC_SERVER_CONNECT_TIMEOUT; use crate::client_api::DEFAULT_REMOTE_EXEC_SERVER_INITIALIZE_TIMEOUT; use crate::client_api::ExecServerClientConnectOptions; +use crate::client_api::ExecServerTransportParams; use crate::client_api::NoiseRendezvousConnectArgs; use crate::client_api::NoiseRendezvousConnectBundle; use crate::client_api::NoiseRendezvousConnectProvider; @@ -91,27 +92,30 @@ impl ExecServerClient { /// Noise connection details are fetched here so reconnects get a fresh URL /// and authorization without replacing the harness identity. pub(crate) async fn connect_for_transport( - transport_params: crate::client_api::ExecServerTransportParams, + transport_params: ExecServerTransportParams, ) -> Result { - match transport_params { - crate::client_api::ExecServerTransportParams::WebSocketUrl { + let (websocket_url, connect_timeout, initialize_timeout) = match transport_params { + ExecServerTransportParams::WebSocketUrl { websocket_url, connect_timeout, initialize_timeout, - } => { - Self::connect_websocket(RemoteExecServerConnectArgs { + } => (websocket_url, connect_timeout, initialize_timeout), + ExecServerTransportParams::PendingWebSocketUrl(websocket_url) => { + let websocket_url = websocket_url + .await + .unwrap_or_else(|_| { + Err("environment registration ended before completion".to_string()) + }) + .map_err(|message| { + ExecServerError::Disconnected(format!("environment unavailable: {message}")) + })?; + ( websocket_url, - client_name: ENVIRONMENT_CLIENT_NAME.to_string(), - connect_timeout, - initialize_timeout, - resume_session_id: None, - }) - .await + DEFAULT_REMOTE_EXEC_SERVER_CONNECT_TIMEOUT, + DEFAULT_REMOTE_EXEC_SERVER_INITIALIZE_TIMEOUT, + ) } - crate::client_api::ExecServerTransportParams::NoiseRendezvous { - provider, - identity, - } => { + ExecServerTransportParams::NoiseRendezvous { provider, identity } => { let reconnect_strategy = ExecServerReconnectStrategy::NoiseRendezvous { provider: Arc::clone(&provider), identity: identity.clone(), @@ -121,21 +125,30 @@ impl ExecServerClient { }; let (connection, options) = Self::open_initial_noise_rendezvous_connection(&provider, &identity).await?; - Self::connect_with_recovery(connection, options, Some(reconnect_strategy)).await + return Self::connect_with_recovery(connection, options, Some(reconnect_strategy)) + .await; } - crate::client_api::ExecServerTransportParams::StdioCommand { + ExecServerTransportParams::StdioCommand { command, initialize_timeout, } => { - Self::connect_stdio_command(StdioExecServerConnectArgs { + return Self::connect_stdio_command(StdioExecServerConnectArgs { command, client_name: ENVIRONMENT_CLIENT_NAME.to_string(), initialize_timeout, resume_session_id: None, }) - .await + .await; } - } + }; + Self::connect_websocket(RemoteExecServerConnectArgs { + websocket_url, + client_name: ENVIRONMENT_CLIENT_NAME.to_string(), + connect_timeout, + initialize_timeout, + resume_session_id: None, + }) + .await } async fn open_initial_noise_rendezvous_connection( diff --git a/codex-rs/exec-server/src/environment.rs b/codex-rs/exec-server/src/environment.rs index d070be6786..e2499927cd 100644 --- a/codex-rs/exec-server/src/environment.rs +++ b/codex-rs/exec-server/src/environment.rs @@ -3,6 +3,8 @@ use std::sync::Arc; use std::sync::Mutex; use std::sync::RwLock; +use futures::FutureExt; + use crate::ExecServerError; use crate::ExecServerRuntimePaths; use crate::ExecutorFileSystem; @@ -26,6 +28,7 @@ use crate::protocol::EnvironmentInfo; use crate::remote::NoiseRendezvousEnvironmentConfig; use crate::remote_file_system::RemoteFileSystem; use crate::remote_process::RemoteProcess; +use tokio::sync::oneshot; use tokio_util::task::AbortOnDropHandle; pub const CODEX_EXEC_SERVER_URL_ENV_VAR: &str = "CODEX_EXEC_SERVER_URL"; @@ -59,6 +62,10 @@ pub struct EnvironmentManager { local_runtime_paths: Option, } +/// The one-shot capability to complete a pending environment registration. +#[must_use = "the pending environment cannot connect until registration is completed"] +pub struct PendingEnvironmentRegistration(oneshot::Sender>); + pub const LOCAL_ENVIRONMENT_ID: &str = "local"; pub const REMOTE_ENVIRONMENT_ID: &str = "remote"; @@ -295,22 +302,8 @@ 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, @@ -326,6 +319,25 @@ impl EnvironmentManager { Ok(()) } + /// Adds or replaces a remote environment whose stable URL will be supplied later. + pub fn register_pending_environment( + &self, + environment_id: String, + ) -> Result { + validate_environment_id(&environment_id)?; + let (completion, websocket_url) = oneshot::channel(); + let environment = Arc::new(Environment::remote_with_transport( + ExecServerTransportParams::PendingWebSocketUrl(websocket_url.shared()), + self.local_runtime_paths.clone(), + )); + environment.start_connecting(); + self.environments + .write() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .insert(environment_id, environment); + Ok(PendingEnvironmentRegistration(completion)) + } + /// Adds or replaces a named remote environment that connects through an /// authenticated, end-to-end encrypted rendezvous stream. /// @@ -336,11 +348,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}" @@ -359,6 +367,46 @@ impl EnvironmentManager { } } +impl PendingEnvironmentRegistration { + /// Completes provisioning with the stable URL or a terminal error message. + pub fn complete(self, result: Result) -> Result<(), ExecServerError> { + let result = match result { + Ok(exec_server_url) => match validate_remote_exec_server_url(exec_server_url) { + Ok(exec_server_url) => Ok(exec_server_url), + Err(error) => { + let _ = self.0.send(Err(error.to_string())); + return Err(error); + } + }, + Err(message) => Err(message), + }; + self.0.send(result).map_err(|_| { + ExecServerError::Disconnected("pending environment registration is inactive".into()) + }) + } +} + +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() -> Result, ExecServerError> { noise_environment_config_from_values( @@ -412,7 +460,6 @@ fn optional_environment_value(name: &str) -> Option { /// paths used by filesystem helpers. #[derive(Clone)] pub struct Environment { - exec_server_url: Option, remote_client: Option, // Dropping the environment stops unfinished background startup work. startup_task: Arc>>>, @@ -426,7 +473,6 @@ 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_client: None, startup_task: Arc::new(Mutex::new(None)), exec_backend: Arc::new(LocalProcess::default()), @@ -439,9 +485,7 @@ impl Environment { impl std::fmt::Debug for Environment { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - f.debug_struct("Environment") - .field("exec_server_url", &self.exec_server_url) - .finish_non_exhaustive() + f.debug_struct("Environment").finish_non_exhaustive() } } @@ -483,7 +527,6 @@ impl Environment { pub(crate) fn local(local_runtime_paths: ExecServerRuntimePaths) -> Self { Self { - exec_server_url: None, remote_client: None, startup_task: Arc::new(Mutex::new(None)), exec_backend: Arc::new(LocalProcess::with_local_runtime_paths( @@ -514,21 +557,12 @@ impl Environment { remote_transport: ExecServerTransportParams, local_runtime_paths: Option, ) -> Self { - let exec_server_url = match &remote_transport { - ExecServerTransportParams::WebSocketUrl { - websocket_url: exec_server_url, - .. - } => Some(exec_server_url.clone()), - ExecServerTransportParams::NoiseRendezvous { .. } => None, - ExecServerTransportParams::StdioCommand { .. } => None, - }; 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_client: Some(client.clone()), startup_task: Arc::new(Mutex::new(None)), exec_backend, @@ -542,11 +576,6 @@ impl Environment { self.remote_client.is_some() } - /// Returns the remote exec-server URL when this environment is remote. - pub fn exec_server_url(&self) -> Option<&str> { - self.exec_server_url.as_deref() - } - pub fn local_runtime_paths(&self) -> Option<&ExecServerRuntimePaths> { self.local_runtime_paths.as_ref() } @@ -718,7 +747,6 @@ mod tests { let environment = Environment::create(/*exec_server_url*/ None, test_runtime_paths()) .expect("create environment"); - assert_eq!(environment.exec_server_url(), None); assert!(!environment.is_remote()); assert!(environment.info().await.is_ok()); } @@ -758,7 +786,7 @@ mod tests { } #[tokio::test] - async fn environment_manager_reports_remote_url() { + async fn environment_manager_creates_remote_environment_for_url() { let manager = EnvironmentManager::create_for_tests( Some("ws://127.0.0.1:8765".to_string()), Some(test_runtime_paths()), @@ -771,7 +799,6 @@ mod tests { Some(REMOTE_ENVIRONMENT_ID) ); assert!(environment.is_remote()); - assert_eq!(environment.exec_server_url(), Some("ws://127.0.0.1:8765")); assert!(Arc::ptr_eq( &environment, &manager @@ -961,7 +988,7 @@ mod tests { assert_eq!(environment.local_runtime_paths(), Some(&runtime_paths)); let manager = EnvironmentManager::create_for_tests( - environment.exec_server_url().map(str::to_owned), + /*exec_server_url*/ None, Some( environment .local_runtime_paths() @@ -1030,7 +1057,6 @@ mod tests { .get_environment("executor-a") .expect("first remote environment"); assert!(first.is_remote()); - assert_eq!(first.exec_server_url(), Some("ws://127.0.0.1:8765")); assert_eq!(manager.default_environment_id(), None); manager @@ -1044,7 +1070,6 @@ mod tests { .get_environment("executor-a") .expect("second remote environment"); assert!(second.is_remote()); - assert_eq!(second.exec_server_url(), Some("ws://127.0.0.1:9876")); assert!(!Arc::ptr_eq(&first, &second)); } diff --git a/codex-rs/exec-server/src/environment_provider.rs b/codex-rs/exec-server/src/environment_provider.rs index ffb5cbd116..69db34cf65 100644 --- a/codex-rs/exec-server/src/environment_provider.rs +++ b/codex-rs/exec-server/src/environment_provider.rs @@ -177,25 +177,17 @@ mod tests { assert!(!environments.contains_key(LOCAL_ENVIRONMENT_ID)); let remote_environment = &environments[REMOTE_ENVIRONMENT_ID]; assert!(remote_environment.is_remote()); - assert_eq!( - remote_environment.exec_server_url(), - Some("ws://127.0.0.1:8765") - ); assert_eq!( default, EnvironmentDefault::EnvironmentId(REMOTE_ENVIRONMENT_ID.to_string()) ); } - #[tokio::test] - async fn default_provider_normalizes_exec_server_url() { - let provider = DefaultEnvironmentProvider::new(Some(" ws://127.0.0.1:8765 ".to_string())); - let snapshot = provider.snapshot().await.expect("environments"); - let environments: HashMap<_, _> = snapshot.environments.into_iter().collect(); - + #[test] + fn normalizes_exec_server_url() { assert_eq!( - environments[REMOTE_ENVIRONMENT_ID].exec_server_url(), - Some("ws://127.0.0.1:8765") + normalize_exec_server_url(Some(" ws://127.0.0.1:8765 ".to_string())), + (Some("ws://127.0.0.1:8765".to_string()), false) ); } } diff --git a/codex-rs/exec-server/src/environment_toml.rs b/codex-rs/exec-server/src/environment_toml.rs index bfc25873ba..51e97b9bc9 100644 --- a/codex-rs/exec-server/src/environment_toml.rs +++ b/codex-rs/exec-server/src/environment_toml.rs @@ -389,12 +389,8 @@ mod tests { assert!(include_local); assert!(!environments.contains_key(LOCAL_ENVIRONMENT_ID)); - assert_eq!( - environments["devbox"].exec_server_url(), - Some("ws://127.0.0.1:8765") - ); + assert!(environments["devbox"].is_remote()); assert!(environments["ssh-dev"].is_remote()); - assert_eq!(environments["ssh-dev"].exec_server_url(), None); assert_eq!( default, EnvironmentDefault::EnvironmentId("ssh-dev".to_string()) diff --git a/codex-rs/exec-server/src/lib.rs b/codex-rs/exec-server/src/lib.rs index c2c8cc5194..9dafe1182e 100644 --- a/codex-rs/exec-server/src/lib.rs +++ b/codex-rs/exec-server/src/lib.rs @@ -68,6 +68,7 @@ pub use environment::CODEX_EXEC_SERVER_URL_ENV_VAR; pub use environment::Environment; pub use environment::EnvironmentManager; pub use environment::LOCAL_ENVIRONMENT_ID; +pub use environment::PendingEnvironmentRegistration; pub use environment::REMOTE_ENVIRONMENT_ID; pub use environment_provider::DefaultEnvironmentProvider; pub use environment_provider::EnvironmentProvider; 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..7503ec2300 --- /dev/null +++ b/codex-rs/exec-server/tests/pending_environment.rs @@ -0,0 +1,70 @@ +mod common; + +use codex_exec_server::EnvironmentManager; +use common::exec_server::exec_server; +use futures::poll; + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn pending_environment_connects_and_reconnects_after_completion() -> anyhow::Result<()> { + let mut server = exec_server().await?; + let mut proxy = server.disconnectable_websocket_proxy().await?; + let manager = EnvironmentManager::without_environments(); + let registration = manager.register_pending_environment("tools".to_string())?; + let environment = manager.get_environment("tools").expect("environment"); + + registration.complete(Ok(proxy.websocket_url().to_string()))?; + environment.wait_until_ready().await?; + proxy.pause_and_disconnect().await?; + proxy.resume()?; + environment.info().await?; + server.shutdown().await?; + Ok(()) +} + +#[tokio::test] +async fn failure_and_dropped_registration_are_terminal() -> anyhow::Result<()> { + let manager = EnvironmentManager::without_environments(); + let failed = manager.register_pending_environment("failed".to_string())?; + let failed_environment = manager.get_environment("failed").expect("environment"); + failed.complete(Err("provisioning failed".to_string()))?; + let error = failed_environment.wait_until_ready().await.unwrap_err(); + let message = error.to_string(); + assert!(message.ends_with("environment unavailable: provisioning failed")); + + let dropped = manager.register_pending_environment("dropped".to_string())?; + let dropped_environment = manager.get_environment("dropped").expect("environment"); + drop(dropped); + let error = dropped_environment.wait_until_ready().await.unwrap_err(); + let message = error.to_string(); + assert!(message.contains("registration ended before completion")); + assert!(manager.get_environment("failed").is_some()); + assert!(manager.get_environment("dropped").is_some()); + + let invalid = manager.register_pending_environment("invalid".to_string())?; + let invalid_environment = manager.get_environment("invalid").expect("environment"); + let error = invalid.complete(Ok(String::new())).unwrap_err(); + assert!(error.to_string().contains("requires an exec-server url")); + let error = invalid_environment.wait_until_ready().await.unwrap_err(); + assert!(error.to_string().contains("requires an exec-server url")); + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn late_completion_is_isolated_from_replacement() -> anyhow::Result<()> { + let mut server = exec_server().await?; + let manager = EnvironmentManager::without_environments(); + let old_registration = manager.register_pending_environment("tools".to_string())?; + let old_environment = manager.get_environment("tools").expect("old environment"); + let current_registration = manager.register_pending_environment("tools".to_string())?; + let current = manager.get_environment("tools").expect("current"); + + old_registration.complete(Ok(server.websocket_url().to_string()))?; + old_environment.wait_until_ready().await?; + let mut current_readiness = Box::pin(current.wait_until_ready()); + assert!(poll!(&mut current_readiness).is_pending()); + + current_registration.complete(Ok(server.websocket_url().to_string()))?; + current_readiness.await?; + server.shutdown().await?; + Ok(()) +}