mirror of
https://github.com/openai/codex.git
synced 2026-10-20 17:38:45 +00:00
Remove WebSocket transport from code-mode hosts (#40692)
## What changed - Accept only `http://` and `https://` gRPC endpoints for app server `--code-mode-host` connections. - Limit the standalone code-mode host to stdio and gRPC listeners. - Remove the WebSocket session provider, dual-WebSocket negotiation, transport lanes, and related dependencies and tests. - Keep coverage for shared and prewarmed gRPC hosts and reject WebSocket host URLs in CLI parsing tests. GitOrigin-RevId: 8792604e79ffd54ce776074a4dc3f7d55e63752a
This commit is contained in:
Ruslan Nigmatullin
committed by
copyberry
copyberry
parent
9be8d6e1c3
commit
725b3a44f8
@@ -18,7 +18,7 @@ workspace = true
|
||||
|
||||
[dependencies]
|
||||
anyhow = { workspace = true }
|
||||
axum = { workspace = true, features = ["http1", "tokio", "ws"] }
|
||||
axum = { workspace = true, features = ["http1", "tokio"] }
|
||||
clap = { workspace = true, features = ["derive"] }
|
||||
codex-code-mode-protocol = { workspace = true }
|
||||
codex-code-mode-runtime = { workspace = true }
|
||||
@@ -39,4 +39,3 @@ codex-code-mode = { workspace = true }
|
||||
codex-utils-cargo-bin = { workspace = true }
|
||||
pretty_assertions = { workspace = true }
|
||||
tempfile = { workspace = true }
|
||||
tokio-tungstenite = { workspace = true }
|
||||
|
||||
@@ -13,7 +13,6 @@ 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;
|
||||
@@ -36,7 +35,6 @@ use tokio::sync::mpsc;
|
||||
use tokio::sync::oneshot;
|
||||
use tokio_util::sync::CancellationToken;
|
||||
use tokio_util::task::TaskTracker;
|
||||
use uuid::Uuid;
|
||||
|
||||
use super::HostLimits;
|
||||
use super::HostState;
|
||||
@@ -45,12 +43,8 @@ use super::MAX_RECENT_REQUEST_IDS;
|
||||
use super::RequestKind;
|
||||
use super::RequestRegistry;
|
||||
use super::SeenSessionIds;
|
||||
use super::negotiate;
|
||||
use super::peer::HostPeer;
|
||||
use super::run;
|
||||
use super::transport::BulkConnectionRegistry;
|
||||
use super::transport::ConnectionReader;
|
||||
use super::transport::ConnectionWriter;
|
||||
|
||||
fn client_hello(
|
||||
versions: impl IntoIterator<Item = ProtocolVersion>,
|
||||
@@ -260,111 +254,8 @@ impl AsyncWrite for BlockingWriter {
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn failed_host_hello_removes_the_bulk_pairing_reservation() {
|
||||
let (host_reader, client_writer) = tokio::io::duplex(/*max_buf_size*/ 4096);
|
||||
let capability = Capability::new(DUAL_WEBSOCKET_CAPABILITY).expect("dual websocket capability");
|
||||
let hello = ClientToHost::ClientHello(
|
||||
ClientHello::new(
|
||||
SupportedProtocolVersions::try_new([ProtocolVersion::V1]).expect("supported versions"),
|
||||
CapabilitySet::empty(),
|
||||
CapabilitySet::try_new([capability]).expect("optional capabilities"),
|
||||
)
|
||||
.expect("client hello"),
|
||||
);
|
||||
FramedWriter::new(client_writer)
|
||||
.write(&hello)
|
||||
.await
|
||||
.expect("write client hello");
|
||||
|
||||
let bytes = Arc::new(Mutex::new(Vec::new()));
|
||||
let registry = BulkConnectionRegistry::default();
|
||||
let mut reader = ConnectionReader::from_reader(host_reader);
|
||||
let mut writer = ConnectionWriter::from_writer(FailingHandshakeWriter {
|
||||
bytes: Arc::clone(&bytes),
|
||||
});
|
||||
let result = negotiate(&mut reader, &mut writer, Some(®istry)).await;
|
||||
assert!(result.is_err());
|
||||
|
||||
let message = EncodedFrame::decode_framed::<HostToClient>(
|
||||
&bytes.lock().unwrap_or_else(PoisonError::into_inner),
|
||||
)
|
||||
.expect("decode failed host hello");
|
||||
let HostToClient::HostHello(hello) = message else {
|
||||
panic!("expected a host hello");
|
||||
};
|
||||
let token = hello
|
||||
.bulk_connection_token()
|
||||
.expect("dual websocket pairing token");
|
||||
let token = Uuid::parse_str(token).expect("UUID pairing token");
|
||||
assert!(registry.remove(token).is_none());
|
||||
}
|
||||
|
||||
struct FailingHandshakeWriter {
|
||||
bytes: Arc<Mutex<Vec<u8>>>,
|
||||
}
|
||||
|
||||
impl AsyncWrite for FailingHandshakeWriter {
|
||||
fn poll_write(
|
||||
self: Pin<&mut Self>,
|
||||
_cx: &mut Context<'_>,
|
||||
bytes: &[u8],
|
||||
) -> Poll<std::io::Result<usize>> {
|
||||
self.bytes
|
||||
.lock()
|
||||
.unwrap_or_else(PoisonError::into_inner)
|
||||
.extend_from_slice(bytes);
|
||||
Poll::Ready(Ok(bytes.len()))
|
||||
}
|
||||
|
||||
fn poll_flush(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<std::io::Result<()>> {
|
||||
Poll::Ready(Err(std::io::Error::other("host hello write failed")))
|
||||
}
|
||||
|
||||
fn poll_shutdown(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<std::io::Result<()>> {
|
||||
Poll::Ready(Ok(()))
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn optional_dual_websocket_capability_falls_back_to_a_single_connection() {
|
||||
let (host_stream, client_stream) = tokio::io::duplex(/*max_buf_size*/ 1024);
|
||||
let (host_reader, host_writer) = tokio::io::split(host_stream);
|
||||
let (client_reader, client_writer) = tokio::io::split(client_stream);
|
||||
let host = tokio::spawn(run(host_reader, host_writer));
|
||||
let mut reader = FramedReader::new(client_reader);
|
||||
let mut writer = FramedWriter::new(client_writer);
|
||||
let capability = Capability::new(DUAL_WEBSOCKET_CAPABILITY).expect("dual websocket capability");
|
||||
|
||||
writer
|
||||
.write(&ClientToHost::ClientHello(
|
||||
ClientHello::new(
|
||||
SupportedProtocolVersions::try_new([ProtocolVersion::V1])
|
||||
.expect("supported versions"),
|
||||
CapabilitySet::empty(),
|
||||
CapabilitySet::try_new([capability]).expect("optional capabilities"),
|
||||
)
|
||||
.expect("client hello"),
|
||||
))
|
||||
.await
|
||||
.expect("write hello");
|
||||
assert_eq!(
|
||||
reader.read::<HostToClient>().await.expect("host hello"),
|
||||
Some(HostToClient::HostHello(HostHello::new(
|
||||
ProtocolVersion::V1,
|
||||
CapabilitySet::empty(),
|
||||
)))
|
||||
);
|
||||
|
||||
drop(writer);
|
||||
drop(reader);
|
||||
host.await.expect("host task").expect("host connection");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn session_resource_limits_are_negotiated_when_optional_or_required() {
|
||||
let dual_capability =
|
||||
Capability::new(DUAL_WEBSOCKET_CAPABILITY).expect("dual websocket capability");
|
||||
let resource_limits_capability =
|
||||
Capability::new(SESSION_RESOURCE_LIMITS_CAPABILITY).expect("resource limits capability");
|
||||
let resource_limits =
|
||||
@@ -373,8 +264,7 @@ async fn session_resource_limits_are_negotiated_when_optional_or_required() {
|
||||
for (required, optional) in [
|
||||
(
|
||||
CapabilitySet::empty(),
|
||||
CapabilitySet::try_new([dual_capability, resource_limits_capability])
|
||||
.expect("optional capabilities"),
|
||||
CapabilitySet::try_new([resource_limits_capability]).expect("optional capabilities"),
|
||||
),
|
||||
(resource_limits.clone(), CapabilitySet::empty()),
|
||||
] {
|
||||
|
||||
@@ -14,20 +14,19 @@ use codex_code_mode_protocol::CodeModeSessionCellExecutionLimits;
|
||||
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::FramedReader;
|
||||
use codex_code_mode_protocol::host::FramedWriter;
|
||||
use codex_code_mode_protocol::host::HandshakeRejectReason;
|
||||
use codex_code_mode_protocol::host::HostHello;
|
||||
use codex_code_mode_protocol::host::HostRequest;
|
||||
use codex_code_mode_protocol::host::HostResponse;
|
||||
use codex_code_mode_protocol::host::HostToClient;
|
||||
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::SESSION_RESOURCE_LIMITS_CAPABILITY;
|
||||
use codex_code_mode_protocol::host::SessionId;
|
||||
use codex_code_mode_protocol::host::SupportedProtocolVersions;
|
||||
use codex_code_mode_protocol::host::TransportLane;
|
||||
use codex_code_mode_runtime::InProcessCodeModeSession;
|
||||
use tokio::io::AsyncRead;
|
||||
use tokio::io::AsyncWrite;
|
||||
@@ -40,10 +39,6 @@ use tokio_util::task::TaskTracker;
|
||||
|
||||
use self::delegate::RemoteDelegate;
|
||||
use self::peer::HostPeer;
|
||||
use self::transport::BulkConnectionRegistration;
|
||||
use self::transport::BulkConnectionRegistry;
|
||||
use self::transport::ConnectionReader;
|
||||
use self::transport::ConnectionWriter;
|
||||
|
||||
pub use self::grpc::GrpcCodeModeHost;
|
||||
pub use self::transport::DEFAULT_LISTEN_URL;
|
||||
@@ -60,12 +55,10 @@ const MAX_RECENT_REQUEST_IDS: usize = 4096;
|
||||
const MAX_RECENT_SESSION_IDS: usize = 4096;
|
||||
const SHUTDOWN_TIMEOUT: Duration = Duration::from_secs(5);
|
||||
const OUTGOING_CHANNEL_CAPACITY: usize = 128;
|
||||
const BULK_PAIRING_TIMEOUT: Duration = Duration::from_secs(10);
|
||||
|
||||
enum NegotiatedConnection {
|
||||
Rejected,
|
||||
Single,
|
||||
Dual(BulkConnectionRegistration),
|
||||
Accepted,
|
||||
}
|
||||
|
||||
struct HostLimits {
|
||||
@@ -90,7 +83,7 @@ impl HostLimits {
|
||||
}
|
||||
}
|
||||
|
||||
/// Runs the code-mode host on its configured stdio or WebSocket transport.
|
||||
/// Runs the code-mode host on its configured stdio or gRPC transport.
|
||||
pub async fn run_main(listen_url: &str) -> Result<()> {
|
||||
transport::run_transport(listen_url).await
|
||||
}
|
||||
@@ -106,55 +99,17 @@ where
|
||||
R: AsyncRead + Send + Unpin + 'static,
|
||||
W: AsyncWrite + Send + Unpin + 'static,
|
||||
{
|
||||
run_connection(
|
||||
ConnectionReader::from_reader(reader),
|
||||
ConnectionWriter::from_writer(writer),
|
||||
Arc::new(HostLimits::new()),
|
||||
/*bulk_connections*/ None,
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
async fn run_connection(
|
||||
mut reader: ConnectionReader,
|
||||
mut writer: ConnectionWriter,
|
||||
limits: Arc<HostLimits>,
|
||||
bulk_connections: Option<BulkConnectionRegistry>,
|
||||
) -> Result<()> {
|
||||
let negotiated = negotiate(&mut reader, &mut writer, bulk_connections.as_ref()).await?;
|
||||
let bulk_connection = match negotiated {
|
||||
let mut reader = FramedReader::new(reader);
|
||||
let mut writer = FramedWriter::new(writer);
|
||||
match negotiate(&mut reader, &mut writer).await? {
|
||||
NegotiatedConnection::Rejected => return Ok(()),
|
||||
NegotiatedConnection::Single => None,
|
||||
NegotiatedConnection::Dual(mut registration) => {
|
||||
match tokio::time::timeout(BULK_PAIRING_TIMEOUT, registration.receive()).await {
|
||||
Ok(Ok(connection)) => Some(connection),
|
||||
Ok(Err(_)) => {
|
||||
anyhow::bail!("code-mode host bulk websocket pairing was abandoned");
|
||||
}
|
||||
Err(_) => {
|
||||
anyhow::bail!("timed out pairing code-mode host bulk websocket");
|
||||
}
|
||||
}
|
||||
}
|
||||
};
|
||||
let (mut bulk_reader, bulk_writer) = match bulk_connection {
|
||||
Some(connection) => (Some(connection.reader), Some(connection.writer)),
|
||||
None => (None, None),
|
||||
};
|
||||
NegotiatedConnection::Accepted => {}
|
||||
}
|
||||
let (outgoing_tx, outgoing_rx) = mpsc::channel::<EncodedFrame>(OUTGOING_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 peer = match bulk_tx {
|
||||
Some(sender) => Arc::new(HostPeer::new(outgoing_tx).with_bulk_sender(sender)),
|
||||
None => Arc::new(HostPeer::new(outgoing_tx)),
|
||||
};
|
||||
let peer = Arc::new(HostPeer::new(outgoing_tx));
|
||||
let state = Arc::new(HostState {
|
||||
sessions: Mutex::new(HashMap::new()),
|
||||
limits,
|
||||
limits: Arc::new(HostLimits::new()),
|
||||
seen_session_ids: Mutex::new(SeenSessionIds::default()),
|
||||
requests: Mutex::new(RequestRegistry::default()),
|
||||
request_tasks: TaskTracker::new(),
|
||||
@@ -162,17 +117,8 @@ async fn run_connection(
|
||||
peer: Arc::clone(&peer),
|
||||
});
|
||||
let writer_disconnected = peer.disconnection_token();
|
||||
let writer_task = tokio::spawn(async move {
|
||||
if let (Some(bulk_writer), Some(bulk_rx)) = (bulk_writer, bulk_rx) {
|
||||
tokio::try_join!(
|
||||
drive_writer(writer, outgoing_rx, writer_disconnected.clone()),
|
||||
drive_writer(bulk_writer, bulk_rx, writer_disconnected)
|
||||
)?;
|
||||
Ok(())
|
||||
} else {
|
||||
drive_writer(writer, outgoing_rx, writer_disconnected).await
|
||||
}
|
||||
});
|
||||
let writer_task =
|
||||
tokio::spawn(async move { drive_writer(writer, outgoing_rx, writer_disconnected).await });
|
||||
let writer_peer = Arc::clone(&peer);
|
||||
let writer_supervisor = tokio::spawn(async move {
|
||||
match writer_task.await {
|
||||
@@ -191,30 +137,14 @@ async fn run_connection(
|
||||
|
||||
let input_result = async {
|
||||
loop {
|
||||
let (message, lane) = tokio::select! {
|
||||
// Session operations and shutdown must make progress even when bulk callbacks are ready.
|
||||
let message = tokio::select! {
|
||||
biased;
|
||||
_ = peer.disconnected() => break,
|
||||
message = reader.read() => (
|
||||
message.context("failed to read code-mode client control message")?,
|
||||
TransportLane::Control,
|
||||
),
|
||||
message = async {
|
||||
match &mut bulk_reader {
|
||||
Some(reader) => reader.read().await,
|
||||
None => Ok(None),
|
||||
}
|
||||
}, if bulk_reader.is_some() => (
|
||||
message.context("failed to read code-mode client bulk message")?,
|
||||
TransportLane::Bulk,
|
||||
),
|
||||
message = reader.read() => message.context("failed to read code-mode client message")?,
|
||||
};
|
||||
let Some(message) = message else {
|
||||
break;
|
||||
};
|
||||
if bulk_reader.is_some() && !message.allows_transport_lane(lane) {
|
||||
anyhow::bail!("code-mode client sent a message on the wrong websocket lane");
|
||||
}
|
||||
match message {
|
||||
ClientToHost::ClientHello(_) => {
|
||||
anyhow::bail!("received a second code-mode client hello");
|
||||
@@ -255,8 +185,8 @@ async fn run_connection(
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn drive_writer(
|
||||
mut writer: ConnectionWriter,
|
||||
async fn drive_writer<W: AsyncWrite + Unpin>(
|
||||
mut writer: FramedWriter<W>,
|
||||
mut outgoing: mpsc::Receiver<EncodedFrame>,
|
||||
disconnected: CancellationToken,
|
||||
) -> Result<()> {
|
||||
@@ -269,7 +199,7 @@ async fn drive_writer(
|
||||
};
|
||||
tokio::select! {
|
||||
_ = disconnected.cancelled() => return Ok(()),
|
||||
result = writer.write_frame(frame) => {
|
||||
result = writer.write_frame(&frame) => {
|
||||
result.context("failed to write code-mode host message")?;
|
||||
}
|
||||
}
|
||||
@@ -278,10 +208,9 @@ async fn drive_writer(
|
||||
}
|
||||
}
|
||||
|
||||
async fn negotiate(
|
||||
reader: &mut ConnectionReader,
|
||||
writer: &mut ConnectionWriter,
|
||||
bulk_connections: Option<&BulkConnectionRegistry>,
|
||||
async fn negotiate<R: AsyncRead + Unpin, W: AsyncWrite + Unpin>(
|
||||
reader: &mut FramedReader<R>,
|
||||
writer: &mut FramedWriter<W>,
|
||||
) -> Result<NegotiatedConnection> {
|
||||
let Some(first_message) = reader
|
||||
.read()
|
||||
@@ -316,18 +245,6 @@ async fn negotiate(
|
||||
return Ok(NegotiatedConnection::Rejected);
|
||||
}
|
||||
|
||||
let dual_capability = Capability::new(DUAL_WEBSOCKET_CAPABILITY)?;
|
||||
let dual_requested = client_hello
|
||||
.required_capabilities()
|
||||
.contains(&dual_capability)
|
||||
|| client_hello
|
||||
.optional_capabilities()
|
||||
.contains(&dual_capability);
|
||||
let registration = if dual_requested {
|
||||
bulk_connections.and_then(BulkConnectionRegistry::reserve)
|
||||
} else {
|
||||
None
|
||||
};
|
||||
let resource_limits_capability = Capability::new(SESSION_RESOURCE_LIMITS_CAPABILITY)?;
|
||||
let resource_limits_requested = client_hello
|
||||
.required_capabilities()
|
||||
@@ -335,14 +252,8 @@ async fn negotiate(
|
||||
|| client_hello
|
||||
.optional_capabilities()
|
||||
.contains(&resource_limits_capability);
|
||||
let host_capabilities = CapabilitySet::try_new(
|
||||
[
|
||||
registration.is_some().then_some(dual_capability),
|
||||
resource_limits_requested.then_some(resource_limits_capability),
|
||||
]
|
||||
.into_iter()
|
||||
.flatten(),
|
||||
)?;
|
||||
let host_capabilities =
|
||||
CapabilitySet::try_new(resource_limits_requested.then_some(resource_limits_capability))?;
|
||||
if let Some(capability) = client_hello
|
||||
.required_capabilities()
|
||||
.iter()
|
||||
@@ -359,23 +270,12 @@ async fn negotiate(
|
||||
return Ok(NegotiatedConnection::Rejected);
|
||||
}
|
||||
|
||||
let (hello, negotiated) = if let Some(registration) = registration {
|
||||
(
|
||||
HostHello::new(ProtocolVersion::V1, host_capabilities)
|
||||
.with_bulk_connection_token(registration.token().to_string()),
|
||||
NegotiatedConnection::Dual(registration),
|
||||
)
|
||||
} else {
|
||||
(
|
||||
HostHello::new(ProtocolVersion::V1, host_capabilities),
|
||||
NegotiatedConnection::Single,
|
||||
)
|
||||
};
|
||||
let hello = HostHello::new(ProtocolVersion::V1, host_capabilities);
|
||||
writer
|
||||
.write(&HostToClient::HostHello(hello))
|
||||
.await
|
||||
.context("failed to write code-mode host hello")?;
|
||||
Ok(negotiated)
|
||||
Ok(NegotiatedConnection::Accepted)
|
||||
}
|
||||
|
||||
struct HostState {
|
||||
|
||||
@@ -2,7 +2,7 @@ use clap::Parser;
|
||||
|
||||
#[derive(Debug, Parser)]
|
||||
struct Cli {
|
||||
/// Transport endpoint: `stdio`, `stdio://`, `ws://IP:PORT`, or `grpc://IP:PORT`.
|
||||
/// Transport endpoint: `stdio`, `stdio://`, or `grpc://IP:PORT`.
|
||||
#[arg(
|
||||
long,
|
||||
value_name = "URL",
|
||||
|
||||
@@ -16,7 +16,6 @@ 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;
|
||||
@@ -30,7 +29,6 @@ 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>>,
|
||||
@@ -64,7 +62,6 @@ 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()),
|
||||
@@ -75,16 +72,10 @@ 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()))?;
|
||||
let lane = message.transport_lane();
|
||||
self.send_frame(frame, lane)
|
||||
self.send_frame(frame)
|
||||
}
|
||||
|
||||
pub(super) fn respond(
|
||||
@@ -398,12 +389,8 @@ impl HostPeer {
|
||||
});
|
||||
}
|
||||
|
||||
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) {
|
||||
fn send_frame(&self, frame: EncodedFrame) -> Result<(), PeerSendError> {
|
||||
match self.outgoing_tx.try_send(frame) {
|
||||
Ok(()) => Ok(()),
|
||||
Err(mpsc::error::TrySendError::Full(_)) => {
|
||||
self.disconnect();
|
||||
|
||||
@@ -1,211 +1,22 @@
|
||||
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;
|
||||
use axum::extract::ws::WebSocketUpgrade;
|
||||
use axum::http::Request;
|
||||
use axum::http::StatusCode;
|
||||
use axum::http::header::ORIGIN;
|
||||
use axum::middleware;
|
||||
use axum::middleware::Next;
|
||||
use axum::response::IntoResponse;
|
||||
use axum::response::Response;
|
||||
use axum::routing::any;
|
||||
use axum::routing::get;
|
||||
use axum::serve::Listener;
|
||||
use axum::serve::ListenerExt;
|
||||
use axum::serve::TapIo;
|
||||
use codex_code_mode_protocol::host::ClientToHost;
|
||||
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 futures::SinkExt;
|
||||
use futures::StreamExt;
|
||||
use futures::stream::SplitSink;
|
||||
use futures::stream::SplitStream;
|
||||
use tokio::io::AsyncRead;
|
||||
use tokio::io::AsyncWrite;
|
||||
use tokio::net::TcpListener;
|
||||
use tokio::net::TcpStream;
|
||||
use tokio::sync::oneshot;
|
||||
use tracing::info;
|
||||
use tracing::warn;
|
||||
use uuid::Uuid;
|
||||
|
||||
use crate::HostLimits;
|
||||
use crate::MAX_IN_FLIGHT_REQUESTS;
|
||||
use crate::grpc_transport;
|
||||
|
||||
/// 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>;
|
||||
|
||||
#[derive(Debug, Clone, Eq, PartialEq)]
|
||||
enum ListenTransport {
|
||||
Stdio,
|
||||
WebSocket(SocketAddr),
|
||||
Grpc(SocketAddr),
|
||||
}
|
||||
|
||||
pub(crate) enum ConnectionReader {
|
||||
Framed(FramedReader<BoxedReader>),
|
||||
WebSocket(SplitStream<WebSocket>),
|
||||
}
|
||||
|
||||
pub(crate) enum ConnectionWriter {
|
||||
Framed(FramedWriter<BoxedWriter>),
|
||||
WebSocket(SplitSink<WebSocket, Message>),
|
||||
}
|
||||
|
||||
#[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 {
|
||||
pub(crate) fn from_reader<R>(reader: R) -> Self
|
||||
where
|
||||
R: AsyncRead + Send + Unpin + 'static,
|
||||
{
|
||||
Self::Framed(FramedReader::new(Box::new(reader)))
|
||||
}
|
||||
|
||||
pub(crate) async fn read(&mut self) -> io::Result<Option<ClientToHost>> {
|
||||
match self {
|
||||
Self::Framed(reader) => reader.read().await,
|
||||
Self::WebSocket(reader) => loop {
|
||||
match reader.next().await {
|
||||
Some(Ok(Message::Binary(bytes))) => {
|
||||
return EncodedFrame::decode_framed(&bytes).map(Some);
|
||||
}
|
||||
Some(Ok(Message::Text(_))) => {
|
||||
return Err(io::Error::new(
|
||||
io::ErrorKind::InvalidData,
|
||||
"code-mode websocket messages must be binary framed messages",
|
||||
));
|
||||
}
|
||||
Some(Ok(Message::Ping(_) | Message::Pong(_))) => {}
|
||||
Some(Ok(Message::Close(_))) | None => return Ok(None),
|
||||
Some(Err(err)) => {
|
||||
return Err(io::Error::other(format!(
|
||||
"failed to read code-mode websocket message: {err}"
|
||||
)));
|
||||
}
|
||||
}
|
||||
},
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl ConnectionWriter {
|
||||
pub(crate) fn from_writer<W>(writer: W) -> Self
|
||||
where
|
||||
W: AsyncWrite + Send + Unpin + 'static,
|
||||
{
|
||||
Self::Framed(FramedWriter::new(Box::new(writer)))
|
||||
}
|
||||
|
||||
pub(crate) async fn write(&mut self, message: &HostToClient) -> io::Result<()> {
|
||||
self.write_frame(EncodedFrame::encode(message)?).await
|
||||
}
|
||||
|
||||
pub(crate) async fn write_frame(&mut self, frame: EncodedFrame) -> io::Result<()> {
|
||||
match self {
|
||||
Self::Framed(writer) => writer.write_frame(&frame).await,
|
||||
Self::WebSocket(writer) => writer
|
||||
.send(Message::Binary(frame.into_framed_bytes().into()))
|
||||
.await
|
||||
.map_err(|err| {
|
||||
io::Error::other(format!(
|
||||
"failed to write code-mode websocket message: {err}"
|
||||
))
|
||||
}),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn run_transport(listen_url: &str) -> Result<()> {
|
||||
match parse_listen_url(listen_url)? {
|
||||
ListenTransport::Stdio => crate::run_stdio().await,
|
||||
ListenTransport::WebSocket(bind_address) => run_websocket_listener(bind_address).await,
|
||||
ListenTransport::Grpc(bind_address) => grpc_transport::run_tcp_listener(bind_address).await,
|
||||
}
|
||||
}
|
||||
@@ -215,15 +26,6 @@ fn parse_listen_url(listen_url: &str) -> Result<ListenTransport> {
|
||||
return Ok(ListenTransport::Stdio);
|
||||
}
|
||||
|
||||
if let Some(socket_addr) = listen_url.strip_prefix("ws://") {
|
||||
return socket_addr
|
||||
.parse::<SocketAddr>()
|
||||
.map(ListenTransport::WebSocket)
|
||||
.with_context(|| {
|
||||
format!("invalid websocket --listen URL `{listen_url}`; expected `ws://IP:PORT`")
|
||||
});
|
||||
}
|
||||
|
||||
if let Some(socket_addr) = listen_url.strip_prefix("grpc://") {
|
||||
return socket_addr
|
||||
.parse::<SocketAddr>()
|
||||
@@ -234,129 +36,10 @@ fn parse_listen_url(listen_url: &str) -> Result<ListenTransport> {
|
||||
}
|
||||
|
||||
anyhow::bail!(
|
||||
"unsupported --listen URL `{listen_url}`; expected `ws://IP:PORT`, `grpc://IP:PORT`, `stdio`, or `stdio://`"
|
||||
"unsupported --listen URL `{listen_url}`; expected `grpc://IP:PORT`, `stdio`, or `stdio://`"
|
||||
);
|
||||
}
|
||||
|
||||
async fn bind_websocket_listener(
|
||||
bind_address: SocketAddr,
|
||||
) -> Result<TapIo<TcpListener, impl FnMut(&mut TcpStream) + Send + 'static>> {
|
||||
let listener = TcpListener::bind(bind_address)
|
||||
.await
|
||||
.with_context(|| format!("failed to bind code-mode host websocket to {bind_address}"))?;
|
||||
|
||||
Ok(listener.tap_io(|stream| {
|
||||
if let Err(error) = stream.set_nodelay(/*nodelay*/ true) {
|
||||
warn!(%error, "failed to enable TCP_NODELAY for code-mode host connection");
|
||||
}
|
||||
}))
|
||||
}
|
||||
|
||||
async fn run_websocket_listener(bind_address: SocketAddr) -> Result<()> {
|
||||
let listener = bind_websocket_listener(bind_address).await?;
|
||||
let local_addr = listener
|
||||
.local_addr()
|
||||
.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}");
|
||||
io::stdout()
|
||||
.flush()
|
||||
.context("failed to publish code-mode host websocket listen address")?;
|
||||
|
||||
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);
|
||||
axum::serve(
|
||||
listener,
|
||||
router.into_make_service_with_connect_info::<SocketAddr>(),
|
||||
)
|
||||
.await
|
||||
.context("code-mode host websocket listener failed")
|
||||
}
|
||||
|
||||
async fn readiness_handler() -> StatusCode {
|
||||
StatusCode::OK
|
||||
}
|
||||
|
||||
async fn reject_requests_with_origin_header(
|
||||
request: Request<Body>,
|
||||
next: Next,
|
||||
) -> Result<Response, StatusCode> {
|
||||
if request.headers().contains_key(ORIGIN) {
|
||||
warn!(
|
||||
method = %request.method(),
|
||||
uri = %request.uri(),
|
||||
"rejecting code-mode host websocket request with Origin header"
|
||||
);
|
||||
Err(StatusCode::FORBIDDEN)
|
||||
} else {
|
||||
Ok(next.run(request).await)
|
||||
}
|
||||
}
|
||||
|
||||
async fn websocket_upgrade_handler(
|
||||
websocket: WebSocketUpgrade,
|
||||
ConnectInfo(peer_addr): ConnectInfo<SocketAddr>,
|
||||
State(state): State<WebSocketListenerState>,
|
||||
) -> impl IntoResponse {
|
||||
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 websocket client connected");
|
||||
let (writer, reader) = stream.split();
|
||||
if let Err(err) = crate::run_connection(
|
||||
ConnectionReader::WebSocket(reader),
|
||||
ConnectionWriter::WebSocket(writer),
|
||||
state.limits,
|
||||
Some(state.bulk_connections),
|
||||
)
|
||||
.await
|
||||
{
|
||||
warn!(%peer_addr, "code-mode host session failed: {err:#}");
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
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;
|
||||
|
||||
@@ -1,40 +1,11 @@
|
||||
use std::net::SocketAddr;
|
||||
|
||||
use axum::serve::Listener;
|
||||
use futures::StreamExt;
|
||||
use pretty_assertions::assert_eq;
|
||||
use tokio::net::TcpStream;
|
||||
|
||||
use super::BulkConnectionRegistry;
|
||||
use super::ListenTransport;
|
||||
use super::MAX_PENDING_BULK_CONNECTIONS;
|
||||
use super::bind_websocket_listener;
|
||||
use super::parse_listen_url;
|
||||
use crate::grpc_transport::bind_tcp_listener;
|
||||
|
||||
#[tokio::test]
|
||||
async fn websocket_listener_disables_nagle() {
|
||||
let bind_address = "127.0.0.1:0"
|
||||
.parse()
|
||||
.expect("websocket test listener should have a valid bind address");
|
||||
let mut listener = bind_websocket_listener(bind_address)
|
||||
.await
|
||||
.expect("websocket test listener should bind");
|
||||
let local_addr = listener
|
||||
.local_addr()
|
||||
.expect("websocket test listener should have a local address");
|
||||
let _client = TcpStream::connect(local_addr)
|
||||
.await
|
||||
.expect("websocket test client should connect");
|
||||
let (stream, _) = listener.accept().await;
|
||||
|
||||
let stream = stream
|
||||
.nodelay()
|
||||
.expect("accepted websocket socket should expose TCP_NODELAY");
|
||||
|
||||
assert!(stream);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn grpc_listener_disables_nagle() {
|
||||
let bind_address = "127.0.0.1:0"
|
||||
@@ -60,29 +31,6 @@ async fn grpc_listener_disables_nagle() {
|
||||
assert!(nodelay);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn bulk_connection_registration_cleans_up_when_dropped() {
|
||||
let registry = BulkConnectionRegistry::default();
|
||||
let registration = registry.reserve().expect("bulk connection registration");
|
||||
let token = registration.token();
|
||||
|
||||
drop(registration);
|
||||
|
||||
assert!(registry.remove(token).is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn bulk_connection_registrations_are_bounded_and_released() {
|
||||
let registry = BulkConnectionRegistry::default();
|
||||
let mut registrations = (0..MAX_PENDING_BULK_CONNECTIONS)
|
||||
.map(|_| registry.reserve().expect("bulk connection registration"))
|
||||
.collect::<Vec<_>>();
|
||||
|
||||
assert!(registry.reserve().is_none());
|
||||
drop(registrations.pop());
|
||||
assert!(registry.reserve().is_some());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parse_listen_url_accepts_stdio_transports() {
|
||||
assert_eq!(
|
||||
@@ -95,34 +43,14 @@ fn parse_listen_url_accepts_stdio_transports() {
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parse_listen_url_accepts_websocket_addresses() {
|
||||
assert_eq!(
|
||||
parse_listen_url("ws://127.0.0.1:0").expect("websocket listen URL should parse"),
|
||||
ListenTransport::WebSocket(
|
||||
"127.0.0.1:0"
|
||||
.parse::<SocketAddr>()
|
||||
.expect("valid socket address")
|
||||
)
|
||||
);
|
||||
assert_eq!(
|
||||
parse_listen_url("ws://[::1]:9000").expect("IPv6 websocket listen URL should parse"),
|
||||
ListenTransport::WebSocket(
|
||||
"[::1]:9000"
|
||||
.parse::<SocketAddr>()
|
||||
.expect("valid IPv6 socket address")
|
||||
)
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parse_listen_url_rejects_invalid_transports() {
|
||||
let invalid_address = parse_listen_url("ws://localhost:9000")
|
||||
.expect_err("websocket listener requires an IP address");
|
||||
let invalid_address = parse_listen_url("grpc://localhost:9000")
|
||||
.expect_err("gRPC listener requires an IP address");
|
||||
assert!(
|
||||
invalid_address
|
||||
.to_string()
|
||||
.contains("expected `ws://IP:PORT`")
|
||||
.contains("expected `grpc://IP:PORT`")
|
||||
);
|
||||
|
||||
let unsupported =
|
||||
|
||||
@@ -1,842 +0,0 @@
|
||||
use std::sync::Arc;
|
||||
use std::time::Duration;
|
||||
|
||||
use anyhow::Context;
|
||||
use anyhow::Result;
|
||||
use codex_code_mode::CellId;
|
||||
use codex_code_mode::CodeModeSessionCellExecutionLimits;
|
||||
use codex_code_mode::CodeModeSessionProvider;
|
||||
use codex_code_mode::CodeModeToolKind;
|
||||
use codex_code_mode::ExecuteRequest;
|
||||
use codex_code_mode::FunctionCallOutputContentItem;
|
||||
use codex_code_mode::NoopCodeModeSessionDelegate;
|
||||
use codex_code_mode::RuntimeResponse;
|
||||
use codex_code_mode::ToolDefinition;
|
||||
use codex_code_mode::WebSocketCodeModeSessionProvider;
|
||||
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::DelegateRequest;
|
||||
use codex_code_mode_protocol::host::DelegateResponse;
|
||||
use codex_code_mode_protocol::host::EncodedFrame;
|
||||
use codex_code_mode_protocol::host::HostHello;
|
||||
use codex_code_mode_protocol::host::HostRequest;
|
||||
use codex_code_mode_protocol::host::HostResponse;
|
||||
use codex_code_mode_protocol::host::HostToClient;
|
||||
use codex_code_mode_protocol::host::MAX_FRAME_BYTES;
|
||||
use codex_code_mode_protocol::host::ProtocolVersion;
|
||||
use codex_code_mode_protocol::host::RequestId;
|
||||
use codex_code_mode_protocol::host::SESSION_RESOURCE_LIMITS_CAPABILITY;
|
||||
use codex_code_mode_protocol::host::SessionId;
|
||||
use codex_code_mode_protocol::host::SupportedProtocolVersions;
|
||||
use codex_code_mode_protocol::host::WireContentItem;
|
||||
use codex_code_mode_protocol::host::WireExecuteRequest;
|
||||
use codex_code_mode_protocol::host::WireResult;
|
||||
use codex_code_mode_protocol::host::WireRuntimeResponse;
|
||||
use codex_code_mode_protocol::host::WireToolDefinition;
|
||||
use codex_code_mode_protocol::host::WireToolKind;
|
||||
use codex_code_mode_protocol::host::WireToolName;
|
||||
use codex_code_mode_protocol::host::WireWaitRequest;
|
||||
use codex_protocol::ToolName;
|
||||
use futures::SinkExt;
|
||||
use futures::StreamExt;
|
||||
use pretty_assertions::assert_eq;
|
||||
use serde_json::json;
|
||||
use tokio::io::AsyncBufReadExt;
|
||||
use tokio::io::AsyncReadExt;
|
||||
use tokio::io::AsyncWriteExt;
|
||||
use tokio::io::BufReader;
|
||||
use tokio::net::TcpStream;
|
||||
use tokio::sync::Semaphore;
|
||||
use tokio::time::timeout;
|
||||
use tokio_tungstenite::MaybeTlsStream;
|
||||
use tokio_tungstenite::WebSocketStream;
|
||||
use tokio_tungstenite::connect_async;
|
||||
use tokio_tungstenite::connect_async_with_config;
|
||||
use tokio_tungstenite::tungstenite::Error as WebSocketError;
|
||||
use tokio_tungstenite::tungstenite::Message;
|
||||
use tokio_tungstenite::tungstenite::client::IntoClientRequest;
|
||||
use tokio_tungstenite::tungstenite::http::HeaderValue;
|
||||
use tokio_tungstenite::tungstenite::http::StatusCode;
|
||||
use tokio_tungstenite::tungstenite::http::header::ORIGIN;
|
||||
use tokio_tungstenite::tungstenite::protocol::WebSocketConfig;
|
||||
use uuid::Uuid;
|
||||
|
||||
#[path = "support/host.rs"]
|
||||
mod host;
|
||||
#[path = "support/large_tool_delegate.rs"]
|
||||
mod large_tool_delegate;
|
||||
|
||||
use host::HostHarness;
|
||||
use large_tool_delegate::LargeToolResultDelegate;
|
||||
|
||||
const TEST_TIMEOUT: Duration = Duration::from_secs(10);
|
||||
const MAX_WEBSOCKET_FRAME_BYTES: usize = MAX_FRAME_BYTES + std::mem::size_of::<u32>();
|
||||
|
||||
struct HostClient {
|
||||
websocket: WebSocketStream<MaybeTlsStream<TcpStream>>,
|
||||
}
|
||||
|
||||
impl HostHarness {
|
||||
async fn connect(&self) -> Result<HostClient> {
|
||||
let config = WebSocketConfig::default()
|
||||
.max_frame_size(Some(MAX_WEBSOCKET_FRAME_BYTES))
|
||||
.max_message_size(Some(MAX_WEBSOCKET_FRAME_BYTES));
|
||||
let (websocket, _) = timeout(
|
||||
TEST_TIMEOUT,
|
||||
connect_async_with_config(
|
||||
self.endpoint.as_str(),
|
||||
Some(config),
|
||||
/*disable_nagle*/ false,
|
||||
),
|
||||
)
|
||||
.await
|
||||
.context("timed out connecting to code-mode host websocket")??;
|
||||
Ok(HostClient { websocket })
|
||||
}
|
||||
}
|
||||
|
||||
impl HostClient {
|
||||
async fn send(&mut self, message: &ClientToHost) -> Result<()> {
|
||||
let frame = EncodedFrame::encode(message)?;
|
||||
self.send_binary(frame.into_framed_bytes()).await
|
||||
}
|
||||
|
||||
async fn send_binary(&mut self, bytes: Vec<u8>) -> Result<()> {
|
||||
timeout(
|
||||
TEST_TIMEOUT,
|
||||
self.websocket.send(Message::Binary(bytes.into())),
|
||||
)
|
||||
.await
|
||||
.context("timed out writing code-mode websocket message")?
|
||||
.context("failed to write code-mode websocket message")
|
||||
}
|
||||
|
||||
async fn read(&mut self) -> Result<HostToClient> {
|
||||
loop {
|
||||
let message = timeout(TEST_TIMEOUT, self.websocket.next())
|
||||
.await
|
||||
.context("timed out waiting for code-mode websocket message")?
|
||||
.context("code-mode websocket closed before returning a message")?
|
||||
.context("failed to read code-mode websocket message")?;
|
||||
match message {
|
||||
Message::Binary(bytes) => {
|
||||
return EncodedFrame::decode_framed(&bytes)
|
||||
.context("failed to decode code-mode websocket frame");
|
||||
}
|
||||
Message::Ping(_) | Message::Pong(_) => {}
|
||||
Message::Close(frame) => {
|
||||
anyhow::bail!("code-mode websocket closed unexpectedly: {frame:?}");
|
||||
}
|
||||
Message::Text(text) => {
|
||||
anyhow::bail!("code-mode host returned a text websocket message: {text}");
|
||||
}
|
||||
Message::Frame(_) => {
|
||||
anyhow::bail!("code-mode host returned an unexpected raw websocket frame");
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async fn negotiate(&mut self, optional_capabilities: CapabilitySet) -> Result<()> {
|
||||
let hello = ClientHello::new(
|
||||
SupportedProtocolVersions::try_new([ProtocolVersion::V1])?,
|
||||
CapabilitySet::empty(),
|
||||
optional_capabilities,
|
||||
)?;
|
||||
self.send(&ClientToHost::ClientHello(hello)).await?;
|
||||
assert_eq!(
|
||||
self.read().await?,
|
||||
HostToClient::HostHello(HostHello::new(ProtocolVersion::V1, CapabilitySet::empty()))
|
||||
);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn negotiate_dual(&mut self, websocket_url: &str) -> Result<HostClient> {
|
||||
let capability = Capability::new(DUAL_WEBSOCKET_CAPABILITY)?;
|
||||
let resource_limits = Capability::new(SESSION_RESOURCE_LIMITS_CAPABILITY)?;
|
||||
let hello = ClientHello::new(
|
||||
SupportedProtocolVersions::try_new([ProtocolVersion::V1])?,
|
||||
CapabilitySet::empty(),
|
||||
CapabilitySet::try_new([capability.clone(), resource_limits.clone()])?,
|
||||
)?;
|
||||
self.send(&ClientToHost::ClientHello(hello)).await?;
|
||||
let HostToClient::HostHello(hello) = self.read().await? else {
|
||||
anyhow::bail!("expected code-mode host hello");
|
||||
};
|
||||
assert!(hello.capabilities().contains(&capability));
|
||||
assert!(hello.capabilities().contains(&resource_limits));
|
||||
let token = hello
|
||||
.bulk_connection_token()
|
||||
.context("dual websocket handshake omitted its pairing token")?;
|
||||
let bulk_url = format!("{}/bulk/{token}", websocket_url.trim_end_matches('/'));
|
||||
let config = WebSocketConfig::default()
|
||||
.max_frame_size(Some(MAX_WEBSOCKET_FRAME_BYTES))
|
||||
.max_message_size(Some(MAX_WEBSOCKET_FRAME_BYTES));
|
||||
let (websocket, _) = timeout(
|
||||
TEST_TIMEOUT,
|
||||
connect_async_with_config(bulk_url, Some(config), /*disable_nagle*/ false),
|
||||
)
|
||||
.await
|
||||
.context("timed out connecting to code-mode host bulk websocket")??;
|
||||
Ok(HostClient { websocket })
|
||||
}
|
||||
|
||||
async fn open_session(&mut self, session_id: SessionId) -> Result<()> {
|
||||
let id = RequestId::new(/*value*/ 1);
|
||||
self.send(&ClientToHost::Request {
|
||||
id,
|
||||
request: HostRequest::OpenSession {
|
||||
session_id: session_id.clone(),
|
||||
cell_execution_limits: None,
|
||||
},
|
||||
})
|
||||
.await?;
|
||||
assert_eq!(
|
||||
self.read().await?,
|
||||
HostToClient::Response {
|
||||
id,
|
||||
result: WireResult::Ok {
|
||||
value: HostResponse::SessionReady { session_id },
|
||||
},
|
||||
}
|
||||
);
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn websocket_listener_serves_readiness_endpoint() -> Result<()> {
|
||||
let host = HostHarness::start("ws://127.0.0.1:0").await?;
|
||||
let address = host
|
||||
.endpoint
|
||||
.strip_prefix("ws://")
|
||||
.context("code-mode host websocket URL should use ws://")?;
|
||||
|
||||
let response = timeout(TEST_TIMEOUT, async {
|
||||
let mut stream = TcpStream::connect(address)
|
||||
.await
|
||||
.context("failed to connect to code-mode host readiness endpoint")?;
|
||||
let request =
|
||||
format!("GET /readyz HTTP/1.1\r\nHost: {address}\r\nConnection: close\r\n\r\n");
|
||||
stream
|
||||
.write_all(request.as_bytes())
|
||||
.await
|
||||
.context("failed to request code-mode host readiness")?;
|
||||
|
||||
let mut response = String::new();
|
||||
stream
|
||||
.read_to_string(&mut response)
|
||||
.await
|
||||
.context("failed to read code-mode host readiness response")?;
|
||||
Ok::<_, anyhow::Error>(response)
|
||||
})
|
||||
.await
|
||||
.context("timed out requesting code-mode host readiness")??;
|
||||
|
||||
let status_line = response
|
||||
.lines()
|
||||
.next()
|
||||
.context("code-mode host readiness response is missing a status line")?;
|
||||
assert_eq!(status_line, "HTTP/1.1 200 OK");
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn websocket_listener_executes_cells_and_forwards_tool_callbacks() -> Result<()> {
|
||||
let host = HostHarness::start("ws://127.0.0.1:0").await?;
|
||||
let mut client = host.connect().await?;
|
||||
client.negotiate(CapabilitySet::empty()).await?;
|
||||
|
||||
let session_id = SessionId::new("websocket-session")?;
|
||||
client.open_session(session_id.clone()).await?;
|
||||
|
||||
let execute_id = RequestId::new(/*value*/ 2);
|
||||
client
|
||||
.send(&ClientToHost::Request {
|
||||
id: execute_id,
|
||||
request: HostRequest::Execute {
|
||||
session_id: session_id.clone(),
|
||||
request: WireExecuteRequest {
|
||||
tool_call_id: "websocket-call".to_string(),
|
||||
enabled_tools: vec![WireToolDefinition {
|
||||
name: "echo".to_string(),
|
||||
tool_name: WireToolName {
|
||||
name: "echo".to_string(),
|
||||
namespace: None,
|
||||
},
|
||||
description: String::new(),
|
||||
kind: WireToolKind::Function,
|
||||
input_schema: None,
|
||||
output_schema: None,
|
||||
}],
|
||||
source:
|
||||
r#"const result = await tools.echo({ value: "ping" }); text(result.value);"#
|
||||
.to_string(),
|
||||
yield_time_ms: Some(5_000),
|
||||
max_output_tokens: Some(1_000),
|
||||
},
|
||||
},
|
||||
})
|
||||
.await?;
|
||||
|
||||
let started = client.read().await?;
|
||||
let HostToClient::Response {
|
||||
id,
|
||||
result:
|
||||
WireResult::Ok {
|
||||
value: HostResponse::ExecutionStarted { cell_id },
|
||||
},
|
||||
} = started
|
||||
else {
|
||||
anyhow::bail!("expected execution-started response, got {started:?}");
|
||||
};
|
||||
assert_eq!(id, execute_id);
|
||||
|
||||
let callback = client.read().await?;
|
||||
let HostToClient::DelegateRequest {
|
||||
id: delegate_id,
|
||||
session_id: callback_session_id,
|
||||
request: DelegateRequest::InvokeTool { invocation },
|
||||
} = callback
|
||||
else {
|
||||
anyhow::bail!("expected tool callback, got {callback:?}");
|
||||
};
|
||||
assert_eq!(callback_session_id, session_id);
|
||||
assert_eq!(invocation.input, Some(json!({ "value": "ping" })));
|
||||
|
||||
client
|
||||
.send(&ClientToHost::DelegateResponse {
|
||||
id: delegate_id,
|
||||
result: WireResult::Ok {
|
||||
value: DelegateResponse::ToolResult {
|
||||
result: json!({ "value": "pong" }),
|
||||
},
|
||||
},
|
||||
})
|
||||
.await?;
|
||||
|
||||
assert_eq!(
|
||||
client.read().await?,
|
||||
HostToClient::InitialResponse {
|
||||
id: execute_id,
|
||||
result: WireResult::Ok {
|
||||
value: WireRuntimeResponse::Result {
|
||||
cell_id,
|
||||
content_items: vec![WireContentItem::InputText {
|
||||
text: "pong".to_string(),
|
||||
}],
|
||||
error_text: None,
|
||||
},
|
||||
},
|
||||
}
|
||||
);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn production_websocket_client_runs_nested_tools_while_other_sessions_progress() -> Result<()>
|
||||
{
|
||||
let host = HostHarness::start("ws://127.0.0.1:0").await?;
|
||||
let provider = WebSocketCodeModeSessionProvider::new(host.endpoint.clone());
|
||||
let delegate = Arc::new(LargeToolResultDelegate {
|
||||
started: Semaphore::new(/*permits*/ 0),
|
||||
release: Semaphore::new(/*permits*/ 0),
|
||||
});
|
||||
let slow_session = provider
|
||||
.create_session(delegate.clone())
|
||||
.await
|
||||
.map_err(anyhow::Error::msg)?;
|
||||
let fast_session = provider
|
||||
.create_session_with_limits(
|
||||
Arc::new(NoopCodeModeSessionDelegate),
|
||||
CodeModeSessionCellExecutionLimits {
|
||||
max_yield_time_ms: Some(5_000),
|
||||
max_heap_size_bytes: None,
|
||||
},
|
||||
)
|
||||
.await
|
||||
.map_err(anyhow::Error::msg)?;
|
||||
|
||||
let slow_cell = slow_session
|
||||
.execute(ExecuteRequest {
|
||||
tool_call_id: "large-tool".to_string(),
|
||||
enabled_tools: vec![ToolDefinition {
|
||||
name: "large".to_string(),
|
||||
tool_name: ToolName::plain("large"),
|
||||
description: String::new(),
|
||||
kind: CodeModeToolKind::Function,
|
||||
input_schema: None,
|
||||
output_schema: None,
|
||||
}],
|
||||
source: r#"const result = await tools.large({ value: "request" }); text(String(result.value.length));"#
|
||||
.to_string(),
|
||||
yield_time_ms: Some(20_000),
|
||||
max_output_tokens: Some(1_000),
|
||||
})
|
||||
.await
|
||||
.map_err(anyhow::Error::msg)?;
|
||||
let started = timeout(TEST_TIMEOUT, delegate.started.acquire())
|
||||
.await
|
||||
.context("large tool callback did not start")??;
|
||||
started.forget();
|
||||
|
||||
let fast_response = timeout(TEST_TIMEOUT, async {
|
||||
fast_session
|
||||
.execute(ExecuteRequest {
|
||||
tool_call_id: "fast-before-transfer".to_string(),
|
||||
enabled_tools: Vec::new(),
|
||||
source: r#"text("fast-before");"#.to_string(),
|
||||
yield_time_ms: Some(5_000),
|
||||
max_output_tokens: Some(1_000),
|
||||
})
|
||||
.await
|
||||
.map_err(anyhow::Error::msg)?
|
||||
.initial_response()
|
||||
.await
|
||||
.map_err(anyhow::Error::msg)
|
||||
})
|
||||
.await
|
||||
.context("unrelated execution was blocked by the pending tool callback")??;
|
||||
assert_eq!(
|
||||
fast_response,
|
||||
RuntimeResponse::Result {
|
||||
cell_id: CellId::new("1".to_string()),
|
||||
content_items: vec![FunctionCallOutputContentItem::InputText {
|
||||
text: "fast-before".to_string(),
|
||||
}],
|
||||
error_text: None,
|
||||
}
|
||||
);
|
||||
|
||||
delegate.release.add_permits(1);
|
||||
let slow_response = async {
|
||||
timeout(TEST_TIMEOUT, slow_cell.initial_response())
|
||||
.await
|
||||
.context("large tool result did not finish")?
|
||||
.map_err(anyhow::Error::msg)
|
||||
};
|
||||
let concurrent_response = async {
|
||||
timeout(TEST_TIMEOUT, async {
|
||||
fast_session
|
||||
.execute(ExecuteRequest {
|
||||
tool_call_id: "fast-during-transfer".to_string(),
|
||||
enabled_tools: Vec::new(),
|
||||
source: r#"text("fast-during");"#.to_string(),
|
||||
yield_time_ms: Some(5_000),
|
||||
max_output_tokens: Some(1_000),
|
||||
})
|
||||
.await
|
||||
.map_err(anyhow::Error::msg)?
|
||||
.initial_response()
|
||||
.await
|
||||
.map_err(anyhow::Error::msg)
|
||||
})
|
||||
.await
|
||||
.context("unrelated execution was blocked by the large tool transfer")?
|
||||
};
|
||||
tokio::pin!(slow_response);
|
||||
tokio::pin!(concurrent_response);
|
||||
let concurrent_response = tokio::select! {
|
||||
response = &mut concurrent_response => response?,
|
||||
response = &mut slow_response => {
|
||||
response?;
|
||||
anyhow::bail!("large tool result completed before the unrelated control response");
|
||||
}
|
||||
};
|
||||
let slow_response = slow_response.await?;
|
||||
assert_eq!(
|
||||
slow_response,
|
||||
RuntimeResponse::Result {
|
||||
cell_id: CellId::new("1".to_string()),
|
||||
content_items: vec![FunctionCallOutputContentItem::InputText {
|
||||
text: "8388608".to_string(),
|
||||
}],
|
||||
error_text: None,
|
||||
}
|
||||
);
|
||||
assert_eq!(
|
||||
concurrent_response,
|
||||
RuntimeResponse::Result {
|
||||
cell_id: CellId::new("2".to_string()),
|
||||
content_items: vec![FunctionCallOutputContentItem::InputText {
|
||||
text: "fast-during".to_string(),
|
||||
}],
|
||||
error_text: None,
|
||||
}
|
||||
);
|
||||
|
||||
slow_session.shutdown().await.map_err(anyhow::Error::msg)?;
|
||||
fast_session.shutdown().await.map_err(anyhow::Error::msg)?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn websocket_dual_connections_route_notifications_and_tool_callbacks_to_separate_lanes()
|
||||
-> Result<()> {
|
||||
let host = HostHarness::start("ws://127.0.0.1:0").await?;
|
||||
let mut control = host.connect().await?;
|
||||
let mut bulk = control.negotiate_dual(&host.endpoint).await?;
|
||||
let session_id = SessionId::new("dual-websocket-session")?;
|
||||
control.open_session(session_id.clone()).await?;
|
||||
|
||||
let execute_id = RequestId::new(/*value*/ 2);
|
||||
control
|
||||
.send(&ClientToHost::Request {
|
||||
id: execute_id,
|
||||
request: HostRequest::Execute {
|
||||
session_id: session_id.clone(),
|
||||
request: WireExecuteRequest {
|
||||
tool_call_id: "dual-websocket-call".to_string(),
|
||||
enabled_tools: vec![WireToolDefinition {
|
||||
name: "echo".to_string(),
|
||||
tool_name: WireToolName {
|
||||
name: "echo".to_string(),
|
||||
namespace: None,
|
||||
},
|
||||
description: String::new(),
|
||||
kind: WireToolKind::Function,
|
||||
input_schema: None,
|
||||
output_schema: None,
|
||||
}],
|
||||
source: r#"notify("important"); const result = await tools.echo({ value: "ping" }); text(result.value);"#
|
||||
.to_string(),
|
||||
yield_time_ms: Some(5_000),
|
||||
max_output_tokens: Some(1_000),
|
||||
},
|
||||
},
|
||||
})
|
||||
.await?;
|
||||
|
||||
let started = control.read().await?;
|
||||
let HostToClient::Response {
|
||||
id,
|
||||
result:
|
||||
WireResult::Ok {
|
||||
value: HostResponse::ExecutionStarted { cell_id },
|
||||
},
|
||||
} = started
|
||||
else {
|
||||
anyhow::bail!("expected execution-started response on control lane, got {started:?}");
|
||||
};
|
||||
assert_eq!(id, execute_id);
|
||||
|
||||
let notification = control.read().await?;
|
||||
let HostToClient::DelegateRequest {
|
||||
id: notification_id,
|
||||
..
|
||||
} = ¬ification
|
||||
else {
|
||||
anyhow::bail!("expected notification on control lane, got {notification:?}");
|
||||
};
|
||||
let notification_id = *notification_id;
|
||||
assert_eq!(
|
||||
notification,
|
||||
HostToClient::DelegateRequest {
|
||||
id: notification_id,
|
||||
session_id: session_id.clone(),
|
||||
request: DelegateRequest::Notify {
|
||||
call_id: "dual-websocket-call".to_string(),
|
||||
cell_id: cell_id.clone(),
|
||||
text: "important".to_string(),
|
||||
},
|
||||
}
|
||||
);
|
||||
control
|
||||
.send(&ClientToHost::DelegateResponse {
|
||||
id: notification_id,
|
||||
result: WireResult::Ok {
|
||||
value: DelegateResponse::NotificationDelivered,
|
||||
},
|
||||
})
|
||||
.await?;
|
||||
|
||||
let callback = bulk.read().await?;
|
||||
let HostToClient::DelegateRequest {
|
||||
id: delegate_id,
|
||||
session_id: callback_session_id,
|
||||
request: DelegateRequest::InvokeTool { invocation },
|
||||
} = callback
|
||||
else {
|
||||
anyhow::bail!("expected tool callback on bulk lane, got {callback:?}");
|
||||
};
|
||||
assert_eq!(callback_session_id, session_id);
|
||||
assert_eq!(invocation.input, Some(json!({ "value": "ping" })));
|
||||
|
||||
bulk.send(&ClientToHost::DelegateResponse {
|
||||
id: delegate_id,
|
||||
result: WireResult::Ok {
|
||||
value: DelegateResponse::ToolResult {
|
||||
result: json!({ "value": "pong" }),
|
||||
},
|
||||
},
|
||||
})
|
||||
.await?;
|
||||
|
||||
assert_eq!(
|
||||
control.read().await?,
|
||||
HostToClient::InitialResponse {
|
||||
id: execute_id,
|
||||
result: WireResult::Ok {
|
||||
value: WireRuntimeResponse::Result {
|
||||
cell_id: cell_id.clone(),
|
||||
content_items: vec![WireContentItem::InputText {
|
||||
text: "pong".to_string(),
|
||||
}],
|
||||
error_text: None,
|
||||
},
|
||||
},
|
||||
}
|
||||
);
|
||||
|
||||
assert_eq!(
|
||||
control.read().await?,
|
||||
HostToClient::CellClosed {
|
||||
session_id: session_id.clone(),
|
||||
cell_id,
|
||||
}
|
||||
);
|
||||
|
||||
let shutdown_id = RequestId::new(/*value*/ 3);
|
||||
control
|
||||
.send(&ClientToHost::Request {
|
||||
id: shutdown_id,
|
||||
request: HostRequest::ShutdownSession {
|
||||
session_id: session_id.clone(),
|
||||
},
|
||||
})
|
||||
.await?;
|
||||
assert_eq!(
|
||||
control.read().await?,
|
||||
HostToClient::Response {
|
||||
id: shutdown_id,
|
||||
result: WireResult::Ok {
|
||||
value: HostResponse::SessionClosed { session_id },
|
||||
},
|
||||
}
|
||||
);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn websocket_control_operations_bypass_an_incomplete_bulk_frame() -> Result<()> {
|
||||
let host = HostHarness::start("ws://127.0.0.1:0").await?;
|
||||
let mut control = host.connect().await?;
|
||||
let mut bulk = control.negotiate_dual(&host.endpoint).await?;
|
||||
let session_id = SessionId::new("bulk-priority-session")?;
|
||||
control.open_session(session_id.clone()).await?;
|
||||
|
||||
let bulk_message = ClientToHost::DelegateResponse {
|
||||
id: codex_code_mode_protocol::host::DelegateRequestId::new(/*value*/ 999),
|
||||
result: WireResult::Ok {
|
||||
value: DelegateResponse::ToolResult {
|
||||
result: json!({ "image": "x".repeat(1024 * 1024) }),
|
||||
},
|
||||
},
|
||||
};
|
||||
let payload = EncodedFrame::encode(&bulk_message)?.into_framed_bytes();
|
||||
let mask = [0x13_u8, 0x37, 0xc0, 0xde];
|
||||
let mut partial_frame = vec![0x82_u8, 0xff];
|
||||
partial_frame.extend_from_slice(&(payload.len() as u64).to_be_bytes());
|
||||
partial_frame.extend_from_slice(&mask);
|
||||
partial_frame.extend(
|
||||
payload[..64 * 1024]
|
||||
.iter()
|
||||
.enumerate()
|
||||
.map(|(index, byte)| byte ^ mask[index % mask.len()]),
|
||||
);
|
||||
bulk.websocket
|
||||
.get_mut()
|
||||
.write_all(&partial_frame)
|
||||
.await
|
||||
.context("failed to start the intentionally incomplete bulk frame")?;
|
||||
|
||||
let request_id = RequestId::new(/*value*/ 42);
|
||||
control
|
||||
.send(&ClientToHost::Request {
|
||||
id: request_id,
|
||||
request: HostRequest::Wait {
|
||||
session_id: session_id.clone(),
|
||||
request: WireWaitRequest {
|
||||
cell_id: codex_code_mode_protocol::host::WireCellId::new("missing-cell"),
|
||||
yield_time_ms: 10,
|
||||
},
|
||||
},
|
||||
})
|
||||
.await?;
|
||||
let response = timeout(TEST_TIMEOUT, control.read())
|
||||
.await
|
||||
.context("control wait was blocked behind the incomplete bulk frame")??;
|
||||
assert!(matches!(
|
||||
response,
|
||||
HostToClient::Response {
|
||||
id,
|
||||
result: WireResult::Ok {
|
||||
value: HostResponse::WaitCompleted { .. },
|
||||
},
|
||||
} if id == request_id
|
||||
));
|
||||
|
||||
let execute_id = RequestId::new(/*value*/ 43);
|
||||
control
|
||||
.send(&ClientToHost::Request {
|
||||
id: execute_id,
|
||||
request: HostRequest::Execute {
|
||||
session_id,
|
||||
request: WireExecuteRequest {
|
||||
tool_call_id: "control-execute".to_string(),
|
||||
enabled_tools: Vec::new(),
|
||||
source: r#"text("fast");"#.to_string(),
|
||||
yield_time_ms: Some(5_000),
|
||||
max_output_tokens: Some(1_000),
|
||||
},
|
||||
},
|
||||
})
|
||||
.await?;
|
||||
for _ in 0..2 {
|
||||
let response = timeout(TEST_TIMEOUT, control.read())
|
||||
.await
|
||||
.context("control execute was blocked behind the incomplete bulk frame")??;
|
||||
assert!(matches!(
|
||||
response,
|
||||
HostToClient::Response { id, .. } | HostToClient::InitialResponse { id, .. }
|
||||
if id == execute_id
|
||||
));
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn websocket_bulk_lane_rejects_control_messages() -> Result<()> {
|
||||
let host = HostHarness::start("ws://127.0.0.1:0").await?;
|
||||
let mut control = host.connect().await?;
|
||||
let mut bulk = control.negotiate_dual(&host.endpoint).await?;
|
||||
let session_id = SessionId::new("wrong-lane-session")?;
|
||||
control.open_session(session_id.clone()).await?;
|
||||
|
||||
bulk.send(&ClientToHost::Request {
|
||||
id: RequestId::new(/*value*/ 42),
|
||||
request: HostRequest::Wait {
|
||||
session_id,
|
||||
request: WireWaitRequest {
|
||||
cell_id: codex_code_mode_protocol::host::WireCellId::new("missing-cell"),
|
||||
yield_time_ms: 10,
|
||||
},
|
||||
},
|
||||
})
|
||||
.await?;
|
||||
|
||||
let result = timeout(TEST_TIMEOUT, control.websocket.next())
|
||||
.await
|
||||
.context("wrong-lane control message did not disconnect the paired sockets")?;
|
||||
assert!(
|
||||
matches!(result, None | Some(Ok(Message::Close(_))) | Some(Err(_))),
|
||||
"wrong-lane message unexpectedly returned {result:?}"
|
||||
);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn websocket_bulk_pairing_rejects_unknown_tokens() -> Result<()> {
|
||||
let host = HostHarness::start("ws://127.0.0.1:0").await?;
|
||||
let unknown_token = Uuid::new_v4();
|
||||
let url = format!("{}/bulk/{unknown_token}", host.endpoint);
|
||||
let error = match connect_async(url).await {
|
||||
Ok(_) => anyhow::bail!("unknown bulk pairing token should be rejected"),
|
||||
Err(error) => error,
|
||||
};
|
||||
let WebSocketError::Http(response) = error else {
|
||||
anyhow::bail!("bulk pairing failed unexpectedly: {error}");
|
||||
};
|
||||
assert_eq!(response.status(), StatusCode::NOT_FOUND);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn websocket_listener_accepts_frames_larger_than_default_websocket_limit() -> Result<()> {
|
||||
let host = HostHarness::start("ws://127.0.0.1:0").await?;
|
||||
let mut client = host.connect().await?;
|
||||
let capability = Capability::new("x".repeat((16 * 1024 * 1024) + 1))?;
|
||||
|
||||
client
|
||||
.negotiate(CapabilitySet::try_new([capability])?)
|
||||
.await
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn websocket_listener_keeps_connections_and_session_ids_isolated() -> Result<()> {
|
||||
let host = HostHarness::start("ws://127.0.0.1:0").await?;
|
||||
let mut first = host.connect().await?;
|
||||
let mut second = host.connect().await?;
|
||||
first.negotiate(CapabilitySet::empty()).await?;
|
||||
second.negotiate(CapabilitySet::empty()).await?;
|
||||
|
||||
let session_id = SessionId::new("shared-session-name")?;
|
||||
first.open_session(session_id.clone()).await?;
|
||||
second.open_session(session_id).await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn malformed_websocket_frame_does_not_stop_the_listener() -> Result<()> {
|
||||
let mut host = HostHarness::start("ws://127.0.0.1:0").await?;
|
||||
let stderr = host
|
||||
._child
|
||||
.stderr
|
||||
.take()
|
||||
.context("code-mode host stderr was not captured")?;
|
||||
let mut stderr_lines = BufReader::new(stderr).lines();
|
||||
let mut malformed = host.connect().await?;
|
||||
malformed.send_binary(vec![1, 0, 0, 0, b'{']).await?;
|
||||
|
||||
let close = timeout(TEST_TIMEOUT, malformed.websocket.next())
|
||||
.await
|
||||
.context("timed out waiting for malformed websocket connection to close")?;
|
||||
if let Some(Ok(message)) = close
|
||||
&& !matches!(message, Message::Close(_))
|
||||
{
|
||||
anyhow::bail!("malformed websocket returned an unexpected message: {message:?}");
|
||||
}
|
||||
|
||||
let diagnostic = timeout(TEST_TIMEOUT, async {
|
||||
loop {
|
||||
let line = stderr_lines
|
||||
.next_line()
|
||||
.await?
|
||||
.context("code-mode host exited before reporting the malformed frame")?;
|
||||
if line.contains("code-mode host session failed") {
|
||||
return Ok::<_, anyhow::Error>(line);
|
||||
}
|
||||
}
|
||||
})
|
||||
.await
|
||||
.context("timed out waiting for the malformed websocket diagnostic")??;
|
||||
assert!(
|
||||
diagnostic.contains("failed to read code-mode client hello"),
|
||||
"unexpected malformed websocket diagnostic: {diagnostic}",
|
||||
);
|
||||
|
||||
let mut recovered = host.connect().await?;
|
||||
recovered.negotiate(CapabilitySet::empty()).await
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn websocket_listener_rejects_browser_origin_handshakes() -> Result<()> {
|
||||
let host = HostHarness::start("ws://127.0.0.1:0").await?;
|
||||
let mut request = host.endpoint.as_str().into_client_request()?;
|
||||
request
|
||||
.headers_mut()
|
||||
.insert(ORIGIN, HeaderValue::from_static("https://evil.example"));
|
||||
|
||||
let error = match connect_async(request).await {
|
||||
Ok(_) => anyhow::bail!("browser-origin websocket handshake should be rejected"),
|
||||
Err(error) => error,
|
||||
};
|
||||
let WebSocketError::Http(response) = error else {
|
||||
anyhow::bail!("browser-origin websocket handshake failed unexpectedly: {error}");
|
||||
};
|
||||
assert_eq!(response.status(), StatusCode::FORBIDDEN);
|
||||
Ok(())
|
||||
}
|
||||
Reference in New Issue
Block a user