6krrt/tests/test_context_prune.py

"""Pure relevance-based context pruning.

Ported from the MIT-licensed llmrouter "pinch" module, reduced to the safe
invariants 6krrt requires: user/assistant/system messages are always kept
verbatim, only old tool results are trimmed, and message order / role pairing
survives so the result still parses as a conversation.
"""

from __future__ import annotations

from context_prune import (
    estimate_tokens,
    extract_text,
    order_by_relevance,
    prune_context,
    trim_candidates,
)
from dispatcher import estimate_prompt_tokens


def _user(text: str) -> dict:
    return {"role": "user", "content": text}


def _assistant(text: str) -> dict:
    return {"role": "assistant", "content": text}


def _tool(name: str, content: str) -> dict:
    return {"role": "tool", "name": name, "content": content}


def test_short_conversation_is_untouched():
    messages = [_user("hi"), _assistant("hello")]
    out, stats = prune_context(messages, budget_tokens=50000)
    assert out == messages
    assert stats["pruned"] is False


def test_user_and_assistant_messages_always_kept():
    messages = [
        _user("big user " * 5000),
        _assistant("big assistant " * 5000),
        _user("tail"),
    ]
    out, stats = prune_context(messages, budget_tokens=100)
    assert stats["pruned"] is True
    assert len(out) == len(messages)
    assert out[0]["role"] == "user"
    assert out[1]["role"] == "assistant"
    assert out[2]["role"] == "user"
    # All three survived verbatim; only their absolute content is untouched.
    assert out[0]["content"] == messages[0]["content"]
    assert out[2]["content"] == "tail"


def test_old_tool_results_beyond_window_are_trimmed():
    messages = [
        _tool("read", "huge tool output " * 2000),        # old, should be trimmed
        _assistant("step one"),
        _user("continue"),
        _assistant("ok"),
    ]
    out, stats = prune_context(messages, budget_tokens=100, keep_last_turns=1)
    assert stats["pruned"] is True
    assert stats["tokens_saved"] > 0
    # The old tool result got shortened (head + tail + elision marker).
    tool_out = out[0]
    assert tool_out["role"] == "tool"
    assert "chars trimmed" in tool_out["content"]
    assert out[1]["content"] == "step one"
    assert out[2]["content"] == "continue"
    assert out[3]["content"] == "ok"


def test_recent_turn_is_protected_from_pruning():
    # keep_last_turns=1 protects the last user turn AND its tool results.
    messages = [
        _user("first"),
        _tool("search", "old result " * 3000),
        _assistant("first answer"),
        _user("second"),
        _tool("search", "recent result " * 3000),
        _assistant("second answer"),
    ]
    out, stats = prune_context(messages, budget_tokens=100, keep_last_turns=1)
    assert stats["pruned"] is True
    # The recent tool result (after the 2nd user) survives verbatim.
    recent_tool = [m for m in out if m.get("role") == "tool"][-1]
    assert recent_tool["content"] == messages[4]["content"]
    # Old tool result got trimmed.
    old_tool = [m for m in out if m.get("role") == "tool"][0]
    assert "chars trimmed" in old_tool["content"]


def test_protected_max_chars_caps_outsized_recent_tool_result():
    # A protected (recent) tool result far larger than protected_max_chars
    # -- e.g. a full verbose test-suite run -- still gets elided, closing
    # the gap keep_last_turns otherwise leaves: no size limit inside the
    # protected window at all.
    messages = [
        _user("first"),
        _tool("search", "old result " * 3000),
        _assistant("first answer"),
        _user("second"),
        _tool("pytest", "PASSED " * 50000),  # huge, but inside the protected window
        _assistant("second answer"),
    ]
    out, stats = prune_context(
        messages, budget_tokens=100, keep_last_turns=1, protected_max_chars=2000,
    )
    recent_tool = [m for m in out if m.get("role") == "tool"][-1]
    assert "chars trimmed" in recent_tool["content"]
    assert len(recent_tool["content"]) < len(messages[4]["content"])
    assert stats["tokens_saved"] > 0


def test_protected_max_chars_none_preserves_historical_behavior():
    # The pure function's default (None) must reproduce the exact historical
    # behavior new tests weren't written against -- a huge recent tool result
    # ships verbatim when the cap isn't configured.
    messages = [
        _user("first"),
        _tool("search", "old result " * 3000),
        _assistant("first answer"),
        _user("second"),
        _tool("pytest", "PASSED " * 50000),
        _assistant("second answer"),
    ]
    out, _ = prune_context(messages, budget_tokens=100, keep_last_turns=1)
    recent_tool = [m for m in out if m.get("role") == "tool"][-1]
    assert recent_tool["content"] == messages[4]["content"]


