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)