openswarm/backend/tests/test_distill_history.py

"""Distilled-history summary invariant.

On a rebuild the recap hard-drops everything before the cutoff, losing the thread of a
long chat. distilled_history_summary replaces that void with a cached aux-LLM summary of
the dropped span. These pin: it summarizes the dropped span, caches against the cutoff id,
recomputes when the cutoff advances, and fails open (no provider / kill switch / aux error
-> "", so the caller keeps today's hard-drop).
"""

import asyncio

import backend.apps.agents.manager.session.distill_history as dh
from backend.apps.agents.core.models import AgentSession, Message
from backend.apps.settings.settings import load_settings


def p_session(n: int) -> AgentSession:
    s = AgentSession(name="t", model="sonnet")
    s.messages = [Message(role="user", content=f"turn {i}") for i in range(n)]
    return s


def p_stub_distiller(monkeypatch, calls: list) -> None:
    async def fake(session, settings, body):
        calls.append(body)
        return f"SUMMARY[{len(body)} chars]"
    monkeypatch.setattr(dh, "p_call_distiller", fake)


def test_no_cutoff_returns_empty(monkeypatch) -> None:
    calls: list = []
    p_stub_distiller(monkeypatch, calls)
    s = p_session(8)
    out = asyncio.run(dh.distilled_history_summary(s, load_settings()))
    assert out == ""
    assert calls == []


def test_summarizes_dropped_span_and_caches(monkeypatch) -> None:
    calls: list = []
    p_stub_distiller(monkeypatch, calls)
    s = p_session(8)
    s.compacted_through_msg_id = s.messages[3].id  # drop turns 0..3
    out = asyncio.run(dh.distilled_history_summary(s, load_settings()))
    assert out.startswith("SUMMARY[")
    assert s.compacted_summary == out
    assert s.compacted_summary_through == s.messages[3].id
    assert "turn 0" in calls[0] and "turn 3" in calls[0]
    assert "turn 4" not in calls[0]  # surviving turns aren't distilled
    # Second call at the same cutoff reuses the cache, no new aux call.
    again = asyncio.run(dh.distilled_history_summary(s, load_settings()))
    assert again == out
    assert len(calls) == 1


def test_recomputes_when_cutoff_advances(monkeypatch) -> None:
    calls: list = []
    p_stub_distiller(monkeypatch, calls)
    s = p_session(10)
    s.compacted_through_msg_id = s.messages[3].id
    asyncio.run(dh.distilled_history_summary(s, load_settings()))
    s.compacted_through_msg_id = s.messages[6].id  # cutoff moved forward
    asyncio.run(dh.distilled_history_summary(s, load_settings()))
    assert len(calls) == 2
    assert "turn 6" in calls[1]


def test_fail_open_on_aux_error(monkeypatch) -> None:
    async def boom(session, settings, body):
        raise RuntimeError("provider down")
    monkeypatch.setattr(dh, "p_call_distiller", boom)
    s = p_session(8)
    s.compacted_through_msg_id = s.messages[3].id
    out = asyncio.run(dh.distilled_history_summary(s, load_settings()))
    assert out == ""
    assert s.compacted_summary is None


def test_stale_cache_not_served_when_cutoff_left_the_branch(monkeypatch) -> None:
    calls: list = []
    p_stub_distiller(monkeypatch, calls)
    s = p_session(8)
    s.compacted_through_msg_id = s.messages[3].id
    asyncio.run(dh.distilled_history_summary(s, load_settings()))  # caches
    assert s.compacted_summary is not None
    # Simulate a branch edit that dropped the cutoff message from the active branch.
    s.messages = [m for m in s.messages if m.id != s.messages[3].id]
    out = asyncio.run(dh.distilled_history_summary(s, load_settings()))
    assert out == ""  # membership check fires before the cache, so the stale summary is not served


def test_kill_switch_disables(monkeypatch) -> None:
    calls: list = []
    p_stub_distiller(monkeypatch, calls)
    monkeypatch.setattr(dh, "DISTILL_ENABLED", False)
    s = p_session(8)
    s.compacted_through_msg_id = s.messages[3].id
    out = asyncio.run(dh.distilled_history_summary(s, load_settings()))
    assert out == ""
    assert calls == []


def p_edit_session(paths: list) -> AgentSession:
    s = AgentSession(name="t", model="sonnet")
    s.messages = [Message(role="user", content="refactor them")]
    for p in paths:
        s.messages.append(Message(role="tool_call", content={"tool": "Edit", "input": {"file_path": p}}))
        s.messages.append(Message(role="tool_result", content={"text": f"edited {p}", "tool_name": "Edit"}))
    s.messages.append(Message(role="assistant", content="done"))
    return s


def test_touched_paths_come_from_tool_inputs_not_prose() -> None:
    s = p_edit_session(["/a/one.py", "/b/two.py", "/a/one.py"])
    s.messages.append(Message(role="tool_call", content={"tool": "NotebookEdit", "input": {"notebook_path": "/c/n.ipynb"}}))
    assert dh.touched_file_paths(s.messages) == ["/a/one.py", "/b/two.py", "/c/n.ipynb"]


def test_paths_the_summary_paraphrased_away_are_pinned_back(monkeypatch) -> None:
    """Live-measured: the aux model collapsed 12 literal paths into '/repo/src/module_N/handler_N.py'
    and 0 survived verbatim. Prose may paraphrase; the paths get re-attached deterministically."""
    paths = [f"/repo/src/module_{i}/handler_{i}.py" for i in range(12)]

    async def paraphrase(session, settings, body):
        return "The agent edited the twelve handler files under /repo/src."
    monkeypatch.setattr(dh, "p_call_distiller", paraphrase)
    s = p_edit_session(paths)
    s.compacted_through_msg_id = s.messages[-1].id
    out = asyncio.run(dh.distilled_history_summary(s, load_settings()))
    for p in paths:
        assert p in out


def test_paths_already_in_the_summary_are_not_repeated(monkeypatch) -> None:
    async def verbatim(session, settings, body):
        return "Edited /a/one.py and /b/two.py."
    monkeypatch.setattr(dh, "p_call_distiller", verbatim)
    s = p_edit_session(["/a/one.py", "/b/two.py"])
    s.compacted_through_msg_id = s.messages[-1].id
    out = asyncio.run(dh.distilled_history_summary(s, load_settings()))
    assert out == "Edited /a/one.py and /b/two.py."


def test_pinned_path_list_is_capped(monkeypatch) -> None:
    paths = [f"/repo/f{i}.py" for i in range(dh.MAX_PINNED_PATHS + 25)]

    async def paraphrase(session, settings, body):
        return "Edited a lot of files."
    monkeypatch.setattr(dh, "p_call_distiller", paraphrase)
    s = p_edit_session(paths)
    s.compacted_through_msg_id = s.messages[-1].id
    out = asyncio.run(dh.distilled_history_summary(s, load_settings()))
    assert sum(1 for p in paths if p in out) == dh.MAX_PINNED_PATHS