From 28ce32edc7ad0e714134055215aaee2c15a1f636 Mon Sep 17 00:00:00 2001 From: Nick Hollon Date: Thu, 16 Apr 2026 14:32:45 -0400 Subject: [PATCH] Bound EventLog / StreamChannel memory with drop-oldest semantics MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - EventLog(maxlen=N) caps retention. When the buffer is full, push evicts the oldest item and advances an absolute _first_seq so cursors can detect they've fallen off the back. A lagging cursor raises BufferOverflowError on its next read — mirrors the restored=false signal from the reconnection scenario (§06). - New cursors start at the current head of the buffer, not seq 0. For unbounded logs this is indistinguishable from the old behavior; for bounded logs, new consumers see whatever is still retained. - StreamChannel(name, *, maxlen=N) forwards maxlen to its inner log. - StreamMux(..., max_events=N) sets a default maxlen for every log / channel it binds (main event log plus each transformer projection). Explicit per-log maxlen wins over the mux default. - StreamingHandler.stream() / astream() expose max_events: caller sets the run-wide memory budget; transformer authors can override per-log when they know better. Default unbounded, matching §15 Q3. --- libs/langgraph/langgraph/stream/__init__.py | 3 +- libs/langgraph/langgraph/stream/_event_log.py | 75 +++++++-- libs/langgraph/langgraph/stream/_mux.py | 17 ++- .../langgraph/stream/stream_channel.py | 4 +- .../langgraph/stream/streaming_handler.py | 14 ++ .../langgraph/tests/test_streaming_handler.py | 143 ++++++++++++++++++ 6 files changed, 239 insertions(+), 17 deletions(-) diff --git a/libs/langgraph/langgraph/stream/__init__.py b/libs/langgraph/langgraph/stream/__init__.py index 386deeb1a..9f0c161c5 100644 --- a/libs/langgraph/langgraph/stream/__init__.py +++ b/libs/langgraph/langgraph/stream/__init__.py @@ -4,7 +4,7 @@ Provides a ``StreamingHandler`` that wraps a compiled graph and exposes ergonomic streaming projections through a transformer pipeline. """ -from langgraph.stream._event_log import EventLog +from langgraph.stream._event_log import BufferOverflowError, EventLog from langgraph.stream._types import ProtocolEvent, StreamTransformer from langgraph.stream.run_stream import AsyncGraphRunStream, GraphRunStream from langgraph.stream.stream_channel import StreamChannel @@ -12,6 +12,7 @@ from langgraph.stream.streaming_handler import StreamingHandler __all__ = [ "AsyncGraphRunStream", + "BufferOverflowError", "EventLog", "GraphRunStream", "ProtocolEvent", diff --git a/libs/langgraph/langgraph/stream/_event_log.py b/libs/langgraph/langgraph/stream/_event_log.py index 5776d9f7d..6008d53af 100644 --- a/libs/langgraph/langgraph/stream/_event_log.py +++ b/libs/langgraph/langgraph/stream/_event_log.py @@ -1,12 +1,23 @@ from __future__ import annotations import asyncio +from collections import deque from collections.abc import AsyncIterator, Callable, Iterator from typing import Generic, TypeVar T = TypeVar("T") +class BufferOverflowError(RuntimeError): + """Raised when an EventLog cursor falls off the back of a bounded buffer. + + Mirrors the ``restored: false`` signal from the protocol's reconnection + story (§ 06): consumers that fall behind the retention window get an + explicit error and can decide to rebuild from a snapshot rather than + silently losing events. + """ + + class EventLog(Generic[T]): """Append-only buffer that supports multiple independent consumers. @@ -28,10 +39,26 @@ class EventLog(Generic[T]): Async iteration uses a shared ``asyncio.Event`` — cursors await the event when they catch up, and the producer sets it on each push. + + Bounded mode + ------------ + Pass ``maxlen=N`` to cap memory. When the buffer is full, ``push`` + drops the oldest item to make room. Cursors track an absolute + sequence number; a cursor that falls off the back of the retention + window raises ``BufferOverflowError`` on its next read. + + New cursors start at the current head of the buffer, not at seq 0 + — they see whatever is still retained. This matches the protocol's + reconnection semantics (§ 06: "missed events can be replayed from + a bounded buffer"). """ - def __init__(self) -> None: - self._items: list[T] = [] + def __init__(self, maxlen: int | None = None) -> None: + if maxlen is not None and maxlen <= 0: + raise ValueError("EventLog maxlen must be a positive int or None") + self._items: deque[T] = deque() + self._maxlen: int | None = maxlen + self._first_seq = 0 # absolute seq of _items[0] self._closed = False self._error: BaseException | None = None @@ -65,9 +92,17 @@ class EventLog(Generic[T]): # ------------------------------------------------------------------ def push(self, item: T) -> None: - """Append *item* and wake all waiting cursors.""" + """Append *item* and wake all waiting cursors. + + In bounded mode, evicts the oldest item first if the buffer + is full, advancing ``_first_seq`` so cursors can detect that + they've fallen off the back of the retention window. + """ if self._closed: raise RuntimeError("Cannot push to a closed EventLog") + if self._maxlen is not None and len(self._items) >= self._maxlen: + self._items.popleft() + self._first_seq += 1 self._items.append(item) self._notify() @@ -112,11 +147,19 @@ class EventLog(Generic[T]): return self._sync_cursor() def _sync_cursor(self) -> Iterator[T]: - cursor = 0 + # Start at the current head — if maxlen is None this is 0 (seen everything), + # if bounded this is wherever retention currently begins. + seq = self._first_seq while True: - if cursor < len(self._items): - item = self._items[cursor] - cursor += 1 + if seq < self._first_seq: + raise BufferOverflowError( + f"Cursor fell {self._first_seq - seq} items behind the " + f"bounded EventLog's retention window (maxlen={self._maxlen})" + ) + idx = seq - self._first_seq + if idx < len(self._items): + item = self._items[idx] + seq += 1 yield item elif self._closed: if self._error is not None: @@ -125,7 +168,7 @@ class EventLog(Generic[T]): elif self._request_more is not None: # Pull from the producer until this log gets a new item # or the graph is exhausted (which closes the log). - while cursor >= len(self._items) and not self._closed: + while (seq - self._first_seq) >= len(self._items) and not self._closed: if not self._request_more(): break else: @@ -149,16 +192,22 @@ class EventLog(Generic[T]): async def _async_cursor(self) -> AsyncIterator[T]: assert self._event is not None - cursor = 0 + seq = self._first_seq while True: - if cursor < len(self._items): - yield self._items[cursor] - cursor += 1 + if seq < self._first_seq: + raise BufferOverflowError( + f"Cursor fell {self._first_seq - seq} items behind the " + f"bounded EventLog's retention window (maxlen={self._maxlen})" + ) + idx = seq - self._first_seq + if idx < len(self._items): + yield self._items[idx] + seq += 1 elif self._closed: if self._error is not None: raise self._error return else: self._event.clear() - if cursor >= len(self._items) and not self._closed: + if (seq - self._first_seq) >= len(self._items) and not self._closed: await self._event.wait() diff --git a/libs/langgraph/langgraph/stream/_mux.py b/libs/langgraph/langgraph/stream/_mux.py index 8ef8b50e8..8677e863c 100644 --- a/libs/langgraph/langgraph/stream/_mux.py +++ b/libs/langgraph/langgraph/stream/_mux.py @@ -33,6 +33,7 @@ class StreamMux: transformers: list[StreamTransformer] | None = None, *, is_async: bool = False, + max_events: int | None = None, ) -> None: """Initialize the mux and register *transformers* in order. @@ -42,11 +43,18 @@ class StreamMux: keys are recorded in ``self.native_keys``, and any ``EventLog`` / ``StreamChannel`` instances are bound/wired. + *max_events* sets a default capacity for every ``EventLog`` / + ``StreamChannel`` the mux binds, including the main event log. + Logs that were constructed with an explicit ``maxlen`` keep + their own setting — the mux only fills in ``None`` defaults. + Unbounded when ``max_events`` is ``None``. + Raises ``RuntimeError`` if any transformer requires an async run under sync mode, and ``ValueError`` on projection-key conflicts. """ self._is_async = is_async - self._events: EventLog[ProtocolEvent] = EventLog() + self._default_maxlen = max_events + self._events: EventLog[ProtocolEvent] = EventLog(maxlen=max_events) self._events._bind(is_async=is_async) self._transformers: list[StreamTransformer] = [] self._channels: list[StreamChannel[Any]] = [] @@ -275,6 +283,7 @@ class StreamMux: """Bind and wire EventLog / StreamChannel instances in *projection*.""" for value in projection.values(): if isinstance(value, StreamChannel): + self._apply_default_maxlen(value._log) value._bind(is_async=self._is_async) self._channels.append(value) channel_name = value.name @@ -287,9 +296,15 @@ class StreamMux: value._wire(_make_forward(channel_name)) elif isinstance(value, EventLog): + self._apply_default_maxlen(value) value._bind(is_async=self._is_async) self._logs.append(value) + def _apply_default_maxlen(self, log: EventLog[Any]) -> None: + """Fill in the mux's default maxlen if the log hasn't set its own.""" + if log._maxlen is None and self._default_maxlen is not None: + log._maxlen = self._default_maxlen + def _forward(self, channel_name: str, item: Any) -> None: """Inject a ProtocolEvent for a StreamChannel push. diff --git a/libs/langgraph/langgraph/stream/stream_channel.py b/libs/langgraph/langgraph/stream/stream_channel.py index d1800da60..c5b079e29 100644 --- a/libs/langgraph/langgraph/stream/stream_channel.py +++ b/libs/langgraph/langgraph/stream/stream_channel.py @@ -29,9 +29,9 @@ class StreamChannel(Generic[T]): using only StreamChannels don't need ``finalize`` / ``fail`` hooks. """ - def __init__(self, name: str) -> None: + def __init__(self, name: str, *, maxlen: int | None = None) -> None: self.name = name - self._log: EventLog[T] = EventLog() + self._log: EventLog[T] = EventLog(maxlen=maxlen) self._wire_fn: Callable[[T], None] | None = None def _bind(self, *, is_async: bool) -> None: diff --git a/libs/langgraph/langgraph/stream/streaming_handler.py b/libs/langgraph/langgraph/stream/streaming_handler.py index 4c26b0c41..db88ff26b 100644 --- a/libs/langgraph/langgraph/stream/streaming_handler.py +++ b/libs/langgraph/langgraph/stream/streaming_handler.py @@ -56,6 +56,7 @@ class StreamingHandler: interrupt_before: All | Sequence[str] | None = None, interrupt_after: All | Sequence[str] | None = None, transformers: list[StreamTransformer] | None = None, + max_events: int | None = None, ) -> GraphRunStream: """Start a sync streaming run. @@ -63,11 +64,19 @@ class StreamingHandler: any projection drives the graph forward — no background thread is used. This matches v1's model where the caller's ``for`` loop is the pump. + + *max_events* caps the retention of every ``EventLog`` / + ``StreamChannel`` the mux binds (main event log plus each + transformer's projection logs) to the given number of items, + dropping the oldest when full. Transformers that constructed + their own logs with an explicit ``maxlen`` keep their setting. + Unbounded when ``None``. """ values_t = ValuesTransformer() mux = StreamMux( [values_t, MessagesTransformer(), *(transformers or ())], is_async=False, + max_events=max_events, ) graph_iter = iter( @@ -92,16 +101,21 @@ class StreamingHandler: interrupt_before: All | Sequence[str] | None = None, interrupt_after: All | Sequence[str] | None = None, transformers: list[StreamTransformer] | None = None, + max_events: int | None = None, ) -> AsyncGraphRunStream: """Start an async streaming run. Returns an `AsyncGraphRunStream` immediately. A background asyncio task pumps events from the graph into the transformer pipeline. + + *max_events* caps retention of every ``EventLog`` / ``StreamChannel`` + the mux binds — see ``stream()`` for the full semantics. """ values_t = ValuesTransformer() mux = StreamMux( [values_t, MessagesTransformer(), *(transformers or ())], is_async=True, + max_events=max_events, ) async def pump() -> None: diff --git a/libs/langgraph/tests/test_streaming_handler.py b/libs/langgraph/tests/test_streaming_handler.py index d9beadc6b..58020ea0f 100644 --- a/libs/langgraph/tests/test_streaming_handler.py +++ b/libs/langgraph/tests/test_streaming_handler.py @@ -15,6 +15,7 @@ from typing_extensions import TypedDict from langgraph.constants import END, START from langgraph.graph import StateGraph from langgraph.stream import ( + BufferOverflowError, EventLog, StreamChannel, StreamingHandler, @@ -1444,3 +1445,145 @@ class TestAsyncTransformerLane: _ = await run.output scores = [x async for x in run.extensions["scores"]] assert scores and all(s == 42 for s in scores) + + +# --------------------------------------------------------------------------- +# Bounded EventLog / StreamChannel — memory caps and overflow semantics +# --------------------------------------------------------------------------- + + +class TestBoundedEventLog: + def test_invalid_maxlen_raises(self) -> None: + with pytest.raises(ValueError, match="positive int or None"): + EventLog(maxlen=0) + with pytest.raises(ValueError, match="positive int or None"): + EventLog(maxlen=-3) + + def test_unbounded_default_preserves_replay(self) -> None: + """Default EventLog (maxlen=None) still replays from seq 0.""" + log: EventLog[int] = EventLog() + log._bind(is_async=False) + for i in range(100): + log.push(i) + log.close() + assert list(log) == list(range(100)) + + def test_bounded_drops_oldest_on_overflow(self) -> None: + """When bounded, pushing past maxlen evicts the oldest item.""" + log: EventLog[int] = EventLog(maxlen=3) + log._bind(is_async=False) + for i in range(5): + log.push(i) + log.close() + # Only the last 3 survive; new cursors start at the current head. + assert list(log) == [2, 3, 4] + + def test_new_cursor_starts_at_head_not_zero(self) -> None: + """New cursors see the retained window, not the evicted prefix.""" + log: EventLog[int] = EventLog(maxlen=2) + log._bind(is_async=False) + log.push(1) + log.push(2) + log.push(3) # evicts 1 + log.close() + assert list(log) == [2, 3] + + @pytest.mark.anyio + async def test_async_cursor_overflow_raises(self) -> None: + """An async cursor that falls behind the retention window raises.""" + log: EventLog[int] = EventLog(maxlen=2) + log._bind(is_async=True) + + log.push(1) + cursor = aiter(log) + # Advance cursor to seq 1, reading item 1. + first = await anext(cursor) + assert first == 1 + # Now push enough to roll the cursor off the back. + log.push(2) # buffer: [1, 2] _first_seq=0, cursor at seq=1 + log.push(3) # buffer: [2, 3] _first_seq=1, cursor at seq=1 still OK + log.push(4) # buffer: [3, 4] _first_seq=2, cursor at seq=1 — overflow + with pytest.raises(BufferOverflowError, match="fell"): + await anext(cursor) + + def test_sync_cursor_sees_all_while_bounded_but_under_cap(self) -> None: + """Bounded mode with pushes under cap behaves identically to unbounded.""" + log: EventLog[int] = EventLog(maxlen=100) + log._bind(is_async=False) + log.push(1) + log.push(2) + log.close() + assert list(log) == [1, 2] + + +class TestStreamChannelMaxlen: + def test_maxlen_passes_through_to_inner_log(self) -> None: + ch: StreamChannel[int] = StreamChannel("ch", maxlen=2) + ch._bind(is_async=False) + ch.push(1) + ch.push(2) + ch.push(3) + ch._close() + assert list(ch) == [2, 3] + + +class TestMuxMaxEventsDefault: + def test_mux_fills_in_default_when_log_has_none(self) -> None: + class Simple(StreamTransformer): + def __init__(self) -> None: + self.log: EventLog[int] = EventLog() + + def init(self) -> dict[str, Any]: + return {"out": self.log} + + def process(self, event: ProtocolEvent) -> bool: + return True + + t = Simple() + StreamMux([t], max_events=10) + assert t.log._maxlen == 10 + + def test_explicit_log_maxlen_wins_over_mux_default(self) -> None: + class Explicit(StreamTransformer): + def __init__(self) -> None: + self.log: EventLog[int] = EventLog(maxlen=3) + + def init(self) -> dict[str, Any]: + return {"out": self.log} + + def process(self, event: ProtocolEvent) -> bool: + return True + + t = Explicit() + StreamMux([t], max_events=1000) + assert t.log._maxlen == 3 # transformer author's setting stands + + def test_main_event_log_respects_max_events(self) -> None: + mux = StreamMux([], max_events=5) + assert mux._events._maxlen == 5 + + def test_max_events_default_cascades_to_channels(self) -> None: + class WithChannel(StreamTransformer): + def __init__(self) -> None: + self.ch: StreamChannel[int] = StreamChannel("out") + + def init(self) -> dict[str, Any]: + return {"out": self.ch} + + def process(self, event: ProtocolEvent) -> bool: + return True + + t = WithChannel() + StreamMux([t], max_events=7) + assert t.ch._log._maxlen == 7 + + def test_handler_propagates_max_events(self) -> None: + graph = _build_simple_graph() + handler = StreamingHandler(graph) + run = handler.stream({"value": "x", "items": []}, max_events=50) + # Main log inherits the default. + assert run._mux._events._maxlen == 50 + # Native projections (values log, messages log) inherit too. + for log in run._mux._logs: + assert log._maxlen == 50 + _ = run.output