From cb2faa7dda0f2ea49d2729361e1802fb233feaa5 Mon Sep 17 00:00:00 2001
From: Sidharth Rajmohan <138385787+dumko2001@users.noreply.github.com>
Date: Tue, 13 Jan 2026 00:13:45 +0530
Subject: [PATCH] fix(prebuilt): support generic type arguments for ToolRuntime
 injection (#6509)

**Description:**
This PR fixes an issue where injection types (like `ToolRuntime`) were
not recognized by `ToolNode` when used with generic type arguments
(e.g., `ToolRuntime[MyContext]`).

Previously, the `_is_injection` check relied solely on `isinstance` and
`issubclass`, which fail for `typing._GenericAlias` objects. This update
adds a check using `typing.get_origin()` to correctly identify the base
class of generic types, ensuring the runtime is injected correctly even
when type hints are present.

**Issue:** Fixes #6465

**Dependencies:** None

**Twitter handle:** @SidharthRajmoh2

---------

Co-authored-by: Sydney Runkle <54324534+sydney-runkle@users.noreply.github.com>
---
 libs/prebuilt/langgraph/prebuilt/tool_node.py |  6 +++
 libs/prebuilt/tests/test_tool_node.py         | 42 +++++++++++++++++++
 2 files changed, 48 insertions(+)

diff --git a/libs/prebuilt/langgraph/prebuilt/tool_node.py b/libs/prebuilt/langgraph/prebuilt/tool_node.py
index aee8f8b02..88d7c62d6 100644
--- a/libs/prebuilt/langgraph/prebuilt/tool_node.py
+++ b/libs/prebuilt/langgraph/prebuilt/tool_node.py
@@ -1761,6 +1761,12 @@ def _is_injection(
     origin_ = get_origin(type_arg)
     if origin_ is Union or origin_ is Annotated:
         return any(_is_injection(ta, injection_type) for ta in get_args(type_arg))
+
+    if origin_ is not None and (
+        origin_ is injection_type
+        or (isinstance(origin_, type) and issubclass(origin_, injection_type))
+    ):
+        return True
     return False
 
 
diff --git a/libs/prebuilt/tests/test_tool_node.py b/libs/prebuilt/tests/test_tool_node.py
index 00845d398..d1d83382a 100644
--- a/libs/prebuilt/tests/test_tool_node.py
+++ b/libs/prebuilt/tests/test_tool_node.py
@@ -1860,3 +1860,45 @@ async def test_tool_node_inject_async_all_types_with_schema() -> None:
         "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"