Add some pr feedback

Author: Unshure

src/strands/agent/agent.py

@@ -322,7 +322,7 @@ def __init__(
 
         # Setup session callback handler if session is enabled
         if self.session_manager:
-            self.session_manager.initialize_agent(self)
+            self.session_manager.initialize(self)
 
     @property
     def tool(self) -> ToolCaller:
@@ -500,7 +500,7 @@ def _run_loop(self, prompt: str, kwargs: dict[str, Any]) -> Generator[dict[str,
 
             # Save message if session manager is available
             if self.session_manager:
-                self.session_manager.append_message_to_agent_session(self, new_message)
+                self.session_manager.append_message(self, new_message)
 
             # Execute the event loop cycle with retry logic for context limits
             yield from self._execute_event_loop_cycle(kwargs)
@@ -597,10 +597,10 @@ def _record_tool_execution(
 
         # 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)
+            self.session_manager.append_message(self, user_msg)
+            self.session_manager.append_message(self, tool_use_msg)
+            self.session_manager.append_message(self, tool_result_msg)
+            self.session_manager.append_message(self, assistant_msg)
 
     def _start_agent_trace_span(self, prompt: str) -> None:
         """Starts a trace span for the agent.

src/strands/session/agent_session_manager.py

@@ -1,13 +1,13 @@
 """File-based implementation of session manager."""
 
 import logging
-from typing import TYPE_CHECKING, Any, Optional
+from typing import TYPE_CHECKING, Any, List
+from uuid import uuid4
 
 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 ..types.exceptions import SessionException
 from .session_dao import SessionDAO
 from .session_manager import SessionManager
 from .session_models import Session, SessionAgent, SessionMessage, SessionType
@@ -28,13 +28,13 @@ class AgentSessionManager(SessionManager):
     def __init__(
         self,
         session_id: str,
-        session_dao: Optional[SessionDAO] = None,
+        session_dao: SessionDAO,
     ):
         """Initialize the FileSessionManager."""
-        self.session_dao = session_dao or FileSessionDAO()
+        self.session_dao = session_dao
         self.session_id = session_id
 
-    def append_message_to_agent_session(self, agent: "Agent", message: Message) -> None:
+    def append_message(self, agent: "Agent", message: Message) -> None:
         """Append a message to the agent's session.
 
         Args:
@@ -56,7 +56,7 @@ def append_message_to_agent_session(self, agent: "Agent", message: Message) -> N
             ),
         )
 
-    def initialize_agent(self, agent: "Agent") -> None:
+    def initialize(self, agent: "Agent") -> None:
         """Restore agent data from the current session.
 
         Args:
@@ -65,20 +65,33 @@ def initialize_agent(self, agent: "Agent") -> None:
         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 agent.id is None:
+                agents: List[SessionAgent] = self.session_dao.list_agents(self.session_id)
+                if len(agents) == 0:
+                    agent_id = str(uuid4())
+                if len(agents) == 1:
+                    agent_id = agents[0].agent_id
+                    logger.debug(
+                        "session_id=<%s> | agent_id=<%s> | Restoring agent data from session", self.session_id, agent_id
+                    )
+                else:
+                    raise ValueError(
+                        "If there is more than one agent in a session, agent.agent_id must be set manually."
+                    )
+            else:
+                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}")
+                agent_id = agent.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.id = agent_id
             agent.messages = [
                 session_message.to_message()
                 for session_message in self.session_dao.list_messages(self.session_id, agent.id)
@@ -87,8 +100,12 @@ def initialize_agent(self, agent: "Agent") -> None:
 
         except SessionException:
             # Session doesn't exist, create new one
-            logger.debug("Session not found, creating new session")
+            logger.debug("session_id=<%s> | Session not found, creating new session")
             # Session doesn't exist, create new one
+            if agent.id is None:
+                agent_id = str(uuid4())
+                logger.debug("agent_id=<%s> | Creating agent_id for agent since none was set.", agent_id)
+                agent.id = agent_id
             session = Session(session_id=self.session_id, session_type=SessionType.AGENT)
             session_agent = SessionAgent(
                 agent_id=agent.id,
@@ -110,7 +127,7 @@ def session_callback(**kwargs: Any) -> None:
                 # Handle message persistence
                 if "message" in kwargs:
                     message = kwargs["message"]
-                    self.append_message_to_agent_session(kwargs["agent"], message)
+                    self.append_message(kwargs["agent"], message)
             except Exception as e:
                 logger.error("Persistence operation failed", e)
 

src/strands/session/exceptions.py

@@ -1,4 +1,4 @@
-"""File-based session DAO for local filesystem storage."""
+"""File-based session manager for local filesystem storage."""
 
 import json
 import os
@@ -7,7 +7,8 @@
 from datetime import datetime, timezone
 from typing import Any, Dict, List, Optional, cast
 
-from .exceptions import SessionException
+from ..types.exceptions import SessionException
+from .agent_session_manager import AgentSessionManager
 from .session_dao import SessionDAO
 from .session_models import Session, SessionAgent, SessionMessage
 
@@ -16,18 +17,21 @@
 MESSAGE_PREFIX = "message_"
 
 
-class FileSessionDAO(SessionDAO):
-    """File-based session DAO for local filesystem storage."""
+class FileSessionManager(AgentSessionManager, SessionDAO):
+    """File-based session manager for local filesystem storage."""
 
-    def __init__(self, storage_dir: Optional[str] = None):
-        """Initialize FileSessionDAO with 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_dao=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}")

src/strands/session/file_session_dao.py …/strands/session/file_session_manager.pysrc/strands/session/file_session_dao.py renamed to src/strands/session/file_session_manager.py src/strands/session/file_session_dao.py renamed to src/strands/session/file_session_manager.py

@@ -1,4 +1,4 @@
-"""S3-based session DAO for cloud storage."""
+"""S3-based session manager for AWS S3 cloud storage."""
 
 import json
 from typing import Any, Dict, List, Optional, cast
@@ -7,7 +7,8 @@
 from botocore.config import Config as BotocoreConfig
 from botocore.exceptions import ClientError
 
-from .exceptions import SessionException
+from ..types.exceptions import SessionException
+from .agent_session_manager import AgentSessionManager
 from .session_dao import SessionDAO
 from .session_models import Session, SessionAgent, SessionMessage
 
@@ -16,20 +17,22 @@
 MESSAGE_PREFIX = "message_"
 
 
-class S3SessionDAO(SessionDAO):
-    """S3-based session DAO for cloud storage."""
+class S3SessionManager(AgentSessionManager, SessionDAO):
+    """S3-based session manager for cloud storage."""
 
     def __init__(
         self,
+        session_id: str,
         bucket: str,
         prefix: str = "",
         boto_session: Optional[boto3.Session] = None,
         boto_client_config: Optional[BotocoreConfig] = None,
         region_name: Optional[str] = None,
     ):
-        """Initialize S3SessionDAO with S3 storage.
+        """Initialize S3SessionManager with S3 storage.
 
         Args:
+            session_id: ID for the session
             bucket: S3 bucket name (required)
             prefix: S3 key prefix for storage organization
             boto_session: Optional boto3 session
@@ -55,6 +58,8 @@ def __init__(
 
         self.client = session.client(service_name="s3", config=client_config)
 
+        super().__init__(session_id=session_id, session_dao=self)
+
     def _get_session_path(self, session_id: str) -> str:
         """Get session S3 prefix."""
         return f"{self.prefix}{SESSION_PREFIX}{session_id}/"