fixing Parallel Stepping Conflict

mesa_llm/parallel_stepping.py

@@ -24,6 +24,10 @@ async def step_agents_parallel(agents: list[Agent | LLMAgent]) -> None:
     """Step all agents in parallel using async/await."""
     tasks = []
     for agent in agents:
+        # Skip agents that are components of a Meta Agent
+        if getattr(agent, "is_component", False):
+            continue
+
         if hasattr(agent, "astep"):
             tasks.append(agent.astep())
         elif hasattr(agent, "step"):
@@ -41,6 +45,10 @@ def step_agents_multithreaded(agents: list[Agent | LLMAgent]) -> None:
     with concurrent.futures.ThreadPoolExecutor() as executor:
         futures = []
         for agent in agents:
+            # Skip agents that are components of a Meta Agent
+            if getattr(agent, "is_component", False):
+                continue
+
             if hasattr(agent, "astep"):
                 # run async steps in the event loop in a thread
                 futures.append(
@@ -75,6 +83,13 @@ def step_agents_parallel_sync(agents: list[Agent | LLMAgent]) -> None:
 
 # Patch Mesa's shuffle_do for automatic parallel detection
 _original_shuffle_do = AgentSet.shuffle_do
+try:
+    from mesa.agentset import _HardKeyAgentSet
+
+    _original_hardkey_shuffle_do = _HardKeyAgentSet.shuffle_do
+except ImportError:
+    _HardKeyAgentSet = None
+    _original_hardkey_shuffle_do = None
 
 
 def _enhanced_shuffle_do(self, method: str, *args, **kwargs):
@@ -84,7 +99,11 @@ def _enhanced_shuffle_do(self, method: str, *args, **kwargs):
         if hasattr(agent, "model") and getattr(agent.model, "parallel_stepping", False):
             step_agents_parallel_sync(list(self))
             return
-    _original_shuffle_do(self, method, *args, **kwargs)
+
+    if _HardKeyAgentSet and isinstance(self, _HardKeyAgentSet):
+        _original_hardkey_shuffle_do(self, method, *args, **kwargs)
+    else:
+        _original_shuffle_do(self, method, *args, **kwargs)
 
 
 def enable_automatic_parallel_stepping(mode: str = "asyncio"):
@@ -94,11 +113,15 @@ def enable_automatic_parallel_stepping(mode: str = "asyncio"):
         raise ValueError("mode must be either 'asyncio' or 'threading'")
     _PARALLEL_STEPPING_MODE = mode
     AgentSet.shuffle_do = _enhanced_shuffle_do
+    if _HardKeyAgentSet:
+        _HardKeyAgentSet.shuffle_do = _enhanced_shuffle_do
 
 
 def disable_automatic_parallel_stepping():
     """Restore original shuffle_do behavior."""
     AgentSet.shuffle_do = _original_shuffle_do
+    if _HardKeyAgentSet:
+        _HardKeyAgentSet.shuffle_do = _original_hardkey_shuffle_do
 
 
 # --- Monkey-patch AgentSet with do_async for async parallel method calls ---
@@ -114,6 +137,10 @@ def _agentset_do_async(self, method: str, *args, **kwargs):
     async def _run():
         tasks = []
         for agent in self:
+            # Skip agents that are components of a Meta Agent
+            if getattr(agent, "is_component", False):
+                continue
+
             fn = getattr(agent, method, None)
             if fn is not None and asyncio.iscoroutinefunction(fn):
                 tasks.append(fn(*args, **kwargs))

tests/test_parallel_stepping.py

@@ -2,6 +2,7 @@
 
 import pytest
 from mesa.agent import Agent, AgentSet
+from mesa.experimental.meta_agents.meta_agent import MetaAgent
 from mesa.model import Model
 
 from mesa_llm.parallel_stepping import (
@@ -15,7 +16,7 @@
 
 class DummyModel(Model):
     def __init__(self):
-        super().__init__(seed=42)
+        super().__init__(rng=42)
         self.parallel_stepping = False
 
 
@@ -99,3 +100,84 @@ async def wrapper():
     asyncio.run(wrapper())
     assert a1.counter == 1
     assert a2.counter == 1
+
+
+@pytest.fixture(autouse=True)
+def manage_parallel_stepping_patch():
+    # Helper to clean up parallel stepping state if tests fail
+    yield
+    disable_automatic_parallel_stepping()
+
+
+# --- Meta Agent Parallel Conflict Tests ---
+
+
+class ConflictWorker(Agent):
+    def __init__(self, model):
+        super().__init__(model)
+        self.step_count = 0
+
+    def step(self):
+        self.step_count += 1
+
+
+class ConflictManager(MetaAgent):
+    def step(self):
+        for agent in self.agents:
+            agent.step()
+
+
+class ConflictBusinessModel(Model):
+    def __init__(self):
+        super().__init__(rng=42)
+        self.parallel_stepping = True
+        self.worker = ConflictWorker(self)
+        self.manager = ConflictManager(self, agents={self.worker})
+
+    def step(self):
+        self.agents.shuffle_do("step")
+
+
+def test_meta_agent_parallel_conflict_fix():
+    """
+    Test that a constituent agent is only stepped ONCE when
+    parallel stepping is enabled, because it is skipped by
+    the scheduler and only stepped by its MetaAgent.
+    """
+    enable_automatic_parallel_stepping("asyncio")
+    model = ConflictBusinessModel()
+
+    # Run one step
+    model.step()
+
+    # Assert worker was only stepped once
+    # If the fix failed, this would be 2
+    assert model.worker.step_count == 1
+
+
+def test_meta_agent_multithreaded_conflict_fix():
+    """Test same logic but with multithreaded mode."""
+    disable_automatic_parallel_stepping()
+    enable_automatic_parallel_stepping("threading")
+
+    model = ConflictBusinessModel()
+    model.step()
+
+    assert model.worker.step_count == 1
+
+
+def test_agent_becomes_independent_again():
+    """Test that an agent removed from MetaAgent is stepped by scheduler again."""
+    enable_automatic_parallel_stepping("asyncio")
+    model = ConflictBusinessModel()
+
+    # 1. Initially it's a component
+    assert model.worker.is_component is True
+
+    # 2. Remove it from manager
+    model.manager.remove_constituting_agents({model.worker})
+    assert model.worker.is_component is False
+
+    # 3. Model step should now step the worker directly
+    model.step()
+    assert model.worker.step_count == 1