Import AITURK IDE 1.0.0-beta.1 from Hermes 63279301; preserve MIT license

tests/agent/test_compression_attempt_ownership.py

@@ -0,0 +1,244 @@
+"""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 TestSummaryRoutePinSingleUse:
+    """The pin is single-use: consumed by the one summary call per attempt."""
+
+    def test_pin_is_consumed_once(self):
+        import contextvars
+
+        def _probe():
+            from agent.context_compressor import (
+                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
+                # ...and the main-model retry never re-issues it.
+                assert take_pinned_summary_route() is None
+            return True
+
+        # Fresh context per test: state must not leak in from any other
+        # test that touched the contextvars.
+        assert contextvars.copy_context().run(_probe) is True
+
+    def test_consume_is_context_local(self):
+        """A consume in one context cannot leak into an unrelated attempt's."""
+        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.
+        from agent.context_compressor import take_pinned_summary_route
+
+        assert take_pinned_summary_route() is None
+
+
+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"