mirror of
https://github.com/andrewyng/openworker.git
synced 2026-11-04 21:37:21 +00:00
Audited the checks added in this PR by building the defects they are meant to stop and confirming each one goes red. The ones that did not are closed here; the changes are test-only. The attachment assertions were one-sided. `est < 5_000` certifies an estimate that collapsed to 26 exactly as readily as the correct ~1,500, and that collapse is the inverse defect: a history that really is over the cap, held permanently under it. Both attachment checks are now bracketed on both sides, and a differential check asserts the estimate is *independent* of data-URL length across the range attachments.py actually admits (12 M / 15 M), rather than sampling one size and inferring the property. Also newly covered: text and metadata inside a multi-part message (only image_url and file parts get the flat charge; everything else must still be measured), the `or` -> `is not None` boundary fix, which changed real behaviour with nothing pinning it, and the `assistants[-1]` fallback. The linearity claim now has a clock-free check that counts per-message measurements (<= 3n; the quadratic version ran ~150x that), so the complexity guarantee no longer rests on a wall-clock bound evaluated on a shared CI runner. The timing test stays as a companion — it catches a regression that keeps the call count linear but makes each call expensive. The module now carries a written statement of what it does not cover — the absolute accuracy of the 1,500 / 3,000 constants, which are bracketed to an order of magnitude but not anchored to a published provider formula. Co-Authored-By: Claude <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01Wevp5RAtiXBunmUuXfA6hk
621 lines
25 KiB
Python
621 lines
25 KiB
Python
"""OPE-27 — auto-compaction pure functions: trigger math, boundary picking, mechanical
|
|
extraction, summarizer seam, trim fallback, outbound view. No engine involved."""
|
|
|
|
import json
|
|
import time
|
|
|
|
import pytest
|
|
|
|
from coworker import compaction
|
|
from coworker import compaction as C
|
|
from coworker.compaction import (
|
|
CompactionState,
|
|
DEFAULT_CAP_TOKENS,
|
|
DEFAULT_CONTEXT_WINDOW,
|
|
apply_to_outbound,
|
|
build_state,
|
|
compacted_block,
|
|
estimate_tokens,
|
|
extract_user_messages,
|
|
extract_working_state,
|
|
is_context_overflow,
|
|
pick_boundary,
|
|
should_compact,
|
|
summarize_span,
|
|
summarizer_messages,
|
|
trigger_tokens,
|
|
trim_state,
|
|
)
|
|
|
|
|
|
# -- message builders ---------------------------------------------------------
|
|
|
|
|
|
def user(text):
|
|
return {"role": "user", "content": text, "ts": 1.0}
|
|
|
|
|
|
_call_seq = 0
|
|
|
|
|
|
def assistant(text="", tool_calls=None):
|
|
global _call_seq
|
|
msg = {"role": "assistant", "content": text, "ts": 1.0}
|
|
if tool_calls:
|
|
calls = []
|
|
for name, args in tool_calls:
|
|
calls.append(
|
|
{
|
|
"id": f"c{_call_seq}",
|
|
"type": "function",
|
|
"function": {"name": name, "arguments": json.dumps(args)},
|
|
}
|
|
)
|
|
_call_seq += 1
|
|
msg["tool_calls"] = calls
|
|
return msg
|
|
|
|
|
|
def tool(call_id, content):
|
|
return {
|
|
"role": "tool",
|
|
"tool_call_id": call_id,
|
|
"content": content if isinstance(content, str) else json.dumps(content),
|
|
"ts": 1.0,
|
|
}
|
|
|
|
|
|
def tool_turn(name, args, result):
|
|
"""[assistant tool-call, matching tool result] with a properly paired call id."""
|
|
a = assistant(tool_calls=[(name, args)])
|
|
return [a, tool(a["tool_calls"][0]["id"], result)]
|
|
|
|
|
|
def convo(turns=6, bulk=2000):
|
|
"""system + N user/assistant turns with bulky assistant text."""
|
|
msgs = [{"role": "system", "content": "You are a coworker."}]
|
|
for i in range(turns):
|
|
msgs.append(user(f"request {i}"))
|
|
msgs.append(assistant(f"answer {i} " + "x" * bulk))
|
|
return msgs
|
|
|
|
|
|
class FakeSummarizer:
|
|
# The default must satisfy the OPE-189 quality gate (all eight sections, past the
|
|
# minimum length) — an unusable summary is now refused, not stored. GOOD_SUMMARY is
|
|
# defined further down the module and resolved when a fake is constructed.
|
|
def __init__(self, text=None, fail_times=0):
|
|
self.text = GOOD_SUMMARY if text is None else text
|
|
self.fail_times = fail_times
|
|
self.calls = []
|
|
|
|
def complete(self, *, model, messages, tools=None, **settings):
|
|
self.calls.append({"model": model, "messages": messages, "tools": tools, **settings})
|
|
if self.fail_times > 0:
|
|
self.fail_times -= 1
|
|
raise RuntimeError("summarizer down")
|
|
|
|
class Turn:
|
|
pass
|
|
|
|
t = Turn()
|
|
t.text = self.text
|
|
return t
|
|
|
|
|
|
# -- trigger math -------------------------------------------------------------
|
|
|
|
|
|
def test_trigger_is_min_of_pct_and_cap():
|
|
assert trigger_tokens(100_000) == 80_000
|
|
assert trigger_tokens(1_000_000) == DEFAULT_CAP_TOKENS # the 250k cap wins
|
|
assert trigger_tokens(None) == int(0.8 * DEFAULT_CONTEXT_WINDOW)
|
|
# both knobs are user-overridable
|
|
assert trigger_tokens(100_000, threshold_pct=0.5, cap_tokens=40_000) == 40_000
|
|
assert trigger_tokens(100_000, threshold_pct=0.5, cap_tokens=999_999) == 50_000
|
|
|
|
|
|
def test_should_compact_crosses_threshold():
|
|
assert not should_compact(79_999, 100_000)
|
|
assert should_compact(80_000, 100_000)
|
|
|
|
|
|
def test_estimate_tokens_is_chars_over_four():
|
|
msgs = [user("a" * 400)]
|
|
est = estimate_tokens(msgs)
|
|
assert 100 <= est <= 120 # 400 chars of content + json overhead, /4
|
|
|
|
|
|
# -- boundary -----------------------------------------------------------------
|
|
|
|
|
|
def test_boundary_prefers_earliest_user_turn_that_fits():
|
|
msgs = convo(turns=6)
|
|
per_turn = estimate_tokens(msgs[1:3])
|
|
boundary = pick_boundary(msgs, keep_tokens=per_turn * 2 + 10)
|
|
assert msgs[boundary]["role"] == "user"
|
|
assert msgs[boundary]["content"] == "request 4" # newest two turns survive
|
|
|
|
|
|
def test_boundary_falls_inside_a_giant_final_turn():
|
|
# One user turn followed by a huge tool loop: the turn alone exceeds the budget,
|
|
# so the cut lands on an assistant (iteration) boundary inside it — never a tool row.
|
|
msgs = [{"role": "system", "content": "s"}, user("go")]
|
|
for i in range(8):
|
|
a = assistant("step " + "y" * 3000, tool_calls=[("run_shell", {"command": f"cmd{i}"})])
|
|
msgs += [a, tool(a["tool_calls"][0]["id"], {"exit_code": 0, "out": "z" * 3000})]
|
|
boundary = pick_boundary(msgs, keep_tokens=estimate_tokens(msgs[-3:]))
|
|
assert msgs[boundary]["role"] == "assistant"
|
|
|
|
|
|
def test_boundary_none_when_nothing_to_summarize():
|
|
msgs = [{"role": "system", "content": "s"}, user("hi"), assistant("hello")]
|
|
assert pick_boundary(msgs, keep_tokens=10_000_000) is None
|
|
|
|
|
|
# -- mechanical extraction ----------------------------------------------------
|
|
|
|
|
|
def test_working_state_files_commands_tools():
|
|
span = [
|
|
user("write it"),
|
|
*tool_turn("write_file", {"path": "a.py", "content": "x"}, {"ok": True}),
|
|
*tool_turn("run_shell", {"command": "pytest -q"}, {"exit_code": 1}),
|
|
*tool_turn("write_file", {"path": "b.py", "content": "y"}, {"ok": True}),
|
|
*tool_turn("write_file", {"path": "a.py", "content": "x2"}, {"ok": True}),
|
|
]
|
|
block = extract_working_state(span)
|
|
# deduped, most recent first
|
|
assert block.index("- a.py") < block.index("- b.py")
|
|
assert block.count("a.py") == 1
|
|
assert "pytest -q" in block and "[exit 1]" in block
|
|
assert "run_shell" in block and "write_file" in block
|
|
|
|
|
|
def test_working_state_empty_span():
|
|
assert extract_working_state([user("hi"), assistant("yo")]) == ""
|
|
|
|
|
|
def test_user_messages_extracted_verbatim_and_uncut():
|
|
# OPE-189: extraction no longer clips — a 1,400-char task statement reached the block
|
|
# cut mid-sentence, losing the output path and the field names the grader checks.
|
|
span = [
|
|
user("first ask"),
|
|
assistant("a"),
|
|
user([{"type": "text", "text": "second"}, {"type": "image_url", "image_url": {}}]),
|
|
assistant("b"),
|
|
user("bulk " + "z" * 2000),
|
|
]
|
|
out = extract_user_messages(span)
|
|
assert out[0] == "first ask"
|
|
assert out[1] == "second [image]"
|
|
assert out[2] == "bulk " + "z" * 2000 # verbatim; fitting is fit_user_messages' job
|
|
|
|
|
|
# -- summarizer seam ----------------------------------------------------------
|
|
|
|
|
|
def test_summarizer_messages_clip_tool_results_and_fold_prior():
|
|
span = [user("go"), *tool_turn("read_file", {"path": "big.txt"}, "huge " * 500)]
|
|
msgs = summarizer_messages(span, prior_summary="OLD SUMMARY")
|
|
body = msgs[1]["content"]
|
|
assert "OLD SUMMARY" in body
|
|
# The 2500-char tool result got clipped hard; the tail instruction is the only other bulk.
|
|
assert len(body) < 3000 + len(C.SUMMARY_TAIL_INSTRUCTION)
|
|
assert msgs[0]["role"] == "system" and "Primary request and intent" in msgs[0]["content"]
|
|
|
|
|
|
def test_summarizer_instruction_comes_after_the_transcript():
|
|
"""OPE-189: recency decides. With the instruction only up front, the last thing the
|
|
model read was the coworker mid-task and it continued that voice instead of
|
|
summarizing. The transcript is delimited and the instruction restated last."""
|
|
span = [user("go"), *tool_turn("run_shell", {"command": "ls"}, "a.txt b.txt")]
|
|
body = summarizer_messages(span)[1]["content"]
|
|
assert body.startswith("**--- BEGIN TRANSCRIPT TO SUMMARIZE ---**")
|
|
assert body.rstrip().endswith(C.SUMMARY_TAIL_INSTRUCTION)
|
|
assert body.index("[tool result]") < body.index("**--- END OF TRANSCRIPT ---**")
|
|
assert "## 1. Primary request and intent" in C.SUMMARY_TAIL_INSTRUCTION
|
|
|
|
|
|
GOOD_SUMMARY = "\n".join(
|
|
f"## {i}. {name}\ndetail line with enough substance to clear the minimum length bar"
|
|
for i, name in enumerate(
|
|
[
|
|
"Primary request and intent",
|
|
"Key concepts and decisions",
|
|
"Artifacts and files",
|
|
"Errors and fixes",
|
|
"All user messages",
|
|
"Pending tasks",
|
|
"Current work",
|
|
"Next step",
|
|
],
|
|
start=1,
|
|
)
|
|
)
|
|
|
|
|
|
def test_summarize_span_passes_model_and_budget():
|
|
fake = FakeSummarizer(text=GOOD_SUMMARY)
|
|
out = summarize_span(fake, "prov:model-x", [user("hi")], max_tokens=9999)
|
|
assert out == GOOD_SUMMARY
|
|
assert fake.calls[0]["model"] == "prov:model-x"
|
|
assert fake.calls[0]["tools"] is None
|
|
assert fake.calls[0]["max_tokens"] == 9999
|
|
|
|
|
|
def test_summarize_span_refuses_an_unusable_summary():
|
|
"""OPE-189: 'not empty' accepted a 61-char next-action sentence as the coworker's whole
|
|
memory of 75 turns. Empty, too short, and role-slipped replies are all refused so the
|
|
caller's retry (then trim) policy runs."""
|
|
for bad in (
|
|
" ",
|
|
"The output got truncated. Let me re-run the remaining checks.",
|
|
"## Artifacts and files\n" + "detail " * 200, # long, but not the anchor section
|
|
):
|
|
with pytest.raises(RuntimeError):
|
|
summarize_span(FakeSummarizer(text=bad), "m", [user("hi")])
|
|
|
|
|
|
def test_summary_quality_problem_names_the_fault():
|
|
assert C.summary_quality_problem(GOOD_SUMMARY) is None
|
|
assert C.summary_quality_problem("") == "empty"
|
|
assert "too short" in C.summary_quality_problem("## 1. Primary request and intent")
|
|
slipped = "Let me re-run the remaining checks. " * 20
|
|
assert "does not open the sections" in C.summary_quality_problem(slipped)
|
|
|
|
|
|
# -- build + repeated compaction ----------------------------------------------
|
|
|
|
|
|
def test_build_state_and_outbound_view():
|
|
msgs = convo(turns=6)
|
|
fake = FakeSummarizer(text=GOOD_SUMMARY + "\nthe gist")
|
|
state = build_state(
|
|
msgs, provider=fake, model="m", keep_tokens=estimate_tokens(msgs[-4:]) + 10
|
|
)
|
|
assert state is not None and not state.trimmed
|
|
assert state.user_messages[0] == "request 0"
|
|
|
|
out = apply_to_outbound(msgs, state)
|
|
assert out[0]["role"] == "system" # instructions survive
|
|
assert "<compacted-history>" in out[1]["content"]
|
|
assert "the gist" in out[1]["content"]
|
|
assert "request 0" in out[1]["content"] # mechanical user-message list
|
|
assert out[2] is msgs[state.boundary_index] # verbatim tail, canonical untouched
|
|
assert len(msgs) == 13 # canonical history unchanged
|
|
|
|
|
|
def test_repeated_compaction_summarizes_prior_plus_new_turns():
|
|
msgs = convo(turns=4)
|
|
fake = FakeSummarizer()
|
|
first = build_state(msgs, provider=fake, model="m", keep_tokens=estimate_tokens(msgs[-4:]) + 10)
|
|
# session grows
|
|
for i in range(4, 8):
|
|
msgs.append(user(f"request {i}"))
|
|
msgs.append(assistant(f"answer {i} " + "x" * 2000))
|
|
second = build_state(
|
|
msgs, provider=fake, model="m",
|
|
keep_tokens=estimate_tokens(msgs[-4:]) + 10, prior=first,
|
|
)
|
|
assert second is not None and second.boundary_index > first.boundary_index
|
|
# the second summarizer call folds the prior summary in
|
|
assert "previous compaction summary" in fake.calls[1]["messages"][1]["content"]
|
|
# user messages accumulate across compactions
|
|
assert "request 0" in second.user_messages[0]
|
|
assert any("request 5" in u for u in second.user_messages)
|
|
|
|
|
|
def test_build_state_none_when_boundary_stale():
|
|
msgs = convo(turns=3)
|
|
fake = FakeSummarizer()
|
|
state = build_state(msgs, provider=fake, model="m", keep_tokens=estimate_tokens(msgs[-2:]) + 10)
|
|
again = build_state(
|
|
msgs, provider=fake, model="m",
|
|
keep_tokens=10_000_000, prior=state,
|
|
)
|
|
assert again is None # nothing new fits below the prior boundary
|
|
|
|
|
|
# -- trim fallback ------------------------------------------------------------
|
|
|
|
|
|
def test_trim_advances_boundary_and_keeps_user_messages():
|
|
msgs = convo(turns=10)
|
|
state = trim_state(msgs)
|
|
assert state is not None and state.trimmed
|
|
assert msgs[state.boundary_index]["role"] in ("user", "assistant")
|
|
assert state.user_messages # preserved mechanically even without a summary
|
|
assert "trimmed" in state.summary_text
|
|
out = apply_to_outbound(msgs, state)
|
|
assert len(out) < len(msgs) + 1
|
|
|
|
|
|
def test_trim_from_prior_state_never_lands_on_tool_row():
|
|
msgs = [{"role": "system", "content": "s"}, user("go")]
|
|
for i in range(10):
|
|
msgs += tool_turn("run_shell", {"command": f"c{i}"}, {"exit_code": 0})
|
|
prior = trim_state(msgs)
|
|
later = trim_state(msgs, prior=prior)
|
|
assert later.boundary_index > prior.boundary_index
|
|
assert msgs[later.boundary_index]["role"] != "tool"
|
|
|
|
|
|
def test_trim_none_when_too_small():
|
|
assert trim_state([user("hi"), assistant("yo")]) is None
|
|
|
|
|
|
# -- state round-trip + overflow detection ------------------------------------
|
|
|
|
|
|
def test_state_dict_round_trip():
|
|
state = CompactionState(
|
|
boundary_index=7, summary_text="s", working_state="w",
|
|
user_messages=["u1"], created_at=1.5, model_used="m", trimmed=True,
|
|
)
|
|
assert CompactionState.from_dict(state.as_dict()) == state
|
|
assert CompactionState.from_dict(None) is None
|
|
assert CompactionState.from_dict({}) is None
|
|
|
|
|
|
def test_apply_to_outbound_noop_on_stale_or_missing_state():
|
|
msgs = convo(turns=2)
|
|
assert apply_to_outbound(msgs, None) is msgs
|
|
stale = CompactionState(boundary_index=999, summary_text="s", working_state="")
|
|
assert apply_to_outbound(msgs, stale) is msgs
|
|
|
|
|
|
def test_is_context_overflow():
|
|
assert is_context_overflow(Exception("Error 400: maximum context length is 128000 tokens"))
|
|
assert is_context_overflow(Exception("context_length_exceeded"))
|
|
assert is_context_overflow(Exception("Prompt is too long: 210000 tokens > limit"))
|
|
assert not is_context_overflow(Exception("rate limit exceeded"))
|
|
assert not is_context_overflow(Exception("connection reset"))
|
|
|
|
|
|
def test_user_messages_pin_the_first_and_budget_the_rest():
|
|
"""OPE-189: the list must not grow forever, but the OLD rule (newest 40) silently threw
|
|
away the opening message — the one that says why the session exists. Now the first is
|
|
pinned and the newest fill a token budget; the gap is in the middle."""
|
|
msgs = [{"role": "system", "content": "s"}]
|
|
msgs.append({"role": "user", "content": "THE ORIGINAL TASK: " + "spec " * 100})
|
|
msgs.append({"role": "assistant", "content": "starting"})
|
|
for i in range(300):
|
|
msgs.append({"role": "user", "content": f"ask {i} " + "filler " * 40})
|
|
msgs.append({"role": "assistant", "content": f"answer {i}"})
|
|
|
|
state = None
|
|
while True:
|
|
nxt = trim_state(msgs, prior=state, fraction=0.4, user_budget_tokens=2_000)
|
|
if nxt is None:
|
|
break
|
|
state = nxt
|
|
|
|
assert state is not None
|
|
assert state.user_messages[0].startswith("THE ORIGINAL TASK:") # pinned, never dropped
|
|
# The newest preserved message is from the end of the run (the very last turns stay in
|
|
# the verbatim tail rather than the block, so this is the newest that was compacted).
|
|
assert state.user_messages[-1].startswith("ask 29")
|
|
assert state.user_messages_dropped > 0 # the middle is what gives
|
|
block = compacted_block(state)
|
|
assert "older user messages omitted here" in block
|
|
assert block.index("THE ORIGINAL TASK") < block.index("omitted here")
|
|
|
|
restored = CompactionState.from_dict(state.as_dict())
|
|
assert restored.user_messages == state.user_messages
|
|
assert restored.user_messages_dropped == state.user_messages_dropped
|
|
|
|
|
|
def test_fit_user_messages_pins_first_clips_giants_and_counts_drops():
|
|
budget = 100 # tokens, i.e. ~400 chars
|
|
msgs = ["FIRST " + "a" * 200, "old " + "b" * 400, "mid", "newest"]
|
|
kept, dropped = C.fit_user_messages(msgs, prior_dropped=3, budget_tokens=budget)
|
|
assert kept[0] == msgs[0] # pinned whole: well under the pin cap
|
|
assert kept[-1] == "newest"
|
|
assert dropped >= 3 # running total carries forward
|
|
assert "mid" in kept # cheap messages still fit
|
|
|
|
# A first message far over the pin cap is clipped, never dropped.
|
|
huge = "X" * 100_000
|
|
kept, _ = C.fit_user_messages([huge, "later"], prior_dropped=0, budget_tokens=budget)
|
|
assert kept[0].startswith("XXX") and kept[0].endswith("…")
|
|
assert len(kept[0]) <= C._USER_PIN_MAX_TOKENS * 4
|
|
|
|
# A single newest message over budget is clipped to fit, not dropped.
|
|
kept, _ = C.fit_user_messages(["first", huge], prior_dropped=0, budget_tokens=budget)
|
|
assert len(kept) == 2 and kept[1].endswith("…")
|
|
|
|
|
|
def test_user_message_budget_scales_with_the_trigger():
|
|
assert C.user_message_budget(250_000) == 20_000 # at the default trigger
|
|
assert C.user_message_budget(60_000) == 4_800 # a lowered trigger keeps its saving
|
|
assert C.user_message_budget(1_000) == C._USER_BUDGET_MIN # floor
|
|
assert C.user_message_budget(10_000_000) == C._USER_BUDGET_MAX # ceiling
|
|
|
|
|
|
# -- attachments --------------------------------------------------------------
|
|
|
|
|
|
def _image_turn(chars: int = 1_500_000, text: str = "what's in this screenshot?") -> dict:
|
|
return {
|
|
"role": "user",
|
|
"content": [
|
|
{"type": "text", "text": text},
|
|
{
|
|
"type": "image_url",
|
|
"image_url": {"url": "data:image/png;base64," + "A" * chars},
|
|
},
|
|
],
|
|
}
|
|
|
|
|
|
def _pdf_turn(chars: int = 2_000_000) -> dict:
|
|
return {
|
|
"role": "user",
|
|
"content": [
|
|
{"type": "text", "text": "summarize this"},
|
|
{
|
|
"type": "file",
|
|
"file": {
|
|
"filename": "report.pdf",
|
|
"file_data": "data:application/pdf;base64," + "B" * chars,
|
|
},
|
|
},
|
|
],
|
|
}
|
|
|
|
|
|
def test_image_part_is_not_charged_its_base64_length():
|
|
"""A data-URL image must not dominate the signal: chars/4 read a 1.5 MB screenshot as
|
|
~375k tokens, so a single attachment sat permanently above cap_tokens and re-triggered
|
|
compaction on every iteration of the turn loop.
|
|
|
|
Bracketed on both sides. The upper bound is the defect being fixed; the lower bound
|
|
guards the inverse defect, an attachment charged ~nothing, which would keep a history
|
|
that really is over the cap permanently under it. A one-sided `< 5_000` certifies an
|
|
estimate that collapsed to 26 just as readily as the correct ~1.5k.
|
|
"""
|
|
msgs = [{"role": "system", "content": "s"}, _image_turn()]
|
|
est = estimate_tokens(msgs)
|
|
assert 1_000 < est < 5_000
|
|
assert not should_compact(est, 200_000)
|
|
|
|
|
|
def test_pdf_part_is_not_charged_its_base64_length():
|
|
msgs = [{"role": "system", "content": "s"}, _pdf_turn()]
|
|
est = estimate_tokens(msgs)
|
|
assert 2_000 < est < 10_000 # bracketed both sides; see the image test
|
|
|
|
|
|
def test_attachment_estimate_does_not_track_data_url_length():
|
|
"""The property the fix establishes is *independence* from base64 length, so assert it
|
|
differentially rather than at one sampled size. The range spans what attachments.py
|
|
actually admits (MAX_IMAGE_CHARS 12 M, MAX_PDF_CHARS 15 M) — a bound checked only at
|
|
1.5 MB says nothing about the sizes the GUI's own picker allows through.
|
|
"""
|
|
images = {estimate_tokens([_image_turn(n)]) for n in (100_000, 1_500_000, 12_000_000)}
|
|
assert len(images) == 1, f"image estimate varied with data-URL length: {sorted(images)}"
|
|
pdfs = {estimate_tokens([_pdf_turn(n)]) for n in (100_000, 2_000_000, 15_000_000)}
|
|
assert len(pdfs) == 1, f"file estimate varied with data-URL length: {sorted(pdfs)}"
|
|
|
|
|
|
def test_text_alongside_an_attachment_is_still_counted():
|
|
"""Only `image_url` and `file` parts get the flat charge; every other part falls to the
|
|
`else` branch and must still be measured. A long pasted prompt next to a screenshot is
|
|
the ordinary case, and charging it zero is the same under-count in a different place.
|
|
"""
|
|
small = estimate_tokens([_image_turn(text="hi")])
|
|
large = estimate_tokens([_image_turn(text="x" * 40_000)])
|
|
assert large - small >= 9_000 # 40k chars / 4, less serialization slack
|
|
|
|
|
|
def test_message_keys_outside_content_are_still_counted():
|
|
"""`_message_chars` measures the non-content keys separately once it takes the list
|
|
branch. `tool_call_id`, `name` and friends are real tokens on the wire."""
|
|
base = {"role": "user", "content": [{"type": "text", "text": "hi"}]}
|
|
tagged = {**base, "name": "z" * 4_000}
|
|
assert estimate_tokens([tagged]) - estimate_tokens([base]) >= 900
|
|
|
|
|
|
def test_text_only_estimate_is_unchanged_by_part_handling():
|
|
"""Text-only histories must measure exactly as before — the whole trigger threshold is
|
|
calibrated against that number."""
|
|
msgs = convo(turns=6)
|
|
assert estimate_tokens(msgs) == sum(len(json.dumps(m, default=str)) for m in msgs) // 4
|
|
|
|
|
|
def _long_history(turns: int = 300) -> list[dict]:
|
|
msgs = [{"role": "system", "content": "s"}]
|
|
for i in range(turns):
|
|
a = assistant("done " + "z" * 300, tool_calls=[("read_file", {"path": "a.py"})])
|
|
msgs += [
|
|
user(f"task {i} " + "x" * 400),
|
|
a,
|
|
tool(a["tool_calls"][0]["id"], {"exit_code": 0, "out": "y" * 4000}),
|
|
]
|
|
return msgs
|
|
|
|
|
|
def test_boundary_search_measures_each_message_a_bounded_number_of_times(monkeypatch):
|
|
"""The complexity claim, asserted without a clock.
|
|
|
|
Measuring each candidate with `estimate_tokens(messages[i:])` re-serialized the whole
|
|
tail per candidate — ~1 s of blocking CPU on a 1 200-message session, on the turn loop.
|
|
Counting the per-message measurements separates O(n) from O(n²) exactly, on any
|
|
machine and under any CI load, where a wall-clock bound only does so probabilistically.
|
|
"""
|
|
msgs = _long_history()
|
|
calls = 0
|
|
real = compaction._message_chars
|
|
|
|
def counting(msg):
|
|
nonlocal calls
|
|
calls += 1
|
|
return real(msg)
|
|
|
|
monkeypatch.setattr(compaction, "_message_chars", counting)
|
|
boundary = pick_boundary(msgs, keep_tokens=40_000)
|
|
assert boundary is not None
|
|
assert calls <= 3 * len(msgs) # one backward pass; quadratic ran ~150x this
|
|
|
|
|
|
def test_boundary_search_is_linear_in_history_length():
|
|
"""Wall-clock companion to the call-count test above: catches a regression that keeps
|
|
the call count linear but makes each call expensive."""
|
|
msgs = _long_history()
|
|
started = time.perf_counter()
|
|
boundary = pick_boundary(msgs, keep_tokens=40_000)
|
|
elapsed = time.perf_counter() - started
|
|
assert boundary is not None
|
|
assert estimate_tokens(msgs[boundary:]) <= 40_000
|
|
assert elapsed < 0.25 # linear lands ~5 ms here; the quadratic version took ~700 ms
|
|
|
|
|
|
def test_falls_back_to_the_last_assistant_when_no_candidate_fits():
|
|
"""The `assistants[-1]` fallback, on the one path that reaches it: a history with no
|
|
user turns at all (the `if users:` block short-circuits every other route there).
|
|
|
|
Pins which end of the list the fallback takes. `assistants[0]` would cut at the very
|
|
front and summarize the whole history; the intent is the opposite — keep as little as
|
|
possible when even the newest turn will not fit.
|
|
"""
|
|
msgs = [{"role": "system", "content": "s"}]
|
|
for i in range(6):
|
|
a = assistant("step " + "y" * 3000, tool_calls=[("run_shell", {"command": f"c{i}"})])
|
|
msgs += [a, tool(a["tool_calls"][0]["id"], {"exit_code": 0, "out": "z" * 3000})]
|
|
_, assistants = compaction._turn_starts(msgs, start=1)
|
|
assert len(assistants) > 2 # otherwise [-1] and [-2] coincide and the check is inert
|
|
boundary = pick_boundary(msgs, keep_tokens=10) # nothing fits in 10 tokens
|
|
assert boundary == assistants[-1]
|
|
assert msgs[boundary]["role"] == "assistant"
|
|
|
|
|
|
def test_history_that_wholly_fits_is_not_compacted():
|
|
"""`_fit` returning a legitimate index 0 means the entire history fits inside
|
|
keep_tokens, so there is nothing to summarize. The `or` this replaced coerced that 0 to
|
|
the `assistants[-1]` fallback and compacted a history of a few dozen tokens.
|
|
"""
|
|
msgs = [{"role": "assistant", "content": f"a{i}"} for i in range(3)]
|
|
assert estimate_tokens(msgs) < 100
|
|
assert pick_boundary(msgs, keep_tokens=10_000) is None
|
|
|
|
|
|
# -- what these checks do NOT cover -------------------------------------------
|
|
#
|
|
# * **The absolute accuracy of the flat estimates.** 1_500 / 3_000 are bracketed to an
|
|
# order of magnitude, not anchored to a provider's published image-token formula. A
|
|
# systematically wrong-but-plausible constant (say 600, or 4_000) passes everything here.
|
|
# Anchoring to the documented formula for the vision models actually in use would turn
|
|
# this from a range check into a real oracle.
|
|
# * **Multiple attachments in one message.** Every case above carries a single image or
|
|
# file; N-attachment accumulation is unexercised.
|
|
# * **Providers that report real usage.** estimate_tokens is the fallback signal only, so
|
|
# none of this constrains behaviour on providers that meter properly.
|
|
# * **Boundary quality, as opposed to boundary cost.** The complexity checks assert the
|
|
# search is linear and that the kept tail fits; they do not assert the boundary chosen is
|
|
# the *best* one, and the earliest-first preference is covered only by the pre-existing
|
|
# test_boundary_prefers_earliest_user_turn_that_fits.
|
|
# * **The exact-threshold boundary.** Nothing pins `suffix[i] // 4 <= keep_tokens` at
|
|
# equality, so an off-by-one to `<` survives; every case here sits clear of the threshold.
|
|
# Known from a mutation pass rather than assumed.
|