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"