mirror of
https://github.com/openai/codex.git
synced 2026-08-23 13:09:46 +00:00
core: preserve Responses WebSockets with system proxy (#31441)
## Why Responses WebSockets are the normal lower-latency transport for WebSocket-capable providers. They must not bypass an OS-selected proxy when `features.respect_system_proxy` is enabled, but disabling WebSockets whenever the feature is enabled would impose a substantial performance penalty. Merged PR #31622 introduced the reusable proxy-aware WebSocket transport. This PR makes the Responses API its first consumer so the existing fast path uses the same effective proxy and trust policy as HTTP. ## What changed - Register `codex-websocket-client` as a workspace dependency and use it from `codex-api`. - Feed the shared crate’s route-independent `WebSocketConnection` into the existing Responses message pump. - Require a configured `HttpClientFactory` for normal Responses WebSocket connections and the CLI doctor probe, so neither path can open a connection without consulting the effective proxy policy. - Pass the session factory from `core` and the effective configuration factory from `doctor`. - Add an end-to-end Responses test that enables `RespectSystemProxy`, asserts the resolved policy, completes a turn over WebSocket, and verifies the connection and request counts. - Keep the existing Responses protocol handling, ping/pong pump, and session-scoped HTTP fallback unchanged. The DNS, proxy, TLS, custom-CA, and Happy Eyeballs implementation and its transport tests live in merged PR #31622. This PR deliberately contains only the Responses integration and does not duplicate that transport code. ## Review guide 1. `codex-rs/codex-api/src/endpoint/responses_websocket.rs` constructs the shared connector and adapts its uniform stream to the existing pump. 2. `codex-rs/core/src/client.rs` supplies the session-scoped factory for production Responses connections. 3. `codex-rs/cli/src/doctor.rs` supplies the effective configuration factory to the handshake probe. 4. `codex-rs/core/tests/suite/client_websockets.rs` covers the enabled-feature path end to end. ## Test plan - `cargo check --tests -p codex-api -p codex-core -p codex-cli` - `just test -p codex-api` - `just test -p codex-core responses_websocket_streams_with_system_proxy_feature` - `cargo shear` - `just bazel-lock-check` --- [//]: # (BEGIN SAPLING FOOTER) Stack created with [Sapling](https://sapling-scm.com). Best reviewed with [ReviewStack](https://reviewstack.dev/openai/codex/pull/31441). * #31637 * #31431 * #31363 * #31362 * #31361 * __->__ #31441
This commit is contained in:
1
codex-rs/Cargo.lock
generated
1
codex-rs/Cargo.lock
generated
@@ -1940,6 +1940,7 @@ dependencies = [
|
||||
"codex-http-client",
|
||||
"codex-protocol",
|
||||
"codex-utils-rustls-provider",
|
||||
"codex-websocket-client",
|
||||
"eventsource-stream",
|
||||
"futures",
|
||||
"http 1.4.0",
|
||||
|
||||
@@ -168,6 +168,7 @@ codex-code-mode = { path = "code-mode" }
|
||||
codex-code-mode-protocol = { path = "code-mode-protocol" }
|
||||
codex-home = { path = "codex-home" }
|
||||
codex-http-client = { path = "http-client" }
|
||||
codex-websocket-client = { path = "websocket-client" }
|
||||
codex-config = { path = "config" }
|
||||
codex-connectors = { path = "connectors" }
|
||||
codex-connectors-extension = { path = "ext/connectors" }
|
||||
|
||||
@@ -2383,9 +2383,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,
|
||||
|
||||
@@ -13,6 +13,7 @@ codex-client = { workspace = true }
|
||||
codex-http-client = { workspace = true }
|
||||
codex-protocol = { workspace = true }
|
||||
codex-utils-rustls-provider = { workspace = true }
|
||||
codex-websocket-client = { workspace = true }
|
||||
futures = { workspace = true }
|
||||
http = { workspace = true }
|
||||
reqwest = { workspace = true, features = ["json", "stream"] }
|
||||
|
||||
@@ -11,8 +11,9 @@ use crate::sse::ResponsesStreamEvent;
|
||||
use crate::sse::process_responses_event;
|
||||
use crate::telemetry::WebsocketTelemetry;
|
||||
use codex_client::TransportError;
|
||||
use codex_http_client::maybe_build_rustls_client_config_with_custom_ca;
|
||||
use codex_utils_rustls_provider::ensure_rustls_crypto_provider;
|
||||
use codex_http_client::HttpClientFactory;
|
||||
use codex_websocket_client::WebSocketConnection;
|
||||
use codex_websocket_client::WebSocketConnector;
|
||||
use futures::SinkExt;
|
||||
use futures::StreamExt;
|
||||
use http::HeaderMap;
|
||||
@@ -25,14 +26,10 @@ 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::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 +59,7 @@ enum WsCommand {
|
||||
}
|
||||
|
||||
impl WsStream {
|
||||
fn new(inner: WebSocketStream<MaybeTlsStream<TcpStream>>) -> Self {
|
||||
fn new(inner: WebSocketConnection) -> Self {
|
||||
let (tx_command, mut rx_command) = mpsc::channel::<WsCommand>(32);
|
||||
let (tx_message, rx_message) = mpsc::unbounded_channel::<Result<Message, WsError>>();
|
||||
|
||||
@@ -334,6 +331,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 +347,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 +367,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 +382,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,9 +442,9 @@ 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();
|
||||
info!("connecting to websocket: {url}");
|
||||
|
||||
let mut request = url
|
||||
@@ -448,20 +453,9 @@ 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);
|
||||
|
||||
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 connector = WebSocketConnector::new(http_client_factory)
|
||||
.map_err(|err| ApiError::Stream(format!("failed to configure websocket TLS: {err}")))?;
|
||||
let response = connector.connect(request, websocket_config()).await;
|
||||
|
||||
let (stream, response) = match response {
|
||||
Ok((stream, response)) => {
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -8,6 +8,7 @@ use codex_core::Prompt;
|
||||
use codex_core::ResponseEvent;
|
||||
use codex_core::X_RESPONSESAPI_INCLUDE_TIMING_METRICS_HEADER;
|
||||
use codex_features::Feature;
|
||||
use codex_http_client::OutboundProxyPolicy;
|
||||
use codex_login::CodexAuth;
|
||||
use codex_login::auth::AgentIdentityAuthPolicy;
|
||||
use codex_model_provider_info::ModelProviderInfo;
|
||||
@@ -102,6 +103,7 @@ fn assert_request_trace_matches(body: &serde_json::Value, expected_trace: &W3cTr
|
||||
struct WebsocketTestHarness {
|
||||
_codex_home: TempDir,
|
||||
client: ModelClient,
|
||||
outbound_proxy_policy: OutboundProxyPolicy,
|
||||
session_id: SessionId,
|
||||
thread_id: ThreadId,
|
||||
model_info: ModelInfo,
|
||||
@@ -230,6 +232,38 @@ 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;
|
||||
assert_eq!(
|
||||
harness.outbound_proxy_policy,
|
||||
OutboundProxyPolicy::RespectSystemProxy
|
||||
);
|
||||
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 +423,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 +1041,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 +2254,7 @@ async fn websocket_harness_with_options(
|
||||
websocket_provider(server),
|
||||
runtime_metrics_enabled,
|
||||
/*concurrent_reasoning_summaries_enabled*/ false,
|
||||
/*enabled_features*/ &[],
|
||||
)
|
||||
.await
|
||||
}
|
||||
@@ -2224,6 +2263,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 +2280,15 @@ 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");
|
||||
}
|
||||
config.respect_system_proxy = config.features.enabled(Feature::RespectSystemProxy);
|
||||
let http_client_factory = config.http_client_factory();
|
||||
let outbound_proxy_policy = http_client_factory.outbound_proxy_policy();
|
||||
let config = Arc::new(config);
|
||||
let model_info = codex_core::test_support::construct_model_info_offline(MODEL, &config);
|
||||
let thread_id = ThreadId::new();
|
||||
@@ -2284,12 +2333,13 @@ async fn websocket_harness_with_provider_options(
|
||||
.features
|
||||
.enabled(Feature::ConcurrentReasoningSummaries),
|
||||
/*attestation_provider*/ None,
|
||||
config.http_client_factory(),
|
||||
http_client_factory,
|
||||
);
|
||||
|
||||
WebsocketTestHarness {
|
||||
_codex_home: codex_home,
|
||||
client,
|
||||
outbound_proxy_policy,
|
||||
session_id,
|
||||
thread_id,
|
||||
model_info,
|
||||
|
||||
Reference in New Issue
Block a user