fix(stream): cap concurrent SSE streams per tenant (S-5)

src/serving/api/routers/stream.py

@@ -2,16 +2,22 @@
 
 import asyncio
 import json
+import os
 from collections.abc import AsyncIterator
 from datetime import datetime
+from typing import TYPE_CHECKING
 
-from fastapi import APIRouter, Request
+from fastapi import APIRouter, HTTPException, Request
 from fastapi.responses import StreamingResponse
 from opentelemetry import trace
+from starlette.background import BackgroundTask
 from starlette.concurrency import run_in_threadpool
 
 from src.serving.seen_events import BoundedSeenSet
 
+if TYPE_CHECKING:
+    from fastapi import FastAPI
+
 router = APIRouter(prefix="/v1/stream", tags=["stream"])
 tracer = trace.get_tracer("agentflow.api")
 
@@ -24,6 +30,59 @@
 # by then it can never re-enter the window.
 SEEN_CACHE_SIZE = 10_000
 
+DEFAULT_MAX_STREAMS_PER_TENANT = 5
+
+
+def max_streams_per_tenant() -> int:
+    """Per-tenant cap on concurrent SSE connections (security pre-audit S-5).
+
+    Each open stream runs a journal scan every second for as long as it stays
+    open, but the rate limiter only charges the one request that opened it —
+    without a cap a tenant's scan load grows by rate_limit_rpm scanners per
+    minute, unbounded. The cap is per process, like the failed-auth throttle
+    (S-7 accepted risk): N replicas give an N× effective cap, which still
+    bounds the growth the finding is about.
+    """
+    default = str(DEFAULT_MAX_STREAMS_PER_TENANT)
+    return int(os.getenv("AGENTFLOW_SSE_MAX_STREAMS_PER_TENANT", default))
+
+
+def _active_stream_counts(app: "FastAPI") -> dict[str, int]:
+    counts = getattr(app.state, "sse_active_streams", None)
+    if counts is None:
+        counts = {}
+        app.state.sse_active_streams = counts
+    return counts
+
+
+class _StreamSlot:
+    """One tenant's claim on a concurrent-stream slot.
+
+    ``release()`` is idempotent and wired to BOTH the generator's ``finally``
+    and the response's background task: a started generator releases on close
+    (client disconnect included, via ``aclose()``), while a generator the
+    server never iterates skips its ``finally`` entirely — there the
+    background task, which Starlette runs after the response finishes, is the
+    release path. Check-then-claim runs with no ``await`` in between, so it is
+    atomic on the event loop.
+    """
+
+    def __init__(self, counts: dict[str, int], tenant_id: str) -> None:
+        self._counts = counts
+        self._tenant_id = tenant_id
+        self._released = False
+        counts[tenant_id] = counts.get(tenant_id, 0) + 1
+
+    def release(self) -> None:
+        if self._released:
+            return
+        self._released = True
+        remaining = self._counts.get(self._tenant_id, 0) - 1
+        if remaining > 0:
+            self._counts[self._tenant_id] = remaining
+        else:
+            self._counts.pop(self._tenant_id, None)
+
 
 async def fetch_recent_events(
     request: Request,
@@ -61,55 +120,77 @@ async def stream_events(
     entity_id: str | None = None,
 ) -> StreamingResponse:
     """Server-Sent Events stream of validated pipeline events."""
+    # tenant_id is None only with auth disabled (dev/demo mode) — then the cap
+    # is skipped, same convention as the /v1/batch rate-limit charge (S-4).
+    tenant_id = getattr(request.state, "tenant_id", None)
+    slot: _StreamSlot | None = None
+    if tenant_id is not None:
+        cap = max_streams_per_tenant()
+        counts = _active_stream_counts(request.app)
+        if counts.get(tenant_id, 0) >= cap:
+            raise HTTPException(
+                status_code=429,
+                detail=(
+                    f"Too many concurrent event streams for this tenant "
+                    f"(limit {cap}). Close an open stream and retry."
+                ),
+                headers={"Retry-After": "1"},
+            )
+        slot = _StreamSlot(counts, tenant_id)
 
     async def event_generator() -> AsyncIterator[str]:
-        seen_event_ids = BoundedSeenSet(maxlen=SEEN_CACHE_SIZE)
-        events_sent = 0
-
-        with tracer.start_as_current_span("sse_stream") as span:
-            span.set_attribute("stream.event_type", event_type or "all")
-            if entity_id is not None:
-                span.set_attribute("stream.entity_id", entity_id)
-
-            while True:
-                if await request.is_disconnected():
-                    break
-
-                events = await fetch_recent_events(
-                    request=request,
-                    event_type=event_type,
-                    entity_id=entity_id,
-                    limit=10,
-                )
-
-                emitted = False
-                for event in reversed(events):
-                    if await request.is_disconnected():
-                        span.set_attribute("stream.events_sent", events_sent)
-                        return
-
-                    event_id = str(event.get("event_id", ""))
-                    if event_id in seen_event_ids:
-                        continue
-
-                    seen_event_ids.add(event_id)
-                    payload = {
-                        key: value.isoformat() if isinstance(value, datetime) else value
-                        for key, value in event.items()
-                    }
-                    emitted = True
-                    events_sent += 1
-                    yield f"data: {json.dumps(payload)}\n\n"
+        try:
+            seen_event_ids = BoundedSeenSet(maxlen=SEEN_CACHE_SIZE)
+            events_sent = 0
 
-                if not emitted:
-                    yield ": keepalive\n\n"
+            with tracer.start_as_current_span("sse_stream") as span:
+                span.set_attribute("stream.event_type", event_type or "all")
+                if entity_id is not None:
+                    span.set_attribute("stream.entity_id", entity_id)
 
-                await asyncio.sleep(1.0)
-
-            span.set_attribute("stream.events_sent", events_sent)
+                while True:
+                    if await request.is_disconnected():
+                        break
+
+                    events = await fetch_recent_events(
+                        request=request,
+                        event_type=event_type,
+                        entity_id=entity_id,
+                        limit=10,
+                    )
+
+                    emitted = False
+                    for event in reversed(events):
+                        if await request.is_disconnected():
+                            span.set_attribute("stream.events_sent", events_sent)
+                            return
+
+                        event_id = str(event.get("event_id", ""))
+                        if event_id in seen_event_ids:
+                            continue
+
+                        seen_event_ids.add(event_id)
+                        payload = {
+                            key: value.isoformat() if isinstance(value, datetime) else value
+                            for key, value in event.items()
+                        }
+                        emitted = True
+                        events_sent += 1
+                        yield f"data: {json.dumps(payload)}\n\n"
+
+                    if not emitted:
+                        yield ": keepalive\n\n"
+
+                    await asyncio.sleep(1.0)
+
+                span.set_attribute("stream.events_sent", events_sent)
+        finally:
+            if slot is not None:
+                slot.release()
 
     return StreamingResponse(
         event_generator(),
         media_type="text/event-stream",
         headers={"Cache-Control": "no-cache"},
+        background=BackgroundTask(slot.release) if slot is not None else None,
     )

tests/unit/test_stream_router_unit.py

@@ -17,6 +17,7 @@
 
 import duckdb
 import pytest
+from fastapi import HTTPException
 
 from src.serving.api.routers import stream as stream_module
 from src.serving.api.routers.stream import fetch_recent_events, stream_events
@@ -164,11 +165,17 @@ async def test_fetch_synthesizes_topic_when_column_absent() -> None:
 
 class _StreamReq:
     """Request whose ``is_disconnected`` flips to True after N checks, so the
-    SSE loop runs a bounded number of iterations."""
-
-    def __init__(self, *, disconnect_after: int) -> None:
-        self.app = SimpleNamespace(state=SimpleNamespace(query_engine=SimpleNamespace(_conn=None)))
-        self.state = SimpleNamespace(tenant_id=None)
+    SSE loop runs a bounded number of iterations. Pass the same ``app`` to
+    several requests to model connections sharing one process (the S-5
+    concurrent-stream counter lives on ``app.state``)."""
+
+    def __init__(
+        self, *, disconnect_after: int, tenant_id: str | None = None, app: Any = None
+    ) -> None:
+        self.app = app or SimpleNamespace(
+            state=SimpleNamespace(query_engine=SimpleNamespace(_conn=None))
+        )
+        self.state = SimpleNamespace(tenant_id=tenant_id)
         self._checks = 0
         self._disconnect_after = disconnect_after
 
@@ -310,6 +317,112 @@ async def fake_fetch(**_kwargs: Any) -> list[dict[str, object]]:
     assert created == [stream_module.SEEN_CACHE_SIZE]
 
 
+# ── concurrent-stream cap (security pre-audit S-5) ───────────────
+
+
+def _capped_app() -> SimpleNamespace:
+    return SimpleNamespace(state=SimpleNamespace(query_engine=SimpleNamespace(_conn=None)))
+
+
+@pytest.fixture
+def _empty_fetch(monkeypatch: pytest.MonkeyPatch) -> None:
+    async def fake_fetch(**_kwargs: Any) -> list[dict[str, object]]:
+        return []
+
+    monkeypatch.setattr(stream_module, "fetch_recent_events", fake_fetch)
+
+
+@pytest.mark.usefixtures("_empty_fetch")
+async def test_stream_cap_rejects_excess_connection_with_429(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    monkeypatch.setenv("AGENTFLOW_SSE_MAX_STREAMS_PER_TENANT", "2")
+    app = _capped_app()
+
+    r1 = await stream_events(_StreamReq(disconnect_after=1, tenant_id="acme", app=app))
+    r2 = await stream_events(_StreamReq(disconnect_after=1, tenant_id="acme", app=app))
+    # Slots are claimed at accept time, before the generator ever runs.
+    assert app.state.sse_active_streams == {"acme": 2}
+
+    with pytest.raises(HTTPException) as excinfo:
+        await stream_events(_StreamReq(disconnect_after=1, tenant_id="acme", app=app))
+    assert excinfo.value.status_code == 429
+    assert excinfo.value.headers == {"Retry-After": "1"}
+    # The rejected request must not consume a slot.
+    assert app.state.sse_active_streams == {"acme": 2}
+
+    # Closing both streams frees the slots and the next connection is accepted.
+    await _drain(r1)
+    await _drain(r2)
+    assert app.state.sse_active_streams == {}
+    r4 = await stream_events(_StreamReq(disconnect_after=1, tenant_id="acme", app=app))
+    assert r4.media_type == "text/event-stream"
+    await _drain(r4)
+
+
+@pytest.mark.usefixtures("_empty_fetch")
+async def test_stream_cap_is_per_tenant(monkeypatch: pytest.MonkeyPatch) -> None:
+    monkeypatch.setenv("AGENTFLOW_SSE_MAX_STREAMS_PER_TENANT", "1")
+    app = _capped_app()
+
+    r_acme = await stream_events(_StreamReq(disconnect_after=1, tenant_id="acme", app=app))
+    # acme is at its cap; globex still gets its own budget.
+    r_globex = await stream_events(_StreamReq(disconnect_after=1, tenant_id="globex", app=app))
+    assert app.state.sse_active_streams == {"acme": 1, "globex": 1}
+
+    with pytest.raises(HTTPException):
+        await stream_events(_StreamReq(disconnect_after=1, tenant_id="acme", app=app))
+
+    await _drain(r_acme)
+    await _drain(r_globex)
+    assert app.state.sse_active_streams == {}
+
+
+@pytest.mark.usefixtures("_empty_fetch")
+async def test_stream_cap_skipped_when_auth_disabled(monkeypatch: pytest.MonkeyPatch) -> None:
+    # tenant_id=None means auth is off (dev/demo) — no cap, matching the
+    # /v1/batch S-4 convention.
+    monkeypatch.setenv("AGENTFLOW_SSE_MAX_STREAMS_PER_TENANT", "1")
+    app = _capped_app()
+
+    responses = [
+        await stream_events(_StreamReq(disconnect_after=1, tenant_id=None, app=app))
+        for _ in range(3)
+    ]
+    assert getattr(app.state, "sse_active_streams", None) is None
+    for response in responses:
+        await _drain(response)
+
+
+async def test_stream_cap_default_and_env_override(monkeypatch: pytest.MonkeyPatch) -> None:
+    monkeypatch.delenv("AGENTFLOW_SSE_MAX_STREAMS_PER_TENANT", raising=False)
+    assert stream_module.max_streams_per_tenant() == stream_module.DEFAULT_MAX_STREAMS_PER_TENANT
+    monkeypatch.setenv("AGENTFLOW_SSE_MAX_STREAMS_PER_TENANT", "7")
+    assert stream_module.max_streams_per_tenant() == 7
+
+
+@pytest.mark.usefixtures("_empty_fetch")
+async def test_stream_slot_released_by_background_task_when_never_iterated(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    # If the server never iterates the generator (response setup fails before
+    # streaming), its `finally` never runs — the response's background task is
+    # the release path, and release stays idempotent if both paths do run.
+    monkeypatch.setenv("AGENTFLOW_SSE_MAX_STREAMS_PER_TENANT", "1")
+    app = _capped_app()
+
+    response = await stream_events(_StreamReq(disconnect_after=1, tenant_id="acme", app=app))
+    assert app.state.sse_active_streams == {"acme": 1}
+    assert response.background is not None
+
+    await response.background()
+    assert app.state.sse_active_streams == {}
+
+    # Draining after the background release must not push the count negative.
+    await _drain(response)
+    assert app.state.sse_active_streams == {}
+
+
 async def test_stream_dedup_survives_eviction_of_older_ids(
     monkeypatch: pytest.MonkeyPatch,
 ) -> None: