Refresh signed exec-server URLs on reconnect (#28374)

## Summary

- add a provider API that supplies a fresh signed WebSocket URL for each
remote exec-server connection
- refresh the signed URL after disconnects and retry once when a
handshake returns `401 Unauthorized`
- allow `EnvironmentManager` consumers to register remote environments
backed by the URL provider

## Tests

- `just test -p codex-exec-server -E
'test(remote_websocket_client_refreshes_url_after_unauthorized_handshake)
| test(remote_websocket_client_refreshes_url_after_disconnect)'` — 2
passed
- `cargo check -p codex-core-api` — passed
- `just fix -p codex-exec-server` — passed
- `just fix -p codex-core-api` — no test targets; no-op
- `just fmt` — passed
- `just test -p codex-exec-server` — 187 passed; 32 unrelated macOS
sandbox tests could not invoke nested `sandbox-exec` (`Operation not
permitted`)

codex-rs/exec-server/src/client_transport.rs

@@ -117,17 +117,8 @@ impl ExecServerClient {
                     connect_timeout: DEFAULT_REMOTE_EXEC_SERVER_CONNECT_TIMEOUT,
                     initialize_timeout: DEFAULT_REMOTE_EXEC_SERVER_INITIALIZE_TIMEOUT,
                 };
-                let bundle = provider.connect_bundle(identity.public_key()).await?;
                 let (connection, options) =
-                    Self::open_noise_rendezvous_connection(NoiseRendezvousConnectArgs {
-                        bundle,
-                        harness_identity: identity,
-                        client_name: ENVIRONMENT_CLIENT_NAME.to_string(),
-                        connect_timeout: DEFAULT_REMOTE_EXEC_SERVER_CONNECT_TIMEOUT,
-                        initialize_timeout: DEFAULT_REMOTE_EXEC_SERVER_INITIALIZE_TIMEOUT,
-                        resume_session_id: None,
-                    })
-                    .await?;
+                    Self::open_initial_noise_rendezvous_connection(&provider, &identity).await?;
                 Self::connect_with_recovery(connection, options, Some(reconnect_strategy)).await
             }
             crate::client_api::ExecServerTransportParams::StdioCommand {
@@ -145,6 +136,40 @@ impl ExecServerClient {
         }
     }
 
