fix: enforce websocket first token retries

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

backend/src/runtime/services.rs

@@ -175,11 +175,14 @@ impl Services {
             }))
         });
         let codex = {
+            let websocket_first_token_timeout = (config.ws_pool.first_token_timeout_ms > 0)
+                .then(|| std::time::Duration::from_millis(config.ws_pool.first_token_timeout_ms));
             let client = CodexBackendClient::new(
                 build_reqwest_client(config.tls.force_http11)?,
                 config.api.base_url.clone(),
                 fingerprint.clone(),
-            );
+            )
+            .with_websocket_first_token_timeout(websocket_first_token_timeout);
             if let Some(pool) = &websocket_pool {
                 StdArc::new(client.with_websocket_pool(pool.clone()))
             } else {

backend/src/upstream/transport/client.rs

@@ -46,7 +46,9 @@ use super::websocket::{
     write_websocket_audit_artifact_from_env, CodexWebSocketConnection, CodexWebSocketExchangeError,
     CodexWebSocketRateLimitHeaderUpdates, CodexWebSocketTurnStateUpdate,
 };
-use super::websocket_pool::{CodexWebSocketPool, CodexWebSocketPoolKey, WebSocketPoolDecision};
+use super::websocket_pool::{
+    CodexWebSocketPool, CodexWebSocketPoolKey, WebSocketPoolDecision, DEFAULT_FIRST_TOKEN_TIMEOUT,
+};
 
 // ---------------------------------------------------------------------------
 // Constants
@@ -324,6 +326,7 @@ pub struct CodexBackendClient {
     pub(super) base_url: String,
     fingerprint: Fingerprint,
     websocket_pool: Option<Arc<CodexWebSocketPool>>,
+    websocket_first_token_timeout: Option<Duration>,
 }
 
 impl CodexBackendClient {
@@ -334,9 +337,15 @@ impl CodexBackendClient {
             base_url: base_url.into().trim_end_matches('/').to_string(),
             fingerprint,
             websocket_pool: None,
+            websocket_first_token_timeout: Some(DEFAULT_FIRST_TOKEN_TIMEOUT),
         }
     }
 
+    pub fn with_websocket_first_token_timeout(mut self, timeout: Option<Duration>) -> Self {
+        self.websocket_first_token_timeout = timeout.filter(|timeout| !timeout.is_zero());
+        self
+    }
+
     /// 为 Responses WebSocket 请求启用连接池。
     pub fn with_websocket_pool(mut self, pool: Arc<CodexWebSocketPool>) -> Self {
         self.websocket_pool = Some(pool);
@@ -590,10 +599,23 @@ impl CodexBackendClient {
         let pool_log_context = pool_key.as_ref().map(WebSocketPoolLogContext::from_key);
         let exchange = match (self.websocket_pool.as_deref(), pool_key) {
             (Some(pool), Some(key)) => {
-                execute_response_create_request_with_pool(&prepared, Some((pool, key)), started_at)
-                    .await
+                execute_response_create_request_with_pool(
+                    &prepared,
+                    Some((pool, key)),
+                    started_at,
+                    self.websocket_first_token_timeout,
+                )
+                .await
+            }
+            _ => {
+                execute_response_create_request_with_pool(
+                    &prepared,
+                    None,
+                    started_at,
+                    self.websocket_first_token_timeout,
+                )
+                .await
             }
-            _ => execute_response_create_request_with_pool(&prepared, None, started_at).await,
         }
         .map_err(websocket_exchange_error_to_client_error)?;
         log_websocket_pool_decision(
@@ -642,9 +664,21 @@ impl CodexBackendClient {
         let pool_log_context = pool_key.as_ref().map(WebSocketPoolLogContext::from_key);
         let exchange = match (self.websocket_pool.as_deref(), pool_key) {
             (Some(pool), Some(key)) => {
-                execute_response_create_request_stream_with_pool(&prepared, Some((pool, key))).await
+                execute_response_create_request_stream_with_pool(
+                    &prepared,
+                    Some((pool, key)),
+                    self.websocket_first_token_timeout,
+                )
+                .await
+            }
+            _ => {
+                execute_response_create_request_stream_with_pool(
+                    &prepared,
+                    None,
+                    self.websocket_first_token_timeout,
+                )
+                .await
             }
-            _ => execute_response_create_request_stream_with_pool(&prepared, None).await,
         }
         .map_err(websocket_exchange_error_to_client_error)?;
         log_websocket_pool_decision(

backend/src/upstream/transport/websocket.rs

@@ -518,9 +518,16 @@ pub(crate) async fn execute_response_create_request_with_pool(
     request: &CodexWebSocketRequest,
     pool: Option<(&CodexWebSocketPool, CodexWebSocketPoolKey)>,
     started_at: Instant,
+    fallback_first_token_timeout: Option<Duration>,
 ) -> Result<CodexWebSocketExchange, CodexWebSocketExchangeError> {
     let Some((pool, key)) = pool else {
-        return execute_fresh_response_create_request(request, started_at, None).await;
+        return execute_fresh_response_create_request_with_retries(
+            request,
+            started_at,
+            fallback_first_token_timeout,
+            WEBSOCKET_FIRST_TOKEN_FRESH_RETRY_ATTEMPTS,
+        )
+        .await;
     };
 
     match pool.acquire(&key).await {
@@ -582,10 +589,11 @@ pub(crate) async fn execute_response_create_request_with_pool(
             }
         }
         WebSocketPoolAcquire::Bypass(reason) => {
-            let mut exchange = execute_fresh_response_create_request(
+            let mut exchange = execute_fresh_response_create_request_with_retries(
                 request,
                 started_at,
                 pool.first_token_timeout(),
+                WEBSOCKET_FIRST_TOKEN_FRESH_RETRY_ATTEMPTS,
             )
             .await?;
             exchange.pool_decision = Some(WebSocketPoolDecision::bypass(reason));
@@ -597,9 +605,15 @@ pub(crate) async fn execute_response_create_request_with_pool(
 pub(crate) async fn execute_response_create_request_stream_with_pool(
     request: &CodexWebSocketRequest,
     pool: Option<(&CodexWebSocketPool, CodexWebSocketPoolKey)>,
+    fallback_first_token_timeout: Option<Duration>,
 ) -> Result<CodexWebSocketStreamingExchange, CodexWebSocketExchangeError> {
     let Some((pool, key)) = pool else {
-        return execute_fresh_response_create_request_stream(request, None).await;
+        return execute_fresh_response_create_request_stream_with_retries(
+            request,
+            fallback_first_token_timeout,
+            WEBSOCKET_FIRST_TOKEN_FRESH_RETRY_ATTEMPTS,
+        )
+        .await;
     };
 
     match pool.acquire(&key).await {
@@ -662,9 +676,12 @@ pub(crate) async fn execute_response_create_request_stream_with_pool(
             }
         }
         WebSocketPoolAcquire::Bypass(reason) => {
-            let mut exchange =
-                execute_fresh_response_create_request_stream(request, pool.first_token_timeout())
-                    .await?;
+            let mut exchange = execute_fresh_response_create_request_stream_with_retries(
+                request,
+                pool.first_token_timeout(),
+                WEBSOCKET_FIRST_TOKEN_FRESH_RETRY_ATTEMPTS,
+            )
+            .await?;
             exchange.pool_decision = Some(WebSocketPoolDecision::bypass(reason));
             Ok(exchange)
         }
@@ -856,7 +873,7 @@ async fn execute_pooled_response_create_request_stream(
 pub async fn execute_response_create_request(
     request: &CodexWebSocketRequest,
 ) -> Result<CodexWebSocketExchange, CodexWebSocketExchangeError> {
-    execute_response_create_request_with_pool(request, None, Instant::now()).await
+    execute_response_create_request_with_pool(request, None, Instant::now(), None).await
 }
 
 // ---------------------------------------------------------------------------

backend/src/upstream/transport/websocket_pool.rs

@@ -17,7 +17,7 @@ const DEFAULT_MAX_AGE: Duration = Duration::from_mins(55);
 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_FIRST_TOKEN_TIMEOUT: Duration = Duration::from_secs(15);
+pub(crate) const DEFAULT_FIRST_TOKEN_TIMEOUT: Duration = Duration::from_secs(15);
 
 /// WebSocket 连接池 key。
 #[derive(Debug, Clone, PartialEq, Eq, Hash)]

backend/tests/upstream/transport/websocket.rs

@@ -1041,6 +1041,87 @@ async fn codex_backend_client_stream_should_reject_binary_websocket_event() {
     ));
 }
 
+#[tokio::test]
+async fn codex_backend_client_stream_should_retry_without_pool_when_first_content_times_out() {
+    let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
+    let addr = listener.local_addr().unwrap();
+    let accepted_connections = Arc::new(AtomicUsize::new(0));
+    let accepted_connections_for_server = Arc::clone(&accepted_connections);
+    let server = tokio::spawn(async move {
+        let (first_stream, _) = listener.accept().await.unwrap();
+        accepted_connections_for_server.fetch_add(1, Ordering::SeqCst);
+        let mut first_websocket = accept_codex_test_websocket(first_stream).await;
+        let _first_message = first_websocket.next().await.unwrap().unwrap();
+        first_websocket
+            .send(Message::Text(
+                json!({
+                    "type": "response.created",
+                    "response": {
+                        "id": "resp_no_pool_first_token_stalled",
+                        "object": "response"
+                    }
+                })
+                .to_string()
+                .into(),
+            ))
+            .await
+            .unwrap();
+
+        let (second_stream, _) = listener.accept().await.unwrap();
+        accepted_connections_for_server.fetch_add(1, Ordering::SeqCst);
+        let mut second_websocket = accept_codex_test_websocket(second_stream).await;
+        let _second_message = second_websocket.next().await.unwrap().unwrap();
+        second_websocket
+            .send(Message::Text(
+                json!({
+                    "type": "response.output_text.delta",
+                    "delta": "no-pool retry recovered"
+                })
+                .to_string()
+                .into(),
+            ))
+            .await
+            .unwrap();
+        second_websocket
+            .send(Message::Text(
+                completed_websocket_response("resp_no_pool_retry_recovered", 3, 1).into(),
+            ))
+            .await
+            .unwrap();
+        second_websocket.close(None).await.unwrap();
+        drop(first_websocket);
+    });
+    let backend = CodexBackendClient::new(
+        reqwest::Client::builder().no_proxy().build().unwrap(),
+        format!("http://{addr}"),
+        crate::support::fingerprint::test_fingerprint(),
+    )
+    .with_websocket_first_token_timeout(Some(Duration::from_millis(30)));
+    let request = pooled_websocket_request("conversation-first-token-no-pool-timeout");
+
+    let response = backend
+        .create_response_stream(
+            &request,
+            request_context("req_first_token_no_pool_timeout", Some("chatgpt-account")),
+        )
+        .await
+        .expect("no-pool first-token timeout should retry before returning stream");
+    let decision = response.websocket_pool_decision;
+    let mut stream = response.body;
+    let mut body = String::new();
+    while let Some(chunk) = stream.next().await {
+        let chunk = chunk.expect("retried no-pool stream chunk should be valid");
+        body.push_str(std::str::from_utf8(&chunk).unwrap());
+    }
+    server.await.unwrap();
+
+    assert!(body.contains("no-pool retry recovered"));
+    assert!(body.contains("resp_no_pool_retry_recovered"));
+    assert!(!body.contains("resp_no_pool_first_token_stalled"));
+    assert!(decision.is_none());
+    assert_eq!(accepted_connections.load(Ordering::SeqCst), 2);
+}
+
 #[tokio::test]
 async fn codex_backend_client_stream_should_error_when_websocket_closes_before_terminal() {
     let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();

backend/tests/upstream/transport/websocket_pool.rs

@@ -341,6 +341,94 @@ async fn codex_backend_client_stream_should_retry_reused_socket_when_first_conte
     assert_eq!(accepted_connections.load(Ordering::SeqCst), 2);
 }
 
+#[tokio::test]
+async fn codex_backend_client_stream_should_retry_bypass_socket_when_first_content_times_out() {
+    let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
+    let addr = listener.local_addr().unwrap();
+    let accepted_connections = Arc::new(AtomicUsize::new(0));
+    let accepted_connections_for_server = Arc::clone(&accepted_connections);
+    let server = tokio::spawn(async move {
+        let (first_stream, _) = listener.accept().await.unwrap();
+        accepted_connections_for_server.fetch_add(1, Ordering::SeqCst);
+        let mut first_websocket = accept_codex_test_websocket(first_stream).await;
+        let _first_message = first_websocket.next().await.unwrap().unwrap();
+        first_websocket
+            .send(Message::Text(
+                json!({
+                    "type": "response.created",
+                    "response": {
+                        "id": "resp_bypass_first_token_stalled",
+                        "object": "response"
+                    }
+                })
+                .to_string()
+                .into(),
+            ))
+            .await
+            .unwrap();
+
+        let (second_stream, _) = listener.accept().await.unwrap();
+        accepted_connections_for_server.fetch_add(1, Ordering::SeqCst);
+        let mut second_websocket = accept_codex_test_websocket(second_stream).await;
+        let _second_message = second_websocket.next().await.unwrap().unwrap();
+        second_websocket
+            .send(Message::Text(
+                json!({
+                    "type": "response.output_text.delta",
+                    "delta": "bypass retry recovered"
+                })
+                .to_string()
+                .into(),
+            ))
+            .await
+            .unwrap();
+        second_websocket
+            .send(Message::Text(
+                completed_websocket_response("resp_bypass_retry_recovered", 3, 1).into(),
+            ))
+            .await
+            .unwrap();
+        second_websocket.close(None).await.unwrap();
+        drop(first_websocket);
+    });
+    let pool = Arc::new(CodexWebSocketPool::with_config(CodexWebSocketPoolConfig {
+        max_per_account: 0,
+        first_token_timeout: Some(Duration::from_millis(30)),
+        ..websocket_pool_config_for_tests(None, None, None)
+    }));
+    let backend = CodexBackendClient::new(
+        reqwest::Client::builder().no_proxy().build().unwrap(),
+        format!("http://{addr}"),
+        crate::support::fingerprint::test_fingerprint(),
+    )
+    .with_websocket_pool(Arc::clone(&pool));
+    let request = pooled_websocket_request("conversation-first-token-bypass-timeout");
+
+    let response = backend
+        .create_response_stream(
+            &request,
+            request_context("req_first_token_bypass_timeout", Some("chatgpt-account")),
+        )
+        .await
+        .expect("bypass first-token timeout should retry before returning stream");
+    let decision = response.websocket_pool_decision;
+    let mut stream = response.body;
+    let mut body = String::new();
+    while let Some(chunk) = stream.next().await {
+        let chunk = chunk.expect("retried bypass stream chunk should be valid");
+        body.push_str(std::str::from_utf8(&chunk).unwrap());
+    }
+    server.await.unwrap();
+
+    assert!(body.contains("bypass retry recovered"));
+    assert!(body.contains("resp_bypass_retry_recovered"));
+    assert!(!body.contains("resp_bypass_first_token_stalled"));
+    let decision = decision.expect("bypass decision");
+    assert_eq!(decision.kind(), "bypass");
+    assert_eq!(decision.reason(), Some("cap"));
+    assert_eq!(accepted_connections.load(Ordering::SeqCst), 2);
+}
+
 #[tokio::test]
 async fn codex_backend_client_should_not_reuse_pooled_websocket_across_local_accounts() {
     let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();