agent-relay 3.2.14 → 3.2.16

package/packages/sdk-py/tests/communicate/conftest.py

@@ -16,20 +16,7 @@ from agent_relay.communicate.types import RelayConfig
 
 
 class MockRelayServer:
-    """Minimal in-process Relaycast mock that mirrors the real API surface.
-
-    Endpoints:
-        POST   /v1/agents                          — register (workspace key auth)
-        POST   /v1/agents/disconnect                — unregister (agent token auth)
-        POST   /v1/agents/heartbeat                 — presence heartbeat (agent token)
-        GET    /v1/agents                           — list agents (workspace key auth)
-        POST   /v1/dm                               — send DM (agent token auth)
-        POST   /v1/channels/{name}/messages         — channel message (agent token auth)
-        POST   /v1/messages/{id}/replies            — reply (agent token auth)
-        GET    /v1/inbox                            — inbox summary (agent token auth)
-        GET    /v1/dm/{conv_id}/messages            — DM conversation messages (agent token)
-        GET    /v1/ws                               — websocket (agent token via query param)
-    """
+    """Minimal in-process Relaycast mock for transport and integration tests."""
 
     def __init__(self, *, api_key: str = "test-key", workspace: str = "test-workspace") -> None:
         self.api_key = api_key
@@ -44,27 +31,20 @@ class MockRelayServer:
         self.received_ws_messages: list[dict[str, Any]] = []
         self.ws_connection_counts: dict[str, int] = defaultdict(int)
 
-        # Map agent tokens → agent IDs for auth
-        self._token_to_agent_id: dict[str, str] = {}
-        # DM conversations: conv_id → list of messages
-        self._dm_conversations: dict[str, list[dict[str, Any]]] = defaultdict(list)
-
         self._active_websockets: dict[str, web.WebSocketResponse] = {}
         self._queued_errors: dict[str, deque[tuple[int, dict[str, Any]]]] = defaultdict(deque)
         self._agent_ids = count(1)
         self._message_ids = count(1)
 
         self._app = web.Application()
-        self._app.router.add_post("/v1/agents", self._handle_register)
+        self._app.router.add_post("/v1/agents/register", self._handle_register)
+        self._app.router.add_delete("/v1/agents/{agent_id}", self._handle_unregister)
+        self._app.router.add_post("/v1/messages/dm", self._handle_dm)
+        self._app.router.add_post("/v1/messages/channel", self._handle_channel)
+        self._app.router.add_post("/v1/messages/reply", self._handle_reply)
+        self._app.router.add_get("/v1/inbox/{agent_id}", self._handle_inbox)
         self._app.router.add_get("/v1/agents", self._handle_agents)
-        self._app.router.add_post("/v1/agents/disconnect", self._handle_disconnect)
-        self._app.router.add_post("/v1/agents/heartbeat", self._handle_heartbeat)
-        self._app.router.add_post("/v1/dm", self._handle_dm)
-        self._app.router.add_get("/v1/dm/{conv_id}/messages", self._handle_dm_messages)
-        self._app.router.add_post("/v1/channels/{channel}/messages", self._handle_channel)
-        self._app.router.add_post("/v1/messages/{message_id}/replies", self._handle_reply)
-        self._app.router.add_get("/v1/inbox", self._handle_inbox)
-        self._app.router.add_get("/v1/ws", self._handle_ws)
+        self._app.router.add_get("/v1/ws/{agent_id}", self._handle_ws)
 
         self._runner: web.AppRunner | None = None
         self._site: web.TCPSite | None = None
@@ -89,7 +69,7 @@ class MockRelayServer:
         message: str,
         repeat: int = 1,
     ) -> None:
-        body = {"ok": False, "error": {"code": "error", "message": message}}
+        body = {"message": message}
         for _ in range(repeat):
             self._queued_errors[operation].append((status, body))
 
@@ -121,30 +101,15 @@ class MockRelayServer:
         message_id: str | None = None,
         timestamp: float | None = None,
     ) -> dict[str, Any]:
-        msg_id = message_id or f"message-{next(self._message_ids)}"
         payload = {
             "sender": sender,
             "text": text,
             "channel": channel,
             "thread_id": thread_id,
-            "message_id": msg_id,
+            "message_id": message_id or f"message-{next(self._message_ids)}",
             "timestamp": timestamp,
         }
