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

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

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`"
+    );
+}