fix: close ws_bridge send loop on client disconnect

Previously the ws_bridge send loop blocked forever on queue.get(),
never learning the client left. When the server service stopped,
uvicorn waited for open WebSocket tasks until systemd SIGKILLed
after TimeoutStopSec=90s. Now each iteration races recv+send; a
completed receive means the client disconnected => break the loop.
This commit is contained in:
Ari
2026-08-10 11:04:12 -07:00
committed by Elliot Slusky
parent ebf370595d
commit 9498adc7c4
+21 -2
View File
@@ -81,8 +81,27 @@ def create_ws_router(event_bus: EventBus) -> Any:
clients[websocket] = (queue, loop)
try:
while True:
payload = await queue.get()
await websocket.send_json(payload)
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).
message = await recv
if message.get("type") == "websocket.disconnect":
break
else:
await websocket.send_json(payload.result())
except WebSocketDisconnect:
pass
finally: