fix: Fix session manager agent init (#458)

Author: Unshure

src/strands/session/repository_session_manager.py

@@ -115,17 +115,20 @@ def initialize(self, agent: Agent) -> None:
             session_agent = SessionAgent.from_agent(agent)
             self.session_repository.create_agent(self.session_id, session_agent)
             # Initialize messages with sequential indices
+            session_message = None
             for i, message in enumerate(agent.messages):
                 session_message = SessionMessage.from_message(message, i)
                 self.session_repository.create_message(self.session_id, agent.agent_id, session_message)
+            self._latest_agent_message[agent.agent_id] = session_message
         else:
             logger.debug(
                 "agent_id=<%s> | session_id=<%s> | restoring agent",
                 agent.agent_id,
                 self.session_id,
             )
-            agent.messages = [
-                session_message.to_message()
-                for session_message in self.session_repository.list_messages(self.session_id, agent.agent_id)
-            ]
+            session_messages = self.session_repository.list_messages(self.session_id, agent.agent_id)
+            if len(session_messages) > 0:
+                self._latest_agent_message[agent.agent_id] = session_messages[-1]
+            agent.messages = [session_message.to_message() for session_message in session_messages]
+
             agent.state = AgentState(session_agent.state)

tests/fixtures/mock_session_repository.py

@@ -33,7 +33,7 @@ def create_agent(self, session_id, session_agent):
         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, [])
+        self.messages.setdefault(session_id, {}).setdefault(agent_id, {})
         return session_agent
 
     def read_agent(self, session_id, agent_id):
@@ -53,37 +53,34 @@ def update_agent(self, session_id, session_agent):
 
     def create_message(self, session_id, agent_id, session_message):
         """Create 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}")
-        self.messages.setdefault(session_id, {}).setdefault(agent_id, []).append(session_message)
+            raise SessionException(f"Agent {agent_id} does not exists in session {session_id}")
+        if message_id in self.messages.get(session_id, {}).get(agent_id, {}):
+            raise SessionException(f"Message {message_id} already exists in agent {agent_id} in session {session_id}")
+        self.messages.setdefault(session_id, {}).setdefault(agent_id, {})[message_id] = 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
+        return self.messages.get(session_id, {}).get(agent_id, {}).get(message_id)
 
     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")
+        if message_id not in self.messages.get(session_id, {}).get(agent_id, {}):
+            raise SessionException(f"Message {message_id} does not exist in session {session_id}")
+        self.messages[session_id][agent_id][message_id] = session_message
 
     def list_messages(self, session_id, agent_id, limit=None, offset=0):
         """List messages."""
@@ -92,7 +89,9 @@ def list_messages(self, session_id, agent_id, limit=None, offset=0):
         if agent_id not in self.agents.get(session_id, {}):
             return []
 
-        messages = self.messages.get(session_id, {}).get(agent_id, [])
+        messages = self.messages.get(session_id, {}).get(agent_id, {})
+        sorted_messages = [messages[key] for key in sorted(messages.keys())]
+
         if limit is not None:
-            return messages[offset : offset + limit]
-        return messages[offset:]
+            return sorted_messages[offset : offset + limit]
+        return sorted_messages[offset:]

tests/strands/agent/test_agent.py

@@ -20,7 +20,7 @@
 from strands.session.repository_session_manager import RepositorySessionManager
 from strands.types.content import Messages
 from strands.types.exceptions import ContextWindowOverflowException, EventLoopException
-from strands.types.session import Session, SessionAgent, SessionType
+from strands.types.session import Session, SessionAgent, SessionMessage, SessionType
 from tests.fixtures.mock_session_repository import MockedSessionRepository
 from tests.fixtures.mocked_model_provider import MockedModelProvider
 
@@ -1428,6 +1428,26 @@ def test_agent_restored_from_session_management():
     assert agent.state.get("foo") == "bar"
 
 
+def test_agent_restored_from_session_management_with_message():
+    mock_session_repository = MockedSessionRepository()
+    mock_session_repository.create_session(Session(session_id="123", session_type=SessionType.AGENT))
+    mock_session_repository.create_agent(
+        "123",
+        SessionAgent(
+            agent_id="default",
+            state={"foo": "bar"},
+        ),
+    )
+    mock_session_repository.create_message(
+        "123", "default", SessionMessage({"role": "user", "content": [{"text": "Hello!"}]}, 0)
+    )
+    session_manager = RepositorySessionManager(session_id="123", session_repository=mock_session_repository)
+
+    agent = Agent(session_manager=session_manager)
+
+    assert agent.state.get("foo") == "bar"
+
+
 def test_agent_redacts_input_on_triggered_guardrail():
     mocked_model = MockedModelProvider(
         [{"redactedUserContent": "BLOCKED!", "redactedAssistantContent": "INPUT BLOCKED!"}]
@@ -1484,3 +1504,29 @@ def test_agent_restored_from_session_management_with_redacted_input():
 
     # Assert that the restored agent redacted message is equal to the original agent
     assert agent.messages[0] == agent_2.messages[0]
+
+
+def test_agent_restored_from_session_management_with_correct_index():
+    mock_model_provider = MockedModelProvider(
+        [{"role": "assistant", "content": [{"text": "hello!"}]}, {"role": "assistant", "content": [{"text": "world!"}]}]
+    )
+    mock_session_repository = MockedSessionRepository()
+    session_manager = RepositorySessionManager(session_id="test", session_repository=mock_session_repository)
+    agent = Agent(session_manager=session_manager, model=mock_model_provider)
+    agent("Hello!")
+
+    assert len(mock_session_repository.list_messages("test", agent.agent_id)) == 2
+
+    session_manager_2 = RepositorySessionManager(session_id="test", session_repository=mock_session_repository)
+    agent_2 = Agent(session_manager=session_manager_2, model=mock_model_provider)
+
+    assert len(agent_2.messages) == 2
+    assert agent_2.messages[1]["content"][0]["text"] == "hello!"
+
+    agent_2("Hello!")
+
+    assert len(agent_2.messages) == 4
+    session_messages = mock_session_repository.list_messages("test", agent_2.agent_id)
+    assert (len(session_messages)) == 4
+    assert session_messages[1].message["content"][0]["text"] == "hello!"
+    assert session_messages[3].message["content"][0]["text"] == "world!"