From 5cfc6dee104d0577a62fefe721daec859e2eb948 Mon Sep 17 00:00:00 2001 From: Zyy Date: Sat, 24 Oct 2026 18:45:04 +0800 Subject: [PATCH] fix(openai): abort stalled websocket pumps during error cleanup Stop the send pump before returning terminal errors or discarding an exchange. Preserve socket cleanup ordering and cover stalled data and control writes. Co-authored-by: Codex --- .../src/transport/websocket/coordinator.rs | 4 +- .../transport/websocket/exchange/stream.rs | 2 +- .../openai/src/transport/websocket/pump.rs | 9 +- .../contract/response_interrupt/cleanup.rs | 106 ++++++++++++++++ .../mod.rs} | 2 + .../tests/transport/websocket/cleanup.rs | 113 ++++++++++++++++++ .../openai/tests/transport/websocket/mod.rs | 10 +- .../openai/tests/transport/websocket_pool.rs | 13 +- 8 files changed, 247 insertions(+), 12 deletions(-) create mode 100644 backend/crates/providers/openai/tests/provider/contract/response_interrupt/cleanup.rs rename backend/crates/providers/openai/tests/provider/contract/{response_interrupt.rs => response_interrupt/mod.rs} (99%) create mode 100644 backend/crates/providers/openai/tests/transport/websocket/cleanup.rs diff --git a/backend/crates/providers/openai/src/transport/websocket/coordinator.rs b/backend/crates/providers/openai/src/transport/websocket/coordinator.rs index 74afc9024..6049a2bae 100644 --- a/backend/crates/providers/openai/src/transport/websocket/coordinator.rs +++ b/backend/crates/providers/openai/src/transport/websocket/coordinator.rs @@ -394,7 +394,7 @@ fn start_pooled_websocket_connect( ); } Err(connection) => { - connection.websocket.close().await; + connection.websocket.abort_and_wait().await; let _ = sender.send(Err(continuation_unavailable( PreviousResponseUnavailableReason::PoolUnavailable, ))); @@ -683,7 +683,7 @@ async fn discard_after_send( if let Some(lease) = lease { lease.discard_with_observation(observation).await; } - websocket.close().await; + websocket.abort_and_wait().await; } pub(crate) fn post_send_ambiguous( 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 55aa538d0..6da84fd67 100644 --- a/backend/crates/providers/openai/src/transport/websocket/exchange/stream.rs +++ b/backend/crates/providers/openai/src/transport/websocket/exchange/stream.rs @@ -470,7 +470,7 @@ async fn discard_stream_websocket( .discard_with_observation(tombstone_observation) .await; } - websocket.close().await; + websocket.abort_and_wait().await; } fn connection_observation( diff --git a/backend/crates/providers/openai/src/transport/websocket/pump.rs b/backend/crates/providers/openai/src/transport/websocket/pump.rs index db4e2f816..44a2a5e58 100644 --- a/backend/crates/providers/openai/src/transport/websocket/pump.rs +++ b/backend/crates/providers/openai/src/transport/websocket/pump.rs @@ -405,10 +405,15 @@ impl PumpedWebSocket { .observation(self.connection_id, Instant::now()) } - pub(crate) async fn close_and_wait(mut self) { + pub(crate) async fn close_and_wait(self) { self.close().await; + self.abort_and_wait().await; + } + + /// 异常退出不等待 Close 发送，避免已阻塞的 pump 掩盖原始错误 + pub(crate) async fn abort_and_wait(mut self) { if let Some(pump) = self.pump.as_mut() { - // 沿用 Drop 的停止语义，并等待 socket 析构完成后再让容量等待者继续 + // 只等待任务析构，确保先释放物理 socket，再归还代理名额 pump.abort(); let _ = pump.await; } diff --git a/backend/crates/providers/openai/tests/provider/contract/response_interrupt/cleanup.rs b/backend/crates/providers/openai/tests/provider/contract/response_interrupt/cleanup.rs new file mode 100644 index 000000000..4392e8913 --- /dev/null +++ b/backend/crates/providers/openai/tests/provider/contract/response_interrupt/cleanup.rs @@ -0,0 +1,106 @@ +//! 验证控制帧写入阻塞时，接收超时和下游取消仍能结束响应并释放连接 + +use super::*; +use crate::transport::accept_codex_test_websocket_with; + +#[tokio::test] +async fn receive_timeout_should_abort_stalled_control_write() { + assert_stalled_control_cleanup(false).await; +} + +#[tokio::test] +async fn cancellation_should_abort_stalled_control_write() { + assert_stalled_control_cleanup(true).await; +} + +async fn assert_stalled_control_cleanup(cancel_response: bool) { + let store = Arc::new(MemoryAccountStore::default()); + create_account(&store, "acct_provider_contract").await; + let socket = tokio::net::TcpSocket::new_v4().unwrap(); + socket.set_recv_buffer_size(8192).unwrap(); + socket.bind("127.0.0.1:0".parse().unwrap()).unwrap(); + let listener = socket.listen(1).unwrap(); + let base_url = format!("http://{}", listener.local_addr().unwrap()); + let (started, started_rx) = oneshot::channel(); + let (drain, drain_rx) = oneshot::channel(); + let server = tokio::spawn(async move { + let (socket, _) = listener.accept().await.unwrap(); + // 不协商压缩，保证控制帧足够大，能够填满 TCP 发送缓冲 + let mut ws = accept_codex_test_websocket_with(socket, |_, response| { + response.headers_mut().remove("sec-websocket-extensions"); + }) + .await; + assert_eq!(next_json(&mut ws).await["type"], "response.create"); + // 直接写入未压缩的文本帧，与握手中拒绝扩展的结果一致 + let delta = br#"{"type":"response.output_text.delta","output_index":0,"content_index":0,"delta":"partial"}"#; + assert!(delta.len() < 126); + ws.get_mut() + .write_all(&[0x81, u8::try_from(delta.len()).unwrap()]) + .await + .unwrap(); + ws.get_mut().write_all(delta).await.unwrap(); + ws.get_mut().read_u8().await.unwrap(); + started.send(()).unwrap(); + drain_rx.await.unwrap(); + tokio::io::copy(ws.get_mut(), &mut tokio::io::sink()) + .await + .unwrap(); + }); + let control = ResponseControl::default(); + let mut stream = provider_with_base_url(&store, base_url) + .execute( + planned_request("openai", generate_operation()), + interrupt_context(control.clone(), None), + ) + .await + .unwrap(); + loop { + if stream.next().await.unwrap().unwrap().has_client_event() { + break; + } + } + let payload = + json!({"type": "future.control", "data": "x".repeat(16 * 1024 * 1024)}).to_string(); + let pending = control.send(&payload); + tokio::pin!(pending); + tokio::select! { + result = &mut pending => panic!("control write should stall: {result:?}"), + ready = started_rx => ready.unwrap(), + () = tokio::time::sleep(Duration::from_secs(5)) => panic!("control write should start"), + } + if cancel_response { + drop(stream); + } else { + // 已有输出后才快进空闲时限，错误必须通过流返回，不能被 Close 发送阻塞 + tokio::time::pause(); + tokio::time::advance(Duration::from_secs(301)).await; + let result = timeout(Duration::from_secs(1), async { + loop { + if let Err(error) = stream.next().await.expect("timeout should yield an error") { + break error; + } + } + }) + .await; + tokio::time::resume(); + let error = result.expect("receive timeout must not wait for a stalled control write"); + assert_eq!( + error.diagnostic().unwrap().code(), + Some("receive_idle_timeout") + ); + assert!(error.pre_delivery_retry().is_none()); + drop(stream); + } + assert_eq!( + timeout(Duration::from_secs(2), pending) + .await + .expect("cleanup must stop the pending control write"), + Err(ResponseControlUnavailable) + ); + assert_eq!(control.send("{}").await, Err(ResponseControlUnavailable)); + drain.send(()).unwrap(); + timeout(Duration::from_secs(5), server) + .await + .expect("cleanup must release the physical socket") + .unwrap(); +} diff --git a/backend/crates/providers/openai/tests/provider/contract/response_interrupt.rs b/backend/crates/providers/openai/tests/provider/contract/response_interrupt/mod.rs similarity index 99% rename from backend/crates/providers/openai/tests/provider/contract/response_interrupt.rs rename to backend/crates/providers/openai/tests/provider/contract/response_interrupt/mod.rs index 4bd8810f1..d45e898cf 100644 --- a/backend/crates/providers/openai/tests/provider/contract/response_interrupt.rs +++ b/backend/crates/providers/openai/tests/provider/contract/response_interrupt/mod.rs @@ -1,5 +1,7 @@ //! 验证未知控制帧和迟到中断使用原连接，并在复用时撤销旧执行控制权 +mod cleanup; + use super::*; use gateway_core::engine::response_control::{ResponseControl, ResponseControlUnavailable}; diff --git a/backend/crates/providers/openai/tests/transport/websocket/cleanup.rs b/backend/crates/providers/openai/tests/transport/websocket/cleanup.rs new file mode 100644 index 000000000..1a7d6ec1a --- /dev/null +++ b/backend/crates/providers/openai/tests/transport/websocket/cleanup.rs @@ -0,0 +1,113 @@ +//! 验证上游停止读请求时，发送超时仍能返回并释放连接 + +use super::*; +use std::assert_matches; + +#[tokio::test] +async fn send_timeout_should_abort_stalled_pump_and_allow_next_request() { + let socket = tokio::net::TcpSocket::new_v4().unwrap(); + socket.set_recv_buffer_size(8192).unwrap(); + socket.bind("127.0.0.1:0".parse().unwrap()).unwrap(); + let listener = socket.listen(2).unwrap(); + let addr = listener.local_addr().unwrap(); + let (started, started_rx) = tokio::sync::oneshot::channel(); + let (drain, drain_rx) = tokio::sync::oneshot::channel(); + let server = tokio::spawn(async move { + let (stream, _) = listener.accept().await.unwrap(); + // 不协商压缩，保证大请求填满 TCP 发送缓冲，而不是压缩后一次写完 + let mut websocket = accept_codex_test_websocket_with(stream, |_, response| { + response.headers_mut().remove("sec-websocket-extensions"); + }) + .await; + websocket.get_mut().read_u8().await.unwrap(); + started.send(()).unwrap(); + drain_rx.await.unwrap(); + // 只排空已进入内核的字节；必须看到 EOF，证明客户端已释放物理 socket + tokio::io::copy(websocket.get_mut(), &mut tokio::io::sink()) + .await + .unwrap(); + + let (stream, _) = listener.accept().await.unwrap(); + let mut websocket = accept_codex_test_websocket(stream).await; + assert_matches!(websocket.next().await.unwrap().unwrap(), Message::Text(_)); + websocket + .send(Message::Text( + completed_websocket_response("resp_after_send_timeout", 2, 1).into(), + )) + .await + .unwrap(); + }); + let pool = Arc::new(CodexWebSocketPool::with_config(CodexWebSocketPoolConfig { + ping_interval: None, + liveness_timeout: None, + maintenance_interval: None, + ..CodexWebSocketPoolConfig::default() + })); + let backend = CodexBackendClient::new( + reqwest::Client::builder().no_proxy().build().unwrap(), + format!("http://{addr}"), + test_wire_profile(), + ) + .with_websocket_pool(Arc::clone(&pool)); + let mut request = codex_request_with_prompt_cache_key( + "gpt-test", + "be brief", + vec![json!({"role": "user", "content": "x".repeat(16 * 1024 * 1024)})], + "stalled-send", + ); + request.set_previous_response_id(Some("resp_previous".to_owned())); + request.previous_response_scope = Some(PreviousResponseScope::Persisted); + request.local_conversation_id = Some("stalled-send".to_owned()); + let error = { + let pending = backend.create_response_stream( + &request, + request_context("stalled_send", Some("chatgpt-account")), + ); + tokio::pin!(pending); + tokio::select! { + result = &mut pending => panic!("send should stall: {}", result.is_ok()), + ready = started_rx => ready.unwrap(), + () = tokio::time::sleep(Duration::from_secs(5)) => panic!("payload send should start"), + } + // 握手和首次写入用真实时间，随后仅快进发送 deadline，避免等待五分钟 + tokio::time::pause(); + tokio::time::advance(Duration::from_secs(301)).await; + let result = timeout(Duration::from_secs(1), &mut pending).await; + tokio::time::resume(); + result + .expect("send timeout must not wait for Close on the stalled pump") + .err() + .expect("stalled send should fail") + }; + let CodexClientError::WebSocket(error) = error else { + panic!("expected WebSocket send failure: {error:?}"); + }; + let CodexWebSocketExchangeError::ConnectionObserved { source, .. } = error else { + panic!("send timeout must retain its connection observation: {error:?}"); + }; + let CodexWebSocketExchangeError::PostSendAmbiguous { + source: Some(source), + .. + } = source.as_ref() + else { + panic!("send timeout must retain its post-send ambiguity: {source:?}"); + }; + assert_matches!( + source.as_ref(), + CodexWebSocketExchangeError::SendTimeout { .. } + ); + drain.send(()).unwrap(); + let response = timeout( + Duration::from_secs(5), + backend.create_response( + &pooled_websocket_request("stalled-send"), + request_context("after_stalled_send", Some("chatgpt-account")), + ), + ) + .await + .expect("failed request must release the pool reservation and socket") + .unwrap(); + assert!(response.body.contains("resp_after_send_timeout")); + server.await.unwrap(); + pool.shutdown().await; +} diff --git a/backend/crates/providers/openai/tests/transport/websocket/mod.rs b/backend/crates/providers/openai/tests/transport/websocket/mod.rs index 83ceb1291..ddf5ed895 100644 --- a/backend/crates/providers/openai/tests/transport/websocket/mod.rs +++ b/backend/crates/providers/openai/tests/transport/websocket/mod.rs @@ -7,6 +7,7 @@ use provider_openai::transport::{ }; mod breaker; +mod cleanup; mod diagnostics; mod fingerprint; mod metadata; @@ -1979,10 +1980,13 @@ async fn codex_backend_client_stream_should_cancel_while_inbound_is_backpressure timeout(Duration::from_secs(1), async { loop { match websocket.next().await { - Some(Ok(Message::Close(_))) => return, + None + | Some(Ok(Message::Close(_))) + | Some(Err(tungstenite::Error::Protocol( + tungstenite::error::ProtocolError::ResetWithoutClosingHandshake, + ))) => return, Some(Ok(_)) => continue, - Some(Err(error)) => panic!("websocket close should be graceful: {error}"), - None => panic!("websocket ended without a close frame"), + Some(Err(error)) => panic!("unexpected websocket cleanup error: {error}"), } } }) diff --git a/backend/crates/providers/openai/tests/transport/websocket_pool.rs b/backend/crates/providers/openai/tests/transport/websocket_pool.rs index 3a286e899..e2c6ba05f 100644 --- a/backend/crates/providers/openai/tests/transport/websocket_pool.rs +++ b/backend/crates/providers/openai/tests/transport/websocket_pool.rs @@ -992,10 +992,15 @@ async fn websocket_pool_shutdown_should_cancel_and_join_active_stream() { .unwrap(); let close = timeout(Duration::from_secs(2), websocket.next()) .await - .expect("pool shutdown should close the active websocket") - .expect("shutdown should produce a close frame") - .expect("close frame should be valid"); - std::assert_matches!(close, Message::Close(_)); + .expect("pool shutdown should close the active websocket"); + // 取消活跃响应直接中止 pump，不依赖关闭握手完成 + std::assert_matches!( + close, + None | Some(Ok(Message::Close(_))) + | Some(Err(tungstenite::Error::Protocol( + tungstenite::error::ProtocolError::ResetWithoutClosingHandshake + ))) + ); }); let pool = Arc::new(CodexWebSocketPool::default()); let backend = CodexBackendClient::new(