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>
This commit is contained in:
Sidharth Rajmohan
2026-01-12 18:43:45 +00:00
committed by GitHub
co-authored by Sydney Runkle
parent a5827c5c61
commit cb2faa7dda
2 changed files with 48 additions and 0 deletions
+42
View File
@@ -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"