From 98a8475d10371d3db3836f6a73d463f717f67656 Mon Sep 17 00:00:00 2001 From: Vito Sansevero Date: Sat, 17 Oct 2026 17:58:38 -0700 Subject: [PATCH] feat(history): bound context payloads beside lifecycle finalization --- tests/test_execution_context_budget.py | 115 +++++++++ tests/test_execution_history_writer.py | 307 +++++++++++++++++++++++++ utils/execution_context_budget.py | 111 +++++++++ utils/execution_contexts.py | 6 + utils/execution_history_writer.py | 137 ++++++++++- 5 files changed, 673 insertions(+), 3 deletions(-) create mode 100644 tests/test_execution_context_budget.py create mode 100644 utils/execution_context_budget.py diff --git a/tests/test_execution_context_budget.py b/tests/test_execution_context_budget.py new file mode 100644 index 0000000..55f4449 --- /dev/null +++ b/tests/test_execution_context_budget.py @@ -0,0 +1,115 @@ +"""Ownership and accounting behavior for shared context payloads.""" + +import unittest + + +class ContextBudgetTests(unittest.TestCase): + def test_reservation_transfer_shrink_and_idempotent_release(self): + from utils.execution_context_budget import ContextBudget + + pool = ContextBudget(100) + lease = pool.reserve(100) + self.assertIsNotNone(lease) + self.assertIsNone(pool.reserve(1)) + owner = lease.transfer() + lease.release() + self.assertEqual(pool.charged_bytes, 100) + with self.assertRaises(ValueError): + lease.transfer() + with self.assertRaises(ValueError): + lease.shrink_to(1) + owner.shrink_to(60) + self.assertEqual(owner.byte_count, 60) + self.assertEqual(pool.charged_bytes, 60) + for invalid in (-1, 61, True, 1.0): + with self.assertRaises(ValueError): + owner.shrink_to(invalid) + owner.release() + owner.release() + self.assertEqual(pool.charged_bytes, 0) + + def test_default_pool_and_invalid_reservations(self): + from utils.execution_context_budget import ContextBudget + + pool = ContextBudget() + lease = pool.reserve(8388608) + self.assertIsNotNone(lease) + self.assertIsNone(pool.reserve(1)) + lease.release() + for value in (-1, True, 1.0): + with self.assertRaises(ValueError): + pool.reserve(value) + with self.assertRaises(ValueError): + ContextBudget(value) + + def test_only_pool_issued_leases_can_charge_or_release(self): + from utils.execution_context_budget import ContextBudget, PayloadLease + + pool = ContextBudget(100) + with self.assertRaises(ValueError): + PayloadLease(pool, 100) + self.assertEqual(pool.charged_bytes, 0) + + def test_measured_payload_field_context_and_attempt_boundaries(self): + from dataclasses import replace + from utils.execution_contexts import ( + ContextBundle, + RecipeContext, + PromptSnapshot, + TextEvidence, + bundle_payload_size, + validate_bundle, + ) + + def recipe(anchor, fourth): + values = ( + "x" * 65536, + "x" * 65536, + "x" * 65536, + "x" * fourth, + "", + "", + "", + "", + ) + fields = tuple( + TextEvidence(v, "recorded", "runtime_return", None) for v in values + ) + return RecipeContext( + anchor, + PromptSnapshot(*fields[:4]), + PromptSnapshot(*fields[4:]), + "inferred", + ) + + # 8 * (8 provenance + 14 source) + 8 role + 1 anchor = 185. + context = recipe("1", 65351) + bundle = ContextBundle((context,), "complete", "captured") + self.assertEqual(bundle_payload_size(bundle), 262160) + self.assertIsNone(validate_bundle(bundle)) + with self.assertRaises(ValueError): + bundle_payload_size(replace(bundle, contexts=(recipe("1", 65352),))) + field = replace(context.positive.raw, value="é" * 32769) + bad = replace(context, positive=replace(context.positive, raw=field)) + with self.assertRaises(ValueError): + bundle_payload_size(replace(bundle, contexts=(bad,))) + attempt = replace( + bundle, contexts=tuple(recipe(str(i), 65347) for i in range(4)) + ) + self.assertEqual(bundle_payload_size(attempt), 1048576) + with self.assertRaises(ValueError): + bundle_payload_size( + replace(attempt, contexts=(recipe("0", 65348), *attempt.contexts[1:])) + ) + small = replace(context, positive=None, negative=None) + sixteen = replace( + bundle, contexts=tuple(replace(small, anchor_id=str(i)) for i in range(16)) + ) + validate_bundle(sixteen) + with self.assertRaises(ValueError): + bundle_payload_size( + replace( + sixteen, + contexts=(*sixteen.contexts, replace(small, anchor_id="17")), + ) + ) diff --git a/tests/test_execution_history_writer.py b/tests/test_execution_history_writer.py index 3ad2150..246f556 100644 --- a/tests/test_execution_history_writer.py +++ b/tests/test_execution_history_writer.py @@ -105,6 +105,313 @@ class HistoryWriterTests(unittest.TestCase): self.releases.append(release) return entered, release + def sealed(self, writer, value="exact text", generation=None): + from utils.execution_context_budget import SealedContext + from utils.execution_contexts import ( + ContextBundle, + RecipeContext, + PromptSnapshot, + TextEvidence, + bundle_payload_size, + ) + + field = TextEvidence(value, "recorded", "runtime_return", None) + snapshot = PromptSnapshot(field, field, field, field) + bundle = ContextBundle( + (RecipeContext("1", snapshot, snapshot, "inferred"),), + "complete", + "captured", + ) + lease = writer.context_budget.reserve(bundle_payload_size(bundle)) + return SealedContext( + bundle, + lease, + "complete", + "captured", + writer.status().capture_generation if generation is None else generation, + ) + + def test_context_success_releases_and_legacy_finish_preserves_capture(self): + writer = self.writer() + start = self.start_command() + sealed = self.sealed(writer) + self.assertTrue(writer.submit(start)) + self.assertTrue(writer.submit_finalization(self.finish_command(start), sealed)) + self.wait_for(lambda: writer.status().pending == 0) + self.assertEqual(writer.context_budget.charged_bytes, 0) + self.assertEqual( + self.db.get_execution_contexts(start.run_id)["state"], "complete" + ) + self.assertTrue(writer.submit(self.finish_command(start))) + self.wait_for(lambda: writer.status().pending == 0) + self.assertEqual( + self.db.get_execution_contexts(start.run_id)["state"], "complete" + ) + + def test_context_saturation_keeps_lifecycle_and_status(self): + from utils.execution_context_budget import SealedContext + from utils.execution_history_events import encode_command + + writer = self.writer() + lease = writer.context_budget.reserve(8388608) + self.assertIsNone(writer.context_budget.reserve(1)) + start = self.start_command() + self.assertTrue(writer.submit(start)) + finish = self.finish_command(start) + self.assertLessEqual(len(encode_command(finish)), 2048) + self.assertTrue( + writer.submit_finalization( + finish, SealedContext(None, None, "unavailable", "budget_exhausted") + ) + ) + self.wait_for(lambda: writer.status().pending == 0) + self.assertEqual( + self.db.get_execution_run(start.run_id)["outcome"], "succeeded" + ) + self.assertEqual( + self.db.get_execution_contexts(start.run_id)["reason"], "budget_exhausted" + ) + lease.release() + self.assertEqual(writer.context_budget.charged_bytes, 0) + + def test_outcome_survives_context_transaction_failure(self): + for error, ready in ( + (ValueError("invalid evidence"), True), + (sqlite3.OperationalError("unavailable"), False), + ): + with self.subTest(error=type(error)): + writer = self.writer() + start = self.start_command(10 if ready else 11) + self.assertTrue(writer.submit(start)) + sealed = self.sealed(writer) + entered, release = self.block() + + def fail_context(*args): + entered.set() + release.wait(3) + raise error + + with patch.object(self.db, "write_execution_contexts", fail_context): + self.assertTrue( + writer.submit_finalization(self.finish_command(start), sealed) + ) + self.assertTrue(entered.wait(1)) + self.assertEqual( + self.db.get_execution_run(start.run_id)["outcome"], "succeeded" + ) + self.assertGreater(writer.context_budget.charged_bytes, 0) + release.set() + self.wait_for(lambda: writer.status().pending == 0) + self.assertEqual( + self.db.get_execution_run(start.run_id)["outcome"], "succeeded" + ) + self.assertEqual(writer.status().accepting, ready) + self.assertEqual(writer.context_budget.charged_bytes, 0) + if not ready: + self.clock.advance(5) + self.wait_for(lambda: writer.status().accepting) + self.assertEqual( + self.db.get_execution_contexts(start.run_id)["state"], + "not_captured", + ) + self.assertTrue(writer.close()) + + def test_same_epoch_recovery_changes_capture_generation(self): + writer = self.writer() + start = self.start_command() + old_generation = writer.status().capture_generation + sealed = self.sealed(writer) + writer.invalidate() + self.clock.advance(5) + self.wait_for(lambda: writer.status().accepting) + self.assertEqual(writer.status().epoch, start.storage_epoch) + self.assertGreater(writer.status().capture_generation, old_generation) + self.assertTrue(writer.submit(start)) + self.assertTrue(writer.submit_finalization(self.finish_command(start), sealed)) + self.wait_for(lambda: writer.status().pending == 0) + self.assertEqual( + self.db.get_execution_run(start.run_id)["outcome"], "succeeded" + ) + self.assertEqual(writer.context_budget.charged_bytes, 0) + self.assertEqual(self.db.get_execution_contexts(start.run_id)["contexts"], []) + + def test_restore_between_transactions_rejects_context(self): + writer = self.writer() + start = self.start_command() + self.db.write_execution_start(start) + backup = os.path.join(self.directory.name, "between.db") + self.assertTrue(self.db.model.backup_database(backup)) + original = self.db.write_execution_finish + + def finish(command): + result = original(command) + self.db.model.restore_from_file(backup) + return result + + with patch.object(self.db, "write_execution_finish", finish): + self.assertTrue( + writer.submit_finalization( + self.finish_command(start), self.sealed(writer) + ) + ) + self.wait_for(lambda: writer.status().pending == 0) + self.assertEqual( + self.db.get_execution_contexts(start.run_id)["state"], "not_captured" + ) + self.assertEqual(writer.context_budget.charged_bytes, 0) + self.assertFalse(writer.status().accepting) + + def test_reject_discard_and_bounded_close_release_leases(self): + writer = self.writer() + start = self.start_command() + sealed = self.sealed(writer) + self.assertFalse(writer.submit_finalization(object(), sealed)) + self.assertEqual(writer.context_budget.charged_bytes, 0) + entered, release = self.block() + original = self.db.write_execution_start + + def write(command): + entered.set() + release.wait(3) + return original(command) + + with patch.object(self.db, "write_execution_start", write): + self.assertTrue(writer.submit(start)) + self.assertTrue(entered.wait(1)) + for _ in range(1023): + self.assertTrue( + writer.submit_finalization( + self.finish_command(start), self.sealed(writer) + ) + ) + self.assertEqual(writer.status().pending, 1024) + self.assertFalse( + writer.submit_finalization( + self.finish_command(start), self.sealed(writer) + ) + ) + self.assertFalse(writer.close(0.01)) + release.set() + self.wait_for(lambda: not writer.status().alive) + self.assertEqual(writer.context_budget.charged_bytes, 0) + self.assertEqual( + self.db.get_execution_contexts(start.run_id)["state"], "not_captured" + ) + + def test_bounded_close_keeps_inflight_charge_and_releases_queued(self): + from utils.execution_context_budget import ContextBudget + + pool = ContextBudget(1000) + writer = HistoryWriter(lambda: self.db, clock=self.clock, context_budget=pool) + self.writers.append(writer) + writer.start() + self.wait_for(lambda: writer.status().accepting) + start = self.start_command() + self.db.write_execution_start(start) + entered, release = self.block() + original = self.db.write_execution_contexts + first = self.sealed(writer) + charge = first.lease.byte_count + + def write(*args): + entered.set() + release.wait(3) + return original(*args) + + with patch.object(self.db, "write_execution_contexts", write): + self.assertTrue( + writer.submit_finalization(self.finish_command(start), first) + ) + self.assertTrue(entered.wait(1)) + self.assertTrue( + writer.submit_finalization( + self.finish_command(start), self.sealed(writer) + ) + ) + before = time.monotonic() + self.assertFalse(writer.close(0.02)) + self.assertLess(time.monotonic() - before, 0.2) + self.assertEqual(pool.charged_bytes, charge) + replacement = HistoryWriter(lambda: self.db, context_budget=pool) + self.writers.append(replacement) + self.assertIsNone(replacement.context_budget.reserve(1000)) + release.set() + self.wait_for(lambda: not writer.status().alive) + self.assertEqual(pool.charged_bytes, 0) + + def test_same_epoch_invalidation_between_transactions_drops_payload(self): + writer = self.writer() + start = self.start_command() + self.db.write_execution_start(start) + original = self.db.write_execution_finish + + def finish(command): + result = original(command) + writer.invalidate() + return result + + with patch.object(self.db, "write_execution_finish", finish): + self.assertTrue( + writer.submit_finalization( + self.finish_command(start), self.sealed(writer) + ) + ) + self.wait_for(lambda: writer.status().pending == 0) + self.assertEqual( + self.db.get_execution_run(start.run_id)["outcome"], "succeeded" + ) + self.assertEqual( + self.db.get_execution_contexts(start.run_id)["state"], "not_captured" + ) + self.assertEqual(writer.context_budget.charged_bytes, 0) + self.clock.advance(5) + self.wait_for(lambda: writer.status().accepting) + self.assertEqual( + self.db.get_execution_contexts(start.run_id)["state"], "not_captured" + ) + + def test_invalid_seals_release_and_keep_valid_outcome(self): + from utils.execution_context_budget import ContextBudget + + writer = self.writer() + cases = [ + dict(state="private text" * 100), + dict(reason="arbitrary_exception"), + dict(capture_generation=True), + dict(capture_generation=-1), + dict(capture_generation=None), + dict(lease=None), + ] + for item, changes in enumerate(cases, 20): + start = self.start_command(item) + self.db.write_execution_start(start) + sealed = self.sealed(writer) + malformed = replace(sealed, **changes) + if changes == {"lease": None}: + sealed.lease.release() + self.assertTrue( + writer.submit_finalization(self.finish_command(start), malformed) + ) + self.wait_for(lambda: writer.status().pending == 0) + self.assertEqual( + self.db.get_execution_run(start.run_id)["outcome"], "succeeded" + ) + self.assertEqual( + self.db.get_execution_contexts(start.run_id)["reason"], + "invalid_evidence", + ) + self.assertEqual(writer.context_budget.charged_bytes, 0) + foreign = ContextBudget(1000) + sealed = self.sealed(writer) + sealed.lease.release() + sealed = replace(sealed, lease=foreign.reserve(1000)) + start = self.start_command(30) + self.db.write_execution_start(start) + self.assertTrue(writer.submit_finalization(self.finish_command(start), sealed)) + self.wait_for(lambda: writer.status().pending == 0) + self.assertEqual(foreign.charged_bytes, 0) + self.assertEqual(self.db.get_execution_contexts(start.run_id)["contexts"], []) + def test_fifo_single_worker_and_thread_local_connection_cleanup(self): caller_conn = self.db.model.get_connection() threads, timeouts = [], [] diff --git a/utils/execution_context_budget.py b/utils/execution_context_budget.py new file mode 100644 index 0000000..0b319fa --- /dev/null +++ b/utils/execution_context_budget.py @@ -0,0 +1,111 @@ +"""Shared payload reservations with explicit, single-owner transport leases.""" + +from dataclasses import dataclass +import json +import threading + +from .execution_contexts import BUNDLE_STATES, CAPTURE_REASONS, ContextBundle + + +def _byte_count(value: int) -> None: + if type(value) is not int or value < 0: + raise ValueError("Byte count must be a nonnegative integer") + + +class ContextBudget: + def __init__(self, limit_bytes: int = 8388608) -> None: + _byte_count(limit_bytes) + self._limit = limit_bytes + self._charged = 0 + self._lock = threading.Lock() + + @property + def charged_bytes(self) -> int: + with self._lock: + return self._charged + + def reserve(self, byte_count: int) -> "PayloadLease | None": + _byte_count(byte_count) + with self._lock: + if byte_count > self._limit - self._charged: + return None + self._charged += byte_count + return PayloadLease._new(self, byte_count) + + +class PayloadLease: + def __init__(self, pool: ContextBudget, byte_count: int) -> None: + raise ValueError("Reserve payload leases through ContextBudget") + + @classmethod + def _new(cls, pool: ContextBudget, byte_count: int) -> "PayloadLease": + owner = object.__new__(cls) + owner._pool = pool + owner._bytes = byte_count + owner._active = True + return owner + + @property + def byte_count(self) -> int: + with self._pool._lock: + self._require_owner() + return self._bytes + + def _require_owner(self) -> None: + if not self._active: + raise ValueError("Inactive payload lease") + + def transfer(self) -> "PayloadLease": + with self._pool._lock: + self._require_owner() + owner = PayloadLease._new(self._pool, self._bytes) + self._active = False + return owner + + def shrink_to(self, byte_count: int) -> None: + _byte_count(byte_count) + with self._pool._lock: + self._require_owner() + if byte_count > self._bytes: + raise ValueError("Cannot grow payload lease") + self._pool._charged -= self._bytes - byte_count + self._bytes = byte_count + + def release(self) -> None: + with self._pool._lock: + if self._active: + self._pool._charged -= self._bytes + self._active = False + + +SHARED_CONTEXT_BUDGET = ContextBudget() + + +@dataclass(frozen=True) +class SealedContext: + bundle: ContextBundle | None + lease: PayloadLease | None + state: str + reason: str + capture_generation: int | None = None + + def validate_header(self) -> None: + if ( + type(self.state) is not str + or self.state not in BUNDLE_STATES + or type(self.reason) is not str + or self.reason not in CAPTURE_REASONS + or ( + self.capture_generation is not None + and ( + type(self.capture_generation) is not int + or not 0 <= self.capture_generation <= 2**63 - 1 + ) + ) + ): + raise ValueError("Invalid capture status") + header = json.dumps( + [self.state, self.reason, self.capture_generation], separators=(",", ":") + ).encode("utf-8") + if len(header) > 256: + raise ValueError("Oversized capture status") diff --git a/utils/execution_contexts.py b/utils/execution_contexts.py index 1a1a492..64c537f 100644 --- a/utils/execution_contexts.py +++ b/utils/execution_contexts.py @@ -129,6 +129,11 @@ EVIDENCE_SOURCES = { def validate_bundle(bundle: ContextBundle) -> None: + bundle_payload_size(bundle) + + +def bundle_payload_size(bundle: ContextBundle) -> int: + """Validate scalar contracts and measure the charged UTF-8 payload.""" if ( type(bundle) is not ContextBundle or type(bundle.format_version) is not int @@ -193,3 +198,4 @@ def validate_bundle(bundle: ContextBundle) -> None: total += size if total > MAX_ATTEMPT_BYTES: raise ValueError("Attempt payload limit reached") + return total diff --git a/utils/execution_history_writer.py b/utils/execution_history_writer.py index 3cba705..4122cea 100644 --- a/utils/execution_history_writer.py +++ b/utils/execution_history_writer.py @@ -6,6 +6,7 @@ import queue import threading import time from typing import Callable +from uuid import UUID, uuid5 try: from ..database import PromptDatabase @@ -20,7 +21,14 @@ except ImportError: subscribe_history_restore, ) -from .execution_history_events import RunCommand, RunStart, encode_command +from .execution_context_budget import ( + SHARED_CONTEXT_BUDGET, + ContextBudget, + PayloadLease, + SealedContext, +) +from .execution_contexts import ContextBundle, bundle_payload_size +from .execution_history_events import RunCommand, RunFinish, RunStart, encode_command from .logging_config import get_logger _logger = get_logger("prompt_manager.execution_history_writer") @@ -36,6 +44,13 @@ class HistoryWriterStatus: pending: int recording_gaps: int alive: bool + capture_generation: int + + +@dataclass(frozen=True) +class _Finalization: + command: RunFinish + context: SealedContext class HistoryWriter: @@ -46,12 +61,18 @@ class HistoryWriter: db_factory: Callable[[], PromptDatabase], *, clock: Callable[[], float] = time.monotonic, + context_budget: ContextBudget | None = None, ) -> None: + self.context_budget = ( + SHARED_CONTEXT_BUDGET if context_budget is None else context_budget + ) self._db_factory = db_factory self._clock = clock self._lock = threading.Lock() self._wake = threading.Event() - self._queue: queue.Queue[RunCommand] = queue.Queue(maxsize=_MAX_PENDING) + self._queue: queue.Queue[RunCommand | _Finalization] = queue.Queue( + maxsize=_MAX_PENDING + ) self._thread: threading.Thread | None = None self._epoch: str | None = None self._accepting = False @@ -94,6 +115,69 @@ class HistoryWriter: self._wake.set() return accepted + @staticmethod + def _release_context(context: SealedContext) -> None: + if isinstance(context, SealedContext) and type(context.lease) is PayloadLease: + context.lease.release() + + def _consume_context_locked(self, context: SealedContext) -> SealedContext: + try: + if type(context) is not SealedContext: + raise ValueError("Invalid sealed context") + context.validate_header() + if context.bundle is None: + if context.lease is not None: + raise ValueError("Unexpected payload lease") + return context + size = bundle_payload_size(context.bundle) + if ( + context.capture_generation != self._generation + or context.capture_generation is None + or type(context.lease) is not PayloadLease + or context.lease._pool is not self.context_budget + or context.lease.byte_count < size + or context.state != context.bundle.state + or context.reason != context.bundle.reason + ): + raise ValueError("Invalid context ownership or generation") + owner = context.lease.transfer() + return SealedContext( + context.bundle, + owner, + context.state, + context.reason, + context.capture_generation, + ) + except ValueError: + self._release_context(context) + return SealedContext(None, None, "unavailable", "invalid_evidence") + + def submit_finalization(self, command: RunFinish, context: SealedContext) -> bool: + try: + if type(command) is not RunFinish: + raise ValueError("Finalization requires finish") + encode_command(command) + except ValueError: + self._release_context(context) + self.record_gap() + return False + with self._lock: + if ( + not self._accepting + or command.storage_epoch != self._epoch + or self._pending >= _MAX_PENDING + ): + self._release_context(context) + self._record_gap_locked() + accepted = False + else: + owned = self._consume_context_locked(context) + self._queue.put_nowait(_Finalization(command, owned)) + self._pending += 1 + accepted = True + self._wake.set() + return accepted + def _record_gap_locked(self) -> None: self._gaps = min(_MAX_GAPS, self._gaps + 1) @@ -120,6 +204,7 @@ class HistoryWriter: self._pending, self._gaps, self._thread is not None and self._thread.is_alive(), + self._generation, ) def close(self, timeout: float = 2.0) -> bool: @@ -135,14 +220,19 @@ class HistoryWriter: self._wake.set() if thread is not None: thread.join(max(0.0, deadline - time.monotonic())) + with self._lock: + if self._closing and time.monotonic() >= self._close_deadline: + self._discard_locked() return thread is None or not thread.is_alive() def _discard_locked(self) -> None: while True: try: - self._queue.get_nowait() + item = self._queue.get_nowait() except queue.Empty: return + if isinstance(item, _Finalization): + self._release_context(item.context) self._pending -= 1 self._record_gap_locked() @@ -200,6 +290,10 @@ class HistoryWriter: self._record_gap_locked() if command is not None: failed = False + context = None + if isinstance(command, _Finalization): + context = command.context + command = command.command try: # Restore may have reopened this thread's managed connection. db.model.get_connection().execute("PRAGMA busy_timeout=250") @@ -207,12 +301,49 @@ class HistoryWriter: db.write_execution_start(command) else: db.write_execution_finish(command) + if context is not None: + with self._lock: + current = ( + generation == self._generation + and command.storage_epoch == self._epoch + and ( + context.capture_generation is None + or context.capture_generation + == self._generation + ) + ) + if current: + try: + db.write_execution_contexts( + str( + uuid5( + UUID(command.session_id), + str(command.queue_item_id), + ) + ), + command.storage_epoch, + ( + context.bundle + if context.bundle is not None + else ContextBundle( + (), context.state, context.reason + ) + ), + ) + except ValueError: + # Validation/collision cannot revoke the outcome. + with self._lock: + self._record_gap_locked() except HistoryEpochChanged: # The exception's storage_epoch is only a hint; readiness # comes from a later probe and never replays this command. failed = True except Exception: failed = True + finally: + if context is not None: + self._release_context(context) + context = None with self._lock: self._pending -= 1 if failed: