Add Aave quote-push collateral supply and toggle hooks.

Extends AaveQuotePushFlashReceiver with before/after swap collateral steps, env-driven run scripts, forkproof parity, and scoped forge tests for supply/toggle callback ordering.

Co-authored-by: Cursor <cursoragent@cursor.com>

contracts/flash/AaveQuotePushFlashReceiver.sol

@@ -26,6 +26,10 @@ interface IAavePoolLike {
         bytes calldata params,
         uint16 referralCode
     ) external;
+
+    function supply(address asset, uint256 amount, address onBehalfOf, uint16 referralCode) external;
+
+    function setUserUseReserveAsCollateral(address asset, bool useAsCollateral) external;
 }
 
 /// @dev Callback for `flashLoan` (multi-asset API).
@@ -84,13 +88,39 @@ interface IAaveAtomicBridgeCoordinator {
 /**
  * @title AaveQuotePushFlashReceiver
  * @notice Aave V3 flash-loan receiver for the quote-push workflow:
- *         flash borrow quote (`flashLoan` single-asset) -> buy PMM base -> unwind base externally -> repay lender, retaining any surplus.
+ *         flash borrow quote (`flashLoan` single-asset) -> optional Aave collateral supply/toggle hooks ->
+ *         buy PMM base -> optional atomic bridge -> unwind base externally -> repay lender, retaining any surplus.
+ * @dev Collateral hooks apply to **this contract's** Aave position (`onBehalfOf` = address(this)). Pre-fund the
+ *      receiver with collateral tokens before `supplyBeforeSwap` steps, or supply from flash proceeds in later txs.
  */
 contract AaveQuotePushFlashReceiver is IAaveFlashLoanSimpleReceiver, IAaveFlashLoanReceiver, Ownable {
     using SafeERC20 for IERC20;
 
     address public immutable pool;
 
+    enum CollateralHook {
+        BeforeSwap,
+        AfterSwap,
+        BeforeUnwind
+    }
+
+    struct CollateralSupplyStep {
+        address asset;
+        uint256 amount;
+    }
+
+    struct CollateralToggleStep {
+        address asset;
+        bool useAsCollateral;
+    }
+
+    struct CollateralParams {
+        CollateralSupplyStep[] supplyBeforeSwap;
+        CollateralToggleStep[] toggleBeforeSwap;
+        CollateralToggleStep[] toggleAfterSwap;
+        CollateralToggleStep[] toggleBeforeUnwind;
+    }
+
     struct QuotePushParams {
         address integration;
         address pmmPool;
@@ -100,6 +130,7 @@ contract AaveQuotePushFlashReceiver is IAaveFlashLoanSimpleReceiver, IAaveFlashL
         uint256 minOutUnwind;
         bytes unwindData;
         AtomicBridgeParams atomicBridge;
+        CollateralParams collateral;
     }
 
     struct AtomicBridgeParams {
@@ -141,6 +172,18 @@ contract AaveQuotePushFlashReceiver is IAaveFlashLoanSimpleReceiver, IAaveFlashL
     );
     event TokenSwept(address indexed token, address indexed to, uint256 amount);
     event SurplusSwept(address indexed token, address indexed to, uint256 amount, uint256 reserveRetained);
+    event CollateralSupplied(address indexed asset, uint256 amount);
+    event CollateralToggled(address indexed asset, bool useAsCollateral, CollateralHook hook);
+
+    /// @notice Empty `collateral` block for quote-push scripts that do not use Aave position hooks.
+    function emptyCollateralParams() external pure returns (CollateralParams memory params) {
+        return params;
+    }
+
+    /// @notice Non-zero when this deployment supports `QuotePushParams.collateral` atomic hooks.
+    function collateralHooksVersion() external pure returns (uint256) {
+        return 1;
+    }
 
     constructor(address pool_, address initialOwner) Ownable(initialOwner) {
         if (pool_ == address(0) || initialOwner == address(0)) revert BadParams();
@@ -225,13 +268,20 @@ contract AaveQuotePushFlashReceiver is IAaveFlashLoanSimpleReceiver, IAaveFlashL
         ) revert BadParams();
         if (p.baseToken == asset) revert BadParams();
 
+        _applyCollateralSupplies(p.collateral.supplyBeforeSwap);
+        _applyCollateralToggles(p.collateral.toggleBeforeSwap, CollateralHook.BeforeSwap);
+
         uint256 baseOut = _swapQuoteForBase(asset, amount, p.integration, p.pmmPool, p.minOutPmm);
 
+        _applyCollateralToggles(p.collateral.toggleAfterSwap, CollateralHook.AfterSwap);
+
         uint256 baseBal = IERC20(p.baseToken).balanceOf(address(this));
         if (p.atomicBridge.coordinator != address(0)) {
             _triggerAtomicBridge(p.baseToken, baseBal, p.atomicBridge);
         }
 
+        _applyCollateralToggles(p.collateral.toggleBeforeUnwind, CollateralHook.BeforeUnwind);
+
         uint256 unwindOut = _unwindBaseIntoQuote(p.baseToken, asset, p.externalUnwinder, p.minOutUnwind, p.unwindData);
         uint256 surplus = _approveRepayment(asset, amount + premium);
         emit QuotePushExecuted(asset, p.baseToken, amount, premium, baseOut, unwindOut, surplus);
@@ -258,6 +308,27 @@ contract AaveQuotePushFlashReceiver is IAaveFlashLoanSimpleReceiver, IAaveFlashL
             IAaveExternalUnwinder(externalUnwinder).unwind(baseToken, quoteToken, remainingBase, minOutUnwind, unwindData);
     }
 
+    function _applyCollateralSupplies(CollateralSupplyStep[] memory steps) internal {
+        uint256 n = steps.length;
+        for (uint256 i = 0; i < n; ++i) {
+            CollateralSupplyStep memory step = steps[i];
+            if (step.asset == address(0) || step.amount == 0) continue;
+            IERC20(step.asset).forceApprove(pool, step.amount);
+            IAavePoolLike(pool).supply(step.asset, step.amount, address(this), 0);
+            emit CollateralSupplied(step.asset, step.amount);
+        }
+    }
+
+    function _applyCollateralToggles(CollateralToggleStep[] memory steps, CollateralHook hook) internal {
+        uint256 n = steps.length;
+        for (uint256 i = 0; i < n; ++i) {
+            CollateralToggleStep memory step = steps[i];
+            if (step.asset == address(0)) continue;
+            IAavePoolLike(pool).setUserUseReserveAsCollateral(step.asset, step.useAsCollateral);
+            emit CollateralToggled(step.asset, step.useAsCollateral, hook);
+        }
+    }
+
     function _approveRepayment(address quoteToken, uint256 need) internal returns (uint256 surplus) {
         IERC20 quote = IERC20(quoteToken);
         uint256 quoteBal = quote.balanceOf(address(this));

forkproof/src/AaveQuotePushFlashReceiver.sol

@@ -22,6 +22,10 @@ interface IAavePoolLike {
         bytes calldata params,
         uint16 referralCode
     ) external;
+
+    function supply(address asset, uint256 amount, address onBehalfOf, uint16 referralCode) external;
+
+    function setUserUseReserveAsCollateral(address asset, bool useAsCollateral) external;
 }
 
 interface IAaveFlashLoanReceiver {
@@ -85,6 +89,29 @@ contract AaveQuotePushFlashReceiver is IAaveFlashLoanSimpleReceiver, IAaveFlashL
 
     address public immutable pool;
 
+    enum CollateralHook {
+        BeforeSwap,
+        AfterSwap,
+        BeforeUnwind
+    }
+
+    struct CollateralSupplyStep {
+        address asset;
+        uint256 amount;
+    }
+
+    struct CollateralToggleStep {
+        address asset;
+        bool useAsCollateral;
+    }
+
+    struct CollateralParams {
+        CollateralSupplyStep[] supplyBeforeSwap;
+        CollateralToggleStep[] toggleBeforeSwap;
+        CollateralToggleStep[] toggleAfterSwap;
+        CollateralToggleStep[] toggleBeforeUnwind;
+    }
+
     struct QuotePushParams {
         address integration;
         address pmmPool;
@@ -94,6 +121,7 @@ contract AaveQuotePushFlashReceiver is IAaveFlashLoanSimpleReceiver, IAaveFlashL
         uint256 minOutUnwind;
         bytes unwindData;
         AtomicBridgeParams atomicBridge;
+        CollateralParams collateral;
     }
 
     struct AtomicBridgeParams {
@@ -132,6 +160,16 @@ contract AaveQuotePushFlashReceiver is IAaveFlashLoanSimpleReceiver, IAaveFlashL
         uint256 bridgeAmount,
         uint256 minDestinationAmount
     );
+    event CollateralSupplied(address indexed asset, uint256 amount);
+    event CollateralToggled(address indexed asset, bool useAsCollateral, CollateralHook hook);
+
+    function emptyCollateralParams() external pure returns (CollateralParams memory params) {
+        return params;
+    }
+
+    function collateralHooksVersion() external pure returns (uint256) {
+        return 1;
+    }
 
     constructor(address pool_) {
         pool = pool_;
@@ -189,13 +227,20 @@ contract AaveQuotePushFlashReceiver is IAaveFlashLoanSimpleReceiver, IAaveFlashL
         ) revert BadParams();
         if (p.baseToken == asset) revert BadParams();
 
+        _applyCollateralSupplies(p.collateral.supplyBeforeSwap);
+        _applyCollateralToggles(p.collateral.toggleBeforeSwap, CollateralHook.BeforeSwap);
+
         uint256 baseOut = _swapQuoteForBase(asset, amount, p.integration, p.pmmPool, p.minOutPmm);
 
+        _applyCollateralToggles(p.collateral.toggleAfterSwap, CollateralHook.AfterSwap);
+
         uint256 baseBal = IERC20(p.baseToken).balanceOf(address(this));
         if (p.atomicBridge.coordinator != address(0)) {
             _triggerAtomicBridge(p.baseToken, baseBal, p.atomicBridge);
         }
 
+        _applyCollateralToggles(p.collateral.toggleBeforeUnwind, CollateralHook.BeforeUnwind);
+
         uint256 unwindOut = _unwindBaseIntoQuote(p.baseToken, asset, p.externalUnwinder, p.minOutUnwind, p.unwindData);
         uint256 surplus = _approveRepayment(asset, amount + premium);
         emit QuotePushExecuted(asset, p.baseToken, amount, premium, baseOut, unwindOut, surplus);
@@ -222,6 +267,27 @@ contract AaveQuotePushFlashReceiver is IAaveFlashLoanSimpleReceiver, IAaveFlashL
             IAaveExternalUnwinder(externalUnwinder).unwind(baseToken, quoteToken, remainingBase, minOutUnwind, unwindData);
     }
 
+    function _applyCollateralSupplies(CollateralSupplyStep[] memory steps) internal {
+        uint256 n = steps.length;
+        for (uint256 i = 0; i < n; ++i) {
+            CollateralSupplyStep memory step = steps[i];
+            if (step.asset == address(0) || step.amount == 0) continue;
+            IERC20(step.asset).forceApprove(pool, step.amount);
+            IAavePoolLike(pool).supply(step.asset, step.amount, address(this), 0);
+            emit CollateralSupplied(step.asset, step.amount);
+        }
+    }
+
+    function _applyCollateralToggles(CollateralToggleStep[] memory steps, CollateralHook hook) internal {
+        uint256 n = steps.length;
+        for (uint256 i = 0; i < n; ++i) {
+            CollateralToggleStep memory step = steps[i];
+            if (step.asset == address(0)) continue;
+            IAavePoolLike(pool).setUserUseReserveAsCollateral(step.asset, step.useAsCollateral);
+            emit CollateralToggled(step.asset, step.useAsCollateral, hook);
+        }
+    }
+
     function _approveRepayment(address quoteToken, uint256 need) internal returns (uint256 surplus) {
         IERC20 quote = IERC20(quoteToken);
         uint256 quoteBal = quote.balanceOf(address(this));

forkproof/test/AaveQuotePushFlashReceiverMainnetFork.t.sol

@@ -186,7 +186,8 @@ contract AaveQuotePushFlashReceiverMainnetForkTest is Test {
                 routeId: bytes32(0),
                 settlementMode: bytes32(0),
                 submitCommitment: false
-            })
+            }),
+            collateral: receiver.emptyCollateralParams()
         });
 
         receiver.flashQuotePush(USDC, amount, p);
