x

Move unit tests for ToolNode into the tool node testing code

libs/prebuilt/Makefile

@@ -30,6 +30,9 @@ test_watch:
 	make stop-services; \
 	exit $$EXIT_CODE
 
+snapshot_upate:
+	LANGGRAPH_TEST_FAST=1 uv run pytest --snapshot-update $(TEST)
+
 ######################
 # LINTING AND FORMATTING
 ######################

libs/prebuilt/langgraph/prebuilt/tool_node.py

@@ -340,14 +340,22 @@ class ToolNode(RunnableCallable):
         self.tools_by_name: dict[str, BaseTool] = {}
         self.tool_to_state_args: dict[str, dict[str, Optional[str]]] = {}
         self.tool_to_store_arg: dict[str, Optional[str]] = {}
+        self.structured_output_tools: list[str] = []
         self.handle_tool_errors = handle_tool_errors
         self.messages_key = messages_key
         for tool_ in tools:
-            if not isinstance(tool_, BaseTool):
-                tool_ = create_tool(tool_)
-            self.tools_by_name[tool_.name] = tool_
-            self.tool_to_state_args[tool_.name] = _get_state_args(tool_)
-            self.tool_to_store_arg[tool_.name] = _get_store_arg(tool_)
+            if inspect.isclass(tool_) and issubclass(tool_, BaseModel):
+                # Handle Pydantic model classes as structured output tools
+                self.tools_by_name[tool_.__name__] = tool_
+                self.tool_to_state_args[tool_.__name__] = {}
+                self.tool_to_store_arg[tool_.__name__] = None
+                self.structured_output_tools.append(tool_.__name__)
+            else:
+                if not isinstance(tool_, BaseTool):
+                    tool_ = create_tool(tool_)
+                self.tools_by_name[tool_.name] = tool_
+                self.tool_to_state_args[tool_.name] = _get_state_args(tool_)
+                self.tool_to_store_arg[tool_.name] = _get_store_arg(tool_)
 
     def _func(
         self,
@@ -390,7 +398,7 @@ class ToolNode(RunnableCallable):
 
     def _combine_tool_outputs(
         self,
-        outputs: list[ToolMessage],
+        outputs: list[Union[ToolMessage, Command]],
         input_type: Literal["list", "dict", "tool_calls"],
     ) -> list[Union[Command, list[ToolMessage], dict[str, list[ToolMessage]]]]:
         # preserve existing behavior for non-command tool outputs for backwards
@@ -437,10 +445,27 @@ class ToolNode(RunnableCallable):
         call: ToolCall,
         input_type: Literal["list", "dict", "tool_calls"],
         config: RunnableConfig,
-    ) -> ToolMessage:
+    ) -> Union[ToolMessage, Command]:
         """Run a single tool call synchronously."""
         if invalid_tool_message := self._validate_tool_call(call):
             return invalid_tool_message
+
+        # Handle structured output tools
+        if call["name"] in self.structured_output_tools:
+            response_schema = self.tools_by_name[call["name"]]
+            return Command(
+                update={
+                    "messages": [
+                        ToolMessage(
+                            content=msg_content_output(call["args"]),
+                            name=call["name"],
+                            tool_call_id=call["id"],
+                        )
+                    ],
+                    "structured_response": response_schema(**call["args"]),
+                }
+            )
+
         try:
             call_args = {**call, **{"type": "tool_call"}}
             response = self.tools_by_name[call["name"]].invoke(call_args, config)
@@ -493,11 +518,27 @@ class ToolNode(RunnableCallable):
         call: ToolCall,
         input_type: Literal["list", "dict", "tool_calls"],
         config: RunnableConfig,
-    ) -> ToolMessage:
+    ) -> Union[ToolMessage, Command]:
         """Run a single tool call asynchronously."""
         if invalid_tool_message := self._validate_tool_call(call):
             return invalid_tool_message
 
+        # Handle structured output tools
+        if call["name"] in self.structured_output_tools:
+            response_schema = self.tools_by_name[call["name"]]
+            return Command(
+                update={
+                    "messages": [
+                        ToolMessage(
+                            content=msg_content_output(call["args"]),
+                            name=call["name"],
+                            tool_call_id=call["id"],
+                        )
+                    ],
+                    "structured_response": response_schema(**call["args"]),
+                }
+            )
+
         try:
             call_args = {**call, **{"type": "tool_call"}}
             response = await self.tools_by_name[call["name"]].ainvoke(call_args, config)

libs/prebuilt/tests/__snapshots__/test_react_agent_graph.ambr

