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:
@@ -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"
|
||||
|
||||
Reference in New Issue
Block a user