@@ -239,7 +240,8 @@ contract AaveQuotePushFlashReceiverMainnetForkTest is Test {
                 routeId: atomicCorridorId,
                 settlementMode: bytes32(0),
                 submitCommitment: true
-            })
+            }),
+            collateral: receiver.emptyCollateralParams()
         });
 
         uint256 receiverQuoteBefore = IERC20(USDC).balanceOf(address(receiver));
@@ -297,7 +299,8 @@ contract AaveQuotePushFlashReceiverMainnetForkTest is Test {
                 routeId: atomicCorridorId,
                 settlementMode: bytes32(0),
                 submitCommitment: true
-            })
+            }),
+            collateral: receiver.emptyCollateralParams()
         });
 
         receiver.flashQuotePush(USDC, amount, p);

script/deploy/DeployAaveQuotePushFlashReceiver.s.sol

@@ -13,6 +13,9 @@ import {AaveQuotePushFlashReceiver} from "../../contracts/flash/AaveQuotePushFla
  *   AAVE_POOL_ADDRESS        optional; defaults to Aave V3 mainnet Pool
  *   QUOTE_PUSH_RECEIVER_OWNER optional; defaults to deployer derived from PRIVATE_KEY
  *
+ * Post-deploy: set AAVE_QUOTE_PUSH_RECEIVER_MAINNET; verify collateralHooksVersion() == 1.
+ * Redeploy wrapper: bash scripts/deployment/redeploy-aave-quote-push-receiver-mainnet.sh
+ *
  * Usage:
  *   forge script script/deploy/DeployAaveQuotePushFlashReceiver.s.sol:DeployAaveQuotePushFlashReceiver \
  *     --rpc-url $ETHEREUM_MAINNET_RPC --broadcast -vvvv
