Merge branch 'feat/multiple-receivers' into 'dev'

feat: able to send to multiple receivers

See merge request ics/sp/2025/n25b/pepperplus-cb!42

src/control_backend/core/agent_system.py

@@ -60,6 +60,9 @@ class BaseAgent(ABC):
         self._tasks: set[asyncio.Task] = set()
         self._running = False
 
+        self._internal_pub_socket: None | azmq.Socket = None
+        self._internal_sub_socket: None | azmq.Socket = None
+
         # Register immediately
         AgentDirectory.register(name, self)
 
@@ -130,16 +133,22 @@ class BaseAgent(ABC):
 
         :param message: The message to send.
         """
-        target = AgentDirectory.get(message.to)
+        message.sender = self.name
-        if target:
+        to = message.to
-            await target.inbox.put(message)
+        receivers = [to] if isinstance(to, str) else to
-            self.logger.debug(f"Sent message {message.body} to {message.to} via regular inbox.")
+
-        else:
+        for receiver in receivers:
-            # Apparently target agent is on a different process, send via ZMQ
+            target = AgentDirectory.get(receiver)
-            topic = f"internal/{message.to}".encode()
+
-            body = message.model_dump_json().encode()
+            if target:
-            await self._internal_pub_socket.send_multipart([topic, body])
+                await target.inbox.put(message)
-            self.logger.debug(f"Sent message {message.body} to {message.to} via ZMQ.")
+                self.logger.debug(f"Sent message {message.body} to {message.to} via regular inbox.")
+            else:
+                # Apparently target agent is on a different process, send via ZMQ
+                topic = f"internal/{receiver}".encode()
+                body = message.model_dump_json().encode()
+                await self._internal_pub_socket.send_multipart([topic, body])
+                self.logger.debug(f"Sent message {message.body} to {message.to} via ZMQ.")
 
     async def _process_inbox(self):
         """

src/control_backend/schemas/internal_message.py

@@ -1,3 +1,5 @@
+from collections.abc import Iterable
+
 from pydantic import BaseModel
 
 
@@ -5,13 +7,13 @@ class InternalMessage(BaseModel):
     """
     Standard message envelope for communication between agents within the Control Backend.
 
-    :ivar to: The name of the destination agent.
+    :ivar to: The name(s) of the destination agent(s).
     :ivar sender: The name of the sending agent.
     :ivar body: The string payload (often a JSON-serialized model).
     :ivar thread: An optional thread identifier/topic to categorize the message (e.g., 'beliefs').
     """
 
-    to: str
+    to: str | Iterable[str]
     sender: str
     body: str
     thread: str | None = None

test/unit/core/test_agent_system.py

@@ -99,12 +99,75 @@ async def test_send_to_local_agent(monkeypatch):
     # Patch inbox.put
     target.inbox.put = AsyncMock()
 
-    message = InternalMessage(to="receiver", sender="sender", body="hello")
+    message = InternalMessage(to=target.name, sender=sender.name, body="hello")
 
     await sender.send(message)
 
     target.inbox.put.assert_awaited_once_with(message)
-    sender.logger.debug.assert_called_once()
+
+
+@pytest.mark.asyncio
+async def test_send_to_zmq_agent(monkeypatch):
+    sender = DummyAgent("sender")
+    target = "remote_receiver"
+
+    # Fake logger
+    sender.logger = MagicMock()
+
+    # Fake zmq
+    sender._internal_pub_socket = AsyncMock()
+
+    message = InternalMessage(to=target, sender=sender.name, body="hello")
+
+    await sender.send(message)
+
+    zmq_calls = sender._internal_pub_socket.send_multipart.call_args[0][0]
+    assert zmq_calls[0] == f"internal/{target}".encode()
+
+
+@pytest.mark.asyncio
+async def test_send_to_multiple_local_agents(monkeypatch):
+    sender = DummyAgent("sender")
+    target1 = DummyAgent("receiver1")
+    target2 = DummyAgent("receiver2")
+
+    # Fake logger
+    sender.logger = MagicMock()
+
+    # Patch inbox.put
+    target1.inbox.put = AsyncMock()
+    target2.inbox.put = AsyncMock()
+
+    message = InternalMessage(to=[target1.name, target2.name], sender=sender.name, body="hello")
+
+    await sender.send(message)
+
+    target1.inbox.put.assert_awaited_once_with(message)
+    target2.inbox.put.assert_awaited_once_with(message)
+
+
+@pytest.mark.asyncio
+async def test_send_to_multiple_agents(monkeypatch):
+    sender = DummyAgent("sender")
+    target1 = DummyAgent("receiver1")
+    target2 = "remote_receiver"
+
+    # Fake logger
+    sender.logger = MagicMock()
+
+    # Fake zmq
+    sender._internal_pub_socket = AsyncMock()
+
+    # Patch inbox.put
+    target1.inbox.put = AsyncMock()
+
+    message = InternalMessage(to=[target1.name, target2], sender=sender.name, body="hello")
+
+    await sender.send(message)
+
+    target1.inbox.put.assert_awaited_once_with(message)
+    zmq_calls = sender._internal_pub_socket.send_multipart.call_args[0][0]
+    assert zmq_calls[0] == f"internal/{target2}".encode()
 
 
 @pytest.mark.asyncio