From f4388df77f4095a4db3560a2922e12139e7233e8 Mon Sep 17 00:00:00 2001 From: Nick Hollon Date: Tue, 28 Apr 2026 16:29:21 -0400 Subject: [PATCH] feat(langgraph): add streaming transformer infrastructure and tests (#7519) --- .../langgraph/_internal/_constants.py | 4 + libs/langgraph/langgraph/graph/state.py | 9 + libs/langgraph/langgraph/pregel/_messages.py | 128 ++ libs/langgraph/langgraph/pregel/_tools.py | 268 +++ libs/langgraph/langgraph/pregel/main.py | 394 ++++- libs/langgraph/langgraph/stream/__init__.py | 37 + libs/langgraph/langgraph/stream/_convert.py | 32 + libs/langgraph/langgraph/stream/_event_log.py | 306 ++++ libs/langgraph/langgraph/stream/_mux.py | 510 ++++++ libs/langgraph/langgraph/stream/_types.py | 312 ++++ libs/langgraph/langgraph/stream/run_stream.py | 521 ++++++ .../langgraph/stream/stream_channel.py | 113 ++ .../langgraph/stream/transformers.py | 748 ++++++++ libs/langgraph/pyproject.toml | 2 +- libs/langgraph/tests/test_pregel_stream_v2.py | 1512 +++++++++++++++++ .../test_stream_lifecycle_transformer.py | 401 +++++ .../tests/test_stream_messages_transformer.py | 872 ++++++++++ .../tests/test_stream_subgraph_transformer.py | 865 ++++++++++ libs/langgraph/tests/test_stream_v2_e2e.py | 792 +++++++++ .../tests/test_tool_stream_handler.py | 290 ++++ libs/langgraph/uv.lock | 21 +- libs/prebuilt/langgraph/prebuilt/__init__.py | 2 + .../langgraph/prebuilt/_tool_call_stream.py | 117 ++ .../prebuilt/_tool_call_transformer.py | 128 ++ libs/prebuilt/langgraph/prebuilt/tool_node.py | 21 + .../tests/test_tool_call_transformer.py | 307 ++++ libs/prebuilt/uv.lock | 21 +- libs/sdk-py/uv.lock | 21 +- 28 files changed, 8696 insertions(+), 58 deletions(-) create mode 100644 libs/langgraph/langgraph/pregel/_tools.py create mode 100644 libs/langgraph/langgraph/stream/__init__.py create mode 100644 libs/langgraph/langgraph/stream/_convert.py create mode 100644 libs/langgraph/langgraph/stream/_event_log.py create mode 100644 libs/langgraph/langgraph/stream/_mux.py create mode 100644 libs/langgraph/langgraph/stream/_types.py create mode 100644 libs/langgraph/langgraph/stream/run_stream.py create mode 100644 libs/langgraph/langgraph/stream/stream_channel.py create mode 100644 libs/langgraph/langgraph/stream/transformers.py create mode 100644 libs/langgraph/tests/test_pregel_stream_v2.py create mode 100644 libs/langgraph/tests/test_stream_lifecycle_transformer.py create mode 100644 libs/langgraph/tests/test_stream_messages_transformer.py create mode 100644 libs/langgraph/tests/test_stream_subgraph_transformer.py create mode 100644 libs/langgraph/tests/test_stream_v2_e2e.py create mode 100644 libs/langgraph/tests/test_tool_stream_handler.py create mode 100644 libs/prebuilt/langgraph/prebuilt/_tool_call_stream.py create mode 100644 libs/prebuilt/langgraph/prebuilt/_tool_call_transformer.py create mode 100644 libs/prebuilt/tests/test_tool_call_transformer.py diff --git a/libs/langgraph/langgraph/_internal/_constants.py b/libs/langgraph/langgraph/_internal/_constants.py index 68cb48fe8..d28289053 100644 --- a/libs/langgraph/langgraph/_internal/_constants.py +++ b/libs/langgraph/langgraph/_internal/_constants.py @@ -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, diff --git a/libs/langgraph/langgraph/graph/state.py b/libs/langgraph/langgraph/graph/state.py index b1c24de2b..f545d4535 100644 --- a/libs/langgraph/langgraph/graph/state.py +++ b/libs/langgraph/langgraph/graph/state.py @@ -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[[tuple[str, ...]], Any]] | None = None, ) -> CompiledStateGraph[StateT, ContextT, InputT, OutputT]: """Compiles the `StateGraph` into a `CompiledStateGraph` object. @@ -1077,6 +1078,13 @@ 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 `StreamTransformer` classes or + configured factories. Classes and factories are instantiated + per run whenever `stream_v2` / `astream_v2` is called and are + propagated to subgraph scopes. Custom factories should follow + the standard `StreamTransformer` constructor shape by + accepting `scope` as their first argument. Appended after the + built-in stream transformers. Returns: CompiledStateGraph: The compiled `StateGraph`. @@ -1159,6 +1167,7 @@ class StateGraph(Generic[StateT, ContextT, InputT, OutputT]): store=store, cache=cache, name=name or "LangGraph", + stream_transformers=transformers, ) compiled._serde_allowlist = serde_allowlist diff --git a/libs/langgraph/langgraph/pregel/_messages.py b/libs/langgraph/langgraph/pregel/_messages.py index acab06098..aa1db1ac9 100644 --- a/libs/langgraph/langgraph/pregel/_messages.py +++ b/libs/langgraph/langgraph/pregel/_messages.py @@ -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]] @@ -248,3 +253,126 @@ 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 __init__( + self, + stream: Callable[[StreamChunk], None], + subgraphs: bool, + *, + parent_ns: tuple[str, ...] | None = None, + ) -> None: + super().__init__(stream, subgraphs, parent_ns=parent_ns) + self._streamed_run_ids: set[UUID] = set() + + def on_llm_end( + self, + response: LLMResult, + *, + run_id: UUID, + parent_run_id: UUID | None = None, + **kwargs: Any, + ) -> Any: + if meta := self.metadata.get(run_id): + if response.generations and response.generations[0]: + gen = response.generations[0][0] + if isinstance(gen, ChatGeneration): + if run_id in self._streamed_run_ids: + if gen.message.id is None: + gen.message.id = str(uuid4()) + self.seen.add(gen.message.id) + else: + self._emit(meta, gen.message, dedupe=True) + self._streamed_run_ids.discard(run_id) + self.metadata.pop(run_id, None) + + def on_llm_error( + self, + error: BaseException, + *, + run_id: UUID, + parent_run_id: UUID | None = None, + **kwargs: Any, + ) -> Any: + self._streamed_run_ids.discard(run_id) + super().on_llm_error( + error, + run_id=run_id, + parent_run_id=parent_run_id, + **kwargs, + ) + + 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": + self._streamed_run_ids.add(run_id) + 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))) diff --git a/libs/langgraph/langgraph/pregel/_tools.py b/libs/langgraph/langgraph/pregel/_tools.py new file mode 100644 index 000000000..fd7bb63bc --- /dev/null +++ b/libs/langgraph/langgraph/pregel/_tools.py @@ -0,0 +1,268 @@ +from __future__ import annotations + +from collections.abc import AsyncIterator, Callable, Iterator +from contextvars import ContextVar, Token +from typing import Any, TypeVar, cast +from uuid import UUID + +from langchain_core.callbacks import BaseCallbackHandler + +from langgraph._internal._constants import NS_SEP +from langgraph.constants import TAG_NOSTREAM +from langgraph.pregel.protocol import StreamChunk + +try: + from langchain_core.tracers._streaming import _StreamingCallbackHandler +except ImportError: + _StreamingCallbackHandler = object # type: ignore[assignment,misc] + + +T = TypeVar("T") + +ToolCallWriter = Callable[[Any], None] +"""A closure bound to a single tool call that emits `tool-output-delta` events.""" + +_tool_call_writer: ContextVar[ToolCallWriter | None] = ContextVar( + "langgraph_tool_call_writer", default=None +) +"""ContextVar holding the writer for the currently-executing tool call. + +Set by `StreamToolCallHandler.on_tool_start` and reset on end/error. +Read by `ToolRuntime.emit_output_delta` (in `langgraph.prebuilt`). +""" + + +class StreamToolCallHandler(BaseCallbackHandler, _StreamingCallbackHandler): + """Callback handler that emits tool-call lifecycle events on the stream. + + Fires on LangChain's `on_tool_*` callbacks and pushes to the `tools` + stream mode. Emits `tool-started` / `tool-output-delta` / + `tool-finished` / `tool-error` payloads keyed by `tool_call_id`. + + While a tool is executing, this handler sets `_tool_call_writer` to a + closure bound to that call's namespace and `tool_call_id`. + `ToolRuntime.emit_output_delta` reads that ContextVar so tool bodies + can stream partial output without threading the writer through their + own signature. + + Attached by `Pregel.stream` / `astream` when `"tools"` is in + `stream_modes`. `run_inline = True` keeps event ordering + deterministic. + """ + + run_inline = True + + def __init__( + self, + stream: Callable[[StreamChunk], None], + subgraphs: bool, + *, + parent_ns: tuple[str, ...] | None = None, + ) -> None: + """Configure the handler to stream tool-call events. + + Args: + stream: Callable that accepts a `StreamChunk` tuple + `(namespace, mode, payload)` and enqueues it. + subgraphs: Whether to emit events from tools called inside + nested subgraphs. When False, only tools at the + handler's own scope (`parent_ns`) emit. + parent_ns: Namespace where the handler was attached. + Mirrors the `StreamMessagesHandler` escape hatch: + tools whose containing namespace equals `parent_ns` + still emit even with `subgraphs=False`, so a node that + explicitly streams a subgraph with `stream_mode="tools"` + sees its own tools. + """ + self.stream = stream + self.subgraphs = subgraphs + self.parent_ns = parent_ns + # run_id → (namespace, tool_call_id, ContextVar token) + # `on_tool_end` does not receive `tool_call_id` in kwargs, so + # we correlate by `run_id` which is present on every callback. + self._run_to_call: dict[ + UUID, tuple[tuple[str, ...], str, Token[ToolCallWriter | None]] + ] = {} + + def _ns_for_emit( + self, + metadata: dict[str, Any] | None, + tags: list[str] | None, + ) -> tuple[str, ...] | None: + """Resolve the namespace this tool call should emit at, or `None` to skip. + + Mirrors `StreamMessagesHandler.on_chat_model_start`'s namespace + derivation: parses `langgraph_checkpoint_ns` (which ends with + the `node_name:task_id` of the calling node), drops that + trailing segment, and returns the containing subgraph's own + namespace. Returns `None` when the call should be silently + suppressed: + + - `metadata` is missing — handler is attached to a context + without Pregel routing info. + - `TAG_NOSTREAM` is in `tags` — caller explicitly opted out. + - Tool runs in a subgraph (`len(ns) > 0`) and the handler was + attached with `subgraphs=False` and a different `parent_ns` + than the call's containing subgraph. + """ + if not metadata: + return None + if tags and TAG_NOSTREAM in tags: + return None + nskey = metadata.get("langgraph_checkpoint_ns") + if not nskey: + ns: tuple[str, ...] = () + else: + ns = tuple(cast(str, nskey).split(NS_SEP))[:-1] + if not self.subgraphs and len(ns) > 0 and ns != self.parent_ns: + return None + return ns + + def _start( + self, + serialized: dict[str, Any] | None, + input_str: str, + *, + run_id: UUID, + metadata: dict[str, Any] | None, + tags: list[str] | None, + inputs: dict[str, Any] | None, + kwargs: dict[str, Any], + ) -> None: + ns = self._ns_for_emit(metadata, tags) + if ns is None: + return + tool_call_id = cast("str | None", kwargs.get("tool_call_id")) or str(run_id) + tool_name = ( + (serialized or {}).get("name") + or cast("str | None", kwargs.get("name")) + or "" + ) + + def writer(delta: Any) -> None: + self.stream( + ( + ns, + "tools", + { + "event": "tool-output-delta", + "tool_call_id": tool_call_id, + "delta": delta, + }, + ) + ) + + token = _tool_call_writer.set(writer) + self._run_to_call[run_id] = (ns, tool_call_id, token) + + payload: dict[str, Any] = { + "event": "tool-started", + "tool_call_id": tool_call_id, + "tool_name": tool_name, + } + if inputs is not None: + payload["input"] = inputs + self.stream((ns, "tools", payload)) + + def _end(self, output: Any, *, run_id: UUID) -> None: + info = self._run_to_call.pop(run_id, None) + if info is None: + return + ns, tool_call_id, token = info + self._reset_writer(token) + self.stream( + ( + ns, + "tools", + { + "event": "tool-finished", + "tool_call_id": tool_call_id, + "output": output, + }, + ) + ) + + def _error(self, error: BaseException, *, run_id: UUID) -> None: + info = self._run_to_call.pop(run_id, None) + if info is None: + return + ns, tool_call_id, token = info + self._reset_writer(token) + self.stream( + ( + ns, + "tools", + { + "event": "tool-error", + "tool_call_id": tool_call_id, + "message": str(error), + }, + ) + ) + + def tap_output_aiter( + self, run_id: UUID, output: AsyncIterator[T] + ) -> AsyncIterator[T]: + """Pass-through — required by the `_StreamingCallbackHandler` protocol.""" + return output + + def tap_output_iter(self, run_id: UUID, output: Iterator[T]) -> Iterator[T]: + """Pass-through — sync counterpart to `tap_output_aiter`.""" + return output + + @staticmethod + def _reset_writer(token: Token[ToolCallWriter | None]) -> None: + # Token is invalid if `on_tool_end` runs in a different context + # than `on_tool_start` (e.g. langchain may hand off to a thread + # worker without copying the context). Swallow that case; the + # ContextVar lifetime is bounded by the enclosing task anyway. + try: + _tool_call_writer.reset(token) + except ValueError: + pass + + # ------------------------------------------------------------------ + # Sync callbacks + # ------------------------------------------------------------------ + + def on_tool_start( + self, + serialized: dict[str, Any], + input_str: str, + *, + run_id: UUID, + parent_run_id: UUID | None = None, + tags: list[str] | None = None, + metadata: dict[str, Any] | None = None, + inputs: dict[str, Any] | None = None, + **kwargs: Any, + ) -> Any: + self._start( + serialized, + input_str, + run_id=run_id, + metadata=metadata, + tags=tags, + inputs=inputs, + kwargs=kwargs, + ) + + def on_tool_end( + self, + output: Any, + *, + run_id: UUID, + parent_run_id: UUID | None = None, + **kwargs: Any, + ) -> Any: + self._end(output, run_id=run_id) + + def on_tool_error( + self, + error: BaseException, + *, + run_id: UUID, + parent_run_id: UUID | None = None, + **kwargs: Any, + ) -> Any: + self._error(error, run_id=run_id) diff --git a/libs/langgraph/langgraph/pregel/main.py b/libs/langgraph/langgraph/pregel/main.py index c440e77f9..33a135d25 100644 --- a/libs/langgraph/langgraph/pregel/main.py +++ b/libs/langgraph/langgraph/pregel/main.py @@ -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,10 +134,14 @@ 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 +from langgraph.pregel._tools import StreamToolCallHandler from langgraph.pregel._utils import get_new_channel_versions from langgraph.pregel._validate import validate_graph, validate_keys from langgraph.pregel._write import ChannelWrite, ChannelWriteEntry @@ -148,6 +153,15 @@ from langgraph.runtime import ( Runtime, ServerInfo, ) +from langgraph.stream._mux import StreamMux +from langgraph.stream._types import StreamTransformer +from langgraph.stream.run_stream import AsyncGraphRunStream, GraphRunStream +from langgraph.stream.transformers import ( + LifecycleTransformer, + MessagesTransformer, + SubgraphTransformer, + ValuesTransformer, +) from langgraph.types import ( All, CachePolicy, @@ -340,6 +354,58 @@ class NodeBuilder: ) +def _collect_stream_modes(mux: Any) -> list[StreamMode]: + """Return the union of `required_stream_modes` across registered transformers. + + Transformers declare the stream modes they need to function, and + `stream_v2` asks the graph for exactly that union — no hardcoded + default set. If zero transformers declare a given mode, the graph + does not stream events for it. + """ + modes: set[StreamMode] = set() + for transformer in mux._transformers: + modes.update( + cast( + "tuple[StreamMode, ...]", + getattr(transformer, "required_stream_modes", ()), + ) + ) + return list(modes) + + +def _normalize_stream_transformer_factories( + specs: Sequence[Callable[[tuple[str, ...]], Any]] | None, +) -> list[Callable[[tuple[str, ...]], Any]]: + """Normalize stream transformer specs to scoped factories. + + A stream transformer spec is a callable that accepts + `scope: tuple[str, ...]` and returns a fresh `StreamTransformer`. + Transformer classes work when their constructor follows the same + shape. Pre-built instances are rejected because they cannot be + cloned into subgraph scopes. + """ + factories: list[Callable[[tuple[str, ...]], Any]] = [] + for spec in specs or (): + if isinstance(spec, StreamTransformer): + raise TypeError( + "stream_v2 transformers must be scope-aware callables, " + f"got pre-built instance {type(spec).__name__}. Pass the " + "transformer class or a factory like " + "`lambda scope: MyTransformer(scope, ...)`." + ) + if not callable(spec): + raise TypeError( + "stream_v2 transformers must be scope-aware callables, " + f"got {type(spec).__name__}." + ) + + def factory(scope: tuple[str, ...], _spec: Callable[..., Any] = spec) -> Any: + return _spec(scope) + + factories.append(factory) + return factories + + class Pregel( PregelProtocol[StateT, ContextT, InputT, OutputT], Generic[StateT, ContextT, InputT, OutputT], @@ -671,6 +737,7 @@ class Pregel( config: RunnableConfig | None = None, trigger_to_nodes: Mapping[str, Sequence[str]] | None = None, name: str = "LangGraph", + stream_transformers: Sequence[Callable[[tuple[str, ...]], Any]] | None = None, **deprecated_kwargs: Unpack[DeprecatedKwargs], ) -> None: if ( @@ -717,6 +784,9 @@ class Pregel( self.config = config self.trigger_to_nodes = trigger_to_nodes or {} self.name = name + self.stream_transformers: tuple[Callable[[tuple[str, ...]], Any], ...] = tuple( + stream_transformers or () + ) self._serde_allowlist: set[tuple[str, ...]] | None = None if auto_validate: self.validate() @@ -2582,19 +2652,7 @@ class Pregel( stream = SyncQueue() config = ensure_config(self.config, config) - callback_manager = get_callback_manager_for_config(config) - if "ls_integration" not in callback_manager.metadata: - callback_manager.add_metadata({"ls_integration": "langgraph"}) - run_manager = callback_manager.on_chain_start( - None, - input, - name=config.get("run_name", self.get_name()), - run_id=config.get("run_id"), - ) - graph_callback_manager = get_sync_graph_callback_manager_for_config( - config, - run_id=run_manager.run_id, - ) + run_manager = None try: # assign defaults ( @@ -2615,6 +2673,36 @@ class Pregel( interrupt_after=interrupt_after, durability=durability, ) + callback_manager = get_callback_manager_for_config(config) + if "messages" in stream_modes and version != "v2": + # Strip any inherited v2 messages handler so a v1 stream + # does not get routed through the content-block event + # protocol. Leave v1 handlers in place — an outer + # stream(stream_mode="messages", subgraphs=True) relies + # on its inheritable handler to observe events emitted + # by inner stream(stream_mode="messages") calls. + callback_manager.handlers = [ + h + for h in callback_manager.handlers + if not isinstance(h, StreamMessagesHandlerV2) + ] + callback_manager.inheritable_handlers = [ + h + for h in callback_manager.inheritable_handlers + if not isinstance(h, StreamMessagesHandlerV2) + ] + if "ls_integration" not in callback_manager.metadata: + callback_manager.add_metadata({"ls_integration": "langgraph"}) + run_manager = callback_manager.on_chain_start( + None, + input, + name=config.get("run_name", self.get_name()), + run_id=config.get("run_id"), + ) + graph_callback_manager = get_sync_graph_callback_manager_for_config( + config, + run_id=run_manager.run_id, + ) if checkpointer is None and durability is not None: warnings.warn( "`durability` has no effect when no checkpointer is present.", @@ -2626,14 +2714,33 @@ class Pregel( # set up messages stream mode if "messages" in stream_modes: ns_ = cast(str | None, config[CONF].get(CONFIG_KEY_CHECKPOINT_NS)) + use_stream_messages_v2 = bool( + version == "v2" and config[CONF].get(CONFIG_KEY_STREAM_MESSAGES_V2) + ) + messages_handler_cls = ( + StreamMessagesHandlerV2 + if use_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, ) ) + # set up tools stream mode + if "tools" in stream_modes: + ns_tools = cast(str | None, config[CONF].get(CONFIG_KEY_CHECKPOINT_NS)) + run_manager.inheritable_handlers.append( + StreamToolCallHandler( + stream.put, + subgraphs, + parent_ns=tuple(ns_tools.split(NS_SEP)) if ns_tools else None, + ) + ) + # set up custom stream mode if "custom" in stream_modes: @@ -2804,7 +2911,8 @@ class Pregel( # set final channel values as run output run_manager.on_chain_end(loop.output) except BaseException as e: - run_manager.on_chain_error(e) + if run_manager is not None: + run_manager.on_chain_error(e) raise @overload @@ -2944,33 +3052,7 @@ class Pregel( ) config = ensure_config(self.config, config) - callback_manager = get_async_callback_manager_for_config(config) - if "ls_integration" not in callback_manager.metadata: - callback_manager.add_metadata({"ls_integration": "langgraph"}) - run_manager = await callback_manager.on_chain_start( - None, - input, - name=config.get("run_name", self.get_name()), - run_id=config.get("run_id"), - ) - graph_callback_manager = get_async_graph_callback_manager_for_config( - config, - run_id=run_manager.run_id, - ) - # if running from astream_log() run each proc with streaming - do_stream = ( - next( - ( - True - for h in run_manager.handlers - if isinstance(h, _StreamingCallbackHandler) - and not isinstance(h, StreamMessagesHandler) - ), - False, - ) - if _StreamingCallbackHandler is not None - else False - ) + run_manager = None try: # assign defaults ( @@ -2991,6 +3073,50 @@ class Pregel( interrupt_after=interrupt_after, durability=durability, ) + callback_manager = get_async_callback_manager_for_config(config) + if "messages" in stream_modes and version != "v2": + # Strip any inherited v2 messages handler so a v1 stream + # does not get routed through the content-block event + # protocol. Leave v1 handlers in place — an outer + # astream(stream_mode="messages", subgraphs=True) relies + # on its inheritable handler to observe events emitted + # by inner astream(stream_mode="messages") calls. + callback_manager.handlers = [ + h + for h in callback_manager.handlers + if not isinstance(h, StreamMessagesHandlerV2) + ] + callback_manager.inheritable_handlers = [ + h + for h in callback_manager.inheritable_handlers + if not isinstance(h, StreamMessagesHandlerV2) + ] + if "ls_integration" not in callback_manager.metadata: + callback_manager.add_metadata({"ls_integration": "langgraph"}) + run_manager = await callback_manager.on_chain_start( + None, + input, + name=config.get("run_name", self.get_name()), + run_id=config.get("run_id"), + ) + graph_callback_manager = get_async_graph_callback_manager_for_config( + config, + run_id=run_manager.run_id, + ) + # if running from astream_log() run each proc with streaming + do_stream = ( + next( + ( + True + for h in run_manager.handlers + if isinstance(h, _StreamingCallbackHandler) + and not isinstance(h, StreamMessagesHandler) + ), + False, + ) + if _StreamingCallbackHandler is not None + else False + ) if checkpointer is None and durability is not None: warnings.warn( "`durability` has no effect when no checkpointer is present.", @@ -3003,14 +3129,33 @@ 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)) + use_stream_messages_v2 = bool( + version == "v2" and config[CONF].get(CONFIG_KEY_STREAM_MESSAGES_V2) + ) + messages_handler_cls = ( + StreamMessagesHandlerV2 + if use_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, ) ) + # set up tools stream mode + if "tools" in stream_modes: + ns_tools = cast(str | None, config[CONF].get(CONFIG_KEY_CHECKPOINT_NS)) + run_manager.inheritable_handlers.append( + StreamToolCallHandler( + stream_put, + subgraphs, + parent_ns=tuple(ns_tools.split(NS_SEP)) if ns_tools else None, + ) + ) + # set up custom stream mode def stream_writer(c: Any) -> None: aioloop.call_soon_threadsafe( @@ -3234,9 +3379,150 @@ class Pregel( # set final channel values as run output await run_manager.on_chain_end(loop.output) except BaseException as e: - await asyncio.shield(run_manager.on_chain_error(e)) + if run_manager is not None: + 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[Callable[[tuple[str, ...]], Any]] | None = None, + ) -> Any: + """Start a sync v2 streaming run driven by transformer projections. + + Builds a `StreamMux` from the built-in `ValuesTransformer` / + `MessagesTransformer`, 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. + + Note: + Nesting v1 `stream(stream_mode="messages")` inside a node + of a `stream_v2` run is not fully supported. The outer v2 + messages handler is inheritable, so it sits in the inner + chat model's callback chain; `BaseChatModel.invoke` then + routes through the v2 event protocol and the inner v1 + messages handler does not see `on_llm_new_token` chunks. + The inner stream still yields a finalized message via + `on_llm_end`, but token-by-token output is lost. Use + `stream_v2` for the inner graph as well, or call + `chat_model.stream(...)` explicitly inside the node, to + get token-level streaming. + + 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 classes or configured factories + appended after compile-time `stream_transformers`. Factories + are called as `factory(scope)` so they can propagate to + subgraph scopes. + + Returns: + A `GraphRunStream` the caller iterates to drive the run. + """ + parent_ns = _resolve_parent_ns(self.config, config) + compiled_factories = _normalize_stream_transformer_factories( + self.stream_transformers + ) + extra_factories = _normalize_stream_transformer_factories(transformers) + mux = StreamMux( + factories=[ + ValuesTransformer, + MessagesTransformer, + LifecycleTransformer, + SubgraphTransformer, + *compiled_factories, + *extra_factories, + ], + scope=parent_ns, + is_async=False, + ) + values_t = cast(ValuesTransformer, mux.transformer_by_key("values")) + graph_iter = iter( + self.stream( + input, + patch_configurable(config, {CONFIG_KEY_STREAM_MESSAGES_V2: True}), + stream_mode=_collect_stream_modes(mux), + subgraphs=True, + version="v2", + interrupt_before=interrupt_before, + interrupt_after=interrupt_after, + ) + ) + return GraphRunStream(graph_iter, mux, values_t) + + 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[Callable[[tuple[str, ...]], 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. + + Note: + Same nesting limitation as `stream_v2`: nesting v1 + `astream(stream_mode="messages")` inside a node of an + `astream_v2` run drops `on_llm_new_token` chunks because + the outer v2 handler reroutes `BaseChatModel.invoke` + through the v2 event protocol. The inner stream still + yields a finalized message at end-of-call. Use + `astream_v2` for the inner graph as well, or call + `chat_model.astream(...)` explicitly inside the node, to + get token-level streaming. + + 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 classes or configured factories + appended after compile-time `stream_transformers`. Factories + are called as `factory(scope)` so they can propagate to + subgraph scopes. + """ + parent_ns = _resolve_parent_ns(self.config, config) + compiled_factories = _normalize_stream_transformer_factories( + self.stream_transformers + ) + extra_factories = _normalize_stream_transformer_factories(transformers) + mux = StreamMux( + factories=[ + ValuesTransformer, + MessagesTransformer, + LifecycleTransformer, + SubgraphTransformer, + *compiled_factories, + *extra_factories, + ], + scope=parent_ns, + is_async=True, + ) + values_t = cast(ValuesTransformer, mux.transformer_by_key("values")) + graph_aiter = self.astream( + input, + patch_configurable(config, {CONFIG_KEY_STREAM_MESSAGES_V2: True}), + stream_mode=_collect_stream_modes(mux), + subgraphs=True, + version="v2", + interrupt_before=interrupt_before, + interrupt_after=interrupt_after, + ).__aiter__() + return AsyncGraphRunStream(graph_aiter, mux, values_t) + @overload def invoke( self, @@ -3712,6 +3998,24 @@ def _coerce_checkpoint_values(payload: Any, mapper: Callable[[Any], Any]) -> Non payload["values"] = mapper(payload["values"]) +def _resolve_parent_ns( + graph_config: RunnableConfig | None, call_config: RunnableConfig | None +) -> tuple[str, ...]: + """Return the checkpoint namespace the caller is running under. + + `stream_v2` uses this to scope its native projections + (`ValuesTransformer`, `MessagesTransformer`) to events emitted at + the run's own level. A root call resolves to `()`; a call made + from inside a node carries the outer graph's task namespace so the + projection still matches its own root-level events. + """ + merged = ensure_config(graph_config, call_config) + ns = merged.get(CONF, {}).get(CONFIG_KEY_CHECKPOINT_NS) + if not ns: + return () + return tuple(ns.split(NS_SEP)) + + def _build_server_info( config: RunnableConfig, parent_runtime: Runtime[Any] ) -> ServerInfo | None: diff --git a/libs/langgraph/langgraph/stream/__init__.py b/libs/langgraph/langgraph/stream/__init__.py new file mode 100644 index 000000000..6e1ae7161 --- /dev/null +++ b/libs/langgraph/langgraph/stream/__init__.py @@ -0,0 +1,37 @@ +"""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, + AsyncSubgraphRunStream, + GraphRunStream, + SubgraphRunStream, +) +from langgraph.stream.stream_channel import StreamChannel +from langgraph.stream.transformers import ( + LifecyclePayload, + LifecycleTransformer, + SubgraphStatus, + SubgraphTransformer, +) + +__all__ = [ + "AsyncGraphRunStream", + "AsyncSubgraphRunStream", + "EventLog", + "GraphRunStream", + "LifecyclePayload", + "LifecycleTransformer", + "ProtocolEvent", + "StreamChannel", + "StreamTransformer", + "SubgraphRunStream", + "SubgraphStatus", + "SubgraphTransformer", +] diff --git a/libs/langgraph/langgraph/stream/_convert.py b/libs/langgraph/langgraph/stream/_convert.py new file mode 100644 index 000000000..c1a14b5b0 --- /dev/null +++ b/libs/langgraph/langgraph/stream/_convert.py @@ -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, + } diff --git a/libs/langgraph/langgraph/stream/_event_log.py b/libs/langgraph/langgraph/stream/_event_log.py new file mode 100644 index 000000000..05bc10504 --- /dev/null +++ b/libs/langgraph/langgraph/stream/_event_log.py @@ -0,0 +1,306 @@ +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) -> 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. + + 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. + self._subscribed = False + + # 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. + + Raises: + RuntimeError: If the log is closed (and subscribed). + """ + if not self._subscribed: + 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)) diff --git a/libs/langgraph/langgraph/stream/_mux.py b/libs/langgraph/langgraph/stream/_mux.py new file mode 100644 index 000000000..34578628d --- /dev/null +++ b/libs/langgraph/langgraph/stream/_mux.py @@ -0,0 +1,510 @@ +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` with the mux's scope (typically `()` for +the root). Standard transformer classes 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, ...] = (), + _assign_seq: bool = True, + ) -> None: + """Initialize the mux and register transformers in order. + + Callers pass either `transformers` (pre-built instances) or + `factories` (callables producing fresh instances per mux). Each + transformer's `init()` is called, 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. Registered + only on this mux — they are NOT cloned into child + mini-muxes built by `_make_child`. Use `factories` for + transformers that should propagate to nested scopes. + is_async: True for async dispatch (`apush` / `aclose` / + `afail`), False for the sync path. + factories: One-argument callables `(scope) -> StreamTransformer`. + Called once with this mux's `scope` here, and cloned + again per child scope by `_make_child` so each + sub-mux gets fresh instances. + scope: The namespace the mux operates within. The root mux + is `()`. + _assign_seq: Internal flag for child muxes. Root muxes assign + monotonic `seq` numbers when appending to their main event + log; child muxes share forwarded event objects and must not + mutate their envelopes. + + 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. + """ + self.is_async = is_async + self.scope: tuple[str, ...] = scope + self._assign_seq = _assign_seq + 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] = {} + + # Stored only when constructed from factories — used by + # `_make_child` to clone the transformer pipeline at a deeper + # scope. Pre-built transformers can't be cloned, so a mux + # built with `transformers=` rejects child construction. + self._factories: list[TransformerFactory] | None = ( + list(factories) if factories is not None else None + ) + self._pump_fn: Callable[[], bool] | None = None + self._apump_fn: Callable[[], Awaitable[bool]] | None = None + + # Factories run first (they propagate to child mini-muxes + # via `_make_child`), then any pre-built `transformers=` + # instances are registered as root-only — they aren't cloned + # for child scopes. + if factories is not None: + for factory in factories: + self._register(factory(scope)) + for transformer in transformers or (): + self._register(transformer) + + def transformer_by_key(self, key: str) -> StreamTransformer | None: + """Return the transformer that contributed `key` to the projection.""" + return self._transformer_by_key.get(key) + + # ------------------------------------------------------------------ + # Pump wiring + mini-mux nesting + # ------------------------------------------------------------------ + + def bind_pump(self, fn: Callable[[], bool]) -> None: + """Wire the sync pull callback onto every projection in this mux. + + Records the pump on the mux so child mini-muxes built by + `_make_child` can inherit it. Propagates to: + - the main event log (`self._events`) + - every projection EventLog and StreamChannel in `extensions` + - any registered transformer that exposes `_bind_pump` (e.g. + `MessagesTransformer` so `ChatModelStream` instances drive the + shared pump from their 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 _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 bindings (so cursors on its + projection logs drive the root pump), carries the same factory + list forward to any grandchild subgraphs, and does not assign + `seq` numbers so forwarded events can be shared without + mutating their envelope. + + Raises: + RuntimeError: If the mux was not constructed with + `factories=`. Mini-muxes require factories so each scope + gets its own fresh transformer instances. + """ + if self._factories is None: + 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, + _assign_seq=False, + ) + 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 _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}" + ) + is_native = bool(getattr(transformer, "_native", False)) + self._transformers.append(transformer) + self._bind_and_wire(projection, 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()) + transformer._on_register(self) + + 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. + If any transformer returns False, the event is suppressed from + the main log, but transformers that already saw it keep their + side effects. + + On the root mux, `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 root log. Child muxes do + not assign `seq`, so subgraph forwarding can share event objects + without mutating their envelopes. + + Args: + event: The protocol event to dispatch. + """ + keep = True + for transformer in self._transformers: + if not transformer.process(event): + keep = False + if keep: + if self._assign_seq: + 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. 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. The root mux assigns `seq`; child muxes do + not, so forwarded subgraph events can be shared without copying. + 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. + """ + keep = True + for transformer in self._transformers: + if not await transformer.aprocess(event): + keep = False + if keep: + if self._assign_seq: + 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], *, native: bool = False + ) -> None: + """Bind and wire EventLog / StreamChannel instances in a projection. + + Args: + projection: The projection dict returned by a transformer's + `init()`. + native: True when the owning transformer is `_native`. + Channels owned by a native transformer use the channel + name directly as the protocol method; user-defined + channels are prefixed with `custom:`. + """ + for value in projection.values(): + if isinstance(value, StreamChannel): + value._bind(is_async=self.is_async) + self._channels.append(value) + method = value.name if native else f"custom:{value.name}" + + def _make_forward(method_name: str) -> Callable[[Any], None]: + def _forward(item: Any) -> None: + self._forward(method_name, item) + + return _forward + + value._wire(_make_forward(method)) + elif isinstance(value, EventLog): + value._bind(is_async=self.is_async) + self._logs.append(value) + + def _forward(self, method: str, item: Any) -> None: + """Inject a ProtocolEvent for a StreamChannel push. + + Forwarded events bypass the transformer pipeline to avoid + infinite recursion (a transformer that pushes to a channel + during `process()` would re-trigger itself). These events are + visible in this mux's main event log but are not passed through + transformers' `process()` methods. Only the root mux assigns + `seq` to forwarded channel events. + + Args: + method: The full protocol method (already with or without + the `custom:` prefix; resolved by `_bind_and_wire`). + item: The payload pushed onto the channel. + """ + event: ProtocolEvent = { + "type": "event", + "method": method, + "params": { + "namespace": [], + "timestamp": int(time.time() * 1000), + "data": item, + }, + } + if self._assign_seq: + self._seq += 1 + event["seq"] = self._seq + self._events.push(event) diff --git a/libs/langgraph/langgraph/stream/_types.py b/libs/langgraph/langgraph/stream/_types.py new file mode 100644 index 000000000..530dd7c87 --- /dev/null +++ b/libs/langgraph/langgraph/stream/_types.py @@ -0,0 +1,312 @@ +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 root StreamMux. + Consumers that need a total order across root 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. Set at construction from the mux's scope (each + factory is called as `factory(scope)`). + 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. + supports_sync: Set True only for transformers that override + async-lane hooks while still fully supporting the sync lane. + Such transformers may be registered under `stream()`. + required_stream_modes: Stream modes the graph must emit for + this transformer to have anything to process. Computed as + the union across all registered transformers to determine + which modes a `stream_v2` run requests from the graph. + Empty tuple means the transformer consumes only synthetic + events (or is purely passive). + """ + + requires_async: ClassVar[bool] = False + supports_sync: ClassVar[bool] = False + required_stream_modes: ClassVar[tuple[str, ...]] = () + + 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. 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 _on_register(self, mux: Any) -> None: + """Called by `StreamMux._register` after this transformer is wired in. + + Default is a no-op. Override to capture a reference to the + owning mux — needed for transformers that build mini-muxes + via `mux._make_child(...)` (e.g. `SubgraphTransformer`). + """ + + 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`) without also declaring that it + supports the sync lane. + + Args: + transformer: The transformer to inspect. + + Returns: + True if the transformer cannot run under sync `stream()`. + """ + if transformer.requires_async: + return True + if transformer.supports_sync: + return False + cls = type(transformer) + for name in ("aprocess", "afinalize", "afail"): + if getattr(cls, name) is not getattr(StreamTransformer, name): + return True + return False diff --git a/libs/langgraph/langgraph/stream/run_stream.py b/libs/langgraph/langgraph/stream/run_stream.py new file mode 100644 index 000000000..e712fedee --- /dev/null +++ b/libs/langgraph/langgraph/stream/run_stream.py @@ -0,0 +1,521 @@ +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 SubgraphStatus, 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 GraphRunStream: + """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. + + All transformer projections live in `extensions`. Native transformer + projections (those with `_native = True`) are also set as direct + attributes on this instance (e.g. `run.values`, `run.messages`). + """ + + def __init__( + self, + graph_iter: Iterator[Any] | None, + mux: StreamMux, + values_transformer: ValuesTransformer, + *, + wire_pump: bool = True, + ) -> None: + """Initialize the run stream. + + Args: + graph_iter: Pull-based iterator over the graph's stream, + or `None` for nested run streams whose pump is driven + by an outer run (e.g. `SubgraphRunStream`). + mux: The StreamMux owning projections and the main log. + values_transformer: The built-in values transformer + providing `output` / `interrupted` / `interrupts`. + wire_pump: When True (default), bind `_pump_next` as the + mux's pump callable. Subclasses that inherit a parent + pump via `StreamMux._make_child` should pass False to + preserve the parent binding. + """ + self._graph_iter = graph_iter + self._mux = mux + self.extensions: Mapping[str, Any] = MappingProxyType(mux.extensions) + self._values_transformer = values_transformer + self._exhausted = False + for key in mux.native_keys: + setattr(self, key, mux.extensions[key]) + if wire_pump: + self._wire_request_more(mux) + + def _wire_request_more(self, mux: StreamMux) -> None: + """Wire the sync pull callback through the mux. + + Routing through `mux.bind_pump` (rather than walking + projections directly here) lets child mini-muxes built by + `mux._make_child(...)` inherit the same pump callable, so + cursors on a subgraph handle's projections drive the root + pump just like cursors on `run.values` do. + """ + 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. Always False when constructed with + `graph_iter=None` (the run is driven by an outer pump). + """ + if self._exhausted or self._graph_iter is None: + return False + try: + part = next(self._graph_iter) + self._mux.push(convert_to_protocol_event(part)) + return True + except StopIteration: + self._mux.close() + self._exhausted = True + return False + except Exception as e: + self._mux.fail(e) + self._exhausted = True + return False + + 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) + if (err := self._values_transformer.error) 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) + if (err := self._values_transformer.error) 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) + if (err := self._values_transformer.error) is not None: + raise err + return self._values_transformer._interrupts + + def __iter__(self) -> Iterator[ProtocolEvent]: + """Subscribe to the main event log and iterate protocol events.""" + return iter(self._mux._events) + + def interleave(self, *names: str) -> Iterator[tuple[str, Any]]: + """Iterate multiple projections round-robin, yielding ``(name, item)``. + + Each turn advances one projection's cursor; when a cursor's buffer + is empty, pulling from it drives the pump once, which fans out to + every subscribed projection log. Projections whose items aren't + consumed on this turn sit in their own buffers only until the next + turn reaches them, bounding memory by the skew between projection + rates rather than letting any single log grow to the full run + length. + + Projections are exhausted independently; a projection that finishes + early drops out of the rotation while others continue. The overall + iterator ends once all named projections are done. + + Args: + *names: Projection keys to interleave. Must match keys in + ``extensions``. + + Yields: + ``(name, item)`` tuples in round-robin order across the named + projections. + + Raises: + KeyError: If a name doesn't match a registered projection. + + Example: + ```python + for name, item in run.interleave("messages", "values"): + if name == "messages": + print("msg:", item) + else: + print("val:", item) + ``` + """ + cursors: dict[str, Iterator[Any]] = { + name: iter(self.extensions[name]) for name in names + } + done: set[str] = set() + while len(done) < len(cursors): + for name, cursor in cursors.items(): + if name in done: + continue + try: + item = next(cursor) + except StopIteration: + done.add(name) + continue + yield (name, item) + + +class AsyncGraphRunStream: + """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] | None, + mux: StreamMux, + values_transformer: ValuesTransformer, + *, + wire_pump: bool = True, + ) -> None: + """Initialize the async run stream. + + Args: + graph_aiter: Async iterator over the graph's stream, or + `None` for nested run streams whose pump is driven by + an outer run (e.g. `AsyncSubgraphRunStream`). + mux: The StreamMux owning projections and the main log. + values_transformer: The built-in values transformer + providing `output` / `interrupted` / `interrupts`. + wire_pump: When True (default), bind `_apump_next` as the + mux's async pump callable. Subclasses that inherit a + parent pump via `StreamMux._make_child` should pass + False to preserve the parent binding. + """ + self._graph_aiter = graph_aiter + self._mux = mux + self.extensions: Mapping[str, Any] = MappingProxyType(mux.extensions) + self._values_transformer = values_transformer + self._exhausted = False + self._pump_cond = asyncio.Condition() + self._pumping = False + for key in mux.native_keys: + setattr(self, key, mux.extensions[key]) + if wire_pump: + self._wire_arequest_more(mux) + + def _wire_arequest_more(self, mux: StreamMux) -> None: + """Wire the async pull callback through the mux. + + Mirrors `_wire_request_more`: routing through + `mux.bind_apump` lets child mini-muxes inherit the pump + callable so cursors on subgraph handles drive the root + pump. + """ + 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 or self._graph_aiter is None: + 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__() + await self._mux.apush(convert_to_protocol_event(part)) + return True + 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 + 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 + + def __aiter__(self) -> AsyncIterator[ProtocolEvent]: + """Subscribe to the main event log and iterate protocol events.""" + return self._mux._events.__aiter__() + + +class _SubgraphRunStreamMixin: + """Subgraph metadata + parent-pump delegation shared by both lanes. + + Inherits from `GraphRunStream` (or `AsyncGraphRunStream`) with + `graph_iter=None` + `wire_pump=False` — the mini-mux is driven + by the parent's pump (inherited via `StreamMux._make_child`), and + the handle never pulls upstream itself. Pump-driving methods + delegate to the parent pump so `handle.output` and friends drive + the root run. + + Subclasses set the parent pump function captured at construction + (`_parent_pump_fn` / `_parent_apump_fn`) and override + `_pump_next` / `_apump_next` to delegate to it. + + Status is updated in place by `SubgraphTransformer`. Iterate + `run.subgraphs` to receive handles as subgraphs spawn, then + drill into projections inside the loop body **before** the next + pump cycle — same lazy-subscribe constraint as root projections. + """ + + path: tuple[str, ...] + graph_name: str | None + trigger_call_id: str | None + status: SubgraphStatus + error: str | None + _seen_terminal: bool + + +class SubgraphRunStream(GraphRunStream, _SubgraphRunStreamMixin): + """Sync handle for a discovered subgraph (extends `GraphRunStream`).""" + + def __init__( + self, + mux: StreamMux, + values_transformer: ValuesTransformer, + *, + path: tuple[str, ...], + graph_name: str | None = None, + trigger_call_id: str | None = None, + ) -> None: + # Capture the parent-inherited pump before super().__init__ + # touches anything; we delegate to it from `_pump_next`. + self._parent_pump_fn: Callable[[], bool] | None = mux._pump_fn + super().__init__( + graph_iter=None, + mux=mux, + values_transformer=values_transformer, + wire_pump=False, + ) + self.path = path + self.graph_name = graph_name + self.trigger_call_id = trigger_call_id + self.status = "started" + self.error = None + self._seen_terminal = False + + def _pump_next(self) -> bool: + """Delegate to the parent's pump. + + Cursors on this handle's projections call here when their + buffers empty. Driving the parent fans events into our + mini-mux, transparently advancing the whole run. + """ + if ( + self._exhausted + or self._seen_terminal + or self._mux._events._closed + or self._parent_pump_fn is None + ): + return False + return self._parent_pump_fn() + + +class AsyncSubgraphRunStream(AsyncGraphRunStream, _SubgraphRunStreamMixin): + """Async handle for a discovered subgraph (extends `AsyncGraphRunStream`).""" + + def __init__( + self, + mux: StreamMux, + values_transformer: ValuesTransformer, + *, + path: tuple[str, ...], + graph_name: str | None = None, + trigger_call_id: str | None = None, + ) -> None: + self._parent_apump_fn: Callable[[], Awaitable[bool]] | None = mux._apump_fn + super().__init__( + graph_aiter=None, + mux=mux, + values_transformer=values_transformer, + wire_pump=False, + ) + self.path = path + self.graph_name = graph_name + self.trigger_call_id = trigger_call_id + self.status = "started" + self.error = None + self._seen_terminal = False + + async def _apump_next(self) -> bool: + """Delegate to the parent's async pump.""" + if ( + self._exhausted + or self._seen_terminal + or self._mux._events._closed + or self._parent_apump_fn is None + ): + return False + return await self._parent_apump_fn() diff --git a/libs/langgraph/langgraph/stream/stream_channel.py b/libs/langgraph/langgraph/stream/stream_channel.py new file mode 100644 index 000000000..26533beb4 --- /dev/null +++ b/libs/langgraph/langgraph/stream/stream_channel.py @@ -0,0 +1,113 @@ +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:` 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:")`. + + 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) -> None: + """Initialize the channel with an empty inner log. + + Args: + name: The protocol channel name used for auto-forwarded + events. Surfaced on the wire as `custom:` for + user-defined transformers, or as `` for channels + owned by a native transformer (`_native = True`). The + StreamMux derives the prefix from the owning + transformer at registration time. + maxlen: Optional retention cap on the inner EventLog. See + `EventLog.__init__` for semantics. + """ + self.name = name + self._log: EventLog[T] = EventLog(maxlen=maxlen) + 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) diff --git a/libs/langgraph/langgraph/stream/transformers.py b/libs/langgraph/langgraph/stream/transformers.py new file mode 100644 index 000000000..90313f2e0 --- /dev/null +++ b/libs/langgraph/langgraph/stream/transformers.py @@ -0,0 +1,748 @@ +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 MessagesData +from typing_extensions import NotRequired, 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 AsyncSubgraphRunStream, SubgraphRunStream +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__) + + +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`). + + Only values events at the run's own level are captured; snapshots + from deeper subgraphs are left in the main event log but excluded + from the projection. "Own level" is defined by `scope`, which + `stream_v2` / `astream_v2` populate from the caller's checkpoint + namespace so that a nested `stream_v2` call still sees its own + root snapshots. + """ + + _native = True + required_stream_modes = ("values",) + + 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] = [] + # Cached as a list once for cheap equality with the protocol + # event's `namespace` field, which is `list[str]`. + self._scope_list: list[str] = list(scope) + + 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: + if event["method"] != "values": + return True + params = event["params"] + if params["namespace"] != self._scope_list: + return True + 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. + + Only events at the run's own level are projected; tokens from + deeper subgraphs are left in the main event log but excluded from + `.messages`. "Own level" is defined by `scope`, which + `stream_v2` / `astream_v2` populate from the caller's checkpoint + namespace so that a `stream_v2` call inside a node still sees its + own root chat model streams on `.messages`. Consumers that need + subgraph tokens should iterate the raw event stream or register a + custom transformer. + + Native transformer — the `messages` projection is exposed as a + direct attribute on the run stream. + """ + + _native = True + required_stream_modes = ("messages",) + + 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 + # Cached as a list once for cheap equality with the protocol + # event's `namespace` field, which is `list[str]`. + self._scope_list: list[str] = list(scope) + + 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: + if event["method"] != "messages": + return True + params = event["params"] + if params["namespace"] != self._scope_list: + return True + + 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=[], + 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=[], 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() + + +SubgraphStatus = Literal["started", "completed", "failed", "interrupted"] + + +def _parse_ns_segment(segment: str) -> tuple[str, str | None]: + """Split a namespace segment into `(graph_name, trigger_call_id)`. + + Segments are formatted `node_name:task_id` by `prepare_next_tasks`. + Returns `(segment, None)` if no `:` is present. + """ + name, sep, task_id = segment.partition(":") + return name, task_id if sep else None + + +class LifecyclePayload(TypedDict, total=False): + """Payload of a lifecycle event surfaced on the `lifecycle` channel. + + Auto-forwarded as `lifecycle` protocol events (no `custom:` prefix + because `LifecycleTransformer` is a native transformer) so remote + SDK clients receive the same data in-process consumers see via + `run.lifecycle`. + """ + + event: SubgraphStatus + namespace: list[str] + graph_name: NotRequired[str] + trigger_call_id: NotRequired[str] + error: NotRequired[str] + + +class _TasksLifecycleBase(StreamTransformer): + """Shared bookkeeping for `tasks`-event-driven lifecycle inference. + + Both `LifecycleTransformer` (wire-serializable channel) and + `SubgraphTransformer` (in-process navigation handles) discover + subgraphs by watching the same `tasks` stream — `started` on the + first event at a tracked namespace, terminal status when the + parent's `TaskResultPayload` arrives. Centralizing the dispatch + + open-set bookkeeping here keeps the inference rules from + drifting between the two surfaces. + + Subclasses provide three template-method hooks: + + - `_should_track(ns)` — scope filter (e.g. multi-depth vs + direct-children-only). + - `_on_started(ns, graph_name, trigger_call_id)` — first sighting + action (push payload / build handle / etc.). Called once per + discovered namespace. + - `_on_terminal(ns, status, error)` — terminal action (push + terminal payload / mark handle status). Called once per + tracked namespace at result time, or via `finalize` / `fail` + sweeps if no parent result arrived. + + Tasks events are suppressed from the main event log (`process` + returns False) — they're folded into whichever projection the + subclass populates; consumers iterating the raw protocol stream + see the higher-level view. + """ + + required_stream_modes = ("tasks",) + + def __init__(self, scope: tuple[str, ...] = ()) -> None: + super().__init__(scope) + self._seen: set[tuple[str, ...]] = set() + # Maps tracked namespace -> task_id of the parent task whose + # `TaskResultPayload` will close it. + self._open: dict[tuple[str, ...], str] = {} + + # --- Template-method hooks (subclass overrides) --- + + def _should_track(self, ns: tuple[str, ...]) -> bool: + """Scope filter — return True iff `ns` is in this transformer's region.""" + raise NotImplementedError + + def _on_started( + self, + ns: tuple[str, ...], + graph_name: str | None, + trigger_call_id: str | None, + ) -> None: + """Fired once per discovered namespace (first observed task event).""" + raise NotImplementedError + + def _on_terminal( + self, + ns: tuple[str, ...], + status: SubgraphStatus, + error: str | None, + ) -> None: + """Fired once per tracked namespace when its parent's result arrives, + or via finalize/fail safety-net sweeps. + """ + raise NotImplementedError + + # --- Dispatch + bookkeeping (shared) --- + + def process(self, event: ProtocolEvent) -> bool: + if event["method"] != "tasks": + return True + ns = tuple(event["params"]["namespace"]) + data = event["params"]["data"] + if "result" in data: + self._handle_task_result(ns, data) + else: + self._handle_task_start(ns) + # Tasks events are folded into the synthesized projections; + # suppress from the main event log so iterators don't double-see + # the same information in two shapes. + return False + + def _handle_task_start(self, ns: tuple[str, ...]) -> None: + if not self._should_track(ns) or ns in self._seen: + return + self._seen.add(ns) + graph_name, trigger_call_id = _parse_ns_segment(ns[-1]) + self._on_started(ns, graph_name or None, trigger_call_id) + if trigger_call_id is not None: + self._open[ns] = trigger_call_id + + def _pop_terminal_transitions( + self, ns: tuple[str, ...], data: dict[str, Any] + ) -> list[tuple[tuple[str, ...], SubgraphStatus, str | None]]: + """Return and remove tracked children closed by this task result.""" + result_id = data.get("id") + if not result_id: + return [] + transitions: list[tuple[tuple[str, ...], SubgraphStatus, str | None]] = [] + for child_ns, parent_task_id in list(self._open.items()): + if child_ns[:-1] != ns or parent_task_id != result_id: + continue + status, error = _terminal_from_result(data) + transitions.append((child_ns, status, error)) + del self._open[child_ns] + return transitions + + def _handle_task_result(self, ns: tuple[str, ...], data: dict[str, Any]) -> None: + for child_ns, status, error in self._pop_terminal_transitions(ns, data): + self._on_terminal(child_ns, status, error) + + def finalize(self) -> None: + """Emit `completed` for any tracked namespace still open at run end.""" + for ns in list(self._open): + self._on_terminal(ns, "completed", None) + self._open.clear() + + def fail(self, err: BaseException) -> None: + """Emit `failed` / `interrupted` for any tracked namespace still open.""" + is_interrupt = isinstance(err, GraphInterrupt) + status: SubgraphStatus = "interrupted" if is_interrupt else "failed" + error_str = None if is_interrupt else str(err) + for ns in list(self._open): + self._on_terminal(ns, status, error_str) + self._open.clear() + + +def _status_from_exception(err: BaseException) -> tuple[SubgraphStatus, str | None]: + """Map a run exception to a subgraph terminal status and error string.""" + if isinstance(err, GraphInterrupt): + return "interrupted", None + return "failed", str(err) + + +def _terminal_from_result( + payload: dict[str, Any], +) -> tuple[SubgraphStatus, str | None]: + """Map a `TaskResultPayload` to a `(status, error)` pair. + + Order matters: a result with both `error` and `interrupts` prefers + the interrupt classification, since `GraphInterrupt` manifests as + a populated `interrupts` list, not as `error`. + """ + if payload.get("interrupts"): + return "interrupted", None + error = payload.get("error") + if error: + return "failed", str(error) + return "completed", None + + +class LifecycleTransformer(_TasksLifecycleBase): + """Surface subgraph lifecycle as `lifecycle` protocol events. + + Pushes `LifecyclePayload` to a `StreamChannel` named `lifecycle`. + The channel is auto-forwarded by the mux so payloads land in the + main event log under `method = "lifecycle"` (native transformer — + no `custom:` prefix) — visible to remote SDK clients over the + wire and to in-process consumers via `run.lifecycle`. + + Tracks subgraphs at every depth strictly below the transformer's + scope, so a graph → subgraph → subgraph chain produces lifecycle + events for both nested levels in a flat stream. + + Native transformer — projection key `lifecycle` is exposed as + `run.lifecycle`. + """ + + _native = True + + def __init__(self, scope: tuple[str, ...] = ()) -> None: + super().__init__(scope) + self._channel: StreamChannel[LifecyclePayload] = StreamChannel("lifecycle") + + def init(self) -> dict[str, Any]: + return {"lifecycle": self._channel} + + def _should_track(self, ns: tuple[str, ...]) -> bool: + depth = len(self.scope) + return len(ns) > depth and ns[:depth] == self.scope + + def _on_started( + self, + ns: tuple[str, ...], + graph_name: str | None, + trigger_call_id: str | None, + ) -> None: + if trigger_call_id is None: + # Without a task id we can't correlate a parent-result + # event back to this namespace — skip the started payload + # and rely on finalize/fail to close. + return + payload: LifecyclePayload = {"event": "started", "namespace": list(ns)} + if graph_name: + payload["graph_name"] = graph_name + payload["trigger_call_id"] = trigger_call_id + self._channel.push(payload) + + def _on_terminal( + self, + ns: tuple[str, ...], + status: SubgraphStatus, + error: str | None, + ) -> None: + payload: LifecyclePayload = {"event": status, "namespace": list(ns)} + if error is not None: + payload["error"] = error + self._channel.push(payload) + + +class SubgraphTransformer(_TasksLifecycleBase): + """Discover subgraph invocations as in-process navigation handles. + + Per discovered direct-child subgraph, builds a `SubgraphRunStream` + (or `AsyncSubgraphRunStream`) wrapping a child mini-mux scoped to + the subgraph's namespace. Consumers iterate `run.subgraphs` to + receive handles, then drill into `handle.values` / `handle.messages` + / `handle.subgraphs` (recursive grandchildren) / `handle.lifecycle`. + + Each mini-mux owns its own scope and uses its own + `SubgraphTransformer` to discover its direct children, so + grandchildren live on the child handle — never on the root's + `subgraphs` log. Forwarding events into the matching child mini-mux + is what keeps the child's projections populated. + + Native transformer — `subgraphs` is exposed as `run.subgraphs`. + """ + + _native = True + supports_sync = True + + def __init__(self, scope: tuple[str, ...] = ()) -> None: + super().__init__(scope) + self._log: EventLog[SubgraphRunStream | AsyncSubgraphRunStream] = EventLog() + self._handles: dict[ + tuple[str, ...], SubgraphRunStream | AsyncSubgraphRunStream + ] = {} + self._mux: StreamMux | None = None + + def init(self) -> dict[str, Any]: + return {"subgraphs": self._log} + + def _on_register(self, mux: Any) -> None: + self._mux = mux + + def _should_track(self, ns: tuple[str, ...]) -> bool: + # Direct children only — grandchildren are picked up by the + # child mini-mux's own SubgraphTransformer. + depth = len(self.scope) + return len(ns) == depth + 1 and ns[:depth] == self.scope + + def _on_started( + self, + ns: tuple[str, ...], + graph_name: str | None, + trigger_call_id: str | None, + ) -> None: + if self._mux is None: + return + try: + child_mux = self._mux._make_child(ns) + except RuntimeError: + # Mux wasn't built from factories — no mini-mux navigation + # available. Skip; LifecycleTransformer still tracks the + # subgraph via the flat event stream. + return + values_t = child_mux.transformer_by_key("values") + if not isinstance(values_t, ValuesTransformer): + return + handle_cls = AsyncSubgraphRunStream if child_mux.is_async else SubgraphRunStream + handle = handle_cls( + mux=child_mux, + values_transformer=values_t, + path=ns, + graph_name=graph_name, + trigger_call_id=trigger_call_id, + ) + self._handles[ns] = handle + self._log.push(handle) + + def _on_terminal( + self, + ns: tuple[str, ...], + status: SubgraphStatus, + error: str | None, + ) -> None: + handle = self._handles.get(ns) + if handle is None or not self._mark_terminal(handle, status, error): + return + self._close_or_fail_handle(handle, status, error) + + async def _aon_terminal( + self, + ns: tuple[str, ...], + status: SubgraphStatus, + error: str | None, + ) -> None: + handle = self._handles.get(ns) + if handle is None or not self._mark_terminal(handle, status, error): + return + await self._aclose_or_fail_handle(handle, status, error) + + def _mark_terminal( + self, + handle: SubgraphRunStream | AsyncSubgraphRunStream, + status: SubgraphStatus, + error: str | None, + ) -> bool: + """Mark a handle terminal once. Returns True on first transition.""" + if handle._seen_terminal: + return False + handle.status = status + if error is not None and handle.error is None: + handle.error = error + handle._seen_terminal = True + return True + + def _close_or_fail_handle( + self, + handle: SubgraphRunStream | AsyncSubgraphRunStream, + status: SubgraphStatus, + error: str | None, + ) -> None: + if handle._mux is None or handle._mux._events._closed: + return + if status == "failed": + handle._mux.fail(RuntimeError(error or "Subgraph failed")) + else: + handle._mux.close() + + async def _aclose_or_fail_handle( + self, + handle: SubgraphRunStream | AsyncSubgraphRunStream, + status: SubgraphStatus, + error: str | None, + ) -> None: + if handle._mux is None or handle._mux._events._closed: + return + if status == "failed": + await handle._mux.afail(RuntimeError(error or "Subgraph failed")) + else: + await handle._mux.aclose() + + def _child_mux_for_event(self, event: ProtocolEvent) -> StreamMux | None: + ns = tuple(event["params"]["namespace"]) + depth = len(self.scope) + if len(ns) < depth + 1: + return None + handle = self._handles.get(ns[: depth + 1]) + if handle is None or handle._mux is None or handle._mux._events._closed: + return None + return handle._mux + + def process(self, event: ProtocolEvent) -> bool: + # Discover / update terminal status before forwarding so a + # `started` handle exists by the time the child mini-mux sees + # its own first event. + keep = super().process(event) + child_mux = self._child_mux_for_event(event) + if child_mux is not None: + child_mux.push(event) + return keep + + async def aprocess(self, event: ProtocolEvent) -> bool: + # Async counterpart to `process`: repeat the tasks bookkeeping + # here instead of delegating to `process`, so child mini-muxes + # receive events through their async lane. + if event["method"] == "tasks": + ns = tuple(event["params"]["namespace"]) + data = event["params"]["data"] + if "result" in data: + for child_ns, status, error in self._pop_terminal_transitions(ns, data): + await self._aon_terminal(child_ns, status, error) + else: + self._handle_task_start(ns) + keep = False + else: + keep = True + child_mux = self._child_mux_for_event(event) + if child_mux is not None: + await child_mux.apush(event) + return keep + + def _complete_open_handles(self) -> BaseException | None: + first_error: BaseException | None = None + for ns in list(self._open): + try: + self._on_terminal(ns, "completed", None) + except BaseException as e: + if first_error is None: + first_error = e + self._open.clear() + for handle in self._handles.values(): + if self._mark_terminal(handle, "completed", None): + try: + self._close_or_fail_handle(handle, "completed", None) + except BaseException as e: + if first_error is None: + first_error = e + return first_error + + async def _acomplete_open_handles(self) -> BaseException | None: + first_error: BaseException | None = None + for ns in list(self._open): + try: + await self._aon_terminal(ns, "completed", None) + except BaseException as e: + if first_error is None: + first_error = e + self._open.clear() + for handle in self._handles.values(): + if self._mark_terminal(handle, "completed", None): + try: + await self._aclose_or_fail_handle(handle, "completed", None) + except BaseException as e: + if first_error is None: + first_error = e + return first_error + + def finalize(self) -> None: + first_error = self._complete_open_handles() + if first_error is not None: + raise first_error + + async def afinalize(self) -> None: + first_error = await self._acomplete_open_handles() + if first_error is not None: + raise first_error + + def fail(self, err: BaseException) -> None: + status, error_str = _status_from_exception(err) + self._open.clear() + for handle in self._handles.values(): + self._mark_terminal(handle, status, error_str) + if handle._mux is not None and 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, + ) + + async def afail(self, err: BaseException) -> None: + status, error_str = _status_from_exception(err) + self._open.clear() + for handle in self._handles.values(): + self._mark_terminal(handle, status, error_str) + if handle._mux is not None and not handle._mux._events._closed: + try: + await handle._mux.afail(err) + except Exception: + _logger.warning( + "Error failing subgraph mini-mux at %s; " + "subscribers may not see the terminal error.", + handle.path, + exc_info=True, + ) diff --git a/libs/langgraph/pyproject.toml b/libs/langgraph/pyproject.toml index 6200a601a..d8664c6b6 100644 --- a/libs/langgraph/pyproject.toml +++ b/libs/langgraph/pyproject.toml @@ -24,7 +24,7 @@ classifiers = [ 'Programming Language :: Python :: 3.13', ] dependencies = [ - "langchain-core>=1.3.0,<2", + "langchain-core>=1.3.2,<2", "langgraph-checkpoint>=2.1.0,<5.0.0", "langgraph-sdk>=0.3.0,<0.4.0", "langgraph-prebuilt>=1.0.12,<1.1.0", diff --git a/libs/langgraph/tests/test_pregel_stream_v2.py b/libs/langgraph/tests/test_pregel_stream_v2.py new file mode 100644 index 000000000..c76e05f1a --- /dev/null +++ b/libs/langgraph/tests/test_pregel_stream_v2.py @@ -0,0 +1,1512 @@ +"""Tests for Pregel.stream_v2 / astream_v2 and the transformer pipeline.""" + +from __future__ import annotations + +import asyncio +import operator +import sys +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.graph import StateGraph +from langgraph.stream import ( + EventLog, + StreamChannel, + StreamTransformer, +) +from langgraph.stream._convert import convert_to_protocol_event +from langgraph.stream._mux import StreamMux +from langgraph.stream._types import ProtocolEvent +from langgraph.stream.transformers import MessagesTransformer, ValuesTransformer +from langgraph.types import StreamWriter, interrupt + +NEEDS_CONTEXTVARS = pytest.mark.skipif( + sys.version_info < (3, 11), + reason="Python 3.11+ is required for async contextvars support", +) + +TS = int(time.time() * 1000) + + +def _event( + method: str, + data: Any = None, + *, + namespace: list[str] | None = None, + interrupts: tuple[Any, ...] | None = None, +) -> ProtocolEvent: + params: dict[str, Any] = { + "namespace": namespace or [], + "timestamp": TS, + "data": data if data is not None else {}, + } + if interrupts is not None: + params["interrupts"] = interrupts + return {"type": "event", "method": method, "params": params} + + +# --------------------------------------------------------------------------- +# Shared graph builders +# --------------------------------------------------------------------------- + + +class SimpleState(TypedDict): + value: str + items: Annotated[list[str], operator.add] + + +def _build_simple_graph(): + def node_a(state: SimpleState) -> dict: + return {"value": state["value"] + "A", "items": ["a"]} + + def node_b(state: SimpleState) -> dict: + return {"value": state["value"] + "B", "items": ["b"]} + + builder = StateGraph(SimpleState) + builder.add_node("node_a", node_a) + builder.add_node("node_b", node_b) + builder.add_edge(START, "node_a") + builder.add_edge("node_a", "node_b") + builder.add_edge("node_b", END) + return builder.compile() + + +def _build_interrupt_graph(): + def node_a(state: SimpleState) -> dict: + return {"value": state["value"] + "A", "items": ["a"]} + + def node_b(state: SimpleState) -> dict: + interrupt("need approval") + return {"value": state["value"] + "B", "items": ["b"]} + + builder = StateGraph(SimpleState) + builder.add_node("node_a", node_a) + builder.add_node("node_b", node_b) + builder.add_edge(START, "node_a") + builder.add_edge("node_a", "node_b") + builder.add_edge("node_b", END) + return builder.compile(checkpointer=InMemorySaver()) + + +def _build_error_graph(): + def node_a(state: SimpleState) -> dict: + return {"value": state["value"] + "A", "items": ["a"]} + + def node_b(state: SimpleState) -> dict: + raise ValueError("boom") + + builder = StateGraph(SimpleState) + builder.add_node("node_a", node_a) + builder.add_node("node_b", node_b) + builder.add_edge(START, "node_a") + builder.add_edge("node_a", "node_b") + builder.add_edge("node_b", END) + return builder.compile() + + +def _build_custom_stream_graph(): + def node_a(state: SimpleState, *, writer: StreamWriter) -> dict: + writer({"step": "start"}) + writer({"step": "end"}) + return {"value": state["value"] + "A", "items": ["a"]} + + builder = StateGraph(SimpleState) + builder.add_node("node_a", node_a) + builder.add_edge(START, "node_a") + builder.add_edge("node_a", END) + return builder.compile() + + +class _CustomPassthroughTransformer(StreamTransformer): + """Opts a run into the `custom` stream mode without building a projection. + + `stream_v2` requests only the modes that registered transformers + declare via `required_stream_modes`. Custom events are raw user + emissions from `StreamWriter`, so tests that want them visible on + the main event log register this pass-through transformer. + """ + + required_stream_modes = ("custom",) + + def init(self) -> dict[str, Any]: + return {} + + def process(self, event: ProtocolEvent) -> bool: + return True + + +# --------------------------------------------------------------------------- +# EventLog unit tests +# --------------------------------------------------------------------------- + + +class TestEventLog: + def test_sync_iteration(self) -> None: + log: EventLog[int] = EventLog() + log._bind(is_async=False) + it = iter(log) + log.push(1) + log.push(2) + log.push(3) + log.close() + assert list(it) == [1, 2, 3] + + def test_drain_on_consume(self) -> None: + log: EventLog[str] = EventLog() + log._bind(is_async=False) + it = iter(log) + log.push("a") + log.push("b") + log.close() + assert list(it) == ["a", "b"] + assert list(log._items) == [] + + def test_second_subscribe_raises(self) -> None: + log: EventLog[str] = EventLog() + log._bind(is_async=False) + log.close() + _ = iter(log) + with pytest.raises(RuntimeError, match="already has a subscriber"): + iter(log) + + def test_pre_subscription_push_is_noop(self) -> None: + # Lazy-subscribe: pushes before subscription are dropped silently. + log: EventLog[int] = EventLog() + log._bind(is_async=False) + log.push(1) + log.push(2) + it = iter(log) + log.push(3) + log.close() + assert list(it) == [3] + + def test_fail_propagation(self) -> None: + log: EventLog[int] = EventLog() + log._bind(is_async=False) + it = iter(log) + log.push(1) + log.fail(ValueError("test error")) + with pytest.raises(ValueError, match="test error"): + list(it) + + def test_sync_cursor_yields_items_before_error(self) -> None: + log: EventLog[int] = EventLog() + log._bind(is_async=False) + it = iter(log) + log.push(1) + log.push(2) + log.push(3) + log.fail(ValueError("late error")) + items: list[int] = [] + with pytest.raises(ValueError, match="late error"): + for item in it: + items.append(item) + assert items == [1, 2, 3] + + def test_push_after_close_raises(self) -> None: + log: EventLog[int] = EventLog() + log._bind(is_async=False) + it = iter(log) + log.push(1) + log.close() + with pytest.raises(RuntimeError, match="Cannot push to a closed EventLog"): + log.push(2) + _ = list(it) + + def test_push_after_fail_raises(self) -> None: + log: EventLog[int] = EventLog() + log._bind(is_async=False) + it = iter(log) + log.fail(ValueError("err")) + with pytest.raises(RuntimeError, match="Cannot push to a closed EventLog"): + log.push(1) + with pytest.raises(ValueError, match="err"): + list(it) + + def test_empty_log_sync(self) -> None: + log: EventLog[int] = EventLog() + log._bind(is_async=False) + log.close() + assert list(log) == [] + + def test_empty_log_fail_sync(self) -> None: + log: EventLog[int] = EventLog() + log._bind(is_async=False) + log.fail(ValueError("empty fail")) + with pytest.raises(ValueError, match="empty fail"): + list(log) + + def test_unbound_iter_raises(self) -> None: + log: EventLog[int] = EventLog() + log.close() + with pytest.raises(TypeError, match="has not been bound"): + list(log) + + def test_sync_bound_aiter_raises(self) -> None: + log: EventLog[int] = EventLog() + log._bind(is_async=False) + log.close() + with pytest.raises(TypeError, match="bound to sync mode"): + log.__aiter__() + + def test_double_bind_raises(self) -> None: + log: EventLog[int] = EventLog() + log._bind(is_async=False) + with pytest.raises(RuntimeError, match="already bound"): + log._bind(is_async=True) + + @pytest.mark.anyio + async def test_async_iteration(self) -> None: + log: EventLog[int] = EventLog() + log._bind(is_async=True) + cursor = aiter(log) + for i in range(3): + log.push(i) + log.close() + assert [item async for item in cursor] == [0, 1, 2] + + @pytest.mark.anyio + async def test_async_second_subscribe_raises(self) -> None: + log: EventLog[str] = EventLog() + log._bind(is_async=True) + log.close() + _ = log.__aiter__() + with pytest.raises(RuntimeError, match="already has a subscriber"): + log.__aiter__() + + @pytest.mark.anyio + async def test_async_fail(self) -> None: + log: EventLog[int] = EventLog() + log._bind(is_async=True) + cursor = aiter(log) + log.push(1) + log.fail(RuntimeError("async error")) + with pytest.raises(RuntimeError, match="async error"): + async for _ in cursor: + pass + + @pytest.mark.anyio + async def test_async_cursor_yields_items_before_error(self) -> None: + log: EventLog[int] = EventLog() + log._bind(is_async=True) + cursor = aiter(log) + log.push(1) + log.push(2) + log.push(3) + log.fail(ValueError("late error")) + items: list[int] = [] + with pytest.raises(ValueError, match="late error"): + async for item in cursor: + items.append(item) + assert items == [1, 2, 3] + + @pytest.mark.anyio + async def test_empty_log_async(self) -> None: + log: EventLog[int] = EventLog() + log._bind(is_async=True) + log.close() + assert [item async for item in log] == [] + + @pytest.mark.anyio + async def test_empty_log_fail_async(self) -> None: + log: EventLog[int] = EventLog() + log._bind(is_async=True) + log.fail(ValueError("empty fail")) + with pytest.raises(ValueError, match="empty fail"): + async for _ in log: + pass + + @pytest.mark.anyio + async def test_async_bound_iter_raises(self) -> None: + log: EventLog[int] = EventLog() + log._bind(is_async=True) + log.close() + with pytest.raises(TypeError, match="bound to async mode"): + iter(log) + + +# --------------------------------------------------------------------------- +# StreamChannel unit tests +# --------------------------------------------------------------------------- + + +class TestStreamChannel: + def test_push_and_iterate(self) -> None: + ch: StreamChannel[str] = StreamChannel("test") + ch._bind(is_async=False) + it = iter(ch) + ch.push("a") + ch.push("b") + ch._close() + assert list(it) == ["a", "b"] + + def test_wire_callback(self) -> None: + forwarded: list[str] = [] + ch: StreamChannel[str] = StreamChannel("test") + ch._bind(is_async=False) + ch._wire(lambda item: forwarded.append(item)) + it = iter(ch) + ch.push("x") + ch.push("y") + ch._close() + assert forwarded == ["x", "y"] + assert list(it) == ["x", "y"] + + def test_fail_propagation(self) -> None: + ch: StreamChannel[str] = StreamChannel("test") + ch._bind(is_async=False) + it = iter(ch) + ch.push("a") + ch._fail(ValueError("channel error")) + items: list[str] = [] + with pytest.raises(ValueError, match="channel error"): + for item in it: + items.append(item) + assert items == ["a"] + + def test_push_without_wire(self) -> None: + ch: StreamChannel[int] = StreamChannel("test") + ch._bind(is_async=False) + assert ch._wire_fn is None + it = iter(ch) + ch.push(42) + ch._close() + assert list(it) == [42] + + @pytest.mark.anyio + async def test_async_iteration(self) -> None: + ch: StreamChannel[str] = StreamChannel("test") + ch._bind(is_async=True) + cursor = ch.__aiter__() + ch._log.push("x") + ch._log.push("y") + ch._close() + assert [item async for item in cursor] == ["x", "y"] + + +# --------------------------------------------------------------------------- +# stream_v2 sync tests +# --------------------------------------------------------------------------- + + +class TestStreamV2Sync: + def test_values_projection(self) -> None: + run = _build_simple_graph().stream_v2({"value": "x", "items": []}) + snapshots = list(run.values) + assert len(snapshots) >= 1 + last = snapshots[-1] + assert "A" in last["value"] and "B" in last["value"] + + def test_output(self) -> None: + run = _build_simple_graph().stream_v2({"value": "x", "items": []}) + output = run.output + assert output == {"value": "xAB", "items": ["a", "b"]} + + def test_raw_event_iteration(self) -> None: + run = _build_simple_graph().stream_v2({"value": "x", "items": []}) + events = list(run) + assert len(events) > 0 + for event in events: + assert event["type"] == "event" + assert "method" in event + assert "seq" in event + assert isinstance(event["params"]["timestamp"], int) + + def test_extensions_has_native_keys(self) -> None: + run = _build_simple_graph().stream_v2({"value": "x", "items": []}) + _ = run.output + assert "values" in run.extensions and "messages" in run.extensions + assert run.values is run.extensions["values"] + assert run.messages is run.extensions["messages"] + + def test_extensions_is_read_only(self) -> None: + run = _build_simple_graph().stream_v2({"value": "x", "items": []}) + with pytest.raises(TypeError): + run.extensions["new_key"] = object() # type: ignore[index] + with pytest.raises(TypeError): + del run.extensions["values"] # type: ignore[attr-defined] + + def test_custom_stream_events(self) -> None: + run = _build_custom_stream_graph().stream_v2( + {"value": "x", "items": []}, + transformers=[_CustomPassthroughTransformer], + ) + custom_events = [e for e in run if e["method"] == "custom"] + assert len(custom_events) == 2 + assert custom_events[0]["params"]["data"] == {"step": "start"} + assert custom_events[1]["params"]["data"] == {"step": "end"} + + def test_custom_events_suppressed_without_transformer(self) -> None: + """Without a transformer declaring `"custom"`, no custom events flow. + + `stream_v2` asks the graph only for the modes that registered + transformers require. Built-ins cover `values` / `messages`; + consumers that want raw custom events surface them by + registering a transformer whose `required_stream_modes` + includes `"custom"`. + """ + run = _build_custom_stream_graph().stream_v2({"value": "x", "items": []}) + custom_events = [e for e in run if e["method"] == "custom"] + assert custom_events == [] + + def test_interleave_values_and_messages(self) -> None: + run = _build_simple_graph().stream_v2({"value": "x", "items": []}) + tagged = list(run.interleave("values", "messages")) + names = [name for name, _ in tagged] + assert set(names).issubset({"values", "messages"}) + assert names.count("values") >= 1 + with pytest.raises(RuntimeError, match="already has a subscriber"): + list(run.values) + + def test_abort_marks_exhausted_and_closes_mux(self) -> None: + run = _build_simple_graph().stream_v2({"value": "x", "items": []}) + values_iter = iter(run.values) + _ = next(values_iter) + run.abort() + list(values_iter) + assert run._exhausted is True + run.abort() # idempotent + + def test_context_manager_calls_abort_on_exit(self) -> None: + with _build_simple_graph().stream_v2({"value": "x", "items": []}) as run: + _ = next(iter(run.values)) + assert run._exhausted is True + + def test_interleave_unknown_projection(self) -> None: + run = _build_simple_graph().stream_v2({"value": "x", "items": []}) + with pytest.raises(KeyError): + list(run.interleave("values", "does_not_exist")) + + +class TestStreamV2SyncErrors: + def test_error_propagation_output(self) -> None: + run = _build_error_graph().stream_v2({"value": "x", "items": []}) + with pytest.raises(ValueError, match="boom"): + _ = run.output + + def test_error_propagation_values(self) -> None: + run = _build_error_graph().stream_v2({"value": "x", "items": []}) + with pytest.raises(ValueError, match="boom"): + list(run.values) + + def test_error_propagation_raw_events(self) -> None: + run = _build_error_graph().stream_v2({"value": "x", "items": []}) + with pytest.raises(ValueError, match="boom"): + list(run) + + def test_error_propagation_interrupted(self) -> None: + run = _build_error_graph().stream_v2({"value": "x", "items": []}) + with pytest.raises(ValueError, match="boom"): + _ = run.interrupted + + def test_error_propagation_interrupts(self) -> None: + run = _build_error_graph().stream_v2({"value": "x", "items": []}) + with pytest.raises(ValueError, match="boom"): + _ = run.interrupts + + +class TestStreamV2SyncInterrupt: + def test_interrupted(self) -> None: + run = _build_interrupt_graph().stream_v2( + {"value": "x", "items": []}, + {"configurable": {"thread_id": "t1"}}, + ) + _ = run.output + assert run.interrupted is True + assert len(run.interrupts) > 0 + + +# --------------------------------------------------------------------------- +# astream_v2 async tests +# --------------------------------------------------------------------------- + + +@pytest.mark.anyio +@NEEDS_CONTEXTVARS +class TestStreamV2Async: + async def test_values_projection(self) -> None: + run = await _build_simple_graph().astream_v2({"value": "x", "items": []}) + snapshots = [s async for s in run.values] + assert len(snapshots) >= 1 + last = snapshots[-1] + assert "A" in last["value"] and "B" in last["value"] + + async def test_output(self) -> None: + run = await _build_simple_graph().astream_v2({"value": "x", "items": []}) + output = await run.output() + assert output == {"value": "xAB", "items": ["a", "b"]} + + async def test_raw_event_iteration(self) -> None: + run = await _build_simple_graph().astream_v2({"value": "x", "items": []}) + events = [e async for e in run] + assert len(events) > 0 + for event in events: + assert event["type"] == "event" + + async def test_abort_marks_exhausted_and_closes_mux(self) -> None: + run = await _build_simple_graph().astream_v2({"value": "x", "items": []}) + values_iter = aiter(run.values) + _ = await anext(values_iter) + await run.abort() + async for _item in values_iter: + pass + assert run._exhausted is True + await run.abort() # idempotent + + async def test_context_manager_calls_abort_on_exit(self) -> None: + run = await _build_simple_graph().astream_v2({"value": "x", "items": []}) + async with run: + _ = await anext(aiter(run.values)) + assert run._exhausted is True + + async def test_extensions_has_native_keys(self) -> None: + run = await _build_simple_graph().astream_v2({"value": "x", "items": []}) + _ = await run.output() + assert "values" in run.extensions and "messages" in run.extensions + assert run.values is run.extensions["values"] + assert run.messages is run.extensions["messages"] + + async def test_custom_stream_events(self) -> None: + run = await _build_custom_stream_graph().astream_v2( + {"value": "x", "items": []}, + transformers=[_CustomPassthroughTransformer], + ) + events = [e async for e in run] + custom_events = [e for e in events if e["method"] == "custom"] + assert len(custom_events) == 2 + assert custom_events[0]["params"]["data"] == {"step": "start"} + assert custom_events[1]["params"]["data"] == {"step": "end"} + + +@pytest.mark.anyio +@NEEDS_CONTEXTVARS +class TestStreamV2AsyncErrors: + async def test_error_propagation_output(self) -> None: + run = await _build_error_graph().astream_v2({"value": "x", "items": []}) + with pytest.raises(ValueError, match="boom"): + await run.output() + + async def test_error_propagation_values(self) -> None: + run = await _build_error_graph().astream_v2({"value": "x", "items": []}) + with pytest.raises(ValueError, match="boom"): + async for _ in run.values: + pass + + async def test_error_propagation_raw_events(self) -> None: + run = await _build_error_graph().astream_v2({"value": "x", "items": []}) + with pytest.raises(ValueError, match="boom"): + async for _ in run: + pass + + async def test_error_propagation_interrupted(self) -> None: + run = await _build_error_graph().astream_v2({"value": "x", "items": []}) + with pytest.raises(ValueError, match="boom"): + await run.interrupted() + + async def test_error_propagation_interrupts(self) -> None: + run = await _build_error_graph().astream_v2({"value": "x", "items": []}) + with pytest.raises(ValueError, match="boom"): + await run.interrupts() + + +@pytest.mark.anyio +@NEEDS_CONTEXTVARS +class TestStreamV2AsyncInterrupt: + async def test_interrupted(self) -> None: + run = await _build_interrupt_graph().astream_v2( + {"value": "x", "items": []}, + {"configurable": {"thread_id": "t2"}}, + ) + _ = await run.output() + assert await run.interrupted() is True + assert len(await run.interrupts()) > 0 + + +# --------------------------------------------------------------------------- +# convert_to_protocol_event unit tests +# --------------------------------------------------------------------------- + + +class TestConvertToProtocolEvent: + def test_basic_conversion(self) -> None: + before = int(time.time() * 1000) + event = convert_to_protocol_event( + {"type": "values", "ns": ("sub", "graph"), "data": {"key": "val"}} + ) + after = int(time.time() * 1000) + assert event["type"] == "event" + assert event["method"] == "values" + assert event["params"]["namespace"] == ["sub", "graph"] + assert event["params"]["data"] == {"key": "val"} + assert "interrupts" not in event["params"] + assert before <= event["params"]["timestamp"] <= after + + def test_conversion_with_interrupts(self) -> None: + event = convert_to_protocol_event( + { + "type": "values", + "ns": (), + "data": {"k": 1}, + "interrupts": ({"value": "pause"},), + } + ) + assert event["params"]["interrupts"] == ({"value": "pause"},) + + def test_namespace_tuple_becomes_list(self) -> None: + event = convert_to_protocol_event( + {"type": "updates", "ns": ("a", "b", "c"), "data": {}} + ) + assert event["params"]["namespace"] == ["a", "b", "c"] + + +# --------------------------------------------------------------------------- +# StreamMux unit tests +# --------------------------------------------------------------------------- + + +class TestStreamMux: + def test_register_non_dict_raises(self) -> None: + class BadTransformer(StreamTransformer): + def init(self) -> Any: + return ["not", "a", "dict"] + + def process(self, event: ProtocolEvent) -> bool: + return True + + with pytest.raises(TypeError, match="must return a dict"): + StreamMux([BadTransformer()]) + + def test_event_suppression(self) -> None: + class FilterTransformer(StreamTransformer): + def init(self) -> dict[str, Any]: + return {} + + def process(self, event: ProtocolEvent) -> bool: + return event["method"] != "updates" + + mux = StreamMux([FilterTransformer()]) + it = iter(mux._events) + mux.push(_event("values", {"a": 1})) + mux.push(_event("updates", {"b": 2})) + mux.push(_event("custom", {"c": 3})) + mux.close() + assert [e["method"] for e in it] == ["values", "custom"] + + def test_suppression_all_transformers_still_see_event(self) -> None: + """If any transformer returns False, the event is suppressed from the main + log, but all transformers still receive it.""" + seen_by_second: list[str] = [] + + class PassTransformer(StreamTransformer): + def init(self) -> dict[str, Any]: + return {} + + def process(self, event: ProtocolEvent) -> bool: + return True + + class RejectTransformer(StreamTransformer): + def init(self) -> dict[str, Any]: + return {} + + def process(self, event: ProtocolEvent) -> bool: + seen_by_second.append(event["method"]) + return False + + mux = StreamMux([PassTransformer(), RejectTransformer()]) + mux.push(_event("values")) + mux.close() + assert seen_by_second == ["values"] + assert list(mux._events) == [] + + def test_empty_mux(self) -> None: + mux = StreamMux() + it = iter(mux._events) + mux.push(_event("values", {"x": 1})) + mux.close() + events = list(it) + assert len(events) == 1 + assert events[0]["method"] == "values" + + def test_empty_mux_fail(self) -> None: + mux = StreamMux() + mux.fail(ValueError("boom")) + with pytest.raises(ValueError, match="boom"): + list(mux._events) + + +# --------------------------------------------------------------------------- +# ValuesTransformer / MessagesTransformer unit tests +# --------------------------------------------------------------------------- + + +class TestValuesTransformer: + def test_ignores_non_root_namespace(self) -> None: + t = ValuesTransformer() + t.init() + t._log._bind(is_async=False) + it = iter(t._log) + t.process(_event("values", {"val": "root"})) + t.process(_event("values", {"val": "sub"}, namespace=["sub"])) + t._log.close() + items = list(it) + assert len(items) == 1 + assert items[0]["val"] == "root" + + def test_ignores_non_values_methods(self) -> None: + t = ValuesTransformer() + t.init() + t._log._bind(is_async=False) + it = iter(t._log) + assert t.process(_event("updates", {"x": 1})) is True + t._log.close() + assert list(it) == [] + + def test_tracks_interrupts(self) -> None: + t = ValuesTransformer() + t.init() + t.process( + _event( + "values", + {"v": 1}, + interrupts=({"value": "pause1"}, {"value": "pause2"}), + ) + ) + assert t._interrupted is True + assert len(t._interrupts) == 2 + + +class TestMessagesTransformer: + def test_captures_root_messages(self) -> None: + t = MessagesTransformer() + t.init() + t._log._bind(is_async=False) + t._bind_pump(lambda: False) + it = iter(t._log) + meta = {"langgraph_node": "llm", "run_id": "run-1"} + for evt in ( + {"event": "message-start", "role": "ai", "message_id": "run-1"}, + {"event": "message-finish", "reason": "stop"}, + ): + t.process(_event("messages", (evt, meta))) + t._log.close() + items = list(it) + assert len(items) == 1 + assert hasattr(items[0], "dispatch") + assert items[0].message_id == "run-1" + + def test_ignores_non_root_namespace(self) -> None: + t = MessagesTransformer() + t.init() + t._log._bind(is_async=False) + t._bind_pump(lambda: False) + it = iter(t._log) + meta = {"langgraph_node": "llm", "run_id": "run-1"} + t.process( + _event( + "messages", + ({"event": "message-start", "message_id": "run-1"}, meta), + namespace=["sub"], + ) + ) + t._log.close() + assert list(it) == [] + + def test_ignores_non_messages_methods(self) -> None: + t = MessagesTransformer() + t.init() + t._log._bind(is_async=False) + it = iter(t._log) + assert t.process(_event("values", {"v": 1})) is True + t._log.close() + assert list(it) == [] + + def test_fail_propagates(self) -> None: + t = MessagesTransformer() + t.init() + t._log._bind(is_async=False) + it = iter(t._log) + t._log.fail(ValueError("msg error")) + with pytest.raises(ValueError, match="msg error"): + list(it) + + +# --------------------------------------------------------------------------- +# StreamMux resilience: close/fail continue cleanup on transformer errors +# --------------------------------------------------------------------------- + + +class TestStreamMuxResilience: + def test_close_continues_after_finalize_error(self) -> None: + class BrokenFinalizer(StreamTransformer): + def init(self) -> dict[str, Any]: + return {} + + def process(self, event: ProtocolEvent) -> bool: + return True + + def finalize(self) -> None: + raise RuntimeError("finalize broke") + + class GoodTransformer(StreamTransformer): + def __init__(self) -> None: + self.finalized = False + + def init(self) -> dict[str, Any]: + return {} + + def process(self, event: ProtocolEvent) -> bool: + return True + + def finalize(self) -> None: + self.finalized = True + + good = GoodTransformer() + mux = StreamMux([BrokenFinalizer(), good]) + mux.push(_event("values")) + with pytest.raises(RuntimeError, match="finalize broke"): + mux.close() + assert good.finalized + assert mux._events._closed + + def test_fail_continues_after_transformer_error(self) -> None: + class BrokenFailer(StreamTransformer): + def init(self) -> dict[str, Any]: + return {} + + def process(self, event: ProtocolEvent) -> bool: + return True + + def fail(self, err: BaseException) -> None: + raise RuntimeError("fail handler broke") + + class GoodTransformer(StreamTransformer): + def __init__(self) -> None: + self.failed_with: BaseException | None = None + + def init(self) -> dict[str, Any]: + return {} + + def process(self, event: ProtocolEvent) -> bool: + return True + + def fail(self, err: BaseException) -> None: + self.failed_with = err + + good = GoodTransformer() + mux = StreamMux([BrokenFailer(), good]) + original_error = ValueError("original") + mux.fail(original_error) + assert good.failed_with is original_error + assert mux._events._error is original_error + + def test_channels_closed_after_finalize_error(self) -> None: + class BrokenWithChannel(StreamTransformer): + def __init__(self) -> None: + self._channel: StreamChannel[str] = StreamChannel("ch") + + def init(self) -> dict[str, Any]: + return {"ch": self._channel} + + def process(self, event: ProtocolEvent) -> bool: + return True + + def finalize(self) -> None: + raise RuntimeError("finalize broke") + + t = BrokenWithChannel() + mux = StreamMux([t]) + with pytest.raises(RuntimeError, match="finalize broke"): + mux.close() + assert t._channel._log._closed + + +# --------------------------------------------------------------------------- +# Custom transformer tests +# --------------------------------------------------------------------------- + + +class TestCustomTransformer: + def test_extension_transformer_with_stream_channel(self) -> None: + class CounterTransformer(StreamTransformer): + def __init__(self, scope: tuple[str, ...] = ()) -> None: + super().__init__(scope) + self._channel: StreamChannel[int] = StreamChannel("counter") + self._count = 0 + + def init(self) -> dict[str, Any]: + return {"counter": self._channel} + + def process(self, event: ProtocolEvent) -> bool: + if event["method"] == "values": + self._count += 1 + self._channel.push(self._count) + return True + + run = _build_simple_graph().stream_v2( + {"value": "x", "items": []}, transformers=[CounterTransformer] + ) + assert "counter" in run.extensions + counter_iter = iter(run.extensions["counter"]) + _ = run.output + counts = list(counter_iter) + assert len(counts) > 0 + assert not hasattr(run, "counter") # non-native: no direct attribute + + def test_native_transformer_gets_direct_attr(self) -> None: + class FooTransformer(StreamTransformer): + _native = True + + def __init__(self, scope: tuple[str, ...] = ()) -> None: + super().__init__(scope) + self._log: EventLog[str] = EventLog() + + def init(self) -> dict[str, Any]: + return {"foo": self._log} + + def process(self, event: ProtocolEvent) -> bool: + if event["method"] == "values": + self._log.push("saw_values") + return True + + run = _build_simple_graph().stream_v2( + {"value": "x", "items": []}, transformers=[FooTransformer] + ) + foo_iter = iter(run.foo) + _ = run.output + assert "foo" in run.extensions and run.foo is run.extensions["foo"] + assert "saw_values" in list(foo_iter) + + def test_stream_v2_rejects_transformer_instances(self) -> None: + class InstanceTransformer(StreamTransformer): + def init(self) -> dict[str, Any]: + return {} + + def process(self, event: ProtocolEvent) -> bool: + return True + + with pytest.raises(TypeError, match="pre-built instance"): + _build_simple_graph().stream_v2( + {"value": "x", "items": []}, transformers=[InstanceTransformer()] + ) + + def test_stream_channel_auto_forward(self) -> None: + """StreamChannel pushes inject ProtocolEvents into the main log.""" + + class EmitterTransformer(StreamTransformer): + def __init__(self, scope: tuple[str, ...] = ()) -> None: + super().__init__(scope) + self._channel: StreamChannel[str] = StreamChannel("emitter") + + def init(self) -> dict[str, Any]: + return {"emitter": self._channel} + + def process(self, event: ProtocolEvent) -> bool: + if event["method"] == "values": + self._channel.push("emitted") + return True + + run = _build_simple_graph().stream_v2( + {"value": "x", "items": []}, transformers=[EmitterTransformer] + ) + custom_events = [e for e in run if e["method"] == "custom:emitter"] + assert len(custom_events) > 0 + assert custom_events[0]["params"]["data"] == "emitted" + + def test_stream_channel_seq_ordering(self) -> None: + """Seq numbers must be monotonically increasing even when a channel push + auto-forwards an event mid-pipeline.""" + + class ChannelPusher(StreamTransformer): + def __init__(self) -> None: + super().__init__() + self._channel: StreamChannel[str] = StreamChannel("ch") + + def init(self) -> dict[str, Any]: + return {"ch": self._channel} + + def process(self, event: ProtocolEvent) -> bool: + self._channel.push(f"saw:{event['method']}") + return True + + mux = StreamMux([ChannelPusher()]) + it = iter(mux._events) + mux.push(_event("values")) + mux.push(_event("updates")) + mux.close() + seqs = [e["seq"] for e in it] + for i in range(1, len(seqs)): + assert seqs[i] > seqs[i - 1], f"Seq out of order at index {i}: {seqs}" + + def test_projection_key_conflict_raises(self) -> None: + class ConflictTransformer(StreamTransformer): + def __init__(self, scope: tuple[str, ...] = ()) -> None: + super().__init__(scope) + self._log: EventLog[str] = EventLog() + + def init(self) -> dict[str, Any]: + return {"values": self._log} + + def process(self, event: ProtocolEvent) -> bool: + return True + + with pytest.raises(ValueError, match=r"conflict.*'values'.*ValuesTransformer"): + _build_simple_graph().stream_v2( + {"value": "x", "items": []}, transformers=[ConflictTransformer] + ) + + +# --------------------------------------------------------------------------- +# EventLog auto-lifecycle via StreamMux +# --------------------------------------------------------------------------- + + +class TestEventLogAutoLifecycle: + def test_mux_auto_closes_event_logs(self) -> None: + class SimpleTransformer(StreamTransformer): + def __init__(self) -> None: + super().__init__() + self._log: EventLog[str] = EventLog() + + def init(self) -> dict[str, Any]: + return {"items": self._log} + + def process(self, event: ProtocolEvent) -> bool: + self._log.push("saw_event") + return True + + mux = StreamMux([SimpleTransformer()]) + it = iter(mux._events) + mux.push(_event("values")) + mux.close() + assert len(list(it)) == 1 + + def test_mux_auto_fails_event_logs(self) -> None: + class SimpleTransformer(StreamTransformer): + def __init__(self) -> None: + super().__init__() + self._log: EventLog[str] = EventLog() + + def init(self) -> dict[str, Any]: + return {"items": self._log} + + def process(self, event: ProtocolEvent) -> bool: + self._log.push("saw_event") + return True + + t = SimpleTransformer() + mux = StreamMux([t]) + it = iter(t._log) + mux.push(_event("values")) + mux.fail(ValueError("boom")) + with pytest.raises(ValueError, match="boom"): + list(it) + + def test_no_double_close_if_transformer_closes_own_log(self) -> None: + class ManualCloseTransformer(StreamTransformer): + def __init__(self) -> None: + super().__init__() + self._log: EventLog[str] = EventLog() + + def init(self) -> dict[str, Any]: + return {"items": self._log} + + def process(self, event: ProtocolEvent) -> bool: + return True + + def finalize(self) -> None: + self._log.close() + + mux = StreamMux([ManualCloseTransformer()]) + mux.close() # should not raise even with double-close + + def test_transformer_without_finalize_works(self) -> None: + class MinimalTransformer(StreamTransformer): + def __init__(self, scope: tuple[str, ...] = ()) -> None: + super().__init__(scope) + self._log: EventLog[str] = EventLog() + + def init(self) -> dict[str, Any]: + return {"minimal": self._log} + + def process(self, event: ProtocolEvent) -> bool: + if event["method"] == "values": + self._log.push("got_it") + return True + + run = _build_simple_graph().stream_v2( + {"value": "x", "items": []}, transformers=[MinimalTransformer] + ) + minimal_iter = iter(run.extensions["minimal"]) + _ = run.output + assert len(list(minimal_iter)) > 0 + + +class TestStreamTransformerSchedule: + def test_schedule_without_running_loop_raises(self) -> None: + class Sched(StreamTransformer): + requires_async = True + + def init(self) -> dict[str, Any]: + return {} + + def process(self, event: ProtocolEvent) -> bool: + return True + + t = Sched() + + async def noop() -> None: + pass + + coro = noop() + try: + with pytest.raises(RuntimeError, match="requires a running event loop"): + t.schedule(coro) + finally: + coro.close() + + +# --------------------------------------------------------------------------- +# Async transformer lane +# --------------------------------------------------------------------------- + + +@pytest.mark.anyio +class TestAsyncTransformerLane: + async def test_aprocess_is_awaited_before_next_transformer(self) -> None: + """aprocess must complete before the next transformer sees the event — + load-bearing guarantee for mutating transformers like PII redaction.""" + order: list[str] = [] + + class RedactTransformer(StreamTransformer): + requires_async = True + + def init(self) -> dict[str, Any]: + return {} + + async def aprocess(self, event: ProtocolEvent) -> bool: + await asyncio.sleep(0.01) + order.append("redact") + event["params"]["data"]["redacted"] = True + return True + + class ObserverTransformer(StreamTransformer): + def init(self) -> dict[str, Any]: + return {} + + def process(self, event: ProtocolEvent) -> bool: + if event["method"] == "values": + order.append(f"observe:{event['params']['data'].get('redacted')}") + return True + + mux = StreamMux([RedactTransformer(), ObserverTransformer()], is_async=True) + await mux.apush(_event("values", {"secret": "x"})) + await mux.aclose() + assert order == ["redact", "observe:True"] + + async def test_schedule_joins_tasks_before_afinalize(self) -> None: + """Every scheduled task must complete before afinalize runs.""" + phase: list[str] = [] + + class SchedTransformer(StreamTransformer): + requires_async = True + + def __init__(self) -> None: + self._log: EventLog[str] = EventLog() + + def init(self) -> dict[str, Any]: + return {"out": self._log} + + def process(self, event: ProtocolEvent) -> bool: + if event["method"] == "values": + + async def work() -> None: + await asyncio.sleep(0.01) + phase.append("task") + self._log.push("done") + + self.schedule(work()) + return True + + async def afinalize(self) -> None: + phase.append("afinalize") + self._log.close() + + t = SchedTransformer() + mux = StreamMux([t], is_async=True) + await mux.apush(_event("values", {})) + await mux.apush(_event("values", {})) + await mux.aclose() + assert phase.count("task") == 2 + assert phase[-1] == "afinalize" + + async def test_sync_stream_rejects_async_transformer(self) -> None: + class NeedsAsync(StreamTransformer): + requires_async = True + + def init(self) -> dict[str, Any]: + return {} + + def process(self, event: ProtocolEvent) -> bool: + return True + + with pytest.raises(RuntimeError, match="requires an async run"): + StreamMux([NeedsAsync()], is_async=False) + + async def test_sync_stream_rejects_aprocess_override(self) -> None: + class HasAprocess(StreamTransformer): + def init(self) -> dict[str, Any]: + return {} + + async def aprocess(self, event: ProtocolEvent) -> bool: + return True + + with pytest.raises(RuntimeError, match="requires an async run"): + StreamMux([HasAprocess()], is_async=False) + + async def test_schedule_on_error_log_swallows_exceptions(self) -> None: + class Bad(StreamTransformer): + requires_async = True + + def __init__(self) -> None: + self._log: EventLog[str] = EventLog() + + def init(self) -> dict[str, Any]: + return {"out": self._log} + + def process(self, event: ProtocolEvent) -> bool: + async def work() -> None: + raise ValueError("boom") + + self.schedule(work()) # default on_error="log" + return True + + async def afinalize(self) -> None: + self._log.close() + + mux = StreamMux([Bad()], is_async=True) + await mux.apush(_event("values", {})) + await mux.aclose() # should not raise; exception is logged + + async def test_schedule_on_error_raise_fails_the_run(self) -> None: + class Strict(StreamTransformer): + requires_async = True + + def init(self) -> dict[str, Any]: + return {} + + def process(self, event: ProtocolEvent) -> bool: + async def work() -> None: + raise ValueError("strict boom") + + self.schedule(work(), on_error="raise") + return True + + mux = StreamMux([Strict()], is_async=True) + await mux.apush(_event("values", {})) + with pytest.raises(ValueError, match="strict boom"): + await mux.aclose() + + async def test_afail_cancels_pending_scheduled_tasks(self) -> None: + cancelled = asyncio.Event() + + class Sched(StreamTransformer): + requires_async = True + + def init(self) -> dict[str, Any]: + return {} + + def process(self, event: ProtocolEvent) -> bool: + async def work() -> None: + try: + await asyncio.sleep(5) + except asyncio.CancelledError: + cancelled.set() + raise + + self.schedule(work()) + return True + + mux = StreamMux([Sched()], is_async=True) + await mux.apush(_event("values", {})) + # Yield so the task actually starts before we cancel it. + await asyncio.sleep(0) + await mux.afail(RuntimeError("run died")) + assert cancelled.is_set() + + async def test_mixed_sync_and_async_transformers(self) -> None: + seen_sync: list[str] = [] + + class SyncOne(StreamTransformer): + def init(self) -> dict[str, Any]: + return {} + + def process(self, event: ProtocolEvent) -> bool: + seen_sync.append(event["method"]) + return True + + class AsyncOne(StreamTransformer): + requires_async = True + + def __init__(self) -> None: + self._log: EventLog[str] = EventLog() + + def init(self) -> dict[str, Any]: + return {"seen": self._log} + + async def aprocess(self, event: ProtocolEvent) -> bool: + await asyncio.sleep(0) + self._log.push(event["method"]) + return True + + async def afinalize(self) -> None: + self._log.close() + + async_t = AsyncOne() + mux = StreamMux([SyncOne(), async_t], is_async=True) + seen_cursor = aiter(async_t._log) + await mux.apush(_event("values", {})) + await mux.apush(_event("updates", {})) + await mux.aclose() + assert seen_sync == ["values", "updates"] + assert [x async for x in seen_cursor] == ["values", "updates"] + + async def test_handler_astream_with_scheduled_work(self) -> None: + class Scorer(StreamTransformer): + requires_async = True + + def __init__(self, scope: tuple[str, ...] = ()) -> None: + super().__init__(scope) + self._log: EventLog[int] = EventLog() + + def init(self) -> dict[str, Any]: + return {"scores": self._log} + + def process(self, event: ProtocolEvent) -> bool: + if event["method"] == "values": + + async def work() -> None: + await asyncio.sleep(0.01) + self._log.push(42) + + self.schedule(work()) + return True + + async def afinalize(self) -> None: + self._log.close() + + run = await _build_simple_graph().astream_v2( + {"value": "x", "items": []}, transformers=[Scorer] + ) + scores_cursor = aiter(run.extensions["scores"]) + _ = await run.output() + scores = [x async for x in scores_cursor] + assert scores and all(s == 42 for s in scores) + + +# --------------------------------------------------------------------------- +# Memory bounds: drain-on-consume semantics +# --------------------------------------------------------------------------- + + +@NEEDS_CONTEXTVARS +class TestMemoryBounds: + def test_sync_subscribed_buffer_stays_at_most_one_between_yields(self) -> None: + run = _build_simple_graph().stream_v2({"value": "x", "items": []}) + events_iter = iter(run) + max_buffered = 0 + count = 0 + for _ in events_iter: + max_buffered = max(max_buffered, len(run._mux._events._items)) + count += 1 + assert count > 0 + assert max_buffered == 0, ( + f"drain-on-consume violated, observed max {max_buffered}" + ) + + def test_unsubscribed_projections_never_accumulate(self) -> None: + run = _build_simple_graph().stream_v2({"value": "x", "items": []}) + list(run) + values_log = run.extensions["values"] + messages_log = run.extensions["messages"] + assert len(values_log._items) == 0 and not values_log._subscribed + assert len(messages_log._items) == 0 and not messages_log._subscribed + + def test_output_path_does_not_retain_values(self) -> None: + run = _build_simple_graph().stream_v2({"value": "x", "items": []}) + _ = run.output + values_log = run.extensions["values"] + assert len(values_log._items) == 0 and not values_log._subscribed + + def test_drained_subscriber_buffer_returns_to_empty(self) -> None: + run = _build_simple_graph().stream_v2({"value": "x", "items": []}) + list(run.values) + assert len(run.extensions["values"]._items) == 0 + + @pytest.mark.anyio + async def test_async_single_consumer_buffer_stays_at_most_one(self) -> None: + run = await _build_simple_graph().astream_v2({"value": "x", "items": []}) + max_buffered = 0 + count = 0 + async for _ in run: + max_buffered = max(max_buffered, len(run._mux._events._items)) + count += 1 + assert count > 0 + assert max_buffered == 0 + + @pytest.mark.anyio + async def test_async_unsubscribed_projections_never_accumulate(self) -> None: + run = await _build_simple_graph().astream_v2({"value": "x", "items": []}) + _ = await run.output() + values_log = run.extensions["values"] + messages_log = run.extensions["messages"] + assert len(values_log._items) == 0 and not values_log._subscribed + assert len(messages_log._items) == 0 and not messages_log._subscribed + + +# --------------------------------------------------------------------------- +# DrainOnConsume: EventLog capacity semantics +# --------------------------------------------------------------------------- + + +class TestDrainOnConsume: + def test_invalid_maxlen_raises(self) -> None: + with pytest.raises(ValueError, match="positive int or None"): + EventLog(maxlen=0) + with pytest.raises(ValueError, match="positive int or None"): + EventLog(maxlen=-3) + + def test_push_unbounded_by_design(self) -> None: + """Push is non-blocking; the caller-driven pump bounds memory via iteration pace.""" + log: EventLog[int] = EventLog() + log._bind(is_async=False) + it = iter(log) + for i in range(100): + log.push(i) + log.close() + assert list(it) == list(range(100)) + + def test_tee_fans_out_sync(self) -> None: + log: EventLog[int] = EventLog() + log._bind(is_async=False) + a, b = log.tee(2) + for i in range(3): + log.push(i) + log.close() + assert list(a) == [0, 1, 2] + assert list(b) == [0, 1, 2] + + @pytest.mark.anyio + async def test_atee_fans_out(self) -> None: + log: EventLog[int] = EventLog() + log._bind(is_async=True) + a, b = log.atee(2) + for i in range(3): + log.push(i) + log.close() + assert [x async for x in a] == [0, 1, 2] + assert [x async for x in b] == [0, 1, 2] diff --git a/libs/langgraph/tests/test_stream_lifecycle_transformer.py b/libs/langgraph/tests/test_stream_lifecycle_transformer.py new file mode 100644 index 000000000..df5c7f1c4 --- /dev/null +++ b/libs/langgraph/tests/test_stream_lifecycle_transformer.py @@ -0,0 +1,401 @@ +"""Tests for LifecycleTransformer. + +Consumes the `tasks` stream mode and emits subgraph lifecycle payloads +on the `lifecycle` channel for both in-process iteration via +`run.lifecycle` and wire delivery via `custom:lifecycle` protocol +events. Most tests dispatch synthetic protocol events through a +`StreamMux` to keep the inference logic isolated; the end-of-file +group exercises the path through real graphs (multi-depth +discovery, nested `stream_v2` calls with non-empty `parent_ns`). +""" + +from __future__ import annotations + +import operator +import time +from typing import Annotated, Any + +from typing_extensions import TypedDict + +from langgraph._internal._constants import CONF, CONFIG_KEY_CHECKPOINT_NS +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.transformers import ( + LifecyclePayload, + LifecycleTransformer, +) + +TS = int(time.time() * 1000) + + +def _tasks_start( + namespace: list[str], + *, + task_id: str, + name: str, +) -> dict[str, Any]: + """Build a `tasks` ProtocolEvent carrying a TaskPayload (start).""" + return { + "type": "event", + "method": "tasks", + "params": { + "namespace": namespace, + "timestamp": TS, + "data": { + "id": task_id, + "name": name, + "input": None, + "triggers": [], + }, + }, + } + + +def _tasks_result( + namespace: list[str], + *, + task_id: str, + name: str, + error: str | None = None, + interrupts: list[dict[str, Any]] | None = None, +) -> dict[str, Any]: + """Build a `tasks` ProtocolEvent carrying a TaskResultPayload (finish).""" + return { + "type": "event", + "method": "tasks", + "params": { + "namespace": namespace, + "timestamp": TS, + "data": { + "id": task_id, + "name": name, + "error": error, + "interrupts": interrupts or [], + "result": {}, + }, + }, + } + + +def _arm(mux: StreamMux) -> None: + """Force projection logs to accept pushes (skip lazy-subscribe gate). + + `EventLog.push` is a no-op until a subscriber attaches. Tests that + inspect `_items` directly need the gate flipped before any event + is dispatched. + """ + mux._events._subscribed = True + for transformer in mux._transformers: + if isinstance(transformer, LifecycleTransformer): + transformer._channel._log._subscribed = True + + +def _drain_lifecycle(mux: StreamMux) -> list[LifecyclePayload]: + """Snapshot the lifecycle channel's underlying log.""" + transformer = mux.transformer_by_key("lifecycle") + assert isinstance(transformer, LifecycleTransformer) + return list(transformer._channel._log._items) + + +def _build_lifecycle_mux(*, scope: tuple[str, ...] = ()) -> StreamMux: + mux = StreamMux([LifecycleTransformer(scope=scope)], is_async=False) + _arm(mux) + return mux + + +# --------------------------------------------------------------------------- +# LifecycleTransformer +# --------------------------------------------------------------------------- + + +def test_started_emitted_on_first_direct_child_task() -> None: + mux = _build_lifecycle_mux() + mux.push(_tasks_start(["agent:abc123"], task_id="t1", name="tool")) + + [payload] = _drain_lifecycle(mux) + assert payload["event"] == "started" + assert payload["namespace"] == ["agent:abc123"] + assert payload["graph_name"] == "agent" + assert payload["trigger_call_id"] == "abc123" + + +def test_started_dedup_on_repeat_namespace() -> None: + mux = _build_lifecycle_mux() + mux.push(_tasks_start(["agent:abc"], task_id="t1", name="a")) + mux.push(_tasks_start(["agent:abc"], task_id="t2", name="b")) + + payloads = _drain_lifecycle(mux) + assert [p["event"] for p in payloads] == ["started"] + + +def test_grandchild_namespace_discovered() -> None: + """Subgraphs at any depth below scope are tracked, not just direct children.""" + mux = _build_lifecycle_mux() + # First-seen task at length-2 ns means a 2nd-level subgraph started. + mux.push(_tasks_start(["agent:abc", "tool:def"], task_id="t1", name="x")) + + [payload] = _drain_lifecycle(mux) + assert payload["event"] == "started" + assert payload["namespace"] == ["agent:abc", "tool:def"] + + +def test_nested_chain_emits_started_at_each_depth() -> None: + """A graph → subgraph → subgraph chain produces a started event per level.""" + mux = _build_lifecycle_mux() + # Subgraph1 starts emitting tasks (events tagged with its own ns). + mux.push(_tasks_start(["agent:abc"], task_id="t1", name="tool")) + # Subgraph1 invokes subgraph2; subgraph2's first task event arrives. + mux.push(_tasks_start(["agent:abc", "tool:def"], task_id="t2", name="deep")) + + payloads = _drain_lifecycle(mux) + assert [p["namespace"] for p in payloads] == [ + ["agent:abc"], + ["agent:abc", "tool:def"], + ] + assert all(p["event"] == "started" for p in payloads) + + +def test_nested_chain_emits_completed_at_each_depth() -> None: + """Each subgraph in a nested chain closes when its parent task result arrives.""" + mux = _build_lifecycle_mux() + mux.push(_tasks_start(["agent:abc"], task_id="t1", name="tool")) + mux.push(_tasks_start(["agent:abc", "tool:def"], task_id="t2", name="deep")) + + # Subgraph2's owning task (id=def, inside subgraph1) finishes. + mux.push(_tasks_result(["agent:abc"], task_id="def", name="tool")) + # Subgraph1's owning task (id=abc, at root) finishes. + mux.push(_tasks_result([], task_id="abc", name="agent")) + + payloads = _drain_lifecycle(mux) + events = [(p["event"], p["namespace"]) for p in payloads] + assert events == [ + ("started", ["agent:abc"]), + ("started", ["agent:abc", "tool:def"]), + ("completed", ["agent:abc", "tool:def"]), + ("completed", ["agent:abc"]), + ] + + +def test_completed_on_parent_task_result() -> None: + mux = _build_lifecycle_mux() + mux.push(_tasks_start(["agent:abc"], task_id="t1", name="tool")) + mux.push(_tasks_result([], task_id="abc", name="agent")) + + events = [p["event"] for p in _drain_lifecycle(mux)] + assert events == ["started", "completed"] + + +def test_failed_on_parent_task_result_with_error() -> None: + mux = _build_lifecycle_mux() + mux.push(_tasks_start(["agent:abc"], task_id="t1", name="tool")) + mux.push(_tasks_result([], task_id="abc", name="agent", error="boom")) + + payloads = _drain_lifecycle(mux) + assert [p["event"] for p in payloads] == ["started", "failed"] + assert payloads[1]["error"] == "boom" + + +def test_interrupted_on_parent_task_result_with_interrupts() -> None: + mux = _build_lifecycle_mux() + mux.push(_tasks_start(["agent:abc"], task_id="t1", name="tool")) + mux.push( + _tasks_result( + [], + task_id="abc", + name="agent", + interrupts=[{"value": "pause"}], + ) + ) + + payloads = _drain_lifecycle(mux) + assert [p["event"] for p in payloads] == ["started", "interrupted"] + + +def test_interrupt_takes_precedence_over_error() -> None: + mux = _build_lifecycle_mux() + mux.push(_tasks_start(["agent:abc"], task_id="t1", name="tool")) + mux.push( + _tasks_result( + [], + task_id="abc", + name="agent", + error="should-be-suppressed", + interrupts=[{"value": "pause"}], + ) + ) + + last = _drain_lifecycle(mux)[-1] + assert last["event"] == "interrupted" + assert "error" not in last + + +def test_finalize_completes_open_subgraphs() -> None: + mux = _build_lifecycle_mux() + mux.push(_tasks_start(["agent:abc"], task_id="t1", name="tool")) + + mux.close() + payloads = _drain_lifecycle(mux) + assert [p["event"] for p in payloads] == ["started", "completed"] + + +def test_fail_emits_interrupted_for_graph_interrupt() -> None: + mux = _build_lifecycle_mux() + mux.push(_tasks_start(["agent:abc"], task_id="t1", name="tool")) + + mux.fail(GraphInterrupt()) + payloads = _drain_lifecycle(mux) + assert [p["event"] for p in payloads] == ["started", "interrupted"] + assert "error" not in payloads[1] + + +def test_fail_emits_failed_for_other_exceptions() -> None: + mux = _build_lifecycle_mux() + mux.push(_tasks_start(["agent:abc"], task_id="t1", name="tool")) + + mux.fail(RuntimeError("boom")) + payloads = _drain_lifecycle(mux) + assert [p["event"] for p in payloads] == ["started", "failed"] + assert payloads[1]["error"] == "boom" + + +def test_unrelated_methods_pass_through() -> None: + """Non-`tasks` events are not consumed and don't emit lifecycle.""" + mux = _build_lifecycle_mux() + mux.push( + { + "type": "event", + "method": "values", + "params": {"namespace": ["agent:abc"], "timestamp": TS, "data": {}}, + } + ) + assert _drain_lifecycle(mux) == [] + + +def test_scoped_transformer_filters_outside_scope_but_tracks_all_depths() -> None: + """Scope filters the prefix; subgraphs at any depth below scope are tracked.""" + mux = _build_lifecycle_mux(scope=("agent:abc",)) + # Root-level task — out of scope (no shared prefix). + mux.push(_tasks_start(["other:1"], task_id="t1", name="other")) + # Direct child of agent:abc — in scope. + mux.push(_tasks_start(["agent:abc", "tool:def"], task_id="t2", name="tool")) + # Grandchild of agent:abc — also in scope, tracked at its own depth. + mux.push( + _tasks_start(["agent:abc", "tool:def", "deep:ghi"], task_id="t3", name="deep") + ) + + payloads = _drain_lifecycle(mux) + assert [p["namespace"] for p in payloads] == [ + ["agent:abc", "tool:def"], + ["agent:abc", "tool:def", "deep:ghi"], + ] + + +def test_required_stream_modes_declared() -> None: + assert LifecycleTransformer.required_stream_modes == ("tasks",) + + +def test_protocol_event_method_is_native() -> None: + """Native transformer — auto-forwarded events use `lifecycle`, not `custom:lifecycle`.""" + mux = _build_lifecycle_mux() + mux.push(_tasks_start(["agent:abc"], task_id="t1", name="tool")) + + methods = {evt["method"] for evt in mux._events._items} + assert "lifecycle" in methods + assert "custom:lifecycle" not in methods + + +def test_tasks_events_suppressed_from_main_log() -> None: + """Tasks events are folded into lifecycle and don't appear on the main log.""" + mux = _build_lifecycle_mux() + mux.push(_tasks_start(["agent:abc"], task_id="t1", name="tool")) + mux.push(_tasks_result([], task_id="abc", name="agent")) + + methods = [evt["method"] for evt in mux._events._items] + assert "tasks" not in methods + # Lifecycle events did make it through, though. + assert "lifecycle" in methods + + +# --------------------------------------------------------------------------- +# End-to-end: real graphs through stream_v2 +# --------------------------------------------------------------------------- + + +class _State(TypedDict): + value: str + items: Annotated[list[str], operator.add] + + +def _passthrough(state: _State) -> dict[str, Any]: + return {"value": state["value"] + "!", "items": ["x"]} + + +def _make_two_level_nested() -> Any: + """Build outer → middle → inner. Three Pregel instances, two nesting levels.""" + inner_b: StateGraph = StateGraph(_State, input_schema=_State) + inner_b.add_node("inner_node", _passthrough) + inner_b.add_edge(START, "inner_node") + inner_b.add_edge("inner_node", END) + inner = inner_b.compile() + + middle_b: StateGraph = StateGraph(_State, input_schema=_State) + middle_b.add_node("inner", inner) + middle_b.add_edge(START, "inner") + middle_b.add_edge("inner", END) + middle = middle_b.compile() + + outer_b: StateGraph = StateGraph(_State, input_schema=_State) + outer_b.add_node("middle", middle) + outer_b.add_edge(START, "middle") + outer_b.add_edge("middle", END) + return outer_b.compile() + + +def test_stream_v2_real_graph_emits_lifecycle_at_each_depth() -> None: + """Outer graph with two nested subgraphs surfaces lifecycle for both.""" + graph = _make_two_level_nested() + run = graph.stream_v2({"value": "x", "items": []}) + + # Iterating the projection drives the pump and drains synthesized + # lifecycle events at the same time. + payloads = list(run.lifecycle) + # Each subgraph instance produces a started + a terminal event. Two + # nested instances, so four payloads total in some interleaving. + by_event = {p["event"] for p in payloads} + assert "started" in by_event + assert "completed" in by_event + # Two distinct namespaces — direct child of root, and grandchild. + namespaces = {tuple(p["namespace"]) for p in payloads} + direct_children = {ns for ns in namespaces if len(ns) == 1} + grandchildren = {ns for ns in namespaces if len(ns) == 2} + assert direct_children, f"expected a level-1 lifecycle namespace, got {namespaces}" + assert grandchildren, f"expected a level-2 lifecycle namespace, got {namespaces}" + # Every direct-child namespace has a matching grandchild whose path extends it. + for parent in direct_children: + assert any(gc[: len(parent)] == parent for gc in grandchildren), ( + f"grandchild does not extend parent {parent}: {grandchildren}" + ) + + +def test_stream_v2_with_nested_parent_ns_scopes_lifecycle() -> None: + """When `stream_v2` is called with a non-empty checkpoint_ns in config, + `_resolve_parent_ns` returns that namespace and the registered + `LifecycleTransformer` is constructed with `scope=parent_ns`. This + exercises the path that exists today purely for nested-stream_v2 + callers; the test simulates such a caller by injecting a + checkpoint_ns into the config. + """ + graph = _make_two_level_nested() + config = {CONF: {CONFIG_KEY_CHECKPOINT_NS: "outer:abc"}} + run = graph.stream_v2({"value": "x", "items": []}, config=config) + + payloads = list(run.lifecycle) + # Every emitted lifecycle namespace must extend the caller's scope — + # nothing at root-level, nothing under a sibling prefix. + for p in payloads: + ns = tuple(p["namespace"]) + assert ns[:1] == ("outer:abc",), ( + f"namespace {ns} not within scoped prefix ('outer:abc',)" + ) diff --git a/libs/langgraph/tests/test_stream_messages_transformer.py b/libs/langgraph/tests/test_stream_messages_transformer.py new file mode 100644 index 000000000..b2f4f9c65 --- /dev/null +++ b/libs/langgraph/tests/test_stream_messages_transformer.py @@ -0,0 +1,872 @@ +"""Tests for MessagesTransformer: protocol event routing, whole-message fallback, +legacy v1 chunk filtering, and end-to-end via stream_v2 / astream_v2.""" + +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 langchain_core.runnables import RunnableConfig +from typing_extensions import TypedDict + +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).""" + return { + "type": "event", + "method": "messages", + "params": { + "namespace": [], + "timestamp": TS, + "data": (event, {"langgraph_node": node, "run_id": run_id}), + }, + } + + +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] = {"finish_reason": "stop"} if finish else {} + return { + "type": "event", + "method": "messages", + "params": { + "namespace": [], + "timestamp": TS, + "data": ( + AIMessageChunk(content=text, id=msg_id, response_metadata=rm), + {"langgraph_node": node}, + ), + }, + } + + +def _whole_msg( + text: str, + msg_id: str = "msg-10", + *, + node: str = "node", +) -> dict[str, Any]: + """Build a messages ProtocolEvent carrying a completed AIMessage.""" + return { + "type": "event", + "method": "messages", + "params": { + "namespace": [], + "timestamp": TS, + "data": (AIMessage(content=text, id=msg_id), {"langgraph_node": node}), + }, + } + + +def _make_sync_transformer() -> tuple[MessagesTransformer, EventLog[ChatModelStream]]: + t = MessagesTransformer() + log: EventLog[ChatModelStream] = t.init()["messages"] + log._bind(is_async=False) + # Subscribe up front so pushes during process() are retained. + log._subscribed = True + t._bind_pump(lambda: False) + return t, log + + +def _make_async_transformer() -> tuple[MessagesTransformer, EventLog[ChatModelStream]]: + t = MessagesTransformer() + log: EventLog[ChatModelStream] = t.init()["messages"] + log._bind(is_async=True) + log._subscribed = True + return t, log + + +def _lifecycle( + *, text: str = "hello world", message_id: str = "run-1" +) -> list[dict[str, Any]]: + """Produce a valid protocol event lifecycle: start, delta, finish.""" + 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"}, + ] + + +def _simple_graph(): + def call_model(state: MessagesState) -> dict[str, Any]: + model = GenericFakeChatModel(messages=iter(["hello world"])) + stream = model.stream_v2(state["messages"]) + return {"messages": stream.output} + + return ( + StateGraph(MessagesState) + .add_node("call_model", call_model) + .add_edge(START, "call_model") + .add_edge("call_model", END) + .compile() + ) + + +# --------------------------------------------------------------------------- +# 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", + ) + ) + log.close() + (stream,) = list(log._items) + assert isinstance(stream, ChatModelStream) + assert stream.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: + 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: + t, log = _make_sync_transformer() + 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) + assert "".join(stream._text_proj._deltas) == "abcdef" + + def test_stream_pushed_on_message_start_not_finish(self) -> None: + # Consumer can see the stream before message-finish arrives. + t, log = _make_sync_transformer() + t.process( + _proto_event( + {"event": "message-start", "role": "ai", "message_id": "run-1"}, + run_id="run-1", + ) + ) + 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,) = list(log._items) + assert stream.node == "my_llm" + + +# --------------------------------------------------------------------------- +# Whole-message 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) + assert [e["event"] for e in stream._events] == [ + "message-start", + "content-block-start", + "content-block-delta", + "content-block-finish", + "message-finish", + ] + + +# --------------------------------------------------------------------------- +# Filtering +# --------------------------------------------------------------------------- + + +class TestFiltering: + def test_non_messages_events_pass_through(self) -> None: + t, _ = _make_sync_transformer() + assert ( + t.process( + { + "type": "event", + "method": "values", + "params": {"namespace": [], "timestamp": TS, "data": {"x": 1}}, + } + ) + is True + ) + + def test_subgraph_namespace_dropped(self) -> None: + t, log = _make_sync_transformer() + t.process( + { + "type": "event", + "method": "messages", + "params": { + "namespace": ["subgraph"], + "timestamp": TS, + "data": ( + {"event": "message-start", "message_id": "run-x"}, + {"run_id": "run-x"}, + ), + }, + } + ) + log.close() + assert list(log._items) == [] + + def test_legacy_v1_chunks_ignored(self) -> None: + # v1 AIMessageChunk tuples (from on_llm_new_token) are not streamed + # into this projection; callers must migrate to stream_v2. + t, log = _make_sync_transformer() + t.process(_v1_chunk("hello")) + t.process(_v1_chunk(" world", finish=True)) + log.close() + assert list(log._items) == [] + + +# --------------------------------------------------------------------------- +# Lifecycle: fail / finalize +# --------------------------------------------------------------------------- + + +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 +# --------------------------------------------------------------------------- + + +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)) + assert isinstance(list(log._items)[0], AsyncChatModelStream) + + @pytest.mark.anyio + async def test_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) + assert "".join([d async for d in stream.text]) == "hello world" + + @pytest.mark.anyio + async def test_output_awaitable(self) -> None: + t, log = _make_async_transformer() + for evt in _lifecycle(text="async"): + t.process(_proto_event(evt)) + (stream,) = list(log._items) + assert (await stream.output).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, values_t) + assert messages_t._pump_fn is not None + 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, values_t) + + log: EventLog[ChatModelStream] = mux.extensions["messages"] + log._subscribed = True + for evt in _lifecycle(): + messages_t.process(_proto_event(evt)) + + (stream,) = list(log._items) + assert stream._request_more is messages_t._pump_fn + + +# --------------------------------------------------------------------------- +# End-to-end via StreamMux +# --------------------------------------------------------------------------- + + +class TestViaMux: + def _make_mux( + self, + ) -> tuple[MessagesTransformer, StreamMux, EventLog[ChatModelStream]]: + 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 + return t, mux, log + + def test_streaming_via_mux(self) -> None: + t, mux, log = self._make_mux() + 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, mux, log = self._make_mux() + 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)) + + (stream,) = list(log._items) + assert (await stream.output).text == "async mux" + await mux.aclose() + + +# --------------------------------------------------------------------------- +# End-to-end: graph → stream_v2 → run.messages (node calls stream_v2) +# --------------------------------------------------------------------------- + + +class TestEndToEnd: + """stream_v2 path: node calls model.stream_v2() explicitly.""" + + 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"}) + (stream,) = list(run.messages) + assert isinstance(stream, ChatModelStream) + assert stream.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"}) + (stream,) = list(run.messages) + assert "".join(stream.text) == "streamed answer" + + def test_non_llm_message_returned_from_node(self) -> None: + """Whole-message fallback: node returns a finalized AIMessage directly.""" + + 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"}) + (stream,) = list(run.messages) + assert stream.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"]) + return {"messages": await stream} + + 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 = [s async for s in run.messages] + assert len(streams) == 1 + assert isinstance(streams[0], AsyncChatModelStream) + assert (await streams[0].output).text == "async answer" + + @pytest.mark.anyio + async def test_nested_async_iteration_yields_text_deltas(self) -> None: + """Inner stream.text drives the shared graph pump via the async pump binding.""" + import asyncio + + model = GenericFakeChatModel(messages=iter(["hello world"])) + + async def call_model(state: MessagesState) -> dict[str, Any]: + stream = await model.astream_v2(state["messages"]) + return {"messages": await stream} + + 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() -> list[str]: + collected: list[str] = [] + async for stream in run.messages: + async for delta in stream.text: + collected.append(delta) + return collected + + assert "".join(await asyncio.wait_for(consume(), timeout=2.0)) == "hello world" + + +# --------------------------------------------------------------------------- +# End-to-end: graph → stream_v2 → run.messages (node calls invoke) +# --------------------------------------------------------------------------- + + +class TestEndToEndV2Invoke: + """Auto-routing path: stream_v2 injects CONFIG_KEY_STREAM_MESSAGES_V2, + causing BaseChatModel to drive the v2 protocol event generator even for + model.invoke().""" + + def _graph(self, model): + def call_model(state: MessagesState) -> dict[str, Any]: + return {"messages": model.invoke(state["messages"])} + + return ( + StateGraph(MessagesState) + .add_node("call_model", call_model) + .add_edge(START, "call_model") + .add_edge("call_model", END) + .compile() + ) + + def test_invoke_populates_messages(self) -> None: + run = self._graph( + GenericFakeChatModel(messages=iter(["hello world"])) + ).stream_v2({"messages": "hi"}) + (stream,) = list(run.messages) + assert isinstance(stream, ChatModelStream) + assert stream.output.text == "hello world" + + def test_invoke_emits_protocol_events(self) -> None: + """Iterating the stream yields the full v2 lifecycle, not v1 chunks.""" + run = self._graph( + GenericFakeChatModel(messages=iter(["streamed answer"])) + ).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(self) -> None: + run = self._graph( + GenericFakeChatModel(messages=iter(["delta streaming works"])) + ).stream_v2({"messages": "hi"}) + (stream,) = list(run.messages) + assert "".join(stream.text) == "delta streaming works" + + def test_invoke_two_nodes_two_streams(self) -> None: + 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() + ) + + streams = list(graph.stream_v2({"messages": "hi"}).messages) + assert len(streams) == 2 + assert {s.output.text for s in streams} == {"alpha", "beta"} + + def test_invoke_plus_constructed_message_two_streams(self) -> None: + """Live-streamed node + constructed-message node → two ChatModelStreams.""" + 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_populates_messages(self) -> None: + 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 = [s async for s in run.messages] + assert len(streams) == 1 + assert isinstance(streams[0], AsyncChatModelStream) + assert (await streams[0].output).text == "async invoke" + + +# --------------------------------------------------------------------------- +# Regression: direct stream_mode="messages" must stay v1 +# --------------------------------------------------------------------------- + + +class TestDirectMessagesModeStaysV1: + def test_direct_graph_stream_messages_yields_ai_message_chunks(self) -> None: + """graph.stream(stream_mode="messages") must not leak v2 event dicts — + the v2 flag is only injected by stream_v2 / astream_v2.""" + 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")) + assert parts, "expected stream_mode='messages' to emit tuples" + for payload, _metadata in parts: + assert isinstance(payload, AIMessageChunk) + assert ( + "".join(p[0].content for p in parts if isinstance(p[0].content, str)) + == "legacy path" + ) + + def test_nested_graph_stream_messages_stays_v1_under_outer_stream_v2(self) -> None: + """An outer `stream_v2()` run must not flip an inner direct + `stream_mode="messages"` call onto the v2 event protocol.""" + model = GenericFakeChatModel(messages=iter(["nested legacy path"])) + + def call_model(state: MessagesState) -> dict[str, Any]: + return {"messages": model.invoke(state["messages"])} + + inner = ( + StateGraph(MessagesState) + .add_node("call_model", call_model) + .add_edge(START, "call_model") + .add_edge("call_model", END) + .compile() + ) + + class OuterState(TypedDict, total=False): + saw_only_chunks: bool + first_payload_type: str + text: str + + def call_subgraph(state: OuterState, config: RunnableConfig) -> dict[str, Any]: + parts = list( + inner.stream( + {"messages": "hi"}, + config, + stream_mode="messages", + ) + ) + assert parts + payloads = [payload for payload, _metadata in parts] + return { + "saw_only_chunks": all( + isinstance(payload, AIMessageChunk) for payload in payloads + ), + "first_payload_type": type(payloads[0]).__name__, + "text": "".join( + payload.content + for payload in payloads + if isinstance(payload, AIMessageChunk) + and isinstance(payload.content, str) + ), + } + + outer = ( + StateGraph(OuterState) + .add_node("call_subgraph", call_subgraph) + .add_edge(START, "call_subgraph") + .add_edge("call_subgraph", END) + .compile() + ) + + result = outer.stream_v2({}).output + + assert result is not None + assert result["saw_only_chunks"] is True + assert result["first_payload_type"] == "AIMessageChunk" + assert result["text"] == "nested legacy path" + + +# --------------------------------------------------------------------------- +# StreamMessagesHandlerV2 unit +# --------------------------------------------------------------------------- + + +class TestStreamMessagesHandlerV2Unit: + def test_on_llm_new_token_is_noop(self) -> None: + """v2 handler must not emit v1 chunks even when on_llm_new_token fires.""" + 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() + 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 == [] + + def test_on_llm_end_dedupes_when_final_message_id_differs(self) -> None: + """A streamed v2 message should not be emitted again from the final + AIMessage fallback when its final id does not match `message-start`.""" + from uuid import uuid4 + + from langchain_core.outputs import ChatGeneration, LLMResult + + from langgraph.pregel._messages import StreamMessagesHandlerV2 + + emitted: list[Any] = [] + handler = StreamMessagesHandlerV2(emitted.append, subgraphs=False) + run_id = uuid4() + handler.metadata[run_id] = ((), {"langgraph_node": "x"}) + + handler.on_stream_event( + {"event": "message-start", "message_id": "stream-msg-1"}, + run_id=run_id, + ) + handler.on_llm_end( + LLMResult( + generations=[ + [ + ChatGeneration( + message=AIMessage(content="hello", id="final-msg-1") + ) + ] + ] + ), + run_id=run_id, + ) + + assert len(emitted) == 1 diff --git a/libs/langgraph/tests/test_stream_subgraph_transformer.py b/libs/langgraph/tests/test_stream_subgraph_transformer.py new file mode 100644 index 000000000..71b161243 --- /dev/null +++ b/libs/langgraph/tests/test_stream_subgraph_transformer.py @@ -0,0 +1,865 @@ +"""Tests for SubgraphTransformer. + +Subscribes to `tasks` events and produces in-process `SubgraphRunStream` +handles backed by mini-muxes (built via `StreamMux._make_child`). The +synthetic-event tests isolate the inference / mini-mux wiring; the +real-graph tests exercise the end-to-end navigation path through +`stream_v2`. +""" + +from __future__ import annotations + +import operator +import time +from collections.abc import AsyncIterator +from functools import partial +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.pregel.main import _normalize_stream_transformer_factories +from langgraph.stream._mux import StreamMux +from langgraph.stream._types import ProtocolEvent, StreamTransformer +from langgraph.stream.run_stream import ( + AsyncGraphRunStream, + AsyncSubgraphRunStream, + GraphRunStream, + SubgraphRunStream, +) +from langgraph.stream.transformers import ( + LifecycleTransformer, + MessagesTransformer, + SubgraphTransformer, + ValuesTransformer, +) + +TS = int(time.time() * 1000) + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +def _tasks_start( + namespace: list[str], + *, + task_id: str, + name: str, +) -> dict[str, Any]: + return { + "type": "event", + "method": "tasks", + "params": { + "namespace": namespace, + "timestamp": TS, + "data": { + "id": task_id, + "name": name, + "input": None, + "triggers": [], + }, + }, + } + + +def _tasks_result( + namespace: list[str], + *, + task_id: str, + name: str, + error: str | None = None, + interrupts: list[dict[str, Any]] | None = None, +) -> dict[str, Any]: + return { + "type": "event", + "method": "tasks", + "params": { + "namespace": namespace, + "timestamp": TS, + "data": { + "id": task_id, + "name": name, + "error": error, + "interrupts": interrupts or [], + "result": {}, + }, + }, + } + + +def _native_factories() -> list[Any]: + """Mirror the factory list `Pregel.stream_v2` registers.""" + return [ + ValuesTransformer, + MessagesTransformer, + LifecycleTransformer, + SubgraphTransformer, + ] + + +def _stream_part( + method: str, + namespace: tuple[str, ...], + data: Any, +) -> dict[str, Any]: + return {"type": method, "ns": namespace, "data": data} + + +async def _astream_parts(*parts: dict[str, Any]) -> AsyncIterator[dict[str, Any]]: + for part in parts: + yield part + + +def _arm(mux: StreamMux) -> None: + """Pre-subscribe every projection in the mux so synthetic pushes accumulate. + + Real consumer code subscribes by iterating the projection; tests + inspect `_items` directly, so the lazy-subscribe gate has to be + flipped manually before any synthetic events are pushed. + """ + mux._events._subscribed = True + for value in mux.extensions.values(): + if hasattr(value, "_subscribed"): # EventLog + value._subscribed = True + elif hasattr(value, "_log"): # StreamChannel + value._log._subscribed = True + + +def _arm_recursive(mux: StreamMux) -> None: + """Arm `mux` and every mini-mux currently held by SubgraphTransformer handles. + + Mini-muxes are created during `mux.push(...)` when a new direct + child is discovered. Tests must call this after each push that + might have created a new mini-mux so subsequent pushes' projection + side effects accumulate (rather than dropping silently against an + unsubscribed log). + """ + _arm(mux) + for handle in _subgraph_transformer(mux)._handles.values(): + if handle._mux is not None: + _arm_recursive(handle._mux) + + +def _build_root_mux(*, scope: tuple[str, ...] = ()) -> StreamMux: + mux = StreamMux( + factories=_native_factories(), + scope=scope, + is_async=False, + ) + _arm(mux) + return mux + + +def _subgraph_transformer(mux: StreamMux) -> SubgraphTransformer: + transformer = mux.transformer_by_key("subgraphs") + assert isinstance(transformer, SubgraphTransformer) + return transformer + + +def _drain_subgraphs(mux: StreamMux) -> list[SubgraphRunStream]: + return list(_subgraph_transformer(mux)._log._items) + + +def _child_mux(handle: SubgraphRunStream | AsyncSubgraphRunStream) -> StreamMux: + assert handle._mux is not None + return handle._mux + + +def _event_items(mux: StreamMux) -> list[ProtocolEvent]: + return list(mux._events._items) + + +def _lifecycle_payloads(mux: StreamMux) -> list[dict[str, Any]]: + lifecycle_t = mux.transformer_by_key("lifecycle") + assert isinstance(lifecycle_t, LifecycleTransformer) + return list(lifecycle_t._channel._log._items) + + +# --------------------------------------------------------------------------- +# Synthetic-event tests +# --------------------------------------------------------------------------- + + +def test_handle_created_on_first_direct_child_task() -> None: + mux = _build_root_mux() + mux.push(_tasks_start(["agent:abc"], task_id="t1", name="tool")) + + [handle] = _drain_subgraphs(mux) + assert handle.path == ("agent:abc",) + assert handle.graph_name == "agent" + assert handle.trigger_call_id == "abc" + assert handle.status == "started" + _child_mux(handle) # mini-mux backed + + +def test_handle_status_completes_on_parent_result() -> None: + mux = _build_root_mux() + mux.push(_tasks_start(["agent:abc"], task_id="t1", name="tool")) + mux.push(_tasks_result([], task_id="abc", name="agent")) + + [handle] = _drain_subgraphs(mux) + assert handle.status == "completed" + assert handle.error is None + + +def test_handle_status_failed_with_error() -> None: + mux = _build_root_mux() + mux.push(_tasks_start(["agent:abc"], task_id="t1", name="tool")) + mux.push(_tasks_result([], task_id="abc", name="agent", error="boom")) + + [handle] = _drain_subgraphs(mux) + assert handle.status == "failed" + assert handle.error == "boom" + + +def test_handle_status_interrupted() -> None: + mux = _build_root_mux() + mux.push(_tasks_start(["agent:abc"], task_id="t1", name="tool")) + mux.push( + _tasks_result( + [], + task_id="abc", + name="agent", + interrupts=[{"value": "pause"}], + ) + ) + + [handle] = _drain_subgraphs(mux) + assert handle.status == "interrupted" + + +def test_grandchild_discovered_via_child_mini_mux() -> None: + """Each mini-mux owns its own scope; grandchildren live on the child handle.""" + mux = _build_root_mux() + # Direct child started — creates the mini-mux. + mux.push(_tasks_start(["agent:abc"], task_id="t1", name="tool")) + # Pre-subscribe the freshly-created mini-mux so subsequent + # forwarded events land on its projections (consumer would + # subscribe naturally by iterating handle.subgraphs, but the + # test inspects `_items` directly). + _arm_recursive(mux) + # Grandchild's first task event flows down into the child mini-mux. + mux.push(_tasks_start(["agent:abc", "tool:def"], task_id="t2", name="deep")) + + [child_handle] = _drain_subgraphs(mux) + assert child_handle.path == ("agent:abc",) + # The grandchild appears on the CHILD'S subgraphs projection. + grandchildren = list(child_handle.subgraphs._items) + assert len(grandchildren) == 1 + assert grandchildren[0].path == ("agent:abc", "tool:def") + + +def test_finalize_completes_open_handles() -> None: + mux = _build_root_mux() + mux.push(_tasks_start(["agent:abc"], task_id="t1", name="tool")) + + mux.close() + [handle] = _drain_subgraphs(mux) + assert handle.status == "completed" + + +def test_fail_marks_open_handles_interrupted_for_graph_interrupt() -> None: + mux = _build_root_mux() + mux.push(_tasks_start(["agent:abc"], task_id="t1", name="tool")) + + mux.fail(GraphInterrupt()) + [handle] = _drain_subgraphs(mux) + assert handle.status == "interrupted" + + +def test_fail_marks_open_handles_failed_for_other_errors() -> None: + mux = _build_root_mux() + mux.push(_tasks_start(["agent:abc"], task_id="t1", name="tool")) + + mux.fail(RuntimeError("boom")) + [handle] = _drain_subgraphs(mux) + assert handle.status == "failed" + assert handle.error == "boom" + + +def test_child_mux_requires_factories() -> None: + """A mux constructed only from `transformers=` can't clone factories.""" + transformer = SubgraphTransformer() + mux = StreamMux(transformers=[transformer], is_async=False) + with pytest.raises(RuntimeError, match="factories"): + mux._make_child(("anything",)) + + +def test_subgraph_and_lifecycle_agree_on_terminal_status() -> None: + """Both transformers consume the same tasks signal — no drift.""" + mux = _build_root_mux() + mux.push(_tasks_start(["agent:abc"], task_id="t1", name="tool")) + mux.push(_tasks_result([], task_id="abc", name="agent", error="boom")) + + [handle] = _drain_subgraphs(mux) + payloads = _lifecycle_payloads(mux) + assert handle.status == "failed" + assert payloads[-1]["event"] == "failed" + assert handle.error == payloads[-1]["error"] + + +def test_required_stream_modes_declared() -> None: + assert SubgraphTransformer.required_stream_modes == ("tasks",) + + +def test_tasks_events_suppressed_from_main_log() -> None: + """Tasks events are folded into discovery and don't appear on the main log.""" + mux = _build_root_mux() + mux.push(_tasks_start(["agent:abc"], task_id="t1", name="tool")) + mux.push(_tasks_result([], task_id="abc", name="agent")) + + methods = [evt["method"] for evt in _event_items(mux)] + assert "tasks" not in methods + + +class _ChildEventObserver(StreamTransformer): + """Records child-scope event identity without mutating it.""" + + records: list[tuple[tuple[str, ...], int, int, bool]] = [] + + def init(self) -> dict[str, Any]: + return {} + + def process(self, event: ProtocolEvent) -> bool: + if self.scope and event["method"] == "values": + self.records.append( + ( + self.scope, + id(event), + id(event["params"]["data"]), + "seq" in event, + ) + ) + return True + + +def test_child_forwarding_reuses_event_without_assigning_seq() -> None: + _ChildEventObserver.records = [] + mux = StreamMux( + factories=[ + ValuesTransformer, + MessagesTransformer, + LifecycleTransformer, + SubgraphTransformer, + _ChildEventObserver, + ], + is_async=False, + ) + _arm(mux) + mux.push(_tasks_start(["agent:abc"], task_id="t1", name="tool")) + data = {"x": 1} + event: ProtocolEvent = { + "type": "event", + "method": "values", + "params": { + "namespace": ["agent:abc"], + "timestamp": TS, + "data": data, + }, + } + mux.push(event) + + assert _ChildEventObserver.records == [(("agent:abc",), id(event), id(data), False)] + [root_event] = [evt for evt in _event_items(mux) if evt["method"] == "values"] + assert root_event is event + assert "seq" in root_event + + +class _AsyncProbeTransformer(StreamTransformer): + """Async-only transformer used to verify mini-mux async dispatch.""" + + required_stream_modes = ("tasks",) + + def __init__(self, scope: tuple[str, ...] = ()) -> None: + super().__init__(scope) + self.seen: list[tuple[str, ...]] = [] + self.finalized = False + self.failed: BaseException | None = None + + def init(self) -> dict[str, Any]: + return {"async_probe": self} + + async def aprocess(self, event: ProtocolEvent) -> bool: + self.seen.append(tuple(event["params"]["namespace"])) + return True + + async def afinalize(self) -> None: + self.finalized = True + + async def afail(self, err: BaseException) -> None: + self.failed = err + + +@pytest.mark.anyio +async def test_async_child_mini_mux_uses_async_lane() -> None: + mux = StreamMux( + factories=[ + ValuesTransformer, + MessagesTransformer, + LifecycleTransformer, + SubgraphTransformer, + _AsyncProbeTransformer, + ], + is_async=True, + ) + await mux.apush(_tasks_start(["agent:abc"], task_id="t1", name="tool")) + + handle = _subgraph_transformer(mux)._handles[("agent:abc",)] + assert isinstance(handle, AsyncSubgraphRunStream) + probe = _child_mux(handle).transformer_by_key("async_probe") + assert isinstance(probe, _AsyncProbeTransformer) + assert probe.seen == [("agent:abc",)] + + await mux.apush(_tasks_result([], task_id="abc", name="agent")) + assert probe.finalized is True + + +@pytest.mark.anyio +async def test_async_child_mini_mux_fail_uses_async_lane() -> None: + mux = StreamMux( + factories=[ + ValuesTransformer, + MessagesTransformer, + LifecycleTransformer, + SubgraphTransformer, + _AsyncProbeTransformer, + ], + is_async=True, + ) + await mux.apush(_tasks_start(["agent:abc"], task_id="t1", name="tool")) + + handle = _subgraph_transformer(mux)._handles[("agent:abc",)] + probe = _child_mux(handle).transformer_by_key("async_probe") + assert isinstance(probe, _AsyncProbeTransformer) + + err = RuntimeError("boom") + await mux.afail(err) + assert probe.failed is err + + +class _StandardCtorTransformer(StreamTransformer): + """Transformer class that inherits the standard scoped constructor.""" + + def init(self) -> dict[str, Any]: + return {"standard_ctor": self} + + def process(self, event: ProtocolEvent) -> bool: + return True + + +class _ScopedTransformer(StreamTransformer): + """Transformer class that uses the inherited scoped construction.""" + + def init(self) -> dict[str, Any]: + return {"scoped": self} + + def process(self, event: ProtocolEvent) -> bool: + return True + + +class _ConfigurableFactoryTransformer(StreamTransformer): + """Transformer built by a configured per-scope factory.""" + + def __init__(self, scope: tuple[str, ...] = (), *, label: str) -> None: + super().__init__(scope) + self.label = label + + def init(self) -> dict[str, Any]: + return {"configurable": self} + + def process(self, event: ProtocolEvent) -> bool: + return True + + +class _ChildExploder(StreamTransformer): + """Raise from child mini-muxes to verify errors propagate upstream.""" + + def init(self) -> dict[str, Any]: + return {} + + def process(self, event: ProtocolEvent) -> bool: + if self.scope and event["method"] == "values": + raise RuntimeError("child boom") + return True + + +class _ChildFinalizeExploder(StreamTransformer): + """Raise from child mini-mux finalization.""" + + supports_sync = True + + def init(self) -> dict[str, Any]: + return {} + + def process(self, event: ProtocolEvent) -> bool: + return True + + def finalize(self) -> None: + if self.scope: + raise RuntimeError("child finalize boom") + + async def afinalize(self) -> None: + if self.scope: + raise RuntimeError("child afinalize boom") + + +def test_normalize_transformer_factories_supports_scoped_classes() -> None: + factories = _normalize_stream_transformer_factories( + [_StandardCtorTransformer, _ScopedTransformer] + ) + + standard_ctor = factories[0](("child",)) + scoped = factories[1](("child",)) + assert isinstance(standard_ctor, _StandardCtorTransformer) + assert standard_ctor.scope == ("child",) + assert isinstance(scoped, _ScopedTransformer) + assert scoped.scope == ("child",) + + +def test_normalize_transformer_factories_supports_configured_factories() -> None: + factories = _normalize_stream_transformer_factories( + [partial(_ConfigurableFactoryTransformer, label="configured")] + ) + + built = factories[0](("child",)) + assert isinstance(built, _ConfigurableFactoryTransformer) + assert built.label == "configured" + assert built.scope == ("child",) + + +def test_normalize_transformer_factories_rejects_instances() -> None: + with pytest.raises(TypeError, match="pre-built instance"): + _normalize_stream_transformer_factories([_StandardCtorTransformer()]) + + +def test_child_forwarding_errors_fail_sync_run() -> None: + mux = StreamMux( + factories=[ + ValuesTransformer, + MessagesTransformer, + LifecycleTransformer, + SubgraphTransformer, + _ChildExploder, + ], + is_async=False, + ) + values_t = mux.transformer_by_key("values") + assert isinstance(values_t, ValuesTransformer) + run = GraphRunStream( + iter( + [ + _stream_part( + "tasks", + ("agent:abc",), + { + "id": "t1", + "name": "tool", + "input": None, + "triggers": [], + }, + ), + _stream_part("values", ("agent:abc",), {"x": 1}), + ] + ), + mux, + values_t, + ) + + handle = next(iter(run.subgraphs)) + assert handle.path == ("agent:abc",) + with pytest.raises(RuntimeError, match="child boom"): + _ = run.output + assert run._mux._events._error is not None + + +@pytest.mark.anyio +async def test_child_forwarding_errors_fail_async_run() -> None: + mux = StreamMux( + factories=[ + ValuesTransformer, + MessagesTransformer, + LifecycleTransformer, + SubgraphTransformer, + _ChildExploder, + ], + is_async=True, + ) + values_t = mux.transformer_by_key("values") + assert isinstance(values_t, ValuesTransformer) + run = AsyncGraphRunStream( + _astream_parts( + _stream_part( + "tasks", + ("agent:abc",), + { + "id": "t1", + "name": "tool", + "input": None, + "triggers": [], + }, + ), + _stream_part("values", ("agent:abc",), {"x": 1}), + ), + mux, + values_t, + ) + + handle = await run.subgraphs.__aiter__().__anext__() + assert handle.path == ("agent:abc",) + with pytest.raises(RuntimeError, match="child boom"): + await run.output() + assert run._mux._events._error is not None + + +def test_child_finalize_errors_propagate_to_sync_run() -> None: + mux = StreamMux( + factories=[ + ValuesTransformer, + MessagesTransformer, + LifecycleTransformer, + SubgraphTransformer, + _ChildFinalizeExploder, + ], + is_async=False, + ) + values_t = mux.transformer_by_key("values") + assert isinstance(values_t, ValuesTransformer) + run = GraphRunStream( + iter( + [ + _stream_part( + "tasks", + ("agent:abc",), + { + "id": "t1", + "name": "tool", + "input": None, + "triggers": [], + }, + ) + ] + ), + mux, + values_t, + ) + + with pytest.raises(RuntimeError, match="child finalize boom"): + _ = run.output + + +@pytest.mark.anyio +async def test_child_finalize_errors_propagate_to_async_run() -> None: + mux = StreamMux( + factories=[ + ValuesTransformer, + MessagesTransformer, + LifecycleTransformer, + SubgraphTransformer, + _ChildFinalizeExploder, + ], + is_async=True, + ) + values_t = mux.transformer_by_key("values") + assert isinstance(values_t, ValuesTransformer) + run = AsyncGraphRunStream( + _astream_parts( + _stream_part( + "tasks", + ("agent:abc",), + { + "id": "t1", + "name": "tool", + "input": None, + "triggers": [], + }, + ) + ), + mux, + values_t, + ) + + with pytest.raises(RuntimeError, match="child afinalize boom"): + await run.output() + + +# --------------------------------------------------------------------------- +# End-to-end real-graph tests +# --------------------------------------------------------------------------- + + +class _State(TypedDict): + value: str + items: Annotated[list[str], operator.add] + + +def _passthrough(state: _State) -> dict[str, Any]: + return {"value": state["value"] + "!", "items": ["x"]} + + +def _make_two_level_nested() -> Any: + """outer → middle → inner. Three Pregel instances, two nesting levels.""" + inner_b: StateGraph = StateGraph(_State, input_schema=_State) + inner_b.add_node("inner_node", _passthrough) + inner_b.add_edge(START, "inner_node") + inner_b.add_edge("inner_node", END) + inner = inner_b.compile() + + middle_b: StateGraph = StateGraph(_State, input_schema=_State) + middle_b.add_node("inner", inner) + middle_b.add_edge(START, "inner") + middle_b.add_edge("inner", END) + middle = middle_b.compile() + + outer_b: StateGraph = StateGraph(_State, input_schema=_State) + outer_b.add_node("middle", middle) + outer_b.add_edge(START, "middle") + outer_b.add_edge("middle", END) + return outer_b.compile() + + +def _item_node(item: str): + def node(state: _State) -> dict[str, Any]: + return {"items": [item]} + + return node + + +def _make_two_sibling_subgraphs() -> Any: + """outer → one → two, where both nodes are compiled subgraphs.""" + one_b: StateGraph = StateGraph(_State, input_schema=_State) + one_b.add_node("add_one", _item_node("one")) + one_b.add_edge(START, "add_one") + one_b.add_edge("add_one", END) + one = one_b.compile() + + two_b: StateGraph = StateGraph(_State, input_schema=_State) + two_b.add_node("add_two", _item_node("two")) + two_b.add_edge(START, "add_two") + two_b.add_edge("add_two", END) + two = two_b.compile() + + outer_b: StateGraph = StateGraph(_State, input_schema=_State) + outer_b.add_node("one", one) + outer_b.add_node("two", two) + outer_b.add_edge(START, "one") + outer_b.add_edge("one", "two") + outer_b.add_edge("two", END) + return outer_b.compile() + + +def _failing_node(state: _State) -> dict[str, Any]: + raise ValueError("child boom") + + +def _make_failing_nested() -> Any: + inner_b: StateGraph = StateGraph(_State, input_schema=_State) + inner_b.add_node("fail", _failing_node) + inner_b.add_edge(START, "fail") + inner_b.add_edge("fail", END) + inner = inner_b.compile() + + outer_b: StateGraph = StateGraph(_State, input_schema=_State) + outer_b.add_node("inner", inner) + outer_b.add_edge(START, "inner") + outer_b.add_edge("inner", END) + return outer_b.compile() + + +def test_stream_v2_real_graph_yields_subgraph_handles() -> None: + """Iterating `run.subgraphs` yields handles for direct-child subgraphs.""" + graph = _make_two_level_nested() + run = graph.stream_v2({"value": "x", "items": []}) + + handle_paths: list[tuple[str, ...]] = [] + final_status: dict[tuple[str, ...], str] = {} + for handle in run.subgraphs: + # Drill into the handle's projections inside the loop body so + # the mini-mux is subscribed before the next pump cycle. + list(handle.values) + handle_paths.append(handle.path) + final_status[handle.path] = handle.status + + assert len(handle_paths) == 1 + assert handle_paths[0][0].startswith("middle:") + assert final_status[handle_paths[0]] == "completed" + + +def test_stream_v2_grandchild_visible_on_child_handle() -> None: + """Drilling into `handle.subgraphs` surfaces nested grandchildren.""" + graph = _make_two_level_nested() + run = graph.stream_v2({"value": "x", "items": []}) + + grandchild_paths: list[tuple[str, ...]] = [] + middle_path: tuple[str, ...] | None = None + for middle_handle in run.subgraphs: + # Subscribe to grandchildren before the next pump cycle. + for inner_handle in middle_handle.subgraphs: + # Subscribe to inner.values so its mini-mux drains. + list(inner_handle.values) + grandchild_paths.append(inner_handle.path) + middle_path = middle_handle.path + + assert middle_path is not None + assert len(grandchild_paths) == 1 + inner_path = grandchild_paths[0] + assert inner_path[1].startswith("inner:") + assert inner_path[: len(middle_path)] == middle_path + + +def test_subgraph_output_stops_at_own_terminal_without_draining_siblings() -> None: + """A handle's `output` must not pump past its terminal event. + + If it over-pumps the root run, the second sibling handle is yielded + only after it has already completed, so subscribing to `values` + inside the loop body misses its events. + """ + graph = _make_two_sibling_subgraphs() + run = graph.stream_v2({"value": "x", "items": []}) + + paths: list[tuple[str, ...]] = [] + second_values: list[dict[str, Any]] = [] + for handle in run.subgraphs: + paths.append(handle.path) + if handle.graph_name == "one": + assert handle.output is not None + assert handle.status == "completed" + elif handle.graph_name == "two": + second_values = list(handle.values) + + assert [path[0].split(":", 1)[0] for path in paths] == ["one", "two"] + assert second_values + assert second_values[-1]["items"] == ["one", "two"] + + +def test_aborted_subgraph_handle_does_not_fail_parent_forwarding() -> None: + graph = _make_two_sibling_subgraphs() + run = graph.stream_v2({"value": "x", "items": []}) + + seen: list[str | None] = [] + for handle in run.subgraphs: + seen.append(handle.graph_name) + if handle.graph_name == "one": + # Subscribe before aborting to ensure forwarding into the + # closed mini-mux would have raised without the closed check. + iter(handle.values) + handle.abort() + elif handle.graph_name == "two": + assert list(handle.values) + + assert seen == ["one", "two"] + + +def test_failed_subgraph_output_raises_terminal_error() -> None: + graph = _make_failing_nested() + run = graph.stream_v2({"value": "x", "items": []}) + + handle = next(iter(run.subgraphs)) + with pytest.raises(RuntimeError, match="child boom"): + _ = handle.output + assert handle.status == "failed" + assert handle.error == "child boom" diff --git a/libs/langgraph/tests/test_stream_v2_e2e.py b/libs/langgraph/tests/test_stream_v2_e2e.py new file mode 100644 index 000000000..bfbc9278a --- /dev/null +++ b/libs/langgraph/tests/test_stream_v2_e2e.py @@ -0,0 +1,792 @@ +"""End-to-end tests exercising all stream_v2 projections together. + +Each test builds a realistic graph (subgraphs, LLM calls, custom writers, +interrupts) and verifies that every projection — values, messages, lifecycle, +subgraphs, raw events, output, interleave — produces correct, consistent +results through a single stream_v2 / astream_v2 run. +""" + +from __future__ import annotations + +import operator +import sys +from typing import Annotated, 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 +from langgraph.checkpoint.memory import InMemorySaver +from typing_extensions import TypedDict + +from langgraph.constants import END, START +from langgraph.graph import MessagesState, StateGraph +from langgraph.stream import StreamChannel, StreamTransformer +from langgraph.stream._types import ProtocolEvent +from langgraph.types import StreamWriter, interrupt + +NEEDS_CONTEXTVARS = pytest.mark.skipif( + sys.version_info < (3, 11), + reason="Python 3.11+ is required for async contextvars support", +) + + +# --------------------------------------------------------------------------- +# State and graph builders +# --------------------------------------------------------------------------- + + +class AgentState(TypedDict): + value: str + items: Annotated[list[str], operator.add] + + +def _make_nested_graph(): + """Build a two-level graph with pure state transforms. + + Structure: + outer: + router_node (state transform) + inner_graph (compiled subgraph) + + inner_graph: + process_node (state transform) + """ + + def process_node(state: AgentState) -> dict[str, Any]: + return {"value": state["value"] + "_processed", "items": ["processed"]} + + inner_builder: StateGraph = StateGraph(AgentState, input_schema=AgentState) + inner_builder.add_node("process_node", process_node) + inner_builder.add_edge(START, "process_node") + inner_builder.add_edge("process_node", END) + inner_graph = inner_builder.compile() + + def router_node(state: AgentState) -> dict[str, Any]: + return {"value": state["value"] + "_routed", "items": ["routed"]} + + outer_builder: StateGraph = StateGraph(AgentState, input_schema=AgentState) + outer_builder.add_node("router", router_node) + outer_builder.add_node("inner", inner_graph) + outer_builder.add_edge(START, "router") + outer_builder.add_edge("router", "inner") + outer_builder.add_edge("inner", END) + return outer_builder.compile() + + +def _make_messages_graph(): + """Flat graph with an LLM call for messages projection testing.""" + model = GenericFakeChatModel(messages=iter(["hello world"])) + + def call_model(state: MessagesState) -> dict[str, Any]: + return {"messages": model.invoke(state["messages"])} + + return ( + StateGraph(MessagesState) + .add_node("call_model", call_model) + .add_edge(START, "call_model") + .add_edge("call_model", END) + .compile() + ) + + +def _make_messages_subgraph(): + """Outer graph with a MessagesState subgraph that returns an AIMessage. + + Uses the whole-message fallback path (node returns AIMessage directly) + to exercise messages through a subgraph boundary. + """ + + def return_message(state: MessagesState) -> dict[str, Any]: + return {"messages": AIMessage(content="from subgraph", id="sub-msg-1")} + + inner = ( + StateGraph(MessagesState) + .add_node("return_message", return_message) + .add_edge(START, "return_message") + .add_edge("return_message", END) + .compile() + ) + + class OuterState(TypedDict): + messages: Annotated[list[Any], operator.add] + done: bool + + def pre_node(state: OuterState) -> dict[str, Any]: + return {"done": False} + + return ( + StateGraph(OuterState) + .add_node("pre", pre_node) + .add_node("inner", inner) + .add_edge(START, "pre") + .add_edge("pre", "inner") + .add_edge("inner", END) + .compile() + ) + + +def _make_custom_writer_graph(): + """Graph where a node emits custom stream events via StreamWriter.""" + + def writer_node(state: AgentState, *, writer: StreamWriter) -> dict[str, Any]: + writer({"step": "start", "detail": "beginning work"}) + writer({"step": "middle", "detail": "processing"}) + writer({"step": "end", "detail": "done"}) + return {"value": state["value"] + "_custom", "items": ["custom"]} + + builder = StateGraph(AgentState) + builder.add_node("writer_node", writer_node) + builder.add_edge(START, "writer_node") + builder.add_edge("writer_node", END) + return builder.compile() + + +def _make_interrupt_graph(): + """Graph that interrupts after the first node.""" + + def step_one(state: AgentState) -> dict[str, Any]: + return {"value": state["value"] + "_step1", "items": ["step1"]} + + def step_two(state: AgentState) -> dict[str, Any]: + answer = interrupt("need approval") + return {"value": state["value"] + f"_{answer}", "items": ["step2"]} + + builder = StateGraph(AgentState) + builder.add_node("step_one", step_one) + builder.add_node("step_two", step_two) + builder.add_edge(START, "step_one") + builder.add_edge("step_one", "step_two") + builder.add_edge("step_two", END) + return builder.compile(checkpointer=InMemorySaver()) + + +def _make_error_subgraph(): + """Graph with a subgraph that raises.""" + + def failing_node(state: AgentState) -> dict[str, Any]: + raise ValueError("subgraph explosion") + + inner_builder = StateGraph(AgentState) + inner_builder.add_node("fail", failing_node) + inner_builder.add_edge(START, "fail") + inner_builder.add_edge("fail", END) + inner = inner_builder.compile() + + outer_builder = StateGraph(AgentState) + outer_builder.add_node("inner", inner) + outer_builder.add_edge(START, "inner") + outer_builder.add_edge("inner", END) + return outer_builder.compile() + + +class _CustomPassthroughTransformer(StreamTransformer): + required_stream_modes = ("custom",) + + def init(self) -> dict[str, Any]: + return {} + + def process(self, event: ProtocolEvent) -> bool: + return True + + +class _CounterTransformer(StreamTransformer): + """Custom transformer that counts values events via a StreamChannel.""" + + def __init__(self, scope: tuple[str, ...] = ()) -> None: + super().__init__(scope) + self._channel: StreamChannel[int] = StreamChannel("counter") + self._count = 0 + + def init(self) -> dict[str, Any]: + return {"counter": self._channel} + + def process(self, event: ProtocolEvent) -> bool: + if event["method"] == "values": + self._count += 1 + self._channel.push(self._count) + return True + + +# --------------------------------------------------------------------------- +# Sync end-to-end: all projections on nested graph +# --------------------------------------------------------------------------- + + +class TestStreamV2E2ESync: + def test_all_projections_nested_graph(self) -> None: + """Run a nested graph through stream_v2 and verify values + lifecycle.""" + graph = _make_nested_graph() + run = graph.stream_v2({"value": "x", "items": []}) + + values_snapshots: list[dict[str, Any]] = [] + lifecycle_events: list[dict[str, Any]] = [] + for name, item in run.interleave("values", "lifecycle"): + if name == "values": + values_snapshots.append(item) + elif name == "lifecycle": + lifecycle_events.append(item) + + assert len(values_snapshots) >= 1 + final = values_snapshots[-1] + assert "routed" in final["items"] + assert "processed" in final["items"] + assert "_routed" in final["value"] + assert "_processed" in final["value"] + + assert len(lifecycle_events) >= 2 + started = [e for e in lifecycle_events if e["event"] == "started"] + completed = [e for e in lifecycle_events if e["event"] == "completed"] + assert len(started) >= 1 + assert len(completed) >= 1 + + def test_subgraph_handles_with_drill_down(self) -> None: + """Subgraph handles yield and support values drill-down.""" + graph = _make_nested_graph() + run = graph.stream_v2({"value": "x", "items": []}) + + handles = [] + for handle in run.subgraphs: + child_values = list(handle.values) + handles.append( + { + "path": handle.path, + "graph_name": handle.graph_name, + "values_count": len(child_values), + } + ) + + assert len(handles) >= 1 + assert handles[0]["values_count"] >= 1 + + output = run.output + assert output is not None + assert "_routed" in output["value"] + assert "_processed" in output["value"] + + def test_raw_events_have_monotonic_seq(self) -> None: + """Raw protocol events have monotonically increasing seq numbers.""" + graph = _make_nested_graph() + run = graph.stream_v2({"value": "x", "items": []}) + events = list(run) + assert len(events) > 0 + + seqs = [e["seq"] for e in events] + for i in range(1, len(seqs)): + assert seqs[i] > seqs[i - 1], f"seq not monotonic at {i}: {seqs}" + + for event in events: + assert event["type"] == "event" + assert "method" in event + assert isinstance(event["params"]["timestamp"], int) + + def test_output_matches_final_values_snapshot(self) -> None: + """output property returns the same state as the last values snapshot.""" + run1 = _make_nested_graph().stream_v2({"value": "x", "items": []}) + snapshots = list(run1.values) + final_via_values = snapshots[-1] + + run2 = _make_nested_graph().stream_v2({"value": "x", "items": []}) + final_via_output = run2.output + + assert final_via_values == final_via_output + + def test_context_manager_and_abort(self) -> None: + """Context manager calls abort, marking the stream exhausted.""" + graph = _make_nested_graph() + with graph.stream_v2({"value": "x", "items": []}) as run: + first_val = next(iter(run.values)) + assert isinstance(first_val, dict) + assert run._exhausted is True + + def test_extensions_has_all_native_keys(self) -> None: + """Extensions dict exposes all native projection keys.""" + graph = _make_nested_graph() + run = graph.stream_v2({"value": "x", "items": []}) + _ = run.output + + assert "values" in run.extensions + assert "messages" in run.extensions + assert "lifecycle" in run.extensions + assert "subgraphs" in run.extensions + assert run.values is run.extensions["values"] + assert run.messages is run.extensions["messages"] + assert run.lifecycle is run.extensions["lifecycle"] + assert run.subgraphs is run.extensions["subgraphs"] + + +# --------------------------------------------------------------------------- +# Sync: messages projection +# --------------------------------------------------------------------------- + + +class TestStreamV2E2EMessages: + def test_messages_projection_from_invoke(self) -> None: + """Messages projection captures LLM calls via model.invoke() auto-routing.""" + graph = _make_messages_graph() + run = graph.stream_v2({"messages": "hi"}) + streams = list(run.messages) + + assert len(streams) >= 1 + for stream in streams: + assert isinstance(stream, ChatModelStream) + assert streams[0].output.text == "hello world" + + def test_messages_text_deltas(self) -> None: + """Text deltas from the messages projection concatenate correctly.""" + 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) + assert "".join(stream.text) == "streamed answer" + + def test_messages_from_whole_ai_message(self) -> None: + """Node returning AIMessage directly produces a complete stream.""" + + def return_msg(state: MessagesState) -> dict[str, Any]: + return {"messages": AIMessage(content="hardcoded", id="msg-1")} + + graph = ( + StateGraph(MessagesState) + .add_node("return_msg", return_msg) + .add_edge(START, "return_msg") + .add_edge("return_msg", END) + .compile() + ) + + run = graph.stream_v2({"messages": "hi"}) + (stream,) = list(run.messages) + assert stream.output.text == "hardcoded" + assert stream.message_id == "msg-1" + + def test_root_messages_only_shows_root_scope(self) -> None: + """Root messages projection doesn't surface subgraph-scoped messages.""" + graph = _make_messages_subgraph() + run = graph.stream_v2({"messages": ["hi"], "done": False}) + root_streams = list(run.messages) + # The message is emitted inside the subgraph, so the root + # messages projection (scoped to root namespace) doesn't see it. + assert root_streams == [] + + def test_subgraph_handle_messages_drill_down(self) -> None: + """Drilling into subgraph handle's messages surfaces subgraph messages.""" + graph = _make_messages_subgraph() + run = graph.stream_v2({"messages": ["hi"], "done": False}) + + found_messages = False + for handle in run.subgraphs: + child_messages = list(handle.messages) + if child_messages: + found_messages = True + assert isinstance(child_messages[0], ChatModelStream) + assert child_messages[0].output.text == "from subgraph" + assert found_messages + + +# --------------------------------------------------------------------------- +# Sync: custom stream writer + custom transformer +# --------------------------------------------------------------------------- + + +class TestStreamV2E2ECustom: + def test_custom_events_with_passthrough_transformer(self) -> None: + """Custom StreamWriter events appear on the main log when a + transformer declares the custom mode.""" + graph = _make_custom_writer_graph() + run = graph.stream_v2( + {"value": "x", "items": []}, + transformers=[_CustomPassthroughTransformer], + ) + events = list(run) + custom = [e for e in events if e["method"] == "custom"] + assert len(custom) == 3 + steps = [e["params"]["data"]["step"] for e in custom] + assert steps == ["start", "middle", "end"] + + def test_custom_events_suppressed_without_transformer(self) -> None: + """Without a custom-mode transformer, custom events don't flow.""" + graph = _make_custom_writer_graph() + run = graph.stream_v2({"value": "x", "items": []}) + events = list(run) + custom = [e for e in events if e["method"] == "custom"] + assert custom == [] + + def test_custom_transformer_with_stream_channel(self) -> None: + """A custom transformer with a StreamChannel produces extension data.""" + graph = _make_nested_graph() + run = graph.stream_v2( + {"value": "x", "items": []}, + transformers=[_CounterTransformer], + ) + + assert "counter" in run.extensions + counter_iter = iter(run.extensions["counter"]) + _ = run.output + counts = list(counter_iter) + assert len(counts) >= 1 + assert all(isinstance(c, int) for c in counts) + assert counts == sorted(counts) + + def test_custom_channel_events_on_main_log(self) -> None: + """StreamChannel auto-forward injects custom: events into the main log.""" + graph = _make_nested_graph() + run = graph.stream_v2( + {"value": "x", "items": []}, + transformers=[_CounterTransformer], + ) + events = list(run) + counter_events = [e for e in events if e["method"] == "custom:counter"] + assert len(counter_events) >= 1 + assert all(isinstance(e["params"]["data"], int) for e in counter_events) + + +# --------------------------------------------------------------------------- +# Sync: interrupt handling +# --------------------------------------------------------------------------- + + +class TestStreamV2E2EInterrupt: + def test_interrupt_sets_flags_and_surfaces_interrupts(self) -> None: + """Interrupted run has correct flags and interrupt payloads.""" + graph = _make_interrupt_graph() + config: dict[str, Any] = {"configurable": {"thread_id": "int-1"}} + run = graph.stream_v2({"value": "x", "items": []}, config) + + output = run.output + assert output is not None + assert run.interrupted is True + assert len(run.interrupts) > 0 + assert output["items"] == ["step1"] + assert "_step1" in output["value"] + + def test_interrupt_values_snapshot_has_partial_state(self) -> None: + """Values snapshots captured before the interrupt reflect partial state.""" + graph = _make_interrupt_graph() + config: dict[str, Any] = {"configurable": {"thread_id": "int-2"}} + run = graph.stream_v2({"value": "x", "items": []}, config) + + snapshots = list(run.values) + assert len(snapshots) >= 1 + last = snapshots[-1] + assert "step1" in last["items"] + + +# --------------------------------------------------------------------------- +# Sync: error propagation +# --------------------------------------------------------------------------- + + +class TestStreamV2E2EErrors: + def test_subgraph_error_propagates_through_output(self) -> None: + """Error in a subgraph propagates through output.""" + graph = _make_error_subgraph() + run = graph.stream_v2({"value": "x", "items": []}) + + with pytest.raises(ValueError, match="subgraph explosion"): + _ = run.output + + def test_subgraph_error_propagates_through_raw_events(self) -> None: + graph = _make_error_subgraph() + run = graph.stream_v2({"value": "x", "items": []}) + + with pytest.raises(ValueError, match="subgraph explosion"): + list(run) + + def test_error_subgraph_handle_status(self) -> None: + """Subgraph handle surfaces the error status.""" + graph = _make_error_subgraph() + run = graph.stream_v2({"value": "x", "items": []}) + + handle = next(iter(run.subgraphs)) + with pytest.raises(RuntimeError, match="subgraph explosion"): + _ = handle.output + assert handle.status == "failed" + assert handle.error == "subgraph explosion" + + +# --------------------------------------------------------------------------- +# Async end-to-end +# --------------------------------------------------------------------------- + + +@pytest.mark.anyio +@NEEDS_CONTEXTVARS +class TestStreamV2E2EAsync: + async def test_all_projections_async(self) -> None: + """Async run exercises values projection.""" + graph = _make_nested_graph() + run = await graph.astream_v2({"value": "x", "items": []}) + + values_snapshots = [s async for s in run.values] + assert len(values_snapshots) >= 1 + final = values_snapshots[-1] + assert "_routed" in final["value"] + assert "_processed" in final["value"] + + async def test_async_output(self) -> None: + """Async output returns the final state.""" + graph = _make_nested_graph() + run = await graph.astream_v2({"value": "x", "items": []}) + output = await run.output() + assert output is not None + assert output["value"] == "x_routed_processed" + assert "routed" in output["items"] + assert "processed" in output["items"] + + async def test_async_raw_events(self) -> None: + """Async raw event iteration yields well-formed ProtocolEvents.""" + graph = _make_nested_graph() + run = await graph.astream_v2({"value": "x", "items": []}) + events = [e async for e in run] + assert len(events) > 0 + seqs = [e["seq"] for e in events] + for i in range(1, len(seqs)): + assert seqs[i] > seqs[i - 1] + + async def test_async_messages_projection(self) -> None: + """Async messages projection captures LLM streams.""" + model = GenericFakeChatModel(messages=iter(["async answer"])) + + 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 = [s async for s in run.messages] + assert len(streams) >= 1 + for s in streams: + assert isinstance(s, AsyncChatModelStream) + assert (await streams[0].output).text == "async answer" + + async def test_async_interrupt(self) -> None: + """Async interrupted run has correct flags.""" + graph = _make_interrupt_graph() + config: dict[str, Any] = {"configurable": {"thread_id": "async-int-1"}} + run = await graph.astream_v2({"value": "x", "items": []}, config) + + output = await run.output() + assert output is not None + assert await run.interrupted() is True + assert len(await run.interrupts()) > 0 + + async def test_async_error_propagation(self) -> None: + """Async error from subgraph propagates through output.""" + graph = _make_error_subgraph() + run = await graph.astream_v2({"value": "x", "items": []}) + with pytest.raises(ValueError, match="subgraph explosion"): + await run.output() + + async def test_async_context_manager(self) -> None: + """Async context manager calls abort on exit.""" + graph = _make_nested_graph() + run = await graph.astream_v2({"value": "x", "items": []}) + async with run: + _ = await anext(aiter(run.values)) + assert run._exhausted is True + + async def test_async_extensions_present(self) -> None: + """Async run has all native extensions.""" + graph = _make_nested_graph() + run = await graph.astream_v2({"value": "x", "items": []}) + _ = await run.output() + assert "values" in run.extensions + assert "messages" in run.extensions + assert "lifecycle" in run.extensions + assert "subgraphs" in run.extensions + + async def test_async_custom_transformer(self) -> None: + """Async custom transformer with StreamChannel works.""" + graph = _make_nested_graph() + run = await graph.astream_v2( + {"value": "x", "items": []}, + transformers=[_CounterTransformer], + ) + assert "counter" in run.extensions + counter_cursor = aiter(run.extensions["counter"]) + _ = await run.output() + counts = [c async for c in counter_cursor] + assert len(counts) >= 1 + assert counts == sorted(counts) + + +# --------------------------------------------------------------------------- +# Sync: combined projections stress test +# --------------------------------------------------------------------------- + + +class TestStreamV2E2ECombined: + def test_interleave_all_native_projections(self) -> None: + """Interleave values + messages + lifecycle without deadlock.""" + graph = _make_nested_graph() + run = graph.stream_v2({"value": "x", "items": []}) + + seen_names: set[str] = set() + for name, _item in run.interleave("values", "messages", "lifecycle"): + seen_names.add(name) + + assert "values" in seen_names + assert "lifecycle" in seen_names + + def test_multiple_custom_transformers(self) -> None: + """Multiple custom transformers can coexist.""" + + class TagTransformer(StreamTransformer): + def __init__(self, scope: tuple[str, ...] = ()) -> None: + super().__init__(scope) + self._channel: StreamChannel[str] = StreamChannel("tags") + + def init(self) -> dict[str, Any]: + return {"tags": self._channel} + + def process(self, event: ProtocolEvent) -> bool: + if event["method"] == "values": + self._channel.push( + f"tag:{event['params']['data'].get('value', '')}" + ) + return True + + graph = _make_nested_graph() + run = graph.stream_v2( + {"value": "x", "items": []}, + transformers=[_CounterTransformer, TagTransformer], + ) + + assert "counter" in run.extensions + assert "tags" in run.extensions + + counter_iter = iter(run.extensions["counter"]) + tags_iter = iter(run.extensions["tags"]) + _ = run.output + counts = list(counter_iter) + tags = list(tags_iter) + + assert len(counts) >= 1 + assert len(tags) >= 1 + assert all(t.startswith("tag:") for t in tags) + + def test_two_sibling_subgraphs_both_discoverable(self) -> None: + """Two sequential subgraph invocations produce two handles.""" + + class _S(TypedDict): + items: Annotated[list[str], operator.add] + + def _item(name: str): + def node(state: _S) -> dict[str, Any]: + return {"items": [name]} + + return node + + inner_a = ( + StateGraph(_S) + .add_node("add_a", _item("a")) + .add_edge(START, "add_a") + .add_edge("add_a", END) + .compile() + ) + inner_b = ( + StateGraph(_S) + .add_node("add_b", _item("b")) + .add_edge(START, "add_b") + .add_edge("add_b", END) + .compile() + ) + + outer = ( + StateGraph(_S) + .add_node("sub_a", inner_a) + .add_node("sub_b", inner_b) + .add_edge(START, "sub_a") + .add_edge("sub_a", "sub_b") + .add_edge("sub_b", END) + .compile() + ) + + run = outer.stream_v2({"items": []}) + handles = [] + for handle in run.subgraphs: + list(handle.values) + handles.append(handle) + + assert len(handles) == 2 + names = [h.graph_name for h in handles] + assert "sub_a" in names + assert "sub_b" in names + assert all(h.status == "completed" for h in handles) + + output = run.output + assert output is not None + assert set(output["items"]) == {"a", "b"} + + def test_lifecycle_matches_subgraph_handles(self) -> None: + """Lifecycle events and subgraph handles agree on discovered subgraphs.""" + run1 = _make_nested_graph().stream_v2({"value": "x", "items": []}) + handle_paths: list[tuple[str, ...]] = [] + for handle in run1.subgraphs: + list(handle.values) + handle_paths.append(handle.path) + + run2 = _make_nested_graph().stream_v2({"value": "x", "items": []}) + lifecycle = list(run2.lifecycle) + + started_ns = [ + tuple(e["namespace"]) for e in lifecycle if e["event"] == "started" + ] + # Handle paths use format "graph_name:call_id", lifecycle namespaces + # use the same format. Both should have the same graph_name prefix. + handle_prefixes = {p[0].split(":")[0] for p in handle_paths} + lifecycle_prefixes = {ns[0].split(":")[0] for ns in started_ns} + assert handle_prefixes == lifecycle_prefixes + + def test_values_plus_messages_plus_custom(self) -> None: + """Values, messages, and a custom transformer all produce data in one run.""" + model = GenericFakeChatModel(messages=iter(["combined test"])) + + 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"}, + transformers=[_CounterTransformer], + ) + + counter_iter = iter(run.extensions["counter"]) + values_iter = iter(run.values) + messages_iter = iter(run.messages) + + values = list(values_iter) + messages = list(messages_iter) + counts = list(counter_iter) + + assert len(values) >= 1 + assert len(messages) >= 1 + assert len(counts) >= 1 + assert messages[0].output.text == "combined test" diff --git a/libs/langgraph/tests/test_tool_stream_handler.py b/libs/langgraph/tests/test_tool_stream_handler.py new file mode 100644 index 000000000..722c70b9a --- /dev/null +++ b/libs/langgraph/tests/test_tool_stream_handler.py @@ -0,0 +1,290 @@ +"""Tests for StreamToolCallHandler and ToolRuntime.emit_output_delta. + +These tests exercise the langgraph-core piece in isolation — the prebuilt +`ToolCallTransformer` has its own test file. Here we feed real graphs +through `Pregel.stream(stream_mode=["tools", ...])` and inspect the raw +`(ns, mode, payload)` tuples on the `tools` channel. +""" + +from __future__ import annotations + +from typing import Annotated, Any + +import pytest +from langchain_core.messages import AIMessage +from langchain_core.tools import tool +from langgraph.prebuilt import ToolNode, ToolRuntime +from typing_extensions import TypedDict + +from langgraph.constants import END, START +from langgraph.graph import StateGraph +from langgraph.graph.message import add_messages +from langgraph.pregel._tools import _tool_call_writer + + +class _State(TypedDict): + messages: Annotated[list, add_messages] + + +def _caller_sync(tool_name: str, tool_args: dict[str, Any], tc_id: str = "tc1"): + def caller(state: _State) -> dict: + return { + "messages": [ + AIMessage( + content="", + tool_calls=[{"name": tool_name, "args": tool_args, "id": tc_id}], + ) + ] + } + + return caller + + +def _caller_async(tool_name: str, tool_args: dict[str, Any], tc_id: str = "tc1"): + async def caller(state: _State) -> dict: + return { + "messages": [ + AIMessage( + content="", + tool_calls=[{"name": tool_name, "args": tool_args, "id": tc_id}], + ) + ] + } + + return caller + + +def _build_graph(caller, tools) -> Any: + sg = StateGraph(_State) + sg.add_node("caller", caller) + sg.add_node("tools", ToolNode(tools)) + sg.add_edge(START, "caller") + sg.add_edge("caller", "tools") + sg.add_edge("tools", END) + return sg.compile() + + +def _tool_events(stream) -> list[tuple[tuple[str, ...], dict]]: + """Collect `(ns, payload)` for every `tools`-mode chunk.""" + out: list[tuple[tuple[str, ...], dict]] = [] + for ns, mode, payload in stream: + if mode == "tools": + out.append((tuple(ns), payload)) + return out + + +class TestSyncGraphSyncTool: + def test_started_finished_cycle(self) -> None: + @tool + def echo(text: str) -> str: + """echo.""" + return f"echoed:{text}" + + graph = _build_graph(_caller_sync("echo", {"text": "hi"}), [echo]) + events = _tool_events( + graph.stream( + {"messages": []}, + stream_mode=["tools"], + subgraphs=True, + ) + ) + + assert [p["event"] for _, p in events] == [ + "tool-started", + "tool-finished", + ] + assert events[0][1]["tool_call_id"] == "tc1" + assert events[0][1]["tool_name"] == "echo" + assert events[0][1]["input"] == {"text": "hi"} + # ToolNode wraps the return in a ToolMessage. + assert events[1][1]["tool_call_id"] == "tc1" + + def test_emit_output_delta_produces_delta_events(self) -> None: + @tool + def streaming_echo(text: str, runtime: ToolRuntime) -> str: + """stream chunks.""" + for chunk in ("a", "b", "c"): + runtime.emit_output_delta(chunk) + return text + + graph = _build_graph( + _caller_sync("streaming_echo", {"text": "x"}), [streaming_echo] + ) + events = _tool_events( + graph.stream( + {"messages": []}, + stream_mode=["tools"], + subgraphs=True, + ) + ) + + deltas = [p["delta"] for _, p in events if p["event"] == "tool-output-delta"] + assert deltas == ["a", "b", "c"] + # The deltas must be bracketed by started and finished. + ordered = [p["event"] for _, p in events] + assert ordered[0] == "tool-started" + assert ordered[-1] == "tool-finished" + + def test_tool_error_event(self) -> None: + @tool + def boom() -> str: + """raises.""" + raise ValueError("nope") + + graph = _build_graph(_caller_sync("boom", {}), [boom]) + events: list[tuple[tuple[str, ...], dict]] = [] + with pytest.raises(ValueError, match="nope"): + for ns, mode, payload in graph.stream( + {"messages": []}, + stream_mode=["tools"], + subgraphs=True, + ): + if mode == "tools": + events.append((tuple(ns), payload)) + + kinds = [p["event"] for _, p in events] + assert kinds == ["tool-started", "tool-error"] + assert events[1][1]["message"] == "nope" + + def test_writer_unset_outside_tool(self) -> None: + # Outside any tool body the ContextVar that ToolRuntime reads + # is unset — emitting from there would be a no-op. + assert _tool_call_writer.get() is None + + def test_no_events_without_tools_mode(self) -> None: + @tool + def echo(text: str) -> str: + """echo.""" + return text + + graph = _build_graph(_caller_sync("echo", {"text": "hi"}), [echo]) + # No "tools" in stream_mode — handler is not attached and zero + # `tools`-method events fire. + chunks = list( + graph.stream( + {"messages": []}, + stream_mode=["values"], + subgraphs=True, + ) + ) + assert all( + not (isinstance(c, tuple) and len(c) == 3 and c[1] == "tools") + for c in chunks + ) + + +class TestAsyncGraphAsyncTool: + @pytest.mark.anyio + async def test_async_tool_produces_events(self) -> None: + @tool + async def aecho(text: str, runtime: ToolRuntime) -> str: + """async echo.""" + runtime.emit_output_delta(text) + return f"got:{text}" + + graph = _build_graph(_caller_async("aecho", {"text": "hi"}), [aecho]) + events: list[tuple[tuple[str, ...], dict]] = [] + async for ns, mode, payload in graph.astream( + {"messages": []}, + stream_mode=["tools"], + subgraphs=True, + ): + if mode == "tools": + events.append((tuple(ns), payload)) + + kinds = [p["event"] for _, p in events] + assert kinds == ["tool-started", "tool-output-delta", "tool-finished"] + assert events[1][1]["delta"] == "hi" + + +class TestConcurrentToolCalls: + def test_parallel_tool_calls_do_not_bleed(self) -> None: + @tool + def streamer(marker: str, runtime: ToolRuntime) -> str: + """emits marker twice.""" + runtime.emit_output_delta(f"{marker}-1") + runtime.emit_output_delta(f"{marker}-2") + return marker + + def caller(state: _State) -> dict: + return { + "messages": [ + AIMessage( + content="", + tool_calls=[ + {"name": "streamer", "args": {"marker": "A"}, "id": "a"}, + {"name": "streamer", "args": {"marker": "B"}, "id": "b"}, + ], + ) + ] + } + + graph = _build_graph(caller, [streamer]) + events = _tool_events( + graph.stream( + {"messages": []}, + stream_mode=["tools"], + subgraphs=True, + ) + ) + + # Group deltas by tool_call_id. + by_id: dict[str, list[str]] = {} + for _, p in events: + if p["event"] == "tool-output-delta": + by_id.setdefault(p["tool_call_id"], []).append(p["delta"]) + assert by_id["a"] == ["A-1", "A-2"] + assert by_id["b"] == ["B-1", "B-2"] + + +class TestSubgraphNamespacePropagation: + def test_tool_inside_subgraph_emits_with_subgraph_ns(self) -> None: + @tool + def inner_tool(text: str) -> str: + """inner tool.""" + return text + + def sub_caller(state: _State) -> dict: + return { + "messages": [ + AIMessage( + content="", + tool_calls=[ + { + "name": "inner_tool", + "args": {"text": "x"}, + "id": "tc1", + } + ], + ) + ] + } + + inner = StateGraph(_State) + inner.add_node("sub_caller", sub_caller) + inner.add_node("sub_tools", ToolNode([inner_tool])) + inner.add_edge(START, "sub_caller") + inner.add_edge("sub_caller", "sub_tools") + inner.add_edge("sub_tools", END) + inner_graph = inner.compile() + + outer = StateGraph(_State) + outer.add_node("sub", inner_graph) + outer.add_edge(START, "sub") + outer.add_edge("sub", END) + graph = outer.compile() + + events = _tool_events( + graph.stream( + {"messages": []}, + stream_mode=["tools"], + subgraphs=True, + ) + ) + + # All `tools` events should carry a non-empty namespace rooted + # at the `sub` node. + assert events, "expected at least one tools event" + for ns, _ in events: + assert ns # non-empty + assert ns[0].startswith("sub:") diff --git a/libs/langgraph/uv.lock b/libs/langgraph/uv.lock index 87636bce7..9a3ca4cb7 100644 --- a/libs/langgraph/uv.lock +++ b/libs/langgraph/uv.lock @@ -1348,10 +1348,11 @@ wheels = [ [[package]] name = "langchain-core" -version = "1.3.1" +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/f4/fe/abeae8d0d2899e191d67c6c7f065f7e52a953f30b21ef327fa49084e4af9/langchain_core-1.3.1.tar.gz", hash = "sha256:41b384055799f93f34520df6bf7b80e2e5e23153cdfd46874251c6c9916ea030", size = 862403, upload-time = "2026-04-23T18:54:01.857Z" } +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/a1/c2/8493be505921857988db068b7c027f28a9b1587b4425c6a32b1221c9c9fe/langchain_core-1.3.1-py3-none-any.whl", hash = "sha256:8b13d19d3bed3f4768df12c7f6932d2ada715f3ac9fd020c63d28c693968269e", size = 515879, upload-time = "2026-04-23T18:53:59.94Z" }, + { 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.12" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "typing-extensions" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/5c/51/1157009b6f94e6e58be58fa8b620187d657909a8b36a6bf5b0c52a2711f6/langchain_protocol-0.0.12.tar.gz", hash = "sha256:5e14c434290a705c9510fdb1a83ecf7561a5e6e0dfd053930ade80dba069269f", size = 6408, upload-time = "2026-04-25T01:05:01.489Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/95/82/3431e3061c917439589fa88a6b23c9bc0e154cba0f05d2e895a68c76ff74/langchain_protocol-0.0.12-py3-none-any.whl", hash = "sha256:402b61f42d4139692528cf37226c367bb6efc8ff8165b29380accb0abfece7b2", size = 6639, upload-time = "2026-04-25T01:05:00.487Z" }, ] [[package]] @@ -1439,7 +1452,7 @@ test = [ [package.metadata] requires-dist = [ - { name = "langchain-core", specifier = ">=1.3.0,<2" }, + { name = "langchain-core", specifier = ">=1.3.2,<2" }, { name = "langgraph-checkpoint", editable = "../checkpoint" }, { name = "langgraph-prebuilt", editable = "../prebuilt" }, { name = "langgraph-sdk", editable = "../sdk-py" }, diff --git a/libs/prebuilt/langgraph/prebuilt/__init__.py b/libs/prebuilt/langgraph/prebuilt/__init__.py index a93cc0219..0d8795262 100644 --- a/libs/prebuilt/langgraph/prebuilt/__init__.py +++ b/libs/prebuilt/langgraph/prebuilt/__init__.py @@ -1,5 +1,6 @@ """langgraph.prebuilt exposes a higher-level API for creating and executing agents and tools.""" +from langgraph.prebuilt._tool_call_transformer import ToolCallTransformer from langgraph.prebuilt.chat_agent_executor import create_react_agent from langgraph.prebuilt.tool_node import ( InjectedState, @@ -13,6 +14,7 @@ from langgraph.prebuilt.tool_validator import ValidationNode __all__ = [ "create_react_agent", "ToolNode", + "ToolCallTransformer", "tools_condition", "ValidationNode", "InjectedState", diff --git a/libs/prebuilt/langgraph/prebuilt/_tool_call_stream.py b/libs/prebuilt/langgraph/prebuilt/_tool_call_stream.py new file mode 100644 index 000000000..2df12e823 --- /dev/null +++ b/libs/prebuilt/langgraph/prebuilt/_tool_call_stream.py @@ -0,0 +1,117 @@ +"""In-process handle for a single tool call's streaming execution. + +Mirrors the shape of `ChatModelStream` from langchain-core but simpler — +a tool has one output channel, no content-block multiplexing. Populated +by `ToolCallTransformer` as `tool-started` / `tool-output-delta` / +`tool-finished` / `tool-error` events flow in on the `tools` channel. +""" + +from __future__ import annotations + +from collections.abc import AsyncIterator, Iterator +from typing import Any + +from langgraph.stream._event_log import EventLog + + +class ToolCallStream: + """Scoped view of a single tool call's lifecycle. + + Yielded on `run.tool_calls` once per `tool-started` event. Fields + are populated as events arrive: + + - `tool_call_id`, `tool_name`, `input`: stable from the start event. + - `output_deltas`: an `EventLog` of delta chunks. Iterate (sync or + async) to consume partial output in arrival order. + - `output`: terminal payload from `tool-finished`, or `None` if the + call failed or is still in flight. + - `error`: terminal error string from `tool-error`, or `None` if the + call succeeded or is still in flight. + - `completed`: True once a terminal event (`tool-finished` or + `tool-error`) has been observed. + + `ToolCallStream` is not meant to be constructed by end users — it's + produced by `ToolCallTransformer` as events flow through the mux. + """ + + def __init__( + self, + tool_call_id: str, + tool_name: str, + input: dict[str, Any] | None = None, + ) -> None: + """Initialize a fresh handle for a tool call. + + Args: + tool_call_id: The `tool_call_id` from the AIMessage. + tool_name: The tool's name. + input: The tool's input arguments (as reported by + `on_tool_start`), or `None` if none were captured. + """ + self.tool_call_id = tool_call_id + self.tool_name = tool_name + self.input = input + self._output_deltas: EventLog[Any] = EventLog() + self.output: Any = None + self.error: str | None = None + self.completed = False + + @property + def output_deltas(self) -> EventLog[Any]: + """The EventLog of streamed `tool-output-delta` payloads. + + Iterate (sync or async depending on how the run was started) + to consume partial output in arrival order. The log closes when + the tool finishes or errors. + """ + return self._output_deltas + + def _bind(self, *, is_async: bool) -> None: + """Bind the deltas log to sync or async iteration. + + Called by `ToolCallTransformer` when constructing this handle so + the log matches the enclosing mux's mode. + """ + self._output_deltas._bind(is_async=is_async) + + def _push_delta(self, delta: Any) -> None: + self._output_deltas.push(delta) + + def _finish(self, output: Any) -> None: + self.output = output + self.completed = True + self._output_deltas.close() + + def _fail(self, message: str) -> None: + self.error = message + self.completed = True + self._output_deltas.close() + + def __iter__(self) -> Iterator[Any]: + """Iterate delta chunks synchronously. + + Equivalent to `iter(self.output_deltas)`. Raises `TypeError` if + the underlying log is bound to async mode. + """ + return iter(self._output_deltas) + + def __aiter__(self) -> AsyncIterator[Any]: + """Iterate delta chunks asynchronously. + + Equivalent to `aiter(self.output_deltas)`. Raises `TypeError` + if the underlying log is bound to sync mode. + """ + return self._output_deltas.__aiter__() + + def __repr__(self) -> str: + status = ( + "completed" + if self.completed and self.error is None + else "failed" + if self.completed + else "running" + ) + return ( + f"ToolCallStream(tool_call_id={self.tool_call_id!r}, " + f"tool_name={self.tool_name!r}, status={status})" + ) diff --git a/libs/prebuilt/langgraph/prebuilt/_tool_call_transformer.py b/libs/prebuilt/langgraph/prebuilt/_tool_call_transformer.py new file mode 100644 index 000000000..949a473cc --- /dev/null +++ b/libs/prebuilt/langgraph/prebuilt/_tool_call_transformer.py @@ -0,0 +1,128 @@ +"""Transformer that projects `tools` channel events into `ToolCallStream`s.""" + +from __future__ import annotations + +from collections.abc import Awaitable, Callable +from typing import Any + +from langgraph.stream._event_log import EventLog +from langgraph.stream._types import ProtocolEvent, StreamTransformer + +from langgraph.prebuilt._tool_call_stream import ToolCallStream + + +class ToolCallTransformer(StreamTransformer): + """Project `tools` channel events into `ToolCallStream` handles. + + Each `tool-started` event spawns a `ToolCallStream`, pushed onto + `run.tool_calls`. Subsequent `tool-output-delta` events append to + that stream's deltas log; `tool-finished` and `tool-error` close it. + + Native transformer — the `tool_calls` projection is exposed as a + direct attribute on the run stream. + + `EventLog[ToolCallStream]` is used (not `StreamChannel`) because the + live handles are not serializable and should not be auto-forwarded + onto the main event log. Wire consumers subscribe to the `tools` + channel instead, where the raw protocol events flow through + untouched by this transformer (`process` returns `True`). + + Registered explicitly by users at compile time via + `builder.compile(transformers=[ToolCallTransformer])` — not a + default built-in, so the `tools` channel is user-opt-in. + """ + + _native = True + required_stream_modes = ("tools",) + + def __init__(self, scope: tuple[str, ...] = ()) -> None: + super().__init__(scope) + self._log: EventLog[ToolCallStream] = EventLog() + self._active: dict[str, ToolCallStream] = {} + self._is_async = False + self._pump_fn: Callable[[], bool] | None = None + self._apump_fn: Callable[[], Awaitable[bool]] | None = None + + def init(self) -> dict[str, Any]: + return {"tool_calls": self._log} + + def _bind_pump(self, fn: Callable[[], bool]) -> None: + """Wire the sync pull callback onto this transformer. + + Called by `StreamMux.bind_pump`. Stored so each new + `ToolCallStream` created by `process` can wire its deltas log + for pump-driven iteration. + """ + self._pump_fn = fn + self._is_async = False + + def _bind_apump(self, fn: Callable[[], Awaitable[bool]]) -> None: + """Async counterpart to `_bind_pump`.""" + self._apump_fn = fn + self._is_async = True + + def _new_stream( + self, + tool_call_id: str, + tool_name: str, + tool_input: dict[str, Any] | None, + ) -> ToolCallStream: + stream = ToolCallStream(tool_call_id, tool_name, tool_input) + stream._bind(is_async=self._is_async) + if self._apump_fn is not None: + stream._output_deltas._arequest_more = self._apump_fn + if self._pump_fn is not None: + stream._output_deltas._request_more = self._pump_fn + return stream + + def process(self, event: ProtocolEvent) -> bool: + # Namespace filtering is handled by the mux via `scope_exact`. + if event["method"] != "tools": + return True + + data = event["params"]["data"] + tool_call_id = data.get("tool_call_id") + if tool_call_id is None: + return True + event_type = data.get("event") + + stream: ToolCallStream | None + if event_type == "tool-started": + stream = self._new_stream( + tool_call_id, + data.get("tool_name", ""), + data.get("input"), + ) + self._active[tool_call_id] = stream + self._log.push(stream) + elif event_type == "tool-output-delta": + stream = self._active.get(tool_call_id) + if stream is not None: + stream._push_delta(data.get("delta")) + elif event_type == "tool-finished": + stream = self._active.pop(tool_call_id, None) + if stream is not None: + stream._finish(data.get("output")) + elif event_type == "tool-error": + stream = self._active.pop(tool_call_id, None) + if stream is not None: + stream._fail(data.get("message", "")) + + # Pass-through — wire consumers subscribe to the `tools` channel + # directly and reconstruct handles client-side. + return True + + def finalize(self) -> None: + """Close any still-active tool streams left open at run end.""" + for stream in self._active.values(): + if not stream.completed: + stream._finish(None) + self._active.clear() + + def fail(self, err: BaseException) -> None: + """Fail any still-active tool streams when the run errors.""" + message = str(err) + for stream in self._active.values(): + if not stream.completed: + stream._fail(message) + self._active.clear() diff --git a/libs/prebuilt/langgraph/prebuilt/tool_node.py b/libs/prebuilt/langgraph/prebuilt/tool_node.py index cd0a5e01b..3137309da 100644 --- a/libs/prebuilt/langgraph/prebuilt/tool_node.py +++ b/libs/prebuilt/langgraph/prebuilt/tool_node.py @@ -86,6 +86,7 @@ from langgraph._internal._constants import CONF, CONFIG_KEY_READ from langgraph._internal._runnable import RunnableCallable from langgraph.errors import GraphBubbleUp from langgraph.graph.message import REMOVE_ALL_MESSAGES +from langgraph.pregel._tools import _tool_call_writer from langgraph.runtime import ExecutionInfo, ServerInfo # noqa: TC002 from langgraph.store.base import BaseStore # noqa: TC002 from langgraph.types import Command, Send, StreamWriter @@ -1728,6 +1729,26 @@ class ToolRuntime(_DirectlyInjectedToolArg, Generic[ContextT, StateT]): execution_info: ExecutionInfo | None = None server_info: ServerInfo | None = None + def emit_output_delta(self, delta: Any) -> None: + """Stream a partial output chunk on the `tools` stream channel. + + Reads the per-tool-call writer that `StreamToolCallHandler` + installs on a ContextVar at `on_tool_start` and forwards `delta` + through it. Silent no-op when the graph was not run with + `"tools"` in `stream_mode` (no writer is set), so tool authors + can leave `emit_output_delta` calls in place without gating + them on stream mode. + + Args: + delta: Partial output chunk. Any JSON-serializable value; + surfaced as-is on the `tools` channel's + `tool-output-delta` payload under `"delta"`. + """ + writer = _tool_call_writer.get() + if writer is None: + return + writer(delta) + class InjectedState(InjectedToolArg): """Annotation for injecting graph state into tool arguments. diff --git a/libs/prebuilt/tests/test_tool_call_transformer.py b/libs/prebuilt/tests/test_tool_call_transformer.py new file mode 100644 index 000000000..b7c88d65f --- /dev/null +++ b/libs/prebuilt/tests/test_tool_call_transformer.py @@ -0,0 +1,307 @@ +"""Tests for ToolCallTransformer and the ToolCallStream projection.""" + +from __future__ import annotations + +import time +from typing import Annotated, Any + +import pytest +from langchain_core.messages import AIMessage +from langchain_core.tools import tool +from langgraph.constants import END, START +from langgraph.graph import StateGraph +from langgraph.graph.message import add_messages +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, + ValuesTransformer, +) +from typing_extensions import TypedDict + +from langgraph.prebuilt import ( + ToolCallTransformer, + ToolNode, + ToolRuntime, +) +from langgraph.prebuilt._tool_call_stream import ToolCallStream + +TS = int(time.time() * 1000) + + +def _tool_event( + event: str, + tool_call_id: str, + *, + tool_name: str = "", + input: dict[str, Any] | None = None, + delta: Any = None, + output: Any = None, + message: str = "", + namespace: list[str] | None = None, +) -> ProtocolEvent: + data: dict[str, Any] = {"event": event, "tool_call_id": tool_call_id} + if event == "tool-started": + data["tool_name"] = tool_name + if input is not None: + data["input"] = input + elif event == "tool-output-delta": + data["delta"] = delta + elif event == "tool-finished": + data["output"] = output + elif event == "tool-error": + data["message"] = message + return { + "type": "event", + "method": "tools", + "params": { + "namespace": namespace or [], + "timestamp": TS, + "data": data, + }, + } + + +def _subscribe(log: EventLog) -> None: + log._subscribed = True + + +def _mux() -> tuple[StreamMux, ToolCallTransformer]: + transformer = ToolCallTransformer() + mux = StreamMux( + [ + ValuesTransformer(), + MessagesTransformer(), + transformer, + ], + is_async=False, + ) + _subscribe(transformer._log) + return mux, transformer + + +class TestToolCallTransformerUnit: + def test_required_stream_modes_declares_tools(self) -> None: + assert ToolCallTransformer.required_stream_modes == ("tools",) + + def test_tool_started_yields_handle(self) -> None: + mux, transformer = _mux() + mux.push( + _tool_event( + "tool-started", + "tc1", + tool_name="echo", + input={"text": "hi"}, + ) + ) + handles = list(transformer._log._items) + assert len(handles) == 1 + h = handles[0] + assert isinstance(h, ToolCallStream) + assert h.tool_call_id == "tc1" + assert h.tool_name == "echo" + assert h.input == {"text": "hi"} + assert h.completed is False + + def test_delta_accumulates_on_active_stream(self) -> None: + mux, transformer = _mux() + mux.push(_tool_event("tool-started", "tc1", tool_name="echo")) + _subscribe(transformer._active["tc1"]._output_deltas) + mux.push(_tool_event("tool-output-delta", "tc1", delta="a")) + mux.push(_tool_event("tool-output-delta", "tc1", delta="b")) + stream = transformer._active["tc1"] + assert list(stream._output_deltas._items) == ["a", "b"] + + def test_finish_closes_stream(self) -> None: + mux, transformer = _mux() + mux.push(_tool_event("tool-started", "tc1", tool_name="echo")) + stream = transformer._active["tc1"] + mux.push(_tool_event("tool-finished", "tc1", output="done")) + assert stream.completed is True + assert stream.output == "done" + assert stream.error is None + assert "tc1" not in transformer._active + + def test_error_closes_stream(self) -> None: + mux, transformer = _mux() + mux.push(_tool_event("tool-started", "tc1", tool_name="boom")) + stream = transformer._active["tc1"] + mux.push(_tool_event("tool-error", "tc1", message="nope")) + assert stream.completed is True + assert stream.output is None + assert stream.error == "nope" + assert "tc1" not in transformer._active + + def test_concurrent_tool_calls_do_not_bleed(self) -> None: + mux, transformer = _mux() + mux.push(_tool_event("tool-started", "a", tool_name="t")) + mux.push(_tool_event("tool-started", "b", tool_name="t")) + for tc in ("a", "b"): + _subscribe(transformer._active[tc]._output_deltas) + mux.push(_tool_event("tool-output-delta", "a", delta="A1")) + mux.push(_tool_event("tool-output-delta", "b", delta="B1")) + mux.push(_tool_event("tool-output-delta", "a", delta="A2")) + assert list(transformer._active["a"]._output_deltas._items) == ["A1", "A2"] + assert list(transformer._active["b"]._output_deltas._items) == ["B1"] + + def test_tools_event_passes_through_main_log(self) -> None: + mux, transformer = _mux() + _subscribe(mux._events) + mux.push(_tool_event("tool-started", "tc1", tool_name="echo")) + kept = [e for e in mux._events._items if e["method"] == "tools"] + assert len(kept) == 1 + + +# --------------------------------------------------------------------------- +# End-to-end tests with a real graph +# --------------------------------------------------------------------------- + + +class _State(TypedDict): + messages: Annotated[list, add_messages] + + +def _build_graph(caller, tools): + sg = StateGraph(_State) + sg.add_node("caller", caller) + sg.add_node("tools", ToolNode(tools)) + sg.add_edge(START, "caller") + sg.add_edge("caller", "tools") + sg.add_edge("tools", END) + return sg.compile() + + +class TestToolCallTransformerEndToEnd: + def test_sync_streaming_tool_populates_tool_calls(self) -> None: + @tool + def streamer(text: str, runtime: ToolRuntime) -> str: + """streams chunks.""" + for chunk in ("one", "two"): + runtime.emit_output_delta(chunk) + return text + + def caller(state: _State) -> dict: + return { + "messages": [ + AIMessage( + content="", + tool_calls=[ + {"name": "streamer", "args": {"text": "x"}, "id": "tc1"} + ], + ) + ] + } + + graph = _build_graph(caller, [streamer]) + run = graph.stream_v2({"messages": []}, transformers=[ToolCallTransformer]) + + tool_calls: list[ToolCallStream] = [] + for tc in run.tool_calls: + tool_calls.append(tc) + deltas = list(tc.output_deltas) + assert deltas == ["one", "two"] + assert len(tool_calls) == 1 + tc = tool_calls[0] + assert tc.tool_call_id == "tc1" + assert tc.tool_name == "streamer" + assert tc.completed is True + assert tc.error is None + + def test_stream_modes_union_includes_tools(self) -> None: + @tool + def echo(text: str) -> str: + """echo.""" + return text + + def caller(state: _State) -> dict: + return { + "messages": [ + AIMessage( + content="", + tool_calls=[ + {"name": "echo", "args": {"text": "x"}, "id": "tc1"} + ], + ) + ] + } + + graph = _build_graph(caller, [echo]) + # Without ToolCallTransformer, no tool_calls projection is + # exposed and no `tools` events flow through (required_stream_modes + # omits it). + run_no_tc = graph.stream_v2({"messages": []}) + assert "tool_calls" not in run_no_tc._mux.extensions # type: ignore[attr-defined] + + # With ToolCallTransformer, the projection is present. + run = graph.stream_v2({"messages": []}, transformers=[ToolCallTransformer]) + assert "tool_calls" in run._mux.extensions # type: ignore[attr-defined] + # Drain so the run closes cleanly. + list(run.tool_calls) + + @pytest.mark.anyio + async def test_async_streaming_tool_populates_tool_calls(self) -> None: + @tool + async def astreamer(text: str, runtime: ToolRuntime) -> str: + """async streams.""" + runtime.emit_output_delta(text) + runtime.emit_output_delta(text + "!") + return text + + async def caller(state: _State) -> dict: + return { + "messages": [ + AIMessage( + content="", + tool_calls=[ + {"name": "astreamer", "args": {"text": "hi"}, "id": "tc1"} + ], + ) + ] + } + + graph = _build_graph(caller, [astreamer]) + run = await graph.astream_v2( + {"messages": []}, transformers=[ToolCallTransformer] + ) + + collected: list[ToolCallStream] = [] + async for tc in run.tool_calls: + collected.append(tc) + deltas = [d async for d in tc.output_deltas] + assert deltas == ["hi", "hi!"] + assert len(collected) == 1 + assert collected[0].completed is True + assert collected[0].error is None + + def test_tool_error_populates_error_field(self) -> None: + @tool + def boom() -> str: + """raises.""" + raise ValueError("nope") + + def caller(state: _State) -> dict: + return { + "messages": [ + AIMessage( + content="", + tool_calls=[{"name": "boom", "args": {}, "id": "tc1"}], + ) + ] + } + + graph = _build_graph(caller, [boom]) + run = graph.stream_v2({"messages": []}, transformers=[ToolCallTransformer]) + + collected: list[ToolCallStream] = [] + with pytest.raises(ValueError, match="nope"): + for tc in run.tool_calls: + collected.append(tc) + # Drain deltas so the error field is populated before we + # inspect it below. + list(tc.output_deltas) + + assert len(collected) == 1 + assert collected[0].error == "nope" + assert collected[0].output is None + assert collected[0].completed is True diff --git a/libs/prebuilt/uv.lock b/libs/prebuilt/uv.lock index f6b81827b..b781f1176 100644 --- a/libs/prebuilt/uv.lock +++ b/libs/prebuilt/uv.lock @@ -249,10 +249,11 @@ wheels = [ [[package]] name = "langchain-core" -version = "1.3.1" +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/f4/fe/abeae8d0d2899e191d67c6c7f065f7e52a953f30b21ef327fa49084e4af9/langchain_core-1.3.1.tar.gz", hash = "sha256:41b384055799f93f34520df6bf7b80e2e5e23153cdfd46874251c6c9916ea030", size = 862403, upload-time = "2026-04-23T18:54:01.857Z" } +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/a1/c2/8493be505921857988db068b7c027f28a9b1587b4425c6a32b1221c9c9fe/langchain_core-1.3.1-py3-none-any.whl", hash = "sha256:8b13d19d3bed3f4768df12c7f6932d2ada715f3ac9fd020c63d28c693968269e", size = 515879, upload-time = "2026-04-23T18:53:59.94Z" }, + { 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.12" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "typing-extensions" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/5c/51/1157009b6f94e6e58be58fa8b620187d657909a8b36a6bf5b0c52a2711f6/langchain_protocol-0.0.12.tar.gz", hash = "sha256:5e14c434290a705c9510fdb1a83ecf7561a5e6e0dfd053930ade80dba069269f", size = 6408, upload-time = "2026-04-25T01:05:01.489Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/95/82/3431e3061c917439589fa88a6b23c9bc0e154cba0f05d2e895a68c76ff74/langchain_protocol-0.0.12-py3-none-any.whl", hash = "sha256:402b61f42d4139692528cf37226c367bb6efc8ff8165b29380accb0abfece7b2", size = 6639, upload-time = "2026-04-25T01:05:00.487Z" }, ] [[package]] @@ -281,7 +294,7 @@ dependencies = [ [package.metadata] requires-dist = [ - { name = "langchain-core", specifier = ">=1.3.0,<2" }, + { name = "langchain-core", specifier = ">=1.3.2,<2" }, { name = "langgraph-checkpoint", editable = "../checkpoint" }, { name = "langgraph-prebuilt", editable = "." }, { name = "langgraph-sdk", editable = "../sdk-py" }, diff --git a/libs/sdk-py/uv.lock b/libs/sdk-py/uv.lock index 2180f6511..468451c08 100644 --- a/libs/sdk-py/uv.lock +++ b/libs/sdk-py/uv.lock @@ -266,10 +266,11 @@ wheels = [ [[package]] name = "langchain-core" -version = "1.3.1" +version = "1.3.2" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "jsonpatch" }, + { name = "langchain-protocol" }, { name = "langsmith" }, { name = "packaging" }, { name = "pydantic" }, @@ -278,9 +279,21 @@ dependencies = [ { name = "typing-extensions" }, { name = "uuid-utils" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/f4/fe/abeae8d0d2899e191d67c6c7f065f7e52a953f30b21ef327fa49084e4af9/langchain_core-1.3.1.tar.gz", hash = "sha256:41b384055799f93f34520df6bf7b80e2e5e23153cdfd46874251c6c9916ea030", size = 862403, upload-time = "2026-04-23T18:54:01.857Z" } +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/a1/c2/8493be505921857988db068b7c027f28a9b1587b4425c6a32b1221c9c9fe/langchain_core-1.3.1-py3-none-any.whl", hash = "sha256:8b13d19d3bed3f4768df12c7f6932d2ada715f3ac9fd020c63d28c693968269e", size = 515879, upload-time = "2026-04-23T18:53:59.94Z" }, + { 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.12" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "typing-extensions" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/5c/51/1157009b6f94e6e58be58fa8b620187d657909a8b36a6bf5b0c52a2711f6/langchain_protocol-0.0.12.tar.gz", hash = "sha256:5e14c434290a705c9510fdb1a83ecf7561a5e6e0dfd053930ade80dba069269f", size = 6408, upload-time = "2026-04-25T01:05:01.489Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/95/82/3431e3061c917439589fa88a6b23c9bc0e154cba0f05d2e895a68c76ff74/langchain_protocol-0.0.12-py3-none-any.whl", hash = "sha256:402b61f42d4139692528cf37226c367bb6efc8ff8165b29380accb0abfece7b2", size = 6639, upload-time = "2026-04-25T01:05:00.487Z" }, ] [[package]] @@ -298,7 +311,7 @@ dependencies = [ [package.metadata] requires-dist = [ - { name = "langchain-core", specifier = ">=1.3.0,<2" }, + { name = "langchain-core", specifier = ">=1.3.2,<2" }, { name = "langgraph-checkpoint", editable = "../checkpoint" }, { name = "langgraph-prebuilt", editable = "../prebuilt" }, { name = "langgraph-sdk", editable = "." },