tests/test_context_middleware.py

"""Tests for the active-path context-management middleware (langgraph-context plan).

Red/green TDD oracle for ``src/loosecanvas/context_middleware.py``. The middleware
implements the staged context policy from ``plan/langgraph-context``:

* Phase A — instrumentation: measure context growth per model call.
* Phase B — ``before_model`` compaction: keep recent dialogue verbatim, collapse
  stale tool chatter into compact structured receipts, drop superseded
  ``get_canvas_state`` payloads, enforce an optional token budget.

The load-bearing correctness invariant is tested explicitly: compaction never
splits an assistant ``tool_call`` from its matching ``ToolMessage`` (a model API
rejects a dangling pair).
"""

from __future__ import annotations

from typing import Any

from langchain_core.messages import (
    AIMessage,
    BaseMessage,
    HumanMessage,
    RemoveMessage,
    SystemMessage,
    ToolMessage,
)
from langgraph.graph.message import REMOVE_ALL_MESSAGES
from loosecanvas.context_middleware import (
    SOURCE_PREAMBLE_MARKER,
    ContextCompactionMiddleware,
    ContextInstrumentationMiddleware,
    ContextStats,
    approx_tokens,
    compact_messages,
    measure_context,
    message_text,
    summarize_tool_call,
)

# ── Message builders ──────────────────────────────────────────────────────────


def _human(text: str, mid: str) -> HumanMessage:
    return HumanMessage(content=text, id=mid)


def _ai(
    text: str = "",
    tool_calls: list[dict[str, Any]] | None = None,
    *,
    mid: str,
) -> AIMessage:
    return AIMessage(content=text, tool_calls=tool_calls or [], id=mid)


def _tool(content: str, tool_call_id: str, name: str, mid: str) -> ToolMessage:
    return ToolMessage(content=content, tool_call_id=tool_call_id, name=name, id=mid)


def _call(name: str, args: dict[str, Any], call_id: str) -> dict[str, Any]:
    return {"name": name, "args": args, "id": call_id, "type": "tool_call"}


def _preamble(mid: str = "src") -> HumanMessage:
    text = f"{SOURCE_PREAMBLE_MARKER} below. Treat it as background.\n--- BEGIN ---\nlots of source text\n--- END ---"
    return HumanMessage(content=text, id=mid)


def _completed_turn(
    n: int,
    *,
    ask: str = "add a node",
    tool_name: str = "create_node",
    tool_args: dict[str, Any] | None = None,
    tool_result: str = "ok",
    final: str = "done",
) -> list[BaseMessage]:
    """A full user turn: human ask → AI tool call → tool result → AI prose."""
    args = (
        tool_args if tool_args is not None else {"node_id": f"n{n}", "label": f"N{n}"}
    )
    cid = f"call-{n}"
    return [
        _human(f"{ask} {n}", f"h{n}"),
        _ai("", [_call(tool_name, args, cid)], mid=f"a{n}"),
        _tool(tool_result, cid, tool_name, f"t{n}"),
        _ai(final, mid=f"a{n}b"),
    ]


def _tool_call_ids(messages: list[BaseMessage]) -> set[str]:
    ids: set[str] = set()
    for m in messages:
        if isinstance(m, AIMessage):
            for c in m.tool_calls:
                ids.add(str(c.get("id")))
    return ids


def _tool_result_ids(messages: list[BaseMessage]) -> set[str]:
    return {m.tool_call_id for m in messages if isinstance(m, ToolMessage)}


# ── approx_tokens / summarize_tool_call ───────────────────────────────────────


def test_approx_tokens_empty_is_zero_and_scales() -> None:
    assert approx_tokens("") == 0
    assert approx_tokens("a" * 4) >= 1
    assert approx_tokens("a" * 400) > approx_tokens("a" * 40)


def test_summarize_tool_call_uses_activity_label_convention() -> None:
    assert (
        summarize_tool_call(
            "create_edge", {"source_id": "a", "target_id": "b", "label": "x"}
        )
        == "create_edge:a->b"
    )
    assert (
        summarize_tool_call("create_node", {"node_id": "n1", "label": "N1"})
        == "create_node:n1"
    )
    assert summarize_tool_call("reveal_node", {"node_id": "x"}) == "reveal_node:x"
    assert summarize_tool_call("get_canvas_state", {}) == "get_canvas_state"
    assert summarize_tool_call("fog_all_except", {"node_ids": ["a", "b"]}).startswith(
        "fog_all_except"
    )


