fix: rebase from main and address comments

Author: JackYPCOnline

src/strands/multiagent/graph.py

@@ -29,6 +29,7 @@
 from ..experimental.hooks.multiagent import (
     AfterMultiAgentInvocationEvent,
     AfterNodeCallEvent,
+    BeforeNodeCallEvent,
     MultiAgentInitializedEvent,
 )
 from ..hooks import HookProvider, HookRegistry
@@ -409,7 +410,6 @@ def __init__(
         reset_on_revisit: bool = False,
         session_manager: Optional[SessionManager] = None,
         hooks: Optional[list[HookProvider]] = None,
-        *,
         id: str = _DEFAULT_GRAPH_ID,
     ) -> None:
         """Initialize Graph with execution limits and reset behavior.
@@ -430,7 +430,6 @@ def __init__(
 
         # Validate nodes for duplicate instances
         self._validate_graph(nodes)
-        self.id = id or _DEFAULT_GRAPH_ID
 
         self.nodes = nodes
         self.edges = edges
@@ -451,6 +450,7 @@ def __init__(
 
         self._resume_next_nodes: list[GraphNode] = []
         self._resume_from_session = False
+        self.id = id
 
         self.hooks.invoke_callbacks(MultiAgentInitializedEvent(self))
 
@@ -773,7 +773,9 @@ def _is_node_ready_with_conditions(self, node: GraphNode, completed_batch: list[
 
     async def _execute_node(self, node: GraphNode, invocation_state: dict[str, Any]) -> AsyncIterator[Any]:
         """Execute a single node and yield TypedEvent objects."""
-        # Reset the node's state if reset_on_revisit is enabled and it's being revisited
+        self.hooks.invoke_callbacks(BeforeNodeCallEvent(self, node.node_id, invocation_state))
+
+        # Reset the node's state if reset_on_revisit is enabled, and it's being revisited
         if self.reset_on_revisit and node in self.state.completed_nodes:
             logger.debug("node_id=<%s> | resetting node state for revisit", node.node_id)
             node.reset_executor_state()
@@ -914,6 +916,9 @@ async def _execute_node(self, node: GraphNode, invocation_state: dict[str, Any])
             # Re-raise to stop graph execution (fail-fast behavior)
             raise
 
+        finally:
+            self.hooks.invoke_callbacks(AfterNodeCallEvent(self, node.node_id, invocation_state))
+
     def _accumulate_metrics(self, node_result: NodeResult) -> None:
         """Accumulate metrics from a node result."""
         self.state.accumulated_usage["inputTokens"] += node_result.accumulated_usage.get("inputTokens", 0)
@@ -1001,13 +1006,12 @@ def _build_result(self) -> GraphResult:
 
     def serialize_state(self) -> dict[str, Any]:
         """Serialize the current graph state to a dictionary."""
-        status_str = self.state.status.value
         compute_nodes = self._compute_ready_nodes_for_resume()
         next_nodes = [n.node_id for n in compute_nodes] if compute_nodes else []
         return {
             "type": "graph",
             "id": self.id,
-            "status": status_str,
+            "status": self.state.status.value,
             "completed_nodes": [n.node_id for n in self.state.completed_nodes],
             "failed_nodes": [n.node_id for n in self.state.failed_nodes],
             "node_results": {k: v.to_dict() for k, v in (self.state.results or {}).items()},
@@ -1020,10 +1024,10 @@ def deserialize_state(self, payload: dict[str, Any]) -> None:
         """Restore graph state from a session dict and prepare for execution.
 
         This method handles two scenarios:
-        1. If the persisted status is COMPLETED, FAILED resets all nodes and graph state
-           to allow re-execution from the beginning.
-        2. Otherwise, restores the persisted state and prepares to resume execution
-           from the next ready nodes.
+        1. If the graph execution ended (no next_nodes_to_execute, eg: Completed, or Failed with dead end nodes),
+        resets all nodes and graph state to allow re-execution from the beginning.
+        2. If the graph execution was interrupted mid-execution (has next_nodes_to_execute),
+           restores the persisted state and prepares to resume execution from the next ready nodes.
 
         Args:
             payload: Dictionary containing persisted state data including status,
@@ -1041,7 +1045,6 @@ def deserialize_state(self, payload: dict[str, Any]) -> None:
             self._from_dict(payload)
             self._resume_from_session = True
 
-    # Helper functions for serialize and deserialize
     def _compute_ready_nodes_for_resume(self) -> list[GraphNode]:
         if self.state.status == Status.PENDING:
             return []
@@ -1073,7 +1076,9 @@ def _from_dict(self, payload: dict[str, Any]) -> None:
                 raise
         self.state.results = results
 
-        self.state.failed_nodes = set(payload.get("failed_nodes") or [])
+        self.state.failed_nodes = set(
+            self.nodes[node_id] for node_id in (payload.get("failed_nodes") or []) if node_id in self.nodes
+        )
 
         # Restore completed nodes from persisted data
         completed_node_ids = payload.get("completed_nodes") or []

src/strands/multiagent/swarm.py

@@ -18,7 +18,7 @@
 import logging
 import time
 from dataclasses import dataclass, field
-from typing import Any, AsyncIterator, Callable, Tuple, cast,Optional
+from typing import Any, AsyncIterator, Callable, Optional, Tuple, cast
 
 from opentelemetry import trace as trace_api
 
