Surface vanished websocket peers as WebSocketDisconnect
CI / Build Pip package (push) Successful in 4m55s

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.
This commit is contained in:
2026-09-29 22:10:16 +08:00
parent 148a35b71c
commit 7c46badebb
7 changed files with 357 additions and 26 deletions
+2 -1
View File
@@ -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'
]
+7 -1
View File
@@ -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
+14 -5
View File
@@ -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()
@@ -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
+113 -3
View File
@@ -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': [],
}