# ── measure_context ───────────────────────────────────────────────────────────


def test_measure_context_classifies_messages_and_detects_source_seed() -> None:
    msgs: list[BaseMessage] = [
        _preamble(),
        _human("hi", "h1"),
        _ai("", [_call("get_canvas_state", {}, "c1")], mid="a1"),
        _tool('{"nodes": []}', "c1", "get_canvas_state", "t1"),
        _ai("here you go", mid="a1b"),
    ]
    stats = measure_context(msgs)
    assert isinstance(stats, ContextStats)
    assert stats.total_messages == 5
    assert stats.human_messages == 2
    assert stats.ai_messages == 2
    assert stats.tool_messages == 1
    assert stats.ai_tool_call_messages == 1
    assert stats.source_seed_tokens > 0
    assert stats.approx_tokens >= stats.source_seed_tokens


def test_measure_context_no_source_seed_is_zero() -> None:
    stats = measure_context([_human("hi", "h1"), _ai("hey", mid="a1")])
    assert stats.source_seed_tokens == 0
    assert stats.total_messages == 2


# ── compact_messages — keep recent verbatim ───────────────────────────────────


def test_compact_keeps_recent_turns_verbatim() -> None:
    msgs: list[BaseMessage] = [_preamble()]
    for n in (1, 2, 3):
        msgs += _completed_turn(n)
    out = compact_messages(msgs, keep_recent_turns=2)
    out_ids = [m.id for m in out]
    # Turns 2 and 3 (recent) survive verbatim — their tool messages remain.
    assert "t2" in out_ids and "t3" in out_ids  # recent tool results kept
    assert "t1" not in out_ids  # old tool result collapsed away
    assert "h1" in out_ids  # old human ask kept verbatim


def test_compact_collapses_old_tool_chatter_into_receipt() -> None:
    msgs: list[BaseMessage] = [_preamble()]
    msgs += _completed_turn(
        1,
        tool_name="create_node",
        tool_args={"node_id": "x", "label": "X"},
        final="added X",
    )
    msgs += _completed_turn(2)  # recent, kept verbatim
    out = compact_messages(msgs, keep_recent_turns=1)
    # The old turn keeps its human ask, and folds prose + a receipt into one AIMessage
    # that has NO tool_calls; its ToolMessage is gone.
    assert "t1" not in [m.id for m in out]
    ai_after_h1 = next(
        m
        for i, m in enumerate(out)
        if isinstance(m, AIMessage) and out[i - 1].id == "h1"
    )
    assert ai_after_h1.tool_calls == []
    text = str(ai_after_h1.content)
    assert "added X" in text
    assert "create_node:x" in text


def test_compact_never_leaves_dangling_tool_pairs() -> None:
    msgs: list[BaseMessage] = [_preamble()]
    for n in (1, 2, 3, 4):
        msgs += _completed_turn(n)
    out = compact_messages(msgs, keep_recent_turns=2)
    # Every surviving tool_call has a matching tool result and vice-versa.
    assert _tool_call_ids(out) == _tool_result_ids(out)


def test_compact_preserves_source_preamble_even_under_tiny_budget() -> None:
    msgs: list[BaseMessage] = [_preamble()]
    for n in (1, 2, 3, 4, 5):
        msgs += _completed_turn(n)
    out = compact_messages(msgs, keep_recent_turns=1, token_budget=5)
    assert out[0].id == "src"
    assert SOURCE_PREAMBLE_MARKER in str(out[0].content)


def test_compact_drops_superseded_canvas_state_payload() -> None:
    big_payload = '{"visible_node_ids": ' + str(list(range(200))) + "}"
    msgs: list[BaseMessage] = [
        _preamble(),
        _human("show me", "h1"),
        _ai("", [_call("get_canvas_state", {}, "c1")], mid="a1"),
        _tool(big_payload, "c1", "get_canvas_state", "t1"),
        _ai("here", mid="a1b"),
    ]
    msgs += _completed_turn(2)  # recent verbatim turn
    out = compact_messages(msgs, keep_recent_turns=1)
    blob = "".join(str(m.content) for m in out)
    assert big_payload not in blob  # the stale payload is gone
    assert "get_canvas_state" in blob  # but the receipt name survives


