diff --git a/strands-py/src/strands/_identifier.py b/strands-py/src/strands/_identifier.py index e8b12635ce..91541e9684 100644 --- a/strands-py/src/strands/_identifier.py +++ b/strands-py/src/strands/_identifier.py @@ -2,6 +2,10 @@ import enum import os +import re +import threading +import time +import uuid class Identifier(enum.Enum): @@ -28,3 +32,63 @@ def validate(id_: str, type_: Identifier) -> str: raise ValueError(f"{type_.value}_id={id_} | id cannot contain path separators") return id_ + + +# UUIDv7 monotonic counter state (RFC 9562 method 2), persisted across calls so that ids minted +# inside the same millisecond are still ordered by an incrementing counter. +_uuid7_state_lock = threading.Lock() +_uuid7_last_timestamp_ms = -1 +_uuid7_intra_ms_counter = 0 + + +def new_uuid7() -> str: + """Return a fresh, monotonic UUID version 7 (RFC 9562). + + The 48-bit millisecond timestamp in the high bits makes lexicographic order equal creation + order, so a set of these ids can be listed oldest-first by sorting alone — no separate index. + Monotonic across calls, including bursts that mint more than 4096 ids inside one millisecond. + + ``uuid.uuid7`` is only available in Python 3.14+, below the SDK's floor, hence this + implementation. + """ + global _uuid7_last_timestamp_ms, _uuid7_intra_ms_counter + + with _uuid7_state_lock: + now_ms = int(time.time() * 1000) + if now_ms > _uuid7_last_timestamp_ms: + _uuid7_last_timestamp_ms = now_ms + # Seed in the low half of the 12-bit field (fresh randomness for uniqueness) so at + # least 2048 increments fit before overflow. + _uuid7_intra_ms_counter = int.from_bytes(os.urandom(2), "big") & 0x07FF + else: + _uuid7_intra_ms_counter += 1 + if _uuid7_intra_ms_counter > 0x0FFF: + # Counter exhausted within this millisecond: borrow from the clock by advancing + # the timestamp 1ms (RFC 9562 method 1), keeping ids strictly increasing rather + # than wrapping the counter field backwards. + _uuid7_last_timestamp_ms += 1 + _uuid7_intra_ms_counter = int.from_bytes(os.urandom(2), "big") & 0x07FF + counter = _uuid7_intra_ms_counter + timestamp_ms = _uuid7_last_timestamp_ms + + # Layout: 48-bit ms timestamp | version(7) | 12-bit counter | variant(0b10) | 62 random bits. + value = (timestamp_ms & ((1 << 48) - 1)) << 80 + value |= 0x7 << 76 + value |= (counter & 0x0FFF) << 64 + value |= 0b10 << 62 + value |= int.from_bytes(os.urandom(8), "big") & ((1 << 62) - 1) + return str(uuid.UUID(int=value)) + + +# ``\A...\Z`` (not ``^...$``) so a trailing newline cannot slip past into a value derived from +# the id, such as a storage key. +_UUID7_PATTERN = re.compile(r"\A[0-9a-f]{8}-[0-9a-f]{4}-7[0-9a-f]{3}-[89ab][0-9a-f]{3}-[0-9a-f]{12}\Z") + + +def is_uuid7(value: str) -> bool: + """Return True if the string is a canonical lowercase UUID version 7. + + Args: + value: The string to check. + """ + return _UUID7_PATTERN.match(value) is not None diff --git a/strands-py/src/strands/session/__init__.py b/strands-py/src/strands/session/__init__.py index 7b53101908..bd582ca0b9 100644 --- a/strands-py/src/strands/session/__init__.py +++ b/strands-py/src/strands/session/__init__.py @@ -1,6 +1,12 @@ """Session module. This module provides session management functionality. + +``SnapshotSessionManager`` persists a single agent as a versioned snapshot to a unified +:class:`~strands.storage.storage.Storage` backend and supports checkpointing (time-travel +restore). The repository-based managers (``FileSessionManager``, ``S3SessionManager``, +``RepositorySessionManager``) persist each message individually and also cover multi-agent +orchestrators. """ from .file_session_manager import FileSessionManager @@ -8,11 +14,15 @@ from .s3_session_manager import S3SessionManager from .session_manager import SessionManager from .session_repository import SessionRepository +from .snapshot_session_manager import SaveLatestStrategy, SnapshotSessionManager, SnapshotTrigger __all__ = [ "FileSessionManager", "RepositorySessionManager", "S3SessionManager", + "SaveLatestStrategy", "SessionManager", "SessionRepository", + "SnapshotSessionManager", + "SnapshotTrigger", ] diff --git a/strands-py/src/strands/session/snapshot_session_manager.py b/strands-py/src/strands/session/snapshot_session_manager.py new file mode 100644 index 0000000000..ce6a3a7700 --- /dev/null +++ b/strands-py/src/strands/session/snapshot_session_manager.py @@ -0,0 +1,500 @@ +"""Snapshot-based session manager. + +Persists an agent as a single versioned :class:`~strands.types._snapshot.Snapshot` +blob on each lifecycle event, mirroring the TypeScript SDK's ``SessionManager``: + +- A mutable ``snapshot_latest`` is overwritten on each save, for crash/restart resume. +- Append-only immutable snapshots (time-ordered keys) are written when a ``snapshot_trigger`` + fires, enabling checkpointing — restore to any prior state, not just the latest. + +The manager persists snapshots through the unified :class:`~strands.storage.storage.Storage` +primitive (``write``/``read``/``delete``/``list`` over byte blobs). It owns the key layout, +snapshot-id scheme, and serialization; the storage backend only moves bytes, so the same +``Storage`` instance can back sessions, memory, and other subsystems. + +This is distinct from the older message-log session managers +(:class:`~strands.session.repository_session_manager.RepositorySessionManager` and its +subclasses), which persist each message individually. Snapshots capture the whole agent +in one atomic blob and are the recommended path for new agents. +""" + +import asyncio +import json +import logging +import re +from typing import TYPE_CHECKING, Any, Literal, Protocol, get_args, runtime_checkable + +from .._async import run_async +from .._identifier import Identifier, is_uuid7 +from .._identifier import new_uuid7 as _new_snapshot_id +from .._identifier import validate as validate_identifier +from ..experimental.hooks.events import BidiAgentInitializedEvent +from ..hooks.events import ( + AfterInvocationEvent, + AgentInitializedEvent, + MessageAddedEvent, + MultiAgentInitializedEvent, +) +from ..hooks.registry import HookRegistry +from ..storage.local_file_storage import LocalFileStorage +from ..storage.storage import _NAMESPACED, Storage, _NamespacedStorage +from ..types._snapshot import Snapshot +from ..types.content import Message +from ..types.exceptions import SnapshotException +from ..types.session import decode_bytes_values, encode_bytes_values +from .session_manager import SessionManager + +if TYPE_CHECKING: + from ..agent.agent import Agent + +logger = logging.getLogger(__name__) + +SaveLatestStrategy = Literal["message", "invocation", "trigger"] +"""Controls how often ``snapshot_latest`` is saved automatically. + +- ``"invocation"``: after every agent invocation completes (default; balances durability and I/O). +- ``"message"``: after every message added (most durable, highest I/O). +- ``"trigger"``: only when ``snapshot_trigger`` fires (or manually via ``save_snapshot``). + +Guardrail redactions are flushed immediately under every strategy, including ``"trigger"``, +so pre-redaction content never sits at rest. This diverges from the TypeScript SDK, which +does not flush redactions under ``"trigger"``; see :meth:`SnapshotSessionManager.redact_latest_message`. +""" + +# Derived from the Literal above so the accepted runtime values cannot drift from the type. +_SAVE_LATEST_STRATEGIES = get_args(SaveLatestStrategy) + +# 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 +# the caller passed an already-namespaced view) and builds keys relative to it, so a caller who +# pre-namespaces under "session" does not get a doubled "session/session/..." prefix. +_SESSIONS_NAMESPACE = "session" + +_SNAPSHOT_LATEST = "snapshot_latest.json" +_IMMUTABLE_HISTORY = "immutable_history" +_SNAPSHOT_REGEX = re.compile(r"snapshot_([\w-]+)\.json\Z") + +# Cap concurrent deletes so a session with many immutable checkpoints does not spawn thousands +# of simultaneous blocking storage calls (e.g. S3's per-object delete via asyncio.to_thread). +_DELETE_CONCURRENCY = 100 + +# -- Key layout, relative to the "session" storage namespace applied in __init__ -- +# +# /scopes/agent//snapshots/ +# snapshot_latest.json +# immutable_history/snapshot_.json +# +# The namespaced storage view prepends "session/", so the full on-disk key is +# session//... — byte-identical to the TypeScript SDK. These are module-level so the +# migration utility builds the same keys the manager reads. + + +def _resolve_storage(storage: Storage) -> Storage: + """Namespace raw storage under ``"session"``; pass an already-namespaced view through. + + Mirrors the TypeScript SDK: a view the caller already scoped (marked with ``_NAMESPACED``) + is used as-is so its prefix is not doubled, otherwise raw storage is wrapped under the + ``"session"`` namespace. Manager keys are built relative to the result. + """ + if getattr(storage, "_namespaced", None) is _NAMESPACED: + return storage + return _NamespacedStorage(storage, _SESSIONS_NAMESPACE) + + +def _validate_snapshot_id(snapshot_id: str) -> None: + """Validate that a string is an SDK-vended snapshot id (a UUIDv7). + + Snapshot ids are opaque handles callers get from ``list_snapshot_ids`` and never construct. + The fixed shape also guards the immutable-history key against traversal. + + Args: + snapshot_id: The string to validate. + + Raises: + ValueError: If the string is not a valid snapshot id. + """ + if not is_uuid7(snapshot_id): + raise ValueError(f"'{snapshot_id}' is not a valid snapshot id") + + +def _session_prefix(session_id: str) -> str: + """Return the namespace-relative key prefix covering an entire session.""" + session_id = validate_identifier(session_id, Identifier.SESSION) + return f"{session_id}/" + + +def _snapshots_prefix(session_id: str, agent_id: str) -> str: + """Return the namespace-relative key prefix for an agent's snapshots directory.""" + agent_id = validate_identifier(agent_id, Identifier.AGENT) + return f"{_session_prefix(session_id)}scopes/agent/{agent_id}/snapshots/" + + +def _snapshot_key(session_id: str, agent_id: str, *, snapshot_id: str | None) -> str: + """Return the storage key for a snapshot; ``None`` targets ``snapshot_latest``.""" + prefix = _snapshots_prefix(session_id, agent_id) + if snapshot_id is None: + return f"{prefix}{_SNAPSHOT_LATEST}" + _validate_snapshot_id(snapshot_id) + return f"{prefix}{_IMMUTABLE_HISTORY}/snapshot_{snapshot_id}.json" + + +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") + + +def _deserialize_snapshot(data: bytes) -> Snapshot: + """Deserialize JSON bytes into a snapshot, decoding any base64 bytes content. + + Raises: + SnapshotException: If the bytes are not a valid, current-schema snapshot (malformed + JSON, wrong encoding, wrong shape, or an unsupported schema/scope). + """ + try: + decoded = decode_bytes_values(json.loads(data)) + except (ValueError, UnicodeDecodeError) as error: + raise SnapshotException(f"Failed to deserialize snapshot: {error}") from error + if not isinstance(decoded, dict): + raise SnapshotException(f"Snapshot is not an object: got {type(decoded).__name__}") + try: + return Snapshot.from_dict(decoded) + except (KeyError, TypeError, AttributeError) as error: + raise SnapshotException(f"Failed to deserialize snapshot: {error}") from error + + +@runtime_checkable +class SnapshotTrigger(Protocol): + """Decides whether to write an immutable checkpoint after an invocation.""" + + def __call__(self, *, agent_data: "Agent", **kwargs: Any) -> bool: + """Return True to append an immutable snapshot for the given agent. + + Args: + agent_data: The agent that just completed an invocation. + **kwargs: Additional keyword arguments for future extensibility. + + Returns: + True to create an immutable checkpoint, False otherwise. + """ + ... + + +class SnapshotSessionManager(SessionManager): + """Persists agent snapshots to a :class:`~strands.storage.storage.Storage` across invocations. + + On agent initialization the latest snapshot is restored automatically. On each + qualifying lifecycle event the agent is re-captured and ``snapshot_latest`` is + 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. + + Example: + ```python + from strands import Agent + from strands.session import SnapshotSessionManager + from strands.storage import LocalFileStorage + + session = SnapshotSessionManager("my-session", storage=LocalFileStorage()) + agent = Agent(session_manager=session) + ``` + """ + + def __init__( + self, + session_id: str = "default-session", + *, + storage: Storage | None = None, + save_latest_on: SaveLatestStrategy = "invocation", + snapshot_trigger: SnapshotTrigger | None = None, + **kwargs: Any, + ) -> None: + """Initialize the snapshot session manager. + + Args: + session_id: Unique session identifier. Must not contain path separators. + storage: Unified storage backend that persists snapshot blobs. Defaults to + :class:`~strands.storage.local_file_storage.LocalFileStorage`, which writes + under the local filesystem. + save_latest_on: When to overwrite ``snapshot_latest``. See :data:`SaveLatestStrategy`. + 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`. + **kwargs: Additional keyword arguments for future extensibility. + + Raises: + ValueError: If ``session_id`` is empty, is a relative-path segment (``.`` or ``..``), + normalizes to empty, or contains a path separator; or if ``save_latest_on`` is + not a recognized strategy. + """ + self.session_id = validate_identifier(session_id, Identifier.SESSION) + # validate_identifier permits "."/".."/whitespace, which either collapse the session + # prefix (silently writing to a shared/wrong key on which delete_session then no-ops) + # or explode deep in the storage layer. Reject anything that isn't a usable id here. + if not self.session_id.strip() or self.session_id in (".", ".."): + raise ValueError(f"session_id is not a valid session identifier: {session_id!r}") + if save_latest_on not in _SAVE_LATEST_STRATEGIES: + # 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}") + self._storage = _resolve_storage(storage if storage is not None else LocalFileStorage()) + self._save_latest_on: SaveLatestStrategy = save_latest_on + self._snapshot_trigger = snapshot_trigger + + def register_hooks(self, registry: HookRegistry, **kwargs: Any) -> None: + """Register lifecycle callbacks for snapshot persistence. + + Overrides the base wiring: the message-log callbacks are replaced with + snapshot save/restore handlers. + """ + # Restore must be synchronous — AgentInitializedEvent forbids async callbacks. + registry.add_callback(AgentInitializedEvent, lambda event: self.initialize(event.agent)) + + # The save paths run under invoke_callbacks_async, so register them as native + # async handlers and avoid the sync bridge. + if self._save_latest_on == "message": + 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) + 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 _reject_bidi_agent(self, event: BidiAgentInitializedEvent) -> None: + """Raise on BidiAgent init; bidirectional-streaming snapshot persistence is not supported yet.""" + raise NotImplementedError( + f"{type(self).__name__} does not support BidiAgent persistence. " + "Use a message-log session manager (FileSessionManager, S3SessionManager) for " + "bidirectional-streaming agents." + ) + + # -- ABC methods (invoked synchronously by the Agent; bridge to async storage) -- + + def initialize(self, agent: "Agent", **kwargs: Any) -> None: + """Restore the agent from its latest snapshot, if one exists. + + Args: + agent: Agent to restore. + **kwargs: Additional keyword arguments for future extensibility. + """ + run_async(lambda: self._initialize_async(agent)) + + def sync_agent(self, agent: "Agent", **kwargs: Any) -> None: + """Capture the agent and overwrite ``snapshot_latest``. + + Args: + agent: Agent to persist. + **kwargs: Additional keyword arguments for future extensibility. + """ + run_async(lambda: self._save_latest(agent)) + + def redact_latest_message(self, redact_message: Message, agent: "Agent", **kwargs: Any) -> None: + """Persist immediately after a guardrail redaction, under every strategy. + + The Agent has already applied the redaction to ``agent.messages[-1]`` before + calling this, so re-capturing the agent flushes pre-redaction content out of + the persisted latest snapshot. This flush happens regardless of ``save_latest_on`` + (including ``"trigger"``) because the Agent invokes this method directly, not through + a hook the manager could decline to register — so pre-redaction content never sits at + rest. This diverges from the TypeScript SDK, which gates redaction persistence behind + an ``AfterModelCall`` hook it skips under ``"trigger"`` and therefore does not flush there. + + Args: + redact_message: The redacted replacement message (already applied by the Agent). + agent: Agent whose latest message was redacted. + **kwargs: Additional keyword arguments for future extensibility. + """ + run_async(lambda: self._save_latest(agent)) + + def append_message(self, message: Message, agent: "Agent", **kwargs: Any) -> None: + """No-op — snapshots capture the whole agent. + + Per-message persistence under the ``"message"`` strategy is handled by the + ``MessageAddedEvent`` hook, not by this method. + + Args: + message: The message that was appended (unused). + agent: The agent the message was appended to (unused). + **kwargs: Additional keyword arguments for future extensibility. + """ + + # -- Public time-travel API -- + + async def list_snapshot_ids( + self, agent: "Agent", *, limit: int | None = None, start_after: str | None = None + ) -> list[str]: + """List immutable snapshot ids for an agent, oldest first. + + Args: + agent: Agent whose snapshots to list. + limit: Optional cap on the number of ids returned. + start_after: Exclusive cursor; a snapshot id from a prior page. + + Returns: + Immutable snapshot ids in chronological order. + + Raises: + ValueError: If ``start_after`` is not a valid snapshot id. + """ + if limit is not None and limit <= 0: + return [] + if start_after is not None: + _validate_snapshot_id(start_after) + + history_prefix = f"{_snapshots_prefix(self.session_id, agent.agent_id)}{_IMMUTABLE_HISTORY}/" + keys = await self._storage.list(history_prefix) + ids = sorted(match.group(1) for key in keys if (match := _SNAPSHOT_REGEX.search(key))) + if start_after is not None: + ids = [snapshot_id for snapshot_id in ids if snapshot_id > start_after] + if limit is not None: + ids = ids[:limit] + return ids + + async def restore_snapshot(self, agent: "Agent", *, snapshot_id: str | None = None) -> bool: + """Restore an agent from a stored snapshot. + + Args: + agent: Agent to restore into. + snapshot_id: The immutable snapshot id to restore (time travel). Omit to restore + ``snapshot_latest``, the same snapshot restore-on-init loads. + + Returns: + True if the snapshot existed and was restored, False otherwise. + + Raises: + ValueError: If ``snapshot_id`` is given and is not a valid snapshot id. + """ + return await self._restore(agent, snapshot_id=snapshot_id) + + async def save_snapshot(self, agent: "Agent", *, is_latest: bool) -> str | None: + """Save a snapshot of the agent's current state on demand. + + Use ``is_latest=False`` to force an immutable checkpoint at an arbitrary point (independent + of ``snapshot_trigger``), so it can later be restored with :meth:`restore_snapshot`; use + ``is_latest=True`` to overwrite ``snapshot_latest``. + + Args: + agent: Agent whose state to capture. + is_latest: When True, overwrite ``snapshot_latest`` (a single mutable snapshot). When + False, append a new immutable snapshot under a fresh, time-ordered id. + + Returns: + The new immutable snapshot id, ready to pass to :meth:`restore_snapshot`, or ``None`` + when ``is_latest=True`` (``snapshot_latest`` is not addressed by id). + """ + data = _serialize_snapshot(self._capture(agent)) + snapshot_id = None if is_latest else _new_snapshot_id() + await self._storage.write(_snapshot_key(self.session_id, agent.agent_id, snapshot_id=snapshot_id), data) + return snapshot_id + + async def delete_session(self) -> None: + """Delete all snapshots for this session.""" + keys = await self._storage.list(_session_prefix(self.session_id)) + semaphore = asyncio.Semaphore(_DELETE_CONCURRENCY) + + async def _delete(key: str) -> None: + async with semaphore: + await self._storage.delete(key) + + await asyncio.gather(*(_delete(key) for key in keys)) + + # -- Async internals -- + + async def _initialize_async(self, agent: "Agent") -> None: + """Restore latest snapshot on init, warning on overwrite and handling stateful models.""" + had_messages = len(agent.messages) > 0 + restored = await self._restore(agent) + + if restored and had_messages: + logger.warning( + "agent_id=<%s>, session_id=<%s> | agent had existing messages that were overwritten by session restore", + agent.agent_id, + self.session_id, + ) + + # Stateful models manage history server-side, so restored messages would drift + # from the server's view. Keep the restored model_state and drop the messages. + if restored and agent.model.stateful and len(agent.messages) > 0: + logger.warning( + "agent_id=<%s>, message_count=<%s> | discarding restored messages for stateful model", + agent.agent_id, + len(agent.messages), + ) + agent.messages = [] + + async def _restore(self, agent: "Agent", *, snapshot_id: str | None = None) -> bool: + """Load a snapshot into the agent. Returns False if none exists.""" + data = await self._storage.read(_snapshot_key(self.session_id, agent.agent_id, snapshot_id=snapshot_id)) + if data is None: + return False + agent.load_snapshot(_deserialize_snapshot(data)) + return True + + async def _save_latest(self, agent: "Agent") -> None: + """Capture the agent and overwrite ``snapshot_latest``.""" + await self.save_snapshot(agent, is_latest=True) + + async def _save_immutable_and_latest(self, agent: "Agent") -> None: + """Capture once and write the immutable snapshot, then ``snapshot_latest``. + + Ordered, not concurrent: writing the immutable checkpoint first means a partial failure + can leave an orphaned immutable snapshot (harmless — the next list simply includes it) + but never a ``snapshot_latest`` pointing at history that was never written. + """ + data = _serialize_snapshot(self._capture(agent)) + await self._storage.write(_snapshot_key(self.session_id, agent.agent_id, snapshot_id=_new_snapshot_id()), data) + await self._storage.write(_snapshot_key(self.session_id, agent.agent_id, snapshot_id=None), data) + + async def _on_message_added(self, event: MessageAddedEvent) -> None: + """Save latest after each message under the ``"message"`` strategy.""" + await self._save_latest(event.agent) + + async def _on_after_invocation(self, event: AfterInvocationEvent) -> None: + """Save latest on invocation and fire the immutable-checkpoint trigger. + + When the trigger fires, the immutable+latest write subsumes the invocation save, so the + agent is captured only once even under ``save_latest_on="invocation"``. + + ``"message"`` also saves here, not just per message: the Agent runs + ``conversation_manager.apply_management`` (trimming/summarizing) after the last + ``MessageAddedEvent`` but before this event, so the per-message saves would otherwise + persist pre-management messages and a stale ``removed_message_count``. + """ + triggered = False + trigger_failed = False + if self._snapshot_trigger is not None: + try: + triggered = self._snapshot_trigger(agent_data=event.agent) + except Exception: + # A caller's trigger raising must not discard the completed turn's latest save + trigger_failed = True + logger.exception( + "agent_id=<%s>, session_id=<%s> | snapshot_trigger raised; skipping immutable checkpoint", + event.agent.agent_id, + self.session_id, + ) + if triggered: + await self._save_immutable_and_latest(event.agent) + elif trigger_failed or self._save_latest_on in ("invocation", "message"): + await self._save_latest(event.agent) + + def _capture(self, agent: "Agent") -> Snapshot: + """Capture a full session snapshot including the system prompt. + + The shared ``"session"`` preset omits ``system_prompt`` (opt-in for callers like + the goal plugin); session persistence includes it so a rehydrated agent behaves + identically to the original, matching the TypeScript SDK's session preset. + """ + return agent.take_snapshot(preset="session", include=["system_prompt"]) diff --git a/strands-py/tests/strands/session/test_snapshot_session_manager.py b/strands-py/tests/strands/session/test_snapshot_session_manager.py new file mode 100644 index 0000000000..26991ff443 --- /dev/null +++ b/strands-py/tests/strands/session/test_snapshot_session_manager.py @@ -0,0 +1,813 @@ +"""Tests for SnapshotSessionManager.""" + +import asyncio +import tempfile +import uuid +from unittest.mock import AsyncMock, Mock + +import pytest + +from strands.agent.agent import Agent +from strands.agent.conversation_manager.sliding_window_conversation_manager import SlidingWindowConversationManager +from strands.experimental.hooks.events import BidiAgentInitializedEvent +from strands.hooks.registry import HookRegistry +from strands.multiagent import GraphBuilder, Swarm +from strands.session.snapshot_session_manager import ( + SnapshotSessionManager, + _new_snapshot_id, + _session_prefix, + _snapshot_key, +) +from strands.storage import LocalFileStorage +from strands.types.content import ContentBlock +from strands.types.exceptions import ContextWindowOverflowException, SnapshotException +from tests.fixtures.mocked_model_provider import MockedModelProvider + + +@pytest.fixture +def temp_dir(): + """Create a temporary directory for testing.""" + with tempfile.TemporaryDirectory() as temp_dir: + yield temp_dir + + +@pytest.fixture +def storage(temp_dir): + """A file-backed unified storage.""" + return LocalFileStorage(temp_dir) + + +def _model(*texts): + """Build a mock model that replies with the given texts in sequence.""" + return MockedModelProvider([{"role": "assistant", "content": [{"text": text}]} for text in texts]) + + +def _on_disk_key(session_id: str, agent_id: str) -> str: + """The full raw-storage key for a session's latest snapshot (namespace + relative key).""" + return f"session/{_snapshot_key(session_id, agent_id, snapshot_id=None)}" + + +def _texts(agent) -> list[str]: + """Flatten an agent's message text content, for asserting which turns are present.""" + return [content["text"] for message in agent.messages for content in message["content"] if "text" in content] + + +def test_new_session_starts_empty(storage): + """A brand-new session leaves a fresh agent's messages untouched.""" + manager = SnapshotSessionManager("s1", storage=storage) + agent = Agent(model=_model("hi"), session_manager=manager, agent_id="a1") + + assert agent.messages == [] + + +def test_empty_session_id_is_rejected(storage): + """An empty session id is rejected; otherwise its prefix would broaden to all sessions.""" + with pytest.raises(ValueError, match="not a valid session identifier"): + SnapshotSessionManager("", storage=storage) + + +@pytest.mark.parametrize("bad_id", [".", "..", " "]) +def test_relative_or_blank_session_id_is_rejected(storage, bad_id): + """'.'/'..'/whitespace pass validate_identifier but collapse or explode the key — reject them.""" + with pytest.raises(ValueError, match="not a valid session identifier"): + SnapshotSessionManager(bad_id, storage=storage) + + +def test_unknown_save_latest_on_is_rejected(storage): + """A mistyped save_latest_on is rejected, rather than silently persisting nothing.""" + with pytest.raises(ValueError, match="save_latest_on must be one of"): + SnapshotSessionManager("s1", storage=storage, save_latest_on="Invocation") # type: ignore[arg-type] + + +def test_graph_is_rejected_rather_than_silently_not_persisted(storage): + """Attaching this single-agent manager to a Graph fails loudly instead of persisting nothing.""" + builder = GraphBuilder() + builder.add_node(Agent(model=_model("done"), agent_id="n1"), "n1") + builder.set_session_manager(SnapshotSessionManager("g1", storage=storage)) + with pytest.raises(NotImplementedError, match="does not support multi-agent"): + builder.build() + + +def test_swarm_is_rejected_rather_than_silently_not_persisted(storage): + """Attaching this single-agent manager to a Swarm fails loudly instead of persisting nothing.""" + with pytest.raises(NotImplementedError, match="does not support multi-agent"): + Swarm( + nodes=[Agent(model=_model("done"), agent_id="n1")], + session_manager=SnapshotSessionManager("sw1", storage=storage), + ) + + +def test_bidi_agent_is_rejected_rather_than_silently_not_persisted(storage): + """Attaching this single-agent manager to a BidiAgent fails loudly instead of persisting nothing. + + The base SessionManager wires BidiAgent hooks to message-log methods; this manager replaces + that wiring, so without an explicit rejection a BidiAgent would appear persisted while + nothing was ever written. + """ + manager = SnapshotSessionManager("b1", storage=storage) + registry = HookRegistry() + manager.register_hooks(registry) + + with pytest.raises(NotImplementedError, match="does not support BidiAgent"): + registry.invoke_callbacks(BidiAgentInitializedEvent(agent=Mock())) + + +def test_child_agent_session_manager_still_blocked(storage): + """Child agents inside a Graph still may not carry their own session manager.""" + builder = GraphBuilder() + child = Agent(model=_model("hi"), agent_id="n1", session_manager=SnapshotSessionManager("child", storage=storage)) + with pytest.raises(ValueError, match="not supported for Graph"): + builder.add_node(child, "n1") + + +def test_raising_snapshot_trigger_still_saves_latest(storage): + """A snapshot_trigger that raises does not discard the completed turn's latest save.""" + + def boom(*, agent_data, **kwargs): + raise RuntimeError("trigger blew up") + + manager = SnapshotSessionManager("s1", storage=storage, snapshot_trigger=boom) + agent = Agent(model=_model("saved"), session_manager=manager, agent_id="a1") + agent("go") # trigger raises here, but the invocation-end latest save must still happen + + manager_2 = SnapshotSessionManager("s1", storage=storage) + agent_2 = Agent(model=_model("x"), session_manager=manager_2, agent_id="a1") + tru_texts = [content["text"] for message in agent_2.messages for content in message["content"] if "text" in content] + assert "go" in tru_texts # the turn survived despite the raising trigger + + +def test_raising_snapshot_trigger_still_saves_latest_under_trigger_strategy(storage): + """Under ``save_latest_on="trigger"`` a raising trigger must not lose the turn entirely. + + The trigger is the only save under this strategy, so a failing trigger has to fall back to a + latest save; otherwise the whole invocation is silently dropped with only a log line. + """ + + def boom(*, agent_data, **kwargs): + raise RuntimeError("trigger blew up") + + manager = SnapshotSessionManager("s1", storage=storage, save_latest_on="trigger", snapshot_trigger=boom) + agent = Agent(model=_model("saved"), session_manager=manager, agent_id="a1") + agent("go") + + manager_2 = SnapshotSessionManager("s1", storage=storage, save_latest_on="trigger") + agent_2 = Agent(model=_model("x"), session_manager=manager_2, agent_id="a1") + tru_texts = [content["text"] for message in agent_2.messages for content in message["content"] if "text" in content] + assert "go" in tru_texts + + +@pytest.mark.asyncio +async def test_empty_session_id_cannot_delete_other_sessions(temp_dir): + """Guard against the destructive prefix broadening: an empty id must not reach delete_session.""" + storage = LocalFileStorage(temp_dir) + # Populate an unrelated, real session. + other = SnapshotSessionManager("real-session", storage=storage) + Agent(model=_model("hi"), session_manager=other, agent_id="a1")("keep me") + assert await storage.read(_on_disk_key("real-session", "a1")) is not None + + # Constructing with an empty id must fail rather than yield a manager whose delete_session + # would list/delete the whole "session/" namespace (every session). + with pytest.raises(ValueError, match="not a valid session identifier"): + SnapshotSessionManager("", storage=storage) + + # The unrelated session is untouched. + assert await storage.read(_on_disk_key("real-session", "a1")) is not None + + +def test_restore_across_instances(storage): + """A fresh agent with the same session id rehydrates prior conversation.""" + manager = SnapshotSessionManager("s1", storage=storage) + agent = Agent(model=_model("The answer is 42."), session_manager=manager, agent_id="a1") + agent("What is the answer?") + + # Simulate process restart: new manager + new agent over the same storage. + manager_2 = SnapshotSessionManager("s1", storage=storage) + agent_2 = Agent(model=_model("Still 42."), session_manager=manager_2, agent_id="a1") + + tru_texts = [content["text"] for message in agent_2.messages for content in message["content"] if "text" in content] + assert "What is the answer?" in tru_texts + assert "The answer is 42." in tru_texts + + +def test_restore_warns_on_overwrite(storage, caplog): + """Restoring over an agent that already had messages logs a warning.""" + manager = SnapshotSessionManager("s1", storage=storage) + agent = Agent(model=_model("saved"), session_manager=manager, agent_id="a1") + agent("first turn") + + manager_2 = SnapshotSessionManager("s1", storage=storage) + Agent( + model=_model("x"), + session_manager=manager_2, + agent_id="a1", + messages=[{"role": "user", "content": [{"text": "pre-existing"}]}], + ) + + assert "overwritten by session restore" in caplog.text + + +def test_state_round_trips(storage): + """Agent state persists and restores across instances.""" + manager = SnapshotSessionManager("s1", storage=storage) + agent = Agent(model=_model("ok"), session_manager=manager, agent_id="a1") + agent.state.set("favorite", "blue") + agent("remember my favorite") + + manager_2 = SnapshotSessionManager("s1", storage=storage) + agent_2 = Agent(model=_model("ok"), session_manager=manager_2, agent_id="a1") + + assert agent_2.state.get("favorite") == "blue" + + +def test_system_prompt_round_trips(storage): + """The system prompt persists and restores (session preset opt-in).""" + manager = SnapshotSessionManager("s1", storage=storage) + agent = Agent( + model=_model("ok"), session_manager=manager, agent_id="a1", system_prompt="You are a helpful assistant." + ) + agent("hi") + + manager_2 = SnapshotSessionManager("s1", storage=storage) + agent_2 = Agent(model=_model("ok"), session_manager=manager_2, agent_id="a1") + + assert agent_2.system_prompt == "You are a helpful assistant." + + +def test_bytes_content_round_trips(storage): + """Image bytes in messages survive JSON serialization via base64 encoding.""" + manager = SnapshotSessionManager("s1", storage=storage) + agent = Agent(model=_model("saw it"), session_manager=manager, agent_id="a1") + image_block: ContentBlock = {"image": {"format": "png", "source": {"bytes": b"\x89PNG\r\n\x1a\n"}}} + agent([image_block]) + + manager_2 = SnapshotSessionManager("s1", storage=storage) + agent_2 = Agent(model=_model("x"), session_manager=manager_2, agent_id="a1") + + tru_bytes = agent_2.messages[0]["content"][0]["image"]["source"]["bytes"] + assert tru_bytes == b"\x89PNG\r\n\x1a\n" + + +@pytest.mark.asyncio +async def test_save_latest_on_message_writes_each_message(temp_dir): + """The ``message`` strategy persists after every message added.""" + storage = LocalFileStorage(temp_dir) + save_keys = [] + original = storage.write + + async def _spy(key, data, **kwargs): + save_keys.append(key) + await original(key, data, **kwargs) + + storage.write = _spy # type: ignore[method-assign] + + manager = SnapshotSessionManager("s1", storage=storage, save_latest_on="message") + agent = Agent(model=_model("reply"), session_manager=manager, agent_id="a1") + save_keys.clear() + await agent.invoke_async("hello") + + # Two messages added (user + assistant) each trigger a per-message save, plus one final + # invocation-end save that captures post-conversation-management state. + assert len(save_keys) == 3 + assert all(key.endswith("snapshot_latest.json") for key in save_keys) + + +@pytest.mark.asyncio +async def test_message_mode_persists_post_management_state(temp_dir): + """``message`` mode restores the trimmed conversation, not the pre-management one. + + The Agent runs conversation management after the last MessageAddedEvent but before the + AfterInvocationEvent, so per-message saves alone would persist untrimmed messages and a + stale removed_message_count. + """ + storage = LocalFileStorage(temp_dir) + manager = SnapshotSessionManager("s1", storage=storage, save_latest_on="message") + agent = Agent( + model=_model("a1", "a2"), + session_manager=manager, + agent_id="a1", + conversation_manager=SlidingWindowConversationManager(window_size=2), + ) + agent("u1") + agent("u2") + + # The live agent has been trimmed to the window and tracks the removed count. + assert agent.conversation_manager.removed_message_count > 0 + live_texts = [content["text"] for message in agent.messages for content in message["content"] if "text" in content] + + manager_2 = SnapshotSessionManager("s1", storage=storage, save_latest_on="message") + agent_2 = Agent( + model=_model("x"), + session_manager=manager_2, + agent_id="a1", + conversation_manager=SlidingWindowConversationManager(window_size=2), + ) + restored_texts = [ + content["text"] for message in agent_2.messages for content in message["content"] if "text" in content + ] + + assert restored_texts == live_texts + assert agent_2.conversation_manager.removed_message_count == agent.conversation_manager.removed_message_count + + +@pytest.mark.asyncio +async def test_invocation_strategy_saves_once_not_per_message(temp_dir): + """The default ``invocation`` strategy saves once at invocation end, not per message.""" + storage = LocalFileStorage(temp_dir) + save_keys = [] + original = storage.write + + async def _spy(key, data, **kwargs): + save_keys.append(key) + await original(key, data, **kwargs) + + storage.write = _spy # type: ignore[method-assign] + + # Default save_latest_on="invocation": MessageAddedEvent must not be registered. + manager = SnapshotSessionManager("s1", storage=storage) + agent = Agent(model=_model("reply"), session_manager=manager, agent_id="a1") + save_keys.clear() + await agent.invoke_async("hello") + + # One save at invocation end, despite two messages being added during the turn. + assert len(save_keys) == 1 + assert save_keys[0].endswith("snapshot_latest.json") + + +def test_snapshot_trigger_creates_immutable(storage): + """When the trigger fires, an immutable snapshot is appended.""" + manager = SnapshotSessionManager("s1", storage=storage, snapshot_trigger=lambda *, agent_data, **_: True) + agent = Agent(model=_model("turn one"), session_manager=manager, agent_id="a1") + agent("go") + + ids = asyncio.run(manager.list_snapshot_ids(agent)) + assert len(ids) == 1 + + +@pytest.mark.asyncio +async def test_save_snapshot_forces_immutable_checkpoint_without_a_trigger(storage): + """save_snapshot(is_latest=False) appends an immutable checkpoint on demand and is restorable.""" + manager = SnapshotSessionManager("s1", storage=storage) # no snapshot_trigger + agent = Agent(model=_model("first", "second"), session_manager=manager, agent_id="a1") + agent("turn 1") + + # No trigger fired, so nothing immutable exists yet. + assert await manager.list_snapshot_ids(agent) == [] + + checkpoint_id = await manager.save_snapshot(agent, is_latest=False) + agent("turn 2") + + ids = await manager.list_snapshot_ids(agent) + assert len(ids) == 1 # the manual checkpoint, not the second turn + assert checkpoint_id == ids[0] # the returned id addresses the snapshot just written + + # The returned id restores that checkpoint directly, with no list_snapshot_ids round trip. + restored = SnapshotSessionManager("s1", storage=storage) + agent_2 = Agent(model=_model("x"), session_manager=restored, agent_id="a1") + assert await restored.restore_snapshot(agent_2, snapshot_id=checkpoint_id) is True + tru_texts = [content["text"] for message in agent_2.messages for content in message["content"] if "text" in content] + assert "turn 1" in tru_texts + assert "turn 2" not in tru_texts + + +@pytest.mark.asyncio +async def test_save_snapshot_is_latest_overwrites_latest_only(storage): + """save_snapshot(is_latest=True) overwrites snapshot_latest and appends no immutable snapshot.""" + manager = SnapshotSessionManager("s1", storage=storage, save_latest_on="trigger") + agent = Agent(model=_model("only turn"), session_manager=manager, agent_id="a1") + agent("go") + + assert await manager.save_snapshot(agent, is_latest=True) is None # latest has no id + + assert await manager.list_snapshot_ids(agent) == [] + assert await storage.read(_on_disk_key("s1", "a1")) is not None + + +@pytest.mark.asyncio +async def test_triggered_turn_captures_and_writes_latest_once(temp_dir): + """A triggered turn under ``invocation`` writes latest once (immutable + latest), not twice.""" + storage = LocalFileStorage(temp_dir) + latest_writes = [] + original = storage.write + + async def _spy(key, data, **kwargs): + if key.endswith("snapshot_latest.json"): + latest_writes.append(key) + await original(key, data, **kwargs) + + storage.write = _spy # type: ignore[method-assign] + + # Default save_latest_on="invocation" with a trigger that always fires. + manager = SnapshotSessionManager("s1", storage=storage, snapshot_trigger=lambda *, agent_data, **_: True) + agent = Agent(model=_model("reply"), session_manager=manager, agent_id="a1") + latest_writes.clear() + await agent.invoke_async("go") + + # The immutable+latest write subsumes the invocation save: one latest write, not two. + assert len(latest_writes) == 1 + ids = await manager.list_snapshot_ids(agent) + assert len(ids) == 1 + + +def test_time_travel_restore(storage): + """restore_snapshot rewinds an agent to an earlier immutable checkpoint.""" + manager = SnapshotSessionManager("s1", storage=storage, snapshot_trigger=lambda *, agent_data, **_: True) + agent = Agent(model=_model("first", "second"), session_manager=manager, agent_id="a1") + agent("turn 1") + agent("turn 2") + + ids = asyncio.run(manager.list_snapshot_ids(agent)) + assert len(ids) == 2 + + restored = asyncio.run(manager.restore_snapshot(agent, snapshot_id=ids[0])) + assert restored is True + + tru_texts = [content["text"] for message in agent.messages for content in message["content"] if "text" in content] + assert "turn 1" in tru_texts + assert "turn 2" not in tru_texts + + +@pytest.mark.asyncio +async def test_restore_snapshot_without_id_restores_latest(storage): + """Omitting snapshot_id restores ``snapshot_latest``, undoing an in-memory time-travel rewind.""" + manager = SnapshotSessionManager("s1", storage=storage, snapshot_trigger=lambda *, agent_data, **_: True) + agent = Agent(model=_model("first", "second"), session_manager=manager, agent_id="a1") + agent("turn 1") + agent("turn 2") + + ids = await manager.list_snapshot_ids(agent) + assert await manager.restore_snapshot(agent, snapshot_id=ids[0]) is True # rewind to turn 1 + assert "turn 2" not in _texts(agent) + + # No id: back to latest, which still holds both turns. + assert await manager.restore_snapshot(agent) is True + tru_texts = _texts(agent) + assert "turn 1" in tru_texts + assert "turn 2" in tru_texts + + +@pytest.mark.asyncio +async def test_restore_snapshot_without_id_returns_false_for_new_session(storage): + """Omitting snapshot_id on a session that has never been saved reports no snapshot.""" + manager = SnapshotSessionManager("never-saved", storage=storage) + agent = Agent(model=_model("hi"), agent_id="a1") + + assert await manager.restore_snapshot(agent) is False + + +def _stateful_model(*texts): + """Build a mock model that reports itself as stateful (server-managed history).""" + model = _model(*texts) + # Stateful models manage conversation history server-side; the constructor swaps in a + # NullConversationManager for them, so both the saved and restored agents match. + object.__setattr__(model, "_force_stateful", True) + type(model).stateful = property(lambda self: getattr(self, "_force_stateful", False)) + return model + + +def test_stateful_model_discards_restored_messages(storage): + """Restore keeps model_state but drops messages for a stateful model.""" + original = _stateful_model("hi") + try: + manager = SnapshotSessionManager("s1", storage=storage) + agent = Agent(model=original, session_manager=manager, agent_id="a1") + # Persist a snapshot that contains messages (a stateful model normally clears + # local history mid-turn, so set them explicitly to exercise the discard branch). + agent.messages = [{"role": "user", "content": [{"text": "hello"}]}] + manager.sync_agent(agent) + + manager_2 = SnapshotSessionManager("s1", storage=storage) + agent_2 = Agent(model=_stateful_model("x"), session_manager=manager_2, agent_id="a1") + assert agent_2.messages == [] + finally: + del type(original).stateful + + +def test_redaction_flush_persists_redacted_content(temp_dir): + """A guardrail redaction is flushed to the latest snapshot immediately.""" + storage = LocalFileStorage(temp_dir) + manager = SnapshotSessionManager("s1", storage=storage) + redaction_model = MockedModelProvider( + [{"redactedUserContent": "REDACTED", "redactedAssistantContent": "I can't help with that."}] + ) + agent = Agent(model=redaction_model, session_manager=manager, agent_id="a1") + agent("sensitive prompt") + + manager_2 = SnapshotSessionManager("s1", storage=storage) + agent_2 = Agent(model=_model("x"), session_manager=manager_2, agent_id="a1") + + tru_texts = [content["text"] for message in agent_2.messages for content in message["content"] if "text" in content] + assert "sensitive prompt" not in tru_texts + assert "REDACTED" in tru_texts + + +@pytest.mark.asyncio +async def test_delete_session_removes_snapshots(storage): + """delete_session clears persisted snapshots.""" + manager = SnapshotSessionManager("s1", storage=storage, snapshot_trigger=lambda *, agent_data, **_: True) + agent = Agent(model=_model("hi"), session_manager=manager, agent_id="a1") + agent("go") + + # Seed a key under the session's namespace to confirm delete clears the whole subtree. + assert await storage.read(_on_disk_key("s1", "a1")) is not None + + await manager.delete_session() + + assert await storage.read(_on_disk_key("s1", "a1")) is None + assert await storage.list(f"session/{_session_prefix('s1')}") == [] + + +def test_restore_by_id_missing_returns_false(storage): + """Restoring a non-existent immutable snapshot returns False.""" + manager = SnapshotSessionManager("s1", storage=storage) + agent = Agent(model=_model("hi"), session_manager=manager, agent_id="a1") + + assert asyncio.run(manager.restore_snapshot(agent, snapshot_id=_new_snapshot_id())) is False + + +def test_trigger_strategy_skips_latest_without_trigger(temp_dir): + """Under ``trigger`` strategy with no trigger, nothing is persisted on invocation.""" + storage = LocalFileStorage(temp_dir) + storage.write = AsyncMock() # type: ignore[method-assign] + + manager = SnapshotSessionManager("s1", storage=storage, save_latest_on="trigger") + agent = Agent(model=_model("hi"), session_manager=manager, agent_id="a1") + storage.write.reset_mock() + agent("go") + + storage.write.assert_not_called() + + +def test_no_warning_when_restoring_into_empty_agent(storage, caplog): + """Restoring into a fresh agent with no messages does not log the overwrite warning.""" + manager = SnapshotSessionManager("s1", storage=storage) + agent = Agent(model=_model("saved"), session_manager=manager, agent_id="a1") + agent("first turn") + + # The second agent starts empty, so restore should populate it silently. + manager_2 = SnapshotSessionManager("s1", storage=storage) + agent_2 = Agent(model=_model("x"), session_manager=manager_2, agent_id="a1") + + assert len(agent_2.messages) > 0 + assert "overwritten by session restore" not in caplog.text + + +class _OverflowThenAnswerModel(MockedModelProvider): + """A model that raises a context-overflow on its first stream, then answers normally. + + This drives the Agent's reactive overflow-recovery path (reduce_context), which in turn + makes the direct ``session_manager.sync_agent`` call at the overflow catch site. + """ + + def __init__(self, *texts): + super().__init__([{"role": "assistant", "content": [{"text": text}]} for text in texts]) + self._overflowed = False + + async def stream(self, *args, **kwargs): + if not self._overflowed: + self._overflowed = True + raise ContextWindowOverflowException("Input is too long for requested model") + async for event in super().stream(*args, **kwargs): + yield event + + +def test_context_overflow_syncs_reduced_conversation(storage): + """A context-window overflow persists the reduced conversation via the direct sync_agent call. + + On ContextWindowOverflowException the Agent calls ``session_manager.sync_agent(agent)`` + directly (outside the hook system) after trimming context. To prove this specific path — + and not the invocation-end save — the manager runs under ``save_latest_on="trigger"`` with + no trigger, so ``_on_after_invocation`` writes nothing and ``MessageAddedEvent`` is not + registered. The only thing that can persist ``snapshot_latest`` is the overflow-time + ``sync_agent``. Deleting the direct call at the overflow catch site would fail this test. + """ + # window_size=2 forces the reactive trim to actually drop the seeded backlog. + conversation_manager = SlidingWindowConversationManager(window_size=2) + manager = SnapshotSessionManager("s1", storage=storage, save_latest_on="trigger") + + sync_calls = [] + original_sync = manager.sync_agent + + def _spy_sync(agent, **kwargs): + sync_calls.append(len(agent.messages)) + return original_sync(agent, **kwargs) + + manager.sync_agent = _spy_sync # type: ignore[method-assign] + + seeded = [ + {"role": "user", "content": [{"text": "one"}]}, + {"role": "assistant", "content": [{"text": "1"}]}, + {"role": "user", "content": [{"text": "two"}]}, + {"role": "assistant", "content": [{"text": "2"}]}, + ] + agent = Agent( + model=_OverflowThenAnswerModel("recovered."), + session_manager=manager, + agent_id="a1", + conversation_manager=conversation_manager, + messages=seeded, + ) + agent("three") + + # The overflow catch site invoked sync_agent exactly once, on the trimmed conversation. + assert len(sync_calls) == 1 + + # A fresh instance restores only what the overflow-path sync persisted: the reduced + # window, not the full seeded backlog (proving the reduced conversation was the thing saved). + manager_2 = SnapshotSessionManager("s1", storage=storage, save_latest_on="trigger") + agent_2 = Agent( + model=_model("x"), + session_manager=manager_2, + agent_id="a1", + conversation_manager=SlidingWindowConversationManager(window_size=2), + ) + + tru_texts = [content["text"] for message in agent_2.messages for content in message["content"] if "text" in content] + assert "one" not in tru_texts # the oldest seeded messages were trimmed before the sync + assert tru_texts # but the reduced conversation was persisted (not empty) + + +def test_redaction_flushes_even_under_trigger_strategy(temp_dir): + """A guardrail redaction is flushed immediately even under the ``trigger`` strategy. + + This is a deliberate divergence from the TypeScript SDK. TS gates its redaction flush on an + AfterModelCall hook that it does not register under ``saveLatestOn: 'trigger'``, so TS does not + flush redactions under that strategy. Python has no redaction signal on AfterModelCallEvent; + redaction arrives through the Agent's direct ``redact_latest_message`` call, which always + persists so pre-redaction content never sits at rest. We assert the safer always-flush here. + """ + storage = LocalFileStorage(temp_dir) + manager = SnapshotSessionManager("s1", storage=storage, save_latest_on="trigger") + redaction_model = MockedModelProvider( + [{"redactedUserContent": "REDACTED", "redactedAssistantContent": "I can't help with that."}] + ) + agent = Agent(model=redaction_model, session_manager=manager, agent_id="a1") + agent("sensitive prompt") + + manager_2 = SnapshotSessionManager("s1", storage=storage, save_latest_on="trigger") + agent_2 = Agent(model=_model("x"), session_manager=manager_2, agent_id="a1") + + tru_texts = [content["text"] for message in agent_2.messages for content in message["content"] if "text" in content] + assert "sensitive prompt" not in tru_texts + assert "REDACTED" in tru_texts + + +def test_snapshot_trigger_returning_false_appends_nothing(storage): + """A present trigger that returns False creates no immutable snapshot and receives the agent.""" + seen_agents = [] + + def trigger(*, agent_data, **kwargs): + seen_agents.append(agent_data) + return False + + manager = SnapshotSessionManager("s1", storage=storage, snapshot_trigger=trigger) + agent = Agent(model=_model("hi"), session_manager=manager, agent_id="a1") + agent("go") + + assert asyncio.run(manager.list_snapshot_ids(agent)) == [] + # The trigger was invoked with the agent as the agent_data keyword argument. + assert seen_agents and seen_agents[0] is agent + + +@pytest.mark.asyncio +async def test_list_snapshot_ids_pagination(storage): + """limit and start_after page the immutable id list; invalid start_after raises.""" + manager = SnapshotSessionManager("s1", storage=storage, snapshot_trigger=lambda *, agent_data, **_: True) + agent = Agent(model=_model("a", "b", "c"), session_manager=manager, agent_id="a1") + agent("one") + agent("two") + agent("three") + + all_ids = await manager.list_snapshot_ids(agent) + assert len(all_ids) == 3 + assert all_ids == sorted(all_ids) + + assert await manager.list_snapshot_ids(agent, limit=2) == all_ids[:2] + assert await manager.list_snapshot_ids(agent, limit=0) == [] + assert await manager.list_snapshot_ids(agent, start_after=all_ids[0]) == all_ids[1:] + + with pytest.raises(ValueError, match="not a valid snapshot id"): + await manager.list_snapshot_ids(agent, start_after="not-an-id") + + +@pytest.mark.asyncio +async def test_restore_snapshot_rejects_malformed_id(storage): + """restore_snapshot with a malformed id raises rather than silently missing.""" + manager = SnapshotSessionManager("s1", storage=storage) + agent = Agent(model=_model("hi"), session_manager=manager, agent_id="a1") + + with pytest.raises(ValueError, match="not a valid snapshot id"): + await manager.restore_snapshot(agent, snapshot_id="../escape") + + +@pytest.mark.asyncio +async def test_delete_session_is_scoped_to_its_own_session(temp_dir): + """delete_session removes only its session, leaving other sessions and keys intact.""" + storage = LocalFileStorage(temp_dir) + + manager_a = SnapshotSessionManager("sess-a", storage=storage) + Agent(model=_model("a"), session_manager=manager_a, agent_id="a1")("hi") + manager_b = SnapshotSessionManager("sess-b", storage=storage) + Agent(model=_model("b"), session_manager=manager_b, agent_id="a1")("hi") + # An unrelated key from another subsystem sharing the same storage. + await storage.write("memory/note.json", b"keep me") + + await manager_a.delete_session() + + assert await storage.read(_on_disk_key("sess-a", "a1")) is None + assert await storage.read(_on_disk_key("sess-b", "a1")) is not None + assert await storage.read("memory/note.json") == b"keep me" + + +def test_stateful_model_restore_keeps_model_state(storage): + """Restoring a stateful-model session drops messages but preserves model_state.""" + original = _stateful_model("hi") + try: + manager = SnapshotSessionManager("s1", storage=storage) + agent = Agent(model=original, session_manager=manager, agent_id="a1") + agent.messages = [{"role": "user", "content": [{"text": "hello"}]}] + agent._model_state = {"response_id": "resp-123"} + manager.sync_agent(agent) + + manager_2 = SnapshotSessionManager("s1", storage=storage) + agent_2 = Agent(model=_stateful_model("x"), session_manager=manager_2, agent_id="a1") + + assert agent_2.messages == [] # messages dropped for the stateful model + assert agent_2._model_state == {"response_id": "resp-123"} # but model_state survives + finally: + del type(original).stateful + + +@pytest.mark.asyncio +async def test_raw_storage_is_namespaced_under_session(temp_dir): + """Raw storage is auto-namespaced under 'session/', matching the TS key layout.""" + storage = LocalFileStorage(temp_dir) + manager = SnapshotSessionManager("sid", storage=storage) + Agent(model=_model("hi"), session_manager=manager, agent_id="a1")("go") + + keys = await storage.list("") + assert keys == ["session/sid/scopes/agent/a1/snapshots/snapshot_latest.json"] + + +@pytest.mark.asyncio +async def test_prenamespaced_storage_is_not_double_prefixed(temp_dir): + """A caller-namespaced view is used as-is; its 'session' prefix is not doubled.""" + storage = LocalFileStorage(temp_dir) + scoped = storage.namespace("session") # caller pre-namespaces under the same prefix + manager = SnapshotSessionManager("sid", storage=scoped) + Agent(model=_model("hi"), session_manager=manager, agent_id="a1")("go") + + # On raw storage the key is session/sid/... — a single "session/", not session/session/... + keys = await storage.list("") + assert keys == ["session/sid/scopes/agent/a1/snapshots/snapshot_latest.json"] + + +def test_snapshot_ids_are_monotonic_uuidv7(storage): + """Immutable ids are UUIDv7 and sort in creation order even within one millisecond.""" + manager = SnapshotSessionManager("s1", storage=storage, snapshot_trigger=lambda *, agent_data, **_: True) + agent = Agent(model=_model(*[f"t{index}" for index in range(6)]), session_manager=manager, agent_id="a1") + for index in range(6): + agent(f"turn {index}") + + ids = asyncio.run(manager.list_snapshot_ids(agent)) + assert len(ids) == 6 + assert all(uuid.UUID(snapshot_id).version == 7 for snapshot_id in ids) + # list_snapshot_ids sorts lexicographically; that must equal creation order. + assert ids == sorted(ids) + + +@pytest.mark.asyncio +async def test_corrupt_snapshot_raises_typed_error_on_restore(temp_dir): + """A corrupt/truncated stored snapshot surfaces a typed SnapshotException, not a raw decode error. + + Restore runs in the agent constructor, so a partially written or tampered blob would + otherwise crash construction with a JSONDecodeError leaking out of the session manager. + """ + storage = LocalFileStorage(temp_dir) + await storage.write(f"session/{_snapshot_key('s1', 'a1', snapshot_id=None)}", b'{"scope": "agent", "data": {') + + with pytest.raises(SnapshotException, match="Failed to deserialize snapshot"): + Agent(model=_model("hi"), session_manager=SnapshotSessionManager("s1", storage=storage), agent_id="a1") + + +@pytest.mark.parametrize( + "blob", + [ + b"{}", # object missing required keys -> KeyError in Snapshot.from_dict + b"42", # non-object scalar -> would AttributeError on .get() + b'"a string"', + b"[]", + b"null", + ], +) +@pytest.mark.asyncio +async def test_wrong_shape_snapshot_raises_typed_error_on_restore(temp_dir, blob): + """A valid-JSON but wrong-shape stored snapshot surfaces a typed SnapshotException. + + Valid JSON that is not a well-formed snapshot (missing keys, or a non-object scalar/array) + must not leak a raw KeyError/AttributeError out of the agent constructor's restore path. + """ + storage = LocalFileStorage(temp_dir) + await storage.write(f"session/{_snapshot_key('s1', 'a1', snapshot_id=None)}", blob) + + with pytest.raises(SnapshotException): + Agent(model=_model("hi"), session_manager=SnapshotSessionManager("s1", storage=storage), agent_id="a1") diff --git a/strands-py/tests/strands/test_identifier.py b/strands-py/tests/strands/test_identifier.py index df673baa8b..8442af3515 100644 --- a/strands-py/tests/strands/test_identifier.py +++ b/strands-py/tests/strands/test_identifier.py @@ -1,3 +1,6 @@ +import uuid +from unittest.mock import patch + import pytest from strands import _identifier @@ -15,3 +18,52 @@ def test_validate_invalid(type_): id_ = "a/../b" with pytest.raises(ValueError, match=f"{type_.value}={id_} | id cannot contain path separators"): _identifier.validate(id_, type_) + + +def test_new_uuid7_is_valid_uuidv7(): + """Generated ids are valid version-7 UUIDs and pass the predicate.""" + generated_id = _identifier.new_uuid7() + parsed = uuid.UUID(generated_id) + assert parsed.version == 7 + assert (parsed.int >> 62) & 0b11 == 0b10 # RFC 4122 variant + assert _identifier.is_uuid7(generated_id) + + +def test_uuid7_ids_sort_in_creation_order(): + """Ids minted in sequence sort lexicographically into creation order.""" + ids = [_identifier.new_uuid7() for _ in range(1000)] + assert len(set(ids)) == len(ids) + assert ids == sorted(ids) + + +def test_uuid7_ids_stay_monotonic_past_counter_overflow(): + """Minting >4096 ids in one millisecond keeps them unique, valid v7, and strictly increasing. + + Guards the 12-bit intra-millisecond counter: on overflow it must borrow from the clock, not + wrap backwards, which would break the sort-equals-creation-order property callers rely on. + """ + with patch.object(_identifier.time, "time", return_value=1_000.0): + ids = [_identifier.new_uuid7() for _ in range(20_000)] + + assert len(set(ids)) == len(ids) # unique + assert all(uuid.UUID(generated_id).version == 7 for generated_id in ids) # valid v7 + assert all((uuid.UUID(generated_id).int >> 62) & 0b11 == 0b10 for generated_id in ids) # variant intact + assert all(ids[index] < ids[index + 1] for index in range(len(ids) - 1)) # strictly increasing + + +@pytest.mark.parametrize( + "bad_id", + [ + "not-a-uuid", + "00000000-0000-4000-8000-000000000000", # version 4, not 7 + "", + ], +) +def test_is_uuid7_rejects_non_uuidv7(bad_id): + """Malformed strings and non-v7 UUIDs are rejected.""" + assert not _identifier.is_uuid7(bad_id) + + +def test_is_uuid7_rejects_trailing_newline(): + """A trailing newline cannot slip past the pattern (\\A...\\Z, not ^...$).""" + assert not _identifier.is_uuid7(_identifier.new_uuid7() + "\n") diff --git a/strands-py/tests_integ/test_snapshot_session.py b/strands-py/tests_integ/test_snapshot_session.py new file mode 100644 index 0000000000..77cffa9234 --- /dev/null +++ b/strands-py/tests_integ/test_snapshot_session.py @@ -0,0 +1,137 @@ +"""Integration tests for snapshot-based session management.""" + +import asyncio +import os +import tempfile +from uuid import uuid4 + +import boto3 +import pytest +from botocore.client import ClientError + +from strands import Agent +from strands.models.openai_responses import OpenAIResponsesModel +from strands.session.snapshot_session_manager import SnapshotSessionManager +from strands.storage import LocalFileStorage, S3Storage +from tests_integ.models.providers import openai as openai_provider + +# yellow_img imported from conftest + + +@pytest.fixture +def temp_dir(): + """Create a temporary directory for testing.""" + with tempfile.TemporaryDirectory() as temp_dir: + yield temp_dir + + +@pytest.fixture +def bucket_name(): + # Shares the bucket the message-log session tests use, because the integ-test IAM role only + # grants S3 access to an explicit allowlist of bucket names. Snapshots key under "session/" + # while the message log keys under "session_", so the two cannot collide. + bucket_name = f"test-strands-session-bucket-{boto3.client('sts').get_caller_identity()['Account']}" + s3_client = boto3.resource("s3", region_name="us-west-2") + try: + s3_client.create_bucket(Bucket=bucket_name, CreateBucketConfiguration={"LocationConstraint": "us-west-2"}) + except ClientError as error: + if "BucketAlreadyOwnedByYou" not in str(error): + raise error + yield bucket_name + + +def test_agent_with_file_snapshot_session(temp_dir): + """A fresh agent rehydrates its conversation from a file-backed snapshot.""" + test_session_id = str(uuid4()) + manager = SnapshotSessionManager(test_session_id, storage=LocalFileStorage(temp_dir)) + agent = Agent(session_manager=manager) + agent("Hello!") + assert len(agent.messages) == 2 + + # Simulate process restart: new manager + agent over the same storage. + manager_2 = SnapshotSessionManager(test_session_id, storage=LocalFileStorage(temp_dir)) + agent_2 = Agent(session_manager=manager_2) + assert len(agent_2.messages) == 2 + agent_2("Hello again!") + assert len(agent_2.messages) == 4 + + +def test_agent_with_file_snapshot_session_with_image(temp_dir, yellow_img): + """Image bytes survive a snapshot round-trip across instances.""" + test_session_id = str(uuid4()) + manager = SnapshotSessionManager(test_session_id, storage=LocalFileStorage(temp_dir)) + agent = Agent(session_manager=manager) + agent([{"image": {"format": "png", "source": {"bytes": yellow_img}}}]) + assert len(agent.messages) == 2 + + manager_2 = SnapshotSessionManager(test_session_id, storage=LocalFileStorage(temp_dir)) + agent_2 = Agent(session_manager=manager_2) + assert agent_2.messages[0]["content"][0]["image"]["source"]["bytes"] == yellow_img + + +def test_agent_with_s3_snapshot_session(bucket_name): + """A fresh agent rehydrates its conversation from an S3-backed snapshot.""" + test_session_id = str(uuid4()) + store = S3Storage(bucket=bucket_name, region_name="us-west-2") + manager = SnapshotSessionManager(test_session_id, storage=store) + try: + agent = Agent(session_manager=manager) + agent("Hello!") + assert len(agent.messages) == 2 + + manager_2 = SnapshotSessionManager( + test_session_id, storage=S3Storage(bucket=bucket_name, region_name="us-west-2") + ) + agent_2 = Agent(session_manager=manager_2) + assert len(agent_2.messages) == 2 + agent_2("Hello again!") + assert len(agent_2.messages) == 4 + finally: + asyncio.run(manager.delete_session()) + + +def test_snapshot_session_time_travel(temp_dir): + """Immutable checkpoints allow restoring an agent to an earlier turn.""" + test_session_id = str(uuid4()) + store = LocalFileStorage(temp_dir) + manager = SnapshotSessionManager(test_session_id, storage=store, snapshot_trigger=lambda *, agent_data, **_: True) + agent = Agent(session_manager=manager) + agent("My favorite color is blue.") + agent("My favorite number is seven.") + + ids = asyncio.run(manager.list_snapshot_ids(agent)) + assert len(ids) == 2 + + # Restore to the first checkpoint: only the first turn should be present. + restored = asyncio.run(manager.restore_snapshot(agent, snapshot_id=ids[0])) + assert restored is True + assert len(agent.messages) == 2 + + +@openai_provider.mark +def test_agent_with_snapshot_session_server_side_conversation(temp_dir): + """Server-side conversation state survives snapshot save/restore for a stateful model.""" + test_session_id = str(uuid4()) + store = LocalFileStorage(temp_dir) + manager = SnapshotSessionManager(test_session_id, storage=store) + + model = OpenAIResponsesModel( + model_id="gpt-4o-mini", + stateful=True, + client_args={"api_key": os.getenv("OPENAI_API_KEY")}, + ) + agent = Agent(model=model, system_prompt="Reply in one short sentence.", session_manager=manager) + agent("My name is Alice.") + assert len(agent.messages) == 0 + + # Simulate process restart. + manager_2 = SnapshotSessionManager(test_session_id, storage=LocalFileStorage(temp_dir)) + model_2 = OpenAIResponsesModel( + model_id="gpt-4o-mini", + stateful=True, + client_args={"api_key": os.getenv("OPENAI_API_KEY")}, + ) + agent_2 = Agent(model=model_2, system_prompt="Reply in one short sentence.", session_manager=manager_2) + assert len(agent_2.messages) == 0 + result = agent_2("What is my name?") + assert "alice" in result.message["content"][0]["text"].lower()