feat: enable multiagent session persistence

Author: JackYPCOnline

src/strands/experimental/multiagent_session/__init__.py renamed to src/strands/experimental/multiagent_hooks/__init__.py

@@ -5,22 +5,19 @@
 """
 
 from .multiagent_events import (
-    AfterGraphInvocationEvent,
+    AfterMultiAgentInvocationEvent,
     AfterNodeInvocationEvent,
-    BeforeGraphInvocationEvent,
+    BeforeMultiAgentInvocationEvent,
     BeforeNodeInvocationEvent,
     MultiAgentInitializationEvent,
 )
-from .multiagent_state import MultiAgentState, MultiAgentType
-from .multiagent_state_adapter import MultiAgentAdapter
+from .persistence_hooks import PersistentHook
 
 __all__ = [
-    "BeforeGraphInvocationEvent",
-    "AfterGraphInvocationEvent",
+    "BeforeMultiAgentInvocationEvent",
+    "AfterMultiAgentInvocationEvent",
     "MultiAgentInitializationEvent",
     "BeforeNodeInvocationEvent",
     "AfterNodeInvocationEvent",
-    "MultiAgentState",
-    "MultiAgentAdapter",
-    "MultiAgentType",
+    "PersistentHook",
 ]

src/strands/experimental/multiagent_session/multiagent_events.py renamed to src/strands/experimental/multiagent_hooks/multiagent_events.py

@@ -1,87 +1,82 @@
 """Multi-agent execution lifecycle events for hook system integration.
 
-This module defines event classes that are triggered at key points during
-multi-agent orchestrator execution, enabling hooks to respond to lifecycle
-events for purposes like persistence, monitoring, and debugging.
-
-Event Types:
-- Initialization: When orchestrator starts up
-- Before/After Graph: Start/end of overall execution
-- Before/After Node: Start/end of individual node execution
+These events are fired by orchestrators (Graph/Swarm) at key points so
+hooks can persist, monitor, or debug execution. No intermediate state model
+is used—hooks read from the orchestrator directly.
 """
 
 from dataclasses import dataclass
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Any
 
-from ...hooks.registry import HookEvent
-from .multiagent_state import MultiAgentState
+from ...hooks.registry import BaseHookEvent
 
 if TYPE_CHECKING:
     from ...multiagent.base import MultiAgentBase
 
 
 @dataclass
-class MultiAgentInitializationEvent(HookEvent):
+class MultiAgentInitializationEvent(BaseHookEvent):
     """Event triggered when multi-agent orchestrator initializes.
 
     Attributes:
         orchestrator: The multi-agent orchestrator instance
-        state: Current state of the orchestrator
+        invocation_state: Configuration that user pass in
     """
 
     orchestrator: "MultiAgentBase"
-    state: MultiAgentState
+    invocation_state: dict[str, Any] | None = None
 
 
 @dataclass
-class BeforeGraphInvocationEvent(HookEvent):
+class BeforeMultiAgentInvocationEvent(BaseHookEvent):
     """Event triggered before orchestrator execution begins.
 
     Attributes:
         orchestrator: The multi-agent orchestrator instance
-        state: Current state before execution starts
+        invocation_state: Configuration that user pass in
     """
 
     orchestrator: "MultiAgentBase"
-    state: MultiAgentState
+    invocation_state: dict[str, Any] | None = None
 
 
 @dataclass
-class BeforeNodeInvocationEvent(HookEvent):
+class BeforeNodeInvocationEvent(BaseHookEvent):
     """Event triggered before individual node execution.
 
     Attributes:
         orchestrator: The multi-agent orchestrator instance
-        next_node_to_execute: ID of the node about to be executed
+        invocation_state: Configuration that user pass in
     """
 
     orchestrator: "MultiAgentBase"
     next_node_to_execute: str
+    invocation_state: dict[str, Any] | None = None
 
 
 @dataclass
-class AfterNodeInvocationEvent(HookEvent):
+class AfterNodeInvocationEvent(BaseHookEvent):
     """Event triggered after individual node execution completes.
 
     Attributes:
         orchestrator: The multi-agent orchestrator instance
         executed_node: ID of the node that just completed execution
-        state: Updated state after node execution
+        invocation_state: Configuration that user pass in
     """
 
     orchestrator: "MultiAgentBase"
     executed_node: str
-    state: MultiAgentState
+    invocation_state: dict[str, Any] | None = None
 
 
 @dataclass
-class AfterGraphInvocationEvent(HookEvent):
+class AfterMultiAgentInvocationEvent(BaseHookEvent):
     """Event triggered after orchestrator execution completes.
 
     Attributes:
         orchestrator: The multi-agent orchestrator instance
-        state: Final state after execution completes
+        invocation_state: Configuration that user pass in
     """
 
     orchestrator: "MultiAgentBase"
-    state: MultiAgentState
+    invocation_state: dict[str, Any] | None = None

src/strands/experimental/multiagent_session/persistence_hooks.py renamed to src/strands/experimental/multiagent_hooks/persistence_hooks.py

@@ -11,52 +11,37 @@
 """
 
 import threading
