feat(prebuilt): add support for handle_tool_errors with single exception type (#6021)

This PR extends the `ToolNode` `handle_tool_errors` property to also
accept a single exception type: `type[Exception]`.

Example:
```python
tool_node_with_error_handling = ToolNode(
    [tool1, tool2, tool3], handle_tool_errors=ValueError
).invoke(...)
"""A tool node that will only retry a failed tool call if a ValueError is raised"""
```

libs/langgraph/bench/react_agent.py

@@ -10,7 +10,7 @@ from langchain_core.outputs import ChatGeneration, ChatResult
 from langchain_core.tools import StructuredTool
 
 from langgraph.checkpoint.base import BaseCheckpointSaver
-from langgraph.prebuilt.chat_agent_executor import create_react_agent
+from langgraph.prebuilt.chat_agent_executor import create_agent
 from langgraph.pregel import Pregel
 
 
@@ -60,7 +60,7 @@ def react_agent(n_tools: int, checkpointer: Optional[BaseCheckpointSaver]) -> Pr
         ]
     )
 
-    return create_react_agent(model, [tool], checkpointer=checkpointer)
+    return create_agent(model, [tool], checkpointer=checkpointer)
 
 
 if __name__ == "__main__":

libs/prebuilt/langgraph/prebuilt/_internal/_typing.py

@@ -0,0 +1,10 @@
+from __future__ import annotations
+
+from typing import Awaitable, Callable, TypeVar, Union
+
+from typing_extensions import ParamSpec
+
+P = ParamSpec("P")
+R = TypeVar("R")
+
+SyncOrAsync = Callable[P, Union[R, Awaitable[R]]]

libs/prebuilt/langgraph/prebuilt/responses.py

@@ -0,0 +1,397 @@
+"""Types for setting agent response formats."""
+
+from __future__ import annotations
+
+import sys
+import uuid
+from dataclasses import dataclass, is_dataclass
+from typing import (
+    Any,
+    Callable,
+    Generic,
+    Iterable,
+    Literal,
+    TypeVar,
+    Union,
+    get_args,
+    get_origin,
+)
+
+from langchain_core.messages import AIMessage
+from langchain_core.tools import BaseTool, StructuredTool
+from pydantic import BaseModel, TypeAdapter
+from typing_extensions import Self, is_typeddict
+
+# Supported schema types: Pydantic models, dataclasses, TypedDict, JSON schema dicts
+SchemaT = TypeVar("SchemaT")
+
+
+if sys.version_info >= (3, 10):
+    from types import UnionType
+else:
+    UnionType = Union
+
+SchemaKind = Literal["pydantic", "dataclass", "typeddict", "json_schema"]
+
+
+class StructuredOutputError(Exception):
+    """Base class for structured output errors."""
+
+
+class MultipleStructuredOutputsError(StructuredOutputError):
+    """Raised when model returns multiple structured output tool calls when only one is expected."""
+
+    def __init__(self, tool_names: list[str]):
+        self.tool_names = tool_names
+        super().__init__(
+            f"Model incorrectly returned multiple structured responses ({', '.join(tool_names)}) when only one is expected."
+        )
+
+
+class StructuredOutputParsingError(StructuredOutputError):
+    """Raised when structured output tool call arguments fail to parse according to the schema."""
+
+    def __init__(self, tool_name: str, parse_error: Exception):
+        self.tool_name = tool_name
+        self.parse_error = parse_error
+        super().__init__(
+            f"Failed to parse structured output for tool '{tool_name}': {parse_error}."
+        )
+
+
+def _parse_with_schema(
+    schema: Union[type[SchemaT], dict], schema_kind: SchemaKind, data: dict[str, Any]
+) -> Any:
+    """Parse data using for any supported schema type.
+
+    Args:
+        schema: The schema type (Pydantic model, dataclass, or TypedDict)
+        data: The data to parse
+
+    Returns:
+        The parsed instance according to the schema type
+
+    Raises:
+        ValueError: If parsing fails
+    """
+    if schema_kind == "json_schema":
+        return data
+    else:
+        try:
+            adapter: TypeAdapter[SchemaT] = TypeAdapter(schema)
+            return adapter.validate_python(data)
+        except Exception as e:
+            schema_name = getattr(schema, "__name__", str(schema))
+            raise ValueError(f"Failed to parse data to {schema_name}: {e}") from e
+
+
+@dataclass(init=False)
+class _SchemaSpec(Generic[SchemaT]):
+    """Describes a structured output schema."""
+
+    schema: type[SchemaT]
+    """The schema for the response, can be a Pydantic model, dataclass, TypedDict, or JSON schema dict."""
+
+    name: str
+    """Name of the schema, used for tool calling.
+    
+    If not provided, the name will be the model name or "response_format" if it's a JSON schema.
+    """
+
+    description: str
+    """Custom description of the schema. 
+    
+    If not provided, provided will use the model's docstring.
+    """
+
+    schema_kind: SchemaKind
+    """The kind of schema."""
+
+    json_schema: dict[str, Any]
+    """JSON schema associated with the schema."""
+
+    strict: bool = False
+    """Whether to enforce strict validation of the schema."""
+
+    def __init__(
+        self,
+        schema: type[SchemaT],
+        *,
+        name: str | None = None,
+        description: str | None = None,
+        strict: bool = False,
+    ) -> None:
+        """Initialize SchemaSpec with schema and optional parameters."""
+        self.schema = schema
+
+        if name:
+            self.name = name
+        elif isinstance(schema, dict):
+            self.name = str(
+                schema.get("title", f"response_format_{str(uuid.uuid4())[:4]}")
+            )
+        else:
+            self.name = str(
+                getattr(schema, "__name__", f"response_format_{str(uuid.uuid4())[:4]}")
+            )
+
+        self.description = description or (
+            schema.get("description", "")
+            if isinstance(schema, dict)
+            else getattr(schema, "__doc__", None) or ""
+        )
+
+        self.strict = strict
+
+        if isinstance(schema, dict):
+            self.schema_kind = "json_schema"
+            self.json_schema = schema
+        elif isinstance(schema, type) and issubclass(schema, BaseModel):
+            self.schema_kind = "pydantic"
+            self.json_schema = schema.model_json_schema()
+        elif is_dataclass(schema):
+            self.schema_kind = "dataclass"
+            self.json_schema = TypeAdapter(schema).json_schema()
+        elif is_typeddict(schema):
+            self.schema_kind = "typeddict"
+            self.json_schema = TypeAdapter(schema).json_schema()
+        else:
+            raise ValueError(
+                f"Unsupported schema type: {type(schema)}. "
+                f"Supported types: Pydantic models, dataclasses, TypedDicts, and JSON schema dicts."
+            )
+
+
+@dataclass(init=False)
+class ToolOutput(Generic[SchemaT]):
+    """Use a tool calling strategy for model responses."""
+
+    schema: type[SchemaT]
+    """Schema for the tool calls."""
+
+    schema_specs: list[_SchemaSpec[SchemaT]]
+    """Schema specs for the tool calls."""
+
+    tool_message_content: str | None
+    """The content of the tool message to be returned when the model calls an artificial structured output tool."""
+
+    handle_errors: Union[
+        bool,
+        str,
+        type[Exception],
+        tuple[type[Exception], ...],
+        Callable[[Exception], str],
+    ]
+    """Error handling strategy for structured output via ToolOutput. Default is True.
+    
+    - True: Catch all errors with default error template
+    - str: Catch all errors with this custom message
+    - type[Exception]: Only catch this exception type with default message
+    - tuple[type[Exception], ...]: Only catch these exception types with default message
+    - Callable[[Exception], str]: Custom function that returns error message
+    - False: No retry, let exceptions propagate
+    """
+
+    def __init__(
+        self,
+        schema: type[SchemaT],
+        tool_message_content: str | None = None,
+        handle_errors: Union[
+            bool,
+            str,
+            type[Exception],
+            tuple[type[Exception], ...],
+            Callable[[Exception], str],
+        ] = True,
+    ) -> None:
+        """Initialize ToolOutput with schemas, tool message content, and error handling strategy."""
+        self.schema = schema
+        self.tool_message_content = tool_message_content
+        self.handle_errors = handle_errors
+
+        def _iter_variants(schema: Any) -> Iterable[Any]:
+            """Yield leaf variants from Union and JSON Schema oneOf."""
+
+            if get_origin(schema) in (UnionType, Union):
+                for arg in get_args(schema):
+                    yield from _iter_variants(arg)
+                return
+
+            if isinstance(schema, dict) and "oneOf" in schema:
+                for sub in schema.get("oneOf", []):
+                    yield from _iter_variants(sub)
+                return
+
+            yield schema
+
+        self.schema_specs = [_SchemaSpec(s) for s in _iter_variants(schema)]
+
+
+@dataclass(init=False)
+class NativeOutput(Generic[SchemaT]):
+    """Use the model provider's native structured output method."""
+
+    schema: type[SchemaT]
+    """Schema for native mode."""
+
+    schema_spec: _SchemaSpec[SchemaT]
+    """Schema spec for native mode."""
+
+    def __init__(
+        self,
+        schema: type[SchemaT],
+    ) -> None:
+        self.schema = schema
+        self.schema_spec = _SchemaSpec(schema)
+
+    def to_model_kwargs(self) -> dict[str, Any]:
+        # OpenAI:
+        # - see https://platform.openai.com/docs/guides/structured-outputs
+        response_format = {
+            "type": "json_schema",
+            "json_schema": {
+                "name": self.schema_spec.name,
+                "schema": self.schema_spec.json_schema,
+            },
+        }
+        return {"response_format": response_format}
+
+
+@dataclass
+class OutputToolBinding(Generic[SchemaT]):
+    """Information for tracking structured output tool metadata.
+
+    This contains all necessary information to handle structured responses
+    generated via tool calls, including the original schema, its type classification,
+    and the corresponding tool implementation used by the tools strategy.
+    """
+
+    schema: type[SchemaT]
+    """The original schema provided for structured output (Pydantic model, dataclass, TypedDict, or JSON schema dict)."""
+
+    schema_kind: SchemaKind
+    """Classification of the schema type for proper response construction."""
+
+    tool: BaseTool
+    """LangChain tool instance created from the schema for model binding."""
+
+    @classmethod
+    def from_schema_spec(cls, schema_spec: _SchemaSpec[SchemaT]) -> Self:
+        """Create an OutputToolBinding instance from a SchemaSpec.
+
+        Args:
+            schema_spec: The SchemaSpec to convert
+
+        Returns:
+            An OutputToolBinding instance with the appropriate tool created
+        """
+        return cls(
+            schema=schema_spec.schema,
+            schema_kind=schema_spec.schema_kind,
+            tool=StructuredTool(
+                args_schema=schema_spec.json_schema,
+                name=schema_spec.name,
+                description=schema_spec.description,
+            ),
+        )
+
+    def parse(self, tool_args: dict[str, Any]) -> SchemaT:
+        """Parse tool arguments according to the schema.
+
+        Args:
+            tool_args: The arguments from the tool call
+
+        Returns:
+            The parsed response according to the schema type
+
+        Raises:
+            ValueError: If parsing fails
+        """
+        return _parse_with_schema(self.schema, self.schema_kind, tool_args)
+
+
+@dataclass
+class NativeOutputBinding(Generic[SchemaT]):
+    """Information for tracking native structured output metadata.
+
+    This contains all necessary information to handle structured responses
+    generated via native provider output, including the original schema,
+    its type classification, and parsing logic for provider-enforced JSON.
+    """
+
+    schema: type[SchemaT]
+    """The original schema provided for structured output (Pydantic model, dataclass, TypedDict, or JSON schema dict)."""
+
+    schema_kind: SchemaKind
+    """Classification of the schema type for proper response construction."""
+
+    @classmethod
+    def from_schema_spec(cls, schema_spec: _SchemaSpec[SchemaT]) -> Self:
+        """Create a NativeOutputBinding instance from a SchemaSpec.
+
+        Args:
+            schema_spec: The SchemaSpec to convert
+
+        Returns:
+            A NativeOutputBinding instance for parsing native structured output
+        """
+        return cls(
+            schema=schema_spec.schema,
+            schema_kind=schema_spec.schema_kind,
+        )
+
+    def parse(self, response: AIMessage) -> SchemaT:
+        """Parse AIMessage content according to the schema.
+
+        Args:
+            response: The AI message containing the structured output
+
+        Returns:
+            The parsed response according to the schema
+
+        Raises:
+            ValueError: If text extraction, JSON parsing or schema validation fails
+        """
+        # Extract text content from AIMessage and parse as JSON
+        raw_text = self._extract_text_content_from_message(response)
+
+        import json
+
+        try:
+            data = json.loads(raw_text)
+        except Exception as e:
+            schema_name = getattr(self.schema, "__name__", "response_format")
+            raise ValueError(
+                f"Native structured output expected valid JSON for {schema_name}, but parsing failed: {e}."
+            ) from e
+
+        # Parse according to schema
+        return _parse_with_schema(self.schema, self.schema_kind, data)
+
+    def _extract_text_content_from_message(self, message: AIMessage) -> str:
+        """Extract text content from an AIMessage.
+
+        Args:
+            message: The AI message to extract text from
+
+        Returns:
+            The extracted text content
+        """
+        content = message.content
+        if isinstance(content, str):
+            return content
+        if isinstance(content, list):
+            parts: list[str] = []
+            for c in content:
+                if isinstance(c, dict):
+                    if c.get("type") == "text" and "text" in c:
+                        parts.append(str(c["text"]))
+                    elif "content" in c and isinstance(c["content"], str):
+                        parts.append(c["content"])
+                else:
+                    parts.append(str(c))
+            return "".join(parts)
+        return str(content)
+
+
+ResponseFormat = Union[ToolOutput[SchemaT], NativeOutput[SchemaT]]

