diff --git a/agent/context_compressor.py b/agent/context_compressor.py index 73a7060eef..1d7f9d626e 100644 --- a/agent/context_compressor.py +++ b/agent/context_compressor.py @@ -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 diff --git a/agent/conversation_compression.py b/agent/conversation_compression.py index ee27095cd3..f55f70056e 100644 --- a/agent/conversation_compression.py +++ b/agent/conversation_compression.py @@ -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 diff --git a/tests/agent/test_compression_attempt_ownership.py b/tests/agent/test_compression_attempt_ownership.py new file mode 100644 index 0000000000..3917ae127e --- /dev/null +++ b/tests/agent/test_compression_attempt_ownership.py @@ -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() == {}