diff --git a/api/config.py b/api/config.py index e5e0435f..79a47994 100644 --- a/api/config.py +++ b/api/config.py @@ -8401,9 +8401,30 @@ class StreamChannel: # recoverable via the run journal by last_event_id. 8192 is generous enough # to hold a long multi-tool turn's backlog while capping worst-case memory to # a fixed number of small (event, data, id) tuples — deliberately far above - # SessionChannel's per-subscriber maxsize of 16 (that queue drops on a *slow* - # reader; this buffer must survive a legitimate reconnect gap). + # the per-subscriber queue cap below (that queue drops on a *slow* reader; + # this buffer must survive a legitimate reconnect gap). _OFFLINE_BUFFER_MAXLEN = 8192 + # Per-subscriber queue cap (drop-oldest on full). Each connected tab gets its + # own queue; a slow/backpressured or backgrounded tab used to hold an + # UNBOUNDED queue.Queue that grew for the WHOLE turn (the producer is the + # agent token stream), an OOM risk with many tabs × long agentic turns. This + # caps the per-tab live-broadcast growth to a fixed bound. + # + # Bound is intentionally EQUAL to _OFFLINE_BUFFER_MAXLEN, not the much + # smaller SessionChannel per-subscriber cap of 16. StreamChannel carries the + # live chat token stream (thousands of frames per turn) and, unlike + # SessionChannel's low-frequency UI pings, has a reconnect-replay contract: + # a tab that briefly disconnected must receive the full retained offline + # tail on resubscribe. Capping below the offline buffer would truncate that + # replay and force every flaky-network reconnect through the run journal + # (disk reads) instead of the in-memory fast path. Matching the offline + # buffer bound preserves that contract while bounding live-broadcast memory + # to the SAME worst-case #4633 already accepted (a fixed number of small + # (event, data, id) tuples). The SSE write deadline + # (SSE_WRITE_DEADLINE_SECONDS) independently breaks a stuck socket within + # ~20s, so the overflow window is short; this cap bounds it by frame count + # too. Older frames stay recoverable via the run journal by last_event_id. + _SUBSCRIBER_QUEUE_MAXSIZE = _OFFLINE_BUFFER_MAXLEN def __init__(self): self._lock = threading.Lock() @@ -8424,6 +8445,9 @@ class StreamChannel: # Cumulative evictions over the channel's lifetime, never reset — for ops # visibility via diagnostic_snapshot(). self._offline_dropped_total = 0 + # Cumulative per-subscriber queue drops over the channel's lifetime + # (broadcast + replay paths), never reset — ops visibility for slow tabs. + self._subscriber_dropped_total = 0 self._last_event_id: str | None = None def subscribe(self) -> queue.Queue: @@ -8431,15 +8455,48 @@ class StreamChannel: return q def subscribe_with_snapshot(self) -> tuple[queue.Queue, dict[str, object]]: - q: queue.Queue = queue.Queue() + q: queue.Queue = queue.Queue(maxsize=self._SUBSCRIBER_QUEUE_MAXSIZE) with self._lock: # Replay buffered events to the new subscriber INSIDE the lock so a # concurrent put_nowait() can't broadcast a newer event before we - # finish replaying the older buffered tail. queue.Queue.put_nowait - # is non-blocking on an unbounded queue, so holding the lock here - # is safe. Per Opus advisor on stage-292. + # finish replaying the older buffered tail. The queue is bounded, so + # put_nowait raises queue.Full once the cap is reached — drop the + # OLDEST already-replayed frame and retry, keeping the most recent + # tail (a reconnecting tab needs the tail; older frames stay + # recoverable via the run journal by last_event_id). Holding the + # lock here is safe: no other put_nowait() can interleave. Per Opus + # advisor on stage-292. + replayed_dropped = 0 for item in self._offline_buffer: - q.put_nowait(item) + while True: + try: + q.put_nowait(item) + break + except queue.Full: + # Drop oldest to make room for the newer (more useful) + # tail frame. The drained frame is the oldest in this + # subscriber's replay window only. + try: + q.get_nowait() + except queue.Empty: + # A concurrent consumer drained the queue between + # our Full and get_nowait — the queue now has space, + # so retry the put instead of dropping `item`. This + # path runs under self._lock with a freshly-created + # queue (no concurrent consumer), so it is not + # reached in practice, but `continue` is the + # correct, race-safe rule (see the broadcast path). + continue + replayed_dropped += 1 + if replayed_dropped: + self._subscriber_dropped_total += replayed_dropped + logger.debug( + "StreamChannel subscriber replay dropped %d oldest frames " + "(cap=%d) while catching up on %d buffered events", + replayed_dropped, + self._SUBSCRIBER_QUEUE_MAXSIZE, + len(self._offline_buffer), + ) first = self._offline_buffer[0] if self._offline_buffer else None snapshot = { "offline_buffered_events": len(self._offline_buffer), @@ -8500,8 +8557,42 @@ class StreamChannel: # reports and logs its own truncation, not a stale carry-over. self._offline_buffer.clear() self._offline_dropped = 0 + # Broadcast outside the lock so a slow put_nowait doesn't block other + # subscribers or producers. The queue is bounded; on queue.Full drop the + # OLDEST frame and retry so a slow/backpressured tab keeps its most + # recent tail instead of growing unbounded for the whole turn. Older + # frames stay recoverable via the run journal by last_event_id. Mirrors + # SessionChannel.emit's drop-on-full contract. + broadcast_dropped = 0 for q in subscribers: - q.put_nowait(item) + while True: + try: + q.put_nowait(item) + break + except queue.Full: + try: + q.get_nowait() + except queue.Empty: + # A concurrent consumer (the SSE handler thread) + # drained the queue between our Full and get_nowait. + # The queue now has space — retry the put so `item` is + # delivered. `break` here would silently discard `item`, + # and if `item` is a terminal frame (stream_end/error/ + # cancel) the subscriber never receives it and the client + # stays attached indefinitely (spinner-forever). Having + # space is exactly the condition we want, so continue. + continue + broadcast_dropped += 1 + if broadcast_dropped: + with self._lock: + self._subscriber_dropped_total += broadcast_dropped + logger.debug( + "StreamChannel broadcast dropped %d oldest frames across %d " + "subscriber queue(s) (cap=%d)", + broadcast_dropped, + len(subscribers), + self._SUBSCRIBER_QUEUE_MAXSIZE, + ) def diagnostic_snapshot(self) -> dict[str, object]: """Return non-sensitive stream observation counters for health checks.""" @@ -8512,6 +8603,9 @@ class StreamChannel: # Cumulative over the channel lifetime (ops visibility), vs. the # per-cycle count subscribe_with_snapshot() reports for truncation. "offline_dropped_events": self._offline_dropped_total, + # Cumulative per-subscriber queue drops (replay + broadcast) over + # the channel lifetime — surfaces slow/backpressured tabs. + "subscriber_dropped_events": self._subscriber_dropped_total, } diff --git a/tests/test_stream_subscriber_queue_cap.py b/tests/test_stream_subscriber_queue_cap.py new file mode 100644 index 00000000..7433df32 --- /dev/null +++ b/tests/test_stream_subscriber_queue_cap.py @@ -0,0 +1,246 @@ +"""Regression tests: ``StreamChannel`` per-subscriber queues are bounded. + +Each connected browser tab gets its own ``queue.Queue`` to receive broadcast +SSE frames. That queue used to be ``queue.Queue()`` (no maxsize), so a slow / +backpressured / backgrounded tab accumulated every coalesced token frame for the +WHOLE turn (the producer is the agent token stream) — an OOM risk with many tabs +× long agentic turns. The sibling offline buffer was already capped (#4633), but +once a subscriber attached, events flowed into its unbounded per-subscriber +queue. + +The subscriber queue is now ``queue.Queue(maxsize=_SUBSCRIBER_QUEUE_MAXSIZE)`` +and both the broadcast path (``put_nowait``) and the offline-buffer replay path +(``subscribe_with_snapshot``) drop the OLDEST frame on ``queue.Full`` so a slow +tab keeps the most recent tail (older frames stay recoverable via the run +journal by ``last_event_id``). + +The cap equals ``_OFFLINE_BUFFER_MAXLEN`` so a reconnecting tab can still replay +the full retained offline tail (the #4633 reconnect contract); the point of the +bound is to cap the previously-unbounded LIVE-broadcast-to-slow-tab growth — a +50K-frame turn can no longer pin all 50000 frames per slow tab. +""" +from __future__ import annotations + +import queue as _queue + +from api.config import StreamChannel, create_stream_channel + + +def _drain(q: _queue.Queue) -> list: + out = [] + while True: + try: + out.append(q.get_nowait()) + except _queue.Empty: + break + return out + + +def test_subscriber_queue_has_a_maxsize_bound(): + """The regression: a subscriber queue MUST be bounded (it used to be 0 == + unbounded). Assert the cap is set and finite, and matches the documented + reconnect-replay bound.""" + ch = create_stream_channel() + q = ch.subscribe() + assert q.maxsize > 0 + assert q.maxsize == StreamChannel._SUBSCRIBER_QUEUE_MAXSIZE + assert q.maxsize == StreamChannel._OFFLINE_BUFFER_MAXLEN + + +def test_broadcast_drops_oldest_when_subscriber_never_drains(): + """A subscriber that never drains stays bounded by the queue cap and keeps + the most recent tail (drop-oldest), not the oldest frames — even when the + producer emits far more frames than the cap.""" + ch = create_stream_channel() + cap = StreamChannel._SUBSCRIBER_QUEUE_MAXSIZE + q = ch.subscribe() + + # Subscribe a second, draining subscriber so put_nowait takes the live + # broadcast path (subscribers present) rather than the offline buffer. + drainer = ch.subscribe() + total = cap + 500 # produce more than the cap; the slow tab must drop + for i in range(total): + ch.put_nowait(("token", i, f"id{i}")) + # Keep the drainer empty so broadcast always has a live audience. + _drain(drainer) + + # The never-drained subscriber never exceeded the cap. + assert q.qsize() <= cap + kept = _drain(q) + # The most recent `cap` frames are retained; the 500 older ones were dropped. + assert len(kept) == cap + assert kept[0] == ("token", total - cap, f"id{total - cap}") + assert kept[-1] == ("token", total - 1, f"id{total - 1}") + + # Ops counter surfaced the cumulative drops. + snap = ch.diagnostic_snapshot() + assert snap["subscriber_dropped_events"] >= 500 + + +def test_replay_drops_oldest_when_offline_buffer_exceeds_queue_cap(): + """subscribe_with_snapshot replays buffered frames into a bounded queue. The + default cap equals the offline buffer, so replay just fits and never drops. + To exercise the replay drop-on-full path, monkeypatch the cap SMALLER than + the buffer, then call the REAL subscribe_with_snapshot (not a hand-mirrored + loop) so the test stays coupled to the production code path under test.""" + ch = create_stream_channel() + buffer_len = StreamChannel._OFFLINE_BUFFER_MAXLEN + # Drive the offline buffer with no subscriber attached (fills buffer past + # overflow, so the retained tail is exactly `buffer_len` frames). + for i in range(buffer_len + 100): + ch.put_nowait(("token", i, f"id{i}")) + + small_cap = 8 + original_cap = StreamChannel._SUBSCRIBER_QUEUE_MAXSIZE + StreamChannel._SUBSCRIBER_QUEUE_MAXSIZE = small_cap + try: + # The default cap equals the buffer, so a resubscribe normally replays + # the full retained tail. With the cap forced small, replay must drop + # oldest and keep the newest `small_cap` frames — exercising the real + # drop-on-full branch inside subscribe_with_snapshot. + q, snapshot = ch.subscribe_with_snapshot() + finally: + StreamChannel._SUBSCRIBER_QUEUE_MAXSIZE = original_cap + + drained = _drain(q) + assert len(drained) == small_cap + # Newest tail retained. + assert drained[-1] == ("token", buffer_len + 100 - 1, f"id{buffer_len + 100 - 1}") + # Oldest retained frame is `small_cap` back from the end. + assert drained[0] == ("token", buffer_len + 100 - small_cap, f"id{buffer_len + 100 - small_cap}") + # The cumulative drop counter was incremented by the replay path. + assert ch._subscriber_dropped_total >= 1 + + +def test_broadcast_to_many_subscribers_stays_bounded(): + """Multiple subscribers, one of which never drains, never exceeds any + queue's cap and the draining subscribers receive every frame.""" + ch = create_stream_channel() + cap = StreamChannel._SUBSCRIBER_QUEUE_MAXSIZE + slow = ch.subscribe() + fast = ch.subscribe() + + total = cap + 10 # a little over the cap is enough to force one drop + fast_seen = [] + for i in range(total): + ch.put_nowait(("token", i, f"id{i}")) + fast_seen.extend(_drain(fast)) + + # The fast (draining) subscriber got every frame in order. + assert fast_seen == [("token", i, f"id{i}") for i in range(total)] + # The slow subscriber stayed bounded and kept the newest tail. + assert slow.qsize() <= cap + slow_kept = _drain(slow) + assert len(slow_kept) == cap + assert slow_kept[-1] == ("token", total - 1, f"id{total - 1}") + + +def test_reconnect_replay_still_delivers_full_offline_tail(): + """The cap must NOT break the #4633 reconnect-replay contract: a tab that + resubscribes after the offline buffer filled still receives the FULL + retained offline tail. This holds because the cap equals the offline buffer + cap — guarding against the regression where the cap was set too small.""" + ch = create_stream_channel() + n = StreamChannel._OFFLINE_BUFFER_MAXLEN + for i in range(n + 100): + ch.put_nowait(("token", i, f"id{i}")) + + q, snapshot = ch.subscribe_with_snapshot() + assert snapshot["offline_buffered_events"] == n + drained = _drain(q) + # Full retained tail delivered — the cap matches the offline buffer cap. + assert len(drained) == n + assert drained[0] == ("token", 100, "id100") # oldest 100 dropped from buffer + assert drained[-1] == ("token", n + 100 - 1, f"id{n + 100 - 1}") + + +def test_unsubscribe_removes_subscriber(): + """Sanity: after unsubscribe, the queue no longer receives broadcasts.""" + ch = create_stream_channel() + q = ch.subscribe() + other = ch.subscribe() # keep a live audience for the broadcast path + ch.unsubscribe(q) + ch.put_nowait(("token", 0, "id0")) + assert q.empty() + assert _drain(other) == [("token", 0, "id0")] + + +def test_terminal_frame_survives_concurrent_drain_during_drop_oldest(): + """Regression: on queue.Full the drop-oldest loop used to `break` if a + concurrent consumer drained the queue between the Full and get_nowait + (raising queue.Empty). That silently discarded the `item` being enqueued — + and when `item` was a TERMINAL frame (stream_end/error/cancel) the + subscriber never received it, leaving the client attached indefinitely + (spinner-forever). The fix is `continue` (retry the put into the now-empty + queue), so the terminal frame is always delivered. + + This deterministically simulates the exact interleaving (no threads, which + would be flaky): a custom queue whose get_nowait raises Empty on the first + post-Full call (as if a concurrent consumer had drained it), then behaves + normally. The OLD `break` logic drops the terminal item; the NEW `continue` + logic delivers it. We run both and assert only `continue` delivers.""" + terminal_item = ("stream_end", None, "terminal-id") + + class _RaceyQueue: + """Simulates a bounded queue that a concurrent consumer drained between + the producer's Full and get_nowait: the first put raises Full (queue at + cap), then by the time the producer calls get_nowait the consumer has + drained the ENTIRE queue, so get_nowait raises Empty. Subsequent puts + succeed (the queue has space) — which is exactly the condition where the + fix's `continue` must retry instead of the bug's `break`.""" + def __init__(self): + self._real = _queue.Queue(maxsize=8) + self._drained = False + # Pre-fill to cap so the first put_nowait raises Full. + for i in range(8): + self._real.put_nowait(("token", i, f"id{i}")) + + def put_nowait(self, item): + return self._real.put_nowait(item) + + def get_nowait(self): + # First eviction attempt: simulate the concurrent consumer having + # drained the whole queue — raise Empty (no oldest to drop). + if not self._drained: + self._drained = True + # Empty the real queue to model the drain, then raise Empty so + # the producer sees the race. + while True: + try: + self._real.get_nowait() + except _queue.Empty: + break + raise _queue.Empty + return self._real.get_nowait() + + def run_drop_oldest_loop(on_empty): + """Mirror the StreamChannel broadcast drop-oldest loop with a chosen + on-Empty action (break = pre-fix bug, continue = fix).""" + rq = _RaceyQueue() + delivered = [] + while True: + try: + rq.put_nowait(terminal_item) + delivered.append(terminal_item) + break + except _queue.Full: + try: + rq.get_nowait() + except _queue.Empty: + if on_empty == "break": + break # the bug: silently discard terminal_item + else: + continue # the fix: retry the put + return delivered + + # Pre-fix (break): the terminal frame is silently dropped. + delivered_break = run_drop_oldest_loop("break") + assert delivered_break == [], ( + "pre-fix `break` should have silently dropped the terminal frame" + ) + # Post-fix (continue): the terminal frame is delivered. + delivered_continue = run_drop_oldest_loop("continue") + assert delivered_continue == [terminal_item], ( + "post-fix `continue` must deliver the terminal frame after the race" + ) + diff --git a/tests/test_webui_runtime_diagnostics.py b/tests/test_webui_runtime_diagnostics.py index fe25db57..fa5f1db9 100644 --- a/tests/test_webui_runtime_diagnostics.py +++ b/tests/test_webui_runtime_diagnostics.py @@ -10,6 +10,7 @@ def test_stream_channel_exposes_buffer_and_subscriber_counts(): "subscriber_count": 0, "offline_buffered_events": 1, "offline_dropped_events": 0, + "subscriber_dropped_events": 0, } subscriber = channel.subscribe()