feat: Session persistence

Author: Unshure

src/strands/agent/agent.py

@@ -23,6 +23,7 @@
 from ..handlers.callback_handler import PrintingCallbackHandler, null_callback_handler
 from ..handlers.tool_handler import AgentToolHandler
 from ..models.bedrock import BedrockModel
+from ..session.session_manager import SessionManager
 from ..telemetry.metrics import EventLoopMetrics
 from ..telemetry.tracer import get_tracer
 from ..tools.registry import ToolRegistry
@@ -192,8 +193,10 @@ def __init__(
         trace_attributes: Optional[Mapping[str, AttributeValue]] = None,
         *,
         name: Optional[str] = None,
+        id: Optional[str] = None,
         description: Optional[str] = None,
         state: Optional[Union[AgentState, dict]] = None,
+        session_manager: Optional["SessionManager"] = None,
     ):
         """Initialize the Agent with the specified configuration.
 
@@ -228,10 +231,15 @@ def __init__(
             trace_attributes: Custom trace attributes to apply to the agent's trace span.
             name: name of the Agent
                 Defaults to None.
+            id: identifier for the agent, used by session manager.
+                Defaults to uuid4().
             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.
+            session_manager: Manager for handling agent sessions including conversation history and state.
+                If provided, enables session-based persistence and state management. The session manager
+                handles agent_id, and session identification internally.
 
         Raises:
             ValueError: If max_parallel_tools is less than 1.
@@ -304,10 +312,18 @@ def __init__(
         else:
             self.state = AgentState()
 
+        # Initialize session management functionality
+        self.session_manager = session_manager
+
         self.tool_caller = Agent.ToolCaller(self)
         self.name = name
+        self.id = id
         self.description = description
 
+        # Setup session callback handler if session is enabled
+        if self.session_manager:
+            self.session_manager.initialize_agent(self)
+
     @property
     def tool(self) -> ToolCaller:
         """Call tool as a function.
