Surface vanished websocket peers as WebSocketDisconnect
CI / Build Pip package (push) Successful in 4m55s
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:
@@ -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'
|
||||
]
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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': [],
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user