Python: Fix compaction message-id collisions and tool-loop summary persistence (#6299)

* Fix compaction message-id collisions and tool-loop summary persistence Fixes two bugs in the compaction strategies: - #5237: incremental group annotation assigned message ids by position within the re-annotated slice, so moving the re-annotation start back to a previous group start restarted ids at 0 and produced collisions (e.g. a user message reusing an assistant message's id), merging groups and causing tool-result compaction to wrongly exclude messages. group_messages/_ensure_message_ids now take an id_offset and guard against existing-id collisions; annotate_message_groups threads the slice start index through as the offset. - #4991: the function-invocation loop copied the message list each iteration, so summaries inserted by compaction landed in a throwaway copy and were lost across tool-loop iterations (only the persistent excluded flags survived). _prepare_messages_for_model_call now compacts the list in place when messages is a list, so inserted summaries persist. Adds regression tests (incremental id uniqueness, existing-id collision avoidance, idempotency, and tool-loop summary persistence including streaming and conversation-id modes). Also adds a summarization.py sample demonstrating SummarizationStrategy directly with a real client, and reworks advanced.py with tool-call groups and a real summarizer. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> * Guard incremental message-id assignment against prefix-id collisions Addresses PR review on #5237: _ensure_message_ids only guarded against collisions within the re-annotated slice. A preexisting (e.g. user-supplied) id in the preserved prefix could still be reassigned in the suffix when the id was numerically out of position, merging groups across the re-annotation boundary again. group_messages/_ensure_message_ids now accept reserved_ids, and annotate_message_groups passes the preserved prefix's ids so auto-assigned suffix ids never collide across the full list. Adds a regression test reproducing the out-of-position prefix-id collision. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --------- Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
2026-06-16 21:04:09 +08:00 · 2026-06-04 10:37:59 +02:00
parent f29bae8fbc
commit f970a699d8
7 changed files with 633 additions and 54 deletions
@@ -380,8 +380,15 @@ class BaseChatClient(SerializationMixin, ABC, Generic[OptionsCoT]):
            return prepared_messages
        from ._compaction import apply_compaction

+        # Compact the caller's list in place when possible. A compaction operation has
+        # two halves: exclusion flags (mutated on shared Message objects) and inserted
+        # summary messages. Operating on the original list keeps both halves on the list
+        # the function-invocation tool loop reuses across iterations; otherwise inserted
+        # summaries would be lost on a throwaway copy while exclusions persisted, silently
+        # dropping older groups (issue #4991).
+        working_messages = messages if isinstance(messages, list) else prepared_messages
        return await apply_compaction(
-            prepared_messages,
+            working_messages,
            strategy=compaction_strategy,
            tokenizer=tokenizer,
        )
@@ -4,7 +4,7 @@ from __future__ import annotations

 import json
 import logging
-from collections.abc import Mapping, Sequence
+from collections.abc import Iterable, Mapping, Sequence
 from typing import (
    TYPE_CHECKING,
    Any,
@@ -92,10 +92,23 @@ def _is_reasoning_only_assistant(message: Message) -> bool:
    return all(content.type == "text_reasoning" for content in message.contents)


-def _ensure_message_ids(messages: list[Message]) -> None:
+def _ensure_message_ids(
+    messages: list[Message], *, id_offset: int = 0, reserved_ids: Iterable[str] | None = None
+) -> None:
+    existing_ids: set[str] = set(reserved_ids) if reserved_ids is not None else set()
+    existing_ids.update(message.message_id for message in messages if message.message_id)
    for index, message in enumerate(messages):
-        if not message.message_id:
-            message.message_id = f"msg_{index}"
+        if message.message_id:
+            continue
+        candidate = f"msg_{id_offset + index}"
+        if candidate in existing_ids:
+            counter = id_offset + len(messages)
+            candidate = f"msg_{counter}"
+            while candidate in existing_ids:
+                counter += 1
+                candidate = f"msg_{counter}"
+        message.message_id = candidate
+        existing_ids.add(candidate)


 def _group_id_for(message: Message, group_index: int) -> str:
@@ -104,14 +117,27 @@ def _group_id_for(message: Message, group_index: int) -> str:
    return f"group_index_{group_index}"


-def group_messages(messages: list[Message]) -> list[dict[str, Any]]:
+def group_messages(
+    messages: list[Message], *, id_offset: int = 0, reserved_ids: Iterable[str] | None = None
+) -> list[dict[str, Any]]:
    """Compute group spans and metadata for annotation.

+    Args:
+        messages: The messages (or a slice of them) to group.
+
+    Keyword Args:
+        id_offset: Absolute starting index used when auto-assigning ``message_id``
+            values, so incremental annotation of a list slice produces ids that
+            stay unique across the full list.
+        reserved_ids: Message ids that already exist outside ``messages`` (for
+            example in a preserved prefix). Auto-assigned ids are guaranteed not
+            to collide with these, preventing duplicate ids across the full list.
+
    Returns:
        Ordered list of lightweight span dicts with keys:
        ``group_id``, ``kind``, ``start_index``, ``end_index``, ``has_reasoning``.
    """
-    _ensure_message_ids(messages)
+    _ensure_message_ids(messages, id_offset=id_offset, reserved_ids=reserved_ids)
    spans: list[dict[str, Any]] = []
    i = 0
    group_index = 0
@@ -439,7 +465,8 @@ def annotate_message_groups(
        if previous_group_index is not None:
            group_index_offset = previous_group_index + 1

-    spans = group_messages(messages[start_index:])
+    reserved_ids = {message.message_id for message in messages[:start_index] if message.message_id}
+    spans = group_messages(messages[start_index:], id_offset=start_index, reserved_ids=reserved_ids)
    for span_index, span in enumerate(spans):
        group_id = str(span["group_id"])
        kind = _coerce_group_kind(span["kind"])
@@ -11,10 +11,14 @@ from agent_framework import (
    GROUP_TOKEN_COUNT_KEY,
    BaseChatClient,
    ChatResponse,
+    ChatResponseUpdate,
+    Content,
    Message,
    SlidingWindowStrategy,
    SupportsChatGetResponse,
+    ToolResultCompactionStrategy,
    TruncationStrategy,
+    tool,
 )


@@ -258,6 +262,196 @@ async def test_base_client_default_tokenizer_without_strategy_annotates_messages
    assert captured_token_counts == [[19, 19]]


+def _tool_call_response(call_id: str, location: str) -> ChatResponse:
+    return ChatResponse(
+        messages=Message(
+            role="assistant",
+            contents=[
+                Content.from_function_call(
+                    call_id=call_id,
+                    name="lookup_weather",
+                    arguments=f'{{"location": "{location}"}}',
+                )
+            ],
+        ),
+        response_id=f"resp_{call_id}",
+    )
+
+
+def _is_tool_result_summary(message: Message) -> bool:
+    text = message.text or ""
+    return message.role == "assistant" and text.startswith("[Tool results:")
+
+
+async def test_function_loop_persists_inserted_summaries_across_iterations(
+    chat_client_base: SupportsChatGetResponse,
+) -> None:
+    # Regression test for #4991: compaction inserts summary messages and excludes the
+    # originals. Across tool-loop iterations the exclusion flags persisted (shared Message
+    # objects) but the inserted summaries were dropped (they only lived on a throwaway copy),
+    # so older tool groups were silently lost with no summary representing them.
+    chat_client_base.function_invocation_configuration["enabled"] = True  # type: ignore[attr-defined]
+    chat_client_base.function_invocation_configuration["max_iterations"] = 3  # type: ignore[attr-defined]
+    chat_client_base.compaction_strategy = ToolResultCompactionStrategy(keep_last_tool_call_groups=1)  # type: ignore[attr-defined]
+
+    @tool(name="lookup_weather", approval_mode="never_require")
+    def lookup_weather(location: str) -> str:
+        return f"Weather in {location}: sunny"
+
+    chat_client_base.run_responses = [  # type: ignore[attr-defined]
+        _tool_call_response("call_1", "London"),
+        _tool_call_response("call_2", "Paris"),
+        _tool_call_response("call_3", "Tokyo"),
+    ]
+
+    captured_inputs: list[list[Message]] = []
+    original = chat_client_base._get_non_streaming_response  # type: ignore[attr-defined]
+
+    async def _capture(
+        *,
+        messages: list[Message],
+        options: dict[str, Any],
+        **kwargs: Any,
+    ) -> ChatResponse:
+        captured_inputs.append(list(messages))
+        return await original(messages=messages, options=options, **kwargs)
+
+    chat_client_base._get_non_streaming_response = _capture  # type: ignore[attr-defined,method-assign]
+
+    await chat_client_base.get_response(
+        [Message(role="user", contents=["What is the weather in London?"])],
+        options={"tools": [lookup_weather]},  # type: ignore[typeddict-unknown-key]
+    )
+
+    # The final model call should represent every compacted tool group with a summary.
+    # Two older tool groups get collapsed (London, Paris) while the last (Tokyo) is kept.
+    final_input = captured_inputs[-1]
+    summaries = [message for message in final_input if _is_tool_result_summary(message)]
+    summary_text = " ".join(message.text or "" for message in summaries)
+
+    assert len(summaries) == 2, [message.text for message in final_input]
+    assert "London" in summary_text
+    assert "Paris" in summary_text
+
+
+def _tool_call_update(call_id: str, location: str) -> list[ChatResponseUpdate]:
+    return [
+        ChatResponseUpdate(
+            contents=[
+                Content.from_function_call(
+                    call_id=call_id,
+                    name="lookup_weather",
+                    arguments=f'{{"location": "{location}"}}',
+                )
+            ],
+            role="assistant",
+            finish_reason="stop",
+            response_id=f"resp_{call_id}",
+        )
+    ]
+
+
+async def test_function_loop_persists_inserted_summaries_across_iterations_streaming(
+    chat_client_base: SupportsChatGetResponse,
+) -> None:
+    # Streaming counterpart of the #4991 regression test: the summary persistence fix in
+    # ``_prepare_messages_for_model_call`` must cover the streaming tool loop too.
+    chat_client_base.function_invocation_configuration["enabled"] = True  # type: ignore[attr-defined]
+    chat_client_base.function_invocation_configuration["max_iterations"] = 3  # type: ignore[attr-defined]
+    chat_client_base.compaction_strategy = ToolResultCompactionStrategy(keep_last_tool_call_groups=1)  # type: ignore[attr-defined]
+
+    @tool(name="lookup_weather", approval_mode="never_require")
+    def lookup_weather(location: str) -> str:
+        return f"Weather in {location}: sunny"
+
+    chat_client_base.streaming_responses = [  # type: ignore[attr-defined]
+        _tool_call_update("call_1", "London"),
+        _tool_call_update("call_2", "Paris"),
+        _tool_call_update("call_3", "Tokyo"),
+    ]
+
+    captured_inputs: list[list[Message]] = []
+    original = chat_client_base._get_streaming_response  # type: ignore[attr-defined]
+
+    def _capture(
+        *,
+        messages: list[Message],
+        options: dict[str, Any],
+        **kwargs: Any,
+    ):
+        captured_inputs.append(list(messages))
+        return original(messages=messages, options=options, **kwargs)
+
+    chat_client_base._get_streaming_response = _capture  # type: ignore[attr-defined,method-assign]
+
+    stream = chat_client_base.get_response(
+        [Message(role="user", contents=["What is the weather in London?"])],
+        stream=True,
+        options={"tools": [lookup_weather]},  # type: ignore[typeddict-unknown-key]
+    )
+    async for _ in stream:
+        pass
+
+    final_input = captured_inputs[-1]
+    summaries = [message for message in final_input if _is_tool_result_summary(message)]
+    summary_text = " ".join(message.text or "" for message in summaries)
+
+    assert len(summaries) == 2, [message.text for message in final_input]
+    assert "London" in summary_text
+    assert "Paris" in summary_text
+
+
+async def test_function_loop_compaction_conversation_id_mode_does_not_resend_history(
+    chat_client_base: SupportsChatGetResponse,
+) -> None:
+    # In conversation-id mode the server owns prior context, so the tool loop clears
+    # ``prepped_messages`` and only sends the latest message. Compaction must not fight that
+    # by re-inserting summaries or re-sending earlier turns.
+    chat_client_base.function_invocation_configuration["enabled"] = True  # type: ignore[attr-defined]
+    chat_client_base.function_invocation_configuration["max_iterations"] = 3  # type: ignore[attr-defined]
+    chat_client_base.compaction_strategy = ToolResultCompactionStrategy(keep_last_tool_call_groups=1)  # type: ignore[attr-defined]
+
+    @tool(name="lookup_weather", approval_mode="never_require")
+    def lookup_weather(location: str) -> str:
+        return f"Weather in {location}: sunny"
+
+    def _conversation_tool_call(call_id: str, location: str) -> ChatResponse:
+        response = _tool_call_response(call_id, location)
+        response.conversation_id = "conv_1"
+        return response
+
+    chat_client_base.run_responses = [  # type: ignore[attr-defined]
+        _conversation_tool_call("call_1", "London"),
+        _conversation_tool_call("call_2", "Paris"),
+        _conversation_tool_call("call_3", "Tokyo"),
+    ]
+
+    captured_inputs: list[list[Message]] = []
+    original = chat_client_base._get_non_streaming_response  # type: ignore[attr-defined]
+
+    async def _capture(
+        *,
+        messages: list[Message],
+        options: dict[str, Any],
+        **kwargs: Any,
+    ) -> ChatResponse:
+        captured_inputs.append(list(messages))
+        return await original(messages=messages, options=options, **kwargs)
+
+    chat_client_base._get_non_streaming_response = _capture  # type: ignore[attr-defined,method-assign]
+
+    await chat_client_base.get_response(
+        [Message(role="user", contents=["What is the weather in London?"])],
+        options={"tools": [lookup_weather]},  # type: ignore[typeddict-unknown-key]
+    )
+
+    # After the conversation id is established the loop only forwards the latest message,
+    # so subsequent model calls never receive the full history or summary messages.
+    for sent in captured_inputs[1:]:
+        assert len(sent) <= 1, [message.text for message in sent]
+        assert not any(_is_tool_result_summary(message) for message in sent)
+
+
 def test_base_client_as_agent_does_not_copy_client_compaction_defaults(
    chat_client_base: SupportsChatGetResponse,
 ) -> None:
@@ -196,6 +196,64 @@ def test_append_compaction_message_annotates_new_message() -> None:
    assert isinstance(_group_id(messages[1]), str)


+def test_incremental_annotation_assigns_unique_message_ids() -> None:
+    # Regression test for #5237: ``_ensure_message_ids`` assigned ``msg_{index}``
+    # using the position within the slice handed to ``group_messages``. Successive
+    # incremental annotations restart the index at 0, so distinct messages collided
+    # on the same ``message_id``.
+    messages: list[Message] = []
+    for turn in range(4):
+        messages.append(Message(role="user", contents=[f"user {turn}"]))
+        annotate_message_groups(messages)
+        messages.append(Message(role="assistant", contents=[f"assistant {turn}"]))
+        annotate_message_groups(messages)
+
+    message_ids = [message.message_id for message in messages]
+    assert all(message_ids), "every message should receive an id"
+    assert len(set(message_ids)) == len(message_ids), f"duplicate message ids: {message_ids}"
+
+
+def test_ensure_message_ids_avoids_existing_id_collisions() -> None:
+    # An auto-generated ``msg_{index}`` must not collide with an id already present
+    # on another message (user-supplied or assigned by an earlier annotation pass).
+    messages = [
+        Message(role="user", contents=["zero"]),
+        Message(role="assistant", contents=["one"], message_id="msg_2"),
+        Message(role="user", contents=["two"]),
+    ]
+    annotate_message_groups(messages)
+
+    message_ids = [message.message_id for message in messages]
+    assert message_ids[1] == "msg_2"
+    assert len(set(message_ids)) == len(message_ids), f"duplicate message ids: {message_ids}"
+
+
+def test_incremental_annotation_avoids_prefix_id_collision() -> None:
+    # Regression for the PR review on #5237: when only a suffix is re-annotated,
+    # an auto-assigned ``msg_{index}`` in the suffix must not collide with a
+    # preexisting id carried by a message in the *preserved prefix* (a group
+    # before the one re-annotation pulls back to). Otherwise ``_group_id_for``
+    # derives the same group id and merges groups across the boundary.
+    messages = [
+        # Out-of-position, user-supplied id that matches the ``msg_{index}`` the
+        # suffix pass would assign to the appended message below. This message is
+        # two groups back, so it stays outside the re-annotated slice.
+        Message(role="user", contents=["zero"], message_id="msg_2"),
+        Message(role="user", contents=["one"]),
+    ]
+    annotate_message_groups(messages)
+    assert messages[0].message_id == "msg_2"
+    assert messages[1].message_id == "msg_1"
+
+    messages.append(Message(role="user", contents=["two"]))
+    annotate_message_groups(messages, from_index=2)
+
+    message_ids = [message.message_id for message in messages]
+    assert all(message_ids), "every message should receive an id"
+    assert len(set(message_ids)) == len(message_ids), f"duplicate message ids: {message_ids}"
+    assert messages[0].message_id == "msg_2"
+
+
 async def test_truncation_strategy_keeps_system_anchor() -> None:
    messages = [
        Message(role="system", contents=["you are helpful"]),
@@ -484,6 +542,44 @@ async def test_tool_result_compaction_collapses_old_groups_into_summary() -> Non
    assert any(m.role == "tool" for m in projected)


+async def test_tool_result_compaction_is_idempotent_after_summary_insertion() -> None:
+    """Re-running compaction after a mid-list summary insertion must not duplicate it.
+
+    Mirrors a subsequent tool-loop iteration (issue #4991): the inserted summary and the
+    excluded originals now persist on the same list, so a second annotate + compaction pass
+    over the same groups should be a no-op rather than collapsing the group again.
+    """
+    messages = [
+        Message(role="user", contents=["u"]),
+        _assistant_function_call("call-1"),
+        _tool_result("call-1", "r1"),
+        _assistant_function_call("call-2"),
+        _tool_result("call-2", "r2"),
+        Message(role="assistant", contents=["done"]),
+    ]
+    strategy = ToolResultCompactionStrategy(keep_last_tool_call_groups=1)
+    annotate_message_groups(messages)
+    assert await strategy(messages) is True
+
+    summaries_after_first = [m for m in messages if (m.text or "").startswith("[Tool results:")]
+    assert len(summaries_after_first) == 1
+    summary = summaries_after_first[0]
+    summary_group_ids = _group_unknown_value(summary, SUMMARY_OF_GROUP_IDS_KEY)
+
+    # Second pass over the same (now partially compacted) list.
+    annotate_message_groups(messages)
+    changed = await strategy(messages)
+
+    assert changed is False
+    summaries_after_second = [m for m in messages if (m.text or "").startswith("[Tool results:")]
+    assert len(summaries_after_second) == 1
+    assert _group_unknown_value(summaries_after_second[0], SUMMARY_OF_GROUP_IDS_KEY) == summary_group_ids
+
+    # The kept tool-call group stays atomic and included.
+    projected = included_messages(messages)
+    assert any(m.role == "tool" for m in projected)
+
+
 async def test_tool_result_compaction_zero_collapses_all() -> None:
    """With keep=0, all tool-call groups are collapsed into summaries."""
    messages = [