ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Anchor full tool-pair; close pair on assistant boundary; skip compaction-part anchors

Browse source 
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.
This commit is contained in:
Daniel Han 2026-05-25 13:38:44 +00:00
commit 071fef372a
2 changed files with 139 additions and 18 deletions
Download patch file Download diff file Expand all files Collapse all files

48
studio/backend/core/inference/context_compaction.py
View file

@ -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():

111
studio/backend/tests/test_context_compaction_edge.py
View file

@ -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