Address review feedback on subgraph lifecycle streaming

- Robustify nested-Pregel detection with a parent_run_id fallback so subgraphs compiled with name equal to their node name are still recognized. - Narrow bare excepts in SubgraphTransformer close/fail paths; log at warning with exc_info instead of silently swallowing. - Assert mux registration in SubgraphTransformer._on_started instead of silently dropping events. - Warn when RemoteGraph strips an unsupported "lifecycle" stream mode so callers aren't left wondering why no events arrive. - Reject pre-built transformer instances in StreamingHandler; factories are required so transformers propagate into every subgraph scope. - Comment the forward-before-close ordering in SubgraphTransformer. - Trim duplicated pump/projection docstrings across run-stream classes. - Add end-to-end tests for trigger_call_id, subgraph interrupt, and the name-collision detection fallback.
2026-09-06 01:37:49 +02:00 · 2026-04-18 18:59:23 -04:00
parent 2ad30132a3
commit bebcd20815
7 changed files with 253 additions and 70 deletions
@@ -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
@@ -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:
@@ -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:
@@ -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.

@@ -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,
+                    )
@@ -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"
@@ -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"])