feat(compression): proactive tool-result pruning for large-window models

The phase-1 tool-result prune only runs inside compress(), which fires
near 50% of the context window, so it never triggers on large-window
models; old tool outputs then ride in history and are re-sent every turn.

Add prune_tool_results_only(): the same no-LLM prune on a separate, low
proactive_prune_tokens trigger, run as an elif to the compression branch.
Opt-in (default 0), protects the recent tail by message count.

Add the method to the ContextEngine base as a no-op default so pluggable
engines inherit it safely (the post-tool-call path never AttributeErrors on
a non-built-in engine); the built-in compressor supplies the real prune.
Register both keys under the top-level compression config with defaults and
document them.
This commit is contained in:
Kolektori 2026-07-11 16:37:39 +04:00 committed by Teknium
parent 66fdcfa3bd
commit cb481e2f2b
8 changed files with 407 additions and 1 deletions

View file

@ -1838,6 +1838,39 @@ def init_agent(
if compression_max_attempts < 1:
compression_max_attempts = 3
compression_max_attempts = min(compression_max_attempts, 10)
def _parse_prune_int(raw, default):
# Same parser semantics as compression.max_attempts above: reject
# booleans (bool subclasses int — YAML `true` would coerce to 1),
# reject fractional floats rather than truncating them, accept
# integral floats and numeric strings, fall back to the default on
# anything else.
if isinstance(raw, bool):
return default
if isinstance(raw, int):
return raw
if isinstance(raw, float):
return int(raw) if raw.is_integer() else default
try:
return int(str(raw).strip())
except (TypeError, ValueError):
return default
# Opt-in proactive tool-result prune trigger (0 = disabled — the
# default, so an unset key is behavior-neutral). Negative values are
# treated as disabled rather than erroring.
compression_proactive_prune_tokens = max(
0, _parse_prune_int(_compression_cfg.get("proactive_prune_tokens", 0), 0)
)
compression_proactive_prune_min_chars = _parse_prune_int(
_compression_cfg.get("proactive_prune_min_result_chars", 8000), 8000
)
compression_proactive_prune_min_reclaim = max(
0,
_parse_prune_int(
_compression_cfg.get("proactive_prune_min_reclaim_tokens", 4096), 4096
),
)
# protect_first_n is the number of non-system messages to protect at
# the head, in addition to the system prompt (which is always
# implicitly protected by the compressor). Floor at 0 — a value of
@ -2312,6 +2345,9 @@ def init_agent(
max_tokens=agent.max_tokens,
model_thresholds=compression_model_thresholds,
threshold_tokens_cap=compression_threshold_tokens,
proactive_prune_tokens=compression_proactive_prune_tokens,
proactive_prune_min_result_chars=compression_proactive_prune_min_chars,
proactive_prune_min_reclaim_tokens=compression_proactive_prune_min_reclaim,
)
_bind_session_state = getattr(agent.context_compressor, "bind_session_state", None)
if callable(_bind_session_state):

View file

@ -1629,6 +1629,9 @@ class ContextCompressor(ContextEngine):
max_tokens: int | None = None,
model_thresholds: dict[str, float] | None = None,
threshold_tokens_cap: Any = None,
proactive_prune_tokens: int = 0,
proactive_prune_min_result_chars: int = 8000,
proactive_prune_min_reclaim_tokens: int = 4096,
):
self.model = model
self.base_url = base_url
@ -1658,6 +1661,31 @@ class ContextCompressor(ContextEngine):
)
self.protect_first_n = protect_first_n
self.protect_last_n = protect_last_n
# Proactive tool-result pruning (cost-oriented; runs INDEPENDENTLY of the
# full-compression trigger, via prune_tool_results_only()). 0 = disabled.
self.proactive_prune_tokens = int(proactive_prune_tokens or 0)
# Floor the summarize threshold at 200 chars (matching
# _prune_old_tool_results' dedup floor). Below ~200 a generated summary
# can be longer than the floor it replaces, so Pass 2 would re-summarize
# its own output every turn (corrupting it and never converging); a
# negative value would strip every non-tail tool result outright. A
# configured 0 keeps the 8000 default via `or`. Keep the floor well above
# typical summary length (default 8000) to stay idempotent.
self.proactive_prune_min_result_chars = max(
200, int(proactive_prune_min_result_chars or 8000)
)
# Minimum estimated token reclaim before a proactive prune COMMITS.
# Every commit rewrites messages the provider has already seen, which
# invalidates the prompt-cache prefix from the earliest rewritten
# message forward. Without this gate a busy tool loop would re-fire
# the prune nearly every iteration (each new tool pair ages an old one
# out of the protected tail), breaking the cache per turn. Requiring a
# meaningful batch of reclaimable tokens makes fires episodic and
# amortized — the same way full compression is the one sanctioned
# cache break. 0 disables the gate (commit any non-zero prune).
self.proactive_prune_min_reclaim_tokens = max(
0, int(proactive_prune_min_reclaim_tokens or 0)
)
self.summary_target_ratio = max(0.10, min(summary_target_ratio, 0.80))
self.quiet_mode = quiet_mode
# Output-token reservation: the provider carves max_tokens out of the
@ -2059,6 +2087,7 @@ class ContextCompressor(ContextEngine):
def _prune_old_tool_results(
self, messages: List[Dict[str, Any]], protect_tail_count: int,
protect_tail_tokens: int | None = None,
min_prune_chars: int = 200,
) -> tuple[List[Dict[str, Any]], int]:
"""Replace old tool result contents with informative 1-line summaries.
@ -2201,7 +2230,9 @@ class ContextCompressor(ContextEngine):
return False
if content.startswith("[screenshot removed"):
return False
if len(content) <= 200:
# Only prune if the content is substantial (default >200 chars; the
# proactive path raises this floor via min_prune_chars).
if len(content) <= min_prune_chars:
return False
call_id = msg.get("tool_call_id", "")
tool_name, tool_args = call_id_to_tool.get(call_id, ("unknown", ""))
@ -2317,6 +2348,75 @@ class ContextCompressor(ContextEngine):
return result, pruned
def prune_tool_results_only(
self, messages: List[Dict[str, Any]], current_tokens: int | None = None,
) -> tuple[List[Dict[str, Any]], int]:
"""Deterministic, no-LLM tool-result prune for the cost-oriented path.
Runs the Phase-1 prune (``_prune_old_tool_results``) WITHOUT the
compression summary phase, gated on ``proactive_prune_tokens`` rather
than the (much higher) full-compression threshold. On large-window
models ``should_compress()`` (50% of the window) rarely fires, so old
tool outputs otherwise ride in history and are re-sent verbatim on every
subsequent turn; this reclaims them early with no quality-risky LLM
summarization.
Protects the recent tail by message COUNT (``protect_last_n``), never by
``tail_token_budget`` the latter is derived from the 50% compression
threshold (100K tokens on a 1M window) and would protect the entire
session, pruning nothing.
``_prune_old_tool_results`` runs all three deterministic passes:
(1) dedup byte-identical tool results keeps the newest full copy and
back-references older exact duplicates ANYWHERE in the list (including
the protected tail), so no unique content is ever lost; (2) summarize
non-tail tool results larger than ``min_prune_chars``; (3) truncate
oversized tool_call arguments on non-tail assistant messages. Only
pass (2)'s floor is raised by ``proactive_prune_min_result_chars``;
passes (1) and (3) keep their own fixed floors. The recent-tail
protection applies to passes (2) and (3); pass (1) is tail-agnostic by
design because dedup is lossless.
PROMPT-CACHE CONTRACT: a committed prune rewrites message bodies the
provider has already seen, invalidating the cached prefix from the
earliest rewritten message forward exactly like a compression
boundary. To keep that break episodic rather than per-turn, the prune
only COMMITS when the estimated reclaim meets
``proactive_prune_min_reclaim_tokens`` (measured on the actual pruned
output, not guessed up front). Below the gate the INPUT list object is
returned unchanged the standard no-op caller contract (callers gate
bookkeeping on ``result is not input``).
Returns ``(messages, 0)`` the input object when disabled, below
the trigger, or when the reclaim gate rejects the commit.
"""
if self.proactive_prune_tokens <= 0:
return messages, 0
if current_tokens is not None and current_tokens < self.proactive_prune_tokens:
return messages, 0
# Nothing to reclaim until there are messages outside the protected tail.
if len(messages) <= self.protect_last_n + self._protect_head_size(messages) + 1:
return messages, 0
pruned_msgs, pruned_count = self._prune_old_tool_results(
messages,
protect_tail_count=self.protect_last_n,
protect_tail_tokens=None,
min_prune_chars=self.proactive_prune_min_result_chars,
)
if not pruned_count:
# Standard no-op contract: hand back the INPUT object so callers
# can gate bookkeeping on `result is not input`.
return messages, 0
# Measured-savings gate (prompt-cache hysteresis): only commit when
# the prune reclaims a meaningful batch of tokens. Estimated on the
# real before/after messages so dedup + arg truncation count too.
if self.proactive_prune_min_reclaim_tokens > 0:
before = sum(_estimate_msg_budget_tokens(m) for m in messages)
after = sum(_estimate_msg_budget_tokens(m) for m in pruned_msgs)
if (before - after) < self.proactive_prune_min_reclaim_tokens:
return messages, 0
return pruned_msgs, pruned_count
# ------------------------------------------------------------------
# Summarization
# ------------------------------------------------------------------

View file

@ -189,6 +189,27 @@ class ContextEngine(ABC):
host filters unsupported optional arguments by signature.
"""
# -- Optional: proactive tool-result prune -----------------------------
def prune_tool_results_only(
self,
messages: List[Dict[str, Any]],
current_tokens: int | None = None,
) -> tuple[List[Dict[str, Any]], int]:
"""Deterministically trim old tool-result payloads without an LLM call.
Runs on a low, cost-oriented trigger independent of ``should_compress``
so large-window engines can reclaim re-sent tool output long before full
compaction would fire. Returns ``(messages, n_pruned)``.
Default is a safe no-op: the list is returned unchanged with ``0``
pruned. Engines that don't implement a cheap prune — and any engine that
predates this hook inherit this default, so the agent loop's
post-tool-call prune path never raises ``AttributeError`` on them. The
built-in ContextCompressor overrides this with the real implementation.
"""
return messages, 0
# -- Optional: pre-flight check ----------------------------------------
def should_compress_preflight(self, messages: List[Dict[str, Any]]) -> bool:

