Fix cross-subnet received amount display using sim_swap and add missing proxy check

Author: bitloi

bittensor_cli/src/commands/stake/move.py

@@ -28,7 +28,7 @@
 
 if TYPE_CHECKING:
     from bittensor_cli.src.bittensor.subtensor_interface import SubtensorInterface
-    from bittensor_cli.src.bittensor.chain_data import DynamicInfo
+    from bittensor_cli.src.bittensor.chain_data import DynamicInfo, SimSwapResult
 
 MIN_STAKE_FEE = Balance.from_rao(50_000)
 
@@ -99,6 +99,7 @@ async def display_stake_movement_cross_subnets(
     destination_hotkey: str,
     amount_to_move: Balance,
     pricing: MovementPricing,
+    sim_swap: "SimSwapResult",
     stake_fee: Balance,
     extrinsic_fee: Balance,
     safe_staking: bool = False,
@@ -116,6 +117,7 @@ async def display_stake_movement_cross_subnets(
         destination_hotkey: The destination hotkey SS58 address.
         amount_to_move: The amount of stake to move/swap.
         pricing: Pricing information including rates and limits.
+        sim_swap: SimSwapResult from the runtime API with accurate swap amounts.
         stake_fee: The fee for the stake transaction.
         extrinsic_fee: The fee for the extrinsic execution.
         safe_staking: Whether to enable safe staking.
@@ -147,12 +149,12 @@ async def display_stake_movement_cross_subnets(
             + f"({Balance.get_unit(0)}/{Balance.get_unit(origin_netuid)})"
         )
     else:
-        dynamic_origin = pricing.origin_subnet
-        dynamic_destination = pricing.destination_subnet
-        received_amount_tao = (
-            dynamic_origin.alpha_to_tao(amount_to_move - stake_fee) - extrinsic_fee
-        )
-        received_amount = dynamic_destination.tao_to_alpha(received_amount_tao)
+        received_amount = sim_swap.alpha_amount
+        if not proxy:
+            extrinsic_fee_as_alpha = pricing.destination_subnet.tao_to_alpha(
+                extrinsic_fee
+            )
+            received_amount = received_amount - extrinsic_fee_as_alpha
         received_amount.set_unit(destination_netuid)
 
         if received_amount < Balance.from_tao(0).set_unit(destination_netuid):
@@ -670,6 +672,7 @@ async def move_stake(
                 destination_hotkey=destination_hotkey,
                 amount_to_move=amount_to_move_as_balance,
                 pricing=pricing,
+                sim_swap=sim_swap,
                 stake_fee=sim_swap.alpha_fee
                 if origin_netuid != 0
                 else sim_swap.tao_fee,
@@ -889,6 +892,7 @@ async def transfer_stake(
                 destination_hotkey=origin_hotkey,
                 amount_to_move=amount_to_transfer,
                 pricing=pricing,
+                sim_swap=sim_swap,
                 stake_fee=sim_swap.alpha_fee
                 if origin_netuid != 0
                 else sim_swap.tao_fee,
@@ -1116,6 +1120,7 @@ async def swap_stake(
                 destination_hotkey=hotkey_ss58,
                 amount_to_move=amount_to_swap,
                 pricing=pricing,
+                sim_swap=sim_swap,
                 stake_fee=sim_swap.alpha_fee
                 if origin_netuid != 0
                 else sim_swap.tao_fee,

tests/unit_tests/test_stake_move_display.py

@@ -0,0 +1,234 @@
+"""
+Unit tests for display_stake_movement_cross_subnets in stake/move.py.
+
+Covers:
+  - Cross-subnet received amount uses sim_swap.alpha_amount (not linear math)
+  - Cross-subnet with proxy does not deduct extrinsic fee from received
+  - Cross-subnet without proxy deducts extrinsic fee from received
+  - Same-subnet still uses existing linear pricing
+  - Cross-subnet raises ValueError when received amount is negative
+"""
+
+from types import SimpleNamespace
+from unittest.mock import MagicMock, patch
+
+import pytest
+
+from bittensor_cli.src.bittensor.balances import Balance
+from bittensor_cli.src.commands.stake.move import (
+    display_stake_movement_cross_subnets,
+    MovementPricing,
+)
+
+MODULE = "bittensor_cli.src.commands.stake.move"
+
+
+def _make_subnet(netuid: int, price_tao: float):
+    """Build a mock DynamicInfo with working alpha_to_tao / tao_to_alpha."""
+    subnet = MagicMock()
+    subnet.price = Balance.from_tao(price_tao)
+    subnet.is_dynamic = netuid != 0
+    subnet.netuid = netuid
+
+    def alpha_to_tao(alpha: Balance) -> Balance:
+        return Balance.from_tao(alpha.tao * price_tao)
+
+    def tao_to_alpha(tao: Balance) -> Balance:
+        if price_tao == 0:
+            return Balance.from_tao(0)
+        return Balance.from_tao(tao.tao / price_tao).set_unit(netuid)
+
+    subnet.alpha_to_tao = alpha_to_tao
+    subnet.tao_to_alpha = tao_to_alpha
+    return subnet
+
+
+def _make_sim_swap(
+    alpha_amount_tao: float, dest_netuid: int, alpha_fee_tao: float = 1.0
+):
+    """Build a SimpleNamespace matching SimSwapResult shape."""
+    return SimpleNamespace(
+        alpha_amount=Balance.from_tao(alpha_amount_tao).set_unit(dest_netuid),
+        tao_amount=Balance.from_tao(alpha_amount_tao),
+        alpha_fee=Balance.from_tao(alpha_fee_tao).set_unit(dest_netuid),
+        tao_fee=Balance.from_tao(alpha_fee_tao),
+    )
+
+
+# ---------------------------------------------------------------------------
+# Cross-subnet tests
+# ---------------------------------------------------------------------------
+
+
+class TestCrossSubnetDisplay:
+    @pytest.mark.asyncio
+    async def test_received_amount_uses_sim_swap_not_linear_math(self):
+        """The cross-subnet received amount must come from sim_swap.alpha_amount,
+        not from linear alpha_to_tao/tao_to_alpha calculations."""
+        origin_netuid, dest_netuid = 1, 2
+        # Price deliberately set so linear math would give a very different result
+        origin_subnet = _make_subnet(origin_netuid, price_tao=2.0)
+        dest_subnet = _make_subnet(dest_netuid, price_tao=0.5)
+        pricing = MovementPricing(
+            origin_subnet=origin_subnet,
+            destination_subnet=dest_subnet,
+            rate=4.0,
+            rate_with_tolerance=None,
+        )
+        amount = Balance.from_tao(10.0).set_unit(origin_netuid)
+        stake_fee = Balance.from_tao(0.5).set_unit(origin_netuid)
+        extrinsic_fee = Balance.from_tao(0.0)
+        # sim_swap says user receives 35 alpha on dest — linear math would give ~38
+        sim_swap = _make_sim_swap(alpha_amount_tao=35.0, dest_netuid=dest_netuid)
+
+        with patch(f"{MODULE}.console"):
+            received, _ = await display_stake_movement_cross_subnets(
+                subtensor=MagicMock(network="test"),
+                origin_netuid=origin_netuid,
+                destination_netuid=dest_netuid,
+                origin_hotkey="5C" + "a" * 46,
+                destination_hotkey="5C" + "b" * 46,
+                amount_to_move=amount,
+                pricing=pricing,
+                sim_swap=sim_swap,
+                stake_fee=stake_fee,
+                extrinsic_fee=extrinsic_fee,
+                proxy="5C" + "p" * 46,  # proxy → no extrinsic_fee deduction
+            )
+
+        assert received.tao == pytest.approx(35.0, abs=1e-6)
+
+    @pytest.mark.asyncio
+    async def test_proxy_does_not_deduct_extrinsic_fee(self):
+        """With a proxy, the extrinsic fee should not reduce the received amount."""
+        origin_netuid, dest_netuid = 1, 2
+        dest_subnet = _make_subnet(dest_netuid, price_tao=1.0)
+        pricing = MovementPricing(
+            origin_subnet=_make_subnet(origin_netuid, price_tao=1.0),
+            destination_subnet=dest_subnet,
+            rate=1.0,
+            rate_with_tolerance=None,
+        )
+        sim_swap = _make_sim_swap(alpha_amount_tao=50.0, dest_netuid=dest_netuid)
+        extrinsic_fee = Balance.from_tao(0.5)
+
+        with patch(f"{MODULE}.console"):
+            received, _ = await display_stake_movement_cross_subnets(
+                subtensor=MagicMock(network="test"),
+                origin_netuid=origin_netuid,
+                destination_netuid=dest_netuid,
+                origin_hotkey="5C" + "a" * 46,
+                destination_hotkey="5C" + "b" * 46,
+                amount_to_move=Balance.from_tao(50).set_unit(origin_netuid),
+                pricing=pricing,
+                sim_swap=sim_swap,
+                stake_fee=Balance.from_tao(0),
+                extrinsic_fee=extrinsic_fee,
+                proxy="5C" + "p" * 46,
+            )
+
+        # Full sim_swap amount — extrinsic fee NOT deducted
+        assert received.tao == pytest.approx(50.0, abs=1e-6)
+
+    @pytest.mark.asyncio
+    async def test_no_proxy_deducts_extrinsic_fee(self):
+        """Without a proxy, the extrinsic fee should reduce the received amount."""
+        origin_netuid, dest_netuid = 1, 2
+        dest_subnet = _make_subnet(dest_netuid, price_tao=1.0)
+        pricing = MovementPricing(
+            origin_subnet=_make_subnet(origin_netuid, price_tao=1.0),
+            destination_subnet=dest_subnet,
+            rate=1.0,
+            rate_with_tolerance=None,
+        )
+        sim_swap = _make_sim_swap(alpha_amount_tao=50.0, dest_netuid=dest_netuid)
+        extrinsic_fee = Balance.from_tao(0.5)
+
+        with patch(f"{MODULE}.console"):
+            received, _ = await display_stake_movement_cross_subnets(
+                subtensor=MagicMock(network="test"),
+                origin_netuid=origin_netuid,
+                destination_netuid=dest_netuid,
+                origin_hotkey="5C" + "a" * 46,
+                destination_hotkey="5C" + "b" * 46,
+                amount_to_move=Balance.from_tao(50).set_unit(origin_netuid),
+                pricing=pricing,
+                sim_swap=sim_swap,
+                stake_fee=Balance.from_tao(0),
+                extrinsic_fee=extrinsic_fee,
+            )
+
+        # Extrinsic fee converted to dest alpha (price=1.0 so 0.5 TAO → 0.5 alpha)
+        assert received.tao == pytest.approx(49.5, abs=1e-6)
+
+    @pytest.mark.asyncio
+    async def test_negative_received_raises_value_error(self):
+        """When fees exceed the swap result, ValueError must be raised."""
+        origin_netuid, dest_netuid = 1, 2
+        dest_subnet = _make_subnet(dest_netuid, price_tao=1.0)
+        pricing = MovementPricing(
+            origin_subnet=_make_subnet(origin_netuid, price_tao=1.0),
+            destination_subnet=dest_subnet,
+            rate=1.0,
+            rate_with_tolerance=None,
+        )
+        # Tiny swap result, large extrinsic fee → negative received
+        sim_swap = _make_sim_swap(alpha_amount_tao=0.001, dest_netuid=dest_netuid)
+        extrinsic_fee = Balance.from_tao(1.0)
+
+        with patch(f"{MODULE}.console"), pytest.raises(ValueError):
+            await display_stake_movement_cross_subnets(
+                subtensor=MagicMock(network="test"),
+                origin_netuid=origin_netuid,
+                destination_netuid=dest_netuid,
+                origin_hotkey="5C" + "a" * 46,
+                destination_hotkey="5C" + "b" * 46,
+                amount_to_move=Balance.from_tao(1).set_unit(origin_netuid),
+                pricing=pricing,
+                sim_swap=sim_swap,
+                stake_fee=Balance.from_tao(0),
+                extrinsic_fee=extrinsic_fee,
+            )
+
+
+# ---------------------------------------------------------------------------
+# Same-subnet tests (behaviour must be unchanged)
+# ---------------------------------------------------------------------------
+
+
+class TestSameSubnetDisplay:
+    @pytest.mark.asyncio
+    async def test_same_subnet_uses_linear_pricing(self):
+        """Same-subnet moves use linear alpha_to_tao/tao_to_alpha, not sim_swap."""
+        netuid = 3
+        price = 2.0
+        subnet = _make_subnet(netuid, price_tao=price)
+        pricing = MovementPricing(
+            origin_subnet=subnet,
+            destination_subnet=subnet,
+            rate=1.0,
+            rate_with_tolerance=None,
+        )
+        amount = Balance.from_tao(10.0).set_unit(netuid)
+        stake_fee = Balance.from_tao(0.5).set_unit(netuid)
+        extrinsic_fee = Balance.from_tao(0.0)
+        # sim_swap with a wildly different alpha_amount to prove it's not used
+        sim_swap = _make_sim_swap(alpha_amount_tao=999.0, dest_netuid=netuid)
+
+        with patch(f"{MODULE}.console"):
+            received, _ = await display_stake_movement_cross_subnets(
+                subtensor=MagicMock(network="test"),
+                origin_netuid=netuid,
+                destination_netuid=netuid,
+                origin_hotkey="5C" + "a" * 46,
+                destination_hotkey="5C" + "b" * 46,
+                amount_to_move=amount,
+                pricing=pricing,
+                sim_swap=sim_swap,
+                stake_fee=stake_fee,
+                extrinsic_fee=extrinsic_fee,
+                proxy="5C" + "p" * 46,
+            )
+
+        # Linear: (10 - 0.5) * 2.0 / 2.0 = 9.5 (proxy → no extrinsic fee deduction)
+        assert received.tao == pytest.approx(9.5, abs=1e-6)