refactor: channels ws handler disconnect handling (#2691)

* refactor: channels ws handler disconnect handling

Currently, we call `WebSocket.receive()` to block, however that method doesn't raise `WebSocketDisconnect`, instead it returns the `websocket.disconnect` message.

* Adds test to ensure subscription survives receipt of arbitrary messages.

litestar/channels/plugin.py

@@ -288,11 +288,10 @@ async def _ws_handler_func(self, channel_name: str, socket: WebSocket) -> None:
             if self._handler_should_send_history:
                 await self.put_subscriber_history(subscriber, channels=channel_name, limit=self._history_limit)
 
-            # use the background task, so we can receive(), raising a WebSocketDisconnect
-            # when a connection closes, breaking the loop
+            # use the background task, so we can block on receive(), breaking the loop when a connection closes
             async with subscriber.run_in_background(on_event):
-                while True:
-                    await socket.receive()
+                while (await socket.receive())["type"] != "websocket.disconnect":
+                    continue
 
     def _create_ws_handler_func(self, channel_name: str) -> Callable[[WebSocket], Awaitable[None]]:
         async def ws_handler_func(socket: WebSocket) -> None:

tests/unit/test_channels/test_plugin.py

@@ -118,6 +118,27 @@ def test_create_ws_route_handlers(
         assert ws.receive_json(mode=socket_send_mode, timeout=2) == ["foo"]
 
 
+@pytest.mark.flaky(reruns=5)
+def test_ws_route_handlers_receive_arbitrary_message(channels_backend: ChannelsBackend) -> None:
+    """The websocket handlers await `WebSocket.receive()` to detect disconnection and stop the subscription.
+
+    This test ensures that the subscription is only stopped in the case of receiving a `websocket.disconnect` message.
+    """
+    channels_plugin = ChannelsPlugin(
+        backend=channels_backend,
+        create_ws_route_handlers=True,
+        channels=["something"],
+    )
+    app = Litestar(plugins=[channels_plugin])
+
+    with TestClient(app) as client, client.websocket_connect("/something") as ws:
+        channels_plugin.publish(["foo"], "something")
+        # send some arbitrary message
+        ws.send("bar")
+        # the subscription should still be alive
+        assert ws.receive_json(timeout=2) == ["foo"]
+
+
 @pytest.mark.flaky(reruns=5)
 async def test_create_ws_route_handlers_arbitrary_channels_allowed(channels_backend: ChannelsBackend) -> None:
     channels_plugin = ChannelsPlugin(