View file

@ -5622,6 +5622,48 @@ def run_conversation(
_real_tokens,
int(getattr(_compressor, "threshold_tokens", 0) or 0),
)
# Proactive tool-result prune: reclaim re-sent history on
# large-window models long before should_compress() (≈50% of
# the window) would ever fire. Deterministic, no LLM call;
# protects the recent tail. No-op unless proactive_prune_tokens
# is configured and _real_tokens is above it — and even then
# the prune only commits when it reclaims at least
# proactive_prune_min_reclaim_tokens, so prompt-cache breaks
# stay episodic like compression's (the one sanctioned cache
# break) instead of firing every tool iteration. See
# ContextCompressor.prune_tool_results_only.
# getattr guard: plugin context engines predating the hook and
# minimal test doubles (SimpleNamespace compressors) lack the
# method — treat absence as a no-op.
_prune = getattr(_compressor, "prune_tool_results_only", None)
if callable(_prune):
try:
_pruned_msgs, _pruned_n = _prune(
messages, current_tokens=_real_tokens
)
except Exception:
logger.debug(
"proactive tool-result prune failed; skipping",
exc_info=True,
)
_pruned_msgs, _pruned_n = messages, 0
# Standard no-op caller contract: only commit when the
# engine returned a NEW list object with a non-zero count.
if _pruned_n and _pruned_msgs is not messages:
# Do NOT rebuild conversation_history here. Unlike the
# compression branch, the prune neither rotates the session
# nor calls archive_and_compact(), so there is no new
# persistence baseline to establish. _prune_old_tool_results
# returns per-message copies that preserve the
# _DB_PERSISTED_MARKER, so the marker-based flush dedup (see
# _flush_messages_to_session_db) already prevents both
# duplicate writes and dropped rows. Calling
# conversation_history_after_compression (a compaction-only
# helper keyed on the _last_compaction_in_place flag) would be
# a no-op at best, and on a stale in-place flag could seed
# this turn's fresh, not-yet-persisted rows into history_ids
# and skip writing them.
messages = _pruned_msgs
# Save session log incrementally (so progress is visible even if interrupted)
agent._session_messages = messages

