simbarag/tests/unit/test_context_trim.py

"""Regression tests for complete-turn context budgeting."""

import pytest
from langchain_core.messages import AIMessage, HumanMessage, ToolMessage

from blueprints.conversation import context
from blueprints.conversation.context import (
    IMAGE_PLACEHOLDER,
    build_context,
    message_tokens,
    trim_history,
)


def call(*ids):
    return AIMessage(
        content="",
        tool_calls=[
            {"id": i, "name": "search", "args": {"path": "notes/budget.md"}}
            for i in ids
        ],
    )


def result(i, content="found"):
    return ToolMessage(content=content, tool_call_id=i)


def cost(messages):
    return sum(message_tokens(m) for m in messages)


def test_long_current_turn_keeps_question_and_all_exchanges():
    question = HumanMessage(content="Compare my spending with my plan")
    messages = [question]
    for i in range(8):
        messages += [call(str(i)), result(str(i), "data " * 1000)]
    out = trim_history(messages, token_budget=3000)
    assert out[0] is question
    assert len(out) == 17
    assert cost(out) <= 3000
    assert [m.tool_call_id for m in out if isinstance(m, ToolMessage)] == [
        str(i) for i in range(8)
    ]
    assert messages[2].content == "data " * 1000
    assert out[1].tool_calls == messages[1].tool_calls


def test_parallel_tool_results_remain_paired():
    messages = [
        HumanMessage(content="q"),
        call("a", "b"),
        result("a", "x" * 20000),
        result("b", "y" * 20000),
    ]
    out = trim_history(messages, token_budget=700)
    assert len(out) == 4
    assert out[1].tool_calls == messages[1].tool_calls
    assert cost(out) <= 700


def test_drops_whole_turns_and_includes_bounded_excerpts():
    messages = []
    for i in range(12):
        messages += [
            HumanMessage(content=f"Question {i}"),
            call(str(i)),
            result(str(i), "data " * 1000),
            AIMessage(content=f"Decision {i}"),
        ]
    question = HumanMessage(content="What did we decide?")
    messages.append(question)
    out = trim_history(messages, token_budget=3000)
    assert isinstance(out[0], AIMessage)
    assert "Decision" in out[0].content
    assert "historical data" in out[0].content
    assert cost(out) <= 3000
    assert out[-1] is question
    # Every retained result still has its original call preceding it.
    seen = set()
    for m in out:
        if isinstance(m, AIMessage):
            seen.update(c["id"] for c in m.tool_calls)
        if isinstance(m, ToolMessage):
            assert m.tool_call_id in seen


def test_earlier_images_replaced_current_image_preserved():
    image = {"type": "image_url", "image_url": {"url": "data:image/png;base64,AAAA"}}
    messages = [
        HumanMessage(content=[{"type": "text", "text": "old"}, image]),
        AIMessage(content="photo"),
        HumanMessage(content=[{"type": "text", "text": "new"}, image]),
    ]
    out = trim_history(messages)
    assert out[0].content[1] == {"type": "text", "text": IMAGE_PLACEHOLDER}
    assert out[-1].content[1] == image
    assert messages[0].content[1] == image


def test_short_history_unchanged():
    messages = [HumanMessage(content="hi"), AIMessage(content="hello")]
    assert trim_history(messages) == messages


def test_question_that_cannot_fit_fails_without_truncation():
    messages = [HumanMessage(content="important request " * 1000)]
    with pytest.raises(ValueError, match="current question"):
        trim_history(messages, token_budget=100)
    assert messages[0].content == "important request " * 1000


def test_system_and_tool_schemas_are_charged_to_budget(monkeypatch):
    monkeypatch.setattr(context, "PROMPT_BUDGET_TOKENS", 2500)
    monkeypatch.setattr(context, "CONTEXT_WINDOW_TOKENS", 10000)
    schemas = [{"name": "search", "description": "tool docs " * 100}]
    messages = [HumanMessage(content="q"), call("a"), result("a", "result " * 3000)]
    out = build_context(messages, "system " * 100, schemas)
    assert cost(out) + context.estimate_tokens(schemas) <= 2500
    assert out[1] is messages[0]


def test_fixed_overhead_that_cannot_fit_fails(monkeypatch):
    monkeypatch.setattr(context, "PROMPT_BUDGET_TOKENS", 100)
    with pytest.raises(ValueError, match="tool definitions"):
        build_context([HumanMessage(content="q")], "system" * 1000, [])


def test_response_and_safety_reserves_reduce_prompt_limit(monkeypatch):
    monkeypatch.setattr(context, "CONTEXT_WINDOW_TOKENS", 8000)
    monkeypatch.setattr(context, "PROMPT_BUDGET_TOKENS", 16000)
    messages = [HumanMessage(content="q"), call("a"), result("a", "data " * 5000)]
    out = build_context(messages, "system", [])
    assert (
        cost(out) + context.estimate_tokens([])
        <= 8000 - context.RESPONSE_RESERVE_TOKENS - context.SAFETY_RESERVE_TOKENS
    )


