Add exec-server forwarding mode (#39249)

## What changed

- Add `codex exec-server forward --connect ws://HOST:PORT` to register an
  existing WebSocket exec-server as a remote environment.
- Open a separate destination WebSocket for each authenticated Noise relay
  stream and forward complete payloads unchanged in both directions.
- Preserve large messages through WebSocket fragmentation and coordinate
  disconnects, close acknowledgements, and pending output delivery.

## Testing

- Cover CLI validation, opaque payload forwarding, disconnect handling, and
  destination close behavior.
- Verify commands and large file transfers through the complete relay and
  forwarding path.

GitOrigin-RevId: 6e2dccff05b9e3571865611f4d432b0e2b4cdbac

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

@@ -0,0 +1,195 @@
+use std::time::Duration;
+
+use bytes::Bytes;
+use codex_http_client::HttpClientFactory;
+use codex_websocket_client::WebSocketConnector;
+use futures::Sink;
+use futures::SinkExt;
+use futures::StreamExt;
+use tokio::time::timeout;
+use tokio_tungstenite::tungstenite::Message;
+use tokio_tungstenite::tungstenite::client::IntoClientRequest;
+use tokio_tungstenite::tungstenite::protocol::WebSocketConfig;
+use tokio_tungstenite::tungstenite::protocol::frame::Frame;
+use tokio_tungstenite::tungstenite::protocol::frame::coding::Data;
+use tokio_tungstenite::tungstenite::protocol::frame::coding::OpCode;
+use tokio_util::task::AbortOnDropHandle;
+use tracing::warn;
+
+use crate::ExecServerError;
+use crate::ExecServerTelemetry;
+use crate::client_api::DEFAULT_REMOTE_EXEC_SERVER_CONNECT_TIMEOUT;
+use crate::noise_relay::message_framing::MAX_NOISE_JSONRPC_MESSAGE_LEN;
+use crate::noise_relay::message_framing::frame_message;
+use crate::noise_relay::stream_handler::NoiseOutboundMessage;
+use crate::noise_relay::stream_handler::NoiseStreamConnection;
+use crate::noise_relay::stream_handler::NoiseStreamHandler;
+use crate::telemetry::ConnectionTransport;
+
+// Existing exec-server listeners accept 64 MiB messages but only 16 MiB frames.
+const WEBSOCKET_FRAGMENT_LEN: usize = 8 * 1024 * 1024;
+const WEBSOCKET_CLOSE_TIMEOUT: Duration = Duration::from_secs(1);
+
+async fn send_websocket_message<S>(websocket: &mut S, mut payload: Bytes) -> Result<(), S::Error>
+where
+    S: Sink<Message> + Unpin,
+{
+    if payload.len() <= WEBSOCKET_FRAGMENT_LEN {
+        return websocket.send(Message::Binary(payload)).await;
+    }
+
+    let mut opcode = OpCode::Data(Data::Binary);
+    while !payload.is_empty() {
+        let chunk = payload.split_to(payload.len().min(WEBSOCKET_FRAGMENT_LEN));
+        let is_final = payload.is_empty();
+        websocket
+            .send(Message::Frame(Frame::message(chunk, opcode, is_final)))
+            .await?;
+        opcode = OpCode::Data(Data::Continue);
+    }
+    Ok(())
+}
+
+/// Copies authenticated remote messages to an independently owned executor.
+#[derive(Clone)]
+pub(crate) struct Forwarder {
+    websocket_url: String,
+    connector: WebSocketConnector,
+    telemetry: ExecServerTelemetry,
+}
+
+impl Forwarder {
+    pub(crate) fn new(
+        websocket_url: String,
+        http_client_factory: &HttpClientFactory,
+        telemetry: ExecServerTelemetry,
+    ) -> Result<Self, ExecServerError> {
+        let url = url::Url::parse(&websocket_url)
+            .map_err(|error| ExecServerError::WebSocketConfiguration(error.to_string()))?;
+        if !matches!(url.scheme(), "ws" | "wss") || url.host_str().is_none() {
+            return Err(ExecServerError::WebSocketConfiguration(
+                "forward destination must be a ws:// or wss:// URL".to_string(),
+            ));
+        }
+        let connector = WebSocketConnector::new(http_client_factory)
+            .map_err(|error| ExecServerError::WebSocketConfiguration(error.to_string()))?
+            .with_tcp_nodelay();
+        Ok(Self {
+            websocket_url,
+            connector,
+            telemetry,
+        })
+    }
+
+    pub(crate) async fn run_connection(self, mut remote: NoiseStreamConnection<Bytes, Bytes>) {
+        let mut writer_task = AbortOnDropHandle::new(remote.writer_task);
+        let _metrics = self
+            .telemetry
+            .connection_started(ConnectionTransport::Relay);
+        let connect = async {
+            let request = self.websocket_url.as_str().into_client_request()?;
+            self.connector
+                .connect(
+                    request,
+                    WebSocketConfig::default()
+                        .max_message_size(Some(MAX_NOISE_JSONRPC_MESSAGE_LEN))
+                        .max_frame_size(Some(MAX_NOISE_JSONRPC_MESSAGE_LEN)),
+                )
+                .await
+        };
+        let connected = tokio::select! {
+            biased;
+            _ = remote.disconnected_rx.wait_for(|disconnected| *disconnected) => None,
+            result = timeout(DEFAULT_REMOTE_EXEC_SERVER_CONNECT_TIMEOUT, connect) => {
+                match result {
+                    Ok(Ok((websocket, _))) => Some(websocket),
+                    Ok(Err(_)) => {
+                        warn!("failed to connect to forwarded exec-server");
+                        None
+                    }
+                    Err(_) => {
+                        warn!("timed out connecting to forwarded exec-server");
+                        None
+                    }
+                }
+            }
+        };
+        let drain_outgoing = if let Some(websocket) = connected {
+            let (mut destination_tx, mut destination_rx) = websocket.split();
+            let to_destination = async {
+                while let Some(payload) = remote.incoming_rx.recv().await {
+                    if send_websocket_message(&mut destination_tx, payload)
+                        .await
+                        .is_err()
+                    {
+                        break;
+                    }
+                }
+            };
+            let from_destination = async {
+                while let Some(Ok(message)) = destination_rx.next().await {
+                    let payload = match message {
+                        Message::Text(_) | Message::Binary(_) => message.into_data(),
+                        Message::Close(_) => return true,
+                        Message::Ping(_) | Message::Pong(_) | Message::Frame(_) => continue,
+                    };
+                    if remote.outgoing_tx.send(payload).await.is_err() {
+                        break;
+                    }
+                }
+                false
+            };
+            let (drain_outgoing, received_close) = tokio::select! {
+                _ = remote.disconnected_rx.wait_for(|disconnected| *disconnected) => (false, false),
+                _ = to_destination => (false, false),
+                received_close = from_destination => (true, received_close),
+            };
+            if received_close && let Ok(mut websocket) = destination_tx.reunite(destination_rx) {
+                // Reuniting drops any canceled application send before flushing
+                // Tungstenite's automatically queued Close acknowledgement.
+                tokio::select! {
+                    _ = timeout(WEBSOCKET_CLOSE_TIMEOUT, websocket.flush()) => {},
+                    _ = remote.disconnected_rx.wait_for(|disconnected| *disconnected) => {},
+                }
+            }
+            drain_outgoing
+        } else {
+            false
+        };
+        // Preserve messages received before the destination's Close. Never wait
+        // for a dead remote, and let owner shutdown interrupt a blocked drain.
+        drop(remote.outgoing_tx);
+        if drain_outgoing {
+            tokio::select! {
+                _ = &mut writer_task => return,
+                _ = remote.disconnected_rx.wait_for(|disconnected| *disconnected) => {},
+            }
+        }
+        writer_task.abort();
+        let _ = writer_task.await;
+    }
+}
+
+impl NoiseStreamHandler for Forwarder {
+    type Incoming = Bytes;
+    type Outgoing = Bytes;
+
+    fn decode(payload: Bytes) -> Result<Bytes, ExecServerError> {
+        Ok(payload)
+    }
+
+    fn encode(payload: Bytes) -> Result<NoiseOutboundMessage, ExecServerError> {
+        Ok(NoiseOutboundMessage {
+            framed: frame_message(&payload)?,
+            trace: None,
+        })
+    }
+
+    async fn run_connection(self, connection: NoiseStreamConnection<Bytes, Bytes>) {
+        Forwarder::run_connection(self, connection).await;
+    }
+}
+
+#[cfg(test)]
+#[path = "forward_tests.rs"]
+mod tests;

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

@@ -0,0 +1,87 @@
+use std::time::Duration;
+
+use anyhow::Context;
+use anyhow::Result;
+use bytes::Bytes;
+use codex_http_client::HttpClientFactory;
+use codex_http_client::OutboundProxyPolicy;
+use futures::SinkExt;
+use futures::StreamExt;
+use pretty_assertions::assert_eq;
+use tokio::io::AsyncReadExt;
+use tokio::net::TcpListener;
+use tokio::sync::mpsc;
+use tokio::sync::watch;
+use tokio::time::timeout;
+use tokio_tungstenite::accept_async;
+use tokio_tungstenite::tungstenite::Message;
+use tokio_tungstenite::tungstenite::protocol::CloseFrame;
+use tokio_tungstenite::tungstenite::protocol::frame::coding::CloseCode;
+
+use super::Forwarder;
+use crate::ExecServerTelemetry;
+use crate::noise_relay::stream_handler::NoiseStreamConnection;
+
+#[tokio::test]
+async fn transport_disconnect_cancels_an_unfinished_websocket_handshake() -> Result<()> {
+    let listener = TcpListener::bind("127.0.0.1:0").await?;
+    let forwarder = Forwarder::new(
+        format!("ws://{}", listener.local_addr()?),
+        &HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault),
+        ExecServerTelemetry::default(),
+    )?;
+    let (_incoming, incoming_rx) = mpsc::channel(1);
+    let (outgoing_tx, _outgoing) = mpsc::channel(1);
+    let (disconnected, disconnected_rx) = watch::channel(false);
+    let connection = NoiseStreamConnection {
+        incoming_rx,
+        outgoing_tx,
+        disconnected_rx,
+        writer_task: tokio::spawn(async {}),
+    };
+    let task = tokio::spawn(forwarder.run_connection(connection));
+    let deadline = Duration::from_secs(5);
+    let (mut socket, _) = timeout(deadline, listener.accept()).await??;
+    disconnected.send(true)?;
+    timeout(deadline, task).await??;
+    timeout(deadline, socket.read_to_end(&mut Vec::new())).await??;
+    Ok(())
+}
+
+#[tokio::test]
+async fn destination_close_is_acknowledged() -> Result<()> {
+    let listener = TcpListener::bind("127.0.0.1:0").await?;
+    let forwarder = Forwarder::new(
+        format!("ws://{}", listener.local_addr()?),
+        &HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault),
+        ExecServerTelemetry::default(),
+    )?;
+    let (_incoming, incoming_rx) = mpsc::channel(1);
+    let (outgoing_tx, mut outgoing) = mpsc::channel(1);
+    let (_disconnected, disconnected_rx) = watch::channel(false);
+    let task = tokio::spawn(forwarder.run_connection(NoiseStreamConnection {
+        incoming_rx,
+        outgoing_tx,
+        disconnected_rx,
+        writer_task: tokio::spawn(async {}),
+    }));
+    let deadline = Duration::from_secs(5);
+    let (socket, _) = timeout(deadline, listener.accept()).await??;
+    let mut socket = timeout(deadline, accept_async(socket)).await??;
+    let response = Bytes::from_static(b"final response");
+    socket.send(Message::Binary(response.clone())).await?;
+    let close = CloseFrame {
+        code: CloseCode::Normal,
+        reason: "finished".into(),
+    };
+    socket.close(Some(close.clone())).await?;
+
+    let reply = timeout(deadline, socket.next())
+        .await?
+        .context("destination should receive the Close acknowledgement")??;
+    assert_eq!(reply, Message::Close(Some(close)));
+    timeout(deadline, task).await??;
+    assert_eq!(outgoing.recv().await, Some(response));
+    assert_eq!(outgoing.recv().await, None);
+    Ok(())
+}

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