def test_protected_max_chars_does_not_touch_small_recent_results():
    # A recent tool result under the cap is left alone even when the cap is
    # configured -- this only catches true outliers, never ordinary output.
    messages = [
        _user("first"),
        _tool("search", "old result " * 3000),
        _assistant("first answer"),
        _user("second"),
        _tool("search", "small recent result"),
        _assistant("second answer"),
    ]
    out, _ = prune_context(
        messages, budget_tokens=100, keep_last_turns=1, protected_max_chars=2000,
    )
    recent_tool = [m for m in out if m.get("role") == "tool"][-1]
    assert recent_tool["content"] == "small recent result"


def test_tool_result_role_pairing_preserved():
    # The tool message keeps its name and position, so the API still parses.
    messages = [
        _user("q"),
        _tool("search", "x" * 6000),
        _assistant("a"),
    ]
    out, stats = prune_context(messages, budget_tokens=50, keep_last_turns=0)
    tool_out = [m for m in out if m.get("role") == "tool"][0]
    assert tool_out["name"] == "search"
    assert out[0]["role"] == "user"
    assert out[1]["role"] == "tool"
    assert out[2]["role"] == "assistant"


def test_short_tool_result_is_dropped_to_placeholder():
    # A large user message pushes the conversation over budget, so the SHORT
    # old tool result (under max_summarize_chars but longer than the marker) is
    # dropped to a placeholder that is shorter than the result it replaces.
    result_text = "result " * 40           # 280 chars > placeholder, < 4000 cap
    messages = [
        _user("big user " * 5000),
        _tool("search", result_text),
        _assistant("a"),
    ]
    out, stats = prune_context(messages, budget_tokens=100, keep_last_turns=0)
    assert stats["pruned"] is True
    tool_out = [m for m in out if m.get("role") == "tool"][0]
    assert "[search: result omitted]" in tool_out["content"]
    assert len(tool_out["content"]) < len(result_text)


def test_tiny_tool_result_is_kept_when_placeholder_loses_tokens():
    # A result shorter than the placeholder marker is kept verbatim: replacing
    # it would ADD tokens, which is not the point of pruning.
    messages = [
        _user("big user " * 5000),
        _tool("search", "hi"),
        _assistant("a"),
    ]
    out, stats = prune_context(messages, budget_tokens=100, keep_last_turns=0)
    tool_out = [m for m in out if m.get("role") == "tool"][0]
    assert tool_out["content"] == "hi"


def test_stats_report_tokens_saved():
    messages = [_user("q"), _tool("search", "y" * 9000), _assistant("a")]
    out, stats = prune_context(messages, budget_tokens=10, keep_last_turns=0)
    assert stats["pruned"] is True
    assert stats["original_tokens"] > stats["final_tokens"]
    assert stats["tokens_saved"] > 0
    assert stats["summarized"] >= 1


def test_extract_text_handles_content_blocks():
    msg = {
        "role": "tool",
        "content": [
            {"type": "text", "text": "hello"},
            {"type": "tool_result", "content": "world"},
        ],
    }
    assert "hello" in extract_text(msg)
    assert "world" in extract_text(msg)


def test_estimate_tokens_is_crude_but_consistent():
    assert estimate_tokens("x" * 30) == 10
    assert estimate_tokens("") == 0
    assert estimate_tokens(None) == 0


def test_max_summarize_chars_below_3000_never_grows_or_goes_negative():
    # Regression: a max_summarize_chars well below 3000 (here 100) with a
    # moderately long tool result used to head/tail-trim into something LONGER
    # than the input and a negative tokens_saved. It must never grow or negate.
    content = "tool output " * 100                 # 1200 chars
    messages = [
        _user("big user " * 5000),
        _tool("read", content),
        _assistant("a"),
    ]
    out, stats = prune_context(
        messages,
        budget_tokens=100,
        keep_last_turns=0,
        max_summarize_chars=100,
    )
    tool_out = [m for m in out if m.get("role") == "tool"][0]
    assert len(tool_out["content"]) < len(content)   # strictly shorter
    assert stats["tokens_saved"] >= 0
    assert "[-" not in tool_out["content"]           # no negative trim marker