def test_cancelled_turn_does_not_replay_incomplete_exchange():
    messages = [
        HumanMessage(content="old"),
        call("a", "b"),
        result("a"),
        HumanMessage(content="new"),
    ]
    out = trim_history(messages)
    assert out == [messages[0], messages[-1]]


def test_orphan_results_are_removed():
    messages = [
        result("orphan"),
        HumanMessage(content="q"),
        result("another orphan"),
        AIMessage(content="answer"),
    ]
    assert trim_history(messages) == [messages[1], messages[3]]


def test_large_assistant_answer_is_bounded_in_earlier_turn():
    messages = [
        HumanMessage(content="old"),
        AIMessage(content="long " * 10000),
        HumanMessage(content="new"),
    ]
    out = trim_history(messages, token_budget=4000)
    assert cost(out) <= 4000
    assert len(out[1].content) < len(messages[1].content)
    assert messages[1].content == "long " * 10000


def test_multiblock_attachment_shares_one_cap():
    messages = [
        HumanMessage(content=[{"type": "text", "text": "x" * 10000} for _ in range(5)]),
        AIMessage(content="ok"),
        HumanMessage(content="next"),
    ]
    out = trim_history(messages, token_budget=5000)
    assert cost(out) <= 5000
    assert sum(len(p["text"]) for p in out[0].content) < 5000


async def test_saved_memories_are_bounded_without_deleting(monkeypatch):
    from blueprints.conversation import memory

    saved = ["newest preference", "big " * 5000, "older preference"]

    async def get_memories(user_id):
        return saved

    monkeypatch.setattr(memory, "get_memories_for_user", get_memories)
    prompt = await memory.build_system_prompt_with_memories("user")
    block = prompt.split("USER MEMORIES", 1)[1]
    assert "newest preference" in block
    assert "older preference" in block
    assert "omitted" in block
    assert "big big" not in block
    assert saved[1] == "big " * 5000


def test_prompt_accounts_for_selected_tool_schemas(monkeypatch):
    monkeypatch.setenv("TAVILY_API_KEY", "test-key")
    from blueprints.conversation import agents

    extra = {
        "type": "function",
        "function": {
            "name": "extra",
            "description": "extra docs",
            "parameters": {"type": "object", "properties": {}},
        },
    }
    seen = {}

    def build(messages, system, schemas, **kwargs):
        seen["schemas"] = schemas
        return messages

    monkeypatch.setattr(agents, "build_context", build)
    messages = [HumanMessage(content="hello")]
    assert (
        agents._build_prompt({"messages": messages}, tool_schemas=[extra]) == messages
    )
    assert extra in seen["schemas"]


def test_heavy_current_turn_leaves_room_for_previous_decisions():
    question = HumanMessage(content="Apply that plan to my spending")
    messages = [
        HumanMessage(content="Plan?"),
        AIMessage(content="Keep dining under $300."),
        question,
        call("a"),
        result("a", "spending " * 10000),
    ]
    out = trim_history(messages, token_budget=2500)
    assert cost(out) <= 2500
    assert question in out
    assert any(
        "Keep dining under $300." in m.content
        for m in out
        if isinstance(m.content, str)
    )


def test_large_tool_arguments_fail_explicitly_instead_of_breaking_pair():
    exchange = call("a")
    exchange.tool_calls[0]["args"] = {"query": "huge " * 10000}
    messages = [HumanMessage(content="q"), exchange, result("a")]
    with pytest.raises(ValueError, match="tool calls exceed"):
        trim_history(messages, token_budget=300)
    assert exchange.tool_calls[0]["args"]["query"] == "huge " * 10000


def test_unicode_tool_result_fits_estimated_budget():
    messages = [HumanMessage(content="q"), call("a"), result("a", "猫😺预算" * 10000)]
    out = trim_history(messages, token_budget=600)
    assert cost(out) <= 600
    assert out[0] is messages[0]
    assert out[1].tool_calls == messages[1].tool_calls


