From 503e47fac28e479dd069c92a6006e62f773c5a73 Mon Sep 17 00:00:00 2001 From: Zyy Date: Sat, 31 Oct 2026 01:09:05 +0800 Subject: [PATCH] fix(openai): restore upstream websocket keepalive pings Send upstream Ping frames every 25 seconds without a Pong deadline. Keep downstream cancellation, business idle timeout and pool eviction independent from heartbeat traffic, with bounded heartbeat writes. Co-authored-by: Codex --- .../transport/websocket/exchange/stream.rs | 5 +- .../openai/src/transport/websocket/pump.rs | 33 +++++++++- .../openai/tests/transport/websocket/mod.rs | 64 ++++++++++++++++--- .../openai/tests/transport/websocket_pool.rs | 56 +++++++++------- deploy/README.md | 3 +- docs/api.md | 2 +- 6 files changed, 126 insertions(+), 37 deletions(-) diff --git a/backend/crates/providers/openai/src/transport/websocket/exchange/stream.rs b/backend/crates/providers/openai/src/transport/websocket/exchange/stream.rs index 4133a6969..3327cbca2 100644 --- a/backend/crates/providers/openai/src/transport/websocket/exchange/stream.rs +++ b/backend/crates/providers/openai/src/transport/websocket/exchange/stream.rs @@ -404,7 +404,10 @@ async fn forward_websocket_response_stream(state: WebSocketStreamForwardState) { reason => CodexWebSocketExchangeError::StreamEndedBeforeTerminal { reason: reason.map_or("stream_eof", PumpExitReason::as_str), timeout: match reason { - Some(PumpExitReason::PongWriteTimeout { timeout }) => Some(*timeout), + Some( + PumpExitReason::PingWriteTimeout { timeout } + | PumpExitReason::PongWriteTimeout { timeout }, + ) => Some(*timeout), _ => None, }, last_event_type, diff --git a/backend/crates/providers/openai/src/transport/websocket/pump.rs b/backend/crates/providers/openai/src/transport/websocket/pump.rs index 68288e2aa..6c0b23a1c 100644 --- a/backend/crates/providers/openai/src/transport/websocket/pump.rs +++ b/backend/crates/providers/openai/src/transport/websocket/pump.rs @@ -1,4 +1,4 @@ -//! 单 owner 驱动上游 WebSocket，被动回复心跳并在终止后有界释放 socket +//! 单 owner 驱动上游 WebSocket，定时保活、回复心跳并在终止后有界释放 socket //! //! 业务消息通过有界队列交付，Ping/Pong 由 pump 消费，不刷新生成期间的消息等待超时 //! 入站背压暂停读取，仍可接收关闭命令；已观察到的终止先释放 socket，再等待终止消息交付 @@ -42,6 +42,8 @@ pub(crate) type RawWsStream = Box; const PUMP_COMMAND_BUFFER: usize = 32; const CLOSE_TIMEOUT: Duration = Duration::from_secs(1); +const PING_INTERVAL: Duration = Duration::from_secs(25); +const PING_WRITE_TIMEOUT: Duration = Duration::from_secs(5); const PONG_WRITE_TIMEOUT: Duration = Duration::from_secs(5); /// pump 日志的账号级关联上下文；空闲连接被上游重置等场景用于回溯归属 @@ -78,6 +80,9 @@ pub(crate) enum PumpExitReason { message: String, metric_reason: &'static str, }, + PingWriteTimeout { + timeout: Duration, + }, PongWriteTimeout { timeout: Duration, }, @@ -97,6 +102,7 @@ impl PumpExitReason { Self::UpstreamCloseFrame { .. } => "upstream_close_frame", Self::UpstreamEof => "upstream_eof", Self::InboundTransportError { .. } => "inbound_transport_error", + Self::PingWriteTimeout { .. } => "ping_write_timeout", Self::PongWriteTimeout { .. } => "pong_write_timeout", Self::KeepaliveTransportError { .. } => "keepalive_transport_error", Self::MessageReceiverClosed => "message_receiver_closed", @@ -109,7 +115,9 @@ impl PumpExitReason { | Self::InboundTransportError { message, .. } | Self::KeepaliveTransportError { message, .. } => Some(message.clone()), Self::UpstreamCloseFrame { close } => close.as_ref().map(ToString::to_string), - Self::PongWriteTimeout { timeout } => Some(format!("{timeout:?}")), + Self::PingWriteTimeout { timeout } | Self::PongWriteTimeout { timeout } => { + Some(format!("{timeout:?}")) + } Self::CommandChannelClosed | Self::LocalClose | Self::UpstreamEof @@ -132,6 +140,7 @@ impl PumpExitReason { } } Self::UpstreamEof => "eof", + Self::PingWriteTimeout { .. } => "ping_write_timeout", Self::PongWriteTimeout { .. } => "pong_write_timeout", Self::MessageReceiverClosed => "message_receiver_closed", } @@ -150,6 +159,7 @@ impl PumpExitReason { Self::OutboundTransportError { .. } | Self::UpstreamEof | Self::InboundTransportError { .. } + | Self::PingWriteTimeout { .. } | Self::PongWriteTimeout { .. } | Self::KeepaliveTransportError { .. } ) @@ -400,6 +410,9 @@ async fn pump_loop( ) { let mut pending_inbound = None; let mut backpressure_events = 0_u64; + // 首次保活留出完整间隔，背压恢复后跳过错过的 tick，避免集中补发 + let mut heartbeat = tokio::time::interval_at(Instant::now() + PING_INTERVAL, PING_INTERVAL); + heartbeat.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip); let reason = 'pump: loop { if pending_inbound.is_some() { @@ -434,6 +447,22 @@ async fn pump_loop( break 'pump reason; } } + _ = heartbeat.tick() => { + // Ping 只保活链路，不设 Pong deadline，也不刷新业务空闲或连接池回收时间 + match tokio::time::timeout( + PING_WRITE_TIMEOUT, + inner.stream.send(Message::Ping(Default::default())), + ).await { + Ok(Ok(())) => {} + Ok(Err(error)) => break 'pump PumpExitReason::KeepaliveTransportError { + message: error.to_string(), + metric_reason: transport_metric_reason(&error), + }, + Err(_) => break 'pump PumpExitReason::PingWriteTimeout { + timeout: PING_WRITE_TIMEOUT, + }, + } + } message = inner.stream.next() => { match message { None => break 'pump PumpExitReason::UpstreamEof, diff --git a/backend/crates/providers/openai/tests/transport/websocket/mod.rs b/backend/crates/providers/openai/tests/transport/websocket/mod.rs index e5c69dcac..6d36e86e9 100644 --- a/backend/crates/providers/openai/tests/transport/websocket/mod.rs +++ b/backend/crates/providers/openai/tests/transport/websocket/mod.rs @@ -1493,10 +1493,12 @@ async fn codex_backend_client_should_timeout_when_upstream_is_silent() { } #[tokio::test] -async fn websocket_stream_should_allow_silence_below_idle_timeout() { +async fn websocket_stream_should_keep_alive_without_requiring_pong_below_idle_timeout() { let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); let addr = listener.local_addr().unwrap(); let (created_tx, created_rx) = tokio::sync::oneshot::channel(); + let (ping_tx, ping_rx) = tokio::sync::oneshot::channel(); + let (finish_tx, finish_rx) = tokio::sync::oneshot::channel(); let server = tokio::spawn(async move { let (stream, _) = listener.accept().await.unwrap(); let mut websocket = accept_codex_test_websocket(stream).await; @@ -1513,7 +1515,10 @@ async fn websocket_stream_should_allow_silence_below_idle_timeout() { .await .unwrap(); created_tx.send(()).unwrap(); - tokio::time::sleep(Duration::from_secs(30)).await; + std::assert_matches!(websocket.next().await.unwrap().unwrap(), Message::Ping(_)); + ping_tx.send(()).unwrap(); + // 暂停读取和 flush，避免 tungstenite 自动回复 Pong + finish_rx.await.unwrap(); websocket .send(Message::Text( completed_websocket_response("resp_structural", 3, 1).into(), @@ -1529,7 +1534,21 @@ async fn websocket_stream_should_allow_silence_below_idle_timeout() { created_rx.await.unwrap(); tokio::task::yield_now().await; tokio::time::pause(); - tokio::time::advance(Duration::from_secs(30)).await; + tokio::time::advance(Duration::from_secs(26)).await; + tokio::time::resume(); + timeout(Duration::from_secs(2), ping_rx) + .await + .unwrap() + .unwrap(); + tokio::time::pause(); + tokio::time::advance(Duration::from_secs(60)).await; + tokio::time::resume(); + tokio::task::yield_now().await; + assert!( + !response_task.is_finished(), + "missing Pong must not end the response" + ); + finish_tx.send(()).unwrap(); let response = response_task .await .expect("websocket task should finish") @@ -1725,6 +1744,12 @@ async fn codex_backend_client_stream_should_keep_socket_after_structural_activit .await .unwrap(); first_websocket.close(None).await.unwrap(); + // 完成关闭握手，避免未读的保活 Ping 导致 TCP reset 丢失已发送的响应 + while let Some(Ok(message)) = first_websocket.next().await { + if matches!(message, Message::Close(_)) { + break; + } + } }); let backend = CodexBackendClient::new( reqwest::Client::builder().no_proxy().build().unwrap(), @@ -1744,7 +1769,9 @@ async fn codex_backend_client_stream_should_keep_socket_after_structural_activit let mut stream = response.body; let mut body = String::new(); while let Some(chunk) = stream.next().await { - let chunk = chunk.expect("delayed websocket stream chunk should be valid"); + let chunk = chunk.unwrap_or_else(|error| { + panic!("delayed websocket stream chunk should be valid: {error}") + }); body.push_str(std::str::from_utf8(&chunk).unwrap()); } server.await.unwrap(); @@ -2087,10 +2114,11 @@ async fn codex_backend_client_stream_should_wait_for_terminal_after_active_webso ); } -#[tokio::test(start_paused = true)] -async fn codex_backend_client_stream_should_timeout_when_active_websocket_stalls() { +#[tokio::test] +async fn codex_backend_client_stream_should_timeout_when_active_websocket_stalls_despite_pong() { let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); let addr = listener.local_addr().unwrap(); + let (pong_tx, mut pong_rx) = tokio::sync::mpsc::channel(1); let server = tokio::spawn(async move { let (stream, _) = listener.accept().await.unwrap(); let mut websocket = accept_codex_test_websocket(stream).await; @@ -2106,7 +2134,16 @@ async fn codex_backend_client_stream_should_timeout_when_active_websocket_stalls )) .await .unwrap(); - futures::future::pending::<()>().await; + while let Some(Ok(message)) = websocket.next().await { + match message { + Message::Ping(_) => { + websocket.flush().await.unwrap(); + pong_tx.send(()).await.unwrap(); + } + Message::Close(_) => break, + other => panic!("unexpected frame during stalled response: {other:?}"), + } + } }); let mut request = codex_request("gpt-5.5", "be brief", Vec::new()); request.set_previous_response_id(Some("resp_previous".to_string())); @@ -2132,7 +2169,18 @@ async fn codex_backend_client_stream_should_timeout_when_active_websocket_stalls .expect("stream should yield partial frame") .expect("partial frame should be valid"); assert!(std::str::from_utf8(&first).unwrap().contains("partial")); - tokio::time::advance(Duration::from_secs(5 * 60)).await; + for _ in 0..11 { + tokio::time::pause(); + tokio::time::advance(Duration::from_secs(25)).await; + tokio::time::resume(); + timeout(Duration::from_secs(2), pong_rx.recv()) + .await + .expect("gateway should send keepalive Ping") + .expect("upstream should reply with Pong"); + } + tokio::time::pause(); + tokio::time::advance(Duration::from_secs(25)).await; + tokio::time::resume(); tokio::task::yield_now().await; let error = response .body diff --git a/backend/crates/providers/openai/tests/transport/websocket_pool.rs b/backend/crates/providers/openai/tests/transport/websocket_pool.rs index 504681f6e..0598771d0 100644 --- a/backend/crates/providers/openai/tests/transport/websocket_pool.rs +++ b/backend/crates/providers/openai/tests/transport/websocket_pool.rs @@ -365,7 +365,7 @@ async fn codex_backend_client_should_open_fresh_socket_when_idle_pooled_websocke .unwrap(); second_websocket.close(None).await.unwrap(); }); - // 无 maintenance、无主动 ping：完全依赖 pump 后台读取感知连接死亡 + // 不启动维护，验证 pump 读取到上游关闭后即可阻止复用 let pool = Arc::new(CodexWebSocketPool::with_config( websocket_pool_config_for_tests(None), )); @@ -1210,11 +1210,10 @@ async fn concurrent_pool_transport_counts(max_connecting: usize) -> (usize, usiz } #[tokio::test] -async fn websocket_pool_should_reject_expired_idle_continuation_despite_upstream_ping() { +async fn websocket_pool_should_reject_expired_idle_continuation_despite_keepalive() { let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); let addr = listener.local_addr().unwrap(); - let (send_ping_tx, send_ping_rx) = tokio::sync::oneshot::channel(); - let (pong_seen_tx, pong_seen_rx) = tokio::sync::oneshot::channel(); + let (ping_seen_tx, ping_seen_rx) = tokio::sync::oneshot::channel(); let server = tokio::spawn(async move { let (stream, _) = listener.accept().await.unwrap(); let mut first = accept_codex_test_websocket(stream).await; @@ -1225,16 +1224,16 @@ async fn websocket_pool_should_reject_expired_idle_continuation_despite_upstream )) .await .unwrap(); - send_ping_rx.await.unwrap(); - let payload = b"upstream-idle-probe".to_vec(); - first - .send(Message::Ping(payload.clone().into())) - .await - .unwrap(); - let pong = first.next().await.unwrap().unwrap(); - assert_eq!(pong, Message::Pong(payload.into())); - pong_seen_tx.send(()).unwrap(); - std::assert_matches!(first.next().await.unwrap().unwrap(), Message::Close(_)); + std::assert_matches!(first.next().await.unwrap().unwrap(), Message::Ping(_)); + first.flush().await.unwrap(); + ping_seen_tx.send(()).unwrap(); + loop { + match first.next().await.unwrap().unwrap() { + Message::Ping(_) => first.flush().await.unwrap(), + Message::Close(_) => break, + other => panic!("unexpected idle frame: {other:?}"), + } + } let (stream, _) = listener.accept().await.unwrap(); let mut second = accept_codex_test_websocket(stream).await; @@ -1269,8 +1268,7 @@ async fn websocket_pool_should_reject_expired_idle_continuation_despite_upstream tokio::time::pause(); tokio::time::advance(Duration::from_secs(299)).await; tokio::time::resume(); - send_ping_tx.send(()).unwrap(); - timeout(Duration::from_secs(2), pong_seen_rx) + timeout(Duration::from_secs(2), ping_seen_rx) .await .unwrap() .unwrap(); @@ -1332,12 +1330,17 @@ async fn websocket_pool_should_gc_expired_idle_connections() { )) .await .unwrap(); - let close = timeout(Duration::from_secs(60), first_websocket.next()) - .await - .expect("gc sweep should close the expired idle websocket") - .expect("gc sweep should send a close frame") - .expect("close frame should be valid"); - std::assert_matches!(close, Message::Close(_)); + timeout(Duration::from_secs(60), async { + loop { + match first_websocket.next().await.unwrap().unwrap() { + Message::Ping(_) => first_websocket.flush().await.unwrap(), + Message::Close(_) => break, + other => panic!("unexpected idle frame: {other:?}"), + } + } + }) + .await + .expect("gc sweep should close the expired idle websocket"); let (second_stream, _) = listener.accept().await.unwrap(); accepted_connections_for_server.fetch_add(1, Ordering::SeqCst); @@ -1791,8 +1794,13 @@ async fn websocket_pool_should_reclaim_idle_connections_on_maintenance_ticks() { )) .await .unwrap(); - let message = websocket.next().await.unwrap().unwrap(); - std::assert_matches!(message, Message::Close(_)); + loop { + match websocket.next().await.unwrap().unwrap() { + Message::Ping(_) => websocket.flush().await.unwrap(), + Message::Close(_) => break, + other => panic!("unexpected idle frame: {other:?}"), + } + } closed_tx.send(()).unwrap(); }); let pool = Arc::new(CodexWebSocketPool::default()); diff --git a/deploy/README.md b/deploy/README.md index a325e05db..5a5868c19 100644 --- a/deploy/README.md +++ b/deploy/README.md @@ -305,7 +305,8 @@ Responses 流式请求完成准入后，每 15 秒无输出会发送一次 SSE `api.request_timeout_seconds` 默认 `null`，只控制 HTTP 路由返回响应前的等待，不是流式正文或 WebSocket 每轮执行的总时限。 排队、插件显式执行期限与客户端断开的边界见 [请求期限](../docs/api.md#请求期限) -OpenAI 上游 WebSocket 被动回应 Ping，不主动发送心跳；Pong 写入最多等待 5 秒。 +OpenAI Responses 上游 WebSocket 每 25 秒发送 Ping 保活，并回应上游 Ping；Ping/Pong 写入最多等待 5 秒， +不因等待 Pong 超时主动断开。入站背压期间暂停主动 Ping，恢复后不集中补发。 下游 Responses WebSocket 每 25 秒发送 Ping，第一条未响应的 Ping 发出后 30 秒没有入站活动即断开； 接收中的消息字节、Pong 和其他有效入站消息均算活动，慢速上传不必等整帧收齐才刷新探活状态。 后续 Ping 不延长这次截止，心跳入队及写入最多等待 5 秒。两段心跳各自处理，不跨 WS 转发或转换为 SSE 事件。 diff --git a/docs/api.md b/docs/api.md index fcd9fbfca..911816965 100644 --- a/docs/api.md +++ b/docs/api.md @@ -466,7 +466,7 @@ Responses WebSocket 接受文本消息，`response.create` 在同一连接串行 留在有界接收队列中，待当前响应完成终结和写出后再逐条校验、准入与执行，不因请求提前到达而断开。 接收队列容量为 32 个事件，超载仍关闭连接；Ping/Pong、客户端关闭和服务关闭不等待队列中的请求执行 -网关分别维护两段链路的心跳：上游 WS 被动回复 Pong，下游 WS 独立探活；Ping/Pong 不跨段转发，也不转换为 SSE 事件。 +网关分别维护两段链路的心跳：上游 WS 定时发送 Ping 保活并回复 Pong，下游 WS 独立探活；Ping/Pong 不跨段转发，也不转换为 SSE 事件。 下游探活包含消息接收中的字节进度，完整消息收齐后才进入中间件与请求处理。 客户端断开时取消当前执行，并释放仍归该客户端所有的上游连接。`response.completed` 只结束本轮响应，保留客户端 WS 等待下一轮请求