Compare commits

...
Author SHA1 Message Date
Nick Hollon 0818530f07 Merge branch 'nh/streaming-transformer' into nh/subgraph-lifecycle 2026-04-27 09:40:51 -04:00
Nick HollonandGitHub 0eac6626b3 feat(langgraph): route invoke messages through v2 via StreamingHandler (#7546) 2026-04-24 16:59:07 -04:00
Nick Hollon 67b2512857 move toward lifecycle transformer instead of callback handler 2026-04-24 16:55:41 -04:00
Nick Hollon 2982512adb Merge branch 'nh/messages-content-blocks' into nh/subgraph-lifecycle
# Conflicts:
#	libs/langgraph/langgraph/stream/streaming_handler.py
#	libs/langgraph/tests/test_pregel_stream_v2.py
2026-04-20 16:22:18 -04:00
Nick Hollon 4956512692 Merge branch 'nh/streaming-transformer' into nh/messages-content-blocks
# Conflicts:
#	libs/langgraph/langgraph/stream/streaming_handler.py
2026-04-20 16:13:08 -04:00
Nick Hollon 40055e92cc feat(langgraph): move stream_v2/astream_v2 onto Pregel, drop StreamingHandler
Compile-time `transformers=` on `StateGraph.compile` now registers
transformer factories directly on the compiled graph. `stream_v2` and
`astream_v2` live on Pregel and read the stashed list, so callers no
longer need a separate wrapper to drive the transformer pipeline.
2026-04-20 16:08:19 -04:00
Nick Hollon f4a5d56535 fix(stream): replace pump lock with condition-based take-a-number
The async pump serialized `graph_aiter.__anext__()` with an
`asyncio.Lock`, held across the full await. When two cursors read
different projections concurrently, the "losing" task slept inside
`_apump_next` on the lock itself — so when the active pumper pushed
its data onto the losing task's buffer, the loser couldn't observe it
until another graph event forced the lock to change hands. Each
passive consumer saw its deltas one graph event late; bursts
coalesced at turn boundaries instead of streaming live.

Switch to an `asyncio.Condition` + `_pumping` flag. Exactly one task
is the active pumper; others do `cond.wait()` and are notified after
every pump step. Passive consumers wake as soon as their buffer
fills, drop out of `_apump_next`, and let the iterator's buffer check
yield the data. Single-consumer behavior is unchanged; multi-
consumer throughput improves ~5x on bursty tools and no events are
lost.
2026-04-20 13:58:25 -04:00
Nick Hollon 742488d5a0 fix(stream): replace pump lock with condition-based take-a-number
The async pump serialized `graph_aiter.__anext__()` with an
`asyncio.Lock`, held across the full await. When two cursors read
different projections concurrently, the "losing" task slept inside
`_apump_next` on the lock itself — so when the active pumper pushed
its data onto the losing task's buffer, the loser couldn't observe it
until another graph event forced the lock to change hands. Each
passive consumer saw its deltas one graph event late; bursts
coalesced at turn boundaries instead of streaming live.

Switch to an `asyncio.Condition` + `_pumping` flag. Exactly one task
is the active pumper; others do `cond.wait()` and are notified after
every pump step. Passive consumers wake as soon as their buffer
fills, drop out of `_apump_next`, and let the iterator's buffer check
yield the data. Single-consumer behavior is unchanged; multi-
consumer throughput improves ~5x on bursty tools and no events are
lost.
2026-04-20 13:58:15 -04:00
Nick Hollon b319426725 fix(stream): replace pump lock with condition-based take-a-number
The async pump serialized `graph_aiter.__anext__()` with an
`asyncio.Lock`, held across the full await. When two cursors read
different projections concurrently, the "losing" task slept inside
`_apump_next` on the lock itself — so when the active pumper pushed
its data onto the losing task's buffer, the loser couldn't observe it
until another graph event forced the lock to change hands. Each
passive consumer saw its deltas one graph event late; bursts
coalesced at turn boundaries instead of streaming live.

Switch to an `asyncio.Condition` + `_pumping` flag. Exactly one task
is the active pumper; others do `cond.wait()` and are notified after
every pump step. Passive consumers wake as soon as their buffer
fills, drop out of `_apump_next`, and let the iterator's buffer check
yield the data. Single-consumer behavior is unchanged; multi-
consumer throughput improves ~5x on bursty tools and no events are
lost.
2026-04-20 13:57:37 -04:00
Nick Hollon bebcd20815 Address review feedback on subgraph lifecycle streaming
- Robustify nested-Pregel detection with a parent_run_id fallback so
  subgraphs compiled with name equal to their node name are still
  recognized.
- Narrow bare excepts in SubgraphTransformer close/fail paths; log at
  warning with exc_info instead of silently swallowing.
- Assert mux registration in SubgraphTransformer._on_started instead
  of silently dropping events.
- Warn when RemoteGraph strips an unsupported "lifecycle" stream mode
  so callers aren't left wondering why no events arrive.
- Reject pre-built transformer instances in StreamingHandler; factories
  are required so transformers propagate into every subgraph scope.
- Comment the forward-before-close ordering in SubgraphTransformer.
- Trim duplicated pump/projection docstrings across run-stream classes.
- Add end-to-end tests for trigger_call_id, subgraph interrupt, and the
  name-collision detection fallback.
2026-04-18 18:59:23 -04:00
Nick Hollon 2ad30132a3 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.
2026-04-18 18:15:31 -04:00
Nick Hollon 7715239e3b Merge branch 'nh/messages-content-blocks' into nh/subgraph-lifecycle 2026-04-18 13:30:02 -04:00
Nick Hollon ad0146a4de Wire async pump into MessagesTransformer streams
AsyncChatModelStream projections deadlocked when iterated inside
the outer run.messages cursor: the inner stream awaited an
asyncio.Event that nothing was driving while the outer cursor was
suspended.

Plumb the langchain-core async pump hook down to each stream:
- MessagesTransformer gains _bind_apump (mirror of _bind_pump) and
  prefers async wiring in _make_stream.
- AsyncGraphRunStream._wire_arequest_more calls _bind_apump on any
  transformer that exposes it.

Test helpers: EventLog.push is a no-op before subscription; the
unit-test helpers and TestViaMux setups now pre-subscribe the log
(simulating what run.messages iteration does in production) and
verify pushed items via log._items. Flip the known-failure nested
iteration test to pass.
2026-04-18 13:20:06 -04:00
Nick Hollon b6a196fac6 Merge streaming-transformer drain-on-consume redesign
# Conflicts:
#	libs/langgraph/langgraph/stream/run_stream.py
#	libs/langgraph/langgraph/stream/streaming_handler.py
#	libs/langgraph/langgraph/stream/transformers.py
#	libs/langgraph/tests/test_streaming_handler.py
2026-04-18 12:38:46 -04:00
Nick Hollon 910240a930 Merge streaming-transformer drain-on-consume redesign
# Conflicts:
#	libs/langgraph/langgraph/stream/streaming_handler.py
2026-04-18 12:37:16 -04:00
Nick Hollon ab1d6980b5 Drain-on-consume streaming with caller-driven async pump
Collapse the eager async pump task into the same caller-driven model
as sync: each cursor's advance drives one graph event through the
mux. Concurrent async consumers serialize through an asyncio.Lock so
each acquisition produces exactly one event, matching sync semantics.

EventLog becomes a single-consumer drainable queue — items pop off
as the cursor advances, a second __iter__ / __aiter__ raises.
Fan-out moves to explicit tee(n) / atee(n) helpers. Retention
windows, BufferOverflowError, and max_events are gone; pre-
subscription pushes are silent no-ops so unsubscribed projections
don't accumulate.

Both run streams gain abort() and context-manager support; the
pump's BaseException catch is narrowed to Exception so
CancelledError propagates per asyncio contract.

TestMemoryBounds locks in the drain-on-consume invariants:
subscribed buffers drop back to empty after each yield, unsubscribed
projections never accumulate, and run.output leaves the values log
untouched.
2026-04-18 12:34:10 -04:00
Nick Hollon 7e5df56688 Produce ChatModelStream objects from MessagesTransformer
Replace the passthrough (chunk, metadata) tuple projection with one
that yields a ChatModelStream per LLM call, routed by run_id. Handle
both v2 protocol-event payloads (message-start/chunk/message-finish)
and whole AIMessage payloads from on_chain_end (replayed via
message_to_events). Wire _bind_pump from GraphRunStream so nested
sync streams share the caller-driven pump.
2026-04-18 10:56:45 -04:00
Nick Hollon 0f2f66fc8f refactor(langgraph): scope on_stream_event to StreamMessagesHandlerV2
Moves `on_stream_event` off the v1 `StreamMessagesHandler` base and
onto the v2 subclass. Content-block events are a v2-only concept, so
forwarding them only when the v2 handler is attached keeps the
messages channel's shape predictable for v1 callers: plain
`graph.stream(stream_mode="messages")` now ignores `on_stream_event`
entirely, even if a node explicitly calls `model.stream_v2()` on a
v1-flagged run. Dedupe of the returned AIMessage still works in that
case via `_find_and_emit_messages` / `on_chain_end`.

Also makes the v2 handler's `on_llm_new_token` override an explicit
pass-through with a comment rather than `return None`, so it reads as
an intentional no-op instead of a missing return value.
2026-04-17 16:22:50 -04:00
Nick Hollon acaa767542 feat(langgraph): route invoke messages through v2 via StreamingHandler
When `StreamingHandler(graph).stream()` is used, content-block (v2)
protocol events now flow through `stream_mode="messages"` for every
`model.invoke()` call inside a node — with no node-level code changes.

Adds `StreamMessagesHandlerV2`, a `StreamMessagesHandler` subclass that
also inherits `_V2StreamingCallbackHandler` from langchain-core. The
marker base flips `BaseChatModel.invoke` to drive the protocol event
generator (firing `on_stream_event`) instead of `_stream` (firing
`on_llm_new_token`). The handler inherits `on_stream_event` from the
parent — events forward onto the messages channel unchanged — and
overrides `on_llm_new_token` to no-op so a node calling `model.stream()`
directly on a v2-flagged run can't leak AIMessageChunks onto the same
channel.

Opt-in is scoped to `StreamingHandler`: it merges a new internal
`CONFIG_KEY_STREAM_MESSAGES_V2=True` into `config.configurable` before
dispatching to `graph.stream` / `graph.astream`. Pregel reads the flag
at handler-construction time in both sync and async stream paths and
attaches the v2 subclass only when set. Direct
`graph.stream(stream_mode="messages")` callers keep the v1
`(AIMessageChunk, metadata)` shape — confirmed by a regression test.

Existing dedupe between the streamed v2 lifecycle and a node returning
the same assembled `AIMessage` transfers for free: the handler populates
`self.seen` from `message-start` events (via the inherited
`on_stream_event` body), and `on_chain_end`'s `_find_and_emit_messages`
already gates on `seen` — so an invoking node surfaces as exactly one
`ChatModelStream`, not two.

Test coverage in `tests/test_stream_messages_transformer.py`:

- `TestEndToEndV2Invoke` — node calling `model.invoke()` produces a
  single `ChatModelStream` with the full v2 event lifecycle, text
  projection accumulates correctly, multi-node graphs produce one
  stream per model call, constructed-message nodes still replay via
  `message_to_events`, async mirror via `ainvoke` + `astream`.
- `TestDirectMessagesModeStaysV1` — regression guard: direct
  `graph.stream(stream_mode="messages")` still yields AIMessageChunk
  tuples (not event dicts).
- `TestStreamMessagesHandlerV2Unit` — direct unit test that the v2
  handler's `on_llm_new_token` does not emit.
2026-04-17 16:13:11 -04:00
Nick Hollon 5f24a0356a Tighten streaming run stream API and close review footguns
- AsyncGraphRunStream.output/interrupted/interrupts are now methods
  (await run.output()), not properties returning coroutines. Forgetting
  `await` now fails at type-check time and at runtime on the common
  operations (bool/len/iter), instead of silently yielding a live
  coroutine that's truthy, lenless, and never awaited.
- interrupted/interrupts re-raise the run's error on both lanes so a
  failed run doesn't silently return the last-known interrupt state.
- Narrow the async pump catch from BaseException to Exception so
  CancelledError / KeyboardInterrupt propagate.
- Wrap run.extensions with types.MappingProxyType so users can't add
  or remove projection keys behind the mux's back.
- Add ValuesTransformer.error accessor; run stream stops reaching into
  _log._error.
- Tighten StreamingHandler graph type from Any to Pregel and widen
  convert_to_protocol_event to accept StreamPart.
- Projection-conflict ValueError now names the transformer that owns
  each colliding key, not just the incoming transformer.
- Replace deprecated asyncio.get_event_loop() in the async iteration
  test with asyncio.create_task.
- Document wall-clock semantics of ProtocolEvent.params.timestamp,
  the subgraph-namespace drop in MessagesTransformer, and the
  transformer-pipeline bypass for StreamChannel auto-forwarded events.
- Add tests for the new error-raising behavior on interrupted /
  interrupts and for the read-only extensions contract.
2026-04-16 15:54:14 -04:00
Nick Hollon adda5f0341 Convert stream module docstrings to Google style
Per repo convention (CLAUDE.md) and general project style: use single
backticks for inline code, Google-style Args/Returns/Raises sections,
and triple-backtick fenced code blocks instead of Sphinx double
backticks, :param: markers, or Usage:: blocks.

No behavior changes — docs only.
2026-04-16 14:45:02 -04:00
Nick Hollon 6fcca359df Remove streaming comparison example — moved into PR description
The v1-vs-v2 comparison was the whole purpose of the runnable script,
and a condensed v2-only usage + transformer example now lives in the
PR description where reviewers will see it.
2026-04-16 14:37:15 -04:00
Nick Hollon 28ce32edc7 Bound EventLog / StreamChannel memory with drop-oldest semantics
- EventLog(maxlen=N) caps retention. When the buffer is full, push
  evicts the oldest item and advances an absolute _first_seq so
  cursors can detect they've fallen off the back. A lagging cursor
  raises BufferOverflowError on its next read — mirrors the
  restored=false signal from the reconnection scenario (§06).
- New cursors start at the current head of the buffer, not seq 0.
  For unbounded logs this is indistinguishable from the old behavior;
  for bounded logs, new consumers see whatever is still retained.
- StreamChannel(name, *, maxlen=N) forwards maxlen to its inner log.
- StreamMux(..., max_events=N) sets a default maxlen for every log /
  channel it binds (main event log plus each transformer projection).
  Explicit per-log maxlen wins over the mux default.
- StreamingHandler.stream() / astream() expose max_events: caller
  sets the run-wide memory budget; transformer authors can override
  per-log when they know better. Default unbounded, matching §15 Q3.
2026-04-16 14:32:45 -04:00
Nick Hollon 119847f80f Add async lane to StreamTransformer; roll registration into StreamMux
- StreamTransformer: aprocess/afinalize/afail + schedule() helper with
  on_error="log"|"raise". requires_async flag (plus override detection)
  makes sync stream() raise at registration rather than at first event.
- StreamMux: apush/aclose/afail for the async dispatch path. aclose
  awaits all scheduled tasks across transformers before afinalize;
  afail cancels and awaits them before afail hooks.
- StreamMux now takes transformers in __init__ and owns extensions /
  native_keys aggregation and conflict detection — register() is gone.
- GraphRunStream / AsyncGraphRunStream read extensions and native keys
  off the mux directly; StreamingHandler._setup() inlined.
2026-04-16 14:20:03 -04:00
Nick Hollon f43743c3e7 Use asyncio.Event for async notification instead of per-cursor futures
Replace the _async_waiters list and manual future management with a
single shared asyncio.Event. Simpler notification (just event.set()),
no per-cursor future allocation, no get_running_loop/create_future
in our code.
2026-04-16 11:31:35 -04:00
Nick Hollon dbded7a59e Drop threading from EventLog — single-threaded by design
Remove threading.Lock and call_soon_threadsafe. Both sync and async
paths are single-threaded (caller-driven sync, event-loop-bound
async), so there is no concurrent access to the buffer. Direct
fut.set_result() replaces call_soon_threadsafe for async notification
since the producer always runs on the event loop thread.
2026-04-16 11:21:26 -04:00
Nick Hollon 986c1cc2e3 Auto-close EventLogs, reject projection key conflicts, fix async interrupted/interrupts
Three usability fixes:

- Mux now auto-closes/fails EventLogs in projections (like StreamChannels),
  so transformers no longer need finalize/fail boilerplate
- StreamingHandler._setup() raises ValueError if a user transformer
  returns projection keys that collide with already-registered keys
- AsyncGraphRunStream.interrupted and .interrupts now await the pump
  task before returning, matching the output property's behavior
2026-04-16 10:24:52 -04:00
Nick Hollon 28cf5ed78d Unify EventLog — remove sync/async split from transformer API
Merge EventLog and AsyncEventLog into a single class with a _bind()
mechanism. EventLog starts unbound; the StreamMux calls _bind(is_async)
after transformer registration so only the correct iteration protocol
is available. This removes the is_async parameter from EventLog,
StreamChannel, and all transformer constructors — transformers just
create EventLog() and never need to know whether they run in sync or
async context.
2026-04-16 09:32:47 -04:00
Nick Hollon ca5d9a6bd7 Fix mypy errors, tighten init() return type, move import to top level 2026-04-15 19:21:08 -04:00
Nick Hollon ae3c823499 Make sync streaming caller-driven, no background thread
Replace the daemon thread pump with a pull-based model where the
caller's iteration on any projection drives the graph forward. EventLog
uses a _request_more callback instead of threading.Condition. Matches
v1's model where the caller's for loop is the pump. Async path is
unchanged (background task on the event loop).
2026-04-15 19:14:18 -04:00
Nick Hollon b72b5fefd0 Fix async transformer example to use async iteration on channel 2026-04-15 19:02:24 -04:00
Nick Hollon 5b1f86facc Split EventLog into sync and async classes
Separate EventLog (sync, __iter__) and AsyncEventLog (async, __aiter__)
with a shared _EventLogBase for the producer API. Thread is_async through
StreamMux, StreamChannel, and transformers so the right log type is
created based on whether stream() or astream() is called. Prevents
accidentally mixing sync and async iteration on the same log.
2026-04-15 18:59:47 -04:00
Nick Hollon cf966419d5 Add streaming comparison example (v1 vs v2 with custom transformer)
Side-by-side comparison of token-level LLM streaming using v1
graph.stream() and v2 StreamingHandler, both sync and async. Includes
a TokenMetrics custom transformer to demonstrate extensibility vs the
equivalent inline bookkeeping in v1.
2026-04-15 18:50:41 -04:00
Nick Hollon 8f03bf9f15 Add timestamp and eventId to ProtocolEvent
Add timestamp (ms epoch) to event params and eventId to the event
envelope, aligning the in-process event shape with the protocol spec.
Update tests to include timestamps and verify their presence.
2026-04-15 18:26:09 -04:00
Nick Hollon 0076da9008 feat(langgraph): add streaming transformer infrastructure and tests
Introduces the StreamingHandler, StreamMux, EventLog, StreamChannel,
and StreamTransformer abstractions for ergonomic streaming projections
over compiled graphs. Includes ValuesTransformer and MessagesTransformer
as built-in native projections, plus support for user-defined custom
transformers.
2026-04-15 18:12:13 -04:00
22 changed files with 5990 additions and 20 deletions
@@ -66,6 +66,9 @@ CONFIG_KEY_RUNTIME = sys.intern("__pregel_runtime")
# holds a `Runtime` instance with context, store, stream writer, etc.
CONFIG_KEY_RESUME_MAP = sys.intern("__pregel_resume_map")
# holds a mapping of task ns -> resume value for resuming tasks
CONFIG_KEY_STREAM_MESSAGES_V2 = sys.intern("__pregel_stream_messages_v2")
# when True, attach StreamMessagesHandlerV2 so content-block (v2) events
# flow through stream_mode="messages"; set by StreamingHandler only.
# --- Other constants ---
PUSH = sys.intern("__pregel_push")
@@ -107,6 +110,7 @@ RESERVED = {
CONFIG_KEY_CHECKPOINT_ID,
CONFIG_KEY_CHECKPOINT_NS,
CONFIG_KEY_RESUME_MAP,
CONFIG_KEY_STREAM_MESSAGES_V2,
# other constants
PUSH,
PULL,
+7
View File
@@ -1045,6 +1045,7 @@ class StateGraph(Generic[StateT, ContextT, InputT, OutputT]):
interrupt_after: All | list[str] | None = None,
debug: bool = False,
name: str | None = None,
transformers: Sequence[Callable[[], Any]] | None = None,
) -> CompiledStateGraph[StateT, ContextT, InputT, OutputT]:
"""Compiles the `StateGraph` into a `CompiledStateGraph` object.
@@ -1077,6 +1078,11 @@ class StateGraph(Generic[StateT, ContextT, InputT, OutputT]):
interrupt_after: An optional list of node names to interrupt after.
debug: A flag indicating whether to enable debug mode.
name: The name to use for the compiled graph.
transformers: Optional sequence of zero-arg factories returning
`StreamTransformer` instances. Registered on the compiled
graph and instantiated per-run whenever `stream_v2` /
`astream_v2` is called. Appended after the built-in
`ValuesTransformer` and `MessagesTransformer`.
Returns:
CompiledStateGraph: The compiled `StateGraph`.
@@ -1159,6 +1165,7 @@ class StateGraph(Generic[StateT, ContextT, InputT, OutputT]):
store=store,
cache=cache,
name=name or "LangGraph",
stream_transformers=transformers,
)
compiled._serde_allowlist = serde_allowlist
@@ -24,6 +24,11 @@ try:
except ImportError:
_StreamingCallbackHandler = object # type: ignore
try:
from langchain_core.tracers._streaming import _V2StreamingCallbackHandler
except ImportError:
_V2StreamingCallbackHandler = object # type: ignore
T = TypeVar("T")
Meta = tuple[tuple[str, ...], dict[str, Any]]
@@ -256,3 +261,78 @@ class StreamMessagesHandler(BaseCallbackHandler, _StreamingCallbackHandler):
**kwargs: Any,
) -> Any:
self.metadata.pop(run_id, None)
class StreamMessagesHandlerV2(StreamMessagesHandler, _V2StreamingCallbackHandler):
"""v2 variant of `StreamMessagesHandler`.
Declaring `_V2StreamingCallbackHandler` as a base flips
`BaseChatModel.invoke` to route through `_stream_chat_model_events`
(firing `on_stream_event`) instead of `_stream` (firing
`on_llm_new_token`). Inherits `on_stream_event` from the parent,
which forwards protocol events onto the messages stream channel.
Pregel attaches this class instead of the v1 handler only when
`StreamingHandler` opts in via the internal
`CONFIG_KEY_STREAM_MESSAGES_V2` config key; direct
`graph.stream(stream_mode="messages")` callers keep the v1
AIMessageChunk shape.
"""
def on_llm_new_token(
self,
token: str,
*,
chunk: ChatGenerationChunk | None = None,
run_id: UUID,
parent_run_id: UUID | None = None,
tags: list[str] | None = None,
**kwargs: Any,
) -> Any:
"""Intentional no-op — v1 chunks are not used on v2-flagged runs.
The v2 marker already steers `invoke` to the event generator, so
`on_llm_new_token` should not fire under normal routing. This
override stays a pass-through (no call to `super()`) to make
the intent explicit and to guard against any caller (e.g. a
node that calls `model.stream()` directly, which still fires
the v1 callback) leaking AIMessageChunks onto a v2-flagged
messages stream.
"""
# Intentionally empty: v2 handler does not forward v1 chunks.
def on_stream_event(
self,
event: dict[str, Any],
*,
run_id: UUID,
parent_run_id: UUID | None = None,
tags: list[str] | None = None,
**kwargs: Any,
) -> Any:
"""Forward a protocol event from `stream_v2` as a messages stream part.
Fires once per `MessagesData` event (`message-start`, per-block
`content-block-*`, `message-finish`). The transformer layer
correlates events back to a single `ChatModelStream` via
`metadata["run_id"]` — attached here so the v1
`stream_mode="messages"` output (which emits
`(AIMessageChunk, metadata)` via `on_llm_new_token`) keeps its
original metadata shape.
Lives on the v2 handler rather than the v1 base: content-block
events are a v2-only concept, and forwarding them only when the
v2 handler is attached keeps the message channel's shape
predictable for v1 callers.
"""
if meta := self.metadata.get(run_id):
# Record message_id on message-start so on_chain_end's
# dedupe skips the finalized AIMessage the node returns
# (otherwise the messages projection double-counts: once
# from streaming, once from the chain output).
if event.get("event") == "message-start":
msg_id = event.get("message_id")
if msg_id:
self.seen.add(msg_id)
v2_meta = {**meta[1], "run_id": str(run_id)}
self.stream((meta[0], "messages", (event, v2_meta)))
+176 -3
View File
@@ -73,6 +73,7 @@ from langgraph._internal._constants import (
CONFIG_KEY_RUNTIME,
CONFIG_KEY_SEND,
CONFIG_KEY_STREAM,
CONFIG_KEY_STREAM_MESSAGES_V2,
CONFIG_KEY_TASK_ID,
CONFIG_KEY_THREAD_ID,
ERROR,
@@ -133,7 +134,10 @@ from langgraph.pregel._loop import (
AsyncPregelLoop,
SyncPregelLoop,
)
from langgraph.pregel._messages import StreamMessagesHandler
from langgraph.pregel._messages import (
StreamMessagesHandler,
StreamMessagesHandlerV2,
)
from langgraph.pregel._read import DEFAULT_BOUND, PregelNode
from langgraph.pregel._retry import RetryPolicy
from langgraph.pregel._runner import PregelRunner
@@ -340,6 +344,69 @@ class NodeBuilder:
)
_STREAM_V2_MODES: list[StreamMode] = [
"values",
"updates",
"messages",
"custom",
"checkpoints",
"tasks",
"debug",
]
def _build_stream_factories(
compile_time: Sequence[Callable[..., Any]],
call_site: Sequence[Any] | None,
) -> list[Callable[..., Any]]:
"""Assemble the factory list handed to `StreamMux(factories=...)`.
Prepends the auto-registered built-ins — `ValuesTransformer`
(state snapshots backing `run.output` / `run.interrupted`),
`MessagesTransformer` (LLM token streaming), and
`SubgraphTransformer` (in-process subgraph handle discovery) —
then appends the graph's compile-time `stream_transformers`
followed by any call-site additions. Factories flow down into
subgraph mini-muxes, so per-scope instances propagate
automatically.
`LifecycleTransformer` is opt-in: add it via compile-time
`stream_transformers=[...]` or the per-call `transformers=[...]`
kwarg on `stream_v2()` / `astream_v2()`. Without it, no
`lifecycle` wire events are emitted and `run.lifecycle` is
absent.
"""
from langgraph.stream.transformers import (
MessagesTransformer,
SubgraphTransformer,
ValuesTransformer,
)
builtins: list[Callable[..., Any]] = [
ValuesTransformer,
MessagesTransformer,
SubgraphTransformer,
]
return [*builtins, *compile_time, *(call_site or ())]
def _merge_v2_messages_flag(
config: RunnableConfig | None,
) -> RunnableConfig:
"""Return a config with the v2 messages flag set in `configurable`.
Signals to pregel that `stream_mode="messages"` should attach
`StreamMessagesHandlerV2` for this call so invoke-time model runs
route through the v2 event generator and their protocol events
reach the messages channel.
"""
merged: RunnableConfig = dict(config or {}) # type: ignore[assignment]
configurable = dict(merged.get(CONF) or {})
configurable[CONFIG_KEY_STREAM_MESSAGES_V2] = True
merged[CONF] = configurable
return merged
class Pregel(
PregelProtocol[StateT, ContextT, InputT, OutputT],
Generic[StateT, ContextT, InputT, OutputT],
@@ -671,6 +738,7 @@ class Pregel(
config: RunnableConfig | None = None,
trigger_to_nodes: Mapping[str, Sequence[str]] | None = None,
name: str = "LangGraph",
stream_transformers: Sequence[Callable[..., Any]] | None = None,
**deprecated_kwargs: Unpack[DeprecatedKwargs],
) -> None:
if (
@@ -717,6 +785,9 @@ class Pregel(
self.config = config
self.trigger_to_nodes = trigger_to_nodes or {}
self.name = name
self._stream_transformers: tuple[Callable[..., Any], ...] = tuple(
stream_transformers or ()
)
self._serde_allowlist: set[tuple[str, ...]] | None = None
if auto_validate:
self.validate()
@@ -2626,8 +2697,13 @@ class Pregel(
# set up messages stream mode
if "messages" in stream_modes:
ns_ = cast(str | None, config[CONF].get(CONFIG_KEY_CHECKPOINT_NS))
messages_handler_cls = (
StreamMessagesHandlerV2
if config[CONF].get(CONFIG_KEY_STREAM_MESSAGES_V2)
else StreamMessagesHandler
)
run_manager.inheritable_handlers.append(
StreamMessagesHandler(
messages_handler_cls(
stream.put,
subgraphs,
parent_ns=tuple(ns_.split(NS_SEP)) if ns_ else None,
@@ -3003,8 +3079,13 @@ class Pregel(
if "messages" in stream_modes:
# namespace can be None in a root level graph?
ns_ = cast(str | None, config[CONF].get(CONFIG_KEY_CHECKPOINT_NS))
messages_handler_cls = (
StreamMessagesHandlerV2
if config[CONF].get(CONFIG_KEY_STREAM_MESSAGES_V2)
else StreamMessagesHandler
)
run_manager.inheritable_handlers.append(
StreamMessagesHandler(
messages_handler_cls(
stream_put,
subgraphs,
parent_ns=tuple(ns_.split(NS_SEP)) if ns_ else None,
@@ -3237,6 +3318,98 @@ class Pregel(
await asyncio.shield(run_manager.on_chain_error(e))
raise
def stream_v2(
self,
input: InputT | Command | None,
config: RunnableConfig | None = None,
*,
interrupt_before: All | Sequence[str] | None = None,
interrupt_after: All | Sequence[str] | None = None,
transformers: Sequence[Any] | None = None,
) -> Any:
"""Start a sync v2 streaming run driven by transformer projections.
Builds a `StreamMux` from the auto-registered built-ins
(`ValuesTransformer`, `MessagesTransformer`,
`SubgraphTransformer`), this graph's compile-time
`stream_transformers`, and any additional `transformers=`
supplied at the call site. Returns a `GraphRunStream` that
the caller drives by iterating any projection — no background
thread.
`LifecycleTransformer` (emits `lifecycle` wire events,
exposes `run.lifecycle`) is opt-in — add it via
`stream_transformers` at compile time or via the
`transformers=` kwarg here.
Args:
input: Graph input.
config: Optional runnable config forwarded to the graph.
interrupt_before: Nodes to interrupt before, if any.
interrupt_after: Nodes to interrupt after, if any.
transformers: Extra transformer instances appended after
compile-time `stream_transformers`.
Returns:
A `GraphRunStream` the caller iterates to drive the run.
"""
from langgraph.stream._mux import StreamMux
from langgraph.stream.run_stream import GraphRunStream
factories = _build_stream_factories(self._stream_transformers, transformers)
mux = StreamMux(factories=factories, is_async=False)
graph_iter = iter(
self.stream(
input,
_merge_v2_messages_flag(config),
stream_mode=_STREAM_V2_MODES,
subgraphs=True,
version="v2",
interrupt_before=interrupt_before,
interrupt_after=interrupt_after,
)
)
return GraphRunStream(graph_iter, mux)
async def astream_v2(
self,
input: InputT | Command | None,
config: RunnableConfig | None = None,
*,
interrupt_before: All | Sequence[str] | None = None,
interrupt_after: All | Sequence[str] | None = None,
transformers: Sequence[Any] | None = None,
) -> Any:
"""Async counterpart to `stream_v2`.
Returns an `AsyncGraphRunStream` whose projections can be awaited
concurrently; each subscribed cursor drives the pump when its
buffer is empty.
Args:
input: Graph input.
config: Optional runnable config forwarded to the graph.
interrupt_before: Nodes to interrupt before, if any.
interrupt_after: Nodes to interrupt after, if any.
transformers: Extra transformer instances appended after
compile-time `stream_transformers`.
"""
from langgraph.stream._mux import StreamMux
from langgraph.stream.run_stream import AsyncGraphRunStream
factories = _build_stream_factories(self._stream_transformers, transformers)
mux = StreamMux(factories=factories, is_async=True)
graph_aiter = self.astream(
input,
_merge_v2_messages_flag(config),
stream_mode=_STREAM_V2_MODES,
subgraphs=True,
version="v2",
interrupt_before=interrupt_before,
interrupt_after=interrupt_after,
).__aiter__()
return AsyncGraphRunStream(graph_aiter, mux)
@overload
def invoke(
self,
+5 -3
View File
@@ -653,10 +653,11 @@ 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)
updated_stream_modes.append(cast(StreamModeSDK, stream_mode))
else:
req_single = False
updated_stream_modes.extend(stream_mode)
for m in stream_mode:
updated_stream_modes.append(cast(StreamModeSDK, m))
else:
updated_stream_modes.append(default)
requested_stream_modes = updated_stream_modes.copy()
@@ -665,7 +666,8 @@ class RemoteGraph(PregelProtocol):
(config or {}).get(CONF, {}).get(CONFIG_KEY_STREAM)
)
if stream:
updated_stream_modes.extend(stream.modes)
for m in stream.modes:
updated_stream_modes.append(cast(StreamModeSDK, m))
# map "messages" to "messages-tuple"
if "messages" in updated_stream_modes:
updated_stream_modes.remove("messages")
@@ -0,0 +1,20 @@
"""Streaming infrastructure for LangGraph.
Compile a graph with `transformers=[...]` and call `graph.stream_v2()` /
`graph.astream_v2()` to drive a transformer pipeline that projects the
graph's raw events into ergonomic per-channel streams.
"""
from langgraph.stream._event_log import EventLog
from langgraph.stream._types import ProtocolEvent, StreamTransformer
from langgraph.stream.run_stream import AsyncGraphRunStream, GraphRunStream
from langgraph.stream.stream_channel import StreamChannel
__all__ = [
"AsyncGraphRunStream",
"EventLog",
"GraphRunStream",
"ProtocolEvent",
"StreamChannel",
"StreamTransformer",
]
@@ -0,0 +1,32 @@
from __future__ import annotations
import time
from typing import Any, cast
from langgraph.stream._types import ProtocolEvent, _ProtocolEventParams
from langgraph.types import StreamPart
def convert_to_protocol_event(part: StreamPart) -> ProtocolEvent:
"""Convert a v2 StreamPart to a ProtocolEvent.
Args:
part: A stream part with keys `type`, `ns`, `data`, and
optionally `interrupts` (present on values events).
Returns:
The equivalent ProtocolEvent.
"""
part_dict = cast(dict[str, Any], part)
params: _ProtocolEventParams = {
"namespace": list(part_dict["ns"]),
"timestamp": int(time.time() * 1000),
"data": part_dict["data"],
}
if "interrupts" in part_dict:
params["interrupts"] = part_dict["interrupts"]
return {
"type": "event",
"method": part_dict["type"],
"params": params,
}
@@ -0,0 +1,318 @@
from __future__ import annotations
import asyncio
from collections import deque
from collections.abc import AsyncIterator, Awaitable, Callable, Iterator
from typing import Generic, TypeVar
T = TypeVar("T")
class EventLog(Generic[T]):
"""Single-consumer drainable queue for streaming events.
Items are popped off the front as the consumer advances there is
no retention beyond what's currently queued. A log accepts exactly
one subscriber; a second `__iter__` / `__aiter__` call raises. Use
`tee(n)` / `atee(n)` for fan-out.
Starts unbound neither `__iter__` nor `__aiter__` is available
until the StreamMux calls `_bind(is_async)`. After binding, only
the matching iteration protocol works; the other raises `TypeError`.
Pump wiring (set by the run stream, not by `_bind`):
- `_request_more`: sync pump callable, returns True if a new
event was produced.
- `_arequest_more`: async pump coroutine factory, same contract.
Memory is bounded by caller pace: both sync and async use caller-
driven pumps, so each cursor advance produces at most one event.
The only shape where a log can accumulate meaningfully is
concurrent async consumers at unequal rates a slow consumer's
log grows while fast consumers drive the shared pump. That's the
documented tradeoff for concurrent consumption; consume at similar
rates or use a single consumer if memory matters.
Lazy-subscribe: `push` is a no-op when no subscriber has registered.
Transformers still execute `process()` (so scalar state like
`ValuesTransformer._latest` stays current); only the log append is
skipped.
"""
def __init__(self, maxlen: int | None = None, *, retain: bool = False) -> None:
"""Initialize an empty, unbound log.
Args:
maxlen: Accepted for forward compatibility; currently unused.
The caller-driven pump bounds memory naturally for
single-consumer use.
retain: If True, `push()` retains items regardless of whether
a consumer has subscribed yet. Used for projections
whose consumer only becomes visible after events have
already flowed (e.g. mini-mux logs inside dynamically
discovered subgraph handles, or the `lifecycle` channel
iterated after draining `values`). Subscription
exclusivity on `__iter__` is unchanged.
Raises:
ValueError: If `maxlen` is not a positive integer or `None`.
"""
if maxlen is not None and maxlen <= 0:
raise ValueError("EventLog maxlen must be a positive int or None")
self._items: deque[T] = deque()
self._maxlen: int | None = maxlen
self._closed = False
self._error: BaseException | None = None
# Binding state — None means unbound.
self._is_async: bool | None = None
# Flipped on first __iter__ / __aiter__. Pre-subscription
# pushes are silent no-ops unless `_retain` is True.
self._subscribed = False
self._retain = retain
# Pump wiring set by the run stream after bind.
self._request_more: Callable[[], bool] | None = None
self._arequest_more: Callable[[], Awaitable[bool]] | None = None
# ------------------------------------------------------------------
# Binding
# ------------------------------------------------------------------
def _bind(self, *, is_async: bool) -> None:
"""Bind this log to sync or async mode.
Called by the StreamMux after transformer registration. Must be
called exactly once before any iteration.
Args:
is_async: True to enable async iteration, False for sync.
Raises:
RuntimeError: If the log has already been bound.
"""
if self._is_async is not None:
raise RuntimeError("EventLog is already bound")
self._is_async = is_async
# ------------------------------------------------------------------
# Producer API
# ------------------------------------------------------------------
def push(self, item: T) -> None:
"""Append an item. No-op when no subscriber is registered.
Non-blocking in both sync and async matches v1's
`put_nowait` producer shape. Memory is bounded by caller pace
via the caller-driven pump.
When `retain=True` was set at construction, items are appended
regardless of subscription for projections whose consumer
only reaches them after events have already flowed.
Raises:
RuntimeError: If the log is closed (and subscribed).
"""
if not self._subscribed and not self._retain:
return
if self._closed:
raise RuntimeError("Cannot push to a closed EventLog")
self._items.append(item)
def close(self) -> None:
"""Mark the log as complete."""
self._closed = True
def fail(self, err: BaseException) -> None:
"""Mark the log as errored.
Args:
err: The exception to surface to the subscriber.
"""
self._error = err
self._closed = True
# ------------------------------------------------------------------
# Sync iteration (caller-driven pump)
# ------------------------------------------------------------------
def __iter__(self) -> Iterator[T]:
"""Subscribe and return a sync cursor. Can be called only once.
Raises:
TypeError: If the log is unbound or bound to async mode.
RuntimeError: If the log already has a subscriber.
"""
if self._is_async is None:
raise TypeError(
"EventLog has not been bound yet. "
"Register the transformer with a StreamMux first."
)
if self._is_async:
raise TypeError(
"This EventLog is bound to async mode — use 'async for' instead."
)
if self._subscribed:
raise RuntimeError(
"EventLog already has a subscriber; use .tee(n) for fan-out."
)
self._subscribed = True
return self._sync_cursor()
def _sync_cursor(self) -> Iterator[T]:
while True:
if self._items:
yield self._items.popleft()
elif self._closed:
if self._error is not None:
raise self._error
return
elif self._request_more is not None:
if not self._request_more():
if not self._items and not self._closed:
return
else:
return
# ------------------------------------------------------------------
# Async iteration (caller-driven pump)
# ------------------------------------------------------------------
def __aiter__(self) -> AsyncIterator[T]:
"""Subscribe and return an async cursor. Can be called only once.
Raises:
TypeError: If the log is unbound or bound to sync mode.
RuntimeError: If the log already has a subscriber.
"""
if self._is_async is None:
raise TypeError(
"EventLog has not been bound yet. "
"Register the transformer with a StreamMux first."
)
if not self._is_async:
raise TypeError("This EventLog is bound to sync mode — use 'for' instead.")
if self._subscribed:
raise RuntimeError(
"EventLog already has a subscriber; use .atee(n) for fan-out."
)
self._subscribed = True
return self._async_cursor()
async def _async_cursor(self) -> AsyncIterator[T]:
while True:
if self._items:
yield self._items.popleft()
elif self._closed:
if self._error is not None:
raise self._error
return
elif self._arequest_more is not None:
if not await self._arequest_more():
if not self._items and not self._closed:
return
else:
return
# ------------------------------------------------------------------
# Fan-out via tee
# ------------------------------------------------------------------
def tee(self, n: int = 2) -> tuple[Iterator[T], ...]:
"""Subscribe and return `n` independent sync iterators.
Each branch has its own buffer; items pulled from the
underlying cursor are copied into every branch. Branches are
naturally bounded by caller pace since the sync pump is
caller-driven.
Args:
n: Number of branches to create. Must be >= 1.
Returns:
A tuple of `n` iterators over the same underlying stream.
Raises:
TypeError: If the log is unbound or bound to async mode.
RuntimeError: If the log already has a subscriber.
ValueError: If `n` < 1.
"""
if n < 1:
raise ValueError("tee() requires n >= 1")
source = self.__iter__()
buffers: list[deque[T]] = [deque() for _ in range(n)]
exhausted = [False]
def branch(i: int) -> Iterator[T]:
buf = buffers[i]
while True:
if buf:
yield buf.popleft()
elif exhausted[0]:
return
else:
try:
item = next(source)
except StopIteration:
exhausted[0] = True
return
for b in buffers:
b.append(item)
return tuple(branch(i) for i in range(n))
def atee(self, n: int = 2) -> tuple[AsyncIterator[T], ...]:
"""Subscribe and return `n` independent async iterators.
Caller-driven fan-out: each branch's `__anext__` either pops
from its own buffer or, under a shared `asyncio.Lock`, pulls
one item from the underlying cursor and distributes it to
every branch's buffer.
Args:
n: Number of branches to create. Must be >= 1.
Returns:
A tuple of `n` async iterators over the same underlying
stream.
Raises:
TypeError: If the log is unbound or bound to sync mode.
RuntimeError: If the log already has a subscriber.
ValueError: If `n` < 1.
"""
if n < 1:
raise ValueError("atee() requires n >= 1")
source = self.__aiter__()
buffers: list[deque[T]] = [deque() for _ in range(n)]
exhausted = [False]
error: list[BaseException | None] = [None]
lock = asyncio.Lock()
async def branch(i: int) -> AsyncIterator[T]:
buf = buffers[i]
while True:
if buf:
yield buf.popleft()
continue
if exhausted[0]:
if error[0] is not None:
raise error[0]
return
async with lock:
if buf or exhausted[0]:
continue
try:
item = await source.__anext__()
except StopAsyncIteration:
exhausted[0] = True
continue
except Exception as e:
error[0] = e
exhausted[0] = True
continue
for b in buffers:
b.append(item)
return tuple(branch(i) for i in range(n))
+514
View File
@@ -0,0 +1,514 @@
from __future__ import annotations
import asyncio
import time
from collections.abc import Awaitable, Callable
from typing import Any
from langgraph.stream._event_log import EventLog
from langgraph.stream._types import (
ProtocolEvent,
StreamTransformer,
transformer_requires_async,
)
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.
Owns the main event log and routes events through a transformer
pipeline. StreamChannels discovered in transformer projections are
auto-wired so that every `push()` also injects a `ProtocolEvent`
into the main log.
Pass `is_async=True` when the mux will be consumed via async
iteration (`handler.astream()`). All EventLog and StreamChannel
instances discovered during registration are automatically bound
to the matching mode.
Attributes:
extensions: Merged projection dict across all registered
transformers. Treat as read-only mutations won't be
reflected back in individual transformers' state.
native_keys: Projection keys contributed by transformers with
`_native = True`.
"""
def __init__(
self,
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.
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: 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, 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] = []
self._channels: list[StreamChannel[Any]] = []
self._logs: list[EventLog[Any]] = []
self._seq = 0
self.extensions: dict[str, Any] = {}
self.native_keys: set[str] = set()
self._projection_owners: dict[str, str] = {}
self._transformer_by_key: dict[str, StreamTransformer] = {}
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,
)
# Mini-muxes are created during the pump, after the first
# event at the child's scope has already been dispatched.
# Consumers reach child projections via the parent's
# `subgraphs` handle — necessarily after that first event.
# Flip retain on every log and channel so pushes are buffered
# until the consumer subscribes.
child._events._retain = True
for value in child.extensions.values():
if isinstance(value, EventLog):
value._retain = True
elif isinstance(value, StreamChannel):
value._log._retain = True
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.
Calls `transformer.init()`, stores the transformer for event
processing, binds any EventLog or StreamChannel instances in
the projection, and merges the projection into `extensions`.
"""
if transformer_requires_async(transformer) and not self._is_async:
raise RuntimeError(
f"{type(transformer).__name__} requires an async run — "
"it overrides aprocess/afinalize/afail or sets "
"requires_async=True. Use astream(), not stream()."
)
projection = transformer.init()
if not isinstance(projection, dict):
raise TypeError(
f"StreamTransformer.init() must return a dict, "
f"got {type(projection).__name__}"
)
conflicts = set(projection) & set(self.extensions)
if conflicts:
attributions = ", ".join(
f"{key!r} (owned by {self._projection_owners[key]})"
for key in sorted(conflicts)
)
raise ValueError(
f"Transformer {type(transformer).__name__} returned "
f"projection keys that conflict with already-registered "
f"keys: {attributions}"
)
self._transformers.append(transformer)
is_native = bool(getattr(transformer, "_native", False))
self._bind_and_wire(projection, is_native=is_native)
self.extensions.update(projection)
owner_name = type(transformer).__name__
for key in projection:
self._projection_owners[key] = owner_name
self._transformer_by_key[key] = transformer
if is_native:
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
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.
Seq is assigned right before an event enters the main log, not
before the transformer pipeline runs. This ensures that events
auto-forwarded from StreamChannels during `process()` get
earlier seq numbers than the original event, preserving
monotonic ordering in the log.
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:
self._seq += 1
event["seq"] = self._seq
self._events.push(event)
def close(self) -> None:
"""Finalize all transformers, close all projections and the main log.
EventLogs and StreamChannels discovered in transformer
projections are auto-closed after `finalize()` runs
transformers don't need to close them manually. If any
transformer's `finalize()` raises, the remaining transformers,
projections, and the main log are still closed; the first error
is re-raised after cleanup completes.
Raises:
BaseException: The first error raised by a transformer's
`finalize()`, re-raised after cleanup finishes.
"""
first_error: BaseException | None = None
for transformer in self._transformers:
try:
transformer.finalize()
except BaseException as e:
if first_error is None:
first_error = e
for log in self._logs:
if not log._closed:
log.close()
for ch in self._channels:
if not ch._log._closed:
ch._close()
self._events.close()
if first_error is not None:
raise first_error
def fail(self, err: BaseException) -> None:
"""Fail all transformers, projections, and the main log.
EventLogs and StreamChannels discovered in transformer
projections are auto-failed transformers don't need to fail
them manually. If any transformer's `fail()` raises, the
remaining transformers, projections, and the main log are still
failed.
Args:
err: The exception that ended the run.
"""
for transformer in self._transformers:
try:
transformer.fail(err)
except BaseException:
pass
for log in self._logs:
if not log._closed:
log.fail(err)
for ch in self._channels:
if not ch._log._closed:
ch._fail(err)
self._events.fail(err)
# ------------------------------------------------------------------
# Async dispatch
# ------------------------------------------------------------------
async def apush(self, event: ProtocolEvent) -> None:
"""Dispatch an event on the async lane.
Awaits each transformer's `aprocess` in registration order
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
caller-driven pump; see `EventLog` for the full tradeoff story.
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:
self._seq += 1
event["seq"] = self._seq
self._events.push(event)
async def aclose(self) -> None:
"""Finalize on the async lane.
Awaits every task started via `StreamTransformer.schedule()`
across all transformers, then calls `afinalize()` on each,
then auto-closes logs, channels, and the main event log.
If any scheduled task raised under `on_error="raise"`, or any
transformer's `afinalize` raises, the exception propagates.
The caller (the pump) handles it by routing into `afail`.
Raises:
BaseException: The first scheduled-task or `afinalize`
error, re-raised after cleanup.
"""
pending = self._collect_scheduled_tasks()
if pending:
results = await asyncio.gather(*pending, return_exceptions=True)
first_err = next(
(
r
for r in results
if isinstance(r, BaseException)
and not isinstance(r, asyncio.CancelledError)
),
None,
)
if first_err is not None:
raise first_err
first_error: BaseException | None = None
for transformer in self._transformers:
try:
await transformer.afinalize()
except BaseException as e:
if first_error is None:
first_error = e
for log in self._logs:
if not log._closed:
log.close()
for ch in self._channels:
if not ch._log._closed:
ch._close()
self._events.close()
if first_error is not None:
raise first_error
async def afail(self, err: BaseException) -> None:
"""Fail on the async lane.
Cancels every scheduled task across all transformers, awaits
them to completion, then runs each transformer's `afail` hook
and auto-fails logs, channels, and the main event log.
Args:
err: The exception that ended the run.
"""
pending = self._collect_scheduled_tasks()
for task in pending:
task.cancel()
if pending:
await asyncio.gather(*pending, return_exceptions=True)
for transformer in self._transformers:
try:
await transformer.afail(err)
except BaseException:
pass
for log in self._logs:
if not log._closed:
log.fail(err)
for ch in self._channels:
if not ch._log._closed:
ch._fail(err)
if not self._events._closed:
self._events.fail(err)
def _collect_scheduled_tasks(self) -> list[asyncio.Task[Any]]:
"""Return a snapshot of in-flight tasks scheduled via transformers."""
return [
task
for transformer in self._transformers
for task in getattr(transformer, "_stream_scheduled_tasks", ())
if not task.done()
]
# ------------------------------------------------------------------
# Binding and StreamChannel auto-wiring
# ------------------------------------------------------------------
def _bind_and_wire(
self, projection: dict[str, Any], *, is_native: bool = False
) -> None:
"""Bind and wire EventLog / StreamChannel instances in a projection.
`is_native` controls wire naming: native transformer channels
emit events with `method` equal to the channel name, while
non-native channels get a `custom:` prefix to keep user-defined
projections from colliding with built-in method names.
"""
for value in projection.values():
if isinstance(value, StreamChannel):
value._bind(is_async=self._is_async)
self._channels.append(value)
channel_name = value.name
def _make_forward(name: str, native: bool) -> Callable[[Any], None]:
def _forward(item: Any) -> None:
self._forward(name, item, native=native)
return _forward
value._wire(_make_forward(channel_name, is_native))
elif isinstance(value, EventLog):
value._bind(is_async=self._is_async)
self._logs.append(value)
def _forward(self, channel_name: str, item: Any, *, native: bool) -> None:
"""Inject a ProtocolEvent for a StreamChannel push.
Forwarded events bypass the transformer pipeline to avoid
infinite recursion (a transformer that pushes to a channel
during `process()` would re-trigger itself). These events are
visible in the main event log but are not passed through
transformers' `process()` methods.
Native transformers emit with `method` equal to the channel
name (e.g. `"lifecycle"`); non-native transformers get a
`custom:` prefix so user-defined projections can't collide
with built-in method names.
"""
self._seq += 1
method = channel_name if native else f"custom:{channel_name}"
event: ProtocolEvent = {
"type": "event",
"seq": self._seq,
"method": method,
"params": {
"namespace": [],
"timestamp": int(time.time() * 1000),
"data": item,
},
}
self._events.push(event)
+303
View File
@@ -0,0 +1,303 @@
from __future__ import annotations
import asyncio
import logging
from abc import ABC, abstractmethod
from collections.abc import Coroutine
from typing import Any, ClassVar, Literal
from typing_extensions import NotRequired, TypedDict
_logger = logging.getLogger(__name__)
class _ProtocolEventParams(TypedDict):
"""Parameters for a protocol event.
`timestamp` is wall-clock milliseconds since the epoch and can go
backwards across NTP adjustments use `ProtocolEvent.seq` for
ordering.
"""
namespace: list[str]
timestamp: int
data: Any
interrupts: NotRequired[tuple[Any, ...]]
class ProtocolEvent(TypedDict):
"""A protocol event emitted by the streaming infrastructure.
Wraps a raw stream part (values, messages, custom, etc.) in a uniform
envelope with a monotonic sequence number assigned by the StreamMux.
Consumers that need a total order across events should use `seq`, not
`params.timestamp` (which is wall-clock and not monotonic).
"""
type: Literal["event"]
eventId: NotRequired[str]
seq: NotRequired[int]
method: str # StreamMode value: "values", "messages", "custom", etc.
params: _ProtocolEventParams
class StreamTransformer(ABC):
"""Extension point for custom stream projections.
Transformers observe protocol events flowing through the StreamMux and
build typed derived projections (EventLogs, StreamChannels, promises,
etc.).
Set `_native = True` on a transformer to have its projection keys
exposed as direct attributes on the run stream (in addition to
appearing in `run.extensions`).
Subclasses must implement `init` and override at least one of
`process` / `aprocess`. The `finalize` / `afinalize` and `fail` /
`afail` hooks are optional the default implementations are no-ops.
EventLog and StreamChannel instances in the projection dict are
auto-closed / auto-failed by the mux, so most transformers don't
need `finalize` or `fail` at all.
Transformers that need async work pick the async lane by:
1. Overriding `aprocess` (and optionally `afinalize` / `afail`), or
2. Calling `self.schedule(coro)` from inside a sync `process`, or
3. Setting `requires_async = True` explicitly.
The mux detects these cases at registration and raises if they're
used under sync `stream()` they only work under `astream()`.
Use `aprocess` when the pump must wait for async work before the
next transformer sees the event (e.g. PII redaction that mutates
`event` in place). Use `schedule()` for decoupled async work whose
result lands on an independent projection (e.g. async moderation
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
`process`). The mux also auto-detects the async lane when
`aprocess`, `afinalize`, or `afail` is overridden.
"""
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]:
"""Return the projection dict.
Keys become entries in `run.extensions`. If the transformer has
`_native = True`, keys are also set as direct attributes on the
run stream.
StreamChannel instances in the return value are automatically
wired by the StreamMux for protocol event auto-forwarding.
"""
...
def process(self, event: ProtocolEvent) -> bool:
"""Handle an event on the sync lane.
Called for every event before it is appended to the main event
log. Subclasses must override either `process` or `aprocess`.
The default raises so a missing override fails loudly rather
than silently passing every event through.
Args:
event: The protocol event to observe.
Returns:
True to keep the event in the main log, False to suppress it.
"""
raise NotImplementedError(
f"{type(self).__name__} must override process() or aprocess()"
)
async def aprocess(self, event: ProtocolEvent) -> bool:
"""Handle an event on the async lane.
The mux awaits this before dispatching to the next transformer,
so a slow `aprocess` serializes the pipeline. Use it only when
a later transformer or a consumer reading the event
synchronously must see the result of the async work (e.g.
PII redaction that mutates `event` in place).
The default delegates to `process`, so purely-sync transformers
run unchanged under `astream()`.
Args:
event: The protocol event to observe.
Returns:
True to keep the event in the main log, False to suppress it.
"""
return self.process(event)
def finalize(self) -> None:
"""Called when the run ends normally (sync lane).
Override to close EventLogs, resolve promises, or perform other
teardown. StreamChannel instances are auto-closed by the mux.
"""
async def afinalize(self) -> None:
"""Called when the run ends normally (async lane).
By the time this runs, the mux has already awaited every task
started via `schedule()`, so EventLogs can be closed here
without a last-task-wins race.
The default delegates to `finalize`.
"""
self.finalize()
def fail(self, err: BaseException) -> None:
"""Called when the run ends with an error (sync lane).
Override to fail EventLogs, reject promises, or perform other
teardown. StreamChannel instances are auto-failed by the mux.
Args:
err: The exception that ended the run.
"""
async def afail(self, err: BaseException) -> None:
"""Called when the run ends with an error (async lane).
The mux cancels and awaits every task started via `schedule()`
before calling this, so cleanup doesn't race with in-flight work.
The default delegates to `fail`.
Args:
err: The exception that ended the run.
"""
self.fail(err)
# ------------------------------------------------------------------
# Scheduled async work
# ------------------------------------------------------------------
def schedule(
self,
coro: Coroutine[Any, Any, Any],
*,
on_error: Literal["log", "raise"] = "log",
) -> asyncio.Task[Any]:
"""Schedule a coroutine tied to this transformer's lifecycle.
The mux holds the task reference, awaits all scheduled tasks
during `aclose()` before calling `afinalize()`, and cancels
them on `afail()`. Authors don't need to track tasks or
implement the last-task-closes-the-log dance.
Requires a running event loop call only under `astream()`.
Set `requires_async = True` on the class so registration under
sync `stream()` fails fast with a clear message.
Args:
coro: The coroutine to run. Its lifecycle is owned by the
mux from this point on.
on_error: `"log"` (default) catches and logs any exception
the coroutine raises, so a single failure doesn't tear
down the run. `"raise"` lets the exception propagate
when the mux joins pendings, converting the close path
into the fail path.
Returns:
The asyncio Task. Authors rarely need to await it directly
consumers read results from whatever projection the
coroutine pushes into.
Raises:
RuntimeError: If called without a running event loop (i.e.
under sync `stream()` rather than `astream()`).
"""
try:
asyncio.get_running_loop()
except RuntimeError:
raise RuntimeError(
f"{type(self).__name__}.schedule() requires a running "
"event loop; this transformer must run under astream(), "
"not stream(). Set requires_async=True on the class so "
"this fails at registration rather than at first event."
) from None
wrapped = self._wrap_scheduled(coro) if on_error == "log" else coro
task = asyncio.create_task(wrapped)
tasks = self._scheduled_task_set()
tasks.add(task)
task.add_done_callback(tasks.discard)
return task
@staticmethod
async def _wrap_scheduled(coro: Coroutine[Any, Any, Any]) -> Any:
try:
return await coro
except asyncio.CancelledError:
raise
except BaseException:
_logger.exception("Scheduled StreamTransformer task failed")
def _scheduled_task_set(self) -> set[asyncio.Task[Any]]:
"""Return the lazily-allocated task set.
Avoids requiring subclasses to call `super().__init__()`.
"""
tasks: set[asyncio.Task[Any]] | None = getattr(
self, "_stream_scheduled_tasks", None
)
if tasks is None:
tasks = set()
self._stream_scheduled_tasks = tasks
return tasks
def transformer_requires_async(transformer: StreamTransformer) -> bool:
"""Return True if the transformer needs a running event loop.
A transformer requires async if it explicitly opts in
(`requires_async = True`) or overrides any of the async-lane methods
(`aprocess`, `afinalize`, `afail`).
Args:
transformer: The transformer to inspect.
Returns:
True if the transformer cannot run under sync `stream()`.
"""
if transformer.requires_async:
return True
cls = type(transformer)
for name in ("aprocess", "afinalize", "afail"):
if getattr(cls, name) is not getattr(StreamTransformer, name):
return True
return False
@@ -0,0 +1,408 @@
from __future__ import annotations
import asyncio
from collections.abc import AsyncIterator, Awaitable, Callable, Iterator, Mapping
from types import MappingProxyType, TracebackType
from typing import TYPE_CHECKING, Any
from langgraph.stream._convert import convert_to_protocol_event
from langgraph.stream._mux import StreamMux
from langgraph.stream._types import ProtocolEvent
if TYPE_CHECKING:
from langgraph.stream.transformers import ValuesTransformer
def _drive_until_done(pump: Callable[[], bool]) -> None:
"""Call the sync pump until it returns False."""
while pump():
pass
async def _adrive_until_done(pump: Callable[[], Awaitable[bool]]) -> None:
"""Call the async pump until it returns False."""
while await pump():
pass
class BaseRunStream:
"""Shared shape for any object that wraps a `StreamMux`.
Root (`GraphRunStream` / `AsyncGraphRunStream`) and scoped
(`SubgraphRunStream`) streams both compose a `StreamMux`. The mux
owns the projections `values`, `messages`, `subgraphs`, and any
user-registered keys all exposed via `extensions`. Native
projections (`_native = True`) are also bound as direct attributes
(`run.values`, `run.messages`, ) for ergonomics.
Raw iteration (`for event in run` / `async for event in run`) and
the `interleave(...)` helper both live here so every subclass
behaves consistently. Subclasses only add pump ownership, scope
metadata, or sync/async flavor.
"""
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])
@property
def _values_transformer(self) -> ValuesTransformer:
"""Look up the `ValuesTransformer` registered on this mux.
`output` / `interrupted` / `interrupts` need scalar state from
the `ValuesTransformer` without threading it through the
constructor. Raises if none is registered `stream_v2` /
`astream_v2` always register one, so hitting this path means
the caller assembled the mux themselves and forgot.
"""
from langgraph.stream.transformers import ValuesTransformer
for t in self._mux._transformers:
if isinstance(t, ValuesTransformer):
return t
raise RuntimeError(
"No ValuesTransformer is registered on this mux — "
"`output` / `interrupted` / `interrupts` are unavailable."
)
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`,
raw events, or `output`) drives the graph forward. No background
thread is used the caller's `for` loop is the pump.
Projections are single-consumer iterating `run.values` twice
raises. Use `projection.tee(n)` if you genuinely need fan-out.
"""
def __init__(
self,
graph_iter: Iterator[Any],
mux: StreamMux,
) -> None:
"""Initialize the run stream.
Args:
graph_iter: Pull-based iterator over the graph's stream.
mux: The StreamMux owning projections and the main log.
Must have a `ValuesTransformer` registered for
`output` / `interrupted` / `interrupts` to work.
"""
super().__init__(mux)
self._graph_iter = graph_iter
self._exhausted = False
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.
"""
if self._exhausted:
return False
try:
part = next(self._graph_iter)
except StopIteration:
self._mux.close()
self._exhausted = True
return False
except Exception as e:
self._mux.fail(e)
self._exhausted = True
return False
self._mux.push(convert_to_protocol_event(part))
return True
def abort(self) -> None:
"""Stop the run early.
Closes the mux and marks the stream exhausted. The graph
iterator is dropped; any in-flight nodes see the closure on
their next yield point. Idempotent.
"""
if self._exhausted:
return
self._exhausted = True
try:
self._mux.close()
except Exception:
pass
def __enter__(self) -> GraphRunStream:
return self
def __exit__(
self,
exc_type: type[BaseException] | None,
exc: BaseException | None,
tb: TracebackType | None,
) -> None:
self.abort()
@property
def output(self) -> dict[str, Any] | None:
"""Drive the run to completion and return the final state."""
_drive_until_done(self._pump_next)
err = self._values_transformer.error
if err is not None:
raise err
return self._values_transformer._latest
@property
def interrupted(self) -> bool:
"""Drive the run to completion, then return whether it was interrupted.
Raises:
BaseException: If the run ended with an error.
"""
_drive_until_done(self._pump_next)
err = self._values_transformer.error
if err is not None:
raise err
return self._values_transformer._interrupted
@property
def interrupts(self) -> list[Any]:
"""Drive the run to completion, then return interrupt payloads.
Raises:
BaseException: If the run ended with an error.
"""
_drive_until_done(self._pump_next)
err = self._values_transformer.error
if err is not None:
raise err
return self._values_transformer._interrupts
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.
Projections are single-consumer a second `aiter(run.values)`
raises. Use `projection.tee(n)` for fan-out.
Use as an async context manager to guarantee clean shutdown on
early exit:
```python
async with await handler.astream(input) as run:
async for msg in run.messages:
...
```
"""
def __init__(
self,
graph_aiter: AsyncIterator[Any],
mux: StreamMux,
) -> None:
"""Initialize the async run stream.
Args:
graph_aiter: Async iterator over the graph's stream.
mux: The StreamMux owning projections and the main log.
Must have a `ValuesTransformer` registered for
`output` / `interrupted` / `interrupts` to work.
"""
super().__init__(mux)
self._graph_aiter = graph_aiter
self._exhausted = False
self._pump_cond = asyncio.Condition()
self._pumping = False
mux.bind_apump(self._apump_next)
async def _apump_next(self) -> bool:
"""Drive one pump step, or wait for the active pumper to drive one.
"Take-a-number" semantics: at most one task at a time calls
`graph_aiter.__anext__()` (asyncio iterators can't be advanced
concurrently). Other callers wait on a Condition that the
active pumper notifies after each step. This lets a "passive"
consumer one whose projection's buffer is being filled by the
active pumper's push — wake up as soon as its data lands,
instead of queueing on the pump and only observing its data one
graph event late.
`except Exception` is intentional `CancelledError` and other
`BaseException` subclasses propagate, matching asyncio's
cancellation contract.
Returns:
True if a pump step completed (by this task or another),
False if the graph is exhausted.
"""
async with self._pump_cond:
if self._exhausted:
return False
if self._pumping:
# Another task is pumping; wait for its progress signal.
await self._pump_cond.wait()
return not self._exhausted
self._pumping = True
try:
try:
part = await self._graph_aiter.__anext__()
except StopAsyncIteration:
self._exhausted = True
await self._mux.aclose()
return False
except Exception as e:
self._exhausted = True
await self._mux.afail(e)
return False
await self._mux.apush(convert_to_protocol_event(part))
return True
finally:
async with self._pump_cond:
self._pumping = False
self._pump_cond.notify_all()
async def abort(self) -> None:
"""Stop the run early.
Marks the stream exhausted, wakes any pump-waiters, and closes
the mux. Any `apush` blocked on backpressure wakes and returns
without appending. Idempotent.
"""
async with self._pump_cond:
if self._exhausted:
return
self._exhausted = True
self._pump_cond.notify_all()
try:
await self._mux.aclose()
except Exception:
pass
async def __aenter__(self) -> AsyncGraphRunStream:
return self
async def __aexit__(
self,
exc_type: type[BaseException] | None,
exc: BaseException | None,
tb: TracebackType | None,
) -> None:
await self.abort()
async def output(self) -> dict[str, Any] | None:
"""Drive the run to completion and return the final state.
Methods (not properties) on the async lane so `run.output`
without `await` raises at type-check time instead of silently
yielding a coroutine object.
Example:
```python
output = await run.output()
```
Raises:
BaseException: If the run ended with an error.
"""
await _adrive_until_done(self._apump_next)
if (err := self._values_transformer.error) is not None:
raise err
return self._values_transformer._latest
async def interrupted(self) -> bool:
"""Drive the run to completion and return whether it was interrupted.
Raises:
BaseException: If the run ended with an error.
"""
await _adrive_until_done(self._apump_next)
if (err := self._values_transformer.error) is not None:
raise err
return self._values_transformer._interrupted
async def interrupts(self) -> list[Any]:
"""Drive the run to completion and return interrupt payloads.
Raises:
BaseException: If the run ended with an error.
"""
await _adrive_until_done(self._apump_next)
if (err := self._values_transformer.error) is not None:
raise err
return self._values_transformer._interrupts
@@ -0,0 +1,115 @@
from __future__ import annotations
from collections.abc import AsyncIterator, Callable, Iterator
from typing import Generic, TypeVar
from langgraph.stream._event_log import EventLog
T = TypeVar("T")
class StreamChannel(Generic[T]):
"""A named projection channel with optional protocol auto-forwarding.
Wraps an event log and declares a protocol channel name. When the
StreamMux detects a StreamChannel in a transformer's `init()`
return value, it automatically wires every `push()` to inject a
`ProtocolEvent` into the main event stream using the channel's
name as the method.
Auto-forwarded events bypass the transformer pipeline other
transformers' `process()` / `aprocess()` methods do not see
`custom:<name>` events produced by a channel push. This prevents a
transformer that pushes to its own channel during `process()` from
re-triggering itself, but it also means filter- or tap-style
transformers cannot observe channel output from peer transformers.
Consumers that need that should iterate the main event stream.
In-process consumers iterate the channel directly (`for item in ch`
or `async for item in ch`). Remote SDK clients subscribe via
`session.subscribe("custom:<channelName>")`.
Like EventLog, a StreamChannel starts unbound. The mux calls
`_bind(is_async)` during registration so the correct iteration
protocol is available by the time user code sees it.
Lifecycle (`_close` / `_fail`) is managed by the mux transformers
using only StreamChannels don't need `finalize` or `fail` hooks.
"""
def __init__(
self, name: str, *, maxlen: int | None = None, retain: bool = False
) -> None:
"""Initialize the channel with an empty inner log.
Args:
name: The protocol channel name used for auto-forwarded
events (`custom:<name>` on the wire).
maxlen: Optional retention cap on the inner EventLog. See
`EventLog.__init__` for semantics.
retain: If True, the inner log retains pushes before any
consumer subscribes needed for channels whose
consumer iterates after events have already flowed
(e.g. `lifecycle` inspected after draining `values`).
"""
self.name = name
self._log: EventLog[T] = EventLog(maxlen=maxlen, retain=retain)
self._wire_fn: Callable[[T], None] | None = None
def _bind(self, *, is_async: bool) -> None:
"""Bind the underlying event log to sync or async mode.
Args:
is_async: True for async iteration, False for sync.
"""
self._log._bind(is_async=is_async)
def push(self, item: T) -> None:
"""Append an item to the log and auto-forward if wired.
Args:
item: The item to push.
"""
self._log.push(item)
if self._wire_fn is not None:
self._wire_fn(item)
# ------------------------------------------------------------------
# Mux lifecycle hooks (not called by transformers directly)
# ------------------------------------------------------------------
def _wire(self, fn: Callable[[T], None]) -> None:
"""Install the auto-forward callback (called by StreamMux)."""
self._wire_fn = fn
def _close(self) -> None:
"""Close the underlying log (called by StreamMux on run end)."""
self._log.close()
def _fail(self, err: BaseException) -> None:
"""Fail the underlying log (called by StreamMux on run error)."""
self._log.fail(err)
# ------------------------------------------------------------------
# Iteration — delegates to the inner event log (multi-cursor)
# ------------------------------------------------------------------
def __iter__(self) -> Iterator[T]:
return iter(self._log)
def __aiter__(self) -> AsyncIterator[T]:
return self._log.__aiter__()
def tee(self, n: int = 2) -> tuple[Iterator[T], ...]:
"""Fan out the channel into `n` independent sync iterators.
Delegates to the underlying EventLog's `tee()`.
"""
return self._log.tee(n)
def atee(self, n: int = 2) -> tuple[AsyncIterator[T], ...]:
"""Fan out the channel into `n` independent async iterators.
Delegates to the underlying EventLog's `atee()`.
"""
return self._log.atee(n)
@@ -0,0 +1,585 @@
from __future__ import annotations
import logging
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 (
AsyncChatModelStream,
ChatModelStream,
)
from langchain_core.messages import AIMessageChunk, BaseMessage
from langchain_protocol.protocol import CheckpointRef, MessagesData
from typing_extensions import TypedDict
from langgraph.errors import GraphInterrupt
from langgraph.stream._event_log import EventLog
from langgraph.stream._types import ProtocolEvent, StreamTransformer
from langgraph.stream.run_stream import BaseRunStream
from langgraph.stream.stream_channel import StreamChannel
if TYPE_CHECKING:
from collections.abc import Awaitable, Callable
from langgraph.stream._mux import StreamMux
logger = logging.getLogger(__name__)
SubgraphStatus = Literal["started", "completed", "failed", "interrupted"]
_TERMINAL_STATUSES: frozenset[SubgraphStatus] = frozenset(
{"completed", "failed", "interrupted"}
)
def _is_new_direct_child(
ns: tuple[str, ...],
scope: tuple[str, ...],
seen: set[tuple[str, ...]] | dict[tuple[str, ...], Any],
) -> bool:
"""Return True iff `ns` is a direct child of `scope` not yet seen.
Shared by `SubgraphTransformer` (in-process handle discovery) and
`LifecycleTransformer` (wire event emission) so the two can't
disagree on what counts as a new subgraph.
"""
return len(ns) == len(scope) + 1 and ns[:-1] == scope and ns not in seen
def _parse_ns_segment(segment: str) -> tuple[str, str | None]:
"""Split `node_name:task_id` into (node_name, task_id).
Task ids are present when Pregel spawned the subgraph as a task;
absent on synthesized namespaces (tests, hand-crafted events).
"""
node_name, sep, task_id = segment.partition(":")
if not sep:
return segment, None
return node_name, task_id or None
class ValuesTransformer(StreamTransformer):
"""Capture values events as a drainable stream of state snapshots.
Keeps `_latest` / `_interrupted` / `_interrupts` as scalar state
regardless of whether the log has a subscriber so `run.output()`
and `run.interrupted` work without forcing the caller to iterate
`run.values`. Log pushes are silent no-ops when unsubscribed.
Native transformer projection keys are exposed as direct
attributes on the run stream (e.g. `run.values`).
`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, scope: tuple[str, ...] = ()) -> None:
super().__init__(scope)
self._log: EventLog[dict[str, Any]] = EventLog()
self._latest: dict[str, Any] | None = None
self._interrupted = False
self._interrupts: list[Any] = []
def init(self) -> dict[str, Any]:
return {"values": self._log}
@property
def error(self) -> BaseException | None:
"""The error that ended the run, or `None` if it succeeded.
Set by the mux when it auto-fails the projection log.
"""
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"]
self._latest = params["data"]
interrupts = params.get("interrupts", ())
if interrupts:
self._interrupted = True
self._interrupts.extend(interrupts)
self._log.push(params["data"])
return True
class MessagesTransformer(StreamTransformer):
"""Capture messages events as ChatModelStream objects.
The messages projection yields one `ChatModelStream` (or
`AsyncChatModelStream`) per LLM call. Consumers iterate
`run.messages` to get stream handles, then use each handle's typed
projections (`.text`, `.reasoning`, `.tool_calls`, `.usage`,
`.output`) for per-message content.
Two input shapes are handled (via `params["data"] = (payload,
metadata)` from `StreamMessagesHandler`):
1. Protocol event (dict with `"event"` key) emitted by
`stream_v2()` / `astream_v2()` via the `on_stream_event`
callback. Routed to an existing `ChatModelStream` by
`metadata["run_id"]`. A `message-start` event creates a new
stream; `message-finish` closes it.
2. Whole `AIMessage` emitted from `on_chain_end` when a node
returns a finalized message. Replayed as a synthetic protocol
event lifecycle via `message_to_events`, then the
already-complete stream is pushed to the log.
V1 `AIMessageChunk` tuples (from `on_llm_new_token`) are not
streamed into this projection: chat models that want to populate
`run.messages` with content-block streaming must use
`stream_v2()` / `astream_v2()`. Models called via the legacy
`stream()` method still surface their final `AIMessage` via
`on_chain_end` when a node returns it as state.
`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.
"""
_native = True
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).
self._by_run: dict[str, ChatModelStream] = {}
self._pump_fn: Callable[[], bool] | None = None
self._apump_fn: Callable[[], Awaitable[bool]] | None = None
def init(self) -> dict[str, Any]:
return {"messages": self._log}
def _bind_pump(self, fn: Callable[[], bool]) -> None:
"""Wire the sync pull callback. Called by GraphRunStream._wire_request_more."""
self._pump_fn = fn
def _bind_apump(self, fn: Callable[[], Awaitable[bool]]) -> None:
"""Wire the async pull callback.
Called by `AsyncGraphRunStream._wire_arequest_more` so each
`AsyncChatModelStream` this transformer creates can drive the
shared graph pump from its projection cursors.
"""
self._apump_fn = fn
def _make_stream(
self,
*,
namespace: list[str],
node: str | None,
message_id: str | None,
) -> ChatModelStream:
"""Create a ChatModelStream (sync) or AsyncChatModelStream (async).
Wires whichever pump is bound. Prefers the async pump so nested
iteration under `AsyncGraphRunStream` drives the graph forward
without a background task. The unwired fallback (no pump bound)
is used by unit tests that dispatch events manually.
"""
if self._apump_fn is not None:
astream = AsyncChatModelStream(
namespace=namespace,
node=node,
message_id=message_id,
)
astream.set_arequest_more(self._apump_fn)
return astream
if self._pump_fn is not None:
stream: ChatModelStream = ChatModelStream(
namespace=namespace,
node=node,
message_id=message_id,
)
stream.set_request_more(self._pump_fn)
return stream
return AsyncChatModelStream(
namespace=namespace,
node=node,
message_id=message_id,
)
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"]
payload, metadata = params["data"]
node: str | None = metadata.get("langgraph_node")
run_id = str(metadata.get("run_id", "")) if metadata else ""
if isinstance(payload, dict) and "event" in payload:
self._route_protocol_event(
cast("MessagesData", payload), run_id=run_id, node=node
)
elif isinstance(payload, BaseMessage) and not isinstance(
payload, AIMessageChunk
):
self._route_whole_message(payload, node=node)
# Legacy AIMessageChunk tuples (from on_llm_new_token) are ignored;
# v1 streaming callers must switch to stream_v2() to populate this
# projection.
return True
def _route_protocol_event(
self,
event: MessagesData,
*,
run_id: str,
node: str | None,
) -> None:
event_type = event.get("event")
if event_type == "message-start":
message_id = event.get("message_id")
stream = self._make_stream(
namespace=list(self.scope),
node=node,
message_id=str(message_id) if message_id is not None else None,
)
self._by_run[run_id] = stream
self._log.push(stream)
stream.dispatch(event)
elif run_id in self._by_run:
stream = self._by_run[run_id]
stream.dispatch(event)
if event_type == "message-finish":
del self._by_run[run_id]
def _route_whole_message(self, message: BaseMessage, *, node: str | None) -> None:
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)
def finalize(self) -> None:
"""Clear any routing state — streams close themselves via `message-finish`."""
self._by_run.clear()
def fail(self, err: BaseException) -> None:
"""Propagate run error to any streams still open when the graph fails."""
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. 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.
Handle fields:
- `path`: the namespace tuple stable for the life of the handle.
- `graph_name` / `trigger_call_id`: parsed from the namespace
segment at discovery (`node_name:task_id`).
- `status`: `started` on discovery; advances to `completed` when
the parent mux closes, or `failed` / `interrupted` when it
errors.
- `error`: set on terminal error.
- `checkpoint`: unused by the current discovery path kept for
compatibility with consumers that inspect it.
`.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 dispatcher. At its own `scope` (inherited from
`StreamTransformer`, determined by the enclosing mux), it watches
for the first event at exactly one namespace level deeper to
discover a direct child. 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.
Discovery is method-agnostic: the first event of any mode whose
namespace places it directly below `scope` spawns the handle.
`graph_name` and `trigger_call_id` are parsed from the namespace
segment, which encodes `node_name:task_id`.
Terminal status for each handle is set by the parent mux's
`close` / `fail` path. `finalize` transitions still-open handles
to `completed`; `fail` transitions them to `failed` or
`interrupted` depending on the error.
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"])
depth = len(self.scope)
# 1. Discover: first-seen direct-child namespace registers a
# handle. Any event method triggers discovery — no dedicated
# channel.
if _is_new_direct_child(ns, self.scope, self._by_ns):
self._on_started(ns)
# 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)
return True
def _on_started(self, ns: tuple[str, ...]) -> None:
# `_on_register` is called by the mux during registration, which
# happens before any event can be dispatched — so this should
# always be set by the time we process an event.
assert self._mux is not None, (
"SubgraphTransformer processed an event before _on_register; "
"transformer registration ordering is broken."
)
graph_name, trigger_call_id = _parse_ns_segment(ns[-1])
child_mux = self._mux.make_child(ns)
handle = SubgraphRunStream(
path=ns,
mux=child_mux,
graph_name=graph_name,
trigger_call_id=trigger_call_id,
)
self._by_ns[ns] = handle
self._root_log.push(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:
logger.warning(
"Error closing subgraph mini-mux at %s; subscribers "
"may not see a clean close.",
handle.path,
exc_info=True,
)
def finalize(self) -> None:
"""Transition any still-open direct children to `completed`.
Subgraph interrupts surface as a values event with a populated
`interrupts` field rather than as an exception at the graph
boundary the parent pump exhausts normally and `finalize`
runs the close path. Inspect each child's `ValuesTransformer`
to distinguish "completed cleanly" from "interrupted".
"""
for handle in self._by_ns.values():
if handle.status not in _TERMINAL_STATUSES:
values_t = handle._mux.transformer_by_key("values")
if isinstance(values_t, ValuesTransformer) and values_t._interrupted:
handle.status = "interrupted"
else:
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:
logger.warning(
"Error failing subgraph mini-mux at %s; subscribers "
"may not see the terminal error.",
handle.path,
exc_info=True,
)
class LifecyclePayload(TypedDict, total=False):
"""Payload of a lifecycle event emitted by `LifecycleTransformer`."""
event: SubgraphStatus
namespace: list[str]
graph_name: str | None
trigger_call_id: str | None
error: str | None
class LifecycleTransformer(StreamTransformer):
"""Synthesize subgraph lifecycle events from observed namespaces.
Observes the same namespace signal `SubgraphTransformer` uses for
in-process discovery and emits `started` / `completed` / `failed`
/ `interrupted` payloads onto its `lifecycle` channel. Consumers
subscribed to that channel see the events in-process; wire
consumers receive them as protocol events with `method:
"lifecycle"` (unprefixed because this transformer is `_native`).
No `running` event: the ns-discovery signal only fires once a
subgraph has emitted output, so `started` already implies
execution. Consumers needing finer-grained task-start visibility
should read the `tasks` stream mode alongside.
No root `started`: the run object itself signals run start.
Terminal events are synthesized `finalize` emits `completed` for
still-open handles; `fail` emits `failed` or `interrupted`
depending on whether the error is a `GraphInterrupt`.
`scope_exact = False` so the transformer sees events at any
namespace (needed for discovery of direct children).
"""
_native = True
scope_exact = False
def __init__(self, scope: tuple[str, ...] = ()) -> None:
super().__init__(scope)
# retain=True: lifecycle events are low-volume and consumers
# commonly inspect them after draining `values`; without
# retention those pushes would be dropped.
self._channel: StreamChannel[LifecyclePayload] = StreamChannel(
"lifecycle", retain=True
)
self._seen: set[tuple[str, ...]] = set()
self._open: set[tuple[str, ...]] = set()
def init(self) -> dict[str, Any]:
return {"lifecycle": self._channel}
def process(self, event: ProtocolEvent) -> bool:
ns = tuple(event["params"]["namespace"])
if _is_new_direct_child(ns, self.scope, self._seen):
self._emit_started(ns)
return True
def _emit_started(self, ns: tuple[str, ...]) -> None:
graph_name, trigger_call_id = _parse_ns_segment(ns[-1])
self._seen.add(ns)
self._open.add(ns)
payload: LifecyclePayload = {
"event": "started",
"namespace": list(ns),
}
if graph_name:
payload["graph_name"] = graph_name
if trigger_call_id is not None:
payload["trigger_call_id"] = trigger_call_id
self._channel.push(payload)
def finalize(self) -> None:
"""Emit `completed` for every still-open direct child."""
for ns in list(self._open):
self._channel.push({"event": "completed", "namespace": list(ns)})
self._open.clear()
def fail(self, err: BaseException) -> None:
"""Emit `failed` / `interrupted` for every still-open direct child.
Closes the channel after emitting rather than letting the mux
auto-fail it the "failed" payload is the signal to
consumers, so they should be able to iterate it. A failed
channel would raise on iteration and hide the events that just
got pushed.
"""
is_interrupt = isinstance(err, GraphInterrupt)
event_type: SubgraphStatus = "interrupted" if is_interrupt else "failed"
error_str = None if is_interrupt else str(err)
for ns in list(self._open):
payload: LifecyclePayload = {
"event": event_type,
"namespace": list(ns),
}
if error_str is not None:
payload["error"] = error_str
self._channel.push(payload)
self._open.clear()
self._channel._close()
+7 -1
View File
@@ -116,7 +116,13 @@ def ensure_valid_checkpointer(checkpointer: Checkpointer) -> Checkpointer:
StreamMode = Literal[
"values", "updates", "checkpoints", "tasks", "debug", "messages", "custom"
"values",
"updates",
"checkpoints",
"tasks",
"debug",
"messages",
"custom",
]
"""How the stream method should emit outputs.
+1 -1
View File
@@ -24,7 +24,7 @@ classifiers = [
'Programming Language :: Python :: 3.13',
]
dependencies = [
"langchain-core==1.3.0a2",
"langchain-core==1.3.2",
"langgraph-checkpoint>=2.1.0,<5.0.0",
"langgraph-sdk>=0.3.0,<0.4.0",
"langgraph-prebuilt>=1.0.9,<1.1.0",
File diff suppressed because it is too large Load Diff
@@ -0,0 +1,258 @@
"""Tests for LifecycleTransformer — derives subgraph lifecycle from ns discovery."""
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._mux import StreamMux
from langgraph.stream._types import ProtocolEvent
from langgraph.stream.transformers import (
LifecycleTransformer,
ValuesTransformer,
)
TS = int(time.time() * 1000)
def _event(method: str, data: Any, *, namespace: list[str]) -> ProtocolEvent:
return {
"type": "event",
"method": method,
"params": {
"namespace": namespace,
"timestamp": TS,
"data": data,
},
}
def _drain_channel(mux: StreamMux) -> list[dict[str, Any]]:
ch = mux.extensions["lifecycle"]
return list(ch._log._items) # type: ignore[attr-defined]
def _drain_main_events(mux: StreamMux) -> list[ProtocolEvent]:
return list(mux._events._items) # type: ignore[attr-defined]
def _subscribe(mux: StreamMux) -> None:
ch = mux.extensions["lifecycle"]
ch._log._subscribed = True # type: ignore[attr-defined]
mux._events._subscribed = True # type: ignore[attr-defined]
class TestLifecycleTransformerUnit:
def _mux(self) -> StreamMux:
mux = StreamMux(
factories=[ValuesTransformer, LifecycleTransformer], is_async=False
)
_subscribe(mux)
return mux
def test_first_event_at_child_ns_emits_started(self) -> None:
mux = self._mux()
mux.push(_event("values", {"v": 1}, namespace=["child:task_a"]))
events = _drain_channel(mux)
assert events == [
{
"event": "started",
"namespace": ["child:task_a"],
"graph_name": "child",
"trigger_call_id": "task_a",
}
]
def test_no_started_at_root_ns(self) -> None:
mux = self._mux()
mux.push(_event("values", {"v": 1}, namespace=[]))
assert _drain_channel(mux) == []
def test_method_agnostic_discovery(self) -> None:
mux = self._mux()
mux.push(_event("messages", "x", namespace=["c:t"]))
(started,) = _drain_channel(mux)
assert started["event"] == "started"
assert started["namespace"] == ["c:t"]
def test_repeated_events_single_started(self) -> None:
mux = self._mux()
mux.push(_event("values", {"v": 1}, namespace=["c:t"]))
mux.push(_event("values", {"v": 2}, namespace=["c:t"]))
mux.push(_event("updates", {"n": "x"}, namespace=["c:t"]))
events = _drain_channel(mux)
assert len(events) == 1
assert events[0]["event"] == "started"
def test_ns_without_task_id(self) -> None:
mux = self._mux()
mux.push(_event("values", {"v": 1}, namespace=["child"]))
(started,) = _drain_channel(mux)
assert started["graph_name"] == "child"
assert "trigger_call_id" not in started
def test_finalize_emits_completed(self) -> None:
mux = self._mux()
mux.push(_event("values", {"v": 1}, namespace=["c:t"]))
mux.close()
events = _drain_channel(mux)
assert events == [
{
"event": "started",
"namespace": ["c:t"],
"graph_name": "c",
"trigger_call_id": "t",
},
{"event": "completed", "namespace": ["c:t"]},
]
def test_fail_with_graph_interrupt(self) -> None:
mux = self._mux()
mux.push(_event("values", {"v": 1}, namespace=["c:t"]))
mux.fail(GraphInterrupt())
events = _drain_channel(mux)
assert events[-1] == {"event": "interrupted", "namespace": ["c:t"]}
def test_fail_with_generic_error(self) -> None:
mux = self._mux()
mux.push(_event("values", {"v": 1}, namespace=["c:t"]))
mux.fail(RuntimeError("boom"))
events = _drain_channel(mux)
assert events[-1] == {
"event": "failed",
"namespace": ["c:t"],
"error": "boom",
}
class TestLifecycleWireFormat:
"""Native transformer: method on the wire is `"lifecycle"`, no `custom:` prefix."""
def test_wire_method_is_lifecycle_unprefixed(self) -> None:
mux = StreamMux(factories=[LifecycleTransformer], is_async=False)
_subscribe(mux)
mux.push(_event("values", {"v": 1}, namespace=["c:t"]))
# Find the lifecycle event in the main log.
lifecycle_events = [
ev for ev in _drain_main_events(mux) if ev["method"] == "lifecycle"
]
assert len(lifecycle_events) == 1
assert lifecycle_events[0]["params"]["data"]["event"] == "started"
# Also verify no `custom:lifecycle` leaks through.
assert not any(
ev["method"].startswith("custom:") for ev in _drain_main_events(mux)
)
def test_started_precedes_originating_event_on_wire(self) -> None:
"""Seq ordering: synthesized lifecycle event lands before the event that triggered it."""
mux = StreamMux(
factories=[ValuesTransformer, LifecycleTransformer], is_async=False
)
_subscribe(mux)
mux.push(_event("values", {"v": 1}, namespace=["c:t"]))
wire = _drain_main_events(mux)
methods = [ev["method"] for ev in wire]
# Lifecycle's synthetic event is forwarded during process() and
# gets an earlier seq than the originating values event.
assert methods.index("lifecycle") < methods.index("values")
# ---------------------------------------------------------------------------
# End-to-end via stream_v2
# ---------------------------------------------------------------------------
class SimpleState(TypedDict):
value: str
items: Annotated[list[str], operator.add]
def _build_nested_graph():
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()
outer_builder = StateGraph(SimpleState)
outer_builder.add_node("sub", inner)
outer_builder.add_edge(START, "sub")
outer_builder.add_edge("sub", END)
return outer_builder.compile()
class TestLifecycleEndToEnd:
def test_real_run_emits_started_and_completed(self) -> None:
graph = _build_nested_graph()
run = graph.stream_v2(
{"value": "", "items": []}, transformers=[LifecycleTransformer]
)
# Drain the run so finalize fires.
list(run.values)
lifecycle = list(run.lifecycle) # type: ignore[attr-defined]
events = [e["event"] for e in lifecycle]
assert "started" in events
assert "completed" in events
def test_real_run_error_emits_failed(self) -> None:
def boom(state: SimpleState) -> dict:
raise RuntimeError("kaboom")
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()
run = graph.stream_v2(
{"value": "", "items": []}, transformers=[LifecycleTransformer]
)
with pytest.raises(RuntimeError):
list(run.values)
lifecycle = list(run.lifecycle) # type: ignore[attr-defined]
events = [e["event"] for e in lifecycle]
assert "failed" in events
def test_lifecycle_and_subgraphs_agree(self) -> None:
"""SubgraphTransformer and LifecycleTransformer share the discovery predicate."""
graph = _build_nested_graph()
run = graph.stream_v2(
{"value": "", "items": []}, transformers=[LifecycleTransformer]
)
subs = list(run.subgraphs)
lifecycle = list(run.lifecycle) # type: ignore[attr-defined]
sub_paths = {tuple(s.path) for s in subs}
started_paths = {
tuple(e["namespace"]) for e in lifecycle if e["event"] == "started"
}
assert sub_paths == started_paths
@@ -0,0 +1,907 @@
"""Tests for the MessagesTransformer content-block upgrade (B2).
Verifies that `MessagesTransformer` routes protocol events (emitted by
`stream_v2` via `on_stream_event`) to `ChatModelStream` objects keyed by
run_id, and replays whole `AIMessage` payloads via `message_to_events`.
Legacy v1 `AIMessageChunk` tuples (from `on_llm_new_token`) are ignored.
"""
from __future__ import annotations
import time
from typing import Any
import pytest
from langchain_core.language_models import GenericFakeChatModel
from langchain_core.language_models.chat_model_stream import (
AsyncChatModelStream,
ChatModelStream,
)
from langchain_core.messages import AIMessage, AIMessageChunk
from langgraph.constants import END, START
from langgraph.graph import MessagesState, StateGraph
from langgraph.stream._event_log import EventLog
from langgraph.stream._mux import StreamMux
from langgraph.stream.run_stream import GraphRunStream
from langgraph.stream.transformers import MessagesTransformer, ValuesTransformer
TS = int(time.time() * 1000)
# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------
def _proto_event(
event: dict[str, Any],
*,
run_id: str = "run-1",
node: str = "llm",
) -> dict[str, Any]:
"""Build a messages ProtocolEvent carrying a protocol event dict (v2 path)."""
metadata: dict[str, Any] = {"langgraph_node": node, "run_id": run_id}
return {
"type": "event",
"method": "messages",
"params": {
"namespace": [],
"timestamp": TS,
"data": (event, metadata),
},
}
def _v1_chunk(
text: str,
msg_id: str = "msg-1",
*,
finish: bool = False,
node: str = "llm",
) -> dict[str, Any]:
"""Build a messages ProtocolEvent carrying a v1 AIMessageChunk tuple."""
rm: dict[str, Any] = {}
if finish:
rm["finish_reason"] = "stop"
message = AIMessageChunk(content=text, id=msg_id, response_metadata=rm)
metadata: dict[str, Any] = {"langgraph_node": node}
return {
"type": "event",
"method": "messages",
"params": {
"namespace": [],
"timestamp": TS,
"data": (message, metadata),
},
}
def _whole_msg(
text: str,
msg_id: str = "msg-10",
*,
node: str = "node",
) -> dict[str, Any]:
"""Build a messages ProtocolEvent carrying a completed AIMessage."""
message = AIMessage(content=text, id=msg_id)
metadata: dict[str, Any] = {"langgraph_node": node}
return {
"type": "event",
"method": "messages",
"params": {
"namespace": [],
"timestamp": TS,
"data": (message, metadata),
},
}
def _make_sync_transformer() -> tuple[MessagesTransformer, EventLog[ChatModelStream]]:
t = MessagesTransformer()
proj = t.init()
log: EventLog[ChatModelStream] = proj["messages"]
log._bind(is_async=False)
# Production subscribes via `iter(log)` from the graph consumer — do that
# up front so `push` during `process` isn't a no-op. Tests read buffered
# items via `log._items` directly rather than re-iterating.
log._subscribed = True
t._bind_pump(lambda: False)
return t, log
def _make_async_transformer() -> tuple[MessagesTransformer, EventLog[ChatModelStream]]:
t = MessagesTransformer()
proj = t.init()
log: EventLog[ChatModelStream] = proj["messages"]
log._bind(is_async=True)
log._subscribed = True
return t, log
# Standard lifecycle events for one streaming LLM call.
def _lifecycle(
*,
text: str = "hello world",
message_id: str = "run-1",
) -> list[dict[str, Any]]:
"""Produce a valid protocol event lifecycle: start, delta, finish, end."""
# Split text into two deltas to exercise delta accumulation.
half = len(text) // 2
first, second = text[:half], text[half:]
return [
{"event": "message-start", "role": "ai", "message_id": message_id},
{
"event": "content-block-start",
"index": 0,
"content_block": {"type": "text", "text": ""},
},
{
"event": "content-block-delta",
"index": 0,
"content_block": {"type": "text", "text": first},
},
{
"event": "content-block-delta",
"index": 0,
"content_block": {"type": "text", "text": second},
},
{
"event": "content-block-finish",
"index": 0,
"content_block": {"type": "text", "text": text},
},
{"event": "message-finish", "reason": "stop"},
]
# ---------------------------------------------------------------------------
# Primary path: protocol event routing
# ---------------------------------------------------------------------------
class TestProtocolEventRouting:
def test_message_start_creates_stream(self) -> None:
t, log = _make_sync_transformer()
t.process(
_proto_event(
{"event": "message-start", "role": "ai", "message_id": "run-1"},
run_id="run-1",
)
)
# Stream is in the log immediately.
log.close()
streams = list(log._items)
assert len(streams) == 1
assert isinstance(streams[0], ChatModelStream)
assert streams[0].message_id == "run-1"
def test_full_lifecycle_yields_done_stream(self) -> None:
t, log = _make_sync_transformer()
for evt in _lifecycle(text="hello world"):
t.process(_proto_event(evt, run_id="run-1"))
log.close()
(stream,) = list(log._items)
assert stream.done
assert stream.output.text == "hello world"
def test_message_finish_cleans_up_routing(self) -> None:
t, log = _make_sync_transformer()
for evt in _lifecycle():
t.process(_proto_event(evt, run_id="run-1"))
assert t._by_run == {}
def test_events_without_prior_start_are_ignored(self) -> None:
"""Orphan delta events (no preceding message-start) are dropped silently."""
t, log = _make_sync_transformer()
t.process(
_proto_event(
{
"event": "content-block-delta",
"index": 0,
"content_block": {"type": "text", "text": "orphan"},
},
run_id="unknown",
)
)
log.close()
assert list(log._items) == []
def test_concurrent_streams_routed_by_run_id(self) -> None:
"""Two interleaved LLM calls each produce their own stream."""
t, log = _make_sync_transformer()
# Interleave events from two different run_ids.
life_a = _lifecycle(text="aaaa", message_id="run-a")
life_b = _lifecycle(text="bbbb", message_id="run-b")
for a, b in zip(life_a, life_b):
t.process(_proto_event(a, run_id="run-a"))
t.process(_proto_event(b, run_id="run-b"))
log.close()
streams = list(log._items)
assert len(streams) == 2
by_id = {s.message_id: s for s in streams}
assert by_id["run-a"].output.text == "aaaa"
assert by_id["run-b"].output.text == "bbbb"
def test_text_deltas_accumulated_on_stream(self) -> None:
t, log = _make_sync_transformer()
for evt in _lifecycle(text="abcdef"):
t.process(_proto_event(evt))
log.close()
(stream,) = list(log._items)
deltas = list(stream._text_proj._deltas)
assert "".join(deltas) == "abcdef"
def test_stream_pushed_on_message_start_not_finish(self) -> None:
"""Consumer can see the stream before it finishes."""
t, log = _make_sync_transformer()
t.process(
_proto_event(
{"event": "message-start", "role": "ai", "message_id": "run-1"},
run_id="run-1",
)
)
# The log has the stream immediately — even though message-finish
# hasn't arrived yet.
assert len(log._items) == 1
def test_node_metadata_set_on_stream(self) -> None:
t, log = _make_sync_transformer()
t.process(
_proto_event(
{"event": "message-start", "role": "ai", "message_id": "run-1"},
run_id="run-1",
node="my_llm",
)
)
(stream,) = [*log._items]
assert stream.node == "my_llm"
# ---------------------------------------------------------------------------
# Non-streaming (whole AIMessage) fallback
# ---------------------------------------------------------------------------
class TestWholeMessageFallback:
def test_whole_ai_message_produces_complete_stream(self) -> None:
t, log = _make_sync_transformer()
t.process(_whole_msg("the full answer"))
log.close()
(stream,) = list(log._items)
assert stream.done
assert stream.output.text == "the full answer"
def test_whole_message_has_full_lifecycle(self) -> None:
t, log = _make_sync_transformer()
t.process(_whole_msg("full"))
log.close()
(stream,) = list(log._items)
event_types = [e["event"] for e in stream._events]
assert event_types == [
"message-start",
"content-block-start",
"content-block-delta",
"content-block-finish",
"message-finish",
]
# ---------------------------------------------------------------------------
# Legacy v1 chunks are ignored (users must migrate to stream_v2)
# ---------------------------------------------------------------------------
class TestLegacyChunksIgnored:
def test_aimessage_chunk_tuple_is_dropped(self) -> None:
t, log = _make_sync_transformer()
t.process(_v1_chunk("hello"))
t.process(_v1_chunk(" world", finish=True))
log.close()
assert list(log._items) == []
# ---------------------------------------------------------------------------
# Filtering behaviors
# ---------------------------------------------------------------------------
class TestFiltering:
def test_non_messages_events_pass_through(self) -> None:
t, _ = _make_sync_transformer()
values_event = {
"type": "event",
"method": "values",
"params": {"namespace": [], "timestamp": TS, "data": {"x": 1}},
}
assert t.process(values_event) is True
def test_subgraph_namespace_dropped(self) -> None:
"""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",
"params": {
"namespace": ["subgraph"],
"timestamp": TS,
"data": (
{"event": "message-start", "message_id": "run-x"},
{"run_id": "run-x"},
),
},
}
)
t._log.close()
assert list(t._log._items) == []
# ---------------------------------------------------------------------------
# Lifecycle: finalize / fail
# ---------------------------------------------------------------------------
class TestLifecycle:
def test_fail_propagates_to_open_streams(self) -> None:
t, log = _make_sync_transformer()
t.process(
_proto_event(
{"event": "message-start", "message_id": "run-1"},
run_id="run-1",
)
)
streams = list(log._items)
err = RuntimeError("graph died")
t.fail(err)
assert t._by_run == {}
assert streams[0]._error is err
def test_finalize_clears_routing_state(self) -> None:
t, _ = _make_sync_transformer()
t.process(
_proto_event(
{"event": "message-start", "message_id": "run-1"},
run_id="run-1",
)
)
assert "run-1" in t._by_run
t.finalize()
assert t._by_run == {}
# ---------------------------------------------------------------------------
# Async mode (AsyncChatModelStream)
# ---------------------------------------------------------------------------
class TestAsyncMode:
def test_async_mode_creates_async_stream(self) -> None:
t, log = _make_async_transformer()
for evt in _lifecycle(text="async stream"):
t.process(_proto_event(evt))
streams = list(log._items)
assert len(streams) == 1
assert isinstance(streams[0], AsyncChatModelStream)
@pytest.mark.anyio
async def test_async_text_projection_yields_deltas(self) -> None:
t, log = _make_async_transformer()
for evt in _lifecycle(text="hello world"):
t.process(_proto_event(evt))
(stream,) = list(log._items)
assert isinstance(stream, AsyncChatModelStream)
collected = []
async for delta in stream.text:
collected.append(delta)
assert "".join(collected) == "hello world"
@pytest.mark.anyio
async def test_async_output_awaitable(self) -> None:
t, log = _make_async_transformer()
for evt in _lifecycle(text="async"):
t.process(_proto_event(evt))
(stream,) = list(log._items)
msg = await stream.output
assert msg.text == "async"
# ---------------------------------------------------------------------------
# GraphRunStream integration
# ---------------------------------------------------------------------------
class TestWireRequestMore:
def test_bind_pump_called_on_wire(self) -> None:
values_t = ValuesTransformer()
messages_t = MessagesTransformer()
mux = StreamMux([values_t, messages_t], is_async=False)
assert messages_t._pump_fn is None
run = GraphRunStream(iter([]), mux)
# After wire, the transformer's pump callback is set.
assert messages_t._pump_fn is not None
# And calling it invokes GraphRunStream._pump_next (drains an empty
# graph_iter, returns False).
assert messages_t._pump_fn() is False
assert run._exhausted
def test_created_streams_have_request_more(self) -> None:
values_t = ValuesTransformer()
messages_t = MessagesTransformer()
mux = StreamMux([values_t, messages_t], is_async=False)
GraphRunStream(iter([]), mux)
log: EventLog[ChatModelStream] = mux.extensions["messages"]
log._subscribed = True
for evt in _lifecycle():
messages_t.process(_proto_event(evt))
(stream,) = list(log._items)
# Pump was threaded through: the stream's _request_more points at
# the same callable the transformer was bound with.
assert stream._request_more is messages_t._pump_fn
# ---------------------------------------------------------------------------
# End-to-end via StreamMux
# ---------------------------------------------------------------------------
class TestViaMux:
def test_streaming_via_mux(self) -> None:
t = MessagesTransformer()
v = ValuesTransformer()
mux = StreamMux([v, t], is_async=False)
t._bind_pump(lambda: False)
log: EventLog[ChatModelStream] = mux.extensions["messages"]
# Simulate a consumer subscribing (as `run.messages` iteration would).
log._subscribed = True
for evt in _lifecycle(text="mux stream"):
mux.push(_proto_event(evt))
mux.close()
(stream,) = list(log._items)
assert stream.output.text == "mux stream"
def test_whole_message_via_mux(self) -> None:
t = MessagesTransformer()
v = ValuesTransformer()
mux = StreamMux([v, t], is_async=False)
t._bind_pump(lambda: False)
log: EventLog[ChatModelStream] = mux.extensions["messages"]
log._subscribed = True
mux.push(_whole_msg("result"))
mux.close()
(stream,) = list(log._items)
assert stream.output.text == "result"
@pytest.mark.anyio
async def test_async_streaming_via_mux(self) -> None:
t = MessagesTransformer()
v = ValuesTransformer()
mux = StreamMux([v, t], is_async=True)
log: EventLog[ChatModelStream] = mux.extensions["messages"]
log._subscribed = True
for evt in _lifecycle(text="async mux"):
await mux.apush(_proto_event(evt))
streams = list(log._items)
assert len(streams) == 1
msg = await streams[0].output
assert msg.text == "async mux"
await mux.aclose()
# ---------------------------------------------------------------------------
# End-to-end: full graph → stream_v2 → run.messages
# ---------------------------------------------------------------------------
class TestEndToEnd:
"""Prove the full pipeline works when a node calls `model.stream_v2()`.
These tests exercise the path that the new messages projection is
designed for: a user node invokes `stream_v2` on a chat model,
`on_stream_event` fires on `StreamMessagesHandler`, the handler
forwards to the mux, and the transformer routes events into a
`ChatModelStream` exposed on `run.messages`.
Nothing in Pregel calls `stream_v2` automatically yet; the planned
`graph.stream_v2()` API (B4) and the `create_react_agent`
integration (C2) will wire that up. Until then, populating the
messages projection is opt-in at the node level.
"""
def test_node_calling_stream_v2_populates_messages(self) -> None:
model = GenericFakeChatModel(messages=iter(["hello world"]))
def call_model(state: MessagesState) -> dict[str, Any]:
stream = model.stream_v2(state["messages"])
return {"messages": stream.output}
graph = (
StateGraph(MessagesState)
.add_node("call_model", call_model)
.add_edge(START, "call_model")
.add_edge("call_model", END)
.compile()
)
run = graph.stream_v2({"messages": "hi"})
streams = list(run.messages)
assert len(streams) == 1
assert isinstance(streams[0], ChatModelStream)
assert streams[0].output.text == "hello world"
def test_node_stream_v2_text_deltas_iterate(self) -> None:
"""Consumer can iterate `.text` on the streamed message in real time."""
model = GenericFakeChatModel(messages=iter(["streamed answer"]))
def call_model(state: MessagesState) -> dict[str, Any]:
stream = model.stream_v2(state["messages"])
return {"messages": stream.output}
graph = (
StateGraph(MessagesState)
.add_node("call_model", call_model)
.add_edge(START, "call_model")
.add_edge("call_model", END)
.compile()
)
run = graph.stream_v2({"messages": "go"})
# Pull the stream handle out, then iterate its text deltas.
(stream,) = list(run.messages)
text = "".join(stream.text)
assert text == "streamed answer"
def test_non_llm_message_returned_from_node(self) -> None:
"""Node returns a finalized AIMessage directly — whole-message fallback."""
def return_message(state: MessagesState) -> dict[str, Any]:
return {"messages": AIMessage(content="hardcoded", id="msg-abc")}
graph = (
StateGraph(MessagesState)
.add_node("return_message", return_message)
.add_edge(START, "return_message")
.add_edge("return_message", END)
.compile()
)
run = graph.stream_v2({"messages": "hi"})
streams = list(run.messages)
assert len(streams) == 1
assert streams[0].output.text == "hardcoded"
@pytest.mark.anyio
async def test_async_node_calling_astream_v2(self) -> None:
model = GenericFakeChatModel(messages=iter(["async answer"]))
async def call_model(state: MessagesState) -> dict[str, Any]:
stream = await model.astream_v2(state["messages"])
msg = await stream
return {"messages": msg}
graph = (
StateGraph(MessagesState)
.add_node("call_model", call_model)
.add_edge(START, "call_model")
.add_edge("call_model", END)
.compile()
)
run = await graph.astream_v2({"messages": "hi"})
streams = []
async for stream in run.messages:
streams.append(stream)
assert len(streams) == 1
assert isinstance(streams[0], AsyncChatModelStream)
msg = await streams[0].output
assert msg.text == "async answer"
@pytest.mark.anyio
async def test_nested_async_iteration_yields_text_deltas(self) -> None:
"""Iterate `stream.text` inside `async for stream in run.messages`.
The inner `stream.text` cursor drives the shared graph pump via
`AsyncProjection._arequest_more`, wired by
`MessagesTransformer._bind_apump` and
`AsyncGraphRunStream._wire_arequest_more`.
"""
import asyncio
model = GenericFakeChatModel(messages=iter(["hello world"]))
async def call_model(state: MessagesState) -> dict[str, Any]:
stream = await model.astream_v2(state["messages"])
msg = await stream
return {"messages": msg}
graph = (
StateGraph(MessagesState)
.add_node("call_model", call_model)
.add_edge(START, "call_model")
.add_edge("call_model", END)
.compile()
)
run = await graph.astream_v2({"messages": "hi"})
async def consume_nested() -> list[str]:
collected: list[str] = []
async for stream in run.messages:
async for delta in stream.text:
collected.append(delta)
return collected
deltas = await asyncio.wait_for(consume_nested(), timeout=2.0)
assert "".join(deltas) == "hello world"
class TestEndToEndV2Invoke:
"""Nodes call `model.invoke()`; `stream_v2` routes through v2.
Exercises the auto-routing path added in
`feat(core): route invoke through v2 event path for
_V2StreamingCallbackHandler`: `stream_v2` injects
`CONFIG_KEY_STREAM_MESSAGES_V2` into the config, pregel attaches
`StreamMessagesHandlerV2`, `BaseChatModel._should_stream_v2` sees the
v2 marker and drives the protocol event generator, and
`on_stream_event` forwards each event onto the messages channel.
"""
def test_invoke_with_v2_marker_populates_messages(self) -> None:
"""Node calling `model.invoke()` produces one ChatModelStream with v2 events."""
model = GenericFakeChatModel(messages=iter(["hello world"]))
def call_model(state: MessagesState) -> dict[str, Any]:
return {"messages": model.invoke(state["messages"])}
graph = (
StateGraph(MessagesState)
.add_node("call_model", call_model)
.add_edge(START, "call_model")
.add_edge("call_model", END)
.compile()
)
run = graph.stream_v2({"messages": "hi"})
streams = list(run.messages)
assert len(streams) == 1, (
"Expected exactly one ChatModelStream — the streamed invoke and "
"the node's return of the same AIMessage must dedupe."
)
stream = streams[0]
assert isinstance(stream, ChatModelStream)
assert stream.output.text == "hello world"
def test_invoke_v2_emits_protocol_events(self) -> None:
"""Iterating the stream yields the full v2 lifecycle (not v1 chunks)."""
model = GenericFakeChatModel(messages=iter(["streamed answer"]))
def call_model(state: MessagesState) -> dict[str, Any]:
return {"messages": model.invoke(state["messages"])}
graph = (
StateGraph(MessagesState)
.add_node("call_model", call_model)
.add_edge(START, "call_model")
.add_edge("call_model", END)
.compile()
)
run = graph.stream_v2({"messages": "go"})
(stream,) = list(run.messages)
events = list(stream)
event_types = [e.get("event") for e in events]
assert "message-start" in event_types
assert "content-block-start" in event_types
assert "content-block-delta" in event_types
assert "content-block-finish" in event_types
assert "message-finish" in event_types
# Sanity: every event is a dict carrying an "event" key — not an
# AIMessageChunk tuple from the v1 path.
for event in events:
assert isinstance(event, dict)
assert "event" in event
# Typed projection still assembles the final text.
assert stream.output.text == "streamed answer"
def test_invoke_text_deltas_iterate_live(self) -> None:
"""`.text` projection yields deltas in order."""
model = GenericFakeChatModel(messages=iter(["delta streaming works"]))
def call_model(state: MessagesState) -> dict[str, Any]:
return {"messages": model.invoke(state["messages"])}
graph = (
StateGraph(MessagesState)
.add_node("call_model", call_model)
.add_edge(START, "call_model")
.add_edge("call_model", END)
.compile()
)
run = graph.stream_v2({"messages": "hi"})
(stream,) = list(run.messages)
assembled = "".join(stream.text)
assert assembled == "delta streaming works"
def test_invoke_dedupe_survives_multi_node_graph(self) -> None:
"""Two model-invoking nodes produce exactly two streams, each once."""
model_a = GenericFakeChatModel(messages=iter(["alpha"]))
model_b = GenericFakeChatModel(messages=iter(["beta"]))
def node_a(state: MessagesState) -> dict[str, Any]:
return {"messages": model_a.invoke(state["messages"])}
def node_b(state: MessagesState) -> dict[str, Any]:
return {"messages": model_b.invoke(state["messages"])}
graph = (
StateGraph(MessagesState)
.add_node("node_a", node_a)
.add_node("node_b", node_b)
.add_edge(START, "node_a")
.add_edge("node_a", "node_b")
.add_edge("node_b", END)
.compile()
)
run = graph.stream_v2({"messages": "hi"})
streams = list(run.messages)
assert len(streams) == 2
contents = {s.output.text for s in streams}
assert contents == {"alpha", "beta"}
def test_invoke_plus_constructed_message_two_streams(self) -> None:
"""A v2-streamed node + a node that returns a constructed AIMessage
produces two ChatModelStreams one from the live event lifecycle,
one synthesized from the constructed message via `message_to_events`.
"""
model = GenericFakeChatModel(messages=iter(["live stream"]))
def streaming_node(state: MessagesState) -> dict[str, Any]:
return {"messages": model.invoke(state["messages"])}
def constructed_node(state: MessagesState) -> dict[str, Any]:
return {"messages": [AIMessage(content="hardcoded", id="constructed-1")]}
graph = (
StateGraph(MessagesState)
.add_node("streaming_node", streaming_node)
.add_node("constructed_node", constructed_node)
.add_edge(START, "streaming_node")
.add_edge("streaming_node", "constructed_node")
.add_edge("constructed_node", END)
.compile()
)
run = graph.stream_v2({"messages": "hi"})
streams = list(run.messages)
assert len(streams) == 2
assert streams[0].node == "streaming_node"
assert streams[0].output.text == "live stream"
assert streams[1].node == "constructed_node"
assert streams[1].output.text == "hardcoded"
assert streams[1].message_id == "constructed-1"
@pytest.mark.anyio
async def test_ainvoke_with_v2_marker_populates_messages(self) -> None:
"""Async mirror: `model.ainvoke()` + `astream_v2`."""
model = GenericFakeChatModel(messages=iter(["async invoke"]))
async def call_model(state: MessagesState) -> dict[str, Any]:
return {"messages": await model.ainvoke(state["messages"])}
graph = (
StateGraph(MessagesState)
.add_node("call_model", call_model)
.add_edge(START, "call_model")
.add_edge("call_model", END)
.compile()
)
run = await graph.astream_v2({"messages": "hi"})
streams = []
async for stream in run.messages:
streams.append(stream)
assert len(streams) == 1
assert isinstance(streams[0], AsyncChatModelStream)
msg = await streams[0].output
assert msg.text == "async invoke"
class TestDirectMessagesModeStaysV1:
"""Regression guard: direct `graph.stream(stream_mode="messages")`
(no `stream_v2`) must keep the v1 `(AIMessageChunk, metadata)`
tuple shape. The v2 flag is only injected by `stream_v2` / `astream_v2`.
"""
def test_direct_graph_stream_messages_yields_ai_message_chunks(self) -> None:
model = GenericFakeChatModel(messages=iter(["legacy path"]))
def call_model(state: MessagesState) -> dict[str, Any]:
return {"messages": model.invoke(state["messages"])}
graph = (
StateGraph(MessagesState)
.add_node("call_model", call_model)
.add_edge(START, "call_model")
.add_edge("call_model", END)
.compile()
)
parts = list(graph.stream({"messages": "hi"}, stream_mode="messages"))
# Should have at least one streamed chunk; each part is
# (AIMessageChunk, metadata) — not a v2 event dict.
assert parts, "expected stream_mode='messages' to emit tuples"
for part in parts:
payload, _metadata = part
assert isinstance(payload, AIMessageChunk), (
"direct graph.stream(stream_mode='messages') leaked v2 "
"event dicts — stream_v2 flag bled through."
)
assembled = "".join(
p[0].content for p in parts if isinstance(p[0].content, str)
)
assert assembled == "legacy path"
class TestStreamMessagesHandlerV2Unit:
"""Unit tests on the handler class itself."""
def test_on_llm_new_token_is_noop(self) -> None:
"""v2 handler must not emit v1 chunks even if `on_llm_new_token` fires
(e.g. from a node calling `model.stream()` directly on a v2-flagged run).
"""
from uuid import uuid4
from langchain_core.outputs import ChatGenerationChunk
from langgraph.pregel._messages import StreamMessagesHandlerV2
emitted: list[Any] = []
handler = StreamMessagesHandlerV2(emitted.append, subgraphs=False)
run_id = uuid4()
# Register a fake run so `self.metadata.get(run_id)` would succeed for
# other callbacks — this makes sure the no-op is unconditional, not a
# side effect of missing metadata.
handler.metadata[run_id] = ((), {"langgraph_node": "x"})
handler.on_llm_new_token(
"hello",
chunk=ChatGenerationChunk(message=AIMessageChunk(content="hello")),
run_id=run_id,
)
assert emitted == [], (
"StreamMessagesHandlerV2.on_llm_new_token must not push to the "
"messages stream — it's the v2 marker's guarantee."
)
@@ -0,0 +1,411 @@
"""Tests for SubgraphTransformer namespace-based discovery."""
from __future__ import annotations
import operator
import time
from typing import Annotated, Any
import pytest
from langgraph.checkpoint.memory import InMemorySaver
from typing_extensions import TypedDict
from langgraph.constants import END, START
from langgraph.errors import GraphInterrupt
from langgraph.graph import StateGraph
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,
)
from langgraph.types import interrupt
TS = int(time.time() * 1000)
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:
(handle,) = list(transformer._root_log._items)
return handle
def test_root_event_does_not_create_handle(self) -> None:
mux, transformer = self._mux()
mux.push(_values({"value": 1}, namespace=[]))
assert list(transformer._root_log._items) == []
assert transformer._by_ns == {}
def test_first_event_at_child_depth_yields_handle(self) -> None:
mux, transformer = self._mux()
mux.push(_values({"value": 1}, namespace=["child:task_a"]))
handle = self._handle(transformer)
assert handle.path == ("child:task_a",)
assert handle.graph_name == "child"
assert handle.trigger_call_id == "task_a"
assert handle.status == "started"
def test_handle_without_task_id_suffix(self) -> None:
mux, transformer = self._mux()
mux.push(_values({"value": 1}, namespace=["child"]))
handle = self._handle(transformer)
assert handle.path == ("child",)
assert handle.graph_name == "child"
assert handle.trigger_call_id is None
def test_discovery_is_method_agnostic(self) -> None:
mux, transformer = self._mux()
# Method-agnostic means "any event method triggers discovery".
# Use `updates` — neither ValuesTransformer nor
# MessagesTransformer care about it, so the test only exercises
# SubgraphTransformer's discovery path.
mux.push(
{
"type": "event",
"method": "updates",
"params": {"namespace": ["c:t"], "timestamp": TS, "data": "x"},
}
)
handle = self._handle(transformer)
assert handle.path == ("c:t",)
def test_grandchild_surfaces_under_child(self) -> None:
mux, transformer = self._mux()
mux.push(_values({"value": 1}, namespace=["child:t"]))
child = self._handle(transformer)
_pre_subscribe_handle(child)
mux.push(_values({"value": 2}, namespace=["child:t", "grand:u"]))
(grand,) = _handle_subgraphs_items(child)
assert grand.path == ("child:t", "grand:u")
assert grand.graph_name == "grand"
assert grand.trigger_call_id == "u"
def test_values_routed_into_handle(self) -> None:
mux, transformer = self._mux()
mux.push(_values({"value": 1}, namespace=["c:t"]))
handle = self._handle(transformer)
_pre_subscribe_handle(handle)
mux.push(_values({"value": 2}, namespace=["c:t"]))
# Both the discovery event and subsequent values land in the child.
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(_values({"value": 1}, namespace=["c:t"]))
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) == [{"value": 1}]
def test_finalize_closes_dangling(self) -> None:
mux, transformer = self._mux()
mux.push(_values({"value": 1}, namespace=["c:t"]))
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(_values({"value": 1}, namespace=["c:t"]))
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(_values({"value": 1}, namespace=["c:t"]))
handle = self._handle(transformer)
mux.fail(RuntimeError("explode"))
assert handle.status == "failed"
assert handle.error == "explode"
def test_repeated_events_same_ns_single_handle(self) -> None:
mux, transformer = self._mux()
mux.push(_values({"value": 1}, namespace=["c:t"]))
mux.push(_values({"value": 2}, namespace=["c:t"]))
mux.push(_values({"value": 3}, namespace=["c:t"]))
handles = list(transformer._root_log._items)
assert len(handles) == 1
assert handles[0].path == ("c:t",)
# ---------------------------------------------------------------------------
# End-to-end tests via stream_v2 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()
run = graph.stream_v2({"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()
run = graph.stream_v2({"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()
run = graph.stream_v2({"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()
run = await graph.astream_v2({"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"
class TestSubgraphTriggerCallId:
"""Confirm `trigger_call_id` flows from real pregel metadata."""
def test_trigger_call_id_populated_end_to_end(self) -> None:
graph = _build_nested_graph()
run = graph.stream_v2({"value": "", "items": []})
collected: list[SubgraphRunStream] = list(run.subgraphs)
assert len(collected) == 1
child = collected[0]
# The child's single-segment path encodes `node_name:task_id`.
# Both the parsed task_id (`trigger_call_id`) and the segment
# should match the same task_id suffix.
assert ":" in child.path[0]
node_name, _, task_id = child.path[0].partition(":")
assert node_name == "sub"
assert task_id # non-empty
assert child.trigger_call_id == task_id
class TestSubgraphInterrupt:
"""Interrupts raised inside a subgraph surface as status=interrupted."""
def _build_interrupt_subgraph(self):
def inner_node(state: SimpleState) -> dict:
interrupt("need approval")
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()
outer_builder = StateGraph(SimpleState)
outer_builder.add_node("sub", inner)
outer_builder.add_edge(START, "sub")
outer_builder.add_edge("sub", END)
return outer_builder.compile(checkpointer=InMemorySaver())
def test_interrupt_in_subgraph_marks_handle_interrupted(self) -> None:
graph = self._build_interrupt_subgraph()
run = graph.stream_v2(
{"value": "", "items": []},
config={"configurable": {"thread_id": "t1"}},
)
collected: list[SubgraphRunStream] = list(run.subgraphs)
assert run.interrupted is True
assert len(collected) == 1
assert collected[0].status == "interrupted"
class TestSubgraphNameCollision:
"""The subgraph's compiled `name` equaling its node name is detected.
Primary detector `name != langgraph_node` fails here; the
parent_run_id fallback in `_is_nested_pregel_start` is what keeps
the subgraph visible.
"""
def test_name_equals_node_name_still_detected(self) -> None:
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)
# Compile with the same name as the node it will be registered as.
inner = inner_builder.compile(name="sub")
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()
run = graph.stream_v2({"value": "", "items": []})
collected: list[SubgraphRunStream] = list(run.subgraphs)
assert len(collected) == 1
child = collected[0]
assert child.graph_name == "sub"
assert child.status == "completed"
+17 -4
View File
@@ -1348,10 +1348,11 @@ wheels = [
[[package]]
name = "langchain-core"
version = "1.3.0a2"
version = "1.3.2"
source = { registry = "https://pypi.org/simple" }
dependencies = [
{ name = "jsonpatch" },
{ name = "langchain-protocol" },
{ name = "langsmith" },
{ name = "packaging" },
{ name = "pydantic" },
@@ -1360,9 +1361,21 @@ dependencies = [
{ name = "typing-extensions" },
{ name = "uuid-utils" },
]
sdist = { url = "https://files.pythonhosted.org/packages/af/bc/0bff31fcaff174d86031cc713471a3e85ed4ec8e5cd95ad0217f2aced20e/langchain_core-1.3.0a2.tar.gz", hash = "sha256:52d978c84552b74b9a3f16c1fced84f9e27cc96d7a67c601925ce6cbc4ea3cf9", size = 854580, upload-time = "2026-04-13T14:37:55.745Z" }
sdist = { url = "https://files.pythonhosted.org/packages/a8/03/7219502e8ca728d65eb44d7a3eb60239230742a70dbfc9241b9bfd61c4ab/langchain_core-1.3.2.tar.gz", hash = "sha256:fd7a50b2f28ba561fd9d7f5d2760bc9e06cf00cdf820a3ccafe88a94ffa8d5b7", size = 911813, upload-time = "2026-04-24T15:49:23.699Z" }
wheels = [
{ url = "https://files.pythonhosted.org/packages/0e/14/03c09686602567059f26af29de0c44546a83af2f2aa29925e61040e43ea2/langchain_core-1.3.0a2-py3-none-any.whl", hash = "sha256:9e929a34f0b0c6c1255e395a1de34f8626893ceb4cdae550a22a0bd18c87be54", size = 510233, upload-time = "2026-04-13T14:37:54.277Z" },
{ url = "https://files.pythonhosted.org/packages/7d/d5/8fa4431007cbb7cfed7590f4d6a5dea3ad724f4174d248f6642ef5ce7d05/langchain_core-1.3.2-py3-none-any.whl", hash = "sha256:d44a66127f9f8db735bdfd0ab9661bccb47a97113cfd3f2d89c74864422b7274", size = 542390, upload-time = "2026-04-24T15:49:21.991Z" },
]
[[package]]
name = "langchain-protocol"
version = "0.0.11"
source = { registry = "https://pypi.org/simple" }
dependencies = [
{ name = "typing-extensions" },
]
sdist = { url = "https://files.pythonhosted.org/packages/d0/bb/38b5eaefa41c67735eedd9f9a2568b11c9eb376fa129a5edd7cc3dcde071/langchain_protocol-0.0.11.tar.gz", hash = "sha256:c276e2373b5ac691fc7ac9a72019d55182444ce8e89385c3f7e9f0185d0aace7", size = 6622, upload-time = "2026-04-23T22:13:16.771Z" }
wheels = [
{ url = "https://files.pythonhosted.org/packages/0b/fa/6a8ecad8472b182f2caf9d83fd89f40fc1590cb96546d90089b7869b7f5e/langchain_protocol-0.0.11-py3-none-any.whl", hash = "sha256:364da1faf6f5d3001413bede792c1a822c0f23ae55d1ce1266ca7d8e80e79011", size = 6778, upload-time = "2026-04-23T22:13:15.712Z" },
]
[[package]]
@@ -1439,7 +1452,7 @@ test = [
[package.metadata]
requires-dist = [
{ name = "langchain-core", specifier = "==1.3.0a2" },
{ name = "langchain-core", specifier = "==1.3.2" },
{ name = "langgraph-checkpoint", editable = "../checkpoint" },
{ name = "langgraph-prebuilt", editable = "../prebuilt" },
{ name = "langgraph-sdk", editable = "../sdk-py" },
+17 -4
View File
@@ -249,10 +249,11 @@ wheels = [
[[package]]
name = "langchain-core"
version = "1.3.0a2"
version = "1.3.2"
source = { registry = "https://pypi.org/simple" }
dependencies = [
{ name = "jsonpatch" },
{ name = "langchain-protocol" },
{ name = "langsmith" },
{ name = "packaging" },
{ name = "pydantic" },
@@ -261,9 +262,21 @@ dependencies = [
{ name = "typing-extensions" },
{ name = "uuid-utils" },
]
sdist = { url = "https://files.pythonhosted.org/packages/af/bc/0bff31fcaff174d86031cc713471a3e85ed4ec8e5cd95ad0217f2aced20e/langchain_core-1.3.0a2.tar.gz", hash = "sha256:52d978c84552b74b9a3f16c1fced84f9e27cc96d7a67c601925ce6cbc4ea3cf9", size = 854580, upload-time = "2026-04-13T14:37:55.745Z" }
sdist = { url = "https://files.pythonhosted.org/packages/a8/03/7219502e8ca728d65eb44d7a3eb60239230742a70dbfc9241b9bfd61c4ab/langchain_core-1.3.2.tar.gz", hash = "sha256:fd7a50b2f28ba561fd9d7f5d2760bc9e06cf00cdf820a3ccafe88a94ffa8d5b7", size = 911813, upload-time = "2026-04-24T15:49:23.699Z" }
wheels = [
{ url = "https://files.pythonhosted.org/packages/0e/14/03c09686602567059f26af29de0c44546a83af2f2aa29925e61040e43ea2/langchain_core-1.3.0a2-py3-none-any.whl", hash = "sha256:9e929a34f0b0c6c1255e395a1de34f8626893ceb4cdae550a22a0bd18c87be54", size = 510233, upload-time = "2026-04-13T14:37:54.277Z" },
{ url = "https://files.pythonhosted.org/packages/7d/d5/8fa4431007cbb7cfed7590f4d6a5dea3ad724f4174d248f6642ef5ce7d05/langchain_core-1.3.2-py3-none-any.whl", hash = "sha256:d44a66127f9f8db735bdfd0ab9661bccb47a97113cfd3f2d89c74864422b7274", size = 542390, upload-time = "2026-04-24T15:49:21.991Z" },
]
[[package]]
name = "langchain-protocol"
version = "0.0.11"
source = { registry = "https://pypi.org/simple" }
dependencies = [
{ name = "typing-extensions" },
]
sdist = { url = "https://files.pythonhosted.org/packages/d0/bb/38b5eaefa41c67735eedd9f9a2568b11c9eb376fa129a5edd7cc3dcde071/langchain_protocol-0.0.11.tar.gz", hash = "sha256:c276e2373b5ac691fc7ac9a72019d55182444ce8e89385c3f7e9f0185d0aace7", size = 6622, upload-time = "2026-04-23T22:13:16.771Z" }
wheels = [
{ url = "https://files.pythonhosted.org/packages/0b/fa/6a8ecad8472b182f2caf9d83fd89f40fc1590cb96546d90089b7869b7f5e/langchain_protocol-0.0.11-py3-none-any.whl", hash = "sha256:364da1faf6f5d3001413bede792c1a822c0f23ae55d1ce1266ca7d8e80e79011", size = 6778, upload-time = "2026-04-23T22:13:15.712Z" },
]
[[package]]
@@ -281,7 +294,7 @@ dependencies = [
[package.metadata]
requires-dist = [
{ name = "langchain-core", specifier = "==1.3.0a2" },
{ name = "langchain-core", specifier = "==1.3.2" },
{ name = "langgraph-checkpoint", editable = "../checkpoint" },
{ name = "langgraph-prebuilt", editable = "." },
{ name = "langgraph-sdk", editable = "../sdk-py" },
+17 -4
View File
@@ -262,10 +262,11 @@ wheels = [
[[package]]
name = "langchain-core"
version = "1.3.0a2"
version = "1.3.2"
source = { registry = "https://pypi.org/simple" }
dependencies = [
{ name = "jsonpatch" },
{ name = "langchain-protocol" },
{ name = "langsmith" },
{ name = "packaging" },
{ name = "pydantic" },
@@ -274,9 +275,21 @@ dependencies = [
{ name = "typing-extensions" },
{ name = "uuid-utils" },
]
sdist = { url = "https://files.pythonhosted.org/packages/af/bc/0bff31fcaff174d86031cc713471a3e85ed4ec8e5cd95ad0217f2aced20e/langchain_core-1.3.0a2.tar.gz", hash = "sha256:52d978c84552b74b9a3f16c1fced84f9e27cc96d7a67c601925ce6cbc4ea3cf9", size = 854580, upload-time = "2026-04-13T14:37:55.745Z" }
sdist = { url = "https://files.pythonhosted.org/packages/a8/03/7219502e8ca728d65eb44d7a3eb60239230742a70dbfc9241b9bfd61c4ab/langchain_core-1.3.2.tar.gz", hash = "sha256:fd7a50b2f28ba561fd9d7f5d2760bc9e06cf00cdf820a3ccafe88a94ffa8d5b7", size = 911813, upload-time = "2026-04-24T15:49:23.699Z" }
wheels = [
{ url = "https://files.pythonhosted.org/packages/0e/14/03c09686602567059f26af29de0c44546a83af2f2aa29925e61040e43ea2/langchain_core-1.3.0a2-py3-none-any.whl", hash = "sha256:9e929a34f0b0c6c1255e395a1de34f8626893ceb4cdae550a22a0bd18c87be54", size = 510233, upload-time = "2026-04-13T14:37:54.277Z" },
{ url = "https://files.pythonhosted.org/packages/7d/d5/8fa4431007cbb7cfed7590f4d6a5dea3ad724f4174d248f6642ef5ce7d05/langchain_core-1.3.2-py3-none-any.whl", hash = "sha256:d44a66127f9f8db735bdfd0ab9661bccb47a97113cfd3f2d89c74864422b7274", size = 542390, upload-time = "2026-04-24T15:49:21.991Z" },
]
[[package]]
name = "langchain-protocol"
version = "0.0.11"
source = { registry = "https://pypi.org/simple" }
dependencies = [
{ name = "typing-extensions" },
]
sdist = { url = "https://files.pythonhosted.org/packages/d0/bb/38b5eaefa41c67735eedd9f9a2568b11c9eb376fa129a5edd7cc3dcde071/langchain_protocol-0.0.11.tar.gz", hash = "sha256:c276e2373b5ac691fc7ac9a72019d55182444ce8e89385c3f7e9f0185d0aace7", size = 6622, upload-time = "2026-04-23T22:13:16.771Z" }
wheels = [
{ url = "https://files.pythonhosted.org/packages/0b/fa/6a8ecad8472b182f2caf9d83fd89f40fc1590cb96546d90089b7869b7f5e/langchain_protocol-0.0.11-py3-none-any.whl", hash = "sha256:364da1faf6f5d3001413bede792c1a822c0f23ae55d1ce1266ca7d8e80e79011", size = 6778, upload-time = "2026-04-23T22:13:15.712Z" },
]
[[package]]
@@ -294,7 +307,7 @@ dependencies = [
[package.metadata]
requires-dist = [
{ name = "langchain-core", specifier = "==1.3.0a2" },
{ name = "langchain-core", specifier = "==1.3.2" },
{ name = "langgraph-checkpoint", editable = "../checkpoint" },
{ name = "langgraph-prebuilt", editable = "../prebuilt" },
{ name = "langgraph-sdk", editable = "." },