bug: Double pre_step/post_step calls when astep() delegates to wrapped step() causes orphaned memory entries

[yashhzd]

Bug Description

When a subclass of LLMAgent defines only step() (no astep()), the default astep() method calls apre_step() and apost_step() — but then delegates to self.step(), which is already wrapped by __init_subclass__ with its own pre_step() and post_step(). This causes process_step(pre_step=True) and process_step(pre_step=False) to each be called twice per step during parallel execution.

For STLTMemory (the default memory), this creates an orphaned MemoryEntry with step=None that is never finalized or cleaned up.

Environment

• mesa-llm: v0.3.0 (commit a8161c7)
• mesa: 3.5.0
• Python: 3.12

Root Cause

In llm_agent.py, the default astep() (line 318) and __init_subclass__ wrapping (line 331) both inject pre/post step hooks:

Default astep() (line 318–329):

async def astep(self):
    await self.apre_step()          # ← First pre_step call
    if hasattr(self, "step") and self.__class__.step != LLMAgent.step:
        self.step()                 # ← This is the WRAPPED step
    await self.apost_step()         # ← Second post_step call

__init_subclass__ wrapping (line 340–351):

if user_step:
    def wrapped(self, *args, **kwargs):
        LLMAgent.pre_step(self)     # ← Second pre_step call
        result = user_step(self)
        LLMAgent.post_step(self)    # ← First post_step call
        return result
    cls.step = wrapped

When astep() calls the wrapped self.step(), the actual call sequence is:

#	Call	Effect on STLTMemory
1	apre_step()	Creates entry1 (with step_content data, step=None), resets step_content={}
2	wrapped → pre_step()	Creates entry2 (empty content, step=None)
3	wrapped → user_step()	User's logic runs, may add to step_content
4	wrapped → post_step()	Pops entry2 (last), merges with step_content, creates real_entry
5	apost_step()	Sees last entry has step != None → early return, no-op

Result: entry1 is permanently orphaned in short_term_memory with step=None.

Impact

• Memory pollution: Orphaned entries with step=None accumulate in short-term memory (one per step during parallel execution)
• Premature consolidation: Orphaned entries count toward len(short_term_memory), triggering early popleft() eviction of real entries
• LLM context confusion: format_short_term() renders orphaned entries as "Step None: {data}", polluting the prompt sent to the LLM
• Data loss: The first process_step(pre_step=True) captures the original step_content but that entry is never finalized — the actual data ends up orphaned while the second (empty) entry gets processed

Reproduction

import asyncio
from mesa.model import Model
from mesa_llm.llm_agent import LLMAgent
from mesa_llm.reasoning.cot import CoTReasoning

class DummyModel(Model):
    def __init__(self):
        super().__init__(seed=42)

class MyAgent(LLMAgent):
    def __init__(self, model):
        super().__init__(model, reasoning=CoTReasoning, llm_model="gemini/gemini-2.0-flash")

    def step(self):  # Only defines step(), no astep()
        pass

async def main():
    m = DummyModel()
    agent = MyAgent(m)

    # Simulate what parallel stepping does
    await agent.astep()

    # Check for orphaned entries
    orphaned = [e for e in agent.memory.short_term_memory if e.step is None]
    print(f"Orphaned entries: {len(orphaned)}")  # Expected: 0, Actual: 1

    # After 10 steps, there are 10 orphaned entries
    for _ in range(9):
        await agent.astep()
    orphaned = [e for e in agent.memory.short_term_memory if e.step is None]
    print(f"Orphaned entries after 10 steps: {len(orphaned)}")  # 10!

asyncio.run(main())

Who Is Affected

Any LLMAgent subclass that:

1. Defines only step() (no custom astep()) — the most common pattern
2. Is executed via astep() during parallel stepping — the default execution mode

This means the bug affects most mesa-llm users running parallel simulations with the default STLTMemory.

Expected Behavior

• pre_step and post_step should be called exactly once per step, regardless of whether execution goes through step() or astep()
• No orphaned MemoryEntry objects with step=None should persist in memory