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/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<NoiseRendezvousConnectBundle, ExecServerError>> {
+        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(())
+}

codex-rs/exec-server/tests/pending_environment.rs

@@ -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(())
-}

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<NoiseRendezvousConnectBundle, ExecServerError>> {
+        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(&registry)
+        .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(&registry)
+        .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(&registry).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?;