safe stake added to async

Author: ibraheem-abe

bittensor/core/extrinsics/asyncex/move_stake.py

@@ -160,6 +160,9 @@ async def swap_stake_extrinsic(
     amount: Balance,
     wait_for_inclusion: bool = True,
     wait_for_finalization: bool = False,
+    safe_staking: bool = False,
+    allow_partial_stake: bool = False,
+    rate_threshold: float = 0.005,
 ) -> bool:
     """
     Swaps stake from one subnet to another for a given hotkey in the Bittensor network.
@@ -173,6 +176,9 @@ async def swap_stake_extrinsic(
         amount (Balance): The amount of stake to swap as a `Balance` object.
         wait_for_inclusion (bool): If True, waits for transaction inclusion in a block. Defaults to True.
         wait_for_finalization (bool): If True, waits for transaction finalization. Defaults to False.
+        safe_staking (bool): If true, enables price safety checks to protect against price impact.
+        allow_partial_stake (bool): If true, allows partial stake swaps when the full amount would exceed the price threshold.
+        rate_threshold (float): Maximum allowed increase in price ratio (0.005 = 0.5%).
 
     Returns:
         bool: True if the swap was successful, False otherwise.
@@ -205,20 +211,47 @@ async def swap_stake_extrinsic(
         return False
 
     try:
-        logging.info(
-            f"Swapping stake for hotkey [blue]{hotkey_ss58}[/blue]\n"
-            f"Amount: [green]{amount}[/green] from netuid [yellow]{origin_netuid}[/yellow] to netuid "
-            f"[yellow]{destination_netuid}[/yellow]"
-        )
+        call_params = {
+            "hotkey": hotkey_ss58,
+            "origin_netuid": origin_netuid,
+            "destination_netuid": destination_netuid,
+            "alpha_amount": amount.rao,
+        }
+
+        if safe_staking:
+            origin_pool, destination_pool = await asyncio.gather(
+                subtensor.subnet(netuid=origin_netuid),
+                subtensor.subnet(netuid=destination_netuid),
+            )
+            swap_rate_ratio = origin_pool.price.rao / destination_pool.price.rao
+            swap_rate_ratio_with_tolerance = swap_rate_ratio * (1 + rate_threshold)
+
+            logging.info(
+                f"Swapping stake with safety for hotkey [blue]{hotkey_ss58}[/blue]\n"
+                f"Amount: [green]{amount}[/green] from netuid [green]{origin_netuid}[/green] to netuid "
+                f"[green]{destination_netuid}[/green]\n"
+                f"Current price ratio: [green]{swap_rate_ratio:.4f}[/green], "
+                f"Ratio with tolerance: [green]{swap_rate_ratio_with_tolerance:.4f}[/green]"
+            )
+            call_params.update(
+                {
+                    "limit_price": swap_rate_ratio_with_tolerance,
+                    "allow_partial": allow_partial_stake,
+                }
+            )
+            call_function = "swap_stake_limit"
+        else:
+            logging.info(
+                f"Swapping stake for hotkey [blue]{hotkey_ss58}[/blue]\n"
+                f"Amount: [green]{amount}[/green] from netuid [green]{origin_netuid}[/green] to netuid "
+                f"[green]{destination_netuid}[/green]"
+            )
+            call_function = "swap_stake"
+
         call = await subtensor.substrate.compose_call(
             call_module="SubtensorModule",
-            call_function="swap_stake",
-            call_params={
-                "hotkey": hotkey_ss58,
-                "origin_netuid": origin_netuid,
-                "destination_netuid": destination_netuid,
-                "alpha_amount": amount.rao,
-            },
+            call_function=call_function,
+            call_params=call_params,
         )
 
         success, err_msg = await subtensor.sign_and_send_extrinsic(
@@ -253,7 +286,12 @@ async def swap_stake_extrinsic(
 
             return True
         else:
-            logging.error(f":cross_mark: [red]Failed[/red]: {err_msg}")
+            if safe_staking and "Custom error: 8" in err_msg:
+                logging.error(
+                    ":cross_mark: [red]Failed[/red]: Price ratio exceeded tolerance limit. Either increase price tolerance or enable partial staking."
+                )
+            else:
+                logging.error(f":cross_mark: [red]Failed[/red]: {err_msg}")
             return False
 
     except Exception as e:

bittensor/core/extrinsics/asyncex/staking.py

@@ -21,6 +21,9 @@ async def add_stake_extrinsic(
     amount: Optional[Balance] = None,
     wait_for_inclusion: bool = True,
     wait_for_finalization: bool = False,
+    safe_staking: bool = False,
+    allow_partial_stake: bool = False,
+    rate_threshold: float = 0.005,
 ) -> bool:
     """
     Adds the specified amount of stake to passed hotkey `uid`.
