From 5c9710712e3d33d6a2e6cf05672130058bebfb88 Mon Sep 17 00:00:00 2001 From: ciregenz Date: Mon, 3 Aug 2026 19:48:55 -0700 Subject: [PATCH] [eric] agents: mid-turn context breaker, one giant turn compacts and continues instead of starving --- backend/apps/agents/manager/context_budget.py | 58 ++++++++++++++++--- .../streaming/handle_assistant_message.py | 19 ++++++ .../apps/agents/manager/streaming/state.py | 3 + backend/tests/test_context_budget.py | 52 +++++++++++++++++ backend/tests/test_context_pressure_valve.py | 31 ++++++++++ 5 files changed, 156 insertions(+), 7 deletions(-) diff --git a/backend/apps/agents/manager/context_budget.py b/backend/apps/agents/manager/context_budget.py index 8e4832d2..0d4cacfe 100644 --- a/backend/apps/agents/manager/context_budget.py +++ b/backend/apps/agents/manager/context_budget.py @@ -7,13 +7,62 @@ 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 typing import Dict, 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 +from backend.apps.agents.manager.streaming.state import TurnState + + +@typechecked +def compact_trigger_tokens(session: AgentSession) -> int: + """The token count where compaction fires: the TIGHTER of the pct threshold and the + absolute ceiling (on a 200K window the pct wins at 130K; on a 1M window the ceiling + wins at 180K, not 650K).""" + window = max(1, session.context_window) + abs_pct = min(1.0, session.compact_abs_ceiling_tokens / window) + return int(window * min(session.compact_threshold_pct, abs_pct)) + + +CONTINUATION_PROMPT = ( + "Continue the task exactly where you left off. Your earlier progress in this chat is " + "summarized above; do not redo completed steps, pick up at the next unfinished one." +) + + +@typechecked +def maybe_break_midturn(session: AgentSession, turn: TurnState, msg_usage: Dict) -> bool: + """Mid-turn context breaker: one giant turn (dozens of tool calls off a single ask) can + blow past every turn-boundary wall, so when a request's input usage crosses the compact + trigger MID-turn, end the turn at the next message boundary (the pending_continuation + break the MCPActivate flow already uses), force-compact, and auto-continue fresh. + Live incident: 925K/1M with zero CLI compact_boundary events, task abandoned mid-way.""" + try: + total = ( + int(msg_usage.get("input_tokens") or 0) + + int(msg_usage.get("cache_creation_input_tokens") or 0) + + int(msg_usage.get("cache_read_input_tokens") or 0) + ) + except Exception: + return False + if total <= 0: + return False + # Keep the session's counter honest mid-turn: a broken turn never gets its ResultMessage accounting, and the next pre-send guard reads this. + session.tokens["input"] = total + if total < compact_trigger_tokens(session): + turn.saw_input_below_trigger = True + return False + if turn.context_break_fired or not turn.saw_input_below_trigger: + return False + turn.context_break_fired = True + maybe_compact(session, force=True) + session.needs_fresh_session = True + session.pending_continuation = True + session.pending_continuation_prompt = CONTINUATION_PROMPT + return True @typechecked @@ -22,12 +71,7 @@ def maybe_compact(session: AgentSession, force: bool = False) -> bool: 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.""" - window = max(1, session.context_window) - # Fire at the TIGHTER of the pct or the absolute ceiling: on a 200K window the pct wins (130K), on a 1M window the ceiling wins (180K, not 650K). Not "just 65%". - abs_pct = min(1.0, session.compact_abs_ceiling_tokens / window) - trigger = min(session.compact_threshold_pct, abs_pct) - ctx_used = session.tokens.get("input", 0) / window - if not force and ctx_used < trigger: + if not force and session.tokens.get("input", 0) < compact_trigger_tokens(session): return False msgs = get_branch_messages(session) if len(msgs) < 4: diff --git a/backend/apps/agents/manager/streaming/handle_assistant_message.py b/backend/apps/agents/manager/streaming/handle_assistant_message.py index 3465237b..e7a4dfc4 100644 --- a/backend/apps/agents/manager/streaming/handle_assistant_message.py +++ b/backend/apps/agents/manager/streaming/handle_assistant_message.py @@ -5,6 +5,7 @@ Lifted out of the agent loop; mutates the passed TurnState / ThinkingState by re through the manager's live-partial mirror + session registry, exactly as it did inline.""" import asyncio +import logging from typing import Dict, Optional from uuid import uuid4 @@ -74,6 +75,24 @@ async def handle_assistant_message( ot = int(msg_usage.get("output_tokens", 0) or 0) if ot > 0: turn.output_tokens += ot + from backend.apps.agents.manager.context_budget import maybe_break_midturn + if maybe_break_midturn(session, turn, msg_usage): + logging.getLogger(__name__).warning( + f"[context-break] session {session_id}: mid-turn input " + f"{session.tokens.get('input')} crossed the compact trigger; breaking at the " + "next message boundary and continuing on a fresh compacted session" + ) + try: + from backend.apps.service.client import submit_diagnostic + submit_diagnostic({ + "kind": "context_midturn_break", + "session_id": session_id, + "model": session.model, + "input_tokens": session.tokens.get("input"), + "context_window": session.context_window, + }) + except Exception: + pass except Exception: pass diff --git a/backend/apps/agents/manager/streaming/state.py b/backend/apps/agents/manager/streaming/state.py index f7537cc0..6bb33398 100644 --- a/backend/apps/agents/manager/streaming/state.py +++ b/backend/apps/agents/manager/streaming/state.py @@ -56,3 +56,6 @@ class TurnState(BaseModel): baseline_captured: bool = False # CLI compact_boundary events seen this turn; one plus a ProcessError = the autocompact-thrash death the context-pressure valve retries. compact_boundaries: int = 0 + # Mid-turn context breaker: fires once per turn, and only after a below-trigger reading (a turn that STARTS over the trigger must run, or a failed shrink would break-loop forever). + context_break_fired: bool = False + saw_input_below_trigger: bool = False diff --git a/backend/tests/test_context_budget.py b/backend/tests/test_context_budget.py index 7f67931a..3f541ca6 100644 --- a/backend/tests/test_context_budget.py +++ b/backend/tests/test_context_budget.py @@ -180,3 +180,55 @@ def test_recap_leaves_normal_messages_verbatim(): assert clamp_recap_text(text) == text recap = build_history_prefix([Message(role="user", content=text)]) assert text in recap and "elided" not in recap + +# ---- mid-turn breaker: one giant turn can't blow past every wall ------------ + +from backend.apps.agents.manager.streaming.state import TurnState + +def p_usage(total: int) -> dict: + return {"input_tokens": total - 200, "cache_creation_input_tokens": 100, "cache_read_input_tokens": 100, "output_tokens": 5} + + +def test_midturn_break_fires_on_crossing_and_arms_the_continuation(): + s = p_session_with(messages=10, input_tokens=0, context_window=1_000_000) + t = TurnState() + assert cb.maybe_break_midturn(s, t, p_usage(50_000)) is False + assert t.saw_input_below_trigger is True + assert cb.maybe_break_midturn(s, t, p_usage(200_000)) is True + assert t.context_break_fired is True + assert s.needs_fresh_session is True + assert s.pending_continuation is True + assert s.compacted_through_msg_id is not None + assert s.tokens["input"] == 200_000 + + +def test_midturn_break_fires_once_per_turn(): + s = p_session_with(messages=10, input_tokens=0, context_window=1_000_000) + t = TurnState() + cb.maybe_break_midturn(s, t, p_usage(50_000)) + assert cb.maybe_break_midturn(s, t, p_usage(200_000)) is True + assert cb.maybe_break_midturn(s, t, p_usage(300_000)) is False + + +def test_turn_already_over_trigger_at_start_never_breaks(): + # A rebuild that failed to shrink must RUN, not break-loop forever. + s = p_session_with(messages=10, input_tokens=0, context_window=1_000_000) + t = TurnState() + assert cb.maybe_break_midturn(s, t, p_usage(500_000)) is False + assert cb.maybe_break_midturn(s, t, p_usage(600_000)) is False + assert t.context_break_fired is False + + +def test_midturn_break_zero_or_garbage_usage_is_inert(): + s = p_session_with(messages=10, input_tokens=7, context_window=1_000_000) + t = TurnState() + assert cb.maybe_break_midturn(s, t, {}) is False + assert cb.maybe_break_midturn(s, t, {"input_tokens": "nope"}) is False + assert s.tokens["input"] == 7 + + +def test_trigger_formula_matches_maybe_compact(): + s = p_session_with(messages=10, input_tokens=0, context_window=1_000_000) + assert cb.compact_trigger_tokens(s) == 180_000 + s2 = p_session_with(messages=10, input_tokens=0, context_window=200_000) + assert cb.compact_trigger_tokens(s2) == 130_000 diff --git a/backend/tests/test_context_pressure_valve.py b/backend/tests/test_context_pressure_valve.py index da19199e..ca02a4f4 100644 --- a/backend/tests/test_context_pressure_valve.py +++ b/backend/tests/test_context_pressure_valve.py @@ -167,3 +167,34 @@ def test_valve_never_loops(monkeypatch) -> None: assert len(calls) == 2 assert session.status == "error" assert [m for m in session.messages if m.role == "system" and str(m.content).startswith("Error:")] + + +def test_midturn_break_completes_and_fires_the_hidden_continuation(monkeypatch) -> None: + from backend.apps.agents.manager.context_budget import CONTINUATION_PROMPT + session = p_seed_session() + continues: list = [] + + async def fake_run_turn(sess, session_id, prompt_content, options, options_kwargs, + turn, thinking, stderr, resolved_model, api_type, + global_settings, force_respawn=False): + # Simulate maybe_break_midturn firing inside the stream loop: flags set, turn returns at the boundary. + turn.context_break_fired = True + sess.needs_fresh_session = True + sess.pending_continuation = True + sess.pending_continuation_prompt = CONTINUATION_PROMPT + + async def fake_send_message(session_id, prompt, hidden=False, **kwargs): + continues.append({"prompt": prompt, "hidden": hidden}) + + p_install_run_fakes(monkeypatch, fake_run_turn) + monkeypatch.setattr(agent_manager, "send_message", fake_send_message) + + async def main(): + await agent_manager.run_agent_loop(session.id, "audit everything") + await asyncio.sleep(0) + + asyncio.run(main()) + assert session.status == "completed" + assert session.needs_fresh_session is True + assert session.pending_continuation is False + assert continues == [{"prompt": CONTINUATION_PROMPT, "hidden": True}]