feat: Session persistence

Author: Unshure

src/strands/agent/agent.py

@@ -15,7 +15,6 @@
 import random
 from concurrent.futures import ThreadPoolExecutor
 from typing import Any, AsyncGenerator, AsyncIterator, Callable, Mapping, Optional, Type, TypeVar, Union, cast
-from uuid import uuid4
 
 from opentelemetry import trace
 from pydantic import BaseModel
@@ -32,6 +31,7 @@
 )
 from ..models.bedrock import BedrockModel
 from ..models.model import Model
+from ..session.session_manager import SessionManager
 from ..telemetry.metrics import EventLoopMetrics
 from ..telemetry.tracer import get_tracer
 from ..tools.registry import ToolRegistry
@@ -62,6 +62,7 @@ class _DefaultCallbackHandlerSentinel:
 
 _DEFAULT_CALLBACK_HANDLER = _DefaultCallbackHandlerSentinel()
 _DEFAULT_AGENT_NAME = "Strands Agents"
+_DEFAULT_AGENT_ID = "default"
 
 
 class Agent:
@@ -207,6 +208,7 @@ def __init__(
         description: Optional[str] = None,
         state: Optional[Union[AgentState, dict]] = None,
         hooks: Optional[list[HookProvider]] = None,
+        session_manager: Optional[SessionManager] = None,
     ):
         """Initialize the Agent with the specified configuration.
 
@@ -237,22 +239,24 @@ def __init__(
             load_tools_from_directory: Whether to load and automatically reload tools in the `./tools/` directory.
                 Defaults to False.
             trace_attributes: Custom trace attributes to apply to the agent's trace span.
-            agent_id: Optional ID for the agent, useful for multi-agent scenarios.
-                If None, a UUID is generated.
+            agent_id: Optional ID for the agent, useful for session management and multi-agent scenarios.
+                Defaults to "default".
             name: name of the Agent
-                Defaults to None.
+                Defaults to "Strands Agents".
             description: description of what the Agent does
                 Defaults to None.
             state: stateful information for the agent. Can be either an AgentState object, or a json serializable dict.
                 Defaults to an empty AgentState object.
             hooks: hooks to be added to the agent hook registry
                 Defaults to None.
+            session_manager: Manager for handling agent sessions including conversation history and state.
+                If provided, enables session-based persistence and state management.
         """
         self.model = BedrockModel() if not model else BedrockModel(model_id=model) if isinstance(model, str) else model
         self.messages = messages if messages is not None else []
 
         self.system_prompt = system_prompt
-        self.agent_id = agent_id or str(uuid4())
+        self.agent_id = agent_id or _DEFAULT_AGENT_ID
         self.name = name or _DEFAULT_AGENT_NAME
         self.description = description
 
