From 0076da90086f1b181c1bea0f3d6dc5e8d5363f8a Mon Sep 17 00:00:00 2001 From: Nick Hollon Date: Wed, 15 Apr 2026 18:12:03 -0400 Subject: [PATCH] feat(langgraph): add streaming transformer infrastructure and tests Introduces the StreamingHandler, StreamMux, EventLog, StreamChannel, and StreamTransformer abstractions for ergonomic streaming projections over compiled graphs. Includes ValuesTransformer and MessagesTransformer as built-in native projections, plus support for user-defined custom transformers. --- libs/langgraph/langgraph/stream/__init__.py | 21 + libs/langgraph/langgraph/stream/_convert.py | 24 + libs/langgraph/langgraph/stream/_event_log.py | 138 +++ libs/langgraph/langgraph/stream/_mux.py | 139 +++ libs/langgraph/langgraph/stream/_types.py | 79 ++ libs/langgraph/langgraph/stream/run_stream.py | 119 ++ .../langgraph/stream/stream_channel.py | 63 + .../langgraph/stream/streaming_handler.py | 161 +++ .../langgraph/stream/transformers.py | 82 ++ .../langgraph/tests/test_streaming_handler.py | 1051 +++++++++++++++++ 10 files changed, 1877 insertions(+) create mode 100644 libs/langgraph/langgraph/stream/__init__.py create mode 100644 libs/langgraph/langgraph/stream/_convert.py create mode 100644 libs/langgraph/langgraph/stream/_event_log.py create mode 100644 libs/langgraph/langgraph/stream/_mux.py create mode 100644 libs/langgraph/langgraph/stream/_types.py create mode 100644 libs/langgraph/langgraph/stream/run_stream.py create mode 100644 libs/langgraph/langgraph/stream/stream_channel.py create mode 100644 libs/langgraph/langgraph/stream/streaming_handler.py create mode 100644 libs/langgraph/langgraph/stream/transformers.py create mode 100644 libs/langgraph/tests/test_streaming_handler.py diff --git a/libs/langgraph/langgraph/stream/__init__.py b/libs/langgraph/langgraph/stream/__init__.py new file mode 100644 index 000000000..386deeb1a --- /dev/null +++ b/libs/langgraph/langgraph/stream/__init__.py @@ -0,0 +1,21 @@ +"""Streaming infrastructure for LangGraph. + +Provides a ``StreamingHandler`` that wraps a compiled graph and exposes +ergonomic streaming projections through a transformer pipeline. +""" + +from langgraph.stream._event_log import EventLog +from langgraph.stream._types import ProtocolEvent, StreamTransformer +from langgraph.stream.run_stream import AsyncGraphRunStream, GraphRunStream +from langgraph.stream.stream_channel import StreamChannel +from langgraph.stream.streaming_handler import StreamingHandler + +__all__ = [ + "AsyncGraphRunStream", + "EventLog", + "GraphRunStream", + "ProtocolEvent", + "StreamChannel", + "StreamTransformer", + "StreamingHandler", +] diff --git a/libs/langgraph/langgraph/stream/_convert.py b/libs/langgraph/langgraph/stream/_convert.py new file mode 100644 index 000000000..aeaf3f17a --- /dev/null +++ b/libs/langgraph/langgraph/stream/_convert.py @@ -0,0 +1,24 @@ +from __future__ import annotations + +from typing import Any + +from langgraph.stream._types import ProtocolEvent, _ProtocolEventParams + + +def convert_to_protocol_event(part: dict[str, Any]) -> ProtocolEvent: + """Convert a v2 StreamPart dict to a ProtocolEvent. + + Expects a dict with keys ``type``, ``ns``, ``data``, and optionally + ``interrupts`` (present on values events). + """ + params: _ProtocolEventParams = { + "namespace": list(part["ns"]), + "data": part["data"], + } + if "interrupts" in part: + params["interrupts"] = part["interrupts"] + return { + "type": "event", + "method": part["type"], + "params": params, + } diff --git a/libs/langgraph/langgraph/stream/_event_log.py b/libs/langgraph/langgraph/stream/_event_log.py new file mode 100644 index 000000000..4bd6f775f --- /dev/null +++ b/libs/langgraph/langgraph/stream/_event_log.py @@ -0,0 +1,138 @@ +from __future__ import annotations + +import asyncio +import threading +from collections.abc import AsyncIterator, Iterator +from typing import Generic, TypeVar + +T = TypeVar("T") + + +class EventLog(Generic[T]): + """Append-only buffer with multi-cursor sync and async iteration. + + Supports multiple independent consumers iterating the same log + concurrently. Each call to ``__iter__`` or ``__aiter__`` creates a + new cursor starting from the beginning. + + A given instance should be used in either sync or async mode — the + two notification paths are independent and do not interfere, but + mixing them on one instance is not tested. + + Producer API (thread-safe): + push(item) — append an item, notify all waiting cursors + close() — mark the log as done + fail(err) — mark the log as errored + + Consumer API: + __iter__() — new sync cursor (blocks via threading.Condition) + __aiter__() — new async cursor (awaits via asyncio.Future) + """ + + def __init__(self) -> None: + self._items: list[T] = [] + self._closed = False + self._error: BaseException | None = None + # Sync notification + self._lock = threading.Lock() + self._cond = threading.Condition(self._lock) + # Async notification — futures created lazily by async cursors + self._async_waiters: list[asyncio.Future[None]] = [] + + # ------------------------------------------------------------------ + # Producer API + # ------------------------------------------------------------------ + + def push(self, item: T) -> None: + """Append *item* and wake all waiting cursors.""" + with self._lock: + if self._closed: + raise RuntimeError("Cannot push to a closed EventLog") + self._items.append(item) + self._cond.notify_all() + self._wake_async() + + def close(self) -> None: + """Mark the log as complete — open cursors will finish cleanly.""" + with self._lock: + self._closed = True + self._cond.notify_all() + self._wake_async() + + def fail(self, err: BaseException) -> None: + """Mark the log as errored — open cursors will raise *err*.""" + with self._lock: + self._error = err + self._closed = True + self._cond.notify_all() + self._wake_async() + + # ------------------------------------------------------------------ + # Sync iteration + # ------------------------------------------------------------------ + + def __iter__(self) -> Iterator[T]: + """Return a new independent sync cursor over the log.""" + return self._sync_cursor() + + def _sync_cursor(self) -> Iterator[T]: + cursor = 0 + while True: + with self._lock: + # Wait until data is available or the log is done. + while cursor >= len(self._items) and not self._closed: + self._cond.wait() + # Yield available items before raising errors, matching + # the async cursor's behavior. + if cursor < len(self._items): + item = self._items[cursor] + cursor += 1 + elif self._error is not None: + raise self._error + else: + # closed and no more items + return + yield item + + # ------------------------------------------------------------------ + # Async iteration + # ------------------------------------------------------------------ + + def __aiter__(self) -> AsyncIterator[T]: + """Return a new independent async cursor over the log.""" + return self._async_cursor() + + async def _async_cursor(self) -> AsyncIterator[T]: + cursor = 0 + while True: + if cursor < len(self._items): + yield self._items[cursor] + cursor += 1 + elif self._closed: + if self._error is not None: + raise self._error + return + else: + # Wait for notification from push/close/fail. + loop = asyncio.get_running_loop() + fut: asyncio.Future[None] = loop.create_future() + self._async_waiters.append(fut) + await fut + + # ------------------------------------------------------------------ + # Internal helpers + # ------------------------------------------------------------------ + + def _wake_async(self) -> None: + """Resolve all pending async futures (safe from any thread).""" + waiters = self._async_waiters + if not waiters: + return + self._async_waiters = [] + for fut in waiters: + if not fut.done(): + try: + fut.get_loop().call_soon_threadsafe(fut.set_result, None) + except RuntimeError: + # Event loop already closed — nothing to notify. + pass diff --git a/libs/langgraph/langgraph/stream/_mux.py b/libs/langgraph/langgraph/stream/_mux.py new file mode 100644 index 000000000..8ed740e8d --- /dev/null +++ b/libs/langgraph/langgraph/stream/_mux.py @@ -0,0 +1,139 @@ +from __future__ import annotations + +from collections.abc import 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 + + +class StreamMux: + """Central event dispatcher for the streaming infrastructure. + + Owns the main ``EventLog[ProtocolEvent]`` 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. + """ + + def __init__(self) -> None: + self._events: EventLog[ProtocolEvent] = EventLog() + self._transformers: list[StreamTransformer] = [] + self._channels: list[StreamChannel[Any]] = [] + self._seq = 0 + + def register(self, transformer: StreamTransformer) -> dict[str, Any]: + """Register a transformer and return its projection dict. + + Calls ``transformer.init()``, stores the transformer for event + processing, and returns the projection. StreamChannels in the + projection are auto-wired. + """ + projection = transformer.init() + if not isinstance(projection, dict): + raise TypeError( + f"StreamTransformer.init() must return a dict, " + f"got {type(projection).__name__}" + ) + self._transformers.append(transformer) + self._wire_channels(projection) + return projection + + def push(self, event: ProtocolEvent) -> None: + """Route *event* through all transformers, then append to the main log. + + Each transformer's ``process()`` is called in registration order. + If any transformer returns ``False``, the event is suppressed + from the main log (but transformers that already saw it keep + their side-effects). + + Seq is assigned right before an event enters the main log, not + before the transformer pipeline runs. This ensures that events + auto-forwarded from StreamChannels during ``process()`` get + earlier seq numbers than the original event, preserving + monotonic ordering in the log. + """ + keep = True + for transformer in self._transformers: + if not transformer.process(event): + keep = False + if keep: + self._seq += 1 + event["seq"] = self._seq + self._events.push(event) + + def close(self) -> None: + """Finalize all transformers, close all channels and the main log. + + If any transformer's ``finalize()`` raises, the remaining + transformers, channels, and the main log are still closed. + The first error is re-raised after cleanup completes. + """ + first_error: BaseException | None = None + for transformer in self._transformers: + try: + transformer.finalize() + except BaseException as e: + if first_error is None: + first_error = e + for ch in self._channels: + ch._close() + self._events.close() + if first_error is not None: + raise first_error + + def fail(self, err: BaseException) -> None: + """Fail all transformers, channels, and the main log. + + If any transformer's ``fail()`` raises, the remaining + transformers, channels, and the main log are still failed. + """ + for transformer in self._transformers: + try: + transformer.fail(err) + except BaseException: + pass + for ch in self._channels: + ch._fail(err) + self._events.fail(err) + + # ------------------------------------------------------------------ + # StreamChannel auto-wiring + # ------------------------------------------------------------------ + + def _wire_channels(self, projection: dict[str, Any]) -> None: + """Find StreamChannel instances in *projection* and wire them.""" + for value in projection.values(): + if isinstance(value, StreamChannel): + self._channels.append(value) + channel_name = value.name + + def _make_forward(name: str) -> Callable[[Any], None]: + def _forward(item: Any) -> None: + self._forward(name, item) + + return _forward + + value._wire(_make_forward(channel_name)) + + def _forward(self, channel_name: str, item: Any) -> None: + """Inject a ProtocolEvent for a StreamChannel push. + + Forwarded events bypass the transformer pipeline to avoid + infinite recursion (a transformer that pushes to a channel + during ``process()`` would re-trigger itself). These events + are visible in the main event log but are not passed through + transformers' ``process()`` methods. + """ + self._seq += 1 + event: ProtocolEvent = { + "type": "event", + "seq": self._seq, + "method": f"custom:{channel_name}", + "params": { + "namespace": [], + "data": item, + }, + } + self._events.push(event) diff --git a/libs/langgraph/langgraph/stream/_types.py b/libs/langgraph/langgraph/stream/_types.py new file mode 100644 index 000000000..f5764b751 --- /dev/null +++ b/libs/langgraph/langgraph/stream/_types.py @@ -0,0 +1,79 @@ +from __future__ import annotations + +from abc import ABC, abstractmethod +from typing import Any, Literal + +from typing_extensions import NotRequired, TypedDict + + +class _ProtocolEventParams(TypedDict): + """Parameters for a protocol event.""" + + namespace: list[str] + data: Any + interrupts: NotRequired[tuple[Any, ...]] + + +class ProtocolEvent(TypedDict): + """A protocol event emitted by the streaming infrastructure. + + Wraps a raw stream part (values, messages, custom, etc.) in a uniform + envelope with a monotonic sequence number assigned by the StreamMux. + """ + + type: Literal["event"] + seq: NotRequired[int] + method: str # StreamMode value: "values", "messages", "custom", etc. + params: _ProtocolEventParams + + +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.). + + Set `_native = True` on a transformer to have its projection keys + exposed as direct attributes on the run stream (in addition to + appearing in `run.extensions`). + + Subclasses must implement `init` and `process`. The `finalize` and + `fail` hooks are optional — the default implementations are no-ops. + StreamChannel instances are auto-closed/failed by the mux regardless. + """ + + @abstractmethod + def init(self) -> Any: + """Return the projection dict. + + Keys become entries in `run.extensions`. If the transformer has + `_native = True`, keys are also set as direct attributes on the + run stream. + + StreamChannel instances in the return value are automatically + wired by the StreamMux for protocol event auto-forwarding. + """ + ... + + @abstractmethod + def process(self, event: ProtocolEvent) -> bool: + """Process a protocol event. + + Called for every event before it is appended to the main event log. + Return False to suppress the event from the main log. + """ + ... + + def finalize(self) -> None: + """Called when the run ends normally. + + Override to close EventLogs, resolve promises, or perform other + teardown. StreamChannel instances are auto-closed by the mux. + """ + + def fail(self, err: BaseException) -> None: + """Called when the run ends with an error. + + Override to fail EventLogs, reject promises, or perform other + teardown. StreamChannel instances are auto-failed by the mux. + """ diff --git a/libs/langgraph/langgraph/stream/run_stream.py b/libs/langgraph/langgraph/stream/run_stream.py new file mode 100644 index 000000000..51058305b --- /dev/null +++ b/libs/langgraph/langgraph/stream/run_stream.py @@ -0,0 +1,119 @@ +from __future__ import annotations + +import asyncio +import threading +from collections.abc import AsyncIterator, Iterator +from typing import Any + +from langgraph.stream._mux import StreamMux +from langgraph.stream._types import ProtocolEvent +from langgraph.stream.transformers import ValuesTransformer + + +class GraphRunStream: + """Sync run stream with transformer-driven projections. + + All transformer projections live in ``extensions``. Native transformer + projections (those with ``_native = True``) are also set as direct + attributes on this instance (e.g. ``run.values``, ``run.messages``). + + Iterating the run stream directly yields raw ``ProtocolEvent`` objects + from the mux's main event log. + """ + + def __init__( + self, + mux: StreamMux, + extensions: dict[str, Any], + values_transformer: ValuesTransformer, + pump_thread: threading.Thread, + ) -> None: + self._mux = mux + self.extensions = extensions + self._values_transformer = values_transformer + self._pump_thread = pump_thread + + @property + def output(self) -> dict[str, Any] | None: + """Block until the run completes and return the final state.""" + self._pump_thread.join() + if self._values_transformer._log._error is not None: + raise self._values_transformer._log._error + return self._values_transformer._latest + + @property + def interrupted(self) -> bool: + """Block until the run completes, then return whether it was interrupted.""" + self._pump_thread.join() + return self._values_transformer._interrupted + + @property + def interrupts(self) -> list[Any]: + """Block until the run completes, then return interrupt payloads.""" + self._pump_thread.join() + return self._values_transformer._interrupts + + def __iter__(self) -> Iterator[ProtocolEvent]: + """Iterate all protocol events from the mux's main event log.""" + return iter(self._mux._events) + + +class AsyncGraphRunStream: + """Async run stream with transformer-driven projections. + + All transformer projections live in ``extensions``. Native transformer + projections (those with ``_native = True``) are also set as direct + attributes on this instance (e.g. ``run.values``, ``run.messages``). + + Async-iterating the run stream yields raw ``ProtocolEvent`` objects + from the mux's main event log. + """ + + def __init__( + self, + mux: StreamMux, + extensions: dict[str, Any], + values_transformer: ValuesTransformer, + pump_task: asyncio.Task[None], + ) -> None: + self._mux = mux + self.extensions = extensions + self._values_transformer = values_transformer + self._pump_task = pump_task + + @property + def output(self) -> Any: + """Return an awaitable that resolves to the final state. + + Usage:: + + output = await run.output + """ + return self._get_output() + + async def _get_output(self) -> dict[str, Any] | None: + try: + await self._pump_task + except BaseException: + pass + if self._values_transformer._log._error is not None: + raise self._values_transformer._log._error + return self._values_transformer._latest + + @property + def interrupted(self) -> bool: + """Whether the run was interrupted. + + Only meaningful after the run has completed (after consuming the + stream or awaiting ``output``). + """ + return self._values_transformer._interrupted + + @property + def interrupts(self) -> list[Any]: + """Interrupt payloads, populated when interrupted is True.""" + return self._values_transformer._interrupts + + def __aiter__(self) -> AsyncIterator[ProtocolEvent]: + """Iterate all protocol events from the mux's main event log.""" + return self._mux._events.__aiter__() diff --git a/libs/langgraph/langgraph/stream/stream_channel.py b/libs/langgraph/langgraph/stream/stream_channel.py new file mode 100644 index 000000000..d5ca8d070 --- /dev/null +++ b/libs/langgraph/langgraph/stream/stream_channel.py @@ -0,0 +1,63 @@ +from __future__ import annotations + +from collections.abc import AsyncIterator, 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. + + Wraps an `EventLog[T]` 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``. + + 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:")``. + + Lifecycle (``_close`` / ``_fail``) is managed by the mux — transformers + using only StreamChannels don't need ``finalize`` / ``fail`` hooks. + """ + + def __init__(self, name: str) -> None: + self.name = name + self._log: EventLog[T] = EventLog() + self._wire_fn: Callable[[T], None] | None = None + + def push(self, item: T) -> None: + """Append *item* to the log and auto-forward if wired.""" + self._log.push(item) + if self._wire_fn is not None: + self._wire_fn(item) + + # ------------------------------------------------------------------ + # Mux lifecycle hooks (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() + + def _fail(self, err: BaseException) -> None: + """Fail the underlying log (called by StreamMux on run error).""" + self._log.fail(err) + + # ------------------------------------------------------------------ + # Iteration — delegates to the inner EventLog (multi-cursor) + # ------------------------------------------------------------------ + + def __iter__(self) -> Iterator[T]: + return iter(self._log) + + def __aiter__(self) -> AsyncIterator[T]: + return self._log.__aiter__() diff --git a/libs/langgraph/langgraph/stream/streaming_handler.py b/libs/langgraph/langgraph/stream/streaming_handler.py new file mode 100644 index 000000000..ff2365873 --- /dev/null +++ b/libs/langgraph/langgraph/stream/streaming_handler.py @@ -0,0 +1,161 @@ +from __future__ import annotations + +import asyncio +import threading +from collections.abc import Sequence +from typing import Any + +from langchain_core.runnables import RunnableConfig + +from langgraph.stream._convert import convert_to_protocol_event +from langgraph.stream._mux import StreamMux +from langgraph.stream._types import StreamTransformer +from langgraph.stream.run_stream import AsyncGraphRunStream, GraphRunStream +from langgraph.stream.transformers import MessagesTransformer, ValuesTransformer +from langgraph.types import All, StreamMode + +# All stream modes to request from the graph. +STREAM_V2_MODES: list[StreamMode] = [ + "values", + "updates", + "messages", + "custom", + "checkpoints", + "tasks", + "debug", +] + + +class StreamingHandler: + """Wraps a compiled graph and provides ergonomic streaming projections. + + Usage:: + + handler = StreamingHandler(graph) + + # Sync + run = handler.stream(input_data) + for state in run.values: + print(state) + output = run.output + + # Async + run = await handler.astream(input_data) + async for state in run.values: + print(state) + output = await run.output + """ + + def __init__(self, graph: Any) -> None: + self._graph = graph + + def stream( + self, + input: Any, + config: RunnableConfig | None = None, + *, + interrupt_before: All | Sequence[str] | None = None, + interrupt_after: All | Sequence[str] | None = None, + transformers: list[StreamTransformer] | None = None, + ) -> GraphRunStream: + """Start a sync streaming run. + + Returns a `GraphRunStream` immediately. A background daemon thread + pumps events from the graph into the transformer pipeline. + """ + mux, extensions, native_keys, values_t = self._setup(transformers) + + def pump() -> None: + try: + for part in self._graph.stream( + input, + config, + stream_mode=STREAM_V2_MODES, + subgraphs=True, + version="v2", + interrupt_before=interrupt_before, + interrupt_after=interrupt_after, + ): + mux.push(convert_to_protocol_event(part)) + mux.close() + except BaseException as e: + mux.fail(e) + + thread = threading.Thread(target=pump, daemon=True) + thread.start() + + run = GraphRunStream(mux, extensions, values_t, thread) + for key in native_keys: + setattr(run, key, extensions[key]) + return run + + async def astream( + self, + input: Any, + config: RunnableConfig | None = None, + *, + interrupt_before: All | Sequence[str] | None = None, + interrupt_after: All | Sequence[str] | None = None, + transformers: list[StreamTransformer] | None = None, + ) -> AsyncGraphRunStream: + """Start an async streaming run. + + Returns an `AsyncGraphRunStream` immediately. A background asyncio + task pumps events from the graph into the transformer pipeline. + """ + mux, extensions, native_keys, values_t = self._setup(transformers) + + async def pump() -> None: + try: + async for part in self._graph.astream( + input, + config, + stream_mode=STREAM_V2_MODES, + subgraphs=True, + version="v2", + interrupt_before=interrupt_before, + interrupt_after=interrupt_after, + ): + mux.push(convert_to_protocol_event(part)) + mux.close() + except BaseException as e: + mux.fail(e) + + task = asyncio.create_task(pump()) + + run = AsyncGraphRunStream(mux, extensions, values_t, task) + for key in native_keys: + setattr(run, key, extensions[key]) + return run + + # ------------------------------------------------------------------ + # Internal helpers + # ------------------------------------------------------------------ + + @staticmethod + def _setup( + user_transformers: list[StreamTransformer] | None, + ) -> tuple[StreamMux, dict[str, Any], set[str], ValuesTransformer]: + """Create the mux, register all transformers. + + Returns (mux, extensions, native_keys, values_transformer). + """ + mux = StreamMux() + + values_t = ValuesTransformer() + messages_t = MessagesTransformer() + + all_transformers: list[StreamTransformer] = [values_t, messages_t] + if user_transformers: + all_transformers.extend(user_transformers) + + extensions: dict[str, Any] = {} + native_keys: set[str] = set() + + for t in all_transformers: + projection = mux.register(t) + extensions.update(projection) + if getattr(t, "_native", False): + native_keys.update(projection.keys()) + + return mux, extensions, native_keys, values_t diff --git a/libs/langgraph/langgraph/stream/transformers.py b/libs/langgraph/langgraph/stream/transformers.py new file mode 100644 index 000000000..9c7bd8444 --- /dev/null +++ b/libs/langgraph/langgraph/stream/transformers.py @@ -0,0 +1,82 @@ +from __future__ import annotations + +from typing import Any + +from langgraph.stream._event_log import EventLog +from langgraph.stream._types import ProtocolEvent, StreamTransformer + + +class ValuesTransformer(StreamTransformer): + """Captures values events and projects them into an iterable of state snapshots. + + Native transformer — projection keys are exposed as direct attributes + on the run stream (e.g. ``run.values``). + """ + + _native = True + + def __init__(self) -> None: + self._log: EventLog[dict[str, Any]] = EventLog() + self._latest: dict[str, Any] | None = None + self._interrupted = False + self._interrupts: list[Any] = [] + + def init(self) -> dict[str, Any]: + return {"values": self._log} + + def process(self, event: ProtocolEvent) -> bool: + if event["method"] != "values": + return True + params = event["params"] + # Only capture root namespace events + if params["namespace"]: + return True + self._latest = params["data"] + self._log.push(params["data"]) + interrupts = params.get("interrupts", ()) + if interrupts: + self._interrupted = True + self._interrupts.extend(interrupts) + return True + + def finalize(self) -> None: + self._log.close() + + def fail(self, err: BaseException) -> None: + self._log.fail(err) + + +class MessagesTransformer(StreamTransformer): + """Captures messages events and passes through raw (chunk, metadata) tuples. + + This is the same shape as today's ``stream_mode="messages"`` output. + A follow-on PR will replace this with a richer transformer that + produces ChatModelStream objects using the protocol handler. + + Native transformer — projection keys are exposed as direct attributes + on the run stream (e.g. ``run.messages``). + """ + + _native = True + + def __init__(self) -> None: + self._log: EventLog[tuple[Any, dict[str, Any]]] = EventLog() + + def init(self) -> dict[str, Any]: + return {"messages": self._log} + + def process(self, event: ProtocolEvent) -> bool: + if event["method"] != "messages": + return True + params = event["params"] + # Only capture root namespace events + if params["namespace"]: + return True + self._log.push(params["data"]) + return True + + def finalize(self) -> None: + self._log.close() + + def fail(self, err: BaseException) -> None: + self._log.fail(err) diff --git a/libs/langgraph/tests/test_streaming_handler.py b/libs/langgraph/tests/test_streaming_handler.py new file mode 100644 index 000000000..4657cef53 --- /dev/null +++ b/libs/langgraph/tests/test_streaming_handler.py @@ -0,0 +1,1051 @@ +"""Tests for the StreamingHandler and its supporting infrastructure.""" + +from __future__ import annotations + +import asyncio +import operator +import sys +from typing import Annotated, Any + +import pytest +from langgraph.checkpoint.memory import InMemorySaver +from typing_extensions import TypedDict + +from langgraph.constants import END, START +from langgraph.graph import StateGraph +from langgraph.stream import ( + EventLog, + StreamChannel, + StreamingHandler, + StreamTransformer, +) +from langgraph.stream._convert import convert_to_protocol_event +from langgraph.stream._mux import StreamMux +from langgraph.stream._types import ProtocolEvent +from langgraph.stream.transformers import MessagesTransformer, ValuesTransformer +from langgraph.types import StreamWriter, interrupt + +NEEDS_CONTEXTVARS = pytest.mark.skipif( + sys.version_info < (3, 11), + reason="Python 3.11+ is required for async contextvars support", +) + + +# --------------------------------------------------------------------------- +# Shared state and graph builders +# --------------------------------------------------------------------------- + + +class SimpleState(TypedDict): + value: str + items: Annotated[list[str], operator.add] + + +def _build_simple_graph(): + """Two-node graph: node_a appends 'a', node_b appends 'b'.""" + + def node_a(state: SimpleState) -> dict: + return {"value": state["value"] + "A", "items": ["a"]} + + def node_b(state: SimpleState) -> dict: + return {"value": state["value"] + "B", "items": ["b"]} + + builder = StateGraph(SimpleState) + builder.add_node("node_a", node_a) + builder.add_node("node_b", node_b) + builder.add_edge(START, "node_a") + builder.add_edge("node_a", "node_b") + builder.add_edge("node_b", END) + return builder.compile() + + +def _build_interrupt_graph(): + """Graph that interrupts before node_b.""" + + def node_a(state: SimpleState) -> dict: + return {"value": state["value"] + "A", "items": ["a"]} + + def node_b(state: SimpleState) -> dict: + interrupt("need approval") + return {"value": state["value"] + "B", "items": ["b"]} + + builder = StateGraph(SimpleState) + builder.add_node("node_a", node_a) + builder.add_node("node_b", node_b) + builder.add_edge(START, "node_a") + builder.add_edge("node_a", "node_b") + builder.add_edge("node_b", END) + return builder.compile(checkpointer=InMemorySaver()) + + +def _build_error_graph(): + """Graph where node_b raises.""" + + def node_a(state: SimpleState) -> dict: + return {"value": state["value"] + "A", "items": ["a"]} + + def node_b(state: SimpleState) -> dict: + raise ValueError("boom") + + builder = StateGraph(SimpleState) + builder.add_node("node_a", node_a) + builder.add_node("node_b", node_b) + builder.add_edge(START, "node_a") + builder.add_edge("node_a", "node_b") + builder.add_edge("node_b", END) + return builder.compile() + + +def _build_custom_stream_graph(): + """Graph that emits custom stream events.""" + + def node_a(state: SimpleState, *, writer: StreamWriter) -> dict: + writer({"step": "start"}) + writer({"step": "end"}) + return {"value": state["value"] + "A", "items": ["a"]} + + builder = StateGraph(SimpleState) + builder.add_node("node_a", node_a) + builder.add_edge(START, "node_a") + builder.add_edge("node_a", END) + return builder.compile() + + +# --------------------------------------------------------------------------- +# EventLog unit tests +# --------------------------------------------------------------------------- + + +class TestEventLog: + def test_sync_iteration(self) -> None: + log: EventLog[int] = EventLog() + log.push(1) + log.push(2) + log.push(3) + log.close() + assert list(log) == [1, 2, 3] + + def test_multi_cursor(self) -> None: + log: EventLog[str] = EventLog() + log.push("a") + log.push("b") + log.close() + # Two independent cursors see all items. + assert list(log) == ["a", "b"] + assert list(log) == ["a", "b"] + + def test_fail_propagation(self) -> None: + log: EventLog[int] = EventLog() + log.push(1) + log.fail(ValueError("test error")) + with pytest.raises(ValueError, match="test error"): + list(log) + + @pytest.mark.anyio + async def test_async_iteration(self) -> None: + log: EventLog[int] = EventLog() + + async def producer(): + for i in range(3): + log.push(i) + log.close() + + asyncio.get_event_loop().call_soon(lambda: asyncio.ensure_future(producer())) + items = [item async for item in log] + assert items == [0, 1, 2] + + @pytest.mark.anyio + async def test_async_multi_cursor(self) -> None: + log: EventLog[str] = EventLog() + log.push("x") + log.push("y") + log.close() + items1 = [item async for item in log] + items2 = [item async for item in log] + assert items1 == ["x", "y"] + assert items2 == ["x", "y"] + + @pytest.mark.anyio + async def test_async_fail(self) -> None: + log: EventLog[int] = EventLog() + log.push(1) + log.fail(RuntimeError("async error")) + with pytest.raises(RuntimeError, match="async error"): + async for _ in log: + pass + + def test_sync_cursor_yields_items_before_error(self) -> None: + """Sync cursor should yield all buffered items before raising.""" + log: EventLog[int] = EventLog() + log.push(1) + log.push(2) + log.push(3) + log.fail(ValueError("late error")) + items: list[int] = [] + with pytest.raises(ValueError, match="late error"): + for item in log: + items.append(item) + assert items == [1, 2, 3] + + @pytest.mark.anyio + async def test_async_cursor_yields_items_before_error(self) -> None: + """Async cursor should yield all buffered items before raising.""" + log: EventLog[int] = EventLog() + log.push(1) + log.push(2) + log.push(3) + log.fail(ValueError("late error")) + items: list[int] = [] + with pytest.raises(ValueError, match="late error"): + async for item in log: + items.append(item) + assert items == [1, 2, 3] + + def test_push_after_close_raises(self) -> None: + """Push after close should raise RuntimeError.""" + log: EventLog[int] = EventLog() + log.push(1) + log.close() + with pytest.raises(RuntimeError, match="Cannot push to a closed EventLog"): + log.push(2) + + def test_push_after_fail_raises(self) -> None: + """Fail closes the log, so push after fail should also raise.""" + log: EventLog[int] = EventLog() + log.fail(ValueError("err")) + with pytest.raises(RuntimeError, match="Cannot push to a closed EventLog"): + log.push(1) + + def test_empty_log_sync(self) -> None: + """Iterating a closed empty log should yield nothing.""" + log: EventLog[int] = EventLog() + log.close() + assert list(log) == [] + + @pytest.mark.anyio + async def test_empty_log_async(self) -> None: + """Async-iterating a closed empty log should yield nothing.""" + log: EventLog[int] = EventLog() + log.close() + assert [item async for item in log] == [] + + def test_empty_log_fail_sync(self) -> None: + """Failing an empty log should raise immediately with no items.""" + log: EventLog[int] = EventLog() + log.fail(ValueError("empty fail")) + with pytest.raises(ValueError, match="empty fail"): + list(log) + + @pytest.mark.anyio + async def test_empty_log_fail_async(self) -> None: + """Failing an empty log should raise immediately with no items (async).""" + log: EventLog[int] = EventLog() + log.fail(ValueError("empty fail")) + with pytest.raises(ValueError, match="empty fail"): + async for _ in log: + pass + + +# --------------------------------------------------------------------------- +# StreamChannel unit tests +# --------------------------------------------------------------------------- + + +class TestStreamChannel: + def test_push_and_iterate(self) -> None: + ch: StreamChannel[str] = StreamChannel("test") + ch.push("a") + ch.push("b") + ch._close() + assert list(ch) == ["a", "b"] + + def test_wire_callback(self) -> None: + forwarded: list[str] = [] + ch: StreamChannel[str] = StreamChannel("test") + ch._wire(lambda item: forwarded.append(item)) + ch.push("x") + ch.push("y") + ch._close() + assert forwarded == ["x", "y"] + assert list(ch) == ["x", "y"] + + def test_fail_propagation(self) -> None: + """_fail() should propagate the error through the underlying log.""" + ch: StreamChannel[str] = StreamChannel("test") + ch.push("a") + ch._fail(ValueError("channel error")) + items: list[str] = [] + with pytest.raises(ValueError, match="channel error"): + for item in ch: + items.append(item) + assert items == ["a"] + + @pytest.mark.anyio + async def test_async_iteration(self) -> None: + """Async iteration should delegate to the inner EventLog.""" + ch: StreamChannel[str] = StreamChannel("test") + ch.push("x") + ch.push("y") + ch._close() + items = [item async for item in ch] + assert items == ["x", "y"] + + def test_push_without_wire(self) -> None: + """Push without a wire callback should still append to the log.""" + ch: StreamChannel[int] = StreamChannel("test") + assert ch._wire_fn is None + ch.push(42) + ch._close() + assert list(ch) == [42] + + +# --------------------------------------------------------------------------- +# StreamingHandler sync tests +# --------------------------------------------------------------------------- + + +class TestStreamingHandlerSync: + def test_values_projection(self) -> None: + graph = _build_simple_graph() + handler = StreamingHandler(graph) + run = handler.stream({"value": "x", "items": []}) + snapshots = list(run.values) + # Should have at least the initial + per-node snapshots. + assert len(snapshots) >= 1 + # Last snapshot should have both nodes' effects. + last = snapshots[-1] + assert "A" in last["value"] + assert "B" in last["value"] + + def test_output(self) -> None: + graph = _build_simple_graph() + handler = StreamingHandler(graph) + run = handler.stream({"value": "x", "items": []}) + output = run.output + assert output is not None + assert output["value"] == "xAB" + assert output["items"] == ["a", "b"] + + def test_raw_event_iteration(self) -> None: + graph = _build_simple_graph() + handler = StreamingHandler(graph) + run = handler.stream({"value": "x", "items": []}) + events = list(run) + assert len(events) > 0 + for event in events: + assert event["type"] == "event" + assert "method" in event + assert "seq" in event + + def test_extensions_has_native_keys(self) -> None: + graph = _build_simple_graph() + handler = StreamingHandler(graph) + run = handler.stream({"value": "x", "items": []}) + # Drain events so the run completes. + _ = run.output + assert "values" in run.extensions + assert "messages" in run.extensions + # Native keys should also be direct attributes. + assert run.values is run.extensions["values"] + assert run.messages is run.extensions["messages"] + + def test_custom_stream_events(self) -> None: + graph = _build_custom_stream_graph() + handler = StreamingHandler(graph) + run = handler.stream({"value": "x", "items": []}) + custom_events = [e for e in run if e["method"] == "custom"] + assert len(custom_events) == 2 + assert custom_events[0]["params"]["data"] == {"step": "start"} + assert custom_events[1]["params"]["data"] == {"step": "end"} + + +class TestStreamingHandlerSyncErrors: + def test_error_propagation_output(self) -> None: + graph = _build_error_graph() + handler = StreamingHandler(graph) + run = handler.stream({"value": "x", "items": []}) + with pytest.raises(ValueError, match="boom"): + _ = run.output + + def test_error_propagation_values(self) -> None: + graph = _build_error_graph() + handler = StreamingHandler(graph) + run = handler.stream({"value": "x", "items": []}) + with pytest.raises(ValueError, match="boom"): + list(run.values) + + def test_error_propagation_raw_events(self) -> None: + graph = _build_error_graph() + handler = StreamingHandler(graph) + run = handler.stream({"value": "x", "items": []}) + with pytest.raises(ValueError, match="boom"): + list(run) + + + +class TestStreamingHandlerSyncInterrupt: + def test_interrupted(self) -> None: + graph = _build_interrupt_graph() + handler = StreamingHandler(graph) + run = handler.stream( + {"value": "x", "items": []}, + {"configurable": {"thread_id": "t1"}}, + ) + _ = run.output + assert run.interrupted is True + assert len(run.interrupts) > 0 + + +# --------------------------------------------------------------------------- +# StreamingHandler async tests +# --------------------------------------------------------------------------- + + +class TestStreamingHandlerAsync: + @pytest.mark.anyio + @NEEDS_CONTEXTVARS + async def test_values_projection(self) -> None: + graph = _build_simple_graph() + handler = StreamingHandler(graph) + run = await handler.astream({"value": "x", "items": []}) + snapshots = [s async for s in run.values] + assert len(snapshots) >= 1 + last = snapshots[-1] + assert "A" in last["value"] + assert "B" in last["value"] + + @pytest.mark.anyio + @NEEDS_CONTEXTVARS + async def test_output(self) -> None: + graph = _build_simple_graph() + handler = StreamingHandler(graph) + run = await handler.astream({"value": "x", "items": []}) + output = await run.output + assert output is not None + assert output["value"] == "xAB" + assert output["items"] == ["a", "b"] + + @pytest.mark.anyio + @NEEDS_CONTEXTVARS + async def test_raw_event_iteration(self) -> None: + graph = _build_simple_graph() + handler = StreamingHandler(graph) + run = await handler.astream({"value": "x", "items": []}) + events = [e async for e in run] + assert len(events) > 0 + for event in events: + assert event["type"] == "event" + + @pytest.mark.anyio + @NEEDS_CONTEXTVARS + async def test_extensions_has_native_keys(self) -> None: + graph = _build_simple_graph() + handler = StreamingHandler(graph) + run = await handler.astream({"value": "x", "items": []}) + _ = await run.output + assert "values" in run.extensions + assert "messages" in run.extensions + assert run.values is run.extensions["values"] + assert run.messages is run.extensions["messages"] + + +class TestStreamingHandlerAsyncErrors: + @pytest.mark.anyio + @NEEDS_CONTEXTVARS + async def test_error_propagation_output(self) -> None: + graph = _build_error_graph() + handler = StreamingHandler(graph) + run = await handler.astream({"value": "x", "items": []}) + with pytest.raises(ValueError, match="boom"): + await run.output + + @pytest.mark.anyio + @NEEDS_CONTEXTVARS + async def test_error_propagation_values(self) -> None: + graph = _build_error_graph() + handler = StreamingHandler(graph) + run = await handler.astream({"value": "x", "items": []}) + with pytest.raises(ValueError, match="boom"): + async for _ in run.values: + pass + + @pytest.mark.anyio + @NEEDS_CONTEXTVARS + async def test_error_propagation_raw_events(self) -> None: + graph = _build_error_graph() + handler = StreamingHandler(graph) + run = await handler.astream({"value": "x", "items": []}) + with pytest.raises(ValueError, match="boom"): + async for _ in run: + pass + + +class TestStreamingHandlerAsyncInterrupt: + @pytest.mark.anyio + @NEEDS_CONTEXTVARS + async def test_interrupted(self) -> None: + graph = _build_interrupt_graph() + handler = StreamingHandler(graph) + run = await handler.astream( + {"value": "x", "items": []}, + {"configurable": {"thread_id": "t2"}}, + ) + _ = await run.output + assert run.interrupted is True + assert len(run.interrupts) > 0 + + +class TestStreamingHandlerAsyncCustom: + @pytest.mark.anyio + @NEEDS_CONTEXTVARS + async def test_custom_stream_events(self) -> None: + graph = _build_custom_stream_graph() + handler = StreamingHandler(graph) + run = await handler.astream({"value": "x", "items": []}) + events = [e async for e in run] + custom_events = [e for e in events if e["method"] == "custom"] + assert len(custom_events) == 2 + assert custom_events[0]["params"]["data"] == {"step": "start"} + assert custom_events[1]["params"]["data"] == {"step": "end"} + + +# --------------------------------------------------------------------------- +# Custom transformer tests +# --------------------------------------------------------------------------- + + +# --------------------------------------------------------------------------- +# convert_to_protocol_event unit tests +# --------------------------------------------------------------------------- + + +class TestConvertToProtocolEvent: + def test_basic_conversion(self) -> None: + part = {"type": "values", "ns": ("sub", "graph"), "data": {"key": "val"}} + event = convert_to_protocol_event(part) + assert event["type"] == "event" + assert event["method"] == "values" + assert event["params"]["namespace"] == ["sub", "graph"] + assert event["params"]["data"] == {"key": "val"} + assert "interrupts" not in event["params"] + + def test_conversion_with_interrupts(self) -> None: + part = { + "type": "values", + "ns": (), + "data": {"k": 1}, + "interrupts": ({"value": "pause"},), + } + event = convert_to_protocol_event(part) + assert event["params"]["interrupts"] == ({"value": "pause"},) + + def test_namespace_tuple_becomes_list(self) -> None: + """ns tuple should be converted to a list.""" + part = {"type": "updates", "ns": ("a", "b", "c"), "data": {}} + event = convert_to_protocol_event(part) + assert isinstance(event["params"]["namespace"], list) + assert event["params"]["namespace"] == ["a", "b", "c"] + + +# --------------------------------------------------------------------------- +# StreamMux unit tests +# --------------------------------------------------------------------------- + + +class TestStreamMux: + def test_register_non_dict_raises(self) -> None: + """init() returning a non-dict should raise TypeError.""" + + class BadTransformer(StreamTransformer): + def init(self) -> Any: + return ["not", "a", "dict"] + + def process(self, event: ProtocolEvent) -> bool: + return True + + mux = StreamMux() + with pytest.raises(TypeError, match="must return a dict"): + mux.register(BadTransformer()) + + def test_event_suppression(self) -> None: + """When process() returns False, the event should not appear in the main log.""" + + class FilterTransformer(StreamTransformer): + def init(self) -> dict[str, Any]: + return {} + + def process(self, event: ProtocolEvent) -> bool: + # Suppress "updates" events + return event["method"] != "updates" + + mux = StreamMux() + mux.register(FilterTransformer()) + + mux.push( + { + "type": "event", + "method": "values", + "params": {"namespace": [], "data": {"a": 1}}, + } + ) + mux.push( + { + "type": "event", + "method": "updates", + "params": {"namespace": [], "data": {"b": 2}}, + } + ) + mux.push( + { + "type": "event", + "method": "custom", + "params": {"namespace": [], "data": {"c": 3}}, + } + ) + mux.close() + + events = list(mux._events) + methods = [e["method"] for e in events] + assert "updates" not in methods + assert methods == ["values", "custom"] + + def test_suppression_partial_transformers(self) -> None: + """If any transformer returns False, the event is suppressed, + but all transformers still see it.""" + + seen_by_second: list[str] = [] + + class PassTransformer(StreamTransformer): + def init(self) -> dict[str, Any]: + return {} + + def process(self, event: ProtocolEvent) -> bool: + return True + + class RejectTransformer(StreamTransformer): + def init(self) -> dict[str, Any]: + return {} + + def process(self, event: ProtocolEvent) -> bool: + seen_by_second.append(event["method"]) + return False + + mux = StreamMux() + mux.register(PassTransformer()) + mux.register(RejectTransformer()) + + mux.push( + { + "type": "event", + "method": "values", + "params": {"namespace": [], "data": {}}, + } + ) + mux.close() + + # RejectTransformer saw the event even though it rejected it + assert seen_by_second == ["values"] + # But nothing in the main log + assert list(mux._events) == [] + + def test_empty_mux(self) -> None: + """Push/close/fail on a mux with no transformers should work.""" + mux = StreamMux() + mux.push( + { + "type": "event", + "method": "values", + "params": {"namespace": [], "data": {"x": 1}}, + } + ) + mux.close() + events = list(mux._events) + assert len(events) == 1 + assert events[0]["method"] == "values" + + def test_empty_mux_fail(self) -> None: + """Fail on an empty mux should propagate to the event log.""" + mux = StreamMux() + mux.fail(ValueError("boom")) + with pytest.raises(ValueError, match="boom"): + list(mux._events) + + +# --------------------------------------------------------------------------- +# ValuesTransformer / MessagesTransformer unit tests +# --------------------------------------------------------------------------- + + +class TestValuesTransformer: + def test_ignores_non_root_namespace(self) -> None: + """Values events from subgraphs (non-empty namespace) should be ignored.""" + t = ValuesTransformer() + t.init() + + # Root namespace — should be captured + root_event: ProtocolEvent = { + "type": "event", + "method": "values", + "params": {"namespace": [], "data": {"val": "root"}}, + } + t.process(root_event) + + # Subgraph namespace — should be ignored + sub_event: ProtocolEvent = { + "type": "event", + "method": "values", + "params": {"namespace": ["sub"], "data": {"val": "sub"}}, + } + t.process(sub_event) + + t.finalize() + items = list(t._log) + assert len(items) == 1 + assert items[0]["val"] == "root" + + def test_ignores_non_values_methods(self) -> None: + """Non-values events should be passed through but not captured.""" + t = ValuesTransformer() + t.init() + + updates_event: ProtocolEvent = { + "type": "event", + "method": "updates", + "params": {"namespace": [], "data": {"x": 1}}, + } + result = t.process(updates_event) + assert result is True # passed through + t.finalize() + assert list(t._log) == [] # but not captured + + def test_tracks_interrupts(self) -> None: + """Interrupts should be accumulated across events.""" + t = ValuesTransformer() + t.init() + + event: ProtocolEvent = { + "type": "event", + "method": "values", + "params": { + "namespace": [], + "data": {"v": 1}, + "interrupts": ({"value": "pause1"}, {"value": "pause2"}), + }, + } + t.process(event) + assert t._interrupted is True + assert len(t._interrupts) == 2 + + +class TestMessagesTransformer: + def test_captures_root_messages(self) -> None: + t = MessagesTransformer() + t.init() + + event: ProtocolEvent = { + "type": "event", + "method": "messages", + "params": {"namespace": [], "data": ("chunk", {"meta": True})}, + } + t.process(event) + t.finalize() + items = list(t._log) + assert len(items) == 1 + assert items[0] == ("chunk", {"meta": True}) + + def test_ignores_non_root_namespace(self) -> None: + t = MessagesTransformer() + t.init() + + event: ProtocolEvent = { + "type": "event", + "method": "messages", + "params": {"namespace": ["sub"], "data": ("chunk", {})}, + } + t.process(event) + t.finalize() + assert list(t._log) == [] + + def test_ignores_non_messages_methods(self) -> None: + t = MessagesTransformer() + t.init() + + event: ProtocolEvent = { + "type": "event", + "method": "values", + "params": {"namespace": [], "data": {"v": 1}}, + } + result = t.process(event) + assert result is True + t.finalize() + assert list(t._log) == [] + + def test_fail_propagates(self) -> None: + t = MessagesTransformer() + t.init() + t.fail(ValueError("msg error")) + with pytest.raises(ValueError, match="msg error"): + list(t._log) + + +# --------------------------------------------------------------------------- +# StreamMux resilience tests +# --------------------------------------------------------------------------- + + +class TestStreamMuxResilience: + """StreamMux.close() and fail() must complete cleanup even if a transformer raises.""" + + def test_close_continues_after_finalize_error(self) -> None: + """If a transformer's finalize() raises, the main event log and + remaining transformers should still be closed/finalized.""" + from langgraph.stream._mux import StreamMux + + class BrokenFinalizer(StreamTransformer): + def init(self) -> dict[str, Any]: + return {} + + def process(self, event: ProtocolEvent) -> bool: + return True + + def finalize(self) -> None: + raise RuntimeError("finalize broke") + + class GoodTransformer(StreamTransformer): + def __init__(self) -> None: + self.finalized = False + + def init(self) -> dict[str, Any]: + return {} + + def process(self, event: ProtocolEvent) -> bool: + return True + + def finalize(self) -> None: + self.finalized = True + + mux = StreamMux() + mux.register(BrokenFinalizer()) + good = GoodTransformer() + mux.register(good) + + mux.push( + { + "type": "event", + "method": "values", + "params": {"namespace": [], "data": {}}, + } + ) + + with pytest.raises(RuntimeError, match="finalize broke"): + mux.close() + + assert good.finalized + assert mux._events._closed + + def test_fail_continues_after_transformer_error(self) -> None: + """If a transformer's fail() raises, the main event log and + remaining transformers should still be failed.""" + from langgraph.stream._mux import StreamMux + + class BrokenFailer(StreamTransformer): + def init(self) -> dict[str, Any]: + return {} + + def process(self, event: ProtocolEvent) -> bool: + return True + + def fail(self, err: BaseException) -> None: + raise RuntimeError("fail handler broke") + + class GoodTransformer(StreamTransformer): + def __init__(self) -> None: + self.failed_with: BaseException | None = None + + def init(self) -> dict[str, Any]: + return {} + + def process(self, event: ProtocolEvent) -> bool: + return True + + def fail(self, err: BaseException) -> None: + self.failed_with = err + + mux = StreamMux() + mux.register(BrokenFailer()) + good = GoodTransformer() + mux.register(good) + + original_error = ValueError("original") + mux.fail(original_error) + + assert good.failed_with is original_error + assert mux._events._error is original_error + + def test_close_still_closes_channels_after_finalize_error(self) -> None: + """Channels should be closed even if a transformer's finalize raises.""" + from langgraph.stream._mux import StreamMux + + class BrokenWithChannel(StreamTransformer): + def __init__(self) -> None: + self._channel: StreamChannel[str] = StreamChannel("ch") + + def init(self) -> dict[str, Any]: + return {"ch": self._channel} + + def process(self, event: ProtocolEvent) -> bool: + return True + + def finalize(self) -> None: + raise RuntimeError("finalize broke") + + t = BrokenWithChannel() + mux = StreamMux() + mux.register(t) + + with pytest.raises(RuntimeError, match="finalize broke"): + mux.close() + + assert t._channel._log._closed + + +class TestCustomTransformer: + def test_extension_transformer_with_stream_channel(self) -> None: + """User transformer with StreamChannel appears in extensions.""" + + class CounterTransformer(StreamTransformer): + def __init__(self) -> None: + super().__init__() + self._channel: StreamChannel[int] = StreamChannel("counter") + self._count = 0 + + def init(self) -> dict[str, Any]: + return {"counter": self._channel} + + def process(self, event: ProtocolEvent) -> bool: + if event["method"] == "values": + self._count += 1 + self._channel.push(self._count) + return True + + graph = _build_simple_graph() + handler = StreamingHandler(graph) + counter_t = CounterTransformer() + run = handler.stream({"value": "x", "items": []}, transformers=[counter_t]) + _ = run.output + assert "counter" in run.extensions + # Counter channel should have been pushed to. + counts = list(run.extensions["counter"]) + assert len(counts) > 0 + # Non-native transformer should not set direct attributes. + assert not hasattr(run, "counter") + + def test_native_transformer_gets_direct_attr(self) -> None: + """A transformer with _native=True gets its keys as run attributes.""" + + class FooTransformer(StreamTransformer): + _native = True + + def __init__(self) -> None: + super().__init__() + self._log: EventLog[str] = EventLog() + + def init(self) -> dict[str, Any]: + return {"foo": self._log} + + def process(self, event: ProtocolEvent) -> bool: + if event["method"] == "values": + self._log.push("saw_values") + return True + + def finalize(self) -> None: + self._log.close() + + def fail(self, err: BaseException) -> None: + self._log.fail(err) + + graph = _build_simple_graph() + handler = StreamingHandler(graph) + foo_t = FooTransformer() + run = handler.stream({"value": "x", "items": []}, transformers=[foo_t]) + _ = run.output + # foo should be both in extensions and as a direct attribute. + assert "foo" in run.extensions + assert hasattr(run, "foo") + assert run.foo is run.extensions["foo"] + items = list(run.foo) + assert "saw_values" in items + + def test_stream_channel_auto_forward(self) -> None: + """StreamChannel pushes inject ProtocolEvents into main log.""" + + class EmitterTransformer(StreamTransformer): + def __init__(self) -> None: + super().__init__() + self._channel: StreamChannel[str] = StreamChannel("emitter") + + def init(self) -> dict[str, Any]: + return {"emitter": self._channel} + + def process(self, event: ProtocolEvent) -> bool: + if event["method"] == "values": + self._channel.push("emitted") + return True + + graph = _build_simple_graph() + handler = StreamingHandler(graph) + run = handler.stream( + {"value": "x", "items": []}, transformers=[EmitterTransformer()] + ) + events = list(run) + custom_events = [e for e in events if e["method"] == "custom:emitter"] + assert len(custom_events) > 0 + assert custom_events[0]["params"]["data"] == "emitted" + + def test_stream_channel_seq_ordering(self) -> None: + """Seq numbers in the main event log must be monotonically increasing. + + When a transformer pushes to a StreamChannel during process(), the + auto-forwarded event enters the main log before the original event. + The seq numbers must still be in order. + """ + from langgraph.stream._mux import StreamMux + + class ChannelPusher(StreamTransformer): + def __init__(self) -> None: + super().__init__() + self._channel: StreamChannel[str] = StreamChannel("ch") + + def init(self) -> dict[str, Any]: + return {"ch": self._channel} + + def process(self, event: ProtocolEvent) -> bool: + # Push to channel during process — this triggers auto-forward + # which injects an event into the main log mid-pipeline. + self._channel.push(f"saw:{event['method']}") + return True + + mux = StreamMux() + mux.register(ChannelPusher()) + + mux.push( + { + "type": "event", + "method": "values", + "params": {"namespace": [], "data": {}}, + } + ) + mux.push( + { + "type": "event", + "method": "updates", + "params": {"namespace": [], "data": {}}, + } + ) + mux.close() + + events = list(mux._events) + seqs = [e["seq"] for e in events] + # Seq numbers must be strictly increasing. + for i in range(1, len(seqs)): + assert seqs[i] > seqs[i - 1], f"Seq out of order at index {i}: {seqs}"