ouyangbo/hermes-agent
1
0
Fork 0
Code Issues Pull Requests Actions 10 Packages Projects Releases Wiki Activity

fix(compression): attempt-generation ownership for overlapping stall-fallback attempts

Browse Source 
Follow-up to #96634 (stall-fallback retry, #78981) addressing
donovan-yohan's post-merge adversarial review. The stall path detaches a
timed-out primary worker (fence cancel wins; future stays on the pool)
and immediately runs the fallback against the SAME ContextCompressor,
creating two verified races:

1. Late-primary snapshot restore: the detached primary's unwind called
   _restore_compressor_attempt_state with the PRIMARY's pre-attempt
   snapshot. Landing after the fallback's commit it rolled
   _previous_summary/cooldown/provenance/telemetry back to pre-primary
   values, silently discarding fallback-owned state.
2. Shared _compression_cancelled_check: the late primary's `finally`
   cleared the callback the fallback had just installed, so the
   fallback's F4 cancellation consult read None.

Fix: a monotonic per-compressor attempt generation claimed under one
module lock (_claim_compressor_attempt). Snapshot restores carry their
claiming generation and no-op when stale; the cancelled-check set/clear
moves into owner-stamped helpers (_install_compression_cancelled_check /
_clear_compression_cancelled_check_if_owner) so only the installing
attempt can clear it. The commit fence keeps owning COMMIT admission;
the generation owns compressor-ATTRIBUTE writes — two boundaries.
Legacy callers (attempt_generation=None) and slotted third-party
compressors (generation 0) keep the historical unconditional behavior.

Secondary review items:
- Lean chunk digests during a stall-fallback retry now follow the
  summary onto the pinned healthy route: take_pinned_summary_route()
  echoes the consumed route into a context-local
  _SUMMARY_ROUTE_CONSUMED, and _build_chunk_digests passes
  attempt_summary_route_kwargs() (non-consuming) to call_llm. The pin's
  single-use contract for the SUMMARY call is unchanged — the
  main-model retry still never re-issues the pinned route.
- Worker re-run repeating pre-compression callbacks: documented as an
  accepted limitation on _retry_compression_on_fallback_chain
  (built-ins idempotent; resuming mid-pipeline would couple the retry
  to host callback ordering).

