langgraph/libs/prebuilt/tests/test_tool_node.py

import contextlib
import dataclasses
import json
import sys
import warnings
from functools import partial
from typing import (
    Annotated,
    Any,
    NoReturn,
    TypeVar,
)
from unittest.mock import Mock

import pytest
from langchain_core.messages import (
    AIMessage,
    AnyMessage,
    HumanMessage,
    RemoveMessage,
    ToolCall,
    ToolMessage,
)
from langchain_core.runnables.config import RunnableConfig
from langchain_core.tools import BaseTool, InjectedToolArg, ToolException
from langchain_core.tools import tool as dec_tool
from langchain_core.tools.base import InjectedToolCallId
from langgraph.config import get_stream_writer
from langgraph.errors import GraphBubbleUp, GraphInterrupt
from langgraph.graph import START, MessagesState, StateGraph
from langgraph.graph.message import REMOVE_ALL_MESSAGES, add_messages
from langgraph.runtime import ExecutionInfo, ServerInfo
from langgraph.store.base import BaseStore
from langgraph.store.memory import InMemoryStore
from langgraph.types import Command, Send
from pydantic import BaseModel
from pydantic.v1 import BaseModel as BaseModelV1
from typing_extensions import TypedDict

from langgraph.prebuilt import (
    InjectedState,
    InjectedStore,
    ToolNode,
)
from langgraph.prebuilt.tool_node import (
    TOOL_CALL_ERROR_TEMPLATE,
    ToolCallRequest,
    ToolInvocationError,
    ToolRuntime,
    tools_condition,
)

from .messages import _AnyIdHumanMessage, _AnyIdToolMessage
from .model import FakeToolCallingModel

pytestmark = pytest.mark.anyio


def _create_mock_runtime(store: BaseStore | None = None) -> Mock:
    """Create a mock Runtime object for testing ToolNode outside of graph context.

    This helper is needed because ToolNode._func expects a Runtime parameter
    which is injected by RunnableCallable from config["configurable"]["__pregel_runtime"].
    When testing ToolNode directly (outside a graph), we need to provide this manually.
    """

    mock_runtime = Mock()
    mock_runtime.store = store
    mock_runtime.context = None
    mock_runtime.stream_writer = lambda *args, **kwargs: None
    mock_runtime.execution_info = ExecutionInfo(
        checkpoint_id="test-cp", checkpoint_ns="", task_id="test-task"
    )
    mock_runtime.server_info = None
    return mock_runtime


def _create_config_with_runtime(store: BaseStore | None = None) -> RunnableConfig:
    """Create a RunnableConfig with mock Runtime for testing ToolNode.

    Returns:
        RunnableConfig with __pregel_runtime in configurable dict.
    """
    return {"configurable": {"__pregel_runtime": _create_mock_runtime(store)}}


def tool1(some_val: int, some_other_val: str) -> str:
    """Tool 1 docstring."""
    if some_val == 0:
        msg = "Test error"
        raise ValueError(msg)
    return f"{some_val} - {some_other_val}"


async def tool2(some_val: int, some_other_val: str) -> str:
    """Tool 2 docstring."""
    if some_val == 0:
        msg = "Test error"
        raise ToolException(msg)
    return f"tool2: {some_val} - {some_other_val}"


async def tool3(some_val: int, some_other_val: str) -> str:
    """Tool 3 docstring."""
    return [
        {"key_1": some_val, "key_2": "foo"},
        {"key_1": some_other_val, "key_2": "baz"},
    ]


async def tool4(some_val: int, some_other_val: str) -> str:
    """Tool 4 docstring."""
    return [
        {"type": "image_url", "image_url": {"url": "abdc"}},
    ]


@dec_tool
def tool5(some_val: int) -> NoReturn:
    """Tool 5 docstring."""
    msg = "Test error"
    raise ToolException(msg)


tool5.handle_tool_error = "foo"


async def test_tool_node() -> None:
    """Test tool node."""
    result = ToolNode([tool1]).invoke(
        {
            "messages": [
                AIMessage(
                    "hi?",
                    tool_calls=[
                        {
                            "name": "tool1",
                            "args": {"some_val": 1, "some_other_val": "foo"},
                            "id": "some 0",
                        }
                    ],
                )
            ]
        },
        config=_create_config_with_runtime(),
    )

    tool_message: ToolMessage = result["messages"][-1]
    assert tool_message.type == "tool"
    assert tool_message.content == "1 - foo"
    assert tool_message.tool_call_id == "some 0"

    result2 = await ToolNode([tool2]).ainvoke(
        {
            "messages": [
                AIMessage(
                    "hi?",
                    tool_calls=[
                        {
                            "name": "tool2",
                            "args": {"some_val": 2, "some_other_val": "bar"},
                            "id": "some 1",
                        }
                    ],
                )
            ]
        },
        config=_create_config_with_runtime(),
    )

    tool_message: ToolMessage = result2["messages"][-1]
    assert tool_message.type == "tool"
    assert tool_message.content == "tool2: 2 - bar"

    # list of dicts tool content
    result3 = await ToolNode([tool3]).ainvoke(
        {
            "messages": [
                AIMessage(
                    "hi?",
                    tool_calls=[
                        {
                            "name": "tool3",
                            "args": {"some_val": 2, "some_other_val": "bar"},
                            "id": "some 2",
                        }
                    ],
                )
            ]
        },
        config=_create_config_with_runtime(),
    )
    tool_message: ToolMessage = result3["messages"][-1]
    assert tool_message.type == "tool"
    assert (
        tool_message.content
        == '[{"key_1": 2, "key_2": "foo"}, {"key_1": "bar", "key_2": "baz"}]'
    )
    assert tool_message.tool_call_id == "some 2"

    # list of content blocks tool content
    result4 = await ToolNode([tool4]).ainvoke(
        {
            "messages": [
                AIMessage(
                    "hi?",
                    tool_calls=[
                        {
                            "name": "tool4",
                            "args": {"some_val": 2, "some_other_val": "bar"},
                            "id": "some 3",
                        }
                    ],
                )
            ]
        },
        config=_create_config_with_runtime(),
    )
    tool_message: ToolMessage = result4["messages"][-1]
    assert tool_message.type == "tool"
    assert tool_message.content == [{"type": "image_url", "image_url": {"url": "abdc"}}]
    assert tool_message.tool_call_id == "some 3"


async def test_tool_node_tool_call_input() -> None:
    # Single tool call
    tool_call_1 = {
        "name": "tool1",
        "args": {"some_val": 1, "some_other_val": "foo"},
        "id": "some 0",
        "type": "tool_call",
    }
    result = ToolNode([tool1]).invoke(
        [tool_call_1], config=_create_config_with_runtime()
    )
    assert result["messages"] == [
        ToolMessage(content="1 - foo", tool_call_id="some 0", name="tool1"),
    ]

    # Multiple tool calls
    tool_call_2 = {
        "name": "tool1",
        "args": {"some_val": 2, "some_other_val": "bar"},
        "id": "some 1",
        "type": "tool_call",
    }
    result = ToolNode([tool1]).invoke(
        [tool_call_1, tool_call_2], config=_create_config_with_runtime()
    )
    assert result["messages"] == [
        ToolMessage(content="1 - foo", tool_call_id="some 0", name="tool1"),
        ToolMessage(content="2 - bar", tool_call_id="some 1", name="tool1"),
    ]

    # Test with unknown tool
    tool_call_3 = tool_call_1.copy()
    tool_call_3["name"] = "tool2"
    result = ToolNode([tool1]).invoke(
        [tool_call_1, tool_call_3], config=_create_config_with_runtime()
    )
    assert result["messages"] == [
        ToolMessage(content="1 - foo", tool_call_id="some 0", name="tool1"),
        ToolMessage(
            content="Error: tool2 is not a valid tool, try one of [tool1].",
            name="tool2",
            tool_call_id="some 0",
            status="error",
        ),
    ]


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",
                        },
                    ],
                )
            ]
        },
        config=_create_config_with_runtime(),
    )

    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",
                            },
                        ],
                    )
                ]
            },
            config=_create_config_with_runtime(),
        )


