From 725b3a44f8d13c8d2f019e1ecbc8ab686466b950 Mon Sep 17 00:00:00 2001 From: Ruslan Nigmatullin Date: Tue, 25 Aug 2026 20:28:35 +0000 Subject: [PATCH] Remove WebSocket transport from code-mode hosts (#40692) ## What changed - Accept only `http://` and `https://` gRPC endpoints for app server `--code-mode-host` connections. - Limit the standalone code-mode host to stdio and gRPC listeners. - Remove the WebSocket session provider, dual-WebSocket negotiation, transport lanes, and related dependencies and tests. - Keep coverage for shared and prewarmed gRPC hosts and reject WebSocket host URLs in CLI parsing tests. GitOrigin-RevId: 8792604e79ffd54ce776074a4dc3f7d55e63752a --- codex-rs/Cargo.lock | 3 - codex-rs/app-server/README.md | 2 +- codex-rs/app-server/src/code_mode_host.rs | 24 +- .../app-server/src/code_mode_host_tests.rs | 36 +- codex-rs/app-server/src/lib.rs | 15 - codex-rs/app-server/src/main_tests.rs | 25 +- .../tests/suite/v2/code_mode_host.rs | 69 +- codex-rs/cli/src/main.rs | 35 +- codex-rs/code-mode-host/Cargo.toml | 3 +- codex-rs/code-mode-host/src/host_tests.rs | 112 +-- codex-rs/code-mode-host/src/lib.rs | 150 +--- codex-rs/code-mode-host/src/main.rs | 2 +- codex-rs/code-mode-host/src/peer.rs | 19 +- codex-rs/code-mode-host/src/transport.rs | 319 +------ .../code-mode-host/src/transport_tests.rs | 78 +- codex-rs/code-mode-host/tests/websocket.rs | 842 ------------------ .../code-mode-protocol/src/host/host_tests.rs | 125 --- .../code-mode-protocol/src/host/message.rs | 64 -- codex-rs/code-mode-protocol/src/host/mod.rs | 13 +- codex-rs/code-mode/Cargo.toml | 2 - codex-rs/code-mode/src/lib.rs | 1 - codex-rs/code-mode/src/remote_session.rs | 90 +- .../src/remote_session/connection.rs | 319 +------ .../src/remote_session/connection/driver.rs | 16 +- .../connection/driver/commands.rs | 9 +- .../connection/driver/delegate_runtime.rs | 3 +- .../connection/driver/responses.rs | 3 +- .../remote_session/connection/driver_tests.rs | 126 --- .../src/remote_session/connection/reader.rs | 12 +- .../remote_session/connection/transport.rs | 103 --- .../code-mode/src/remote_session_tests.rs | 299 ------- 31 files changed, 153 insertions(+), 2766 deletions(-) delete mode 100644 codex-rs/code-mode-host/tests/websocket.rs delete mode 100644 codex-rs/code-mode/src/remote_session/connection/transport.rs diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index d4c6f28438..a72f837ac4 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -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", diff --git a/codex-rs/app-server/README.md b/codex-rs/app-server/README.md index 741f52659f..bda8322766 100644 --- a/codex-rs/app-server/README.md +++ b/codex-rs/app-server/README.md @@ -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` diff --git a/codex-rs/app-server/src/code_mode_host.rs b/codex-rs/app-server/src/code_mode_host.rs index 6a0ceae5e9..97b250a59a 100644 --- a/codex-rs/app-server/src/code_mode_host.rs +++ b/codex-rs/app-server/src/code_mode_host.rs @@ -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 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 { 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) } diff --git a/codex-rs/app-server/src/code_mode_host_tests.rs b/codex-rs/app-server/src/code_mode_host_tests.rs index 6fe3c7f8e3..964a382b7e 100644 --- a/codex-rs/app-server/src/code_mode_host_tests.rs +++ b/codex-rs/app-server/src/code_mode_host_tests.rs @@ -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"); diff --git a/codex-rs/app-server/src/lib.rs b/codex-rs/app-server/src/lib.rs index 3aedcf9500..fcf9ecff1b 100644 --- a/codex-rs/app-server/src/lib.rs +++ b/codex-rs/app-server/src/lib.rs @@ -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> = 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( diff --git a/codex-rs/app-server/src/main_tests.rs b/codex-rs/app-server/src/main_tests.rs index 3334fe9ac1..e48ec9e743 100644 --- a/codex-rs/app-server/src/main_tests.rs +++ b/codex-rs/app-server/src/main_tests.rs @@ -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::::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", diff --git a/codex-rs/app-server/tests/suite/v2/code_mode_host.rs b/codex-rs/app-server/tests/suite/v2/code_mode_host.rs index 82d2a7fa50..ad849c70ac 100644 --- a/codex-rs/app-server/tests/suite/v2/code_mode_host.rs +++ b/codex-rs/app-server/tests/suite/v2/code_mode_host.rs @@ -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(()) +} diff --git a/codex-rs/cli/src/main.rs b/codex-rs/cli/src/main.rs index 15b769f6f3..0bfea04498 100644 --- a/codex-rs/cli/src/main.rs +++ b/codex-rs/cli/src/main.rs @@ -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", diff --git a/codex-rs/code-mode-host/Cargo.toml b/codex-rs/code-mode-host/Cargo.toml index 685397aadf..0045279df5 100644 --- a/codex-rs/code-mode-host/Cargo.toml +++ b/codex-rs/code-mode-host/Cargo.toml @@ -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 } diff --git a/codex-rs/code-mode-host/src/host_tests.rs b/codex-rs/code-mode-host/src/host_tests.rs index f5ac27a5db..7a2215cb7b 100644 --- a/codex-rs/code-mode-host/src/host_tests.rs +++ b/codex-rs/code-mode-host/src/host_tests.rs @@ -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, @@ -260,111 +254,8 @@ impl AsyncWrite for BlockingWriter { } } -#[tokio::test] -async fn failed_host_hello_removes_the_bulk_pairing_reservation() { - let (host_reader, client_writer) = tokio::io::duplex(/*max_buf_size*/ 4096); - let capability = Capability::new(DUAL_WEBSOCKET_CAPABILITY).expect("dual websocket capability"); - let hello = ClientToHost::ClientHello( - ClientHello::new( - SupportedProtocolVersions::try_new([ProtocolVersion::V1]).expect("supported versions"), - CapabilitySet::empty(), - CapabilitySet::try_new([capability]).expect("optional capabilities"), - ) - .expect("client hello"), - ); - FramedWriter::new(client_writer) - .write(&hello) - .await - .expect("write client hello"); - - let bytes = Arc::new(Mutex::new(Vec::new())); - let registry = BulkConnectionRegistry::default(); - let mut reader = ConnectionReader::from_reader(host_reader); - let mut writer = ConnectionWriter::from_writer(FailingHandshakeWriter { - bytes: Arc::clone(&bytes), - }); - let result = negotiate(&mut reader, &mut writer, Some(®istry)).await; - assert!(result.is_err()); - - let message = EncodedFrame::decode_framed::( - &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>>, -} - -impl AsyncWrite for FailingHandshakeWriter { - fn poll_write( - self: Pin<&mut Self>, - _cx: &mut Context<'_>, - bytes: &[u8], - ) -> Poll> { - 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> { - Poll::Ready(Err(std::io::Error::other("host hello write failed"))) - } - - fn poll_shutdown(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll> { - 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::().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()), ] { diff --git a/codex-rs/code-mode-host/src/lib.rs b/codex-rs/code-mode-host/src/lib.rs index 34e13af94c..d61df805fa 100644 --- a/codex-rs/code-mode-host/src/lib.rs +++ b/codex-rs/code-mode-host/src/lib.rs @@ -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, - bulk_connections: Option, -) -> 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::(OUTGOING_CHANNEL_CAPACITY); - let (bulk_tx, bulk_rx) = if bulk_writer.is_some() { - let (sender, receiver) = mpsc::channel::(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( + mut writer: FramedWriter, mut outgoing: mpsc::Receiver, 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( + reader: &mut FramedReader, + writer: &mut FramedWriter, ) -> Result { 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 { diff --git a/codex-rs/code-mode-host/src/main.rs b/codex-rs/code-mode-host/src/main.rs index c7da2080c3..52e2346513 100644 --- a/codex-rs/code-mode-host/src/main.rs +++ b/codex-rs/code-mode-host/src/main.rs @@ -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", diff --git a/codex-rs/code-mode-host/src/peer.rs b/codex-rs/code-mode-host/src/peer.rs index 02b0c263dc..7912051802 100644 --- a/codex-rs/code-mode-host/src/peer.rs +++ b/codex-rs/code-mode-host/src/peer.rs @@ -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, - bulk_tx: Option>, pending: Mutex>, delegate_permits: Arc, cell_routes: StdMutex>, @@ -64,7 +62,6 @@ impl HostPeer { pub(super) fn new(outgoing_tx: mpsc::Sender) -> 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) -> 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(); diff --git a/codex-rs/code-mode-host/src/transport.rs b/codex-rs/code-mode-host/src/transport.rs index 3817f3b605..2c56855726 100644 --- a/codex-rs/code-mode-host/src/transport.rs +++ b/codex-rs/code-mode-host/src/transport.rs @@ -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::(); -const MAX_PENDING_BULK_CONNECTIONS: usize = MAX_IN_FLIGHT_REQUESTS; - -type BoxedReader = Box; -type BoxedWriter = Box; - #[derive(Debug, Clone, Eq, PartialEq)] enum ListenTransport { Stdio, - WebSocket(SocketAddr), Grpc(SocketAddr), } -pub(crate) enum ConnectionReader { - Framed(FramedReader), - WebSocket(SplitStream), -} - -pub(crate) enum ConnectionWriter { - Framed(FramedWriter), - WebSocket(SplitSink), -} - -#[derive(Clone)] -struct WebSocketListenerState { - limits: Arc, - bulk_connections: BulkConnectionRegistry, -} - -pub(crate) struct BulkConnection { - pub(crate) reader: ConnectionReader, - pub(crate) writer: ConnectionWriter, -} - -#[derive(Clone, Default)] -pub(crate) struct BulkConnectionRegistry { - pending: Arc>>>, -} - -pub(crate) struct BulkConnectionRegistration { - registry: BulkConnectionRegistry, - token: Uuid, - receiver: oneshot::Receiver, -} - -impl BulkConnectionRegistry { - pub(crate) fn reserve(&self) -> Option { - 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> { - 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 { - (&mut self.receiver).await - } -} - -impl Drop for BulkConnectionRegistration { - fn drop(&mut self) { - self.registry.remove(self.token); - } -} - -impl ConnectionReader { - pub(crate) fn from_reader(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> { - 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(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 { return Ok(ListenTransport::Stdio); } - if let Some(socket_addr) = listen_url.strip_prefix("ws://") { - return socket_addr - .parse::() - .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::() @@ -234,129 +36,10 @@ fn parse_listen_url(listen_url: &str) -> Result { } 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> { - 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::(), - ) - .await - .context("code-mode host websocket listener failed") -} - -async fn readiness_handler() -> StatusCode { - StatusCode::OK -} - -async fn reject_requests_with_origin_header( - request: Request, - next: Next, -) -> Result { - 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, - State(state): State, -) -> 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, - ConnectInfo(peer_addr): ConnectInfo, - State(state): State, -) -> 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; diff --git a/codex-rs/code-mode-host/src/transport_tests.rs b/codex-rs/code-mode-host/src/transport_tests.rs index e1982e1864..d09728fda6 100644 --- a/codex-rs/code-mode-host/src/transport_tests.rs +++ b/codex-rs/code-mode-host/src/transport_tests.rs @@ -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::>(); - - 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::() - .expect("valid socket address") - ) - ); - assert_eq!( - parse_listen_url("ws://[::1]:9000").expect("IPv6 websocket listen URL should parse"), - ListenTransport::WebSocket( - "[::1]:9000" - .parse::() - .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 = diff --git a/codex-rs/code-mode-host/tests/websocket.rs b/codex-rs/code-mode-host/tests/websocket.rs deleted file mode 100644 index b4a1dfaaf0..0000000000 --- a/codex-rs/code-mode-host/tests/websocket.rs +++ /dev/null @@ -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::(); - -struct HostClient { - websocket: WebSocketStream>, -} - -impl HostHarness { - async fn connect(&self) -> Result { - 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) -> 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 { - 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 { - let capability = Capability::new(DUAL_WEBSOCKET_CAPABILITY)?; - let resource_limits = Capability::new(SESSION_RESOURCE_LIMITS_CAPABILITY)?; - let hello = ClientHello::new( - SupportedProtocolVersions::try_new([ProtocolVersion::V1])?, - CapabilitySet::empty(), - CapabilitySet::try_new([capability.clone(), resource_limits.clone()])?, - )?; - self.send(&ClientToHost::ClientHello(hello)).await?; - let HostToClient::HostHello(hello) = self.read().await? else { - anyhow::bail!("expected code-mode host hello"); - }; - assert!(hello.capabilities().contains(&capability)); - assert!(hello.capabilities().contains(&resource_limits)); - let token = hello - .bulk_connection_token() - .context("dual websocket handshake omitted its pairing token")?; - let bulk_url = format!("{}/bulk/{token}", websocket_url.trim_end_matches('/')); - let config = WebSocketConfig::default() - .max_frame_size(Some(MAX_WEBSOCKET_FRAME_BYTES)) - .max_message_size(Some(MAX_WEBSOCKET_FRAME_BYTES)); - let (websocket, _) = timeout( - TEST_TIMEOUT, - connect_async_with_config(bulk_url, Some(config), /*disable_nagle*/ false), - ) - .await - .context("timed out connecting to code-mode host bulk websocket")??; - Ok(HostClient { websocket }) - } - - async fn open_session(&mut self, session_id: SessionId) -> Result<()> { - let id = RequestId::new(/*value*/ 1); - self.send(&ClientToHost::Request { - id, - request: HostRequest::OpenSession { - session_id: session_id.clone(), - cell_execution_limits: None, - }, - }) - .await?; - assert_eq!( - self.read().await?, - HostToClient::Response { - id, - result: WireResult::Ok { - value: HostResponse::SessionReady { session_id }, - }, - } - ); - Ok(()) - } -} - -#[tokio::test] -async fn websocket_listener_serves_readiness_endpoint() -> Result<()> { - let host = HostHarness::start("ws://127.0.0.1:0").await?; - let address = host - .endpoint - .strip_prefix("ws://") - .context("code-mode host websocket URL should use ws://")?; - - let response = timeout(TEST_TIMEOUT, async { - let mut stream = TcpStream::connect(address) - .await - .context("failed to connect to code-mode host readiness endpoint")?; - let request = - format!("GET /readyz HTTP/1.1\r\nHost: {address}\r\nConnection: close\r\n\r\n"); - stream - .write_all(request.as_bytes()) - .await - .context("failed to request code-mode host readiness")?; - - let mut response = String::new(); - stream - .read_to_string(&mut response) - .await - .context("failed to read code-mode host readiness response")?; - Ok::<_, anyhow::Error>(response) - }) - .await - .context("timed out requesting code-mode host readiness")??; - - let status_line = response - .lines() - .next() - .context("code-mode host readiness response is missing a status line")?; - assert_eq!(status_line, "HTTP/1.1 200 OK"); - Ok(()) -} - -#[tokio::test] -async fn websocket_listener_executes_cells_and_forwards_tool_callbacks() -> Result<()> { - let host = HostHarness::start("ws://127.0.0.1:0").await?; - let mut client = host.connect().await?; - client.negotiate(CapabilitySet::empty()).await?; - - let session_id = SessionId::new("websocket-session")?; - client.open_session(session_id.clone()).await?; - - let execute_id = RequestId::new(/*value*/ 2); - client - .send(&ClientToHost::Request { - id: execute_id, - request: HostRequest::Execute { - session_id: session_id.clone(), - request: WireExecuteRequest { - tool_call_id: "websocket-call".to_string(), - enabled_tools: vec![WireToolDefinition { - name: "echo".to_string(), - tool_name: WireToolName { - name: "echo".to_string(), - namespace: None, - }, - description: String::new(), - kind: WireToolKind::Function, - input_schema: None, - output_schema: None, - }], - source: - r#"const result = await tools.echo({ value: "ping" }); text(result.value);"# - .to_string(), - yield_time_ms: Some(5_000), - max_output_tokens: Some(1_000), - }, - }, - }) - .await?; - - let started = client.read().await?; - let HostToClient::Response { - id, - result: - WireResult::Ok { - value: HostResponse::ExecutionStarted { cell_id }, - }, - } = started - else { - anyhow::bail!("expected execution-started response, got {started:?}"); - }; - assert_eq!(id, execute_id); - - let callback = client.read().await?; - let HostToClient::DelegateRequest { - id: delegate_id, - session_id: callback_session_id, - request: DelegateRequest::InvokeTool { invocation }, - } = callback - else { - anyhow::bail!("expected tool callback, got {callback:?}"); - }; - assert_eq!(callback_session_id, session_id); - assert_eq!(invocation.input, Some(json!({ "value": "ping" }))); - - client - .send(&ClientToHost::DelegateResponse { - id: delegate_id, - result: WireResult::Ok { - value: DelegateResponse::ToolResult { - result: json!({ "value": "pong" }), - }, - }, - }) - .await?; - - assert_eq!( - client.read().await?, - HostToClient::InitialResponse { - id: execute_id, - result: WireResult::Ok { - value: WireRuntimeResponse::Result { - cell_id, - content_items: vec![WireContentItem::InputText { - text: "pong".to_string(), - }], - error_text: None, - }, - }, - } - ); - Ok(()) -} - -#[tokio::test] -async fn production_websocket_client_runs_nested_tools_while_other_sessions_progress() -> Result<()> -{ - let host = HostHarness::start("ws://127.0.0.1:0").await?; - let provider = WebSocketCodeModeSessionProvider::new(host.endpoint.clone()); - let delegate = Arc::new(LargeToolResultDelegate { - started: Semaphore::new(/*permits*/ 0), - release: Semaphore::new(/*permits*/ 0), - }); - let slow_session = provider - .create_session(delegate.clone()) - .await - .map_err(anyhow::Error::msg)?; - let fast_session = provider - .create_session_with_limits( - Arc::new(NoopCodeModeSessionDelegate), - CodeModeSessionCellExecutionLimits { - max_yield_time_ms: Some(5_000), - max_heap_size_bytes: None, - }, - ) - .await - .map_err(anyhow::Error::msg)?; - - let slow_cell = slow_session - .execute(ExecuteRequest { - tool_call_id: "large-tool".to_string(), - enabled_tools: vec![ToolDefinition { - name: "large".to_string(), - tool_name: ToolName::plain("large"), - description: String::new(), - kind: CodeModeToolKind::Function, - input_schema: None, - output_schema: None, - }], - source: r#"const result = await tools.large({ value: "request" }); text(String(result.value.length));"# - .to_string(), - yield_time_ms: Some(20_000), - max_output_tokens: Some(1_000), - }) - .await - .map_err(anyhow::Error::msg)?; - let started = timeout(TEST_TIMEOUT, delegate.started.acquire()) - .await - .context("large tool callback did not start")??; - started.forget(); - - let fast_response = timeout(TEST_TIMEOUT, async { - fast_session - .execute(ExecuteRequest { - tool_call_id: "fast-before-transfer".to_string(), - enabled_tools: Vec::new(), - source: r#"text("fast-before");"#.to_string(), - yield_time_ms: Some(5_000), - max_output_tokens: Some(1_000), - }) - .await - .map_err(anyhow::Error::msg)? - .initial_response() - .await - .map_err(anyhow::Error::msg) - }) - .await - .context("unrelated execution was blocked by the pending tool callback")??; - assert_eq!( - fast_response, - RuntimeResponse::Result { - cell_id: CellId::new("1".to_string()), - content_items: vec![FunctionCallOutputContentItem::InputText { - text: "fast-before".to_string(), - }], - error_text: None, - } - ); - - delegate.release.add_permits(1); - let slow_response = async { - timeout(TEST_TIMEOUT, slow_cell.initial_response()) - .await - .context("large tool result did not finish")? - .map_err(anyhow::Error::msg) - }; - let concurrent_response = async { - timeout(TEST_TIMEOUT, async { - fast_session - .execute(ExecuteRequest { - tool_call_id: "fast-during-transfer".to_string(), - enabled_tools: Vec::new(), - source: r#"text("fast-during");"#.to_string(), - yield_time_ms: Some(5_000), - max_output_tokens: Some(1_000), - }) - .await - .map_err(anyhow::Error::msg)? - .initial_response() - .await - .map_err(anyhow::Error::msg) - }) - .await - .context("unrelated execution was blocked by the large tool transfer")? - }; - tokio::pin!(slow_response); - tokio::pin!(concurrent_response); - let concurrent_response = tokio::select! { - response = &mut concurrent_response => response?, - response = &mut slow_response => { - response?; - anyhow::bail!("large tool result completed before the unrelated control response"); - } - }; - let slow_response = slow_response.await?; - assert_eq!( - slow_response, - RuntimeResponse::Result { - cell_id: CellId::new("1".to_string()), - content_items: vec![FunctionCallOutputContentItem::InputText { - text: "8388608".to_string(), - }], - error_text: None, - } - ); - assert_eq!( - concurrent_response, - RuntimeResponse::Result { - cell_id: CellId::new("2".to_string()), - content_items: vec![FunctionCallOutputContentItem::InputText { - text: "fast-during".to_string(), - }], - error_text: None, - } - ); - - slow_session.shutdown().await.map_err(anyhow::Error::msg)?; - fast_session.shutdown().await.map_err(anyhow::Error::msg)?; - Ok(()) -} - -#[tokio::test] -async fn websocket_dual_connections_route_notifications_and_tool_callbacks_to_separate_lanes() --> Result<()> { - let host = HostHarness::start("ws://127.0.0.1:0").await?; - let mut control = host.connect().await?; - let mut bulk = control.negotiate_dual(&host.endpoint).await?; - let session_id = SessionId::new("dual-websocket-session")?; - control.open_session(session_id.clone()).await?; - - let execute_id = RequestId::new(/*value*/ 2); - control - .send(&ClientToHost::Request { - id: execute_id, - request: HostRequest::Execute { - session_id: session_id.clone(), - request: WireExecuteRequest { - tool_call_id: "dual-websocket-call".to_string(), - enabled_tools: vec![WireToolDefinition { - name: "echo".to_string(), - tool_name: WireToolName { - name: "echo".to_string(), - namespace: None, - }, - description: String::new(), - kind: WireToolKind::Function, - input_schema: None, - output_schema: None, - }], - source: r#"notify("important"); const result = await tools.echo({ value: "ping" }); text(result.value);"# - .to_string(), - yield_time_ms: Some(5_000), - max_output_tokens: Some(1_000), - }, - }, - }) - .await?; - - let started = control.read().await?; - let HostToClient::Response { - id, - result: - WireResult::Ok { - value: HostResponse::ExecutionStarted { cell_id }, - }, - } = started - else { - anyhow::bail!("expected execution-started response on control lane, got {started:?}"); - }; - assert_eq!(id, execute_id); - - let notification = control.read().await?; - let HostToClient::DelegateRequest { - id: notification_id, - .. - } = ¬ification - else { - anyhow::bail!("expected notification on control lane, got {notification:?}"); - }; - let notification_id = *notification_id; - assert_eq!( - notification, - HostToClient::DelegateRequest { - id: notification_id, - session_id: session_id.clone(), - request: DelegateRequest::Notify { - call_id: "dual-websocket-call".to_string(), - cell_id: cell_id.clone(), - text: "important".to_string(), - }, - } - ); - control - .send(&ClientToHost::DelegateResponse { - id: notification_id, - result: WireResult::Ok { - value: DelegateResponse::NotificationDelivered, - }, - }) - .await?; - - let callback = bulk.read().await?; - let HostToClient::DelegateRequest { - id: delegate_id, - session_id: callback_session_id, - request: DelegateRequest::InvokeTool { invocation }, - } = callback - else { - anyhow::bail!("expected tool callback on bulk lane, got {callback:?}"); - }; - assert_eq!(callback_session_id, session_id); - assert_eq!(invocation.input, Some(json!({ "value": "ping" }))); - - bulk.send(&ClientToHost::DelegateResponse { - id: delegate_id, - result: WireResult::Ok { - value: DelegateResponse::ToolResult { - result: json!({ "value": "pong" }), - }, - }, - }) - .await?; - - assert_eq!( - control.read().await?, - HostToClient::InitialResponse { - id: execute_id, - result: WireResult::Ok { - value: WireRuntimeResponse::Result { - cell_id: cell_id.clone(), - content_items: vec![WireContentItem::InputText { - text: "pong".to_string(), - }], - error_text: None, - }, - }, - } - ); - - assert_eq!( - control.read().await?, - HostToClient::CellClosed { - session_id: session_id.clone(), - cell_id, - } - ); - - let shutdown_id = RequestId::new(/*value*/ 3); - control - .send(&ClientToHost::Request { - id: shutdown_id, - request: HostRequest::ShutdownSession { - session_id: session_id.clone(), - }, - }) - .await?; - assert_eq!( - control.read().await?, - HostToClient::Response { - id: shutdown_id, - result: WireResult::Ok { - value: HostResponse::SessionClosed { session_id }, - }, - } - ); - Ok(()) -} - -#[tokio::test] -async fn websocket_control_operations_bypass_an_incomplete_bulk_frame() -> Result<()> { - let host = HostHarness::start("ws://127.0.0.1:0").await?; - let mut control = host.connect().await?; - let mut bulk = control.negotiate_dual(&host.endpoint).await?; - let session_id = SessionId::new("bulk-priority-session")?; - control.open_session(session_id.clone()).await?; - - let bulk_message = ClientToHost::DelegateResponse { - id: codex_code_mode_protocol::host::DelegateRequestId::new(/*value*/ 999), - result: WireResult::Ok { - value: DelegateResponse::ToolResult { - result: json!({ "image": "x".repeat(1024 * 1024) }), - }, - }, - }; - let payload = EncodedFrame::encode(&bulk_message)?.into_framed_bytes(); - let mask = [0x13_u8, 0x37, 0xc0, 0xde]; - let mut partial_frame = vec![0x82_u8, 0xff]; - partial_frame.extend_from_slice(&(payload.len() as u64).to_be_bytes()); - partial_frame.extend_from_slice(&mask); - partial_frame.extend( - payload[..64 * 1024] - .iter() - .enumerate() - .map(|(index, byte)| byte ^ mask[index % mask.len()]), - ); - bulk.websocket - .get_mut() - .write_all(&partial_frame) - .await - .context("failed to start the intentionally incomplete bulk frame")?; - - let request_id = RequestId::new(/*value*/ 42); - control - .send(&ClientToHost::Request { - id: request_id, - request: HostRequest::Wait { - session_id: session_id.clone(), - request: WireWaitRequest { - cell_id: codex_code_mode_protocol::host::WireCellId::new("missing-cell"), - yield_time_ms: 10, - }, - }, - }) - .await?; - let response = timeout(TEST_TIMEOUT, control.read()) - .await - .context("control wait was blocked behind the incomplete bulk frame")??; - assert!(matches!( - response, - HostToClient::Response { - id, - result: WireResult::Ok { - value: HostResponse::WaitCompleted { .. }, - }, - } if id == request_id - )); - - let execute_id = RequestId::new(/*value*/ 43); - control - .send(&ClientToHost::Request { - id: execute_id, - request: HostRequest::Execute { - session_id, - request: WireExecuteRequest { - tool_call_id: "control-execute".to_string(), - enabled_tools: Vec::new(), - source: r#"text("fast");"#.to_string(), - yield_time_ms: Some(5_000), - max_output_tokens: Some(1_000), - }, - }, - }) - .await?; - for _ in 0..2 { - let response = timeout(TEST_TIMEOUT, control.read()) - .await - .context("control execute was blocked behind the incomplete bulk frame")??; - assert!(matches!( - response, - HostToClient::Response { id, .. } | HostToClient::InitialResponse { id, .. } - if id == execute_id - )); - } - Ok(()) -} - -#[tokio::test] -async fn websocket_bulk_lane_rejects_control_messages() -> Result<()> { - let host = HostHarness::start("ws://127.0.0.1:0").await?; - let mut control = host.connect().await?; - let mut bulk = control.negotiate_dual(&host.endpoint).await?; - let session_id = SessionId::new("wrong-lane-session")?; - control.open_session(session_id.clone()).await?; - - bulk.send(&ClientToHost::Request { - id: RequestId::new(/*value*/ 42), - request: HostRequest::Wait { - session_id, - request: WireWaitRequest { - cell_id: codex_code_mode_protocol::host::WireCellId::new("missing-cell"), - yield_time_ms: 10, - }, - }, - }) - .await?; - - let result = timeout(TEST_TIMEOUT, control.websocket.next()) - .await - .context("wrong-lane control message did not disconnect the paired sockets")?; - assert!( - matches!(result, None | Some(Ok(Message::Close(_))) | Some(Err(_))), - "wrong-lane message unexpectedly returned {result:?}" - ); - Ok(()) -} - -#[tokio::test] -async fn websocket_bulk_pairing_rejects_unknown_tokens() -> Result<()> { - let host = HostHarness::start("ws://127.0.0.1:0").await?; - let unknown_token = Uuid::new_v4(); - let url = format!("{}/bulk/{unknown_token}", host.endpoint); - let error = match connect_async(url).await { - Ok(_) => anyhow::bail!("unknown bulk pairing token should be rejected"), - Err(error) => error, - }; - let WebSocketError::Http(response) = error else { - anyhow::bail!("bulk pairing failed unexpectedly: {error}"); - }; - assert_eq!(response.status(), StatusCode::NOT_FOUND); - Ok(()) -} - -#[tokio::test] -async fn websocket_listener_accepts_frames_larger_than_default_websocket_limit() -> Result<()> { - let host = HostHarness::start("ws://127.0.0.1:0").await?; - let mut client = host.connect().await?; - let capability = Capability::new("x".repeat((16 * 1024 * 1024) + 1))?; - - client - .negotiate(CapabilitySet::try_new([capability])?) - .await -} - -#[tokio::test] -async fn websocket_listener_keeps_connections_and_session_ids_isolated() -> Result<()> { - let host = HostHarness::start("ws://127.0.0.1:0").await?; - let mut first = host.connect().await?; - let mut second = host.connect().await?; - first.negotiate(CapabilitySet::empty()).await?; - second.negotiate(CapabilitySet::empty()).await?; - - let session_id = SessionId::new("shared-session-name")?; - first.open_session(session_id.clone()).await?; - second.open_session(session_id).await?; - Ok(()) -} - -#[tokio::test] -async fn malformed_websocket_frame_does_not_stop_the_listener() -> Result<()> { - let mut host = HostHarness::start("ws://127.0.0.1:0").await?; - let stderr = host - ._child - .stderr - .take() - .context("code-mode host stderr was not captured")?; - let mut stderr_lines = BufReader::new(stderr).lines(); - let mut malformed = host.connect().await?; - malformed.send_binary(vec![1, 0, 0, 0, b'{']).await?; - - let close = timeout(TEST_TIMEOUT, malformed.websocket.next()) - .await - .context("timed out waiting for malformed websocket connection to close")?; - if let Some(Ok(message)) = close - && !matches!(message, Message::Close(_)) - { - anyhow::bail!("malformed websocket returned an unexpected message: {message:?}"); - } - - let diagnostic = timeout(TEST_TIMEOUT, async { - loop { - let line = stderr_lines - .next_line() - .await? - .context("code-mode host exited before reporting the malformed frame")?; - if line.contains("code-mode host session failed") { - return Ok::<_, anyhow::Error>(line); - } - } - }) - .await - .context("timed out waiting for the malformed websocket diagnostic")??; - assert!( - diagnostic.contains("failed to read code-mode client hello"), - "unexpected malformed websocket diagnostic: {diagnostic}", - ); - - let mut recovered = host.connect().await?; - recovered.negotiate(CapabilitySet::empty()).await -} - -#[tokio::test] -async fn websocket_listener_rejects_browser_origin_handshakes() -> Result<()> { - let host = HostHarness::start("ws://127.0.0.1:0").await?; - let mut request = host.endpoint.as_str().into_client_request()?; - request - .headers_mut() - .insert(ORIGIN, HeaderValue::from_static("https://evil.example")); - - let error = match connect_async(request).await { - Ok(_) => anyhow::bail!("browser-origin websocket handshake should be rejected"), - Err(error) => error, - }; - let WebSocketError::Http(response) = error else { - anyhow::bail!("browser-origin websocket handshake failed unexpectedly: {error}"); - }; - assert_eq!(response.status(), StatusCode::FORBIDDEN); - Ok(()) -} diff --git a/codex-rs/code-mode-protocol/src/host/host_tests.rs b/codex-rs/code-mode-protocol/src/host/host_tests.rs index 7404cb6d4d..e215d2a30b 100644 --- a/codex-rs/code-mode-protocol/src/host/host_tests.rs +++ b/codex-rs/code-mode-protocol/src/host/host_tests.rs @@ -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(), diff --git a/codex-rs/code-mode-protocol/src/host/message.rs b/codex-rs/code-mode-protocol/src/host/message.rs index 4b922f5ce2..bc442426a3 100644 --- a/codex-rs/code-mode-protocol/src/host/message.rs +++ b/codex-rs/code-mode-protocol/src/host/message.rs @@ -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, } 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 { diff --git a/codex-rs/code-mode-protocol/src/host/mod.rs b/codex-rs/code-mode-protocol/src/host/mod.rs index 4059230213..43bb711162 100644 --- a/codex-rs/code-mode-protocol/src/host/mod.rs +++ b/codex-rs/code-mode-protocol/src/host/mod.rs @@ -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; diff --git a/codex-rs/code-mode/Cargo.toml b/codex-rs/code-mode/Cargo.toml index 0ef5079beb..8bc8b84933 100644 --- a/codex-rs/code-mode/Cargo.toml +++ b/codex-rs/code-mode/Cargo.toml @@ -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"] } diff --git a/codex-rs/code-mode/src/lib.rs b/codex-rs/code-mode/src/lib.rs index efbe08e1fa..8bed201d33 100644 --- a/codex-rs/code-mode/src/lib.rs +++ b/codex-rs/code-mode/src/lib.rs @@ -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; diff --git a/codex-rs/code-mode/src/remote_session.rs b/codex-rs/code-mode/src/remote_session.rs index 5433b88f38..fb38e2899e 100644 --- a/codex-rs/code-mode/src/remote_session.rs +++ b/codex-rs/code-mode/src/remote_session.rs @@ -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, -} - 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, - ) -> CodeModeSessionProviderFuture<'a> { - self.create_session_with_limits(delegate, CodeModeSessionCellExecutionLimits::default()) - } - - fn create_session_with_limits<'a>( - &'a self, - delegate: Arc, - limits: CodeModeSessionCellExecutionLimits, - ) -> CodeModeSessionProviderFuture<'a> { - Box::pin(create_host_session( - delegate, - Arc::clone(&self.host), - limits, - )) - } -} - async fn create_host_session( delegate: Arc, host: Arc, @@ -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>>, 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>, } -/// 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, } diff --git a/codex-rs/code-mode/src/remote_session/connection.rs b/codex-rs/code-mode/src/remote_session/connection.rs index 4fcd188f3a..248b3b0171 100644 --- a/codex-rs/code-mode/src/remote_session/connection.rs +++ b/codex-rs/code-mode/src/remote_session/connection.rs @@ -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::(); // 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, cancellation: CancellationToken, alive: Arc, @@ -143,50 +127,6 @@ struct ConnectionSupervisor { writer_task: JoinHandle>, } -enum ConnectionOwner { - Process(Box), - 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::().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 { - 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, + mut reader: FramedReader, + mut writer: FramedWriter, + mut child: Child, ) -> Result { 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::(IPC_CHANNEL_CAPACITY); - let (bulk_tx, bulk_rx) = if bulk_writer.is_some() { - let (sender, receiver) = mpsc::channel::(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, mut outgoing: mpsc::Receiver, 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; } } } diff --git a/codex-rs/code-mode/src/remote_session/connection/driver.rs b/codex-rs/code-mode/src/remote_session/connection/driver.rs index 667a435805..a4c8ab8749 100644 --- a/codex-rs/code-mode/src/remote_session/connection/driver.rs +++ b/codex-rs/code-mode/src/remote_session/connection/driver.rs @@ -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, execute_claim_rx: mpsc::UnboundedReceiver, outgoing_tx: mpsc::Sender, - bulk_tx: Option>, requests: RequestTracker, deferred_host_messages: VecDeque, 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) -> 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()); diff --git a/codex-rs/code-mode/src/remote_session/connection/driver/commands.rs b/codex-rs/code-mode/src/remote_session/connection/driver/commands.rs index 446c9c6abc..ce50e06881 100644 --- a/codex-rs/code-mode/src/remote_session/connection/driver/commands.rs +++ b/codex-rs/code-mode/src/remote_session/connection/driver/commands.rs @@ -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) } } diff --git a/codex-rs/code-mode/src/remote_session/connection/driver/delegate_runtime.rs b/codex-rs/code-mode/src/remote_session/connection/driver/delegate_runtime.rs index fb01f8783e..a7353a3f75 100644 --- a/codex-rs/code-mode/src/remote_session/connection/driver/delegate_runtime.rs +++ b/codex-rs/code-mode/src/remote_session/connection/driver/delegate_runtime.rs @@ -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 { diff --git a/codex-rs/code-mode/src/remote_session/connection/driver/responses.rs b/codex-rs/code-mode/src/remote_session/connection/driver/responses.rs index 17f697362e..7f0067e5ed 100644 --- a/codex-rs/code-mode/src/remote_session/connection/driver/responses.rs +++ b/codex-rs/code-mode/src/remote_session/connection/driver/responses.rs @@ -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 { diff --git a/codex-rs/code-mode/src/remote_session/connection/driver_tests.rs b/codex-rs/code-mode/src/remote_session/connection/driver_tests.rs index 3f51607b44..85571c1eec 100644 --- a/codex-rs/code-mode/src/remote_session/connection/driver_tests.rs +++ b/codex-rs/code-mode/src/remote_session/connection/driver_tests.rs @@ -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::(&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(); diff --git a/codex-rs/code-mode/src/remote_session/connection/reader.rs b/codex-rs/code-mode/src/remote_session/connection/reader.rs index 8bb218e67c..cfef435d78 100644 --- a/codex-rs/code-mode/src/remote_session/connection/reader.rs +++ b/codex-rs/code-mode/src/remote_session/connection/reader.rs @@ -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, events: mpsc::Sender, cancellation: CancellationToken, - lane: Option, ) -> 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 diff --git a/codex-rs/code-mode/src/remote_session/connection/transport.rs b/codex-rs/code-mode/src/remote_session/connection/transport.rs deleted file mode 100644 index 9ce520b5f2..0000000000 --- a/codex-rs/code-mode/src/remote_session/connection/transport.rs +++ /dev/null @@ -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), - WebSocket(SplitStream), -} - -pub(super) enum ConnectionWriter { - Stdio(FramedWriter), - WebSocket(SplitSink), -} - -impl ConnectionReader { - pub(super) async fn read(&mut self) -> io::Result> { - 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}" - )) - }) - } - } - } -} diff --git a/codex-rs/code-mode/src/remote_session_tests.rs b/codex-rs/code-mode/src/remote_session_tests.rs index d5fd559648..a2e0a08ff8 100644 --- a/codex-rs/code-mode/src/remote_session_tests.rs +++ b/codex-rs/code-mode/src/remote_session_tests.rs @@ -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::(&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(