diff --git a/src/bedrock_agentcore/memory/integrations/strands/session_manager.py b/src/bedrock_agentcore/memory/integrations/strands/session_manager.py index 5e6ac7bc..c7ea08cc 100644 --- a/src/bedrock_agentcore/memory/integrations/strands/session_manager.py +++ b/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 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 9a68530e..3ac7fa90 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 @@ -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.