Add explicit remote executor connection refresh (#40710)

## Why Planned executor replacement needs a fresh session without waiting for the old session's transient-disconnect recovery to finish. ## What changed - Add `Environment::refresh_connection` for remote Noise registry-backed environments. It performs a fresh registry lookup, reuses a healthy session when the executor identity is unchanged, and connects to a replacement when it has changed. - Retire superseded sessions and connection attempts so they cannot publish stale state, accept late RPC results, or replay outstanding work. - Preserve the existing environment and filesystem handles while replacing the underlying client, and require a live status probe before refresh succeeds. ## Testing Add coverage for replacement and session reuse, recovery and connection races, lookup and handshake failures, handle preservation, and late RPC responses. GitOrigin-RevId: f1d11208cbfe8af8feb25f6b6b8100da82169a99
2026-10-25 18:28:38 +00:00 · 2026-08-25 21:23:17 +00:00
parent 5ca4175295
commit eb49f491c6
8 changed files with 1029 additions and 83 deletions
@@ -0,0 +1,620 @@
+use std::sync::Arc;
+use std::sync::Mutex;
+use std::time::Duration;
+
+use anyhow::Result;
+use futures::future::BoxFuture;
+use pretty_assertions::assert_eq;
+use tokio::net::TcpListener;
+use tokio::sync::Notify;
+use tokio::sync::oneshot;
+use tokio::task::JoinSet;
+use tokio_util::task::AbortOnDropHandle;
+
+use super::*;
+use crate::ExecServerRuntimePaths;
+use crate::NoiseChannelIdentity;
+use crate::ProcessId;
+use crate::relay::HarnessKeyValidator;
+use crate::relay::run_multiplexed_environment;
+use crate::server::ConnectionProcessor;
+use codex_http_client::HttpClientFactory;
+use codex_http_client::OutboundProxyPolicy;
+
+// These tests exercise real sockets, not keepalive deadlines. The shared unit-test
+// Pong timeout is only 100 ms and can expire during Noise handshakes under load.
+// A blocking task prevents paused time from auto-advancing while socket I/O is
+// pending; dropping the returned sender releases it without polling or sleeping.
+fn freeze_clock() -> std::sync::mpsc::Sender<()> {
+    tokio::time::pause();
+    let (guard, dropped) = std::sync::mpsc::channel();
+    tokio::task::spawn_blocking(move || {
+        let _ = dropped.recv();
+    });
+    guard
+}
+
+#[derive(Clone)]
+struct Target {
+    url: String,
+    identity: NoiseChannelIdentity,
+    registration: String,
+}
+
+struct Registry {
+    target: Mutex<Target>,
+    next_lookup: Mutex<Option<oneshot::Receiver<()>>>,
+    lookup_started: Notify,
+}
+
+impl Registry {
+    fn block_next_lookup(&self) -> oneshot::Sender<()> {
+        let (tx, rx) = oneshot::channel();
+        *self.next_lookup.lock().unwrap() = Some(rx);
+        tx
+    }
+}
+
+impl NoiseRendezvousConnectProvider for Registry {
+    fn connect_bundle(
+        &self,
+        _: NoiseChannelPublicKey,
+    ) -> BoxFuture<'_, Result<NoiseRendezvousConnectBundle, ExecServerError>> {
+        Box::pin(async move {
+            let target = self.target.lock().unwrap().clone();
+            let block = self.next_lookup.lock().unwrap().take();
+            if let Some(block) = block {
+                self.lookup_started.notify_one();
+                block
+                    .await
+                    .map_err(|_| ExecServerError::Protocol("test lookup failed".to_owned()))?;
+            }
+            Ok(NoiseRendezvousConnectBundle {
+                websocket_url: target.url,
+                environment_id: "environment".to_owned(),
+                executor_registration_id: target.registration,
+                executor_public_key: target.identity.public_key(),
+                harness_key_authorization: "authorization".to_owned(),
+            })
+        })
+    }
+}
+
+#[derive(Clone, Default)]
+struct Validator {
+    handshake: Option<Arc<Notify>>,
+    started: Arc<Notify>,
+}
+
+impl HarnessKeyValidator for Validator {
+    async fn validate_harness_key(
+        &self,
+        _: &NoiseChannelPublicKey,
+        _: &str,
+    ) -> Result<(), ExecServerError> {
+        if let Some(handshake) = &self.handshake {
+            self.started.notify_one();
+            handshake.notified().await;
+        }
+        Ok(())
+    }
+}
+
+struct Executor {
+    target: Target,
+    _server: AbortOnDropHandle<()>,
+}
+
+impl Executor {
+    async fn start(validator: Validator) -> Result<Self> {
+        let listener = TcpListener::bind("127.0.0.1:0").await?;
+        let target = Target {
+            url: format!("ws://{}", listener.local_addr()?),
+            identity: NoiseChannelIdentity::generate()?,
+            registration: uuid::Uuid::new_v4().to_string(),
+        };
+        let executor = target.clone();
+        let processor = ConnectionProcessor::new(ExecServerRuntimePaths::new(
+            std::env::current_exe()?,
+            /*codex_linux_sandbox_exe*/ None,
+        )?);
+        let server = tokio::spawn(async move {
+            let mut connections = JoinSet::new();
+            loop {
+                tokio::select! {
+                    socket = listener.accept() => {
+                        let (socket, _) = socket.unwrap();
+                        let executor = executor.clone();
+                        let processor = processor.clone();
+                        let validator = validator.clone();
+                        connections.spawn(async move {
+                            let socket = tokio_tungstenite::accept_async(socket).await.unwrap();
+                            run_multiplexed_environment(socket, processor, "environment".to_owned(), executor.registration, executor.identity, validator).await;
+                        });
+                    }
+                    _ = connections.join_next(), if !connections.is_empty() => {}
+                }
+            }
+        });
+        Ok(Self {
+            target,
+            _server: AbortOnDropHandle::new(server),
+        })
+    }
+
+    fn client(&self) -> Result<(LazyRemoteExecServerClient, Arc<Registry>)> {
+        let registry = Arc::new(Registry {
+            target: Mutex::new(self.target.clone()),
+            next_lookup: Mutex::new(None),
+            lookup_started: Notify::new(),
+        });
+        let client = LazyRemoteExecServerClient::new(
+            ExecServerTransportParams::NoiseRendezvous {
+                provider: registry.clone(),
+                identity: NoiseChannelIdentity::generate()?,
+            },
+            HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault),
+        );
+        Ok((client, registry))
+    }
+}
+
+async fn disconnect(client: &ExecServerClient) {
+    let rpc_client = {
+        let connection = client.inner.connection.lock().unwrap();
+        let ConnectionStatus::Connected(rpc_client) = &connection.status else {
+            panic!("expected connected session")
+        };
+        Arc::clone(rpc_client)
+    };
+    rpc_client.close_transport().await;
+}
+
+#[tokio::test]
+async fn refresh_cancels_old_recovery_and_connects_without_resuming_old_session() -> Result<()> {
+    let _clock = freeze_clock();
+    let old = Executor::start(Validator::default()).await?;
+    let new = Executor::start(Validator::default()).await?;
+    let (client, registry) = old.client()?;
+    let original = client.get().await?;
+    let process = original
+        .register_session(&ProcessId::from("old-process"))
+        .await?;
+    let blocked_recovery = registry.block_next_lookup();
+    disconnect(&original).await;
+    registry.lookup_started.notified().await;
+    *registry.target.lock().unwrap() = new.target.clone();
+    let refresh = client.refresh_connection();
+    tokio::pin!(refresh);
+    let started = tokio::time::Instant::now();
+    assert!(futures::poll!(refresh.as_mut()).is_pending());
+    assert!(original.inner.retired.is_cancelled());
+    assert!(original.is_disconnected());
+    assert_eq!(started.elapsed(), Duration::ZERO);
+    refresh.await?;
+    let replacement = client.get().await?;
+    assert_ne!(original.session_id(), replacement.session_id());
+    assert_eq!(
+        *client.environment_connection_state_tx.borrow(),
+        EnvironmentConnectionState::Connected
+    );
+    assert!(matches!(
+        process.write(b"never replay".to_vec()).await,
+        Err(ExecServerError::Disconnected(_))
+    ));
+    // Releasing an old registry response cannot reinstall or disconnect the replacement.
+    let _ = blocked_recovery.send(());
+    tokio::task::yield_now().await;
+    assert!(Arc::ptr_eq(&client.get().await?.inner, &replacement.inner));
+    replacement.environment_status().await?;
+    Ok(())
+}
+
+#[tokio::test]
+async fn refresh_retires_a_still_connected_old_executor() -> Result<()> {
+    let _clock = freeze_clock();
+    let old = Executor::start(Validator::default()).await?;
+    let new = Executor::start(Validator::default()).await?;
+    let (client, registry) = old.client()?;
+    let original = client.get().await?;
+    *registry.target.lock().unwrap() = new.target.clone();
+    client.refresh_connection().await?;
+    assert!(original.is_disconnected());
+    assert_ne!(original.session_id(), client.get().await?.session_id());
+    Ok(())
+}
+
+#[tokio::test]
+async fn refresh_preserves_a_current_session_across_registration_renewal() -> Result<()> {
+    let _clock = freeze_clock();
+    let executor = Executor::start(Validator::default()).await?;
+    let (client, registry) = executor.client()?;
+    let original = client.get().await?;
+    registry.target.lock().unwrap().registration = "renewed-registration".to_owned();
+    let concurrent = client.clone();
+    let concurrent = tokio::spawn(async move { concurrent.refresh_connection().await });
+    client.refresh_connection().await?;
+    concurrent.await??;
+    assert!(Arc::ptr_eq(&original.inner, &client.get().await?.inner));
+    original.environment_status().await?;
+    Ok(())
+}
+
+#[tokio::test]
+async fn refresh_cancels_a_stalled_initial_lookup() -> Result<()> {
+    let _clock = freeze_clock();
+    let old = Executor::start(Validator::default()).await?;
+    let new = Executor::start(Validator::default()).await?;
+    let (client, registry) = old.client()?;
+    let release = registry.block_next_lookup();
+    let initial = client.get();
+    tokio::pin!(initial);
+    assert!(futures::poll!(initial.as_mut()).is_pending());
+    *registry.target.lock().unwrap() = new.target.clone();
+    client.refresh_connection().await?;
+    assert!(initial.await.is_err());
+    let _ = release.send(());
+    client.get().await?.environment_status().await?;
+    Ok(())
+}
+
+#[tokio::test]
+async fn refresh_cancels_a_stalled_noise_handshake() -> Result<()> {
+    let _clock = freeze_clock();
+    let validator = Validator {
+        handshake: Some(Arc::new(Notify::new())),
+        ..Default::default()
+    };
+    let old = Executor::start(validator.clone()).await?;
+    let new = Executor::start(Validator::default()).await?;
+    let (client, registry) = old.client()?;
+    let connecting = client.clone();
+    let initial = tokio::spawn(async move { connecting.get().await });
+    validator.started.notified().await;
+    *registry.target.lock().unwrap() = new.target.clone();
+    client.refresh_connection().await?;
+    assert!(initial.await?.is_err());
+    validator.handshake.unwrap().notify_one();
+    client.get().await?.environment_status().await?;
+    assert_eq!(
+        *client.environment_connection_state_tx.borrow(),
+        EnvironmentConnectionState::Connected
+    );
+    Ok(())
+}
+
+#[tokio::test]
+async fn failed_refresh_lookup_leaves_the_existing_session_usable() -> Result<()> {
+    let _clock = freeze_clock();
+    let executor = Executor::start(Validator::default()).await?;
+    let (client, registry) = executor.client()?;
+    let original = client.get().await?;
+    drop(registry.block_next_lookup());
+    assert!(client.refresh_connection().await.is_err());
+    assert!(Arc::ptr_eq(&original.inner, &client.get().await?.inner));
+    original.environment_status().await?;
+    Ok(())
+}
+
+#[tokio::test]
+async fn failed_replacement_connection_keeps_old_handles_retired_and_get_retries() -> Result<()> {
+    let _clock = freeze_clock();
+    let old = Executor::start(Validator::default()).await?;
+    let new = Executor::start(Validator::default()).await?;
+    let (client, registry) = old.client()?;
+    let original = client.get().await?;
+    let process = original
+        .register_session(&ProcessId::from("old-process"))
+        .await?;
+
+    // The registry knows the replacement, but its endpoint drops the new connection.
+    let unavailable = TcpListener::bind("127.0.0.1:0").await?;
+    let target = Target {
+        url: format!("ws://{}", unavailable.local_addr()?),
+        ..new.target.clone()
+    };
+    let _rejected_connection = AbortOnDropHandle::new(tokio::spawn(async move {
+        drop(unavailable.accept().await.unwrap());
+    }));
+    *registry.target.lock().unwrap() = target;
+    assert!(matches!(
+        client.refresh_connection().await,
+        Err(ExecServerError::ConnectionAttempt(_))
+    ));
+    assert!(original.inner.retired.is_cancelled());
+    assert!(matches!(
+        original.environment_status().await,
+        Err(ExecServerError::Disconnected(_))
+    ));
+    assert!(matches!(
+        process.write(b"never replay".to_vec()).await,
+        Err(ExecServerError::Disconnected(_))
+    ));
+    assert_eq!(
+        *client.environment_connection_state_tx.borrow(),
+        EnvironmentConnectionState::Disconnected
+    );
+
+    // A later caller retries against the registry instead of reviving the retired client.
+    *registry.target.lock().unwrap() = new.target.clone();
+    let replacement = client.get().await?;
+    assert_ne!(original.session_id(), replacement.session_id());
+    replacement.environment_status().await?;
+    assert_eq!(
+        *client.environment_connection_state_tx.borrow(),
+        EnvironmentConnectionState::Connected
+    );
+    assert!(matches!(
+        process.write(b"still retired".to_vec()).await,
+        Err(ExecServerError::Disconnected(_))
+    ));
+    Ok(())
+}
+
+#[tokio::test]
+async fn superseded_refresh_lookup_does_not_retire_a_newer_session() -> Result<()> {
+    let _clock = freeze_clock();
+    let old = Executor::start(Validator::default()).await?;
+    let new = Executor::start(Validator::default()).await?;
+    let (client, registry) = old.client()?;
+    let original = client.get().await?;
+    let release = registry.block_next_lookup();
+    let refreshing = client.refresh_connection();
+    tokio::pin!(refreshing);
+    assert!(futures::poll!(refreshing.as_mut()).is_pending());
+    *registry.target.lock().unwrap() = new.target.clone();
+    original.inner.retire().await;
+    let replacement = client.get().await?;
+    release.send(()).unwrap();
+    refreshing.await?;
+    assert!(Arc::ptr_eq(&replacement.inner, &client.get().await?.inner));
+    assert!(!replacement.inner.retired.is_cancelled());
+    replacement.environment_status().await?;
+    Ok(())
+}
+
+#[tokio::test]
+async fn environment_refresh_preserves_environment_and_filesystem_handles() -> Result<()> {
+    let _clock = freeze_clock();
+    let old = Executor::start(Validator::default()).await?;
+    let new = Executor::start(Validator::default()).await?;
+    let (_, registry) = old.client()?;
+    let manager = crate::EnvironmentManager::from_snapshot(
+        crate::environment_provider::EnvironmentProviderSnapshot {
+            environments: Vec::new(),
+            default: crate::environment_provider::EnvironmentDefault::Disabled,
+            include_local: false,
+        },
+        /*local_runtime_paths*/ None,
+        HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault),
+    )?;
+    let environment = manager
+        .materialize_pending_noise_environment("environment".to_owned(), registry.clone())?;
+    manager.report_environment_provisioning_status(
+        "environment".to_owned(),
+        Ok(crate::EnvironmentReadyInfo {
+            selected_capability_roots: Vec::new(),
+        }),
+        registry.clone(),
+    )?;
+    environment.info().await?;
+    let filesystem = environment.get_filesystem();
+    *registry.target.lock().unwrap() = new.target.clone();
+    environment.refresh_connection().await?;
+    assert!(Arc::ptr_eq(
+        &environment,
+        &manager.get_environment("environment").unwrap()
+    ));
+    assert!(Arc::ptr_eq(&filesystem, &environment.get_filesystem()));
+    environment.info().await?;
+    Ok(())
+}
+
+#[tokio::test]
+async fn refresh_does_not_accept_cached_metadata_during_recovery() -> Result<()> {
+    let _clock = freeze_clock();
+    let executor = Executor::start(Validator::default()).await?;
+    let (client, registry) = executor.client()?;
+    let original = client.get().await?;
+    original.environment_info().await?;
+    let blocked_recovery = registry.block_next_lookup();
+    disconnect(&original).await;
+    registry.lookup_started.notified().await;
+    let started = tokio::time::Instant::now();
+    assert!(matches!(
+        client.refresh_connection().await,
+        Err(ExecServerError::Disconnected(_))
+    ));
+    assert_eq!(started.elapsed(), Duration::ZERO);
+    assert!(!original.inner.retired.is_cancelled());
+    drop(blocked_recovery);
+    Ok(())
+}
+
+#[tokio::test]
+async fn refresh_before_startup_marks_startup_finished() -> Result<()> {
+    let _clock = freeze_clock();
+    let executor = Executor::start(Validator::default()).await?;
+    let (client, _) = executor.client()?;
+    assert!(!client.startup_finished());
+    client.refresh_connection().await?;
+    assert!(client.startup_finished());
+    assert!(matches!(client.readiness_result(), Some(Ok(()))));
+    Ok(())
+}
+
+struct ControlledRpc {
+    client: ExecServerClient,
+    requests: tokio::sync::mpsc::Receiver<codex_exec_server_protocol::JSONRPCMessage>,
+    responses: tokio::sync::mpsc::Sender<crate::connection::JsonRpcConnectionEvent>,
+}
+
+async fn controlled_rpc() -> Result<ControlledRpc> {
+    use crate::connection::JsonRpcConnection;
+    use crate::connection::JsonRpcConnectionEvent;
+    use crate::connection::JsonRpcTransport;
+    use codex_exec_server_protocol::JSONRPCMessage;
+    use codex_exec_server_protocol::JSONRPCResponse;
+    let (outgoing_tx, mut requests) = tokio::sync::mpsc::channel(/*buffer*/ 8);
+    let (responses, incoming_rx) = tokio::sync::mpsc::channel(/*buffer*/ 8);
+    let connection = JsonRpcConnection {
+        outgoing_tx,
+        incoming_rx,
+        disconnected_rx: tokio::sync::watch::channel(/*init*/ false).1,
+        task_handles: Vec::new(),
+        transport: JsonRpcTransport::Plain,
+    };
+    let connecting = ExecServerClient::connect(connection, /*options*/ Default::default());
+    tokio::pin!(connecting);
+    assert!(futures::poll!(connecting.as_mut()).is_pending());
+    let Some(JSONRPCMessage::Request(initialize)) = requests.recv().await else {
+        anyhow::bail!("expected initialize request");
+    };
+    responses
+        .send(JsonRpcConnectionEvent::message(JSONRPCMessage::Response(
+            JSONRPCResponse {
+                id: initialize.id,
+                result: serde_json::json!({"sessionId": "controlled-session"}),
+            },
+        )))
+        .await?;
+    let client = connecting.await?;
+    assert!(matches!(
+        requests.recv().await,
+        Some(JSONRPCMessage::Notification(_))
+    ));
+    Ok(ControlledRpc {
+        client,
+        requests,
+        responses,
+    })
+}
+
+#[tokio::test]
+#[expect(
+    clippy::await_holding_invalid_type,
+    reason = "hold stream cleanup pending to exercise retirement ordering"
+)]
+async fn retirement_rejects_pending_mutation_before_stream_cleanup() -> Result<()> {
+    use crate::connection::JsonRpcConnectionEvent;
+    use codex_exec_server_protocol::JSONRPCMessage;
+    use codex_exec_server_protocol::JSONRPCResponse;
+    for response_queued in [false, true] {
+        let mut rpc = controlled_rpc().await?;
+        let call = rpc.client.fs_remove(crate::protocol::FsRemoveParams {
+            path: "file:///retired-file".parse()?,
+            recursive: None,
+            force: None,
+            follow_symlinks: None,
+            sandbox: None,
+        });
+        tokio::pin!(call);
+        assert!(futures::poll!(call.as_mut()).is_pending());
+        let Some(JSONRPCMessage::Request(request)) = rpc.requests.recv().await else {
+            anyhow::bail!("expected filesystem request");
+        };
+        let response = JSONRPCMessage::Response(JSONRPCResponse {
+            id: request.id,
+            result: serde_json::json!({}),
+        });
+        if response_queued {
+            rpc.responses
+                .send(JsonRpcConnectionEvent::message(response.clone()))
+                .await?;
+            let transport = rpc.client.rpc_client_without_recovery()?;
+            tokio::time::timeout(Duration::from_secs(5), async {
+                while transport.pending_request_count().await != 0 {
+                    tokio::task::yield_now().await;
+                }
+            })
+            .await?;
+        }
+        // Hold stream cleanup. Cover both a late response and one already queued
+        // for the caller when retirement begins; closing the socket alone misses the latter.
+        let streams = rpc.client.inner.http_body_streams_write_lock.lock().await;
+        let retirement = rpc.client.inner.retire();
+        tokio::pin!(retirement);
+        assert!(futures::poll!(retirement.as_mut()).is_pending());
+        if !response_queued {
+            rpc.responses
+                .send(JsonRpcConnectionEvent::message(response))
+                .await?;
+        }
+        assert!(matches!(call.await, Err(ExecServerError::Disconnected(_))));
+        drop(streams);
+        retirement.await;
+    }
+    Ok(())
+}
+
+#[tokio::test]
+#[expect(
+    clippy::await_holding_invalid_type,
+    reason = "hold stream cleanup pending to exercise retirement ordering"
+)]
+async fn retirement_rejects_pending_process_start_before_stream_cleanup() -> Result<()> {
+    use crate::connection::JsonRpcConnectionEvent;
+    use codex_exec_server_protocol::JSONRPCMessage;
+    use codex_exec_server_protocol::JSONRPCResponse;
+    for response_queued in [false, true] {
+        let mut rpc = controlled_rpc().await?;
+        let process_id = ProcessId::from("retired-process");
+        let call = rpc.client.start_process(
+            crate::protocol::ExecParams {
+                process_id: process_id.clone(),
+                argv: vec!["unused".to_owned()],
+                cwd: "file:///".parse()?,
+                shell_snapshot: None,
+                env_policy: None,
+                env: Default::default(),
+                tty: false,
+                pipe_stdin: false,
+                arg0: None,
+                sandbox: None,
+                enforce_managed_network: false,
+                managed_network: None,
+                network_proxy: None,
+            },
+            /*network_policy_decider*/ None,
+        );
+        tokio::pin!(call);
+        assert!(futures::poll!(call.as_mut()).is_pending());
+        let Some(JSONRPCMessage::Request(request)) = rpc.requests.recv().await else {
+            anyhow::bail!("expected process start request");
+        };
+        let response = JSONRPCMessage::Response(JSONRPCResponse {
+            id: request.id,
+            result: serde_json::json!({"processId": "retired-process"}),
+        });
+        if response_queued {
+            rpc.responses
+                .send(JsonRpcConnectionEvent::message(response.clone()))
+                .await?;
+            let state = rpc
+                .client
+                .inner
+                .get_session(&process_id)
+                .expect("pending process");
+            // The start task has a second response channel to the original caller.
+            tokio::time::timeout(Duration::from_secs(5), async {
+                while !state.recoverable.load(std::sync::atomic::Ordering::Acquire) {
+                    tokio::task::yield_now().await;
+                }
+            })
+            .await?;
+        }
+        let streams = rpc.client.inner.http_body_streams_write_lock.lock().await;
+        let retirement = rpc.client.inner.retire();
+        tokio::pin!(retirement);
+        assert!(futures::poll!(retirement.as_mut()).is_pending());
+        if !response_queued {
+            rpc.responses
+                .send(JsonRpcConnectionEvent::message(response))
+                .await?;
+        }
+        assert!(matches!(call.await, Err(ExecServerError::Disconnected(_))));
+        drop(streams);
+        retirement.await;
+    }
+    Ok(())
+}