Skip to content

Commit 5b8fae0

Browse files
fix(memory): order get_last_k_turns messages chronologically
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
1 parent 0a8a486 commit 5b8fae0

4 files changed

Lines changed: 81 additions & 0 deletions

File tree

src/bedrock_agentcore/memory/client.py

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1214,6 +1214,11 @@ def get_last_k_turns(
12141214

12151215
total_fetched += len(events)
12161216

1217+
# ListEvents returns events newest-first; group them oldest-first so a
1218+
# turn's USER message precedes its ASSISTANT response(s). See issue #320.
1219+
if all(e.get("eventTimestamp") is not None for e in events):
1220+
events = sorted(events, key=lambda e: e["eventTimestamp"])
1221+
12171222
for event in events:
12181223
if len(turns) >= k:
12191224
break

src/bedrock_agentcore/memory/session.py

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -841,6 +841,11 @@ def get_last_k_turns(
841841

842842
total_fetched += len(events)
843843

844+
# ListEvents returns events newest-first; group them oldest-first so a
845+
# turn's USER message precedes its ASSISTANT response(s). See issue #320.
846+
if all(e.get("eventTimestamp") is not None for e in events):
847+
events = sorted(events, key=lambda e: e["eventTimestamp"])
848+
844849
for event in events:
845850
if len(turns) >= k:
846851
break

tests/bedrock_agentcore/memory/test_client.py

Lines changed: 36 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3569,6 +3569,42 @@ def test_get_last_k_turns_explicit_max_results():
35693569
assert total_fetched <= 50 # Should stop after fetching 50 events worth of calls
35703570

35713571

3572+
def test_get_last_k_turns_orders_messages_chronologically():
3573+
"""Turns are grouped chronologically even when list_events returns newest-first.
3574+
3575+
Regression for #320/#253: the ListEvents API returns events newest-first, but
3576+
get_last_k_turns grouped them in arrival order, yielding turns whose messages
3577+
were misaligned (ASSISTANT before its USER, orphaned leading turn). Single-message
3578+
events (e.g. a toolUse split across events) make the misalignment visible.
3579+
"""
3580+
with patch("boto3.Session"):
3581+
client = MemoryClient()
3582+
3583+
mock_gmdp = MagicMock()
3584+
client.gmdp_client = mock_gmdp
3585+
3586+
# As returned by the real API: newest event first.
3587+
mock_events = [
3588+
{
3589+
"eventId": "event-2",
3590+
"eventTimestamp": datetime(2023, 1, 1, 10, 5, 0),
3591+
"payload": [{"conversational": {"role": "ASSISTANT", "content": {"text": "It's sunny"}}}],
3592+
},
3593+
{
3594+
"eventId": "event-1",
3595+
"eventTimestamp": datetime(2023, 1, 1, 10, 0, 0),
3596+
"payload": [{"conversational": {"role": "USER", "content": {"text": "Weather?"}}}],
3597+
},
3598+
]
3599+
mock_gmdp.list_events.return_value = {"events": mock_events, "nextToken": None}
3600+
3601+
turns = client.get_last_k_turns(memory_id="mem-123", actor_id="user-123", session_id="session-456", k=5)
3602+
3603+
# One turn: USER then ASSISTANT, in chronological order.
3604+
assert len(turns) == 1
3605+
assert [m["role"] for m in turns[0]] == ["USER", "ASSISTANT"]
3606+
3607+
35723608
# ============================================================================
35733609
# LTM Metadata: indexed_keys and metadata_filters tests
35743610
# ============================================================================

tests/bedrock_agentcore/memory/test_session.py

Lines changed: 35 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2174,6 +2174,41 @@ def test_get_last_k_turns_turn_grouping(self):
21742174
assert len(result[0]) == 2 # First turn: USER + ASSISTANT
21752175
assert len(result[1]) == 2 # Second turn: USER + ASSISTANT
21762176

2177+
def test_get_last_k_turns_orders_messages_chronologically(self):
2178+
"""Turns are grouped chronologically even when list_events returns newest-first.
2179+
2180+
Regression for #320/#253: the ListEvents API returns events newest-first, but
2181+
get_last_k_turns grouped them in arrival order, yielding turns whose messages
2182+
were misaligned (ASSISTANT before its USER, orphaned leading turn).
2183+
"""
2184+
with patch("boto3.Session") as mock_boto_client:
2185+
mock_client_instance = MagicMock()
2186+
mock_boto_client.return_value = mock_client_instance
2187+
2188+
manager = MemorySessionManager(memory_id="testMemory-1234567890", region_name="us-west-2")
2189+
2190+
# As returned by the real API: newest event first.
2191+
manager._data_plane_client.list_events.return_value = {
2192+
"events": [
2193+
{
2194+
"eventId": "event-2",
2195+
"eventTimestamp": datetime(2023, 1, 1, 10, 5, 0),
2196+
"payload": [{"conversational": {"role": "ASSISTANT", "content": {"text": "It's sunny"}}}],
2197+
},
2198+
{
2199+
"eventId": "event-1",
2200+
"eventTimestamp": datetime(2023, 1, 1, 10, 0, 0),
2201+
"payload": [{"conversational": {"role": "USER", "content": {"text": "Weather?"}}}],
2202+
},
2203+
]
2204+
}
2205+
2206+
result = manager.get_last_k_turns(actor_id="user-123", session_id="session-456", k=5)
2207+
2208+
# One turn: USER then ASSISTANT, in chronological order.
2209+
assert len(result) == 1
2210+
assert [m["role"] for m in result[0]] == ["USER", "ASSISTANT"]
2211+
21772212
def test_session_delegation_with_optional_parameters(self):
21782213
"""Test MemorySession methods properly pass optional parameters."""
21792214
with patch("boto3.Session"):

0 commit comments

Comments
 (0)