Adds safe to stake swap

Author: ibraheem-abe

bittensor/core/extrinsics/move_stake.py

@@ -157,6 +157,9 @@ def swap_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:
     """
     Moves stake between subnets while keeping the same coldkey-hotkey pair ownership.
@@ -170,6 +173,9 @@ def swap_stake_extrinsic(
         amount (Union[Balance, float]): Amount to swap.
         wait_for_inclusion (bool): If true, waits for inclusion before returning.
         wait_for_finalization (bool): If true, waits for finalization before returning.
+        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:
         success (bool): True if the swap was successful.
@@ -203,20 +209,45 @@ 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 = subtensor.subnet(netuid=origin_netuid)
+            destination_pool = 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 [yellow]{origin_netuid}[/yellow] to netuid "
+                f"[yellow]{destination_netuid}[/yellow]\n"
+                f"Current price ratio: [green]{swap_rate_ratio:.4f}[/green], "
+                f"Ratio with tolerance: [yellow]{swap_rate_ratio_with_tolerance:.4f}[/yellow]"
+            )
+            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 [yellow]{origin_netuid}[/yellow] to netuid "
+                f"[yellow]{destination_netuid}[/yellow]"
+            )
+            call_function = "swap_stake"
+
         call = 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 = subtensor.sign_and_send_extrinsic(
@@ -251,7 +282,12 @@ 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: