core: preserve Responses WebSockets with system proxy

codex-rs/codex-api/src/endpoint/responses_websocket.rs

@@ -10,7 +10,9 @@ use crate::safety_buffering::treatment_from_headers;
 use crate::sse::ResponsesStreamEvent;
 use crate::sse::process_responses_event;
 use crate::telemetry::WebsocketTelemetry;
+use crate::websocket_connector;
 use codex_client::TransportError;
+use codex_http_client::HttpClientFactory;
 use codex_http_client::maybe_build_rustls_client_config_with_custom_ca;
 use codex_utils_rustls_provider::ensure_rustls_crypto_provider;
 use futures::SinkExt;
@@ -32,7 +34,6 @@ use tokio::sync::oneshot;
 use tokio::time::Instant;
 use tokio_tungstenite::MaybeTlsStream;
 use tokio_tungstenite::WebSocketStream;
-use tokio_tungstenite::connect_async_tls_with_config;
 use tokio_tungstenite::tungstenite::Error as WsError;
 use tokio_tungstenite::tungstenite::Message;
 use tokio_tungstenite::tungstenite::client::IntoClientRequest;
@@ -334,6 +335,7 @@ impl ResponsesWebsocketClient {
     )]
     pub async fn connect(
         &self,
+        http_client_factory: &HttpClientFactory,
         extra_headers: HeaderMap,
         default_headers: HeaderMap,
         turn_state: Option<Arc<OnceLock<String>>>,
@@ -349,7 +351,7 @@ impl ResponsesWebsocketClient {
         self.auth.add_auth_headers(&mut headers);
 
         let (stream, _status, server_reasoning_included, models_etag, server_model) =
-            connect_websocket(ws_url, headers, turn_state.clone()).await?;
+            connect_websocket(ws_url, headers, http_client_factory, turn_state.clone()).await?;
         Ok(ResponsesWebsocketConnection::new(
             stream,
             self.provider.stream_idle_timeout,
@@ -369,6 +371,7 @@ impl ResponsesWebsocketClient {
     /// a usable connection from a policy rejection that closes right away.
     pub async fn probe_handshake(
         &self,
+        http_client_factory: &HttpClientFactory,
         extra_headers: HeaderMap,
         default_headers: HeaderMap,
         immediate_close_timeout: Duration,
@@ -383,7 +386,13 @@ impl ResponsesWebsocketClient {
         self.auth.add_auth_headers(&mut headers);
 
         let (mut stream, status, reasoning_included, models_etag, server_model) =
-            connect_websocket(ws_url.clone(), headers, /*turn_state*/ None).await?;
+            connect_websocket(
+                ws_url.clone(),
+                headers,
+                http_client_factory,
+                /*turn_state*/ None,
+            )
+            .await?;
         let immediate_close = tokio::time::timeout(immediate_close_timeout, stream.next())
             .await
             .ok()
@@ -437,6 +446,7 @@ fn merge_request_headers(
 async fn connect_websocket(
     url: Url,
     headers: HeaderMap,
+    http_client_factory: &HttpClientFactory,
     turn_state: Option<Arc<OnceLock<String>>>,
 ) -> Result<(WsStream, StatusCode, bool, Option<String>, Option<String>), ApiError> {
     ensure_rustls_crypto_provider();
@@ -455,13 +465,10 @@ async fn connect_websocket(
         .map_err(|err| ApiError::Stream(format!("failed to configure websocket TLS: {err}")))?
         .map(tokio_tungstenite::Connector::Rustls);
 
-    let response = connect_async_tls_with_config(
-        request,
-        Some(websocket_config()),
-        false, // `false` means "do not disable Nagle", which is tungstenite's recommended default.
-        connector,
-    )
-    .await;
+    let proxy_route = http_client_factory.resolve_proxy_route(url.as_str());
+    let response =
+        websocket_connector::connect(request, Some(websocket_config()), connector, proxy_route)
+            .await;
 
     let (stream, response) = match response {
         Ok((stream, response)) => {

codex-rs/codex-api/src/lib.rs

@@ -12,6 +12,7 @@ pub(crate) mod safety_buffering;
 pub(crate) mod search;
 pub(crate) mod sse;
 pub(crate) mod telemetry;
+mod websocket_connector;
 
 pub use crate::requests::headers::build_session_headers;
 pub use codex_client::RequestTelemetry;

codex-rs/codex-api/src/websocket_connector.rs

@@ -0,0 +1,83 @@
+//! Route-aware WebSocket connection setup.
+
+use codex_http_client::OutboundProxyRoute;
+use tokio::net::TcpStream;
+use tokio_tungstenite::Connector;
+use tokio_tungstenite::MaybeTlsStream;
+use tokio_tungstenite::WebSocketStream;
+use tokio_tungstenite::client_async_tls_with_config;
+use tokio_tungstenite::connect_async_tls_with_config;
+use tokio_tungstenite::proxy::connect_via_proxy;
+use tokio_tungstenite::tungstenite::Error as WsError;
+use tokio_tungstenite::tungstenite::error::UrlError;
+use tokio_tungstenite::tungstenite::handshake::client::Request;
+use tokio_tungstenite::tungstenite::handshake::client::Response;
+use tokio_tungstenite::tungstenite::protocol::WebSocketConfig;
+use tokio_tungstenite::tungstenite::proxy::ProxyConfig;
+
+/// Connects a WebSocket using the resolved outbound proxy route.
+pub(crate) async fn connect(
+    request: Request,
+    config: Option<WebSocketConfig>,
+    connector: Option<Connector>,
+    proxy_route: OutboundProxyRoute,
+) -> Result<(WebSocketStream<MaybeTlsStream<TcpStream>>, Response), WsError> {
+    match proxy_route {
+        OutboundProxyRoute::TransportDefault => {
+            connect_async_tls_with_config(
+                request, config, false, // Preserve tungstenite's recommended Nagle default.
+                connector,
+            )
+            .await
+        }
+        OutboundProxyRoute::Direct => {
+            let host = request
+                .uri()
+                .host()
+                .ok_or(WsError::Url(UrlError::NoHostName))?;
+            let port = websocket_port(&request)?;
+            let address = host_port(host, port);
+            let stream = TcpStream::connect(address).await.map_err(WsError::Io)?;
+            client_async_tls_with_config(request, stream, config, connector).await
+        }
+        OutboundProxyRoute::Proxy { url } => {
+            let proxy = ProxyConfig::parse(&url).map_err(|_| {
+                WsError::Url(UrlError::InvalidProxyConfig("<redacted>".to_string()))
+            })?;
+            let host = request
+                .uri()
+                .host()
+                .ok_or(WsError::Url(UrlError::NoHostName))?;
+            let port = websocket_port(&request)?;
+            let stream = TcpStream::connect(proxy.authority())
+                .await
+                .map_err(WsError::Io)?;
+            let stream = connect_via_proxy(stream, &proxy, host, port).await?;
+            client_async_tls_with_config(request, stream, config, connector).await
+        }
+    }
+}
+
+fn websocket_port(request: &Request) -> Result<u16, WsError> {
+    request
+        .uri()
+        .port_u16()
+        .or_else(|| match request.uri().scheme_str() {
+            Some("ws") => Some(80),
+            Some("wss") => Some(443),
+            _ => None,
+        })
+        .ok_or(WsError::Url(UrlError::UnsupportedUrlScheme))
+}
+
+fn host_port(host: &str, port: u16) -> String {
+    if host.contains(':') && !host.starts_with('[') {
+        format!("[{host}]:{port}")
+    } else {
+        format!("{host}:{port}")
+    }
+}
+
+#[cfg(test)]
+#[path = "websocket_connector_tests.rs"]
+mod tests;

codex-rs/codex-api/src/websocket_connector_tests.rs

@@ -0,0 +1,89 @@
+use std::sync::Arc;
+
+use pretty_assertions::assert_eq;
+use tokio::io::AsyncReadExt;
+use tokio::io::AsyncWriteExt;
+use tokio::net::TcpListener;
+use tokio::sync::Mutex;
+use tokio_tungstenite::accept_async;
+use tokio_tungstenite::tungstenite::client::IntoClientRequest;
+
+use super::*;
+
+#[tokio::test]
+async fn proxy_route_establishes_connect_tunnel_before_websocket_handshake() {
+    let target_listener = TcpListener::bind("127.0.0.1:0")
+        .await
+        .expect("target listener should bind");
+    let target_addr = target_listener
+        .local_addr()
+        .expect("target listener should have an address");
+    let target_task = tokio::spawn(async move {
+        let (stream, _) = target_listener
+            .accept()
+            .await
+            .expect("target should accept");
+        let mut websocket = accept_async(stream)
+            .await
+            .expect("target websocket handshake should succeed");
+        let _ = websocket.close(None).await;
+    });
+
+    let proxy_listener = TcpListener::bind("127.0.0.1:0")
+        .await
+        .expect("proxy listener should bind");
+    let proxy_addr = proxy_listener
+        .local_addr()
+        .expect("proxy listener should have an address");
+    let connect_request = Arc::new(Mutex::new(None));
+    let proxy_connect_request = Arc::clone(&connect_request);
+    let proxy_task = tokio::spawn(async move {
+        let (mut client, _) = proxy_listener.accept().await.expect("proxy should accept");
+        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
+                .expect("proxy should read CONNECT request");
+            request.push(byte[0]);
+        }
+        *proxy_connect_request.lock().await =
+            Some(String::from_utf8(request).expect("CONNECT request should contain valid UTF-8"));
+
+        let mut target = tokio::net::TcpStream::connect(target_addr)
+            .await
+            .expect("proxy should connect to target");
+        client
+            .write_all(b"HTTP/1.1 200 Connection Established\r\n\r\n")
+            .await
+            .expect("proxy should acknowledge CONNECT");
+        let _ = tokio::io::copy_bidirectional(&mut client, &mut target).await;
+    });
+
+    let request = format!("ws://{target_addr}/v1/responses")
+        .into_client_request()
+        .expect("websocket request should build");
+    let (mut websocket, _) = connect(
+        request,
+        /*config*/ None,
+        /*connector*/ None,
+        OutboundProxyRoute::Proxy {
+            url: format!("http://{proxy_addr}"),
+        },
+    )
+    .await
+    .expect("proxied websocket handshake should succeed");
+    let _ = websocket.close(None).await;
+    drop(websocket);
+
+    target_task.await.expect("target task should finish");
+    proxy_task.await.expect("proxy task should finish");
+    let request = connect_request
+        .lock()
+        .await
+        .clone()
+        .expect("proxy should record CONNECT request");
+    let expected_request_line = format!("CONNECT {target_addr} HTTP/1.1");
+    assert_eq!(request.lines().next(), Some(expected_request_line.as_str()));
+}