Plug state resume into deletion unschedule

litecord/gateway/state_manager.py

@@ -245,6 +245,8 @@ class StateManager:
     async def _future_cleanup(self, state: GatewayState):
         await asyncio.sleep(30)
         self.remove(state)
+        state.ws.state = None
+        state.ws = None
 
     async def schedule_deletion(self, state: GatewayState):
         task = app.loop.create_task(self._future_cleanup(state))
@@ -253,6 +255,7 @@ class StateManager:
     async def unschedule_deletion(self, state: GatewayState):
         try:
             task = self.tasks.pop(state.session_id)
-            task.cancel()
         except KeyError:
-            pass
+            return
+
+        task.cancel()

litecord/gateway/websocket.py

@@ -824,6 +824,7 @@ class GatewayWebsocket:
             return await self.invalidate_session(False)
 
         # relink this connection
+        await self.app.state_manager.unschedule_deletion(state)
         self.state = state
         state.ws = self
 
@@ -1085,9 +1086,10 @@ class GatewayWebsocket:
             task.cancel()
 
         if self.state:
-            self.app.state_manager.remove(self.state)
-            self.state.ws = None
-            self.state = None
+            self.app.state_manager.schedule_deletion(self.state)
+            # self.app.state_manager.remove(self.state)
+            # self.state.ws = None
+            # self.state = None
 
     async def _check_conns(self, user_id):
         """Check if there are any existing connections.