Cancel WebSocket handlers on disconnect

[humancto]

Summary

• Wires WebSocket connections to a connection-scoped cancellation token so long-running message handlers stop when the client closes.
• Splits socket receive/send handling so disconnects are still observed while Forge handler code runs on the blocking pool.
• Adds an integration regression test that proves a running WS handler stops making progress after client disconnect.

Closes #111.

Test plan

• cargo fmt
• cargo test --test server_concurrency websocket_handler_cancelled_on_client_disconnect
• cargo test --test server_concurrency
• cargo test

Made with Cursor

Summary by CodeRabbit

• Bug Fixes

  • WebSocket handlers now properly observe client disconnections and cancel long-running operations, ensuring handlers terminate at the next safe point rather than continuing after the client disconnects.

• Tests

  • Added integration tests verifying that WebSocket handler cancellation properly occurs when clients disconnect, including validation that operations cease within expected timeouts.

[coderabbitai]

Caution

Review failed

Pull request was closed or merged during review

📝 Walkthrough

Walkthrough

PR implements WebSocket handler cancellation aligned with HTTP behavior introduced in PR #108. Per-connection cancellation tokens wire into forked interpreters. A dedicated receiver task monitors socket and signals disconnect. The main loop runs handlers on the blocking pool, respecting cancellation at safe points.

Changes

WebSocket Handler Cancellation

Layer / File(s)	Summary
Documentation & Planning .planning/ws-disconnect-cancellation.plan.md, CHANGELOG.md	Plan document specifies per-connection cancellation token architecture, receiver/main-loop split, and blocking-pool integration. Changelog entry notes fixed WebSocket handler cancellation on client disconnect.
Handler Execution Helper src/runtime/server.rs (lines 242–253)	New call_ws_handler(interp, handler_name, text) -> String helper centralizes Forge handler lookup, execution, and error-to-string formatting.
WebSocket Receiver & Main Loop Refactor src/runtime/server.rs (lines 451–565, 59)	Socket is split into recv/send halves. Spawned receiver task forwards Message::Text frames through bounded mpsc channel and sets cancellation flag on close/error. Main loop spawns spawn_blocking tasks that run handlers within tracing span, lock interpreter, and call helper. Receiver aborted on exit. Imports updated for StreamExt and SinkExt.
Integration Tests & Helpers tests/server_concurrency.rs (lines 23–26, 87–109, 272–274, 337–444)	New unique_temp_file() and wait_for_path() helpers support test infrastructure. Refactored schedule_mutations_do_not_leak_into_handler_forks to use new temp-file helper. New websocket_handler_cancelled_on_client_disconnect test verifies that closing the client connection stops a long-running handler loop and prevents completion sentinel from being written.

Sequence Diagram

sequenceDiagram
    participant Client
    participant Receiver as Receiver Task
    participant Channel as Message Channel
    participant MainLoop as Main Loop
    participant BlockPool as Blocking Pool
    participant Interp as Interpreter

    rect rgba(100, 150, 255, 0.5)
    Note over Client,Interp: WebSocket Connection Established
    Client->>Receiver: WebSocket upgrade
    Receiver->>Interp: Fork interpreter, set Arc<AtomicBool>
    end

    rect rgba(100, 200, 150, 0.5)
    Note over Client,Interp: Message Processing Loop
    loop Active Connection
        Receiver->>Receiver: Poll socket for Text/Close/Error
        Receiver->>Channel: try_send(text_message)
        MainLoop->>Channel: recv message
        MainLoop->>BlockPool: spawn_blocking task
        BlockPool->>Interp: enter span, lock interp
        BlockPool->>Interp: call_ws_handler(text)
        Interp-->>BlockPool: response
        BlockPool->>Client: send response
    end
    end

    rect rgba(255, 150, 100, 0.5)
    Note over Client,Interp: Client Disconnect
    Client->>Receiver: Close frame or drop
    Receiver->>Interp: set Arc<AtomicBool> (cancelled = true)
    Receiver->>Receiver: propagate disconnect
    MainLoop->>BlockPool: running task detects cancelled
    BlockPool->>Interp: exit handler at safe point
    MainLoop->>Receiver: abort receiver task
    end

Loading

Estimated Code Review Effort

🎯 4 (Complex) | ⏱️ ~45 minutes

Poem

🐰 A hop, a skip, and a clean disconnect!
WebSockets now bow to the cancellation check—
Long loops must exit when clients depart,
No more blocked interpreters, just a swift heart. ✨

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 60.00% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly and concisely summarizes the main change: adding cancellation support for WebSocket handlers when clients disconnect.
Description check	✅ Passed	The description covers the key aspects (summary, changes, and testing), though it lacks explicit detail on some implementation aspects mentioned in the plan.
Linked Issues check	✅ Passed	All three objectives from issue #111 are met: per-connection interpreter receives Arc, Drop guard sets cancellation flag on disconnect, and Message::Close is respected.
Out of Scope Changes check	✅ Passed	All changes directly support the linked issue objectives; test helpers and changelog updates are appropriately scoped.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing Touches

📝 Generate docstrings

• Create stacked PR
• Commit on current branch

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Commit unit tests in branch feat/ws-disconnect-cancellation

---

_{Review rate limit: 8/10 reviews remaining, refill in 6 minutes and 7 seconds.}

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: dcdcfcdc8f

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

[chatgpt-codex-connector]

Avoid dropping WS connections on brief message bursts

Using try_send on a channel of capacity 1 makes the server treat a temporarily full queue as a disconnect and immediately set the cancel flag. In practice, a client that sends two text frames back-to-back (or before the async loop has polled text_rx) can have its connection cancelled even though the server could have processed both sequentially; this is a behavioral regression from the previous loop that handled all incoming messages in order.

Useful? React with 👍 / 👎.