Tests (tests/agent/test_compression_attempt_ownership.py, 10 cases):
deterministic interleavings for both races (late-primary restore
no-ops + preserves fallback state; stale finally cannot clear the
fallback's callback), legacy/slotted compatibility, digest route
follow + context-locality of the consumed echo. Mutation-checked:
reverting only the two prod files to origin/main fails the suite;
restored stack green (21 passed incl. the original #78981 suite).

The one red in the wider sweep
(test_silence_cannot_approach_double_idle_timeout) is pre-existing on
clean origin/main — verified independently.
This commit is contained in:
kshitijk4poor
2026-08-28 12:34:34 +05:30
committed by kshitij
parent d3ccc09dc2
commit 61cd299c6e
3 changed files with 398 additions and 17 deletions
Download Patch File Download Diff File Expand all files Collapse all files
agent/context_compressor.py
+41 -5
View File
@@ -77,15 +77,23 @@ def _safe_int(value: Any) -> int | None:
# per compression run (its only non-recursive call site is the compress path;
# the two recursive calls are the deliberate main-model retry that must NOT
# re-issue the pin). Lean ``tail_mode`` additionally runs
# ``_build_chunk_digests``, which issues its own ``call_llm`` calls directly
# and never consults the pin — during a stall-fallback retry those digests
# still target the stalled primary and degrade to per-segment placeholders.
# Deliberate: the digest path is a best-effort augmentation, not the summary,
# and pinning it would require weakening the single-use contract below.
# ``_build_chunk_digests``, which issues its own ``call_llm`` calls directly.
# Those digests consult ``attempt_summary_route_kwargs()`` (non-consuming):
# during a stall-fallback retry they follow the summary onto the healthy
# fallback backend instead of returning to the stalled primary. The consumed
# echo below preserves the pin's single-use contract for the SUMMARY call —
# the main-model retry still never re-issues the pinned route.
_SUMMARY_ROUTE_PIN: contextvars.ContextVar[Optional[Dict[str, Any]]] = (
contextvars.ContextVar("hermes_summary_route_pin", default=None)
)
# Echo of the route the summary call consumed, for SIBLING aux calls of the
# same attempt (lean digests). Context-local like the pin itself, so it can
# never leak across threads or into an unrelated compression attempt.
_SUMMARY_ROUTE_CONSUMED: contextvars.ContextVar[Optional[Dict[str, Any]]] = (
contextvars.ContextVar("hermes_summary_route_consumed", default=None)
)
# call_llm kwargs a pinned route may set. ``timeout`` lets a fallback entry
# keep its own deadline instead of inheriting one the primary already burned
# (same per-entry semantics the aux client applies to chain candidates).
@@ -120,14 +128,38 @@ def take_pinned_summary_route() -> Optional[Dict[str, Any]]:
Single use by design. ``_generate_summary`` retries itself on the main
model when the summary route fails; re-issuing the pinned route there
would spend a second full deadline on the backend that just failed.
The consumed route is echoed into ``_SUMMARY_ROUTE_CONSUMED`` so that
SIBLING auxiliary calls in the same attempt (the lean chunk digests,
which run after the summary) can keep addressing the healthy fallback
backend instead of silently returning to the stalled task route
(#96634 post-merge review, secondary item).
"""
route = _SUMMARY_ROUTE_PIN.get()
if route is None:
return None
_SUMMARY_ROUTE_PIN.set(None)
_SUMMARY_ROUTE_CONSUMED.set(route)
return route
def attempt_summary_route_kwargs() -> Dict[str, Any]:
"""Route kwargs for sibling aux calls of the CURRENT summary attempt.
Non-consuming. Prefers a still-pending pin (digest paths that run before
the summary), else the route the summary call just consumed. Empty when
no stall-fallback pin is active — normal task routing applies.
"""
route = _SUMMARY_ROUTE_PIN.get() or _SUMMARY_ROUTE_CONSUMED.get()
if not route:
return {}
return {
field: route[field]
for field in _PINNED_ROUTE_FIELDS
if route.get(field) not in (None, "")
}
def _pinned_summary_call_kwargs() -> Dict[str, Any]:
"""Consume the pinned route as explicit ``call_llm`` keyword arguments."""
route = take_pinned_summary_route()
@@ -4646,6 +4678,9 @@ Summary generation was unavailable, so this is a best-effort deterministic fallb
try:
from agent.auxiliary_client import call_llm
# During a stall-fallback retry, follow the summary onto the
# pinned healthy route (non-consuming read) instead of
# re-addressing the stalled task backend (#96634 follow-up).
resp = call_llm(
messages=[{
"role": "user",
@@ -4653,6 +4688,7 @@ Summary generation was unavailable, so this is a best-effort deterministic fallb
}],
task="compression",
max_tokens=_LEAN_DIGEST_MAX_TOKENS,
**attempt_summary_route_kwargs(),
)
body = (
resp.choices[0].message.content
agent/conversation_compression.py
+131 -12
View File
@@ -329,14 +329,121 @@ def _snapshot_compressor_attempt_state(compressor: Any) -> dict[str, Any]:
return copy.deepcopy(selected)
# ---------------------------------------------------------------------------
# Attempt ownership (#96634 follow-up).
#
# The stall-fallback path deliberately DETACHES a timed-out primary worker
# (fence cancel wins; the future stays on the shared pool) and immediately
# starts a fallback attempt against the SAME ContextCompressor. Two races
# follow from that overlap:
#
# 1. The late primary's unwind still calls _restore_compressor_attempt_state
# with the PRIMARY's pre-attempt snapshot. Landing after the fallback's
# commit, it rolls _previous_summary / cooldown / provenance / telemetry
# back to pre-primary values — silently discarding fallback-owned state.
# 2. _compression_cancelled_check is one shared attribute: the late
# primary's ``finally`` clears the callback the fallback just installed,
# so the fallback's F4 cancellation consult reads None.
#
# Both are fixed with a monotonic per-compressor attempt generation, claimed
# under one module lock. Restores and callback set/clear are keyed to the
# claiming generation and no-op when a newer attempt owns the compressor.
# The commit fence still owns COMMIT admission; the generation owns
# compressor-ATTRIBUTE writes — two different boundaries.
# ---------------------------------------------------------------------------
_COMPRESSOR_ATTEMPT_LOCK = threading.Lock()
def _claim_compressor_attempt(compressor: Any) -> int:
"""Claim the compressor for a new attempt; returns its generation id.
Monotonic per compressor instance. Any restore or cancelled-check
mutation stamped with an OLDER generation becomes a no-op, so a
detached, late-unwinding attempt cannot clobber its successor's state.
"""
with _COMPRESSOR_ATTEMPT_LOCK:
generation = int(getattr(compressor, "_compression_attempt_generation", 0) or 0) + 1
try:
compressor._compression_attempt_generation = generation
except Exception:
# Slotted/frozen third-party compressor: ownership tracking is
# unavailable; generation 0 disables the guard (legacy behavior).
return 0
return generation
def _compressor_attempt_is_current(compressor: Any, generation: int) -> bool:
"""True when *generation* still owns the compressor (or guard disabled)."""
if not generation:
return True
with _COMPRESSOR_ATTEMPT_LOCK:
return (
int(getattr(compressor, "_compression_attempt_generation", 0) or 0)
== generation
)
def _install_compression_cancelled_check(
compressor: Any, check: Any, generation: int
) -> None:
"""Install the F4 cancellation consult, stamped with its owner attempt."""
with _COMPRESSOR_ATTEMPT_LOCK:
try:
compressor._compression_cancelled_check = check
compressor._compression_cancelled_check_owner = generation
except Exception:
pass
def _clear_compression_cancelled_check_if_owner(
compressor: Any, generation: int
) -> bool:
"""Clear the cancellation consult only when *generation* installed it.
A detached late primary's ``finally`` must not tear down the callback a
newer fallback attempt just installed. Returns True when cleared.
"""
with _COMPRESSOR_ATTEMPT_LOCK:
owner = getattr(compressor, "_compression_cancelled_check_owner", None)
if owner is not None and generation and owner != generation:
return False
try:
compressor._compression_cancelled_check = None
compressor._compression_cancelled_check_owner = None
except Exception:
pass
return True
def _restore_compressor_attempt_state(
compressor: Any,
snapshot: dict[str, Any],
*,
durable_cooldown_authoritative: Optional[bool] = None,
durable_cooldown_state: Optional[dict[str, Any]] = None,
attempt_generation: Optional[int] = None,
) -> None:
"""Restore the safe per-attempt snapshot after a pre-commit hard cancel."""
"""Restore the safe per-attempt snapshot after a pre-commit hard cancel.
``attempt_generation`` (when provided) is the claim the calling attempt
took via :func:`_claim_compressor_attempt`. A restore stamped with a
stale generation no-ops: the stall-fallback path detaches a timed-out
primary and hands the compressor to a fallback attempt, and the late
primary's unwind must not roll fallback-owned state back to the
primary's pre-attempt snapshot (#96634 post-merge review, claim 1).
"""
if attempt_generation is not None and not _compressor_attempt_is_current(
compressor, attempt_generation
):
logger.warning(
"Skipping stale compressor attempt-state restore: attempt "
"generation %s no longer owns the compressor (current: %s). A "
"newer (stall-fallback) attempt's state is preserved.",
attempt_generation,
getattr(compressor, "_compression_attempt_generation", None),
)
return
# A successful summary clears the durable cooldown before the outer commit
# boundary. Recreate (or clear) that row before restoring exact in-memory
# values, otherwise the next refresh would overwrite this rollback. Unknown
@@ -930,6 +1037,13 @@ def _retry_compression_on_fallback_chain(
entry's own ``timeout`` (when declared) sets that idle window, so a
fallback tuned for a slower-but-healthy backend is not held to a deadline
the stalled primary defined (#62452 semantics, applied to the stall path).
Known limitation (accepted, #96634 review): the retry re-runs the COMPLETE
worker, which repeats memory/plugin pre-compression callbacks. Built-in
callbacks are idempotent (re-reads and overwrites of attempt-scoped
state); third-party plugin callbacks are advised to be. Splitting the
worker to resume mid-pipeline would couple this path to every host's
callback ordering — deliberately out of scope.
"""
# An explicit stop is not a stalled route. The retry worker would abort on
# the same event anyway, but starting one at all makes /stop look ignored.
@@ -2580,6 +2694,10 @@ def compress_context(
_compressor_attempt_snapshot = _snapshot_compressor_attempt_state(
agent.context_compressor
)
# Claim attempt ownership: a detached, late-unwinding sibling attempt
# (stall-fallback overlap) must not restore its snapshot over ours or
# clear our cancellation consult (#96634 post-merge review).
_attempt_generation = _claim_compressor_attempt(agent.context_compressor)
_durable_cooldown_authoritative: Optional[bool] = None
_durable_cooldown_state: Optional[dict[str, Any]] = None
if (
@@ -2643,7 +2761,8 @@ def compress_context(
)
if not _codex_fence_entered:
_restore_compressor_attempt_state(
agent.context_compressor, _compressor_attempt_snapshot
agent.context_compressor, _compressor_attempt_snapshot,
attempt_generation=_attempt_generation,
)
existing_prompt = getattr(agent, "_cached_system_prompt", None)
if not existing_prompt:
@@ -3381,12 +3500,11 @@ def compress_context(
# removed in the finally below so it cannot leak into later attempts
# (e.g. a manual /compress force-clear).
if commit_fence is not None:
try:
agent.context_compressor._compression_cancelled_check = (
lambda: commit_fence.is_cancelled
)
except Exception:
pass
_install_compression_cancelled_check(
agent.context_compressor,
lambda: commit_fence.is_cancelled,
_attempt_generation,
)
# Incoming-message interrupts and active-turn redirects must not tear an
# atomic summary in half (#23975). Explicit stop surfaces set a separate
# Event atomically; never infer cause from the racy message fields.
@@ -3416,10 +3534,9 @@ def compress_context(
raise AuxiliaryExplicitCancellation()
finally:
if commit_fence is not None:
try:
agent.context_compressor._compression_cancelled_check = None
except Exception:
pass
_clear_compression_cancelled_check_if_owner(
agent.context_compressor, _attempt_generation
)
except AuxiliaryExplicitCancellation:
try:
_restore_compressor_attempt_state(
@@ -3427,6 +3544,7 @@ def compress_context(
_compressor_attempt_snapshot,
durable_cooldown_authoritative=_durable_cooldown_authoritative,
durable_cooldown_state=_durable_cooldown_state,
attempt_generation=_attempt_generation,
)
except BaseException as _rollback_exc:
# Compensation failure must surface, but it must not strand the
@@ -3594,6 +3712,7 @@ def compress_context(
_compressor_attempt_snapshot,
durable_cooldown_authoritative=_durable_cooldown_authoritative,
durable_cooldown_state=_durable_cooldown_state,
attempt_generation=_attempt_generation,
)
if (
messages_before_compression is not None
tests/agent/test_compression_attempt_ownership.py
+226
View File
@@ -0,0 +1,226 @@
"""Attempt-ownership guards for overlapping compression attempts (#96634).
The stall-fallback path (#78981, PR #96634) deliberately DETACHES a
timed-out primary compression worker — the fence cancel wins and the
future stays on the shared pool — then immediately runs a fallback attempt
against the SAME ``ContextCompressor``. donovan-yohan's post-merge
adversarial review identified two interleavings where the still-unwinding
primary clobbers fallback-owned state:
1. The late primary's unwind calls ``_restore_compressor_attempt_state``
with the PRIMARY's pre-attempt snapshot. Landing after the fallback's
commit, it rolls ``_previous_summary`` / cooldown / provenance /
telemetry back to pre-primary values.
2. ``_compression_cancelled_check`` is one shared attribute: the late
primary's ``finally`` clears the callback the fallback just installed.
Both are now guarded by a monotonic per-compressor attempt generation
(``_claim_compressor_attempt``): restores and callback set/clear are keyed
to the claiming generation and no-op when a newer attempt owns the
compressor. These tests drive both interleavings deterministically —
no timing, no threads.
"""
from types import SimpleNamespace
from agent.conversation_compression import (
_claim_compressor_attempt,
_clear_compression_cancelled_check_if_owner,
_compressor_attempt_is_current,
_install_compression_cancelled_check,
_restore_compressor_attempt_state,
_snapshot_compressor_attempt_state,
)
def _compressor(**overrides):
"""Bare compressor stand-in carrying only attempt-state fields."""
base = {
"_previous_summary": "primary-era summary",
"_summary_failure_cooldown_until": 0.0,
"_last_summary_error": None,
"_cooldown_persist_failed": False,
}
base.update(overrides)
return SimpleNamespace(**base)
class TestLatePrimaryRestoreAfterFallbackCommit:
"""Claim 1: a stale attempt's snapshot restore must no-op."""
def test_stale_restore_noops_and_preserves_fallback_state(self):
compressor = _compressor()
# Primary attempt starts: snapshot + claim.
primary_snapshot = _snapshot_compressor_attempt_state(compressor)
primary_gen = _claim_compressor_attempt(compressor)
# Primary "runs" and mutates state, then stalls; the host detaches
# it and starts the fallback attempt, which claims a NEWER generation
# and commits its own state.
compressor._previous_summary = "primary partial work"
fallback_gen = _claim_compressor_attempt(compressor)
assert fallback_gen > primary_gen
compressor._previous_summary = "FALLBACK COMMITTED SUMMARY"
compressor._summary_failure_cooldown_until = 123.0
# The detached primary finally unwinds and tries to restore its
# pre-attempt snapshot. Stale generation → must no-op.
_restore_compressor_attempt_state(
compressor, primary_snapshot, attempt_generation=primary_gen
)
assert compressor._previous_summary == "FALLBACK COMMITTED SUMMARY"
assert compressor._summary_failure_cooldown_until == 123.0
def test_current_attempt_restore_still_works(self):
"""The guard must not break the legitimate same-attempt rollback."""
compressor = _compressor()
snapshot = _snapshot_compressor_attempt_state(compressor)
gen = _claim_compressor_attempt(compressor)
compressor._previous_summary = "attempt scribbles"
_restore_compressor_attempt_state(
compressor, snapshot, attempt_generation=gen
)
assert compressor._previous_summary == "primary-era summary"
def test_legacy_callers_without_generation_are_unchanged(self):
"""attempt_generation=None preserves the historical always-restore."""
compressor = _compressor()
snapshot = _snapshot_compressor_attempt_state(compressor)
_claim_compressor_attempt(compressor) # someone else claims
compressor._previous_summary = "scribbles"
_restore_compressor_attempt_state(compressor, snapshot)
assert compressor._previous_summary == "primary-era summary"
def test_slotted_compressor_disables_guard_gracefully(self):
"""A compressor that rejects attribute writes yields generation 0
(guard off) rather than raising into the compression path."""
class Frozen:
__slots__ = ()
gen = _claim_compressor_attempt(Frozen())
assert gen == 0
# Generation 0 always reports current — legacy behavior.
assert _compressor_attempt_is_current(Frozen(), 0) is True
class TestCancelledCheckOwnership:
"""Claim 2: only the installing attempt may clear the shared callback."""
def test_stale_primary_finally_cannot_clear_fallback_callback(self):
compressor = _compressor()
primary_gen = _claim_compressor_attempt(compressor)
_install_compression_cancelled_check(
compressor, lambda: "primary", primary_gen
)
# Fallback claims and installs ITS callback while the primary is
# still unwinding.
fallback_gen = _claim_compressor_attempt(compressor)
fallback_check = lambda: "fallback" # noqa: E731
_install_compression_cancelled_check(
compressor, fallback_check, fallback_gen
)
# Detached primary's ``finally`` fires late — must be refused.
cleared = _clear_compression_cancelled_check_if_owner(
compressor, primary_gen
)
assert cleared is False
assert compressor._compression_cancelled_check is fallback_check
# The owner can still clear its own callback afterwards.
assert _clear_compression_cancelled_check_if_owner(
compressor, fallback_gen
) is True
assert compressor._compression_cancelled_check is None
def test_owner_clear_roundtrip(self):
compressor = _compressor()
gen = _claim_compressor_attempt(compressor)
check = lambda: True # noqa: E731
_install_compression_cancelled_check(compressor, check, gen)
assert compressor._compression_cancelled_check is check
assert _clear_compression_cancelled_check_if_owner(compressor, gen)
assert compressor._compression_cancelled_check is None
assert compressor._compression_cancelled_check_owner is None
def test_generation_zero_clear_is_unconditional(self):
"""Guard-disabled (legacy/slotted) attempts keep the old clear."""
compressor = _compressor()
gen = _claim_compressor_attempt(compressor)
_install_compression_cancelled_check(compressor, lambda: True, gen)
# generation 0 = guard disabled → clears like the historical code.
assert _clear_compression_cancelled_check_if_owner(compressor, 0)
assert compressor._compression_cancelled_check is None
class TestDigestCallsFollowThePinnedRoute:
"""Secondary item: lean digests must not return to the stalled route."""
def test_consumed_pin_still_routes_sibling_digest_calls(self):
import contextvars
def _probe():
from agent.context_compressor import (
attempt_summary_route_kwargs,
pin_summary_route,
take_pinned_summary_route,
)
route = {"provider": "fallback-prov", "model": "fallback-model"}
with pin_summary_route(route):
# Summary call consumes the pin (single-use preserved)...
consumed = take_pinned_summary_route()
assert consumed == route
assert take_pinned_summary_route() is None
# ...but sibling digest calls still see the attempt route.
kwargs = attempt_summary_route_kwargs()
assert kwargs.get("provider") == "fallback-prov"
assert kwargs.get("model") == "fallback-model"
return True
# Fresh context per test: the consumed echo must not leak in from
# any other test that touched the contextvars.
assert contextvars.copy_context().run(_probe) is True
def test_no_pin_means_no_route_override(self):
import contextvars
def _probe():
from agent.context_compressor import attempt_summary_route_kwargs
return attempt_summary_route_kwargs()
assert contextvars.copy_context().run(_probe) == {}
def test_consumed_echo_is_context_local(self):
"""The echo cannot leak into an unrelated attempt's context."""
import contextvars
def _consume_in_isolated_context():
from agent.context_compressor import (
pin_summary_route,
take_pinned_summary_route,
)
with pin_summary_route({"provider": "p", "model": "m"}):
take_pinned_summary_route()
ctx = contextvars.copy_context()
ctx.run(_consume_in_isolated_context)
# Outer context never saw the pin or its echo.
from agent.context_compressor import attempt_summary_route_kwargs
assert attempt_summary_route_kwargs() == {}