Add support for stateOverrideSet when calling eth_estimateUserOp

src/ethproto/aa_bundler.py

@@ -14,7 +14,7 @@
 from hexbytes import HexBytes
 from web3 import Web3
 from web3.constants import ADDRESS_ZERO
-from web3.types import TxParams
+from web3.types import StateOverride, TxParams
 
 from .contracts import RevertError
 
@@ -30,6 +30,8 @@
 AA_BUNDLER_BASE_GAS_PRICE_FACTOR = env.float("AA_BUNDLER_BASE_GAS_PRICE_FACTOR", 1)
 AA_BUNDLER_VERIFICATION_GAS_FACTOR = env.float("AA_BUNDLER_VERIFICATION_GAS_FACTOR", 1)
 
+AA_BUNDLER_STATE_OVERRIDES = env.json("AA_BUNDLER_STATE_OVERRIDES", default={})
+
 NonceMode = Enum(
     "NonceMode",
     [
@@ -68,6 +70,16 @@
 )
 
 
+class BundlerRevertError(RevertError):
+    """Bundler specific revert error"""
+
+    def __init__(self, message, userop=None, response=None):
+        super().__init__(message)
+        self.message = message
+        self.userop = userop
+        self.response = response
+
+
 @dataclass(frozen=True)
 class UserOpEstimation:
     """eth_estimateUserOperationGas response"""
@@ -288,11 +300,11 @@ def check_nonce_error(resp, retry_nonce):
     if "AA25" in resp["error"]["message"] and AA_BUNDLER_MAX_GETNONCE_RETRIES > 0:
         # Retry fetching the nonce
         if retry_nonce == AA_BUNDLER_MAX_GETNONCE_RETRIES:
-            raise RevertError(resp["error"]["message"])
+            raise BundlerRevertError(resp["error"]["message"], response=resp)
         warn(f'{resp["error"]["message"]} error, I will retry fetching the nonce')
         return (retry_nonce or 0) + 1
     else:
-        raise RevertError(resp["error"]["message"])
+        raise BundlerRevertError(resp["error"]["message"], response=resp)
 
 
 def get_sender(tx):
@@ -318,6 +330,7 @@ def __init__(
         priority_gas_price_factor: float = AA_BUNDLER_PRIORITY_GAS_PRICE_FACTOR,
         base_gas_price_factor: float = AA_BUNDLER_BASE_GAS_PRICE_FACTOR,
         executor_pk: HexBytes = AA_BUNDLER_EXECUTOR_PK,
+        overrides: StateOverride = AA_BUNDLER_STATE_OVERRIDES,
     ):
         self.w3 = w3
         self.bundler = Web3(Web3.HTTPProvider(bundler_url), middleware=[])
@@ -331,6 +344,10 @@ def __init__(
         self.base_gas_price_factor = base_gas_price_factor
         self.executor_pk = executor_pk
 
+        # stateOverrideSet mapping to use when calling eth_estimateUserOperationGas
+        # https://docs.alchemy.com/reference/eth-estimateuseroperationgas
+        self.overrides = overrides
+
     def __str__(self):
         return (
             f"Bundler(type={self.bundler_type}, entrypoint={self.entrypoint}, nonce_mode={self.nonce_mode}, "
@@ -364,10 +381,11 @@ def get_base_fee(self):
 
     def estimate_user_operation_gas(self, user_operation: UserOperation) -> UserOpEstimation:
         resp = self.bundler.provider.make_request(
-            "eth_estimateUserOperationGas", [user_operation.as_reduced_dict(), self.entrypoint]
+            "eth_estimateUserOperationGas",
+            [user_operation.as_reduced_dict(), self.entrypoint, self.overrides],
         )
         if "error" in resp:
-            raise RevertError(resp["error"]["message"])
+            raise BundlerRevertError(resp["error"]["message"], user_operation, resp)
 
         paymaster_verification_gas_limit = resp["result"].get("paymasterVerificationGasLimit", "0x00")
         return UserOpEstimation(
@@ -386,7 +404,7 @@ def estimate_user_operation_gas(self, user_operation: UserOperation) -> UserOpEs
     def alchemy_gas_price(self):
         resp = self.bundler.provider.make_request("rundler_maxPriorityFeePerGas", [])
         if "error" in resp:
-            raise RevertError(resp["error"]["message"])
+            raise BundlerRevertError(resp["error"]["message"], response=resp)
         max_priority_fee_per_gas = int(int(resp["result"], 16) * self.priority_gas_price_factor)
         max_fee_per_gas = max_priority_fee_per_gas + self.get_base_fee()
 
@@ -424,7 +442,14 @@ def send_transaction(self, tx: Tx, retry_nonce=None):
             "eth_sendUserOperation", [user_operation.as_dict(), self.entrypoint]
         )
         if "error" in resp:
-            next_nonce = check_nonce_error(resp, retry_nonce)
+            try:
+                next_nonce = check_nonce_error(resp, retry_nonce)
+            except BundlerRevertError as e:
+                raise BundlerRevertError(
+                    e.message,
+                    userop=user_operation,
+                    response=e.response,
+                )
             return self.send_transaction(tx, retry_nonce=next_nonce)
 
         return {"userOpHash": resp["result"]}

tests/cassettes/test_aa_bundler/test_build_user_operation.yaml

@@ -23,7 +23,7 @@ interactions:
         [{"sender": "0xE8B412158c205B0F605e0FC09dCdA27d3F140FE9", "nonce": "0xae85c374ae0606ed34d0ee009a9ca43a757a8a46a324510000000000000000",
         "callData": "0xb61d27f60000000000000000000000002791bca1f2de4661ed88a30c99a7a9449aa84174000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000600000000000000000000000000000000000000000000000000000000000000044095ea7b30000000000000000000000007ace242f32208d836a2245df957c08547059bf45ffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffff00000000000000000000000000000000000000000000000000000000",
         "signature": "0xfffffffffffffffffffffffffffffff0000000000000000000000000000000007aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa1c"},
-        "0x0000000071727De22E5E9d8BAf0edAc6f37da032"], "id": 0}'
+        "0x0000000071727De22E5E9d8BAf0edAc6f37da032", {}], "id": 0}'
       headers:
         Content-Length:
           - "885"

tests/test_aa_bundler.py

@@ -182,7 +182,8 @@ def test_send_transaction():
 
     def make_request(method, params):
         if method == "eth_estimateUserOperationGas":
-            assert len(params) == 2
+            assert len(params) == 3
+            assert params[2] == {}
             assert params[1] == ENTRYPOINT
             assert params[0] == {
                 "sender": "0xE8B412158c205B0F605e0FC09dCdA27d3F140FE9",