@@ -312,6 +316,12 @@ def __init__(
         self.tool_caller = Agent.ToolCaller(self)
 
         self.hooks = HookRegistry()
+
+        # Initialize session management functionality
+        self.session_manager = session_manager
+        if self.session_manager:
+            self.hooks.add_hook(self.session_manager)
+
         if hooks:
             for hook in hooks:
                 self.hooks.add_hook(hook)

src/strands/session/agent_session_manager.py

@@ -0,0 +1,102 @@
+"""Agent session manager implementation."""
+
+import logging
+
+from ..agent.agent import _DEFAULT_AGENT_ID, Agent
+from ..agent.state import AgentState
+from ..types.content import Message
+from ..types.exceptions import SessionException
+from ..types.session import (
+    Session,
+    SessionAgent,
+    SessionMessage,
+    SessionType,
+)
+from .session_manager import SessionManager
+from .session_repository import SessionRepository
+
+logger = logging.getLogger(__name__)
+
+DEFAULT_SESSION_AGENT_ID = "default"
+
+
+class AgentSessionManager(SessionManager):
+    """Session manager for persisting agent's in a Session."""
+
+    def __init__(
+        self,
+        session_id: str,
+        session_repository: SessionRepository,
+    ):
+        """Initialize the AgentSessionManager."""
+        self.session_repository = session_repository
+        self.session_id = session_id
+        session = session_repository.read_session(session_id)
+        # Create a session if it does not exist yet
+        if session is None:
+            logger.debug("session_id=<%s> | Session not found, creating new session.", self.session_id)
+            session = Session(session_id=session_id, session_type=SessionType.AGENT)
+            session_repository.create_session(session)
+
+        self.session = session
+        self._default_agent_initialized = False
+
+    def append_message(self, message: Message, agent: Agent) -> None:
+        """Append a message to the agent's session.
+
+        Args:
+            message: Message to add to the agent in the session
+            agent: Agent to append the message to
+        """
+        session_message = SessionMessage.from_message(message)
+        if agent.agent_id is None:
+            raise ValueError("`agent.agent_id` must be set before appending message to session.")
+        self.session_repository.create_message(self.session_id, agent.agent_id, session_message)
+
+    def sync_agent(self, agent: Agent) -> None:
+        """Sync agent to the session.
+
+        Args:
+            agent: Agent to sync to the session.
+        """
+        self.session_repository.update_agent(
+            self.session_id,
+            SessionAgent.from_agent(agent),
+        )
+
+    def initialize(self, agent: Agent) -> None:
+        """Initialize an agent with a session.
+
+        Args:
+            agent: Agent to initialize from the session
+        """
+        if agent.agent_id is _DEFAULT_AGENT_ID:
+            if self._default_agent_initialized:
+                raise SessionException("Set `agent_id` to support more than one agent in a session.")
+            self._default_agent_initialized = True
+
+        session_agent = self.session_repository.read_agent(self.session_id, agent.agent_id)
+
+        if session_agent is None:
+            logger.debug(
+                "agent_id=<%s> | session_id=<%s> | Creating agent.",
+                agent.agent_id,
+                self.session_id,
+            )
+
+            session_agent = SessionAgent.from_agent(agent)
+            self.session_repository.create_agent(self.session_id, session_agent)
+            for message in agent.messages:
+                session_message = SessionMessage.from_message(message)
+                self.session_repository.create_message(self.session_id, agent.agent_id, session_message)
+        else:
+            logger.debug(
+                "agent_id=<%s> | session_id=<%s> | Restoring agent.",
+                agent.agent_id,
+                self.session_id,
+            )
+            agent.messages = [
+                session_message.to_message()
+                for session_message in self.session_repository.list_messages(self.session_id, agent.agent_id)
+            ]
+            agent.state = AgentState(session_agent.state)

src/strands/session/file_session_manager.py

@@ -0,0 +1,187 @@
+"""File-based session manager for local filesystem storage."""
+
+import json
+import logging
+import os
+import shutil
+import tempfile
+from dataclasses import asdict
+from typing import Any, Optional, cast
+
+from ..types.exceptions import SessionException
+from ..types.session import Session, SessionAgent, SessionMessage
+from .agent_session_manager import AgentSessionManager
+from .session_repository import SessionRepository
+
+logger = logging.getLogger(__name__)
+
+SESSION_PREFIX = "session_"
+AGENT_PREFIX = "agent_"
+MESSAGE_PREFIX = "message_"
+
+
+class FileSessionManager(AgentSessionManager, SessionRepository):
+    """File-based session manager for local filesystem storage."""
+
+    def __init__(self, session_id: str, storage_dir: Optional[str] = None):
+        """Initialize FileSession with filesystem storage.
+
+        Args:
+            session_id: ID for the session
+            storage_dir: Directory for local filesystem storage (defaults to temp dir)
+        """
+        self.storage_dir = storage_dir or os.path.join(tempfile.gettempdir(), "strands/sessions")
+        os.makedirs(self.storage_dir, exist_ok=True)
+
+        super().__init__(session_id=session_id, session_repository=self)
+
+    def _get_session_path(self, session_id: str) -> str:
+        """Get session directory path."""
+        return os.path.join(self.storage_dir, f"{SESSION_PREFIX}{session_id}")
+
+    def _get_agent_path(self, session_id: str, agent_id: str) -> str:
+        """Get agent directory path."""
+        session_path = self._get_session_path(session_id)
+        return os.path.join(session_path, "agents", f"{AGENT_PREFIX}{agent_id}")
+
+    def _get_message_path(self, session_id: str, agent_id: str, message_id: str) -> str:
+        """Get message file path."""
+        agent_path = self._get_agent_path(session_id, agent_id)
+        return os.path.join(agent_path, "messages", f"{MESSAGE_PREFIX}{message_id}.json")
+
+    def _read_file(self, path: str) -> dict[str, Any]:
+        """Read JSON file."""
+        try:
+            with open(path, "r", encoding="utf-8") as f:
+                return cast(dict[str, Any], json.load(f))
+        except json.JSONDecodeError as e:
+            raise SessionException(f"Invalid JSON in file {path}: {str(e)}") from e
+
+    def _write_file(self, path: str, data: dict[str, Any]) -> None:
+        """Write JSON file."""
+        os.makedirs(os.path.dirname(path), exist_ok=True)
+        with open(path, "w", encoding="utf-8") as f:
+            json.dump(data, f, indent=2, ensure_ascii=False)
+
+    def create_session(self, session: Session) -> Session:
+        """Create a new session."""
+        session_dir = self._get_session_path(session.session_id)
+        if os.path.exists(session_dir):
+            raise SessionException(f"Session {session.session_id} already exists")
+
+        # Create directory structure
+        os.makedirs(session_dir, exist_ok=True)
+        os.makedirs(os.path.join(session_dir, "agents"), exist_ok=True)
+
+        # Write session file
+        session_file = os.path.join(session_dir, "session.json")
+        session_dict = asdict(session)
+        self._write_file(session_file, session_dict)
+
+        return session
+
+    def read_session(self, session_id: str) -> Optional[Session]:
+        """Read session data."""
+        session_file = os.path.join(self._get_session_path(session_id), "session.json")
+        if not os.path.exists(session_file):
+            return None
+
+        session_data = self._read_file(session_file)
+        return Session.from_dict(session_data)
+
+    def create_agent(self, session_id: str, session_agent: SessionAgent) -> None:
+        """Create a new agent in the session."""
+        agent_id = session_agent.agent_id
+
+        agent_dir = self._get_agent_path(session_id, agent_id)
+        os.makedirs(agent_dir, exist_ok=True)
+        os.makedirs(os.path.join(agent_dir, "messages"), exist_ok=True)
+
+        agent_file = os.path.join(agent_dir, "agent.json")
+        session_data = asdict(session_agent)
+        self._write_file(agent_file, session_data)
+
+    def delete_session(self, session_id: str) -> None:
+        """Delete session and all associated data."""
+        session_dir = self._get_session_path(session_id)
+        if not os.path.exists(session_dir):
+            raise SessionException(f"Session {session_id} does not exist")
+
+        shutil.rmtree(session_dir)
+
+    def read_agent(self, session_id: str, agent_id: str) -> Optional[SessionAgent]:
+        """Read agent data."""
+        agent_file = os.path.join(self._get_agent_path(session_id, agent_id), "agent.json")
+        if not os.path.exists(agent_file):
+            return None
+
+        agent_data = self._read_file(agent_file)
+        return SessionAgent.from_dict(agent_data)
+
+    def update_agent(self, session_id: str, session_agent: SessionAgent) -> None:
+        """Update agent data."""
+        agent_id = session_agent.agent_id
+        previous_agent = self.read_agent(session_id=session_id, agent_id=agent_id)
+        if previous_agent is None:
+            raise SessionException(f"Agent {agent_id} in session {session_id} does not exist")
+
+        session_agent.created_at = previous_agent.created_at
+        agent_file = os.path.join(self._get_agent_path(session_id, agent_id), "agent.json")
+        self._write_file(agent_file, asdict(session_agent))
+
+    def create_message(self, session_id: str, agent_id: str, session_message: SessionMessage) -> None:
+        """Create a new message for the agent."""
+        message_file = self._get_message_path(
+            session_id,
+            agent_id,
+            session_message.message_id,
+        )
+        session_dict = asdict(session_message)
+        self._write_file(message_file, session_dict)
+
+    def read_message(self, session_id: str, agent_id: str, message_id: str) -> Optional[SessionMessage]:
+        """Read message data."""
+        message_file = self._get_message_path(session_id, agent_id, message_id)
+        if not os.path.exists(message_file):
+            return None
+        message_data = self._read_file(message_file)
+        return SessionMessage.from_dict(message_data)
+
+    def update_message(self, session_id: str, agent_id: str, session_message: SessionMessage) -> None:
+        """Update message data."""
+        message_id = session_message.message_id
+        previous_message = self.read_message(session_id=session_id, agent_id=agent_id, message_id=message_id)
+        if previous_message is None:
+            raise SessionException(f"Message {message_id} does not exist")
+
+        # Preserve the original created_at timestamp
+        session_message.created_at = previous_message.created_at
+        message_file = self._get_message_path(session_id, agent_id, message_id)
+        self._write_file(message_file, asdict(session_message))
+
+    def list_messages(
+        self, session_id: str, agent_id: str, limit: Optional[int] = None, offset: int = 0
+    ) -> list[SessionMessage]:
+        """List messages for an agent with pagination."""
+        messages_dir = os.path.join(self._get_agent_path(session_id, agent_id), "messages")
+        if not os.path.exists(messages_dir):
+            raise SessionException(f"Messages directory missing from agent: {agent_id} in session {session_id}")
+
+        # Read all message files
+        messages: list[SessionMessage] = []
+        for filename in os.listdir(messages_dir):
+            if filename.startswith(MESSAGE_PREFIX) and filename.endswith(".json"):
+                file_path = os.path.join(messages_dir, filename)
+                message_data = self._read_file(file_path)
+                messages.append(SessionMessage.from_dict(message_data))
+
+        # Sort by created_at timestamp (oldest first)
+        messages.sort(key=lambda x: x.created_at)
+
+        # Apply pagination
+        if limit is not None:
+            messages = messages[offset : offset + limit]
+        else:
+            messages = messages[offset:]
+
+        return messages