Files
Jan-Stefan Janetzky 7fe434ee70 fix(memory): suppress the proactive tool-result prune under the checkpoint gate
The proactive prune is a lossy rewrite of the durable transcript — the
built-in commits it through archive_and_compact — and it runs BETWEEN
compactions, so no pre-compress checkpoint precedes it. With
compression.checkpoint_required armed, tool evidence could still be
dropped without ever reaching the durable provider: the same authority
class already closed for native compaction, codex app-server, and
post-turn micro-compaction.

- turn_preflight: the prune arm of compress_after_tool_results() moves
  into _proactive_tool_result_prune() and grows an `is True` gate on
  compression_checkpoint_required. Gating the call site rather than
  ContextCompressor's implementation covers plugin context engines that
  override prune_tool_results_only too — both reach the transcript
  through this one seam. The caller keeps the read-dedup reset on a
  committed prune; publish_prune moves with the commit.
- Suppression, not refusal (the prune fires per tool iteration; a raise
  would abort the turn). It is also lossless: the material the prune
  would have dropped stays in context, and the eventual threshold
  compaction — which IS checkpointed — handles it.
- conversation_compression_checkpoint: _record_checkpoint_gate_suppression
  bumps a suppression counter on the agent and logs once per withheld
  authority. The log names the availability trade the operator made:
  with the prune suppressed and compaction fail-closed, a checkpoint
  provider outage leaves no reclamation path at all and the session
  halts instead of compacting unarchived context. That is the gate's
  promise rather than a side effect, so it belongs in the line that
  fires when it becomes true.
- Count and log only a prune that would have run
  (ProactivePruneMixin.would_proactively_prune, _prune_would_have_run).
  A default install ships proactive_prune_tokens: 0, and a configured
  prune below its trigger returns the input untouched, so neither is a
  withheld reclamation path. An engine that overrides
  prune_tool_results_only owns its trigger policy, so occupying the hook
  is the precondition there; the inherited ABC no-op is not an authority.
- Tests (tests/agent/test_proactive_prune_checkpoint_gate.py): the call
  site swallows compressor exceptions with logger.debug, so the doubles
  count calls instead of asserting. Gate off is the sabotage control
  proving the harness reaches the prune; the engine-override case proves
  the seam covers both authorities; the log test pins one warning and its
  wording; five tests pin the default, the threshold (with a
  same-compressor control one token above it), the real compressor's
  predicate, the hookless and raising shapes, and the inherited-vs-
  overridden hook. test_proactive_prune_loop_wiring drives a real
  run_conversation with the gate armed: the prune is never consulted.

(cherry picked from commit 6904b46b55ac125f0b816f705e23f4c4a9c69991)
(cherry picked from commit 9e505a834742df11045da3b809a3a49c6178354c)
2026-11-05 19:55:33 -06:00

419 lines
23 KiB
Python

