fix: make websocket bridge race cancellation-safe

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
 

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())