def test_extract_text_counts_image_url_blocks():
    # Regression: image_url blocks used to read as "", undercounting the
    # message and letting pruning miss it. They must contribute their url.
    msg = {
        "content": [
            {"type": "text", "text": "describe this"},
            {"type": "image_url", "image_url": {"url": "data:image/png;base64,AAAA"}},
        ],
    }
    text = extract_text(msg)
    assert text != ""
    assert "describe this" in text
    assert "AAAA" in text


def test_extract_text_flattens_nested_tool_result_content_list():
    # Regression: a tool_result whose content is itself a list of blocks was
    # str()'d as a repr instead of being flattened like a top-level list.
    msg = {
        "content": [
            {
                "type": "tool_result",
                "content": [
                    {"type": "text", "text": "nested text"},
                    {"type": "image_url", "image_url": {"url": "http://x/img.png"}},
                ],
            },
        ],
    }
    text = extract_text(msg)
    assert "nested text" in text
    assert "img.png" in text


def test_stats_use_summarized_and_length_is_always_preserved():
    # Regression: the stat was named "dropped" but nothing is ever removed —
    # content is replaced in place, so len(pruned) == len(messages) always.
    messages = [
        _user("q"),
        _tool("search", "x" * 200),
        _assistant("a"),
    ]
    out, stats = prune_context(
        messages,
        budget_tokens=50,
        keep_last_turns=0,
        max_summarize_chars=5000,
    )
    assert "dropped" not in stats
    assert "summarized" in stats
    assert stats["summarized"] >= 1
    assert len(out) == len(messages)


def test_zero_user_messages_protect_newest_tool_result():
    # Regression: with no user message the recency guard collapsed to
    # protected_from = len(messages), so even the newest tool result (the one
    # the next turn needs) was trimmed. It must be protected.
    messages = [
        {"role": "system", "content": "sys"},
        {"role": "assistant", "content": "call tool", "tool_calls": [{"id": "t1"}]},
        {"role": "tool", "name": "read", "content": "x" * 20000},
    ]
    out, stats = prune_context(messages, budget_tokens=100, keep_last_turns=4)
    assert stats["pruned"] is True
    tool_out = [m for m in out if m.get("role") == "tool"][0]
    assert tool_out["content"] == messages[2]["content"]


def test_structured_tool_result_keeps_image_block_when_trimmed():
    # Regression: trimming flattened structured content to a string, dropping
    # the image_url part permanently. The non-text block must survive — kept
    # as an image_url block (so the model knows an image was present), with
    # its payload shrunk to a stub rather than shipped full-size.
    messages = [
        _user("big user " * 5000),
        {
            "role": "tool",
            "name": "capture",
            "content": [
                {"type": "text", "text": "some result " * 300},  # > 3000 chars
                {"type": "image_url", "image_url": {"url": "data:image/png;base64,AAAA"}},
            ],
        },
        _assistant("a"),
    ]
    out, stats = prune_context(
        messages,
        budget_tokens=100,
        keep_last_turns=0,
        max_summarize_chars=100,
    )
    tool_out = [m for m in out if m.get("role") == "tool"][0]
    # The content is still a structured list, not a flattened string.
    assert isinstance(tool_out["content"], list)
    text_blocks = [
        b for b in tool_out["content"]
        if isinstance(b, dict) and b.get("type") == "text"
    ]
    assert text_blocks
    # Non-empty prose and non-trivial image means it was shrunk (not kept full).
    img_blocks = [
        b for b in tool_out["content"]
        if isinstance(b, dict) and b.get("type") == "image_url"
    ]
    assert len(img_blocks) == 1
    # The block survives (type preserved) but its payload is stubbed, so the
    # actual shipped bytes drop rather than staying full-size (review #new).
    assert img_blocks[0]["image_url"]["url"] == "[image omitted]"
    assert img_blocks[0]["image_url"]["url"] != "data:image/png;base64,AAAA"


