From 9806aeb83179d0d1e90d903e30db7e69e0d492e5 Mon Sep 17 00:00:00 2001
From: Michael H <michael@michaelhall.tech>
Date: Sun, 1 Dec 2024 16:19:09 -0500
Subject: [PATCH] Add public method to get session start limits

---
 discord/http.py  |  5 +++--
 discord/shard.py | 58 +++++++++++++++++++++++++++++++++++++++++++++++-
 docs/api.rst     | 10 ++++++++-
 3 files changed, 69 insertions(+), 4 deletions(-)

diff --git a/discord/http.py b/discord/http.py
index ac1254e3f..7c3f198c9 100644
--- a/discord/http.py
+++ b/discord/http.py
@@ -97,6 +97,7 @@ if TYPE_CHECKING:
         subscription,
     )
     from .types.snowflake import Snowflake, SnowflakeList
+    from .types.gateway import SessionStartLimit
 
     from types import TracebackType
 
@@ -2753,13 +2754,13 @@ class HTTPClient:
 
     # Misc
 
-    async def get_bot_gateway(self) -> Tuple[int, str]:
+    async def get_bot_gateway(self) -> Tuple[int, str, SessionStartLimit]:
         try:
             data = await self.request(Route('GET', '/gateway/bot'))
         except HTTPException as exc:
             raise GatewayNotFound() from exc
 
-        return data['shards'], data['url']
+        return data['shards'], data['url'], data['session_start_limit']
 
     def get_user(self, user_id: Snowflake) -> Response[user.User]:
         return self.request(Route('GET', '/users/{user_id}', user_id=user_id))
diff --git a/discord/shard.py b/discord/shard.py
index 52155aa24..eeb240c95 100644
--- a/discord/shard.py
+++ b/discord/shard.py
@@ -47,13 +47,16 @@ from .enums import Status
 from typing import TYPE_CHECKING, Any, Callable, Tuple, Type, Optional, List, Dict
 
 if TYPE_CHECKING:
+    from typing_extensions import Unpack
     from .gateway import DiscordWebSocket
     from .activity import BaseActivity
     from .flags import Intents
+    from .types.gateway import SessionStartLimit
 
 __all__ = (
     'AutoShardedClient',
     'ShardInfo',
+    'SessionStartLimits',
 )
 
 _log = logging.getLogger(__name__)
@@ -293,6 +296,32 @@ class ShardInfo:
         return self._parent.ws.is_ratelimited()
 
 
+class SessionStartLimits:
+    """A class that holds info about session start limits
+
+    Attributes
+    ----------
+    total: :class:`int`
+        The total number of session starts the current user is allowed
+    remaining: :class:`int`
+        Remaining remaining number of session starts the current user is allowed
+    reset_after: :class:`int`
+        The number of milliseconds until the limit resets
+    max_concurrency: :class:`int`
+        The number of identify requests allowed per 5 seconds
+
+    .. versionadded:: 2.5
+    """
+
+    __slots__ = ("total", "remaining", "reset_after", "max_concurrency")
+
+    def __init__(self, **kwargs: Unpack[SessionStartLimit]):
+        self.total: int = kwargs['total']
+        self.remaining: int = kwargs['remaining']
+        self.reset_after: int = kwargs['reset_after']
+        self.max_concurrency: int = kwargs['max_concurrency']
+
+
 class AutoShardedClient(Client):
     """A client similar to :class:`Client` except it handles the complications
     of sharding for the user into a more manageable and transparent single
@@ -415,6 +444,33 @@ class AutoShardedClient(Client):
         """Mapping[int, :class:`ShardInfo`]: Returns a mapping of shard IDs to their respective info object."""
         return {shard_id: ShardInfo(parent, self.shard_count) for shard_id, parent in self.__shards.items()}
 
+    async def fetch_session_start_limits(self) -> SessionStartLimits:
+        """|coro|
+
+        Get the session start limits.
+
+        This is not typically needed, and will be handled for you by default.
+
+        At the point where you are launching multiple instances
+        with manual shard ranges and are considered required to use large bot
+        sharding by Discord, this function when used along IPC and a
+        before_identity_hook can speed up session start.
+
+        .. versionadded:: 2.5
+
+        Returns
+        -------
+        :class:`SessionStartLimits`
+            A class containing the session start limits
+
+        Raises
+        ------
+        GatewayNotFound
+            The gateway was unreachable
+        """
+        _, _, limits = await self.http.get_bot_gateway()
+        return SessionStartLimits(**limits)
+
     async def launch_shard(self, gateway: yarl.URL, shard_id: int, *, initial: bool = False) -> None:
         try:
             coro = DiscordWebSocket.from_client(self, initial=initial, gateway=gateway, shard_id=shard_id)
@@ -434,7 +490,7 @@ class AutoShardedClient(Client):
 
         if self.shard_count is None:
             self.shard_count: int
-            self.shard_count, gateway_url = await self.http.get_bot_gateway()
+            self.shard_count, gateway_url, _session_start_limit = await self.http.get_bot_gateway()
             gateway = yarl.URL(gateway_url)
         else:
             gateway = DiscordWebSocket.DEFAULT_GATEWAY
diff --git a/docs/api.rst b/docs/api.rst
index 338368910..ca0fb4ef4 100644
--- a/docs/api.rst
+++ b/docs/api.rst
@@ -1381,7 +1381,7 @@ Subscriptions
 .. function:: on_subscription_delete(subscription)
 
     Called when a subscription is deleted.
-    
+
     .. versionadded:: 2.5
 
     :param subscription: The subscription that was deleted.
@@ -5209,6 +5209,14 @@ ShardInfo
 .. autoclass:: ShardInfo()
     :members:
 
+SessionStartLimits
+~~~~~~~~~~~~~~~~~~~~
+
+.. attributetable:: SessionStartLimits
+
+.. autoclass:: SessionStartLimits()
+    :members:
+
 SKU
 ~~~~~~~~~~~