fix: memoise ToolResolver in resolve_tools to avoid per-agent rebuild#3123
Conversation
…fixes #3121) Callers invoke resolve_tools() once per agent/step over the same registry object, and each call rebuilt a fresh ToolRegistry + ToolResolver, discarding per-instance caches and re-executing local tools.py once per agent. Extract resolver construction into _get_resolver(), memoised on (id(registry), template_dir, autoload_enabled), so one resolver is reused across a workflow build. Resolution order, results, and the security gate are unchanged. Co-authored-by: MervinPraison <MervinPraison@users.noreply.github.com>
|
@coderabbitai review |
|
/review |
Qodo reviews are paused for this user.Troubleshooting steps vary by plan Learn more → On a Teams plan? Using GitHub Enterprise Server, GitLab Self-Managed, or Bitbucket Data Center? |
✅ Action performedReview finished.
|
MervinPraison
added
pipeline/blocked:ci
|
Important Review skippedBot user detected. To trigger a single review, invoke the ⚙️ Run configurationConfiguration used: defaults Review profile: CHILL Plan: Pro Run ID: You can disable this status message by setting the Use the checkbox below for a quick retry:
📝 WalkthroughWalkthrough
ChangesResolver reuse
Estimated code review effort: 3 (Moderate) | ~20 minutes Possibly related issues
Suggested labels: Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
Greptile SummaryThis PR reuses tool resolvers across repeated workflow calls. The main changes are:
Confidence Score: 5/5This looks safe to merge.
Important Files Changed
Reviews (2): Last reviewed commit: "fix: bound resolver memo cache and skip ..." | Re-trigger Greptile |
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
| cache_key = (id(registry), template_dir, autoload) | ||
|
|
||
| cached = _resolver_cache.get(cache_key) | ||
| # Guard against id() reuse after GC: verify the stored registry is the same | ||
| # object we were passed before returning the cached resolver. | ||
| if cached is not None and cached[0] is registry: |
There was a problem hiding this comment.
When the same registry dictionary is changed between calls, this identity-only key returns a resolver containing the original snapshot. Added tools remain unavailable, while removed or replaced tools can still resolve to old callables; callers that update a shared registry therefore receive stale results.
|
@claude You are the FINAL architecture reviewer. If the branch is under MervinPraison/PraisonAI (not a fork), you are able to make modifications to this branch and push directly. SCOPE: Focus ONLY on Python packages (praisonaiagents, praisonai). Do NOT modify praisonai-rust or praisonai-ts. Read ALL comments above from Gemini, Qodo, CodeRabbit, and Copilot carefully before responding. Phase 1: Review per AGENTS.md
Phase 2: FIX Valid Issues Phase 3: Final Verdict |
MervinPraison
added
pipeline/awaiting-merge-gate
|
Claude finished @MervinPraison's task —— View job Final Architecture Review — PR #3123 ✅ Approved (fixes pushed)
Phase 1 — Architecture (per AGENTS.md)
Phase 2 — Valid reviewer issues fixedEvaluated Greptile's three P2 findings:
File modified: Phase 3 — Validation
Verdict: ✅ Approve. The optimization is sound and the two valid memory/cache-lifecycle issues raised by Greptile are now fixed. Fixes pushed to
|
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 4
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@src/praisonai/praisonai/templates/tool_override.py`:
- Around line 490-491: Update the cache flow around _autoload_tools_enabled and
ToolResolver construction to resolve the effective tools.py path, including
current-working-directory resolution for relative template_dir values, before
creating cache_key or constructing the resolver. Key the cache by that resolved
path so cwd changes and equivalent relative paths cannot reuse a registry for
the wrong project.
- Around line 493-529: Protect the entire cache lookup, resolver construction,
and `_resolver_cache` store sequence with a module-level lock, including the
`ToolResolver` creation and optional `tools.py` loading. Keep the existing
registry identity check and cache key behavior, and add a concurrent-call test
confirming callers receive the same resolver instance and opted-in template
tools are loaded only once.
- Around line 18-24: The _resolver_cache and resolve_tools flow need bounded
lifecycle management and stable default-registry memoization. Use a stable
sentinel/key for registry=None so the default registry is created only on a
cache miss, and bound or explicitly clean up cache entries; apply the same
changes to the related cache accesses and registry creation at the referenced
resolve_tools paths.
- Around line 493-497: Update the resolver cache logic surrounding
_resolver_cache and its lookup so cached entries also reflect the current
registry contents, not only registry identity. Add a registry-content version or
fingerprint to cache validation/key generation, ensuring additions,
replacements, and removals trigger resolver rebuilding while unchanged
registries still reuse cached resolvers; preserve existing APIs and add coverage
for mutations between calls.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: 42545730-719a-4202-b9cf-f7e5a375d455
📒 Files selected for processing (1)
src/praisonai/praisonai/templates/tool_override.py
| # Memoisation cache for ToolResolver instances built by resolve_tools(). | ||
| # Keyed on (id(registry), template_dir, autoload_enabled) so that repeated | ||
| # calls within a single workflow build (once per agent/step, over the same | ||
| # registry object) reuse one resolver and its per-instance caches instead of | ||
| # rebuilding the registry+resolver and re-executing local tools.py per agent. | ||
| # The registry object is also stored to guard against id() reuse after GC. | ||
| _resolver_cache: Dict[Any, Any] = {} |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
Bound the cache and reuse the default-registry entry.
The cache permanently retains every registry and resolver. Worse, registry=None creates a new dictionary before keying, so every such call misses and leaks another entry rather than being memoized.
Use a bounded cache or explicit lifecycle cleanup, and use a stable sentinel/key so the default registry is only created on a cache miss.
Also applies to: 486-491, 529-530
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@src/praisonai/praisonai/templates/tool_override.py` around lines 18 - 24, The
_resolver_cache and resolve_tools flow need bounded lifecycle management and
stable default-registry memoization. Use a stable sentinel/key for registry=None
so the default registry is created only on a cache miss, and bound or explicitly
clean up cache entries; apply the same changes to the related cache accesses and
registry creation at the referenced resolve_tools paths.
| autoload = _autoload_tools_enabled() | ||
| cache_key = (id(registry), template_dir, autoload) |
There was a problem hiding this comment.
🔒 Security & Privacy | 🟠 Major | ⚡ Quick win
Key the cache by the resolver’s effective resolved tools path.
ToolResolver resolves its path against the current working directory during construction, but this key stores raw template_dir. Reusing the registry after a cwd change can therefore resolve tools from the previous project; relative template paths have the same problem.
Resolve the effective tools.py path before both keying and construction.
Proposed direction
- cache_key = (id(registry), template_dir, autoload)
+ tools_path = None
+ if autoload:
+ base = Path(template_dir) if template_dir else Path.cwd()
+ tools_path = str((base / "tools.py").resolve())
+ cache_key = (id(registry), tools_path, autoload)
- template_tools_path = None
- if template_dir and autoload:
- template_tools_path = str(Path(template_dir) / "tools.py")
+ template_tools_path = tools_pathAlso applies to: 517-526
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@src/praisonai/praisonai/templates/tool_override.py` around lines 490 - 491,
Update the cache flow around _autoload_tools_enabled and ToolResolver
construction to resolve the effective tools.py path, including
current-working-directory resolution for relative template_dir values, before
creating cache_key or constructing the resolver. Key the cache by that resolved
path so cwd changes and equivalent relative paths cannot reuse a registry for
the wrong project.
| cached = _resolver_cache.get(cache_key) | ||
| # Guard against id() reuse after GC: verify the stored registry is the same | ||
| # object we were passed before returning the cached resolver. | ||
| if cached is not None and cached[0] is registry: | ||
| return cached[1] |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | 🏗️ Heavy lift
Invalidate cached resolvers when the registry dictionary changes.
Identity-only lookup returns a resolver containing the original copied entries after callers add, replace, or remove tools from the same registry dictionary. Previously, rebuilding observed those mutations.
Include a registry-content version/fingerprint in the key or provide mutation-aware invalidation, with tests covering updates between calls.
As per coding guidelines, “Ensure backward compatibility with existing APIs.”
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@src/praisonai/praisonai/templates/tool_override.py` around lines 493 - 497,
Update the resolver cache logic surrounding _resolver_cache and its lookup so
cached entries also reflect the current registry contents, not only registry
identity. Add a registry-content version or fingerprint to cache validation/key
generation, ensuring additions, replacements, and removals trigger resolver
rebuilding while unchanged registries still reuse cached resolvers; preserve
existing APIs and add coverage for mutations between calls.
Source: Coding guidelines
| cached = _resolver_cache.get(cache_key) | ||
| # Guard against id() reuse after GC: verify the stored registry is the same | ||
| # object we were passed before returning the cached resolver. | ||
| if cached is not None and cached[0] is registry: | ||
| return cached[1] | ||
|
|
||
| # Create a ToolRegistry instance for high-priority overrides | ||
| tool_registry = ToolRegistry() | ||
|
|
||
| # Don't manually load template tools here; let ToolResolver handle it to | ||
| # avoid double-execution. Populate tool_registry ONLY with registry | ||
| # overrides, which have HIGHER priority than template tools in ToolResolver. | ||
| if registry: | ||
| # Filter out lazy-loaded tuples from registry and register callables | ||
| for name, tool in registry.items(): | ||
| if not isinstance(tool, tuple): | ||
| if callable(tool): | ||
| tool_registry.register_function(name, tool) | ||
| elif hasattr(tool, "run") and callable(getattr(tool, "run", None)): | ||
| # Support non-callable tools with a run method | ||
| def make_callable(t): | ||
| return lambda *args, **kwargs: t.run(*args, **kwargs) | ||
| tool_registry.register_function(name, make_callable(tool)) | ||
|
|
||
| # Only pass template tools path if autoload is enabled (security gate) | ||
| template_tools_path = None | ||
| if template_dir and autoload: | ||
| template_tools_path = str(Path(template_dir) / "tools.py") | ||
|
|
||
| # Create ToolResolver with the registry having highest priority. | ||
| # Template tools will only be loaded if autoload is enabled. | ||
| resolver = ToolResolver( | ||
| tools_py_path=template_tools_path, | ||
| registry=tool_registry | ||
| ) | ||
|
|
||
| _resolver_cache[cache_key] = (registry, resolver) |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
Synchronize cache lookup and resolver construction.
Concurrent calls can both miss, construct separate resolvers, and return different instances. This defeats memoization and may execute an opted-in tools.py more than once.
Guard the complete lookup/build/store sequence with a module-level lock and add a concurrent-call test.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@src/praisonai/praisonai/templates/tool_override.py` around lines 493 - 529,
Protect the entire cache lookup, resolver construction, and `_resolver_cache`
store sequence with a module-level lock, including the `ToolResolver` creation
and optional `tools.py` loading. Keep the existing registry identity check and
cache key behavior, and add a concurrent-call test confirming callers receive
the same resolver instance and opted-in template tools are loaded only once.
Addresses reviewer feedback on the ToolResolver memoisation: - Cap the resolver cache as a small LRU (OrderedDict) so a long-running process that builds many workflows no longer retains every registry, its tool callables/modules and resolver caches for the process life. - Skip caching entirely when registry is None: a fresh registry is built per call so its id() is never stable, meaning caching it only leaked memory and never produced a hit. Co-authored-by: Mervin Praison <MervinPraison@users.noreply.github.com>
MervinPraison
added
pipeline/blocked:stale-final
praisonai-triage-agent
Bot
removed
the
pipeline/blocked:cooldown
2 participants
Fixes #3121
Summary
resolve_tools()rebuilt a freshToolRegistryandToolResolveron every call. Callers loop over agents/steps passing the same registry object, so for a workflow with N agents the resolver was built N times — discarding its per-instance caches (_resolve_cache,_local_tools_cache) and re-executing the user's localtools.pyonce per agent.Change
_get_resolver()helper, memoised on(id(registry), template_dir, autoload_enabled).isthe same object before returning the cache (protects againstid()reuse after GC).resolve_tools()now reuses one resolver across the whole workflow build.Wrapper-only change (
praisonai/templates/tool_override.py); no caller signatures changed, canonicalpraisonai_code.ToolResolveruntouched.Capability preserved
--tools/--tools-dir,TEMPLATE.yamltools_sources, template-localtools.py, and thePRAISONAI_ALLOW_TEMPLATE_TOOLSsecurity gate all intact (autoload state is part of the cache key).Validation
tests/unit/test_templates.py+test_tool_override_autoload_gate.py: 84 passed.Generated with Claude Code
Summary by CodeRabbit
Performance Improvements
Compatibility
run()method.