fix: withdraw limits

contracts/VaultV3.vy

@@ -551,13 +551,88 @@ def _max_deposit(receiver: address) -> uint256:
 
 @view
 @internal
-def _max_redeem(owner: address) -> uint256:
-    return self.balance_of[owner]
+def _max_withdraw(
+    owner: address,
+    max_loss: uint256,
+    strategies: DynArray[address, MAX_QUEUE]
+) -> uint256:
+    """
+    @dev Returns the max amount of `asset` an `owner` can withdraw.
 
-@view
-@internal
-def _max_withdraw(owner: address) -> uint256:
-    return self._convert_to_assets(self.balance_of[owner], Rounding.ROUND_DOWN)
+    This will do a full simulation of the withdraw in order to determine
+    how much is currently liquid and if the `max_loss` would allow for the 
+    tx to not revert.
+
+    This will track any expected loss to check if the tx will revert, but
+    not account for it in the amount returned since it is unrealised and 
+    therefore will not be accounted for in the conversion rates.
+
+    i.e. If we have 100 debt and 10 of unrealised loss, the max we can get
+    out is 90, but a user of the vault will need to call withdraw with 100
+    in order to get the full 90 out.
+    """
+    # Get the max amount for the owner if fully liquid.
+    max_assets: uint256 = self._convert_to_assets(self.balance_of[owner], Rounding.ROUND_DOWN)
+
+    # See if we have enough idle to service the withdraw.
+    current_idle: uint256 = self.total_idle
+    if max_assets > current_idle:
+        # Amount left that we need.
+        needed: uint256 = unsafe_sub(max_assets, current_idle)
+        # Track how much we can pull.
+        have: uint256 = current_idle
+        loss: uint256 = 0
+
+        # If no queue was passed use the default one.
+        _strategies: DynArray[address, MAX_QUEUE] = strategies
+        if len(_strategies) == 0:
+            _strategies = self.default_queue
+
+        for strategy in _strategies:
+            # Can't use an invalid strategy.
+            assert self.strategies[strategy].activation != 0, "inactive strategy"
+
+            # Get the maximum amount the vault can withdraw from the strategy.
+            max_withdraw: uint256 = min(needed, self.strategies[strategy].current_debt)
+
+            # Get any unrealised loss for the strategy.
+            unrealised_loss: uint256 = self._assess_share_of_unrealised_losses(strategy, max_withdraw)
+
+            # See if any limit is enforced by the strategy.
+            strategy_limit: uint256 = IStrategy(strategy).maxWithdraw(self)
+
+            # Adjust accordingly if there is a max withdraw limit.
+            if strategy_limit < max_withdraw - unrealised_loss:
+                # lower unrealised loss to the proportion to the limit.
+                unrealised_loss = unrealised_loss * strategy_limit / max_withdraw
+                # Still count the unrealised loss as withdrawable.
+                max_withdraw = strategy_limit + unrealised_loss
+
+            # If 0 move on to the next strategy.
+            if max_withdraw == 0:
+                continue
+
+            # Add to what we can pull.
+            have += max_withdraw
+            # Reduce how much we have left.
+            needed -= max_withdraw
+            # Add any unrealised loss to the total
+            loss += unrealised_loss
+
+            if needed == 0:
+                break
+
+        # Update the max after going through the queue.
+        max_assets = have
+
+        # Check if there is a loss and a non-default value was set.
+        if loss > 0 and max_loss < MAX_BPS:
+            # If the loss is not within the allowed range.
+            if loss > max_assets * max_loss / MAX_BPS:
+                # The tx will revert. Can only withdraw idle.
+                max_assets = current_idle
+
+    return max_assets
 
 @internal
 def _deposit(sender: address, recipient: address, assets: uint256) -> uint256:
@@ -1798,23 +1873,41 @@ def maxMint(receiver: address) -> uint256:
 
 @view
 @external
-def maxWithdraw(owner: address) -> uint256:
+def maxWithdraw(
+    owner: address,
+    max_loss: uint256 = 0,
+    strategies: DynArray[address, MAX_QUEUE] = []
+) -> uint256:
     """
     @notice Get the maximum amount of assets that can be withdrawn.
+    @dev Complies to normal 4626 interface and takes custom params.
     @param owner The address that owns the shares.
+    @param max_loss Custom max_loss if any.
+    @param strategies Custom strategies queue if any.
     @return The maximum amount of assets that can be withdrawn.
     """
-    return self._max_withdraw(owner)
+    return self._max_withdraw(owner, max_loss, strategies)
 
 @view
 @external