@@ -35,5 +38,6 @@ contract DeployAaveQuotePushFlashReceiver is Script {
         vm.stopBroadcast();
 
         console.log("AaveQuotePushFlashReceiver:", address(receiver));
+        console.log("collateralHooksVersion:", receiver.collateralHooksVersion());
     }
 }

script/flash/RunMainnetAaveCwusdcUsdcQuotePushOnce.s.sol

@@ -3,6 +3,7 @@ pragma solidity ^0.8.20;
 
 import {Script, console} from "forge-std/Script.sol";
 import {AaveQuotePushFlashReceiver} from "../../contracts/flash/AaveQuotePushFlashReceiver.sol";
+import {QuotePushCollateralEnv} from "./lib/QuotePushCollateralEnv.sol";
 
 interface IDODOPMMPoolQuote {
     function querySellQuote(address trader, uint256 payQuoteAmount) external view returns (uint256 receiveBaseAmount, uint256 mtFee);
@@ -45,6 +46,9 @@ interface IDODOPMMPoolQuote {
  *                                   set UNWIND_TWO_HOP_POOL_A, UNWIND_TWO_HOP_POOL_B, UNWIND_TWO_HOP_MID_TOKEN,
  *                                   optional UNWIND_MIN_MID_OUT_RAW, then UNWIND_INTERMEDIATE_TOKEN,
  *                                   UNWIND_MIN_INTERMEDIATE_OUT_RAW, UNWIND_V3_PATH_HEX
+ *
+ * Optional collateral hooks (receiver must expose collateralHooksVersion() == 1):
+ *   See script/flash/lib/QuotePushCollateralEnv.sol and docs/runbooks/AAVE_ATOMIC_COLLATERAL_TOGGLE_RUNBOOK.md
  */
 contract RunMainnetAaveCwusdcUsdcQuotePushOnce is Script {
     address internal constant DEFAULT_POOL = 0x69776fc607e9edA8042e320e7e43f54d06c68f0E;
@@ -141,7 +145,8 @@ contract RunMainnetAaveCwusdcUsdcQuotePushOnce is Script {
                 routeId: bytes32(0),
                 settlementMode: bytes32(0),
                 submitCommitment: false
-            })
+            }),
+            collateral: QuotePushCollateralEnv.loadCollateralParams()
         });
 
         console.log("receiver", receiver);

