Restore update_vaults_state with batched multicall

Author: evgeny-stakewise

src/config/settings.py

@@ -410,6 +410,9 @@ def network_config(self) -> NetworkConfig:
 )
 EVENTS_CONCURRENCY_LIMIT: int = decouple_config('EVENTS_CONCURRENCY_LIMIT', default=10, cast=int)
 
+# Maximum number of updateState calls batched into a single multicall transaction.
+MULTICALL_CHUNK_SIZE: int = decouple_config('MULTICALL_CHUNK_SIZE', default=20, cast=int)
+
 # Backoff retries
 DEFAULT_RETRY_TIME = 60
 

src/redemptions/commands/process_redeemer.py

@@ -32,6 +32,7 @@
     simulate_redeem_position,
     tx_process_exit_queue,
     tx_redeem_position,
+    update_vaults_state,
 )
 from src.redemptions.fetch_positions import (
     cached_fetch_positions_from_ipfs,
@@ -273,6 +274,15 @@ async def _redeem_os_token_positions(
         logger.info('No redeemable positions found. Skipping to next interval.')
         return
 
+    if not dry_run:
+        # Bring vaults up to date on-chain so withdrawable assets and position LTV
+        # are computed from fresh state rather than skipping unharvested vaults.
+        vaults = list({position.vault for position in os_token_positions})
+        await update_vaults_state(vaults=vaults)
+
+        # Re-fetch the block number so the freshly-updated state is visible downstream.
+        block_number = await execution_client.eth.block_number
+
     tree = PositionsMerkleTree(all_positions, nonce)
     await redeem_positions(
         tree=tree,
@@ -323,7 +333,7 @@ async def redeem_positions(
                 unharvested_vaults.add(position.vault)
                 continue
             vault_to_withdrawable[position.vault] = await get_withdrawable_assets(
-                position.vault, harvest_params=None
+                position.vault, block_number=block_number
             )
         withdrawable = vault_to_withdrawable[position.vault]
 

src/redemptions/commands/tests/test_process_redeemer.py

@@ -392,13 +392,18 @@ def _mock_process(
             f'{MODULE}.redeem_positions',
             new=AsyncMock(),
         ) as mock_redeem,
+        patch(
+            f'{MODULE}.update_vaults_state',
+            new=AsyncMock(),
+        ) as mock_update_state,
         patch(f'{MODULE}.execution_client', new=mock_client),
     ):
         mock_settings.network_config.VAULT_BALANCE_SYMBOL = 'ETH'
         mock_redeemer.can_process_exit_queue = AsyncMock(return_value=False)
         yield {
             'mock_redeemer': mock_redeemer,
             'mock_redeem': mock_redeem,
+            'mock_update_state': mock_update_state,
         }
 
 

src/redemptions/execution.py

@@ -1,22 +1,78 @@
 import logging
+from itertools import batched
 
+from eth_typing import ChecksumAddress, HexStr
+from hexbytes import HexBytes
 from web3 import Web3
 from web3.contract.async_contract import AsyncContractFunction
 from web3.exceptions import Web3Exception
 
 from src.common.clients import execution_client
+from src.common.contracts import VaultContract, multicall_contract
 from src.common.execution import (
     transaction_gas_wrapper,
     wait_for_execution_endpoints_synced,
 )
-from src.config.settings import settings
+from src.common.harvest import get_multiple_harvest_params
+from src.config.settings import MULTICALL_CHUNK_SIZE, settings
+from src.meta_vault.service import is_meta_vault
 from src.redemptions.contracts import os_token_redeemer_contract
 from src.redemptions.merkle_tree import PositionsMerkleTree
 from src.redemptions.typings import OsTokenPosition
 
 logger = logging.getLogger(__name__)
 
 
+async def update_vaults_state(
+    vaults: list[ChecksumAddress],
+) -> None:
+    """Bring every regular vault in ``vaults`` up to date on-chain via batched
+    updateState multicalls.
+
+    Meta vaults are skipped. Vaults are selected via ``can_harvest`` at the
+    latest state. Fresh on-chain state lets ``get_withdrawable_assets`` (and
+    position LTV) read accurate values without threading harvest_params through
+    every call.
+    """
+    regular_vaults: list[ChecksumAddress] = []
+    for vault in vaults:
+        if await is_meta_vault(vault):
+            continue
+        regular_vaults.append(vault)
+
+    if not regular_vaults:
+        return
+
+    vault_to_harvest_params = await get_multiple_harvest_params(regular_vaults)
+    calls: list[tuple[ChecksumAddress, HexStr]] = []
+    for vault in regular_vaults:
+        harvest_params = vault_to_harvest_params.get(vault)
+        if harvest_params is None:
+            continue
+        vault_contract = VaultContract(vault)
+        calls.append(
+            (vault_contract.contract_address, vault_contract.get_update_state_call(harvest_params))
+        )
+
+    if not calls:
+        return
+
+    for chunk in batched(calls, MULTICALL_CHUNK_SIZE):
+        await tx_update_vaults_state(list(chunk))
+
+
+async def tx_update_vaults_state(calls: list[tuple[ChecksumAddress, HexStr]]) -> None:
+    """Submit a single updateState multicall and wait for its receipt."""
+    tx_hash = await multicall_contract.tx_aggregate(calls)
+    logger.info('Waiting for updateState multicall tx %s confirmation', tx_hash)
+    tx_receipt = await execution_client.eth.wait_for_transaction_receipt(
+        HexBytes(Web3.to_bytes(hexstr=tx_hash)), timeout=settings.execution_transaction_timeout
+    )
+    if not tx_receipt['status']:
+        raise RuntimeError(f'updateState multicall tx failed. Tx Hash: {tx_hash}')
+    logger.info('updateState multicall confirmed. Tx Hash: %s', tx_hash)
+
+
 async def tx_process_exit_queue() -> None:
     """Call processExitQueue() on the redeemer contract."""
     tx_hash = await os_token_redeemer_contract.process_exit_queue()

src/redemptions/tests/test_execution.py

@@ -13,13 +13,15 @@
     simulate_redeem_position,
     tx_process_exit_queue,
     tx_redeem_position,
+    update_vaults_state,
 )
 from src.redemptions.merkle_tree import PositionsMerkleTree
 from src.redemptions.typings import OsTokenPosition
 
 MODULE = 'src.redemptions.execution'
 
 VAULT_1 = Web3.to_checksum_address('0x' + '11' * 20)
+VAULT_2 = Web3.to_checksum_address('0x' + '22' * 20)
 OWNER_1 = Web3.to_checksum_address('0x' + '33' * 20)
 
 
@@ -111,9 +113,86 @@ async def test_process_exit_queue(self, tx_status: int) -> None:
             mock_redeemer.process_exit_queue.assert_called_once()
 
 
+class TestUpdateVaultsState:
+    async def test_submits_multicall_for_harvestable_vaults(self) -> None:
+        with _mock_update_vaults_state() as mocks:
+            await update_vaults_state(vaults=[VAULT_1, VAULT_2])
+        mocks['tx_aggregate'].assert_awaited_once()
+        await_args = mocks['tx_aggregate'].await_args
+        assert await_args is not None
+        # One updateState call per harvestable vault.
+        assert len(await_args.args[0]) == 2
+
+    async def test_skips_meta_vaults(self) -> None:
+        with _mock_update_vaults_state(is_meta_vault=True) as mocks:
+            await update_vaults_state(vaults=[VAULT_1])
+        mocks['tx_aggregate'].assert_not_awaited()
+
+    async def test_skips_when_not_harvestable(self) -> None:
+        """can_harvest returns no harvest params, so there is nothing to update."""
+        with _mock_update_vaults_state(harvest_params=None) as mocks:
+            await update_vaults_state(vaults=[VAULT_1])
+        mocks['tx_aggregate'].assert_not_awaited()
+
+    async def test_batches_calls_into_chunks(self) -> None:
+        """Calls are split into MULTICALL_CHUNK_SIZE-sized transactions."""
+        with _mock_update_vaults_state(chunk_size=1) as mocks:
+            await update_vaults_state(vaults=[VAULT_1, VAULT_2])
+        # Two vaults, chunk size 1 → one multicall transaction per vault.
+        assert mocks['tx_aggregate'].await_count == 2
+        for call in mocks['tx_aggregate'].await_args_list:
+            assert len(call.args[0]) == 1
+
+    async def test_raises_on_failed_receipt(self) -> None:
+        with _mock_update_vaults_state(tx_status=0):
+            with pytest.raises(RuntimeError, match='updateState multicall tx failed'):
+                await update_vaults_state(vaults=[VAULT_1])
+
+
 # --- Helpers ---
 
 
+@contextmanager
+def _mock_update_vaults_state(
+    is_meta_vault: bool = False,
+    harvest_params: object = object(),
+    tx_status: int = 1,
+    chunk_size: int | None = None,
+) -> Iterator[dict[str, AsyncMock]]:
+    """Mock setup for update_vaults_state tests. VaultContract, harvest params,
+    the multicall contract and the execution client are all stubbed so only the
+    orchestration logic is exercised. ``harvest_params=None`` models a vault that
+    can_harvest filtered out. ``chunk_size`` overrides MULTICALL_CHUNK_SIZE."""
+    vault_contract = MagicMock()
+    vault_contract.contract_address = VAULT_1
+    vault_contract.get_update_state_call = MagicMock(return_value='0xdeadbeef')
+
+    tx_aggregate = AsyncMock(return_value='0x' + 'ab' * 32)
+    mock_multicall = MagicMock()
+    mock_multicall.tx_aggregate = tx_aggregate
+
+    mock_client = MagicMock()
+    mock_client.eth.wait_for_transaction_receipt = AsyncMock(return_value={'status': tx_status})
+
+    def harvest_params_for(vaults: list, _block: BlockNumber | None = None) -> dict:
+        return {vault: harvest_params for vault in vaults}
+
+    with (
+        patch(f'{MODULE}.is_meta_vault', new=AsyncMock(return_value=is_meta_vault)),
+        patch(f'{MODULE}.VaultContract', return_value=vault_contract),
+        patch(
+            f'{MODULE}.get_multiple_harvest_params',
+            new=AsyncMock(side_effect=harvest_params_for),
+        ),
+        patch(f'{MODULE}.multicall_contract', new=mock_multicall),
+        patch(f'{MODULE}.execution_client', new=mock_client),
+        patch(f'{MODULE}.MULTICALL_CHUNK_SIZE', new=chunk_size if chunk_size is not None else 20),
+        patch(f'{MODULE}.settings') as mock_settings,
+    ):
+        mock_settings.execution_transaction_timeout = 60
+        yield {'tx_aggregate': tx_aggregate}
+
+
 @contextmanager
 def _mock_tx_redeem_position(
     tx_status: int = 1,

src/validators/execution.py

@@ -140,7 +140,7 @@ async def tx_fund_validators(
 
 async def get_withdrawable_assets(
     vault: ChecksumAddress,
-    harvest_params: HarvestParams | None,
+    harvest_params: HarvestParams | None = None,
     block_number: BlockNumber | None = None,
 ) -> Wei:
     """Fetches vault's available assets for staking."""