-def maxRedeem(owner: address) -> uint256:
+def maxRedeem(
+    owner: address,
+    max_loss: uint256 = MAX_BPS,
+    strategies: DynArray[address, MAX_QUEUE] = []
+) -> uint256:
     """
     @notice Get the maximum amount of shares that can be redeemed.
+    @dev Complies to normal 4626 interface and takes custom params.
     @param owner The address that owns the shares.
+    @param max_loss Custom max_loss if any.
+    @param strategies Custom strategies queue if any.
     @return The maximum amount of shares that can be redeemed.
     """
-    return self._max_redeem(owner)
+    return min(
+        # Max witdraw is rounding so we check against the full balance.
+        self._convert_to_shares(self._max_withdraw(owner, max_loss, strategies), Rounding.ROUND_UP),
+        self.balance_of[owner]
+    )
 
 @view
 @external

tests/unit/vault/test_erc4626.py

@@ -1,6 +1,6 @@
 import ape
 import pytest
-from utils.constants import ROLES
+from utils.constants import ROLES, DAY
 
 
 def test_total_assets(asset, fish, fish_amount, create_vault, user_deposit):
@@ -156,6 +156,106 @@ def test_max_withdraw__with_balance_less_than_or_equal_to_total_idle__returns_ba
     assert vault.maxWithdraw(fish.address) == assets
 
 
+def test_max_withdraw__with_custom_params(
+    asset,
+    fish,
+    fish_amount,
+    gov,
+    create_vault,
+    create_strategy,
+    add_debt_to_strategy,
+    add_strategy_to_vault,
+    user_deposit,
+):
+    vault = create_vault(asset)
+    shares = fish_amount
+    assets = shares
+    strategy = create_strategy(vault)
+    strategy_deposit = assets // 2
+    total_idle = assets - strategy_deposit
+
+    vault.set_role(
+        gov.address,
+        ROLES.ADD_STRATEGY_MANAGER | ROLES.DEBT_MANAGER | ROLES.MAX_DEBT_MANAGER,
+        sender=gov,
+    )
+    user_deposit(fish, vault, asset, assets)
+    add_strategy_to_vault(gov, strategy, vault)
+    add_debt_to_strategy(gov, strategy, vault, strategy_deposit)
+
+    assert vault.maxWithdraw(fish.address, 22, [strategy]) == assets
+
+
+def test_max_withdraw__with_lossy_strategy(
+    asset,
+    fish,
+    fish_amount,
+    gov,
+    create_vault,
+    create_lossy_strategy,
+    add_debt_to_strategy,
+    add_strategy_to_vault,
+    user_deposit,
+):
+    vault = create_vault(asset)
+    shares = fish_amount
+    assets = shares
+    strategy = create_lossy_strategy(vault)
+    strategy_deposit = assets // 2
+    loss = strategy_deposit // 2
+    total_idle = assets - strategy_deposit
+
+    vault.set_role(
+        gov.address,
+        ROLES.ADD_STRATEGY_MANAGER | ROLES.DEBT_MANAGER | ROLES.MAX_DEBT_MANAGER,
+        sender=gov,
+    )
+    user_deposit(fish, vault, asset, assets)
+    add_strategy_to_vault(gov, strategy, vault)
+    add_debt_to_strategy(gov, strategy, vault, strategy_deposit)
+
+    strategy.setLoss(gov.address, loss, sender=gov)
+
+    # Should not effect the default returned value
+    assert vault.maxWithdraw(fish.address, 10_000) == assets
+
+    # Should return just idle if max_loss == 0.
+    assert vault.maxWithdraw(fish.address) == total_idle
+
+
+def test_max_withdraw__with_locked_strategy(
+    asset,
+    fish,
+    fish_amount,
+    gov,
+    create_vault,
+    create_locked_strategy,
+    add_debt_to_strategy,
+    add_strategy_to_vault,
+    user_deposit,
+):
+    vault = create_vault(asset)
+    shares = fish_amount
+    assets = shares
+    strategy = create_locked_strategy(vault)
+    strategy_deposit = assets // 2
+    locked = strategy_deposit // 2
+    total_idle = assets - strategy_deposit
+
+    vault.set_role(
+        gov.address,
+        ROLES.ADD_STRATEGY_MANAGER | ROLES.DEBT_MANAGER | ROLES.MAX_DEBT_MANAGER,
+        sender=gov,
+    )
+    user_deposit(fish, vault, asset, assets)
+    add_strategy_to_vault(gov, strategy, vault)
+    add_debt_to_strategy(gov, strategy, vault, strategy_deposit)
+
+    strategy.setLockedFunds(locked, DAY, sender=gov)
+
+    assert vault.maxWithdraw(fish.address) == assets - locked
+
+
 def test_preview_redeem(asset, fish, fish_amount, create_vault, user_deposit):
     vault = create_vault(asset)
     shares = fish_amount
