core: preserve Responses WebSockets with system proxy

This commit is contained in:
Michael Bolin
2026-07-07 15:37:58 -07:00
parent 8139255b46
commit 487d1b65b5
9 changed files with 472 additions and 20 deletions
+3
View File
@@ -1944,13 +1944,16 @@ dependencies = [
"futures",
"http 1.4.0",
"pretty_assertions",
"rcgen",
"regex-lite",
"reqwest 0.12.28",
"rustls",
"schemars 0.8.22",
"serde",
"serde_json",
"thiserror 2.0.18",
"tokio",
"tokio-rustls",
"tokio-test",
"tokio-tungstenite",
"tokio-util",
+1
View File
@@ -425,6 +425,7 @@ thiserror = "2.0.17"
time = "0.3.47"
tiny_http = "0.12"
tokio = "1"
tokio-rustls = "0.26.4"
tokio-stream = "0.1.18"
tokio-test = "0.4"
tokio-tungstenite = { version = "0.28.0", features = [
+2
View File
@@ -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,
+3
View File
@@ -16,11 +16,13 @@ codex-utils-rustls-provider = { workspace = true }
futures = { workspace = true }
http = { workspace = true }
reqwest = { workspace = true, features = ["json", "stream"] }
rustls = { workspace = true }
schemars = { workspace = true }
serde = { workspace = true, features = ["derive"] }
serde_json = { workspace = true }
thiserror = { workspace = true }
tokio = { workspace = true, features = ["fs", "macros", "net", "rt", "sync", "time"] }
tokio-rustls = { workspace = true }
tokio-tungstenite = { workspace = true }
tungstenite = { workspace = true }
tracing = { workspace = true }
@@ -33,6 +35,7 @@ url = { workspace = true }
anyhow = { workspace = true }
assert_matches = { workspace = true }
pretty_assertions = { workspace = true }
rcgen = { workspace = true }
tokio-test = { workspace = true }
wiremock = { workspace = true }
reqwest = { workspace = true }
@@ -10,8 +10,11 @@ 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 crate::websocket_connector::ConnectedWebSocket;
use codex_client::TransportError;
use codex_http_client::maybe_build_rustls_client_config_with_custom_ca;
use codex_http_client::HttpClientFactory;
use codex_http_client::build_rustls_client_config_with_custom_ca;
use codex_utils_rustls_provider::ensure_rustls_crypto_provider;
use futures::SinkExt;
use futures::StreamExt;
@@ -25,14 +28,13 @@ use serde_json::map::Map as JsonMap;
use std::sync::Arc;
use std::sync::OnceLock;
use std::time::Duration;
use tokio::net::TcpStream;
use tokio::io::AsyncRead;
use tokio::io::AsyncWrite;
use tokio::sync::Mutex;
use tokio::sync::mpsc;
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;
@@ -62,7 +64,17 @@ enum WsCommand {
}
impl WsStream {
fn new(inner: WebSocketStream<MaybeTlsStream<TcpStream>>) -> Self {
fn new(inner: ConnectedWebSocket) -> Self {
match inner {
ConnectedWebSocket::TransportDefault(inner) => Self::new_inner(inner),
ConnectedWebSocket::Routed(inner) => Self::new_inner(inner),
}
}
fn new_inner<S>(inner: WebSocketStream<S>) -> Self
where
S: AsyncRead + AsyncWrite + Send + Unpin + 'static,
{
let (tx_command, mut rx_command) = mpsc::channel::<WsCommand>(32);
let (tx_message, rx_message) = mpsc::unbounded_channel::<Result<Message, WsError>>();
@@ -334,6 +346,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 +362,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 +382,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 +397,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 +457,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();
@@ -448,20 +469,15 @@ async fn connect_websocket(
.map_err(|err| ApiError::Stream(format!("failed to build websocket request: {err}")))?;
request.headers_mut().extend(headers);
// Secure websocket traffic needs the same custom-CA policy as reqwest-based HTTPS traffic.
// If a Codex-specific CA bundle is configured, build an explicit rustls connector so this
// websocket path does not fall back to tungstenite's default native-roots-only behavior.
let connector = maybe_build_rustls_client_config_with_custom_ca()
.map_err(|err| ApiError::Stream(format!("failed to configure websocket TLS: {err}")))?
.map(tokio_tungstenite::Connector::Rustls);
// Secure websocket traffic and TLS-encrypted proxy connections need the same custom-CA policy
// as reqwest-based HTTPS traffic, so build one rustls config for both handshakes.
let tls_config = build_rustls_client_config_with_custom_ca()
.map_err(|err| ApiError::Stream(format!("failed to configure websocket TLS: {err}")))?;
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()), tls_config, proxy_route)
.await;
let (stream, response) = match response {
Ok((stream, response)) => {
+1
View File
@@ -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;
@@ -0,0 +1,251 @@
//! Route-aware WebSocket connection setup.
use std::collections::VecDeque;
use std::io;
use std::net::SocketAddr;
use std::sync::Arc;
use std::time::Duration;
use codex_http_client::OutboundProxyRoute;
use futures::StreamExt;
use futures::stream::FuturesUnordered;
use rustls::ClientConfig;
use rustls::pki_types::ServerName;
use tokio::io::AsyncRead;
use tokio::io::AsyncWrite;
use tokio::net::TcpStream;
use tokio::time::Instant;
use tokio::time::sleep_until;
use tokio_rustls::TlsConnector;
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::TlsError;
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;
const HAPPY_EYEBALLS_DELAY: Duration = Duration::from_millis(250);
/// An established WebSocket whose concrete network stream depends on route selection.
pub(crate) enum ConnectedWebSocket {
/// A connection established by `tokio-tungstenite`'s default transport.
TransportDefault(WebSocketStream<MaybeTlsStream<TcpStream>>),
/// A connection established over a caller-selected direct or proxy stream.
Routed(WebSocketStream<MaybeTlsStream<Box<dyn AsyncIo>>>),
}
/// Async network I/O that can be carried through a proxy tunnel and target TLS handshake.
pub(crate) trait AsyncIo: AsyncRead + AsyncWrite + Send + Unpin {}
impl<T> AsyncIo for T where T: AsyncRead + AsyncWrite + Send + Unpin {}
struct ProxyEndpoint {
config: ProxyConfig,
tls: bool,
}
impl ProxyEndpoint {
fn parse(url: &str) -> Result<Self, WsError> {
let mut parsed_url = url::Url::parse(url).map_err(|_| invalid_proxy_config())?;
let tls = parsed_url.scheme() == "https";
if tls {
if parsed_url.port().is_none() {
parsed_url
.set_port(Some(443))
.map_err(|_| invalid_proxy_config())?;
}
parsed_url
.set_scheme("http")
.map_err(|_| invalid_proxy_config())?;
}
let config = ProxyConfig::parse(parsed_url.as_str()).map_err(|error| match error {
WsError::Url(UrlError::UnsupportedProxyScheme) => error,
_ => invalid_proxy_config(),
})?;
Ok(Self { config, tls })
}
}
/// Connects a WebSocket using the resolved outbound proxy route.
pub(crate) async fn connect(
request: Request,
config: Option<WebSocketConfig>,
tls_config: Arc<ClientConfig>,
proxy_route: OutboundProxyRoute,
) -> Result<(ConnectedWebSocket, Response), WsError> {
let stream: Box<dyn AsyncIo> = match proxy_route {
OutboundProxyRoute::TransportDefault => {
let (stream, response) = connect_async_tls_with_config(
request,
config,
false, // Preserve tungstenite's recommended Nagle default.
Some(Connector::Rustls(tls_config)),
)
.await?;
return Ok((ConnectedWebSocket::TransportDefault(stream), response));
}
OutboundProxyRoute::Direct => {
let host = request
.uri()
.host()
.ok_or(WsError::Url(UrlError::NoHostName))?;
let port = websocket_port(&request)?;
let address = host_port(host, port);
Box::new(connect_tcp(address).await.map_err(WsError::Io)?)
}
OutboundProxyRoute::Proxy { url } => {
let proxy = ProxyEndpoint::parse(&url)?;
let host = request
.uri()
.host()
.ok_or(WsError::Url(UrlError::NoHostName))?;
let port = websocket_port(&request)?;
let stream = connect_tcp(proxy.config.authority())
.await
.map_err(WsError::Io)?;
let stream: Box<dyn AsyncIo> = if proxy.tls {
let server_name = ServerName::try_from(proxy.config.host.clone())
.map_err(|_| WsError::Tls(TlsError::InvalidDnsName))?;
let stream = TlsConnector::from(Arc::clone(&tls_config))
.connect(server_name, stream)
.await
.map_err(WsError::Io)?;
Box::new(stream)
} else {
Box::new(stream)
};
connect_via_proxy(stream, &proxy.config, host, port).await?
}
};
let (stream, response) =
client_async_tls_with_config(request, stream, config, Some(Connector::Rustls(tls_config)))
.await?;
Ok((ConnectedWebSocket::Routed(stream), response))
}
fn invalid_proxy_config() -> WsError {
WsError::Url(UrlError::InvalidProxyConfig("<redacted>".to_string()))
}
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}")
}
}
/// Preserves the Happy Eyeballs behavior of `tokio-tungstenite`'s built-in dialer.
///
/// Explicit routes require a caller-provided TCP stream, so they cannot use that private dialer
/// directly. Keep the same family interleaving and 250 ms attempt delay here to avoid regressing
/// direct or proxy connections when one address family is unreachable.
async fn connect_tcp(address: String) -> io::Result<TcpStream> {
let mut resolved_addresses = tokio::net::lookup_host(address).await?;
let Some(first_address) = resolved_addresses.next() else {
return Err(io::Error::new(
io::ErrorKind::InvalidInput,
"could not resolve to any address",
));
};
let first_is_ipv4 = first_address.is_ipv4();
let mut preferred = VecDeque::new();
let mut alternate = VecDeque::new();
for address in resolved_addresses {
if address.is_ipv4() == first_is_ipv4 {
preferred.push_back(address);
} else {
alternate.push_back(address);
}
}
let mut addresses = VecDeque::<SocketAddr>::new();
while !preferred.is_empty() || !alternate.is_empty() {
if let Some(address) = alternate.pop_front() {
addresses.push_back(address);
}
if let Some(address) = preferred.pop_front() {
addresses.push_back(address);
}
}
let mut attempts = FuturesUnordered::new();
attempts.push(TcpStream::connect(first_address));
let mut next_attempt_at = Instant::now() + HAPPY_EYEBALLS_DELAY;
let mut last_error = None;
loop {
if addresses.is_empty() {
match attempts.next().await {
Some(Ok(stream)) => return Ok(stream),
Some(Err(error)) => {
if attempts.is_empty() {
return Err(error);
}
last_error = Some(error);
}
None => {
return Err(last_error.unwrap_or_else(|| {
io::Error::other("connection attempts ended without an error")
}));
}
}
continue;
}
tokio::select! {
result = attempts.next() => {
match result {
Some(Ok(stream)) => return Ok(stream),
Some(Err(error)) => {
last_error = Some(error);
let address = take_next_address(&mut addresses)?;
attempts.push(TcpStream::connect(address));
next_attempt_at = Instant::now() + HAPPY_EYEBALLS_DELAY;
}
None => {
let address = take_next_address(&mut addresses)?;
attempts.push(TcpStream::connect(address));
next_attempt_at = Instant::now() + HAPPY_EYEBALLS_DELAY;
}
}
}
_ = sleep_until(next_attempt_at) => {
let address = take_next_address(&mut addresses)?;
attempts.push(TcpStream::connect(address));
next_attempt_at = Instant::now() + HAPPY_EYEBALLS_DELAY;
}
}
}
}
fn take_next_address(addresses: &mut VecDeque<SocketAddr>) -> io::Result<SocketAddr> {
addresses
.pop_front()
.ok_or_else(|| io::Error::other("connection address queue unexpectedly empty"))
}
#[cfg(test)]
#[path = "websocket_connector_tests.rs"]
mod tests;
@@ -0,0 +1,174 @@
use std::net::SocketAddr;
use std::sync::Arc;
use codex_utils_rustls_provider::ensure_rustls_crypto_provider;
use pretty_assertions::assert_eq;
use rcgen::CertifiedKey;
use rcgen::generate_simple_self_signed;
use rustls::ClientConfig;
use rustls::RootCertStore;
use rustls::ServerConfig;
use rustls::pki_types::PrivateKeyDer;
use rustls::pki_types::PrivatePkcs8KeyDer;
use tokio::io::AsyncReadExt;
use tokio::io::AsyncWriteExt;
use tokio::net::TcpListener;
use tokio::sync::Mutex;
use tokio::task::JoinHandle;
use tokio_rustls::TlsAcceptor;
use tokio_tungstenite::accept_async;
use tokio_tungstenite::tungstenite::client::IntoClientRequest;
use super::*;
#[tokio::test]
async fn direct_route_connects_secure_websocket() {
let (tls_config, acceptor) = test_tls_configs();
let (target_addr, target_task) = start_tls_websocket_server(acceptor).await;
let request = format!("wss://localhost:{}/v1/responses", target_addr.port())
.into_client_request()
.expect("websocket request should build");
let (websocket, _) = connect(
request,
/*config*/ None,
tls_config,
OutboundProxyRoute::Direct,
)
.await
.expect("direct websocket handshake should succeed");
drop(websocket);
target_task.await.expect("target task should finish");
}
#[tokio::test]
async fn http_proxy_tunnels_secure_websocket_before_handshake() {
assert_proxy_tunnels_secure_websocket(/*proxy_tls*/ false).await;
}
#[tokio::test]
async fn https_proxy_tunnels_secure_websocket_before_handshake() {
assert_proxy_tunnels_secure_websocket(/*proxy_tls*/ true).await;
}
async fn assert_proxy_tunnels_secure_websocket(proxy_tls: bool) {
let (tls_config, acceptor) = test_tls_configs();
let (target_addr, target_task) = start_tls_websocket_server(acceptor.clone()).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 (client, _) = proxy_listener.accept().await.expect("proxy should accept");
let mut client: Box<dyn AsyncIo> = if proxy_tls {
Box::new(
acceptor
.accept(client)
.await
.expect("proxy TLS handshake should succeed"),
)
} else {
Box::new(client)
};
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 target_authority = format!("localhost:{}", target_addr.port());
let proxy_scheme = if proxy_tls { "https" } else { "http" };
let request = format!("wss://{target_authority}/v1/responses")
.into_client_request()
.expect("websocket request should build");
let (websocket, _) = connect(
request,
/*config*/ None,
tls_config,
OutboundProxyRoute::Proxy {
url: format!("{proxy_scheme}://localhost:{}", proxy_addr.port()),
},
)
.await
.expect("proxied websocket handshake should succeed");
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_authority} HTTP/1.1");
assert_eq!(request.lines().next(), Some(expected_request_line.as_str()));
}
async fn start_tls_websocket_server(acceptor: TlsAcceptor) -> (SocketAddr, JoinHandle<()>) {
let listener = TcpListener::bind("127.0.0.1:0")
.await
.expect("target listener should bind");
let address = listener
.local_addr()
.expect("target listener should have an address");
let task = tokio::spawn(async move {
let (stream, _) = listener.accept().await.expect("target should accept");
let stream = acceptor
.accept(stream)
.await
.expect("target TLS handshake should succeed");
let mut websocket = accept_async(stream)
.await
.expect("target websocket handshake should succeed");
let _ = websocket.close(None).await;
});
(address, task)
}
fn test_tls_configs() -> (Arc<ClientConfig>, TlsAcceptor) {
ensure_rustls_crypto_provider();
let CertifiedKey { cert, signing_key } =
generate_simple_self_signed(vec!["localhost".to_string()])
.expect("test certificate should generate");
let certificate = cert.der().clone();
let private_key = PrivateKeyDer::Pkcs8(PrivatePkcs8KeyDer::from(signing_key.serialize_der()));
let server_config = ServerConfig::builder()
.with_no_client_auth()
.with_single_cert(vec![certificate.clone()], private_key)
.expect("test server config should build");
let mut roots = RootCertStore::empty();
roots
.add(certificate)
.expect("test certificate should be trusted");
let client_config = ClientConfig::builder()
.with_root_certificates(roots)
.with_no_client_auth();
(
Arc::new(client_config),
TlsAcceptor::from(Arc::new(server_config)),
)
}
+1
View File
@@ -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,