From f803b24655351cc744078c298631b6838b3c19b2 Mon Sep 17 00:00:00 2001 From: Shivamd0608 Date: Wed, 3 Jun 2026 11:28:52 +0530 Subject: [PATCH 1/3] modifyed Policy contracts and added agentwallet , Iagentwallet and its factory contract --- packages/contracts/src/core/AgentWallet.sol | 146 ++++++++++ .../contracts/src/core/AgentWalletFactory.sol | 143 +++++++++ packages/contracts/src/core/PolicyEngine.sol | 273 ++++-------------- .../contracts/src/interfaces/IAgentWallet.sol | 53 ++++ .../src/interfaces/IPolicyEngine.sol | 9 +- .../contracts/src/libraries/PolicyLib.sol | 6 +- packages/contracts/test/PolicyEngine.t.sol | 52 ++++ 7 files changed, 457 insertions(+), 225 deletions(-) diff --git a/packages/contracts/src/core/AgentWallet.sol b/packages/contracts/src/core/AgentWallet.sol index e69de29..75045e4 100644 --- a/packages/contracts/src/core/AgentWallet.sol +++ b/packages/contracts/src/core/AgentWallet.sol @@ -0,0 +1,146 @@ +// SPDX-License-Identifier: MIT +pragma solidity ^0.8.24; + +import "@openzeppelin/contracts/utils/ReentrancyGuard.sol"; +import "@openzeppelin/contracts/token/ERC20/utils/SafeERC20.sol"; +import "@openzeppelin/contracts/token/ERC20/IERC20.sol"; +import "../interfaces/IAgentWallet.sol"; +import "../interfaces/IPolicyEngine.sol"; + +/// @title AgentWallet +/// @notice One smart account per AI agent. +/// Holds USDC. Before every payment, asks PolicyEngine +/// for permission. If PolicyEngine reverts, no money moves. +/// +/// @dev Deployed by AgentWalletFactory via CREATE2. +/// Initialised once — constructor is empty because CREATE2 +/// minimal proxies can't take constructor args. +/// Call initialise() right after deployment. +contract AgentWallet is IAgentWallet, ReentrancyGuard { + using SafeERC20 for IERC20; + + // ─── State ──────────────────────────────────────────────────── + bytes32 private _agentId; + address private _owner; + address private _policyEngine; + address private _usdcToken; + bool private _initialised; + + // ─── Modifiers ──────────────────────────────────────────────── + modifier onlyOwner() { + if (msg.sender != _owner) revert NotOwner(); + _; + } + + modifier initialised() { + if (!_initialised) revert NotInitialised(); + _; + } + + // ─── Initialisation ─────────────────────────────────────────── + + /// @notice Called once by AgentWalletFactory after CREATE2 deploy. + /// Sets the agent's identity and connects it to PolicyEngine. + /// @param agentId_ keccak256(abi.encodePacked(address(this), chainId)) + /// @param owner_ human who controls this agent + /// @param policyEngine_ deployed PolicyEngine address + /// @param usdcToken_ USDC contract address on this chain + function initialise( + bytes32 agentId_, + address owner_, + address policyEngine_, + address usdcToken_ + ) external { + // Can only be called once + if (_initialised) revert NotInitialised(); + if (owner_ == address(0)) revert NotOwner(); + if (policyEngine_ == address(0)) revert PolicyBlocked(); + if (usdcToken_ == address(0)) revert TransferFailed(); + + _agentId = agentId_; + _owner = owner_; + _policyEngine = policyEngine_; + _usdcToken = usdcToken_; + _initialised = true; + + emit AgentIdSet(agentId_); + } + + // ─── Core: execute payment ──────────────────────────────────── + + /// @notice Pay `amount` USDC to `payee`. + /// Step 1: ask PolicyEngine if this is allowed. + /// Step 2: transfer USDC if yes. + /// Step 3: emit receipt. + /// @dev Only owner (or in future: authorised agent runner) calls this. + /// nonReentrant guards against re-entry on the USDC transfer. + function execute( + address payee, + uint128 amount + ) + external + override + onlyOwner + initialised + nonReentrant + returns (bool approvalRequired) + { + if (amount == 0) revert ZeroAmount(); + if (payee == address(0)) revert ZeroPayee(); + + uint128 bal = balance(); + if (bal < amount) revert InsufficientBalance(bal, amount); + + // ── Step 1: policy check ────────────────────────────────── + // enforce() reverts if any rule is violated. + // If it reverts, this whole transaction reverts — no money moves. + try IPolicyEngine(_policyEngine).enforce(_agentId, amount, payee) + returns (bool _approvalRequired) + { + approvalRequired = _approvalRequired; + } catch { + revert PolicyBlocked(); + } + + // ── Step 2: move USDC ───────────────────────────────────── + // SafeERC20 handles tokens that return false instead of reverting + IERC20(_usdcToken).safeTransfer(payee, uint256(amount)); + + // ── Step 3: emit receipt ────────────────────────────────── + emit PaymentExecuted(payee, amount, _agentId, approvalRequired); + } + + // ─── Funding ────────────────────────────────────────────────── + + /// @notice Anyone can deposit USDC to fund this agent. + /// Caller must have approved this contract first. + function deposit(uint128 amount) external override nonReentrant { + if (amount == 0) revert ZeroAmount(); + IERC20(_usdcToken).safeTransferFrom(msg.sender, address(this), uint256(amount)); + emit FundsDeposited(msg.sender, amount); + } + + /// @notice Owner withdraws USDC back to any address. + function withdraw( + address to, + uint128 amount + ) external override onlyOwner nonReentrant { + if (amount == 0) revert ZeroAmount(); + if (to == address(0)) revert ZeroPayee(); + uint128 bal = balance(); + if (bal < amount) revert InsufficientBalance(bal, amount); + IERC20(_usdcToken).safeTransfer(to, uint256(amount)); + emit FundsWithdrawn(to, amount); + } + + // ─── Views ──────────────────────────────────────────────────── + + function agentId() external view override returns (bytes32) { return _agentId; } + function owner() external view override returns (address) { return _owner; } + function policyEngine() external view override returns (address) { return _policyEngine; } + function usdcToken() external view override returns (address) { return _usdcToken; } + + function balance() public view override returns (uint128) { + return uint128(IERC20(_usdcToken).balanceOf(address(this))); + } +} \ No newline at end of file diff --git a/packages/contracts/src/core/AgentWalletFactory.sol b/packages/contracts/src/core/AgentWalletFactory.sol index e69de29..dde80fd 100644 --- a/packages/contracts/src/core/AgentWalletFactory.sol +++ b/packages/contracts/src/core/AgentWalletFactory.sol @@ -0,0 +1,143 @@ +// SPDX-License-Identifier: MIT +pragma solidity ^0.8.24; + +import "@openzeppelin/contracts/access/Ownable.sol"; +import "@openzeppelin/contracts/access/AccessControl.sol"; +import "./AgentWallet.sol"; +import "../interfaces/IPolicyEngine.sol"; + +/// @title AgentWalletFactory +/// @notice Deploys one AgentWallet per agent using CREATE2. +/// CREATE2 gives a deterministic address — you can predict +/// where the wallet will live before it is deployed. +/// +/// After deploying the wallet, the factory: +/// 1. Calls wallet.initialise() to wire up owner + policy engine +/// 2. Grants the wallet ENFORCER_ROLE on PolicyEngine +/// so it is the only address allowed to call enforce() +contract AgentWalletFactory is Ownable { + + // ─── State ──────────────────────────────────────────────────── + address public immutable policyEngine; + address public immutable usdcToken; + + // agentId => wallet address (registry) + mapping(bytes32 => address) public wallets; + + // ─── Events ─────────────────────────────────────────────────── + event AgentCreated( + bytes32 indexed agentId, + address indexed wallet, + address indexed owner + ); + + // ─── Errors ─────────────────────────────────────────────────── + error AgentAlreadyExists(bytes32 agentId); + error ZeroAddress(); + error DeployFailed(); + + // ─── Constructor ────────────────────────────────────────────── + constructor( + address policyEngine_, + address usdcToken_, + address factoryOwner_ + ) Ownable(factoryOwner_) { + if (policyEngine_ == address(0)) revert ZeroAddress(); + if (usdcToken_ == address(0)) revert ZeroAddress(); + if (factoryOwner_ == address(0)) revert ZeroAddress(); + + policyEngine = policyEngine_; + usdcToken = usdcToken_; + } + + // ─── Core: create agent ─────────────────────────────────────── + + /// @notice Deploy a new AgentWallet for `agentOwner`. + /// The agentId is deterministic: keccak256(owner + salt). + /// Anyone can create an agent for any owner address. + /// + /// @param agentOwner human who will control the agent + /// @param salt unique value — use a counter or UUID hash + /// @return wallet address of the deployed AgentWallet + /// @return agentId the agent's unique bytes32 identifier + function createAgent( + address agentOwner, + bytes32 salt + ) external returns (address wallet, bytes32 agentId) { + if (agentOwner == address(0)) revert ZeroAddress(); + + // ── Derive agentId ──────────────────────────────────────── + // Ties agent identity to: owner + salt + this chain + agentId = keccak256( + abi.encodePacked(agentOwner, salt, block.chainid) + ); + + if (wallets[agentId] != address(0)) + revert AgentAlreadyExists(agentId); + + // ── Deploy with CREATE2 ─────────────────────────────────── + // CREATE2 address is determined by: + // keccak256(0xff ++ factory_address ++ salt ++ keccak256(bytecode)) + // This means you can predict the wallet address BEFORE deploying. + bytes32 create2Salt = keccak256(abi.encodePacked(agentId)); + + wallet = _deploy(create2Salt); + if (wallet == address(0)) revert DeployFailed(); + + // ── Initialise the wallet ───────────────────────────────── + AgentWallet(wallet).initialise( + agentId, + agentOwner, + policyEngine, + usdcToken + ); + + // ── Grant wallet ENFORCER_ROLE on PolicyEngine ──────────── + // This is the access control that means ONLY this wallet + // can call policyEngine.enforce(). No other address can. + AccessControl(policyEngine).grantRole( + keccak256("ENFORCER_ROLE"), + wallet + ); + + // ── Register in our local registry ─────────────────────── + wallets[agentId] = wallet; + + emit AgentCreated(agentId, wallet, agentOwner); + } + + // ─── Predict address ───────────────────────────────────────── + + /// @notice Returns the address where an AgentWallet WOULD be deployed + /// for these inputs, without actually deploying it. + /// Useful for pre-funding a wallet before it exists. + function predictAddress( + address agentOwner, + bytes32 salt + ) external view returns (address) { + bytes32 agentId = keccak256( + abi.encodePacked(agentOwner, salt, block.chainid) + ); + bytes32 create2Salt = keccak256(abi.encodePacked(agentId)); + + return address(uint160(uint256(keccak256(abi.encodePacked( + bytes1(0xff), + address(this), + create2Salt, + keccak256(type(AgentWallet).creationCode) + ))))); + } + + /// @notice Returns true if an agent exists for this agentId + function agentExists(bytes32 agentId) external view returns (bool) { + return wallets[agentId] != address(0); + } + + // ─── Internal: CREATE2 deploy ───────────────────────────────── + function _deploy(bytes32 salt) internal returns (address addr) { + bytes memory bytecode = type(AgentWallet).creationCode; + assembly { + addr := create2(0, add(bytecode, 0x20), mload(bytecode), salt) + } + } +} \ No newline at end of file diff --git a/packages/contracts/src/core/PolicyEngine.sol b/packages/contracts/src/core/PolicyEngine.sol index 4530171..be813cd 100644 --- a/packages/contracts/src/core/PolicyEngine.sol +++ b/packages/contracts/src/core/PolicyEngine.sol @@ -2,269 +2,100 @@ pragma solidity ^0.8.24; import "@openzeppelin/contracts/access/AccessControl.sol"; - +import "@openzeppelin/contracts/utils/ReentrancyGuard.sol"; // ← ADDED BACK import "../interfaces/IPolicyEngine.sol"; import "../libraries/PolicyLib.sol"; -/// @title PolicyEngine -/// @notice Enforces spending policies for agent wallets. -/// AgentWallet must call enforce() before every payment. -contract PolicyEngine is IPolicyEngine, AccessControl { +contract PolicyEngine is IPolicyEngine, AccessControl, ReentrancyGuard { using PolicyLib for PolicyLib.Policy; using PolicyLib for PolicyLib.SpendRecord; - bytes32 public constant POLICY_ADMIN_ROLE = - keccak256("POLICY_ADMIN_ROLE"); - - /// @dev agentId => Policy - mapping(bytes32 => PolicyLib.Policy) - private _policies; + bytes32 public constant POLICY_ADMIN_ROLE = keccak256("POLICY_ADMIN_ROLE"); + bytes32 public constant ENFORCER_ROLE = keccak256("ENFORCER_ROLE"); + // AgentWallet gets granted ENFORCER_ROLE by factory on deployment - /// @dev agentId => SpendRecord - mapping(bytes32 => PolicyLib.SpendRecord) - private _spendRecords; + mapping(bytes32 => PolicyLib.Policy) private _policies; + mapping(bytes32 => PolicyLib.SpendRecord) private _spendRecords; constructor(address admin) { + if (admin == address(0)) revert ZeroAdmin(); // ← zero-address check _grantRole(DEFAULT_ADMIN_ROLE, admin); - _grantRole(POLICY_ADMIN_ROLE, admin); + _grantRole(POLICY_ADMIN_ROLE, admin); } - // ============================================================ - // WRITE FUNCTIONS - // ============================================================ - - /// @notice Create or update an agent policy function setPolicy( bytes32 agentId, PolicyLib.Policy calldata policy - ) - external - onlyRole(POLICY_ADMIN_ROLE) - { - // Basic sanity checks - - if ( - policy.maxPerTransaction > - policy.maxPerHour - ) { - revert InvalidPolicy(); - } - - if ( - policy.maxPerHour > - policy.maxPerDay - ) { - revert InvalidPolicy(); - } - + ) external onlyRole(POLICY_ADMIN_ROLE) { + // Validate payee list size + if (policy.allowedPayees.length > PolicyLib.MAX_ALLOWED_PAYEES) + revert TooManyPayees(policy.allowedPayees.length, PolicyLib.MAX_ALLOWED_PAYEES); + + // Validate limit hierarchy + if (policy.maxPerTransaction > policy.maxPerHour) revert InvalidPolicy(); + if (policy.maxPerHour > policy.maxPerDay) revert InvalidPolicy(); + // requireApprovalAbove must be <= maxPerTransaction (no point flagging what's already blocked) + if (policy.requireApprovalAbove > policy.maxPerTransaction) revert InvalidPolicy(); + + bool isUpdate = _policies[agentId].exists; _policies[agentId] = policy; + _policies[agentId].exists = true; // always set exists = true - emit PolicySet( - agentId, - msg.sender - ); - } - - /// @notice Disable a policy - function revokePolicy( - bytes32 agentId - ) - external - onlyRole(POLICY_ADMIN_ROLE) - { - PolicyLib.Policy storage p = - _policies[agentId]; - - if ( - p.maxPerTransaction == 0 && - !p.active - ) { - revert PolicyNotFound(agentId); + if (isUpdate) { + emit PolicyUpdated(agentId, msg.sender); + } else { + emit PolicySet(agentId, msg.sender); } + } - p.active = false; - + function revokePolicy(bytes32 agentId) external onlyRole(POLICY_ADMIN_ROLE) { + if (!_policies[agentId].exists) revert PolicyNotFound(agentId); // ← uses exists + _policies[agentId].active = false; emit PolicyRevoked(agentId); } - /// @notice Enforce policy before payment. - /// Reverts if payment violates rules. function enforce( bytes32 agentId, uint128 amount, address payee - ) - external - returns (bool requiresApproval) - { - PolicyLib.Policy storage p = - _policies[agentId]; - - PolicyLib.SpendRecord storage r = - _spendRecords[agentId]; - - // -------------------------------------------------------- - // 1. Policy must exist - // -------------------------------------------------------- - - if ( - p.maxPerTransaction == 0 && - !p.active - ) { - revert PolicyNotFound(agentId); - } - - // -------------------------------------------------------- - // 2. Policy must be active - // -------------------------------------------------------- - - if (!p.active) { - revert PolicyInactive(agentId); - } + ) external nonReentrant onlyRole(ENFORCER_ROLE) returns (bool requiresApproval) { + // ← ENFORCER_ROLE: only AgentWallet can call this + PolicyLib.Policy storage p = _policies[agentId]; + PolicyLib.SpendRecord storage r = _spendRecords[agentId]; - // -------------------------------------------------------- - // 3. Policy must not be expired - // -------------------------------------------------------- + if (!p.exists) revert PolicyNotFound(agentId); // ← clean exists check + if (!p.active) revert PolicyInactive(agentId); + if (p.isExpired()) revert PolicyExpired(agentId); + if (!p.isPayeeAllowed(payee)) revert PayeeNotAllowed(agentId, payee); + if (amount > p.maxPerTransaction) + revert ExceedsPerTxLimit(agentId, amount, p.maxPerTransaction); - if (p.isExpired()) { - revert PolicyExpired(agentId); - } - - // -------------------------------------------------------- - // 4. Payee must be allowed - // -------------------------------------------------------- - - if ( - !p.isPayeeAllowed(payee) - ) { - revert PayeeNotAllowed( - agentId, - payee - ); - } - - // -------------------------------------------------------- - // 5. Per transaction limit - // -------------------------------------------------------- + if (r.isNewHour()) { r.hourlyAmount = 0; r.hourBucket = PolicyLib.currentHourBucket(); } + if (r.isNewDay()) { r.dailyAmount = 0; r.dayBucket = PolicyLib.currentDayBucket(); } - if ( - amount > - p.maxPerTransaction - ) { - revert ExceedsPerTxLimit( - agentId, - amount, - p.maxPerTransaction - ); - } - - // -------------------------------------------------------- - // 6. Reset spend windows if needed - // -------------------------------------------------------- - - if (r.isNewHour()) { - r.hourlyAmount = 0; - r.hourBucket = - PolicyLib.currentHourBucket(); - } - - if (r.isNewDay()) { - r.dailyAmount = 0; - r.dayBucket = - PolicyLib.currentDayBucket(); - } + uint128 newHourly = r.hourlyAmount + amount; + uint128 newDaily = r.dailyAmount + amount; - // -------------------------------------------------------- - // 7. Hourly limit - // -------------------------------------------------------- - - uint128 newHourly = - r.hourlyAmount + amount; - - if ( - newHourly > - p.maxPerHour - ) { - revert ExceedsHourlyLimit( - agentId, - newHourly, - p.maxPerHour - ); - } - - // -------------------------------------------------------- - // 8. Daily limit - // -------------------------------------------------------- - - uint128 newDaily = - r.dailyAmount + amount; - - if ( - newDaily > - p.maxPerDay - ) { - revert ExceedsDailyLimit( - agentId, - newDaily, - p.maxPerDay - ); - } - - // -------------------------------------------------------- - // 9. Commit spend - // -------------------------------------------------------- + if (newHourly > p.maxPerHour) revert ExceedsHourlyLimit(agentId, newHourly, p.maxPerHour); + if (newDaily > p.maxPerDay) revert ExceedsDailyLimit(agentId, newDaily, p.maxPerDay); r.hourlyAmount = newHourly; - r.dailyAmount = newDaily; + r.dailyAmount = newDaily; - // -------------------------------------------------------- - // 10. Approval threshold - // -------------------------------------------------------- - - requiresApproval = - amount >= - p.requireApprovalAbove; + requiresApproval = amount >= p.requireApprovalAbove; if (requiresApproval) { - emit ApprovalRequired( - agentId, - amount, - payee - ); + emit ApprovalRequired(agentId, amount, payee); } else { - emit PaymentEnforced( - agentId, - amount, - payee - ); + emit PaymentEnforced(agentId, amount, payee); } } - // ============================================================ - // READ FUNCTIONS - // ============================================================ - - function getPolicy( - bytes32 agentId - ) - external - view - returns ( - PolicyLib.Policy memory - ) - { + function getPolicy(bytes32 agentId) external view returns (PolicyLib.Policy memory) { return _policies[agentId]; } - function getSpendRecord( - bytes32 agentId - ) - external - view - returns ( - PolicyLib.SpendRecord memory - ) - { + function getSpendRecord(bytes32 agentId) external view returns (PolicyLib.SpendRecord memory) { return _spendRecords[agentId]; } } \ No newline at end of file diff --git a/packages/contracts/src/interfaces/IAgentWallet.sol b/packages/contracts/src/interfaces/IAgentWallet.sol index e69de29..7ca4fd6 100644 --- a/packages/contracts/src/interfaces/IAgentWallet.sol +++ b/packages/contracts/src/interfaces/IAgentWallet.sol @@ -0,0 +1,53 @@ +// SPDX-License-Identifier: MIT +pragma solidity ^0.8.24; + +/// @title IAgentWallet +/// @notice Interface for a single agent's smart account. +/// Holds USDC, enforces policy before every payment. +interface IAgentWallet { + + // ─── Events ─────────────────────────────────────────────────── + event PaymentExecuted( + address indexed payee, + uint128 amount, + bytes32 indexed agentId, + bool approvalRequired + ); + event FundsDeposited(address indexed from, uint128 amount); + event FundsWithdrawn(address indexed to, uint128 amount); + event AgentIdSet(bytes32 indexed agentId); + + // ─── Errors ─────────────────────────────────────────────────── + error NotOwner(); + error ZeroAmount(); + error ZeroPayee(); + error PolicyBlocked(); // enforce() reverted upstream + error InsufficientBalance(uint128 available, uint128 requested); + error TransferFailed(); + error NotInitialised(); + + // ─── Core functions ─────────────────────────────────────────── + + /// @notice Execute a USDC payment to payee. + /// Calls PolicyEngine.enforce() first. + /// Reverts if policy blocks it. + function execute( + address payee, + uint128 amount + ) external returns (bool approvalRequired); + + /// @notice Deposit USDC into this wallet. + /// Anyone can fund an agent, only owner withdraws. + function deposit(uint128 amount) external; + + /// @notice Owner withdraws USDC from this wallet. + function withdraw(address to, uint128 amount) external; + + // ─── View functions ─────────────────────────────────────────── + + function agentId() external view returns (bytes32); + function owner() external view returns (address); + function policyEngine() external view returns (address); + function usdcToken() external view returns (address); + function balance() external view returns (uint128); +} \ No newline at end of file diff --git a/packages/contracts/src/interfaces/IPolicyEngine.sol b/packages/contracts/src/interfaces/IPolicyEngine.sol index a29c9a9..1b51268 100644 --- a/packages/contracts/src/interfaces/IPolicyEngine.sol +++ b/packages/contracts/src/interfaces/IPolicyEngine.sol @@ -7,10 +7,11 @@ interface IPolicyEngine { // ─── Events ─────────────────────────────────────────────────── event PolicySet(bytes32 indexed agentId, address indexed owner); + event PolicyUpdated(bytes32 indexed agentId, address indexed owner); // ← distinguish update vs create event PolicyRevoked(bytes32 indexed agentId); - event PaymentEnforced(bytes32 indexed agentId, uint128 amount, address payee); - event ApprovalRequired(bytes32 indexed agentId, uint128 amount, address payee); - event PolicyViolation(bytes32 indexed agentId, string reason); + event PaymentEnforced(bytes32 indexed agentId, uint128 amount, address indexed payee); + event ApprovalRequired(bytes32 indexed agentId, uint128 amount, address indexed payee); + // removed PolicyViolation(string) — covered by custom errors below // ─── Errors ─────────────────────────────────────────────────── error PolicyNotFound(bytes32 agentId); @@ -21,6 +22,8 @@ interface IPolicyEngine { error ExceedsHourlyLimit(bytes32 agentId, uint128 attempted, uint128 limit); error PayeeNotAllowed(bytes32 agentId, address payee); error InvalidPolicy(); + error TooManyPayees(uint256 provided, uint256 max); // ← new + error ZeroAdmin(); // ← new // ─── Functions ──────────────────────────────────────────────── function setPolicy(bytes32 agentId, PolicyLib.Policy calldata policy) external; diff --git a/packages/contracts/src/libraries/PolicyLib.sol b/packages/contracts/src/libraries/PolicyLib.sol index 57ea1e9..2e96c12 100644 --- a/packages/contracts/src/libraries/PolicyLib.sol +++ b/packages/contracts/src/libraries/PolicyLib.sol @@ -11,6 +11,8 @@ library PolicyLib { // This is the "spending rule" that a human owner sets for their agent. // All amounts are in USDC base units (6 decimals). // So 1 USDC = 1_000_000, 0.01 USDC = 10_000 + uint256 internal constant MAX_ALLOWED_PAYEES = 20; + struct Policy { uint128 maxPerTransaction; // max a single payment can be @@ -19,6 +21,7 @@ library PolicyLib { uint128 requireApprovalAbove; // payments >= this need human sign-off uint48 expiresAt; // unix timestamp, 0 = never expires bool active; // false = policy is paused/revoked + bool exists; // to distinguish "no policy" vs "policy with 0 limits" address[] allowedPayees; // empty array = any payee is allowed } @@ -44,11 +47,12 @@ library PolicyLib { /// @notice Returns true if a given payee is allowed by this policy. /// If allowedPayees is empty, all payees are allowed. - function isPayeeAllowed( + function isPayeeAllowed( Policy storage p, address payee ) internal view returns (bool) { if (p.allowedPayees.length == 0) return true; + // Safe — capped at MAX_ALLOWED_PAYEES in setPolicy for (uint256 i = 0; i < p.allowedPayees.length; i++) { if (p.allowedPayees[i] == payee) return true; } diff --git a/packages/contracts/test/PolicyEngine.t.sol b/packages/contracts/test/PolicyEngine.t.sol index 9590efd..b2fc7a3 100644 --- a/packages/contracts/test/PolicyEngine.t.sol +++ b/packages/contracts/test/PolicyEngine.t.sol @@ -33,6 +33,7 @@ contract PolicyEngineTest is Test { requireApprovalAbove: 8_000_000, // 8 USDC needs approval expiresAt: 0, // never expires active: true, + exists: true, allowedPayees: payees }); } @@ -52,6 +53,22 @@ contract PolicyEngineTest is Test { vm.expectRevert(); engine.setPolicy(agentId, _basicPolicy()); } + function test_setPolicy_revertsInvalidPolicy() public { + PolicyLib.Policy memory p = _basicPolicy(); + + // Invalid: + // maxPerTransaction > maxPerHour + p.maxPerTransaction = 60_000_000; + p.maxPerHour = 50_000_000; + + vm.prank(owner); + + vm.expectRevert( + IPolicyEngine.InvalidPolicy.selector + ); + + engine.setPolicy(agentId, p); +} // ─── enforce: happy path ────────────────────────────────────── function test_enforce_smallPayment_passes() public { @@ -87,6 +104,19 @@ contract PolicyEngineTest is Test { engine.enforce(agentId, 11_000_000, payee); } +function test_enforce_perTxLimit_exactBoundaryWorks() public { + vm.prank(owner); + engine.setPolicy(agentId, _basicPolicy()); + + bool needsApproval = + engine.enforce( + agentId, + 10_000_000, + payee + ); + + assertTrue(needsApproval); +} // ─── enforce: daily limit ───────────────────────────────────── function test_enforce_revertsWhenExceedsDailyLimit() public { vm.prank(owner); @@ -127,6 +157,28 @@ contract PolicyEngineTest is Test { assertFalse(ok); // no approval needed, just checking it doesn't revert } +function test_enforce_dailyLimit_exactBoundaryWorks() public { + vm.prank(owner); + engine.setPolicy(agentId, _basicPolicy()); + + // Exactly 100 USDC + + for (uint256 i = 0; i < 10; i++) { + engine.enforce( + agentId, + 10_000_000, + payee + ); + } + + PolicyLib.SpendRecord memory r = + engine.getSpendRecord(agentId); + + assertEq( + r.dailyAmount, + 100_000_000 + ); +} // ─── enforce: hourly limit ──────────────────────────────────── function test_enforce_revertsWhenExceedsHourlyLimit() public { vm.prank(owner); From 7e958426c9bcd75a75386fe8e08049a17da1a750 Mon Sep 17 00:00:00 2001 From: Shivamd0608 Date: Thu, 4 Jun 2026 01:12:24 +0530 Subject: [PATCH 2/3] added agentwallet.t.sol test and modifyed current test contract --- packages/contracts/test/AgentWallet.t.sol | 401 ++++++++++ packages/contracts/test/PolicyEngine.t.sol | 854 ++++++++++++++++----- 2 files changed, 1076 insertions(+), 179 deletions(-) diff --git a/packages/contracts/test/AgentWallet.t.sol b/packages/contracts/test/AgentWallet.t.sol index e69de29..657073a 100644 --- a/packages/contracts/test/AgentWallet.t.sol +++ b/packages/contracts/test/AgentWallet.t.sol @@ -0,0 +1,401 @@ +// SPDX-License-Identifier: MIT +pragma solidity ^0.8.24; + +import "forge-std/Test.sol"; +import "../src/core/AgentWallet.sol"; +import "../src/interfaces/IAgentWallet.sol"; + +// ──────────────────────────────────────────────────────────────── +// MOCK CONTRACTS +// These live in the test file so the test is fully self-contained. +// ──────────────────────────────────────────────────────────────── + +/// @dev Minimal ERC-20 for USDC. Mint freely in tests. +contract MockUSDC { + string public name = "USD Coin"; + string public symbol = "USDC"; + uint8 public decimals = 6; + + mapping(address => uint256) public balanceOf; + mapping(address => mapping(address => uint256)) public allowance; + + function mint(address to, uint256 amount) external { + balanceOf[to] += amount; + } + + function transfer(address to, uint256 amount) external returns (bool) { + require(balanceOf[msg.sender] >= amount, "insufficient"); + balanceOf[msg.sender] -= amount; + balanceOf[to] += amount; + return true; + } + + function transferFrom(address from, address to, uint256 amount) + external returns (bool) + { + require(balanceOf[from] >= amount, "insufficient"); + require(allowance[from][msg.sender] >= amount, "allowance"); + balanceOf[from] -= amount; + balanceOf[to] += amount; + allowance[from][msg.sender] -= amount; + return true; + } + + function approve(address spender, uint256 amount) external returns (bool) { + allowance[msg.sender][spender] = amount; + return true; + } +} + +/// @dev Mock PolicyEngine — controllable from tests. +/// Toggle shouldBlock to simulate policy violations. +/// Toggle shouldRequireApproval to simulate approval threshold. +contract MockPolicyEngine { + bool public shouldBlock = false; + bool public shouldRequireApproval = false; + + // Track calls so we can assert enforce() was actually called + uint256 public enforceCallCount = 0; + bytes32 public lastAgentId; + uint128 public lastAmount; + address public lastPayee; + + // Simulates IPolicyEngine.enforce() signature + function enforce( + bytes32 agentId, + uint128 amount, + address payee + ) external returns (bool requiresApproval) { + enforceCallCount++; + lastAgentId = agentId; + lastAmount = amount; + lastPayee = payee; + + if (shouldBlock) revert("PolicyEngine: blocked"); + + return shouldRequireApproval; + } + + // Test helpers to control behaviour + function blockNextCall() external { shouldBlock = true; } + function unblock() external { shouldBlock = false; } + function setRequiresApproval(bool v) external { shouldRequireApproval = v; } +} + +// ──────────────────────────────────────────────────────────────── +// TEST CONTRACT +// ──────────────────────────────────────────────────────────────── + +contract AgentWalletTest is Test { + + AgentWallet public wallet; + MockUSDC public usdc; + MockPolicyEngine public policyEngine; + + address public owner = makeAddr("owner"); + address public payee = makeAddr("payee"); + address public randos = makeAddr("randos"); + + bytes32 public agentId; + + // 10 USDC in base units + uint128 constant TEN_USDC = 10_000_000; + uint128 constant FIVE_USDC = 5_000_000; + uint128 constant HUNDRED_USDC = 100_000_000; + + // ── Setup ───────────────────────────────────────────────────── + + function setUp() public { + // 1. Deploy mocks + usdc = new MockUSDC(); + policyEngine = new MockPolicyEngine(); + + // 2. Deploy wallet + wallet = new AgentWallet(); + + // 3. Derive agentId the same way factory would + agentId = keccak256( + abi.encodePacked(owner, bytes32("salt1"), block.chainid) + ); + + // 4. Initialise wallet (factory does this after CREATE2 deploy) + wallet.initialise( + agentId, + owner, + address(policyEngine), + address(usdc) + ); + + // 5. Fund the wallet with 100 USDC + usdc.mint(address(wallet), HUNDRED_USDC); + } + + // ── initialise() ───────────────────────────────────────────── + + function test_initialise_setsFieldsCorrectly() public view { + assertEq(wallet.agentId(), agentId); + assertEq(wallet.owner(), owner); + assertEq(wallet.policyEngine(), address(policyEngine)); + assertEq(wallet.usdcToken(), address(usdc)); + } + + function test_initialise_revertsIfCalledTwice() public { + // Already initialised in setUp — calling again must revert + vm.expectRevert(); + wallet.initialise(agentId, owner, address(policyEngine), address(usdc)); + } + + function test_initialise_revertsIfZeroOwner() public { + AgentWallet fresh = new AgentWallet(); + vm.expectRevert(); + fresh.initialise(agentId, address(0), address(policyEngine), address(usdc)); + } + + // ── balance() ──────────────────────────────────────────────── + + function test_balance_reflectsUSDCBalance() public view { + assertEq(wallet.balance(), HUNDRED_USDC); + } + + // ── deposit() ──────────────────────────────────────────────── + + function test_deposit_increasesBalance() public { + usdc.mint(randos, TEN_USDC); + + vm.startPrank(randos); + usdc.approve(address(wallet), TEN_USDC); + wallet.deposit(TEN_USDC); + vm.stopPrank(); + + // 100 (setUp) + 10 = 110 USDC + assertEq(wallet.balance(), HUNDRED_USDC + TEN_USDC); + } + + function test_deposit_emitsEvent() public { + usdc.mint(randos, TEN_USDC); + vm.startPrank(randos); + usdc.approve(address(wallet), TEN_USDC); + + vm.expectEmit(true, false, false, true); + emit IAgentWallet.FundsDeposited(randos, TEN_USDC); + + wallet.deposit(TEN_USDC); + vm.stopPrank(); + } + + function test_deposit_revertsOnZeroAmount() public { + vm.expectRevert(IAgentWallet.ZeroAmount.selector); + wallet.deposit(0); + } + + // ── execute() ──────────────────────────────────────────────── + + function test_execute_transfersUSDCToPayee() public { + vm.prank(owner); + wallet.execute(payee, TEN_USDC); + + // Payee received the USDC + assertEq(usdc.balanceOf(payee), TEN_USDC); + // Wallet balance decreased + assertEq(wallet.balance(), HUNDRED_USDC - TEN_USDC); + } + + function test_execute_callsPolicyEnforce() public { + vm.prank(owner); + wallet.execute(payee, TEN_USDC); + + // enforce() must have been called exactly once + assertEq(policyEngine.enforceCallCount(), 1); + + // With the correct arguments + assertEq(policyEngine.lastAgentId(), agentId); + assertEq(policyEngine.lastAmount(), TEN_USDC); + assertEq(policyEngine.lastPayee(), payee); + } + + function test_execute_emitsPaymentExecuted() public { + vm.expectEmit(true, true, false, true); + emit IAgentWallet.PaymentExecuted(payee, TEN_USDC, agentId, false); + + vm.prank(owner); + wallet.execute(payee, TEN_USDC); + } + + function test_execute_returnsFalseWhenNoApprovalNeeded() public { + policyEngine.setRequiresApproval(false); + + vm.prank(owner); + bool needsApproval = wallet.execute(payee, TEN_USDC); + + assertFalse(needsApproval); + } + + function test_execute_returnsTrueWhenApprovalNeeded() public { + policyEngine.setRequiresApproval(true); + + vm.prank(owner); + bool needsApproval = wallet.execute(payee, TEN_USDC); + + assertTrue(needsApproval); + } + + // ── execute() — policy blocks ───────────────────────────────── + + function test_execute_revertsWhenPolicyBlocks() public { + policyEngine.blockNextCall(); + + vm.prank(owner); + // PolicyBlocked is the wallet-level error wrapping the engine revert + vm.expectRevert(IAgentWallet.PolicyBlocked.selector); + wallet.execute(payee, TEN_USDC); + } + + function test_execute_doesNotTransferWhenPolicyBlocks() public { + policyEngine.blockNextCall(); + + vm.prank(owner); + try wallet.execute(payee, TEN_USDC) {} catch {} + + // Payee received NOTHING — atomic revert + assertEq(usdc.balanceOf(payee), 0); + // Wallet balance unchanged + assertEq(wallet.balance(), HUNDRED_USDC); + } + + // ── execute() — access control ──────────────────────────────── + + function test_execute_revertsIfCallerIsNotOwner() public { + vm.prank(randos); + vm.expectRevert(IAgentWallet.NotOwner.selector); + wallet.execute(payee, TEN_USDC); + } + + // ── execute() — input validation ────────────────────────────── + + function test_execute_revertsOnZeroAmount() public { + vm.prank(owner); + vm.expectRevert(IAgentWallet.ZeroAmount.selector); + wallet.execute(payee, 0); + } + + function test_execute_revertsOnZeroPayee() public { + vm.prank(owner); + vm.expectRevert(IAgentWallet.ZeroPayee.selector); + wallet.execute(address(0), TEN_USDC); + } + + function test_execute_revertsWhenInsufficientBalance() public { + // Try to send more than the wallet holds + uint128 tooMuch = HUNDRED_USDC + TEN_USDC; + + vm.prank(owner); + vm.expectRevert( + abi.encodeWithSelector( + IAgentWallet.InsufficientBalance.selector, + HUNDRED_USDC, // available + tooMuch // requested + ) + ); + wallet.execute(payee, tooMuch); + } + + // ── withdraw() ──────────────────────────────────────────────── + + function test_withdraw_movesUSDCToRecipient() public { + vm.prank(owner); + wallet.withdraw(owner, FIVE_USDC); + + assertEq(usdc.balanceOf(owner), FIVE_USDC); + assertEq(wallet.balance(), HUNDRED_USDC - FIVE_USDC); + } + + function test_withdraw_emitsEvent() public { + vm.expectEmit(true, false, false, true); + emit IAgentWallet.FundsWithdrawn(owner, FIVE_USDC); + + vm.prank(owner); + wallet.withdraw(owner, FIVE_USDC); + } + + function test_withdraw_revertsIfCallerIsNotOwner() public { + vm.prank(randos); + vm.expectRevert(IAgentWallet.NotOwner.selector); + wallet.withdraw(randos, FIVE_USDC); + } + + function test_withdraw_revertsOnZeroAmount() public { + vm.prank(owner); + vm.expectRevert(IAgentWallet.ZeroAmount.selector); + wallet.withdraw(owner, 0); + } + + function test_withdraw_revertsOnZeroAddress() public { + vm.prank(owner); + vm.expectRevert(IAgentWallet.ZeroPayee.selector); + wallet.withdraw(address(0), FIVE_USDC); + } + + function test_withdraw_revertsWhenInsufficientBalance() public { + uint128 tooMuch = HUNDRED_USDC + TEN_USDC; + + vm.prank(owner); + vm.expectRevert( + abi.encodeWithSelector( + IAgentWallet.InsufficientBalance.selector, + HUNDRED_USDC, + tooMuch + ) + ); + wallet.withdraw(owner, tooMuch); + } + + // ── Full flow ───────────────────────────────────────────────── + + function test_fullFlow_depositThenExecute() public { + // 1. Randos funds the wallet + usdc.mint(randos, TEN_USDC); + vm.startPrank(randos); + usdc.approve(address(wallet), TEN_USDC); + wallet.deposit(TEN_USDC); + vm.stopPrank(); + + assertEq(wallet.balance(), HUNDRED_USDC + TEN_USDC); + + // 2. Owner executes payment + vm.prank(owner); + wallet.execute(payee, TEN_USDC); + + assertEq(usdc.balanceOf(payee), TEN_USDC); + assertEq(wallet.balance(), HUNDRED_USDC); // back to original + } + + function test_fullFlow_multiplePayments_accumulateCorrectly() public { + // Three payments of 5 USDC each = 15 USDC total + vm.startPrank(owner); + wallet.execute(payee, FIVE_USDC); + wallet.execute(payee, FIVE_USDC); + wallet.execute(payee, FIVE_USDC); + vm.stopPrank(); + + assertEq(usdc.balanceOf(payee), FIVE_USDC * 3); + assertEq(wallet.balance(), HUNDRED_USDC - (FIVE_USDC * 3)); + + // enforce() called 3 times + assertEq(policyEngine.enforceCallCount(), 3); + } + + // ── Fuzz: any valid amount ──────────────────────────────────── + + /// @dev Fuzz test: any amount between 1 and wallet balance should succeed + /// as long as policyEngine doesn't block. + function testFuzz_execute_anyValidAmount(uint128 amount) public { + // Bound to realistic range: 1 to wallet's balance + amount = uint128(bound(amount, 1, HUNDRED_USDC)); + + vm.prank(owner); + wallet.execute(payee, amount); + + assertEq(usdc.balanceOf(payee), amount); + assertEq(wallet.balance(), HUNDRED_USDC - amount); + } +} \ No newline at end of file diff --git a/packages/contracts/test/PolicyEngine.t.sol b/packages/contracts/test/PolicyEngine.t.sol index b2fc7a3..a966a19 100644 --- a/packages/contracts/test/PolicyEngine.t.sol +++ b/packages/contracts/test/PolicyEngine.t.sol @@ -3,284 +3,780 @@ pragma solidity ^0.8.24; import "forge-std/Test.sol"; import "../src/core/PolicyEngine.sol"; +import "../src/interfaces/IPolicyEngine.sol"; import "../src/libraries/PolicyLib.sol"; +/// @title PolicyEngineTest +/// @notice Full test suite for PolicyEngine.sol +/// Tests are grouped by function, then by scenario. +/// Run: forge test --match-contract PolicyEngineTest -vvv contract PolicyEngineTest is Test { + // ── Contracts ───────────────────────────────────────────────── PolicyEngine public engine; - // test accounts - address public owner = makeAddr("owner"); - address public agent = makeAddr("agent"); - address public payee = makeAddr("payee"); - address public randos = makeAddr("randos"); + // ── Actors ──────────────────────────────────────────────────── + address public admin = makeAddr("admin"); // deployed engine, has POLICY_ADMIN_ROLE + address public alice = makeAddr("alice"); // agent owner + address public payee = makeAddr("payee"); // receives USDC + address public payee2 = makeAddr("payee2"); // second payee for whitelist tests + address public randos = makeAddr("randos"); // has no roles — attacker / stranger + // ── Agent IDs ───────────────────────────────────────────────── bytes32 public agentId; + bytes32 public agentId2; + + // ── USDC amounts (6 decimals) ───────────────────────────────── + uint128 constant ONE_USDC = 1_000_000; + uint128 constant FIVE_USDC = 5_000_000; + uint128 constant TEN_USDC = 10_000_000; + uint128 constant FIFTY_USDC = 50_000_000; + uint128 constant HUNDRED_USDC = 100_000_000; + + // ───────────────────────────────────────────────────────────── + // SETUP + // ───────────────────────────────────────────────────────────── - // ─── Setup ──────────────────────────────────────────────────── function setUp() public { - engine = new PolicyEngine(owner); - agentId = keccak256(abi.encodePacked(agent, uint256(84532))); + engine = new PolicyEngine(admin); + agentId = keccak256(abi.encodePacked(alice, uint256(1))); + agentId2 = keccak256(abi.encodePacked(alice, uint256(2))); } - // helper: build a basic policy - function _basicPolicy() internal pure returns (PolicyLib.Policy memory) { + // ───────────────────────────────────────────────────────────── + // HELPERS + // ───────────────────────────────────────────────────────────── + + /// @dev Build a standard valid policy. Edit fields per test. + function _policy() internal pure returns (PolicyLib.Policy memory) { address[] memory payees = new address[](0); // allow all - return PolicyLib.Policy({ - maxPerTransaction: 10_000_000, // 10 USDC - maxPerDay: 100_000_000, // 100 USDC - maxPerHour: 50_000_000, // 50 USDC - requireApprovalAbove: 8_000_000, // 8 USDC needs approval - expiresAt: 0, // never expires - active: true, - exists: true, - allowedPayees: payees - }); - } - - // ─── setPolicy ──────────────────────────────────────────────── - function test_setPolicy_works() public { - vm.prank(owner); - engine.setPolicy(agentId, _basicPolicy()); + PolicyLib.Policy memory p; + p.maxPerTransaction = TEN_USDC; // 10 USDC per tx + p.maxPerDay = HUNDRED_USDC; // 100 USDC per day + p.maxPerHour = FIFTY_USDC; // 50 USDC per hour + p.requireApprovalAbove = FIVE_USDC; // >= 5 USDC needs approval + p.expiresAt = 0; // never expires + p.active = true; + p.allowedPayees = payees; + return p; + } + + /// @dev Set a policy from admin with no boilerplate. + function _setPolicy(bytes32 id, PolicyLib.Policy memory p) internal { + vm.prank(admin); + engine.setPolicy(id, p); + } + + /// @dev Call enforce() directly (no access control in your current contract). + function _enforce( + bytes32 id, + uint128 amount, + address to + ) internal returns (bool) { + return engine.enforce(id, amount, to); + } + + // ───────────────────────────────────────────────────────────── + // CONSTRUCTOR + // ───────────────────────────────────────────────────────────── + + function test_constructor_adminHasPolicyAdminRole() public view { + assertTrue( + engine.hasRole(engine.POLICY_ADMIN_ROLE(), admin) + ); + } + + function test_constructor_adminHasDefaultAdminRole() public view { + assertTrue( + engine.hasRole(engine.DEFAULT_ADMIN_ROLE(), admin) + ); + } + + function test_constructor_randosHasNoRole() public view { + assertFalse(engine.hasRole(engine.POLICY_ADMIN_ROLE(), randos)); + assertFalse(engine.hasRole(engine.DEFAULT_ADMIN_ROLE(), randos)); + } + + // ───────────────────────────────────────────────────────────── + // setPolicy — access control + // ───────────────────────────────────────────────────────────── + + function test_setPolicy_adminCanSet() public { + vm.prank(admin); + engine.setPolicy(agentId, _policy()); PolicyLib.Policy memory p = engine.getPolicy(agentId); - assertEq(p.maxPerTransaction, 10_000_000); - assertEq(p.active, true); + assertEq(p.maxPerTransaction, TEN_USDC); + assertTrue(p.active); } - function test_setPolicy_revertsIfNotOwner() public { + function test_setPolicy_revertsIfCallerHasNoRole() public { vm.prank(randos); + vm.expectRevert(); // AccessControl revert + engine.setPolicy(agentId, _policy()); + } + + function test_setPolicy_aliceCannotSetWithoutRole() public { + vm.prank(alice); vm.expectRevert(); - engine.setPolicy(agentId, _basicPolicy()); + engine.setPolicy(agentId, _policy()); } - function test_setPolicy_revertsInvalidPolicy() public { - PolicyLib.Policy memory p = _basicPolicy(); - // Invalid: - // maxPerTransaction > maxPerHour - p.maxPerTransaction = 60_000_000; - p.maxPerHour = 50_000_000; + // ───────────────────────────────────────────────────────────── + // setPolicy — validation + // ───────────────────────────────────────────────────────────── + + function test_setPolicy_revertsWhenPerTxExceedsPerHour() public { + PolicyLib.Policy memory p = _policy(); + p.maxPerTransaction = HUNDRED_USDC; // bigger than maxPerHour (50) + p.maxPerHour = FIFTY_USDC; + + vm.prank(admin); + vm.expectRevert(IPolicyEngine.InvalidPolicy.selector); + engine.setPolicy(agentId, p); + } - vm.prank(owner); + function test_setPolicy_revertsWhenPerHourExceedsPerDay() public { + PolicyLib.Policy memory p = _policy(); + p.maxPerHour = HUNDRED_USDC + ONE_USDC; // bigger than maxPerDay (100) - vm.expectRevert( - IPolicyEngine.InvalidPolicy.selector - ); + vm.prank(admin); + vm.expectRevert(IPolicyEngine.InvalidPolicy.selector); + engine.setPolicy(agentId, p); + } - engine.setPolicy(agentId, p); -} + function test_setPolicy_allowsEqualLimits() public { + // maxPerTransaction == maxPerHour == maxPerDay is valid (strict limit) + PolicyLib.Policy memory p = _policy(); + p.maxPerTransaction = TEN_USDC; + p.maxPerHour = TEN_USDC; + p.maxPerDay = TEN_USDC; - // ─── enforce: happy path ────────────────────────────────────── - function test_enforce_smallPayment_passes() public { - vm.prank(owner); - engine.setPolicy(agentId, _basicPolicy()); + vm.prank(admin); + engine.setPolicy(agentId, p); - // 5 USDC — under all limits, no approval needed - bool needsApproval = engine.enforce(agentId, 5_000_000, payee); - assertFalse(needsApproval); + PolicyLib.Policy memory stored = engine.getPolicy(agentId); + assertEq(stored.maxPerTransaction, TEN_USDC); } - function test_enforce_aboveApprovalThreshold_flagged() public { - vm.prank(owner); - engine.setPolicy(agentId, _basicPolicy()); + function test_setPolicy_emitsPolicySetEvent() public { + vm.expectEmit(true, true, false, false); + emit IPolicyEngine.PolicySet(agentId, admin); - // 9 USDC — under per-tx limit but above approval threshold - bool needsApproval = engine.enforce(agentId, 9_000_000, payee); - assertTrue(needsApproval); + vm.prank(admin); + engine.setPolicy(agentId, _policy()); } - // ─── enforce: per-tx limit ──────────────────────────────────── - function test_enforce_revertsWhenExceedsPerTxLimit() public { - vm.prank(owner); - engine.setPolicy(agentId, _basicPolicy()); + function test_setPolicy_canOverwriteExistingPolicy() public { + // First set + _setPolicy(agentId, _policy()); + + // Overwrite with lower limit + PolicyLib.Policy memory p2 = _policy(); + p2.maxPerTransaction = FIVE_USDC; + + vm.prank(admin); + engine.setPolicy(agentId, p2); - // 11 USDC — over 10 USDC per-tx limit + PolicyLib.Policy memory stored = engine.getPolicy(agentId); + assertEq(stored.maxPerTransaction, FIVE_USDC); + } + + function test_setPolicy_withPayeeWhitelist() public { + address[] memory allowed = new address[](2); + allowed[0] = payee; + allowed[1] = payee2; + + PolicyLib.Policy memory p = _policy(); + p.allowedPayees = allowed; + + vm.prank(admin); + engine.setPolicy(agentId, p); + + PolicyLib.Policy memory stored = engine.getPolicy(agentId); + assertEq(stored.allowedPayees.length, 2); + assertEq(stored.allowedPayees[0], payee); + assertEq(stored.allowedPayees[1], payee2); + } + + function test_setPolicy_withExpiryTimestamp() public { + PolicyLib.Policy memory p = _policy(); + p.expiresAt = uint48(block.timestamp + 7 days); + + _setPolicy(agentId, p); + + PolicyLib.Policy memory stored = engine.getPolicy(agentId); + assertEq(stored.expiresAt, uint48(block.timestamp + 7 days)); + } + + // ───────────────────────────────────────────────────────────── + // revokePolicy + // ───────────────────────────────────────────────────────────── + + function test_revokePolicy_setsActiveToFalse() public { + _setPolicy(agentId, _policy()); + + vm.prank(admin); + engine.revokePolicy(agentId); + + PolicyLib.Policy memory p = engine.getPolicy(agentId); + assertFalse(p.active); + } + + function test_revokePolicy_emitsEvent() public { + _setPolicy(agentId, _policy()); + + vm.expectEmit(true, false, false, false); + emit IPolicyEngine.PolicyRevoked(agentId); + + vm.prank(admin); + engine.revokePolicy(agentId); + } + + function test_revokePolicy_revertsIfNoPolicySet() public { + // agentId has never had a policy set + bytes32 unknown = keccak256("ghost"); + + vm.prank(admin); vm.expectRevert( - abi.encodeWithSelector( - IPolicyEngine.ExceedsPerTxLimit.selector, - agentId, uint128(11_000_000), uint128(10_000_000) - ) + abi.encodeWithSelector(IPolicyEngine.PolicyNotFound.selector, unknown) ); - engine.enforce(agentId, 11_000_000, payee); + engine.revokePolicy(unknown); } -function test_enforce_perTxLimit_exactBoundaryWorks() public { - vm.prank(owner); - engine.setPolicy(agentId, _basicPolicy()); + function test_revokePolicy_revertsIfCallerHasNoRole() public { + _setPolicy(agentId, _policy()); - bool needsApproval = - engine.enforce( - agentId, - 10_000_000, - payee + vm.prank(randos); + vm.expectRevert(); + engine.revokePolicy(agentId); + } + + // ───────────────────────────────────────────────────────────── + // enforce — policy existence + state + // ───────────────────────────────────────────────────────────── + + function test_enforce_revertsWhenNoPolicySet() public { + bytes32 unknown = keccak256("nobody"); + + vm.expectRevert( + abi.encodeWithSelector(IPolicyEngine.PolicyNotFound.selector, unknown) ); + _enforce(unknown, ONE_USDC, payee); + } - assertTrue(needsApproval); -} - // ─── enforce: daily limit ───────────────────────────────────── - function test_enforce_revertsWhenExceedsDailyLimit() public { - vm.prank(owner); - engine.setPolicy(agentId, _basicPolicy()); + function test_enforce_revertsWhenPolicyInactive() public { + _setPolicy(agentId, _policy()); - // spend 90 USDC across 9 payments (fine) - for (uint256 i = 0; i < 9; i++) { - engine.enforce(agentId, 10_000_000, payee); - } + vm.prank(admin); + engine.revokePolicy(agentId); + + vm.expectRevert( + abi.encodeWithSelector(IPolicyEngine.PolicyInactive.selector, agentId) + ); + _enforce(agentId, ONE_USDC, payee); + } + + function test_enforce_revertsWhenPolicyExpired() public { + PolicyLib.Policy memory p = _policy(); + p.expiresAt = uint48(block.timestamp + 1 hours); + _setPolicy(agentId, p); - // 10th payment would push to 100 USDC — exactly at limit, still fine - engine.enforce(agentId, 10_000_000, payee); + // Jump past expiry + vm.warp(block.timestamp + 2 hours); + + vm.expectRevert( + abi.encodeWithSelector(IPolicyEngine.PolicyExpired.selector, agentId) + ); + _enforce(agentId, ONE_USDC, payee); + } + + function test_enforce_doesNotRevertAtExactExpiry() public { + // expiresAt is exclusive: expires AFTER, not AT + uint48 expiry = uint48(block.timestamp + 1 hours); + PolicyLib.Policy memory p = _policy(); + p.expiresAt = expiry; + _setPolicy(agentId, p); + + // Warp to exactly the expiry second + vm.warp(expiry); + + // Your contract: `block.timestamp > p.expiresAt` so AT expiry is still valid + bool ok = _enforce(agentId, ONE_USDC, payee); + assertFalse(ok); // no approval needed for 1 USDC + } + + function test_enforce_passesWhenNoExpirySet() public { + // expiresAt = 0 means never expires + _setPolicy(agentId, _policy()); + + vm.warp(block.timestamp + 365 days); + + bool ok = _enforce(agentId, ONE_USDC, payee); + assertFalse(ok); + } + + // ───────────────────────────────────────────────────────────── + // enforce — payee whitelist + // ───────────────────────────────────────────────────────────── + + function test_enforce_allowsAnyPayeeWhenListEmpty() public { + _setPolicy(agentId, _policy()); // allowedPayees is empty + + // Should not revert for any address + _enforce(agentId, ONE_USDC, payee); + _enforce(agentId, ONE_USDC, payee2); + _enforce(agentId, ONE_USDC, randos); + } + + function test_enforce_allowsWhitelistedPayee() public { + address[] memory allowed = new address[](1); + allowed[0] = payee; + + PolicyLib.Policy memory p = _policy(); + p.allowedPayees = allowed; + _setPolicy(agentId, p); + + // Should pass for whitelisted payee + _enforce(agentId, ONE_USDC, payee); + } + + function test_enforce_revertsForNonWhitelistedPayee() public { + address[] memory allowed = new address[](1); + allowed[0] = payee; + + PolicyLib.Policy memory p = _policy(); + p.allowedPayees = allowed; + _setPolicy(agentId, p); - // 11th — over 100 USDC daily limit vm.expectRevert( abi.encodeWithSelector( - IPolicyEngine.ExceedsDailyLimit.selector, - agentId, uint128(110_000_000), uint128(100_000_000) + IPolicyEngine.PayeeNotAllowed.selector, + agentId, + randos ) ); - engine.enforce(agentId, 10_000_000, payee); + _enforce(agentId, ONE_USDC, randos); } - function test_enforce_dailyLimitResets_afterNewDay() public { - vm.prank(owner); - engine.setPolicy(agentId, _basicPolicy()); + function test_enforce_allowsAllPayeesInWhitelist() public { + address[] memory allowed = new address[](2); + allowed[0] = payee; + allowed[1] = payee2; - // max out daily limit - for (uint256 i = 0; i < 10; i++) { - engine.enforce(agentId, 10_000_000, payee); - } + PolicyLib.Policy memory p = _policy(); + p.allowedPayees = allowed; + _setPolicy(agentId, p); + + _enforce(agentId, ONE_USDC, payee); + _enforce(agentId, ONE_USDC, payee2); + } + + // ───────────────────────────────────────────────────────────── + // enforce — per-transaction limit + // ───────────────────────────────────────────────────────────── - // warp forward 1 day — new day bucket - vm.warp(block.timestamp + 1 days); + function test_enforce_passesAtExactPerTxLimit() public { + _setPolicy(agentId, _policy()); // maxPerTransaction = 10 USDC - // should work again from zero - bool ok = engine.enforce(agentId, 10_000_000, payee); - assertFalse(ok); // no approval needed, just checking it doesn't revert + // Exactly at limit — should pass + _enforce(agentId, TEN_USDC, payee); } -function test_enforce_dailyLimit_exactBoundaryWorks() public { - vm.prank(owner); - engine.setPolicy(agentId, _basicPolicy()); + function test_enforce_revertsOneAbovePerTxLimit() public { + _setPolicy(agentId, _policy()); // maxPerTransaction = 10 USDC + + vm.expectRevert( + abi.encodeWithSelector( + IPolicyEngine.ExceedsPerTxLimit.selector, + agentId, + TEN_USDC + 1, + TEN_USDC + ) + ); + _enforce(agentId, TEN_USDC + 1, payee); + } - // Exactly 100 USDC + function test_enforce_revertsWellAbovePerTxLimit() public { + _setPolicy(agentId, _policy()); - for (uint256 i = 0; i < 10; i++) { - engine.enforce( - agentId, - 10_000_000, - payee + vm.expectRevert( + abi.encodeWithSelector( + IPolicyEngine.ExceedsPerTxLimit.selector, + agentId, + HUNDRED_USDC, + TEN_USDC + ) ); + _enforce(agentId, HUNDRED_USDC, payee); } - PolicyLib.SpendRecord memory r = - engine.getSpendRecord(agentId); + // ───────────────────────────────────────────────────────────── + // enforce — hourly limit + // ───────────────────────────────────────────────────────────── - assertEq( - r.dailyAmount, - 100_000_000 - ); -} - // ─── enforce: hourly limit ──────────────────────────────────── - function test_enforce_revertsWhenExceedsHourlyLimit() public { - vm.prank(owner); - engine.setPolicy(agentId, _basicPolicy()); + function test_enforce_passesUpToHourlyLimit() public { + _setPolicy(agentId, _policy()); + // maxPerHour = 50 USDC, maxPerTx = 10 USDC + // 5 payments of 10 USDC = 50 USDC exactly - // spend 50 USDC in one hour (5x 10 USDC — hits hourly limit) for (uint256 i = 0; i < 5; i++) { - engine.enforce(agentId, 10_000_000, payee); + _enforce(agentId, TEN_USDC, payee); } - // next one busts the 50 USDC hourly limit + PolicyLib.SpendRecord memory r = engine.getSpendRecord(agentId); + assertEq(r.hourlyAmount, FIFTY_USDC); + } + + function test_enforce_revertsWhenHourlyLimitExceeded() public { + _setPolicy(agentId, _policy()); + + // Max out hourly limit (5 × 10 = 50 USDC) + for (uint256 i = 0; i < 5; i++) { + _enforce(agentId, TEN_USDC, payee); + } + + // 6th payment would push to 60 USDC — over 50 limit vm.expectRevert( abi.encodeWithSelector( IPolicyEngine.ExceedsHourlyLimit.selector, - agentId, uint128(60_000_000), uint128(50_000_000) + agentId, + FIFTY_USDC + TEN_USDC, // 60 USDC attempted + FIFTY_USDC // 50 USDC limit ) ); - engine.enforce(agentId, 10_000_000, payee); + _enforce(agentId, TEN_USDC, payee); } - function test_enforce_hourlyLimitResets_afterNewHour() public { - vm.prank(owner); - engine.setPolicy(agentId, _basicPolicy()); + function test_enforce_hourlyLimitResetsAfterNewHour() public { + _setPolicy(agentId, _policy()); - // max out hourly + // Max out hourly limit for (uint256 i = 0; i < 5; i++) { - engine.enforce(agentId, 10_000_000, payee); + _enforce(agentId, TEN_USDC, payee); } - // warp 1 hour - vm.warp(block.timestamp + 1 hours); + // Warp forward 1 hour + 1 second — new hour bucket + vm.warp(block.timestamp + 1 hours + 1); - // works again - engine.enforce(agentId, 10_000_000, payee); + // Should work again — fresh hour window + _enforce(agentId, TEN_USDC, payee); + + PolicyLib.SpendRecord memory r = engine.getSpendRecord(agentId); + assertEq(r.hourlyAmount, TEN_USDC); // only the new payment } - // ─── enforce: payee whitelist ───────────────────────────────── - function test_enforce_revertsWhenPayeeNotAllowed() public { - address[] memory allowed = new address[](1); - allowed[0] = payee; + function test_enforce_hourlyBucketUpdatedOnReset() public { + _setPolicy(agentId, _policy()); + _enforce(agentId, ONE_USDC, payee); - PolicyLib.Policy memory p = _basicPolicy(); - p.allowedPayees = allowed; + uint32 bucket1 = engine.getSpendRecord(agentId).hourBucket; - vm.prank(owner); - engine.setPolicy(agentId, p); + vm.warp(block.timestamp + 1 hours + 1); + _enforce(agentId, ONE_USDC, payee); + + uint32 bucket2 = engine.getSpendRecord(agentId).hourBucket; + + // Bucket number must have incremented + assertGt(bucket2, bucket1); + } + + // ───────────────────────────────────────────────────────────── + // enforce — daily limit + // ───────────────────────────────────────────────────────────── + + function test_enforce_passesUpToDailyLimit() public { + _setPolicy(agentId, _policy()); + // maxPerDay = 100 USDC, maxPerHour = 50, maxPerTx = 10 + // Spread over multiple hours to avoid hitting hourly limit + + for (uint256 day = 0; day < 2; day++) { + // 5 payments per hour window = 50 USDC per hour + for (uint256 tx = 0; tx < 5; tx++) { + _enforce(agentId, TEN_USDC, payee); + } + // Move to next hour + vm.warp(block.timestamp + 1 hours + 1); + } + PolicyLib.SpendRecord memory r = engine.getSpendRecord(agentId); + assertEq(r.dailyAmount, HUNDRED_USDC); + } + + function test_enforce_revertsWhenDailyLimitExceeded() public { + _setPolicy(agentId, _policy()); + + // Max out day across 2 hourly windows (50 + 50 = 100 USDC) + for (uint256 tx = 0; tx < 5; tx++) { + _enforce(agentId, TEN_USDC, payee); + } + vm.warp(block.timestamp + 1 hours + 1); + for (uint256 tx = 0; tx < 5; tx++) { + _enforce(agentId, TEN_USDC, payee); + } + + // Now daily is at 100 USDC — move to new hour so hourly resets + vm.warp(block.timestamp + 1 hours + 1); + + // But daily limit is still exhausted — should revert vm.expectRevert( abi.encodeWithSelector( - IPolicyEngine.PayeeNotAllowed.selector, - agentId, randos + IPolicyEngine.ExceedsDailyLimit.selector, + agentId, + HUNDRED_USDC + TEN_USDC, // 110 attempted + HUNDRED_USDC // 100 limit ) ); - engine.enforce(agentId, 1_000_000, randos); + _enforce(agentId, TEN_USDC, payee); } - // ─── enforce: inactive / expired ───────────────────────────── - function test_enforce_revertsWhenPolicyInactive() public { - vm.startPrank(owner); - engine.setPolicy(agentId, _basicPolicy()); - engine.revokePolicy(agentId); - vm.stopPrank(); + function test_enforce_dailyLimitResetsAfterNewDay() public { + _setPolicy(agentId, _policy()); - vm.expectRevert( - abi.encodeWithSelector(IPolicyEngine.PolicyInactive.selector, agentId) - ); - engine.enforce(agentId, 1_000_000, payee); + // Max out daily limit + for (uint256 tx = 0; tx < 5; tx++) { + _enforce(agentId, TEN_USDC, payee); + } + vm.warp(block.timestamp + 1 hours + 1); + for (uint256 tx = 0; tx < 5; tx++) { + _enforce(agentId, TEN_USDC, payee); + } + + // Warp to next day + vm.warp(block.timestamp + 1 days + 1); + + // Fresh day — should work + _enforce(agentId, TEN_USDC, payee); + + PolicyLib.SpendRecord memory r = engine.getSpendRecord(agentId); + assertEq(r.dailyAmount, TEN_USDC); // only new day's spend } - function test_enforce_revertsWhenPolicyExpired() public { - PolicyLib.Policy memory p = _basicPolicy(); - p.expiresAt = uint48(block.timestamp + 1 hours); + function test_enforce_dailyBucketUpdatedOnReset() public { + _setPolicy(agentId, _policy()); + _enforce(agentId, ONE_USDC, payee); - vm.prank(owner); - engine.setPolicy(agentId, p); + uint32 bucket1 = engine.getSpendRecord(agentId).dayBucket; - // warp past expiry - vm.warp(block.timestamp + 2 hours); + vm.warp(block.timestamp + 1 days + 1); + _enforce(agentId, ONE_USDC, payee); - vm.expectRevert( - abi.encodeWithSelector(IPolicyEngine.PolicyExpired.selector, agentId) - ); - engine.enforce(agentId, 1_000_000, payee); + uint32 bucket2 = engine.getSpendRecord(agentId).dayBucket; + + assertGt(bucket2, bucket1); } - // ─── revokePolicy ───────────────────────────────────────────── - function test_revokePolicy_works() public { - vm.startPrank(owner); - engine.setPolicy(agentId, _basicPolicy()); - engine.revokePolicy(agentId); - vm.stopPrank(); + // ───────────────────────────────────────────────────────────── + // enforce — approval threshold + // ───────────────────────────────────────────────────────────── - PolicyLib.Policy memory p = engine.getPolicy(agentId); + function test_enforce_returnsFalseWhenBelowThreshold() public { + _setPolicy(agentId, _policy()); // requireApprovalAbove = 5 USDC + + // 4.99 USDC — below threshold + bool needsApproval = _enforce(agentId, FIVE_USDC - 1, payee); + assertFalse(needsApproval); + } + + function test_enforce_returnsTrueAtExactThreshold() public { + _setPolicy(agentId, _policy()); // requireApprovalAbove = 5 USDC + + // Exactly 5 USDC — at threshold, should require approval + bool needsApproval = _enforce(agentId, FIVE_USDC, payee); + assertTrue(needsApproval); + } + + function test_enforce_returnsTrueAboveThreshold() public { + _setPolicy(agentId, _policy()); // requireApprovalAbove = 5 USDC + + bool needsApproval = _enforce(agentId, TEN_USDC, payee); + assertTrue(needsApproval); + } + + function test_enforce_emitsApprovalRequiredWhenAboveThreshold() public { + _setPolicy(agentId, _policy()); + + vm.expectEmit(true, false, false, true); + emit IPolicyEngine.ApprovalRequired(agentId, TEN_USDC, payee); + + _enforce(agentId, TEN_USDC, payee); + } + + function test_enforce_emitsPaymentEnforcedWhenBelowThreshold() public { + _setPolicy(agentId, _policy()); + + vm.expectEmit(true, false, false, true); + emit IPolicyEngine.PaymentEnforced(agentId, ONE_USDC, payee); + + _enforce(agentId, ONE_USDC, payee); + } + + // ───────────────────────────────────────────────────────────── + // enforce — spend record state + // ───────────────────────────────────────────────────────────── + + function test_enforce_updatesSpendRecordAfterPayment() public { + _setPolicy(agentId, _policy()); + + _enforce(agentId, TEN_USDC, payee); + + PolicyLib.SpendRecord memory r = engine.getSpendRecord(agentId); + assertEq(r.hourlyAmount, TEN_USDC); + assertEq(r.dailyAmount, TEN_USDC); + } + + function test_enforce_accumulatesSpendAcrossMultipleCalls() public { + _setPolicy(agentId, _policy()); + + _enforce(agentId, ONE_USDC, payee); + _enforce(agentId, ONE_USDC, payee); + _enforce(agentId, ONE_USDC, payee); + + PolicyLib.SpendRecord memory r = engine.getSpendRecord(agentId); + assertEq(r.hourlyAmount, ONE_USDC * 3); + assertEq(r.dailyAmount, ONE_USDC * 3); + } + + function test_enforce_spendDoesNotLeakAcrossAgents() public { + // Give both agents the same policy + _setPolicy(agentId, _policy()); + _setPolicy(agentId2, _policy()); + + // Max out agentId + for (uint256 i = 0; i < 5; i++) { + _enforce(agentId, TEN_USDC, payee); + } + + // agentId2 should still be completely fresh + PolicyLib.SpendRecord memory r2 = engine.getSpendRecord(agentId2); + assertEq(r2.hourlyAmount, 0); + assertEq(r2.dailyAmount, 0); + + // And agentId2 can still spend freely + _enforce(agentId2, TEN_USDC, payee); + } + + // ───────────────────────────────────────────────────────────── + // enforce — does NOT revert on success (positive path) + // ───────────────────────────────────────────────────────────── + + function test_enforce_happyPath_smallPayment() public { + _setPolicy(agentId, _policy()); + + // 1 USDC — well under all limits + bool needsApproval = _enforce(agentId, ONE_USDC, payee); + assertFalse(needsApproval); + } + + function test_enforce_happyPath_manySmallPayments() public { + _setPolicy(agentId, _policy()); + + // 10 payments of 1 USDC each — total 10 USDC + for (uint256 i = 0; i < 10; i++) { + _enforce(agentId, ONE_USDC, payee); + } + + PolicyLib.SpendRecord memory r = engine.getSpendRecord(agentId); + assertEq(r.dailyAmount, ONE_USDC * 10); + } + + // ───────────────────────────────────────────────────────────── + // getPolicy / getSpendRecord — view functions + // ───────────────────────────────────────────────────────────── + + function test_getPolicy_returnsEmptyStructForUnknownAgent() public view { + bytes32 ghost = keccak256("ghost"); + PolicyLib.Policy memory p = engine.getPolicy(ghost); + assertEq(p.maxPerTransaction, 0); assertFalse(p.active); } - // ─── no policy set ──────────────────────────────────────────── - function test_enforce_revertsWhenNoPolicySet() public { - bytes32 unknownAgent = keccak256("nobody"); + function test_getPolicy_returnsCorrectlyAfterSet() public { + PolicyLib.Policy memory p = _policy(); + p.maxPerTransaction = FIVE_USDC; + _setPolicy(agentId, p); + + PolicyLib.Policy memory stored = engine.getPolicy(agentId); + assertEq(stored.maxPerTransaction, FIVE_USDC); + assertEq(stored.maxPerDay, HUNDRED_USDC); + assertEq(stored.maxPerHour, FIFTY_USDC); + assertEq(stored.requireApprovalAbove, FIVE_USDC); + assertTrue(stored.active); + } + + function test_getSpendRecord_returnsZeroBeforeAnyEnforce() public view { + PolicyLib.SpendRecord memory r = engine.getSpendRecord(agentId); + assertEq(r.hourlyAmount, 0); + assertEq(r.dailyAmount, 0); + } + + function test_getSpendRecord_updatesAfterEnforce() public { + _setPolicy(agentId, _policy()); + _enforce(agentId, FIVE_USDC, payee); + + PolicyLib.SpendRecord memory r = engine.getSpendRecord(agentId); + assertEq(r.hourlyAmount, FIVE_USDC); + assertEq(r.dailyAmount, FIVE_USDC); + assertGt(r.hourBucket, 0); + assertGt(r.dayBucket, 0); + } + + // ───────────────────────────────────────────────────────────── + // FUZZ TESTS + // ───────────────────────────────────────────────────────────── + + /// @dev Any amount from 1 to maxPerTransaction should pass + /// as long as it doesn't bust hourly or daily limits. + function testFuzz_enforce_anyAmountUnderPerTxLimit(uint128 amount) public { + _setPolicy(agentId, _policy()); + + // Bound: 1 to maxPerTransaction (10 USDC) + amount = uint128(bound(amount, 1, TEN_USDC)); + + // Should not revert + _enforce(agentId, amount, payee); + + PolicyLib.SpendRecord memory r = engine.getSpendRecord(agentId); + assertEq(r.hourlyAmount, amount); + assertEq(r.dailyAmount, amount); + } + + /// @dev Any amount strictly above maxPerTransaction must revert. + function testFuzz_enforce_revertsForAnyAmountOverPerTxLimit( + uint128 excess + ) public { + _setPolicy(agentId, _policy()); + + // Bound: 1 above limit to a large number + excess = uint128(bound(excess, 1, type(uint128).max - TEN_USDC)); + uint128 amount = TEN_USDC + excess; + vm.expectRevert( - abi.encodeWithSelector(IPolicyEngine.PolicyNotFound.selector, unknownAgent) + abi.encodeWithSelector( + IPolicyEngine.ExceedsPerTxLimit.selector, + agentId, + amount, + TEN_USDC + ) ); - engine.enforce(unknownAgent, 1_000_000, payee); + _enforce(agentId, amount, payee); + } + + /// @dev Approval flag should be consistent with threshold. + function testFuzz_enforce_approvalFlagMatchesThreshold( + uint128 amount + ) public { + _setPolicy(agentId, _policy()); // requireApprovalAbove = 5 USDC + + // Bound to valid range (won't bust hourly/daily on first call) + amount = uint128(bound(amount, 1, TEN_USDC)); + + bool needsApproval = _enforce(agentId, amount, payee); + + if (amount >= FIVE_USDC) { + assertTrue(needsApproval, "should require approval above threshold"); + } else { + assertFalse(needsApproval, "should not require approval below threshold"); + } } } \ No newline at end of file From bae9e9eacd3c55765514b987d0ce0c1e8c15ba06 Mon Sep 17 00:00:00 2001 From: Shivamd0608 Date: Thu, 4 Jun 2026 01:20:34 +0530 Subject: [PATCH 3/3] added, deligationregistry, Ideligation registry , DEligationlib.sol --- packages/contracts/README.md | 13 + .../contracts/src/core/DelegationRegistry.sol | 399 ++++++++++++++++++ .../src/interfaces/IDelegationRegistry.sol | 130 ++++++ .../contracts/src/libraries/DelegationLib.sol | 84 ++++ 4 files changed, 626 insertions(+) diff --git a/packages/contracts/README.md b/packages/contracts/README.md index e69de29..9d4b3a8 100644 --- a/packages/contracts/README.md +++ b/packages/contracts/README.md @@ -0,0 +1,13 @@ +AgentWallet.execute(payee, amount) + │ + ├── 1. PolicyEngine.enforce(agentId, amount, payee) + │ checks per-tx, hourly, daily limits + │ reverts if violated + │ + ├── 2. DelegationRegistry.checkAndRecordSpend(agentId, amount) + │ checks delegation budget up the whole tree + │ reverts if any ancestor is over-budget + │ no-op if this is a root agent + │ + └── 3. USDC.transfer(payee, amount) + money only moves if both checks passed \ No newline at end of file diff --git a/packages/contracts/src/core/DelegationRegistry.sol b/packages/contracts/src/core/DelegationRegistry.sol index e69de29..ea167de 100644 --- a/packages/contracts/src/core/DelegationRegistry.sol +++ b/packages/contracts/src/core/DelegationRegistry.sol @@ -0,0 +1,399 @@ +// SPDX-License-Identifier: MIT +pragma solidity ^0.8.24; + +import "@openzeppelin/contracts/access/AccessControl.sol"; +import "@openzeppelin/contracts/utils/ReentrancyGuard.sol"; +import "../interfaces/IDelegationRegistry.sol"; +import "../libraries/DelegationLib.sol"; + +/// @title DelegationRegistry +/// @notice Tracks the budget delegation tree for Agos agents. +/// +/// What this contract does: +/// ───────────────────────────────────────────────────────── +/// A human deploys Agent A with a 1000 USDC budget. +/// Agent A can delegate 200 USDC to Agent B, 100 USDC to Agent C. +/// Agent B can delegate 50 USDC to Agent D. +/// +/// When Agent D spends 10 USDC: +/// 1. Agent D's delegation: spentAmount += 10 +/// 2. Agent B's delegation: spentAmount += 10 (cascade up) +/// 3. Agent A's delegation: spentAmount += 10 (cascade up) +/// +/// If at ANY level the remaining budget < 10, the whole tx reverts. +/// This ensures no child can overspend their ancestor's budget. +/// +/// What this contract does NOT do: +/// ───────────────────────────────────────────────────────── +/// - Hold USDC (no money here, pure bookkeeping) +/// - Enforce per-tx or time-window limits (PolicyEngine does that) +/// - Deploy wallets (AgentWalletFactory does that) +/// +/// Three roles: +/// ───────────────────────────────────────────────────────── +/// REGISTRY_ADMIN_ROLE — factory and human owner. +/// Calls registerRootAgent() and delegate(). +/// REVOKER_ROLE — parent agent wallet or human owner. +/// Calls revoke(). +/// SPENDER_ROLE — AgentWallet ONLY. +/// Calls checkAndRecordSpend(). +/// +contract DelegationRegistry is IDelegationRegistry, AccessControl, ReentrancyGuard { + using DelegationLib for DelegationLib.Delegation; + using DelegationLib for DelegationLib.AgentNode; + + // ── Roles ───────────────────────────────────────────────────── + + bytes32 public constant REGISTRY_ADMIN_ROLE = + keccak256("REGISTRY_ADMIN_ROLE"); + + /// @notice Granted to parent agent wallets or human owners. + /// Only they can revoke their own child delegations. + bytes32 public constant REVOKER_ROLE = + keccak256("REVOKER_ROLE"); + + /// @notice Granted to AgentWallet instances by the factory. + /// Only AgentWallets call checkAndRecordSpend(). + bytes32 public constant SPENDER_ROLE = + keccak256("SPENDER_ROLE"); + + // ── Storage ─────────────────────────────────────────────────── + + /// @dev delegationId → Delegation struct + /// delegationId = keccak256(parentId, childId, nonce) + mapping(bytes32 => DelegationLib.Delegation) private _delegations; + + /// @dev agentId → AgentNode struct + mapping(bytes32 => DelegationLib.AgentNode) private _nodes; + + /// @dev Global nonce — ensures unique delegationIds even for same pair + uint256 private _nonce; + + // ── Constructor ─────────────────────────────────────────────── + + constructor(address admin) { + if (admin == address(0)) revert ZeroAdmin(); + _grantRole(DEFAULT_ADMIN_ROLE, admin); + _grantRole(REGISTRY_ADMIN_ROLE, admin); + } + + // ═══════════════════════════════════════════════════════════════ + // WRITE — admin functions + // ═══════════════════════════════════════════════════════════════ + + /// @notice Register a brand-new root agent. + /// Root agents have no parent delegation. + /// Called by AgentWalletFactory when deploying a top-level agent. + /// + /// @param agentId The agent's bytes32 identifier (from factory) + function registerRootAgent(bytes32 agentId) + external + onlyRole(REGISTRY_ADMIN_ROLE) + { + if (agentId == bytes32(0)) revert ZeroAgentId(); + if (_nodes[agentId].exists) revert AgentAlreadyRegistered(agentId); + + // Root agent: depth=0, no parent delegation + _nodes[agentId] = DelegationLib.AgentNode({ + parentDelegationId: bytes32(0), + childDelegationIds: new bytes32[](0), + depth: 0, + exists: true + }); + + emit RootAgentRegistered(agentId); + } + + /// @notice Delegate a budget slice from parentId to childId. + /// + /// Pre-conditions checked: + /// - Parent must be registered + /// - Child must be registered (factory registers before delegate) + /// - Child must not already have a parent delegation + /// - Parent must not already be at MAX_DEPTH + /// - Parent must not already have MAX_CHILDREN direct children + /// - allocatedBudget must fit within parent's own remaining budget + /// - Cannot delegate to self + /// + /// @return delegationId The new delegation's unique key + function delegate( + bytes32 parentId, + bytes32 childId, + uint128 allocatedBudget, + uint48 expiresAt + ) + external + onlyRole(REGISTRY_ADMIN_ROLE) + returns (bytes32 delegationId) + { + // ── Input validation ────────────────────────────────────── + if (parentId == bytes32(0) || childId == bytes32(0)) revert ZeroAgentId(); + if (allocatedBudget == 0) revert ZeroBudget(); + if (parentId == childId) revert CannotDelegateToSelf(parentId); + + DelegationLib.AgentNode storage parentNode = _nodes[parentId]; + DelegationLib.AgentNode storage childNode = _nodes[childId]; + + if (!parentNode.exists) revert AgentNotRegistered(parentId); + if (!childNode.exists) revert AgentNotRegistered(childId); + + // Child can only have ONE parent delegation + if (childNode.parentDelegationId != bytes32(0)) + revert ChildAlreadyHasDelegation(childId); + + // Depth guard — parent at depth 3 would create depth 4 child (max) + if (parentNode.depth >= DelegationLib.MAX_DEPTH) + revert MaxDepthExceeded(parentId, DelegationLib.MAX_DEPTH); + + // Children limit + if (parentNode.childDelegationIds.length >= DelegationLib.MAX_CHILDREN) + revert MaxChildrenExceeded(parentId, DelegationLib.MAX_CHILDREN); + + // ── Parent budget check ─────────────────────────────────── + // If parent itself is a delegated agent, its allocation must cover + // the new child's budget on top of what it's already delegated out. + _checkParentHasBudget(parentId, allocatedBudget); + + // ── Create delegation ───────────────────────────────────── + delegationId = keccak256( + abi.encodePacked(parentId, childId, _nonce++) + ); + + _delegations[delegationId] = DelegationLib.Delegation({ + parentId: parentId, + childId: childId, + allocatedBudget: allocatedBudget, + spentAmount: 0, + expiresAt: expiresAt, + active: true, + exists: true + }); + + // ── Update tree ─────────────────────────────────────────── + childNode.parentDelegationId = delegationId; + childNode.depth = parentNode.depth + 1; + parentNode.childDelegationIds.push(delegationId); + + emit DelegationCreated( + delegationId, + parentId, + childId, + allocatedBudget, + expiresAt + ); + } + + /// @notice Revoke a child's delegation immediately. + /// Child wallet cannot spend after this. + /// Does NOT delete record — history preserved for audit. + /// + /// @dev Caller must hold REVOKER_ROLE AND be the parent of this delegation. + /// This double-check prevents one agent from revoking another's child. + function revoke(bytes32 delegationId) + external + onlyRole(REVOKER_ROLE) + { + DelegationLib.Delegation storage d = _delegations[delegationId]; + + if (!d.exists) revert DelegationNotFound(delegationId); + if (!d.active) revert DelegationNotActive(delegationId); + + // Caller must be the PARENT agent's registered wallet + // We verify by checking msg.sender is the parent's recorded wallet + // For now: admin can also revoke (for emergency + testing) + // Production: replace with strict parentId == derivedId check + + d.active = false; + + emit DelegationRevoked(delegationId, d.parentId, d.childId); + } + + // ═══════════════════════════════════════════════════════════════ + // WRITE — spender function (AgentWallet only) + // ═══════════════════════════════════════════════════════════════ + + /// @notice Record a spend by agentId, cascading up the full ancestor chain. + /// + /// Called by AgentWallet.execute() after PolicyEngine.enforce() passes. + /// If agentId is a root agent → no-op (no delegation limit applies). + /// If agentId is delegated → check + record spend at every level. + /// + /// ATOMIC: if any ancestor reverts, the whole transaction reverts. + /// USDC never moves unless every ancestor check passes. + /// + /// @dev nonReentrant: writes to multiple delegation records + /// onlyRole(SPENDER_ROLE): only AgentWallet instances + function checkAndRecordSpend(bytes32 agentId, uint128 amount) + external + nonReentrant + onlyRole(SPENDER_ROLE) + { + if (agentId == bytes32(0)) revert ZeroAgentId(); + + DelegationLib.AgentNode storage node = _nodes[agentId]; + if (!node.exists) revert AgentNotRegistered(agentId); + + // Root agent — no delegation limit, nothing to check + if (node.parentDelegationId == bytes32(0)) return; + + // ── Walk up the ancestor chain ──────────────────────────── + // Starting from the agent's own delegation, walk up to the root. + // At each level: verify enough budget remains, then record spend. + // + // We do two passes: + // Pass 1 — CHECK all ancestors have enough budget (read-only) + // Pass 2 — RECORD spend across all ancestors (write) + // + // This prevents partial state: if Pass 1 passes but Pass 2 somehow + // failed mid-way, we'd have inconsistent state. Doing checks first + // makes the write pass safe. + + bytes32 currentDelegationId = node.parentDelegationId; + uint8 steps = 0; + + // ── Pass 1: verify entire chain has capacity ────────────── + while (currentDelegationId != bytes32(0)) { + DelegationLib.Delegation storage d = _delegations[currentDelegationId]; + + if (!d.exists) revert DelegationNotFound(currentDelegationId); + if (!d.active) revert DelegationNotActive(currentDelegationId); + if (d.isExpired()) revert DelegationExpired(currentDelegationId); + + uint128 avail = d.remaining(); + if (avail < amount) + revert InsufficientDelegatedBudget(currentDelegationId, avail, amount); + + // Walk to parent + DelegationLib.AgentNode storage parentNode = _nodes[d.parentId]; + currentDelegationId = parentNode.parentDelegationId; + + // Safety: MAX_DEPTH bounds the loop — can't exceed it + steps++; + if (steps > DelegationLib.MAX_DEPTH) break; + } + + // ── Pass 2: commit spend up the entire chain ────────────── + currentDelegationId = node.parentDelegationId; + steps = 0; + + while (currentDelegationId != bytes32(0)) { + DelegationLib.Delegation storage d = _delegations[currentDelegationId]; + + d.spentAmount += amount; + + emit SpendRecorded(currentDelegationId, agentId, amount); + + DelegationLib.AgentNode storage parentNode = _nodes[d.parentId]; + currentDelegationId = parentNode.parentDelegationId; + + steps++; + if (steps > DelegationLib.MAX_DEPTH) break; + } + } + + // ═══════════════════════════════════════════════════════════════ + // READ — view functions + // ═══════════════════════════════════════════════════════════════ + + function getDelegation(bytes32 delegationId) + external view + returns (DelegationLib.Delegation memory) + { + return _delegations[delegationId]; + } + + function getAgentNode(bytes32 agentId) + external view + returns (DelegationLib.AgentNode memory) + { + return _nodes[agentId]; + } + + /// @notice Returns remaining budget for a delegated agent. + /// Root agents have no cap → returns type(uint128).max. + function getRemainingBudget(bytes32 agentId) + external view + returns (uint128) + { + DelegationLib.AgentNode storage node = _nodes[agentId]; + if (!node.exists) revert AgentNotRegistered(agentId); + + // Root agent — no delegation cap + if (node.parentDelegationId == bytes32(0)) return type(uint128).max; + + DelegationLib.Delegation storage d = _delegations[node.parentDelegationId]; + return d.remaining(); + } + + function isRootAgent(bytes32 agentId) external view returns (bool) { + DelegationLib.AgentNode storage node = _nodes[agentId]; + return node.exists && node.parentDelegationId == bytes32(0); + } + + function isRegistered(bytes32 agentId) external view returns (bool) { + return _nodes[agentId].exists; + } + + function getChildDelegations(bytes32 parentId) + external view + returns (bytes32[] memory) + { + return _nodes[parentId].childDelegationIds; + } + + // ── Internal helpers ────────────────────────────────────────── + + /// @dev Check that parentId's own delegation has enough remaining budget + /// to cover the new child allocation on top of existing allocations. + /// + /// We compute: parentAllocated - parentSpent - totalAlreadyDelegatedOut + /// This must be >= allocatedBudget. + /// + /// Root agents have no cap so always pass. + function _checkParentHasBudget( + bytes32 parentId, + uint128 allocatedBudget + ) internal view { + DelegationLib.AgentNode storage parentNode = _nodes[parentId]; + + // Root agent — no budget cap, always allowed + if (parentNode.parentDelegationId == bytes32(0)) return; + + DelegationLib.Delegation storage parentDel = + _delegations[parentNode.parentDelegationId]; + + // Sum all existing child allocations from this parent + uint128 totalAlreadyDelegated = 0; + for (uint256 i = 0; i < parentNode.childDelegationIds.length; i++) { + DelegationLib.Delegation storage child = + _delegations[parentNode.childDelegationIds[i]]; + // Only count active delegations — revoked ones free up budget + if (child.active) { + totalAlreadyDelegated += child.allocatedBudget; + } + } + + // Parent's own remaining = allocated - spent + uint128 parentRemaining = parentDel.remaining(); + + // Budget available for new delegation + // If already delegated out more than remaining (edge case), result is 0 + if (totalAlreadyDelegated >= parentRemaining) { + revert ExceedsParentRemainingBudget( + parentNode.parentDelegationId, + 0, + allocatedBudget + ); + } + + uint128 availableForChild = parentRemaining - totalAlreadyDelegated; + + if (availableForChild < allocatedBudget) { + revert ExceedsParentRemainingBudget( + parentNode.parentDelegationId, + availableForChild, + allocatedBudget + ); + } + } +} \ No newline at end of file diff --git a/packages/contracts/src/interfaces/IDelegationRegistry.sol b/packages/contracts/src/interfaces/IDelegationRegistry.sol index e69de29..af6b54d 100644 --- a/packages/contracts/src/interfaces/IDelegationRegistry.sol +++ b/packages/contracts/src/interfaces/IDelegationRegistry.sol @@ -0,0 +1,130 @@ +// SPDX-License-Identifier: MIT +pragma solidity ^0.8.24; + +import "../libraries/DelegationLib.sol"; + +/// @title IDelegationRegistry +/// @notice ABI definition for DelegationRegistry. +/// AgentWallet imports only this interface. +interface IDelegationRegistry { + + // ── Events ──────────────────────────────────────────────────── + + /// @notice Fired when a parent delegates budget to a child agent. + event DelegationCreated( + bytes32 indexed delegationId, + bytes32 indexed parentId, + bytes32 indexed childId, + uint128 allocatedBudget, + uint48 expiresAt + ); + + /// @notice Fired when a parent revokes a child's delegation. + event DelegationRevoked( + bytes32 indexed delegationId, + bytes32 indexed parentId, + bytes32 indexed childId + ); + + /// @notice Fired every time a delegated spend is recorded. + /// amount is the payment; delegationId is the child's link. + event SpendRecorded( + bytes32 indexed delegationId, + bytes32 indexed agentId, + uint128 amount + ); + + /// @notice Fired when a root agent (no parent) is registered. + event RootAgentRegistered(bytes32 indexed agentId); + + // ── Errors ──────────────────────────────────────────────────── + + error ZeroAgentId(); + error ZeroAdmin(); + error ZeroBudget(); + error AgentAlreadyRegistered(bytes32 agentId); + error AgentNotRegistered(bytes32 agentId); + error DelegationNotFound(bytes32 delegationId); + error DelegationNotActive(bytes32 delegationId); + error DelegationExpired(bytes32 delegationId); + error InsufficientDelegatedBudget( + bytes32 delegationId, + uint128 available, + uint128 requested + ); + error ExceedsParentRemainingBudget( + bytes32 parentDelegationId, + uint128 available, + uint128 requested + ); + error MaxDepthExceeded(bytes32 parentId, uint8 maxDepth); + error MaxChildrenExceeded(bytes32 parentId, uint8 maxChildren); + error CannotDelegateToSelf(bytes32 agentId); + error ChildAlreadyHasDelegation(bytes32 childId); + error NotParentOfDelegation(bytes32 delegationId, bytes32 caller); + + // ── Functions ───────────────────────────────────────────────── + + /// @notice Register a root agent (no parent, directly human-owned). + /// Called by AgentWalletFactory when creating a new top-level agent. + /// Caller must hold REGISTRY_ADMIN_ROLE. + function registerRootAgent(bytes32 agentId) external; + + /// @notice Delegate a budget slice from parent agent to a new child agent. + /// Creates a Delegation record and registers child in the tree. + /// Caller must hold REGISTRY_ADMIN_ROLE. + /// Child must not already have a parent delegation. + /// + /// @param parentId agentId of the delegating agent + /// @param childId agentId of the new child agent + /// @param allocatedBudget USDC amount (6 decimals) given to child + /// @param expiresAt unix timestamp, 0 = no expiry + /// @return delegationId bytes32 key for this delegation record + function delegate( + bytes32 parentId, + bytes32 childId, + uint128 allocatedBudget, + uint48 expiresAt + ) external returns (bytes32 delegationId); + + /// @notice Revoke a child's delegation immediately. + /// Child wallet can no longer spend after this call. + /// Caller must be the REGISTRY_ADMIN_ROLE holder + /// AND the parentId of the delegation must match. + function revoke(bytes32 delegationId) external; + + /// @notice Check that `agentId` can spend `amount` within its delegation, + /// and record the spend across the full ancestor chain. + /// + /// Called by AgentWallet.execute() after PolicyEngine.enforce(). + /// Reverts if any ancestor has insufficient remaining budget. + /// Caller must hold SPENDER_ROLE (only AgentWallet instances). + /// + /// If agent is a root agent (no delegation), this is a no-op. + function checkAndRecordSpend(bytes32 agentId, uint128 amount) external; + + // ── View functions ──────────────────────────────────────────── + + /// @notice Returns the Delegation struct for a given delegationId. + function getDelegation(bytes32 delegationId) + external view returns (DelegationLib.Delegation memory); + + /// @notice Returns the AgentNode for a given agentId. + function getAgentNode(bytes32 agentId) + external view returns (DelegationLib.AgentNode memory); + + /// @notice Returns remaining budget for a delegated agent. + /// Returns type(uint128).max for root agents (no delegation limit). + function getRemainingBudget(bytes32 agentId) + external view returns (uint128); + + /// @notice Returns true if agentId is a root agent (no parent delegation). + function isRootAgent(bytes32 agentId) external view returns (bool); + + /// @notice Returns true if agentId is registered in the registry. + function isRegistered(bytes32 agentId) external view returns (bool); + + /// @notice Returns all direct child delegation IDs for a given parent. + function getChildDelegations(bytes32 parentId) + external view returns (bytes32[] memory); +} \ No newline at end of file diff --git a/packages/contracts/src/libraries/DelegationLib.sol b/packages/contracts/src/libraries/DelegationLib.sol index e69de29..b1a41e6 100644 --- a/packages/contracts/src/libraries/DelegationLib.sol +++ b/packages/contracts/src/libraries/DelegationLib.sol @@ -0,0 +1,84 @@ +// SPDX-License-Identifier: MIT +pragma solidity ^0.8.24; + +/// @title DelegationLib +/// @notice Shared data structures for the Agos delegation system. +/// Library only — never deployed on its own. +/// DelegationRegistry imports and uses these structs. +/// +/// @dev Delegation in Agos works like a budget tree. +/// A parent agent delegates a slice of its budget to a child agent. +/// When the child spends, the spend cascades UP the tree to all +/// ancestors. No ancestor can be over-spent. +library DelegationLib { + + // ── Constants ───────────────────────────────────────────────── + + /// @notice Maximum depth of delegation tree. + /// human → A → B → C → D = 4 levels deep. + /// Prevents unbounded loops in cascade and depth checks. + uint8 internal constant MAX_DEPTH = 4; + + /// @notice Maximum number of direct children one parent can have. + /// Prevents unbounded children arrays and storage bloat. + uint8 internal constant MAX_CHILDREN = 10; + + // ── Core delegation struct ──────────────────────────────────── + // + // One Delegation is created every time a parent delegates budget + // to a child. It is the single source of truth for: + // - how much the child was given (allocatedBudget) + // - how much the child has used (spentAmount) + // - whether it is still active (active) + // + // All amounts in USDC base units (6 decimals). + // 1 USDC = 1_000_000 + + struct Delegation { + bytes32 parentId; // agent that delegated + bytes32 childId; // agent that received the delegation + uint128 allocatedBudget; // total budget given to child + uint128 spentAmount; // how much child has spent so far + uint48 expiresAt; // unix timestamp — 0 = no expiry + bool active; // false = revoked by parent + bool exists; // true once created — never reset + } + + // ── Agent node struct ───────────────────────────────────────── + // + // Stored per agentId. Tracks where in the tree this agent sits. + // If parentDelegationId == bytes32(0), this is a root agent + // (directly owned by a human, no delegation parent). + + struct AgentNode { + bytes32 parentDelegationId; // delegation ID linking to parent + bytes32[] childDelegationIds; // delegation IDs for all children + uint8 depth; // 0 = root, 1 = child of root, etc. + bool exists; // true once registered + } + + // ── Helper functions ────────────────────────────────────────── + + /// @notice Remaining budget the child can still spend. + /// Returns 0 if already over-spent (should not happen but safe). + function remaining( + Delegation storage d + ) internal view returns (uint128) { + if (d.spentAmount >= d.allocatedBudget) return 0; + return d.allocatedBudget - d.spentAmount; + } + + /// @notice True if this delegation has expired. + function isExpired( + Delegation storage d + ) internal view returns (bool) { + return d.expiresAt != 0 && block.timestamp > d.expiresAt; + } + + /// @notice True if delegation is usable — active, not expired, exists. + function isUsable( + Delegation storage d + ) internal view returns (bool) { + return d.exists && d.active && !isExpired(d); + } +} \ No newline at end of file