fix(memory): serialize batched flushes so retries and late sends can't overtake newer state

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

@@ -166,6 +166,9 @@ def __init__(
         self._agent_state_buffer: list[tuple[str, SessionAgent]] = []
         self._agent_state_lock = threading.Lock()
 
+        # Serializes flushes so a retry or a late send cannot overtake a newer flush
+        self._flush_lock = threading.RLock()
+
         # Cache for agent created_at timestamps to avoid fetching on every update
         self._agent_created_at_cache: dict[str, datetime] = {}
 
@@ -1071,67 +1074,68 @@ def _flush_messages_only(self) -> list[dict[str, Any]]:
         if self.persistence_mode is PersistenceMode.NONE:
             return []
 
-        with self._message_lock:
-            messages_to_send = list(self._message_buffer)
-            self._message_buffer.clear()
-
-        if not messages_to_send:
-            return []
+        with self._flush_lock:
+            with self._message_lock:
+                messages_to_send = list(self._message_buffer)
+                self._message_buffer.clear()
+
+            if not messages_to_send:
+                return []
+
+            # Group all messages by session_id, combining conversational and blob messages
+            # Structure: {session_id: {"payload": [...], "timestamp": latest_timestamp, "metadata": {...}}}
+            session_groups: dict[str, dict[str, Any]] = {}
+
+            for buffered_msg in messages_to_send:
+                sid = buffered_msg.session_id
+                if sid not in session_groups:
+                    session_groups[sid] = {"payload": [], "timestamp": buffered_msg.timestamp, "metadata": {}}
+
+                if buffered_msg.is_blob:
+                    for msg in buffered_msg.messages:
+                        session_groups[sid]["payload"].append({"blob": json.dumps(msg)})
+                else:
+                    for text, role in buffered_msg.messages:
+                        session_groups[sid]["payload"].append(
+                            {"conversational": {"content": {"text": text}, "role": role.upper()}}
+                        )
 
-        # Group all messages by session_id, combining conversational and blob messages
-        # Structure: {session_id: {"payload": [...], "timestamp": latest_timestamp, "metadata": {...}}}
-        session_groups: dict[str, dict[str, Any]] = {}
+                # Use the latest timestamp for the combined event
+                if buffered_msg.timestamp > session_groups[sid]["timestamp"]:
+                    session_groups[sid]["timestamp"] = buffered_msg.timestamp
 
-        for buffered_msg in messages_to_send:
-            sid = buffered_msg.session_id
-            if sid not in session_groups:
-                session_groups[sid] = {"payload": [], "timestamp": buffered_msg.timestamp, "metadata": {}}
+                # Merge metadata (later entries override earlier for same key)
+                if buffered_msg.metadata:
+                    session_groups[sid]["metadata"].update(buffered_msg.metadata)
 
-            if buffered_msg.is_blob:
-                for msg in buffered_msg.messages:
-                    session_groups[sid]["payload"].append({"blob": json.dumps(msg)})
-            else:
-                for text, role in buffered_msg.messages:
-                    session_groups[sid]["payload"].append(
-                        {"conversational": {"content": {"text": text}, "role": role.upper()}}
+            results = []
+            try:
+                # Send one create_event per session_id with all messages (conversational + blob)
+                for session_id, group in session_groups.items():
+                    create_event_kwargs: dict[str, Any] = {
+                        "memoryId": self.config.memory_id,
+                        "actorId": self.config.actor_id,
+                        "sessionId": session_id,
+                        "payload": group["payload"],
+                        "eventTimestamp": group["timestamp"],
+                    }
+                    if group["metadata"]:
+                        create_event_kwargs["metadata"] = group["metadata"]
+                    event = self.memory_client.gmdp_client.create_event(**create_event_kwargs)
+                    results.append(event)
+                    logger.debug(
+                        "Flushed batched event for session %s with %d messages: %s",
+                        session_id,
+                        len(group["payload"]),
+                        event.get("eventId"),
                     )
 
-            # Use the latest timestamp for the combined event
-            if buffered_msg.timestamp > session_groups[sid]["timestamp"]:
-                session_groups[sid]["timestamp"] = buffered_msg.timestamp
-
-            # Merge metadata (later entries override earlier for same key)
-            if buffered_msg.metadata:
-                session_groups[sid]["metadata"].update(buffered_msg.metadata)
-
-        results = []
-        try:
-            # Send one create_event per session_id with all messages (conversational + blob)
-            for session_id, group in session_groups.items():
-                create_event_kwargs: dict[str, Any] = {
-                    "memoryId": self.config.memory_id,
-                    "actorId": self.config.actor_id,
-                    "sessionId": session_id,
-                    "payload": group["payload"],
-                    "eventTimestamp": group["timestamp"],
-                }
-                if group["metadata"]:
-                    create_event_kwargs["metadata"] = group["metadata"]
-                event = self.memory_client.gmdp_client.create_event(**create_event_kwargs)
-                results.append(event)
-                logger.debug(
-                    "Flushed batched event for session %s with %d messages: %s",
-                    session_id,
-                    len(group["payload"]),
-                    event.get("eventId"),
-                )
-
-        except Exception as e:
-            # Restore messages to buffer so they aren't lost
-            with self._message_lock:
-                self._message_buffer.extend(messages_to_send)
-            logger.error("Failed to flush messages to AgentCore Memory: %s", e)
-            raise SessionException(f"Failed to flush messages: {e}") from e
+            except Exception as e:
+                # Prepend so they aren't lost: messages buffered during the send are newer
+                with self._message_lock:
+                    self._message_buffer[:0] = messages_to_send
+                logger.error("Failed to flush messages to AgentCore Memory: %s", e)
+                raise SessionException(f"Failed to flush messages: {e}") from e
 
         logger.info("Flushed %d message events to AgentCore Memory", len(results))
         return results
@@ -1152,14 +1156,26 @@ def _flush_agent_states_only(self) -> list[dict[str, Any]]:
         if self.persistence_mode is PersistenceMode.NONE:
             return []
 
+        with self._flush_lock:
+            return self._send_agent_states(self._take_agent_states())
+
+    def _take_agent_states(self) -> list[tuple[str, SessionAgent]]:
         with self._agent_state_lock:
-            agent_states_to_send = list(self._agent_state_buffer)
+            agent_states = list(self._agent_state_buffer)
             self._agent_state_buffer.clear()
+        return agent_states
+
+    def _restore_agent_states(self, agent_states: list[tuple[str, SessionAgent]]) -> None:
+        # Prepend: states buffered since they were taken are newer, and read_agent restores the last
+        with self._agent_state_lock:
+            self._agent_state_buffer[:0] = agent_states
 
+    def _send_agent_states(self, agent_states_to_send: list[tuple[str, SessionAgent]]) -> list[dict[str, Any]]:
         if not agent_states_to_send:
             return []
 
         results = []
+        sent_agent_ids: set[str] = set()
         try:
             # Group agent states by agent_id
             agent_groups: dict[str, list[dict]] = {}
@@ -1182,13 +1198,14 @@ def _flush_agent_states_only(self) -> list[dict[str, Any]]:
                         AGENT_ID_KEY: {"stringValue": agent_id},
                     },
                 )
