From bebcd20815a0bb6387764b5d4efcac6e35d0a392 Mon Sep 17 00:00:00 2001 From: Nick Hollon Date: Sat, 18 Apr 2026 18:59:23 -0400 Subject: [PATCH] Address review feedback on subgraph lifecycle streaming - Robustify nested-Pregel detection with a parent_run_id fallback so subgraphs compiled with name equal to their node name are still recognized. - Narrow bare excepts in SubgraphTransformer close/fail paths; log at warning with exc_info instead of silently swallowing. - Assert mux registration in SubgraphTransformer._on_started instead of silently dropping events. - Warn when RemoteGraph strips an unsupported "lifecycle" stream mode so callers aren't left wondering why no events arrive. - Reject pre-built transformer instances in StreamingHandler; factories are required so transformers propagate into every subgraph scope. - Comment the forward-before-close ordering in SubgraphTransformer. - Trim duplicated pump/projection docstrings across run-stream classes. - Add end-to-end tests for trigger_call_id, subgraph interrupt, and the name-collision detection fallback. --- libs/langgraph/langgraph/pregel/_lifecycle.py | 64 +++++++++++-- libs/langgraph/langgraph/pregel/remote.py | 29 +++++- libs/langgraph/langgraph/stream/run_stream.py | 24 ++--- .../langgraph/stream/streaming_handler.py | 43 ++++----- .../langgraph/stream/transformers.py | 46 ++++++--- .../tests/test_stream_subgraph_transformer.py | 94 +++++++++++++++++++ .../langgraph/tests/test_streaming_handler.py | 23 +++-- 7 files changed, 253 insertions(+), 70 deletions(-) diff --git a/libs/langgraph/langgraph/pregel/_lifecycle.py b/libs/langgraph/langgraph/pregel/_lifecycle.py index 6f2beb07c..7875a748b 100644 --- a/libs/langgraph/langgraph/pregel/_lifecycle.py +++ b/libs/langgraph/langgraph/pregel/_lifecycle.py @@ -22,7 +22,12 @@ T = TypeVar("T") _LANGGRAPH_SENTINEL_NODES = frozenset({"__start__", "__end__"}) -def _is_nested_pregel_start(name: str | None, metadata: dict[str, Any] | None) -> bool: +def _is_nested_pregel_start( + name: str | None, + metadata: dict[str, Any] | None, + parent_run_id: UUID | None, + task_run_ids: set[UUID], +) -> bool: """Recognize a nested `Pregel` invocation from its `on_chain_start` metadata. When a compiled graph is added as a node, pregel fires two @@ -31,12 +36,21 @@ def _is_nested_pregel_start(name: str | None, metadata: dict[str, Any] | None) - the inner `Pregel` chain (whose `name` is the graph's `name`, not the node name). Both share the same `langgraph_checkpoint_ns`. - A nested `Pregel` start is therefore identified by having a - `langgraph_checkpoint_ns` AND a `name` that does NOT match the - owning task's `langgraph_node`. Regular node chains are skipped; - the root `Pregel` (which has no `langgraph_node` metadata) isn't - observed by this handler because the root's start fires before the - handler is attached. + Primary signal: a `langgraph_checkpoint_ns` is set AND `name` + differs from the owning task's `langgraph_node`. This covers the + common case where the compiled subgraph's name differs from the + node name it was registered under. + + Fallback for name collisions (subgraph compiled with + `name == node_name`): the inner `Pregel` start's `parent_run_id` + is the run_id of the node chain's start event, which the handler + records in `task_run_ids` on the first start. Matching + `parent_run_id` to that set identifies the second start as the + nested `Pregel` even when names coincide. + + Regular node chains are skipped; the root `Pregel` (which has no + `langgraph_node` metadata) isn't observed by this handler because + the root's start fires before the handler is attached. Metadata-based detection is used because `on_chain_start`'s `serialized` argument is `None` for compiled graphs in this @@ -48,6 +62,13 @@ def _is_nested_pregel_start(name: str | None, metadata: dict[str, Any] | None) - and the router function's name as `name`, which would otherwise match the discriminator without representing an actual nested `Pregel`. + + Args: + name: The `name` kwarg from `on_chain_start`. + metadata: The `metadata` kwarg from `on_chain_start`. + parent_run_id: The `parent_run_id` kwarg from `on_chain_start`. + task_run_ids: The set of run_ids the handler has already seen + as node-chain starts (i.e. `name == langgraph_node`). """ if not metadata: return False @@ -56,7 +77,12 @@ def _is_nested_pregel_start(name: str | None, metadata: dict[str, Any] | None) - lg_node = metadata.get("langgraph_node") if lg_node is None or lg_node in _LANGGRAPH_SENTINEL_NODES: return False - return name != lg_node + if name != lg_node: + return True + # Name collision fallback: the inner Pregel's parent_run_id is + # the node chain's run_id, which we recorded when that node + # chain's start fired. + return parent_run_id is not None and parent_run_id in task_run_ids class StreamLifecycleHandler(BaseCallbackHandler, _StreamingCallbackHandler): @@ -98,6 +124,10 @@ class StreamLifecycleHandler(BaseCallbackHandler, _StreamingCallbackHandler): self._pending_running: set[tuple[str, ...]] = set() # run_id → subgraph namespace; populated only for Pregel chains. self._run_to_ns: dict[UUID, tuple[str, ...]] = {} + # run_ids of node-chain starts (name == langgraph_node); used + # as the parent_run_id fallback when a subgraph's name equals + # its node name. Cleared as each chain ends. + self._task_run_ids: set[UUID] = set() root_payload: dict[str, Any] = {"event": "started"} if root_graph_name is not None: @@ -193,7 +223,21 @@ class StreamLifecycleHandler(BaseCallbackHandler, _StreamingCallbackHandler): self._fire_running_if_pending(containing) name = cast(str | None, kwargs.get("name")) - if not _is_nested_pregel_start(name, metadata): + lg_node = (metadata or {}).get("langgraph_node") + + # Record node-chain starts so the name-collision fallback in + # `_is_nested_pregel_start` can match the inner Pregel's + # parent_run_id to them. + if ( + lg_node is not None + and lg_node not in _LANGGRAPH_SENTINEL_NODES + and name == lg_node + ): + self._task_run_ids.add(run_id) + + if not _is_nested_pregel_start( + name, metadata, parent_run_id, self._task_run_ids + ): return ns = self._subgraph_ns_from_metadata(metadata) @@ -218,6 +262,7 @@ class StreamLifecycleHandler(BaseCallbackHandler, _StreamingCallbackHandler): parent_run_id: UUID | None = None, **kwargs: Any, ) -> Any: + self._task_run_ids.discard(run_id) ns = self._run_to_ns.pop(run_id, None) if ns is None: return @@ -235,6 +280,7 @@ class StreamLifecycleHandler(BaseCallbackHandler, _StreamingCallbackHandler): parent_run_id: UUID | None = None, **kwargs: Any, ) -> Any: + self._task_run_ids.discard(run_id) ns = self._run_to_ns.pop(run_id, None) if ns is None: return diff --git a/libs/langgraph/langgraph/pregel/remote.py b/libs/langgraph/langgraph/pregel/remote.py index 8919be4cd..6dd8a0346 100644 --- a/libs/langgraph/langgraph/pregel/remote.py +++ b/libs/langgraph/langgraph/pregel/remote.py @@ -650,16 +650,27 @@ class RemoteGraph(PregelProtocol): """ updated_stream_modes: list[StreamModeSDK] = [] req_single = True + # `"lifecycle"` is emitted locally by the `StreamLifecycleHandler` + # attached inside `Pregel.stream` / `astream`. The remote graph + # API has no corresponding mode, so requests for it against a + # `RemoteGraph` are silently stripped here and a warning is + # logged so the caller isn't left wondering why no lifecycle + # events arrive. + dropped_lifecycle = False # coerce to list, or add default stream mode if stream_mode: if isinstance(stream_mode, str): if stream_mode != "lifecycle": updated_stream_modes.append(cast(StreamModeSDK, stream_mode)) + else: + dropped_lifecycle = True else: req_single = False - updated_stream_modes.extend( - cast(StreamModeSDK, m) for m in stream_mode if m != "lifecycle" - ) + for m in stream_mode: + if m == "lifecycle": + dropped_lifecycle = True + else: + updated_stream_modes.append(cast(StreamModeSDK, m)) else: updated_stream_modes.append(default) # type: ignore[arg-type] requested_stream_modes = updated_stream_modes.copy() @@ -668,8 +679,16 @@ class RemoteGraph(PregelProtocol): (config or {}).get(CONF, {}).get(CONFIG_KEY_STREAM) ) if stream: - updated_stream_modes.extend( - cast(StreamModeSDK, m) for m in stream.modes if m != "lifecycle" + for m in stream.modes: + if m == "lifecycle": + dropped_lifecycle = True + else: + updated_stream_modes.append(cast(StreamModeSDK, m)) + if dropped_lifecycle: + logger.warning( + "Stream mode 'lifecycle' is not supported by RemoteGraph " + "and was stripped from the request; no lifecycle events " + "will be emitted for this remote run." ) # map "messages" to "messages-tuple" if "messages" in updated_stream_modes: diff --git a/libs/langgraph/langgraph/stream/run_stream.py b/libs/langgraph/langgraph/stream/run_stream.py index fb3152ff5..f12c642f9 100644 --- a/libs/langgraph/langgraph/stream/run_stream.py +++ b/libs/langgraph/langgraph/stream/run_stream.py @@ -28,21 +28,17 @@ async def _adrive_until_done(pump: Callable[[], Awaitable[bool]]) -> None: class BaseRunStream: """Shared shape for any object that wraps a `StreamMux`. - Both the root `GraphRunStream` / `AsyncGraphRunStream` and the - scoped `SubgraphRunStream` compose a `StreamMux` and expose its - projections (`values`, `messages`, `subgraphs`, user-registered - keys). The root additionally owns the graph iterator and drives - the pump; a subgraph's mini-mux borrows the root pump via - `make_child`'s pump inheritance. + Root (`GraphRunStream` / `AsyncGraphRunStream`) and scoped + (`SubgraphRunStream`) streams both compose a `StreamMux`. The mux + owns the projections — `values`, `messages`, `subgraphs`, and any + user-registered keys — all exposed via `extensions`. Native + projections (`_native = True`) are also bound as direct attributes + (`run.values`, `run.messages`, …) for ergonomics. - Projections registered on the mux show up in `extensions`, and - native ones (those with `_native = True`) are also bound directly - as attributes (`run.values`, `run.messages`, …). - - Raw iteration (`for event in run:` / `async for event in run`) - yields every `ProtocolEvent` that reached this mux's main log — - for the root that's every event in the run, for a subgraph that's - every event forwarded into its subtree. + Raw iteration (`for event in run` / `async for event in run`) and + the `interleave(...)` helper both live here so every subclass + behaves consistently. Subclasses only add pump ownership, scope + metadata, or sync/async flavor. """ def __init__(self, mux: StreamMux) -> None: diff --git a/libs/langgraph/langgraph/stream/streaming_handler.py b/libs/langgraph/langgraph/stream/streaming_handler.py index c0c2c265e..bbc764dfe 100644 --- a/libs/langgraph/langgraph/stream/streaming_handler.py +++ b/libs/langgraph/langgraph/stream/streaming_handler.py @@ -19,31 +19,32 @@ from langgraph.types import All, StreamMode def _coerce_factories( - transformers: list[StreamTransformer | TransformerFactory] | None, + transformers: list[TransformerFactory] | None, ) -> list[TransformerFactory]: - """Normalize caller-supplied transformers into scope-taking factories. + """Validate caller-supplied factories. - Accepts already-built instances (wrapped as single-use factories, - with the caveat that they won't be re-instantiated in subgraph - mini-muxes) or proper factories (classes / callables taking a - scope). The built-in root transformers are always factories so - they propagate into every subgraph scope automatically. + Each factory must be callable — a transformer class (which accepts + a positional `scope` argument) or a callable that returns a fresh + instance per scope. Already-built instances are rejected because + they can't be re-instantiated in subgraph mini-muxes, which would + silently disable per-subagent scoping for that transformer. """ - - def _wrap_instance(t: StreamTransformer) -> TransformerFactory: - # Single-use: only wires at root scope. A user that wants - # subgraph propagation should pass the class (or a lambda). - def _factory(_scope: tuple[str, ...]) -> StreamTransformer: - return t - - return _factory - coerced: list[TransformerFactory] = [] for item in transformers or (): if isinstance(item, StreamTransformer): - coerced.append(_wrap_instance(item)) - else: - coerced.append(item) + raise TypeError( + "StreamingHandler.transformers takes factories, not " + "pre-built instances. Pass the transformer class " + "(e.g. `MyTransformer`) or a callable taking `scope` " + "(e.g. `lambda scope: MyTransformer(scope, foo=...)`), " + "so fresh instances can be built for each subgraph." + ) + if not callable(item): + raise TypeError( + f"StreamingHandler.transformers entries must be callable; " + f"got {type(item).__name__}." + ) + coerced.append(item) return coerced @@ -121,7 +122,7 @@ class StreamingHandler: *, interrupt_before: All | Sequence[str] | None = None, interrupt_after: All | Sequence[str] | None = None, - transformers: list[StreamTransformer | TransformerFactory] | None = None, + transformers: list[TransformerFactory] | None = None, ) -> GraphRunStream: """Start a sync streaming run. @@ -169,7 +170,7 @@ class StreamingHandler: *, interrupt_before: All | Sequence[str] | None = None, interrupt_after: All | Sequence[str] | None = None, - transformers: list[StreamTransformer | TransformerFactory] | None = None, + transformers: list[TransformerFactory] | None = None, ) -> AsyncGraphRunStream: """Start an async streaming run. diff --git a/libs/langgraph/langgraph/stream/transformers.py b/libs/langgraph/langgraph/stream/transformers.py index aca437cee..8bc53ab1f 100644 --- a/libs/langgraph/langgraph/stream/transformers.py +++ b/libs/langgraph/langgraph/stream/transformers.py @@ -1,5 +1,6 @@ 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 @@ -21,6 +22,9 @@ if TYPE_CHECKING: from langgraph.stream._mux import StreamMux +logger = logging.getLogger(__name__) + + SubgraphStatus = Literal["started", "running", "completed", "failed", "interrupted"] _TERMINAL_STATUSES: frozenset[SubgraphStatus] = frozenset( {"completed", "failed", "interrupted"} @@ -256,14 +260,9 @@ class SubgraphRunStream(BaseRunStream): the same transformer factories as the root mux, so `.values`, `.messages`, `.subgraphs` are populated by the standard transformers scoped to this handle's namespace — no duplicated - routing logic. - - Inherits `BaseRunStream`, so all shared shape works uniformly: - raw iteration (`for event in sub`), `.interleave(...)`, native - projection attributes, and the `extensions` mapping — identical - to the root `GraphRunStream`. The mini-mux borrows the root's - pump via `make_child`'s pump inheritance, so any cursor on a - subagent projection drives the whole run forward. + routing logic. The mini-mux borrows the root's pump via + `make_child`'s pump inheritance, so any cursor on a subagent + projection drives the whole run forward. Lifecycle fields update in place as events arrive: @@ -368,8 +367,11 @@ class SubgraphTransformer(StreamTransformer): if data.get("event") == "started": self._on_started(ns, data) - # 2. Forward the event to the matching direct-child mini-mux. - # Prefix-match: ns must start with some child's path. + # 2. Forward the event to the matching direct-child mini-mux + # before the status-change step below so that terminal events + # reach the child's log and grandchild transformers *before* + # the child's mini-mux is closed. Prefix-match: ns must start + # with some child's path. direct_child_ns = ns[: depth + 1] if len(ns) > depth else None if direct_child_ns is not None and direct_child_ns in self._by_ns: self._by_ns[direct_child_ns]._mux.push(event) @@ -394,9 +396,13 @@ class SubgraphTransformer(StreamTransformer): if ns in self._by_ns: # Duplicate started — ignore. return - if self._mux is None: - # Not registered yet; can't build a mini-mux. - return + # `_on_register` is called by the mux during registration, which + # happens before any event can be dispatched — so this should + # always be set by the time we process an event. + assert self._mux is not None, ( + "SubgraphTransformer processed an event before _on_register; " + "transformer registration ordering is broken." + ) child_mux = self._mux.make_child(ns) handle = SubgraphRunStream( path=ns, @@ -432,7 +438,12 @@ class SubgraphTransformer(StreamTransformer): try: handle._mux.close() except Exception: - pass + logger.warning( + "Error closing subgraph mini-mux at %s; subscribers " + "may not see a clean close.", + handle.path, + exc_info=True, + ) def finalize(self) -> None: """Transition any still-open direct children to `completed`.""" @@ -455,4 +466,9 @@ class SubgraphTransformer(StreamTransformer): try: handle._mux.fail(err) except Exception: - pass + 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/tests/test_stream_subgraph_transformer.py b/libs/langgraph/tests/test_stream_subgraph_transformer.py index 67fa25826..b4d7dadcb 100644 --- a/libs/langgraph/tests/test_stream_subgraph_transformer.py +++ b/libs/langgraph/tests/test_stream_subgraph_transformer.py @@ -7,6 +7,7 @@ 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 @@ -22,6 +23,7 @@ from langgraph.stream.transformers import ( SubgraphTransformer, ValuesTransformer, ) +from langgraph.types import interrupt TS = int(time.time() * 1000) @@ -365,3 +367,95 @@ class TestSubgraphTransformerAsyncEndToEnd: assert len(collected) == 1 child = collected[0] assert child.status == "completed" + + +class TestSubgraphTriggerCallId: + """Confirm `trigger_call_id` flows from real pregel metadata.""" + + def test_trigger_call_id_populated_end_to_end(self) -> None: + graph = _build_nested_graph() + handler = StreamingHandler(graph) + run = handler.stream({"value": "", "items": []}) + + collected: list[SubgraphRunStream] = list(run.subgraphs) + assert len(collected) == 1 + child = collected[0] + + # The child's single-segment path encodes `node_name:task_id`. + # Both the parsed task_id (`trigger_call_id`) and the segment + # should match the same task_id suffix. + assert ":" in child.path[0] + node_name, _, task_id = child.path[0].partition(":") + assert node_name == "sub" + assert task_id # non-empty + assert child.trigger_call_id == task_id + + +class TestSubgraphInterrupt: + """Interrupts raised inside a subgraph surface as status=interrupted.""" + + def _build_interrupt_subgraph(self): + def inner_node(state: SimpleState) -> dict: + interrupt("need approval") + return {"value": state["value"] + "X", "items": ["x"]} + + inner_builder = StateGraph(SimpleState) + inner_builder.add_node("inner_node", inner_node) + inner_builder.add_edge(START, "inner_node") + inner_builder.add_edge("inner_node", END) + inner = inner_builder.compile() + + outer_builder = StateGraph(SimpleState) + outer_builder.add_node("sub", inner) + outer_builder.add_edge(START, "sub") + outer_builder.add_edge("sub", END) + return outer_builder.compile(checkpointer=InMemorySaver()) + + def test_interrupt_in_subgraph_marks_handle_interrupted(self) -> None: + graph = self._build_interrupt_subgraph() + handler = StreamingHandler(graph) + run = handler.stream( + {"value": "", "items": []}, + config={"configurable": {"thread_id": "t1"}}, + ) + + collected: list[SubgraphRunStream] = list(run.subgraphs) + + assert run.interrupted is True + assert len(collected) == 1 + assert collected[0].status == "interrupted" + + +class TestSubgraphNameCollision: + """The subgraph's compiled `name` equaling its node name is detected. + + Primary detector `name != langgraph_node` fails here; the + parent_run_id fallback in `_is_nested_pregel_start` is what keeps + the subgraph visible. + """ + + def test_name_equals_node_name_still_detected(self) -> None: + def inner_node(state: SimpleState) -> dict: + return {"value": state["value"] + "X", "items": ["x"]} + + inner_builder = StateGraph(SimpleState) + inner_builder.add_node("inner_node", inner_node) + inner_builder.add_edge(START, "inner_node") + inner_builder.add_edge("inner_node", END) + # Compile with the same name as the node it will be registered as. + inner = inner_builder.compile(name="sub") + + outer_builder = StateGraph(SimpleState) + outer_builder.add_node("sub", inner) + outer_builder.add_edge(START, "sub") + outer_builder.add_edge("sub", END) + graph = outer_builder.compile() + + handler = StreamingHandler(graph) + run = handler.stream({"value": "", "items": []}) + + collected: list[SubgraphRunStream] = list(run.subgraphs) + assert len(collected) == 1 + child = collected[0] + assert child.graph_name == "sub" + assert child.status == "completed" diff --git a/libs/langgraph/tests/test_streaming_handler.py b/libs/langgraph/tests/test_streaming_handler.py index 4e40c74fd..a70181589 100644 --- a/libs/langgraph/tests/test_streaming_handler.py +++ b/libs/langgraph/tests/test_streaming_handler.py @@ -1107,7 +1107,10 @@ class TestCustomTransformer: graph = _build_simple_graph() handler = StreamingHandler(graph) counter_t = CounterTransformer() - run = handler.stream({"value": "x", "items": []}, transformers=[counter_t]) + run = handler.stream( + {"value": "x", "items": []}, + transformers=[lambda _scope: counter_t], + ) assert "counter" in run.extensions # Subscribe before driving the run so channel pushes are retained. counter_iter = iter(run.extensions["counter"]) @@ -1138,7 +1141,10 @@ class TestCustomTransformer: graph = _build_simple_graph() handler = StreamingHandler(graph) foo_t = FooTransformer() - run = handler.stream({"value": "x", "items": []}, transformers=[foo_t]) + run = handler.stream( + {"value": "x", "items": []}, + transformers=[lambda _scope: foo_t], + ) # Subscribe before driving the run. foo_iter = iter(run.foo) _ = run.output @@ -1168,7 +1174,8 @@ class TestCustomTransformer: graph = _build_simple_graph() handler = StreamingHandler(graph) run = handler.stream( - {"value": "x", "items": []}, transformers=[EmitterTransformer()] + {"value": "x", "items": []}, + transformers=[lambda _scope: EmitterTransformer()], ) events = list(run) custom_events = [e for e in events if e["method"] == "custom:emitter"] @@ -1232,7 +1239,7 @@ class TestCustomTransformer: ): handler.stream( {"value": "x", "items": []}, - transformers=[ConflictTransformer()], + transformers=[lambda _scope: ConflictTransformer()], ) @@ -1329,7 +1336,10 @@ class TestEventLogAutoLifecycle: graph = _build_simple_graph() handler = StreamingHandler(graph) t = MinimalTransformer() - run = handler.stream({"value": "x", "items": []}, transformers=[t]) + run = handler.stream( + {"value": "x", "items": []}, + transformers=[lambda _scope: t], + ) minimal_iter = iter(run.extensions["minimal"]) _ = run.output items = list(minimal_iter) @@ -1629,9 +1639,10 @@ class TestAsyncTransformerLane: graph = _build_simple_graph() handler = StreamingHandler(graph) + scorer = Scorer() run = await handler.astream( {"value": "x", "items": []}, - transformers=[Scorer()], + transformers=[lambda _scope: scorer], ) # Subscribe before driving the run so scheduled pushes are retained. scores_cursor = aiter(run.extensions["scores"])