perf(prebuilt): cache _get_all_injected_args by tool object identity

`_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 <noreply@anthropic.com>
This commit is contained in:
Sydney Runkle
2026-04-22 16:13:24 -04:00
co-authored by Claude Sonnet 4.6
parent 8cd53c408d
commit aa992adf57
+17 -1
View File
@@ -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