diff --git a/studio/backend/core/inference/context_compaction.py b/studio/backend/core/inference/context_compaction.py index 18c36bb07c..f4a52c675a 100644 --- a/studio/backend/core/inference/context_compaction.py +++ b/studio/backend/core/inference/context_compaction.py @@ -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] diff --git a/studio/backend/tests/test_context_compaction.py b/studio/backend/tests/test_context_compaction.py index c42b39d9ad..aa629093d7 100644 --- a/studio/backend/tests/test_context_compaction.py +++ b/studio/backend/tests/test_context_compaction.py @@ -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