From 32cd5d4eab898cfadb593eee61400812758386e7 Mon Sep 17 00:00:00 2001 From: TAFOYA-OAI Date: Tue, 14 Jul 2026 21:02:20 +0000 Subject: [PATCH] Defer Noise environment connections until registration (#33166) ## What changed - Replace pending WebSocket URL registration with deferred Noise environment registration that gates connection attempts on an explicit readiness signal. - Reuse the Noise rendezvous transport after readiness so reconnects request a fresh connection bundle. - Preserve terminal errors for failed or dropped registrations and keep late completion isolated from replacement environments. ## Testing - Add coverage for readiness gating, registration failure and replacement, eager Noise connections, and reconnection through a fresh rendezvous bundle. GitOrigin-RevId: 83e23fa03a02e3b2bdf1a83fe26d7ac461f55cf3 --- codex-rs/core/tests/suite/remote_env.rs | 138 ++++++------------ codex-rs/exec-server/src/client.rs | 4 +- codex-rs/exec-server/src/client_api.rs | 17 ++- codex-rs/exec-server/src/client_transport.rs | 38 +++-- codex-rs/exec-server/src/environment.rs | 55 ++++--- codex-rs/exec-server/src/lib.rs | 2 +- .../exec-server/tests/deferred_environment.rs | 135 +++++++++++++++++ .../exec-server/tests/pending_environment.rs | 70 --------- codex-rs/exec-server/tests/relay.rs | 126 ++++++++++++++++ 9 files changed, 369 insertions(+), 216 deletions(-) create mode 100644 codex-rs/exec-server/tests/deferred_environment.rs delete 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 89901cc1f9..abaace0c8d 100644 --- a/codex-rs/core/tests/suite/remote_env.rs +++ b/codex-rs/core/tests/suite/remote_env.rs @@ -7,8 +7,12 @@ use codex_core::compact::SUMMARIZATION_PROMPT; use codex_core::config::Constrained; use codex_exec_server::CopyOptions; use codex_exec_server::CreateDirectoryOptions; +use codex_exec_server::ExecServerError; use codex_exec_server::FileSystemSandboxContext; use codex_exec_server::LOCAL_ENVIRONMENT_ID; +use codex_exec_server::NoiseChannelPublicKey; +use codex_exec_server::NoiseRendezvousConnectBundle; +use codex_exec_server::NoiseRendezvousConnectProvider; use codex_exec_server::REMOTE_ENVIRONMENT_ID; use codex_exec_server::RemoveOptions; use codex_features::Feature; @@ -69,6 +73,7 @@ use core_test_support::wait_for_event; use core_test_support::wait_for_event_match; use futures::SinkExt; use futures::StreamExt; +use futures::future::BoxFuture; use pretty_assertions::assert_eq; use serde_json::Value; use serde_json::json; @@ -76,6 +81,9 @@ use std::collections::HashMap; use std::fs; use std::path::PathBuf; use std::process::Command; +use std::sync::Arc; +use std::sync::atomic::AtomicUsize; +use std::sync::atomic::Ordering; use std::time::Duration; use std::time::SystemTime; use std::time::UNIX_EPOCH; @@ -483,6 +491,25 @@ fn tool_names(body: &Value) -> Vec { .collect() } +#[derive(Default)] +struct FailingNoiseConnectProvider { + calls: AtomicUsize, +} + +impl NoiseRendezvousConnectProvider for FailingNoiseConnectProvider { + fn connect_bundle( + &self, + _: NoiseChannelPublicKey, + ) -> BoxFuture<'_, std::result::Result> { + self.calls.fetch_add(1, Ordering::Relaxed); + Box::pin(async { + Err(ExecServerError::Protocol( + "test Noise connection failed".to_string(), + )) + }) + } +} + async fn wait_for_response_request_count(response_mock: &ResponseMock, expected_count: usize) { timeout(Duration::from_secs(5), async { while response_mock.requests().len() < expected_count { @@ -494,8 +521,7 @@ async fn wait_for_response_request_count(response_mock: &ResponseMock, expected_ } #[tokio::test(flavor = "multi_thread", worker_threads = 2)] -async fn deferred_executor_updates_context_and_tools_after_startup() -> Result<()> { - let listener = TcpListener::bind("127.0.0.1:0").await?; +async fn deferred_executor_starts_noise_connection_after_registration() -> Result<()> { let server = start_mock_server().await; let wait_call_id = "wait-for-startup"; let response_mock = mount_sse_sequence( @@ -515,47 +541,24 @@ async fn deferred_executor_updates_context_and_tools_after_startup() -> Result<( ]), sse(vec![ ev_response_created("resp-2"), - ev_function_call( - "request-permissions", - "request_permissions", - &json!({ - "reason": "Verify that the ready environment is used.", - "permissions": { - "network": { "enabled": true } - } - }) - .to_string(), - ), + ev_assistant_message("msg-2", "done"), ev_completed("resp-2"), ]), - sse(vec![ - ev_response_created("resp-3"), - ev_assistant_message("msg-3", "done"), - ev_completed("resp-3"), - ]), ], ) .await; 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())?; + let provider = Arc::new(FailingNoiseConnectProvider::default()); + let registration = environment_manager + .register_deferred_noise_environment(REMOTE_ENVIRONMENT_ID.to_string(), provider.clone())?; test.codex .submit(Op::UserInput { @@ -581,83 +584,26 @@ async fn deferred_executor_updates_context_and_tools_after_startup() -> Result<( .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!( - event, - EventMsg::RequestPermissions(_) | EventMsg::TurnComplete(_) - ) - }) - .await; - let EventMsg::RequestPermissions(permission_request) = event else { - panic!("ready environment should be available to request_permissions: {event:?}"); - }; - assert_eq!( - permission_request.environment_id.as_deref(), - Some(REMOTE_ENVIRONMENT_ID) - ); - test.codex - .submit(Op::RequestPermissionsResponse { - id: permission_request.call_id, - response: RequestPermissionsResponse { - permissions: RequestPermissionProfile::default(), - scope: PermissionGrantScope::Turn, - strict_auto_review: false, - }, - }) - .await?; + assert_eq!(provider.calls.load(Ordering::Relaxed), 0); + registration.complete(Ok(()))?; wait_for_event(&test.codex, |event| { matches!(event, EventMsg::TurnComplete(_)) }) .await; + assert_eq!(provider.calls.load(Ordering::Relaxed), 1); let requests = response_mock.requests(); - assert_eq!(requests.len(), 3); + assert_eq!(requests.len(), 2); let starting_tools = tool_names(&requests[0].body_json()); - let ready_tools = tool_names(&requests[1].body_json()); assert!(starting_tools.contains(&"wait_for_environment".to_string())); assert!(!starting_tools.contains(&"exec_command".to_string())); - assert!(ready_tools.contains(&"exec_command".to_string())); - assert!(ready_tools.contains(&"wait_for_environment".to_string())); let (wait_output, _) = requests[1] .function_call_output_content_and_success(wait_call_id) .context("wait_for_environment output should be present")?; - assert_eq!( - serde_json::from_str::(&wait_output.context("wait output should contain text")?)?, - json!({ - "environment_id": REMOTE_ENVIRONMENT_ID, - "status": "ready", - }) - ); assert!( - requests[0] - .message_input_texts("user") - .iter() - .any(|text| text.contains("starting")) - ); - let ready_user_context = requests[1].message_input_texts("user"); - assert_eq!( - ready_user_context - .iter() - .filter(|text| text.contains("zsh")) - .count(), - 1 - ); - let final_user_context = requests[2].message_input_texts("user"); - assert_eq!( - final_user_context - .iter() - .filter(|text| text.contains("starting")) - .count(), - 1 - ); - assert_eq!( - final_user_context - .iter() - .filter(|text| text.contains("zsh")) - .count(), - 1 + wait_output + .context("wait output should contain text")? + .contains("failed to start") ); Ok(()) @@ -792,8 +738,10 @@ async fn deferred_executor_wait_reports_startup_failure() -> Result<()> { .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())?; + let registration = environment_manager.register_deferred_noise_environment( + REMOTE_ENVIRONMENT_ID.to_string(), + Arc::new(FailingNoiseConnectProvider::default()), + )?; test.codex .submit(Op::UserInput { diff --git a/codex-rs/exec-server/src/client.rs b/codex-rs/exec-server/src/client.rs index 4475ffe45b..654db07023 100644 --- a/codex-rs/exec-server/src/client.rs +++ b/codex-rs/exec-server/src/client.rs @@ -453,8 +453,8 @@ impl LazyRemoteExecServerClient { fn can_reconnect(&self) -> bool { matches!( self.transport_params, - ExecServerTransportParams::WebSocketUrl { .. } - | ExecServerTransportParams::PendingWebSocketUrl(..) + ExecServerTransportParams::Deferred(_) + | ExecServerTransportParams::WebSocketUrl { .. } | ExecServerTransportParams::NoiseRendezvous { .. } ) } diff --git a/codex-rs/exec-server/src/client_api.rs b/codex-rs/exec-server/src/client_api.rs index 5e72508df6..cd284a35e8 100644 --- a/codex-rs/exec-server/src/client_api.rs +++ b/codex-rs/exec-server/src/client_api.rs @@ -90,17 +90,23 @@ pub(crate) struct StdioExecServerCommand { pub cwd: Option, } -pub(crate) type PendingExecServerUrl = Shared>>; +pub(crate) type DeferredEnvironmentReadiness = Shared>>; + +#[derive(Clone)] +pub(crate) struct Deferred { + pub readiness: DeferredEnvironmentReadiness, + pub transport: T, +} /// Parameters used to connect to a remote exec-server environment. #[derive(Clone)] pub(crate) enum ExecServerTransportParams { + Deferred(Box>), WebSocketUrl { websocket_url: String, connect_timeout: Duration, initialize_timeout: Duration, }, - PendingWebSocketUrl(PendingExecServerUrl), NoiseRendezvous { provider: Arc, identity: NoiseChannelIdentity, @@ -115,6 +121,10 @@ pub(crate) enum ExecServerTransportParams { impl std::fmt::Debug for ExecServerTransportParams { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { match self { + Self::Deferred(deferred) => f + .debug_struct("Deferred") + .field("transport", &deferred.transport) + .finish_non_exhaustive(), Self::WebSocketUrl { websocket_url, connect_timeout, @@ -125,9 +135,6 @@ 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 99cb8a9aa8..14cc4cc352 100644 --- a/codex-rs/exec-server/src/client_transport.rs +++ b/codex-rs/exec-server/src/client_transport.rs @@ -94,27 +94,35 @@ impl ExecServerClient { pub(crate) async fn connect_for_transport( transport_params: ExecServerTransportParams, ) -> Result { + let (transport_params, deferred_readiness) = match transport_params { + ExecServerTransportParams::Deferred(deferred) => { + (deferred.transport, Some(deferred.readiness)) + } + transport_params => (transport_params, None), + }; + + if let Some(readiness) = deferred_readiness { + readiness + .await + .unwrap_or_else(|_| { + Err("environment registration ended before completion".to_string()) + }) + .map_err(|message| { + ExecServerError::Disconnected(format!("environment unavailable: {message}")) + })?; + } + let (websocket_url, connect_timeout, initialize_timeout) = match transport_params { + ExecServerTransportParams::Deferred(_) => { + return Err(ExecServerError::Protocol( + "nested deferred exec-server transports are unsupported".to_string(), + )); + } ExecServerTransportParams::WebSocketUrl { websocket_url, connect_timeout, initialize_timeout, } => (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, - DEFAULT_REMOTE_EXEC_SERVER_CONNECT_TIMEOUT, - DEFAULT_REMOTE_EXEC_SERVER_INITIALIZE_TIMEOUT, - ) - } ExecServerTransportParams::NoiseRendezvous { provider, identity } => { let reconnect_strategy = ExecServerReconnectStrategy::NoiseRendezvous { provider: Arc::clone(&provider), diff --git a/codex-rs/exec-server/src/environment.rs b/codex-rs/exec-server/src/environment.rs index c840268e21..cae42e1eff 100644 --- a/codex-rs/exec-server/src/environment.rs +++ b/codex-rs/exec-server/src/environment.rs @@ -62,9 +62,9 @@ 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>); +/// The one-shot capability to complete a deferred environment registration. +#[must_use = "the deferred environment cannot connect until registration is completed"] +pub struct DeferredEnvironmentRegistration(oneshot::Sender>); pub const LOCAL_ENVIRONMENT_ID: &str = "local"; pub const REMOTE_ENVIRONMENT_ID: &str = "remote"; @@ -352,15 +352,20 @@ impl EnvironmentManager { Ok(()) } - /// Adds or replaces a remote environment whose stable URL will be supplied later. - pub fn register_pending_environment( + /// Adds or replaces a Noise rendezvous environment that will become ready later. + pub fn register_deferred_noise_environment( &self, environment_id: String, - ) -> Result { + provider: Arc, + ) -> Result { validate_environment_id(&environment_id)?; - let (completion, websocket_url) = oneshot::channel(); + let identity = noise_channel_identity()?; + let (completion, readiness) = oneshot::channel(); let environment = Arc::new(Environment::remote_with_transport( - ExecServerTransportParams::PendingWebSocketUrl(websocket_url.shared()), + ExecServerTransportParams::Deferred(Box::new(crate::client_api::Deferred { + readiness: readiness.shared(), + transport: ExecServerTransportParams::NoiseRendezvous { provider, identity }, + })), self.local_runtime_paths.clone(), )); environment.start_connecting(); @@ -368,7 +373,7 @@ impl EnvironmentManager { .write() .unwrap_or_else(std::sync::PoisonError::into_inner) .insert(environment_id, environment); - Ok(PendingEnvironmentRegistration(completion)) + Ok(DeferredEnvironmentRegistration(completion)) } /// Adds or replaces a named remote environment that connects through an @@ -382,11 +387,7 @@ impl EnvironmentManager { provider: Arc, ) -> Result<(), ExecServerError> { validate_environment_id(&environment_id)?; - let identity = NoiseChannelIdentity::generate().map_err(|error| { - ExecServerError::Protocol(format!( - "failed to generate Noise harness identity: {error}" - )) - })?; + let identity = noise_channel_identity()?; let environment = Arc::new(Environment::remote_with_transport( ExecServerTransportParams::NoiseRendezvous { provider, identity }, self.local_runtime_paths.clone(), @@ -400,25 +401,23 @@ 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), - }; +impl DeferredEnvironmentRegistration { + /// Completes provisioning with readiness or a terminal error message. + pub fn complete(self, result: Result<(), String>) -> Result<(), ExecServerError> { self.0.send(result).map_err(|_| { - ExecServerError::Disconnected("pending environment registration is inactive".into()) + ExecServerError::Disconnected("deferred environment registration is inactive".into()) }) } } +fn noise_channel_identity() -> Result { + NoiseChannelIdentity::generate().map_err(|error| { + ExecServerError::Protocol(format!( + "failed to generate Noise harness identity: {error}" + )) + }) +} + fn validate_environment_id(environment_id: &str) -> Result<(), ExecServerError> { if environment_id.is_empty() { return Err(ExecServerError::Protocol( diff --git a/codex-rs/exec-server/src/lib.rs b/codex-rs/exec-server/src/lib.rs index 4de83da035..0f8a02b376 100644 --- a/codex-rs/exec-server/src/lib.rs +++ b/codex-rs/exec-server/src/lib.rs @@ -65,11 +65,11 @@ pub use environment::CODEX_EXEC_SERVER_NOISE_CHATGPT_ACCOUNT_ID_ENV_VAR; pub use environment::CODEX_EXEC_SERVER_NOISE_ENVIRONMENT_ID_ENV_VAR; pub use environment::CODEX_EXEC_SERVER_NOISE_REGISTRY_URL_ENV_VAR; pub use environment::CODEX_EXEC_SERVER_URL_ENV_VAR; +pub use environment::DeferredEnvironmentRegistration; pub use environment::Environment; pub use environment::EnvironmentManager; pub use environment::EnvironmentObservedStatus; 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/deferred_environment.rs b/codex-rs/exec-server/tests/deferred_environment.rs new file mode 100644 index 0000000000..847a3d7767 --- /dev/null +++ b/codex-rs/exec-server/tests/deferred_environment.rs @@ -0,0 +1,135 @@ +use std::sync::Arc; +use std::sync::atomic::AtomicUsize; +use std::sync::atomic::Ordering; + +use codex_exec_server::EnvironmentManager; +use codex_exec_server::ExecServerError; +use codex_exec_server::NoiseChannelPublicKey; +use codex_exec_server::NoiseRendezvousConnectBundle; +use codex_exec_server::NoiseRendezvousConnectProvider; +use futures::FutureExt; +use futures::future::BoxFuture; +use futures::poll; +use pretty_assertions::assert_eq; + +#[derive(Default)] +struct FailingNoiseConnectProvider { + calls: AtomicUsize, +} + +impl FailingNoiseConnectProvider { + fn calls(&self) -> usize { + self.calls.load(Ordering::Relaxed) + } +} + +impl NoiseRendezvousConnectProvider for FailingNoiseConnectProvider { + fn connect_bundle( + &self, + _: NoiseChannelPublicKey, + ) -> BoxFuture<'_, Result> { + self.calls.fetch_add(1, Ordering::Relaxed); + async { + Err(ExecServerError::Protocol( + "test Noise provider called".to_string(), + )) + } + .boxed() + } +} + +#[tokio::test] +async fn deferred_environment_waits_before_connecting() -> anyhow::Result<()> { + let manager = EnvironmentManager::without_environments(); + let provider = Arc::new(FailingNoiseConnectProvider::default()); + let registration = + manager.register_deferred_noise_environment("tools".to_string(), provider.clone())?; + let environment = manager.get_environment("tools").expect("environment"); + let mut readiness = Box::pin(environment.wait_until_ready()); + + assert!(poll!(&mut readiness).is_pending()); + assert_eq!(provider.calls(), 0); + + registration.complete(Ok(()))?; + let error = readiness.await.unwrap_err(); + assert!(error.to_string().contains("test Noise provider called")); + assert_eq!(provider.calls(), 1); + Ok(()) +} + +#[tokio::test] +async fn failure_and_dropped_registration_are_terminal() -> anyhow::Result<()> { + let manager = EnvironmentManager::without_environments(); + let failed_provider = Arc::new(FailingNoiseConnectProvider::default()); + let failed = manager + .register_deferred_noise_environment("failed".to_string(), failed_provider.clone())?; + 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(); + assert!( + error + .to_string() + .ends_with("environment unavailable: provisioning failed") + ); + assert_eq!(failed_provider.calls(), 0); + + let dropped_provider = Arc::new(FailingNoiseConnectProvider::default()); + let dropped = manager + .register_deferred_noise_environment("dropped".to_string(), dropped_provider.clone())?; + let dropped_environment = manager.get_environment("dropped").expect("environment"); + drop(dropped); + let error = dropped_environment.wait_until_ready().await.unwrap_err(); + assert!( + error + .to_string() + .contains("registration ended before completion") + ); + assert_eq!(dropped_provider.calls(), 0); + assert!(manager.get_environment("failed").is_some()); + assert!(manager.get_environment("dropped").is_some()); + Ok(()) +} + +#[tokio::test] +async fn late_completion_is_isolated_from_replacement() -> anyhow::Result<()> { + let manager = EnvironmentManager::without_environments(); + let old_provider = Arc::new(FailingNoiseConnectProvider::default()); + let old_registration = + manager.register_deferred_noise_environment("tools".to_string(), old_provider.clone())?; + let old_environment = manager.get_environment("tools").expect("old environment"); + let current_provider = Arc::new(FailingNoiseConnectProvider::default()); + let current_registration = manager + .register_deferred_noise_environment("tools".to_string(), current_provider.clone())?; + let current = manager.get_environment("tools").expect("current"); + + old_registration.complete(Ok(()))?; + let old_error = old_environment.wait_until_ready().await.unwrap_err(); + assert!(old_error.to_string().contains("test Noise provider called")); + assert_eq!(old_provider.calls(), 1); + let mut current_readiness = Box::pin(current.wait_until_ready()); + assert!(poll!(&mut current_readiness).is_pending()); + assert_eq!(current_provider.calls(), 0); + + current_registration.complete(Ok(()))?; + let current_error = current_readiness.await.unwrap_err(); + assert!( + current_error + .to_string() + .contains("test Noise provider called") + ); + assert_eq!(current_provider.calls(), 1); + Ok(()) +} + +#[tokio::test] +async fn eager_noise_environment_connects_without_registration() -> anyhow::Result<()> { + let manager = EnvironmentManager::without_environments(); + let provider = Arc::new(FailingNoiseConnectProvider::default()); + manager.upsert_noise_environment("tools".to_string(), provider.clone())?; + let environment = manager.get_environment("tools").expect("environment"); + + let error = environment.wait_until_ready().await.unwrap_err(); + assert!(error.to_string().contains("test Noise provider called")); + assert_eq!(provider.calls(), 1); + Ok(()) +} diff --git a/codex-rs/exec-server/tests/pending_environment.rs b/codex-rs/exec-server/tests/pending_environment.rs deleted file mode 100644 index 7503ec2300..0000000000 --- a/codex-rs/exec-server/tests/pending_environment.rs +++ /dev/null @@ -1,70 +0,0 @@ -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(()) -} diff --git a/codex-rs/exec-server/tests/relay.rs b/codex-rs/exec-server/tests/relay.rs index 918cba7393..a5e44cf73c 100644 --- a/codex-rs/exec-server/tests/relay.rs +++ b/codex-rs/exec-server/tests/relay.rs @@ -6,6 +6,8 @@ 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; @@ -13,20 +15,24 @@ 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::SinkExt; use futures::StreamExt; +use futures::future::BoxFuture; use http::HeaderMap; use http::HeaderValue; use pretty_assertions::assert_eq; @@ -69,6 +75,126 @@ fn static_registry_auth_provider() -> codex_api::SharedAuthProvider { Arc::new(StaticRegistryAuthProvider) } +struct FreshBundleNoiseConnectProvider { + websocket_url: String, + executor_public_key: NoiseChannelPublicKey, + calls: AtomicUsize, +} + +impl FreshBundleNoiseConnectProvider { + fn calls(&self) -> usize { + self.calls.load(Ordering::Relaxed) + } +} + +impl NoiseRendezvousConnectProvider for FreshBundleNoiseConnectProvider { + fn connect_bundle( + &self, + _: NoiseChannelPublicKey, + ) -> BoxFuture<'_, Result> { + let call = self.calls.fetch_add(1, Ordering::Relaxed) + 1; + let bundle = NoiseRendezvousConnectBundle { + websocket_url: self.websocket_url.clone(), + environment_id: ENVIRONMENT_ID.to_string(), + executor_registration_id: EXECUTOR_REGISTRATION_ID.to_string(), + executor_public_key: self.executor_public_key.clone(), + harness_key_authorization: format!("{HARNESS_KEY_AUTHORIZATION}-{call}"), + }; + Box::pin(async move { Ok(bundle) }) + } +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn deferred_noise_environment_connects_and_reconnects_with_fresh_bundle() -> Result<()> { + let listener = TcpListener::bind("127.0.0.1:0").await?; + let rendezvous_url = format!("ws://{}", listener.local_addr()?); + let registry = MockServer::start().await; + Mock::given(method("POST")) + .and(path(format!( + "/cloud/environment/{ENVIRONMENT_ID}/register" + ))) + .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({ + "environment_id": ENVIRONMENT_ID, + "url": format!("{rendezvous_url}/relay?role=environment"), + "security_profile": "noise_hybrid_ik_v1", + "executor_registration_id": EXECUTOR_REGISTRATION_ID, + }))) + .expect(1) + .mount(®istry) + .await; + Mock::given(method("POST")) + .and(path(format!( + "/cloud/environment/{ENVIRONMENT_ID}/validate" + ))) + .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({ + "valid": true, + }))) + .expect(2) + .mount(®istry) + .await; + + let (codex_exe, codex_linux_sandbox_exe) = common::current_test_binary_helper_paths()?; + let runtime_paths = ExecServerRuntimePaths::new(codex_exe, codex_linux_sandbox_exe)?; + let config = RemoteEnvironmentConfig::new( + registry.uri(), + ENVIRONMENT_ID.to_string(), + static_registry_auth_provider(), + )?; + let remote_environment = tokio::spawn(codex_exec_server::run_remote_environment( + config, + runtime_paths, + )); + + let environment_websocket = accept_websocket(&listener, "environment").await?; + let provider = Arc::new(FreshBundleNoiseConnectProvider { + websocket_url: format!("{rendezvous_url}/relay?role=harness"), + executor_public_key: registered_executor_public_key(®istry).await?, + calls: AtomicUsize::new(0), + }); + let manager = EnvironmentManager::without_environments(); + let registration = manager + .register_deferred_noise_environment(ENVIRONMENT_ID.to_string(), provider.clone())?; + let environment = manager + .get_environment(ENVIRONMENT_ID) + .context("deferred Noise environment")?; + + assert_eq!(provider.calls(), 0); + registration.complete(Ok(()))?; + let harness_websocket = accept_websocket(&listener, "harness").await?; + let first_relay = tokio::spawn(proxy_relay_frames( + environment_websocket, + harness_websocket, + Arc::new(Mutex::new(Vec::new())), + )); + let initial_info = timeout(TEST_TIMEOUT, environment.info()) + .await + .context("deferred Noise environment should become ready")??; + assert_eq!(provider.calls(), 1); + + first_relay.abort(); + let _ = first_relay.await; + let first_reconnected_websocket = accept_websocket(&listener, "reconnected peer").await?; + let second_reconnected_websocket = accept_websocket(&listener, "reconnected peer").await?; + let second_relay = tokio::spawn(proxy_relay_frames( + first_reconnected_websocket, + second_reconnected_websocket, + Arc::new(Mutex::new(Vec::new())), + )); + let recovered_info = timeout(TEST_TIMEOUT, environment.info()) + .await + .context("deferred Noise environment should reconnect")??; + + assert_eq!(recovered_info, initial_info); + assert_eq!(provider.calls(), 2); + registry.verify().await; + + second_relay.abort(); + remote_environment.abort(); + let _ = second_relay.await; + let _ = remote_environment.await; + 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?;