From 5af4c5addf42e651fa59bfba7187235935b8a4ae Mon Sep 17 00:00:00 2001 From: Nick Hollon Date: Tue, 28 Apr 2026 18:43:51 -0400 Subject: [PATCH] refactor(langgraph,prebuilt): merge EventLog into StreamChannel with optional name (#7637) --- libs/langgraph/langgraph/stream/__init__.py | 2 - libs/langgraph/langgraph/stream/_event_log.py | 306 ---------------- libs/langgraph/langgraph/stream/_mux.py | 121 +++---- libs/langgraph/langgraph/stream/_types.py | 21 +- .../langgraph/stream/stream_channel.py | 340 ++++++++++++++---- .../langgraph/stream/transformers.py | 9 +- libs/langgraph/tests/test_pregel_stream_v2.py | 109 +++--- .../test_stream_lifecycle_transformer.py | 14 +- .../tests/test_stream_messages_transformer.py | 22 +- .../tests/test_stream_subgraph_transformer.py | 6 +- .../langgraph/prebuilt/_tool_call_stream.py | 10 +- .../prebuilt/_tool_call_transformer.py | 15 +- .../tests/test_tool_call_transformer.py | 4 +- 13 files changed, 435 insertions(+), 544 deletions(-) delete mode 100644 libs/langgraph/langgraph/stream/_event_log.py diff --git a/libs/langgraph/langgraph/stream/__init__.py b/libs/langgraph/langgraph/stream/__init__.py index 6e1ae7161..1830597d0 100644 --- a/libs/langgraph/langgraph/stream/__init__.py +++ b/libs/langgraph/langgraph/stream/__init__.py @@ -5,7 +5,6 @@ Compile a graph with `transformers=[...]` and call `graph.stream_v2()` / graph's raw events into ergonomic per-channel streams. """ -from langgraph.stream._event_log import EventLog from langgraph.stream._types import ProtocolEvent, StreamTransformer from langgraph.stream.run_stream import ( AsyncGraphRunStream, @@ -24,7 +23,6 @@ from langgraph.stream.transformers import ( __all__ = [ "AsyncGraphRunStream", "AsyncSubgraphRunStream", - "EventLog", "GraphRunStream", "LifecyclePayload", "LifecycleTransformer", diff --git a/libs/langgraph/langgraph/stream/_event_log.py b/libs/langgraph/langgraph/stream/_event_log.py deleted file mode 100644 index 05bc10504..000000000 --- a/libs/langgraph/langgraph/stream/_event_log.py +++ /dev/null @@ -1,306 +0,0 @@ -from __future__ import annotations - -import asyncio -from collections import deque -from collections.abc import AsyncIterator, Awaitable, Callable, Iterator -from typing import Generic, TypeVar - -T = TypeVar("T") - - -class EventLog(Generic[T]): - """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`. - - 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. - - 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. - - 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: 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`. - """ - if maxlen is not None and maxlen <= 0: - raise ValueError("EventLog maxlen must be a positive int or None") - self._items: deque[T] = deque() - self._maxlen: int | None = maxlen - self._closed = False - self._error: BaseException | None = None - - # Binding state — None means unbound. - self._is_async: bool | None = None - - # Flipped on first __iter__ / __aiter__. Pre-subscription - # pushes are silent no-ops. - self._subscribed = False - - # 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 - # ------------------------------------------------------------------ - - def _bind(self, *, is_async: bool) -> None: - """Bind this log to sync or async mode. - - Called by the StreamMux after transformer registration. Must be - called exactly once before any iteration. - - Args: - is_async: True to enable async iteration, False for sync. - - Raises: - RuntimeError: If the log has already been bound. - """ - if self._is_async is not None: - raise RuntimeError("EventLog is already bound") - self._is_async = is_async - - # ------------------------------------------------------------------ - # Producer API - # ------------------------------------------------------------------ - - def push(self, item: T) -> None: - """Append an item. No-op when no subscriber is registered. - - 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 (and subscribed). - """ - if not self._subscribed: - return - if self._closed: - raise RuntimeError("Cannot push to a closed EventLog") - self._items.append(item) - - def close(self) -> None: - """Mark the log as complete.""" - self._closed = True - - def fail(self, err: BaseException) -> None: - """Mark the log as errored. - - Args: - err: The exception to surface to the subscriber. - """ - self._error = err - self._closed = True - - # ------------------------------------------------------------------ - # Sync iteration (caller-driven pump) - # ------------------------------------------------------------------ - - def __iter__(self) -> Iterator[T]: - """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( - "EventLog has not been bound yet. " - "Register the transformer with a StreamMux first." - ) - if self._is_async: - 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]: - while True: - 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: - if not self._request_more(): - if not self._items and not self._closed: - return - else: - return - - # ------------------------------------------------------------------ - # Async iteration (caller-driven pump) - # ------------------------------------------------------------------ - - def __aiter__(self) -> AsyncIterator[T]: - """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( - "EventLog has not been bound yet. " - "Register the transformer with a StreamMux first." - ) - 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]: - while True: - 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: - 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 34578628d..a31e325e9 100644 --- a/libs/langgraph/langgraph/stream/_mux.py +++ b/libs/langgraph/langgraph/stream/_mux.py @@ -5,7 +5,6 @@ import time from collections.abc import Awaitable, Callable from typing import Any -from langgraph.stream._event_log import EventLog from langgraph.stream._types import ( ProtocolEvent, StreamTransformer, @@ -28,14 +27,15 @@ class StreamMux: """Central event dispatcher for the streaming infrastructure. Owns the main event log and routes events through a transformer - pipeline. StreamChannels discovered in transformer projections are - auto-wired so that every `push()` also injects a `ProtocolEvent` - into the main log. + pipeline. StreamChannels with a name discovered in transformer + projections are auto-wired so that every `push()` also injects a + `ProtocolEvent` into the main log. StreamChannels without a name + are local-only. Pass `is_async=True` when the mux will be consumed via async - iteration (`handler.astream()`). All EventLog and StreamChannel - instances discovered during registration are automatically bound - to the matching mode. + iteration (`handler.astream()`). All StreamChannel instances + discovered during registration are automatically bound to the + matching mode. Attributes: extensions: Merged projection dict across all registered @@ -60,7 +60,7 @@ class StreamMux: `factories` (callables producing fresh instances per mux). Each transformer's `init()` is called, projections are merged into `extensions`, `_native` keys are recorded in `native_keys`, and - any EventLog / StreamChannel instances are bound and wired. + any StreamChannel instances are bound and (if named) wired. Args: transformers: Already-built transformer instances. Registered @@ -89,11 +89,10 @@ class StreamMux: self.is_async = is_async self.scope: tuple[str, ...] = scope self._assign_seq = _assign_seq - self._events: EventLog[ProtocolEvent] = EventLog() + self._events: StreamChannel[ProtocolEvent] = StreamChannel() self._events._bind(is_async=is_async) self._transformers: list[StreamTransformer] = [] self._channels: list[StreamChannel[Any]] = [] - self._logs: list[EventLog[Any]] = [] self._seq = 0 self.extensions: dict[str, Any] = {} @@ -135,18 +134,15 @@ class StreamMux: Records the pump on the mux so child mini-muxes built by `_make_child` can inherit it. Propagates to: - the main event log (`self._events`) - - every projection EventLog and StreamChannel in `extensions` + - every projection StreamChannel in `extensions` - any registered transformer that exposes `_bind_pump` (e.g. `MessagesTransformer` so `ChatModelStream` instances drive the shared pump from their cursors) """ self._pump_fn = fn self._events._request_more = fn - for value in self.extensions.values(): - if isinstance(value, EventLog): - value._request_more = fn - elif isinstance(value, StreamChannel): - value._log._request_more = fn + for ch in self._channels: + ch._request_more = fn for transformer in self._transformers: bind = getattr(transformer, "_bind_pump", None) if bind is not None: @@ -156,11 +152,8 @@ class StreamMux: """Async counterpart to `bind_pump`.""" self._apump_fn = fn self._events._arequest_more = fn - for value in self.extensions.values(): - if isinstance(value, EventLog): - value._arequest_more = fn - elif isinstance(value, StreamChannel): - value._log._arequest_more = fn + for ch in self._channels: + ch._arequest_more = fn for transformer in self._transformers: abind = getattr(transformer, "_bind_apump", None) if abind is not None: @@ -204,8 +197,8 @@ class StreamMux: """Register a single transformer. Calls `transformer.init()`, stores the transformer for event - processing, binds any EventLog or StreamChannel instances in - the projection, and merges the projection into `extensions`. + processing, binds any StreamChannel instances in the projection, + and merges the projection into `extensions`. """ if transformer_requires_async(transformer) and not self.is_async: raise RuntimeError( @@ -274,12 +267,12 @@ class StreamMux: def close(self) -> None: """Finalize all transformers, close all projections and the main log. - EventLogs and StreamChannels discovered in transformer - projections are auto-closed after `finalize()` runs — - transformers don't need to close them manually. If any - transformer's `finalize()` raises, the remaining transformers, - projections, and the main log are still closed; the first error - is re-raised after cleanup completes. + StreamChannels discovered in transformer projections are + auto-closed after `finalize()` runs — transformers don't need + to close them manually. If any transformer's `finalize()` raises, + the remaining transformers, projections, and the main log are + still closed; the first error is re-raised after cleanup + completes. Raises: BaseException: The first error raised by a transformer's @@ -292,12 +285,9 @@ class StreamMux: except BaseException as e: if first_error is None: first_error = e - for log in self._logs: - if not log._closed: - log.close() for ch in self._channels: - if not ch._log._closed: - ch._close() + if not ch._closed: + ch.close() self._events.close() if first_error is not None: raise first_error @@ -305,11 +295,10 @@ class StreamMux: def fail(self, err: BaseException) -> None: """Fail all transformers, projections, and the main log. - EventLogs and StreamChannels discovered in transformer - projections are auto-failed — transformers don't need to fail - them manually. If any transformer's `fail()` raises, the - remaining transformers, projections, and the main log are still - failed. + StreamChannels discovered in transformer projections are + auto-failed — transformers don't need to fail them manually. + If any transformer's `fail()` raises, the remaining + transformers, projections, and the main log are still failed. Args: err: The exception that ended the run. @@ -319,12 +308,9 @@ class StreamMux: transformer.fail(err) except BaseException: pass - for log in self._logs: - if not log._closed: - log.fail(err) for ch in self._channels: - if not ch._log._closed: - ch._fail(err) + if not ch._closed: + ch.fail(err) self._events.fail(err) # ------------------------------------------------------------------ @@ -345,7 +331,7 @@ class StreamMux: `put_nowait` shape. The root mux assigns `seq`; child muxes do not, so forwarded subgraph events can be shared without copying. Memory is bounded by caller pace via the caller-driven pump; see - `EventLog` for the full tradeoff story. + `StreamChannel` for the full tradeoff story. Args: event: The protocol event to dispatch. @@ -365,7 +351,7 @@ class StreamMux: Awaits every task started via `StreamTransformer.schedule()` across all transformers, then calls `afinalize()` on each, - then auto-closes logs, channels, and the main event log. + then auto-closes channels and the main event log. If any scheduled task raised under `on_error="raise"`, or any transformer's `afinalize` raises, the exception propagates. @@ -397,12 +383,9 @@ class StreamMux: except BaseException as e: if first_error is None: first_error = e - for log in self._logs: - if not log._closed: - log.close() for ch in self._channels: - if not ch._log._closed: - ch._close() + if not ch._closed: + ch.close() self._events.close() if first_error is not None: raise first_error @@ -412,7 +395,7 @@ class StreamMux: Cancels every scheduled task across all transformers, awaits them to completion, then runs each transformer's `afail` hook - and auto-fails logs, channels, and the main event log. + and auto-fails channels and the main event log. Args: err: The exception that ended the run. @@ -428,12 +411,9 @@ class StreamMux: await transformer.afail(err) except BaseException: pass - for log in self._logs: - if not log._closed: - log.fail(err) for ch in self._channels: - if not ch._log._closed: - ch._fail(err) + if not ch._closed: + ch.fail(err) if not self._events._closed: self._events.fail(err) @@ -453,32 +433,33 @@ class StreamMux: def _bind_and_wire( self, projection: dict[str, Any], *, native: bool = False ) -> None: - """Bind and wire EventLog / StreamChannel instances in a projection. + """Bind and optionally wire StreamChannel instances in a projection. + + All StreamChannels are bound and tracked. Channels with a name + are additionally wired for protocol auto-forwarding. Args: projection: The projection dict returned by a transformer's `init()`. native: True when the owning transformer is `_native`. - Channels owned by a native transformer use the channel - name directly as the protocol method; user-defined - channels are prefixed with `custom:`. + Named channels owned by a native transformer use the + channel name directly as the protocol method; + user-defined channels are prefixed with `custom:`. """ for value in projection.values(): if isinstance(value, StreamChannel): value._bind(is_async=self.is_async) self._channels.append(value) - method = value.name if native else f"custom:{value.name}" + if value.name is not None: + method = value.name if native else f"custom:{value.name}" - def _make_forward(method_name: str) -> Callable[[Any], None]: - def _forward(item: Any) -> None: - self._forward(method_name, item) + def _make_forward(method_name: str) -> Callable[[Any], None]: + def _forward(item: Any) -> None: + self._forward(method_name, item) - return _forward + return _forward - value._wire(_make_forward(method)) - elif isinstance(value, EventLog): - value._bind(is_async=self.is_async) - self._logs.append(value) + value._wire(_make_forward(method)) def _forward(self, method: str, item: Any) -> None: """Inject a ProtocolEvent for a StreamChannel push. diff --git a/libs/langgraph/langgraph/stream/_types.py b/libs/langgraph/langgraph/stream/_types.py index 530dd7c87..4001c8860 100644 --- a/libs/langgraph/langgraph/stream/_types.py +++ b/libs/langgraph/langgraph/stream/_types.py @@ -45,8 +45,7 @@ class StreamTransformer(ABC): """Extension point for custom stream projections. Transformers observe protocol events flowing through the StreamMux and - build typed derived projections (EventLogs, StreamChannels, promises, - etc.). + build typed derived projections (StreamChannels, promises, etc.). Set `_native = True` on a transformer to have its projection keys exposed as direct attributes on the run stream (in addition to @@ -55,9 +54,9 @@ class StreamTransformer(ABC): Subclasses must implement `init` and override at least one of `process` / `aprocess`. The `finalize` / `afinalize` and `fail` / `afail` hooks are optional — the default implementations are no-ops. - EventLog and StreamChannel instances in the projection dict are - auto-closed / auto-failed by the mux, so most transformers don't - need `finalize` or `fail` at all. + StreamChannel instances in the projection dict are auto-closed / + auto-failed by the mux, so most transformers don't need `finalize` + or `fail` at all. Transformers that need async work pick the async lane by: @@ -170,15 +169,16 @@ class StreamTransformer(ABC): def finalize(self) -> None: """Called when the run ends normally (sync lane). - Override to close EventLogs, resolve promises, or perform other - teardown. StreamChannel instances are auto-closed by the mux. + Override to close StreamChannels, resolve promises, or perform + other teardown. StreamChannel instances in the projection dict + are auto-closed by the mux. """ async def afinalize(self) -> None: """Called when the run ends normally (async lane). By the time this runs, the mux has already awaited every task - started via `schedule()`, so EventLogs can be closed here + started via `schedule()`, so StreamChannels can be closed here without a last-task-wins race. The default delegates to `finalize`. @@ -188,8 +188,9 @@ class StreamTransformer(ABC): def fail(self, err: BaseException) -> None: """Called when the run ends with an error (sync lane). - Override to fail EventLogs, reject promises, or perform other - teardown. StreamChannel instances are auto-failed by the mux. + Override to fail StreamChannels, reject promises, or perform + other teardown. StreamChannel instances in the projection dict + are auto-failed by the mux. Args: err: The exception that ended the run. diff --git a/libs/langgraph/langgraph/stream/stream_channel.py b/libs/langgraph/langgraph/stream/stream_channel.py index 26533beb4..8f4ed821c 100644 --- a/libs/langgraph/langgraph/stream/stream_channel.py +++ b/libs/langgraph/langgraph/stream/stream_channel.py @@ -1,113 +1,327 @@ from __future__ import annotations -from collections.abc import AsyncIterator, Callable, Iterator +import asyncio +from collections import deque +from collections.abc import AsyncIterator, Awaitable, Callable, Iterator from typing import Generic, TypeVar -from langgraph.stream._event_log import EventLog - T = TypeVar("T") class StreamChannel(Generic[T]): - """A named projection channel with optional protocol auto-forwarding. + """Single-consumer drainable queue for streaming events, with optional + protocol auto-forwarding. - Wraps an event log and declares a protocol channel name. When the - StreamMux detects a StreamChannel in a transformer's `init()` - return value, it automatically wires every `push()` to inject a - `ProtocolEvent` into the main event stream using the channel's - name as the method. + When constructed with a `name`, the StreamMux auto-wires every + `push()` to also inject a `ProtocolEvent` into the main event stream + using the channel's name as the method. When constructed without a + name, the channel is local-only — items are only visible to + in-process consumers that iterate the channel directly. - Auto-forwarded events bypass the transformer pipeline — other - transformers' `process()` / `aprocess()` methods do not see - `custom:` events produced by a channel push. This prevents a - transformer that pushes to its own channel during `process()` from - re-triggering itself, but it also means filter- or tap-style - transformers cannot observe channel output from peer transformers. - Consumers that need that should iterate the main event stream. + Items are popped off the front as the consumer advances — there is + no retention beyond what's currently queued. A channel accepts + exactly one subscriber; a second `__iter__` / `__aiter__` call + raises. Use `tee(n)` / `atee(n)` for fan-out. - In-process consumers iterate the channel directly (`for item in ch` - or `async for item in ch`). Remote SDK clients subscribe via - `session.subscribe("custom:")`. + 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`. - Like EventLog, a StreamChannel starts unbound. The mux calls - `_bind(is_async)` during registration so the correct iteration - protocol is available by the time user code sees it. + 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. - Lifecycle (`_close` / `_fail`) is managed by the mux — transformers - using only StreamChannels don't need `finalize` or `fail` hooks. + Memory is bounded by caller pace: both sync and async use caller- + driven pumps, so each cursor advance produces at most one event. + + Lazy-subscribe: `push` appends to the local buffer only when a + subscriber has registered. Auto-forward via `_wire_fn` always fires + regardless of subscription state. + + Lifecycle (`close` / `fail`) is managed by the mux — transformers + don't need to close their channels manually. """ - def __init__(self, name: str, *, maxlen: int | None = None) -> None: - """Initialize the channel with an empty inner log. + def __init__(self, name: str | None = None, *, maxlen: int | None = None) -> None: + """Initialize the channel. Args: - name: The protocol channel name used for auto-forwarded - events. Surfaced on the wire as `custom:` for - user-defined transformers, or as `` for channels - owned by a native transformer (`_native = True`). The - StreamMux derives the prefix from the owning - transformer at registration time. - maxlen: Optional retention cap on the inner EventLog. See - `EventLog.__init__` for semantics. + name: Optional protocol channel name. When set, the + StreamMux wires every `push()` to also inject a + `ProtocolEvent` into the main event stream. Surfaced + on the wire as `custom:` for user-defined + transformers, or as `` for channels owned by a + native transformer (`_native = True`). When `None`, + the channel is local-only. + 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`. """ + if maxlen is not None and maxlen <= 0: + raise ValueError("StreamChannel maxlen must be a positive int or None") self.name = name - self._log: EventLog[T] = EventLog(maxlen=maxlen) + self._items: deque[T] = deque() + self._maxlen: int | None = maxlen + self._closed = False + self._error: BaseException | None = None + + self._is_async: bool | None = None + + self._subscribed = False + + self._request_more: Callable[[], bool] | None = None + self._arequest_more: Callable[[], Awaitable[bool]] | None = None + self._wire_fn: Callable[[T], None] | None = None + # ------------------------------------------------------------------ + # Binding + # ------------------------------------------------------------------ + def _bind(self, *, is_async: bool) -> None: - """Bind the underlying event log to sync or async mode. + """Bind this channel to sync or async mode. + + Called by the StreamMux after transformer registration. Must be + called exactly once before any iteration. Args: - is_async: True for async iteration, False for sync. - """ - self._log._bind(is_async=is_async) + is_async: True to enable async iteration, False for sync. - def push(self, item: T) -> None: - """Append an item to the log and auto-forward if wired. - - Args: - item: The item to push. + Raises: + RuntimeError: If the channel has already been bound. """ - self._log.push(item) - if self._wire_fn is not None: - self._wire_fn(item) + if self._is_async is not None: + raise RuntimeError("StreamChannel is already bound") + self._is_async = is_async # ------------------------------------------------------------------ - # Mux lifecycle hooks (not called by transformers directly) + # Mux wiring (not called by transformers directly) # ------------------------------------------------------------------ def _wire(self, fn: Callable[[T], None]) -> None: """Install the auto-forward callback (called by StreamMux).""" self._wire_fn = fn - def _close(self) -> None: - """Close the underlying log (called by StreamMux on run end).""" - self._log.close() + # ------------------------------------------------------------------ + # Producer API + # ------------------------------------------------------------------ - def _fail(self, err: BaseException) -> None: - """Fail the underlying log (called by StreamMux on run error).""" - self._log.fail(err) + def push(self, item: T) -> None: + """Append an item. Auto-forwards if wired. + + The local buffer append is a no-op when no subscriber is + registered, but auto-forwarding always fires so wired events + reach the main event log regardless of subscription state. + + Raises: + RuntimeError: If the channel is closed (and subscribed). + """ + if self._subscribed: + if self._closed: + raise RuntimeError("Cannot push to a closed StreamChannel") + self._items.append(item) + if self._wire_fn is not None: + self._wire_fn(item) + + def close(self) -> None: + """Mark the channel as complete.""" + self._closed = True + + def fail(self, err: BaseException) -> None: + """Mark the channel as errored. + + Args: + err: The exception to surface to the subscriber. + """ + self._error = err + self._closed = True # ------------------------------------------------------------------ - # Iteration — delegates to the inner event log (multi-cursor) + # Sync iteration (caller-driven pump) # ------------------------------------------------------------------ def __iter__(self) -> Iterator[T]: - return iter(self._log) + """Subscribe and return a sync cursor. Can be called only once. + + Raises: + TypeError: If the channel is unbound or bound to async mode. + RuntimeError: If the channel already has a subscriber. + """ + if self._is_async is None: + raise TypeError( + "StreamChannel has not been bound yet. " + "Register the transformer with a StreamMux first." + ) + if self._is_async: + raise TypeError( + "This StreamChannel is bound to async mode — use 'async for' instead." + ) + if self._subscribed: + raise RuntimeError( + "StreamChannel already has a subscriber; use .tee(n) for fan-out." + ) + self._subscribed = True + return self._sync_cursor() + + def _sync_cursor(self) -> Iterator[T]: + while True: + 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: + if not self._request_more(): + if not self._items and not self._closed: + return + else: + return + + # ------------------------------------------------------------------ + # Async iteration (caller-driven pump) + # ------------------------------------------------------------------ def __aiter__(self) -> AsyncIterator[T]: - return self._log.__aiter__() + """Subscribe and return an async cursor. Can be called only once. + + Raises: + TypeError: If the channel is unbound or bound to sync mode. + RuntimeError: If the channel already has a subscriber. + """ + if self._is_async is None: + raise TypeError( + "StreamChannel has not been bound yet. " + "Register the transformer with a StreamMux first." + ) + if not self._is_async: + raise TypeError( + "This StreamChannel is bound to sync mode — use 'for' instead." + ) + if self._subscribed: + raise RuntimeError( + "StreamChannel already has a subscriber; use .atee(n) for fan-out." + ) + self._subscribed = True + return self._async_cursor() + + async def _async_cursor(self) -> AsyncIterator[T]: + while True: + 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: + return + + # ------------------------------------------------------------------ + # Fan-out via tee + # ------------------------------------------------------------------ def tee(self, n: int = 2) -> tuple[Iterator[T], ...]: - """Fan out the channel into `n` independent sync iterators. + """Subscribe and return `n` independent sync iterators. - Delegates to the underlying EventLog's `tee()`. + 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 channel is unbound or bound to async mode. + RuntimeError: If the channel already has a subscriber. + ValueError: If `n` < 1. """ - return self._log.tee(n) + 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], ...]: - """Fan out the channel into `n` independent async iterators. + """Subscribe and return `n` independent async iterators. - Delegates to the underlying EventLog's `atee()`. + 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 channel is unbound or bound to sync mode. + RuntimeError: If the channel already has a subscriber. + ValueError: If `n` < 1. """ - return self._log.atee(n) + 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/transformers.py b/libs/langgraph/langgraph/stream/transformers.py index 90313f2e0..a72664ff1 100644 --- a/libs/langgraph/langgraph/stream/transformers.py +++ b/libs/langgraph/langgraph/stream/transformers.py @@ -13,7 +13,6 @@ from langchain_protocol.protocol import MessagesData from typing_extensions import NotRequired, TypedDict from langgraph.errors import GraphInterrupt -from langgraph.stream._event_log import EventLog from langgraph.stream._types import ProtocolEvent, StreamTransformer from langgraph.stream.run_stream import AsyncSubgraphRunStream, SubgraphRunStream from langgraph.stream.stream_channel import StreamChannel @@ -50,7 +49,7 @@ class ValuesTransformer(StreamTransformer): def __init__(self, scope: tuple[str, ...] = ()) -> None: super().__init__(scope) - self._log: EventLog[dict[str, Any]] = EventLog() + self._log: StreamChannel[dict[str, Any]] = StreamChannel() self._latest: dict[str, Any] | None = None self._interrupted = False self._interrupts: list[Any] = [] @@ -131,7 +130,7 @@ class MessagesTransformer(StreamTransformer): def __init__(self, scope: tuple[str, ...] = ()) -> None: super().__init__(scope) - self._log: EventLog[ChatModelStream] = EventLog() + self._log: StreamChannel[ChatModelStream] = StreamChannel() # Correlate protocol events back to a ChatModelStream by run_id # (attached to the event's metadata by StreamMessagesHandler). self._by_run: dict[str, ChatModelStream] = {} @@ -519,7 +518,9 @@ class SubgraphTransformer(_TasksLifecycleBase): def __init__(self, scope: tuple[str, ...] = ()) -> None: super().__init__(scope) - self._log: EventLog[SubgraphRunStream | AsyncSubgraphRunStream] = EventLog() + self._log: StreamChannel[SubgraphRunStream | AsyncSubgraphRunStream] = ( + StreamChannel() + ) self._handles: dict[ tuple[str, ...], SubgraphRunStream | AsyncSubgraphRunStream ] = {} diff --git a/libs/langgraph/tests/test_pregel_stream_v2.py b/libs/langgraph/tests/test_pregel_stream_v2.py index c76e05f1a..e4a62a6f4 100644 --- a/libs/langgraph/tests/test_pregel_stream_v2.py +++ b/libs/langgraph/tests/test_pregel_stream_v2.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 ( - EventLog, StreamChannel, StreamTransformer, ) @@ -141,13 +140,13 @@ class _CustomPassthroughTransformer(StreamTransformer): # --------------------------------------------------------------------------- -# EventLog unit tests +# StreamChannel (local, unnamed) unit tests # --------------------------------------------------------------------------- -class TestEventLog: +class TestStreamChannelLocal: def test_sync_iteration(self) -> None: - log: EventLog[int] = EventLog() + log: StreamChannel[int] = StreamChannel() log._bind(is_async=False) it = iter(log) log.push(1) @@ -157,7 +156,7 @@ class TestEventLog: assert list(it) == [1, 2, 3] def test_drain_on_consume(self) -> None: - log: EventLog[str] = EventLog() + log: StreamChannel[str] = StreamChannel() log._bind(is_async=False) it = iter(log) log.push("a") @@ -167,7 +166,7 @@ class TestEventLog: assert list(log._items) == [] def test_second_subscribe_raises(self) -> None: - log: EventLog[str] = EventLog() + log: StreamChannel[str] = StreamChannel() log._bind(is_async=False) log.close() _ = iter(log) @@ -176,7 +175,7 @@ class TestEventLog: def test_pre_subscription_push_is_noop(self) -> None: # Lazy-subscribe: pushes before subscription are dropped silently. - log: EventLog[int] = EventLog() + log: StreamChannel[int] = StreamChannel() log._bind(is_async=False) log.push(1) log.push(2) @@ -186,7 +185,7 @@ class TestEventLog: assert list(it) == [3] def test_fail_propagation(self) -> None: - log: EventLog[int] = EventLog() + log: StreamChannel[int] = StreamChannel() log._bind(is_async=False) it = iter(log) log.push(1) @@ -195,7 +194,7 @@ class TestEventLog: list(it) def test_sync_cursor_yields_items_before_error(self) -> None: - log: EventLog[int] = EventLog() + log: StreamChannel[int] = StreamChannel() log._bind(is_async=False) it = iter(log) log.push(1) @@ -209,60 +208,60 @@ class TestEventLog: assert items == [1, 2, 3] def test_push_after_close_raises(self) -> None: - log: EventLog[int] = EventLog() + log: StreamChannel[int] = StreamChannel() log._bind(is_async=False) it = iter(log) log.push(1) log.close() - with pytest.raises(RuntimeError, match="Cannot push to a closed EventLog"): + with pytest.raises(RuntimeError, match="Cannot push to a closed StreamChannel"): log.push(2) _ = list(it) def test_push_after_fail_raises(self) -> None: - log: EventLog[int] = EventLog() + log: StreamChannel[int] = StreamChannel() log._bind(is_async=False) it = iter(log) log.fail(ValueError("err")) - with pytest.raises(RuntimeError, match="Cannot push to a closed EventLog"): + with pytest.raises(RuntimeError, match="Cannot push to a closed StreamChannel"): log.push(1) with pytest.raises(ValueError, match="err"): list(it) def test_empty_log_sync(self) -> None: - log: EventLog[int] = EventLog() + log: StreamChannel[int] = StreamChannel() log._bind(is_async=False) log.close() assert list(log) == [] def test_empty_log_fail_sync(self) -> None: - log: EventLog[int] = EventLog() + log: StreamChannel[int] = StreamChannel() log._bind(is_async=False) log.fail(ValueError("empty fail")) with pytest.raises(ValueError, match="empty fail"): list(log) def test_unbound_iter_raises(self) -> None: - log: EventLog[int] = EventLog() + log: StreamChannel[int] = StreamChannel() log.close() with pytest.raises(TypeError, match="has not been bound"): list(log) def test_sync_bound_aiter_raises(self) -> None: - log: EventLog[int] = EventLog() + log: StreamChannel[int] = StreamChannel() log._bind(is_async=False) log.close() with pytest.raises(TypeError, match="bound to sync mode"): log.__aiter__() def test_double_bind_raises(self) -> None: - log: EventLog[int] = EventLog() + log: StreamChannel[int] = StreamChannel() log._bind(is_async=False) with pytest.raises(RuntimeError, match="already bound"): log._bind(is_async=True) @pytest.mark.anyio async def test_async_iteration(self) -> None: - log: EventLog[int] = EventLog() + log: StreamChannel[int] = StreamChannel() log._bind(is_async=True) cursor = aiter(log) for i in range(3): @@ -272,7 +271,7 @@ class TestEventLog: @pytest.mark.anyio async def test_async_second_subscribe_raises(self) -> None: - log: EventLog[str] = EventLog() + log: StreamChannel[str] = StreamChannel() log._bind(is_async=True) log.close() _ = log.__aiter__() @@ -281,7 +280,7 @@ class TestEventLog: @pytest.mark.anyio async def test_async_fail(self) -> None: - log: EventLog[int] = EventLog() + log: StreamChannel[int] = StreamChannel() log._bind(is_async=True) cursor = aiter(log) log.push(1) @@ -292,7 +291,7 @@ class TestEventLog: @pytest.mark.anyio async def test_async_cursor_yields_items_before_error(self) -> None: - log: EventLog[int] = EventLog() + log: StreamChannel[int] = StreamChannel() log._bind(is_async=True) cursor = aiter(log) log.push(1) @@ -307,14 +306,14 @@ class TestEventLog: @pytest.mark.anyio async def test_empty_log_async(self) -> None: - log: EventLog[int] = EventLog() + log: StreamChannel[int] = StreamChannel() log._bind(is_async=True) log.close() assert [item async for item in log] == [] @pytest.mark.anyio async def test_empty_log_fail_async(self) -> None: - log: EventLog[int] = EventLog() + log: StreamChannel[int] = StreamChannel() log._bind(is_async=True) log.fail(ValueError("empty fail")) with pytest.raises(ValueError, match="empty fail"): @@ -323,7 +322,7 @@ class TestEventLog: @pytest.mark.anyio async def test_async_bound_iter_raises(self) -> None: - log: EventLog[int] = EventLog() + log: StreamChannel[int] = StreamChannel() log._bind(is_async=True) log.close() with pytest.raises(TypeError, match="bound to async mode"): @@ -331,18 +330,18 @@ class TestEventLog: # --------------------------------------------------------------------------- -# StreamChannel unit tests +# StreamChannel (named, wired) unit tests # --------------------------------------------------------------------------- -class TestStreamChannel: +class TestStreamChannelNamed: 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() + ch.close() assert list(it) == ["a", "b"] def test_wire_callback(self) -> None: @@ -353,7 +352,7 @@ class TestStreamChannel: it = iter(ch) ch.push("x") ch.push("y") - ch._close() + ch.close() assert forwarded == ["x", "y"] assert list(it) == ["x", "y"] @@ -362,7 +361,7 @@ class TestStreamChannel: ch._bind(is_async=False) it = iter(ch) ch.push("a") - ch._fail(ValueError("channel error")) + ch.fail(ValueError("channel error")) items: list[str] = [] with pytest.raises(ValueError, match="channel error"): for item in it: @@ -375,7 +374,7 @@ class TestStreamChannel: assert ch._wire_fn is None it = iter(ch) ch.push(42) - ch._close() + ch.close() assert list(it) == [42] @pytest.mark.anyio @@ -383,9 +382,9 @@ class TestStreamChannel: ch: StreamChannel[str] = StreamChannel("test") ch._bind(is_async=True) cursor = ch.__aiter__() - ch._log.push("x") - ch._log.push("y") - ch._close() + ch.push("x") + ch.push("y") + ch.close() assert [item async for item in cursor] == ["x", "y"] @@ -922,7 +921,7 @@ class TestStreamMuxResilience: mux = StreamMux([t]) with pytest.raises(RuntimeError, match="finalize broke"): mux.close() - assert t._channel._log._closed + assert t._channel._closed # --------------------------------------------------------------------------- @@ -963,7 +962,7 @@ class TestCustomTransformer: def __init__(self, scope: tuple[str, ...] = ()) -> None: super().__init__(scope) - self._log: EventLog[str] = EventLog() + self._log: StreamChannel[str] = StreamChannel() def init(self) -> dict[str, Any]: return {"foo": self._log} @@ -1046,7 +1045,7 @@ class TestCustomTransformer: class ConflictTransformer(StreamTransformer): def __init__(self, scope: tuple[str, ...] = ()) -> None: super().__init__(scope) - self._log: EventLog[str] = EventLog() + self._log: StreamChannel[str] = StreamChannel() def init(self) -> dict[str, Any]: return {"values": self._log} @@ -1061,16 +1060,16 @@ class TestCustomTransformer: # --------------------------------------------------------------------------- -# EventLog auto-lifecycle via StreamMux +# StreamChannel auto-lifecycle via StreamMux # --------------------------------------------------------------------------- -class TestEventLogAutoLifecycle: - def test_mux_auto_closes_event_logs(self) -> None: +class TestStreamChannelAutoLifecycle: + def test_mux_auto_closes_channels(self) -> None: class SimpleTransformer(StreamTransformer): def __init__(self) -> None: super().__init__() - self._log: EventLog[str] = EventLog() + self._log: StreamChannel[str] = StreamChannel() def init(self) -> dict[str, Any]: return {"items": self._log} @@ -1085,11 +1084,11 @@ class TestEventLogAutoLifecycle: mux.close() assert len(list(it)) == 1 - def test_mux_auto_fails_event_logs(self) -> None: + def test_mux_auto_fails_channels(self) -> None: class SimpleTransformer(StreamTransformer): def __init__(self) -> None: super().__init__() - self._log: EventLog[str] = EventLog() + self._log: StreamChannel[str] = StreamChannel() def init(self) -> dict[str, Any]: return {"items": self._log} @@ -1110,7 +1109,7 @@ class TestEventLogAutoLifecycle: class ManualCloseTransformer(StreamTransformer): def __init__(self) -> None: super().__init__() - self._log: EventLog[str] = EventLog() + self._log: StreamChannel[str] = StreamChannel() def init(self) -> dict[str, Any]: return {"items": self._log} @@ -1128,7 +1127,7 @@ class TestEventLogAutoLifecycle: class MinimalTransformer(StreamTransformer): def __init__(self, scope: tuple[str, ...] = ()) -> None: super().__init__(scope) - self._log: EventLog[str] = EventLog() + self._log: StreamChannel[str] = StreamChannel() def init(self) -> dict[str, Any]: return {"minimal": self._log} @@ -1216,7 +1215,7 @@ class TestAsyncTransformerLane: requires_async = True def __init__(self) -> None: - self._log: EventLog[str] = EventLog() + self._log: StreamChannel[str] = StreamChannel() def init(self) -> dict[str, Any]: return {"out": self._log} @@ -1273,7 +1272,7 @@ class TestAsyncTransformerLane: requires_async = True def __init__(self) -> None: - self._log: EventLog[str] = EventLog() + self._log: StreamChannel[str] = StreamChannel() def init(self) -> dict[str, Any]: return {"out": self._log} @@ -1353,7 +1352,7 @@ class TestAsyncTransformerLane: requires_async = True def __init__(self) -> None: - self._log: EventLog[str] = EventLog() + self._log: StreamChannel[str] = StreamChannel() def init(self) -> dict[str, Any]: return {"seen": self._log} @@ -1381,7 +1380,7 @@ class TestAsyncTransformerLane: def __init__(self, scope: tuple[str, ...] = ()) -> None: super().__init__(scope) - self._log: EventLog[int] = EventLog() + self._log: StreamChannel[int] = StreamChannel() def init(self) -> dict[str, Any]: return {"scores": self._log} @@ -1469,20 +1468,20 @@ class TestMemoryBounds: # --------------------------------------------------------------------------- -# DrainOnConsume: EventLog capacity semantics +# DrainOnConsume: StreamChannel capacity semantics # --------------------------------------------------------------------------- class TestDrainOnConsume: def test_invalid_maxlen_raises(self) -> None: with pytest.raises(ValueError, match="positive int or None"): - EventLog(maxlen=0) + StreamChannel(maxlen=0) with pytest.raises(ValueError, match="positive int or None"): - EventLog(maxlen=-3) + StreamChannel(maxlen=-3) def test_push_unbounded_by_design(self) -> None: """Push is non-blocking; the caller-driven pump bounds memory via iteration pace.""" - log: EventLog[int] = EventLog() + log: StreamChannel[int] = StreamChannel() log._bind(is_async=False) it = iter(log) for i in range(100): @@ -1491,7 +1490,7 @@ class TestDrainOnConsume: assert list(it) == list(range(100)) def test_tee_fans_out_sync(self) -> None: - log: EventLog[int] = EventLog() + log: StreamChannel[int] = StreamChannel() log._bind(is_async=False) a, b = log.tee(2) for i in range(3): @@ -1502,7 +1501,7 @@ class TestDrainOnConsume: @pytest.mark.anyio async def test_atee_fans_out(self) -> None: - log: EventLog[int] = EventLog() + log: StreamChannel[int] = StreamChannel() log._bind(is_async=True) a, b = log.atee(2) for i in range(3): diff --git a/libs/langgraph/tests/test_stream_lifecycle_transformer.py b/libs/langgraph/tests/test_stream_lifecycle_transformer.py index df5c7f1c4..ef4b735d1 100644 --- a/libs/langgraph/tests/test_stream_lifecycle_transformer.py +++ b/libs/langgraph/tests/test_stream_lifecycle_transformer.py @@ -80,23 +80,23 @@ def _tasks_result( def _arm(mux: StreamMux) -> None: - """Force projection logs to accept pushes (skip lazy-subscribe gate). + """Force projection channels to accept pushes (skip lazy-subscribe gate). - `EventLog.push` is a no-op until a subscriber attaches. Tests that - inspect `_items` directly need the gate flipped before any event - is dispatched. + `StreamChannel.push` only appends to the local buffer when a + subscriber is attached. Tests that inspect `_items` directly need + the gate flipped before any event is dispatched. """ mux._events._subscribed = True for transformer in mux._transformers: if isinstance(transformer, LifecycleTransformer): - transformer._channel._log._subscribed = True + transformer._channel._subscribed = True def _drain_lifecycle(mux: StreamMux) -> list[LifecyclePayload]: - """Snapshot the lifecycle channel's underlying log.""" + """Snapshot the lifecycle channel's buffer.""" transformer = mux.transformer_by_key("lifecycle") assert isinstance(transformer, LifecycleTransformer) - return list(transformer._channel._log._items) + return list(transformer._channel._items) def _build_lifecycle_mux(*, scope: tuple[str, ...] = ()) -> StreamMux: diff --git a/libs/langgraph/tests/test_stream_messages_transformer.py b/libs/langgraph/tests/test_stream_messages_transformer.py index b2f4f9c65..0a767dcad 100644 --- a/libs/langgraph/tests/test_stream_messages_transformer.py +++ b/libs/langgraph/tests/test_stream_messages_transformer.py @@ -18,9 +18,9 @@ from typing_extensions import TypedDict from langgraph.constants import END, START from langgraph.graph import MessagesState, StateGraph -from langgraph.stream._event_log import EventLog from langgraph.stream._mux import StreamMux from langgraph.stream.run_stream import GraphRunStream +from langgraph.stream.stream_channel import StreamChannel from langgraph.stream.transformers import MessagesTransformer, ValuesTransformer TS = int(time.time() * 1000) @@ -90,9 +90,11 @@ def _whole_msg( } -def _make_sync_transformer() -> tuple[MessagesTransformer, EventLog[ChatModelStream]]: +def _make_sync_transformer() -> tuple[ + MessagesTransformer, StreamChannel[ChatModelStream] +]: t = MessagesTransformer() - log: EventLog[ChatModelStream] = t.init()["messages"] + log: StreamChannel[ChatModelStream] = t.init()["messages"] log._bind(is_async=False) # Subscribe up front so pushes during process() are retained. log._subscribed = True @@ -100,9 +102,11 @@ def _make_sync_transformer() -> tuple[MessagesTransformer, EventLog[ChatModelStr return t, log -def _make_async_transformer() -> tuple[MessagesTransformer, EventLog[ChatModelStream]]: +def _make_async_transformer() -> tuple[ + MessagesTransformer, StreamChannel[ChatModelStream] +]: t = MessagesTransformer() - log: EventLog[ChatModelStream] = t.init()["messages"] + log: StreamChannel[ChatModelStream] = t.init()["messages"] log._bind(is_async=True) log._subscribed = True return t, log @@ -410,7 +414,7 @@ class TestWireRequestMore: mux = StreamMux([values_t, messages_t], is_async=False) GraphRunStream(iter([]), mux, values_t) - log: EventLog[ChatModelStream] = mux.extensions["messages"] + log: StreamChannel[ChatModelStream] = mux.extensions["messages"] log._subscribed = True for evt in _lifecycle(): messages_t.process(_proto_event(evt)) @@ -427,12 +431,12 @@ class TestWireRequestMore: class TestViaMux: def _make_mux( self, - ) -> tuple[MessagesTransformer, StreamMux, EventLog[ChatModelStream]]: + ) -> tuple[MessagesTransformer, StreamMux, StreamChannel[ChatModelStream]]: t = MessagesTransformer() v = ValuesTransformer() mux = StreamMux([v, t], is_async=False) t._bind_pump(lambda: False) - log: EventLog[ChatModelStream] = mux.extensions["messages"] + log: StreamChannel[ChatModelStream] = mux.extensions["messages"] log._subscribed = True return t, mux, log @@ -456,7 +460,7 @@ class TestViaMux: t = MessagesTransformer() v = ValuesTransformer() mux = StreamMux([v, t], is_async=True) - log: EventLog[ChatModelStream] = mux.extensions["messages"] + log: StreamChannel[ChatModelStream] = mux.extensions["messages"] log._subscribed = True for evt in _lifecycle(text="async mux"): diff --git a/libs/langgraph/tests/test_stream_subgraph_transformer.py b/libs/langgraph/tests/test_stream_subgraph_transformer.py index 71b161243..ce46d0c3a 100644 --- a/libs/langgraph/tests/test_stream_subgraph_transformer.py +++ b/libs/langgraph/tests/test_stream_subgraph_transformer.py @@ -124,10 +124,8 @@ def _arm(mux: StreamMux) -> None: """ mux._events._subscribed = True for value in mux.extensions.values(): - if hasattr(value, "_subscribed"): # EventLog + if hasattr(value, "_subscribed"): value._subscribed = True - elif hasattr(value, "_log"): # StreamChannel - value._log._subscribed = True def _arm_recursive(mux: StreamMux) -> None: @@ -177,7 +175,7 @@ def _event_items(mux: StreamMux) -> list[ProtocolEvent]: def _lifecycle_payloads(mux: StreamMux) -> list[dict[str, Any]]: lifecycle_t = mux.transformer_by_key("lifecycle") assert isinstance(lifecycle_t, LifecycleTransformer) - return list(lifecycle_t._channel._log._items) + return list(lifecycle_t._channel._items) # --------------------------------------------------------------------------- diff --git a/libs/prebuilt/langgraph/prebuilt/_tool_call_stream.py b/libs/prebuilt/langgraph/prebuilt/_tool_call_stream.py index 2df12e823..da2e1c737 100644 --- a/libs/prebuilt/langgraph/prebuilt/_tool_call_stream.py +++ b/libs/prebuilt/langgraph/prebuilt/_tool_call_stream.py @@ -11,7 +11,7 @@ from __future__ import annotations from collections.abc import AsyncIterator, Iterator from typing import Any -from langgraph.stream._event_log import EventLog +from langgraph.stream.stream_channel import StreamChannel class ToolCallStream: @@ -21,7 +21,7 @@ class ToolCallStream: are populated as events arrive: - `tool_call_id`, `tool_name`, `input`: stable from the start event. - - `output_deltas`: an `EventLog` of delta chunks. Iterate (sync or + - `output_deltas`: a `StreamChannel` of delta chunks. Iterate (sync or async) to consume partial output in arrival order. - `output`: terminal payload from `tool-finished`, or `None` if the call failed or is still in flight. @@ -51,14 +51,14 @@ class ToolCallStream: self.tool_call_id = tool_call_id self.tool_name = tool_name self.input = input - self._output_deltas: EventLog[Any] = EventLog() + self._output_deltas: StreamChannel[Any] = StreamChannel() self.output: Any = None self.error: str | None = None self.completed = False @property - def output_deltas(self) -> EventLog[Any]: - """The EventLog of streamed `tool-output-delta` payloads. + def output_deltas(self) -> StreamChannel[Any]: + """The channel of streamed `tool-output-delta` payloads. Iterate (sync or async depending on how the run was started) to consume partial output in arrival order. The log closes when diff --git a/libs/prebuilt/langgraph/prebuilt/_tool_call_transformer.py b/libs/prebuilt/langgraph/prebuilt/_tool_call_transformer.py index 949a473cc..6f3213210 100644 --- a/libs/prebuilt/langgraph/prebuilt/_tool_call_transformer.py +++ b/libs/prebuilt/langgraph/prebuilt/_tool_call_transformer.py @@ -5,8 +5,8 @@ from __future__ import annotations from collections.abc import Awaitable, Callable from typing import Any -from langgraph.stream._event_log import EventLog from langgraph.stream._types import ProtocolEvent, StreamTransformer +from langgraph.stream.stream_channel import StreamChannel from langgraph.prebuilt._tool_call_stream import ToolCallStream @@ -21,11 +21,12 @@ class ToolCallTransformer(StreamTransformer): Native transformer — the `tool_calls` projection is exposed as a direct attribute on the run stream. - `EventLog[ToolCallStream]` is used (not `StreamChannel`) because the - live handles are not serializable and should not be auto-forwarded - onto the main event log. Wire consumers subscribe to the `tools` - channel instead, where the raw protocol events flow through - untouched by this transformer (`process` returns `True`). + A nameless `StreamChannel[ToolCallStream]` is used (no protocol + auto-forwarding) because the live handles are not serializable and + should not be injected into the main event log. Wire consumers + subscribe to the `tools` channel instead, where the raw protocol + events flow through untouched by this transformer (`process` + returns `True`). Registered explicitly by users at compile time via `builder.compile(transformers=[ToolCallTransformer])` — not a @@ -37,7 +38,7 @@ class ToolCallTransformer(StreamTransformer): def __init__(self, scope: tuple[str, ...] = ()) -> None: super().__init__(scope) - self._log: EventLog[ToolCallStream] = EventLog() + self._log: StreamChannel[ToolCallStream] = StreamChannel() self._active: dict[str, ToolCallStream] = {} self._is_async = False self._pump_fn: Callable[[], bool] | None = None diff --git a/libs/prebuilt/tests/test_tool_call_transformer.py b/libs/prebuilt/tests/test_tool_call_transformer.py index b7c88d65f..df843dd44 100644 --- a/libs/prebuilt/tests/test_tool_call_transformer.py +++ b/libs/prebuilt/tests/test_tool_call_transformer.py @@ -11,9 +11,9 @@ from langchain_core.tools import tool from langgraph.constants import END, START from langgraph.graph import StateGraph from langgraph.graph.message import add_messages -from langgraph.stream._event_log import EventLog from langgraph.stream._mux import StreamMux from langgraph.stream._types import ProtocolEvent +from langgraph.stream.stream_channel import StreamChannel from langgraph.stream.transformers import ( MessagesTransformer, ValuesTransformer, @@ -63,7 +63,7 @@ def _tool_event( } -def _subscribe(log: EventLog) -> None: +def _subscribe(log: StreamChannel) -> None: log._subscribed = True