feat: Add kwargs to session interfaces for future extensibility (#464)

Author: Unshure

src/strands/session/file_session_manager.py

@@ -35,12 +35,13 @@ class FileSessionManager(RepositorySessionManager, SessionRepository):
 
     """
 
-    def __init__(self, session_id: str, storage_dir: Optional[str] = None):
+    def __init__(self, session_id: str, storage_dir: Optional[str] = None, **kwargs: Any):
         """Initialize FileSession with filesystem storage.
 
         Args:
             session_id: ID for the session
             storage_dir: Directory for local filesystem storage (defaults to temp dir)
+            **kwargs: Additional keyword arguments for future extensibility.
         """
         self.storage_dir = storage_dir or os.path.join(tempfile.gettempdir(), "strands/sessions")
         os.makedirs(self.storage_dir, exist_ok=True)
@@ -83,7 +84,7 @@ def _write_file(self, path: str, data: dict[str, Any]) -> None:
         with open(path, "w", encoding="utf-8") as f:
             json.dump(data, f, indent=2, ensure_ascii=False)
 
-    def create_session(self, session: Session) -> Session:
+    def create_session(self, session: Session, **kwargs: Any) -> Session:
         """Create a new session."""
         session_dir = self._get_session_path(session.session_id)
         if os.path.exists(session_dir):
@@ -100,7 +101,7 @@ def create_session(self, session: Session) -> Session:
 
         return session
 
-    def read_session(self, session_id: str) -> Optional[Session]:
+    def read_session(self, session_id: str, **kwargs: Any) -> Optional[Session]:
         """Read session data."""
         session_file = os.path.join(self._get_session_path(session_id), "session.json")
         if not os.path.exists(session_file):
@@ -109,7 +110,15 @@ def read_session(self, session_id: str) -> Optional[Session]:
         session_data = self._read_file(session_file)
         return Session.from_dict(session_data)
 
-    def create_agent(self, session_id: str, session_agent: SessionAgent) -> None:
+    def delete_session(self, session_id: str, **kwargs: Any) -> None:
+        """Delete session and all associated data."""
+        session_dir = self._get_session_path(session_id)
+        if not os.path.exists(session_dir):
+            raise SessionException(f"Session {session_id} does not exist")
+
+        shutil.rmtree(session_dir)
+
+    def create_agent(self, session_id: str, session_agent: SessionAgent, **kwargs: Any) -> None:
         """Create a new agent in the session."""
         agent_id = session_agent.agent_id
 
@@ -121,15 +130,7 @@ def create_agent(self, session_id: str, session_agent: SessionAgent) -> None:
         session_data = session_agent.to_dict()
         self._write_file(agent_file, session_data)
 
-    def delete_session(self, session_id: str) -> None:
-        """Delete session and all associated data."""
-        session_dir = self._get_session_path(session_id)
-        if not os.path.exists(session_dir):
-            raise SessionException(f"Session {session_id} does not exist")
-
-        shutil.rmtree(session_dir)
-
-    def read_agent(self, session_id: str, agent_id: str) -> Optional[SessionAgent]:
+    def read_agent(self, session_id: str, agent_id: str, **kwargs: Any) -> Optional[SessionAgent]:
         """Read agent data."""
         agent_file = os.path.join(self._get_agent_path(session_id, agent_id), "agent.json")
         if not os.path.exists(agent_file):
@@ -138,7 +139,7 @@ def read_agent(self, session_id: str, agent_id: str) -> Optional[SessionAgent]:
         agent_data = self._read_file(agent_file)
         return SessionAgent.from_dict(agent_data)
 
-    def update_agent(self, session_id: str, session_agent: SessionAgent) -> None:
+    def update_agent(self, session_id: str, session_agent: SessionAgent, **kwargs: Any) -> None:
         """Update agent data."""
         agent_id = session_agent.agent_id
         previous_agent = self.read_agent(session_id=session_id, agent_id=agent_id)
@@ -149,7 +150,7 @@ def update_agent(self, session_id: str, session_agent: SessionAgent) -> None:
         agent_file = os.path.join(self._get_agent_path(session_id, agent_id), "agent.json")
         self._write_file(agent_file, session_agent.to_dict())
 
-    def create_message(self, session_id: str, agent_id: str, session_message: SessionMessage) -> None:
+    def create_message(self, session_id: str, agent_id: str, session_message: SessionMessage, **kwargs: Any) -> None:
         """Create a new message for the agent."""
         message_file = self._get_message_path(
             session_id,
@@ -159,15 +160,15 @@ def create_message(self, session_id: str, agent_id: str, session_message: Sessio
         session_dict = session_message.to_dict()
         self._write_file(message_file, session_dict)
 
-    def read_message(self, session_id: str, agent_id: str, message_id: int) -> Optional[SessionMessage]:
+    def read_message(self, session_id: str, agent_id: str, message_id: int, **kwargs: Any) -> Optional[SessionMessage]:
         """Read message data."""
         message_path = self._get_message_path(session_id, agent_id, message_id)
         if not os.path.exists(message_path):
             return None
         message_data = self._read_file(message_path)
         return SessionMessage.from_dict(message_data)
 
-    def update_message(self, session_id: str, agent_id: str, session_message: SessionMessage) -> None:
+    def update_message(self, session_id: str, agent_id: str, session_message: SessionMessage, **kwargs: Any) -> None:
         """Update message data."""
         message_id = session_message.message_id
         previous_message = self.read_message(session_id=session_id, agent_id=agent_id, message_id=message_id)
@@ -180,7 +181,7 @@ def update_message(self, session_id: str, agent_id: str, session_message: Sessio
         self._write_file(message_file, session_message.to_dict())
 
     def list_messages(
-        self, session_id: str, agent_id: str, limit: Optional[int] = None, offset: int = 0
+        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."""
         messages_dir = os.path.join(self._get_agent_path(session_id, agent_id), "messages")

src/strands/session/repository_session_manager.py

@@ -1,7 +1,7 @@
 """Repository session manager implementation."""
 
 import logging
-from typing import Optional
+from typing import Any, Optional
 
 from ..agent.agent import Agent
 from ..agent.state import AgentState
@@ -22,20 +22,18 @@
 class RepositorySessionManager(SessionManager):
     """Session manager for persisting agents in a SessionRepository."""
 
-    def __init__(
-        self,
-        session_id: str,
-        session_repository: SessionRepository,
-    ):
+    def __init__(self, session_id: str, session_repository: SessionRepository, **kwargs: Any):
         """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.
+            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.
+            **kwargs: Additional keyword arguments for future extensibility.
+
         """
         self.session_repository = session_repository
         self.session_id = session_id
@@ -51,12 +49,13 @@ def __init__(
         # Keep track of the latest message of each agent in case we need to redact it.
         self._latest_agent_message: dict[str, Optional[SessionMessage]] = {}
 
-    def append_message(self, message: Message, agent: Agent) -> None:
+    def append_message(self, message: Message, agent: Agent, **kwargs: Any) -> None:
         """Append a message to the agent's session.
 
         Args:
             message: Message to add to the agent in the session
             agent: Agent to append the message to
+            **kwargs: Additional keyword arguments for future extensibility.
         """
         # Calculate the next index (0 if this is the first message, otherwise increment the previous index)
         latest_agent_message = self._latest_agent_message[agent.agent_id]
@@ -69,35 +68,38 @@ def append_message(self, message: Message, agent: Agent) -> None:
         self._latest_agent_message[agent.agent_id] = session_message
         self.session_repository.create_message(self.session_id, agent.agent_id, session_message)
 
-    def redact_latest_message(self, redact_message: Message, agent: Agent) -> None:
+    def redact_latest_message(self, redact_message: Message, agent: Agent, **kwargs: Any) -> None:
         """Redact the latest message appended to the session.
 
         Args:
             redact_message: New message to use that contains the redact content
             agent: Agent to apply the message redaction to
+            **kwargs: Additional keyword arguments for future extensibility.
         """
         latest_agent_message = self._latest_agent_message[agent.agent_id]
         if latest_agent_message is None:
             raise SessionException("No message to redact.")
         latest_agent_message.redact_message = redact_message
         return self.session_repository.update_message(self.session_id, agent.agent_id, latest_agent_message)
 
-    def sync_agent(self, agent: Agent) -> None:
+    def sync_agent(self, agent: Agent, **kwargs: Any) -> None:
         """Serialize and update the agent into the session repository.
 
         Args:
             agent: Agent to sync to the session.
+            **kwargs: Additional keyword arguments for future extensibility.
         """
         self.session_repository.update_agent(
             self.session_id,
             SessionAgent.from_agent(agent),
         )
 
-    def initialize(self, agent: Agent) -> None:
+    def initialize(self, agent: Agent, **kwargs: Any) -> None:
         """Initialize an agent with a session.
 
         Args:
             agent: Agent to initialize from the session
+            **kwargs: Additional keyword arguments for future extensibility.
         """
         if agent.agent_id in self._latest_agent_message:
             raise SessionException("The `agent_id` of an agent must be unique in a session.")

src/strands/session/s3_session_manager.py

@@ -44,6 +44,7 @@ def __init__(
         boto_session: Optional[boto3.Session] = None,
         boto_client_config: Optional[BotocoreConfig] = None,
         region_name: Optional[str] = None,
+        **kwargs: Any,
     ):
         """Initialize S3SessionManager with S3 storage.
 
@@ -54,6 +55,7 @@ def __init__(
             boto_session: Optional boto3 session
             boto_client_config: Optional boto3 client configuration
             region_name: AWS region for S3 storage
+            **kwargs: Additional keyword arguments for future extensibility.
         """
         self.bucket = bucket
         self.prefix = prefix
@@ -91,6 +93,8 @@ def _get_message_path(self, session_id: str, agent_id: str, message_id: int) ->
             session_id: ID of the session
             agent_id: ID of the agent
             message_id: Index of the message
+            **kwargs: Additional keyword arguments for future extensibility.
+
         Returns:
             The key for the message
         """
@@ -121,7 +125,7 @@ def _write_s3_object(self, key: str, data: Dict[str, Any]) -> None:
         except ClientError as e:
             raise SessionException(f"Failed to write S3 object {key}: {e}") from e
 
-    def create_session(self, session: Session) -> Session:
+    def create_session(self, session: Session, **kwargs: Any) -> Session:
         """Create a new session in S3."""
         session_key = f"{self._get_session_path(session.session_id)}session.json"
 
@@ -138,15 +142,15 @@ def create_session(self, session: Session) -> Session:
         self._write_s3_object(session_key, session_dict)
         return session
 
-    def read_session(self, session_id: str) -> Optional[Session]:
+    def read_session(self, session_id: str, **kwargs: Any) -> Optional[Session]:
         """Read session data from S3."""
         session_key = f"{self._get_session_path(session_id)}session.json"
         session_data = self._read_s3_object(session_key)
         if session_data is None:
             return None
         return Session.from_dict(session_data)
 
-    def delete_session(self, session_id: str) -> None:
+    def delete_session(self, session_id: str, **kwargs: Any) -> None:
         """Delete session and all associated data from S3."""
         session_prefix = self._get_session_path(session_id)
         try:
@@ -169,22 +173,22 @@ def delete_session(self, session_id: str) -> None:
         except ClientError as e:
             raise SessionException(f"S3 error deleting session {session_id}: {e}") from e
 
-    def create_agent(self, session_id: str, session_agent: SessionAgent) -> None:
+    def create_agent(self, session_id: str, session_agent: SessionAgent, **kwargs: Any) -> None:
         """Create a new agent in S3."""
         agent_id = session_agent.agent_id
         agent_dict = session_agent.to_dict()
         agent_key = f"{self._get_agent_path(session_id, agent_id)}agent.json"
         self._write_s3_object(agent_key, agent_dict)
 
-    def read_agent(self, session_id: str, agent_id: str) -> Optional[SessionAgent]:
+    def read_agent(self, session_id: str, agent_id: str, **kwargs: Any) -> Optional[SessionAgent]:
         """Read agent data from S3."""
         agent_key = f"{self._get_agent_path(session_id, agent_id)}agent.json"
         agent_data = self._read_s3_object(agent_key)
         if agent_data is None:
             return None
         return SessionAgent.from_dict(agent_data)
 
-    def update_agent(self, session_id: str, session_agent: SessionAgent) -> None:
+    def update_agent(self, session_id: str, session_agent: SessionAgent, **kwargs: Any) -> None:
         """Update agent data in S3."""
         agent_id = session_agent.agent_id
         previous_agent = self.read_agent(session_id=session_id, agent_id=agent_id)
@@ -196,22 +200,22 @@ def update_agent(self, session_id: str, session_agent: SessionAgent) -> None:
         agent_key = f"{self._get_agent_path(session_id, agent_id)}agent.json"
         self._write_s3_object(agent_key, session_agent.to_dict())
 
-    def create_message(self, session_id: str, agent_id: str, session_message: SessionMessage) -> None:
+    def create_message(self, session_id: str, agent_id: str, session_message: SessionMessage, **kwargs: Any) -> None:
         """Create a new message in S3."""
         message_id = session_message.message_id
         message_dict = session_message.to_dict()
         message_key = self._get_message_path(session_id, agent_id, message_id)
         self._write_s3_object(message_key, message_dict)
 
-    def read_message(self, session_id: str, agent_id: str, message_id: int) -> Optional[SessionMessage]:
+    def read_message(self, session_id: str, agent_id: str, message_id: int, **kwargs: Any) -> 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:
             return None
         return SessionMessage.from_dict(message_data)
 
-    def update_message(self, session_id: str, agent_id: str, session_message: SessionMessage) -> None:
+    def update_message(self, session_id: str, agent_id: str, session_message: SessionMessage, **kwargs: Any) -> None:
         """Update message data in S3."""
         message_id = session_message.message_id
         previous_message = self.read_message(session_id=session_id, agent_id=agent_id, message_id=message_id)
@@ -224,7 +228,7 @@ def update_message(self, session_id: str, agent_id: str, session_message: Sessio
         self._write_s3_object(message_key, session_message.to_dict())
 
     def list_messages(
-        self, session_id: str, agent_id: str, limit: Optional[int] = None, offset: int = 0
+        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."""
         messages_prefix = f"{self._get_agent_path(session_id, agent_id)}messages/"

src/strands/session/session_manager.py

@@ -35,35 +35,39 @@ def register_hooks(self, registry: HookRegistry, **kwargs: Any) -> None:
         registry.add_callback(AfterInvocationEvent, lambda event: self.sync_agent(event.agent))
 
     @abstractmethod
-    def redact_latest_message(self, redact_message: Message, agent: "Agent") -> None:
+    def redact_latest_message(self, redact_message: Message, agent: "Agent", **kwargs: Any) -> None:
         """Redact the message most recently appended to the agent in the session.
 
         Args:
             redact_message: New message to use that contains the redact content
             agent: Agent to apply the message redaction to
+            **kwargs: Additional keyword arguments for future extensibility.
         """
 
     @abstractmethod
-    def append_message(self, message: Message, agent: "Agent") -> None:
+    def append_message(self, message: Message, agent: "Agent", **kwargs: Any) -> None:
         """Append a message to the agent's session.
 
         Args:
             message: Message to add to the agent in the session
             agent: Agent to append the message to
+            **kwargs: Additional keyword arguments for future extensibility.
         """
 
     @abstractmethod
-    def sync_agent(self, agent: "Agent") -> None:
+    def sync_agent(self, agent: "Agent", **kwargs: Any) -> None:
         """Serialize and sync the agent with the session storage.
 
         Args:
             agent: Agent who should be synchronized with the session storage
+            **kwargs: Additional keyword arguments for future extensibility.
         """
 
     @abstractmethod
-    def initialize(self, agent: "Agent") -> None:
+    def initialize(self, agent: "Agent", **kwargs: Any) -> None:
         """Initialize an agent with a session.
 
         Args:
             agent: Agent to initialize
+            **kwargs: Additional keyword arguments for future extensibility.
         """