Add a dual-WebSocket transport for code mode (#36812)

## Why

Large nested-tool callbacks can occupy a WebSocket and delay unrelated session
operations on the same code-mode connection.

## What changed

- Negotiate the optional `dual-websocket-v1` capability and pair a second,
  token-scoped WebSocket with the control connection.
- Route nested-tool callbacks and their results over the bulk socket while
  keeping session operations, notifications, and execution responses on the
  control socket. Reject messages sent on the wrong lane.
- Preserve the single-connection transport when the capability is unavailable,
  and bound pairing, queued callbacks, and deferred cross-socket messages.
- Defer callbacks that arrive before their execution-started response, and
  return delegate errors without disconnecting the connection.

## Testing

Add protocol, transport, driver, and WebSocket integration coverage for
capability negotiation, lane routing, pairing failures, out-of-order messages,
and progress during large concurrent tool results.

GitOrigin-RevId: fa4504653e7cbf3c4ec930ae57aa0a41345bad66

codex-rs/code-mode-host/src/transport.rs

@@ -1,13 +1,17 @@
+use std::collections::HashMap;
 use std::io;
 use std::io::Write as _;
 use std::net::SocketAddr;
 use std::sync::Arc;
+use std::sync::Mutex;
+use std::sync::PoisonError;
 
 use anyhow::Context;
 use anyhow::Result;
 use axum::Router;
 use axum::body::Body;
 use axum::extract::ConnectInfo;
+use axum::extract::Path;
 use axum::extract::State;
 use axum::extract::ws::Message;
 use axum::extract::ws::WebSocket;
@@ -34,15 +38,19 @@ use futures::stream::SplitStream;
 use tokio::io::AsyncRead;
 use tokio::io::AsyncWrite;
 use tokio::net::TcpListener;
+use tokio::sync::oneshot;
 use tracing::info;
 use tracing::warn;
+use uuid::Uuid;
 
 use crate::HostLimits;
+use crate::MAX_IN_FLIGHT_REQUESTS;
 
 /// The default transport retains the standalone host's original stdio behavior.
 pub const DEFAULT_LISTEN_URL: &str = "stdio";
 
 const MAX_WEBSOCKET_FRAME_BYTES: usize = MAX_FRAME_BYTES + std::mem::size_of::<u32>();
+const MAX_PENDING_BULK_CONNECTIONS: usize = MAX_IN_FLIGHT_REQUESTS;
 
 type BoxedReader = Box<dyn AsyncRead + Send + Unpin>;
 type BoxedWriter = Box<dyn AsyncWrite + Send + Unpin>;
@@ -66,6 +74,64 @@ pub(crate) enum ConnectionWriter {
 #[derive(Clone)]
 struct WebSocketListenerState {
     limits: Arc<HostLimits>,
+    bulk_connections: BulkConnectionRegistry,
+}
+
+pub(crate) struct BulkConnection {
+    pub(crate) reader: ConnectionReader,
+    pub(crate) writer: ConnectionWriter,
+}
+
+#[derive(Clone, Default)]
+pub(crate) struct BulkConnectionRegistry {
+    pending: Arc<Mutex<HashMap<Uuid, oneshot::Sender<BulkConnection>>>>,
+}
+
+pub(crate) struct BulkConnectionRegistration {
+    registry: BulkConnectionRegistry,
+    token: Uuid,
+    receiver: oneshot::Receiver<BulkConnection>,
+}
+
+impl BulkConnectionRegistry {
+    pub(crate) fn reserve(&self) -> Option<BulkConnectionRegistration> {
+        let mut pending = self.pending.lock().unwrap_or_else(PoisonError::into_inner);
+        if pending.len() >= MAX_PENDING_BULK_CONNECTIONS {
+            return None;
+        }
+
+        let token = Uuid::new_v4();
+        let (sender, receiver) = oneshot::channel();
+        pending.insert(token, sender);
+        Some(BulkConnectionRegistration {
+            registry: self.clone(),
+            token,
+            receiver,
+        })
+    }
+
+    pub(crate) fn remove(&self, token: Uuid) -> Option<oneshot::Sender<BulkConnection>> {
+        self.pending
+            .lock()
+            .unwrap_or_else(PoisonError::into_inner)
+            .remove(&token)
+    }
+}
+
+impl BulkConnectionRegistration {
+    pub(crate) fn token(&self) -> Uuid {
+        self.token
+    }
+
+    pub(crate) async fn receive(&mut self) -> Result<BulkConnection, oneshot::error::RecvError> {
+        (&mut self.receiver).await
+    }
+}
+
+impl Drop for BulkConnectionRegistration {
+    fn drop(&mut self) {
+        self.registry.remove(self.token);
+    }
 }
 
 impl ConnectionReader {
@@ -165,6 +231,7 @@ async fn run_websocket_listener(bind_address: SocketAddr) -> Result<()> {
         .context("failed to read code-mode host websocket listen address")?;
     let state = WebSocketListenerState {
         limits: Arc::new(HostLimits::new()),
+        bulk_connections: BulkConnectionRegistry::default(),
     };
     info!("codex-code-mode-host listening on ws://{local_addr}");
     println!("ws://{local_addr}");
@@ -174,6 +241,7 @@ async fn run_websocket_listener(bind_address: SocketAddr) -> Result<()> {
 
     let router = Router::new()
         .route("/", any(websocket_upgrade_handler))
+        .route("/bulk/{token}", any(bulk_websocket_upgrade_handler))
         .route("/readyz", get(readiness_handler))
         .layer(middleware::from_fn(reject_requests_with_origin_header))
         .with_state(state);
@@ -220,6 +288,7 @@ async fn websocket_upgrade_handler(
                 ConnectionReader::WebSocket(reader),
                 ConnectionWriter::WebSocket(writer),
                 state.limits,
+                Some(state.bulk_connections),
             )
             .await
             {
@@ -228,6 +297,38 @@ async fn websocket_upgrade_handler(
         })
 }
 
+async fn bulk_websocket_upgrade_handler(
+    websocket: WebSocketUpgrade,
+    Path(token): Path<String>,
+    ConnectInfo(peer_addr): ConnectInfo<SocketAddr>,
+    State(state): State<WebSocketListenerState>,
+) -> Response {
+    let Ok(token) = Uuid::parse_str(&token) else {
+        return StatusCode::NOT_FOUND.into_response();
+    };
+    let Some(pending) = state.bulk_connections.remove(token) else {
+        return StatusCode::NOT_FOUND.into_response();
+    };
+
+    websocket
+        .max_frame_size(MAX_WEBSOCKET_FRAME_BYTES)
+        .max_message_size(MAX_WEBSOCKET_FRAME_BYTES)
+        .on_upgrade(move |stream| async move {
+            info!(%peer_addr, "code-mode host bulk websocket client connected");
+            let (writer, reader) = stream.split();
+            if pending
+                .send(BulkConnection {
+                    reader: ConnectionReader::WebSocket(reader),
+                    writer: ConnectionWriter::WebSocket(writer),
+                })
+                .is_err()
+            {
+                warn!(%peer_addr, "code-mode host bulk websocket pairing expired");
+            }
+        })
+        .into_response()
+}
+
 #[cfg(test)]
 #[path = "transport_tests.rs"]
 mod tests;