Fix exec-server in-order request handling and stdio transport

Co-authored-by: Codex <noreply@openai.com>

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

@@ -11,11 +11,13 @@ use tokio::sync::mpsc;
 use tokio_tungstenite::WebSocketStream;
 use tokio_tungstenite::tungstenite::Message;
 
+
 pub(crate) const CHANNEL_CAPACITY: usize = 128;
 
 #[derive(Debug)]
 pub(crate) enum JsonRpcConnectionEvent {
     Message(JSONRPCMessage),
+    MalformedMessage { reason: String },
     Disconnected { reason: Option<String> },
 }
 
@@ -55,14 +57,13 @@ impl JsonRpcConnection {
                                 }
                             }
                             Err(err) => {
-                                send_disconnected(
+                                send_malformed_message(
                                     &incoming_tx_for_reader,
                                     Some(format!(
                                         "failed to parse JSON-RPC message from {reader_label}: {err}"
                                     )),
                                 )
                                 .await;
-                                break;
                             }
                         }
                     }
@@ -132,14 +133,13 @@ impl JsonRpcConnection {
                                 }
                             }
                             Err(err) => {
-                                send_disconnected(
+                                send_malformed_message(
                                     &incoming_tx_for_reader,
                                     Some(format!(
                                         "failed to parse websocket JSON-RPC message from {reader_label}: {err}"
                                     )),
                                 )
                                 .await;
-                                break;
                             }
                         }
                     }
@@ -155,14 +155,13 @@ impl JsonRpcConnection {
                                 }
                             }
                             Err(err) => {
-                                send_disconnected(
+                                send_malformed_message(
                                     &incoming_tx_for_reader,
                                     Some(format!(
                                         "failed to parse websocket JSON-RPC message from {reader_label}: {err}"
                                     )),
                                 )
                                 .await;
-                                break;
                             }
                         }
                     }
@@ -247,6 +246,17 @@ async fn send_disconnected(
         .await;
 }
 
