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

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

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;