diff --git a/python/packages/core/agent_framework/_compaction.py b/python/packages/core/agent_framework/_compaction.py index e8b46a008a..43c6c007f8 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__( @@ -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 @@ -1632,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 @@ -1663,6 +1725,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 +1744,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 +1771,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 +1791,28 @@ 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 + 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 __all__ = [ diff --git a/python/packages/core/agent_framework/_middleware.py b/python/packages/core/agent_framework/_middleware.py index 45a3dcbf82..562d9d7df8 100644 --- a/python/packages/core/agent_framework/_middleware.py +++ b/python/packages/core/agent_framework/_middleware.py @@ -72,6 +72,85 @@ 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} + candidates: list[tuple[Message, set[str]]] = [] + 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) + 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 + annotation = cast("dict[str, Any]", annotation_value) + if annotation.get(SUMMARIZED_BY_SUMMARY_ID_KEY) not in pending_ids: + continue + 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): """Control-flow exception to terminate middleware execution early.""" @@ -1374,12 +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: - return await pipeline.execute( - context=context, - final_handler=self._middleware_handler, - ) + try: + return await pipeline.execute( + context=context, + final_handler=self._middleware_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 @@ -1408,14 +1497,19 @@ def _middleware_handler( 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, + 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=working_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, + ), ) diff --git a/python/packages/core/tests/core/test_clients.py b/python/packages/core/tests/core/test_clients.py index b57d5eba48..75251fc6e4 100644 --- a/python/packages/core/tests/core/test_clients.py +++ b/python/packages/core/tests/core/test_clients.py @@ -10,18 +10,37 @@ GROUP_ANNOTATION_KEY, GROUP_TOKEN_COUNT_KEY, BaseChatClient, + ChatMiddleware, ChatResponse, ChatResponseUpdate, Content, Message, SlidingWindowStrategy, + SummarizationStrategy, SupportsChatGetResponse, ToolResultCompactionStrategy, TruncationStrategy, + apply_compaction, tool, ) +class _NoOpChatMiddleware(ChatMiddleware): + 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"])]) + + class _FixedTokenizer: def __init__(self, token_count: int) -> None: self.token_count = token_count @@ -293,8 +312,14 @@ def _is_tool_result_summary(message: Message) -> bool: return message.role == "assistant" and text.startswith("[Tool results:") +@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, + 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 @@ -303,6 +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 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: @@ -361,14 +388,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 +442,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 +453,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 +497,124 @@ 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) + + +@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: + 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 3582985997..b0e82ae94f 100644 --- a/python/packages/core/tests/core/test_compaction.py +++ b/python/packages/core/tests/core/test_compaction.py @@ -1149,6 +1149,34 @@ 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() + 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 == [] + + # --- ToolResultCompactionStrategy tests --- @@ -1634,7 +1662,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 +1680,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 +1689,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 +1871,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 +1915,29 @@ 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(caplog: Any) -> 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, + ) + + 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: """The keep_last_tool_call_groups parameter controls how many groups are retained.""" # Create enough tokens to trigger tool eviction (>50% of input budget)