diff --git a/python/packages/core/agent_framework/_compaction.py b/python/packages/core/agent_framework/_compaction.py index 5140f6fc17..59abb10a46 100644 --- a/python/packages/core/agent_framework/_compaction.py +++ b/python/packages/core/agent_framework/_compaction.py @@ -650,6 +650,17 @@ def _included_group_ids(messages: list[Message], ordered_group_ids: list[str]) - return included_ids +def _minimum_retained_group_ids( + messages: list[Message], ordered_group_ids: list[str], kinds: dict[str, GroupKind] +) -> set[str]: + """Return the group ids needed to keep a compaction projection non-empty.""" + included_ids = _included_group_ids(messages, ordered_group_ids) + non_system_ids = [group_id for group_id in included_ids if kinds.get(group_id) != "system"] + if non_system_ids: + return {non_system_ids[-1]} + return {included_ids[-1]} if included_ids else set() + + def _count_included_messages(messages: list[Message]) -> int: return len(included_messages(messages)) @@ -665,8 +676,9 @@ class TruncationStrategy: groups (never partial tool-call groups). The metric is: - token count when ``tokenizer`` is provided - included message count when ``tokenizer`` is not provided - Compaction triggers when the metric exceeds ``max_n`` and trims to - ``compact_to``. + Compaction triggers when the metric exceeds ``max_n`` and trims toward + ``compact_to``. The minimum retained group is never excluded, so the + result may remain above ``compact_to`` when that group alone exceeds it. """ def __init__( @@ -713,6 +725,7 @@ async def __call__(self, messages: list[Message]) -> bool: protected_ids: set[str] = set() if self.preserve_system: protected_ids = {group_id for group_id in ordered_group_ids if kinds.get(group_id) == "system"} + protected_ids.update(_minimum_retained_group_ids(messages, ordered_group_ids, kinds)) changed = False for group_id in ordered_group_ids: @@ -1146,7 +1159,9 @@ class TokenBudgetComposedStrategy: Strategies run in the provided order over shared message annotations. After each step, token counts are refreshed. If no strategy reaches budget, a deterministic fallback excludes oldest groups (and finally anchors when - necessary) to enforce the limit. + necessary) to enforce the limit, while retaining the minimum group needed + for a non-empty projection. The result may remain above the budget when + that group alone exceeds it. """ def __init__( @@ -1191,8 +1206,9 @@ async def __call__(self, messages: list[Message]) -> bool: ordered_group_ids = annotate_message_groups(messages) grouped = _group_messages_by_id(messages) kinds = _group_kind_map(messages) + minimum_retained_ids = _minimum_retained_group_ids(messages, ordered_group_ids, kinds) for group_id in ordered_group_ids: - if kinds.get(group_id) == "system": + if kinds.get(group_id) == "system" or group_id in minimum_retained_ids: continue for message in grouped.get(group_id, []): changed = set_excluded(message, excluded=True, reason="token_budget_fallback") or changed @@ -1203,7 +1219,7 @@ async def __call__(self, messages: list[Message]) -> bool: # Strict budget enforcement fallback: if anchors alone exceed budget, exclude remaining groups. for group_id in ordered_group_ids: - if kinds.get(group_id) != "system": + if kinds.get(group_id) != "system" or group_id in minimum_retained_ids: continue for message in grouped.get(group_id, []): changed = set_excluded(message, excluded=True, reason="token_budget_fallback_strict") or changed diff --git a/python/packages/core/tests/core/test_compaction.py b/python/packages/core/tests/core/test_compaction.py index 25ab21a9de..4507d39946 100644 --- a/python/packages/core/tests/core/test_compaction.py +++ b/python/packages/core/tests/core/test_compaction.py @@ -348,12 +348,12 @@ async def test_truncation_strategy_compacts_when_token_limit_exceeded() -> None: tokenizer = CharacterEstimatorTokenizer() messages = [ Message(role="system", contents=["you are helpful"]), - Message(role="user", contents=["u1 " * 200]), - Message(role="assistant", contents=["a1 " * 200]), + Message(role="user", contents=["u1 " * 5]), + Message(role="assistant", contents=["a1 " * 5]), ] strategy = TruncationStrategy( max_n=80, - compact_to=40, + compact_to=70, tokenizer=tokenizer, preserve_system=True, ) @@ -364,7 +364,23 @@ async def test_truncation_strategy_compacts_when_token_limit_exceeded() -> None: assert changed is True projected = included_messages(messages) assert projected[0].role == "system" - assert included_token_count(messages) <= 40 + assert included_token_count(messages) <= 70 + + +async def test_truncation_strategy_keeps_latest_group_when_it_exceeds_target() -> None: + tokenizer = CharacterEstimatorTokenizer() + messages = [Message(role="user", contents=["latest " * 200])] + strategy = TruncationStrategy( + max_n=20, + compact_to=10, + tokenizer=tokenizer, + ) + annotate_message_groups(messages, tokenizer=tokenizer) + + changed = await strategy(messages) + + assert changed is False + assert included_messages(messages) == messages def test_truncation_strategy_validates_token_targets() -> None: @@ -539,11 +555,11 @@ async def test_summarization_strategy_returns_false_when_summary_is_empty( async def test_token_budget_composed_strategy_meets_budget_or_falls_back() -> None: messages = [ Message(role="system", contents=["system"]), - Message(role="user", contents=["user " * 200]), - Message(role="assistant", contents=["assistant " * 200]), + Message(role="user", contents=["user " * 10]), + Message(role="assistant", contents=["assistant " * 2]), ] strategy = TokenBudgetComposedStrategy( - token_budget=20, + token_budget=70, tokenizer=CharacterEstimatorTokenizer(), strategies=[SlidingWindowStrategy(keep_last_groups=1)], ) @@ -551,7 +567,39 @@ async def test_token_budget_composed_strategy_meets_budget_or_falls_back() -> No changed = await strategy(messages) assert changed is True - assert included_token_count(messages) <= 20 + assert included_token_count(messages) <= 70 + + +async def test_token_budget_composed_strategy_keeps_latest_group_when_all_groups_exceed_budget() -> None: + messages = [ + Message(role="system", contents=["system " * 100]), + Message(role="user", contents=["latest " * 100]), + ] + strategy = TokenBudgetComposedStrategy( + token_budget=1, + tokenizer=CharacterEstimatorTokenizer(), + strategies=[], + ) + + changed = await strategy(messages) + + assert changed is True + projected = included_messages(messages) + assert projected == [messages[-1]] + + +async def test_token_budget_composed_strategy_keeps_last_system_group_when_no_user_group_exists() -> None: + messages = [Message(role="system", contents=["system " * 100])] + strategy = TokenBudgetComposedStrategy( + token_budget=1, + tokenizer=CharacterEstimatorTokenizer(), + strategies=[], + ) + + changed = await strategy(messages) + + assert changed is False + assert included_messages(messages) == messages class _ExcludeOldestNonSystem: