fix(memory): restore latest agent snapshot from batched events

src/bedrock_agentcore/memory/integrations/strands/session_manager.py

@@ -400,9 +400,10 @@ def create_agent(self, session_id: str, session_agent: SessionAgent, **kwargs: A
                 self._agent_state_buffer.append((session_id, session_agent))
                 should_flush = len(self._agent_state_buffer) >= self.config.batch_size
 
-            # Flush only agent states outside the lock to prevent deadlock
+            # Flush outside the lock to prevent deadlock. Flush messages too, not just agent states: agent state
+            # depends on the messages (removed_message_count indexes into them), so they must be persisted together.
             if should_flush:
-                self._flush_agent_states_only()
+                self._flush_messages()
 
             logger.info(
                 "Buffered agent creation: %s in session: %s",
@@ -471,7 +472,8 @@ def read_agent(self, session_id: str, agent_id: str, **kwargs: Any) -> Optional[
             )
 
             if events:
-                agent_data = json.loads(events[0].get("payload", {})[0].get("blob"))
+                # A batched flush packs several snapshots into one event, oldest first.
+                agent_data = json.loads(events[0].get("payload", {})[-1].get("blob"))
                 agent = SessionAgent.from_dict(agent_data)
                 # Cache the created_at timestamp to avoid re-fetching on updates
                 if agent.created_at:
@@ -1184,9 +1186,9 @@ def _flush_agent_states_only(self) -> list[dict[str, Any]]:
                 logger.debug("Flushed %d agent states for agent %s: %s", len(payloads), agent_id, event.get("eventId"))
 
         except Exception as e:
-            # Restore agent states to buffer so they aren't lost
+            # Prepend so they aren't lost: states buffered during the send are newer, and read_agent restores the last
             with self._agent_state_lock:
-                self._agent_state_buffer.extend(agent_states_to_send)
+                self._agent_state_buffer[:0] = agent_states_to_send
             logger.error("Failed to flush agent states to AgentCore Memory: %s", e)
             raise SessionException(f"Failed to flush agent states: {e}") from e
 

tests/bedrock_agentcore/memory/integrations/strands/test_agentcore_memory_session_manager.py

@@ -277,6 +277,20 @@ def test_read_agent_valid(self, session_manager, mock_memory_client):
         assert result.agent_id == "test-agent-123"
         assert result.agent_id == "test-agent-123"
 
+    def test_read_agent_batched_event_returns_latest_snapshot(self, batching_session_manager, mock_memory_client):
+        """Test that an agent restored from a batched event gets the last snapshot buffered (#698)."""
+        for state in ({}, {"n": 1}, {"n": 2}):
+            batching_session_manager.create_agent(
+                "test-session-456", SessionAgent(agent_id="test-agent-123", state=state, conversation_manager_state={})
+            )
+        batching_session_manager._flush_agent_states_only()
+        written = mock_memory_client.gmdp_client.create_event.call_args.kwargs
+        mock_memory_client.list_events.return_value = [{"eventId": "event-1", "payload": written["payload"]}]
+
+        result = batching_session_manager.read_agent("test-session-456", "test-agent-123")
+
+        assert result.state == {"n": 2}
+
     def test_read_agent_no_events(self, session_manager, mock_memory_client):
         """Test reading an agent with no events."""
         mock_memory_client.list_events.return_value = []
@@ -3882,6 +3896,50 @@ def test_flush_agent_states_failure_restores_buffer(self, batching_session_manag
         # State must be back in the buffer for retry.
         assert batching_session_manager.pending_agent_state_count() == 1
 
+    def test_flush_agent_states_failure_keeps_newer_concurrent_state_last(
+        self, batching_session_manager, mock_memory_client
+    ):
+        """A state buffered during a failed flush must stay last, since read_agent restores the last state."""
+
+        def fail_after_concurrent_append(**kwargs):
+            batching_session_manager.create_agent(
+                "test-session-456", SessionAgent(agent_id="agent-1", state={"v": 2}, conversation_manager_state={})
+            )
+            raise Exception("API Error")
+
+        batching_session_manager.create_agent(
+            "test-session-456", SessionAgent(agent_id="agent-1", state={"v": 1}, conversation_manager_state={})
+        )
+        mock_memory_client.gmdp_client.create_event.side_effect = fail_after_concurrent_append
+        with pytest.raises(SessionException):
+            batching_session_manager._flush_agent_states_only()
+
+        assert [a.state for _, a in batching_session_manager._agent_state_buffer] == [{"v": 1}, {"v": 2}]
+
+    def test_agent_state_auto_flush_persists_buffered_messages_first(
+        self, batching_session_manager, mock_memory_client
+    ):
+        """An agent snapshot must never be persisted ahead of the messages it accounts for."""
+        mock_memory_client.gmdp_client.create_event.return_value = {"event": {"eventId": "e"}}
+        batching_session_manager.create_message(
+            "test-session-456",
+            "agent-1",
+            SessionMessage(message={"role": "user", "content": [{"text": "hi"}]}, message_id=0),
+        )
+
+        for i in range(batching_session_manager.config.batch_size):
+            batching_session_manager.create_agent(
+                "test-session-456", SessionAgent(agent_id="agent-1", state={"v": i}, conversation_manager_state={})
+            )
+
+        written = [
+            c.kwargs.get("metadata", {}).get("stateType")
+            for c in mock_memory_client.gmdp_client.create_event.call_args_list
+        ]
+        assert written[-1] == {"stringValue": "AGENT"}
+        assert batching_session_manager.pending_message_count() == 0
+        assert len(written) == 2
+
 
 class TestMonotonicTimestamp:
     """Tests for monotonic event-timestamp ordering across processes/pods.