fix(gateway): retry-next-message semantics for compression_deferred + regression suite

Gateway half of the #49874 salvage: pass compression_deferred through
both _run_agent_inner result dicts and guard the compression-exhausted
auto-reset block with it — a lock-contended defer keeps the session
intact (the concurrent compressor is actively shrinking it) instead of
wiping it via reset_session.

Regression tests:
- tests/run_agent/test_compression_lock_defer.py — provider-mock 413 and
  400-overflow turns whose compression pass lost the lock end as
  compression_deferred (failed=False, no compression_exhausted); flag
  unset keeps the terminal exhaustion path byte-identical; type-pin
  tests vs MagicMock agents and junk flag values; cap=1 e2e proving the
  refunded pre-API defer leaves the budget for the provider-proven
  413 retry.
- tests/agent/test_preflight_lock_defer.py — a lock-skipped preflight
  pass stops the loop WITHOUT arming preflight_compression_blocked;
  plain no-op still arms it; MagicMock junk does not defer.
- tests/gateway/test_compression_deferred_soft_result.py — AST pin that
  the deferred branch guards the auto-reset chain and performs no
  session mutation (mirrors test_35809_auto_reset_clean_context.py).
This commit is contained in:
Teknium 2026-07-23 12:59:33 -07:00
parent 056a40aa4d
commit eebc2286fc
4 changed files with 615 additions and 1 deletions

View file

