feat(exec-server): use protobuf relay frames (#22343)

## Why Remote exec-server now needs one executor websocket to serve multiple harness JSON-RPC sessions. Rendezvous routes by `stream_id`, and the exec-server side needs to use the same stable relay frame contract instead of a hand-rolled JSON shape. The relay protocol also needs to make ownership boundaries clear: harness and executor endpoints own sequencing, acks, retries, duplicate suppression, segmentation, and reassembly; rendezvous only routes frames. ## What Changed - Add the checked-in `codex.exec_server.relay.v1.RelayMessageFrame` proto plus generated prost bindings for `codex-exec-server`. - Encode remote harness/executor relay traffic as binary protobuf websocket frames while keeping local websocket JSON-RPC unchanged. - Demux executor-side relay streams into independent `ConnectionProcessor` sessions keyed by `stream_id`. - Add a programmatic `RemoteExecutorConfig::with_bearer_token(...)` constructor for non-CLI callers and integration tests. - Add an integration test that starts the remote executor against a fake registry/rendezvous websocket and verifies two virtual streams share one executor websocket without cross-talk, including per-stream reset behavior. - Document the remote relay envelope, sequence ranges, `ack`/`ack_bits`, and endpoint responsibilities in `exec-server/README.md`. ## Verification - `cargo test -p codex-exec-server --test relay multiplexed_remote_executor_routes_independent_virtual_streams -- --exact` - `cargo test -p codex-exec-server --test relay` - `cargo test -p codex-exec-server` passed outside the sandbox. The sandboxed run hit macOS `sandbox-exec: sandbox_apply: Operation not permitted` in filesystem sandbox tests.
2026-07-01 00:31:56 +08:00 · 2026-05-12 16:50:45 -07:00
parent 6dc3b3d7c8
commit ac466c0dbd
12 changed files with 993 additions and 17 deletions
@@ -15,6 +15,7 @@ use crate::client_api::RemoteExecServerConnectArgs;
 use crate::client_api::StdioExecServerCommand;
 use crate::client_api::StdioExecServerConnectArgs;
 use crate::connection::JsonRpcConnection;
+use crate::relay::harness_connection_from_websocket;

 const ENVIRONMENT_CLIENT_NAME: &str = "codex-environment";

@@ -69,14 +70,13 @@ impl ExecServerClient {
                source,
            })?;

-        Self::connect(
-            JsonRpcConnection::from_websocket(
-                stream,
-                format!("exec-server websocket {websocket_url}"),
-            ),
-            args.into(),
-        )
-        .await
+        let connection_label = format!("exec-server websocket {websocket_url}");
+        let connection = if is_rendezvous_harness_url(&websocket_url) {
+            harness_connection_from_websocket(stream, connection_label)
+        } else {
+            JsonRpcConnection::from_websocket(stream, connection_label)
+        };
+        Self::connect(connection, args.into()).await
    }

    pub(crate) async fn connect_stdio_command(
@@ -120,6 +120,16 @@ impl ExecServerClient {
    }
 }

