Assign redemption budget lazily in redeem loop

src/redemptions/commands/process_redeemer.py

@@ -45,7 +45,7 @@
     OsTokenConverter,
     create_os_token_converter,
 )
-from src.redemptions.tasks import assign_shares_to_redeem, is_position_ltv_exceeded
+from src.redemptions.tasks import is_position_ltv_exceeded
 from src.redemptions.typings import OsTokenPosition
 from src.validators.execution import get_withdrawable_assets
 
@@ -271,19 +271,22 @@ async def _redeem_os_token_positions(
     positions_with_processed_shares = await fetch_positions_with_processed_shares(
         nonce=nonce, block_number=block_number
     )
-    logger.info('Assigning shares to redeem')
-    os_token_positions = await assign_shares_to_redeem(
-        positions_with_processed_shares,
-        total_redemption_shares=Wei(queued_shares),
-    )
-    if not os_token_positions:
+    if queued_shares <= 0 or not any(
+        position.unprocessed_shares > 1 for position in positions_with_processed_shares
+    ):
         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})
+        vaults = list(
+            {
+                position.vault
+                for position in positions_with_processed_shares
+                if position.unprocessed_shares > 1
+            }
+        )
         if not await update_vaults_state(vaults=vaults):
             logger.error('Some vaults were left with stale state. Skipping to next interval.')
             return
