Execute multiple contract calls atomically with optional sweep-on-failure

contracts/src/AgentExecutor.sol

@@ -5,6 +5,7 @@ pragma solidity 0.8.33;
 import {IERC20} from "./interfaces/IERC20.sol";
 import {SafeTokenTransfer, SafeNativeTransfer} from "./utils/SafeTransfer.sol";
 import {Call} from "./utils/Call.sol";
+import {CallContractParams} from "./v2/Types.sol";
 
 /// @title Code which will run within an `Agent` using `delegatecall`.
 /// @dev This is a singleton contract, meaning that all agents will execute the same code.
@@ -29,4 +30,33 @@ contract AgentExecutor {
             revert();
         }
     }
+
+    // Call multiple contracts with Ether values; reverts on the first failure
+    function callContracts(CallContractParams[] calldata params) external {
+        uint256 len = params.length;
+        for (uint256 i; i < len; ++i) {
+            bool success = Call.safeCall(params[i].target, params[i].data, params[i].value);
+            if (!success) {
+                revert();
+            }
+        }
+    }
+
+    // Sweep remaining assets when specified
+    function sweep(address recipient, address[] calldata tokens) external {
+        for (uint256 i; i < tokens.length; ++i) {
+            address token = tokens[i];
+            if (token == address(0)) {
+                uint256 balance = address(this).balance;
+                if (balance > 0) {
+                    payable(recipient).safeNativeTransfer(balance);
+                }
+            } else {
+                uint256 balance = IERC20(token).balanceOf(address(this));
+                if (balance > 0) {
+                    IERC20(token).safeTransfer(recipient, balance);
+                }
+            }
+        }
+    }
 }

contracts/src/Gateway.sol

@@ -25,7 +25,7 @@ import {
     IGatewayV1,
     IGatewayV2
 } from "./Types.sol";
-import {Network} from "./v2/Types.sol";
+import {Network, CallContractsParams} from "./v2/Types.sol";
 import {Upgrade} from "./Upgrade.sol";
 import {IInitializable} from "./interfaces/IInitializable.sol";
 import {IUpgradable} from "./interfaces/IUpgradable.sol";
@@ -500,6 +500,12 @@ contract Gateway is IGatewayBase, IGatewayV1, IGatewayV2, IInitializable, IUpgra
             catch {
                 success = false;
                 emit IGatewayV2.CommandFailed(nonce, i);
+                if (command.kind == CommandKind.CallContracts) {
+                    try this.trySweepOnFailure(origin, command.payload) {}
+                    catch {
+                        emit IGatewayV2.SweepAfterCallContractsFailed(nonce, i);
+                    }
+                }
             }
         }
         return (false, success);
@@ -523,11 +529,25 @@ contract Gateway is IGatewayBase, IGatewayV1, IGatewayV2, IInitializable, IUpgra
             HandlersV2.mintForeignToken(command.payload);
         } else if (command.kind == CommandKind.CallContract) {
             HandlersV2.callContract(origin, AGENT_EXECUTOR, command.payload);
+        } else if (command.kind == CommandKind.CallContracts) {
+            HandlersV2.callContracts(origin, AGENT_EXECUTOR, command.payload);
         } else {
             revert IGatewayV2.InvalidCommand();
         }
     }
 
+    /// @dev Decode CallContractsParams and run sweep when configured; used in CallContracts catch block.
+    ///      Caller wraps in try-catch so malformed payloads or sweep failures don't revert dispatch.
+    function trySweepOnFailure(bytes32 origin, bytes calldata payload)
+        external
+        onlySelf
+    {
+        CallContractsParams memory params = abi.decode(payload, (CallContractsParams));
+        if (params.sweepRecipient != address(0) && params.tokensToSweep.length > 0) {
+            HandlersV2.sweepAfterCallContracts(origin, AGENT_EXECUTOR, payload);
+        }
+    }
+
     /**
      * Upgrades
      */

contracts/src/v2/Handlers.sol

@@ -19,7 +19,8 @@ import {
     UnlockNativeTokenParams,
     RegisterForeignTokenParams,
     MintForeignTokenParams,
-    CallContractParams
+    CallContractParams,
+    CallContractsParams
 } from "./Types.sol";
 
 library HandlersV2 {
@@ -71,4 +72,21 @@ library HandlersV2 {
             abi.encodeCall(AgentExecutor.callContract, (params.target, params.data, params.value));
         Functions.invokeOnAgent(agent, executor, call);
     }
+
+    function callContracts(bytes32 origin, address executor, bytes calldata data) external {
+        CallContractsParams memory params = abi.decode(data, (CallContractsParams));
+        address agent = Functions.ensureAgent(origin);
+        bytes memory call = abi.encodeCall(AgentExecutor.callContracts, params.calls);
+        Functions.invokeOnAgent(agent, executor, call);
+    }
+
+    function sweepAfterCallContracts(bytes32 origin, address executor, bytes calldata data)
+        external
+    {
+        CallContractsParams memory params = abi.decode(data, (CallContractsParams));
+        address agent = Functions.ensureAgent(origin);
+        bytes memory call =
+            abi.encodeCall(AgentExecutor.sweep, (params.sweepRecipient, params.tokensToSweep));
+        Functions.invokeOnAgent(agent, executor, call);
+    }
 }

contracts/src/v2/IGateway.sol

@@ -38,6 +38,9 @@ interface IGatewayV2 {
     /// Emitted when a command at `index` within an inbound message identified by `nonce` fails to execute
     event CommandFailed(uint64 indexed nonce, uint256 index);
 
+    /// Emitted when sweep-after-CallContracts fails (e.g. malformed payload or sweep revert)
+    event SweepAfterCallContractsFailed(uint64 indexed nonce, uint256 index);
+
     /// Emitted when an outbound message has been accepted for delivery to a Polkadot parachain
     event OutboundMessageAccepted(uint64 nonce, Payload payload);
 

contracts/src/v2/Types.sol

@@ -36,6 +36,8 @@ library CommandKind {
     uint8 constant MintForeignToken = 4;
     // Call an arbitrary solidity contract
     uint8 constant CallContract = 5;
+    // Call multiple arbitrary solidity contracts
+    uint8 constant CallContracts = 6;
 }
 
 // Payload for outbound messages destined for Polkadot
@@ -185,6 +187,16 @@ struct CallContractParams {
     uint256 value;
 }
 
+// Payload for CallContracts command. Reverts on first call failure; optional sweep when calls fail.
+struct CallContractsParams {
+    // Sub-calls to execute (reverts on first failure)
+    CallContractParams[] calls;
+    // Recipient for sweep when calls fail; address(0) = no sweep
+    address sweepRecipient;
+    // Tokens to sweep full balance to sweepRecipient; address(0) = sweep ETH
+    address[] tokensToSweep;
+}
+
 enum Network {
     Polkadot
 }