Add locking for multiple reciever

src/nanobot.py

@@ -14,6 +14,7 @@ class Nanobot:
     def __init__(self):
         self._ws = None
         self._chat_id = None
+        self._lock = asyncio.Lock()
 
     async def __aenter__(self):
         await self.connect()
@@ -24,39 +25,47 @@ class Nanobot:
 
     async def connect(self):
         uri = f"{NANOBOT_WS_ENDPOINT}/?token={NANOBOT_WS_TOKEN}"
-        self._ws = await websockets.connect(uri)
+        self._ws = await websockets.connect(uri, ping_interval=None)
         ready = json.loads(await self._ws.recv())
         self._chat_id = ready["chat_id"]
         main_logger.info("connected chat_id=%s", self._chat_id)
 
-    async def chat(self, message: str, timeout: float = 10):
+    async def chat(self, message: str, timeout: float = 60):
-        await self._ws.send(json.dumps({"content": message}))
+        async with self._lock:
-        parts = []
-        while True:
             try:
-                raw = await asyncio.wait_for(self._ws.recv(), timeout=timeout)
+                await self._ws.send(json.dumps({"content": message}))
-            except asyncio.TimeoutError:
+            except websockets.ConnectionClosed:
-                main_logger.info("ws: timeout (no data for %.1fs), generator done", timeout)
+                main_logger.info("ws: reconnecting...")
-                return
+                await self.connect()
-            frame = json.loads(raw)
+                await self._ws.send(json.dumps({"content": message}))
-            event = frame.get("event")
+            parts = []
-            text = frame.get("text", "")
+            while True:
-            main_logger.debug("ws raw: %s", raw)
+                try:
-            if event == "message":
+                    raw = await asyncio.wait_for(self._ws.recv(), timeout=timeout)
-                if frame.get("kind") == "progress":
+                except asyncio.TimeoutError:
-                    main_logger.debug("ws: skipping progress message")
+                    main_logger.info("ws: timeout (no data for %.1fs), reconnecting", timeout)
-                    continue
+                    await self.close()
-                yield text
+                    await self.connect()
-                return
+                    return
-            if event == "delta":
+                frame = json.loads(raw)
-                parts.append(text)
+                event = frame.get("event")
-            elif event == "stream_end":
+                text = frame.get("text", "")
-                t = "".join(parts)
+                main_logger.debug("ws raw: %s", raw)
-                if t:
+                if event == "message":
-                    yield t
+                    if frame.get("kind") == "progress":
-                parts = []
+                        main_logger.debug("ws: skipping progress message")
-            elif event == "error":
+                        continue
-                raise RuntimeError(frame.get("detail", "unknown error"))
+                    yield text
+                    return
+                if event == "delta":
+                    parts.append(text)
+                elif event == "stream_end":
+                    t = "".join(parts)
+                    if t:
+                        yield t
+                    parts = []
+                elif event == "error":
+                    raise RuntimeError(frame.get("detail", "unknown error"))
 
     async def close(self):
         if self._ws: