From 94ebae725e5e8f22b5d86773d9223047f57b6118 Mon Sep 17 00:00:00 2001 From: Celia Chen Date: Fri, 24 Jul 2026 00:11:39 +0000 Subject: [PATCH] Route exec-server WebSockets through configured proxies (#35056) ## Why Remote environment connections need to honor Codex's effective outbound proxy policy, including when a rendezvous connection reconnects. ## What changed - Pass the configured `HttpClientFactory` into remote environment transports and use `WebSocketConnector` for exec-server and rendezvous WebSockets. - Resolve proxy routes asynchronously so these connections can use the configured system proxy. - Add connector options that preserve Tungstenite's default TLS behavior and enable `TCP_NODELAY` for latency-sensitive rendezvous traffic. ## Testing - Verify prepared remote environments connect through a configured system proxy. - Verify initial and reconnected encrypted relay peers use the system proxy. - Cover default TLS selection and opt-in `TCP_NODELAY` behavior in the WebSocket client. GitOrigin-RevId: 8a8da2116e37cb3a891269d0c0b037986fecdd3c --- codex-rs/Cargo.lock | 1 + codex-rs/exec-server/Cargo.toml | 1 + codex-rs/exec-server/src/client.rs | 107 ++++++++---- codex-rs/exec-server/src/client_api.rs | 3 + codex-rs/exec-server/src/client_transport.rs | 66 +++++-- .../exec-server/src/client_transport_tests.rs | 10 +- codex-rs/exec-server/src/connection.rs | 6 +- codex-rs/exec-server/src/environment.rs | 141 +++++++-------- .../exec-server/src/environment_bootstrap.rs | 3 +- .../src/environment_bootstrap_tests.rs | 50 +++--- .../exec-server/src/environment_provider.rs | 40 +++-- codex-rs/exec-server/src/environment_toml.rs | 24 +-- codex-rs/exec-server/src/relay.rs | 10 +- codex-rs/exec-server/src/remote.rs | 37 ++-- .../src/remote_file_system_path_uri_tests.rs | 3 + .../exec-server/src/server/handler/tests.rs | 16 +- codex-rs/exec-server/tests/environment.rs | 82 +++++++++ codex-rs/exec-server/tests/file_stream.rs | 5 + codex-rs/exec-server/tests/http_client.rs | 3 + codex-rs/exec-server/tests/relay.rs | 88 +++++++++- .../tests/streamable_http_test_support.rs | 3 + codex-rs/websocket-client/src/dialer.rs | 35 ++-- codex-rs/websocket-client/src/dialer_tests.rs | 161 +++++++++++++++++- codex-rs/websocket-client/src/lib.rs | 58 ++++++- 24 files changed, 727 insertions(+), 226 deletions(-) diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index 34626523b5..7e391d84d9 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -2929,6 +2929,7 @@ dependencies = [ "codex-utils-path-uri", "codex-utils-pty", "codex-utils-rustls-provider", + "codex-websocket-client", "ctor 0.6.3", "futures", "http 1.4.0", diff --git a/codex-rs/exec-server/Cargo.toml b/codex-rs/exec-server/Cargo.toml index 076689bc77..834584e74a 100644 --- a/codex-rs/exec-server/Cargo.toml +++ b/codex-rs/exec-server/Cargo.toml @@ -28,6 +28,7 @@ codex-utils-absolute-path = { workspace = true } codex-utils-path-uri = { workspace = true } codex-utils-pty = { workspace = true } codex-utils-rustls-provider = { workspace = true } +codex-websocket-client = { workspace = true } futures = { workspace = true } http = { workspace = true } reqwest = { workspace = true, features = ["json", "rustls-tls", "stream"] } diff --git a/codex-rs/exec-server/src/client.rs b/codex-rs/exec-server/src/client.rs index 7be935c01f..1ff98fcc7c 100644 --- a/codex-rs/exec-server/src/client.rs +++ b/codex-rs/exec-server/src/client.rs @@ -112,6 +112,7 @@ use crate::protocol::WriteParams; use crate::protocol::WriteResponse; use crate::rpc::RpcCallError; use crate::rpc::RpcClient; +use codex_http_client::HttpClientFactory; pub(crate) mod http_client; #[path = "client_recovery.rs"] @@ -157,13 +158,18 @@ impl From for ExecServerClientConnectOptions { } impl RemoteExecServerConnectArgs { - pub fn new(websocket_url: String, client_name: String) -> Self { + pub fn new( + websocket_url: String, + client_name: String, + http_client_factory: HttpClientFactory, + ) -> Self { Self { websocket_url, client_name, connect_timeout: CONNECT_TIMEOUT, initialize_timeout: INITIALIZE_TIMEOUT, resume_session_id: None, + http_client_factory, } } } @@ -288,7 +294,8 @@ type ConnectionAttempt = OnceCell; #[derive(Clone)] pub(crate) struct LazyRemoteExecServerClient { - transport_params: ExecServerTransportParams, + pub(crate) transport_params: ExecServerTransportParams, + http_client_factory: HttpClientFactory, recovery_policy: RecoveryPolicy, // Saves the first startup result so callers share it and failures remain final. startup: Arc, @@ -299,9 +306,13 @@ pub(crate) struct LazyRemoteExecServerClient { } impl LazyRemoteExecServerClient { - pub(crate) fn new(transport_params: ExecServerTransportParams) -> Self { + pub(crate) fn new( + transport_params: ExecServerTransportParams, + http_client_factory: HttpClientFactory, + ) -> Self { Self { transport_params, + http_client_factory, recovery_policy: RecoveryPolicy::Wait, startup: Arc::new(ConnectionAttempt::new()), current_client: Arc::new(StdMutex::new(None)), @@ -499,9 +510,12 @@ impl LazyRemoteExecServerClient { } async fn connect_once(&self) -> ConnectionResult { - let result = ExecServerClient::connect_for_transport(self.transport_params.clone()) - .await - .map_err(Arc::new); + let result = ExecServerClient::connect_for_transport( + self.transport_params.clone(), + self.http_client_factory.clone(), + ) + .await + .map_err(Arc::new); if let Ok(client) = &result { client .attach_environment_connection_state(self.environment_connection_state_tx.clone()); @@ -547,6 +561,8 @@ pub enum ExecServerError { #[source] source: tokio_tungstenite::tungstenite::Error, }, + #[error("failed to configure exec-server websocket: {0}")] + WebSocketConfiguration(String), #[error("timed out waiting for exec-server initialize handshake after {timeout:?}")] InitializeTimedOut { timeout: Duration }, #[error("exec-server transport closed")] @@ -1493,6 +1509,8 @@ mod tests { use codex_exec_server_protocol::JSONRPCMessage; use codex_exec_server_protocol::JSONRPCNotification; use codex_exec_server_protocol::JSONRPCResponse; + use codex_http_client::HttpClientFactory; + use codex_http_client::OutboundProxyPolicy; use codex_utils_path_uri::PathUri; use futures::SinkExt; use futures::StreamExt; @@ -1809,6 +1827,9 @@ mod tests { }, initialize_timeout: DEFAULT_REMOTE_EXEC_SERVER_INITIALIZE_TIMEOUT, }, + codex_http_client::HttpClientFactory::new( + codex_http_client::OutboundProxyPolicy::ReqwestDefault, + ), ) .await .expect("stdio transport should connect"); @@ -2221,11 +2242,14 @@ mod tests { finish_rx.await.expect("test should finish"); }); - let client = LazyRemoteExecServerClient::new(ExecServerTransportParams::WebSocketUrl { - websocket_url, - connect_timeout: Duration::from_secs(1), - initialize_timeout: Duration::from_secs(1), - }); + let client = LazyRemoteExecServerClient::new( + ExecServerTransportParams::WebSocketUrl { + websocket_url, + connect_timeout: Duration::from_secs(1), + initialize_timeout: Duration::from_secs(1), + }, + HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault), + ); let stable_client = client.get().await.expect("client should connect"); timeout(Duration::from_secs(1), resumed_rx) .await @@ -2328,11 +2352,14 @@ mod tests { finish_rx.await.expect("test should finish"); }); - let client = LazyRemoteExecServerClient::new(ExecServerTransportParams::WebSocketUrl { - websocket_url, - connect_timeout: Duration::from_secs(1), - initialize_timeout: Duration::from_secs(1), - }); + let client = LazyRemoteExecServerClient::new( + ExecServerTransportParams::WebSocketUrl { + websocket_url, + connect_timeout: Duration::from_secs(1), + initialize_timeout: Duration::from_secs(1), + }, + HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault), + ); let stable_client = client.get().await.expect("client should connect"); let session = stable_client .register_session(&ProcessId::from("proc-write")) @@ -2427,6 +2454,9 @@ mod tests { connect_timeout: Duration::from_secs(1), initialize_timeout: Duration::from_secs(1), resume_session_id: Some("session-1".to_string()), + http_client_factory: codex_http_client::HttpClientFactory::new( + codex_http_client::OutboundProxyPolicy::ReqwestDefault, + ), }), ) .await @@ -2463,11 +2493,14 @@ mod tests { .await .expect("client should close after the test"); }); - let client = LazyRemoteExecServerClient::new(ExecServerTransportParams::WebSocketUrl { - websocket_url, - connect_timeout: Duration::from_secs(1), - initialize_timeout: Duration::from_secs(1), - }); + let client = LazyRemoteExecServerClient::new( + ExecServerTransportParams::WebSocketUrl { + websocket_url, + connect_timeout: Duration::from_secs(1), + initialize_timeout: Duration::from_secs(1), + }, + HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault), + ); assert!(!client.startup_finished()); let _startup_task = client.start_connecting(); @@ -2489,15 +2522,18 @@ mod tests { #[tokio::test] async fn terminal_stdio_startup_failure_is_remembered() { - let client = LazyRemoteExecServerClient::new(ExecServerTransportParams::StdioCommand { - command: StdioExecServerCommand { - program: "codex-missing-exec-server-for-test".to_string(), - args: Vec::new(), - env: HashMap::new(), - cwd: None, + let client = LazyRemoteExecServerClient::new( + ExecServerTransportParams::StdioCommand { + command: StdioExecServerCommand { + program: "codex-missing-exec-server-for-test".to_string(), + args: Vec::new(), + env: HashMap::new(), + cwd: None, + }, + initialize_timeout: Duration::from_secs(1), }, - initialize_timeout: Duration::from_secs(1), - }); + HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault), + ); assert!(client.start_connecting().is_none()); assert!(!client.startup_finished()); @@ -2576,11 +2612,14 @@ mod tests { .await .expect("client should close after the test"); }); - let client = LazyRemoteExecServerClient::new(ExecServerTransportParams::WebSocketUrl { - websocket_url, - connect_timeout: Duration::from_secs(1), - initialize_timeout: Duration::from_secs(1), - }); + let client = LazyRemoteExecServerClient::new( + ExecServerTransportParams::WebSocketUrl { + websocket_url, + connect_timeout: Duration::from_secs(1), + initialize_timeout: Duration::from_secs(1), + }, + HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault), + ); let initial = client.get().await.expect("startup should connect"); timeout(Duration::from_secs(1), async { diff --git a/codex-rs/exec-server/src/client_api.rs b/codex-rs/exec-server/src/client_api.rs index cd284a35e8..a8ce40c8e8 100644 --- a/codex-rs/exec-server/src/client_api.rs +++ b/codex-rs/exec-server/src/client_api.rs @@ -3,6 +3,7 @@ use std::path::PathBuf; use std::sync::Arc; use std::time::Duration; +use codex_http_client::HttpClientFactory; use futures::future::BoxFuture; use futures::future::Shared; use tokio::sync::oneshot; @@ -33,6 +34,7 @@ pub struct RemoteExecServerConnectArgs { pub connect_timeout: Duration, pub initialize_timeout: Duration, pub resume_session_id: Option, + pub http_client_factory: HttpClientFactory, } /// Registry-authorized material for one Noise rendezvous connection attempt. @@ -61,6 +63,7 @@ pub struct NoiseRendezvousConnectArgs { pub connect_timeout: Duration, pub initialize_timeout: Duration, pub resume_session_id: Option, + pub http_client_factory: HttpClientFactory, } /// Supplies fresh registry-authorized material for Noise rendezvous connections. diff --git a/codex-rs/exec-server/src/client_transport.rs b/codex-rs/exec-server/src/client_transport.rs index 14cc4cc352..87c238e55b 100644 --- a/codex-rs/exec-server/src/client_transport.rs +++ b/codex-rs/exec-server/src/client_transport.rs @@ -6,13 +6,14 @@ use tokio::io::AsyncBufReadExt; use tokio::io::BufReader; use tokio::process::Command; use tokio::time::timeout; -use tokio_tungstenite::connect_async; -use tokio_tungstenite::connect_async_with_config; use tokio_tungstenite::tungstenite::client::IntoClientRequest; use tracing::debug; use tracing::warn; +use codex_http_client::HttpClientFactory; use codex_utils_rustls_provider::ensure_rustls_crypto_provider; +use codex_websocket_client::WebSocketConnector; +use codex_websocket_client::WebSocketTlsMode; use crate::ExecServerClient; use crate::ExecServerError; @@ -50,6 +51,7 @@ pub(crate) enum ExecServerReconnectStrategy { client_name: String, connect_timeout: Duration, initialize_timeout: Duration, + http_client_factory: HttpClientFactory, }, } @@ -71,6 +73,7 @@ impl ExecServerReconnectStrategy { client_name, connect_timeout, initialize_timeout, + http_client_factory, } => { let bundle = provider.connect_bundle(identity.public_key()).await?; ExecServerClient::open_noise_rendezvous_connection(NoiseRendezvousConnectArgs { @@ -80,6 +83,7 @@ impl ExecServerReconnectStrategy { connect_timeout: *connect_timeout, initialize_timeout: *initialize_timeout, resume_session_id: Some(session_id.to_string()), + http_client_factory: http_client_factory.clone(), }) .await } @@ -93,6 +97,7 @@ impl ExecServerClient { /// and authorization without replacing the harness identity. pub(crate) async fn connect_for_transport( transport_params: ExecServerTransportParams, + http_client_factory: HttpClientFactory, ) -> Result { let (transport_params, deferred_readiness) = match transport_params { ExecServerTransportParams::Deferred(deferred) => { @@ -130,9 +135,14 @@ impl ExecServerClient { client_name: ENVIRONMENT_CLIENT_NAME.to_string(), connect_timeout: DEFAULT_REMOTE_EXEC_SERVER_CONNECT_TIMEOUT, initialize_timeout: DEFAULT_REMOTE_EXEC_SERVER_INITIALIZE_TIMEOUT, + http_client_factory: http_client_factory.clone(), }; - let (connection, options) = - Self::open_initial_noise_rendezvous_connection(&provider, &identity).await?; + let (connection, options) = Self::open_initial_noise_rendezvous_connection( + &provider, + &identity, + http_client_factory, + ) + .await?; return Self::connect_with_recovery(connection, options, Some(reconnect_strategy)) .await; } @@ -155,6 +165,7 @@ impl ExecServerClient { connect_timeout, initialize_timeout, resume_session_id: None, + http_client_factory, }) .await } @@ -162,6 +173,7 @@ impl ExecServerClient { async fn open_initial_noise_rendezvous_connection( provider: &Arc, identity: &NoiseChannelIdentity, + http_client_factory: HttpClientFactory, ) -> Result<(JsonRpcConnection, ExecServerClientConnectOptions), ExecServerError> { let open_connection = |bundle: NoiseRendezvousConnectBundle| { Self::open_noise_rendezvous_connection(NoiseRendezvousConnectArgs { @@ -171,6 +183,7 @@ impl ExecServerClient { connect_timeout: DEFAULT_REMOTE_EXEC_SERVER_CONNECT_TIMEOUT, initialize_timeout: DEFAULT_REMOTE_EXEC_SERVER_INITIALIZE_TIMEOUT, resume_session_id: None, + http_client_factory: http_client_factory.clone(), }) }; let bundle = provider.connect_bundle(identity.public_key()).await?; @@ -212,16 +225,34 @@ impl ExecServerClient { ensure_rustls_crypto_provider(); let websocket_url = args.websocket_url.clone(); let connect_timeout = args.connect_timeout; - let (stream, _) = timeout(connect_timeout, connect_async(websocket_url.as_str())) - .await - .map_err(|_| ExecServerError::WebSocketConnectTimeout { - url: websocket_url.clone(), - timeout: connect_timeout, - })? + let request = websocket_url + .as_str() + .into_client_request() .map_err(|source| ExecServerError::WebSocketConnect { url: websocket_url.clone(), source, })?; + let connector = WebSocketConnector::new_with_tls_mode( + &args.http_client_factory, + WebSocketTlsMode::TungsteniteDefault, + ) + .map_err(|error| ExecServerError::WebSocketConfiguration(error.to_string()))?; + let (stream, _) = timeout( + connect_timeout, + connector.connect( + request, + tokio_tungstenite::tungstenite::protocol::WebSocketConfig::default(), + ), + ) + .await + .map_err(|_| ExecServerError::WebSocketConnectTimeout { + url: websocket_url.clone(), + timeout: connect_timeout, + })? + .map_err(|source| ExecServerError::WebSocketConnect { + url: websocket_url.clone(), + source, + })?; let connection_label = format!("exec-server websocket {websocket_url}"); let connection = if is_rendezvous_harness_url(&websocket_url) { @@ -267,6 +298,7 @@ impl ExecServerClient { connect_timeout, initialize_timeout, resume_session_id, + http_client_factory, } = args; let NoiseRendezvousConnectBundle { websocket_url, @@ -292,13 +324,13 @@ impl ExecServerClient { .extend(current_trace_context_headers()); let (stream, _) = timeout( connect_timeout, - connect_async_with_config( - request, - Some(noise_relay_websocket_config()), - // Rendezvous sends small, latency-sensitive frames, so avoid Nagle's coalescing delay. - /*disable_nagle*/ - true, - ), + WebSocketConnector::new_with_tls_mode( + &http_client_factory, + WebSocketTlsMode::TungsteniteDefault, + ) + .map_err(|error| ExecServerError::WebSocketConfiguration(error.to_string()))? + .with_tcp_nodelay() + .connect(request, noise_relay_websocket_config()), ) .await .map_err(|_| ExecServerError::WebSocketConnectTimeout { diff --git a/codex-rs/exec-server/src/client_transport_tests.rs b/codex-rs/exec-server/src/client_transport_tests.rs index 7dab5fa218..cce85b3331 100644 --- a/codex-rs/exec-server/src/client_transport_tests.rs +++ b/codex-rs/exec-server/src/client_transport_tests.rs @@ -99,8 +99,14 @@ async fn initial_noise_connection_refreshes_bundle_after_unauthorized_handshake( let provider: Arc = sequence.clone(); let identity = NoiseChannelIdentity::generate()?; - let _connection = - ExecServerClient::open_initial_noise_rendezvous_connection(&provider, &identity).await?; + let _connection = ExecServerClient::open_initial_noise_rendezvous_connection( + &provider, + &identity, + codex_http_client::HttpClientFactory::new( + codex_http_client::OutboundProxyPolicy::ReqwestDefault, + ), + ) + .await?; assert_eq!( sequence.returned_urls(), diff --git a/codex-rs/exec-server/src/connection.rs b/codex-rs/exec-server/src/connection.rs index b0805f74de..4c2af0813d 100644 --- a/codex-rs/exec-server/src/connection.rs +++ b/codex-rs/exec-server/src/connection.rs @@ -18,6 +18,7 @@ use tokio::process::Child; use tokio::sync::mpsc; use tokio::sync::watch; use tokio::time::timeout; +#[cfg(test)] use tokio_tungstenite::WebSocketStream; use tokio_tungstenite::tungstenite::Message; use tracing::debug; @@ -385,9 +386,10 @@ impl JsonRpcConnection { } } - pub(crate) fn from_websocket(stream: WebSocketStream, connection_label: String) -> Self + pub(crate) fn from_websocket(stream: T, connection_label: String) -> Self where - S: AsyncRead + AsyncWrite + Unpin + Send + 'static, + T: Sink + Stream> + Unpin + Send + 'static, + E: std::fmt::Display + Send + 'static, { Self::from_websocket_stream(stream, connection_label, /*ping_interval*/ None) } diff --git a/codex-rs/exec-server/src/environment.rs b/codex-rs/exec-server/src/environment.rs index 7c1cacf9bb..ff39714049 100644 --- a/codex-rs/exec-server/src/environment.rs +++ b/codex-rs/exec-server/src/environment.rs @@ -287,7 +287,7 @@ impl EnvironmentManager { } else { None }; - for (id, environment) in environments { + for (id, transport) in environments { if id.is_empty() { return Err(ExecServerError::Protocol( "environment id cannot be empty".to_string(), @@ -298,6 +298,11 @@ impl EnvironmentManager { "environment id `{LOCAL_ENVIRONMENT_ID}` is reserved for EnvironmentManager" ))); } + let environment = Environment::remote_with_transport( + transport, + /*local_runtime_paths*/ None, + http_client_factory.clone(), + ); if environment_map .insert(id.clone(), Arc::new(environment)) .is_some() @@ -408,6 +413,7 @@ impl EnvironmentManager { connect_timeout.unwrap_or(DEFAULT_REMOTE_EXEC_SERVER_CONNECT_TIMEOUT), ), self.local_runtime_paths.clone(), + self.http_client_factory.clone(), )); self.insert_environment(environment_id, environment); Ok(()) @@ -429,6 +435,7 @@ impl EnvironmentManager { transport: ExecServerTransportParams::NoiseRendezvous { provider, identity }, })), self.local_runtime_paths.clone(), + self.http_client_factory.clone(), ); environment.ready_info = Some(Arc::clone(&ready_info)); let environment = Arc::new(environment); @@ -455,6 +462,7 @@ impl EnvironmentManager { let environment = Arc::new(Environment::remote_with_transport( ExecServerTransportParams::NoiseRendezvous { provider, identity }, self.local_runtime_paths.clone(), + self.http_client_factory.clone(), )); self.insert_environment(environment_id, environment); Ok(()) @@ -639,17 +647,26 @@ impl std::fmt::Debug for Environment { } impl Environment { - /// Builds an environment from the raw `CODEX_EXEC_SERVER_URL` value. + /// Builds an environment using the caller's effective outbound HTTP policy. pub fn create( exec_server_url: Option, local_runtime_paths: ExecServerRuntimePaths, + http_client_factory: HttpClientFactory, ) -> Result { - Self::create_inner(exec_server_url, Some(local_runtime_paths)) + Self::create_inner( + exec_server_url, + Some(local_runtime_paths), + http_client_factory, + ) } /// Builds a test-only environment without configured sandbox helper paths. pub fn create_for_tests(exec_server_url: Option) -> Result { - Self::create_inner(exec_server_url, /*local_runtime_paths*/ None) + Self::create_inner( + exec_server_url, + /*local_runtime_paths*/ None, + HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault), + ) } /// Builds an environment from the raw `CODEX_EXEC_SERVER_URL` value and @@ -657,6 +674,7 @@ impl Environment { fn create_inner( exec_server_url: Option, local_runtime_paths: Option, + http_client_factory: HttpClientFactory, ) -> Result { let (exec_server_url, disabled) = normalize_exec_server_url(exec_server_url); if disabled { @@ -666,13 +684,16 @@ impl Environment { } Ok(match exec_server_url { - Some(exec_server_url) => Self::remote_inner(exec_server_url, local_runtime_paths), - None => match local_runtime_paths { - Some(local_runtime_paths) => Self::local( - local_runtime_paths, - // This legacy constructor has no resolved application proxy configuration. - HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault), + Some(exec_server_url) => Self::remote_with_transport( + ExecServerTransportParams::websocket_url( + exec_server_url, + DEFAULT_REMOTE_EXEC_SERVER_CONNECT_TIMEOUT, ), + local_runtime_paths, + http_client_factory, + ), + None => match local_runtime_paths { + Some(local_runtime_paths) => Self::local(local_runtime_paths, http_client_factory), None => Self::default_for_tests(), }, }) @@ -697,24 +718,12 @@ impl Environment { } } - pub(crate) fn remote_inner( - exec_server_url: String, - local_runtime_paths: Option, - ) -> Self { - Self::remote_with_transport( - ExecServerTransportParams::websocket_url( - exec_server_url, - DEFAULT_REMOTE_EXEC_SERVER_CONNECT_TIMEOUT, - ), - local_runtime_paths, - ) - } - pub(crate) fn remote_with_transport( remote_transport: ExecServerTransportParams, local_runtime_paths: Option, + http_client_factory: HttpClientFactory, ) -> Self { - let client = LazyRemoteExecServerClient::new(remote_transport); + let client = LazyRemoteExecServerClient::new(remote_transport, http_client_factory); let exec_backend: Arc = Arc::new(RemoteProcess::new(client.clone())); let filesystem: Arc = Arc::new(RemoteFileSystem::new(client.clone())); @@ -883,6 +892,7 @@ mod tests { use super::noise_environment_config_from_values; use crate::ExecServerRuntimePaths; use crate::ProcessId; + use crate::client_api::DEFAULT_REMOTE_EXEC_SERVER_CONNECT_TIMEOUT; use crate::client_api::ExecServerTransportParams; use crate::client_api::StdioExecServerCommand; use crate::environment_provider::EnvironmentDefault; @@ -898,6 +908,13 @@ mod tests { HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault) } + fn prepared_websocket_environment() -> ExecServerTransportParams { + ExecServerTransportParams::websocket_url( + "ws://127.0.0.1:8765".to_string(), + DEFAULT_REMOTE_EXEC_SERVER_CONNECT_TIMEOUT, + ) + } + fn test_runtime_paths() -> ExecServerRuntimePaths { ExecServerRuntimePaths::new( std::env::current_exe().expect("current exe"), @@ -960,8 +977,12 @@ mod tests { #[tokio::test] async fn create_local_environment_does_not_connect() { - let environment = Environment::create(/*exec_server_url*/ None, test_runtime_paths()) - .expect("create environment"); + let environment = Environment::create( + /*exec_server_url*/ None, + test_runtime_paths(), + legacy_http_client_factory(), + ) + .expect("create environment"); assert!(!environment.is_remote()); assert!(environment.info().await.is_ok()); @@ -1050,8 +1071,7 @@ mod tests { let snapshot = EnvironmentProviderSnapshot { environments: vec![( REMOTE_ENVIRONMENT_ID.to_string(), - Environment::create_for_tests(Some("ws://127.0.0.1:8765".to_string())) - .expect("remote environment"), + prepared_websocket_environment(), )], default: EnvironmentDefault::EnvironmentId(REMOTE_ENVIRONMENT_ID.to_string()), include_local: false, @@ -1080,7 +1100,7 @@ mod tests { #[tokio::test] async fn environment_manager_rejects_empty_environment_id() { let snapshot = EnvironmentProviderSnapshot { - environments: vec![("".to_string(), Environment::default_for_tests())], + environments: vec![("".to_string(), prepared_websocket_environment())], default: EnvironmentDefault::Disabled, include_local: false, }; @@ -1102,7 +1122,7 @@ mod tests { let snapshot = EnvironmentProviderSnapshot { environments: vec![( LOCAL_ENVIRONMENT_ID.to_string(), - Environment::default_for_tests(), + prepared_websocket_environment(), )], default: EnvironmentDefault::Disabled, include_local: false, @@ -1123,11 +1143,7 @@ mod tests { #[tokio::test] async fn environment_manager_uses_explicit_provider_default() { let snapshot = EnvironmentProviderSnapshot { - environments: vec![( - "devbox".to_string(), - Environment::create_for_tests(Some("ws://127.0.0.1:8765".to_string())) - .expect("remote environment"), - )], + environments: vec![("devbox".to_string(), prepared_websocket_environment())], default: EnvironmentDefault::EnvironmentId("devbox".to_string()), include_local: true, }; @@ -1149,11 +1165,7 @@ mod tests { #[tokio::test] async fn environment_manager_disables_provider_default() { let snapshot = EnvironmentProviderSnapshot { - environments: vec![( - "devbox".to_string(), - Environment::create_for_tests(Some("ws://127.0.0.1:8765".to_string())) - .expect("remote environment"), - )], + environments: vec![("devbox".to_string(), prepared_websocket_environment())], default: EnvironmentDefault::Disabled, include_local: true, }; @@ -1177,11 +1189,7 @@ mod tests { #[tokio::test] async fn environment_manager_rejects_unknown_provider_default() { let snapshot = EnvironmentProviderSnapshot { - environments: vec![( - "devbox".to_string(), - Environment::create_for_tests(Some("ws://127.0.0.1:8765".to_string())) - .expect("remote environment"), - )], + environments: vec![("devbox".to_string(), prepared_websocket_environment())], default: EnvironmentDefault::EnvironmentId("missing".to_string()), include_local: true, }; @@ -1356,6 +1364,7 @@ mod tests { initialize_timeout: Duration::from_secs(1), }, /*local_runtime_paths*/ None, + legacy_http_client_factory(), ); assert_eq!( @@ -1367,21 +1376,18 @@ mod tests { #[tokio::test] async fn environment_manager_leaves_stdio_environment_lazy() { - let environment = Environment::remote_with_transport( - ExecServerTransportParams::StdioCommand { - command: StdioExecServerCommand { - program: "codex-missing-exec-server-for-test".to_string(), - args: Vec::new(), - env: HashMap::new(), - cwd: None, - }, - initialize_timeout: Duration::from_secs(1), + let transport = ExecServerTransportParams::StdioCommand { + command: StdioExecServerCommand { + program: "codex-missing-exec-server-for-test".to_string(), + args: Vec::new(), + env: HashMap::new(), + cwd: None, }, - /*local_runtime_paths*/ None, - ); + initialize_timeout: Duration::from_secs(1), + }; let manager = EnvironmentManager::from_snapshot( EnvironmentProviderSnapshot { - environments: vec![("stdio".to_string(), environment)], + environments: vec![("stdio".to_string(), transport)], default: EnvironmentDefault::Disabled, include_local: false, }, @@ -1401,21 +1407,18 @@ mod tests { use codex_protocol::capabilities::CapabilityRootLocation; use codex_protocol::capabilities::SelectedCapabilityRoot; - let environment = Environment::remote_with_transport( - ExecServerTransportParams::StdioCommand { - command: StdioExecServerCommand { - program: "codex-missing-exec-server-for-test".to_string(), - args: Vec::new(), - env: HashMap::new(), - cwd: None, - }, - initialize_timeout: Duration::from_secs(1), + let transport = ExecServerTransportParams::StdioCommand { + command: StdioExecServerCommand { + program: "codex-missing-exec-server-for-test".to_string(), + args: Vec::new(), + env: HashMap::new(), + cwd: None, }, - /*local_runtime_paths*/ None, - ); + initialize_timeout: Duration::from_secs(1), + }; let manager = EnvironmentManager::from_snapshot( EnvironmentProviderSnapshot { - environments: vec![("stdio".to_string(), environment)], + environments: vec![("stdio".to_string(), transport)], default: EnvironmentDefault::Disabled, include_local: false, }, diff --git a/codex-rs/exec-server/src/environment_bootstrap.rs b/codex-rs/exec-server/src/environment_bootstrap.rs index 3427321cac..ad557c3b29 100644 --- a/codex-rs/exec-server/src/environment_bootstrap.rs +++ b/codex-rs/exec-server/src/environment_bootstrap.rs @@ -33,8 +33,7 @@ impl PreparedEnvironmentManager { EnvironmentDefault::EnvironmentId(default_id) => snapshot .environments .iter() - .find(|(environment_id, _)| environment_id == default_id) - .is_some_and(|(_, environment)| environment.is_remote()), + .any(|(environment_id, _)| environment_id == default_id), }, } } diff --git a/codex-rs/exec-server/src/environment_bootstrap_tests.rs b/codex-rs/exec-server/src/environment_bootstrap_tests.rs index f4c01df664..d877871b22 100644 --- a/codex-rs/exec-server/src/environment_bootstrap_tests.rs +++ b/codex-rs/exec-server/src/environment_bootstrap_tests.rs @@ -5,35 +5,37 @@ use pretty_assertions::assert_eq; use super::PreparedEnvironmentManager; use super::PreparedEnvironmentSource; use crate::DefaultEnvironmentProvider; -use crate::Environment; -use crate::EnvironmentConnectionState; -use crate::ExecServerRuntimePaths; use crate::LOCAL_ENVIRONMENT_ID; use crate::REMOTE_ENVIRONMENT_ID; +use crate::client_api::DEFAULT_REMOTE_EXEC_SERVER_CONNECT_TIMEOUT; +use crate::client_api::ExecServerTransportParams; use crate::environment_provider::EnvironmentDefault; use crate::environment_provider::EnvironmentProviderSnapshot; use crate::remote::NoiseRendezvousEnvironmentConfig; #[test] -fn prepared_remote_environment_is_detected_without_starting_a_connection() { - let environment = Environment::create_for_tests(Some("ws://127.0.0.1:8765".to_string())) - .expect("remote environment"); - let connection_state = environment - .subscribe_connection_state() - .expect("remote connection state"); +fn prepared_remote_environment_is_detected_without_constructing_a_connection() { let prepared = PreparedEnvironmentManager { source: PreparedEnvironmentSource::Snapshot(EnvironmentProviderSnapshot { - environments: vec![(REMOTE_ENVIRONMENT_ID.to_string(), environment)], + environments: vec![( + REMOTE_ENVIRONMENT_ID.to_string(), + ExecServerTransportParams::websocket_url( + "ws://username:password@executor.example/private?token=secret".to_string(), + DEFAULT_REMOTE_EXEC_SERVER_CONNECT_TIMEOUT, + ), + )], default: EnvironmentDefault::EnvironmentId(REMOTE_ENVIRONMENT_ID.to_string()), include_local: false, }), }; assert!(prepared.default_environment_is_remote()); - assert_eq!( - *connection_state.borrow(), - EnvironmentConnectionState::Disconnected - ); + + let debug = format!("{prepared:?}"); + assert!(!debug.contains("username")); + assert!(!debug.contains("password")); + assert!(!debug.contains("executor.example")); + assert!(!debug.contains("secret")); } #[test] @@ -130,21 +132,27 @@ fn prepared_local_and_disabled_environments_are_not_remote() { async fn prepared_environment_manager_builds_with_the_explicit_http_policy() { let prepared = PreparedEnvironmentManager { source: PreparedEnvironmentSource::Snapshot( - DefaultEnvironmentProvider::new(/*exec_server_url*/ None).snapshot_inner(), + DefaultEnvironmentProvider::new(Some("ws://127.0.0.1:8765".to_string())) + .snapshot_inner(), ), }; - let runtime_paths = ExecServerRuntimePaths::new( - std::env::current_exe().expect("current exe"), - /*codex_linux_sandbox_exe*/ None, - ) - .expect("runtime paths"); let manager = prepared .build( - Some(runtime_paths), + /*local_runtime_paths*/ None, HttpClientFactory::new(OutboundProxyPolicy::RespectSystemProxy), ) .expect("environment manager"); + assert_eq!( + manager.default_environment_id(), + Some(REMOTE_ENVIRONMENT_ID) + ); + assert!( + manager + .default_environment() + .expect("remote environment") + .is_remote() + ); assert_eq!( manager.http_client_factory().outbound_proxy_policy(), OutboundProxyPolicy::RespectSystemProxy diff --git a/codex-rs/exec-server/src/environment_provider.rs b/codex-rs/exec-server/src/environment_provider.rs index 69db34cf65..2366eb8be2 100644 --- a/codex-rs/exec-server/src/environment_provider.rs +++ b/codex-rs/exec-server/src/environment_provider.rs @@ -1,19 +1,20 @@ use std::future::Future; use std::pin::Pin; -use crate::Environment; use crate::ExecServerError; +use crate::client_api::DEFAULT_REMOTE_EXEC_SERVER_CONNECT_TIMEOUT; +use crate::client_api::ExecServerTransportParams; use crate::environment::CODEX_EXEC_SERVER_URL_ENV_VAR; use crate::environment::LOCAL_ENVIRONMENT_ID; use crate::environment::REMOTE_ENVIRONMENT_ID; -/// Lists the concrete environments available to Codex. +/// Lists the remote environment transports available to Codex. /// /// Implementations own a startup snapshot containing both the available -/// environment list in configured order and the default environment -/// selection. Providers should only return provider-owned remote environments; -/// `include_local` controls whether `EnvironmentManager` should add the local -/// environment to the snapshot. +/// environment transport list in configured order and the default environment +/// selection. Providers return transport descriptions before the effective HTTP +/// policy is available; `include_local` controls whether `EnvironmentManager` +/// should add the local environment when the snapshot is built. pub trait EnvironmentProvider: Send + Sync { /// Returns the provider-owned environment startup snapshot. fn snapshot(&self) -> EnvironmentProviderFuture<'_>; @@ -22,13 +23,24 @@ pub trait EnvironmentProvider: Send + Sync { pub type EnvironmentProviderFuture<'a> = Pin> + Send + 'a>>; -#[derive(Clone, Debug)] +#[derive(Clone)] pub struct EnvironmentProviderSnapshot { - pub environments: Vec<(String, Environment)>, + pub(crate) environments: Vec<(String, ExecServerTransportParams)>, pub default: EnvironmentDefault, pub include_local: bool, } +impl std::fmt::Debug for EnvironmentProviderSnapshot { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + let environment_ids: Vec<_> = self.environments.iter().map(|(id, _)| id).collect(); + f.debug_struct("EnvironmentProviderSnapshot") + .field("environments", &environment_ids) + .field("default", &self.default) + .field("include_local", &self.include_local) + .finish() + } +} + #[derive(Clone, Debug, PartialEq, Eq)] pub enum EnvironmentDefault { Disabled, @@ -59,7 +71,10 @@ impl DefaultEnvironmentProvider { if let Some(exec_server_url) = exec_server_url { environments.push(( REMOTE_ENVIRONMENT_ID.to_string(), - Environment::remote_inner(exec_server_url, /*local_runtime_paths*/ None), + ExecServerTransportParams::websocket_url( + exec_server_url, + DEFAULT_REMOTE_EXEC_SERVER_CONNECT_TIMEOUT, + ), )); } @@ -175,8 +190,11 @@ mod tests { assert!(!include_local); assert!(!environments.contains_key(LOCAL_ENVIRONMENT_ID)); - let remote_environment = &environments[REMOTE_ENVIRONMENT_ID]; - assert!(remote_environment.is_remote()); + assert!(matches!( + &environments[REMOTE_ENVIRONMENT_ID], + ExecServerTransportParams::WebSocketUrl { websocket_url, .. } + if websocket_url == "ws://127.0.0.1:8765" + )); assert_eq!( default, EnvironmentDefault::EnvironmentId(REMOTE_ENVIRONMENT_ID.to_string()) diff --git a/codex-rs/exec-server/src/environment_toml.rs b/codex-rs/exec-server/src/environment_toml.rs index 51e97b9bc9..d144c65a95 100644 --- a/codex-rs/exec-server/src/environment_toml.rs +++ b/codex-rs/exec-server/src/environment_toml.rs @@ -8,7 +8,6 @@ use serde::Deserialize; use tokio_tungstenite::tungstenite::client::IntoClientRequest; use crate::DefaultEnvironmentProvider; -use crate::Environment; use crate::EnvironmentProvider; use crate::EnvironmentProviderFuture; use crate::ExecServerError; @@ -94,19 +93,8 @@ impl TomlEnvironmentProvider { } async fn snapshot(&self) -> Result { - let mut environments = Vec::with_capacity(self.environments.len()); - for (id, transport_params) in &self.environments { - environments.push(( - id.clone(), - Environment::remote_with_transport( - transport_params.clone(), - /*local_runtime_paths*/ None, - ), - )); - } - Ok(EnvironmentProviderSnapshot { - environments, + environments: self.environments.clone(), default: self.default.clone(), include_local: self.include_local, }) @@ -389,8 +377,14 @@ mod tests { assert!(include_local); assert!(!environments.contains_key(LOCAL_ENVIRONMENT_ID)); - assert!(environments["devbox"].is_remote()); - assert!(environments["ssh-dev"].is_remote()); + assert!(matches!( + &environments["devbox"], + ExecServerTransportParams::WebSocketUrl { .. } + )); + assert!(matches!( + &environments["ssh-dev"], + ExecServerTransportParams::StdioCommand { .. } + )); assert_eq!( default, EnvironmentDefault::EnvironmentId("ssh-dev".to_string()) diff --git a/codex-rs/exec-server/src/relay.rs b/codex-rs/exec-server/src/relay.rs index b8d477f5db..3d0be695b4 100644 --- a/codex-rs/exec-server/src/relay.rs +++ b/codex-rs/exec-server/src/relay.rs @@ -7,13 +7,10 @@ use futures::SinkExt; use futures::Stream; use futures::StreamExt; use prost::Message as ProstMessage; -use tokio::io::AsyncRead; -use tokio::io::AsyncWrite; use tokio::sync::mpsc; use tokio::sync::watch; use tokio::task::JoinSet; use tokio::time::timeout; -use tokio_tungstenite::WebSocketStream; use tokio_tungstenite::tungstenite::Message; use tracing::debug; use tracing::info; @@ -465,8 +462,8 @@ pub(crate) trait HarnessKeyValidator: Send + Sync { /// Parsing the first Noise message authenticates the harness key. Only a /// successful registry check turns that pending handshake into a virtual stream. #[tracing::instrument(level = "debug", skip_all, fields(noise_side = "executor"))] -pub(crate) async fn run_multiplexed_environment( - stream: WebSocketStream, +pub(crate) async fn run_multiplexed_environment( + stream: T, processor: ConnectionProcessor, environment_id: String, executor_registration_id: String, @@ -474,7 +471,8 @@ pub(crate) async fn run_multiplexed_environment( validator: V, ) -> RendezvousDisconnectReason where - S: AsyncRead + AsyncWrite + Unpin + Send + 'static, + T: Sink + Stream> + Unpin + Send + 'static, + E: std::fmt::Display + Send + 'static, V: HarnessKeyValidator + Clone + 'static, { debug!( diff --git a/codex-rs/exec-server/src/remote.rs b/codex-rs/exec-server/src/remote.rs index 8a3c1703e0..731c67c293 100644 --- a/codex-rs/exec-server/src/remote.rs +++ b/codex-rs/exec-server/src/remote.rs @@ -15,15 +15,15 @@ use http::HeaderValue; use reqwest::StatusCode; use serde::Deserialize; use tokio::time::sleep; -use tokio_tungstenite::MaybeTlsStream; -use tokio_tungstenite::WebSocketStream; -use tokio_tungstenite::connect_async_with_config; use tokio_tungstenite::tungstenite::client::IntoClientRequest; use tracing::debug; use tracing::info; use tracing::warn; use codex_utils_rustls_provider::ensure_rustls_crypto_provider; +use codex_websocket_client::WebSocketConnection; +use codex_websocket_client::WebSocketConnector; +use codex_websocket_client::WebSocketTlsMode; use crate::EnvironmentRegistryConnectRequest; use crate::EnvironmentRegistryConnectResponse; @@ -514,7 +514,13 @@ pub async fn run_remote_environment( .await?; loop { - match connect_rendezvous(&response.url, &config.telemetry).await { + match connect_rendezvous( + &response.url, + &config.telemetry, + &config.http_client_factory, + ) + .await + { Ok(websocket) => { backoff = Duration::from_secs(1); let executor_registration_id = response.executor_registration_id.clone(); @@ -587,25 +593,24 @@ pub async fn run_remote_environment( async fn connect_rendezvous( url: &str, telemetry: &ExecServerTelemetry, -) -> Result< - WebSocketStream>, - tokio_tungstenite::tungstenite::Error, -> { + http_client_factory: &HttpClientFactory, +) -> Result { let started_at = Instant::now(); let result = async { let mut request = url.into_client_request()?; request .headers_mut() .extend(current_trace_context_headers()); - connect_async_with_config( - request, - Some(noise_relay_websocket_config()), - // Rendezvous sends small, latency-sensitive frames, so avoid Nagle's coalescing delay. - /*disable_nagle*/ - true, + let connector = WebSocketConnector::new_with_tls_mode( + http_client_factory, + WebSocketTlsMode::TungsteniteDefault, ) - .await - .map(|(websocket, _)| websocket) + .map_err(|error| tokio_tungstenite::tungstenite::Error::Io(std::io::Error::other(error)))?; + connector + .with_tcp_nodelay() + .connect(request, noise_relay_websocket_config()) + .await + .map(|(websocket, _)| websocket) } .await; let result_name = if result.is_ok() { "success" } else { "error" }; diff --git a/codex-rs/exec-server/src/remote_file_system_path_uri_tests.rs b/codex-rs/exec-server/src/remote_file_system_path_uri_tests.rs index 1c989704a6..2f4f99ac45 100644 --- a/codex-rs/exec-server/src/remote_file_system_path_uri_tests.rs +++ b/codex-rs/exec-server/src/remote_file_system_path_uri_tests.rs @@ -2,6 +2,8 @@ use codex_exec_server_protocol::JSONRPCMessage; use codex_exec_server_protocol::JSONRPCResponse; +use codex_http_client::HttpClientFactory; +use codex_http_client::OutboundProxyPolicy; use codex_protocol::models::PermissionProfile; use codex_protocol::permissions::FileSystemAccessMode; use codex_protocol::permissions::FileSystemPath; @@ -41,6 +43,7 @@ async fn remote_file_system_sends_path_and_sandbox_cwd_uris_without_native_conve websocket_url, DEFAULT_REMOTE_EXEC_SERVER_CONNECT_TIMEOUT, ), + HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault), )); let paths = vec![ PathUri::parse("file:///C:/Users/Alice/src/main.rs").expect("valid drive URI"), diff --git a/codex-rs/exec-server/src/server/handler/tests.rs b/codex-rs/exec-server/src/server/handler/tests.rs index 0d3b41768a..3ecfd7dddc 100644 --- a/codex-rs/exec-server/src/server/handler/tests.rs +++ b/codex-rs/exec-server/src/server/handler/tests.rs @@ -83,6 +83,10 @@ fn test_runtime_paths() -> ExecServerRuntimePaths { .expect("runtime paths") } +fn test_http_client_factory() -> HttpClientFactory { + HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault) +} + async fn initialized_handler() -> Arc { let (outgoing_tx, _outgoing_rx) = mpsc::channel(16); let registry = SessionRegistry::new(crate::ExecServerTelemetry::default()); @@ -90,7 +94,7 @@ async fn initialized_handler() -> Arc { registry, RpcNotificationSender::new(outgoing_tx), test_runtime_paths(), - HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault), + test_http_client_factory(), )); let initialize_response = handler .initialize(InitializeParams { @@ -169,7 +173,7 @@ async fn long_poll_read_fails_after_session_resume() { Arc::clone(®istry), RpcNotificationSender::new(first_tx), test_runtime_paths(), - HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault), + test_http_client_factory(), )); let initialize_response = first_handler .initialize(InitializeParams { @@ -210,7 +214,7 @@ async fn long_poll_read_fails_after_session_resume() { registry, RpcNotificationSender::new(second_tx), test_runtime_paths(), - HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault), + test_http_client_factory(), )); second_handler .initialize(InitializeParams { @@ -244,7 +248,7 @@ async fn active_session_resume_is_rejected() { Arc::clone(®istry), RpcNotificationSender::new(first_tx), test_runtime_paths(), - HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault), + test_http_client_factory(), )); let initialize_response = first_handler .initialize(InitializeParams { @@ -259,7 +263,7 @@ async fn active_session_resume_is_rejected() { registry, RpcNotificationSender::new(second_tx), test_runtime_paths(), - HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault), + test_http_client_factory(), )); let err = second_handler .initialize(InitializeParams { @@ -288,7 +292,7 @@ async fn output_and_exit_are_retained_after_notification_receiver_closes() { SessionRegistry::new(crate::ExecServerTelemetry::default()), RpcNotificationSender::new(outgoing_tx), test_runtime_paths(), - HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault), + test_http_client_factory(), )); handler .initialize(InitializeParams { diff --git a/codex-rs/exec-server/tests/environment.rs b/codex-rs/exec-server/tests/environment.rs index 4b97361ac2..01cb6a16df 100644 --- a/codex-rs/exec-server/tests/environment.rs +++ b/codex-rs/exec-server/tests/environment.rs @@ -6,13 +6,95 @@ use anyhow::Context; use codex_exec_server::EnvironmentManager; use codex_exec_server::REMOTE_ENVIRONMENT_ID; use codex_exec_server::SelectedCapabilityRootsStatus; +use codex_http_client::HttpClientFactory; +use codex_http_client::OutboundProxyPolicy; +use codex_http_client::cache_system_proxy_route_for_test; use codex_protocol::capabilities::CapabilityRootLocation; use codex_protocol::capabilities::SelectedCapabilityRoot; use codex_utils_path_uri::PathUri; use common::exec_server::exec_server; use pretty_assertions::assert_eq; +use tokio::io::AsyncReadExt; +use tokio::io::AsyncWriteExt; +use tokio::net::TcpListener; +use tokio::net::TcpStream; +use tokio::sync::oneshot; use tokio::time::sleep; use tokio::time::timeout; +use tokio_util::task::AbortOnDropHandle; + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +#[serial_test::serial(remote_exec_server)] +async fn prepared_remote_environment_uses_configured_system_proxy() -> anyhow::Result<()> { + let server = exec_server().await?; + let upstream = server + .websocket_url() + .strip_prefix("ws://") + .context("exec-server websocket should use ws://")? + .to_string(); + let proxy_listener = TcpListener::bind("127.0.0.1:0").await?; + let proxy_url = format!("http://{}", proxy_listener.local_addr()?); + let websocket_url = "ws://exec-server-system-proxy.invalid:8765/"; + let proxy_resolution_url = "http://exec-server-system-proxy.invalid:8765/"; + cache_system_proxy_route_for_test(proxy_resolution_url, proxy_url); + + let (request_tx, request_rx) = oneshot::channel(); + let _proxy_task = AbortOnDropHandle::new(tokio::spawn(async move { + let (mut client, _) = proxy_listener.accept().await?; + let mut request = Vec::new(); + let mut byte = [0_u8; 1]; + while !request.ends_with(b"\r\n\r\n") { + client.read_exact(&mut byte).await?; + request.push(byte[0]); + } + let request_line = String::from_utf8(request)? + .lines() + .next() + .context("system proxy should receive a CONNECT request")? + .to_string(); + request_tx + .send(request_line) + .map_err(|_| anyhow::anyhow!("system proxy request receiver was dropped"))?; + + let mut target = TcpStream::connect(upstream).await?; + client + .write_all(b"HTTP/1.1 200 Connection Established\r\n\r\n") + .await?; + tokio::io::copy_bidirectional(&mut client, &mut target).await?; + Ok::<(), anyhow::Error>(()) + })); + + let codex_home = tempfile::tempdir()?; + std::fs::write( + codex_home.path().join("environments.toml"), + format!( + "default = \"{REMOTE_ENVIRONMENT_ID}\"\ninclude_local = false\n\n[[environments]]\nid = \"{REMOTE_ENVIRONMENT_ID}\"\nurl = \"{websocket_url}\"\n" + ), + )?; + + let prepared = EnvironmentManager::prepare_from_codex_home(codex_home.path()).await?; + assert!(prepared.default_environment_is_remote()); + let manager = prepared.build( + /*local_runtime_paths*/ None, + HttpClientFactory::new(OutboundProxyPolicy::RespectSystemProxy), + )?; + + let request_line = timeout(Duration::from_secs(5), request_rx) + .await + .context("prepared environment did not connect through the system proxy")??; + assert_eq!( + request_line, + "CONNECT exec-server-system-proxy.invalid:8765 HTTP/1.1" + ); + let environment = manager + .default_environment() + .context("prepared remote environment")?; + timeout(Duration::from_secs(5), environment.info()) + .await + .context("prepared remote environment did not initialize through the system proxy")??; + + Ok(()) +} #[tokio::test(flavor = "multi_thread", worker_threads = 2)] #[serial_test::serial(remote_exec_server)] diff --git a/codex-rs/exec-server/tests/file_stream.rs b/codex-rs/exec-server/tests/file_stream.rs index 9376792b6b..af00b161a9 100644 --- a/codex-rs/exec-server/tests/file_stream.rs +++ b/codex-rs/exec-server/tests/file_stream.rs @@ -11,6 +11,8 @@ use codex_exec_server::FsOpenParams; use codex_exec_server::FsReadBlockParams; use codex_exec_server::FsReadBlockResponse; use codex_exec_server::RemoteExecServerConnectArgs; +use codex_http_client::HttpClientFactory; +use codex_http_client::OutboundProxyPolicy; use codex_protocol::models::PermissionProfile; use codex_protocol::permissions::FileSystemAccessMode; use codex_protocol::permissions::FileSystemPath; @@ -234,6 +236,7 @@ async fn read_block_supports_non_sequential_offsets_and_lengths() -> Result<()> let client = ExecServerClient::connect_websocket(RemoteExecServerConnectArgs::new( server.websocket_url().to_string(), "file-stream-protocol-test".to_string(), + HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault), )) .await?; let tmp = TempDir::new()?; @@ -296,6 +299,7 @@ async fn open_enforces_the_per_connection_limit_and_close_releases_capacity() -> let client = ExecServerClient::connect_websocket(RemoteExecServerConnectArgs::new( server.websocket_url().to_string(), "file-stream-protocol-test".to_string(), + HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault), )) .await?; let tmp = TempDir::new()?; @@ -356,6 +360,7 @@ async fn open_rejects_handle_ids_longer_than_32_bytes() -> Result<()> { let client = ExecServerClient::connect_websocket(RemoteExecServerConnectArgs::new( server.websocket_url().to_string(), "file-stream-protocol-test".to_string(), + HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault), )) .await?; let tmp = TempDir::new()?; diff --git a/codex-rs/exec-server/tests/http_client.rs b/codex-rs/exec-server/tests/http_client.rs index 8fd3e4ea3f..b9593c2f6b 100644 --- a/codex-rs/exec-server/tests/http_client.rs +++ b/codex-rs/exec-server/tests/http_client.rs @@ -19,6 +19,8 @@ use codex_exec_server_protocol::JSONRPCRequest; use codex_exec_server_protocol::JSONRPCResponse; use codex_exec_server_protocol::MAX_HTTP_BODY_DELTA_BYTES; use codex_exec_server_protocol::RequestId; +use codex_http_client::HttpClientFactory; +use codex_http_client::OutboundProxyPolicy; use futures::SinkExt; use futures::StreamExt; use pretty_assertions::assert_eq; @@ -1369,6 +1371,7 @@ impl ScriptedExecServer { ExecServerClient::connect_websocket(RemoteExecServerConnectArgs::new( self.websocket_url.clone(), CLIENT_NAME.to_string(), + HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault), )) .await .context("client should connect to fake exec-server") diff --git a/codex-rs/exec-server/tests/relay.rs b/codex-rs/exec-server/tests/relay.rs index 6ae4d226b1..1b03096df6 100644 --- a/codex-rs/exec-server/tests/relay.rs +++ b/codex-rs/exec-server/tests/relay.rs @@ -16,6 +16,7 @@ use base64::Engine as _; use base64::engine::general_purpose::STANDARD; use codex_api::AuthProvider; use codex_exec_server::EnvironmentConnectionState; +use codex_exec_server::EnvironmentManager; use codex_exec_server::EnvironmentReadyInfo; use codex_exec_server::ExecParams; use codex_exec_server::ExecResponse; @@ -30,9 +31,9 @@ use codex_exec_server::NoiseRendezvousConnectBundle; use codex_exec_server::NoiseRendezvousConnectProvider; use codex_exec_server::ProcessId; use codex_exec_server::RemoteEnvironmentConfig; -use codex_exec_server_test_support::environment_manager_without_environments; use codex_http_client::HttpClientFactory; use codex_http_client::OutboundProxyPolicy; +use codex_http_client::cache_system_proxy_route_for_test; use codex_protocol::capabilities::CapabilityRootLocation; use codex_protocol::capabilities::SelectedCapabilityRoot; use codex_utils_path_uri::PathUri; @@ -46,13 +47,18 @@ use prost::Message as ProstMessage; use relay_proto::RelayMessageFrame; use relay_proto::relay_message_frame; use tempfile::TempDir; +use tokio::io::AsyncReadExt; +use tokio::io::AsyncWriteExt; use tokio::net::TcpListener; use tokio::net::TcpStream; +use tokio::sync::mpsc; use tokio::sync::watch; +use tokio::task::JoinSet; use tokio::time::timeout; use tokio_tungstenite::WebSocketStream; use tokio_tungstenite::accept_async; use tokio_tungstenite::tungstenite::Message; +use tokio_util::task::AbortOnDropHandle; use wiremock::Mock; use wiremock::MockServer; use wiremock::ResponseTemplate; @@ -114,7 +120,46 @@ impl NoiseRendezvousConnectProvider for FreshBundleNoiseConnectProvider { #[tokio::test(flavor = "multi_thread", worker_threads = 4)] async fn deferred_noise_environment_connects_and_reconnects_with_fresh_bundle() -> Result<()> { let listener = TcpListener::bind("127.0.0.1:0").await?; - let rendezvous_url = format!("ws://{}", listener.local_addr()?); + let rendezvous_address = listener.local_addr()?; + let environment_rendezvous_url = + "ws://environment-noise-relay-system-proxy.invalid:8765/relay?role=environment"; + let harness_rendezvous_url = + "ws://harness-noise-relay-system-proxy.invalid:8765/relay?role=harness"; + let proxy_listener = TcpListener::bind("127.0.0.1:0").await?; + let proxy_url = format!("http://{}", proxy_listener.local_addr()?); + for rendezvous_url in [environment_rendezvous_url, harness_rendezvous_url] { + let proxy_resolution_url = rendezvous_url.replacen("ws://", "http://", /*count*/ 1); + cache_system_proxy_route_for_test(&proxy_resolution_url, proxy_url.clone()); + } + let (proxy_request_tx, mut proxy_request_rx) = mpsc::unbounded_channel(); + let _proxy_task = AbortOnDropHandle::new(tokio::spawn(async move { + let mut proxy_connections = JoinSet::new(); + while let Ok((mut client, _)) = proxy_listener.accept().await { + let proxy_request_tx = proxy_request_tx.clone(); + proxy_connections.spawn(async move { + let mut request = Vec::new(); + let mut byte = [0_u8; 1]; + while !request.ends_with(b"\r\n\r\n") { + client.read_exact(&mut byte).await?; + request.push(byte[0]); + } + let request_line = String::from_utf8(request)? + .lines() + .next() + .context("system proxy should receive a CONNECT request")? + .to_string(); + proxy_request_tx + .send(request_line) + .map_err(|_| anyhow::anyhow!("system proxy request receiver was dropped"))?; + let mut target = TcpStream::connect(rendezvous_address).await?; + client + .write_all(b"HTTP/1.1 200 Connection Established\r\n\r\n") + .await?; + tokio::io::copy_bidirectional(&mut client, &mut target).await?; + Ok::<(), anyhow::Error>(()) + }); + } + })); let registry = MockServer::start().await; Mock::given(method("POST")) .and(path(format!( @@ -122,7 +167,7 @@ async fn deferred_noise_environment_connects_and_reconnects_with_fresh_bundle() ))) .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({ "environment_id": ENVIRONMENT_ID, - "url": format!("{rendezvous_url}/relay?role=environment"), + "url": environment_rendezvous_url, "security_profile": "noise_hybrid_ik_v1", "executor_registration_id": EXECUTOR_REGISTRATION_ID, }))) @@ -142,11 +187,12 @@ async fn deferred_noise_environment_connects_and_reconnects_with_fresh_bundle() let (codex_exe, codex_linux_sandbox_exe) = common::current_test_binary_helper_paths()?; let runtime_paths = ExecServerRuntimePaths::new(codex_exe, codex_linux_sandbox_exe)?; + let http_client_factory = HttpClientFactory::new(OutboundProxyPolicy::RespectSystemProxy); let config = RemoteEnvironmentConfig::new( registry.uri(), ENVIRONMENT_ID.to_string(), static_registry_auth_provider(), - HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault), + http_client_factory.clone(), )?; let remote_environment = tokio::spawn(codex_exec_server::run_remote_environment( config, @@ -154,12 +200,19 @@ async fn deferred_noise_environment_connects_and_reconnects_with_fresh_bundle() )); let environment_websocket = accept_websocket(&listener, "environment").await?; + let environment_proxy_request = + "CONNECT environment-noise-relay-system-proxy.invalid:8765 HTTP/1.1"; + let harness_proxy_request = "CONNECT harness-noise-relay-system-proxy.invalid:8765 HTTP/1.1"; + assert_eq!( + timeout(TEST_TIMEOUT, proxy_request_rx.recv()).await?, + Some(environment_proxy_request.to_string()) + ); let provider = Arc::new(FreshBundleNoiseConnectProvider { - websocket_url: format!("{rendezvous_url}/relay?role=harness"), + websocket_url: harness_rendezvous_url.to_string(), executor_public_key: registered_executor_public_key(®istry).await?, calls: AtomicUsize::new(0), }); - let manager = environment_manager_without_environments(); + let manager = EnvironmentManager::without_environments(http_client_factory); let registration = manager .register_deferred_noise_environment(ENVIRONMENT_ID.to_string(), provider.clone())?; let environment = manager @@ -181,6 +234,10 @@ async fn deferred_noise_environment_connects_and_reconnects_with_fresh_bundle() selected_capability_roots: selected_capability_roots.clone(), }))?; let harness_websocket = accept_websocket(&listener, "harness").await?; + assert_eq!( + timeout(TEST_TIMEOUT, proxy_request_rx.recv()).await?, + Some(harness_proxy_request.to_string()) + ); let first_relay = tokio::spawn(proxy_relay_frames( environment_websocket, harness_websocket, @@ -207,6 +264,22 @@ async fn deferred_noise_environment_connects_and_reconnects_with_fresh_bundle() ); let first_reconnected_websocket = accept_websocket(&listener, "reconnected peer").await?; let second_reconnected_websocket = accept_websocket(&listener, "reconnected peer").await?; + let mut reconnect_proxy_requests = vec![ + timeout(TEST_TIMEOUT, proxy_request_rx.recv()) + .await? + .context("first reconnected peer should use the system proxy")?, + timeout(TEST_TIMEOUT, proxy_request_rx.recv()) + .await? + .context("second reconnected peer should use the system proxy")?, + ]; + reconnect_proxy_requests.sort(); + assert_eq!( + reconnect_proxy_requests, + vec![ + environment_proxy_request.to_string(), + harness_proxy_request.to_string(), + ] + ); let second_relay = tokio::spawn(proxy_relay_frames( first_reconnected_websocket, second_reconnected_websocket, @@ -300,6 +373,9 @@ async fn remote_environment_routes_encrypted_exec_server_rpc() -> Result<()> { connect_timeout: TEST_TIMEOUT, initialize_timeout: TEST_TIMEOUT, resume_session_id: None, + http_client_factory: codex_http_client::HttpClientFactory::new( + codex_http_client::OutboundProxyPolicy::ReqwestDefault, + ), }; let client_task = tokio::spawn(async move { ExecServerClient::connect_noise_rendezvous(client_args).await }); diff --git a/codex-rs/rmcp-client/tests/streamable_http_test_support.rs b/codex-rs/rmcp-client/tests/streamable_http_test_support.rs index 87a52f55df..a765665b48 100644 --- a/codex-rs/rmcp-client/tests/streamable_http_test_support.rs +++ b/codex-rs/rmcp-client/tests/streamable_http_test_support.rs @@ -23,6 +23,8 @@ use codex_exec_server::Environment; use codex_exec_server::ExecServerClient; use codex_exec_server::HttpClient; use codex_exec_server::RemoteExecServerConnectArgs; +use codex_http_client::HttpClientFactory; +use codex_http_client::OutboundProxyPolicy; use codex_rmcp_client::ElicitationAction; use codex_rmcp_client::ElicitationResponse; use codex_rmcp_client::RmcpClient; @@ -345,6 +347,7 @@ pub(crate) async fn spawn_exec_server() -> anyhow::Result { let client = ExecServerClient::connect_websocket(RemoteExecServerConnectArgs::new( websocket_url, "rmcp-client-remote-http-test".to_string(), + HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault), )) .await?; diff --git a/codex-rs/websocket-client/src/dialer.rs b/codex-rs/websocket-client/src/dialer.rs index 93fedd3e64..53d5b763b5 100644 --- a/codex-rs/websocket-client/src/dialer.rs +++ b/codex-rs/websocket-client/src/dialer.rs @@ -6,6 +6,7 @@ use std::sync::Arc; use std::time::Duration; use codex_http_client::OutboundProxyRoute; +use codex_http_client::build_rustls_client_config_with_custom_ca; use futures::StreamExt; use futures::stream::FuturesUnordered; use rustls::ClientConfig; @@ -28,15 +29,18 @@ use tokio_tungstenite::tungstenite::proxy::ProxyConfig; use crate::AsyncIo; use crate::ConnectionInner; +use crate::TcpNodelay; const HAPPY_EYEBALLS_DELAY: Duration = Duration::from_millis(250); pub(crate) async fn connect( request: Request, config: WebSocketConfig, - tls_config: Arc, + tls_config: Option>, proxy_route: OutboundProxyRoute, + tcp_nodelay: TcpNodelay, ) -> Result<(ConnectionInner, Response), WebSocketError> { + let disable_nagle = tcp_nodelay == TcpNodelay::Enabled; let proxy_url = match proxy_route { OutboundProxyRoute::TransportDefault => { // The workspace enables tokio-tungstenite's `proxy` feature, so its default dialer @@ -44,8 +48,8 @@ pub(crate) async fn connect( let (stream, response) = connect_async_tls_with_config( request, Some(config), - false, // Preserve Tungstenite's recommended Nagle default. - Some(Connector::Rustls(tls_config)), + disable_nagle, + tls_config.map(Connector::Rustls), ) .await?; return Ok((ConnectionInner::TransportDefault(stream), response)); @@ -65,8 +69,8 @@ pub(crate) async fn connect( match connect_async_tls_with_config( request.clone(), Some(config), - false, // Preserve Tungstenite's recommended Nagle default. - Some(Connector::Rustls(Arc::clone(&tls_config))), + disable_nagle, + tls_config.clone().map(Connector::Rustls), ) .await { @@ -84,7 +88,7 @@ pub(crate) async fn connect( let host = websocket_host(&request)?; let port = websocket_port(&request)?; Box::new( - connect_tcp(host_port(host, port)) + connect_tcp(host_port(host, port), tcp_nodelay) .await .map_err(WebSocketError::Io)?, ) @@ -93,13 +97,18 @@ pub(crate) async fn connect( let proxy = ProxyEndpoint::parse(&url)?; let host = websocket_host(&request)?; let port = websocket_port(&request)?; - let stream = connect_tcp(proxy.config.authority()) + let stream = connect_tcp(proxy.config.authority(), tcp_nodelay) .await .map_err(WebSocketError::Io)?; let stream: Box = if proxy.tls { + let proxy_tls_config = match &tls_config { + Some(tls_config) => Arc::clone(tls_config), + None => build_rustls_client_config_with_custom_ca() + .map_err(|error| WebSocketError::Io(error.into()))?, + }; let server_name = ServerName::try_from(proxy.config.host.clone()) .map_err(|_| WebSocketError::Tls(TlsError::InvalidDnsName))?; - let stream = TlsConnector::from(Arc::clone(&tls_config)) + let stream = TlsConnector::from(proxy_tls_config) .connect(server_name, stream) .await .map_err(WebSocketError::Io)?; @@ -115,7 +124,7 @@ pub(crate) async fn connect( request, stream, Some(config), - Some(Connector::Rustls(tls_config)), + tls_config.map(Connector::Rustls), ) .await?; Ok((ConnectionInner::Routed(stream), response)) @@ -183,9 +192,13 @@ fn host_port(host: &str, port: u16) -> String { } } -async fn connect_tcp(address: String) -> io::Result { +async fn connect_tcp(address: String, tcp_nodelay: TcpNodelay) -> io::Result { let addresses = tokio::net::lookup_host(address).await?.collect::>(); - connect_happy_eyeballs(addresses, TcpStream::connect).await + let stream = connect_happy_eyeballs(addresses, TcpStream::connect).await?; + if tcp_nodelay == TcpNodelay::Enabled { + stream.set_nodelay(/*nodelay*/ true)?; + } + Ok(stream) } async fn connect_happy_eyeballs( diff --git a/codex-rs/websocket-client/src/dialer_tests.rs b/codex-rs/websocket-client/src/dialer_tests.rs index eae86d90a0..b8c3574305 100644 --- a/codex-rs/websocket-client/src/dialer_tests.rs +++ b/codex-rs/websocket-client/src/dialer_tests.rs @@ -24,6 +24,7 @@ use tokio::net::TcpListener; use tokio::sync::Mutex; use tokio::task::JoinHandle; use tokio_rustls::TlsAcceptor; +use tokio_tungstenite::MaybeTlsStream; use tokio_tungstenite::accept_async; use tokio_tungstenite::tungstenite::Message; use tokio_tungstenite::tungstenite::client::IntoClientRequest; @@ -34,6 +35,7 @@ use super::*; use crate::AsyncIo; use crate::WebSocketConnection; use crate::WebSocketConnector; +use crate::WebSocketTlsMode; #[tokio::test] async fn public_connector_uses_factory_and_exposes_stream_and_sink() { @@ -48,6 +50,7 @@ async fn public_connector_uses_factory_and_exposes_stream_and_sink() { .connect(request, WebSocketConfig::default()) .await .expect("websocket handshake should succeed"); + assert!(!websocket_tcp_nodelay(&websocket)); let expected = Message::Text("hello".into()); websocket .send(expected.clone()) @@ -63,6 +66,120 @@ async fn public_connector_uses_factory_and_exposes_stream_and_sink() { target_task.await.expect("target task should finish"); } +#[tokio::test] +async fn public_connector_enables_tcp_nodelay_when_requested() { + let (target_addr, target_task) = start_echo_websocket_server(/*acceptor*/ None).await; + let request = format!("ws://localhost:{}/v1/responses", target_addr.port()) + .into_client_request() + .expect("websocket request should build"); + let factory = HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault); + let connector = WebSocketConnector::new(&factory) + .expect("connector should build") + .with_tcp_nodelay(); + + let (mut websocket, _) = connector + .connect(request, WebSocketConfig::default()) + .await + .expect("websocket handshake should succeed"); + assert!(websocket_tcp_nodelay(&websocket)); + let expected = Message::Text("latency-sensitive".into()); + websocket + .send(expected.clone()) + .await + .expect("websocket should send"); + assert_eq!( + websocket + .next() + .await + .expect("websocket should receive a message") + .expect("websocket message should be valid"), + expected + ); + + target_task.await.expect("target task should finish"); +} + +#[tokio::test] +async fn tungstenite_default_tls_mode_ignores_invalid_custom_ca_in_a_subprocess() { + let (target_addr, target_task) = start_echo_websocket_server(/*acceptor*/ None).await; + let target_url = format!("ws://127.0.0.1:{}/v1/responses", target_addr.port()); + let executable = std::env::current_exe().expect("test executable should be available"); + let output = tokio::task::spawn_blocking(move || { + let mut command = Command::new(executable); + command.args([ + "--exact", + "dialer::tests::tungstenite_default_tls_mode_subprocess_probe", + "--nocapture", + ]); + for key in [ + "HTTP_PROXY", + "http_proxy", + "HTTPS_PROXY", + "https_proxy", + "ALL_PROXY", + "all_proxy", + "NO_PROXY", + "no_proxy", + "SSL_CERT_FILE", + ] { + command.env_remove(key); + } + command + .env( + "CODEX_CA_CERTIFICATE", + "/codex-websocket-client-nonexistent-custom-ca.pem", + ) + .env("CODEX_WEBSOCKET_DEFAULT_TLS_PROBE_URL", target_url) + .output() + .expect("WebSocket default-TLS subprocess should run") + }) + .await + .expect("WebSocket default-TLS subprocess should join"); + + assert!( + output.status.success(), + "WebSocket default-TLS subprocess failed\nstdout:\n{}\nstderr:\n{}", + String::from_utf8_lossy(&output.stdout), + String::from_utf8_lossy(&output.stderr) + ); + target_task.await.expect("target task should finish"); +} + +#[tokio::test] +async fn tungstenite_default_tls_mode_subprocess_probe() { + let Ok(url) = std::env::var("CODEX_WEBSOCKET_DEFAULT_TLS_PROBE_URL") else { + return; + }; + let factory = HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault); + assert!( + WebSocketConnector::new(&factory).is_err(), + "explicit Codex TLS should reject the invalid custom CA" + ); + let connector = + WebSocketConnector::new_with_tls_mode(&factory, WebSocketTlsMode::TungsteniteDefault) + .expect("Tungstenite-default TLS should ignore custom CA configuration"); + let request = url + .into_client_request() + .expect("websocket request should build"); + let (mut websocket, _) = connector + .connect(request, WebSocketConfig::default()) + .await + .expect("WebSocket should connect without constructing Codex TLS"); + let expected = Message::Text("Tungstenite default TLS".into()); + websocket + .send(expected.clone()) + .await + .expect("WebSocket should send"); + assert_eq!( + websocket + .next() + .await + .expect("WebSocket should receive a message") + .expect("WebSocket message should be valid"), + expected + ); +} + #[tokio::test] async fn direct_route_connects_secure_websocket() { let (tls_config, acceptor, _) = test_tls_configs(); @@ -74,8 +191,9 @@ async fn direct_route_connects_secure_websocket() { let (inner, _) = connect( request, WebSocketConfig::default(), - tls_config, + Some(tls_config), OutboundProxyRoute::Direct, + TcpNodelay::Enabled, ) .await .expect("direct websocket handshake should succeed"); @@ -157,11 +275,12 @@ async fn no_proxy_subprocess_probe() { let (inner, _) = connect( request, WebSocketConfig::default(), - tls_config, + Some(tls_config), OutboundProxyRoute::Proxy { url: proxy_url, no_proxy: Some(no_proxy), }, + TcpNodelay::Enabled, ) .await .expect("websocket handshake should succeed"); @@ -238,6 +357,41 @@ async fn happy_eyeballs_does_not_wait_for_stalled_preferred_family() { assert_eq!(connected, reachable); } +#[tokio::test] +async fn routed_tcp_connections_only_enable_nodelay_when_requested() { + let listener = TcpListener::bind("127.0.0.1:0") + .await + .expect("listener should bind"); + let address = listener + .local_addr() + .expect("listener should have an address") + .to_string(); + + let default_stream = connect_tcp(address.clone(), TcpNodelay::Default) + .await + .expect("default TCP connection should succeed"); + assert!(!default_stream.nodelay().expect("TCP_NODELAY should read")); + + let low_latency_stream = connect_tcp(address, TcpNodelay::Enabled) + .await + .expect("low-latency TCP connection should succeed"); + assert!( + low_latency_stream + .nodelay() + .expect("TCP_NODELAY should read") + ); +} + +fn websocket_tcp_nodelay(websocket: &WebSocketConnection) -> bool { + let ConnectionInner::TransportDefault(stream) = &websocket.inner else { + panic!("default connector should use Tungstenite's transport"); + }; + let MaybeTlsStream::Plain(stream) = stream.get_ref() else { + panic!("test websocket should use a plain TCP stream"); + }; + stream.nodelay().expect("TCP_NODELAY should read") +} + async fn start_echo_websocket_server( acceptor: Option, ) -> (SocketAddr, JoinHandle<()>) { @@ -453,11 +607,12 @@ async fn assert_proxy_tunnels_secure_websocket(proxy_tls: bool) { let (inner, _) = connect( request, WebSocketConfig::default(), - tls_config, + Some(tls_config), OutboundProxyRoute::Proxy { url: format!("{proxy_scheme}://localhost:{}", proxy_addr.port()), no_proxy: None, }, + TcpNodelay::Enabled, ) .await .expect("proxied websocket handshake should succeed"); diff --git a/codex-rs/websocket-client/src/lib.rs b/codex-rs/websocket-client/src/lib.rs index b49046b6a0..932cc6498a 100644 --- a/codex-rs/websocket-client/src/lib.rs +++ b/codex-rs/websocket-client/src/lib.rs @@ -32,7 +32,23 @@ use tokio_tungstenite::tungstenite::protocol::WebSocketConfig; #[derive(Clone)] pub struct WebSocketConnector { http_client_factory: HttpClientFactory, - tls_config: Arc, + tls_config: Option>, + tcp_nodelay: TcpNodelay, +} + +/// Selects whether WebSocket TLS follows Codex custom-CA policy or Tungstenite defaults. +#[derive(Clone, Copy, PartialEq, Eq)] +pub enum WebSocketTlsMode { + /// Build an explicit TLS configuration from native roots and configured Codex custom CAs. + ExplicitCodexTls, + /// Let Tungstenite build its default TLS configuration when the target requires TLS. + TungsteniteDefault, +} + +#[derive(Clone, Copy, PartialEq, Eq)] +pub(crate) enum TcpNodelay { + Default, + Enabled, } impl WebSocketConnector { @@ -40,12 +56,36 @@ impl WebSocketConnector { pub fn new( http_client_factory: &HttpClientFactory, ) -> Result { + Self::new_with_tls_mode(http_client_factory, WebSocketTlsMode::ExplicitCodexTls) + } + + /// Creates a connector with explicit Codex TLS or the transport's existing TLS defaults. + /// + /// HTTPS proxy connections still build Codex TLS configuration when they establish their + /// proxy tunnel; default-mode target connections otherwise remain entirely with Tungstenite. + pub fn new_with_tls_mode( + http_client_factory: &HttpClientFactory, + tls_mode: WebSocketTlsMode, + ) -> Result { + let tls_config = match tls_mode { + WebSocketTlsMode::ExplicitCodexTls => { + Some(build_rustls_client_config_with_custom_ca()?) + } + WebSocketTlsMode::TungsteniteDefault => None, + }; Ok(Self { http_client_factory: http_client_factory.clone(), - tls_config: build_rustls_client_config_with_custom_ca()?, + tls_config, + tcp_nodelay: TcpNodelay::Default, }) } + /// Disables Nagle's algorithm for latency-sensitive WebSocket connections. + pub fn with_tcp_nodelay(mut self) -> Self { + self.tcp_nodelay = TcpNodelay::Enabled; + self + } + /// Connects a WebSocket after resolving the request destination through the configured proxy /// policy. pub async fn connect( @@ -55,9 +95,17 @@ impl WebSocketConnector { ) -> Result<(WebSocketConnection, Response), WebSocketError> { let proxy_route = self .http_client_factory - .resolve_proxy_route(&request.uri().to_string()); - let (inner, response) = - dialer::connect(request, config, Arc::clone(&self.tls_config), proxy_route).await?; + .resolve_proxy_route_async(request.uri().to_string()) + .await + .map_err(WebSocketError::Io)?; + let (inner, response) = dialer::connect( + request, + config, + self.tls_config.clone(), + proxy_route, + self.tcp_nodelay, + ) + .await?; Ok((WebSocketConnection { inner }, response)) } }