From bfc481655f2be7f6c677a20caed3c377fc37e240 Mon Sep 17 00:00:00 2001 From: JuliaEdom Date: Mon, 20 Jul 2026 17:01:39 +0300 Subject: [PATCH] fix(stream): cap concurrent SSE streams per tenant (S-5) Each open /v1/stream/events connection runs a journal scan every second but the rate limiter only charges the request that opened it, so a tenant's scan load grew without bound (security pre-audit S-5). New per-tenant cap on concurrent streams: the (N+1)-th connection is rejected with 429 + Retry-After before the stream opens; already-open streams are never dropped. Cap defaults to 5 and is tunable via AGENTFLOW_SSE_MAX_STREAMS_PER_TENANT. The counter is per process, like the failed-auth throttle (S-7 accepted risk): N replicas give an NxCap effective cap, which still bounds the growth the finding is about. Slot release is idempotent and wired to both the generator's finally (normal close/disconnect) and the response background task (generator never iterated). With auth disabled (tenant_id=None) the cap is skipped, matching the /v1/batch S-4 convention. Co-Authored-By: Claude Fable 5 --- src/serving/api/routers/stream.py | 167 +++++++++++++++++++------- tests/unit/test_stream_router_unit.py | 123 ++++++++++++++++++- 2 files changed, 242 insertions(+), 48 deletions(-) diff --git a/src/serving/api/routers/stream.py b/src/serving/api/routers/stream.py index 274910d5..ff39461f 100644 --- a/src/serving/api/routers/stream.py +++ b/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, ) diff --git a/tests/unit/test_stream_router_unit.py b/tests/unit/test_stream_router_unit.py index b98a8646..8943fc5b 100644 --- a/tests/unit/test_stream_router_unit.py +++ b/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: