From 7115fcd9de02e062d472c8e6040fa669b9f24f5f Mon Sep 17 00:00:00 2001 From: jariy17 Date: Mon, 12 Oct 2026 21:44:10 +0000 Subject: [PATCH 1/3] fix(memory): restore latest agent snapshot from batched events With batch_size > 1, one AGENT event holds several SessionAgent snapshots, oldest first. read_agent read payload[0], so restored agents lost the state saved during their previous invocation. Fixes #698 --- .../memory/integrations/strands/session_manager.py | 3 ++- .../test_agentcore_memory_session_manager.py | 14 ++++++++++++++ 2 files changed, 16 insertions(+), 1 deletion(-) diff --git a/src/bedrock_agentcore/memory/integrations/strands/session_manager.py b/src/bedrock_agentcore/memory/integrations/strands/session_manager.py index 5e6ac7bc..267267ad 100644 --- a/src/bedrock_agentcore/memory/integrations/strands/session_manager.py +++ b/src/bedrock_agentcore/memory/integrations/strands/session_manager.py @@ -471,7 +471,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: 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..a78ad28f 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 = [] From 578b8d1c0bcc276afca142be65ae3a14be3009dc Mon Sep 17 00:00:00 2001 From: jariy17 Date: Mon, 12 Oct 2026 21:53:27 +0000 Subject: [PATCH 2/3] fix(memory): keep batched agent snapshots ordered and behind their messages Restoring payload[-1] relies on the newest snapshot being last and on messages being persisted no later than the snapshot that indexes them. - Prepend agent states re-buffered after a failed flush, so states buffered during the failed send stay newest. - Flush buffered messages before agent states when the agent state buffer fills, so a crash can't leave a snapshot whose removed_message_count points past the saved history. --- .../integrations/strands/session_manager.py | 8 ++-- .../test_agentcore_memory_session_manager.py | 44 +++++++++++++++++++ 2 files changed, 48 insertions(+), 4 deletions(-) diff --git a/src/bedrock_agentcore/memory/integrations/strands/session_manager.py b/src/bedrock_agentcore/memory/integrations/strands/session_manager.py index 267267ad..0f951b31 100644 --- a/src/bedrock_agentcore/memory/integrations/strands/session_manager.py +++ b/src/bedrock_agentcore/memory/integrations/strands/session_manager.py @@ -400,9 +400,9 @@ 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, messages first so a snapshot never lands ahead of its history if should_flush: - self._flush_agent_states_only() + self._flush_messages() logger.info( "Buffered agent creation: %s in session: %s", @@ -1185,9 +1185,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 a78ad28f..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 @@ -3896,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. From 50aab6b78a56b6231bfd25c5c797326083d346ed Mon Sep 17 00:00:00 2001 From: jariy17 Date: Tue, 13 Oct 2026 19:47:05 +0000 Subject: [PATCH 3/3] docs(memory): explain why the agent-state auto-flush also flushes messages --- .../memory/integrations/strands/session_manager.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/bedrock_agentcore/memory/integrations/strands/session_manager.py b/src/bedrock_agentcore/memory/integrations/strands/session_manager.py index 0f951b31..c7ea08cc 100644 --- a/src/bedrock_agentcore/memory/integrations/strands/session_manager.py +++ b/src/bedrock_agentcore/memory/integrations/strands/session_manager.py @@ -400,7 +400,8 @@ 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 outside the lock to prevent deadlock, messages first so a snapshot never lands ahead of its history + # 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_messages()