From 3e7ea6be432ab0de09a911c626d888116362b24d Mon Sep 17 00:00:00 2001
From: Roman <roman@opentensor.dev>
Date: Thu, 8 May 2025 17:57:00 -0500
Subject: [PATCH 1/2] improve get_next_epoch_start_block

---
 bittensor/core/async_subtensor.py | 31 ++++++++++++++++++++++++-------
 bittensor/core/subtensor.py       |  6 +++++-
 2 files changed, 29 insertions(+), 8 deletions(-)

diff --git a/bittensor/core/async_subtensor.py b/bittensor/core/async_subtensor.py
index 7b8a87b37f..edf83cf78b 100644
--- a/bittensor/core/async_subtensor.py
+++ b/bittensor/core/async_subtensor.py
@@ -616,19 +616,31 @@ async def all_subnets(
         return subnets
 
     async def blocks_since_last_step(
-        self, netuid: int, block: Optional[int] = None
+        self,
+        netuid: int,
+        block: Optional[int] = None,
+        block_hash: Optional[str] = None,
+        reuse_block: bool = False,
     ) -> Optional[int]:
         """Returns number of blocks since the last epoch of the subnet.
 
         Arguments:
             netuid (int): The unique identifier of the subnetwork.
             block: the block number for this query.
+            block_hash: The hash of the blockchain block number for the query. Do not specify if using reuse_block or
+                block.
+            reuse_block: Whether to reuse the last-used blockchain block hash. Do not set if using block_hash or block.
+
 
         Returns:
             block number of the last step in the subnet.
         """
         query = await self.query_subtensor(
-            name="BlocksSinceLastStep", block=block, params=[netuid]
+            name="BlocksSinceLastStep",
+            block=block,
+            block_hash=block_hash,
+            reuse_block=reuse_block,
+            params=[netuid],
         )
         return query.value if query is not None and hasattr(query, "value") else query
 
@@ -1803,11 +1815,16 @@ async def get_next_epoch_start_block(
             int: The block number at which the next epoch will start.
         """
         block_hash = await self.determine_block_hash(block, block_hash, reuse_block)
-        if not block_hash and reuse_block:
-            block_hash = self.substrate.last_block_hash
-        block = await self.substrate.get_block_number(block_hash=block_hash)
-        tempo = await self.tempo(netuid=netuid, block_hash=block_hash)
-        return (((block // tempo) + 1) * tempo) + 1 if tempo else None
+        blocks_since_last_step = await self.blocks_since_last_step(
+            netuid=netuid, block=block, block_hash=block_hash, reuse_block=reuse_block
+        )
+        tempo = await self.tempo(
+            netuid=netuid, block=block, block_hash=block_hash, reuse_block=reuse_block
+        )
+
+        if block and blocks_since_last_step and tempo:
+            return block - blocks_since_last_step + tempo + 1
+        return None
 
     async def get_owned_hotkeys(
         self,
diff --git a/bittensor/core/subtensor.py b/bittensor/core/subtensor.py
index 47d5741113..796440cb90 100644
--- a/bittensor/core/subtensor.py
+++ b/bittensor/core/subtensor.py
@@ -1410,8 +1410,12 @@ def get_next_epoch_start_block(
             int: The block number at which the next epoch will start.
         """
         block = block or self.block
+        blocks_since_last_step = self.blocks_since_last_step(netuid=netuid, block=block)
         tempo = self.tempo(netuid=netuid, block=block)
-        return (((block // tempo) + 1) * tempo) + 1 if tempo else None
+
+        if block and blocks_since_last_step and tempo:
+            return block - blocks_since_last_step + tempo + 1
+        return None
 
     def get_owned_hotkeys(
         self,

From 29736b0ab4a303e2b36ab134e6234370e14666f0 Mon Sep 17 00:00:00 2001
From: Roman <roman@opentensor.dev>
Date: Thu, 8 May 2025 17:57:27 -0500
Subject: [PATCH 2/2] fix test

---
 tests/unit_tests/test_async_subtensor.py | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/tests/unit_tests/test_async_subtensor.py b/tests/unit_tests/test_async_subtensor.py
index b61df6dfaa..8c0c5c6288 100644
--- a/tests/unit_tests/test_async_subtensor.py
+++ b/tests/unit_tests/test_async_subtensor.py
@@ -3192,6 +3192,8 @@ async def test_blocks_since_last_step_with_value(subtensor, mocker):
     mocked_query_subtensor.assert_awaited_once_with(
         name="BlocksSinceLastStep",
         block=block,
+        block_hash=None,
+        reuse_block=False,
         params=[netuid],
     )
 
@@ -3214,6 +3216,8 @@ async def test_blocks_since_last_step_is_none(subtensor, mocker):
     mocked_query_subtensor.assert_awaited_once_with(
         name="BlocksSinceLastStep",
         block=block,
+        block_hash=None,
+        reuse_block=False,
         params=[netuid],
     )