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

Studio: edge-case coverage for context-compaction module

Browse source 
Adds an adversarial test file alongside the happy-path tests for the
new sliding-window compaction strategy. Pins down the structural
invariants under malformed and boundary inputs that the original
suite does not exercise:

* Empty / single-message / over-budget-anchor inputs.
* keep_recent boundary cases (>> len(messages), == group count).
* compact_threshold == 1.0 boundary and negative-threshold validation.
* Assistant messages carrying both content and tool_calls.
* Assistant content == None (OpenAI schema when only tool_calls set).
* Out-of-order tool messages (malformed, must not crash).
* Orphan tool messages and duplicated tool_call_ids across assistants.
* Non-string content (int, dict) defensive paths in estimate_tokens.
* Histories without a system message or with the first user not at
  index 1.
* First-user message that is itself multimodal (anchor overlap).
* Multi-tool-call single-assistant grouping.
* Interleaved user message between assistant tool_call and its tool
  response.
* NoCompact applied to an empty list.

Module-only PR, still unwired on main; no behavior change.
This commit is contained in:
Daniel Han 2026-05-24 13:53:53 +00:00
commit 15aa69dac8
1 changed files with 523 additions and 0 deletions
Download patch file Download diff file Expand all files Collapse all files

523
studio/backend/tests/test_context_compaction_edge.py Normal file
View file