@@ -36,6 +39,9 @@ async def add_stake_extrinsic(
             `False` if the extrinsic fails to enter the block within the timeout.
         wait_for_finalization: If set, waits for the extrinsic to be finalized on the chain before returning `True`,
             or returns `False` if the extrinsic fails to be finalized within the timeout.
+        safe_staking: If set, uses safe staking logic
+        allow_partial_stake: If set, allows partial stake
+        rate_threshold: The rate threshold for safe staking
 
     Returns:
         success: Flag is `True` if extrinsic was finalized or included in the block. If we did not wait for
@@ -97,19 +103,48 @@ async def add_stake_extrinsic(
         return False
 
     try:
-        logging.info(
-            f":satellite: [magenta]Staking to:[/magenta] "
-            f"[blue]netuid: {netuid}, amount: {staking_balance} "
-            f"on {subtensor.network}[/blue] [magenta]...[/magenta]"
-        )
+        call_params = {
+            "hotkey": hotkey_ss58,
+            "netuid": netuid,
+            "amount_staked": staking_balance.rao,
+        }
+
+        if safe_staking:
+            pool = await subtensor.subnet(netuid=netuid)
+            base_price = pool.price.rao
+            price_with_tolerance = base_price * (1 + rate_threshold)
+            call_params.update(
+                {
+                    "limit_price": price_with_tolerance,
+                    "allow_partial": allow_partial_stake,
+                }
+            )
+            call_function = "add_stake_limit"
+
+            # For logging
+            base_rate = pool.price.tao
+            rate_with_tolerance = base_rate * (1 + rate_threshold)
+            logging.info(
+                f":satellite: [magenta]Safe Staking to:[/magenta] "
+                f"[blue]netuid: [green]{netuid}[/green], amount: [green]{staking_balance}[/green], "
+                f"tolerance percentage: [green]{rate_threshold*100}%[/green], "
+                f"price limit: [green]{rate_with_tolerance}[/green], "
+                f"original price: [green]{base_rate}[/green], "
+                f"with partial stake: [green]{allow_partial_stake}[/green] "
+                f"on [blue]{subtensor.network}[/blue][/magenta]...[/magenta]"
+            )
+        else:
+            logging.info(
+                f":satellite: [magenta]Staking to:[/magenta] "
+                f"[blue]netuid: [green]{netuid}[/green], amount: [green]{staking_balance}[/green] "
+                f"on [blue]{subtensor.network}[/blue][magenta]...[/magenta]"
+            )
+            call_function = "add_stake"
+
         call = await subtensor.substrate.compose_call(
             call_module="SubtensorModule",
-            call_function="add_stake",
-            call_params={
-                "hotkey": hotkey_ss58,
-                "amount_staked": staking_balance.rao,
-                "netuid": netuid,
-            },
+            call_function=call_function,
+            call_params=call_params,
         )
         staking_response, err_msg = await subtensor.sign_and_send_extrinsic(
             call,
@@ -152,7 +187,12 @@ async def add_stake_extrinsic(
             )
             return True
         else:
-            logging.error(f":cross_mark: [red]Failed: {err_msg}.[/red]")
+            if safe_staking and "Custom error: 8" in err_msg:
+                logging.error(
+                    ":cross_mark: [red]Failed[/red]: Price exceeded tolerance limit. Either increase price tolerance or enable partial staking."
+                )
+            else:
+                logging.error(f":cross_mark: [red]Failed: {err_msg}.[/red]")
             return False
 
     except NotRegisteredError:

bittensor/core/extrinsics/asyncex/unstaking.py

@@ -20,6 +20,9 @@ async def unstake_extrinsic(
     amount: Optional[Balance] = None,
     wait_for_inclusion: bool = True,
     wait_for_finalization: bool = False,
+    safe_staking: bool = False,
+    allow_partial_stake: bool = False,
+    rate_threshold: float = 0.005,
 ) -> bool:
     """Removes stake into the wallet coldkey from the specified hotkey ``uid``.
 
@@ -34,6 +37,9 @@ async def unstake_extrinsic(
             returns ``False`` if the extrinsic fails to enter the block within the timeout.
         wait_for_finalization (bool): If set, waits for the extrinsic to be finalized on the chain before returning
             ``True``, or returns ``False`` if the extrinsic fails to be finalized within the timeout.
+        safe_staking: If true, enables price safety checks
+        allow_partial_stake: If true, allows partial unstaking if price threshold exceeded
+        rate_threshold: Maximum allowed price decrease percentage (0.005 = 0.5%)
 
     Returns:
         success (bool): Flag is ``True`` if extrinsic was finalized or included in the block. If we did not wait for
@@ -83,19 +89,49 @@ async def unstake_extrinsic(
         return False
 
     try:
-        logging.info(
-            f"Unstaking [blue]{unstaking_balance}[/blue] from hotkey: [magenta]{hotkey_ss58}[/magenta] on netuid: "
-            f"[blue]{netuid}[/blue]"
-        )
+        call_params = {
+            "hotkey": hotkey_ss58,
+            "netuid": netuid,
+            "amount_unstaked": unstaking_balance.rao,
+        }
+        if safe_staking:
+            pool = await subtensor.subnet(netuid=netuid)
+            base_price = pool.price.rao
+            price_with_tolerance = base_price * (1 - rate_threshold)
+
+            # For logging
+            base_rate = pool.price.tao
+            rate_with_tolerance = base_rate * (1 - rate_threshold)
+
+            logging.info(
+                f":satellite: [magenta]Safe Unstaking from:[/magenta] "
+                f"netuid: [green]{netuid}[/green], amount: [green]{unstaking_balance}[/green], "
+                f"tolerance percentage: [green]{rate_threshold*100}%[/green], "
+                f"price limit: [green]{rate_with_tolerance}[/green], "
+                f"original price: [green]{base_rate}[/green], "
+                f"with partial unstake: [green]{allow_partial_stake}[/green] "
+                f"on [blue]{subtensor.network}[/blue][magenta]...[/magenta]"
+            )
+
+            call_params.update(
+                {
+                    "limit_price": price_with_tolerance,
+                    "allow_partial": allow_partial_stake,
+                }
+            )
+            call_function = "remove_stake_limit"
+        else:
+            logging.info(
+                f":satellite: [magenta]Unstaking from:[/magenta] "
+                f"netuid: [green]{netuid}[/green], amount: [green]{unstaking_balance}[/green] "
+                f"on [blue]{subtensor.network}[/blue][magenta]...[/magenta]"
+            )
+            call_function = "remove_stake"
 
         call = await subtensor.substrate.compose_call(
             call_module="SubtensorModule",
-            call_function="remove_stake",
-            call_params={
-                "hotkey": hotkey_ss58,
-                "amount_unstaked": unstaking_balance.rao,
-                "netuid": netuid,
-            },
+            call_function=call_function,
+            call_params=call_params,
         )
         staking_response, err_msg = await subtensor.sign_and_send_extrinsic(
             call,
@@ -138,7 +174,12 @@ async def unstake_extrinsic(
             )
             return True
         else:
-            logging.error(f":cross_mark: [red]Failed: {err_msg}.[/red]")
+            if safe_staking and "Custom error: 8" in err_msg:
+                logging.error(
+                    ":cross_mark: [red]Failed[/red]: Price exceeded tolerance limit. Either increase price tolerance or enable partial staking."
+                )
+            else:
+                logging.error(f":cross_mark: [red]Failed: {err_msg}.[/red]")
             return False
 
     except NotRegisteredError: