From 6f6efc2b7527fd5844ed9575e1e14d01d59a483d Mon Sep 17 00:00:00 2001 From: Phil Merrell Date: Sat, 2 May 2026 10:42:33 -0600 Subject: [PATCH] feat(chat): add BFF chat SSE proxy (Phase 4) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Adds POST /chat/proxy-stream on app-api, the cookie-authenticated SSE proxy the SPA will adopt in Phase 6 once it cuts over from Bearer-in- localStorage to the __Host-bff_session cookie. Dormant in Phase 4 — no SPA caller exists yet. Mechanics: - Authenticated via get_current_user_from_session (Phase 2 dep). SessionRefreshMiddleware refreshes the stored Cognito access token if near expiry before the handler runs. - Forwards request body verbatim to {INFERENCE_API_URL}/invocations with Authorization: Bearer . Inference-api already accepts Cognito Bearer tokens via get_current_user_trusted — no inference-api changes (architecture decision #4). - Sets X-Accel-Buffering: no and Cache-Control: no-cache so late SSE events (notably oauth_required after message_stop) reach the browser without an intermediary holding them back. - Guarded by CSRFMiddleware via the existing exempt-list; double- submit X-CSRF-Token must match the __Host-bff_csrf cookie. - Maps upstream connect/timeout errors to 502/504 to match the existing converse proxy. Lifecycle subtlety: the upstream httpx.AsyncClient is built via a _build_upstream_client() seam and closed in the streaming generator's finally block — closing the client via async with while a stream is in flight makes httpx drain the upstream during __aexit__, defeating the whole point of the streaming proxy. The seam also lets tests inject a MockTransport-backed client without mutating global httpx. Tests (16 new): - auth gate (no session → 401), - header/body/URL relay, - SSE and non-SSE response paths, - upstream 4xx/5xx propagation, - ConnectError → 502, TimeoutException → 504, - CSRF: missing → 403, header/cookie mismatch → 403, valid → 200, plus a guard that /chat/proxy-stream is not in CSRFMiddleware._EXEMPT_PATHS, - TTFB < 200ms integration test backed by a real uvicorn server (httpx ASGITransport buffers the body before headers, so it can't measure TTFB) with a slow upstream that yields, sleeps 300ms, then yields more — asserts headers arrive fast and the late oauth_required event makes it through. Co-Authored-By: Claude Opus 4.7 --- backend/src/apis/app_api/chat/proxy_routes.py | 151 ++++++++++ backend/src/apis/app_api/main.py | 2 + backend/tests/apis/app_api/chat/__init__.py | 0 .../apis/app_api/chat/test_proxy_routes.py | 280 ++++++++++++++++++ .../app_api/chat/test_proxy_routes_csrf.py | 153 ++++++++++ .../chat/test_proxy_routes_streaming.py | 171 +++++++++++ 6 files changed, 757 insertions(+) create mode 100644 backend/src/apis/app_api/chat/proxy_routes.py create mode 100644 backend/tests/apis/app_api/chat/__init__.py create mode 100644 backend/tests/apis/app_api/chat/test_proxy_routes.py create mode 100644 backend/tests/apis/app_api/chat/test_proxy_routes_csrf.py create mode 100644 backend/tests/apis/app_api/chat/test_proxy_routes_streaming.py diff --git a/backend/src/apis/app_api/chat/proxy_routes.py b/backend/src/apis/app_api/chat/proxy_routes.py new file mode 100644 index 00000000..c48d15d7 --- /dev/null +++ b/backend/src/apis/app_api/chat/proxy_routes.py @@ -0,0 +1,151 @@ +"""BFF chat proxy — forwards browser SSE chat requests to inference-api. + +`POST /chat/proxy-stream` is the cookie-authenticated counterpart to the +SPA's current direct-to-inference-api chat call. The flow: + + Browser → CloudFront `/api/*` → app-api → inference-api `/invocations` + (httpOnly session cookie) (Authorization: Bearer ) + +`SessionRefreshMiddleware` resolves the cookie and, if the stored Cognito +access token is near expiry, refreshes it before this handler runs. The +handler then forwards `current_user.raw_token` — the freshly-validated +access token — to inference-api, which already accepts Cognito Bearer +tokens via `get_current_user_trusted` on `/invocations`. No inference-api +changes needed (architecture decision #4 in the BFF migration plan). + +Phase 4 ships this dormant — no SPA caller exists yet. Phase 6b will +rename this to `/chat/stream` once the SPA cuts over to cookie auth. +""" + +from __future__ import annotations + +import logging +import os + +import httpx +from fastapi import APIRouter, Depends, HTTPException, Request +from fastapi.responses import StreamingResponse + +from apis.shared.auth.dependencies import get_current_user_from_session +from apis.shared.auth.models import User + +logger = logging.getLogger(__name__) + +router = APIRouter(prefix="/chat", tags=["bff-chat-proxy"]) + +_INFERENCE_API_URL = os.environ.get("INFERENCE_API_URL", "http://localhost:8001") + +# Long enough to cover a full agent turn (model + tool calls), bounded so a +# wedged upstream eventually surfaces. +_PROXY_TIMEOUT_SECONDS = 300.0 + + +def _build_upstream_client() -> httpx.AsyncClient: + """Single seam where the proxy's upstream client is constructed. + + Tests substitute a MockTransport-backed client here without having to + monkey-patch the global `httpx.AsyncClient` symbol — which would also + intercept any test-side httpx clients running in the same process. + """ + return httpx.AsyncClient(timeout=httpx.Timeout(_PROXY_TIMEOUT_SECONDS)) + + +@router.post( + "/proxy-stream", + summary="Cookie-authenticated SSE proxy to inference-api /invocations", + responses={ + 401: {"description": "No active BFF session"}, + 403: {"description": "CSRF token missing or invalid"}, + 502: {"description": "Inference API unreachable"}, + 504: {"description": "Inference API request timed out"}, + }, +) +async def chat_proxy_stream( + request: Request, + current_user: User = Depends(get_current_user_from_session), +): + """Relay the request body verbatim to inference-api `/invocations`. + + The body is opaque bytes — validation lives on inference-api so this + handler stays decoupled from the InvocationRequest schema. SSE chunks + flow back unmodified; `X-Accel-Buffering: no` defeats proxy buffering + so streaming events (notably `oauth_required` after `message_stop`) + reach the browser without being held back by an intermediary. + """ + target_url = f"{_INFERENCE_API_URL}/invocations" + body = await request.body() + + headers = { + "Content-Type": "application/json", + "Authorization": f"Bearer {current_user.raw_token}", + } + + # The client lifecycle must outlive this handler — closing it via + # `async with` while a stream is in flight makes httpx drain the upstream + # response during `__aexit__`, buffering the entire SSE stream before + # headers reach the browser. Open the client manually and tie its + # cleanup to the streaming generator's `finally` (or to the early-exit + # paths below) so headers can flush as soon as the upstream's first + # response message arrives. + client = _build_upstream_client() + try: + response = await client.send( + client.build_request("POST", target_url, headers=headers, content=body), + stream=True, + ) + except httpx.ConnectError: + await client.aclose() + logger.error(f"Cannot reach Inference API at {target_url}") + raise HTTPException(status_code=502, detail="Inference API is unreachable") + except httpx.TimeoutException: + await client.aclose() + logger.error(f"Inference API request timed out: {target_url}") + raise HTTPException(status_code=504, detail="Inference API request timed out") + except Exception as exc: + await client.aclose() + logger.error(f"BFF chat proxy error: {exc}", exc_info=True) + raise HTTPException( + status_code=502, + detail="An unexpected error occurred while proxying to the Inference API", + ) + + if response.status_code >= 400: + try: + error_body = await response.aread() + finally: + await response.aclose() + await client.aclose() + raise HTTPException( + status_code=response.status_code, + detail=error_body.decode("utf-8", errors="replace"), + ) + + content_type = response.headers.get("content-type", "") + if "text/event-stream" in content_type: + async def stream_relay(): + try: + async for chunk in response.aiter_bytes(): + yield chunk + finally: + await response.aclose() + await client.aclose() + + return StreamingResponse( + stream_relay(), + media_type="text/event-stream", + headers={ + "Cache-Control": "no-cache", + "X-Accel-Buffering": "no", + }, + ) + + try: + response_body = await response.aread() + finally: + await response.aclose() + await client.aclose() + return StreamingResponse( + iter([response_body]), + media_type=content_type or "application/json", + status_code=response.status_code, + ) diff --git a/backend/src/apis/app_api/main.py b/backend/src/apis/app_api/main.py index 56509106..e9af8dda 100644 --- a/backend/src/apis/app_api/main.py +++ b/backend/src/apis/app_api/main.py @@ -116,6 +116,7 @@ async def lifespan(app: FastAPI): from apis.app_api.costs.routes import router as costs_router from apis.app_api.chat.routes import router as chat_router from apis.app_api.chat.converse_routes import router as converse_router +from apis.app_api.chat.proxy_routes import router as bff_chat_proxy_router from apis.app_api.memory.routes import router as memory_router from apis.app_api.tools.routes import router as tools_router from apis.app_api.files.routes import router as files_router @@ -142,6 +143,7 @@ async def lifespan(app: FastAPI): app.include_router(costs_router) app.include_router(chat_router) # Application-specific chat endpoints app.include_router(converse_router) # Proxies to Inference API for cost accounting +app.include_router(bff_chat_proxy_router) # Cookie-authenticated SSE proxy (Phase 4, dormant until SPA cutover) app.include_router(memory_router) # AgentCore Memory access endpoints app.include_router(tools_router) # Tool discovery and permissions app.include_router(files_router) # File upload via pre-signed URLs diff --git a/backend/tests/apis/app_api/chat/__init__.py b/backend/tests/apis/app_api/chat/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/backend/tests/apis/app_api/chat/test_proxy_routes.py b/backend/tests/apis/app_api/chat/test_proxy_routes.py new file mode 100644 index 00000000..baed3efe --- /dev/null +++ b/backend/tests/apis/app_api/chat/test_proxy_routes.py @@ -0,0 +1,280 @@ +"""Tests for `POST /chat/proxy-stream` (Phase 4 BFF chat proxy). + +Covers the proxy mechanics in isolation — auth gate, body/header relay, +SSE streaming, and error mapping. The full SessionRefreshMiddleware ↔ +CSRFMiddleware stack is exercised separately in +`test_proxy_routes_csrf.py`. +""" + +from __future__ import annotations + +import time +from typing import Callable, Optional + +import httpx +import pytest +from fastapi import Depends, FastAPI +from fastapi.testclient import TestClient +from starlette.middleware.base import BaseHTTPMiddleware + +from apis.app_api.chat import proxy_routes +from apis.app_api.chat.proxy_routes import router as proxy_router +from apis.shared.auth.dependencies import get_current_user_from_session +from apis.shared.auth.models import User +from apis.shared.sessions_bff.models import SessionRecord + + +def _record() -> SessionRecord: + now = int(time.time()) + return SessionRecord( + session_id="sess-001", + user_id="user-sub", + username="alice", + cognito_access_token="access.token.value", + cognito_refresh_token="refresh.token.value", + id_token="id.token.value", + access_token_exp=now + 3600, + csrf_secret="csrf-secret", + created_at=now, + last_seen_at=now, + ttl=now + 28800, + ) + + +def _user(*, raw_token: str = "access.token.value") -> User: + user = User( + email="alice@example.com", + user_id="user-sub", + name="Alice", + roles=["user"], + ) + user.raw_token = raw_token + return user + + +class _AttachSession(BaseHTTPMiddleware): + """Minimal stand-in for SessionRefreshMiddleware — sets bff_session.""" + + def __init__(self, app, record: Optional[SessionRecord]) -> None: + super().__init__(app) + self._record = record + + async def dispatch(self, request, call_next): + if self._record is not None: + request.state.bff_session = self._record + return await call_next(request) + + +def _build_app( + *, + record: Optional[SessionRecord] = None, + user_override: Optional[User] = None, +) -> FastAPI: + app = FastAPI() + app.add_middleware(_AttachSession, record=record) + app.include_router(proxy_router) + if user_override is not None: + app.dependency_overrides[get_current_user_from_session] = lambda: user_override + return app + + +def _patch_upstream( + monkeypatch: pytest.MonkeyPatch, + handler: Callable[[httpx.Request], httpx.Response], +) -> None: + """Replace the proxy's upstream-client builder with a MockTransport- + backed one. The seam lives in `proxy_routes._build_upstream_client` + so we don't have to mutate global `httpx.AsyncClient`.""" + transport = httpx.MockTransport(handler) + monkeypatch.setattr( + proxy_routes, + "_build_upstream_client", + lambda: httpx.AsyncClient(transport=transport), + ) + + +# ── Auth gate ───────────────────────────────────────────────────────────── + + +def test_returns_401_when_no_session_attached() -> None: + app = _build_app(record=None) + response = TestClient(app).post("/chat/proxy-stream", json={"message": "hi"}) + assert response.status_code == 401 + + +# ── Happy path: SSE relay ───────────────────────────────────────────────── + + +def test_relays_sse_response_verbatim(monkeypatch: pytest.MonkeyPatch) -> None: + sse_body = ( + b'event: message_start\ndata: {"role": "assistant"}\n\n' + b'event: content_block_delta\ndata: {"text": "hello"}\n\n' + b'event: done\ndata: {}\n\n' + ) + + def handler(request: httpx.Request) -> httpx.Response: + assert request.url.path == "/invocations" + assert request.method == "POST" + return httpx.Response( + 200, + content=sse_body, + headers={"content-type": "text/event-stream"}, + ) + + _patch_upstream(monkeypatch, handler) + app = _build_app(record=_record(), user_override=_user()) + + response = TestClient(app).post("/chat/proxy-stream", json={"message": "hi"}) + + assert response.status_code == 200 + assert response.headers["content-type"].startswith("text/event-stream") + assert response.headers["x-accel-buffering"] == "no" + assert response.headers["cache-control"] == "no-cache" + assert response.content == sse_body + + +def test_forwards_authorization_bearer_from_session(monkeypatch: pytest.MonkeyPatch) -> None: + captured: dict = {} + + def handler(request: httpx.Request) -> httpx.Response: + captured["authorization"] = request.headers.get("authorization") + return httpx.Response( + 200, content=b"event: done\ndata: {}\n\n", + headers={"content-type": "text/event-stream"}, + ) + + _patch_upstream(monkeypatch, handler) + app = _build_app(record=_record(), user_override=_user(raw_token="the-stored-token")) + + TestClient(app).post("/chat/proxy-stream", json={"message": "hi"}) + assert captured["authorization"] == "Bearer the-stored-token" + + +def test_forwards_request_body_verbatim(monkeypatch: pytest.MonkeyPatch) -> None: + captured: dict = {} + + def handler(request: httpx.Request) -> httpx.Response: + captured["body"] = request.content + captured["content_type"] = request.headers.get("content-type") + return httpx.Response( + 200, content=b"event: done\ndata: {}\n\n", + headers={"content-type": "text/event-stream"}, + ) + + _patch_upstream(monkeypatch, handler) + app = _build_app(record=_record(), user_override=_user()) + + payload = b'{"session_id":"s1","message":"hello there","enabled_tools":["foo"]}' + TestClient(app).post( + "/chat/proxy-stream", + content=payload, + headers={"Content-Type": "application/json"}, + ) + assert captured["body"] == payload + assert captured["content_type"] == "application/json" + + +def test_targets_invocations_path_on_inference_api( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(proxy_routes, "_INFERENCE_API_URL", "http://upstream:9999") + captured: dict = {} + + def handler(request: httpx.Request) -> httpx.Response: + captured["url"] = str(request.url) + return httpx.Response( + 200, content=b"event: done\ndata: {}\n\n", + headers={"content-type": "text/event-stream"}, + ) + + _patch_upstream(monkeypatch, handler) + app = _build_app(record=_record(), user_override=_user()) + + TestClient(app).post("/chat/proxy-stream", json={"message": "hi"}) + assert captured["url"] == "http://upstream:9999/invocations" + + +# ── Non-SSE relay (e.g. inference-api returns JSON validation error pre-stream) ── + + +def test_relays_non_sse_response_with_status_and_content_type( + monkeypatch: pytest.MonkeyPatch, +) -> None: + def handler(request: httpx.Request) -> httpx.Response: + # Simulate inference-api returning a non-streaming success — a small + # JSON body that should be passed through, not re-wrapped as SSE. + return httpx.Response( + 200, + content=b'{"ok": true}', + headers={"content-type": "application/json"}, + ) + + _patch_upstream(monkeypatch, handler) + app = _build_app(record=_record(), user_override=_user()) + + response = TestClient(app).post("/chat/proxy-stream", json={"message": "hi"}) + assert response.status_code == 200 + assert response.headers["content-type"].startswith("application/json") + assert response.content == b'{"ok": true}' + + +# ── Upstream error propagation ──────────────────────────────────────────── + + +def test_propagates_upstream_4xx(monkeypatch: pytest.MonkeyPatch) -> None: + def handler(request: httpx.Request) -> httpx.Response: + return httpx.Response(401, content=b"token rejected by inference-api") + + _patch_upstream(monkeypatch, handler) + app = _build_app(record=_record(), user_override=_user()) + + response = TestClient(app).post("/chat/proxy-stream", json={"message": "hi"}) + assert response.status_code == 401 + assert "token rejected" in response.text + + +def test_propagates_upstream_5xx(monkeypatch: pytest.MonkeyPatch) -> None: + def handler(request: httpx.Request) -> httpx.Response: + return httpx.Response(503, content=b"upstream overloaded") + + _patch_upstream(monkeypatch, handler) + app = _build_app(record=_record(), user_override=_user()) + + response = TestClient(app).post("/chat/proxy-stream", json={"message": "hi"}) + assert response.status_code == 503 + + +def test_returns_502_when_upstream_unreachable(monkeypatch: pytest.MonkeyPatch) -> None: + def handler(request: httpx.Request) -> httpx.Response: + raise httpx.ConnectError("Connection refused") + + _patch_upstream(monkeypatch, handler) + app = _build_app(record=_record(), user_override=_user()) + + response = TestClient(app).post("/chat/proxy-stream", json={"message": "hi"}) + assert response.status_code == 502 + assert response.json()["detail"] == "Inference API is unreachable" + + +def test_returns_504_on_upstream_timeout(monkeypatch: pytest.MonkeyPatch) -> None: + def handler(request: httpx.Request) -> httpx.Response: + raise httpx.ReadTimeout("Read timed out") + + _patch_upstream(monkeypatch, handler) + app = _build_app(record=_record(), user_override=_user()) + + response = TestClient(app).post("/chat/proxy-stream", json={"message": "hi"}) + assert response.status_code == 504 + + +def test_returns_502_on_unexpected_upstream_error( + monkeypatch: pytest.MonkeyPatch, +) -> None: + def handler(request: httpx.Request) -> httpx.Response: + raise RuntimeError("something is on fire") + + _patch_upstream(monkeypatch, handler) + app = _build_app(record=_record(), user_override=_user()) + + response = TestClient(app).post("/chat/proxy-stream", json={"message": "hi"}) + assert response.status_code == 502 diff --git a/backend/tests/apis/app_api/chat/test_proxy_routes_csrf.py b/backend/tests/apis/app_api/chat/test_proxy_routes_csrf.py new file mode 100644 index 00000000..0761981c --- /dev/null +++ b/backend/tests/apis/app_api/chat/test_proxy_routes_csrf.py @@ -0,0 +1,153 @@ +"""CSRF integration tests for `POST /chat/proxy-stream`. + +`/chat/proxy-stream` is a POST that rides the BFF session cookie, so it +MUST be guarded by CSRFMiddleware. This file confirms (a) the route is +not accidentally listed in the middleware's exempt set, and (b) a valid +double-submit token allows the request through. +""" + +from __future__ import annotations + +import time +from typing import Optional + +import httpx +import pytest +from fastapi import FastAPI +from fastapi.testclient import TestClient +from starlette.middleware.base import BaseHTTPMiddleware + +from apis.app_api.chat import proxy_routes +from apis.app_api.chat.proxy_routes import router as proxy_router +from apis.shared.auth.dependencies import get_current_user_from_session +from apis.shared.auth.models import User +from apis.shared.middleware.csrf import CSRFMiddleware +from apis.shared.sessions_bff.config import CSRF_COOKIE_NAME, CSRF_HEADER_NAME +from apis.shared.sessions_bff.csrf import CSRFHelper +from apis.shared.sessions_bff.models import SessionRecord + + +def _record() -> SessionRecord: + now = int(time.time()) + return SessionRecord( + session_id="sess-csrf-001", + user_id="user-sub", + username="alice", + cognito_access_token="access.token", + cognito_refresh_token="refresh.token", + id_token="id.token", + access_token_exp=now + 3600, + csrf_secret="csrf-secret", + created_at=now, + last_seen_at=now, + ttl=now + 28800, + ) + + +def _user() -> User: + user = User( + email="alice@example.com", + user_id="user-sub", + name="Alice", + roles=["user"], + ) + user.raw_token = "access.token" + return user + + +class _AttachSession(BaseHTTPMiddleware): + """Stand-in for SessionRefreshMiddleware that just sets bff_session.""" + + def __init__(self, app, record: Optional[SessionRecord]) -> None: + super().__init__(app) + self._record = record + + async def dispatch(self, request, call_next): + if self._record is not None: + request.state.bff_session = self._record + return await call_next(request) + + +def _build_app(record: Optional[SessionRecord]) -> FastAPI: + app = FastAPI() + # Order matters: Starlette runs the LAST-added middleware first, so this + # reflects the production stack in main.py — CSRFMiddleware sees the + # session that AttachSession (a stand-in for SessionRefreshMiddleware) + # populated. + app.add_middleware(CSRFMiddleware) + app.add_middleware(_AttachSession, record=record) + app.include_router(proxy_router) + app.dependency_overrides[get_current_user_from_session] = _user + return app + + +def _ok_handler(request: httpx.Request) -> httpx.Response: + return httpx.Response( + 200, + content=b"event: done\ndata: {}\n\n", + headers={"content-type": "text/event-stream"}, + ) + + +def _patch_upstream(monkeypatch: pytest.MonkeyPatch) -> None: + transport = httpx.MockTransport(_ok_handler) + monkeypatch.setattr( + proxy_routes, + "_build_upstream_client", + lambda: httpx.AsyncClient(transport=transport), + ) + + +def test_proxy_stream_without_csrf_returns_403(monkeypatch: pytest.MonkeyPatch) -> None: + """Session attached but no CSRF cookie/header → CSRF middleware rejects.""" + _patch_upstream(monkeypatch) + app = _build_app(record=_record()) + + response = TestClient(app).post("/chat/proxy-stream", json={"message": "hi"}) + assert response.status_code == 403 + assert "CSRF" in response.json()["detail"] + + +def test_proxy_stream_with_valid_csrf_passes(monkeypatch: pytest.MonkeyPatch) -> None: + record = _record() + csrf_token = CSRFHelper.derive_token(record.csrf_secret, record.session_id) + + _patch_upstream(monkeypatch) + app = _build_app(record=record) + + response = TestClient(app).post( + "/chat/proxy-stream", + json={"message": "hi"}, + cookies={CSRF_COOKIE_NAME: csrf_token}, + headers={CSRF_HEADER_NAME: csrf_token}, + ) + assert response.status_code == 200 + assert response.headers["content-type"].startswith("text/event-stream") + + +def test_proxy_stream_csrf_header_mismatch_returns_403( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Cookie present but header value differs — classic forged-form scenario.""" + record = _record() + csrf_token = CSRFHelper.derive_token(record.csrf_secret, record.session_id) + + _patch_upstream(monkeypatch) + app = _build_app(record=record) + + response = TestClient(app).post( + "/chat/proxy-stream", + json={"message": "hi"}, + cookies={CSRF_COOKIE_NAME: csrf_token}, + headers={CSRF_HEADER_NAME: "different-token"}, + ) + assert response.status_code == 403 + + +def test_proxy_stream_path_not_in_csrf_exempt_set() -> None: + """Future-proofing: surface a regression if someone adds /chat/proxy-stream + to the CSRF exempt list. The point of this proxy is to be guarded — + exempting it would defeat the whole BFF-cookie security model.""" + from apis.shared.middleware.csrf import _EXEMPT_PATHS + + assert "/chat/proxy-stream" not in _EXEMPT_PATHS diff --git a/backend/tests/apis/app_api/chat/test_proxy_routes_streaming.py b/backend/tests/apis/app_api/chat/test_proxy_routes_streaming.py new file mode 100644 index 00000000..5c3b6a1c --- /dev/null +++ b/backend/tests/apis/app_api/chat/test_proxy_routes_streaming.py @@ -0,0 +1,171 @@ +"""Streaming-behavior integration test for `POST /chat/proxy-stream`. + +The migration plan calls out a buffering risk on SSE events that arrive +late in the stream — most importantly the `oauth_required` event, which +is emitted *after* `message_stop` and would be hidden from the SPA if +the proxy (or any intermediary) read the upstream response to completion +before flushing headers downstream. + +This test runs the proxy behind a real uvicorn server (NOT +`httpx.ASGITransport`, which buffers the entire body before exposing +headers — useless for measuring TTFB) and points it at a slow upstream +that yields one SSE chunk, sleeps, then yields more. We assert: + + - response headers reach the client well under 200ms (TTFB), + - `X-Accel-Buffering: no` is set on the response, + - total stream consumption takes at least the upstream gap, proving + the body is actually streamed (not buffered then flushed). +""" + +from __future__ import annotations + +import asyncio +import socket +import threading +import time + +import httpx +import pytest +import uvicorn +from fastapi import FastAPI + +from apis.app_api.chat import proxy_routes +from apis.app_api.chat.proxy_routes import router as proxy_router +from apis.shared.auth.dependencies import get_current_user_from_session +from apis.shared.auth.models import User + + +# Upstream delay between SSE chunks — short enough to keep the test fast, +# long enough that buffering vs. streaming is unambiguously distinguishable. +_UPSTREAM_GAP_SECONDS = 0.3 +_TTFB_BUDGET_SECONDS = 0.2 + + +def _user() -> User: + user = User( + email="alice@example.com", + user_id="user-sub", + name="Alice", + roles=["user"], + ) + user.raw_token = "access.token" + return user + + +def _free_port() -> int: + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: + s.bind(("127.0.0.1", 0)) + return s.getsockname()[1] + + +class _UvicornInThread: + """Tiny harness: run uvicorn in a daemon thread on a free port, block + until the socket accepts connections, and tear it down on exit.""" + + def __init__(self, app: FastAPI) -> None: + self.port = _free_port() + self.url = f"http://127.0.0.1:{self.port}" + self._server = uvicorn.Server( + uvicorn.Config( + app, + host="127.0.0.1", + port=self.port, + log_level="warning", + lifespan="off", + access_log=False, + ) + ) + self._thread: threading.Thread | None = None + + def __enter__(self) -> "_UvicornInThread": + self._thread = threading.Thread(target=self._server.run, daemon=True) + self._thread.start() + + # Wait for socket to accept connections — bounded so we fail fast + # if uvicorn never came up. + deadline = time.monotonic() + 5.0 + while time.monotonic() < deadline: + try: + with socket.create_connection(("127.0.0.1", self.port), timeout=0.1): + return self + except OSError: + time.sleep(0.05) + raise RuntimeError("uvicorn server failed to start within 5s") + + def __exit__(self, exc_type, exc, tb) -> None: + self._server.should_exit = True + if self._thread is not None: + self._thread.join(timeout=5.0) + + +def _build_app() -> FastAPI: + app = FastAPI() + app.include_router(proxy_router) + app.dependency_overrides[get_current_user_from_session] = _user + return app + + +def _patch_slow_upstream(monkeypatch: pytest.MonkeyPatch) -> None: + """Replace the proxy's upstream client with one whose MockTransport + yields one SSE chunk, sleeps, then yields more.""" + + async def slow_content(): + yield b'event: message_start\ndata: {"role": "assistant"}\n\n' + await asyncio.sleep(_UPSTREAM_GAP_SECONDS) + yield b'event: content_block_delta\ndata: {"text": "hi"}\n\n' + yield b'event: message_stop\ndata: {"stopReason": "end_turn"}\n\n' + yield b'event: oauth_required\ndata: {"providerId":"slack"}\n\n' + yield b"event: done\ndata: {}\n\n" + + async def handler(request: httpx.Request) -> httpx.Response: + return httpx.Response( + 200, + content=slow_content(), + headers={"content-type": "text/event-stream"}, + ) + + transport = httpx.MockTransport(handler) + monkeypatch.setattr( + proxy_routes, + "_build_upstream_client", + lambda: httpx.AsyncClient(transport=transport), + ) + + +@pytest.mark.asyncio +async def test_ttfb_under_200ms_with_x_accel_buffering( + monkeypatch: pytest.MonkeyPatch, +) -> None: + _patch_slow_upstream(monkeypatch) + app = _build_app() + + with _UvicornInThread(app) as server: + async with httpx.AsyncClient(base_url=server.url, timeout=10.0) as client: + t0 = time.monotonic() + async with client.stream( + "POST", "/chat/proxy-stream", json={"message": "hi"} + ) as response: + ttfb = time.monotonic() - t0 + assert response.status_code == 200 + assert response.headers["x-accel-buffering"] == "no" + assert response.headers["cache-control"] == "no-cache" + assert response.headers["content-type"].startswith( + "text/event-stream" + ) + assert ttfb < _TTFB_BUDGET_SECONDS, ( + f"TTFB {ttfb:.3f}s exceeded {_TTFB_BUDGET_SECONDS}s budget — " + "the proxy is buffering upstream before flushing headers." + ) + + chunks = [] + async for chunk in response.aiter_bytes(): + chunks.append(chunk) + body = b"".join(chunks) + total = time.monotonic() - t0 + + assert total >= _UPSTREAM_GAP_SECONDS, ( + f"Total {total:.3f}s shorter than upstream gap {_UPSTREAM_GAP_SECONDS}s " + "— upstream stream did not actually slow-yield." + ) + assert b"oauth_required" in body + assert b"event: done" in body