@@ -171,3 +171,191 @@
   
   '''
 # ---
+# name: test_react_agent_graph_structure_with_individual_nodes[no_response_format-no_post_hook-no_pre_hook-no_tools]
+  '''
+  graph TD;
+  	__start__ --> agent;
+  	agent --> __end__;
+  
+  '''
+# ---
+# name: test_react_agent_graph_structure_with_individual_nodes[no_response_format-no_post_hook-no_pre_hook-two_tools]
+  '''
+  graph TD;
+  	__start__ --> agent;
+  	agent -.-> __end__;
+  	agent -.-> tool;
+  	agent -.-> tool2;
+  	tool --> agent;
+  	tool2 --> agent;
+  
+  '''
+# ---
+# name: test_react_agent_graph_structure_with_individual_nodes[no_response_format-no_post_hook-with_pre_hook-no_tools]
+  '''
+  graph TD;
+  	__start__ --> pre_model_hook;
+  	pre_model_hook --> agent;
+  	agent --> __end__;
+  
+  '''
+# ---
+# name: test_react_agent_graph_structure_with_individual_nodes[no_response_format-no_post_hook-with_pre_hook-two_tools]
+  '''
+  graph TD;
+  	__start__ --> pre_model_hook;
+  	agent -.-> __end__;
+  	agent -.-> tool;
+  	agent -.-> tool2;
+  	pre_model_hook --> agent;
+  	tool --> pre_model_hook;
+  	tool2 --> pre_model_hook;
+  
+  '''
+# ---
+# name: test_react_agent_graph_structure_with_individual_nodes[no_response_format-with_post_hook-no_pre_hook-no_tools]
+  '''
+  graph TD;
+  	__start__ --> agent;
+  	agent --> post_model_hook;
+  	post_model_hook --> __end__;
+  
+  '''
+# ---
+# name: test_react_agent_graph_structure_with_individual_nodes[no_response_format-with_post_hook-no_pre_hook-two_tools]
+  '''
+  graph TD;
+  	__start__ --> agent;
+  	agent --> post_model_hook;
+  	post_model_hook -.-> __end__;
+  	post_model_hook -.-> agent;
+  	post_model_hook -.-> tool;
+  	post_model_hook -.-> tool2;
+  	tool --> agent;
+  	tool2 --> agent;
+  
+  '''
+# ---
+# name: test_react_agent_graph_structure_with_individual_nodes[no_response_format-with_post_hook-with_pre_hook-no_tools]
+  '''
+  graph TD;
+  	__start__ --> pre_model_hook;
+  	agent --> post_model_hook;
+  	pre_model_hook --> agent;
+  	post_model_hook --> __end__;
+  
+  '''
+# ---
+# name: test_react_agent_graph_structure_with_individual_nodes[no_response_format-with_post_hook-with_pre_hook-two_tools]
+  '''
+  graph TD;
+  	__start__ --> pre_model_hook;
+  	agent --> post_model_hook;
+  	post_model_hook -.-> __end__;
+  	post_model_hook -.-> pre_model_hook;
+  	post_model_hook -.-> tool;
+  	post_model_hook -.-> tool2;
+  	pre_model_hook --> agent;
+  	tool --> pre_model_hook;
+  	tool2 --> pre_model_hook;
+  
+  '''
+# ---
+# name: test_react_agent_graph_structure_with_individual_nodes[with_response_format-no_post_hook-no_pre_hook-no_tools]
+  '''
+  graph TD;
+  	__start__ --> agent;
+  	agent --> generate_structured_response;
+  	generate_structured_response --> __end__;
+  
+  '''
+# ---
+# name: test_react_agent_graph_structure_with_individual_nodes[with_response_format-no_post_hook-no_pre_hook-two_tools]
+  '''
+  graph TD;
+  	__start__ --> agent;
+  	agent -.-> generate_structured_response;
+  	agent -.-> tool;
+  	agent -.-> tool2;
+  	tool --> agent;
+  	tool2 --> agent;
+  	generate_structured_response --> __end__;
+  
+  '''
+# ---
+# name: test_react_agent_graph_structure_with_individual_nodes[with_response_format-no_post_hook-with_pre_hook-no_tools]
+  '''
+  graph TD;
+  	__start__ --> pre_model_hook;
+  	agent --> generate_structured_response;
+  	pre_model_hook --> agent;
+  	generate_structured_response --> __end__;
+  
+  '''
+# ---
+# name: test_react_agent_graph_structure_with_individual_nodes[with_response_format-no_post_hook-with_pre_hook-two_tools]
+  '''
+  graph TD;
+  	__start__ --> pre_model_hook;
+  	agent -.-> generate_structured_response;
+  	agent -.-> tool;
+  	agent -.-> tool2;
+  	pre_model_hook --> agent;
+  	tool --> pre_model_hook;
+  	tool2 --> pre_model_hook;
+  	generate_structured_response --> __end__;
+  
+  '''
+# ---
+# name: test_react_agent_graph_structure_with_individual_nodes[with_response_format-with_post_hook-no_pre_hook-no_tools]
+  '''
+  graph TD;
+  	__start__ --> agent;
+  	agent --> post_model_hook;
+  	post_model_hook --> generate_structured_response;
+  	generate_structured_response --> __end__;
+  
+  '''
+# ---
+# name: test_react_agent_graph_structure_with_individual_nodes[with_response_format-with_post_hook-no_pre_hook-two_tools]
+  '''
+  graph TD;
+  	__start__ --> agent;
+  	agent --> post_model_hook;
+  	post_model_hook -.-> agent;
+  	post_model_hook -.-> generate_structured_response;
+  	post_model_hook -.-> tool;
+  	post_model_hook -.-> tool2;
+  	tool --> agent;
+  	tool2 --> agent;
+  	generate_structured_response --> __end__;
+  
+  '''
+# ---
+# name: test_react_agent_graph_structure_with_individual_nodes[with_response_format-with_post_hook-with_pre_hook-no_tools]
+  '''
+  graph TD;
+  	__start__ --> pre_model_hook;
+  	agent --> post_model_hook;
+  	post_model_hook --> generate_structured_response;
+  	pre_model_hook --> agent;
+  	generate_structured_response --> __end__;
+  
+  '''
+# ---
+# name: test_react_agent_graph_structure_with_individual_nodes[with_response_format-with_post_hook-with_pre_hook-two_tools]
+  '''
+  graph TD;
+  	__start__ --> pre_model_hook;
+  	agent --> post_model_hook;
+  	post_model_hook -.-> generate_structured_response;
+  	post_model_hook -.-> pre_model_hook;
+  	post_model_hook -.-> tool;
+  	post_model_hook -.-> tool2;
+  	pre_model_hook --> agent;
+  	tool --> pre_model_hook;
+  	tool2 --> pre_model_hook;
+  	generate_structured_response --> __end__;
+  
+  '''
+# ---

libs/prebuilt/tests/test_react_agent.py

@@ -1,14 +1,9 @@
 import dataclasses
 import inspect
-import json
-from functools import partial
 from typing import (
     Annotated,
-    List,
     Literal,
     Optional,
-    Type,
-    TypeVar,
     Union,
 )
 
@@ -16,7 +11,6 @@ import pytest
 from langchain_core.language_models import BaseChatModel
 from langchain_core.messages import (
     AIMessage,
-    AnyMessage,
     HumanMessage,
     MessageLikeRepresentation,
     RemoveMessage,
@@ -28,17 +22,14 @@ from langchain_core.runnables import RunnableConfig, RunnableLambda
 from langchain_core.tools import InjectedToolCallId, ToolException
 from langchain_core.tools import tool as dec_tool
 from pydantic import BaseModel, Field
-from pydantic.v1 import BaseModel as BaseModelV1
 from typing_extensions import TypedDict
 
 from langgraph.checkpoint.base import BaseCheckpointSaver
-from langgraph.config import get_stream_writer
-from langgraph.graph import START, MessagesState, StateGraph, add_messages
+from langgraph.graph import START, MessagesState, StateGraph
 from langgraph.graph.message import REMOVE_ALL_MESSAGES
 from langgraph.prebuilt import (
     ToolNode,
     create_react_agent,
-    tools_condition,
 )
 from langgraph.prebuilt.chat_agent_executor import (
     AgentState,
@@ -184,7 +175,7 @@ def test_runnable_prompt():
 
 
 @pytest.mark.parametrize("version", REACT_TOOL_CALL_VERSIONS)
-def test_prompt_with_store(version: str):
+def test_prompt_with_store(version: Literal["v1", "v2"]):
     def add(a: int, b: int):
         """Adds a and b"""
         return a + b
@@ -654,124 +645,6 @@ def test_react_agent_parallel_tool_calls(
         assert get_weather_execution_count == 1
 
 
-class _InjectStateSchema(TypedDict):
-    messages: list
-    foo: str
-
-
-class _InjectedStatePydanticSchema(BaseModelV1):
-    messages: list
-    foo: str
-
-
-class _InjectedStatePydanticV2Schema(BaseModel):
-    messages: list
-    foo: str
-
-
-@dataclasses.dataclass
-class _InjectedStateDataclassSchema:
-    messages: list
-    foo: str
-
-
-T = TypeVar("T")
-
-
-@pytest.mark.parametrize(
-    "schema_",
-    [
-        _InjectStateSchema,
-        _InjectedStatePydanticSchema,
-        _InjectedStatePydanticV2Schema,
-        _InjectedStateDataclassSchema,
-    ],
-)
-def test_tool_node_inject_state(schema_: Type[T]) -> None:
-    def tool1(some_val: int, state: Annotated[T, InjectedState]) -> str:
-        """Tool 1 docstring."""
-        if isinstance(state, dict):
-            return state["foo"]
-        else:
-            return getattr(state, "foo")
-
-    def tool2(some_val: int, state: Annotated[T, InjectedState()]) -> str:
-        """Tool 2 docstring."""
-        if isinstance(state, dict):
-            return state["foo"]
-        else:
-            return getattr(state, "foo")
-
-    def tool3(
-        some_val: int,
-        foo: Annotated[str, InjectedState("foo")],
-        msgs: Annotated[List[AnyMessage], InjectedState("messages")],
-    ) -> str:
-        """Tool 1 docstring."""
-        return foo
-
-    def tool4(
-        some_val: int, msgs: Annotated[List[AnyMessage], InjectedState("messages")]
-    ) -> str:
-        """Tool 1 docstring."""
-        return msgs[0].content
-
-    node = ToolNode([tool1, tool2, tool3, tool4])
-    for tool_name in ("tool1", "tool2", "tool3"):
-        tool_call = {
-            "name": tool_name,
-            "args": {"some_val": 1},
-            "id": "some 0",
-            "type": "tool_call",
-        }
-        msg = AIMessage("hi?", tool_calls=[tool_call])
-        result = node.invoke(schema_(**{"messages": [msg], "foo": "bar"}))
-        tool_message = result["messages"][-1]
-        assert tool_message.content == "bar", f"Failed for tool={tool_name}"
-
-        if tool_name == "tool3":
-            failure_input = None
-            try:
-                failure_input = schema_(**{"messages": [msg], "notfoo": "bar"})
-            except Exception:
-                pass
-            if failure_input is not None:
-                with pytest.raises(KeyError):
-                    node.invoke(failure_input)
-
-                with pytest.raises(ValueError):
-                    node.invoke([msg])
-        else:
-            failure_input = None
-            try:
-                failure_input = schema_(**{"messages": [msg], "notfoo": "bar"})
-            except Exception:
-                # We'd get a validation error from pydantic state and wouldn't make it to the node
-                # anyway
-                pass
-            if failure_input is not None:
-                messages_ = node.invoke(failure_input)
-                tool_message = messages_["messages"][-1]
-                assert "KeyError" in tool_message.content
-                tool_message = node.invoke([msg])[-1]
-                assert "KeyError" in tool_message.content
-
-    tool_call = {
-        "name": "tool4",
-        "args": {"some_val": 1},
-        "id": "some 0",
-        "type": "tool_call",
-    }
-    msg = AIMessage("hi?", tool_calls=[tool_call])
-    result = node.invoke(schema_(**{"messages": [msg], "foo": ""}))
-    tool_message = result["messages"][-1]
-    assert tool_message.content == "hi?"
-
-    result = node.invoke([msg])
-    tool_message = result[-1]
-    assert tool_message.content == "hi?"
-
-
 class AgentStateExtraKey(AgentState):
     foo: int
 
@@ -780,14 +653,24 @@ class AgentStateExtraKeyPydantic(AgentStatePydantic):
     foo: int
 
 
-@pytest.mark.parametrize("version", REACT_TOOL_CALL_VERSIONS)
+@pytest.mark.parametrize("version", ["v1", "v2"])
 @pytest.mark.parametrize(
     "state_schema", [AgentStateExtraKey, AgentStateExtraKeyPydantic]
 )
+@pytest.mark.parametrize(
+    "use_individual_tool_nodes",
+    [False, True],
+    ids=["single_tool_node", "node_per_tool"],
+)
 def test_create_react_agent_inject_vars(
-    version: Literal["v1", "v2"], state_schema: StateSchemaType
+    version: Literal["v1", "v2"],
+    state_schema: StateSchemaType,
+    use_individual_tool_nodes: bool,
 ) -> None:
     """Test that the agent can inject state and store into tool functions."""
+    if version == "v1" and use_individual_tool_nodes:
+        pytest.skip("v1 does not support individual tool nodes")
+
     store = InMemoryStore()
     namespace = ("test",)
     store.put(namespace, "test_key", {"bar": 3})
@@ -826,6 +709,7 @@ def test_create_react_agent_inject_vars(
         state_schema=state_schema,
         store=store,
         version=version,
+        use_individual_tool_nodes=use_individual_tool_nodes,
     )
     result = agent.invoke({"messages": [{"role": "user", "content": "hi"}], "foo": 2})
     assert result["messages"] == [
@@ -837,137 +721,18 @@ def test_create_react_agent_inject_vars(
     assert result["foo"] == 2
 
 
-def test_tool_node_inject_store() -> None:
-    store = InMemoryStore()
-    namespace = ("test",)
-
-    def tool1(some_val: int, store: Annotated[BaseStore, InjectedStore()]) -> str:
-        """Tool 1 docstring."""
-        store_val = store.get(namespace, "test_key").value["foo"]
-        return f"Some val: {some_val}, store val: {store_val}"
-
-    def tool2(some_val: int, store: Annotated[BaseStore, InjectedStore()]) -> str:
-        """Tool 2 docstring."""
-        store_val = store.get(namespace, "test_key").value["foo"]
-        return f"Some val: {some_val}, store val: {store_val}"
-
-    def tool3(
-        some_val: int,
-        bar: Annotated[str, InjectedState("bar")],
-        store: Annotated[BaseStore, InjectedStore()],
-    ) -> str:
-        """Tool 3 docstring."""
-        store_val = store.get(namespace, "test_key").value["foo"]
-        return f"Some val: {some_val}, store val: {store_val}, state val: {bar}"
-
-    node = ToolNode([tool1, tool2, tool3], handle_tool_errors=True)
-    store.put(namespace, "test_key", {"foo": "bar"})
-
-    class State(MessagesState):
-        bar: str
-
-    builder = StateGraph(State)
-    builder.add_node("tools", node)
-    builder.add_edge(START, "tools")
-    graph = builder.compile(store=store)
-
-    for tool_name in ("tool1", "tool2"):
-        tool_call = {
-            "name": tool_name,
-            "args": {"some_val": 1},
-            "id": "some 0",
-            "type": "tool_call",
-        }
-        msg = AIMessage("hi?", tool_calls=[tool_call])
-        node_result = node.invoke({"messages": [msg]}, store=store)
-        graph_result = graph.invoke({"messages": [msg]})
-        for result in (node_result, graph_result):
-            result["messages"][-1]
-            tool_message = result["messages"][-1]
-            assert tool_message.content == "Some val: 1, store val: bar", (
-                f"Failed for tool={tool_name}"
-            )
-
-    tool_call = {
-        "name": "tool3",
-        "args": {"some_val": 1},
-        "id": "some 0",
-        "type": "tool_call",
-    }
-    msg = AIMessage("hi?", tool_calls=[tool_call])
-    node_result = node.invoke({"messages": [msg], "bar": "baz"}, store=store)
-    graph_result = graph.invoke({"messages": [msg], "bar": "baz"})
-    for result in (node_result, graph_result):
-        result["messages"][-1]
-        tool_message = result["messages"][-1]
-        assert tool_message.content == "Some val: 1, store val: bar, state val: baz", (
-            f"Failed for tool={tool_name}"
-        )
-
-    # test injected store without passing store to compiled graph
-    failing_graph = builder.compile()
-    with pytest.raises(ValueError):
-        failing_graph.invoke({"messages": [msg], "bar": "baz"})
-
-
-def test_tool_node_ensure_utf8() -> None:
-    @dec_tool
-    def get_day_list(days: list[str]) -> list[str]:
-        """choose days"""
-        return days
-
-    data = ["星期一", "水曜日", "목요일", "Friday"]
-    tools = [get_day_list]
-    tool_calls = [ToolCall(name=get_day_list.name, args={"days": data}, id="test_id")]
-    outputs: list[ToolMessage] = ToolNode(tools).invoke(
-        [AIMessage(content="", tool_calls=tool_calls)]
-    )
-    assert outputs[0].content == json.dumps(data, ensure_ascii=False)
-
-
-def test_tool_node_messages_key() -> None:
-    @dec_tool
-    def add(a: int, b: int):
-        """Adds a and b."""
-        return a + b
-
-    model = FakeToolCallingModel(
-        tool_calls=[[ToolCall(name=add.name, args={"a": 1, "b": 2}, id="test_id")]]
-    )
-
-    class State(TypedDict):
-        subgraph_messages: Annotated[list[AnyMessage], add_messages]
-
-    def call_model(state: State):
-        response = model.invoke(state["subgraph_messages"])
-        model.tool_calls = []
-        return {"subgraph_messages": response}
-
-    builder = StateGraph(State)
-    builder.add_node("agent", call_model)
-    builder.add_node("tools", ToolNode([add], messages_key="subgraph_messages"))
-    builder.add_conditional_edges(
-        "agent", partial(tools_condition, messages_key="subgraph_messages")
-    )
-    builder.add_edge(START, "agent")
-    builder.add_edge("tools", "agent")
-
-    graph = builder.compile()
-    result = graph.invoke({"subgraph_messages": [HumanMessage(content="hi")]})
-    assert result["subgraph_messages"] == [
-        _AnyIdHumanMessage(content="hi"),
-        AIMessage(
-            content="hi",
-            id="0",
-            tool_calls=[ToolCall(name=add.name, args={"a": 1, "b": 2}, id="test_id")],
-        ),
-        _AnyIdToolMessage(content="3", name=add.name, tool_call_id="test_id"),
-        AIMessage(content="hi-hi-3", id="1"),
-    ]
-
-
 @pytest.mark.parametrize("version", REACT_TOOL_CALL_VERSIONS)
-async def test_return_direct(version: str) -> None:
+@pytest.mark.parametrize(
+    "use_individual_tool_nodes",
+    [False, True],
+    ids=["single_tool_node", "node_per_tool"],
+)
+async def test_return_direct(
+    version: Literal["v1", "v2"], use_individual_tool_nodes: bool
+) -> None:
+    if version == "v1" and use_individual_tool_nodes:
+        pytest.skip("v1 does not support individual tool nodes")
+
     @dec_tool(return_direct=True)
     def tool_return_direct(input: str) -> str:
         """A tool that returns directly."""
@@ -995,6 +760,7 @@ async def test_return_direct(version: str) -> None:
         model,
         [tool_return_direct, tool_normal],
         version=version,
+        use_individual_tool_nodes=use_individual_tool_nodes,
     )
 
     # Test direct return for tool_return_direct
@@ -1088,15 +854,27 @@ def test__get_state_args() -> None:
 
 
 def test_inspect_react() -> None:
+    """Test that we can inspect the agent and its nodes."""
     model = FakeToolCallingModel(tool_calls=[])
     agent = create_react_agent(model, [])
     inspect.getclosurevars(agent.nodes["agent"].bound.func)
 
 
 @pytest.mark.parametrize("version", REACT_TOOL_CALL_VERSIONS)
+@pytest.mark.parametrize(
+    "use_individual_tool_nodes",
+    [False, True],
+    ids=["single_tool_node", "node_per_tool"],
+)
 def test_react_with_subgraph_tools(
-    sync_checkpointer: BaseCheckpointSaver, version: Literal["v1", "v2"]
+    sync_checkpointer: BaseCheckpointSaver,
+    version: Literal["v1", "v2"],
+    use_individual_tool_nodes: bool,
 ) -> None:
+    """Test React agent with subgraph tools."""
+    if version == "v1" and use_individual_tool_nodes:
+        pytest.skip("v1 does not support individual tool nodes")
+
     class State(TypedDict):
         a: int
         b: int
@@ -1152,6 +930,7 @@ def test_react_with_subgraph_tools(
         tool_node,
         checkpointer=sync_checkpointer,
         version=version,
+        use_individual_tool_nodes=use_individual_tool_nodes,
     )
     result = agent.invoke(
         {"messages": [HumanMessage(content="What's 2 + 3 and 2 * 3?")]},
@@ -1182,63 +961,18 @@ def test_react_with_subgraph_tools(
     ]
 
 
-def test_tool_node_stream_writer() -> None:
-    @dec_tool
-    def streaming_tool(x: int) -> str:
-        """Do something with writer."""
-        my_writer = get_stream_writer()
-        for value in ["foo", "bar", "baz"]:
-            my_writer({"custom_tool_value": value})
-
-        return x
-
-    tool_node = ToolNode([streaming_tool])
-    graph = (
-        StateGraph(MessagesState)
-        .add_node("tools", tool_node)
-        .add_edge(START, "tools")
-        .compile()
-    )
-
-    tool_call = {
-        "name": "streaming_tool",
-        "args": {"x": 1},
-        "id": "1",
-        "type": "tool_call",
-    }
-    inputs = {
-        "messages": [AIMessage("", tool_calls=[tool_call])],
-    }
-
-    assert list(graph.stream(inputs, stream_mode="custom")) == [
-        {"custom_tool_value": "foo"},
-        {"custom_tool_value": "bar"},
-        {"custom_tool_value": "baz"},
-    ]
-    assert list(graph.stream(inputs, stream_mode=["custom", "updates"])) == [
-        ("custom", {"custom_tool_value": "foo"}),
-        ("custom", {"custom_tool_value": "bar"}),
-        ("custom", {"custom_tool_value": "baz"}),
-        (
-            "updates",
-            {
-                "tools": {
-                    "messages": [
-                        _AnyIdToolMessage(
-                            content="1",
-                            name="streaming_tool",
-                            tool_call_id="1",
-                        ),
-                    ],
-                },
-            },
-        ),
-    ]
-
-
 @pytest.mark.parametrize("version", REACT_TOOL_CALL_VERSIONS)
-def test_react_agent_subgraph_streaming_sync(version: Literal["v1", "v2"]) -> None:
+@pytest.mark.parametrize(
+    "use_individual_tool_nodes",
+    [False, True],
+    ids=["single_tool_node", "node_per_tool"],
+)
+def test_react_agent_subgraph_streaming_sync(
+    version: Literal["v1", "v2"], use_individual_tool_nodes: bool
+) -> None:
     """Test React agent streaming when used as a subgraph node sync version"""
+    if version == "v1" and use_individual_tool_nodes:
+        pytest.skip("v1 does not support individual tool nodes")
 
     @dec_tool
     def get_weather(city: str) -> str:
@@ -1258,6 +992,7 @@ def test_react_agent_subgraph_streaming_sync(version: Literal["v1", "v2"]) -> No
         tools=[get_weather],
         prompt="You are a helpful travel assistant.",
         version=version,
+        use_individual_tool_nodes=use_individual_tool_nodes,
     )
 
     # Create a subgraph that uses the React agent as a node

libs/prebuilt/tests/test_react_agent_graph.py

@@ -15,6 +15,11 @@ def tool() -> None:
     ...
 
 
+def tool2() -> None:
+    """Another testing tool."""
+    ...
+
+
 def pre_model_hook() -> None:
     """Pre-model hook."""
     ...
@@ -49,4 +54,44 @@ def test_react_agent_graph_structure(
         post_model_hook=post_model_hook,
         response_format=response_format,
     )
+    try:
+        assert agent.get_graph().draw_mermaid(with_styles=False) == snapshot
+    except Exception as e:
+        raise ValueError(
+            "The graph structure has changed. Please update the snapshot."
+            "Configuration used:\n"
+            f"tools: {tools}, "
+            f"pre_model_hook: {pre_model_hook}, "
+            f"post_model_hook: {post_model_hook}, "
+            f"response_format: {response_format}"
+        ) from e
+
+
+@pytest.mark.parametrize("tools", [[], [tool, tool2]], ids=["no_tools", "two_tools"])
+@pytest.mark.parametrize(
+    "pre_model_hook", [None, pre_model_hook], ids=["no_pre_hook", "with_pre_hook"]
+)
+@pytest.mark.parametrize(
+    "post_model_hook", [None, post_model_hook], ids=["no_post_hook", "with_post_hook"]
+)
+@pytest.mark.parametrize(
+    "response_format",
+    [None, ResponseFormat],
+    ids=["no_response_format", "with_response_format"],
+)
+def test_react_agent_graph_structure_with_individual_nodes(
+    snapshot: SnapshotAssertion,
+    tools: list[Callable],
+    pre_model_hook: Union[Callable, None],
+    post_model_hook: Union[Callable, None],
+    response_format: Union[type[BaseModel], None],
+) -> None:
+    agent = create_react_agent(
+        model,
+        tools=tools,
+        pre_model_hook=pre_model_hook,
+        post_model_hook=post_model_hook,
+        response_format=response_format,
+        use_individual_tool_nodes=True,
+    )
     assert agent.get_graph().draw_mermaid(with_styles=False) == snapshot

libs/prebuilt/tests/test_tool_node.py

@@ -1,25 +1,49 @@
+import dataclasses
+import json
+from functools import partial
 from typing import (
     Annotated,
     Any,
+    List,
+    Type,
+    TypeVar,
     Union,
 )
 
 import pytest
 from langchain_core.messages import (
     AIMessage,
+    AnyMessage,
+    HumanMessage,
     RemoveMessage,
+    ToolCall,
     ToolMessage,
 )
 from langchain_core.tools import BaseTool, ToolException
 from langchain_core.tools import tool as dec_tool
 from pydantic import BaseModel, ValidationError
+from pydantic.v1 import BaseModel as BaseModelV1
 from pydantic.v1 import ValidationError as ValidationErrorV1
+from typing_extensions import TypedDict
 
+from langgraph.config import get_stream_writer
 from langgraph.errors import GraphBubbleUp, GraphInterrupt
-from langgraph.graph.message import REMOVE_ALL_MESSAGES
-from langgraph.prebuilt import ToolNode
-from langgraph.prebuilt.tool_node import TOOL_CALL_ERROR_TEMPLATE
+from langgraph.graph import START, MessagesState, StateGraph
+from langgraph.graph.message import REMOVE_ALL_MESSAGES, add_messages
+from langgraph.prebuilt import (
+    ToolNode,
+)
+from langgraph.prebuilt.tool_node import (
+    TOOL_CALL_ERROR_TEMPLATE,
+    InjectedState,
+    InjectedStore,
+    tools_condition,
+)
+from langgraph.store.base import BaseStore
+from langgraph.store.memory import InMemoryStore
 from langgraph.types import Command, Send
+from tests.messages import _AnyIdHumanMessage, _AnyIdToolMessage
+from tests.model import FakeToolCallingModel
 
 pytestmark = pytest.mark.anyio
 
@@ -1156,3 +1180,372 @@ async def test_tool_node_command_remove_all_messages():
     command = result[0]
     assert isinstance(command, Command)
     assert command.update == {"messages": [RemoveMessage(id=REMOVE_ALL_MESSAGES)]}
+
+
+class _InjectStateSchema(TypedDict):
+    messages: list
+    foo: str
+
+
+class _InjectedStatePydanticSchema(BaseModelV1):
+    messages: list
+    foo: str
+
+
+class _InjectedStatePydanticV2Schema(BaseModel):
+    messages: list
+    foo: str
+
+
+@dataclasses.dataclass
+class _InjectedStateDataclassSchema:
+    messages: list
+    foo: str
+
+
+T = TypeVar("T")
+
+
+@pytest.mark.parametrize(
+    "schema_",
+    [
+        _InjectStateSchema,
+        _InjectedStatePydanticSchema,
+        _InjectedStatePydanticV2Schema,
+        _InjectedStateDataclassSchema,
+    ],
+)
+def test_tool_node_inject_state(schema_: Type[T]) -> None:
+    def tool1(some_val: int, state: Annotated[T, InjectedState]) -> str:
+        """Tool 1 docstring."""
+        if isinstance(state, dict):
+            return state["foo"]
+        else:
+            return getattr(state, "foo")
+
+    def tool2(some_val: int, state: Annotated[T, InjectedState()]) -> str:
+        """Tool 2 docstring."""
+        if isinstance(state, dict):
+            return state["foo"]
+        else:
+            return getattr(state, "foo")
+
+    def tool3(
+        some_val: int,
+        foo: Annotated[str, InjectedState("foo")],
+        msgs: Annotated[List[AnyMessage], InjectedState("messages")],
+    ) -> str:
+        """Tool 1 docstring."""
+        return foo
+
+    def tool4(
+        some_val: int, msgs: Annotated[List[AnyMessage], InjectedState("messages")]
+    ) -> str:
+        """Tool 1 docstring."""
+        return msgs[0].content
+
+    node = ToolNode([tool1, tool2, tool3, tool4])
+    for tool_name in ("tool1", "tool2", "tool3"):
+        tool_call = {
+            "name": tool_name,
+            "args": {"some_val": 1},
+            "id": "some 0",
+            "type": "tool_call",
+        }
+        msg = AIMessage("hi?", tool_calls=[tool_call])
+        result = node.invoke(schema_(**{"messages": [msg], "foo": "bar"}))
+        tool_message = result["messages"][-1]
+        assert tool_message.content == "bar", f"Failed for tool={tool_name}"
+
+        if tool_name == "tool3":
+            failure_input = None
+            try:
+                failure_input = schema_(**{"messages": [msg], "notfoo": "bar"})
+            except Exception:
+                pass
+            if failure_input is not None:
+                with pytest.raises(KeyError):
+                    node.invoke(failure_input)
+
+                with pytest.raises(ValueError):
+                    node.invoke([msg])
+        else:
+            failure_input = None
+            try:
+                failure_input = schema_(**{"messages": [msg], "notfoo": "bar"})
+            except Exception:
+                # We'd get a validation error from pydantic state and wouldn't make it to the node
+                # anyway
+                pass
+            if failure_input is not None:
+                messages_ = node.invoke(failure_input)
+                tool_message = messages_["messages"][-1]
+                assert "KeyError" in tool_message.content
+                tool_message = node.invoke([msg])[-1]
+                assert "KeyError" in tool_message.content
+
+    tool_call = {
+        "name": "tool4",
+        "args": {"some_val": 1},
+        "id": "some 0",
+        "type": "tool_call",
+    }
+    msg = AIMessage("hi?", tool_calls=[tool_call])
+    result = node.invoke(schema_(**{"messages": [msg], "foo": ""}))
+    tool_message = result["messages"][-1]
+    assert tool_message.content == "hi?"
+
+    result = node.invoke([msg])
+    tool_message = result[-1]
+    assert tool_message.content == "hi?"
+
+
+def test_tool_node_inject_store() -> None:
+    store = InMemoryStore()
+    namespace = ("test",)
+
+    def tool1(some_val: int, store: Annotated[BaseStore, InjectedStore()]) -> str:
+        """Tool 1 docstring."""
+        store_val = store.get(namespace, "test_key").value["foo"]
+        return f"Some val: {some_val}, store val: {store_val}"
+
+    def tool2(some_val: int, store: Annotated[BaseStore, InjectedStore()]) -> str:
+        """Tool 2 docstring."""
+        store_val = store.get(namespace, "test_key").value["foo"]
+        return f"Some val: {some_val}, store val: {store_val}"
+
+    def tool3(
+        some_val: int,
+        bar: Annotated[str, InjectedState("bar")],
+        store: Annotated[BaseStore, InjectedStore()],
+    ) -> str:
+        """Tool 3 docstring."""
+        store_val = store.get(namespace, "test_key").value["foo"]
+        return f"Some val: {some_val}, store val: {store_val}, state val: {bar}"
+
+    node = ToolNode([tool1, tool2, tool3], handle_tool_errors=True)
+    store.put(namespace, "test_key", {"foo": "bar"})
+
+    class State(MessagesState):
+        bar: str
+
+    builder = StateGraph(State)
+    builder.add_node("tools", node)
+    builder.add_edge(START, "tools")
+    graph = builder.compile(store=store)
+
+    for tool_name in ("tool1", "tool2"):
+        tool_call = {
+            "name": tool_name,
+            "args": {"some_val": 1},
+            "id": "some 0",
+            "type": "tool_call",
+        }
+        msg = AIMessage("hi?", tool_calls=[tool_call])
+        node_result = node.invoke({"messages": [msg]}, store=store)
+        graph_result = graph.invoke({"messages": [msg]})
+        for result in (node_result, graph_result):
+            result["messages"][-1]
+            tool_message = result["messages"][-1]
+            assert tool_message.content == "Some val: 1, store val: bar", (
+                f"Failed for tool={tool_name}"
+            )
+
+    tool_call = {
+        "name": "tool3",
+        "args": {"some_val": 1},
+        "id": "some 0",
+        "type": "tool_call",
+    }
+    msg = AIMessage("hi?", tool_calls=[tool_call])
+    node_result = node.invoke({"messages": [msg], "bar": "baz"}, store=store)
+    graph_result = graph.invoke({"messages": [msg], "bar": "baz"})
+    for result in (node_result, graph_result):
+        result["messages"][-1]
+        tool_message = result["messages"][-1]
+        assert tool_message.content == "Some val: 1, store val: bar, state val: baz", (
+            f"Failed for tool={tool_name}"
+        )
+
+    # test injected store without passing store to compiled graph
+    failing_graph = builder.compile()
+    with pytest.raises(ValueError):
+        failing_graph.invoke({"messages": [msg], "bar": "baz"})
+
+
+def test_tool_node_ensure_utf8() -> None:
+    @dec_tool
+    def get_day_list(days: list[str]) -> list[str]:
+        """choose days"""
+        return days
+
+    data = ["星期一", "水曜日", "목요일", "Friday"]
+    tools = [get_day_list]
+    tool_calls = [ToolCall(name=get_day_list.name, args={"days": data}, id="test_id")]
+    outputs: list[ToolMessage] = ToolNode(tools).invoke(
+        [AIMessage(content="", tool_calls=tool_calls)]
+    )
+    assert outputs[0].content == json.dumps(data, ensure_ascii=False)
+
+
+def test_tool_node_messages_key() -> None:
+    @dec_tool
+    def add(a: int, b: int):
+        """Adds a and b."""
+        return a + b
+
+    model = FakeToolCallingModel(
+        tool_calls=[[ToolCall(name=add.name, args={"a": 1, "b": 2}, id="test_id")]]
+    )
+
+    class State(TypedDict):
+        subgraph_messages: Annotated[list[AnyMessage], add_messages]
+
+    def call_model(state: State):
+        response = model.invoke(state["subgraph_messages"])
+        model.tool_calls = []
+        return {"subgraph_messages": response}
+
+    builder = StateGraph(State)
+    builder.add_node("agent", call_model)
+    builder.add_node("tools", ToolNode([add], messages_key="subgraph_messages"))
+    builder.add_conditional_edges(
+        "agent", partial(tools_condition, messages_key="subgraph_messages")
+    )
+    builder.add_edge(START, "agent")
+    builder.add_edge("tools", "agent")
+
+    graph = builder.compile()
+    result = graph.invoke({"subgraph_messages": [HumanMessage(content="hi")]})
+    assert result["subgraph_messages"] == [
+        _AnyIdHumanMessage(content="hi"),
+        AIMessage(
+            content="hi",
+            id="0",
+            tool_calls=[ToolCall(name=add.name, args={"a": 1, "b": 2}, id="test_id")],
+        ),
+        _AnyIdToolMessage(content="3", name=add.name, tool_call_id="test_id"),
+        AIMessage(content="hi-hi-3", id="1"),
+    ]
+
+
+def test_tool_node_stream_writer() -> None:
+    @dec_tool
+    def streaming_tool(x: int) -> str:
+        """Do something with writer."""
+        my_writer = get_stream_writer()
+        for value in ["foo", "bar", "baz"]:
+            my_writer({"custom_tool_value": value})
+
+        return x
+
+    tool_node = ToolNode([streaming_tool])
+    graph = (
+        StateGraph(MessagesState)
+        .add_node("tools", tool_node)
+        .add_edge(START, "tools")
+        .compile()
+    )
+
+    tool_call = {
+        "name": "streaming_tool",
+        "args": {"x": 1},
+        "id": "1",
+        "type": "tool_call",
+    }
+    inputs = {
+        "messages": [AIMessage("", tool_calls=[tool_call])],
+    }
+
+    assert list(graph.stream(inputs, stream_mode="custom")) == [
+        {"custom_tool_value": "foo"},
+        {"custom_tool_value": "bar"},
+        {"custom_tool_value": "baz"},
+    ]
+    assert list(graph.stream(inputs, stream_mode=["custom", "updates"])) == [
+        ("custom", {"custom_tool_value": "foo"}),
+        ("custom", {"custom_tool_value": "bar"}),
+        ("custom", {"custom_tool_value": "baz"}),
+        (
+            "updates",
+            {
+                "tools": {
+                    "messages": [
+                        _AnyIdToolMessage(
+                            content="1",
+                            name="streaming_tool",
+                            tool_call_id="1",
+                        ),
+                    ],
+                },
+            },
+        ),
+    ]
+
+
+
+async def test_structured_output_tools():
+    """Test that ToolNode handles Pydantic model classes as structured output tools."""
+
+    class OutputSchema(BaseModel):
+        name: str
+        age: int
+        location: str
+
+    tool_node = ToolNode([OutputSchema])
+
+    # Test that the structured output tool is registered correctly
+    assert "OutputSchema" in tool_node.tools_by_name
+    assert "OutputSchema" in tool_node.structured_output_tools
+
+    # Create a tool call that matches the schema
+    tool_call = {
+        "name": "OutputSchema",
+        "args": {"name": "Alice", "age": 30, "location": "NYC"},
+        "id": "call_123",
+        "type": "tool_call",
+    }
+
+    # Test sync execution
+    result = tool_node.invoke(
+        {"messages": [AIMessage(content="", tool_calls=[tool_call])]}
+    )
+
+    # Should return a Command with structured response
+    assert isinstance(result, list)
+    assert len(result) == 1
+    command = result[0]
+    assert isinstance(command, Command)
+
+    # Check the update structure
+    assert "messages" in command.update
+    assert "structured_response" in command.update
+
+    # Check the tool message
+    tool_message = command.update["messages"][0]
+    assert isinstance(tool_message, ToolMessage)
+    assert tool_message.name == "OutputSchema"
+    assert tool_message.tool_call_id == "call_123"
+
+    # Check the structured response
+    structured_response = command.update["structured_response"]
+    assert isinstance(structured_response, OutputSchema)
+    assert structured_response.name == "Alice"
+    assert structured_response.age == 30
+    assert structured_response.location == "NYC"
+
+    # Test async execution
+    result_async = await tool_node.ainvoke(
+        {"messages": [AIMessage(content="", tool_calls=[tool_call])]}
+    )
+
+    # Should produce the same result
+    assert isinstance(result_async, list)
+    assert len(result_async) == 1
+    command_async = result_async[0]
+    assert isinstance(command_async, Command)
+    assert "structured_response" in command_async.update
+
+    structured_response_async = command_async.update["structured_response"]
+    assert isinstance(structured_response_async, OutputSchema)
+    assert structured_response_async.name == "Alice"