Agent vaults

// SPDX-License-Identifier: MIT
pragma solidity ^0.8.28;

interface IAgentIdentity {
    function isTrusted(address account) external view returns (bool);
}

// SPDX-License-Identifier: MIT
pragma solidity ^0.8.28;

import {Ownable} from "@openzeppelin/contracts/access/Ownable.sol";
import {Pausable} from "@openzeppelin/contracts/utils/Pausable.sol";
import {ReentrancyGuard} from "@openzeppelin/contracts/utils/ReentrancyGuard.sol";
import {IERC20} from "@openzeppelin/contracts/token/ERC20/IERC20.sol";
import {SafeERC20} from "@openzeppelin/contracts/token/ERC20/utils/SafeERC20.sol";
import {IAgentIdentity} from "./IAgentIdentity.sol";

/// @notice Holds USDT for an AI agent. The agent key is the operator and can only call `pay`.
/// The owner sets the rules, can pause and can withdraw at any time. Not audited.
contract AgentVault is Ownable, Pausable, ReentrancyGuard {
    using SafeERC20 for IERC20;

    IERC20 public immutable usdt;
    address public operator;
    /// @notice Most USDT (6 decimals) the operator can send per UTC day.
    uint256 public dailyLimit;
    /// @notice Optional. When set, `pay` also requires `isTrusted(to)`.
    IAgentIdentity public identity;

    mapping(address => bool) public isRecipientAllowed;

    uint256 private _spentDay;
    uint256 private _spent;

    event Paid(address indexed to, uint256 amount);
    event OperatorChanged(address indexed operator);
    event DailyLimitChanged(uint256 limit);
    event RecipientAllowed(address indexed account, bool allowed);
    event IdentityChanged(address indexed identity);

    error NotOperator();
    error ZeroAddress();
    error ZeroAmount();
    error RecipientNotAllowed(address to);
    error RecipientNotTrusted(address to);
    error DailyLimitExceeded(uint256 amount, uint256 remaining);

    modifier onlyOperator() {
        if (msg.sender != operator) revert NotOperator();
        _;
    }

    constructor(address owner_, address operator_, address usdt_, uint256 dailyLimit_) Ownable(owner_) {
        if (operator_ == address(0) || usdt_ == address(0)) revert ZeroAddress();
        operator = operator_;
        usdt = IERC20(usdt_);
        dailyLimit = dailyLimit_;
    }

    // Operator

    function pay(address to, uint256 amount) external nonReentrant onlyOperator whenNotPaused {
        if (!isRecipientAllowed[to]) revert RecipientNotAllowed(to);
        if (address(identity) != address(0) && !identity.isTrusted(to)) revert RecipientNotTrusted(to);
        _spend(amount);
        usdt.safeTransfer(to, amount);
        emit Paid(to, amount);
    }

    // Owner

    function setOperator(address newOperator) external onlyOwner {
        if (newOperator == address(0)) revert ZeroAddress();
        operator = newOperator;
        emit OperatorChanged(newOperator);
    }

    function setDailyLimit(uint256 newLimit) external onlyOwner {
        dailyLimit = newLimit;
        emit DailyLimitChanged(newLimit);
    }

    function setRecipientAllowed(address account, bool allowed) external onlyOwner {
        if (account == address(0)) revert ZeroAddress();
        isRecipientAllowed[account] = allowed;
        emit RecipientAllowed(account, allowed);
    }

    /// @notice Set to address(0) to turn the identity check off.
    function setIdentity(address newIdentity) external onlyOwner {
        identity = IAgentIdentity(newIdentity);
        emit IdentityChanged(newIdentity);
    }

    function pause() external onlyOwner {
        _pause();
    }

    function unpause() external onlyOwner {
        _unpause();
    }

    /// @notice Sends tokens back to the owner. Works while paused.
    function withdraw(address token, uint256 amount) external onlyOwner {
        IERC20(token).safeTransfer(owner(), amount);
    }

    // Views

    /// @notice The UTC day number: whole days since 1 January 1970.
    function today() public view returns (uint256) {
        return block.timestamp / 1 days;
    }

    function spentToday() public view returns (uint256) {
        return _spentDay == today() ? _spent : 0;
    }

    function remainingToday() public view returns (uint256) {
        uint256 spent = spentToday();
        return spent >= dailyLimit ? 0 : dailyLimit - spent;
    }

    // Internal

    function _spend(uint256 amount) private {
        if (amount == 0) revert ZeroAmount();
        uint256 day = today();
        if (_spentDay != day) {
            _spentDay = day;
            _spent = 0;
        }
        uint256 remaining = _spent >= dailyLimit ? 0 : dailyLimit - _spent;
        if (amount > remaining) revert DailyLimitExceeded(amount, remaining);
        _spent += amount;
    }
}

