From 70d495aa5bf5f417613996ef171d654fa8aca7b6 Mon Sep 17 00:00:00 2001 From: Patrick Buckley Date: Tue, 10 Mar 2026 13:25:16 -0700 Subject: [PATCH] =?UTF-8?q?fix:=20per-workstream=20SSE=20fan-out=20?= =?UTF-8?q?=E2=80=94=20multiple=20consumers=20no=20longer=20steal=20?= =?UTF-8?q?=E2=80=A6=20(#38)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * fix: per-workstream SSE fan-out — multiple consumers no longer steal each other's tokens After 4d665a5 removed the single-consumer SSE lock, the shared _event_queue let concurrent consumers (browser, bridge, console proxy) race on Queue.get(), each receiving ~1/N of content tokens and producing garbled streaming text. Replace the single queue with per-client fan-out: each SSE connection registers its own bounded queue (maxsize=500) on WebUI._listeners, and _enqueue() copies every event to all registered queues. On eviction or close, a ws_closed sentinel is injected so SSE generators exit promptly. * fix: address CI failures and Copilot review feedback - Handle ws_closed sentinel in events_sse generator (break on close) - Guarantee sentinel delivery by evicting one item when queue is full - Clear listeners list after injecting sentinels on cleanup - Fix test_slow_consumer to fill only slow queue directly - Fix ruff SIM117 (nested with), unused import, mypy unused-ignore --- docs/api-reference.md | 12 +- docs/architecture.md | 2 +- docs/diagrams/03-core-engine-classes.puml | 2 +- docs/diagrams/png/03-core-engine-classes.png | 4 +- tests/test_workstream.py | 111 +++++++++++++++++++ turnstone/core/workstream.py | 15 +++ turnstone/server.py | 37 +++++-- 7 files changed, 161 insertions(+), 22 deletions(-) diff --git a/docs/api-reference.md b/docs/api-reference.md index 497a5f23..f8f41579 100644 --- a/docs/api-reference.md +++ b/docs/api-reference.md @@ -460,13 +460,13 @@ The server sends an SSE comment every 5 seconds when no events are pending: This prevents proxies and browsers from closing the connection due to inactivity. -#### Generation mechanism +#### Multi-consumer fan-out -Each new SSE connection to a workstream increments an internal -`_sse_generation` counter. The previous SSE handler detects the generation -mismatch and exits its event loop, ensuring only one active SSE connection per -workstream at a time. The event queue is drained of stale events before the new -connection begins streaming. +Each SSE connection to a workstream receives its own delivery queue. Events +produced by the worker thread are fanned out to all registered listener queues, +so multiple consumers (browser, bridge, console proxy, SDK) can connect +simultaneously and each receives every event. On reconnect the client receives +a full history replay, so no catch-up mechanism is needed. --- diff --git a/docs/architecture.md b/docs/architecture.md index 649d6ded..aab2f20a 100644 --- a/docs/architecture.md +++ b/docs/architecture.md @@ -1113,7 +1113,7 @@ context manager handles startup/shutdown (health monitor, MCP client, registry). Each workstream's `WebUI` has: -- `_event_queue` (per-workstream SSE events, `queue.Queue`) +- `_listeners` (per-client SSE queues, fan-out on `_enqueue()`) - `_approval_event` / `_plan_event` (`threading.Event` for blocking) - `_global_queue` (class variable, shared, for state broadcasts) diff --git a/docs/diagrams/03-core-engine-classes.puml b/docs/diagrams/03-core-engine-classes.puml index 8937e890..f32f4034 100644 --- a/docs/diagrams/03-core-engine-classes.puml +++ b/docs/diagrams/03-core-engine-classes.puml @@ -41,7 +41,7 @@ class "WorkstreamTerminalUI" as WsTermUI { } class "WebUI" as WebUI { - - _event_queue: Queue + - _listeners: list[Queue] - _approval_event: Event - _plan_event: Event - _ws_prompt_tokens: int diff --git a/docs/diagrams/png/03-core-engine-classes.png b/docs/diagrams/png/03-core-engine-classes.png index c102247c..065fbb4e 100644 --- a/docs/diagrams/png/03-core-engine-classes.png +++ b/docs/diagrams/png/03-core-engine-classes.png @@ -1,3 +1,3 @@ version https://git-lfs.github.com/spec/v1 -oid sha256:c53ddce800c59f9432d7a016c7d66282a449555d452b7fe9dd393f4282f08c46 -size 554721 +oid sha256:760c37e67736588dadee21d500419a48e9fc50f8bdc5667e686c580022bd40e2 +size 554869 diff --git a/tests/test_workstream.py b/tests/test_workstream.py index 74d1ff36..95913bce 100644 --- a/tests/test_workstream.py +++ b/tests/test_workstream.py @@ -683,6 +683,117 @@ class TestWebUI: t.join() +# --------------------------------------------------------------------------- +# WebUI SSE fan-out +# --------------------------------------------------------------------------- + + +class TestWebUIFanOut: + """Verify per-client SSE fan-out on WebUI._enqueue / _register_listener.""" + + def test_enqueue_no_listeners(self): + """Events silently dropped when no listeners are registered.""" + from turnstone.server import WebUI + + ui = WebUI(ws_id="test") + ui._enqueue({"type": "content", "text": "hello"}) # should not raise + + def test_enqueue_single_listener(self): + """Single listener receives the event.""" + from turnstone.server import WebUI + + ui = WebUI(ws_id="test") + q = ui._register_listener() + ui._enqueue({"type": "content", "text": "hello"}) + assert q.get_nowait() == {"type": "content", "text": "hello"} + + def test_enqueue_multiple_listeners(self): + """All registered listeners receive an identical copy.""" + from turnstone.server import WebUI + + ui = WebUI(ws_id="test") + q1 = ui._register_listener() + q2 = ui._register_listener() + q3 = ui._register_listener() + + event = {"type": "content", "text": "world"} + ui._enqueue(event) + + assert q1.get_nowait() == event + assert q2.get_nowait() == event + assert q3.get_nowait() == event + + def test_unregister_stops_delivery(self): + """After unregister, the queue receives no further events.""" + import queue as queue_mod + + from turnstone.server import WebUI + + ui = WebUI(ws_id="test") + q = ui._register_listener() + ui._unregister_listener(q) + ui._enqueue({"type": "content", "text": "gone"}) + + with pytest.raises(queue_mod.Empty): + q.get_nowait() + + def test_slow_consumer_does_not_block(self): + """A full queue doesn't block the producer or starve other listeners.""" + from turnstone.server import WebUI + + ui = WebUI(ws_id="test") + slow = ui._register_listener() + fast = ui._register_listener() + + # Fill only the slow consumer's queue directly to capacity + for i in range(500): + slow.put_nowait({"type": "content", "text": f"fill-{i}"}) + + assert slow.qsize() == 500 + assert fast.qsize() == 0 + + # Enqueue via fan-out — slow drops (full), fast receives + event = {"type": "content", "text": "overflow"} + ui._enqueue(event) + assert slow.qsize() == 500 # still full, overflow dropped + assert fast.qsize() == 1 + assert fast.get_nowait() == event + + def test_unregister_idempotent(self): + """Double unregister does not raise.""" + + from turnstone.server import WebUI + + ui = WebUI(ws_id="test") + q = ui._register_listener() + ui._unregister_listener(q) + ui._unregister_listener(q) # should not raise + + def test_concurrent_enqueue_and_register(self): + """Concurrent register/unregister and enqueue should not crash.""" + from turnstone.server import WebUI + + ui = WebUI(ws_id="test") + stop = threading.Event() + + def register_loop(): + while not stop.is_set(): + q = ui._register_listener() + ui._unregister_listener(q) + + def enqueue_loop(): + for i in range(500): + ui._enqueue({"type": "content", "text": f"tok-{i}"}) + + t1 = threading.Thread(target=register_loop) + t2 = threading.Thread(target=enqueue_loop) + t1.start() + t2.start() + t2.join() + stop.set() + t1.join() + + # --------------------------------------------------------------------------- # Integration: WorkstreamManager + session state transitions # --------------------------------------------------------------------------- diff --git a/turnstone/core/workstream.py b/turnstone/core/workstream.py index f6c542e0..7a12c0ed 100644 --- a/turnstone/core/workstream.py +++ b/turnstone/core/workstream.py @@ -213,6 +213,21 @@ class WorkstreamManager: ws.ui._plan_event.set() if hasattr(ws.ui, "_fg_event"): ws.ui._fg_event.set() + # Notify SSE listeners so generators exit promptly + if hasattr(ws.ui, "_listeners_lock"): + import contextlib + import queue as _queue + + with ws.ui._listeners_lock: + for lq in ws.ui._listeners: # type: ignore[attr-defined] + try: + lq.put_nowait({"type": "ws_closed"}) + except _queue.Full: + with contextlib.suppress(_queue.Empty): + lq.get_nowait() + with contextlib.suppress(_queue.Full): + lq.put_nowait({"type": "ws_closed"}) + ws.ui._listeners.clear() # type: ignore[attr-defined] # Release MCP listener registration if ws.session and hasattr(ws.session, "close"): ws.session.close() diff --git a/turnstone/server.py b/turnstone/server.py index b3840461..7badcb56 100644 --- a/turnstone/server.py +++ b/turnstone/server.py @@ -82,8 +82,8 @@ class WebUI: def __init__(self, ws_id: str = "") -> None: self.ws_id = ws_id - self._event_queue: queue.Queue[dict[str, Any]] = queue.Queue() - self._sse_generation = 0 # incremented on each new SSE connection + self._listeners: list[queue.Queue[dict[str, Any]]] = [] + self._listeners_lock = threading.Lock() self._approval_event = threading.Event() self._approval_result: tuple[bool, str | None] = (False, None) self._pending_approval: dict[str, Any] | None = None # re-sent on SSE reconnect @@ -102,7 +102,23 @@ class WebUI: self._ws_activity_state: str = "" # "tool" | "approval" | "thinking" | "" def _enqueue(self, data: dict[str, Any]) -> None: - self._event_queue.put(data) + with self._listeners_lock: + snapshot = list(self._listeners) + for lq in snapshot: + with contextlib.suppress(queue.Full): + lq.put_nowait(data) + + def _register_listener(self) -> queue.Queue[dict[str, Any]]: + """Create a per-client queue and register it as a listener.""" + client_queue: queue.Queue[dict[str, Any]] = queue.Queue(maxsize=500) + with self._listeners_lock: + self._listeners.append(client_queue) + return client_queue + + def _unregister_listener(self, client_queue: queue.Queue[dict[str, Any]]) -> None: + """Remove a client queue from the listeners list.""" + with self._listeners_lock, contextlib.suppress(ValueError): + self._listeners.remove(client_queue) def _broadcast_state(self, state: str) -> None: """Send a state-change event to the global SSE channel.""" @@ -473,14 +489,8 @@ async def events_sse(request: Request) -> Response: if not ws or not ui: return JSONResponse({"error": "Unknown workstream"}, status_code=404) - # Drain stale events so this client starts fresh. A race with the - # worker thread is acceptable: worst case we discard one fresh event, - # and the client catches up via the history replay below. - while not ui._event_queue.empty(): - try: - ui._event_queue.get_nowait() - except queue.Empty: - break + # Each client gets its own queue — no drain needed. + client_queue = ui._register_listener() async def event_generator() -> AsyncGenerator[dict[str, str], None]: assert ws.session is not None @@ -510,13 +520,16 @@ async def events_sse(request: Request) -> Response: while True: try: event = await loop.run_in_executor( - None, functools.partial(ui._event_queue.get, timeout=5) + None, functools.partial(client_queue.get, timeout=5) ) + if event.get("type") == "ws_closed": + return yield {"data": json.dumps(event)} except queue.Empty: pass finally: _metrics.record_sse_disconnect() + ui._unregister_listener(client_queue) return EventSourceResponse(event_generator(), ping=5)