From 87bff31520dda2beffdb4a6aedccbbc704de6c95 Mon Sep 17 00:00:00 2001 From: Jorge Villacorta Date: Mon, 13 Apr 2026 21:46:37 -0400 Subject: [PATCH] Python: move _ensure_message_ids higher in the chain to avoid duplicate message ids --- .../core/agent_framework/_compaction.py | 3 ++- .../core/tests/core/test_compaction.py | 18 ++++++++++++++++++ 2 files changed, 20 insertions(+), 1 deletion(-) diff --git a/python/packages/core/agent_framework/_compaction.py b/python/packages/core/agent_framework/_compaction.py index dd76d1f0f4..6f5e725162 100644 --- a/python/packages/core/agent_framework/_compaction.py +++ b/python/packages/core/agent_framework/_compaction.py @@ -111,7 +111,6 @@ def group_messages(messages: list[Message]) -> list[dict[str, Any]]: Ordered list of lightweight span dicts with keys: ``group_id``, ``kind``, ``start_index``, ``end_index``, ``has_reasoning``. """ - _ensure_message_ids(messages) spans: list[dict[str, Any]] = [] i = 0 group_index = 0 @@ -416,6 +415,8 @@ def annotate_message_groups( if not messages: return [] + _ensure_message_ids(messages) + if force_reannotate: start_index = 0 elif from_index is not None: diff --git a/python/packages/core/tests/core/test_compaction.py b/python/packages/core/tests/core/test_compaction.py index 1db3260822..170147ddff 100644 --- a/python/packages/core/tests/core/test_compaction.py +++ b/python/packages/core/tests/core/test_compaction.py @@ -153,6 +153,24 @@ def test_group_annotations_handle_same_message_reasoning_and_function_calls() -> assert _group_has_reasoning(messages[1]) is True +def test_annotate_message_groups_assigns_unique_ids_after_incremental_calls() -> None: + """Regression: ``_ensure_message_ids`` previously ran inside + ``group_messages`` against a slice and assigned ``msg_{slice_index}``, + which collided with ids already present in the preserved prefix. + """ + messages: list[Message] = [] + for turn in range(4): + messages.append(Message(role="user", contents=[f"u{turn}"])) + messages.append(Message(role="assistant", contents=[f"a{turn}"])) + annotate_message_groups(messages) + + ids = [m.message_id for m in messages] + assert None not in ids + assert len(set(ids)) == len(ids), ( + f"expected unique message_ids across incremental annotation, got {ids}" + ) + + def test_annotate_message_groups_with_tokenizer_adds_token_counts() -> None: messages = [ Message(role="user", contents=["hello"]),