fix(stream): cap StreamChannel per-subscriber queues (drop-oldest on full)
Memory: each connected tab's subscriber queue was queue.Queue() with no maxsize. A slow/backpressured/backgrounded tab accumulated every coalesced token frame for the WHOLE turn (producer = agent token stream) — unbounded per-tab growth, OOM risk with many tabs x long agentic turns. The offline buffer was already capped (#4633), but once a subscriber attached the live-broadcast path flowed into an unbounded per-subscriber queue. Cap per-subscriber queues at _SUBSCRIBER_QUEUE_MAXSIZE (== _OFFLINE_BUFFER_MAXLEN) and drop the OLDEST frame on queue.Full in both the broadcast (put_nowait) and offline-buffer replay (subscribe_with_snapshot) paths, so a slow tab keeps the most recent tail (older frames stay recoverable via the run journal by last_event_id). Mirrors SessionChannel.emit's drop-on-full contract. The cap intentionally equals the offline buffer bound, not SessionChannel's smaller 16: StreamChannel carries the live token stream and has a #4633 reconnect-replay contract (a resubscribing tab must receive the full retained offline tail), so capping below the buffer would force every flaky-network reconnect through the run journal. Matching the buffer bound preserves that contract while bounding live-broadcast memory to the same worst-case #4633 already accepts. The SSE write deadline (SSE_WRITE_DEADLINE_SECONDS) breaks a stuck socket within ~20s, so the overflow window is short; this also bounds it by frame count. Adds a cumulative subscriber_dropped_events diagnostic counter (ops visibility). Test: tests/test_stream_subscriber_queue_cap.py — bounded queue never exceeds cap, drop-oldest keeps newest tail, draining vs non-draining subscribers, and the #4633 reconnect-replay contract still delivers the full offline tail. Existing offline-buffer/gap-recovery/multitab tests unchanged. Contract Routing - Task type: memory-bound fix (unbounded per-subscriber queue). - Touched areas: live chat SSE broadcast (StreamChannel). - Relevant public docs: AGENTS.md, CONTRIBUTING.md, docs/CONTRACTS.md, docs/rfcs/webui-run-state-consistency-contract.md (Live stream/SSE layer), docs/rfcs/session-sse-contract-v1.md (distinct: that RFC governs the PROPOSED per-session endpoint /api/sessions/{id}/events, NOT StreamChannel, which is the /api/chat/stream live token broadcast; no text in that RFC is changed). - State layer mutated: Live stream / SSE layer — specifically the per-subscriber queue inside StreamChannel (api/config.py). NOT the durable run journal, NOT the transcript, NOT replay. - Invariant preserved: the #4633 reconnect-replay contract — a resubscribing tab must receive the full retained offline tail. The cap is set EQUAL to _OFFLINE_BUFFER_MAXLEN so replay delivery is unchanged; only the previously- unbounded live-broadcast-to-slow-tab growth is bounded. Regression: test_reconnect_replay_still_delivers_full_offline_tail (and the pre-existing test_stream_offline_buffer_cap.py suite) pass unchanged. - Scope boundaries: does not touch replay idempotency (Invariant 5), the turn journal, compression, or sidebar metadata. Slow-tab frames dropped at the cap stay recoverable via the run journal by last_event_id. Release note: StreamChannel per-subscriber SSE queues are now bounded (drop-oldest), capping memory for slow/backgrounded tabs during long turns. Model Used: builtin:zai-coding-plan/GLM-5.2 (ZCode agent). No AI-specific tool use mattered beyond the repo's standard read/edit/test workflow.
This commit is contained in:
1 parent
6ed474ca54
commit
bbda885aa6
2 files changed
+252
-8
No files matched your search
@@ -0,0 +1,165 @@
|
||||
"""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")]
|
||||
Reference in new issue
Block a user