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-protocol/src/host/host_tests.rs

@@ -22,6 +22,7 @@ use super::ProtocolVersion;
 use super::RequestId;
 use super::SessionId;
 use super::SupportedProtocolVersions;
+use super::TransportLane;
 use super::WireCellId;
 use super::WireContentItem;
 use super::WireExecuteRequest;
@@ -72,6 +73,130 @@ where
     );
 }
 
+#[test]
+fn dual_websocket_hello_preserves_the_pairing_token() {
+    assert_wire_round_trip(
+        HostToClient::HostHello(
+            HostHello::new(
+                ProtocolVersion::V1,
+                CapabilitySet::try_new([capability("dual-websocket-v1")])
+                    .expect("valid capabilities"),
+            )
+            .with_bulk_connection_token("pairing-token".to_string()),
+        ),
+        json!({
+            "type": "connection/ready",
+            "selectedVersion": 1,
+            "capabilities": ["dual-websocket-v1"],
+            "bulkConnectionToken": "pairing-token",
+        }),
+    );
+}
+
+#[test]
+fn message_families_use_dedicated_transport_lanes() {
+    for (message, lane) in [
+        (
+            ClientToHost::CancelRequest {
+                id: request_id(/*value*/ 1),
+            },
+            TransportLane::Control,
+        ),
+        (
+            ClientToHost::DelegateResponse {
+                id: delegate_request_id(/*value*/ 1),
+                result: WireResult::Ok {
+                    value: DelegateResponse::NotificationDelivered,
+                },
+            },
+            TransportLane::Control,
+        ),
+        (
+            ClientToHost::DelegateResponse {
+                id: delegate_request_id(/*value*/ 2),
+                result: WireResult::Ok {
+                    value: DelegateResponse::ToolResult {
+                        result: json!({ "value": "tool result" }),
+                    },
+                },
+            },
+            TransportLane::Bulk,
+        ),
+        (
+            ClientToHost::DelegateResponse {
+                id: delegate_request_id(/*value*/ 3),
+                result: WireResult::Err {
+                    message: "delegate failed".to_string(),
+                },
+            },
+            TransportLane::Bulk,
+        ),
+    ] {
+        assert_eq!(message.transport_lane(), lane);
+        assert!(message.allows_transport_lane(lane));
+        assert!(!message.allows_transport_lane(match lane {
+            TransportLane::Control => TransportLane::Bulk,
+            TransportLane::Bulk => TransportLane::Control,
+        }));
+    }
+
+    for (message, lane) in [
+        (
+            HostToClient::Response {
+                id: request_id(/*value*/ 1),
+                result: WireResult::Err {
+                    message: "x".repeat(128 * 1024),
+                },
+            },
+            TransportLane::Control,
+        ),
+        (
+            HostToClient::DelegateRequest {
+                id: delegate_request_id(/*value*/ 1),
+                session_id: session_id(),
+                request: DelegateRequest::Notify {
+                    call_id: "call-1".to_string(),
+                    cell_id: cell_id("cell-1"),
+                    text: "important".to_string(),
+                },
+            },
+            TransportLane::Control,
+        ),
+        (
+            HostToClient::DelegateRequest {
+                id: delegate_request_id(/*value*/ 2),
+                session_id: session_id(),
+                request: DelegateRequest::InvokeTool {
+                    invocation: WireNestedToolCall {
+                        cell_id: cell_id("cell-1"),
+                        runtime_tool_call_id: "runtime-call-1".to_string(),
+                        tool_name: WireToolName {
+                            name: "tool".to_string(),
+                            namespace: None,
+                        },
+                        tool_kind: WireToolKind::Function,
+                        input: None,
+                    },
+                },
+            },
+            TransportLane::Bulk,
+        ),
+        (
+            HostToClient::CancelDelegateRequest {
+                id: delegate_request_id(/*value*/ 1),
+            },
+            TransportLane::Bulk,
+        ),
+    ] {
+        assert_eq!(message.transport_lane(), lane);
+        assert!(message.allows_transport_lane(lane));
+        assert!(!message.allows_transport_lane(match lane {
+            TransportLane::Control => TransportLane::Bulk,
+            TransportLane::Bulk => TransportLane::Control,
+        }));
+    }
+}
+
 fn execute_request() -> WireExecuteRequest {
     WireExecuteRequest {
         tool_call_id: "call-1".to_string(),

codex-rs/code-mode-protocol/src/host/message.rs

@@ -12,6 +12,7 @@ use super::ProtocolVersion;
 use super::RequestId;
 use super::SessionId;
 use super::SupportedProtocolVersions;
+use super::TransportLane;
 use super::WireCellId;
 use super::WireExecuteRequest;
 use super::WireNestedToolCall;
@@ -105,6 +106,8 @@ impl std::error::Error for ClientHelloError {}
 pub struct HostHello {
     selected_version: ProtocolVersion,
     capabilities: CapabilitySet,
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    bulk_connection_token: Option<String>,
 }
 
 impl HostHello {
@@ -112,9 +115,15 @@ impl HostHello {
         Self {
             selected_version,
             capabilities,
+            bulk_connection_token: None,
         }
     }
 
+    pub fn with_bulk_connection_token(mut self, token: String) -> Self {
+        self.bulk_connection_token = Some(token);
+        self
+    }
+
     pub fn selected_version(&self) -> ProtocolVersion {
         self.selected_version
     }
@@ -122,6 +131,10 @@ impl HostHello {
     pub fn capabilities(&self) -> &CapabilitySet {
         &self.capabilities
     }
+
+    pub fn bulk_connection_token(&self) -> Option<&str> {
+        self.bulk_connection_token.as_deref()
+    }
 }
 
 /// Messages sent from a client to the code-mode host.
@@ -141,6 +154,30 @@ pub enum ClientToHost {
     },
 }
 
+impl ClientToHost {
+    /// Keeps notification acknowledgments with control traffic and tool results on the bulk lane.
+    pub fn transport_lane(&self) -> TransportLane {
+        match self {
+            Self::DelegateResponse {
+                result:
+                    WireResult::Ok {
+                        value: DelegateResponse::NotificationDelivered,
+                    },
+                ..
+            }
+            | Self::ClientHello(_)
+            | Self::Request { .. }
+            | Self::CancelRequest { .. } => TransportLane::Control,
+            Self::DelegateResponse { .. } => TransportLane::Bulk,
+        }
+    }
+
+    /// Validates the message families accepted by each paired socket.
+    pub fn allows_transport_lane(&self, lane: TransportLane) -> bool {
+        self.transport_lane() == lane
+    }
+}
+
 /// Messages sent from the code-mode host to a client.
 #[derive(Debug, Deserialize, PartialEq, Serialize)]
 #[serde(deny_unknown_fields, tag = "type", rename_all_fields = "camelCase")]
@@ -174,6 +211,33 @@ pub enum HostToClient {
     },
 }
 
+impl HostToClient {
+    /// Keeps notifications with control traffic and nested-tool callbacks on the bulk lane.
+    pub fn transport_lane(&self) -> TransportLane {
+        match self {
+            Self::DelegateRequest {
+                request: DelegateRequest::InvokeTool { .. },
+                ..
+            }
+            | Self::CancelDelegateRequest { .. } => TransportLane::Bulk,
+            Self::DelegateRequest {
+                request: DelegateRequest::Notify { .. },
+                ..
+            }
+            | Self::HostHello(_)
+            | Self::HandshakeRejected { .. }
+            | Self::Response { .. }
+            | Self::InitialResponse { .. }
+            | Self::CellClosed { .. } => TransportLane::Control,
+        }
+    }
+
+    /// Rejects messages received on the wrong paired socket.
+    pub fn allows_transport_lane(&self, lane: TransportLane) -> bool {
+        self.transport_lane() == lane
+    }
+}
+
 #[derive(Clone, Debug, Deserialize, PartialEq, Serialize)]
 #[serde(deny_unknown_fields, tag = "method", rename_all_fields = "camelCase")]
 pub enum HostRequest {

codex-rs/code-mode-protocol/src/host/mod.rs

@@ -1,9 +1,8 @@
 //! Messages and framing for the code-mode host boundary.
 //!
 //! Protocol version 1 multiplexes session operations and delegate callbacks by
-//! request ID over one ordered connection. It defines no optional capabilities
-//! yet; capability names provide an extension point for later versions without
-//! weakening the v1 decoder.
+//! request ID over one ordered connection. WebSocket peers can negotiate a
+//! separate bulk connection without changing the existing inner messages.
 
 mod codec;
 mod error;
@@ -14,6 +13,16 @@ mod types;
 /// Maximum number of unresolved delegate callbacks allowed per host connection.
 pub const MAX_PENDING_DELEGATE_CALLS: usize = 1_024;
 
+/// Optional second WebSocket carrying delegate callbacks and their responses.
+pub const DUAL_WEBSOCKET_CAPABILITY: &str = "dual-websocket-v1";
+
+/// Selects one socket of a negotiated dual-WebSocket connection.
+#[derive(Clone, Copy, Debug, Eq, PartialEq)]
+pub enum TransportLane {
+    Control,
+    Bulk,
+}
+
 pub use codec::EncodedFrame;
 pub use codec::FramedReader;
 pub use codec::FramedWriter;