diff --git a/backend/src/runtime/services.rs b/backend/src/runtime/services.rs index 7be83b8ab..ba37c64bb 100644 --- a/backend/src/runtime/services.rs +++ b/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 { diff --git a/backend/src/upstream/transport/client.rs b/backend/src/upstream/transport/client.rs index 453282d64..64f7a9512 100644 --- a/backend/src/upstream/transport/client.rs +++ b/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>, + websocket_first_token_timeout: Option, } 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) -> Self { + self.websocket_first_token_timeout = timeout.filter(|timeout| !timeout.is_zero()); + self + } + /// 为 Responses WebSocket 请求启用连接池。 pub fn with_websocket_pool(mut self, pool: Arc) -> 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( diff --git a/backend/src/upstream/transport/websocket.rs b/backend/src/upstream/transport/websocket.rs index c3eec22f9..9fd87b7e5 100644 --- a/backend/src/upstream/transport/websocket.rs +++ b/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, ) -> Result { 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, ) -> Result { 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 { - execute_response_create_request_with_pool(request, None, Instant::now()).await + execute_response_create_request_with_pool(request, None, Instant::now(), None).await } // --------------------------------------------------------------------------- diff --git a/backend/src/upstream/transport/websocket_pool.rs b/backend/src/upstream/transport/websocket_pool.rs index e726b0890..9103a6c26 100644 --- a/backend/src/upstream/transport/websocket_pool.rs +++ b/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)] diff --git a/backend/tests/upstream/transport/websocket.rs b/backend/tests/upstream/transport/websocket.rs index 2b97cfbe2..e740b3b1c 100644 --- a/backend/tests/upstream/transport/websocket.rs +++ b/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(); diff --git a/backend/tests/upstream/transport/websocket_pool.rs b/backend/tests/upstream/transport/websocket_pool.rs index e93e71967..c7fc7d372 100644 --- a/backend/tests/upstream/transport/websocket_pool.rs +++ b/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();