async def test_tool_node_error_handling() -> None:
    def handle_all(e: ValueError | ToolException | ToolInvocationError):
        return TOOL_CALL_ERROR_TEMPLATE.format(error=repr(e))

    # test catching all exceptions, via:
    # - handle_tool_errors = True
    # - passing a tuple of all exceptions
    # - passing a callable with all exceptions in the signature
    for handle_tool_errors in (
        True,
        (ValueError, ToolException, ToolInvocationError),
        handle_all,
    ):
        result_error = await ToolNode(
            [tool1, tool2, tool3], handle_tool_errors=handle_tool_errors
        ).ainvoke(
            {
                "messages": [
                    AIMessage(
                        "hi?",
                        tool_calls=[
                            {
                                "name": "tool1",
                                "args": {"some_val": 0, "some_other_val": "foo"},
                                "id": "some id",
                            },
                            {
                                "name": "tool2",
                                "args": {"some_val": 0, "some_other_val": "bar"},
                                "id": "some other id",
                            },
                            {
                                "name": "tool3",
                                "args": {"some_val": 0},
                                "id": "another id",
                            },
                        ],
                    )
                ]
            },
            config=_create_config_with_runtime(),
        )

        assert all(m.type == "tool" for m in result_error["messages"])
        assert all(m.status == "error" for m in result_error["messages"])
        assert (
            result_error["messages"][0].content
            == f"Error: {ValueError('Test error')!r}\n Please fix your mistakes."
        )
        assert (
            result_error["messages"][1].content
            == f"Error: {ToolException('Test error')!r}\n Please fix your mistakes."
        )
        # Check that the validation error contains the field name
        assert "some_other_val" in result_error["messages"][2].content

        assert result_error["messages"][0].tool_call_id == "some id"
        assert result_error["messages"][1].tool_call_id == "some other id"
        assert result_error["messages"][2].tool_call_id == "another id"


async def test_tool_node_error_handling_callable() -> None:
    def handle_value_error(e: ValueError) -> str:
        return "Value error"

    def handle_tool_exception(e: ToolException) -> str:
        return "Tool exception"

    for handle_tool_errors in ("Value error", handle_value_error):
        result_error = await ToolNode(
            [tool1], handle_tool_errors=handle_tool_errors
        ).ainvoke(
            {
                "messages": [
                    AIMessage(
                        "hi?",
                        tool_calls=[
                            {
                                "name": "tool1",
                                "args": {"some_val": 0, "some_other_val": "foo"},
                                "id": "some id",
                            },
                        ],
                    )
                ]
            },
            config=_create_config_with_runtime(),
        )
        tool_message: ToolMessage = result_error["messages"][-1]
        assert tool_message.type == "tool"
        assert tool_message.status == "error"
        assert tool_message.content == "Value error"

    # test raising for an unhandled exception, via:
    # - passing a tuple of all exceptions
    # - passing a callable with all exceptions in the signature
    for handle_tool_errors in ((ValueError,), handle_value_error):
        with pytest.raises(ToolException) as exc_info:
            await ToolNode(
                [tool1, tool2], handle_tool_errors=handle_tool_errors
            ).ainvoke(
                {
                    "messages": [
                        AIMessage(
                            "hi?",
                            tool_calls=[
                                {
                                    "name": "tool1",
                                    "args": {"some_val": 0, "some_other_val": "foo"},
                                    "id": "some id",
                                },
                                {
                                    "name": "tool2",
                                    "args": {"some_val": 0, "some_other_val": "bar"},
                                    "id": "some other id",
                                },
                            ],
                        )
                    ]
                },
                config=_create_config_with_runtime(),
            )
        assert str(exc_info.value) == "Test error"

    for handle_tool_errors in ((ToolException,), handle_tool_exception):
        with pytest.raises(ValueError) as exc_info:
            await ToolNode(
                [tool1, tool2], handle_tool_errors=handle_tool_errors
            ).ainvoke(
                {
                    "messages": [
                        AIMessage(
                            "hi?",
                            tool_calls=[
                                {
                                    "name": "tool1",
                                    "args": {"some_val": 0, "some_other_val": "foo"},
                                    "id": "some id",
                                },
                                {
                                    "name": "tool2",
                                    "args": {"some_val": 0, "some_other_val": "bar"},
                                    "id": "some other id",
                                },
                            ],
                        )
                    ]
                },
                config=_create_config_with_runtime(),
            )
        assert str(exc_info.value) == "Test error"


async def test_tool_node_handle_tool_errors_false() -> None:
    with pytest.raises(ValueError) as exc_info:
        ToolNode([tool1], handle_tool_errors=False).invoke(
            {
                "messages": [
                    AIMessage(
                        "hi?",
                        tool_calls=[
                            {
                                "name": "tool1",
                                "args": {"some_val": 0, "some_other_val": "foo"},
                                "id": "some id",
                            }
                        ],
                    )
                ]
            },
            config=_create_config_with_runtime(),
        )

    assert str(exc_info.value) == "Test error"

    with pytest.raises(ToolException):
        await ToolNode([tool2], handle_tool_errors=False).ainvoke(
            {
                "messages": [
                    AIMessage(
                        "hi?",
                        tool_calls=[
                            {
                                "name": "tool2",
                                "args": {"some_val": 0, "some_other_val": "bar"},
                                "id": "some id",
                            }
                        ],
                    )
                ]
            },
            config=_create_config_with_runtime(),
        )

    assert str(exc_info.value) == "Test error"

    # test validation errors get raised if handle_tool_errors is False
    with pytest.raises(ToolInvocationError):
        ToolNode([tool1], handle_tool_errors=False).invoke(
            {
                "messages": [
                    AIMessage(
                        "hi?",
                        tool_calls=[
                            {
                                "name": "tool1",
                                "args": {"some_val": 0},
                                "id": "some id",
                            }
                        ],
                    )
                ]
            },
            config=_create_config_with_runtime(),
        )


def test_tool_node_individual_tool_error_handling() -> None:
    # test error handling on individual tools (and that it overrides overall error handling!)
    result_individual_tool_error_handler = ToolNode(
        [tool5], handle_tool_errors="bar"
    ).invoke(
        {
            "messages": [
                AIMessage(
                    "hi?",
                    tool_calls=[
                        {
                            "name": "tool5",
                            "args": {"some_val": 0},
                            "id": "some 0",
                        }
                    ],
                )
            ]
        },
        config=_create_config_with_runtime(),
    )

    tool_message: ToolMessage = result_individual_tool_error_handler["messages"][-1]
    assert tool_message.type == "tool"
    assert tool_message.status == "error"
    assert tool_message.content == "foo"
    assert tool_message.tool_call_id == "some 0"


def test_tool_node_incorrect_tool_name() -> None:
    result_incorrect_name = ToolNode([tool1, tool2]).invoke(
        {
            "messages": [
                AIMessage(
                    "hi?",
                    tool_calls=[
                        {
                            "name": "tool3",
                            "args": {"some_val": 1, "some_other_val": "foo"},
                            "id": "some 0",
                        }
                    ],
                )
            ]
        },
        config=_create_config_with_runtime(),
    )

    tool_message: ToolMessage = result_incorrect_name["messages"][-1]
    assert tool_message.type == "tool"
    assert tool_message.status == "error"
    assert (
        tool_message.content
        == "Error: tool3 is not a valid tool, try one of [tool1, tool2]."
    )
    assert tool_message.tool_call_id == "some 0"


def test_tool_node_node_interrupt() -> None:
    def tool_interrupt(some_val: int) -> None:
        """Tool docstring."""
        msg = "foo"
        raise GraphBubbleUp(msg)

    def handle(e: GraphInterrupt) -> str:
        return "handled"

    for handle_tool_errors in (True, (GraphBubbleUp,), "handled", handle, False):
        node = ToolNode([tool_interrupt], handle_tool_errors=handle_tool_errors)
        with pytest.raises(GraphBubbleUp) as exc_info:
            node.invoke(
                {
                    "messages": [
                        AIMessage(
                            "hi?",
                            tool_calls=[
                                {
                                    "name": "tool_interrupt",
                                    "args": {"some_val": 0},
                                    "id": "some 0",
                                }
                            ],
                        )
                    ]
                },
                config=_create_config_with_runtime(),
            )
            assert exc_info.value == "foo"


