From a3e30b1af34f437a31808548d99559690d735b45 Mon Sep 17 00:00:00 2001 From: JimStenstrom <30080538+JimStenstrom@users.noreply.github.com> Date: Sat, 6 Jun 2026 10:57:37 -0500 Subject: [PATCH] fix(compaction): add mid-turn in-flight compression safety valve MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Proactive context compression runs only once per turn, in the preflight before the tool-calling loop, so a single long autonomous turn that makes many tool calls can grow context unbounded and overflow the model's window mid-turn — today only a reactive post-error path and the Ollama hard-abort catch it. Observed on a local model: a 6-tool-call turn reached 2.6x the preflight threshold with zero compaction. Add _maybe_inflight_compress(), called at the top of the tool-calling loop, which re-checks compression but only once the request nears the model's real context window — an emergency fraction (HERMES_INFLIGHT_COMPRESS_FRACTION, default 0.85), well above the 50% preflight threshold. Normal turns stay below it and keep their prompt-cache prefix; the one-time mid-turn cache bust is paid only when a turn would otherwise overflow. Mirrors the preflight's post-compression resets and clears the caller's history ref so the session-DB flush writes the compacted messages. Complements (not replaces) the preflight and reactive-on-error compaction paths. Refs #36624 --- agent/conversation_loop.py | 124 +++++++++++++++++++++++++++++ tests/test_inflight_compression.py | 122 ++++++++++++++++++++++++++++ 2 files changed, 246 insertions(+) create mode 100644 tests/test_inflight_compression.py diff --git a/agent/conversation_loop.py b/agent/conversation_loop.py index a07d0ed4b0b62..ad8f3eb1ccc94 100644 --- a/agent/conversation_loop.py +++ b/agent/conversation_loop.py @@ -361,6 +361,115 @@ def _get_continuation_prompt(is_partial_stub: bool, dropped_tools: Optional[List ) +def _maybe_inflight_compress( + agent: Any, + messages: List[Dict[str, Any]], + system_message: Any, + active_system_prompt: Any, + task_id: str, + *, + api_call_count: int, +) -> tuple: + """Mid-turn emergency context-compression safety valve. + + Preflight compression runs once per turn (before the tool-calling loop) + and is intentionally per-turn so it never mutates the conversation + mid-turn and break the prompt-cache prefix. But a single long autonomous + turn that makes many tool calls can grow context unbounded and overflow + the model's window before the turn ends — the proactive compressor never + gets another look until the *next* user turn, and today only a reactive + post-error path (and the Ollama hard-abort) catches it. + + This valve re-checks compression inside the loop, but ONLY once the + request approaches the model's real context window — an EMERGENCY fraction + (``HERMES_INFLIGHT_COMPRESS_FRACTION``, default 0.85), well above the 50% + preflight threshold. Normal turns stay below it and keep their cached + prefix; we pay the one-time mid-turn cache bust only when a turn would + otherwise overflow. Complements (does not replace) the preflight and + reactive-on-error compaction paths. + + Returns ``(messages, active_system_prompt, fired)``. When ``fired`` is + True the caller must reset its ``conversation_history`` reference so the + session-DB flush writes the compacted messages (mirrors preflight). + """ + compressor = getattr(agent, "context_compressor", None) + if ( + not getattr(agent, "compression_enabled", False) + or compressor is None + or api_call_count <= 1 + ): + return messages, active_system_prompt, False + + protect = compressor.protect_first_n + compressor.protect_last_n + 1 + ctx_len = getattr(compressor, "context_length", 0) or 0 + if len(messages) <= protect or ctx_len <= 0: + return messages, active_system_prompt, False + + try: + frac = float(os.getenv("HERMES_INFLIGHT_COMPRESS_FRACTION", "0.85")) + except ValueError: + frac = 0.85 + frac = min(max(frac, 0.5), 0.98) + emergency = int(ctx_len * frac) + + # Prefer the real provider prompt count from the previous API call + # (kept current by compressor.update_from_response); fall back to a rough + # estimate only when we have no real number yet. + inflight_tokens = getattr(compressor, "last_prompt_tokens", 0) or 0 + if inflight_tokens <= 0: + inflight_tokens = estimate_request_tokens_rough( + messages, + system_prompt=active_system_prompt or "", + tools=getattr(agent, "tools", None) or None, + ) + + if inflight_tokens < emergency or not compressor.should_compress(inflight_tokens): + return messages, active_system_prompt, False + + logger.warning( + "In-flight compression (mid-turn safety valve): ~%s tokens >= %s " + "emergency threshold (%.0f%% of %s ctx) at API call #%d — preflight " + "only runs once per turn; compacting to avoid a window overflow.", + f"{inflight_tokens:,}", f"{emergency:,}", frac * 100, + f"{ctx_len:,}", api_call_count, + ) + try: + agent._emit_status( + f"📦 In-flight compression at ~{inflight_tokens:,} tokens " + f"(nearing the {ctx_len:,}-token context limit)." + ) + except Exception: + pass + + fired = False + for _pass in range(3): + orig_len = len(messages) + messages, active_system_prompt = agent._compress_context( + messages, system_message, + approx_tokens=inflight_tokens, + task_id=task_id, + ) + if len(messages) >= orig_len: + break # cannot compress further + fired = True + # Mirror preflight's post-compression resets so the model gets a fresh + # budget on the compacted context. + agent._empty_content_retries = 0 + agent._thinking_prefill_retries = 0 + agent._last_content_with_tools = None + agent._last_content_tools_all_housekeeping = False + agent._mute_post_response = False + inflight_tokens = estimate_request_tokens_rough( + messages, + system_prompt=active_system_prompt or "", + tools=getattr(agent, "tools", None) or None, + ) + if inflight_tokens < emergency: + break + + return messages, active_system_prompt, fired + + def run_conversation( agent, user_message: str, @@ -838,6 +947,21 @@ def run_conversation( agent._safe_print(f"\n⚠️ Iteration budget exhausted ({agent.iteration_budget.used}/{agent.iteration_budget.max_total} iterations used)") break + # ── In-flight emergency context compression (mid-turn safety valve) ── + # Preflight compression runs only once per turn (before this loop), so a + # single long autonomous turn that makes many tool calls can overflow the + # model's window mid-turn. Re-check near the context limit (an emergency + # fraction, not the 50% preflight threshold) so normal turns keep their + # prompt-cache prefix. Pre-empts the reactive post-error / Ollama-abort paths. + messages, active_system_prompt, _inflight_compressed = _maybe_inflight_compress( + agent, messages, system_message, active_system_prompt, + effective_task_id, api_call_count=api_call_count, + ) + if _inflight_compressed: + # Compaction may have created a new session — clear the history + # reference so the session-DB flush writes the compacted messages. + conversation_history = None + # Fire step_callback for gateway hooks (agent:step event) if agent.step_callback is not None: try: diff --git a/tests/test_inflight_compression.py b/tests/test_inflight_compression.py new file mode 100644 index 0000000000000..46999fd23eb6f --- /dev/null +++ b/tests/test_inflight_compression.py @@ -0,0 +1,122 @@ +"""Unit tests for the mid-turn in-flight compression safety valve. + +``_maybe_inflight_compress`` complements the once-per-turn preflight +compression: it re-checks context size *inside* the tool-calling loop, but +only once the request nears the model's real context window (an emergency +fraction, not the 50% preflight threshold), so normal turns keep their +prompt-cache prefix and only a turn that would otherwise overflow pays the +mid-turn cache bust. +""" + +from agent.conversation_loop import _maybe_inflight_compress + + +class FakeCompressor: + def __init__(self, context_length, last_prompt_tokens, *, block=False): + self.context_length = context_length + self.last_prompt_tokens = last_prompt_tokens + self.threshold_tokens = int(context_length * 0.5) # preflight threshold + self.protect_first_n = 1 + self.protect_last_n = 1 + self._block = block + + def should_compress(self, tokens): + # Mirrors the real anti-thrash gate: above threshold unless backed off. + return tokens >= self.threshold_tokens and not self._block + + +class FakeAgent: + def __init__(self, compressor, *, enabled=True): + self.compression_enabled = enabled + self.context_compressor = compressor + self.tools = None + self.compress_calls = 0 + # Pre-set to non-zero so we can assert the helper resets them. + self._empty_content_retries = 7 + self._thinking_prefill_retries = 7 + self._last_content_with_tools = "stale" + self._last_content_tools_all_housekeeping = True + self._mute_post_response = True + + def _emit_status(self, *_a, **_k): + pass + + def _compress_context(self, messages, system_message, *, approx_tokens=None, task_id="default"): + self.compress_calls += 1 + # Shrink hard: keep the first message and the last two. + return [messages[0]] + messages[-2:], "compressed-sys" + + +def _msgs(n=12): + return [{"role": "user" if i % 2 else "assistant", "content": f"message {i}"} for i in range(n)] + + +def test_fires_when_request_nears_context_limit(): + comp = FakeCompressor(context_length=40000, last_prompt_tokens=36000) # > 0.85*40000 + agent = FakeAgent(comp) + msgs, asp, fired = _maybe_inflight_compress( + agent, _msgs(), "sys", "orig-sys", "default", api_call_count=3, + ) + assert fired is True + assert agent.compress_calls == 1 + assert len(msgs) < 12 + assert asp == "compressed-sys" + # post-compression resets mirror preflight + assert agent._empty_content_retries == 0 + assert agent._thinking_prefill_retries == 0 + assert agent._last_content_with_tools is None + + +def test_skips_on_first_iteration(): + comp = FakeCompressor(context_length=40000, last_prompt_tokens=39000) + agent = FakeAgent(comp) + msgs, asp, fired = _maybe_inflight_compress( + agent, _msgs(), "sys", "orig-sys", "default", api_call_count=1, + ) + assert fired is False + assert agent.compress_calls == 0 + assert asp == "orig-sys" + + +def test_skips_below_emergency_even_if_above_preflight_threshold(): + # 24000 is above the 20000 preflight threshold but below the 34000 (85%) + # emergency line — preflight handles this between turns, not the valve. + comp = FakeCompressor(context_length=40000, last_prompt_tokens=24000) + agent = FakeAgent(comp) + _msgs_out, _asp, fired = _maybe_inflight_compress( + agent, _msgs(), "sys", "orig-sys", "default", api_call_count=5, + ) + assert fired is False + assert agent.compress_calls == 0 + + +def test_skips_when_compression_disabled(): + comp = FakeCompressor(context_length=40000, last_prompt_tokens=39000) + agent = FakeAgent(comp, enabled=False) + _m, _a, fired = _maybe_inflight_compress( + agent, _msgs(), "sys", "orig-sys", "default", api_call_count=3, + ) + assert fired is False + assert agent.compress_calls == 0 + + +def test_anti_thrash_blocks_repeat_compaction(): + comp = FakeCompressor(context_length=40000, last_prompt_tokens=39000, block=True) + agent = FakeAgent(comp) + _m, _a, fired = _maybe_inflight_compress( + agent, _msgs(), "sys", "orig-sys", "default", api_call_count=4, + ) + assert fired is False + assert agent.compress_calls == 0 + + +def test_emergency_fraction_env_override(monkeypatch): + # Lower the emergency line to 50% — now 24000 (60% of 40000) trips it. + monkeypatch.setenv("HERMES_INFLIGHT_COMPRESS_FRACTION", "0.5") + comp = FakeCompressor(context_length=40000, last_prompt_tokens=24000) + agent = FakeAgent(comp) + _m, _a, fired = _maybe_inflight_compress( + agent, _msgs(), "sys", "orig-sys", "default", api_call_count=3, + ) + assert fired is True + assert agent.compress_calls == 1