script/flash/RunManagedMainnetAaveCwusdcUsdcQuotePushCycle.s.sol

@@ -4,6 +4,7 @@ pragma solidity ^0.8.20;
 import {Script, console} from "forge-std/Script.sol";
 import {AaveQuotePushFlashReceiver} from "../../contracts/flash/AaveQuotePushFlashReceiver.sol";
 import {QuotePushTreasuryManager} from "../../contracts/flash/QuotePushTreasuryManager.sol";
+import {QuotePushCollateralEnv} from "./lib/QuotePushCollateralEnv.sol";
 
 interface IDODOPMMPoolQuoteManaged {
     function querySellQuote(address trader, uint256 payQuoteAmount) external view returns (uint256 receiveBaseAmount, uint256 mtFee);
@@ -164,7 +165,8 @@ contract RunManagedMainnetAaveCwusdcUsdcQuotePushCycle is Script {
                 routeId: bytes32(0),
                 settlementMode: bytes32(0),
                 submitCommitment: false
-            })
+            }),
+            collateral: QuotePushCollateralEnv.loadCollateralParams()
         });
 
         console.log("minOutPmm", minOutPmm);

script/flash/RunManagedMainnetAaveCwusdtUsdtQuotePushCycle.s.sol

@@ -4,6 +4,7 @@ pragma solidity ^0.8.20;
 import {Script, console} from "forge-std/Script.sol";
 import {AaveQuotePushFlashReceiver} from "../../contracts/flash/AaveQuotePushFlashReceiver.sol";
 import {QuotePushTreasuryManager} from "../../contracts/flash/QuotePushTreasuryManager.sol";