@pytest.mark.parametrize("input_type", ["dict", "tool_calls"])
async def test_tool_node_command(input_type: str) -> None:

    @dec_tool
    def transfer_to_bob(tool_call_id: Annotated[str, InjectedToolCallId]):
        """Transfer to Bob"""
        return Command(
            update={
                "messages": [
                    ToolMessage(content="Transferred to Bob", tool_call_id=tool_call_id)
                ]
            },
            goto="bob",
            graph=Command.PARENT,
        )

    @dec_tool
    async def async_transfer_to_bob(tool_call_id: Annotated[str, InjectedToolCallId]):
        """Transfer to Bob"""
        return Command(
            update={
                "messages": [
                    ToolMessage(content="Transferred to Bob", tool_call_id=tool_call_id)
                ]
            },
            goto="bob",
            graph=Command.PARENT,
        )

    class CustomToolSchema(BaseModel):
        tool_call_id: Annotated[str, InjectedToolCallId]

    class MyCustomTool(BaseTool):
        def _run(*args: Any, **kwargs: Any):
            return Command(
                update={
                    "messages": [
                        ToolMessage(
                            content="Transferred to Bob",
                            tool_call_id=kwargs["tool_call_id"],
                        )
                    ]
                },
                goto="bob",
                graph=Command.PARENT,
            )

        async def _arun(*args: Any, **kwargs: Any):
            return Command(
                update={
                    "messages": [
                        ToolMessage(
                            content="Transferred to Bob",
                            tool_call_id=kwargs["tool_call_id"],
                        )
                    ]
                },
                goto="bob",
                graph=Command.PARENT,
            )

    custom_tool = MyCustomTool(
        name="custom_transfer_to_bob",
        description="Transfer to bob",
        args_schema=CustomToolSchema,
    )
    async_custom_tool = MyCustomTool(
        name="async_custom_transfer_to_bob",
        description="Transfer to bob",
        args_schema=CustomToolSchema,
    )

    # test mixing regular tools and tools returning commands
    def add(a: int, b: int) -> int:
        """Add two numbers"""
        return a + b

    tool_calls = [
        {"args": {"a": 1, "b": 2}, "id": "1", "name": "add", "type": "tool_call"},
        {"args": {}, "id": "2", "name": "transfer_to_bob", "type": "tool_call"},
    ]
    if input_type == "dict":
        input_ = {"messages": [AIMessage("", tool_calls=tool_calls)]}
    elif input_type == "tool_calls":
        input_ = tool_calls
    result = ToolNode([add, transfer_to_bob]).invoke(
        input_, config=_create_config_with_runtime()
    )

    assert result == [
        {
            "messages": [
                ToolMessage(
                    content="3",
                    tool_call_id="1",
                    name="add",
                )
            ]
        },
        Command(
            update={
                "messages": [
                    ToolMessage(
                        content="Transferred to Bob",
                        tool_call_id="2",
                        name="transfer_to_bob",
                    )
                ]
            },
            goto="bob",
            graph=Command.PARENT,
        ),
    ]

    # test tools returning commands

    # test sync tools
    for tool in [transfer_to_bob, custom_tool]:
        result = ToolNode([tool]).invoke(
            {
                "messages": [
                    AIMessage(
                        "", tool_calls=[{"args": {}, "id": "1", "name": tool.name}]
                    )
                ]
            },
            config=_create_config_with_runtime(),
        )
        assert result == [
            Command(
                update={
                    "messages": [
                        ToolMessage(
                            content="Transferred to Bob",
                            tool_call_id="1",
                            name=tool.name,
                        )
                    ]
                },
                goto="bob",
                graph=Command.PARENT,
            )
        ]

    # test async tools
    for tool in [async_transfer_to_bob, async_custom_tool]:
        result = await ToolNode([tool]).ainvoke(
            {
                "messages": [
                    AIMessage(
                        "", tool_calls=[{"args": {}, "id": "1", "name": tool.name}]
                    )
                ]
            },
            config=_create_config_with_runtime(),
        )
        assert result == [
            Command(
                update={
                    "messages": [
                        ToolMessage(
                            content="Transferred to Bob",
                            tool_call_id="1",
                            name=tool.name,
                        )
                    ]
                },
                goto="bob",
                graph=Command.PARENT,
            )
        ]

    # test multiple commands
    result = ToolNode([transfer_to_bob, custom_tool]).invoke(
        {
            "messages": [
                AIMessage(
                    "",
                    tool_calls=[
                        {"args": {}, "id": "1", "name": "transfer_to_bob"},
                        {"args": {}, "id": "2", "name": "custom_transfer_to_bob"},
                    ],
                )
            ]
        },
        config=_create_config_with_runtime(),
    )
    assert result == [
        Command(
            update={
                "messages": [
                    ToolMessage(
                        content="Transferred to Bob",
                        tool_call_id="1",
                        name="transfer_to_bob",
                    )
                ]
            },
            goto="bob",
            graph=Command.PARENT,
        ),
        Command(
            update={
                "messages": [
                    ToolMessage(
                        content="Transferred to Bob",
                        tool_call_id="2",
                        name="custom_transfer_to_bob",
                    )
                ]
            },
            goto="bob",
            graph=Command.PARENT,
        ),
    ]

    # test validation (mismatch between input type and command.update type)
    with pytest.raises(ValueError):

        @dec_tool
        def list_update_tool(tool_call_id: Annotated[str, InjectedToolCallId]):
            """My tool"""
            return Command(
                update=[ToolMessage(content="foo", tool_call_id=tool_call_id)]
            )

        ToolNode([list_update_tool]).invoke(
            {
                "messages": [
                    AIMessage(
                        "",
                        tool_calls=[
                            {"args": {}, "id": "1", "name": "list_update_tool"}
                        ],
                    )
                ]
            },
            config=_create_config_with_runtime(),
        )

    # test validation (missing tool message in the update for current graph)
    with pytest.raises(ValueError):

        @dec_tool
        def no_update_tool():
            """My tool"""
            return Command(update={"messages": []})

        ToolNode([no_update_tool]).invoke(
            {
                "messages": [
                    AIMessage(
                        "",
                        tool_calls=[{"args": {}, "id": "1", "name": "no_update_tool"}],
                    )
                ]
            },
            config=_create_config_with_runtime(),
        )

    # test validation (tool message with a wrong tool call ID)
    with pytest.raises(ValueError):

        @dec_tool
        def mismatching_tool_call_id_tool():
            """My tool"""
            return Command(
                update={"messages": [ToolMessage(content="foo", tool_call_id="2")]}
            )

        ToolNode([mismatching_tool_call_id_tool]).invoke(
            {
                "messages": [
                    AIMessage(
                        "",
                        tool_calls=[
                            {
                                "args": {},
                                "id": "1",
                                "name": "mismatching_tool_call_id_tool",
                            }
                        ],
                    )
                ]
            },
            config=_create_config_with_runtime(),
        )

    # test validation (missing tool message in the update for parent graph is OK)
    @dec_tool
    def node_update_parent_tool():
        """No update"""
        return Command(update={"messages": []}, graph=Command.PARENT)

    assert ToolNode([node_update_parent_tool]).invoke(
        {
            "messages": [
                AIMessage(
                    "",
                    tool_calls=[
                        {"args": {}, "id": "1", "name": "node_update_parent_tool"}
                    ],
                )
            ]
        },
        config=_create_config_with_runtime(),
    ) == [Command(update={"messages": []}, graph=Command.PARENT)]