libs/prebuilt/langgraph/prebuilt/tool_node.py

@@ -31,6 +31,8 @@ Typical Usage:
     ```
 """
 
+from __future__ import annotations
+
 import asyncio
 import inspect
 import json
@@ -68,7 +70,7 @@ from langchain_core.tools.base import (
     TOOL_MESSAGE_BLOCK_TYPES,
     get_all_basemodel_annotations,
 )
-from pydantic import BaseModel
+from pydantic import BaseModel, ValidationError
 from typing_extensions import Annotated, get_args, get_origin
 
 from langgraph._internal._runnable import RunnableCallable
@@ -81,6 +83,8 @@ INVALID_TOOL_NAME_ERROR_TEMPLATE = (
     "Error: {requested_tool} is not a valid tool, try one of [{available_tools}]."
 )
 TOOL_CALL_ERROR_TEMPLATE = "Error: {error}\n Please fix your mistakes."
+TOOL_EXECUTION_ERROR_TEMPLATE = "Error executing tool '{tool_name}' with kwargs {tool_kwargs} with error:\n {error}\n Please fix the error and try again."
+TOOL_INVOCATION_ERROR_TEMPLATE = "Error invoking tool '{tool_name}' with kwargs {tool_kwargs} with error:\n {error}\n Please fix the error and try again."
 
 
 def msg_content_output(output: Any) -> Union[str, list[dict]]:
@@ -122,6 +126,30 @@ def msg_content_output(output: Any) -> Union[str, list[dict]]:
             return str(output)
 
 
+class ToolInvocationError(Exception):
+    """Exception raised when a tool invocation fails due to invalid arguments."""
+
+    def __init__(self, tool_name: str, error: Exception, tool_kwargs: dict[str, Any]):
+        self.message = TOOL_INVOCATION_ERROR_TEMPLATE.format(
+            tool_name=tool_name, tool_kwargs=tool_kwargs, error=error
+        )
+        self.tool_name = tool_name
+        self.tool_kwargs = tool_kwargs
+        self.error = error
+        super().__init__(self.message)
+
+
+def _default_handle_tool_errors(e: Exception) -> str:
+    """Default error handler for tool errors.
+
+    If the tool is a tool invocation error, return its message.
+    Otherwise, raise the error.
+    """
+    if isinstance(e, ToolInvocationError):
+        return e.message
+    raise e
+
+
 def _handle_tool_error(
     e: Exception,
     *,
@@ -129,6 +157,7 @@ def _handle_tool_error(
         bool,
         str,
         Callable[..., str],
+        type[Exception],
         tuple[type[Exception], ...],
     ],
 ) -> str:
@@ -156,12 +185,14 @@ def _handle_tool_error(
         The tuple case is handled by the caller through exception type checking,
         not by this function directly.
     """
-    if isinstance(flag, (bool, tuple)):
+    if isinstance(flag, (bool, tuple)) or (
+        isinstance(flag, type) and issubclass(flag, Exception)
+    ):
         content = TOOL_CALL_ERROR_TEMPLATE.format(error=repr(e))
     elif isinstance(flag, str):
         content = flag
     elif callable(flag):
-        content = flag(e)
+        content = flag(e)  # type: ignore [assignment, call-arg]
     else:
         raise ValueError(
             f"Got unexpected type of `handle_tool_error`. Expected bool, str "
@@ -237,39 +268,69 @@ 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
+
+    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
 
     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
         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
             and organization. Defaults to None.
         handle_tool_errors: Configuration for error handling during tool execution.
-            Defaults to True. Supports multiple strategies:
+            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
+            - **type[Exception]**: Only catch exceptions with the specified type and return the default error message for it.
+            - **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.
+
+            Defaults to a callable that:
+                - catches tool invocation errors (due to invalid arguments provided by the model) and returns a descriptive error message
+                - ignores tool execution errors (they will be re-raised)
 
         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,38 +344,31 @@ 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,
@@ -323,8 +377,8 @@ class ToolNode(RunnableCallable):
         name: str = "tools",
         tags: Optional[list[str]] = None,
         handle_tool_errors: Union[
-            bool, str, Callable[..., str], tuple[type[Exception], ...]
-        ] = True,
+            bool, str, Callable[..., str], type[Exception], tuple[type[Exception], ...]
+        ] = _default_handle_tool_errors,
         messages_key: str = "messages",
     ) -> None:
         """Initialize the ToolNode with the provided tools and configuration.
@@ -337,17 +391,24 @@ 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._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(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
 
     def _func(
         self,
@@ -390,14 +451,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 +486,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 +498,19 @@ 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
+
         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"]]
+
+            try:
+                response = tool.invoke(call_args, config)
+            except ValidationError as exc:
+                raise ToolInvocationError(call["name"], exc, call["args"])
 
         # GraphInterrupt is a special exception that will always be raised.
         # It can be triggered in the following scenarios,
@@ -455,20 +522,27 @@ 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)
+            handled_types: tuple[type[Exception], ...]
+            if isinstance(self._handle_tool_errors, type) and issubclass(
+                self._handle_tool_errors, Exception
+            ):
+                handled_types = (self._handle_tool_errors,)
+            elif isinstance(self._handle_tool_errors, tuple):
+                handled_types = self._handle_tool_errors
+            elif callable(self._handle_tool_errors) and not isinstance(
+                self._handle_tool_errors, type
+            ):
+                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,14 +567,19 @@ 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
 
         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"]]
+
+            try:
+                response = await tool.ainvoke(call_args, config)
+            except ValidationError as exc:
+                raise ToolInvocationError(call["name"], exc, call["args"])
 
         # GraphInterrupt is a special exception that will always be raised.
         # It can be triggered in the following scenarios,
@@ -512,20 +591,27 @@ 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)
+            handled_types: tuple[type[Exception], ...]
+            if isinstance(self._handle_tool_errors, type) and issubclass(
+                self._handle_tool_errors, Exception
+            ):
+                handled_types = (self._handle_tool_errors,)
+            elif isinstance(self._handle_tool_errors, tuple):
+                handled_types = self._handle_tool_errors
+            elif callable(self._handle_tool_errors) and not isinstance(
+                self._handle_tool_errors, type
+            ):
+                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 +650,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 +674,12 @@ 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:
+            all_tool_names = list(self.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 +696,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 +735,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
 
@@ -722,15 +812,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

@@ -1,5 +1,5 @@
 # serializer version: 1
-# name: test_react_agent_graph_structure[None-None-None-tools0]
+# name: test_react_agent_graph_structure[None-None-tools0]
   '''
   graph TD;
   	__start__ --> agent;
@@ -7,7 +7,7 @@
   
   '''
 # ---
-# name: test_react_agent_graph_structure[None-None-None-tools1]
+# name: test_react_agent_graph_structure[None-None-tools1]
   '''
   graph TD;
   	__start__ --> agent;
@@ -17,7 +17,7 @@
   
   '''
 # ---
-# name: test_react_agent_graph_structure[None-None-pre_model_hook-tools0]
+# name: test_react_agent_graph_structure[None-pre_model_hook-tools0]
   '''
   graph TD;
   	__start__ --> pre_model_hook;
@@ -26,7 +26,7 @@
   
   '''
 # ---
-# name: test_react_agent_graph_structure[None-None-pre_model_hook-tools1]
+# name: test_react_agent_graph_structure[None-pre_model_hook-tools1]
   '''
   graph TD;
   	__start__ --> pre_model_hook;
@@ -37,7 +37,7 @@
   
   '''
 # ---
-# name: test_react_agent_graph_structure[None-post_model_hook-None-tools0]
+# name: test_react_agent_graph_structure[post_model_hook-None-tools0]
   '''
   graph TD;
   	__start__ --> agent;
@@ -46,7 +46,7 @@
   
   '''
 # ---
-# name: test_react_agent_graph_structure[None-post_model_hook-None-tools1]
+# name: test_react_agent_graph_structure[post_model_hook-None-tools1]
   '''
   graph TD;
   	__start__ --> agent;
@@ -58,7 +58,7 @@
   
   '''
 # ---
-# name: test_react_agent_graph_structure[None-post_model_hook-pre_model_hook-tools0]
+# name: test_react_agent_graph_structure[post_model_hook-pre_model_hook-tools0]
   '''
   graph TD;
   	__start__ --> pre_model_hook;
@@ -68,7 +68,7 @@
   
   '''
 # ---
-# name: test_react_agent_graph_structure[None-post_model_hook-pre_model_hook-tools1]
+# name: test_react_agent_graph_structure[post_model_hook-pre_model_hook-tools1]
   '''
   graph TD;
   	__start__ --> pre_model_hook;
@@ -81,93 +81,3 @@
   
   '''
 # ---
-# name: test_react_agent_graph_structure[ResponseFormat-None-None-tools0]
-  '''
-  graph TD;
-  	__start__ --> agent;
-  	agent --> generate_structured_response;
-  	generate_structured_response --> __end__;
-  
-  '''
-# ---
-# name: test_react_agent_graph_structure[ResponseFormat-None-None-tools1]
-  '''
-  graph TD;
-  	__start__ --> agent;
-  	agent -.-> generate_structured_response;
-  	agent -.-> tools;
-  	tools --> agent;
-  	generate_structured_response --> __end__;
-  
-  '''
-# ---
-# name: test_react_agent_graph_structure[ResponseFormat-None-pre_model_hook-tools0]
-  '''
-  graph TD;
-  	__start__ --> pre_model_hook;
-  	agent --> generate_structured_response;
-  	pre_model_hook --> agent;
-  	generate_structured_response --> __end__;
-  
-  '''
-# ---
-# name: test_react_agent_graph_structure[ResponseFormat-None-pre_model_hook-tools1]
-  '''
-  graph TD;
-  	__start__ --> pre_model_hook;
-  	agent -.-> generate_structured_response;
-  	agent -.-> tools;
-  	pre_model_hook --> agent;
-  	tools --> pre_model_hook;
-  	generate_structured_response --> __end__;
-  
-  '''
-# ---
-# name: test_react_agent_graph_structure[ResponseFormat-post_model_hook-None-tools0]
-  '''
-  graph TD;
-  	__start__ --> agent;
-  	agent --> post_model_hook;
-  	post_model_hook --> generate_structured_response;
-  	generate_structured_response --> __end__;
-  
-  '''
-# ---
-# name: test_react_agent_graph_structure[ResponseFormat-post_model_hook-None-tools1]
-  '''
-  graph TD;
-  	__start__ --> agent;
-  	agent --> post_model_hook;
-  	post_model_hook -.-> agent;
-  	post_model_hook -.-> generate_structured_response;
-  	post_model_hook -.-> tools;
-  	tools --> agent;
-  	generate_structured_response --> __end__;
-  
-  '''
-# ---
-# name: test_react_agent_graph_structure[ResponseFormat-post_model_hook-pre_model_hook-tools0]
-  '''
-  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[ResponseFormat-post_model_hook-pre_model_hook-tools1]
-  '''
-  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 -.-> tools;
-  	pre_model_hook --> agent;
-  	tools --> pre_model_hook;
-  	generate_structured_response --> __end__;
-  
-  '''
-# ---

libs/prebuilt/tests/model.py

@@ -1,7 +1,10 @@
+import json
+from dataclasses import asdict, is_dataclass
 from typing import (
     Any,
     Callable,
     Dict,
+    Generic,
     List,
     Literal,
     Optional,
@@ -18,16 +21,16 @@ from langchain_core.messages import (
     ToolCall,
 )
 from langchain_core.outputs import ChatGeneration, ChatResult
-from langchain_core.runnables import Runnable, RunnableLambda
+from langchain_core.runnables import Runnable
 from langchain_core.tools import BaseTool
 from pydantic import BaseModel
 
-from langgraph.prebuilt.chat_agent_executor import StructuredResponse
+from langgraph.prebuilt.chat_agent_executor import StructuredResponseT
 
 
-class FakeToolCallingModel(BaseChatModel):
-    tool_calls: Optional[list[list[ToolCall]]] = None
-    structured_response: Optional[StructuredResponse] = None
+class FakeToolCallingModel(BaseChatModel, Generic[StructuredResponseT]):
+    tool_calls: Optional[Union[list[list[ToolCall]], list[list[dict]]]] = None
+    structured_response: Optional[StructuredResponseT] = None
     index: int = 0
     tool_style: Literal["openai", "anthropic"] = "openai"
 
@@ -39,15 +42,36 @@ class FakeToolCallingModel(BaseChatModel):
         **kwargs: Any,
     ) -> ChatResult:
         """Top Level call"""
-        messages_string = "-".join([m.content for m in messages])
-        tool_calls = (
-            self.tool_calls[self.index % len(self.tool_calls)]
-            if self.tool_calls
-            else []
-        )
-        message = AIMessage(
-            content=messages_string, id=str(self.index), tool_calls=tool_calls.copy()
-        )
+        rf = kwargs.get("response_format")
+        is_native = isinstance(rf, dict) and rf.get("type") == "json_schema"
+
+        if self.tool_calls:
+            if is_native:
+                tool_calls = (
+                    self.tool_calls[self.index]
+                    if self.index < len(self.tool_calls)
+                    else []
+                )
+            else:
+                tool_calls = self.tool_calls[self.index % len(self.tool_calls)]
+        else:
+            tool_calls = []
+
+        if is_native and not tool_calls:
+            if isinstance(self.structured_response, BaseModel):
+                content_obj = self.structured_response.model_dump()
+            elif is_dataclass(self.structured_response):
+                content_obj = asdict(self.structured_response)
+            elif isinstance(self.structured_response, dict):
+                content_obj = self.structured_response
+            message = AIMessage(content=json.dumps(content_obj), id=str(self.index))
+        else:
+            messages_string = "-".join([m.content for m in messages])
+            message = AIMessage(
+                content=messages_string,
+                id=str(self.index),
+                tool_calls=tool_calls.copy(),
+            )
         self.index += 1
         return ChatResult(generations=[ChatGeneration(message=message)])
 
@@ -55,14 +79,6 @@ class FakeToolCallingModel(BaseChatModel):
     def _llm_type(self) -> str:
         return "fake-tool-call-model"
 
-    def with_structured_output(
-        self, schema: Type[BaseModel]
-    ) -> Runnable[LanguageModelInput, StructuredResponse]:
-        if self.structured_response is None:
-            raise ValueError("Structured response is not set")
-
-        return RunnableLambda(lambda x: self.structured_response)
-
     def bind_tools(
         self,
         tools: Sequence[Union[Dict[str, Any], Type[BaseModel], Callable, BaseTool]],

libs/prebuilt/tests/specifications/responses.json

@@ -0,0 +1,87 @@
+[
+    {
+      "name": "updated structured response",
+      "responseFormat": [
+        {
+          "title": "role_schema_structured_output",
+          "type": "object",
+          "properties": {
+            "name": { "type": "string" },
+            "role": { "type": "string" }
+          },
+          "required": ["name", "role"]
+        },
+        {
+          "title": "department_schema_structured_output",
+          "type": "object",
+          "properties": {
+            "name": { "type": "string" },
+            "department": { "type": "string" }
+          },
+          "required": ["name", "department"]
+        }
+      ],
+      "assertionsByInvocation": [
+        {
+          "prompt": "What is the role of Sabine?",
+          "toolsWithExpectedCalls": {
+            "getEmployeeRole": 1,
+            "getEmployeeDepartment": 0
+          },
+          "expectedLastMessage": "Returning structured response: {'name': 'Sabine', 'role': 'Developer'}",
+          "expectedStructuredResponse": { "name": "Sabine", "role": "Developer" },
+          "llmRequestCount": 2
+        },
+        {
+          "prompt": "In which department does Henrik work?",
+          "toolsWithExpectedCalls": {
+            "getEmployeeRole": 1,
+            "getEmployeeDepartment": 1
+          },
+          "expectedLastMessage": "Returning structured response: {'name': 'Henrik', 'department': 'IT'}",
+          "expectedStructuredResponse": { "name": "Henrik", "department": "IT" },
+          "llmRequestCount": 4
+        }
+      ]
+    },
+    {
+      "name": "asking for information that does not fit into the response format",
+      "responseFormat": [
+        {
+          "schema": {
+            "type": "object",
+            "properties": {
+              "name": { "type": "string" },
+              "role": { "type": "string" }
+            },
+            "required": ["name", "role"]
+          }
+        },
+        {
+          "schema": {
+            "type": "object",
+            "properties": {
+              "name": { "type": "string" },
+              "department": { "type": "string" }
+            },
+            "required": ["name", "department"]
+          }
+        }
+      ],
+      "assertionsByInvocation": [
+        {
+          "prompt": "How much does Saskia earn?",
+          "toolsWithExpectedCalls": {
+            "getEmployeeRole": 1,
+            "getEmployeeDepartment": 0
+          },
+          "expectedLastMessage": "Returning structured response: {'name': 'Saskia', 'role': 'Software Engineer'}",
+          "expectedStructuredResponse": {
+            "name": "Saskia",
+            "role": "Software Engineer"
+          },
+          "llmRequestCount": 2
+        }
+      ]
+    }
+  ]

libs/prebuilt/tests/specifications/return_direct.json

@@ -0,0 +1,48 @@
+[
+  {
+    "name": "Scenario: NO return_direct, NO response_format",
+    "returnDirect": false,
+    "responseFormat": null,
+    "expectedToolCalls": 10,
+    "expectedLastMessage": "Attempts: 10",
+    "expectedStructuredResponse": null
+  },
+  {
+    "name": "Scenario: NO return_direct, YES response_format",
+    "returnDirect": false,
+    "responseFormat": {
+      "type": "object",
+      "properties": {
+        "attempts": { "type": "number" },
+        "succeeded": { "type": "boolean" }
+      },
+      "required": ["attempts", "succeeded"]
+    },
+    "expectedToolCalls": 10,
+    "expectedLastMessage": "Returning structured response: {'attempts': 10, 'succeeded': True}",
+    "expectedStructuredResponse": { "attempts": 10, "succeeded": true }
+  },
+  {
+    "name": "Scenario: YES return_direct, NO response_format",
+    "returnDirect": true,
+    "responseFormat": null,
+    "expectedToolCalls": 1,
+    "expectedLastMessage": "{\"status\": \"pending\", \"attempts\": 1}",
+    "expectedStructuredResponse": null
+  },
+  {
+    "name": "Scenario: YES return_direct, YES response_format",
+    "returnDirect": true,
+    "responseFormat": {
+      "type": "object",
+      "properties": {
+        "attempts": { "type": "number" },
+        "succeeded": { "type": "boolean" }
+      },
+      "required": ["attempts", "succeeded"]
+    },
+    "expectedToolCalls": 1,
+    "expectedLastMessage": "{\"status\": \"pending\", \"attempts\": 1}",
+    "expectedStructuredResponse": null
+  }
+]

libs/prebuilt/tests/test_deprecation.py

@@ -14,7 +14,7 @@ class Config(TypedDict):
 @pytest.mark.filterwarnings("ignore:`get_config_jsonschema` is deprecated")
 def test_config_schema_deprecation() -> None:
     with pytest.warns(
-        LangGraphDeprecatedSinceV10,
+        DeprecationWarning,
         match="`config_schema` is deprecated and will be removed. Please use `context_schema` instead.",
     ):
         agent = create_react_agent(FakeToolCallingModel(), [], config_schema=Config)

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,
@@ -25,27 +19,22 @@ from langchain_core.messages import (
     ToolMessage,
 )
 from langchain_core.runnables import RunnableConfig, RunnableLambda
-from langchain_core.tools import InjectedToolCallId, ToolException
+from langchain_core.tools import BaseTool, 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,
     AgentStatePydantic,
-    StateSchemaType,
-    _get_model,
-    _should_bind_tools,
+    StateT,
     _validate_chat_history,
 )
 from langgraph.prebuilt.tool_node import (
@@ -184,7 +173,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
@@ -275,7 +264,7 @@ async def test_prompt_with_store_async():
 @pytest.mark.parametrize("tool_style", ["openai", "anthropic"])
 @pytest.mark.parametrize("version", REACT_TOOL_CALL_VERSIONS)
 @pytest.mark.parametrize("include_builtin", [True, False])
-def test_model_with_tools(tool_style: str, version: str, include_builtin: bool):
+def test_model_with_tools(tool_style: str, version: str, include_builtin: bool) -> None:
     model = FakeToolCallingModel(tool_style=tool_style)
 
     @dec_tool
@@ -288,7 +277,7 @@ def test_model_with_tools(tool_style: str, version: str, include_builtin: bool):
         """Tool 2 docstring."""
         return f"Tool 2: {some_val}"
 
-    tools = [tool1, tool2]
+    tools: list[BaseTool | dict] = [tool1, tool2]
     if include_builtin:
         tools.append(
             {
@@ -306,45 +295,12 @@ def test_model_with_tools(tool_style: str, version: str, include_builtin: bool):
             }
         )
     # check valid agent constructor
-    agent = create_react_agent(
-        model.bind_tools(tools),
-        tools,
-        version=version,
-    )
-    result = agent.nodes["tools"].invoke(
-        {
-            "messages": [
-                AIMessage(
-                    "hi?",
-                    tool_calls=[
-                        {
-                            "name": "tool1",
-                            "args": {"some_val": 2},
-                            "id": "some 1",
-                        },
-                        {
-                            "name": "tool2",
-                            "args": {"some_val": 2},
-                            "id": "some 2",
-                        },
-                    ],
-                )
-            ]
-        }
-    )
-    tool_messages: ToolMessage = result["messages"][-2:]
-    for tool_message in tool_messages:
-        assert tool_message.type == "tool"
-        assert tool_message.content in {"Tool 1: 2", "Tool 2: 2"}
-        assert tool_message.tool_call_id in {"some 1", "some 2"}
-
-    # test mismatching tool lengths
     with pytest.raises(ValueError):
-        create_react_agent(model.bind_tools([tool1]), [tool1, tool2])
-
-    # test missing bound tools
-    with pytest.raises(ValueError):
-        create_react_agent(model.bind_tools([tool1]), [tool2])
+        create_react_agent(
+            model.bind_tools(tools),
+            tools,
+            version=version,
+        )
 
 
 def test__validate_messages():
@@ -478,27 +434,46 @@ def test_react_agent_with_structured_response(version: str) -> None:
     class WeatherResponse(BaseModel):
         temperature: float = Field(description="The temperature in fahrenheit")
 
-    tool_calls = [[{"args": {}, "id": "1", "name": "get_weather"}], []]
+    tool_calls = [
+        [{"args": {}, "id": "1", "name": "get_weather"}],
+        [{"name": "WeatherResponse", "id": "2", "args": {"temperature": 75}}],
+    ]
 
     def get_weather():
         """Get the weather"""
         return "The weather is sunny and 75°F."
 
     expected_structured_response = WeatherResponse(temperature=75)
-    model = FakeToolCallingModel(
+    model = FakeToolCallingModel[WeatherResponse](
         tool_calls=tool_calls, structured_response=expected_structured_response
     )
-    for response_format in (WeatherResponse, ("Meow", WeatherResponse)):
-        agent = create_react_agent(
-            model,
-            [get_weather],
-            response_format=response_format,
-            version=version,
-        )
-        response = agent.invoke({"messages": [HumanMessage("What's the weather?")]})
-        assert response["structured_response"] == expected_structured_response
-        assert len(response["messages"]) == 4
-        assert response["messages"][-2].content == "The weather is sunny and 75°F."
+    agent = create_react_agent(
+        model,
+        [get_weather],
+        response_format=WeatherResponse,
+        version=version,
+    )
+    response = agent.invoke({"messages": [HumanMessage("What's the weather?")]})
+    assert response["structured_response"] == expected_structured_response
+    assert len(response["messages"]) == 5
+
+    # Check message types in message history
+    msg_types = [m.type for m in response["messages"]]
+    assert msg_types == [
+        "human",  # "What's the weather?"
+        "ai",  # "What's the weather?"
+        "tool",  # "The weather is sunny and 75°F."
+        "ai",  # structured response
+        "tool",  # artificial tool message
+    ]
+
+    assert [m.content for m in response["messages"]] == [
+        "What's the weather?",
+        "What's the weather?",
+        "The weather is sunny and 75°F.",
+        "What's the weather?-What's the weather?-The weather is sunny and 75°F.",
+        "Returning structured response: {'temperature': 75.0}",
+    ]
 
 
 class CustomState(AgentState):
@@ -514,7 +489,7 @@ class CustomStatePydantic(AgentStatePydantic):
 def test_react_agent_update_state(
     sync_checkpointer: BaseCheckpointSaver,
     version: Literal["v1", "v2"],
-    state_schema: StateSchemaType,
+    state_schema: StateT,
 ) -> None:
     @dec_tool
     def get_user_name(tool_call_id: Annotated[str, InjectedToolCallId]):
@@ -654,124 +629,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
 
@@ -785,7 +642,7 @@ class AgentStateExtraKeyPydantic(AgentStatePydantic):
     "state_schema", [AgentStateExtraKey, AgentStateExtraKeyPydantic]
 )
 def test_create_react_agent_inject_vars(
-    version: Literal["v1", "v2"], state_schema: StateSchemaType
+    version: Literal["v1", "v2"], state_schema: StateT
 ) -> None:
     """Test that the agent can inject state and store into tool functions."""
     store = InMemoryStore()
@@ -837,135 +694,6 @@ 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:
     @dec_tool(return_direct=True)
@@ -1182,60 +910,6 @@ 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:
     """Test React agent streaming when used as a subgraph node sync version"""
@@ -1491,69 +1165,6 @@ def test_tool_node_node_interrupt(
     )
 
 
-@pytest.mark.parametrize("tool_style", ["openai", "anthropic"])
-def test_should_bind_tools(tool_style: str) -> None:
-    @dec_tool
-    def some_tool(some_val: int) -> str:
-        """Tool docstring."""
-        return "meow"
-
-    @dec_tool
-    def some_other_tool(some_val: int) -> str:
-        """Tool docstring."""
-        return "meow"
-
-    model = FakeToolCallingModel(tool_style=tool_style)
-    # should bind when a regular model
-    assert _should_bind_tools(model, [])
-    assert _should_bind_tools(model, [some_tool])
-
-    # should bind when a seq
-    seq = model | RunnableLambda(lambda message: message)
-    assert _should_bind_tools(seq, [])
-    assert _should_bind_tools(seq, [some_tool])
-
-    # should not bind when a model with tools
-    assert not _should_bind_tools(model.bind_tools([some_tool]), [some_tool])
-    # should not bind when a seq with tools
-    seq_with_tools = model.bind_tools([some_tool]) | RunnableLambda(
-        lambda message: message
-    )
-    assert not _should_bind_tools(seq_with_tools, [some_tool])
-
-    # should raise on invalid inputs
-    with pytest.raises(ValueError):
-        _should_bind_tools(model.bind_tools([some_tool]), [])
-    with pytest.raises(ValueError):
-        _should_bind_tools(model.bind_tools([some_tool]), [some_other_tool])
-    with pytest.raises(ValueError):
-        _should_bind_tools(model.bind_tools([some_tool]), [some_tool, some_other_tool])
-
-
-def test_get_model() -> None:
-    model = FakeToolCallingModel(tool_calls=[])
-    assert _get_model(model) == model
-
-    @dec_tool
-    def some_tool(some_val: int) -> str:
-        """Tool docstring."""
-        return "meow"
-
-    model_with_tools = model.bind_tools([some_tool])
-    assert _get_model(model_with_tools) == model
-
-    seq = model | RunnableLambda(lambda message: message)
-    assert _get_model(seq) == model
-
-    seq_with_tools = model.bind_tools([some_tool]) | RunnableLambda(
-        lambda message: message
-    )
-    assert _get_model(seq_with_tools) == model
-
-    with pytest.raises(TypeError):
-        _get_model(RunnableLambda(lambda message: message))
-
-
 @pytest.mark.parametrize("version", REACT_TOOL_CALL_VERSIONS)
 def test_dynamic_model_basic(version: str) -> None:
     """Test basic dynamic model functionality."""
@@ -1730,9 +1341,17 @@ def test_dynamic_model_with_structured_response(version: str) -> None:
         confidence: float
 
     def dynamic_model(state, runtime: Runtime):
-        expected_response = TestResponse(message="dynamic response", confidence=0.9)
         return FakeToolCallingModel(
-            tool_calls=[], structured_response=expected_response
+            tool_calls=[
+                [
+                    ToolCall(
+                        name="TestResponse",
+                        args={"message": "dynamic response", "confidence": 0.9},
+                        id="1",
+                        type="tool_call",
+                    )
+                ]
+            ],
         )
 
     agent = create_react_agent(
@@ -1997,16 +1616,16 @@ def test_post_model_hook_with_structured_output() -> None:
     class WeatherResponse(BaseModel):
         temperature: float = Field(description="The temperature in fahrenheit")
 
-    tool_calls = [[{"args": {}, "id": "1", "name": "get_weather"}]]
+    tool_calls: list[list[ToolCall]] = [
+        [{"args": {}, "id": "1", "name": "get_weather"}],
+        [{"args": {"temperature": 75}, "id": "2", "name": "WeatherResponse"}],
+    ]
 
     def get_weather():
         """Get the weather"""
         return "The weather is sunny and 75°F."
 
     expected_structured_response = WeatherResponse(temperature=75)
-    model = FakeToolCallingModel(
-        tool_calls=tool_calls, structured_response=expected_structured_response
-    )
 
     class State(AgentState):
         flag: bool
@@ -2015,6 +1634,7 @@ def test_post_model_hook_with_structured_output() -> None:
     def post_model_hook(state: State) -> Union[dict[str, bool], Command]:
         return {"flag": True}
 
+    model = FakeToolCallingModel(tool_calls=tool_calls)
     agent = create_react_agent(
         model,
         [get_weather],
@@ -2024,7 +1644,6 @@ def test_post_model_hook_with_structured_output() -> None:
     )
 
     assert "post_model_hook" in agent.nodes
-    assert "generate_structured_response" in agent.nodes
 
     response = agent.invoke(
         {"messages": [HumanMessage("What's the weather?")], "flag": False}
@@ -2032,10 +1651,19 @@ def test_post_model_hook_with_structured_output() -> None:
     assert response["flag"] is True
     assert response["structured_response"] == expected_structured_response
 
+    # Reset the state of the model
+    model = FakeToolCallingModel(tool_calls=tool_calls)
+    agent = create_react_agent(
+        model,
+        [get_weather],
+        response_format=WeatherResponse,
+        post_model_hook=post_model_hook,
+        state_schema=State,
+    )
+
     events = list(
         agent.stream({"messages": [HumanMessage("What's the weather?")], "flag": False})
     )
-    assert "generate_structured_response" in events[-1]
     assert events == [
         {
             "agent": {
@@ -2044,7 +1672,7 @@ def test_post_model_hook_with_structured_output() -> None:
                         content="What's the weather?",
                         additional_kwargs={},
                         response_metadata={},
-                        id="2",
+                        id="0",
                         tool_calls=[
                             {
                                 "name": "get_weather",
@@ -2065,7 +1693,7 @@ def test_post_model_hook_with_structured_output() -> None:
                         content="The weather is sunny and 75°F.",
                         name="get_weather",
                         tool_call_id="1",
-                    ),
+                    )
                 ]
             }
         },
@@ -2076,25 +1704,26 @@ def test_post_model_hook_with_structured_output() -> None:
                         content="What's the weather?-What's the weather?-The weather is sunny and 75°F.",
                         additional_kwargs={},
                         response_metadata={},
-                        id="3",
+                        id="1",
                         tool_calls=[
                             {
-                                "name": "get_weather",
-                                "args": {},
-                                "id": "1",
+                                "name": "WeatherResponse",
+                                "args": {"temperature": 75},
+                                "id": "2",
                                 "type": "tool_call",
                             }
                         ],
-                    )
-                ]
+                    ),
+                    _AnyIdToolMessage(
+                        content="Returning structured response: {'temperature': 75.0}",
+                        name="WeatherResponse",
+                        tool_call_id="2",
+                    ),
+                ],
+                "structured_response": WeatherResponse(temperature=75.0),
             }
         },
         {"post_model_hook": {"flag": True}},
-        {
-            "generate_structured_response": {
-                "structured_response": WeatherResponse(temperature=75.0)
-            }
-        },
     ]
 
 
@@ -2102,7 +1731,7 @@ def test_post_model_hook_with_structured_output() -> None:
     "state_schema", [AgentStateExtraKey, AgentStateExtraKeyPydantic]
 )
 def test_create_react_agent_inject_vars_with_post_model_hook(
-    state_schema: StateSchemaType,
+    state_schema: StateT,
 ) -> None:
     store = InMemoryStore()
     namespace = ("test",)
@@ -2157,3 +1786,40 @@ def test_create_react_agent_inject_vars_with_post_model_hook(
         AIMessage("hi-hi-6", id="1"),
     ]
     assert result["foo"] == 2
+
+
+@pytest.mark.parametrize("version", REACT_TOOL_CALL_VERSIONS)
+def test_response_format_using_tool_choice(version: Literal["v1", "v2"]) -> None:
+    """Test response format using tool choice."""
+
+    class WeatherResponse(BaseModel):
+        temperature: float = Field(description="The temperature in fahrenheit")
+
+    tool_calls: list[list[ToolCall]] = [
+        [{"args": {}, "id": "1", "name": "get_weather"}],
+        [{"args": {"temperature": "75"}, "id": "2", "name": "WeatherResponse"}],
+    ]
+
+    def get_weather() -> str:
+        """Get the weather"""
+        return "The weather is sunny and 75°F."
+
+    expected_structured_response = WeatherResponse(temperature=75)
+    model = FakeToolCallingModel(tool_calls=tool_calls)
+    agent = create_react_agent(
+        model,
+        [get_weather],
+        response_format=WeatherResponse,
+        version=version,
+    )
+    response = agent.invoke(
+        {
+            "messages": [
+                {
+                    "role": "user",
+                    "content": "What's the weather?",
+                }
+            ]
+        }
+    )
+    assert response.get("structured_response") == expected_structured_response

libs/prebuilt/tests/test_react_agent_graph.py

@@ -34,19 +34,25 @@ class ResponseFormat(BaseModel):
 @pytest.mark.parametrize("tools", [[], [tool]])
 @pytest.mark.parametrize("pre_model_hook", [None, pre_model_hook])
 @pytest.mark.parametrize("post_model_hook", [None, post_model_hook])
-@pytest.mark.parametrize("response_format", [None, ResponseFormat])
 def test_react_agent_graph_structure(
     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,
     )
-    assert agent.get_graph().draw_mermaid(with_styles=False) == snapshot
+    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}, "
+        ) from e

libs/prebuilt/tests/test_response_format.py

@@ -0,0 +1,793 @@
+"""Test suite for create_react_agent with structured output response_format permutations."""
+
+from dataclasses import dataclass
+from typing import Union
+
+import pytest
+from langchain_core.messages import HumanMessage
+from pydantic import BaseModel, Field
+from typing_extensions import TypedDict
+
+from langgraph.prebuilt import create_react_agent
+from langgraph.prebuilt.responses import (
+    MultipleStructuredOutputsError,
+    NativeOutput,
+    StructuredOutputParsingError,
+    ToolOutput,
+)
+from tests.model import FakeToolCallingModel
+
+try:
+    from langchain_openai import ChatOpenAI
+except ImportError:
+    skip_openai_integration_tests = True
+else:
+    skip_openai_integration_tests = False
+
+
+# Test data models
+class WeatherBaseModel(BaseModel):
+    """Weather response."""
+
+    temperature: float = Field(description="The temperature in fahrenheit")
+    condition: str = Field(description="Weather condition")
+
+
+@dataclass
+class WeatherDataclass:
+    """Weather response."""
+
+    temperature: float
+    condition: str
+
+
+class WeatherTypedDict(TypedDict):
+    """Weather response."""
+
+    temperature: float
+    condition: str
+
+
+weather_json_schema = {
+    "type": "object",
+    "properties": {
+        "temperature": {"type": "number", "description": "Temperature in fahrenheit"},
+        "condition": {"type": "string", "description": "Weather condition"},
+    },
+    "title": "weather_schema",
+    "required": ["temperature", "condition"],
+}
+
+
+class LocationResponse(BaseModel):
+    city: str = Field(description="The city name")
+    country: str = Field(description="The country name")
+
+
+class LocationTypedDict(TypedDict):
+    city: str
+    country: str
+
+
+location_json_schema = {
+    "type": "object",
+    "properties": {
+        "city": {"type": "string", "description": "The city name"},
+        "country": {"type": "string", "description": "The country name"},
+    },
+    "title": "location_schema",
+    "required": ["city", "country"],
+}
+
+
+def get_weather() -> str:
+    """Get the weather."""
+
+    return "The weather is sunny and 75°F."
+
+
+def get_location() -> str:
+    """Get the current location."""
+
+    return "You are in New York, USA."
+
+
+# Standardized test data
+WEATHER_DATA = {"temperature": 75.0, "condition": "sunny"}
+LOCATION_DATA = {"city": "New York", "country": "USA"}
+
+# Standardized expected responses
+EXPECTED_WEATHER_PYDANTIC = WeatherBaseModel(**WEATHER_DATA)
+EXPECTED_WEATHER_DATACLASS = WeatherDataclass(**WEATHER_DATA)
+EXPECTED_WEATHER_DICT: WeatherTypedDict = {"temperature": 75.0, "condition": "sunny"}
+EXPECTED_LOCATION = LocationResponse(**LOCATION_DATA)
+EXPECTED_LOCATION_DICT: LocationTypedDict = {"city": "New York", "country": "USA"}
+
+
+class TestResponseFormatAsModel:
+    def test_pydantic_model(self) -> None:
+        """Test response_format as Pydantic model."""
+        tool_calls = [
+            [{"args": {}, "id": "1", "name": "get_weather"}],
+            [
+                {
+                    "name": "WeatherBaseModel",
+                    "id": "2",
+                    "args": WEATHER_DATA,
+                }
+            ],
+        ]
+
+        model = FakeToolCallingModel(tool_calls=tool_calls)
+
+        agent = create_react_agent(
+            model, [get_weather], response_format=WeatherBaseModel
+        )
+        response = agent.invoke({"messages": [HumanMessage("What's the weather?")]})
+
+        assert response["structured_response"] == EXPECTED_WEATHER_PYDANTIC
+        assert len(response["messages"]) == 5
+
+    def test_dataclass(self) -> None:
+        """Test response_format as dataclass."""
+        tool_calls = [
+            [{"args": {}, "id": "1", "name": "get_weather"}],
+            [
+                {
+                    "name": "WeatherDataclass",
+                    "id": "2",
+                    "args": WEATHER_DATA,
+                }
+            ],
+        ]
+
+        model = FakeToolCallingModel(tool_calls=tool_calls)
+
+        agent = create_react_agent(
+            model, [get_weather], response_format=WeatherDataclass
+        )
+        response = agent.invoke({"messages": [HumanMessage("What's the weather?")]})
+
+        assert response["structured_response"] == EXPECTED_WEATHER_DATACLASS
+        assert len(response["messages"]) == 5
+
+    def test_typed_dict(self) -> None:
+        """Test response_format as TypedDict."""
+        tool_calls = [
+            [{"args": {}, "id": "1", "name": "get_weather"}],
+            [
+                {
+                    "name": "WeatherTypedDict",
+                    "id": "2",
+                    "args": WEATHER_DATA,
+                }
+            ],
+        ]
+
+        model = FakeToolCallingModel(tool_calls=tool_calls)
+
+        agent = create_react_agent(
+            model, [get_weather], response_format=WeatherTypedDict
+        )
+        response = agent.invoke({"messages": [HumanMessage("What's the weather?")]})
+
+        assert response["structured_response"] == EXPECTED_WEATHER_DICT
+        assert len(response["messages"]) == 5
+
+    def test_json_schema(self) -> None:
+        """Test response_format as JSON schema."""
+        tool_calls = [
+            [{"args": {}, "id": "1", "name": "get_weather"}],
+            [
+                {
+                    "name": "weather_schema",
+                    "id": "2",
+                    "args": WEATHER_DATA,
+                }
+            ],
+        ]
+
+        model = FakeToolCallingModel(tool_calls=tool_calls)
+
+        agent = create_react_agent(
+            model, [get_weather], response_format=weather_json_schema
+        )
+        response = agent.invoke({"messages": [HumanMessage("What's the weather?")]})
+
+        assert response["structured_response"] == EXPECTED_WEATHER_DICT
+        assert len(response["messages"]) == 5
+
+
+class TestResponseFormatAsToolOutput:
+    def test_pydantic_model(self) -> None:
+        """Test response_format as ToolOutput with Pydantic model."""
+        tool_calls = [
+            [{"args": {}, "id": "1", "name": "get_weather"}],
+            [
+                {
+                    "name": "WeatherBaseModel",
+                    "id": "2",
+                    "args": WEATHER_DATA,
+                }
+            ],
+        ]
+
+        model = FakeToolCallingModel(tool_calls=tool_calls)
+
+        agent = create_react_agent(
+            model, [get_weather], response_format=ToolOutput(WeatherBaseModel)
+        )
+        response = agent.invoke({"messages": [HumanMessage("What's the weather?")]})
+
+        assert response["structured_response"] == EXPECTED_WEATHER_PYDANTIC
+        assert len(response["messages"]) == 5
+
+    def test_dataclass(self) -> None:
+        """Test response_format as ToolOutput with dataclass."""
+        tool_calls = [
+            [{"args": {}, "id": "1", "name": "get_weather"}],
+            [
+                {
+                    "name": "WeatherDataclass",
+                    "id": "2",
+                    "args": WEATHER_DATA,
+                }
+            ],
+        ]
+
+        model = FakeToolCallingModel(tool_calls=tool_calls)
+
+        agent = create_react_agent(
+            model, [get_weather], response_format=ToolOutput(WeatherDataclass)
+        )
+        response = agent.invoke({"messages": [HumanMessage("What's the weather?")]})
+
+        assert response["structured_response"] == EXPECTED_WEATHER_DATACLASS
+        assert len(response["messages"]) == 5
+
+    def test_typed_dict(self) -> None:
+        """Test response_format as ToolOutput with TypedDict."""
+        tool_calls = [
+            [{"args": {}, "id": "1", "name": "get_weather"}],
+            [
+                {
+                    "name": "WeatherTypedDict",
+                    "id": "2",
+                    "args": WEATHER_DATA,
+                }
+            ],
+        ]
+
+        model = FakeToolCallingModel(tool_calls=tool_calls)
+
+        agent = create_react_agent(
+            model, [get_weather], response_format=ToolOutput(WeatherTypedDict)
+        )
+        response = agent.invoke({"messages": [HumanMessage("What's the weather?")]})
+
+        assert response["structured_response"] == EXPECTED_WEATHER_DICT
+        assert len(response["messages"]) == 5
+
+    def test_json_schema(self) -> None:
+        """Test response_format as ToolOutput with JSON schema."""
+        tool_calls = [
+            [{"args": {}, "id": "1", "name": "get_weather"}],
+            [
+                {
+                    "name": "weather_schema",
+                    "id": "2",
+                    "args": WEATHER_DATA,
+                }
+            ],
+        ]
+
+        model = FakeToolCallingModel(tool_calls=tool_calls)
+
+        agent = create_react_agent(
+            model, [get_weather], response_format=ToolOutput(weather_json_schema)
+        )
+        response = agent.invoke({"messages": [HumanMessage("What's the weather?")]})
+
+        assert response["structured_response"] == EXPECTED_WEATHER_DICT
+        assert len(response["messages"]) == 5
+
+    def test_union_of_json_schemas(self) -> None:
+        """Test response_format as ToolOutput with union of JSON schemas."""
+        tool_calls = [
+            [{"args": {}, "id": "1", "name": "get_weather"}],
+            [
+                {
+                    "name": "weather_schema",
+                    "id": "2",
+                    "args": WEATHER_DATA,
+                }
+            ],
+        ]
+
+        model = FakeToolCallingModel(tool_calls=tool_calls)
+
+        agent = create_react_agent(
+            model,
+            [get_weather, get_location],
+            response_format=ToolOutput(
+                {"oneOf": [weather_json_schema, location_json_schema]}
+            ),
+        )
+        response = agent.invoke({"messages": [HumanMessage("What's the weather?")]})
+
+        assert response["structured_response"] == EXPECTED_WEATHER_DICT
+        assert len(response["messages"]) == 5
+
+        # Test with LocationResponse
+        tool_calls_location = [
+            [{"args": {}, "id": "1", "name": "get_location"}],
+            [
+                {
+                    "name": "location_schema",
+                    "id": "2",
+                    "args": LOCATION_DATA,
+                }
+            ],
+        ]
+
+        model_location = FakeToolCallingModel(tool_calls=tool_calls_location)
+
+        agent_location = create_react_agent(
+            model_location,
+            [get_weather, get_location],
+            response_format=ToolOutput(
+                {"oneOf": [weather_json_schema, location_json_schema]}
+            ),
+        )
+        response_location = agent_location.invoke(
+            {"messages": [HumanMessage("Where am I?")]}
+        )
+
+        assert response_location["structured_response"] == EXPECTED_LOCATION_DICT
+        assert len(response_location["messages"]) == 5
+
+    def test_union_of_types(self) -> None:
+        """Test response_format as ToolOutput with Union of various types."""
+        # Test with WeatherBaseModel
+        tool_calls = [
+            [{"args": {}, "id": "1", "name": "get_weather"}],
+            [
+                {
+                    "name": "WeatherBaseModel",
+                    "id": "2",
+                    "args": WEATHER_DATA,
+                }
+            ],
+        ]
+
+        model = FakeToolCallingModel[Union[WeatherBaseModel, LocationResponse]](
+            tool_calls=tool_calls
+        )
+
+        agent = create_react_agent(
+            model,
+            [get_weather, get_location],
+            response_format=ToolOutput(Union[WeatherBaseModel, LocationResponse]),
+        )
+        response = agent.invoke({"messages": [HumanMessage("What's the weather?")]})
+
+        assert response["structured_response"] == EXPECTED_WEATHER_PYDANTIC
+        assert len(response["messages"]) == 5
+
+        # Test with LocationResponse
+        tool_calls_location = [
+            [{"args": {}, "id": "1", "name": "get_location"}],
+            [
+                {
+                    "name": "LocationResponse",
+                    "id": "2",
+                    "args": LOCATION_DATA,
+                }
+            ],
+        ]
+
+        model_location = FakeToolCallingModel(tool_calls=tool_calls_location)
+
+        agent_location = create_react_agent(
+            model_location,
+            [get_weather, get_location],
+            response_format=ToolOutput(Union[WeatherBaseModel, LocationResponse]),
+        )
+        response_location = agent_location.invoke(
+            {"messages": [HumanMessage("Where am I?")]}
+        )
+
+        assert response_location["structured_response"] == EXPECTED_LOCATION
+        assert len(response_location["messages"]) == 5
+
+    def test_multiple_structured_outputs_error_without_retry(self) -> None:
+        """Test that MultipleStructuredOutputsError is raised when model returns multiple structured tool calls without retry."""
+        tool_calls = [
+            [
+                {
+                    "name": "WeatherBaseModel",
+                    "id": "1",
+                    "args": WEATHER_DATA,
+                },
+                {
+                    "name": "LocationResponse",
+                    "id": "2",
+                    "args": LOCATION_DATA,
+                },
+            ],
+        ]
+
+        model = FakeToolCallingModel(tool_calls=tool_calls)
+
+        agent = create_react_agent(
+            model,
+            [],
+            response_format=ToolOutput(
+                Union[WeatherBaseModel, LocationResponse],
+                handle_errors=False,
+            ),
+        )
+
+        with pytest.raises(
+            MultipleStructuredOutputsError,
+            match=".*WeatherBaseModel.*LocationResponse.*",
+        ):
+            agent.invoke({"messages": [HumanMessage("Give me weather and location")]})
+
+    def test_multiple_structured_outputs_with_retry(self) -> None:
+        """Test that retry handles multiple structured output tool calls."""
+        tool_calls = [
+            [
+                {
+                    "name": "WeatherBaseModel",
+                    "id": "1",
+                    "args": WEATHER_DATA,
+                },
+                {
+                    "name": "LocationResponse",
+                    "id": "2",
+                    "args": LOCATION_DATA,
+                },
+            ],
+            [
+                {
+                    "name": "WeatherBaseModel",
+                    "id": "3",
+                    "args": WEATHER_DATA,
+                },
+            ],
+        ]
+
+        model = FakeToolCallingModel(tool_calls=tool_calls)
+
+        agent = create_react_agent(
+            model,
+            [],
+            response_format=ToolOutput(
+                Union[WeatherBaseModel, LocationResponse],
+                handle_errors=True,
+            ),
+        )
+
+        response = agent.invoke({"messages": [HumanMessage("Give me weather")]})
+
+        # HumanMessage, AIMessage, ToolMessage, ToolMessage, AI, ToolMessage
+        assert len(response["messages"]) == 6
+        assert response["structured_response"] == EXPECTED_WEATHER_PYDANTIC
+
+    def test_structured_output_parsing_error_without_retry(self) -> None:
+        """Test that StructuredOutputParsingError is raised when tool args fail to parse without retry."""
+        tool_calls = [
+            [
+                {
+                    "name": "WeatherBaseModel",
+                    "id": "1",
+                    "args": {"invalid": "data"},
+                },
+            ],
+        ]
+
+        model = FakeToolCallingModel(tool_calls=tool_calls)
+
+        agent = create_react_agent(
+            model,
+            [],
+            response_format=ToolOutput(
+                WeatherBaseModel,
+                handle_errors=False,
+            ),
+        )
+
+        with pytest.raises(
+            StructuredOutputParsingError,
+            match=".*WeatherBaseModel.*",
+        ):
+            agent.invoke({"messages": [HumanMessage("What's the weather?")]})
+
+    def test_structured_output_parsing_error_with_retry(self) -> None:
+        """Test that retry handles parsing errors for structured output."""
+        tool_calls = [
+            [
+                {
+                    "name": "WeatherBaseModel",
+                    "id": "1",
+                    "args": {"invalid": "data"},
+                },
+            ],
+            [
+                {
+                    "name": "WeatherBaseModel",
+                    "id": "2",
+                    "args": WEATHER_DATA,
+                },
+            ],
+        ]
+
+        model = FakeToolCallingModel(tool_calls=tool_calls)
+
+        agent = create_react_agent(
+            model,
+            [],
+            response_format=ToolOutput(
+                WeatherBaseModel,
+                handle_errors=(StructuredOutputParsingError,),
+            ),
+        )
+
+        response = agent.invoke({"messages": [HumanMessage("What's the weather?")]})
+
+        # HumanMessage, AIMessage, ToolMessage, AIMessage, ToolMessage
+        assert len(response["messages"]) == 5
+        assert response["structured_response"] == EXPECTED_WEATHER_PYDANTIC
+
+    def test_retry_with_custom_function(self) -> None:
+        """Test retry with custom message generation."""
+        tool_calls = [
+            [
+                {
+                    "name": "WeatherBaseModel",
+                    "id": "1",
+                    "args": WEATHER_DATA,
+                },
+                {
+                    "name": "LocationResponse",
+                    "id": "2",
+                    "args": LOCATION_DATA,
+                },
+            ],
+            [
+                {
+                    "name": "WeatherBaseModel",
+                    "id": "3",
+                    "args": WEATHER_DATA,
+                },
+            ],
+        ]
+
+        model = FakeToolCallingModel(tool_calls=tool_calls)
+
+        def custom_message(exception: Exception) -> str:
+            if isinstance(exception, MultipleStructuredOutputsError):
+                return "Custom error: Multiple outputs not allowed"
+            return "Custom error"
+
+        agent = create_react_agent(
+            model,
+            [],
+            response_format=ToolOutput(
+                Union[WeatherBaseModel, LocationResponse],
+                handle_errors=custom_message,
+            ),
+        )
+
+        response = agent.invoke({"messages": [HumanMessage("Give me weather")]})
+
+        # HumanMessage, AIMessage, ToolMessage, ToolMessage, AI, ToolMessage
+        assert len(response["messages"]) == 6
+        assert (
+            response["messages"][2].content
+            == "Custom error: Multiple outputs not allowed"
+        )
+        assert (
+            response["messages"][3].content
+            == "Custom error: Multiple outputs not allowed"
+        )
+        assert response["structured_response"] == EXPECTED_WEATHER_PYDANTIC
+
+    def test_retry_with_custom_string_message(self) -> None:
+        """Test retry with custom static string message."""
+        tool_calls = [
+            [
+                {
+                    "name": "WeatherBaseModel",
+                    "id": "1",
+                    "args": {"invalid": "data"},
+                },
+            ],
+            [
+                {
+                    "name": "WeatherBaseModel",
+                    "id": "2",
+                    "args": WEATHER_DATA,
+                },
+            ],
+        ]
+
+        model = FakeToolCallingModel(tool_calls=tool_calls)
+
+        agent = create_react_agent(
+            model,
+            [],
+            response_format=ToolOutput(
+                WeatherBaseModel,
+                handle_errors="Please provide valid weather data with temperature and condition.",
+            ),
+        )
+
+        response = agent.invoke({"messages": [HumanMessage("What's the weather?")]})
+
+        assert len(response["messages"]) == 5
+        assert (
+            response["messages"][2].content
+            == "Please provide valid weather data with temperature and condition."
+        )
+        assert response["structured_response"] == EXPECTED_WEATHER_PYDANTIC
+
+
+class TestResponseFormatAsNativeOutput:
+    def test_pydantic_model(self) -> None:
+        """Test response_format as NativeOutput with Pydantic model."""
+        tool_calls = [
+            [{"args": {}, "id": "1", "name": "get_weather"}],
+        ]
+
+        model = FakeToolCallingModel[WeatherBaseModel](
+            tool_calls=tool_calls, structured_response=EXPECTED_WEATHER_PYDANTIC
+        )
+
+        agent = create_react_agent(
+            model, [get_weather], response_format=NativeOutput(WeatherBaseModel)
+        )
+        response = agent.invoke({"messages": [HumanMessage("What's the weather?")]})
+
+        assert response["structured_response"] == EXPECTED_WEATHER_PYDANTIC
+        assert len(response["messages"]) == 4
+
+    def test_dataclass(self) -> None:
+        """Test response_format as NativeOutput with dataclass."""
+        tool_calls = [
+            [{"args": {}, "id": "1", "name": "get_weather"}],
+        ]
+
+        model = FakeToolCallingModel[WeatherDataclass](
+            tool_calls=tool_calls, structured_response=EXPECTED_WEATHER_DATACLASS
+        )
+
+        agent = create_react_agent(
+            model, [get_weather], response_format=NativeOutput(WeatherDataclass)
+        )
+        response = agent.invoke(
+            {"messages": [HumanMessage("What's the weather?")]},
+        )
+
+        assert response["structured_response"] == EXPECTED_WEATHER_DATACLASS
+        assert len(response["messages"]) == 4
+
+    def test_typed_dict(self) -> None:
+        """Test response_format as NativeOutput with TypedDict."""
+        tool_calls = [
+            [{"args": {}, "id": "1", "name": "get_weather"}],
+        ]
+
+        model = FakeToolCallingModel[WeatherTypedDict](
+            tool_calls=tool_calls, structured_response=EXPECTED_WEATHER_DICT
+        )
+
+        agent = create_react_agent(
+            model, [get_weather], response_format=NativeOutput(WeatherTypedDict)
+        )
+        response = agent.invoke({"messages": [HumanMessage("What's the weather?")]})
+
+        assert response["structured_response"] == EXPECTED_WEATHER_DICT
+        assert len(response["messages"]) == 4
+
+    def test_json_schema(self) -> None:
+        """Test response_format as NativeOutput with JSON schema."""
+        tool_calls = [
+            [{"args": {}, "id": "1", "name": "get_weather"}],
+        ]
+
+        model = FakeToolCallingModel[dict](
+            tool_calls=tool_calls, structured_response=EXPECTED_WEATHER_DICT
+        )
+
+        agent = create_react_agent(
+            model, [get_weather], response_format=NativeOutput(weather_json_schema)
+        )
+        response = agent.invoke({"messages": [HumanMessage("What's the weather?")]})
+
+        assert response["structured_response"] == EXPECTED_WEATHER_DICT
+        assert len(response["messages"]) == 4
+
+
+def test_union_of_types() -> None:
+    """Test response_format as NativeOutput with Union (if supported)."""
+    tool_calls = [
+        [{"args": {}, "id": "1", "name": "get_weather"}],
+        [
+            {
+                "name": "WeatherBaseModel",
+                "id": "2",
+                "args": WEATHER_DATA,
+            }
+        ],
+    ]
+
+    model = FakeToolCallingModel[Union[WeatherBaseModel, LocationResponse]](
+        tool_calls=tool_calls, structured_response=EXPECTED_WEATHER_PYDANTIC
+    )
+
+    agent = create_react_agent(
+        model,
+        [get_weather, get_location],
+        response_format=ToolOutput(Union[WeatherBaseModel, LocationResponse]),
+    )
+    response = agent.invoke({"messages": [HumanMessage("What's the weather?")]})
+
+    assert response["structured_response"] == EXPECTED_WEATHER_PYDANTIC
+    assert len(response["messages"]) == 5
+
+
+@pytest.mark.skipif(
+    skip_openai_integration_tests, reason="OpenAI integration tests are disabled."
+)
+def test_inference_to_native_output() -> None:
+    """Test that native output is inferred when a model supports it."""
+    model = ChatOpenAI(model="gpt-5")
+    agent = create_react_agent(
+        model,
+        prompt="You are a helpful weather assistant. Please call the get_weather tool, then use the WeatherReport tool to generate the final response.",
+        tools=[get_weather],
+        response_format=WeatherBaseModel,
+    )
+    response = agent.invoke({"messages": [HumanMessage("What's the weather?")]})
+
+    assert isinstance(response["structured_response"], WeatherBaseModel)
+    assert response["structured_response"].temperature == 75.0
+    assert response["structured_response"].condition.lower() == "sunny"
+    assert len(response["messages"]) == 4
+
+    assert [m.type for m in response["messages"]] == [
+        "human",  # "What's the weather?"
+        "ai",  # "What's the weather?"
+        "tool",  # "The weather is sunny and 75°F."
+        "ai",  # structured response
+    ]
+
+
+@pytest.mark.skipif(
+    skip_openai_integration_tests, reason="OpenAI integration tests are disabled."
+)
+def test_inference_to_tool_output() -> None:
+    """Test that tool output is inferred when a model supports it."""
+    model = ChatOpenAI(model="gpt-4")
+    agent = create_react_agent(
+        model,
+        prompt="You are a helpful weather assistant. Please call the get_weather tool, then use the WeatherReport tool to generate the final response.",
+        tools=[get_weather],
+        response_format=ToolOutput(WeatherBaseModel),
+    )
+    response = agent.invoke({"messages": [HumanMessage("What's the weather?")]})
+
+    assert isinstance(response["structured_response"], WeatherBaseModel)
+    assert response["structured_response"].temperature == 75.0
+    assert response["structured_response"].condition.lower() == "sunny"
+    assert len(response["messages"]) == 5
+
+    assert [m.type for m in response["messages"]] == [
+        "human",  # "What's the weather?"
+        "ai",  # "What's the weather?"
+        "tool",  # "The weather is sunny and 75°F."
+        "ai",  # structured response
+        "tool",  # artificial tool message
+    ]

libs/prebuilt/tests/test_responses.py

@@ -0,0 +1,154 @@
+"""Unit tests for langgraph.prebuilt.responses module."""
+
+from typing import Union
+
+import pytest
+from pydantic import BaseModel
+
+from langgraph.prebuilt.responses import (
+    OutputToolBinding,
+    ToolOutput,
+    _SchemaSpec,
+)
+
+
+class _TestModel(BaseModel):
+    """A test model for structured output."""
+
+    name: str
+    age: int
+    email: str = "default@example.com"
+
+
+class CustomModel(BaseModel):
+    """Custom model with a custom docstring."""
+
+    value: float
+    description: str
+
+
+class EmptyDocModel(BaseModel):
+    # No custom docstring, should have no description in tool
+    data: str
+
+
+class TestUsingToolStrategy:
+    """Test UsingToolStrategy dataclass."""
+
+    def test_basic_creation(self):
+        """Test basic UsingToolStrategy creation."""
+        strategy = ToolOutput(schema=_TestModel)
+        assert strategy.schema == _TestModel
+        assert strategy.tool_message_content is None
+        assert len(strategy.schema_specs) == 1
+
+    def test_multiple_schemas(self):
+        """Test UsingToolStrategy with multiple schemas."""
+        strategy = ToolOutput(schema=Union[_TestModel, CustomModel])
+        assert len(strategy.schema_specs) == 2
+        assert strategy.schema_specs[0].schema == _TestModel
+        assert strategy.schema_specs[1].schema == CustomModel
+
+    def test_schema_with_tool_message_content(self):
+        """Test UsingToolStrategy with tool message content."""
+        strategy = ToolOutput(schema=_TestModel, tool_message_content="custom message")
+        assert strategy.schema == _TestModel
+        assert strategy.tool_message_content == "custom message"
+        assert len(strategy.schema_specs) == 1
+
+
+class TestOutputToolBinding:
+    """Test OutputToolBinding dataclass and its methods."""
+
+    def test_from_schema_spec_basic(self):
+        """Test basic OutputToolBinding creation from SchemaSpec."""
+        schema_spec = _SchemaSpec(schema=_TestModel)
+        tool_binding = OutputToolBinding.from_schema_spec(schema_spec)
+
+        assert tool_binding.schema == _TestModel
+        assert tool_binding.schema_kind == "pydantic"
+        assert tool_binding.tool is not None
+        assert tool_binding.tool.name == "_TestModel"
+
+    def test_from_schema_spec_with_custom_name(self):
+        """Test OutputToolBinding creation with custom name."""
+        schema_spec = _SchemaSpec(schema=_TestModel, name="custom_tool_name")
+        tool_binding = OutputToolBinding.from_schema_spec(schema_spec)
+        assert tool_binding.tool.name == "custom_tool_name"
+
+    def test_from_schema_spec_with_custom_description(self):
+        """Test OutputToolBinding creation with custom description."""
+        schema_spec = _SchemaSpec(
+            schema=_TestModel, description="Custom tool description"
+        )
+        tool_binding = OutputToolBinding.from_schema_spec(schema_spec)
+
+        assert tool_binding.tool.description == "Custom tool description"
+
+    def test_from_schema_spec_with_model_docstring(self):
+        """Test OutputToolBinding creation using model docstring as description."""
+        schema_spec = _SchemaSpec(schema=CustomModel)
+        tool_binding = OutputToolBinding.from_schema_spec(schema_spec)
+
+        assert tool_binding.tool.description == "Custom model with a custom docstring."
+
+    @pytest.mark.skip(
+        reason="Need to fix bug in langchain-core for inheritance of doc-strings."
+    )
+    def test_from_schema_spec_empty_docstring(self):
+        """Test OutputToolBinding creation with model that has default docstring."""
+
+        # Create a model with the same docstring as BaseModel
+        class DefaultDocModel(BaseModel):
+            # This should have the same docstring as BaseModel
+            pass
+
+        schema_spec = _SchemaSpec(schema=DefaultDocModel)
+        tool_binding = OutputToolBinding.from_schema_spec(schema_spec)
+
+        # Should use empty description when model has default BaseModel docstring
+        assert tool_binding.tool.description == ""
+
+    def test_parse_payload_pydantic_success(self):
+        """Test successful parsing for Pydantic model."""
+        schema_spec = _SchemaSpec(schema=_TestModel)
+        tool_binding = OutputToolBinding.from_schema_spec(schema_spec)
+
+        tool_args = {"name": "John", "age": 30}
+        result = tool_binding.parse(tool_args)
+
+        assert isinstance(result, _TestModel)
+        assert result.name == "John"
+        assert result.age == 30
+        assert result.email == "default@example.com"  # default value
+
+    def test_parse_payload_pydantic_validation_error(self):
+        """Test parsing failure for invalid Pydantic data."""
+        schema_spec = _SchemaSpec(schema=_TestModel)
+        tool_binding = OutputToolBinding.from_schema_spec(schema_spec)
+
+        # Missing required field 'name'
+        tool_args = {"age": 30}
+
+        with pytest.raises(ValueError, match="Failed to parse data to _TestModel"):
+            tool_binding.parse(tool_args)
+
+
+class TestEdgeCases:
+    """Test edge cases and error conditions."""
+
+    def test_empty_schemas_list(self) -> None:
+        """Test UsingToolStrategy with empty schemas list."""
+        strategy = ToolOutput(EmptyDocModel)
+        assert len(strategy.schema_specs) == 1
+
+    @pytest.mark.skip(
+        reason="Need to fix bug in langchain-core for inheritance of doc-strings."
+    )
+    def test_base_model_doc_constant(self) -> None:
+        """Test that BASE_MODEL_DOC constant is set correctly."""
+        binding = OutputToolBinding.from_schema_spec(_SchemaSpec(EmptyDocModel))
+        assert binding.tool.name == "EmptyDocModel"
+        assert (
+            binding.tool.description[:5] == ""
+        )  # Should be empty for default docstring

libs/prebuilt/tests/test_responses_spec.py

@@ -0,0 +1,160 @@
+from __future__ import annotations
+
+from typing import Any, Dict, List, Optional, Union
+from unittest.mock import MagicMock
+
+import httpx
+import pytest
+from langchain_core.messages import HumanMessage
+from langchain_core.tools import tool
+from pydantic import BaseModel, create_model
+
+from langgraph.prebuilt import create_react_agent
+from langgraph.prebuilt.responses import ToolOutput
+from tests.utils import BaseSchema, load_spec
+
+try:
+    from langchain_openai import ChatOpenAI
+except ImportError:
+    skip_openai_integration_tests = True
+else:
+    skip_openai_integration_tests = False
+
+AGENT_PROMPT = "You are an HR assistant."
+
+
+class ToolCalls(BaseSchema):
+    get_employee_role: int
+    get_employee_department: int
+
+
+class AssertionByInvocation(BaseSchema):
+    prompt: str
+    tools_with_expected_calls: ToolCalls
+    expected_last_message: str
+    expected_structured_response: Optional[Dict[str, Any]]
+    llm_request_count: int
+
+
+class TestCase(BaseSchema):
+    name: str
+    response_format: Union[Dict[str, Any], List[Dict[str, Any]]]
+    assertions_by_invocation: List[AssertionByInvocation]
+
+
+class Employee(BaseModel):
+    name: str
+    role: str
+    department: str
+
+
+EMPLOYEES: list[Employee] = [
+    Employee(name="Sabine", role="Developer", department="IT"),
+    Employee(name="Henrik", role="Product Manager", department="IT"),
+    Employee(name="Jessica", role="HR", department="People"),
+]
+
+TEST_CASES = load_spec("responses", as_model=TestCase)
+
+
+def _make_tool(fn, *, name: str, description: str):
+    mock = MagicMock(side_effect=lambda *, name: fn(name=name))
+    InputModel = create_model(f"{name}_input", name=(str, ...))
+
+    @tool(name, description=description, args_schema=InputModel)
+    def _wrapped(name: str):
+        return mock(name=name)
+
+    return {"tool": _wrapped, "mock": mock}
+
+
+@pytest.mark.skipif(
+    skip_openai_integration_tests, reason="OpenAI integration tests are disabled."
+)
+@pytest.mark.parametrize("case", TEST_CASES, ids=[c.name for c in TEST_CASES])
+def test_responses_integration_matrix(case: TestCase) -> None:
+    if case.name == "asking for information that does not fit into the response format":
+        pytest.xfail(
+            "currently failing due to undefined behavior when model cannot conform to any of the structured response formats."
+        )
+
+    def get_employee_role(*, name: str) -> Optional[str]:
+        for e in EMPLOYEES:
+            if e.name == name:
+                return e.role
+        return None
+
+    def get_employee_department(*, name: str) -> Optional[str]:
+        for e in EMPLOYEES:
+            if e.name == name:
+                return e.department
+        return None
+
+    role_tool = _make_tool(
+        get_employee_role,
+        name="get_employee_role",
+        description="Get the employee role by name",
+    )
+    dept_tool = _make_tool(
+        get_employee_department,
+        name="get_employee_department",
+        description="Get the employee department by name",
+    )
+
+    response_format_spec = case.response_format
+    if isinstance(response_format_spec, dict):
+        response_format_spec = [response_format_spec]
+    # Unwrap nested schema objects
+    response_format_spec = [item.get("schema", item) for item in response_format_spec]
+    if len(response_format_spec) == 1:
+        tool_output = ToolOutput(response_format_spec[0])
+    else:
+        tool_output = ToolOutput({"oneOf": response_format_spec})
+
+    llm_request_count = 0
+
+    for assertion in case.assertions_by_invocation:
+
+        def on_request(request: httpx.Request) -> None:
+            nonlocal llm_request_count
+            llm_request_count += 1
+
+        http_client = httpx.Client(
+            event_hooks={"request": [on_request]},
+        )
+
+        model = ChatOpenAI(
+            model="gpt-4o",
+            temperature=0,
+            http_client=http_client,
+        )
+
+        agent = create_react_agent(
+            model,
+            tools=[role_tool["tool"], dept_tool["tool"]],
+            prompt=AGENT_PROMPT,
+            response_format=tool_output,
+        )
+
+        result = agent.invoke({"messages": [HumanMessage(assertion.prompt)]})
+
+        # Count tool calls
+        assert (
+            role_tool["mock"].call_count
+            == assertion.tools_with_expected_calls.get_employee_role
+        )
+        assert (
+            dept_tool["mock"].call_count
+            == assertion.tools_with_expected_calls.get_employee_department
+        )
+
+        # Count LLM calls
+        assert llm_request_count == assertion.llm_request_count
+
+        # Check last message content
+        last_message = result["messages"][-1]
+        assert last_message.content == assertion.expected_last_message
+
+        # Check structured response
+        structured_response_json = result["structured_response"]
+        assert structured_response_json == assertion.expected_structured_response

libs/prebuilt/tests/test_return_direct_spec.py

@@ -0,0 +1,117 @@
+from __future__ import annotations
+
+from typing import Any, Dict, Optional
+from unittest.mock import MagicMock
+
+import pytest
+from langchain_core.messages import HumanMessage
+from langchain_core.tools import tool
+
+from langgraph.prebuilt import create_react_agent
+from langgraph.prebuilt.responses import ToolOutput
+from tests.utils import BaseSchema, load_spec
+
+try:
+    from langchain_openai import ChatOpenAI
+except ImportError:
+    skip_openai_integration_tests = True
+else:
+    skip_openai_integration_tests = False
+
+AGENT_PROMPT = """
+You are a strict polling bot.
+
+- Only use the "poll_job" tool until it returns { status: "succeeded" }.
+- If status is "pending", call the tool again. Do not produce a final answer.
+- When it is "succeeded", return exactly: "Attempts: <number>" with no extra text.
+"""
+
+
+class TestCase(BaseSchema):
+    name: str
+    return_direct: bool
+    response_format: Optional[Dict[str, Any]]
+    expected_tool_calls: int
+    expected_last_message: str
+    expected_structured_response: Optional[Dict[str, Any]]
+
+
+TEST_CASES = load_spec("return_direct", as_model=TestCase)
+
+
+def _make_tool(return_direct: bool):
+    attempts = 0
+
+    def _side_effect():
+        nonlocal attempts
+        attempts += 1
+        return {
+            "status": "succeeded" if attempts >= 10 else "pending",
+            "attempts": attempts,
+        }
+
+    mock = MagicMock(side_effect=_side_effect)
+
+    @tool(
+        "pollJob",
+        description=(
+            "Check the status of a long-running job. "
+            "Returns { status: 'pending' | 'succeeded', attempts: number }."
+        ),
+        return_direct=return_direct,
+    )
+    def _wrapped():
+        return mock()
+
+    return {"tool": _wrapped, "mock": mock}
+
+
+@pytest.mark.skipif(
+    skip_openai_integration_tests, reason="OpenAI integration tests are disabled."
+)
+@pytest.mark.parametrize("case", TEST_CASES, ids=[c.name for c in TEST_CASES])
+def test_return_direct_integration_matrix(case: TestCase) -> None:
+    poll_tool = _make_tool(case.return_direct)
+
+    model = ChatOpenAI(
+        model="gpt-4o",
+        temperature=0,
+    )
+
+    if case.response_format:
+        agent = create_react_agent(
+            model,
+            tools=[poll_tool["tool"]],
+            prompt=AGENT_PROMPT,
+            response_format=ToolOutput(case.response_format),
+        )
+    else:
+        agent = create_react_agent(
+            model,
+            tools=[poll_tool["tool"]],
+            prompt=AGENT_PROMPT,
+        )
+
+    result = agent.invoke(
+        {
+            "messages": [
+                HumanMessage(
+                    "Poll the job until it's done and tell me how many attempts it took."
+                )
+            ]
+        }
+    )
+
+    # Count tool calls
+    assert poll_tool["mock"].call_count == case.expected_tool_calls
+
+    # Check last message content
+    last_message = result["messages"][-1]
+    assert last_message.content == case.expected_last_message
+
+    # Check structured response
+    if case.expected_structured_response is not None:
+        structured_response_json = result["structured_response"]
+        assert structured_response_json == case.expected_structured_response
+    else:
+        assert "structured_response" not in result

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 ValidationError as ValidationErrorV1
+from pydantic import BaseModel
+from pydantic.v1 import BaseModel as BaseModelV1
+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,
+    ToolInvocationError,
+    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,17 +220,67 @@ async def test_tool_node_tool_call_input():
     ]
 
 
-async def test_tool_node_error_handling():
-    def handle_all(e: Union[ValueError, ToolException, ValidationError]):
+def test_tool_node_error_handling_default_invocation() -> None:
+    tn = ToolNode([tool1])
+    result = tn.invoke(
+        {
+            "messages": [
+                AIMessage(
+                    "hi?",
+                    tool_calls=[
+                        {
+                            "name": "tool1",
+                            "args": {"invalid": 0, "args": "foo"},
+                            "id": "some id",
+                        },
+                    ],
+                )
+            ]
+        }
+    )
+
+    assert all(m.type == "tool" for m in result["messages"])
+    assert all(m.status == "error" for m in result["messages"])
+    assert (
+        "Error invoking tool 'tool1' with kwargs {'invalid': 0, 'args': 'foo'} with error:\n"
+        in result["messages"][0].content
+    )
+
+
+def test_tool_node_error_handling_default_exception() -> None:
+    tn = ToolNode([tool1])
+    with pytest.raises(ValueError):
+        tn.invoke(
+            {
+                "messages": [
+                    AIMessage(
+                        "hi?",
+                        tool_calls=[
+                            {
+                                "name": "tool1",
+                                "args": {"some_val": 0, "some_other_val": "foo"},
+                                "id": "some id",
+                            },
+                        ],
+                    )
+                ]
+            }
+        )
+
+
+async def test_tool_node_error_handling() -> None:
+    def handle_all(e: Union[ValueError, ToolException, ToolInvocationError]):
         return TOOL_CALL_ERROR_TEMPLATE.format(error=repr(e))
 
     # test catching all exceptions, via:
     # - handle_tool_errors = True
+    # - passing a single exception
     # - passing a tuple of all exceptions
     # - passing a callable with all exceptions in the signature
     for handle_tool_errors in (
         True,
-        (ValueError, ToolException, ValidationError),
+        Exception,
+        (ValueError, ToolException, ToolInvocationError),
         handle_all,
     ):
         result_error = await ToolNode(
@@ -257,7 +332,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"
 
@@ -388,7 +463,7 @@ async def test_tool_node_handle_tool_errors_false():
     assert str(exc_info.value) == "Test error"
 
     # test validation errors get raised if handle_tool_errors is False
-    with pytest.raises((ValidationError, ValidationErrorV1)):
+    with pytest.raises((ToolInvocationError)):
         ToolNode([tool1], handle_tool_errors=False).invoke(
             {
                 "messages": [
@@ -1156,3 +1231,304 @@ 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], handle_tool_errors=True)
+    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",
+                        ),
+                    ],
+                },
+            },
+        ),
+    ]

libs/prebuilt/tests/utils.py

@@ -0,0 +1,22 @@
+import json
+from pathlib import Path
+from typing import Type
+
+from pydantic import BaseModel, ConfigDict
+from pydantic.alias_generators import to_camel
+
+
+class BaseSchema(BaseModel):
+    model_config = ConfigDict(
+        alias_generator=to_camel,
+        populate_by_name=True,
+        from_attributes=True,
+    )
+
+
+def load_spec(spec_name: str, as_model: Type[BaseModel]) -> list[BaseModel]:
+    with (Path(__file__).parent / "specifications" / f"{spec_name}.json").open(
+        "r", encoding="utf-8"
+    ) as f:
+        data = json.load(f)
+        return [as_model(**item) for item in data]