def test_compact_token_budget_drops_oldest_turns_first() -> None:
    msgs: list[BaseMessage] = [_preamble()]
    for n in (1, 2, 3, 4, 5):
        msgs += _completed_turn(n)
    out = compact_messages(msgs, keep_recent_turns=1, token_budget=200)
    ids = [m.id for m in out]
    # Recent turn (5) and anchor always survive.
    assert "src" in ids
    assert "h5" in ids and "t5" in ids
    # If anything old was dropped, the oldest (turn 1) goes before a newer one (turn 4).
    if "h1" not in ids:
        assert ids.index("src") == 0
    # Either we fit the budget, or we're already down to anchor + recent turns.
    assert measure_context(out).approx_tokens <= 200 or "h2" not in ids


def test_compact_is_noop_when_only_recent_turns() -> None:
    msgs: list[BaseMessage] = [_preamble()]
    msgs += _completed_turn(1)
    out = compact_messages(msgs, keep_recent_turns=4)
    assert [m.id for m in out] == [m.id for m in msgs]


# ── Middleware classes ────────────────────────────────────────────────────────


def test_instrumentation_middleware_measures_without_mutating() -> None:
    captured: list[ContextStats] = []
    mw = ContextInstrumentationMiddleware(on_measure=captured.append)
    msgs: list[BaseMessage] = [_preamble(), _human("hi", "h1"), _ai("hey", mid="a1")]
    result = mw.before_model({"messages": msgs}, None)
    assert result is None  # instrumentation never changes state
    assert len(captured) == 1
    assert captured[0].total_messages == 3


def test_compaction_middleware_emits_remove_all_then_rebuilt() -> None:
    msgs: list[BaseMessage] = [_preamble()]
    for n in (1, 2, 3):
        msgs += _completed_turn(n)
    mw = ContextCompactionMiddleware(keep_recent_turns=1)
    result = mw.before_model({"messages": msgs}, None)
    assert result is not None
    new_msgs = result["messages"]
    assert isinstance(new_msgs[0], RemoveMessage)
    assert new_msgs[0].id == REMOVE_ALL_MESSAGES
    # The rest are the rebuilt history (no RemoveMessage among them).
    assert all(not isinstance(m, RemoveMessage) for m in new_msgs[1:])


def test_compaction_middleware_noop_returns_none() -> None:
    msgs: list[BaseMessage] = [_preamble(), *_completed_turn(1)]
    mw = ContextCompactionMiddleware(keep_recent_turns=4)
    assert mw.before_model({"messages": msgs}, None) is None


def test_compaction_middleware_handles_leading_system_message() -> None:
    msgs: list[BaseMessage] = [SystemMessage(content="sys", id="s0"), _preamble()]
    for n in (1, 2, 3):
        msgs += _completed_turn(n)
    mw = ContextCompactionMiddleware(keep_recent_turns=1)
    result = mw.before_model({"messages": msgs}, None)
    assert result is not None
    rebuilt = result["messages"][1:]  # skip RemoveMessage
    assert isinstance(rebuilt[0], SystemMessage)
    assert rebuilt[0].id == "s0"


# ── Drift guard: marker must match the harness preamble ───────────────────────


def test_source_marker_matches_agent_harness_preamble() -> None:
    from loosecanvas.agent_harness import _format_source_preamble

    preamble = _format_source_preamble("Title", "body text")
    assert preamble.startswith(SOURCE_PREAMBLE_MARKER)


# ── Harness wiring: build_context_middleware ──────────────────────────────────