async def test_tool_node_command_list_input() -> None:

    @dec_tool
    def transfer_to_bob(tool_call_id: Annotated[str, InjectedToolCallId]):
        """Transfer to Bob"""
        return Command(
            update=[
                ToolMessage(content="Transferred to Bob", tool_call_id=tool_call_id)
            ],
            goto="bob",
            graph=Command.PARENT,
        )

    @dec_tool
    async def async_transfer_to_bob(tool_call_id: Annotated[str, InjectedToolCallId]):
        """Transfer to Bob"""
        return Command(
            update=[
                ToolMessage(content="Transferred to Bob", tool_call_id=tool_call_id)
            ],
            goto="bob",
            graph=Command.PARENT,
        )

    class CustomToolSchema(BaseModel):
        tool_call_id: Annotated[str, InjectedToolCallId]

    class MyCustomTool(BaseTool):
        def _run(*args: Any, **kwargs: Any):
            return Command(
                update=[
                    ToolMessage(
                        content="Transferred to Bob",
                        tool_call_id=kwargs["tool_call_id"],
                    )
                ],
                goto="bob",
                graph=Command.PARENT,
            )

        async def _arun(*args: Any, **kwargs: Any):
            return Command(
                update=[
                    ToolMessage(
                        content="Transferred to Bob",
                        tool_call_id=kwargs["tool_call_id"],
                    )
                ],
                goto="bob",
                graph=Command.PARENT,
            )

    custom_tool = MyCustomTool(
        name="custom_transfer_to_bob",
        description="Transfer to bob",
        args_schema=CustomToolSchema,
    )
    async_custom_tool = MyCustomTool(
        name="async_custom_transfer_to_bob",
        description="Transfer to bob",
        args_schema=CustomToolSchema,
    )

    # test mixing regular tools and tools returning commands
    def add(a: int, b: int) -> int:
        """Add two numbers"""
        return a + b

    result = ToolNode([add, transfer_to_bob]).invoke(
        [
            AIMessage(
                "",
                tool_calls=[
                    {"args": {"a": 1, "b": 2}, "id": "1", "name": "add"},
                    {"args": {}, "id": "2", "name": "transfer_to_bob"},
                ],
            )
        ],
        config=_create_config_with_runtime(),
    )

    assert result == [
        [
            ToolMessage(
                content="3",
                tool_call_id="1",
                name="add",
            )
        ],
        Command(
            update=[
                ToolMessage(
                    content="Transferred to Bob",
                    tool_call_id="2",
                    name="transfer_to_bob",
                )
            ],
            goto="bob",
            graph=Command.PARENT,
        ),
    ]

    # test tools returning commands

    # test sync tools
    for tool in [transfer_to_bob, custom_tool]:
        result = ToolNode([tool]).invoke(
            [AIMessage("", tool_calls=[{"args": {}, "id": "1", "name": tool.name}])],
            config=_create_config_with_runtime(),
        )
        assert result == [
            Command(
                update=[
                    ToolMessage(
                        content="Transferred to Bob",
                        tool_call_id="1",
                        name=tool.name,
                    )
                ],
                goto="bob",
                graph=Command.PARENT,
            )
        ]

    # test async tools
    for tool in [async_transfer_to_bob, async_custom_tool]:
        result = await ToolNode([tool]).ainvoke(
            [AIMessage("", tool_calls=[{"args": {}, "id": "1", "name": tool.name}])],
            config=_create_config_with_runtime(),
        )
        assert result == [
            Command(
                update=[
                    ToolMessage(
                        content="Transferred to Bob",
                        tool_call_id="1",
                        name=tool.name,
                    )
                ],
                goto="bob",
                graph=Command.PARENT,
            )
        ]

    # test multiple commands
    result = ToolNode([transfer_to_bob, custom_tool]).invoke(
        [
            AIMessage(
                "",
                tool_calls=[
                    {"args": {}, "id": "1", "name": "transfer_to_bob"},
                    {"args": {}, "id": "2", "name": "custom_transfer_to_bob"},
                ],
            )
        ],
        config=_create_config_with_runtime(),
    )
    assert result == [
        Command(
            update=[
                ToolMessage(
                    content="Transferred to Bob",
                    tool_call_id="1",
                    name="transfer_to_bob",
                )
            ],
            goto="bob",
            graph=Command.PARENT,
        ),
        Command(
            update=[
                ToolMessage(
                    content="Transferred to Bob",
                    tool_call_id="2",
                    name="custom_transfer_to_bob",
                )
            ],
            goto="bob",
            graph=Command.PARENT,
        ),
    ]

    # test validation (mismatch between input type and command.update type)
    with pytest.raises(ValueError):

        @dec_tool
        def list_update_tool(tool_call_id: Annotated[str, InjectedToolCallId]):
            """My tool"""
            return Command(
                update={
                    "messages": [ToolMessage(content="foo", tool_call_id=tool_call_id)]
                }
            )

        ToolNode([list_update_tool]).invoke(
            [
                AIMessage(
                    "",
                    tool_calls=[{"args": {}, "id": "1", "name": "list_update_tool"}],
                )
            ],
            config=_create_config_with_runtime(),
        )

    # test validation (missing tool message in the update for current graph)
    with pytest.raises(ValueError):

        @dec_tool
        def no_update_tool():
            """My tool"""
            return Command(update=[])

        ToolNode([no_update_tool]).invoke(
            [
                AIMessage(
                    "",
                    tool_calls=[{"args": {}, "id": "1", "name": "no_update_tool"}],
                )
            ],
            config=_create_config_with_runtime(),
        )

    # test validation (tool message with a wrong tool call ID)
    with pytest.raises(ValueError):

        @dec_tool
        def mismatching_tool_call_id_tool():
            """My tool"""
            return Command(update=[ToolMessage(content="foo", tool_call_id="2")])

        ToolNode([mismatching_tool_call_id_tool]).invoke(
            [
                AIMessage(
                    "",
                    tool_calls=[
                        {"args": {}, "id": "1", "name": "mismatching_tool_call_id_tool"}
                    ],
                )
            ],
            config=_create_config_with_runtime(),
        )

    # test validation (missing tool message in the update for parent graph is OK)
    @dec_tool
    def node_update_parent_tool():
        """No update"""
        return Command(update=[], graph=Command.PARENT)

    assert ToolNode([node_update_parent_tool]).invoke(
        [
            AIMessage(
                "",
                tool_calls=[{"args": {}, "id": "1", "name": "node_update_parent_tool"}],
            )
        ],
        config=_create_config_with_runtime(),
    ) == [Command(update=[], graph=Command.PARENT)]


def test_tool_node_parent_command_with_send() -> None:

    @dec_tool
    def transfer_to_alice(tool_call_id: Annotated[str, InjectedToolCallId]):
        """Transfer to Alice"""
        return Command(
            goto=[
                Send(
                    "alice",
                    {
                        "messages": [
                            ToolMessage(
                                content="Transferred to Alice",
                                name="transfer_to_alice",
                                tool_call_id=tool_call_id,
                            )
                        ]
                    },
                )
            ],
            graph=Command.PARENT,
        )

    @dec_tool
    def transfer_to_bob(tool_call_id: Annotated[str, InjectedToolCallId]):
        """Transfer to Bob"""
        return Command(
            goto=[
                Send(
                    "bob",
                    {
                        "messages": [
                            ToolMessage(
                                content="Transferred to Bob",
                                name="transfer_to_bob",
                                tool_call_id=tool_call_id,
                            )
                        ]
                    },
                )
            ],
            graph=Command.PARENT,
        )

    tool_calls = [
        {"args": {}, "id": "1", "name": "transfer_to_alice", "type": "tool_call"},
        {"args": {}, "id": "2", "name": "transfer_to_bob", "type": "tool_call"},
    ]

    result = ToolNode([transfer_to_alice, transfer_to_bob]).invoke(
        [AIMessage("", tool_calls=tool_calls)],
        config=_create_config_with_runtime(),
    )

    assert result == [
        Command(
            goto=[
                Send(
                    "alice",
                    {
                        "messages": [
                            ToolMessage(
                                content="Transferred to Alice",
                                name="transfer_to_alice",
                                tool_call_id="1",
                            )
                        ]
                    },
                ),
                Send(
                    "bob",
                    {
                        "messages": [
                            ToolMessage(
                                content="Transferred to Bob",
                                name="transfer_to_bob",
                                tool_call_id="2",
                            )
                        ]
                    },
                ),
            ],
            graph=Command.PARENT,
        )
    ]


