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

Strip orphan tool_calls from surviving assistants

Browse source 
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.
This commit is contained in:
Daniel Han 2026-05-25 13:24:24 +00:00
commit 71cf01b171
2 changed files with 107 additions and 6 deletions
Download patch file Download diff file Expand all files Collapse all files

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

@ -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] = {

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

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