Makes things work with the proxy logic

Author: thewhaleking

bittensor_cli/src/bittensor/extrinsics/mev_shield.py

@@ -57,56 +57,6 @@ async def encrypt_extrinsic(
     return encrypted_call
 
 
-async def create_mev_protected_extrinsic(
-    subtensor: "SubtensorInterface",
-    keypair: "Keypair",
-    call: "GenericCall",
-    nonce: int,
-    era: Optional[int] = None,
-) -> tuple["GenericExtrinsic", str]:
-    """
-    Create a MEV-protected extrinsic.
-
-    This function handles MEV Shield wrapping:
-    1. Fetches a future nonce (current_nonce + 1) & signs the inner call with it
-    2. Encrypts it into a wrapper call
-    3. Signs the wrapper with the current nonce
-
-    Args:
-        subtensor: The SubtensorInterface instance.
-        keypair: Keypair for signing.
-        call: The call to protect (e.g., add_stake).
-        nonce: The current account nonce (wrapper will use this, inner uses nonce+1).
-        era: The era period for the extrinsic.
-
-    Returns:
-        Tuple of (signed_shield_extrinsic, inner_extrinsic_hash).
-        The inner_extrinsic_hash is used for tracking actual extrinsic.
-    """
-
-    async def create_signed(call_to_sign, n):
-        kwargs = {
-            "call": call_to_sign,
-            "keypair": keypair,
-            "nonce": n,
-        }
-        if era is not None:
-            kwargs["era"] = {"period": era}
-        return await subtensor.substrate.create_signed_extrinsic(**kwargs)
-
-    next_nonce = await subtensor.substrate.get_account_next_index(keypair.ss58_address)
-
-    # Actual call: Sign with future nonce (current_nonce + 1)
-    inner_extrinsic = await create_signed(call, next_nonce)
-    inner_hash = f"0x{inner_extrinsic.extrinsic_hash.hex()}"
-
-    # MeV Shield wrapper: Sign with current nonce
-    shield_call = await encrypt_extrinsic(subtensor, inner_extrinsic)
-    shield_extrinsic = await create_signed(shield_call, nonce)
-
-    return shield_extrinsic, inner_hash
-
-
 async def extract_mev_shield_id(response: "AsyncExtrinsicReceipt") -> Optional[str]:
     """
     Extract the MEV Shield wrapper ID from an extrinsic response.

bittensor_cli/src/bittensor/subtensor_interface.py

@@ -34,6 +34,7 @@
 from bittensor_cli.src import DelegatesDetails
 from bittensor_cli.src.bittensor.balances import Balance, fixed_to_float
 from bittensor_cli.src import Constants, defaults, TYPE_REGISTRY
+from bittensor_cli.src.bittensor.extrinsics.mev_shield import encrypt_extrinsic
 from bittensor_cli.src.bittensor.utils import (
     format_error_message,
     console,
@@ -2406,6 +2407,7 @@ async def get_all_subnet_ema_tao_inflow(
             else:
                 _, raw_ema_value = value
                 ema_value = fixed_to_float(raw_ema_value)
+                # TODO @abe is this intentional: float passed as int for from_rao
                 ema_map[netuid] = Balance.from_rao(ema_value)
         return ema_map
 
@@ -2437,12 +2439,13 @@ async def get_subnet_ema_tao_inflow(
             return Balance.from_rao(0)
         _, raw_ema_value = value
         ema_value = fixed_to_float(raw_ema_value)
+        # TODO @abe this is a float, but we're passing it as an int for from_rao, is this intentional?
         return Balance.from_rao(ema_value)
 
     async def get_mev_shield_next_key(
         self,
         block_hash: Optional[str] = None,
-    ) -> Optional[tuple[bytes, int]]:
+    ) -> bytes:
         """
         Get the next MEV Shield public key and epoch from chain storage.
 
@@ -2470,7 +2473,7 @@ async def get_mev_shield_next_key(
     async def get_mev_shield_current_key(
         self,
         block_hash: Optional[str] = None,
-    ) -> Optional[tuple[bytes, int]]:
+    ) -> bytes:
         """
         Get the current MEV Shield public key and epoch from chain storage.
 

bittensor_cli/src/commands/stake/add.py

@@ -11,7 +11,6 @@
 from bittensor_cli.src import COLOR_PALETTE
 from bittensor_cli.src.bittensor.balances import Balance
 from bittensor_cli.src.bittensor.extrinsics.mev_shield import (
-    create_mev_protected_extrinsic,
     extract_mev_shield_id,
     wait_for_extrinsic_by_hash,
 )
@@ -141,14 +140,13 @@ async def safe_stake_extrinsic(
                 },
             ),
         )
-        if mev_protection:
-            call = await encrypt_call(subtensor, wallet, call)
         success_, err_msg, response = await subtensor.sign_and_send_extrinsic(
             call=call,
             wallet=wallet,
             nonce=next_nonce,
             era={"period": era},
             proxy=proxy,
+            mev_protection=mev_protection,
         )
         if not success_:
             if "Custom error: 8" in err_msg:
@@ -164,16 +162,20 @@ async def safe_stake_extrinsic(
             return False, err_msg, None
         else:
             if mev_protection:
+                inner_hash = err_msg
                 mev_shield_id = await extract_mev_shield_id(response)
-                if mev_shield_id:
-                    mev_success, mev_error, response = await wait_for_mev_execution(
-                        subtensor, mev_shield_id, response.block_hash, status=status_
-                    )
-                    if not mev_success:
-                        status_.stop()
-                        err_msg = f"{failure_prelude}: {mev_error}"
-                        err_out("\n" + err_msg)
-                        return False, err_msg, None
+                mev_success, mev_error, response = await wait_for_extrinsic_by_hash(
+                    subtensor=subtensor,
+                    extrinsic_hash=inner_hash,
+                    shield_id=mev_shield_id,
+                    submit_block_hash=response.block_hash,
+                    status=status_,
+                )
+                if not mev_success:
+                    status_.stop()
+                    err_msg = f"{failure_prelude}: {mev_error}"
+                    err_out("\n" + err_msg)
+                    return False, err_msg, None
             if json_output:
                 # the rest of this checking is not necessary if using json_output
                 return True, "", response
@@ -239,31 +241,34 @@ async def stake_extrinsic(
         failure_prelude = (
             f":cross_mark: [red]Failed[/red] to stake {amount} on Netuid {netuid_i}"
         )
-        if mev_protection:
-            call = await encrypt_call(subtensor, wallet, call)
         success_, err_msg, response = await subtensor.sign_and_send_extrinsic(
             call=call,
             wallet=wallet,
             nonce=next_nonce,
             era={"period": era},
             proxy=proxy,
+            mev_protection=mev_protection,
         )
         if not success_:
             err_msg = f"{failure_prelude} with error: {err_msg}"
             err_out("\n" + err_msg)
             return False, err_msg, None
         else:
             if mev_protection:
+                inner_hash = err_msg
                 mev_shield_id = await extract_mev_shield_id(response)
-                if mev_shield_id:
-                    mev_success, mev_error, response = await wait_for_mev_execution(
-                        subtensor, mev_shield_id, response.block_hash, status=status_
-                    )
-                    if not mev_success:
-                        status_.stop()
-                        err_msg = f"{failure_prelude}: {mev_error}"
-                        err_out("\n" + err_msg)
-                        return False, err_msg, None
+                mev_success, mev_error, response = await wait_for_extrinsic_by_hash(
+                    subtensor=subtensor,
+                    extrinsic_hash=inner_hash,
+                    shield_id=mev_shield_id,
+                    submit_block_hash=response.block_hash,
+                    status=status_,
+                )
+                if not mev_success:
+                    status_.stop()
+                    err_msg = f"{failure_prelude}: {mev_error}"
+                    err_out("\n" + err_msg)
+                    return False, err_msg, None
             if json_output:
                 # the rest of this is not necessary if using json_output
                 return True, "", response