Surface vanished websocket peers as WebSocketDisconnect

Granian's RSGI transport raises RSGIProtocolClosed when the peer is gone
without a close handshake, which escaped receive() and was logged as an
application error, and sends to the dead transport raised the raw Granian
exception as well.

Treat a closed RSGI transport like the ASGI websocket.disconnect message:
receive() reports a close message so async-for loops end normally, while
writes (accept/send/close) raise a new kaya.core.WebSocketDisconnect
carrying close code 1006. KayaApp.handle_websocket swallows that exception
as normal termination, so after hooks still run and disconnects are no
longer reported as errors. The ASGI adapter maps its server's OSError
disconnect (e.g. uvicorn's ClientDisconnected) to the same exception for
consistent semantics across backends.

packages/kaya-core/src/kaya/core/__init__.py

@@ -4,7 +4,7 @@ from ._http_context import HttpContext
 from ._mixin import KayaMixin
 from ._tree import Tree, PathIterator
 from ._path_handler import PathHandler, Matches
-from ._websocket import WebSocket, WebSocketMessage
+from ._websocket import WebSocket, WebSocketDisconnect, WebSocketMessage
 
 
 __all__ = [
@@ -18,5 +18,6 @@ __all__ = [
     'Matches',
     'PathIterator',
     'WebSocket',
+    'WebSocketDisconnect',
     'WebSocketMessage'
 ]

packages/kaya-core/src/kaya/core/_app.py

@@ -9,7 +9,7 @@ from ._http_method import HttpMethod
 from ._mixin import KayaMixin
 from ._path_handler import Context
 from ._types import StrOrStrings
-from ._websocket import WebSocket, WebSocketMessage
+from ._websocket import WebSocket, WebSocketDisconnect, WebSocketMessage
 from ._asgi import AsgiContext, AsgiWebSocket
 from ._tree import Tree
 from ._types.asgi import LifespanScope, HTTPScope as ASGIHTTPScope, WebSocketScope as ASGIWebSocketScope
@@ -164,6 +164,12 @@ class KayaApp(AbstractKayaApp):
                 await handler.handle_request(ws, captured)
             else:
                 await ws.close(1000)
+        except WebSocketDisconnect:
+            # The peer went away while the handler was using the websocket
+            # (e.g. sending on a vanished transport). This is normal
+            # termination, not an application error, so it is not recorded
+            # on the socket nor re-raised.
+            pass
         except Exception as exc:
             ws.exception = exc
             raise

packages/kaya-core/src/kaya/core/_asgi.py

@@ -17,7 +17,7 @@ from pwo import Maybe
 from pathlib import Path
 from ._http_method import HttpMethod
 from ._http_context import HttpContext
-from ._websocket import WebSocket, WebSocketMessage
+from ._websocket import WebSocket, WebSocketDisconnect, WebSocketMessage
 from ._types import StrOrStrings
 from ._types.asgi import HTTPScope, WebSocketScope
 
@@ -179,7 +179,16 @@ class AsgiWebSocket(WebSocket):
             # Emit a list rather than a tuple: wsproto's AcceptConnection (used
             # by httpx_ws and others) requires list concatenation.
             message['headers'] = list(encode_headers(headers))
+        await self._send_message(message)
+
+    async def _send_message(self, message: Mapping[str, Any]) -> None:
+        try:
             await self._send(message)
+        except OSError as exc:
+            # ASGI servers (e.g. uvicorn) report a vanished peer with an
+            # OSError subclass; surface it as a kaya-level disconnect so
+            # handlers do not need to know about server internals.
+            raise WebSocketDisconnect(code=1006) from exc
 
     async def receive(self) -> WebSocketMessage:
         message = await self._receive()
@@ -200,13 +209,13 @@ class AsgiWebSocket(WebSocket):
             return WebSocketMessage(kind='close')
 
     async def send_text(self, data: str) -> None:
-        await self._send({'type': 'websocket.send', 'text': data})
+        await self._send_message({'type': 'websocket.send', 'text': data})
 
     async def send_bytes(self, data: bytes) -> None:
-        await self._send({'type': 'websocket.send', 'bytes': data})
+        await self._send_message({'type': 'websocket.send', 'bytes': data})
 
     async def close(self, code: int = 1000) -> None:
-        await self._send({'type': 'websocket.close', 'code': code})
+        await self._send_message({'type': 'websocket.close', 'code': code})
 
     async def __anext__(self) -> WebSocketMessage:
         message = await self.receive()

packages/kaya-core/src/kaya/core/_websocket.py

@@ -13,6 +13,41 @@ class WebSocketMessage:
     data: WebSocketData = None
 
 
+class WebSocketDisconnect(Exception):
+    """Raised when using a websocket whose peer has disconnected.
+
+    Concrete websocket implementations raise this from ``send_text()`` and
+    ``send_bytes()`` when the transport reports that the peer is gone (for
+    example when a browser tab is closed without a close handshake), so
+    handlers can catch it and stop producing output:
+
+    .. code-block:: python
+
+        try:
+            await ws.send_text('hello')
+        except WebSocketDisconnect:
+            pass
+
+    Reading from a disconnected websocket is not an error: ``receive()``
+    reports a ``close`` message and iterating with ``async for`` simply ends,
+    which is what most handlers should rely on instead.
+
+    ``code`` is a websocket close code; ``1006`` means the connection was
+    closed abnormally, without a close handshake.
+    """
+
+    code: int
+    reason: str
+
+    def __init__(self, code: int = 1006, reason: str = '') -> None:
+        self.code = code
+        self.reason = reason
+        message = f'WebSocket disconnected (code={code})'
+        if reason:
+            message = f'{message}: {reason}'
+        super().__init__(message)
+
+
 class WebSocket(ABC):
     path: str
     query_string: str

packages/kaya-core/tests/test_websocket.py

@@ -1,9 +1,9 @@
 import unittest
 import httpx
 from pwo import async_test
-from httpx_ws import aconnect_ws, WebSocketDisconnect
+from httpx_ws import aconnect_ws, WebSocketDisconnect as HttpxWebSocketDisconnect
 from httpx_ws.transport import ASGIWebSocketTransport
-from kaya.core import KayaApp, WebSocket
+from kaya.core import KayaApp, WebSocket, WebSocketDisconnect, WebSocketMessage
 from kaya.core._asgi import AsgiWebSocket
 
 
@@ -66,7 +66,7 @@ class WebSocketTest(unittest.TestCase):
         transport = ASGIWebSocketTransport(app=self.app)
 
         async with httpx.AsyncClient(transport=transport, base_url="http://testserver") as client:
-            with self.assertRaises(WebSocketDisconnect) as cm:
+            with self.assertRaises(HttpxWebSocketDisconnect) as cm:
                 async with aconnect_ws("/unknown", client) as ws:
                     pass
 
@@ -151,3 +151,113 @@ class WebSocketTest(unittest.TestCase):
         }
         ws = AsgiWebSocket(scope, receive, send)
         self.assertEqual('ws', ws.scheme)
+
+
+class FakeWebSocket(WebSocket):
+    """Minimal in-memory websocket for exercising ``KayaApp`` directly."""
+
+    def __init__(self, path: str) -> None:
+        self.path = path
+        self.query_string = ''
+        self.scheme = 'ws'
+        self.headers: dict[str, list[str]] = {}
+        self.client = ('127.0.0.1', 12345)
+        self.server = ('127.0.0.1', 80)
+        self.closed: list[int] = []
+
+    async def accept(self, headers=None) -> None:
+        pass
+
+    async def receive(self) -> WebSocketMessage:
+        return WebSocketMessage(kind='close')
+
+    async def send_text(self, data: str) -> None:
+        pass
+
+    async def send_bytes(self, data: bytes) -> None:
+        pass
+
+    async def close(self, code: int = 1000) -> None:
+        self.closed.append(code)
+
+    async def __anext__(self) -> WebSocketMessage:
+        raise StopAsyncIteration
+
+
+class WebSocketDisconnectTest(unittest.TestCase):
+
+    def setUp(self):
+        self.app = KayaApp()
+
+        @self.app.websocket('/disconnect')
+        async def disconnect(ws: WebSocket) -> None:
+            await ws.accept()
+            await ws.send_text('before')
+            raise WebSocketDisconnect()
+
+        @self.app.websocket('/failure')
+        async def failure(ws: WebSocket) -> None:
+            await ws.accept()
+            raise RuntimeError('boom')
+
+    @async_test
+    async def test_disconnect_terminates_quietly(self):
+        after_hook_calls = []
+
+        async def after(ws: WebSocket) -> None:
+            after_hook_calls.append(ws)
+
+        self.app.add_after_websocket_hook(after)
+        ws = FakeWebSocket('/disconnect')
+        await self.app.handle_websocket(ws)
+        self.assertEqual([ws], after_hook_calls)
+        self.assertIsNone(ws.exception)
+
+    @async_test
+    async def test_other_errors_are_recorded_and_reraised(self):
+        ws = FakeWebSocket('/failure')
+        with self.assertRaises(RuntimeError):
+            await self.app.handle_websocket(ws)
+        self.assertIsInstance(ws.exception, RuntimeError)
+
+    @async_test
+    async def test_asgi_send_after_disconnect_raises_websocket_disconnect(self):
+        async def send(message):
+            raise OSError('client gone')
+
+        async def receive():
+            return {'type': 'websocket.disconnect', 'code': 1006}
+
+        ws = AsgiWebSocket(self._scope(), receive, send)
+        with self.assertRaises(WebSocketDisconnect) as ctx:
+            await ws.send_text('hello')
+        self.assertEqual(1006, ctx.exception.code)
+        with self.assertRaises(WebSocketDisconnect):
+            await ws.send_bytes(b'hello')
+        with self.assertRaises(WebSocketDisconnect):
+            await ws.close()
+
+    @async_test
+    async def test_asgi_accept_after_disconnect_raises_websocket_disconnect(self):
+        async def send(message):
+            raise OSError('client gone')
+
+        async def receive():
+            return {'type': 'websocket.connect'}
+
+        ws = AsgiWebSocket(self._scope(), receive, send)
+        with self.assertRaises(WebSocketDisconnect) as ctx:
+            await ws.accept()
+        self.assertEqual(1006, ctx.exception.code)
+
+    @staticmethod
+    def _scope() -> dict:
+        return {
+            'type': 'websocket',
+            'path': '/echo',
+            'query_string': b'',
+            'scheme': 'ws',
+            'client': ('127.0.0.1', 12345),
+            'server': ('127.0.0.1', 80),
+            'headers': [],
+        }

packages/kaya-rsgi/src/kaya/rsgi/_rsgi.py

@@ -17,13 +17,21 @@ from typing import (
 from granian._granian import (  # type: ignore[attr-defined]
     RSGIHTTPProtocol,
     RSGIHTTPScope,
+    RSGIProtocolClosed,
     RSGIWebsocketProtocol,
     RSGIWebsocketScope,
     RSGIWebsocketTransport,
 )
 from pwo import Maybe
 
-from kaya.core import AbstractKayaApp, HttpContext, HttpMethod, WebSocket, WebSocketMessage
+from kaya.core import (
+    AbstractKayaApp,
+    HttpContext,
+    HttpMethod,
+    WebSocket,
+    WebSocketDisconnect,
+    WebSocketMessage,
+)
 from kaya.core._types import StrOrStrings
 
 
@@ -109,6 +117,7 @@ class RsgiContext(HttpContext):
 class RsgiWebSocket(WebSocket):
     _protocol: RSGIWebsocketProtocol
     _transport: Optional[RSGIWebsocketTransport]
+    _closed: bool
     scheme: str
     path: str
     query_string: str
@@ -124,6 +133,7 @@ class RsgiWebSocket(WebSocket):
             )
         self._protocol = protocol
         self._transport = None
+        self._closed = False
         self.scheme = scope.scheme
         self.path = scope.path
         self.query_string = scope.query_string
@@ -144,13 +154,28 @@ class RsgiWebSocket(WebSocket):
     async def accept(self, headers: Optional[Mapping[str, StrOrStrings]] = None) -> None:
         # RSGI's websocket accept() takes no arguments (https://github.com/emmett-framework/granian/blob/master/docs/spec/RSGI.md),
         # so handshake headers (e.g. Set-Cookie) cannot be sent on RSGI; they are silently ignored here.
+        try:
             self._transport = await self._protocol.accept()
+        except RSGIProtocolClosed as exc:
+            self._closed = True
+            raise WebSocketDisconnect(code=1006) from exc
 
     async def receive(self) -> WebSocketMessage:
         if self._transport is None:
             raise RuntimeError('WebSocket connection has not been accepted yet')
+        if self._closed:
+            return WebSocketMessage(kind='close')
+        try:
             message = await self._transport.receive()
+        except RSGIProtocolClosed:
+            # Granian raises this when the peer is gone without a close
+            # handshake (browser tab closed, network dropped, ...). Surface it
+            # as a close message so receive loops terminate normally instead
+            # of seeing an exception.
+            self._closed = True
+            return WebSocketMessage(kind='close')
         if message.kind == 0:
+            self._closed = True
             return WebSocketMessage(kind='close')
         elif message.kind == 1:
             return WebSocketMessage(kind='binary', data=message.data)
@@ -162,15 +187,32 @@ class RsgiWebSocket(WebSocket):
     async def send_text(self, data: str) -> None:
         if self._transport is None:
             raise RuntimeError('WebSocket connection has not been accepted yet')
+        if self._closed:
+            raise WebSocketDisconnect(code=1006)
+        try:
             await self._transport.send_str(data)
+        except RSGIProtocolClosed as exc:
+            self._closed = True
+            raise WebSocketDisconnect(code=1006) from exc
 
     async def send_bytes(self, data: bytes) -> None:
         if self._transport is None:
             raise RuntimeError('WebSocket connection has not been accepted yet')
+        if self._closed:
+            raise WebSocketDisconnect(code=1006)
+        try:
             await self._transport.send_bytes(data)
+        except RSGIProtocolClosed as exc:
+            self._closed = True
+            raise WebSocketDisconnect(code=1006) from exc
 
     async def close(self, code: int = 1000) -> None:
+        if self._closed:
+            return
+        try:
             self._protocol.close(code)
+        except RSGIProtocolClosed:
+            self._closed = True
 
     async def __anext__(self) -> WebSocketMessage:
         message = await self.receive()

packages/kaya-rsgi/tests/test_rsgi.py

@@ -1,22 +1,150 @@
 import unittest
+from types import SimpleNamespace
+from typing import Any
+
+from granian.rsgi import ProtocolClosed
+from pwo import async_test
+
+from kaya.core import WebSocketDisconnect
 from kaya.rsgi import RsgiWebSocket
 
 
+class FakeScope:
+    scheme = 'ws'
+    path = '/ws'
+    query_string = ''
+    headers: dict[str, list[str]] = {}
+    client = '127.0.0.1:12345'
+    server = '127.0.0.1:80'
+
+
+class FakeTransport:
+    """A transport whose peer vanished without a close handshake."""
+
+    def __init__(self) -> None:
+        self.receive_calls = 0
+        self.sent: list[Any] = []
+
+    async def receive(self) -> Any:
+        self.receive_calls += 1
+        raise ProtocolClosed()
+
+    async def send_str(self, data: str) -> None:
+        self.sent.append(data)
+        raise ProtocolClosed()
+
+    async def send_bytes(self, data: bytes) -> None:
+        self.sent.append(data)
+        raise ProtocolClosed()
+
+
+class CleanCloseTransport(FakeTransport):
+    """A transport that reports a close message (kind 0)."""
+
+    async def receive(self) -> Any:
+        self.receive_calls += 1
+        return SimpleNamespace(kind=0, data=None)
+
+
+class FakeProtocol:
+    def __init__(self, transport: FakeTransport) -> None:
+        self._transport = transport
+        self.closed: list[int] = []
+
+    async def accept(self) -> FakeTransport:
+        return self._transport
+
+    def close(self, code: int) -> None:
+        self.closed.append(code)
+
+
+class ClosedOnAcceptProtocol(FakeProtocol):
+    """A protocol whose peer vanished during the handshake."""
+
+    async def accept(self) -> FakeTransport:
+        raise ProtocolClosed()
+
+
 class RsgiWebSocketTest(unittest.TestCase):
 
     def test_misconfigured_granian(self):
-        class FakeScope:
+        class FakeProtocolWithoutWebsockets:
-            scheme = 'ws'
-            path = '/ws'
-            query_string = ''
-            headers = {}
-            client = '127.0.0.1:12345'
-            server = '127.0.0.1:80'
-
-        class FakeProtocol:
             pass
 
         with self.assertRaises(RuntimeError) as ctx:
-            RsgiWebSocket(FakeScope(), FakeProtocol())  # type: ignore[arg-type]
+            RsgiWebSocket(FakeScope(), FakeProtocolWithoutWebsockets())  # type: ignore[arg-type]
 
         self.assertIn('Granian was not configured for websockets', str(ctx.exception))
+
+    async def _socket(self, transport: FakeTransport) -> tuple[RsgiWebSocket, FakeProtocol]:
+        protocol = FakeProtocol(transport)
+        ws = RsgiWebSocket(FakeScope(), protocol)  # type: ignore[arg-type]
+        await ws.accept()
+        return ws, protocol
+
+    @async_test
+    async def test_receive_returns_close_when_transport_is_gone(self):
+        ws, _ = await self._socket(FakeTransport())
+        message = await ws.receive()
+        self.assertEqual('close', message.kind)
+
+    @async_test
+    async def test_iteration_ends_when_transport_is_gone(self):
+        ws, _ = await self._socket(FakeTransport())
+        messages = [message async for message in ws]
+        self.assertEqual([], messages)
+
+    @async_test
+    async def test_receive_remembers_the_close(self):
+        transport = FakeTransport()
+        ws, _ = await self._socket(transport)
+        self.assertEqual('close', (await ws.receive()).kind)
+        self.assertEqual('close', (await ws.receive()).kind)
+        self.assertEqual(1, transport.receive_calls)
+
+    @async_test
+    async def test_send_text_raises_disconnect_when_transport_is_gone(self):
+        ws, _ = await self._socket(FakeTransport())
+        with self.assertRaises(WebSocketDisconnect) as ctx:
+            await ws.send_text('hello')
+        self.assertEqual(1006, ctx.exception.code)
+
+    @async_test
+    async def test_send_bytes_raises_disconnect_when_transport_is_gone(self):
+        ws, _ = await self._socket(FakeTransport())
+        with self.assertRaises(WebSocketDisconnect) as ctx:
+            await ws.send_bytes(b'hello')
+        self.assertEqual(1006, ctx.exception.code)
+
+    @async_test
+    async def test_sends_after_disconnect_do_not_touch_the_transport(self):
+        transport = FakeTransport()
+        ws, _ = await self._socket(transport)
+        with self.assertRaises(WebSocketDisconnect):
+            await ws.send_text('first')
+        with self.assertRaises(WebSocketDisconnect):
+            await ws.send_text('second')
+        self.assertEqual(['first'], transport.sent)
+
+    @async_test
+    async def test_send_after_a_clean_close_raises_disconnect(self):
+        transport = CleanCloseTransport()
+        ws, _ = await self._socket(transport)
+        self.assertEqual('close', (await ws.receive()).kind)
+        with self.assertRaises(WebSocketDisconnect):
+            await ws.send_text('hello')
+        self.assertEqual([], transport.sent)
+
+    @async_test
+    async def test_accept_raises_disconnect_when_transport_is_gone(self):
+        ws = RsgiWebSocket(FakeScope(), ClosedOnAcceptProtocol(FakeTransport()))  # type: ignore[arg-type]
+        with self.assertRaises(WebSocketDisconnect) as ctx:
+            await ws.accept()
+        self.assertEqual(1006, ctx.exception.code)
+
+    @async_test
+    async def test_close_is_noop_once_the_transport_is_gone(self):
+        ws, protocol = await self._socket(FakeTransport())
+        await ws.receive()
+        await ws.close(1000)
+        self.assertEqual([], protocol.closed)