openai/codex
1
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-10-19 17:28:46 +00:00
Code Issues Packages Projects Releases Wiki Activity

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

Browse Source 
## 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
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/code-mode/src/remote_session/connection.rs
+217 -63
View File
@@ -14,17 +14,21 @@ use codex_code_mode_protocol::ExecuteRequest;
use codex_code_mode_protocol::StartedCell;
use codex_code_mode_protocol::WaitOutcome;
use codex_code_mode_protocol::WaitRequest;
use codex_code_mode_protocol::host::Capability;
use codex_code_mode_protocol::host::CapabilitySet;
use codex_code_mode_protocol::host::ClientHello;
use codex_code_mode_protocol::host::ClientToHost;
use codex_code_mode_protocol::host::DUAL_WEBSOCKET_CAPABILITY;
use codex_code_mode_protocol::host::EncodedFrame;
use codex_code_mode_protocol::host::FramedReader;
use codex_code_mode_protocol::host::FramedWriter;
use codex_code_mode_protocol::host::HostToClient;
use codex_code_mode_protocol::host::MAX_FRAME_BYTES;
use codex_code_mode_protocol::host::MAX_PENDING_DELEGATE_CALLS;
use codex_code_mode_protocol::host::ProtocolVersion;
use codex_code_mode_protocol::host::RequestId;
use codex_code_mode_protocol::host::SupportedProtocolVersions;
use codex_code_mode_protocol::host::TransportLane;
use codex_http_client::HttpClientFactory;
use codex_websocket_client::WebSocketConnector;
use futures::StreamExt;
@@ -36,6 +40,7 @@ use tokio::sync::mpsc;
use tokio::sync::oneshot;
use tokio::task::JoinHandle;
use tokio_tungstenite::tungstenite::client::IntoClientRequest;
use tokio_tungstenite::tungstenite::http::Uri;
use tokio_tungstenite::tungstenite::protocol::WebSocketConfig;
use tokio_util::sync::CancellationToken;
use tracing::debug;
@@ -68,6 +73,7 @@ pub(super) enum ConnectionError {
host_program: PathBuf,
error: io::Error,
},
BulkConnectionUnavailable(String),
Other(String),
}
@@ -98,7 +104,9 @@ impl fmt::Display for ConnectionError {
&host_program[suffix_start..]
)
}
Self::Other(message) => formatter.write_str(message),
Self::BulkConnectionUnavailable(message) | Self::Other(message) => {
formatter.write_str(message)
}
}
}
}
@@ -132,6 +140,45 @@ enum ConnectionOwner {
WebSocket,
}
struct BulkConnectionOptions {
websocket_url: String,
http_client_factory: HttpClientFactory,
}
impl BulkConnectionOptions {
async fn connect(
&self,
token: &str,
) -> Result<(ConnectionReader, ConnectionWriter), ConnectionError> {
let control_uri = self.websocket_url.parse::<Uri>().map_err(|error| {
ConnectionError::Other(format!(
"failed to build code-mode host bulk websocket URL: {error}"
))
})?;
let bulk_path = format!("{}/bulk/{token}", control_uri.path().trim_end_matches('/'));
let bulk_path_and_query = match control_uri.query() {
Some(query) => format!("{bulk_path}?{query}"),
None => bulk_path,
};
let mut bulk_uri_parts = control_uri.into_parts();
bulk_uri_parts.path_and_query = Some(bulk_path_and_query.parse().map_err(|error| {
ConnectionError::Other(format!(
"failed to build code-mode host bulk websocket path: {error}"
))
})?);
let bulk_url = Uri::from_parts(bulk_uri_parts)
.map_err(|error| {
ConnectionError::Other(format!(
"failed to build code-mode host bulk websocket URL: {error}"
))
})?
.to_string();
connect_websocket_transport(&bulk_url, &self.http_client_factory)
.await
.map_err(|error| ConnectionError::BulkConnectionUnavailable(error.to_string()))
}
}
impl CallerCancellation {
fn new() -> Self {
Self {
@@ -202,6 +249,7 @@ impl Connection {
ConnectionReader::Stdio(FramedReader::new(stdout)),
ConnectionWriter::Stdio(FramedWriter::new(stdin)),
ConnectionOwner::Process(Box::new(child)),
/*bulk_connection_options*/ None,
)
.await
}
@@ -210,38 +258,17 @@ impl Connection {
websocket_url: &str,
http_client_factory: &HttpClientFactory,
) -> Result<Self, ConnectionError> {
let request = websocket_url.into_client_request().map_err(|error| {
ConnectionError::Other(format!(
"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 websocket_config = WebSocketConfig::default()
.max_frame_size(Some(MAX_WEBSOCKET_FRAME_BYTES))
.max_message_size(Some(MAX_WEBSOCKET_FRAME_BYTES));
let (websocket, _) = tokio::time::timeout(
HOST_HANDSHAKE_TIMEOUT,
connector.connect(request, websocket_config),
)
.await
.map_err(|_| {
ConnectionError::Other("timed out connecting to the code-mode host websocket".into())
})?
.map_err(|error| {
ConnectionError::Other(format!(
"failed to connect to the code-mode host websocket: {error}"
))
})?;
let (writer, reader) = websocket.split();
let (reader, writer) =
connect_websocket_transport(websocket_url, http_client_factory).await?;
Self::establish(
ConnectionReader::WebSocket(reader),
ConnectionWriter::WebSocket(writer),
reader,
writer,
ConnectionOwner::WebSocket,
Some(BulkConnectionOptions {
websocket_url: websocket_url.to_string(),
http_client_factory: http_client_factory.clone(),
}),
)
.await
}
@@ -250,13 +277,22 @@ impl Connection {
mut reader: ConnectionReader,
mut writer: ConnectionWriter,
mut owner: ConnectionOwner,
bulk_connection_options: Option<BulkConnectionOptions>,
) -> Result<Self, ConnectionError> {
let handshake = async {
let dual_capability =
Capability::new(DUAL_WEBSOCKET_CAPABILITY).map_err(|error| error.to_string())?;
let optional_capabilities = if bulk_connection_options.is_some() {
CapabilitySet::try_new([dual_capability.clone()])
.map_err(|error| error.to_string())?
} else {
CapabilitySet::empty()
};
let hello = ClientHello::new(
SupportedProtocolVersions::try_new([ProtocolVersion::V1])
.map_err(|err| err.to_string())?,
CapabilitySet::empty(),
CapabilitySet::empty(),
optional_capabilities,
)
.map_err(|err| err.to_string())?;
writer
@@ -271,7 +307,20 @@ impl Connection {
Some(HostToClient::HostHello(hello))
if hello.selected_version() == ProtocolVersion::V1 =>
{
Ok(())
if hello.capabilities().contains(&dual_capability) {
hello
.bulk_connection_token()
.map(str::to_string)
.ok_or_else(|| {
"code-mode host advertised dual websockets without a pairing token"
.to_string()
})
.map(Some)
} else if hello.bulk_connection_token().is_some() {
Err("code-mode host returned an unexpected bulk pairing token".to_string())
} else {
Ok(None)
}
}
Some(HostToClient::HandshakeRejected { reason }) => {
Err(format!("code-mode host rejected the handshake: {reason:?}"))
@@ -292,51 +341,85 @@ impl Connection {
));
}
};
if let Err(err) = handshake_result {
let _ = writer.close().await;
owner.close().await;
return Err(ConnectionError::Other(err));
}
let bulk_token = match handshake_result {
Ok(token) => token,
Err(err) => {
let _ = writer.close().await;
owner.close().await;
return Err(ConnectionError::Other(err));
}
};
let (bulk_reader, bulk_writer) = if let Some(token) = bulk_token {
let Some(options) = bulk_connection_options else {
let _ = writer.close().await;
owner.close().await;
return Err(ConnectionError::Other(
"code-mode host negotiated an unsupported bulk websocket".to_string(),
));
};
match options.connect(&token).await {
Ok((reader, writer)) => (Some(reader), Some(writer)),
Err(error) => {
let _ = writer.close().await;
owner.close().await;
return Err(error);
}
}
} else {
(None, None)
};
let (command_tx, command_rx) = mpsc::channel(IPC_CHANNEL_CAPACITY);
let (event_tx, event_rx) = mpsc::channel(IPC_CHANNEL_CAPACITY);
let (outgoing_tx, mut outgoing_rx) = mpsc::channel::<EncodedFrame>(IPC_CHANNEL_CAPACITY);
let (outgoing_tx, outgoing_rx) = mpsc::channel::<EncodedFrame>(IPC_CHANNEL_CAPACITY);
let (bulk_tx, bulk_rx) = if bulk_writer.is_some() {
let (sender, receiver) = mpsc::channel::<EncodedFrame>(MAX_PENDING_DELEGATE_CALLS);
(Some(sender), Some(receiver))
} else {
(None, None)
};
let dual_websocket = bulk_writer.is_some();
let cancellation = CancellationToken::new();
let alive = Arc::new(AtomicBool::new(true));
let failure = Arc::new(std::sync::Mutex::new(None));
let writer_cancellation = cancellation.clone();
let writer_task = tokio::spawn(async move {
loop {
tokio::select! {
_ = writer_cancellation.cancelled() => {
return writer
.close()
.await
.map_err(|error| format!("failed to close code-mode host connection: {error}"));
}
frame = outgoing_rx.recv() => {
let Some(frame) = frame else {
return Err("code-mode host outgoing stream closed".to_string());
};
let result = tokio::select! {
_ = writer_cancellation.cancelled() => return Ok(()),
result = writer.write_frame(frame) => result,
};
if let Err(err) = result {
return Err(format!("failed to write code-mode host message: {err}"));
}
}
}
if let (Some(bulk_writer), Some(bulk_rx)) = (bulk_writer, bulk_rx) {
tokio::try_join!(
drive_writer(writer, outgoing_rx, writer_cancellation.clone()),
drive_writer(bulk_writer, bulk_rx, writer_cancellation)
)?;
Ok(())
} else {
drive_writer(writer, outgoing_rx, writer_cancellation).await
}
});
let reader_events = event_tx.clone();
let reader_cancellation = cancellation.clone();
let reader_task =
tokio::spawn(
async move { drive_reader(reader, reader_events, reader_cancellation).await },
);
let reader_task = tokio::spawn(async move {
let lane = dual_websocket.then_some(TransportLane::Control);
if let Some(bulk_reader) = bulk_reader {
tokio::try_join!(
drive_reader(
reader,
reader_events.clone(),
reader_cancellation.clone(),
lane,
),
drive_reader(
bulk_reader,
reader_events,
reader_cancellation,
Some(TransportLane::Bulk),
)
)?;
Ok(())
} else {
drive_reader(reader, reader_events, reader_cancellation, lane).await
}
});
let (driver, execute_claim_tx) = ConnectionDriver::new(
command_rx,
@@ -349,6 +432,10 @@ impl Connection {
cancellation: cancellation.clone(),
},
);
let driver = match bulk_tx {
Some(sender) => driver.with_bulk_sender(sender),
None => driver,
};
let driver_task = tokio::spawn(driver.run());
tokio::spawn(
ConnectionSupervisor {
@@ -504,6 +591,73 @@ impl Connection {
}
}
async fn connect_websocket_transport(
websocket_url: &str,
http_client_factory: &HttpClientFactory,
) -> Result<(ConnectionReader, ConnectionWriter), ConnectionError> {
let request = websocket_url.into_client_request().map_err(|error| {
ConnectionError::Other(format!(
"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 websocket_config = WebSocketConfig::default()
.max_frame_size(Some(MAX_WEBSOCKET_FRAME_BYTES))
.max_message_size(Some(MAX_WEBSOCKET_FRAME_BYTES));
let (websocket, _) = tokio::time::timeout(
HOST_HANDSHAKE_TIMEOUT,
connector.connect(request, websocket_config),
)
.await
.map_err(|_| {
ConnectionError::Other("timed out connecting to the code-mode host websocket".into())
})?
.map_err(|error| {
ConnectionError::Other(format!(
"failed to connect to the code-mode host websocket: {error}"
))
})?;
let (writer, reader) = websocket.split();
Ok((
ConnectionReader::WebSocket(reader),
ConnectionWriter::WebSocket(writer),
))
}
async fn drive_writer(
mut writer: ConnectionWriter,
mut outgoing: mpsc::Receiver<EncodedFrame>,
cancellation: CancellationToken,
) -> Result<(), String> {
loop {
tokio::select! {
_ = cancellation.cancelled() => {
return writer
.close()
.await
.map_err(|error| format!("failed to close code-mode host connection: {error}"));
}
frame = outgoing.recv() => {
let Some(frame) = frame else {
return Err("code-mode host outgoing stream closed".to_string());
};
tokio::select! {
_ = cancellation.cancelled() => return Ok(()),
result = writer.write_frame(frame) => {
result.map_err(|error| {
format!("failed to write code-mode host message: {error}")
})?;
}
}
}
}
}
}
impl Drop for Connection {
fn drop(&mut self) {
mark_connection_dead(
codex-rs/code-mode/src/remote_session/connection/driver.rs
+18 -2
View File
@@ -1,3 +1,4 @@
use std::collections::VecDeque;
use std::panic::AssertUnwindSafe;
use std::sync::Arc;
use std::sync::atomic::AtomicBool;
@@ -6,7 +7,9 @@ use std::sync::atomic::Ordering;
use codex_code_mode_protocol::CellId;
use codex_code_mode_protocol::CodeModeSessionDelegate;
use codex_code_mode_protocol::host::EncodedFrame;
use codex_code_mode_protocol::host::HostToClient;
use codex_code_mode_protocol::host::RequestId;
use codex_code_mode_protocol::host::TransportLane;
use tokio::sync::mpsc;
use tokio_util::sync::CancellationToken;
@@ -39,7 +42,9 @@ pub(super) struct ConnectionDriver {
event_tx: mpsc::Sender<DriverEvent>,
execute_claim_rx: mpsc::UnboundedReceiver<RequestId>,
outgoing_tx: mpsc::Sender<EncodedFrame>,
bulk_tx: Option<mpsc::Sender<EncodedFrame>>,
requests: RequestTracker,
deferred_host_messages: VecDeque<HostToClient>,
sessions: SessionRegistry,
delegates: DelegateRuntime,
alive: Arc<AtomicBool>,
@@ -64,7 +69,9 @@ impl ConnectionDriver {
event_tx: event_tx.clone(),
execute_claim_rx,
outgoing_tx,
bulk_tx: None,
requests: RequestTracker::new(),
deferred_host_messages: VecDeque::new(),
sessions: SessionRegistry::new(),
delegates: DelegateRuntime::new(event_tx),
alive: lifecycle.alive,
@@ -130,8 +137,17 @@ impl ConnectionDriver {
}
}
fn queue_frame(&mut self, frame: EncodedFrame) -> bool {
match self.outgoing_tx.try_send(frame) {
pub(super) fn with_bulk_sender(mut self, sender: mpsc::Sender<EncodedFrame>) -> Self {
self.bulk_tx = Some(sender);
self
}
fn queue_frame(&mut self, frame: EncodedFrame, lane: TransportLane) -> bool {
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(()) => true,
Err(mpsc::error::TrySendError::Full(_)) => {
self.fail("code-mode host outgoing queue is full".to_string());
codex-rs/code-mode/src/remote_session/connection/driver/commands.rs
+6 -3
View File
@@ -105,7 +105,8 @@ impl ConnectionDriver {
},
&self.event_tx,
);
self.queue_frame(frame)
let lane = message.transport_lane();
self.queue_frame(frame, lane)
}
fn execute(
@@ -164,7 +165,8 @@ impl ConnectionDriver {
},
&self.event_tx,
);
self.queue_frame(frame)
let lane = message.transport_lane();
self.queue_frame(frame, lane)
}
fn wait(
@@ -293,6 +295,7 @@ impl ConnectionDriver {
};
self.requests
.insert_pending(request_id, pending, &self.event_tx);
self.queue_frame(frame)
let lane = message.transport_lane();
self.queue_frame(frame, lane)
}
}
codex-rs/code-mode/src/remote_session/connection/driver/delegate_runtime.rs
+3 -5
View File
@@ -261,10 +261,7 @@ impl ConnectionDriver {
};
let target = match self.sessions.delegate_target(&session_id, wire_cell_id) {
Ok(target) => target,
Err(err) => {
self.fail(err);
return false;
}
Err(err) => return self.send_delegate_response(id, Err(err)),
};
match self.delegates.start(id, target, request) {
Ok(()) => true,
@@ -321,7 +318,8 @@ impl ConnectionDriver {
}
}
};
self.queue_frame(frame)
let lane = message.transport_lane();
self.queue_frame(frame, lane)
}
pub(super) fn close_cell(&mut self, session_id: SessionId, cell_id: WireCellId) -> bool {
codex-rs/code-mode/src/remote_session/connection/driver/request_tracker.rs
+9
View File
@@ -52,6 +52,15 @@ impl RequestTracker {
})
}
pub(super) fn has_pending_execute_for_session(&self, session_id: &SessionId) -> bool {
self.pending.values().any(|request| {
matches!(
request,
PendingRequest::Execute { session, .. } if session.id == *session_id
)
})
}
pub(super) fn allocate_id(&mut self) -> Result<RequestId, String> {
let id = self.next_request_id;
self.next_request_id = self
codex-rs/code-mode/src/remote_session/connection/driver/responses.rs
+66 -2
View File
@@ -56,6 +56,62 @@ impl ConnectionDriver {
}
pub(super) fn handle_host_message(&mut self, message: HostToClient) -> bool {
if self.should_defer_host_message(&message) {
if self.deferred_host_messages.len()
>= codex_code_mode_protocol::host::MAX_PENDING_DELEGATE_CALLS
{
self.fail(
"code-mode host exceeded deferred cross-socket message limit".to_string(),
);
return false;
}
self.deferred_host_messages.push_back(message);
return true;
}
if !self.dispatch_host_message(message) {
return false;
}
for _ in 0..self.deferred_host_messages.len() {
let Some(message) = self.deferred_host_messages.pop_front() else {
break;
};
if self.should_defer_host_message(&message) {
self.deferred_host_messages.push_back(message);
} else if !self.dispatch_host_message(message) {
return false;
}
}
true
}
fn should_defer_host_message(&self, message: &HostToClient) -> bool {
match message {
HostToClient::DelegateRequest {
session_id,
request,
..
} => {
let cell_id = match request {
codex_code_mode_protocol::host::DelegateRequest::InvokeTool { invocation } => {
&invocation.cell_id
}
codex_code_mode_protocol::host::DelegateRequest::Notify { cell_id, .. } => {
cell_id
}
};
!self.sessions.contains_cell(session_id, cell_id)
&& self.requests.has_pending_execute_for_session(session_id)
}
HostToClient::Response { .. }
| HostToClient::InitialResponse { .. }
| HostToClient::CellClosed { .. }
| HostToClient::CancelDelegateRequest { .. }
| HostToClient::HostHello(_)
| HostToClient::HandshakeRejected { .. } => false,
}
}
fn dispatch_host_message(&mut self, message: HostToClient) -> bool {
match message {
HostToClient::Response { id, result } => {
self.complete_request(id, result.into_result())
@@ -69,6 +125,12 @@ impl ConnectionDriver {
request,
} => self.start_delegate(id, session_id, request),
HostToClient::CancelDelegateRequest { id } => {
self.deferred_host_messages.retain(|message| {
!matches!(
message,
HostToClient::DelegateRequest { id: deferred_id, .. } if *deferred_id == id
)
});
self.delegates.cancel(id);
true
}
@@ -302,7 +364,8 @@ impl ConnectionDriver {
}
fn send_cancel_request(&mut self, id: RequestId) -> bool {
let frame = match EncodedFrame::encode(&ClientToHost::CancelRequest { id }) {
let message = ClientToHost::CancelRequest { id };
let frame = match EncodedFrame::encode(&message) {
Ok(frame) => frame,
Err(err) => {
self.fail(format!(
@@ -311,7 +374,8 @@ impl ConnectionDriver {
return false;
}
};
self.queue_frame(frame)
let lane = message.transport_lane();
self.queue_frame(frame, lane)
}
fn shutdown_abandoned_session(&mut self, session: RemoteSession) -> bool {
codex-rs/code-mode/src/remote_session/connection/driver/session_registry.rs
+6
View File
@@ -61,6 +61,12 @@ impl SessionRegistry {
self.records.contains_key(session_id)
}
pub(super) fn contains_cell(&self, session_id: &SessionId, cell_id: &WireCellId) -> bool {
self.records
.get(session_id)
.is_some_and(|session| session.cells.contains_key(cell_id))
}
pub(super) fn insert_ready(
&mut self,
session: RemoteSession,
codex-rs/code-mode/src/remote_session/connection/driver_tests.rs
+300 -9
View File
@@ -15,6 +15,7 @@ use codex_code_mode_protocol::WaitRequest;
use codex_code_mode_protocol::host::ClientToHost;
use codex_code_mode_protocol::host::DelegateRequest;
use codex_code_mode_protocol::host::DelegateRequestId;
use codex_code_mode_protocol::host::DelegateResponse;
use codex_code_mode_protocol::host::EncodedFrame;
use codex_code_mode_protocol::host::HostResponse;
use codex_code_mode_protocol::host::HostToClient;
@@ -193,6 +194,11 @@ struct RecordingDelegate {
struct PanickingDelegate;
struct LargeResultBurstDelegate {
started: AtomicUsize,
release: CancellationToken,
}
#[derive(Debug, Eq, PartialEq)]
enum HeldDelegateEvent {
Started,
@@ -282,6 +288,35 @@ impl CodeModeSessionDelegate for PanickingDelegate {
fn cell_closed(&self, _cell_id: &CellId) {}
}
impl CodeModeSessionDelegate for LargeResultBurstDelegate {
fn invoke_tool<'a>(
&'a self,
_invocation: CodeModeNestedToolCall,
cancellation_token: CancellationToken,
) -> ToolInvocationFuture<'a> {
self.started.fetch_add(1, Ordering::Release);
let release = self.release.clone();
Box::pin(async move {
tokio::select! {
_ = cancellation_token.cancelled() => Err("cancelled".to_string()),
_ = release.cancelled() => Ok("x".repeat(256 * 1024).into()),
}
})
}
fn notify<'a>(
&'a self,
_call_id: String,
_cell_id: CellId,
_text: String,
_cancellation_token: CancellationToken,
) -> NotificationFuture<'a> {
Box::pin(async { Ok(()) })
}
fn cell_closed(&self, _cell_id: &CellId) {}
}
impl CodeModeSessionDelegate for RecordingDelegate {
fn invoke_tool<'a>(
&'a self,
@@ -330,6 +365,143 @@ async fn next_held_delegate_event(
.expect("delegate event stream")
}
#[tokio::test]
async fn deferred_delegates_follow_cell_readiness_and_cancellation() {
let mut harness = DriverHarness::start();
let session = remote_session();
let delegate = Arc::new(RecordingDelegate::default());
harness.open(session.clone(), delegate.clone()).await;
let (first_response_tx, first_response_rx) = oneshot::channel();
let (second_response_tx, second_response_rx) = oneshot::channel();
for (tool_call_id, response_tx) in [
("first-cell", first_response_tx),
("second-cell", second_response_tx),
] {
harness
.command_tx
.send(DriverCommand::Execute {
session: session.clone(),
request: ExecuteRequest {
tool_call_id: tool_call_id.to_string(),
enabled_tools: Vec::new(),
source: "text('done')".to_string(),
yield_time_ms: Some(/*yield_time_ms*/ 1),
max_output_tokens: None,
},
caller_cancellation: CancellationToken::new(),
response_tx,
})
.await
.expect("execute command");
harness.outgoing_rx.recv().await.expect("execute frame");
}
let first_cell_id = CellId::new("first-cell".to_string());
let second_cell_id = CellId::new("second-cell".to_string());
let first_delegate_id = DelegateRequestId::new(/*value*/ 7);
let second_delegate_id = DelegateRequestId::new(/*value*/ 8);
let cancelled_delegate_id = DelegateRequestId::new(/*value*/ 9);
for (delegate_id, cell_id) in [
(first_delegate_id, first_cell_id.clone()),
(second_delegate_id, second_cell_id.clone()),
(cancelled_delegate_id, first_cell_id.clone()),
] {
harness
.event_tx
.send(DriverEvent::HostMessage(HostToClient::DelegateRequest {
id: delegate_id,
session_id: session.id.clone(),
request: DelegateRequest::Notify {
call_id: format!("notify-{}", cell_id.as_str()),
cell_id: (&cell_id).into(),
text: "hello".to_string(),
},
}))
.await
.expect("early delegate request");
}
harness
.event_tx
.send(DriverEvent::HostMessage(
HostToClient::CancelDelegateRequest {
id: cancelled_delegate_id,
},
))
.await
.expect("deferred delegate cancellation");
harness
.event_tx
.send(DriverEvent::HostMessage(HostToClient::Response {
id: RequestId::new(/*value*/ 3),
result: WireResult::Ok {
value: HostResponse::ExecutionStarted {
cell_id: (&second_cell_id).into(),
},
},
}))
.await
.expect("second execution-started response");
let _second_started = second_response_rx
.await
.expect("second execute response")
.expect("second started cell");
let second_response = tokio::time::timeout(Duration::from_secs(1), harness.outgoing_rx.recv())
.await
.expect("second delegate response timeout")
.expect("second delegate response frame");
assert_eq!(
EncodedFrame::decode_framed::<ClientToHost>(&second_response.into_framed_bytes())
.expect("decode second delegate response"),
ClientToHost::DelegateResponse {
id: second_delegate_id,
result: WireResult::Ok {
value: codex_code_mode_protocol::host::DelegateResponse::NotificationDelivered,
},
}
);
assert_eq!(delegate.notifications.load(Ordering::Relaxed), 1);
harness
.event_tx
.send(DriverEvent::HostMessage(HostToClient::Response {
id: RequestId::new(/*value*/ 2),
result: WireResult::Ok {
value: HostResponse::ExecutionStarted {
cell_id: (&first_cell_id).into(),
},
},
}))
.await
.expect("first execution-started response");
let _first_started = first_response_rx
.await
.expect("first execute response")
.expect("first started cell");
let first_response = tokio::time::timeout(Duration::from_secs(1), harness.outgoing_rx.recv())
.await
.expect("first delegate response timeout")
.expect("first delegate response frame");
assert_eq!(
EncodedFrame::decode_framed::<ClientToHost>(&first_response.into_framed_bytes())
.expect("decode first delegate response"),
ClientToHost::DelegateResponse {
id: first_delegate_id,
result: WireResult::Ok {
value: codex_code_mode_protocol::host::DelegateResponse::NotificationDelivered,
},
}
);
assert_eq!(delegate.notifications.load(Ordering::Relaxed), 2);
assert!(matches!(
harness.outgoing_rx.try_recv(),
Err(mpsc::error::TryRecvError::Empty)
));
assert!(harness.alive.load(Ordering::Relaxed));
}
#[tokio::test]
async fn dropped_open_waiter_shuts_down_committed_session() {
let mut harness = DriverHarness::start();
@@ -465,6 +637,95 @@ async fn delegate_cancel_is_best_effort_and_sends_no_late_response() {
assert_eq!(delegate.notifications.load(Ordering::Relaxed), 0);
}
#[tokio::test]
async fn concurrent_large_delegate_results_do_not_disconnect_a_backpressured_bulk_lane() {
const CONCURRENT_RESULTS: usize = 129;
let (command_tx, command_rx) = mpsc::channel(/*max_capacity*/ 16);
let (event_tx, event_rx) = mpsc::channel(/*max_capacity*/ 16);
let (outgoing_tx, outgoing_rx) = mpsc::channel(/*max_capacity*/ 16);
let (bulk_tx, mut bulk_rx) = mpsc::channel(MAX_PENDING_DELEGATE_CALLS);
let cancellation = CancellationToken::new();
let alive = Arc::new(AtomicBool::new(true));
let (driver, execute_claim_tx) = ConnectionDriver::new(
command_rx,
event_rx,
event_tx.clone(),
outgoing_tx,
DriverLifecycle {
alive: Arc::clone(&alive),
failure: Arc::new(StdMutex::new(None)),
cancellation: cancellation.clone(),
},
);
let driver_task = tokio::spawn(driver.with_bulk_sender(bulk_tx).run());
let mut harness = DriverHarness {
command_tx,
event_tx,
execute_claim_tx,
outgoing_rx,
cancellation,
alive,
driver_task,
};
let session = remote_session();
let delegate = Arc::new(LargeResultBurstDelegate {
started: AtomicUsize::new(0),
release: CancellationToken::new(),
});
harness.open(session.clone(), delegate.clone()).await;
let _started = harness
.start_cell(session.clone(), /*request_id*/ 2, "1")
.await;
for value in 1..=CONCURRENT_RESULTS {
harness
.start_tool_delegate(&session, DelegateRequestId::new(value as i64))
.await;
}
tokio::time::timeout(Duration::from_secs(10), async {
while delegate.started.load(Ordering::Acquire) < CONCURRENT_RESULTS {
tokio::task::yield_now().await;
}
})
.await
.expect("concurrent delegate calls should all start");
delegate.release.cancel();
tokio::time::timeout(Duration::from_secs(10), async {
while bulk_rx.len() < CONCURRENT_RESULTS {
assert!(
harness.alive.load(Ordering::Acquire),
"bulk queue disconnected before accepting all concurrent tool results"
);
tokio::task::yield_now().await;
}
})
.await
.expect("concurrent large results should queue behind the blocked bulk writer");
let _unrelated = harness.start_cell(session, /*request_id*/ 3, "2").await;
assert!(harness.alive.load(Ordering::Acquire));
for _ in 0..CONCURRENT_RESULTS {
let frame = bulk_rx.recv().await.expect("queued bulk delegate result");
let message = EncodedFrame::decode_framed::<ClientToHost>(&frame.into_framed_bytes())
.expect("decode queued delegate result");
let ClientToHost::DelegateResponse {
result:
WireResult::Ok {
value: DelegateResponse::ToolResult { result },
},
..
} = message
else {
panic!("expected a successful large delegate result");
};
assert_eq!(result.as_str().map(str::len), Some(256 * 1024));
}
assert!(harness.alive.load(Ordering::Acquire));
}
#[tokio::test]
async fn delegate_limit_returns_an_error_without_disconnecting() {
let mut harness = DriverHarness::start();
@@ -796,16 +1057,17 @@ async fn delegate_task_panic_becomes_tool_error_without_killing_connection() {
}
#[tokio::test]
async fn delegate_for_unknown_cell_fails_connection_without_invocation() {
async fn delegate_for_unknown_cell_returns_error_without_invocation() {
let mut harness = DriverHarness::start();
let session = remote_session();
let delegate = Arc::new(RecordingDelegate::default());
harness.open(session.clone(), delegate.clone()).await;
let id = DelegateRequestId::new(/*value*/ 7);
harness
.event_tx
.send(DriverEvent::HostMessage(HostToClient::DelegateRequest {
id: DelegateRequestId::new(/*value*/ 7),
id,
session_id: session.id,
request: DelegateRequest::InvokeTool {
invocation: WireNestedToolCall {
@@ -819,14 +1081,28 @@ async fn delegate_for_unknown_cell_fails_connection_without_invocation() {
}))
.await
.expect("delegate request");
tokio::task::yield_now().await;
let response = tokio::time::timeout(Duration::from_secs(1), harness.outgoing_rx.recv())
.await
.expect("delegate response timeout")
.expect("delegate response frame");
assert!(!harness.alive.load(Ordering::Acquire));
assert_eq!(
EncodedFrame::decode_framed::<ClientToHost>(&response.into_framed_bytes())
.expect("decode delegate response"),
ClientToHost::DelegateResponse {
id,
result: WireResult::Err {
message: "code-mode host delegated for unknown cell missing in session session-1"
.to_string(),
},
}
);
assert!(harness.alive.load(Ordering::Acquire));
assert_eq!(delegate.invocations.load(Ordering::Relaxed), 0);
}
#[tokio::test]
async fn delegate_after_cell_close_fails_connection_without_invocation() {
async fn delegate_after_cell_close_returns_error_without_invocation() {
let mut harness = DriverHarness::start();
let session = remote_session();
let delegate = Arc::new(RecordingDelegate::default());
@@ -842,10 +1118,11 @@ async fn delegate_after_cell_close_fails_connection_without_invocation() {
}))
.await
.expect("cell close");
let id = DelegateRequestId::new(/*value*/ 7);
harness
.event_tx
.send(DriverEvent::HostMessage(HostToClient::DelegateRequest {
id: DelegateRequestId::new(/*value*/ 7),
id,
session_id: session.id,
request: DelegateRequest::Notify {
call_id: "notify-1".to_string(),
@@ -855,10 +1132,24 @@ async fn delegate_after_cell_close_fails_connection_without_invocation() {
}))
.await
.expect("delegate request");
tokio::task::yield_now().await;
let response = tokio::time::timeout(Duration::from_secs(1), harness.outgoing_rx.recv())
.await
.expect("delegate response timeout")
.expect("delegate response frame");
assert!(!harness.alive.load(Ordering::Acquire));
assert_eq!(delegate.invocations.load(Ordering::Relaxed), 0);
assert_eq!(
EncodedFrame::decode_framed::<ClientToHost>(&response.into_framed_bytes())
.expect("decode delegate response"),
ClientToHost::DelegateResponse {
id,
result: WireResult::Err {
message: "code-mode host delegated for unknown cell 1 in session session-1"
.to_string(),
},
}
);
assert!(harness.alive.load(Ordering::Acquire));
assert_eq!(delegate.notifications.load(Ordering::Relaxed), 0);
}
#[tokio::test]
codex-rs/code-mode/src/remote_session/connection/reader.rs
+7
View File
@@ -1,3 +1,4 @@
use codex_code_mode_protocol::host::TransportLane;
use tokio::sync::mpsc;
use tokio_util::sync::CancellationToken;
@@ -8,6 +9,7 @@ pub(super) async fn drive_reader(
mut reader: ConnectionReader,
events: mpsc::Sender<DriverEvent>,
cancellation: CancellationToken,
lane: Option<TransportLane>,
) -> Result<(), String> {
loop {
let message = tokio::select! {
@@ -19,6 +21,11 @@ pub(super) async fn drive_reader(
Ok(None) => return Err("code-mode host closed its stdout".to_string()),
Err(err) => return Err(format!("failed to read code-mode host message: {err}")),
};
if let Some(lane) = lane
&& !message.allows_transport_lane(lane)
{
return Err("code-mode host sent a message on the wrong websocket lane".to_string());
}
events
.send(DriverEvent::HostMessage(message))
.await
codex-rs/code-mode/src/remote_session_tests.rs
+89
View File
@@ -7,8 +7,10 @@ use codex_code_mode_protocol::CodeModeSessionProvider;
use codex_code_mode_protocol::ExecuteRequest;
use codex_code_mode_protocol::FunctionCallOutputContentItem;
use codex_code_mode_protocol::RuntimeResponse;
use codex_code_mode_protocol::host::Capability;
use codex_code_mode_protocol::host::CapabilitySet;
use codex_code_mode_protocol::host::ClientToHost;
use codex_code_mode_protocol::host::DUAL_WEBSOCKET_CAPABILITY;
use codex_code_mode_protocol::host::EncodedFrame;
use codex_code_mode_protocol::host::HostHello;
use codex_code_mode_protocol::host::HostRequest;
@@ -24,6 +26,8 @@ use codex_http_client::OutboundProxyPolicy;
use futures::SinkExt;
use futures::StreamExt;
use pretty_assertions::assert_eq;
use tokio::io::AsyncReadExt;
use tokio::io::AsyncWriteExt;
use tokio::net::TcpListener;
use tokio::time::timeout;
use tokio_tungstenite::accept_async;
@@ -224,6 +228,91 @@ async fn websocket_provider_executes_over_shared_connector() {
.expect("websocket test host task should succeed");
}
#[tokio::test]
async fn websocket_provider_fails_when_a_negotiated_bulk_connection_is_unavailable() {
let listener = TcpListener::bind("127.0.0.1:0")
.await
.expect("websocket test listener should bind");
let websocket_url = format!(
"ws://{}/?access_token=shared-token",
listener
.local_addr()
.expect("websocket test listener should have an address")
);
let server = tokio::spawn(async move {
let (stream, _) = listener
.accept()
.await
.expect("websocket test host should accept the first connection");
let mut control = accept_async(stream)
.await
.expect("first control websocket should connect");
let frame = control
.next()
.await
.expect("first client hello")
.expect("first websocket frame")
.into_data();
let ClientToHost::ClientHello(hello) =
EncodedFrame::decode_framed(&frame).expect("decode first client hello")
else {
panic!("expected first client hello");
};
let capability =
Capability::new(DUAL_WEBSOCKET_CAPABILITY).expect("dual websocket capability");
assert!(hello.optional_capabilities().contains(&capability));
let hello = HostToClient::HostHello(
HostHello::new(
ProtocolVersion::V1,
CapabilitySet::try_new([capability.clone()]).expect("host capabilities"),
)
.with_bulk_connection_token("fallback-token".to_string()),
);
let frame = EncodedFrame::encode(&hello).expect("encode dual host hello");
control
.send(Message::Binary(frame.into_framed_bytes().into()))
.await
.expect("send dual host hello");
let (mut bulk, _) = listener
.accept()
.await
.expect("websocket test host should accept the bulk connection");
let mut request = [0_u8; 1024];
let request_len = bulk
.read(&mut request)
.await
.expect("read bulk websocket handshake");
let request = std::str::from_utf8(&request[..request_len]).expect("bulk HTTP request");
assert!(
request.starts_with("GET /bulk/fallback-token?access_token=shared-token HTTP/1.1\r\n")
);
bulk.write_all(b"HTTP/1.1 404 Not Found\r\nContent-Length: 0\r\nConnection: close\r\n\r\n")
.await
.expect("reject unavailable bulk websocket");
drop(bulk);
drop(control);
});
let provider = WebSocketCodeModeSessionProvider::new(websocket_url);
let error = match provider
.create_session(Arc::new(NoopCodeModeSessionDelegate))
.await
{
Ok(_) => panic!("provider should reject an unavailable negotiated bulk websocket"),
Err(error) => error,
};
assert!(
error.contains("404"),
"unexpected negotiated bulk websocket error: {error}"
);
drop(provider);
timeout(Duration::from_secs(5), server)
.await
.expect("negotiated bulk websocket test host should disconnect promptly")
.expect("negotiated bulk websocket test host task should succeed");
}
#[tokio::test]
async fn provider_returns_missing_host_error() {
let provider = ProcessOwnedCodeModeSessionProvider::with_host_program(