openai/codex
1
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-11-02 19:43:24 +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/Cargo.lock
Generated
-3
View File
@@ -2573,7 +2573,6 @@ dependencies = [
"codex-http-client",
"codex-install-context",
"codex-protocol",
"codex-websocket-client",
"futures",
"http-body-util",
"pretty_assertions",
@@ -2581,7 +2580,6 @@ dependencies = [
"reqwest 0.12.28",
"serde_json",
"tokio",
"tokio-tungstenite",
"tokio-util",
"tonic",
"tower",
@@ -2608,7 +2606,6 @@ dependencies = [
"tempfile",
"tokio",
"tokio-stream",
"tokio-tungstenite",
"tokio-util",
"tonic",
"tracing",
codex-rs/app-server/README.md
+1 -1
View File
@@ -36,7 +36,7 @@ When running with `--listen ws://IP:PORT`, the same listener also serves basic H
Websocket transport is currently experimental and unsupported. Do not rely on it for production workloads.
Pass `--code-mode-host URL` to connect this app-server process to a remote code-mode host instead of starting a local host. Use `ws://` or `wss://` for the WebSocket protocol, or a root `http://` or `https://` URL without a path or query for gRPC. Remote hosts require the `code_mode_host` feature. This outbound connection is independent of `--listen` and is shared by the process's threads.
Pass `--code-mode-host URL` to connect this app-server process to a remote code-mode host instead of starting a local host. Use a root `http://` or `https://` URL without a path or query for gRPC. Remote hosts require the `code_mode_host` feature. This outbound connection is independent of `--listen` and is shared by the process's threads.
The unix socket transport is intended for local app-server control-plane clients. `codex app-server proxy`
opens exactly one raw stream connection to `$CODEX_HOME/app-server-control/app-server-control.sock`
codex-rs/app-server/src/code_mode_host.rs
+8 -16
View File
@@ -23,8 +23,6 @@ pub enum CodeModeHostTransport {
/// Start and own the default local code-mode host.
#[default]
Local,
/// Share a connection to the specified remote code-mode host.
WebSocket(Url),
/// Share an HTTP/2 gRPC connection to the specified remote code-mode host.
Grpc(Url),
}
@@ -32,8 +30,7 @@ pub enum CodeModeHostTransport {
impl From<AppServerCodeModeHostArgs> for CodeModeHostTransport {
fn from(args: AppServerCodeModeHostArgs) -> Self {
match args.code_mode_host {
Some(url) if matches!(url.scheme(), "http" | "https") => Self::Grpc(url),
Some(url) => Self::WebSocket(url),
Some(url) => Self::Grpc(url),
None => Self::Local,
}
}
@@ -66,22 +63,17 @@ impl TypedValueParser for RedactedHostUrlParser {
fn parse_host_url(value: &str) -> Result<Url, String> {
let url = Url::parse(value).map_err(|error| format!("invalid code-mode host URL: {error}"))?;
if !matches!(url.scheme(), "ws" | "wss" | "http" | "https") || url.host_str().is_none() {
return Err(
"code-mode host URL must use ws://, wss://, http://, or https:// with a host"
.to_string(),
);
if !matches!(url.scheme(), "http" | "https") || url.host_str().is_none() {
return Err("code-mode host URL must use http:// or https:// with a host".to_string());
}
if url.fragment().is_some() {
return Err("code-mode host URL must not contain a fragment".to_string());
}
if matches!(url.scheme(), "http" | "https") {
if !url.username().is_empty() || url.password().is_some() {
return Err("gRPC code-mode host URL must not contain credentials".to_string());
}
if url.path() != "/" || url.query().is_some() {
return Err("gRPC code-mode host URL must not contain a path or query".to_string());
}
if !url.username().is_empty() || url.password().is_some() {
return Err("gRPC code-mode host URL must not contain credentials".to_string());
}
if url.path() != "/" || url.query().is_some() {
return Err("gRPC code-mode host URL must not contain a path or query".to_string());
}
Ok(url)
}
codex-rs/app-server/src/code_mode_host_tests.rs
+7 -29
View File
@@ -4,21 +4,6 @@ use super::parse_host_url;
use pretty_assertions::assert_eq;
use url::Url;
#[test]
fn websocket_host_accepts_local_and_secure_endpoints() {
for endpoint in [
"ws://127.0.0.1:8765",
"wss://example.test/code-mode",
"ws://alice:secret@example.test/code-mode",
"wss://alice:secret@example.test/code-mode",
] {
assert_eq!(
parse_host_url(endpoint),
Ok(Url::parse(endpoint).expect("test endpoint should parse"))
);
}
}
#[test]
fn grpc_host_accepts_local_and_secure_endpoints() {
for endpoint in ["http://127.0.0.1:8765", "https://example.test"] {
@@ -50,8 +35,13 @@ fn code_mode_host_rejects_invalid_endpoints() {
for endpoint in [
"ftp://127.0.0.1:8765",
"ws://",
"not a host endpoint",
"ws://127.0.0.1:8765",
"wss://example.test/code-mode",
"ws://alice:secret@example.test/code-mode",
"wss://alice:secret@example.test/code-mode",
"wss://example.test/code-mode#fragment",
"http://",
"not a host endpoint",
"https://example.test/code-mode#fragment",
"https://example.test/code-mode",
"http://example.test/?token=secret",
@@ -64,25 +54,13 @@ fn code_mode_host_rejects_invalid_endpoints() {
}
#[test]
fn omitted_websocket_host_selects_local_transport() {
fn omitted_host_selects_local_transport() {
assert_eq!(
CodeModeHostTransport::from(AppServerCodeModeHostArgs::default()),
CodeModeHostTransport::Local
);
}
#[test]
fn explicit_websocket_host_selects_remote_transport() {
let url = Url::parse("wss://example.test/code-mode").expect("test endpoint should parse");
assert_eq!(
CodeModeHostTransport::from(AppServerCodeModeHostArgs {
code_mode_host: Some(url.clone()),
}),
CodeModeHostTransport::WebSocket(url)
);
}
#[test]
fn explicit_grpc_host_selects_remote_transport() {
let url = Url::parse("https://example.test").expect("test endpoint should parse");
codex-rs/app-server/src/lib.rs
-15
View File
@@ -4,7 +4,6 @@
use codex_arg0::Arg0DispatchPaths;
use codex_code_mode::CodeModeSessionProvider;
use codex_code_mode::GrpcCodeModeSessionProvider;
use codex_code_mode::WebSocketCodeModeSessionProvider;
use codex_config::LoaderOverrides;
use codex_config::NoopThreadConfigLoader;
use codex_core::config::Config;
@@ -552,20 +551,6 @@ pub async fn run_main_with_transport_options(
let code_mode_session_provider: Option<Arc<dyn CodeModeSessionProvider>> =
match &runtime_options.code_mode_host_transport {
CodeModeHostTransport::Local => None,
CodeModeHostTransport::WebSocket(url) => {
if !config.features.enabled(Feature::CodeModeHost) {
return Err(std::io::Error::new(
ErrorKind::InvalidInput,
"remote code-mode host requires the code_mode_host feature to be enabled",
));
}
Some(Arc::new(
WebSocketCodeModeSessionProvider::with_http_client_factory(
url.to_string(),
config.http_client_factory(),
),
))
}
CodeModeHostTransport::Grpc(url) => {
if !config.features.enabled(Feature::CodeModeHost) {
return Err(std::io::Error::new(
codex-rs/app-server/src/main_tests.rs
+6 -19
View File
@@ -38,25 +38,6 @@ fn app_server_accepts_cli_config_overrides() {
);
}
#[test]
fn app_server_accepts_process_scoped_code_mode_host() {
let args = AppServerArgs::try_parse_from([
"codex-app-server",
"--code-mode-host",
"wss://example.test/code-mode",
"--listen",
"off",
])
.expect("parse app-server args");
assert_eq!(
args.code_mode_host.code_mode_host,
Some(Url::parse("wss://example.test/code-mode").expect("test endpoint should parse"))
);
assert_eq!(args.listen, AppServerTransport::Off);
assert_eq!(args.config_overrides.raw_overrides, Vec::<String>::new());
}
#[test]
fn app_server_accepts_process_scoped_grpc_code_mode_host() {
let args = AppServerArgs::try_parse_from([
@@ -80,7 +61,13 @@ fn app_server_rejects_invalid_code_mode_host() {
for endpoint in [
"ftp://127.0.0.1:8765",
"ws://",
"ws://127.0.0.1:8765",
"wss://example.test/code-mode",
"ws://alice:secret@example.test/code-mode",
"wss://alice:secret@example.test/code-mode",
"wss://example.test/code-mode#fragment",
"http://",
"https://example.test/#fragment",
"https://example.test/code-mode",
"http://alice:secret@example.test",
"https://alice:secret@example.test",
codex-rs/app-server/tests/suite/v2/code_mode_host.rs
+30 -39
View File
@@ -24,49 +24,11 @@ const DEFAULT_READ_TIMEOUT: Duration = Duration::from_secs(/*secs*/ 60);
#[cfg(not(any(target_os = "macos", windows)))]
const DEFAULT_READ_TIMEOUT: Duration = Duration::from_secs(/*secs*/ 10);
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn app_server_shares_flag_selected_code_mode_host_across_threads() -> Result<()> {
assert_shared_remote_code_mode_host("ws://127.0.0.1:0").await
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn app_server_shares_flag_selected_grpc_code_mode_host_across_threads() -> Result<()> {
assert_shared_remote_code_mode_host("grpc://127.0.0.1:0").await
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn app_server_prewarms_flag_selected_code_mode_host_before_first_turn() -> Result<()> {
let model_server = responses::start_mock_server().await;
let codex_home = TempDir::new()?;
MockResponsesConfig::new(&model_server.uri())
.enable_feature(Feature::CodeModePrewarm)
.write(codex_home.path())?;
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await?;
let host_url = format!("ws://{}", listener.local_addr()?);
let mut app_server = TestAppServer::builder()
.with_codex_home(codex_home.path())
.with_args(&["--code-mode-host", &host_url])
.build_initialized_with_timeout(DEFAULT_READ_TIMEOUT)
.await?;
app_server
.start_thread(ThreadStartParams::default())
.await?;
let (_stalled_connection, _) = timeout(DEFAULT_READ_TIMEOUT, listener.accept())
.await
.context("code-mode host was not contacted before the first turn")??;
let status = timeout(Duration::from_secs(5), app_server.shutdown_gracefully())
.await
.context("stalled code-mode prewarm blocked thread shutdown")??;
assert!(status.success(), "app-server did not exit successfully");
Ok(())
}
async fn assert_shared_remote_code_mode_host(listen_url: &str) -> Result<()> {
let host_program = codex_utils_cargo_bin::cargo_bin("codex-code-mode-host")?;
let mut code_mode_host = Command::new(host_program)
.args(["--listen", listen_url])
.args(["--listen", "grpc://127.0.0.1:0"])
.stdin(Stdio::null())
.stdout(Stdio::piped())
.stderr(Stdio::inherit())
@@ -174,3 +136,32 @@ async fn assert_shared_remote_code_mode_host(listen_url: &str) -> Result<()> {
Ok(())
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn app_server_prewarms_flag_selected_grpc_code_mode_host_before_first_turn() -> Result<()> {
let model_server = responses::start_mock_server().await;
let codex_home = TempDir::new()?;
MockResponsesConfig::new(&model_server.uri())
.enable_feature(Feature::CodeModePrewarm)
.write(codex_home.path())?;
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await?;
let host_url = format!("http://{}", listener.local_addr()?);
let mut app_server = TestAppServer::builder()
.with_codex_home(codex_home.path())
.with_args(&["--code-mode-host", &host_url])
.build_initialized_with_timeout(DEFAULT_READ_TIMEOUT)
.await?;
app_server
.start_thread(ThreadStartParams::default())
.await?;
let (_stalled_connection, _) = timeout(DEFAULT_READ_TIMEOUT, listener.accept())
.await
.context("code-mode host was not contacted before the first turn")??;
let status = timeout(Duration::from_secs(5), app_server.shutdown_gracefully())
.await
.context("stalled code-mode prewarm blocked thread shutdown")??;
assert!(status.success(), "app-server did not exit successfully");
Ok(())
}
codex-rs/cli/src/main.rs
+6 -29
View File
@@ -4487,35 +4487,6 @@ mod tests {
assert!(err.to_string().contains("is empty"));
}
#[test]
fn app_server_code_mode_host_url_parses_independently_of_listen_transport() {
let app_server = app_server_from_args(
[
"codex",
"app-server",
"--code-mode-host",
"wss://example.test/code-mode",
"--listen",
"ws://127.0.0.1:4500",
]
.as_ref(),
);
assert_eq!(
app_server.code_mode_host.code_mode_host,
Some(
url::Url::parse("wss://example.test/code-mode")
.expect("test endpoint should parse")
)
);
assert_eq!(
app_server.listen,
codex_app_server::AppServerTransport::WebSocket {
bind_address: "127.0.0.1:4500".parse().expect("valid socket address"),
}
);
}
#[test]
fn app_server_grpc_code_mode_host_url_parses_independently_of_listen_transport() {
let app_server = app_server_from_args(
@@ -4541,7 +4512,13 @@ mod tests {
for endpoint in [
"ftp://127.0.0.1:8765",
"ws://",
"ws://127.0.0.1:8765",
"wss://example.test/code-mode",
"ws://alice:secret@example.test/code-mode",
"wss://alice:secret@example.test/code-mode",
"wss://example.test/code-mode#fragment",
"http://",
"https://example.test/#fragment",
"https://example.test/code-mode",
"http://alice:secret@example.test",
"https://alice:secret@example.test",
codex-rs/code-mode-host/Cargo.toml
+1 -2
View File
@@ -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 }
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 =
codex-rs/code-mode-host/tests/websocket.rs
-842
View File
@@ -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,
..
} = &notification
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(())
}
codex-rs/code-mode-protocol/src/host/host_tests.rs
-125
View File
@@ -22,7 +22,6 @@ use super::ProtocolVersion;
use super::RequestId;
use super::SessionId;
use super::SupportedProtocolVersions;
use super::TransportLane;
use super::WireCellId;
use super::WireContentItem;
use super::WireExecuteRequest;
@@ -75,130 +74,6 @@ where
);
}
#[test]
fn dual_websocket_hello_preserves_the_pairing_token() {
assert_wire_round_trip(
HostToClient::HostHello(
HostHello::new(
ProtocolVersion::V1,
CapabilitySet::try_new([capability("dual-websocket-v1")])
.expect("valid capabilities"),
)
.with_bulk_connection_token("pairing-token".to_string()),
),
json!({
"type": "connection/ready",
"selectedVersion": 1,
"capabilities": ["dual-websocket-v1"],
"bulkConnectionToken": "pairing-token",
}),
);
}
#[test]
fn message_families_use_dedicated_transport_lanes() {
for (message, lane) in [
(
ClientToHost::CancelRequest {
id: request_id(/*value*/ 1),
},
TransportLane::Control,
),
(
ClientToHost::DelegateResponse {
id: delegate_request_id(/*value*/ 1),
result: WireResult::Ok {
value: DelegateResponse::NotificationDelivered,
},
},
TransportLane::Control,
),
(
ClientToHost::DelegateResponse {
id: delegate_request_id(/*value*/ 2),
result: WireResult::Ok {
value: DelegateResponse::ToolResult {
result: json!({ "value": "tool result" }),
},
},
},
TransportLane::Bulk,
),
(
ClientToHost::DelegateResponse {
id: delegate_request_id(/*value*/ 3),
result: WireResult::Err {
message: "delegate failed".to_string(),
},
},
TransportLane::Bulk,
),
] {
assert_eq!(message.transport_lane(), lane);
assert!(message.allows_transport_lane(lane));
assert!(!message.allows_transport_lane(match lane {
TransportLane::Control => TransportLane::Bulk,
TransportLane::Bulk => TransportLane::Control,
}));
}
for (message, lane) in [
(
HostToClient::Response {
id: request_id(/*value*/ 1),
result: WireResult::Err {
message: "x".repeat(128 * 1024),
},
},
TransportLane::Control,
),
(
HostToClient::DelegateRequest {
id: delegate_request_id(/*value*/ 1),
session_id: session_id(),
request: DelegateRequest::Notify {
call_id: "call-1".to_string(),
cell_id: cell_id("cell-1"),
text: "important".to_string(),
},
},
TransportLane::Control,
),
(
HostToClient::DelegateRequest {
id: delegate_request_id(/*value*/ 2),
session_id: session_id(),
request: DelegateRequest::InvokeTool {
invocation: WireNestedToolCall {
cell_id: cell_id("cell-1"),
runtime_tool_call_id: "runtime-call-1".to_string(),
tool_name: WireToolName {
name: "tool".to_string(),
namespace: None,
},
tool_kind: WireToolKind::Function,
input: None,
},
},
},
TransportLane::Bulk,
),
(
HostToClient::CancelDelegateRequest {
id: delegate_request_id(/*value*/ 1),
},
TransportLane::Bulk,
),
] {
assert_eq!(message.transport_lane(), lane);
assert!(message.allows_transport_lane(lane));
assert!(!message.allows_transport_lane(match lane {
TransportLane::Control => TransportLane::Bulk,
TransportLane::Bulk => TransportLane::Control,
}));
}
}
fn execute_request() -> WireExecuteRequest {
WireExecuteRequest {
tool_call_id: "call-1".to_string(),
codex-rs/code-mode-protocol/src/host/message.rs
-64
View File
@@ -12,7 +12,6 @@ use super::ProtocolVersion;
use super::RequestId;
use super::SessionId;
use super::SupportedProtocolVersions;
use super::TransportLane;
use super::WireCellId;
use super::WireExecuteRequest;
use super::WireNestedToolCall;
@@ -107,8 +106,6 @@ impl std::error::Error for ClientHelloError {}
pub struct HostHello {
selected_version: ProtocolVersion,
capabilities: CapabilitySet,
#[serde(default, skip_serializing_if = "Option::is_none")]
bulk_connection_token: Option<String>,
}
impl HostHello {
@@ -116,15 +113,9 @@ impl HostHello {
Self {
selected_version,
capabilities,
bulk_connection_token: None,
}
}
pub fn with_bulk_connection_token(mut self, token: String) -> Self {
self.bulk_connection_token = Some(token);
self
}
pub fn selected_version(&self) -> ProtocolVersion {
self.selected_version
}
@@ -132,10 +123,6 @@ impl HostHello {
pub fn capabilities(&self) -> &CapabilitySet {
&self.capabilities
}
pub fn bulk_connection_token(&self) -> Option<&str> {
self.bulk_connection_token.as_deref()
}
}
/// Messages sent from a client to the code-mode host.
@@ -155,30 +142,6 @@ pub enum ClientToHost {
},
}
impl ClientToHost {
/// Keeps notification acknowledgments with control traffic and tool results on the bulk lane.
pub fn transport_lane(&self) -> TransportLane {
match self {
Self::DelegateResponse {
result:
WireResult::Ok {
value: DelegateResponse::NotificationDelivered,
},
..
}
| Self::ClientHello(_)
| Self::Request { .. }
| Self::CancelRequest { .. } => TransportLane::Control,
Self::DelegateResponse { .. } => TransportLane::Bulk,
}
}
/// Validates the message families accepted by each paired socket.
pub fn allows_transport_lane(&self, lane: TransportLane) -> bool {
self.transport_lane() == lane
}
}
/// Messages sent from the code-mode host to a client.
#[derive(Debug, Deserialize, PartialEq, Serialize)]
#[serde(deny_unknown_fields, tag = "type", rename_all_fields = "camelCase")]
@@ -212,33 +175,6 @@ pub enum HostToClient {
},
}
impl HostToClient {
/// Keeps notifications with control traffic and nested-tool callbacks on the bulk lane.
pub fn transport_lane(&self) -> TransportLane {
match self {
Self::DelegateRequest {
request: DelegateRequest::InvokeTool { .. },
..
}
| Self::CancelDelegateRequest { .. } => TransportLane::Bulk,
Self::DelegateRequest {
request: DelegateRequest::Notify { .. },
..
}
| Self::HostHello(_)
| Self::HandshakeRejected { .. }
| Self::Response { .. }
| Self::InitialResponse { .. }
| Self::CellClosed { .. } => TransportLane::Control,
}
}
/// Rejects messages received on the wrong paired socket.
pub fn allows_transport_lane(&self, lane: TransportLane) -> bool {
self.transport_lane() == lane
}
}
#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)]
#[serde(deny_unknown_fields, tag = "method", rename_all_fields = "camelCase")]
pub enum HostRequest {
codex-rs/code-mode-protocol/src/host/mod.rs
+1 -12
View File
@@ -1,8 +1,7 @@
//! Messages and framing for the code-mode host boundary.
//!
//! Protocol version 1 multiplexes session operations and delegate callbacks by
//! request ID over one ordered connection. WebSocket peers can negotiate a
//! separate bulk connection without changing the existing inner messages.
//! request ID over one ordered connection.
mod codec;
mod error;
@@ -13,19 +12,9 @@ mod types;
/// Maximum number of unresolved delegate callbacks allowed per host connection.
pub const MAX_PENDING_DELEGATE_CALLS: usize = 1_024;
/// Optional second WebSocket carrying delegate callbacks and their responses.
pub const DUAL_WEBSOCKET_CAPABILITY: &str = "dual-websocket-v1";
/// Negotiated support for cell execution resource limits on `session/open`.
pub const SESSION_RESOURCE_LIMITS_CAPABILITY: &str = "session-cell-execution-resource-limits";
/// Selects one socket of a negotiated dual-WebSocket connection.
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum TransportLane {
Control,
Bulk,
}
pub use codec::EncodedFrame;
pub use codec::FramedReader;
pub use codec::FramedWriter;
codex-rs/code-mode/Cargo.toml
-2
View File
@@ -17,14 +17,12 @@ codex-code-mode-protocol = { workspace = true }
codex-http-client = { workspace = true }
codex-install-context = { workspace = true }
codex-protocol = { workspace = true }
codex-websocket-client = { workspace = true }
futures = { workspace = true }
http-body-util = "0.1.3"
prost = "0.14.3"
reqwest = { workspace = true, features = ["stream"] }
serde_json = { workspace = true }
tokio = { workspace = true, features = ["io-util", "macros", "net", "process", "rt", "sync", "time"] }
tokio-tungstenite = { workspace = true }
tokio-util = { workspace = true, features = ["rt"] }
tonic = { workspace = true }
tower = { version = "0.5.3", features = ["util"] }
codex-rs/code-mode/src/lib.rs
-1
View File
@@ -6,4 +6,3 @@ pub use grpc_session::GrpcCodeModeSessionProvider;
pub use remote_session::DisabledCodeModeSessionProvider;
pub use remote_session::ProcessOwnedCodeModeSession;
pub use remote_session::ProcessOwnedCodeModeSessionProvider;
pub use remote_session::WebSocketCodeModeSessionProvider;
codex-rs/code-mode/src/remote_session.rs
+5 -85
View File
@@ -18,8 +18,6 @@ use codex_code_mode_protocol::StartedCell;
use codex_code_mode_protocol::WaitOutcome;
use codex_code_mode_protocol::WaitRequest;
use codex_code_mode_protocol::host::SessionId;
use codex_http_client::HttpClientFactory;
use codex_http_client::OutboundProxyPolicy;
use codex_install_context::InstallContext;
use tokio::sync::Semaphore;
use tokio::sync::watch;
@@ -43,11 +41,6 @@ pub struct ProcessOwnedCodeModeSessionProvider {
#[derive(Default)]
pub struct DisabledCodeModeSessionProvider;
/// Creates code-mode sessions backed by one shared remote WebSocket connection.
pub struct WebSocketCodeModeSessionProvider {
host: Arc<OwnedCodeModeHost>,
}
impl ProcessOwnedCodeModeSessionProvider {
pub fn with_host_program(host_program: PathBuf) -> Self {
Self {
@@ -68,9 +61,7 @@ impl Default for ProcessOwnedCodeModeSessionProvider {
impl CodeModeSessionProvider for ProcessOwnedCodeModeSessionProvider {
fn availability(&self) -> Result<(), String> {
let HostEndpoint::Process(host_program) = &self.host.endpoint else {
unreachable!("a process-owned provider always has a process endpoint");
};
let host_program = &self.host.host_program;
if host_program.is_file() {
Ok(())
} else {
@@ -119,49 +110,6 @@ impl CodeModeSessionProvider for DisabledCodeModeSessionProvider {
}
}
impl WebSocketCodeModeSessionProvider {
pub fn new(websocket_url: String) -> Self {
Self::with_http_client_factory(
websocket_url,
HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault),
)
}
/// Creates a remote host using the application's effective proxy and TLS policy.
pub fn with_http_client_factory(
websocket_url: String,
http_client_factory: HttpClientFactory,
) -> Self {
Self {
host: Arc::new(OwnedCodeModeHost::websocket(
websocket_url,
http_client_factory,
)),
}
}
}
impl CodeModeSessionProvider for WebSocketCodeModeSessionProvider {
fn create_session<'a>(
&'a self,
delegate: Arc<dyn CodeModeSessionDelegate>,
) -> CodeModeSessionProviderFuture<'a> {
self.create_session_with_limits(delegate, CodeModeSessionCellExecutionLimits::default())
}
fn create_session_with_limits<'a>(
&'a self,
delegate: Arc<dyn CodeModeSessionDelegate>,
limits: CodeModeSessionCellExecutionLimits,
) -> CodeModeSessionProviderFuture<'a> {
Box::pin(create_host_session(
delegate,
Arc::clone(&self.host),
limits,
))
}
}
async fn create_host_session(
delegate: Arc<dyn CodeModeSessionDelegate>,
host: Arc<OwnedCodeModeHost>,
@@ -172,16 +120,8 @@ async fn create_host_session(
Ok(Arc::new(session))
}
enum HostEndpoint {
Process(PathBuf),
WebSocket {
websocket_url: String,
http_client_factory: HttpClientFactory,
},
}
struct OwnedCodeModeHost {
endpoint: HostEndpoint,
host_program: PathBuf,
connection: StdMutex<Option<Arc<Connection>>>,
connect_permit: Semaphore,
connection_generation: AtomicU64,
@@ -192,21 +132,7 @@ struct OwnedCodeModeHost {
impl OwnedCodeModeHost {
fn new(host_program: PathBuf) -> Self {
Self {
endpoint: HostEndpoint::Process(host_program),
connection: StdMutex::new(None),
connect_permit: Semaphore::new(/*permits*/ 1),
connection_generation: AtomicU64::new(0),
last_connection_error: StdMutex::new(None),
next_session_id: AtomicU64::new(1),
}
}
fn websocket(websocket_url: String, http_client_factory: HttpClientFactory) -> Self {
Self {
endpoint: HostEndpoint::WebSocket {
websocket_url,
http_client_factory,
},
host_program,
connection: StdMutex::new(None),
connect_permit: Semaphore::new(/*permits*/ 1),
connection_generation: AtomicU64::new(0),
@@ -238,13 +164,7 @@ impl OwnedCodeModeHost {
{
return Err(ConnectionError::Other(error.clone()));
}
let connection = match &self.endpoint {
HostEndpoint::Process(host_program) => Connection::spawn(host_program).await,
HostEndpoint::WebSocket {
websocket_url,
http_client_factory,
} => Connection::connect_websocket(websocket_url, http_client_factory).await,
};
let connection = Connection::spawn(&self.host_program).await;
let new_connection = match connection {
Ok(connection) => connection,
Err(error) => {
@@ -312,7 +232,7 @@ struct SessionInner {
retired_cleanups: StdMutex<Vec<SessionCleanup>>,
}
/// A logical code-mode session assigned to a process or WebSocket host.
/// A logical code-mode session assigned to a process host.
pub struct ProcessOwnedCodeModeSession {
inner: Arc<SessionInner>,
}
codex-rs/code-mode/src/remote_session/connection.rs
+44 -275
View File
@@ -20,32 +20,24 @@ use codex_code_mode_protocol::host::Capability;
use codex_code_mode_protocol::host::CapabilitySet;
use codex_code_mode_protocol::host::ClientHello;
use codex_code_mode_protocol::host::ClientToHost;
use codex_code_mode_protocol::host::DUAL_WEBSOCKET_CAPABILITY;
use codex_code_mode_protocol::host::EncodedFrame;
use codex_code_mode_protocol::host::FramedReader;
use codex_code_mode_protocol::host::FramedWriter;
use codex_code_mode_protocol::host::HostToClient;
use codex_code_mode_protocol::host::MAX_FRAME_BYTES;
use codex_code_mode_protocol::host::MAX_PENDING_DELEGATE_CALLS;
use codex_code_mode_protocol::host::ProtocolVersion;
use codex_code_mode_protocol::host::RequestId;
use codex_code_mode_protocol::host::SESSION_RESOURCE_LIMITS_CAPABILITY;
use codex_code_mode_protocol::host::SupportedProtocolVersions;
use codex_code_mode_protocol::host::TransportLane;
use codex_http_client::HttpClientFactory;
use codex_protocol::shell_environment::scrub_non_inheritable_env_vars;
use codex_websocket_client::WebSocketConnector;
use futures::StreamExt;
use tokio::io::AsyncBufReadExt;
use tokio::io::BufReader;
use tokio::process::Child;
use tokio::process::ChildStdin;
use tokio::process::ChildStdout;
use tokio::process::Command;
use tokio::sync::mpsc;
use tokio::sync::oneshot;
use tokio::task::JoinHandle;
use tokio_tungstenite::tungstenite::client::IntoClientRequest;
use tokio_tungstenite::tungstenite::http::Uri;
use tokio_tungstenite::tungstenite::protocol::WebSocketConfig;
use tokio_util::sync::CancellationToken;
use tracing::debug;
use tracing::warn;
@@ -57,19 +49,14 @@ use self::driver::DriverLifecycle;
pub(super) use self::driver::RemoteSession;
pub(super) use self::driver::SessionCleanup;
use self::reader::drive_reader;
use self::transport::ConnectionReader;
use self::transport::ConnectionWriter;
mod driver;
mod reader;
mod transport;
const IPC_CHANNEL_CAPACITY: usize = 128;
const HOST_HANDSHAKE_TIMEOUT: Duration = Duration::from_secs(10);
const LOCAL_HOST_STARTUP_TIMEOUT: Duration = Duration::from_secs(30);
// TODO(anp) make this timeout configurable if 60 seconds is insufficient.
const DEFAULT_HOST_WAIT_TRANSPORT_TIMEOUT: Duration = Duration::from_secs(60);
const MAX_WEBSOCKET_FRAME_BYTES: usize = MAX_FRAME_BYTES + std::mem::size_of::<u32>();
// Host spawn errors become model-visible tool output. Bound configured paths
// while preserving the executable-bearing suffix needed to diagnose failures.
const MAX_DISPLAYED_HOST_PROGRAM_BYTES: usize = 512;
@@ -80,7 +67,6 @@ pub(super) enum ConnectionError {
host_program: PathBuf,
error: io::Error,
},
BulkConnectionUnavailable(String),
Other(String),
}
@@ -111,9 +97,7 @@ impl fmt::Display for ConnectionError {
&host_program[suffix_start..]
)
}
Self::BulkConnectionUnavailable(message) | Self::Other(message) => {
formatter.write_str(message)
}
Self::Other(message) => formatter.write_str(message),
}
}
}
@@ -133,7 +117,7 @@ struct CallerCancellation {
}
struct ConnectionSupervisor {
owner: ConnectionOwner,
child: Child,
event_tx: mpsc::Sender<DriverEvent>,
cancellation: CancellationToken,
alive: Arc<AtomicBool>,
@@ -143,50 +127,6 @@ struct ConnectionSupervisor {
writer_task: JoinHandle<Result<(), String>>,
}
enum ConnectionOwner {
Process(Box<Child>),
WebSocket,
}
struct BulkConnectionOptions {
websocket_url: String,
http_client_factory: HttpClientFactory,
}
impl BulkConnectionOptions {
async fn connect(
&self,
token: &str,
) -> Result<(ConnectionReader, ConnectionWriter), ConnectionError> {
let control_uri = self.websocket_url.parse::<Uri>().map_err(|error| {
ConnectionError::Other(format!(
"failed to build code-mode host bulk websocket URL: {error}"
))
})?;
let bulk_path = format!("{}/bulk/{token}", control_uri.path().trim_end_matches('/'));
let bulk_path_and_query = match control_uri.query() {
Some(query) => format!("{bulk_path}?{query}"),
None => bulk_path,
};
let mut bulk_uri_parts = control_uri.into_parts();
bulk_uri_parts.path_and_query = Some(bulk_path_and_query.parse().map_err(|error| {
ConnectionError::Other(format!(
"failed to build code-mode host bulk websocket path: {error}"
))
})?);
let bulk_url = Uri::from_parts(bulk_uri_parts)
.map_err(|error| {
ConnectionError::Other(format!(
"failed to build code-mode host bulk websocket URL: {error}"
))
})?
.to_string();
connect_websocket_transport(&bulk_url, &self.http_client_factory)
.await
.map_err(|error| ConnectionError::BulkConnectionUnavailable(error.to_string()))
}
}
impl CallerCancellation {
fn new() -> Self {
Self {
@@ -253,52 +193,19 @@ impl Connection {
.take()
.ok_or_else(|| ConnectionError::Other("spawned code-mode host has no stdout".into()))?;
Self::establish(
ConnectionReader::Stdio(FramedReader::new(stdout)),
ConnectionWriter::Stdio(FramedWriter::new(stdin)),
ConnectionOwner::Process(Box::new(child)),
/*bulk_connection_options*/ None,
)
.await
}
pub(super) async fn connect_websocket(
websocket_url: &str,
http_client_factory: &HttpClientFactory,
) -> Result<Self, ConnectionError> {
let (reader, writer) =
connect_websocket_transport(websocket_url, http_client_factory).await?;
Self::establish(
reader,
writer,
ConnectionOwner::WebSocket,
Some(BulkConnectionOptions {
websocket_url: websocket_url.to_string(),
http_client_factory: http_client_factory.clone(),
}),
)
.await
Self::establish(FramedReader::new(stdout), FramedWriter::new(stdin), child).await
}
async fn establish(
mut reader: ConnectionReader,
mut writer: ConnectionWriter,
mut owner: ConnectionOwner,
bulk_connection_options: Option<BulkConnectionOptions>,
mut reader: FramedReader<ChildStdout>,
mut writer: FramedWriter<ChildStdin>,
mut child: Child,
) -> Result<Self, ConnectionError> {
let handshake = async {
let dual_capability =
Capability::new(DUAL_WEBSOCKET_CAPABILITY).map_err(|error| error.to_string())?;
let session_limits_capability = Capability::new(SESSION_RESOURCE_LIMITS_CAPABILITY)
.map_err(|error| error.to_string())?;
let optional_capabilities = if bulk_connection_options.is_some() {
CapabilitySet::try_new([dual_capability.clone(), session_limits_capability])
.map_err(|error| error.to_string())?
} else {
CapabilitySet::try_new([session_limits_capability])
.map_err(|error| error.to_string())?
};
let optional_capabilities = CapabilitySet::try_new([session_limits_capability])
.map_err(|error| error.to_string())?;
let hello = ClientHello::new(
SupportedProtocolVersions::try_new([ProtocolVersion::V1])
.map_err(|err| err.to_string())?,
@@ -318,24 +225,7 @@ impl Connection {
Some(HostToClient::HostHello(hello))
if hello.selected_version() == ProtocolVersion::V1 =>
{
let capabilities = hello.capabilities().clone();
let bulk_token = if capabilities.contains(&dual_capability) {
hello
.bulk_connection_token()
.map(str::to_string)
.ok_or_else(|| {
"code-mode host advertised dual websockets without a pairing token"
.to_string()
})
.map(Some)?
} else if hello.bulk_connection_token().is_some() {
return Err(
"code-mode host returned an unexpected bulk pairing token".to_string()
);
} else {
None
};
Ok((capabilities, bulk_token))
Ok(hello.capabilities().clone())
}
Some(HostToClient::HandshakeRejected { reason }) => {
Err(format!("code-mode host rejected the handshake: {reason:?}"))
@@ -346,99 +236,42 @@ impl Connection {
None => Err("code-mode host exited during handshake".to_string()),
}
};
let handshake_timeout = match &owner {
ConnectionOwner::Process(_) => LOCAL_HOST_STARTUP_TIMEOUT,
ConnectionOwner::WebSocket => HOST_HANDSHAKE_TIMEOUT,
};
let handshake_result = match tokio::time::timeout(handshake_timeout, handshake).await {
Ok(result) => result,
Err(_) => {
let _ = writer.close().await;
owner.close().await;
return Err(ConnectionError::Other(
"timed out negotiating with the code-mode host".into(),
));
}
};
let (capabilities, bulk_token) = match handshake_result {
let handshake_result =
match tokio::time::timeout(LOCAL_HOST_STARTUP_TIMEOUT, handshake).await {
Ok(result) => result,
Err(_) => {
kill_and_reap(&mut child).await;
return Err(ConnectionError::Other(
"timed out negotiating with the code-mode host".into(),
));
}
};
let capabilities = match handshake_result {
Ok(negotiated) => negotiated,
Err(err) => {
let _ = writer.close().await;
owner.close().await;
kill_and_reap(&mut child).await;
return Err(ConnectionError::Other(err));
}
};
let (bulk_reader, bulk_writer) = if let Some(token) = bulk_token {
let Some(options) = bulk_connection_options else {
let _ = writer.close().await;
owner.close().await;
return Err(ConnectionError::Other(
"code-mode host negotiated an unsupported bulk websocket".to_string(),
));
};
match options.connect(&token).await {
Ok((reader, writer)) => (Some(reader), Some(writer)),
Err(error) => {
let _ = writer.close().await;
owner.close().await;
return Err(error);
}
}
} else {
(None, None)
};
let (command_tx, command_rx) = mpsc::channel(IPC_CHANNEL_CAPACITY);
let (event_tx, event_rx) = mpsc::channel(IPC_CHANNEL_CAPACITY);
let (outgoing_tx, outgoing_rx) = mpsc::channel::<EncodedFrame>(IPC_CHANNEL_CAPACITY);
let (bulk_tx, bulk_rx) = if bulk_writer.is_some() {
let (sender, receiver) = mpsc::channel::<EncodedFrame>(MAX_PENDING_DELEGATE_CALLS);
(Some(sender), Some(receiver))
} else {
(None, None)
};
let dual_websocket = bulk_writer.is_some();
let cancellation = CancellationToken::new();
let alive = Arc::new(AtomicBool::new(true));
let failure = Arc::new(std::sync::Mutex::new(None));
let writer_cancellation = cancellation.clone();
let writer_task = tokio::spawn(async move {
if let (Some(bulk_writer), Some(bulk_rx)) = (bulk_writer, bulk_rx) {
tokio::try_join!(
drive_writer(writer, outgoing_rx, writer_cancellation.clone()),
drive_writer(bulk_writer, bulk_rx, writer_cancellation)
)?;
Ok(())
} else {
drive_writer(writer, outgoing_rx, writer_cancellation).await
}
});
let writer_task =
tokio::spawn(
async move { drive_writer(writer, outgoing_rx, writer_cancellation).await },
);
let reader_events = event_tx.clone();
let reader_cancellation = cancellation.clone();
let reader_task = tokio::spawn(async move {
let lane = dual_websocket.then_some(TransportLane::Control);
if let Some(bulk_reader) = bulk_reader {
tokio::try_join!(
drive_reader(
reader,
reader_events.clone(),
reader_cancellation.clone(),
lane,
),
drive_reader(
bulk_reader,
reader_events,
reader_cancellation,
Some(TransportLane::Bulk),
)
)?;
Ok(())
} else {
drive_reader(reader, reader_events, reader_cancellation, lane).await
}
});
let reader_task =
tokio::spawn(
async move { drive_reader(reader, reader_events, reader_cancellation).await },
);
let (driver, execute_claim_tx) = ConnectionDriver::new(
command_rx,
@@ -451,14 +284,10 @@ impl Connection {
cancellation: cancellation.clone(),
},
);
let driver = match bulk_tx {
Some(sender) => driver.with_bulk_sender(sender),
None => driver,
};
let driver_task = tokio::spawn(driver.run());
tokio::spawn(
ConnectionSupervisor {
owner,
child,
event_tx,
cancellation: cancellation.clone(),
alive: Arc::clone(&alive),
@@ -653,65 +482,21 @@ impl Connection {
}
}
async fn connect_websocket_transport(
websocket_url: &str,
http_client_factory: &HttpClientFactory,
) -> Result<(ConnectionReader, ConnectionWriter), ConnectionError> {
let request = websocket_url.into_client_request().map_err(|error| {
ConnectionError::Other(format!(
"failed to build code-mode host websocket request: {error}"
))
})?;
let connector = WebSocketConnector::new(http_client_factory)
.map_err(|error| {
ConnectionError::Other(format!(
"failed to configure code-mode host websocket TLS: {error}"
))
})?
.with_tcp_nodelay();
let websocket_config = WebSocketConfig::default()
.max_frame_size(Some(MAX_WEBSOCKET_FRAME_BYTES))
.max_message_size(Some(MAX_WEBSOCKET_FRAME_BYTES));
let (websocket, _) = tokio::time::timeout(
HOST_HANDSHAKE_TIMEOUT,
connector.connect(request, websocket_config),
)
.await
.map_err(|_| {
ConnectionError::Other("timed out connecting to the code-mode host websocket".into())
})?
.map_err(|error| {
ConnectionError::Other(format!(
"failed to connect to the code-mode host websocket: {error}"
))
})?;
let (writer, reader) = websocket.split();
Ok((
ConnectionReader::WebSocket(reader),
ConnectionWriter::WebSocket(writer),
))
}
async fn drive_writer(
mut writer: ConnectionWriter,
mut writer: FramedWriter<ChildStdin>,
mut outgoing: mpsc::Receiver<EncodedFrame>,
cancellation: CancellationToken,
) -> Result<(), String> {
loop {
tokio::select! {
_ = cancellation.cancelled() => {
return writer
.close()
.await
.map_err(|error| format!("failed to close code-mode host connection: {error}"));
}
_ = cancellation.cancelled() => return Ok(()),
frame = outgoing.recv() => {
let Some(frame) = frame else {
return Err("code-mode host outgoing stream closed".to_string());
};
tokio::select! {
_ = cancellation.cancelled() => return Ok(()),
result = writer.write_frame(frame) => {
result = writer.write_frame(&frame) => {
result.map_err(|error| {
format!("failed to write code-mode host message: {error}")
})?;
@@ -735,7 +520,7 @@ impl Drop for Connection {
impl ConnectionSupervisor {
async fn run(mut self) {
let mut owner_exited = false;
let mut child_exited = false;
let reason = tokio::select! {
biased;
_ = self.cancellation.cancelled() => failure_message(&self.failure),
@@ -745,35 +530,19 @@ impl ConnectionSupervisor {
},
result = &mut self.reader_task => task_failure("reader", result),
result = &mut self.writer_task => task_failure("writer", result),
reason = self.owner.wait() => {
owner_exited = true;
reason
result = self.child.wait() => {
child_exited = true;
match result {
Ok(status) => format!("code-mode host exited with status {status}"),
Err(error) => format!("failed waiting for code-mode host: {error}"),
}
}
};
mark_connection_dead(&self.alive, &self.failure, reason.clone());
let _ = self.event_tx.try_send(DriverEvent::Failed(reason));
self.cancellation.cancel();
if !owner_exited {
self.owner.close().await;
}
}
}
impl ConnectionOwner {
async fn wait(&mut self) -> String {
match self {
Self::Process(child) => match child.wait().await {
Ok(status) => format!("code-mode host exited with status {status}"),
Err(error) => format!("failed waiting for code-mode host: {error}"),
},
Self::WebSocket => std::future::pending().await,
}
}
async fn close(&mut self) {
match self {
Self::Process(child) => kill_and_reap(child).await,
Self::WebSocket => {}
if !child_exited {
kill_and_reap(&mut self.child).await;
}
}
}
codex-rs/code-mode/src/remote_session/connection/driver.rs
+2 -14
View File
@@ -9,7 +9,6 @@ use codex_code_mode_protocol::CodeModeSessionDelegate;
use codex_code_mode_protocol::host::EncodedFrame;
use codex_code_mode_protocol::host::HostToClient;
use codex_code_mode_protocol::host::RequestId;
use codex_code_mode_protocol::host::TransportLane;
use tokio::sync::mpsc;
use tokio_util::sync::CancellationToken;
@@ -42,7 +41,6 @@ pub(super) struct ConnectionDriver {
event_tx: mpsc::Sender<DriverEvent>,
execute_claim_rx: mpsc::UnboundedReceiver<RequestId>,
outgoing_tx: mpsc::Sender<EncodedFrame>,
bulk_tx: Option<mpsc::Sender<EncodedFrame>>,
requests: RequestTracker,
deferred_host_messages: VecDeque<HostToClient>,
sessions: SessionRegistry,
@@ -69,7 +67,6 @@ impl ConnectionDriver {
event_tx: event_tx.clone(),
execute_claim_rx,
outgoing_tx,
bulk_tx: None,
requests: RequestTracker::new(),
deferred_host_messages: VecDeque::new(),
sessions: SessionRegistry::new(),
@@ -137,17 +134,8 @@ impl ConnectionDriver {
}
}
pub(super) fn with_bulk_sender(mut self, sender: mpsc::Sender<EncodedFrame>) -> Self {
self.bulk_tx = Some(sender);
self
}
fn queue_frame(&mut self, frame: EncodedFrame, lane: TransportLane) -> bool {
let sender = match lane {
TransportLane::Control => &self.outgoing_tx,
TransportLane::Bulk => self.bulk_tx.as_ref().unwrap_or(&self.outgoing_tx),
};
match sender.try_send(frame) {
fn queue_frame(&mut self, frame: EncodedFrame) -> bool {
match self.outgoing_tx.try_send(frame) {
Ok(()) => true,
Err(mpsc::error::TrySendError::Full(_)) => {
self.fail("code-mode host outgoing queue is full".to_string());
codex-rs/code-mode/src/remote_session/connection/driver/commands.rs
+3 -6
View File
@@ -127,8 +127,7 @@ impl ConnectionDriver {
},
&self.event_tx,
);
let lane = message.transport_lane();
self.queue_frame(frame, lane)
self.queue_frame(frame)
}
fn execute(
@@ -187,8 +186,7 @@ impl ConnectionDriver {
},
&self.event_tx,
);
let lane = message.transport_lane();
self.queue_frame(frame, lane)
self.queue_frame(frame)
}
fn wait(
@@ -317,7 +315,6 @@ impl ConnectionDriver {
};
self.requests
.insert_pending(request_id, pending, &self.event_tx);
let lane = message.transport_lane();
self.queue_frame(frame, lane)
self.queue_frame(frame)
}
}
codex-rs/code-mode/src/remote_session/connection/driver/delegate_runtime.rs
+1 -2
View File
@@ -318,8 +318,7 @@ impl ConnectionDriver {
}
}
};
let lane = message.transport_lane();
self.queue_frame(frame, lane)
self.queue_frame(frame)
}
pub(super) fn close_cell(&mut self, session_id: SessionId, cell_id: WireCellId) -> bool {
codex-rs/code-mode/src/remote_session/connection/driver/responses.rs
+1 -2
View File
@@ -374,8 +374,7 @@ impl ConnectionDriver {
return false;
}
};
let lane = message.transport_lane();
self.queue_frame(frame, lane)
self.queue_frame(frame)
}
fn shutdown_abandoned_session(&mut self, session: RemoteSession) -> bool {
codex-rs/code-mode/src/remote_session/connection/driver_tests.rs
-126
View File
@@ -17,7 +17,6 @@ use codex_code_mode_protocol::host::CapabilitySet;
use codex_code_mode_protocol::host::ClientToHost;
use codex_code_mode_protocol::host::DelegateRequest;
use codex_code_mode_protocol::host::DelegateRequestId;
use codex_code_mode_protocol::host::DelegateResponse;
use codex_code_mode_protocol::host::EncodedFrame;
use codex_code_mode_protocol::host::HostRequest;
use codex_code_mode_protocol::host::HostResponse;
@@ -248,11 +247,6 @@ struct RecordingDelegate {
struct PanickingDelegate;
struct LargeResultBurstDelegate {
started: AtomicUsize,
release: CancellationToken,
}
#[derive(Debug, Eq, PartialEq)]
enum HeldDelegateEvent {
Started,
@@ -342,35 +336,6 @@ impl CodeModeSessionDelegate for PanickingDelegate {
fn cell_closed(&self, _cell_id: &CellId) {}
}
impl CodeModeSessionDelegate for LargeResultBurstDelegate {
fn invoke_tool<'a>(
&'a self,
_invocation: CodeModeNestedToolCall,
cancellation_token: CancellationToken,
) -> ToolInvocationFuture<'a> {
self.started.fetch_add(1, Ordering::Release);
let release = self.release.clone();
Box::pin(async move {
tokio::select! {
_ = cancellation_token.cancelled() => Err("cancelled".to_string()),
_ = release.cancelled() => Ok("x".repeat(256 * 1024).into()),
}
})
}
fn notify<'a>(
&'a self,
_call_id: String,
_cell_id: CellId,
_text: String,
_cancellation_token: CancellationToken,
) -> NotificationFuture<'a> {
Box::pin(async { Ok(()) })
}
fn cell_closed(&self, _cell_id: &CellId) {}
}
impl CodeModeSessionDelegate for RecordingDelegate {
fn invoke_tool<'a>(
&'a self,
@@ -692,97 +657,6 @@ async fn delegate_cancel_is_best_effort_and_sends_no_late_response() {
assert_eq!(delegate.notifications.load(Ordering::Relaxed), 0);
}
#[tokio::test]
async fn concurrent_large_delegate_results_do_not_disconnect_a_backpressured_bulk_lane() {
const CONCURRENT_RESULTS: usize = 129;
let (command_tx, command_rx) = mpsc::channel(/*max_capacity*/ 16);
let (event_tx, event_rx) = mpsc::channel(/*max_capacity*/ 16);
let (outgoing_tx, outgoing_rx) = mpsc::channel(/*max_capacity*/ 16);
let (bulk_tx, mut bulk_rx) = mpsc::channel(MAX_PENDING_DELEGATE_CALLS);
let cancellation = CancellationToken::new();
let alive = Arc::new(AtomicBool::new(true));
let failure = Arc::new(StdMutex::new(None));
let (driver, execute_claim_tx) = ConnectionDriver::new(
command_rx,
event_rx,
event_tx.clone(),
outgoing_tx,
DriverLifecycle {
alive: Arc::clone(&alive),
failure: Arc::clone(&failure),
cancellation: cancellation.clone(),
},
);
let driver_task = tokio::spawn(driver.with_bulk_sender(bulk_tx).run());
let mut harness = DriverHarness {
command_tx,
event_tx,
execute_claim_tx,
outgoing_rx,
cancellation,
alive,
failure,
driver_task,
};
let session = remote_session();
let delegate = Arc::new(LargeResultBurstDelegate {
started: AtomicUsize::new(0),
release: CancellationToken::new(),
});
harness.open(session.clone(), delegate.clone()).await;
let _started = harness
.start_cell(session.clone(), /*request_id*/ 2, "1")
.await;
for value in 1..=CONCURRENT_RESULTS {
harness
.start_tool_delegate(&session, DelegateRequestId::new(value as i64))
.await;
}
tokio::time::timeout(Duration::from_secs(10), async {
while delegate.started.load(Ordering::Acquire) < CONCURRENT_RESULTS {
tokio::task::yield_now().await;
}
})
.await
.expect("concurrent delegate calls should all start");
delegate.release.cancel();
tokio::time::timeout(Duration::from_secs(10), async {
while bulk_rx.len() < CONCURRENT_RESULTS {
assert!(
harness.alive.load(Ordering::Acquire),
"bulk queue disconnected before accepting all concurrent tool results"
);
tokio::task::yield_now().await;
}
})
.await
.expect("concurrent large results should queue behind the blocked bulk writer");
let _unrelated = harness.start_cell(session, /*request_id*/ 3, "2").await;
assert!(harness.alive.load(Ordering::Acquire));
for _ in 0..CONCURRENT_RESULTS {
let frame = bulk_rx.recv().await.expect("queued bulk delegate result");
let message = EncodedFrame::decode_framed::<ClientToHost>(&frame.into_framed_bytes())
.expect("decode queued delegate result");
let ClientToHost::DelegateResponse {
result:
WireResult::Ok {
value: DelegateResponse::ToolResult { result },
},
..
} = message
else {
panic!("expected a successful large delegate result");
};
assert_eq!(result.as_str().map(str::len), Some(256 * 1024));
}
assert!(harness.alive.load(Ordering::Acquire));
}
#[tokio::test]
async fn delegate_limit_returns_an_error_without_disconnecting() {
let mut harness = DriverHarness::start();
codex-rs/code-mode/src/remote_session/connection/reader.rs
+3 -9
View File
@@ -1,15 +1,14 @@
use codex_code_mode_protocol::host::TransportLane;
use codex_code_mode_protocol::host::FramedReader;
use tokio::process::ChildStdout;
use tokio::sync::mpsc;
use tokio_util::sync::CancellationToken;
use super::driver::DriverEvent;
use super::transport::ConnectionReader;
pub(super) async fn drive_reader(
mut reader: ConnectionReader,
mut reader: FramedReader<ChildStdout>,
events: mpsc::Sender<DriverEvent>,
cancellation: CancellationToken,
lane: Option<TransportLane>,
) -> Result<(), String> {
loop {
let message = tokio::select! {
@@ -21,11 +20,6 @@ pub(super) async fn drive_reader(
Ok(None) => return Err("code-mode host closed its stdout".to_string()),
Err(err) => return Err(format!("failed to read code-mode host message: {err}")),
};
if let Some(lane) = lane
&& !message.allows_transport_lane(lane)
{
return Err("code-mode host sent a message on the wrong websocket lane".to_string());
}
events
.send(DriverEvent::HostMessage(message))
.await
codex-rs/code-mode/src/remote_session/connection/transport.rs
-103
View File
@@ -1,103 +0,0 @@
use std::io;
use std::time::Duration;
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_websocket_client::WebSocketConnection;
use futures::SinkExt;
use futures::StreamExt;
use futures::stream::SplitSink;
use futures::stream::SplitStream;
use tokio::process::ChildStdin;
use tokio::process::ChildStdout;
use tokio_tungstenite::tungstenite::Message;
const WEBSOCKET_CLOSE_TIMEOUT: Duration = Duration::from_secs(5);
pub(super) enum ConnectionReader {
Stdio(FramedReader<ChildStdout>),
WebSocket(SplitStream<WebSocketConnection>),
}
pub(super) enum ConnectionWriter {
Stdio(FramedWriter<ChildStdin>),
WebSocket(SplitSink<WebSocketConnection, Message>),
}
impl ConnectionReader {
pub(super) async fn read(&mut self) -> io::Result<Option<HostToClient>> {
match self {
Self::Stdio(reader) => reader.read().await,
Self::WebSocket(reader) => loop {
match reader.next().await {
Some(Ok(Message::Binary(frame))) => {
return EncodedFrame::decode_framed(&frame).map(Some);
}
Some(Ok(Message::Ping(_) | Message::Pong(_))) => {}
Some(Ok(Message::Close(_))) | None => return Ok(None),
Some(Ok(Message::Text(_))) => {
return Err(io::Error::new(
io::ErrorKind::InvalidData,
"code-mode host websocket messages must be binary framed messages",
));
}
Some(Ok(Message::Frame(_))) => {
return Err(io::Error::new(
io::ErrorKind::InvalidData,
"code-mode host websocket returned an unexpected raw frame",
));
}
Some(Err(error)) => {
return Err(io::Error::other(format!(
"failed to read code-mode host websocket message: {error}"
)));
}
}
},
}
}
}
impl ConnectionWriter {
pub(super) async fn write(&mut self, message: &ClientToHost) -> io::Result<()> {
self.write_frame(EncodedFrame::encode(message)?).await
}
pub(super) async fn write_frame(&mut self, frame: EncodedFrame) -> io::Result<()> {
match self {
Self::Stdio(writer) => writer.write_frame(&frame).await,
Self::WebSocket(writer) => writer
.send(Message::Binary(frame.into_framed_bytes().into()))
.await
.map_err(|error| {
io::Error::other(format!(
"failed to write code-mode host websocket message: {error}"
))
}),
}
}
pub(super) async fn close(&mut self) -> io::Result<()> {
match self {
Self::Stdio(_) => Ok(()),
Self::WebSocket(writer) => {
tokio::time::timeout(WEBSOCKET_CLOSE_TIMEOUT, writer.close())
.await
.map_err(|_| {
io::Error::new(
io::ErrorKind::TimedOut,
"timed out closing code-mode host websocket connection",
)
})?
.map_err(|error| {
io::Error::other(format!(
"failed to close code-mode host websocket connection: {error}"
))
})
}
}
}
}
codex-rs/code-mode/src/remote_session_tests.rs
-299
View File
@@ -1,43 +1,12 @@
use std::io;
use std::path::PathBuf;
use std::sync::Arc;
use std::time::Duration;
use codex_code_mode_protocol::CodeModeSessionCellExecutionLimits;
use codex_code_mode_protocol::CodeModeSessionProvider;
use codex_code_mode_protocol::ExecuteRequest;
use codex_code_mode_protocol::FunctionCallOutputContentItem;
use codex_code_mode_protocol::RuntimeResponse;
use codex_code_mode_protocol::host::Capability;
use codex_code_mode_protocol::host::CapabilitySet;
use codex_code_mode_protocol::host::ClientToHost;
use codex_code_mode_protocol::host::DUAL_WEBSOCKET_CAPABILITY;
use codex_code_mode_protocol::host::EncodedFrame;
use codex_code_mode_protocol::host::HostHello;
use codex_code_mode_protocol::host::HostRequest;
use codex_code_mode_protocol::host::HostResponse;
use codex_code_mode_protocol::host::HostToClient;
use codex_code_mode_protocol::host::ProtocolVersion;
use codex_code_mode_protocol::host::SESSION_RESOURCE_LIMITS_CAPABILITY;
use codex_code_mode_protocol::host::WireCellId;
use codex_code_mode_protocol::host::WireContentItem;
use codex_code_mode_protocol::host::WireResult;
use codex_code_mode_protocol::host::WireRuntimeResponse;
use codex_http_client::HttpClientFactory;
use codex_http_client::OutboundProxyPolicy;
use futures::SinkExt;
use futures::StreamExt;
use pretty_assertions::assert_eq;
use tokio::io::AsyncReadExt;
use tokio::io::AsyncWriteExt;
use tokio::net::TcpListener;
use tokio::time::timeout;
use tokio_tungstenite::accept_async;
use tokio_tungstenite::tungstenite::Message;
use super::ProcessOwnedCodeModeSession;
use super::ProcessOwnedCodeModeSessionProvider;
use super::WebSocketCodeModeSessionProvider;
use super::connection::ConnectionError;
use crate::NoopCodeModeSessionDelegate;
@@ -86,274 +55,6 @@ fn missing_host_error_preserves_utf8_boundaries_when_truncating_the_path() {
assert!(displayed_path.len() <= 512);
}
#[tokio::test]
async fn websocket_provider_executes_over_shared_connector() {
let listener = TcpListener::bind("127.0.0.1:0")
.await
.expect("websocket test listener should bind");
let websocket_url = format!(
"ws://{}",
listener
.local_addr()
.expect("websocket test listener should have an address")
);
let server = tokio::spawn(async move {
let (stream, _) = listener
.accept()
.await
.expect("websocket test host should accept a connection");
let mut websocket = accept_async(stream)
.await
.expect("websocket test host should complete the HTTP handshake");
while let Some(message) = websocket.next().await {
let message = message.expect("websocket test host should receive a valid message");
let frame = match message {
Message::Binary(frame) => frame,
Message::Ping(_) | Message::Pong(_) => continue,
Message::Close(_) => break,
Message::Text(_) | Message::Frame(_) => {
panic!("websocket test host received an unexpected message: {message:?}");
}
};
let request = EncodedFrame::decode_framed::<ClientToHost>(&frame)
.expect("websocket test host should decode a framed protocol message");
let responses = match request {
ClientToHost::ClientHello(hello) => {
let capability = Capability::new(SESSION_RESOURCE_LIMITS_CAPABILITY)
.expect("session-limit capability");
assert!(hello.optional_capabilities().contains(&capability));
assert_eq!(hello.required_capabilities(), &CapabilitySet::empty());
vec![HostToClient::HostHello(HostHello::new(
ProtocolVersion::V1,
CapabilitySet::empty(),
))]
}
ClientToHost::Request {
id,
request: HostRequest::OpenSession { session_id, .. },
} => vec![HostToClient::Response {
id,
result: WireResult::Ok {
value: HostResponse::SessionReady { session_id },
},
}],
ClientToHost::Request {
id,
request: HostRequest::Execute { request, .. },
} => {
assert_eq!(request.source, "text('shared connector')");
let cell_id = WireCellId::new("1");
vec![
HostToClient::Response {
id,
result: WireResult::Ok {
value: HostResponse::ExecutionStarted {
cell_id: cell_id.clone(),
},
},
},
HostToClient::InitialResponse {
id,
result: WireResult::Ok {
value: WireRuntimeResponse::Result {
cell_id,
content_items: vec![WireContentItem::InputText {
text: "shared connector".to_string(),
}],
error_text: None,
},
},
},
]
}
ClientToHost::Request {
id,
request: HostRequest::ShutdownSession { session_id },
} => vec![HostToClient::Response {
id,
result: WireResult::Ok {
value: HostResponse::SessionClosed { session_id },
},
}],
request => {
panic!("websocket test host received an unexpected request: {request:?}")
}
};
for response in responses {
let frame = EncodedFrame::encode(&response)
.expect("websocket test host should encode a framed response");
websocket
.send(Message::Binary(frame.into_framed_bytes().into()))
.await
.expect("websocket test host should send its response");
}
}
});
let provider = WebSocketCodeModeSessionProvider::with_http_client_factory(
websocket_url,
HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault),
);
let session = provider
.create_session(Arc::new(NoopCodeModeSessionDelegate))
.await
.expect("shared websocket connector should open a code-mode session");
let error = provider
.create_session_with_limits(
Arc::new(NoopCodeModeSessionDelegate),
CodeModeSessionCellExecutionLimits {
max_yield_time_ms: Some(250),
max_heap_size_bytes: None,
},
)
.await
.err()
.expect("legacy host should reject a limited session");
assert_eq!(
error,
format!(
"code-mode host does not support session resource limits: missing `{SESSION_RESOURCE_LIMITS_CAPABILITY}` capability"
)
);
let second_session = provider
.create_session(Arc::new(NoopCodeModeSessionDelegate))
.await
.expect("rejecting limited sessions should preserve the shared legacy-host connection");
second_session
.shutdown()
.await
.expect("second unlimited session should shut down");
drop(second_session);
let response = session
.execute(ExecuteRequest {
tool_call_id: "shared-websocket".to_string(),
enabled_tools: Vec::new(),
source: "text('shared connector')".to_string(),
yield_time_ms: None,
max_output_tokens: None,
})
.await
.expect("shared websocket connector should start a cell")
.initial_response()
.await
.expect("shared websocket connector should return a cell result");
assert_eq!(
response,
RuntimeResponse::Result {
cell_id: codex_code_mode_protocol::CellId::new("1".to_string()),
content_items: vec![FunctionCallOutputContentItem::InputText {
text: "shared connector".to_string(),
}],
error_text: None,
}
);
session
.shutdown()
.await
.expect("shared websocket connector should shut down its session");
drop(session);
drop(provider);
timeout(Duration::from_secs(5), server)
.await
.expect("websocket test host should disconnect promptly")
.expect("websocket test host task should succeed");
}
#[tokio::test]
async fn websocket_provider_fails_when_a_negotiated_bulk_connection_is_unavailable() {
let listener = TcpListener::bind("127.0.0.1:0")
.await
.expect("websocket test listener should bind");
let websocket_url = format!(
"ws://{}/?access_token=shared-token",
listener
.local_addr()
.expect("websocket test listener should have an address")
);
let server = tokio::spawn(async move {
let (stream, _) = listener
.accept()
.await
.expect("websocket test host should accept the first connection");
let mut control = accept_async(stream)
.await
.expect("first control websocket should connect");
let frame = control
.next()
.await
.expect("first client hello")
.expect("first websocket frame")
.into_data();
let ClientToHost::ClientHello(hello) =
EncodedFrame::decode_framed(&frame).expect("decode first client hello")
else {
panic!("expected first client hello");
};
let capability =
Capability::new(DUAL_WEBSOCKET_CAPABILITY).expect("dual websocket capability");
assert!(hello.optional_capabilities().contains(&capability));
let session_limits_capability =
Capability::new(SESSION_RESOURCE_LIMITS_CAPABILITY).expect("session-limit capability");
assert!(
hello
.optional_capabilities()
.contains(&session_limits_capability)
);
let hello = HostToClient::HostHello(
HostHello::new(
ProtocolVersion::V1,
CapabilitySet::try_new([capability.clone()]).expect("host capabilities"),
)
.with_bulk_connection_token("fallback-token".to_string()),
);
let frame = EncodedFrame::encode(&hello).expect("encode dual host hello");
control
.send(Message::Binary(frame.into_framed_bytes().into()))
.await
.expect("send dual host hello");
let (mut bulk, _) = listener
.accept()
.await
.expect("websocket test host should accept the bulk connection");
let mut request = [0_u8; 1024];
let request_len = bulk
.read(&mut request)
.await
.expect("read bulk websocket handshake");
let request = std::str::from_utf8(&request[..request_len]).expect("bulk HTTP request");
assert!(
request.starts_with("GET /bulk/fallback-token?access_token=shared-token HTTP/1.1\r\n")
);
bulk.write_all(b"HTTP/1.1 404 Not Found\r\nContent-Length: 0\r\nConnection: close\r\n\r\n")
.await
.expect("reject unavailable bulk websocket");
drop(bulk);
drop(control);
});
let provider = WebSocketCodeModeSessionProvider::new(websocket_url);
let error = match provider
.create_session(Arc::new(NoopCodeModeSessionDelegate))
.await
{
Ok(_) => panic!("provider should reject an unavailable negotiated bulk websocket"),
Err(error) => error,
};
assert!(
error.contains("404"),
"unexpected negotiated bulk websocket error: {error}"
);
drop(provider);
timeout(Duration::from_secs(5), server)
.await
.expect("negotiated bulk websocket test host should disconnect promptly")
.expect("negotiated bulk websocket test host task should succeed");
}
#[tokio::test]
async fn provider_returns_missing_host_error() {
let provider = ProcessOwnedCodeModeSessionProvider::with_host_program(