"""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)