diff --git a/studio/backend/core/inference/context_compaction.py b/studio/backend/core/inference/context_compaction.py new file mode 100644 index 0000000000..f8ae4c3289 --- /dev/null +++ b/studio/backend/core/inference/context_compaction.py @@ -0,0 +1,328 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 +# +# SlidingWindowCompact is adapted from forge +# (https://github.com/antoinezambelli/forge), Copyright (c) 2025-2026 +# Antoine Zambelli, used under the MIT License. + +"""Context-window compaction for OpenAI-style chat messages. + +Long-running Studio Chat sessions can outgrow a model's context window. +This module ships a strategy that trims older messages from the prompt +sent to the model while preserving the persisted transcript shown in +the UI. The strategy returns a NEW list; ``messages`` is never mutated. + +Invariants preserved by all strategies: + +1. The system message (when present at index 0) is never dropped. +2. The first user message (when present) is never dropped: it carries + the task prompt the rest of the conversation references. +3. Tool-call <-> tool-result pair linkage stays valid. An assistant + message that carries ``tool_calls`` and the matching tool-role + messages are kept or dropped as a unit. Dropping one side leaves + the OpenAI chat template invalid and llama-server returns 400. +4. Multimodal turns (any message whose ``content`` is a list of parts) + are treated as non-droppable. There is no tested compacted-media + representation today; the strategies leave such turns intact. + +Only one strategy is shipped here. ``TieredCompact`` from forge depends +on per-message metadata tags Studio's flat message dicts do not carry; +it will land in a follow-up once the message model gets the tags. +""" + +from __future__ import annotations + +from abc import ABC, abstractmethod +from typing import Any + + +# Conservative char-to-token ratio so compaction fires early. +_CHARS_PER_TOKEN = 4 + + +def estimate_tokens(messages: list[dict]) -> int: + """Rough token count via 4-char-per-token heuristic. + + Counts visible content, multimodal text parts, ``compaction`` + summary text, and serialized tool_calls arguments. + """ + total_chars = 0 + for m in messages: + content = m.get("content") + if isinstance(content, str): + total_chars += len(content) + elif isinstance(content, list): + for part in content: + if not isinstance(part, dict): + continue + t = part.get("text") + if isinstance(t, str): + total_chars += len(t) + # Count compaction summary text. + if part.get("type") == "compaction": + summary = part.get("content") + if isinstance(summary, str): + total_chars += len(summary) + tcs = m.get("tool_calls") + if isinstance(tcs, list): + for tc in tcs: + # Skip malformed pre-pydantic entries. + if not isinstance(tc, dict): + continue + fn = tc.get("function") + args = (fn or {}).get("arguments") if isinstance(fn, dict) else None + if isinstance(args, str): + total_chars += len(args) + # Ceil-divide: floor would let just-over-budget prompts bypass compaction. + return -(-total_chars // _CHARS_PER_TOKEN) + + +# Text-only parts (don't anchor); ``compaction`` is Anthropic round-trip +# state so pinning it would defeat compaction. +_TEXT_ONLY_PART_TYPES = {"text", "compaction"} + + +def _is_multimodal(msg: dict) -> bool: + content = msg.get("content") + if not isinstance(content, list): + return False + for part in content: + # Unknown shape => conservatively treat as multimodal. + if not isinstance(part, dict): + return True + if part.get("type") not in _TEXT_ONLY_PART_TYPES: + return True + return False + + +def _assistant_tool_call_ids(msg: dict) -> set[str]: + """Return tool_call ids on an assistant message; skip malformed entries.""" + out: set[str] = set() + tcs = msg.get("tool_calls") + if isinstance(tcs, list): + for tc in tcs: + if not isinstance(tc, dict): + continue + tcid = tc.get("id") + if isinstance(tcid, str) and tcid: + out.add(tcid) + return out + + +def _pair_linked_indices(messages: list[dict]) -> dict[int, set[int]]: + """Map asst index -> indices of its tool-role follow-ups so the + compactor drops or keeps the group as a unit. + """ + out: dict[int, set[int]] = {} + # OpenAI schema: tool messages must follow their asst directly. + # ANY non-tool boundary (user, system, another asst) clears the + # pending window so a stale asst can't snap to a later-turn tool. + pending_ids: dict[str, int] = {} + for i, m in enumerate(messages): + role = m.get("role") + if role == "assistant": + pending_ids.clear() + ids = _assistant_tool_call_ids(m) + out.setdefault(i, set()) + for tid in ids: + pending_ids[tid] = i + elif role == "tool": + tcid = m.get("tool_call_id") + if isinstance(tcid, str) and tcid in pending_ids: + out.setdefault(pending_ids[tcid], set()).add(i) + else: + pending_ids.clear() + return out + + +class CompactStrategy(ABC): + """Interface for context-compaction strategies.""" + + @abstractmethod + def compact(self, messages: list[dict], budget_tokens: int) -> list[dict]: + """Return a (possibly shorter) list of messages within + ``budget_tokens``. Returns ``messages`` unchanged when no + compaction is needed or possible. + """ + ... + + +class NoCompact(CompactStrategy): + """Passthrough strategy. Returns ``messages`` unchanged.""" + + def compact(self, messages: list[dict], budget_tokens: int) -> list[dict]: + return list(messages) + + +class SlidingWindowCompact(CompactStrategy): + """Keep system, first user, multimodal, and last ``keep_recent`` + groups. Asst+tools form one group. No-op when within budget. + """ + + def __init__( + self, + keep_recent: int = 2, + compact_threshold: float = 0.85, + ) -> None: + if keep_recent < 0: + raise ValueError("keep_recent must be >= 0") + if not (0.0 < compact_threshold <= 1.0): + raise ValueError("compact_threshold must be in (0, 1]") + self.keep_recent = keep_recent + self.compact_threshold = compact_threshold + + def compact(self, messages: list[dict], budget_tokens: int) -> list[dict]: + if budget_tokens <= 0 or not messages: + return list(messages) + + threshold = int(budget_tokens * self.compact_threshold) + if estimate_tokens(messages) <= threshold: + return list(messages) + + # Anchor indices that may never be dropped. + anchor_idx: set[int] = set() + # System message at index 0. + if messages and messages[0].get("role") == "system": + anchor_idx.add(0) + # First user message (the task prompt). + for i, m in enumerate(messages): + if m.get("role") == "user": + anchor_idx.add(i) + break + # Multimodal turns. + for i, m in enumerate(messages): + if _is_multimodal(m): + anchor_idx.add(i) + + # Group asst tool-call with its tool responses so they drop together. + pair_map = _pair_linked_indices(messages) + # If any pair member is anchored, anchor the whole group; otherwise + # an anchored multimodal asst/tool could orphan its partner. + for asst_idx, tool_idxs in pair_map.items(): + if not tool_idxs: + continue + pair_idxs = {asst_idx, *tool_idxs} + if pair_idxs & anchor_idx: + anchor_idx |= pair_idxs + group_id: list[int] = list(range(len(messages))) + next_g = len(messages) + for asst_idx, tool_idxs in pair_map.items(): + if not tool_idxs: + continue + g = next_g + next_g += 1 + group_id[asst_idx] = g + for ti in tool_idxs: + group_id[ti] = g + + # Last ``keep_recent`` distinct groups from the end. The limit + # check runs BEFORE appending so ``keep_recent == 0`` keeps zero. + 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) + recent_groups_set = set(recent_groups) + + # Droppable = non-anchor indices outside the recent window. + droppable: list[int] = [] + for i in range(len(messages)): + if i in anchor_idx: + continue + if group_id[i] in recent_groups_set: + continue + droppable.append(i) + + # Drop oldest-first until under threshold or nothing left. + dropped: set[int] = set() + for i in droppable: + kept = [m for j, m in enumerate(messages) if j not in dropped and j != i] + if estimate_tokens(kept) <= threshold: + dropped.add(i) + break + dropped.add(i) + + # Drop the partner when one side of a pair is dropped, except + # anchored asst whose tools all died -- keep their content but + # strip the dangling tool_calls in the final pass. + rewrite_strip_tool_calls: set[int] = set() + for asst_idx, tool_idxs in pair_map.items(): + if asst_idx in dropped: + dropped.update(t for t in tool_idxs if t not in anchor_idx) + elif tool_idxs and tool_idxs <= dropped: + if asst_idx not in anchor_idx: + dropped.add(asst_idx) + else: + rewrite_strip_tool_calls.add(asst_idx) + + # Two-pass invariant sweep: + # 1) drop tools whose tcid has no prior surviving asst + # (anchored ones too -- pair-validity beats anchor rule); + # 2) recompute responded_ids from survivors, then mark assts + # with unanswered tool_calls for strip. + # Order matters: a stale orphan tcid in responded_ids would let + # the asst keep dangling tool_calls (OpenAI 400). + seen_ids: set[str] = set() + for i, m in enumerate(messages): + if i in dropped: + continue + if m.get("role") == "assistant": + seen_ids |= _assistant_tool_call_ids(m) + elif m.get("role") == "tool": + tcid = m.get("tool_call_id") + if isinstance(tcid, str) and tcid and tcid not in seen_ids: + dropped.add(i) + responded_ids: set[str] = set() + for i, m in enumerate(messages): + if i in dropped: + continue + if m.get("role") == "tool": + tcid = m.get("tool_call_id") + if isinstance(tcid, str) and tcid: + responded_ids.add(tcid) + for i, m in enumerate(messages): + if i in dropped: + continue + if m.get("role") == "assistant": + ids = _assistant_tool_call_ids(m) + if ids and not (ids <= responded_ids): + rewrite_strip_tool_calls.add(i) + + out: list[dict] = [] + for i, m in enumerate(messages): + if i in dropped: + continue + if i in rewrite_strip_tool_calls: + # Keep content but strip dangling tool_calls. + kept_tcs = [ + tc + for tc in (m.get("tool_calls") or []) + if isinstance(tc, dict) + and isinstance(tc.get("id"), str) + and tc["id"] in responded_ids + ] + copy = dict(m) + if kept_tcs: + copy["tool_calls"] = kept_tcs + else: + copy.pop("tool_calls", None) + out.append(copy) + else: + out.append(m) + return out + + +_STRATEGIES: dict[str, CompactStrategy] = { + "none": NoCompact(), + "sliding": SlidingWindowCompact(), +} + + +def get_strategy(name: str) -> CompactStrategy: + """Return strategy by name; unknown names fall back to ``NoCompact``.""" + return _STRATEGIES.get(name, _STRATEGIES["none"]) diff --git a/studio/backend/tests/test_context_compaction.py b/studio/backend/tests/test_context_compaction.py new file mode 100644 index 0000000000..9a95ca5561 --- /dev/null +++ b/studio/backend/tests/test_context_compaction.py @@ -0,0 +1,377 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Tests for the context-compaction module. + +Coverage focuses on the structural invariants the strategies promise: + +* System and first-user messages are never dropped. +* Tool-call <-> tool-result pair linkage stays valid. +* Multimodal turns are never dropped. +* Persisted history is not mutated; the strategy returns a new list. +* ``NoCompact`` is a true passthrough. +* ``SlidingWindowCompact`` is a no-op when the budget already fits. +""" + +from core.inference.context_compaction import ( + NoCompact, + SlidingWindowCompact, + estimate_tokens, + get_strategy, +) + + +def _msg( + role, + content = "", + tool_calls = None, + tool_call_id = None, + name = None, +): + m = {"role": role} + if content is not None: + m["content"] = content + if tool_calls is not None: + m["tool_calls"] = tool_calls + if tool_call_id is not None: + m["tool_call_id"] = tool_call_id + if name is not None: + m["name"] = name + return m + + +def _long( + role, + length, + *, + content_prefix = "x", +): + return _msg(role, content_prefix * length) + + +class TestEstimateTokens: + def test_empty_returns_zero(self): + assert estimate_tokens([]) == 0 + + def test_simple_string_content(self): + # 4 chars per token: "abcd" -> 1. + msgs = [_msg("user", "abcd")] + assert estimate_tokens(msgs) == 1 + + def test_non_multiple_of_4_rounds_up(self): + # Regression: floor (`// 4`) would estimate 1 for "abcde" and + # let an over-budget message slip past the threshold check. + # Ceil keeps the heuristic conservative. + msgs = [_msg("user", "abcde")] # 5 chars + assert estimate_tokens(msgs) == 2 + + def test_single_char_rounds_up_to_one_token(self): + msgs = [_msg("user", "a")] + assert estimate_tokens(msgs) == 1 + + def test_multimodal_text_part_counts(self): + msgs = [ + { + "role": "user", + "content": [ + {"type": "text", "text": "abcdefgh"}, + {"type": "image_url", "image_url": {"url": "x"}}, + ], + } + ] + # 8 chars text, image URL ignored. + assert estimate_tokens(msgs) == 2 + + def test_tool_call_arguments_count(self): + msgs = [ + _msg( + "assistant", + content = "", + tool_calls = [ + { + "id": "c1", + "type": "function", + "function": {"name": "web_search", "arguments": '{"q":"hi"}'}, + } + ], + ) + ] + # 10 char arguments -> 2 tokens. + assert estimate_tokens(msgs) >= 2 + + +class TestNoCompact: + def test_passthrough_returns_copy(self): + msgs = [_msg("user", "hi")] + out = NoCompact().compact(msgs, budget_tokens = 10) + assert out == msgs + # Mutating the result does not affect the original. + out.append(_msg("assistant", "new")) + assert len(msgs) == 1 + + +class TestSlidingWindowUnderBudget: + def test_no_op_when_already_fits(self): + msgs = [ + _msg("system", "be concise"), + _msg("user", "hi"), + _msg("assistant", "bye"), + ] + out = SlidingWindowCompact(keep_recent = 2).compact(msgs, budget_tokens = 1000) + assert out == msgs + + def test_no_op_when_budget_is_zero(self): + # A zero/negative budget collapses to no-op (defensive). + msgs = [_msg("user", "abcd" * 1000)] + out = SlidingWindowCompact(keep_recent = 2).compact(msgs, budget_tokens = 0) + assert out == msgs + + +class TestSlidingWindowInvariants: + def _make_long_chat( + self, + n_turns, + length_per_turn = 1000, + ): + msgs = [_msg("system", "system prompt")] + msgs.append(_msg("user", "the original task: " + "x" * length_per_turn)) + # Alternating assistant/user follow-ups. + for i in range(n_turns): + msgs.append(_long("assistant", length_per_turn)) + msgs.append(_long("user", length_per_turn)) + return msgs + + def test_keeps_system_and_first_user(self): + msgs = self._make_long_chat(n_turns = 20) + out = SlidingWindowCompact(keep_recent = 2).compact(msgs, budget_tokens = 200) + assert out[0]["role"] == "system" + # First user message must survive. + assert any(m.get("role") == "user" and "original task" in m.get("content", "") for m in out) + + def test_keeps_last_n_turns(self): + msgs = self._make_long_chat(n_turns = 20) + out = SlidingWindowCompact(keep_recent = 3).compact(msgs, budget_tokens = 200) + # The last assistant and user pair must survive. + assert out[-1] == msgs[-1] + assert out[-2] == msgs[-2] + + def test_does_not_mutate_input(self): + msgs = self._make_long_chat(n_turns = 10) + snap = [dict(m) for m in msgs] + _ = SlidingWindowCompact(keep_recent = 2).compact(msgs, budget_tokens = 100) + assert msgs == snap + + def test_multimodal_turn_never_dropped(self): + msgs = [ + _msg("system", "sys"), + _msg("user", "first task"), + _long("assistant", 1000), + _long("user", 1000), + { + "role": "user", + "content": [ + {"type": "text", "text": "look at this"}, + {"type": "image_url", "image_url": {"url": "x"}}, + ], + }, + _long("assistant", 1000), + _long("user", 1000), + ] + out = SlidingWindowCompact(keep_recent = 2).compact(msgs, budget_tokens = 200) + # The multimodal turn must still be present. + assert any(isinstance(m.get("content"), list) for m in out) + + +class TestSlidingWindowToolPairs: + def _make_chat_with_tool_pair(self): + return [ + _msg("system", "sys"), + _msg("user", "first task"), + _long("assistant", 800), + _long("user", 800), + # The tool pair we want to test linkage on. + _msg( + "assistant", + content = "", + tool_calls = [ + { + "id": "call_42", + "type": "function", + "function": { + "name": "web_search", + "arguments": '{"q":"x"}', + }, + } + ], + ), + _msg( + "tool", + content = "result for x", + tool_call_id = "call_42", + name = "web_search", + ), + _long("assistant", 800), + _long("user", 800), + _long("assistant", 800), + _long("user", 800), + ] + + def test_tool_pair_dropped_together(self): + msgs = self._make_chat_with_tool_pair() + # Force aggressive compaction. + out = SlidingWindowCompact(keep_recent = 2).compact(msgs, budget_tokens = 50) + kept_assistant_with_calls = [ + m for m in out if m.get("role") == "assistant" and isinstance(m.get("tool_calls"), list) + ] + kept_tool_msgs = [m for m in out if m.get("role") == "tool"] + # If the assistant tool-call message survives, every matching + # tool-role message must also survive (and vice versa). + kept_ids = set() + for m in kept_assistant_with_calls: + for tc in m.get("tool_calls", []): + if isinstance(tc, dict) and tc.get("id"): + kept_ids.add(tc["id"]) + for m in kept_tool_msgs: + assert m.get("tool_call_id") in kept_ids + + def test_tool_pair_kept_when_in_recent_window(self): + # Build a chat where the tool pair sits inside the recent + # window so it must survive even when the rest is far over + # budget. Layout: system, first-user, long pair x2, tool pair, + # one final user turn. keep_recent=2 catches the tool pair as + # the second-to-last group and the final user turn as the last. + msgs = [ + _msg("system", "sys"), + _msg("user", "first task"), + _long("assistant", 800), + _long("user", 800), + _long("assistant", 800), + _long("user", 800), + _msg( + "assistant", + content = "", + tool_calls = [ + { + "id": "call_42", + "type": "function", + "function": { + "name": "web_search", + "arguments": '{"q":"x"}', + }, + } + ], + ), + _msg("tool", content = "result", tool_call_id = "call_42", name = "web_search"), + _msg("user", "thanks"), + ] + out = SlidingWindowCompact(keep_recent = 2).compact(msgs, budget_tokens = 50) + ids = [ + ( + m.get("role"), + m.get("tool_call_id") or (m.get("tool_calls") and m["tool_calls"][0].get("id")), + ) + for m in out + ] + assert ("assistant", "call_42") in ids + assert ("tool", "call_42") in ids + + +class TestStrategyRegistry: + def test_get_strategy_known(self): + assert isinstance(get_strategy("none"), NoCompact) + assert isinstance(get_strategy("sliding"), SlidingWindowCompact) + + def test_get_strategy_unknown_falls_back_to_none(self): + # Unknown strategy names must degrade to no-op rather than raise. + 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 + with pytest.raises(ValueError): + SlidingWindowCompact(keep_recent = -1) + + def test_invalid_threshold_raises(self): + import pytest + with pytest.raises(ValueError): + SlidingWindowCompact(compact_threshold = 0.0) + with pytest.raises(ValueError): + SlidingWindowCompact(compact_threshold = 1.5) diff --git a/studio/backend/tests/test_context_compaction_edge.py b/studio/backend/tests/test_context_compaction_edge.py new file mode 100644 index 0000000000..e4329deae7 --- /dev/null +++ b/studio/backend/tests/test_context_compaction_edge.py @@ -0,0 +1,947 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Edge-case coverage for the context-compaction module. + +These complement ``test_context_compaction.py`` with adversarial / +boundary inputs the happy-path tests do not exercise. The goal is to +pin down the structural invariants under malformed or extreme inputs: + +* Empty / tiny / oversize inputs. +* Floating / boundary constructor arguments. +* Assistant messages whose ``content`` is None or carries both text and + tool_calls at the same time. +* Tool messages without a matching assistant, or that arrive before + their assistant (malformed but must not crash). +* Duplicate or non-string ``tool_call_id`` values. +* ``content`` that is neither str nor list (defensive). +* No-system / no-first-user histories. +* Threshold and recent-window boundaries (``compact_threshold == 1.0``, + ``keep_recent > len(messages)``). +* The tool-pair "drop one side, drop the other" rule under awkward + inputs (orphan tool, tool-without-id, anchored multimodal user with + a tool message tied to a different assistant). +""" + +import pytest + +from core.inference.context_compaction import ( + NoCompact, + SlidingWindowCompact, + _pair_linked_indices, + estimate_tokens, + get_strategy, +) + + +def _msg( + role, + content = "", + tool_calls = None, + tool_call_id = None, + name = None, +): + m = {"role": role} + if content is not None: + m["content"] = content + if tool_calls is not None: + m["tool_calls"] = tool_calls + if tool_call_id is not None: + m["tool_call_id"] = tool_call_id + if name is not None: + m["name"] = name + return m + + +def _long( + role, + length, + *, + content_prefix = "x", +): + return _msg(role, content_prefix * length) + + +def _tool_call( + tcid, + name = "web_search", + args = '{"q":"x"}', +): + return { + "id": tcid, + "type": "function", + "function": {"name": name, "arguments": args}, + } + + +def _surviving_tool_ids(messages): + """Tool-call ids that a chat template would expect to match. + + Returns (asst_call_ids, tool_response_ids). For a structurally valid + OpenAI chat-completions request, the two sets must be equal. + """ + asst_ids = set() + tool_ids = set() + for m in messages: + if m.get("role") == "assistant": + for tc in m.get("tool_calls") or []: + if isinstance(tc, dict) and isinstance(tc.get("id"), str): + asst_ids.add(tc["id"]) + elif m.get("role") == "tool": + tcid = m.get("tool_call_id") + if isinstance(tcid, str): + tool_ids.add(tcid) + return asst_ids, tool_ids + + +class TestEmptyAndTrivialInputs: + def test_empty_messages_returns_empty(self): + # estimate_tokens([]) == 0 already covered; here check compact(). + out = SlidingWindowCompact().compact([], budget_tokens = 100) + assert out == [] + + def test_single_message_under_budget_passthrough(self): + msgs = [_msg("user", "hi")] + out = SlidingWindowCompact().compact(msgs, budget_tokens = 100) + assert out == msgs + # Defensive: result is a copy, not the same list object. + assert out is not msgs + + def test_single_anchor_over_budget_stays(self): + # Only message is the first-user anchor. Anchors are never + # dropped, even when alone they exceed budget. Result will be + # over budget — that is acceptable per the strategy's docstring + # ("nothing left to drop"), and the compactor must not loop + # forever or crash trying. + msgs = [_long("user", 10000)] + out = SlidingWindowCompact(keep_recent = 0).compact(msgs, budget_tokens = 5) + assert out == msgs + + +class TestKeepRecentBoundaries: + def test_keep_recent_larger_than_message_count(self): + # keep_recent >> len(messages): recent window covers everything, + # so the strategy is effectively a no-op even when over budget. + msgs = [ + _msg("system", "sys"), + _msg("user", "first task"), + _long("assistant", 5000), + _long("user", 5000), + ] + out = SlidingWindowCompact(keep_recent = 100).compact(msgs, budget_tokens = 10) + assert out == msgs + + def test_keep_recent_equals_message_count_after_anchors(self): + # keep_recent exactly equal to the non-anchor group count is the + # boundary between "everything kept" and "drop at least one". + msgs = [ + _msg("system", "sys"), + _msg("user", "first task"), + _long("assistant", 5000), + _long("user", 5000), + _long("assistant", 5000), + ] + out = SlidingWindowCompact(keep_recent = 3).compact(msgs, budget_tokens = 10) + # All non-anchor groups (the trailing 3) are in the recent + # window. Result is the full list. + assert out == msgs + + +class TestThresholdBoundary: + def test_compact_threshold_one_exact(self): + # threshold == 1.0 is the inclusive boundary; estimate exactly + # at budget should be considered "fits" and skip compaction. + msgs = [_msg("user", "abcd")] # 1 token + out = SlidingWindowCompact( + keep_recent = 0, + compact_threshold = 1.0, + ).compact(msgs, budget_tokens = 1) + assert out == msgs + + def test_threshold_just_below_estimate_triggers_compaction(self): + # threshold (== int(budget * compact_threshold)) below estimate + # should compact. + msgs = [ + _msg("system", "sys"), + _msg("user", "first task"), + _long("assistant", 4000), + _long("user", 4000), + _long("assistant", 4000), + _long("user", 4000), + ] + before = estimate_tokens(msgs) + assert before > 100 + out = SlidingWindowCompact(keep_recent = 1).compact(msgs, budget_tokens = 100) + assert len(out) < len(msgs) + + def test_constructor_rejects_negative_threshold(self): + with pytest.raises(ValueError): + SlidingWindowCompact(compact_threshold = -0.1) + + +class TestAssistantWithBothContentAndToolCalls: + def test_assistant_with_text_and_tool_calls_kept_together(self): + # OpenAI chat completions schema allows an assistant message to + # carry BOTH non-empty content and tool_calls. The compactor + # should treat the (content + tool_calls) message and its tool + # response as one group. + msgs = [ + _msg("system", "sys"), + _msg("user", "first task"), + _long("assistant", 4000), + _long("user", 4000), + _msg( + "assistant", + content = "thinking out loud while calling a tool", + tool_calls = [_tool_call("call_x")], + ), + _msg("tool", content = "result for x", tool_call_id = "call_x"), + _msg("user", "thanks"), + ] + out = SlidingWindowCompact(keep_recent = 2).compact(msgs, budget_tokens = 50) + asst_ids, tool_ids = _surviving_tool_ids(out) + assert asst_ids == tool_ids + + +class TestNoneContentAssistant: + def test_none_content_with_tool_calls(self): + # OpenAI permits assistant content=None when tool_calls is set. + msgs = [ + _msg("system", "sys"), + _msg("user", "first task"), + _long("assistant", 4000), + _long("user", 4000), + { + "role": "assistant", + "content": None, + "tool_calls": [_tool_call("call_n")], + }, + _msg("tool", content = "ok", tool_call_id = "call_n"), + _msg("user", "thanks"), + ] + # Must not raise. estimate_tokens treats None content as 0. + out = SlidingWindowCompact(keep_recent = 2).compact(msgs, budget_tokens = 50) + asst_ids, tool_ids = _surviving_tool_ids(out) + assert asst_ids == tool_ids + + +class TestOutOfOrderToolMessage: + def test_tool_message_before_its_assistant_does_not_crash(self): + # Malformed input: tool message arrives before any assistant. + # _pair_linked_indices walks in order; the tool message has no + # pending id to match, so it ends up unpaired. The compactor + # must not crash. Use distinguishable contents so we can verify + # output order directly (identical _long() messages would + # collide under list.index()). + msgs = [ + _msg("system", "sys"), + _msg("user", "first task"), + _msg("tool", content = "orphan-tool", tool_call_id = "call_z"), + _msg("assistant", content = "a-" + "x" * 5000), + _msg("user", "u-" + "x" * 5000), + _msg("assistant", content = "b-" + "x" * 5000), + _msg("user", "v-" + "x" * 5000), + ] + # Identity-track each output back to the source index. + out = SlidingWindowCompact(keep_recent = 2).compact(msgs, budget_tokens = 30) + idxs = [msgs.index(m) for m in out] + assert idxs == sorted(idxs) + + +class TestOrphanToolMessageInRecentWindow: + def test_orphan_tool_message_alone_does_not_break(self): + # An orphan tool message (no preceding assistant tool_calls) + # ends up in its own group. If it lands in the recent window + # the compactor will keep it. That is malformed input on the + # caller's side, not a compactor bug — but it must not crash. + msgs = [ + _msg("system", "sys"), + _msg("user", "first task"), + _long("assistant", 4000), + _long("user", 4000), + _msg("tool", content = "lonely", tool_call_id = "ghost"), + _msg("user", "thanks"), + ] + out = SlidingWindowCompact(keep_recent = 2).compact(msgs, budget_tokens = 30) + # Compactor returns something and does not raise. + assert isinstance(out, list) + + +class TestDuplicateToolCallIds: + def test_duplicate_ids_across_assistants(self): + # Two assistants reuse the same tool_call_id. Only the second + # gets the tool match (pending_ids[tid] = i overwrites). The + # earlier assistant's tool_call has no matching tool. The + # compactor must still produce a coherent result when the + # second assistant + its tool fall inside the recent window. + msgs = [ + _msg("system", "sys"), + _msg("user", "first task"), + _msg("assistant", content = "", tool_calls = [_tool_call("dup")]), + # No tool message for the first call (malformed input). + _long("assistant", 4000), + _long("user", 4000), + _msg("assistant", content = "", tool_calls = [_tool_call("dup")]), + _msg("tool", content = "second", tool_call_id = "dup"), + _msg("user", "thanks"), + ] + out = SlidingWindowCompact(keep_recent = 2).compact(msgs, budget_tokens = 30) + # The recent-window assistant+tool group must stay linked. + kept_asst_tool = [m for m in out if m.get("role") == "assistant" and m.get("tool_calls")] + kept_tools = [m for m in out if m.get("role") == "tool"] + if kept_tools: + tool_ids_in_out = {m["tool_call_id"] for m in kept_tools} + asst_ids_in_out = set() + for m in kept_asst_tool: + for tc in m.get("tool_calls") or []: + if isinstance(tc.get("id"), str): + asst_ids_in_out.add(tc["id"]) + assert tool_ids_in_out <= asst_ids_in_out + + +class TestPairMapInternals: + def test_pair_map_empty_for_plain_chat(self): + msgs = [ + _msg("system", "sys"), + _msg("user", "hi"), + _msg("assistant", "hello"), + ] + pm = _pair_linked_indices(msgs) + # Assistant has no tool_calls so the set entry is empty. + assert pm == {2: set()} + + def test_pair_map_tool_without_string_id_ignored(self): + # tool_call_id is not a string — should be ignored, not crash. + msgs = [ + _msg("assistant", content = "", tool_calls = [_tool_call("ok")]), + {"role": "tool", "content": "x", "tool_call_id": 42}, + {"role": "tool", "content": "y", "tool_call_id": None}, + ] + pm = _pair_linked_indices(msgs) + # Only the string-id tool would match, and we did not include + # one — so the assistant entry is empty. + assert pm.get(0) == set() + + +class TestNonStringContent: + def test_int_or_dict_content_does_not_crash(self): + # Defensive: content that is neither str nor list (a stray int + # or dict from a misbehaving caller) should be ignored by + # estimate_tokens and treated as non-multimodal by compact(). + msgs = [ + _msg("system", "sys"), + _msg("user", "first task"), + {"role": "assistant", "content": 12345}, + {"role": "user", "content": {"unexpected": "shape"}}, + _long("assistant", 4000), + _long("user", 4000), + ] + out = SlidingWindowCompact(keep_recent = 2).compact(msgs, budget_tokens = 30) + assert isinstance(out, list) + + +class TestNoSystemNoFirstUser: + def test_assistant_only_history(self): + # No system, no user. The first-user anchor loop simply does + # not match anything. Compactor must not crash. + msgs = [ + _long("assistant", 4000), + _long("assistant", 4000), + _long("assistant", 4000), + _long("assistant", 4000), + ] + out = SlidingWindowCompact(keep_recent = 1).compact(msgs, budget_tokens = 30) + # At least one message survives (the last assistant in the + # recent window). + assert len(out) >= 1 + + def test_first_user_not_at_index_one(self): + # First user is not adjacent to a system message. The + # first-user anchor must still find it by iteration order. + msgs = [ + _msg("assistant", content = "stray pre-amble"), + _msg("user", "the real first task"), + _long("assistant", 4000), + _long("user", 4000), + _long("assistant", 4000), + ] + out = SlidingWindowCompact(keep_recent = 1).compact(msgs, budget_tokens = 30) + # The first user must be present. + assert any( + m.get("role") == "user" and "real first task" in m.get("content", "") for m in out + ) + + +class TestMultimodalAsAnchor: + def test_multimodal_first_user(self): + # The very first user message is itself multimodal. Both the + # first-user anchor and the multimodal anchor refer to the + # same index — set semantics make the overlap a no-op. + msgs = [ + _msg("system", "sys"), + { + "role": "user", + "content": [ + {"type": "text", "text": "describe"}, + {"type": "image_url", "image_url": {"url": "x"}}, + ], + }, + _long("assistant", 4000), + _long("user", 4000), + _long("assistant", 4000), + ] + out = SlidingWindowCompact(keep_recent = 1).compact(msgs, budget_tokens = 20) + # The multimodal first user must still be present. + assert any(isinstance(m.get("content"), list) for m in out) + + +class TestPairCleanupConsistency: + def test_drop_assistant_drops_its_tools(self): + # The assistant tool-call message is forced out of the recent + # window by keep_recent and is not an anchor. Its tool message + # must be dropped with it. + msgs = [ + _msg("system", "sys"), + _msg("user", "first task"), + _msg("assistant", content = "", tool_calls = [_tool_call("call_a")]), + _msg("tool", content = "old result", tool_call_id = "call_a"), + _long("assistant", 5000), + _long("user", 5000), + _long("assistant", 5000), + _long("user", 5000), + ] + out = SlidingWindowCompact(keep_recent = 1).compact(msgs, budget_tokens = 20) + asst_ids, tool_ids = _surviving_tool_ids(out) + # No orphans either way. + assert asst_ids == tool_ids + # And the old pair is gone. + assert "call_a" not in tool_ids + + def test_drop_tools_drops_assistant(self): + # Hand-craft a scenario where the tool-side gets removed by + # the threshold loop but the assistant would be kept. The + # pair-cleanup pass must then drop the assistant too. + # We exploit the fact that the threshold-driven loop runs in + # `droppable` order (system-anchored + first-user-anchored + # excluded). The assistant tool-call carries little text; + # the tool response carries a lot. Dropping the tool may + # already get us under threshold, leaving the assistant. + # The cleanup pass must then drop the orphaned assistant. + msgs = [ + _msg("system", "sys"), + _msg("user", "first task" + "y" * 100), + _msg("assistant", content = "", tool_calls = [_tool_call("call_b")]), + _msg("tool", content = "x" * 4000, tool_call_id = "call_b"), + _long("assistant", 50), + _long("user", 50), + ] + out = SlidingWindowCompact(keep_recent = 1).compact(msgs, budget_tokens = 100) + asst_ids, tool_ids = _surviving_tool_ids(out) + # Both halves of the pair are either kept or dropped together. + assert asst_ids == tool_ids + + +class TestImmutability: + def test_input_messages_never_mutated_under_pressure(self): + msgs = [ + _msg("system", "sys"), + _msg("user", "first task"), + _msg("assistant", content = "", tool_calls = [_tool_call("call_m")]), + _msg("tool", content = "ok", tool_call_id = "call_m"), + _long("assistant", 5000), + _long("user", 5000), + ] + # Deep-ish snapshot: messages themselves + their tool_calls list + # are what the compactor could conceivably mutate. + snap = [dict(m) for m in msgs] + snap_tcs = [list(m.get("tool_calls") or []) for m in msgs] + _ = SlidingWindowCompact(keep_recent = 1).compact(msgs, budget_tokens = 30) + assert msgs == snap + for m, original_tcs in zip(msgs, snap_tcs): + assert (m.get("tool_calls") or []) == original_tcs + + +class TestStrategyRegistryAliases: + def test_get_strategy_empty_name_falls_back(self): + # The fallback path matters: a misconfigured request should + # degrade to no-op rather than raise. + assert isinstance(get_strategy(""), NoCompact) + + +class TestEstimateTokensWithBadToolCalls: + def test_tool_call_without_function_dict(self): + # Defensive: tool_calls list where an entry lacks "function". + # estimate_tokens does `(tc.get("function") or {}).get(...)`, + # which is safe — must not crash. + msgs = [ + { + "role": "assistant", + "content": "", + "tool_calls": [{"id": "x", "type": "function"}], + } + ] + assert estimate_tokens(msgs) == 0 + + +class TestMultiToolCallSingleAssistant: + def test_assistant_with_multiple_tool_calls_grouped(self): + # One assistant message carries two tool_calls; both tool + # responses must be grouped with it so the trio stays linked. + msgs = [ + _msg("system", "sys"), + _msg("user", "first task"), + _long("assistant", 4000), + _long("user", 4000), + _msg( + "assistant", + content = "", + tool_calls = [_tool_call("a"), _tool_call("b")], + ), + _msg("tool", content = "ra", tool_call_id = "a"), + _msg("tool", content = "rb", tool_call_id = "b"), + _msg("user", "thanks"), + ] + out = SlidingWindowCompact(keep_recent = 2).compact(msgs, budget_tokens = 50) + asst_ids, tool_ids = _surviving_tool_ids(out) + # All-or-nothing: the multi-call assistant brings both tool + # responses along, or none of the three survive. + assert asst_ids == tool_ids + + +class TestInterleavedUserBetweenAsstAndTool: + def test_user_between_asst_and_tool_does_not_break_grouping(self): + # Malformed (user message between asst tool_call and its tool + # response). pair_map still walks in order so the tool matches + # the most recent assistant carrying its id. The compactor + # must not raise; downstream rendering of an "orphan" user + # is acceptable since it was already orphaned in the input. + msgs = [ + _msg("system", "sys"), + _msg("user", "first task"), + _msg("assistant", content = "", tool_calls = [_tool_call("g")]), + _msg("user", "intermediate"), + _msg("tool", content = "result", tool_call_id = "g"), + _long("assistant", 5000), + _long("user", 5000), + ] + out = SlidingWindowCompact(keep_recent = 1).compact(msgs, budget_tokens = 30) + asst_ids, tool_ids = _surviving_tool_ids(out) + assert asst_ids == tool_ids + + +class TestNoCompactWithEmpty: + def test_nocompact_empty_messages_returns_empty(self): + out = NoCompact().compact([], budget_tokens = 100) + assert out == [] + + +# ── User-boundary breaks the pair window ───────────────────── + + +def test_tool_message_after_user_does_not_pair_with_earlier_assistant(): + """An intervening user message ends the pending pair window. A + later tool message arriving after the user is malformed input per + the OpenAI chat schema and must NOT be grouped with the earlier + assistant tool_call; otherwise the compactor would drop them + together and corrupt the surviving template. + """ + from core.inference.context_compaction import _pair_linked_indices + + msgs = [ + {"role": "system", "content": "sys"}, + {"role": "user", "content": "first task"}, + { + "role": "assistant", + "content": "", + "tool_calls": [ + { + "id": "call_X", + "type": "function", + "function": {"name": "search", "arguments": "{}"}, + } + ], + }, + # User interrupts before any tool message lands. The next tool + # message (malformed) must not pair back to assistant[2]. + {"role": "user", "content": "changed my mind, ask something else"}, + { + "role": "tool", + "content": "stale result", + "tool_call_id": "call_X", + "name": "search", + }, + ] + out = _pair_linked_indices(msgs) + # Assistant at index 2 keeps an entry (it was seen) but with NO + # tool follow-ups paired across the user boundary. + assert 2 in out, out + assert out[2] == set(), out + + +def test_paired_window_resumes_after_new_assistant_with_tool_call(): + """A fresh assistant after a user message starts a new pending + window. Its own tool messages pair correctly, the prior assistant + stays unpaired. + """ + from core.inference.context_compaction import _pair_linked_indices + + msgs = [ + {"role": "user", "content": "first"}, + { + "role": "assistant", + "content": "", + "tool_calls": [ + { + "id": "A", + "type": "function", + "function": {"name": "f", "arguments": "{}"}, + } + ], + }, + {"role": "user", "content": "wait, do this instead"}, + { + "role": "assistant", + "content": "", + "tool_calls": [ + { + "id": "B", + "type": "function", + "function": {"name": "g", "arguments": "{}"}, + } + ], + }, + { + "role": "tool", + "content": "result B", + "tool_call_id": "B", + "name": "g", + }, + ] + out = _pair_linked_indices(msgs) + assert out[1] == set(), out + assert out[3] == {4}, out + + +# ── Defensive estimate_tokens ──────────────────────────────── + + +def test_estimate_tokens_does_not_crash_on_non_dict_tool_calls(): + """OpenAI dicts can carry non-dict tool_calls entries before + pydantic validation; the heuristic must skip them rather than + raise AttributeError mid-compaction. + """ + msgs = [ + { + "role": "assistant", + "content": "", + # First entry is malformed (string), second is fine. + "tool_calls": [ + "bad-entry", + None, + { + "id": "ok", + "function": {"name": "f", "arguments": "{}"}, + }, + ], + } + ] + assert estimate_tokens(msgs) >= 0 + + +def test_estimate_tokens_counts_compaction_part(): + """Studio's compaction content part carries an Anthropic summary + that can be multi-KB; counting it prevents the threshold check + from skipping compaction when the real prompt is huge. + """ + summary = "x" * 4000 + msgs = [ + { + "role": "assistant", + "content": [ + {"type": "text", "text": "intro"}, + {"type": "compaction", "content": summary}, + ], + } + ] + # 5 + 4000 chars -> ceil-divided by 4 ~= 1002 tokens. + assert estimate_tokens(msgs) >= 1000, estimate_tokens(msgs) + + +# ── Anomaly regressions surfaced by sim_pr5710.py fuzzing ──── + + +def test_pair_linked_indices_skips_non_dict_tool_call_entries(): + """``_assistant_tool_call_ids`` ran inside ``_pair_linked_indices``, + so a malformed string / None entry in ``tool_calls`` used to crash + the compactor mid-call. Mirrors the ``estimate_tokens`` guard. + """ + msgs = [ + {"role": "assistant", "content": "x", "tool_calls": ["bare-string"]}, + {"role": "assistant", "content": "y", "tool_calls": [None, {"id": "t1"}]}, + {"role": "tool", "tool_call_id": "t1", "content": "ok"}, + ] + pm = _pair_linked_indices(msgs) + assert pm[0] == set() + assert pm[1] == {2} + + +def test_compact_drops_orphan_tool_left_by_user_boundary(): + """Tool message arrives after a user / system boundary, references + an assistant tool_call_id that the boundary-clearing logic no + longer treats as the pair root. The assistant gets dropped by the + main sweep; without the final invariant pass the orphan tool + survives into the output and llama-server rejects the template. + Repro distilled from sim_pr5710.py fuzz iter 1407. + """ + msgs = [ + {"role": "system", "content": "sys"}, + {"role": "user", "content": "task " * 20}, + {"role": "assistant", "content": "a", "tool_calls": [_tool_call("t1")]}, + {"role": "tool", "tool_call_id": "t1", "content": "first"}, + {"role": "user", "content": "next"}, + {"role": "user", "content": "again"}, + {"role": "tool", "tool_call_id": "t1", "content": "stale"}, + ] + out = SlidingWindowCompact(keep_recent = 2).compact(msgs, budget_tokens = 1) + # Any surviving tool message must have its assistant earlier in + # the output. + seen_ids: set[str] = set() + for m in out: + if m.get("role") == "assistant": + for tc in m.get("tool_calls") or []: + tcid = tc.get("id") if isinstance(tc, dict) else None + if isinstance(tcid, str) and tcid: + seen_ids.add(tcid) + elif m.get("role") == "tool": + tcid = m.get("tool_call_id") + assert ( + isinstance(tcid, str) and tcid in seen_ids + ), f"orphan tool {tcid!r} survived; seen={seen_ids}" + + +def test_compact_keeps_full_pair_when_tool_is_multimodal_anchor(): + """An anchored multimodal tool implies the whole asst+tools pair + gets anchored. Earlier behavior leaked the orphan tool alone, which + OpenAI 400s on ("orphan tool message"). Anchor propagation across + paired indices keeps both halves so the chat template stays valid. + """ + msgs = [ + {"role": "system", "content": "sys"}, + {"role": "user", "content": "task " * 20}, + {"role": "assistant", "content": "a", "tool_calls": [_tool_call("t1")]}, + # Real multimodal tool message (image part) -- anchored. + { + "role": "tool", + "tool_call_id": "t1", + "content": [ + {"type": "text", "text": "ok"}, + {"type": "image_url", "image_url": {"url": "x"}}, + ], + }, + {"role": "user", "content": "more"}, + {"role": "user", "content": "even more"}, + ] + out = SlidingWindowCompact(keep_recent = 2).compact(msgs, budget_tokens = 1) + # Both halves of the pair survive: anchored tool keeps its asst. + asst_ids, tool_ids = _surviving_tool_ids(out) + assert "t1" in asst_ids + assert "t1" in tool_ids + assert asst_ids == tool_ids + + +def test_anchored_multimodal_asst_orphan_tool_calls_stripped(): + """Multimodal assistant carrying tool_calls whose tool follow-ups + all get dropped: the anchor invariant keeps the assistant, but the + leftover tool_calls field references ids with no matching tool + response and OpenAI 400s on that shape. Strip the orphan tool_calls + from a copy so the multimodal content survives and the request + stays well-formed. + """ + multimodal_asst = { + "role": "assistant", + "content": [ + {"type": "text", "text": "look"}, + {"type": "image_url", "image_url": {"url": "x"}}, + ], + "tool_calls": [_tool_call("t1"), _tool_call("t2")], + } + msgs = [ + _msg("system", "sys"), + _msg("user", "task"), + multimodal_asst, + _msg("tool", content = "A" * 50, tool_call_id = "t1"), + _msg("tool", content = "B" * 4000, tool_call_id = "t2"), + _long("assistant", 50), + _long("user", 50), + ] + out = SlidingWindowCompact(keep_recent = 1).compact(msgs, budget_tokens = 100) + asst_ids, tool_ids = _surviving_tool_ids(out) + assert asst_ids == tool_ids, (asst_ids, tool_ids) + # Multimodal content preserved. + assert any(isinstance(m.get("content"), list) for m in out) + # Original input unchanged (we copy on rewrite). + assert multimodal_asst["tool_calls"] == [_tool_call("t1"), _tool_call("t2")] + + +def test_intervening_assistant_breaks_pair_window(): + """A later assistant turn arriving before the matching tool message + must end the pending pair window. Otherwise the compactor groups a + stale assistant + a malformed late tool together. Mirrors the + user/system boundary rule for the assistant boundary. + """ + from core.inference.context_compaction import _pair_linked_indices + + msgs = [ + {"role": "system", "content": "sys"}, + {"role": "user", "content": "task"}, + { + "role": "assistant", + "content": "", + "tool_calls": [ + {"id": "a", "function": {"name": "f", "arguments": "{}"}}, + ], + }, + {"role": "assistant", "content": "intervening"}, + {"role": "tool", "tool_call_id": "a", "content": "stale"}, + ] + pm = _pair_linked_indices(msgs) + # Asst at index 2 had its window closed by asst at index 3. + assert pm.get(2) == set(), pm + # The intervening assistant has no tool_calls and so an empty set. + assert pm.get(3) == set(), pm + + +def test_compaction_content_part_is_not_multimodal_anchor(): + """Studio's ``{"type":"compaction","content":"..."}`` parts are the + summary the compactor is supposed to compact away -- they must not + pin the carrier message as a multimodal anchor. A text-only list + (only text + compaction parts) collapses to "not multimodal" and + stays droppable. + """ + from core.inference.context_compaction import _is_multimodal + + # Pure compaction part. + msg = { + "role": "assistant", + "content": [{"type": "compaction", "content": "OLD " * 200}], + } + assert _is_multimodal(msg) is False + + # Mixed text + compaction (still no media payload). + msg2 = { + "role": "assistant", + "content": [ + {"type": "text", "text": "intro"}, + {"type": "compaction", "content": "OLD " * 200}, + ], + } + assert _is_multimodal(msg2) is False + + # Real multimodal part wins anchoring even alongside compaction. + msg3 = { + "role": "assistant", + "content": [ + {"type": "compaction", "content": "OLD"}, + {"type": "image_url", "image_url": {"url": "x"}}, + ], + } + assert _is_multimodal(msg3) is True + + # End-to-end: an old compaction-only message is droppable. + msgs = [ + _msg("system", "sys"), + _msg("user", "first task"), + { + "role": "assistant", + "content": [{"type": "compaction", "content": "x" * 8000}], + }, + _long("assistant", 100), + _long("user", 100), + ] + out = SlidingWindowCompact(keep_recent = 1).compact(msgs, budget_tokens = 20) + # The old compaction-only assistant should have been dropped. + assert not any( + m.get("role") == "assistant" + and isinstance(m.get("content"), list) + and any(isinstance(p, dict) and p.get("type") == "compaction" for p in m["content"]) + for m in out + ) + + +def test_partial_tool_drop_strips_orphan_tool_call_id(): + """Same shape but a plain-content assistant: the per-index budget + loop dropped only one of the two tool follow-ups. The surviving + tool_call_id on the assistant points at a dropped tool message -- + strip it so the chat template stays valid. Plain assistant is not + anchored, so the original "drop the assistant when all tools are + gone" rule would have caught this if BOTH tools had been dropped. + Here only one tool was dropped, which is the gap this guard closes. + """ + msgs = [ + _msg("system", "sys"), + _msg("user", "task"), + _msg( + "assistant", + content = "thinking", + tool_calls = [_tool_call("t1"), _tool_call("t2")], + ), + _msg("tool", content = "A" * 50, tool_call_id = "t1"), + _msg("tool", content = "B" * 4000, tool_call_id = "t2"), + _long("assistant", 50), + _long("user", 50), + ] + out = SlidingWindowCompact(keep_recent = 1).compact(msgs, budget_tokens = 200) + asst_ids, tool_ids = _surviving_tool_ids(out) + assert asst_ids == tool_ids, (asst_ids, tool_ids) + + +def test_orphan_tool_before_asst_does_not_leave_dangling_tool_calls(): + """Malformed input: a tool message arrives BEFORE the assistant + that references its id. The final sweep must drop the orphan tool + AND strip the now-dangling tool_call from the later assistant so + the chat template stays valid. Earlier behavior computed + responded_ids once before the orphan-tool drop, so the post-drop + asst still matched ids <= responded_ids and kept its tool_calls. + """ + msgs = [ + _msg("system", "sys"), + _msg("user", "first task"), + _msg("tool", content = "early", tool_call_id = "call_a"), + _msg( + "assistant", + content = "thinking", + tool_calls = [_tool_call("call_a")], + ), + _long("user", 50), + ] + out = SlidingWindowCompact(keep_recent = 5).compact(msgs, budget_tokens = 1000) + asst_ids, tool_ids = _surviving_tool_ids(out) + assert asst_ids == tool_ids, (asst_ids, tool_ids) + + +def test_anchored_multimodal_orphan_tool_dropped(): + """Anchored multimodal tool message references a tool_call_id no + surviving assistant declares. Keeping the anchor would leave a + dangling `tool_call_id` in the output and 400 upstream. Pair + validity is the hard invariant; the multimodal anchor is the soft + quality preference, so we drop the orphan rather than violate the + template. Earlier behavior preserved the anchor and produced an + invalid chat template. + """ + msgs = [ + _msg("system", "sys"), + _msg("user", "first task"), + { + "role": "tool", + "tool_call_id": "stale_call", + "content": [ + {"type": "text", "text": "image description"}, + {"type": "image_url", "image_url": {"url": "x"}}, + ], + }, + _msg("user", "follow up"), + _long("assistant", 100), + ] + out = SlidingWindowCompact(keep_recent = 2).compact(msgs, budget_tokens = 5) + asst_ids, tool_ids = _surviving_tool_ids(out) + assert asst_ids == tool_ids, (asst_ids, tool_ids)