From aa992adf57ec1249d46d4023daad9e20fae3bafb Mon Sep 17 00:00:00 2001 From: Sydney Runkle Date: Tue, 21 Apr 2026 12:19:44 -0400 Subject: [PATCH] perf(prebuilt): cache _get_all_injected_args by tool object identity MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `_get_all_injected_args` runs `get_input_schema()`, `get_type_hints`, and `inspect.signature` on every tool for every `ToolNode` construction. When the same tool objects are passed to repeated `create_agent` calls (the common case), this work was repeated from scratch each time. Add a module-level cache keyed by `id(tool)`. The cache entry stores a strong reference to the tool alongside the result, which prevents GC from reusing the id for a different object and making an incorrect cache hit. The `entry[0] is tool` identity check provides a second safety guard. `BaseTool` (Pydantic v2 BaseModel) is not hashable, so `@lru_cache` cannot be used directly — hence the id-based dict approach. Co-Authored-By: Claude Sonnet 4.6 --- libs/prebuilt/langgraph/prebuilt/tool_node.py | 18 +++++++++++++++++- 1 file changed, 17 insertions(+), 1 deletion(-) diff --git a/libs/prebuilt/langgraph/prebuilt/tool_node.py b/libs/prebuilt/langgraph/prebuilt/tool_node.py index d8a1f3182..35611d318 100644 --- a/libs/prebuilt/langgraph/prebuilt/tool_node.py +++ b/libs/prebuilt/langgraph/prebuilt/tool_node.py @@ -40,6 +40,7 @@ Typical Usage: from __future__ import annotations import asyncio +import functools import inspect import json from collections.abc import Awaitable, Callable @@ -1842,9 +1843,17 @@ def _get_injection_from_type( return None +# Cache keyed by tool object identity. Stores (tool, result) to keep a strong +# reference that prevents GC from reusing the id for a different object. +_INJECTED_ARGS_CACHE: dict[int, tuple[BaseTool, "_InjectedArgs"]] = {} + + def _get_all_injected_args(tool: BaseTool) -> _InjectedArgs: """Extract all injected arguments from tool in a single pass. + Results are cached by tool identity so the expensive type-hint and schema + inspection only runs once per unique tool object across ToolNode instances. + This function analyzes both the tool's input schema and function signature to identify all arguments that should be injected (state, store, runtime). @@ -1854,6 +1863,11 @@ def _get_all_injected_args(tool: BaseTool) -> _InjectedArgs: Returns: _InjectedArgs structure containing all detected injections. """ + tool_id = id(tool) + entry = _INJECTED_ARGS_CACHE.get(tool_id) + if entry is not None and entry[0] is tool: + return entry[1] + # Get annotations from both schema and function signature full_schema = tool.get_input_schema() schema_annotations = get_all_basemodel_annotations(full_schema) @@ -1899,10 +1913,12 @@ def _get_all_injected_args(tool: BaseTool) -> _InjectedArgs: if _get_injection_from_type(type_, ToolRuntime): runtime_arg = name - return _InjectedArgs( + result = _InjectedArgs( state=state_args, store=store_arg, runtime=runtime_arg, all_injected_keys=all_injected_keys, _optional_state_args=_optional_state_args, ) + _INJECTED_ARGS_CACHE[tool_id] = (tool, result) + return result