diff --git a/libs/langgraph/langgraph/pregel/_lifecycle.py b/libs/langgraph/langgraph/pregel/_lifecycle.py new file mode 100644 index 000000000..6f2beb07c --- /dev/null +++ b/libs/langgraph/langgraph/pregel/_lifecycle.py @@ -0,0 +1,245 @@ +from __future__ import annotations + +from collections.abc import AsyncIterator, Callable, Iterator +from typing import Any, TypeVar, cast +from uuid import UUID + +from langchain_core.callbacks import BaseCallbackHandler + +from langgraph._internal._constants import NS_SEP +from langgraph.errors import GraphInterrupt +from langgraph.pregel.protocol import StreamChunk + +try: + from langchain_core.tracers._streaming import _StreamingCallbackHandler +except ImportError: + _StreamingCallbackHandler = object # type: ignore[assignment,misc] + + +T = TypeVar("T") + + +_LANGGRAPH_SENTINEL_NODES = frozenset({"__start__", "__end__"}) + + +def _is_nested_pregel_start(name: str | None, metadata: dict[str, Any] | None) -> bool: + """Recognize a nested `Pregel` invocation from its `on_chain_start` metadata. + + When a compiled graph is added as a node, pregel fires two + `on_chain_start` callbacks at that task: first for the node chain + (whose `name` matches `metadata["langgraph_node"]`) and second for + the inner `Pregel` chain (whose `name` is the graph's `name`, not + the node name). Both share the same `langgraph_checkpoint_ns`. + + A nested `Pregel` start is therefore identified by having a + `langgraph_checkpoint_ns` AND a `name` that does NOT match the + owning task's `langgraph_node`. Regular node chains are skipped; + the root `Pregel` (which has no `langgraph_node` metadata) isn't + observed by this handler because the root's start fires before the + handler is attached. + + Metadata-based detection is used because `on_chain_start`'s + `serialized` argument is `None` for compiled graphs in this + version of langchain-core, so class-based detection via + `serialized["id"]` isn't available. + + Sentinel nodes (`__start__` / `__end__`) are excluded: conditional + edges from `START` fire an `on_chain_start` with `lg_node=__start__` + and the router function's name as `name`, which would otherwise + match the discriminator without representing an actual nested + `Pregel`. + """ + if not metadata: + return False + if not metadata.get("langgraph_checkpoint_ns"): + return False + lg_node = metadata.get("langgraph_node") + if lg_node is None or lg_node in _LANGGRAPH_SENTINEL_NODES: + return False + return name != lg_node + + +class StreamLifecycleHandler(BaseCallbackHandler, _StreamingCallbackHandler): + """Callback handler that emits subgraph lifecycle events on the stream. + + Pushes `LifecycleData`-shaped payloads onto the pregel stream under + the `"lifecycle"` mode, keyed by the subgraph's namespace tuple. + Drives the `started` → `running` → `completed` / `failed` / + `interrupted` state machine. + + The handler is attached to `run_manager.inheritable_handlers` inside + a `Pregel.stream` / `astream` call, so it sees callbacks for every + descendant chain (nodes, nested `Pregel` subgraphs) but *not* for + the root `Pregel` whose start event has already fired. The root's + `started` event is emitted eagerly at construction; its terminal + state is emitted by `SubgraphTransformer.finalize` / `fail`. + + `run_inline = True` keeps event ordering deterministic. + """ + + run_inline = True + + def __init__( + self, + stream: Callable[[StreamChunk], None], + *, + root_graph_name: str | None = None, + ) -> None: + """Initialize the handler and emit the root graph's `started` event. + + Args: + stream: Callable that accepts a `StreamChunk` tuple + `(namespace, mode, payload)` and enqueues it. + root_graph_name: The root `Pregel` instance's `name`, emitted + with the root's `started` lifecycle payload. + """ + self.stream = stream + # Namespaces awaiting the started→running transition. + self._pending_running: set[tuple[str, ...]] = set() + # run_id → subgraph namespace; populated only for Pregel chains. + self._run_to_ns: dict[UUID, tuple[str, ...]] = {} + + root_payload: dict[str, Any] = {"event": "started"} + if root_graph_name is not None: + root_payload["graph_name"] = root_graph_name + self.stream(((), "lifecycle", root_payload)) + self._pending_running.add(()) + + @staticmethod + def _subgraph_ns_from_metadata(metadata: dict[str, Any] | None) -> tuple[str, ...]: + """Return the running subgraph's own namespace from task metadata. + + For a nested `Pregel` invoked as a node, `langgraph_checkpoint_ns` + ends at the node segment (no inner task appended yet), so + splitting on `NS_SEP` gives the subgraph's own namespace. + """ + if not metadata: + return () + nskey = metadata.get("langgraph_checkpoint_ns") + if not nskey: + return () + return tuple(cast(str, nskey).split(NS_SEP)) + + @staticmethod + def _containing_ns_from_metadata( + metadata: dict[str, Any] | None, + ) -> tuple[str, ...]: + """Return the namespace of the subgraph that contains this task. + + For an inner task with `langgraph_checkpoint_ns` + `"seg_a|seg_b"`, the containing subgraph is `("seg_a",)`. + """ + if not metadata: + return () + nskey = metadata.get("langgraph_checkpoint_ns") + if not nskey: + return () + return tuple(cast(str, nskey).split(NS_SEP))[:-1] + + @staticmethod + def _trigger_call_id(metadata: dict[str, Any] | None) -> str | None: + """Extract `trigger_call_id` from task metadata if present. + + The task that spawned a nested `Pregel` has its task id encoded + in `langgraph_checkpoint_ns`'s last segment as + `node_name:task_id`. Returns the `task_id` portion, which + parents can correlate with their `tools` / `tasks` events. + """ + if not metadata: + return None + nskey = cast(str | None, metadata.get("langgraph_checkpoint_ns")) + if not nskey: + return None + last = nskey.split(NS_SEP)[-1] + _, sep, task_id = last.rpartition(":") + return task_id if sep else None + + def _emit(self, ns: tuple[str, ...], payload: dict[str, Any]) -> None: + self.stream((ns, "lifecycle", payload)) + + def tap_output_aiter( + self, run_id: UUID, output: AsyncIterator[T] + ) -> AsyncIterator[T]: + """Pass-through — required by the `_StreamingCallbackHandler` protocol. + + Returns the iterator unchanged. A missing implementation lets + langchain's default `Protocol` body return `None`, which breaks + the `_consume_aiter` code path in `_runnable.py:900`. + """ + return output + + def tap_output_iter(self, run_id: UUID, output: Iterator[T]) -> Iterator[T]: + """Pass-through — sync counterpart to `tap_output_aiter`.""" + return output + + def _fire_running_if_pending(self, ns: tuple[str, ...]) -> None: + if ns in self._pending_running: + self._pending_running.discard(ns) + self._emit(ns, {"event": "running"}) + + def on_chain_start( + self, + serialized: dict[str, Any], + inputs: dict[str, Any], + *, + run_id: UUID, + parent_run_id: UUID | None = None, + tags: list[str] | None = None, + metadata: dict[str, Any] | None = None, + **kwargs: Any, + ) -> Any: + # Any descendant activity transitions the containing subgraph to running. + containing = self._containing_ns_from_metadata(metadata) + self._fire_running_if_pending(containing) + + name = cast(str | None, kwargs.get("name")) + if not _is_nested_pregel_start(name, metadata): + return + + ns = self._subgraph_ns_from_metadata(metadata) + if not ns: + return + + self._run_to_ns[run_id] = ns + payload: dict[str, Any] = {"event": "started"} + if name: + payload["graph_name"] = name + trigger_call_id = self._trigger_call_id(metadata) + if trigger_call_id: + payload["trigger_call_id"] = trigger_call_id + self._emit(ns, payload) + self._pending_running.add(ns) + + def on_chain_end( + self, + response: Any, + *, + run_id: UUID, + parent_run_id: UUID | None = None, + **kwargs: Any, + ) -> Any: + ns = self._run_to_ns.pop(run_id, None) + if ns is None: + return + # Ensure started→running fired even for empty subgraphs. + if ns in self._pending_running: + self._pending_running.discard(ns) + self._emit(ns, {"event": "running"}) + self._emit(ns, {"event": "completed"}) + + def on_chain_error( + self, + error: BaseException, + *, + run_id: UUID, + parent_run_id: UUID | None = None, + **kwargs: Any, + ) -> Any: + ns = self._run_to_ns.pop(run_id, None) + if ns is None: + return + self._pending_running.discard(ns) + if isinstance(error, GraphInterrupt): + self._emit(ns, {"event": "interrupted"}) + else: + self._emit(ns, {"event": "failed", "error": str(error)}) diff --git a/libs/langgraph/langgraph/pregel/main.py b/libs/langgraph/langgraph/pregel/main.py index 6c0471970..6d5b7cc48 100644 --- a/libs/langgraph/langgraph/pregel/main.py +++ b/libs/langgraph/langgraph/pregel/main.py @@ -130,6 +130,7 @@ from langgraph.pregel._checkpoint import ( ) from langgraph.pregel._draw import draw_graph from langgraph.pregel._io import map_input, read_channels +from langgraph.pregel._lifecycle import StreamLifecycleHandler from langgraph.pregel._loop import ( AsyncPregelLoop, SyncPregelLoop, @@ -2643,6 +2644,15 @@ class Pregel( ) ) + # set up lifecycle stream mode + if "lifecycle" in stream_modes: + run_manager.inheritable_handlers.append( + StreamLifecycleHandler( + stream.put, + root_graph_name=self.name, + ) + ) + # set up custom stream mode if "custom" in stream_modes: @@ -3025,6 +3035,15 @@ class Pregel( ) ) + # set up lifecycle stream mode + if "lifecycle" in stream_modes: + run_manager.inheritable_handlers.append( + StreamLifecycleHandler( + stream_put, + root_graph_name=self.name, + ) + ) + # set up custom stream mode def stream_writer(c: Any) -> None: aioloop.call_soon_threadsafe( diff --git a/libs/langgraph/langgraph/pregel/remote.py b/libs/langgraph/langgraph/pregel/remote.py index 973768a4a..8919be4cd 100644 --- a/libs/langgraph/langgraph/pregel/remote.py +++ b/libs/langgraph/langgraph/pregel/remote.py @@ -653,19 +653,24 @@ class RemoteGraph(PregelProtocol): # coerce to list, or add default stream mode if stream_mode: if isinstance(stream_mode, str): - updated_stream_modes.append(stream_mode) + if stream_mode != "lifecycle": + updated_stream_modes.append(cast(StreamModeSDK, stream_mode)) else: req_single = False - updated_stream_modes.extend(stream_mode) + updated_stream_modes.extend( + cast(StreamModeSDK, m) for m in stream_mode if m != "lifecycle" + ) else: - updated_stream_modes.append(default) + updated_stream_modes.append(default) # type: ignore[arg-type] requested_stream_modes = updated_stream_modes.copy() # add any from parent graph stream: StreamProtocol | None = ( (config or {}).get(CONF, {}).get(CONFIG_KEY_STREAM) ) if stream: - updated_stream_modes.extend(stream.modes) + updated_stream_modes.extend( + cast(StreamModeSDK, m) for m in stream.modes if m != "lifecycle" + ) # map "messages" to "messages-tuple" if "messages" in updated_stream_modes: updated_stream_modes.remove("messages") diff --git a/libs/langgraph/langgraph/stream/_mux.py b/libs/langgraph/langgraph/stream/_mux.py index e5e911e6d..ea9196942 100644 --- a/libs/langgraph/langgraph/stream/_mux.py +++ b/libs/langgraph/langgraph/stream/_mux.py @@ -2,7 +2,7 @@ from __future__ import annotations import asyncio import time -from collections.abc import Callable +from collections.abc import Awaitable, Callable from typing import Any from langgraph.stream._event_log import EventLog @@ -13,6 +13,17 @@ from langgraph.stream._types import ( ) from langgraph.stream.stream_channel import StreamChannel +TransformerFactory = Callable[["tuple[str, ...]"], StreamTransformer] +"""Factory that builds a scoped transformer for a mux. + +Called once per `StreamMux` (root or mini-mux) with the mux's scope +— typically a subgraph's namespace or `()` for the root. Standard +transformer classes (`ValuesTransformer`, `MessagesTransformer`, +`SubgraphTransformer`) accept a single positional scope argument, so +the class itself is a valid factory. User transformers can close over +their config: `lambda scope: MyTransformer(scope, foo=...)`. +""" + class StreamMux: """Central event dispatcher for the streaming infrastructure. @@ -40,28 +51,50 @@ class StreamMux: transformers: list[StreamTransformer] | None = None, *, is_async: bool = False, + factories: list[TransformerFactory] | None = None, + scope: tuple[str, ...] = (), ) -> None: """Initialize the mux and register transformers in order. - Transformers are fixed at construction time — there is no - post-init `register()`. Each transformer's `init()` is called, + Callers pass either `transformers` (pre-built instances) or + `factories` (callables producing fresh instances per mux). A + factory list is preferred — mini-muxes built by `make_child()` + inherit the factory list, so transformers propagate naturally + into every subgraph's scope. `transformers` is kept for + back-compat tests that exercise the mux directly. + + Each transformer's `init()` is called once during registration, projections are merged into `extensions`, `_native` keys are recorded in `native_keys`, and any EventLog / StreamChannel instances are bound and wired. Args: - transformers: Transformers to register, in dispatch order. - `None` or empty gives a mux with no projections. + transformers: Already-built transformer instances. Mutually + exclusive with `factories`. is_async: True for async dispatch (`apush` / `aclose` / `afail`), False for the sync path. + factories: Zero-or-one-argument callables producing + transformers. Called with this mux's `scope`. + scope: The namespace the mux operates within. The root mux + is `()`; mini-muxes for subgraphs use the subgraph's + namespace tuple. Raises: RuntimeError: If any transformer requires an async run but the mux is in sync mode. TypeError: If a transformer's `init()` doesn't return a dict. - ValueError: If transformers' projection keys collide. + ValueError: If transformers' projection keys collide, or if + both `transformers` and `factories` are supplied. """ + if transformers is not None and factories is not None: + raise ValueError("Pass either `transformers` or `factories`, not both.") + self._is_async = is_async + self._factories: list[TransformerFactory] = list(factories or ()) + self.scope: tuple[str, ...] = scope + self._pump_fn: Callable[[], bool] | None = None + self._apump_fn: Callable[[], Awaitable[bool]] | None = None + self._events: EventLog[ProtocolEvent] = EventLog() self._events._bind(is_async=is_async) self._transformers: list[StreamTransformer] = [] @@ -72,9 +105,80 @@ class StreamMux: self.extensions: dict[str, Any] = {} self.native_keys: set[str] = set() self._projection_owners: dict[str, str] = {} + self._transformer_by_key: dict[str, StreamTransformer] = {} - for transformer in transformers or (): - self._register(transformer) + if factories is not None: + for factory in factories: + self._register(factory(scope)) + else: + for transformer in transformers or (): + self._register(transformer) + + def make_child(self, scope: tuple[str, ...]) -> StreamMux: + """Build a mini-mux with the same factories scoped to `scope`. + + Used by `SubgraphTransformer` to attach a fresh transformer + pipeline to each discovered subgraph handle. The child mux + inherits the current pump binding (so cursors on its projection + logs drive the root pump) and carries the same factory list + forward to any grandchild subgraphs. + + Raises: + RuntimeError: If the mux was not built from a factory list + (i.e., constructed with `transformers=`). Mini-muxes + require factories so each scope gets its own fresh + transformer instances. + """ + if not self._factories: + raise RuntimeError( + "StreamMux.make_child requires the mux to be constructed " + "with factories; pre-built transformers can't be cloned " + "to a new scope." + ) + child = StreamMux( + factories=self._factories, + is_async=self._is_async, + scope=scope, + ) + if self._pump_fn is not None: + child.bind_pump(self._pump_fn) + if self._apump_fn is not None: + child.bind_apump(self._apump_fn) + return child + + def bind_pump(self, fn: Callable[[], bool]) -> None: + """Wire the sync pull callback onto every EventLog in the mux. + + Also propagates to transformers that expose `_bind_pump` so + nested handles (e.g., `ChatModelStream` instances produced by + `MessagesTransformer`) can drive the graph pump from their + projection cursors. + """ + self._pump_fn = fn + self._events._request_more = fn + for value in self.extensions.values(): + if isinstance(value, EventLog): + value._request_more = fn + elif isinstance(value, StreamChannel): + value._log._request_more = fn + for transformer in self._transformers: + bind = getattr(transformer, "_bind_pump", None) + if bind is not None: + bind(fn) + + def bind_apump(self, fn: Callable[[], Awaitable[bool]]) -> None: + """Async counterpart to `bind_pump`.""" + self._apump_fn = fn + self._events._arequest_more = fn + for value in self.extensions.values(): + if isinstance(value, EventLog): + value._arequest_more = fn + elif isinstance(value, StreamChannel): + value._log._arequest_more = fn + for transformer in self._transformers: + abind = getattr(transformer, "_bind_apump", None) + if abind is not None: + abind(fn) def _register(self, transformer: StreamTransformer) -> None: """Register a single transformer. @@ -112,13 +216,27 @@ class StreamMux: owner_name = type(transformer).__name__ for key in projection: self._projection_owners[key] = owner_name + self._transformer_by_key[key] = transformer if getattr(transformer, "_native", False): self.native_keys.update(projection.keys()) + on_register = getattr(transformer, "_on_register", None) + if on_register is not None: + on_register(self) + + def transformer_by_key(self, key: str) -> StreamTransformer | None: + """Return the transformer that owns the projection at `key`, if any.""" + return self._transformer_by_key.get(key) def push(self, event: ProtocolEvent) -> None: """Route an event through all transformers, then append to the main log. - Each transformer's `process()` is called in registration order. + Each transformer's `process()` is called in registration order + — except when the transformer has `scope_exact = True` (the + default) and the event's namespace differs from the mux's + `scope`, in which case the transformer is skipped. Transformers + that need to see cross-scope events opt out by setting + `scope_exact = False` (e.g. `SubgraphTransformer`). + If any transformer returns False, the event is suppressed from the main log, but transformers that already saw it keep their side effects. @@ -132,8 +250,12 @@ class StreamMux: Args: event: The protocol event to dispatch. """ + ns = tuple(event["params"]["namespace"]) + in_scope = ns == self.scope keep = True for transformer in self._transformers: + if transformer.scope_exact and not in_scope: + continue if not transformer.process(event): keep = False if keep: @@ -205,11 +327,13 @@ class StreamMux: """Dispatch an event on the async lane. Awaits each transformer's `aprocess` in registration order - before appending to the main log. A slow `aprocess` serializes - the pipeline by design — that's the guarantee that lets a later - transformer (or a synchronous consumer) see the result of the - async work. For decoupled work, use `schedule()` from inside - `process` / `aprocess` instead. + before appending to the main log — except when the transformer + has `scope_exact = True` and the event's namespace differs from + `self.scope`, in which case it is skipped. A slow `aprocess` + serializes the pipeline by design — that's the guarantee that + lets a later transformer (or a synchronous consumer) see the + result of the async work. For decoupled work, use `schedule()` + from inside `process` / `aprocess` instead. The main log append is a non-blocking `push` — matching v1's `put_nowait` shape. Memory is bounded by caller pace via the @@ -218,8 +342,12 @@ class StreamMux: Args: event: The protocol event to dispatch. """ + ns = tuple(event["params"]["namespace"]) + in_scope = ns == self.scope keep = True for transformer in self._transformers: + if transformer.scope_exact and not in_scope: + continue if not await transformer.aprocess(event): keep = False if keep: diff --git a/libs/langgraph/langgraph/stream/_types.py b/libs/langgraph/langgraph/stream/_types.py index 80fb630d4..797be7e76 100644 --- a/libs/langgraph/langgraph/stream/_types.py +++ b/libs/langgraph/langgraph/stream/_types.py @@ -75,6 +75,20 @@ class StreamTransformer(ABC): scoring, cost lookup, external tracing). Attributes: + scope: Namespace the transformer operates within — `()` for the + root mux, a subgraph's namespace tuple inside a mini-mux. + Set at construction from the mux's scope (each factory is + called as `factory(scope)`). Transformers that only care + about events at their own namespace compare against + `self.scope`; subgraph-aware transformers can treat it as + a parent path. + scope_exact: If True (the default), the mux only calls + `process` / `aprocess` for events whose namespace equals + `self.scope` — user transformers get scope-scoped events + for free with no boilerplate. Set False for transformers + that need to see events across scopes (e.g. + `SubgraphTransformer` forwards deeper events into child + mini-muxes). requires_async: Explicit opt-in for transformers that need a running event loop but don't override any async method (for example, transformers that call `schedule()` from a sync @@ -83,6 +97,18 @@ class StreamTransformer(ABC): """ requires_async: ClassVar[bool] = False + scope_exact: ClassVar[bool] = True + + def __init__(self, scope: tuple[str, ...] = ()) -> None: + """Initialize the transformer with its mux's scope. + + Args: + scope: The namespace tuple the owning mux is scoped to. + `()` for the root, the subgraph's namespace inside a + mini-mux. Factories receive this at construction time + (`factory(scope)` in `StreamMux`). + """ + self.scope: tuple[str, ...] = scope @abstractmethod def init(self) -> dict[str, Any]: diff --git a/libs/langgraph/langgraph/stream/run_stream.py b/libs/langgraph/langgraph/stream/run_stream.py index 3f2d5b0ac..fb3152ff5 100644 --- a/libs/langgraph/langgraph/stream/run_stream.py +++ b/libs/langgraph/langgraph/stream/run_stream.py @@ -3,14 +3,14 @@ from __future__ import annotations import asyncio from collections.abc import AsyncIterator, Awaitable, Callable, Iterator, Mapping from types import MappingProxyType, TracebackType -from typing import Any +from typing import TYPE_CHECKING, 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._types import ProtocolEvent -from langgraph.stream.stream_channel import StreamChannel -from langgraph.stream.transformers import ValuesTransformer + +if TYPE_CHECKING: + from langgraph.stream.transformers import ValuesTransformer def _drive_until_done(pump: Callable[[], bool]) -> None: @@ -25,7 +25,99 @@ async def _adrive_until_done(pump: Callable[[], Awaitable[bool]]) -> None: pass -class GraphRunStream: +class BaseRunStream: + """Shared shape for any object that wraps a `StreamMux`. + + Both the root `GraphRunStream` / `AsyncGraphRunStream` and the + scoped `SubgraphRunStream` compose a `StreamMux` and expose its + projections (`values`, `messages`, `subgraphs`, user-registered + keys). The root additionally owns the graph iterator and drives + the pump; a subgraph's mini-mux borrows the root pump via + `make_child`'s pump inheritance. + + Projections registered on the mux show up in `extensions`, and + native ones (those with `_native = True`) are also bound directly + as attributes (`run.values`, `run.messages`, …). + + Raw iteration (`for event in run:` / `async for event in run`) + yields every `ProtocolEvent` that reached this mux's main log — + for the root that's every event in the run, for a subgraph that's + every event forwarded into its subtree. + """ + + def __init__(self, mux: StreamMux) -> None: + self._mux = mux + self.extensions: Mapping[str, Any] = MappingProxyType(mux.extensions) + for key in mux.native_keys: + setattr(self, key, mux.extensions[key]) + + def __iter__(self) -> Iterator[ProtocolEvent]: + """Sync iteration of protocol events on this mux's main log. + + Raises at the EventLog level if the mux is async-bound. + """ + return iter(self._mux._events) + + def __aiter__(self) -> AsyncIterator[ProtocolEvent]: + """Async iteration of protocol events on this mux's main log. + + Raises at the EventLog level if the mux is sync-bound. + """ + return self._mux._events.__aiter__() + + def interleave(self, *names: str) -> Iterator[tuple[str, Any]]: + """Iterate multiple projections round-robin, yielding ``(name, item)``. + + Each turn advances one projection's cursor; when a cursor's + buffer is empty, pulling from it drives the pump once, which + fans out to every subscribed projection log. Projections whose + items aren't consumed on this turn sit in their own buffers + only until the next turn reaches them, bounding memory by the + skew between projection rates rather than letting any single + log grow to the full run length. + + Projections are exhausted independently; a projection that + finishes early drops out of the rotation while others + continue. The overall iterator ends once all named projections + are done. + + Args: + *names: Projection keys to interleave. Must match keys in + `extensions`. + + Yields: + `(name, item)` tuples in round-robin order across the named + projections. + + Raises: + KeyError: If a name doesn't match a registered projection. + + Example: + ```python + for name, item in run.interleave("messages", "values"): + if name == "messages": + print("msg:", item) + else: + print("val:", item) + ``` + """ + cursors: dict[str, Iterator[Any]] = { + name: iter(self.extensions[name]) for name in names + } + done: set[str] = set() + while len(done) < len(cursors): + for name, cursor in cursors.items(): + if name in done: + continue + try: + item = next(cursor) + except StopIteration: + done.add(name) + continue + yield (name, item) + + +class GraphRunStream(BaseRunStream): """Sync run stream with caller-driven pumping. The caller's iteration on any projection (`values`, `messages`, @@ -34,10 +126,6 @@ class GraphRunStream: Projections are single-consumer — iterating `run.values` twice raises. Use `projection.tee(n)` if you genuinely need fan-out. - - All transformer projections live in `extensions`. Native transformer - projections (those with `_native = True`) are also set as direct - attributes on this instance (e.g. `run.values`, `run.messages`). """ def __init__( @@ -54,40 +142,18 @@ class GraphRunStream: values_transformer: The built-in values transformer providing `output` / `interrupted` / `interrupts`. """ + super().__init__(mux) self._graph_iter = graph_iter - self._mux = mux - self.extensions: Mapping[str, Any] = MappingProxyType(mux.extensions) self._values_transformer = values_transformer self._exhausted = False - for key in mux.native_keys: - setattr(self, key, mux.extensions[key]) - self._wire_request_more(mux) - - def _wire_request_more(self, mux: StreamMux) -> None: - """Install `_request_more` on every sync EventLog so cursors can - drive the pump when their buffer catches up. - - Also calls `_bind_pump` on any transformer that exposes it, so - transformers producing ChatModelStream objects (e.g. - MessagesTransformer) can wire the pull callback on each stream - as it's created. - """ - mux._events._request_more = self._pump_next - for value in mux.extensions.values(): - if isinstance(value, EventLog): - value._request_more = self._pump_next - elif isinstance(value, StreamChannel): - value._log._request_more = self._pump_next - for transformer in mux._transformers: - if hasattr(transformer, "_bind_pump"): - transformer._bind_pump(self._pump_next) + mux.bind_pump(self._pump_next) def _pump_next(self) -> bool: """Pull one event from the graph and push it through the mux. Returns: - True if an event was pulled, False if the graph is exhausted - or has raised. + True if an event was pulled, False if the graph is + exhausted or has raised. """ if self._exhausted: return False @@ -141,8 +207,7 @@ class GraphRunStream: @property def interrupted(self) -> bool: - """Drive the run to completion, then return whether it was - interrupted. + """Drive the run to completion, then return whether it was interrupted. Raises: BaseException: If the run ended with an error. @@ -166,70 +231,17 @@ class GraphRunStream: raise err return self._values_transformer._interrupts - def __iter__(self) -> Iterator[ProtocolEvent]: - """Subscribe to the main event log and iterate protocol events.""" - return iter(self._mux._events) - def interleave(self, *names: str) -> Iterator[tuple[str, Any]]: - """Iterate multiple projections round-robin, yielding ``(name, item)``. - - Each turn advances one projection's cursor; when a cursor's buffer - is empty, pulling from it drives the pump once, which fans out to - every subscribed projection log. Projections whose items aren't - consumed on this turn sit in their own buffers only until the next - turn reaches them, bounding memory by the skew between projection - rates rather than letting any single log grow to the full run - length. - - Projections are exhausted independently; a projection that finishes - early drops out of the rotation while others continue. The overall - iterator ends once all named projections are done. - - Args: - *names: Projection keys to interleave. Must match keys in - ``extensions``. - - Yields: - ``(name, item)`` tuples in round-robin order across the named - projections. - - Raises: - KeyError: If a name doesn't match a registered projection. - - Example: - ```python - for name, item in run.interleave("messages", "values"): - if name == "messages": - print("msg:", item) - else: - print("val:", item) - ``` - """ - cursors: dict[str, Iterator[Any]] = { - name: iter(self.extensions[name]) for name in names - } - done: set[str] = set() - while len(done) < len(cursors): - for name, cursor in cursors.items(): - if name in done: - continue - try: - item = next(cursor) - except StopIteration: - done.add(name) - continue - yield (name, item) - - -class AsyncGraphRunStream: +class AsyncGraphRunStream(BaseRunStream): """Async run stream with caller-driven pumping. Async iteration on any projection drives the graph forward — there is no background task. Concurrent consumers share a single-flight - pump via an `asyncio.Lock`, so each awaiting cursor contributes one - event per acquisition. Backpressure comes from the logs: when a - subscribed log's buffer reaches `maxlen`, `apush` awaits the - subscriber to drain, which holds back the pump and paces the graph. + pump via an `asyncio.Lock`, so each awaiting cursor contributes + one event per acquisition. Backpressure comes from the logs: when + a subscribed log's buffer reaches `maxlen`, `apush` awaits the + subscriber to drain, which holds back the pump and paces the + graph. Projections are single-consumer — a second `aiter(run.values)` raises. Use `projection.tee(n)` for fan-out. @@ -258,35 +270,12 @@ class AsyncGraphRunStream: values_transformer: The built-in values transformer providing `output` / `interrupted` / `interrupts`. """ + super().__init__(mux) self._graph_aiter = graph_aiter - self._mux = mux - self.extensions: Mapping[str, Any] = MappingProxyType(mux.extensions) self._values_transformer = values_transformer self._exhausted = False self._pump_lock = asyncio.Lock() - for key in mux.native_keys: - setattr(self, key, mux.extensions[key]) - self._wire_arequest_more(mux) - - def _wire_arequest_more(self, mux: StreamMux) -> None: - """Install `_arequest_more` on every async EventLog so cursors - can drive the pump when their buffer catches up. - - Also calls `_bind_apump` on any transformer that exposes it, - so transformers producing `AsyncChatModelStream` objects (e.g. - `MessagesTransformer`) can fan the pull callback out to each - stream's projections. Mirrors the sync `_wire_request_more` - plumbing. - """ - mux._events._arequest_more = self._apump_next - for value in mux.extensions.values(): - if isinstance(value, EventLog): - value._arequest_more = self._apump_next - elif isinstance(value, StreamChannel): - value._log._arequest_more = self._apump_next - for transformer in mux._transformers: - if hasattr(transformer, "_bind_apump"): - transformer._bind_apump(self._apump_next) + mux.bind_apump(self._apump_next) async def _apump_next(self) -> bool: """Pull one event from the graph and push it through the mux. @@ -324,7 +313,8 @@ class AsyncGraphRunStream: Closes the mux and marks the stream exhausted. Any awaiting cursors wake up and see the closed state; any `apush` blocked - on backpressure wakes and returns without appending. Idempotent. + on backpressure wakes and returns without appending. + Idempotent. """ async with self._pump_lock: if self._exhausted: @@ -367,8 +357,7 @@ class AsyncGraphRunStream: return self._values_transformer._latest async def interrupted(self) -> bool: - """Drive the run to completion and return whether it was - interrupted. + """Drive the run to completion and return whether it was interrupted. Raises: BaseException: If the run ended with an error. @@ -388,7 +377,3 @@ class AsyncGraphRunStream: if (err := self._values_transformer.error) is not None: raise err return self._values_transformer._interrupts - - def __aiter__(self) -> AsyncIterator[ProtocolEvent]: - """Subscribe to the main event log and iterate protocol events.""" - return self._mux._events.__aiter__() diff --git a/libs/langgraph/langgraph/stream/streaming_handler.py b/libs/langgraph/langgraph/stream/streaming_handler.py index c5f2015c2..c0c2c265e 100644 --- a/libs/langgraph/langgraph/stream/streaming_handler.py +++ b/libs/langgraph/langgraph/stream/streaming_handler.py @@ -7,13 +7,53 @@ from langchain_core.runnables import RunnableConfig from langgraph._internal._constants import CONF, CONFIG_KEY_STREAM_MESSAGES_V2 from langgraph.pregel import Pregel -from langgraph.stream._mux import StreamMux +from langgraph.stream._mux import StreamMux, TransformerFactory from langgraph.stream._types import StreamTransformer from langgraph.stream.run_stream import AsyncGraphRunStream, GraphRunStream -from langgraph.stream.transformers import MessagesTransformer, ValuesTransformer +from langgraph.stream.transformers import ( + MessagesTransformer, + SubgraphTransformer, + ValuesTransformer, +) from langgraph.types import All, StreamMode +def _coerce_factories( + transformers: list[StreamTransformer | TransformerFactory] | None, +) -> list[TransformerFactory]: + """Normalize caller-supplied transformers into scope-taking factories. + + Accepts already-built instances (wrapped as single-use factories, + with the caveat that they won't be re-instantiated in subgraph + mini-muxes) or proper factories (classes / callables taking a + scope). The built-in root transformers are always factories so + they propagate into every subgraph scope automatically. + """ + + def _wrap_instance(t: StreamTransformer) -> TransformerFactory: + # Single-use: only wires at root scope. A user that wants + # subgraph propagation should pass the class (or a lambda). + def _factory(_scope: tuple[str, ...]) -> StreamTransformer: + return t + + return _factory + + coerced: list[TransformerFactory] = [] + for item in transformers or (): + if isinstance(item, StreamTransformer): + coerced.append(_wrap_instance(item)) + else: + coerced.append(item) + return coerced + + +_BUILTIN_FACTORIES: list[TransformerFactory] = [ + ValuesTransformer, + MessagesTransformer, + SubgraphTransformer, +] + + def _merge_v2_messages_flag( config: RunnableConfig | None, ) -> RunnableConfig: @@ -40,6 +80,7 @@ STREAM_V2_MODES: list[StreamMode] = [ "checkpoints", "tasks", "debug", + "lifecycle", ] @@ -80,7 +121,7 @@ class StreamingHandler: *, interrupt_before: All | Sequence[str] | None = None, interrupt_after: All | Sequence[str] | None = None, - transformers: list[StreamTransformer] | None = None, + transformers: list[StreamTransformer | TransformerFactory] | None = None, ) -> GraphRunStream: """Start a sync streaming run. @@ -100,11 +141,12 @@ class StreamingHandler: Returns: A GraphRunStream the caller can iterate to drive the run. """ - values_t = ValuesTransformer() mux = StreamMux( - [values_t, MessagesTransformer(), *(transformers or ())], + factories=_BUILTIN_FACTORIES + _coerce_factories(transformers), is_async=False, ) + values_t = mux.transformer_by_key("values") + assert isinstance(values_t, ValuesTransformer) graph_iter = iter( self._graph.stream( @@ -127,7 +169,7 @@ class StreamingHandler: *, interrupt_before: All | Sequence[str] | None = None, interrupt_after: All | Sequence[str] | None = None, - transformers: list[StreamTransformer] | None = None, + transformers: list[StreamTransformer | TransformerFactory] | None = None, ) -> AsyncGraphRunStream: """Start an async streaming run. @@ -149,11 +191,12 @@ class StreamingHandler: concurrently; each subscribed cursor drives the pump when its buffer is empty. """ - values_t = ValuesTransformer() mux = StreamMux( - [values_t, MessagesTransformer(), *(transformers or ())], + factories=_BUILTIN_FACTORIES + _coerce_factories(transformers), is_async=True, ) + values_t = mux.transformer_by_key("values") + assert isinstance(values_t, ValuesTransformer) graph_aiter = self._graph.astream( input, diff --git a/libs/langgraph/langgraph/stream/transformers.py b/libs/langgraph/langgraph/stream/transformers.py index 70a981310..aca437cee 100644 --- a/libs/langgraph/langgraph/stream/transformers.py +++ b/libs/langgraph/langgraph/stream/transformers.py @@ -1,6 +1,6 @@ from __future__ import annotations -from typing import TYPE_CHECKING, Any, cast +from typing import TYPE_CHECKING, Any, Literal, cast from langchain_core.language_models._compat_bridge import message_to_events from langchain_core.language_models.chat_model_stream import ( @@ -8,14 +8,24 @@ from langchain_core.language_models.chat_model_stream import ( ChatModelStream, ) from langchain_core.messages import AIMessageChunk, BaseMessage -from langchain_protocol.protocol import MessagesData +from langchain_protocol.protocol import CheckpointRef, LifecycleData, MessagesData +from langgraph.errors import GraphInterrupt from langgraph.stream._event_log import EventLog from langgraph.stream._types import ProtocolEvent, StreamTransformer +from langgraph.stream.run_stream import BaseRunStream if TYPE_CHECKING: from collections.abc import Awaitable, Callable + from langgraph.stream._mux import StreamMux + + +SubgraphStatus = Literal["started", "running", "completed", "failed", "interrupted"] +_TERMINAL_STATUSES: frozenset[SubgraphStatus] = frozenset( + {"completed", "failed", "interrupted"} +) + class ValuesTransformer(StreamTransformer): """Capture values events as a drainable stream of state snapshots. @@ -28,13 +38,16 @@ class ValuesTransformer(StreamTransformer): Native transformer — projection keys are exposed as direct attributes on the run stream (e.g. `run.values`). - Only root-namespace values events are captured; subgraph state - snapshots are ignored. + `scope` (inherited from `StreamTransformer`) is the namespace the + transformer captures values for. `()` matches the root graph; + subgraph mini-muxes pass their subgraph's namespace, so each + instance sees only its own level. """ _native = True - def __init__(self) -> None: + def __init__(self, scope: tuple[str, ...] = ()) -> None: + super().__init__(scope) self._log: EventLog[dict[str, Any]] = EventLog() self._latest: dict[str, Any] | None = None self._interrupted = False @@ -52,11 +65,10 @@ class ValuesTransformer(StreamTransformer): return self._log._error def process(self, event: ProtocolEvent) -> bool: + # Namespace filtering is handled by the mux via `scope_exact`. if event["method"] != "values": return True params = event["params"] - if params["namespace"]: - return True self._latest = params["data"] interrupts = params.get("interrupts", ()) if interrupts: @@ -95,9 +107,10 @@ class MessagesTransformer(StreamTransformer): `stream()` method still surface their final `AIMessage` via `on_chain_end` when a node returns it as state. - Only root-namespace events are captured; tokens from subgraphs are - dropped. Consumers that need subgraph tokens should iterate the raw - event stream or register a custom transformer. + `scope` (inherited from `StreamTransformer`) is the namespace the + transformer captures messages for. `()` matches the root graph; + subgraph mini-muxes pass their subgraph's namespace, so each + instance sees only its own level. Native transformer — the `messages` projection is exposed as a direct attribute on the run stream. @@ -105,7 +118,8 @@ class MessagesTransformer(StreamTransformer): _native = True - def __init__(self) -> None: + def __init__(self, scope: tuple[str, ...] = ()) -> None: + super().__init__(scope) self._log: EventLog[ChatModelStream] = EventLog() # Correlate protocol events back to a ChatModelStream by run_id # (attached to the event's metadata by StreamMessagesHandler). @@ -166,11 +180,10 @@ class MessagesTransformer(StreamTransformer): ) def process(self, event: ProtocolEvent) -> bool: + # Namespace filtering is handled by the mux via `scope_exact`. if event["method"] != "messages": return True params = event["params"] - if params["namespace"]: - return True payload, metadata = params["data"] node: str | None = metadata.get("langgraph_node") @@ -201,7 +214,7 @@ class MessagesTransformer(StreamTransformer): if event_type == "message-start": message_id = event.get("message_id") stream = self._make_stream( - namespace=[], + namespace=list(self.scope), node=node, message_id=str(message_id) if message_id is not None else None, ) @@ -215,7 +228,11 @@ class MessagesTransformer(StreamTransformer): del self._by_run[run_id] def _route_whole_message(self, message: BaseMessage, *, node: str | None) -> None: - stream = self._make_stream(namespace=[], node=node, message_id=message.id) + stream = self._make_stream( + namespace=list(self.scope), + node=node, + message_id=message.id, + ) for evt in message_to_events(message, message_id=message.id): stream.dispatch(evt) self._log.push(stream) @@ -229,3 +246,213 @@ class MessagesTransformer(StreamTransformer): for stream in list(self._by_run.values()): stream.fail(err) self._by_run.clear() + + +class SubgraphRunStream(BaseRunStream): + """Scoped view of a single nested subgraph execution. + + Yielded on `run.subgraphs` (or `parent.subgraphs` for grandchildren) + when a nested `Pregel` spawns. Wraps a mini-`StreamMux` built with + the same transformer factories as the root mux, so `.values`, + `.messages`, `.subgraphs` are populated by the standard + transformers scoped to this handle's namespace — no duplicated + routing logic. + + Inherits `BaseRunStream`, so all shared shape works uniformly: + raw iteration (`for event in sub`), `.interleave(...)`, native + projection attributes, and the `extensions` mapping — identical + to the root `GraphRunStream`. The mini-mux borrows the root's + pump via `make_child`'s pump inheritance, so any cursor on a + subagent projection drives the whole run forward. + + Lifecycle fields update in place as events arrive: + + - `path`: the namespace tuple — stable for the life of the handle. + - `graph_name` / `trigger_call_id`: set once from the `started` + payload. + - `status`: advances `started` → `running` → `completed` / + `failed` / `interrupted`. + - `error` / `checkpoint`: set on the terminal event when present. + + `.output` is a snapshot of the latest values seen at this + namespace — it doesn't drive the pump (unlike root's + `GraphRunStream.output`), because advancing a subgraph to + completion is only meaningful as part of advancing the whole run. + """ + + def __init__( + self, + path: tuple[str, ...], + mux: StreamMux, + *, + graph_name: str | None = None, + trigger_call_id: str | None = None, + ) -> None: + super().__init__(mux) + self.path: tuple[str, ...] = path + self.graph_name: str | None = graph_name + self.trigger_call_id: str | None = trigger_call_id + self.status: SubgraphStatus = "started" + self.error: str | None = None + self.checkpoint: CheckpointRef | None = None + + @property + def output(self) -> dict[str, Any] | None: + """Latest values snapshot at this namespace, or `None`. + + Snapshot-only — iterating other projections or the root's + `.output` is what drives the pump. + """ + values_t = self._mux.transformer_by_key("values") + if isinstance(values_t, ValuesTransformer): + return values_t._latest + return None + + +class SubgraphTransformer(StreamTransformer): + """Discover subgraphs and route events into per-subgraph mini-muxes. + + Thin state-machine + dispatcher. At its own `scope` (inherited + from `StreamTransformer`, determined by the enclosing mux), it + watches for `lifecycle` events at exactly one level deeper to + discover direct children. Each discovered child gets its own + `SubgraphRunStream` backed by a mini-`StreamMux` — built via + `parent_mux.make_child(path)`, so the same factory list produces + fresh transformer instances at the child's scope. + + Every incoming event that falls under one of the direct children + (ns starts with a child's `path`) is forwarded into that child's + mini-mux via `push`. The standard transformers in that mini-mux + (`ValuesTransformer`, `MessagesTransformer`, and another + `SubgraphTransformer` for grandchildren) handle the rest. No + duplicated routing or assembly logic. + + Lifecycle state for each handle (running / completed / failed / + interrupted) is updated in place as events fire. On terminal + events, the handle's mini-mux is closed so any subscribed cursors + unblock. `finalize` / `fail` handle dangling handles left mid-run. + + Native transformer — `subgraphs` exposes the direct-children log. + + `scope_exact = False`: this transformer sees events at any + namespace, because it forwards out-of-scope events to the matching + direct-child mini-mux. + """ + + _native = True + scope_exact = False + + def __init__(self, scope: tuple[str, ...] = ()) -> None: + super().__init__(scope) + self._root_log: EventLog[SubgraphRunStream] = EventLog() + # Direct children only (namespace = scope + one segment). + self._by_ns: dict[tuple[str, ...], SubgraphRunStream] = {} + self._mux: StreamMux | None = None + + def init(self) -> dict[str, Any]: + return {"subgraphs": self._root_log} + + def _on_register(self, mux: StreamMux) -> None: + """Capture the enclosing mux so we can build child mini-muxes.""" + self._mux = mux + + def process(self, event: ProtocolEvent) -> bool: + ns = tuple(event["params"]["namespace"]) + method = event["method"] + depth = len(self.scope) + + # 1. On `started` for a direct child (ns depth = mine + 1 and + # ns prefix matches mine), register the handle. + if method == "lifecycle" and len(ns) == depth + 1 and ns[:-1] == self.scope: + data = cast(LifecycleData, event["params"]["data"]) + if data.get("event") == "started": + self._on_started(ns, data) + + # 2. Forward the event to the matching direct-child mini-mux. + # Prefix-match: ns must start with some child's path. + direct_child_ns = ns[: depth + 1] if len(ns) > depth else None + if direct_child_ns is not None and direct_child_ns in self._by_ns: + self._by_ns[direct_child_ns]._mux.push(event) + + # 3. Status change for a direct child (ns = child's path, method + # = lifecycle). Update handle fields, close mini-mux on + # terminal. + if ( + method == "lifecycle" + and ns in self._by_ns + and len(ns) == depth + 1 + and ns[:-1] == self.scope + ): + data = cast(LifecycleData, event["params"]["data"]) + event_type = data.get("event") + if event_type in ("running", "completed", "failed", "interrupted"): + self._on_status_change(ns, event_type, data) + + return True + + def _on_started(self, ns: tuple[str, ...], data: LifecycleData) -> None: + if ns in self._by_ns: + # Duplicate started — ignore. + return + if self._mux is None: + # Not registered yet; can't build a mini-mux. + return + child_mux = self._mux.make_child(ns) + handle = SubgraphRunStream( + path=ns, + mux=child_mux, + graph_name=data.get("graph_name"), + trigger_call_id=data.get("trigger_call_id"), + ) + self._by_ns[ns] = handle + self._root_log.push(handle) + + def _on_status_change( + self, + ns: tuple[str, ...], + event_type: SubgraphStatus, + data: LifecycleData, + ) -> None: + handle = self._by_ns[ns] + handle.status = event_type + err = data.get("error") + if err is not None: + handle.error = err + checkpoint = data.get("checkpoint") + if checkpoint is not None: + handle.checkpoint = checkpoint + if event_type in _TERMINAL_STATUSES: + self._close_handle_mux(handle) + + @staticmethod + def _close_handle_mux(handle: SubgraphRunStream) -> None: + # Idempotent close — mux.close() runs finalize on its transformers + # (which cascades through grandchildren) and closes projection logs. + if not handle._mux._events._closed: + try: + handle._mux.close() + except Exception: + pass + + def finalize(self) -> None: + """Transition any still-open direct children to `completed`.""" + for handle in self._by_ns.values(): + if handle.status not in _TERMINAL_STATUSES: + handle.status = "completed" + self._close_handle_mux(handle) + + def fail(self, err: BaseException) -> None: + """Transition any still-open direct children to `failed` / `interrupted`.""" + is_interrupt = isinstance(err, GraphInterrupt) + terminal: SubgraphStatus = "interrupted" if is_interrupt else "failed" + error_str = None if is_interrupt else str(err) + for handle in self._by_ns.values(): + if handle.status not in _TERMINAL_STATUSES: + handle.status = terminal + if error_str is not None and handle.error is None: + handle.error = error_str + if not handle._mux._events._closed: + try: + handle._mux.fail(err) + except Exception: + pass diff --git a/libs/langgraph/langgraph/types.py b/libs/langgraph/langgraph/types.py index d04d82da7..73954e5eb 100644 --- a/libs/langgraph/langgraph/types.py +++ b/libs/langgraph/langgraph/types.py @@ -116,7 +116,14 @@ def ensure_valid_checkpointer(checkpointer: Checkpointer) -> Checkpointer: StreamMode = Literal[ - "values", "updates", "checkpoints", "tasks", "debug", "messages", "custom" + "values", + "updates", + "checkpoints", + "tasks", + "debug", + "messages", + "custom", + "lifecycle", ] """How the stream method should emit outputs. @@ -129,6 +136,7 @@ StreamMode = Literal[ - `"checkpoints"`: Emit an event when a checkpoint is created, in the same format as returned by `get_state()`. - `"tasks"`: Emit events when tasks start and finish, including their results and errors. - `"debug"`: Emit `"checkpoints"` and `"tasks"` events for debugging purposes. +- `"lifecycle"`: Emit subgraph lifecycle events (`started`, `running`, `completed`, `failed`, `interrupted`) with payloads matching `LifecycleData`. """ StreamWriter = Callable[[Any], None] diff --git a/libs/langgraph/tests/test_stream_messages_transformer.py b/libs/langgraph/tests/test_stream_messages_transformer.py index 35e62f8d1..27492ac40 100644 --- a/libs/langgraph/tests/test_stream_messages_transformer.py +++ b/libs/langgraph/tests/test_stream_messages_transformer.py @@ -318,8 +318,16 @@ class TestFiltering: assert t.process(values_event) is True def test_subgraph_namespace_dropped(self) -> None: - t, log = _make_sync_transformer() - t.process( + """Root MessagesTransformer (via the mux) ignores non-root events.""" + from langgraph.stream._mux import StreamMux + + mux = StreamMux([MessagesTransformer()], is_async=False) + t = mux.transformer_by_key("messages") + assert isinstance(t, MessagesTransformer) + t._log._subscribed = True + t._bind_pump(lambda: False) + + mux.push( { "type": "event", "method": "messages", @@ -333,8 +341,8 @@ class TestFiltering: }, } ) - log.close() - assert list(log._items) == [] + t._log.close() + assert list(t._log._items) == [] # --------------------------------------------------------------------------- diff --git a/libs/langgraph/tests/test_stream_subgraph_transformer.py b/libs/langgraph/tests/test_stream_subgraph_transformer.py new file mode 100644 index 000000000..67fa25826 --- /dev/null +++ b/libs/langgraph/tests/test_stream_subgraph_transformer.py @@ -0,0 +1,367 @@ +"""Tests for subgraph lifecycle events and the SubgraphTransformer.""" + +from __future__ import annotations + +import operator +import time +from typing import Annotated, Any + +import pytest +from typing_extensions import TypedDict + +from langgraph.constants import END, START +from langgraph.errors import GraphInterrupt +from langgraph.graph import StateGraph +from langgraph.stream import StreamingHandler +from langgraph.stream._event_log import EventLog +from langgraph.stream._mux import StreamMux +from langgraph.stream._types import ProtocolEvent +from langgraph.stream.transformers import ( + MessagesTransformer, + SubgraphRunStream, + SubgraphTransformer, + ValuesTransformer, +) + +TS = int(time.time() * 1000) + + +def _lifecycle( + event: str, + *, + namespace: list[str] | None = None, + graph_name: str | None = None, + trigger_call_id: str | None = None, + error: str | None = None, +) -> ProtocolEvent: + data: dict[str, Any] = {"event": event} + if graph_name is not None: + data["graph_name"] = graph_name + if trigger_call_id is not None: + data["trigger_call_id"] = trigger_call_id + if error is not None: + data["error"] = error + return { + "type": "event", + "method": "lifecycle", + "params": { + "namespace": namespace or [], + "timestamp": TS, + "data": data, + }, + } + + +def _values(payload: dict[str, Any], *, namespace: list[str]) -> ProtocolEvent: + return { + "type": "event", + "method": "values", + "params": { + "namespace": namespace, + "timestamp": TS, + "data": payload, + }, + } + + +def _subscribe(log: EventLog) -> None: + """Flip `_subscribed = True` so pushes retain items for test inspection.""" + log._subscribed = True + + +# --------------------------------------------------------------------------- +# Unit tests: feed events directly into the transformer +# --------------------------------------------------------------------------- + + +_FACTORIES = [ValuesTransformer, MessagesTransformer, SubgraphTransformer] + + +def _handle_values_items(handle: SubgraphRunStream) -> list: + return list(handle._mux.extensions["values"]._items) # type: ignore[attr-defined] + + +def _handle_subgraphs_items(handle: SubgraphRunStream) -> list: + return list(handle._mux.extensions["subgraphs"]._items) # type: ignore[attr-defined] + + +def _pre_subscribe_handle(handle: SubgraphRunStream) -> None: + """Flip `_subscribed` on every EventLog inside the handle's mini-mux. + + The mini-mux is built via `make_child` with the full factory list, + so values / messages / subgraphs logs all exist as projections. + Tests that feed events directly need them subscribed so pushes + retain items in the deque for `_items` inspection. + """ + for value in handle._mux.extensions.values(): + if isinstance(value, EventLog): + _subscribe(value) + + +class TestSubgraphTransformerUnit: + def _mux(self) -> tuple[StreamMux, SubgraphTransformer]: + mux = StreamMux(factories=_FACTORIES, is_async=False) + transformer = mux.transformer_by_key("subgraphs") + assert isinstance(transformer, SubgraphTransformer) + _subscribe(transformer._root_log) + return mux, transformer + + def _handle(self, transformer: SubgraphTransformer) -> SubgraphRunStream: + """Return the single root handle after pushing one lifecycle started.""" + (handle,) = list(transformer._root_log._items) + return handle + + def test_root_started_is_ignored(self) -> None: + mux, transformer = self._mux() + mux.push(_lifecycle("started", graph_name="root")) + assert list(transformer._root_log._items) == [] + assert transformer._by_ns == {} + + def test_child_started_yields_handle(self) -> None: + mux, transformer = self._mux() + mux.push( + _lifecycle( + "started", + namespace=["task_a:child"], + graph_name="child", + trigger_call_id="task_a", + ) + ) + + handle = self._handle(transformer) + assert handle.path == ("task_a:child",) + assert handle.graph_name == "child" + assert handle.trigger_call_id == "task_a" + assert handle.status == "started" + + def test_status_transitions(self) -> None: + mux, transformer = self._mux() + mux.push(_lifecycle("started", namespace=["t:c"], graph_name="c")) + mux.push(_lifecycle("running", namespace=["t:c"])) + mux.push(_lifecycle("completed", namespace=["t:c"])) + + handle = self._handle(transformer) + assert handle.status == "completed" + + def test_grandchild_surfaces_under_child(self) -> None: + mux, transformer = self._mux() + mux.push(_lifecycle("started", namespace=["t:child"], graph_name="child")) + + child = self._handle(transformer) + _pre_subscribe_handle(child) + + mux.push( + _lifecycle( + "started", + namespace=["t:child", "u:grand"], + graph_name="grand", + ) + ) + + (grand,) = _handle_subgraphs_items(child) + assert grand.path == ("t:child", "u:grand") + assert grand.graph_name == "grand" + + def test_failed_stores_error(self) -> None: + mux, transformer = self._mux() + mux.push(_lifecycle("started", namespace=["t:c"], graph_name="c")) + mux.push(_lifecycle("failed", namespace=["t:c"], error="boom")) + + handle = self._handle(transformer) + assert handle.status == "failed" + assert handle.error == "boom" + + def test_values_routed_into_handle(self) -> None: + mux, transformer = self._mux() + mux.push(_lifecycle("started", namespace=["t:c"], graph_name="c")) + + handle = self._handle(transformer) + _pre_subscribe_handle(handle) + + mux.push(_values({"value": 1}, namespace=["t:c"])) + mux.push(_values({"value": 2}, namespace=["t:c"])) + + assert _handle_values_items(handle) == [{"value": 1}, {"value": 2}] + assert handle.output == {"value": 2} + + def test_root_values_not_routed(self) -> None: + mux, transformer = self._mux() + mux.push(_lifecycle("started", namespace=["t:c"], graph_name="c")) + handle = self._handle(transformer) + _pre_subscribe_handle(handle) + + # Values event at root namespace — must not leak into child handle. + mux.push(_values({"value": "root"}, namespace=[])) + assert _handle_values_items(handle) == [] + + def test_finalize_closes_dangling(self) -> None: + mux, transformer = self._mux() + mux.push(_lifecycle("started", namespace=["t:c"], graph_name="c")) + handle = self._handle(transformer) + + mux.close() + assert handle.status == "completed" + assert handle._mux.extensions["values"]._closed + assert handle._mux.extensions["subgraphs"]._closed + + def test_fail_with_graph_interrupt_marks_interrupted(self) -> None: + mux, transformer = self._mux() + mux.push(_lifecycle("started", namespace=["t:c"], graph_name="c")) + handle = self._handle(transformer) + + mux.fail(GraphInterrupt()) + assert handle.status == "interrupted" + + def test_fail_with_generic_error_marks_failed(self) -> None: + mux, transformer = self._mux() + mux.push(_lifecycle("started", namespace=["t:c"], graph_name="c")) + handle = self._handle(transformer) + + mux.fail(RuntimeError("explode")) + assert handle.status == "failed" + assert handle.error == "explode" + + def test_duplicate_started_ignored(self) -> None: + mux, transformer = self._mux() + mux.push(_lifecycle("started", namespace=["t:c"], graph_name="c")) + mux.push(_lifecycle("started", namespace=["t:c"], graph_name="other")) + + handles = list(transformer._root_log._items) + assert len(handles) == 1 + assert handles[0].graph_name == "c" + + def test_non_lifecycle_non_values_passthrough(self) -> None: + mux, transformer = self._mux() + mux.push( + { + "type": "event", + "method": "messages", + "params": {"namespace": ["t:c"], "timestamp": TS, "data": "x"}, + } + ) + assert list(transformer._root_log._items) == [] + + +# --------------------------------------------------------------------------- +# End-to-end tests via StreamingHandler on real graphs +# --------------------------------------------------------------------------- + + +class SimpleState(TypedDict): + value: str + items: Annotated[list[str], operator.add] + + +def _build_nested_graph(): + """Parent graph with a compiled subgraph node.""" + + def inner_node(state: SimpleState) -> dict: + return {"value": state["value"] + "X", "items": ["x"]} + + inner_builder = StateGraph(SimpleState) + inner_builder.add_node("inner_node", inner_node) + inner_builder.add_edge(START, "inner_node") + inner_builder.add_edge("inner_node", END) + inner = inner_builder.compile() + + def outer_node(state: SimpleState) -> dict: + return {"value": state["value"] + "Y", "items": ["y"]} + + outer_builder = StateGraph(SimpleState) + outer_builder.add_node("outer_node", outer_node) + outer_builder.add_node("sub", inner) + outer_builder.add_edge(START, "outer_node") + outer_builder.add_edge("outer_node", "sub") + outer_builder.add_edge("sub", END) + return outer_builder.compile() + + +class TestSubgraphTransformerEndToEnd: + def test_flat_graph_yields_no_subgraphs(self) -> None: + builder = StateGraph(SimpleState) + builder.add_node("n", lambda s: {"value": s["value"] + "!", "items": ["!"]}) + builder.add_edge(START, "n") + builder.add_edge("n", END) + graph = builder.compile() + + handler = StreamingHandler(graph) + run = handler.stream({"value": "", "items": []}) + + collected: list[SubgraphRunStream] = [] + for sub in run.subgraphs: + collected.append(sub) + assert collected == [] + # Output still resolves. + assert run.output is not None + + def test_nested_graph_yields_one_child(self) -> None: + graph = _build_nested_graph() + handler = StreamingHandler(graph) + run = handler.stream({"value": "", "items": []}) + + collected: list[SubgraphRunStream] = [] + for sub in run.subgraphs: + collected.append(sub) + + assert len(collected) == 1 + child = collected[0] + assert len(child.path) == 1 + assert child.path[0].startswith("sub:") + assert child.status == "completed" + + def test_error_in_subgraph_fails_child(self) -> None: + def boom(state: SimpleState) -> dict: + raise RuntimeError("subgraph_failed") + + inner_builder = StateGraph(SimpleState) + inner_builder.add_node("inner", boom) + inner_builder.add_edge(START, "inner") + inner_builder.add_edge("inner", END) + inner = inner_builder.compile() + + outer_builder = StateGraph(SimpleState) + outer_builder.add_node("sub", inner) + outer_builder.add_edge(START, "sub") + outer_builder.add_edge("sub", END) + graph = outer_builder.compile() + + handler = StreamingHandler(graph) + run = handler.stream({"value": "", "items": []}) + + collected: list[SubgraphRunStream] = [] + with pytest.raises(RuntimeError): + for sub in run.subgraphs: + collected.append(sub) + + assert len(collected) == 1 + assert collected[0].status == "failed" + + +class TestSubgraphTransformerAsyncEndToEnd: + @pytest.mark.anyio + async def test_nested_graph_yields_one_child(self) -> None: + async def inner(state: SimpleState) -> dict: + return {"value": state["value"] + "X", "items": ["x"]} + + inner_builder = StateGraph(SimpleState) + inner_builder.add_node("inner", inner) + inner_builder.add_edge(START, "inner") + inner_builder.add_edge("inner", END) + inner_graph = inner_builder.compile() + + outer_builder = StateGraph(SimpleState) + outer_builder.add_node("sub", inner_graph) + outer_builder.add_edge(START, "sub") + outer_builder.add_edge("sub", END) + graph = outer_builder.compile() + + handler = StreamingHandler(graph) + run = await handler.astream({"value": "", "items": []}) + + collected: list[SubgraphRunStream] = [] + async for sub in run.subgraphs: + collected.append(sub) + + assert len(collected) == 1 + child = collected[0] + assert child.status == "completed" diff --git a/libs/langgraph/tests/test_streaming_handler.py b/libs/langgraph/tests/test_streaming_handler.py index 3f5cdc8e2..4e40c74fd 100644 --- a/libs/langgraph/tests/test_streaming_handler.py +++ b/libs/langgraph/tests/test_streaming_handler.py @@ -869,14 +869,18 @@ class TestStreamMux: class TestValuesTransformer: def test_ignores_non_root_namespace(self) -> None: - """Values events from subgraphs (non-empty namespace) should be ignored.""" - t = ValuesTransformer() - t.init() - t._log._bind(is_async=False) + """The root mux only dispatches root-ns values events to its ValuesTransformer. + + Namespace filtering is enforced by the mux via `scope_exact` + — the transformer itself no longer filters. + """ + mux = StreamMux([ValuesTransformer()], is_async=False) + t = mux.transformer_by_key("values") + assert isinstance(t, ValuesTransformer) it = iter(t._log) - t.process(_event("values", {"val": "root"})) - t.process(_event("values", {"val": "sub"}, namespace=["sub"])) + mux.push(_event("values", {"val": "root"})) + mux.push(_event("values", {"val": "sub"}, namespace=["sub"])) t._log.close() items = list(it) @@ -935,14 +939,15 @@ class TestMessagesTransformer: assert items[0].message_id == "run-1" def test_ignores_non_root_namespace(self) -> None: - t = MessagesTransformer() - t.init() - t._log._bind(is_async=False) + """Namespace filtering is enforced by the mux via `scope_exact`.""" + mux = StreamMux([MessagesTransformer()], is_async=False) + t = mux.transformer_by_key("messages") + assert isinstance(t, MessagesTransformer) t._bind_pump(lambda: False) it = iter(t._log) meta = {"langgraph_node": "llm", "run_id": "run-1"} - t.process( + mux.push( _event( "messages", ({"event": "message-start", "message_id": "run-1"}, meta),