diff --git a/src/bedrock_agentcore/memory/integrations/strands/session_manager.py b/src/bedrock_agentcore/memory/integrations/strands/session_manager.py index c7ea08cc..fe79c9e5 100644 --- a/src/bedrock_agentcore/memory/integrations/strands/session_manager.py +++ b/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. diff --git a/tests/bedrock_agentcore/memory/integrations/strands/test_agentcore_memory_session_manager.py b/tests/bedrock_agentcore/memory/integrations/strands/test_agentcore_memory_session_manager.py index 3ac7fa90..50309980 100644 --- a/tests/bedrock_agentcore/memory/integrations/strands/test_agentcore_memory_session_manager.py +++ b/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.