diff --git a/libs/langgraph/langgraph/stream/__init__.py b/libs/langgraph/langgraph/stream/__init__.py index d882191b5..56db44c03 100644 --- a/libs/langgraph/langgraph/stream/__init__.py +++ b/libs/langgraph/langgraph/stream/__init__.py @@ -1,8 +1,7 @@ """Stream protocol types and infrastructure for LangGraph.""" from langgraph.stream._convert import STREAM_V2_MODES, convert_to_protocol_event -from langgraph.stream._event_log import EventLog -from langgraph.stream._mux import StreamMux +from langgraph.stream._mux import AsyncStreamMux, StreamMux from langgraph.stream._types import ( InterruptPayload, ProtocolEvent, @@ -27,9 +26,9 @@ __all__ = [ "STREAM_V2_MODES", "AsyncChatModelStream", "AsyncGraphRunStream", + "AsyncStreamMux", "AsyncSubgraphRunStream", "ChatModelStream", - "EventLog", "GraphRunStream", "InterruptPayload", "MessagesTransformer", diff --git a/libs/langgraph/langgraph/stream/_convert.py b/libs/langgraph/langgraph/stream/_convert.py index 6d199b789..7eb6ca8f1 100644 --- a/libs/langgraph/langgraph/stream/_convert.py +++ b/libs/langgraph/langgraph/stream/_convert.py @@ -7,7 +7,6 @@ original payload. from __future__ import annotations -import time from typing import Any from langgraph.stream._types import ProtocolEvent, _ProtocolEventParams @@ -58,7 +57,6 @@ def convert_to_protocol_event( params: _ProtocolEventParams = { "namespace": list(ns), - "timestamp": _now_ms(), "data": payload, } if node is not None: @@ -71,9 +69,4 @@ def convert_to_protocol_event( ) -def _now_ms() -> int: - """Current time in milliseconds since epoch.""" - return int(time.time() * 1000) - - __all__ = ["STREAM_V2_MODES", "convert_to_protocol_event"] diff --git a/libs/langgraph/langgraph/stream/_event_log.py b/libs/langgraph/langgraph/stream/_event_log.py deleted file mode 100644 index 660a05531..000000000 --- a/libs/langgraph/langgraph/stream/_event_log.py +++ /dev/null @@ -1,132 +0,0 @@ -"""Replayable append-only event buffer for StreamingHandler. - -``EventLog`` stores protocol events in an ordered list and supports -multiple independent async iterators, each with their own cursor -offset. Subscribers that join mid-stream replay from a given offset -without losing earlier events. -""" - -from __future__ import annotations - -import asyncio -import threading -from typing import Generic, TypeVar - -T = TypeVar("T") - - -class EventLog(Generic[T]): - """Append-only event buffer with cursor-based async iteration. - - Multiple consumers can subscribe independently and each will see - every event from their starting offset onward. - """ - - __slots__ = ("_items", "_closed", "_error", "_waiters", "_lock") - - def __init__(self) -> None: - self._items: list[T] = [] - self._closed = False - self._error: BaseException | None = None - self._waiters: list[asyncio.Future[None]] = [] - self._lock = threading.Lock() - - # -- Producer API ------------------------------------------------------- - - def append(self, item: T) -> None: - """Append an event and wake all waiting consumers.""" - with self._lock: - if self._closed: - raise RuntimeError("EventLog is closed") - self._items.append(item) - self._wake_all() - - def close(self) -> None: - """Mark the log as complete. Iterators will end gracefully.""" - with self._lock: - self._closed = True - self._wake_all() - - def fail(self, error: BaseException) -> None: - """Mark the log as failed. Iterators will raise *error*.""" - with self._lock: - self._error = error - self._closed = True - self._wake_all() - - # -- Consumer API ------------------------------------------------------- - - def subscribe(self, offset: int = 0) -> _Cursor[T]: - """Create a new cursor starting at *offset*. - - If *offset* is 0 the cursor replays the entire log. If - *offset* equals ``len(self)`` the cursor starts from the - current tip and only sees future events. - """ - return _Cursor(self, offset) - - # -- Inspection --------------------------------------------------------- - - def __len__(self) -> int: - return len(self._items) - - def __getitem__(self, index: int) -> T: - return self._items[index] - - @property - def closed(self) -> bool: - return self._closed - - # -- Internal ----------------------------------------------------------- - - def _wake_all(self) -> None: - for fut in self._waiters: - if not fut.done(): - try: - fut.get_loop().call_soon_threadsafe(fut.set_result, None) - except RuntimeError: - # Loop already closed — ignore. - pass - self._waiters.clear() - - -class _Cursor(Generic[T]): - """An independent async iterator over an :class:`EventLog`.""" - - __slots__ = ("_log", "_offset") - - def __init__(self, log: EventLog[T], offset: int) -> None: - self._log = log - self._offset = offset - - def __aiter__(self) -> _Cursor[T]: - return self - - async def __anext__(self) -> T: - while True: - with self._log._lock: - if self._offset < len(self._log._items): - item = self._log._items[self._offset] - self._offset += 1 - return item - if self._log._error is not None: - raise self._log._error - if self._log._closed: - raise StopAsyncIteration - # Nothing available yet — register a waiter - loop = asyncio.get_running_loop() - fut: asyncio.Future[None] = loop.create_future() - self._log._waiters.append(fut) - # Wait outside the lock - try: - await fut - except asyncio.CancelledError: - # Remove our future so it doesn't accumulate in the list. - try: - self._log._waiters.remove(fut) - except ValueError: - pass # Already removed by _wake_all - raise - - -__all__ = ["EventLog"] diff --git a/libs/langgraph/langgraph/stream/_mux.py b/libs/langgraph/langgraph/stream/_mux.py index 83082d143..28f559e19 100644 --- a/libs/langgraph/langgraph/stream/_mux.py +++ b/libs/langgraph/langgraph/stream/_mux.py @@ -1,33 +1,36 @@ """Central event dispatcher with transformer pipeline for StreamingHandler. -``StreamMux`` is the core coordination point: it holds the main -:class:`EventLog`, tracks discovered namespaces for subgraph stream +``StreamMux`` is the synchronous core: it holds the main +event log (a plain list), tracks discovered namespaces for subgraph stream creation, and pipes every event through the registered :class:`StreamTransformer` pipeline before appending it to the log. + +``AsyncStreamMux`` extends ``StreamMux`` with async consumer APIs +(output futures, async event subscriptions, subgraph discovery). """ from __future__ import annotations import asyncio -import time from collections.abc import AsyncIterator from typing import Any -from langgraph.stream._event_log import EventLog from langgraph.stream._types import InterruptPayload, ProtocolEvent, StreamTransformer from langgraph.stream.stream_channel import StreamChannel, is_stream_channel class StreamMux: - """Central event dispatcher for the StreamingHandler infrastructure. + """Synchronous event dispatcher for the StreamingHandler infrastructure. The mux owns the main event log, applies the transformer pipeline to - every incoming event, and provides subscription endpoints for - filtered event iteration and subgraph discovery. + every incoming event, and tracks namespace discovery and latest values. + + For async consumer APIs (output futures, async event subscriptions, + subgraph discovery) use :class:`AsyncStreamMux`. """ def __init__(self, transformers: list[StreamTransformer] | None = None) -> None: - self._event_log: EventLog[ProtocolEvent] = EventLog() + self._event_log: list[ProtocolEvent] = [] self._transformers: list[StreamTransformer] = list(transformers or []) self._channels: list[StreamChannel[Any]] = [] self._current_namespace: list[str] = [] @@ -35,8 +38,6 @@ class StreamMux: # Namespace discovery: maps top-level ns segment → True self._discovered_ns: dict[str, bool] = {} - # Waiters for new namespace discovery - self._ns_waiters: list[asyncio.Future[None]] = [] # Latest values per namespace (list-of-strings key) self._latest_values: dict[str, Any] = {} @@ -45,9 +46,6 @@ class StreamMux: self._interrupts: list[InterruptPayload] = [] self._interrupted = False - # Output promise tracking - self._output_futures: dict[str, asyncio.Future[Any]] = {} - # Closed state self._closed = False self._error: BaseException | None = None @@ -76,7 +74,6 @@ class StreamMux: top_segment = ns[0] if top_segment not in self._discovered_ns: self._discovered_ns[top_segment] = True - self._wake_ns_waiters() # Track values if event["method"] == "values": @@ -113,23 +110,162 @@ class StreamMux: self._event_log.append(event) def close(self, output: Any = None) -> None: - """Close the mux, resolving all output futures.""" + """Close the mux and finalize all transformers.""" if self._closed: return self._closed = True - # Finalize transformers (optional method) + # Finalize transformers for transformer in self._transformers: - if hasattr(transformer, "finalize"): - transformer.finalize() + transformer.finalize() # Close wired channels for channel in self._channels: channel._close() - # Close the event log - self._event_log.close() + def fail(self, error: BaseException) -> None: + """Fail the mux and propagate the error to all consumers.""" + if self._closed: + return + self._closed = True + self._error = error + # Fail transformers + for transformer in self._transformers: + transformer.fail(error) + + # Fail wired channels + for channel in self._channels: + channel._fail(error) + + # -- Inspection --------------------------------------------------------- + + @property + def interrupted(self) -> bool: + return self._interrupted + + @property + def interrupts(self) -> list[InterruptPayload]: + return list(self._interrupts) + + @property + def event_log(self) -> list[ProtocolEvent]: + return self._event_log + + def get_latest_values(self, ns: list[str] | None = None) -> Any: + """Return the most recent values for a namespace.""" + return self._latest_values.get(_ns_key(ns or [])) + + # -- Internal ----------------------------------------------------------- + + def register_transformer(self, transformer: StreamTransformer) -> None: + """Register a new transformer and replay all buffered events through it. + + This is the safe way to add a late-arriving transformer after the mux + has already started processing events. The sequence is: + + 1. Snapshot the current log length. + 2. Append the transformer so future ``push()`` calls reach it. + 3. Replay events ``[0, snapshot)`` through the transformer. + 4. If the mux is already closed, call ``finalize()`` immediately so + the transformer's log/channel terminates cleanly. + + No namespace filtering is applied — all buffered events are + replayed. Transformers that need namespace filtering should do + so inside their ``process()`` implementation. + """ + snapshot = len(self._event_log) + self._transformers.append(transformer) + for i in range(snapshot): + transformer.process(self._event_log[i]) + if self._closed: + transformer.finalize() + + def wire_channels(self, projection: Any) -> None: + """Scan *projection* for :class:`StreamChannel` instances and wire them. + + For each ``StreamChannel`` found, registers a push callback that + appends a :class:`ProtocolEvent` directly to the main event log + with ``method`` set to the channel's name. + + Channel events bypass the transformer pipeline (matching the JS + implementation). They are visible to raw event iteration and + remote SDK clients but not to other transformers' ``process()``. + """ + if projection is None: + return + items: dict[str, Any] = {} + if isinstance(projection, dict): + items = projection + elif hasattr(projection, "__dict__"): + items = vars(projection) + for _key, value in items.items(): + if is_stream_channel(value): + channel: StreamChannel[Any] = value + self._channels.append(channel) + + def _make_forwarder(ch: StreamChannel[Any]) -> Any: + def _forward(item: Any) -> None: + if self._closed: + return + # Append directly to the event log, bypassing + # the transformer pipeline. This matches the JS + # implementation and avoids re-entrancy bugs + # (namespace clobbering, infinite recursion). + self._event_log.append( + ProtocolEvent( + type="event", + seq=self._next_emit_seq, + method=ch.channel_name, + params={ + "namespace": list(self._current_namespace), + "data": item, + }, + ) + ) + self._next_emit_seq += 1 + + return _forward + + channel._wire(_make_forwarder(channel)) + + +# --------------------------------------------------------------------------- +# AsyncStreamMux — async consumer APIs on top of the sync core +# --------------------------------------------------------------------------- + + +class AsyncStreamMux(StreamMux): + """Async extension of :class:`StreamMux`. + + Adds output futures, async event subscriptions, and subgraph + discovery on top of the synchronous producer/transformer core. + """ + + def __init__(self, transformers: list[StreamTransformer] | None = None) -> None: + super().__init__(transformers=transformers) + # Notification event — set on every push/close/fail to wake async consumers + self._notify: asyncio.Event = asyncio.Event() + # Waiters for new namespace discovery + self._ns_waiters: list[asyncio.Future[None]] = [] + # Output promise tracking + self._output_futures: dict[str, asyncio.Future[Any]] = {} + + # -- Producer overrides (extend to resolve async primitives) ------------- + + def push(self, event: ProtocolEvent) -> None: + # Peek at namespace before super().push() so we can detect new + # discoveries and wake waiters. + ns = event["params"].get("namespace", []) + is_new_ns = bool(ns) and ns[0] not in self._discovered_ns + super().push(event) + if is_new_ns and ns[0] in self._discovered_ns: + self._wake_ns_waiters() + self._notify.set() + + def close(self, output: Any = None) -> None: + super().close(output) + self._notify.set() # Resolve output futures for ns_key, fut in self._output_futures.items(): if not fut.done(): @@ -138,29 +274,12 @@ class StreamMux: fut.get_loop().call_soon_threadsafe(fut.set_result, value) except RuntimeError: pass - # Wake namespace waiters self._wake_ns_waiters() def fail(self, error: BaseException) -> None: - """Fail the mux, rejecting all output futures.""" - if self._closed: - return - self._closed = True - self._error = error - - # Fail transformers (optional method) - for transformer in self._transformers: - if hasattr(transformer, "fail"): - transformer.fail(error) - - # Fail wired channels - for channel in self._channels: - channel._fail(error) - - # Fail the event log - self._event_log.fail(error) - + super().fail(error) + self._notify.set() # Reject output futures for fut in self._output_futures.values(): if not fut.done(): @@ -168,25 +287,38 @@ class StreamMux: fut.get_loop().call_soon_threadsafe(fut.set_exception, error) except RuntimeError: pass - # Wake namespace waiters self._wake_ns_waiters() - # -- Consumer API ------------------------------------------------------- + # -- Async consumer API ------------------------------------------------- - def subscribe_events( + async def subscribe_events( self, path: list[str] | None = None, offset: int = 0 ) -> AsyncIterator[ProtocolEvent]: - """Return an async iterator over events matching *path*. + """Async iterate over events matching *path*. If *path* is ``None`` or empty, all events are yielded. Otherwise, only events whose namespace starts with *path* are yielded. + + Uses the list + ``asyncio.Event`` notification pattern: poll + the event log, yield what's new, await the notify event for more. """ - cursor = self._event_log.subscribe(offset) - if not path: - return cursor - return _FilteredEventIterator(cursor, path) + cursor = offset + while True: + while cursor < len(self._event_log): + event = self._event_log[cursor] + cursor += 1 + if not path or _ns_starts_with( + event["params"].get("namespace", []), path + ): + yield event + if self._closed: + if self._error is not None: + raise self._error + return + self._notify.clear() + await self._notify.wait() async def subscribe_subgraphs( self, path: list[str] | None = None, offset: int = 0 @@ -238,101 +370,8 @@ class StreamMux: return self._output_futures[ns_key] - # -- Inspection --------------------------------------------------------- - - @property - def interrupted(self) -> bool: - return self._interrupted - - @property - def interrupts(self) -> list[InterruptPayload]: - return list(self._interrupts) - - @property - def event_log(self) -> EventLog[ProtocolEvent]: - return self._event_log - - def get_latest_values(self, ns: list[str] | None = None) -> Any: - """Return the most recent values for a namespace.""" - return self._latest_values.get(_ns_key(ns or [])) - # -- Internal ----------------------------------------------------------- - def register_transformer(self, transformer: StreamTransformer) -> None: - """Register a new transformer and replay all buffered events through it. - - This is the safe way to add a late-arriving transformer after the mux - has already started processing events. The sequence is: - - 1. Snapshot the current log length (no await → no gap possible in - asyncio's cooperative threading model). - 2. Append the transformer so future ``push()`` calls reach it. - 3. Replay events ``[0, snapshot)`` through the transformer. - 4. If the mux is already closed, call ``finalize()`` immediately so - the transformer's log/channel terminates cleanly. - - ``process()`` is only called for events whose namespace starts with - any prefix — callers that need namespace filtering should do so inside - their ``process()`` implementation, or wrap this call with their own - filtering logic. - """ - snapshot = len(self._event_log) - self._transformers.append(transformer) - for i in range(snapshot): - transformer.process(self._event_log[i]) - if self._closed: - if hasattr(transformer, "finalize"): - transformer.finalize() - - def wire_channels(self, projection: Any) -> None: - """Scan *projection* for :class:`StreamChannel` instances and wire them. - - For each ``StreamChannel`` found, registers a push callback that - appends a :class:`ProtocolEvent` directly to the main event log - with ``method`` set to the channel's name. - - Channel events bypass the transformer pipeline (matching the JS - implementation). They are visible to raw event iteration and - remote SDK clients but not to other transformers' ``process()``. - """ - if projection is None: - return - items: dict[str, Any] = {} - if isinstance(projection, dict): - items = projection - elif hasattr(projection, "__dict__"): - items = vars(projection) - for _key, value in items.items(): - if is_stream_channel(value): - channel: StreamChannel[Any] = value - self._channels.append(channel) - - def _make_forwarder(ch: StreamChannel[Any]) -> Any: - def _forward(item: Any) -> None: - if self._closed: - return - # Append directly to the event log, bypassing - # the transformer pipeline. This matches the JS - # implementation and avoids re-entrancy bugs - # (namespace clobbering, infinite recursion). - self._event_log.append( - ProtocolEvent( - type="event", - seq=self._next_emit_seq, - method=ch.channel_name, - params={ - "namespace": list(self._current_namespace), - "timestamp": int(time.time() * 1000), - "data": item, - }, - ) - ) - self._next_emit_seq += 1 - - return _forward - - channel._wire(_make_forwarder(channel)) - def _wake_ns_waiters(self) -> None: for fut in self._ns_waiters: if not fut.done(): @@ -343,25 +382,6 @@ class StreamMux: self._ns_waiters.clear() -class _FilteredEventIterator: - """Async iterator that filters events by namespace prefix.""" - - __slots__ = ("_cursor", "_path") - - def __init__(self, cursor: AsyncIterator[ProtocolEvent], path: list[str]) -> None: - self._cursor = cursor - self._path = path - - def __aiter__(self) -> _FilteredEventIterator: - return self - - async def __anext__(self) -> ProtocolEvent: - while True: - event = await self._cursor.__anext__() - ns = event["params"].get("namespace", []) - if _ns_starts_with(ns, self._path): - return event - def _ns_key(ns: list[str] | tuple[str, ...]) -> str: """Convert a namespace list to a hashable key.""" @@ -375,4 +395,4 @@ def _ns_starts_with(ns: list[str], prefix: list[str]) -> bool: return ns[: len(prefix)] == prefix -__all__ = ["StreamMux"] +__all__ = ["AsyncStreamMux", "StreamMux"] diff --git a/libs/langgraph/langgraph/stream/_types.py b/libs/langgraph/langgraph/stream/_types.py index f4a29e227..610a9374f 100644 --- a/libs/langgraph/langgraph/stream/_types.py +++ b/libs/langgraph/langgraph/stream/_types.py @@ -6,7 +6,8 @@ in-process-only types needed by the LangGraph streaming infrastructure. from __future__ import annotations -from typing import Any, Protocol, runtime_checkable +from abc import ABC, abstractmethod +from typing import Any # --------------------------------------------------------------------------- # Re-exports from langchain-protocol (CDDL-derived) @@ -55,7 +56,6 @@ class _ProtocolEventParams(TypedDict): """Payload envelope for a :class:`ProtocolEvent`.""" namespace: Namespace - timestamp: int node: NotRequired[str] data: Any @@ -74,8 +74,7 @@ class ProtocolEvent(TypedDict): params: _ProtocolEventParams -@runtime_checkable -class StreamTransformer(Protocol): +class StreamTransformer(ABC): """Extension point for custom stream projections. Implementations are registered with ``StreamingHandler`` and receive every @@ -87,6 +86,7 @@ class StreamTransformer(Protocol): """ + @abstractmethod def init(self) -> Any: """Return the initial projection value. @@ -96,6 +96,7 @@ class StreamTransformer(Protocol): """ ... + @abstractmethod def process(self, event: ProtocolEvent) -> bool: """Process an event. @@ -110,7 +111,6 @@ class StreamTransformer(Protocol): Optional — the mux auto-closes any :class:`StreamChannel` instances, so transformers that only use channels can omit this. """ - ... def fail(self, err: BaseException) -> None: """Called once when the run fails. @@ -118,7 +118,6 @@ class StreamTransformer(Protocol): Optional — the mux auto-fails any :class:`StreamChannel` instances, so transformers that only use channels can omit this. """ - ... class InterruptPayload(TypedDict): diff --git a/libs/langgraph/langgraph/stream/chat_model_stream.py b/libs/langgraph/langgraph/stream/chat_model_stream.py index 1905696bc..dfff09989 100644 --- a/libs/langgraph/langgraph/stream/chat_model_stream.py +++ b/libs/langgraph/langgraph/stream/chat_model_stream.py @@ -130,22 +130,22 @@ class ChatModelStream: # --------------------------------------------------------------------------- -# Async dual-projection helpers +# Dual-projection helpers — sync data container + async notification layer # --------------------------------------------------------------------------- class _DualProjection: - """Async iterable of deltas that is also awaitable for the final value. + """Sync data container for incremental deltas and a final value. - When iterated, yields delta values (e.g. text fragments) as they arrive. - When awaited, returns the accumulated final value (e.g. full text string). + Stores deltas as they arrive and tracks the final accumulated value. + No async primitives — see :class:`_AsyncDualProjection` for the + async-iterable + awaitable extension. """ def __init__(self) -> None: self._deltas: list[Any] = [] self._done = False self._error: BaseException | None = None - self._waiters: list[asyncio.Future[None]] = [] self._final_value: Any = None self._final_set = False @@ -154,33 +154,48 @@ class _DualProjection: def _push(self, delta: Any) -> None: """Add a new delta value.""" self._deltas.append(delta) - self._wake() def _finish(self, accumulated: Any) -> None: """Set the final accumulated value and mark as done.""" self._final_value = accumulated self._final_set = True self._done = True - self._wake() def _fail(self, error: BaseException) -> None: self._error = error self._done = True - self._wake() - def _wake(self) -> None: - for fut in self._waiters: - if not fut.done(): - try: - fut.get_loop().call_soon_threadsafe(fut.set_result, None) - except RuntimeError: - pass - self._waiters.clear() + +class _AsyncDualProjection(_DualProjection): + """Async extension of :class:`_DualProjection`. + + Async iterable of deltas that is also awaitable for the final value. + Uses an ``asyncio.Event`` to notify async consumers when new data + arrives — the same pattern as :class:`AsyncStreamMux`. + """ + + def __init__(self) -> None: + super().__init__() + self._notify: asyncio.Event = asyncio.Event() + + # -- Producer overrides (extend to notify) ----------------------------- + + def _push(self, delta: Any) -> None: + super()._push(delta) + self._notify.set() + + def _finish(self, accumulated: Any) -> None: + super()._finish(accumulated) + self._notify.set() + + def _fail(self, error: BaseException) -> None: + super()._fail(error) + self._notify.set() # -- Async iterable (yields deltas) ------------------------------------ - def __aiter__(self) -> _DualProjectionIterator: - return _DualProjectionIterator(self) + def __aiter__(self) -> _AsyncDualProjectionIterator: + return _AsyncDualProjectionIterator(self) # -- Awaitable (returns final value) ----------------------------------- @@ -191,25 +206,23 @@ class _DualProjection: while not self._final_set: if self._error is not None: raise self._error - loop = asyncio.get_running_loop() - fut: asyncio.Future[None] = loop.create_future() - self._waiters.append(fut) - await fut + self._notify.clear() + await self._notify.wait() if self._error is not None: raise self._error return self._final_value -class _DualProjectionIterator: - """Async iterator over a :class:`_DualProjection`'s deltas.""" +class _AsyncDualProjectionIterator: + """Async iterator over an :class:`_AsyncDualProjection`'s deltas.""" __slots__ = ("_proj", "_offset") - def __init__(self, proj: _DualProjection) -> None: + def __init__(self, proj: _AsyncDualProjection) -> None: self._proj = proj self._offset = 0 - def __aiter__(self) -> _DualProjectionIterator: + def __aiter__(self) -> _AsyncDualProjectionIterator: return self async def __anext__(self) -> Any: @@ -222,10 +235,8 @@ class _DualProjectionIterator: raise self._proj._error if self._proj._done: raise StopAsyncIteration - loop = asyncio.get_running_loop() - fut: asyncio.Future[None] = loop.create_future() - self._proj._waiters.append(fut) - await fut + self._proj._notify.clear() + await self._proj._notify.wait() # --------------------------------------------------------------------------- @@ -257,24 +268,24 @@ class AsyncChatModelStream(ChatModelStream): message_id: str | None = None, ) -> None: super().__init__(namespace=namespace, node=node, message_id=message_id) - self._text_proj = _DualProjection() - self._reasoning_proj = _DualProjection() - self._usage_proj = _DualProjection() + self._text_proj = _AsyncDualProjection() + self._reasoning_proj = _AsyncDualProjection() + self._usage_proj = _AsyncDualProjection() # -- Public projections (override sync properties) --------------------- @property - def text(self) -> _DualProjection: + def text(self) -> _AsyncDualProjection: """Text content — async iterable of deltas, awaitable for full text.""" return self._text_proj @property - def reasoning(self) -> _DualProjection: + def reasoning(self) -> _AsyncDualProjection: """Reasoning content — async iterable of deltas, awaitable for full text.""" return self._reasoning_proj @property - def usage(self) -> _DualProjection: + def usage(self) -> _AsyncDualProjection: """Usage info — awaitable for :class:`UsageInfo`.""" return self._usage_proj diff --git a/libs/langgraph/langgraph/stream/run_stream.py b/libs/langgraph/langgraph/stream/run_stream.py index 8044f1238..2e171b06b 100644 --- a/libs/langgraph/langgraph/stream/run_stream.py +++ b/libs/langgraph/langgraph/stream/run_stream.py @@ -2,9 +2,13 @@ These are the top-level objects returned by ``StreamingHandler.stream()`` / ``StreamingHandler.astream()``. -They wrap a :class:`StreamMux` and expose named -projections (``.values``, ``.messages``, ``.subgraphs``, ``.output``) -for ergonomic consumption. + +``AsyncGraphRunStream`` wraps an :class:`AsyncStreamMux` and exposes +``.values``, ``.messages``, ``.subgraphs``, ``.output``, and +``.messages_from()``. + +``GraphRunStream`` wraps a :class:`StreamMux` and exposes the sync +equivalents: ``.values``, ``.messages``, and ``.output``. """ from __future__ import annotations @@ -14,8 +18,7 @@ from collections.abc import AsyncIterator, Callable, Iterator from typing import Any from langgraph.stream._convert import convert_to_protocol_event -from langgraph.stream._event_log import EventLog -from langgraph.stream._mux import StreamMux +from langgraph.stream._mux import AsyncStreamMux, StreamMux from langgraph.stream._types import InterruptPayload, ProtocolEvent, StreamTransformer from langgraph.stream.chat_model_stream import AsyncChatModelStream, ChatModelStream from langgraph.stream.transformers import MessagesTransformer, ValuesTransformer @@ -30,7 +33,7 @@ class _ValuesProjection: def __init__( self, - mux: StreamMux, + mux: AsyncStreamMux, values_transformer: ValuesTransformer, ns: list[str], mapper: Callable[[Any], Any] | None = None, @@ -40,8 +43,23 @@ class _ValuesProjection: self._ns = ns self._mapper = mapper - def __aiter__(self) -> AsyncIterator[Any]: - return _ValuesIterator(self._values_transformer, self._ns, self._mapper) + async def __aiter__(self) -> AsyncIterator[Any]: + log = self._values_transformer.values_log + cursor = 0 + while True: + while cursor < len(log): + item = log[cursor] + cursor += 1 + if item.get("namespace", []) == self._ns: + data = item["data"] + if data is not None and self._mapper is not None: + yield self._mapper(data) + else: + yield data + if self._mux._closed: + return + self._mux._notify.clear() + await self._mux._notify.wait() def __await__(self) -> Any: return self._await_impl().__await__() @@ -53,33 +71,6 @@ class _ValuesProjection: return value -class _ValuesIterator: - """Filters the values log to events matching a namespace.""" - - def __init__( - self, - transformer: ValuesTransformer, - ns: list[str], - mapper: Callable[[Any], Any] | None = None, - ) -> None: - self._cursor = transformer.values_log.subscribe(0) - self._ns = ns - self._mapper = mapper - - def __aiter__(self) -> _ValuesIterator: - return self - - async def __anext__(self) -> Any: - while True: - item = await self._cursor.__anext__() - item_ns = item.get("namespace", []) - if item_ns == self._ns: - data = item["data"] - if data is not None and self._mapper is not None: - return self._mapper(data) - return data - - # --------------------------------------------------------------------------- # Messages projection # --------------------------------------------------------------------------- @@ -88,11 +79,21 @@ class _ValuesIterator: class _MessagesProjection: """Async iterable of :class:`AsyncChatModelStream` instances.""" - def __init__(self, messages_transformer: MessagesTransformer) -> None: + def __init__(self, mux: AsyncStreamMux, messages_transformer: MessagesTransformer) -> None: + self._mux = mux self._transformer = messages_transformer - def __aiter__(self) -> AsyncIterator[AsyncChatModelStream]: - return self._transformer.messages_log.subscribe(0) + async def __aiter__(self) -> AsyncIterator[AsyncChatModelStream]: + log = self._transformer.messages_log + cursor = 0 + while True: + while cursor < len(log): + yield log[cursor] + cursor += 1 + if self._mux._closed: + return + self._mux._notify.clear() + await self._mux._notify.wait() # --------------------------------------------------------------------------- @@ -103,7 +104,7 @@ class _MessagesProjection: class _SubgraphsProjection: """Async iterable yielding :class:`AsyncSubgraphRunStream` for each discovered subgraph.""" - def __init__(self, mux: StreamMux, ns: list[str]) -> None: + def __init__(self, mux: AsyncStreamMux, ns: list[str]) -> None: self._mux = mux self._ns = ns @@ -143,7 +144,7 @@ class AsyncGraphRunStream: def __init__( self, *, - mux: StreamMux, + mux: AsyncStreamMux, namespace: list[str] | None = None, transformers: list[StreamTransformer], abort_event: asyncio.Event | None = None, @@ -186,7 +187,7 @@ class AsyncGraphRunStream: def messages(self) -> _MessagesProjection: """Async iterable of :class:`AsyncChatModelStream` instances.""" t = self._find_transformer("messages") - return _MessagesProjection(t) + return _MessagesProjection(self._mux, t) def messages_from(self, node: str) -> _MessagesProjection: """Async iterable of messages from a specific node.""" @@ -196,7 +197,7 @@ class AsyncGraphRunStream: stream_cls=AsyncChatModelStream, ) self._mux.register_transformer(filtered) - return _MessagesProjection(filtered) + return _MessagesProjection(self._mux, filtered) @property def subgraphs(self) -> _SubgraphsProjection: @@ -308,7 +309,7 @@ async def create_async_graph_run_stream( if projection is not None: projections.append(projection) - mux = StreamMux(transformers=all_transformers) + mux = AsyncStreamMux(transformers=all_transformers) # Wire any StreamChannel instances found in transformer projections for projection in projections: @@ -355,15 +356,16 @@ async def create_async_graph_run_stream( class _PumpDrivenLog: - """Wraps an ``EventLog`` so that iteration drives the sync pump. + """Wraps a list so that iteration drives the sync pump. - Used by :attr:`GraphRunStream.extensions` to make extension logs - iterable without requiring the caller to drain the stream first. + Used by all :class:`GraphRunStream` projections (``__iter__``, + ``.values``, ``.messages``, ``.extensions``) so that iterating + any projection lazily consumes the source. """ __slots__ = ("_log", "_pump_one") - def __init__(self, log: EventLog, pump_one: Callable[[], bool]) -> None: + def __init__(self, log: list, pump_one: Callable[[], bool]) -> None: self._log = log self._pump_one = pump_one @@ -527,7 +529,7 @@ class GraphRunStream: name = getattr(t, "name", None) value = getattr(t, "value", None) if name is not None and value is not None: - if isinstance(value, EventLog): + if isinstance(value, list): result[name] = _PumpDrivenLog(value, self._pump_one) else: result[name] = value diff --git a/libs/langgraph/langgraph/stream/stream_channel.py b/libs/langgraph/langgraph/stream/stream_channel.py index fade0c2a1..97849b8eb 100644 --- a/libs/langgraph/langgraph/stream/stream_channel.py +++ b/libs/langgraph/langgraph/stream/stream_channel.py @@ -1,24 +1,21 @@ """StreamChannel — typed push-based channel for StreamTransformer projections. -A ``StreamChannel`` wraps an :class:`EventLog` and declares a protocol -channel name. When the :class:`StreamMux` detects a ``StreamChannel`` -in a transformer's ``init()`` return, it wires every ``push()`` call to -inject a :class:`ProtocolEvent` into the main event stream using the -channel's name as the ``method``. +A ``StreamChannel`` wraps a list and declares a protocol channel name. +When the :class:`StreamMux` detects a ``StreamChannel`` in a transformer's +``init()`` return, it wires every ``push()`` call to inject a +:class:`ProtocolEvent` into the main event stream using the channel's +name as the ``method``. -In-process consumers iterate the channel directly (it is an async -iterable). Remote SDK clients subscribe via -``session.subscribe("custom:")``. +In-process consumers iterate the channel directly. Remote SDK clients +subscribe via ``session.subscribe("custom:")``. """ from __future__ import annotations -from collections.abc import AsyncIterator, Callable +from collections.abc import Callable from typing import Any, Generic, TypeVar -from langgraph.stream._event_log import EventLog - T = TypeVar("T") @@ -31,28 +28,19 @@ class StreamChannel(Generic[T]): channel on run completion. """ - __slots__ = ("channel_name", "_log", "_on_push") + __slots__ = ("channel_name", "_items", "_on_push") def __init__(self, name: str) -> None: self.channel_name = name - self._log: EventLog[T] = EventLog() + self._items: list[T] = [] self._on_push: Callable[[Any], None] | None = None def push(self, item: T) -> None: - """Push an item to the channel. - - If the mux has wired this channel, the push also injects a - protocol event into the main event stream. - """ - self._log.append(item) + """Push an item to the channel.""" + self._items.append(item) if self._on_push is not None: self._on_push(item) - # -- Async iteration (in-process consumption) --------------------------- - - def __aiter__(self) -> AsyncIterator[T]: - return self._log.subscribe(0) - # -- Internal (called by the mux) --------------------------------------- def _wire(self, fn: Callable[[Any], None]) -> None: @@ -60,12 +48,12 @@ class StreamChannel(Generic[T]): self._on_push = fn def _close(self) -> None: - """Close the underlying log. Called by the mux on normal completion.""" - self._log.close() + """No-op for compatibility. Called by the mux on normal completion.""" + pass def _fail(self, err: BaseException) -> None: - """Fail the underlying log. Called by the mux on failure.""" - self._log.fail(err) + """No-op for compatibility. Called by the mux on failure.""" + pass def is_stream_channel(value: object) -> bool: diff --git a/libs/langgraph/langgraph/stream/transformers.py b/libs/langgraph/langgraph/stream/transformers.py index cc9771791..3f4edf3b6 100644 --- a/libs/langgraph/langgraph/stream/transformers.py +++ b/libs/langgraph/langgraph/stream/transformers.py @@ -9,36 +9,33 @@ from __future__ import annotations from typing import Any -from langgraph.stream._event_log import EventLog -from langgraph.stream._types import ProtocolEvent +from langgraph.stream._types import ProtocolEvent, StreamTransformer from langgraph.stream.chat_model_stream import ChatModelStream # Type alias for the stream class constructor signature _StreamCls = type[ChatModelStream] -class ValuesTransformer: - """Extracts ``values`` events and populates a values event log. +class ValuesTransformer(StreamTransformer): + """Extracts ``values`` events and populates a values log. Maintains the latest state per namespace and provides a separate - event log that :class:`AsyncGraphRunStream` / :class:`GraphRunStream` uses for ``.values`` + log that :class:`AsyncGraphRunStream` / :class:`GraphRunStream` uses for ``.values`` iteration. - - Implements the :class:`StreamTransformer` protocol. """ name = "values" def __init__(self) -> None: - self._values_log: EventLog[dict[str, Any]] = EventLog() + self._values_log: list[dict[str, Any]] = [] self._latest: dict[str, Any] = {} @property - def value(self) -> EventLog[dict[str, Any]]: + def value(self) -> list[dict[str, Any]]: return self._values_log @property - def values_log(self) -> EventLog[dict[str, Any]]: + def values_log(self) -> list[dict[str, Any]]: return self._values_log def get_latest(self, ns_key: str = "") -> Any: @@ -61,20 +58,18 @@ class ValuesTransformer: return True def finalize(self) -> None: - self._values_log.close() + pass def fail(self, err: BaseException) -> None: - self._values_log.fail(err) + pass -class MessagesTransformer: +class MessagesTransformer(StreamTransformer): """Groups ``messages`` events into :class:`ChatModelStream` instances. One ``ChatModelStream`` is created per ``message-start`` event. Content-block events are routed to the active stream until ``message-finish`` or ``message-error`` closes it. - - Implements the :class:`StreamTransformer` protocol. """ name = "messages" @@ -91,17 +86,17 @@ class MessagesTransformer: self._stream_cls: _StreamCls = stream_cls or ChatModelStream # Message log for .messages iteration - self._messages_log: EventLog[ChatModelStream] = EventLog() + self._messages_log: list[ChatModelStream] = [] # Current active stream per namespace key self._active: dict[str, ChatModelStream] = {} @property - def value(self) -> EventLog[ChatModelStream]: + def value(self) -> list[ChatModelStream]: return self._messages_log @property - def messages_log(self) -> EventLog[ChatModelStream]: + def messages_log(self) -> list[ChatModelStream]: return self._messages_log def init(self) -> Any: @@ -160,17 +155,15 @@ class MessagesTransformer: return True def finalize(self) -> None: - # Close any remaining active streams + # Finish any remaining active streams for stream in self._active.values(): stream._finish({"reason": "stop"}) self._active.clear() - self._messages_log.close() def fail(self, err: BaseException) -> None: for stream in self._active.values(): stream._fail(err) self._active.clear() - self._messages_log.fail(err) __all__ = [ diff --git a/libs/langgraph/tests/test_graph_run_stream.py b/libs/langgraph/tests/test_graph_run_stream.py index a0583ac2a..5e95e617d 100644 --- a/libs/langgraph/tests/test_graph_run_stream.py +++ b/libs/langgraph/tests/test_graph_run_stream.py @@ -9,7 +9,7 @@ from typing_extensions import TypedDict from langgraph.config import get_stream_writer from langgraph.graph import END, START, MessagesState, StateGraph from langgraph.stream import AsyncChatModelStream, StreamingHandler -from langgraph.stream._types import ProtocolEvent +from langgraph.stream._types import ProtocolEvent, StreamTransformer from tests.fake_chat import FakeChatModel @@ -379,7 +379,7 @@ async def test_subgraph_child_output(): # --------------------------------------------------------------------------- -class _CountTransformer: +class _CountTransformer(StreamTransformer): """Counts events. Exposes count via .value for extensions.""" name = "event_count" @@ -427,6 +427,73 @@ def test_sync_custom_reducer_extensions(): assert run.extensions["event_count"] == counter.value +# --------------------------------------------------------------------------- +# Double iteration over .values +# --------------------------------------------------------------------------- + + +@pytest.mark.anyio +async def test_async_values_double_iteration(): + """Iterating over run.values twice should yield the same snapshots both times.""" + graph = make_simple_graph() + run = await StreamingHandler(graph).astream({"value": "x", "items": []}) + await asyncio.sleep(0.1) + + first = [] + async for v in run.values: + first.append(v) + + second = [] + async for v in run.values: + second.append(v) + + assert len(first) == 3 + assert first == second + + +def test_sync_values_double_iteration(): + """Iterating over run.values twice should yield the same snapshots both times.""" + graph = make_simple_graph() + run = StreamingHandler(graph).stream({"value": "x", "items": []}) + + first = list(run.values) + second = list(run.values) + + assert len(first) == 3 + assert first == second + + +@pytest.mark.anyio +async def test_async_raw_events_double_iteration(): + """Iterating over the raw event stream twice should yield the same events.""" + graph = make_simple_graph() + run = await StreamingHandler(graph).astream({"value": "x", "items": []}) + await asyncio.sleep(0.1) + + first = [] + async for event in run: + first.append(event) + + second = [] + async for event in run: + second.append(event) + + assert len(first) > 0 + assert first == second + + +def test_sync_raw_events_double_iteration(): + """Iterating over the raw event stream twice should yield the same events.""" + graph = make_simple_graph() + run = StreamingHandler(graph).stream({"value": "x", "items": []}) + + first = list(run) + second = list(run) + + assert len(first) > 0 + assert first == second + + # --------------------------------------------------------------------------- # Tool transformer via extensions # --------------------------------------------------------------------------- @@ -440,15 +507,13 @@ class _ToolExecution: self.output = output -class _ToolsTransformer: +class _ToolsTransformer(StreamTransformer): """Groups tool-started/tool-finished custom events into _ToolExecution objects.""" name = "tools" def __init__(self) -> None: - from langgraph.stream._event_log import EventLog - - self._log: EventLog[_ToolExecution] = EventLog() + self._log: list[_ToolExecution] = [] self._pending: dict[str, dict] = {} self.value = self._log @@ -483,10 +548,10 @@ class _ToolsTransformer: return True def finalize(self) -> None: - self._log.close() + pass def fail(self, err: BaseException) -> None: - self._log.fail(err) + pass def _make_tool_graph(): diff --git a/libs/langgraph/tests/test_stream_convert.py b/libs/langgraph/tests/test_stream_convert.py index 091363854..e7db2f529 100644 --- a/libs/langgraph/tests/test_stream_convert.py +++ b/libs/langgraph/tests/test_stream_convert.py @@ -51,13 +51,6 @@ def test_namespace_passthrough(): assert evt["params"]["namespace"] == ["agent", "0"] -def test_timestamp_populated(): - evt = convert_to_protocol_event((), "values", {}) - assert evt is not None - assert isinstance(evt["params"]["timestamp"], int) - assert evt["params"]["timestamp"] > 0 - - def test_unknown_mode_returns_none(): assert convert_to_protocol_event((), "unknown_mode", {}) is None diff --git a/libs/langgraph/tests/test_stream_event_log.py b/libs/langgraph/tests/test_stream_event_log.py deleted file mode 100644 index 7bec10156..000000000 --- a/libs/langgraph/tests/test_stream_event_log.py +++ /dev/null @@ -1,144 +0,0 @@ -import asyncio - -import pytest - -from langgraph.stream._event_log import EventLog - - -@pytest.mark.anyio -async def test_push_and_iterate_in_order(): - log = EventLog() - log.append("a") - log.append("b") - log.append("c") - log.close() - items = [item async for item in log.subscribe(0)] - assert items == ["a", "b", "c"] - - -@pytest.mark.anyio -async def test_multiple_independent_cursors(): - log = EventLog() - log.append("x") - log.append("y") - log.close() - items1 = [item async for item in log.subscribe(0)] - items2 = [item async for item in log.subscribe(0)] - assert items1 == ["x", "y"] - assert items2 == ["x", "y"] - - -@pytest.mark.anyio -async def test_close_ends_iteration(): - log = EventLog() - log.close() - items = [item async for item in log.subscribe(0)] - assert items == [] - - -@pytest.mark.anyio -async def test_fail_raises_error(): - log = EventLog() - log.fail(RuntimeError("boom")) - with pytest.raises(RuntimeError, match="boom"): - async for _ in log.subscribe(0): - pass - - -@pytest.mark.anyio -async def test_nonzero_offset(): - log = EventLog() - log.append("x") - log.append("y") - log.append("z") - log.close() - items = [item async for item in log.subscribe(2)] - assert items == ["z"] - - -@pytest.mark.anyio -async def test_concurrent_push_and_iterate(): - log = EventLog() - received = [] - - async def consumer(): - async for item in log.subscribe(0): - received.append(item) - - async def producer(): - for i in range(5): - log.append(i) - await asyncio.sleep(0.01) - log.close() - - await asyncio.gather(producer(), consumer()) - assert received == [0, 1, 2, 3, 4] - - -@pytest.mark.anyio -async def test_items_before_cursor_visible(): - log = EventLog() - log.append("a") - log.append("b") - cursor = log.subscribe(0) - log.append("c") - log.close() - items = [item async for item in cursor] - assert items == ["a", "b", "c"] - - -@pytest.mark.anyio -async def test_empty_log_closed_yields_nothing(): - log = EventLog() - log.close() - items = [item async for item in log.subscribe(0)] - assert items == [] - - -@pytest.mark.anyio -async def test_fail_mid_iteration(): - """A cursor that has consumed some items should raise when fail() is called.""" - log = EventLog() - received = [] - - async def consumer(): - async for item in log.subscribe(0): - received.append(item) - - async def producer(): - log.append("a") - log.append("b") - await asyncio.sleep(0.02) - log.fail(RuntimeError("mid-stream error")) - - with pytest.raises(RuntimeError, match="mid-stream error"): - await asyncio.gather(producer(), consumer()) - assert received == ["a", "b"] - - -@pytest.mark.anyio -async def test_abandoned_cursor_cleans_up_waiters(): - """Abandoned async cursors should not leave stale futures in the - EventLog waiter list. - - When a cursor's __anext__ is cancelled (e.g. consumer breaks out of - ``async for``), the Future it registered in ``_waiters`` should be - cleaned up. Otherwise the list grows without bound until the next - append/close/fail triggers ``_wake_all()``. - """ - log: EventLog[str] = EventLog() - - for _ in range(10): - cursor = log.subscribe(0) - task = asyncio.ensure_future(cursor.__anext__()) - await asyncio.sleep(0) # let task register its waiter - task.cancel() - try: - await task - except asyncio.CancelledError: - pass - - assert len(log._waiters) == 0, ( - f"Expected 0 waiters after abandoning 10 cursors, " - f"got {len(log._waiters)}. Abandoned cursors leak futures." - ) diff --git a/libs/langgraph/tests/test_stream_mux.py b/libs/langgraph/tests/test_stream_mux.py index d2160e5f1..7655c4d35 100644 --- a/libs/langgraph/tests/test_stream_mux.py +++ b/libs/langgraph/tests/test_stream_mux.py @@ -3,8 +3,8 @@ from typing import Any import pytest from langgraph.stream._convert import convert_to_protocol_event -from langgraph.stream._mux import StreamMux -from langgraph.stream._types import ProtocolEvent +from langgraph.stream._mux import AsyncStreamMux, StreamMux +from langgraph.stream._types import ProtocolEvent, StreamTransformer from langgraph.stream.stream_channel import StreamChannel @@ -14,7 +14,7 @@ def _event(mode: str, data: Any, ns: list[str] | None = None) -> ProtocolEvent: return ev -class _MockTransformer: +class _MockTransformer(StreamTransformer): def __init__(self, *, suppress: bool = False): self.calls: list[ProtocolEvent] = [] self._suppress = suppress @@ -70,7 +70,7 @@ async def test_top_level_ns_only(): @pytest.mark.anyio async def test_subscribe_events_filter(): - mux = StreamMux() + mux = AsyncStreamMux() mux.push(_event("values", {"a": 1}, ns=["child:0"])) mux.push(_event("values", {"b": 2}, ns=["other:1"])) mux.push(_event("values", {"c": 3}, ns=["child:0"])) @@ -86,7 +86,7 @@ async def test_subscribe_events_filter(): @pytest.mark.anyio async def test_close_resolves_output(): - mux = StreamMux() + mux = AsyncStreamMux() fut = mux.get_output_future() mux.push(_event("values", {"v": 1})) mux.push(_event("values", {"v": 2})) @@ -97,7 +97,7 @@ async def test_close_resolves_output(): @pytest.mark.anyio async def test_fail_rejects_output(): - mux = StreamMux() + mux = AsyncStreamMux() fut = mux.get_output_future() mux.fail(ValueError("boom")) with pytest.raises(ValueError, match="boom"): @@ -155,7 +155,7 @@ async def test_push_after_close_ignored(): @pytest.mark.anyio async def test_fail_rejects_all_futures(): - mux = StreamMux() + mux = AsyncStreamMux() fut1 = mux.get_output_future([]) fut2 = mux.get_output_future(["child:0"]) mux.fail(ValueError("boom")) @@ -172,7 +172,7 @@ async def test_channel_events_bypass_transformer_pipeline(): the JS implementation and avoids re-entrancy bugs. """ mock = _MockTransformer() - mux = StreamMux(transformers=[mock]) + mux = AsyncStreamMux(transformers=[mock]) channel: StreamChannel[str] = StreamChannel("my_channel") mux.wire_channels({"ch": channel}) @@ -208,7 +208,7 @@ async def test_event_log_has_monotonic_seq_numbers(): while channel-emitted events use a separate counter (``_next_emit_seq``). When interleaved, seq numbers can duplicate. """ - mux = StreamMux() + mux = AsyncStreamMux() channel: StreamChannel[str] = StreamChannel("test_ch") mux.wire_channels({"ch": channel}) @@ -243,7 +243,7 @@ async def test_channel_push_during_process_preserves_namespace(): ``namespace: []`` instead of the original. """ - class _ChannelTransformer: + class _ChannelTransformer(StreamTransformer): """Pushes to its channel whenever it sees a ``values`` event.""" def __init__(self, name: str) -> None: @@ -266,7 +266,7 @@ async def test_channel_push_during_process_preserves_namespace(): t1 = _ChannelTransformer("first") t2 = _ChannelTransformer("second") - mux = StreamMux(transformers=[t1, t2]) + mux = AsyncStreamMux(transformers=[t1, t2]) mux.wire_channels({"first": t1.channel}) mux.wire_channels({"second": t2.channel}) diff --git a/libs/langgraph/tests/test_stream_reducers.py b/libs/langgraph/tests/test_stream_reducers.py index c0d119168..ca703b1e8 100644 --- a/libs/langgraph/tests/test_stream_reducers.py +++ b/libs/langgraph/tests/test_stream_reducers.py @@ -22,38 +22,29 @@ def _event( # -- ValuesTransformer --------------------------------------------------------- -@pytest.mark.anyio -async def test_values_captures_values_events(): +def test_values_captures_values_events(): reducer = ValuesTransformer() reducer.init() reducer.process(_event("values", {"a": 1})) reducer.process(_event("values", {"b": 2})) reducer.finalize() - collected = [] - async for item in reducer.values_log.subscribe(0): - collected.append(item) - assert len(collected) == 2 - assert collected[0]["data"] == {"a": 1} - assert collected[1]["data"] == {"b": 2} + assert len(reducer.values_log) == 2 + assert reducer.values_log[0]["data"] == {"a": 1} + assert reducer.values_log[1]["data"] == {"b": 2} -@pytest.mark.anyio -async def test_values_ignores_other_modes(): +def test_values_ignores_other_modes(): reducer = ValuesTransformer() reducer.init() reducer.process(_event("updates", {"x": 1})) reducer.process(_event("messages", {"event": "message-start"})) reducer.finalize() - collected = [] - async for item in reducer.values_log.subscribe(0): - collected.append(item) - assert len(collected) == 0 + assert len(reducer.values_log) == 0 -@pytest.mark.anyio -async def test_values_latest_per_namespace(): +def test_values_latest_per_namespace(): reducer = ValuesTransformer() reducer.init() reducer.process(_event("values", {"v": 1}, ns=["child:0"])) @@ -61,15 +52,6 @@ async def test_values_latest_per_namespace(): assert reducer.get_latest("child:0") == {"v": 2} -@pytest.mark.anyio -async def test_values_finalize_closes_log(): - reducer = ValuesTransformer() - reducer.init() - reducer.process(_event("values", {"a": 1})) - reducer.finalize() - assert reducer.values_log.closed - - # -- MessagesTransformer ------------------------------------------------------- @@ -103,8 +85,7 @@ def _msg_finish(ns=None, node=None): ) -@pytest.mark.anyio -async def test_messages_groups_lifecycle(): +def test_messages_groups_lifecycle(): reducer = MessagesTransformer() reducer.init() reducer.process(_msg_start()) @@ -112,16 +93,12 @@ async def test_messages_groups_lifecycle(): reducer.process(_msg_finish()) reducer.finalize() - collected = [] - async for stream in reducer.messages_log.subscribe(0): - collected.append(stream) - assert len(collected) == 1 - assert isinstance(collected[0], ChatModelStream) - assert collected[0].done + assert len(reducer.messages_log) == 1 + assert isinstance(reducer.messages_log[0], ChatModelStream) + assert reducer.messages_log[0].done -@pytest.mark.anyio -async def test_messages_multiple_sequential(): +def test_messages_multiple_sequential(): reducer = MessagesTransformer() reducer.init() reducer.process(_msg_start(message_id="m1")) @@ -130,14 +107,10 @@ async def test_messages_multiple_sequential(): reducer.process(_msg_finish()) reducer.finalize() - collected = [] - async for stream in reducer.messages_log.subscribe(0): - collected.append(stream) - assert len(collected) == 2 + assert len(reducer.messages_log) == 2 -@pytest.mark.anyio -async def test_messages_namespace_filter(): +def test_messages_namespace_filter(): reducer = MessagesTransformer(namespace=["root"]) reducer.init() reducer.process(_msg_start(ns=["root"])) @@ -146,14 +119,10 @@ async def test_messages_namespace_filter(): reducer.process(_msg_finish(ns=["other"])) reducer.finalize() - collected = [] - async for stream in reducer.messages_log.subscribe(0): - collected.append(stream) - assert len(collected) == 1 + assert len(reducer.messages_log) == 1 -@pytest.mark.anyio -async def test_messages_node_filter(): +def test_messages_node_filter(): reducer = MessagesTransformer(node_filter="agent") reducer.init() reducer.process(_msg_start(node="agent")) @@ -162,14 +131,10 @@ async def test_messages_node_filter(): reducer.process(_msg_finish(node="tools")) reducer.finalize() - collected = [] - async for stream in reducer.messages_log.subscribe(0): - collected.append(stream) - assert len(collected) == 1 + assert len(reducer.messages_log) == 1 -@pytest.mark.anyio -async def test_messages_error_event(): +def test_messages_error_event(): """An error event should fail the active ChatModelStream.""" reducer = MessagesTransformer() reducer.init() @@ -180,35 +145,17 @@ async def test_messages_error_event(): ) reducer.finalize() - collected: list[ChatModelStream] = [] - async for stream in reducer.messages_log.subscribe(0): - collected.append(stream) - assert len(collected) == 1 - assert collected[0].done + assert len(reducer.messages_log) == 1 + assert reducer.messages_log[0].done -@pytest.mark.anyio -async def test_messages_fail_propagates_to_active(): - """transformer.fail() should propagate the error to any active streams.""" +def test_messages_fail_propagates_to_active(): + """transformer.fail() should mark active streams as done.""" reducer = MessagesTransformer() reducer.init() reducer.process(_msg_start()) reducer.process(_content_delta("partial")) reducer.fail(RuntimeError("graph failed")) - # The messages log should be failed too - with pytest.raises(RuntimeError, match="graph failed"): - async for _ in reducer.messages_log.subscribe(0): - pass - - -@pytest.mark.anyio -async def test_values_fail_propagates(): - reducer = ValuesTransformer() - reducer.init() - reducer.process(_event("values", {"a": 1})) - reducer.fail(RuntimeError("graph failed")) - - with pytest.raises(RuntimeError, match="graph failed"): - async for _ in reducer.values_log.subscribe(0): - pass + assert len(reducer.messages_log) == 1 + assert reducer.messages_log[0].done diff --git a/libs/langgraph/tests/test_stream_run_stream.py b/libs/langgraph/tests/test_stream_run_stream.py index 09cb5b67b..d73bb20e2 100644 --- a/libs/langgraph/tests/test_stream_run_stream.py +++ b/libs/langgraph/tests/test_stream_run_stream.py @@ -4,7 +4,7 @@ from typing import Any import pytest -from langgraph.stream._mux import StreamMux +from langgraph.stream._mux import AsyncStreamMux, StreamMux from langgraph.stream._types import ProtocolEvent from langgraph.stream.chat_model_stream import ChatModelStream from langgraph.stream.run_stream import ( @@ -44,7 +44,7 @@ async def test_aiter_yields_all_events(): @pytest.mark.anyio async def test_subgraph_name_and_index(): vr, mr = ValuesTransformer(), MessagesTransformer() - mux = StreamMux(transformers=[vr, mr]) + mux = AsyncStreamMux(transformers=[vr, mr]) sub = AsyncSubgraphRunStream( mux=mux, namespace=["researcher:2"], @@ -57,7 +57,7 @@ async def test_subgraph_name_and_index(): @pytest.mark.anyio async def test_subgraph_name_no_index(): vr, mr = ValuesTransformer(), MessagesTransformer() - mux = StreamMux(transformers=[vr, mr]) + mux = AsyncStreamMux(transformers=[vr, mr]) sub = AsyncSubgraphRunStream( mux=mux, namespace=["agent"], transformers=[vr, mr] ) @@ -132,7 +132,7 @@ async def test_messages_yields_streams(): @pytest.mark.anyio async def test_interrupted_false_by_default(): vr, mr = ValuesTransformer(), MessagesTransformer() - mux = StreamMux(transformers=[vr, mr]) + mux = AsyncStreamMux(transformers=[vr, mr]) run = AsyncGraphRunStream(mux=mux, transformers=[vr, mr]) assert run.interrupted is False @@ -140,7 +140,7 @@ async def test_interrupted_false_by_default(): @pytest.mark.anyio async def test_abort_sets_signal(): vr, mr = ValuesTransformer(), MessagesTransformer() - mux = StreamMux(transformers=[vr, mr]) + mux = AsyncStreamMux(transformers=[vr, mr]) run = AsyncGraphRunStream(mux=mux, transformers=[vr, mr]) assert not run.signal.is_set() run.abort()