feat: Updating session tests

Author: Unshure

src/strands/agent/agent.py

@@ -15,7 +15,6 @@
 import random
 from concurrent.futures import ThreadPoolExecutor
 from typing import Any, AsyncGenerator, AsyncIterator, Callable, Mapping, Optional, Type, TypeVar, Union, cast
-from uuid import uuid4
 
 from opentelemetry import trace
 from pydantic import BaseModel
@@ -259,7 +258,7 @@ def __init__(
         self.messages = messages if messages is not None else []
 
         self.system_prompt = system_prompt
-        self.agent_id = agent_id or str(uuid4())
+        self.agent_id = agent_id
         self.name = name or _DEFAULT_AGENT_NAME
         self.description = description
 
@@ -330,7 +329,6 @@ def __init__(
                 self.hooks.add_hook(hook)
         self.hooks.invoke_callbacks(AgentInitializedEvent(agent=self))
 
-
     @property
     def tool(self) -> ToolCaller:
         """Call tool as a function.

src/strands/event_loop/event_loop.py

@@ -454,6 +454,7 @@ def tool_handler(tool_use: ToolUse) -> ToolGenerator:
         "content": [{"toolResult": result} for result in tool_results],
     }
 
+    agent.messages.append(tool_result_message)
     agent.hooks.invoke_callbacks(MessageAddedEvent(agent=agent, message=tool_result_message))
     yield {"callback": {"message": tool_result_message}}
 

src/strands/session/agent_session_manager.py

@@ -1,10 +1,9 @@
 """Agent session manager implementation."""
 
 import logging
-from typing import TYPE_CHECKING
 
 from ..agent.state import AgentState
-from ..experimental.hooks.events import AgentInitializedEvent, MessageAddedEvent
+from ..hooks.events import AgentInitializedEvent, MessageAddedEvent
 from ..telemetry.metrics import EventLoopMetrics
 from ..types.session import (
     SessionType,
@@ -18,9 +17,6 @@
 
 logger = logging.getLogger(__name__)
 
-if TYPE_CHECKING:
-    pass
-
 DEFAULT_SESSION_AGENT_ID = "default"
 
 
@@ -41,9 +37,6 @@ def __init__(
             logger.debug("session_id=<%s> | Session not found, creating new session.", self.session_id)
             session = create_session(session_id=session_id, session_type=SessionType.AGENT)
             session_repository.create_session(session)
-        else:
-            if session["session_type"] != SessionType.AGENT:
-                raise ValueError(f"Invalid session type: {session.session_type}")
 
         self.session = session
         self._default_agent_initialized = False

src/strands/session/file_session_manager.py

@@ -50,19 +50,14 @@ def _read_file(self, path: str) -> dict[str, Any]:
         try:
             with open(path, "r", encoding="utf-8") as f:
                 return cast(dict[str, Any], json.load(f))
-        except FileNotFoundError as e:
-            raise SessionException(f"File not found: {path}") from e
         except json.JSONDecodeError as e:
             raise SessionException(f"Invalid JSON in file {path}: {e}") from e
 
     def _write_file(self, path: str, data: dict[str, Any]) -> None:
         """Write JSON file."""
-        try:
-            os.makedirs(os.path.dirname(path), exist_ok=True)
-            with open(path, "w", encoding="utf-8") as f:
-                json.dump(data, f, indent=2, ensure_ascii=False)
-        except Exception as e:
-            raise SessionException(f"Failed to write file {path}: {e}") from e
+        os.makedirs(os.path.dirname(path), exist_ok=True)
+        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:
         """Create a new session."""
@@ -132,12 +127,11 @@ def create_message(self, session_id: str, agent_id: str, session_message: Sessio
         session_dict = cast(dict, session_message)
         self._write_file(message_file, session_dict)
 
-    def read_message(self, session_id: str, agent_id: str, message_id: str) -> SessionMessage:
+    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):
-            raise SessionException(f"Message {message_id} does not exist for agent {agent_id} in session {session_id}")
-
+            return None
         message_data = self._read_file(message_file)
         return SessionMessage(**message_data)  # type: ignore
 
@@ -154,7 +148,7 @@ def list_messages(
         """List messages for an agent with pagination."""
         messages_dir = os.path.join(self._get_agent_path(session_id, agent_id), "messages")
         if not os.path.exists(messages_dir):
-            return []
+            raise SessionException("messages directory missing from agent: %s in session %s", agent_id, session_id)
 
         # Get all message files and sort by creation time (newest first)
         message_files = []
@@ -175,6 +169,8 @@ def list_messages(
         # Read message data
         messages: list[SessionMessage] = []
         for file_path, _ in message_files:
+            if not os.path.exists(file_path):
+                continue
             message_data = self._read_file(file_path)
             messages.append(SessionMessage(**message_data))  # type: ignore
 

src/strands/session/s3_session_manager.py

@@ -1,4 +1,4 @@
-"""S3-based session DAO for cloud storage."""
+"""S3-based session manager for cloud storage."""
 
 import json
 from typing import Any, Dict, List, Optional, cast
@@ -18,7 +18,7 @@
 
 
 class S3SessionManager(AgentSessionManager, SessionRepository):
-    """S3-based session DAO for cloud storage."""
+    """S3-based session manager for cloud storage."""
 
     def __init__(
         self,
@@ -29,7 +29,7 @@ def __init__(
         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
@@ -163,6 +163,10 @@ def read_agent(self, session_id: str, agent_id: str) -> Optional[SessionAgent]:
     def update_agent(self, session_id: str, session_agent: SessionAgent) -> 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)
+        if previous_agent is None:
+            raise SessionException(f"Agent {agent_id} in session {session_id} does not exist")
+        session_agent["created_at"] = previous_agent["created_at"]
         agent_dict = cast(dict, session_agent)
         agent_key = f"{self._get_agent_path(session_id, agent_id)}agent.json"
         self._write_s3_object(agent_key, agent_dict)
@@ -185,6 +189,12 @@ def read_message(self, session_id: str, agent_id: str, message_id: str) -> Optio
     def update_message(self, session_id: str, agent_id: str, session_message: SessionMessage) -> 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=session_message["message_id"]
+        )
+        if previous_message is None:
+            raise SessionException(f"Message {message_id} does not exist")
+        session_message["created_at"] = previous_message["created_at"]
         message_dict = cast(dict, session_message)
         message_key = self._get_message_path(session_id, agent_id, message_id)
         self._write_s3_object(message_key, message_dict)

src/strands/session/session_manager.py

@@ -1,9 +1,10 @@
 """Session manager interface for agent session management."""
 
 from abc import ABC, abstractmethod
+from typing import Any
 
-from ..experimental.hooks.events import AgentInitializedEvent, MessageAddedEvent
-from ..experimental.hooks.registry import HookProvider, HookRegistry
+from ..hooks.events import AgentInitializedEvent, MessageAddedEvent
+from ..hooks.registry import HookProvider, HookRegistry
 
 
 class SessionManager(HookProvider, ABC):
@@ -14,7 +15,7 @@ class SessionManager(HookProvider, ABC):
     methods to manage sessions and their associated data.
     """
 
-    def register_hooks(self, registry: HookRegistry) -> None:
+    def register_hooks(self, registry: HookRegistry, **kwargs: Any) -> None:
         """Register initialize and append_message as hooks for the Agent."""
         registry.add_callback(AgentInitializedEvent, self.initialize)
         registry.add_callback(MessageAddedEvent, self.append_message)

tests/fixtures/mock_session_repository.py

@@ -0,0 +1,98 @@
+from strands.session.session_repository import SessionRepository
+from strands.types.exceptions import SessionException
+
+
+class MockedSessionRepository(SessionRepository):
+    """Mock repository for testing."""
+
+    def __init__(self):
+        """Initialize with empty storage."""
+        self.sessions = {}
+        self.agents = {}
+        self.messages = {}
+
+    def create_session(self, session):
+        """Create a session."""
+        session_id = session["session_id"]
+        if session_id in self.sessions:
+            raise SessionException(f"Session {session_id} already exists")
+        self.sessions[session_id] = session
+        self.agents[session_id] = {}
+        self.messages[session_id] = {}
+        return session
+
+    def read_session(self, session_id):
+        """Read a session."""
+        return self.sessions.get(session_id)
+
+    def create_agent(self, session_id, session_agent):
+        """Create an agent."""
+        agent_id = session_agent["agent_id"]
+        if session_id not in self.sessions:
+            raise SessionException(f"Session {session_id} does not exist")
+        if agent_id in self.agents.get(session_id, {}):
+            raise SessionException(f"Agent {agent_id} already exists in session {session_id}")
+        self.agents.setdefault(session_id, {})[agent_id] = session_agent
+        self.messages.setdefault(session_id, {}).setdefault(agent_id, [])
+        return session_agent
+
+    def read_agent(self, session_id, agent_id):
+        """Read an agent."""
+        if session_id not in self.sessions:
+            return None
+        return self.agents.get(session_id, {}).get(agent_id)
+
+    def update_agent(self, session_id, session_agent):
+        """Update an agent."""
+        agent_id = session_agent["agent_id"]
+        if session_id not in self.sessions:
+            raise SessionException(f"Session {session_id} does not exist")
+        if agent_id not in self.agents.get(session_id, {}):
+            raise SessionException(f"Agent {agent_id} does not exist in session {session_id}")
+        self.agents[session_id][agent_id] = session_agent
+
+    def create_message(self, session_id, agent_id, session_message):
+        """Create a message."""
+        if session_id not in self.sessions:
+            raise SessionException(f"Session {session_id} does not exist")
+        if agent_id not in self.agents.get(session_id, {}):
+            raise SessionException(f"Agent {agent_id} does not exist in session {session_id}")
+        self.messages.setdefault(session_id, {}).setdefault(agent_id, []).append(session_message)
+
+    def read_message(self, session_id, agent_id, message_id):
+        """Read a message."""
+        if session_id not in self.sessions:
+            return None
+        if agent_id not in self.agents.get(session_id, {}):
+            return None
+        for message in self.messages.get(session_id, {}).get(agent_id, []):
+            if message["message_id"] == message_id:
+                return message
+        return None
+
+    def update_message(self, session_id, agent_id, session_message):
+        """Update a message."""
+        message_id = session_message["message_id"]
+        if session_id not in self.sessions:
+            raise SessionException(f"Session {session_id} does not exist")
+        if agent_id not in self.agents.get(session_id, {}):
+            raise SessionException(f"Agent {agent_id} does not exist in session {session_id}")
+
+        for i, message in enumerate(self.messages.get(session_id, {}).get(agent_id, [])):
+            if message["message_id"] == message_id:
+                self.messages[session_id][agent_id][i] = session_message
+                return
+
+        raise SessionException(f"Message {message_id} does not exist")
+
+    def list_messages(self, session_id, agent_id, limit=None, offset=0):
+        """List messages."""
+        if session_id not in self.sessions:
+            return []
+        if agent_id not in self.agents.get(session_id, {}):
+            return []
+
+        messages = self.messages.get(session_id, {}).get(agent_id, [])
+        if limit is not None:
+            return messages[offset : offset + limit]
+        return messages[offset:]

tests/strands/agent/test_agent.py

@@ -13,10 +13,16 @@
 from strands.agent import AgentResult
 from strands.agent.conversation_manager.null_conversation_manager import NullConversationManager
 from strands.agent.conversation_manager.sliding_window_conversation_manager import SlidingWindowConversationManager
+from strands.agent.state import AgentState
 from strands.handlers.callback_handler import PrintingCallbackHandler, null_callback_handler
 from strands.models.bedrock import DEFAULT_BEDROCK_MODEL_ID, BedrockModel
+from strands.session.agent_session_manager import DEFAULT_SESSION_AGENT_ID, AgentSessionManager
+from strands.telemetry.metrics import EventLoopMetrics
 from strands.types.content import Messages
 from strands.types.exceptions import ContextWindowOverflowException, EventLoopException
+from strands.types.session import Session, SessionAgent, SessionType
+from tests.fixtures.mock_session_repository import MockedSessionRepository
+from tests.fixtures.mocked_model_provider import MockedModelProvider
 
 
 @pytest.fixture
@@ -1313,6 +1319,11 @@ async def test_agent_stream_async_creates_and_ends_span_on_exception(mock_get_tr
     mock_tracer.end_agent_span.assert_called_once_with(span=mock_span, error=test_exception)
 
 
+def test_agent_init_with_state_object():
+    agent = Agent(state=AgentState({"foo": "bar"}))
+    assert agent.state.get("foo") == "bar"
+
+
 def test_non_dict_throws_error():
     with pytest.raises(ValueError, match="state must be an AgentState object or a dict"):
         agent = Agent(state={"object", object()})
@@ -1366,3 +1377,38 @@ def test_agent_state_get_breaks_deep_dict_reference():
 
     # This will fail if AgentState reflects the updated reference
     json.dumps(agent.state.get())
+
+
+def test_agent_session_management():
+    mock_session_repository = MockedSessionRepository()
+    session_manager = AgentSessionManager(session_id="123", session_repository=mock_session_repository)
+    model = MockedModelProvider([{"role": "assistant", "content": [{"text": "hello!"}]}])
+    agent = Agent(session_manager=session_manager, model=model)
+    agent("Hello!")
+
+
+def test_agent_restored_from_session_management():
+    mock_session_repository = MockedSessionRepository()
+    mock_session_repository.create_session(
+        Session(
+            session_id="123",
+            session_type=SessionType.AGENT,
+            created_at="2025-01-01T00:00:00Z",
+            updated_at="2025-01-01T00:00:00Z",
+        )
+    )
+    mock_session_repository.create_agent(
+        "123",
+        SessionAgent(
+            agent_id=DEFAULT_SESSION_AGENT_ID,
+            event_loop_metrics=EventLoopMetrics().to_dict(),
+            state={"foo": "bar"},
+            created_at="2025-01-01T00:00:00Z",
+            updated_at="2025-01-01T00:00:00Z",
+        ),
+    )
+    session_manager = AgentSessionManager(session_id="123", session_repository=mock_session_repository)
+
+    agent = Agent(session_manager=session_manager)
+
+    assert agent.state.get("foo") == "bar"

tests/strands/session/__init__.py

@@ -0,0 +1 @@
+"""Tests for session management."""