kshitijk4poor
e078b2fe7c
Post-review hardening on the attempt-ownership commit: - Write-time re-validation: the entry staleness check in _restore_compressor_attempt_state runs before the durable-cooldown DB I/O, so a fallback could claim the compressor in that window and the stale setattr loop would still clobber its state. The in-memory writes now re-validate AND execute under _COMPRESSOR_ATTEMPT_LOCK — the same lock claims are taken under. The DB rollback stays outside the lock (safe: the dangerous direction requires a prior claim, which the entry check rejects). Both the quality reviewer and the lead's independent pre-verification converged on this window. New deterministic test: TestMidRestoreClaimRace (claim injected between entry check and write via instrumented deepcopy). - Documented gen-0 semantics on _claim_compressor_attempt: per-compressor all-or-nothing, never mixed with gen>0 on one instance (reviewer finding 2, verified unreachable — comment hardens against future confusion). Dropped after verification (reviewer finding 3): resetting _SUMMARY_ROUTE_CONSUMED on pin_summary_route exit — the echo lives in the worker thread's COPIED context (propagate_context_to_thread) and dies with it; a probe confirmed the next attempt's context is clean. Resetting it would break digests running after the with-block.
259 lines
10 KiB
Python
259 lines
10 KiB
Python
"""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() == {}
|
|
|
|
|
|
class TestMidRestoreClaimRace:
|
|
"""TOCTOU: a claim landing between entry check and write must void the restore."""
|
|
|
|
def test_claim_during_restore_body_voids_the_write(self, monkeypatch):
|
|
import agent.conversation_compression as cc
|
|
|
|
compressor = _compressor()
|
|
snapshot = _snapshot_compressor_attempt_state(compressor)
|
|
primary_gen = _claim_compressor_attempt(compressor)
|
|
compressor._previous_summary = "FALLBACK STATE"
|
|
|
|
# Interleave deterministically: the fallback claims the compressor
|
|
# while the primary is inside the restore body (during deepcopy of
|
|
# the snapshot, i.e. after the entry check passed).
|
|
real_deepcopy = cc.copy.deepcopy
|
|
state = {"claimed": False}
|
|
|
|
def claiming_deepcopy(obj, *a, **kw):
|
|
if not state["claimed"] and obj is snapshot:
|
|
state["claimed"] = True
|
|
_claim_compressor_attempt(compressor) # fallback arrives NOW
|
|
return real_deepcopy(obj, *a, **kw)
|
|
|
|
monkeypatch.setattr(cc.copy, "deepcopy", claiming_deepcopy)
|
|
_restore_compressor_attempt_state(
|
|
compressor, snapshot, attempt_generation=primary_gen
|
|
)
|
|
|
|
# The write-time re-check must have refused the stale restore.
|
|
assert compressor._previous_summary == "FALLBACK STATE"
|