diff --git a/tests/test_routes_a2a_bus_stream.py b/tests/test_routes_a2a_bus_stream.py new file mode 100644 index 000000000..fa079ca40 --- /dev/null +++ b/tests/test_routes_a2a_bus_stream.py @@ -0,0 +1,201 @@ +"""Tests for the authenticated A2A SSE stream proxy (Slice S3). + +Covers: an agent-JWT (bound to a project) can open the stream and receives +forwarded SSE frames relayed from the raw bus; the messages proxy forwards the +`since` cursor; an unauthenticated request is rejected 401; and the raw :7900 +bus is never exposed directly (the upstream call is made by the proxy only). +""" +from __future__ import annotations + +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest +import pytest_asyncio +from httpx import ASGITransport, AsyncClient + +from tinyagentos.agent_registry_store import mint_registry_token + + +async def _mint_agent(app, *, scopes=("a2a_receive",), project_id="prj-1"): + """Register an active agent with *scopes*, return (canonical_id, jwt).""" + registry = app.state.agent_registry + grants = app.state.agent_grants + priv, _pub = app.state.agent_registry_keypair + + rec = await registry.register( + framework="grok", + display_name="Grok", + origin="external-selfjoin", + handle="@grok", + ) + cid = rec["canonical_id"] + await registry.set_status(cid, "active") + for scope in scopes: + await grants.add_grant(cid, scope, project_id=project_id) + token = mint_registry_token( + cid, priv, user_id="u", framework="grok", project_id=project_id + ) + return cid, token + + +@pytest_asyncio.fixture +async def agent_app(app): + for attr in ("agent_registry", "agent_grants"): + store = getattr(app.state, attr) + if store._db is None: # noqa: SLF001 + await store.init() + yield app + for attr in ("agent_registry", "agent_grants"): + store = getattr(app.state, attr) + if store._db is not None: + await store.close() + + +@pytest_asyncio.fixture +async def noauth_client(app): + """An httpx client with NO session cookie, so requests are unauthenticated.""" + transport = ASGITransport(app=app) + async with AsyncClient(transport=transport, base_url="http://test") as c: + yield c + + +def _fake_sse_lines(lines): + """Return an async iterator over *lines* usable as aiter_lines().""" + + async def _gen(): + for ln in lines: + yield ln + + return _gen() + + +@pytest.mark.asyncio +class TestBusStreamProxy: + async def test_agent_jwt_receives_forwarded_sse_frames(self, agent_app, client): + _, token = await _mint_agent(agent_app, scopes=("a2a_receive",)) + sse_body = [ + "event: message", + 'data: {"id":"m1","ts":1,"from":"a","body":"hi"}', + "", + 'data: {"id":"m2","ts":2,"from":"b","body":"yo"}', + "", + ] + + upstream_resp = MagicMock() + upstream_resp.aiter_lines = MagicMock(return_value=_fake_sse_lines(sse_body)) + + upstream_ctx = AsyncMock() + upstream_ctx.__aenter__ = AsyncMock(return_value=upstream_resp) + upstream_ctx.__aexit__ = AsyncMock(return_value=False) + + client_ctx = AsyncMock() + client_ctx.stream = MagicMock(return_value=upstream_ctx) + client_ctx.__aenter__ = AsyncMock(return_value=client_ctx) + client_ctx.__aexit__ = AsyncMock(return_value=False) + + with patch( + "tinyagentos.routes.a2a_bus.httpx.AsyncClient", + return_value=client_ctx, + ): + resp = await client.get( + "/api/a2a/bus/stream", + params={"channel": "general"}, + headers={"Authorization": f"Bearer {token}"}, + ) + assert resp.status_code == 200 + assert resp.headers["content-type"].startswith("text/event-stream") + body = resp.text + assert "m1" in body + assert "m2" in body + # The bus stream was reached via the proxy with the channel mapped to thread. + assert client_ctx.stream.called + call_args = client_ctx.stream.call_args + assert call_args.args[0] == "GET" + assert call_args.args[1].endswith("/a2a/stream") + + async def test_stream_forwards_since_cursor(self, agent_app, client): + _, token = await _mint_agent(agent_app, scopes=("a2a_receive",)) + upstream_resp = MagicMock() + upstream_resp.aiter_lines = MagicMock(return_value=_fake_sse_lines([])) + upstream_ctx = AsyncMock() + upstream_ctx.__aenter__ = AsyncMock(return_value=upstream_resp) + upstream_ctx.__aexit__ = AsyncMock(return_value=False) + client_ctx = AsyncMock() + client_ctx.stream = MagicMock(return_value=upstream_ctx) + client_ctx.__aenter__ = AsyncMock(return_value=client_ctx) + client_ctx.__aexit__ = AsyncMock(return_value=False) + + with patch( + "tinyagentos.routes.a2a_bus.httpx.AsyncClient", + return_value=client_ctx, + ): + resp = await client.get( + "/api/a2a/bus/stream", + params={"channel": "general", "since": "1234.5"}, + headers={"Authorization": f"Bearer {token}"}, + ) + assert resp.status_code == 200 + params = client_ctx.stream.call_args.kwargs["params"] + assert params["thread"] == "general" + assert float(params["since"]) == 1234.5 + + async def test_stream_requires_channel(self, agent_app, client): + _, token = await _mint_agent(agent_app, scopes=("a2a_receive",)) + resp = await client.get( + "/api/a2a/bus/stream", + headers={"Authorization": f"Bearer {token}"}, + ) + assert resp.status_code == 400 + + async def test_unauthenticated_stream_is_401(self, noauth_client): + resp = await noauth_client.get( + "/api/a2a/bus/stream", + params={"channel": "general"}, + ) + assert resp.status_code == 401 + + async def test_stream_forbidden_without_a2a_receive(self, agent_app, client): + _, token = await _mint_agent(agent_app, scopes=("project_tasks",)) + resp = await client.get( + "/api/a2a/bus/stream", + params={"channel": "general"}, + headers={"Authorization": f"Bearer {token}"}, + ) + assert resp.status_code == 403 + + +@pytest.mark.asyncio +class TestMessagesSincePassthrough: + async def test_messages_forwards_since(self, agent_app, client): + payload = {"messages": [{"id": "m1", "ts": 1, "from": "a", "body": "hi"}]} + mock_resp = MagicMock() + mock_resp.raise_for_status = MagicMock() + mock_resp.json.return_value = payload + mock_client = AsyncMock() + mock_client.get.return_value = mock_resp + mock_ctx = AsyncMock() + mock_ctx.__aenter__ = AsyncMock(return_value=mock_client) + mock_ctx.__aexit__ = AsyncMock(return_value=False) + + with patch( + "tinyagentos.routes.a2a_bus.httpx.AsyncClient", + return_value=mock_ctx, + ): + resp = await client.get( + "/api/a2a/bus/messages", + params={"channel": "general", "since": "42.0"}, + ) + assert resp.status_code == 200 + call_kwargs = mock_client.get.call_args.kwargs + assert call_kwargs["params"]["thread"] == "general" + assert float(call_kwargs["params"]["since"]) == 42.0 + + +@pytest.mark.asyncio +class TestRawBusNeverExposed: + async def test_raw_stream_endpoint_is_not_a_proxy_route(self, client): + # The :7900 bus must not be reachable through a taOS route. Our stream + # route is /api/a2a/bus/stream; a raw /a2a/stream on this host must 404 + # (it is not a registered route), proving the proxy is the only path in. + resp = await client.get("/a2a/stream", params={"thread": "general"}) + assert resp.status_code == 404 diff --git a/tinyagentos/auth_middleware.py b/tinyagentos/auth_middleware.py index 343c942f0..ff9ed59af 100644 --- a/tinyagentos/auth_middleware.py +++ b/tinyagentos/auth_middleware.py @@ -23,6 +23,7 @@ _A2A_BUS_READ_PATHS = frozenset({ "/api/a2a/bus/channels", "/api/a2a/bus/messages", + "/api/a2a/bus/stream", }) # Authenticated A2A bus WRITE path: an agent may POST here with its own registry # JWT (scope a2a_send, verified by the route, which forces the bus `from` to the diff --git a/tinyagentos/routes/a2a_bus.py b/tinyagentos/routes/a2a_bus.py index 9a654fdaf..e2e023357 100644 --- a/tinyagentos/routes/a2a_bus.py +++ b/tinyagentos/routes/a2a_bus.py @@ -28,8 +28,10 @@ import os import httpx +import asyncio + from fastapi import APIRouter, HTTPException, Request -from fastapi.responses import JSONResponse +from fastapi.responses import JSONResponse, StreamingResponse from pydantic import BaseModel, Field from tinyagentos.agent_token_auth import check_agent_scope @@ -40,6 +42,10 @@ _DEFAULT_BUS_URL = "http://127.0.0.1:7900" +# Idle-stream heartbeat so intermediaries (and the taOSgo relay) do not reap a +# quiet SSE connection. The design mandates a `: ping` comment every 25s. +_STREAM_HEARTBEAT_SEC = 25 + def _bus_url() -> str: """Resolve the bus base URL from the environment, trailing slash stripped.""" @@ -95,27 +101,37 @@ async def bus_channels(request: Request): @router.get("/api/a2a/bus/messages") -async def bus_messages(request: Request, channel: str = "", limit: int = 100): +async def bus_messages( + request: Request, + channel: str = "", + limit: int = 100, + since: float | None = None, +): """Read messages from one bus channel, oldest-first as the bus returns them. Authorized readers: an admin session, the host local token, or an active agent registry JWT holding the ``a2a_receive`` scope. ``channel`` is required and maps to the bus ``thread`` query param. ``limit`` - is clamped to 1..500. On a bus error this returns an empty list with - ``available: false`` and HTTP 200. + is clamped to 1..500. ``since`` is forwarded verbatim to the bus as the + cursor (a message ``ts``); the bus replays everything after it, so an agent + can resume from the highest ``ts`` it has processed. On a bus error this + returns an empty list with ``available: false`` and HTTP 200. """ await _authorize_bus_read(request) if not channel: return JSONResponse({"error": "channel required"}, status_code=400) limit = max(1, min(500, limit)) + params: dict = {"thread": channel, "limit": limit} + if since is not None: + params["since"] = since bus = _bus_url() try: async with httpx.AsyncClient(timeout=5.0) as client: resp = await client.get( f"{bus}/a2a/messages", - params={"thread": channel, "limit": limit}, + params=params, ) resp.raise_for_status() data = resp.json() @@ -127,6 +143,75 @@ async def bus_messages(request: Request, channel: str = "", limit: int = 100): return {"messages": messages, "available": True} +@router.get("/api/a2a/bus/stream") +async def bus_stream( + request: Request, + channel: str = "", + since: float | None = None, +): + """Authenticated SSE proxy to the raw bus stream. + + Mirrors the existing messages proxy's auth gate exactly: an admin session, + the host local token, or an active agent registry JWT holding ``a2a_receive`` + (fail closed). It holds ONE upstream SSE connection to the bus per client and + relays events verbatim, injecting a ``: ping`` heartbeat comment every 25s so + idle streams are distinguishable from dead ones and intermediaries do not reap + them. The raw :7900 bus is never exposed directly: ``channel`` maps to the + bus ``thread`` query param and ``since`` to the bus cursor. + """ + await _authorize_bus_read(request) + if not channel: + return JSONResponse({"error": "channel required"}, status_code=400) + + bus = _bus_url() + + async def event_stream(): + try: + async with httpx.AsyncClient(timeout=None) as client: + async with client.stream( + "GET", + f"{bus}/a2a/stream", + params={"thread": channel, **({"since": since} if since is not None else {})}, + ) as upstream: + heartbeat = asyncio.create_task( + _stream_sleep(_STREAM_HEARTBEAT_SEC) + ) + try: + async for line in upstream.aiter_lines(): + if await request.is_disconnected(): + break + # Drain any pending heartbeat that fired while we were + # forwarding real data: emit it before the next event. + if heartbeat.done(): + heartbeat.cancel() + yield ": ping\n\n" + heartbeat = asyncio.create_task( + _stream_sleep(_STREAM_HEARTBEAT_SEC) + ) + if line == "": + continue + yield f"{line}\n\n" + finally: + heartbeat.cancel() + except Exception as exc: # noqa: BLE001 (surface a final SSE comment) + logger.warning("A2A bus stream proxy failed (%s): %s", bus, exc) + yield f": stream error\n\n" + + return StreamingResponse( + event_stream(), + media_type="text/event-stream", + headers={ + "Cache-Control": "no-cache", + "X-Accel-Buffering": "no", + }, + ) + + +async def _stream_sleep(secs: float) -> None: + """Sleep for *secs*; cancelled cleanly when a real event arrives first.""" + await asyncio.sleep(secs) + + class BusSendBody(BaseModel): """A message to post to a coordination-bus thread.