@ -0,0 +1,523 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Edge-case coverage for the context-compaction module.
These complement ``test_context_compaction.py`` with adversarial /
boundary inputs the happy-path tests do not exercise. The goal is to
pin down the structural invariants under malformed or extreme inputs:
* Empty / tiny / oversize inputs.
* Floating / boundary constructor arguments.
* Assistant messages whose ``content`` is None or carries both text and
tool_calls at the same time.
* Tool messages without a matching assistant, or that arrive before
their assistant (malformed but must not crash).
* Duplicate or non-string ``tool_call_id`` values.
* ``content`` that is neither str nor list (defensive).
* No-system / no-first-user histories.
* Threshold and recent-window boundaries (``compact_threshold == 1.0``,
``keep_recent > len(messages)``).
* The tool-pair "drop one side, drop the other" rule under awkward
inputs (orphan tool, tool-without-id, anchored multimodal user with
a tool message tied to a different assistant).
"""
import pytest
from core.inference.context_compaction import (
NoCompact,
SlidingWindowCompact,
_pair_linked_indices,
estimate_tokens,
get_strategy,
)
def _msg(role, content = "", tool_calls = None, tool_call_id = None, name = None):
m = {"role": role}
if content is not None:
m["content"] = content
if tool_calls is not None:
m["tool_calls"] = tool_calls
if tool_call_id is not None:
m["tool_call_id"] = tool_call_id
if name is not None:
m["name"] = name
return m
def _long(role, length, *, content_prefix = "x"):
return _msg(role, content_prefix * length)
def _tool_call(tcid, name = "web_search", args = '{"q":"x"}'):
return {
"id": tcid,
"type": "function",
"function": {"name": name, "arguments": args},
}
def _surviving_tool_ids(messages):
"""Tool-call ids that a chat template would expect to match.
Returns (asst_call_ids, tool_response_ids). For a structurally valid
OpenAI chat-completions request, the two sets must be equal.
"""
asst_ids = set()
tool_ids = set()
for m in messages:
if m.get("role") == "assistant":
for tc in m.get("tool_calls") or []:
if isinstance(tc, dict) and isinstance(tc.get("id"), str):
asst_ids.add(tc["id"])
elif m.get("role") == "tool":
tcid = m.get("tool_call_id")
if isinstance(tcid, str):
tool_ids.add(tcid)
return asst_ids, tool_ids
class TestEmptyAndTrivialInputs:
def test_empty_messages_returns_empty(self):
# estimate_tokens([]) == 0 already covered; here check compact().
out = SlidingWindowCompact().compact([], budget_tokens = 100)
assert out == []
def test_single_message_under_budget_passthrough(self):
msgs = [_msg("user", "hi")]
out = SlidingWindowCompact().compact(msgs, budget_tokens = 100)
assert out == msgs
# Defensive: result is a copy, not the same list object.
assert out is not msgs
def test_single_anchor_over_budget_stays(self):
# Only message is the first-user anchor. Anchors are never
# dropped, even when alone they exceed budget. Result will be
# over budget — that is acceptable per the strategy's docstring
# ("nothing left to drop"), and the compactor must not loop
# forever or crash trying.
msgs = [_long("user", 10000)]
out = SlidingWindowCompact(keep_recent = 0).compact(msgs, budget_tokens = 5)
assert out == msgs
class TestKeepRecentBoundaries:
def test_keep_recent_larger_than_message_count(self):
# keep_recent >> len(messages): recent window covers everything,
# so the strategy is effectively a no-op even when over budget.
msgs = [
_msg("system", "sys"),
_msg("user", "first task"),
_long("assistant", 5000),
_long("user", 5000),
]
out = SlidingWindowCompact(keep_recent = 100).compact(msgs, budget_tokens = 10)
assert out == msgs
def test_keep_recent_equals_message_count_after_anchors(self):
# keep_recent exactly equal to the non-anchor group count is the
# boundary between "everything kept" and "drop at least one".
msgs = [
_msg("system", "sys"),
_msg("user", "first task"),
_long("assistant", 5000),
_long("user", 5000),
_long("assistant", 5000),
]
out = SlidingWindowCompact(keep_recent = 3).compact(msgs, budget_tokens = 10)
# All non-anchor groups (the trailing 3) are in the recent
# window. Result is the full list.
assert out == msgs
class TestThresholdBoundary:
def test_compact_threshold_one_exact(self):
# threshold == 1.0 is the inclusive boundary; estimate exactly
# at budget should be considered "fits" and skip compaction.
msgs = [_msg("user", "abcd")] # 1 token
out = SlidingWindowCompact(
keep_recent = 0,
compact_threshold = 1.0,
).compact(msgs, budget_tokens = 1)
assert out == msgs
def test_threshold_just_below_estimate_triggers_compaction(self):
# threshold (== int(budget * compact_threshold)) below estimate
# should compact.
msgs = [
_msg("system", "sys"),
_msg("user", "first task"),
_long("assistant", 4000),
_long("user", 4000),
_long("assistant", 4000),
_long("user", 4000),
]
before = estimate_tokens(msgs)
assert before > 100
out = SlidingWindowCompact(keep_recent = 1).compact(msgs, budget_tokens = 100)
assert len(out) < len(msgs)
def test_constructor_rejects_negative_threshold(self):
with pytest.raises(ValueError):
SlidingWindowCompact(compact_threshold = -0.1)
class TestAssistantWithBothContentAndToolCalls:
def test_assistant_with_text_and_tool_calls_kept_together(self):
# OpenAI chat completions schema allows an assistant message to
# carry BOTH non-empty content and tool_calls. The compactor
# should treat the (content + tool_calls) message and its tool
# response as one group.
msgs = [
_msg("system", "sys"),
_msg("user", "first task"),
_long("assistant", 4000),
_long("user", 4000),
_msg(
"assistant",
content = "thinking out loud while calling a tool",
tool_calls = [_tool_call("call_x")],
),
_msg("tool", content = "result for x", tool_call_id = "call_x"),
_msg("user", "thanks"),
]
out = SlidingWindowCompact(keep_recent = 2).compact(msgs, budget_tokens = 50)
asst_ids, tool_ids = _surviving_tool_ids(out)
assert asst_ids == tool_ids
class TestNoneContentAssistant:
def test_none_content_with_tool_calls(self):
# OpenAI permits assistant content=None when tool_calls is set.
msgs = [
_msg("system", "sys"),
_msg("user", "first task"),
_long("assistant", 4000),
_long("user", 4000),
{
"role": "assistant",
"content": None,
"tool_calls": [_tool_call("call_n")],
},
_msg("tool", content = "ok", tool_call_id = "call_n"),
_msg("user", "thanks"),
]
# Must not raise. estimate_tokens treats None content as 0.
out = SlidingWindowCompact(keep_recent = 2).compact(msgs, budget_tokens = 50)
asst_ids, tool_ids = _surviving_tool_ids(out)
assert asst_ids == tool_ids
class TestOutOfOrderToolMessage:
def test_tool_message_before_its_assistant_does_not_crash(self):
# Malformed input: tool message arrives before any assistant.
# _pair_linked_indices walks in order; the tool message has no
# pending id to match, so it ends up unpaired. The compactor
# must not crash. Use distinguishable contents so we can verify
# output order directly (identical _long() messages would
# collide under list.index()).
msgs = [
_msg("system", "sys"),
_msg("user", "first task"),
_msg("tool", content = "orphan-tool", tool_call_id = "call_z"),
_msg("assistant", content = "a-" + "x" * 5000),
_msg("user", "u-" + "x" * 5000),
_msg("assistant", content = "b-" + "x" * 5000),
_msg("user", "v-" + "x" * 5000),
]
# Identity-track each output back to the source index.
out = SlidingWindowCompact(keep_recent = 2).compact(msgs, budget_tokens = 30)
idxs = [msgs.index(m) for m in out]
assert idxs == sorted(idxs)
class TestOrphanToolMessageInRecentWindow:
def test_orphan_tool_message_alone_does_not_break(self):
# An orphan tool message (no preceding assistant tool_calls)
# ends up in its own group. If it lands in the recent window
# the compactor will keep it. That is malformed input on the
# caller's side, not a compactor bug — but it must not crash.
msgs = [
_msg("system", "sys"),
_msg("user", "first task"),
_long("assistant", 4000),
_long("user", 4000),
_msg("tool", content = "lonely", tool_call_id = "ghost"),
_msg("user", "thanks"),
]
out = SlidingWindowCompact(keep_recent = 2).compact(msgs, budget_tokens = 30)
# Compactor returns something and does not raise.
assert isinstance(out, list)
class TestDuplicateToolCallIds:
def test_duplicate_ids_across_assistants(self):
# Two assistants reuse the same tool_call_id. Only the second
# gets the tool match (pending_ids[tid] = i overwrites). The
# earlier assistant's tool_call has no matching tool. The
# compactor must still produce a coherent result when the
# second assistant + its tool fall inside the recent window.
msgs = [
_msg("system", "sys"),
_msg("user", "first task"),
_msg("assistant", content = "", tool_calls = [_tool_call("dup")]),
# No tool message for the first call (malformed input).
_long("assistant", 4000),
_long("user", 4000),
_msg("assistant", content = "", tool_calls = [_tool_call("dup")]),
_msg("tool", content = "second", tool_call_id = "dup"),
_msg("user", "thanks"),
]
out = SlidingWindowCompact(keep_recent = 2).compact(msgs, budget_tokens = 30)
# The recent-window assistant+tool group must stay linked.
kept_asst_tool = [
m for m in out if m.get("role") == "assistant" and m.get("tool_calls")
]
kept_tools = [m for m in out if m.get("role") == "tool"]
if kept_tools:
tool_ids_in_out = {m["tool_call_id"] for m in kept_tools}
asst_ids_in_out = set()
for m in kept_asst_tool:
for tc in m.get("tool_calls") or []:
if isinstance(tc.get("id"), str):
asst_ids_in_out.add(tc["id"])
assert tool_ids_in_out <= asst_ids_in_out
class TestPairMapInternals:
def test_pair_map_empty_for_plain_chat(self):
msgs = [
_msg("system", "sys"),
_msg("user", "hi"),
_msg("assistant", "hello"),
]
pm = _pair_linked_indices(msgs)
# Assistant has no tool_calls so the set entry is empty.
assert pm == {2: set()}
def test_pair_map_tool_without_string_id_ignored(self):
# tool_call_id is not a string — should be ignored, not crash.
msgs = [
_msg("assistant", content = "", tool_calls = [_tool_call("ok")]),
{"role": "tool", "content": "x", "tool_call_id": 42},
{"role": "tool", "content": "y", "tool_call_id": None},
]
pm = _pair_linked_indices(msgs)
# Only the string-id tool would match, and we did not include
# one — so the assistant entry is empty.
assert pm.get(0) == set()
class TestNonStringContent:
def test_int_or_dict_content_does_not_crash(self):
# Defensive: content that is neither str nor list (a stray int
# or dict from a misbehaving caller) should be ignored by
# estimate_tokens and treated as non-multimodal by compact().
msgs = [
_msg("system", "sys"),
_msg("user", "first task"),
{"role": "assistant", "content": 12345},
{"role": "user", "content": {"unexpected": "shape"}},
_long("assistant", 4000),
_long("user", 4000),
]
out = SlidingWindowCompact(keep_recent = 2).compact(msgs, budget_tokens = 30)
assert isinstance(out, list)
class TestNoSystemNoFirstUser:
def test_assistant_only_history(self):
# No system, no user. The first-user anchor loop simply does
# not match anything. Compactor must not crash.
msgs = [
_long("assistant", 4000),
_long("assistant", 4000),
_long("assistant", 4000),
_long("assistant", 4000),
]
out = SlidingWindowCompact(keep_recent = 1).compact(msgs, budget_tokens = 30)
# At least one message survives (the last assistant in the
# recent window).
assert len(out) >= 1
def test_first_user_not_at_index_one(self):
# First user is not adjacent to a system message. The
# first-user anchor must still find it by iteration order.
msgs = [
_msg("assistant", content = "stray pre-amble"),
_msg("user", "the real first task"),
_long("assistant", 4000),
_long("user", 4000),
_long("assistant", 4000),
]
out = SlidingWindowCompact(keep_recent = 1).compact(msgs, budget_tokens = 30)
# The first user must be present.
assert any(
m.get("role") == "user" and "real first task" in m.get("content", "")
for m in out
)
class TestMultimodalAsAnchor:
def test_multimodal_first_user(self):
# The very first user message is itself multimodal. Both the
# first-user anchor and the multimodal anchor refer to the
# same index — set semantics make the overlap a no-op.
msgs = [
_msg("system", "sys"),
{
"role": "user",
"content": [
{"type": "text", "text": "describe"},
{"type": "image_url", "image_url": {"url": "x"}},
],
},
_long("assistant", 4000),
_long("user", 4000),
_long("assistant", 4000),
]
out = SlidingWindowCompact(keep_recent = 1).compact(msgs, budget_tokens = 20)
# The multimodal first user must still be present.
assert any(isinstance(m.get("content"), list) for m in out)
class TestPairCleanupConsistency:
def test_drop_assistant_drops_its_tools(self):
# The assistant tool-call message is forced out of the recent
# window by keep_recent and is not an anchor. Its tool message
# must be dropped with it.
msgs = [
_msg("system", "sys"),
_msg("user", "first task"),
_msg("assistant", content = "", tool_calls = [_tool_call("call_a")]),
_msg("tool", content = "old result", tool_call_id = "call_a"),
_long("assistant", 5000),
_long("user", 5000),
_long("assistant", 5000),
_long("user", 5000),
]
out = SlidingWindowCompact(keep_recent = 1).compact(msgs, budget_tokens = 20)
asst_ids, tool_ids = _surviving_tool_ids(out)
# No orphans either way.
assert asst_ids == tool_ids
# And the old pair is gone.
assert "call_a" not in tool_ids
def test_drop_tools_drops_assistant(self):
# Hand-craft a scenario where the tool-side gets removed by
# the threshold loop but the assistant would be kept. The
# pair-cleanup pass must then drop the assistant too.
# We exploit the fact that the threshold-driven loop runs in
# `droppable` order (system-anchored + first-user-anchored
# excluded). The assistant tool-call carries little text;
# the tool response carries a lot. Dropping the tool may
# already get us under threshold, leaving the assistant.
# The cleanup pass must then drop the orphaned assistant.
msgs = [
_msg("system", "sys"),
_msg("user", "first task" + "y" * 100),
_msg("assistant", content = "", tool_calls = [_tool_call("call_b")]),
_msg("tool", content = "x" * 4000, tool_call_id = "call_b"),
_long("assistant", 50),
_long("user", 50),
]
out = SlidingWindowCompact(keep_recent = 1).compact(msgs, budget_tokens = 100)
asst_ids, tool_ids = _surviving_tool_ids(out)
# Both halves of the pair are either kept or dropped together.
assert asst_ids == tool_ids
class TestImmutability:
def test_input_messages_never_mutated_under_pressure(self):
msgs = [
_msg("system", "sys"),
_msg("user", "first task"),
_msg("assistant", content = "", tool_calls = [_tool_call("call_m")]),
_msg("tool", content = "ok", tool_call_id = "call_m"),
_long("assistant", 5000),
_long("user", 5000),
]
# Deep-ish snapshot: messages themselves + their tool_calls list
# are what the compactor could conceivably mutate.
snap = [dict(m) for m in msgs]
snap_tcs = [list(m.get("tool_calls") or []) for m in msgs]
_ = SlidingWindowCompact(keep_recent = 1).compact(msgs, budget_tokens = 30)
assert msgs == snap
for m, original_tcs in zip(msgs, snap_tcs):
assert (m.get("tool_calls") or []) == original_tcs
class TestStrategyRegistryAliases:
def test_get_strategy_empty_name_falls_back(self):
# The fallback path matters: a misconfigured request should
# degrade to no-op rather than raise.
assert isinstance(get_strategy(""), NoCompact)
class TestEstimateTokensWithBadToolCalls:
def test_tool_call_without_function_dict(self):
# Defensive: tool_calls list where an entry lacks "function".
# estimate_tokens does `(tc.get("function") or {}).get(...)`,
# which is safe — must not crash.
msgs = [
{
"role": "assistant",
"content": "",
"tool_calls": [{"id": "x", "type": "function"}],
}
]
assert estimate_tokens(msgs) == 0
class TestMultiToolCallSingleAssistant:
def test_assistant_with_multiple_tool_calls_grouped(self):
# One assistant message carries two tool_calls; both tool
# responses must be grouped with it so the trio stays linked.
msgs = [
_msg("system", "sys"),
_msg("user", "first task"),
_long("assistant", 4000),
_long("user", 4000),
_msg(
"assistant",
content = "",
tool_calls = [_tool_call("a"), _tool_call("b")],
),
_msg("tool", content = "ra", tool_call_id = "a"),
_msg("tool", content = "rb", tool_call_id = "b"),
_msg("user", "thanks"),
]
out = SlidingWindowCompact(keep_recent = 2).compact(msgs, budget_tokens = 50)
asst_ids, tool_ids = _surviving_tool_ids(out)
# All-or-nothing: the multi-call assistant brings both tool
# responses along, or none of the three survive.
assert asst_ids == tool_ids
class TestInterleavedUserBetweenAsstAndTool:
def test_user_between_asst_and_tool_does_not_break_grouping(self):
# Malformed (user message between asst tool_call and its tool
# response). pair_map still walks in order so the tool matches
# the most recent assistant carrying its id. The compactor
# must not raise; downstream rendering of an "orphan" user
# is acceptable since it was already orphaned in the input.
msgs = [
_msg("system", "sys"),
_msg("user", "first task"),
_msg("assistant", content = "", tool_calls = [_tool_call("g")]),
_msg("user", "intermediate"),
_msg("tool", content = "result", tool_call_id = "g"),
_long("assistant", 5000),
_long("user", 5000),
]
out = SlidingWindowCompact(keep_recent = 1).compact(msgs, budget_tokens = 30)
asst_ids, tool_ids = _surviving_tool_ids(out)
assert asst_ids == tool_ids
class TestNoCompactWithEmpty:
def test_nocompact_empty_messages_returns_empty(self):
out = NoCompact().compact([], budget_tokens = 100)
assert out == []