refactor: add pr feedback

Author: Unshure

src/strands/agent/agent.py

@@ -318,9 +318,9 @@ def __init__(
         self.hooks = HookRegistry()
 
         # Initialize session management functionality
-        self.session_manager = session_manager
-        if self.session_manager:
-            self.hooks.add_hook(self.session_manager)
+        self._session_manager = session_manager
+        if self._session_manager:
+            self.hooks.add_hook(self._session_manager)
 
         if hooks:
             for hook in hooks:

src/strands/session/file_session_manager.py

@@ -10,7 +10,7 @@
 
 from ..types.exceptions import SessionException
 from ..types.session import Session, SessionAgent, SessionMessage
-from .agent_session_manager import AgentSessionManager
+from .repository_session_manager import RepositorySessionManager
 from .session_repository import SessionRepository
 
 logger = logging.getLogger(__name__)
@@ -20,8 +20,21 @@
 MESSAGE_PREFIX = "message_"
 
 
-class FileSessionManager(AgentSessionManager, SessionRepository):
-    """File-based session manager for local filesystem storage."""
+class FileSessionManager(RepositorySessionManager, SessionRepository):
+    """File-based session manager for local filesystem storage.
+
+    Creates the following filesystem structure for the session storage:
+    /<sessions_dir>/
+    └── session_<session_id>/
+        ├── session.json                # Session metadata
+        └── agents/
+            └── agent_<agent_id>/
+                ├── agent.json          # Agent metadata
+                └── messages/
+                    ├── message_<created_timestamp>_<id1>.json
+                    └── message_<created_timestamp>_<id2>.json
+
+    """
 
     def __init__(self, session_id: str, storage_dir: Optional[str] = None):
         """Initialize FileSession with filesystem storage.
@@ -44,10 +57,22 @@ def _get_agent_path(self, session_id: str, agent_id: str) -> str:
         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."""
+    def _get_message_path(self, session_id: str, agent_id: str, message_id: str, timestamp: str) -> str:
+        """Get message file path.
+
+        Args:
+            session_id: ID of the session
+            agent_id: ID of the agent
+            message_id: ID of the message
+            timestamp: ISO format timestamp to include in filename for sorting
+        Returns:
+            The filename for the message
+        """
         agent_path = self._get_agent_path(session_id, agent_id)
-        return os.path.join(agent_path, "messages", f"{MESSAGE_PREFIX}{message_id}.json")
+        # Use timestamp for sortable filenames
+        # Replace colons and periods in ISO format with underscores for filesystem compatibility
+        filename_timestamp = timestamp.replace(":", "_").replace(".", "_")
+        return os.path.join(agent_path, "messages", f"{MESSAGE_PREFIX}{filename_timestamp}_{message_id}.json")
 
     def _read_file(self, path: str) -> dict[str, Any]:
         """Read JSON file."""
@@ -135,17 +160,26 @@ def create_message(self, session_id: str, agent_id: str, session_message: Sessio
             session_id,
             agent_id,
             session_message.message_id,
+            session_message.created_at,
         )
         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):
+        # Get the messages directory
+        messages_dir = os.path.join(self._get_agent_path(session_id, agent_id), "messages")
+        if not os.path.exists(messages_dir):
             return None
-        message_data = self._read_file(message_file)
-        return SessionMessage.from_dict(message_data)
+
+        # List files in messages directory, and check if the filename ends with the message id
+        for filename in os.listdir(messages_dir):
+            if filename.endswith(f"{message_id}.json"):
+                file_path = os.path.join(messages_dir, filename)
+                message_data = self._read_file(file_path)
+                return SessionMessage.from_dict(message_data)
+
+        return None
 
     def update_message(self, session_id: str, agent_id: str, session_message: SessionMessage) -> None:
         """Update message data."""