+    async fn open_initial_noise_rendezvous_connection(
+        provider: &Arc<dyn NoiseRendezvousConnectProvider>,
+        identity: &NoiseChannelIdentity,
+    ) -> Result<(JsonRpcConnection, ExecServerClientConnectOptions), ExecServerError> {
+        let open_connection = |bundle: NoiseRendezvousConnectBundle| {
+            Self::open_noise_rendezvous_connection(NoiseRendezvousConnectArgs {
+                bundle,
+                harness_identity: identity.clone(),
+                client_name: ENVIRONMENT_CLIENT_NAME.to_string(),
+                connect_timeout: DEFAULT_REMOTE_EXEC_SERVER_CONNECT_TIMEOUT,
+                initialize_timeout: DEFAULT_REMOTE_EXEC_SERVER_INITIALIZE_TIMEOUT,
+                resume_session_id: None,
+            })
+        };
+        let bundle = provider.connect_bundle(identity.public_key()).await?;
+        match open_connection(bundle).await {
+            Err(error)
+                if matches!(
+                    &error,
+                    ExecServerError::WebSocketConnect { source, .. }
+                        if matches!(
+                            source,
+                            tokio_tungstenite::tungstenite::Error::Http(response)
+                                if response.status().as_u16() == 401
+                        )
+                ) =>
+            {
+                let bundle = provider.connect_bundle(identity.public_key()).await?;
+                open_connection(bundle).await
+            }
+            result => result,
+        }
+    }
+
     pub async fn connect_websocket(
         args: RemoteExecServerConnectArgs,
     ) -> Result<Self, ExecServerError> {
@@ -326,3 +351,7 @@ fn stdio_command_process(stdio_command: &StdioExecServerCommand) -> Command {
     command.process_group(0);
     command
 }
+
+#[cfg(test)]
+#[path = "client_transport_tests.rs"]
+mod tests;

codex-rs/exec-server/src/client_transport_tests.rs

@@ -0,0 +1,112 @@
+use std::collections::VecDeque;
+use std::sync::Arc;
+use std::sync::Mutex;
+
+use anyhow::Result;
+use futures::future::BoxFuture;
+use pretty_assertions::assert_eq;
+use tokio::io::AsyncReadExt;
+use tokio::io::AsyncWriteExt;
+use tokio::net::TcpListener;
+use tokio_tungstenite::accept_async;
+
+use super::ExecServerClient;
+use crate::ExecServerError;
+use crate::NoiseChannelIdentity;
+use crate::NoiseChannelPublicKey;
+use crate::NoiseRendezvousConnectBundle;
+use crate::NoiseRendezvousConnectProvider;
+
+struct SequenceNoiseConnectProvider {
+    bundles: Mutex<VecDeque<NoiseRendezvousConnectBundle>>,
+    returned_urls: Mutex<Vec<String>>,
+}
+
+impl SequenceNoiseConnectProvider {
+    fn new(bundles: Vec<NoiseRendezvousConnectBundle>) -> Self {
+        Self {
+            bundles: Mutex::new(bundles.into()),
+            returned_urls: Mutex::new(Vec::new()),
+        }
+    }
+
+    fn returned_urls(&self) -> Vec<String> {
+        self.returned_urls
+            .lock()
+            .unwrap_or_else(std::sync::PoisonError::into_inner)
+            .clone()
+    }
+}
+
+impl NoiseRendezvousConnectProvider for SequenceNoiseConnectProvider {
+    fn connect_bundle(
+        &self,
+        _: NoiseChannelPublicKey,
+    ) -> BoxFuture<'_, Result<NoiseRendezvousConnectBundle, ExecServerError>> {
+        let result = self
+            .bundles
+            .lock()
+            .unwrap_or_else(std::sync::PoisonError::into_inner)
+            .pop_front()
+            .ok_or_else(|| ExecServerError::Protocol("test Noise provider exhausted".to_string()));
+        if let Ok(bundle) = &result {
+            self.returned_urls
+                .lock()
+                .unwrap_or_else(std::sync::PoisonError::into_inner)
+                .push(bundle.websocket_url.clone());
+        }
+        Box::pin(async move { result })
+    }
+}
+
+fn test_bundle(websocket_url: String) -> Result<NoiseRendezvousConnectBundle> {
+    Ok(NoiseRendezvousConnectBundle {
+        websocket_url,
+        environment_id: "environment".to_string(),
+        executor_registration_id: "registration".to_string(),
+        executor_public_key: NoiseChannelIdentity::generate()?.public_key(),
+        harness_key_authorization: "authorization".to_string(),
+    })
+}
+
+#[tokio::test]
+async fn initial_noise_connection_refreshes_bundle_after_unauthorized_handshake() -> Result<()> {
+    let unauthorized_listener = TcpListener::bind("127.0.0.1:0").await?;
+    let unauthorized_url = format!("ws://{}", unauthorized_listener.local_addr()?);
+    let unauthorized_server = tokio::spawn(async move {
+        let (mut socket, _) = unauthorized_listener.accept().await?;
+        let mut request = [0_u8; 4096];
+        let _ = socket.read(&mut request).await?;
+        socket
+            .write_all(
+                b"HTTP/1.1 401 Unauthorized\r\nContent-Length: 0\r\nConnection: close\r\n\r\n",
+            )
+            .await?;
+        socket.shutdown().await?;
+        anyhow::Ok(())
+    });
+    let accepted_listener = TcpListener::bind("127.0.0.1:0").await?;
+    let accepted_url = format!("ws://{}", accepted_listener.local_addr()?);
+    let accepted_server = tokio::spawn(async move {
+        let (socket, _) = accepted_listener.accept().await?;
+        let _websocket = accept_async(socket).await?;
+        anyhow::Ok(())
+    });
+    let sequence = Arc::new(SequenceNoiseConnectProvider::new(vec![
+        test_bundle(unauthorized_url.clone())?,
+        test_bundle(accepted_url.clone())?,
+    ]));
+    let provider: Arc<dyn NoiseRendezvousConnectProvider> = sequence.clone();
+    let identity = NoiseChannelIdentity::generate()?;
+
+    let _connection =
+        ExecServerClient::open_initial_noise_rendezvous_connection(&provider, &identity).await?;
+
+    assert_eq!(
+        sequence.returned_urls(),
+        vec![unauthorized_url, accepted_url]
+    );
+    unauthorized_server.await??;
+    accepted_server.await??;
+    Ok(())
+}