def test_build_context_middleware_default(monkeypatch: Any) -> None:
    from loosecanvas import agent_harness

    for var in (
        "LOOSECANVAS_CONTEXT_COMPACTION",
        "LOOSECANVAS_CONTEXT_KEEP_TURNS",
        "LOOSECANVAS_CONTEXT_TOKEN_BUDGET",
    ):
        monkeypatch.delenv(var, raising=False)
    mw = agent_harness.build_context_middleware()
    assert isinstance(mw[0], ContextInstrumentationMiddleware)
    assert any(isinstance(m, ContextCompactionMiddleware) for m in mw)
    compaction = next(m for m in mw if isinstance(m, ContextCompactionMiddleware))
    assert compaction._keep_recent_turns == 4
    # W0b: unset budget now defaults to a generous runaway backstop (not unbounded).
    assert compaction._token_budget == agent_harness.DEFAULT_CONTEXT_TOKEN_BUDGET


def test_build_context_middleware_compaction_disabled(monkeypatch: Any) -> None:
    from loosecanvas import agent_harness

    monkeypatch.setenv("LOOSECANVAS_CONTEXT_COMPACTION", "0")
    mw = agent_harness.build_context_middleware()
    assert len(mw) == 1
    assert isinstance(mw[0], ContextInstrumentationMiddleware)


def test_build_context_middleware_env_overrides(monkeypatch: Any) -> None:
    from loosecanvas import agent_harness

    monkeypatch.setenv("LOOSECANVAS_CONTEXT_COMPACTION", "1")
    monkeypatch.setenv("LOOSECANVAS_CONTEXT_KEEP_TURNS", "2")
    monkeypatch.setenv("LOOSECANVAS_CONTEXT_TOKEN_BUDGET", "12000")
    mw = agent_harness.build_context_middleware()
    compaction = next(m for m in mw if isinstance(m, ContextCompactionMiddleware))
    assert compaction._keep_recent_turns == 2
    assert compaction._token_budget == 12000


def test_build_context_middleware_budget_optout(monkeypatch: Any) -> None:
    """W0b: an explicit '0' opts out of the default budget (unbounded collapse-only)."""
    from loosecanvas import agent_harness

    monkeypatch.delenv("LOOSECANVAS_CONTEXT_COMPACTION", raising=False)
    monkeypatch.setenv("LOOSECANVAS_CONTEXT_TOKEN_BUDGET", "0")
    mw = agent_harness.build_context_middleware()
    compaction = next(m for m in mw if isinstance(m, ContextCompactionMiddleware))
    assert compaction._token_budget is None


def test_default_budget_bounds_a_long_history(monkeypatch: Any) -> None:
    """W0b: with the default budget (env unset), a very long session compacts to a
    bounded working set — old receipts drop, the anchor + recent turns survive."""
    from loosecanvas import agent_harness

    for var in (
        "LOOSECANVAS_CONTEXT_COMPACTION",
        "LOOSECANVAS_CONTEXT_KEEP_TURNS",
        "LOOSECANVAS_CONTEXT_TOKEN_BUDGET",
    ):
        monkeypatch.delenv(var, raising=False)
    mw = agent_harness.build_context_middleware()
    compaction = next(m for m in mw if isinstance(m, ContextCompactionMiddleware))

    # A long session: a system anchor + 200 turns of (human ask, AI reply). Each
    # reply is long enough that the raw history far exceeds the default budget.
    long_reply = "detail " * 80  # ~140 tokens
    msgs: list[BaseMessage] = [SystemMessage(content="sys", id="sys")]
    for i in range(200):
        msgs.append(_human(f"question number {i} about the graph", mid=f"h{i}"))
        msgs.append(_ai(f"{long_reply} (reply {i})", mid=f"a{i}"))
    raw_tokens = sum(approx_tokens(message_text(m)) for m in msgs)
    assert raw_tokens > agent_harness.DEFAULT_CONTEXT_TOKEN_BUDGET  # precondition

    out = compact_messages(
        msgs,
        keep_recent_turns=compaction._keep_recent_turns,
        token_budget=compaction._token_budget,
    )
    out_tokens = sum(approx_tokens(message_text(m)) for m in out)
    assert out_tokens <= agent_harness.DEFAULT_CONTEXT_TOKEN_BUDGET + 500
    assert len(out) < len(msgs)
    # The most recent turn survives verbatim.
    assert any(message_text(m) == "question number 199 about the graph" for m in out)