import contextlib
import dataclasses
import json
import sys
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 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.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,
    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.
    """
    from langgraph.runtime import ExecutionInfo

    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:
    from langchain_core.tools.base import InjectedToolCallId

    @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:
    from langchain_core.tools.base import InjectedToolCallId

    @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:
    from langchain_core.tools.base import InjectedToolCallId

    @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:
    from langchain_core.tools.base import InjectedToolCallId

    @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."""
    import warnings

    from langgraph.prebuilt.tool_node import ToolCallRequest

    # 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."""
    from langgraph.runtime import ExecutionInfo, ServerInfo

    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."""
    from langgraph.runtime import ExecutionInfo, ServerInfo

    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