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 <noreply@openai.com>

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(

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(

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;
         }

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();
+}

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};
 

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;
+}

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}"),
                 }
             }
         })

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(