fix(compression): close the restore TOCTOU; fold review findings

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.
This commit is contained in:
kshitijk4poor
2026-08-28 12:47:53 +05:30
committed by kshitij
parent 61cd299c6e
commit e078b2fe7c
2 changed files with 56 additions and 2 deletions
+24 -2
View File
@@ -369,6 +369,9 @@ def _claim_compressor_attempt(compressor: Any) -> int:
except Exception:
# Slotted/frozen third-party compressor: ownership tracking is
# unavailable; generation 0 disables the guard (legacy behavior).
# Per-compressor all-or-nothing, NOT per-attempt: a compressor
# that rejects the setattr rejects it for EVERY claim, so gen-0
# and gen>0 attempts can never coexist on one instance.
return 0
return generation
@@ -514,8 +517,27 @@ def _restore_compressor_attempt_state(
exc_info=True,
)
restored = copy.deepcopy(snapshot)
for name, value in restored.items():
setattr(compressor, name, value)
# Close the check-then-act window: the entry check above runs before the
# (potentially slow) durable-cooldown rollback, so a fallback attempt can
# claim the compressor in between. Re-validate and write the in-memory
# fields under the SAME lock claims are taken under — a stale attempt can
# never interleave its setattr loop with a newer attempt's writes. The
# durable rollback above is safe either way: the dangerous direction
# (restore landing AFTER the fallback's writes) requires the fallback to
# have claimed first, which the entry check already rejects.
with _COMPRESSOR_ATTEMPT_LOCK:
if attempt_generation is not None and attempt_generation and (
int(getattr(compressor, "_compression_attempt_generation", 0) or 0)
!= attempt_generation
):
logger.warning(
"Skipping stale compressor attempt-state restore at write "
"time: attempt generation %s lost the compressor mid-restore.",
attempt_generation,
)
return
for name, value in restored.items():
setattr(compressor, name, value)
def _capture_authoritative_cooldown_under_lease(
@@ -224,3 +224,35 @@ class TestDigestCallsFollowThePinnedRoute:
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"