+import {QuotePushCollateralEnv} from "./lib/QuotePushCollateralEnv.sol";
 
 interface IDODOPMMPoolQuoteManagedUsdt {
     function querySellQuote(address trader, uint256 payQuoteAmount) external view returns (uint256 receiveBaseAmount, uint256 mtFee);
@@ -95,7 +96,8 @@ contract RunManagedMainnetAaveCwusdtUsdtQuotePushCycle is Script {
                 routeId: bytes32(0),
                 settlementMode: bytes32(0),
                 submitCommitment: false
-            })
+            }),
+            collateral: QuotePushCollateralEnv.loadCollateralParams()
         });
     }
 }

script/flash/lib/QuotePushCollateralEnv.sol

@@ -0,0 +1,67 @@
+// SPDX-License-Identifier: MIT
+pragma solidity ^0.8.20;
+
+import {Vm} from "forge-std/Vm.sol";
+import {Strings} from "@openzeppelin/contracts/utils/Strings.sol";
+import {AaveQuotePushFlashReceiver} from "../../../contracts/flash/AaveQuotePushFlashReceiver.sol";
+
+/**
+ * @title QuotePushCollateralEnv
+ * @notice Build `QuotePushParams.collateral` from optional env vars for mainnet quote-push scripts.
+ *
+ * Env (all optional — empty when counts are zero / unset):
+ *   QUOTE_PUSH_COLLATERAL_SUPPLY_COUNT
+ *   QUOTE_PUSH_COLLATERAL_SUPPLY_{i}_ASSET
+ *   QUOTE_PUSH_COLLATERAL_SUPPLY_{i}_AMOUNT_RAW
+ *
+ *   QUOTE_PUSH_COLLATERAL_TOGGLE_BEFORE_COUNT
+ *   QUOTE_PUSH_COLLATERAL_TOGGLE_BEFORE_{i}_ASSET
+ *   QUOTE_PUSH_COLLATERAL_TOGGLE_BEFORE_{i}_ENABLE   (1 = true, 0 = false)
+ *
+ *   QUOTE_PUSH_COLLATERAL_TOGGLE_AFTER_COUNT
+ *   QUOTE_PUSH_COLLATERAL_TOGGLE_AFTER_{i}_ASSET
+ *   QUOTE_PUSH_COLLATERAL_TOGGLE_AFTER_{i}_ENABLE
+ *
+ *   QUOTE_PUSH_COLLATERAL_TOGGLE_BEFORE_UNWIND_COUNT
+ *   QUOTE_PUSH_COLLATERAL_TOGGLE_BEFORE_UNWIND_{i}_ASSET
+ *   QUOTE_PUSH_COLLATERAL_TOGGLE_BEFORE_UNWIND_{i}_ENABLE
+ */
+library QuotePushCollateralEnv {
+    Vm private constant VM = Vm(address(uint160(uint256(keccak256("hevm cheat code")))));
+
+    function loadCollateralParams() internal view returns (AaveQuotePushFlashReceiver.CollateralParams memory params) {
+        params.supplyBeforeSwap = _loadSupplies();
+        params.toggleBeforeSwap = _loadToggles("QUOTE_PUSH_COLLATERAL_TOGGLE_BEFORE");
+        params.toggleAfterSwap = _loadToggles("QUOTE_PUSH_COLLATERAL_TOGGLE_AFTER");
+        params.toggleBeforeUnwind = _loadToggles("QUOTE_PUSH_COLLATERAL_TOGGLE_BEFORE_UNWIND");
+    }
+
+    function _loadSupplies()
+        private
+        view
+        returns (AaveQuotePushFlashReceiver.CollateralSupplyStep[] memory steps)
+    {
+        uint256 n = VM.envOr("QUOTE_PUSH_COLLATERAL_SUPPLY_COUNT", uint256(0));
+        steps = new AaveQuotePushFlashReceiver.CollateralSupplyStep[](n);
+        for (uint256 i = 0; i < n; ++i) {
+            string memory prefix = string.concat("QUOTE_PUSH_COLLATERAL_SUPPLY_", Strings.toString(i), "_");
+            steps[i].asset = VM.envAddress(string.concat(prefix, "ASSET"));
+            steps[i].amount = VM.envUint(string.concat(prefix, "AMOUNT_RAW"));
+        }
+    }
+
+    function _loadToggles(string memory groupPrefix)
+        private
+        view
+        returns (AaveQuotePushFlashReceiver.CollateralToggleStep[] memory steps)
+    {
+        string memory countKey = string.concat(groupPrefix, "_COUNT");
+        uint256 n = VM.envOr(countKey, uint256(0));
+        steps = new AaveQuotePushFlashReceiver.CollateralToggleStep[](n);
+        for (uint256 i = 0; i < n; ++i) {
+            string memory prefix = string.concat(groupPrefix, "_", Strings.toString(i), "_");
+            steps[i].asset = VM.envAddress(string.concat(prefix, "ASSET"));
+            steps[i].useAsCollateral = VM.envUint(string.concat(prefix, "ENABLE")) != 0;
+        }
+    }
+}

