Simple reconnection logic

revolt/client.py

@@ -107,7 +107,7 @@ class Client:
             except:
                 raise RevoltError(f"Cant fetch api info:\n{text}")
 
-    async def start(self) -> None:
+    async def start(self, *, reconnect: bool = True) -> None:
         """Starts the client"""
         api_info = await self.get_api_info()
 
@@ -116,7 +116,7 @@ class Client:
         self.state = State(self.http, api_info, self.max_messages)
         self.websocket = WebsocketHandler(self.session, self.token, api_info["ws"], self.dispatch, self.state)
 
-        await self.websocket.start()
+        await self.websocket.start(reconnect)
 
     async def stop(self) -> None:
         await self.websocket.websocket.close()

revolt/websocket.py

@@ -465,26 +465,32 @@ class WebsocketHandler:
 
             self.dispatch("bulk_message_delete", messages)
 
-    async def start(self) -> None:
+    async def start(self, reconnect: bool) -> None:
         if use_msgpack:
             url = f"{self.ws_url}?format=msgpack"
         else:
             url = f"{self.ws_url}?format=json"
 
-        self.websocket = await self.session.ws_connect(url)  # type: ignore
-        await self.send_authenticate()
-        asyncio.create_task(self.heartbeat())
+        while True:
+            self.websocket = await self.session.ws_connect(url)  # type: ignore
+            await self.send_authenticate()
+            hb = asyncio.create_task(self.heartbeat())
 
-        async for msg in self.websocket:
-            msg = cast(WSMessage, msg)  # aiohttp doesnt use NamedTuple so the type info is missing
+            async for msg in self.websocket:
+                msg = cast(WSMessage, msg)  # aiohttp doesnt use NamedTuple so the type info is missing
 
-            if use_msgpack:
-                data = cast(bytes, msg.data)
+                if use_msgpack:
+                    data = cast(bytes, msg.data)
 
-                payload = msgpack.unpackb(data)
-            else:
-                data = cast(str, msg.data)
+                    payload = msgpack.unpackb(data)
+                else:
+                    data = cast(str, msg.data)
 
-                payload = json.loads(data)
+                    payload = json.loads(data)
 
-            self.loop.create_task(self.handle_event(payload))
+                self.loop.create_task(self.handle_event(payload))
+
+            hb.cancel()
+
+            if not reconnect:
+                return