Corrects the stake fee calculation

bittensor_cli/src/bittensor/chain_data.py

@@ -1193,3 +1193,20 @@ def _fix_decoded(cls, decoded: dict) -> "MetagraphInfo":
                 for adphk in decoded["alpha_dividends_per_hotkey"]
             ],
         )
+
+
+@dataclass
+class SimSwapResult:
+    tao_amount: Balance
+    alpha_amount: Balance
+    tao_fee: Balance
+    alpha_fee: Balance
+
+    @classmethod
+    def from_dict(cls, d: dict, netuid: int) -> "SimSwapResult":
+        return cls(
+            tao_amount=Balance.from_rao(d["tao_amount"]).set_unit(0),
+            alpha_amount=Balance.from_rao(d["alpha_amount"]).set_unit(netuid),
+            tao_fee=Balance.from_rao(d["tao_fee"]).set_unit(0),
+            alpha_fee=Balance.from_rao(d["alpha_fee"]).set_unit(netuid),
+        )

bittensor_cli/src/bittensor/subtensor_interface.py

@@ -28,6 +28,7 @@
     DynamicInfo,
     SubnetState,
     MetagraphInfo,
+    SimSwapResult,
 )
 from bittensor_cli.src import DelegatesDetails
 from bittensor_cli.src.bittensor.balances import Balance, fixed_to_float
@@ -1502,59 +1503,79 @@ async def get_extrinsic_fee(self, call: GenericCall, keypair: Keypair) -> Balanc
         fee_dict = await self.substrate.get_payment_info(call, keypair)
         return Balance.from_rao(fee_dict["partial_fee"])
 
-    async def get_stake_fee(
+    async def sim_swap(
         self,
-        origin_hotkey_ss58: Optional[str],
-        origin_netuid: Optional[int],
-        origin_coldkey_ss58: str,
-        destination_hotkey_ss58: Optional[str],
-        destination_netuid: Optional[int],
-        destination_coldkey_ss58: str,
+        origin_netuid: int,
+        destination_netuid: int,
         amount: int,
         block_hash: Optional[str] = None,
-    ) -> Balance:
+    ) -> SimSwapResult:
         """
-        Calculates the fee for a staking operation.
-
-        :param origin_hotkey_ss58: SS58 address of source hotkey (None for new stake)
-        :param origin_netuid: Netuid of source subnet (None for new stake)
-        :param origin_coldkey_ss58: SS58 address of source coldkey
-        :param destination_hotkey_ss58: SS58 address of destination hotkey (None for removing stake)
-        :param destination_netuid: Netuid of destination subnet (None for removing stake)
-        :param destination_coldkey_ss58: SS58 address of destination coldkey
-        :param amount: Amount of stake to transfer in RAO
-        :param block_hash: Optional block hash at which to perform the calculation
-
-        :return: The calculated stake fee as a Balance object
-
-        When to use None:
+        Hits the SimSwap Runtime API to calculate the fee and result for a given transaction. This should be used
+        instead of get_stake_fee for staking fee calculations. The SimSwapResult contains the staking fees and expected
+        returned amounts of a given transaction. This does not include the transaction (extrinsic) fee.
 
-        1. Adding new stake (default fee):
-        - origin_hotkey_ss58 = None
-        - origin_netuid = None
-        - All other fields required
-
-        2. Removing stake (default fee):
-        - destination_hotkey_ss58 = None
-        - destination_netuid = None
-        - All other fields required
+        Args:
+            origin_netuid: Netuid of the source subnet (0 if new stake)
+            destination_netuid: Netuid of the destination subnet
+            amount: Amount to transfer in Rao
+            block_hash: The hash of the blockchain block number for the query.
 
-        For all other operations, no None values - provide all parameters:
-        3. Moving between subnets
-        4. Moving between hotkeys
-        5. Moving between coldkeys
+        Returns:
+            SimSwapResult object representing the result
         """
