From e078b2fe7c02ae08902704f573e64ab68d57a70f Mon Sep 17 00:00:00 2001
From: kshitijk4poor <82637225+kshitijk4poor@users.noreply.github.com>
Date: Fri, 28 Aug 2026 12:47:53 +0530
Subject: [PATCH] fix(compression): close the restore TOCTOU; fold review
 findings
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

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.
---
 agent/conversation_compression.py             | 26 +++++++++++++--
 .../test_compression_attempt_ownership.py     | 32 +++++++++++++++++++
 2 files changed, 56 insertions(+), 2 deletions(-)

diff --git a/agent/conversation_compression.py b/agent/conversation_compression.py
index f55f70056e..08f70eb31c 100644
--- a/agent/conversation_compression.py
+++ b/agent/conversation_compression.py
@@ -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(
diff --git a/tests/agent/test_compression_attempt_ownership.py b/tests/agent/test_compression_attempt_ownership.py
index 3917ae127e..62d6c37f10 100644
--- a/tests/agent/test_compression_attempt_ownership.py
+++ b/tests/agent/test_compression_attempt_ownership.py
@@ -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"