chore: import upstream snapshot with attribution

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:")