def test_large_image_tool_result_actually_shrinks_and_stats_are_honest():
    # Regression (follow-up review): fixing image_url counting (old #5) and
    # preserving non-text blocks (old #11) composed into a case where the text
    # block was sliced from prose+base64 joined together and the image_url
    # block was copied through UNCHANGED — so the dominant size was never
    # reduced while tokens_saved reported a large "saving" that didn't happen,
    # and a raw base64 fragment leaked into the text field. The payload must
    # genuinely shrink, the image payload must be stubbed, base64 must not
    # leak into the text block, and tokens_saved must equal the real shrink.
    huge_b64 = "A" * 200000
    messages = [
        _user("describe this"),
        {"role": "assistant", "content": None, "tool_calls": [{"id": "1"}]},
        {
            "role": "tool",
            "name": "screenshot",
            "content": [
                {"type": "text", "text": "here is the screenshot"},
                {"type": "image_url", "image_url": {"url": f"data:image/png;base64,{huge_b64}"}},
            ],
        },
        _user("now what"),
    ]
    out, stats = prune_context(messages, budget_tokens=10, keep_last_turns=0)
    assert stats["pruned"] is True
    # The payload genuinely shrank — final below original.
    assert stats["final_tokens"] < stats["original_tokens"]
    # Savings equal the real reduction, not a per-message estimate.
    assert stats["tokens_saved"] == (
        stats["original_tokens"] - stats["final_tokens"]
    )
    tool_out = [m for m in out if m.get("role") == "tool"][0]
    for part in tool_out["content"]:
        if isinstance(part, dict) and part.get("type") == "image_url":
            assert "AAAA" not in part["image_url"]["url"]
        elif isinstance(part, dict) and part.get("type") == "text":
            assert "base64" not in part.get("text", "")


def test_pruned_list_measures_smaller_than_the_raw_conversation():
    """Regression for review #4: routing must measure the PRUNED size.

    chat_completions used to route on the raw conversation's token count and
    only shrink the upstream payload later, so a long tool-heavy session could
    be sent to a pricier large-context model based on its pre-pruned size even
    though pinch would bring the outgoing request well under budget. This test
    pins the invariant that enables the fix: after one prune, the router's own
    ``estimate_prompt_tokens`` sees a smaller count, so the window/tier/cost
    decision reflects what actually ships upstream.
    """
    messages = [
        {"role": "user", "content": "inspect the repo"},
        {"role": "assistant", "content": "reading files", "tool_calls": [{"id": "t1"}]},
        {"role": "tool", "name": "read", "content": "file1: " + "x" * 4000},
        {"role": "assistant", "content": "still going", "tool_calls": [{"id": "t2"}]},
        {"role": "tool", "name": "read", "content": "file2: " + "y" * 4000},
        {"role": "assistant", "content": "ok, fixing"},
        {"role": "user", "content": "fix it"},
    ]
    raw_tokens = estimate_prompt_tokens(messages)

    pruned, stats = prune_context(
        list(messages),
        budget_tokens=100,
        keep_last_turns=0,
        max_summarize_chars=4000,
    )
    assert stats["pruned"] is True

    pruned_tokens = estimate_prompt_tokens(pruned)
    # The outgoing payload measures strictly smaller than the raw conversation,
    # so routing on the pruned list picks the cheaper/smaller-window model.
    assert pruned_tokens < raw_tokens


def test_pinch_disabled_leaves_the_measured_size_unchanged():
    """When pinch is off the routed decision measures the full conversation.

    Mirrors the byte-for-byte no-op contract: with ``prune_context`` not run,
    ``estimate_prompt_tokens`` sees the raw messages exactly as before the
    reordering fix.
    """
    messages = [
        {"role": "user", "content": "q"},
        {"role": "assistant", "content": "a", "tool_calls": [{"id": "t1"}]},
        {"role": "tool", "name": "read", "content": "z" * 4000},
        {"role": "user", "content": "done"},
    ]
    # unpinned: identical to calling estimate_prompt_tokens(messages) directly
    assert estimate_prompt_tokens(messages) == estimate_prompt_tokens(list(messages))


def test_order_by_relevance_near_identical_sorts_last():
    # A candidate nearly identical to the query is the MOST relevant, so it
    # must sort LAST (compressed last).
    query = [1.0, 0.0, 0.0]
    candidates = [
        [-1.0, 0.0, 0.0],
        [0.0, 1.0, 0.0],
        [0.9999, 0.0, 0.0],
    ]
    order = order_by_relevance(query, candidates)
    assert order == [0, 1, 2] or order[-1] == 2


def test_order_by_relevance_opposite_sorts_first():
    query = [1.0, 0.0, 0.0]
    candidates = [
        [0.0, 1.0, 0.0],
        [-1.0, 0.0, 0.0],
        [1.0, 0.0, 0.0],
    ]
    order = order_by_relevance(query, candidates)
    # The least-relevant (opposite, dot<0) is index 1, which must come first.
    assert order[0] == 1


