WEBDIENSTE/langgraph
1
0
Fork 0
mirror of https://github.com/langchain-ai/langgraph.git synced 2026-08-24 08:32:24 +02:00
Code Issues Packages Projects Releases Wiki Activity

Merge streaming-transformer drain-on-consume redesign

Browse Source 
# Conflicts:
#	libs/langgraph/langgraph/stream/run_stream.py
#	libs/langgraph/langgraph/stream/streaming_handler.py
#	libs/langgraph/langgraph/stream/transformers.py
#	libs/langgraph/tests/test_streaming_handler.py
This commit is contained in:
Nick Hollon
2026-04-18 12:38:46 -04:00
parent 7e5df56688 ab1d6980b5
commit b6a196fac6
8 changed files with 710 additions and 418 deletions
Download Patch File Download Diff File Expand all files Collapse all files
libs/langgraph/langgraph/stream/__init__.py
+1 -2
View File
@@ -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",
libs/langgraph/langgraph/stream/_event_log.py
+170 -115
View File
@@ -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))
libs/langgraph/langgraph/stream/_mux.py
+5 -15
View File
@@ -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.
libs/langgraph/langgraph/stream/run_stream.py
+208 -75
View File
@@ -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,23 +59,18 @@ 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.
Also calls `_bind_pump` on any transformer that exposes it, so
that transformers producing ChatModelStream objects (e.g.
MessagesTransformer) can wire the pull callback on each stream as
it's created.
transformers producing ChatModelStream objects (e.g.
MessagesTransformer) can wire the pull callback on each stream
as it's created.
"""
mux._events._request_more = self._pump_next
for value in mux.extensions.values():
@@ -91,22 +97,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
@@ -114,12 +141,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
@@ -127,70 +155,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
@@ -200,52 +351,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__()
libs/langgraph/langgraph/stream/stream_channel.py
+14
View File
@@ -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)
libs/langgraph/langgraph/stream/streaming_handler.py
+16 -36
View File
@@ -1,6 +1,5 @@
from __future__ import annotations
import asyncio
from collections.abc import Sequence
from typing import Any
@@ -8,7 +7,6 @@ from langchain_core.runnables import RunnableConfig
from langgraph._internal._constants import CONF, CONFIG_KEY_STREAM_MESSAGES_V2
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
@@ -83,7 +81,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.
@@ -99,12 +96,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.
@@ -113,7 +104,6 @@ class StreamingHandler:
mux = StreamMux(
[values_t, MessagesTransformer(), *(transformers or ())],
is_async=False,
max_events=max_events,
)
graph_iter = iter(
@@ -138,12 +128,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.
@@ -152,37 +143,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,
_merge_v2_messages_flag(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,
_merge_v2_messages_flag(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)
libs/langgraph/langgraph/stream/transformers.py
+7 -2
View File
@@ -18,7 +18,12 @@ if TYPE_CHECKING:
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`).
@@ -53,11 +58,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
libs/langgraph/tests/test_streaming_handler.py
+289 -173
View File
@@ -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."""
@@ -812,6 +918,7 @@ class TestMessagesTransformer:
t.init()
t._log._bind(is_async=False)
t._bind_pump(lambda: False)
it = iter(t._log)
meta = {"langgraph_node": "llm", "run_id": "run-1"}
for evt in (
@@ -820,7 +927,7 @@ class TestMessagesTransformer:
):
t.process(_event("messages", (evt, meta)))
t._log.close()
items = list(t._log)
items = list(it)
assert len(items) == 1
# Items in the messages log are ChatModelStream objects, not raw
# tuples — the content-block-centric projection.
@@ -832,6 +939,7 @@ class TestMessagesTransformer:
t.init()
t._log._bind(is_async=False)
t._bind_pump(lambda: False)
it = iter(t._log)
meta = {"langgraph_node": "llm", "run_id": "run-1"}
t.process(
@@ -842,25 +950,27 @@ class TestMessagesTransformer:
)
)
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)
# ---------------------------------------------------------------------------
@@ -993,10 +1103,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")
@@ -1023,12 +1134,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:
@@ -1080,12 +1193,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)):
@@ -1134,12 +1248,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:
@@ -1159,13 +1274,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."""
@@ -1209,8 +1325,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
@@ -1469,13 +1586,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
@@ -1510,148 +1628,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]