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

Remove WebSocket transport from code-mode hosts (#40692)

Browse Source 
## 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
2026-08-25 20:37:15 +00:00
committed by copyberry
parent 9be8d6e1c3
commit 725b3a44f8
31 changed files with 153 additions and 2766 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/code-mode-host/src/host_tests.rs
+1 -111
View File
@@ -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(&registry)).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()),
] {
codex-rs/code-mode-host/src/lib.rs
+25 -125
View File
@@ -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 {
codex-rs/code-mode-host/src/main.rs
+1 -1
View File
@@ -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",
codex-rs/code-mode-host/src/peer.rs
+3 -16
View File
@@ -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();
codex-rs/code-mode-host/src/transport.rs
+1 -318
View File
@@ -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;
codex-rs/code-mode-host/src/transport_tests.rs
+3 -75
View File
@@ -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 =