-from typing import Optional
+from typing import TYPE_CHECKING
 
 from ...hooks.registry import HookProvider, HookRegistry
-from ...multiagent.base import MultiAgentBase
 from ...session import SessionManager
 from .multiagent_events import (
-    AfterGraphInvocationEvent,
+    AfterMultiAgentInvocationEvent,
     AfterNodeInvocationEvent,
-    BeforeGraphInvocationEvent,
+    BeforeMultiAgentInvocationEvent,
     BeforeNodeInvocationEvent,
     MultiAgentInitializationEvent,
-    MultiAgentState,
 )
-from .multiagent_state_adapter import MultiAgentAdapter
 
+if TYPE_CHECKING:
+    from ...multiagent.base import MultiAgentBase
 
-def _get_multiagent_state(
-    multiagent_state: Optional[MultiAgentState],
-    orchestrator: MultiAgentBase,
-) -> MultiAgentState:
-    if multiagent_state is not None:
-        return multiagent_state
 
-    return MultiAgentAdapter.create_multi_agent_state(orchestrator=orchestrator)
-
-
-class MultiAgentHook(HookProvider):
+class PersistentHook(HookProvider):
     """Hook provider for automatic multi-agent session persistence.
 
     This hook automatically persists multi-agent orchestrator state at key
     execution points to enable resumable execution after interruptions.
 
-    Args:
-        session_manager: SessionManager instance for state persistence
-        session_id: Unique identifier for the session
     """
 
-    def __init__(self, session_manager: SessionManager, session_id: str):
+    def __init__(self, session_manager: SessionManager):
         """Initialize the multi-agent persistence hook.
 
         Args:
             session_manager: SessionManager instance for state persistence
-            session_id: Unique identifier for the session
         """
         self._session_manager = session_manager
-        self._session_id = session_id
         self._lock = threading.RLock()
 
     def register_hooks(self, registry: HookRegistry, **kwargs: object) -> None:
@@ -67,40 +52,40 @@ def register_hooks(self, registry: HookRegistry, **kwargs: object) -> None:
             **kwargs: Additional keyword arguments (unused)
         """
         registry.add_callback(MultiAgentInitializationEvent, self._on_initialization)
-        registry.add_callback(BeforeGraphInvocationEvent, self._on_before_graph)
+        registry.add_callback(BeforeMultiAgentInvocationEvent, self._on_before_multiagent)
         registry.add_callback(BeforeNodeInvocationEvent, self._on_before_node)
         registry.add_callback(AfterNodeInvocationEvent, self._on_after_node)
-        registry.add_callback(AfterGraphInvocationEvent, self._on_after_graph)
+        registry.add_callback(AfterMultiAgentInvocationEvent, self._on_after_multiagent)
 
-    def _on_initialization(self, event: MultiAgentInitializationEvent):
+    # TODO: We can add **kwarg or invocation_state later if we need to persist
+    def _on_initialization(self, event: MultiAgentInitializationEvent) -> None:
         """Persist state when multi-agent orchestrator initializes."""
-        self._persist(_get_multiagent_state(event.state, event.orchestrator))
+        self._persist(event.orchestrator)
 
-    def _on_before_graph(self, event: BeforeGraphInvocationEvent):
-        """Hook called before graph execution starts."""
+    def _on_before_multiagent(self, event: BeforeMultiAgentInvocationEvent) -> None:
+        """Persist state when multi-agent orchestrator initializes."""
         pass
 
-    def _on_before_node(self, event: BeforeNodeInvocationEvent):
+    def _on_before_node(self, event: BeforeNodeInvocationEvent) -> None:
         """Hook called before individual node execution."""
         pass
 
-    def _on_after_node(self, event: AfterNodeInvocationEvent):
+    def _on_after_node(self, event: AfterNodeInvocationEvent) -> None:
         """Persist state after each node completes execution."""
-        multi_agent_state = _get_multiagent_state(multiagent_state=event.state, orchestrator=event.orchestrator)
-        self._persist(multi_agent_state)
+        self._persist(event.orchestrator)
 
-    def _on_after_graph(self, event: AfterGraphInvocationEvent):
+    def _on_after_multiagent(self, event: AfterMultiAgentInvocationEvent) -> None:
         """Persist final state after graph execution completes."""
-        multiagent_state = _get_multiagent_state(multiagent_state=event.state, orchestrator=event.orchestrator)
-        self._persist(multiagent_state)
+        self._persist(event.orchestrator)
 
-    def _persist(self, multiagent_state: MultiAgentState) -> None:
+    def _persist(self, orchestrator: "MultiAgentBase") -> None:
         """Persist the provided MultiAgentState using the configured SessionManager.
 
         This method is synchronized across threads/tasks to avoid write races.
 
         Args:
-            multiagent_state: State to persist
+            orchestrator: State to persist
         """
+        current_state = orchestrator.get_state_from_orchestrator()
         with self._lock:
-            self._session_manager.write_multi_agent_state(multiagent_state)
+            self._session_manager.write_multi_agent_json(current_state)