From 5e0570f8414709c11e5b1be567c044182ca0c4e4 Mon Sep 17 00:00:00 2001
From: Benjamin Himes <benhimes@opentensor.dev>
Date: Fri, 13 Jun 2025 14:31:01 +0200
Subject: [PATCH] Adds websocket shutdown timer exposure

---
 bittensor/core/async_subtensor.py        | 8 ++++++++
 bittensor/core/subtensor_api/__init__.py | 8 ++++++--
 2 files changed, 14 insertions(+), 2 deletions(-)

diff --git a/bittensor/core/async_subtensor.py b/bittensor/core/async_subtensor.py
index f91df67952..4dd34bb71f 100644
--- a/bittensor/core/async_subtensor.py
+++ b/bittensor/core/async_subtensor.py
@@ -119,6 +119,7 @@ def __init__(
         retry_forever: bool = False,
         _mock: bool = False,
         archive_endpoints: Optional[list[str]] = None,
+        websocket_shutdown_timer: float = 5.0
     ):
         """
         Initializes an instance of the AsyncSubtensor class.
@@ -134,6 +135,8 @@ def __init__(
             archive_endpoints: Similar to fallback_endpoints, but specifically only archive nodes. Will be used in cases
                 where you are requesting a block that is too old for your current (presumably lite) node. Defaults to
                 `None`
+            websocket_shutdown_timer: Number of seconds to wait after last request before shutting down connection to
+                the node
 
         Raises:
             Any exceptions raised during the setup, configuration, or connection process.
@@ -157,6 +160,7 @@ def __init__(
             retry_forever=retry_forever,
             _mock=_mock,
             archive_endpoints=archive_endpoints,
+            ws_shutdown_timer=websocket_shutdown_timer
         )
         if self.log_verbose:
             logging.info(
@@ -296,6 +300,7 @@ def _get_substrate(
         retry_forever: bool = False,
         _mock: bool = False,
         archive_endpoints: Optional[list[str]] = None,
+        ws_shutdown_timer: float = 5.0
     ) -> Union[AsyncSubstrateInterface, RetryAsyncSubstrate]:
         """Creates the Substrate instance based on provided arguments.
 
@@ -307,6 +312,7 @@ def _get_substrate(
             archive_endpoints: Similar to fallback_endpoints, but specifically only archive nodes. Will be used in cases
                 where you are requesting a block that is too old for your current (presumably lite) node. Defaults to
                 `None`
+            ws_shutdown_timer: Number of seconds to wait after last request before shutting down connection to the node
 
         Returns:
             the instance of the SubstrateInterface or RetrySyncSubstrate class.
@@ -322,6 +328,7 @@ def _get_substrate(
                 chain_name="Bittensor",
                 _mock=_mock,
                 archive_nodes=archive_endpoints,
+                ws_shutdown_timer=ws_shutdown_timer
             )
         return AsyncSubstrateInterface(
             url=self.chain_endpoint,
@@ -330,6 +337,7 @@ def _get_substrate(
             use_remote_preset=True,
             chain_name="Bittensor",
             _mock=_mock,
+            ws_shutdown_timer=ws_shutdown_timer
         )
 
     # Subtensor queries ===========================================================================================
diff --git a/bittensor/core/subtensor_api/__init__.py b/bittensor/core/subtensor_api/__init__.py
index 62ad7c1cda..bee892acd9 100644
--- a/bittensor/core/subtensor_api/__init__.py
+++ b/bittensor/core/subtensor_api/__init__.py
@@ -30,8 +30,9 @@ class SubtensorApi:
         log_verbose: Enables or disables verbose logging.
         mock: Whether this is a mock instance. Mainly just for use in testing.
         archive_endpoints: Similar to fallback_endpoints, but specifically only archive nodes. Will be used in cases
-                where you are requesting a block that is too old for your current (presumably lite) node. Defaults to
-                `None`
+            where you are requesting a block that is too old for your current (presumably lite) node. Defaults to `None`
+        websocket_shutdown_timer: Number of seconds to wait after last request before shutting down connection to
+            the node. Only applicable to AsyncSubtensor (and RetryAsyncSubtensor).
 
     Example:
         # sync version
@@ -79,6 +80,7 @@ def __init__(
         log_verbose: bool = False,
         mock: bool = False,
         archive_endpoints: Optional[list[str]] = None,
+        websocket_shutdown_timer: float = 5.0,
     ):
         self.network = network
         self._fallback_endpoints = fallback_endpoints
@@ -88,6 +90,7 @@ def __init__(
         self.log_verbose = log_verbose
         self.is_async = async_subtensor
         self._config = config
+        self._ws_shutdown_timer = websocket_shutdown_timer
 
         # assigned only for async instance
         self.initialize = None
@@ -125,6 +128,7 @@ def _get_subtensor(self) -> Union["_Subtensor", "_AsyncSubtensor"]:
                 retry_forever=self._retry_forever,
                 _mock=self._mock,
                 archive_endpoints=self._archive_endpoints,
+                websocket_shutdown_timer=self._ws_shutdown_timer,
             )
             self.initialize = _subtensor.initialize
             return _subtensor