diff --git a/api/config.py b/api/config.py index 5f0ffc06..36ac1806 100644 --- a/api/config.py +++ b/api/config.py @@ -14,6 +14,7 @@ import copy import json import logging import os +import queue import sys import threading import time @@ -2745,6 +2746,53 @@ _INDEX_HTML_PATH = REPO_ROOT / "static" / "index.html" LOCK = threading.Lock() SESSIONS_MAX = 100 CHAT_LOCK = threading.Lock() + + +class StreamChannel: + """Broadcast SSE events to every connected browser tab for a stream. + + While no tab is connected, events are buffered so the first/reconnected + subscriber still receives the stream tail that arrived during the gap. + Once one or more subscribers are attached, new events are broadcast to all + of them instead of being consumed destructively by a single queue reader. + """ + + def __init__(self): + self._lock = threading.Lock() + self._subscribers: list[queue.Queue] = [] + self._offline_buffer: list[tuple[str, object]] = [] + + def subscribe(self) -> queue.Queue: + q: queue.Queue = queue.Queue() + with self._lock: + snapshot = list(self._offline_buffer) + self._subscribers.append(q) + for item in snapshot: + q.put_nowait(item) + return q + + def unsubscribe(self, q: queue.Queue) -> None: + with self._lock: + try: + self._subscribers.remove(q) + except ValueError: + pass + + def put_nowait(self, item: tuple[str, object]) -> None: + with self._lock: + subscribers = list(self._subscribers) + if not subscribers: + self._offline_buffer.append(item) + return + self._offline_buffer.clear() + for q in subscribers: + q.put_nowait(item) + + +def create_stream_channel() -> StreamChannel: + return StreamChannel() + + STREAMS: dict = {} STREAMS_LOCK = threading.Lock() CANCEL_FLAGS: dict = {} diff --git a/api/routes.py b/api/routes.py index 3c8b8086..c642678f 100644 --- a/api/routes.py +++ b/api/routes.py @@ -332,6 +332,7 @@ from api.config import ( get_reasoning_status, set_reasoning_display, set_reasoning_effort, + create_stream_channel, ) from api.helpers import ( require, @@ -3649,9 +3650,10 @@ def _handle_list_dir(handler, parsed): def _handle_sse_stream(handler, parsed): stream_id = parse_qs(parsed.query).get("stream_id", [""])[0] - q = STREAMS.get(stream_id) - if q is None: + stream = STREAMS.get(stream_id) + if stream is None: return j(handler, {"error": "stream not found"}, status=404) + subscriber = stream.subscribe() if hasattr(stream, "subscribe") else stream handler.send_response(200) handler.send_header("Content-Type", "text/event-stream; charset=utf-8") handler.send_header("Cache-Control", "no-cache") @@ -3661,7 +3663,7 @@ def _handle_sse_stream(handler, parsed): try: while True: try: - event, data = q.get(timeout=30) + event, data = subscriber.get(timeout=30) except queue.Empty: handler.wfile.write(b": heartbeat\n\n") handler.wfile.flush() @@ -3671,6 +3673,12 @@ def _handle_sse_stream(handler, parsed): break except _CLIENT_DISCONNECT_ERRORS: pass + finally: + if subscriber is not stream and hasattr(stream, "unsubscribe"): + try: + stream.unsubscribe(subscriber) + except Exception: + pass return True @@ -4812,9 +4820,9 @@ def _handle_btw(handler, body): stream_id = uuid.uuid4().hex ephemeral.active_stream_id = stream_id ephemeral.save() - q = queue.Queue() + stream = create_stream_channel() with STREAMS_LOCK: - STREAMS[stream_id] = q + STREAMS[stream_id] = stream from api.background import track_btw track_btw(body["session_id"], ephemeral.session_id, stream_id, question) thr = threading.Thread( @@ -4858,9 +4866,9 @@ def _handle_background(handler, body): stream_id = uuid.uuid4().hex bg.active_stream_id = stream_id bg.save() - q = queue.Queue() + stream = create_stream_channel() with STREAMS_LOCK: - STREAMS[stream_id] = q + STREAMS[stream_id] = stream task_id = uuid.uuid4().hex[:8] from api.background import track_background, complete_background parent_sid = body["session_id"] @@ -4974,9 +4982,9 @@ def _handle_chat_start(handler, body): s.pending_started_at = time.time() s.save() set_last_workspace(workspace) - q = queue.Queue() + stream = create_stream_channel() with STREAMS_LOCK: - STREAMS[stream_id] = q + STREAMS[stream_id] = stream thr = threading.Thread( target=_run_agent_streaming, args=(s.session_id, msg, model, workspace, stream_id, attachments), diff --git a/tests/test_issue_1584_multitab_sse.py b/tests/test_issue_1584_multitab_sse.py new file mode 100644 index 00000000..8341d317 --- /dev/null +++ b/tests/test_issue_1584_multitab_sse.py @@ -0,0 +1,83 @@ +import io +import threading +from types import SimpleNamespace + +from api.config import STREAMS, STREAMS_LOCK, create_stream_channel +from api.routes import _handle_sse_stream + + +class _FakeHandler: + def __init__(self): + self.status = None + self.headers = [] + self.wfile = io.BytesIO() + + def send_response(self, status): + self.status = status + + def send_header(self, key, value): + self.headers.append((key, value)) + + def end_headers(self): + return None + + +def test_stream_channel_broadcasts_each_event_to_every_subscriber(): + stream = create_stream_channel() + q1 = stream.subscribe() + q2 = stream.subscribe() + + try: + stream.put_nowait(("token", {"text": "H"})) + stream.put_nowait(("token", {"text": "allo"})) + stream.put_nowait(("stream_end", {"status": "done"})) + + assert q1.get(timeout=1) == ("token", {"text": "H"}) + assert q1.get(timeout=1) == ("token", {"text": "allo"}) + assert q1.get(timeout=1) == ("stream_end", {"status": "done"}) + + assert q2.get(timeout=1) == ("token", {"text": "H"}) + assert q2.get(timeout=1) == ("token", {"text": "allo"}) + assert q2.get(timeout=1) == ("stream_end", {"status": "done"}) + finally: + stream.unsubscribe(q1) + stream.unsubscribe(q2) + + +def test_same_stream_in_two_tabs_receives_identical_token_sequence(): + stream_id = "multitab-stream" + stream = create_stream_channel() + with STREAMS_LOCK: + STREAMS[stream_id] = stream + + handlers = [_FakeHandler(), _FakeHandler()] + threads = [ + threading.Thread( + target=_handle_sse_stream, + args=(handler, SimpleNamespace(query=f"stream_id={stream_id}")), + daemon=True, + ) + for handler in handlers + ] + + try: + for thread in threads: + thread.start() + + stream.put_nowait(("token", {"text": "H"})) + stream.put_nowait(("token", {"text": "allo"})) + stream.put_nowait(("stream_end", {"status": "done"})) + + for thread in threads: + thread.join(timeout=1) + assert not thread.is_alive(), "every tab should finish the same SSE stream" + + for handler in handlers: + payload = handler.wfile.getvalue().decode("utf-8") + assert handler.status == 200 + assert '"text": "H"' in payload + assert '"text": "allo"' in payload + assert "event: stream_end" in payload + finally: + with STREAMS_LOCK: + STREAMS.pop(stream_id, None)