+fn is_rendezvous_harness_url(websocket_url: &str) -> bool {
+    let Some((_path, query)) = websocket_url.split_once('?') else {
+        return false;
+    };
+    query
+        .split('&')
+        .filter_map(|pair| pair.split_once('='))
+        .any(|(key, value)| key == "role" && value == "harness")
+}
+
 fn stdio_command_process(stdio_command: &StdioExecServerCommand) -> Command {
    let mut command = Command::new(&stdio_command.program);
    command.args(&stdio_command.args);
@@ -13,6 +13,8 @@ mod local_process;
 mod process;
 mod process_id;
 mod protocol;
+mod relay;
+mod relay_proto;
 mod remote;
 mod remote_file_system;
 mod remote_process;
@@ -0,0 +1,37 @@
+syntax = "proto3";
+
+package codex.exec_server.relay.v1;
+
+message RelayMessageFrame {
+  uint32 version = 1;
+  string stream_id = 2;
+  uint32 ack = 3;
+  uint32 ack_bits = 4;
+
+  oneof body {
+    RelayData data = 5;
+    RelayAck ack_frame = 6;
+    RelayResume resume = 7;
+    RelayReset reset = 8;
+    RelayHeartbeat heartbeat = 9;
+  }
+}
+
+message RelayData {
+  uint32 seq = 1;
+  uint32 segment_index = 2;
+  uint32 segment_count = 3;
+  bytes payload = 4;
+}
+
+message RelayAck {}
+
+message RelayResume {
+  uint32 next_seq = 1;
+}
+
+message RelayReset {
+  string reason = 1;
+}
+
+message RelayHeartbeat {}
@@ -0,0 +1,54 @@
+// This file is @generated by prost-build.
+#[derive(Clone, PartialEq, ::prost::Message)]
+pub struct RelayMessageFrame {
+    #[prost(uint32, tag = "1")]
+    pub version: u32,
+    #[prost(string, tag = "2")]
+    pub stream_id: ::prost::alloc::string::String,
+    #[prost(uint32, tag = "3")]
+    pub ack: u32,
+    #[prost(uint32, tag = "4")]
+    pub ack_bits: u32,
+    #[prost(oneof = "relay_message_frame::Body", tags = "5, 6, 7, 8, 9")]
+    pub body: ::core::option::Option<relay_message_frame::Body>,
+}
+pub mod relay_message_frame {
+    #[derive(Clone, PartialEq, ::prost::Oneof)]
+    pub enum Body {
+        #[prost(message, tag = "5")]
+        Data(super::RelayData),
+        #[prost(message, tag = "6")]
+        AckFrame(super::RelayAck),
+        #[prost(message, tag = "7")]
+        Resume(super::RelayResume),
+        #[prost(message, tag = "8")]
+        Reset(super::RelayReset),
+        #[prost(message, tag = "9")]
+        Heartbeat(super::RelayHeartbeat),
+    }
+}
+#[derive(Clone, PartialEq, ::prost::Message)]
+pub struct RelayData {
+    #[prost(uint32, tag = "1")]
+    pub seq: u32,
+    #[prost(uint32, tag = "2")]
+    pub segment_index: u32,
+    #[prost(uint32, tag = "3")]
+    pub segment_count: u32,
+    #[prost(bytes = "vec", tag = "4")]
+    pub payload: ::prost::alloc::vec::Vec<u8>,
+}
+#[derive(Clone, PartialEq, Eq, Hash, ::prost::Message)]
+pub struct RelayAck {}
+#[derive(Clone, Copy, PartialEq, Eq, Hash, ::prost::Message)]
+pub struct RelayResume {
+    #[prost(uint32, tag = "1")]
+    pub next_seq: u32,
+}
+#[derive(Clone, PartialEq, Eq, Hash, ::prost::Message)]
+pub struct RelayReset {
+    #[prost(string, tag = "1")]
+    pub reason: ::prost::alloc::string::String,
+}
+#[derive(Clone, PartialEq, Eq, Hash, ::prost::Message)]
+pub struct RelayHeartbeat {}
@@ -0,0 +1,455 @@
+use std::collections::HashMap;
+
+use codex_app_server_protocol::JSONRPCMessage;
+use futures::SinkExt;
+use futures::StreamExt;
+use prost::Message as ProstMessage;
+use tokio::io::AsyncRead;
+use tokio::io::AsyncWrite;
+use tokio::sync::mpsc;
+use tokio::sync::watch;
+use tokio_tungstenite::WebSocketStream;
+use tokio_tungstenite::tungstenite::Message;
+use tracing::debug;
+use tracing::warn;
+use uuid::Uuid;
+
+use crate::ExecServerError;
+use crate::connection::CHANNEL_CAPACITY;
+use crate::connection::JsonRpcConnection;
+use crate::connection::JsonRpcConnectionEvent;
+use crate::connection::JsonRpcTransport;
+use crate::relay_proto::RelayData;
+use crate::relay_proto::RelayMessageFrame;
+use crate::relay_proto::RelayResume;
+use crate::relay_proto::relay_message_frame;
+use crate::server::ConnectionProcessor;
+
+const RELAY_MESSAGE_FRAME_VERSION: u32 = 1;
+
+#[derive(Debug, Clone, Copy, Eq, PartialEq)]
+enum RelayFrameBodyKind {
+    Data,
+    Ack,
+    Resume,
+    Reset,
+    Heartbeat,
+}
+
+impl RelayMessageFrame {
+    fn data(stream_id: String, seq: u32, payload: Vec<u8>) -> Self {
+        Self {
+            version: RELAY_MESSAGE_FRAME_VERSION,
+            stream_id,
+            ack: 0,
+            ack_bits: 0,
+            body: Some(relay_message_frame::Body::Data(RelayData {
+                seq,
+                segment_index: 0,
+                segment_count: 1,
+                payload,
+            })),
+        }
+    }
+
+    fn resume(stream_id: String) -> Self {
+        Self {
+            version: RELAY_MESSAGE_FRAME_VERSION,
+            stream_id,
+            ack: 0,
+            ack_bits: 0,
+            body: Some(relay_message_frame::Body::Resume(RelayResume {
+                next_seq: 0,
+            })),
+        }
+    }
+
+    fn validate(&self) -> Result<RelayFrameBodyKind, ExecServerError> {
+        if self.version != RELAY_MESSAGE_FRAME_VERSION {
+            return Err(ExecServerError::Protocol(format!(
+                "unsupported relay message frame version {}",
+                self.version
+            )));
+        }
+        if self.stream_id.trim().is_empty() {
+            return Err(ExecServerError::Protocol(
+                "relay message frame is missing stream_id".to_string(),
+            ));
+        }
+        match self.body.as_ref() {
+            Some(relay_message_frame::Body::Data(data)) => {
+                if data.segment_index != 0 || data.segment_count != 1 || data.payload.is_empty() {
+                    return Err(ExecServerError::Protocol(
+                        "relay data message frame is missing required fields".to_string(),
+                    ));
+                }
+                Ok(RelayFrameBodyKind::Data)
+            }
+            Some(relay_message_frame::Body::AckFrame(_)) => Ok(RelayFrameBodyKind::Ack),
+            Some(relay_message_frame::Body::Resume(_)) => Ok(RelayFrameBodyKind::Resume),
+            Some(relay_message_frame::Body::Reset(reset)) => {
+                if reset.reason.is_empty() {
+                    return Err(ExecServerError::Protocol(
+                        "relay reset message frame is missing reason".to_string(),
+                    ));
+                }
+                Ok(RelayFrameBodyKind::Reset)
+            }
+            Some(relay_message_frame::Body::Heartbeat(_)) => Ok(RelayFrameBodyKind::Heartbeat),
+            None => Err(ExecServerError::Protocol(
+                "relay message frame is missing body".to_string(),
+            )),
+        }
+    }
+
+    fn into_jsonrpc_message(self) -> Result<JSONRPCMessage, ExecServerError> {
+        let kind = self.validate()?;
+        if kind != RelayFrameBodyKind::Data {
+            return Err(ExecServerError::Protocol(
+                "expected relay data message frame".to_string(),
+            ));
+        }
+        let payload = match self.body {
+            Some(relay_message_frame::Body::Data(data)) => data.payload,
+            _ => Vec::new(),
+        };
+        serde_json::from_slice(&payload).map_err(ExecServerError::Json)
+    }
+
+    fn into_reset_reason(self) -> Option<String> {
+        match self.body {
+            Some(relay_message_frame::Body::Reset(reset)) if !reset.reason.is_empty() => {
+                Some(reset.reason)
+            }
+            _ => None,
+        }
+    }
+}
+
+fn encode_relay_message_frame(frame: &RelayMessageFrame) -> Vec<u8> {
+    frame.encode_to_vec()
+}
+
+fn decode_relay_message_frame(payload: &[u8]) -> Result<RelayMessageFrame, ExecServerError> {
+    RelayMessageFrame::decode(payload)
+        .map_err(|err| ExecServerError::Protocol(format!("invalid relay message frame: {err}")))
+}
+
+fn jsonrpc_payload(message: &JSONRPCMessage) -> Result<Vec<u8>, ExecServerError> {
+    serde_json::to_vec(message).map_err(ExecServerError::Json)
+}
+
+pub(crate) fn harness_connection_from_websocket<S>(
+    stream: WebSocketStream<S>,
+    connection_label: String,
+) -> JsonRpcConnection
+where
+    S: AsyncRead + AsyncWrite + Unpin + Send + 'static,
+{
+    let stream_id = Uuid::new_v4().to_string();
+    let (mut websocket_writer, mut websocket_reader) = stream.split();
+    let (outgoing_tx, mut outgoing_rx) = mpsc::channel(CHANNEL_CAPACITY);
+    let (incoming_tx, incoming_rx) = mpsc::channel(CHANNEL_CAPACITY);
+    let (disconnected_tx, disconnected_rx) = watch::channel(false);
+
+    let reader_label = connection_label;
+    let reader_stream_id = stream_id.clone();
+    let incoming_tx_for_reader = incoming_tx;
+    let disconnected_tx_for_reader = disconnected_tx.clone();
+    let reader_task = tokio::spawn(async move {
+        loop {
+            match websocket_reader.next().await {
+                Some(Ok(Message::Binary(payload))) => {
+                    let frame = match decode_relay_message_frame(payload.as_ref()) {
+                        Ok(frame) => frame,
+                        Err(err) => {
+                            let _ = incoming_tx_for_reader
+                                .send(JsonRpcConnectionEvent::MalformedMessage {
+                                    reason: format!(
+                                        "failed to parse relay message frame from {reader_label}: {err}"
+                                    ),
+                                })
+                                .await;
+                            continue;
+                        }
+                    };
+                    if frame.stream_id != reader_stream_id {
+                        continue;
+                    }
+                    let kind = match frame.validate() {
+                        Ok(kind) => kind,
+                        Err(err) => {
+                            let _ = incoming_tx_for_reader
+                                .send(JsonRpcConnectionEvent::MalformedMessage {
+                                    reason: err.to_string(),
+                                })
+                                .await;
+                            continue;
+                        }
+                    };
+                    match kind {
+                        RelayFrameBodyKind::Data => match frame.into_jsonrpc_message() {
+                            Ok(message) => {
+                                if incoming_tx_for_reader
+                                    .send(JsonRpcConnectionEvent::Message(message))
+                                    .await
+                                    .is_err()
+                                {
+                                    break;
+                                }
+                            }
+                            Err(err) => {
+                                let _ = incoming_tx_for_reader
+                                    .send(JsonRpcConnectionEvent::MalformedMessage {
+                                        reason: err.to_string(),
+                                    })
+                                    .await;
+                            }
+                        },
+                        RelayFrameBodyKind::Reset => {
+                            let _ = disconnected_tx_for_reader.send(true);
+                            let _ = incoming_tx_for_reader
+                                .send(JsonRpcConnectionEvent::Disconnected {
+                                    reason: frame.into_reset_reason(),
+                                })
+                                .await;
+                            break;
+                        }
+                        RelayFrameBodyKind::Ack
+                        | RelayFrameBodyKind::Resume
+                        | RelayFrameBodyKind::Heartbeat => {}
+                    }
+                }
+                Some(Ok(Message::Close(_))) | None => {
+                    let _ = disconnected_tx_for_reader.send(true);
+                    let _ = incoming_tx_for_reader
+                        .send(JsonRpcConnectionEvent::Disconnected { reason: None })
+                        .await;
+                    break;
+                }
+                Some(Ok(Message::Ping(_) | Message::Pong(_) | Message::Frame(_))) => {}
+                Some(Ok(Message::Text(_))) => {
+                    let _ = incoming_tx_for_reader
+                        .send(JsonRpcConnectionEvent::MalformedMessage {
+                            reason: "relay exec-server transport expects binary protobuf frames"
+                                .to_string(),
+                        })
+                        .await;
+                }
+                Some(Err(err)) => {
+                    let _ = disconnected_tx_for_reader.send(true);
+                    let _ = incoming_tx_for_reader
+                        .send(JsonRpcConnectionEvent::Disconnected {
+                            reason: Some(format!(
+                                "failed to read relay websocket frame from {reader_label}: {err}"
+                            )),
+                        })
+                        .await;
+                    break;
+                }
+            }
+        }
+    });
+
+    let writer_task = tokio::spawn(async move {
+        let resume = RelayMessageFrame::resume(stream_id.clone());
+        if websocket_writer
+            .send(Message::Binary(encode_relay_message_frame(&resume).into()))
+            .await
+            .is_err()
+        {
+            let _ = disconnected_tx.send(true);
+            return;
+        }
+
+        let mut next_seq = 0u32;
+        while let Some(message) = outgoing_rx.recv().await {
+            let payload = match jsonrpc_payload(&message) {
+                Ok(payload) => payload,
+                Err(err) => {
+                    warn!("failed to serialize JSON-RPC payload for relay transport: {err}");
+                    break;
+                }
+            };
+            let frame = RelayMessageFrame::data(stream_id.clone(), next_seq, payload);
+            next_seq = next_seq.wrapping_add(1);
+            if websocket_writer
+                .send(Message::Binary(encode_relay_message_frame(&frame).into()))
+                .await
+                .is_err()
+            {
+                let _ = disconnected_tx.send(true);
+                break;
+            }
+        }
+    });
+
+    JsonRpcConnection {
+        outgoing_tx,
+        incoming_rx,
+        disconnected_rx,
+        task_handles: vec![reader_task, writer_task],
+        transport: JsonRpcTransport::Plain,
+    }
+}
+
+pub(crate) async fn run_multiplexed_executor<S>(
+    stream: WebSocketStream<S>,
+    processor: ConnectionProcessor,
+) where
+    S: AsyncRead + AsyncWrite + Unpin + Send + 'static,
+{
+    let (mut websocket_writer, mut websocket_reader) = stream.split();
+    let (physical_outgoing_tx, mut physical_outgoing_rx) =
+        mpsc::channel::<Vec<u8>>(CHANNEL_CAPACITY);
+    let writer_task = tokio::spawn(async move {
+        while let Some(encoded) = physical_outgoing_rx.recv().await {
+            if websocket_writer
+                .send(Message::Binary(encoded.into()))
+                .await
+                .is_err()
+            {
+                break;
+            }
+        }
+    });
+
+    let mut streams: HashMap<String, VirtualStream> = HashMap::new();
+    loop {
+        let frame = match websocket_reader.next().await {
+            Some(Ok(Message::Binary(payload))) => {
+                match decode_relay_message_frame(payload.as_ref()) {
+                    Ok(frame) => frame,
+                    Err(err) => {
+                        warn!("dropping malformed relay message frame from harness: {err}");
+                        continue;
+                    }
+                }
+            }
+            Some(Ok(Message::Close(_))) | None => break,
+            Some(Ok(Message::Ping(_) | Message::Pong(_) | Message::Frame(_))) => continue,
+            Some(Ok(Message::Text(_))) => {
+                warn!("dropping non-binary relay message frame from harness");
+                continue;
+            }
+            Some(Err(err)) => {
+                debug!("multiplexed executor websocket read failed: {err}");
+                break;
+            }
+        };
+
+        let kind = match frame.validate() {
+            Ok(kind) => kind,
+            Err(err) => {
+                warn!("dropping invalid relay message frame: {err}");
+                continue;
+            }
+        };
+
+        match kind {
+            RelayFrameBodyKind::Data => {
+                let stream_id = frame.stream_id.clone();
+                let message = match frame.into_jsonrpc_message() {
+                    Ok(message) => message,
+                    Err(err) => {
+                        warn!("dropping malformed relay data message frame: {err}");
+                        continue;
+                    }
+                };
+                let stream = streams.entry(stream_id.clone()).or_insert_with(|| {
+                    spawn_virtual_stream(
+                        stream_id.clone(),
+                        processor.clone(),
+                        physical_outgoing_tx.clone(),
+                    )
+                });
+                if stream
+                    .incoming_tx
+                    .send(JsonRpcConnectionEvent::Message(message))
+                    .await
+                    .is_err()
+                {
+                    streams.remove(&stream_id);
+                }
+            }
+            RelayFrameBodyKind::Reset => {
+                if let Some(stream) = streams.remove(&frame.stream_id) {
+                    stream.disconnect(frame.into_reset_reason()).await;
+                }
+            }
+            RelayFrameBodyKind::Ack
+            | RelayFrameBodyKind::Resume
+            | RelayFrameBodyKind::Heartbeat => {}
+        }
+    }
+
+    for (_stream_id, stream) in streams {
+        stream.disconnect(/*reason*/ None).await;
+    }
+    drop(physical_outgoing_tx);
+    let _ = writer_task.await;
+}
+
+struct VirtualStream {
+    incoming_tx: mpsc::Sender<JsonRpcConnectionEvent>,
+    disconnected_tx: watch::Sender<bool>,
+}
+
+impl VirtualStream {
+    async fn disconnect(self, reason: Option<String>) {
+        let _ = self.disconnected_tx.send(true);
+        let _ = self
+            .incoming_tx
+            .send(JsonRpcConnectionEvent::Disconnected { reason })
+            .await;
+    }
+}
+
+fn spawn_virtual_stream(
+    stream_id: String,
+    processor: ConnectionProcessor,
+    physical_outgoing_tx: mpsc::Sender<Vec<u8>>,
+) -> VirtualStream {
+    let (json_outgoing_tx, mut json_outgoing_rx) = mpsc::channel(CHANNEL_CAPACITY);
+    let (incoming_tx, incoming_rx) = mpsc::channel(CHANNEL_CAPACITY);
+    let (disconnected_tx, disconnected_rx) = watch::channel(false);
+
+    let writer_stream_id = stream_id;
+    let writer_task = tokio::spawn(async move {
+        let mut next_seq = 0u32;
+        while let Some(message) = json_outgoing_rx.recv().await {
+            let payload = match jsonrpc_payload(&message) {
+                Ok(payload) => payload,
+                Err(err) => {
+                    warn!("failed to serialize virtual stream JSON-RPC payload: {err}");
+                    break;
+                }
+            };
+            let frame = RelayMessageFrame::data(writer_stream_id.clone(), next_seq, payload);
+            next_seq = next_seq.wrapping_add(1);
+            if physical_outgoing_tx
+                .send(encode_relay_message_frame(&frame))
+                .await
+                .is_err()
+            {
+                break;
+            }
+        }
+    });
+
+    let connection = JsonRpcConnection {
+        outgoing_tx: json_outgoing_tx,
+        incoming_rx,
+        disconnected_rx,
+        task_handles: vec![writer_task],
+        transport: JsonRpcTransport::Plain,
+    };
+    tokio::spawn(async move {
+        processor.run_connection(connection).await;
+    });
+
+    VirtualStream {
+        incoming_tx,
+        disconnected_tx,
+    }
+}
@@ -0,0 +1,7 @@
+#[path = "proto/codex.exec_server.relay.v1.rs"]
+mod generated;
+
+pub(crate) use generated::RelayData;
+pub(crate) use generated::RelayMessageFrame;
+pub(crate) use generated::RelayResume;
+pub(crate) use generated::relay_message_frame;
@@ -11,7 +11,7 @@ use codex_utils_rustls_provider::ensure_rustls_crypto_provider;

 use crate::ExecServerError;
 use crate::ExecServerRuntimePaths;
-use crate::connection::JsonRpcConnection;
+use crate::relay::run_multiplexed_executor;
 use crate::server::ConnectionProcessor;

 pub const CODEX_EXEC_SERVER_REMOTE_BEARER_TOKEN_ENV_VAR: &str =
@@ -113,7 +113,7 @@ impl RemoteExecutorConfig {
        Self::with_bearer_token(base_url, executor_id, read_remote_bearer_token_from_env()?)
    }

-    fn with_bearer_token(
+    pub fn with_bearer_token(
        base_url: String,
        executor_id: String,
        bearer_token: String,
@@ -150,12 +150,7 @@ pub async fn run_remote_executor(
        match connect_async(response.url.as_str()).await {
            Ok((websocket, _)) => {
                backoff = Duration::from_secs(1);
-                processor
-                    .run_connection(JsonRpcConnection::from_websocket(
-                        websocket,
-                        "remote exec-server websocket".to_string(),
-                    ))
-                    .await;
+                run_multiplexed_executor(websocket, processor.clone()).await;
            }
            Err(err) => {
                warn!("failed to connect remote exec-server websocket: {err}");