fix(compression): rearm budget from verified usage
This commit is contained in:
parent
72c828ca2c
commit
380e4da3d1
|
|
@ -168,33 +168,29 @@ _HANDOFF_SKIP_FINAL_RESPONSE = (
|
|||
# to treat it as cancellation metadata rather than assistant prose.
|
||||
INTERRUPT_WAITING_FOR_MODEL_PREFIX = "Operation interrupted: waiting for model response ("
|
||||
|
||||
# Refund the per-turn compression budget only when the assembled request has
|
||||
# dropped this far below threshold_tokens. The gap between this margin and
|
||||
# 1.0 is the anti-thrash belt: a compaction that lands barely under threshold
|
||||
# keeps its burnt attempt, so borderline shrink/regrow cycles cannot recycle
|
||||
# the budget indefinitely.
|
||||
_COMPRESSION_BUDGET_REFUND_MARGIN = 0.8
|
||||
|
||||
|
||||
def _should_refund_compression_budget(
|
||||
def _should_rearm_compression_budget(
|
||||
compression_attempts: int,
|
||||
request_pressure_tokens: int,
|
||||
*,
|
||||
completed_compaction_pending: bool,
|
||||
prompt_tokens: int,
|
||||
threshold_tokens: int,
|
||||
) -> bool:
|
||||
"""True when burnt compression attempts should be returned to the turn.
|
||||
"""Return True after a provider proves a completed compaction worked.
|
||||
|
||||
A compaction that brought the assembled request comfortably back under
|
||||
threshold made real progress, so it must not permanently consume the
|
||||
per-turn overflow-recovery backstop. No-progress passes never get here
|
||||
below the margin, so they still exhaust the budget (#11529).
|
||||
Rough estimates cannot safely rearm the anti-thrash budget: they can dip
|
||||
below the threshold while the provider-visible prompt remains too large.
|
||||
Require the completed-compaction latch plus a positive, normalized prompt
|
||||
count below the threshold from the next successful provider response.
|
||||
"""
|
||||
return bool(
|
||||
compression_attempts
|
||||
and completed_compaction_pending
|
||||
and threshold_tokens > 0
|
||||
and request_pressure_tokens
|
||||
< threshold_tokens * _COMPRESSION_BUDGET_REFUND_MARGIN
|
||||
and 0 < prompt_tokens < threshold_tokens
|
||||
)
|
||||
|
||||
|
||||
# Modules that indicate a deterministic local processing error when they
|
||||
# appear in an exception traceback WITHOUT any API-call module. Used by the
|
||||
# outer-loop error classifier to avoid retrying bugs that will fail
|
||||
|
|
@ -1655,7 +1651,10 @@ def run_conversation(
|
|||
compression_attempts = 0
|
||||
# One resolved per-turn compression attempt cap, shared by every site that
|
||||
# consumes ``compression_attempts``: the pre-API pressure gate, the
|
||||
# overflow/413 retry handlers, and the post-tool compaction gate.
|
||||
# overflow/413 retry handlers, and the post-tool compaction gate. The
|
||||
# counter is a consecutive unverified/ineffective-attempt backstop: a
|
||||
# completed compaction rearms it only after a successful provider response
|
||||
# reports a prompt below the threshold.
|
||||
# Config-driven via compression.max_attempts (parsed + validated in
|
||||
# agent_init); default 3 preserves the prior hardcoded behavior for
|
||||
# objects without the attribute (older pickles / minimal stubs).
|
||||
|
|
@ -2316,31 +2315,6 @@ def run_conversation(
|
|||
f"{_previous_preflight_pressure:,}",
|
||||
f"{request_pressure_tokens:,}",
|
||||
)
|
||||
# Refund the shared overflow-recovery budget once the assembled
|
||||
# request is comfortably back under threshold: a compaction that got
|
||||
# us here made real progress, so it must not permanently consume the
|
||||
# per-turn backstop. Without this, a long tool-heavy turn burns all
|
||||
# attempts on *successful* pre-API compactions, the gate below goes
|
||||
# permanently dark, and the context grows unchecked until the
|
||||
# provider rejects the request terminally (root cause of the
|
||||
# max-compression-attempts dead end on marathon turns). The
|
||||
# compressor's own should_compress() must agree the pressure is gone —
|
||||
# when it and the local estimate disagree (#36718 noise, runaway-
|
||||
# compaction stubs), the budget stays burnt and the hard per-turn cap
|
||||
# keeps its original meaning.
|
||||
if _should_refund_compression_budget(
|
||||
compression_attempts, request_pressure_tokens, _preflight_threshold
|
||||
) and not _compressor.should_compress(request_pressure_tokens):
|
||||
logger.info(
|
||||
"Compression budget refunded: ~%s request tokens < %s%% of "
|
||||
"%s threshold (attempts were %s/%s)",
|
||||
f"{request_pressure_tokens:,}",
|
||||
int(_COMPRESSION_BUDGET_REFUND_MARGIN * 100),
|
||||
f"{_preflight_threshold:,}",
|
||||
compression_attempts,
|
||||
max_compression_attempts,
|
||||
)
|
||||
compression_attempts = 0
|
||||
_defer_preflight = getattr(
|
||||
_compressor, "should_defer_preflight_to_real_usage", lambda _t: False
|
||||
)
|
||||
|
|
@ -3736,7 +3710,37 @@ def run_conversation(
|
|||
"cache_write_tokens": canonical_usage.cache_write_tokens,
|
||||
"reasoning_tokens": canonical_usage.reasoning_tokens,
|
||||
}
|
||||
# Capture the boundary latch before update_from_response()
|
||||
# consumes it. Only a real provider prompt count for the
|
||||
# request immediately following a completed compaction can
|
||||
# prove that attempt effective and rearm the shared budget.
|
||||
_completed_compaction_pending = bool(
|
||||
getattr(
|
||||
agent.context_compressor,
|
||||
"_verify_compaction_cleared_threshold",
|
||||
False,
|
||||
)
|
||||
)
|
||||
agent.context_compressor.update_from_response(usage_dict)
|
||||
_compression_threshold = int(
|
||||
getattr(agent.context_compressor, "threshold_tokens", 0)
|
||||
or 0
|
||||
)
|
||||
if _should_rearm_compression_budget(
|
||||
compression_attempts,
|
||||
completed_compaction_pending=_completed_compaction_pending,
|
||||
prompt_tokens=prompt_tokens,
|
||||
threshold_tokens=_compression_threshold,
|
||||
):
|
||||
logger.info(
|
||||
"Compression budget rearmed after provider-confirmed "
|
||||
"recovery: prompt=%s < threshold=%s (attempts were %s/%s)",
|
||||
f"{prompt_tokens:,}",
|
||||
f"{_compression_threshold:,}",
|
||||
compression_attempts,
|
||||
max_compression_attempts,
|
||||
)
|
||||
compression_attempts = 0
|
||||
|
||||
# Stash this response's canonical usage so the post-turn
|
||||
# on_turn_complete() observation hook can forward it (the
|
||||
|
|
|
|||
|
|
@ -1,4 +1,4 @@
|
|||
"""Behavioral tests for the per-turn compression budget refund.
|
||||
"""Behavioral tests for provider-confirmed compression-budget rearming.
|
||||
|
||||
``compression_attempts`` is a shared per-turn backstop (pre-API gate,
|
||||
overflow/413 handlers, post-tool gate). Before the refund fix, *successful*
|
||||
|
|
@ -7,15 +7,9 @@ attempts on compactions that worked, the pre-API gate went dark for the rest
|
|||
of the turn, and the context grew unchecked until the provider rejected the
|
||||
request terminally ("max compression attempts (N) reached").
|
||||
|
||||
The refund returns the budget when BOTH hold at the top of a loop pass:
|
||||
|
||||
* the assembled request sits below ``threshold_tokens *
|
||||
_COMPRESSION_BUDGET_REFUND_MARGIN`` (real progress, not a borderline
|
||||
shrink), and
|
||||
* the compressor's own ``should_compress()`` agrees there is no pressure
|
||||
(divergent-signal guard — a compressor that still demands compression
|
||||
keeps the hard cap's original meaning, see
|
||||
test_post_tool_compression_attempt_cap.py).
|
||||
The budget is rearmed only when a completed compaction is followed by a real
|
||||
provider prompt count below the configured threshold. Rough estimates and
|
||||
usage-less responses cannot reopen the anti-thrash cap.
|
||||
|
||||
These tests drive ``run_conversation()`` through real tool iterations — no
|
||||
source inspection, only observable compaction counts.
|
||||
|
|
@ -29,7 +23,7 @@ from unittest.mock import MagicMock, patch
|
|||
|
||||
import pytest
|
||||
|
||||
from agent.conversation_loop import _should_refund_compression_budget
|
||||
from agent.conversation_loop import _should_rearm_compression_budget
|
||||
from run_agent import AIAgent
|
||||
|
||||
|
||||
|
|
@ -38,23 +32,35 @@ from run_agent import AIAgent
|
|||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestRefundDecision:
|
||||
def test_no_attempts_no_refund(self):
|
||||
assert not _should_refund_compression_budget(0, 100, 10_000)
|
||||
class TestRearmDecision:
|
||||
def test_provider_confirmed_recovery_rearms(self):
|
||||
assert _should_rearm_compression_budget(
|
||||
2,
|
||||
completed_compaction_pending=True,
|
||||
prompt_tokens=7_999,
|
||||
threshold_tokens=10_000,
|
||||
)
|
||||
|
||||
def test_zero_threshold_no_refund(self):
|
||||
assert not _should_refund_compression_budget(2, 100, 0)
|
||||
|
||||
def test_barely_under_threshold_no_refund(self):
|
||||
# 9,999 of 10,000 is inside the anti-thrash belt (margin 0.8).
|
||||
assert not _should_refund_compression_budget(2, 9_999, 10_000)
|
||||
|
||||
def test_at_margin_no_refund(self):
|
||||
assert not _should_refund_compression_budget(2, 8_000, 10_000)
|
||||
|
||||
def test_comfortably_under_margin_refunds(self):
|
||||
assert _should_refund_compression_budget(2, 7_999, 10_000)
|
||||
assert _should_refund_compression_budget(1, 100, 10_000)
|
||||
@pytest.mark.parametrize(
|
||||
("attempts", "pending", "prompt_tokens", "threshold_tokens"),
|
||||
[
|
||||
(0, True, 7_999, 10_000),
|
||||
(2, False, 7_999, 10_000),
|
||||
(2, True, 0, 10_000),
|
||||
(2, True, 10_000, 10_000),
|
||||
(2, True, 10_001, 10_000),
|
||||
(2, True, 7_999, 0),
|
||||
],
|
||||
)
|
||||
def test_unverified_or_pressured_response_keeps_budget_burned(
|
||||
self, attempts, pending, prompt_tokens, threshold_tokens
|
||||
):
|
||||
assert not _should_rearm_compression_budget(
|
||||
attempts,
|
||||
completed_compaction_pending=pending,
|
||||
prompt_tokens=prompt_tokens,
|
||||
threshold_tokens=threshold_tokens,
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
|
|
@ -70,7 +76,17 @@ def _tool_call(i: int):
|
|||
)
|
||||
|
||||
|
||||
def _tool_response(i: int):
|
||||
def _usage(prompt_tokens: int | None):
|
||||
if prompt_tokens is None:
|
||||
return None
|
||||
return SimpleNamespace(
|
||||
prompt_tokens=prompt_tokens,
|
||||
completion_tokens=1,
|
||||
total_tokens=prompt_tokens + 1,
|
||||
)
|
||||
|
||||
|
||||
def _tool_response(i: int, prompt_tokens: int | None):
|
||||
msg = SimpleNamespace(
|
||||
content=None,
|
||||
reasoning_content=None,
|
||||
|
|
@ -78,10 +94,12 @@ def _tool_response(i: int):
|
|||
tool_calls=[_tool_call(i)],
|
||||
)
|
||||
choice = SimpleNamespace(message=msg, finish_reason="tool_calls")
|
||||
return SimpleNamespace(choices=[choice], model="test/model", usage=None)
|
||||
return SimpleNamespace(
|
||||
choices=[choice], model="test/model", usage=_usage(prompt_tokens)
|
||||
)
|
||||
|
||||
|
||||
def _stop_response():
|
||||
def _stop_response(prompt_tokens: int | None):
|
||||
msg = SimpleNamespace(
|
||||
content="done",
|
||||
reasoning_content=None,
|
||||
|
|
@ -89,7 +107,9 @@ def _stop_response():
|
|||
tool_calls=None,
|
||||
)
|
||||
choice = SimpleNamespace(message=msg, finish_reason="stop")
|
||||
return SimpleNamespace(choices=[choice], model="test/model", usage=None)
|
||||
return SimpleNamespace(
|
||||
choices=[choice], model="test/model", usage=_usage(prompt_tokens)
|
||||
)
|
||||
|
||||
|
||||
def _make_tool_defs(*names: str) -> list:
|
||||
|
|
@ -129,9 +149,18 @@ def _coherent_compressor() -> MagicMock:
|
|||
compressor.threshold_tokens = THRESHOLD
|
||||
compressor.context_length = 200_000
|
||||
compressor.last_prompt_tokens = 0
|
||||
compressor._verify_compaction_cleared_threshold = False
|
||||
compressor.awaiting_real_usage_after_compression = False
|
||||
compressor.should_compress.side_effect = lambda t=None: (t or 0) >= THRESHOLD
|
||||
compressor.should_defer_preflight_to_real_usage.return_value = False
|
||||
compressor.get_active_compression_failure_cooldown.return_value = None
|
||||
|
||||
def _update_from_response(usage):
|
||||
compressor.last_prompt_tokens = int(usage.get("prompt_tokens", 0) or 0)
|
||||
compressor._verify_compaction_cleared_threshold = False
|
||||
compressor.awaiting_real_usage_after_compression = False
|
||||
|
||||
compressor.update_from_response.side_effect = _update_from_response
|
||||
return compressor
|
||||
|
||||
|
||||
|
|
@ -161,10 +190,14 @@ def agent():
|
|||
return a
|
||||
|
||||
|
||||
def _run_marathon_turn(agent, n_tool_iterations: int):
|
||||
def _run_marathon_turn(
|
||||
agent, n_tool_iterations: int, *, provider_prompt_tokens: int | None
|
||||
):
|
||||
"""Drive one turn of ``n_tool_iterations`` oversized tool results."""
|
||||
responses = [_tool_response(i) for i in range(n_tool_iterations)]
|
||||
responses.append(_stop_response())
|
||||
responses = [
|
||||
_tool_response(i, provider_prompt_tokens) for i in range(n_tool_iterations)
|
||||
]
|
||||
responses.append(_stop_response(provider_prompt_tokens))
|
||||
agent.client.chat.completions.create.side_effect = responses
|
||||
|
||||
compress_calls = []
|
||||
|
|
@ -172,8 +205,11 @@ def _run_marathon_turn(agent, n_tool_iterations: int):
|
|||
def _fake_compress(messages, system_message, **_kwargs):
|
||||
# Model a compaction that works: blank out every oversized payload,
|
||||
# keeping roles and tool-call pairing intact so sanitization is
|
||||
# unaffected. Pressure drops far below the refund margin.
|
||||
# unaffected. Arm the same provider-verification boundary as the real
|
||||
# compression path.
|
||||
compress_calls.append(len(messages))
|
||||
agent.context_compressor._verify_compaction_cleared_threshold = True
|
||||
agent.context_compressor.awaiting_real_usage_after_compression = True
|
||||
compacted = [
|
||||
dict(m, content="[summarized]")
|
||||
if isinstance(m, dict) and len(str(m.get("content") or "")) > 5_000
|
||||
|
|
@ -209,7 +245,11 @@ class TestCompressionBudgetRefund:
|
|||
completes.
|
||||
"""
|
||||
assert agent.max_compression_attempts == 3 # config default
|
||||
result, compress_calls = _run_marathon_turn(agent, n_tool_iterations=8)
|
||||
result, compress_calls = _run_marathon_turn(
|
||||
agent,
|
||||
n_tool_iterations=8,
|
||||
provider_prompt_tokens=THRESHOLD - 1,
|
||||
)
|
||||
|
||||
assert result["completed"] is True
|
||||
assert len(compress_calls) > 3, (
|
||||
|
|
@ -217,19 +257,19 @@ class TestCompressionBudgetRefund:
|
|||
f"got only {len(compress_calls)} compactions for 8 pressure spikes"
|
||||
)
|
||||
|
||||
def test_no_refund_when_compressor_still_reports_pressure(self, agent):
|
||||
"""Divergent signals: compressor demands compression regardless of
|
||||
the local estimate → budget stays burnt at the hard cap.
|
||||
|
||||
Mirrors the always-True stub of the attempt-cap regression tests —
|
||||
the refund must not reopen that runaway."""
|
||||
agent.context_compressor.should_compress.side_effect = None
|
||||
agent.context_compressor.should_compress.return_value = True
|
||||
|
||||
result, compress_calls = _run_marathon_turn(agent, n_tool_iterations=8)
|
||||
@pytest.mark.parametrize("provider_prompt_tokens", [None, THRESHOLD])
|
||||
def test_unverified_or_pressured_compaction_stays_capped(
|
||||
self, agent, provider_prompt_tokens
|
||||
):
|
||||
"""Missing usage or real usage at threshold cannot recycle the cap."""
|
||||
result, compress_calls = _run_marathon_turn(
|
||||
agent,
|
||||
n_tool_iterations=8,
|
||||
provider_prompt_tokens=provider_prompt_tokens,
|
||||
)
|
||||
|
||||
assert result["completed"] is True
|
||||
assert len(compress_calls) <= agent.max_compression_attempts, (
|
||||
"with should_compress pinned True the per-turn cap must hold; "
|
||||
"without provider-confirmed headroom the per-turn cap must hold; "
|
||||
f"got {len(compress_calls)} compactions"
|
||||
)
|
||||
|
|
|
|||
Loading…
Reference in New Issue