Fix ToolManager instance retention

[Nandith-0777]

Pre-PR Checklist

• This PR is a bug fix, not a new feature or enhancement.

Summary

ToolManager.instances held a strong reference to every manager ever created (one per LLMAgent), so none could be garbage-collected. This PR makes the registry weak and keeps add_tool_to_all() safe under concurrent construction, which a bare WeakSet does not.

Bug / Issue

Fixes #336.

Measured on main: after building and discarding 20 models × 25 agents, 500 managers were still registered (agents and models themselves were freed). Memory grows without bound across long runs and batch sweeps, and every add_tool_to_all() call also iterates all dead managers.

Implementation

• instances is now a weakref.WeakSet, so managers are released as soon as the application drops them.
• Instance registration moved to the end of __init__, after self.tools is fully populated, so a concurrent add_tool_to_all() snapshot can never observe a half-constructed manager.
• add_tool_to_all() takes a snapshot of the registry and retries only if that snapshot is interrupted by a concurrent .add() — no lock anywhere, no overhead added to __init__.

Why not just swap in a bare WeakSet (my own first commit here, and the approach in #338): iterating a WeakSet while another thread adds to it raises RuntimeError: Set changed size during iteration. I also found something one level worse — registering the instance at the start of __init__ (before self.tools is set) means a concurrent snapshot can capture a half-built manager, so add_tool_to_all() crashes with AttributeError: 'ToolManager' object has no attribute 'tools' instead. Moving registration to the end of __init__ fixes both at once: no half-constructed instance is ever visible, and the retry loop handles the ordinary "new manager added mid-snapshot" case without a lock.

Why no lock: add_tool_to_all() can be called frequently and ToolManager is constructed per-agent, so a lock there adds contention to a hot path. The retry-on-interrupted-snapshot approach gets the same correctness without taking a lock on every construction or every broadcast.

Alternatives considered:

• A threading.Lock guarding both add() and the broadcast snapshot (the approach in Hold ToolManager instances weakly in the class registry #338): correct, but adds synchronization overhead to every ToolManager() construction, which happens once per agent.
• List of weakref.ref with removal callbacks: callbacks fire during GC, possibly mid-iteration, and mutating a list during iteration silently skips entries.
• Removing the registry: a breaking API change, out of scope for a bug fix.

Testing

New regression tests in tests/test_tools/test_tool_manager.py:

Test	main	bare WeakSet (my 1st commit)	this PR
test_tool_manager_can_be_garbage_collected	❌	✅	✅
test_add_tool_to_all_skips_collected_managers	❌	✅	✅
test_add_tool_to_all_tolerates_manager_created_mid_broadcast	✅	❌ intermittent AttributeError/RuntimeError (~1 in 12 runs under a tightened thread-switch interval)	✅

• Leak measurement (20 × 25 agents discarded): main 500 → 500 retained; this PR 500 → 0 retained.
• Full suite: 1591 passed, 15 failed. The 15 failures are all in test_async_generator_cleanup.py, pre-existing and timing-sensitive — they fail identically on unmodified main, confirmed unrelated to this change.
• pre-commit run --all-files passes (ruff check, ruff format).

Additional Notes

• instances keeps the operations used in the codebase (len, in, iteration, .clear()). .append() and indexing are no longer available; nothing in the repo uses them.
• A WeakSet is unordered. Broadcast order has no effect because each register() call is independent.
• A manager constructed during a broadcast is not guaranteed to receive that broadcast's tool (point-in-time semantics) — same as Hold ToolManager instances weakly in the class registry #338.
• Relation to Hold ToolManager instances weakly in the class registry #338: Hold ToolManager instances weakly in the class registry #338 was opened against my first commit here (bare WeakSet, before I'd found and fixed the race). It correctly identifies that a bare WeakSet isn't safe under concurrent construction, but that's now fixed on this branch via snapshot-and-retry rather than a lock, so the two PRs currently overlap almost entirely. Happy to consolidate in whatever way maintainers prefer — main difference is lock vs. lock-free retry for add_tool_to_all().

Summary by CodeRabbit

• Bug Fixes

  • Improved tool manager lifecycle handling so unused instances can be garbage-collected.
  • Prevented concurrent tool registration from interacting with partially initialized managers.
  • Improved reliability when adding tools across multiple managers during concurrent activity.

• Tests

  • Added coverage verifying that unused tool managers are properly released from memory.

[coderabbitai]

Navigate logical layers of code changes, visualize relationships, and explore their blast radius.

Important

Review skipped

Auto reviews are disabled on this repository. Please check the settings in the CodeRabbit UI or the .coderabbit.yaml file in this repository. To trigger a single review, invoke the @coderabbitai review command.

⚙️ Run configuration

Configuration used: Organization UI

Review profile: CHILL

Plan: Advanced

Run ID: 5748a128-1dcf-4de1-ad1b-eaaf6d9fb7b4

You can disable this status message by setting the reviews.review_status to false in the CodeRabbit configuration file.

Use the checkbox below for a quick retry:

• 🔍 Trigger review

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Organization UI

Review profile: CHILL

Plan: Advanced

Run ID: 7c4d098d-abb4-42da-b1a0-fcfa9ab4dcf6

📥 Commits

Reviewing files that changed from the base of the PR and between d358ebd and c237a3e.

📒 Files selected for processing (2)

• mesa_llm/tools/tool_manager.py
• tests/test_tools/test_tool_manager.py

🚧 Files skipped from review as they are similar to previous changes (1)

• tests/test_tools/test_tool_manager.py

Included review availability: Your plan provides up to 1 included review per hour; 0 remain after this review.

---

Walkthrough

ToolManager.instances now uses weakref.WeakSet. Managers register after construction completes. Bulk tool updates use retryable snapshots. A regression test verifies garbage collection after all strong references are removed.

Changes

ToolManager Weak Instance Registry

Layer / File(s)	Summary
Weak registry and snapshot updates mesa_llm/tools/tool_manager.py	ToolManager.instances uses weakref.WeakSet. The constructor registers each manager after setup. add_tool_to_all copies the registry and retries if concurrent mutation raises RuntimeError.
Garbage collection validation tests/test_tools/test_tool_manager.py	The test suite verifies that an unreferenced ToolManager is collected after gc.collect() runs. The fixture comment now refers to the registry.

Priority: ➖ Normal

Estimated code review effort: 2 (Simple) | ~10 minutes

Change: Bug fix · Severity of issue fixed: Medium

Merge Risk: ⚪ Minimal · up to c237a

Unused ToolManager instances can be collected, and bulk updates use construction-safe snapshots. The change is ready to merge.

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly and concisely describes the main change: preventing ToolManager instance retention.
Description check	✅ Passed	The description identifies this as a bug fix, references issue #336, explains the cause and implementation, documents alternatives, and includes testing results. It satisfies the repository template r…
Linked Issues check	✅ Passed	Issue #336 requires garbage collection of unreferenced ToolManager instances and a regression test. The PR changes ToolManager.instances to weakref.WeakSet, registers each instance after initial…
Out of Scope Changes check	✅ Passed	The source changes implement weak instance retention and safe registry access. The test changes cover the garbage-collection requirement and update related registry wording. These changes support issu…
Docstring Coverage	✅ Passed	Docstring coverage is 87.50% which is sufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 8 functions across 2 files.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create a new PR

---

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.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{A rabbit sees the weak set glow
Built tools join when ready to go
Snapshots hop through threads with care
Old managers fade into the air
GC finds no strong root to spare}

_{Comment @coderabbitai help to get the list of available commands.}

[Nandith-0777]

@coderabbitai review

[coderabbitai]

✅ Action performed

Review finished.

Note: CodeRabbit is an incremental review system and does not re-review already reviewed commits. This command is applicable only when automatic reviews are paused.

[Nandith-0777]

@coderabbitai review

[coderabbitai]

✅ Action performed

Review finished.

Note: CodeRabbit is an incremental review system and does not re-review already reviewed commits. This command is applicable only when automatic reviews are paused.

[Lokkhita]

Hi @Nandith-0777 ,I opened #338 for the same issue a day after this one, before your Sep 22 update. Your current branch covers both the leak and the RuntimeError from a ToolManager being created while add_tool_to_all() iterates the WeakSet,so I'll close #338 in favour of this PR.