async def test_tool_node_command_remove_all_messages() -> None:

    @dec_tool
    def remove_all_messages_tool(tool_call_id: Annotated[str, InjectedToolCallId]):
        """A tool that removes all messages."""
        return Command(update={"messages": [RemoveMessage(id=REMOVE_ALL_MESSAGES)]})

    tool_node = ToolNode([remove_all_messages_tool])
    tool_call = {
        "name": "remove_all_messages_tool",
        "args": {},
        "id": "tool_call_123",
    }
    result = await tool_node.ainvoke(
        {"messages": [AIMessage(content="", tool_calls=[tool_call])]},
        config=_create_config_with_runtime(),
    )

    assert isinstance(result, list)
    assert len(result) == 1
    command = result[0]
    assert isinstance(command, Command)
    assert command.update == {"messages": [RemoveMessage(id=REMOVE_ALL_MESSAGES)]}


class _InjectStateSchema(TypedDict):
    messages: list
    foo: str


class _InjectedStatePydanticV2Schema(BaseModel):
    messages: list
    foo: str


@dataclasses.dataclass
class _InjectedStateDataclassSchema:
    messages: list
    foo: str


_INJECTED_STATE_SCHEMAS = [
    _InjectStateSchema,
    _InjectedStatePydanticV2Schema,
    _InjectedStateDataclassSchema,
]

if sys.version_info < (3, 14):

    class _InjectedStatePydanticSchema(BaseModelV1):
        messages: list
        foo: str

    _INJECTED_STATE_SCHEMAS.append(_InjectedStatePydanticSchema)

T = TypeVar("T")


@pytest.mark.parametrize("schema_", _INJECTED_STATE_SCHEMAS)
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"]
        return state.foo

    def tool2(some_val: int, state: Annotated[T, InjectedState()]) -> str:
        """Tool 2 docstring."""
        if isinstance(state, dict):
            return state["foo"]
        return 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"), config=_create_config_with_runtime()
        )
        tool_message = result["messages"][-1]
        assert tool_message.content == "bar", f"Failed for tool={tool_name}"

        if tool_name == "tool3":
            failure_input = None
            with contextlib.suppress(Exception):
                failure_input = schema_(messages=[msg], notfoo="bar")
            if failure_input is not None:
                with pytest.raises(KeyError):
                    node.invoke(failure_input, config=_create_config_with_runtime())

                with pytest.raises(ValueError):
                    node.invoke([msg], config=_create_config_with_runtime())
        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, config=_create_config_with_runtime()
                )
                tool_message = messages_["messages"][-1]
                assert "KeyError" in tool_message.content
                tool_message = node.invoke([msg], config=_create_config_with_runtime())[
                    -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=""), config=_create_config_with_runtime()
    )
    tool_message = result["messages"][-1]
    assert tool_message.content == "hi?"

    result = node.invoke([msg], config=_create_config_with_runtime())
    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]}, config=_create_config_with_runtime(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"},
        config=_create_config_with_runtime(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)],
        config=_create_config_with_runtime(),
    )
    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) -> 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) -> dict[str, Any]:
        response = model.invoke(state["subgraph_messages"])
        model.tool_calls = []
        return {"subgraph_messages": response}

    builder = StateGraph(State)
    builder.add_node("agent", call_model)
    builder.add_node("tools", ToolNode([add], messages_key="subgraph_messages"))
    builder.add_conditional_edges(
        "agent", partial(tools_condition, messages_key="subgraph_messages")
    )
    builder.add_edge(START, "agent")
    builder.add_edge("tools", "agent")

    graph = builder.compile()
    result = graph.invoke({"subgraph_messages": [HumanMessage(content="hi")]})
    assert result["subgraph_messages"] == [
        _AnyIdHumanMessage(content="hi"),
        AIMessage(
            content="hi",
            id="0",
            tool_calls=[ToolCall(name=add.name, args={"a": 1, "b": 2}, id="test_id")],
        ),
        _AnyIdToolMessage(content="3", name=add.name, tool_call_id="test_id"),
        AIMessage(content="hi-hi-3", id="1"),
    ]


def test_tool_node_stream_writer() -> None:
    @dec_tool
    def streaming_tool(x: int) -> str:
        """Do something with writer."""
        my_writer = get_stream_writer()
        for value in ["foo", "bar", "baz"]:
            my_writer({"custom_tool_value": value})

        return x

    tool_node = ToolNode([streaming_tool])
    graph = (
        StateGraph(MessagesState)
        .add_node("tools", tool_node)
        .add_edge(START, "tools")
        .compile()
    )

    tool_call = {
        "name": "streaming_tool",
        "args": {"x": 1},
        "id": "1",
        "type": "tool_call",
    }
    inputs = {
        "messages": [AIMessage("", tool_calls=[tool_call])],
    }

    assert list(graph.stream(inputs, stream_mode="custom")) == [
        {"custom_tool_value": "foo"},
        {"custom_tool_value": "bar"},
        {"custom_tool_value": "baz"},
    ]
    assert list(graph.stream(inputs, stream_mode=["custom", "updates"])) == [
        ("custom", {"custom_tool_value": "foo"}),
        ("custom", {"custom_tool_value": "bar"}),
        ("custom", {"custom_tool_value": "baz"}),
        (
            "updates",
            {
                "tools": {
                    "messages": [
                        _AnyIdToolMessage(
                            content="1",
                            name="streaming_tool",
                            tool_call_id="1",
                        ),
                    ],
                },
            },
        ),
    ]


def test_tool_call_request_setattr_deprecation_warning():
    """Test that ToolCallRequest raises a deprecation warning on direct attribute modification."""

    # Create a mock ToolCall
    tool_call = {"name": "test", "args": {"a": 1}, "id": "call_1", "type": "tool_call"}

    # Create a ToolCallRequest
    request = ToolCallRequest(
        tool_call=tool_call,
        tool=None,
        state={"messages": []},
        runtime=None,
    )

    # Test 1: Direct attribute assignment should raise deprecation warning but still work
    with pytest.warns(DeprecationWarning, match="deprecated.*override"):
        request.tool_call = {"name": "other", "args": {}, "id": "call_2"}

    # Verify the attribute was actually modified
    assert request.tool_call == {"name": "other", "args": {}, "id": "call_2"}

    # Reset for further tests
    with warnings.catch_warnings():
        warnings.simplefilter("ignore")
        request.tool_call = tool_call

    # Test 2: override method should work without warnings
    with warnings.catch_warnings(record=True) as w:
        warnings.simplefilter("always")
        new_tool_call = {
            "name": "new_tool",
            "args": {"b": 2},
            "id": "call_3",
            "type": "tool_call",
        }
        new_request = request.override(tool_call=new_tool_call)

        # Verify no warning was raised
        assert len(w) == 0

    # Verify original is unchanged
    assert request.tool_call == tool_call

    # Verify new request has updated values
    assert new_request.tool_call == new_tool_call

    # Test 3: Initialization should not trigger warning
    with warnings.catch_warnings(record=True) as w:
        warnings.simplefilter("always")
        ToolCallRequest(
            tool_call=tool_call,
            tool=None,
            state={"messages": []},
            runtime=None,
        )
        # Verify no warning was raised during initialization
        assert len(w) == 0


