fix(langgraph): keep only name, description and forked in __subagents

libs/langgraph/langgraph/graph/message.py

@@ -22,7 +22,7 @@ from langchain_core.messages import (
 )
 from typing_extensions import TypedDict, deprecated
 
-from langgraph._internal._constants import CONF, CONFIG_KEY_SEND
+from langgraph._internal._constants import CONF, CONFIG_KEY_SEND, NS_SEP
 from langgraph.graph.state import StateGraph
 from langgraph.pregel._messages import (
     _messages_delta_reducer as _messages_delta_reducer,
@@ -343,7 +343,7 @@ def push_message(
     )
 
     from langgraph.config import get_config
-    from langgraph.pregel._messages import StreamMessagesHandler, graph_namespace
+    from langgraph.pregel._messages import StreamMessagesHandler
 
     config = get_config()
     message = next(x for x in convert_to_messages([message]))
@@ -359,12 +359,15 @@ def push_message(
     ):
         handlers = config["callbacks"]
 
-    metadata = config["metadata"]
-    for handler in handlers:
-        if isinstance(handler, StreamMessagesHandler):
-            ns = graph_namespace(metadata)
-            if handler.subgraphs or not ns or ns == handler.parent_ns:
-                handler._emit((ns, metadata), message, dedupe=False)
+    if stream_handler := next(
+        (x for x in handlers if isinstance(x, StreamMessagesHandler)), None
+    ):
+        metadata = config["metadata"]
+        message_meta = (
+            tuple(cast(str, metadata["langgraph_checkpoint_ns"]).split(NS_SEP)),
+            metadata,
+        )
+        stream_handler._emit(message_meta, message, dedupe=False)
 
     if state_key:
         config[CONF][CONFIG_KEY_SEND]([(state_key, message)])

libs/langgraph/langgraph/pregel/_draw.py

@@ -232,6 +232,8 @@ def draw_graph(
         tools_by_name = getattr(node.bound, "tools_by_name", None)
         if isinstance(tools_by_name, Mapping) and tools_by_name:
             metadata["__tools"] = list(tools_by_name)
+            if subagents := _get_tool_subagents(tools_by_name):
+                metadata["__subagents"] = subagents
         graph.add_node(node.bound, name, metadata=metadata or None)
     # add start node
     if START not in nodes:
@@ -281,6 +283,38 @@ def draw_graph(
     return graph
 
 
+def _get_tool_subagents(
+    tools_by_name: Mapping[str, Any],
+) -> dict[str, list[dict[str, Any]]]:
+    """Collect the subagents each tool can run, keyed by tool name."""
+    subagents: dict[str, list[dict[str, Any]]] = {}
+    for tool_name, tool in tools_by_name.items():
+        tool_subagents = getattr(tool, "subagents", None)
+        if (
+            isinstance(tool_subagents, Sequence)
+            and not isinstance(tool_subagents, str)
+            and tool_subagents
+            and all(
+                isinstance(subagent, Mapping) and isinstance(subagent.get("name"), str)
+                for subagent in tool_subagents
+            )
+        ):
+            subagents[tool_name] = [
+                _subagent_entry(subagent) for subagent in tool_subagents
+            ]
+    return subagents
+
+
+def _subagent_entry(subagent: Mapping[str, Any]) -> dict[str, Any]:
+    """Keep only the subagent fields graph viewers display."""
+    entry: dict[str, Any] = {"name": subagent["name"]}
+    if isinstance(description := subagent.get("description"), str):
+        entry["description"] = description
+    if subagent.get("forked") is True:
+        entry["forked"] = True
+    return entry
+
+
 def add_edge(
     graph: Graph,
     source: str,

libs/langgraph/langgraph/pregel/_messages.py

@@ -23,7 +23,7 @@ from langchain_core.outputs import ChatGeneration, ChatGenerationChunk, LLMResul
 from pydantic import BaseModel
 
 from langgraph._internal._config import filter_to_user_tags
-from langgraph._internal._constants import NS_SEP, PUSH
+from langgraph._internal._constants import NS_SEP
 from langgraph.channels.binop import _get_overwrite
 from langgraph.channels.delta import DeltaChannel
 from langgraph.constants import TAG_HIDDEN, TAG_NOSTREAM
@@ -55,23 +55,6 @@ def _state_values(obj: Any) -> Sequence[Any]:
     return ()
 
 
-def graph_namespace(metadata: Mapping[str, Any]) -> tuple[str, ...]:
-    """Namespace of the graph a task runs in, from the task's run metadata.
-
-    Drops the task's own segment, plus one per enclosing `@task` call: a call
-    runs in its caller's graph but nests its segment under the caller's. A
-    call's `langgraph_path` is `(PUSH, <caller's path>, <index>, True)`, while
-    a `Send` task's ends in `False`.
-    """
-    ns = tuple(cast(str, metadata["langgraph_checkpoint_ns"]).split(NS_SEP))
-    path = metadata.get("langgraph_path", ())
-    drop = 1
-    while len(path) == 4 and path[0] == PUSH and path[-1] is True:
-        drop += 1
-        path = path[1]
-    return ns[:-drop]
-
-
 class StreamMessagesHandler(BaseCallbackHandler, _StreamingCallbackHandler):
     """A callback handler that implements stream_mode=messages.
 
@@ -165,7 +148,9 @@ class StreamMessagesHandler(BaseCallbackHandler, _StreamingCallbackHandler):
         **kwargs: Any,
     ) -> Any:
         if metadata and (not tags or (TAG_NOSTREAM not in tags)):
-            ns = graph_namespace(metadata)
+            ns = tuple(cast(str, metadata["langgraph_checkpoint_ns"]).split(NS_SEP))[
+                :-1
+            ]
             if not self.subgraphs and len(ns) > 0 and ns != self.parent_ns:
                 return
             if (filtered_tags := filter_to_user_tags(tags)) is not None:
@@ -228,8 +213,10 @@ class StreamMessagesHandler(BaseCallbackHandler, _StreamingCallbackHandler):
             and kwargs.get("name") == metadata.get("langgraph_node")
             and (not tags or TAG_HIDDEN not in tags)
         ):
-            ns = graph_namespace(metadata)
-            if not self.subgraphs and len(ns) > 0 and ns != self.parent_ns:
+            ns = tuple(cast(str, metadata["langgraph_checkpoint_ns"]).split(NS_SEP))[
+                :-1
+            ]
+            if not self.subgraphs and len(ns) > 0:
                 return
             self.metadata[run_id] = (ns, metadata)
             for value in _state_values(inputs):

libs/langgraph/langgraph/pregel/_tools.py

@@ -7,9 +7,9 @@ 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.errors import GraphDrained, GraphInterrupt
-from langgraph.pregel._messages import graph_namespace
 from langgraph.pregel.protocol import StreamChunk
 
 try:
@@ -92,9 +92,12 @@ class StreamToolCallHandler(BaseCallbackHandler, _StreamingCallbackHandler):
     ) -> tuple[str, ...] | None:
         """Resolve the namespace this tool call should emit at, or `None` to skip.
 
-        Mirrors `StreamMessagesHandler.on_chat_model_start`: the namespace
-        is that of the graph the tool runs in (`graph_namespace`). Returns
-        `None` when the call should be silently suppressed:
+        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.
@@ -107,9 +110,11 @@ class StreamToolCallHandler(BaseCallbackHandler, _StreamingCallbackHandler):
             return None
         if tags and TAG_NOSTREAM in tags:
             return None
-        ns = (
-            graph_namespace(metadata) if metadata.get("langgraph_checkpoint_ns") else ()
-        )
+        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

libs/langgraph/langgraph/pregel/main.py

@@ -2880,20 +2880,6 @@ class Pregel(
                     for h in callback_manager.inheritable_handlers
                     if not isinstance(h, StreamMessagesHandlerV2)
                 ]
-            if CONFIG_KEY_TASK_ID not in config[CONF]:
-                # Not run as a task (a child invoked with its own thread, say),
-                # this graph's namespace starts at the root, so a caller's
-                # handler would take this graph's events for its own graph's.
-                callback_manager.handlers = [
-                    h
-                    for h in callback_manager.handlers
-                    if not _streams_only_its_graph(h)
-                ]
-                callback_manager.inheritable_handlers = [
-                    h
-                    for h in callback_manager.inheritable_handlers
-                    if not _streams_only_its_graph(h)
-                ]
             if "ls_integration" not in callback_manager.metadata:
                 callback_manager.add_metadata({"ls_integration": "langgraph"})
             run_manager = callback_manager.on_chain_start(
@@ -2912,9 +2898,9 @@ class Pregel(
                 )
             # set up subgraph checkpointing
             config = self._own_checkpoint_config(config)
-            stream_handlers: list[StreamMessagesHandler | StreamToolCallHandler] = []
             # 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)
                 )
@@ -2923,12 +2909,24 @@ class Pregel(
                     if use_stream_messages_v2
                     else StreamMessagesHandler
                 )
-                stream_handlers.append(messages_handler_cls(stream.put, subgraphs))
+                run_manager.inheritable_handlers.append(
+                    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:
-                stream_handlers.append(StreamToolCallHandler(stream.put, subgraphs))
-            run_manager.inheritable_handlers.extend(stream_handlers)
+                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:
@@ -3011,10 +3009,6 @@ class Pregel(
                 cache_policy=self.cache_policy,
                 has_graph_lifecycle_callbacks=bool(graph_callback_manager.handlers),
             ) as loop:
-                # Known only now: the loop numbers the namespace of every
-                # subgraph a task runs after its first.
-                for handler in stream_handlers:
-                    handler.parent_ns = loop.checkpoint_ns
                 emit_graph_lifecycle_events(loop)
                 # create runner
                 runner = PregelRunner(
@@ -3297,20 +3291,6 @@ class Pregel(
                     for h in callback_manager.inheritable_handlers
                     if not isinstance(h, StreamMessagesHandlerV2)
                 ]
-            if CONFIG_KEY_TASK_ID not in config[CONF]:
-                # Not run as a task (a child invoked with its own thread, say),
-                # this graph's namespace starts at the root, so a caller's
-                # handler would take this graph's events for its own graph's.
-                callback_manager.handlers = [
-                    h
-                    for h in callback_manager.handlers
-                    if not _streams_only_its_graph(h)
-                ]
-                callback_manager.inheritable_handlers = [
-                    h
-                    for h in callback_manager.inheritable_handlers
-                    if not _streams_only_its_graph(h)
-                ]
             if "ls_integration" not in callback_manager.metadata:
                 callback_manager.add_metadata({"ls_integration": "langgraph"})
             run_manager = await callback_manager.on_chain_start(
@@ -3345,9 +3325,10 @@ class Pregel(
                 )
             # set up subgraph checkpointing
             config = self._own_checkpoint_config(config)
-            stream_handlers: list[StreamMessagesHandler | StreamToolCallHandler] = []
             # set up messages stream mode
             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)
                 )
@@ -3356,12 +3337,24 @@ class Pregel(
                     if use_stream_messages_v2
                     else StreamMessagesHandler
                 )
-                stream_handlers.append(messages_handler_cls(stream_put, subgraphs))
+                run_manager.inheritable_handlers.append(
+                    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:
-                stream_handlers.append(StreamToolCallHandler(stream_put, subgraphs))
-            run_manager.inheritable_handlers.extend(stream_handlers)
+                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:
@@ -3470,10 +3463,6 @@ class Pregel(
                 cache_policy=self.cache_policy,
                 has_graph_lifecycle_callbacks=bool(graph_callback_manager.handlers),
             ) as loop:
-                # Known only now: the loop numbers the namespace of every
-                # subgraph a task runs after its first.
-                for handler in stream_handlers:
-                    handler.parent_ns = loop.checkpoint_ns
                 await aemit_graph_lifecycle_events(loop)
                 # create runner
                 runner = PregelRunner(
@@ -4471,13 +4460,6 @@ def _coerce_checkpoint_values(payload: Any, mapper: Callable[[Any], Any]) -> Non
         payload["values"] = mapper(payload["values"])
 
 
-def _streams_only_its_graph(handler: object) -> bool:
-    return (
-        isinstance(handler, (StreamMessagesHandler, StreamToolCallHandler))
-        and not handler.subgraphs
-    )
-
-
 def _resolve_parent_ns(
     graph_config: RunnableConfig | None, call_config: RunnableConfig | None
 ) -> tuple[str, ...]:

libs/langgraph/tests/test_pregel.py

@@ -36,7 +36,7 @@ from langchain_core.runnables import (
     RunnablePassthrough,
 )
 from langchain_core.runnables.graph import Edge
-from langchain_core.tools import tool
+from langchain_core.tools import StructuredTool, tool
 from langchain_core.version import VERSION as LANGCHAIN_CORE_VERSION
 from langgraph.cache.base import BaseCache
 from langgraph.checkpoint.base import (
@@ -3105,6 +3105,49 @@ def test_simple_multi_edge(snapshot: SnapshotAssertion) -> None:
     )
 
 
+def test_tool_node_lists_tool_subagents_in_graph_metadata() -> None:
+    class TaskTool(StructuredTool):
+        subagents: list[dict[str, Any]] = []
+
+    def run_task(description: str, subagent_type: str) -> str:
+        return subagent_type
+
+    task = TaskTool.from_function(run_task, name="task", description="Run a subagent.")
+    task.subagents = [
+        {"name": "researcher", "description": "Researches a topic."},
+        {
+            "name": "critic",
+            "description": "Reviews drafts.",
+            "forked": True,
+            "system_prompt": "You are a critic.",
+            "runnable": object(),
+        },
+        {"name": "planner", "description": None, "forked": "yes"},
+    ]
+
+    @tool
+    def search(query: str) -> str:
+        """Search the web."""
+        return query
+
+    builder = StateGraph(MessagesState)
+    builder.add_node("agent", lambda state: state)
+    builder.add_node("tools", ToolNode([search, task]))
+    builder.add_edge(START, "agent")
+    builder.add_edge("agent", "tools")
+    graph = builder.compile().get_graph().to_json()
+    nodes = {node["id"]: node for node in graph["nodes"]}
+
+    assert nodes["tools"]["metadata"]["__subagents"] == {
+        "task": [
+            {"name": "researcher", "description": "Researches a topic."},
+            {"name": "critic", "description": "Reviews drafts.", "forked": True},
+            {"name": "planner"},
+        ]
+    }
+    assert "__subagents" not in (nodes["agent"].get("metadata") or {})
+
+
 def test_nested_graph_xray(snapshot: SnapshotAssertion) -> None:
     class State(TypedDict):
         my_key: Annotated[str, operator.add]

libs/langgraph/tests/test_stream_namespaces.py

@@ -1,291 +0,0 @@
-from typing import Any
-
-import pytest
-from langchain_core.language_models.fake_chat_models import GenericFakeChatModel
-from langchain_core.messages import AIMessage
-from langchain_core.tools import tool
-
-from langgraph.func import entrypoint, task
-from langgraph.graph import START, MessagesState, StateGraph
-from langgraph.graph.message import push_message
-from langgraph.graph.state import CompiledStateGraph
-
-pytestmark = pytest.mark.anyio
-
-
-@tool
-def echo(text: str) -> str:
-    """Echo the text."""
-    return text
-
-
-def _one_node_graph(node: Any) -> CompiledStateGraph:
-    return (
-        StateGraph(MessagesState)
-        .add_node(node)
-        .add_edge(START, node.__name__)
-        .compile()
-    )
-
-
-def _chat_and_echo(state: MessagesState) -> dict:
-    echo.invoke({"text": "done"})
-    return {"messages": [GenericFakeChatModel(messages=iter(["one two"])).invoke("hi")]}
-
-
-async def _achat_and_echo(state: MessagesState) -> dict:
-    await echo.ainvoke({"text": "done"})
-    model = GenericFakeChatModel(messages=iter(["one two"]))
-    return {"messages": [await model.ainvoke("hi")]}
-
-
-def _segment_names(ns: tuple[str, ...]) -> tuple[str, ...]:
-    return tuple(segment.split(":")[0] for segment in ns)
-
-
-def _summary(part: dict) -> tuple[tuple[str, ...], str]:
-    if part["type"] == "messages":
-        return _segment_names(part["ns"]), part["data"][0].content
-    return _segment_names(part["ns"]), part["data"]["event"]
-
-
-_CHAT_AND_ECHO_EVENTS = [
-    ((), "tool-started"),
-    ((), "tool-finished"),
-    ((), "one"),
-    ((), " "),
-    ((), "two"),
-]
-
-_CHILD_CONFIGS = pytest.mark.parametrize(
-    "child_config",
-    [
-        None,
-        {"configurable": {"thread_id": "child"}},
-        {"configurable": {"thread_id": "parent"}},
-    ],
-    ids=["no_config", "own_thread", "caller_thread"],
-)
-
-
-_TASK_CALLERS = pytest.mark.parametrize(
-    "nested", [False, True], ids=["from_entrypoint", "from_task_in_task"]
-)
-
-
-@_TASK_CALLERS
-@pytest.mark.parametrize("subgraphs", [False, True])
-def test_task_events_stream_under_the_graph_namespace(
-    nested: bool, subgraphs: bool
-) -> None:
-    @task
-    def chat_and_echo() -> None:
-        _chat_and_echo({"messages": []})
-
-    @task
-    def outer() -> None:
-        chat_and_echo().result()
-
-    @entrypoint()
-    def workflow(inputs: dict) -> None:
-        (outer() if nested else chat_and_echo()).result()
-
-    parts = workflow.stream(
-        {}, stream_mode=["messages", "tools"], subgraphs=subgraphs, version="v2"
-    )
-
-    assert [_summary(part) for part in parts] == _CHAT_AND_ECHO_EVENTS
-
-
-@_TASK_CALLERS
-@pytest.mark.parametrize("subgraphs", [False, True])
-async def test_task_events_stream_under_the_graph_namespace_async(
-    nested: bool, subgraphs: bool
-) -> None:
-    @task
-    async def chat_and_echo() -> None:
-        await _achat_and_echo({"messages": []})
-
-    @task
-    async def outer() -> None:
-        await chat_and_echo()
-
-    @entrypoint()
-    async def workflow(inputs: dict) -> None:
-        await (outer() if nested else chat_and_echo())
-
-    parts = workflow.astream(
-        {}, stream_mode=["messages", "tools"], subgraphs=subgraphs, version="v2"
-    )
-
-    assert [_summary(part) async for part in parts] == _CHAT_AND_ECHO_EVENTS
-
-
-def test_task_events_stream_under_the_namespace_of_their_subgraph() -> None:
-    @task
-    def chat_and_echo() -> None:
-        _chat_and_echo({"messages": []})
-
-    def node(state: MessagesState) -> dict:
-        chat_and_echo().result()
-        return {}
-
-    sub = _one_node_graph(node)
-    parent = StateGraph(MessagesState).add_node("sub", sub).add_edge(START, "sub")
-
-    parts = parent.compile().stream(
-        {"messages": []},
-        stream_mode=["messages", "tools"],
-        subgraphs=True,
-        version="v2",
-    )
-
-    assert [_summary(part) for part in parts] == [
-        (("sub",), event) for _, event in _CHAT_AND_ECHO_EVENTS
-    ]
-
-
-@_CHILD_CONFIGS
-def test_child_graph_streams_only_to_the_node_streaming_it(
-    child_config: dict | None,
-) -> None:
-    child = _one_node_graph(_chat_and_echo)
-    received = []
-
-    def node(state: MessagesState) -> dict:
-        received.extend(
-            child.stream(
-                {"messages": []}, child_config, stream_mode=["messages", "tools"]
-            )
-        )
-        return {}
-
-    parts = _one_node_graph(node).stream(
-        {"messages": []},
-        {"configurable": {"thread_id": "parent"}},
-        stream_mode=["messages", "tools"],
-    )
-
-    assert list(parts) == []
-    assert len(received) == len(_CHAT_AND_ECHO_EVENTS)
-
-
-@_CHILD_CONFIGS
-async def test_child_graph_streams_only_to_the_node_streaming_it_async(
-    child_config: dict | None,
-) -> None:
-    child = _one_node_graph(_achat_and_echo)
-    received = []
-
-    async def node(state: MessagesState) -> dict:
-        stream = child.astream(
-            {"messages": []}, child_config, stream_mode=["messages", "tools"]
-        )
-        received.extend([part async for part in stream])
-        return {}
-
-    parts = _one_node_graph(node).astream(
-        {"messages": []},
-        {"configurable": {"thread_id": "parent"}},
-        stream_mode=["messages", "tools"],
-    )
-
-    assert [part async for part in parts] == []
-    assert len(received) == len(_CHAT_AND_ECHO_EVENTS)
-
-
-def _reply(state: MessagesState) -> dict:
-    return {"messages": [AIMessage(id="reply", content="hi")]}
-
-
-def test_node_gets_the_messages_returned_in_the_graph_it_streams() -> None:
-    sub = _one_node_graph(_reply)
-    received = []
-
-    def node(state: MessagesState) -> dict:
-        for message, _ in sub.stream({"messages": []}, stream_mode="messages"):
-            received.append(message.id)
-        return {}
-
-    parts = _one_node_graph(node).stream({"messages": []}, stream_mode="messages")
-
-    assert list(parts) == []
-    assert received == ["reply"]
-
-
-def test_node_gets_the_events_of_every_subgraph_it_streams() -> None:
-    sub = _one_node_graph(_chat_and_echo)
-    counts = []
-
-    def node(state: MessagesState) -> dict:
-        for _ in range(2):
-            stream = sub.stream({"messages": []}, stream_mode=["messages", "tools"])
-            counts.append(len(list(stream)))
-        return {}
-
-    _one_node_graph(node).invoke({"messages": []})
-
-    assert counts == [len(_CHAT_AND_ECHO_EVENTS)] * 2
-
-
-async def test_node_gets_the_events_of_every_subgraph_it_streams_async() -> None:
-    sub = _one_node_graph(_achat_and_echo)
-    counts = []
-
-    async def node(state: MessagesState) -> dict:
-        for _ in range(2):
-            stream = sub.astream({"messages": []}, stream_mode=["messages", "tools"])
-            counts.append(len([part async for part in stream]))
-        return {}
-
-    await _one_node_graph(node).ainvoke({"messages": []})
-
-    assert counts == [len(_CHAT_AND_ECHO_EVENTS)] * 2
-
-
-def _push(state: MessagesState) -> dict:
-    push_message(AIMessage(id="pushed", content="hi"), state_key=None)
-    return {}
-
-
-def _namespaces(graph: CompiledStateGraph, subgraphs: bool) -> list[tuple[str, ...]]:
-    parts = graph.stream(
-        {"messages": []}, stream_mode="messages", subgraphs=subgraphs, version="v2"
-    )
-    return [_segment_names(part["ns"]) for part in parts]
-
-
-def _streams_pushing_graph(received: list) -> Any:
-    pushing = _one_node_graph(_push)
-
-    def node(state: MessagesState) -> dict:
-        for message, _ in pushing.stream({"messages": []}, stream_mode="messages"):
-            received.append(message.id)
-        return {}
-
-    return node
-
-
-def test_pushed_message_streams_under_the_namespace_of_its_graph() -> None:
-    pushing = _one_node_graph(_push)
-    parent = StateGraph(MessagesState).add_node("sub", pushing).add_edge(START, "sub")
-
-    assert _namespaces(pushing, subgraphs=True) == [()]
-    assert _namespaces(parent.compile(), subgraphs=False) == []
-    assert _namespaces(parent.compile(), subgraphs=True) == [("sub",)]
-
-
-def test_pushed_message_streams_to_the_node_that_streams_its_graph() -> None:
-    received: list[str] = []
-    graph = _one_node_graph(_streams_pushing_graph(received))
-
-    assert _namespaces(graph, subgraphs=False) == []
-    assert received == ["pushed"]
-
-
-def test_pushed_message_streams_to_every_stream_that_takes_it() -> None:
-    received: list[str] = []
-    graph = _one_node_graph(_streams_pushing_graph(received))
-
-    assert _namespaces(graph, subgraphs=True) == [("node",)]
-    assert received == ["pushed"]