Add subgraph lifecycle streaming with scoped per-subagent projections

StreamingHandler now yields SubgraphRunStream handles for each nested
Pregel as it spawns. Each handle is a BaseRunStream wrapping a
mini-mux built from the same transformer factories as the root — so
sub.values, sub.messages, sub.subgraphs are populated by standard
ValuesTransformer / MessagesTransformer / SubgraphTransformer
instances at that subagent's scope. No routing or ChatModelStream
assembly duplicated across transformers.

Key pieces:

- StreamLifecycleHandler (pregel/_lifecycle.py): callback handler
  attached at pregel stream / astream sites when "lifecycle" is in
  stream_modes. Emits started / running / completed / failed /
  interrupted events per nested Pregel via metadata-based detection
  (langgraph_checkpoint_ns + name != langgraph_node, excluding
  __start__/__end__ sentinels). Carries trigger_call_id from the
  parent task id. Root subgraph terminal state is emitted eagerly at
  __init__ and by SubgraphTransformer.finalize / fail.

- StreamMux.make_child(scope) + factory-based construction. Mux takes
  a factory list; make_child produces a mini-mux at a new scope with
  fresh instances. bind_pump / bind_apump cascade through children so
  any subagent cursor drives the root pump.

- StreamTransformer.scope (base attribute) + scope_exact class flag.
  Mux skips process() for out-of-scope events when scope_exact=True
  (default), so user transformers get scope-filtered events with no
  boilerplate. SubgraphTransformer opts out to receive cross-scope
  events for forwarding.

- BaseRunStream shared base for GraphRunStream, AsyncGraphRunStream,
  and SubgraphRunStream. Provides extensions, native attrs, raw
  __iter__ / __aiter__, interleave. Root subclasses own the
  graph iterator and pump; SubgraphRunStream adds lifecycle metadata
  (path, status, error, checkpoint, trigger_call_id, graph_name).

- SubgraphTransformer (stream/transformers.py) is now a thin
  discovery + forwarding dispatcher. On lifecycle.started at its
  scope + 1, it creates a SubgraphRunStream via
  parent_mux.make_child(ns). It forwards each event matching a
  direct-child's path into that child's mini-mux. Terminal lifecycle
  closes the mini-mux.

- StreamMode literal extended with "lifecycle". pregel/remote.py
  filters it out of the SDK stream-mode list since wire-protocol
  support isn't landed yet.

Tests: 16 new in test_stream_subgraph_transformer.py (unit + sync/async
end-to-end including error + grandchild); existing Values/Messages
namespace filter tests migrated through mux.push to reflect the new
scope_exact contract.
This commit is contained in:
Nick Hollon
2026-04-18 18:15:31 -04:00
parent 7715239e3b
commit 2ad30132a3
12 changed files with 1250 additions and 184 deletions
@@ -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)})
+19
View File
@@ -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(
+9 -4
View File
@@ -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")
+142 -14
View File
@@ -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:
+26
View File
@@ -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]:
+113 -128
View File
@@ -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__()
@@ -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,
+242 -15
View File
@@ -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
+9 -1
View File
@@ -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]
@@ -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) == []
# ---------------------------------------------------------------------------
@@ -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"
+15 -10
View File
@@ -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),