test/flash/AaveQuotePushCollateralHooks.t.sol

@@ -0,0 +1,247 @@
+// SPDX-License-Identifier: MIT
+pragma solidity ^0.8.20;
+
+import {Test} from "forge-std/Test.sol";
+import {ERC20} from "@openzeppelin/contracts/token/ERC20/ERC20.sol";
+import {IERC20} from "@openzeppelin/contracts/token/ERC20/IERC20.sol";
+import {
+    AaveQuotePushFlashReceiver,
+    IAaveDODOQuotePushSwapExactIn,
+    IAaveExternalUnwinder,
+    IAavePoolLike
+} from "../../contracts/flash/AaveQuotePushFlashReceiver.sol";
+
+contract MockQuoteToken is ERC20 {
+    constructor() ERC20("Mock Quote", "MQ") {}
+
+    function mint(address to, uint256 amount) external {
+        _mint(to, amount);
+    }
+}
+
+contract MockBaseToken is ERC20 {
+    constructor() ERC20("Mock Base", "MB") {}
+
+    function mint(address to, uint256 amount) external {
+        _mint(to, amount);
+    }
+}
+
+contract MockPmmIntegration is IAaveDODOQuotePushSwapExactIn {
+    MockQuoteToken internal immutable quote;
+    MockBaseToken internal immutable base;
+
+    constructor(MockQuoteToken quote_, MockBaseToken base_) {
+        quote = quote_;
+        base = base_;
+    }
+
+    function swapExactIn(address, address tokenIn, uint256 amountIn, uint256 minAmountOut)
+        external
+        override
+        returns (uint256 amountOut)
+    {
+        require(tokenIn == address(quote), "quote only");
+        quote.transferFrom(msg.sender, address(this), amountIn);
+        amountOut = amountIn;
+        require(amountOut >= minAmountOut, "minOut");
+        base.mint(msg.sender, amountOut);
+    }
+}
+
+contract MockUnwinder is IAaveExternalUnwinder {
+    MockQuoteToken internal immutable quote;
+    MockBaseToken internal immutable base;
+
+    constructor(MockQuoteToken quote_, MockBaseToken base_) {
+        quote = quote_;
+        base = base_;
+    }
+
+    function unwind(address tokenIn, address tokenOut, uint256 amountIn, uint256 minAmountOut, bytes calldata)
+        external
+        override
+        returns (uint256 amountOut)
+    {
+        require(tokenIn == address(base), "base only");
+        require(tokenOut == address(quote), "quote only");
+        base.transferFrom(msg.sender, address(this), amountIn);
+        amountOut = amountIn;
+        require(amountOut >= minAmountOut, "min unwind");
+        quote.mint(msg.sender, amountOut);
+    }
+}
+
+contract MockAavePool is IAavePoolLike {
+    struct ToggleCall {
+        address asset;
+        bool useAsCollateral;
+    }
+
+    uint256 public premiumBps = 5;
+    uint256 public toggleCallCount;
+    ToggleCall[] internal _toggleCalls;
+    uint256 public supplyCount;
+
+    function flashLoan(
+        address receiverAddress,
+        address[] calldata assets,
+        uint256[] calldata amounts,
+        uint256[] calldata,
+        address,
+        bytes calldata params,
+        uint16
+    ) external override {
+        uint256 premium = amounts[0] * premiumBps / 10_000;
+        uint256[] memory premiums = new uint256[](1);
+        premiums[0] = premium;
+        IERC20(assets[0]).transfer(receiverAddress, amounts[0]);
+        bool ok = IAaveFlashLoanReceiver(receiverAddress).executeOperation(
+            assets, amounts, premiums, receiverAddress, params
+        );
+        require(ok, "callback failed");
+        IERC20(assets[0]).transferFrom(receiverAddress, address(this), amounts[0] + premium);
+    }
+
+    function flashLoanSimple(address, address, uint256, bytes calldata, uint16) external pure override {
+        revert("use flashLoan");
+    }
+
+    function supply(address asset, uint256 amount, address onBehalfOf, uint16) external override {
+        IERC20(asset).transferFrom(msg.sender, address(this), amount);
+        supplyCount++;
+        onBehalfOf;
+    }
+
+    function setUserUseReserveAsCollateral(address asset, bool useAsCollateral) external override {
+        _toggleCalls.push(ToggleCall({asset: asset, useAsCollateral: useAsCollateral}));
+        toggleCallCount++;
+    }
+
+    function toggleCallAt(uint256 index) external view returns (address asset, bool useAsCollateral) {
+        ToggleCall memory t = _toggleCalls[index];
+        return (t.asset, t.useAsCollateral);
+    }
+}
+
+interface IAaveFlashLoanReceiver {
+    function executeOperation(
+        address[] calldata assets,
+        uint256[] calldata amounts,
+        uint256[] calldata premiums,
+        address initiator,
+        bytes calldata params
+    ) external returns (bool);
+}
+
+contract AaveQuotePushCollateralHooksTest is Test {
+    MockQuoteToken internal quote;
+    MockBaseToken internal base;
+    MockPmmIntegration internal integration;
+    MockUnwinder internal unwinder;
+    MockAavePool internal pool;
+    AaveQuotePushFlashReceiver internal receiver;
+
+    address internal constant WETH = address(0xC02aaA39b223FE8D0A0e5C4F27eAD9083C756Cc2);
+
+    function setUp() public {
+        quote = new MockQuoteToken();
+        base = new MockBaseToken();
+        integration = new MockPmmIntegration(quote, base);
+        unwinder = new MockUnwinder(quote, base);
+        pool = new MockAavePool();
+        receiver = new AaveQuotePushFlashReceiver(address(pool), address(this));
+        quote.mint(address(pool), 10_000_000);
+        quote.mint(address(receiver), 1_000_000);
+    }
+
+    function testCollateralToggleRunsInFlashCallbackBeforeAndAfterSwap() public {
+        AaveQuotePushFlashReceiver.CollateralToggleStep[] memory beforeSwap =
+            new AaveQuotePushFlashReceiver.CollateralToggleStep[](1);
+        beforeSwap[0] = AaveQuotePushFlashReceiver.CollateralToggleStep({asset: WETH, useAsCollateral: false});
+
+        AaveQuotePushFlashReceiver.CollateralToggleStep[] memory afterSwap =
+            new AaveQuotePushFlashReceiver.CollateralToggleStep[](1);
+        afterSwap[0] = AaveQuotePushFlashReceiver.CollateralToggleStep({asset: WETH, useAsCollateral: true});
+
+        AaveQuotePushFlashReceiver.QuotePushParams memory p = AaveQuotePushFlashReceiver.QuotePushParams({
+            integration: address(integration),
+            pmmPool: address(0xBEEF),
+            baseToken: address(base),
+            externalUnwinder: address(unwinder),
+            minOutPmm: 1,
+            minOutUnwind: 1,
+            unwindData: "",
+            atomicBridge: AaveQuotePushFlashReceiver.AtomicBridgeParams({
+                coordinator: address(0),
+                sourceChain: 0,
+                destinationChain: 0,
+                destinationAsset: address(0),
+                bridgeAmount: 0,
+                minDestinationAmount: 0,
+                destinationRecipient: address(0),
+                destinationDeadline: 0,
+                routeId: bytes32(0),
+                settlementMode: bytes32(0),
+                submitCommitment: false
+            }),
+            collateral: AaveQuotePushFlashReceiver.CollateralParams({
+                supplyBeforeSwap: new AaveQuotePushFlashReceiver.CollateralSupplyStep[](0),
+                toggleBeforeSwap: beforeSwap,
+                toggleAfterSwap: afterSwap,
+                toggleBeforeUnwind: new AaveQuotePushFlashReceiver.CollateralToggleStep[](0)
+            })
+        });
+
+        uint256 borrow = 1_000_000;
+        receiver.flashQuotePush(address(quote), borrow, p);
+
+        assertEq(pool.toggleCallCount(), 2, "two toggles");
+        (address asset0, bool use0) = pool.toggleCallAt(0);
+        (address asset1, bool use1) = pool.toggleCallAt(1);
+        assertEq(asset0, WETH);
+        assertFalse(use0);
+        assertEq(asset1, WETH);
+        assertTrue(use1);
+    }
+
+    function testCollateralSupplyBeforeSwapUsesReceiverBalance() public {
+        quote.mint(address(receiver), 500_000);
+
+        AaveQuotePushFlashReceiver.CollateralSupplyStep[] memory supplies =
+            new AaveQuotePushFlashReceiver.CollateralSupplyStep[](1);
+        supplies[0] = AaveQuotePushFlashReceiver.CollateralSupplyStep({asset: address(quote), amount: 100_000});
+
+        AaveQuotePushFlashReceiver.QuotePushParams memory p = AaveQuotePushFlashReceiver.QuotePushParams({
+            integration: address(integration),
+            pmmPool: address(0xBEEF),
+            baseToken: address(base),
+            externalUnwinder: address(unwinder),
+            minOutPmm: 1,
+            minOutUnwind: 1,
+            unwindData: "",
+            atomicBridge: AaveQuotePushFlashReceiver.AtomicBridgeParams({
+                coordinator: address(0),
+                sourceChain: 0,
+                destinationChain: 0,
+                destinationAsset: address(0),
+                bridgeAmount: 0,
+                minDestinationAmount: 0,
+                destinationRecipient: address(0),
+                destinationDeadline: 0,
+                routeId: bytes32(0),
+                settlementMode: bytes32(0),
+                submitCommitment: false
+            }),
+            collateral: AaveQuotePushFlashReceiver.CollateralParams({
+                supplyBeforeSwap: supplies,
+                toggleBeforeSwap: new AaveQuotePushFlashReceiver.CollateralToggleStep[](0),
+                toggleAfterSwap: new AaveQuotePushFlashReceiver.CollateralToggleStep[](0),
+                toggleBeforeUnwind: new AaveQuotePushFlashReceiver.CollateralToggleStep[](0)
+            })
+        });
+
+        receiver.flashQuotePush(address(quote), 1_000_000, p);
+        assertEq(pool.supplyCount(), 1, "supply called once");
+    }
+}