def test_eight_tool_finance_scenario_preserves_structured_evidence():
    import copy
    import json

    question = HumanMessage(
        content="Compare September spending with my plan; exclude transfers."
    )
    messages = [question]
    kinds = [
        "ynab.summary",
        "ynab.transactions",
        "ynab.spending_report",
        "ynab.category_budget",
        "search.documents",
        "obsidian.note",
        "ynab.transactions",
        "ynab.spending_report",
    ]
    for i, kind in enumerate(kinds):
        data = {
            "start_date": "2026-09-01",
            "end_date": "2026-09-30",
            "filters": {"exclude_transfers": True, "budget_id": "budget-123"},
            "total": -1234.56,
            "count": 30,
        }
        if kind == "obsidian.note":
            data.update(
                path="notes/budget.md", markdown="Keep dining under $300.\n" * 2000
            )
        else:
            key = (
                "results"
                if kind == "search.documents"
                else "transactions"
                if kind == "ynab.transactions"
                else "rows"
            )
            data[key] = [
                {
                    "id": f"{i}-{n}",
                    "path": f"notes/{n}.md",
                    "date": "2026-09-15",
                    "amount": -41.152,
                    "snippet": "finance evidence " * 200,
                }
                for n in range(30)
            ]
        artifact = {
            "kind": kind,
            "data": data,
            "source_ref": {"system": "ynab", "id": "budget-123"},
        }
        messages += [
            call(str(i)),
            ToolMessage(
                content="formatted report " * 5000,
                artifact=artifact,
                tool_call_id=str(i),
            ),
        ]
    original = copy.deepcopy(messages)
    out = trim_history(messages, token_budget=10000)
    assert cost(out) <= 10000
    assert len(out) == 17
    assert out[0] is question
    assert messages == original
    for old, projected in zip(messages[2::2], out[2::2]):
        assert projected.artifact == old.artifact
        parsed = json.loads(projected.content)
        assert parsed["source_ref"] == old.artifact["source_ref"]
        for key in ("total", "start_date", "end_date", "filters", "count"):
            assert parsed["data"][key] == old.artifact["data"][key]
        if parsed["kind"] == "obsidian.note":
            assert parsed["data"]["path"] == "notes/budget.md"
            assert parsed["data"]["markdown_compaction"]["omitted_characters"] > 0
        else:
            data = parsed["data"]
            key = next(k for k in ("rows", "transactions", "results") if k in data)
            omitted = data[key + "_compaction"]
            assert omitted["retained_rows"] + omitted["omitted_rows"] == 30
            assert {r["id"] for r in data[key]} | {
                r["id"] for r in omitted["omitted_source_refs"]
            } == {f"{old.tool_call_id}-{n}" for n in range(30)}


def test_json_content_without_artifact_stays_parseable():
    import json

    payload = {
        "total": 10,
        "filters": {"query": "food"},
        "rows": [{"id": str(i), "memo": "x" * 2000} for i in range(20)],
    }
    source = result("a", json.dumps(payload))
    projected = context.compact_tool_result(source, 900)
    parsed = json.loads(projected.content)
    assert parsed["total"] == 10
    assert parsed["rows_compaction"]["omitted_rows"] > 0
    assert json.loads(source.content) == payload


def test_unknown_result_is_marked_excerpt():
    projected = context.compact_tool_result(result("a", "unknown prose " * 1000), 100)
    assert "tool result excerpt" in projected.content


def test_structural_metadata_floor_fails_instead_of_corrupting_json():
    message = ToolMessage(
        content="short", artifact={"source_ref": {"id": "x" * 2000}}, tool_call_id="a"
    )
    with pytest.raises(ValueError, match="required tool metadata"):
        trim_history([HumanMessage(content="q"), call("a"), message], token_budget=400)


def test_summary_role_and_budget_are_enforced():
    from langchain_core.messages import SystemMessage

    with pytest.raises(ValueError, match="assistant-role"):
        trim_history(
            [HumanMessage(content="q")],
            summary_message=SystemMessage(content="untrusted"),
        )
    question = HumanMessage(content="Current request " * 10)
    assert trim_history(
        [question], token_budget=100, summary_message=AIMessage(content="x" * 1000)
    ) == [question]
    with pytest.raises(ValueError, match="without tool calls"):
        trim_history([HumanMessage(content="q")], summary_message=call("a"))


def test_nested_source_refs_and_dates_survive_row_omission():
    import json

    payload = {
        "rows": [
            {
                "metadata": {
                    "source_ids": ["one", "two"],
                    "start_date": "2026-09-01",
                    "path": "notes/plan.md",
                },
                "snippet": "x" * 1000,
            }
        ]
    }
    message = ToolMessage(content="display", artifact=payload, tool_call_id="a")
    parsed = json.loads(context.compact_tool_result(message, 0).content)
    assert (
        parsed["rows_compaction"]["omitted_source_refs"][0]["metadata"]
        == payload["rows"][0]["metadata"]
    )


