Fix exec-server compile breakage and websocket test helpers

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

@@ -1,5 +1,5 @@
 use clap::Parser;
-use codex_exec_server::ExecServerTransport;
+use codex_exec_server::DEFAULT_LISTEN_URL;
 
 #[derive(Debug, Parser)]
 struct ExecServerArgs {
@@ -8,15 +8,15 @@ struct ExecServerArgs {
     #[arg(
         long = "listen",
         value_name = "URL",
-        default_value = ExecServerTransport::DEFAULT_LISTEN_URL
+        default_value = DEFAULT_LISTEN_URL
     )]
-    listen: ExecServerTransport,
+    listen: String,
 }
 
 #[tokio::main]
 async fn main() {
     let args = ExecServerArgs::parse();
-    if let Err(err) = codex_exec_server::run_main_with_transport(args.listen).await {
+    if let Err(err) = codex_exec_server::run_main_with_transport(&args.listen).await {
         eprintln!("{err}");
         std::process::exit(1);
     }

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

@@ -42,7 +42,8 @@ pub use protocol::TerminateParams;
 pub use protocol::TerminateResponse;
 pub use protocol::WriteParams;
 pub use protocol::WriteResponse;
-pub use server::ExecServerTransport;
+pub use server::DEFAULT_LISTEN_URL;
 pub use server::ExecServerTransportParseError;
 pub use server::run_main;
+pub use server::run_main_with_listen_url;
 pub use server::run_main_with_transport;

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

@@ -4,7 +4,6 @@ use std::pin::Pin;
 use std::sync::Arc;
 use std::sync::atomic::AtomicI64;
 use std::sync::atomic::Ordering;
-use std::pin::Pin;
 
 use codex_app_server_protocol::JSONRPCError;
 use codex_app_server_protocol::JSONRPCErrorError;
@@ -197,6 +196,10 @@ impl RpcClient {
                             break;
                         }
                     }
+                    JsonRpcConnectionEvent::MalformedMessage { reason } => {
+                        warn!("JSON-RPC client closing after malformed message: {reason}");
+                        break;
+                    }
                     JsonRpcConnectionEvent::Disconnected { reason } => {
                         let _ = event_tx.send(RpcClientEvent::Disconnected { reason }).await;
                         drain_pending(&pending_for_reader).await;
@@ -442,217 +445,6 @@ async fn drain_pending(pending: &Mutex<HashMap<RequestId, PendingRequest>>) {
     }
 }
 
-#[derive(Debug)]
-pub(crate) enum RpcServerOutboundMessage {
-    Response {
-        request_id: RequestId,
-        result: Value,
-    },
-    Error {
-        request_id: RequestId,
-        error: JSONRPCErrorError,
-    },
-    Notification(JSONRPCNotification),
-}
-
-impl RpcServerOutboundMessage {
-    fn response(request_id: RequestId, result: Result<Value, JSONRPCErrorError>) -> Self {
-        match result {
-            Ok(result) => Self::Response {
-                request_id,
-                result,
-            },
-            Err(error) => Self::Error {
-                request_id,
-                error,
-            },
-        }
-    }
-}
-
-pub(crate) fn invalid_request(message: String) -> JSONRPCErrorError {
-    JSONRPCErrorError {
-        code: -32600,
-        data: None,
-        message,
-    }
-}
-
-pub(crate) fn invalid_params(message: String) -> JSONRPCErrorError {
-    JSONRPCErrorError {
-        code: -32602,
-        data: None,
-        message,
-    }
-}
-
-pub(crate) fn method_not_found(message: String) -> JSONRPCErrorError {
-    JSONRPCErrorError {
-        code: -32601,
-        data: None,
-        message,
-    }
-}
-
-pub(crate) fn internal_error(message: String) -> JSONRPCErrorError {
-    JSONRPCErrorError {
-        code: -32603,
-        data: None,
-        message,
-    }
-}
-
-pub(crate) fn encode_server_message(
-    message: RpcServerOutboundMessage,
-) -> Result<JSONRPCMessage, serde_json::Error> {
-    Ok(match message {
-        RpcServerOutboundMessage::Response { request_id, result } => {
-            JSONRPCMessage::Response(JSONRPCResponse { id: request_id, result })
-        }
-        RpcServerOutboundMessage::Error { request_id, error } => {
-            JSONRPCMessage::Error(JSONRPCError { id: request_id, error })
-        }
-        RpcServerOutboundMessage::Notification(notification) => {
-            JSONRPCMessage::Notification(notification)
-        }
-    })
-}
-
-#[derive(Clone)]
-pub(crate) struct RpcNotificationSender {
-    tx: mpsc::Sender<RpcServerOutboundMessage>,
-}
-
-impl RpcNotificationSender {
-    pub(crate) fn new(tx: mpsc::Sender<RpcServerOutboundMessage>) -> Self {
-        Self { tx }
-    }
-
-    pub(crate) async fn notify<P: Serialize>(
-        &self,
-        method: &str,
-        params: &P,
-    ) -> Result<(), serde_json::Error> {
-        let params = serde_json::to_value(params)?;
-        self.tx
-            .send(RpcServerOutboundMessage::Notification(JSONRPCNotification {
-                method: method.to_string(),
-                params: Some(params),
-            }))
-            .await
-            .map_err(|_| {
-                serde_json::Error::io(std::io::Error::new(
-                    std::io::ErrorKind::BrokenPipe,
-                    "JSON-RPC transport closed",
-                ))
-            })
-    }
-}
-
-type RpcRequestRoute<H> = dyn Fn(
-        Arc<H>,
-        codex_app_server_protocol::JSONRPCRequest,
-    ) -> Pin<Box<dyn std::future::Future<Output = RpcServerOutboundMessage> + Send>>
-    + Send
-    + Sync;
-
-type RpcNotificationRoute<H> = dyn Fn(
-        Arc<H>,
-        codex_app_server_protocol::JSONRPCNotification,
-    ) -> Pin<Box<dyn std::future::Future<Output = Result<(), String>> + Send>>
-    + Send
-    + Sync;
-
-pub(crate) struct RpcRouter<H> {
-    request_routes: HashMap<String, Box<RpcRequestRoute<H>>>,
-    notification_routes: HashMap<String, Box<RpcNotificationRoute<H>>>,
-}
-
-impl<H: Send + Sync + 'static> RpcRouter<H> {
-    pub(crate) fn new() -> Self {
-        Self {
-            request_routes: HashMap::new(),
-            notification_routes: HashMap::new(),
-        }
-    }
-
-    pub(crate) fn request<P, F, Fut, R>(&mut self, method: &str, handler: F)
-    where
-        P: DeserializeOwned + Send + 'static,
-        R: Serialize + Send + 'static,
-        F: Fn(Arc<H>, P) -> Fut + Send + Sync + 'static,
-        Fut: std::future::Future<Output = Result<R, JSONRPCErrorError>> + Send + 'static,
-    {
-        let method = method.to_string();
-        let handler = std::sync::Arc::new(handler);
-        self.request_routes.insert(
-            method,
-            Box::new(
-                move |server_handler: Arc<H>, request: codex_app_server_protocol::JSONRPCRequest| {
-                    let handler = std::sync::Arc::clone(&handler);
-                    let params = serde_json::from_value::<P>(request.params.unwrap_or(Value::Null))
-                        .map_err(|error| invalid_params(error.to_string()));
-                    let request_id = request.id;
-                    Box::pin(async move {
-                        let result = match params {
-                            Ok(params) => handler(server_handler.clone(), params)
-                                .await
-                                .and_then(|value| {
-                                serde_json::to_value(value)
-                                    .map_err(|error| invalid_params(error.to_string()))
-                            }),
-                            Err(error) => Err(error),
-                        };
-                        RpcServerOutboundMessage::response(request_id, result)
-                    })
-                },
-            ),
-        );
-    }
-
-    pub(crate) fn notification<P, F, Fut>(&mut self, method: &str, handler: F)
-    where
-        P: DeserializeOwned + Send + 'static,
-        F: Fn(Arc<H>, P) -> Fut + Send + Sync + 'static,
-        Fut: std::future::Future<Output = Result<(), String>> + Send + 'static,
-    {
-        let method = method.to_string();
-        let handler = std::sync::Arc::new(handler);
-        self.notification_routes.insert(
-            method,
-            Box::new(
-                move |
-                    server_handler: Arc<H>,
-                    notification: codex_app_server_protocol::JSONRPCNotification| {
-                    let handler = std::sync::Arc::clone(&handler);
-                    let params = serde_json::from_value::<P>(notification.params.unwrap_or(Value::Null))
-                        .map_err(|err| err.to_string());
-                    Box::pin(async move {
-                        match params {
-                            Ok(params) => handler(server_handler.clone(), params).await,
-                            Err(error) => Err(error),
-                        }
-                    })
-                },
-            ),
-        );
-    }
-
-    pub(crate) fn request_route(
-        &self,
-        method: &str,
-    ) -> Option<&Box<RpcRequestRoute<H>>> {
-        self.request_routes.get(method)
-    }
-
-    pub(crate) fn notification_route(
-        &self,
-        method: &str,
-    ) -> Option<&Box<RpcNotificationRoute<H>>> {
-        self.notification_routes.get(method)
-    }
-}
-
 #[cfg(test)]
 mod tests {
     use std::time::Duration;

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

@@ -1,16 +1,15 @@
 mod filesystem;
 mod handler;
-mod jsonrpc;
 mod registry;
 mod processor;
 mod transport;
 
 pub(crate) use handler::ExecServerHandler;
 pub use transport::DEFAULT_LISTEN_URL;
-pub use transport::ExecServerListenUrlParseError;
+pub use transport::ExecServerListenUrlParseError as ExecServerTransportParseError;
 
 pub async fn run_main() -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
-    run_main_with_listen_url(DEFAULT_LISTEN_URL).await
+    run_main_with_transport(DEFAULT_LISTEN_URL).await
 }
 
 pub async fn run_main_with_listen_url(
@@ -18,3 +17,9 @@ pub async fn run_main_with_listen_url(
 ) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
     transport::run_transport(listen_url).await
 }
