From 91b2fa57a8031c590781761129e9444a8eae032c Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Fri, 22 May 2026 14:46:27 +0000 Subject: [PATCH 01/18] Studio: sliding-window context compaction module Adds a standalone context-compaction module so long Studio Chat sessions can trim older turns from the prompt sent to the model without losing them from the persisted transcript. This PR ships the module and tests only; the wire-up into the agentic and chat loops lands in a follow-up so behavior on main is unchanged. Two strategies: * NoCompact: passthrough, returns a copy. * SlidingWindowCompact: keeps the system message, the first user message, and the last N "groups" (an assistant tool-call message plus its matching tool-role responses count as one group). Invariants preserved by every strategy: * System message and first user message are never dropped. * Tool-call <-> tool-result pair linkage stays valid: an assistant message that carries tool_calls is kept or dropped together with every tool-role message matching its tool_call_ids. * Multimodal turns (list-typed content) are never dropped (no tested compacted-media representation today). * The input list is never mutated; the strategy returns a new list. Token estimation uses a 4-char-per-token char-count heuristic; tokenizer-aware estimates are a follow-up. Adapted from forge (https://github.com/antoinezambelli/forge, MIT). --- .../core/inference/context_compaction.py | 262 ++++++++++++++++++ .../backend/tests/test_context_compaction.py | 261 +++++++++++++++++ 2 files changed, 523 insertions(+) create mode 100644 studio/backend/core/inference/context_compaction.py create mode 100644 studio/backend/tests/test_context_compaction.py diff --git a/studio/backend/core/inference/context_compaction.py b/studio/backend/core/inference/context_compaction.py new file mode 100644 index 0000000000..3920ee3e3e --- /dev/null +++ b/studio/backend/core/inference/context_compaction.py @@ -0,0 +1,262 @@ +# 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 + + +# Char-to-token heuristic. Conservative on the high side so the +# compactor triggers earlier rather than later. Tokenizer-aware +# estimates may land in a follow-up. +_CHARS_PER_TOKEN = 4 + + +def estimate_tokens(messages: list[dict]) -> int: + """Rough token count for ``messages``. Uses a 4-char-per-token + char-count heuristic on the visible ``content`` (str) and on + serialized ``tool_calls`` arguments. Multimodal parts (list-typed + content) contribute only their text parts. + """ + 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 isinstance(part, dict): + t = part.get("text") + if isinstance(t, str): + total_chars += len(t) + tcs = m.get("tool_calls") + if isinstance(tcs, list): + for tc in tcs: + args = (tc.get("function") or {}).get("arguments") + if isinstance(args, str): + total_chars += len(args) + return total_chars // _CHARS_PER_TOKEN + + +def _is_multimodal(msg: dict) -> bool: + return isinstance(msg.get("content"), list) + + +def _is_tool_message(msg: dict) -> bool: + return msg.get("role") == "tool" + + +def _assistant_tool_call_ids(msg: dict) -> set[str]: + """Return the set of ``id`` values from an assistant message's + ``tool_calls``. Empty set when the message has no tool calls. + """ + out: set[str] = set() + tcs = msg.get("tool_calls") + if isinstance(tcs, list): + for tc in tcs: + 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 an assistant-message index to the indices of its tool-role + follow-ups (matching ``tool_call_id``). Used so the compactor drops + or keeps an assistant+tool group as a unit. + """ + out: dict[int, set[int]] = {} + # Walk in order so the next tool-role messages after an assistant + # call are the natural matches. A tool message is matched to the + # most recent prior assistant whose ``tool_calls`` contain that id. + pending_ids: dict[str, int] = {} + for i, m in enumerate(messages): + if m.get("role") == "assistant": + ids = _assistant_tool_call_ids(m) + out.setdefault(i, set()) + for tid in ids: + pending_ids[tid] = i + elif _is_tool_message(m): + tcid = m.get("tool_call_id") + if isinstance(tcid, str) and tcid in pending_ids: + out.setdefault(pending_ids[tcid], set()).add(i) + 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 the system message, the first user message, and the last + ``keep_recent`` non-droppable turns. Multimodal turns are never + dropped (no compacted-media representation today). Assistant + messages with ``tool_calls`` are grouped with their matching + tool-role responses and treated as one unit. + + The strategy is a no-op when the estimated token count is already + within ``budget_tokens`` or when there is nothing left to drop. + """ + + 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) + + # Tool-call / tool-result grouping. An assistant tool-call + # message and its matching tool-role responses get the same + # group id so we keep or drop them as a unit. + pair_map = _pair_linked_indices(messages) + 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 + + # The "recent window" is the last ``keep_recent`` distinct + # groups encountered scanning from the end. + recent_groups: list[int] = [] + seen_groups: set[int] = set() + for i in range(len(messages) - 1, -1, -1): + 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 + # window AND that is not an anchor. + 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) + + # 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. + for asst_idx, tool_idxs in pair_map.items(): + if asst_idx in dropped: + dropped.update(tool_idxs) + elif tool_idxs and tool_idxs <= dropped: + # All matching tool messages were dropped: drop the + # assistant tool-call message too. + dropped.add(asst_idx) + + return [m for i, m in enumerate(messages) if i not in dropped] + + +_STRATEGIES: dict[str, CompactStrategy] = { + "none": NoCompact(), + "sliding": SlidingWindowCompact(), +} + + +def get_strategy(name: str) -> CompactStrategy: + """Return the compaction strategy registered under ``name``. + + Falls back to ``NoCompact`` for unknown names so a misconfigured + request degrades to no-op rather than raising. + """ + 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..f95c0499fe --- /dev/null +++ b/studio/backend/tests/test_context_compaction.py @@ -0,0 +1,261 @@ +# 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_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 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) From 75e32f20f7c5461a813a604689bd4aaa7cbe4739 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Fri, 22 May 2026 14:47:06 +0000 Subject: [PATCH 02/18] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- .../core/inference/context_compaction.py | 21 ++---- .../backend/tests/test_context_compaction.py | 73 ++++++++++++------- 2 files changed, 52 insertions(+), 42 deletions(-) diff --git a/studio/backend/core/inference/context_compaction.py b/studio/backend/core/inference/context_compaction.py index 3920ee3e3e..8428d0cd53 100644 --- a/studio/backend/core/inference/context_compaction.py +++ b/studio/backend/core/inference/context_compaction.py @@ -117,9 +117,7 @@ class CompactStrategy(ABC): """Interface for context-compaction strategies.""" @abstractmethod - def compact( - self, messages: list[dict], budget_tokens: int - ) -> list[dict]: + 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. @@ -130,9 +128,7 @@ class CompactStrategy(ABC): class NoCompact(CompactStrategy): """Passthrough strategy. Returns ``messages`` unchanged.""" - def compact( - self, messages: list[dict], budget_tokens: int - ) -> list[dict]: + def compact(self, messages: list[dict], budget_tokens: int) -> list[dict]: return list(messages) @@ -147,9 +143,7 @@ class SlidingWindowCompact(CompactStrategy): within ``budget_tokens`` or when there is nothing left to drop. """ - def __init__( - self, keep_recent: int = 2, compact_threshold: float = 0.85 - ) -> None: + 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): @@ -157,9 +151,7 @@ class SlidingWindowCompact(CompactStrategy): self.keep_recent = keep_recent self.compact_threshold = compact_threshold - def compact( - self, messages: list[dict], budget_tokens: int - ) -> list[dict]: + def compact(self, messages: list[dict], budget_tokens: int) -> list[dict]: if budget_tokens <= 0 or not messages: return list(messages) @@ -224,10 +216,7 @@ class SlidingWindowCompact(CompactStrategy): # 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 - ] + 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 diff --git a/studio/backend/tests/test_context_compaction.py b/studio/backend/tests/test_context_compaction.py index f95c0499fe..03e688e78a 100644 --- a/studio/backend/tests/test_context_compaction.py +++ b/studio/backend/tests/test_context_compaction.py @@ -65,11 +65,13 @@ class TestEstimateTokens: _msg( "assistant", content = "", - tool_calls = [{ - "id": "c1", - "type": "function", - "function": {"name": "web_search", "arguments": '{"q":"hi"}'}, - }], + tool_calls = [ + { + "id": "c1", + "type": "function", + "function": {"name": "web_search", "arguments": '{"q":"hi"}'}, + } + ], ) ] # 10 char arguments -> 2 tokens. @@ -88,7 +90,11 @@ class TestNoCompact: class TestSlidingWindowUnderBudget: def test_no_op_when_already_fits(self): - msgs = [_msg("system", "be concise"), _msg("user", "hi"), _msg("assistant", "bye")] + msgs = [ + _msg("system", "be concise"), + _msg("user", "hi"), + _msg("assistant", "bye"), + ] out = SlidingWindowCompact(keep_recent = 2).compact(msgs, budget_tokens = 1000) assert out == msgs @@ -164,16 +170,23 @@ class TestSlidingWindowToolPairs: _msg( "assistant", content = "", - tool_calls = [{ - "id": "call_42", - "type": "function", - "function": { - "name": "web_search", - "arguments": '{"q":"x"}', - }, - }], + 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", ), - _msg("tool", content = "result for x", tool_call_id = "call_42", name = "web_search"), _long("assistant", 800), _long("user", 800), _long("assistant", 800), @@ -185,7 +198,8 @@ class TestSlidingWindowToolPairs: # Force aggressive compaction. out = SlidingWindowCompact(keep_recent = 2).compact(msgs, budget_tokens = 50) kept_assistant_with_calls = [ - m for m in out + 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"] @@ -215,22 +229,27 @@ class TestSlidingWindowToolPairs: _msg( "assistant", content = "", - tool_calls = [{ - "id": "call_42", - "type": "function", - "function": { - "name": "web_search", - "arguments": '{"q":"x"}', - }, - }], + 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"))) + ( + 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 @@ -250,11 +269,13 @@ class TestStrategyRegistry: 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): From 12270ca9fab67f17617be4fbc5ead975bdcb7e6c Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Fri, 22 May 2026 15:48:28 +0000 Subject: [PATCH 03/18] Round estimate_tokens up so over-budget prompts don't skip compaction MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `estimate_tokens` used floor division (`// 4`) on the char count, which systematically underestimates any non-multiple-of-4 message length. A 5-char message returned 1 token instead of 2, so a prompt sitting just over the budget could appear under threshold and bypass compaction — the exact failure the module exists to prevent. Switch to ceil-style rounding so the heuristic stays conservative. Adds regression tests covering the off-by-one and the single-char case (where floor returned 0 tokens for non-empty content). --- studio/backend/core/inference/context_compaction.py | 7 ++++++- studio/backend/tests/test_context_compaction.py | 11 +++++++++++ 2 files changed, 17 insertions(+), 1 deletion(-) diff --git a/studio/backend/core/inference/context_compaction.py b/studio/backend/core/inference/context_compaction.py index 8428d0cd53..18c36bb07c 100644 --- a/studio/backend/core/inference/context_compaction.py +++ b/studio/backend/core/inference/context_compaction.py @@ -65,7 +65,12 @@ def estimate_tokens(messages: list[dict]) -> int: args = (tc.get("function") or {}).get("arguments") if isinstance(args, str): total_chars += len(args) - return total_chars // _CHARS_PER_TOKEN + # Ceil-style division: floor (`// 4`) would systematically + # underestimate non-multiple-of-4 lengths, letting just-over-budget + # prompts appear under threshold and bypass compaction — the exact + # failure this module exists to prevent. Round up so the heuristic + # stays on the conservative side described above. + return -(-total_chars // _CHARS_PER_TOKEN) def _is_multimodal(msg: dict) -> bool: diff --git a/studio/backend/tests/test_context_compaction.py b/studio/backend/tests/test_context_compaction.py index 03e688e78a..c42b39d9ad 100644 --- a/studio/backend/tests/test_context_compaction.py +++ b/studio/backend/tests/test_context_compaction.py @@ -47,6 +47,17 @@ class TestEstimateTokens: 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 = [ { From 41192f0452b1fb4c68c0475c8c0305cfa4cd5201 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sun, 24 May 2026 12:45:16 +0000 Subject: [PATCH 04/18] 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. --- .../core/inference/context_compaction.py | 27 ++++-- .../backend/tests/test_context_compaction.py | 82 +++++++++++++++++++ 2 files changed, 103 insertions(+), 6 deletions(-) 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 From db79f192700923d8e601e276838ebb2ed57867d6 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Sun, 24 May 2026 13:44:57 +0000 Subject: [PATCH 05/18] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- studio/backend/tests/test_context_compaction.py | 8 ++------ 1 file changed, 2 insertions(+), 6 deletions(-) diff --git a/studio/backend/tests/test_context_compaction.py b/studio/backend/tests/test_context_compaction.py index aa629093d7..ccce541e84 100644 --- a/studio/backend/tests/test_context_compaction.py +++ b/studio/backend/tests/test_context_compaction.py @@ -296,9 +296,7 @@ class TestKeepRecentZero: _long("assistant", 4000), _long("user", 4000), ] - out = SlidingWindowCompact(keep_recent = 0).compact( - msgs, budget_tokens = 50 - ) + 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"] @@ -349,9 +347,7 @@ class TestAnchoredMultimodalPairCleanup: _long("assistant", 4000), _long("user", 4000), ] - out = SlidingWindowCompact(keep_recent = 1).compact( - msgs, budget_tokens = 50 - ) + 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) From 15aa69dac80ae0b73230da4a505abafbbcd44b54 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sun, 24 May 2026 13:53:53 +0000 Subject: [PATCH 06/18] Studio: edge-case coverage for context-compaction module Adds an adversarial test file alongside the happy-path tests for the new sliding-window compaction strategy. Pins down the structural invariants under malformed and boundary inputs that the original suite does not exercise: * Empty / single-message / over-budget-anchor inputs. * keep_recent boundary cases (>> len(messages), == group count). * compact_threshold == 1.0 boundary and negative-threshold validation. * Assistant messages carrying both content and tool_calls. * Assistant content == None (OpenAI schema when only tool_calls set). * Out-of-order tool messages (malformed, must not crash). * Orphan tool messages and duplicated tool_call_ids across assistants. * Non-string content (int, dict) defensive paths in estimate_tokens. * Histories without a system message or with the first user not at index 1. * First-user message that is itself multimodal (anchor overlap). * Multi-tool-call single-assistant grouping. * Interleaved user message between assistant tool_call and its tool response. * NoCompact applied to an empty list. Module-only PR, still unwired on main; no behavior change. --- .../tests/test_context_compaction_edge.py | 523 ++++++++++++++++++ 1 file changed, 523 insertions(+) create mode 100644 studio/backend/tests/test_context_compaction_edge.py 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..b1397788bb --- /dev/null +++ b/studio/backend/tests/test_context_compaction_edge.py @@ -0,0 +1,523 @@ +# 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 == [] From a12bc349f2db9cc08bae98dd37ee87a9e41b8aec Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Mon, 25 May 2026 07:39:04 +0000 Subject: [PATCH 07/18] Break tool pair window on user boundary, harden estimator Three follow-ups on #5710 before wire-up: 1) _pair_linked_indices kept the pending assistant tool_call window across an intervening user message. Per the OpenAI chat schema tool messages must follow their assistant directly, so a tool message arriving after a user is malformed input. Pairing it back to the earlier assistant just preserved invalid linkage and could force the compactor to drop a stale assistant + later-turn tool as one group. Clear pending_ids on any non-assistant / non-tool role so the window resets cleanly. 2) estimate_tokens crashed on non-dict entries inside ``tool_calls``. OpenAI payloads can carry malformed entries before pydantic validation; skip non-dicts and missing function shapes instead of raising AttributeError mid-compaction. 3) estimate_tokens ignored Studio's ``compaction`` content part. An Anthropic round-trip compaction summary can be multi-KB, so a compaction-only message previously estimated to zero tokens and slipped past the threshold check. Count the summary string. Tests grow by 4 (48 -> 52 total): two for the user-boundary pair reset, two for the defensive heuristic paths. --- .../core/inference/context_compaction.py | 42 ++++-- .../tests/test_context_compaction_edge.py | 124 ++++++++++++++++++ 2 files changed, 158 insertions(+), 8 deletions(-) diff --git a/studio/backend/core/inference/context_compaction.py b/studio/backend/core/inference/context_compaction.py index f4a52c675a..b825b62fd0 100644 --- a/studio/backend/core/inference/context_compaction.py +++ b/studio/backend/core/inference/context_compaction.py @@ -46,7 +46,10 @@ def estimate_tokens(messages: list[dict]) -> int: """Rough token count for ``messages``. Uses a 4-char-per-token char-count heuristic on the visible ``content`` (str) and on serialized ``tool_calls`` arguments. Multimodal parts (list-typed - content) contribute only their text parts. + content) contribute only their text parts. Studio's ``compaction`` + content parts (Anthropic round-trip state) also contribute their + ``content`` string so a multi-KB compaction summary does not + estimate as zero and slip past the threshold. """ total_chars = 0 for m in messages: @@ -55,14 +58,28 @@ def estimate_tokens(messages: list[dict]) -> int: total_chars += len(content) elif isinstance(content, list): for part in content: - if isinstance(part, dict): - t = part.get("text") - if isinstance(t, str): - total_chars += len(t) + if not isinstance(part, dict): + continue + t = part.get("text") + if isinstance(t, str): + total_chars += len(t) + # Studio's compaction content part: {"type":"compaction", + # "content": ""}. Count the summary string. + 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: - args = (tc.get("function") or {}).get("arguments") + # Defensive: pre-pydantic OpenAI payloads occasionally + # carry malformed entries (string, None) before + # validation. Skip non-dict items instead of raising + # AttributeError mid-compaction. + 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-style division: floor (`// 4`) would systematically @@ -104,17 +121,26 @@ def _pair_linked_indices(messages: list[dict]) -> dict[int, set[int]]: # Walk in order so the next tool-role messages after an assistant # call are the natural matches. A tool message is matched to the # most recent prior assistant whose ``tool_calls`` contain that id. + # Any non-assistant / non-tool message (user, system) ends the + # pending window: per the OpenAI chat schema tool messages must + # follow their assistant directly, so a tool message arriving + # after an intervening user is malformed input. Clearing here + # prevents the compactor from treating a stale assistant + a + # later-turn tool as one group and dropping them together. pending_ids: dict[str, int] = {} for i, m in enumerate(messages): - if m.get("role") == "assistant": + role = m.get("role") + if role == "assistant": ids = _assistant_tool_call_ids(m) out.setdefault(i, set()) for tid in ids: pending_ids[tid] = i - elif _is_tool_message(m): + 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 diff --git a/studio/backend/tests/test_context_compaction_edge.py b/studio/backend/tests/test_context_compaction_edge.py index b1397788bb..65f6b5c330 100644 --- a/studio/backend/tests/test_context_compaction_edge.py +++ b/studio/backend/tests/test_context_compaction_edge.py @@ -521,3 +521,127 @@ 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) From ac22ff2d326a68f1de28a8aae07540cab8eae194 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Mon, 25 May 2026 07:39:43 +0000 Subject: [PATCH 08/18] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- studio/backend/tests/test_context_compaction_edge.py | 12 ++++++++++-- 1 file changed, 10 insertions(+), 2 deletions(-) diff --git a/studio/backend/tests/test_context_compaction_edge.py b/studio/backend/tests/test_context_compaction_edge.py index 65f6b5c330..e2046b8efb 100644 --- a/studio/backend/tests/test_context_compaction_edge.py +++ b/studio/backend/tests/test_context_compaction_edge.py @@ -579,7 +579,11 @@ def test_paired_window_resumes_after_new_assistant_with_tool_call(): "role": "assistant", "content": "", "tool_calls": [ - {"id": "A", "type": "function", "function": {"name": "f", "arguments": "{}"}} + { + "id": "A", + "type": "function", + "function": {"name": "f", "arguments": "{}"}, + } ], }, {"role": "user", "content": "wait, do this instead"}, @@ -587,7 +591,11 @@ def test_paired_window_resumes_after_new_assistant_with_tool_call(): "role": "assistant", "content": "", "tool_calls": [ - {"id": "B", "type": "function", "function": {"name": "g", "arguments": "{}"}} + { + "id": "B", + "type": "function", + "function": {"name": "g", "arguments": "{}"}, + } ], }, { From 21c3b55fd923f17c6279fa68569cf8a7f60fd449 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Mon, 25 May 2026 08:00:59 +0000 Subject: [PATCH 09/18] Drop orphan tool messages and guard pair helper Two regressions surfaced by a 5000-iteration fuzz pass on sim_pr5710.py (deterministic seed 0xC0FFEE): 1) _assistant_tool_call_ids called tc.get("id") without an isinstance(tc, dict) guard, so a malformed pre-pydantic tool_calls entry (a bare string, a None) crashed _pair_linked_indices and the whole compactor with AttributeError. Mirrors the same guard the round-2 commit added to estimate_tokens. 2) When a tool message arrives after a user / system boundary and references an assistant tool_call_id that the boundary-clearing logic in _pair_linked_indices no longer treats as the pair root, the assistant gets dropped by the main sweep but the tool survives into the output. The result is a tool message whose tool_call_id has no matching assistant tool_calls earlier in the kept output; llama-server returns 400 on that template. Add a final invariant sweep that walks the surviving indices in order and drops any tool message whose tool_call_id is not introduced by a surviving assistant earlier in the output. Anchored (multimodal-content) tools stay regardless, matching the existing leak-rather-than-violate-anchor rule. Tests grow from 86 to 89. --- .../core/inference/context_compaction.py | 24 ++++++ .../tests/test_context_compaction_edge.py | 77 +++++++++++++++++++ 2 files changed, 101 insertions(+) diff --git a/studio/backend/core/inference/context_compaction.py b/studio/backend/core/inference/context_compaction.py index b825b62fd0..12491a22be 100644 --- a/studio/backend/core/inference/context_compaction.py +++ b/studio/backend/core/inference/context_compaction.py @@ -101,11 +101,16 @@ def _is_tool_message(msg: dict) -> bool: def _assistant_tool_call_ids(msg: dict) -> set[str]: """Return the set of ``id`` values from an assistant message's ``tool_calls``. Empty set when the message has no tool calls. + Mirrors ``estimate_tokens``: skip non-dict entries so malformed + pre-pydantic inputs (a string or ``None`` in the list) don't crash + ``_pair_linked_indices`` mid-compaction. """ 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) @@ -279,6 +284,25 @@ class SlidingWindowCompact(CompactStrategy): if asst_idx not in anchor_idx: dropped.add(asst_idx) + # Final invariant sweep: drop any surviving tool message whose + # tool_call_id has no matching assistant ``tool_calls`` earlier + # in the kept output. This catches orphans pair_map could not + # link -- e.g. a tool that arrives after a user boundary and + # references an assistant that the boundary-clearing logic in + # ``_pair_linked_indices`` no longer treats as the pair root. + # Anchored tool messages (multimodal content) stay regardless, + # matching the existing leak-rather-than-violate-anchor rule. + 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" and i not in anchor_idx: + tcid = m.get("tool_call_id") + if isinstance(tcid, str) and tcid and tcid not in seen_ids: + dropped.add(i) + return [m for i, m in enumerate(messages) if i not in dropped] diff --git a/studio/backend/tests/test_context_compaction_edge.py b/studio/backend/tests/test_context_compaction_edge.py index e2046b8efb..d1e60f662e 100644 --- a/studio/backend/tests/test_context_compaction_edge.py +++ b/studio/backend/tests/test_context_compaction_edge.py @@ -653,3 +653,80 @@ def test_estimate_tokens_counts_compaction_part(): ] # 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_does_not_drop_anchored_multimodal_tool(): + """An anchored tool (multimodal content) still survives even when + its matching assistant gets dropped -- the existing + leak-rather-than-violate-anchor rule applies in this direction too. + """ + msgs = [ + {"role": "system", "content": "sys"}, + {"role": "user", "content": "task " * 20}, + {"role": "assistant", "content": "a", "tool_calls": [_tool_call("t1")]}, + # Multimodal tool message -- anchored. + { + "role": "tool", + "tool_call_id": "t1", + "content": [{"type": "text", "text": "ok"}], + }, + {"role": "user", "content": "more"}, + {"role": "user", "content": "even more"}, + ] + out = SlidingWindowCompact(keep_recent=2).compact(msgs, budget_tokens=1) + # The anchored tool survives. + assert any( + m.get("role") == "tool" and m.get("tool_call_id") == "t1" for m in out + ) From 8508d2aea9ca455edbcb4818ef3c84a88128e21e Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Mon, 25 May 2026 08:01:22 +0000 Subject: [PATCH 10/18] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- .../backend/tests/test_context_compaction_edge.py | 14 ++++++-------- 1 file changed, 6 insertions(+), 8 deletions(-) diff --git a/studio/backend/tests/test_context_compaction_edge.py b/studio/backend/tests/test_context_compaction_edge.py index d1e60f662e..4f9c368d80 100644 --- a/studio/backend/tests/test_context_compaction_edge.py +++ b/studio/backend/tests/test_context_compaction_edge.py @@ -690,7 +690,7 @@ def test_compact_drops_orphan_tool_left_by_user_boundary(): {"role": "user", "content": "again"}, {"role": "tool", "tool_call_id": "t1", "content": "stale"}, ] - out = SlidingWindowCompact(keep_recent=2).compact(msgs, budget_tokens=1) + 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() @@ -702,9 +702,9 @@ def test_compact_drops_orphan_tool_left_by_user_boundary(): 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}" - ) + assert ( + isinstance(tcid, str) and tcid in seen_ids + ), f"orphan tool {tcid!r} survived; seen={seen_ids}" def test_compact_does_not_drop_anchored_multimodal_tool(): @@ -725,8 +725,6 @@ def test_compact_does_not_drop_anchored_multimodal_tool(): {"role": "user", "content": "more"}, {"role": "user", "content": "even more"}, ] - out = SlidingWindowCompact(keep_recent=2).compact(msgs, budget_tokens=1) + out = SlidingWindowCompact(keep_recent = 2).compact(msgs, budget_tokens = 1) # The anchored tool survives. - assert any( - m.get("role") == "tool" and m.get("tool_call_id") == "t1" for m in out - ) + assert any(m.get("role") == "tool" and m.get("tool_call_id") == "t1" for m in out) From 71cf01b171e64dcce5d533b4c145814f04b5850c Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Mon, 25 May 2026 13:24:24 +0000 Subject: [PATCH 11/18] Strip orphan tool_calls from surviving assistants Anchored multimodal assistants kept their tool_calls field even when every matching tool follow-up got dropped. The per-index budget loop could also leave a plain assistant with only one of its tool_calls matched. Both shapes 400 on OpenAI / strict OpenAI-compatible servers ("assistant message with tool_calls must be followed by tool messages responding to each tool_call_id"). Rewrite the surviving assistant in a copy so the multimodal content survives and the request stays valid. --- .../core/inference/context_compaction.py | 52 ++++++++++++++-- .../tests/test_context_compaction_edge.py | 61 +++++++++++++++++++ 2 files changed, 107 insertions(+), 6 deletions(-) diff --git a/studio/backend/core/inference/context_compaction.py b/studio/backend/core/inference/context_compaction.py index 12491a22be..68a31d5d41 100644 --- a/studio/backend/core/inference/context_compaction.py +++ b/studio/backend/core/inference/context_compaction.py @@ -273,16 +273,20 @@ class SlidingWindowCompact(CompactStrategy): # enforce, and llama-server would 400 on the resulting # template (a tool message whose tool_call_id has no # surviving assistant tool_calls entry). + # Assistants whose tool_calls all got orphaned are repaired in + # the final pass below: we keep the multimodal content but + # strip the dangling tool_calls so OpenAI does not 400 on + # "assistant message with tool_calls must be followed by tool + # messages". + 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: - # All matching tool messages were dropped: drop the - # 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) + else: + rewrite_strip_tool_calls.add(asst_idx) # Final invariant sweep: drop any surviving tool message whose # tool_call_id has no matching assistant ``tool_calls`` earlier @@ -292,18 +296,54 @@ class SlidingWindowCompact(CompactStrategy): # ``_pair_linked_indices`` no longer treats as the pair root. # Anchored tool messages (multimodal content) stay regardless, # matching the existing leak-rather-than-violate-anchor rule. + # Also collect assistant indices whose surviving tool_calls + # entries lack a matching tool follow-up so we can strip the + # orphan ids from a copy below (same reason: OpenAI 400s on + # tool_calls without matching tool responses). + 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) 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) + ids = _assistant_tool_call_ids(m) + seen_ids |= ids + if ids and not (ids <= responded_ids): + rewrite_strip_tool_calls.add(i) elif m.get("role") == "tool" and i not in anchor_idx: tcid = m.get("tool_call_id") if isinstance(tcid, str) and tcid and tcid not in seen_ids: dropped.add(i) - return [m for i, m in enumerate(messages) if i not in dropped] + out: list[dict] = [] + for i, m in enumerate(messages): + if i in dropped: + continue + if i in rewrite_strip_tool_calls: + # Keep the message (multimodal content stays) but strip + # tool_calls entries with no surviving tool follow-up. + 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] = { diff --git a/studio/backend/tests/test_context_compaction_edge.py b/studio/backend/tests/test_context_compaction_edge.py index 4f9c368d80..370fa8884e 100644 --- a/studio/backend/tests/test_context_compaction_edge.py +++ b/studio/backend/tests/test_context_compaction_edge.py @@ -728,3 +728,64 @@ def test_compact_does_not_drop_anchored_multimodal_tool(): out = SlidingWindowCompact(keep_recent = 2).compact(msgs, budget_tokens = 1) # The anchored tool survives. assert any(m.get("role") == "tool" and m.get("tool_call_id") == "t1" for m in out) + + +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_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) From 585d7bdc9449da918badb01648e178769e4d471d Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Mon, 25 May 2026 13:27:17 +0000 Subject: [PATCH 12/18] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- studio/backend/core/inference/context_compaction.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/studio/backend/core/inference/context_compaction.py b/studio/backend/core/inference/context_compaction.py index 68a31d5d41..efce545d30 100644 --- a/studio/backend/core/inference/context_compaction.py +++ b/studio/backend/core/inference/context_compaction.py @@ -330,7 +330,8 @@ class SlidingWindowCompact(CompactStrategy): # Keep the message (multimodal content stays) but strip # tool_calls entries with no surviving tool follow-up. kept_tcs = [ - tc for tc in (m.get("tool_calls") or []) + 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 From 071fef372a49912d43e8485cff4cb56401a0812a Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Mon, 25 May 2026 13:38:44 +0000 Subject: [PATCH 13/18] Anchor full tool-pair; close pair on assistant boundary; skip compaction-part anchors Three correctness fixes surfaced by the round-3 audit: 1. Anchor propagation across paired indices. An anchored multimodal assistant or tool now anchors the full asst+tools group; previously only one half survived and the chat template was orphaned. 2. _pair_linked_indices clears its pending window on a new assistant turn (was already clearing on user/system). A malformed history like asst(tc=a), asst("plain"), tool(a) no longer pairs the late tool back to the stale assistant. 3. _is_multimodal excludes text-only and Studio "compaction" content parts. The previous list-content-is-multimodal check pinned old compaction summaries forever and kept oversized prompts above the threshold. --- .../core/inference/context_compaction.py | 48 ++++++-- .../tests/test_context_compaction_edge.py | 111 ++++++++++++++++-- 2 files changed, 140 insertions(+), 19 deletions(-) diff --git a/studio/backend/core/inference/context_compaction.py b/studio/backend/core/inference/context_compaction.py index efce545d30..96d091d2a3 100644 --- a/studio/backend/core/inference/context_compaction.py +++ b/studio/backend/core/inference/context_compaction.py @@ -90,12 +90,25 @@ def estimate_tokens(messages: list[dict]) -> int: return -(-total_chars // _CHARS_PER_TOKEN) +# Content-part types that carry no media payload and shouldn't anchor the +# message. ``compaction`` is Studio's Anthropic round-trip state -- pinning +# it would keep the very thing we're trying to compact away. +_TEXT_ONLY_PART_TYPES = {"text", "compaction"} + + def _is_multimodal(msg: dict) -> bool: - return isinstance(msg.get("content"), list) - - -def _is_tool_message(msg: dict) -> bool: - return msg.get("role") == "tool" + content = msg.get("content") + if not isinstance(content, list): + return False + for part in content: + # Unknown shapes (raw strings, ints, None) keep the conservative + # "treat as multimodal" stance -- there's no test rendering for + # them either. + 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]: @@ -126,16 +139,17 @@ def _pair_linked_indices(messages: list[dict]) -> dict[int, set[int]]: # Walk in order so the next tool-role messages after an assistant # call are the natural matches. A tool message is matched to the # most recent prior assistant whose ``tool_calls`` contain that id. - # Any non-assistant / non-tool message (user, system) ends the - # pending window: per the OpenAI chat schema tool messages must - # follow their assistant directly, so a tool message arriving - # after an intervening user is malformed input. Clearing here - # prevents the compactor from treating a stale assistant + a - # later-turn tool as one group and dropping them together. + # ANY non-tool boundary (user, system, OR another assistant) ends + # the pending window: per the OpenAI chat schema tool messages + # must follow their assistant directly. A later assistant turn + # arriving before the matching tool means that tool is malformed + # input -- treating the late tool as still paired would let the + # compactor drop a stale assistant + a later-turn tool together. 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: @@ -214,6 +228,18 @@ class SlidingWindowCompact(CompactStrategy): # message and its matching tool-role responses get the same # group id so we keep or drop them as a unit. pair_map = _pair_linked_indices(messages) + # If any member of a valid asst+tools group is anchored, + # anchor the whole group. Without this an anchored multimodal + # tool whose assistant is droppable (or an anchored multimodal + # assistant whose tools are droppable) would survive alone, + # leaving the chat template invalid (OpenAI 400s on dangling + # tool_calls / orphan tool messages). + 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(): diff --git a/studio/backend/tests/test_context_compaction_edge.py b/studio/backend/tests/test_context_compaction_edge.py index 370fa8884e..ac4071ef03 100644 --- a/studio/backend/tests/test_context_compaction_edge.py +++ b/studio/backend/tests/test_context_compaction_edge.py @@ -707,27 +707,34 @@ def test_compact_drops_orphan_tool_left_by_user_boundary(): ), f"orphan tool {tcid!r} survived; seen={seen_ids}" -def test_compact_does_not_drop_anchored_multimodal_tool(): - """An anchored tool (multimodal content) still survives even when - its matching assistant gets dropped -- the existing - leak-rather-than-violate-anchor rule applies in this direction too. +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")]}, - # Multimodal tool message -- anchored. + # Real multimodal tool message (image part) -- anchored. { "role": "tool", "tool_call_id": "t1", - "content": [{"type": "text", "text": "ok"}], + "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) - # The anchored tool survives. - assert any(m.get("role") == "tool" and m.get("tool_call_id") == "t1" for m in out) + # 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(): @@ -764,6 +771,94 @@ def test_anchored_multimodal_asst_orphan_tool_calls_stripped(): 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 From ad6a34502e884524a6c90c820f9196ae8cc05737 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Mon, 25 May 2026 13:39:05 +0000 Subject: [PATCH 14/18] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- studio/backend/tests/test_context_compaction_edge.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/studio/backend/tests/test_context_compaction_edge.py b/studio/backend/tests/test_context_compaction_edge.py index ac4071ef03..e7b367573b 100644 --- a/studio/backend/tests/test_context_compaction_edge.py +++ b/studio/backend/tests/test_context_compaction_edge.py @@ -852,8 +852,7 @@ def test_compaction_content_part_is_not_multimodal_anchor(): 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"] + isinstance(p, dict) and p.get("type") == "compaction" for p in m["content"] ) for m in out ) From 61bb76281aac901719ed87e820731b7e3ba5bbc7 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Tue, 26 May 2026 00:47:02 +0000 Subject: [PATCH 15/18] Recompute responded_ids after orphan-tool drop for PR #5710 Before, the final sweep computed `responded_ids` once before pruning orphan tools, so a tool message arriving before its assistant got dropped as orphan but the later assistant still matched `ids <= responded_ids` and kept its dangling `tool_calls` entry, violating the chat-template invariant the sweep exists to enforce. Split the sweep into two passes: drop orphan tools first, then recompute `responded_ids` from the surviving tools and mark assts whose tool_calls aren't all responded for strip. --- .../core/inference/context_compaction.py | 38 ++++++++++--------- .../tests/test_context_compaction_edge.py | 24 ++++++++++++ 2 files changed, 44 insertions(+), 18 deletions(-) diff --git a/studio/backend/core/inference/context_compaction.py b/studio/backend/core/inference/context_compaction.py index 96d091d2a3..4572e702b4 100644 --- a/studio/backend/core/inference/context_compaction.py +++ b/studio/backend/core/inference/context_compaction.py @@ -314,18 +314,26 @@ class SlidingWindowCompact(CompactStrategy): else: rewrite_strip_tool_calls.add(asst_idx) - # Final invariant sweep: drop any surviving tool message whose - # tool_call_id has no matching assistant ``tool_calls`` earlier - # in the kept output. This catches orphans pair_map could not - # link -- e.g. a tool that arrives after a user boundary and - # references an assistant that the boundary-clearing logic in - # ``_pair_linked_indices`` no longer treats as the pair root. - # Anchored tool messages (multimodal content) stay regardless, - # matching the existing leak-rather-than-violate-anchor rule. - # Also collect assistant indices whose surviving tool_calls - # entries lack a matching tool follow-up so we can strip the - # orphan ids from a copy below (same reason: OpenAI 400s on - # tool_calls without matching tool responses). + # Final invariant sweep, two passes so the asst-strip decision + # sees the post-orphan-drop tool set: + # pass 1: drop tools whose tcid has no matching assistant + # ``tool_calls`` earlier in the kept output (anchored tools + # stay; same leak-rather-than-violate-anchor rule); + # pass 2: recompute responded_ids from the surviving tools and + # mark assts whose tool_calls aren't all responded for strip. + # Computing responded_ids before pass 1 would let a stale tcid + # of a just-dropped orphan tool satisfy `ids <= responded_ids`, + # leaving the asst with 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" and i not in anchor_idx: + 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: @@ -334,19 +342,13 @@ class SlidingWindowCompact(CompactStrategy): tcid = m.get("tool_call_id") if isinstance(tcid, str) and tcid: responded_ids.add(tcid) - seen_ids: set[str] = set() for i, m in enumerate(messages): if i in dropped: continue if m.get("role") == "assistant": ids = _assistant_tool_call_ids(m) - seen_ids |= ids if ids and not (ids <= responded_ids): rewrite_strip_tool_calls.add(i) - elif m.get("role") == "tool" and i not in anchor_idx: - tcid = m.get("tool_call_id") - if isinstance(tcid, str) and tcid and tcid not in seen_ids: - dropped.add(i) out: list[dict] = [] for i, m in enumerate(messages): diff --git a/studio/backend/tests/test_context_compaction_edge.py b/studio/backend/tests/test_context_compaction_edge.py index e7b367573b..601c96b544 100644 --- a/studio/backend/tests/test_context_compaction_edge.py +++ b/studio/backend/tests/test_context_compaction_edge.py @@ -883,3 +883,27 @@ def test_partial_tool_drop_strips_orphan_tool_call_id(): 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) From 3c248dca26ac6bf8d6ca289aed84c550efa6ef44 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Tue, 26 May 2026 01:03:10 +0000 Subject: [PATCH 16/18] Drop anchored orphan tools to preserve template validity for PR #5710 The final invariant sweep skipped anchored (multimodal) tool messages whose tcid had no matching prior assistant tool_calls. The anchor survived but the chat template did not -- a tool_call_id pointing at nothing returns 400 from llama-server / OpenAI on the next call. Pair-validity is the hard invariant (the module exists to enforce it upstream of llama-server) and the multimodal-anchor rule is the soft quality preference. Drop the orphan rather than violate the template: lose the image content for that one turn, keep the conversation alive. --- .../core/inference/context_compaction.py | 11 ++++++-- .../tests/test_context_compaction_edge.py | 28 +++++++++++++++++++ 2 files changed, 36 insertions(+), 3 deletions(-) diff --git a/studio/backend/core/inference/context_compaction.py b/studio/backend/core/inference/context_compaction.py index 4572e702b4..e9a5d0ca46 100644 --- a/studio/backend/core/inference/context_compaction.py +++ b/studio/backend/core/inference/context_compaction.py @@ -317,20 +317,25 @@ class SlidingWindowCompact(CompactStrategy): # Final invariant sweep, two passes so the asst-strip decision # sees the post-orphan-drop tool set: # pass 1: drop tools whose tcid has no matching assistant - # ``tool_calls`` earlier in the kept output (anchored tools - # stay; same leak-rather-than-violate-anchor rule); + # ``tool_calls`` earlier in the kept output; # pass 2: recompute responded_ids from the surviving tools and # mark assts whose tool_calls aren't all responded for strip. # Computing responded_ids before pass 1 would let a stale tcid # of a just-dropped orphan tool satisfy `ids <= responded_ids`, # leaving the asst with dangling tool_calls (OpenAI 400). + # Anchored (multimodal) orphan tools are still dropped here: + # keeping them violates the chat-template pair invariant (a hard + # invariant that 400s upstream) to honor the multimodal-anchor + # quality preference, which the docstring describes as a soft + # quality rule. The pair-validity wins; the image content is + # lost rather than the entire turn. 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" and i not in anchor_idx: + 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) diff --git a/studio/backend/tests/test_context_compaction_edge.py b/studio/backend/tests/test_context_compaction_edge.py index 601c96b544..75fa1f0345 100644 --- a/studio/backend/tests/test_context_compaction_edge.py +++ b/studio/backend/tests/test_context_compaction_edge.py @@ -907,3 +907,31 @@ def test_orphan_tool_before_asst_does_not_leave_dangling_tool_calls(): 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) From fcdd5fc88e217bfb0d7c898cb92b053ed82cf40b Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Tue, 26 May 2026 05:46:02 +0000 Subject: [PATCH 17/18] Shorten comments in compaction module for PR #5710 --- .../core/inference/context_compaction.py | 145 +++++------------- 1 file changed, 37 insertions(+), 108 deletions(-) diff --git a/studio/backend/core/inference/context_compaction.py b/studio/backend/core/inference/context_compaction.py index e9a5d0ca46..0aa7e327b0 100644 --- a/studio/backend/core/inference/context_compaction.py +++ b/studio/backend/core/inference/context_compaction.py @@ -36,20 +36,15 @@ from abc import ABC, abstractmethod from typing import Any -# Char-to-token heuristic. Conservative on the high side so the -# compactor triggers earlier rather than later. Tokenizer-aware -# estimates may land in a follow-up. +# Conservative char-to-token ratio so compaction fires early. _CHARS_PER_TOKEN = 4 def estimate_tokens(messages: list[dict]) -> int: - """Rough token count for ``messages``. Uses a 4-char-per-token - char-count heuristic on the visible ``content`` (str) and on - serialized ``tool_calls`` arguments. Multimodal parts (list-typed - content) contribute only their text parts. Studio's ``compaction`` - content parts (Anthropic round-trip state) also contribute their - ``content`` string so a multi-KB compaction summary does not - estimate as zero and slip past the threshold. + """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: @@ -63,8 +58,7 @@ def estimate_tokens(messages: list[dict]) -> int: t = part.get("text") if isinstance(t, str): total_chars += len(t) - # Studio's compaction content part: {"type":"compaction", - # "content": ""}. Count the summary string. + # Count compaction summary text. if part.get("type") == "compaction": summary = part.get("content") if isinstance(summary, str): @@ -72,27 +66,19 @@ def estimate_tokens(messages: list[dict]) -> int: tcs = m.get("tool_calls") if isinstance(tcs, list): for tc in tcs: - # Defensive: pre-pydantic OpenAI payloads occasionally - # carry malformed entries (string, None) before - # validation. Skip non-dict items instead of raising - # AttributeError mid-compaction. + # 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-style division: floor (`// 4`) would systematically - # underestimate non-multiple-of-4 lengths, letting just-over-budget - # prompts appear under threshold and bypass compaction — the exact - # failure this module exists to prevent. Round up so the heuristic - # stays on the conservative side described above. + # Ceil-divide: floor would let just-over-budget prompts bypass compaction. return -(-total_chars // _CHARS_PER_TOKEN) -# Content-part types that carry no media payload and shouldn't anchor the -# message. ``compaction`` is Studio's Anthropic round-trip state -- pinning -# it would keep the very thing we're trying to compact away. +# Text-only parts (don't anchor); ``compaction`` is Anthropic round-trip +# state so pinning it would defeat compaction. _TEXT_ONLY_PART_TYPES = {"text", "compaction"} @@ -101,9 +87,7 @@ def _is_multimodal(msg: dict) -> bool: if not isinstance(content, list): return False for part in content: - # Unknown shapes (raw strings, ints, None) keep the conservative - # "treat as multimodal" stance -- there's no test rendering for - # them either. + # Unknown shape => conservatively treat as multimodal. if not isinstance(part, dict): return True if part.get("type") not in _TEXT_ONLY_PART_TYPES: @@ -112,12 +96,7 @@ def _is_multimodal(msg: dict) -> bool: def _assistant_tool_call_ids(msg: dict) -> set[str]: - """Return the set of ``id`` values from an assistant message's - ``tool_calls``. Empty set when the message has no tool calls. - Mirrors ``estimate_tokens``: skip non-dict entries so malformed - pre-pydantic inputs (a string or ``None`` in the list) don't crash - ``_pair_linked_indices`` mid-compaction. - """ + """Return tool_call ids on an assistant message; skip malformed entries.""" out: set[str] = set() tcs = msg.get("tool_calls") if isinstance(tcs, list): @@ -131,20 +110,13 @@ def _assistant_tool_call_ids(msg: dict) -> set[str]: def _pair_linked_indices(messages: list[dict]) -> dict[int, set[int]]: - """Map an assistant-message index to the indices of its tool-role - follow-ups (matching ``tool_call_id``). Used so the compactor drops - or keeps an assistant+tool group as a unit. + """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]] = {} - # Walk in order so the next tool-role messages after an assistant - # call are the natural matches. A tool message is matched to the - # most recent prior assistant whose ``tool_calls`` contain that id. - # ANY non-tool boundary (user, system, OR another assistant) ends - # the pending window: per the OpenAI chat schema tool messages - # must follow their assistant directly. A later assistant turn - # arriving before the matching tool means that tool is malformed - # input -- treating the late tool as still paired would let the - # compactor drop a stale assistant + a later-turn tool together. + # 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") @@ -183,14 +155,8 @@ class NoCompact(CompactStrategy): class SlidingWindowCompact(CompactStrategy): - """Keep the system message, the first user message, and the last - ``keep_recent`` non-droppable turns. Multimodal turns are never - dropped (no compacted-media representation today). Assistant - messages with ``tool_calls`` are grouped with their matching - tool-role responses and treated as one unit. - - The strategy is a no-op when the estimated token count is already - within ``budget_tokens`` or when there is nothing left to drop. + """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: @@ -224,16 +190,10 @@ class SlidingWindowCompact(CompactStrategy): if _is_multimodal(m): anchor_idx.add(i) - # Tool-call / tool-result grouping. An assistant tool-call - # message and its matching tool-role responses get the same - # group id so we keep or drop them as a unit. + # Group asst tool-call with its tool responses so they drop together. pair_map = _pair_linked_indices(messages) - # If any member of a valid asst+tools group is anchored, - # anchor the whole group. Without this an anchored multimodal - # tool whose assistant is droppable (or an anchored multimodal - # assistant whose tools are droppable) would survive alone, - # leaving the chat template invalid (OpenAI 400s on dangling - # tool_calls / orphan tool 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 @@ -251,13 +211,8 @@ class SlidingWindowCompact(CompactStrategy): for ti in tool_idxs: group_id[ti] = g - # The "recent window" is the last ``keep_recent`` distinct - # 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. + # 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): @@ -270,8 +225,7 @@ class SlidingWindowCompact(CompactStrategy): recent_groups.append(g) recent_groups_set = set(recent_groups) - # Decide drop set: every index whose group is NOT in the recent - # window AND that is not an anchor. + # Droppable = non-anchor indices outside the recent window. droppable: list[int] = [] for i in range(len(messages)): if i in anchor_idx: @@ -289,21 +243,9 @@ class SlidingWindowCompact(CompactStrategy): break dropped.add(i) - # 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). - # Assistants whose tool_calls all got orphaned are repaired in - # the final pass below: we keep the multimodal content but - # strip the dangling tool_calls so OpenAI does not 400 on - # "assistant message with tool_calls must be followed by tool - # messages". + # 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: @@ -314,21 +256,13 @@ class SlidingWindowCompact(CompactStrategy): else: rewrite_strip_tool_calls.add(asst_idx) - # Final invariant sweep, two passes so the asst-strip decision - # sees the post-orphan-drop tool set: - # pass 1: drop tools whose tcid has no matching assistant - # ``tool_calls`` earlier in the kept output; - # pass 2: recompute responded_ids from the surviving tools and - # mark assts whose tool_calls aren't all responded for strip. - # Computing responded_ids before pass 1 would let a stale tcid - # of a just-dropped orphan tool satisfy `ids <= responded_ids`, - # leaving the asst with dangling tool_calls (OpenAI 400). - # Anchored (multimodal) orphan tools are still dropped here: - # keeping them violates the chat-template pair invariant (a hard - # invariant that 400s upstream) to honor the multimodal-anchor - # quality preference, which the docstring describes as a soft - # quality rule. The pair-validity wins; the image content is - # lost rather than the entire turn. + # 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: @@ -360,8 +294,7 @@ class SlidingWindowCompact(CompactStrategy): if i in dropped: continue if i in rewrite_strip_tool_calls: - # Keep the message (multimodal content stays) but strip - # tool_calls entries with no surviving tool follow-up. + # Keep content but strip dangling tool_calls. kept_tcs = [ tc for tc in (m.get("tool_calls") or []) @@ -387,9 +320,5 @@ _STRATEGIES: dict[str, CompactStrategy] = { def get_strategy(name: str) -> CompactStrategy: - """Return the compaction strategy registered under ``name``. - - Falls back to ``NoCompact`` for unknown names so a misconfigured - request degrades to no-op rather than raising. - """ + """Return strategy by name; unknown names fall back to ``NoCompact``.""" return _STRATEGIES.get(name, _STRATEGIES["none"]) From cc96dc8aad49585dad12d459552e54594893c996 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Tue, 16 Jun 2026 07:59:29 +0000 Subject: [PATCH 18/18] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- .../core/inference/context_compaction.py | 6 ++- .../backend/tests/test_context_compaction.py | 38 +++++++++++-------- .../tests/test_context_compaction_edge.py | 32 ++++++++++------ 3 files changed, 48 insertions(+), 28 deletions(-) diff --git a/studio/backend/core/inference/context_compaction.py b/studio/backend/core/inference/context_compaction.py index 0aa7e327b0..f8ae4c3289 100644 --- a/studio/backend/core/inference/context_compaction.py +++ b/studio/backend/core/inference/context_compaction.py @@ -159,7 +159,11 @@ class SlidingWindowCompact(CompactStrategy): groups. Asst+tools form one group. No-op when within budget. """ - def __init__(self, keep_recent: int = 2, compact_threshold: float = 0.85) -> None: + 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): diff --git a/studio/backend/tests/test_context_compaction.py b/studio/backend/tests/test_context_compaction.py index ccce541e84..9a95ca5561 100644 --- a/studio/backend/tests/test_context_compaction.py +++ b/studio/backend/tests/test_context_compaction.py @@ -21,7 +21,13 @@ from core.inference.context_compaction import ( ) -def _msg(role, content = "", tool_calls = None, tool_call_id = None, name = None): +def _msg( + role, + content = "", + tool_calls = None, + tool_call_id = None, + name = None, +): m = {"role": role} if content is not None: m["content"] = content @@ -34,7 +40,12 @@ def _msg(role, content = "", tool_calls = None, tool_call_id = None, name = None return m -def _long(role, length, *, content_prefix = "x"): +def _long( + role, + length, + *, + content_prefix = "x", +): return _msg(role, content_prefix * length) @@ -117,7 +128,11 @@ class TestSlidingWindowUnderBudget: class TestSlidingWindowInvariants: - def _make_long_chat(self, n_turns, length_per_turn = 1000): + 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. @@ -131,10 +146,7 @@ class TestSlidingWindowInvariants: 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 - ) + 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) @@ -209,9 +221,7 @@ class TestSlidingWindowToolPairs: # 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) + 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 @@ -258,8 +268,7 @@ class TestSlidingWindowToolPairs: ids = [ ( m.get("role"), - m.get("tool_call_id") - or (m.get("tool_calls") and m["tool_calls"][0].get("id")), + m.get("tool_call_id") or (m.get("tool_calls") and m["tool_calls"][0].get("id")), ) for m in out ] @@ -350,21 +359,18 @@ class TestAnchoredMultimodalPairCleanup: 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 + 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): diff --git a/studio/backend/tests/test_context_compaction_edge.py b/studio/backend/tests/test_context_compaction_edge.py index 75fa1f0345..e4329deae7 100644 --- a/studio/backend/tests/test_context_compaction_edge.py +++ b/studio/backend/tests/test_context_compaction_edge.py @@ -34,7 +34,13 @@ from core.inference.context_compaction import ( ) -def _msg(role, content = "", tool_calls = None, tool_call_id = None, name = None): +def _msg( + role, + content = "", + tool_calls = None, + tool_call_id = None, + name = None, +): m = {"role": role} if content is not None: m["content"] = content @@ -47,11 +53,20 @@ def _msg(role, content = "", tool_calls = None, tool_call_id = None, name = None return m -def _long(role, length, *, content_prefix = "x"): +def _long( + role, + length, + *, + content_prefix = "x", +): return _msg(role, content_prefix * length) -def _tool_call(tcid, name = "web_search", args = '{"q":"x"}'): +def _tool_call( + tcid, + name = "web_search", + args = '{"q":"x"}', +): return { "id": tcid, "type": "function", @@ -272,9 +287,7 @@ class TestDuplicateToolCallIds: ] 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_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} @@ -355,8 +368,7 @@ class TestNoSystemNoFirstUser: 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 + m.get("role") == "user" and "real first task" in m.get("content", "") for m in out ) @@ -851,9 +863,7 @@ def test_compaction_content_part_is_not_multimodal_anchor(): 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"] - ) + and any(isinstance(p, dict) and p.get("type") == "compaction" for p in m["content"]) for m in out )