Add parallel reading support to S3SessionManager.list_messages()

- Add max_parallel_reads parameter to S3SessionManager.__init__()
- Implement parallel S3 reads using ThreadPoolExecutor in list_messages()
- Support both instance-level and per-call max_parallel_reads configuration
- Add comprehensive tests for parallel reading functionality
- Maintain backward compatibility (default max_parallel_reads=1)

Author: CrysisDeu

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, Dict, List, Optional, cast
 
 import boto3
@@ -50,6 +51,7 @@ def __init__(
         boto_session: Optional[boto3.Session] = None,
         boto_client_config: Optional[BotocoreConfig] = None,
         region_name: Optional[str] = None,
+        max_parallel_reads: int = 1,
         **kwargs: Any,
     ):
         """Initialize S3SessionManager with S3 storage.
@@ -62,11 +64,20 @@ def __init__(
             boto_session: Optional boto3 session
             boto_client_config: Optional boto3 client configuration
             region_name: AWS region for S3 storage
+            max_parallel_reads: Maximum number of parallel S3 read operations for list_messages().
+                Defaults to 1 (sequential) for backward compatibility and safety.
+                Set to a higher value (e.g., 10) for better performance with many messages.
+                Can be overridden per-call via list_messages() kwargs.
             **kwargs: Additional keyword arguments for future extensibility.
         """
         self.bucket = bucket
         self.prefix = prefix
 
+        # Validate max_parallel_reads
+        if not isinstance(max_parallel_reads, int) or max_parallel_reads < 1:
+            raise ValueError(f"max_parallel_reads must be a positive integer, got {max_parallel_reads}")
+        self.max_parallel_reads = max_parallel_reads
+
         session = boto_session or boto3.Session(region_name=region_name)
 
         # Add strands-agents to the request user agent
@@ -259,7 +270,24 @@ def update_message(self, session_id: str, agent_id: str, session_message: Sessio
     def list_messages(
         self, session_id: str, agent_id: str, limit: Optional[int] = 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. Supports:
+                max_parallel_reads: Override the instance-level max_parallel_reads setting
+                    for this call only.
+
+        Returns:
+            List of SessionMessage objects, sorted by message_id.
+
+        Raises:
+            ValueError: If max_parallel_reads override is not a positive integer.
+            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 +315,51 @@ 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
+
+            # Use ThreadPoolExecutor to fetch messages concurrently
+            # Allow per-call override of max_parallel_reads via kwargs, otherwise use instance default
+            max_parallel_reads_override = kwargs.get("max_parallel_reads")
+            if max_parallel_reads_override is not None:
+                if not isinstance(max_parallel_reads_override, int) or max_parallel_reads_override < 1:
+                    raise ValueError(
+                        f"max_parallel_reads must be a positive integer, got {max_parallel_reads_override}"
+                    )
+                max_parallel_reads_value = max_parallel_reads_override
+            else:
+                # Instance default was already validated in __init__, no need to check again
+                max_parallel_reads_value = self.max_parallel_reads
+            
+            max_workers = min(max_parallel_reads_value, len(message_keys))
+
+            with ThreadPoolExecutor(max_workers=max_workers) 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[Optional[Dict[str, Any]]] = [None] * len(message_keys)
+
+                # Process results as they complete
+                for future in as_completed(future_to_key):
+                    key = future_to_key[future]
+                    try:
+                        message_data = future.result()
+                        # Store result at the correct index to maintain order
+                        results[key_to_index[key]] = message_data
+                    except Exception as e:
+                        # Log error but continue processing other messages
+                        # Individual failures shouldn't stop the entire operation
+                        logger.warning("key=<%s> | failed to read message from s3", key, exc_info=e)
+
+            # Convert results to SessionMessage objects, filtering out None values
+            # If SessionMessage.from_dict fails, let it propagate - data corruption should be visible
+            for message_data in results:
                 if message_data:
                     messages.append(SessionMessage.from_dict(message_data))