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
+22 -1
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,
+121
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
@@ -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)."
)
@@ -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")