diff --git a/docs/OPENAI_CODEX_PARITY_AUDIT.md b/docs/OPENAI_CODEX_PARITY_AUDIT.md index 0ad7a3744..bb54ce4a3 100644 --- a/docs/OPENAI_CODEX_PARITY_AUDIT.md +++ b/docs/OPENAI_CODEX_PARITY_AUDIT.md @@ -57,9 +57,11 @@ Rust 状态: - key 以账号和 conversation 维度复用,一条 WS 严格单 in-flight,busy/cap/dead 时旁路 one-shot。 Rust 状态: -- 到位:`src/codex/gateway/transport/websocket.rs` 增加进程内 `CodexWebSocketPool`,默认 55 分钟过期。 -- 到位:池 key 由 base URL、本地账号和派生后的 conversation/prompt cache key 组成;活跃连接从池中移出,terminal frame 后再放回,确保单连接严格单 in-flight。 -- 到位:复用连接如果在首帧前发现已被对端关闭,会丢弃后用 fresh WS 重试一次;账号状态被标记限流/封禁/额度耗尽时会驱逐该账号的空闲 WS。 +- 到位:`src/codex/gateway/transport/websocket.rs` 增加进程内 `CodexWebSocketPool`,默认 55 分钟过期,默认每 25 秒维护空闲连接。 +- 到位:池 key 由 base URL、本地账号和派生后的 conversation/prompt cache key 组成;同 key busy、账号连接数达到上限或 pool disabled/shutdown 时旁路 one-shot,避免排队和死锁。 +- 到位:复用连接如果在首帧前发现已被对端关闭,会丢弃后用 fresh one-shot WS 重试一次;mid-stream 提前关闭会向客户端 stream 报错,不做重试。 +- 到位:空闲连接支持 keepalive ping/pong 和 liveness timeout,周期 sweep 会关闭过期、无响应或被 shutdown/evict 的连接;`websocket_connection_limit_reached` 会按连接级 fatal 错误剔除。 +- 到位:账号状态被标记限流/封禁/额度耗尽时会驱逐该账号的空闲 WS,忙碌中的连接完成后不会再放回池中。 ## 5. 非目标:IP 代理/VPN diff --git a/src/codex/gateway/transport/client.rs b/src/codex/gateway/transport/client.rs index abffe65a9..7b98878ba 100644 --- a/src/codex/gateway/transport/client.rs +++ b/src/codex/gateway/transport/client.rs @@ -724,7 +724,7 @@ fn websocket_error_allows_http_sse_fallback(error: &CodexWebSocketError) -> bool | StatusCode::UPGRADE_REQUIRED | StatusCode::NOT_IMPLEMENTED ), - CodexWebSocketError::InvalidRequest(_) => false, + CodexWebSocketError::InvalidRequest(_) | CodexWebSocketError::ClosedBeforeTerminal => false, } } diff --git a/src/codex/gateway/transport/websocket.rs b/src/codex/gateway/transport/websocket.rs index f38dda879..dabbed167 100644 --- a/src/codex/gateway/transport/websocket.rs +++ b/src/codex/gateway/transport/websocket.rs @@ -6,7 +6,7 @@ use std::{ }; use thiserror::Error; -use tokio::{net::TcpStream, sync::Mutex}; +use tokio::{net::TcpStream, sync::Mutex, time::timeout}; use tokio_tungstenite::{ connect_async, tungstenite::{ @@ -52,6 +52,8 @@ pub enum CodexWebSocketError { }, #[error("websocket response ended before any events")] EmptyResponse, + #[error("websocket closed before terminal event")] + ClosedBeforeTerminal, } #[derive(Debug, Clone, PartialEq, Eq)] @@ -99,41 +101,289 @@ impl CodexWebSocketPoolKey { #[derive(Clone)] pub struct CodexWebSocketPool { - inner: Arc>>, - max_age: Duration, + inner: Arc>, + config: CodexWebSocketPoolConfig, +} + +#[derive(Debug, Clone, Copy)] +pub struct CodexWebSocketPoolConfig { + pub enabled: bool, + pub max_age: Duration, + pub max_per_account: usize, + pub maintenance_interval: Option, + pub ping_interval: Option, + pub ping_timeout: Duration, + pub liveness_timeout: Option, } impl CodexWebSocketPool { const DEFAULT_MAX_AGE: Duration = Duration::from_secs(55 * 60); + const DEFAULT_MAX_PER_ACCOUNT: usize = 8; + const DEFAULT_MAINTENANCE_INTERVAL: Duration = Duration::from_secs(25); + const DEFAULT_PING_INTERVAL: Duration = Duration::from_secs(25); + const DEFAULT_PING_TIMEOUT: Duration = Duration::from_secs(5); + const DEFAULT_LIVENESS_TIMEOUT: Duration = Duration::from_millis(62_500); pub fn with_default_max_age() -> Self { - Self { - inner: Arc::new(Mutex::new(HashMap::new())), - max_age: Self::DEFAULT_MAX_AGE, - } + Self::with_config(CodexWebSocketPoolConfig::default()) + } + + pub fn with_limits(max_age: Duration, max_per_account: usize) -> Self { + Self::with_config(CodexWebSocketPoolConfig { + max_age, + max_per_account, + maintenance_interval: None, + ping_interval: None, + liveness_timeout: None, + ..CodexWebSocketPoolConfig::default() + }) + } + + pub fn with_config(config: CodexWebSocketPoolConfig) -> Self { + let pool = Self { + inner: Arc::new(Mutex::new(WebSocketPoolState::default())), + config, + }; + pool.spawn_maintenance_task(); + pool } pub async fn evict_account(&self, account_id: &str) { - self.inner - .lock() - .await - .retain(|key, _| key.account_id() != account_id); + let mut idle_connections = Vec::new(); + { + let mut state = self.inner.lock().await; + let keys = state + .slots + .keys() + .filter(|key| key.account_id() == account_id) + .cloned() + .collect::>(); + for key in keys { + if let Some(WebSocketPoolSlot::Idle(connection)) = state.slots.remove(&key) { + idle_connections.push(*connection); + } + } + } + close_idle_connections(idle_connections).await; } - async fn take(&self, key: &CodexWebSocketPoolKey) -> Option { - let connection = self.inner.lock().await.remove(key)?; - if connection.created_at.elapsed() <= self.max_age { - Some(connection) - } else { - None + pub async fn shutdown(&self) { + let idle_connections = { + let mut state = self.inner.lock().await; + state.shutting_down = true; + state + .slots + .drain() + .filter_map(|(_, slot)| match slot { + WebSocketPoolSlot::Idle(connection) => Some(*connection), + WebSocketPoolSlot::Busy | WebSocketPoolSlot::Checking => None, + }) + .collect::>() + }; + close_idle_connections(idle_connections).await; + } + + pub async fn gc_sweep(&self) { + self.maintain_idle_connections().await; + } + + async fn acquire(&self, key: &CodexWebSocketPoolKey) -> WebSocketPoolAcquire { + let mut expired_connection = None; + let acquire = { + let mut state = self.inner.lock().await; + if !self.config.enabled || state.shutting_down { + return WebSocketPoolAcquire::Bypass; + } + match state.slots.get(key) { + Some(WebSocketPoolSlot::Busy | WebSocketPoolSlot::Checking) => { + return WebSocketPoolAcquire::Bypass; + } + Some(WebSocketPoolSlot::Idle(_)) => { + let Some(WebSocketPoolSlot::Idle(connection)) = state.slots.remove(key) else { + return WebSocketPoolAcquire::Bypass; + }; + if connection.created_at.elapsed() <= self.config.max_age { + state.slots.insert(key.clone(), WebSocketPoolSlot::Busy); + return WebSocketPoolAcquire::Reused(connection); + } + expired_connection = Some(*connection); + } + None => {} + } + + if self.config.max_per_account == 0 + || account_slot_count(&state.slots, key.account_id()) >= self.config.max_per_account + { + WebSocketPoolAcquire::Bypass + } else { + state.slots.insert(key.clone(), WebSocketPoolSlot::Busy); + WebSocketPoolAcquire::FreshReserved + } + }; + + if let Some(connection) = expired_connection { + close_idle_connections(vec![connection]).await; } + + acquire } async fn put(&self, key: CodexWebSocketPoolKey, connection: PooledWebSocketConnection) { - if connection.created_at.elapsed() > self.max_age { + let mut connection = Some(connection); + { + let mut state = self.inner.lock().await; + let expired = connection + .as_ref() + .is_some_and(|connection| connection.created_at.elapsed() > self.config.max_age); + if expired || state.shutting_down || !self.config.enabled { + state.slots.remove(&key); + } else if matches!(state.slots.get(&key), Some(WebSocketPoolSlot::Busy)) { + if let Some(connection) = connection.take() { + state + .slots + .insert(key, WebSocketPoolSlot::Idle(Box::new(connection))); + } + } + } + if let Some(connection) = connection { + close_pooled_connection(connection).await; + } + } + + async fn discard(&self, key: &CodexWebSocketPoolKey) { + let idle_connection = { + let mut state = self.inner.lock().await; + match state.slots.remove(key) { + Some(WebSocketPoolSlot::Idle(connection)) => Some(*connection), + Some(WebSocketPoolSlot::Busy | WebSocketPoolSlot::Checking) | None => None, + } + }; + if let Some(connection) = idle_connection { + close_idle_connections(vec![connection]).await; + } + } + + fn spawn_maintenance_task(&self) { + let Some(interval_duration) = self.config.maintenance_interval else { return; + }; + let inner = Arc::downgrade(&self.inner); + let config = self.config; + let Ok(_handle) = tokio::runtime::Handle::try_current() else { + return; + }; + tokio::spawn(async move { + let mut interval = tokio::time::interval(interval_duration); + loop { + interval.tick().await; + let Some(inner) = inner.upgrade() else { + break; + }; + let pool = CodexWebSocketPool { inner, config }; + if pool.is_shutting_down().await { + break; + } + pool.maintain_idle_connections().await; + } + }); + } + + async fn is_shutting_down(&self) -> bool { + self.inner.lock().await.shutting_down + } + + pub async fn maintain_idle_connections(&self) { + let maintenance = self.take_idle_connections_for_maintenance().await; + close_idle_connections(maintenance.close).await; + for (key, connection) in maintenance.probe { + let connection = probe_idle_connection(connection, self.config.ping_timeout).await; + self.return_maintained_connection(key, connection).await; + } + } + + async fn take_idle_connections_for_maintenance(&self) -> WebSocketPoolMaintenance { + let mut close = Vec::new(); + let mut probe = Vec::new(); + let now = Instant::now(); + let Some(ping_interval) = self.config.ping_interval else { + let mut state = self.inner.lock().await; + if state.shutting_down || !self.config.enabled { + return WebSocketPoolMaintenance { close, probe }; + } + let expired_keys = expired_idle_keys(&state.slots, self.config.max_age); + for key in expired_keys { + if let Some(WebSocketPoolSlot::Idle(connection)) = state.slots.remove(&key) { + close.push(*connection); + } + } + return WebSocketPoolMaintenance { close, probe }; + }; + + let mut state = self.inner.lock().await; + if state.shutting_down || !self.config.enabled { + return WebSocketPoolMaintenance { close, probe }; + } + let keys = state + .slots + .iter() + .filter_map(|(key, slot)| match slot { + WebSocketPoolSlot::Idle(connection) + if should_close_idle_connection( + connection, + now, + self.config.max_age, + self.config.liveness_timeout, + ) || should_probe_idle_connection(connection, now, ping_interval) => + { + Some(key.clone()) + } + WebSocketPoolSlot::Idle(_) + | WebSocketPoolSlot::Busy + | WebSocketPoolSlot::Checking => None, + }) + .collect::>(); + for key in keys { + let Some(WebSocketPoolSlot::Idle(connection)) = state.slots.remove(&key) else { + continue; + }; + if should_close_idle_connection( + &connection, + now, + self.config.max_age, + self.config.liveness_timeout, + ) { + close.push(*connection); + } else { + state.slots.insert(key.clone(), WebSocketPoolSlot::Checking); + probe.push((key, *connection)); + } + } + WebSocketPoolMaintenance { close, probe } + } + + async fn return_maintained_connection( + &self, + key: CodexWebSocketPoolKey, + connection: Option, + ) { + let mut connection = connection; + { + let mut state = self.inner.lock().await; + if state.shutting_down || !self.config.enabled { + state.slots.remove(&key); + } else if matches!(state.slots.get(&key), Some(WebSocketPoolSlot::Checking)) { + if let Some(connection) = connection.take() { + state + .slots + .insert(key, WebSocketPoolSlot::Idle(Box::new(connection))); + } else { + state.slots.remove(&key); + } + } + } + if let Some(connection) = connection { + close_pooled_connection(connection).await; } - self.inner.lock().await.insert(key, connection); } } @@ -143,6 +393,26 @@ impl Default for CodexWebSocketPool { } } +impl Default for CodexWebSocketPoolConfig { + fn default() -> Self { + Self { + enabled: true, + max_age: CodexWebSocketPool::DEFAULT_MAX_AGE, + max_per_account: CodexWebSocketPool::DEFAULT_MAX_PER_ACCOUNT, + maintenance_interval: Some(CodexWebSocketPool::DEFAULT_MAINTENANCE_INTERVAL), + ping_interval: Some(CodexWebSocketPool::DEFAULT_PING_INTERVAL), + ping_timeout: CodexWebSocketPool::DEFAULT_PING_TIMEOUT, + liveness_timeout: Some(CodexWebSocketPool::DEFAULT_LIVENESS_TIMEOUT), + } + } +} + +#[derive(Default)] +struct WebSocketPoolState { + slots: HashMap, + shutting_down: bool, +} + type CodexWsStream = WebSocketStream>; #[derive(Debug, Clone, PartialEq, Eq)] @@ -166,6 +436,116 @@ struct PooledWebSocketConnection { websocket: CodexWsStream, metadata: CodexWebSocketConnectionMetadata, created_at: Instant, + last_activity_at: Instant, + last_ping_at: Option, +} + +enum WebSocketPoolSlot { + Idle(Box), + Busy, + Checking, +} + +enum WebSocketPoolAcquire { + Reused(Box), + FreshReserved, + Bypass, +} + +fn account_slot_count( + inner: &HashMap, + account_id: &str, +) -> usize { + inner + .keys() + .filter(|key| key.account_id() == account_id) + .count() +} + +async fn close_idle_connections(connections: Vec) { + for connection in connections { + close_pooled_connection(connection).await; + } +} + +async fn close_pooled_connection(connection: PooledWebSocketConnection) { + let mut websocket = connection.websocket; + let _ = websocket.close(None).await; +} + +struct WebSocketPoolMaintenance { + close: Vec, + probe: Vec<(CodexWebSocketPoolKey, PooledWebSocketConnection)>, +} + +fn expired_idle_keys( + slots: &HashMap, + max_age: Duration, +) -> Vec { + slots + .iter() + .filter_map(|(key, slot)| match slot { + WebSocketPoolSlot::Idle(connection) if connection.created_at.elapsed() > max_age => { + Some(key.clone()) + } + WebSocketPoolSlot::Idle(_) | WebSocketPoolSlot::Busy | WebSocketPoolSlot::Checking => { + None + } + }) + .collect() +} + +fn should_close_idle_connection( + connection: &PooledWebSocketConnection, + now: Instant, + max_age: Duration, + liveness_timeout: Option, +) -> bool { + connection.created_at.elapsed() > max_age + || liveness_timeout.is_some_and(|timeout| { + !timeout.is_zero() && now.duration_since(connection.last_activity_at) > timeout + }) +} + +fn should_probe_idle_connection( + connection: &PooledWebSocketConnection, + now: Instant, + ping_interval: Duration, +) -> bool { + !ping_interval.is_zero() + && connection + .last_ping_at + .is_none_or(|last_ping_at| now.duration_since(last_ping_at) >= ping_interval) +} + +async fn probe_idle_connection( + mut connection: PooledWebSocketConnection, + ping_timeout: Duration, +) -> Option { + let now = Instant::now(); + connection.last_ping_at = Some(now); + if connection + .websocket + .send(Message::Ping(Vec::new().into())) + .await + .is_err() + { + close_pooled_connection(connection).await; + return None; + } + if ping_timeout.is_zero() { + return Some(connection); + } + match timeout(ping_timeout, connection.websocket.next()).await { + Ok(Some(Ok(Message::Close(_)))) | Ok(None) | Ok(Some(Err(_))) | Err(_) => { + close_pooled_connection(connection).await; + None + } + Ok(Some(Ok(_message))) => { + connection.last_activity_at = Instant::now(); + Some(connection) + } + } } struct ActiveWebSocket { @@ -173,6 +553,7 @@ struct ActiveWebSocket { metadata: CodexWebSocketConnectionMetadata, pool_return: Option, reused: bool, + last_activity_at: Instant, } struct WebSocketPoolReturn { @@ -188,6 +569,7 @@ impl ActiveWebSocket { metadata, pool_return, reused: _, + last_activity_at, } = self; let Some(pool_return) = pool_return else { let mut websocket = websocket; @@ -202,13 +584,22 @@ impl ActiveWebSocket { websocket, metadata, created_at: pool_return.created_at, + last_activity_at, + last_ping_at: None, }, ) .await; } async fn discard(self) { - let mut websocket = self.websocket; + let Self { + mut websocket, + pool_return, + .. + } = self; + if let Some(pool_return) = pool_return { + pool_return.pool.discard(&pool_return.key).await; + } let _ = websocket.close(None).await; } } @@ -283,9 +674,10 @@ async fn create_response_via_websocket_stream_inner( headers: HeaderMap, pool: Option<(Arc, CodexWebSocketPoolKey)>, ) -> Result { - let mut retry_stale_reuse = pool.is_some(); + let mut pool_context = pool; + let mut retry_stale_reuse = pool_context.is_some(); loop { - let mut active = acquire_websocket(base_url, headers.clone(), pool.clone()).await?; + let mut active = acquire_websocket(base_url, headers.clone(), pool_context.clone()).await?; if let Err(error) = active .websocket .send(Message::Text( @@ -297,6 +689,7 @@ async fn create_response_via_websocket_stream_inner( active.discard().await; if reused && retry_stale_reuse { retry_stale_reuse = false; + pool_context = None; continue; } return Err(CodexWebSocketError::Transport(error)); @@ -309,6 +702,7 @@ async fn create_response_via_websocket_stream_inner( active.discard().await; if reused && retry_stale_reuse { retry_stale_reuse = false; + pool_context = None; break; } return Err(CodexWebSocketError::EmptyResponse); @@ -320,21 +714,27 @@ async fn create_response_via_websocket_stream_inner( active.discard().await; if reused && retry_stale_reuse { retry_stale_reuse = false; + pool_context = None; break; } return Err(CodexWebSocketError::Transport(error)); } }; + active.last_activity_at = Instant::now(); let Some(raw) = websocket_message_text(message) else { continue; }; if is_internal_websocket_event(&raw) { continue; } - if let Some(status) = classify_ws_error_frame(&raw) { - active.finish().await; + if let Some(classified) = classify_ws_error_frame(&raw) { + if classified.connection_fatal { + active.discard().await; + } else { + active.finish().await; + } return Err(CodexWebSocketError::Upstream { - status, + status: classified.status, retry_after_seconds: retry_after_seconds_from_body(&raw), body: raw, }); @@ -373,30 +773,59 @@ async fn acquire_websocket( pool: Option<(Arc, CodexWebSocketPoolKey)>, ) -> Result { let Some((pool, key)) = pool else { + let created_at = Instant::now(); let (websocket, metadata) = connect_websocket(base_url, headers).await?; return Ok(ActiveWebSocket { websocket, metadata, pool_return: None, reused: false, + last_activity_at: created_at, }); }; - if let Some(connection) = pool.take(&key).await { - return Ok(ActiveWebSocket { - websocket: connection.websocket, - metadata: connection.metadata, - pool_return: Some(WebSocketPoolReturn { - pool, - key, - created_at: connection.created_at, - }), - reused: true, - }); + match pool.acquire(&key).await { + WebSocketPoolAcquire::Reused(connection) => { + let PooledWebSocketConnection { + websocket, + metadata, + created_at, + last_activity_at, + last_ping_at: _, + } = *connection; + return Ok(ActiveWebSocket { + websocket, + metadata, + pool_return: Some(WebSocketPoolReturn { + pool, + key, + created_at, + }), + reused: true, + last_activity_at, + }); + } + WebSocketPoolAcquire::FreshReserved => {} + WebSocketPoolAcquire::Bypass => { + let (websocket, metadata) = connect_websocket(base_url, headers).await?; + return Ok(ActiveWebSocket { + websocket, + metadata, + pool_return: None, + reused: false, + last_activity_at: Instant::now(), + }); + } } let created_at = Instant::now(); - let (websocket, metadata) = connect_websocket(base_url, headers).await?; + let (websocket, metadata) = match connect_websocket(base_url, headers).await { + Ok(connection) => connection, + Err(error) => { + pool.discard(&key).await; + return Err(error); + } + }; Ok(ActiveWebSocket { websocket, metadata, @@ -406,6 +835,7 @@ async fn acquire_websocket( created_at, }), reused: false, + last_activity_at: created_at, }) } @@ -535,6 +965,7 @@ async fn forward_websocket_as_sse( while let Some(message) = active.websocket.next().await { match message { Ok(message) => { + active.last_activity_at = Instant::now(); let Some(raw) = websocket_message_text(message) else { continue; }; @@ -546,22 +977,30 @@ async fn forward_websocket_as_sse( let chunk = websocket_sse_chunk(&raw, event.as_deref()); if tx.unbounded_send(Ok(chunk)).is_err() { active.discard().await; - break; + return; } if terminal { active.finish().await; - break; + return; } } Err(error) => { let _ = tx.unbounded_send(Err(CodexWebSocketError::Transport(error))); - break; + active.discard().await; + return; } } } + let _ = tx.unbounded_send(Err(CodexWebSocketError::ClosedBeforeTerminal)); + active.discard().await; } -fn classify_ws_error_frame(raw: &str) -> Option { +struct ClassifiedWebSocketError { + status: StatusCode, + connection_fatal: bool, +} + +fn classify_ws_error_frame(raw: &str) -> Option { let value = serde_json::from_str::(raw).ok()?; let event_type = value.get("type").and_then(Value::as_str)?; if event_type != "error" && event_type != "response.failed" { @@ -574,7 +1013,11 @@ fn classify_ws_error_frame(raw: &str) -> Option { .or_else(|| value.pointer("/error/type")) .and_then(Value::as_str)? .to_ascii_lowercase(); - rotatable_error_status(&code) + let status = rotatable_error_status(&code)?; + Some(ClassifiedWebSocketError { + status, + connection_fatal: code == "websocket_connection_limit_reached", + }) } fn rotatable_error_status(code: &str) -> Option { @@ -588,6 +1031,7 @@ fn rotatable_error_status(code: &str) -> Option { } "forbidden" | "account_banned" | "banned" => Some(StatusCode::FORBIDDEN), "previous_response_not_found" => Some(StatusCode::BAD_REQUEST), + "websocket_connection_limit_reached" => Some(StatusCode::SERVICE_UNAVAILABLE), _ => None, } } diff --git a/tests/codex_gateway/websocket.rs b/tests/codex_gateway/websocket.rs index 499ead2b5..c7988fa5c 100644 --- a/tests/codex_gateway/websocket.rs +++ b/tests/codex_gateway/websocket.rs @@ -1,11 +1,12 @@ -use std::sync::Arc; +use std::{sync::Arc, time::Duration}; use futures::{SinkExt, StreamExt}; use serde_json::{json, Value}; use tokio::{ io::{AsyncReadExt, AsyncWriteExt}, net::TcpListener, - sync::Mutex, + sync::{oneshot, Mutex}, + time::sleep, }; use tokio_tungstenite::{ accept_hdr_async, @@ -18,7 +19,10 @@ use tokio_tungstenite::{ use codex_proxy_rs::codex::gateway::transport::{ client::{build_reqwest_client, CodexBackendClient, CodexClientError, CodexRequestContext}, types::CodexResponsesRequest, - websocket::{http_sse_fallback_allowed, transport_for_request, CodexTransport}, + websocket::{ + http_sse_fallback_allowed, transport_for_request, CodexTransport, CodexWebSocketPool, + CodexWebSocketPoolConfig, + }, }; #[test] @@ -661,10 +665,678 @@ async fn websocket_success_should_capture_handshake_headers_and_rate_limit_heade .any(|(name, value)| { name == "x-ratelimit-remaining-requests" && value == "41" })); } +#[tokio::test] +#[expect( + clippy::result_large_err, + reason = "tokio-tungstenite handshake callbacks use a large error response type" +)] +async fn websocket_pool_should_bypass_busy_key_with_one_shot_connections() { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + let (release_first_tx, release_first_rx) = oneshot::channel(); + let server = tokio::spawn(async move { + let (first_stream, _) = listener.accept().await.unwrap(); + let mut first_ws = + accept_hdr_async(first_stream, |_request: &WsRequest, response| Ok(response)) + .await + .unwrap(); + let _first_request = first_ws.next().await.unwrap().unwrap(); + first_ws + .send(Message::Text( + json!({ + "type": "response.output_text.delta", + "delta": "first connection is still busy" + }) + .to_string() + .into(), + )) + .await + .unwrap(); + + let (second_stream, _) = listener.accept().await.unwrap(); + let mut second_ws = + accept_hdr_async(second_stream, |_request: &WsRequest, response| Ok(response)) + .await + .unwrap(); + let _second_request = second_ws.next().await.unwrap().unwrap(); + second_ws + .send(Message::Text( + websocket_completed_response("resp_busy_second", 2, 1).into(), + )) + .await + .unwrap(); + + let third_on_new_connection = tokio::select! { + message = second_ws.next() => { + if matches!(message, Some(Ok(message)) if message.is_text()) { + false + } else { + let (third_stream, _) = listener.accept().await.unwrap(); + let mut third_ws = accept_hdr_async( + third_stream, + |_request: &WsRequest, response| Ok(response), + ) + .await + .unwrap(); + let _third_request = third_ws.next().await.unwrap().unwrap(); + third_ws + .send(Message::Text( + websocket_completed_response("resp_busy_third", 2, 1).into(), + )) + .await + .unwrap(); + third_ws.close(None).await.unwrap(); + true + } + } + accepted = listener.accept() => { + let (third_stream, _) = accepted.unwrap(); + let mut third_ws = accept_hdr_async( + third_stream, + |_request: &WsRequest, response| Ok(response), + ) + .await + .unwrap(); + let _third_request = third_ws.next().await.unwrap().unwrap(); + third_ws + .send(Message::Text( + websocket_completed_response("resp_busy_third", 2, 1).into(), + )) + .await + .unwrap(); + third_ws.close(None).await.unwrap(); + true + } + }; + + release_first_rx.await.unwrap(); + first_ws + .send(Message::Text( + websocket_completed_response("resp_busy_first", 2, 1).into(), + )) + .await + .unwrap(); + first_ws.close(None).await.unwrap(); + third_on_new_connection + }); + + let pool = Arc::new(CodexWebSocketPool::with_default_max_age()); + let client = CodexBackendClient::new( + build_reqwest_client(false).unwrap(), + format!("http://{addr}"), + codex_proxy_rs::codex::gateway::fingerprint::model::Fingerprint::default_for_tests(), + ) + .with_websocket_pool(pool, "entry-a"); + let mut request = base_request(); + request.prompt_cache_key = Some("conversation-a".to_string()); + + let mut first_stream = client + .websocket_stream_response(&request, request_context("req_busy_first", Some("conv-a"))) + .await + .unwrap() + .body_stream; + let first_chunk = first_stream.next().await.unwrap().unwrap(); + assert!(first_chunk.contains("first connection is still busy")); + + let second = client + .create_response(&request, request_context("req_busy_second", Some("conv-a"))) + .await + .unwrap(); + assert!(second.body.contains("\"id\":\"resp_busy_second\"")); + + let third = client + .create_response(&request, request_context("req_busy_third", Some("conv-a"))) + .await + .unwrap(); + assert!(third.body.contains("\"id\":\"resp_busy_third\"")); + + release_first_tx.send(()).unwrap(); + while first_stream.next().await.transpose().unwrap().is_some() {} + + assert!( + server.await.unwrap(), + "busy pool key reused a second pooled websocket instead of one-shot bypass" + ); +} + +#[tokio::test] +#[expect( + clippy::result_large_err, + reason = "tokio-tungstenite handshake callbacks use a large error response type" +)] +async fn websocket_pool_should_bypass_new_keys_after_account_cap() { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + let server = tokio::spawn(async move { + let (first_stream, _) = listener.accept().await.unwrap(); + let mut first_ws = + accept_hdr_async(first_stream, |_request: &WsRequest, response| Ok(response)) + .await + .unwrap(); + let _first_request = first_ws.next().await.unwrap().unwrap(); + first_ws + .send(Message::Text( + websocket_completed_response("resp_cap_first", 2, 1).into(), + )) + .await + .unwrap(); + + let (second_stream, _) = listener.accept().await.unwrap(); + let mut second_ws = + accept_hdr_async(second_stream, |_request: &WsRequest, response| Ok(response)) + .await + .unwrap(); + let _second_request = second_ws.next().await.unwrap().unwrap(); + second_ws + .send(Message::Text( + websocket_completed_response("resp_cap_second", 2, 1).into(), + )) + .await + .unwrap(); + + let third_on_new_connection = tokio::select! { + message = second_ws.next() => { + if matches!(message, Some(Ok(message)) if message.is_text()) { + false + } else { + let (third_stream, _) = listener.accept().await.unwrap(); + let mut third_ws = accept_hdr_async( + third_stream, + |_request: &WsRequest, response| Ok(response), + ) + .await + .unwrap(); + let _third_request = third_ws.next().await.unwrap().unwrap(); + third_ws + .send(Message::Text( + websocket_completed_response("resp_cap_third", 2, 1).into(), + )) + .await + .unwrap(); + third_ws.close(None).await.unwrap(); + true + } + } + accepted = listener.accept() => { + let (third_stream, _) = accepted.unwrap(); + let mut third_ws = accept_hdr_async( + third_stream, + |_request: &WsRequest, response| Ok(response), + ) + .await + .unwrap(); + let _third_request = third_ws.next().await.unwrap().unwrap(); + third_ws + .send(Message::Text( + websocket_completed_response("resp_cap_third", 2, 1).into(), + )) + .await + .unwrap(); + third_ws.close(None).await.unwrap(); + true + } + }; + + first_ws.close(None).await.unwrap(); + third_on_new_connection + }); + + let pool = Arc::new(CodexWebSocketPool::with_limits(Duration::from_secs(60), 1)); + let client = CodexBackendClient::new( + build_reqwest_client(false).unwrap(), + format!("http://{addr}"), + codex_proxy_rs::codex::gateway::fingerprint::model::Fingerprint::default_for_tests(), + ) + .with_websocket_pool(pool, "entry-cap"); + + let mut first_request = base_request(); + first_request.prompt_cache_key = Some("conversation-one".to_string()); + let first = client + .create_response( + &first_request, + request_context("req_cap_first", Some("conv-one")), + ) + .await + .unwrap(); + assert!(first.body.contains("\"id\":\"resp_cap_first\"")); + + let mut second_request = base_request(); + second_request.prompt_cache_key = Some("conversation-two".to_string()); + let second = client + .create_response( + &second_request, + request_context("req_cap_second", Some("conv-two")), + ) + .await + .unwrap(); + assert!(second.body.contains("\"id\":\"resp_cap_second\"")); + + let third = client + .create_response( + &second_request, + request_context("req_cap_third", Some("conv-two")), + ) + .await + .unwrap(); + assert!(third.body.contains("\"id\":\"resp_cap_third\"")); + + assert!( + server.await.unwrap(), + "account cap bypass still allowed a capped conversation to be pooled" + ); +} + +#[tokio::test] +#[expect( + clippy::result_large_err, + reason = "tokio-tungstenite handshake callbacks use a large error response type" +)] +async fn websocket_stream_should_error_when_connection_closes_before_terminal_frame() { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + let server = tokio::spawn(async move { + let (stream, _) = listener.accept().await.unwrap(); + let mut websocket = accept_hdr_async(stream, |_request: &WsRequest, response| Ok(response)) + .await + .unwrap(); + let _request = websocket.next().await.unwrap().unwrap(); + websocket + .send(Message::Text( + json!({ + "type": "response.output_text.delta", + "delta": "partial" + }) + .to_string() + .into(), + )) + .await + .unwrap(); + websocket.close(None).await.unwrap(); + }); + + let client = CodexBackendClient::new( + build_reqwest_client(false).unwrap(), + format!("http://{addr}"), + codex_proxy_rs::codex::gateway::fingerprint::model::Fingerprint::default_for_tests(), + ); + let mut stream = client + .websocket_stream_response(&base_request(), request_context("req_mid_close", None)) + .await + .unwrap() + .body_stream; + + let first_chunk = stream.next().await.unwrap().unwrap(); + assert!(first_chunk.contains("partial")); + let error = stream.next().await.unwrap().unwrap_err(); + assert!(error + .to_string() + .contains("websocket closed before terminal event")); + + server.await.unwrap(); +} + +#[tokio::test] +#[expect( + clippy::result_large_err, + reason = "tokio-tungstenite handshake callbacks use a large error response type" +)] +async fn websocket_pool_keepalive_should_keep_idle_connection_when_pong_is_received() { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + let server = tokio::spawn(async move { + let (stream, _) = listener.accept().await.unwrap(); + let mut websocket = accept_hdr_async(stream, |_request: &WsRequest, response| Ok(response)) + .await + .unwrap(); + let _first_request = websocket.next().await.unwrap().unwrap(); + websocket + .send(Message::Text( + websocket_completed_response("resp_keepalive_first", 2, 1).into(), + )) + .await + .unwrap(); + + let ping = websocket.next().await.unwrap().unwrap(); + assert!( + ping.is_ping(), + "expected pool maintenance ping, got {ping:?}" + ); + websocket + .send(Message::Pong(Vec::new().into())) + .await + .unwrap(); + + let reused = tokio::select! { + message = websocket.next() => { + let message = message.unwrap().unwrap(); + if !message.is_text() { + return false; + } + websocket + .send(Message::Text( + websocket_completed_response("resp_keepalive_second", 2, 1).into(), + )) + .await + .unwrap(); + true + } + accepted = listener.accept() => { + let (stream, _) = accepted.unwrap(); + let mut websocket = accept_hdr_async( + stream, + |_request: &WsRequest, response| Ok(response), + ) + .await + .unwrap(); + let _request = websocket.next().await.unwrap().unwrap(); + websocket + .send(Message::Text( + websocket_completed_response("resp_keepalive_second", 2, 1).into(), + )) + .await + .unwrap(); + false + } + }; + + websocket.close(None).await.unwrap(); + reused + }); + + let pool = Arc::new(CodexWebSocketPool::with_config(keepalive_pool_config( + Duration::from_millis(50), + ))); + let client = CodexBackendClient::new( + build_reqwest_client(false).unwrap(), + format!("http://{addr}"), + codex_proxy_rs::codex::gateway::fingerprint::model::Fingerprint::default_for_tests(), + ) + .with_websocket_pool(Arc::clone(&pool), "entry-keepalive"); + let mut request = base_request(); + request.prompt_cache_key = Some("conversation-keepalive".to_string()); + + let first = client + .create_response( + &request, + request_context("req_keepalive_first", Some("conv-keepalive")), + ) + .await + .unwrap(); + assert!(first.body.contains("\"id\":\"resp_keepalive_first\"")); + + pool.maintain_idle_connections().await; + + let second = client + .create_response( + &request, + request_context("req_keepalive_second", Some("conv-keepalive")), + ) + .await + .unwrap(); + assert!(second.body.contains("\"id\":\"resp_keepalive_second\"")); + + assert!( + server.await.unwrap(), + "keepalive pong should leave the pooled websocket reusable" + ); +} + +#[tokio::test] +#[expect( + clippy::result_large_err, + reason = "tokio-tungstenite handshake callbacks use a large error response type" +)] +async fn websocket_pool_keepalive_should_evict_idle_connection_without_pong() { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + let server = tokio::spawn(async move { + let (first_stream, _) = listener.accept().await.unwrap(); + let mut first_ws = + accept_hdr_async(first_stream, |_request: &WsRequest, response| Ok(response)) + .await + .unwrap(); + let _first_request = first_ws.next().await.unwrap().unwrap(); + first_ws + .send(Message::Text( + websocket_completed_response("resp_no_pong_first", 2, 1).into(), + )) + .await + .unwrap(); + + let (second_stream, _) = listener.accept().await.unwrap(); + let mut second_ws = + accept_hdr_async(second_stream, |_request: &WsRequest, response| Ok(response)) + .await + .unwrap(); + let _second_request = second_ws.next().await.unwrap().unwrap(); + second_ws + .send(Message::Text( + websocket_completed_response("resp_no_pong_second", 2, 1).into(), + )) + .await + .unwrap(); + second_ws.close(None).await.unwrap(); + }); + + let pool = Arc::new(CodexWebSocketPool::with_config(keepalive_pool_config( + Duration::from_millis(20), + ))); + let client = CodexBackendClient::new( + build_reqwest_client(false).unwrap(), + format!("http://{addr}"), + codex_proxy_rs::codex::gateway::fingerprint::model::Fingerprint::default_for_tests(), + ) + .with_websocket_pool(Arc::clone(&pool), "entry-no-pong"); + let mut request = base_request(); + request.prompt_cache_key = Some("conversation-no-pong".to_string()); + + let first = client + .create_response( + &request, + request_context("req_no_pong_first", Some("conv-no-pong")), + ) + .await + .unwrap(); + assert!(first.body.contains("\"id\":\"resp_no_pong_first\"")); + + pool.maintain_idle_connections().await; + + let second = client + .create_response( + &request, + request_context("req_no_pong_second", Some("conv-no-pong")), + ) + .await + .unwrap(); + assert!(second.body.contains("\"id\":\"resp_no_pong_second\"")); + + server.await.unwrap(); +} + +#[tokio::test] +#[expect( + clippy::result_large_err, + reason = "tokio-tungstenite handshake callbacks use a large error response type" +)] +async fn websocket_pool_gc_sweep_should_close_expired_idle_connections() { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + let server = tokio::spawn(async move { + for response_id in ["resp_gc_first", "resp_gc_second"] { + let (stream, _) = listener.accept().await.unwrap(); + let mut websocket = + accept_hdr_async(stream, |_request: &WsRequest, response| Ok(response)) + .await + .unwrap(); + let _request = websocket.next().await.unwrap().unwrap(); + websocket + .send(Message::Text( + websocket_completed_response(response_id, 2, 1).into(), + )) + .await + .unwrap(); + websocket.close(None).await.unwrap(); + } + }); + + let pool = Arc::new(CodexWebSocketPool::with_config(manual_pool_config( + Duration::from_millis(5), + 8, + ))); + let client = CodexBackendClient::new( + build_reqwest_client(false).unwrap(), + format!("http://{addr}"), + codex_proxy_rs::codex::gateway::fingerprint::model::Fingerprint::default_for_tests(), + ) + .with_websocket_pool(Arc::clone(&pool), "entry-gc"); + let mut request = base_request(); + request.prompt_cache_key = Some("conversation-gc".to_string()); + + let first = client + .create_response(&request, request_context("req_gc_first", Some("conv-gc"))) + .await + .unwrap(); + assert!(first.body.contains("\"id\":\"resp_gc_first\"")); + + sleep(Duration::from_millis(15)).await; + pool.gc_sweep().await; + + let second = client + .create_response(&request, request_context("req_gc_second", Some("conv-gc"))) + .await + .unwrap(); + assert!(second.body.contains("\"id\":\"resp_gc_second\"")); + + server.await.unwrap(); +} + +#[tokio::test] +#[expect( + clippy::result_large_err, + reason = "tokio-tungstenite handshake callbacks use a large error response type" +)] +async fn websocket_pool_shutdown_should_close_idle_connections_and_disable_pooling() { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + let server = tokio::spawn(async move { + for response_id in ["resp_shutdown_first", "resp_shutdown_second"] { + let (stream, _) = listener.accept().await.unwrap(); + let mut websocket = + accept_hdr_async(stream, |_request: &WsRequest, response| Ok(response)) + .await + .unwrap(); + let _request = websocket.next().await.unwrap().unwrap(); + websocket + .send(Message::Text( + websocket_completed_response(response_id, 2, 1).into(), + )) + .await + .unwrap(); + websocket.close(None).await.unwrap(); + } + }); + + let pool = Arc::new(CodexWebSocketPool::with_config(manual_pool_config( + Duration::from_secs(60), + 8, + ))); + let client = CodexBackendClient::new( + build_reqwest_client(false).unwrap(), + format!("http://{addr}"), + codex_proxy_rs::codex::gateway::fingerprint::model::Fingerprint::default_for_tests(), + ) + .with_websocket_pool(Arc::clone(&pool), "entry-shutdown"); + let mut request = base_request(); + request.prompt_cache_key = Some("conversation-shutdown".to_string()); + + let first = client + .create_response( + &request, + request_context("req_shutdown_first", Some("conv-shutdown")), + ) + .await + .unwrap(); + assert!(first.body.contains("\"id\":\"resp_shutdown_first\"")); + + pool.shutdown().await; + + let second = client + .create_response( + &request, + request_context("req_shutdown_second", Some("conv-shutdown")), + ) + .await + .unwrap(); + assert!(second.body.contains("\"id\":\"resp_shutdown_second\"")); + + server.await.unwrap(); +} + fn base_request() -> CodexResponsesRequest { CodexResponsesRequest::new_http_sse("gpt-5.5", "", Vec::new()) } +fn manual_pool_config(max_age: Duration, max_per_account: usize) -> CodexWebSocketPoolConfig { + CodexWebSocketPoolConfig { + enabled: true, + max_age, + max_per_account, + maintenance_interval: None, + ping_interval: None, + ping_timeout: Duration::from_millis(50), + liveness_timeout: None, + } +} + +fn keepalive_pool_config(ping_timeout: Duration) -> CodexWebSocketPoolConfig { + CodexWebSocketPoolConfig { + ping_interval: Some(Duration::from_millis(1)), + ping_timeout, + liveness_timeout: Some(Duration::from_secs(60)), + ..manual_pool_config(Duration::from_secs(60), 8) + } +} + +fn request_context<'a>( + request_id: &'a str, + session_id: Option<&'a str>, +) -> CodexRequestContext<'a> { + CodexRequestContext { + access_token: "access-token", + account_id: Some("chatgpt-account"), + request_id, + turn_state: None, + turn_metadata: None, + beta_features: None, + include_timing_metrics: None, + version: None, + codex_window_id: None, + parent_thread_id: None, + cookie_header: None, + installation_id: None, + session_id, + } +} + +fn websocket_completed_response( + response_id: &str, + input_tokens: i64, + output_tokens: i64, +) -> String { + json!({ + "type": "response.completed", + "response": { + "id": response_id, + "object": "response", + "usage": { + "input_tokens": input_tokens, + "output_tokens": output_tokens + } + } + }) + .to_string() +} + async fn read_http_upgrade_request(stream: &mut tokio::net::TcpStream) -> String { let mut request = Vec::new(); let mut buffer = [0u8; 1024];