def test_order_by_relevance_ascending_cosine():
    query = [1.0, 1.0]
    # cosine to query: (q itself) 1.0, (0.5,0.5) ~1.0, (1,-1) 0.0, (-1,1) 0.0.
    candidates = [
        [1.0, -1.0],
        [0.5, 0.5],
        [0.0, 1.0],
    ]
    order = order_by_relevance(query, candidates)
    plain = [candidates[i] for i in order]
    assert plain[0] == [1.0, -1.0]


def test_order_by_relevance_empty_and_single():
    assert order_by_relevance([1.0, 0.0], []) == []
    assert order_by_relevance([1.0, 0.0], [[1.0, 0.0]]) == [0]


def test_order_by_relevance_zero_norm_candidate_is_least_relevant():
    # A zero-norm candidate (no signal) gets similarity 0.0, so it sorts early
    # (least relevant) without dividing by zero.
    query = [1.0, 0.0]
    candidates = [
        [0.0, 0.0],
        [1.0, 0.0],
    ]
    order = order_by_relevance(query, candidates)
    assert order == [0, 1]


def test_trim_candidates_normal_multi_turn():
    messages = [
        _user("q1"),
        _tool("read", "old"),
        _assistant("a1"),
        _user("q2"),
        _tool("read", "recent"),
        _assistant("a2"),
    ]
    candidates, protected_from, protected_user = trim_candidates(messages, keep_last_turns=1)
    # keep_last_turns=1 protects the last user turn (index 3) and everything
    # at/after it. Tool message 1 (index 1) is the only trim candidate.
    assert candidates == [1]
    assert protected_from == 3
    assert protected_user == 3


def test_trim_candidates_no_user_turn():
    messages = [
        {"role": "system", "content": "sys"},
        {"role": "assistant", "content": "call", "tool_calls": [{"id": "t1"}]},
        {"role": "tool", "name": "read", "content": "x" * 20000},
        {"role": "tool", "name": "read", "content": "y" * 20000},
    ]
    candidates, protected_from, protected_user = trim_candidates(messages, keep_last_turns=1)
    # No user turn: protect the trailing keep_last_turns tool results, so the
    # first tool (index 2) is a candidate, the last is protected.
    assert 2 in candidates
    assert len(messages) - 1 not in candidates
    assert protected_user is None


def test_trim_candidates_keep_last_turns_larger_than_turns():
    messages = [
        _user("q1"),
        _tool("read", "old"),
        _assistant("a1"),
    ]
    candidates, protected_from, _ = trim_candidates(messages, keep_last_turns=10)
    assert candidates == []
    assert protected_from == 0


def test_trim_candidates_zero_keep_last_turns():
    messages = [
        _user("q1"),
        _tool("read", "old"),
        _assistant("a1"),
    ]
    candidates, protected_from, _ = trim_candidates(messages, keep_last_turns=0)
    # Nothing protected: both the tool message and (by position) everything
    # before protected_from=len is a candidate zone; only tool candidates count.
    assert candidates == [1]
    assert protected_from == len(messages)


def test_relevance_order_overrides_recency():
    # Two old candidates; the OLDER one (index 1) is more relevant, so it is
    # ranked LAST in relevance_order and must survive verbatim, while the NEWER
    # (index 3) is less relevant, ranked FIRST, and gets compressed.
    messages = [
        _user("q0"),
        _tool("read", "old-but-relevant " * 4000),
        _assistant("a1"),
        _tool("read", "newer-but-irrelevant " * 4000),
        _assistant("a2"),
        _user("final"),
        _assistant("ok"),
    ]
    # trim_candidates with keep_last_turns=1 protects the "final" turn; both
    # tool messages (index 1 and 3) are candidates in clock order [1, 3].
    candidates, protected_from, _ = trim_candidates(messages, keep_last_turns=1)
    assert candidates == [1, 3]
    # Relevance order indexes into candidates: candidate 3 (newer) is LEAST
    # relevant so it is compressed first; candidate 1 (older) is most relevant.
    relevance_order = [1, 0]
    out, stats = prune_context(
        messages,
        budget_tokens=100,
        keep_last_turns=1,
        max_summarize_chars=4000,
        relevance_order=relevance_order,
    )
    assert stats["pruned"] is True
    # The most-relevant (older) tool result survived verbatim.
    assert out[1]["content"] == messages[1]["content"]
    # The least-relevant (newer) tool result was compressed.
    assert "chars trimmed" in out[3]["content"]