[PASS] test_IdentityHook() (gas: 476568)
[PASS] test_LimitResetsAtUtcMidnight() (gas: 226594)
[PASS] test_PauseStopsAgentButNotOwner() (gas: 137768)
[PASS] test_PayWithinLimit() (gas: 127517)
[PASS] test_RevertWhen_CallerIsNotOperator() (gas: 43154)
[PASS] test_RevertWhen_OverDailyLimit() (gas: 155510)
[PASS] test_RevertWhen_RecipientNotAllowed() (gas: 48503)

AGENT_PRIVATE_KEY=0xTheAgentPrivateKey
VAULT=YOUR_VAULT_ADDRESS
RECIPIENT=YOUR_RECIPIENT_ADDRESS

import { BaseError, ContractFunctionRevertedError, createPublicClient, createWalletClient, formatUnits, getAddress, http, parseAbi, parseUnits } from "viem";
import { privateKeyToAccount } from "viem/accounts";
import { botChainTestnet } from "@uzolabs/sdk/chains";

const VAULT = getAddress(process.env.VAULT as string);
const to = getAddress(process.env.RECIPIENT as string);
const amount = parseUnits(process.env.AMOUNT ?? "1", 6); // USDT has 6 decimals

const vaultAbi = parseAbi([
  "function pay(address to, uint256 amount)",
  "function remainingToday() view returns (uint256)",
  "function isRecipientAllowed(address account) view returns (bool)",
  "function paused() view returns (bool)",
  "error NotOperator()",
  "error ZeroAmount()",
  "error RecipientNotAllowed(address to)",
  "error RecipientNotTrusted(address to)",
  "error DailyLimitExceeded(uint256 amount, uint256 remaining)",
  "error EnforcedPause()",
]);

// The agent's key. It is the vault's operator, not its owner.
const account = privateKeyToAccount(process.env.AGENT_PRIVATE_KEY as `0x${string}`);
const publicClient = createPublicClient({ chain: botChainTestnet, transport: http() });
const walletClient = createWalletClient({ account, chain: botChainTestnet, transport: http() });

const [remaining, allowed, paused] = await Promise.all([
  publicClient.readContract({ address: VAULT, abi: vaultAbi, functionName: "remainingToday" }),
  publicClient.readContract({ address: VAULT, abi: vaultAbi, functionName: "isRecipientAllowed", args: [to] }),
  publicClient.readContract({ address: VAULT, abi: vaultAbi, functionName: "paused" }),
]);
console.log(`Left today: ${formatUnits(remaining, 6)} USDT. Recipient allowed: ${allowed}. Paused: ${paused}`);

try {
  const { request } = await publicClient.simulateContract({ account, address: VAULT, abi: vaultAbi, functionName: "pay", args: [to, amount] });
  const hash = await walletClient.writeContract(request);
  const receipt = await publicClient.waitForTransactionReceipt({ hash });
  console.log(`Paid ${formatUnits(amount, 6)} USDT, ${receipt.status}: ${botChainTestnet.blockExplorers.default.url}/tx/${hash}`);
} catch (error) {
  // The vault refused. Show its reason, for example DailyLimitExceeded.
  const revert = error instanceof BaseError ? error.walk((e) => e instanceof ContractFunctionRevertedError) : null;
  if (revert instanceof ContractFunctionRevertedError) {
    console.log(`Vault refused: ${revert.data?.errorName}`, revert.data?.args ?? []);
  } else {
    throw error;
  }
}