@@ -193,7 +293,7 @@ def test_max_redeem__with_balance_greater_than_total_idle__returns_balance(
     add_strategy_to_vault(gov, strategy, vault)
     add_debt_to_strategy(gov, strategy, vault, strategy_deposit)
 
-    assert vault.maxWithdraw(fish.address) == assets
+    assert vault.maxRedeem(fish.address) == assets
 
 
 def test_max_redeem__with_balance_less_than_or_equal_to_total_idle__returns_balance(
@@ -205,4 +305,104 @@ def test_max_redeem__with_balance_less_than_or_equal_to_total_idle__returns_bala
 
     user_deposit(fish, vault, asset, shares)
 
-    assert vault.maxWithdraw(fish.address) == assets
+    assert vault.maxRedeem(fish.address) == assets
+
+
+def test_max_redeem__with_custom_params(
+    asset,
+    fish,
+    fish_amount,
+    gov,
+    create_vault,
+    create_strategy,
+    add_debt_to_strategy,
+    add_strategy_to_vault,
+    user_deposit,
+):
+    vault = create_vault(asset)
+    shares = fish_amount
+    assets = shares
+    strategy = create_strategy(vault)
+    strategy_deposit = assets // 2
+    total_idle = assets - strategy_deposit
+
+    vault.set_role(
+        gov.address,
+        ROLES.ADD_STRATEGY_MANAGER | ROLES.DEBT_MANAGER | ROLES.MAX_DEBT_MANAGER,
+        sender=gov,
+    )
+    user_deposit(fish, vault, asset, assets)
+    add_strategy_to_vault(gov, strategy, vault)
+    add_debt_to_strategy(gov, strategy, vault, strategy_deposit)
+
+    assert vault.maxRedeem(fish.address, 0, [strategy]) == assets
+
+
+def test_max_redeem__with_lossy_strategy(
+    asset,
+    fish,
+    fish_amount,
+    gov,
+    create_vault,
+    create_lossy_strategy,
+    add_debt_to_strategy,
+    add_strategy_to_vault,
+    user_deposit,
+):
+    vault = create_vault(asset)
+    shares = fish_amount
+    assets = shares
+    strategy = create_lossy_strategy(vault)
+    strategy_deposit = assets // 2
+    loss = strategy_deposit // 2
+    total_idle = assets - strategy_deposit
+
+    vault.set_role(
+        gov.address,
+        ROLES.ADD_STRATEGY_MANAGER | ROLES.DEBT_MANAGER | ROLES.MAX_DEBT_MANAGER,
+        sender=gov,
+    )
+    user_deposit(fish, vault, asset, assets)
+    add_strategy_to_vault(gov, strategy, vault)
+    add_debt_to_strategy(gov, strategy, vault, strategy_deposit)
+
+    strategy.setLoss(gov.address, loss, sender=gov)
+
+    # Should not effect the default returned value
+    assert vault.maxRedeem(fish.address) == assets
+
+    # Should return just idle if max_loss == 0.
+    assert vault.maxRedeem(fish.address, 0) == total_idle
+
+
+def test_max_redeem__with_locked_strategy(
+    asset,
+    fish,
+    fish_amount,
+    gov,
+    create_vault,
+    create_locked_strategy,
+    add_debt_to_strategy,
+    add_strategy_to_vault,
+    user_deposit,
+):
+    vault = create_vault(asset)
+    shares = fish_amount
+    assets = shares
+    strategy = create_locked_strategy(vault)
+    strategy_deposit = assets // 2
+    locked = strategy_deposit // 2
+    total_idle = assets - strategy_deposit
+
+    vault.set_role(
+        gov.address,
+        ROLES.ADD_STRATEGY_MANAGER | ROLES.DEBT_MANAGER | ROLES.MAX_DEBT_MANAGER,
+        sender=gov,
+    )
+    user_deposit(fish, vault, asset, assets)
+    add_strategy_to_vault(gov, strategy, vault)
+    add_debt_to_strategy(gov, strategy, vault, strategy_deposit)
+
+    strategy.setLockedFunds(locked, DAY, sender=gov)
+
+    assert vault.maxRedeem(fish.address) == assets - locked