From 439d92518336af87072d205beabec3e0890b3fef Mon Sep 17 00:00:00 2001 From: ciregenz Date: Tue, 23 Jun 2026 00:01:51 -0700 Subject: [PATCH] [eric] agents: extract token-accounting + compaction trigger into manager/context_budget + 10-test suite --- backend/apps/agents/agent_manager.py | 51 ++------ backend/apps/agents/manager/context_budget.py | 70 +++++++++++ backend/tests/test_context_budget.py | 110 ++++++++++++++++++ 3 files changed, 187 insertions(+), 44 deletions(-) create mode 100644 backend/apps/agents/manager/context_budget.py create mode 100644 backend/tests/test_context_budget.py diff --git a/backend/apps/agents/agent_manager.py b/backend/apps/agents/agent_manager.py index 94daa4ec..566b98b4 100644 --- a/backend/apps/agents/agent_manager.py +++ b/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) diff --git a/backend/apps/agents/manager/context_budget.py b/backend/apps/agents/manager/context_budget.py new file mode 100644 index 00000000..d7c2c17e --- /dev/null +++ b/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), + }) diff --git a/backend/tests/test_context_budget.py b/backend/tests/test_context_budget.py new file mode 100644 index 00000000..085dd483 --- /dev/null +++ b/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