async def test_tool_node_inject_async_all_types_signature_only() -> None:
    """Test all injection types without @tool decorator."""
    store = InMemoryStore()
    namespace = ("test",)
    store.put(namespace, "test_key", {"store_data": "from_store"})

    class TestState(TypedDict):
        messages: list
        foo: str
        bar: int

    async def comprehensive_async_tool(
        x: int,
        whole_state: Annotated[TestState, InjectedState],
        foo_field: Annotated[str, InjectedState("foo")],
        store: Annotated[BaseStore, InjectedStore()],
        runtime: ToolRuntime,
    ) -> str:
        """Async tool that uses all injection types."""
        bar_from_whole = whole_state["bar"]
        foo_value = foo_field
        store_val = store.get(namespace, "test_key").value["store_data"]
        foo_from_runtime = runtime.state["foo"]
        tool_call_id = runtime.tool_call_id

        return (
            f"x={x}, "
            f"bar_from_whole={bar_from_whole}, "
            f"foo_field={foo_value}, "
            f"store={store_val}, "
            f"foo_from_runtime={foo_from_runtime}, "
            f"tool_call_id={tool_call_id}"
        )

    node = ToolNode([comprehensive_async_tool], handle_tool_errors=True)
    tool_call = {
        "name": "comprehensive_async_tool",
        "args": {"x": 42},
        "id": "test_call_123",
        "type": "tool_call",
    }
    msg = AIMessage("hi?", tool_calls=[tool_call])

    config = _create_config_with_runtime(store=store)
    result = await node.ainvoke(
        {"messages": [msg], "foo": "foo_value", "bar": 99}, config=config
    )

    tool_message = result["messages"][-1]
    assert tool_message.content == (
        "x=42, "
        "bar_from_whole=99, "
        "foo_field=foo_value, "
        "store=from_store, "
        "foo_from_runtime=foo_value, "
        "tool_call_id=test_call_123"
    )


async def test_tool_node_inject_async_all_types_with_decorator() -> None:
    """Test all injection types with @tool decorator."""
    store = InMemoryStore()
    namespace = ("test",)
    store.put(namespace, "test_key", {"store_data": "from_store"})

    class TestState(TypedDict):
        messages: list
        foo: str
        bar: int

    @dec_tool
    async def comprehensive_async_tool(
        x: int,
        whole_state: Annotated[TestState, InjectedState],
        foo_field: Annotated[str, InjectedState("foo")],
        store: Annotated[BaseStore, InjectedStore()],
        runtime: ToolRuntime,
    ) -> str:
        """Async tool that uses all injection types."""
        bar_from_whole = whole_state["bar"]
        foo_value = foo_field
        store_val = store.get(namespace, "test_key").value["store_data"]
        foo_from_runtime = runtime.state["foo"]
        tool_call_id = runtime.tool_call_id

        return (
            f"x={x}, "
            f"bar_from_whole={bar_from_whole}, "
            f"foo_field={foo_value}, "
            f"store={store_val}, "
            f"foo_from_runtime={foo_from_runtime}, "
            f"tool_call_id={tool_call_id}"
        )

    node = ToolNode([comprehensive_async_tool], handle_tool_errors=True)
    tool_call = {
        "name": "comprehensive_async_tool",
        "args": {"x": 42},
        "id": "test_call_456",
        "type": "tool_call",
    }
    msg = AIMessage("hi?", tool_calls=[tool_call])

    config = _create_config_with_runtime(store=store)
    result = await node.ainvoke(
        {"messages": [msg], "foo": "foo_value", "bar": 99}, config=config
    )

    tool_message = result["messages"][-1]
    assert tool_message.content == (
        "x=42, "
        "bar_from_whole=99, "
        "foo_field=foo_value, "
        "store=from_store, "
        "foo_from_runtime=foo_value, "
        "tool_call_id=test_call_456"
    )


async def test_tool_node_inject_async_all_types_with_schema() -> None:
    """Test all injection types with explicit schema."""
    store = InMemoryStore()
    namespace = ("test",)
    store.put(namespace, "test_key", {"store_data": "from_store"})

    class TestState(TypedDict):
        messages: list
        foo: str
        bar: int

    class ComprehensiveToolSchema(BaseModel):
        model_config = {"arbitrary_types_allowed": True}
        x: int
        whole_state: Annotated[TestState, InjectedState]
        foo_field: Annotated[str, InjectedState("foo")]
        store: Annotated[BaseStore, InjectedStore()]
        runtime: ToolRuntime

    @dec_tool(args_schema=ComprehensiveToolSchema)
    async def comprehensive_async_tool(
        x: int,
        whole_state: Annotated[TestState, InjectedState],
        foo_field: Annotated[str, InjectedState("foo")],
        store: Annotated[BaseStore, InjectedStore()],
        runtime: ToolRuntime,
    ) -> str:
        """Async tool that uses all injection types."""
        bar_from_whole = whole_state["bar"]
        foo_value = foo_field
        store_val = store.get(namespace, "test_key").value["store_data"]
        foo_from_runtime = runtime.state["foo"]
        tool_call_id = runtime.tool_call_id

        return (
            f"x={x}, "
            f"bar_from_whole={bar_from_whole}, "
            f"foo_field={foo_value}, "
            f"store={store_val}, "
            f"foo_from_runtime={foo_from_runtime}, "
            f"tool_call_id={tool_call_id}"
        )

    node = ToolNode([comprehensive_async_tool], handle_tool_errors=True)
    tool_call = {
        "name": "comprehensive_async_tool",
        "args": {"x": 42},
        "id": "test_call_789",
        "type": "tool_call",
    }
    msg = AIMessage("hi?", tool_calls=[tool_call])

    config = _create_config_with_runtime(store=store)
    result = await node.ainvoke(
        {"messages": [msg], "foo": "foo_value", "bar": 99}, config=config
    )

    tool_message = result["messages"][-1]
    assert tool_message.content == (
        "x=42, "
        "bar_from_whole=99, "
        "foo_field=foo_value, "
        "store=from_store, "
        "foo_from_runtime=foo_value, "
        "tool_call_id=test_call_789"
    )


async def test_tool_node_tool_runtime_generic() -> None:
    """Test that ToolRuntime with generic type arguments is correctly injected."""

    @dataclasses.dataclass
    class MyContext:
        some_info: str

    @dec_tool
    def get_info(rt: ToolRuntime[MyContext]):
        """This tool returns info from context."""
        return rt.context.some_info

    # Create a mock runtime with context
    mock_runtime = _create_mock_runtime()
    mock_runtime.context = MyContext(some_info="test_info")

    config = {"configurable": {"__pregel_runtime": mock_runtime}}

    result = await ToolNode([get_info]).ainvoke(
        {
            "messages": [
                AIMessage(
                    "call tool",
                    tool_calls=[
                        {
                            "name": "get_info",
                            "args": {},
                            "id": "call_1",
                        }
                    ],
                )
            ]
        },
        config=config,
    )

    tool_message = result["messages"][-1]
    assert tool_message.type == "tool"
    assert tool_message.content == "test_info"
    assert tool_message.tool_call_id == "call_1"


def test_tool_node_inject_runtime_dynamic_tool_via_wrap_tool_call() -> None:
    """Test that ToolRuntime is injected for dynamically registered tools.

    Regression test for https://github.com/langchain-ai/langchain/issues/35305.
    When a tool is dynamically provided via wrap_tool_call (not registered at
    ToolNode init time), ToolRuntime should still be injected into the tool.
    """

    @dec_tool
    def static_tool(x: int) -> str:
        """A static tool registered at init."""
        return f"static: {x}"

    @dec_tool
    def dynamic_tool_with_runtime(x: int, runtime: ToolRuntime) -> str:
        """A dynamic tool that needs ToolRuntime injection."""
        return f"dynamic: x={x}, tool_call_id={runtime.tool_call_id}"

    def wrap_tool_call(request, execute):
        """Middleware that swaps in a dynamic tool."""
        if request.tool_call["name"] == "dynamic_tool_with_runtime":
            # Override tool to the dynamic one (not registered at init)
            new_request = request.override(tool=dynamic_tool_with_runtime)
            return execute(new_request)
        return execute(request)

    # ToolNode only knows about static_tool at init time
    tool_node = ToolNode(
        [static_tool],
        wrap_tool_call=wrap_tool_call,
    )

    # Verify the dynamic tool is NOT in the tool node's registered tools
    assert "dynamic_tool_with_runtime" not in tool_node.tools_by_name

    # Call the dynamic tool
    tool_call = {
        "name": "dynamic_tool_with_runtime",
        "args": {"x": 42},
        "id": "call_dynamic_1",
        "type": "tool_call",
    }
    msg = AIMessage("", tool_calls=[tool_call])
    result = tool_node.invoke(
        {"messages": [msg]},
        config=_create_config_with_runtime(),
    )

    # ToolRuntime should be injected and the tool should execute successfully
    tool_message = result["messages"][-1]
    assert tool_message.content == "dynamic: x=42, tool_call_id=call_dynamic_1"
    assert tool_message.tool_call_id == "call_dynamic_1"