@@ -12,6 +12,7 @@ mod environment_provider;
 mod environment_registry;
 mod environment_toml;
 mod file_read;
+mod forward;
 mod fs_helper;
 mod fs_helper_main;
 mod fs_sandbox;
@@ -191,6 +192,7 @@ pub use protocol::WriteStatus;
 pub use regular_file::read_sensitive_file_to_string;
 pub use remote::RemoteEnvironmentConfig;
 pub use remote::run_remote_environment;
+pub use remote::run_remote_environment_forward_until_shutdown;
 pub use remote::run_remote_environment_until_shutdown;
 pub use resolved_capability::ResolvedSelectedCapabilityRoot;
 pub use resolved_capability::SelectedCapabilityRootsStatus;

codex-rs/exec-server/src/noise_relay/executor_stream.rs

@@ -186,3 +186,7 @@ pub(crate) fn spawn_noise_virtual_stream<H: NoiseStreamHandler>(
 #[cfg(test)]
 #[path = "executor_stream_tests.rs"]
 mod tests;
+
+#[cfg(test)]
+#[path = "forward_stream_tests.rs"]
+mod forward_tests;

codex-rs/exec-server/src/noise_relay/forward_stream_tests.rs

@@ -0,0 +1,111 @@
+use std::time::Duration;
+
+use anyhow::Context;
+use anyhow::Result;
+use bytes::Bytes;
+use codex_http_client::HttpClientFactory;
+use codex_http_client::OutboundProxyPolicy;
+use futures::SinkExt;
+use futures::StreamExt;
+use pretty_assertions::assert_eq;
+use tokio::net::TcpListener;
+use tokio::sync::mpsc;
+use tokio::time::timeout;
+use tokio_tungstenite::accept_async;
+use tokio_tungstenite::tungstenite::Message;
+
+use super::ClosedNoiseVirtualStream;
+use super::spawn_noise_virtual_stream;
+use crate::ExecServerTelemetry;
+use crate::forward::Forwarder;
+use crate::noise_channel::InitiatorHandshake;
+use crate::noise_channel::NoiseChannelIdentity;
+use crate::noise_channel::PendingResponderHandshake;
+use crate::noise_relay::message_framing::MessageDecoder;
+use crate::noise_relay::message_framing::frame_message;
+use crate::relay::decode_relay_message_frame;
+use crate::relay_proto::RelayData;
+
+const TEST_TIMEOUT: Duration = Duration::from_secs(5);
+
+#[tokio::test]
+async fn forwards_opaque_noise_payloads_and_drains_before_closing() -> Result<()> {
+    let executor_identity = NoiseChannelIdentity::generate()?;
+    let harness_identity = NoiseChannelIdentity::generate()?;
+    let prologue = b"forwarding-test";
+    let (initiator, request) = InitiatorHandshake::start(
+        &harness_identity,
+        &executor_identity.public_key(),
+        prologue,
+        b"authorization",
+    )?;
+    let pending = PendingResponderHandshake::read_request(&executor_identity, prologue, &request)?;
+    let (executor_transport, response) = pending.complete()?;
+    let mut harness_transport = initiator.finish(&response)?;
+
+    let destination = TcpListener::bind("127.0.0.1:0").await?;
+    let forwarder = Forwarder::new(
+        format!("ws://{}", destination.local_addr()?),
+        &HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault),
+        ExecServerTelemetry::default(),
+    )?;
+    // A single queued record makes the final response exercise writer draining.
+    let (physical_outgoing_tx, mut physical_outgoing_rx) = mpsc::channel(1);
+    let (closed_stream_tx, mut closed_stream_rx) = mpsc::channel(2);
+    let mut stream = spawn_noise_virtual_stream(
+        "stream-1".to_string(),
+        /*instance_id*/ 7,
+        forwarder,
+        physical_outgoing_tx,
+        closed_stream_tx,
+        executor_transport,
+    );
+    let (socket, _) = timeout(TEST_TIMEOUT, destination.accept()).await??;
+    let mut socket = timeout(TEST_TIMEOUT, accept_async(socket)).await??;
+
+    let request = Bytes::from_static(b"not JSON\x00\xff");
+    stream.receive_data(RelayData {
+        seq: 0,
+        segment_index: 0,
+        segment_count: 1,
+        payload: harness_transport.encrypt(&frame_message(&request)?)?,
+    })?;
+    let forwarded = timeout(TEST_TIMEOUT, socket.next())
+        .await?
+        .context("forwarded request")??;
+    assert_eq!(forwarded.into_data(), request);
+
+    let response = Bytes::from(vec![0xff; 128 * 1024]);
+    socket.send(Message::Binary(response.clone())).await?;
+    socket.close(None).await?;
+
+    let mut decoder = MessageDecoder::default();
+    let mut messages = Vec::new();
+    let mut next_seq = 0;
+    let closed = timeout(TEST_TIMEOUT, async {
+        loop {
+            tokio::select! {
+                biased;
+                Some(encoded) = physical_outgoing_rx.recv() => {
+                    let frame = decode_relay_message_frame(&encoded)?;
+                    assert_eq!(frame.stream_id, "stream-1");
+                    let data = frame.into_data()?;
+                    assert_eq!(data.seq, next_seq);
+                    next_seq += 1;
+                    messages.extend(decoder.push(&harness_transport.decrypt(&data.payload)?)?);
+                }
+                closed = closed_stream_rx.recv() => {
+                    break closed.context("stream close notification");
+                }
+            }
+        }
+    })
+    .await??;
+    assert!(next_seq > 1);
+    assert_eq!(messages, vec![response]);
+    assert!(matches!(
+        closed,
+        ClosedNoiseVirtualStream { stream_id, instance_id: 7 } if stream_id == "stream-1"
+    ));
+    Ok(())
+}

