From ed96ba2e462f3dd774ab57755911edca40fbf080 Mon Sep 17 00:00:00 2001 From: Valter Silva Date: Sat, 27 Jun 2026 16:05:30 +0800 Subject: [PATCH] fix(memory): order get_last_k_turns messages chronologically MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The ListEvents API returns events newest-first, but get_last_k_turns grouped them in arrival order. Because a new turn is started at each USER message, processing newest-first misaligns turns — an ASSISTANT reply is orphaned into a leading turn and every subsequent turn pairs a USER message with the previous turn's ASSISTANT response. The effect is most visible when tool use splits a turn across single-message events. Sort each fetched batch by eventTimestamp (oldest-first) before grouping, in both MemoryClient.get_last_k_turns and MemorySessionManager.get_last_k_turns. The sort is skipped when any event in the batch lacks a timestamp, preserving existing behavior for callers that don't supply one. Public signatures and pagination are unchanged. Add regression tests in both modules that feed events newest-first (as the real API does) and assert a single, correctly ordered USER->ASSISTANT turn. Fixes #253 --- src/bedrock_agentcore/memory/client.py | 5 +++ src/bedrock_agentcore/memory/session.py | 5 +++ tests/bedrock_agentcore/memory/test_client.py | 36 +++++++++++++++++++ .../bedrock_agentcore/memory/test_session.py | 35 ++++++++++++++++++ 4 files changed, 81 insertions(+) diff --git a/src/bedrock_agentcore/memory/client.py b/src/bedrock_agentcore/memory/client.py index 421dfcff..c9712b17 100644 --- a/src/bedrock_agentcore/memory/client.py +++ b/src/bedrock_agentcore/memory/client.py @@ -1208,6 +1208,11 @@ def get_last_k_turns( total_fetched += len(events) + # ListEvents returns events newest-first; group them oldest-first so a + # turn's USER message precedes its ASSISTANT response(s). See issue #320. + if all(e.get("eventTimestamp") is not None for e in events): + events = sorted(events, key=lambda e: e["eventTimestamp"]) + for event in events: if len(turns) >= k: break diff --git a/src/bedrock_agentcore/memory/session.py b/src/bedrock_agentcore/memory/session.py index 5f65198a..6d53e2de 100644 --- a/src/bedrock_agentcore/memory/session.py +++ b/src/bedrock_agentcore/memory/session.py @@ -841,6 +841,11 @@ def get_last_k_turns( total_fetched += len(events) + # ListEvents returns events newest-first; group them oldest-first so a + # turn's USER message precedes its ASSISTANT response(s). See issue #320. + if all(e.get("eventTimestamp") is not None for e in events): + events = sorted(events, key=lambda e: e["eventTimestamp"]) + for event in events: if len(turns) >= k: break diff --git a/tests/bedrock_agentcore/memory/test_client.py b/tests/bedrock_agentcore/memory/test_client.py index f7df4c4a..535484e7 100644 --- a/tests/bedrock_agentcore/memory/test_client.py +++ b/tests/bedrock_agentcore/memory/test_client.py @@ -3506,6 +3506,42 @@ def test_get_last_k_turns_explicit_max_results(): assert total_fetched <= 50 # Should stop after fetching 50 events worth of calls +def test_get_last_k_turns_orders_messages_chronologically(): + """Turns are grouped chronologically even when list_events returns newest-first. + + Regression for #320/#253: the ListEvents API returns events newest-first, but + get_last_k_turns grouped them in arrival order, yielding turns whose messages + were misaligned (ASSISTANT before its USER, orphaned leading turn). Single-message + events (e.g. a toolUse split across events) make the misalignment visible. + """ + with patch("boto3.Session"): + client = MemoryClient() + + mock_gmdp = MagicMock() + client.gmdp_client = mock_gmdp + + # As returned by the real API: newest event first. + mock_events = [ + { + "eventId": "event-2", + "eventTimestamp": datetime(2023, 1, 1, 10, 5, 0), + "payload": [{"conversational": {"role": "ASSISTANT", "content": {"text": "It's sunny"}}}], + }, + { + "eventId": "event-1", + "eventTimestamp": datetime(2023, 1, 1, 10, 0, 0), + "payload": [{"conversational": {"role": "USER", "content": {"text": "Weather?"}}}], + }, + ] + mock_gmdp.list_events.return_value = {"events": mock_events, "nextToken": None} + + turns = client.get_last_k_turns(memory_id="mem-123", actor_id="user-123", session_id="session-456", k=5) + + # One turn: USER then ASSISTANT, in chronological order. + assert len(turns) == 1 + assert [m["role"] for m in turns[0]] == ["USER", "ASSISTANT"] + + # ============================================================================ # LTM Metadata: indexed_keys and metadata_filters tests # ============================================================================ diff --git a/tests/bedrock_agentcore/memory/test_session.py b/tests/bedrock_agentcore/memory/test_session.py index 3108587a..f4bcc643 100644 --- a/tests/bedrock_agentcore/memory/test_session.py +++ b/tests/bedrock_agentcore/memory/test_session.py @@ -2174,6 +2174,41 @@ def test_get_last_k_turns_turn_grouping(self): assert len(result[0]) == 2 # First turn: USER + ASSISTANT assert len(result[1]) == 2 # Second turn: USER + ASSISTANT + def test_get_last_k_turns_orders_messages_chronologically(self): + """Turns are grouped chronologically even when list_events returns newest-first. + + Regression for #320/#253: the ListEvents API returns events newest-first, but + get_last_k_turns grouped them in arrival order, yielding turns whose messages + were misaligned (ASSISTANT before its USER, orphaned leading turn). + """ + with patch("boto3.Session") as mock_boto_client: + mock_client_instance = MagicMock() + mock_boto_client.return_value = mock_client_instance + + manager = MemorySessionManager(memory_id="testMemory-1234567890", region_name="us-west-2") + + # As returned by the real API: newest event first. + manager._data_plane_client.list_events.return_value = { + "events": [ + { + "eventId": "event-2", + "eventTimestamp": datetime(2023, 1, 1, 10, 5, 0), + "payload": [{"conversational": {"role": "ASSISTANT", "content": {"text": "It's sunny"}}}], + }, + { + "eventId": "event-1", + "eventTimestamp": datetime(2023, 1, 1, 10, 0, 0), + "payload": [{"conversational": {"role": "USER", "content": {"text": "Weather?"}}}], + }, + ] + } + + result = manager.get_last_k_turns(actor_id="user-123", session_id="session-456", k=5) + + # One turn: USER then ASSISTANT, in chronological order. + assert len(result) == 1 + assert [m["role"] for m in result[0]] == ["USER", "ASSISTANT"] + def test_session_delegation_with_optional_parameters(self): """Test MemorySession methods properly pass optional parameters.""" with patch("boto3.Session"):