@@ -482,6 +498,10 @@ def _run_loop(self, prompt: str, kwargs: dict[str, Any]) -> Generator[dict[str,
             new_message: Message = {"role": "user", "content": message_content}
             self.messages.append(new_message)
 
+            # Save message if session manager is available
+            if self.session_manager:
+                self.session_manager.append_message_to_agent_session(self, new_message)
+
             # Execute the event loop cycle with retry logic for context limits
             yield from self._execute_event_loop_cycle(kwargs)
 
@@ -575,6 +595,13 @@ def _record_tool_execution(
         messages.append(tool_result_msg)
         messages.append(assistant_msg)
 
+        # Save to conversation manager if available
+        if self.session_manager:
+            self.session_manager.append_message_to_agent_session(self, user_msg)
+            self.session_manager.append_message_to_agent_session(self, tool_use_msg)
+            self.session_manager.append_message_to_agent_session(self, tool_result_msg)
+            self.session_manager.append_message_to_agent_session(self, assistant_msg)
+
     def _start_agent_trace_span(self, prompt: str) -> None:
         """Starts a trace span for the agent.
 

src/strands/agent/agent_result.py

@@ -4,7 +4,7 @@
 """
 
 from dataclasses import dataclass
-from typing import Any
+from typing import Any, Optional
 
 from ..telemetry.metrics import EventLoopMetrics
 from ..types.content import Message
@@ -20,12 +20,16 @@ class AgentResult:
         message: The last message generated by the agent.
         metrics: Performance metrics collected during processing.
         state: Additional state information from the event loop.
+        name: (Optional) name of the agent if it is defined
+        id: (Optional) id of the agent if it is defined
     """
 
     stop_reason: StopReason
     message: Message
     metrics: EventLoopMetrics
     state: Any
+    name: Optional[str] = None
+    id: Optional[str] = None
 
     def __str__(self) -> str:
         """Get the agent's last message as a string.

src/strands/agent/state.py

@@ -14,7 +14,11 @@ class AgentState:
     - Get/set/delete operations
     """
 
+<<<<<<< HEAD
     def __init__(self, initial_state: Optional[Dict[str, Any]] = None):
+=======
+    def __init__(self, initial_state: Optional[Dict[str, Dict[str, Any]]] = None):
+>>>>>>> 1dbbfe6 (feat: Session persistence)
         """Initialize AgentState."""
         self._state: Dict[str, Dict[str, Any]]
         if initial_state:

src/strands/event_loop/event_loop.py

@@ -181,7 +181,10 @@ def event_loop_cycle(
 
         # Add the response message to the conversation
         messages.append(message)
-        yield {"callback": {"message": message}}
+        callback_data = {"message": message}
+        if "agent" in kwargs:
+            callback_data["agent"] = kwargs["agent"]
+        yield {"callback": callback_data}
 
         # Update metrics
         event_loop_metrics.update_usage(usage)
@@ -390,7 +393,10 @@ def _handle_tool_execution(
     }
 
     messages.append(tool_result_message)
-    yield {"callback": {"message": tool_result_message}}
+    callback_data = {"message": tool_result_message}
+    if "agent" in kwargs:
+        callback_data["agent"] = kwargs["agent"]
+    yield {"callback": callback_data}
 
     if cycle_span:
         tracer = get_tracer()

src/strands/session/agent_session_manager.py

@@ -0,0 +1,117 @@
+"""File-based implementation of session manager."""
+
+import logging
+from typing import TYPE_CHECKING, Any, Optional
+
+from ..agent.state import AgentState
+from ..handlers.callback_handler import CompositeCallbackHandler
+from ..types.content import Message
+from .exceptions import SessionException
+from .file_session_dao import FileSessionDAO
+from .session_dao import SessionDAO
+from .session_manager import SessionManager
+from .session_models import Session, SessionAgent, SessionMessage, SessionType
+
+logger = logging.getLogger(__name__)
+
+if TYPE_CHECKING:
+    from ..agent.agent import Agent
+
+
+class AgentSessionManager(SessionManager):
+    """Session manager for a single Agent.
+
+    This implementation stores sessions as JSON files in a specified directory.
+    Each session is stored in a separate file named by its session_id.
+    """
+
+    def __init__(
+        self,
+        session_id: str,
+        session_dao: Optional[SessionDAO] = None,
+    ):
+        """Initialize the FileSessionManager."""
+        self.session_dao = session_dao or FileSessionDAO()
+        self.session_id = session_id
+
+    def append_message_to_agent_session(self, agent: "Agent", message: Message) -> None:
+        """Append a message to the agent's session.
+
+        Args:
+            agent: The agent whose session to update
+            message: The message to append
+        """
+        if agent.id is None:
+            raise ValueError("`agent.id` must be set before appending message to session.")
+
+        session_message = SessionMessage.from_dict(dict(message))
+        self.session_dao.create_message(self.session_id, agent.id, session_message)
+        self.session_dao.update_agent(
+            self.session_id,
+            SessionAgent(
+                agent_id=agent.id,
+                session_id=self.session_id,
+                event_loop_metrics=agent.event_loop_metrics.to_dict(),
+                state=agent.state.get(),
+            ),
+        )
+
+    def initialize_agent(self, agent: "Agent") -> None:
+        """Restore agent data from the current session.
+
+        Args:
+            agent: Agent instance to restore session data to
+
+        Raises:
+            SessionException: If restore operation fails
+        """
+        if agent.id is None:
+            raise ValueError("`agent.id` must be set before initializing session.")
+
+        try:
+            # Try to read existing session
+            session = self.session_dao.read_session(self.session_id)
+
+            if session.session_type != SessionType.AGENT:
+                raise ValueError(f"Invalid session type: {session.session_type}")
+
+            if agent.id not in [agent.agent_id for agent in self.session_dao.list_agents(self.session_id)]:
+                raise ValueError(f"Agent {agent.id} not found in session {self.session_id}")
+
+            # Initialize agent
+            agent.messages = [
+                session_message.to_message()
+                for session_message in self.session_dao.list_messages(self.session_id, agent.id)
+            ]
+            agent.state = AgentState(self.session_dao.read_agent(self.session_id, agent.id).state)
+
+        except SessionException:
+            # Session doesn't exist, create new one
+            logger.debug("Session not found, creating new session")
+            # Session doesn't exist, create new one
+            session = Session(session_id=self.session_id, session_type=SessionType.AGENT)
+            session_agent = SessionAgent(
+                agent_id=agent.id,
+                session_id=self.session_id,
+                event_loop_metrics=agent.event_loop_metrics.to_dict(),
+                state=agent.state.get(),
+            )
+            self.session_dao.create_session(session)
+            self.session_dao.create_agent(self.session_id, session_agent)
+            for message in agent.messages:
+                session_message = SessionMessage.from_dict(dict(message))
+                self.session_dao.create_message(self.session_id, agent.id, session_message)
+
+        self.session = session
+
+        # Attach a callback handler for persisting messages
+        def session_callback(**kwargs: Any) -> None:
+            try:
+                # Handle message persistence
+                if "message" in kwargs:
+                    message = kwargs["message"]
+                    self.append_message_to_agent_session(kwargs["agent"], message)
+            except Exception as e:
+                logger.error("Persistence operation failed", e)
+
+        agent.callback_handler = CompositeCallbackHandler(agent.callback_handler, session_callback)

src/strands/session/exceptions.py

@@ -0,0 +1,5 @@
+"""Exception classes for session management operations."""
+
+
+class SessionException(Exception):
+    """Exception raised when session operations fail."""