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, prune_context
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_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))