feat(a2a): authenticated SSE stream proxy + since cursor (#1780 S3)

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

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

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.