cleveragents-core/robot/helper_tui_conversation_pruning.py

"""Helper script for tui_conversation_pruning.robot integration tests.

Each subcommand is a self-contained check that prints a sentinel on success.
Tests ConversationStream pruning behaviour directly — no Textual event loop required.
"""

# ruff: noqa: E402
from __future__ import annotations

import sys
from pathlib import Path

_SRC = str(Path(__file__).resolve().parents[1] / "src")
if _SRC not in sys.path:
    sys.path.insert(0, _SRC)

from cleveragents.tui.conversation import (
    DEFAULT_PRUNE_EXCESS,
    DEFAULT_PRUNE_LOW_MARK,
    ConversationSettings,
    ConversationStream,
    load_conversation_settings,
)


def _make_stream(
    prune_low_mark: int = 100,
    prune_excess: int = 50,
    preserve_recent_lines: int = 10,
) -> ConversationStream:
    settings = ConversationSettings(
        prune_low_mark=prune_low_mark,
        prune_excess=prune_excess,
        preserve_recent_lines=preserve_recent_lines,
    )
    return ConversationStream(settings=settings)


def _multiline_block(n_lines: int, label: str = "line") -> str:
    return "\n".join(f"{label} {i}" for i in range(n_lines))


def _rendered_line_count(stream: ConversationStream) -> int:
    rendered = stream.render()
    if not rendered:
        return 0
    return rendered.count("\n") + 1


# ---------------------------------------------------------------------------
# Subcommands
# ---------------------------------------------------------------------------


def cmd_prune_trigger() -> None:
    """Verify pruning fires and total_lines returns to <= prune_low_mark."""
    stream = _make_stream(prune_low_mark=100, prune_excess=50)
    # trigger_line_count = 150; push enough rendered lines to trigger pruning twice,
    # ensuring total_lines ends at prune_low_mark or below.
    for i in range(24):
        stream.add_block(_multiline_block(10, f"msg{i}"), block_type="message")

    assert stream.total_lines <= 100, (
        f"Expected total_lines <= prune_low_mark (100), got {stream.total_lines}"
    )
    print("tui-prune-trigger-ok")


def cmd_separator_aware_prune() -> None:
    """Verify rendered inter-block separators count toward pruning thresholds."""
    stream = _make_stream(
        prune_low_mark=100,
        prune_excess=50,
        preserve_recent_lines=0,
    )

    for i in range(100):
        stream.add_block(f"one-line-{i}", block_type="message")

    rendered = stream.render()
    assert stream.total_lines == _rendered_line_count(stream), (
        "Expected total_lines to match rendered line count, "
        f"got total_lines={stream.total_lines}, rendered={_rendered_line_count(stream)}"
    )
    assert stream.total_lines <= 150, (
        "Expected rendered lines to stay within trigger threshold, "
        f"got {stream.total_lines}"
    )
    assert "one-line-0" not in rendered, "Expected oldest one-line block to be pruned"
    assert any(block.block_type == "note" for block in stream.blocks), (
        "Expected pruning note after separator-aware pruning"
    )
    print("tui-separator-aware-prune-ok")


def cmd_prune_note_inserted() -> None:
    """Verify a note block is at index 0 with the expected text after pruning."""
    stream = _make_stream(prune_low_mark=100, prune_excess=50)
    for i in range(21):
        stream.add_block(_multiline_block(10, f"msg{i}"), block_type="message")

    blocks = stream.blocks
    assert len(blocks) > 0, "Expected at least one block after pruning"
    first = blocks[0]
    assert first.block_type == "note", (
        f"Expected first block to be 'note', got {first.block_type!r}"
    )
    assert "(Earlier messages pruned" in first.text, (
        f"Expected pruned-note text in first block, got: {first.text!r}"
    )
    print("tui-prune-note-inserted-ok")


def cmd_prune_protected() -> None:
    """Verify protected blocks survive pruning."""
    stream = _make_stream(prune_low_mark=100, prune_excess=50)
    protected_text = "protected welcome block"
    stream.add_block(protected_text, block_type="welcome", protected=True)

    for i in range(20):
        stream.add_block(_multiline_block(10, f"msg{i}"), block_type="message")

    texts = [b.text for b in stream.blocks]
    assert protected_text in texts, (
        f"Protected block was removed during pruning. Block types: "
        f"{[b.block_type for b in stream.blocks]}"
    )
    print("tui-prune-protected-ok")


def cmd_clear_resets_state() -> None:
    """Verify clear() resets blocks, total_lines, and the note-active state."""
    stream = _make_stream(prune_low_mark=100, prune_excess=50)
    for i in range(20):
        stream.add_block(_multiline_block(10, f"msg{i}"), block_type="message")

    assert len(stream.blocks) > 0, "Expected blocks before clear"

    stream.clear()

    assert len(stream.blocks) == 0, (
        f"Expected 0 blocks after clear, got {len(stream.blocks)}"
    )
    assert stream.total_lines == 0, (
        f"Expected total_lines == 0 after clear, got {stream.total_lines}"
    )
    # After clear, a single small block should not trigger pruning
    stream.add_block("hello", block_type="message")
    assert len(stream.blocks) == 1, (
        f"Expected 1 block after clear + add, got {len(stream.blocks)}"
    )
    print("tui-clear-resets-state-ok")


def cmd_settings_defaults() -> None:
    """Verify load_conversation_settings returns defaults when config file is absent."""
    import tempfile

    absent_path = Path(tempfile.mkdtemp()) / "no-such-file.json"
    settings = load_conversation_settings(config_path=absent_path)

    assert settings.prune_low_mark == DEFAULT_PRUNE_LOW_MARK, (
        f"prune_low_mark: {settings.prune_low_mark} != {DEFAULT_PRUNE_LOW_MARK}"
    )
    assert settings.prune_excess == DEFAULT_PRUNE_EXCESS, (
        f"Expected prune_excess={DEFAULT_PRUNE_EXCESS}, got {settings.prune_excess}"
    )
    print("tui-settings-defaults-ok")


# ---------------------------------------------------------------------------
# Dispatch
# ---------------------------------------------------------------------------

COMMANDS = {
    "prune-trigger": cmd_prune_trigger,
    "separator-aware-prune": cmd_separator_aware_prune,
    "prune-note-inserted": cmd_prune_note_inserted,
    "prune-protected": cmd_prune_protected,
    "clear-resets-state": cmd_clear_resets_state,
    "settings-defaults": cmd_settings_defaults,
}

if __name__ == "__main__":
    if len(sys.argv) < 2 or sys.argv[1] not in COMMANDS:
        print(f"Usage: {sys.argv[0]} [{' | '.join(COMMANDS)}]", file=sys.stderr)
        sys.exit(1)
    COMMANDS[sys.argv[1]]()