+async fn send_malformed_message(
+    incoming_tx: &mpsc::Sender<JsonRpcConnectionEvent>,
+    reason: Option<String>,
+) {
+    let _ = incoming_tx
+        .send(JsonRpcConnectionEvent::MalformedMessage {
+            reason: reason.unwrap_or_else(|| "malformed JSON-RPC message".to_string()),
+        })
+        .await;
+}
+
 async fn write_jsonrpc_line_message<W>(
     writer: &mut BufWriter<W>,
     message: &JSONRPCMessage,

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

@@ -42,6 +42,7 @@ pub async fn spawn_local_exec_server(
 ) -> Result<SpawnedExecServer, ExecServerError> {
     let mut child = Command::new(&command.program);
     child.args(&command.args);
+    child.args(["--listen", "stdio://"]);
     child.stdin(Stdio::piped());
     child.stdout(Stdio::piped());
     child.stderr(Stdio::inherit());

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

@@ -4,6 +4,7 @@ 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;
@@ -441,6 +442,217 @@ 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,19 +1,20 @@
 mod filesystem;
 mod handler;
-mod processor;
+mod jsonrpc;
 mod registry;
+mod processor;
 mod transport;
 
 pub(crate) use handler::ExecServerHandler;
-pub use transport::ExecServerTransport;
-pub use transport::ExecServerTransportParseError;
+pub use transport::DEFAULT_LISTEN_URL;
+pub use transport::ExecServerListenUrlParseError;
 
 pub async fn run_main() -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
-    run_main_with_transport(ExecServerTransport::Stdio).await
+    run_main_with_listen_url(DEFAULT_LISTEN_URL).await
 }
 
-pub async fn run_main_with_transport(
-    transport: ExecServerTransport,
+pub async fn run_main_with_listen_url(
+    listen_url: &str,
 ) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
-    transport::run_transport(transport).await
+    transport::run_transport(listen_url).await
 }

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

@@ -10,6 +10,7 @@ use crate::connection::JsonRpcConnectionEvent;
 use crate::rpc::RpcNotificationSender;
 use crate::rpc::RpcServerOutboundMessage;
 use crate::rpc::encode_server_message;
+use crate::rpc::invalid_request;
 use crate::rpc::method_not_found;
 use crate::server::ExecServerHandler;
 use crate::server::registry::build_router;
@@ -39,15 +40,26 @@ pub(crate) async fn run_connection(connection: JsonRpcConnection) {
 
     while let Some(event) = incoming_rx.recv().await {
         match event {
+            JsonRpcConnectionEvent::MalformedMessage { reason } => {
+                warn!("ignoring malformed exec-server message: {reason}");
+                if outgoing_tx
+                    .send(RpcServerOutboundMessage::Error {
+                        request_id: codex_app_server_protocol::RequestId::Integer(-1),
+                        error: invalid_request(reason),
+                    })
+                    .await
+                    .is_err()
+                {
+                    break;
+                }
+            }
             JsonRpcConnectionEvent::Message(message) => match message {
                 codex_app_server_protocol::JSONRPCMessage::Request(request) => {
                     if let Some(route) = router.request_route(request.method.as_str()) {
-                        let route = route(handler.clone(), request);
-                        let outgoing_tx = outgoing_tx.clone();
-                        tokio::spawn(async move {
-                            let message = route.await;
-                            let _ = outgoing_tx.send(message).await;
-                        });
+                        let message = route(handler.clone(), request).await;
+                        if outgoing_tx.send(message).await.is_err() {
+                            break;
+                        }
                     } else if outgoing_tx
                         .send(RpcServerOutboundMessage::Error {
                             request_id: request.id,

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

@@ -38,7 +38,9 @@ pub(crate) fn build_router() -> RpcRouter<ExecServerHandler> {
     );
     router.notification(
         INITIALIZED_METHOD,
-        |handler: Arc<ExecServerHandler>, (): ()| async move { handler.initialized() },
+        |handler: Arc<ExecServerHandler>, _params: serde_json::Value| async move {
+            handler.initialized()
+        },
     );
     router.request(
         EXEC_METHOD,

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

@@ -1,33 +1,29 @@
 use std::net::SocketAddr;
-use std::str::FromStr;
 
 use tokio::net::TcpListener;
+use tokio::io;
 use tokio_tungstenite::accept_async;
 use tracing::warn;
 
 use crate::connection::JsonRpcConnection;
 use crate::server::processor::run_connection;
 
-#[derive(Clone, Copy, Debug, Eq, PartialEq)]
-pub enum ExecServerTransport {
-    Stdio,
-    WebSocket { bind_address: SocketAddr },
-}
+pub const DEFAULT_LISTEN_URL: &str = "ws://127.0.0.1:0";
 
 #[derive(Debug, Clone, Eq, PartialEq)]
-pub enum ExecServerTransportParseError {
+pub enum ExecServerListenUrlParseError {
     UnsupportedListenUrl(String),
     InvalidWebSocketListenUrl(String),
 }
 
-impl std::fmt::Display for ExecServerTransportParseError {
+impl std::fmt::Display for ExecServerListenUrlParseError {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
         match self {
-            ExecServerTransportParseError::UnsupportedListenUrl(listen_url) => write!(
+            ExecServerListenUrlParseError::UnsupportedListenUrl(listen_url) => write!(
                 f,
-                "unsupported --listen URL `{listen_url}`; expected `stdio://` or `ws://IP:PORT`"
+                "unsupported --listen URL `{listen_url}`; expected `ws://IP:PORT`"
             ),
-            ExecServerTransportParseError::InvalidWebSocketListenUrl(listen_url) => write!(
+            ExecServerListenUrlParseError::InvalidWebSocketListenUrl(listen_url) => write!(
                 f,
                 "invalid websocket --listen URL `{listen_url}`; expected `ws://IP:PORT`"
             ),
@@ -35,62 +31,59 @@ impl std::fmt::Display for ExecServerTransportParseError {
     }
 }
 
-impl std::error::Error for ExecServerTransportParseError {}
+impl std::error::Error for ExecServerListenUrlParseError {}
 
-impl ExecServerTransport {
-    pub const DEFAULT_LISTEN_URL: &str = "stdio://";
-
-    pub fn from_listen_url(listen_url: &str) -> Result<Self, ExecServerTransportParseError> {
-        if listen_url == Self::DEFAULT_LISTEN_URL {
-            return Ok(Self::Stdio);
-        }
-
-        if let Some(socket_addr) = listen_url.strip_prefix("ws://") {
-            let bind_address = socket_addr.parse::<SocketAddr>().map_err(|_| {
-                ExecServerTransportParseError::InvalidWebSocketListenUrl(listen_url.to_string())
-            })?;
-            return Ok(Self::WebSocket { bind_address });
-        }
-
-        Err(ExecServerTransportParseError::UnsupportedListenUrl(
-            listen_url.to_string(),
-        ))
-    }
+#[derive(Debug, Clone, Eq, PartialEq)]
+enum ListenAddress {
+    Websocket(SocketAddr),
+    Stdio,
 }
 
-impl FromStr for ExecServerTransport {
-    type Err = ExecServerTransportParseError;
-
-    fn from_str(s: &str) -> Result<Self, Self::Err> {
-        Self::from_listen_url(s)
+fn parse_listen_url(listen_url: &str) -> Result<ListenAddress, ExecServerListenUrlParseError> {
+    if let Some(socket_addr) = listen_url.strip_prefix("ws://") {
+        return socket_addr
+            .parse::<SocketAddr>()
+            .map(ListenAddress::Websocket)
+            .map_err(|_| {
+                ExecServerListenUrlParseError::InvalidWebSocketListenUrl(listen_url.to_string())
+            });
     }
+    if listen_url == "stdio://" {
+        return Ok(ListenAddress::Stdio);
+    }
+
+    Err(ExecServerListenUrlParseError::UnsupportedListenUrl(
+        listen_url.to_string(),
+    ))
 }
 
 pub(crate) async fn run_transport(
-    transport: ExecServerTransport,
+    listen_url: &str,
 ) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
-    match transport {
-        ExecServerTransport::Stdio => {
-            run_connection(JsonRpcConnection::from_stdio(
-                tokio::io::stdin(),
-                tokio::io::stdout(),
-                "exec-server stdio".to_string(),
-            ))
-            .await;
-            Ok(())
-        }
-        ExecServerTransport::WebSocket { bind_address } => {
-            run_websocket_listener(bind_address).await
-        }
+    match parse_listen_url(listen_url)? {
+        ListenAddress::Websocket(bind_address) => run_websocket_listener(bind_address).await,
+        ListenAddress::Stdio => run_stdio_listener().await,
     }
 }
 
+async fn run_stdio_listener() -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
+    run_connection(JsonRpcConnection::from_stdio(
+        io::stdin(),
+        io::stdout(),
+        "exec-server stdio".to_string(),
+    ))
+    .await;
+    Ok(())
+}
+
 async fn run_websocket_listener(
     bind_address: SocketAddr,
 ) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
     let listener = TcpListener::bind(bind_address).await?;
     let local_addr = listener.local_addr()?;
-    print_websocket_startup_banner(local_addr);
+    let listen_message = format!("codex-exec-server listening on ws://{local_addr}");
+    tracing::info!("{}", listen_message);
+    eprintln!("{listen_message}");
 
     loop {
         let (stream, peer_addr) = listener.accept().await?;
@@ -113,54 +106,6 @@ async fn run_websocket_listener(
     }
 }
 
-#[allow(clippy::print_stderr)]
-fn print_websocket_startup_banner(addr: SocketAddr) {
-    eprintln!("codex-exec-server listening on ws://{addr}");
-}
-
 #[cfg(test)]
-mod tests {
-    use pretty_assertions::assert_eq;
-
-    use super::ExecServerTransport;
-
-    #[test]
-    fn exec_server_transport_parses_stdio_listen_url() {
-        let transport =
-            ExecServerTransport::from_listen_url(ExecServerTransport::DEFAULT_LISTEN_URL)
-                .expect("stdio listen URL should parse");
-        assert_eq!(transport, ExecServerTransport::Stdio);
-    }
-
-    #[test]
-    fn exec_server_transport_parses_websocket_listen_url() {
-        let transport = ExecServerTransport::from_listen_url("ws://127.0.0.1:1234")
-            .expect("websocket listen URL should parse");
-        assert_eq!(
-            transport,
-            ExecServerTransport::WebSocket {
-                bind_address: "127.0.0.1:1234".parse().expect("valid socket address"),
-            }
-        );
-    }
-
-    #[test]
-    fn exec_server_transport_rejects_invalid_websocket_listen_url() {
-        let err = ExecServerTransport::from_listen_url("ws://localhost:1234")
-            .expect_err("hostname bind address should be rejected");
-        assert_eq!(
-            err.to_string(),
-            "invalid websocket --listen URL `ws://localhost:1234`; expected `ws://IP:PORT`"
-        );
-    }
-
-    #[test]
-    fn exec_server_transport_rejects_unsupported_listen_url() {
-        let err = ExecServerTransport::from_listen_url("http://127.0.0.1:1234")
-            .expect_err("unsupported scheme should fail");
-        assert_eq!(
-            err.to_string(),
-            "unsupported --listen URL `http://127.0.0.1:1234`; expected `stdio://` or `ws://IP:PORT`"
-        );
-    }
-}
+#[path = "transport_tests.rs"]
+mod transport_tests;

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

@@ -0,0 +1,50 @@
+use pretty_assertions::assert_eq;
+
+use super::DEFAULT_LISTEN_URL;
+use super::parse_listen_url;
+
+#[test]
+fn parse_listen_url_accepts_default_websocket_url() {
+    let bind_address = parse_listen_url(DEFAULT_LISTEN_URL)
+        .expect("default listen URL should parse");
+    assert_eq!(
+        bind_address,
+        super::ListenAddress::Websocket("127.0.0.1:0".parse().expect("valid socket address"))
+    );
+}
+
+#[test]
+fn parse_listen_url_accepts_stdio_url() {
+    let listen_address = parse_listen_url("stdio://").expect("stdio listen URL should parse");
+    assert_eq!(listen_address, super::ListenAddress::Stdio);
+}
+
+#[test]
+fn parse_listen_url_accepts_websocket_url() {
+    let bind_address = parse_listen_url("ws://127.0.0.1:1234")
+        .expect("websocket listen URL should parse");
+    assert_eq!(
+        bind_address,
+        super::ListenAddress::Websocket("127.0.0.1:1234".parse().expect("valid socket address"))
+    );
+}
+
+#[test]
+fn parse_listen_url_rejects_invalid_websocket_url() {
+    let err = parse_listen_url("ws://localhost:1234")
+        .expect_err("hostname bind address should be rejected");
+    assert_eq!(
+        err.to_string(),
+        "invalid websocket --listen URL `ws://localhost:1234`; expected `ws://IP:PORT`"
+    );
+}
+
+#[test]
+fn parse_listen_url_rejects_unsupported_url() {
+    let err =
+        parse_listen_url("http://127.0.0.1:1234").expect_err("unsupported scheme should fail");
+    assert_eq!(
+        err.to_string(),
+        "unsupported --listen URL `http://127.0.0.1:1234`; expected `ws://IP:PORT`"
+    );
+}

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

@@ -0,0 +1,39 @@
+#![cfg(unix)]
+
+mod common;
+
+use codex_app_server_protocol::JSONRPCMessage;
+use codex_app_server_protocol::JSONRPCResponse;
+use codex_exec_server::InitializeParams;
+use codex_exec_server::InitializeResponse;
+use common::exec_server::exec_server;
+use pretty_assertions::assert_eq;
+
+#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+async fn exec_server_accepts_initialize() -> anyhow::Result<()> {
+    let mut server = exec_server().await?;
+    let initialize_id = server
+        .send_request(
+            "initialize",
+            serde_json::to_value(InitializeParams {
+                client_name: "exec-server-test".to_string(),
+            })?,
+        )
+        .await?;
+
+    let response = server.next_event().await?;
+    let JSONRPCMessage::Response(JSONRPCResponse { id, result }) = response else {
+        panic!("expected initialize response");
+    };
+    assert_eq!(id, initialize_id);
+    let initialize_response: InitializeResponse = serde_json::from_value(result)?;
+    assert_eq!(
+        initialize_response,
+        InitializeResponse {
+            protocol_version: "exec-server.v0".to_string()
+        }
+    );
+
+    server.shutdown().await?;
+    Ok(())
+}

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

@@ -0,0 +1,74 @@
+#![cfg(unix)]
+
+mod common;
+
+use codex_app_server_protocol::JSONRPCMessage;
+use codex_app_server_protocol::JSONRPCResponse;
+use codex_exec_server::InitializeParams;
+use codex_exec_server::ExecResponse;
+use common::exec_server::exec_server;
+use pretty_assertions::assert_eq;
+
+#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+async fn exec_server_stubs_process_start_over_websocket() -> anyhow::Result<()> {
+    let mut server = exec_server().await?;
+    let initialize_id = server
+        .send_request(
+            "initialize",
+            serde_json::to_value(InitializeParams {
+                client_name: "exec-server-test".to_string(),
+            })?,
+        )
+        .await?;
+    let _ = server
+        .wait_for_event(|event| {
+            matches!(
+                event,
+                JSONRPCMessage::Response(JSONRPCResponse { id, .. }) if id == &initialize_id
+            )
+        })
+        .await?;
+
+    server
+        .send_notification(
+            "initialized",
+            serde_json::json!({}),
+        )
+        .await?;
+
+    let process_start_id = server
+        .send_request(
+            "process/start",
+            serde_json::json!({
+                "processId": "proc-1",
+                "argv": ["true"],
+                "cwd": std::env::current_dir()?,
+                "env": {},
+                "tty": false,
+                "arg0": null
+            }),
+        )
+        .await?;
+    let response = server
+        .wait_for_event(|event| {
+            matches!(
+                event,
+                JSONRPCMessage::Response(JSONRPCResponse { id, .. }) if id == &process_start_id
+            )
+        })
+        .await?;
+    let JSONRPCMessage::Response(JSONRPCResponse { id, result }) = response else {
+        panic!("expected process/start response");
+    };
+    assert_eq!(id, process_start_id);
+    let process_start_response: ExecResponse = serde_json::from_value(result)?;
+    assert_eq!(
+        process_start_response,
+        ExecResponse {
+            process_id: "proc-1".to_string()
+        }
+    );
+
+    server.shutdown().await?;
+    Ok(())
+}

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

@@ -41,6 +41,7 @@ use tokio::time::timeout;
 async fn exec_server_accepts_initialize_over_stdio() -> anyhow::Result<()> {
     let binary = cargo_bin("codex-exec-server")?;
     let mut child = Command::new(binary);
+    child.args(["--listen", "stdio://"]);
     child.stdin(Stdio::piped());
     child.stdout(Stdio::piped());
     child.stderr(Stdio::inherit());

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

@@ -0,0 +1,65 @@
+#![cfg(unix)]
+
+mod common;
+
+use codex_app_server_protocol::JSONRPCError;
+use codex_app_server_protocol::JSONRPCMessage;
+use codex_app_server_protocol::JSONRPCResponse;
+use codex_exec_server::InitializeParams;
+use codex_exec_server::InitializeResponse;
+use common::exec_server::exec_server;
+use pretty_assertions::assert_eq;
+
+#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+async fn exec_server_reports_malformed_websocket_json_and_keeps_running() -> anyhow::Result<()> {
+    let mut server = exec_server().await?;
+    server.send_raw_text("not-json").await?;
+
+    let response = server
+        .wait_for_event(|event| matches!(event, JSONRPCMessage::Error(_)))
+        .await?;
+    let JSONRPCMessage::Error(JSONRPCError { id, error }) = response else {
+        panic!("expected malformed-message error response");
+    };
+    assert_eq!(id, codex_app_server_protocol::RequestId::Integer(-1));
+    assert_eq!(error.code, -32600);
+    assert!(
+        error
+            .message
+            .starts_with("failed to parse websocket JSON-RPC message from exec-server websocket"),
+        "unexpected malformed-message error: {}",
+        error.message
+    );
+
+    let initialize_id = server
+        .send_request(
+            "initialize",
+            serde_json::to_value(InitializeParams {
+                client_name: "exec-server-test".to_string(),
+            })?,
+        )
+        .await?;
+
+    let response = server
+        .wait_for_event(|event| {
+            matches!(
+                event,
+                JSONRPCMessage::Response(JSONRPCResponse { id, .. }) if id == &initialize_id
+            )
+        })
+        .await?;
+    let JSONRPCMessage::Response(JSONRPCResponse { id, result }) = response else {
+        panic!("expected initialize response after malformed input");
+    };
+    assert_eq!(id, initialize_id);
+    let initialize_response: InitializeResponse = serde_json::from_value(result)?;
+    assert_eq!(
+        initialize_response,
+        InitializeResponse {
+            protocol_version: "exec-server.v0".to_string()
+        }
+    );
+
+    server.shutdown().await?;
+    Ok(())
+}