fix(compression): stop detached stale attempts writing shared compressor state
The stall-fallback detaches a timed-out primary worker and reuses the same ContextCompressor, but the existing attempt-generation guards only covered the unwind-time snapshot restore. Every other summary-state write stayed reachable by the still-running primary after the fallback took over: a late successful summary published _previous_summary and cleared the fallback's cooldown, a late failure armed a shared failure cooldown and stamped error state, the cancel rollback and the abort rollback reverted _previous_summary to the primary's snapshot, and the durable cooldown rollback row could be overwritten mid-restore. Compressor code could not fix this with the shared attributes alone: those cells only name the current owner, never the calling attempt. The calling attempt's generation now rides a ContextVar bound inside _run_summary_dispatch, which every attempt's compress_fn passes through in its own thread, so each attempt reads its own generation. Gates on the working-attempt marker (not the entry claim, so lock sit-outs do not suppress the owner) now cover the cancel rollback, late-success writes, _on_summary_failure, the abort rollback, the deterministic pin, compress() entry, and a Phase-3 choke point. The durable cooldown rollback moved inside the claim lock so the DB row and the in-memory restore are atomic against _claim_compressor_attempt. Regression tests drive the real interleavings deterministically, including two threaded end-to-end arms through _run_summary_dispatch and a real ContextCompressor. (cherry picked from commit 902bfc229e140becfb36679dc33bad550c2ae1e8)
This commit is contained in:
@@ -3614,6 +3614,17 @@ Summary generation was unavailable, so this is a best-effort deterministic fallb
|
||||
summary = self._ground_historical_task_snapshot(summary, turns_to_summarize)
|
||||
summary = self._augment_summary_lean(summary, turns_to_summarize)
|
||||
self._validate_summary_user_provenance(summary, has_user_turn)
|
||||
# A detached stale attempt must not publish its late summary onto shared compressor state:
|
||||
# the fallback already advanced _previous_summary and owns the cooldown/error fields. The
|
||||
# candidate itself is discarded downstream by the working-attempt check; bail here so the
|
||||
# attribute writes never land. Entry-generation claims (lock sit-outs) do not count; the
|
||||
# working marker is the ownership boundary for summary state.
|
||||
from agent.conversation_compression import (
|
||||
_COMPRESSOR_ATTEMPT_GENERATION,
|
||||
_working_attempt_is_current,
|
||||
)
|
||||
if not _working_attempt_is_current(self, _COMPRESSOR_ATTEMPT_GENERATION.get()):
|
||||
raise AuxiliaryExplicitCancellation()
|
||||
self._previous_summary = summary
|
||||
self._clear_compression_failure_cooldown()
|
||||
self._summary_model_fallen_back = False
|
||||
@@ -3761,6 +3772,14 @@ Write only the summary body. Do not include any preamble or prefix."""
|
||||
self, e: Exception, turns_to_summarize: List[Dict[str, Any]], focus_topic: Optional[str], memory_context: str,
|
||||
) -> Optional[str]:
|
||||
"""Classify a summary-call failure; retry once on the main model (returning its result) or arm a cooldown (None)."""
|
||||
# A detached stale attempt must not arm a failure cooldown or stamp error state the fallback
|
||||
# attempt owns; unwind as a cancellation so none of the shared-state writes below can land.
|
||||
from agent.conversation_compression import (
|
||||
_COMPRESSOR_ATTEMPT_GENERATION,
|
||||
_working_attempt_is_current,
|
||||
)
|
||||
if not _working_attempt_is_current(self, _COMPRESSOR_ATTEMPT_GENERATION.get()):
|
||||
raise AuxiliaryExplicitCancellation()
|
||||
# Only a genuine no-provider RuntimeError gets the long cooldown; empty/invalid-response
|
||||
# RuntimeErrors are transient and must get the main-model retry below first.
|
||||
# ``call_llm`` raises ``RuntimeError`` for two very different cases: 1. 2. An empty/invalid response
|
||||
@@ -4801,8 +4820,15 @@ Write only the summary body. Do not include any preamble or prefix."""
|
||||
"%d message(s) preserved unchanged. Conversation is frozen until the next /compress or /new.",
|
||||
)
|
||||
telemetry["failure_class"] = failure_class
|
||||
# Roll back the self-heal rehydration so the aborted attempt is a true no-op (#57835).
|
||||
self._previous_summary = previous_summary_before_scan
|
||||
# Roll back the self-heal rehydration so the aborted attempt is a true no-op (#57835). Only the
|
||||
# attempt still owning summary work may roll back: a detached stale attempt (reachable here via
|
||||
# the deterministic summary pin) must not revert the fallback's _previous_summary.
|
||||
from agent.conversation_compression import (
|
||||
_COMPRESSOR_ATTEMPT_GENERATION,
|
||||
_working_attempt_is_current,
|
||||
)
|
||||
if _working_attempt_is_current(self, _COMPRESSOR_ATTEMPT_GENERATION.get()):
|
||||
self._previous_summary = previous_summary_before_scan
|
||||
if not self.quiet_mode:
|
||||
logger.warning(message, n_skipped)
|
||||
return True
|
||||
@@ -4995,6 +5021,14 @@ Write only the summary body. Do not include any preamble or prefix."""
|
||||
WITHOUT clearing it (#100661). Set by provider-proven overflow recovery, which is already bounded by
|
||||
the caller's attempt budget.
|
||||
"""
|
||||
# A detached stale attempt must not even reset per-call state the fallback owns. Staleness that
|
||||
# arises mid-compress is caught by the write-point gates below; this covers stale-at-entry.
|
||||
from agent.conversation_compression import (
|
||||
_COMPRESSOR_ATTEMPT_GENERATION,
|
||||
_working_attempt_is_current,
|
||||
)
|
||||
if not _working_attempt_is_current(self, _COMPRESSOR_ATTEMPT_GENERATION.get()):
|
||||
raise AuxiliaryExplicitCancellation()
|
||||
telemetry = self._begin_compress_attempt(current_tokens, force)
|
||||
n_messages = len(messages)
|
||||
# Only need head + 3 tail messages minimum (token budget decides the real tail size)
|
||||
@@ -5048,6 +5082,15 @@ Write only the summary body. Do not include any preamble or prefix."""
|
||||
)
|
||||
|
||||
# Phase 3: Generate structured summary (or skip the LLM when the middle is too small to matter)
|
||||
# Choke point for staleness that arose during phases 1-2: everything below writes shared state
|
||||
# (feasibility counters, fallback diagnostics, finalize's cursor/rearm resets), and the inner
|
||||
# _summarize_window/_generate_summary gates cover staleness arising during the LLM call itself.
|
||||
from agent.conversation_compression import (
|
||||
_COMPRESSOR_ATTEMPT_GENERATION,
|
||||
_working_attempt_is_current,
|
||||
)
|
||||
if not _working_attempt_is_current(self, _COMPRESSOR_ATTEMPT_GENERATION.get()):
|
||||
raise AuxiliaryExplicitCancellation()
|
||||
feasibility_skip = not force and self._feasibility_skip(telemetry, turns_to_summarize, compress_start, compress_end)
|
||||
summary = None # feasibility skip: no LLM call; Phase 4 inserts the deterministic fallback
|
||||
if not feasibility_skip:
|
||||
|
||||
@@ -36,6 +36,14 @@ class SummaryDispatchMixin:
|
||||
fallback summary, or abort under ``abort_on_summary_failure`` exactly like a failed summary call."""
|
||||
from agent.context_compressor import take_deterministic_summary_pin
|
||||
if take_deterministic_summary_pin():
|
||||
# A detached stale attempt must not stamp error state or mutate the shared telemetry dict
|
||||
# the fallback owns; unwind as a cancellation before any write lands.
|
||||
from agent.conversation_compression import (
|
||||
_COMPRESSOR_ATTEMPT_GENERATION,
|
||||
_working_attempt_is_current,
|
||||
)
|
||||
if not _working_attempt_is_current(self, _COMPRESSOR_ATTEMPT_GENERATION.get()):
|
||||
raise AuxiliaryExplicitCancellation()
|
||||
# Surfaces through the fallback summary's reason line and the host's one-shot user warning.
|
||||
self._last_summary_error = (
|
||||
"summary model stalled on every route; deterministic fallback summary inserted"
|
||||
@@ -55,6 +63,16 @@ class SummaryDispatchMixin:
|
||||
return self._generate_summary(turns_to_summarize, **summary_kwargs)
|
||||
except AuxiliaryExplicitCancellation:
|
||||
# Cancellation is a true no-op: restore the scan's mutation before the exception escapes.
|
||||
self._previous_summary = scan.previous_summary_before
|
||||
self._summary_has_user_turn = scan.has_user_turn_before
|
||||
# Guard by THIS attempt's ownership: a detached stale attempt (a fallback already claimed
|
||||
# and is working the compressor) must not revert summary state the newer attempt advanced.
|
||||
# The caller's own generation rides a ContextVar because shared compressor attributes can
|
||||
# only name the current owner, never the caller's attempt. Working-attempt comparison (not
|
||||
# the entry generation) so a no-op claim does not suppress the owning attempt's rollback.
|
||||
from agent.conversation_compression import (
|
||||
_COMPRESSOR_ATTEMPT_GENERATION,
|
||||
_working_attempt_is_current,
|
||||
)
|
||||
if _working_attempt_is_current(self, _COMPRESSOR_ATTEMPT_GENERATION.get()):
|
||||
self._previous_summary = scan.previous_summary_before
|
||||
self._summary_has_user_turn = scan.has_user_turn_before
|
||||
raise
|
||||
|
||||
@@ -12,6 +12,7 @@ from __future__ import annotations
|
||||
|
||||
import concurrent.futures
|
||||
import contextlib
|
||||
import contextvars
|
||||
import copy
|
||||
import dataclasses
|
||||
import inspect
|
||||
@@ -214,6 +215,16 @@ def _snapshot_compressor_attempt_state(compressor: Any) -> dict[str, Any]:
|
||||
# ---------------------------------------------------------------------------
|
||||
_COMPRESSOR_ATTEMPT_LOCK = threading.Lock()
|
||||
|
||||
# The calling attempt's generation rides a ContextVar (not a compressor attribute) so compressor code
|
||||
# deep in the call stack can tell ITS OWN attempt apart from whichever attempt currently owns the
|
||||
# compressor. A shared attribute can only answer "who owns now", never "am I stale". Set/reset inside
|
||||
# _run_summary_dispatch around compress_fn, which always runs in the calling attempt's own thread, so
|
||||
# worker and fallback threads each see their own generation. Callers outside the dispatch machinery
|
||||
# (manual compress, legacy paths) read None and keep unguarded historical behavior.
|
||||
_COMPRESSOR_ATTEMPT_GENERATION: contextvars.ContextVar[Any] = contextvars.ContextVar(
|
||||
"hermes_compressor_attempt_generation", default=None
|
||||
)
|
||||
|
||||
|
||||
def _claim_compressor_attempt(compressor: Any) -> int:
|
||||
"""Claim the compressor for a new attempt; return its monotonic generation id.
|
||||
@@ -333,19 +344,13 @@ def _restore_compressor_attempt_state(
|
||||
attempt_generation, getattr(compressor, "_compression_attempt_generation", None),
|
||||
)
|
||||
return
|
||||
# Success clears the durable cooldown pre-commit; recreate/clear that row BEFORE
|
||||
# restoring in-memory values or the next refresh overwrites the rollback. Never
|
||||
# turn unknown durable state / unpersisted local cooldowns into DB writes.
|
||||
if (
|
||||
"_summary_failure_cooldown_until" in snapshot
|
||||
and durable_cooldown_authoritative is not False
|
||||
and (durable_cooldown_authoritative is True or not bool(snapshot.get("_cooldown_persist_failed", False)))
|
||||
):
|
||||
_rollback_durable_cooldown(compressor, snapshot, durable_cooldown_authoritative, durable_cooldown_state)
|
||||
restored = copy.deepcopy(snapshot)
|
||||
# Re-validate under the claim lock: the slow durable rollback above leaves a
|
||||
# window where a fallback may have claimed; stale writes must not interleave.
|
||||
# The rollback itself is safe: landing after a fallback needs a prior claim.
|
||||
# Re-validate under the claim lock AND run the durable rollback inside it: the slow DB
|
||||
# write used to sit between the first ownership check and this re-check, so a fallback
|
||||
# claiming mid-restore could have its freshly written cooldown row overwritten by this
|
||||
# attempt's stale snapshot row. Holding the lock across the durable write serializes
|
||||
# it against _claim_compressor_attempt. The row still lands BEFORE the in-memory
|
||||
# restore so the next refresh cannot overwrite the rollback.
|
||||
with _COMPRESSOR_ATTEMPT_LOCK:
|
||||
if attempt_generation and int(getattr(compressor, "_compression_attempt_generation", 0) or 0) != attempt_generation:
|
||||
logger.warning(
|
||||
@@ -353,6 +358,15 @@ def _restore_compressor_attempt_state(
|
||||
"time: attempt generation %s lost the compressor mid-restore.", attempt_generation,
|
||||
)
|
||||
return
|
||||
# Success clears the durable cooldown pre-commit; recreate/clear that row BEFORE
|
||||
# restoring in-memory values or the next refresh overwrites the rollback. Never
|
||||
# turn unknown durable state / unpersisted local cooldowns into DB writes.
|
||||
if (
|
||||
"_summary_failure_cooldown_until" in snapshot
|
||||
and durable_cooldown_authoritative is not False
|
||||
and (durable_cooldown_authoritative is True or not bool(snapshot.get("_cooldown_persist_failed", False)))
|
||||
):
|
||||
_rollback_durable_cooldown(compressor, snapshot, durable_cooldown_authoritative, durable_cooldown_state)
|
||||
for name, value in restored.items():
|
||||
setattr(compressor, name, value)
|
||||
|
||||
@@ -2941,6 +2955,7 @@ def _run_summary_dispatch(
|
||||
or (commit_fence is not None and commit_fence.is_cancelled)
|
||||
)
|
||||
|
||||
_attempt_ctx_token = _COMPRESSOR_ATTEMPT_GENERATION.set(attempt_generation)
|
||||
try:
|
||||
# F6: never start expensive summary work for an already-cancelled
|
||||
# fence (a stale queued job admitted after host departure).
|
||||
@@ -2964,6 +2979,7 @@ def _run_summary_dispatch(
|
||||
if hard_cancel_event is not None and hard_cancel_event.is_set():
|
||||
raise AuxiliaryExplicitCancellation()
|
||||
finally:
|
||||
_COMPRESSOR_ATTEMPT_GENERATION.reset(_attempt_ctx_token)
|
||||
if commit_fence is not None:
|
||||
_clear_compression_cancelled_check_if_owner(agent.context_compressor, attempt_generation)
|
||||
return compressed
|
||||
|
||||
@@ -242,3 +242,461 @@ class TestMidRestoreClaimRace:
|
||||
|
||||
# The write-time re-check must have refused the stale restore.
|
||||
assert compressor._previous_summary == "FALLBACK STATE"
|
||||
|
||||
|
||||
def _run_as_attempt(generation, fn):
|
||||
"""Run ``fn`` with the attempt-generation ContextVar bound, exactly as
|
||||
``_run_summary_dispatch`` does around ``compress_fn``. Tolerates the var
|
||||
being absent so the tests still exercise behavior on pre-fix revisions."""
|
||||
from agent import conversation_compression as _cc
|
||||
|
||||
var = getattr(_cc, "_COMPRESSOR_ATTEMPT_GENERATION", None)
|
||||
token = var.set(generation) if var is not None else None
|
||||
try:
|
||||
return fn()
|
||||
finally:
|
||||
if token is not None:
|
||||
var.reset(token)
|
||||
|
||||
|
||||
class TestSummarizeWindowCancelRollbackOwnership:
|
||||
"""Gap in the generation guard: ``_summarize_window`` restored
|
||||
``_previous_summary`` / ``_summary_has_user_turn`` unconditionally on
|
||||
``AuxiliaryExplicitCancellation``, so a detached primary's late cancel
|
||||
reverted state a newer fallback attempt already advanced."""
|
||||
|
||||
def _compressor_with_window(self):
|
||||
from agent.context_compressor_summary import SummaryDispatchMixin
|
||||
from agent.auxiliary_client import AuxiliaryExplicitCancellation
|
||||
|
||||
cc = SummaryDispatchMixin.__new__(SummaryDispatchMixin)
|
||||
cc._previous_summary = "fallback-advanced"
|
||||
cc._summary_has_user_turn = True
|
||||
cc._last_summary_error = None
|
||||
cc._active_compression_telemetry = None
|
||||
|
||||
def _generate_summary(turns, **kw):
|
||||
raise AuxiliaryExplicitCancellation()
|
||||
|
||||
cc._generate_summary = _generate_summary
|
||||
return cc
|
||||
|
||||
def _scan(self):
|
||||
return SimpleNamespace(previous_summary_before="primary-era", has_user_turn_before=False)
|
||||
|
||||
def _cancel(self, cc):
|
||||
from agent.auxiliary_client import AuxiliaryExplicitCancellation
|
||||
|
||||
try:
|
||||
cc._summarize_window([], [], self._scan(), "focus", "", False)
|
||||
raise AssertionError("expected AuxiliaryExplicitCancellation")
|
||||
except AuxiliaryExplicitCancellation:
|
||||
pass
|
||||
|
||||
def test_stale_attempt_cancel_does_not_revert_fallback_state(self):
|
||||
cc = self._compressor_with_window()
|
||||
# The real stall-fallback interleaving: attempt 1 detached, fallback
|
||||
# attempt 2 claimed AND published the working marker, then advanced
|
||||
# summary state. Attempt 1's late cancel unwinds after all of it.
|
||||
cc._compression_attempt_generation = 2
|
||||
cc._compression_working_attempt_generation = 2
|
||||
|
||||
_run_as_attempt(1, lambda: self._cancel(cc))
|
||||
|
||||
assert cc._previous_summary == "fallback-advanced"
|
||||
assert cc._summary_has_user_turn is True
|
||||
|
||||
def test_current_attempt_cancel_restores_scan_mutation(self):
|
||||
"""Control: the owning attempt's cancel still rolls back the scan's
|
||||
self-heal mutation so the cancellation stays a true no-op."""
|
||||
cc = self._compressor_with_window()
|
||||
cc._compression_attempt_generation = 1
|
||||
cc._compression_working_attempt_generation = 1
|
||||
|
||||
_run_as_attempt(1, lambda: self._cancel(cc))
|
||||
|
||||
assert cc._previous_summary == "primary-era"
|
||||
assert cc._summary_has_user_turn is False
|
||||
|
||||
def test_noop_entry_claim_does_not_suppress_owning_rollback(self):
|
||||
"""A lock sit-out bumps the entry generation without working the
|
||||
compressor; the working attempt's rollback must still fire."""
|
||||
cc = self._compressor_with_window()
|
||||
cc._compression_attempt_generation = 2
|
||||
cc._compression_working_attempt_generation = 1
|
||||
|
||||
_run_as_attempt(1, lambda: self._cancel(cc))
|
||||
|
||||
assert cc._previous_summary == "primary-era"
|
||||
assert cc._summary_has_user_turn is False
|
||||
|
||||
def test_unguarded_path_restores_normally(self):
|
||||
"""Direct compress() calls outside the attempt machinery carry no
|
||||
generation; the rollback must not be suppressed there."""
|
||||
cc = self._compressor_with_window()
|
||||
|
||||
self._cancel(cc)
|
||||
|
||||
assert cc._previous_summary == "primary-era"
|
||||
|
||||
|
||||
class TestStaleAttemptCompletionWrites:
|
||||
"""The same stale-attempt class at summary-completion time: a detached
|
||||
primary's late success or failure must not publish summary state, arm a
|
||||
failure cooldown, or roll back the fallback's ``_previous_summary``.
|
||||
Guards read the caller's own generation from the dispatch ContextVar and
|
||||
compare it against the working-attempt marker."""
|
||||
|
||||
def _compressor(self):
|
||||
from unittest.mock import patch
|
||||
|
||||
from agent.context_compressor import ContextCompressor
|
||||
|
||||
with patch(
|
||||
"agent.context_compressor.get_model_context_length",
|
||||
return_value=100000,
|
||||
):
|
||||
cc = ContextCompressor(model="test/model", quiet_mode=True)
|
||||
return cc
|
||||
|
||||
def _llm_response(self, content):
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
resp = MagicMock()
|
||||
resp.choices = [MagicMock()]
|
||||
resp.choices[0].message.content = content
|
||||
return resp
|
||||
|
||||
def _own(self, cc, entry=2, working=2):
|
||||
"""Shape the compressor as owned by the fallback (attempt 2)."""
|
||||
cc._compression_attempt_generation = entry
|
||||
cc._compression_working_attempt_generation = working
|
||||
cc._previous_summary = "fallback-advanced"
|
||||
return cc
|
||||
|
||||
def test_stale_attempt_late_success_does_not_publish_summary(self):
|
||||
from unittest.mock import patch
|
||||
|
||||
import pytest
|
||||
|
||||
from agent.auxiliary_client import AuxiliaryExplicitCancellation
|
||||
|
||||
cc = self._own(self._compressor())
|
||||
|
||||
def _call():
|
||||
with patch(
|
||||
"agent.context_compressor.call_llm",
|
||||
return_value=self._llm_response("## Goal\nstale-era summary"),
|
||||
):
|
||||
cc._generate_summary([{"role": "user", "content": "hi"}])
|
||||
|
||||
with pytest.raises(AuxiliaryExplicitCancellation):
|
||||
_run_as_attempt(1, _call)
|
||||
|
||||
assert cc._previous_summary == "fallback-advanced"
|
||||
assert cc._last_summary_error is None
|
||||
|
||||
def test_current_attempt_success_publishes_summary(self):
|
||||
"""Control: the working attempt's identical success still publishes."""
|
||||
from unittest.mock import patch
|
||||
|
||||
cc = self._compressor()
|
||||
cc._compression_attempt_generation = 2
|
||||
cc._compression_working_attempt_generation = 2
|
||||
|
||||
def _call():
|
||||
with patch(
|
||||
"agent.context_compressor.call_llm",
|
||||
return_value=self._llm_response("## Goal\nfresh summary"),
|
||||
):
|
||||
return cc._generate_summary([{"role": "user", "content": "hi"}])
|
||||
|
||||
result = _run_as_attempt(2, _call)
|
||||
|
||||
assert result is not None
|
||||
assert "fresh summary" in cc._previous_summary
|
||||
|
||||
def test_stale_attempt_failure_arms_no_cooldown(self):
|
||||
from unittest.mock import patch
|
||||
|
||||
import pytest
|
||||
|
||||
from agent.auxiliary_client import AuxiliaryExplicitCancellation
|
||||
|
||||
cc = self._own(self._compressor())
|
||||
cc._last_summary_error = "fallback-recorded"
|
||||
|
||||
def _call():
|
||||
with patch(
|
||||
"agent.context_compressor.call_llm",
|
||||
side_effect=ValueError("stale-era transport failure"),
|
||||
):
|
||||
cc._generate_summary([{"role": "user", "content": "hi"}])
|
||||
|
||||
with pytest.raises(AuxiliaryExplicitCancellation):
|
||||
_run_as_attempt(1, _call)
|
||||
|
||||
assert cc._summary_failure_cooldown_until == 0.0
|
||||
assert cc._last_summary_error == "fallback-recorded"
|
||||
|
||||
def test_unguarded_failure_still_records_cooldown(self):
|
||||
"""Control: failures outside the attempt machinery keep legacy
|
||||
cooldown recording (no generation means no guard)."""
|
||||
from unittest.mock import patch
|
||||
|
||||
cc = self._compressor()
|
||||
|
||||
with patch(
|
||||
"agent.context_compressor.call_llm",
|
||||
side_effect=ValueError("transport failure"),
|
||||
):
|
||||
cc._generate_summary([{"role": "user", "content": "hi"}])
|
||||
|
||||
assert cc._summary_failure_cooldown_until > 0.0
|
||||
|
||||
def test_stale_attempt_abort_does_not_restore_previous_summary(self):
|
||||
"""The deterministic-pin path reaches the abort rollback without a
|
||||
cancellation; a stale attempt must not revert the fallback's state."""
|
||||
cc = self._own(self._compressor())
|
||||
cc.abort_on_summary_failure = True
|
||||
|
||||
aborted = _run_as_attempt(
|
||||
1, lambda: cc._abort_on_summary_failure({}, 5, "primary-era")
|
||||
)
|
||||
|
||||
assert aborted is True
|
||||
assert cc._previous_summary == "fallback-advanced"
|
||||
|
||||
def test_current_attempt_abort_restores_previous_summary(self):
|
||||
"""Control: the owning attempt's abort still rolls back the scan's
|
||||
self-heal rehydration (#57835)."""
|
||||
cc = self._compressor()
|
||||
cc.abort_on_summary_failure = True
|
||||
cc._compression_attempt_generation = 1
|
||||
cc._compression_working_attempt_generation = 1
|
||||
cc._previous_summary = "scan-rehydrated"
|
||||
|
||||
aborted = _run_as_attempt(
|
||||
1, lambda: cc._abort_on_summary_failure({}, 5, "pre-scan")
|
||||
)
|
||||
|
||||
assert aborted is True
|
||||
assert cc._previous_summary == "pre-scan"
|
||||
|
||||
|
||||
class TestStaleAttemptEndToEnd:
|
||||
"""E2E: two real attempts through ``_run_summary_dispatch`` on a shared
|
||||
real ``ContextCompressor``. The detached primary's late completion must
|
||||
not write summary state the fallback already owns. Exercises the real
|
||||
plumbing: claim, ContextVar bind inside dispatch, working marker, real
|
||||
``compress()``, and the stale unwind propagating as a cancellation."""
|
||||
|
||||
def _compressor(self):
|
||||
from unittest.mock import patch
|
||||
|
||||
from agent.context_compressor import ContextCompressor
|
||||
|
||||
with patch(
|
||||
"agent.context_compressor.get_model_context_length",
|
||||
return_value=100000,
|
||||
):
|
||||
return ContextCompressor(
|
||||
model="test/model", quiet_mode=True,
|
||||
protect_first_n=2, protect_last_n=2,
|
||||
abort_on_summary_failure=False,
|
||||
)
|
||||
|
||||
def _messages(self, n=12):
|
||||
return [
|
||||
{"role": "user" if i % 2 == 0 else "assistant", "content": f"msg {i}"}
|
||||
for i in range(n)
|
||||
]
|
||||
|
||||
def _llm_response(self, content):
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
resp = MagicMock()
|
||||
resp.choices = [MagicMock()]
|
||||
resp.choices[0].message.content = content
|
||||
return resp
|
||||
|
||||
def test_detached_primary_late_success_cannot_write_after_fallback(self):
|
||||
import threading
|
||||
from unittest.mock import patch
|
||||
|
||||
from agent.auxiliary_client import AuxiliaryExplicitCancellation
|
||||
from agent.conversation_compression import (
|
||||
_claim_compressor_attempt,
|
||||
_run_summary_dispatch,
|
||||
)
|
||||
|
||||
cc = self._compressor()
|
||||
agent = SimpleNamespace(context_compressor=cc, session_id="s1")
|
||||
messages = self._messages()
|
||||
kwargs = {"current_tokens": 999999, "force": True}
|
||||
|
||||
a_in_llm = threading.Event()
|
||||
b_done = threading.Event()
|
||||
outcomes_a = []
|
||||
thread_a = [None]
|
||||
|
||||
def fake_call_llm(**kw):
|
||||
if threading.current_thread() is thread_a[0]:
|
||||
a_in_llm.set()
|
||||
assert b_done.wait(10), "fallback did not complete in time"
|
||||
return self._llm_response("## Goal\nstale-era summary")
|
||||
return self._llm_response("## Goal\nfallback summary")
|
||||
|
||||
def attempt_a():
|
||||
try:
|
||||
_run_summary_dispatch(
|
||||
agent, messages, cc.compress, kwargs,
|
||||
commit_fence=None, attempt_generation=1, hard_cancel_event=None,
|
||||
)
|
||||
except AuxiliaryExplicitCancellation:
|
||||
outcomes_a.append("cancelled")
|
||||
except BaseException as e:
|
||||
outcomes_a.append(f"{type(e).__name__}: {e}")
|
||||
|
||||
gen1 = _claim_compressor_attempt(cc)
|
||||
assert gen1 == 1
|
||||
with patch("agent.context_compressor.call_llm", side_effect=fake_call_llm):
|
||||
t = threading.Thread(target=attempt_a, daemon=True)
|
||||
thread_a[0] = t
|
||||
t.start()
|
||||
assert a_in_llm.wait(10), "primary never reached the provider call"
|
||||
|
||||
# Host detaches the stalled primary and runs the fallback inline.
|
||||
gen2 = _claim_compressor_attempt(cc)
|
||||
assert gen2 == 2
|
||||
_run_summary_dispatch(
|
||||
agent, messages, cc.compress, kwargs,
|
||||
commit_fence=None, attempt_generation=2, hard_cancel_event=None,
|
||||
)
|
||||
assert cc._previous_summary and "fallback summary" in cc._previous_summary
|
||||
|
||||
# The detached primary's provider call finally returns.
|
||||
b_done.set()
|
||||
t.join(10)
|
||||
assert not t.is_alive()
|
||||
|
||||
# The stale attempt unwound as a cancellation and wrote nothing.
|
||||
assert outcomes_a == ["cancelled"]
|
||||
assert "stale-era" not in (cc._previous_summary or "")
|
||||
assert "fallback summary" in cc._previous_summary
|
||||
|
||||
def test_detached_primary_late_cancel_cannot_revert_fallback(self):
|
||||
"""E2E rollback arm: the primary's unwind-time cancellation must not
|
||||
restore its pre-attempt snapshot over the fallback's summary."""
|
||||
import threading
|
||||
from unittest.mock import patch
|
||||
|
||||
from agent.auxiliary_client import AuxiliaryExplicitCancellation
|
||||
from agent.conversation_compression import (
|
||||
_claim_compressor_attempt,
|
||||
_run_summary_dispatch,
|
||||
)
|
||||
|
||||
cc = self._compressor()
|
||||
cc._previous_summary = "pre-attempt summary"
|
||||
agent = SimpleNamespace(context_compressor=cc, session_id="s1")
|
||||
messages = self._messages()
|
||||
kwargs = {"current_tokens": 999999, "force": True}
|
||||
|
||||
a_in_llm = threading.Event()
|
||||
b_done = threading.Event()
|
||||
outcomes_a = []
|
||||
thread_a = [None]
|
||||
|
||||
def fake_call_llm(**kw):
|
||||
if threading.current_thread() is thread_a[0]:
|
||||
a_in_llm.set()
|
||||
assert b_done.wait(10), "fallback did not complete in time"
|
||||
raise AuxiliaryExplicitCancellation()
|
||||
return self._llm_response("## Goal\nfallback summary")
|
||||
|
||||
def attempt_a():
|
||||
try:
|
||||
_run_summary_dispatch(
|
||||
agent, messages, cc.compress, kwargs,
|
||||
commit_fence=None, attempt_generation=1, hard_cancel_event=None,
|
||||
)
|
||||
except AuxiliaryExplicitCancellation:
|
||||
outcomes_a.append("cancelled")
|
||||
except BaseException as e:
|
||||
outcomes_a.append(f"{type(e).__name__}: {e}")
|
||||
|
||||
gen1 = _claim_compressor_attempt(cc)
|
||||
with patch("agent.context_compressor.call_llm", side_effect=fake_call_llm):
|
||||
t = threading.Thread(target=attempt_a, daemon=True)
|
||||
thread_a[0] = t
|
||||
t.start()
|
||||
assert a_in_llm.wait(10)
|
||||
|
||||
gen2 = _claim_compressor_attempt(cc)
|
||||
_run_summary_dispatch(
|
||||
agent, messages, cc.compress, kwargs,
|
||||
commit_fence=None, attempt_generation=2, hard_cancel_event=None,
|
||||
)
|
||||
assert cc._previous_summary and "fallback summary" in cc._previous_summary
|
||||
|
||||
b_done.set()
|
||||
t.join(10)
|
||||
assert not t.is_alive()
|
||||
|
||||
assert outcomes_a == ["cancelled"]
|
||||
# The stale cancel did not roll _previous_summary back to the primary's snapshot.
|
||||
assert "fallback summary" in cc._previous_summary
|
||||
|
||||
|
||||
class TestDurableRollbackAtomicity:
|
||||
"""Gap: the durable cooldown DB write ran BETWEEN the two ownership
|
||||
checks, so a fallback claiming mid-restore had its cooldown row
|
||||
overwritten by the primary's stale snapshot row. The write now sits
|
||||
inside the claim-lock critical section: durable row + in-memory
|
||||
restore are an atomic pair against ``_claim_compressor_attempt``."""
|
||||
|
||||
def test_durable_write_and_restore_are_atomic_against_claims(self):
|
||||
import threading
|
||||
import time
|
||||
|
||||
compressor = _compressor()
|
||||
gen = _claim_compressor_attempt(compressor) # generation 1
|
||||
snapshot = {
|
||||
"_summary_failure_cooldown_until": time.monotonic() + 100.0,
|
||||
"_previous_summary": "primary-era",
|
||||
"_cooldown_persist_failed": False,
|
||||
}
|
||||
written, claims, threads = [], [], []
|
||||
|
||||
class _DB:
|
||||
def record_compression_failure_cooldown(self, sid, until, err):
|
||||
written.append(sid)
|
||||
# Race a fallback claim INTO the durable write. Post-fix the
|
||||
# claim blocks on the same lock until the restore finishes;
|
||||
# pre-fix it lands between the two checks and orphans the write.
|
||||
t = threading.Thread(
|
||||
target=lambda: claims.append(_claim_compressor_attempt(compressor)),
|
||||
daemon=True,
|
||||
)
|
||||
threads.append(t)
|
||||
t.start()
|
||||
t.join(timeout=2.0)
|
||||
|
||||
compressor._session_db = _DB()
|
||||
compressor._session_id = "s1"
|
||||
compressor._previous_summary = "current"
|
||||
|
||||
_restore_compressor_attempt_state(
|
||||
compressor, snapshot, durable_cooldown_authoritative=None,
|
||||
attempt_generation=gen,
|
||||
)
|
||||
for t in threads:
|
||||
t.join(5)
|
||||
|
||||
# Atomicity: the durable row and the in-memory restore are one unit.
|
||||
# Pre-fix the claim slipped between checks: the row was written but
|
||||
# the attribute restore was skipped — the split brain this guards.
|
||||
assert written == ["s1"]
|
||||
assert compressor._previous_summary == "primary-era"
|
||||
assert claims == [2]
|
||||
|
||||
Reference in New Issue
Block a user