From ab1d6980b54abdb6df07232e3b86fbb02c14d96d Mon Sep 17 00:00:00 2001 From: Nick Hollon Date: Sat, 18 Apr 2026 12:34:10 -0400 Subject: [PATCH] Drain-on-consume streaming with caller-driven async pump MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Collapse the eager async pump task into the same caller-driven model as sync: each cursor's advance drives one graph event through the mux. Concurrent async consumers serialize through an asyncio.Lock so each acquisition produces exactly one event, matching sync semantics. EventLog becomes a single-consumer drainable queue — items pop off as the cursor advances, a second __iter__ / __aiter__ raises. Fan-out moves to explicit tee(n) / atee(n) helpers. Retention windows, BufferOverflowError, and max_events are gone; pre- subscription pushes are silent no-ops so unsubscribed projections don't accumulate. Both run streams gain abort() and context-manager support; the pump's BaseException catch is narrowed to Exception so CancelledError propagates per asyncio contract. TestMemoryBounds locks in the drain-on-consume invariants: subscribed buffers drop back to empty after each yield, unsubscribed projections never accumulate, and run.output leaves the values log untouched. --- libs/langgraph/langgraph/stream/__init__.py | 3 +- libs/langgraph/langgraph/stream/_event_log.py | 285 ++++++----- libs/langgraph/langgraph/stream/_mux.py | 20 +- libs/langgraph/langgraph/stream/run_stream.py | 278 ++++++++--- .../langgraph/stream/stream_channel.py | 14 + .../langgraph/stream/streaming_handler.py | 52 +- .../langgraph/stream/transformers.py | 15 +- .../langgraph/tests/test_streaming_handler.py | 462 +++++++++++------- 8 files changed, 710 insertions(+), 419 deletions(-) diff --git a/libs/langgraph/langgraph/stream/__init__.py b/libs/langgraph/langgraph/stream/__init__.py index 294e81f28..5a4633c71 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 BufferOverflowError, EventLog +from langgraph.stream._event_log import EventLog from langgraph.stream._types import ProtocolEvent, StreamTransformer from langgraph.stream.run_stream import AsyncGraphRunStream, GraphRunStream from langgraph.stream.stream_channel import StreamChannel @@ -12,7 +12,6 @@ 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 f87c3b75f..05bc10504 100644 --- a/libs/langgraph/langgraph/stream/_event_log.py +++ b/libs/langgraph/langgraph/stream/_event_log.py @@ -2,60 +2,50 @@ from __future__ import annotations import asyncio from collections import deque -from collections.abc import AsyncIterator, Callable, Iterator +from collections.abc import AsyncIterator, Awaitable, 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. + """Single-consumer drainable queue for streaming events. + + Items are popped off the front as the consumer advances — there is + no retention beyond what's currently queued. A log accepts exactly + one subscriber; a second `__iter__` / `__aiter__` call raises. Use + `tee(n)` / `atee(n)` for fan-out. Starts unbound — neither `__iter__` nor `__aiter__` is available until the StreamMux calls `_bind(is_async)`. After binding, only the matching iteration protocol works; the other raises `TypeError`. - All access is single-threaded: sync mode is caller-driven (no - background thread), async mode runs entirely on the event loop. + Pump wiring (set by the run stream, not by `_bind`): + - `_request_more`: sync pump callable, returns True if a new + event was produced. + - `_arequest_more`: async pump coroutine factory, same contract. - Producer API: - - `push(item)`: append an item, notify all waiting cursors. - - `close()`: mark the log as done. - - `fail(err)`: mark the log as errored. + Memory is bounded by caller pace: both sync and async use caller- + driven pumps, so each cursor advance produces at most one event. + The only shape where a log can accumulate meaningfully is + concurrent async consumers at unequal rates — a slow consumer's + log grows while fast consumers drive the shared pump. That's the + documented tradeoff for concurrent consumption; consume at similar + rates or use a single consumer if memory matters. - Sync iteration is pull-based: when a cursor catches up it calls - `_request_more` to drive the graph forward. Async iteration uses a - shared `asyncio.Event` — cursors await the event when they catch - up, and the producer sets it on each push. - - 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"). + Lazy-subscribe: `push` is a no-op when no subscriber has registered. + Transformers still execute `process()` (so scalar state like + `ValuesTransformer._latest` stays current); only the log append is + skipped. """ def __init__(self, maxlen: int | None = None) -> None: """Initialize an empty, unbound log. Args: - maxlen: Optional cap on retained items. When reached, the - oldest item is dropped on each new push. `None` (the - default) leaves the log unbounded. + maxlen: Accepted for forward compatibility; currently unused. + The caller-driven pump bounds memory naturally for + single-consumer use. Raises: ValueError: If `maxlen` is not a positive integer or `None`. @@ -64,18 +54,19 @@ class EventLog(Generic[T]): 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 # Binding state — None means unbound. self._is_async: bool | None = None - # Sync pull callback (set by the run stream, not by bind). - self._request_more: Callable[[], bool] | None = None + # Flipped on first __iter__ / __aiter__. Pre-subscription + # pushes are silent no-ops. + self._subscribed = False - # Async notification (allocated on bind). - self._event: asyncio.Event | None = None + # Pump wiring set by the run stream after bind. + self._request_more: Callable[[], bool] | None = None + self._arequest_more: Callable[[], Awaitable[bool]] | None = None # ------------------------------------------------------------------ # Binding @@ -96,75 +87,50 @@ class EventLog(Generic[T]): if self._is_async is not None: raise RuntimeError("EventLog is already bound") self._is_async = is_async - if is_async: - self._event = asyncio.Event() # ------------------------------------------------------------------ # Producer API # ------------------------------------------------------------------ def push(self, item: T) -> None: - """Append an item and wake all waiting cursors. + """Append an item. No-op when no subscriber is registered. - 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. - - Args: - item: The item to append. + Non-blocking in both sync and async — matches v1's + `put_nowait` producer shape. Memory is bounded by caller pace + via the caller-driven pump. Raises: - RuntimeError: If the log is closed. + RuntimeError: If the log is closed (and subscribed). """ + if not self._subscribed: + return 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() def close(self) -> None: - """Mark the log as complete. - - Open cursors will finish cleanly once they drain the buffer. - """ + """Mark the log as complete.""" self._closed = True - self._notify() def fail(self, err: BaseException) -> None: """Mark the log as errored. - Open cursors will raise `err` once they drain the buffer. - Args: - err: The exception to surface to consumers. + err: The exception to surface to the subscriber. """ self._error = err self._closed = True - self._notify() # ------------------------------------------------------------------ - # Notification - # ------------------------------------------------------------------ - - def _notify(self) -> None: - """Wake async cursors waiting for data. - - No-op when sync-bound (no event exists). - """ - if self._event is not None: - self._event.set() - - # ------------------------------------------------------------------ - # Sync iteration (pull-based) + # Sync iteration (caller-driven pump) # ------------------------------------------------------------------ def __iter__(self) -> Iterator[T]: - """Return a new independent sync cursor over the log. + """Subscribe and return a sync cursor. Can be called only once. Raises: TypeError: If the log is unbound or bound to async mode. + RuntimeError: If the log already has a subscriber. """ if self._is_async is None: raise TypeError( @@ -175,49 +141,38 @@ class EventLog(Generic[T]): raise TypeError( "This EventLog is bound to async mode — use 'async for' instead." ) + if self._subscribed: + raise RuntimeError( + "EventLog already has a subscriber; use .tee(n) for fan-out." + ) + self._subscribed = True return self._sync_cursor() def _sync_cursor(self) -> Iterator[T]: - # 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 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 + if self._items: + yield self._items.popleft() elif self._closed: if self._error is not None: raise self._error return 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). A push - # may evict the item this cursor was about to read; in that - # case the inner loop breaks and the outer `seq < _first_seq` - # check catches the overflow on the next iteration. - while (seq - self._first_seq) >= len(self._items) and not self._closed: - if not self._request_more(): - break + if not self._request_more(): + if not self._items and not self._closed: + return else: - # No producer callback and not closed — buffer is complete. return # ------------------------------------------------------------------ - # Async iteration + # Async iteration (caller-driven pump) # ------------------------------------------------------------------ def __aiter__(self) -> AsyncIterator[T]: - """Return a new independent async cursor over the log. + """Subscribe and return an async cursor. Can be called only once. Raises: TypeError: If the log is unbound or bound to sync mode. + RuntimeError: If the log already has a subscriber. """ if self._is_async is None: raise TypeError( @@ -226,26 +181,126 @@ class EventLog(Generic[T]): ) if not self._is_async: raise TypeError("This EventLog is bound to sync mode — use 'for' instead.") + if self._subscribed: + raise RuntimeError( + "EventLog already has a subscriber; use .atee(n) for fan-out." + ) + self._subscribed = True return self._async_cursor() async def _async_cursor(self) -> AsyncIterator[T]: - assert self._event is not None - seq = self._first_seq while True: - 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 + if self._items: + yield self._items.popleft() elif self._closed: if self._error is not None: raise self._error return + elif self._arequest_more is not None: + if not await self._arequest_more(): + if not self._items and not self._closed: + return else: - self._event.clear() - if (seq - self._first_seq) >= len(self._items) and not self._closed: - await self._event.wait() + return + + # ------------------------------------------------------------------ + # Fan-out via tee + # ------------------------------------------------------------------ + + def tee(self, n: int = 2) -> tuple[Iterator[T], ...]: + """Subscribe and return `n` independent sync iterators. + + Each branch has its own buffer; items pulled from the + underlying cursor are copied into every branch. Branches are + naturally bounded by caller pace since the sync pump is + caller-driven. + + Args: + n: Number of branches to create. Must be >= 1. + + Returns: + A tuple of `n` iterators over the same underlying stream. + + Raises: + TypeError: If the log is unbound or bound to async mode. + RuntimeError: If the log already has a subscriber. + ValueError: If `n` < 1. + """ + if n < 1: + raise ValueError("tee() requires n >= 1") + source = self.__iter__() + buffers: list[deque[T]] = [deque() for _ in range(n)] + exhausted = [False] + + def branch(i: int) -> Iterator[T]: + buf = buffers[i] + while True: + if buf: + yield buf.popleft() + elif exhausted[0]: + return + else: + try: + item = next(source) + except StopIteration: + exhausted[0] = True + return + for b in buffers: + b.append(item) + + return tuple(branch(i) for i in range(n)) + + def atee(self, n: int = 2) -> tuple[AsyncIterator[T], ...]: + """Subscribe and return `n` independent async iterators. + + Caller-driven fan-out: each branch's `__anext__` either pops + from its own buffer or, under a shared `asyncio.Lock`, pulls + one item from the underlying cursor and distributes it to + every branch's buffer. + + Args: + n: Number of branches to create. Must be >= 1. + + Returns: + A tuple of `n` async iterators over the same underlying + stream. + + Raises: + TypeError: If the log is unbound or bound to sync mode. + RuntimeError: If the log already has a subscriber. + ValueError: If `n` < 1. + """ + if n < 1: + raise ValueError("atee() requires n >= 1") + source = self.__aiter__() + buffers: list[deque[T]] = [deque() for _ in range(n)] + exhausted = [False] + error: list[BaseException | None] = [None] + lock = asyncio.Lock() + + async def branch(i: int) -> AsyncIterator[T]: + buf = buffers[i] + while True: + if buf: + yield buf.popleft() + continue + if exhausted[0]: + if error[0] is not None: + raise error[0] + return + async with lock: + if buf or exhausted[0]: + continue + try: + item = await source.__anext__() + except StopAsyncIteration: + exhausted[0] = True + continue + except Exception as e: + error[0] = e + exhausted[0] = True + continue + for b in buffers: + b.append(item) + + return tuple(branch(i) for i in range(n)) diff --git a/libs/langgraph/langgraph/stream/_mux.py b/libs/langgraph/langgraph/stream/_mux.py index f3fab3965..e5e911e6d 100644 --- a/libs/langgraph/langgraph/stream/_mux.py +++ b/libs/langgraph/langgraph/stream/_mux.py @@ -40,7 +40,6 @@ 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. @@ -55,11 +54,6 @@ class StreamMux: `None` or empty gives a mux with no projections. is_async: True for async dispatch (`apush` / `aclose` / `afail`), False for the sync path. - max_events: Default capacity for every EventLog and - StreamChannel the mux binds, including the main event - log. Logs constructed with an explicit `maxlen` keep - their own setting — the mux only fills in unset - defaults. `None` leaves the logs unbounded. Raises: RuntimeError: If any transformer requires an async run but @@ -68,8 +62,7 @@ class StreamMux: ValueError: If transformers' projection keys collide. """ self._is_async = is_async - self._default_maxlen = max_events - self._events: EventLog[ProtocolEvent] = EventLog(maxlen=max_events) + self._events: EventLog[ProtocolEvent] = EventLog() self._events._bind(is_async=is_async) self._transformers: list[StreamTransformer] = [] self._channels: list[StreamChannel[Any]] = [] @@ -218,6 +211,10 @@ class StreamMux: async work. For decoupled work, use `schedule()` from inside `process` / `aprocess` instead. + The main log append is a non-blocking `push` — matching v1's + `put_nowait` shape. Memory is bounded by caller pace via the + caller-driven pump; see `EventLog` for the full tradeoff story. + Args: event: The protocol event to dispatch. """ @@ -324,7 +321,6 @@ class StreamMux: """Bind and wire EventLog / StreamChannel instances in a 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 @@ -337,15 +333,9 @@ 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 when 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/run_stream.py b/libs/langgraph/langgraph/stream/run_stream.py index dedfd9ab7..13bd6fb86 100644 --- a/libs/langgraph/langgraph/stream/run_stream.py +++ b/libs/langgraph/langgraph/stream/run_stream.py @@ -1,8 +1,8 @@ from __future__ import annotations import asyncio -from collections.abc import AsyncIterator, Iterator, Mapping -from types import MappingProxyType +from collections.abc import AsyncIterator, Awaitable, Callable, Iterator, Mapping +from types import MappingProxyType, TracebackType from typing import Any from langgraph.stream._convert import convert_to_protocol_event @@ -13,20 +13,31 @@ from langgraph.stream.stream_channel import StreamChannel from langgraph.stream.transformers import ValuesTransformer +def _drive_until_done(pump: Callable[[], bool]) -> None: + """Call the sync pump until it returns False.""" + while pump(): + pass + + +async def _adrive_until_done(pump: Callable[[], Awaitable[bool]]) -> None: + """Call the async pump until it returns False.""" + while await pump(): + pass + + class GraphRunStream: """Sync run stream with caller-driven pumping. The caller's iteration on any projection (`values`, `messages`, raw events, or `output`) drives the graph forward. No background - thread is used — this matches v1's model where the caller's `for` - loop is the pump. + thread is used — the caller's `for` loop is the pump. + + Projections are single-consumer — iterating `run.values` twice + raises. Use `projection.tee(n)` if you genuinely need fan-out. All transformer projections live in `extensions`. Native transformer projections (those with `_native = True`) are also set as direct attributes on this instance (e.g. `run.values`, `run.messages`). - - Iterating the run stream directly yields raw `ProtocolEvent` objects - from the mux's main event log. """ def __init__( @@ -48,19 +59,13 @@ class GraphRunStream: self.extensions: Mapping[str, Any] = MappingProxyType(mux.extensions) self._values_transformer = values_transformer self._exhausted = False - # Native-transformer projections also show up as direct attributes. for key in mux.native_keys: setattr(self, key, mux.extensions[key]) - # Wire pull-based iteration: every sync EventLog calls _pump_next - # when its cursor catches up to the buffer. self._wire_request_more(mux) def _wire_request_more(self, mux: StreamMux) -> None: - """Install `_request_more` on every sync EventLog. - - Sync iteration is caller-driven, so a cursor that catches up to - the buffer's tail needs a way to ask the graph for more events. - """ + """Install `_request_more` on every sync EventLog so cursors + can drive the pump when their buffer catches up.""" mux._events._request_more = self._pump_next for value in mux.extensions.values(): if isinstance(value, EventLog): @@ -83,22 +88,43 @@ class GraphRunStream: self._mux.close() self._exhausted = True return False - except BaseException as e: + except Exception as e: self._mux.fail(e) self._exhausted = True return False self._mux.push(convert_to_protocol_event(part)) return True - def _pump_all(self) -> None: - """Drain the graph completely.""" - while self._pump_next(): + def abort(self) -> None: + """Stop the run early. + + Closes the mux and marks the stream exhausted. The graph + iterator is dropped; any in-flight nodes see the closure on + their next yield point. Idempotent. + """ + if self._exhausted: + return + self._exhausted = True + try: + self._mux.close() + except Exception: pass + def __enter__(self) -> GraphRunStream: + return self + + def __exit__( + self, + exc_type: type[BaseException] | None, + exc: BaseException | None, + tb: TracebackType | None, + ) -> None: + self.abort() + @property def output(self) -> dict[str, Any] | None: - """Block until the run completes and return the final state.""" - self._pump_all() + """Drive the run to completion and return the final state.""" + _drive_until_done(self._pump_next) err = self._values_transformer.error if err is not None: raise err @@ -106,12 +132,13 @@ class GraphRunStream: @property def interrupted(self) -> bool: - """Block until the run completes, then return whether it was interrupted. + """Drive the run to completion, then return whether it was + interrupted. Raises: BaseException: If the run ended with an error. """ - self._pump_all() + _drive_until_done(self._pump_next) err = self._values_transformer.error if err is not None: raise err @@ -119,70 +146,193 @@ class GraphRunStream: @property def interrupts(self) -> list[Any]: - """Block until the run completes, then return interrupt payloads. + """Drive the run to completion, then return interrupt payloads. Raises: BaseException: If the run ended with an error. """ - self._pump_all() + _drive_until_done(self._pump_next) err = self._values_transformer.error if err is not None: raise err return self._values_transformer._interrupts def __iter__(self) -> Iterator[ProtocolEvent]: - """Iterate all protocol events from the mux's main event log.""" + """Subscribe to the main event log and iterate protocol events.""" return iter(self._mux._events) + def interleave(self, *names: str) -> Iterator[tuple[str, Any]]: + """Iterate multiple projections round-robin, yielding ``(name, item)``. + + Each turn advances one projection's cursor; when a cursor's buffer + is empty, pulling from it drives the pump once, which fans out to + every subscribed projection log. Projections whose items aren't + consumed on this turn sit in their own buffers only until the next + turn reaches them, bounding memory by the skew between projection + rates rather than letting any single log grow to the full run + length. + + Projections are exhausted independently; a projection that finishes + early drops out of the rotation while others continue. The overall + iterator ends once all named projections are done. + + Args: + *names: Projection keys to interleave. Must match keys in + ``extensions``. + + Yields: + ``(name, item)`` tuples in round-robin order across the named + projections. + + Raises: + KeyError: If a name doesn't match a registered projection. + + Example: + ```python + for name, item in run.interleave("messages", "values"): + if name == "messages": + print("msg:", item) + else: + print("val:", item) + ``` + """ + cursors: dict[str, Iterator[Any]] = { + name: iter(self.extensions[name]) for name in names + } + done: set[str] = set() + while len(done) < len(cursors): + for name, cursor in cursors.items(): + if name in done: + continue + try: + item = next(cursor) + except StopIteration: + done.add(name) + continue + yield (name, item) + class AsyncGraphRunStream: - """Async run stream with transformer-driven projections. + """Async run stream with caller-driven pumping. - A background asyncio task pumps events from the graph into the - transformer pipeline. This is the standard async pattern — the task - runs on the same event loop and async consumers can iterate - multiple projections concurrently. + Async iteration on any projection drives the graph forward — there + is no background task. Concurrent consumers share a single-flight + pump via an `asyncio.Lock`, so each awaiting cursor contributes one + event per acquisition. Backpressure comes from the logs: when a + subscribed log's buffer reaches `maxlen`, `apush` awaits the + subscriber to drain, which holds back the pump and paces the graph. - All transformer projections live in `extensions`. Native transformer - projections (those with `_native = True`) are also set as direct - attributes on this instance (e.g. `run.values`, `run.messages`). + Projections are single-consumer — a second `aiter(run.values)` + raises. Use `projection.tee(n)` for fan-out. - Async-iterating the run stream yields raw `ProtocolEvent` objects - from the mux's main event log. + Use as an async context manager to guarantee clean shutdown on + early exit: + + ```python + async with await handler.astream(input) as run: + async for msg in run.messages: + ... + ``` """ def __init__( self, + graph_aiter: AsyncIterator[Any], mux: StreamMux, values_transformer: ValuesTransformer, - pump_task: asyncio.Task[None], ) -> None: """Initialize the async run stream. Args: + graph_aiter: Async iterator over the graph's stream. mux: The StreamMux owning projections and the main log. values_transformer: The built-in values transformer providing `output` / `interrupted` / `interrupts`. - pump_task: Background task pumping graph events into the mux. """ + self._graph_aiter = graph_aiter self._mux = mux self.extensions: Mapping[str, Any] = MappingProxyType(mux.extensions) self._values_transformer = values_transformer - self._pump_task = pump_task - # Native-transformer projections also show up as direct attributes. + self._exhausted = False + self._pump_lock = asyncio.Lock() for key in mux.native_keys: setattr(self, key, mux.extensions[key]) + self._wire_arequest_more(mux) + + def _wire_arequest_more(self, mux: StreamMux) -> None: + """Install `_arequest_more` on every async EventLog so cursors + can drive the pump when their buffer catches up.""" + mux._events._arequest_more = self._apump_next + for value in mux.extensions.values(): + if isinstance(value, EventLog): + value._arequest_more = self._apump_next + elif isinstance(value, StreamChannel): + value._log._arequest_more = self._apump_next + + async def _apump_next(self) -> bool: + """Pull one event from the graph and push it through the mux. + + Serialized via `self._pump_lock` so concurrent cursors each + produce one event per acquisition rather than racing on the + graph iterator. + + `except Exception` is intentional — `CancelledError` and other + `BaseException` subclasses propagate, matching asyncio's + cancellation contract. + + Returns: + True if an event was pulled, False if the graph is + exhausted or has raised. + """ + async with self._pump_lock: + if self._exhausted: + return False + try: + part = await self._graph_aiter.__anext__() + except StopAsyncIteration: + self._exhausted = True + await self._mux.aclose() + return False + except Exception as e: + self._exhausted = True + await self._mux.afail(e) + return False + await self._mux.apush(convert_to_protocol_event(part)) + return True + + async def abort(self) -> None: + """Stop the run early. + + Closes the mux and marks the stream exhausted. Any awaiting + cursors wake up and see the closed state; any `apush` blocked + on backpressure wakes and returns without appending. Idempotent. + """ + async with self._pump_lock: + if self._exhausted: + return + self._exhausted = True + try: + await self._mux.aclose() + except Exception: + pass + + async def __aenter__(self) -> AsyncGraphRunStream: + return self + + async def __aexit__( + self, + exc_type: type[BaseException] | None, + exc: BaseException | None, + tb: TracebackType | None, + ) -> None: + await self.abort() async def output(self) -> dict[str, Any] | None: - """Wait for the run to complete and return the final state. + """Drive the run to completion and return the final state. - Methods (not properties) on the async lane so `run.output` without - `await` raises at type-check time instead of silently yielding a - coroutine object that's truthy, lenless, and never awaited. - - The pump routes any Exception into `mux.afail`, which surfaces on - `ValuesTransformer.error`. CancelledError / KeyboardInterrupt - propagate so cancellation isn't silently dropped. + Methods (not properties) on the async lane so `run.output` + without `await` raises at type-check time instead of silently + yielding a coroutine object. Example: ```python @@ -192,52 +342,34 @@ class AsyncGraphRunStream: Raises: BaseException: If the run ended with an error. """ - try: - await self._pump_task - except Exception: - pass + await _adrive_until_done(self._apump_next) if (err := self._values_transformer.error) is not None: raise err return self._values_transformer._latest async def interrupted(self) -> bool: - """Wait for the run to complete and return whether it was interrupted. - - Example: - ```python - interrupted = await run.interrupted() - ``` + """Drive the run to completion and return whether it was + interrupted. Raises: BaseException: If the run ended with an error. """ - try: - await self._pump_task - except Exception: - pass + await _adrive_until_done(self._apump_next) if (err := self._values_transformer.error) is not None: raise err return self._values_transformer._interrupted async def interrupts(self) -> list[Any]: - """Wait for the run to complete and return interrupt payloads. - - Example: - ```python - interrupts = await run.interrupts() - ``` + """Drive the run to completion and return interrupt payloads. Raises: BaseException: If the run ended with an error. """ - try: - await self._pump_task - except Exception: - pass + await _adrive_until_done(self._apump_next) if (err := self._values_transformer.error) is not None: raise err return self._values_transformer._interrupts def __aiter__(self) -> AsyncIterator[ProtocolEvent]: - """Iterate all protocol events from the mux's main event log.""" + """Subscribe to the main event log and iterate protocol events.""" return self._mux._events.__aiter__() diff --git a/libs/langgraph/langgraph/stream/stream_channel.py b/libs/langgraph/langgraph/stream/stream_channel.py index 7174002df..e7e397047 100644 --- a/libs/langgraph/langgraph/stream/stream_channel.py +++ b/libs/langgraph/langgraph/stream/stream_channel.py @@ -93,3 +93,17 @@ class StreamChannel(Generic[T]): def __aiter__(self) -> AsyncIterator[T]: return self._log.__aiter__() + + def tee(self, n: int = 2) -> tuple[Iterator[T], ...]: + """Fan out the channel into `n` independent sync iterators. + + Delegates to the underlying EventLog's `tee()`. + """ + return self._log.tee(n) + + def atee(self, n: int = 2) -> tuple[AsyncIterator[T], ...]: + """Fan out the channel into `n` independent async iterators. + + Delegates to the underlying EventLog's `atee()`. + """ + return self._log.atee(n) diff --git a/libs/langgraph/langgraph/stream/streaming_handler.py b/libs/langgraph/langgraph/stream/streaming_handler.py index 9748c48b5..ff031016e 100644 --- a/libs/langgraph/langgraph/stream/streaming_handler.py +++ b/libs/langgraph/langgraph/stream/streaming_handler.py @@ -1,13 +1,11 @@ from __future__ import annotations -import asyncio from collections.abc import Sequence from typing import Any from langchain_core.runnables import RunnableConfig from langgraph.pregel import Pregel -from langgraph.stream._convert import convert_to_protocol_event from langgraph.stream._mux import StreamMux from langgraph.stream._types import StreamTransformer from langgraph.stream.run_stream import AsyncGraphRunStream, GraphRunStream @@ -64,7 +62,6 @@ 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. @@ -80,12 +77,6 @@ class StreamingHandler: interrupt_after: Nodes to interrupt after, if any. transformers: User transformers appended after the built-in `ValuesTransformer` and `MessagesTransformer`. - max_events: Caps the retention of every EventLog and - StreamChannel the mux binds (main event log plus each - transformer's projection logs), dropping the oldest - when full. Transformers that constructed their own - logs with an explicit `maxlen` keep their setting. - Unbounded when `None`. Returns: A GraphRunStream the caller can iterate to drive the run. @@ -94,7 +85,6 @@ class StreamingHandler: mux = StreamMux( [values_t, MessagesTransformer(), *(transformers or ())], is_async=False, - max_events=max_events, ) graph_iter = iter( @@ -119,12 +109,13 @@ 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. + Returns an AsyncGraphRunStream immediately. The caller's + iteration on any projection drives the graph forward — there + is no background task. Concurrent consumers share a + single-flight pump via an internal `asyncio.Lock`. Args: input: Graph input. @@ -133,37 +124,26 @@ class StreamingHandler: interrupt_after: Nodes to interrupt after, if any. transformers: User transformers appended after the built-in `ValuesTransformer` and `MessagesTransformer`. - max_events: Caps retention of every EventLog and - StreamChannel the mux binds — see `stream()` for the - full semantics. Returns: An AsyncGraphRunStream whose projections can be awaited - concurrently while the background pump runs. + concurrently; each subscribed cursor drives the pump when + its buffer is empty. """ values_t = ValuesTransformer() mux = StreamMux( [values_t, MessagesTransformer(), *(transformers or ())], is_async=True, - max_events=max_events, ) - async def pump() -> None: - try: - async for part in self._graph.astream( - input, - config, - stream_mode=STREAM_V2_MODES, - subgraphs=True, - version="v2", - interrupt_before=interrupt_before, - interrupt_after=interrupt_after, - ): - await mux.apush(convert_to_protocol_event(part)) - await mux.aclose() - except BaseException as e: - await mux.afail(e) + graph_aiter = self._graph.astream( + input, + config, + stream_mode=STREAM_V2_MODES, + subgraphs=True, + version="v2", + interrupt_before=interrupt_before, + interrupt_after=interrupt_after, + ).__aiter__() - task = asyncio.create_task(pump()) - - return AsyncGraphRunStream(mux, values_t, task) + return AsyncGraphRunStream(graph_aiter, mux, values_t) diff --git a/libs/langgraph/langgraph/stream/transformers.py b/libs/langgraph/langgraph/stream/transformers.py index 3fe7c5ea9..dec52515f 100644 --- a/libs/langgraph/langgraph/stream/transformers.py +++ b/libs/langgraph/langgraph/stream/transformers.py @@ -7,7 +7,12 @@ from langgraph.stream._types import ProtocolEvent, StreamTransformer class ValuesTransformer(StreamTransformer): - """Capture values events as an iterable of state snapshots. + """Capture values events as a drainable stream of state snapshots. + + Keeps `_latest` / `_interrupted` / `_interrupts` as scalar state + regardless of whether the log has a subscriber — so `run.output()` + and `run.interrupted` work without forcing the caller to iterate + `run.values`. Log pushes are silent no-ops when unsubscribed. Native transformer — projection keys are exposed as direct attributes on the run stream (e.g. `run.values`). @@ -42,11 +47,11 @@ class ValuesTransformer(StreamTransformer): if params["namespace"]: return True self._latest = params["data"] - self._log.push(params["data"]) interrupts = params.get("interrupts", ()) if interrupts: self._interrupted = True self._interrupts.extend(interrupts) + self._log.push(params["data"]) return True @@ -58,9 +63,9 @@ class MessagesTransformer(StreamTransformer): produces ChatModelStream objects using the protocol handler. Only root-namespace messages events are captured; tokens emitted - from subgraphs are dropped from the `messages` projection. Consumers - that need subgraph tokens should iterate the raw event stream or - register a custom transformer. + from subgraphs are dropped from the `messages` projection. + Consumers that need subgraph tokens should iterate the raw event + stream or register a custom transformer. Native transformer — projection keys are exposed as direct attributes on the run stream (e.g. `run.messages`). diff --git a/libs/langgraph/tests/test_streaming_handler.py b/libs/langgraph/tests/test_streaming_handler.py index d4797e4c1..835158ed5 100644 --- a/libs/langgraph/tests/test_streaming_handler.py +++ b/libs/langgraph/tests/test_streaming_handler.py @@ -15,7 +15,6 @@ from typing_extensions import TypedDict from langgraph.constants import END, START from langgraph.graph import StateGraph from langgraph.stream import ( - BufferOverflowError, EventLog, StreamChannel, StreamingHandler, @@ -142,78 +141,98 @@ class TestEventLog: def test_sync_iteration(self) -> None: log: EventLog[int] = EventLog() log._bind(is_async=False) + it = iter(log) # subscribe before pushing log.push(1) log.push(2) log.push(3) log.close() - assert list(log) == [1, 2, 3] + assert list(it) == [1, 2, 3] - def test_multi_cursor(self) -> None: + def test_drain_on_consume(self) -> None: + """Items are popped as consumed — no retention across iterations.""" log: EventLog[str] = EventLog() log._bind(is_async=False) + it = iter(log) log.push("a") log.push("b") log.close() - # Two independent cursors see all items. - assert list(log) == ["a", "b"] - assert list(log) == ["a", "b"] + assert list(it) == ["a", "b"] + # Buffer is drained. + assert list(log._items) == [] + + def test_second_subscribe_raises(self) -> None: + """Only one subscriber allowed; tee() is the fan-out escape hatch.""" + log: EventLog[str] = EventLog() + log._bind(is_async=False) + log.close() + _ = iter(log) + with pytest.raises(RuntimeError, match="already has a subscriber"): + iter(log) + + def test_pre_subscription_push_is_noop(self) -> None: + """Lazy-subscribe: pushes before subscription are dropped silently.""" + log: EventLog[int] = EventLog() + log._bind(is_async=False) + log.push(1) + log.push(2) + it = iter(log) + log.push(3) + log.close() + # Only the post-subscribe push survives. + assert list(it) == [3] def test_fail_propagation(self) -> None: log: EventLog[int] = EventLog() log._bind(is_async=False) + it = iter(log) log.push(1) log.fail(ValueError("test error")) with pytest.raises(ValueError, match="test error"): - list(log) + list(it) @pytest.mark.anyio async def test_async_iteration(self) -> None: log: EventLog[int] = EventLog() log._bind(is_async=True) - - async def producer(): - for i in range(3): - log.push(i) - log.close() - - producer_task = asyncio.create_task(producer()) - items = [item async for item in log] - await producer_task + cursor = aiter(log) + for i in range(3): + log.push(i) + log.close() + items = [item async for item in cursor] assert items == [0, 1, 2] @pytest.mark.anyio - async def test_async_multi_cursor(self) -> None: + async def test_async_second_subscribe_raises(self) -> None: log: EventLog[str] = EventLog() log._bind(is_async=True) - log.push("x") - log.push("y") log.close() - items1 = [item async for item in log] - items2 = [item async for item in log] - assert items1 == ["x", "y"] - assert items2 == ["x", "y"] + _ = log.__aiter__() + with pytest.raises(RuntimeError, match="already has a subscriber"): + log.__aiter__() @pytest.mark.anyio async def test_async_fail(self) -> None: log: EventLog[int] = EventLog() log._bind(is_async=True) + cursor = aiter(log) log.push(1) log.fail(RuntimeError("async error")) with pytest.raises(RuntimeError, match="async error"): - async for _ in log: + async for _ in cursor: pass def test_sync_cursor_yields_items_before_error(self) -> None: """Sync cursor should yield all buffered items before raising.""" log: EventLog[int] = EventLog() log._bind(is_async=False) + it = iter(log) log.push(1) log.push(2) log.push(3) log.fail(ValueError("late error")) items: list[int] = [] with pytest.raises(ValueError, match="late error"): - for item in log: + for item in it: items.append(item) assert items == [1, 2, 3] @@ -222,30 +241,38 @@ class TestEventLog: """Async cursor should yield all buffered items before raising.""" log: EventLog[int] = EventLog() log._bind(is_async=True) + cursor = aiter(log) log.push(1) log.push(2) log.push(3) log.fail(ValueError("late error")) items: list[int] = [] with pytest.raises(ValueError, match="late error"): - async for item in log: + async for item in cursor: items.append(item) assert items == [1, 2, 3] def test_push_after_close_raises(self) -> None: - """Push after close should raise RuntimeError.""" + """Push after close should raise RuntimeError (when subscribed).""" log: EventLog[int] = EventLog() + log._bind(is_async=False) + it = iter(log) log.push(1) log.close() with pytest.raises(RuntimeError, match="Cannot push to a closed EventLog"): log.push(2) + _ = list(it) def test_push_after_fail_raises(self) -> None: - """Fail closes the log, so push after fail should also raise.""" + """Fail closes the log, so push after fail should also raise (when subscribed).""" log: EventLog[int] = EventLog() + log._bind(is_async=False) + it = iter(log) log.fail(ValueError("err")) with pytest.raises(RuntimeError, match="Cannot push to a closed EventLog"): log.push(1) + with pytest.raises(ValueError, match="err"): + list(it) def test_empty_log_sync(self) -> None: """Iterating a closed empty log should yield nothing.""" @@ -321,31 +348,35 @@ class TestStreamChannel: def test_push_and_iterate(self) -> None: ch: StreamChannel[str] = StreamChannel("test") ch._bind(is_async=False) + it = iter(ch) ch.push("a") ch.push("b") ch._close() - assert list(ch) == ["a", "b"] + assert list(it) == ["a", "b"] def test_wire_callback(self) -> None: forwarded: list[str] = [] ch: StreamChannel[str] = StreamChannel("test") ch._bind(is_async=False) ch._wire(lambda item: forwarded.append(item)) + it = iter(ch) ch.push("x") ch.push("y") ch._close() + # Wire callback fires on every push, regardless of subscription. assert forwarded == ["x", "y"] - assert list(ch) == ["x", "y"] + assert list(it) == ["x", "y"] def test_fail_propagation(self) -> None: """_fail() should propagate the error through the underlying log.""" ch: StreamChannel[str] = StreamChannel("test") ch._bind(is_async=False) + it = iter(ch) ch.push("a") ch._fail(ValueError("channel error")) items: list[str] = [] with pytest.raises(ValueError, match="channel error"): - for item in ch: + for item in it: items.append(item) assert items == ["a"] @@ -354,10 +385,11 @@ class TestStreamChannel: """Async iteration should delegate to the inner event log.""" ch: StreamChannel[str] = StreamChannel("test") ch._bind(is_async=True) - ch.push("x") - ch.push("y") + cursor = ch.__aiter__() + ch._log.push("x") + ch._log.push("y") ch._close() - items = [item async for item in ch] + items = [item async for item in cursor] assert items == ["x", "y"] def test_push_without_wire(self) -> None: @@ -365,9 +397,10 @@ class TestStreamChannel: ch: StreamChannel[int] = StreamChannel("test") ch._bind(is_async=False) assert ch._wire_fn is None + it = iter(ch) ch.push(42) ch._close() - assert list(ch) == [42] + assert list(it) == [42] # --------------------------------------------------------------------------- @@ -440,6 +473,49 @@ class TestStreamingHandlerSync: assert custom_events[0]["params"]["data"] == {"step": "start"} assert custom_events[1]["params"]["data"] == {"step": "end"} + def test_interleave_values_and_messages(self) -> None: + graph = _build_simple_graph() + handler = StreamingHandler(graph) + run = handler.stream({"value": "x", "items": []}) + + tagged = list(run.interleave("values", "messages")) + names = [name for name, _ in tagged] + assert set(names).issubset({"values", "messages"}) + # Values projection must have fired at least once. + assert names.count("values") >= 1 + # Values have been drained by the interleave cursor — re-subscribing raises. + with pytest.raises(RuntimeError, match="already has a subscriber"): + list(run.values) + + def test_abort_marks_exhausted_and_closes_mux(self) -> None: + graph = _build_simple_graph() + handler = StreamingHandler(graph) + run = handler.stream({"value": "x", "items": []}) + values_iter = iter(run.values) + # Consume one item so the pump advances. + _ = next(values_iter) + run.abort() + # Remaining iteration yields whatever was buffered, then stops. + list(values_iter) + assert run._exhausted is True + # Second abort is idempotent. + run.abort() + + def test_context_manager_calls_abort_on_exit(self) -> None: + graph = _build_simple_graph() + handler = StreamingHandler(graph) + with handler.stream({"value": "x", "items": []}) as run: + values_iter = iter(run.values) + _ = next(values_iter) + assert run._exhausted is True + + def test_interleave_unknown_projection(self) -> None: + graph = _build_simple_graph() + handler = StreamingHandler(graph) + run = handler.stream({"value": "x", "items": []}) + with pytest.raises(KeyError): + list(run.interleave("values", "does_not_exist")) + class TestStreamingHandlerSyncErrors: def test_error_propagation_output(self) -> None: @@ -533,6 +609,32 @@ class TestStreamingHandlerAsync: for event in events: assert event["type"] == "event" + @pytest.mark.anyio + @NEEDS_CONTEXTVARS + async def test_abort_marks_exhausted_and_closes_mux(self) -> None: + graph = _build_simple_graph() + handler = StreamingHandler(graph) + run = await handler.astream({"value": "x", "items": []}) + values_iter = aiter(run.values) + _ = await anext(values_iter) + await run.abort() + # Drain the rest; should terminate promptly now that mux is closed. + async for _item in values_iter: + pass + assert run._exhausted is True + await run.abort() # idempotent + + @pytest.mark.anyio + @NEEDS_CONTEXTVARS + async def test_async_context_manager_calls_abort_on_exit(self) -> None: + graph = _build_simple_graph() + handler = StreamingHandler(graph) + run = await handler.astream({"value": "x", "items": []}) + async with run: + values_iter = aiter(run.values) + _ = await anext(values_iter) + assert run._exhausted is True + @pytest.mark.anyio @NEEDS_CONTEXTVARS async def test_extensions_has_native_keys(self) -> None: @@ -699,13 +801,14 @@ class TestStreamMux: return event["method"] != "updates" mux = StreamMux([FilterTransformer()]) + it = iter(mux._events) mux.push(_event("values", {"a": 1})) mux.push(_event("updates", {"b": 2})) mux.push(_event("custom", {"c": 3})) mux.close() - events = list(mux._events) + events = list(it) methods = [e["method"] for e in events] assert "updates" not in methods assert methods == ["values", "custom"] @@ -744,9 +847,10 @@ class TestStreamMux: def test_empty_mux(self) -> None: """Push/close/fail on a mux with no transformers should work.""" mux = StreamMux() + it = iter(mux._events) mux.push(_event("values", {"x": 1})) mux.close() - events = list(mux._events) + events = list(it) assert len(events) == 1 assert events[0]["method"] == "values" @@ -769,12 +873,13 @@ class TestValuesTransformer: t = ValuesTransformer() t.init() t._log._bind(is_async=False) + it = iter(t._log) t.process(_event("values", {"val": "root"})) t.process(_event("values", {"val": "sub"}, namespace=["sub"])) t._log.close() - items = list(t._log) + items = list(it) assert len(items) == 1 assert items[0]["val"] == "root" @@ -783,11 +888,12 @@ class TestValuesTransformer: t = ValuesTransformer() t.init() t._log._bind(is_async=False) + it = iter(t._log) result = t.process(_event("updates", {"x": 1})) assert result is True # passed through t._log.close() - assert list(t._log) == [] # but not captured + assert list(it) == [] # but not captured def test_tracks_interrupts(self) -> None: """Interrupts should be accumulated across events.""" @@ -810,10 +916,11 @@ class TestMessagesTransformer: t = MessagesTransformer() t.init() t._log._bind(is_async=False) + it = iter(t._log) t.process(_event("messages", ("chunk", {"meta": True}))) t._log.close() - items = list(t._log) + items = list(it) assert len(items) == 1 assert items[0] == ("chunk", {"meta": True}) @@ -821,28 +928,31 @@ class TestMessagesTransformer: t = MessagesTransformer() t.init() t._log._bind(is_async=False) + it = iter(t._log) t.process(_event("messages", ("chunk", {}), namespace=["sub"])) t._log.close() - assert list(t._log) == [] + assert list(it) == [] def test_ignores_non_messages_methods(self) -> None: t = MessagesTransformer() t.init() t._log._bind(is_async=False) + it = iter(t._log) result = t.process(_event("values", {"v": 1})) assert result is True t._log.close() - assert list(t._log) == [] + assert list(it) == [] def test_fail_propagates(self) -> None: t = MessagesTransformer() t.init() t._log._bind(is_async=False) + it = iter(t._log) t._log.fail(ValueError("msg error")) with pytest.raises(ValueError, match="msg error"): - list(t._log) + list(it) # --------------------------------------------------------------------------- @@ -975,10 +1085,11 @@ class TestCustomTransformer: handler = StreamingHandler(graph) counter_t = CounterTransformer() run = handler.stream({"value": "x", "items": []}, transformers=[counter_t]) - _ = run.output assert "counter" in run.extensions - # Counter channel should have been pushed to. - counts = list(run.extensions["counter"]) + # Subscribe before driving the run so channel pushes are retained. + counter_iter = iter(run.extensions["counter"]) + _ = run.output + counts = list(counter_iter) assert len(counts) > 0 # Non-native transformer should not set direct attributes. assert not hasattr(run, "counter") @@ -1005,12 +1116,14 @@ class TestCustomTransformer: handler = StreamingHandler(graph) foo_t = FooTransformer() run = handler.stream({"value": "x", "items": []}, transformers=[foo_t]) + # Subscribe before driving the run. + foo_iter = iter(run.foo) _ = run.output # foo should be both in extensions and as a direct attribute. assert "foo" in run.extensions assert hasattr(run, "foo") assert run.foo is run.extensions["foo"] - items = list(run.foo) + items = list(foo_iter) assert "saw_values" in items def test_stream_channel_auto_forward(self) -> None: @@ -1062,12 +1175,13 @@ class TestCustomTransformer: return True mux = StreamMux([ChannelPusher()]) + it = iter(mux._events) mux.push(_event("values")) mux.push(_event("updates")) mux.close() - events = list(mux._events) + events = list(it) seqs = [e["seq"] for e in events] # Seq numbers must be strictly increasing. for i in range(1, len(seqs)): @@ -1116,12 +1230,13 @@ class TestEventLogAutoLifecycle: return True mux = StreamMux([SimpleTransformer()]) + it = iter(mux._events) mux.push(_event("values")) mux.close() # The log should have been auto-closed — iteration should work. - items = list(mux._events) + items = list(it) assert len(items) == 1 def test_mux_auto_fails_event_logs(self) -> None: @@ -1141,13 +1256,14 @@ class TestEventLogAutoLifecycle: t = SimpleTransformer() mux = StreamMux([t]) + it = iter(t._log) mux.push(_event("values")) mux.fail(ValueError("boom")) # The transformer's log should have been auto-failed. with pytest.raises(ValueError, match="boom"): - list(t._log) + list(it) def test_no_double_close_if_transformer_closes_own_log(self) -> None: """If a transformer closes its log in finalize(), mux should not error.""" @@ -1191,8 +1307,9 @@ class TestEventLogAutoLifecycle: handler = StreamingHandler(graph) t = MinimalTransformer() run = handler.stream({"value": "x", "items": []}, transformers=[t]) + minimal_iter = iter(run.extensions["minimal"]) _ = run.output - items = list(run.extensions["minimal"]) + items = list(minimal_iter) assert len(items) > 0 @@ -1451,13 +1568,14 @@ class TestAsyncTransformerLane: async_t = AsyncOne() mux = StreamMux([SyncOne(), async_t], is_async=True) + seen_cursor = aiter(async_t._log) await mux.apush(_event("values", {})) await mux.apush(_event("updates", {})) await mux.aclose() assert seen_sync == ["values", "updates"] - items = [x async for x in async_t._log] + items = [x async for x in seen_cursor] assert items == ["values", "updates"] @pytest.mark.anyio @@ -1492,148 +1610,146 @@ class TestAsyncTransformerLane: {"value": "x", "items": []}, transformers=[Scorer()], ) + # Subscribe before driving the run so scheduled pushes are retained. + scores_cursor = aiter(run.extensions["scores"]) _ = await run.output() - scores = [x async for x in run.extensions["scores"]] + scores = [x async for x in scores_cursor] assert scores and all(s == 42 for s in scores) # --------------------------------------------------------------------------- -# Bounded EventLog / StreamChannel — memory caps and overflow semantics +# Drain-on-consume semantics — bounded backpressure, single subscriber # --------------------------------------------------------------------------- -class TestBoundedEventLog: +class TestMemoryBounds: + """Drain-on-consume guarantees memory stays bounded for the common + access patterns. These tests lock in the property — if a change + re-introduces retention, they should fail.""" + + def test_sync_subscribed_buffer_stays_at_most_one_between_yields(self) -> None: + """With a single sync consumer, the pump produces exactly one event + per cursor advance, so the buffer never holds more than one.""" + graph = _build_simple_graph() + handler = StreamingHandler(graph) + run = handler.stream({"value": "x", "items": []}) + events_iter = iter(run) + max_buffered = 0 + count = 0 + for _ in events_iter: + max_buffered = max(max_buffered, len(run._mux._events._items)) + count += 1 + assert count > 0 + assert max_buffered == 0, ( + f"Subscribed buffer should hold 0 items after each yield " + f"(drain-on-consume), observed max {max_buffered}" + ) + + def test_unsubscribed_projections_never_accumulate(self) -> None: + """Projections without a subscriber drop pushes silently — + their buffers stay empty regardless of run length.""" + graph = _build_simple_graph() + handler = StreamingHandler(graph) + run = handler.stream({"value": "x", "items": []}) + # Subscribe to main events only; leave values and messages unsubscribed. + list(run) + values_log = run.extensions["values"] + messages_log = run.extensions["messages"] + assert len(values_log._items) == 0 + assert len(messages_log._items) == 0 + assert values_log._subscribed is False + assert messages_log._subscribed is False + + def test_output_path_does_not_retain_values(self) -> None: + """`run.output` is a scalar accessor — it updates `_latest` from + process() without populating the log, so the values log buffer + stays empty even across a full run.""" + graph = _build_simple_graph() + handler = StreamingHandler(graph) + run = handler.stream({"value": "x", "items": []}) + _ = run.output + values_log = run.extensions["values"] + assert len(values_log._items) == 0 + assert values_log._subscribed is False + + def test_drained_subscriber_buffer_returns_to_empty(self) -> None: + """After fully draining a subscribed log, the internal deque is empty.""" + graph = _build_simple_graph() + handler = StreamingHandler(graph) + run = handler.stream({"value": "x", "items": []}) + values_log = run.extensions["values"] + list(run.values) + assert len(values_log._items) == 0 + + @pytest.mark.anyio + @NEEDS_CONTEXTVARS + async def test_async_single_consumer_buffer_stays_at_most_one(self) -> None: + """Same drain-on-consume guarantee for the async lane.""" + graph = _build_simple_graph() + handler = StreamingHandler(graph) + run = await handler.astream({"value": "x", "items": []}) + max_buffered = 0 + count = 0 + async for _ in run: + max_buffered = max(max_buffered, len(run._mux._events._items)) + count += 1 + assert count > 0 + assert max_buffered == 0 + + @pytest.mark.anyio + @NEEDS_CONTEXTVARS + async def test_async_unsubscribed_projections_never_accumulate(self) -> None: + """Projections with no async subscriber stay empty under astream.""" + graph = _build_simple_graph() + handler = StreamingHandler(graph) + run = await handler.astream({"value": "x", "items": []}) + _ = await run.output() + values_log = run.extensions["values"] + messages_log = run.extensions["messages"] + assert len(values_log._items) == 0 + assert len(messages_log._items) == 0 + assert values_log._subscribed is False + assert messages_log._subscribed is False + + +class TestDrainOnConsume: 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.""" + def test_push_unbounded_by_design(self) -> None: + """Push is non-blocking and doesn't enforce capacity — the + caller-driven pump bounds memory via iteration pace.""" log: EventLog[int] = EventLog() log._bind(is_async=False) + it = iter(log) 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] + assert list(it) == list(range(100)) @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) + async def test_atee_fans_out(self) -> None: + """atee provides the documented fan-out for concurrent consumers.""" + log: EventLog[int] = EventLog() 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) + a, b = log.atee(2) + for i in range(3): + log.push(i) log.close() - assert list(log) == [1, 2] + items_a = [x async for x in a] + items_b = [x async for x in b] + assert items_a == [0, 1, 2] + assert items_b == [0, 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 + def test_tee_fans_out_sync(self) -> None: + log: EventLog[int] = EventLog() + log._bind(is_async=False) + a, b = log.tee(2) + for i in range(3): + log.push(i) + log.close() + assert list(a) == [0, 1, 2] + assert list(b) == [0, 1, 2]