Studio: handle keep_recent=0 + anchored pair cleanup

Two related corner cases on the compaction module, both regression-
tested:

1) keep_recent=0 collected ONE group instead of zero. The recent-
   window loop tested its limit AFTER appending, so the first
   iteration always added one group and broke. Flip the check to
   BEFORE appending so a caller asking for the minimal-anchor regime
   actually gets only the anchor set (system + first-user +
   multimodal) and drops everything else.

2) Pair cleanup could drop an ANCHORED message. When all tool-role
   messages paired with an assistant tool_call message were dropped,
   the cleanup rule forced the assistant to drop too, even if that
   assistant was an anchor (multimodal carrying tool_calls, rare but
   valid for vision-capable tool-calling models). Mirror invariant:
   if the assistant is anchored leak an orphan tool_call shape
   rather than violate "multimodal turns are never dropped". Same
   for the inverse direction: a dropped assistant tool_call message
   should not drag anchored tool-role partners into the drop set.

Tests: 19 -> 21. Add TestKeepRecentZero and
TestAnchoredMultimodalPairCleanup classes covering the two fixes.
This commit is contained in:
Daniel Han 2026-05-24 12:45:16 +00:00
commit 41192f0452
2 changed files with 103 additions and 6 deletions

View file

@ -195,17 +195,22 @@ class SlidingWindowCompact(CompactStrategy):
group_id[ti] = g
# The "recent window" is the last ``keep_recent`` distinct
# groups encountered scanning from the end.
# groups encountered scanning from the end. ``keep_recent == 0``
# must collect ZERO groups so the caller can drop everything
# outside the anchor set (system + first user + multimodal).
# The pre-fix loop tested the limit AFTER appending and so
# always preserved at least one group even when keep_recent
# was 0; flip the check to BEFORE appending so the bound holds.
recent_groups: list[int] = []
seen_groups: set[int] = set()
for i in range(len(messages) - 1, -1, -1):
if len(recent_groups) >= self.keep_recent:
break
g = group_id[i]
if g in seen_groups:
continue
seen_groups.add(g)
recent_groups.append(g)
if len(recent_groups) >= self.keep_recent:
break
recent_groups_set = set(recent_groups)
# Decide drop set: every index whose group is NOT in the recent
@ -230,13 +235,23 @@ class SlidingWindowCompact(CompactStrategy):
# When dropping an assistant-with-tool-calls message we must
# also drop the matching tool-role messages (and vice versa)
# so the chat template stays valid. Iterate pair_map once.
# Anchor indices stay regardless: dragging an anchored
# multimodal assistant or first-user message into the drop
# set just because its tool-pair partner was dropped would
# violate the structural invariant the anchor set exists to
# enforce, and llama-server would 400 on the resulting
# template (a tool message whose tool_call_id has no
# surviving assistant tool_calls entry).
for asst_idx, tool_idxs in pair_map.items():
if asst_idx in dropped:
dropped.update(tool_idxs)
dropped.update(t for t in tool_idxs if t not in anchor_idx)
elif tool_idxs and tool_idxs <= dropped:
# All matching tool messages were dropped: drop the
# assistant tool-call message too.
dropped.add(asst_idx)
# assistant tool-call message too -- unless it's
# anchored, in which case we'd rather leak an
# orphan tool_call shape than violate the invariant.
if asst_idx not in anchor_idx:
dropped.add(asst_idx)
return [m for i, m in enumerate(messages) if i not in dropped]

View file

@ -277,6 +277,88 @@ class TestStrategyRegistry:
assert isinstance(get_strategy("totally-made-up"), NoCompact)
class TestKeepRecentZero:
"""Regression: keep_recent=0 must collect ZERO recent groups, not 1.
Pre-fix, the loop tested the limit AFTER appending, so the first
iteration always added one group and broke. Trim with keep_recent=0
should drop everything outside the anchor set (system + first-user
+ multimodal). This matters when a caller intentionally wants only
the anchors to survive (extreme-pressure regime).
"""
def test_keep_recent_zero_keeps_only_anchors(self):
msgs = [
_msg("system", "sys"),
_msg("user", "task"),
_long("assistant", 4000),
_long("user", 4000),
_long("assistant", 4000),
_long("user", 4000),
]
out = SlidingWindowCompact(keep_recent = 0).compact(
msgs, budget_tokens = 50
)
roles = [m["role"] for m in out]
# System + first-user only.
assert roles == ["system", "user"]
assert out[1]["content"].startswith("task")
class TestAnchoredMultimodalPairCleanup:
"""Regression: an anchored multimodal assistant whose paired tool
messages all get dropped must survive pair-cleanup. Without this
guard the pair_map's "all tool messages dropped -> drop the
assistant too" rule would yank a multimodal anchor out from under
the structural invariant that says multimodal turns are never
dropped, and llama-server would 400 the resulting template.
"""
def test_anchored_multimodal_assistant_survives_pair_cleanup(self):
# Multimodal assistant carrying tool_calls (rare but valid:
# vision models can call tools while emitting image parts).
multimodal_asst = {
"role": "assistant",
"content": [
{"type": "text", "text": "describe and call tool"},
{"type": "image_url", "image_url": {"url": "x"}},
],
"tool_calls": [
{
"id": "call_mm",
"type": "function",
"function": {
"name": "web_search",
"arguments": '{"q":"y"}',
},
}
],
}
msgs = [
_msg("system", "sys"),
_msg("user", "task"),
multimodal_asst,
_msg(
"tool",
content = "result for y" * 200,
tool_call_id = "call_mm",
name = "web_search",
),
_long("assistant", 4000),
_long("user", 4000),
_long("assistant", 4000),
_long("user", 4000),
]
out = SlidingWindowCompact(keep_recent = 1).compact(
msgs, budget_tokens = 50
)
# Anchored multimodal assistant must still be there.
assert any(
m.get("role") == "assistant" and isinstance(m.get("content"), list)
for m in out
), "multimodal assistant got dropped by pair_map cleanup"
class TestConstructorValidation:
def test_negative_keep_recent_raises(self):
import pytest