Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

AgentRegister for AgentCommunication #24

Open
wants to merge 4 commits into
base: main
Choose a base branch
from
+186 −48
Open

AgentRegister for AgentCommunication #24

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
5738b74
Agent register in communication
kongzii Feb 10, 2025
8bf48bf
AgentRegister for AgentCommunication
kongzii Feb 12, 2025
58ce84c
use onlyRegisteredAgent in rest of places
kongzii Feb 12, 2025
00d0713
some mods
kongzii Feb 12, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
91 changes: 73 additions & 18 deletions src/NFT/AgentCommunication.sol
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,19 +4,70 @@ pragma solidity ^0.8.22;
import {Ownable} from "@openzeppelin/contracts/access/Ownable.sol";
import "./DoubleEndedStructQueue.sol";

contract AgentCommunication is Ownable {
address payable public treasury;
uint256 public pctToTreasuryInBasisPoints; //70% becomes 7000
interface IAgentRegistry {
error AgentNotRegistered();

event AgentRegistered(address indexed agent);
event AgentDeregistered(address indexed agent);

function registerAsAgent() external;
function deregisterAsAgent() external;
function isRegisteredAgent(address agent) external view returns (bool);
function getAllRegisteredAgents() external view returns (address[] memory);
}

contract AgentRegistry is IAgentRegistry, Ownable {
mapping(address => bool) public registeredAgents;
address[] private registeredAgentsList;
Comment on lines +20 to +21
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Consider using EnumerableMap (https://docs.openzeppelin.com/contracts/5.x/api/utils#EnumerableMap) instead of mapping + list, more efficient.


constructor() Ownable(msg.sender) {}

modifier onlyRegisteredAgent() {
if (!isRegisteredAgent(msg.sender)) {
revert AgentNotRegistered();
}
_;
}

function registerAsAgent() public {
registeredAgents[msg.sender] = true;
registeredAgentsList.push(msg.sender);
emit AgentRegistered(msg.sender);
}

error MessageNotSentByAgent();
function deregisterAsAgent() public onlyRegisteredAgent {
registeredAgents[msg.sender] = false;
// Remove from list
for (uint256 i = 0; i < registeredAgentsList.length; i++) {
if (registeredAgentsList[i] == msg.sender) {
registeredAgentsList[i] = registeredAgentsList[registeredAgentsList.length - 1];
registeredAgentsList.pop();
break;
}
}
emit AgentDeregistered(msg.sender);
}

function isRegisteredAgent(address agent) public view returns (bool) {
return registeredAgents[agent];
}

function getAllRegisteredAgents() public view returns (address[] memory) {
return registeredAgentsList;
}
}

contract AgentCommunication is AgentRegistry {
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Any reason for coupling these 2 contracts?
In my head, it would make more sense that AgentCommunication receives an AgentRegistry instance (actually the Interface) via constructor arg, and then operate on that.
This way, we keep the contracts independent of each other, since one handles registry and the other handles messaging. Wdty?

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I tried this approach initially but had some errors with it. I'll try it once more and post them here, maybe you will be able to help. (I like it like that more as well)

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Cool, happy to huddle or help async.
What I had in mind (hopefully works) is something like this

contract AgentRegistry {
[...]
}
contract AgentCommunication is Ownable {

IAgentRegistry public agentRegistry;

constructor(IAgentRegistry agentRegistry) Ownable(msg.sender) {
    agentRegistry = agentRegistry
}

// setter - onlyOwner

}

see also https://github.com/seer-pm/demo/blob/4310c1e2b995f675b9a8ef11f9a4f79ed044ecc4/contracts/src/trade/LiquidityManager.sol#L16

address payable public treasury;
uint256 public pctToTreasuryInBasisPoints; // 70% becomes 7000

mapping(address => DoubleEndedStructQueue.Bytes32Deque) public queues;

uint256 public minimumValueForSendingMessageInWei;

event LogMessage(address indexed sender, address indexed agentAddress, bytes message, uint256 value);

constructor(address payable _treasury, uint256 _pctToTreasuryInBasisPoints) Ownable(msg.sender) {
constructor(address payable _treasury, uint256 _pctToTreasuryInBasisPoints) {
treasury = _treasury;
pctToTreasuryInBasisPoints = _pctToTreasuryInBasisPoints;
minimumValueForSendingMessageInWei = 10000000000000; // 0.00001 xDAI
Expand All @@ -35,8 +86,8 @@ contract AgentCommunication is Ownable {
minimumValueForSendingMessageInWei = newValue;
}

function countMessages(address agentAddress) public view returns (uint256) {
return DoubleEndedStructQueue.length(queues[agentAddress]);
function countMessages() public view onlyRegisteredAgent returns (uint256) {
return DoubleEndedStructQueue.length(queues[msg.sender]);
}

// Private function to calculate the amounts
Expand All @@ -46,39 +97,43 @@ contract AgentCommunication is Ownable {
return (amountForTreasury, amountForAgent);
}

// We don't add `onlyRegisteredAgent` modifier here, because anyone should be able to send a message to an agent
Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Everything is guarded by onlyRegisteredAgent, except for this, because I guess we still want to be able to send messages from Streamlit UI to agents

function sendMessage(address agentAddress, bytes memory message) public payable mustPayMoreThanMinimum {
// split message value between treasury and agent
if (!isRegisteredAgent(agentAddress)) {
revert IAgentRegistry.AgentNotRegistered();
}

// Split message value between treasury and agent
(uint256 amountForTreasury, uint256 amountForAgent) = _calculateAmounts(msg.value);

// Transfer the amounts
(bool sentTreasury,) = treasury.call{value: amountForTreasury}("");
require(sentTreasury, "Failed to send Ether");
require(sentTreasury, "Failed to send Ether to treasury");
(bool sentAgent,) = payable(agentAddress).call{value: amountForAgent}("");
require(sentAgent, "Failed to send Ether");
require(sentAgent, "Failed to send Ether to agent");

DoubleEndedStructQueue.MessageContainer memory messageContainer =
DoubleEndedStructQueue.MessageContainer(msg.sender, agentAddress, message, msg.value);
DoubleEndedStructQueue.pushBack(queues[agentAddress], messageContainer);
emit LogMessage(msg.sender, agentAddress, messageContainer.message, msg.value);
}

function getAtIndex(address agentAddress, uint256 idx)
function getAtIndex(uint256 idx)
public
view
onlyRegisteredAgent
returns (DoubleEndedStructQueue.MessageContainer memory)
{
return DoubleEndedStructQueue.at(queues[agentAddress], idx);
return DoubleEndedStructQueue.at(queues[msg.sender], idx);
}
Comment on lines +121 to 128
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I would say reading should be allowed, as it's not modifying state. Messares are public.
Popping (on the other hand) should be guarded (as you implemented).


function popMessageAtIndex(address agentAddress, uint256 idx)
function popMessageAtIndex(uint256 idx)
public
onlyRegisteredAgent
returns (DoubleEndedStructQueue.MessageContainer memory)
{
if (msg.sender != agentAddress) {
revert MessageNotSentByAgent();
}
DoubleEndedStructQueue.MessageContainer memory message = DoubleEndedStructQueue.popAt(queues[agentAddress], idx);
emit LogMessage(message.sender, agentAddress, message.message, message.value);
DoubleEndedStructQueue.MessageContainer memory message = DoubleEndedStructQueue.popAt(queues[msg.sender], idx);
emit LogMessage(message.sender, msg.sender, message.message, message.value);
return message;
}
}
143 changes: 113 additions & 30 deletions test/AgentCommunication.t.sol
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,7 @@ pragma solidity ^0.8.22;

import {Ownable} from "@openzeppelin/contracts/access/Ownable.sol";
import "forge-std/Test.sol";
import {AgentCommunication} from "../src/NFT/AgentCommunication.sol";
import {AgentRegistry, AgentCommunication, IAgentRegistry} from "../src/NFT/AgentCommunication.sol";
import "../src/NFT/DoubleEndedStructQueue.sol";

contract AgentCommunicationTest is Test {
Expand Down Expand Up @@ -59,10 +59,48 @@ contract AgentCommunicationTest is Test {
assertEq(agentComm.minimumValueForSendingMessageInWei(), 10000000000000);
}

function testSendMessage() public {
function testRegisterAsAgent() public {
vm.startPrank(agent);

// Expect AgentRegistered event
vm.expectEmit(true, false, false, false);
emit IAgentRegistry.AgentRegistered(agent);

agentComm.registerAsAgent();
vm.stopPrank();

assertTrue(agentComm.registeredAgents(agent), "Agent should be registered");
}

function testDeregisterAsAgent() public {
vm.startPrank(agent);
agentComm.registerAsAgent();

// Expect AgentDeregistered event
vm.expectEmit(true, false, false, false);
emit IAgentRegistry.AgentDeregistered(agent);

agentComm.deregisterAsAgent();
vm.stopPrank();

assertFalse(agentComm.registeredAgents(agent), "Agent should be deregistered");
}

function testDeregisterAsAgentWhenNotRegistered() public {
vm.startPrank(agent);
vm.expectRevert(IAgentRegistry.AgentNotRegistered.selector);
agentComm.deregisterAsAgent();
vm.stopPrank();
}

function testSendMessageToRegisteredAgent() public {
DoubleEndedStructQueue.MessageContainer memory message = buildMessage("Hello!");
vm.deal(owner, 1 ether);

// Register the agent first
vm.prank(agent);
agentComm.registerAsAgent();

// Record initial balances
uint256 initialBalanceTreasury = treasury.balance;
uint256 initialBalanceAgent = agent.balance;
Expand All @@ -81,10 +119,21 @@ contract AgentCommunicationTest is Test {
assertEq(messageValue * (10000 - pctToTreasuryInBasisPoints) / 10000, diffBalanceAgent);
assertEq(address(agentComm).balance, 0);

DoubleEndedStructQueue.MessageContainer memory storedMessage = agentComm.getAtIndex(agent, 0);
vm.prank(agent);
DoubleEndedStructQueue.MessageContainer memory storedMessage = agentComm.getAtIndex(0);
assertEq(storedMessage.message, message.message);
}

function testSendMessageToUnregisteredAgent() public {
DoubleEndedStructQueue.MessageContainer memory message = buildMessage("Hello!");
vm.deal(owner, 1 ether);

vm.startPrank(owner);
vm.expectRevert(IAgentRegistry.AgentNotRegistered.selector);
agentComm.sendMessage{value: 10000000000000}(agent, message.message);
vm.stopPrank();
}

function testSendMessageInsufficientValue() public {
DoubleEndedStructQueue.MessageContainer memory message = buildMessage("Hello!");
vm.deal(agent, 1 ether);
Expand All @@ -97,14 +146,19 @@ contract AgentCommunicationTest is Test {
function testNewMessageSentEvent() public {
DoubleEndedStructQueue.MessageContainer memory message = buildMessage("Hello!");
vm.deal(agent, 1 ether);

// Register the recipient
vm.prank(message.recipient);
agentComm.registerAsAgent();

vm.startPrank(agent);

// Expect the LogMessage event to be emitted
vm.expectEmit(true, true, true, true);
emit AgentCommunication.LogMessage(message.sender, message.recipient, message.message, message.value);

// Send the message
agentComm.sendMessage{value: message.value}(address(0x789), message.message);
agentComm.sendMessage{value: message.value}(message.recipient, message.message);
vm.stopPrank();
}

Expand All @@ -116,6 +170,18 @@ contract AgentCommunicationTest is Test {
DoubleEndedStructQueue.MessageContainer memory message_4 = buildMessage("Hello 4!");
DoubleEndedStructQueue.MessageContainer memory message_5 = buildMessage("Hello 5!");

// Register the recipients
vm.prank(message_1.recipient);
agentComm.registerAsAgent();
vm.prank(message_2.recipient);
agentComm.registerAsAgent();
vm.prank(message_3.recipient);
agentComm.registerAsAgent();
vm.prank(message_4.recipient);
agentComm.registerAsAgent();
vm.prank(message_5.recipient);
agentComm.registerAsAgent();

// Fund the agent and start the prank
vm.deal(agent, 5 ether);
vm.startPrank(agent);
Expand All @@ -136,8 +202,7 @@ contract AgentCommunicationTest is Test {
emit AgentCommunication.LogMessage(message_1.sender, message_1.recipient, message_1.message, message_1.value);

// Pop the next message
DoubleEndedStructQueue.MessageContainer memory poppedMessage_1 =
agentComm.popMessageAtIndex(message_1.recipient, 0);
DoubleEndedStructQueue.MessageContainer memory poppedMessage_1 = agentComm.popMessageAtIndex(0);
vm.stopPrank();

// Assert that the popped message matches the original message
Expand All @@ -147,12 +212,9 @@ contract AgentCommunicationTest is Test {
assertEq(poppedMessage_1.value, message_1.value);

// Start the prank again for popping another message
vm.startPrank(message_1.recipient);

vm.prank(message_1.recipient);
// Pop the message at specified index
DoubleEndedStructQueue.MessageContainer memory poppedMessage_2 =
agentComm.popMessageAtIndex(message_1.recipient, 2);
vm.stopPrank();
DoubleEndedStructQueue.MessageContainer memory poppedMessage_2 = agentComm.popMessageAtIndex(2);

// Assert that the popped message matches the original message
assertEq(poppedMessage_2.sender, message_4.sender);
Expand All @@ -161,26 +223,15 @@ contract AgentCommunicationTest is Test {
assertEq(poppedMessage_2.value, message_4.value);
}

function testPopMessageNotByAgent() public {
DoubleEndedStructQueue.MessageContainer memory message = buildMessage("Hello!");
vm.deal(agent, 1 ether);
vm.startPrank(agent);
agentComm.sendMessage{value: 10000000000000}(agent, message.message);
vm.stopPrank();

address notAgent = address(0x789);
vm.startPrank(notAgent);
vm.expectRevert(AgentCommunication.MessageNotSentByAgent.selector);
agentComm.popMessageAtIndex(agent, 0);
vm.stopPrank();
}

function testCountMessages() public {
// Initialize a message
DoubleEndedStructQueue.MessageContainer memory message1 = buildMessage("Hello!");

DoubleEndedStructQueue.MessageContainer memory message2 = buildMessage("Hello!");

// Register agent
vm.prank(agent);
agentComm.registerAsAgent();

// Fund the agent and start the prank
vm.deal(agent, 1 ether);
vm.startPrank(agent);
Expand All @@ -189,11 +240,8 @@ contract AgentCommunicationTest is Test {
agentComm.sendMessage{value: 10000000000000}(agent, message1.message);
agentComm.sendMessage{value: 10000000000000}(agent, message2.message);

// Stop the prank
vm.stopPrank();

// Check the count of messages
uint256 messageCount = agentComm.countMessages(agent);
uint256 messageCount = agentComm.countMessages();
assertEq(messageCount, 2, "The message count should be 2");
}

Expand All @@ -218,4 +266,39 @@ contract AgentCommunicationTest is Test {
// Verify that the treasury address has not changed
assertEq(address(agentComm.treasury()), address(treasury));
}

function testGetAllRegisteredAgents() public {
// Initially there should be no registered agents
address[] memory initialAgents = agentComm.getAllRegisteredAgents();
assertEq(initialAgents.length, 0, "Should start with no registered agents");

// Register first agent
vm.prank(agent);
agentComm.registerAsAgent();

// Register second agent
address agent2 = address(0xaaa);
vm.prank(agent2);
agentComm.registerAsAgent();

// Get the list of registered agents
address[] memory registeredAgents = agentComm.getAllRegisteredAgents();

// Verify the list contains both agents
assertEq(registeredAgents.length, 2, "Should have two registered agents");
assertTrue(
(registeredAgents[0] == agent && registeredAgents[1] == agent2)
|| (registeredAgents[0] == agent2 && registeredAgents[1] == agent),
"List should contain both registered agents"
);

// Deregister one agent
vm.prank(agent);
agentComm.deregisterAsAgent();

// Verify the list now only contains one agent
address[] memory remainingAgents = agentComm.getAllRegisteredAgents();
assertEq(remainingAgents.length, 1, "Should have one registered agent");
assertEq(remainingAgents[0], agent2, "Should contain the remaining agent");
}
}