@ -13915,7 +13915,20 @@ class GatewayRunner(GatewayAuthorizationMixin, GatewayKanbanWatchersMixin, Gatew
# large to process. Auto-reset it so the next message starts
# fresh instead of replaying the same oversized context in an
# infinite fail loop. (#9893)
if agent_result.get("compression_exhausted") and session_entry and session_key:
#
# A lock-contended defer is the OPPOSITE case: the session is
# temporarily uncompressible only because a concurrent path holds
# the compression lock and is actively shrinking it. Never wipe
# the session for that — retry-next-message semantics apply
# (#69870 lock-skip consumer; salvaged from #49874).
if agent_result.get("compression_deferred"):
logger.info(
"Compression deferred for session %s — the compression "
"lock is held by a concurrent compressor. Keeping the "
"session intact; the next message retries normally.",
session_entry.session_id if session_entry else "?",
)
elif agent_result.get("compression_exhausted") and session_entry and session_key:
logger.info(
"Auto-resetting session %s after compression exhaustion.",
session_entry.session_id,
@ -21996,6 +22009,7 @@ class GatewayRunner(GatewayAuthorizationMixin, GatewayKanbanWatchersMixin, Gatew
"interrupt_message": result.get("interrupt_message"),
"error": result.get("error"),
"compression_exhausted": result.get("compression_exhausted", False),
"compression_deferred": result.get("compression_deferred", False),
"tools": tools_holder[0] or [],
"history_offset": _effective_history_offset,
"compacted_in_place": _compacted_in_place,
@ -22113,6 +22127,13 @@ class GatewayRunner(GatewayAuthorizationMixin, GatewayKanbanWatchersMixin, Gatew
"partial": result_holder[0].get("partial", False) if result_holder[0] else False,
"error": result_holder[0].get("error") if result_holder[0] else None,
"interrupt_message": result_holder[0].get("interrupt_message") if result_holder[0] else None,
# Soft lock-contention defer (#69870 consumer): distinct from
# compression_exhausted so the gateway never auto-resets a
# session that a concurrent compressor is about to shrink.
"compression_deferred": (
result_holder[0].get("compression_deferred", False)
if result_holder[0] else False
),
"tools": tools_holder[0] or [],
"history_offset": _effective_history_offset,
"compacted_in_place": _compacted_in_place,

View file

@ -0,0 +1,121 @@
"""Preflight lock-defer must not arm the insufficient-progress blocker.
Companion to ``tests/run_agent/test_compression_lock_defer.py`` pins the
``build_turn_context`` preflight loop's handling of a lock-contended
compression no-op (#69870 lock-skip signal, consumer salvaged from #49874):
* the pass stops the preflight loop (the lock winner is already shrinking
the session) WITHOUT setting ``preflight_compression_blocked``, so the
loop's provider-proven error handlers keep their full retry budget;
* a genuine no-progress no-op (flag unset) still arms the blocker exactly
as before;
* a MagicMock-style truthy junk flag value does NOT take the defer branch
(type-pin rule).
"""
from __future__ import annotations
import types
from unittest.mock import MagicMock, patch
import pytest
from tests.agent.test_turn_context import _FakeAgent, _build
@pytest.fixture(autouse=True)
def _stub_runtime_main():
with patch("agent.auxiliary_client.set_runtime_main", lambda *a, **k: None):
yield
def _pressured_compressor():
"""Over-threshold compressor stub that opens the preflight threshold path."""
return types.SimpleNamespace(
protect_first_n=0,
protect_last_n=0,
threshold_tokens=1,
context_length=100_000,
last_prompt_tokens=0,
should_compress=lambda _tokens=None: True,
should_compress_info=lambda _tokens=None: (True, None),
should_defer_preflight_to_real_usage=lambda _t: False,
get_active_compression_failure_cooldown=lambda: None,
)
def _make_agent():
agent = _FakeAgent()
agent.compression_enabled = True
agent.context_compressor = _pressured_compressor()
agent._emit_status = MagicMock()
return agent
_HISTORY = [{"role": "user", "content": "old"}, {"role": "assistant", "content": "older"}]
def test_preflight_lock_skip_does_not_set_blocked_flag():
agent = _make_agent()
calls = []
def _lock_skip_compress(messages, _system_message, **_kwargs):
calls.append(1)
agent._compression_skipped_due_to_lock = "pid=1:tid=2:agent=aa:nonce=bb"
return messages, "SYSTEM"
agent._compress_context = _lock_skip_compress
ctx = _build(agent, conversation_history=list(_HISTORY))
# Exactly one pass: the defer stops the loop without arming the blocker.
assert calls == [1]
assert ctx.preflight_compression_blocked is False
def test_preflight_lock_skip_true_unconfirmed_holder_also_defers():
agent = _make_agent()
calls = []
def _lock_skip_compress(messages, _system_message, **_kwargs):
calls.append(1)
agent._compression_skipped_due_to_lock = True
return messages, "SYSTEM"
agent._compress_context = _lock_skip_compress
ctx = _build(agent, conversation_history=list(_HISTORY))
assert calls == [1]
assert ctx.preflight_compression_blocked is False
def test_preflight_plain_noop_still_arms_blocker():
"""Control: flag unset → unchanged pre-fix behavior (blocker armed)."""
agent = _make_agent()
def _noop_compress(messages, _system_message, **_kwargs):
agent._compression_skipped_due_to_lock = None
return messages, "SYSTEM"
agent._compress_context = _noop_compress
ctx = _build(agent, conversation_history=list(_HISTORY))
assert ctx.preflight_compression_blocked is True
def test_preflight_magicmock_flag_value_is_not_a_defer():
"""Type-pin: truthy junk (MagicMock auto-attribute shape) must not be
treated as lock contention the blocker arms as for a plain no-op."""
agent = _make_agent()
def _junk_flag_compress(messages, _system_message, **_kwargs):
agent._compression_skipped_due_to_lock = MagicMock()
return messages, "SYSTEM"
agent._compress_context = _junk_flag_compress
ctx = _build(agent, conversation_history=list(_HISTORY))
assert ctx.preflight_compression_blocked is True

View file

@ -0,0 +1,98 @@
"""Gateway must treat ``compression_deferred`` as a soft result (#49874).
A lock-contended compression defer means a CONCURRENT compressor is actively
shrinking the session the opposite of ``compression_exhausted`` (session
permanently too large). The gateway's auto-reset (#9893/#35809) must never
fire for a deferred turn: the session stays intact and the next message
retries normally.
AST invariants on ``gateway/run.py`` (mirrors
``test_35809_auto_reset_clean_context.py``'s load-bearing pin style):
* the ``compression_deferred`` branch guards the auto-reset block a
deferred result can never reach ``reset_session``;
* the deferred branch itself performs NO session mutation (no
``reset_session``, no ``_evict_cached_agent``, no
``_clear_conversation_scope``).
"""
from __future__ import annotations
import ast
import inspect
from gateway import run as gateway_run
def _calls(node: ast.AST) -> set[str]:
return {
n.func.attr
for n in ast.walk(node)
if isinstance(n, ast.Call) and isinstance(n.func, ast.Attribute)
}
def _find_deferred_guarded_reset_chain() -> ast.If:
"""Return the ``if agent_result.get('compression_deferred') ... elif
agent_result.get('compression_exhausted') ... reset_session`` chain."""
tree = ast.parse(inspect.getsource(gateway_run))
for node in ast.walk(tree):
if not isinstance(node, ast.If):
continue
test_consts = [
n.value
for n in ast.walk(node.test)
if isinstance(n, ast.Constant) and isinstance(n.value, str)
]
if "compression_deferred" not in test_consts:
continue
# The reset must live in the orelse (elif compression_exhausted ...),
# never in the deferred body.
orelse_calls = set()
for sub in node.orelse:
orelse_calls |= _calls(sub)
if "reset_session" in orelse_calls:
return node
raise AssertionError(
"Could not locate the compression_deferred guard in front of the "
"compression-exhausted auto-reset block in gateway/run.py. The "
"soft-defer contract (#49874: lock-contended defer must never "
"auto-reset the session) is no longer structurally guaranteed."
)
class TestCompressionDeferredIsSoft:
def test_deferred_branch_guards_the_auto_reset(self):
"""The auto-reset (``reset_session``) must be unreachable when
``compression_deferred`` is set: the deferred check comes FIRST and
the reset lives only in its elif chain."""
node = _find_deferred_guarded_reset_chain()
# The exhaustion reset is in the orelse — verified by the finder.
# The deferred body must not mutate the session in any way.
body_calls = set()
for sub in node.body:
body_calls |= _calls(sub)
forbidden = {
"reset_session",
"_evict_cached_agent",
"_clear_conversation_scope",
}
assert not (body_calls & forbidden), (
f"The compression_deferred branch in gateway/run.py performs "
f"session mutation ({body_calls & forbidden}). A lock-contended "
f"defer is transient — the session must stay intact so the next "
f"message retries against the freshly compressed context "
f"(#49874, #69870)."
)
def test_deferred_result_key_is_passed_through_run_agent_inner(self):
"""``_run_agent_inner``'s result dicts must carry the
``compression_deferred`` key so the persistence block can see it
the exact gap that made the exhaustion misclassification possible
(the flag existed but nothing consumed it)."""
src = inspect.getsource(gateway_run)
assert src.count('"compression_deferred"') >= 3, (
"gateway/run.py must read AND pass through compression_deferred "
"(persistence-block guard + both _run_agent_inner result dicts)."
)

View file

@ -0,0 +1,374 @@
"""Lock-contended compression no-ops must soft-DEFER, never exhaust (#49874).
On main before this fix, nothing on the automatic compression paths consumed
the #69870 lock-skip signal (``agent._compression_skipped_due_to_lock``):
* a lock-loser preflight/pre-API no-op counted as "insufficient progress",
* the oversized request went to the provider anyway, and
* the lock-contended 413/overflow retry burned ``compression_attempts`` to
the cap and returned ``compression_exhausted`` which the gateway answers
with a full session auto-reset (#9893/#35809).
A temporary lock defer misclassified as exhaustion == session wipe.
These tests pin the fix: when a compression pass returns its input unchanged
AND the type-pinned lock-skip flag is set, the attempt is refunded and the
turn ends (when it cannot proceed) with a soft ``compression_deferred``
result distinct from ``compression_exhausted``.
Salvaged from PR #49874 (@helix4u), rebuilt on the landed #69870 signal.
"""
from __future__ import annotations
from types import SimpleNamespace
from unittest.mock import MagicMock, patch
import pytest
from agent.conversation_compression import compression_skipped_due_to_lock
from run_agent import AIAgent
import run_agent
LOCK_HOLDER = "pid=4242:tid=1:agent=deadbeef:nonce=abcd1234"
# ---------------------------------------------------------------------------
# Helpers (mirrors tests/run_agent/test_413_compression.py)
# ---------------------------------------------------------------------------
@pytest.fixture(autouse=True)
def _no_compression_sleep(monkeypatch):
import time as _time
monkeypatch.setattr(_time, "sleep", lambda *_a, **_k: None)
monkeypatch.setattr(run_agent, "jittered_backoff", lambda *a, **k: 0.0)
def _make_tool_defs(*names: str) -> list:
return [
{
"type": "function",
"function": {
"name": n,
"description": f"{n} tool",
"parameters": {"type": "object", "properties": {}},
},
}
for n in names
]
def _mock_response(content="Hello", finish_reason="stop"):
msg = SimpleNamespace(
content=content,
tool_calls=None,
reasoning_content=None,
reasoning=None,
)
choice = SimpleNamespace(message=msg, finish_reason=finish_reason)
resp = SimpleNamespace(choices=[choice], model="test/model")
resp.usage = None
return resp
def _make_413_error(message="Request entity too large"):
err = Exception(message)
err.status_code = 413
return err
def _make_overflow_error():
return Exception(
"Error code: 400 - {'type': 'error', 'error': {'type': "
"'invalid_request_error', 'message': 'prompt is too long: "
"233153 tokens > 200000 maximum'}}"
)
@pytest.fixture()
def agent():
with (
patch("run_agent.get_tool_definitions", return_value=_make_tool_defs("web_search")),
patch("run_agent.check_toolset_requirements", return_value={}),
patch("run_agent.OpenAI"),
):
a = AIAgent(
api_key="test-key-1234567890",
base_url="https://openrouter.ai/api/v1",
quiet_mode=True,
skip_context_files=True,
skip_memory=True,
)
a.client = MagicMock()
a._cached_system_prompt = "You are helpful."
a._use_prompt_caching = False
a.tool_delay = 0
a.compression_enabled = True
a.save_trajectories = False
return a
_PREFILL = [
{"role": "user", "content": "previous question"},
{"role": "assistant", "content": "previous answer"},
]
def _lock_skipping_compress(agent, *, holder=LOCK_HOLDER):
"""A compress double that no-ops because 'another path holds the lock'.
Mirrors the real ``compress_context`` lock-contended abort: returns the
INPUT list object unchanged and sets the #69870 lock-skip signal.
"""
def _compress(messages, _system_message, **_kwargs):
agent._compression_skipped_due_to_lock = holder
return messages, "You are helpful."
return _compress
def _plain_noop_compress(agent):
"""A compress double that no-ops WITHOUT lock contention (real no-progress)."""
def _compress(messages, _system_message, **_kwargs):
agent._compression_skipped_due_to_lock = None
return messages, "You are helpful."
return _compress
# ---------------------------------------------------------------------------
# Type-pinned signal read (MagicMock test-double immunity)
# ---------------------------------------------------------------------------
class TestLockSkipSignalTypePin:
def test_true_and_holder_string_are_lock_skips(self):
a = SimpleNamespace(_compression_skipped_due_to_lock=True)
assert compression_skipped_due_to_lock(a) is True
a = SimpleNamespace(_compression_skipped_due_to_lock=LOCK_HOLDER)
assert compression_skipped_due_to_lock(a) is True
def test_none_and_missing_are_not_lock_skips(self):
assert compression_skipped_due_to_lock(
SimpleNamespace(_compression_skipped_due_to_lock=None)
) is False
assert compression_skipped_due_to_lock(SimpleNamespace()) is False
def test_magicmock_agent_auto_attribute_is_not_a_lock_skip(self):
"""MagicMock agents auto-create truthy attributes; bare truthiness
would hijack every mocked agent in sibling suites into the lock-skip
branch (the #69870 × #69840 incident). The read must be type-pinned."""
assert compression_skipped_due_to_lock(MagicMock()) is False
def test_truthy_non_true_non_str_values_are_not_lock_skips(self):
for junk in (1, 1.0, ["holder"], {"holder": True}, object(), MagicMock()):
a = SimpleNamespace(_compression_skipped_due_to_lock=junk)
assert compression_skipped_due_to_lock(a) is False, junk
# ---------------------------------------------------------------------------
# 413 handler: lock-contended no-op → soft defer, no exhaustion
# ---------------------------------------------------------------------------
class TestLockContended413Defer:
def test_lock_contended_413_returns_compression_deferred(self, agent):
"""A 413 whose compression pass lost the lock must end the turn as a
soft ``compression_deferred`` never ``compression_exhausted``."""
agent.client.chat.completions.create.side_effect = _make_413_error()
with (
patch.object(
agent, "_compress_context",
side_effect=_lock_skipping_compress(agent),
) as mock_compress,
patch.object(agent, "_persist_session"),
patch.object(agent, "_save_trajectory"),
patch.object(agent, "_cleanup_task_resources"),
):
result = agent.run_conversation("hello", conversation_history=list(_PREFILL))
mock_compress.assert_called_once()
assert result.get("compression_deferred") is True
assert not result.get("compression_exhausted")
# Soft defer: transient, retry-next-message semantics — the gateway
# persists the user turn (failed=False) and never auto-resets.
assert result.get("failed") is False
assert result.get("completed") is False
assert result.get("partial") is True
def test_lock_contended_overflow_returns_compression_deferred(self, agent):
"""Same contract on the context-length (400 prompt-too-long) handler."""
agent.client.chat.completions.create.side_effect = _make_overflow_error()
with (
patch.object(
agent, "_compress_context",
side_effect=_lock_skipping_compress(agent),
) as mock_compress,
patch.object(agent, "_persist_session"),
patch.object(agent, "_save_trajectory"),
patch.object(agent, "_cleanup_task_resources"),
):
result = agent.run_conversation("hello", conversation_history=list(_PREFILL))
mock_compress.assert_called_once()
assert result.get("compression_deferred") is True
assert not result.get("compression_exhausted")
assert result.get("failed") is False
def test_unconfirmed_lock_skip_true_also_defers(self, agent):
"""``_compression_skipped_due_to_lock = True`` (holder unconfirmed —
``try_acquire`` swallowed a sqlite error) is still a lock skip."""
agent.client.chat.completions.create.side_effect = _make_413_error()
with (
patch.object(
agent, "_compress_context",
side_effect=_lock_skipping_compress(agent, holder=True),
),
patch.object(agent, "_persist_session"),
patch.object(agent, "_save_trajectory"),
patch.object(agent, "_cleanup_task_resources"),
):
result = agent.run_conversation("hello", conversation_history=list(_PREFILL))
assert result.get("compression_deferred") is True
assert not result.get("compression_exhausted")
def test_plain_noop_413_still_exhausts_unchanged(self, agent):
"""Control: flag unset (real no-progress compression) keeps the
pre-fix behavior byte-for-byte terminal ``compression_exhausted``."""
agent.client.chat.completions.create.side_effect = _make_413_error()
with (
patch.object(
agent, "_compress_context",
side_effect=_plain_noop_compress(agent),
),
patch.object(
agent, "_try_strip_image_parts_from_tool_messages",
return_value=False,
),
patch.object(agent, "_persist_session"),
patch.object(agent, "_save_trajectory"),
patch.object(agent, "_cleanup_task_resources"),
):
result = agent.run_conversation("hello", conversation_history=list(_PREFILL))
assert result.get("compression_exhausted") is True
assert not result.get("compression_deferred")
assert result.get("failed") is True
def test_magicmock_flag_value_does_not_defer(self, agent):
"""Type-pin at the consumer site: a truthy non-True/non-str flag value
(e.g. a MagicMock auto-attribute) must NOT take the defer branch."""
agent.client.chat.completions.create.side_effect = _make_413_error()
def _junk_flag_compress(messages, _system_message, **_kwargs):
agent._compression_skipped_due_to_lock = MagicMock() # truthy junk
return messages, "You are helpful."
with (
patch.object(agent, "_compress_context", side_effect=_junk_flag_compress),
patch.object(
agent, "_try_strip_image_parts_from_tool_messages",
return_value=False,
),
patch.object(agent, "_persist_session"),
patch.object(agent, "_save_trajectory"),
patch.object(agent, "_cleanup_task_resources"),
):
result = agent.run_conversation("hello", conversation_history=list(_PREFILL))
assert not result.get("compression_deferred")
assert result.get("compression_exhausted") is True
# ---------------------------------------------------------------------------
# Pre-API gate: a lock-skipped pass must not burn the shared attempt budget
# ---------------------------------------------------------------------------
class TestPreApiLockDeferDoesNotBurnBudget:
def test_lock_loser_turn_recovers_after_lock_release(self, agent):
"""End-to-end shape of the live bug at cap=1:
1. Pre-API pressure gate fires; the compression pass loses the lock
(no-op + lock-skip flag). Pre-fix this burned the single shared
attempt.
2. The oversized request goes to the provider 413.
3. The 413 handler compresses again the lock has been released and
the pass now succeeds and the retry completes.
Pre-fix, step 3 found ``compression_attempts`` already at the cap and
returned ``compression_exhausted`` gateway session wipe. The defer
refund keeps the budget intact for the provider-proven retry.
"""
agent.max_compression_attempts = 1
# Compressor stub: pressure only on the fully-assembled request
# (pre-API site); the turn-context preflight stands down via the
# cheap-gate (small message count) and low turn-context estimate.
agent.context_compressor = SimpleNamespace(
protect_first_n=3,
protect_last_n=20,
threshold_tokens=100_000,
context_length=1_000_000,
last_prompt_tokens=0,
should_compress=lambda t: t >= 100_000,
should_defer_preflight_to_real_usage=lambda _t: False,
get_active_compression_failure_cooldown=lambda: None,
)
agent.client.chat.completions.create.side_effect = [
_make_413_error(),
_mock_response(content="Recovered after lock release"),
]
compress_calls = []
def _lock_then_success(messages, _system_message, **_kwargs):
compress_calls.append(len(messages))
if len(compress_calls) == 1:
# Lock loser: no-op + #69870 signal.
agent._compression_skipped_due_to_lock = LOCK_HOLDER
return messages, "You are helpful."
# Lock released: real compaction (entry clears the signal).
agent._compression_skipped_due_to_lock = None
return (
[{"role": "user", "content": "hello"}],
"You are helpful.",
)
with (
patch(
"agent.turn_context.estimate_request_tokens_rough",
return_value=10,
),
patch(
"agent.conversation_loop.estimate_request_tokens_rough",
return_value=500_000,
),
patch(
"agent.conversation_loop.estimate_messages_tokens_rough",
return_value=500_000,
),
patch.object(agent, "_compress_context", side_effect=_lock_then_success),
patch.object(agent, "_persist_session"),
patch.object(agent, "_save_trajectory"),
patch.object(agent, "_cleanup_task_resources"),
):
result = agent.run_conversation("hello", conversation_history=list(_PREFILL))
# Pass 1: pre-API (lock defer, refunded). Pass 2: 413 handler
# (succeeds within the cap because the defer did not count).
assert len(compress_calls) == 2
assert result.get("completed") is True
assert result["final_response"] == "Recovered after lock release"
assert not result.get("compression_exhausted")
assert not result.get("compression_deferred")