diff --git a/libs/prebuilt/langgraph/prebuilt/tool_node.py b/libs/prebuilt/langgraph/prebuilt/tool_node.py index 31b3f6b98..00385f847 100644 --- a/libs/prebuilt/langgraph/prebuilt/tool_node.py +++ b/libs/prebuilt/langgraph/prebuilt/tool_node.py @@ -922,7 +922,7 @@ class ToolNode(RunnableCallable): raise TypeError(msg) # Inject state, store, and runtime right before invocation - injected_call = self._inject_tool_args(call, request.runtime) + injected_call = self._inject_tool_args(call, request.runtime, tool) call_args = {**injected_call, "type": "tool_call"} try: @@ -1075,7 +1075,7 @@ class ToolNode(RunnableCallable): raise TypeError(msg) # Inject state, store, and runtime right before invocation - injected_call = self._inject_tool_args(call, request.runtime) + injected_call = self._inject_tool_args(call, request.runtime, tool) call_args = {**injected_call, "type": "tool_call"} try: @@ -1281,6 +1281,7 @@ class ToolNode(RunnableCallable): self, tool_call: ToolCall, tool_runtime: ToolRuntime, + tool: BaseTool | None = None, ) -> ToolCall: """Inject graph state, store, and runtime into tool call arguments. @@ -1299,6 +1300,9 @@ class ToolNode(RunnableCallable): Must contain 'name', 'args', 'id', and 'type' fields. tool_runtime: The ToolRuntime instance containing all runtime context (state, config, store, context, stream_writer) to inject into tools. + tool: Optional tool instance. When provided, allows injection for + dynamically registered tools that are not in self.tools_by_name + (e.g., tools added via middleware's wrap_tool_call). Returns: A new ToolCall dictionary with the same structure as the input but with @@ -1312,10 +1316,12 @@ class ToolNode(RunnableCallable): 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 - injected = self._injected_args.get(tool_call["name"]) + if not injected and tool is not None: + # For dynamically registered tools (e.g., added via middleware's + # wrap_tool_call), compute injected args on-the-fly since they + # were not present during ToolNode initialization. + injected = _get_all_injected_args(tool) if not injected: return tool_call diff --git a/libs/prebuilt/tests/test_tool_node.py b/libs/prebuilt/tests/test_tool_node.py index d1d83382a..506572ff3 100644 --- a/libs/prebuilt/tests/test_tool_node.py +++ b/libs/prebuilt/tests/test_tool_node.py @@ -1902,3 +1902,109 @@ async def test_tool_node_tool_runtime_generic() -> None: assert tool_message.type == "tool" assert tool_message.content == "test_info" assert tool_message.tool_call_id == "call_1" + + +def test_tool_node_inject_runtime_dynamic_tool_via_wrap_tool_call() -> None: + """Test that ToolRuntime is injected for dynamically registered tools. + + Regression test for https://github.com/langchain-ai/langchain/issues/35305. + When a tool is dynamically provided via wrap_tool_call (not registered at + ToolNode init time), ToolRuntime should still be injected into the tool. + """ + + @dec_tool + def static_tool(x: int) -> str: + """A static tool registered at init.""" + return f"static: {x}" + + @dec_tool + def dynamic_tool_with_runtime(x: int, runtime: ToolRuntime) -> str: + """A dynamic tool that needs ToolRuntime injection.""" + return f"dynamic: x={x}, tool_call_id={runtime.tool_call_id}" + + def wrap_tool_call(request, execute): + """Middleware that swaps in a dynamic tool.""" + if request.tool_call["name"] == "dynamic_tool_with_runtime": + # Override tool to the dynamic one (not registered at init) + new_request = request.override(tool=dynamic_tool_with_runtime) + return execute(new_request) + return execute(request) + + # ToolNode only knows about static_tool at init time + tool_node = ToolNode( + [static_tool], + wrap_tool_call=wrap_tool_call, + ) + + # Verify the dynamic tool is NOT in the tool node's registered tools + assert "dynamic_tool_with_runtime" not in tool_node.tools_by_name + + # Call the dynamic tool + tool_call = { + "name": "dynamic_tool_with_runtime", + "args": {"x": 42}, + "id": "call_dynamic_1", + "type": "tool_call", + } + msg = AIMessage("", tool_calls=[tool_call]) + result = tool_node.invoke( + {"messages": [msg]}, + config=_create_config_with_runtime(), + ) + + # ToolRuntime should be injected and the tool should execute successfully + tool_message = result["messages"][-1] + assert tool_message.content == "dynamic: x=42, tool_call_id=call_dynamic_1" + assert tool_message.tool_call_id == "call_dynamic_1" + + +async def test_tool_node_inject_runtime_dynamic_tool_via_wrap_tool_call_async() -> None: + """Test that ToolRuntime is injected for dynamically registered tools (async). + + Async version of the regression test for + https://github.com/langchain-ai/langchain/issues/35305. + """ + + @dec_tool + def static_tool(x: int) -> str: + """A static tool registered at init.""" + return f"static: {x}" + + @dec_tool + async def dynamic_tool_with_runtime(x: int, runtime: ToolRuntime) -> str: + """A dynamic async tool that needs ToolRuntime injection.""" + return f"dynamic: x={x}, tool_call_id={runtime.tool_call_id}" + + async def awrap_tool_call(request, execute): + """Async middleware that swaps in a dynamic tool.""" + if request.tool_call["name"] == "dynamic_tool_with_runtime": + new_request = request.override(tool=dynamic_tool_with_runtime) + return await execute(new_request) + return await execute(request) + + # ToolNode only knows about static_tool at init time + tool_node = ToolNode( + [static_tool], + awrap_tool_call=awrap_tool_call, + ) + + # Verify the dynamic tool is NOT in the tool node's registered tools + assert "dynamic_tool_with_runtime" not in tool_node.tools_by_name + + # Call the dynamic tool + tool_call = { + "name": "dynamic_tool_with_runtime", + "args": {"x": 42}, + "id": "call_dynamic_2", + "type": "tool_call", + } + msg = AIMessage("", tool_calls=[tool_call]) + result = await tool_node.ainvoke( + {"messages": [msg]}, + config=_create_config_with_runtime(), + ) + + # ToolRuntime should be injected and the tool should execute successfully + tool_message = result["messages"][-1] + assert tool_message.content == "dynamic: x=42, tool_call_id=call_dynamic_2" + assert tool_message.tool_call_id == "call_dynamic_2"