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

Break tool pair window on user boundary, harden estimator

Browse source 
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.
This commit is contained in:
Daniel Han 2026-05-25 07:39:04 +00:00
commit a12bc349f2
2 changed files with 158 additions and 8 deletions
Download patch file Download diff file Expand all files Collapse all files

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

@ -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": "<summary>"}. 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

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

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