From 1c1c8d489b4d9b67998a10002e095756b04bfb59 Mon Sep 17 00:00:00 2001 From: eavanvalkenburg Date: Wed, 14 Oct 2026 16:27:31 +0200 Subject: [PATCH 1/8] Python: Add agent-triggered compaction --- python/packages/core/AGENTS.md | 13 + .../packages/core/agent_framework/__init__.py | 4 + .../core/agent_framework/__init__.pyi | 4 + .../core/agent_framework/_compaction.py | 365 ++++++++++++++++++ .../packages/core/tests/core/test_agents.py | 50 +++ .../core/tests/core/test_compaction.py | 292 ++++++++++++++ python/samples/02-agents/compaction/README.md | 4 + .../compaction/agent_triggered_compaction.py | 167 ++++++++ .../agent_triggered_compaction_basic.py | 140 +++++++ 9 files changed, 1039 insertions(+) create mode 100644 python/samples/02-agents/compaction/agent_triggered_compaction.py create mode 100644 python/samples/02-agents/compaction/agent_triggered_compaction_basic.py diff --git a/python/packages/core/AGENTS.md b/python/packages/core/AGENTS.md index 895764eb93d..0d7ab4e5763 100644 --- a/python/packages/core/AGENTS.md +++ b/python/packages/core/AGENTS.md @@ -165,6 +165,19 @@ The vector store API is experimental under the shared `VECTOR_STORES` feature ID Shared serialization normalizes stored data; codecs and connector overrides remain trusted Python code - **`SupportsVectorUpsert`** / **`SupportsVectorSearch`** - Structural protocols for vector store capabilities +### Compaction (`_compaction.py`) + +- **`AgentTriggeredCompactionStrategy`** - Wraps one existing compaction strategy and invokes it only after its + explicitly registered trigger tool completes in the latest message group. Trigger consumption is stored on that + group rather than on the strategy instance so shared agents remain safe across concurrent runs. The trigger group + is temporarily excluded while the inner strategy runs, so it does not consume retention counts, then restored by + default; set `exclude_trigger=True` to remove a trigger-only group from the next model call. Mixed groups that + contain unrelated tool calls are always restored atomically. +- **`UnloadSkillsCompactionStrategy`** - Intended as an inner agent-triggered strategy. It supplies an + `unload_skill` trigger-tool factory that `AgentTriggeredCompactionStrategy` adopts automatically, then removes + complete `load_skill` and `read_skill_resource` groups for one named skill, or all skills when no name is + supplied. Mixed tool groups and `run_skill_script` results remain intact. + ### Middleware (`_middleware.py`) - **`AgentMiddleware`** - Intercepts agent `run()` calls diff --git a/python/packages/core/agent_framework/__init__.py b/python/packages/core/agent_framework/__init__.py index 5a0b6800ab5..3895cc5a51f 100644 --- a/python/packages/core/agent_framework/__init__.py +++ b/python/packages/core/agent_framework/__init__.py @@ -69,6 +69,7 @@ "SUMMARIZED_BY_SUMMARY_ID_KEY", "SUMMARY_OF_GROUP_IDS_KEY", "SUMMARY_OF_MESSAGE_IDS_KEY", + "AgentTriggeredCompactionStrategy", "CharacterEstimatorTokenizer", "CompactionProvider", "CompactionStrategy", @@ -80,6 +81,7 @@ "TokenizerProtocol", "ToolResultCompactionStrategy", "TruncationStrategy", + "UnloadSkillsCompactionStrategy", "annotate_message_groups", "apply_compaction", "included_messages", @@ -452,6 +454,7 @@ "AgentResponseUpdate", "AgentRunInputs", "AgentSession", + "AgentTriggeredCompactionStrategy", "AggregatingSkillsSource", "Annotation", "BackgroundAgentsProvider", @@ -645,6 +648,7 @@ "ToolTypes", "TruncationStrategy", "TypeCompatibilityError", + "UnloadSkillsCompactionStrategy", "UpdateT", "UsageDetails", "UserInputRequiredException", diff --git a/python/packages/core/agent_framework/__init__.pyi b/python/packages/core/agent_framework/__init__.pyi index c5d5b0a250e..3eac3908a17 100644 --- a/python/packages/core/agent_framework/__init__.pyi +++ b/python/packages/core/agent_framework/__init__.pyi @@ -31,6 +31,7 @@ from ._compaction import ( SUMMARIZED_BY_SUMMARY_ID_KEY, SUMMARY_OF_GROUP_IDS_KEY, SUMMARY_OF_MESSAGE_IDS_KEY, + AgentTriggeredCompactionStrategy, CharacterEstimatorTokenizer, CompactionProvider, CompactionStrategy, @@ -42,6 +43,7 @@ from ._compaction import ( TokenizerProtocol, ToolResultCompactionStrategy, TruncationStrategy, + UnloadSkillsCompactionStrategy, annotate_message_groups, apply_compaction, included_messages, @@ -407,6 +409,7 @@ __all__ = [ "AgentResponseUpdate", "AgentRunInputs", "AgentSession", + "AgentTriggeredCompactionStrategy", "AggregatingSkillsSource", "Annotation", "BackgroundAgentsProvider", @@ -600,6 +603,7 @@ __all__ = [ "ToolTypes", "TruncationStrategy", "TypeCompatibilityError", + "UnloadSkillsCompactionStrategy", "UpdateT", "UsageDetails", "UserInputRequiredException", diff --git a/python/packages/core/agent_framework/_compaction.py b/python/packages/core/agent_framework/_compaction.py index 6f4078eca35..141637f615c 100644 --- a/python/packages/core/agent_framework/_compaction.py +++ b/python/packages/core/agent_framework/_compaction.py @@ -22,6 +22,7 @@ if TYPE_CHECKING: from ._clients import SupportsChatGetResponse + from ._tools import ApprovalMode, FunctionTool GroupKind: TypeAlias = Literal["system", "user", "assistant_text", "tool_call"] GROUP_ANNOTATION_KEY = "_group" @@ -41,6 +42,8 @@ SUMMARY_OF_MESSAGE_IDS_KEY = "_summary_of_message_ids" SUMMARY_OF_GROUP_IDS_KEY = "_summary_of_group_ids" SUMMARIZED_BY_SUMMARY_ID_KEY = "_summarized_by_summary_id" +_AGENT_COMPACTION_TRIGGER_CONSUMED_KEY = "_agent_compaction_trigger_consumed" +_SKILL_CONTEXT_TOOL_NAMES: Final[frozenset[str]] = frozenset({"load_skill", "read_skill_resource"}) logger = logging.getLogger("agent_framework") @@ -1180,6 +1183,366 @@ def _count_included_tokens(messages: list[Message]) -> int: return included_token_count(messages) +def _latest_included_group(messages: list[Message]) -> list[Message]: + ordered_group_ids = _ordered_group_ids_from_annotations(messages) + grouped = _group_messages_by_id(messages) + for group_id in reversed(ordered_group_ids): + group = grouped.get(group_id, []) + if any(not message.additional_properties.get(EXCLUDED_KEY, False) for message in group): + return group + return [] + + +def _completed_function_calls( + messages: Sequence[Message], + *, + tool_name: str, + include_consumed: bool = False, + include_excluded: bool = False, +) -> list[Content]: + candidate_messages = ( + list(messages) + if include_excluded + else [message for message in messages if not message.additional_properties.get(EXCLUDED_KEY, False)] + ) + successful_result_call_ids = { + content.call_id + for message in candidate_messages + for content in message.contents + if content.type == "function_result" and content.call_id and content.exception is None + } + return [ + content + for message in candidate_messages + for content in message.contents + if ( + content.type == "function_call" + and content.name == tool_name + and content.call_id in successful_result_call_ids + and ( + include_consumed + or message.additional_properties.get(_AGENT_COMPACTION_TRIGGER_CONSUMED_KEY) is not True + ) + ) + ] + + +class AgentTriggeredCompactionStrategy: + """Run an existing compaction strategy after an agent invokes a trigger tool. + + The trigger tool must be registered explicitly on the agent alongside this + strategy. A completed trigger call in the latest included message group is + consumed once, then the wrapped strategy runs before the next model call. + Wrapped strategies may provide ``create_trigger_tool()`` to replace the + default ``compact_context`` tool with one specialized for their inputs. + + The trigger is stored on the function-call content rather than on this + strategy instance, so one strategy can be shared safely by concurrent runs. + """ + + DEFAULT_TOOL_NAME: Final[str] = "compact_context" + + def __init__( + self, + strategy: CompactionStrategy, + *, + tool_name: str = DEFAULT_TOOL_NAME, + approval_mode: ApprovalMode = "never_require", + exclude_trigger: bool = False, + ) -> None: + """Create an agent-triggered wrapper around one compaction strategy. + + Args: + strategy: Existing strategy to invoke after a completed trigger call. + + Keyword Args: + tool_name: Name for the default trigger tool when the wrapped + strategy does not provide ``create_trigger_tool()``. + approval_mode: Approval mode for the generated default trigger tool. + exclude_trigger: Whether to exclude the completed trigger call/result + group before the next model call. Defaults to False so the model + can observe that compaction completed. A trigger group containing + unrelated tool calls is always retained atomically. + """ + if not callable(strategy): + raise TypeError("strategy must be callable.") + if not tool_name.strip(): + raise ValueError("tool_name must not be empty.") + + trigger_tool_factory = getattr(strategy, "create_trigger_tool", None) + tool = ( + trigger_tool_factory() + if callable(trigger_tool_factory) + else self._create_default_tool(tool_name, approval_mode) + ) + from ._tools import FunctionTool + + if not isinstance(tool, FunctionTool): + raise TypeError("create_trigger_tool() must return a FunctionTool.") + if not tool.name.strip(): + raise ValueError("trigger tool name must not be empty.") + + self.strategy = strategy + self.tool = tool + self.exclude_trigger = exclude_trigger + + @property + def tokenizer(self) -> TokenizerProtocol | None: + """Return the wrapped strategy's tokenizer, when it exposes one.""" + tokenizer = getattr(self.strategy, "tokenizer", None) + return tokenizer if isinstance(tokenizer, TokenizerProtocol) else None + + @staticmethod + def _create_default_tool(tool_name: str, approval_mode: ApprovalMode) -> FunctionTool: + from ._tools import FunctionTool + + def request_compaction(reason: str | None = None) -> str: + """Request configured context compaction before the next model call.""" + message = "Compaction requested. The configured strategy will run before the next model call." + normalized_reason = reason.strip() if isinstance(reason, str) else "" + return f"{message} Reason: {normalized_reason}" if normalized_reason else message + + return FunctionTool( + name=tool_name, + description=( + "Request that the configured context compaction strategy run before the next model call. " + "Call this tool by itself, after any other tool calls whose results should be eligible for compaction; " + "do not call it in parallel with other tools." + ), + func=request_compaction, + approval_mode=approval_mode, + input_model={ + "type": "object", + "properties": { + "reason": { + "type": "string", + "description": "Why the agent is requesting compaction before the next model call.", + }, + }, + "additionalProperties": False, + }, + ) + + async def __call__(self, messages: list[Message]) -> bool: + """Run the wrapped strategy once for the latest completed trigger call.""" + annotate_message_groups(messages) + trigger_group = _latest_included_group(messages) + trigger_calls = _completed_function_calls( + trigger_group, + tool_name=self.tool.name, + ) + if not trigger_calls: + return False + + trigger_state = [ + ( + message, + EXCLUDED_KEY in message.additional_properties, + message.additional_properties.get(EXCLUDED_KEY), + EXCLUDE_REASON_KEY in message.additional_properties, + message.additional_properties.get(EXCLUDE_REASON_KEY), + ) + for message in trigger_group + ] + for message in trigger_group: + set_excluded(message, excluded=True, reason="agent_compaction_trigger") + + try: + changed = await self.strategy(messages) + except BaseException: + self._restore_trigger_group(trigger_state) + raise + + for message in trigger_group: + message.additional_properties[_AGENT_COMPACTION_TRIGGER_CONSUMED_KEY] = True + + has_unrelated_tool_calls = any( + content.type in _TOOL_CALL_CONTENT_TYPES + and not (content.type == "function_call" and content.name == self.tool.name) + for message in trigger_group + for content in message.contents + ) + trigger_message_ids = {id(message) for message in trigger_group} + has_other_included_context = any( + id(message) not in trigger_message_ids and not message.additional_properties.get(EXCLUDED_KEY, False) + for message in messages + ) + trigger_changed = self.exclude_trigger and not has_unrelated_tool_calls and has_other_included_context + if not trigger_changed: + self._restore_trigger_group(trigger_state) + return changed or trigger_changed + + @staticmethod + def _restore_trigger_group( + trigger_state: Sequence[tuple[Message, bool, Any, bool, Any]], + ) -> None: + for message, had_excluded, excluded, had_reason, reason in trigger_state: + if had_excluded: + message.additional_properties[EXCLUDED_KEY] = excluded + else: + message.additional_properties.pop(EXCLUDED_KEY, None) + if had_reason: + message.additional_properties[EXCLUDE_REASON_KEY] = reason + else: + message.additional_properties.pop(EXCLUDE_REASON_KEY, None) + + +class UnloadSkillsCompactionStrategy: + """Exclude loaded skill bodies and resources selected by an unload tool. + + This strategy targets complete tool-call groups containing only + ``load_skill`` and ``read_skill_resource`` calls. The latest completed + ``unload_skill`` request selects one skill by name, or all skills when its + ``skill_name`` argument is omitted or blank. + + Use this as the wrapped strategy of :class:`AgentTriggeredCompactionStrategy`. + The wrapper obtains the specialized ``unload_skill`` tool through + :meth:`create_trigger_tool` and exposes that single tool to the caller. + """ + + DEFAULT_TOOL_NAME: Final[str] = "unload_skill" + + def __init__( + self, + *, + tool_name: str = DEFAULT_TOOL_NAME, + approval_mode: ApprovalMode = "never_require", + ) -> None: + """Create a skill-unloading strategy and its trigger tool. + + Keyword Args: + tool_name: Name of the tool that requests skill unloading. + approval_mode: Approval mode for the generated unload tool. + """ + if not tool_name.strip(): + raise ValueError("tool_name must not be empty.") + self.tool_name: str = tool_name + self.approval_mode: ApprovalMode = approval_mode + + def create_trigger_tool(self) -> FunctionTool: + """Create the ``unload_skill`` trigger tool used by the wrapper.""" + from ._tools import FunctionTool + + def request_skill_unload(skill_name: str | Sequence[str] | None = None) -> str: + """Unload the full instructions and resources for selected skills.""" + requested_names = ( + [skill_name] if isinstance(skill_name, str) else skill_name if isinstance(skill_name, Sequence) else [] + ) + normalized_names = [name.strip() for name in requested_names if name.strip()] + if normalized_names: + formatted_names = ", ".join(f"'{name}'" for name in normalized_names) + return f"Skill(s) {formatted_names} will be unloaded before the next model call." + return "All loaded skills will be unloaded before the next model call." + + return FunctionTool( + name=self.tool_name, + description=( + "Unloads the full instructions and resources for one or more loaded skills. " + "Omit skill_name to unload all loaded skills. Call this tool by itself after the relevant skill work " + "is complete; do not call it in parallel with other tools." + ), + func=request_skill_unload, + approval_mode=self.approval_mode, + input_model={ + "type": "object", + "properties": { + "skill_name": { + "oneOf": [ + {"type": "string"}, + {"type": "array", "items": {"type": "string"}}, + ], + "description": ( + "The name of a skill to unload, a list of skill names to unload, " + "or omit to unload all loaded skills." + ), + }, + }, + "additionalProperties": False, + }, + ) + + def _requested_skill_names(self, messages: list[Message]) -> set[str] | None: + ordered_group_ids = _ordered_group_ids_from_annotations(messages) + grouped = _group_messages_by_id(messages) + latest_group = grouped.get(ordered_group_ids[-1], []) if ordered_group_ids else [] + requests = _completed_function_calls( + latest_group, + tool_name=self.tool_name, + include_consumed=True, + include_excluded=True, + ) + if not requests: + return set() + + requested_names: set[str] = set() + for request in requests: + arguments = request.parse_arguments() or {} + raw_skill_names = arguments.get("skill_name") + skill_names = ( + [raw_skill_names] + if isinstance(raw_skill_names, str) + else cast("Sequence[Any]", raw_skill_names) + if isinstance(raw_skill_names, Sequence) + else [] + ) + normalized_names = { + skill_name.strip().casefold() + for skill_name in skill_names + if isinstance(skill_name, str) and skill_name.strip() + } + if not normalized_names: + return None + requested_names.update(normalized_names) + return requested_names + + @staticmethod + def _skill_names_for_removable_group(group: Sequence[Message]) -> set[str] | None: + included_messages = [message for message in group if not message.additional_properties.get(EXCLUDED_KEY, False)] + calls = [ + content for message in included_messages for content in message.contents if content.type == "function_call" + ] + if not calls or any(call.name not in _SKILL_CONTEXT_TOOL_NAMES for call in calls): + return None + + completed_call_ids = { + content.call_id + for message in included_messages + for content in message.contents + if content.type == "function_result" and content.call_id + } + if any(not call.call_id or call.call_id not in completed_call_ids for call in calls): + return None + + skill_names: set[str] = set() + for call in calls: + arguments = call.parse_arguments() or {} + skill_name = arguments.get("skill_name") + if not isinstance(skill_name, str) or not skill_name.strip(): + return None + skill_names.add(skill_name.strip().casefold()) + return skill_names + + async def __call__(self, messages: list[Message]) -> bool: + """Exclude matching loaded skill context groups.""" + ordered_group_ids = annotate_message_groups(messages) + requested_skill_names = self._requested_skill_names(messages) + if requested_skill_names == set(): + return False + + grouped = _group_messages_by_id(messages) + changed = False + for group_id in ordered_group_ids: + group = grouped.get(group_id, []) + skill_names = self._skill_names_for_removable_group(group) + if skill_names is None: + continue + if requested_skill_names is not None and not skill_names.issubset(requested_skill_names): + continue + for message in group: + changed = set_excluded(message, excluded=True, reason="skill_unload") or changed + return changed + + class TruncationStrategy: """Oldest-first compaction using a single metric threshold. @@ -2369,6 +2732,7 @@ async def __call__(self, messages: list[Message]) -> bool: "SUMMARIZED_BY_SUMMARY_ID_KEY", "SUMMARY_OF_GROUP_IDS_KEY", "SUMMARY_OF_MESSAGE_IDS_KEY", + "AgentTriggeredCompactionStrategy", "CharacterEstimatorTokenizer", "CompactionProvider", "CompactionStrategy", @@ -2381,6 +2745,7 @@ async def __call__(self, messages: list[Message]) -> bool: "TokenizerProtocol", "ToolResultCompactionStrategy", "TruncationStrategy", + "UnloadSkillsCompactionStrategy", "annotate_message_groups", "annotate_token_counts", "append_compaction_message", diff --git a/python/packages/core/tests/core/test_agents.py b/python/packages/core/tests/core/test_agents.py index 1481ce8a77b..df83fb923e1 100644 --- a/python/packages/core/tests/core/test_agents.py +++ b/python/packages/core/tests/core/test_agents.py @@ -21,6 +21,7 @@ AgentResponse, AgentResponseUpdate, AgentSession, + AgentTriggeredCompactionStrategy, ChatContext, ChatOptions, ChatResponse, @@ -3446,6 +3447,55 @@ async def capturing_inner( assert captured_token_counts == [[9, 9]] +@pytest.mark.parametrize("stream", [False, True]) +async def test_agent_triggered_compaction_applies_before_next_model_call( + chat_client_base: MockBaseChatClient, + stream: bool, +) -> None: + agent_compaction = AgentTriggeredCompactionStrategy( + TruncationStrategy(max_n=2, compact_to=2, preserve_first_user_group=True), + ) + function_call = Content.from_function_call( + call_id="compact-1", + name=agent_compaction.tool.name, + arguments={"reason": "planning complete"}, + ) + if stream: + chat_client_base.streaming_responses = [ + [ChatResponseUpdate(contents=[function_call], role="assistant", finish_reason="tool_calls")], + [ChatResponseUpdate(contents=[Content.from_text("done")], role="assistant", finish_reason="stop")], + ] + else: + chat_client_base.run_responses = [ + ChatResponse(messages=Message(role="assistant", contents=[function_call])), + ChatResponse(messages=Message(role="assistant", contents=["done"])), + ] + + captured_roles: list[list[str]] = [] + original_inner = chat_client_base._inner_get_response + + def capture_inner( + *, messages: MutableSequence[Message], stream: bool, options: dict[str, Any], **kwargs: Any + ) -> Any: + captured_roles.append([message.role for message in messages]) + return original_inner(messages=messages, stream=stream, options=options, **kwargs) + + agent = Agent( + client=chat_client_base, + tools=[agent_compaction.tool], + compaction_strategy=agent_compaction, + ) + + with patch.object(chat_client_base, "_inner_get_response", side_effect=capture_inner): + if stream: + result_text = "".join([update.text or "" async for update in agent.run("plan", stream=True)]) + else: + result_text = (await agent.run("plan")).text + + assert result_text == "done" + assert captured_roles == [["user"], ["user", "assistant", "tool"]] + + async def test_chat_agent_uses_client_compaction_defaults_when_agent_unset(chat_client_base: Any) -> None: captured_roles: list[list[str]] = [] original_inner = chat_client_base._inner_get_response diff --git a/python/packages/core/tests/core/test_compaction.py b/python/packages/core/tests/core/test_compaction.py index 7c5ed25de54..660cc29446a 100644 --- a/python/packages/core/tests/core/test_compaction.py +++ b/python/packages/core/tests/core/test_compaction.py @@ -19,6 +19,7 @@ SUMMARIZED_BY_SUMMARY_ID_KEY, SUMMARY_OF_GROUP_IDS_KEY, SUMMARY_OF_MESSAGE_IDS_KEY, + AgentTriggeredCompactionStrategy, CharacterEstimatorTokenizer, ChatResponse, CompactionProvider, @@ -31,6 +32,7 @@ TokenBudgetComposedStrategy, ToolResultCompactionStrategy, TruncationStrategy, + UnloadSkillsCompactionStrategy, annotate_message_groups, apply_compaction, included_messages, @@ -52,6 +54,17 @@ def _assistant_function_call(call_id: str) -> Message: ) +def _assistant_named_function_call( + call_id: str, + name: str, + arguments: dict[str, Any] | None = None, +) -> Message: + return Message( + role="assistant", + contents=[Content.from_function_call(call_id=call_id, name=name, arguments=arguments or {})], + ) + + def _assistant_function_call_with_result(call_id: str, result: str) -> Message: return Message( role="assistant", @@ -848,6 +861,285 @@ def test_truncation_strategy_validates_token_targets() -> None: raise AssertionError("Expected ValueError when compact_to is greater than max_n.") +class _RecordingCompactionStrategy: + def __init__(self, *, changed: bool = True) -> None: + self.call_count = 0 + self.changed = changed + + async def __call__(self, messages: list[Message]) -> bool: + self.call_count += 1 + return self.changed + + +async def test_agent_triggered_compaction_runs_only_after_completed_trigger() -> None: + inner = _RecordingCompactionStrategy() + strategy = AgentTriggeredCompactionStrategy(inner) + messages = [ + Message(role="user", contents=["plan"]), + _assistant_named_function_call("compact-1", strategy.tool.name, {"reason": "planning complete"}), + ] + + assert await strategy(messages) is False + assert inner.call_count == 0 + + messages.append(_tool_result("compact-1", "Compaction requested.")) + + assert await strategy(messages) is True + assert inner.call_count == 1 + assert messages[1].additional_properties.get(EXCLUDED_KEY) is not True + assert messages[2].additional_properties.get(EXCLUDED_KEY) is not True + assert await strategy(messages) is False + assert inner.call_count == 1 + + +async def test_agent_triggered_compaction_consumes_request_when_inner_strategy_is_noop() -> None: + inner = _RecordingCompactionStrategy(changed=False) + strategy = AgentTriggeredCompactionStrategy(inner) + messages = [ + _assistant_named_function_call("compact-1", strategy.tool.name), + _tool_result("compact-1", "Compaction requested."), + ] + + assert await strategy(messages) is False + assert await strategy(messages) is False + assert inner.call_count == 1 + + +async def test_agent_triggered_compaction_can_exclude_trigger_group() -> None: + inner = _RecordingCompactionStrategy(changed=False) + strategy = AgentTriggeredCompactionStrategy(inner, exclude_trigger=True) + messages = [ + Message(role="user", contents=["keep working"]), + _assistant_named_function_call("compact-1", strategy.tool.name), + _tool_result("compact-1", "Compaction requested."), + ] + + assert await strategy(messages) is True + assert messages[0].additional_properties.get(EXCLUDED_KEY) is not True + assert messages[1].additional_properties.get(EXCLUDED_KEY) is True + assert messages[2].additional_properties.get(EXCLUDED_KEY) is True + + +async def test_agent_trigger_does_not_count_toward_tool_group_retention() -> None: + strategy = AgentTriggeredCompactionStrategy(SelectiveToolCallCompactionStrategy(keep_last_tool_call_groups=1)) + messages = [ + Message(role="user", contents=["work"]), + _assistant_named_function_call("work-1", "work_tool"), + _tool_result("work-1", "work result"), + _assistant_named_function_call("compact-1", strategy.tool.name), + _tool_result("compact-1", "Compaction requested."), + ] + + assert await strategy(messages) is False + assert all(message.additional_properties.get(EXCLUDED_KEY) is not True for message in messages) + + +@pytest.mark.parametrize("trigger_position", [0, 1, 2]) +@pytest.mark.parametrize("exclude_trigger", [False, True]) +async def test_agent_trigger_in_mixed_tool_batch_is_position_independent_and_retained( + trigger_position: int, + exclude_trigger: bool, +) -> None: + strategy = AgentTriggeredCompactionStrategy( + SelectiveToolCallCompactionStrategy(keep_last_tool_call_groups=0), + exclude_trigger=exclude_trigger, + ) + batch_calls = [ + Content.from_function_call(call_id="work-a", name="work_a", arguments={}), + Content.from_function_call(call_id="work-b", name="work_b", arguments={}), + ] + batch_calls.insert( + trigger_position, + Content.from_function_call(call_id="compact-1", name=strategy.tool.name, arguments={}), + ) + messages = [ + Message(role="user", contents=["work"]), + _assistant_named_function_call("old-work", "old_work"), + _tool_result("old-work", "old result"), + Message(role="assistant", contents=batch_calls), + Message( + role="tool", + contents=[ + Content.from_function_result(call_id=call.call_id or "", result=f"{call.name} result") + for call in batch_calls + ], + ), + ] + + changed = await strategy(messages) + + assert changed is True + assert messages[1].additional_properties.get(EXCLUDED_KEY) is True + assert messages[2].additional_properties.get(EXCLUDED_KEY) is True + assert messages[3].additional_properties.get(EXCLUDED_KEY) is not True + assert messages[4].additional_properties.get(EXCLUDED_KEY) is not True + + +async def test_agent_triggered_compaction_restores_request_when_inner_strategy_fails() -> None: + class _FailingStrategy: + async def __call__(self, messages: list[Message]) -> bool: + raise RuntimeError("compaction failed") + + strategy = AgentTriggeredCompactionStrategy(_FailingStrategy()) + trigger_message = _assistant_named_function_call("compact-1", strategy.tool.name) + messages = [trigger_message, _tool_result("compact-1", "Compaction requested.")] + + with pytest.raises(RuntimeError, match="compaction failed"): + await strategy(messages) + + assert "_agent_compaction_trigger_consumed" not in trigger_message.additional_properties + + +async def test_agent_triggered_compaction_ignores_trigger_outside_latest_group() -> None: + inner = _RecordingCompactionStrategy() + strategy = AgentTriggeredCompactionStrategy(inner) + messages = [ + _assistant_named_function_call("compact-1", strategy.tool.name), + _tool_result("compact-1", "Compaction requested."), + Message(role="assistant", contents=["continued"]), + ] + + assert await strategy(messages) is False + assert inner.call_count == 0 + + +def test_agent_triggered_compaction_tool_describes_reason_parameter() -> None: + strategy = AgentTriggeredCompactionStrategy(_RecordingCompactionStrategy()) + + reason_schema = strategy.tool.parameters()["properties"]["reason"] + + assert reason_schema["description"] == "Why the agent is requesting compaction before the next model call." + assert "Call this tool by itself" in strategy.tool.description + assert "do not call it in parallel" in strategy.tool.description + + +async def test_unload_skills_compaction_removes_named_skill_only() -> None: + skill_strategy = UnloadSkillsCompactionStrategy() + strategy = AgentTriggeredCompactionStrategy(skill_strategy) + messages = [ + Message(role="user", contents=["work"]), + _assistant_named_function_call("load-pdf", "load_skill", {"skill_name": "pdf"}), + _tool_result("load-pdf", "pdf instructions"), + _assistant_named_function_call("load-excel", "load_skill", {"skill_name": "excel"}), + _tool_result("load-excel", "excel instructions"), + _assistant_named_function_call("unload-pdf", strategy.tool.name, {"skill_name": "pdf"}), + _tool_result("unload-pdf", "Unload requested for skill 'pdf'."), + ] + + changed = await strategy(messages) + + assert changed is True + assert messages[1].additional_properties.get(EXCLUDED_KEY) is True + assert messages[2].additional_properties.get(EXCLUDED_KEY) is True + assert messages[3].additional_properties.get(EXCLUDED_KEY) is not True + assert messages[4].additional_properties.get(EXCLUDED_KEY) is not True + assert messages[5].additional_properties.get(EXCLUDED_KEY) is not True + assert messages[6].additional_properties.get(EXCLUDED_KEY) is not True + + +async def test_unload_skills_compaction_removes_multiple_named_skills() -> None: + strategy = AgentTriggeredCompactionStrategy(UnloadSkillsCompactionStrategy()) + messages = [ + _assistant_named_function_call("load-pdf", "load_skill", {"skill_name": "pdf"}), + _tool_result("load-pdf", "pdf instructions"), + _assistant_named_function_call("load-excel", "load_skill", {"skill_name": "excel"}), + _tool_result("load-excel", "excel instructions"), + _assistant_named_function_call("load-word", "load_skill", {"skill_name": "word"}), + _tool_result("load-word", "word instructions"), + _assistant_named_function_call( + "unload-selected", + strategy.tool.name, + {"skill_name": ["pdf", "excel"]}, + ), + _tool_result("unload-selected", "Unload requested for pdf and excel."), + ] + + changed = await strategy(messages) + + assert changed is True + assert all(message.additional_properties.get(EXCLUDED_KEY) is True for message in messages[:4]) + assert all(message.additional_properties.get(EXCLUDED_KEY) is not True for message in messages[4:6]) + assert all(message.additional_properties.get(EXCLUDED_KEY) is not True for message in messages[6:]) + + +def test_unload_skills_tool_describes_skill_name_parameter() -> None: + strategy = AgentTriggeredCompactionStrategy(UnloadSkillsCompactionStrategy()) + + skill_name_schema = strategy.tool.parameters()["properties"]["skill_name"] + + assert skill_name_schema["description"] == ( + "The name of a skill to unload, a list of skill names to unload, or omit to unload all loaded skills." + ) + assert "Call this tool by itself" in strategy.tool.description + assert "do not call it in parallel" in strategy.tool.description + + +async def test_unload_skills_compaction_empty_name_removes_all_skill_context() -> None: + skill_strategy = UnloadSkillsCompactionStrategy() + strategy = AgentTriggeredCompactionStrategy(skill_strategy) + messages = [ + _assistant_named_function_call("load-pdf", "load_skill", {"skill_name": "pdf"}), + _tool_result("load-pdf", "pdf instructions"), + _assistant_named_function_call( + "resource-excel", + "read_skill_resource", + {"skill_name": "excel", "resource_name": "reference"}, + ), + _tool_result("resource-excel", "excel reference"), + _assistant_named_function_call("unload-all", strategy.tool.name, {"skill_name": " "}), + _tool_result("unload-all", "Unload requested for all skills."), + ] + + changed = await strategy(messages) + + assert changed is True + assert all(message.additional_properties.get(EXCLUDED_KEY) is True for message in messages[:4]) + assert all(message.additional_properties.get(EXCLUDED_KEY) is not True for message in messages[4:]) + + +async def test_unload_skills_compaction_preserves_script_results_and_mixed_groups() -> None: + skill_strategy = UnloadSkillsCompactionStrategy() + strategy = AgentTriggeredCompactionStrategy(skill_strategy) + mixed_call = Message( + role="assistant", + contents=[ + Content.from_function_call( + call_id="load-pdf", + name="load_skill", + arguments={"skill_name": "pdf"}, + ), + Content.from_function_call( + call_id="search", + name="search", + arguments={"query": "pdf"}, + ), + ], + ) + messages = [ + mixed_call, + Message( + role="tool", + contents=[ + Content.from_function_result(call_id="load-pdf", result="pdf instructions"), + Content.from_function_result(call_id="search", result="search result"), + ], + ), + _assistant_named_function_call( + "run-pdf", + "run_skill_script", + {"skill_name": "pdf", "script_name": "convert"}, + ), + _tool_result("run-pdf", "converted.pdf"), + _assistant_named_function_call("unload-pdf", strategy.tool.name, {"skill_name": "pdf"}), + _tool_result("unload-pdf", "Unload requested for skill 'pdf'."), + ] + + changed = await strategy(messages) + + assert changed is False + assert all(message.additional_properties.get(EXCLUDED_KEY) is not True for message in messages) + + async def test_selective_tool_call_strategy_excludes_older_tool_groups() -> None: messages = [ Message(role="user", contents=["u"]), diff --git a/python/samples/02-agents/compaction/README.md b/python/samples/02-agents/compaction/README.md index 779893eefeb..420e3791e2c 100644 --- a/python/samples/02-agents/compaction/README.md +++ b/python/samples/02-agents/compaction/README.md @@ -11,6 +11,8 @@ This folder demonstrates context compaction patterns introduced by ADR-0019. - `custom.py` — defines a custom strategy implementing the `CompactionStrategy` protocol. - `tiktoken_tokenizer.py` — shows a `TokenizerProtocol` implementation backed by `tiktoken`. - `compaction_provider.py` — uses `CompactionProvider` with an agent and `InMemoryHistoryProvider`. +- `agent_triggered_compaction.py` — lets an agent trigger in-run compaction to unload one skill's context. +- `agent_triggered_compaction_basic.py` — demonstrates the generic wrapper with an existing tool-call strategy. Run samples with: @@ -22,6 +24,8 @@ uv run samples/02-agents/compaction/agent_client_overrides.py uv run samples/02-agents/compaction/custom.py uv run samples/02-agents/compaction/tiktoken_tokenizer.py uv run samples/02-agents/compaction/compaction_provider.py # requires OPENAI_API_KEY +uv run samples/02-agents/compaction/agent_triggered_compaction.py # requires Foundry configuration +uv run samples/02-agents/compaction/agent_triggered_compaction_basic.py # requires Foundry configuration ``` ## Security Considerations diff --git a/python/samples/02-agents/compaction/agent_triggered_compaction.py b/python/samples/02-agents/compaction/agent_triggered_compaction.py new file mode 100644 index 00000000000..031aef502ff --- /dev/null +++ b/python/samples/02-agents/compaction/agent_triggered_compaction.py @@ -0,0 +1,167 @@ +# Copyright (c) Microsoft. All rights reserved. + +import asyncio +import os +from itertools import count +from textwrap import dedent + +from agent_framework import ( + Agent, + AgentTriggeredCompactionStrategy, + CompactionStrategy, + InlineSkill, + Message, + SkillFrontmatter, + SkillsProvider, + UnloadSkillsCompactionStrategy, + included_messages, +) +from agent_framework.foundry import FoundryChatClient +from azure.identity import AzureCliCredential +from dotenv import load_dotenv + +load_dotenv() + +"""This sample demonstrates agent-triggered in-run compaction. + +The agent follows three steps in one function-calling run: + +1. Call ``load_skill`` to load a planning skill. +2. Use the skill, then call ``unload_skill``. +3. Continue without the loaded skill body in model context. + +``AgentTriggeredCompactionStrategy`` controls when compaction runs. +``UnloadSkillsCompactionStrategy`` controls what context is removed. +The trigger tool is registered explicitly alongside the strategy. + +A small inspection wrapper prints every model input after compaction so the +skill body is visibly present after ``load_skill`` and absent after +``unload_skill``. + +Environment variables: + FOUNDRY_PROJECT_ENDPOINT — Microsoft Foundry project endpoint. + FOUNDRY_MODEL — Model deployment name. + +Run with: + uv run samples/02-agents/compaction/agent_triggered_compaction.py +""" + + +MODEL_CALLS = count(1) + + +class InspectingCompactionStrategy: + """Print the projected messages after applying another strategy.""" + + def __init__(self, strategy: CompactionStrategy) -> None: + self.strategy = strategy + + async def __call__(self, messages: list[Message]) -> bool: + changed = await self.strategy(messages) + visible_messages = included_messages(messages) + + print(f"\n--- Model call {next(MODEL_CALLS)} ---") + print(f"Visible messages: {len(visible_messages)}") + for message in visible_messages: + parts: list[str] = [] + for content in message.contents: + if content.type == "text": + text = (content.text or "").replace("\n", " ") + parts.append(text[:100]) + elif content.type == "function_call": + parts.append(f"{content.name}({content.arguments})") + elif content.type == "function_result": + result = str(content.result).replace("\n", " ") + parts.append(f"result={result[:100]}") + else: + parts.append(content.type) + print(f"[{message.role}] {' | '.join(parts)}") + return changed + + +planning_skill = InlineSkill( + frontmatter=SkillFrontmatter( + name="planning", + description="Create a concise implementation plan before execution.", + ), + instructions=dedent("""\ + # Planning Skill + + Before implementation: + 1. Identify the required behavior. + 2. List the smallest code changes. + 3. Define the targeted validation. + """), +) + + +async def main() -> None: + # 1. Wrap the skill strategy. The wrapper automatically adopts its + # unload_skill trigger and exposes that single tool. + triggered_compaction = AgentTriggeredCompactionStrategy( + UnloadSkillsCompactionStrategy(), + ) + + # 2. Use the standard Foundry chat client. Run `az login` before starting. + client = FoundryChatClient( + project_endpoint=os.environ["FOUNDRY_PROJECT_ENDPOINT"], + model=os.environ["FOUNDRY_MODEL"], + credential=AzureCliCredential(), + ) + + # 3. The trigger tool is registered explicitly. SkillsProvider contributes + # load_skill, while disabling approval keeps this sample automatic. + agent = Agent( + client=client, + name="PlanningAgent", + instructions=dedent("""\ + If the current context contains a completed unload_skill result, the planning + skill is already unloaded: do not call any tools, return the plan, and say + that execution can begin. + + Otherwise, follow this sequence exactly: + 1. Call load_skill with skill_name="planning". + 2. Use the loaded instructions to create a concise plan. + 3. Call unload_skill with skill_name="planning". + + Do not call load_skill and unload_skill in parallel. + """), + tools=[triggered_compaction.tool], + context_providers=[ + SkillsProvider( + planning_skill, + disable_load_skill_approval=True, + ) + ], + compaction_strategy=InspectingCompactionStrategy(triggered_compaction), + default_options={"store": False}, + ) + + # 4. The final model call printed by the middleware no longer contains the + # load_skill call or its large planning-skill result. + result = await agent.run("Plan a small implementation, unload the planning skill, then move to execution.") + print(f"\nAgent: {result.messages[-1].text}") + + +if __name__ == "__main__": + asyncio.run(main()) + +""" +Example output (model wording varies): + +--- Model call 1 --- +[user] Plan a small implementation, unload the planning skill, then move to execution. + +--- Model call 2 --- +[user] Plan a small implementation, unload the planning skill, then move to execution. +[assistant] load_skill({"skill_name":"planning"}) +[tool] result=planning Create a concise implementation plan... + +--- Model call 3 --- +Visible messages: 3 +[user] Plan a small implementation, unload the planning skill, then move to execution. +[assistant] unload_skill({"skill_name":"planning"}) +[tool] result=Skill(s) 'planning' will be unloaded before the next model call. + +Agent: [The generated plan, followed by confirmation that execution can begin.] +""" diff --git a/python/samples/02-agents/compaction/agent_triggered_compaction_basic.py b/python/samples/02-agents/compaction/agent_triggered_compaction_basic.py new file mode 100644 index 00000000000..8ed2e05b59f --- /dev/null +++ b/python/samples/02-agents/compaction/agent_triggered_compaction_basic.py @@ -0,0 +1,140 @@ +# Copyright (c) Microsoft. All rights reserved. + +import asyncio +import os +from itertools import count + +from agent_framework import ( + Agent, + AgentTriggeredCompactionStrategy, + CompactionStrategy, + Message, + ToolResultCompactionStrategy, + included_messages, + tool, +) +from agent_framework.foundry import FoundryChatClient +from azure.identity import AzureCliCredential +from dotenv import load_dotenv + +load_dotenv() + +"""This sample demonstrates generic agent-triggered compaction. + +The agent first calls ``collect_context`` to produce a deliberately large tool +result. It then calls ``compact_context``. Before the next model call, +``ToolResultCompactionStrategy`` replaces the older context-producing tool +group with a compact summary, and the wrapper removes the completed compaction +trigger group. + +Environment variables: + FOUNDRY_PROJECT_ENDPOINT — Microsoft Foundry project endpoint. + FOUNDRY_MODEL — Model deployment name. + +Run with: + uv run samples/02-agents/compaction/agent_triggered_compaction_basic.py +""" + + +MODEL_CALLS = count(1) + + +@tool(approval_mode="never_require") +def collect_context(topic: str) -> str: + """Collect intentionally verbose context for a topic.""" + return "\n".join( + f"Exploration note {index}: {topic} should remain simple, testable, and observable." for index in range(1, 21) + ) + + +class InspectingCompactionStrategy: + """Print the projected messages after applying another strategy.""" + + def __init__(self, strategy: CompactionStrategy) -> None: + self.strategy = strategy + + async def __call__(self, messages: list[Message]) -> bool: + changed = await self.strategy(messages) + visible_messages = included_messages(messages) + + print(f"\n--- Model call {next(MODEL_CALLS)} ---") + print(f"Visible messages: {len(visible_messages)}") + for message in visible_messages: + parts: list[str] = [] + for content in message.contents: + if content.type == "text": + text = (content.text or "").replace("\n", " ") + parts.append(text[:100]) + elif content.type == "function_call": + parts.append(f"{content.name}({content.arguments})") + elif content.type == "function_result": + result = str(content.result).replace("\n", " ") + parts.append(f"result={result[:100]}") + else: + parts.append(content.type) + print(f"[{message.role}] {' | '.join(parts)}") + return changed + + +async def main() -> None: + # 1. Wrap an existing strategy. The trigger does not count toward the + # keep-last setting, so 0 compacts every workload tool-call group. + # A summary remains, allowing this sample to remove the trigger too. + triggered_compaction = AgentTriggeredCompactionStrategy( + ToolResultCompactionStrategy(keep_last_tool_call_groups=0), + exclude_trigger=True, + ) + + # 2. Use Foundry with store=False so Agent Framework owns the full in-run + # message list. Run `az login` before starting. + client = FoundryChatClient( + project_endpoint=os.environ["FOUNDRY_PROJECT_ENDPOINT"], + model=os.environ["FOUNDRY_MODEL"], + credential=AzureCliCredential(), + ) + + # 3. Register both the ordinary context-producing tool and the explicit + # compaction trigger tool. + agent = Agent( + client=client, + name="ResearchAgent", + instructions=( + "If the current context contains an assistant message starting with '[Tool results:', " + "compaction is complete: do not call any tools and answer the user in one sentence. " + "Otherwise, call collect_context once and then call compact_context. " + "Do not call the tools in parallel." + ), + tools=[collect_context, triggered_compaction.tool], + compaction_strategy=InspectingCompactionStrategy(triggered_compaction), + default_options={"store": False}, + ) + + # 4. The third model input contains a compact tool-result summary, but + # neither of the original tool-call groups. + result = await agent.run("Identify the main implementation priorities.") + print(f"\nAgent: {result.messages[-1].text}") + + +if __name__ == "__main__": + asyncio.run(main()) + +""" +Example output (model wording varies): + +--- Model call 1 --- +Visible messages: 1 +[user] Identify the main implementation priorities. + +--- Model call 2 --- +Visible messages: 3 +[user] Identify the main implementation priorities. +[assistant] collect_context({"topic":"implementation priorities"}) +[tool] result=Exploration note 1: implementation priorities should remain simple, testable... + +--- Model call 3 --- +Visible messages: 2 +[user] Identify the main implementation priorities. +[assistant] [Tool results: collect_context: Exploration note 1: implementation priorities...] + +Agent: Keep the implementation simple, testable, and observable. +""" From 1d0d2c1337834e8bbc25d1252dea8cea9bee4550 Mon Sep 17 00:00:00 2001 From: eavanvalkenburg Date: Wed, 14 Oct 2026 17:33:15 +0200 Subject: [PATCH 2/8] Python: Address compaction review feedback --- .../core/agent_framework/_compaction.py | 84 ++++++++++++------- .../packages/core/agent_framework/_tools.py | 13 ++- .../packages/core/agent_framework/_types.py | 1 + .../core/tests/core/test_compaction.py | 68 +++++++++++++++ .../core/test_function_invocation_logic.py | 2 + 5 files changed, 129 insertions(+), 39 deletions(-) diff --git a/python/packages/core/agent_framework/_compaction.py b/python/packages/core/agent_framework/_compaction.py index 141637f615c..cd4696afca3 100644 --- a/python/packages/core/agent_framework/_compaction.py +++ b/python/packages/core/agent_framework/_compaction.py @@ -18,7 +18,7 @@ from ._sessions import ContextProvider from ._telemetry import FeatureIndex, mark_feature_used -from ._types import ChatResponse, Content, Message +from ._types import APPROVAL_REJECTION_RESULT_KEY, ChatResponse, Content, Message if TYPE_CHECKING: from ._clients import SupportsChatGetResponse @@ -1183,16 +1183,21 @@ def _count_included_tokens(messages: list[Message]) -> int: return included_token_count(messages) -def _latest_included_group(messages: list[Message]) -> list[Message]: - ordered_group_ids = _ordered_group_ids_from_annotations(messages) +def _latest_group(messages: list[Message], *, include_excluded: bool) -> list[Message]: grouped = _group_messages_by_id(messages) - for group_id in reversed(ordered_group_ids): - group = grouped.get(group_id, []) - if any(not message.additional_properties.get(EXCLUDED_KEY, False) for message in group): - return group + for message in reversed(messages): + if not include_excluded and message.additional_properties.get(EXCLUDED_KEY, False): + continue + group_id = _group_id(message) + if group_id is not None: + return grouped.get(group_id, []) return [] +def _latest_included_group(messages: list[Message]) -> list[Message]: + return _latest_group(messages, include_excluded=False) + + def _completed_function_calls( messages: Sequence[Message], *, @@ -1209,7 +1214,12 @@ def _completed_function_calls( content.call_id for message in candidate_messages for content in message.contents - if content.type == "function_result" and content.call_id and content.exception is None + if ( + content.type == "function_result" + and content.call_id + and content.exception is None + and content.additional_properties.get(APPROVAL_REJECTION_RESULT_KEY) is not True + ) } return [ content @@ -1236,8 +1246,9 @@ class AgentTriggeredCompactionStrategy: Wrapped strategies may provide ``create_trigger_tool()`` to replace the default ``compact_context`` tool with one specialized for their inputs. - The trigger is stored on the function-call content rather than on this - strategy instance, so one strategy can be shared safely by concurrent runs. + The consumed marker is stored on each trigger-group message rather than on + this strategy instance, so one strategy can be shared safely by concurrent + runs. """ DEFAULT_TOOL_NAME: Final[str] = "compact_context" @@ -1425,9 +1436,16 @@ def create_trigger_tool(self) -> FunctionTool: def request_skill_unload(skill_name: str | Sequence[str] | None = None) -> str: """Unload the full instructions and resources for selected skills.""" - requested_names = ( - [skill_name] if isinstance(skill_name, str) else skill_name if isinstance(skill_name, Sequence) else [] - ) + if skill_name is None: + requested_names: Sequence[str] = [] + elif isinstance(skill_name, str): + requested_names = [skill_name] + elif isinstance(skill_name, Sequence) and not isinstance(skill_name, (bytes, bytearray)): + if any(not isinstance(name, str) for name in skill_name): + raise TypeError("skill_name must contain only strings.") + requested_names = skill_name + else: + raise TypeError("skill_name must be a string, a sequence of strings, or None.") normalized_names = [name.strip() for name in requested_names if name.strip()] if normalized_names: formatted_names = ", ".join(f"'{name}'" for name in normalized_names) @@ -1462,11 +1480,8 @@ def request_skill_unload(skill_name: str | Sequence[str] | None = None) -> str: ) def _requested_skill_names(self, messages: list[Message]) -> set[str] | None: - ordered_group_ids = _ordered_group_ids_from_annotations(messages) - grouped = _group_messages_by_id(messages) - latest_group = grouped.get(ordered_group_ids[-1], []) if ordered_group_ids else [] requests = _completed_function_calls( - latest_group, + _latest_group(messages, include_excluded=True), tool_name=self.tool_name, include_consumed=True, include_excluded=True, @@ -1478,18 +1493,17 @@ def _requested_skill_names(self, messages: list[Message]) -> set[str] | None: for request in requests: arguments = request.parse_arguments() or {} raw_skill_names = arguments.get("skill_name") - skill_names = ( - [raw_skill_names] - if isinstance(raw_skill_names, str) - else cast("Sequence[Any]", raw_skill_names) - if isinstance(raw_skill_names, Sequence) - else [] - ) - normalized_names = { - skill_name.strip().casefold() - for skill_name in skill_names - if isinstance(skill_name, str) and skill_name.strip() - } + if raw_skill_names is None: + skill_names: Sequence[Any] = [] + elif isinstance(raw_skill_names, str): + skill_names = [raw_skill_names] + elif isinstance(raw_skill_names, Sequence) and not isinstance(raw_skill_names, (bytes, bytearray)): + skill_names = cast("Sequence[Any]", raw_skill_names) + else: + return set() + if any(not isinstance(skill_name, str) for skill_name in skill_names): + return set() + normalized_names = {skill_name.strip().casefold() for skill_name in skill_names if skill_name.strip()} if not normalized_names: return None requested_names.update(normalized_names) @@ -1498,11 +1512,17 @@ def _requested_skill_names(self, messages: list[Message]) -> set[str] | None: @staticmethod def _skill_names_for_removable_group(group: Sequence[Message]) -> set[str] | None: included_messages = [message for message in group if not message.additional_properties.get(EXCLUDED_KEY, False)] - calls = [ - content for message in included_messages for content in message.contents if content.type == "function_call" + tool_calls = [ + content + for message in included_messages + for content in message.contents + if content.type in _TOOL_CALL_CONTENT_TYPES ] - if not calls or any(call.name not in _SKILL_CONTEXT_TOOL_NAMES for call in calls): + if not tool_calls or any( + call.type != "function_call" or call.name not in _SKILL_CONTEXT_TOOL_NAMES for call in tool_calls + ): return None + calls = tool_calls completed_call_ids = { content.call_id diff --git a/python/packages/core/agent_framework/_tools.py b/python/packages/core/agent_framework/_tools.py index 08f00fefc11..53b64bce9d3 100644 --- a/python/packages/core/agent_framework/_tools.py +++ b/python/packages/core/agent_framework/_tools.py @@ -416,9 +416,7 @@ def _parse_inputs( # pyright: ignore[reportUnusedFunction] if inputs is None: return [] - from ._types import ( - Content, - ) + from ._types import Content parsed_inputs: list[Content] = [] if not isinstance(inputs, list): @@ -4449,9 +4447,7 @@ def _replace_approval_contents_with_results( Returns: The terminal contents produced while resolving the approval responses, in response order. """ - from ._types import ( - Content, - ) + from ._types import APPROVAL_REJECTION_RESULT_KEY, Content result_groups_by_call_id: dict[str, deque[list[Content]]] = {} for result_group in approved_function_result_groups: @@ -4578,7 +4574,10 @@ def find_approval_occurrence(approval_id: str) -> _ApprovalCallOccurrence | None call_id=call_id, id=occurrence.function_call.id if occurrence is not None else content.function_call.id, result="Error: Tool call invocation was rejected by user.", - additional_properties=content.function_call.additional_properties, + additional_properties={ + **content.function_call.additional_properties, + APPROVAL_REJECTION_RESULT_KEY: True, + }, ) ] if not replacements: diff --git a/python/packages/core/agent_framework/_types.py b/python/packages/core/agent_framework/_types.py index 629593b8839..0909e4fd173 100644 --- a/python/packages/core/agent_framework/_types.py +++ b/python/packages/core/agent_framework/_types.py @@ -459,6 +459,7 @@ def _restore_compaction_annotation_in_additional_properties( "code_interpreter_tool_call", "code_interpreter_tool_result", }) +APPROVAL_REJECTION_RESULT_KEY = "_approval_rejection_result" _CONTENT_TYPES: Final[frozenset[str]] = frozenset(get_args(ContentType)) diff --git a/python/packages/core/tests/core/test_compaction.py b/python/packages/core/tests/core/test_compaction.py index 660cc29446a..24a2ecd0cd9 100644 --- a/python/packages/core/tests/core/test_compaction.py +++ b/python/packages/core/tests/core/test_compaction.py @@ -45,6 +45,7 @@ append_compaction_message, extend_compaction_messages, ) +from agent_framework._types import APPROVAL_REJECTION_RESULT_KEY def _assistant_function_call(call_id: str) -> Message: @@ -1003,6 +1004,37 @@ async def test_agent_triggered_compaction_ignores_trigger_outside_latest_group() assert inner.call_count == 0 +async def test_agent_triggered_compaction_uses_terminal_result_for_nonadjacent_group_recency() -> None: + inner = _RecordingCompactionStrategy() + strategy = AgentTriggeredCompactionStrategy(inner) + messages = [ + Message(role="user", contents=["work"]), + _assistant_named_function_call("compact-1", strategy.tool.name), + Message(role="assistant", contents=["intervening context"]), + _tool_result("compact-1", "Compaction requested."), + ] + + assert await strategy(messages) is True + assert inner.call_count == 1 + + +async def test_agent_triggered_compaction_ignores_rejected_approval_result() -> None: + inner = _RecordingCompactionStrategy() + strategy = AgentTriggeredCompactionStrategy(inner) + rejection_result = Content.from_function_result( + call_id="compact-1", + result="Error: Tool call invocation was rejected by user.", + additional_properties={APPROVAL_REJECTION_RESULT_KEY: True}, + ) + messages = [ + _assistant_named_function_call("compact-1", strategy.tool.name), + Message(role="tool", contents=[rejection_result]), + ] + + assert await strategy(messages) is False + assert inner.call_count == 0 + + def test_agent_triggered_compaction_tool_describes_reason_parameter() -> None: strategy = AgentTriggeredCompactionStrategy(_RecordingCompactionStrategy()) @@ -1074,6 +1106,17 @@ def test_unload_skills_tool_describes_skill_name_parameter() -> None: assert "do not call it in parallel" in strategy.tool.description +@pytest.mark.parametrize( + "skill_name", + [123, [123], ["pdf", 123], b"pdf"], +) +async def test_unload_skills_tool_rejects_malformed_skill_names(skill_name: Any) -> None: + strategy = AgentTriggeredCompactionStrategy(UnloadSkillsCompactionStrategy()) + + with pytest.raises(TypeError, match="skill_name"): + await strategy.tool.invoke(arguments={"skill_name": skill_name}) + + async def test_unload_skills_compaction_empty_name_removes_all_skill_context() -> None: skill_strategy = UnloadSkillsCompactionStrategy() strategy = AgentTriggeredCompactionStrategy(skill_strategy) @@ -1140,6 +1183,31 @@ async def test_unload_skills_compaction_preserves_script_results_and_mixed_group assert all(message.additional_properties.get(EXCLUDED_KEY) is not True for message in messages) +async def test_unload_skills_compaction_preserves_group_with_hosted_tool_call() -> None: + strategy = AgentTriggeredCompactionStrategy(UnloadSkillsCompactionStrategy()) + messages = [ + Message( + role="assistant", + contents=[ + Content.from_function_call( + call_id="load-pdf", + name="load_skill", + arguments={"skill_name": "pdf"}, + ), + Content.from_shell_tool_call(call_id="shell-1", commands=["echo hi"]), + ], + ), + _tool_result("load-pdf", "pdf instructions"), + _assistant_named_function_call("unload-pdf", strategy.tool.name, {"skill_name": "pdf"}), + _tool_result("unload-pdf", "Skill 'pdf' will be unloaded."), + ] + + changed = await strategy(messages) + + assert changed is False + assert all(message.additional_properties.get(EXCLUDED_KEY) is not True for message in messages) + + async def test_selective_tool_call_strategy_excludes_older_tool_groups() -> None: messages = [ Message(role="user", contents=["u"]), diff --git a/python/packages/core/tests/core/test_function_invocation_logic.py b/python/packages/core/tests/core/test_function_invocation_logic.py index f185504d2e3..4b8a34dc829 100644 --- a/python/packages/core/tests/core/test_function_invocation_logic.py +++ b/python/packages/core/tests/core/test_function_invocation_logic.py @@ -7956,6 +7956,7 @@ def test_replace_approval_contents_with_results_correlates_reused_call_id_pendin def test_replace_approval_contents_with_results_resolves_rejected_request() -> None: """A rejected pending call produces one terminal result.""" from agent_framework._tools import _collect_approval_responses, _replace_approval_contents_with_results + from agent_framework._types import APPROVAL_REJECTION_RESULT_KEY function_call, request, _ = _build_approved_tool_roundtrip( call_id="call_rejected", @@ -7977,6 +7978,7 @@ def test_replace_approval_contents_with_results_resolves_rejected_request() -> N results = [content for message in messages for content in message.contents if content.type == "function_result"] assert len(results) == 1 assert results[0].result == "Error: Tool call invocation was rejected by user." + assert results[0].additional_properties[APPROVAL_REJECTION_RESULT_KEY] is True assert resolved_contents == results From 3d2ffa3073a4abd00cc11062578476f9c2a1be5f Mon Sep 17 00:00:00 2001 From: eavanvalkenburg Date: Thu, 15 Oct 2026 11:37:57 +0200 Subject: [PATCH 3/8] Python: Mark agent compaction experimental --- python/packages/core/AGENTS.md | 3 ++- python/packages/core/agent_framework/_compaction.py | 3 +++ .../packages/core/agent_framework/_feature_stage.py | 1 + python/packages/core/tests/core/test_compaction.py | 13 +++++++++++++ 4 files changed, 19 insertions(+), 1 deletion(-) diff --git a/python/packages/core/AGENTS.md b/python/packages/core/AGENTS.md index 0d7ab4e5763..788174af16b 100644 --- a/python/packages/core/AGENTS.md +++ b/python/packages/core/AGENTS.md @@ -172,7 +172,8 @@ The vector store API is experimental under the shared `VECTOR_STORES` feature ID group rather than on the strategy instance so shared agents remain safe across concurrent runs. The trigger group is temporarily excluded while the inner strategy runs, so it does not consume retention counts, then restored by default; set `exclude_trigger=True` to remove a trigger-only group from the next model call. Mixed groups that - contain unrelated tool calls are always restored atomically. + contain unrelated tool calls are always restored atomically. The agent-triggered compaction APIs are experimental + under the shared `AGENT_TRIGGERED_COMPACTION` feature ID. - **`UnloadSkillsCompactionStrategy`** - Intended as an inner agent-triggered strategy. It supplies an `unload_skill` trigger-tool factory that `AgentTriggeredCompactionStrategy` adopts automatically, then removes complete `load_skill` and `read_skill_resource` groups for one named skill, or all skills when no name is diff --git a/python/packages/core/agent_framework/_compaction.py b/python/packages/core/agent_framework/_compaction.py index cd4696afca3..e8929972404 100644 --- a/python/packages/core/agent_framework/_compaction.py +++ b/python/packages/core/agent_framework/_compaction.py @@ -16,6 +16,7 @@ runtime_checkable, ) +from ._feature_stage import ExperimentalFeature, experimental from ._sessions import ContextProvider from ._telemetry import FeatureIndex, mark_feature_used from ._types import APPROVAL_REJECTION_RESULT_KEY, ChatResponse, Content, Message @@ -1237,6 +1238,7 @@ def _completed_function_calls( ] +@experimental(feature_id=ExperimentalFeature.AGENT_TRIGGERED_COMPACTION) class AgentTriggeredCompactionStrategy: """Run an existing compaction strategy after an agent invokes a trigger tool. @@ -1398,6 +1400,7 @@ def _restore_trigger_group( message.additional_properties.pop(EXCLUDE_REASON_KEY, None) +@experimental(feature_id=ExperimentalFeature.AGENT_TRIGGERED_COMPACTION) class UnloadSkillsCompactionStrategy: """Exclude loaded skill bodies and resources selected by an unload tool. diff --git a/python/packages/core/agent_framework/_feature_stage.py b/python/packages/core/agent_framework/_feature_stage.py index 9421c12cc27..b4988db3c19 100644 --- a/python/packages/core/agent_framework/_feature_stage.py +++ b/python/packages/core/agent_framework/_feature_stage.py @@ -50,6 +50,7 @@ class ExperimentalFeature(str, Enum): on enum membership or attribute presence over time. """ + AGENT_TRIGGERED_COMPACTION = "AGENT_TRIGGERED_COMPACTION" AGENT_HOOKS = "AGENT_HOOKS" COMPUTER_USE = "COMPUTER_USE" DECLARATIVE_AGENTS = "DECLARATIVE_AGENTS" diff --git a/python/packages/core/tests/core/test_compaction.py b/python/packages/core/tests/core/test_compaction.py index 24a2ecd0cd9..8f8ad9381a6 100644 --- a/python/packages/core/tests/core/test_compaction.py +++ b/python/packages/core/tests/core/test_compaction.py @@ -45,6 +45,7 @@ append_compaction_message, extend_compaction_messages, ) +from agent_framework._feature_stage import _WARNED_FEATURES, ExperimentalFeature, ExperimentalWarning from agent_framework._types import APPROVAL_REJECTION_RESULT_KEY @@ -872,6 +873,18 @@ async def __call__(self, messages: list[Message]) -> bool: return self.changed +def test_agent_triggered_compaction_apis_are_experimental() -> None: + feature_id = ExperimentalFeature.AGENT_TRIGGERED_COMPACTION.value + for api in (AgentTriggeredCompactionStrategy, UnloadSkillsCompactionStrategy): + assert getattr(api, "__feature_stage__", None) == "experimental" + assert getattr(api, "__feature_id__", None) == feature_id + assert ".. warning:: Experimental" in (api.__doc__ or "") + + _WARNED_FEATURES.discard((ExperimentalWarning, feature_id)) + with pytest.warns(ExperimentalWarning, match=feature_id): + AgentTriggeredCompactionStrategy(UnloadSkillsCompactionStrategy()) + + async def test_agent_triggered_compaction_runs_only_after_completed_trigger() -> None: inner = _RecordingCompactionStrategy() strategy = AgentTriggeredCompactionStrategy(inner) From 95dc714bf2359366d3a5ac2110d0f44383195cea Mon Sep 17 00:00:00 2001 From: eavanvalkenburg Date: Thu, 15 Oct 2026 16:19:19 +0200 Subject: [PATCH 4/8] Python: Reject triggered compaction with service storage --- python/packages/core/AGENTS.md | 3 +- .../packages/core/agent_framework/_clients.py | 22 +++++ .../core/agent_framework/_compaction.py | 38 ++++++++ .../packages/core/tests/core/test_clients.py | 95 +++++++++++++++++++ 4 files changed, 157 insertions(+), 1 deletion(-) diff --git a/python/packages/core/AGENTS.md b/python/packages/core/AGENTS.md index 788174af16b..b98c862894d 100644 --- a/python/packages/core/AGENTS.md +++ b/python/packages/core/AGENTS.md @@ -173,7 +173,8 @@ The vector store API is experimental under the shared `VECTOR_STORES` feature ID is temporarily excluded while the inner strategy runs, so it does not consume retention counts, then restored by default; set `exclude_trigger=True` to remove a trigger-only group from the next model call. Mixed groups that contain unrelated tool calls are always restored atomically. The agent-triggered compaction APIs are experimental - under the shared `AGENT_TRIGGERED_COMPACTION` feature ID. + under the shared `AGENT_TRIGGERED_COMPACTION` feature ID and require client-managed history (`store=False` with no + service-managed conversation ID). - **`UnloadSkillsCompactionStrategy`** - Intended as an inner agent-triggered strategy. It supplies an `unload_skill` trigger-tool factory that `AgentTriggeredCompactionStrategy` adopts automatically, then removes complete `load_skill` and `read_skill_resource` groups for one named skill, or all skills when no name is diff --git a/python/packages/core/agent_framework/_clients.py b/python/packages/core/agent_framework/_clients.py index 564d78aa660..191162d157f 100644 --- a/python/packages/core/agent_framework/_clients.py +++ b/python/packages/core/agent_framework/_clients.py @@ -38,6 +38,7 @@ ResponseStream, validate_chat_options, ) +from .exceptions import ChatClientInvalidRequestException if sys.version_info >= (3, 13): from typing import TypeVar # pragma: no cover @@ -373,8 +374,27 @@ async def _prepare_messages_for_model_call( *, compaction_strategy: CompactionStrategy | None = None, tokenizer: TokenizerProtocol | None = None, + options: Mapping[str, Any] | None = None, ) -> list[Message]: prepared_messages = list(messages) + if compaction_strategy is not None: + from ._compaction import _contains_agent_triggered_compaction_strategy # pyright: ignore[reportPrivateUsage] + from ._sessions import is_local_history_conversation_id + + if _contains_agent_triggered_compaction_strategy(compaction_strategy): + store = options.get("store") if options is not None else None + service_stores_history = ( + getattr(self, "STORES_BY_DEFAULT", False) if store is None else store is not False + ) + conversation_id = options.get("conversation_id") if options is not None else None + has_service_conversation = conversation_id is not None and ( + not isinstance(conversation_id, str) or not is_local_history_conversation_id(conversation_id) + ) + if service_stores_history or has_service_conversation: + raise ChatClientInvalidRequestException( + "Agent-triggered compaction requires client-managed history. " + "Set store=False and do not use a service-managed conversation_id." + ) if compaction_strategy is None: if tokenizer is None: return prepared_messages @@ -531,6 +551,7 @@ async def _get_stream() -> ResponseStream[ChatResponseUpdate, ChatResponse[Any]] prepared_messages = await self._prepare_messages_for_model_call( messages, **compaction_overrides, + options=options, ) stream_response = self._inner_get_response( messages=prepared_messages, @@ -551,6 +572,7 @@ async def _get_response() -> ChatResponse[Any]: prepared_messages = await self._prepare_messages_for_model_call( messages, **compaction_overrides, + options=options, ) return await self._inner_get_response( messages=prepared_messages, diff --git a/python/packages/core/agent_framework/_compaction.py b/python/packages/core/agent_framework/_compaction.py index e8929972404..8b34b9d3973 100644 --- a/python/packages/core/agent_framework/_compaction.py +++ b/python/packages/core/agent_framework/_compaction.py @@ -1238,6 +1238,39 @@ def _completed_function_calls( ] +def _contains_agent_triggered_compaction_strategy( + strategy: object | None, + *, + _seen: set[int] | None = None, +) -> bool: + if strategy is None: + return False + seen: set[int] = _seen if _seen is not None else set() + strategy_id = id(strategy) + if strategy_id in seen: + return False + seen.add(strategy_id) + + if isinstance(strategy, AgentTriggeredCompactionStrategy): + return True + + nested_strategy: object | None = getattr(strategy, "strategy", None) + if ( + nested_strategy is not None + and nested_strategy is not strategy + and _contains_agent_triggered_compaction_strategy(nested_strategy, _seen=seen) + ): + return True + + nested_strategies: object = getattr(strategy, "strategies", None) + if isinstance(nested_strategies, Sequence) and not isinstance(nested_strategies, (str, bytes, bytearray)): + return any( + _contains_agent_triggered_compaction_strategy(nested, _seen=seen) + for nested in cast("Sequence[object]", nested_strategies) + ) + return False + + @experimental(feature_id=ExperimentalFeature.AGENT_TRIGGERED_COMPACTION) class AgentTriggeredCompactionStrategy: """Run an existing compaction strategy after an agent invokes a trigger tool. @@ -1251,6 +1284,11 @@ class AgentTriggeredCompactionStrategy: The consumed marker is stored on each trigger-group message rather than on this strategy instance, so one strategy can be shared safely by concurrent runs. + + This strategy requires client-managed history (for example ``store=False``). + Chat clients reject it when service-side storage or a service-managed + conversation ID is active because local compaction cannot remove remote + conversation context. """ DEFAULT_TOOL_NAME: Final[str] = "compact_context" diff --git a/python/packages/core/tests/core/test_clients.py b/python/packages/core/tests/core/test_clients.py index b122b9c5ebf..37ac55a460a 100644 --- a/python/packages/core/tests/core/test_clients.py +++ b/python/packages/core/tests/core/test_clients.py @@ -14,6 +14,7 @@ GROUP_TOKEN_COUNT_KEY, SUMMARIZED_BY_SUMMARY_ID_KEY, SUMMARY_OF_MESSAGE_IDS_KEY, + AgentTriggeredCompactionStrategy, BaseChatClient, ChatMiddleware, ChatResponse, @@ -30,6 +31,8 @@ included_messages, tool, ) +from agent_framework._sessions import LOCAL_HISTORY_CONVERSATION_ID +from agent_framework.exceptions import ChatClientInvalidRequestException class _NoOpChatMiddleware(ChatMiddleware): @@ -247,6 +250,98 @@ async def _capture( assert captured_roles == [["assistant"]] +@pytest.mark.parametrize("stream", [False, True]) +async def test_agent_triggered_compaction_rejects_default_service_storage( + chat_client_base: SupportsChatGetResponse, + stream: bool, +) -> None: + chat_client_base.function_invocation_configuration["enabled"] = False # type: ignore[attr-defined] # ty: ignore[unresolved-attribute] + chat_client_base.STORES_BY_DEFAULT = True # type: ignore[attr-defined] # ty: ignore[unresolved-attribute] + strategy = AgentTriggeredCompactionStrategy(TruncationStrategy(max_n=2, compact_to=1)) + + response = chat_client_base.get_response( + [Message(role="user", contents=["Hello"])], + stream=stream, + compaction_strategy=strategy, + ) + with pytest.raises(ChatClientInvalidRequestException, match="store=False"): + if stream: + async for _ in response: # type: ignore[union-attr] # ty: ignore[union-attr] + pass + else: + await response # type: ignore[misc] # ty: ignore[not-awaitable] + + +async def test_agent_triggered_compaction_rejects_explicit_store_true( + chat_client_base: SupportsChatGetResponse, +) -> None: + chat_client_base.function_invocation_configuration["enabled"] = False # type: ignore[attr-defined] # ty: ignore[unresolved-attribute] + strategy = AgentTriggeredCompactionStrategy(TruncationStrategy(max_n=2, compact_to=1)) + + with pytest.raises(ChatClientInvalidRequestException, match="store=False"): + await chat_client_base.get_response( + [Message(role="user", contents=["Hello"])], + options={"store": True}, + compaction_strategy=strategy, + ) + + +async def test_nested_agent_triggered_compaction_rejects_service_storage( + chat_client_base: SupportsChatGetResponse, +) -> None: + class StrategyWrapper: + def __init__(self, strategy: Any) -> None: + self.strategy = strategy + + async def __call__(self, messages: list[Message]) -> bool: + return await self.strategy(messages) + + chat_client_base.function_invocation_configuration["enabled"] = False # type: ignore[attr-defined] # ty: ignore[unresolved-attribute] + chat_client_base.STORES_BY_DEFAULT = True # type: ignore[attr-defined] # ty: ignore[unresolved-attribute] + strategy = StrategyWrapper(AgentTriggeredCompactionStrategy(TruncationStrategy(max_n=2, compact_to=1))) + + with pytest.raises(ChatClientInvalidRequestException, match="store=False"): + await chat_client_base.get_response( + [Message(role="user", contents=["Hello"])], + compaction_strategy=strategy, + ) + + +async def test_agent_triggered_compaction_rejects_service_conversation_id( + chat_client_base: SupportsChatGetResponse, +) -> None: + chat_client_base.function_invocation_configuration["enabled"] = False # type: ignore[attr-defined] # ty: ignore[unresolved-attribute] + strategy = AgentTriggeredCompactionStrategy(TruncationStrategy(max_n=2, compact_to=1)) + + with pytest.raises(ChatClientInvalidRequestException, match="service-managed conversation_id"): + await chat_client_base.get_response( + [Message(role="user", contents=["Hello"])], + options={"store": False, "conversation_id": "service-conversation"}, + compaction_strategy=strategy, + ) + + +@pytest.mark.parametrize("conversation_id", [None, LOCAL_HISTORY_CONVERSATION_ID]) +async def test_agent_triggered_compaction_allows_client_managed_history( + chat_client_base: SupportsChatGetResponse, + conversation_id: str | None, +) -> None: + chat_client_base.function_invocation_configuration["enabled"] = False # type: ignore[attr-defined] # ty: ignore[unresolved-attribute] + chat_client_base.STORES_BY_DEFAULT = True # type: ignore[attr-defined] # ty: ignore[unresolved-attribute] + strategy = AgentTriggeredCompactionStrategy(TruncationStrategy(max_n=2, compact_to=1)) + options: dict[str, Any] = {"store": False} + if conversation_id is not None: + options["conversation_id"] = conversation_id + + response = await chat_client_base.get_response( + [Message(role="user", contents=["Hello"])], + options=options, + compaction_strategy=strategy, + ) + + assert response.text == "test response - Hello" + + async def test_base_client_per_call_tokenizer_override_annotates_messages( chat_client_base: SupportsChatGetResponse, ) -> None: From 799c35ce5dffb1ae93dbe9e7dc064d71f2c8ea0b Mon Sep 17 00:00:00 2001 From: eavanvalkenburg Date: Thu, 15 Oct 2026 17:43:36 +0200 Subject: [PATCH 5/8] Python: Fix compaction guard test typing --- .../packages/core/tests/core/test_clients.py | 26 ++++++++++++------- 1 file changed, 16 insertions(+), 10 deletions(-) diff --git a/python/packages/core/tests/core/test_clients.py b/python/packages/core/tests/core/test_clients.py index 37ac55a460a..eb6c69dc9b9 100644 --- a/python/packages/core/tests/core/test_clients.py +++ b/python/packages/core/tests/core/test_clients.py @@ -17,6 +17,7 @@ AgentTriggeredCompactionStrategy, BaseChatClient, ChatMiddleware, + ChatOptions, ChatResponse, ChatResponseUpdate, Content, @@ -259,17 +260,22 @@ async def test_agent_triggered_compaction_rejects_default_service_storage( chat_client_base.STORES_BY_DEFAULT = True # type: ignore[attr-defined] # ty: ignore[unresolved-attribute] strategy = AgentTriggeredCompactionStrategy(TruncationStrategy(max_n=2, compact_to=1)) - response = chat_client_base.get_response( - [Message(role="user", contents=["Hello"])], - stream=stream, - compaction_strategy=strategy, - ) with pytest.raises(ChatClientInvalidRequestException, match="store=False"): if stream: - async for _ in response: # type: ignore[union-attr] # ty: ignore[union-attr] + response_stream = chat_client_base.get_response( + [Message(role="user", contents=["Hello"])], + stream=True, + compaction_strategy=strategy, + ) + async for _ in response_stream: pass else: - await response # type: ignore[misc] # ty: ignore[not-awaitable] + response = chat_client_base.get_response( + [Message(role="user", contents=["Hello"])], + stream=False, + compaction_strategy=strategy, + ) + _ = await response async def test_agent_triggered_compaction_rejects_explicit_store_true( @@ -281,7 +287,7 @@ async def test_agent_triggered_compaction_rejects_explicit_store_true( with pytest.raises(ChatClientInvalidRequestException, match="store=False"): await chat_client_base.get_response( [Message(role="user", contents=["Hello"])], - options={"store": True}, + options=ChatOptions(store=True), compaction_strategy=strategy, ) @@ -316,7 +322,7 @@ async def test_agent_triggered_compaction_rejects_service_conversation_id( with pytest.raises(ChatClientInvalidRequestException, match="service-managed conversation_id"): await chat_client_base.get_response( [Message(role="user", contents=["Hello"])], - options={"store": False, "conversation_id": "service-conversation"}, + options=ChatOptions(store=False, conversation_id="service-conversation"), compaction_strategy=strategy, ) @@ -329,7 +335,7 @@ async def test_agent_triggered_compaction_allows_client_managed_history( chat_client_base.function_invocation_configuration["enabled"] = False # type: ignore[attr-defined] # ty: ignore[unresolved-attribute] chat_client_base.STORES_BY_DEFAULT = True # type: ignore[attr-defined] # ty: ignore[unresolved-attribute] strategy = AgentTriggeredCompactionStrategy(TruncationStrategy(max_n=2, compact_to=1)) - options: dict[str, Any] = {"store": False} + options = ChatOptions(store=False) if conversation_id is not None: options["conversation_id"] = conversation_id From af9f54ace3d01baf81444e9a4dc8d15793380234 Mon Sep 17 00:00:00 2001 From: eavanvalkenburg Date: Mon, 26 Oct 2026 13:18:47 +0100 Subject: [PATCH 6/8] Python: Harden triggered compaction correlation --- .../packages/core/agent_framework/_clients.py | 22 +++++-- .../core/agent_framework/_compaction.py | 58 ++++++++++++++---- .../packages/core/tests/core/test_clients.py | 2 +- .../core/tests/core/test_compaction.py | 61 +++++++++++++++++++ .../agent_framework_openai/_chat_client.py | 5 ++ .../tests/openai/test_openai_chat_client.py | 39 ++++++++++++ 6 files changed, 169 insertions(+), 18 deletions(-) diff --git a/python/packages/core/agent_framework/_clients.py b/python/packages/core/agent_framework/_clients.py index 191162d157f..1eed178ac18 100644 --- a/python/packages/core/agent_framework/_clients.py +++ b/python/packages/core/agent_framework/_clients.py @@ -289,6 +289,8 @@ async def _stream(): When ``True``, the agent skips auto-injecting ``InMemoryHistoryProvider`` unless the user explicitly sets ``store=False``. """ + SERVICE_CONTINUATION_OPTION_NAMES: ClassVar[tuple[str, ...]] = ("conversation_id",) + """Option names that activate provider-managed conversation continuation.""" # OTEL_PROVIDER_NAME is used for OTel setup, should be overridden in subclasses def __init__( @@ -386,14 +388,24 @@ async def _prepare_messages_for_model_call( service_stores_history = ( getattr(self, "STORES_BY_DEFAULT", False) if store is None else store is not False ) - conversation_id = options.get("conversation_id") if options is not None else None - has_service_conversation = conversation_id is not None and ( - not isinstance(conversation_id, str) or not is_local_history_conversation_id(conversation_id) - ) + has_service_conversation = False + if options is not None: + for option_name in self.SERVICE_CONTINUATION_OPTION_NAMES: + value = options.get(option_name) + if value is None or value == "": + continue + if ( + option_name == "conversation_id" + and isinstance(value, str) + and is_local_history_conversation_id(value) + ): + continue + has_service_conversation = True + break if service_stores_history or has_service_conversation: raise ChatClientInvalidRequestException( "Agent-triggered compaction requires client-managed history. " - "Set store=False and do not use a service-managed conversation_id." + "Set store=False and do not use a service-managed continuation option." ) if compaction_strategy is None: if tokenizer is None: diff --git a/python/packages/core/agent_framework/_compaction.py b/python/packages/core/agent_framework/_compaction.py index 8b34b9d3973..40c37e49f39 100644 --- a/python/packages/core/agent_framework/_compaction.py +++ b/python/packages/core/agent_framework/_compaction.py @@ -1211,17 +1211,51 @@ def _completed_function_calls( if include_excluded else [message for message in messages if not message.additional_properties.get(EXCLUDED_KEY, False)] ) - successful_result_call_ids = { - content.call_id - for message in candidate_messages - for content in message.contents - if ( - content.type == "function_result" - and content.call_id - and content.exception is None - and content.additional_properties.get(APPROVAL_REJECTION_RESULT_KEY) is not True - ) - } + open_calls_by_call_id: dict[str, list[Content]] = {} + open_calls_by_occurrence_id: dict[str, list[Content]] = {} + closed_call_identities: set[int] = set() + successful_call_identities: set[int] = set() + + for message in candidate_messages: + for content in message.contents: + if content.type == "function_call" and content.call_id: + open_calls_by_call_id.setdefault(content.call_id, []).append(content) + if content.id: + open_calls_by_occurrence_id.setdefault(content.id, []).append(content) + continue + if content.type != "function_result" or not content.call_id: + continue + + matched_call: Content | None = None + if content.id is not None: + matched_call = next( + ( + call + for call in open_calls_by_occurrence_id.get(content.id, ()) + if id(call) not in closed_call_identities and call.call_id == content.call_id + ), + None, + ) + else: + matched_call = next( + ( + call + for call in open_calls_by_call_id.get(content.call_id, ()) + if id(call) not in closed_call_identities + ), + None, + ) + if matched_call is None: + continue + + matched_identity = id(matched_call) + closed_call_identities.add(matched_identity) + if ( + content.exception is None + and content.additional_properties.get(APPROVAL_REJECTION_RESULT_KEY) is not True + ): + successful_call_identities.add(matched_identity) + return [ content for message in candidate_messages @@ -1229,7 +1263,7 @@ def _completed_function_calls( if ( content.type == "function_call" and content.name == tool_name - and content.call_id in successful_result_call_ids + and id(content) in successful_call_identities and ( include_consumed or message.additional_properties.get(_AGENT_COMPACTION_TRIGGER_CONSUMED_KEY) is not True diff --git a/python/packages/core/tests/core/test_clients.py b/python/packages/core/tests/core/test_clients.py index eb6c69dc9b9..5ee1af25799 100644 --- a/python/packages/core/tests/core/test_clients.py +++ b/python/packages/core/tests/core/test_clients.py @@ -319,7 +319,7 @@ async def test_agent_triggered_compaction_rejects_service_conversation_id( chat_client_base.function_invocation_configuration["enabled"] = False # type: ignore[attr-defined] # ty: ignore[unresolved-attribute] strategy = AgentTriggeredCompactionStrategy(TruncationStrategy(max_n=2, compact_to=1)) - with pytest.raises(ChatClientInvalidRequestException, match="service-managed conversation_id"): + with pytest.raises(ChatClientInvalidRequestException, match="service-managed continuation"): await chat_client_base.get_response( [Message(role="user", contents=["Hello"])], options=ChatOptions(store=False, conversation_id="service-conversation"), diff --git a/python/packages/core/tests/core/test_compaction.py b/python/packages/core/tests/core/test_compaction.py index 8f8ad9381a6..d9ae075b02d 100644 --- a/python/packages/core/tests/core/test_compaction.py +++ b/python/packages/core/tests/core/test_compaction.py @@ -39,6 +39,7 @@ included_token_count, ) from agent_framework._compaction import ( + _completed_function_calls, _format_summary_message, _select_summary_input_groups, _serialize_message, @@ -1048,6 +1049,66 @@ async def test_agent_triggered_compaction_ignores_rejected_approval_result() -> assert inner.call_count == 0 +def test_completed_function_calls_correlates_reused_call_id_by_occurrence() -> None: + rejected_call = Content.from_function_call( + id="occ-rejected", + call_id="shared", + name="compact_context", + arguments={"reason": "rejected"}, + ) + approved_call = Content.from_function_call( + id="occ-approved", + call_id="shared", + name="compact_context", + arguments={"reason": "approved"}, + ) + unresolved_call = Content.from_function_call( + id="occ-unresolved", + call_id="shared", + name="compact_context", + arguments={"reason": "unresolved"}, + ) + messages = [ + Message(role="assistant", contents=[rejected_call, approved_call, unresolved_call]), + Message( + role="tool", + contents=[ + Content.from_function_result( + id="occ-rejected", + call_id="shared", + result="Error: Tool call invocation was rejected by user.", + additional_properties={APPROVAL_REJECTION_RESULT_KEY: True}, + ), + Content.from_function_result( + id="occ-approved", + call_id="shared", + result="Compaction requested.", + ), + ], + ), + ] + + completed = _completed_function_calls(messages, tool_name="compact_context") + + assert completed == [approved_call] + + +def test_completed_function_calls_uses_fifo_for_idless_reused_call_id() -> None: + first_call = Content.from_function_call(call_id="shared", name="compact_context", arguments={}) + second_call = Content.from_function_call(call_id="shared", name="compact_context", arguments={}) + messages = [ + Message(role="assistant", contents=[first_call, second_call]), + Message( + role="tool", + contents=[Content.from_function_result(call_id="shared", result="Compaction requested.")], + ), + ] + + completed = _completed_function_calls(messages, tool_name="compact_context") + + assert completed == [first_call] + + def test_agent_triggered_compaction_tool_describes_reason_parameter() -> None: strategy = AgentTriggeredCompactionStrategy(_RecordingCompactionStrategy()) diff --git a/python/packages/openai/agent_framework_openai/_chat_client.py b/python/packages/openai/agent_framework_openai/_chat_client.py index a5c18ebda85..358855267d0 100644 --- a/python/packages/openai/agent_framework_openai/_chat_client.py +++ b/python/packages/openai/agent_framework_openai/_chat_client.py @@ -601,6 +601,11 @@ class RawOpenAIChatClient( INJECTABLE: ClassVar[set[str]] = {"client"} STORES_BY_DEFAULT: ClassVar[bool] = True + SERVICE_CONTINUATION_OPTION_NAMES: ClassVar[tuple[str, ...]] = ( + "conversation_id", + "previous_response_id", + "conversation", + ) SUPPORTS_RICH_FUNCTION_OUTPUT: ClassVar[bool] = True _FEATURE_USAGE_INDEX: ClassVar[int | None] = FeatureIndex.OPENAI diff --git a/python/packages/openai/tests/openai/test_openai_chat_client.py b/python/packages/openai/tests/openai/test_openai_chat_client.py index 15cadc9d9b9..3b94aa4d912 100644 --- a/python/packages/openai/tests/openai/test_openai_chat_client.py +++ b/python/packages/openai/tests/openai/test_openai_chat_client.py @@ -16,6 +16,7 @@ import pytest from agent_framework import ( Agent, + AgentTriggeredCompactionStrategy, Annotation, ChatOptions, ChatResponse, @@ -30,6 +31,7 @@ SupportsImageGenerationTool, SupportsMCPTool, SupportsWebSearchTool, + TruncationStrategy, WorkflowBuilder, tool, ) @@ -89,6 +91,43 @@ ) +@pytest.mark.parametrize("stream", [False, True]) +@pytest.mark.parametrize( + ("option_name", "option_value"), + [ + ("previous_response_id", "resp_existing"), + ("conversation", "conv_existing"), + ], +) +async def test_agent_triggered_compaction_rejects_native_service_continuation_aliases( + stream: bool, + option_name: str, + option_value: str, +) -> None: + client = OpenAIChatClient(model="test-model", api_key="test-key") + strategy = AgentTriggeredCompactionStrategy(TruncationStrategy(max_n=2, compact_to=1)) + options = cast(OpenAIChatOptions, {"store": False, option_name: option_value}) + + with pytest.raises(ChatClientInvalidRequestException, match="service-managed continuation"): + if stream: + response_stream = client.get_response( + [Message(role="user", contents=["Hello"])], + stream=True, + options=options, + compaction_strategy=strategy, + ) + async for _ in response_stream: + pass + else: + response = client.get_response( + [Message(role="user", contents=["Hello"])], + stream=False, + options=options, + compaction_strategy=strategy, + ) + _ = await response + + class OutputStruct(BaseModel): """A structured output for testing purposes.""" From 18c436f2597c7b34712a28f4ca4e8785e28d4d13 Mon Sep 17 00:00:00 2001 From: eavanvalkenburg Date: Mon, 26 Oct 2026 13:40:47 +0100 Subject: [PATCH 7/8] Python: Preserve triggered compaction context --- .../core/agent_framework/_compaction.py | 50 ++++++-- .../packages/core/tests/core/test_agents.py | 18 ++- .../core/tests/core/test_compaction.py | 113 ++++++++++++++++++ 3 files changed, 169 insertions(+), 12 deletions(-) diff --git a/python/packages/core/agent_framework/_compaction.py b/python/packages/core/agent_framework/_compaction.py index 40c37e49f39..b9d66f8b1f8 100644 --- a/python/packages/core/agent_framework/_compaction.py +++ b/python/packages/core/agent_framework/_compaction.py @@ -56,6 +56,7 @@ "code_interpreter_tool_call", "shell_tool_call", "image_generation_tool_call", + "search_tool_call", } @@ -1411,6 +1412,7 @@ def request_compaction(reason: str | None = None) -> str: async def __call__(self, messages: list[Message]) -> bool: """Run the wrapped strategy once for the latest completed trigger call.""" annotate_message_groups(messages) + grouped = _group_messages_by_id(messages) trigger_group = _latest_included_group(messages) trigger_calls = _completed_function_calls( trigger_group, @@ -1419,6 +1421,21 @@ async def __call__(self, messages: list[Message]) -> bool: if not trigger_calls: return False + groups_to_hide: list[list[Message]] = [trigger_group] + trigger_group_ids = {id(message) for message in trigger_group} + for group in grouped.values(): + if any(id(message) in trigger_group_ids for message in group): + continue + if not any( + message.additional_properties.get(_AGENT_COMPACTION_TRIGGER_CONSUMED_KEY) is True + and not message.additional_properties.get(EXCLUDED_KEY, False) + for message in group + ): + continue + if self._is_trigger_only_group(group): + groups_to_hide.append(group) + + hidden_messages = [message for group in groups_to_hide for message in group] trigger_state = [ ( message, @@ -1427,9 +1444,9 @@ async def __call__(self, messages: list[Message]) -> bool: EXCLUDE_REASON_KEY in message.additional_properties, message.additional_properties.get(EXCLUDE_REASON_KEY), ) - for message in trigger_group + for message in hidden_messages ] - for message in trigger_group: + for message in hidden_messages: set_excluded(message, excluded=True, reason="agent_compaction_trigger") try: @@ -1441,22 +1458,37 @@ async def __call__(self, messages: list[Message]) -> bool: for message in trigger_group: message.additional_properties[_AGENT_COMPACTION_TRIGGER_CONSUMED_KEY] = True - has_unrelated_tool_calls = any( - content.type in _TOOL_CALL_CONTENT_TYPES - and not (content.type == "function_call" and content.name == self.tool.name) - for message in trigger_group - for content in message.contents - ) + has_unrelated_tool_calls = not self._is_trigger_only_group(trigger_group, include_excluded=True) trigger_message_ids = {id(message) for message in trigger_group} has_other_included_context = any( id(message) not in trigger_message_ids and not message.additional_properties.get(EXCLUDED_KEY, False) for message in messages ) trigger_changed = self.exclude_trigger and not has_unrelated_tool_calls and has_other_included_context + prior_trigger_state = [state for state in trigger_state if id(state[0]) not in trigger_message_ids] + self._restore_trigger_group(prior_trigger_state) if not trigger_changed: - self._restore_trigger_group(trigger_state) + current_trigger_state = [state for state in trigger_state if id(state[0]) in trigger_message_ids] + self._restore_trigger_group(current_trigger_state) return changed or trigger_changed + def _is_trigger_only_group( + self, + group: Sequence[Message], + *, + include_excluded: bool = False, + ) -> bool: + tool_calls = [ + content + for message in group + if include_excluded or not message.additional_properties.get(EXCLUDED_KEY, False) + for content in message.contents + if content.type in _TOOL_CALL_CONTENT_TYPES + ] + return bool(tool_calls) and all( + content.type == "function_call" and content.name == self.tool.name for content in tool_calls + ) + @staticmethod def _restore_trigger_group( trigger_state: Sequence[tuple[Message, bool, Any, bool, Any]], diff --git a/python/packages/core/tests/core/test_agents.py b/python/packages/core/tests/core/test_agents.py index df83fb923e1..76cb0d42df8 100644 --- a/python/packages/core/tests/core/test_agents.py +++ b/python/packages/core/tests/core/test_agents.py @@ -3472,12 +3472,14 @@ async def test_agent_triggered_compaction_applies_before_next_model_call( ] captured_roles: list[list[str]] = [] + captured_texts: list[list[str | None]] = [] original_inner = chat_client_base._inner_get_response def capture_inner( *, messages: MutableSequence[Message], stream: bool, options: dict[str, Any], **kwargs: Any ) -> Any: captured_roles.append([message.role for message in messages]) + captured_texts.append([message.text for message in messages]) return original_inner(messages=messages, stream=stream, options=options, **kwargs) agent = Agent( @@ -3487,13 +3489,23 @@ def capture_inner( ) with patch.object(chat_client_base, "_inner_get_response", side_effect=capture_inner): + run_messages = [ + Message(role="user", contents=["original task"]), + Message(role="assistant", contents=["obsolete planning detail"]), + Message(role="user", contents=["plan"]), + ] if stream: - result_text = "".join([update.text or "" async for update in agent.run("plan", stream=True)]) + result_text = "".join([update.text or "" async for update in agent.run(run_messages, stream=True)]) else: - result_text = (await agent.run("plan")).text + result_text = (await agent.run(run_messages)).text assert result_text == "done" - assert captured_roles == [["user"], ["user", "assistant", "tool"]] + assert captured_roles == [ + ["user", "assistant", "user"], + ["user", "user", "assistant", "tool"], + ] + assert "obsolete planning detail" in captured_texts[0] + assert "obsolete planning detail" not in captured_texts[1] async def test_chat_agent_uses_client_compaction_defaults_when_agent_unset(chat_client_base: Any) -> None: diff --git a/python/packages/core/tests/core/test_compaction.py b/python/packages/core/tests/core/test_compaction.py index d9ae075b02d..7e6db7bd17a 100644 --- a/python/packages/core/tests/core/test_compaction.py +++ b/python/packages/core/tests/core/test_compaction.py @@ -990,6 +990,80 @@ async def test_agent_trigger_in_mixed_tool_batch_is_position_independent_and_ret assert messages[4].additional_properties.get(EXCLUDED_KEY) is not True +async def test_agent_trigger_in_mixed_search_batch_is_retained() -> None: + strategy = AgentTriggeredCompactionStrategy( + SelectiveToolCallCompactionStrategy(keep_last_tool_call_groups=0), + exclude_trigger=True, + ) + messages = [ + Message(role="user", contents=["work"]), + _assistant_named_function_call("old-work", "old_work"), + _tool_result("old-work", "old result"), + Message( + role="assistant", + contents=[ + Content.from_function_call(call_id="compact-1", name=strategy.tool.name, arguments={}), + Content.from_search_tool_call( + call_id="search-1", + tool_name="web_search", + arguments={"query": "latest"}, + ), + ], + ), + Message( + role="tool", + contents=[ + Content.from_function_result(call_id="compact-1", result="Compaction requested."), + Content.from_search_tool_result( + call_id="search-1", + tool_name="web_search", + result="search result", + ), + ], + ), + ] + + changed = await strategy(messages) + + assert changed is True + assert messages[1].additional_properties.get(EXCLUDED_KEY) is True + assert messages[2].additional_properties.get(EXCLUDED_KEY) is True + assert messages[3].additional_properties.get(EXCLUDED_KEY) is not True + assert messages[4].additional_properties.get(EXCLUDED_KEY) is not True + + +async def test_successive_triggers_do_not_consume_tool_group_retention() -> None: + strategy = AgentTriggeredCompactionStrategy(SelectiveToolCallCompactionStrategy(keep_last_tool_call_groups=2)) + messages = [ + Message(role="user", contents=["work"]), + _assistant_named_function_call("work-1", "work_tool"), + _tool_result("work-1", "work result 1"), + _assistant_named_function_call("work-2", "work_tool"), + _tool_result("work-2", "work result 2"), + _assistant_named_function_call("compact-1", strategy.tool.name), + _tool_result("compact-1", "Compaction requested."), + ] + + assert await strategy(messages) is False + + messages.extend([ + _assistant_named_function_call("work-3", "work_tool"), + _tool_result("work-3", "work result 3"), + _assistant_named_function_call("compact-2", strategy.tool.name), + _tool_result("compact-2", "Compaction requested."), + ]) + + assert await strategy(messages) is True + assert messages[1].additional_properties.get(EXCLUDED_KEY) is True + assert messages[2].additional_properties.get(EXCLUDED_KEY) is True + assert messages[3].additional_properties.get(EXCLUDED_KEY) is not True + assert messages[4].additional_properties.get(EXCLUDED_KEY) is not True + assert messages[5].additional_properties.get(EXCLUDED_KEY) is not True + assert messages[6].additional_properties.get(EXCLUDED_KEY) is not True + assert messages[7].additional_properties.get(EXCLUDED_KEY) is not True + assert messages[8].additional_properties.get(EXCLUDED_KEY) is not True + + async def test_agent_triggered_compaction_restores_request_when_inner_strategy_fails() -> None: class _FailingStrategy: async def __call__(self, messages: list[Message]) -> bool: @@ -1282,6 +1356,45 @@ async def test_unload_skills_compaction_preserves_group_with_hosted_tool_call() assert all(message.additional_properties.get(EXCLUDED_KEY) is not True for message in messages) +async def test_unload_skills_compaction_preserves_group_with_search_tool_call() -> None: + strategy = AgentTriggeredCompactionStrategy(UnloadSkillsCompactionStrategy()) + messages = [ + Message( + role="assistant", + contents=[ + Content.from_function_call( + call_id="load-pdf", + name="load_skill", + arguments={"skill_name": "pdf"}, + ), + Content.from_search_tool_call( + call_id="search-1", + tool_name="web_search", + arguments={"query": "pdf"}, + ), + ], + ), + Message( + role="tool", + contents=[ + Content.from_function_result(call_id="load-pdf", result="pdf instructions"), + Content.from_search_tool_result( + call_id="search-1", + tool_name="web_search", + result="search result", + ), + ], + ), + _assistant_named_function_call("unload-pdf", strategy.tool.name, {"skill_name": "pdf"}), + _tool_result("unload-pdf", "Skill 'pdf' will be unloaded."), + ] + + changed = await strategy(messages) + + assert changed is False + assert all(message.additional_properties.get(EXCLUDED_KEY) is not True for message in messages) + + async def test_selective_tool_call_strategy_excludes_older_tool_groups() -> None: messages = [ Message(role="user", contents=["u"]), From 2bda4ca1f819867cf4e4422f96838368ade73439 Mon Sep 17 00:00:00 2001 From: eavanvalkenburg Date: Mon, 26 Oct 2026 16:21:10 +0100 Subject: [PATCH 8/8] Python: Harden triggered compaction composition --- python/packages/core/AGENTS.md | 5 +- .../core/agent_framework/_compaction.py | 132 +++++++++------- .../packages/core/tests/core/test_agents.py | 90 +++++++++++ .../core/tests/core/test_compaction.py | 143 +++++++++++++++++- 4 files changed, 310 insertions(+), 60 deletions(-) diff --git a/python/packages/core/AGENTS.md b/python/packages/core/AGENTS.md index b98c862894d..11fe82b46c9 100644 --- a/python/packages/core/AGENTS.md +++ b/python/packages/core/AGENTS.md @@ -168,8 +168,9 @@ The vector store API is experimental under the shared `VECTOR_STORES` feature ID ### Compaction (`_compaction.py`) - **`AgentTriggeredCompactionStrategy`** - Wraps one existing compaction strategy and invokes it only after its - explicitly registered trigger tool completes in the latest message group. Trigger consumption is stored on that - group rather than on the strategy instance so shared agents remain safe across concurrent runs. The trigger group + explicitly registered trigger tool completes in the latest message group. Trigger consumption is stored on the + completed trigger call occurrence rather than on the strategy instance so shared agents remain safe across + concurrent runs and composed trigger tools do not suppress each other. The trigger group is temporarily excluded while the inner strategy runs, so it does not consume retention counts, then restored by default; set `exclude_trigger=True` to remove a trigger-only group from the next model call. Mixed groups that contain unrelated tool calls are always restored atomically. The agent-triggered compaction APIs are experimental diff --git a/python/packages/core/agent_framework/_compaction.py b/python/packages/core/agent_framework/_compaction.py index b9d66f8b1f8..64e37a4fbb2 100644 --- a/python/packages/core/agent_framework/_compaction.py +++ b/python/packages/core/agent_framework/_compaction.py @@ -44,6 +44,7 @@ SUMMARY_OF_GROUP_IDS_KEY = "_summary_of_group_ids" SUMMARIZED_BY_SUMMARY_ID_KEY = "_summarized_by_summary_id" _AGENT_COMPACTION_TRIGGER_CONSUMED_KEY = "_agent_compaction_trigger_consumed" +_SKILL_UNLOAD_SELECTION_KEY = "_skill_unload_selection" _SKILL_CONTEXT_TOOL_NAMES: Final[frozenset[str]] = frozenset({"load_skill", "read_skill_resource"}) @@ -1200,13 +1201,11 @@ def _latest_included_group(messages: list[Message]) -> list[Message]: return _latest_group(messages, include_excluded=False) -def _completed_function_calls( +def _function_call_result_pairs( messages: Sequence[Message], *, - tool_name: str, - include_consumed: bool = False, include_excluded: bool = False, -) -> list[Content]: +) -> list[tuple[Content, Content]]: candidate_messages = ( list(messages) if include_excluded @@ -1215,7 +1214,7 @@ def _completed_function_calls( open_calls_by_call_id: dict[str, list[Content]] = {} open_calls_by_occurrence_id: dict[str, list[Content]] = {} closed_call_identities: set[int] = set() - successful_call_identities: set[int] = set() + pairs: list[tuple[Content, Content]] = [] for message in candidate_messages: for content in message.contents: @@ -1249,26 +1248,45 @@ def _completed_function_calls( if matched_call is None: continue - matched_identity = id(matched_call) - closed_call_identities.add(matched_identity) - if ( - content.exception is None - and content.additional_properties.get(APPROVAL_REJECTION_RESULT_KEY) is not True - ): - successful_call_identities.add(matched_identity) + closed_call_identities.add(id(matched_call)) + pairs.append((matched_call, content)) + + return pairs + +def _completed_function_call_pairs( + messages: Sequence[Message], + *, + tool_name: str, + include_consumed: bool = False, + include_excluded: bool = False, +) -> list[tuple[Content, Content]]: return [ - content - for message in candidate_messages - for content in message.contents + (call, result) + for call, result in _function_call_result_pairs(messages, include_excluded=include_excluded) if ( - content.type == "function_call" - and content.name == tool_name - and id(content) in successful_call_identities - and ( - include_consumed - or message.additional_properties.get(_AGENT_COMPACTION_TRIGGER_CONSUMED_KEY) is not True - ) + call.name == tool_name + and result.exception is None + and result.additional_properties.get(APPROVAL_REJECTION_RESULT_KEY) is not True + and (include_consumed or call.additional_properties.get(_AGENT_COMPACTION_TRIGGER_CONSUMED_KEY) is not True) + ) + ] + + +def _completed_function_calls( + messages: Sequence[Message], + *, + tool_name: str, + include_consumed: bool = False, + include_excluded: bool = False, +) -> list[Content]: + return [ + call + for call, _ in _completed_function_call_pairs( + messages, + tool_name=tool_name, + include_consumed=include_consumed, + include_excluded=include_excluded, ) ] @@ -1316,9 +1334,9 @@ class AgentTriggeredCompactionStrategy: Wrapped strategies may provide ``create_trigger_tool()`` to replace the default ``compact_context`` tool with one specialized for their inputs. - The consumed marker is stored on each trigger-group message rather than on - this strategy instance, so one strategy can be shared safely by concurrent - runs. + The consumed marker is stored on the completed trigger call occurrence + rather than on this strategy instance, so one strategy can be shared safely + by concurrent runs and composed trigger tools do not suppress each other. This strategy requires client-managed history (for example ``store=False``). Chat clients reject it when service-side storage or a service-managed @@ -1427,9 +1445,12 @@ async def __call__(self, messages: list[Message]) -> bool: if any(id(message) in trigger_group_ids for message in group): continue if not any( - message.additional_properties.get(_AGENT_COMPACTION_TRIGGER_CONSUMED_KEY) is True - and not message.additional_properties.get(EXCLUDED_KEY, False) + content.type == "function_call" + and content.name == self.tool.name + and content.additional_properties.get(_AGENT_COMPACTION_TRIGGER_CONSUMED_KEY) is True for message in group + if not message.additional_properties.get(EXCLUDED_KEY, False) + for content in message.contents ): continue if self._is_trigger_only_group(group): @@ -1455,8 +1476,8 @@ async def __call__(self, messages: list[Message]) -> bool: self._restore_trigger_group(trigger_state) raise - for message in trigger_group: - message.additional_properties[_AGENT_COMPACTION_TRIGGER_CONSUMED_KEY] = True + for trigger_call in trigger_calls: + trigger_call.additional_properties[_AGENT_COMPACTION_TRIGGER_CONSUMED_KEY] = True has_unrelated_tool_calls = not self._is_trigger_only_group(trigger_group, include_excluded=True) trigger_message_ids = {id(message) for message in trigger_group} @@ -1541,7 +1562,7 @@ def create_trigger_tool(self) -> FunctionTool: """Create the ``unload_skill`` trigger tool used by the wrapper.""" from ._tools import FunctionTool - def request_skill_unload(skill_name: str | Sequence[str] | None = None) -> str: + def request_skill_unload(skill_name: str | Sequence[str] | None = None) -> list[Content]: """Unload the full instructions and resources for selected skills.""" if skill_name is None: requested_names: Sequence[str] = [] @@ -1556,8 +1577,15 @@ def request_skill_unload(skill_name: str | Sequence[str] | None = None) -> str: normalized_names = [name.strip() for name in requested_names if name.strip()] if normalized_names: formatted_names = ", ".join(f"'{name}'" for name in normalized_names) - return f"Skill(s) {formatted_names} will be unloaded before the next model call." - return "All loaded skills will be unloaded before the next model call." + message = f"Skill(s) {formatted_names} will be unloaded before the next model call." + else: + message = "All loaded skills will be unloaded before the next model call." + return [ + Content.from_text( + message, + additional_properties={_SKILL_UNLOAD_SELECTION_KEY: normalized_names}, + ) + ] return FunctionTool( name=self.tool_name, @@ -1587,7 +1615,7 @@ def request_skill_unload(skill_name: str | Sequence[str] | None = None) -> str: ) def _requested_skill_names(self, messages: list[Message]) -> set[str] | None: - requests = _completed_function_calls( + requests = _completed_function_call_pairs( _latest_group(messages, include_excluded=True), tool_name=self.tool_name, include_consumed=True, @@ -1597,17 +1625,15 @@ def _requested_skill_names(self, messages: list[Message]) -> set[str] | None: return set() requested_names: set[str] = set() - for request in requests: - arguments = request.parse_arguments() or {} - raw_skill_names = arguments.get("skill_name") - if raw_skill_names is None: - skill_names: Sequence[Any] = [] - elif isinstance(raw_skill_names, str): - skill_names = [raw_skill_names] - elif isinstance(raw_skill_names, Sequence) and not isinstance(raw_skill_names, (bytes, bytearray)): - skill_names = cast("Sequence[Any]", raw_skill_names) - else: + for _, result in requests: + selections = [ + item.additional_properties.get(_SKILL_UNLOAD_SELECTION_KEY) + for item in (result.items or []) + if item.type == "text" + ] + if len(selections) != 1 or not isinstance(selections[0], list): return set() + skill_names = cast("list[Any]", selections[0]) if any(not isinstance(skill_name, str) for skill_name in skill_names): return set() normalized_names = {skill_name.strip().casefold() for skill_name in skill_names if skill_name.strip()} @@ -1631,13 +1657,8 @@ def _skill_names_for_removable_group(group: Sequence[Message]) -> set[str] | Non return None calls = tool_calls - completed_call_ids = { - content.call_id - for message in included_messages - for content in message.contents - if content.type == "function_result" and content.call_id - } - if any(not call.call_id or call.call_id not in completed_call_ids for call in calls): + completed_call_identities = {id(call) for call, _ in _function_call_result_pairs(included_messages)} + if any(id(call) not in completed_call_identities for call in calls): return None skill_names: set[str] = set() @@ -1919,7 +1940,11 @@ async def __call__(self, messages: list[Message]) -> bool: for content in msg.contents: if content.type == "function_call" and content.call_id and content.name: call_id_to_name[content.call_id] = content.name - elif content.type == "mcp_server_tool_call" and content.call_id and content.tool_name: + elif ( + content.type in ("mcp_server_tool_call", "search_tool_call") + and content.call_id + and content.tool_name + ): call_id_to_name[content.call_id] = content.tool_name # Collect tool results with the tool name for context. @@ -1938,6 +1963,11 @@ async def __call__(self, messages: list[Message]) -> bool: tool_name = call_id_to_name.get(content.call_id or "", "") label = f"{tool_name}: {result_text}" if tool_name else result_text tool_results.append(label.strip()) + elif content.type == "search_tool_result": + result_text = _tool_result_text(content.result if content.result is not None else content.items) + tool_name = call_id_to_name.get(content.call_id or "", "") + label = f"{tool_name}: {result_text}" if tool_name else result_text + tool_results.append(label.strip()) summary_text = self._summary_text(tool_results) summary_id = f"tool_summary_{group_id}" diff --git a/python/packages/core/tests/core/test_agents.py b/python/packages/core/tests/core/test_agents.py index 76cb0d42df8..a3703420a80 100644 --- a/python/packages/core/tests/core/test_agents.py +++ b/python/packages/core/tests/core/test_agents.py @@ -29,6 +29,7 @@ Content, ContextProvider, Embedding, + FunctionMiddleware, FunctionTool, GeneratedEmbeddings, HistoryProvider, @@ -46,6 +47,7 @@ ToolApprovalMiddleware, ToolResultCompactionStrategy, TruncationStrategy, + UnloadSkillsCompactionStrategy, VectorStoreHistoryProvider, chat_middleware, enqueue_messages, @@ -3508,6 +3510,94 @@ def capture_inner( assert "obsolete planning detail" not in captured_texts[1] +async def test_unload_skills_uses_middleware_rewritten_selection( + chat_client_base: MockBaseChatClient, +) -> None: + strategy = AgentTriggeredCompactionStrategy(UnloadSkillsCompactionStrategy()) + + class RewriteUnloadSelectionMiddleware(FunctionMiddleware): + async def process(self, context: FunctionInvocationContext, call_next: Any) -> None: + if context.function.name == strategy.tool.name: + context.arguments = {"skill_name": "excel"} + await call_next() + + chat_client_base.run_responses = [ + ChatResponse( + messages=Message( + role="assistant", + contents=[ + Content.from_function_call( + call_id="unload-1", + name=strategy.tool.name, + arguments={"skill_name": "pdf"}, + ) + ], + ) + ), + ChatResponse(messages=Message(role="assistant", contents=["done"])), + ] + + captured_calls: list[list[tuple[str | None, Any]]] = [] + original_inner = chat_client_base._inner_get_response + + def capture_inner( + *, messages: MutableSequence[Message], stream: bool, options: dict[str, Any], **kwargs: Any + ) -> Any: + captured_calls.append([ + (content.name, content.parse_arguments()) + for message in messages + for content in message.contents + if content.type == "function_call" + ]) + return original_inner(messages=messages, stream=stream, options=options, **kwargs) + + agent = Agent( + client=chat_client_base, + tools=[strategy.tool], + compaction_strategy=strategy, + middleware=[RewriteUnloadSelectionMiddleware()], + ) + run_messages = [ + Message(role="user", contents=["work"]), + Message( + role="assistant", + contents=[ + Content.from_function_call( + call_id="load-pdf", + name="load_skill", + arguments={"skill_name": "pdf"}, + ) + ], + ), + Message( + role="tool", + contents=[Content.from_function_result(call_id="load-pdf", result="pdf instructions")], + ), + Message( + role="assistant", + contents=[ + Content.from_function_call( + call_id="load-excel", + name="load_skill", + arguments={"skill_name": "excel"}, + ) + ], + ), + Message( + role="tool", + contents=[Content.from_function_result(call_id="load-excel", result="excel instructions")], + ), + Message(role="user", contents=["unload the selected skill"]), + ] + + with patch.object(chat_client_base, "_inner_get_response", side_effect=capture_inner): + result = await agent.run(run_messages) + + assert result.text == "done" + assert ("load_skill", {"skill_name": "pdf"}) in captured_calls[1] + assert ("load_skill", {"skill_name": "excel"}) not in captured_calls[1] + + async def test_chat_agent_uses_client_compaction_defaults_when_agent_unset(chat_client_base: Any) -> None: captured_roles: list[list[str]] = [] original_inner = chat_client_base._inner_get_response diff --git a/python/packages/core/tests/core/test_compaction.py b/python/packages/core/tests/core/test_compaction.py index 7e6db7bd17a..4e55f07fb34 100644 --- a/python/packages/core/tests/core/test_compaction.py +++ b/python/packages/core/tests/core/test_compaction.py @@ -39,6 +39,7 @@ included_token_count, ) from agent_framework._compaction import ( + _SKILL_UNLOAD_SELECTION_KEY, _completed_function_calls, _format_summary_message, _select_summary_input_groups, @@ -128,6 +129,23 @@ def _tool_result(call_id: str, result: str, message_id: str | None = None) -> Me ) +def _skill_unload_result(call_id: str, skill_names: list[str]) -> Message: + return Message( + role="tool", + contents=[ + Content.from_function_result( + call_id=call_id, + result=[ + Content.from_text( + "Skills will be unloaded before the next model call.", + additional_properties={_SKILL_UNLOAD_SELECTION_KEY: skill_names}, + ) + ], + ) + ], + ) + + def _group_id(message: Message) -> str | None: annotation = message.additional_properties.get(GROUP_ANNOTATION_KEY) if not isinstance(annotation, dict): @@ -1076,7 +1094,42 @@ async def __call__(self, messages: list[Message]) -> bool: with pytest.raises(RuntimeError, match="compaction failed"): await strategy(messages) - assert "_agent_compaction_trigger_consumed" not in trigger_message.additional_properties + assert "_agent_compaction_trigger_consumed" not in trigger_message.contents[0].additional_properties + + +@pytest.mark.parametrize("reverse_order", [False, True]) +async def test_composed_trigger_tools_each_execute_once(reverse_order: bool) -> None: + first_inner = _RecordingCompactionStrategy(changed=False) + second_inner = _RecordingCompactionStrategy(changed=False) + first = AgentTriggeredCompactionStrategy(first_inner, tool_name="compact_first") + second = AgentTriggeredCompactionStrategy(second_inner, tool_name="compact_second") + messages = [ + Message(role="user", contents=["work"]), + Message( + role="assistant", + contents=[ + Content.from_function_call(call_id="first", name=first.tool.name, arguments={}), + Content.from_function_call(call_id="second", name=second.tool.name, arguments={}), + ], + ), + Message( + role="tool", + contents=[ + Content.from_function_result(call_id="first", result="First compaction requested."), + Content.from_function_result(call_id="second", result="Second compaction requested."), + ], + ), + ] + strategies = [second, first] if reverse_order else [first, second] + + assert await strategies[0](messages) is False + assert await strategies[1](messages) is False + assert first_inner.call_count == 1 + assert second_inner.call_count == 1 + assert await first(messages) is False + assert await second(messages) is False + assert first_inner.call_count == 1 + assert second_inner.call_count == 1 async def test_agent_triggered_compaction_ignores_trigger_outside_latest_group() -> None: @@ -1203,7 +1256,7 @@ async def test_unload_skills_compaction_removes_named_skill_only() -> None: _assistant_named_function_call("load-excel", "load_skill", {"skill_name": "excel"}), _tool_result("load-excel", "excel instructions"), _assistant_named_function_call("unload-pdf", strategy.tool.name, {"skill_name": "pdf"}), - _tool_result("unload-pdf", "Unload requested for skill 'pdf'."), + _skill_unload_result("unload-pdf", ["pdf"]), ] changed = await strategy(messages) @@ -1231,7 +1284,7 @@ async def test_unload_skills_compaction_removes_multiple_named_skills() -> None: strategy.tool.name, {"skill_name": ["pdf", "excel"]}, ), - _tool_result("unload-selected", "Unload requested for pdf and excel."), + _skill_unload_result("unload-selected", ["pdf", "excel"]), ] changed = await strategy(messages) @@ -1278,7 +1331,7 @@ async def test_unload_skills_compaction_empty_name_removes_all_skill_context() - ), _tool_result("resource-excel", "excel reference"), _assistant_named_function_call("unload-all", strategy.tool.name, {"skill_name": " "}), - _tool_result("unload-all", "Unload requested for all skills."), + _skill_unload_result("unload-all", []), ] changed = await strategy(messages) @@ -1322,7 +1375,7 @@ async def test_unload_skills_compaction_preserves_script_results_and_mixed_group ), _tool_result("run-pdf", "converted.pdf"), _assistant_named_function_call("unload-pdf", strategy.tool.name, {"skill_name": "pdf"}), - _tool_result("unload-pdf", "Unload requested for skill 'pdf'."), + _skill_unload_result("unload-pdf", ["pdf"]), ] changed = await strategy(messages) @@ -1347,7 +1400,7 @@ async def test_unload_skills_compaction_preserves_group_with_hosted_tool_call() ), _tool_result("load-pdf", "pdf instructions"), _assistant_named_function_call("unload-pdf", strategy.tool.name, {"skill_name": "pdf"}), - _tool_result("unload-pdf", "Skill 'pdf' will be unloaded."), + _skill_unload_result("unload-pdf", ["pdf"]), ] changed = await strategy(messages) @@ -1386,7 +1439,47 @@ async def test_unload_skills_compaction_preserves_group_with_search_tool_call() ], ), _assistant_named_function_call("unload-pdf", strategy.tool.name, {"skill_name": "pdf"}), - _tool_result("unload-pdf", "Skill 'pdf' will be unloaded."), + _skill_unload_result("unload-pdf", ["pdf"]), + ] + + changed = await strategy(messages) + + assert changed is False + assert all(message.additional_properties.get(EXCLUDED_KEY) is not True for message in messages) + + +async def test_unload_skills_compaction_preserves_reused_id_group_with_unresolved_call() -> None: + strategy = AgentTriggeredCompactionStrategy(UnloadSkillsCompactionStrategy()) + messages = [ + Message( + role="assistant", + contents=[ + Content.from_function_call( + id="load-pdf", + call_id="shared", + name="load_skill", + arguments={"skill_name": "pdf"}, + ), + Content.from_function_call( + id="load-excel", + call_id="shared", + name="load_skill", + arguments={"skill_name": "excel"}, + ), + ], + ), + Message( + role="tool", + contents=[ + Content.from_function_result( + id="load-pdf", + call_id="shared", + result="pdf instructions", + ) + ], + ), + _assistant_named_function_call("unload-all", strategy.tool.name), + _skill_unload_result("unload-all", []), ] changed = await strategy(messages) @@ -2220,6 +2313,42 @@ async def test_apply_compaction_logs_changed_context_without_content(caplog: Any # --- ToolResultCompactionStrategy tests --- +async def test_tool_result_compaction_preserves_hosted_search_result() -> None: + messages = [ + Message(role="user", contents=["find docs"]), + Message( + role="assistant", + contents=[ + Content.from_search_tool_call( + call_id="search-1", + tool_name="web_search", + arguments={"query": "agent framework"}, + ) + ], + ), + Message( + role="tool", + contents=[ + Content.from_search_tool_result( + call_id="search-1", + tool_name="web_search", + result="Found the Agent Framework documentation.", + ) + ], + ), + Message(role="assistant", contents=["done"]), + ] + annotate_message_groups(messages) + + changed = await ToolResultCompactionStrategy(keep_last_tool_call_groups=0)(messages) + + assert changed is True + summaries = [message.text or "" for message in messages if (message.text or "").startswith("[Tool results:")] + assert len(summaries) == 1 + assert "web_search: Found the Agent Framework documentation." in summaries[0] + assert "no results" not in summaries[0] + + async def test_tool_result_compaction_collapses_old_groups_into_summary() -> None: """Old tool-call groups are collapsed into summary messages, newest kept.""" messages = [