From c9486d84ceb137e664426c72826eb798e91766d2 Mon Sep 17 00:00:00 2001
From: ooliver1 <oliverwilkes2006@icloud.com>
Date: Fri, 16 Sep 2022 17:02:07 +0100
Subject: [PATCH] feat: support resuming

---
 mafic/node.py             | 18 ++++++++++++++++--
 mafic/typings/outgoing.py | 28 ++++++++++++++++++----------
 2 files changed, 34 insertions(+), 12 deletions(-)

diff --git a/mafic/node.py b/mafic/node.py
index 3daf999..a7b7185 100644
--- a/mafic/node.py
+++ b/mafic/node.py
@@ -4,6 +4,7 @@
 
 from asyncio import create_task, sleep
 from logging import getLogger
+from os import urandom
 from typing import TYPE_CHECKING
 
 from aiohttp import ClientSession, WSMsgType
@@ -35,6 +36,7 @@ def __init__(
         heartbeat: int = 30,
         timeout: float = 10,
         session: ClientSession | None = None,
+        resume_key: str | None = None,
     ) -> None:
         self._host = host
         self._port = port
@@ -48,6 +50,7 @@ def __init__(
 
         self._rest_uri = f"http{'s' if secure else ''}://{host}:{port}"
         self._ws_uri = f"ws{'s' if secure else ''}://{host}:{port}"
+        self._resume_key = resume_key or urandom(8).hex()
 
         self._ws: ClientWebSocketResponse | None = None
         self._ws_task: Task[None] | None = None
@@ -88,6 +91,7 @@ async def _connect(self) -> None:
             "Authorization": self.__password,
             "User-Id": str(self._client.user.id),
             "Client-Name": f"Mafic/{__import__('mafic').__version__}",
+            "Resume-Key": self._resume_key,
         }
 
         self._ws = await session.ws_connect(  # pyright: ignore[reportUnknownMemberType]
@@ -96,6 +100,7 @@ async def _connect(self) -> None:
             heartbeat=self._heartbeat,
             headers=headers,
         )
+        create_task(self.send_resume_configuration())
         self._ws_task = create_task(
             self._ws_listener(), name=f"mafic node {self._label}"
         )
@@ -137,6 +142,9 @@ async def __send(self, data: OutgoingMessage) -> None:
 
         await self._ws.send_json(data, dumps=dumps)
 
+    async def _handle_msg(self, data: dict[str, Any]) -> None:
+        raise NotImplementedError
+
     def send_voice_server_update(
         self, guild_id: int, session_id: str, data: VoiceServerUpdatePayload
     ) -> Coro[None]:
@@ -149,5 +157,11 @@ def send_voice_server_update(
             }
         )
 
-    async def _handle_msg(self, data: dict[str, Any]) -> None:
-        raise NotImplementedError
+    def send_resume_configuration(self) -> Coro[None]:
+        return self.__send(
+            {
+                "op": "configureResuming",
+                "key": self._resume_key,
+                "timeout": 60,
+            }
+        )
diff --git a/mafic/typings/outgoing.py b/mafic/typings/outgoing.py
index 8b89377..50105b4 100644
--- a/mafic/typings/outgoing.py
+++ b/mafic/typings/outgoing.py
@@ -11,6 +11,7 @@
 
 __all__ = (
     "ChannelMix",
+    "ConfigureResumingPayload",
     "DestroyPayload",
     "Distortion",
     "EQBand",
@@ -19,7 +20,7 @@
     "LavalinkVoiceState",
     "LowPass",
     "OutgoingMessage",
-    "OutgoingMessagePayload",
+    "PayloadWithGuild",
     "PausePayload",
     "PlayPayload",
     "Rotation",
@@ -33,7 +34,7 @@
 )
 
 
-class OutgoingMessagePayload(TypedDict):
+class PayloadWithGuild(TypedDict):
     guildId: str
 
 
@@ -43,13 +44,13 @@ class LavalinkVoiceState(TypedDict, total=False):
     event: VoiceServerUpdatePayload
 
 
-class VoiceStatePayload(OutgoingMessagePayload):
+class VoiceStatePayload(PayloadWithGuild):
     op: Literal["voiceUpdate"]
     sessionId: str
     event: VoiceServerUpdatePayload
 
 
-class PlayPayload(OutgoingMessagePayload):
+class PlayPayload(PayloadWithGuild):
     op: Literal["play"]
     track: str
     startTime: str
@@ -59,21 +60,21 @@ class PlayPayload(OutgoingMessagePayload):
     pause: bool
 
 
-class StopPayload(OutgoingMessagePayload):
+class StopPayload(PayloadWithGuild):
     op: Literal["stop"]
 
 
-class PausePayload(OutgoingMessagePayload):
+class PausePayload(PayloadWithGuild):
     op: Literal["pause"]
     pause: bool
 
 
-class SeekPayload(OutgoingMessagePayload):
+class SeekPayload(PayloadWithGuild):
     op: Literal["seek"]
     position: int
 
 
-class VolumePayload(OutgoingMessagePayload):
+class VolumePayload(PayloadWithGuild):
     op: Literal["volume"]
     volume: int
 
@@ -132,7 +133,7 @@ class LowPass(TypedDict):
     smoothing: float
 
 
-class FilterPayload(OutgoingMessagePayload):
+class FilterPayload(PayloadWithGuild):
     op: Literal["filters"]
     volume: NotRequired[float]
     equalizer: NotRequired[list[EQBand]]
@@ -146,10 +147,16 @@ class FilterPayload(OutgoingMessagePayload):
     lowPass: NotRequired[LowPass]
 
 
-class DestroyPayload(OutgoingMessagePayload):
+class DestroyPayload(PayloadWithGuild):
     op: Literal["destroy"]
 
 
+class ConfigureResumingPayload(TypedDict):
+    op: Literal["configureResuming"]
+    key: str | None
+    timeout: int
+
+
 OutgoingMessage = Union[
     VoiceStatePayload,
     PlayPayload,
@@ -159,4 +166,5 @@ class DestroyPayload(OutgoingMessagePayload):
     VolumePayload,
     FilterPayload,
     DestroyPayload,
+    ConfigureResumingPayload,
 ]