From 63c11ee5f3dc570461c1228f50cb5e54c25e237c Mon Sep 17 00:00:00 2001 From: jaylfc Date: Thu, 16 Jul 2026 19:06:36 +0100 Subject: [PATCH] feat(a2a): authenticated SSE stream proxy + since cursor (#1780 S3) Add GET /api/a2a/bus/stream, an authenticated SSE proxy to the raw bus stream, gated exactly like the messages proxy. Add the stream path to the middleware read allowlist and forward the since cursor on the messages proxy. --- tests/test_routes_a2a_bus_stream.py | 201 ++++++++++++++++++++++++++++ tinyagentos/auth_middleware.py | 1 + tinyagentos/routes/a2a_bus.py | 95 ++++++++++++- 3 files changed, 292 insertions(+), 5 deletions(-) create mode 100644 tests/test_routes_a2a_bus_stream.py 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.