better typing

libs/prebuilt/langgraph/prebuilt/tool_node.py

@@ -124,7 +124,7 @@ class _ToolCallRequestOverrides(TypedDict, total=False):
 
 
 @dataclass
-class ToolCallRequest:
+class ToolCallRequest(Generic[ContextT, StateT]):
     """Tool execution request passed to tool call interceptors.
 
     Attributes:
@@ -133,14 +133,22 @@ class ToolCallRequest:
             registered with the `ToolNode`. When tool is `None`, interceptors can
             handle the request without validation. If the interceptor calls `execute()`,
             validation will occur and raise an error for unregistered tools.
-        state: Agent state (`dict`, `list`, or `BaseModel`).
         runtime: LangGraph runtime context (optional, `None` if outside graph).
+        state: Agent state (`dict`, `list`, or `BaseModel`). Pulled from `runtime.state`.
     """
 
     tool_call: ToolCall
     tool: BaseTool | None
-    state: Any
-    runtime: ToolRuntime
+    runtime: ToolRuntime[ContextT, StateT]
+
+    @property
+    def state(self) -> StateT:
+        """Get the state from the runtime.
+
+        Returns:
+            The current graph state from the runtime context.
+        """
+        return self.runtime.state
 
     def __setattr__(self, name: str, value: Any) -> None:
         """Raise deprecation warning when setting attributes directly.
@@ -163,7 +171,7 @@ class ToolCallRequest:
 
     def override(
         self, **overrides: Unpack[_ToolCallRequestOverrides]
-    ) -> ToolCallRequest:
+    ) -> ToolCallRequest[ContextT, StateT]:
         """Replace the request with a new request with the given overrides.
 
         Returns a new `ToolCallRequest` instance with the specified attributes replaced.
@@ -947,11 +955,10 @@ class ToolNode(RunnableCallable):
         # to short-circuit requests for unregistered tools
         tool = self.tools_by_name.get(call["name"])
 
-        # Create the tool request with state and runtime
+        # Create the tool request with runtime
         tool_request = ToolCallRequest(
             tool_call=call,
             tool=tool,
-            state=tool_runtime.state,
             runtime=tool_runtime,
         )
 
@@ -1104,11 +1111,10 @@ class ToolNode(RunnableCallable):
         # to short-circuit requests for unregistered tools
         tool = self.tools_by_name.get(call["name"])
 
-        # Create the tool request with state and runtime
+        # Create the tool request with runtime
         tool_request = ToolCallRequest(
             tool_call=call,
             tool=tool,
-            state=tool_runtime.state,
             runtime=tool_runtime,
         )
 

libs/prebuilt/tests/test_on_tool_call.py

@@ -13,6 +13,7 @@ from langgraph.types import Command
 from langgraph.prebuilt.tool_node import (
     ToolCallRequest,
     ToolNode,
+    ToolRuntime,
 )
 
 pytestmark = pytest.mark.anyio
@@ -342,16 +343,23 @@ def test_tool_call_request_dataclass() -> None:
     """Test ToolCallRequest dataclass."""
     tool_call: ToolCall = {"name": "add", "args": {"a": 1, "b": 2}, "id": "call_1"}
     state: dict = {"messages": []}
-    runtime = None
+    tool_runtime = ToolRuntime(
+        state=state,
+        config={},
+        context=None,
+        store=None,
+        stream_writer=Mock(),
+        tool_call_id="call_1",
+    )
 
     request = ToolCallRequest(
-        tool_call=tool_call, tool=add, state=state, runtime=runtime
-    )  # type: ignore[arg-type]
+        tool_call=tool_call, tool=add, runtime=tool_runtime
+    )
 
     assert request.tool_call == tool_call
     assert request.tool == add
     assert request.state == state
-    assert request.runtime is None
+    assert request.runtime is tool_runtime
     assert request.tool_call["name"] == "add"
 
 
