diff --git a/src/openjarvis/server/ws_bridge.py b/src/openjarvis/server/ws_bridge.py index 94c912af..0dc8fc75 100644 --- a/src/openjarvis/server/ws_bridge.py +++ b/src/openjarvis/server/ws_bridge.py @@ -79,33 +79,41 @@ def create_ws_router(event_bus: EventBus) -> Any: queue: asyncio.Queue = asyncio.Queue(maxsize=100) loop = asyncio.get_running_loop() clients[websocket] = (queue, loop) + recv: asyncio.Task | None = None + payload: asyncio.Task | None = None + disconnected = False try: + recv = asyncio.create_task(websocket.receive()) + payload = asyncio.create_task(queue.get()) while True: - recv = asyncio.create_task(websocket.receive()) - payload = asyncio.create_task(queue.get()) done, _ = await asyncio.wait( {recv, payload}, return_when=asyncio.FIRST_COMPLETED ) - for task in (recv, payload): - if task not in done: - task.cancel() if recv in done: - # A completed receive surfaces the client's disconnect — - # Starlette delivers it as a message (or raises - # WebSocketDisconnect for a clean close frame) only when - # the app actually reads from the socket. Without this the - # handler never learns the client left, its task stays - # open forever, and uvicorn's graceful shutdown blocks on - # it until systemd SIGKILLs the unit (TimeoutStopSec). + # Starlette surfaces a disconnect message only when the app + # reads from the socket. Without this receive, the handler + # can stay parked on queue.get() after the client leaves. message = await recv if message.get("type") == "websocket.disconnect": + disconnected = True break - else: + recv = asyncio.create_task(websocket.receive()) + if payload in done: await websocket.send_json(payload.result()) + payload = asyncio.create_task(queue.get()) except WebSocketDisconnect: - pass + disconnected = True finally: clients.pop(websocket, None) + pending = [task for task in (recv, payload) if task is not None] + for task in pending: + task.cancel() + cleanup = asyncio.gather(*pending, return_exceptions=True) + try: + await asyncio.shield(cleanup) + except asyncio.CancelledError: + if not disconnected: + raise return router diff --git a/tests/server/test_ws_bridge.py b/tests/server/test_ws_bridge.py index 1727c1e0..b678bdf8 100644 --- a/tests/server/test_ws_bridge.py +++ b/tests/server/test_ws_bridge.py @@ -2,7 +2,9 @@ from __future__ import annotations +import asyncio import time +from types import SimpleNamespace import pytest @@ -60,3 +62,102 @@ class TestWSBridge: time.sleep(0.05) # Let call_soon_threadsafe deliver to queue data = ws.receive_json() assert data["data"]["agent_id"] == "agent-A" + + def test_client_disconnect_stops_handler(self, event_bus): + async def exercise(): + from openjarvis.server.ws_bridge import create_ws_router + + class FakeWebSocket: + app = SimpleNamespace(state=SimpleNamespace(api_key="")) + query_params = {} + headers = {} + + async def accept(self): + pass + + async def receive(self): + return {"type": "websocket.disconnect"} + + endpoint = create_ws_router(event_bus).routes[0].endpoint + await asyncio.wait_for(endpoint(FakeWebSocket()), timeout=1) + + asyncio.run(exercise()) + + def test_simultaneous_client_message_does_not_drop_event(self, event_bus): + async def exercise(): + from openjarvis.server.ws_bridge import create_ws_router + + class FakeWebSocket: + def __init__(self): + self.app = SimpleNamespace(state=SimpleNamespace(api_key="")) + self.query_params = {} + self.headers = {} + self.sent = [] + self.receive_count = 0 + self.disconnect = asyncio.Event() + + async def accept(self): + pass + + async def receive(self): + self.receive_count += 1 + if self.receive_count == 1: + event_bus.publish( + EventType.AGENT_TICK_START, {"agent_id": "not-dropped"} + ) + return {"type": "websocket.receive", "text": "client message"} + await self.disconnect.wait() + return {"type": "websocket.disconnect"} + + async def send_json(self, payload): + self.sent.append(payload) + self.disconnect.set() + + websocket = FakeWebSocket() + endpoint = create_ws_router(event_bus).routes[0].endpoint + + await asyncio.wait_for(endpoint(websocket), timeout=1) + + assert websocket.sent[0]["data"]["agent_id"] == "not-dropped" + + asyncio.run(exercise()) + + def test_cancelling_handler_cleans_up_child_tasks(self, event_bus): + async def exercise(): + from openjarvis.server.ws_bridge import create_ws_router + + class FakeWebSocket: + def __init__(self): + self.app = SimpleNamespace(state=SimpleNamespace(api_key="")) + self.query_params = {} + self.headers = {} + self.receiving = asyncio.Event() + self.receive_cancelled = asyncio.Event() + + async def accept(self): + pass + + async def receive(self): + self.receiving.set() + try: + await asyncio.Event().wait() + finally: + self.receive_cancelled.set() + + websocket = FakeWebSocket() + endpoint = create_ws_router(event_bus).routes[0].endpoint + handler = asyncio.create_task(endpoint(websocket)) + await websocket.receiving.wait() + + handler.cancel() + with pytest.raises(asyncio.CancelledError): + await handler + + assert websocket.receive_cancelled.is_set() + assert not [ + task + for task in asyncio.all_tasks() + if task is not asyncio.current_task() and not task.done() + ] + + asyncio.run(exercise())