x

* Add structured output tools to ToolNode
* Fix default tool node name to match the actual default ('tools')
* Update doc-strings to explain what inputs/outputs are for the ToolNode.
* Mark internal attributes as private (potentially breaking -- although hopefully users aren't accessing these)


## Decisions points

* OK with two properties? Done since users may be relying on
`tools_by_name` and expanding the return type will break user code.

## Changes in public/private interface

### Marked as public

* Make `tools_by_name` an official public property
* Make `structured_output_tools` a public property

### Marked as private

There should be no reason why users are accessing these attributes

```python
_tool_to_state_args
_tool_to_store_arg
_handle_tool_errors
_messages_key
```


### Usage

```python

    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.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"
```

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/_internal/_injection.py

@@ -0,0 +1,160 @@
+from __future__ import annotations
+
+import inspect
+from dataclasses import dataclass
+from typing import Any, Optional, Callable
+
+from langchain_core.runnables import RunnableConfig
+
+from langgraph.store.base import BaseStore
+from langgraph.types import StreamWriter
+
+
+# Special type to denote any type is accepted
+ANY_TYPE = object()
+
+VALID_KINDS = (inspect.Parameter.POSITIONAL_OR_KEYWORD, inspect.Parameter.KEYWORD_ONLY)
+
+
+
+# List of keyword arguments that can be injected into nodes / tasks / tools at runtime.
+# A named argument may appear multiple times if it appears with distinct types.
+KWARGS_CONFIG_KEYS: tuple[tuple[str, tuple[Any, ...], str, Any], ...] = (
+    (
+        "config",
+        (
+            RunnableConfig,
+            "RunnableConfig",
+            Optional[RunnableConfig],
+            "Optional[RunnableConfig]",
+            inspect.Parameter.empty,
+        ),
+        # for now, use config directly, eventually, will pop off of Runtime
+        "N/A",
+        inspect.Parameter.empty,
+    ),
+    (
+        "writer",
+        (StreamWriter, "StreamWriter", inspect.Parameter.empty),
+        "stream_writer",
+        lambda _: None,
+    ),
+    (
+        "store",
+        (
+            BaseStore,
+            "BaseStore",
+            inspect.Parameter.empty,
+        ),
+        "store",
+        inspect.Parameter.empty,
+    ),
+    (
+        "store",
+        (
+            Optional[BaseStore],
+            "Optional[BaseStore]",
+        ),
+        "store",
+        None,
+    ),
+    (
+        "previous",
+        (ANY_TYPE,),
+        "previous",
+        inspect.Parameter.empty,
+    ),
+    (
+        "runtime",
+        (ANY_TYPE,),
+        # we never hit this block, we just inject runtime directly
+        "N/A",
+        inspect.Parameter.empty,
+    ),
+)
+
+
+@dataclass
+class InjectionInfo:
+    """Information about which injected arguments a function supports.
+
+    Attributes:
+        func_accepts: Dictionary mapping argument names to tuples of (runtime_key, default_value)
+        supported_args: Set of argument names that the function accepts for injection
+        has_config: Whether the function accepts a 'config' argument
+        has_writer: Whether the function accepts a 'writer' argument
+        has_store: Whether the function accepts a 'store' argument
+        has_previous: Whether the function accepts a 'previous' argument
+        has_runtime: Whether the function accepts a 'runtime' argument
+    """
+
+    func_accepts: dict[str, tuple[str, Any]]
+    supported_args: set[str]
+    has_config: bool
+    has_writer: bool
+    has_store: bool
+    has_previous: bool
+    has_runtime: bool
+
+
+def get_function_injection_info(func: Callable) -> InjectionInfo:
+    """Determine which injected arguments are supported by a function.
+
+    This function analyzes a function's signature to determine which runtime arguments
+    it can accept for injection. It uses the same logic as RunnableCallable to check
+    parameter names, types, and annotations against the supported injection types.
+
+    Args:
+        func: The function to analyze for injection support
+
+    Returns:
+        InjectionInfo containing details about which arguments the function supports
+
+    Example:
+        ```python
+        def my_tool(x: int, config: RunnableConfig, store: BaseStore) -> str:
+            return f"x={x}, config={config is not None}, store={store is not None}"
+
+        info = get_function_injection_info(my_tool)
+        print(info.has_config)  # True
+        print(info.has_store)   # True
+        print(info.has_writer)  # False
+        print(info.supported_args)  # {'config', 'store'}
+        ```
+    """
+    func_accepts: dict[str, tuple[str, Any]] = {}
+    params = inspect.signature(func).parameters
+
+    for kw, typ, runtime_key, default in KWARGS_CONFIG_KEYS:
+        p = params.get(kw)
+
+        if p is None or p.kind not in VALID_KINDS:
+            # If parameter is not found or is not a valid kind, skip
+            continue
+
+        if typ != (ANY_TYPE,) and p.annotation not in typ:
+            # A specific type is required, but the function annotation does
+            # not match the expected type.
+
+            # If this is a config parameter with incorrect typing, we still accept it
+            # but could emit a warning (following RunnableCallable behavior)
+            if kw == "config" and p.annotation != inspect.Parameter.empty:
+                # Could add warning here if needed
+                pass
+            else:
+                continue
+
+        # If the kwarg is accepted by the function, store the key / runtime attribute to inject
+        func_accepts[kw] = (runtime_key, default)
+
+    supported_args = set(func_accepts.keys())
+
+    return InjectionInfo(
+        func_accepts=func_accepts,
+        supported_args=supported_args,
+        has_config="config" in supported_args,
+        has_writer="writer" in supported_args,
+        has_store="store" in supported_args,
+        has_previous="previous" in supported_args,
+        has_runtime="runtime" in supported_args,
+    )

libs/prebuilt/langgraph/prebuilt/tool_node.py

@@ -31,6 +31,8 @@ Typical Usage:
     ```
 """
 
+from __future__ import annotations
+
 import asyncio
 import inspect
 import json
@@ -237,17 +239,50 @@ def _infer_handled_types(handler: Callable[..., str]) -> tuple[type[Exception],
 
 
 class ToolNode(RunnableCallable):
-    """A node that runs the tools called in the last AIMessage.
+    """A node for executing tools in LangGraph workflows.
 
-    It can be used either in StateGraph with a "messages" state key (or a custom key passed via ToolNode's 'messages_key').
-    If multiple tool calls are requested, they will be run in parallel. The output will be
-    a list of ToolMessages, one for each tool call.
+    Handles tool execution patterns including function calls, state injection,
+    persistent storage, and control flow. Manages parallel execution,
+    error handling.
 
-    Tool calls can also be passed directly as a list of `ToolCall` dicts.
+    Input Formats:
+        1. Graph state with `messages` key that has a list of messages:
+           - Common representation for agentic workflows
+           - Supports custom messages key via ``messages_key`` parameter
+
+        2. **Message List**: ``[AIMessage(..., tool_calls=[...])]``
+           - List of messages with tool calls in the last AIMessage
+
+        3. **Direct Tool Calls**: ``[{"name": "tool", "args": {...}, "id": "1", "type": "tool_call"}]``
+           - Bypasses message parsing for direct tool execution
+           - For programmatic tool invocation and testing
+
+    Tool Types:
+        1. **Regular tools**: Functions or BaseTool instances that return values or
+            Commands.
+        2. **Structured output tools**: Pydantic model classes for schema-validated
+           responses
+
+    Output Formats:
+        Output format depends on input type and tool behavior:
+
+        **For Regular tools**:
+        - Dict input → ``{"messages": [ToolMessage(...)]}``
+        - List input → ``[ToolMessage(...)]``
+
+        **For Command tools**:
+        - Returns ``[Command(...)]`` or mixed list with regular tool outputs
+        - Commands can update state, trigger navigation, or send messages
+
+        **For Structured output tools**:
+        - Returns ``[Command(update={"messages": [...], "structured_response": schema_instance})]``
+        - Includes both message and structured data in the graph state
 
     Args:
-        tools: A sequence of tools that can be invoked by this node. Tools can be
-            BaseTool instances or plain functions that will be converted to tools.
+        tools: A sequence of tools that can be invoked by this node. Supports:
+            - **BaseTool instances**: Tools with schemas and metadata
+            - **Plain functions**: Automatically converted to tools with inferred schemas
+            - **Pydantic model classes**: Treated as structured output tools
         name: The name identifier for this node in the graph. Used for debugging
             and visualization. Defaults to "tools".
         tags: Optional metadata tags to associate with the node for filtering
@@ -255,21 +290,24 @@ class ToolNode(RunnableCallable):
         handle_tool_errors: Configuration for error handling during tool execution.
             Defaults to True. Supports multiple strategies:
 
-            - True: Catch all errors and return a ToolMessage with the default
+            - **True**: Catch all errors and return a ToolMessage with the default
               error template containing the exception details.
-            - str: Catch all errors and return a ToolMessage with this custom
+            - **str**: Catch all errors and return a ToolMessage with this custom
               error message string.
-            - tuple[type[Exception], ...]: Only catch exceptions of the specified
+            - **tuple[type[Exception], ...]**: Only catch exceptions with the specified
               types and return default error messages for them.
-            - Callable[..., str]: Catch exceptions matching the callable's signature
+            - **Callable[..., str]**: Catch exceptions matching the callable's signature
               and return the string result of calling it with the exception.
-            - False: Disable error handling entirely, allowing exceptions to propagate.
+            - **False**: Disable error handling entirely, allowing exceptions to
+              propagate.
 
         messages_key: The key in the state dictionary that contains the message list.
-            This same key will be used for the output ToolMessages. Defaults to "messages".
+            This same key will be used for the output ToolMessages.
+            Defaults to "messages".
+            Allows custom state schemas with different message field names.
 
-    Example:
-        Basic usage with simple tools:
+    Examples:
+        Basic usage:
 
         ```python
         from langgraph.prebuilt import ToolNode
@@ -283,42 +321,35 @@ class ToolNode(RunnableCallable):
         tool_node = ToolNode([calculator])
         ```
 
-        Custom error handling:
+        State injection:
 
         ```python
-        def handle_math_errors(e: ZeroDivisionError) -> str:
-            return "Cannot divide by zero!"
+        from typing_extensions import Annotated
+        from langgraph.prebuilt import InjectedState
 
-        tool_node = ToolNode([calculator], handle_tool_errors=handle_math_errors)
+        @tool
+        def context_tool(query: str, state: Annotated[dict, InjectedState]) -> str:
+            \"\"\"Some tool that uses state.\"\"\"
+            return f"Query: {query}, Messages: {len(state['messages'])}"
+
+        tool_node = ToolNode([context_tool])
         ```
 
-        Direct tool call execution:
+        Error handling:
 
         ```python
-        tool_calls = [{"name": "calculator", "args": {"a": 5, "b": 3}, "id": "1", "type": "tool_call"}]
-        result = tool_node.invoke(tool_calls)
+        def handle_errors(e: ValueError) -> str:
+            return "Invalid input provided"
+
+        tool_node = ToolNode([my_tool], handle_tool_errors=handle_errors)
         ```
-
-    Note:
-        The ToolNode expects input in one of three formats:
-        1. A dictionary with a messages key containing a list of messages
-        2. A list of messages directly
-        3. A list of tool call dictionaries
-
-        When using message formats, the last message must be an AIMessage with
-        tool_calls populated. The node automatically extracts and processes these
-        tool calls concurrently.
-
-        For advanced use cases involving state injection or store access, tools
-        can be annotated with InjectedState or InjectedStore to receive graph
-        context automatically.
     """
 
-    name: str = "ToolNode"
+    name: str = "tools"
 
     def __init__(
         self,
-        tools: Sequence[Union[BaseTool, Callable]],
+        tools: Sequence[Union[BaseTool, BaseModel, Callable]],
         *,
         name: str = "tools",
         tags: Optional[list[str]] = None,
@@ -337,17 +368,36 @@ class ToolNode(RunnableCallable):
             messages_key: State key containing messages.
         """
         super().__init__(self._func, self._afunc, name=name, tags=tags, trace=False)
-        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.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_)
+        self._tools_by_name: dict[str, BaseTool] = {}
+        self._structured_output_tools_by_name: dict[str, type[BaseModel]] = {}
+        self._tool_to_state_args: dict[str, dict[str, Optional[str]]] = {}
+        self._tool_to_store_arg: dict[str, Optional[str]] = {}
+        self._handle_tool_errors = handle_tool_errors
+        self._messages_key = messages_key
+        for tool in tools:
+            if inspect.isclass(tool) and issubclass(tool, BaseModel):
+                # Handle Pydantic model classes as structured output tools
+                self._structured_output_tools_by_name[tool.__name__] = tool
+                self._tool_to_state_args[tool.__name__] = {}
+                self._tool_to_store_arg[tool.__name__] = None
+            else:
+                if not isinstance(tool, BaseTool):
+                    tool_ = create_tool(cast(Type[BaseTool], tool))
+                else:
+                    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_)
+
+    @property
+    def tools_by_name(self) -> dict[str, BaseTool]:
+        """Mapping from tool name to BaseTool instance."""
+        return self._tools_by_name
+
+    @property
+    def structured_output_tools(self) -> dict[str, type[BaseModel]]:
+        """Mapping from structured output tool name to Pydantic model class."""
+        return self._structured_output_tools_by_name
 
     def _func(
         self,
@@ -390,14 +440,14 @@ 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
         # compatibility
         if not any(isinstance(output, Command) for output in outputs):
             # TypedDict, pydantic, dataclass, etc. should all be able to load from dict
-            return outputs if input_type == "list" else {self.messages_key: outputs}
+            return outputs if input_type == "list" else {self._messages_key: outputs}
 
         # LangGraph will automatically handle list of Command and non-command node
         # updates
@@ -425,7 +475,7 @@ class ToolNode(RunnableCallable):
                     combined_outputs.append(output)
             else:
                 combined_outputs.append(
-                    [output] if input_type == "list" else {self.messages_key: [output]}
+                    [output] if input_type == "list" else {self._messages_key: [output]}
                 )
 
         if parent_command:
@@ -437,13 +487,31 @@ 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._structured_output_tools_by_name[call["name"]]
+            return Command(
+                update={
+                    "messages": [
+                        ToolMessage(
+                            content="ok!",
+                            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)
+            tool = self.tools_by_name[call["name"]]
+            response = tool.invoke(call_args, config)
 
         # GraphInterrupt is a special exception that will always be raised.
         # It can be triggered in the following scenarios,
@@ -455,20 +523,20 @@ class ToolNode(RunnableCallable):
         except GraphBubbleUp as e:
             raise e
         except Exception as e:
-            if isinstance(self.handle_tool_errors, tuple):
-                handled_types: tuple = self.handle_tool_errors
-            elif callable(self.handle_tool_errors):
-                handled_types = _infer_handled_types(self.handle_tool_errors)
+            if isinstance(self._handle_tool_errors, tuple):
+                handled_types: tuple = self._handle_tool_errors
+            elif callable(self._handle_tool_errors):
+                handled_types = _infer_handled_types(self._handle_tool_errors)
             else:
                 # default behavior is catching all exceptions
                 handled_types = (Exception,)
 
             # Unhandled
-            if not self.handle_tool_errors or not isinstance(e, handled_types):
+            if not self._handle_tool_errors or not isinstance(e, handled_types):
                 raise e
             # Handled
             else:
-                content = _handle_tool_error(e, flag=self.handle_tool_errors)
+                content = _handle_tool_error(e, flag=self._handle_tool_errors)
             return ToolMessage(
                 content=content,
                 name=call["name"],
@@ -493,15 +561,31 @@ 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._structured_output_tools_by_name[call["name"]]
+            return Command(
+                update={
+                    "messages": [
+                        ToolMessage(
+                            content="ok!",
+                            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)
-
+            tool = self.tools_by_name[call["name"]]
+            response = await tool.ainvoke(call_args, config)
         # GraphInterrupt is a special exception that will always be raised.
         # It can be triggered in the following scenarios,
         # Where GraphInterrupt(GraphBubbleUp) is raised from an `interrupt` invocation most commonly:
@@ -512,20 +596,20 @@ class ToolNode(RunnableCallable):
         except GraphBubbleUp as e:
             raise e
         except Exception as e:
-            if isinstance(self.handle_tool_errors, tuple):
-                handled_types: tuple = self.handle_tool_errors
-            elif callable(self.handle_tool_errors):
-                handled_types = _infer_handled_types(self.handle_tool_errors)
+            if isinstance(self._handle_tool_errors, tuple):
+                handled_types: tuple = self._handle_tool_errors
+            elif callable(self._handle_tool_errors):
+                handled_types = _infer_handled_types(self._handle_tool_errors)
             else:
                 # default behavior is catching all exceptions
                 handled_types = (Exception,)
 
             # Unhandled
-            if not self.handle_tool_errors or not isinstance(e, handled_types):
+            if not self._handle_tool_errors or not isinstance(e, handled_types):
                 raise e
             # Handled
             else:
-                content = _handle_tool_error(e, flag=self.handle_tool_errors)
+                content = _handle_tool_error(e, flag=self._handle_tool_errors)
 
             return ToolMessage(
                 content=content,
@@ -564,9 +648,11 @@ class ToolNode(RunnableCallable):
             else:
                 input_type = "list"
                 messages = input
-        elif isinstance(input, dict) and (messages := input.get(self.messages_key, [])):
+        elif isinstance(input, dict) and (
+            messages := input.get(self._messages_key, [])
+        ):
             input_type = "dict"
-        elif messages := getattr(input, self.messages_key, []):
+        elif messages := getattr(input, self._messages_key, []):
             # Assume dataclass-like state that can coerce from dict
             input_type = "dict"
         else:
@@ -586,10 +672,17 @@ class ToolNode(RunnableCallable):
         return tool_calls, input_type
 
     def _validate_tool_call(self, call: ToolCall) -> Optional[ToolMessage]:
-        if (requested_tool := call["name"]) not in self.tools_by_name:
+        requested_tool = call["name"]
+        if (
+            requested_tool not in self.tools_by_name
+            and requested_tool not in self._structured_output_tools_by_name
+        ):
+            all_tool_names = list(self.tools_by_name.keys()) + list(
+                self._structured_output_tools_by_name.keys()
+            )
             content = INVALID_TOOL_NAME_ERROR_TEMPLATE.format(
                 requested_tool=requested_tool,
-                available_tools=", ".join(self.tools_by_name.keys()),
+                available_tools=", ".join(all_tool_names),
             )
             return ToolMessage(
                 content, name=requested_tool, tool_call_id=call["id"], status="error"
@@ -606,15 +699,15 @@ class ToolNode(RunnableCallable):
             BaseModel,
         ],
     ) -> ToolCall:
-        state_args = self.tool_to_state_args[tool_call["name"]]
+        state_args = self._tool_to_state_args[tool_call["name"]]
         if state_args and isinstance(input, list):
             required_fields = list(state_args.values())
             if (
                 len(required_fields) == 1
-                and required_fields[0] == self.messages_key
+                and required_fields[0] == self._messages_key
                 or required_fields[0] is None
             ):
-                input = {self.messages_key: input}
+                input = {self._messages_key: input}
             else:
                 err_msg = (
                     f"Invalid input to ToolNode. Tool {tool_call['name']} requires "
@@ -645,7 +738,7 @@ class ToolNode(RunnableCallable):
     def _inject_store(
         self, tool_call: ToolCall, store: Optional[BaseStore]
     ) -> ToolCall:
-        store_arg = self.tool_to_store_arg[tool_call["name"]]
+        store_arg = self._tool_to_store_arg[tool_call["name"]]
         if not store_arg:
             return tool_call
 
@@ -704,7 +797,10 @@ class ToolNode(RunnableCallable):
             The injection is performed on a copy of the tool call to avoid mutating
             the original.
         """
-        if tool_call["name"] not in self.tools_by_name:
+        if (
+            tool_call["name"] not in self.tools_by_name
+            and tool_call["name"] not in self._structured_output_tools_by_name
+        ):
             return tool_call
 
         tool_call_copy: ToolCall = copy(tool_call)
@@ -722,15 +818,15 @@ class ToolNode(RunnableCallable):
             # input type is dict when ToolNode is invoked with a dict input (e.g. {"messages": [AIMessage(..., tool_calls=[...])]})
             if input_type not in ("dict", "tool_calls"):
                 raise ValueError(
-                    f"Tools can provide a dict in Command.update only when using dict with '{self.messages_key}' key as ToolNode input, "
+                    f"Tools can provide a dict in Command.update only when using dict with '{self._messages_key}' key as ToolNode input, "
                     f"got: {command.update} for tool '{call['name']}'"
                 )
 
             updated_command = deepcopy(command)
             state_update = cast(dict[str, Any], updated_command.update) or {}
-            messages_update = state_update.get(self.messages_key, [])
+            messages_update = state_update.get(self._messages_key, [])
         elif isinstance(command.update, list):
-            # input type is list when ToolNode is invoked with a list input (e.g. [AIMessage(..., tool_calls=[...])])
+            # Input type is list when ToolNode is invoked with a list input (e.g. [AIMessage(..., tool_calls=[...])])
             if input_type != "list":
                 raise ValueError(
                     f"Tools can provide a list of messages in Command.update only when using list of messages as ToolNode input, "

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
 
@@ -62,7 +86,8 @@ def tool5(some_val: int):
 tool5.handle_tool_error = "foo"
 
 
-async def test_tool_node():
+async def test_tool_node() -> None:
+    """Test tool node."""
     result = ToolNode([tool1]).invoke(
         {
             "messages": [
@@ -154,7 +179,7 @@ async def test_tool_node():
     assert tool_message.tool_call_id == "some 3"
 
 
-async def test_tool_node_tool_call_input():
+async def test_tool_node_tool_call_input() -> None:
     # Single tool call
     tool_call_1 = {
         "name": "tool1",
@@ -195,7 +220,7 @@ async def test_tool_node_tool_call_input():
     ]
 
 
-async def test_tool_node_error_handling():
+async def test_tool_node_error_handling() -> None:
     def handle_all(e: Union[ValueError, ToolException, ValidationError]):
         return TOOL_CALL_ERROR_TEMPLATE.format(error=repr(e))
 
@@ -257,7 +282,7 @@ async def test_tool_node_error_handling():
         assert result_error["messages"][2].tool_call_id == "another id"
 
 
-async def test_tool_node_error_handling_callable():
+async def test_tool_node_error_handling_callable() -> None:
     def handle_value_error(e: ValueError):
         return "Value error"
 
@@ -1156,3 +1181,395 @@ 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",
+                        ),
+                    ],
+                },
+            },
+        ),
+    ]
+
+
+def test_structured_output_tools_sync() -> None:
+    """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.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"
+
+
+async def test_structured_output_tools_async() -> None:
+    """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" not 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 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"
+    assert structured_response_async.age == 30
+    assert structured_response_async.location == "NYC"