fix: prevent double pre_step/post_step when astep() delegates to wrapped step()

[yashhzd]

Summary

Closes #222

Fixes a bug where pre_step/post_step (and their async variants) are called twice per step when a subclass defines only step() and is executed via astep() during parallel stepping. This is the most common usage pattern in mesa-llm.

Root Cause

The default LLMAgent.astep() calls apre_step() → self.step() → apost_step(). But self.step() is the wrapped version from __init_subclass__, which also calls pre_step() → user_step() → post_step(). Result: memory hooks fire twice.

The Fix

Store the raw user step() function as cls._raw_user_step during __init_subclass__ wrapping. The default astep() calls _raw_user_step directly (bypassing the pre/post wrapper), since it handles its own async pre/post hooks.

Path	Before	After
astep() → step-only subclass	pre_step ×2, post_step ×2	pre_step ×1, post_step ×1
astep() → astep-only subclass	pre_step ×1 (already correct)	unchanged
step() called directly	pre_step ×1 (already correct)	unchanged

Impact of the Bug

• Memory pollution: Orphaned MemoryEntry with step=None accumulates in STLTMemory (one per step)
• Premature consolidation: Orphaned entries count toward memory capacity, evicting real entries early
• LLM confusion: format_short_term() renders orphaned entries as "Step None: {data}" in the prompt

Backward Compatibility

• All 281 existing tests pass unchanged
• Direct step() calls still go through the wrapper (no change)
• Only the astep() → step() delegation path is fixed
• _raw_user_step is a private class attribute, no public API change

Test Plan

• All 281 existing tests pass
• 6 new tests covering:
  • astep() with step-only agent: exactly 1 pre/post call
  • astep() with async-only agent: exactly 1 pre/post call
  • astep() with both-methods agent: exactly 1 pre/post call
  • Direct step() call: exactly 1 pre/post call
  • 5 consecutive astep() calls: counts == 5, not 10
  • Real STLTMemory: no orphaned entries with step=None
• Ruff lint and format clean

[coderabbitai]

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: Pro

Run ID: 85f0d078-09ae-4751-b5af-1b024806663b

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

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

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

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

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 92.03%. Comparing base (a840baf) to head (a61a837).
⚠️ Report is 3 commits behind head on main.

Additional details and impacted files

@@           Coverage Diff           @@
##             main     #224   +/-   ##
=======================================
  Coverage   92.02%   92.03%           
=======================================
  Files          19       19           
  Lines        1681     1683    +2     
=======================================
+ Hits         1547     1549    +2     
  Misses        134      134           

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[wang-boyu]

Thanks for the fix. I have merged main to the PR branch to resolve conflicts introduced by another recently merged PR.

Merging now.