mirror of
https://github.com/turnstonelabs/turnstone.git
synced 2026-08-13 23:42:25 -06:00
Compare commits
2 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
| 136b75fdef | |||
| 4d1107839b |
@@ -207,6 +207,7 @@ services:
|
||||
dockerfile: Dockerfile
|
||||
profiles:
|
||||
- production
|
||||
- cluster
|
||||
command:
|
||||
- sh
|
||||
- -c
|
||||
|
||||
+1
-1
@@ -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
|
||||
|
||||
|
||||
+1
-1
@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
|
||||
|
||||
[project]
|
||||
name = "turnstone"
|
||||
version = "0.4.6"
|
||||
version = "0.5.0"
|
||||
description = "Multi-node AI orchestration platform with tool use, agent routing, and cluster simulation."
|
||||
readme = "README.md"
|
||||
license = "BUSL-1.1"
|
||||
|
||||
@@ -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())
|
||||
|
||||
@@ -1,3 +1,3 @@
|
||||
"""turnstone - Multi-node AI orchestration platform with tool use, agent routing, and cluster simulation."""
|
||||
|
||||
__version__ = "0.4.6"
|
||||
__version__ = "0.5.0"
|
||||
|
||||
+23
-15
@@ -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"},
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
Reference in New Issue
Block a user