def test_repeated_json_only_compaction_preserves_all_source_ids_and_counts():
    import json

    original = {
        "rows": [{"id": f"source-{i}", "memo": "detail " * 1000} for i in range(30)],
        "total": 300,
    }
    message = result("a", json.dumps(original))
    first = context.compact_tool_result(message, 3000)
    second = context.compact_tool_result(first, 1500)
    for projected in (first, second):
        parsed = json.loads(projected.content)
        metadata = parsed["rows_compaction"]
        assert metadata["retained_rows"] == len(parsed["rows"])
        assert metadata["retained_rows"] + metadata["omitted_rows"] == 30
        ids = {r["id"] for r in parsed["rows"]} | {
            r["id"] for r in metadata["omitted_source_refs"]
        }
        assert ids == {f"source-{i}" for i in range(30)}
    out = trim_history(
        [HumanMessage(content="q"), call("a"), message], token_budget=1000
    )
    parsed = json.loads(out[-1].content)
    assert parsed["rows_compaction"]["retained_rows"] == len(parsed["rows"])
    ids = {r["id"] for r in parsed["rows"]} | {
        r["id"] for r in parsed["rows_compaction"]["omitted_source_refs"]
    }
    assert len(ids) == 30
    assert cost(out) <= 1000


def test_focused_note_heading_index_uses_absolute_content_lines():
    import json

    artifact = {
        "kind": "obsidian.note",
        "data": {"path": "late.md", "markdown": "full original"},
        "model_context": {
            "markdown": "# Requested section\n" + "detail\n" * 2000,
            "read_range": {"start_line": 1001, "end_line": 3001, "total_lines": 4000},
        },
    }
    msg = ToolMessage(content="focused read", artifact=artifact, tool_call_id="a")
    projected = context.compact_tool_result(msg, 1500)
    metadata = json.loads(projected.content)["data"]["markdown_compaction"]
    assert metadata["headings"] == [{"line": 1001, "heading": "# Requested section"}]
    assert msg.artifact["data"]["markdown"] == "full original"


def _conversation(turns, size=1000):
    messages = []
    for i in range(turns):
        messages += [
            HumanMessage(content=f"Question {i}"),
            call(str(i)),
            result(str(i), "data " * size),
            AIMessage(content=f"Decision {i}"),
        ]
    return messages


def test_summary_omitted_when_no_turns_dropped():
    messages = _conversation(2, size=10) + [HumanMessage(content="next")]
    summary = AIMessage(content="Earlier: keep dining low")
    assert trim_history(messages, summary_message=summary) == trim_history(messages)
    assert summary not in trim_history(messages, summary_message=summary)


def assert_valid_sequence(out):
    seen = set()
    for m in out:
        if isinstance(m, AIMessage):
            seen.update(c["id"] for c in m.tool_calls)
        if isinstance(m, ToolMessage):
            assert m.tool_call_id in seen


def test_summary_placed_before_current_turn_when_turns_dropped():
    question = HumanMessage(content="next")
    # The current turn has its own tool exchange the summary must not split.
    messages = _conversation(12) + [question, call("now"), result("now")]
    summary = AIMessage(content="Earlier: keep dining low")
    out = trim_history(messages, token_budget=6000, summary_message=summary)
    assert cost(out) <= 6000
    assert out[-4:] == [summary, question, messages[-2], messages[-1]]
    assert out.count(summary) == 1
    # Excerpts of dropped turns stay at the head; they only change with the cut.
    assert out[0].content.startswith("Earlier conversation excerpts")
    assert_valid_sequence(out)


def test_changing_summary_leaves_retained_prefix_identical():
    messages = _conversation(12) + [HumanMessage(content="next")]
    short = AIMessage(content="Earlier: keep dining low")
    long = AIMessage(content="Earlier: " + "constraint " * 400)
    a = trim_history(messages, token_budget=6000, summary_message=short)
    b = trim_history(messages, token_budget=6000, summary_message=long)
    assert a[: a.index(short)] == b[: b.index(long)]


def test_trim_drops_to_low_watermark_and_prefix_stays_stable():
    budget = 8000
    messages, oldest, costs = [], [], []
    for i in range(30):
        messages.append(HumanMessage(content=f"Question {i}"))
        out = trim_history(messages, token_budget=budget)
        assert out[-1] is messages[-1]
        assert cost(out) <= budget
        oldest.append(next(m for m in out if isinstance(m, HumanMessage)).content)
        costs.append(cost(out))
        messages += [call(str(i)), result(str(i), "data " * 300), AIMessage("ok")]
    cuts = [i for i in range(1, 30) if oldest[i] != oldest[i - 1]]
    assert cuts, "history should overflow the budget"
    # Each drop goes well below the budget instead of one turn under it...
    for i in cuts:
        assert costs[i] <= budget * (context.TRIM_LOW_WATERMARK + 0.1)
    # ...so the retained prefix is unchanged for several turns in between.
    assert all(b - a >= 3 for a, b in zip(cuts, cuts[1:]))
    assert len(cuts) <= 5