@@ -28,6 +28,7 @@
 from ..experimental.hooks.multiagent import (
     AfterMultiAgentInvocationEvent,
     AfterNodeCallEvent,
+    BeforeNodeCallEvent,
     MultiAgentInitializedEvent,
 )
 from ..hooks import HookProvider, HookRegistry
@@ -210,8 +211,8 @@ class Swarm(MultiAgentBase):
 
     def __init__(
         self,
-        id:_DEFAULT_SWARM_ID,
         nodes: list[Agent],
+        id: str = _DEFAULT_SWARM_ID,
         *,
         entry_point: Agent | None = None,
         max_handoffs: int = 20,
@@ -351,7 +352,6 @@ async def stream_async(
             self.state.completion_status = Status.EXECUTING
             self.state.start_time = time.time()
 
-        start_time = time.time()
         span = self.tracer.start_multiagent_span(task, "swarm")
         with trace_api.use_span(span, end_on_exit=True):
             try:
@@ -372,7 +372,7 @@ async def stream_async(
                 self.state.completion_status = Status.FAILED
                 raise
             finally:
-                self.state.execution_time = round((time.time() - start_time) * 1000)
+                self.state.execution_time = round((time.time() - self.state.start_time) * 1000)
                 self.hooks.invoke_callbacks(AfterMultiAgentInvocationEvent(self, invocation_state))
                 self._resume_from_session = False
 
@@ -685,6 +685,7 @@ async def _execute_swarm(self, invocation_state: dict[str, Any]) -> AsyncIterato
                 # TODO: Implement cancellation token to stop _execute_node from continuing
                 try:
                     # Execute with timeout wrapper for async generator streaming
+                    self.hooks.invoke_callbacks(BeforeNodeCallEvent(self, current_node.node_id, invocation_state))
                     node_stream = self._stream_with_timeout(
                         self._execute_node(current_node, self.state.task, invocation_state),
                         self.node_timeout,

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

@@ -67,7 +67,7 @@ def test_swarm_complete_hook_lifecycle(swarm, hook_provider):
     result = swarm("test task")
 
     length, events = hook_provider.get_events()
-    assert length == 3
+    assert length == 4
     assert result.status.value == "completed"
 
     events_list = list(events)
@@ -76,20 +76,24 @@ 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], AfterNodeCallEvent)
+    assert isinstance(events_list[1], BeforeNodeCallEvent)
     assert events_list[1].source == swarm
     assert events_list[1].node_id == "agent1"
 
-    assert isinstance(events_list[2], AfterMultiAgentInvocationEvent)
+    assert isinstance(events_list[2], AfterNodeCallEvent)
     assert events_list[2].source == swarm
+    assert events_list[2].node_id == "agent1"
+
+    assert isinstance(events_list[3], AfterMultiAgentInvocationEvent)
+    assert events_list[3].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 == 4
+    assert length == 6
     assert result.status.value == "completed"
 
     events_list = list(events)
@@ -98,13 +102,21 @@ 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], AfterNodeCallEvent)
+    assert isinstance(events_list[1], BeforeNodeCallEvent)
     assert events_list[1].source == graph
     assert events_list[1].node_id == "agent1"
 
     assert isinstance(events_list[2], AfterNodeCallEvent)
     assert events_list[2].source == graph
-    assert events_list[2].node_id == "agent2"
+    assert events_list[2].node_id == "agent1"
 
-    assert isinstance(events_list[3], AfterMultiAgentInvocationEvent)
+    assert isinstance(events_list[3], BeforeNodeCallEvent)
     assert events_list[3].source == graph
+    assert events_list[3].node_id == "agent2"
+
+    assert isinstance(events_list[4], AfterNodeCallEvent)
+    assert events_list[4].source == graph
+    assert events_list[4].node_id == "agent2"
+
+    assert isinstance(events_list[5], AfterMultiAgentInvocationEvent)
+    assert events_list[5].source == graph

tests/strands/multiagent/test_graph.py

@@ -2010,6 +2010,7 @@ async def test_graph_persisted(mock_strands_tracer, mock_use_span):
     persisted_state = {
         "status": "executing",
         "completed_nodes": [],
+        "failed_nodes": [],
         "node_results": {},
         "current_task": "persisted task",
         "execution_order": [],

tests_integ/test_multiagent_graph.py

@@ -1,8 +1,7 @@
+from typing import Any, AsyncIterator
 from unittest.mock import patch
 from uuid import uuid4
 
-from typing import Any, AsyncIterator
-
 import pytest
 
 from strands import Agent, tool
@@ -14,7 +13,6 @@
     BeforeModelCallEvent,
     MessageAddedEvent,
 )
-from strands.multiagent.base import Status
 from strands.multiagent.base import MultiAgentBase, MultiAgentResult, NodeResult, Status
 from strands.multiagent.graph import GraphBuilder
 from strands.session.file_session_manager import FileSessionManager

tests_integ/test_multiagent_swarm.py

@@ -15,9 +15,9 @@
 )
 from strands.multiagent.base import Status
 from strands.multiagent.swarm import Swarm
+from strands.session.file_session_manager import FileSessionManager
 from strands.types.content import ContentBlock
 from tests.fixtures.mock_hook_provider import MockHookProvider
-from strands.session.file_session_manager import FileSessionManager
 
 
 @tool