@@ -156,7 +190,7 @@ def update_message(self, session_id: str, agent_id: str, session_message: Sessio
 
         # 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)
+        message_file = self._get_message_path(session_id, agent_id, message_id, session_message.created_at)
         self._write_file(message_file, asdict(session_message))
 
     def list_messages(
@@ -168,20 +202,25 @@ def list_messages(
             raise SessionException(f"Messages directory missing from agent: {agent_id} in session {session_id}")
 
         # Read all message files
-        messages: list[SessionMessage] = []
+        message_files: list[str] = []
         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))
+                message_files.append(filename)
 
-        # Sort by created_at timestamp (oldest first)
-        messages.sort(key=lambda x: x.created_at)
+        # Sort filenames - the timestamp in the file's name will sort chronologically
+        message_files.sort()
 
-        # Apply pagination
+        # Apply pagination to filenames
         if limit is not None:
-            messages = messages[offset : offset + limit]
+            message_files = message_files[offset : offset + limit]
         else:
-            messages = messages[offset:]
+            message_files = message_files[offset:]
+
+        # Load only the message files
+        messages: list[SessionMessage] = []
+        for filename in message_files:
+            file_path = os.path.join(messages_dir, filename)
+            message_data = self._read_file(file_path)
+            messages.append(SessionMessage.from_dict(message_data))
 
         return messages

src/strands/session/agent_session_manager.py renamed to src/strands/session/repository_session_manager.py

@@ -1,8 +1,8 @@
-"""Agent session manager implementation."""
+"""Repository session manager implementation."""
 
 import logging
 
-from ..agent.agent import _DEFAULT_AGENT_ID, Agent
+from ..agent.agent import Agent
 from ..agent.state import AgentState
 from ..types.content import Message
 from ..types.exceptions import SessionException
@@ -17,29 +17,38 @@
 
 logger = logging.getLogger(__name__)
 
-DEFAULT_SESSION_AGENT_ID = "default"
 
-
-class AgentSessionManager(SessionManager):
-    """Session manager for persisting agent's in a Session."""
+class RepositorySessionManager(SessionManager):
+    """Session manager for persisting agents in a SessionRepository."""
 
     def __init__(
         self,
         session_id: str,
         session_repository: SessionRepository,
     ):
-        """Initialize the AgentSessionManager."""
+        """Initialize the RepositorySessionManager.
+
+        If no session with the specified session_id exists yet, it will be created
+        in the session_repository.
+
+        Args:
+          session_id: ID to use for the session. A new session with this id will be created if it does
+              not exist in the reposiory yet
+          session_repository: Underlying session repository to use to store the sessions state.
+        """
         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)
+            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
+
+        # Keep track of the initialized agent id's so that two agents in a session cannot share an id
+        self._initialized_agent_ids: set[str] = set()
 
     def append_message(self, message: Message, agent: Agent) -> None:
         """Append a message to the agent's session.
@@ -49,12 +58,10 @@ def append_message(self, message: Message, agent: Agent) -> None:
             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.
+        """Serialize and update the agent into the session repository.
 
         Args:
             agent: Agent to sync to the session.
@@ -70,16 +77,15 @@ def initialize(self, agent: Agent) -> None:
         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
+        if agent.agent_id in self._initialized_agent_ids:
+            raise SessionException("The `agent_id` of an agent must be unique in a session.")
+        self._initialized_agent_ids.add(agent.agent_id)
 
         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_id=<%s> | session_id=<%s> | creating agent",
                 agent.agent_id,
                 self.session_id,
             )
@@ -91,7 +97,7 @@ def initialize(self, agent: Agent) -> None:
                 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_id=<%s> | session_id=<%s> | restoring agent",
                 agent.agent_id,
                 self.session_id,
             )

src/strands/session/s3_session_manager.py

@@ -11,7 +11,7 @@
 
 from ..types.exceptions import SessionException
 from ..types.session import Session, SessionAgent, SessionMessage
-from .agent_session_manager import AgentSessionManager
+from .repository_session_manager import RepositorySessionManager
 from .session_repository import SessionRepository
 
 logger = logging.getLogger(__name__)
@@ -21,8 +21,21 @@
 MESSAGE_PREFIX = "message_"
 
 