test/flash/AaveQuotePushFlashReceiver.t.sol

@@ -23,6 +23,10 @@ contract AaveQuotePushFlashReceiverTest is Test {
         token.mint(address(receiver), 1_000_000);
     }
 
+    function testCollateralHooksVersionIsOne() public view {
+        assertEq(receiver.collateralHooksVersion(), 1);
+    }
+
     function testOwnerCanSweepQuoteSurplusAndKeepReserve() public {
         uint256 reserveRetained = 250_000;
 

test/flash/AaveQuotePushFlashReceiverMainnetFork.t.sol

@@ -99,7 +99,8 @@ contract AaveQuotePushFlashReceiverMainnetForkTest is Test {
                 routeId: bytes32(0),
                 settlementMode: bytes32(0),
                 submitCommitment: false
-            })
+            }),
+            collateral: receiver.emptyCollateralParams()
         });
 
         receiver.flashQuotePush(USDC, amount, p);
@@ -138,7 +139,8 @@ contract AaveQuotePushFlashReceiverMainnetForkTest is Test {
                 routeId: bytes32(0),
                 settlementMode: bytes32(0),
                 submitCommitment: false
-            })
+            }),
+            collateral: receiver.emptyCollateralParams()
         });
 
         receiver.flashQuotePush(USDC, amount, p);

test/flash/QuotePushTreasuryManager.t.sol

@@ -181,7 +181,8 @@ contract QuotePushTreasuryManagerTest is Test {
                 routeId: bytes32(0),
                 settlementMode: bytes32(0),
                 submitCommitment: false
-            })
+            }),
+            collateral: AaveQuotePushFlashReceiver(address(cycleReceiver)).emptyCollateralParams()
         });
 
         vm.prank(OPERATOR);
@@ -227,7 +228,8 @@ contract QuotePushTreasuryManagerTest is Test {
                 routeId: bytes32(0),
                 settlementMode: bytes32(0),
                 submitCommitment: false
-            })
+            }),
+            collateral: AaveQuotePushFlashReceiver(address(cycleReceiver)).emptyCollateralParams()
         });
 
         cycleManager.transferManagedReceiverOwnership(address(this));