From 241ed4acac9dfbc6ea8d87b7e996510049be4f9f Mon Sep 17 00:00:00 2001 From: eavanvalkenburg Date: Thu, 27 Aug 2026 14:01:23 +0200 Subject: [PATCH 1/2] Python: fix compaction behavior and observability Preserve compaction summaries across chat middleware boundaries, keep destructive truncation behind its documented threshold, add structured INFO logs, and support preserving the first user group. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: 726d66ea-6cff-4cb8-9090-8d919392271d --- .../core/agent_framework/_compaction.py | 126 ++++++++++++++---- .../core/agent_framework/_middleware.py | 116 ++++++++++++++-- .../packages/core/tests/core/test_clients.py | 77 +++++++++++ .../core/tests/core/test_compaction.py | 73 +++++++++- 4 files changed, 350 insertions(+), 42 deletions(-) diff --git a/python/packages/core/agent_framework/_compaction.py b/python/packages/core/agent_framework/_compaction.py index e8b46a008a3..f382ed2ce64 100644 --- a/python/packages/core/agent_framework/_compaction.py +++ b/python/packages/core/agent_framework/_compaction.py @@ -813,6 +813,7 @@ def __init__( compact_to: int, tokenizer: TokenizerProtocol | None = None, preserve_system: bool = True, + preserve_first_user_group: bool = False, ) -> None: """Create a truncation strategy. @@ -824,6 +825,8 @@ def __init__( tokenizer: Optional tokenizer used for token-based truncation. preserve_system: When True, system groups remain included and only non-system groups are eligible for exclusion. + preserve_first_user_group: When True, the earliest user group + remains included along with the minimum retained group. """ if max_n <= 0: raise ValueError("max_n must be greater than 0.") @@ -835,6 +838,7 @@ def __init__( self.compact_to = compact_to self.tokenizer = tokenizer self.preserve_system = preserve_system + self.preserve_first_user_group = preserve_first_user_group async def __call__(self, messages: list[Message]) -> bool: ordered_group_ids = _ordered_group_ids_from_annotations(messages) @@ -850,6 +854,13 @@ async def __call__(self, messages: list[Message]) -> bool: protected_ids: set[str] = set() if self.preserve_system: protected_ids = {group_id for group_id in ordered_group_ids if kinds.get(group_id) == "system"} + if self.preserve_first_user_group: + first_user_group_id = next( + (group_id for group_id in ordered_group_ids if kinds.get(group_id) == "user"), + None, + ) + if first_user_group_id is not None: + protected_ids.add(first_user_group_id) protected_ids.update(_minimum_retained_group_ids(messages, ordered_group_ids, kinds)) changed = False @@ -1472,6 +1483,50 @@ async def __call__(self, messages: list[Message]) -> bool: return changed +async def _run_compaction_strategy( + messages: list[Message], + *, + strategy: CompactionStrategy, + tokenizer: TokenizerProtocol | None = None, + phase: str, +) -> bool: + """Apply a strategy and log aggregate metrics when it changes context.""" + resolved_tokenizer = tokenizer + if resolved_tokenizer is None: + strategy_tokenizer = getattr(strategy, "tokenizer", None) + if isinstance(strategy_tokenizer, TokenizerProtocol): + resolved_tokenizer = strategy_tokenizer + + annotate_message_groups(messages) + if resolved_tokenizer is not None: + annotate_token_counts(messages, tokenizer=resolved_tokenizer) + before_message_count = _count_included_messages(messages) + before_token_count = included_token_count(messages) if resolved_tokenizer is not None else None + + changed = await strategy(messages) + if not changed: + return False + + annotate_message_groups(messages) + if resolved_tokenizer is not None: + annotate_token_counts(messages, tokenizer=resolved_tokenizer) + after_message_count = _count_included_messages(messages) + after_token_count = included_token_count(messages) if resolved_tokenizer is not None else None + strategy_name = getattr(strategy, "__name__", type(strategy).__name__) + logger.info( + "Compaction applied", + extra={ + "compaction_phase": phase, + "compaction_strategy": strategy_name, + "compaction_included_messages_before": before_message_count, + "compaction_included_messages_after": after_message_count, + "compaction_included_tokens_before": before_token_count, + "compaction_included_tokens_after": after_token_count, + }, + ) + return True + + async def apply_compaction( messages: list[Message], *, @@ -1481,10 +1536,12 @@ async def apply_compaction( """Apply configured compaction and return projected model-input messages.""" if strategy is None: return messages - annotate_message_groups(messages) - if tokenizer is not None: - annotate_token_counts(messages, tokenizer=tokenizer) - await strategy(messages) + await _run_compaction_strategy( + messages, + strategy=strategy, + tokenizer=tokenizer, + phase="in_run", + ) return project_included_messages(messages) @@ -1579,10 +1636,12 @@ async def before_run( if not all_messages: return - annotate_message_groups(all_messages) - if self.tokenizer is not None: - annotate_token_counts(all_messages, tokenizer=self.tokenizer) - await self.before_strategy(all_messages) + await _run_compaction_strategy( + all_messages, + strategy=self.before_strategy, + tokenizer=self.tokenizer, + phase="before_run", + ) projected = project_included_messages(all_messages) projected_set = {id(m) for m in projected} @@ -1611,10 +1670,12 @@ async def after_run( return stored_messages: list[Message] = raw_messages # type: ignore[assignment] - annotate_message_groups(stored_messages) - if self.tokenizer is not None: - annotate_token_counts(stored_messages, tokenizer=self.tokenizer) - await self.after_strategy(stored_messages) + await _run_compaction_strategy( + stored_messages, + strategy=self.after_strategy, + tokenizer=self.tokenizer, + phase="after_run", + ) # Keep all messages (including excluded) in storage so annotations are # preserved. The history provider's ``skip_excluded`` flag controls @@ -1663,6 +1724,7 @@ def __init__( tool_eviction_threshold: float = DEFAULT_TOOL_EVICTION_THRESHOLD, truncation_threshold: float = DEFAULT_TRUNCATION_THRESHOLD, keep_last_tool_call_groups: int = 4, + preserve_first_user_group: bool = False, ) -> None: """Create a context-window compaction strategy. @@ -1681,6 +1743,8 @@ def __init__( keep_last_tool_call_groups: Number of most recent tool-call groups to retain verbatim during tool eviction. Older groups are collapsed into summaries. Defaults to 4. + preserve_first_user_group: Whether destructive truncation preserves + the earliest user group. Defaults to False. Raises: ValueError: If thresholds are out of range or inconsistent. @@ -1706,24 +1770,18 @@ def __init__( self.input_budget_tokens = input_budget self.tool_eviction_threshold = tool_eviction_threshold self.truncation_threshold = truncation_threshold + self.tokenizer = resolved_tokenizer + self._tool_eviction_tokens = tool_eviction_tokens + self._truncation_tokens = truncation_tokens - self._tool_eviction = TokenBudgetComposedStrategy( - token_budget=tool_eviction_tokens, - tokenizer=resolved_tokenizer, - strategies=[ - ToolResultCompactionStrategy(keep_last_tool_call_groups=keep_last_tool_call_groups), - ], + self._tool_eviction = ToolResultCompactionStrategy( + keep_last_tool_call_groups=keep_last_tool_call_groups, ) - self._truncation = TokenBudgetComposedStrategy( - token_budget=truncation_tokens, + self._truncation = TruncationStrategy( + max_n=truncation_tokens, + compact_to=tool_eviction_tokens, tokenizer=resolved_tokenizer, - strategies=[ - TruncationStrategy( - max_n=truncation_tokens, - compact_to=tool_eviction_tokens, - tokenizer=resolved_tokenizer, - ), - ], + preserve_first_user_group=preserve_first_user_group, ) async def __call__(self, messages: list[Message]) -> bool: @@ -1732,8 +1790,18 @@ async def __call__(self, messages: list[Message]) -> bool: Returns: True if compaction changed message inclusion; otherwise False. """ - changed = await self._tool_eviction(messages) - return (await self._truncation(messages)) or changed + annotate_message_groups(messages) + annotate_token_counts(messages, tokenizer=self.tokenizer) + + changed = False + if included_token_count(messages) > self._tool_eviction_tokens: + changed = await self._tool_eviction(messages) + annotate_message_groups(messages) + annotate_token_counts(messages, tokenizer=self.tokenizer) + + if included_token_count(messages) > self._truncation_tokens: + changed = (await self._truncation(messages)) or changed + return changed __all__ = [ diff --git a/python/packages/core/agent_framework/_middleware.py b/python/packages/core/agent_framework/_middleware.py index 45a3dcbf823..ca83168328b 100644 --- a/python/packages/core/agent_framework/_middleware.py +++ b/python/packages/core/agent_framework/_middleware.py @@ -72,6 +72,56 @@ def _empty_async_iterable() -> AsyncIterable[Any]: return _EmptyAsyncIterator() +def _propagate_compaction_summaries( + source_messages: list[Message], + working_messages: Sequence[Message], + previous_message_ids: set[int], +) -> None: + """Propagate summaries of source messages without persisting middleware rewrites.""" + from ._compaction import ( + EXCLUDE_REASON_KEY, + EXCLUDED_KEY, + GROUP_ANNOTATION_KEY, + SUMMARIZED_BY_SUMMARY_ID_KEY, + SUMMARY_OF_MESSAGE_IDS_KEY, + ) + + source_message_ids = {message.message_id for message in source_messages if message.message_id} + for message in working_messages: + if id(message) in previous_message_ids: + continue + + annotation_value: Any = message.additional_properties.get(GROUP_ANNOTATION_KEY) + if not isinstance(annotation_value, Mapping): + continue + annotation = cast("Mapping[str, Any]", annotation_value) + summarized_message_ids: Any = annotation.get(SUMMARY_OF_MESSAGE_IDS_KEY) + summarized_id_set: set[str] = ( + {value for value in cast("list[Any]", summarized_message_ids) if isinstance(value, str)} + if isinstance(summarized_message_ids, list) + else set() + ) + linked_sources: list[tuple[int, Message, dict[str, Any]]] = [] + for source_index, source_message in enumerate(source_messages): + source_annotation_value: Any = source_message.additional_properties.get(GROUP_ANNOTATION_KEY) + if not isinstance(source_annotation_value, dict): + continue + source_annotation = cast("dict[str, Any]", source_annotation_value) + if source_annotation.get(SUMMARIZED_BY_SUMMARY_ID_KEY) == message.message_id: + linked_sources.append((source_index, source_message, source_annotation)) + if not linked_sources: + continue + + if summarized_id_set and summarized_id_set.issubset(source_message_ids): + source_messages.insert(linked_sources[0][0], message) + continue + + for _, source_message, source_annotation in linked_sources: + source_annotation.pop(SUMMARIZED_BY_SUMMARY_ID_KEY, None) + source_message.additional_properties.pop(EXCLUDED_KEY, None) + source_message.additional_properties.pop(EXCLUDE_REASON_KEY, None) + + class MiddlewareTermination(MiddlewareException): """Control-flow exception to terminate middleware execution early.""" @@ -1376,9 +1426,17 @@ def get_response( ) async def _execute() -> ChatResponse | ResponseStream[ChatResponseUpdate, ChatResponse] | None: + def _final_handler( + middleware_context: ChatContext, + ) -> Awaitable[ChatResponse] | ResponseStream[ChatResponseUpdate, ChatResponse]: + return self._middleware_handler( + middleware_context, + source_messages=messages if isinstance(messages, list) else None, + ) + return await pipeline.execute( context=context, - final_handler=self._middleware_handler, + final_handler=_final_handler, ) if stream: @@ -1402,21 +1460,59 @@ async def _execute_stream() -> ResponseStream[ChatResponseUpdate, ChatResponse]: return _execute() # type: ignore[return-value] def _middleware_handler( - self, context: ChatContext + self, + context: ChatContext, + *, + source_messages: list[Message] | None = None, ) -> Awaitable[ChatResponse] | ResponseStream[ChatResponseUpdate, ChatResponse]: """Internal middleware handler to adapt to pipeline.""" handler_kwargs = dict(context.kwargs) compaction_strategy = handler_kwargs.pop("compaction_strategy", None) tokenizer = handler_kwargs.pop("tokenizer", None) - return super().get_response( # type: ignore[misc, no-any-return] - messages=context.messages, - stream=context.stream, - options=context.options or {}, - compaction_strategy=compaction_strategy, - tokenizer=tokenizer, - function_invocation_kwargs=context.function_invocation_kwargs, - client_kwargs=handler_kwargs, + result = cast( + "Awaitable[ChatResponse] | ResponseStream[ChatResponseUpdate, ChatResponse]", + super().get_response( # type: ignore[misc] + messages=context.messages, + stream=context.stream, + options=context.options or {}, + compaction_strategy=compaction_strategy, + tokenizer=tokenizer, + function_invocation_kwargs=context.function_invocation_kwargs, + client_kwargs=handler_kwargs, + ), ) + if source_messages is None: + return result + + if isinstance(result, ResponseStream): + stream_result = cast("ResponseStream[ChatResponseUpdate, ChatResponse]", result) + + async def _resolve_stream() -> ResponseStream[ChatResponseUpdate, ChatResponse]: + previous_message_ids = {id(message) for message in context.messages} + try: + await stream_result + finally: + _propagate_compaction_summaries( + source_messages, + context.messages, + previous_message_ids, + ) + return stream_result + + return ResponseStream[ChatResponseUpdate, ChatResponse].from_awaitable(_resolve_stream()) + + async def _resolve_response() -> ChatResponse: + previous_message_ids = {id(message) for message in context.messages} + try: + return await result + finally: + _propagate_compaction_summaries( + source_messages, + context.messages, + previous_message_ids, + ) + + return _resolve_response() class AgentMiddlewareLayer: diff --git a/python/packages/core/tests/core/test_clients.py b/python/packages/core/tests/core/test_clients.py index b57d5eba48b..dfc98abd42a 100644 --- a/python/packages/core/tests/core/test_clients.py +++ b/python/packages/core/tests/core/test_clients.py @@ -10,11 +10,13 @@ GROUP_ANNOTATION_KEY, GROUP_TOKEN_COUNT_KEY, BaseChatClient, + ChatMiddleware, ChatResponse, ChatResponseUpdate, Content, Message, SlidingWindowStrategy, + SummarizationStrategy, SupportsChatGetResponse, ToolResultCompactionStrategy, TruncationStrategy, @@ -22,6 +24,16 @@ ) +class _NoOpChatMiddleware(ChatMiddleware): + async def process(self, context: Any, call_next: Any) -> None: + await call_next() + + +class _FixedSummarizer: + async def get_response(self, *args: Any, **kwargs: Any) -> ChatResponse: + return ChatResponse(messages=[Message(role="assistant", contents=["SUMMARY"])]) + + class _FixedTokenizer: def __init__(self, token_count: int) -> None: self.token_count = token_count @@ -293,8 +305,10 @@ def _is_tool_result_summary(message: Message) -> bool: return message.role == "assistant" and text.startswith("[Tool results:") +@pytest.mark.parametrize("with_chat_middleware", [False, True]) async def test_function_loop_persists_inserted_summaries_across_iterations( chat_client_base: SupportsChatGetResponse, + with_chat_middleware: bool, ) -> None: # Regression test for #4991: compaction inserts summary messages and excludes the # originals. Across tool-loop iterations the exclusion flags persisted (shared Message @@ -303,6 +317,8 @@ async def test_function_loop_persists_inserted_summaries_across_iterations( chat_client_base.function_invocation_configuration["enabled"] = True # type: ignore[attr-defined] # ty: ignore[unresolved-attribute] chat_client_base.function_invocation_configuration["max_iterations"] = 3 # type: ignore[attr-defined] # ty: ignore[unresolved-attribute] chat_client_base.compaction_strategy = ToolResultCompactionStrategy(keep_last_tool_call_groups=1) # type: ignore[attr-defined] # ty: ignore[unresolved-attribute] + if with_chat_middleware: + chat_client_base.chat_middleware = [_NoOpChatMiddleware()] # type: ignore[attr-defined] # ty: ignore[unresolved-attribute] @tool(name="lookup_weather", approval_mode="never_require") def lookup_weather(location: str) -> str: @@ -361,14 +377,18 @@ def _tool_call_update(call_id: str, location: str) -> list[ChatResponseUpdate]: ] +@pytest.mark.parametrize("with_chat_middleware", [False, True]) async def test_function_loop_persists_inserted_summaries_across_iterations_streaming( chat_client_base: SupportsChatGetResponse, + with_chat_middleware: bool, ) -> 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] # ty: ignore[unresolved-attribute] chat_client_base.function_invocation_configuration["max_iterations"] = 3 # type: ignore[attr-defined] # ty: ignore[unresolved-attribute] chat_client_base.compaction_strategy = ToolResultCompactionStrategy(keep_last_tool_call_groups=1) # type: ignore[attr-defined] # ty: ignore[unresolved-attribute] + if with_chat_middleware: + chat_client_base.chat_middleware = [_NoOpChatMiddleware()] # type: ignore[attr-defined] # ty: ignore[unresolved-attribute] @tool(name="lookup_weather", approval_mode="never_require") def lookup_weather(location: str) -> str: @@ -411,8 +431,10 @@ def _capture( assert "Paris" in summary_text +@pytest.mark.parametrize("with_chat_middleware", [False, True]) async def test_function_loop_compaction_conversation_id_mode_does_not_resend_history( chat_client_base: SupportsChatGetResponse, + with_chat_middleware: bool, ) -> 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 @@ -420,6 +442,8 @@ async def test_function_loop_compaction_conversation_id_mode_does_not_resend_his chat_client_base.function_invocation_configuration["enabled"] = True # type: ignore[attr-defined] # ty: ignore[unresolved-attribute] chat_client_base.function_invocation_configuration["max_iterations"] = 3 # type: ignore[attr-defined] # ty: ignore[unresolved-attribute] chat_client_base.compaction_strategy = ToolResultCompactionStrategy(keep_last_tool_call_groups=1) # type: ignore[attr-defined] # ty: ignore[unresolved-attribute] + if with_chat_middleware: + chat_client_base.chat_middleware = [_NoOpChatMiddleware()] # type: ignore[attr-defined] # ty: ignore[unresolved-attribute] @tool(name="lookup_weather", approval_mode="never_require") def lookup_weather(location: str) -> str: @@ -462,6 +486,59 @@ async def _capture( assert not any(_is_tool_result_summary(message) for message in sent) +async def test_chat_middleware_does_not_persist_summary_of_middleware_messages( + chat_client_base: SupportsChatGetResponse, +) -> None: + class _InsertEphemeralMessage(ChatMiddleware): + async def process(self, context: Any, call_next: Any) -> None: + context.messages.insert(1, Message(role="user", contents=["ephemeral middleware context"])) + await call_next() + + chat_client_base.function_invocation_configuration["enabled"] = False # type: ignore[attr-defined] # ty: ignore[unresolved-attribute] + chat_client_base.chat_middleware = [_InsertEphemeralMessage()] # type: ignore[attr-defined] # ty: ignore[unresolved-attribute] + chat_client_base.compaction_strategy = SummarizationStrategy( # type: ignore[attr-defined] # ty: ignore[unresolved-attribute] + client=_FixedSummarizer(), # type: ignore[arg-type] # pyrefly: ignore[bad-argument-type] # ty: ignore[invalid-argument-type] + target_count=1, + threshold=0, + ) + messages = [ + Message(role="user", contents=["original request"]), + Message(role="assistant", contents=["old response"]), + Message(role="user", contents=["latest request"]), + ] + + await chat_client_base.get_response(messages) + + assert [message.text for message in messages] == ["original request", "old response", "latest request"] + assert all(not message.additional_properties.get("_excluded", False) for message in messages) + + +async def test_chat_middleware_persists_compaction_summary_when_model_call_fails( + chat_client_base: SupportsChatGetResponse, +) -> None: + async def _raise_after_compaction(**kwargs: Any) -> ChatResponse: + raise RuntimeError("model call failed") + + chat_client_base.function_invocation_configuration["enabled"] = False # type: ignore[attr-defined] # ty: ignore[unresolved-attribute] + chat_client_base.chat_middleware = [_NoOpChatMiddleware()] # type: ignore[attr-defined] # ty: ignore[unresolved-attribute] + chat_client_base.compaction_strategy = ToolResultCompactionStrategy(keep_last_tool_call_groups=1) # type: ignore[attr-defined] # ty: ignore[unresolved-attribute] + messages = [ + Message(role="user", contents=["request"]), + _tool_call_response("call_1", "first").messages[0], + Message(role="tool", contents=[Content.from_function_result(call_id="call_1", result="first result")]), + _tool_call_response("call_2", "second").messages[0], + Message(role="tool", contents=[Content.from_function_result(call_id="call_2", result="second result")]), + ] + + with ( + patch.object(chat_client_base, "_inner_get_response", side_effect=_raise_after_compaction), + pytest.raises(RuntimeError, match="model call failed"), + ): + await chat_client_base.get_response(messages) + + assert any(_is_tool_result_summary(message) for message in messages) + + def test_base_client_as_agent_does_not_copy_client_compaction_defaults( chat_client_base: SupportsChatGetResponse, ) -> None: diff --git a/python/packages/core/tests/core/test_compaction.py b/python/packages/core/tests/core/test_compaction.py index 35829859971..dd2547e9a7c 100644 --- a/python/packages/core/tests/core/test_compaction.py +++ b/python/packages/core/tests/core/test_compaction.py @@ -1149,6 +1149,33 @@ async def test_apply_compaction_projects_included_messages_only() -> None: assert projected[0].role == "system" +async def test_apply_compaction_logs_changed_context_without_content(caplog: Any) -> None: + messages = [ + Message(role="user", contents=["sensitive old request"]), + Message(role="user", contents=["latest request"]), + ] + strategy = TruncationStrategy(max_n=1, compact_to=1) + + with caplog.at_level(logging.INFO, logger="agent_framework"): + await apply_compaction(messages, strategy=strategy) + + assert len(caplog.messages) == 1 + assert caplog.messages[0] == "Compaction applied" + record = caplog.records[0] + assert record.compaction_phase == "in_run" + assert record.compaction_strategy == "TruncationStrategy" + assert record.compaction_included_messages_before == 2 + assert record.compaction_included_messages_after == 1 + assert record.compaction_included_tokens_before is record.compaction_included_tokens_after is None + + caplog.clear() + await apply_compaction( + [Message(role="user", contents=["request"])], + strategy=TruncationStrategy(max_n=2, compact_to=1), + ) + assert caplog.messages == [] + + # --- ToolResultCompactionStrategy tests --- @@ -1634,7 +1661,7 @@ def __init__(self) -> None: self.state: dict[str, Any] = {} -async def test_compaction_provider_after_run_compacts_stored_history() -> None: +async def test_compaction_provider_after_run_compacts_stored_history(caplog: Any) -> None: """after_run annotates exclusions on stored messages without removing them.""" provider = CompactionProvider( after_strategy=SelectiveToolCallCompactionStrategy(keep_last_tool_call_groups=0), @@ -1652,8 +1679,8 @@ async def test_compaction_provider_after_run_compacts_stored_history() -> None: ] } - context = _MockSessionContext() - await provider.after_run(agent=None, session=session, context=context, state={}) + with caplog.at_level(logging.INFO, logger="agent_framework"): + await provider.after_run(agent=None, session=session, context=_MockSessionContext(), state={}) stored = session.state["in_memory_history"]["messages"] # All messages are kept; tool-call group is excluded via annotation. @@ -1661,6 +1688,10 @@ async def test_compaction_provider_after_run_compacts_stored_history() -> None: excluded = [m for m in stored if m.additional_properties.get("_excluded", False)] assert len(excluded) == 2 # assistant function_call + tool result assert any(m.text == "final answer" for m in stored if not m.additional_properties.get("_excluded", False)) + assert len(caplog.messages) == 1 + record = caplog.records[0] + assert record.compaction_phase == "after_run" + assert record.compaction_strategy == "SelectiveToolCallCompactionStrategy" async def test_compaction_provider_after_run_noop_without_history() -> None: @@ -1839,6 +1870,22 @@ async def test_context_window_strategy_tool_eviction_triggers_at_threshold() -> assert len(truncation_excluded) == 0 +async def test_context_window_strategy_does_not_truncate_between_thresholds_without_tools() -> None: + messages = [ + Message(role="user", contents=["u " * 500]), + Message(role="assistant", contents=["a " * 500]), + ] + strategy = ContextWindowCompactionStrategy( + max_context_window_tokens=1000, + max_output_tokens=100, + ) + + changed = await strategy(messages) + + assert changed is False + assert included_messages(messages) == messages + + async def test_context_window_strategy_truncation_triggers_above_80_pct() -> None: """Truncation fires when tokens exceed 80% of input budget.""" # input_budget = 1000 - 100 = 900 @@ -1867,6 +1914,26 @@ async def test_context_window_strategy_truncation_triggers_above_80_pct() -> Non assert len(projected) < 5 +async def test_context_window_strategy_can_preserve_first_user_group() -> None: + messages = [ + Message(role="user", contents=["original " * 400]), + Message(role="assistant", contents=["old answer " * 400]), + Message(role="user", contents=["latest " * 400]), + ] + strategy = ContextWindowCompactionStrategy( + max_context_window_tokens=1000, + max_output_tokens=100, + preserve_first_user_group=True, + ) + + changed = await strategy(messages) + + assert changed is True + projected = included_messages(messages) + assert any(message.text == "original " * 400 for message in projected) + assert any(message.text == "latest " * 400 for message in projected) + + async def test_context_window_strategy_keep_last_tool_call_groups_respected() -> None: """The keep_last_tool_call_groups parameter controls how many groups are retained.""" # Create enough tokens to trigger tool eviction (>50% of input budget) From 4a88cae1949a8cf74d350eb99fe757878e11cfe8 Mon Sep 17 00:00:00 2001 From: eavanvalkenburg Date: Thu, 27 Aug 2026 14:36:54 +0200 Subject: [PATCH 2/2] Address compaction review feedback Reconcile summaries at the middleware pipeline boundary, support sequence replacement and nested summaries, strengthen logging coverage, refresh threshold documentation, and warn when protected groups exceed the input budget. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: 726d66ea-6cff-4cb8-9090-8d919392271d --- .../core/agent_framework/_compaction.py | 21 ++- .../core/agent_framework/_middleware.py | 140 +++++++++--------- .../packages/core/tests/core/test_clients.py | 84 ++++++++++- .../core/tests/core/test_compaction.py | 16 +- 4 files changed, 175 insertions(+), 86 deletions(-) diff --git a/python/packages/core/agent_framework/_compaction.py b/python/packages/core/agent_framework/_compaction.py index f382ed2ce64..43c6c007f85 100644 --- a/python/packages/core/agent_framework/_compaction.py +++ b/python/packages/core/agent_framework/_compaction.py @@ -802,8 +802,8 @@ class TruncationStrategy: - token count when ``tokenizer`` is provided - included message count when ``tokenizer`` is not provided Compaction triggers when the metric exceeds ``max_n`` and trims toward - ``compact_to``. The minimum retained group is never excluded, so the - result may remain above ``compact_to`` when that group alone exceeds it. + ``compact_to``. Protected groups are never excluded, so the result may + remain above ``compact_to`` when those groups alone exceed it. """ def __init__( @@ -1693,9 +1693,10 @@ class ContextWindowCompactionStrategy: 2. **Truncation** — removes oldest non-system groups when included tokens exceed ``truncation_threshold`` of the input budget. - The class uses two independent :class:`TokenBudgetComposedStrategy` - instances — one per phase — so each fires only when its own threshold - is exceeded. + Each phase checks its threshold explicitly. Token counts are refreshed + after tool-result eviction before deciding whether destructive truncation + is necessary. If protected groups still exceed the input budget after + truncation, the strategy preserves them and emits a structured warning. Examples: .. code-block:: python @@ -1801,6 +1802,16 @@ async def __call__(self, messages: list[Message]) -> bool: if included_token_count(messages) > self._truncation_tokens: changed = (await self._truncation(messages)) or changed + remaining_tokens = included_token_count(messages) + if remaining_tokens > self.input_budget_tokens: + logger.warning( + "Compaction could not fit protected messages within the input budget", + extra={ + "compaction_strategy": type(self).__name__, + "compaction_included_tokens_after": remaining_tokens, + "compaction_input_budget_tokens": self.input_budget_tokens, + }, + ) return changed diff --git a/python/packages/core/agent_framework/_middleware.py b/python/packages/core/agent_framework/_middleware.py index ca83168328b..562d9d7df85 100644 --- a/python/packages/core/agent_framework/_middleware.py +++ b/python/packages/core/agent_framework/_middleware.py @@ -87,6 +87,7 @@ def _propagate_compaction_summaries( ) source_message_ids = {message.message_id for message in source_messages if message.message_id} + candidates: list[tuple[Message, set[str]]] = [] for message in working_messages: if id(message) in previous_message_ids: continue @@ -96,30 +97,58 @@ def _propagate_compaction_summaries( continue annotation = cast("Mapping[str, Any]", annotation_value) summarized_message_ids: Any = annotation.get(SUMMARY_OF_MESSAGE_IDS_KEY) - summarized_id_set: set[str] = ( - {value for value in cast("list[Any]", summarized_message_ids) if isinstance(value, str)} - if isinstance(summarized_message_ids, list) - else set() - ) - linked_sources: list[tuple[int, Message, dict[str, Any]]] = [] - for source_index, source_message in enumerate(source_messages): - source_annotation_value: Any = source_message.additional_properties.get(GROUP_ANNOTATION_KEY) - if not isinstance(source_annotation_value, dict): - continue - source_annotation = cast("dict[str, Any]", source_annotation_value) - if source_annotation.get(SUMMARIZED_BY_SUMMARY_ID_KEY) == message.message_id: - linked_sources.append((source_index, source_message, source_annotation)) - if not linked_sources: + if ( + message.message_id + and isinstance(summarized_message_ids, list) + and summarized_message_ids + and all(isinstance(value, str) for value in cast("list[Any]", summarized_message_ids)) + ): + candidates.append((message, set(cast("list[str]", summarized_message_ids)))) + + dependencies = {message.message_id: summary_ids for message, summary_ids in candidates if message.message_id} + supported_ids = set(source_message_ids) + pending_ids = set(dependencies) + while pending_ids: + newly_supported = {summary_id for summary_id in pending_ids if dependencies[summary_id].issubset(supported_ids)} + if not newly_supported: + break + supported_ids.update(newly_supported) + pending_ids.difference_update(newly_supported) + + accepted_ids = set(dependencies).difference(pending_ids) + for message in [*source_messages, *(candidate for candidate, _ in candidates)]: + annotation_value = message.additional_properties.get(GROUP_ANNOTATION_KEY) + if not isinstance(annotation_value, dict): continue - - if summarized_id_set and summarized_id_set.issubset(source_message_ids): - source_messages.insert(linked_sources[0][0], message) + annotation = cast("dict[str, Any]", annotation_value) + if annotation.get(SUMMARIZED_BY_SUMMARY_ID_KEY) not in pending_ids: continue - - for _, source_message, source_annotation in linked_sources: - source_annotation.pop(SUMMARIZED_BY_SUMMARY_ID_KEY, None) - source_message.additional_properties.pop(EXCLUDED_KEY, None) - source_message.additional_properties.pop(EXCLUDE_REASON_KEY, None) + annotation.pop(SUMMARIZED_BY_SUMMARY_ID_KEY, None) + message.additional_properties.pop(EXCLUDED_KEY, None) + message.additional_properties.pop(EXCLUDE_REASON_KEY, None) + + def source_dependencies(summary_id: str) -> set[str]: + expanded: set[str] = set() + for dependency_id in dependencies[summary_id]: + if dependency_id in source_message_ids: + expanded.add(dependency_id) + elif dependency_id in accepted_ids: + expanded.update(source_dependencies(dependency_id)) + return expanded + + for message, _ in candidates: + if message.message_id not in accepted_ids: + continue + summarized_source_ids = source_dependencies(message.message_id) + insertion_index = min( + ( + index + for index, source_message in enumerate(source_messages) + if source_message.message_id in summarized_source_ids + ), + default=len(source_messages), + ) + source_messages.insert(insertion_index, message) class MiddlewareTermination(MiddlewareException): @@ -1424,20 +1453,22 @@ def get_response( kwargs=context_kwargs, function_invocation_kwargs=function_invocation_kwargs, ) + source_messages = messages if isinstance(messages, list) else None + baseline_message_ids = {id(message) for message in messages} async def _execute() -> ChatResponse | ResponseStream[ChatResponseUpdate, ChatResponse] | None: - def _final_handler( - middleware_context: ChatContext, - ) -> Awaitable[ChatResponse] | ResponseStream[ChatResponseUpdate, ChatResponse]: - return self._middleware_handler( - middleware_context, - source_messages=messages if isinstance(messages, list) else None, + try: + return await pipeline.execute( + context=context, + final_handler=self._middleware_handler, ) - - return await pipeline.execute( - context=context, - final_handler=_final_handler, - ) + finally: + if source_messages is not None: + _propagate_compaction_summaries( + source_messages, + context.messages, + baseline_message_ids, + ) if stream: # For streaming, wrap execution in ResponseStream.from_awaitable @@ -1460,19 +1491,18 @@ async def _execute_stream() -> ResponseStream[ChatResponseUpdate, ChatResponse]: return _execute() # type: ignore[return-value] def _middleware_handler( - self, - context: ChatContext, - *, - source_messages: list[Message] | None = None, + self, context: ChatContext ) -> Awaitable[ChatResponse] | ResponseStream[ChatResponseUpdate, ChatResponse]: """Internal middleware handler to adapt to pipeline.""" handler_kwargs = dict(context.kwargs) compaction_strategy = handler_kwargs.pop("compaction_strategy", None) tokenizer = handler_kwargs.pop("tokenizer", None) - result = cast( + working_messages = context.messages if isinstance(context.messages, list) else list(context.messages) + context.messages = working_messages + return cast( "Awaitable[ChatResponse] | ResponseStream[ChatResponseUpdate, ChatResponse]", super().get_response( # type: ignore[misc] - messages=context.messages, + messages=working_messages, stream=context.stream, options=context.options or {}, compaction_strategy=compaction_strategy, @@ -1481,38 +1511,6 @@ def _middleware_handler( client_kwargs=handler_kwargs, ), ) - if source_messages is None: - return result - - if isinstance(result, ResponseStream): - stream_result = cast("ResponseStream[ChatResponseUpdate, ChatResponse]", result) - - async def _resolve_stream() -> ResponseStream[ChatResponseUpdate, ChatResponse]: - previous_message_ids = {id(message) for message in context.messages} - try: - await stream_result - finally: - _propagate_compaction_summaries( - source_messages, - context.messages, - previous_message_ids, - ) - return stream_result - - return ResponseStream[ChatResponseUpdate, ChatResponse].from_awaitable(_resolve_stream()) - - async def _resolve_response() -> ChatResponse: - previous_message_ids = {id(message) for message in context.messages} - try: - return await result - finally: - _propagate_compaction_summaries( - source_messages, - context.messages, - previous_message_ids, - ) - - return _resolve_response() class AgentMiddlewareLayer: diff --git a/python/packages/core/tests/core/test_clients.py b/python/packages/core/tests/core/test_clients.py index dfc98abd42a..75251fc6e43 100644 --- a/python/packages/core/tests/core/test_clients.py +++ b/python/packages/core/tests/core/test_clients.py @@ -20,6 +20,7 @@ SupportsChatGetResponse, ToolResultCompactionStrategy, TruncationStrategy, + apply_compaction, tool, ) @@ -29,6 +30,12 @@ async def process(self, context: Any, call_next: Any) -> None: await call_next() +class _TupleMessagesChatMiddleware(ChatMiddleware): + async def process(self, context: Any, call_next: Any) -> None: + context.messages = tuple(context.messages) + await call_next() + + class _FixedSummarizer: async def get_response(self, *args: Any, **kwargs: Any) -> ChatResponse: return ChatResponse(messages=[Message(role="assistant", contents=["SUMMARY"])]) @@ -305,10 +312,14 @@ def _is_tool_result_summary(message: Message) -> bool: return message.role == "assistant" and text.startswith("[Tool results:") -@pytest.mark.parametrize("with_chat_middleware", [False, True]) +@pytest.mark.parametrize( + "chat_middleware", + [None, _NoOpChatMiddleware(), _TupleMessagesChatMiddleware()], + ids=["none", "list", "tuple"], +) async def test_function_loop_persists_inserted_summaries_across_iterations( chat_client_base: SupportsChatGetResponse, - with_chat_middleware: bool, + chat_middleware: ChatMiddleware | None, ) -> None: # Regression test for #4991: compaction inserts summary messages and excludes the # originals. Across tool-loop iterations the exclusion flags persisted (shared Message @@ -317,8 +328,8 @@ async def test_function_loop_persists_inserted_summaries_across_iterations( chat_client_base.function_invocation_configuration["enabled"] = True # type: ignore[attr-defined] # ty: ignore[unresolved-attribute] chat_client_base.function_invocation_configuration["max_iterations"] = 3 # type: ignore[attr-defined] # ty: ignore[unresolved-attribute] chat_client_base.compaction_strategy = ToolResultCompactionStrategy(keep_last_tool_call_groups=1) # type: ignore[attr-defined] # ty: ignore[unresolved-attribute] - if with_chat_middleware: - chat_client_base.chat_middleware = [_NoOpChatMiddleware()] # type: ignore[attr-defined] # ty: ignore[unresolved-attribute] + if chat_middleware is not None: + chat_client_base.chat_middleware = [chat_middleware] # type: ignore[attr-defined] # ty: ignore[unresolved-attribute] @tool(name="lookup_weather", approval_mode="never_require") def lookup_weather(location: str) -> str: @@ -513,6 +524,71 @@ async def process(self, context: Any, call_next: Any) -> None: assert all(not message.additional_properties.get("_excluded", False) for message in messages) +@pytest.mark.parametrize("continue_pipeline", [True, False], ids=["call-next", "terminate"]) +async def test_chat_middleware_reconciles_compaction_before_downstream( + chat_client_base: SupportsChatGetResponse, + continue_pipeline: bool, +) -> None: + class _CompactBeforeDownstream(ChatMiddleware): + async def process(self, context: Any, call_next: Any) -> None: + assert isinstance(context.messages, list) + await apply_compaction( + context.messages, + strategy=ToolResultCompactionStrategy(keep_last_tool_call_groups=1), + ) + if continue_pipeline: + await call_next() + else: + context.result = ChatResponse(messages=[Message(role="assistant", contents=["terminated"])]) + + chat_client_base.function_invocation_configuration["enabled"] = False # type: ignore[attr-defined] # ty: ignore[unresolved-attribute] + chat_client_base.chat_middleware = [_CompactBeforeDownstream()] # type: ignore[attr-defined] # ty: ignore[unresolved-attribute] + messages = [ + Message(role="user", contents=["request"]), + _tool_call_response("call_1", "first").messages[0], + Message(role="tool", contents=[Content.from_function_result(call_id="call_1", result="first result")]), + _tool_call_response("call_2", "second").messages[0], + Message(role="tool", contents=[Content.from_function_result(call_id="call_2", result="second result")]), + ] + + await chat_client_base.get_response(messages) + + assert any(_is_tool_result_summary(message) for message in messages) + + +async def test_chat_middleware_reconciles_nested_compaction_summaries( + chat_client_base: SupportsChatGetResponse, +) -> None: + class _CompactToolsBeforeDownstream(ChatMiddleware): + async def process(self, context: Any, call_next: Any) -> None: + assert isinstance(context.messages, list) + await apply_compaction( + context.messages, + strategy=ToolResultCompactionStrategy(keep_last_tool_call_groups=1), + ) + await call_next() + + chat_client_base.function_invocation_configuration["enabled"] = False # type: ignore[attr-defined] # ty: ignore[unresolved-attribute] + chat_client_base.chat_middleware = [_CompactToolsBeforeDownstream()] # type: ignore[attr-defined] # ty: ignore[unresolved-attribute] + chat_client_base.compaction_strategy = SummarizationStrategy( # type: ignore[attr-defined] # ty: ignore[unresolved-attribute] + client=_FixedSummarizer(), # type: ignore[arg-type] # pyrefly: ignore[bad-argument-type] # ty: ignore[invalid-argument-type] + target_count=1, + threshold=0, + ) + messages = [ + Message(role="user", contents=["request"]), + _tool_call_response("call_1", "first").messages[0], + Message(role="tool", contents=[Content.from_function_result(call_id="call_1", result="first result")]), + _tool_call_response("call_2", "second").messages[0], + Message(role="tool", contents=[Content.from_function_result(call_id="call_2", result="second result")]), + Message(role="assistant", contents=["latest"]), + ] + + await chat_client_base.get_response(messages) + + assert any(message.text == "SUMMARY" and not message.additional_properties.get("_excluded") for message in messages) + + async def test_chat_middleware_persists_compaction_summary_when_model_call_fails( chat_client_base: SupportsChatGetResponse, ) -> None: diff --git a/python/packages/core/tests/core/test_compaction.py b/python/packages/core/tests/core/test_compaction.py index dd2547e9a7c..b0e82ae94fc 100644 --- a/python/packages/core/tests/core/test_compaction.py +++ b/python/packages/core/tests/core/test_compaction.py @@ -1169,10 +1169,11 @@ async def test_apply_compaction_logs_changed_context_without_content(caplog: Any assert record.compaction_included_tokens_before is record.compaction_included_tokens_after is None caplog.clear() - await apply_compaction( - [Message(role="user", contents=["request"])], - strategy=TruncationStrategy(max_n=2, compact_to=1), - ) + with caplog.at_level(logging.INFO, logger="agent_framework"): + await apply_compaction( + [Message(role="user", contents=["request"])], + strategy=TruncationStrategy(max_n=2, compact_to=1), + ) assert caplog.messages == [] @@ -1914,7 +1915,7 @@ async def test_context_window_strategy_truncation_triggers_above_80_pct() -> Non assert len(projected) < 5 -async def test_context_window_strategy_can_preserve_first_user_group() -> None: +async def test_context_window_strategy_can_preserve_first_user_group(caplog: Any) -> None: messages = [ Message(role="user", contents=["original " * 400]), Message(role="assistant", contents=["old answer " * 400]), @@ -1926,12 +1927,15 @@ async def test_context_window_strategy_can_preserve_first_user_group() -> None: preserve_first_user_group=True, ) - changed = await strategy(messages) + with caplog.at_level(logging.WARNING, logger="agent_framework"): + changed = await strategy(messages) assert changed is True projected = included_messages(messages) assert any(message.text == "original " * 400 for message in projected) assert any(message.text == "latest " * 400 for message in projected) + warning = next(record for record in caplog.records if record.levelno == logging.WARNING) + assert warning.compaction_included_tokens_after > warning.compaction_input_budget_tokens async def test_context_window_strategy_keep_last_tool_call_groups_respected() -> None: