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
This commit is contained in:
Channing Conger authored and copyberry committed 2026-08-03 23:48:28 +00:00
1 parent b258c028fe
commit 60c722e075
21 files changed
+1809 -135

No files matched your search

+16 -3
View File
@@ -16,6 +16,7 @@ use codex_code_mode_protocol::host::HostToClient;
use codex_code_mode_protocol::host::MAX_PENDING_DELEGATE_CALLS;
use codex_code_mode_protocol::host::RequestId;
use codex_code_mode_protocol::host::SessionId;
use codex_code_mode_protocol::host::TransportLane;
use codex_code_mode_protocol::host::WireResult;
use tokio::sync::Mutex;
use tokio::sync::Notify;
@@ -29,6 +30,7 @@ const CELL_MESSAGE_CAPACITY: usize = 128;
pub(super) struct HostPeer {
outgoing_tx: mpsc::Sender<EncodedFrame>,
bulk_tx: Option<mpsc::Sender<EncodedFrame>>,
pending: Mutex<HashMap<DelegateRequestId, PendingDelegate>>,
delegate_permits: Arc<Semaphore>,
cell_routes: StdMutex<HashMap<(SessionId, CellId), CellRoute>>,
@@ -62,6 +64,7 @@ impl HostPeer {
pub(super) fn new(outgoing_tx: mpsc::Sender<EncodedFrame>) -> Self {
Self {
outgoing_tx,
bulk_tx: None,
pending: Mutex::new(HashMap::new()),
delegate_permits: Arc::new(Semaphore::new(MAX_PENDING_DELEGATE_CALLS)),
cell_routes: StdMutex::new(HashMap::new()),
@@ -72,10 +75,16 @@ impl HostPeer {
}
}
pub(super) fn with_bulk_sender(mut self, bulk_tx: mpsc::Sender<EncodedFrame>) -> Self {
self.bulk_tx = Some(bulk_tx);
self
}
pub(super) fn send(&self, message: HostToClient) -> Result<(), PeerSendError> {
let frame = EncodedFrame::encode(&message)
.map_err(|err| PeerSendError::Payload(err.to_string()))?;
self.send_frame(frame)
let lane = message.transport_lane();
self.send_frame(frame, lane)
}
pub(super) fn respond(
@@ -389,8 +398,12 @@ impl HostPeer {
});
}
fn send_frame(&self, frame: EncodedFrame) -> Result<(), PeerSendError> {
match self.outgoing_tx.try_send(frame) {
fn send_frame(&self, frame: EncodedFrame, lane: TransportLane) -> Result<(), PeerSendError> {
let sender = match lane {
TransportLane::Control => &self.outgoing_tx,
TransportLane::Bulk => self.bulk_tx.as_ref().unwrap_or(&self.outgoing_tx),
};
match sender.try_send(frame) {
Ok(()) => Ok(()),
Err(mpsc::error::TrySendError::Full(_)) => {
self.disconnect();