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)