Add support for resume_gateway_url

discord/client.py

@@ -682,6 +682,8 @@ class Client:
                 _log.info('Got a request to %s the websocket.', e.op)
                 self.dispatch('disconnect')
                 ws_params.update(sequence=self.ws.sequence, resume=e.resume, session=self.ws.session_id)
+                if e.resume:
+                    ws_params['gateway'] = self.ws.gateway
                 continue
             except (
                 OSError,
@@ -705,7 +707,13 @@ class Client:
 
                 # If we get connection reset by peer then try to RESUME
                 if isinstance(exc, OSError) and exc.errno in (54, 10054):
-                    ws_params.update(sequence=self.ws.sequence, initial=False, resume=True, session=self.ws.session_id)
+                    ws_params.update(
+                        sequence=self.ws.sequence,
+                        gateway=self.ws.gateway,
+                        initial=False,
+                        resume=True,
+                        session=self.ws.session_id,
+                    )
                     continue
 
                 # We should only get this when an unhandled close code happens,
@@ -725,7 +733,12 @@ class Client:
                 # Always try to RESUME the connection
                 # If the connection is not RESUME-able then the gateway will invalidate the session.
                 # This is apparently what the official Discord client does.
-                ws_params.update(sequence=self.ws.sequence, resume=True, session=self.ws.session_id)
+                ws_params.update(
+                    sequence=self.ws.sequence,
+                    gateway=self.ws.gateway,
+                    resume=True,
+                    session=self.ws.session_id,
+                )
 
     async def close(self) -> None:
         """|coro|

discord/gateway.py

@@ -291,6 +291,7 @@ class DiscordWebSocket:
         _max_heartbeat_timeout: float
 
     # fmt: off
+    DEFAULT_GATEWAY    = 'wss://gateway.discord.gg/'
     DISPATCH           = 0
     HEARTBEAT          = 1
     IDENTIFY           = 2
@@ -350,13 +351,24 @@ class DiscordWebSocket:
         session: Optional[str] = None,
         sequence: Optional[int] = None,
         resume: bool = False,
+        encoding: str = 'json',
+        zlib: bool = True,
     ) -> Self:
         """Creates a main websocket for Discord from a :class:`Client`.
 
         This is for internal use only.
         """
-        gateway = gateway or await client.http.get_gateway()
-        socket = await client.http.ws_connect(gateway)
+        # Circular import
+        from .http import INTERNAL_API_VERSION
+
+        gateway = gateway or cls.DEFAULT_GATEWAY
+
+        if zlib:
+            url = f'{gateway}?v={INTERNAL_API_VERSION}&encoding={encoding}&compress=zlib-stream'
+        else:
+            url = f'{gateway}?v={INTERNAL_API_VERSION}&encoding={encoding}'
+
+        socket = await client.http.ws_connect(url)
         ws = cls(socket, loop=client.loop)
 
         # dynamically add attributes needed
@@ -533,6 +545,7 @@ class DiscordWebSocket:
 
                 self.sequence = None
                 self.session_id = None
+                self.gateway = self.DEFAULT_GATEWAY
                 _log.info('Shard ID %s session has been invalidated.', self.shard_id)
                 await self.close(code=1000)
                 raise ReconnectWebSocket(self.shard_id, resume=False)
@@ -543,6 +556,7 @@ class DiscordWebSocket:
         if event == 'READY':
             self.sequence = msg['s']
             self.session_id = data['session_id']
+            self.gateway = data['resume_gateway_url']
             _log.info('Shard ID %s has connected to Gateway (Session ID: %s).', self.shard_id, self.session_id)
 
         elif event == 'RESUMED':

discord/shard.py

@@ -183,6 +183,7 @@ class Shard:
             coro = DiscordWebSocket.from_client(
                 self._client,
                 resume=exc.resume,
+                gateway=None if not exc.resume else self.ws.gateway,
                 shard_id=self.id,
                 session=self.ws.session_id,
                 sequence=self.ws.sequence,

discord/types/gateway.py

@@ -66,6 +66,7 @@ class ReadyEvent(TypedDict):
     user: User
     guilds: List[UnavailableGuild]
     session_id: str
+    resume_gateway_url: str
     shard: List[int]  # shard_id, num_shards
     application: GatewayAppInfo