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
2026-08-03 23:48:28 +00:00
committed by copyberry
parent b258c028fe
commit 60c722e075
21 changed files with 1809 additions and 135 deletions
@@ -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(),