-        # Store as a DM conversation entry for inbox retrieval
-        conv_id = f"dm_{agent_id}_{sender}"
-        dm_msg = {
-            "id": msg_id,
-            "agent_name": sender,
-            "text": text,
-            "created_at": str(timestamp) if timestamp else None,
-        }
-        self._dm_conversations[conv_id].append(dm_msg)
-        self.inboxes[agent_id].append({
-            "conversation_id": conv_id,
-            "from": sender,
-            "unread_count": 1,
-            "last_message": dm_msg,
-        })
+        self.inboxes[agent_id].append(payload)
         return payload
 
     async def push_ws_message(
@@ -162,19 +127,15 @@ class MockRelayServer:
         if ws is None or ws.closed:
             raise AssertionError(f"No active websocket for agent {agent_id!r}")
 
-        payload: dict[str, Any] = {
-            "type": "dm.received" if channel is None else "message.created",
-            "agent_name": sender,
+        payload = {
+            "type": "message",
+            "sender": sender,
             "text": text,
-            "id": message_id or f"message-{next(self._message_ids)}",
+            "channel": channel,
+            "thread_id": thread_id,
+            "message_id": message_id or f"message-{next(self._message_ids)}",
+            "timestamp": timestamp,
         }
-        if channel is not None:
-            payload["channel_name"] = channel
-        if thread_id is not None:
-            payload["thread_id"] = thread_id
-        if timestamp is not None:
-            payload["created_at"] = str(timestamp)
-
         await ws.send_json(payload)
         return payload
 
@@ -228,195 +189,130 @@ class MockRelayServer:
         if self._runner is not None:
             await self._runner.cleanup()
 
-    # --- Handlers ---
-
     async def _handle_register(self, request: web.Request) -> web.StreamResponse:
         payload = await request.json()
         self._record_request("register_agent", request, payload)
 
         if error := self._pop_error("register_agent"):
             return error
-        if error := self._require_workspace_auth(request):
+        if error := self._require_auth(request):
             return error
 
         agent_id = f"agent-{next(self._agent_ids)}"
         token = f"token-{agent_id}"
         self.registered_agents[agent_id] = {"name": payload["name"], "token": token}
-        self._token_to_agent_id[token] = agent_id
-        return web.json_response({
-            "ok": True,
-            "data": {"id": agent_id, "name": payload["name"], "token": token, "status": "online"},
-        })
-
-    async def _handle_disconnect(self, request: web.Request) -> web.StreamResponse:
-        agent_id = self._resolve_agent_from_token(request)
+        return web.json_response({"agent_id": agent_id, "token": token})
+
+    async def _handle_unregister(self, request: web.Request) -> web.StreamResponse:
+        agent_id = request.match_info["agent_id"]
         self._record_request("unregister_agent", request, {"agent_id": agent_id})
 
         if error := self._pop_error("unregister_agent"):
             return error
-        if agent_id is None:
-            return web.json_response({"ok": False, "error": {"message": "Unauthorized"}}, status=401)
+        if error := self._require_auth(request):
+            return error
 
         self.registered_agents.pop(agent_id, None)
         self.inboxes.pop(agent_id, None)
         ws = self._active_websockets.pop(agent_id, None)
         if ws is not None and not ws.closed:
             await ws.close()
-        # Remove token mapping
-        token = self._extract_token(request)
-        if token:
-            self._token_to_agent_id.pop(token, None)
-        return web.json_response({"ok": True})
-
-    async def _handle_heartbeat(self, request: web.Request) -> web.StreamResponse:
-        return web.json_response({"ok": True})
+        return web.Response(status=204)
 
     async def _handle_dm(self, request: web.Request) -> web.StreamResponse:
         payload = await request.json()
-        agent_id = self._resolve_agent_from_token(request)
-        sender_name = self.registered_agents.get(agent_id, {}).get("name", "unknown") if agent_id else "unknown"
         self._record_request("send_dm", request, payload)
 
         if error := self._pop_error("send_dm"):
             return error
-        if agent_id is None:
-            return web.json_response({"ok": False, "error": {"message": "Unauthorized"}}, status=401)
+        if error := self._require_auth(request):
+            return error
 
         message_id = f"message-{next(self._message_ids)}"
-        conv_id = f"dm_{message_id}"
-        self.messages.append({"kind": "dm", "message_id": message_id, **payload, "from": sender_name})
-
-        # Deliver to recipient
-        recipient_ids = self.find_agent_ids(payload["to"])
-        dm_msg = {
-            "id": message_id,
-            "agent_name": sender_name,
-            "text": payload["text"],
-            "created_at": None,
-        }
-        self._dm_conversations[conv_id].append(dm_msg)
-        for rid in recipient_ids:
-            self.inboxes[rid].append({
-                "conversation_id": conv_id,
-                "from": sender_name,
-                "unread_count": 1,
-                "last_message": dm_msg,
-            })
-
+        self.messages.append({"kind": "dm", "message_id": message_id, **payload})
         await self._deliver_to_agents(
-            recipient_ids,
-            sender=sender_name,
+            self.find_agent_ids(payload["to"]),
+            sender=payload["from"],
             text=payload["text"],
             message_id=message_id,
         )
-
-        return web.json_response({
-            "ok": True,
-            "data": {"id": message_id, "conversation_id": conv_id, "text": payload["text"]},
-        }, status=201)
-
-    async def _handle_dm_messages(self, request: web.Request) -> web.StreamResponse:
-        conv_id = request.match_info["conv_id"]
-        messages = list(self._dm_conversations.get(conv_id, []))
-        return web.json_response({"ok": True, "data": messages})
+        return web.json_response({"message_id": message_id})
 
     async def _handle_channel(self, request: web.Request) -> web.StreamResponse:
-        channel = request.match_info["channel"]
         payload = await request.json()
-        agent_id = self._resolve_agent_from_token(request)
-        sender_name = self.registered_agents.get(agent_id, {}).get("name", "unknown") if agent_id else "unknown"
-        self._record_request("post_message", request, {**payload, "channel": channel})
+        self._record_request("post_message", request, payload)
 
         if error := self._pop_error("post_message"):
             return error
-        if agent_id is None:
-            return web.json_response({"ok": False, "error": {"message": "Unauthorized"}}, status=401)
+        if error := self._require_auth(request):
+            return error
 
         message_id = f"message-{next(self._message_ids)}"
-        self.messages.append({"kind": "channel", "message_id": message_id, "channel": channel, **payload, "from": sender_name})
+        self.messages.append({"kind": "channel", "message_id": message_id, **payload})
         await self._deliver_to_agents(
             [
-                aid
-                for aid, registration in self.registered_agents.items()
-                if aid != agent_id
+                agent_id
+                for agent_id, registration in self.registered_agents.items()
+                if registration["name"] != payload["from"]
             ],
-            sender=sender_name,
+            sender=payload["from"],
             text=payload["text"],
-            channel=channel,
+            channel=payload["channel"],
             message_id=message_id,
         )
-        return web.json_response({
-            "ok": True,
-            "data": {"id": message_id, "channel_name": channel, "text": payload["text"]},
-        }, status=201)
+        return web.json_response({"message_id": message_id})
 
     async def _handle_reply(self, request: web.Request) -> web.StreamResponse:
-        parent_id = request.match_info["message_id"]
         payload = await request.json()
-        agent_id = self._resolve_agent_from_token(request)
-        sender_name = self.registered_agents.get(agent_id, {}).get("name", "unknown") if agent_id else "unknown"
-        self._record_request("reply", request, {**payload, "message_id": parent_id})
+        self._record_request("reply", request, payload)
 
         if error := self._pop_error("reply"):
             return error
-        if agent_id is None:
-            return web.json_response({"ok": False, "error": {"message": "Unauthorized"}}, status=401)
+        if error := self._require_auth(request):
+            return error
 
         reply_id = f"message-{next(self._message_ids)}"
-        self.messages.append({
-            "kind": "reply",
-            "reply_id": reply_id,
-            "thread_id": parent_id,
-            "from": sender_name,
-            **payload,
-        })
-        return web.json_response({
-            "ok": True,
-            "data": {"id": reply_id, "text": payload["text"]},
-        }, status=201)
+        self.messages.append(
+            {
+                "kind": "reply",
+                "reply_id": reply_id,
+                "thread_id": payload.get("thread_id") or payload.get("message_id"),
+                **payload,
+            }
+        )
+        return web.json_response({"message_id": reply_id})
 
     async def _handle_inbox(self, request: web.Request) -> web.StreamResponse:
-        agent_id = self._resolve_agent_from_token(request)
+        agent_id = request.match_info["agent_id"]
         self._record_request("check_inbox", request, {"agent_id": agent_id})
 
         if error := self._pop_error("check_inbox"):
             return error
-        if agent_id is None:
-            return web.json_response({"ok": False, "error": {"message": "Unauthorized"}}, status=401)
-
-        unread_dms = list(self.inboxes.get(agent_id, []))
-        self.inboxes[agent_id] = []
-        return web.json_response({
-            "ok": True,
-            "data": {
-                "unread_channels": [],
-                "mentions": [],
-                "unread_dms": unread_dms,
-                "recent_reactions": [],
-            },
-        })
+        if error := self._require_auth(request):
+            return error
+
+        messages = list(self.inboxes[agent_id])
+        self.inboxes[agent_id].clear()
+        return web.json_response({"messages": messages})
 
     async def _handle_agents(self, request: web.Request) -> web.StreamResponse:
         self._record_request("list_agents", request, None)
 
         if error := self._pop_error("list_agents"):
             return error
-        if error := self._require_workspace_auth(request):
+        if error := self._require_auth(request):
             return error
 
-        agent_list = [
-            {"name": agent["name"], "id": agent_id, "status": "online"}
-            for agent_id, agent in self.registered_agents.items()
-        ]
-        for name in self.extra_agents:
-            agent_list.append({"name": name, "id": f"extra-{name}", "status": "online"})
-
-        return web.json_response({"ok": True, "data": agent_list})
+        agents = sorted(
+            {agent["name"] for agent in self.registered_agents.values()} | self.extra_agents
+        )
+        return web.json_response({"agents": agents})
 
     async def _handle_ws(self, request: web.Request) -> web.StreamResponse:
+        agent_id = request.match_info["agent_id"]
         token = request.query.get("token")
-        agent_id = self._token_to_agent_id.get(token) if token else None
-        if agent_id is None or agent_id not in self.registered_agents:
+        registration = self.registered_agents.get(agent_id)
+        if registration is None or token != registration["token"]:
             raise web.HTTPUnauthorized(text="Invalid websocket token")
 
         ws = web.WebSocketResponse()
@@ -437,8 +333,6 @@ class MockRelayServer:
 
         return ws
 
-    # --- Helpers ---
-
     def _record_request(
         self,
         operation: str,
@@ -460,23 +354,11 @@ class MockRelayServer:
         status, body = self._queued_errors[operation].popleft()
         return web.json_response(body, status=status)
 
-    def _extract_token(self, request: web.Request) -> str | None:
-        auth = request.headers.get("Authorization", "")
-        if auth.startswith("Bearer "):
-            return auth[7:]
-        return None
-
-    def _require_workspace_auth(self, request: web.Request) -> web.StreamResponse | None:
-        token = self._extract_token(request)
-        if token == self.api_key:
+    def _require_auth(self, request: web.Request) -> web.StreamResponse | None:
+        auth_header = request.headers.get("Authorization")
+        if auth_header == f"Bearer {self.api_key}":
             return None
-        return web.json_response({"ok": False, "error": {"message": "Unauthorized"}}, status=401)
-
-    def _resolve_agent_from_token(self, request: web.Request) -> str | None:
-        token = self._extract_token(request)
-        if token is None:
-            return None
-        return self._token_to_agent_id.get(token)
+        return web.json_response({"message": "Unauthorized"}, status=401)
 
     async def _deliver_to_agents(
         self,
@@ -492,57 +374,28 @@ class MockRelayServer:
         if not agent_ids:
             return
 
+        payload = {
+            "sender": sender,
+            "text": text,
+            "channel": channel,
+            "thread_id": thread_id,
+            "message_id": message_id or f"message-{next(self._message_ids)}",
+            "timestamp": timestamp,
+        }
+
         for agent_id in agent_ids:
-            await self._deliver_message(agent_id, sender=sender, text=text,
-                                         channel=channel, thread_id=thread_id,
-                                         message_id=message_id, timestamp=timestamp)
+            await self._deliver_message(agent_id, payload)
 
-    async def _deliver_message(
-        self,
-        agent_id: str,
-        *,
-        sender: str,
-        text: str,
-        channel: str | None = None,
-        thread_id: str | None = None,
-        message_id: str | None = None,
-        timestamp: float | None = None,
-    ) -> None:
+    async def _deliver_message(self, agent_id: str, payload: dict[str, Any]) -> None:
         ws = self._active_websockets.get(agent_id)
         if ws is not None and not ws.closed:
-            payload: dict[str, Any] = {
-                "type": "dm.received" if channel is None else "message.created",
-                "agent_name": sender,
-                "text": text,
-                "id": message_id or f"message-{next(self._message_ids)}",
-            }
-            if channel is not None:
-                payload["channel_name"] = channel
-            if thread_id is not None:
-                payload["thread_id"] = thread_id
-            if timestamp is not None:
-                payload["created_at"] = str(timestamp)
             try:
-                await ws.send_json(payload)
+                await ws.send_json({"type": "message", **payload})
                 return
             except Exception:
                 self._active_websockets.pop(agent_id, None)
 
-        # Fall back to inbox buffering
-        conv_id = f"dm_{sender}_{agent_id}"
-        dm_msg = {
-            "id": message_id or f"message-{next(self._message_ids)}",
-            "agent_name": sender,
-            "text": text,
-            "created_at": str(timestamp) if timestamp else None,
-        }
-        self._dm_conversations[conv_id].append(dm_msg)
-        self.inboxes[agent_id].append({
-            "conversation_id": conv_id,
-            "from": sender,
-            "unread_count": 1,
-            "last_message": dm_msg,
-        })
+        self.inboxes[agent_id].append(dict(payload))
 
 
 @pytest_asyncio.fixture

package/packages/sdk-py/tests/communicate/integration/test_cross_framework.py

@@ -218,7 +218,7 @@ async def test_swarms_sender_reaches_claude_sdk_hook_system_message(relay_server
     receiver_options = SimpleNamespace(mcp_servers=[], hooks=Hooks())
 
     swarms_adapter.on_relay(sender_agent, sender_relay)
-    claude_adapter.on_relay("ClaudeReceiver", receiver_options, relay=receiver_relay)
+    claude_adapter.on_relay(receiver_options, relay=receiver_relay, name="ClaudeReceiver")
 
     try:
         await _prime_claude_receiver(
@@ -266,7 +266,7 @@ async def test_multiple_framework_agents_post_to_the_same_channel(relay_server,
     openai_adapter.on_relay(openai_agent, openai_relay)
     google_adapter.on_relay(google_agent, google_relay)
     swarms_adapter.on_relay(swarms_agent, swarms_relay)
-    claude_adapter.on_relay("ClaudeChannelReader", receiver_options, relay=receiver_relay)
+    claude_adapter.on_relay(receiver_options, relay=receiver_relay, name="ClaudeChannelReader")
 
     try:
         await _prime_claude_receiver(

package/packages/sdk-py/tests/communicate/integration/test_end_to_end.py

@@ -64,30 +64,16 @@ async def _wait_for_agent_absent(
     transport: RelayTransport,
     agent_name: str,
     *,
-    timeout: float = 30.0,
+    timeout: float = 15.0,
 ) -> None:
-    """Wait for an agent to disappear from list_agents or go offline.
-
-    Relaycast presence may take a few seconds to propagate after disconnect.
-    We check both absence from the list and offline status.
-    """
     deadline = asyncio.get_running_loop().time() + timeout
 
     while asyncio.get_running_loop().time() < deadline:
-        agents_payload = await transport.send_http("GET", "/v1/agents")
-        data = agents_payload.get("data", agents_payload)
-        if isinstance(data, list):
-            matching = [a for a in data if isinstance(a, dict) and a.get("name") == agent_name]
-            if not matching:
-                return
-            # Also accept "offline" status
-            if all(a.get("status") == "offline" for a in matching):
-                return
-        elif agent_name not in await transport.list_agents():
+        if agent_name not in await transport.list_agents():
             return
-        await asyncio.sleep(0.5)
+        await asyncio.sleep(0.25)
 
-    raise AssertionError(f"Timed out waiting for {agent_name!r} to go offline/absent from list_agents().")
+    raise AssertionError(f"Timed out waiting for {agent_name!r} to disappear from list_agents().")
 
 
 @pytest.mark.asyncio
@@ -110,14 +96,9 @@ async def test_register_send_receive_inbox_and_unregister_round_trip():
         assert message.text == text
         assert message.sender == sender.agent_name
 
-        # Verify disconnect completes without error
         receiver_name = receiver.agent_name
         await receiver.disconnect()
-
-        # Note: Relaycast presence updates are eventually consistent —
-        # agents may remain "online" in list_agents for a heartbeat window
-        # after disconnect. We verify the disconnect call succeeds rather
-        # than waiting for presence propagation.
+        await _wait_for_agent_absent(probe, receiver_name)
     finally:
         await asyncio.gather(
             _safe_disconnect(sender),
@@ -125,6 +106,15 @@ async def test_register_send_receive_inbox_and_unregister_round_trip():
             _safe_disconnect(probe),
         )
 
+    cleanup_probe = RelayTransport(_unique_name("sdk-py-e2e-cleanup"), config)
+    await cleanup_probe.connect()
+    try:
+        agents = await cleanup_probe.list_agents()
+        assert sender.agent_name not in agents
+        assert receiver.agent_name not in agents
+    finally:
+        await _safe_disconnect(cleanup_probe)
+
 
 @pytest.mark.asyncio
 async def test_two_agents_can_exchange_bidirectional_messages():