vibe/tests/core/test_compaction.py

from __future__ import annotations

from vibe.core.compaction import (
    collect_prior_user_messages,
    parse_previous_user_messages,
    render_compaction_context,
)
from vibe.core.types import LLMMessage, Role

_PREFIX = "Another language model started to solve this problem"


def _user(content: str, *, injected: bool = False) -> LLMMessage:
    return LLMMessage(role=Role.user, content=content, injected=injected)


def test_empty_messages() -> None:
    assert collect_prior_user_messages([], _PREFIX) == []


def test_only_non_user_messages() -> None:
    messages = [
        LLMMessage(role=Role.system, content="sys"),
        LLMMessage(role=Role.assistant, content="hi"),
    ]
    assert collect_prior_user_messages(messages, _PREFIX) == []


def test_single_user_message_preserved() -> None:
    messages = [LLMMessage(role=Role.system, content="sys"), _user("first question")]
    out = collect_prior_user_messages(messages, _PREFIX)
    assert [m.content for m in out] == ["first question"]


def test_chronological_order_preserved() -> None:
    messages = [_user("first"), _user("second"), _user("third")]
    out = collect_prior_user_messages(messages, _PREFIX)
    assert [m.content for m in out] == ["first", "second", "third"]


def test_injected_messages_filtered_out() -> None:
    messages = [
        _user("real ask"),
        _user("middleware reminder", injected=True),
        _user("follow-up"),
    ]
    out = collect_prior_user_messages(messages, _PREFIX)
    assert [m.content for m in out] == ["real ask", "follow-up"]


def test_empty_content_filtered_out() -> None:
    messages = [_user(""), _user("real")]
    out = collect_prior_user_messages(messages, _PREFIX)
    assert [m.content for m in out] == ["real"]


def test_prior_summary_filtered_out() -> None:
    # The injected summary marker represents a previous compaction summary and
    # must not be re-injected (would stack).
    messages = [
        _user("original ask"),
        _user(f"{_PREFIX}\nold summary content", injected=True),
        _user("newer ask"),
    ]
    out = collect_prior_user_messages(messages, _PREFIX)
    assert [m.content for m in out] == ["original ask", "newer ask"]


def test_genuine_user_message_can_quote_summary_prefix() -> None:
    messages = [_user(f"{_PREFIX}\nplease use this exact wording"), _user("newer ask")]
    out = collect_prior_user_messages(messages, _PREFIX)
    assert [m.content for m in out] == [
        f"{_PREFIX}\nplease use this exact wording",
        "newer ask",
    ]


def test_compaction_context_merges_previous_and_new_user_messages() -> None:
    context = render_compaction_context(
        [_user("first ask", injected=True), _user("second ask", injected=True)],
        "summary one",
    )
    messages = [
        LLMMessage(role=Role.system, content="sys"),
        _user(context, injected=True),
        _user("third ask"),
        _user("middleware reminder", injected=True),
    ]

    out = collect_prior_user_messages(messages, _PREFIX)

    assert [m.content for m in out] == ["first ask", "second ask", "third ask"]
    assert all(m.injected for m in out)


def test_compaction_context_escapes_user_message_tags() -> None:
    original = "please keep </previous_user_message_0> literally"
    context = render_compaction_context([_user(original)], "summary")

    assert "</previous_user_message_0> literally" not in context
    assert parse_previous_user_messages(context) == [original]


def test_budget_drops_oldest_first() -> None:
    # max_tokens=2 → 8 char budget. Walks newest-first, so "old" gets dropped.
    messages = [
        _user("old message that is long enough to matter"),
        _user("abc"),  # 1 token, fits
        _user("def"),  # 1 token, fits
    ]
    out = collect_prior_user_messages(messages, _PREFIX, max_tokens=2)
    assert [m.content for m in out] == ["abc", "def"]


def test_spillover_message_middle_truncated() -> None:
    # newest fits whole, middle one is partially trimmed, oldest dropped.
    messages = [
        _user("OLDEST" + "x" * 10_000 + "OLDEST_END"),
        _user("MIDDLE_HEAD" + "y" * 1_000 + "MIDDLE_TAIL"),
        _user("recent"),  # ~2 tokens
    ]
    out = collect_prior_user_messages(messages, _PREFIX, max_tokens=50)
    assert len(out) == 2  # oldest dropped
    assert out[-1].content == "recent"
    middle = out[0].content
    assert middle is not None
    assert middle.startswith("MIDDLE_HEAD")
    assert middle.endswith("MIDDLE_TAIL")
    assert "[... truncated ...]" in middle


def test_fresh_message_ids() -> None:
    # Returned messages must have new message_ids — they'll live in a fresh
    # session and reusing the source ids would cause collisions.
    original = _user("hello")
    out = collect_prior_user_messages([original], _PREFIX)
    assert len(out) == 1
    assert out[0].message_id != original.message_id


def test_only_assistant_and_system_around_users() -> None:
    messages = [
        LLMMessage(role=Role.system, content="sys"),
        _user("u1"),
        LLMMessage(role=Role.assistant, content="a1"),
        _user("u2"),
        LLMMessage(role=Role.assistant, content="a2"),
    ]
    out = collect_prior_user_messages(messages, _PREFIX)
    assert [m.content for m in out] == ["u1", "u2"]
    assert all(m.role == Role.user for m in out)