[eric] agents: extract token-accounting + compaction trigger into manager/context_budget + 10-test suite

backend/apps/agents/agent_manager.py

@@ -53,6 +53,7 @@ from backend.apps.agents.manager import metadata
 from backend.apps.agents.manager.session.apply_context_window import apply_context_window
 from backend.apps.agents.manager.session import lifecycle
 from backend.apps.agents.manager.permissions import path_gate
+from backend.apps.agents.manager import context_budget
 from backend.apps.agents.manager.session.workspace_git import _detect_git_identity, _ensure_cwd_git_repo
 from backend.apps.agents.manager.prompt.tool_catalog import (
     FULL_TOOLS,
@@ -370,32 +371,7 @@ class AgentManager:
     # ------------------------------------------------------------------
 
     def _maybe_compact(self, session: AgentSession, force: bool = False) -> bool:
-        """Run summarizer when ctx_used_pct >= compact_threshold_pct (or force).
-
-        Returns True if a new summary was produced. Mutates session state:
-        sets compacted_through_msg_id and emits a context_status event.
-        Never modifies session.messages, originals stay around for the
-        UI drawer; only the history *sent to the SDK* is trimmed (handled
-        in _build_history_prefix lookups).
-        """
-        ctx_used = session.tokens.get("input", 0) / max(1, session.context_window)
-        if not force and ctx_used < session.compact_threshold_pct:
-            return False
-        msgs = _get_branch_messages(session)
-        if len(msgs) < 4:
-            return False
-        # Summarize everything up to (but not including) the last 6
-        # messages, that window keeps recent intent visible to the
-        # model so it doesn't lose its train of thought right after
-        # compaction.
-        cutoff = max(0, len(msgs) - 6)
-        if cutoff == 0:
-            return False
-        last_id = msgs[cutoff - 1].id
-        if session.compacted_through_msg_id == last_id and not force:
-            return False
-        session.compacted_through_msg_id = last_id
-        return True
+        return context_budget.maybe_compact(session, force)
 
     async def _emit_context_update(
         self,
@@ -407,24 +383,11 @@ class AgentManager:
         cache_read_tokens: int = 0,
         cache_read_pct: float = 0.0,
     ) -> None:
-        if input_tokens is None:
-            input_tokens = int(session.tokens.get("input", 0) or 0)
-        if output_tokens is None:
-            output_tokens = int(session.tokens.get("output", 0) or 0)
-        session.tokens["input"] = input_tokens
-        session.tokens["output"] = output_tokens
-        ctx_window = max(1, getattr(session, "context_window", 0) or 200_000)
-        await ws_manager.send_to_session(session_id, "agent:context_update", {
-            "session_id": session_id,
-            "input_tokens": input_tokens,
-            "output_tokens": output_tokens,
-            "cache_read_tokens": cache_read_tokens,
-            "cache_read_pct": cache_read_pct,
-            "ctx_used_pct": round(input_tokens / ctx_window, 4) if input_tokens else 0.0,
-            "context_window": ctx_window,
-            "framework_overhead_tokens": session.framework_overhead_tokens,
-            "active_mcps": list(session.active_mcps),
-        })
+        return await context_budget.emit_context_update(
+            session_id, session,
+            input_tokens=input_tokens, output_tokens=output_tokens,
+            cache_read_tokens=cache_read_tokens, cache_read_pct=cache_read_pct,
+        )
 
     def _build_prompt_content(self, prompt: str, images: list | None = None, context_paths: list | None = None, forced_tools: list[str] | None = None, attached_skills: list | None = None, api_type: str = "anthropic", model: str = ""):
         return _build_prompt_content(prompt, images, context_paths, forced_tools, attached_skills, api_type, model)

backend/apps/agents/manager/context_budget.py

@@ -0,0 +1,70 @@
+"""Token accounting + the context-ratio compaction trigger, lifted out of the agent loop.
+Both operate on a passed AgentSession (no manager state). emit_context_update writes the
+live token counts onto the session and broadcasts them to the UI; maybe_compact decides,
+from the same input_tokens/context_window ratio, whether to mark history for trimming.
+
+Compaction here only MARKS (sets compacted_through_msg_id); it never mutates
+session.messages, the originals stay for the UI drawer and only the history sent to the SDK
+is trimmed downstream (see backend/CLAUDE.md: "compaction must actually trim, not just mark")."""
+
+from typing import Optional
+
+from typeguard import typechecked
+
+from backend.apps.agents.core.models import AgentSession
+from backend.apps.agents.core.ws_manager import ws_manager
+from backend.apps.agents.manager.session.history_compaction import _get_branch_messages as get_branch_messages
+
+
+@typechecked
+def maybe_compact(session: AgentSession, force: bool = False) -> bool:
+    """Mark history for compaction when ctx_used_pct >= compact_threshold_pct (or force).
+    Returns True if a NEW summary boundary was set. Summarizes everything up to (but not
+    including) the last 6 messages so recent intent stays visible to the model. Never
+    touches session.messages."""
+    ctx_used = session.tokens.get("input", 0) / max(1, session.context_window)
+    if not force and ctx_used < session.compact_threshold_pct:
+        return False
+    msgs = get_branch_messages(session)
+    if len(msgs) < 4:
+        return False
+    cutoff = max(0, len(msgs) - 6)
+    if cutoff == 0:
+        return False
+    last_id = msgs[cutoff - 1].id
+    if session.compacted_through_msg_id == last_id and not force:
+        return False
+    session.compacted_through_msg_id = last_id
+    return True
+
+
+@typechecked
+async def emit_context_update(
+    session_id: str,
+    session: AgentSession,
+    *,
+    input_tokens: Optional[int] = None,
+    output_tokens: Optional[int] = None,
+    cache_read_tokens: int = 0,
+    cache_read_pct: float = 0.0,
+) -> None:
+    """Persist the live token counts onto the session and broadcast the context-usage meter
+    to the UI. When input/output aren't supplied, the session's current counts are reused."""
+    if input_tokens is None:
+        input_tokens = int(session.tokens.get("input", 0) or 0)
+    if output_tokens is None:
+        output_tokens = int(session.tokens.get("output", 0) or 0)
+    session.tokens["input"] = input_tokens
+    session.tokens["output"] = output_tokens
+    ctx_window = max(1, getattr(session, "context_window", 0) or 200_000)
+    await ws_manager.send_to_session(session_id, "agent:context_update", {
+        "session_id": session_id,
+        "input_tokens": input_tokens,
+        "output_tokens": output_tokens,
+        "cache_read_tokens": cache_read_tokens,
+        "cache_read_pct": cache_read_pct,
+        "ctx_used_pct": round(input_tokens / ctx_window, 4) if input_tokens else 0.0,
+        "context_window": ctx_window,
+        "framework_overhead_tokens": session.framework_overhead_tokens,
+        "active_mcps": list(session.active_mcps),
+    })

backend/tests/test_context_budget.py

@@ -0,0 +1,110 @@
+"""Rigorous coverage for the token-accounting + compaction-trigger logic lifted into
+manager/context_budget.py. Compaction is correctness-sensitive (backend/CLAUDE.md), so
+every branch of maybe_compact is pinned, plus emit_context_update's token persistence and
+the exact broadcast payload."""
+
+import asyncio
+
+import backend.apps.agents.manager.context_budget as cb
+from backend.apps.agents.core.models import AgentSession, Message
+
+
+def _session_with(messages: int, input_tokens: int, context_window: int = 100, threshold: float = 0.65) -> AgentSession:
+    s = AgentSession(name="t", model="sonnet")
+    s.context_window = context_window
+    s.compact_threshold_pct = threshold
+    s.tokens = {"input": input_tokens, "output": 0}
+    s.messages = [Message(role="user", content=f"m{i}") for i in range(messages)]
+    return s
+
+
+def _capture_ws(monkeypatch):
+    sent = []
+
+    async def fake_send(session_id, event, data):
+        sent.append((event, data))
+
+    monkeypatch.setattr(cb.ws_manager, "send_to_session", fake_send, raising=True)
+    return sent
+
+
+# ---- maybe_compact: every branch -------------------------------------------
+
+def test_compact_skipped_below_threshold():
+    s = _session_with(messages=10, input_tokens=10)  # 0.10 < 0.65
+    assert cb.maybe_compact(s) is False
+    assert s.compacted_through_msg_id is None
+
+
+def test_compact_fires_over_threshold_and_marks_boundary():
+    s = _session_with(messages=7, input_tokens=80)   # 0.80 >= 0.65; cutoff = 7-6 = 1
+    assert cb.maybe_compact(s) is True
+    assert s.compacted_through_msg_id == s.messages[0].id
+
+
+def test_compact_keeps_the_last_six_messages():
+    s = _session_with(messages=10, input_tokens=80)  # cutoff = 10-6 = 4 -> boundary at msgs[3]
+    assert cb.maybe_compact(s) is True
+    assert s.compacted_through_msg_id == s.messages[3].id
+
+
+def test_compact_skipped_with_six_or_fewer_messages():
+    s = _session_with(messages=6, input_tokens=80)   # cutoff = max(0, 6-6) = 0
+    assert cb.maybe_compact(s) is False
+
+
+def test_compact_skipped_under_four_messages():
+    s = _session_with(messages=3, input_tokens=80)
+    assert cb.maybe_compact(s) is False
+
+
+def test_compact_is_idempotent():
+    s = _session_with(messages=7, input_tokens=80)
+    assert cb.maybe_compact(s) is True
+    boundary = s.compacted_through_msg_id
+    assert cb.maybe_compact(s) is False              # already marked through that id
+    assert s.compacted_through_msg_id == boundary
+
+
+def test_force_bypasses_threshold_and_idempotency():
+    s = _session_with(messages=7, input_tokens=1)    # 0.01 < 0.65
+    assert cb.maybe_compact(s, force=True) is True   # force ignores the ratio
+    assert cb.maybe_compact(s, force=True) is True   # force re-marks even when unchanged
+
+
+# ---- emit_context_update ----------------------------------------------------
+
+def test_emit_persists_tokens_and_broadcasts(monkeypatch):
+    sent = _capture_ws(monkeypatch)
+    s = AgentSession(name="t", model="sonnet")
+    s.context_window = 1000
+
+    asyncio.run(cb.emit_context_update("sid", s, input_tokens=250, output_tokens=40, cache_read_tokens=10, cache_read_pct=0.5))
+
+    assert s.tokens["input"] == 250 and s.tokens["output"] == 40
+    assert len(sent) == 1
+    event, data = sent[0]
+    assert event == "agent:context_update"
+    assert data["input_tokens"] == 250 and data["output_tokens"] == 40
+    assert data["cache_read_tokens"] == 10 and data["cache_read_pct"] == 0.5
+    assert data["ctx_used_pct"] == round(250 / 1000, 4)
+    assert data["context_window"] == 1000
+
+
+def test_emit_defaults_to_existing_session_tokens(monkeypatch):
+    sent = _capture_ws(monkeypatch)
+    s = AgentSession(name="t", model="sonnet")
+    s.tokens = {"input": 123, "output": 7}
+
+    asyncio.run(cb.emit_context_update("sid", s))  # no explicit tokens -> reuse the session's
+    _, data = sent[0]
+    assert data["input_tokens"] == 123 and data["output_tokens"] == 7
+
+
+def test_emit_zero_input_yields_zero_ctx_pct(monkeypatch):
+    sent = _capture_ws(monkeypatch)
+    s = AgentSession(name="t", model="sonnet")
+
+    asyncio.run(cb.emit_context_update("sid", s, input_tokens=0))
+    _, data = sent[0]
+    assert data["ctx_used_pct"] == 0.0