From 929c952596a4b8d40c49c96b1f1feb80655ca044 Mon Sep 17 00:00:00 2001
From: Teknium <127238744+teknium1@users.noreply.github.com>
Date: Wed, 22 Jul 2026 21:43:15 -0700
Subject: [PATCH] fix(agent): wire should_compress_preflight into the
 turn-start preflight flow
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Relocates the #20424 wiring: the preflight region moved out of
run_agent.py into agent/turn_context.py (and through the
compression.max_attempts unification, #69315), so the contributor's
elif branch is reapplied at its current home as the else arm of the
threshold dispatch chain.

Integration contracts:
- Byte-identical default: the built-in ContextCompressor inherits
  ContextEngine.should_compress_preflight() -> False, so the default
  path performs no compression and touches no turn bookkeeping
  (pinned by test_builtin_compressor_default_sub_threshold_path_unchanged).
- Attempt-cap: the engine gets exactly ONE compress() pass per turn,
  mutually exclusive with the cap-bounded threshold multi-pass loop,
  so turn-start passes stay within the resolved
  compression.max_attempts budget in every case.
- No-op blocking (#64382 / 377244f7c): an engine pass that no-ops
  (_compress_context returns the input list object) neither sets nor
  clears preflight_compression_blocked and does not re-baseline the
  flush history — a sub-threshold maintenance no-op proves nothing
  about over-threshold compressibility.
- Engine exceptions are swallowed at debug level; cooldown/defer/
  codex-native gates run before the hook is ever consulted.

Salvaged from #20424 by @Beandon13. Fixes #20316.
---
 agent/turn_context.py                     |  67 ++++++
 tests/agent/test_engine_preflight_wire.py | 279 ++++++++++++++++++++++
 2 files changed, 346 insertions(+)
 create mode 100644 tests/agent/test_engine_preflight_wire.py

diff --git a/agent/turn_context.py b/agent/turn_context.py
index c54f8b21a758..fdfee87daa52 100644
--- a/agent/turn_context.py
+++ b/agent/turn_context.py
@@ -828,6 +828,73 @@ def build_turn_context(
                         f"{_preflight_tokens:,}",
                     )
                     break
+        else:
+            # ── Engine-driven sub-threshold preflight maintenance (#20316) ──
+            # None of the threshold-path branches fired (not deferred, no
+            # failure cooldown, not codex-native, and should_compress() said
+            # the request is under pressure). Context engines that override
+            # ``should_compress_preflight()`` (e.g. LCM-style incremental
+            # leaf-chunk compaction) can still request deferred maintenance
+            # below the token threshold. The default
+            # ``ContextEngine.should_compress_preflight()`` returns False, so
+            # the built-in ``ContextCompressor`` path is byte-identical.
+            #
+            # Attempt-cap integration: the engine gets exactly ONE
+            # ``compress()`` pass per turn. It is mutually exclusive with the
+            # threshold multi-pass loop above (if/elif), so turn-start
+            # preflight passes stay bounded by the resolved
+            # ``compression.max_attempts`` cap (floor 1) in every case.
+            #
+            # No-op-blocking integration: a sub-threshold engine pass that
+            # no-ops says nothing about over-threshold compressibility, so it
+            # must neither set nor clear ``_preflight_compression_blocked``
+            # (#64382) — and being in the ``else`` arm it can never run after
+            # the threshold loop has proven a retry ineffective.
+            _engine_preflight = getattr(
+                _compressor, "should_compress_preflight", None
+            )
+            _wants_engine_preflight = False
+            if callable(_engine_preflight):
+                try:
+                    _wants_engine_preflight = bool(_engine_preflight(messages))
+                except Exception as _preflight_exc:
+                    # A buggy engine must never break an otherwise-healthy
+                    # turn: swallow at debug level and skip maintenance.
+                    logger.debug(
+                        "should_compress_preflight raised %s; skipping "
+                        "engine-driven preflight maintenance",
+                        _preflight_exc,
+                    )
+                    _wants_engine_preflight = False
+            if _wants_engine_preflight:
+                logger.info(
+                    "Engine-driven preflight maintenance: %s requested "
+                    "compress() at ~%s tokens (below %s threshold)",
+                    getattr(_compressor, "name", type(_compressor).__name__),
+                    f"{_preflight_tokens:,}",
+                    f"{getattr(_compressor, 'threshold_tokens', 0):,}",
+                )
+                _engine_input = messages
+                messages, active_system_prompt = agent._compress_context(
+                    messages, system_message, approx_tokens=_preflight_tokens,
+                    task_id=effective_task_id,
+                )
+                # ``_compress_context`` returns the INPUT list object on every
+                # skip path (per-session lock held elsewhere, cooldown,
+                # anti-thrash breaker, codex-native routing) and an engine may
+                # legitimately no-op. Only re-baseline the flush history and
+                # re-anchor the user row after a REAL compaction — a skip must
+                # leave the turn's bookkeeping untouched.
+                if messages is not _engine_input:
+                    _preflight_compressed = True
+                    conversation_history = conversation_history_after_compression(
+                        agent, messages
+                    )
+                    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
 
     if _preflight_compressed:
         # Compression rebuilt the list (tail messages are fresh compaction
diff --git a/tests/agent/test_engine_preflight_wire.py b/tests/agent/test_engine_preflight_wire.py
new file mode 100644
index 000000000000..53570fc399ca
--- /dev/null
+++ b/tests/agent/test_engine_preflight_wire.py
@@ -0,0 +1,279 @@
+"""Engine-driven sub-threshold preflight maintenance (#20316, salvaged from #20424).
+
+``build_turn_context`` historically consulted the context engine only through
+``should_compress()`` — the optional ``ContextEngine.should_compress_preflight()``
+hook was never called anywhere in the turn flow, so engines doing LCM-style
+deferred maintenance (incremental leaf-chunk compaction below the token
+threshold) could never run it.  The preflight block now falls through to the
+engine hook when the threshold path does not fire.
+
+Contracts pinned here:
+
+* Byte-identical default: the built-in ``ContextCompressor`` inherits the
+  default ``should_compress_preflight() -> False``, so the sub-threshold path
+  performs no compression, emits no status, and leaves every piece of turn
+  bookkeeping untouched.
+* A True-returning engine gets its ``compress()`` pass exactly ONCE per turn
+  (mutually exclusive with the cap-bounded threshold loop), regardless of the
+  resolved ``compression.max_attempts`` value.
+* No-op interplay (#64382): an engine pass that no-ops (``_compress_context``
+  returns the input list object) must not set ``preflight_compression_blocked``
+  — the sub-threshold pass proves nothing about over-threshold
+  compressibility — and must not re-baseline the flush history.
+* The hook is never consulted when the threshold path fires, when a
+  compression-failure cooldown is active, or when the engine raises.
+"""
+
+from __future__ import annotations
+
+import types
+from unittest.mock import MagicMock, patch
+
+import pytest
+
+from agent.turn_context import TurnContext, build_turn_context
+from tests.agent.test_turn_context import _FakeAgent
+
+
+@pytest.fixture(autouse=True)
+def _stub_runtime_main():
+    """Keep the aux runtime-main global from leaking into sibling tests."""
+    with patch("agent.auxiliary_client.set_runtime_main", lambda *a, **k: None):
+        yield
+
+
+def _history(n_pairs: int = 6) -> list:
+    msgs = []
+    for i in range(n_pairs):
+        msgs.append({"role": "user", "content": f"q{i}"})
+        msgs.append({"role": "assistant", "content": f"a{i}"})
+    return msgs
+
+
+def _stub_compressor(*, preflight=None, threshold_tokens=10**9):
+    """Engine-shaped stub: below threshold, no defer, no cooldown."""
+    comp = types.SimpleNamespace(
+        protect_first_n=2,
+        protect_last_n=2,
+        threshold_tokens=threshold_tokens,
+        last_prompt_tokens=0,
+        should_compress=lambda _tokens=None: False,
+        should_defer_preflight_to_real_usage=lambda _t: False,
+        get_active_compression_failure_cooldown=lambda: None,
+    )
+    if preflight is not None:
+        comp.should_compress_preflight = preflight
+    return comp
+
+
+def _make_agent(compressor):
+    agent = _FakeAgent()
+    agent.compression_enabled = True
+    agent.context_compressor = compressor
+    agent._emit_status = MagicMock()
+    agent._compress_context = MagicMock(
+        side_effect=lambda messages, *_a, **_k: (messages, "SYSTEM")
+    )
+    return agent
+
+
+def _build(agent, **overrides):
+    kwargs = dict(
+        agent=agent,
+        user_message="hello",
+        system_message=None,
+        conversation_history=_history(),
+        task_id=None,
+        stream_callback=None,
+        persist_user_message=None,
+        restore_or_build_system_prompt=lambda *a, **k: None,
+        install_safe_stdio=lambda: None,
+        sanitize_surrogates=lambda s: s,
+        summarize_user_message_for_log=lambda s: s,
+        set_session_context=lambda _sid: None,
+        set_current_write_origin=lambda _o: None,
+        ra=lambda: types.SimpleNamespace(_set_interrupt=lambda *a, **k: None),
+    )
+    kwargs.update(overrides)
+    return build_turn_context(**kwargs)
+
+
+def test_default_false_hook_is_byte_identical_noop():
+    """The default engine (hook returns False) changes nothing sub-threshold."""
+    calls = []
+
+    def _hook(messages):
+        calls.append(list(messages))
+        return False
+
+    agent = _make_agent(_stub_compressor(preflight=_hook))
+
+    ctx = _build(agent)
+
+    assert isinstance(ctx, TurnContext)
+    assert calls, "sub-threshold path must consult the engine hook"
+    agent._compress_context.assert_not_called()
+    agent._emit_status.assert_not_called()
+    assert ctx.preflight_compression_blocked is False
+
+
+def test_engine_without_hook_attribute_is_skipped():
+    """Minimal engines lacking the optional hook must not break the turn."""
+    agent = _make_agent(_stub_compressor(preflight=None))
+
+    ctx = _build(agent)
+
+    assert isinstance(ctx, TurnContext)
+    agent._compress_context.assert_not_called()
+    assert ctx.preflight_compression_blocked is False
+
+
+def test_true_engine_gets_exactly_one_compress_pass():
+    """A True-returning engine triggers exactly one compress() per turn."""
+    hook = MagicMock(return_value=True)
+    agent = _make_agent(_stub_compressor(preflight=hook))
+    # A real compaction: return a NEW list containing this turn's user row.
+    agent._compress_context = MagicMock(
+        side_effect=lambda messages, *_a, **_k: (
+            [dict(m) for m in messages[-3:]],
+            "SYSTEM",
+        )
+    )
+
+    ctx = _build(agent)
+
+    hook.assert_called_once()
+    agent._compress_context.assert_called_once()
+    # Real compaction re-anchors this turn's user message in the new list.
+    assert ctx.messages[ctx.current_turn_user_idx]["content"] == "hello"
+    # A sub-threshold maintenance pass proves nothing about over-threshold
+    # compressibility: the retry-loop blocking flag must stay untouched.
+    assert ctx.preflight_compression_blocked is False
+
+
+def test_true_engine_single_pass_even_with_raised_attempt_cap():
+    """The engine pass is once-per-turn regardless of compression.max_attempts."""
+    hook = MagicMock(return_value=True)
+    agent = _make_agent(_stub_compressor(preflight=hook))
+    agent.max_compression_attempts = 6
+    agent._compress_context = MagicMock(
+        side_effect=lambda messages, *_a, **_k: (
+            [dict(m) for m in messages],
+            "SYSTEM",
+        )
+    )
+
+    _build(agent)
+
+    hook.assert_called_once()
+    agent._compress_context.assert_called_once()
+
+
+def test_true_engine_noop_does_not_defeat_retry_loop_blocking():
+    """#64382 interplay: an engine pass that no-ops must not touch the
+
+    stale-budget blocking machinery — ``preflight_compression_blocked`` stays
+    False (it was never proven ineffective at threshold pressure) and the
+    flush baseline is not re-baselined for a compaction that never happened.
+    """
+    hook = MagicMock(return_value=True)
+    agent = _make_agent(_stub_compressor(preflight=hook))
+    # Skip path: _compress_context returns the INPUT list object.
+    agent._compress_context = MagicMock(
+        side_effect=lambda messages, *_a, **_k: (messages, "SYSTEM")
+    )
+    history = _history()
+
+    ctx = _build(agent, conversation_history=history)
+
+    hook.assert_called_once()
+    agent._compress_context.assert_called_once()
+    assert ctx.preflight_compression_blocked is False
+    # No real compaction -> the caller-provided history object remains the
+    # flush baseline (no re-baseline through
+    # conversation_history_after_compression).
+    assert ctx.conversation_history is history
+
+
+def test_engine_hook_not_consulted_when_threshold_path_fires():
+    """Over-threshold turns take the cap-bounded loop, never the engine arm."""
+    hook = MagicMock(return_value=True)
+    comp = _stub_compressor(preflight=hook, threshold_tokens=100)
+    comp.should_compress = lambda _tokens=None: True
+    comp.context_length = 200_000
+    agent = _make_agent(comp)
+
+    with patch(
+        "agent.turn_context.estimate_request_tokens_rough", return_value=999_999
+    ):
+        _build(agent)
+
+    hook.assert_not_called()
+    # The threshold loop ran instead (progress check breaks after pass 1
+    # because the estimate never shrinks, but the pass itself happened).
+    assert agent._compress_context.call_count >= 1
+
+
+def test_engine_hook_not_consulted_during_failure_cooldown():
+    """An active compression-failure cooldown gates engine maintenance too."""
+    hook = MagicMock(return_value=True)
+    comp = _stub_compressor(preflight=hook)
+    comp.get_active_compression_failure_cooldown = lambda: {
+        "remaining_seconds": 60.0
+    }
+    agent = _make_agent(comp)
+
+    ctx = _build(agent)
+
+    hook.assert_not_called()
+    agent._compress_context.assert_not_called()
+    assert ctx.preflight_compression_blocked is False
+
+
+def test_engine_hook_exception_is_swallowed():
+    """A buggy engine must not break an otherwise-healthy turn."""
+    hook = MagicMock(side_effect=RuntimeError("buggy engine"))
+    agent = _make_agent(_stub_compressor(preflight=hook))
+
+    ctx = _build(agent)
+
+    assert isinstance(ctx, TurnContext)
+    hook.assert_called_once()
+    agent._compress_context.assert_not_called()
+    assert ctx.preflight_compression_blocked is False
+
+
+def test_builtin_compressor_default_sub_threshold_path_unchanged(tmp_path):
+    """Byte-identical-default pin against the REAL ContextCompressor.
+
+    The built-in engine inherits ``should_compress_preflight() -> False``
+    from ``ContextEngine``, so wiring the hook must not change the default
+    sub-threshold behavior at all: no compress call, no status emission,
+    no blocking flag.
+    """
+    from agent.context_compressor import ContextCompressor
+    from agent.context_engine import ContextEngine
+
+    # The default really is False (the dead-code premise of #20316).
+    assert "should_compress_preflight" not in ContextCompressor.__dict__
+    assert ContextEngine.should_compress_preflight(
+        object.__new__(ContextCompressor), []
+    ) is False
+
+    with patch(
+        "agent.context_compressor.get_model_context_length", return_value=1_000_000
+    ):
+        compressor = ContextCompressor(
+            model="test/model",
+            threshold_percent=0.85,
+            protect_first_n=2,
+            protect_last_n=2,
+            quiet_mode=True,
+        )
+    agent = _make_agent(compressor)
+
+    ctx = _build(agent)
+
+    agent._compress_context.assert_not_called()
+    agent._emit_status.assert_not_called()
+    assert ctx.preflight_compression_blocked is False