def test_extra_fixed_tokens_tips_budget_to_prune():
    # A conversation whose message tokens sit just under the budget, but tool
    # overhead pushes the accounting total over. Without overhead it is a no-op;
    # with overhead it prunes.
    messages = [
        _user("question"),
        _tool("read", "x" * 540),  # 540 chars -> 180 tokens
        _assistant("answer"),
        _user("follow up"),
    ]
    msg_tokens = sum(estimate_tokens(extract_text(m)) for m in messages)
    assert msg_tokens == 187

    out, stats = prune_context(
        messages,
        budget_tokens=200,
        keep_last_turns=0,
        extra_fixed_tokens=0,
    )
    assert stats["pruned"] is False
    assert stats["original_tokens"] == msg_tokens

    out, stats = prune_context(
        messages,
        budget_tokens=200,
        keep_last_turns=0,
        extra_fixed_tokens=100,
    )
    assert stats["pruned"] is True
    assert stats["original_tokens"] == msg_tokens + 100
    # The tool result was trimmed to cover the deficit (187 + 100 - 200 = 87).
    assert stats["tokens_saved"] >= 87
    tool_out = [m for m in out if m.get("role") == "tool"][0]
    assert "chars trimmed" in tool_out["content"] or "result omitted" in tool_out["content"]


def test_extra_fixed_tokens_zero_preserves_behavior():
    # The default and explicit zero must be byte-for-byte identical.
    messages = [
        _user("question"),
        _tool("read", "big result " * 500),
        _assistant("answer"),
        _user("follow up"),
        _assistant("ok"),
    ]
    default_out, default_stats = prune_context(
        messages, budget_tokens=100, keep_last_turns=0
    )
    explicit_zero_out, explicit_zero_stats = prune_context(
        messages,
        budget_tokens=100,
        keep_last_turns=0,
        extra_fixed_tokens=0,
    )
    assert explicit_zero_out == default_out
    assert explicit_zero_stats == default_stats


def test_tool_content_never_summarized():
    # Even when extra_fixed_tokens forces pruning, user/assistant/system
    # messages survive verbatim; only tool results are candidates.
    messages = [
        _user("what is the answer"),
        _assistant("i'll check"),
        _tool("search", "big tool result " * 500),
        _assistant("the answer is 42"),
        _user("thanks"),
    ]
    out, stats = prune_context(
        messages,
        budget_tokens=100,
        keep_last_turns=0,
        extra_fixed_tokens=1000,
    )
    assert stats["pruned"] is True
    assert len(out) == len(messages)
    for orig, pruned in zip(messages, out):
        if orig.get("role") in ("user", "assistant", "system"):
            assert pruned["content"] == orig["content"]
    tool_out = [m for m in out if m.get("role") == "tool"][0]
    assert "chars trimmed" in tool_out["content"] or "result omitted" in tool_out["content"]


def test_relevance_order_none_matches_default():
    # relevance_order=None must be byte-for-byte identical to not passing the
    # parameter at all (the historical uniform pass / fallback path).
    messages = [
        _user("q0"),
        _tool("read", "big result " * 5000),
        _assistant("a1"),
        _user("final"),
        _assistant("ok"),
    ]
    default_out, default_stats = prune_context(
        messages, budget_tokens=100, keep_last_turns=0
    )
    explicit_none_out, explicit_none_stats = prune_context(
        messages,
        budget_tokens=100,
        keep_last_turns=0,
        relevance_order=None,
    )
    assert explicit_none_out == default_out
    assert explicit_none_stats["tokens_saved"] == default_stats["tokens_saved"]


def test_relevance_order_indexes_align_with_trim_candidates():
    messages = [
        _user("q0"),
        _tool("search", "first-result " * 5000),
        _assistant("a1"),
        _tool("search", "second-result " * 5000),
        _assistant("a2"),
        _user("final"),
        _assistant("ok"),
    ]
    candidates, protected_from, _ = trim_candidates(messages, keep_last_turns=1)
    assert candidates == [1, 3]
    # relevance_order indexes into `candidates`: index 0 means message 1.
    out, stats = prune_context(
        messages,
        budget_tokens=2000,
        keep_last_turns=1,
        max_summarize_chars=4000,
        relevance_order=[0],
    )
    # Candidate 0 (message 1) was compressed because its content is huge.
    assert "chars trimmed" in out[1]["content"]
    # Candidate 1 (message 3) stayed verbatim (not in the order).
    assert out[3]["content"] == messages[3]["content"]