diff --git a/codex-rs/code-mode-host/src/transport.rs b/codex-rs/code-mode-host/src/transport.rs index c06fca40e3..2f3ed0a7dc 100644 --- a/codex-rs/code-mode-host/src/transport.rs +++ b/codex-rs/code-mode-host/src/transport.rs @@ -25,6 +25,9 @@ use axum::response::IntoResponse; use axum::response::Response; use axum::routing::any; use axum::routing::get; +use axum::serve::Listener; +use axum::serve::ListenerExt; +use axum::serve::TapIo; use codex_code_mode_protocol::host::ClientToHost; use codex_code_mode_protocol::host::EncodedFrame; use codex_code_mode_protocol::host::FramedReader; @@ -38,6 +41,7 @@ use futures::stream::SplitStream; use tokio::io::AsyncRead; use tokio::io::AsyncWrite; use tokio::net::TcpListener; +use tokio::net::TcpStream; use tokio::sync::oneshot; use tracing::info; use tracing::warn; @@ -222,10 +226,22 @@ fn parse_listen_url(listen_url: &str) -> Result { ); } -async fn run_websocket_listener(bind_address: SocketAddr) -> Result<()> { +async fn bind_websocket_listener( + bind_address: SocketAddr, +) -> Result> { let listener = TcpListener::bind(bind_address) .await .with_context(|| format!("failed to bind code-mode host websocket to {bind_address}"))?; + + Ok(listener.tap_io(|stream| { + if let Err(error) = stream.set_nodelay(/*nodelay*/ true) { + warn!(%error, "failed to enable TCP_NODELAY for code-mode host connection"); + } + })) +} + +async fn run_websocket_listener(bind_address: SocketAddr) -> Result<()> { + let listener = bind_websocket_listener(bind_address).await?; let local_addr = listener .local_addr() .context("failed to read code-mode host websocket listen address")?; diff --git a/codex-rs/code-mode-host/src/transport_tests.rs b/codex-rs/code-mode-host/src/transport_tests.rs index 7b4955f488..39a3bc5084 100644 --- a/codex-rs/code-mode-host/src/transport_tests.rs +++ b/codex-rs/code-mode-host/src/transport_tests.rs @@ -1,12 +1,38 @@ use std::net::SocketAddr; +use axum::serve::Listener; use pretty_assertions::assert_eq; +use tokio::net::TcpStream; use super::BulkConnectionRegistry; use super::ListenTransport; use super::MAX_PENDING_BULK_CONNECTIONS; +use super::bind_websocket_listener; use super::parse_listen_url; +#[tokio::test] +async fn websocket_listener_disables_nagle() { + let bind_address = "127.0.0.1:0" + .parse() + .expect("websocket test listener should have a valid bind address"); + let mut listener = bind_websocket_listener(bind_address) + .await + .expect("websocket test listener should bind"); + let local_addr = listener + .local_addr() + .expect("websocket test listener should have a local address"); + let _client = TcpStream::connect(local_addr) + .await + .expect("websocket test client should connect"); + let (stream, _) = listener.accept().await; + + let stream = stream + .nodelay() + .expect("accepted websocket socket should expose TCP_NODELAY"); + + assert!(stream); +} + #[test] fn bulk_connection_registration_cleans_up_when_dropped() { let registry = BulkConnectionRegistry::default(); diff --git a/codex-rs/code-mode/src/remote_session/connection.rs b/codex-rs/code-mode/src/remote_session/connection.rs index ff281e5647..e5c7a4c3d3 100644 --- a/codex-rs/code-mode/src/remote_session/connection.rs +++ b/codex-rs/code-mode/src/remote_session/connection.rs @@ -656,11 +656,13 @@ async fn connect_websocket_transport( "failed to build code-mode host websocket request: {error}" )) })?; - let connector = WebSocketConnector::new(http_client_factory).map_err(|error| { - ConnectionError::Other(format!( - "failed to configure code-mode host websocket TLS: {error}" - )) - })?; + let connector = WebSocketConnector::new(http_client_factory) + .map_err(|error| { + ConnectionError::Other(format!( + "failed to configure code-mode host websocket TLS: {error}" + )) + })? + .with_tcp_nodelay(); let websocket_config = WebSocketConfig::default() .max_frame_size(Some(MAX_WEBSOCKET_FRAME_BYTES)) .max_message_size(Some(MAX_WEBSOCKET_FRAME_BYTES));