diff --git a/tests/test_execution_context_budget.py b/tests/test_execution_context_budget.py index fd38ab9..17a7f81 100644 --- a/tests/test_execution_context_budget.py +++ b/tests/test_execution_context_budget.py @@ -1,6 +1,8 @@ """Ownership and accounting behavior for shared context payloads.""" import unittest +from copy import copy +from unittest.mock import patch class ContextBudgetTests(unittest.TestCase): @@ -158,3 +160,66 @@ class ContextBudgetTests(unittest.TestCase): if lease is not None: lease.release() self.assertEqual(pool.charged_bytes, 0) + + def test_copied_and_forged_handles_cannot_spend_or_release_genuine_credits(self): + from utils.execution_context_budget import ContextBudget, PayloadLease + + pool = ContextBudget(100) + genuine = pool.reserve(80) + forged = object.__new__(PayloadLease) + forged.__dict__.update(genuine.__dict__) + for impostor in (copy(genuine), forged): + with self.subTest(impostor=impostor): + with self.assertRaises(ValueError): + _ = impostor.byte_count + with self.assertRaises(ValueError): + impostor.transfer() + with self.assertRaises(ValueError): + impostor.shrink_to(1) + impostor.release() + self.assertEqual(pool.charged_bytes, 80) + self.assertEqual(genuine.byte_count, 80) + self.assertIsNone(pool.reserve(21)) + owner = genuine.transfer() + genuine.release() + owner.shrink_to(20) + self.assertEqual(pool.charged_bytes, 20) + owner.release() + self.assertEqual(pool.charged_bytes, 0) + last = pool.reserve(100) + self.assertIsNotNone(last) + last.release() + self.assertEqual(pool.charged_bytes, 0) + + def test_failed_owner_allocation_or_registration_never_orphans_credits(self): + from utils.execution_context_budget import ContextBudget, PayloadLease + + pool = ContextBudget(100) + genuine = pool.reserve(20) + with patch.object(PayloadLease, "_new", side_effect=MemoryError): + for operation in ( + lambda: pool.reserve(10), + lambda: pool.reserve_up_to(10), + genuine.transfer, + ): + with self.assertRaises(MemoryError): + operation() + self.assertEqual(pool.charged_bytes, 20) + self.assertEqual(genuine.byte_count, 20) + + class DeniedOwners(set): + def add(self, owner): + raise MemoryError + + pool._owners = DeniedOwners(pool._owners) + for operation in ( + lambda: pool.reserve(10), + lambda: pool.reserve_up_to(10), + genuine.transfer, + ): + with self.assertRaises(MemoryError): + operation() + self.assertEqual(pool.charged_bytes, 20) + self.assertEqual(genuine.byte_count, 20) + genuine.release() + self.assertEqual(pool.charged_bytes, 0) diff --git a/tests/test_execution_context_capture.py b/tests/test_execution_context_capture.py index db0214d..a1151cd 100644 --- a/tests/test_execution_context_capture.py +++ b/tests/test_execution_context_capture.py @@ -73,6 +73,31 @@ class CaptureTests(unittest.TestCase): self.addCleanup(result.lease.release) return result + def test_unexpected_projection_error_releases_payload_but_keeps_capability(self): + with patch.object( + capture, "project_contexts", side_effect=RuntimeError("private") + ): + cap = self.begin() + self.assertTrue(self.manager.valid(cap)) + self.assertEqual(self.pool.charged_bytes, 0) + result = self.seal(cap) + self.assertEqual( + (result.state, result.reason), ("unavailable", "invalid_evidence") + ) + self.assertEqual(self.pool.charged_bytes, 0) + + def test_unexpected_seal_error_discards_popped_attempt_and_lease(self): + cap = self.begin() + with patch.object( + capture, "bundle_payload_size", side_effect=RuntimeError("private") + ): + result = self.seal(cap) + self.assertEqual( + (result.state, result.reason), ("unavailable", "invalid_evidence") + ) + self.assertFalse(self.manager.valid(cap)) + self.assertEqual(self.pool.charged_bytes, 0) + def test_runtime_keeps_raw_affix_whitespace_and_injected_effective_text(self): cap = self.begin() self.manager.observe(cap, SourceRef("p", "PromptManagerPair", 0), 0, snapshot()) diff --git a/utils/execution_context_budget.py b/utils/execution_context_budget.py index 93d3a36..27ec1a6 100644 --- a/utils/execution_context_budget.py +++ b/utils/execution_context_budget.py @@ -17,6 +17,7 @@ class ContextBudget: _byte_count(limit_bytes) self._limit = limit_bytes self._charged = 0 + self._owners: set[PayloadLease] = set() self._lock = threading.Lock() @property @@ -29,8 +30,10 @@ class ContextBudget: with self._lock: if byte_count > self._limit - self._charged: return None + owner = PayloadLease._new(self, byte_count) + self._owners.add(owner) self._charged += byte_count - return PayloadLease._new(self, byte_count) + return owner def reserve_up_to(self, max_bytes: int) -> "PayloadLease | None": """Atomically own up to the remaining capacity in this same pool.""" @@ -39,8 +42,10 @@ class ContextBudget: byte_count = min(max_bytes, self._limit - self._charged) if not byte_count: return None + owner = PayloadLease._new(self, byte_count) + self._owners.add(owner) self._charged += byte_count - return PayloadLease._new(self, byte_count) + return owner class PayloadLease: @@ -62,13 +67,15 @@ class PayloadLease: return self._bytes def _require_owner(self) -> None: - if not self._active: - raise ValueError("Inactive payload lease") + if self not in self._pool._owners or not self._active: + raise ValueError("Inactive or unissued payload lease") def transfer(self) -> "PayloadLease": with self._pool._lock: self._require_owner() owner = PayloadLease._new(self._pool, self._bytes) + self._pool._owners.add(owner) + self._pool._owners.remove(self) self._active = False return owner @@ -83,7 +90,8 @@ class PayloadLease: def release(self) -> None: with self._pool._lock: - if self._active: + if self in self._pool._owners: + self._pool._owners.remove(self) self._pool._charged -= self._bytes self._active = False diff --git a/utils/execution_context_capture.py b/utils/execution_context_capture.py index 05cfaee..3a07985 100644 --- a/utils/execution_context_capture.py +++ b/utils/execution_context_capture.py @@ -187,6 +187,12 @@ class ContextCapture: projection = GraphProjection((), (), "") reason = "graph_limit" retained = 0 + except Exception: + projection = GraphProjection((), (), "") + reason = "invalid_evidence" + retained = 0 + lease.release() + lease = None cap = replace(cap, sources=projection.sources) self._attempts[cap.nonce] = _Attempt( cap, projection, lease, {}, retained, reason @@ -396,6 +402,10 @@ class ContextCapture: self._discard_payload(attempt) lease.release() return SealedContext(None, None, "unavailable", "limit_reached") + except Exception: + self._discard_payload(attempt) + lease.release() + return SealedContext(None, None, "unavailable", "invalid_evidence") def record_runtime_context(