@@ -1324,11 +1332,18 @@ def test_tool_call_request_is_frozen() -> None:
     """Test that ToolCallRequest raises deprecation warnings on direct attribute reassignment."""
     tool_call: ToolCall = {"name": "add", "args": {"a": 1, "b": 2}, "id": "call_1"}
     state: dict = {"messages": []}
-    runtime = None
+    tool_runtime = ToolRuntime(
+        state=state,
+        config={},
+        context=None,
+        store=None,
+        stream_writer=Mock(),
+        tool_call_id="call_1",
+    )
 
     request = ToolCallRequest(
-        tool_call=tool_call, tool=add, state=state, runtime=runtime
-    )  # type: ignore[arg-type]
+        tool_call=tool_call, tool=add, runtime=tool_runtime
+    )
 
     # Test that direct attribute reassignment raises DeprecationWarning
     with pytest.warns(
@@ -1343,11 +1358,14 @@ def test_tool_call_request_is_frozen() -> None:
     ):
         request.tool = None  # type: ignore[misc]
 
+    # state is now a property, so setting it will raise a deprecation warning
+    # (and then fail with AttributeError after the warning)
     with pytest.warns(
         DeprecationWarning,
         match="Setting attribute 'state' on ToolCallRequest is deprecated",
     ):
-        request.state = {}  # type: ignore[misc]
+        with pytest.raises(AttributeError):
+            request.state = {}  # type: ignore[misc]
 
     with pytest.warns(
         DeprecationWarning,
@@ -1365,8 +1383,8 @@ def test_tool_call_request_is_frozen() -> None:
     # Original request should be unchanged (note: it was modified by the warnings tests above)
     # So we create a fresh request to test override properly
     fresh_request = ToolCallRequest(
-        tool_call=tool_call, tool=add, state=state, runtime=runtime
-    )  # type: ignore[arg-type]
+        tool_call=tool_call, tool=add, runtime=tool_runtime
+    )
     fresh_new_request = fresh_request.override(tool_call=new_tool_call)
 
     # Original request should be unchanged
@@ -1378,4 +1396,4 @@ def test_tool_call_request_is_frozen() -> None:
     assert fresh_new_request.tool_call["name"] == "multiply"
     assert fresh_new_request.tool == add  # Other fields should remain the same
     assert fresh_new_request.state == state
-    assert fresh_new_request.runtime is None
+    assert fresh_new_request.runtime is tool_runtime

libs/prebuilt/tests/test_tool_node.py

@@ -1615,18 +1615,28 @@ def test_tool_node_stream_writer() -> None:
 def test_tool_call_request_setattr_deprecation_warning():
     """Test that ToolCallRequest raises a deprecation warning on direct attribute modification."""
     import warnings
+    from unittest.mock import Mock
 
-    from langgraph.prebuilt.tool_node import ToolCallRequest
+    from langgraph.prebuilt.tool_node import ToolCallRequest, ToolRuntime
 
     # Create a mock ToolCall
     tool_call = {"name": "test", "args": {"a": 1}, "id": "call_1", "type": "tool_call"}
 
+    # Create a ToolRuntime
+    tool_runtime = ToolRuntime(
+        state={"messages": []},
+        config={},
+        context=None,
+        store=None,
+        stream_writer=Mock(),
+        tool_call_id="call_1",
+    )
+
     # Create a ToolCallRequest
     request = ToolCallRequest(
         tool_call=tool_call,
         tool=None,
-        state={"messages": []},
-        runtime=None,
+        runtime=tool_runtime,
     )
 
     # Test 1: Direct attribute assignment should raise deprecation warning but still work
@@ -1667,8 +1677,7 @@ def test_tool_call_request_setattr_deprecation_warning():
         ToolCallRequest(
             tool_call=tool_call,
             tool=None,
-            state={"messages": []},
-            runtime=None,
+            runtime=tool_runtime,
         )
         # Verify no warning was raised during initialization
         assert len(w) == 0