+                sent_agent_ids.add(agent_id)
                 results.append(event)
                 logger.debug("Flushed %d agent states for agent %s: %s", len(payloads), agent_id, event.get("eventId"))
 
         except Exception as e:
-            # 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[:0] = agent_states_to_send
+            self._restore_agent_states(
+                [state for state in agent_states_to_send if state[1].agent_id not in sent_agent_ids]
+            )
             logger.error("Failed to flush agent states to AgentCore Memory: %s", e)
             raise SessionException(f"Failed to flush agent states: {e}") from e
 
@@ -1207,10 +1224,18 @@ def _flush_messages(self) -> list[dict[str, Any]]:
         Raises:
             SessionException: If any message or agent state creation fails.
         """
-        results = []
-        results.extend(self._flush_messages_only())
-        results.extend(self._flush_agent_states_only())
-        return results
+        if self.persistence_mode is PersistenceMode.NONE:
+            return []
+
+        with self._flush_lock:
+            # Take agent states before messages: every message a taken state counts was buffered before it
+            agent_states = self._take_agent_states()
+            try:
+                results = self._flush_messages_only()
+            except SessionException:
+                self._restore_agent_states(agent_states)
+                raise
+            return results + self._send_agent_states(agent_states)
 
     def pending_message_count(self) -> int:
         """Return the number of messages pending in the buffer.

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

