diff --git a/transport_ws.py b/transport_ws.py index 2edf437..a8a38d0 100644 --- a/transport_ws.py +++ b/transport_ws.py @@ -13,16 +13,19 @@ class WSTransport: self.thread = None self.app = FastAPI() - + @self.app.websocket(self.path) async def websocket_endpoint(*, websocket: WebSocket, token: str = Query("")): if token != self.token: return await websocket.accept() - while True: - data = await websocket.receive_bytes() - await self._recv_q.put(data) + try: + while True: + data = await websocket.receive_bytes() + await self._recv_q.put(data) + except WebSocketDisconnect: + pass async def start(self): if self.thread != None: