feat(session-py): support snapshot session management for Graph and Swarm

strands-py/src/strands/multiagent/snapshot.py

@@ -0,0 +1,59 @@
+"""Snapshot capture and restore for multi-agent orchestrators (Graph and Swarm).
+
+Wraps an orchestrator's :meth:`~strands.multiagent.base.MultiAgentBase.serialize_state`
+output in the versioned :class:`~strands.types._snapshot.Snapshot` envelope so the persisted
+format can evolve, mirroring the TypeScript SDK's ``multiagent/snapshot.ts``.
+
+Well-known keys in ``data``:
+
+- ``orchestrator_id`` — orchestrator identity, verified on load.
+- ``state`` — the orchestrator's serialized execution state.
+"""
+
+from typing import TYPE_CHECKING, Any
+
+from ..types._snapshot import SNAPSHOT_SCHEMA_VERSION, Snapshot
+from ..types.exceptions import SnapshotException
+
+if TYPE_CHECKING:
+    from .base import MultiAgentBase
+
+
+def take_snapshot(orchestrator: "MultiAgentBase", *, app_data: dict[str, Any] | None = None) -> Snapshot:
+    """Capture an orchestrator's execution state as a versioned snapshot.
+
+    Args:
+        orchestrator: The Graph or Swarm to capture.
+        app_data: Application-owned data. Strands does not read or modify it.
+
+    Returns:
+        A ``multiAgent``-scope snapshot wrapping the orchestrator's serialized state.
+    """
+    return Snapshot(
+        scope="multiAgent",
+        schema_version=SNAPSHOT_SCHEMA_VERSION,
+        data={"orchestrator_id": orchestrator.id, "state": orchestrator.serialize_state()},
+        app_data=app_data or {},
+    )
+
+
+def load_snapshot(orchestrator: "MultiAgentBase", snapshot: Snapshot) -> None:
+    """Restore an orchestrator's execution state from a snapshot.
+
+    Args:
+        orchestrator: The Graph or Swarm to restore into.
+        snapshot: The snapshot to load.
+
+    Raises:
+        SnapshotException: If the snapshot is not a current-schema ``multiAgent`` snapshot,
+            or its ``orchestrator_id`` does not match ``orchestrator.id``.
+    """
+    snapshot.validate()
+    if snapshot.scope != "multiAgent":
+        raise SnapshotException(f"Expected snapshot scope 'multiAgent', got {snapshot.scope!r}")
+
+    snapshot_id = snapshot.data.get("orchestrator_id")
+    if snapshot_id != orchestrator.id:
+        raise SnapshotException(f"Snapshot orchestrator id mismatch: expected {orchestrator.id!r}, got {snapshot_id!r}")
+
+    orchestrator.deserialize_state(snapshot.data["state"])

strands-py/src/strands/session/snapshot_session_manager.py

@@ -31,7 +31,10 @@
 from ..experimental.hooks.events import BidiAgentInitializedEvent
 from ..hooks.events import (
     AfterInvocationEvent,
+    AfterMultiAgentInvocationEvent,
+    AfterNodeCallEvent,
     AgentInitializedEvent,
+    BeforeMultiAgentInvocationEvent,
     MessageAddedEvent,
     MultiAgentInitializedEvent,
 )
@@ -46,6 +49,7 @@
 
 if TYPE_CHECKING:
     from ..agent.agent import Agent
+    from ..multiagent.base import MultiAgentBase
 
 logger = logging.getLogger(__name__)
 
@@ -64,6 +68,19 @@
 # Derived from the Literal above so the accepted runtime values cannot drift from the type.
 _SAVE_LATEST_STRATEGIES = get_args(SaveLatestStrategy)
 
+MultiAgentSaveLatestStrategy = Literal["node", "invocation"]
+"""Controls how often an orchestrator's ``snapshot_latest`` is saved automatically.
+
+- ``"node"``: after every node completes (default; a mid-run crash resumes at the last node).
+- ``"invocation"``: only after the whole orchestrator invocation completes (lower I/O; a crash
+  loses the in-flight run). A large Graph on a remote store can opt down to this.
+
+Orchestrators are latest-only — no immutable history and no ``snapshot_trigger``. Mirrors the
+TypeScript SDK's ``MultiAgentSaveLatestStrategy``.
+"""
+
+_MULTI_AGENT_SAVE_LATEST_STRATEGIES = get_args(MultiAgentSaveLatestStrategy)
+
 # Top-level storage namespace for all session data. Byte-identical to the TypeScript SDK,
 # which namespaces its unified storage under "session" (singular) before the session id, so
 # the on-disk key layout is shared across SDKs. The manager applies this namespace once (unless
@@ -139,6 +156,16 @@ def _snapshot_key(session_id: str, agent_id: str, *, snapshot_id: str | None) ->
     return f"{prefix}{_IMMUTABLE_HISTORY}/snapshot_{snapshot_id}.json"
 
 