@@ -1,9 +1,11 @@
 """Tests for AgentCoreMemorySessionManager."""
 
 import asyncio
+import contextlib
 import inspect
 import json
 import logging
+import threading
 import time
 from datetime import datetime, timedelta, timezone
 from unittest.mock import Mock, patch
@@ -3941,6 +3943,148 @@ def test_agent_state_auto_flush_persists_buffered_messages_first(
         assert len(written) == 2
 
 
+class TestConcurrentFlushes:
+    """Flushes from the interval timer and the agent thread must not overtake each other."""
+
+    SESSION = "test-session-456"
+
+    @staticmethod
+    def _agent(v, agent_id="agent-1"):
+        return SessionAgent(agent_id=agent_id, state={"v": v}, conversation_manager_state={})
+
+    @staticmethod
+    def _message(text):
+        return SessionMessage(message={"role": "user", "content": [{"text": text}]}, message_id=0)
+
+    @staticmethod
+    def _agent_events(mock_memory_client):
+        return [
+            c.kwargs
+            for c in mock_memory_client.gmdp_client.create_event.call_args_list
+            if c.kwargs.get("metadata", {}).get("stateType") == {"stringValue": "AGENT"}
+        ]
+
+    @pytest.mark.parametrize("first_send_fails", [False, True])
+    def test_overlapping_flushes_restore_newest_state(
+        self, batching_session_manager, mock_memory_client, first_send_fails
+    ):
+        """A flush that starts while another is sending waits, so the newest state is written last."""
+        sm = batching_session_manager
+        first_flush_sending = threading.Event()
+        release_first_flush = threading.Event()
+        monotonic_timestamp = sm._get_monotonic_timestamp
+        sent = []
+
+        def first_flush():
+            with contextlib.suppress(SessionException):
+                sm._flush_messages()
+
+        first_flush_thread = threading.Thread(target=first_flush)
+
+        def slow_first_timestamp(*args, **kwargs):
+            if threading.current_thread() is first_flush_thread and not first_flush_sending.is_set():
+                first_flush_sending.set()
+                release_first_flush.wait(timeout=5)
+            return monotonic_timestamp(*args, **kwargs)
+
+        def create_event(**kwargs):
+            if threading.current_thread() is first_flush_thread and first_send_fails:
+                raise Exception("API Error")
+            sent.append(kwargs)
+            return {"event": {"eventId": "e"}}
+
+        mock_memory_client.gmdp_client.create_event.side_effect = create_event
+        sm.create_agent(self.SESSION, self._agent(1))
+        with patch.object(sm, "_get_monotonic_timestamp", side_effect=slow_first_timestamp):
+            first_flush_thread.start()
+            assert first_flush_sending.wait(timeout=5)
+            sm.create_agent(self.SESSION, self._agent(2))
+            second_flush_thread = threading.Thread(target=sm._flush_messages)
+            second_flush_thread.start()
+            time.sleep(0.2)
+            release_first_flush.set()
+            first_flush_thread.join(timeout=5)
+            second_flush_thread.join(timeout=5)
+        sm._flush_messages()
+
+        newest = max(sent, key=lambda kwargs: kwargs["eventTimestamp"])
+        assert json.loads(newest["payload"][-1]["blob"])["state"] == {"v": 2}
+        assert sm.pending_agent_state_count() == 0
+
+    def test_flush_keeps_states_buffered_after_their_messages_were_taken(
+        self, batching_session_manager, mock_memory_client
+    ):
+        """A state buffered while messages are sending stays buffered with its message."""
+        sm = batching_session_manager
+        sm.create_message(self.SESSION, "agent-1", self._message("m1"))
+        sm.create_agent(self.SESSION, self._agent(1))
+
+        def create_event(**kwargs):
+            if "metadata" not in kwargs:
+                sm.create_message(self.SESSION, "agent-1", self._message("m2"))
+                sm.create_agent(self.SESSION, self._agent(2))
+            return {"event": {"eventId": "e"}}
+
+        mock_memory_client.gmdp_client.create_event.side_effect = create_event
+        sm._flush_messages()
+
+        [agent_event] = self._agent_events(mock_memory_client)
+        assert [json.loads(p["blob"])["state"] for p in agent_event["payload"]] == [{"v": 1}]
+        assert sm.pending_message_count() == 1
+        assert sm.pending_agent_state_count() == 1
+
+    def test_failed_message_flush_keeps_agent_states_buffered(self, batching_session_manager, mock_memory_client):
+        """Agent states are not written when their messages failed to flush."""
+        sm = batching_session_manager
+        sm.create_message(self.SESSION, "agent-1", self._message("m1"))
+        sm.create_agent(self.SESSION, self._agent(1))
+        mock_memory_client.gmdp_client.create_event.side_effect = Exception("API Error")
+
+        with pytest.raises(SessionException):
+            sm._flush_messages()
+
+        assert mock_memory_client.gmdp_client.create_event.call_count == 1
+        assert sm.pending_message_count() == 1
+        assert sm.pending_agent_state_count() == 1
+
+    def test_failed_message_flush_keeps_newer_concurrent_message_last(
+        self, batching_session_manager, mock_memory_client
+    ):
+        """A message buffered during a failed flush stays after the messages that failed."""
+        sm = batching_session_manager
+        sm.create_message(self.SESSION, "agent-1", self._message("m1"))
+
+        def fail_after_concurrent_append(**kwargs):
+            sm.create_message(self.SESSION, "agent-1", self._message("m2"))
+            raise Exception("API Error")
+
+        mock_memory_client.gmdp_client.create_event.side_effect = fail_after_concurrent_append
+        with pytest.raises(SessionException):
+            sm._flush_messages_only()
+
+        texts = [json.loads(m.messages[0][0])["message"]["content"][0]["text"] for m in sm._message_buffer]
+        assert texts == ["m1", "m2"]
+
+    def test_partial_agent_state_failure_rebuffers_only_unsent_agents(
+        self, batching_session_manager, mock_memory_client
+    ):
+        """States already written for one agent are not written again when another agent fails."""
+        sm = batching_session_manager
+        sm.create_agent(self.SESSION, self._agent(1, agent_id="agent-a"))
+        sm.create_agent(self.SESSION, self._agent(1, agent_id="agent-b"))
+
+        def fail_for_agent_b(**kwargs):
+            if kwargs["metadata"]["agentId"] == {"stringValue": "agent-b"}:
+                raise Exception("API Error")
+            return {"event": {"eventId": "e"}}
+
+        mock_memory_client.gmdp_client.create_event.side_effect = fail_for_agent_b
+        with pytest.raises(SessionException):
+            sm._flush_agent_states_only()
+
+        assert [a.agent_id for _, a in sm._agent_state_buffer] == ["agent-b"]
+
+
 class TestMonotonicTimestamp:
     """Tests for monotonic event-timestamp ordering across processes/pods.