Rework shielded transactions

bittensor/core/async_subtensor.py

@@ -3361,125 +3361,6 @@ async def get_mev_shield_next_key(
 
         return public_key_bytes
 
-    async def get_mev_shield_submission(
-        self,
-        submission_id: str,
-        block: Optional[int] = None,
-        block_hash: Optional[str] = None,
-        reuse_block: bool = False,
-    ) -> Optional[dict[str, str | int | bytes]]:
-        """
-        Retrieves Submission from the MevShield pallet storage.
-
-        If submission_id is provided, returns a single submission. If submission_id is None, returns all submissions from
-        the storage map.
-
-        Parameters:
-            submission_id: The hash ID of the submission. Can be a hex string with "0x" prefix or bytes. If None,
-                returns all submissions.
-            block: The blockchain block number for the query.
-            block_hash: The hash of the block to retrieve the stake from. Do not specify if using block or reuse_block.
-            reuse_block: Whether to use the last-used block. Do not set if using block_hash or block.
-
-        Returns:
-            If submission_id is provided: A dictionary containing the submission data if found, None otherwise. The
-                dictionary contains:
-                - author: The SS58 address of the account that submitted the encrypted extrinsic
-                - commitment: The blake2_256 hash of the payload_core (as hex string with "0x" prefix)
-                - ciphertext: The encrypted blob as bytes (format: [u16 kem_len][kem_ct][nonce24][aead_ct])
-                - submitted_in: The block number when the submission was created
-
-            If submission_id is None: A dictionary mapping submission IDs (as hex strings) to submission dictionaries.
-
-        Note:
-            If a specific submission does not exist in storage, this function returns None. If querying all submissions
-            and none exist, returns an empty dictionary.
-        """
-        block_hash = await self.determine_block_hash(block, block_hash, reuse_block)
-        submission_id = (
-            submission_id[2:] if submission_id.startswith("0x") else submission_id
-        )
-        submission_id_bytes = bytes.fromhex(submission_id)
-
-        query = await self.substrate.query(
-            module="MevShield",
-            storage_function="Submissions",
-            params=[submission_id_bytes],
-            block_hash=block_hash,
-        )
-
-        query_value = getattr(query, "value", query)
-        if query_value is None or not isinstance(query_value, dict):
-            return None
-
-        author_raw = cast(Union[bytes, str], query_value.get("author"))
-        commitment_raw = cast(list[bytes], query_value.get("commitment"))
-        ciphertext_raw = cast(list[bytes], query_value.get("ciphertext"))
-        submitted_in = cast(int, query_value.get("submitted_in"))
-
-        autor = decode_account_id(author_raw)
-        commitment = bytes(commitment_raw[0])
-        ciphertext = bytes(ciphertext_raw[0])
-
-        return {
-            "author": autor,
-            "commitment": commitment,
-            "ciphertext": ciphertext,
-            "submitted_in": submitted_in,
-        }
-
-    async def get_mev_shield_submissions(
-        self,
-        block: Optional[int] = None,
-        block_hash: Optional[str] = None,
-        reuse_block: bool = False,
-    ) -> Optional[dict[str, dict[str, str | int]]]:
-        """
-        Retrieves all encrypted submissions from the MevShield pallet storage.
-
-        This function queries the MevShield.Submissions storage map and returns all pending encrypted submissions that
-        have been submitted via submit_encrypted but not yet executed via execute_revealed.
-
-        Parameters:
-            block: The blockchain block number for the query. If None, uses the current block.
-            block_hash: The hash of the block to retrieve the submissions from. Do not specify if using block or reuse_block.
-            reuse_block: Whether to use the last-used block. Do not set if using block_hash or block.
-
-        Returns:
-            A dictionary mapping wrapper_id (as hex string with "0x" prefix) to submission data dictionaries. Each
-            submission dictionary contains:
-            - author: The SS58 address of the account that submitted the encrypted extrinsic
-            - commitment: The blake2_256 hash of the payload_core as bytes (32 bytes)
-            - ciphertext: The encrypted blob as bytes (format: [u16 kem_len][kem_ct][nonce24][aead_ct])
-            - submitted_in: The block number when the submission was created
-
-            Returns None if no submissions exist in storage at the specified block.
-
-        Note:
-            Submissions are automatically pruned after KEY_EPOCH_HISTORY blocks (100 blocks) by the pallet's
-            on_initialize hook. Only submissions that have been submitted but not yet executed will be present in
-            storage.
-        """
-        block_hash = await self.determine_block_hash(block, block_hash, reuse_block)
-        query = await self.substrate.query_map(
-            module="MevShield",
-            storage_function="Submissions",
-            block_hash=block_hash,
-        )
-
-        result = {}
-        async for q in query:
-            key, value = q
-            value = value.value
-            result["0x" + bytes(key[0]).hex()] = {
-                "author": decode_account_id(value.get("author")),
-                "commitment": bytes(value.get("commitment")[0]),
-                "ciphertext": bytes(value.get("ciphertext")[0]),
-                "submitted_in": value.get("submitted_in"),
-            }
-
-        return result if result else None
-
     async def get_minimum_required_stake(self):
         """Returns the minimum required stake threshold for nominator cleanup operations.
 

bittensor/core/errors.py

@@ -57,6 +57,8 @@
     "TxRateLimitExceeded",
     "UnknownSynapseError",
     "UnstakeError",
+    "SHIELD_VALIDATION_ERRORS",
+    "map_shield_error",
 ]
 
 
@@ -263,3 +265,32 @@ def __init__(
     ):
         self.message = message
         super().__init__(self.message, synapse)
+
+
+SHIELD_VALIDATION_ERRORS = {
+    "Custom error: 23": (
+        "Failed to parse shielded transaction: the ciphertext has an invalid format."
+    ),
+    "Custom error: 24": (
+        "Invalid encryption key: the key_hash in the ciphertext does not match any known key. "
+        "The key may have rotated between reading NextKey and submitting the transaction."
+    ),
+}
+
+
+def map_shield_error(raw_message: str) -> str:
+    """Map a raw shield validation error to a human-readable description.
+
+    Checks the message against known Custom error codes from CheckShieldedTxValidity,
+    then falls back to detecting a generic ``"invalid"`` subscription status.
+    Returns the original message unchanged if nothing matches.
+    """
+    for marker, description in SHIELD_VALIDATION_ERRORS.items():
+        if marker in raw_message:
+            return description
+    if "'result': 'invalid'" in raw_message.lower():
+        return (
+            "MEV Shield extrinsic rejected as invalid. "
+            "The key may have rotated between reading NextKey and submission."
+        )
+    return raw_message

bittensor/core/extrinsics/asyncex/mev_shield.py

@@ -3,10 +3,12 @@
 from typing import TYPE_CHECKING, Optional
 
 from async_substrate_interface import AsyncExtrinsicReceipt
+from async_substrate_interface.errors import SubstrateRequestException
 
+from bittensor.core.errors import map_shield_error
 from bittensor.core.extrinsics.pallets import MevShield
 from bittensor.core.extrinsics.utils import (
-    get_mev_commitment_and_ciphertext,
+    get_mev_shielded_ciphertext,
     get_event_data_by_event_name,
 )
 from bittensor.core.types import ExtrinsicResponse
@@ -22,23 +24,18 @@
 async def wait_for_extrinsic_by_hash(
     subtensor: "AsyncSubtensor",
     extrinsic_hash: str,
-    shield_id: str,
     submit_block_hash: str,
     timeout_blocks: int = 3,
 ) -> Optional["AsyncExtrinsicReceipt"]:
     """
     Wait for the result of a MeV Shield encrypted extrinsic.
 
     After submit_encrypted succeeds, the block author will decrypt and submit the inner extrinsic directly. This
-    function polls subsequent blocks looking for either:
-    - an extrinsic matching the provided hash (success)
-    OR
-    - a markDecryptionFailed extrinsic with matching shield ID (failure)
+    function polls subsequent blocks looking for an extrinsic matching the provided hash.
 
     Args:
         subtensor: SubtensorInterface instance.
         extrinsic_hash: The hash of the inner extrinsic to find.
-        shield_id: The wrapper ID from EncryptedSubmitted event (for detecting decryption failures).
         submit_block_hash: Block hash where submit_encrypted was included.
         timeout_blocks: Max blocks to wait.
 
@@ -59,34 +56,14 @@ async def wait_for_extrinsic_by_hash(
         block_hash = await subtensor.substrate.get_block_hash(current_block)
         extrinsics = await subtensor.substrate.get_extrinsics(block_hash)
 
-        result_idx = None
         for idx, extrinsic in enumerate(extrinsics):
-            # Success: Inner extrinsic executed
             if f"0x{extrinsic.extrinsic_hash.hex()}" == extrinsic_hash:
-                result_idx = idx
-                break
-
-            # Failure: Decryption failed
-            call = extrinsic.value.get("call", {})
-            if (
-                call.get("call_module") == "MevShield"
-                and call.get("call_function") == "mark_decryption_failed"
-            ):
-                call_args = call.get("call_args", [])
-                for arg in call_args:
-                    if arg.get("name") == "id" and arg.get("value") == shield_id:
-                        result_idx = idx
-                        break
-                if result_idx is not None:
-                    break
-
-        if result_idx is not None:
-            return AsyncExtrinsicReceipt(
-                substrate=subtensor.substrate,
-                block_hash=block_hash,
-                block_number=current_block,
-                extrinsic_idx=result_idx,
-            )
+                return AsyncExtrinsicReceipt(
+                    substrate=subtensor.substrate,
+                    block_hash=block_hash,
+                    block_number=current_block,
+                    extrinsic_idx=idx,
+                )
 
         current_block += 1
 
@@ -125,7 +102,7 @@ async def submit_encrypted_extrinsic(
         wait_for_finalization: Whether to wait for the finalization of the transaction.
         wait_for_revealed_execution: Whether to wait for the executed event, indicating that validators have
             successfully decrypted and executed the inner call. If True, the function will poll subsequent blocks for
-            the event matching this submission's commitment.
+            the extrinsic matching this submission.
         blocks_for_revealed_execution: Maximum number of blocks to poll for the executed event after inclusion.
             The function checks blocks from start_block to start_block + blocks_for_revealed_execution. Returns
             immediately if the event is found before the block limit is reached.
@@ -138,13 +115,8 @@ async def submit_encrypted_extrinsic(
         SubstrateRequestException: If the extrinsic fails to be submitted or included.
 
     Note:
-        The encryption uses the public key from NextKey storage, which rotates every block. The payload structure is:
-        payload_core = signer_bytes (32B) + key_hash (32B Blake2-256 hash of NextKey) + SCALE(call)
-        plaintext = payload_core + b"\\x01" + signature (64B for sr25519)
-        commitment = blake2_256(payload_core)
-
-        The key_hash binds the transaction to the key epoch at submission time and replaces nonce-based replay
-        protection.
+        The encryption uses the public key from NextKey storage, which rotates every block. The ciphertext wire format
+        is: [key_hash(16)][u16 kem_len LE][kem_ct][nonce24][aead_ct], where key_hash = twox_128(NextKey).
     """
     try:
         if sign_with not in ("coldkey", "hotkey"):
@@ -186,15 +158,12 @@ async def submit_encrypted_extrinsic(
             call=call, keypair=inner_signing_keypair, nonce=next_nonce, era=era
         )
 
-        mev_commitment, mev_ciphertext, payload_core = (
-            get_mev_commitment_and_ciphertext(
-                signed_ext=signed_extrinsic,
-                ml_kem_768_public_key=ml_kem_768_public_key,
-            )
+        mev_ciphertext = get_mev_shielded_ciphertext(
+            signed_ext=signed_extrinsic,
+            ml_kem_768_public_key=ml_kem_768_public_key,
         )
 
         extrinsic_call = await MevShield(subtensor).submit_encrypted(
-            commitment=mev_commitment,
             ciphertext=mev_ciphertext,
         )
 
@@ -211,10 +180,8 @@ async def submit_encrypted_extrinsic(
 
         if response.success:
             response.data = {
-                "commitment": mev_commitment,
                 "ciphertext": mev_ciphertext,
                 "ml_kem_768_public_key": ml_kem_768_public_key,
-                "payload_core": payload_core,
                 "signed_extrinsic_hash": f"0x{signed_extrinsic.extrinsic_hash.hex()}",
             }
             if wait_for_revealed_execution:
@@ -230,12 +197,9 @@ async def submit_encrypted_extrinsic(
                         error=RuntimeError("EncryptedSubmitted event not found."),
                     )
 
-                shield_id = event["attributes"]["id"]
-
                 response.mev_extrinsic = await wait_for_extrinsic_by_hash(
                     subtensor=subtensor,
                     extrinsic_hash=f"0x{signed_extrinsic.extrinsic_hash.hex()}",
-                    shield_id=shield_id,
                     submit_block_hash=response.extrinsic_receipt.block_hash,
                     timeout_blocks=blocks_for_revealed_execution,
                 )
@@ -251,7 +215,7 @@ async def submit_encrypted_extrinsic(
                     response.message = format_error_message(
                         await response.mev_extrinsic.error_message
                     )
-                    response.error = RuntimeError(response.message)
+                    response.error = SubstrateRequestException(response.message)
                     response.success = False
                     if raise_error:
                         raise response.error
@@ -260,6 +224,10 @@ async def submit_encrypted_extrinsic(
                         "[green]Encrypted extrinsic submitted successfully.[/green]"
                     )
         else:
+            response.message = map_shield_error(str(response.message))
+            response.error = SubstrateRequestException(response.message)
+            if raise_error:
+                raise response.error
             logging.error(f"[red]{response.message}[/red]")
 
         return response