"""Pre-API preflight compression gate and post-tool compression for the turn loop.
Runs once per API call after the request pressure is measured: grow a managed llama.cpp
window (last resort), or compress when over threshold (deferring on noisy estimates, in
failure cooldown, or while the review fork's first request is pending), handle the
provider-proven overflow re-check fail-closed, and emit the deduped blocked/uncompressed
overflow warnings. Nothing here imports ``agent.conversation_loop`` at module level
(cycle); loop-internal helpers resolve lazily."""
from __future__ import annotations
import logging
from dataclasses import dataclass
from typing import Any, Dict, List, Optional
from agent.context_engine import automatic_compaction_status_message
from agent.conversation_compression import (
PRE_API_COMPRESSION_STATUS_TEMPLATE, _reset_read_dedup_caches, compression_blocked_transiently,
compression_skipped_due_to_lock, context_compression_timed_out, conversation_history_after_compression,
)
from agent.conversation_compression_feasibility import ensure_compression_feasibility_checked
from agent.turn_context import _review_fork_first_request_pending
from agent.turn_context_compaction import (
_apply_grown_window, _blocked_compress_reason, _clear_overflow_warn, _refund_api_call,
_reanchor, _reset_retry_state_after_compaction,
)
logger = logging.getLogger("agent.conversation_loop")
@dataclass
class PreflightGateVerdict:
"""``action``: ``"fallthrough"`` (make the API call), ``"continue"`` (window grown or
history compacted — the call/budget was refunded, re-enter the turn loop and
re-measure), ``"break"`` (turn ends: compression timeout or non-actionable compaction
handoff — ``final_response``/``failed``/``_turn_exit_reason`` set) or ``"return"``
(``result`` is the turn's result dict). The other fields are the loop locals the gate
rebinds; ``_last_preflight_pressure`` is re-armed only by a compression pass."""
action: str
pending_moa_prepared_request: Any
messages: Any
active_system_prompt: Any
conversation_history: Any
api_call_count: Any
compression_attempts: Any
final_response: Any
failed: Any
_turn_exit_reason: Any
_compression_timeout_exhausted: Any
_preflight_compression_blocked: Any
_provider_overflow_recovery_pending: Any
_last_preflight_pressure: Any
result: Optional[dict[str, Any]] = None
def run_preflight_compression(
agent: Any, v: PreflightGateVerdict, *, compressor: Any, request_pressure_tokens: int,
provider_overflow_preflight: bool, defer_preflight: Any, moa_prepared_request: Any,
system_message: Any, user_message: Any, max_compression_attempts: int, effective_task_id: Any,
) -> PreflightGateVerdict:
"""Mirror of the turn-prologue guard chain (defer on noisy estimate → skip in failure
cooldown → ``should_compress``), rebinding the loop locals on ``v`` and setting
``v.action``. A compression pass that never reaches the provider refunds the
call/budget in every branch (skip, re-run, timeout) so ``api_call_count`` never
over-reports; a lock/transient skip refunds the attempt and leaves the progress
blocker unarmed. A forced provider-overflow preflight that any gate blocks fails
closed (llama.cpp may silently truncate)."""
from agent.conversation_loop import (
_COMPRESSION_TIMEOUT_FINAL_RESPONSE, _HANDOFF_SKIP_FINAL_RESPONSE,
_compression_deferred_result, _maybe_grow_local_window, _provider_overflow_exhausted_result,
_should_skip_model_call_for_reference_handoff,
)
def _done(action: str, result: Optional[dict[str, Any]] = None) -> PreflightGateVerdict:
v.action, v.result = action, result
return v
def _exhausted_result() -> dict[str, Any]:
return _provider_overflow_exhausted_result(
agent, v.messages, v.conversation_history, v.api_call_count, request_pressure_tokens,
max_compression_attempts,
)
_eligible = (
agent.compression_enabled
and len(v.messages) > 1
and v.compression_attempts < max_compression_attempts
)
if _eligible:
# Aux clamp must land before the first compaction fires on the main-window threshold (#114707).
# Before the cooldown read: a moved compression route clears the previous summariser's cooldown.
ensure_compression_feasibility_checked(agent, request_pressure_tokens)
_compression_cooldown = getattr(
compressor, "get_active_compression_failure_cooldown", lambda: None
)()
if (
_eligible
and not _review_fork_first_request_pending(agent)
and (not v._preflight_compression_blocked or provider_overflow_preflight)
and (not defer_preflight(request_pressure_tokens) or provider_overflow_preflight)
and not _compression_cooldown
and compressor.should_compress(request_pressure_tokens)
):
# Managed local runtime: grow the context window before compressing (last
# resort). Only for a llamacpp provider at the supervised base_url.
_grown_window = _maybe_grow_local_window(agent, compressor, request_pressure_tokens)
if _grown_window:
# Bigger window granted: recalibrate and skip compression this pass. Never
# reached the provider — refund like the compression path does.
_apply_grown_window(agent, compressor, _grown_window)
v.api_call_count = _refund_api_call(agent, v.api_call_count)
return _done("continue")
if moa_prepared_request is not None:
v.pending_moa_prepared_request = moa_prepared_request
v.compression_attempts += 1
# Compression is running: reset the blocked-overflow warning dedup so a
# later blocked turn warns again.
_clear_overflow_warn(agent)
_threshold = int(getattr(compressor, "threshold_tokens", 0) or 0)
_context_length = int(getattr(compressor, "context_length", 0) or 0)
logger.info(
"Pre-API compression: ~%s request tokens >= %s threshold "
"(context=%s, attempt=%s/%s)",
f"{request_pressure_tokens:,}",
f"{_threshold:,}",
f"{_context_length:,}" if getattr(compressor, "context_length", 0) else "unknown",
v.compression_attempts,
max_compression_attempts,
)
_pre_api_status = automatic_compaction_status_message(
compressor,
phase="pre_api",
default_message=PRE_API_COMPRESSION_STATUS_TEMPLATE.format(
tokens=request_pressure_tokens
),
approx_tokens=request_pressure_tokens,
threshold_tokens=_threshold,
context_length=_context_length,
model=agent.model,
attempt=v.compression_attempts,
max_attempts=max_compression_attempts,
)
if _pre_api_status:
agent._emit_status(_pre_api_status)
v._last_preflight_pressure = request_pressure_tokens
_pre_api_input = v.messages
v.messages, v.active_system_prompt = agent._compress_context(
v.messages, system_message, approx_tokens=request_pressure_tokens,
task_id=effective_task_id, trigger="pre_api",
)
if context_compression_timed_out(agent):
# Progress-aware timeout: never reached the provider — refund the
# call/budget and stop; an overflow retry would only re-compress.
v.api_call_count = _refund_api_call(agent, v.api_call_count)
v.final_response = _COMPRESSION_TIMEOUT_FINAL_RESPONSE
v.failed = True
v._compression_timeout_exhausted = True
v._turn_exit_reason = "context_compression_timeout"
return _done("break")
if v.messages is _pre_api_input and (
compression_skipped_due_to_lock(agent) or compression_blocked_transiently(agent)
):
# Temporary DEFER (lock held / cooldown), not evidence about compressibility:
# refund the attempt, leave the progress blocker unarmed and proceed.
v.compression_attempts -= 1
v._last_preflight_pressure = None
if v.pending_moa_prepared_request is moa_prepared_request:
v.pending_moa_prepared_request = None
else:
_reset_retry_state_after_compaction(agent)
# Re-baseline the flush cursor: rotation returns None (child flushes
# whole); in-place returns list(messages) — None would re-append
# persisted rows. See conversation_history_after_compression().
v.conversation_history = conversation_history_after_compression(
agent, v.messages, v.conversation_history
)
# Never reaches the provider on skip or re-run — refund the call/budget
# in BOTH cases, else budget leaks and api_call_count over-reports.
v.api_call_count = _refund_api_call(agent, v.api_call_count)
if _should_skip_model_call_for_reference_handoff(v.messages, user_message):
# Reference-only handoff must not become the active turn after a
# completed assistant response.
logger.info(
"Skipping post-compaction model call: reference-only "
"handoff would be the sole active user turn (#80622)"
)
if not v.final_response:
v.final_response = _HANDOFF_SKIP_FINAL_RESPONSE
v._turn_exit_reason = "compaction_handoff_not_actionable"
return _done("break")
return _done("continue")
elif provider_overflow_preflight and _compression_cooldown:
# Provider proved the request cannot fit and the compressor is unavailable:
# don't resend; let the next user turn retry after cooldown.
agent._persist_session(v.messages, v.conversation_history)
return _done("return", _compression_deferred_result(
agent, v.messages, v.api_call_count, reason="transient_block"
))
elif provider_overflow_preflight and v.compression_attempts >= max_compression_attempts:
# All recovery passes consumed and still over threshold: fail closed —
# llama.cpp may silently truncate an oversized retry.
return _done("return", _exhausted_result())
elif _eligible and not defer_preflight(request_pressure_tokens) and _compression_cooldown:
# Summary-LLM cooldown blocks compression: deduped warning only when over
# threshold (should_compress_info reason is None below it).
_block_reason = _blocked_compress_reason(compressor, request_pressure_tokens)
if _block_reason:
agent._warn_context_overflow_blocked(
_block_reason, request_pressure_tokens,
int(getattr(compressor, "threshold_tokens", 0) or 0),
)
elif not agent.compression_enabled and len(v.messages) > 1:
# Uncompressed session guard: compression is disabled, so warn (deduped) when
# the request exceeds the context window; the turn-context preflight re-arms.
_ctx_len = getattr(getattr(agent, "context_compressor", None), "context_length", None)
if isinstance(_ctx_len, int) and _ctx_len > 0 and request_pressure_tokens > _ctx_len:
_warn_fn = getattr(agent, "_warn_uncompressed_context_overflow", None)
if callable(_warn_fn):
_warn_fn(request_pressure_tokens, _ctx_len)
if provider_overflow_preflight:
# Any other gate blocking the forced preflight (e.g. uncompressible one-
# message request) must fail closed: the request is proven not to fit.
return _done("return", _exhausted_result())
return _done("fallthrough")
@dataclass
class PostToolCompressionVerdict:
"""``end_turn`` True → a reference-only compaction handoff would be the sole active
user turn: stop without another model call (``final_response`` /
``turn_exit_reason`` set)."""
end_turn: bool
messages: list[dict[str, Any]]
active_system_prompt: Any
conversation_history: Any
compression_attempts: int
final_response: Any
turn_exit_reason: Any
current_turn_user_idx: int
def compress_after_tool_results(
agent: Any, *, messages: list[dict[str, Any]], system_message: Any, user_message: Any,
active_system_prompt: Any, conversation_history: Any, compression_attempts: int,
max_compression_attempts: int, effective_task_id: Any, final_response: Any,
turn_exit_reason: Any, current_turn_user_idx: int,
) -> PostToolCompressionVerdict:
"""Post-tool-call compression decision. Pressure comes from API-reported
``prompt_tokens`` (a tight lower bound; thinking models inflate completion tokens),
``0`` right after compression (no real count yet), else the route-aware
overhead-inclusive estimate. Over threshold but blocked → deduped warning plus the
deterministic tool-result-only prune, committed only when the engine returns a NEW
list (never rebuild ``conversation_history`` for it)."""
from agent.conversation_loop import (
_HANDOFF_SKIP_FINAL_RESPONSE, _midturn_request_pressure_tokens,
_should_skip_model_call_for_reference_handoff,
)
from agent.model_metadata import estimate_request_tokens_rough
def _verdict(end_turn: bool) -> PostToolCompressionVerdict:
return PostToolCompressionVerdict(
end_turn=end_turn, messages=messages, active_system_prompt=active_system_prompt,
conversation_history=conversation_history, compression_attempts=compression_attempts,
final_response=final_response, turn_exit_reason=turn_exit_reason,
current_turn_user_idx=current_turn_user_idx,
)
_compressor = agent.context_compressor
# A new checkpoint must reach the provider before stale usage can trigger
# local compression, overflow warnings, or destructive tool-result pruning.
if bool(getattr(_compressor, "awaiting_real_usage_after_compression", False)):
return _verdict(False)
# Real usage decides: the anchor is the provider's last prompt count plus a rough delta for
# ONLY the tool results appended since (the raw last_prompt_tokens ignores them). Right after
# a compaction (-1 sentinel) there is no real count yet: never treat the schema-heavy rough
# figure as pressure. The whole-request rough estimate is the last resort (usage-less
# provider, post-disconnect, gateway restart), kept route-aware (#96995/#97602).
from agent.usage_anchor import anchored_context_tokens
_anchored = anchored_context_tokens(messages, getattr(agent, "_usage_anchor", None))
if _anchored is not None:
_real_tokens = _anchored
elif _compressor.last_prompt_tokens > 0:
# Only prompt_tokens: thinking models inflate completion_tokens with
# reasoning that uses no context → premature compression. (#12026)
_real_tokens = _compressor.last_prompt_tokens
elif _compressor.last_prompt_tokens == -1:
_real_tokens = 0
else:
_real_tokens = _midturn_request_pressure_tokens(
agent, messages, active_system_prompt or "",
estimate_request_tokens_rough(messages, tools=agent.tools or None),
)
if agent.compression_enabled and compression_attempts < max_compression_attempts:
ensure_compression_feasibility_checked(agent, _real_tokens)
if (
agent.compression_enabled
and compression_attempts < max_compression_attempts
and not bool(
getattr(_compressor, "awaiting_real_usage_after_compression", False)
)
and _compressor.should_compress(_real_tokens)
):
compression_attempts += 1
# Compression is running: reset blocked-overflow warning dedup so a
# future blocked turn can warn again.
_clear_overflow_warn(agent)
agent._safe_print(" ⟳ compacting context…")
_post_tool_input = messages
# Pass overhead-aware _real_tokens, not last_prompt_tokens (0 in the
# no-usage fallback), so the overflow guard sees the true size.
messages, active_system_prompt = agent._compress_context(
messages, system_message, approx_tokens=_real_tokens, task_id=effective_task_id, trigger="post_tool",
)
if messages is _post_tool_input and compression_skipped_due_to_lock(agent):
# Lock-skip no-op is a temporary defer, not evidence about compressibility:
# refund so a lock-loser loop doesn't burn the budget toward exhausted.
# #69870 lock-skip / #97488 transient-block: this pass no-oped for a TEMPORARY reason (another
# path holds the compression lock, or a timed cooldown/backoff guard is active). That is a
# temporary DEFER, not evidence about compressibility — refund the attempt (it must not burn the
# shared overflow-recovery budget toward compression_exhausted → gateway auto-reset,
# #9893/#35809) and leave the insufficient-progress blocker unarmed. Proceed with the current
# request: if it truly does not fit, the provider's 413/overflow handler returns the soft
# compression_deferred result with that stronger signal.
# #69870 lock-skip: the provider proved the request does not fit, but this compression pass
# no-oped only because another path holds the session's compression lock. Temporary defer, not
# exhaustion — refund the attempt and end the turn softly so the gateway does NOT auto-reset the
# session (#9893/#35809).
# #97488 transient-block: compression no-oped because a timed guard (host-timeout cooldown /
# structural backoff) is active — a temporary defer, not evidence of incompressibility. Never
# classify it as compression_exhausted (gateway auto-reset).
# bypass_cooldown=True, # #100661 provider-proven overflow
# #97488: timed transient guard — defer, never exhaustion (gateway auto-reset).
# #69870 lock-skip: the provider proved the request does not fit, but this compression pass
# no-oped only because another path holds the session's compression lock. Temporary defer, not
# exhaustion — refund the attempt and end the turn softly so the gateway does NOT auto-reset the
# session (#9893/#35809).
# #97488 transient-block: a timed guard (host-timeout cooldown / structural backoff) no-oped
# this pass — defer softly, never compression_exhausted (which would auto-reset the session).
# #69870 lock-skip: this pass no-oped because another path holds the session's compression lock
# — a temporary defer, not evidence about compressibility.
compression_attempts -= 1
else:
conversation_history = conversation_history_after_compression(
agent, messages, conversation_history
)
if _should_skip_model_call_for_reference_handoff(messages, user_message):
logger.info(
"Skipping post-tool compaction model call: "
"reference-only handoff would be the sole "
"active user turn (#80622)"
)
if not final_response:
final_response = _HANDOFF_SKIP_FINAL_RESPONSE
turn_exit_reason = "compaction_handoff_not_actionable"
return _verdict(True)
current_turn_user_idx = _reanchor(agent, messages, user_message)
elif agent.compression_enabled:
# Over threshold but compression blocked (cooldown/anti-thrash): deduped
# warning so context can't silently overflow. ``attempts_spent`` names the
# attempts_exhausted lockout when the engine says RUN but the per-turn
# budget is spent (#101889).
_block_reason = _blocked_compress_reason(
_compressor, _real_tokens, attempts_spent=compression_attempts
)
if _block_reason:
agent._warn_context_overflow_blocked(
_block_reason, _real_tokens, int(getattr(_compressor, "threshold_tokens", 0) or 0)
)
_pruned_msgs = _proactive_tool_result_prune(agent, _compressor, messages, _real_tokens)
if _pruned_msgs is not messages:
messages = _pruned_msgs
# A committed prune is a content-loss boundary like compaction: demoted skill_view /
# read_file bodies survive only as one-line markers, so the repeat-read dedup must
# stop answering "unchanged" for them or the reload the marker asks for is refused.
_reset_read_dedup_caches(effective_task_id, session_id=agent.session_id or "")
return _verdict(False)
def _proactive_tool_result_prune(agent: Any, compressor: Any, messages: list, current_tokens: int) -> list:
"""Proactive tool-result prune (deterministic, no LLM, keeps tail): the pruned list when one
committed, else the same ``messages`` object. No-op unless proactive_prune_tokens is exceeded;
commits only past proactive_prune_min_reclaim_tokens so cache breaks stay episodic.
Withheld while compression.checkpoint_required is armed: the prune rewrites the durable
transcript between compactions, so no checkpoint precedes it. Gating here, not in the
compressor, also covers engines that override the hook. Suppression, not refusal (it runs per
tool iteration); the material stays in context for the checkpointed threshold compaction.
"""
if getattr(agent, "compression_checkpoint_required", False) is True:
from agent.conversation_compression_checkpoint import (
_prune_would_have_run, _record_checkpoint_gate_suppression,
)
if _prune_would_have_run(compressor, current_tokens):
_record_checkpoint_gate_suppression(agent, "the proactive tool-result prune")
return messages
prune = getattr(compressor, "prune_tool_results_only", None)
if not callable(prune):
return messages
try:
pruned_msgs, pruned_n = prune(messages, current_tokens=current_tokens)
except Exception:
logger.debug("proactive tool-result prune failed; skipping", exc_info=True)
return messages
# Standard no-op caller contract: only commit when the engine returned a NEW list object
# with a non-zero count. Do NOT rebuild conversation_history: rows already carry
# _DB_PERSISTED_MARKER, and on a stale in-place flag the helper could seed unpersisted rows.
if not pruned_n or pruned_msgs is messages:
return messages
from agent.compaction_events import publish_prune
publish_prune(agent, messages, pruned_msgs, pruned_n)
return pruned_msgs