fix: complete WebSocket shutdown handshake

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