mirror of
https://github.com/openai/codex.git
synced 2026-04-26 23:55:25 +00:00
core: preconnect Responses websocket for first turn (#10698)
## Problem The first user turn can pay websocket handshake latency even when a session has already started. We want to reduce that initial delay while preserving turn semantics and avoiding any prompt send during startup. Reviewer feedback also called out duplicated connect/setup paths and unnecessary preconnect state complexity. ## Mental model `ModelClient` owns session-scoped transport state. During session startup, it can opportunistically warm one websocket handshake slot. A turn-scoped `ModelClientSession` adopts that slot once if available, restores captured sticky turn-state, and otherwise opens a websocket through the same shared connect path. If startup preconnect is still in flight, first turn setup awaits that task and treats it as the first connection attempt for the turn. Preconnect is handshake-only. The first `response.create` is still sent only when a turn starts. ## Non-goals This change does not make preconnect required for correctness and does not change prompt/turn payload semantics. It also does not expand fallback behavior beyond clearing preconnect state when fallback activates. ## Tradeoffs The implementation prioritizes simpler ownership and shared connection code over header-match gating for reuse. The single-slot cache keeps lifecycle straightforward but only benefits the immediate next turn. Awaiting in-flight preconnect has the same app-level connect-timeout semantics as existing websocket connect behavior (no new timeout class introduced by this PR). ## Architecture `core/src/client.rs`: - Added session-level preconnect lifecycle state (`Idle` / `InFlight` / `Ready`) carrying one warmed websocket plus optional captured turn-state. - Added `pre_establish_connection()` startup warmup and `preconnect()` handshake-only setup. - Deduped auth/provider resolution into `current_client_setup()` and websocket handshake wiring into `connect_websocket()` / `build_websocket_headers()`. - Updated turn websocket path to adopt preconnect first, await in-flight preconnect when present, then create a new websocket only when needed. - Ensured fallback activation clears warmed preconnect state. - Added documentation for lifecycle, ownership, sticky-routing invariants, and timeout semantics. `core/src/codex.rs`: - Session startup invokes `model_client.pre_establish_connection(...)`. - Turn metadata resolution uses the shared timeout helper. `core/src/turn_metadata.rs`: - Centralized shared timeout helper used by both turn-time metadata resolution and startup preconnect metadata building. `core/tests/common/responses.rs` + websocket test suites: - Added deterministic handshake waiting helper (`wait_for_handshakes`) with bounded polling. - Added startup preconnect and in-flight preconnect reuse coverage. - Fallback expectations now assert exactly two websocket attempts in covered scenarios (startup preconnect + turn attempt before fallback sticks). ## Observability Preconnect remains best-effort and non-fatal. Existing websocket/fallback telemetry remains in place, and debug logs now make preconnect-await behavior and preconnect failures easier to reason about. ## Tests Validated with: 1. `just fmt` 2. `cargo test -p codex-core websocket_preconnect -- --nocapture` 3. `cargo test -p codex-core websocket_fallback -- --nocapture` 4. `cargo test -p codex-core websocket_first_turn_waits_for_inflight_preconnect -- --nocapture`
This commit is contained in:
@@ -268,6 +268,11 @@ impl WebSocketHandshake {
|
||||
pub struct WebSocketConnectionConfig {
|
||||
pub requests: Vec<Vec<Value>>,
|
||||
pub response_headers: Vec<(String, String)>,
|
||||
/// Optional delay inserted before accepting the websocket handshake.
|
||||
///
|
||||
/// Tests use this to force startup preconnect into an in-flight state so first-turn adoption
|
||||
/// paths can be exercised deterministically.
|
||||
pub accept_delay: Option<Duration>,
|
||||
}
|
||||
|
||||
pub struct WebSocketTestServer {
|
||||
@@ -299,6 +304,29 @@ impl WebSocketTestServer {
|
||||
self.handshakes.lock().unwrap().clone()
|
||||
}
|
||||
|
||||
/// Waits until at least `expected` websocket handshakes have been observed or timeout elapses.
|
||||
///
|
||||
/// Uses a short bounded polling interval so tests can deterministically wait for background
|
||||
/// preconnect activity without busy-spinning.
|
||||
pub async fn wait_for_handshakes(&self, expected: usize, timeout: Duration) -> bool {
|
||||
if self.handshakes.lock().unwrap().len() >= expected {
|
||||
return true;
|
||||
}
|
||||
|
||||
let deadline = tokio::time::Instant::now() + timeout;
|
||||
let poll_interval = Duration::from_millis(10);
|
||||
loop {
|
||||
if self.handshakes.lock().unwrap().len() >= expected {
|
||||
return true;
|
||||
}
|
||||
let now = tokio::time::Instant::now();
|
||||
if now >= deadline {
|
||||
return false;
|
||||
}
|
||||
let sleep_for = std::cmp::min(poll_interval, deadline.saturating_duration_since(now));
|
||||
tokio::time::sleep(sleep_for).await;
|
||||
}
|
||||
}
|
||||
pub fn single_handshake(&self) -> WebSocketHandshake {
|
||||
let handshakes = self.handshakes.lock().unwrap();
|
||||
if handshakes.len() != 1 {
|
||||
@@ -861,6 +889,7 @@ pub async fn start_websocket_server(connections: Vec<Vec<Vec<Value>>>) -> WebSoc
|
||||
.map(|requests| WebSocketConnectionConfig {
|
||||
requests,
|
||||
response_headers: Vec::new(),
|
||||
accept_delay: None,
|
||||
})
|
||||
.collect();
|
||||
start_websocket_server_with_headers(connections).await
|
||||
@@ -900,6 +929,10 @@ pub async fn start_websocket_server_with_headers(
|
||||
continue;
|
||||
};
|
||||
|
||||
if let Some(delay) = connection.accept_delay {
|
||||
tokio::time::sleep(delay).await;
|
||||
}
|
||||
|
||||
let response_headers = connection.response_headers.clone();
|
||||
let handshake_log = Arc::clone(&handshakes);
|
||||
let callback = move |req: &Request, mut response: Response| {
|
||||
|
||||
Reference in New Issue
Block a user