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(