async def test_tool_node_inject_runtime_dynamic_tool_via_wrap_tool_call_async() -> None:
    """Test that ToolRuntime is injected for dynamically registered tools (async).

    Async version of the regression test for
    https://github.com/langchain-ai/langchain/issues/35305.
    """

    @dec_tool
    def static_tool(x: int) -> str:
        """A static tool registered at init."""
        return f"static: {x}"

    @dec_tool
    async def dynamic_tool_with_runtime(x: int, runtime: ToolRuntime) -> str:
        """A dynamic async tool that needs ToolRuntime injection."""
        return f"dynamic: x={x}, tool_call_id={runtime.tool_call_id}"

    async def awrap_tool_call(request, execute):
        """Async middleware that swaps in a dynamic tool."""
        if request.tool_call["name"] == "dynamic_tool_with_runtime":
            new_request = request.override(tool=dynamic_tool_with_runtime)
            return await execute(new_request)
        return await execute(request)

    # ToolNode only knows about static_tool at init time
    tool_node = ToolNode(
        [static_tool],
        awrap_tool_call=awrap_tool_call,
    )

    # Verify the dynamic tool is NOT in the tool node's registered tools
    assert "dynamic_tool_with_runtime" not in tool_node.tools_by_name

    # Call the dynamic tool
    tool_call = {
        "name": "dynamic_tool_with_runtime",
        "args": {"x": 42},
        "id": "call_dynamic_2",
        "type": "tool_call",
    }
    msg = AIMessage("", tool_calls=[tool_call])
    result = await tool_node.ainvoke(
        {"messages": [msg]},
        config=_create_config_with_runtime(),
    )

    # ToolRuntime should be injected and the tool should execute successfully
    tool_message = result["messages"][-1]
    assert tool_message.content == "dynamic: x=42, tool_call_id=call_dynamic_2"
    assert tool_message.tool_call_id == "call_dynamic_2"


def test_tool_runtime_defaults_tools_to_empty_list() -> None:
    runtime = ToolRuntime(
        state={},
        context=None,
        config={},
        stream_writer=lambda *args, **kwargs: None,
        tool_call_id=None,
        store=None,
    )

    assert runtime.tools == []


def test_tool_runtime_forwards_execution_info_server_info_and_tools() -> None:
    """Test that execution_info, server_info, and tools are forwarded from Runtime to ToolRuntime."""

    exec_info = ExecutionInfo(
        thread_id="t-1",
        checkpoint_id="cp-1",
        checkpoint_ns="",
        task_id="tk-1",
        run_id="r-1",
    )
    server_info = ServerInfo(assistant_id="asst-1", graph_id="graph-1")

    mock_runtime = Mock()
    mock_runtime.store = None
    mock_runtime.context = None
    mock_runtime.stream_writer = lambda *args, **kwargs: None
    mock_runtime.execution_info = exec_info
    mock_runtime.server_info = server_info

    captured: dict = {}

    @dec_tool
    def info_tool(x: int, runtime: ToolRuntime) -> str:
        """Tool that captures runtime info."""
        captured["execution_info"] = runtime.execution_info
        captured["server_info"] = runtime.server_info
        captured["tools"] = runtime.tools
        return "ok"

    @dec_tool
    def other_tool(y: int) -> str:
        """Another tool available to the runtime."""
        return str(y)

    node = ToolNode([info_tool, other_tool])
    tool_call = {
        "name": "info_tool",
        "args": {"x": 1},
        "id": "call-1",
        "type": "tool_call",
    }
    msg = AIMessage("", tool_calls=[tool_call])
    config: RunnableConfig = {"configurable": {"__pregel_runtime": mock_runtime}}
    result = node.invoke({"messages": [msg]}, config=config)

    assert result["messages"][-1].content == "ok"
    assert captured["execution_info"] is exec_info
    assert captured["execution_info"].thread_id == "t-1"
    assert captured["execution_info"].task_id == "tk-1"
    assert captured["server_info"] is server_info
    assert captured["server_info"].assistant_id == "asst-1"
    assert [tool.name for tool in captured["tools"]] == ["info_tool", "other_tool"]


async def test_tool_runtime_forwards_execution_info_server_info_and_tools_async() -> (
    None
):
    """Test that execution_info, server_info, and tools are forwarded in async path."""

    exec_info = ExecutionInfo(
        thread_id="t-2",
        checkpoint_id="cp-2",
        checkpoint_ns="",
        task_id="tk-2",
        run_id="r-2",
    )
    server_info = ServerInfo(assistant_id="asst-2", graph_id="graph-2")

    mock_runtime = Mock()
    mock_runtime.store = None
    mock_runtime.context = None
    mock_runtime.stream_writer = lambda *args, **kwargs: None
    mock_runtime.execution_info = exec_info
    mock_runtime.server_info = server_info

    captured: dict = {}

    @dec_tool
    async def info_tool_async(x: int, runtime: ToolRuntime) -> str:
        """Async tool that captures runtime info."""
        captured["execution_info"] = runtime.execution_info
        captured["server_info"] = runtime.server_info
        captured["tools"] = runtime.tools
        return "ok"

    @dec_tool
    async def other_tool_async(y: int) -> str:
        """Another async tool available to the runtime."""
        return str(y)

    node = ToolNode([info_tool_async, other_tool_async])
    tool_call = {
        "name": "info_tool_async",
        "args": {"x": 1},
        "id": "call-2",
        "type": "tool_call",
    }
    msg = AIMessage("", tool_calls=[tool_call])
    config: RunnableConfig = {"configurable": {"__pregel_runtime": mock_runtime}}
    result = await node.ainvoke({"messages": [msg]}, config=config)

    assert result["messages"][-1].content == "ok"
    assert captured["execution_info"] is exec_info
    assert captured["execution_info"].thread_id == "t-2"
    assert captured["server_info"] is server_info
    assert captured["server_info"].graph_id == "graph-2"
    assert [tool.name for tool in captured["tools"]] == [
        "info_tool_async",
        "other_tool_async",
    ]


# --- InjectedToolArg security tests ---


def test_tool_node_strips_plain_injected_tool_arg() -> None:
    """Plain InjectedToolArg values supplied by the LLM should be stripped."""

    @dec_tool
    def read_secret(
        query: str,
        auth: Annotated[dict, InjectedToolArg()],
    ) -> str:
        """Return secret data based on auth role."""
        if auth.get("role") == "admin":
            return "ADMIN_SECRET"
        return "PUBLIC_DATA"

    node = ToolNode([read_secret], handle_tool_errors=True)

    # LLM tries to supply the hidden 'auth' field
    tool_call = {
        "name": "read_secret",
        "args": {"query": "hello", "auth": {"role": "admin"}},
        "id": "call-1",
        "type": "tool_call",
    }
    msg = AIMessage("", tool_calls=[tool_call])
    result = node.invoke({"messages": [msg]}, config=_create_config_with_runtime())
    tool_message = result["messages"][-1]
    # auth should have been stripped, so tool should fail (missing required arg)
    assert "ADMIN_SECRET" not in tool_message.content


def test_tool_node_strips_custom_injected_tool_arg_subclass() -> None:
    """Custom InjectedToolArg subclasses should also be stripped."""

    class InjectedAuth(InjectedToolArg):
        pass

    @dec_tool
    def read_secret(
        query: str,
        auth: Annotated[dict, InjectedAuth()],
    ) -> str:
        """Return secret data based on auth role."""
        if auth.get("role") == "admin":
            return "ADMIN_SECRET"
        return "PUBLIC_DATA"

    node = ToolNode([read_secret], handle_tool_errors=True)

    tool_call = {
        "name": "read_secret",
        "args": {"query": "hello", "auth": {"role": "admin"}},
        "id": "call-1",
        "type": "tool_call",
    }
    msg = AIMessage("", tool_calls=[tool_call])
    result = node.invoke({"messages": [msg]}, config=_create_config_with_runtime())
    tool_message = result["messages"][-1]
    assert "ADMIN_SECRET" not in tool_message.content


