From 7b95ec1b6bdf65d3c4804bf90557f1072437aff9 Mon Sep 17 00:00:00 2001
From: Zomatree <me@zomatree.live>
Date: Fri, 29 Dec 2023 20:11:55 +0000
Subject: [PATCH] Simple reconnection logic

---
 revolt/client.py    |  4 ++--
 revolt/websocket.py | 32 +++++++++++++++++++-------------
 2 files changed, 21 insertions(+), 15 deletions(-)

diff --git a/revolt/client.py b/revolt/client.py
index d8463b2..6870ff8 100755
--- a/revolt/client.py
+++ b/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()
diff --git a/revolt/websocket.py b/revolt/websocket.py
index 8f4b3e9..ae8146b 100755
--- a/revolt/websocket.py
+++ b/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