feat: harden ws pool

Co-authored-by: Codex <noreply@openai.com>

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
 

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,
     }
 }
 

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<Mutex<HashMap<CodexWebSocketPoolKey, PooledWebSocketConnection>>>,
-    max_age: Duration,
+    inner: Arc<Mutex<WebSocketPoolState>>,
+    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<Duration>,
+    pub ping_interval: Option<Duration>,
+    pub ping_timeout: Duration,
+    pub liveness_timeout: Option<Duration>,
 }
 
 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::<Vec<_>>();
+            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<PooledWebSocketConnection> {
-        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::<Vec<_>>()
+        };
+        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::<Vec<_>>();
+        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<PooledWebSocketConnection>,
+    ) {
+        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<CodexWebSocketPoolKey, WebSocketPoolSlot>,
+    shutting_down: bool,
+}
+
 type CodexWsStream = WebSocketStream<MaybeTlsStream<TcpStream>>;
 
 #[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<Instant>,
+}
+
+enum WebSocketPoolSlot {
+    Idle(Box<PooledWebSocketConnection>),
+    Busy,
+    Checking,
+}
+
+enum WebSocketPoolAcquire {
+    Reused(Box<PooledWebSocketConnection>),
+    FreshReserved,
+    Bypass,
+}
+
+fn account_slot_count(
+    inner: &HashMap<CodexWebSocketPoolKey, WebSocketPoolSlot>,
+    account_id: &str,
+) -> usize {
+    inner
+        .keys()
+        .filter(|key| key.account_id() == account_id)
+        .count()
+}
+
+async fn close_idle_connections(connections: Vec<PooledWebSocketConnection>) {
+    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<PooledWebSocketConnection>,
+    probe: Vec<(CodexWebSocketPoolKey, PooledWebSocketConnection)>,
+}
+
+fn expired_idle_keys(
+    slots: &HashMap<CodexWebSocketPoolKey, WebSocketPoolSlot>,
+    max_age: Duration,
+) -> Vec<CodexWebSocketPoolKey> {
+    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<Duration>,
+) -> 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<PooledWebSocketConnection> {
+    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<WebSocketPoolReturn>,
     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<CodexWebSocketPool>, CodexWebSocketPoolKey)>,
 ) -> Result<CodexWebSocketStreamResponse, CodexWebSocketError> {
-    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<CodexWebSocketPool>, CodexWebSocketPoolKey)>,
 ) -> Result<ActiveWebSocket, CodexWebSocketError> {
     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<StatusCode> {
+struct ClassifiedWebSocketError {
+    status: StatusCode,
+    connection_fatal: bool,
+}
+
+fn classify_ws_error_frame(raw: &str) -> Option<ClassifiedWebSocketError> {
     let value = serde_json::from_str::<Value>(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<StatusCode> {
         .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<StatusCode> {
@@ -588,6 +1031,7 @@ fn rotatable_error_status(code: &str) -> Option<StatusCode> {
         }
         "forbidden" | "account_banned" | "banned" => Some(StatusCode::FORBIDDEN),
         "previous_response_not_found" => Some(StatusCode::BAD_REQUEST),
+        "websocket_connection_limit_reached" => Some(StatusCode::SERVICE_UNAVAILABLE),
         _ => None,
     }
 }

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];