Disable Nagle's algorithm for code-mode WebSockets (#37504)

## Why

Code-mode WebSocket connections are latency-sensitive, so buffering small TCP
writes can delay request and response traffic.

## What changed

- Enable `TCP_NODELAY` on outbound remote-session WebSocket connections.
- Enable `TCP_NODELAY` on sockets accepted by the code-mode host, logging a
  warning if the socket option cannot be set.

## Testing

- Add a listener test that connects to the host and verifies the accepted
  socket has `TCP_NODELAY` enabled.

GitOrigin-RevId: e51c781c4b47c6a4ae1c32c93cd79768719a68d9

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<ListenTransport> {
     );
 }
 
-async fn run_websocket_listener(bind_address: SocketAddr) -> Result<()> {
+async fn bind_websocket_listener(
+    bind_address: SocketAddr,
+) -> Result<TapIo<TcpListener, impl FnMut(&mut TcpStream) + Send + 'static>> {
     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")?;

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();

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));