@@ -294,35 +297,46 @@ async def _redeem_os_token_positions(
     tree = PositionsMerkleTree(all_positions, nonce)
     await redeem_positions(
         tree=tree,
-        os_token_positions=os_token_positions,
+        os_token_positions=positions_with_processed_shares,
+        total_redemption_shares=Wei(queued_shares),
         converter=os_token_converter,
         block_number=block_number,
         dry_run=dry_run,
     )
 
 
+# pylint: disable-next=too-many-arguments,too-many-locals,too-many-branches
 async def redeem_positions(
     tree: PositionsMerkleTree,
     os_token_positions: list[OsTokenPosition],
+    total_redemption_shares: Wei,
     converter: OsTokenConverter,
     block_number: BlockNumber,
     dry_run: bool = False,
 ) -> None:
-    """Redeem positions one by one. Each position's shares_to_redeem is already set by
-    assign_shares_to_redeem; this function further caps it by the vault's withdrawable assets.
-
-    Meta-vault positions are skipped entirely. Vaults whose on-chain state is
-    stale (unharvested) are skipped, since their withdrawable assets would be
-    outdated. A position that fails to simulate or submit is skipped, so the
-    remaining positions are still processed.
+    """Redeem positions one by one, assigning each position's shares_to_redeem lazily from
+    the remaining budget as the loop progresses. A position skipped for any reason (LTV > 1,
+    meta vault, unharvested vault, no live minted position, zero withdrawable, failed
+    simulation or submission) frees its share of the budget for later positions in the
+    file, instead of a fixed prefix of the file consuming the whole budget upfront.
+
+    Each position is further capped by the owner's live minted osToken position (which may
+    have shrunk below the file's leafShares since publication) and by the vault's
+    withdrawable assets.
     """
     vault_to_withdrawable: dict[ChecksumAddress, Wei] = {}
     unharvested_vaults: set[ChecksumAddress] = set()
+    remaining_shares = total_redemption_shares
 
     for position in os_token_positions:
+        if remaining_shares <= 0:
+            break
+
+        unprocessed_shares = position.unprocessed_shares
+        if unprocessed_shares <= 1:
+            continue
+
         logger.info('Processing position index=%d', position.index)
-        shares_to_redeem = position.shares_to_redeem
-        assets_to_redeem = converter.to_assets(shares_to_redeem)
 
         if await is_meta_vault(position.vault):
             logger.warning(
@@ -335,19 +349,30 @@ async def redeem_positions(
         if position.vault in unharvested_vaults:
             continue
 
-        if await is_position_ltv_exceeded(position, converter, block_number):
+        ltv_exceeded, minted_shares = await is_position_ltv_exceeded(
+            position, converter, block_number
+        )
+        if ltv_exceeded:
             logger.info('Skipping position index=%d: LTV > 1', position.index)
             continue
 
-        if position.vault not in vault_to_withdrawable:
-            if await VaultContract(position.vault).is_state_update_required(block_number):
-                logger.info('Skipping unharvested vault %s', position.vault)
-                unharvested_vaults.add(position.vault)
-                continue
-            vault_to_withdrawable[position.vault] = await get_withdrawable_assets(
-                position.vault, block_number=block_number
+        # Cap by the live minted position: the owner may have repaid or been
+        # liquidated after the positions file was published.
+        live_shares = Wei(min(unprocessed_shares, minted_shares))
+        if live_shares <= 0:
+            logger.info(
+                'Skipping position index=%d: owner has no live osToken position', position.index
             )
-        withdrawable = vault_to_withdrawable[position.vault]
+            continue
+
+        withdrawable = await _get_vault_withdrawable(
+            position.vault, vault_to_withdrawable, unharvested_vaults, block_number
+        )
+        if withdrawable is None:
+            continue
+
+        shares_to_redeem = Wei(min(live_shares, remaining_shares))
+        assets_to_redeem = converter.to_assets(shares_to_redeem)
 
         if withdrawable < assets_to_redeem:
             shares_to_redeem = converter.to_shares(withdrawable)
@@ -367,9 +392,29 @@ async def redeem_positions(
             if not await tx_redeem_position(position=position_to_redeem, tree=tree):
                 continue
 
+        remaining_shares = Wei(remaining_shares - shares_to_redeem)
         vault_to_withdrawable[position.vault] = Wei(withdrawable - assets_to_redeem)
 
 
+async def _get_vault_withdrawable(
+    vault: ChecksumAddress,
+    vault_to_withdrawable: dict[ChecksumAddress, Wei],
+    unharvested_vaults: set[ChecksumAddress],
+    block_number: BlockNumber,
+) -> Wei | None:
+    """Cached withdrawable assets for a vault, populated on first use. Returns None (and
+    marks the vault unharvested) when its on-chain state requires an update first."""
+    if vault not in vault_to_withdrawable:
+        if await VaultContract(vault).is_state_update_required(block_number):
+            logger.info('Skipping unharvested vault %s', vault)
+            unharvested_vaults.add(vault)
+            return None
+        vault_to_withdrawable[vault] = await get_withdrawable_assets(
+            vault, block_number=block_number
+        )
+    return vault_to_withdrawable[vault]
+
+
 async def _startup_check() -> None:
     logger.info('Checking connection to execution nodes...')
     await wait_for_execution_node()

src/redemptions/commands/tests/test_process_redeemer.py

@@ -32,6 +32,7 @@ async def test_empty_positions(self) -> None:
             await redeem_positions(
                 tree=make_tree(),
                 os_token_positions=[],
+                total_redemption_shares=Wei(1000),
                 converter=make_converter(),
                 block_number=BlockNumber(100),
             )
@@ -44,6 +45,7 @@ async def test_single_position_sufficient_assets(self) -> None:
             await redeem_positions(
                 tree=make_tree([position]),
                 os_token_positions=[position],
+                total_redemption_shares=Wei(500),
                 converter=make_converter(),
                 block_number=BlockNumber(100),
             )
@@ -58,6 +60,7 @@ async def test_single_position_insufficient_assets_partial_fill(self) -> None:
             await redeem_positions(
                 tree=make_tree([position]),
                 os_token_positions=[position],
+                total_redemption_shares=Wei(500),
                 converter=make_converter(100, 100),
                 block_number=BlockNumber(100),
             )
@@ -72,6 +75,7 @@ async def test_single_position_zero_withdrawable_skipped(self) -> None:
             await redeem_positions(
                 tree=make_tree([position]),
                 os_token_positions=[position],
+                total_redemption_shares=Wei(500),
                 converter=make_converter(100, 100),
                 block_number=BlockNumber(100),
             )
@@ -88,6 +92,7 @@ async def test_multiple_positions_share_vault_cache(self) -> None:
             await redeem_positions(
                 tree=make_tree([pos1, pos2]),
                 os_token_positions=[pos1, pos2],
+                total_redemption_shares=Wei(2000),
                 converter=make_converter(100, 100),
                 block_number=BlockNumber(100),
             )
@@ -102,16 +107,16 @@ async def test_multiple_positions_share_vault_cache(self) -> None:
         assert first_position.owner == OWNER_1 and first_position.shares_to_redeem == Wei(500)
         assert second_position.owner == OWNER_2 and second_position.shares_to_redeem == Wei(200)
 
-    async def test_pre_capped_shares_to_redeem_submitted_not_unprocessed(self) -> None:
-        """assign_shares_to_redeem may cap shares_to_redeem below unprocessed_shares.
-        redeem_positions must submit the pre-capped value, not re-derive from unprocessed_shares."""
-        # unprocessed_shares = 1000, but budget was exhausted mid-position
-        pos = make_position(leaf_shares=1000, processed_shares=0, shares_to_redeem=400)
+    async def test_budget_derived_ignoring_incoming_shares_to_redeem(self) -> None:
+        """shares_to_redeem is assigned lazily from unprocessed_shares and the remaining
+        budget; any incoming shares_to_redeem on the position object is ignored."""
+        pos = make_position(leaf_shares=1000, processed_shares=0, shares_to_redeem=999_999)
 
         with _mock_redeem_positions(withdrawable=Wei(10000)) as mocks:
             await redeem_positions(
                 tree=make_tree([pos]),
                 os_token_positions=[pos],
+                total_redemption_shares=Wei(400),
                 converter=make_converter(),
                 block_number=BlockNumber(100),
             )
@@ -126,6 +131,7 @@ async def test_preserves_original_leaf_shares_in_call(self) -> None:
             await redeem_positions(
                 tree=make_tree([pos]),
                 os_token_positions=[pos],
+                total_redemption_shares=Wei(500),
                 converter=make_converter(),
                 block_number=BlockNumber(100),
             )
@@ -143,6 +149,7 @@ async def test_meta_vault_position_skipped(self) -> None:
             await redeem_positions(
                 tree=make_tree([pos]),
                 os_token_positions=[pos],
+                total_redemption_shares=Wei(500),
                 converter=make_converter(100, 100),
                 block_number=BlockNumber(100),
             )
@@ -161,6 +168,7 @@ async def test_unharvested_vault_skipped(self) -> None:
             await redeem_positions(
                 tree=make_tree([pos]),
                 os_token_positions=[pos],
+                total_redemption_shares=Wei(500),
                 converter=make_converter(100, 100),
                 block_number=BlockNumber(100),
             )
@@ -176,6 +184,7 @@ async def test_ltv_exceeded_position_skipped(self) -> None:
             await redeem_positions(
                 tree=make_tree([pos]),
                 os_token_positions=[pos],
+                total_redemption_shares=Wei(500),
                 converter=make_converter(100, 100),
                 block_number=BlockNumber(100),
             )
@@ -184,7 +193,8 @@ async def test_ltv_exceeded_position_skipped(self) -> None:
         mocks['get_withdrawable'].assert_not_called()
 
     async def test_submit_failure_skips_position(self) -> None:
-        """A failed submission skips that position; subsequent positions are still attempted."""
+        """A failed submission skips that position; subsequent positions are still attempted
+        with the full remaining budget, since the failed position's budget is not consumed."""
         pos1 = make_position(vault=VAULT_1, owner=OWNER_1, processed_shares=500)
         pos2 = make_position(vault=VAULT_2, owner=OWNER_2, processed_shares=500)
 
@@ -195,12 +205,37 @@ async def test_submit_failure_skips_position(self) -> None:
             await redeem_positions(
                 tree=make_tree([pos1, pos2]),
                 os_token_positions=[pos1, pos2],
+                total_redemption_shares=Wei(1000),
                 converter=make_converter(100, 100),
                 block_number=BlockNumber(100),
             )
 
         # The first position fails but the round continues to the second
         assert mocks['submit_mock'].await_count == 2
+        assert _submitted_position(mocks, 1).shares_to_redeem == Wei(500)
+
+    async def test_zero_live_position_skipped_budget_reallocated(self) -> None:
+        """A position whose owner has fully repaid or been liquidated (minted_shares == 0)
+        is skipped without consuming budget, so a later file entry still gets redeemed
+        with the freed budget."""
+        pos1 = make_position(vault=VAULT_1, owner=OWNER_1, leaf_shares=1000, processed_shares=0)
+        pos2 = make_position(vault=VAULT_2, owner=OWNER_2, leaf_shares=1000, processed_shares=0)
+
+        with _mock_redeem_positions(
+            withdrawable=Wei(10000), minted_shares=[Wei(0), Wei(1000)]
+        ) as mocks:
+            await redeem_positions(
+                tree=make_tree([pos1, pos2]),
+                os_token_positions=[pos1, pos2],
+                total_redemption_shares=Wei(1000),
+                converter=make_converter(100, 100),
+                block_number=BlockNumber(100),
+            )
+
+        assert mocks['submit_mock'].await_count == 1
+        submitted = _submitted_position(mocks)
+        assert submitted.owner == OWNER_2
+        assert submitted.shares_to_redeem == Wei(1000)
 
 
 # --- Async function tests (with mocks) ---
@@ -272,12 +307,9 @@ async def test_no_positions_from_ipfs(self) -> None:
         mocks['mock_redeem'].assert_not_called()
 
     async def test_no_eligible_positions(self) -> None:
-        """IPFS returns positions but assign_shares_to_redeem filters them all out."""
-        pos = make_position(leaf_shares=1000)
-        with (
-            _mock_process(positions=[pos]) as mocks,
-            patch(f'{MODULE}.assign_shares_to_redeem', new=AsyncMock(return_value=[])),
-        ):
+        """All fetched positions are fully processed (unprocessed_shares <= 1)."""
+        pos = make_position(leaf_shares=1000, processed_shares=1000, shares_to_redeem=0)
+        with _mock_process(positions=[pos]) as mocks:
             mocks['mock_redeemer'].queued_shares = AsyncMock(return_value=Wei(1000))
             mocks['mock_redeemer'].nonce = AsyncMock(return_value=5)
             await process(block_number=BlockNumber(100), min_queued_assets=Gwei(0))
@@ -296,6 +328,7 @@ async def test_successful_redemption(self) -> None:
         redeem_call = mocks['mock_redeem'].await_args
         # The merkle tree is built from the fetched nonce; leaves use nonce - 1 internally
         assert redeem_call.kwargs['tree'].nonce == 5
+        assert redeem_call.kwargs['total_redemption_shares'] == Wei(1000)
 
     async def test_stale_vault_state_skips_redemption(self) -> None:
         """A failed vault state update leaves stale withdrawable assets and LTVs,
@@ -335,6 +368,7 @@ def _mock_redeem_positions(
     state_update_required: bool = False,
     submit_results: list[bool] | None = None,
     ltv_exceeded: bool = False,
+    minted_shares: Wei | list[Wei] | None = None,
 ) -> Iterator[dict[str, MagicMock]]:
     """Mock setup for redeem_positions tests.
 
@@ -346,6 +380,10 @@ def _mock_redeem_positions(
     abort the round. Simulation always succeeds; each live position is simulated first.
     ``ltv_exceeded`` simulates a position where the user's minted osToken loan exceeds
     their vault assets (LTV > 1), causing the position to be skipped.
+    ``minted_shares`` mocks the live vault.osTokenPositions(owner) value returned
+    alongside the LTV check; a constant applies to every call, a list is consumed in
+    call order (one entry per position). Defaults to an effectively unbounded value so
+    the unprocessed_shares cap is the only one exercised unless a test overrides it.
     """
     if isinstance(withdrawable, AsyncMock):
         get_withdrawable = withdrawable
@@ -364,13 +402,22 @@ def _mock_redeem_positions(
     vault_contract = MagicMock()
     vault_contract.is_state_update_required = AsyncMock(return_value=state_update_required)
 
+    if isinstance(minted_shares, list):
+        is_position_ltv_exceeded_mock = AsyncMock(
+            side_effect=[(ltv_exceeded, shares) for shares in minted_shares]
+        )
+    else:
+        is_position_ltv_exceeded_mock = AsyncMock(
+            return_value=(ltv_exceeded, minted_shares if minted_shares is not None else Wei(10**30))
+        )
+
     with (
         patch(f'{MODULE}.get_withdrawable_assets', new=get_withdrawable),
         patch(f'{MODULE}.is_meta_vault', new=AsyncMock(return_value=is_meta_vault)),
         patch(f'{MODULE}.VaultContract', return_value=vault_contract),
         patch(
             f'{MODULE}.is_position_ltv_exceeded',
-            new=AsyncMock(return_value=ltv_exceeded),
+            new=is_position_ltv_exceeded_mock,
         ),
         patch(f'{MODULE}.simulate_redeem_position', new=simulate_mock),
         patch(f'{MODULE}.tx_redeem_position', new=submit_mock),
@@ -417,10 +464,6 @@ def _mock_process(
             f'{MODULE}.fetch_positions_with_processed_shares',
             new=AsyncMock(return_value=positions),
         ),
-        patch(
-            f'{MODULE}.assign_shares_to_redeem',
-            new=AsyncMock(return_value=positions),
-        ),
         patch(
             f'{MODULE}.update_processed_shares_cache',
             new=AsyncMock(),