Add parallel reading support to S3SessionManager.list_messages() (strands-agents#1186)

Co-authored-by: Jack Yuan <jackypc@amazon.com>
Co-authored-by: Nicholas Clegg <ncclegg@amazon.com>

Author: 3 people

src/strands/session/s3_session_manager.py

@@ -2,6 +2,7 @@
 
 import json
 import logging
+from concurrent.futures import ThreadPoolExecutor, as_completed
 from typing import TYPE_CHECKING, Any, cast
 
 import boto3
@@ -259,7 +260,21 @@ def update_message(self, session_id: str, agent_id: str, session_message: Sessio
     def list_messages(
         self, session_id: str, agent_id: str, limit: int | None = None, offset: int = 0, **kwargs: Any
     ) -> list[SessionMessage]:
-        """List messages for an agent with pagination from S3."""
+        """List messages for an agent with pagination from S3.
+
+        Args:
+            session_id: ID of the session
+            agent_id: ID of the agent
+            limit: Optional limit on number of messages to return
+            offset: Optional offset for pagination
+            **kwargs: Additional keyword arguments
+
+        Returns:
+            List of SessionMessage objects, sorted by message_id.
+
+        Raises:
+            SessionException: If S3 error occurs during message retrieval.
+        """
         messages_prefix = f"{self._get_agent_path(session_id, agent_id)}messages/"
         try:
             paginator = self.client.get_paginator("list_objects_v2")
@@ -287,10 +302,38 @@ def list_messages(
             else:
                 message_keys = message_keys[offset:]
 
-            # Load only the required message objects
+            # Load message objects in parallel for better performance
             messages: list[SessionMessage] = []
-            for key in message_keys:
-                message_data = self._read_s3_object(key)
+            if not message_keys:
+                return messages
+
+            # Optimize for single worker case - avoid thread pool overhead
+            if len(message_keys) == 1:
+                for key in message_keys:
+                    message_data = self._read_s3_object(key)
+                    if message_data:
+                        messages.append(SessionMessage.from_dict(message_data))
+                return messages
+
+            with ThreadPoolExecutor() as executor:
+                # Submit all read tasks
+                future_to_key = {executor.submit(self._read_s3_object, key): key for key in message_keys}
+
+                # Create a mapping from key to index to maintain order
+                key_to_index = {key: idx for idx, key in enumerate(message_keys)}
+
+                # Initialize results list with None placeholders to maintain order
+                results: list[dict[str, Any] | None] = [None] * len(message_keys)
+
+                # Process results as they complete
+                for future in as_completed(future_to_key):
+                    key = future_to_key[future]
+                    message_data = future.result()
+                    # Store result at the correct index to maintain order
+                    results[key_to_index[key]] = message_data
+
+            # Convert results to SessionMessage objects, filtering out None values
+            for message_data in results:
                 if message_data:
                     messages.append(SessionMessage.from_dict(message_data))
 

tests/strands/models/test_bedrock.py

@@ -201,10 +201,11 @@ def test__init__region_precedence(mock_client_method, session_cls):
 def test__init__with_endpoint_url(mock_client_method):
     """Test that BedrockModel uses the provided endpoint_url for VPC endpoints."""
     custom_endpoint = "https://vpce-12345-abcde.bedrock-runtime.us-west-2.vpce.amazonaws.com"
-    BedrockModel(endpoint_url=custom_endpoint)
-    mock_client_method.assert_called_with(
-        region_name=DEFAULT_BEDROCK_REGION, config=ANY, service_name=ANY, endpoint_url=custom_endpoint
-    )
+    with unittest.mock.patch.object(os, "environ", {}):
+        BedrockModel(endpoint_url=custom_endpoint)
+        mock_client_method.assert_called_with(
+            region_name=DEFAULT_BEDROCK_REGION, config=ANY, service_name=ANY, endpoint_url=custom_endpoint
+        )
 
 
 def test__init__with_region_and_session_raises_value_error():

tests/strands/session/test_s3_session_manager.py

@@ -282,6 +282,40 @@ def test_list_messages_all(s3_manager, sample_session, sample_agent):
     assert len(result) == 5
 
 
+def test_list_messages_single_message(s3_manager, sample_session, sample_agent):
+    """Test listing all messages from S3."""
+    # Create session and agent
+    s3_manager.create_session(sample_session)
+    s3_manager.create_agent(sample_session.session_id, sample_agent)
+
+    # Create single message
+    message = SessionMessage(
+        {
+            "role": "user",
+            "content": [ContentBlock(text="Single Message")],
+        },
+        0,
+    )
+    s3_manager.create_message(sample_session.session_id, sample_agent.agent_id, message)
+
+    # List all messages
+    result = s3_manager.list_messages(sample_session.session_id, sample_agent.agent_id)
+
+    assert len(result) == 1
+
+
+def test_list_no_messages(s3_manager, sample_session, sample_agent):
+    """Test listing all messages from S3."""
+    # Create session and agent
+    s3_manager.create_session(sample_session)
+    s3_manager.create_agent(sample_session.session_id, sample_agent)
+
+    # List all messages
+    result = s3_manager.list_messages(sample_session.session_id, sample_agent.agent_id)
+
+    assert len(result) == 0
+
+
 def test_list_messages_with_pagination(s3_manager, sample_session, sample_agent):
     """Test listing messages with pagination in S3."""
     # Create session and agent