diff --git a/bittensor_cli/cli.py b/bittensor_cli/cli.py index 9ad6c89df..8d911ffa6 100755 --- a/bittensor_cli/cli.py +++ b/bittensor_cli/cli.py @@ -346,6 +346,12 @@ def edit_help(cls, option_name: str, help_text: str): "--dashboard.path", help="Path to save the dashboard HTML file. For example: `~/.bittensor/dashboard`.", ) + mev_protection = typer.Option( + True, + "--mev-protection/--no-mev-protection", + show_default=False, + help="Enable or disable MEV protection [dim](default: enabled)[/dim].", + ) json_output = typer.Option( False, "--json-output", @@ -3972,6 +3978,7 @@ def stake_add( rate_tolerance: Optional[float] = Options.rate_tolerance, safe_staking: Optional[bool] = Options.safe_staking, allow_partial_stake: Optional[bool] = Options.allow_partial_stake, + mev_protection: bool = Options.mev_protection, period: int = Options.period, prompt: bool = Options.prompt, quiet: bool = Options.quiet, @@ -4006,6 +4013,9 @@ def stake_add( 7. Stake the same amount to multiple subnets: [green]$[/green] btcli stake add --amount 100 --netuids 4,5,6 + 8. Stake without MEV protection: + [green]$[/green] btcli stake add --amount 100 --netuid 1 --no-mev-protection + [bold]Safe Staking Parameters:[/bold] • [blue]--safe[/blue]: Enables rate tolerance checks • [blue]--tolerance[/blue]: Maximum % rate change allowed (0.05 = 5%) @@ -4201,6 +4211,7 @@ def stake_add( f"rate_tolerance: {rate_tolerance}\n" f"allow_partial_stake: {allow_partial_stake}\n" f"period: {period}\n" + f"mev_protection: {mev_protection}\n" ) return self._run_command( add_stake.stake_add( @@ -4218,6 +4229,7 @@ def stake_add( allow_partial_stake, json_output, period, + mev_protection, ) ) @@ -4270,6 +4282,7 @@ def stake_remove( safe_staking: Optional[bool] = Options.safe_staking, allow_partial_stake: Optional[bool] = Options.allow_partial_stake, period: int = Options.period, + mev_protection: bool = Options.mev_protection, prompt: bool = Options.prompt, interactive: bool = typer.Option( False, @@ -4306,6 +4319,9 @@ def stake_remove( 6. Unstake all Alpha from a hotkey and stake to Root: [green]$[/green] btcli stake remove --all-alpha + 7. Unstake without MEV protection: + [green]$[/green] btcli stake remove --amount 100 --netuid 1 --no-mev-protection + [bold]Safe Staking Parameters:[/bold] • [blue]--safe[/blue]: Enables rate tolerance checks during unstaking • [blue]--tolerance[/blue]: Max allowed rate change (0.05 = 5%) @@ -4478,7 +4494,8 @@ def stake_remove( f"all_hotkeys: {all_hotkeys}\n" f"include_hotkeys: {include_hotkeys}\n" f"exclude_hotkeys: {exclude_hotkeys}\n" - f"era: {period}" + f"era: {period}\n" + f"mev_protection: {mev_protection}" ) return self._run_command( remove_stake.unstake_all( @@ -4492,6 +4509,7 @@ def stake_remove( prompt=prompt, json_output=json_output, era=period, + mev_protection=mev_protection, ) ) elif ( @@ -4544,7 +4562,8 @@ def stake_remove( f"safe_staking: {safe_staking}\n" f"rate_tolerance: {rate_tolerance}\n" f"allow_partial_stake: {allow_partial_stake}\n" - f"era: {period}" + f"era: {period}\n" + f"mev_protection: {mev_protection}\n" ) return self._run_command( @@ -4564,6 +4583,7 @@ def stake_remove( allow_partial_stake=allow_partial_stake, json_output=json_output, era=period, + mev_protection=mev_protection, ) ) @@ -4609,6 +4629,7 @@ def stake_move( False, "--stake-all", "--all", help="Stake all", prompt=False ), period: int = Options.period, + mev_protection: bool = Options.mev_protection, prompt: bool = Options.prompt, quiet: bool = Options.quiet, verbose: bool = Options.verbose, @@ -4630,9 +4651,13 @@ def stake_move( If no arguments are provided, an interactive selection menu will be shown. - EXAMPLE + EXAMPLES - [green]$[/green] btcli stake move + 1. Interactive move (guided prompts): + [green]$[/green] btcli stake move + + 2. Move stake without MEV protection: + [green]$[/green] btcli stake move --no-mev-protection """ self.verbosity_handler(quiet, verbose, json_output) if prompt: @@ -4747,6 +4772,7 @@ def stake_move( f"era: {period}\n" f"interactive_selection: {interactive_selection}\n" f"prompt: {prompt}\n" + f"mev_protection: {mev_protection}\n" ) result, ext_id = self._run_command( move_stake.move_stake( @@ -4761,6 +4787,7 @@ def stake_move( era=period, interactive_selection=interactive_selection, prompt=prompt, + mev_protection=mev_protection, ) ) if json_output: @@ -4801,6 +4828,7 @@ def stake_transfer( stake_all: bool = typer.Option( False, "--stake-all", "--all", help="Stake all", prompt=False ), + mev_protection: bool = Options.mev_protection, period: int = Options.period, prompt: bool = Options.prompt, quiet: bool = Options.quiet, @@ -4838,6 +4866,9 @@ def stake_transfer( Transfer all available stake from origin hotkey: [green]$[/green] btcli stake transfer --all --origin-netuid 1 --dest-netuid 2 + + Transfer stake without MEV protection: + [green]$[/green] btcli stake transfer --origin-netuid 1 --dest-netuid 2 --amount 100 --no-mev-protection """ self.verbosity_handler(quiet, verbose, json_output) if prompt: @@ -4933,7 +4964,8 @@ def stake_transfer( f"dest_coldkey_ss58: {dest_ss58}\n" f"amount: {amount}\n" f"era: {period}\n" - f"stake_all: {stake_all}" + f"stake_all: {stake_all}\n" + f"mev_protection: {mev_protection}" ) result, ext_id = self._run_command( move_stake.transfer_stake( @@ -4948,6 +4980,7 @@ def stake_transfer( interactive_selection=interactive_selection, stake_all=stake_all, prompt=prompt, + mev_protection=mev_protection, ) ) if json_output: @@ -4992,6 +5025,7 @@ def stake_swap( prompt: bool = Options.prompt, wait_for_inclusion: bool = Options.wait_for_inclusion, wait_for_finalization: bool = Options.wait_for_finalization, + mev_protection: bool = Options.mev_protection, quiet: bool = Options.quiet, verbose: bool = Options.verbose, json_output: bool = Options.json_output, @@ -5011,10 +5045,13 @@ def stake_swap( If no arguments are provided, an interactive selection menu will be shown. - EXAMPLE + EXAMPLES - Swap 100 TAO from subnet 1 to subnet 2: - [green]$[/green] btcli stake swap --wallet-name default --wallet-hotkey default --origin-netuid 1 --dest-netuid 2 --amount 100 + 1. Swap 100 TAO from subnet 1 to subnet 2: + [green]$[/green] btcli stake swap --wallet-name default --wallet-hotkey default --origin-netuid 1 --dest-netuid 2 --amount 100 + + 2. Swap stake without MEV protection: + [green]$[/green] btcli stake swap --origin-netuid 1 --dest-netuid 2 --amount 100 --no-mev-protection """ self.verbosity_handler(quiet, verbose, json_output) console.print( @@ -5056,6 +5093,7 @@ def stake_swap( f"prompt: {prompt}\n" f"wait_for_inclusion: {wait_for_inclusion}\n" f"wait_for_finalization: {wait_for_finalization}\n" + f"mev_protection: {mev_protection}\n" ) result, ext_id = self._run_command( move_stake.swap_stake( @@ -5070,6 +5108,7 @@ def stake_swap( prompt=prompt, wait_for_inclusion=wait_for_inclusion, wait_for_finalization=wait_for_finalization, + mev_protection=mev_protection, ) ) if json_output: @@ -5085,6 +5124,7 @@ def stake_wizard( wallet_path: Optional[str] = Options.wallet_path, wallet_hotkey: Optional[str] = Options.wallet_hotkey, period: int = Options.period, + mev_protection: bool = Options.mev_protection, prompt: bool = Options.prompt, quiet: bool = Options.quiet, verbose: bool = Options.verbose, @@ -5150,6 +5190,7 @@ def stake_wizard( era=period, interactive_selection=False, prompt=prompt, + mev_protection=mev_protection, ) ) elif operation == "transfer": @@ -5179,6 +5220,7 @@ def stake_wizard( era=period, interactive_selection=False, prompt=prompt, + mev_protection=mev_protection, ) ) elif operation == "swap": @@ -5194,6 +5236,7 @@ def stake_wizard( era=period, interactive_selection=False, prompt=prompt, + mev_protection=mev_protection, ) ) else: @@ -6427,6 +6470,7 @@ def subnets_create( additional_info: Optional[str] = typer.Option( None, "--additional-info", help="Additional information" ), + mev_protection: bool = Options.mev_protection, json_output: bool = Options.json_output, prompt: bool = Options.prompt, quiet: bool = Options.quiet, @@ -6445,6 +6489,9 @@ def subnets_create( 2. Create with GitHub repo and contact email: [green]$[/green] btcli subnets create --subnet-name MySubnet --github-repo https://github.com/myorg/mysubnet --subnet-contact team@mysubnet.net + + 3. Create subnet without MEV protection: + [green]$[/green] btcli subnets create --no-mev-protection """ self.verbosity_handler(quiet, verbose, json_output) wallet = self.wallet_ask( @@ -6472,7 +6519,12 @@ def subnets_create( logger.debug(f"args:\nnetwork: {network}\nidentity: {identity}\n") self._run_command( subnets.create( - wallet, self.initialize_chain(network), identity, json_output, prompt + wallet, + self.initialize_chain(network), + identity, + json_output, + prompt, + mev_protection, ) ) diff --git a/bittensor_cli/src/bittensor/extrinsics/mev_shield.py b/bittensor_cli/src/bittensor/extrinsics/mev_shield.py new file mode 100644 index 000000000..eadce8863 --- /dev/null +++ b/bittensor_cli/src/bittensor/extrinsics/mev_shield.py @@ -0,0 +1,195 @@ +import asyncio +import hashlib +from typing import TYPE_CHECKING, Optional + +from async_substrate_interface import AsyncExtrinsicReceipt +from bittensor_drand import encrypt_mlkem768, mlkem_kdf_id +from bittensor_cli.src.bittensor.utils import encode_account_id, format_error_message + +if TYPE_CHECKING: + from bittensor_wallet import Wallet + from scalecodec import GenericCall + from bittensor_cli.src.bittensor.subtensor_interface import SubtensorInterface + + +async def encrypt_call( + subtensor: "SubtensorInterface", + wallet: "Wallet", + call: "GenericCall", +) -> "GenericCall": + """ + Encrypt a call using MEV Shield. + + Takes any call and returns a MevShield.submit_encrypted call + that can be submitted like any regular extrinsic. + + Args: + subtensor: The SubtensorInterface instance for chain queries. + wallet: The wallet whose coldkey will sign the inner payload. + call: The call to encrypt. + + Returns: + A MevShield.submit_encrypted call. + + Raises: + ValueError: If MEV Shield NextKey is not available on chain. + """ + + next_key_result, genesis_hash = await asyncio.gather( + subtensor.get_mev_shield_next_key(), + subtensor.substrate.get_block_hash(0), + ) + if next_key_result is None: + raise ValueError("MEV Shield NextKey not available on chain") + + ml_kem_768_public_key = next_key_result + + # Create payload_core: signer (32B) + next_key (32B) + SCALE(call) + signer_bytes = encode_account_id(wallet.coldkey.ss58_address) + scale_call_bytes = bytes(call.data.data) + next_key = hashlib.blake2b(next_key_result, digest_size=32).digest() + + payload_core = signer_bytes + next_key + scale_call_bytes + + mev_shield_version = mlkem_kdf_id() + genesis_hash_clean = ( + genesis_hash[2:] if genesis_hash.startswith("0x") else genesis_hash + ) + genesis_hash_bytes = bytes.fromhex(genesis_hash_clean) + + # Sign: coldkey.sign(b"mev-shield:v1" + genesis_hash + payload_core) + message_to_sign = ( + b"mev-shield:" + mev_shield_version + genesis_hash_bytes + payload_core + ) + signature = wallet.coldkey.sign(message_to_sign) + + # Plaintext: payload_core + b"\x01" + signature + plaintext = payload_core + b"\x01" + signature + + # Encrypt using ML-KEM-768 + ciphertext = encrypt_mlkem768(ml_kem_768_public_key, plaintext) + + # Commitment: blake2_256(payload_core) + commitment_hash = hashlib.blake2b(payload_core, digest_size=32).digest() + commitment_hex = "0x" + commitment_hash.hex() + + # Create the MevShield.submit_encrypted call + encrypted_call = await subtensor.substrate.compose_call( + call_module="MevShield", + call_function="submit_encrypted", + call_params={ + "commitment": commitment_hex, + "ciphertext": ciphertext, + }, + ) + + return encrypted_call + + +async def extract_mev_shield_id(response: "AsyncExtrinsicReceipt") -> Optional[str]: + """ + Extract the MEV Shield wrapper ID from an extrinsic response. + + After submitting a MEV Shield encrypted call, the EncryptedSubmitted event + contains the wrapper ID needed to track execution. + + Args: + response: The extrinsic receipt from submit_extrinsic. + + Returns: + The wrapper ID (hex string) or None if not found. + """ + for event in await response.triggered_events: + if event["event_id"] == "EncryptedSubmitted": + return event["attributes"]["id"] + return None + + +async def wait_for_mev_execution( + subtensor: "SubtensorInterface", + wrapper_id: str, + submit_block_hash: str, + timeout_blocks: int = 4, + status=None, +) -> tuple[bool, Optional[str], Optional[AsyncExtrinsicReceipt]]: + """ + Wait for MEV Shield inner call execution. + + After submit_encrypted succeeds, the block author will decrypt and execute + the inner call via execute_revealed. This function polls for the + DecryptedExecuted or DecryptedRejected event. + + Args: + subtensor: SubtensorInterface instance. + wrapper_id: The ID from EncryptedSubmitted event. + submit_block_number: Block number where submit_encrypted was included. + timeout_blocks: Max blocks to wait (default 4). + status: Optional rich.Status object for progress updates. + + Returns: + Tuple of (success: bool, error: Optional[str], receipt: Optional[AsyncExtrinsicReceipt]). + - (True, None, receipt) if DecryptedExecuted was found. + - (False, error_message, None) if the call failed or timeout. + """ + + async def _noop(_): + return True + + starting_block = await subtensor.substrate.get_block_number(submit_block_hash) + current_block = starting_block + 1 + + while current_block - starting_block <= timeout_blocks: + if status: + status.update( + f"Waiting for :shield: MEV Protection " + f"(checking block {current_block - starting_block} of {timeout_blocks})..." + ) + + await subtensor.substrate.wait_for_block( + current_block, + result_handler=_noop, + task_return=False, + ) + + block_hash = await subtensor.substrate.get_block_hash(current_block) + extrinsics = await subtensor.substrate.get_extrinsics(block_hash) + + # Find executeRevealed extrinsic & match ids + execute_revealed_index = None + for idx, extrinsic in enumerate(extrinsics): + call = extrinsic.value.get("call", {}) + call_module = call.get("call_module") + call_function = call.get("call_function") + + if call_module == "MevShield" and call_function in ( + "execute_revealed", + "mark_decryption_failed", + ): + call_args = call.get("call_args", []) + for arg in call_args: + if arg.get("name") == "id": + extrinsic_wrapper_id = arg.get("value") + if extrinsic_wrapper_id == wrapper_id: + execute_revealed_index = idx + break + + if execute_revealed_index is not None: + break + + if execute_revealed_index is None: + current_block += 1 + continue + + receipt = AsyncExtrinsicReceipt( + substrate=subtensor.substrate, + block_hash=block_hash, + extrinsic_idx=execute_revealed_index, + ) + + if not await receipt.is_success: + error_msg = format_error_message(await receipt.error_message) + return False, error_msg, None + + return True, None, receipt + + return False, "Timeout waiting for MEV Shield execution", None diff --git a/bittensor_cli/src/bittensor/subtensor_interface.py b/bittensor_cli/src/bittensor/subtensor_interface.py index 054d67f7a..0566f14e8 100644 --- a/bittensor_cli/src/bittensor/subtensor_interface.py +++ b/bittensor_cli/src/bittensor/subtensor_interface.py @@ -41,7 +41,7 @@ decode_hex_identity_dict, validate_chain_endpoint, u16_normalized_float, - U16_MAX, + MEV_SHIELD_PUBLIC_KEY_SIZE, get_hotkey_pub_ss58, ) @@ -2175,8 +2175,8 @@ async def get_claimable_stakes_for_coldkey( root_stake: Balance claimable_stake: Balance for hotkey, netuid in target_pairs: - root_stake = root_stakes[hotkey] - rate = claimable_rates[hotkey].get(netuid, 0.0) + root_stake = root_stakes.get(hotkey, Balance(0)) + rate = claimable_rates.get(hotkey, {}).get(netuid, 0.0) claimable_stake = rate * root_stake already_claimed = claimed_amounts.get((hotkey, netuid), Balance(0)) net_claimable = max(claimable_stake - already_claimed, Balance(0)) @@ -2299,6 +2299,62 @@ async def get_subnet_ema_tao_inflow( ema_value = fixed_to_float(raw_ema_value) return Balance.from_rao(ema_value) + async def get_mev_shield_next_key( + self, + block_hash: Optional[str] = None, + ) -> Optional[tuple[bytes, int]]: + """ + Get the next MEV Shield public key and epoch from chain storage. + + Args: + block_hash: Optional block hash to query at. + + Returns: + Tuple of (public_key_bytes, epoch) or None if not available. + """ + result = await self.query( + module="MevShield", + storage_function="NextKey", + block_hash=block_hash, + ) + public_key_bytes = bytes(next(iter(result))) + + if len(public_key_bytes) != MEV_SHIELD_PUBLIC_KEY_SIZE: + raise ValueError( + f"Invalid ML-KEM-768 public key size: {len(public_key_bytes)} bytes. " + f"Expected exactly {MEV_SHIELD_PUBLIC_KEY_SIZE} bytes." + ) + + return public_key_bytes + + async def get_mev_shield_current_key( + self, + block_hash: Optional[str] = None, + ) -> Optional[tuple[bytes, int]]: + """ + Get the current MEV Shield public key and epoch from chain storage. + + Args: + block_hash: Optional block hash to query at. + + Returns: + Tuple of (public_key_bytes, epoch) or None if not available. + """ + result = await self.query( + module="MevShield", + storage_function="CurrentKey", + block_hash=block_hash, + ) + public_key_bytes = bytes(next(iter(result))) + + if len(public_key_bytes) != MEV_SHIELD_PUBLIC_KEY_SIZE: + raise ValueError( + f"Invalid ML-KEM-768 public key size: {len(public_key_bytes)} bytes. " + f"Expected exactly {MEV_SHIELD_PUBLIC_KEY_SIZE} bytes." + ) + + return public_key_bytes + async def best_connection(networks: list[str]): """ diff --git a/bittensor_cli/src/bittensor/utils.py b/bittensor_cli/src/bittensor/utils.py index f8e322f01..1ca8325ba 100644 --- a/bittensor_cli/src/bittensor/utils.py +++ b/bittensor_cli/src/bittensor/utils.py @@ -37,6 +37,7 @@ BT_DOCS_LINK = "https://docs.learnbittensor.org" GLOBAL_MAX_SUBNET_COUNT = 4096 +MEV_SHIELD_PUBLIC_KEY_SIZE = 1184 console = Console() json_console = Console() diff --git a/bittensor_cli/src/commands/stake/add.py b/bittensor_cli/src/commands/stake/add.py index 5046981da..d88692261 100644 --- a/bittensor_cli/src/commands/stake/add.py +++ b/bittensor_cli/src/commands/stake/add.py @@ -11,6 +11,11 @@ from async_substrate_interface.errors import SubstrateRequestException from bittensor_cli.src import COLOR_PALETTE from bittensor_cli.src.bittensor.balances import Balance +from bittensor_cli.src.bittensor.extrinsics.mev_shield import ( + encrypt_call, + extract_mev_shield_id, + wait_for_mev_execution, +) from bittensor_cli.src.bittensor.utils import ( console, err_console, @@ -46,6 +51,7 @@ async def stake_add( allow_partial_stake: bool, json_output: bool, era: int, + mev_protection: bool, ): """ Args: @@ -134,6 +140,8 @@ async def safe_stake_extrinsic( }, ), ) + if mev_protection: + call = await encrypt_call(subtensor, wallet, call) extrinsic = await subtensor.substrate.create_signed_extrinsic( call=call, keypair=wallet.coldkey, @@ -160,49 +168,61 @@ async def safe_stake_extrinsic( err_msg = f"{failure_prelude} with error: {format_error_message(await response.error_message)}" err_out("\n" + err_msg) return False, err_msg, None - else: - if json_output: - # the rest of this checking is not necessary if using json_output - return True, "", response - await print_extrinsic_id(response) - block_hash = await subtensor.substrate.get_chain_head() - new_balance, new_stake = await asyncio.gather( - subtensor.get_balance(wallet.coldkeypub.ss58_address, block_hash), - subtensor.get_stake( - hotkey_ss58=hotkey_ss58_, - coldkey_ss58=wallet.coldkeypub.ss58_address, - netuid=netuid_, - block_hash=block_hash, - ), - ) - console.print( - f":white_heavy_check_mark: [dark_sea_green3]Finalized. " - f"Stake added to netuid: {netuid_}[/dark_sea_green3]" - ) - console.print( - f"Balance:\n [blue]{current_balance}[/blue] :arrow_right: " - f"[{COLOR_PALETTE['STAKE']['STAKE_AMOUNT']}]{new_balance}" - ) - amount_staked = current_balance - new_balance - if allow_partial_stake and (amount_staked != amount_): - console.print( - "Partial stake transaction. Staked:\n" - f" [{COLOR_PALETTE['STAKE']['STAKE_AMOUNT']}]{amount_staked}" - f"[/{COLOR_PALETTE['STAKE']['STAKE_AMOUNT']}] " - f"instead of " - f"[blue]{amount_}[/blue]" + if mev_protection: + 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 + + if json_output: + # the rest of this checking is not necessary if using json_output + return True, "", response + await print_extrinsic_id(response) + block_hash = await subtensor.substrate.get_chain_head() + new_balance, new_stake = await asyncio.gather( + subtensor.get_balance(wallet.coldkeypub.ss58_address, block_hash), + subtensor.get_stake( + hotkey_ss58=hotkey_ss58_, + coldkey_ss58=wallet.coldkeypub.ss58_address, + netuid=netuid_, + block_hash=block_hash, + ), + ) + console.print( + f":white_heavy_check_mark: [dark_sea_green3]Finalized. " + f"Stake added to netuid: {netuid_}[/dark_sea_green3]" + ) + console.print( + f"Balance:\n [blue]{current_balance}[/blue] :arrow_right: " + f"[{COLOR_PALETTE['STAKE']['STAKE_AMOUNT']}]{new_balance}" + ) + amount_staked = current_balance - new_balance + if allow_partial_stake and (amount_staked != amount_): console.print( - f"Subnet: [{COLOR_PALETTE['GENERAL']['SUBHEADING']}]" - f"{netuid_}[/{COLOR_PALETTE['GENERAL']['SUBHEADING']}] " - f"Stake:\n" - f" [blue]{current_stake}[/blue] " - f":arrow_right: " - f"[{COLOR_PALETTE['STAKE']['STAKE_AMOUNT']}]{new_stake}\n" + "Partial stake transaction. Staked:\n" + f" [{COLOR_PALETTE['STAKE']['STAKE_AMOUNT']}]{amount_staked}" + f"[/{COLOR_PALETTE['STAKE']['STAKE_AMOUNT']}] " + f"instead of " + f"[blue]{amount_}[/blue]" ) - return True, "", response + + console.print( + f"Subnet: [{COLOR_PALETTE['GENERAL']['SUBHEADING']}]" + f"{netuid_}[/{COLOR_PALETTE['GENERAL']['SUBHEADING']}] " + f"Stake:\n" + f" [blue]{current_stake}[/blue] " + f":arrow_right: " + f"[{COLOR_PALETTE['STAKE']['STAKE_AMOUNT']}]{new_stake}\n" + ) + return True, "", response async def stake_extrinsic( netuid_i, amount_, current, staking_address_ss58, status=None @@ -224,6 +244,8 @@ 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) extrinsic = await subtensor.substrate.create_signed_extrinsic( call=call, keypair=wallet.coldkey, nonce=next_nonce, era={"period": era} ) @@ -235,45 +257,56 @@ async def stake_extrinsic( err_msg = f"{failure_prelude} with error: {format_error_message(e)}" err_out("\n" + err_msg) return False, err_msg, None - else: - if not await response.is_success: - err_msg = f"{failure_prelude} with error: {format_error_message(await response.error_message)}" - err_out("\n" + err_msg) - return False, err_msg, None - else: - if json_output: - # the rest of this is not necessary if using json_output - return True, "", response - await print_extrinsic_id(response) - new_block_hash = await subtensor.substrate.get_chain_head() - new_balance, new_stake = await asyncio.gather( - subtensor.get_balance( - wallet.coldkeypub.ss58_address, block_hash=new_block_hash - ), - subtensor.get_stake( - hotkey_ss58=staking_address_ss58, - coldkey_ss58=wallet.coldkeypub.ss58_address, - netuid=netuid_i, - block_hash=new_block_hash, - ), - ) - console.print( - f":white_heavy_check_mark: " - f"[dark_sea_green3]Finalized. Stake added to netuid: {netuid_i}[/dark_sea_green3]" - ) - console.print( - f"Balance:\n [blue]{current_balance}[/blue] :arrow_right: " - f"[{COLOR_PALETTE['STAKE']['STAKE_AMOUNT']}]{new_balance}" - ) - console.print( - f"Subnet: [{COLOR_PALETTE['GENERAL']['SUBHEADING']}]" - f"{netuid_i}[/{COLOR_PALETTE['GENERAL']['SUBHEADING']}] " - f"Stake:\n" - f" [blue]{current}[/blue] " - f":arrow_right: " - f"[{COLOR_PALETTE['STAKE']['STAKE_AMOUNT']}]{new_stake}\n" + if not await response.is_success: + err_msg = f"{failure_prelude} with error: {format_error_message(await response.error_message)}" + err_out("\n" + err_msg) + return False, err_msg, None + + if mev_protection: + 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 ) - return True, "", response + 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 + await print_extrinsic_id(response) + new_block_hash = await subtensor.substrate.get_chain_head() + new_balance, new_stake = await asyncio.gather( + subtensor.get_balance( + wallet.coldkeypub.ss58_address, block_hash=new_block_hash + ), + subtensor.get_stake( + hotkey_ss58=staking_address_ss58, + coldkey_ss58=wallet.coldkeypub.ss58_address, + netuid=netuid_i, + block_hash=new_block_hash, + ), + ) + console.print( + f":white_heavy_check_mark: " + f"[dark_sea_green3]Finalized. Stake added to netuid: {netuid_i}[/dark_sea_green3]" + ) + console.print( + f"Balance:\n [blue]{current_balance}[/blue] :arrow_right: " + f"[{COLOR_PALETTE['STAKE']['STAKE_AMOUNT']}]{new_balance}" + ) + console.print( + f"Subnet: [{COLOR_PALETTE['GENERAL']['SUBHEADING']}]" + f"{netuid_i}[/{COLOR_PALETTE['GENERAL']['SUBHEADING']}] " + f"Stake:\n" + f" [blue]{current}[/blue] " + f":arrow_right: " + f"[{COLOR_PALETTE['STAKE']['STAKE_AMOUNT']}]{new_stake}\n" + ) + return True, "", response netuids = ( netuids if netuids is not None else await subtensor.get_all_subnet_netuids() @@ -435,47 +468,53 @@ async def stake_extrinsic( if not unlock_key(wallet).success: return - if safe_staking: - stake_coroutines = {} - for i, (ni, am, curr, price_with_tolerance) in enumerate( - zip( - netuids, amounts_to_stake, current_stake_balances, prices_with_tolerance - ) - ): - for _, staking_address in hotkeys_to_stake_to: - # Regular extrinsic for root subnet - if ni == 0: - stake_coroutines[(ni, staking_address)] = stake_extrinsic( - netuid_i=ni, - amount_=am, - current=curr, - staking_address_ss58=staking_address, - ) - else: - stake_coroutines[(ni, staking_address)] = safe_stake_extrinsic( - netuid_=ni, - amount_=am, - current_stake=curr, - hotkey_ss58_=staking_address, - price_limit=price_with_tolerance, - ) - else: - stake_coroutines = { - (ni, staking_address): stake_extrinsic( - netuid_i=ni, - amount_=am, - current=curr, - staking_address_ss58=staking_address, - ) - for i, (ni, am, curr) in enumerate( - zip(netuids, amounts_to_stake, current_stake_balances) - ) - for _, staking_address in hotkeys_to_stake_to - } successes = defaultdict(dict) error_messages = defaultdict(dict) extrinsic_ids = defaultdict(dict) - with console.status(f"\n:satellite: Staking on netuid(s): {netuids} ..."): + with console.status(f"\n:satellite: Staking on netuid(s): {netuids} ...") as status: + if safe_staking: + stake_coroutines = {} + for i, (ni, am, curr, price_with_tolerance) in enumerate( + zip( + netuids, + amounts_to_stake, + current_stake_balances, + prices_with_tolerance, + ) + ): + for _, staking_address in hotkeys_to_stake_to: + # Regular extrinsic for root subnet + if ni == 0: + stake_coroutines[(ni, staking_address)] = stake_extrinsic( + netuid_i=ni, + amount_=am, + current=curr, + staking_address_ss58=staking_address, + status=status, + ) + else: + stake_coroutines[(ni, staking_address)] = safe_stake_extrinsic( + netuid_=ni, + amount_=am, + current_stake=curr, + hotkey_ss58_=staking_address, + price_limit=price_with_tolerance, + status=status, + ) + else: + stake_coroutines = { + (ni, staking_address): stake_extrinsic( + netuid_i=ni, + amount_=am, + current=curr, + staking_address_ss58=staking_address, + status=status, + ) + for i, (ni, am, curr) in enumerate( + zip(netuids, amounts_to_stake, current_stake_balances) + ) + for _, staking_address in hotkeys_to_stake_to + } # We can gather them all at once but balance reporting will be in race-condition. for (ni, staking_address), coroutine in stake_coroutines.items(): success, er_msg, ext_receipt = await coroutine diff --git a/bittensor_cli/src/commands/stake/move.py b/bittensor_cli/src/commands/stake/move.py index 99a0b79ac..a7f3f0782 100644 --- a/bittensor_cli/src/commands/stake/move.py +++ b/bittensor_cli/src/commands/stake/move.py @@ -8,6 +8,11 @@ from bittensor_cli.src import COLOR_PALETTE from bittensor_cli.src.bittensor.balances import Balance +from bittensor_cli.src.bittensor.extrinsics.mev_shield import ( + encrypt_call, + extract_mev_shield_id, + wait_for_mev_execution, +) from bittensor_cli.src.bittensor.utils import ( console, err_console, @@ -458,6 +463,7 @@ async def move_stake( era: int, interactive_selection: bool = False, prompt: bool = True, + mev_protection: bool = True, ) -> tuple[bool, str]: if interactive_selection: try: @@ -578,13 +584,27 @@ async def move_stake( f"\n:satellite: Moving [blue]{amount_to_move_as_balance}[/blue] from [blue]{origin_hotkey}[/blue] on netuid: " f"[blue]{origin_netuid}[/blue] \nto " f"[blue]{destination_hotkey}[/blue] on netuid: [blue]{destination_netuid}[/blue] ..." - ): + ) as status: + if mev_protection: + call = await encrypt_call(subtensor, wallet, call) extrinsic = await subtensor.substrate.create_signed_extrinsic( call=call, keypair=wallet.coldkey, era={"period": era} ) response = await subtensor.substrate.submit_extrinsic( extrinsic, wait_for_inclusion=True, wait_for_finalization=False ) + + if mev_protection: + 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_console.print(f"\n:cross_mark: [red]Failed[/red]: {mev_error}") + return False, "" + ext_id = await response.get_extrinsic_identifier() if not prompt: @@ -644,6 +664,7 @@ async def transfer_stake( interactive_selection: bool = False, stake_all: bool = False, prompt: bool = True, + mev_protection: bool = True, ) -> tuple[bool, str]: """Transfers stake from one network to another. @@ -765,7 +786,9 @@ async def transfer_stake( if not unlock_key(wallet).success: return False, "" - with console.status("\n:satellite: Transferring stake ..."): + with console.status("\n:satellite: Transferring stake ...") as status: + if mev_protection: + call = await encrypt_call(subtensor, wallet, call) extrinsic = await subtensor.substrate.create_signed_extrinsic( call=call, keypair=wallet.coldkey, era={"period": era} ) @@ -773,6 +796,18 @@ async def transfer_stake( response = await subtensor.substrate.submit_extrinsic( extrinsic, wait_for_inclusion=True, wait_for_finalization=False ) + + if mev_protection: + 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_console.print(f"\n:cross_mark: [red]Failed[/red]: {mev_error}") + return False, "" + ext_id = await response.get_extrinsic_identifier() if not prompt: @@ -823,6 +858,7 @@ async def swap_stake( prompt: bool = True, wait_for_inclusion: bool = True, wait_for_finalization: bool = False, + mev_protection: bool = True, ) -> tuple[bool, str]: """Swaps stake between subnets while keeping the same coldkey-hotkey pair ownership. @@ -940,7 +976,9 @@ async def swap_stake( with console.status( f"\n:satellite: Swapping stake from netuid [blue]{origin_netuid}[/blue] " f"to netuid [blue]{destination_netuid}[/blue]..." - ): + ) as status: + if mev_protection: + call = await encrypt_call(subtensor, wallet, call) extrinsic = await subtensor.substrate.create_signed_extrinsic( call=call, keypair=wallet.coldkey, era={"period": era} ) @@ -950,6 +988,18 @@ async def swap_stake( wait_for_inclusion=wait_for_inclusion, wait_for_finalization=wait_for_finalization, ) + + if mev_protection: + 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_console.print(f"\n:cross_mark: [red]Failed[/red]: {mev_error}") + return False, "" + ext_id = await response.get_extrinsic_identifier() if not prompt: diff --git a/bittensor_cli/src/commands/stake/remove.py b/bittensor_cli/src/commands/stake/remove.py index 5d125cc16..afd9310da 100644 --- a/bittensor_cli/src/commands/stake/remove.py +++ b/bittensor_cli/src/commands/stake/remove.py @@ -11,6 +11,11 @@ from async_substrate_interface.errors import SubstrateRequestException from bittensor_cli.src import COLOR_PALETTE +from bittensor_cli.src.bittensor.extrinsics.mev_shield import ( + encrypt_call, + extract_mev_shield_id, + wait_for_mev_execution, +) from bittensor_cli.src.bittensor.balances import Balance from bittensor_cli.src.bittensor.utils import ( console, @@ -48,6 +53,7 @@ async def unstake( allow_partial_stake: bool, json_output: bool, era: int, + mev_protection: bool, ): """Unstake from hotkey(s).""" @@ -326,6 +332,7 @@ async def unstake( "hotkey_ss58": op["hotkey_ss58"], "status": status, "era": era, + "mev_protection": mev_protection, } if safe_staking and op["netuid"] != 0: @@ -370,6 +377,7 @@ async def unstake_all( era: int = 3, prompt: bool = True, json_output: bool = False, + mev_protection: bool = True, ) -> None: """Unstakes all stakes from all hotkeys in all subnets.""" include_hotkeys = include_hotkeys or [] @@ -546,6 +554,7 @@ async def unstake_all( unstake_all_alpha=unstake_all_alpha, status=status, era=era, + mev_protection=mev_protection, ) ext_id = await ext_receipt.get_extrinsic_identifier() if success else None successes[hotkey_ss58] = { @@ -566,6 +575,7 @@ async def _unstake_extrinsic( hotkey_ss58: str, status=None, era: int = 3, + mev_protection: bool = True, ) -> tuple[bool, Optional[AsyncExtrinsicReceipt]]: """Execute a standard unstake extrinsic. @@ -601,6 +611,9 @@ async def _unstake_extrinsic( }, ), ) + + if mev_protection: + call = await encrypt_call(subtensor, wallet, call) extrinsic = await subtensor.substrate.create_signed_extrinsic( call=call, keypair=wallet.coldkey, era={"period": era} ) @@ -615,6 +628,19 @@ async def _unstake_extrinsic( f"{format_error_message(await response.error_message)}" ) return False, None + + if mev_protection: + 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, None + # Fetch latest balance and stake await print_extrinsic_id(response) block_hash = await subtensor.substrate.get_chain_head() @@ -653,6 +679,7 @@ async def _safe_unstake_extrinsic( allow_partial_stake: bool, status=None, era: int = 3, + mev_protection: bool = True, ) -> tuple[bool, Optional[AsyncExtrinsicReceipt]]: """Execute a safe unstake extrinsic with price limit. @@ -678,9 +705,8 @@ async def _safe_unstake_extrinsic( block_hash = await subtensor.substrate.get_chain_head() - current_balance, next_nonce, current_stake, call = await asyncio.gather( + current_balance, current_stake, call = await asyncio.gather( subtensor.get_balance(wallet.coldkeypub.ss58_address, block_hash), - subtensor.substrate.get_account_next_index(wallet.coldkeypub.ss58_address), subtensor.get_stake( hotkey_ss58=hotkey_ss58, coldkey_ss58=wallet.coldkeypub.ss58_address, @@ -701,8 +727,10 @@ async def _safe_unstake_extrinsic( ), ) + if mev_protection: + call = await encrypt_call(subtensor, wallet, call) extrinsic = await subtensor.substrate.create_signed_extrinsic( - call=call, keypair=wallet.coldkey, nonce=next_nonce, era={"period": era} + call=call, keypair=wallet.coldkey, era={"period": era} ) try: @@ -726,6 +754,19 @@ async def _safe_unstake_extrinsic( f"\n{failure_prelude} with error: {format_error_message(await response.error_message)}" ) return False, None + + if mev_protection: + 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, None + await print_extrinsic_id(response) block_hash = await subtensor.substrate.get_chain_head() new_balance, new_stake = await asyncio.gather( @@ -767,6 +808,7 @@ async def _unstake_all_extrinsic( unstake_all_alpha: bool, status=None, era: int = 3, + mev_protection: bool = True, ) -> tuple[bool, Optional[AsyncExtrinsicReceipt]]: """Execute an unstake all extrinsic. @@ -814,6 +856,9 @@ async def _unstake_all_extrinsic( call_params={"hotkey": hotkey_ss58}, ) + if mev_protection: + call = await encrypt_call(subtensor, wallet, call) + try: response = await subtensor.substrate.submit_extrinsic( extrinsic=await subtensor.substrate.create_signed_extrinsic( @@ -829,8 +874,20 @@ async def _unstake_all_extrinsic( f"{format_error_message(await response.error_message)}" ) return False, None - else: - await print_extrinsic_id(response) + + if mev_protection: + 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, None + + await print_extrinsic_id(response) # Fetch latest balance and stake block_hash = await subtensor.substrate.get_chain_head() diff --git a/bittensor_cli/src/commands/subnets/subnets.py b/bittensor_cli/src/commands/subnets/subnets.py index c2346880b..e2cc7bbd7 100644 --- a/bittensor_cli/src/commands/subnets/subnets.py +++ b/bittensor_cli/src/commands/subnets/subnets.py @@ -10,13 +10,18 @@ from rich.table import Column, Table from rich import box -from bittensor_cli.src import COLOR_PALETTE, Constants +from bittensor_cli.src import COLOR_PALETTE from bittensor_cli.src.bittensor.balances import Balance from bittensor_cli.src.bittensor.extrinsics.registration import ( register_extrinsic, burned_register_extrinsic, ) from bittensor_cli.src.bittensor.extrinsics.root import root_register_extrinsic +from bittensor_cli.src.bittensor.extrinsics.mev_shield import ( + encrypt_call, + extract_mev_shield_id, + wait_for_mev_execution, +) from rich.live import Live from bittensor_cli.src.bittensor.minigraph import MiniGraph from bittensor_cli.src.commands.wallets import set_id, get_id @@ -56,6 +61,7 @@ async def register_subnetwork_extrinsic( wait_for_inclusion: bool = False, wait_for_finalization: bool = True, prompt: bool = False, + mev_protection: bool = True, ) -> tuple[bool, Optional[int], Optional[str]]: """Registers a new subnetwork. @@ -166,7 +172,7 @@ async def _find_event_attributes_in_extrinsic_receipt( if not unlock_key(wallet).success: return False, None, None - with console.status(":satellite: Registering subnet...", spinner="earth"): + with console.status(":satellite: Registering subnet...", spinner="earth") as status: substrate = subtensor.substrate # create extrinsic call call = await substrate.compose_call( @@ -174,6 +180,8 @@ async def _find_event_attributes_in_extrinsic_receipt( call_function=call_function, call_params=call_params, ) + if mev_protection: + call = await encrypt_call(subtensor, wallet, call) extrinsic = await substrate.create_signed_extrinsic( call=call, keypair=wallet.coldkey ) @@ -194,17 +202,30 @@ async def _find_event_attributes_in_extrinsic_receipt( await asyncio.sleep(0.5) return False, None, None + # Check for MEV shield execution + if mev_protection: + 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_console.print( + f":cross_mark: [red]Failed[/red]: MEV execution failed: {mev_error}" + ) + return False, None, None + # Successful registration, final check for membership - else: - attributes = await _find_event_attributes_in_extrinsic_receipt( - response, "NetworkAdded" - ) - await print_extrinsic_id(response) - ext_id = await response.get_extrinsic_identifier() - console.print( - f":white_heavy_check_mark: [dark_sea_green3]Registered subnetwork with netuid: {attributes[0]}" - ) - return True, int(attributes[0]), ext_id + attributes = await _find_event_attributes_in_extrinsic_receipt( + response, "NetworkAdded" + ) + await print_extrinsic_id(response) + ext_id = await response.get_extrinsic_identifier() + console.print( + f":white_heavy_check_mark: [dark_sea_green3]Registered subnetwork with netuid: {attributes[0]}" + ) + return True, int(attributes[0]), ext_id # commands @@ -1622,12 +1643,13 @@ async def create( subnet_identity: dict, json_output: bool, prompt: bool, + mev_protection: bool = True, ): """Register a subnetwork""" # Call register command. success, netuid, ext_id = await register_subnetwork_extrinsic( - subtensor, wallet, subnet_identity, prompt=prompt + subtensor, wallet, subnet_identity, prompt=prompt, mev_protection=mev_protection ) if json_output: # technically, netuid can be `None`, but only if not wait for finalization/inclusion. However, as of present diff --git a/pyproject.toml b/pyproject.toml index 9eefb8d4d..fbe506980 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "flit_core.buildapi" [project] name = "bittensor-cli" -version = "9.15.3" +version = "9.16.0rc1" description = "Bittensor CLI" readme = "README.md" authors = [ @@ -30,9 +30,10 @@ classifiers = [ ] dependencies = [ "wheel", - "async-substrate-interface>=1.5.2", + "async-substrate-interface>=1.5.14", "aiohttp~=3.13", "backoff~=2.2.1", + "bittensor-drand>=1.2.0", "GitPython>=3.0.0", "netaddr~=1.3.0", "numpy>=2.0.1,<3.0.0", diff --git a/tests/e2e_tests/test_hyperparams_setting.py b/tests/e2e_tests/test_hyperparams_setting.py index c336f6615..c21008bb5 100644 --- a/tests/e2e_tests/test_hyperparams_setting.py +++ b/tests/e2e_tests/test_hyperparams_setting.py @@ -1,6 +1,6 @@ import asyncio import json - +import pytest from bittensor_cli.src import HYPERPARAMS, RootSudoOnly from .utils import turn_off_hyperparam_freeze_window @@ -13,6 +13,7 @@ """ +@pytest.mark.parametrize("local_chain", [False], indirect=True) def test_hyperparams_setting(local_chain, wallet_setup): netuid = 2 wallet_path_alice = "//Alice" diff --git a/tests/e2e_tests/test_liquidity.py b/tests/e2e_tests/test_liquidity.py index 7a210f0a1..e97e1b6b4 100644 --- a/tests/e2e_tests/test_liquidity.py +++ b/tests/e2e_tests/test_liquidity.py @@ -1,9 +1,7 @@ import asyncio import json -import re import time -from bittensor_cli.src.bittensor.balances import Balance from .utils import turn_off_hyperparam_freeze_window """ @@ -30,7 +28,6 @@ def test_liquidity(local_chain, wallet_setup): print( "Skipping turning off hyperparams freeze window. This indicates the call does not exist on the chain you are testing." ) - time.sleep(10) # Register a subnet with sudo as Alice result = exec_command_alice( @@ -63,6 +60,7 @@ def test_liquidity(local_chain, wallet_setup): "https://testsubnet.com/logo.png", "--no-prompt", "--json-output", + "--no-mev-protection", ], ) result_output = json.loads(result.stdout) @@ -92,6 +90,7 @@ def test_liquidity(local_chain, wallet_setup): assert result_output["success"] is False assert f"Subnet with netuid: {netuid} is not active" in result_output["err_msg"] assert result_output["positions"] == [] + time.sleep(40) # start emissions schedule start_subnet_emissions = exec_command_alice( @@ -137,6 +136,7 @@ def test_liquidity(local_chain, wallet_setup): "--no-prompt", "--era", "144", + "--no-mev-protection", ], ) assert "✅ Finalized" in stake_to_enable_v3.stdout, stake_to_enable_v3.stderr diff --git a/tests/e2e_tests/test_set_identity.py b/tests/e2e_tests/test_set_identity.py index 9c008cdcd..345927785 100644 --- a/tests/e2e_tests/test_set_identity.py +++ b/tests/e2e_tests/test_set_identity.py @@ -23,6 +23,7 @@ def test_set_id(local_chain, wallet_setup): keypair_alice, wallet_alice, wallet_path_alice, exec_command_alice = wallet_setup( wallet_path_alice ) + # Register a subnet with sudo as Alice result = exec_command_alice( command="subnets", @@ -54,6 +55,7 @@ def test_set_id(local_chain, wallet_setup): "https://testsubnet.com/logo.png", "--no-prompt", "--json-output", + "--no-mev-protection", ], ) result_output = json.loads(result.stdout) diff --git a/tests/e2e_tests/test_stake_movement.py b/tests/e2e_tests/test_stake_movement.py new file mode 100644 index 000000000..eb9a22b0e --- /dev/null +++ b/tests/e2e_tests/test_stake_movement.py @@ -0,0 +1,511 @@ +import asyncio +import json +import pytest + +from .utils import find_stake_entries, set_storage_extrinsic + + +@pytest.mark.parametrize("local_chain", [False], indirect=True) +def test_stake_movement(local_chain, wallet_setup): + """ + Exercise stake move, transfer, and swap flows across subnets using Alice and Bob. + + Steps: + 0. Initial setup: Make alice own SN 0, create SN2, SN3, SN4, start emissions on all subnets. + 1. Activation: Register Bob on subnets 2 and 3; add initial stake for V3 activation. + 2. Move: Move stake from Alice's hotkey on netuid 2 to Bob's hotkey on netuid 3. + 3. Transfer: Transfer all root (netuid 0) stake from Alice's coldkey to Bob's coldkey. + 4. Swap: Swap Alice's stake from netuid 4 to the root netuid. + + Note: + - All movement commands executed with mev shield + - Stake commands executed without shield to speed up tests + - Shield for stake commands is already covered in its own test + """ + print("Testing stake movement commands 🧪") + + wallet_path_alice = "//Alice" + wallet_path_bob = "//Bob" + + keypair_alice, wallet_alice, wallet_path_alice, exec_command_alice = wallet_setup( + wallet_path_alice + ) + keypair_bob, wallet_bob, wallet_path_bob, exec_command_bob = wallet_setup( + wallet_path_bob + ) + + # Force Alice to own SN0 by setting storage + sn0_owner_storage_items = [ + ( + bytes.fromhex( + "658faa385070e074c85bf6b568cf055536e3e82152c8758267395fe524fbbd160000" + ), + bytes.fromhex( + "d43593c715fdd31c61141abd04a99fd6822c8558854ccde39a5684e7a56da27d" + ), + ) + ] + asyncio.run( + set_storage_extrinsic( + local_chain, + wallet=wallet_alice, + items=sn0_owner_storage_items, + ) + ) + + # Create SN2, SN3, SN4 for move/transfer/swap checks + subnets_to_create = [2, 3, 4] + for netuid in subnets_to_create: + create_subnet_result = exec_command_alice( + command="subnets", + sub_command="create", + extra_args=[ + "--wallet-path", + wallet_path_alice, + "--chain", + "ws://127.0.0.1:9945", + "--wallet-name", + wallet_alice.name, + "--wallet-hotkey", + wallet_alice.hotkey_str, + "--subnet-name", + "Test Subnet", + "--repo", + "https://github.com/username/repo", + "--contact", + "alice@opentensor.dev", + "--url", + "https://testsubnet.com", + "--discord", + "alice#1234", + "--description", + "A test subnet for e2e testing", + "--additional-info", + "Created by Alice", + "--logo-url", + "https://testsubnet.com/logo.png", + "--no-prompt", + "--json-output", + ], + ) + create_subnet_payload = json.loads(create_subnet_result.stdout) + assert create_subnet_payload["success"] is True + assert create_subnet_payload["netuid"] == netuid + + # Start emission schedule for subnets (including root netuid 0) + for netuid in [0] + subnets_to_create: + start_emission_result = exec_command_alice( + command="subnets", + sub_command="start", + extra_args=[ + "--netuid", + str(netuid), + "--wallet-name", + wallet_alice.name, + "--no-prompt", + "--chain", + "ws://127.0.0.1:9945", + "--wallet-path", + wallet_path_alice, + ], + ) + assert ( + f"Successfully started subnet {netuid}'s emission schedule." + in start_emission_result.stdout + ) + + # Alice is already registered - register Bob on the two non-root subnets + for netuid in [2, 3]: + register_bob_result = exec_command_bob( + command="subnets", + sub_command="register", + extra_args=[ + "--netuid", + str(netuid), + "--wallet-path", + wallet_path_bob, + "--wallet-name", + wallet_bob.name, + "--hotkey", + wallet_bob.hotkey_str, + "--chain", + "ws://127.0.0.1:9945", + "--no-prompt", + ], + ) + assert "✅ Registered" in register_bob_result.stdout, register_bob_result.stderr + assert "Your extrinsic has been included" in register_bob_result.stdout, ( + register_bob_result.stdout + ) + + # Add initial stake to enable V3 (1 TAO) on all created subnets + for netuid in [2, 3, 4]: + add_initial_stake_result = exec_command_alice( + command="stake", + sub_command="add", + extra_args=[ + "--netuid", + str(netuid), + "--wallet-path", + wallet_path_alice, + "--wallet-name", + wallet_alice.name, + "--hotkey", + wallet_alice.hotkey_str, + "--chain", + "ws://127.0.0.1:9945", + "--amount", + "1", + "--unsafe", + "--no-prompt", + "--era", + "144", + "--no-mev-protection", + ], + ) + assert "✅ Finalized" in add_initial_stake_result.stdout, ( + add_initial_stake_result.stderr + ) + + ############################ + # TEST 1: Move stake command + # Move stake between hotkeys while keeping the same coldkey + ############################ + + # Add 25 TAO stake for move test for Alice + add_move_stake_result = exec_command_alice( + command="stake", + sub_command="add", + extra_args=[ + "--netuid", + "2", + "--wallet-path", + wallet_path_alice, + "--wallet-name", + wallet_alice.name, + "--hotkey", + wallet_alice.hotkey_str, + "--chain", + "ws://127.0.0.1:9945", + "--amount", + "25", + "--no-prompt", + "--era", + "144", + "--unsafe", + "--no-mev-protection", + ], + ) + assert "✅ Finalized" in add_move_stake_result.stdout, add_move_stake_result.stderr + + # List Alice's stakes prior to the move + alice_stake_before_move = exec_command_alice( + command="stake", + sub_command="list", + extra_args=[ + "--wallet-path", + wallet_path_alice, + "--wallet-name", + wallet_alice.name, + "--chain", + "ws://127.0.0.1:9945", + "--no-prompt", + "--verbose", + "--json-output", + ], + ) + + # Check Alice's stakes before move to ensure sufficient stake on netuid 2 + alice_stake_list_before_move = json.loads(alice_stake_before_move.stdout) + alice_stakes_before_move = find_stake_entries( + alice_stake_list_before_move, + netuid=2, + hotkey_ss58=wallet_alice.hotkey.ss58_address, + ) + for stake in alice_stakes_before_move: + assert stake["stake_value"] >= int(20) + + # Move stake from Alice's hotkey on netuid 2 -> Bob's hotkey on netuid 3 + move_amount = 20 + move_result = exec_command_alice( + command="stake", + sub_command="move", + extra_args=[ + "--origin-netuid", + "2", + "--wallet-path", + wallet_path_alice, + "--wallet-name", + wallet_alice.name, + "--wallet-hotkey", + wallet_alice.hotkey_str, + "--dest-netuid", + "3", + "--dest", + wallet_bob.hotkey.ss58_address, + "--amount", + move_amount, + "--chain", + "ws://127.0.0.1:9945", + "--no-prompt", + ], + ) + assert "✅ Sent" in move_result.stdout + + # Check Alice's stakes after move + alice_stake_after_move = exec_command_alice( + command="stake", + sub_command="list", + extra_args=[ + "--wallet-path", + wallet_path_alice, + "--wallet-name", + wallet_alice.name, + "--chain", + "ws://127.0.0.1:9945", + "--no-prompt", + "--verbose", + "--json-output", + ], + ) + # Assert stake was moved from Alice's hotkey on netuid 2 -> Bob's hotkey on netuid 3 + alice_stake_list_after_move = json.loads(alice_stake_after_move.stdout) + bob_stakes_after_move = find_stake_entries( + alice_stake_list_after_move, + netuid=3, + hotkey_ss58=wallet_bob.hotkey.ss58_address, + ) + for stake in bob_stakes_after_move: + assert stake["stake_value"] >= move_amount + + ################################ + # TEST 2: Transfer stake command + # Transfer stake between coldkeys while keeping the same hotkey + ################################ + + transfer_amount = 20 + transfer_fund_root_result = exec_command_alice( + command="stake", + sub_command="add", + extra_args=[ + "--netuid", + "0", + "--wallet-path", + wallet_path_alice, + "--wallet-name", + wallet_alice.name, + "--hotkey", + wallet_alice.hotkey_str, + "--chain", + "ws://127.0.0.1:9945", + "--amount", + transfer_amount, + "--no-prompt", + "--era", + "144", + "--unsafe", + "--no-mev-protection", + ], + ) + assert "✅ Finalized" in transfer_fund_root_result.stdout, ( + transfer_fund_root_result.stderr + ) + + # Ensure Bob doesn't have any stake in root netuid before transfer + bob_stake_list_before_transfer = exec_command_bob( + command="stake", + sub_command="list", + extra_args=[ + "--wallet-path", + wallet_path_bob, + "--wallet-name", + wallet_bob.name, + "--chain", + "ws://127.0.0.1:9945", + "--no-prompt", + "--verbose", + "--json-output", + ], + ) + assert bob_stake_list_before_transfer.stdout == "" + + # Transfer stake from Alice's coldkey on netuid 0 -> Bob's coldkey on netuid 0 + transfer_result = exec_command_alice( + command="stake", + sub_command="transfer", + extra_args=[ + "--origin-netuid", + "0", + "--wallet-path", + wallet_path_alice, + "--wallet-name", + wallet_alice.name, + "--wallet-hotkey", + wallet_alice.hotkey_str, + "--dest-netuid", + "0", + "--dest", + wallet_bob.coldkeypub.ss58_address, + "--all", + "--chain", + "ws://127.0.0.1:9945", + "--no-prompt", + ], + ) + assert "✅ Sent" in transfer_result.stdout + + # Check Bob's stakes after transfer + bob_stake_list_after_transfer = exec_command_bob( + command="stake", + sub_command="list", + extra_args=[ + "--wallet-path", + wallet_path_bob, + "--wallet-name", + wallet_bob.name, + "--chain", + "ws://127.0.0.1:9945", + "--no-prompt", + "--verbose", + "--json-output", + ], + ) + bob_stake_list_after_transfer = json.loads(bob_stake_list_after_transfer.stdout) + bob_stakes_after_transfer = find_stake_entries( + bob_stake_list_after_transfer, + netuid=0, + ) + for stake in bob_stakes_after_transfer: + assert stake["stake_value"] >= transfer_amount + + # Check Alice's stakes after transfer + alice_stake_list_after_transfer = exec_command_alice( + command="stake", + sub_command="list", + extra_args=[ + "--wallet-path", + wallet_path_alice, + "--wallet-name", + wallet_alice.name, + "--chain", + "ws://127.0.0.1:9945", + "--no-prompt", + "--verbose", + "--json-output", + ], + ) + + alice_stake_list_after_transfer = json.loads(alice_stake_list_after_transfer.stdout) + alice_stakes_after_transfer = find_stake_entries( + alice_stake_list_after_transfer, + netuid=0, + ) + if alice_stakes_after_transfer: + pytest.fail("Stake found in root netuid after transfer") + + ################################ + # TEST 3: Swap stake command + # Swap stake between subnets while keeping the same coldkey-hotkey pair + ################################ + + swap_seed_stake_result = exec_command_alice( + command="stake", + sub_command="add", + extra_args=[ + "--netuid", + "4", + "--wallet-path", + wallet_path_alice, + "--wallet-name", + wallet_alice.name, + "--hotkey", + wallet_alice.hotkey_str, + "--chain", + "ws://127.0.0.1:9945", + "--amount", + "25", + "--no-prompt", + "--era", + "144", + "--unsafe", + "--no-mev-protection", + ], + ) + assert "✅ Finalized" in swap_seed_stake_result.stdout, ( + swap_seed_stake_result.stderr + ) + + # Ensure stake was added to Alice's hotkey on netuid 4 + alice_stake_list_before_swap_cmd = exec_command_alice( + command="stake", + sub_command="list", + extra_args=[ + "--wallet-path", + wallet_path_alice, + "--wallet-name", + wallet_alice.name, + "--chain", + "ws://127.0.0.1:9945", + "--no-prompt", + "--verbose", + "--json-output", + ], + ) + + alice_stake_list_before_swap = json.loads(alice_stake_list_before_swap_cmd.stdout) + alice_stakes_before_swap = find_stake_entries( + alice_stake_list_before_swap, + netuid=4, + ) + if not alice_stakes_before_swap: + pytest.fail("Stake not found in netuid 4 before swap") + + # Swap stake from Alice's hotkey on netuid 4 -> Bob's hotkey on netuid 0 + swap_result = exec_command_alice( + command="stake", + sub_command="swap", + extra_args=[ + "--origin-netuid", + "4", + "--wallet-path", + wallet_path_alice, + "--wallet-name", + wallet_alice.name, + "--wallet-hotkey", + wallet_alice.hotkey_str, + "--dest-netuid", + "0", + "--all", + "--chain", + "ws://127.0.0.1:9945", + "--no-prompt", + ], + ) + assert "✅ Sent" in swap_result.stdout, swap_result.stderr + + # Check Alice's stakes after swap + alice_stake_list_after_swap_cmd = exec_command_alice( + command="stake", + sub_command="list", + extra_args=[ + "--wallet-path", + wallet_path_alice, + "--wallet-name", + wallet_alice.name, + "--chain", + "ws://127.0.0.1:9945", + "--no-prompt", + "--verbose", + "--json-output", + ], + ) + + alice_stake_list_after_swap = json.loads(alice_stake_list_after_swap_cmd.stdout) + alice_stakes_after_swap = find_stake_entries( + alice_stake_list_after_swap, + netuid=4, + ) + if alice_stakes_after_swap: + pytest.fail("Stake found in netuid 4 after swap") + + print("Passed stake movement commands") diff --git a/tests/e2e_tests/test_staking_sudo.py b/tests/e2e_tests/test_staking_sudo.py index 6bcaa60cc..4f5346207 100644 --- a/tests/e2e_tests/test_staking_sudo.py +++ b/tests/e2e_tests/test_staking_sudo.py @@ -1,6 +1,7 @@ import asyncio import json import re +import pytest from typing import Union from bittensor_cli.src.bittensor.balances import Balance @@ -24,6 +25,7 @@ """ +@pytest.mark.parametrize("local_chain", [False], indirect=True) def test_staking(local_chain, wallet_setup): """ Test staking & sudo commands and inspect their output @@ -402,7 +404,7 @@ def test_staking(local_chain, wallet_setup): for line in show_stake_adding_single.stdout.splitlines() ] stake_added = cleaned_stake[8].split("│")[3].strip().split()[0] - assert Balance.from_tao(float(stake_added)) >= Balance.from_tao(90) + assert Balance.from_tao(float(stake_added)) >= Balance.from_tao(87) show_stake_json = exec_command_alice( command="stake", @@ -419,9 +421,9 @@ def test_staking(local_chain, wallet_setup): ) show_stake_json_output = json.loads(show_stake_json.stdout) alice_stake = show_stake_json_output["stake_info"][keypair_alice.ss58_address][0] - assert Balance.from_tao(alice_stake["stake_value"]) > Balance.from_tao(90.0) + assert Balance.from_tao(alice_stake["stake_value"]) >= Balance.from_tao(87.0) - # Execute remove_stake command and remove all alpha stakes from Alice + # Execute remove_stake command and remove all alpha stakes from Alice's wallet remove_stake = exec_command_alice( command="stake", sub_command="remove", @@ -472,7 +474,7 @@ def test_staking(local_chain, wallet_setup): "--partial", "--no-prompt", "--era", - "144", + "32", "--json-output", ], ) diff --git a/tests/e2e_tests/test_unstaking.py b/tests/e2e_tests/test_unstaking.py index f3173b5a7..072df278c 100644 --- a/tests/e2e_tests/test_unstaking.py +++ b/tests/e2e_tests/test_unstaking.py @@ -1,12 +1,13 @@ import asyncio import json import re - +import pytest from bittensor_cli.src.bittensor.balances import Balance from .utils import set_storage_extrinsic +@pytest.mark.parametrize("local_chain", [False], indirect=True) def test_unstaking(local_chain, wallet_setup): """ Test various unstaking scenarios including partial unstake, unstake all alpha, and unstake all. diff --git a/tests/e2e_tests/test_wallet_interactions.py b/tests/e2e_tests/test_wallet_interactions.py index 3b92c4965..7ed705b65 100644 --- a/tests/e2e_tests/test_wallet_interactions.py +++ b/tests/e2e_tests/test_wallet_interactions.py @@ -74,6 +74,7 @@ def test_wallet_overview_inspect(local_chain, wallet_setup): "--additional-info", "Test subnet", "--no-prompt", + "--no-mev-protection", ], ) assert f"✅ Registered subnetwork with netuid: {netuid}" in result.stdout @@ -393,6 +394,7 @@ def test_wallet_identities(local_chain, wallet_setup): "--logo-url", "https://testsubnet.com/logo.png", "--no-prompt", + "--no-mev-protection", ], ) assert f"✅ Registered subnetwork with netuid: {netuid}" in result.stdout diff --git a/tests/e2e_tests/utils.py b/tests/e2e_tests/utils.py index 323797356..effd4cef9 100644 --- a/tests/e2e_tests/utils.py +++ b/tests/e2e_tests/utils.py @@ -132,6 +132,34 @@ def extract_coldkey_balance( } +def find_stake_entries( + stake_payload: dict, netuid: int, hotkey_ss58: str | None = None +) -> list[dict]: + """ + Return stake entries matching a given netuid, optionally scoped to a specific hotkey. + Requires json payload using `--json-output` flag. + + Args: + stake_payload: Parsed JSON payload containing `stake_info`. + netuid: The subnet identifier to filter on. + hotkey_ss58: Optional hotkey address to further narrow results. + + Returns: + A list of stake dicts matching the criteria (may be empty). + """ + stake_info = stake_payload.get("stake_info", {}) or {} + matching_stakes: list[dict] = [] + + for stake_hotkey, stakes in stake_info.items(): + if hotkey_ss58 and stake_hotkey != hotkey_ss58: + continue + for stake in stakes or []: + if stake.get("netuid") == netuid: + matching_stakes.append(stake) + + return matching_stakes + + def verify_subnet_entry(output_text: str, netuid: str, ss58_address: str) -> bool: """ Verifies the presence of a specific subnet entry subnets list output.