fix: inject ToolRuntime for dynamically registered tools (#6874)

Fixes https://github.com/langchain-ai/langchain/issues/35305

Co-authored-by: Shivangi Sharma <shivangi.sharma7004@gmail.com>

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
 

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"