diff --git a/codex-rs/cli/src/doctor.rs b/codex-rs/cli/src/doctor.rs index ef587cd1d4..f695274675 100644 --- a/codex-rs/cli/src/doctor.rs +++ b/codex-rs/cli/src/doctor.rs @@ -2371,9 +2371,11 @@ async fn websocket_reachability_check( HeaderValue::from_static(RESPONSES_WEBSOCKETS_V2_BETA_HEADER_VALUE), ); let client = ResponsesWebsocketClient::new(api_provider, api_auth); + let http_client_factory = config.http_client_factory(); match tokio::time::timeout( provider.websocket_connect_timeout(), client.probe_handshake( + &http_client_factory, extra_headers, default_headers(), WEBSOCKET_IMMEDIATE_CLOSE_GRACE, diff --git a/codex-rs/codex-api/src/endpoint/responses_websocket.rs b/codex-rs/codex-api/src/endpoint/responses_websocket.rs index cce056f638..71457d6fa7 100644 --- a/codex-rs/codex-api/src/endpoint/responses_websocket.rs +++ b/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>>, @@ -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>>, ) -> Result<(WsStream, StatusCode, bool, Option, Option), 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)) => { diff --git a/codex-rs/codex-api/src/lib.rs b/codex-rs/codex-api/src/lib.rs index 1d1a4c8c6e..4a5c63391e 100644 --- a/codex-rs/codex-api/src/lib.rs +++ b/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; diff --git a/codex-rs/codex-api/src/websocket_connector.rs b/codex-rs/codex-api/src/websocket_connector.rs new file mode 100644 index 0000000000..abd1509b3b --- /dev/null +++ b/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, + connector: Option, + proxy_route: OutboundProxyRoute, +) -> Result<(WebSocketStream>, 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("".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 { + 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; diff --git a/codex-rs/codex-api/src/websocket_connector_tests.rs b/codex-rs/codex-api/src/websocket_connector_tests.rs new file mode 100644 index 0000000000..88d6196e47 --- /dev/null +++ b/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())); +} diff --git a/codex-rs/core/src/client.rs b/codex-rs/core/src/client.rs index 0bf33b3993..5d29979c37 100644 --- a/codex-rs/core/src/client.rs +++ b/codex-rs/core/src/client.rs @@ -1006,6 +1006,7 @@ impl ModelClient { let result = match tokio::time::timeout( websocket_connect_timeout, ApiWebSocketResponsesClient::new(api_provider, api_auth).connect( + &self.http_client_factory, headers, codex_login::default_client::default_headers(), /*turn_state*/ None, diff --git a/codex-rs/core/tests/suite/client_websockets.rs b/codex-rs/core/tests/suite/client_websockets.rs index 05a93807b2..f00a50fa09 100755 --- a/codex-rs/core/tests/suite/client_websockets.rs +++ b/codex-rs/core/tests/suite/client_websockets.rs @@ -230,6 +230,34 @@ async fn responses_websocket_streams_without_feature_flag_when_provider_supports server.shutdown().await; } +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn responses_websocket_streams_with_system_proxy_feature() { + skip_if_no_network!(); + + let server = start_websocket_server(vec![vec![vec![ + ev_response_created("resp-1"), + ev_completed("resp-1"), + ]]]) + .await; + + let harness = websocket_harness_with_provider_options( + websocket_provider(&server), + /*runtime_metrics_enabled*/ false, + /*concurrent_reasoning_summaries_enabled*/ false, + /*enabled_features*/ &[Feature::RespectSystemProxy], + ) + .await; + let mut client_session = harness.client.new_session(); + let prompt = prompt_with_input(vec![message_item("hello")]); + + stream_until_complete(&mut client_session, &harness, &prompt).await; + + assert_eq!(server.handshakes().len(), 1); + assert_eq!(server.single_connection().len(), 1); + + server.shutdown().await; +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn responses_websocket_reuses_connection_with_per_turn_trace_payloads() { skip_if_no_network!(); @@ -389,8 +417,10 @@ async fn responses_websocket_request_prewarm_reuses_connection() { let mut provider = websocket_provider(&server); provider.name = ModelProviderInfo::create_openai_provider(/*base_url*/ None).name; let harness = websocket_harness_with_provider_options( - provider, /*runtime_metrics_enabled*/ true, + provider, + /*runtime_metrics_enabled*/ true, /*concurrent_reasoning_summaries_enabled*/ true, + /*enabled_features*/ &[], ) .await; let mut client_session = harness.client.new_session(); @@ -1005,8 +1035,10 @@ async fn responses_websocket_v2_incremental_requests_are_reused_across_turns() { let mut provider = websocket_provider(&server); provider.name = ModelProviderInfo::create_openai_provider(/*base_url*/ None).name; let harness = websocket_harness_with_provider_options( - provider, /*runtime_metrics_enabled*/ false, + provider, + /*runtime_metrics_enabled*/ false, /*concurrent_reasoning_summaries_enabled*/ false, + /*enabled_features*/ &[], ) .await; @@ -2216,6 +2248,7 @@ async fn websocket_harness_with_options( websocket_provider(server), runtime_metrics_enabled, /*concurrent_reasoning_summaries_enabled*/ false, + /*enabled_features*/ &[], ) .await } @@ -2224,6 +2257,7 @@ async fn websocket_harness_with_provider_options( provider: ModelProviderInfo, runtime_metrics_enabled: bool, concurrent_reasoning_summaries_enabled: bool, + enabled_features: &[Feature], ) -> WebsocketTestHarness { let codex_home = TempDir::new().unwrap(); let mut config = load_default_config_for_test(&codex_home).await; @@ -2240,6 +2274,12 @@ async fn websocket_harness_with_provider_options( .enable(Feature::ConcurrentReasoningSummaries) .expect("test config should allow feature update"); } + for feature in enabled_features { + config + .features + .enable(*feature) + .expect("test config should allow feature update"); + } let config = Arc::new(config); let model_info = codex_core::test_support::construct_model_info_offline(MODEL, &config); let thread_id = ThreadId::new(); diff --git a/codex-rs/http-client/src/lib.rs b/codex-rs/http-client/src/lib.rs index 07e579ecb3..2512948687 100644 --- a/codex-rs/http-client/src/lib.rs +++ b/codex-rs/http-client/src/lib.rs @@ -27,6 +27,7 @@ pub use crate::outbound_proxy::BuildRouteAwareHttpClientError; pub use crate::outbound_proxy::ClientRouteClass; pub use crate::outbound_proxy::HttpClientFactory; pub use crate::outbound_proxy::OutboundProxyPolicy; +pub use crate::outbound_proxy::OutboundProxyRoute; pub use crate::outbound_proxy::RouteFailureClass; pub use crate::request::EncodedJsonBody; pub use crate::request::PreparedRequestBody; diff --git a/codex-rs/http-client/src/outbound_proxy.rs b/codex-rs/http-client/src/outbound_proxy.rs index 5506f3d40b..71793025ed 100644 --- a/codex-rs/http-client/src/outbound_proxy.rs +++ b/codex-rs/http-client/src/outbound_proxy.rs @@ -4,6 +4,7 @@ //! proxies are the fallback, and the final fallback is a direct connection. //! When disabled, callers retain the existing reqwest builder behavior. +use std::borrow::Cow; use std::collections::HashMap; use std::fmt; use std::io; @@ -96,6 +97,30 @@ pub enum OutboundProxyPolicy { RespectSystemProxy, } +/// Resolved proxy route for a concrete outbound destination. +/// +/// `TransportDefault` delegates environment-proxy handling to the underlying transport. Proxy +/// URLs are intentionally redacted from `Debug` output because they may contain credentials. +#[derive(Clone, PartialEq, Eq)] +pub enum OutboundProxyRoute { + /// Preserve the underlying transport's existing proxy behavior. + TransportDefault, + /// Connect directly and bypass transport-level proxy discovery. + Direct, + /// Connect through the selected proxy URL. + Proxy { url: String }, +} + +impl fmt::Debug for OutboundProxyRoute { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + match self { + Self::TransportDefault => f.write_str("TransportDefault"), + Self::Direct => f.write_str("Direct"), + Self::Proxy { .. } => f.debug_struct("Proxy").field("url", &"").finish(), + } + } +} + /// Builds route-specific HTTP clients using one resolved outbound proxy policy. /// /// Construct this once from the effective application configuration and carry it with the @@ -119,6 +144,20 @@ impl HttpClientFactory { self.outbound_proxy_policy } + /// Resolves the proxy route for a concrete destination. + /// + /// WebSocket schemes are resolved through their HTTP equivalents so platform PAC and system + /// proxy APIs apply the same policy to `ws`/`wss` and `http`/`https` destinations. When system + /// resolution is unavailable, the transport retains responsibility for environment-proxy + /// fallback. + pub fn resolve_proxy_route(&self, request_url: &str) -> OutboundProxyRoute { + resolve_proxy_route( + request_url, + self.outbound_proxy_policy, + resolve_system_proxy, + ) + } + /// Builds a reqwest client for a concrete outbound route. pub fn build_reqwest_client( &self, @@ -135,6 +174,37 @@ impl HttpClientFactory { } } +fn resolve_proxy_route( + request_url: &str, + outbound_proxy_policy: OutboundProxyPolicy, + resolve_system_proxy: impl FnOnce(&str, &RequestOrigin) -> SystemProxyDecision, +) -> OutboundProxyRoute { + if matches!(outbound_proxy_policy, OutboundProxyPolicy::ReqwestDefault) { + return OutboundProxyRoute::TransportDefault; + } + + let request_url = proxy_resolution_url(request_url); + let Some(origin) = RequestOrigin::parse(&request_url) else { + return OutboundProxyRoute::TransportDefault; + }; + + match resolve_system_proxy(&request_url, &origin) { + SystemProxyDecision::Direct => OutboundProxyRoute::Direct, + SystemProxyDecision::Proxy { url } => OutboundProxyRoute::Proxy { url }, + SystemProxyDecision::Unavailable { .. } => OutboundProxyRoute::TransportDefault, + } +} + +fn proxy_resolution_url(request_url: &str) -> Cow<'_, str> { + if let Some(suffix) = request_url.strip_prefix("wss://") { + Cow::Owned(format!("https://{suffix}")) + } else if let Some(suffix) = request_url.strip_prefix("ws://") { + Cow::Owned(format!("http://{suffix}")) + } else { + Cow::Borrowed(request_url) + } +} + /// Error while building a resolver-aware reqwest client. #[derive(Debug, Error)] pub enum BuildRouteAwareHttpClientError { @@ -259,8 +329,8 @@ impl RequestOrigin { let scheme = uri.scheme_str()?.to_ascii_lowercase(); let host = uri.host()?.trim_matches(['[', ']']).to_ascii_lowercase(); let port = uri.port_u16().or(match scheme.as_str() { - "http" => Some(80), - "https" => Some(443), + "http" | "ws" => Some(80), + "https" | "wss" => Some(443), _ => None, })?; Some(Self { scheme, host, port }) diff --git a/codex-rs/http-client/src/outbound_proxy_tests.rs b/codex-rs/http-client/src/outbound_proxy_tests.rs index 734d5483ec..3f8bd925c9 100644 --- a/codex-rs/http-client/src/outbound_proxy_tests.rs +++ b/codex-rs/http-client/src/outbound_proxy_tests.rs @@ -9,6 +9,41 @@ struct MapEnv { values: HashMap, } +#[test] +fn websocket_route_uses_http_equivalent_for_system_resolution() { + let route = resolve_proxy_route( + "wss://api.openai.com/v1/responses", + OutboundProxyPolicy::RespectSystemProxy, + |request_url, origin| { + assert_eq!(request_url, "https://api.openai.com/v1/responses"); + assert_eq!(origin.scheme, "https"); + assert_eq!(origin.host, "api.openai.com"); + assert_eq!(origin.port, 443); + SystemProxyDecision::Proxy { + url: "http://proxy.example:8080".to_string(), + } + }, + ); + + assert_eq!( + route, + OutboundProxyRoute::Proxy { + url: "http://proxy.example:8080".to_string(), + } + ); +} + +#[test] +fn reqwest_default_route_preserves_transport_proxy_behavior() { + let route = resolve_proxy_route( + "wss://api.openai.com/v1/responses", + OutboundProxyPolicy::ReqwestDefault, + |_, _| panic!("default policy should not resolve system proxy settings"), + ); + + assert_eq!(route, OutboundProxyRoute::TransportDefault); +} + impl EnvSource for MapEnv { fn var(&self, key: &str) -> Option { self.values.get(key).cloned()