From 7c46badebb968059b19b2f229ad8caf6645ced0a Mon Sep 17 00:00:00 2001 From: Walter Oggioni Date: Tue, 29 Sep 2026 14:00:39 +0000 Subject: [PATCH] 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 | 3 +- packages/kaya-core/src/kaya/core/_app.py | 8 +- packages/kaya-core/src/kaya/core/_asgi.py | 19 ++- .../kaya-core/src/kaya/core/_websocket.py | 35 +++++ packages/kaya-core/tests/test_websocket.py | 116 +++++++++++++- packages/kaya-rsgi/src/kaya/rsgi/_rsgi.py | 54 ++++++- packages/kaya-rsgi/tests/test_rsgi.py | 148 ++++++++++++++++-- 7 files changed, 357 insertions(+), 26 deletions(-) diff --git a/packages/kaya-core/src/kaya/core/__init__.py b/packages/kaya-core/src/kaya/core/__init__.py index 76b56e6..ba76117 100644 --- a/packages/kaya-core/src/kaya/core/__init__.py +++ b/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' ] diff --git a/packages/kaya-core/src/kaya/core/_app.py b/packages/kaya-core/src/kaya/core/_app.py index baf55b6..2d76b19 100644 --- a/packages/kaya-core/src/kaya/core/_app.py +++ b/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 diff --git a/packages/kaya-core/src/kaya/core/_asgi.py b/packages/kaya-core/src/kaya/core/_asgi.py index 0612612..156e858 100644 --- a/packages/kaya-core/src/kaya/core/_asgi.py +++ b/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) + 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() diff --git a/packages/kaya-core/src/kaya/core/_websocket.py b/packages/kaya-core/src/kaya/core/_websocket.py index 2b223a1..9ec1147 100644 --- a/packages/kaya-core/src/kaya/core/_websocket.py +++ b/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 diff --git a/packages/kaya-core/tests/test_websocket.py b/packages/kaya-core/tests/test_websocket.py index 14120f3..36b090e 100644 --- a/packages/kaya-core/tests/test_websocket.py +++ b/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': [], + } diff --git a/packages/kaya-rsgi/src/kaya/rsgi/_rsgi.py b/packages/kaya-rsgi/src/kaya/rsgi/_rsgi.py index 774cbe6..4806388 100644 --- a/packages/kaya-rsgi/src/kaya/rsgi/_rsgi.py +++ b/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. - self._transport = await self._protocol.accept() + 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') - message = await self._transport.receive() + 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') - await self._transport.send_str(data) + 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') - await self._transport.send_bytes(data) + 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: - self._protocol.close(code) + if self._closed: + return + try: + self._protocol.close(code) + except RSGIProtocolClosed: + self._closed = True async def __anext__(self) -> WebSocketMessage: message = await self.receive() diff --git a/packages/kaya-rsgi/tests/test_rsgi.py b/packages/kaya-rsgi/tests/test_rsgi.py index 096850b..26b2a36 100644 --- a/packages/kaya-rsgi/tests/test_rsgi.py +++ b/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: - scheme = 'ws' - path = '/ws' - query_string = '' - headers = {} - client = '127.0.0.1:12345' - server = '127.0.0.1:80' - - class FakeProtocol: + class FakeProtocolWithoutWebsockets: 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)