codex-rs/exec-server/src/noise_relay/message_framing.rs

@@ -6,7 +6,7 @@ use codex_exec_server_protocol::JSONRPCMessage;
 use crate::ExecServerError;
 
 const LENGTH_PREFIX_BYTES: usize = size_of::<u32>();
-const MAX_NOISE_JSONRPC_MESSAGE_LEN: usize = 64 * 1024 * 1024;
+pub(crate) const MAX_NOISE_JSONRPC_MESSAGE_LEN: usize = 64 * 1024 * 1024;
 pub(crate) const NOISE_RECORD_PLAINTEXT_LEN: usize = 60 * 1024;
 
 /// Serialize one JSON-RPC message into the encrypted record byte stream.
@@ -23,6 +23,14 @@ pub(crate) fn frame_jsonrpc_message(message: &JSONRPCMessage) -> Result<Vec<u8>,
     Ok(framed)
 }
 
+pub(crate) fn frame_message(message: &[u8]) -> Result<Vec<u8>, ExecServerError> {
+    let prefix = message_length_prefix(message.len())?;
+    let mut framed = Vec::with_capacity(LENGTH_PREFIX_BYTES + message.len());
+    framed.extend_from_slice(&prefix);
+    framed.extend_from_slice(message);
+    Ok(framed)
+}
+
 fn message_length_prefix(message_len: usize) -> Result<[u8; LENGTH_PREFIX_BYTES], ExecServerError> {
     if message_len == 0 || message_len > MAX_NOISE_JSONRPC_MESSAGE_LEN {
         return Err(ExecServerError::Protocol(

codex-rs/exec-server/src/noise_relay/mod.rs

@@ -1,6 +1,6 @@
 pub(crate) mod executor_stream;
 mod harness;
-mod message_framing;
+pub(crate) mod message_framing;
 mod ordered_ciphertext;
 pub(crate) mod stream_handler;
 

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

@@ -39,7 +39,9 @@ use crate::NoiseChannelPublicKey;
 use crate::NoiseRendezvousConnectBundle;
 use crate::NoiseRendezvousConnectProvider;
 use crate::client_api::DEFAULT_REMOTE_EXEC_SERVER_CONNECT_TIMEOUT;
+use crate::forward::Forwarder;
 use crate::noise_relay::noise_relay_websocket_config;
+use crate::noise_relay::stream_handler::NoiseStreamHandler;
 use crate::relay::HarnessKeyValidator;
 use crate::relay::run_multiplexed_environment;
 use crate::server::ConnectionProcessor;
@@ -550,13 +552,6 @@ pub async fn run_remote_environment_until_shutdown<F>(
 where
     F: std::future::Future<Output = ()>,
 {
-    ensure_rustls_crypto_provider();
-    let client = EnvironmentRegistryClient::new_with_telemetry(
-        config.base_url.clone(),
-        config.auth_provider.clone(),
-        config.telemetry.clone(),
-        config.http_client_factory.clone(),
-    )?;
     let processor = ConnectionProcessor::new_with_telemetry(
         runtime_paths,
         config.telemetry.clone(),
@@ -564,22 +559,56 @@ where
         config.request_dispatch_mode,
     );
 
-    let result = {
-        let run = run_remote_environment_connections(config, client, processor.clone());
-        tokio::pin!(run, shutdown);
-        tokio::select! {
-            result = &mut run => result,
-            _ = &mut shutdown => Ok(()),
-        }
-    };
+    let result = run_remote_transport(config, processor.clone(), shutdown).await;
     processor.shutdown().await;
     result
 }
 
-async fn run_remote_environment_connections(
+/// Register a remote environment backed by an independently owned WebSocket executor.
+pub async fn run_remote_environment_forward_until_shutdown<F>(
+    config: RemoteEnvironmentConfig,
+    websocket_url: String,
+    shutdown: F,
+) -> Result<(), ExecServerError>
+where
+    F: std::future::Future<Output = ()>,
+{
+    let forwarder = Forwarder::new(
+        websocket_url,
+        &config.http_client_factory,
+        config.telemetry.clone(),
+    )?;
+    run_remote_transport(config, forwarder, shutdown).await
+}
+
+async fn run_remote_transport<F, H>(
+    config: RemoteEnvironmentConfig,
+    handler: H,
+    shutdown: F,
+) -> Result<(), ExecServerError>
+where
+    F: std::future::Future<Output = ()>,
+    H: NoiseStreamHandler,
+{
+    ensure_rustls_crypto_provider();
+    let client = EnvironmentRegistryClient::new_with_telemetry(
+        config.base_url.clone(),
+        config.auth_provider.clone(),
+        config.telemetry.clone(),
+        config.http_client_factory.clone(),
+    )?;
+    let run = run_remote_environment_connections(config, client, handler);
+    tokio::pin!(run, shutdown);
+    tokio::select! {
+        result = &mut run => result,
+        _ = &mut shutdown => Ok(()),
+    }
+}
+
+async fn run_remote_environment_connections<H: NoiseStreamHandler>(
     config: RemoteEnvironmentConfig,
     client: EnvironmentRegistryClient,
-    processor: ConnectionProcessor,
+    handler: H,
 ) -> Result<(), ExecServerError> {
     let identity = NoiseChannelIdentity::generate().map_err(|error| {
         ExecServerError::Protocol(format!("failed to generate Noise relay identity: {error}"))
@@ -607,7 +636,7 @@ async fn run_remote_environment_connections(
                 );
                 let disconnect_reason = run_multiplexed_environment(
                     websocket,
-                    processor.clone(),
+                    handler.clone(),
                     response.environment_id.clone(),
                     executor_registration_id.clone(),
                     identity.clone(),