+def _multi_agent_latest_key(session_id: str, orchestrator_id: str) -> str:
+    """Return the ``snapshot_latest`` key for an orchestrator.
+
+    Orchestrators are latest-only (no immutable history), byte-identical to the TypeScript SDK's
+    ``multiAgent`` scope: ``session/<session_id>/scopes/multiAgent/<orchestrator_id>/snapshots/``.
+    """
+    orchestrator_id = validate_identifier(orchestrator_id, Identifier.AGENT)
+    return f"{_session_prefix(session_id)}scopes/multiAgent/{orchestrator_id}/snapshots/{_SNAPSHOT_LATEST}"
+
+
 def _serialize_snapshot(snapshot: Snapshot) -> bytes:
     """Serialize a snapshot to JSON bytes, base64-encoding any bytes content."""
     return json.dumps(encode_bytes_values(snapshot.to_dict()), ensure_ascii=False).encode("utf-8")
@@ -188,8 +215,10 @@ class SnapshotSessionManager(SessionManager):
     overwritten. When ``snapshot_trigger`` returns True after an invocation, an
     additional immutable snapshot is appended for time-travel restore.
 
-    Single agents only. Attaching this manager to a Graph or Swarm raises
-    ``NotImplementedError``; use a message-log session manager for orchestrators.
+    Single agents get immutable time-travel snapshots via ``snapshot_trigger``. Graph and Swarm
+    orchestrators are persisted latest-only: state is captured after each node (or each invocation,
+    per ``multi_agent_save_latest_on``) and restored lazily on their first invocation. Attaching
+    this manager to a BidiAgent raises ``NotImplementedError``.
 
     Example:
         ```python
@@ -208,6 +237,7 @@ def __init__(
         *,
         storage: Storage | None = None,
         save_latest_on: SaveLatestStrategy = "invocation",
+        multi_agent_save_latest_on: MultiAgentSaveLatestStrategy = "node",
         snapshot_trigger: SnapshotTrigger | None = None,
         **kwargs: Any,
     ) -> None:
@@ -220,6 +250,8 @@ def __init__(
                 agent-level storage is available, falls back to
                 :class:`~strands.storage.local_file_storage.LocalFileStorage`.
             save_latest_on: When to overwrite ``snapshot_latest``. See :data:`SaveLatestStrategy`.
+            multi_agent_save_latest_on: For Graph/Swarm orchestrators, when to overwrite the
+                orchestrator's ``snapshot_latest``. See :data:`MultiAgentSaveLatestStrategy`.
             snapshot_trigger: Optional callback invoked after each invocation; when it
                 returns True an immutable snapshot is appended for checkpointing. An immutable
                 snapshot can also be forced at any point via :meth:`save_snapshot`.
@@ -240,9 +272,18 @@ def __init__(
             # Silently accepting an unknown value would register no save hooks — the session
             # would persist nothing with no error.
             raise ValueError(f"save_latest_on must be one of {_SAVE_LATEST_STRATEGIES}, got {save_latest_on!r}")
+        if multi_agent_save_latest_on not in _MULTI_AGENT_SAVE_LATEST_STRATEGIES:
+            raise ValueError(
+                f"multi_agent_save_latest_on must be one of {_MULTI_AGENT_SAVE_LATEST_STRATEGIES}, "
+                f"got {multi_agent_save_latest_on!r}"
+            )
         self._storage: Storage | None = _resolve_storage(storage) if storage is not None else None
         self._save_latest_on: SaveLatestStrategy = save_latest_on
+        self._multi_agent_save_latest_on: MultiAgentSaveLatestStrategy = multi_agent_save_latest_on
         self._snapshot_trigger = snapshot_trigger
+        # Orchestrator ids restored this process, so restore runs once per orchestrator (lazily,
+        # on its first invocation) rather than on every invocation.
+        self._multi_agent_restored_ids: set[str] = set()
 
     @property
     def _resolved_storage(self) -> Storage:
@@ -269,19 +310,62 @@ def register_hooks(self, registry: HookRegistry, **kwargs: Any) -> None:
             registry.add_callback(MessageAddedEvent, self._on_message_added)
         registry.add_callback(AfterInvocationEvent, self._on_after_invocation)
 
-        # Fail loudly rather than silently persisting nothing: this manager handles single agents
-        # only, so an orchestrator or BidiAgent must not be able to attach it and appear to be
-        # persisted. Both are rejected at their initialization event, before any turn runs.
-        registry.add_callback(MultiAgentInitializedEvent, self._reject_multi_agent)
+        # An orchestrator has no AgentInitializedEvent to lazily resolve storage from, so its hooks
+        # are wired at its own init event. A BidiAgent is still unsupported and rejected there.
+        registry.add_callback(MultiAgentInitializedEvent, self._init_multi_agent)
         registry.add_callback(BidiAgentInitializedEvent, self._reject_bidi_agent)
 
-    def _reject_multi_agent(self, event: MultiAgentInitializedEvent) -> None:
-        """Raise on orchestrator init; multi-agent snapshot persistence is not supported yet."""
-        raise NotImplementedError(
-            f"{type(self).__name__} does not support multi-agent (Graph/Swarm) persistence. "
-            "Use a message-log session manager (FileSessionManager, S3SessionManager) for "
-            "orchestrators."
-        )
+    def _init_multi_agent(self, event: MultiAgentInitializedEvent) -> None:
+        """Wire orchestrator snapshot persistence at init.
+
+        Registers a lazy restore before the first invocation and save handlers after work
+        completes. No state is captured here — a freshly built orchestrator has none. Storage
+        must come from the constructor, since an orchestrator has no agent to resolve it from.
+        """
+        orchestrator = event.source
+        if self._storage is None:
+            raise RuntimeError(
+                "SnapshotSessionManager requires a storage backend for multi-agent orchestrators. "
+                "Provide storage in the constructor."
+            )
+        orchestrator.add_hook(self._on_before_multi_agent_invocation, BeforeMultiAgentInvocationEvent)
+        if self._multi_agent_save_latest_on == "node":
+            orchestrator.add_hook(self._on_after_node_call, AfterNodeCallEvent)
+        orchestrator.add_hook(self._on_after_multi_agent_invocation, AfterMultiAgentInvocationEvent)
+
+    async def _on_before_multi_agent_invocation(self, event: BeforeMultiAgentInvocationEvent) -> None:
+        """Restore orchestrator state once, on its first invocation."""
+        orchestrator = event.source
+        if orchestrator.id in self._multi_agent_restored_ids:
+            return
+        self._multi_agent_restored_ids.add(orchestrator.id)
+        await self._restore_multi_agent(orchestrator)
+
+    async def _on_after_node_call(self, event: AfterNodeCallEvent) -> None:
+        """Save latest orchestrator snapshot after each node completes."""
+        await self._save_multi_agent_latest(event.source)
+
+    async def _on_after_multi_agent_invocation(self, event: AfterMultiAgentInvocationEvent) -> None:
+        """Save latest orchestrator snapshot after the invocation completes."""
+        await self._save_multi_agent_latest(event.source)
+
+    async def _restore_multi_agent(self, orchestrator: "MultiAgentBase") -> bool:
+        """Load an orchestrator's latest snapshot into it. Returns False if none exists."""
+        from ..multiagent.snapshot import load_snapshot
+
+        key = _multi_agent_latest_key(self.session_id, orchestrator.id)
+        data = await self._resolved_storage.read(key)
+        if data is None:
+            return False
+        load_snapshot(orchestrator, _deserialize_snapshot(data))
+        return True
+
+    async def _save_multi_agent_latest(self, orchestrator: "MultiAgentBase") -> None:
+        """Capture the orchestrator and overwrite its ``snapshot_latest``."""
+        from ..multiagent.snapshot import take_snapshot
+
+        data = _serialize_snapshot(take_snapshot(orchestrator))
+        await self._resolved_storage.write(_multi_agent_latest_key(self.session_id, orchestrator.id), data)
 
     def _reject_bidi_agent(self, event: BidiAgentInitializedEvent) -> None:
         """Raise on BidiAgent init; bidirectional-streaming snapshot persistence is not supported yet."""

strands-py/src/strands/types/_snapshot.py

@@ -17,7 +17,7 @@
     "model_state",
 ]
 SnapshotPreset = Literal["session"]
-Scope = Literal["agent"]
+Scope = Literal["agent", "multiAgent"]
 
 ALL_SNAPSHOT_FIELDS: tuple[SnapshotField, ...] = (
     "messages",
@@ -28,7 +28,7 @@
     "model_state",
 )
 
-VALID_SCOPES: tuple[Scope, ...] = ("agent",)
+VALID_SCOPES: tuple[Scope, ...] = ("agent", "multiAgent")
 
 SNAPSHOT_SCHEMA_VERSION = "1.0"
 

strands-py/tests/strands/agent/test_snapshot.py

@@ -2,7 +2,7 @@
 
 import json
 import re
-from typing import Any
+from typing import Any, get_args
 from unittest.mock import MagicMock
 
 import pytest
@@ -13,6 +13,7 @@
     SNAPSHOT_PRESETS,
     SNAPSHOT_SCHEMA_VERSION,
     VALID_SCOPES,
+    Scope,
     Snapshot,
     resolve_snapshot_fields,
 )
@@ -300,7 +301,7 @@ def test_take_snapshot_app_data_is_independent_copy():
 
 def test_valid_scopes_constant_matches_scope_type():
     """VALID_SCOPES contains exactly the values from the Scope Literal type."""
-    assert set(VALID_SCOPES) == {"agent"}
+    assert set(VALID_SCOPES) == set(get_args(Scope))
 
 
 def test_snapshot_validate_accepts_valid_scopes():