diff --git a/compose.yaml b/compose.yaml index 2a58dbde..7e059153 100644 --- a/compose.yaml +++ b/compose.yaml @@ -207,6 +207,7 @@ services: dockerfile: Dockerfile profiles: - production + - cluster command: - sh - -c diff --git a/docs/console.md b/docs/console.md index 76a83c64..874b6488 100644 --- a/docs/console.md +++ b/docs/console.md @@ -326,7 +326,7 @@ The server UI uses root-relative URLs (`/v1/api/send`, `/static/app.js`, `/share ### SSE Proxy -SSE streams (`/v1/api/events`, `/v1/api/events/global`) are proxied by creating a per-connection `httpx.AsyncClient(timeout=None)`, streaming the upstream response via `aiter_text()`, parsing SSE framing (`\n\n` delimiters), and re-emitting events through `EventSourceResponse`. Each proxied SSE stream requires its own httpx client since the shared client's 30-second timeout would kill long-lived connections. +SSE streams (`/v1/api/events`, `/v1/api/events/global`) are proxied as raw byte passthrough — the console opens an `httpx.AsyncClient.stream()` to the upstream server (with `read=None` and `pool=None` timeouts since SSE connections are long-lived) and relays every byte via `StreamingResponse`. This preserves server-side ping comments, event framing, and keepalives verbatim without parsing or re-encoding. ### Authentication diff --git a/tests/test_console.py b/tests/test_console.py index dfda851d..c2941f64 100644 --- a/tests/test_console.py +++ b/tests/test_console.py @@ -1,5 +1,6 @@ """Tests for turnstone.console — collector and HTTP server.""" +import asyncio import json import queue from unittest.mock import MagicMock @@ -1299,3 +1300,177 @@ class TestProxySharedStatic: resp = client.get("/node/unknown/shared/base.css") assert resp.status_code == 404 client.close() + + +# --------------------------------------------------------------------------- +# SSE proxy — raw byte passthrough +# --------------------------------------------------------------------------- + + +class TestSSEProxy: + """Verify _proxy_sse forwards raw bytes including ping comments.""" + + def test_proxy_sse_preserves_pings_and_events(self): + """SSE proxy should forward ping comments and events verbatim.""" + from turnstone.console.server import _proxy_sse + + # Simulate an upstream SSE response with a ping comment and a real event + sse_payload = b': ping - 2026-03-08T12:00:00Z\n\nevent: message\ndata: {"type": "test"}\n\n' + + class FakeResponse: + status_code = 200 + headers = {"content-type": "text/event-stream"} + + async def aiter_bytes(self): + yield sse_payload + + async def aclose(self): + pass + + async def __aenter__(self): + return self + + async def __aexit__(self, *args): + pass + + class FakeClient: + def stream(self, method, url, **kwargs): + return FakeResponse() + + class FakeRequest: + class url: # noqa: N801 + query = "ws_id=test123" + + class app: # noqa: N801 + class state: # noqa: N801 + proxy_sse_client = FakeClient() + proxy_auth_token = "" + + headers = {} + + async def is_disconnected(self): + return False + + async def _run(): + response = await _proxy_sse( + FakeRequest(), "http://fake:8080", "events", api_prefix="v1/api" + ) + assert response.media_type == "text/event-stream" + # Collect the streamed bytes + chunks: list[bytes] = [] + async for chunk in response.body_iterator: + chunks.append(chunk if isinstance(chunk, bytes) else chunk.encode()) + body = b"".join(chunks) + # Ping comment must be preserved (not filtered) + assert b": ping" in body + # Real event must be preserved + assert b"event: message" in body + assert b'"type": "test"' in body + + asyncio.run(_run()) + + def test_proxy_sse_upstream_error_status(self): + """Non-200 upstream status should yield an error event.""" + + from turnstone.console.server import _proxy_sse + + class FakeResponse: + status_code = 502 + + async def aiter_bytes(self): + return + yield # make it an async generator + + async def aclose(self): + pass + + async def __aenter__(self): + return self + + async def __aexit__(self, *args): + pass + + class FakeClient: + def stream(self, method, url, **kwargs): + return FakeResponse() + + class FakeRequest: + class url: # noqa: N801 + query = "" + + class app: # noqa: N801 + class state: # noqa: N801 + proxy_sse_client = FakeClient() + proxy_auth_token = "" + + headers = {} + + async def is_disconnected(self): + return False + + async def _run(): + response = await _proxy_sse(FakeRequest(), "http://fake:8080", "events") + chunks: list[bytes] = [] + async for chunk in response.body_iterator: + chunks.append(chunk if isinstance(chunk, bytes) else chunk.encode()) + body = b"".join(chunks) + assert b"event: error" in body + assert b"502" in body + + asyncio.run(_run()) + + def test_proxy_sse_disconnect_handling(self): + """Proxy should stop when browser disconnects.""" + + from turnstone.console.server import _proxy_sse + + class FakeResponse: + status_code = 200 + + async def aiter_bytes(self): + yield b"data: chunk1\n\n" + yield b"data: chunk2\n\n" # should not be reached + yield b"data: chunk3\n\n" + + async def aclose(self): + pass + + async def __aenter__(self): + return self + + async def __aexit__(self, *args): + pass + + class FakeClient: + def stream(self, method, url, **kwargs): + return FakeResponse() + + call_count = 0 + + class FakeRequest: + class url: # noqa: N801 + query = "" + + class app: # noqa: N801 + class state: # noqa: N801 + proxy_sse_client = FakeClient() + proxy_auth_token = "" + + headers = {} + + async def is_disconnected(self): + nonlocal call_count + call_count += 1 + return call_count > 1 # disconnect after first chunk + + async def _run(): + response = await _proxy_sse(FakeRequest(), "http://fake:8080", "events") + chunks: list[bytes] = [] + async for chunk in response.body_iterator: + chunks.append(chunk if isinstance(chunk, bytes) else chunk.encode()) + body = b"".join(chunks) + assert b"chunk1" in body + # Should have stopped before chunk3 + assert b"chunk3" not in body + + asyncio.run(_run()) diff --git a/turnstone/console/server.py b/turnstone/console/server.py index 809df4db..969be692 100644 --- a/turnstone/console/server.py +++ b/turnstone/console/server.py @@ -30,7 +30,7 @@ import httpx from sse_starlette import EventSourceResponse from starlette.applications import Starlette from starlette.middleware import Middleware -from starlette.responses import HTMLResponse, JSONResponse, Response +from starlette.responses import HTMLResponse, JSONResponse, Response, StreamingResponse from starlette.routing import Mount, Route from starlette.staticfiles import StaticFiles @@ -568,7 +568,11 @@ async def _proxy_post( async def _proxy_sse( request: Request, server_url: str, path: str, *, api_prefix: str = "api" ) -> Response: - """Proxy an SSE stream from the target server to the browser.""" + """Proxy an SSE stream from the target server to the browser. + + Relays raw bytes verbatim so server-side ping comments, event framing, + and keepalives all pass through unchanged. + """ target = f"{server_url}/{api_prefix}/{path}" if request.url.query: target += f"?{request.url.query}" @@ -576,30 +580,34 @@ async def _proxy_sse( sse_client: httpx.AsyncClient = request.app.state.proxy_sse_client sse_auth = _proxy_auth_headers(request) - async def sse_generator() -> AsyncGenerator[dict[str, str], None]: - from httpx_sse import aconnect_sse - + async def raw_stream() -> AsyncGenerator[bytes, None]: try: - async with aconnect_sse(sse_client, "GET", target, headers=sse_auth) as source: - if source.response.status_code != 200: + async with sse_client.stream( + "GET", + target, + headers={**sse_auth, "Accept": "text/event-stream", "Cache-Control": "no-store"}, + timeout=httpx.Timeout(connect=10, read=None, write=5, pool=None), + ) as response: + if response.status_code != 200: log.debug( "SSE proxy received status %s from %s", - source.response.status_code, + response.status_code, target, ) - yield { - "event": "error", - "data": f"Upstream returned status {source.response.status_code}", - } + yield f"event: error\ndata: Upstream returned status {response.status_code}\n\n".encode() return - async for sse in source.aiter_sse(): + async for chunk in response.aiter_bytes(): if await request.is_disconnected(): return - yield {"event": sse.event, "data": sse.data} + yield chunk except httpx.HTTPError: log.debug("SSE proxy stream ended for %s", target) - return EventSourceResponse(sse_generator(), ping=5) + return StreamingResponse( + raw_stream(), + media_type="text/event-stream", + headers={"Cache-Control": "no-store", "X-Accel-Buffering": "no"}, + ) # ---------------------------------------------------------------------------