View file

@ -1414,6 +1414,29 @@ DEFAULT_CONFIG = {
# (e.g. 6) for tool-schema-heavy sessions where 3
# rounds cannot clear the request estimate.
# Validated >= 1, hard-capped at 10.
"proactive_prune_tokens": 0, # opt-in trigger (tokens) for the deterministic,
# no-LLM tool-result prune, run independently of
# `threshold` above. On large-window models
# `threshold` (≈50% of the window) rarely fires,
# so old tool output otherwise rides in history
# and is re-sent every turn; a low value like
# 48000 reclaims it early. 0 = off. Recent tail
# protected by `protect_last_n`. Built-in
# compressor only (other engines inherit a no-op).
# NOTE: each committed prune rewrites already-sent
# history, breaking the provider prompt-cache
# prefix — the min_reclaim gate below keeps those
# breaks episodic rather than per-turn.
"proactive_prune_min_result_chars": 8000, # the prune's summarize pass only
# touches tool results larger than this (chars);
# clamped to >= 200 so a generated summary can't
# itself be re-summarized.
"proactive_prune_min_reclaim_tokens": 4096, # a proactive prune only commits
# when it reclaims at least this many tokens
# (measured on the pruned output). Keeps
# prompt-cache invalidation amortized: one big
# episodic break instead of a tiny break every
# tool iteration. 0 = commit any non-zero prune.
"hygiene_hard_message_limit": 5000, # gateway session-hygiene force-compress threshold by message count
"hygiene_timeout_seconds": 30, # max seconds gateway waits for pre-agent hygiene compression
"hygiene_failure_cooldown_seconds": 300, # skip repeated failed hygiene attempts for this session

View file

@ -181,6 +181,21 @@ class TestStubEngine:
assert engine.last_prompt_tokens == 1000
assert engine.last_completion_tokens == 200
def test_prune_tool_results_only_defaults_to_safe_noop(self):
# An engine implementing only the required interface (no prune override)
# must inherit the base no-op instead of raising AttributeError: the
# agent loop calls prune_tool_results_only() on the active engine after a
# tool call whenever full compression does not fire, so every pluggable
# ContextEngine reaches this path (see conversation_loop proactive-prune).
engine = StubEngine()
msgs = [
{"role": "user", "content": "hello"},
{"role": "assistant", "content": "hi"},
]
result, pruned = engine.prune_tool_results_only(msgs, current_tokens=10_000_000)
assert pruned == 0
assert result is msgs
# ---------------------------------------------------------------------------
# ContextCompressor session reset via ABC

View file

@ -0,0 +1,164 @@
"""Tests for proactive tool-result pruning.
``ContextCompressor.prune_tool_results_only`` runs the cheap, deterministic
Phase-1 prune (summarize old tool outputs, dedup repeats) on a cost-oriented
trigger that is INDEPENDENT of the full-compression threshold. On large-window
models ``should_compress()`` (~50% of the window) rarely fires, so without this
the old tool outputs ride in history and are re-sent verbatim every turn.
Mirrors the construction/patching conventions in test_context_compressor.py.
"""
from unittest.mock import patch
from agent.context_compressor import ContextCompressor, _PRUNED_TOOL_PLACEHOLDER
LARGE_WINDOW = 1_000_000
def _compressor(**kw):
defaults = dict(
model="test",
quiet_mode=True,
threshold_percent=0.50,
protect_first_n=2,
protect_last_n=4,
)
defaults.update(kw)
with patch(
"agent.context_compressor.get_model_context_length",
return_value=LARGE_WINDOW,
):
return ContextCompressor(**defaults)
def _assistant_call(cid, name="terminal", args='{"cmd":"ls"}'):
return {
"role": "assistant",
"content": "",
"tool_calls": [
{"id": cid, "type": "function",
"function": {"name": name, "arguments": args}}
],
}
def _tool_msg(cid, content):
return {"role": "tool", "tool_call_id": cid, "content": content}
def _build(n_pairs, big_indices, big_chars=9000, small="ok"):
"""system + n_pairs of (assistant tool_call, tool result).
Tool results whose pair index is in ``big_indices`` get a distinct payload
of ``big_chars`` characters; the rest get a tiny payload.
"""
msgs = [{"role": "system", "content": "sys"}]
for i in range(n_pairs):
cid = f"call_{i}"
msgs.append(_assistant_call(cid))
if i in big_indices:
msgs.append(_tool_msg(cid, chr(65 + (i % 26)) * big_chars))
else:
msgs.append(_tool_msg(cid, small))
return msgs
def _tool_by_id(msgs, cid):
return [m for m in msgs if m.get("role") == "tool" and m.get("tool_call_id") == cid][0]
def test_prunes_below_compression_threshold():
"""The whole point: prune fires at 120k tokens, far below the ~500k
(50% of 1M) full-compression trigger that would otherwise never run."""
c = _compressor(proactive_prune_tokens=48_000, proactive_prune_min_result_chars=8_000)
assert c.should_compress(prompt_tokens=120_000) is False # compression would NOT run
msgs = _build(8, big_indices={0, 1, 2})
result, pruned = c.prune_tool_results_only(msgs, current_tokens=120_000)
assert pruned >= 3
assert len(result) == len(msgs)
for cid in ("call_0", "call_1", "call_2"):
m = _tool_by_id(result, cid)
assert len(m["content"]) < 9000 # summarized
assert m["content"] != _PRUNED_TOOL_PLACEHOLDER # informative, not a blank placeholder
def test_disabled_by_default_is_noop():
c = _compressor() # proactive_prune_tokens defaults to 0
assert c.proactive_prune_tokens == 0
msgs = _build(8, big_indices={0, 1, 2})
result, pruned = c.prune_tool_results_only(msgs, current_tokens=500_000)
assert pruned == 0
assert [m.get("content") for m in result] == [m.get("content") for m in msgs]
def test_below_trigger_is_noop():
c = _compressor(proactive_prune_tokens=48_000)
msgs = _build(8, big_indices={0, 1, 2})
result, pruned = c.prune_tool_results_only(msgs, current_tokens=10_000)
assert pruned == 0
def test_recent_tail_is_protected():
c = _compressor(proactive_prune_tokens=48_000, proactive_prune_min_result_chars=8_000)
# pair 0 tool is old (index 2); pair 7 tool is in the last-4 protected tail (index 16)
msgs = _build(8, big_indices={0, 7})
result, pruned = c.prune_tool_results_only(msgs, current_tokens=120_000)
assert len(_tool_by_id(result, "call_7")["content"]) == 9000 # protected, untouched
assert len(_tool_by_id(result, "call_0")["content"]) < 9000 # old, summarized
def test_size_floor_spares_small_results():
c = _compressor(proactive_prune_tokens=48_000, proactive_prune_min_result_chars=8_000)
msgs = _build(8, big_indices={1}, big_chars=9000)
for m in msgs: # make pair 0's tool 5000 chars (< 8000 floor), still old
if m.get("tool_call_id") == "call_0":
m["content"] = "Z" * 5000
result, pruned = c.prune_tool_results_only(msgs, current_tokens=120_000)
assert len(_tool_by_id(result, "call_0")["content"]) == 5000 # under floor -> untouched
assert len(_tool_by_id(result, "call_1")["content"]) < 9000 # over floor -> summarized
def test_structure_preserved():
c = _compressor(proactive_prune_tokens=48_000, proactive_prune_min_result_chars=8_000)
msgs = _build(8, big_indices={0, 1, 2})
roles_before = [m["role"] for m in msgs]
ids_before = [m.get("tool_call_id") for m in msgs]
result, _ = c.prune_tool_results_only(msgs, current_tokens=120_000)
assert len(result) == len(msgs)
assert [m["role"] for m in result] == roles_before
assert [m.get("tool_call_id") for m in result] == ids_before
def test_idempotent():
c = _compressor(proactive_prune_tokens=48_000, proactive_prune_min_result_chars=8_000)
msgs = _build(8, big_indices={0, 1, 2})
first, n1 = c.prune_tool_results_only(msgs, current_tokens=120_000)
assert n1 >= 3
second, n2 = c.prune_tool_results_only(first, current_tokens=120_000)
assert n2 == 0
assert [m.get("content") for m in second] == [m.get("content") for m in first]
def test_prune_old_tool_results_default_floor_unchanged():
"""Backward-compat: without min_prune_chars, _prune_old_tool_results still
prunes >200-char results (the compression Phase-1 caller's behavior)."""
c = _compressor()
msgs = _build(8, big_indices=set())
for m in msgs: # a 300-char old tool result
if m.get("tool_call_id") == "call_0":
m["content"] = "Q" * 300
result, pruned = c._prune_old_tool_results(msgs, protect_tail_count=4)
assert len(_tool_by_id(result, "call_0")["content"]) < 300
assert pruned >= 1
def test_min_result_chars_floor_is_clamped():
"""Config-robustness: a floor below 200 (or negative) is clamped up to 200,
while a configured 0 falls back to the 8000 default via ``or``. Without the
clamp, a tiny floor lets Pass 2 re-summarize its own (short) summary every
turn, and a negative floor strips every non-tail tool result."""
assert _compressor(proactive_prune_min_result_chars=0).proactive_prune_min_result_chars == 8000
assert _compressor(proactive_prune_min_result_chars=50).proactive_prune_min_result_chars == 200
assert _compressor(proactive_prune_min_result_chars=-1).proactive_prune_min_result_chars == 200
assert _compressor(proactive_prune_min_result_chars=8000).proactive_prune_min_result_chars == 8000

View file

@ -752,6 +752,9 @@ compression:
hygiene_hard_message_limit: 5000 # Gateway safety valve — see below
hygiene_timeout_seconds: 30 # Max seconds gateway waits for pre-agent hygiene compression
hygiene_failure_cooldown_seconds: 300 # Skip repeated failed hygiene attempts for this session
proactive_prune_tokens: 0 # Opt-in tokens trigger for the no-LLM tool-result prune (0 = off; see below)
proactive_prune_min_result_chars: 8000 # Prune's summarize pass only touches tool results larger than this (clamped >= 200)
proactive_prune_min_reclaim_tokens: 4096 # Prune only commits when it reclaims at least this many tokens (0 = commit any)
# The summarization model/provider is configured under auxiliary:
auxiliary:
@ -777,6 +780,8 @@ Older configs with `compression.summary_model`, `compression.summary_provider`,
`idle_compact_after_seconds` is an **opt-in, time-based** trigger that complements the size-based `threshold`. Default `0` (disabled). When set above 0, a session that resumes after at least that many seconds of inactivity compacts its accumulated history up front, before the first reply — so a long-lived thread (e.g. a Telegram conversation you come back to hours later) doesn't re-read its full stale context on every subsequent turn. It never fires when the context is already at or below the post-compression target (`threshold × target_ratio`), and it honors the same failure-cooldown, anti-thrash, and per-session lock guards as every automatic compaction. Example: `idle_compact_after_seconds: 1800` compacts after 30 minutes idle.
`proactive_prune_tokens` enables a deterministic, no-LLM prune of old tool-result payloads that runs independently of `threshold`. On large-window models the `threshold` compaction (≈50% of the window) rarely fires, so bulky tool outputs (terminal dumps, file reads, web extracts) ride along in history and get re-sent on every subsequent turn. When re-sent history exceeds `proactive_prune_tokens` (default `0` = off; try `48000` to enable), the prune dedupes identical results, summarizes older oversized ones, and truncates large tool-call arguments — protecting the most recent `protect_last_n` messages and never calling the model. Full outputs stay recoverable from the session store. `proactive_prune_min_result_chars` (default `8000`, clamped to ≥ 200) sets the size below which a tool result is left untouched. `proactive_prune_min_reclaim_tokens` (default `4096`) prevents a prune from committing unless it reclaims at least that many tokens — a committed prune rewrites already-sent history and invalidates the provider's prompt-cache prefix, so this gate keeps those cache breaks episodic and amortized (one meaningful break, like a compression boundary) instead of firing on every tool iteration. This runs only under the built-in `compressor` engine; other context engines inherit a no-op.
:::tip Gateway hot-reload of compression and context length
As of recent releases, editing `model.context_length` or any `compression.*` key in `config.yaml` on a running gateway takes effect on the next message — no gateway restart, no `/reset`, no session rotation required. The cached-agent signature includes these keys, so the gateway transparently rebuilds the agent when it sees a change. API keys and tool/skill config still require the usual reload paths.
:::