diff --git a/.gitignore b/.gitignore index 0d71fa9..4e15164 100644 --- a/.gitignore +++ b/.gitignore @@ -14,6 +14,9 @@ build/ # Rust target/ +# Local worktrees +.worktrees/ + # IDE .idea/ .vscode/ diff --git a/crates/agentic-server/src/handler/websocket/responses.rs b/crates/agentic-server/src/handler/websocket/responses.rs index cf46efe..b5264a1 100644 --- a/crates/agentic-server/src/handler/websocket/responses.rs +++ b/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( + shutdown_token: &CancellationToken, + receiver: &mut Receiver, +) -> Option +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(shutdown_token: &CancellationToken, value: T) -> Option { + (!shutdown_token.is_cancelled()).then_some(value) +} + +async fn close_ws(sender: &mut Sender, receiver: &mut Receiver) +where + Sender: Sink + Unpin, + Receiver: Stream> + 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> { + self.shutdown_token.cancel(); + Poll::Ready(self.item.take()) + } + } - use super::{ShutdownInput, next_shutdown_input}; + impl Sink for CloseErrorSink { + type Error = &'static str; + + fn poll_ready(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll> { + 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> { + Poll::Ready(Ok(())) + } + + fn poll_close(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll> { + 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::("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() { diff --git a/crates/agentic-server/tests/responses_websocket_test.rs b/crates/agentic-server/tests/responses_websocket_test.rs index e0f6fff..df8bc1c 100644 --- a/crates/agentic-server/tests/responses_websocket_test.rs +++ b/crates/agentic-server/tests/responses_websocket_test.rs @@ -344,6 +344,39 @@ async fn recv_close_or_end(ws: &mut WebSocketStream>) } } +async fn recv_clean_close(ws: &mut WebSocketStream>) { + 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>, 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>, 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); } diff --git a/docs/superpowers/plans/2026-07-22-websocket-shutdown-close-race.md b/docs/superpowers/plans/2026-07-22-websocket-shutdown-close-race.md new file mode 100644 index 0000000..bc63cfc --- /dev/null +++ b/docs/superpowers/plans/2026-07-22-websocket-shutdown-close-race.md @@ -0,0 +1,228 @@ +# WebSocket Shutdown Close-Race Implementation Plan + +> **For agentic workers:** REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (`- [ ]`) syntax for tracking. + +**Goal:** Deliver the active WebSocket response before completing a graceful close handshake, without starting requests received after shutdown. + +**Architecture:** Keep the existing cancellation-aware upstream drain loop. Replace the fire-and-forget sink close with a small internal helper that sends the close frame and then consumes receiver frames until the peer closes, ends, or errors; harden the integration test with a ping/pong receipt barrier before releasing its gated upstream. + +**Tech Stack:** Rust 2024, Tokio, Axum WebSockets, `futures`, `tokio-tungstenite`, Cargo test/Clippy/rustfmt. + +## Global Constraints + +- Preserve the existing eight-second gateway drain timeout as the outer bound. +- Do not start inference for any WebSocket request received after cancellation. +- Do not hold a lock or guard across `.await`. +- Keep receive errors during the close handshake at debug level because shutdown is already handling the connection. +- Do not modify unrelated files or existing user worktree changes. + +--- + +### Task 1: Add the close-handshake regression and implementation + +**Files:** +- Modify: `crates/agentic-server/src/handler/websocket/responses.rs:1-390` +- Test: `crates/agentic-server/src/handler/websocket/responses.rs:362-390` +- Test: `crates/agentic-server/tests/responses_websocket_test.rs:1213-1252` + +**Interfaces:** +- Consumes: the existing `WsSender`, `WsReceiver`, cancellation-aware `stream_ws_response`, and WebSocket `Message` types. +- Produces: internal `close_ws(sender: &mut Sender, receiver: &mut Receiver)` behavior used by `responses_ws_loop`. + +- [ ] **Step 1: Write the failing close-handshake unit test and deterministic integration barrier** + +Add `sink` and `StreamExt` to the unit-test imports, import `Message` and the wished-for `close_ws`, then add: + +```rust +#[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")))); +} +``` + +In `test_websocket_shutdown_drains_active_response_before_closing`, after sending the post-shutdown request and before `release.send(())`, add a ping/pong barrier: + +```rust +let barrier = Bytes::from_static(b"shutdown-request-received"); +ws.send(Message::Ping(barrier.clone())).await.unwrap(); +loop { + let message = tokio::time::timeout(std::time::Duration::from_secs(2), ws.next()) + .await + .expect("timed out waiting for shutdown barrier pong") + .expect("websocket should yield a message") + .expect("websocket message should be ok"); + match message { + Message::Pong(payload) => { + assert_eq!(payload, barrier); + break; + } + Message::Ping(_) | Message::Frame(_) => {} + Message::Text(text) => panic!("unexpected response before upstream release: {text}"), + Message::Close(frame) => panic!("websocket closed before upstream release: {frame:?}"), + Message::Binary(_) => panic!("unexpected binary websocket message"), + } +} +``` + +- [ ] **Step 2: Run the new unit test to verify RED** + +Run: + +```bash +cargo test -p agentic-server handler::websocket::responses::tests::close_ws_ignores_late_frames_until_peer_close +``` + +Expected: compilation fails because `close_ws` does not exist yet. This proves the new test requires the missing close-handshake behavior. + +- [ ] **Step 3: Implement the minimal close handshake** + +Import `Sink`, then add the internal helper: + +```rust +async fn close_ws(sender: &mut Sender, receiver: &mut Receiver) +where + Sender: Sink + Unpin, + Receiver: Stream> + Unpin, + SendError: std::fmt::Display, + ReceiveError: std::fmt::Display, +{ + if let Err(error) = sender.close().await { + 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; + } + } + } +} +``` + +Replace the existing `sender.close().await` block at the end of `responses_ws_loop` with: + +```rust +close_ws(&mut sender, &mut receiver).await; +``` + +- [ ] **Step 4: Run focused tests to verify GREEN** + +Run: + +```bash +cargo test -p agentic-server handler::websocket::responses::tests::close_ws_ignores_late_frames_until_peer_close +cargo test -p agentic-server --test responses_websocket_test test_websocket_shutdown_drains_active_response_before_closing -- --exact +cargo test -p agentic-server --test responses_websocket_test +``` + +Expected: the unit test, targeted integration test, and all WebSocket integration tests pass. + +- [ ] **Step 5: Stress the original failure path** + +Run the exact shutdown integration test 100 times, stopping on the first failure: + +```bash +for iteration in $(seq 1 100); do + cargo test -q -p agentic-server --test responses_websocket_test \ + test_websocket_shutdown_drains_active_response_before_closing -- --exact || exit 1 +done +``` + +Expected: 100 successful iterations and no connection reset. + +- [ ] **Step 6: Commit the verified implementation** + +```bash +git add crates/agentic-server/src/handler/websocket/responses.rs \ + crates/agentic-server/tests/responses_websocket_test.rs \ + docs/superpowers/plans/2026-07-22-websocket-shutdown-close-race.md +git commit -s -m "fix: complete websocket shutdown handshake" +``` + +--- + +### Task 2: Verify and review the branch before publishing + +**Files:** +- Review: every file changed from `origin/main` +- Modify only if a requested review reports an actionable issue in the scoped diff. + +**Interfaces:** +- Consumes: the committed implementation from Task 1 and repository CI commands. +- Produces: a reviewed, pushed branch and a draft GitHub pull request targeting `main`. + +- [ ] **Step 1: Run repository verification** + +```bash +cargo fmt -- --check +cargo clippy --all-targets -- -D warnings +cargo test +pre-commit run --all-files +``` + +Expected: every command exits zero with no warnings promoted to errors. + +- [ ] **Step 2: Run the requested gstack pre-landing review** + +Follow `/Users/farceo/.agents/skills/gstack/review/SKILL.md` against `origin/main`. Apply and re-verify mechanical findings; stop for user input only if the review identifies a non-mechanical choice. + +- [ ] **Step 3: Run the requested Claude read-only review** + +```bash +~/.codex/skills/claude-review/scripts/claude-review "focus on WebSocket shutdown races, cancellation safety, close-handshake liveness, and test determinism" +``` + +Expected: no actionable findings. If Claude reports one, address it with a new failing test first and rerun both verification and Claude review. + +- [ ] **Step 4: Re-verify the final reviewed commit** + +```bash +cargo fmt -- --check +cargo clippy --all-targets -- -D warnings +cargo test +pre-commit run --all-files +git diff --check origin/main...HEAD +``` + +Expected: all commands exit zero and the diff contains only the design, plan, worktree ignore, WebSocket implementation, and its tests. + +- [ ] **Step 5: Push and open the draft pull request** + +```bash +git push -u origin codex/fix-websocket-shutdown-race +``` + +Create a draft PR targeting `main` with these required sections: + +```markdown +## Summary + +- complete the WebSocket close handshake after draining an active response during shutdown +- discard post-shutdown frames without starting another request +- synchronize the shutdown regression test with a ping/pong receipt barrier + +## Test Plan + +- `cargo fmt -- --check` +- `cargo clippy --all-targets -- -D warnings` +- `cargo test` +- `pre-commit run --all-files` +- 100 repeated runs of the targeted WebSocket shutdown test +- gstack `/review` +- Claude read-only review focused on shutdown concurrency +``` diff --git a/docs/superpowers/specs/2026-07-21-websocket-shutdown-close-race-design.md b/docs/superpowers/specs/2026-07-21-websocket-shutdown-close-race-design.md new file mode 100644 index 0000000..0b6ca0c --- /dev/null +++ b/docs/superpowers/specs/2026-07-21-websocket-shutdown-close-race-design.md @@ -0,0 +1,46 @@ +# WebSocket Shutdown Close-Race Design + +## Problem + +The WebSocket shutdown integration test intermittently receives `ConnectionReset` while waiting for the active +response's terminal event. The test cancels the gateway, writes a second request, and releases the gated upstream +response without proving that the server has consumed the second frame. If the upstream finishes first, the handler +can initiate shutdown while client data remains unread, and Linux may reset the connection instead of delivering the +queued terminal event and a clean WebSocket close. + +## Desired behavior + +Once shutdown starts, the gateway must: + +1. finish forwarding the active upstream response, including its terminal event; +2. ignore any new response requests rather than starting more inference; +3. initiate a WebSocket close handshake; and +4. consume remaining client frames until the peer acknowledges the close or disconnects. + +The existing gateway-level eight-second drain timeout remains the outer bound for an unresponsive upstream or peer. + +## Design + +Keep the current active-stream drain loop. After it finishes, send a close frame and continue reading the WebSocket +receiver until it yields a peer close, end-of-stream, or receive error. Text, binary, ping, and pong frames received +during this closing phase are discarded; no new request reaches the executor. A receive error during the close +handshake is logged at debug level because the response has already completed and shutdown is in progress. + +Make the integration test deterministic by sending the post-shutdown request followed by a ping, then waiting for the +matching pong before releasing the upstream response. WebSocket frame ordering means the pong proves the server has +already consumed and rejected the preceding request. The test then asserts that the original response completes, the +connection closes, and the mock upstream observed exactly one request. + +## Alternatives considered + +- **Test-only synchronization:** smallest change, but leaves production clients exposed to the same late-frame reset + window. +- **WebSocket task-ownership redesign:** could centralize shutdown state, but is unnecessary for this localized close + sequencing issue. + +## Verification + +- Demonstrate the close-handshake regression test failing before the production change. +- Run the targeted shutdown test repeatedly after the fix. +- Run the complete WebSocket integration-test binary. +- Run formatting and Clippy for the affected workspace targets.