langchain 1.0.0a12__py3-none-any.whl → 1.0.0a14__py3-none-any.whl

langchain/tools/tool_node.py

@@ -5,7 +5,7 @@ This module provides prebuilt functionality for executing tools in LangGraph.
 Tools are functions that models can call to interact with external systems,
 APIs, databases, or perform computations.
 
-The module implements several key design patterns:
+The module implements design patterns for:
 - Parallel execution of multiple tool calls for efficiency
 - Robust error handling with customizable error messages
 - State injection for tools that need access to graph state
@@ -38,8 +38,9 @@ from __future__ import annotations
 import asyncio
 import inspect
 import json
+from collections.abc import Awaitable, Callable
 from copy import copy, deepcopy
-from dataclasses import replace
+from dataclasses import dataclass, replace
 from types import UnionType
 from typing import (
     TYPE_CHECKING,
@@ -47,6 +48,7 @@ from typing import (
     Any,
     Literal,
     Optional,
+    TypedDict,
     Union,
     cast,
     get_args,
@@ -70,16 +72,18 @@ from langchain_core.tools import BaseTool, InjectedToolArg
 from langchain_core.tools import tool as create_tool
 from langchain_core.tools.base import (
     TOOL_MESSAGE_BLOCK_TYPES,
+    ToolException,
     get_all_basemodel_annotations,
 )
 from langgraph._internal._runnable import RunnableCallable
 from langgraph.errors import GraphBubbleUp
 from langgraph.graph.message import REMOVE_ALL_MESSAGES
+from langgraph.runtime import get_runtime
 from langgraph.types import Command, Send
 from pydantic import BaseModel, ValidationError
 
 if TYPE_CHECKING:
-    from collections.abc import Callable, Sequence
+    from collections.abc import Sequence
 
     from langchain_core.runnables import RunnableConfig
     from langgraph.store.base import BaseStore
@@ -100,24 +104,132 @@ TOOL_INVOCATION_ERROR_TEMPLATE = (
 )
 
 
+@dataclass()
+class ToolCallRequest:
+    """Tool execution request passed to tool call interceptors.
+
+    Attributes:
+        tool_call: Tool call dict with name, args, and id from model output.
+        tool: BaseTool instance to be invoked.
+        state: Agent state (dict, list, or BaseModel).
+        runtime: LangGraph runtime context (optional, None if outside graph).
+    """
+
+    tool_call: ToolCall
+    tool: BaseTool
+    state: Any
+    runtime: Any
+
+
+ToolCallWrapper = Callable[
+    [ToolCallRequest, Callable[[ToolCallRequest], ToolMessage | Command]],
+    ToolMessage | Command,
+]
+"""Wrapper for tool call execution with multi-call support.
+
+Wrapper receives:
+    request: ToolCallRequest with tool_call, tool, state, and runtime.
+    execute: Callable to execute the tool (CAN BE CALLED MULTIPLE TIMES).
+
+Returns:
+    ToolMessage or Command (the final result).
+
+The execute callable can be invoked multiple times for retry logic,
+with potentially modified requests each time. Each call to execute
+is independent and stateless.
+
+Note:
+    When implementing middleware for `create_agent`, use
+    `AgentMiddleware.wrap_tool_call` which provides properly typed
+    state parameter for better type safety.
+
+Examples:
+    Passthrough (execute once):
+
+    def handler(request, execute):
+        return execute(request)
+
+    Modify request before execution:
+
+    def handler(request, execute):
+        request.tool_call["args"]["value"] *= 2
+        return execute(request)
+
+    Retry on error (execute multiple times):
+
+    def handler(request, execute):
+        for attempt in range(3):
+            try:
+                result = execute(request)
+                if is_valid(result):
+                    return result
+            except Exception:
+                if attempt == 2:
+                    raise
+        return result
+
+    Conditional retry based on response:
+
+    def handler(request, execute):
+        for attempt in range(3):
+            result = execute(request)
+            if isinstance(result, ToolMessage) and result.status != "error":
+                return result
+            if attempt < 2:
+                continue
+            return result
+
+    Cache/short-circuit without calling execute:
+
+    def handler(request, execute):
+        if cached := get_cache(request):
+            return ToolMessage(content=cached, tool_call_id=request.tool_call["id"])
+        result = execute(request)
+        save_cache(request, result)
+        return result
+"""
+
+AsyncToolCallWrapper = Callable[
+    [ToolCallRequest, Callable[[ToolCallRequest], Awaitable[ToolMessage | Command]]],
+    Awaitable[ToolMessage | Command],
+]
+"""Async wrapper for tool call execution with multi-call support."""
+
+
+class ToolCallWithContext(TypedDict):
+    """ToolCall with additional context for graph state.
+
+    This is an internal data structure meant to help the ToolNode accept
+    tool calls with additional context (e.g. state) when dispatched using the
+    Send API.
+
+    The Send API is used in create_agent to distribute tool calls in parallel
+    and support human-in-the-loop workflows where graph execution may be paused
+    for an indefinite time.
+    """
+
+    tool_call: ToolCall
+    __type: Literal["tool_call_with_context"]
+    """Type to parameterize the payload.
+
+    Using "__" as a prefix to be defensive against potential name collisions with
+    regular user state.
+    """
+    state: Any
+    """The state is provided as additional context."""
+
+
 def msg_content_output(output: Any) -> str | list[dict]:
-    """Convert tool output to valid message content format.
+    """Convert tool output to ToolMessage content format.
 
-    LangChain ToolMessages accept either string content or a list of content blocks.
-    This function ensures tool outputs are properly formatted for message consumption
-    by attempting to preserve structured data when possible, falling back to JSON
-    serialization or string conversion.
+    Handles str, list[dict] (content blocks), and arbitrary objects by attempting
+    JSON serialization with fallback to str().
 
     Args:
-        output: The raw output from a tool execution. Can be any type.
+        output: Tool execution output of any type.
 
     Returns:
-        Either a string representation of the output or a list of content blocks
-        if the output is already in the correct format for structured content.
-
-    Note:
-        This function prioritizes backward compatibility by defaulting to JSON
-        serialization rather than supporting all possible message content formats.
+        String or list of content blocks suitable for ToolMessage.content.
     """
     if isinstance(output, str) or (
         isinstance(output, list)
@@ -134,8 +246,11 @@ def msg_content_output(output: Any) -> str | list[dict]:
         return str(output)
 
 
-class ToolInvocationError(Exception):
-    """Exception raised when a tool invocation fails due to invalid arguments."""
+class ToolInvocationError(ToolException):
+    """An error occurred while invoking a tool due to invalid arguments.
+
+    This exception is only raised when invoking a tool using the ToolNode!
+    """
 
     def __init__(
         self, tool_name: str, source: ValidationError, tool_kwargs: dict[str, Any]
@@ -181,7 +296,7 @@ def _handle_tool_error(
     Args:
         e: The exception that occurred during tool execution.
         flag: Configuration for how to handle the error. Can be:
-            - bool: If True, use default error template
+            - bool: If `True`, use default error template
             - str: Use this string as the error message
             - Callable: Call this function with the exception to get error message
             - tuple: Not used in this context (handled by caller)
@@ -277,7 +392,7 @@ def _infer_handled_types(handler: Callable[..., str]) -> tuple[type[Exception],
     return (Exception,)
 
 
-class ToolNode(RunnableCallable):
+class _ToolNode(RunnableCallable):
     """A node for executing tools in LangGraph workflows.
 
     Handles tool execution patterns including function calls, state injection,
@@ -286,25 +401,25 @@ class ToolNode(RunnableCallable):
 
     Input Formats:
         1. Graph state with `messages` key that has a list of messages:
-           - Common representation for agentic workflows
-           - Supports custom messages key via ``messages_key`` parameter
+            - Common representation for agentic workflows
+            - Supports custom messages key via `messages_key` parameter
 
-        2. **Message List**: ``[AIMessage(..., tool_calls=[...])]``
-           - List of messages with tool calls in the last AIMessage
+        2. **Message List**: `[AIMessage(..., tool_calls=[...])]`
+            - List of messages with tool calls in the last AIMessage
 
-        3. **Direct Tool Calls**: ``[{"name": "tool", "args": {...}, "id": "1", "type": "tool_call"}]``
-           - Bypasses message parsing for direct tool execution
-           - For programmatic tool invocation and testing
+        3. **Direct Tool Calls**: `[{"name": "tool", "args": {...}, "id": "1", "type": "tool_call"}]`
+            - Bypasses message parsing for direct tool execution
+            - For programmatic tool invocation and testing
 
     Output Formats:
         Output format depends on input type and tool behavior:
 
         **For Regular tools**:
-        - Dict input → ``{"messages": [ToolMessage(...)]}``
-        - List input → ``[ToolMessage(...)]``
+        - Dict input → `{"messages": [ToolMessage(...)]}`
+        - List input → `[ToolMessage(...)]`
 
         **For Command tools**:
-        - Returns ``[Command(...)]`` or mixed list with regular tool outputs
+        - Returns `[Command(...)]` or mixed list with regular tool outputs
         - Commands can update state, trigger navigation, or send messages
 
     Args:
@@ -314,29 +429,29 @@ class ToolNode(RunnableCallable):
         name: The name identifier for this node in the graph. Used for debugging
             and visualization. Defaults to "tools".
         tags: Optional metadata tags to associate with the node for filtering
-            and organization. Defaults to None.
+            and organization. Defaults to `None`.
         handle_tool_errors: Configuration for error handling during tool execution.
             Supports multiple strategies:
 
             - **True**: Catch all errors and return a ToolMessage with the default
-              error template containing the exception details.
+                error template containing the exception details.
             - **str**: Catch all errors and return a ToolMessage with this custom
-              error message string.
+                error message string.
             - **type[Exception]**: Only catch exceptions with the specified type and
-              return the default error message for it.
+                return the default error message for it.
             - **tuple[type[Exception], ...]**: Only catch exceptions with the specified
-              types and return default error messages for them.
+                types and return default error messages for them.
             - **Callable[..., str]**: Catch exceptions matching the callable's signature
-              and return the string result of calling it with the exception.
+                and return the string result of calling it with the exception.
             - **False**: Disable error handling entirely, allowing exceptions to
-              propagate.
+                propagate.
 
             Defaults to a callable that:
                 - catches tool invocation errors (due to invalid arguments provided by the model) and returns a descriptive error message
                 - ignores tool execution errors (they will be re-raised)
 
         messages_key: The key in the state dictionary that contains the message list.
-            This same key will be used for the output ToolMessages.
+            This same key will be used for the output `ToolMessage` objects.
             Defaults to "messages".
             Allows custom state schemas with different message field names.
 
@@ -394,8 +509,10 @@ class ToolNode(RunnableCallable):
         | type[Exception]
         | tuple[type[Exception], ...] = _default_handle_tool_errors,
         messages_key: str = "messages",
+        wrap_tool_call: ToolCallWrapper | None = None,
+        awrap_tool_call: AsyncToolCallWrapper | None = None,
     ) -> None:
-        """Initialize the ToolNode with the provided tools and configuration.
+        """Initialize ToolNode with tools and configuration.
 
         Args:
             tools: Sequence of tools to make available for execution.
@@ -403,6 +520,11 @@ class ToolNode(RunnableCallable):
             tags: Optional metadata tags.
             handle_tool_errors: Error handling configuration.
             messages_key: State key containing messages.
+            wrap_tool_call: Sync wrapper function to intercept tool execution. Receives
+                ToolCallRequest and execute callable, returns ToolMessage or Command.
+                Enables retries, caching, request modification, and control flow.
+            awrap_tool_call: Async wrapper function to intercept tool execution.
+                If not provided, falls back to wrap_tool_call for async execution.
         """
         super().__init__(self._func, self._afunc, name=name, tags=tags, trace=False)
         self._tools_by_name: dict[str, BaseTool] = {}
@@ -410,6 +532,8 @@ class ToolNode(RunnableCallable):
         self._tool_to_store_arg: dict[str, str | None] = {}
         self._handle_tool_errors = handle_tool_errors
         self._messages_key = messages_key
+        self._wrap_tool_call = wrap_tool_call
+        self._awrap_tool_call = awrap_tool_call
         for tool in tools:
             if not isinstance(tool, BaseTool):
                 tool_ = create_tool(cast("type[BaseTool]", tool))
@@ -431,11 +555,23 @@ class ToolNode(RunnableCallable):
         *,
         store: Optional[BaseStore],  # noqa: UP045
     ) -> Any:
-        tool_calls, input_type = self._parse_input(input, store)
+        try:
+            runtime = get_runtime()
+        except RuntimeError:
+            # Running outside of LangGraph runtime context (e.g., unit tests)
+            runtime = None
+
+        tool_calls, input_type = self._parse_input(input)
+        tool_calls = [self._inject_tool_args(call, input, store) for call in tool_calls]
+
         config_list = get_config_list(config, len(tool_calls))
         input_types = [input_type] * len(tool_calls)
+        inputs = [input] * len(tool_calls)
+        runtimes = [runtime] * len(tool_calls)
         with get_executor_for_config(config) as executor:
-            outputs = [*executor.map(self._run_one, tool_calls, input_types, config_list)]
+            outputs = [
+                *executor.map(self._run_one, tool_calls, input_types, config_list, inputs, runtimes)
+            ]
 
         return self._combine_tool_outputs(outputs, input_type)
 
@@ -446,9 +582,16 @@ class ToolNode(RunnableCallable):
         *,
         store: Optional[BaseStore],  # noqa: UP045
     ) -> Any:
-        tool_calls, input_type = self._parse_input(input, store)
+        try:
+            runtime = get_runtime()
+        except RuntimeError:
+            # Running outside of LangGraph runtime context (e.g., unit tests)
+            runtime = None
+
+        tool_calls, input_type = self._parse_input(input)
+        tool_calls = [self._inject_tool_args(call, input, store) for call in tool_calls]
         outputs = await asyncio.gather(
-            *(self._arun_one(call, input_type, config) for call in tool_calls)
+            *(self._arun_one(call, input_type, config, input, runtime) for call in tool_calls)
         )
 
         return self._combine_tool_outputs(outputs, input_type)
@@ -495,20 +638,30 @@ class ToolNode(RunnableCallable):
             combined_outputs.append(parent_command)
         return combined_outputs
 
-    def _run_one(
+    def _execute_tool_sync(
         self,
-        call: ToolCall,
+        request: ToolCallRequest,
         input_type: Literal["list", "dict", "tool_calls"],
         config: RunnableConfig,
     ) -> ToolMessage | Command:
-        """Run a single tool call synchronously."""
-        if invalid_tool_message := self._validate_tool_call(call):
-            return invalid_tool_message
+        """Execute tool call with configured error handling.
 
-        try:
-            call_args = {**call, "type": "tool_call"}
-            tool = self.tools_by_name[call["name"]]
+        Args:
+            request: Tool execution request.
+            input_type: Input format.
+            config: Runnable configuration.
+
+        Returns:
+            ToolMessage or Command.
+
+        Raises:
+            Exception: If tool fails and handle_tool_errors is False.
+        """
+        call = request.tool_call
+        tool = request.tool
+        call_args = {**call, "type": "tool_call"}
 
+        try:
             try:
                 response = tool.invoke(call_args, config)
             except ValidationError as exc:
@@ -526,6 +679,7 @@ class ToolNode(RunnableCallable):
         except GraphBubbleUp:
             raise
         except Exception as e:
+            # Determine which exception types are handled
             handled_types: tuple[type[Exception], ...]
             if isinstance(self._handle_tool_errors, type) and issubclass(
                 self._handle_tool_errors, Exception
@@ -541,10 +695,11 @@ class ToolNode(RunnableCallable):
                 # default behavior is catching all exceptions
                 handled_types = (Exception,)
 
-            # Unhandled
+            # Check if this error should be handled
             if not self._handle_tool_errors or not isinstance(e, handled_types):
                 raise
-            # Handled
+
+            # Error is handled - create error ToolMessage
             content = _handle_tool_error(e, flag=self._handle_tool_errors)
             return ToolMessage(
                 content=content,
@@ -553,28 +708,102 @@ class ToolNode(RunnableCallable):
                 status="error",
             )
 
+        # Process successful response
         if isinstance(response, Command):
-            return self._validate_tool_command(response, call, input_type)
+            # Validate Command before returning to handler
+            return self._validate_tool_command(response, request.tool_call, input_type)
         if isinstance(response, ToolMessage):
             response.content = cast("str | list", msg_content_output(response.content))
             return response
+
         msg = f"Tool {call['name']} returned unexpected type: {type(response)}"
         raise TypeError(msg)
 
-    async def _arun_one(
+    def _run_one(
         self,
         call: ToolCall,
         input_type: Literal["list", "dict", "tool_calls"],
         config: RunnableConfig,
+        input: list[AnyMessage] | dict[str, Any] | BaseModel,
+        runtime: Any,
     ) -> ToolMessage | Command:
-        """Run a single tool call asynchronously."""
+        """Execute single tool call with wrap_tool_call wrapper if configured.
+
+        Args:
+            call: Tool call dict.
+            input_type: Input format.
+            config: Runnable configuration.
+            input: Agent state.
+            runtime: LangGraph runtime or None.
+
+        Returns:
+            ToolMessage or Command.
+        """
         if invalid_tool_message := self._validate_tool_call(call):
             return invalid_tool_message
 
+        tool = self.tools_by_name[call["name"]]
+
+        # Extract state from ToolCallWithContext if present
+        state = self._extract_state(input)
+
+        # Create the tool request with state and runtime
+        tool_request = ToolCallRequest(
+            tool_call=call,
+            tool=tool,
+            state=state,
+            runtime=runtime,
+        )
+
+        if self._wrap_tool_call is None:
+            # No wrapper - execute directly
+            return self._execute_tool_sync(tool_request, input_type, config)
+
+        # Define execute callable that can be called multiple times
+        def execute(req: ToolCallRequest) -> ToolMessage | Command:
+            """Execute tool with given request. Can be called multiple times."""
+            return self._execute_tool_sync(req, input_type, config)
+
+        # Call wrapper with request and execute callable
         try:
-            call_args = {**call, "type": "tool_call"}
-            tool = self.tools_by_name[call["name"]]
+            return self._wrap_tool_call(tool_request, execute)
+        except Exception as e:
+            # Wrapper threw an exception
+            if not self._handle_tool_errors:
+                raise
+            # Convert to error message
+            content = _handle_tool_error(e, flag=self._handle_tool_errors)
+            return ToolMessage(
+                content=content,
+                name=tool_request.tool_call["name"],
+                tool_call_id=tool_request.tool_call["id"],
+                status="error",
+            )
+
+    async def _execute_tool_async(
+        self,
+        request: ToolCallRequest,
+        input_type: Literal["list", "dict", "tool_calls"],
+        config: RunnableConfig,
+    ) -> ToolMessage | Command:
+        """Execute tool call asynchronously with configured error handling.
+
+        Args:
+            request: Tool execution request.
+            input_type: Input format.
+            config: Runnable configuration.
+
+        Returns:
+            ToolMessage or Command.
 
+        Raises:
+            Exception: If tool fails and handle_tool_errors is False.
+        """
+        call = request.tool_call
+        tool = request.tool
+        call_args = {**call, "type": "tool_call"}
+
+        try:
             try:
                 response = await tool.ainvoke(call_args, config)
             except ValidationError as exc:
@@ -592,6 +821,7 @@ class ToolNode(RunnableCallable):
         except GraphBubbleUp:
             raise
         except Exception as e:
+            # Determine which exception types are handled
             handled_types: tuple[type[Exception], ...]
             if isinstance(self._handle_tool_errors, type) and issubclass(
                 self._handle_tool_errors, Exception
@@ -607,12 +837,12 @@ class ToolNode(RunnableCallable):
                 # default behavior is catching all exceptions
                 handled_types = (Exception,)
 
-            # Unhandled
+            # Check if this error should be handled
             if not self._handle_tool_errors or not isinstance(e, handled_types):
                 raise
-            # Handled
-            content = _handle_tool_error(e, flag=self._handle_tool_errors)
 
+            # Error is handled - create error ToolMessage
+            content = _handle_tool_error(e, flag=self._handle_tool_errors)
             return ToolMessage(
                 content=content,
                 name=call["name"],
@@ -620,18 +850,89 @@ class ToolNode(RunnableCallable):
                 status="error",
             )
 
+        # Process successful response
         if isinstance(response, Command):
-            return self._validate_tool_command(response, call, input_type)
+            # Validate Command before returning to handler
+            return self._validate_tool_command(response, request.tool_call, input_type)
         if isinstance(response, ToolMessage):
             response.content = cast("str | list", msg_content_output(response.content))
             return response
+
         msg = f"Tool {call['name']} returned unexpected type: {type(response)}"
         raise TypeError(msg)
 
+    async def _arun_one(
+        self,
+        call: ToolCall,
+        input_type: Literal["list", "dict", "tool_calls"],
+        config: RunnableConfig,
+        input: list[AnyMessage] | dict[str, Any] | BaseModel,
+        runtime: Any,
+    ) -> ToolMessage | Command:
+        """Execute single tool call asynchronously with awrap_tool_call wrapper if configured.
+
+        Args:
+            call: Tool call dict.
+            input_type: Input format.
+            config: Runnable configuration.
+            input: Agent state.
+            runtime: LangGraph runtime or None.
+
+        Returns:
+            ToolMessage or Command.
+        """
+        if invalid_tool_message := self._validate_tool_call(call):
+            return invalid_tool_message
+
+        tool = self.tools_by_name[call["name"]]
+
+        # Extract state from ToolCallWithContext if present
+        state = self._extract_state(input)
+
+        # Create the tool request with state and runtime
+        tool_request = ToolCallRequest(
+            tool_call=call,
+            tool=tool,
+            state=state,
+            runtime=runtime,
+        )
+
+        if self._awrap_tool_call is None and self._wrap_tool_call is None:
+            # No wrapper - execute directly
+            return await self._execute_tool_async(tool_request, input_type, config)
+
+        # Define async execute callable that can be called multiple times
+        async def execute(req: ToolCallRequest) -> ToolMessage | Command:
+            """Execute tool with given request. Can be called multiple times."""
+            return await self._execute_tool_async(req, input_type, config)
+
+        def _sync_execute(req: ToolCallRequest) -> ToolMessage | Command:
+            """Sync execute fallback for sync wrapper."""
+            return self._execute_tool_sync(req, input_type, config)
+
+        # Call wrapper with request and execute callable
+        try:
+            if self._awrap_tool_call is not None:
+                return await self._awrap_tool_call(tool_request, execute)
+            # None check was performed above already
+            self._wrap_tool_call = cast("ToolCallWrapper", self._wrap_tool_call)
+            return self._wrap_tool_call(tool_request, _sync_execute)
+        except Exception as e:
+            # Wrapper threw an exception
+            if not self._handle_tool_errors:
+                raise
+            # Convert to error message
+            content = _handle_tool_error(e, flag=self._handle_tool_errors)
+            return ToolMessage(
+                content=content,
+                name=tool_request.tool_call["name"],
+                tool_call_id=tool_request.tool_call["id"],
+                status="error",
+            )
+
     def _parse_input(
         self,
         input: list[AnyMessage] | dict[str, Any] | BaseModel,
-        store: BaseStore | None,
     ) -> tuple[list[ToolCall], Literal["list", "dict", "tool_calls"]]:
         input_type: Literal["list", "dict", "tool_calls"]
         if isinstance(input, list):
@@ -641,6 +942,14 @@ class ToolNode(RunnableCallable):
                 return tool_calls, input_type
             input_type = "list"
             messages = input
+        elif isinstance(input, dict) and input.get("__type") == "tool_call_with_context":
+            # Handle ToolCallWithContext from Send API
+            # mypy will not be able to type narrow correctly since the signature
+            # for input contains dict[str, Any]. We'd need to narrow dict[str, Any]
+            # before we can apply correct typing.
+            input_with_ctx = cast("ToolCallWithContext", input)
+            input_type = "tool_calls"
+            return [input_with_ctx["tool_call"]], input_type
         elif isinstance(input, dict) and (messages := input.get(self._messages_key, [])):
             input_type = "dict"
         elif messages := getattr(input, self._messages_key, []):
@@ -656,9 +965,7 @@ class ToolNode(RunnableCallable):
             msg = "No AIMessage found in input"
             raise ValueError(msg)
 
-        tool_calls = [
-            self.inject_tool_args(call, input, store) for call in latest_ai_message.tool_calls
-        ]
+        tool_calls = list(latest_ai_message.tool_calls)
         return tool_calls, input_type
 
     def _validate_tool_call(self, call: ToolCall) -> ToolMessage | None:
@@ -674,6 +981,21 @@ class ToolNode(RunnableCallable):
             )
         return None
 
+    def _extract_state(
+        self, input: list[AnyMessage] | dict[str, Any] | BaseModel
+    ) -> list[AnyMessage] | dict[str, Any] | BaseModel:
+        """Extract state from input, handling ToolCallWithContext if present.
+
+        Args:
+            input: The input which may be raw state or ToolCallWithContext.
+
+        Returns:
+            The actual state to pass to wrap_tool_call wrappers.
+        """
+        if isinstance(input, dict) and input.get("__type") == "tool_call_with_context":
+            return input["state"]
+        return input
+
     def _inject_state(
         self,
         tool_call: ToolCall,
@@ -696,14 +1018,20 @@ class ToolNode(RunnableCallable):
                     err_msg += f" State should contain fields {required_fields_str}."
                 raise ValueError(err_msg)
 
-        if isinstance(input, dict):
+        # Extract state from ToolCallWithContext if present
+        if isinstance(input, dict) and input.get("__type") == "tool_call_with_context":
+            state = input["state"]
+        else:
+            state = input
+
+        if isinstance(state, dict):
             tool_state_args = {
-                tool_arg: input[state_field] if state_field else input
+                tool_arg: state[state_field] if state_field else state
                 for tool_arg, state_field in state_args.items()
             }
         else:
             tool_state_args = {
-                tool_arg: getattr(input, state_field) if state_field else input
+                tool_arg: getattr(state, state_field) if state_field else state
                 for tool_arg, state_field in state_args.items()
             }
 
@@ -731,7 +1059,7 @@ class ToolNode(RunnableCallable):
         }
         return tool_call
 
-    def inject_tool_args(
+    def _inject_tool_args(
         self,
         tool_call: ToolCall,
         input: list[AnyMessage] | dict[str, Any] | BaseModel,
@@ -739,10 +1067,11 @@ class ToolNode(RunnableCallable):
     ) -> ToolCall:
         """Inject graph state and store into tool call arguments.
 
-        This method enables tools to access graph context that should not be controlled
-        by the model. Tools can declare dependencies on graph state or persistent storage
-        using InjectedState and InjectedStore annotations. This method automatically
-        identifies these dependencies and injects the appropriate values.
+        This is an internal method that enables tools to access graph context that
+        should not be controlled by the model. Tools can declare dependencies on graph
+        state or persistent storage using InjectedState and InjectedStore annotations.
+        This method automatically identifies these dependencies and injects the
+        appropriate values.
 
         The injection process preserves the original tool call structure while adding
         the necessary context arguments. This allows tools to be both model-callable
@@ -765,10 +1094,8 @@ class ToolNode(RunnableCallable):
                        or if state injection requirements cannot be satisfied.
 
         Note:
-            This method is automatically called during tool execution but can also
-            be used manually when working with the Send API or custom routing logic.
-            The injection is performed on a copy of the tool call to avoid mutating
-            the original.
+            This method is called automatically during tool execution. It should not
+            be called from outside the ToolNode.
         """
         if tool_call["name"] not in self.tools_by_name:
             return tool_call
@@ -940,7 +1267,7 @@ class InjectedState(InjectedToolArg):
     to the model's tool-calling interface.
 
     Args:
-        field: Optional key to extract from the state dictionary. If None, the entire
+        field: Optional key to extract from the state dictionary. If `None`, the entire
             state is injected. If specified, only that field's value is injected.
             This allows tools to request specific state components rather than
             processing the full state structure.
@@ -985,7 +1312,7 @@ class InjectedState(InjectedToolArg):
         node.invoke(state)
         ```
 
-        ```pycon
+        ```python
         [
             ToolMessage(content="not enough messages", name="state_tool", tool_call_id="1"),
             ToolMessage(content="bar2", name="foo_tool", tool_call_id="2"),
@@ -994,12 +1321,12 @@ class InjectedState(InjectedToolArg):
 
     Note:
         - InjectedState arguments are automatically excluded from tool schemas
-          presented to language models
+            presented to language models
         - ToolNode handles the injection process during execution
         - Tools can mix regular arguments (controlled by the model) with injected
-          arguments (controlled by the system)
+            arguments (controlled by the system)
         - State injection occurs after the model generates tool calls but before
-          tool execution
+            tool execution
     """
 
     def __init__(self, field: str | None = None) -> None:
@@ -1074,7 +1401,7 @@ class InjectedStore(InjectedToolArg):
 
     Note:
         - InjectedStore arguments are automatically excluded from tool schemas
-          presented to language models
+            presented to language models
         - The store instance is automatically injected by ToolNode during execution
         - Tools can access namespaced storage using the store's get/put methods
         - Store injection requires the graph to be compiled with a store instance