diff --git a/database/execution_contexts.py b/database/execution_contexts.py index 6f979df..4771f04 100644 --- a/database/execution_contexts.py +++ b/database/execution_contexts.py @@ -1,10 +1,12 @@ """Immutable execution capture and exact, epoch-fenced library navigation hints.""" +from contextlib import ExitStack import json import sqlite3 -from typing import Iterator +from typing import Callable, ContextManager, Iterator try: + from ..utils.execution_history_events import RunFinish, encode_command from ..utils.execution_contexts import ( CAPTURE_REASONS, ContextBundle, @@ -12,6 +14,7 @@ try: validate_bundle, ) except ImportError: + from utils.execution_history_events import RunFinish, encode_command from utils.execution_contexts import ( CAPTURE_REASONS, ContextBundle, @@ -183,6 +186,77 @@ def _values(conn: sqlite3.Connection, snapshot_id: str | None) -> dict | None: return dict(zip(FIELDS, row)) +def _write_capture( + conn: sqlite3.Connection, + run_id: str, + expected_epoch: str, + bundle: ContextBundle, + prepared: list[dict], +) -> bool: + capture = conn.execute( + "SELECT state,reason,format_version FROM execution_context_capture " + "WHERE run_id=?", + (run_id,), + ).fetchone() + if capture is not None: + for item in prepared: + for table, columns, values in _content_rows(item): + _existing_row(conn, table, columns, values) + links = conn.execute( + "SELECT anchor_id,context_id,evidence_json FROM " + "execution_run_contexts WHERE run_id=? ORDER BY anchor_id", + (run_id,), + ).fetchall() + expected = [ + (item["anchor_id"], item["context_id"], item["evidence_json"]) + for item in prepared + ] + # Compare JSON objects so serialization order never creates a conflict. + original = [(row[0], row[1], json.loads(row[2])) for row in links] + proposed = [ + (anchor, context, json.loads(evidence)) + for anchor, context, evidence in expected + ] + if capture[0] != "conflicting" and ( + tuple(capture) != (bundle.state, bundle.reason, bundle.format_version) + or original != proposed + ): + conn.execute( + "UPDATE execution_context_capture SET " + "state='conflicting',reason='conflicting' WHERE run_id=?", + (run_id,), + ) + return True + for item in prepared: + for table, columns, values in _content_rows(item): + _insert_exact(conn, table, columns, values) + raw = [snapshot[1][0] if snapshot else None for snapshot in item["snapshots"]] + positive_id, negative_id, pair_id = _hints(conn, *raw) + conn.execute( + "INSERT INTO " + "execution_run_contexts(run_id,anchor_id,context_id,evidence_json," + "positive_prompt_id,negative_prompt_id,pair_id,reference_epoch," + "capture_version) VALUES(?,?,?,?,?,?,?,?,1)", + ( + run_id, + item["anchor_id"], + item["context_id"], + item["evidence_json"], + positive_id, + negative_id, + pair_id, + expected_epoch, + ), + ) + conn.execute( + "INSERT INTO " + "execution_context_capture(run_id,state,reason,format_version) " + "VALUES(?,?,?,?)", + (run_id, bundle.state, bundle.reason, bundle.format_version), + ) + return True + + class ExecutionContextOperations: """Compose with ExecutionHistoryOperations on PromptDatabase.""" @@ -199,72 +273,36 @@ class ExecutionContextOperations: is None ): return False - capture = conn.execute( - "SELECT state,reason,format_version FROM execution_context_capture " - "WHERE run_id=?", - (run_id,), - ).fetchone() - if capture is not None: - for item in prepared: - for table, columns, values in _content_rows(item): - _existing_row(conn, table, columns, values) - links = conn.execute( - "SELECT anchor_id,context_id,evidence_json FROM " - "execution_run_contexts WHERE run_id=? ORDER BY anchor_id", - (run_id,), - ).fetchall() - expected = [ - (item["anchor_id"], item["context_id"], item["evidence_json"]) - for item in prepared - ] - # Compare JSON objects so serialization order never creates a conflict. - original = [(row[0], row[1], json.loads(row[2])) for row in links] - proposed = [ - (anchor, context, json.loads(evidence)) - for anchor, context, evidence in expected - ] - if capture[0] != "conflicting" and ( - tuple(capture) - != (bundle.state, bundle.reason, bundle.format_version) - or original != proposed + return _write_capture(conn, run_id, expected_epoch, bundle, prepared) + + def write_execution_contexts_for_finish( + self, + command: RunFinish, + bundle: ContextBundle, + *, + generation_guard: Callable[[], ContextManager[bool]] | None = None, + ) -> bool: + encode_command(command) + if type(command) is not RunFinish: + raise ValueError("Expected execution finish command") + prepared = _prepare(bundle) + # Enter the guard after BEGIN, and keep it until COMMIT/ROLLBACK exits. + with ExitStack() as guards: + with self._execution_transaction(command.storage_epoch) as conn: + if generation_guard is not None and not guards.enter_context( + generation_guard() ): - conn.execute( - "UPDATE execution_context_capture SET " - "state='conflicting',reason='conflicting' WHERE run_id=?", - (run_id,), - ) - return True - for item in prepared: - for table, columns, values in _content_rows(item): - _insert_exact(conn, table, columns, values) - raw = [ - snapshot[1][0] if snapshot else None - for snapshot in item["snapshots"] - ] - positive_id, negative_id, pair_id = _hints(conn, *raw) - conn.execute( - "INSERT INTO " - "execution_run_contexts(run_id,anchor_id,context_id,evidence_json," - "positive_prompt_id,negative_prompt_id,pair_id,reference_epoch," - "capture_version) VALUES(?,?,?,?,?,?,?,?,1)", - ( - run_id, - item["anchor_id"], - item["context_id"], - item["evidence_json"], - positive_id, - negative_id, - pair_id, - expected_epoch, - ), + return False + row = conn.execute( + "SELECT run_id FROM execution_runs WHERE session_id=? " + "AND queue_item_id=? AND storage_epoch=?", + (command.session_id, command.queue_item_id, command.storage_epoch), + ).fetchone() + if row is None: + return False + return _write_capture( + conn, row[0], command.storage_epoch, bundle, prepared ) - conn.execute( - "INSERT INTO " - "execution_context_capture(run_id,state,reason,format_version) " - "VALUES(?,?,?,?)", - (run_id, bundle.state, bundle.reason, bundle.format_version), - ) - return True def get_execution_contexts(self, run_id: str) -> dict: """Return first-capture scalars; hints are None unless currently valid. diff --git a/tests/test_execution_context_database.py b/tests/test_execution_context_database.py index d4ab5d3..60036d8 100644 --- a/tests/test_execution_context_database.py +++ b/tests/test_execution_context_database.py @@ -1,6 +1,8 @@ """Immutable scalar capture, deduplication, and conservative navigation hints.""" from concurrent.futures import ThreadPoolExecutor +from contextlib import contextmanager +from uuid import uuid4 from dataclasses import replace from pathlib import Path import subprocess @@ -70,6 +72,70 @@ class ContextDatabaseTests(unittest.TestCase): ) ) + def test_finish_capture_uses_stored_parent_identity(self): + db = self.database() + start = replace(self.start(db), run_id=str(uuid4())) + db.write_execution_start(start) + finish = self.finish(start) + self.assertTrue(db.write_execution_contexts_for_finish(finish, bundle())) + self.assertEqual(db.get_execution_contexts(start.run_id)["state"], "complete") + missing = replace(finish, queue_item_id=999) + self.assertFalse(db.write_execution_contexts_for_finish(missing, bundle())) + self.assertEqual(self.counts(db), (1, 1, 2, 1)) + + def test_finish_generation_guard_encloses_commit_and_rollback(self): + db = self.database() + start = self.start(db) + db.write_execution_start(start) + finish = self.finish(start) + observed = [] + + @contextmanager + def guard(): + observed.append(db.model.get_connection().in_transaction) + try: + yield True + finally: + observed.append(db.model.get_connection().in_transaction) + + self.assertTrue( + db.write_execution_contexts_for_finish( + finish, bundle(), generation_guard=guard + ) + ) + self.assertEqual(observed, [True, False]) + observed.clear() + other = replace(self.start(db, 1), run_id=str(uuid4())) + db.write_execution_start(other) + with patch( + "database.execution_contexts._insert_exact", + side_effect=ValueError("collision"), + ): + with self.assertRaises(ValueError): + db.write_execution_contexts_for_finish( + self.finish(other), bundle(), generation_guard=guard + ) + self.assertEqual(observed, [True, False]) + self.assertEqual( + db.get_execution_contexts(other.run_id)["state"], "not_captured" + ) + + def test_finish_generation_guard_rejects_without_capture(self): + db = self.database() + start = self.start(db) + db.write_execution_start(start) + + @contextmanager + def stale(): + yield False + + self.assertFalse( + db.write_execution_contexts_for_finish( + self.finish(start), bundle(), generation_guard=stale + ) + ) + self.assertEqual(self.counts(db), (1, 0, 0, 0)) + def test_a_a_b_cardinality(self): db = self.database() self.assertTrue( diff --git a/tests/test_execution_history_writer.py b/tests/test_execution_history_writer.py index 246f556..289818c 100644 --- a/tests/test_execution_history_writer.py +++ b/tests/test_execution_history_writer.py @@ -186,12 +186,14 @@ class HistoryWriterTests(unittest.TestCase): sealed = self.sealed(writer) entered, release = self.block() - def fail_context(*args): + def fail_context(*args, **kwargs): entered.set() release.wait(3) raise error - with patch.object(self.db, "write_execution_contexts", fail_context): + with patch.object( + self.db, "write_execution_contexts_for_finish", fail_context + ): self.assertTrue( writer.submit_finalization(self.finish_command(start), sealed) ) @@ -309,16 +311,16 @@ class HistoryWriterTests(unittest.TestCase): start = self.start_command() self.db.write_execution_start(start) entered, release = self.block() - original = self.db.write_execution_contexts + original = self.db.write_execution_contexts_for_finish first = self.sealed(writer) charge = first.lease.byte_count - def write(*args): + def write(*args, **kwargs): entered.set() release.wait(3) - return original(*args) + return original(*args, **kwargs) - with patch.object(self.db, "write_execution_contexts", write): + with patch.object(self.db, "write_execution_contexts_for_finish", write): self.assertTrue( writer.submit_finalization(self.finish_command(start), first) ) @@ -412,6 +414,118 @@ class HistoryWriterTests(unittest.TestCase): self.assertEqual(foreign.charged_bytes, 0) self.assertEqual(self.db.get_execution_contexts(start.run_id)["contexts"], []) + def test_custom_run_identity_receives_finalized_context(self): + writer = self.writer() + start = replace(self.start_command(), run_id=str(uuid4())) + self.assertTrue(writer.submit(start)) + 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"], "complete" + ) + self.assertEqual(writer.context_budget.charged_bytes, 0) + + def test_invalidation_during_context_preparation_fences_actual_transaction(self): + import database.execution_contexts as contexts_db + + writer = self.writer() + start = self.start_command() + self.db.write_execution_start(start) + entered, release = self.block() + original = contexts_db._prepare + + def prepare(bundle): + prepared = original(bundle) + entered.set() + release.wait(3) + return prepared + + with patch.object(contexts_db, "_prepare", prepare): + self.assertTrue( + writer.submit_finalization( + self.finish_command(start), self.sealed(writer) + ) + ) + self.assertTrue(entered.wait(1)) + self.assertEqual( + self.db.get_execution_run(start.run_id)["outcome"], "succeeded" + ) + generation = writer.status().capture_generation + writer.invalidate() + self.assertGreater(writer.status().capture_generation, generation) + release.set() + 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.clock.advance(5) + self.wait_for(lambda: writer.status().accepting) + self.assertEqual(writer.status().epoch, start.storage_epoch) + self.assertEqual( + self.db.get_execution_contexts(start.run_id)["state"], "not_captured" + ) + + def test_context_commit_serializes_invalidation_without_blocking_queue_or_close( + self, + ): + import database.execution_contexts as contexts_db + + writer = self.writer() + start = self.start_command() + self.db.write_execution_start(start) + entered, release = self.block() + invalidating, invalidated = threading.Event(), threading.Event() + original = contexts_db._write_capture + + def capture(*args, **kwargs): + entered.set() + release.wait(3) + return original(*args, **kwargs) + + def invalidate(): + invalidating.set() + writer.invalidate() + invalidated.set() + + with patch.object(contexts_db, "_write_capture", capture): + self.assertTrue( + writer.submit_finalization( + self.finish_command(start), self.sealed(writer) + ) + ) + self.assertTrue(entered.wait(1)) + invalidator = threading.Thread(target=invalidate) + invalidator.start() + try: + self.assertTrue(invalidating.wait(1)) + self.assertFalse(invalidated.wait(0.03)) + before = time.monotonic() + self.assertTrue(writer.status().accepting) + self.assertTrue(writer.submit(self.finish_command(start))) + self.assertLess(time.monotonic() - before, 0.1) + self.assertFalse(writer.close(0.02)) + self.assertGreater(writer.context_budget.charged_bytes, 0) + self.assertEqual(writer.status().pending, 1) + finally: + release.set() + invalidator.join(3) + self.assertFalse(invalidator.is_alive()) + self.assertTrue(invalidated.is_set()) + self.wait_for(lambda: not writer.status().alive) + self.assertEqual( + self.db.get_execution_run(start.run_id)["outcome"], "succeeded" + ) + self.assertEqual( + self.db.get_execution_contexts(start.run_id)["state"], "complete" + ) + self.assertEqual(writer.context_budget.charged_bytes, 0) + 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_history_writer.py b/utils/execution_history_writer.py index 4122cea..3037068 100644 --- a/utils/execution_history_writer.py +++ b/utils/execution_history_writer.py @@ -1,12 +1,13 @@ """A bounded, single-worker transport for best-effort execution history.""" +from contextlib import contextmanager from dataclasses import dataclass +from functools import partial import math import queue import threading import time -from typing import Callable -from uuid import UUID, uuid5 +from typing import Callable, Iterator try: from ..database import PromptDatabase @@ -68,6 +69,7 @@ class HistoryWriter: ) self._db_factory = db_factory self._clock = clock + self._context_commit_gate = threading.Lock() self._lock = threading.Lock() self._wake = threading.Event() self._queue: queue.Queue[RunCommand | _Finalization] = queue.Queue( @@ -187,13 +189,27 @@ class HistoryWriter: self._record_gap_locked() self._wake.set() + @contextmanager + def _context_generation_guard( + self, generation: int, epoch: str, capture_generation: int | None + ) -> Iterator[bool]: + with self._context_commit_gate: + with self._lock: + current = ( + generation == self._generation + and epoch == self._epoch + and (capture_generation is None or capture_generation == generation) + ) + yield current + def invalidate(self) -> None: - with self._lock: - self._accepting = False - self._epoch = None - self._generation += 1 - self._next_probe = self._clock() + _PROBE_INTERVAL - self._discard_locked() + with self._context_commit_gate: + with self._lock: + self._accepting = False + self._epoch = None + self._generation += 1 + self._next_probe = self._clock() + _PROBE_INTERVAL + self._discard_locked() self._wake.set() def status(self) -> HistoryWriterStatus: @@ -314,14 +330,8 @@ class HistoryWriter: ) if current: try: - db.write_execution_contexts( - str( - uuid5( - UUID(command.session_id), - str(command.queue_item_id), - ) - ), - command.storage_epoch, + db.write_execution_contexts_for_finish( + command, ( context.bundle if context.bundle is not None @@ -329,6 +339,12 @@ class HistoryWriter: (), context.state, context.reason ) ), + generation_guard=partial( + self._context_generation_guard, + generation, + command.storage_epoch, + context.capture_generation, + ), ) except ValueError: # Validation/collision cannot revoke the outcome.