-
-        if origin_netuid is None:
-            origin_netuid = 0
-
-        fee_rate = await self.query("Swap", "FeeRate", [origin_netuid])
-        fee = amount * (fee_rate / U16_MAX)
-
-        result = Balance.from_rao(fee)
-        result.set_unit(origin_netuid)
-
-        return result
+        block_hash = block_hash or await self.substrate.get_chain_head()
+        if origin_netuid > 0 and destination_netuid > 0:
+            # for cross-subnet moves where neither origin nor destination is root
+            intermediate_result_, sn_price = await asyncio.gather(
+                self.query_runtime_api(
+                    "SwapRuntimeApi",
+                    "sim_swap_alpha_for_tao",
+                    params={"netuid": origin_netuid, "alpha": amount},
+                    block_hash=block_hash,
+                ),
+                self.get_subnet_price(origin_netuid, block_hash=block_hash),
+            )
+            intermediate_result = SimSwapResult.from_dict(
+                intermediate_result_, origin_netuid
+            )
+            result = SimSwapResult.from_dict(
+                await self.query_runtime_api(
+                    "SwapRuntimeApi",
+                    "sim_swap_tao_for_alpha",
+                    params={
+                        "netuid": destination_netuid,
+                        "tao": intermediate_result.tao_amount,
+                    },
+                    block_hash=block_hash,
+                ),
+                destination_netuid,
+            )
+            secondary_fee = (result.tao_fee * sn_price).set_unit(origin_netuid)
+            result.alpha_fee = result.alpha_fee + secondary_fee
+            return result
+        elif origin_netuid > 0:
+            # dynamic to tao
+            return SimSwapResult.from_dict(
+                await self.query_runtime_api(
+                    "SwapRuntimeApi",
+                    "sim_swap_alpha_for_tao",
+                    params={"netuid": origin_netuid, "alpha": amount},
+                    block_hash=block_hash,
+                ),
+                origin_netuid,
+            )
+        else:
+            # tao to dynamic or unstaked to staked tao (SN0)
+            return SimSwapResult.from_dict(
+                await self.query_runtime_api(
+                    "SwapRuntimeApi",
+                    "sim_swap_tao_for_alpha",
+                    params={"netuid": destination_netuid, "tao": amount},
+                    block_hash=block_hash,
+                ),
+                destination_netuid,
+            )
 
     async def get_scheduled_coldkey_swap(
         self,

bittensor_cli/src/commands/stake/add.py

@@ -347,17 +347,6 @@ async def stake_extrinsic(
                 return False
             remaining_wallet_balance -= amount_to_stake
 
-            # TODO this should be asyncio gathered before the for loop
-            stake_fee = await subtensor.get_stake_fee(
-                origin_hotkey_ss58=None,
-                origin_netuid=None,
-                origin_coldkey_ss58=wallet.coldkeypub.ss58_address,
-                destination_hotkey_ss58=hotkey[1],
-                destination_netuid=netuid,
-                destination_coldkey_ss58=wallet.coldkeypub.ss58_address,
-                amount=amount_to_stake.rao,
-            )
-
             # Calculate slippage
             # TODO: Update for V3, slippage calculation is significantly different in v3
             # try:
@@ -409,7 +398,13 @@ async def stake_extrinsic(
                     safe_staking_=safe_staking,
                 )
                 row_extension = []
-            received_amount = rate * (amount_to_stake - stake_fee - extrinsic_fee)
+            # TODO this should be asyncio gathered before the for loop
+            sim_swap = await subtensor.sim_swap(
+                origin_netuid=0,
+                destination_netuid=netuid,
+                amount=(amount_to_stake - extrinsic_fee).rao,
+            )
+            received_amount = sim_swap.alpha_amount
             # Add rows for the table
             base_row = [
                 str(netuid),  # netuid
@@ -418,7 +413,7 @@ async def stake_extrinsic(
                 str(rate)
                 + f" {Balance.get_unit(netuid)}/{Balance.get_unit(0)} ",  # rate
                 str(received_amount.set_unit(netuid)),  # received
-                str(stake_fee),  # fee
+                str(sim_swap.tao_fee),  # fee
                 str(extrinsic_fee),
                 # str(slippage_pct),  # slippage
             ] + row_extension

bittensor_cli/src/commands/stake/move.py

@@ -520,14 +520,10 @@ async def move_stake(
             "alpha_amount": amount_to_move_as_balance.rao,
         },
     )
-    stake_fee, extrinsic_fee = await asyncio.gather(
-        subtensor.get_stake_fee(
-            origin_hotkey_ss58=origin_hotkey,
+    sim_swap, extrinsic_fee = await asyncio.gather(
+        subtensor.sim_swap(
             origin_netuid=origin_netuid,
-            origin_coldkey_ss58=wallet.coldkeypub.ss58_address,
-            destination_hotkey_ss58=destination_hotkey,
             destination_netuid=destination_netuid,
-            destination_coldkey_ss58=wallet.coldkeypub.ss58_address,
             amount=amount_to_move_as_balance.rao,
         ),
         subtensor.get_extrinsic_fee(call, wallet.coldkeypub),
@@ -543,7 +539,9 @@ async def move_stake(
                 origin_hotkey=origin_hotkey,
                 destination_hotkey=destination_hotkey,
                 amount_to_move=amount_to_move_as_balance,
-                stake_fee=stake_fee,
+                stake_fee=sim_swap.alpha_fee
+                if origin_netuid != 0
+                else sim_swap.tao_fee,
                 extrinsic_fee=extrinsic_fee,
             )
         except ValueError:
@@ -709,14 +707,10 @@ async def transfer_stake(
             "alpha_amount": amount_to_transfer.rao,
         },
     )
-    stake_fee, extrinsic_fee = await asyncio.gather(
-        subtensor.get_stake_fee(
-            origin_hotkey_ss58=origin_hotkey,
+    sim_swap, extrinsic_fee = await asyncio.gather(
+        subtensor.sim_swap(
             origin_netuid=origin_netuid,
-            origin_coldkey_ss58=wallet.coldkeypub.ss58_address,
-            destination_hotkey_ss58=origin_hotkey,
             destination_netuid=dest_netuid,
-            destination_coldkey_ss58=dest_coldkey_ss58,
             amount=amount_to_transfer.rao,
         ),
         subtensor.get_extrinsic_fee(call, wallet.coldkeypub),
@@ -732,7 +726,9 @@ async def transfer_stake(
                 origin_hotkey=origin_hotkey,
                 destination_hotkey=origin_hotkey,
                 amount_to_move=amount_to_transfer,
-                stake_fee=stake_fee,
+                stake_fee=sim_swap.alpha_fee
+                if origin_netuid != 0
+                else sim_swap.tao_fee,
                 extrinsic_fee=extrinsic_fee,
             )
         except ValueError:
@@ -880,14 +876,10 @@ async def swap_stake(
             "alpha_amount": amount_to_swap.rao,
         },
     )
-    stake_fee, extrinsic_fee = await asyncio.gather(
-        subtensor.get_stake_fee(
-            origin_hotkey_ss58=hotkey_ss58,
+    sim_swap, extrinsic_fee = await asyncio.gather(
+        subtensor.sim_swap(
             origin_netuid=origin_netuid,
-            origin_coldkey_ss58=wallet.coldkeypub.ss58_address,
-            destination_hotkey_ss58=hotkey_ss58,
             destination_netuid=destination_netuid,
-            destination_coldkey_ss58=wallet.coldkeypub.ss58_address,
             amount=amount_to_swap.rao,
         ),
         subtensor.get_extrinsic_fee(call, wallet.coldkeypub),
@@ -903,7 +895,9 @@ async def swap_stake(
                 origin_hotkey=hotkey_ss58,
                 destination_hotkey=hotkey_ss58,
                 amount_to_move=amount_to_swap,
-                stake_fee=stake_fee,
+                stake_fee=sim_swap.alpha_fee
+                if origin_netuid != 0
+                else sim_swap.tao_fee,
                 extrinsic_fee=extrinsic_fee,
             )
         except ValueError:

bittensor_cli/src/commands/stake/remove.py

@@ -200,16 +200,6 @@ async def unstake(
                 )
                 continue  # Skip to the next subnet - useful when single amount is specified for all subnets
 
-            stake_fee = await subtensor.get_stake_fee(
-                origin_hotkey_ss58=staking_address_ss58,
-                origin_netuid=netuid,
-                origin_coldkey_ss58=wallet.coldkeypub.ss58_address,
-                destination_hotkey_ss58=None,
-                destination_netuid=None,
-                destination_coldkey_ss58=wallet.coldkeypub.ss58_address,
-                amount=amount_to_unstake_as_balance.rao,
-            )
-
             try:
                 current_price = subnet_info.price.tao
                 if safe_staking:
@@ -240,10 +230,10 @@ async def unstake(
                         netuid=netuid,
                         amount=amount_to_unstake_as_balance,
                     )
-                rate = current_price
-                received_amount = (
-                    (amount_to_unstake_as_balance - stake_fee) * rate
-                ) - extrinsic_fee
+                sim_swap = await subtensor.sim_swap(
+                    netuid, 0, amount_to_unstake_as_balance.rao
+                )
+                received_amount = sim_swap.tao_amount - extrinsic_fee
             except ValueError:
                 continue
             total_received_amount += received_amount
@@ -266,7 +256,7 @@ async def unstake(
                 str(amount_to_unstake_as_balance),  # Amount to Unstake
                 f"{subnet_info.price.tao:.6f}"
                 + f"(τ/{Balance.get_unit(netuid)})",  # Rate
-                str(stake_fee.set_unit(netuid)),  # Fee
+                str(sim_swap.alpha_fee),  # Fee
                 str(extrinsic_fee),  # Extrinsic fee
                 str(received_amount),  # Received Amount
                 # slippage_pct,  # Slippage Percent
@@ -494,15 +484,6 @@ async def unstake_all(
             hotkey_display = hotkey_names.get(stake.hotkey_ss58, stake.hotkey_ss58)
             subnet_info = all_sn_dynamic_info.get(stake.netuid)
             stake_amount = stake.stake
-            stake_fee = await subtensor.get_stake_fee(
-                origin_hotkey_ss58=stake.hotkey_ss58,
-                origin_netuid=stake.netuid,
-                origin_coldkey_ss58=wallet.coldkeypub.ss58_address,
-                destination_hotkey_ss58=None,
-                destination_netuid=None,
-                destination_coldkey_ss58=wallet.coldkeypub.ss58_address,
-                amount=stake_amount.rao,
-            )
 
             try:
                 current_price = subnet_info.price.tao
@@ -515,8 +496,8 @@ async def unstake_all(
                     subtensor,
                     hotkey_ss58=stake.hotkey_ss58,
                 )
-                rate = current_price
-                received_amount = ((stake_amount - stake_fee) * rate) - extrinsic_fee
+                sim_swap = await subtensor.sim_swap(stake.netuid, 0, stake_amount.rao)
+                received_amount = sim_swap.tao_amount - extrinsic_fee
 
                 if received_amount < Balance.from_tao(0):
                     print_error("Not enough Alpha to pay the transaction fee.")
@@ -532,7 +513,7 @@ async def unstake_all(
                 str(stake_amount),
                 f"{float(subnet_info.price):.6f}"
                 + f"({Balance.get_unit(0)}/{Balance.get_unit(stake.netuid)})",
-                str(stake_fee),
+                str(sim_swap.alpha_fee),
                 str(extrinsic_fee),
                 str(received_amount),
             )