+
+pub async fn run_main_with_transport(
+    listen_url: &str,
+) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
+    transport::run_transport(listen_url).await
+}

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

@@ -0,0 +1,182 @@
+use std::sync::atomic::AtomicI64;
+use std::sync::atomic::Ordering;
+
+use codex_app_server_protocol::JSONRPCMessage;
+use codex_app_server_protocol::JSONRPCNotification;
+use codex_app_server_protocol::JSONRPCRequest;
+use codex_app_server_protocol::RequestId;
+use codex_utils_cargo_bin::cargo_bin;
+use futures::{SinkExt, StreamExt};
+use tokio::io::AsyncBufReadExt;
+use tokio::io::BufReader;
+use tokio::process::Command;
+use std::process::Stdio;
+use tokio::sync::mpsc;
+use tokio::task::JoinHandle;
+use tokio::time::timeout;
+use tokio_tungstenite::tungstenite::Message;
+use tokio_tungstenite::connect_async;
+
+enum OutgoingMessage {
+    Json(JSONRPCMessage),
+    RawText(String),
+}
+
+pub struct ExecServer {
+    child: tokio::process::Child,
+    next_request_id: AtomicI64,
+    incoming_rx: mpsc::Receiver<JSONRPCMessage>,
+    outgoing_tx: mpsc::Sender<OutgoingMessage>,
+    reader_task: JoinHandle<()>,
+    writer_task: JoinHandle<()>,
+}
+
+impl ExecServer {
+    pub async fn send_request(
+        &mut self,
+        method: &str,
+        params: serde_json::Value,
+    ) -> anyhow::Result<RequestId> {
+        let request_id = RequestId::Integer(self.next_request_id.fetch_add(1, Ordering::SeqCst));
+        let request = JSONRPCRequest {
+            id: request_id.clone(),
+            method: method.to_string(),
+            params: Some(params),
+            trace: None,
+        };
+        self.outgoing_tx
+            .send(OutgoingMessage::Json(JSONRPCMessage::Request(request)))
+            .await?;
+        Ok(request_id)
+    }
+
+    pub async fn send_notification(
+        &mut self,
+        method: &str,
+        params: serde_json::Value,
+    ) -> anyhow::Result<()> {
+        let notification = JSONRPCNotification {
+            method: method.to_string(),
+            params: Some(params),
+        };
+        self.outgoing_tx
+            .send(OutgoingMessage::Json(JSONRPCMessage::Notification(
+                notification,
+            )))
+            .await?;
+        Ok(())
+    }
+
+    pub async fn send_raw_text(&mut self, text: &str) -> anyhow::Result<()> {
+        self.outgoing_tx
+            .send(OutgoingMessage::RawText(text.to_string()))
+            .await?;
+        Ok(())
+    }
+
+    pub async fn next_event(&mut self) -> anyhow::Result<JSONRPCMessage> {
+        self.incoming_rx
+            .recv()
+            .await
+            .ok_or_else(|| anyhow::anyhow!("exec-server closed before next event"))
+    }
+
+    pub async fn wait_for_event(
+        &mut self,
+        predicate: impl Fn(&JSONRPCMessage) -> bool,
+    ) -> anyhow::Result<JSONRPCMessage> {
+        loop {
+            let event = self.next_event().await?;
+            if predicate(&event) {
+                return Ok(event);
+            }
+        }
+    }
+
+    pub async fn shutdown(&mut self) -> anyhow::Result<()> {
+        self.reader_task.abort();
+        self.writer_task.abort();
+        self.child.start_kill()?;
+        Ok(())
+    }
+}
+
+pub mod exec_server {
+    use super::*;
+
+    pub async fn exec_server() -> anyhow::Result<super::ExecServer> {
+        let binary = cargo_bin("codex-exec-server")?;
+        let mut child = Command::new(binary);
+        child.args(["--listen", "ws://127.0.0.1:0"]);
+        child.stdin(Stdio::null());
+        child.stdout(Stdio::null());
+        child.stderr(Stdio::piped());
+        let mut child = child.spawn()?;
+
+        let stderr = child.stderr.take().expect("stderr should be piped");
+        let mut stderr_lines = BufReader::new(stderr).lines();
+        let websocket_url = read_websocket_url(&mut stderr_lines).await?;
+
+        let (websocket, _) = connect_async(websocket_url).await?;
+        let (mut outgoing_ws, mut incoming_ws) = websocket.split();
+        let (outgoing_tx, mut outgoing_rx) = mpsc::channel::<OutgoingMessage>(128);
+        let (incoming_tx, incoming_rx) = mpsc::channel::<JSONRPCMessage>(128);
+
+        let reader_task = tokio::spawn(async move {
+            while let Some(message) = incoming_ws.next().await {
+                let Ok(message) = message else {
+                    break;
+                };
+                let outgoing = match message {
+                    Message::Text(text) => serde_json::from_str::<JSONRPCMessage>(&text),
+                    Message::Binary(bytes) => serde_json::from_slice::<JSONRPCMessage>(&bytes),
+                    _ => continue,
+                };
+                if let Ok(message) = outgoing && let Err(_err) = incoming_tx.send(message).await {
+                    break;
+                }
+            }
+        });
+
+        let writer_task = tokio::spawn(async move {
+            while let Some(message) = outgoing_rx.recv().await {
+                let outgoing = match message {
+                    OutgoingMessage::Json(message) => {
+                        match serde_json::to_string(&message) {
+                            Ok(json) => Message::Text(json.into()),
+                            Err(_) => continue,
+                        }
+                    }
+                    OutgoingMessage::RawText(message) => Message::Text(message.into()),
+                };
+                if outgoing_ws.send(outgoing).await.is_err() {
+                    break;
+                }
+            }
+        });
+
+        Ok(ExecServer {
+            child,
+            next_request_id: AtomicI64::new(1),
+            incoming_rx,
+            outgoing_tx,
+            reader_task,
+            writer_task,
+        })
+    }
+
+    async fn read_websocket_url<R>(lines: &mut tokio::io::Lines<BufReader<R>>) -> anyhow::Result<String>
+    where
+        R: tokio::io::AsyncRead + Unpin,
+    {
+        let line = timeout(std::time::Duration::from_secs(5), lines.next_line())
+            .await??
+            .ok_or_else(|| anyhow::anyhow!("missing websocket startup banner"))?;
+
+        let websocket_url = line
+            .split_whitespace()
+            .find(|part| part.starts_with("ws://"))
+            .ok_or_else(|| anyhow::anyhow!("missing websocket URL in startup banner: {line}"))?;
+        Ok(websocket_url.to_string())
+    }
+}

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

@@ -24,7 +24,7 @@ async fn exec_server_stubs_process_start_over_websocket() -> anyhow::Result<()>
         .wait_for_event(|event| {
             matches!(
                 event,
-                JSONRPCMessage::Response(JSONRPCResponse { id, .. }) if id == &initialize_id
+                JSONRPCMessage::Response(JSONRPCResponse { id, .. }) if *id == initialize_id
             )
         })
         .await?;
@@ -53,7 +53,7 @@ async fn exec_server_stubs_process_start_over_websocket() -> anyhow::Result<()>
         .wait_for_event(|event| {
             matches!(
                 event,
-                JSONRPCMessage::Response(JSONRPCResponse { id, .. }) if id == &process_start_id
+                JSONRPCMessage::Response(JSONRPCResponse { id, .. }) if *id == process_start_id
             )
         })
         .await?;

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

@@ -44,7 +44,7 @@ async fn exec_server_reports_malformed_websocket_json_and_keeps_running() -> any
         .wait_for_event(|event| {
             matches!(
                 event,
-                JSONRPCMessage::Response(JSONRPCResponse { id, .. }) if id == &initialize_id
+                JSONRPCMessage::Response(JSONRPCResponse { id, .. }) if *id == initialize_id
             )
         })
         .await?;