fix: complete WebSocket shutdown handshake (vllm-project#141)

## Summary

- Complete the WebSocket close handshake after draining an active
response during shutdown.
- Reject post-shutdown WebSocket requests before they can start
inference.
- Make the shutdown regression deterministic with a ping/pong receipt
barrier and strict close acknowledgment.

This fixes the Linux CI race observed at caf6bf1 where unread client
frames could make socket teardown surface ECONNRESET after the active
response completed.

## Test Plan

- `cargo fmt -- --check`
- `cargo clippy --all-targets -- -D warnings`
- `cargo test`
- `uvx pre-commit run --all-files`
- 100 repeated runs of
`test_websocket_shutdown_drains_active_response_before_closing`
- gstack `/review` with specialist, adversarial, and structured
re-review gates
- Claude read-only worktree review

---------

Signed-off-by: Francisco Javier Arceo <farceo@redhat.com>

Author: franciscojavierarceo

.gitignore

@@ -14,6 +14,9 @@ build/
 # Rust
 target/
 
+# Local worktrees
+.worktrees/
+
 # IDE
 .idea/
 .vscode/

crates/agentic-server/src/handler/websocket/responses.rs

@@ -7,7 +7,7 @@ use axum::http::HeaderMap;
 use axum::response::Response;
 use either::Either;
 use futures::stream::{SplitSink, SplitStream};
-use futures::{SinkExt, Stream, StreamExt};
+use futures::{Sink, SinkExt, Stream, StreamExt};
 use serde_json::Value;
 use tokio_util::sync::CancellationToken;
 use tracing::{debug, warn};
@@ -49,10 +49,7 @@ async fn responses_ws_loop(socket: WebSocket, state: AppState, headers: HeaderMa
         let text = if let Some(buffered) = queue.pop_front() {
             buffered
         } else {
-            let message = tokio::select! {
-                () = shutdown_token.cancelled() => break,
-                message = receiver.next() => message,
-            };
+            let message = next_ws_message(&shutdown_token, &mut receiver).await;
 
             let Some(message) = message else {
                 break;
@@ -100,10 +97,56 @@ async fn responses_ws_loop(socket: WebSocket, state: AppState, headers: HeaderMa
             }
         }
     }
+    close_ws(&mut sender, &mut receiver).await;
+    debug!("responses websocket session closed");
+}
+
+async fn next_ws_message<Receiver>(
+    shutdown_token: &CancellationToken,
+    receiver: &mut Receiver,
+) -> Option<Receiver::Item>
+where
+    Receiver: Stream + Unpin,
+{
+    tokio::select! {
+        biased;
+        () = shutdown_token.cancelled() => None,
+        message = receiver.next() => {
+            if shutdown_token.is_cancelled() {
+                None
+            } else {
+                message
+            }
+        },
+    }
+}
+
+fn keep_if_running<T>(shutdown_token: &CancellationToken, value: T) -> Option<T> {
+    (!shutdown_token.is_cancelled()).then_some(value)
+}
+
+async fn close_ws<Sender, Receiver, SendError, ReceiveError>(sender: &mut Sender, receiver: &mut Receiver)
+where
+    Sender: Sink<Message, Error = SendError> + Unpin,
+    Receiver: Stream<Item = Result<Message, ReceiveError>> + Unpin,
+    SendError: std::fmt::Display,
+    ReceiveError: std::fmt::Display,
+{
     if let Err(error) = sender.close().await {
-        debug!(%error, "failed to close responses websocket cleanly");
+        debug!(%error, "failed to send responses websocket close frame");
+        return;
+    }
+
+    while let Some(message) = receiver.next().await {
+        match message {
+            Ok(Message::Close(_)) => break,
+            Ok(Message::Text(_) | Message::Binary(_) | Message::Ping(_) | Message::Pong(_)) => {}
+            Err(error) => {
+                debug!(%error, "responses websocket close handshake receive failed");
+                break;
+            }
+        }
     }
-    debug!("responses websocket session closed");
 }
 
 /// Process one `response.create` message.
@@ -153,6 +196,10 @@ async fn handle_ws_text(
         .with_auth(auth)
         .run()
         .await?;
+    let Some(result) = keep_if_running(shutdown_token, result) else {
+        debug!("discarded websocket response initialized during shutdown");
+        return Ok(());
+    };
     let Either::Right(stream) = result else {
         return Err(WsError::Executor(ExecutorError::InvalidRequest(
             "websocket response.create must produce a stream".to_owned(),
@@ -361,9 +408,116 @@ async fn send_ws_json(sender: &mut WsSender, value: Value) -> Result<(), WsError
 
 #[cfg(test)]
 mod tests {
-    use futures::stream;
+    use std::pin::Pin;
+    use std::task::{Context, Poll};
+
+    use axum::extract::ws::Message;
+    use futures::{Sink, Stream, StreamExt, sink, stream};
+    use tokio_util::sync::CancellationToken;
+
+    use super::{ShutdownInput, close_ws, keep_if_running, next_shutdown_input, next_ws_message};
+
+    struct CloseErrorSink;
+
+    struct CancellingStream {
+        shutdown_token: CancellationToken,
+        item: Option<&'static str>,
+    }
+
+    impl Stream for CancellingStream {
+        type Item = &'static str;
+
+        fn poll_next(mut self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
+            self.shutdown_token.cancel();
+            Poll::Ready(self.item.take())
+        }
+    }
 
-    use super::{ShutdownInput, next_shutdown_input};
+    impl Sink<Message> for CloseErrorSink {
+        type Error = &'static str;
+
+        fn poll_ready(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
+            Poll::Ready(Ok(()))
+        }
+
+        fn start_send(self: Pin<&mut Self>, _item: Message) -> Result<(), Self::Error> {
+            Ok(())
+        }
+
+        fn poll_flush(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
+            Poll::Ready(Ok(()))
+        }
+
+        fn poll_close(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
+            Poll::Ready(Err("close failed"))
+        }
+    }
+
+    #[tokio::test]
+    async fn cancelled_shutdown_wins_over_ready_websocket_message() {
+        let shutdown_token = CancellationToken::new();
+        shutdown_token.cancel();
+        let mut receiver = stream::iter(["must remain unread"]);
+
+        assert!(next_ws_message(&shutdown_token, &mut receiver).await.is_none());
+        assert_eq!(receiver.next().await, Some("must remain unread"));
+    }
+
+    #[tokio::test]
+    async fn cancellation_during_receive_discards_websocket_message() {
+        let shutdown_token = CancellationToken::new();
+        let mut receiver = CancellingStream {
+            shutdown_token: shutdown_token.clone(),
+            item: Some("must be discarded"),
+        };
+
+        assert!(next_ws_message(&shutdown_token, &mut receiver).await.is_none());
+        assert!(shutdown_token.is_cancelled());
+        assert_eq!(receiver.next().await, None);
+    }
+
+    #[test]
+    fn cancellation_after_request_setup_discards_unpolled_stream() {
+        let shutdown_token = CancellationToken::new();
+        shutdown_token.cancel();
+
+        assert_eq!(keep_if_running(&shutdown_token, "unpolled stream"), None);
+    }
+
+    #[tokio::test]
+    async fn close_ws_ignores_late_frames_until_peer_close() {
+        let mut sender = sink::drain();
+        let mut receiver = stream::iter([
+            Ok::<_, &'static str>(Message::Text("late request".into())),
+            Ok(Message::Binary(vec![1].into())),
+            Ok(Message::Close(None)),
+            Err("must remain unread"),
+        ]);
+
+        close_ws(&mut sender, &mut receiver).await;
+
+        assert!(matches!(receiver.next().await, Some(Err("must remain unread"))));
+    }
+
+    #[tokio::test]
+    async fn close_ws_returns_without_reading_when_close_send_fails() {
+        let mut sender = CloseErrorSink;
+        let mut receiver = stream::iter([Ok::<_, &'static str>(Message::Close(None))]);
+
+        close_ws(&mut sender, &mut receiver).await;
+
+        assert!(matches!(receiver.next().await, Some(Ok(Message::Close(None)))));
+    }
+
+    #[tokio::test]
+    async fn close_ws_stops_reading_after_receive_error() {
+        let mut sender = sink::drain();
+        let mut receiver = stream::iter([Err::<Message, _>("receive failed"), Ok(Message::Close(None))]);
+
+        close_ws(&mut sender, &mut receiver).await;
+
+        assert!(matches!(receiver.next().await, Some(Ok(Message::Close(None)))));
+    }
 
     #[tokio::test]
     async fn shutdown_input_priority_alternates_when_both_streams_are_ready() {

crates/agentic-server/tests/responses_websocket_test.rs

@@ -344,6 +344,39 @@ async fn recv_close_or_end(ws: &mut WebSocketStream<MaybeTlsStream<TcpStream>>)
     }
 }
 
+async fn recv_clean_close(ws: &mut WebSocketStream<MaybeTlsStream<TcpStream>>) {
+    let message = tokio::time::timeout(std::time::Duration::from_secs(2), ws.next())
+        .await
+        .expect("timed out waiting for clean websocket close");
+    match message {
+        Some(Ok(Message::Close(_))) => ws.flush().await.expect("failed to acknowledge websocket close"),
+        None => panic!("websocket ended without a close frame"),
+        Some(Err(error)) => panic!("websocket close failed: {error}"),
+        Some(Ok(message)) => panic!("expected websocket close, got {message:?}"),
+    }
+}
+
+async fn send_ping_and_wait_for_pong(ws: &mut WebSocketStream<MaybeTlsStream<TcpStream>>, payload: Bytes) {
+    ws.send(Message::Ping(payload.clone())).await.unwrap();
+    loop {
+        let message = tokio::time::timeout(std::time::Duration::from_secs(2), ws.next())
+            .await
+            .expect("timed out waiting for websocket pong")
+            .expect("websocket should yield a message")
+            .expect("websocket message should be ok");
+        match message {
+            Message::Pong(actual) => {
+                assert_eq!(actual, payload);
+                break;
+            }
+            Message::Ping(_) | Message::Frame(_) => {}
+            Message::Text(text) => panic!("unexpected text before pong: {text}"),
+            Message::Close(frame) => panic!("websocket closed before pong: {frame:?}"),
+            Message::Binary(_) => panic!("unexpected binary websocket message"),
+        }
+    }
+}
+
 async fn send_json(ws: &mut WebSocketStream<MaybeTlsStream<TcpStream>>, value: Value) {
     ws.send(Message::Text(value.to_string().into())).await.unwrap();
 }
@@ -1231,25 +1264,7 @@ async fn test_websocket_ping_returns_pong_without_upstream_request() {
     let (gateway_url, _gateway) = spawn_gateway(fixture.state.clone()).await;
     let mut ws = connect_responses_ws(&gateway_url).await;
 
-    ws.send(Message::Ping(Bytes::from_static(b"ping"))).await.unwrap();
-
-    loop {
-        let message = tokio::time::timeout(std::time::Duration::from_secs(2), ws.next())
-            .await
-            .expect("timed out waiting for websocket pong")
-            .expect("websocket should yield a message")
-            .expect("websocket message should be ok");
-        match message {
-            Message::Pong(payload) => {
-                assert_eq!(payload, Bytes::from_static(b"ping"));
-                break;
-            }
-            Message::Ping(_) | Message::Frame(_) => {}
-            Message::Text(text) => panic!("unexpected text websocket message: {text}"),
-            Message::Close(frame) => panic!("websocket closed before pong: {frame:?}"),
-            Message::Binary(_) => panic!("unexpected binary websocket message"),
-        }
-    }
+    send_ping_and_wait_for_pong(&mut ws, Bytes::from_static(b"ping")).await;
 
     assert!(mock.request_bodies().await.is_empty());
 }
@@ -1274,6 +1289,7 @@ async fn test_websocket_shutdown_drains_active_response_before_closing() {
         MockResponsesServer::start_gated(sse_response("resp_upstream_shutdown", "msg_upstream_shutdown", "DONE")).await;
     let fixture = storage_backed_state(&mock.url).await;
     let shutdown_token = fixture.state.shutdown_token.clone();
+    let websocket_tracker = fixture.state.websocket_tracker.clone();
     let (gateway_url, _gateway) = spawn_gateway(fixture.state.clone()).await;
     let mut ws = connect_responses_ws(&gateway_url).await;
 
@@ -1302,11 +1318,16 @@ async fn test_websocket_shutdown_drains_active_response_before_closing() {
         }),
     )
     .await;
+    let barrier = Bytes::from_static(b"shutdown-request-received");
+    send_ping_and_wait_for_pong(&mut ws, barrier).await;
     release.send(()).unwrap();
 
     let events = recv_until_completed(&mut ws).await;
     assert_eq!(events.last().unwrap()["type"], "response.completed");
-    recv_close_or_end(&mut ws).await;
+    recv_clean_close(&mut ws).await;
+    tokio::time::timeout(std::time::Duration::from_secs(2), websocket_tracker.wait_until_idle())
+        .await
+        .expect("server did not receive the websocket close acknowledgement");
     assert_eq!(mock.request_bodies().await.len(), 1);
 }