Files
codex/codex-rs/codex-api/src/endpoint/responses_websocket/connector.rs
T
jif 73178e7ca6 Keep Guardian v2 WebSocket pools warm with concurrent replenishment (#51378)
## Why

Serial connection opens let a stalled handshake delay healthy sockets. Waiting until the idle pool is empty to refill also leaves classification bursts using HTTP fallback.

## What changed

- Open up to four WebSockets concurrently and start replenishing when the idle pool falls to half capacity. Preserve successful in-flight opens when another open fails or times out.
- Make explicit prewarming wait for an existing refill before checking readiness.
- Add `ResponsesWebsocketClient::connect_with_connector` and reuse one connector per sampler, initializing TLS configuration on a blocking task while retaining per-connection authentication and destination policy checks.
- Record connection lease durations for admission, setup, and total time, tagged by transport.

## Testing

Add regression tests for a stalled handshake alongside a healthy socket and replenishment before the last idle socket is leased. Update cooldown, authentication, and transport retry tests to account for concurrent handshake completion.

GitOrigin-RevId: f4a8258977c1f2a58e00c6badc7b599280ef22ad
2026-10-06 14:42:57 +00:00

50 lines
1.9 KiB
Rust

//! Establishes Responses sockets with caller-owned TLS and proxy configuration.
//! Connector reuse never caches authentication or bypasses per-connection destination policy.
use super::ResponsesWebsocketClient;
use super::ResponsesWebsocketConnection;
use super::connect_websocket;
use super::merge_request_headers;
use crate::ApiError;
use crate::WebsocketTelemetry;
use codex_websocket_client::WebSocketConnector;
use http::HeaderMap;
use std::sync::Arc;
use std::sync::OnceLock;
impl ResponsesWebsocketClient {
/// Opens a socket using the caller's reusable connector and this client's current auth.
/// The caller owns the lifetime of the connector's TLS configuration snapshot.
#[tracing::instrument(
name = "responses_websocket.connect",
level = "info",
skip_all,
fields(transport = "responses_websocket", api.path = "/responses")
)]
pub async fn connect_with_connector(
&self,
connector: &WebSocketConnector,
extra_headers: HeaderMap,
default_headers: HeaderMap,
turn_state: Option<Arc<OnceLock<String>>>,
telemetry: Option<Arc<dyn WebsocketTelemetry>>,
) -> Result<ResponsesWebsocketConnection, ApiError> {
let ws_url = self
.provider
.websocket_url_for_path("/responses")
.map_err(|err| ApiError::Stream(format!("failed to build websocket URL: {err}")))?;
let mut headers =
merge_request_headers(&self.provider.headers, extra_headers, default_headers);
self.auth.add_auth_headers(&mut headers);
let (stream, _status, server_reasoning_included, server_model) =
connect_websocket(ws_url, headers, connector, turn_state).await?;
Ok(ResponsesWebsocketConnection::new(
stream,
self.provider.stream_idle_timeout,
server_reasoning_included,
server_model,
telemetry,
))
}
}