def test_tool_node_injected_state_overwrites_llm_value() -> None:
    """InjectedState should use graph state, not LLM-supplied values."""

    @dec_tool
    def read_secret(
        query: str,
        auth: Annotated[dict, InjectedState("auth")],
    ) -> str:
        """Return secret data based on auth from graph state."""
        if auth.get("role") == "admin":
            return "ADMIN_SECRET"
        return "PUBLIC_DATA"

    node = ToolNode([read_secret])

    # LLM tries to supply auth as admin
    tool_call = {
        "name": "read_secret",
        "args": {"query": "hello", "auth": {"role": "admin"}},
        "id": "call-1",
        "type": "tool_call",
    }
    msg = AIMessage("", tool_calls=[tool_call])

    # Graph state has auth as viewer
    result = node.invoke(
        {"messages": [msg], "auth": {"role": "viewer"}},
        config=_create_config_with_runtime(),
    )
    tool_message = result["messages"][-1]
    assert tool_message.content == "PUBLIC_DATA"


class _ReturningTool(BaseTool):
    """A tool that returns a configured value verbatim."""

    name: str = "list_tool"
    description: str = "Returns a configured value"
    return_value: Any = None

    def _run(self, **kwargs: Any) -> Any:
        return self.return_value

    async def _arun(self, **kwargs: Any) -> Any:
        return self.return_value


def _list_tool_call(outer_id: str = "call-1") -> dict[str, Any]:
    return {"name": "list_tool", "args": {}, "id": outer_id, "type": "tool_call"}


def _invoke_returning(
    return_value: Any,
    *,
    outer_id: str = "call-1",
    handle_tool_errors: bool = True,
) -> Any:
    node = ToolNode(
        [_ReturningTool(return_value=return_value)],
        handle_tool_errors=handle_tool_errors,
    )
    return node.invoke(
        {"messages": [AIMessage("", tool_calls=[_list_tool_call(outer_id)])]},
        config=_create_config_with_runtime(),
    )


def test_tool_node_list_return_command_and_tool_message() -> None:
    """Valid: tool returns [Command(update={...}), ToolMessage(...)]."""
    outer_id = "call-1"
    result = _invoke_returning(
        [
            Command(update={"foo": "bar"}),
            ToolMessage(content="done", tool_call_id=outer_id),
        ]
    )
    assert isinstance(result, list)
    commands = [r for r in result if isinstance(r, Command)]
    assert len(commands) == 1
    assert commands[0].update == {"foo": "bar"}
    non_commands = [r for r in result if not isinstance(r, Command)]
    assert len(non_commands) == 1
    assert isinstance(non_commands[0], dict)
    msgs = non_commands[0]["messages"]
    assert len(msgs) == 1
    assert isinstance(msgs[0], ToolMessage)
    assert msgs[0].content == "done"
    assert msgs[0].tool_call_id == outer_id


def test_tool_node_list_return_nested_terminator() -> None:
    """Valid: terminator nested inside Command.update['messages']."""
    outer_id = "call-1"
    result = _invoke_returning(
        [
            Command(update={"foo": "bar"}),
            Command(
                update={
                    "messages": [ToolMessage(content="done", tool_call_id=outer_id)]
                }
            ),
        ]
    )
    assert isinstance(result, list)
    commands = [r for r in result if isinstance(r, Command)]
    assert len(commands) == 2
    updates = [c.update for c in commands]
    assert {"foo": "bar"} in updates
    msgs_update = next(u for u in updates if "messages" in (u or {}))
    assert any(
        isinstance(m, ToolMessage) and m.tool_call_id == outer_id
        for m in msgs_update["messages"]
    )


def test_tool_node_list_return_parent_goto_with_terminator() -> None:
    """Valid: [Command(graph=PARENT, goto=[Send(...)]), ToolMessage(...)]."""
    outer_id = "call-1"
    result = _invoke_returning(
        [
            Command(graph=Command.PARENT, goto=[Send("child", {})]),
            ToolMessage(content="ok", tool_call_id=outer_id),
        ]
    )
    assert isinstance(result, list)
    parent_cmds = [
        r for r in result if isinstance(r, Command) and r.graph is Command.PARENT
    ]
    assert len(parent_cmds) == 1
    assert isinstance(parent_cmds[0].goto, list)
    assert any(isinstance(s, Send) for s in parent_cmds[0].goto)
    non_commands = [r for r in result if not isinstance(r, Command)]
    assert len(non_commands) == 1


def test_tool_node_list_return_no_terminator_raises() -> None:
    """Invalid: list with no terminating ToolMessage."""
    with pytest.raises(ValueError, match="0 messages bound to tool_call_id"):
        _invoke_returning([Command(update={"foo": "bar"})], handle_tool_errors=False)


def test_tool_node_list_return_multiple_terminators_raises() -> None:
    """Invalid: list with two terminating ToolMessages."""
    outer_id = "call-1"
    with pytest.raises(ValueError, match="2 messages bound to tool_call_id"):
        _invoke_returning(
            [
                ToolMessage(content="a", tool_call_id=outer_id),
                ToolMessage(content="b", tool_call_id=outer_id),
            ],
            handle_tool_errors=False,
        )


def test_tool_node_list_return_validation_error_handled() -> None:
    """handle_tool_errors=True converts validation errors to an error ToolMessage."""
    result = _invoke_returning([Command(update={"foo": "bar"})])
    assert isinstance(result, dict)
    msg = result["messages"][0]
    assert isinstance(msg, ToolMessage)
    assert msg.status == "error"
    assert "0 messages bound to tool_call_id" in msg.content


async def test_tool_node_list_return_async_smoke() -> None:
    """Async path parallels sync for the happy case."""
    outer_id = "call-1"
    node = ToolNode(
        [
            _ReturningTool(
                return_value=[
                    Command(update={"foo": "bar"}),
                    ToolMessage(content="done", tool_call_id=outer_id),
                ]
            )
        ]
    )
    result = await node.ainvoke(
        {"messages": [AIMessage("", tool_calls=[_list_tool_call(outer_id)])]},
        config=_create_config_with_runtime(),
    )
    assert isinstance(result, list)
    commands = [r for r in result if isinstance(r, Command)]
    assert len(commands) == 1 and commands[0].update == {"foo": "bar"}


def test_tool_node_list_return_mixed_with_regular_tool() -> None:
    """List-returning tool and a regular tool dispatched from the same AIMessage."""
    list_tool_id = "call-list"
    regular_tool_id = "call-regular"
    list_tool = _ReturningTool(
        return_value=[
            Command(update={"foo": "bar"}),
            ToolMessage(content="list done", tool_call_id=list_tool_id),
        ]
    )

    def regular_tool(x: int) -> str:
        """A normal tool."""
        return f"regular: {x}"

    tool_calls = [
        {"name": "list_tool", "args": {}, "id": list_tool_id, "type": "tool_call"},
        {
            "name": "regular_tool",
            "args": {"x": 7},
            "id": regular_tool_id,
            "type": "tool_call",
        },
    ]
    node = ToolNode([list_tool, regular_tool])
    result = node.invoke(
        {"messages": [AIMessage("", tool_calls=tool_calls)]},
        config=_create_config_with_runtime(),
    )
    assert isinstance(result, list)
    commands = [r for r in result if isinstance(r, Command)]
    assert len(commands) == 1
    assert commands[0].update == {"foo": "bar"}
    all_msgs = [m for r in result if isinstance(r, dict) for m in r["messages"]]
    tool_call_ids = {m.tool_call_id for m in all_msgs}
    assert list_tool_id in tool_call_ids
    assert regular_tool_id in tool_call_ids