fix: add BeforeMultiAgentInvocation events in Graph/Swarm, add more tests

Author: JackYPCOnline

src/strands/multiagent/graph.py

@@ -29,6 +29,7 @@
 from ..experimental.hooks.multiagent import (
     AfterMultiAgentInvocationEvent,
     AfterNodeCallEvent,
+    BeforeMultiAgentInvocationEvent,
     BeforeNodeCallEvent,
     MultiAgentInitializedEvent,
 )
@@ -468,6 +469,7 @@ def __call__(
         if invocation_state is None:
             invocation_state = {}
 
+        self.hooks.invoke_callbacks(BeforeMultiAgentInvocationEvent(self, invocation_state))
         return run_async(lambda: self.invoke_async(task, invocation_state))
 
     async def invoke_async(

src/strands/multiagent/swarm.py

@@ -28,6 +28,7 @@
 from ..experimental.hooks.multiagent import (
     AfterMultiAgentInvocationEvent,
     AfterNodeCallEvent,
+    BeforeMultiAgentInvocationEvent,
     BeforeNodeCallEvent,
     MultiAgentInitializedEvent,
 )
@@ -287,7 +288,7 @@ def __call__(
         """
         if invocation_state is None:
             invocation_state = {}
-
+        self.hooks.invoke_callbacks(BeforeMultiAgentInvocationEvent(self, invocation_state))
         return run_async(lambda: self.invoke_async(task, invocation_state))
 
     async def invoke_async(

tests/strands/experimental/hooks/multiagent/test_multi_agent_hooks.py

@@ -4,6 +4,7 @@
 from strands.experimental.hooks.multiagent.events import (
     AfterMultiAgentInvocationEvent,
     AfterNodeCallEvent,
+    BeforeMultiAgentInvocationEvent,
     BeforeNodeCallEvent,
     MultiAgentInitializedEvent,
 )
@@ -17,6 +18,7 @@
 def hook_provider():
     return MockMultiAgentHookProvider(
         [
+            BeforeMultiAgentInvocationEvent,
             AfterMultiAgentInvocationEvent,
             AfterNodeCallEvent,
             BeforeNodeCallEvent,
@@ -67,7 +69,7 @@ def test_swarm_complete_hook_lifecycle(swarm, hook_provider):
     result = swarm("test task")
 
     length, events = hook_provider.get_events()
-    assert length == 4
+    assert length == 5
     assert result.status.value == "completed"
 
     events_list = list(events)
@@ -76,24 +78,27 @@ def test_swarm_complete_hook_lifecycle(swarm, hook_provider):
     assert isinstance(events_list[0], MultiAgentInitializedEvent)
     assert events_list[0].source == swarm
 
-    assert isinstance(events_list[1], BeforeNodeCallEvent)
+    assert isinstance(events_list[1], BeforeMultiAgentInvocationEvent)
     assert events_list[1].source == swarm
-    assert events_list[1].node_id == "agent1"
 
-    assert isinstance(events_list[2], AfterNodeCallEvent)
+    assert isinstance(events_list[2], BeforeNodeCallEvent)
     assert events_list[2].source == swarm
     assert events_list[2].node_id == "agent1"
 
-    assert isinstance(events_list[3], AfterMultiAgentInvocationEvent)
+    assert isinstance(events_list[3], AfterNodeCallEvent)
     assert events_list[3].source == swarm
+    assert events_list[3].node_id == "agent1"
+
+    assert isinstance(events_list[4], AfterMultiAgentInvocationEvent)
+    assert events_list[4].source == swarm
 
 
 def test_graph_complete_hook_lifecycle(graph, hook_provider):
     """E2E test verifying complete hook lifecycle for Graph."""
     result = graph("test task")
 
     length, events = hook_provider.get_events()
-    assert length == 6
+    assert length == 7
     assert result.status.value == "completed"
 
     events_list = list(events)
@@ -102,21 +107,24 @@ def test_graph_complete_hook_lifecycle(graph, hook_provider):
     assert isinstance(events_list[0], MultiAgentInitializedEvent)
     assert events_list[0].source == graph
 
-    assert isinstance(events_list[1], BeforeNodeCallEvent)
+    assert isinstance(events_list[1], BeforeMultiAgentInvocationEvent)
     assert events_list[1].source == graph
-    assert events_list[1].node_id == "agent1"
 
-    assert isinstance(events_list[2], AfterNodeCallEvent)
+    assert isinstance(events_list[2], BeforeNodeCallEvent)
     assert events_list[2].source == graph
     assert events_list[2].node_id == "agent1"
 
-    assert isinstance(events_list[3], BeforeNodeCallEvent)
+    assert isinstance(events_list[3], AfterNodeCallEvent)
     assert events_list[3].source == graph
-    assert events_list[3].node_id == "agent2"
+    assert events_list[3].node_id == "agent1"
 
-    assert isinstance(events_list[4], AfterNodeCallEvent)
+    assert isinstance(events_list[4], BeforeNodeCallEvent)
     assert events_list[4].source == graph
     assert events_list[4].node_id == "agent2"
 
-    assert isinstance(events_list[5], AfterMultiAgentInvocationEvent)
+    assert isinstance(events_list[5], AfterNodeCallEvent)
     assert events_list[5].source == graph
+    assert events_list[5].node_id == "agent2"
+
+    assert isinstance(events_list[6], AfterMultiAgentInvocationEvent)
+    assert events_list[6].source == graph

tests_integ/test_multiagent_graph.py

@@ -488,13 +488,14 @@ async def test_graph_interrupt_and_resume():
     graph = builder.build()
 
     # Mock agent2 to fail on first execution
-    async def failing_invoke(*args, **kwargs):
+    async def failing_stream_async(*args, **kwargs):
         raise Exception("Simulated failure in agent2")
+        yield  # This line is never reached, but makes it an async generator
 
-    with patch.object(agent2, "invoke_async", side_effect=failing_invoke):
-        # First execution - should fail at agent2
+    with patch.object(agent2, "stream_async", side_effect=failing_stream_async):
         try:
-            await graph.invoke_async("Test task")
+            await graph.invoke_async("This is a test task, just do it shortly")
+            raise AssertionError("Expected exception was not raised")
         except Exception as e:
             assert "Simulated failure in agent2" in str(e)
 
@@ -530,3 +531,57 @@ async def failing_invoke(*args, **kwargs):
 
     # Clean up
     session_manager.delete_session(session_id)
+
+
+@pytest.mark.asyncio
+async def test_self_loop_resume_from_persisted_state(tmp_path):
+    """Test resuming self-loop from persisted state where next node is itself."""
+
+    session_id = f"self_loop_resume_{uuid4()}"
+    session_manager = FileSessionManager(session_id=session_id, storage_dir=str(tmp_path))
+
+    counter_agent = Agent(
+        model="us.amazon.nova-pro-v1:0",
+        system_prompt="You are a counter. Just respond with 'Count: 1', 'Count: 2', Stop at 5.",
+    )
+
+    def should_continue_loop(state):
+        loop_executions = len([node for node in state.execution_order if node.node_id == "loop_node"])
+        return loop_executions < 5
+
+    builder = GraphBuilder()
+    builder.add_node(counter_agent, "loop_node")
+    builder.add_edge("loop_node", "loop_node", condition=should_continue_loop)
+    builder.set_entry_point("loop_node")
+    builder.set_session_manager(session_manager)
+    builder.reset_on_revisit(True)
+
+    graph = builder.build()
+
+    call_count = 0
+    original_stream = counter_agent.stream_async
+
+    async def failing_after_two(*args, **kwargs):
+        nonlocal call_count
+        call_count += 1
+        if call_count <= 2:
+            async for event in original_stream(*args, **kwargs):
+                yield event
+        else:
+            raise Exception("Simulated failure after two executions")
+
+    with patch.object(counter_agent, "stream_async", side_effect=failing_after_two):
+        try:
+            await graph.invoke_async("Count till 5")
+        except Exception as e:
+            assert "Simulated failure after two executions" in str(e)
+
+    persisted_state = session_manager.read_multi_agent(session_id, graph.id)
+    assert persisted_state["status"] == "failed"
+    assert "loop_node" in persisted_state.get("failed_nodes")
+    assert len(persisted_state.get("execution_order")) == 2
+
+    result = await graph.invoke_async("Continue counting to 5")
+    assert result.status == Status.COMPLETED
+    assert len(result.execution_order) == 5
+    assert all(node.node_id == "loop_node" for node in result.execution_order)

tests_integ/test_multiagent_swarm.py

@@ -337,14 +337,18 @@ async def test_swarm_interrupt_and_resume(researcher_agent, analyst_agent, write
     # Create swarm with session manager
     swarm = Swarm([researcher_agent, analyst_agent, writer_agent], session_manager=session_manager)
 
-    # Mock analyst_agent to fail
+    # Mock analyst_agent's _invoke method to fail
     async def failing_invoke(*args, **kwargs):
         raise Exception("Simulated failure in analyst")
+        yield  # This line is never reached, but makes it an async generator
 
-    with patch.object(analyst_agent, "invoke_async", side_effect=failing_invoke):
+    with patch.object(analyst_agent, "stream_async", side_effect=failing_invoke):
         # First execution - should fail at analyst
         result = await swarm.invoke_async("Research AI trends and create a brief report")
-        assert result.status == Status.FAILED
+        try:
+            assert result.status == Status.FAILED
+        except Exception as e:
+            assert "Simulated failure in analyst" in str(e)
 
     # Verify partial execution was persisted
     persisted_state = session_manager.read_multi_agent(session_id, swarm.id)