-class S3SessionManager(AgentSessionManager, SessionRepository):
-    """S3-based session manager for cloud storage."""
+class S3SessionManager(RepositorySessionManager, SessionRepository):
+    """S3-based session manager for cloud storage.
+
+    Creates the following filesystem structure for the session storage:
+    /<sessions_dir>/
+    └── session_<session_id>/
+        ├── session.json                # Session metadata
+        └── agents/
+            └── agent_<agent_id>/
+                ├── agent.json          # Agent metadata
+                └── messages/
+                    ├── message_<created_timestamp>_<id1>.json
+                    └── message_<created_timestamp>_<id2>.json
+
+    """
 
     def __init__(
         self,
@@ -72,10 +85,22 @@ def _get_agent_path(self, session_id: str, agent_id: str) -> str:
         session_path = self._get_session_path(session_id)
         return f"{session_path}agents/{AGENT_PREFIX}{agent_id}/"
 
-    def _get_message_path(self, session_id: str, agent_id: str, message_id: str) -> str:
-        """Get message S3 key."""
+    def _get_message_path(self, session_id: str, agent_id: str, message_id: str, timestamp: str) -> str:
+        """Get message S3 key.
+
+        Args:
+            session_id: ID of the session
+            agent_id: ID of the agent
+            message_id: ID of the message
+            timestamp: ISO format timestamp to include in key for sorting
+        Returns:
+            The key for the message
+        """
         agent_path = self._get_agent_path(session_id, agent_id)
-        return f"{agent_path}messages/{MESSAGE_PREFIX}{message_id}.json"
+        # Use timestamp for sortable keys
+        # Replace colons and periods in ISO format with underscores for filesystem compatibility
+        filename_timestamp = timestamp.replace(":", "_").replace(".", "_")
+        return f"{agent_path}messages/{MESSAGE_PREFIX}{filename_timestamp}_{message_id}.json"
 
     def _read_s3_object(self, key: str) -> Optional[Dict[str, Any]]:
         """Read JSON object from S3."""
@@ -180,16 +205,29 @@ def create_message(self, session_id: str, agent_id: str, session_message: Sessio
         """Create a new message in S3."""
         message_id = session_message.message_id
         message_dict = asdict(session_message)
-        message_key = self._get_message_path(session_id, agent_id, message_id)
+        message_key = self._get_message_path(session_id, agent_id, message_id, session_message.created_at)
         self._write_s3_object(message_key, message_dict)
 
     def read_message(self, session_id: str, agent_id: str, message_id: str) -> Optional[SessionMessage]:
         """Read message data from S3."""
-        message_key = self._get_message_path(session_id, agent_id, message_id)
-        message_data = self._read_s3_object(message_key)
-        if message_data is None:
+        # Get the messages prefix
+        messages_prefix = f"{self._get_agent_path(session_id, agent_id)}messages/"
+        try:
+            paginator = self.client.get_paginator("list_objects_v2")
+            pages = paginator.paginate(Bucket=self.bucket, Prefix=messages_prefix)
+
+            for page in pages:
+                if "Contents" in page:
+                    for obj in page["Contents"]:
+                        if obj["Key"].endswith(f"{message_id}.json"):
+                            message_data = self._read_s3_object(obj["Key"])
+                            if message_data:
+                                return SessionMessage.from_dict(message_data)
+
             return None
-        return SessionMessage.from_dict(message_data)
+
+        except ClientError as e:
+            raise SessionException(f"S3 error reading message: {e}") from e
 
     def update_message(self, session_id: str, agent_id: str, session_message: SessionMessage) -> None:
         """Update message data in S3."""
@@ -200,7 +238,7 @@ def update_message(self, session_id: str, agent_id: str, session_message: Sessio
 
         # Preserve creation timestamp
         session_message.created_at = previous_message.created_at
-        message_key = self._get_message_path(session_id, agent_id, message_id)
+        message_key = self._get_message_path(session_id, agent_id, message_id, session_message.created_at)
         self._write_s3_object(message_key, asdict(session_message))
 
     def list_messages(
@@ -212,24 +250,29 @@ def list_messages(
             paginator = self.client.get_paginator("list_objects_v2")
             pages = paginator.paginate(Bucket=self.bucket, Prefix=messages_prefix)
 
-            # Read all message objects
-            messages: List[SessionMessage] = []
+            # Collect all message keys first
+            message_keys = []
             for page in pages:
                 if "Contents" in page:
                     for obj in page["Contents"]:
-                        if obj["Key"].endswith(".json"):
-                            message_data = self._read_s3_object(obj["Key"])
-                            if message_data:
-                                messages.append(SessionMessage.from_dict(message_data))
+                        if obj["Key"].endswith(".json") and MESSAGE_PREFIX in obj["Key"]:
+                            message_keys.append(obj["Key"])
 
-            # Sort by created_at timestamp (oldest first)
-            messages.sort(key=lambda x: x.created_at)
+            # Sort keys - timestamp prefixed keys will sort chronologically
+            message_keys.sort()
 
-            # Apply pagination
+            # Apply pagination to keys before loading content
             if limit is not None:
-                messages = messages[offset : offset + limit]
+                message_keys = message_keys[offset : offset + limit]
             else:
-                messages = messages[offset:]
+                message_keys = message_keys[offset:]
+
+            # Load only the required message objects
+            messages: List[SessionMessage] = []
+            for key in message_keys:
+                message_data = self._read_s3_object(key)
+                if message_data:
+                    messages.append(SessionMessage.from_dict(message_data))
 
             return messages