From 39016bc16fc5546a7af957674477da48675021bb Mon Sep 17 00:00:00 2001 From: buildingjoshbetter Date: Fri, 9 Oct 2026 11:55:14 -0500 Subject: [PATCH 1/5] fix: Fence selected source and vector writers --- docs/resource-budgets.md | 59 ++ tests/test-rebuild-source-streaming-756.py | 6 +- tests/test-tier-writer-fence-795.py | 606 +++++++++++++++++++++ truememory/engine.py | 51 +- truememory/storage.py | 14 +- truememory/tier_switch/writer.py | 176 ++++++ truememory/vector_search.py | 30 +- 7 files changed, 917 insertions(+), 25 deletions(-) create mode 100644 tests/test-tier-writer-fence-795.py create mode 100644 truememory/tier_switch/writer.py diff --git a/docs/resource-budgets.md b/docs/resource-budgets.md index e6dd556d..81bb0286 100644 --- a/docs/resource-budgets.md +++ b/docs/resource-budgets.md @@ -2130,3 +2130,62 @@ vectors, stub file observations, and explicit vec0 schema doubles. The opt-in Runtime reconciliation, operation leases, base/pro config-only publication, and manager/MCP/CLI adoption remain separate work. This API has no claimed public memory reduction and does not prove the original 26.24 GB Mac incident fixed. + +### Selected-generation fences for source and vector writers + +`capture_writer_selection(conn)` captures an immutable selected descriptor before +native preparation or source mutation. Its `WriterCapture` binds the exact +accepted connection object; it is not a file-descriptor or reopened-file identity +proof. If the runtime bridge is already loaded, its current operation supplies +the descriptor, including an explicitly admitted legacy `None`. The writer +helper does not import or reconcile runtime state. Unbound callers use a short +read snapshot, or borrow an existing transaction without ending or upgrading it. + +`writer_transaction` acquires BEGIN IMMEDIATE for owned work. Borrowed work uses +a unique outer savepoint and a compatibility inner savepoint, then a zero-row +`UPDATE messages SET id = id WHERE 0` to acquire SQLite's writer before checking +the selection. This changes no source row or metadata. A stale snapshot that +cannot upgrade fails for retry; the helper never commits or rolls back the +caller's outer transaction. `require_writer_selection` itself is read-only and +requires a proven writer acquisition: sqlite3's `in_transaction` alone cannot +distinguish a deferred reader from a writer. + +The comparison rejects a first selection appearing after legacy admission and +any changed generation or frozen descriptor. Config acknowledgement alone may +change. Selected publication also verifies exact registry model, dimension, pair, +and embedder metadata; model-backed paths check their prepared model/dimension +and target tables. No publication rereads a newer generation to legitimize old +vectors. Registry and metadata comparisons use boolean SQL, without copying +unbounded stored values into Python. + +Engine `add`, `update`, `delete`, and `delete_all` fence before their source +mutations, including failed native precomputation. Direct `embed_single` carries +one capture through vector publication. Selected per-user `delete_all` uses the +captured custom pair. Storage `update_message` and `delete_message` add +`commit=False` for caller-owned transactions; their default commit behavior is +unchanged. An owned Engine update still commits its source before vector +publication, preserving the existing committed-source retry message if vectors +cannot publish. Borrowed updates remain pending with the caller's work. + +`_build_streamed_vectors` captures once and checks inside all four existing writer +sites: empty clear, initial clear, each prefix, and final completion. It preserves +source receipts, progress, native batch policy, and the model-generation fence. +Those bulk sites still use the existing `rebuild_source.rebuild_transaction`; +its older commit/rollback uncertainty behavior is unchanged by this patch. + +The new foreground writer helper attempts owned rollback only while a +transaction remains active, and never retries failed cleanup. Borrowed cleanup +targets only its unique savepoint, so an inner RELEASE that succeeded before +raising cannot roll back a same-named caller savepoint. A final unique-savepoint +RELEASE error is ambiguous and propagates without another cleanup attempt. +Post-commit errors propagate without claiming completion or writing status. + +Safe tests exercise actual in-memory SQLite transactions, shared-memory writer +contention, journal parsing, and source/vector method bodies with synthetic +records and model/table doubles. They do not establish native vec0 behavior, +file identity, WAL/file-backed concurrency, or complete runtime adoption. +Raw source-only storage APIs retain their existing defaults. Model-backed +ingest, initialization, compatibility/migration, and maintenance entrypoints +still require the integrated serving boundary or explicit refusal before public +activation. Arbitrary raw SQL is not fenced. This checkpoint adds no activation +call, changes no model or retrieval policy, and makes no measured memory claim. diff --git a/tests/test-rebuild-source-streaming-756.py b/tests/test-rebuild-source-streaming-756.py index bc39586a..3ffb7f5d 100644 --- a/tests/test-rebuild-source-streaming-756.py +++ b/tests/test-rebuild-source-streaming-756.py @@ -34,11 +34,11 @@ def safe_import(name, globals=None, locals=None, fromlist=(), level=0): raise AssertionError("Unexpected application import: " + name) return builtins.__import__(name, globals, locals, fromlist, level) - for name in ("storage", "_platform", "maintenance", "rebuild_source"): - module = types.ModuleType("synthetic_stream_" + name) + for name in ("storage", "_platform", "maintenance", "rebuild_source", "tier_switch.writer"): + module = types.ModuleType("synthetic_stream_" + name.replace(".", "_")) sys.modules[module.__name__] = module module.__dict__["__builtins__"] = dict(vars(builtins), __import__=safe_import) - path = ROOT / "truememory" / (name + ".py") + path = ROOT / "truememory" / (name.replace(".", "/") + ".py") exec(compile(path.read_text(encoding="utf-8"), str(path), "exec"), module.__dict__) modules[name] = module vector = types.ModuleType("synthetic_stream_vector") diff --git a/tests/test-tier-writer-fence-795.py b/tests/test-tier-writer-fence-795.py new file mode 100644 index 00000000..73d61d6f --- /dev/null +++ b/tests/test-tier-writer-fence-795.py @@ -0,0 +1,606 @@ +"""Real in-memory writer transactions with synthetic selected records and models. + +Certificate provenance and native vec0 are tested by their own gates. These +tests exercise actual journal parsing, SQL integrity, storage, Engine method +bodies, and vector writer control flow without a model or file-backed database. +""" + +from __future__ import annotations + +import ast +import builtins +import dataclasses +import json +import logging +import runpy +import sqlite3 +import sys +import threading +import types +import unittest +import uuid +from contextlib import contextmanager +from pathlib import Path +from unittest.mock import patch + +ROOT = Path(__file__).resolve().parents[1] +BASE = runpy.run_path(str(Path(__file__).with_name("test-rebuild-source-streaming-756.py"))) + + +class Cancelled(BaseException): + pass + + +class Connection(sqlite3.Connection): + armed = False + mode = "" + failure: BaseException | None = None + commits = 0 + rollbacks = 0 + after_commit = None + after_execute = None + + def execute(self, sql: str, parameters: tuple = ()): + if hasattr(self, "statements"): + self.statements.append(sql) + if self.armed: + if self.mode == "rollback_to" and sql.startswith("ROLLBACK TO SAVEPOINT tier_writer_"): + raise self.failure + if self.mode == "inner_before" and sql == "RELEASE SAVEPOINT rebuild_source": + raise self.failure + if self.mode == "guard_before" and sql.startswith("RELEASE SAVEPOINT tier_writer_"): + raise self.failure + result = super().execute(sql, parameters) + if self.after_execute is not None: + self.after_execute(sql, parameters) + if self.armed: + if self.mode == "inner_after" and sql == "RELEASE SAVEPOINT rebuild_source": + raise self.failure + if self.mode == "guard_after" and sql.startswith("RELEASE SAVEPOINT tier_writer_"): + raise self.failure + return result + + def commit(self) -> None: + self.commits += 1 + if self.armed and self.mode == "commit_before": + raise self.failure + if self.armed and self.mode == "commit_retained": + return + super().commit() + if self.after_commit is not None: + callback, self.after_commit = self.after_commit, None + callback() + if self.armed and self.mode == "commit_after": + raise self.failure + + def rollback(self) -> None: + self.rollbacks += 1 + if self.armed and self.mode == "rollback_before": + raise self.failure + if self.armed and self.mode == "rollback_retained": + return + super().rollback() + if self.armed and self.mode == "rollback_after": + raise self.failure + + +class TestWriterFence(unittest.TestCase): + def setUp(self) -> None: + self.modules = BASE["load_modules"]() + self.vector = self.modules["vector_search"] + underlying = self.vector.__dict__["__builtins__"]["__import__"] + + def safe_import(name, globals=None, locals=None, fromlist=(), level=0): + if name in {"numpy", "torch", "model2vec", "sentence_transformers", "sqlite_vec"}: + raise AssertionError("Native imports are excluded from writer tests") + if name == "psutil": + return types.SimpleNamespace() + return underlying(name, globals, locals, fromlist, level) + + for name in ("embedding_target", "tier_config", "tier_switch.cache", "tier_switch.job", + "tier_switch.source", "tier_switch.activation"): + module = types.ModuleType("synthetic_writer_" + name.replace(".", "_")) + sys.modules[module.__name__] = module + module.__dict__["__builtins__"] = dict(vars(builtins), __import__=safe_import) + path = ROOT / "truememory" / (name.replace(".", "/") + ".py") + exec(compile(path.read_text(), str(path), "exec"), module.__dict__) + self.modules[name] = module + self.api = self.modules["tier_switch.writer"] + self.activation = self.modules["tier_switch.activation"] + self.conn = self.connect() + self.conn.executescript(self.modules["storage"]._SCHEMA_SQL) + for name in ("vec_messages", "vec_messages_sep", "vec_messages_custom", "vec_messages_sep_custom"): + self.conn.execute(f"CREATE TABLE {name}(embedding BLOB)") + self.conn.commit() + self.calls = [] + self.encode_hook = lambda: None + self.vector._model = object() + self.vector.get_model = lambda: self.vector._model + + def encode(model, texts, **kwargs): + self.calls.append(list(texts)) + self.encode_hook() + return [[1., 0.] for _ in texts] + + self.vector._encode_with_mps_fallback = encode + self.target = self.modules["embedding_target"].EmbeddingTarget("custom", "synthetic/encoder", 2, "custom") + self.generation = 0 + self.addCleanup(patch.stopall) + # No operation is admitted by the fixture unless a test explicitly + # installs the runtime accessor contract. Never import the real bridge. + patch.dict(sys.modules, {"truememory.tier_switch.runtime": types.SimpleNamespace(current_operation=lambda conn: None)}).start() + + def connect(self, uri: str | None = None): + conn = sqlite3.connect(uri or ":memory:", uri=uri is not None, factory=Connection) + conn.statements = [] + self.addCleanup(conn.close) + return conn + + def publish(self, conn=None, *, target=None, acknowledged=True, commit=True): + conn = self.conn if conn is None else conn + target = target or self.target + self.generation += 1 + generation = f"{self.generation:032x}" + intent = self.activation.ActivationIntent( + generation, None, target, "synthetic/reranker", generation, + "rebuild-bridge-v1", "synthetic-epoch", "a" * 64, + state="config_acknowledged" if acknowledged else "db_selected", + ) + selected = self.activation.TierSelection( + generation, generation, target, intent.reranker_id, target.tables, generation, + intent.tracker, intent.source_epoch, intent.source_schema_signature, + generation, "b" * 64, "c" * 64, "d" * 64, "e" * 64, 0, None, acknowledged, + ) + self.activation._put(conn, "tier_activation_v1", self.activation._dump(intent)) + self.activation._put(conn, "tier_selected_v1", self.activation._dump(selected)) + self.activation._put(conn, "embed_model", target.model_id) + self.activation._put(conn, "embed_dim", str(target.dimension)) + self.modules["tier_switch.cache"].VectorCacheRegistry.set( + conn, target.tier_group, model_name=target.model_id, embedding_dim=target.dimension, commit=False, + ) + if commit: + conn.commit() + return selected + + def selected_runtime(self): + selected = self.publish() + self.vector.EMBEDDING_MODEL = selected.target.model_id + self.vector._embedding_dim = selected.target.dimension + self.vector._active_vec_table = lambda conn: selected.tables[0] + self.vector._active_sep_table = lambda conn: selected.tables[1] + return selected + + def seed(self, mid=1, content="synthetic-original") -> None: + self.conn.execute("INSERT INTO messages(id,content,sender) VALUES(?,?,?)", (mid, content, "synthetic-sender")) + self.conn.commit() + + def contents(self, conn=None): + conn = self.conn if conn is None else conn + return conn.execute("SELECT id,content FROM messages ORDER BY id").fetchall() + + def engine(self, *, vectors=True): + tree = ast.parse((ROOT / "truememory/engine.py").read_text()) + original = next(n for n in tree.body if isinstance(n, ast.ClassDef) and n.name == "TrueMemoryEngine") + methods = [n for n in original.body if isinstance(n, ast.FunctionDef) and n.name in {"add", "update", "delete", "delete_all"}] + cls = ast.ClassDef(name="SyntheticEngine", bases=[], keywords=[], body=methods, decorator_list=[]) + constants = [n for n in tree.body if isinstance(n, ast.Assign) and any( + isinstance(t, ast.Name) and t.id in {"_ALLOWED_TABLES", "_ALLOWED_COLUMNS", "_ALL_VEC_TABLES", "_SQLITE_IN_CHUNK"} + for t in n.targets + )] + helpers = [n for n in tree.body if isinstance(n, ast.FunctionDef) and n.name in {"_delete_in_chunks", "_resolve_vec_tables"}] + namespace = dict(__builtins__=self.vector.__dict__["__builtins__"], logger=logging.getLogger(__name__), + MAX_CONTENT_LENGTH=10000, sqlite3=sqlite3, + insert_message=self.modules["storage"].insert_message, + update_message=self.modules["storage"].update_message, + delete_message=self.modules["storage"].delete_message) + exec(compile(ast.fix_missing_locations(ast.Module(body=[*constants, *helpers, cls], type_ignores=[])), + "actual-writer-engine", "exec"), namespace) + engine = namespace["SyntheticEngine"]() + engine.conn = self.conn + engine._write_lock = threading.Lock() + engine._ensure_connection = lambda: None + engine._maybe_auto_consolidate = lambda: None + engine._has_vectors = vectors + engine._has_personality = engine._has_style_vec = False + engine.get = lambda mid: self.modules["storage"].get_message(self.conn, mid) + return engine + + def test_capture_is_frozen_read_only_and_connection_bound(self) -> None: + before = self.conn.total_changes + captured = self.api.capture_writer_selection(self.conn) + self.assertIsNone(captured.selection) + self.assertIs(captured.connection, self.conn) + self.assertNotIn("Connection", repr(captured)) + self.assertEqual(self.conn.total_changes, before) + self.assertFalse(self.conn.in_transaction) + with self.assertRaises(dataclasses.FrozenInstanceError): + captured.selection = "changed" + other = self.connect() + other.executescript("CREATE TABLE messages(id INTEGER PRIMARY KEY, content TEXT)") + with self.assertRaises(self.api.WriterSelectionChanged): + with self.api.writer_transaction(other, captured): + other.execute("INSERT INTO messages VALUES(1,'wrong connection')") + self.assertEqual(self.contents(other), []) + + def test_runtime_legacy_none_never_recaptures_first_selection(self) -> None: + selected = self.selected_runtime() + runtime = types.SimpleNamespace(current_operation=lambda conn: types.SimpleNamespace(selection=None)) + with patch.dict(sys.modules, {"truememory.tier_switch.runtime": runtime}): + captured = self.api.capture_writer_selection(self.conn) + self.assertIsNone(captured.selection) + with self.assertRaises(self.api.WriterSelectionChanged): + with self.api.writer_transaction(self.conn, captured): + self.fail("Legacy admission cannot publish into selected generation") + self.assertEqual(self.activation.read_activation_state(self.conn).selection, selected) + + def test_runtime_wrong_connection_refusal_propagates_without_fallback(self) -> None: + def current(conn): + raise RuntimeError("synthetic wrong accepted connection") + with patch.dict(sys.modules, {"truememory.tier_switch.runtime": types.SimpleNamespace(current_operation=current)}): + before = len(self.conn.statements) + with self.assertRaisesRegex(RuntimeError, "wrong accepted connection"): + self.api.capture_writer_selection(self.conn) + self.assertEqual(len(self.conn.statements), before) + + def test_two_actual_handles_reject_first_selection_and_same_width_generation_change(self) -> None: + uri = "file:synthetic-writer-" + uuid.uuid4().hex + "?mode=memory&cache=shared" + left, right = self.connect(uri), self.connect(uri) + left.executescript(self.modules["storage"]._SCHEMA_SQL) + captured = self.api.capture_writer_selection(left) + first = self.publish(right) + with self.assertRaises(self.api.WriterSelectionChanged): + with self.api.writer_transaction(left, captured): + left.execute("INSERT INTO messages(content) VALUES('stale legacy')") + self.assertEqual(self.contents(left), []) + captured = self.api.capture_writer_selection(left) + second = self.publish(right) + self.assertEqual(first.target.dimension, second.target.dimension) + self.assertNotEqual(first.generation, second.generation) + with self.assertRaises(self.api.WriterSelectionChanged): + with self.api.writer_transaction(left, captured): + left.execute("INSERT INTO messages(content) VALUES('stale selected')") + self.assertEqual(self.contents(left), []) + + def test_borrowed_reader_cannot_upgrade_past_another_writer(self) -> None: + uri = "file:synthetic-lock-" + uuid.uuid4().hex + "?mode=memory&cache=shared" + left, right = self.connect(uri), self.connect(uri) + left.executescript(self.modules["storage"]._SCHEMA_SQL) + right.execute("BEGIN IMMEDIATE") + left.execute("BEGIN") + captured = self.api.capture_writer_selection(left) + with self.assertRaises(sqlite3.OperationalError): + with self.api.writer_transaction(left, captured): + left.execute("INSERT INTO messages(content) VALUES('must not mutate')") + self.assertTrue(left.in_transaction) + self.assertTrue(right.in_transaction) + self.assertEqual(self.contents(left), []) + right.rollback() + left.rollback() + + def test_exact_model_dimension_pair_registry_and_metadata_guard(self) -> None: + self.selected_runtime() + captured = self.api.capture_writer_selection(self.conn) + for options in ({"model_id": "synthetic/wrong"}, {"dimension": 3}, + {"dimension": True}, {"tables": tuple(reversed(self.target.tables))}): + with self.subTest(options=options), self.assertRaises(self.api.WriterSelectionChanged): + with self.api.writer_transaction(self.conn, captured, **options): + self.fail("Mismatched output must not be admitted") + for sql in ("UPDATE vector_cache_registry SET vec_table='wrong'", + "UPDATE metadata SET value='wrong' WHERE key='embed_model'", + "UPDATE metadata SET value='02' WHERE key='embed_dim'"): + self.conn.execute("BEGIN") + self.conn.execute(sql) + with self.subTest(sql=sql), self.assertRaises(self.api.WriterSelectionChanged): + with self.api.writer_transaction(self.conn, captured): + self.fail("Changed metadata must not be admitted") + self.assertTrue(self.conn.in_transaction) + self.conn.rollback() + + def test_acknowledgement_alone_does_not_change_writer_space(self) -> None: + selected = self.publish(acknowledged=False) + captured = self.api.capture_writer_selection(self.conn) + self.activation.acknowledge_config(self.conn, generation=selected.generation) + with self.api.writer_transaction(self.conn, captured): + self.conn.execute("INSERT INTO messages(content) VALUES('same embedding space')") + self.assertEqual(self.contents(), [(1, "same embedding space")]) + + def test_borrowed_success_and_refusal_preserve_outer_work(self) -> None: + captured = self.api.capture_writer_selection(self.conn) + self.conn.execute("INSERT INTO messages(id,content) VALUES(11,'caller')") + before = self.conn.commits + with self.api.writer_transaction(self.conn, captured): + self.conn.execute("INSERT INTO messages(id,content) VALUES(12,'operation')") + self.assertEqual(self.conn.commits, before) + self.assertTrue(self.conn.in_transaction) + self.assertEqual(self.contents(), [(11, "caller"), (12, "operation")]) + self.conn.rollback() + self.assertEqual(self.contents(), []) + self.conn.execute("INSERT INTO messages(id,content) VALUES(11,'caller')") + self.publish(commit=False) + with self.assertRaises(self.api.WriterSelectionChanged): + with self.api.writer_transaction(self.conn, captured): + self.fail("Stale capture must refuse before mutation") + self.assertTrue(self.conn.in_transaction) + self.assertEqual(self.contents(), [(11, "caller")]) + + def test_capture_borrows_without_ending_or_upgrading_outer_transaction(self) -> None: + self.conn.execute("INSERT INTO messages(id,content) VALUES(11,'caller')") + before = self.conn.commits, self.conn.rollbacks + self.assertIsNone(self.api.capture_writer_selection(self.conn).selection) + self.assertEqual((self.conn.commits, self.conn.rollbacks), before) + self.assertTrue(self.conn.in_transaction) + self.conn.rollback() + self.assertEqual(self.contents(), []) + + def test_owned_commit_faults_before_after_and_retained_do_not_fake_success(self) -> None: + for mode in ("commit_before", "commit_after", "commit_retained"): + with self.subTest(mode=mode): + captured = self.api.capture_writer_selection(self.conn) + self.conn.armed, self.conn.mode, self.conn.failure = True, mode, Cancelled("synthetic commit fault") + rollbacks = self.conn.rollbacks + with self.assertRaises((Cancelled, self.api.WriterSelectionChanged)): + with self.api.writer_transaction(self.conn, captured): + self.conn.execute("INSERT INTO messages(id,content) VALUES(12,'operation')") + self.conn.armed = False + self.assertFalse(self.conn.in_transaction) + if mode == "commit_after": + self.assertEqual(self.conn.rollbacks, rollbacks) + self.assertEqual(self.contents(), [(12, "operation")]) + self.conn.execute("DELETE FROM messages") + self.conn.commit() + else: + self.assertEqual(self.conn.rollbacks, rollbacks + 1) + self.assertEqual(self.contents(), []) + + def test_owned_uncertain_rollback_is_attempted_once(self) -> None: + for mode in ("rollback_before", "rollback_after", "rollback_retained"): + with self.subTest(mode=mode): + captured = self.api.capture_writer_selection(self.conn) + self.conn.armed, self.conn.mode = True, mode + self.conn.failure = sqlite3.OperationalError("synthetic rollback uncertainty") + before = self.conn.rollbacks + try: + with self.assertRaises((sqlite3.OperationalError, self.api.WriterSelectionChanged)): + with self.api.writer_transaction(self.conn, captured): + self.conn.execute("INSERT INTO messages(id,content) VALUES(12,'operation')") + raise Cancelled("synthetic body cancellation") + self.assertEqual(self.conn.rollbacks, before + 1) + self.assertEqual(self.conn.in_transaction, mode != "rollback_after") + finally: + self.conn.armed = False + self.conn.rollback() + + def test_repeated_caller_savepoint_survives_inner_release_before_and_after_fault(self) -> None: + for mode in ("inner_before", "inner_after"): + with self.subTest(mode=mode): + captured = self.api.capture_writer_selection(self.conn) + self.conn.execute("BEGIN") + self.conn.execute("SAVEPOINT rebuild_source") + self.conn.execute("INSERT INTO messages(id,content) VALUES(11,'caller')") + self.conn.armed, self.conn.mode, self.conn.failure = True, mode, Cancelled("synthetic inner release fault") + with self.assertRaises(Cancelled): + with self.api.writer_transaction(self.conn, captured): + self.conn.execute("INSERT INTO messages(id,content) VALUES(12,'operation')") + self.conn.armed = False + self.assertTrue(self.conn.in_transaction) + self.assertEqual(self.contents(), [(11, "caller")]) + self.conn.execute("RELEASE SAVEPOINT rebuild_source") + self.assertTrue(self.conn.in_transaction) + self.conn.rollback() + + def test_ambiguous_final_guard_release_never_attempts_cleanup(self) -> None: + for mode in ("guard_before", "guard_after"): + with self.subTest(mode=mode): + captured = self.api.capture_writer_selection(self.conn) + self.conn.execute("INSERT INTO messages(id,content) VALUES(11,'caller')") + self.conn.armed, self.conn.mode, self.conn.failure = True, mode, Cancelled("synthetic guard release fault") + before = len(self.conn.statements) + with self.assertRaises(Cancelled): + with self.api.writer_transaction(self.conn, captured): + self.conn.execute("INSERT INTO messages(id,content) VALUES(12,'operation')") + tail = self.conn.statements[before:] + self.assertTrue(tail[-1].startswith("RELEASE SAVEPOINT tier_writer_")) + self.assertFalse(any(sql.startswith("ROLLBACK") for sql in tail)) + self.assertTrue(self.conn.in_transaction) + self.assertEqual(self.contents(), [(11, "caller"), (12, "operation")]) + self.conn.armed = False + self.conn.rollback() + + def test_borrowed_rollback_to_failure_never_releases_or_discards_outer_work(self) -> None: + captured = self.api.capture_writer_selection(self.conn) + self.conn.execute("INSERT INTO messages(id,content) VALUES(11,'caller')") + self.conn.armed, self.conn.mode = True, "rollback_to" + self.conn.failure = sqlite3.OperationalError("synthetic rollback-to failure") + before = len(self.conn.statements) + with self.assertRaises(sqlite3.OperationalError): + with self.api.writer_transaction(self.conn, captured): + self.conn.execute("INSERT INTO messages(id,content) VALUES(12,'operation')") + raise Cancelled("synthetic cancellation") + tail = self.conn.statements[before:] + self.assertEqual(sum(sql.startswith("ROLLBACK TO") for sql in tail), 1) + self.assertFalse(any(sql.startswith("RELEASE") for sql in tail)) + self.assertTrue(self.conn.in_transaction) + self.assertEqual(self.contents(), [(11, "caller"), (12, "operation")]) + self.conn.armed = False + self.conn.rollback() + + def test_add_stale_capture_refuses_even_when_native_preparation_fails(self) -> None: + self.selected_runtime() + def fail(): + self.publish() + raise RuntimeError("synthetic encoder unavailable") + self.encode_hook = fail + with self.assertRaises(self.api.WriterSelectionChanged): + self.engine().add("must not store") + self.assertEqual(self.contents(), []) + + def test_update_stale_capture_refuses_before_source_change(self) -> None: + self.selected_runtime() + self.seed() + fired = [] + def change(): + if not fired: + fired.append(True) + self.publish() + self.encode_hook = change + with self.assertRaises(self.api.WriterSelectionChanged): + self.engine().update(1, "must not store") + self.assertEqual(self.contents(), [(1, "synthetic-original")]) + + def test_source_commit_then_selection_change_keeps_retry_wording(self) -> None: + self.selected_runtime() + self.seed() + self.conn.after_commit = self.publish + with self.assertRaisesRegex(self.vector.VectorPublicationChanged, "Source update committed; retry update to publish vectors"): + self.engine().update(1, "committed-source") + self.assertEqual(self.contents(), [(1, "committed-source")]) + self.assertEqual(self.conn.execute("SELECT * FROM vec_messages_custom").fetchall(), []) + + def test_delete_and_delete_all_reject_stale_admitted_operation_before_any_delete(self) -> None: + selected = self.selected_runtime() + self.seed() + runtime = types.SimpleNamespace(current_operation=lambda conn: types.SimpleNamespace(selection=selected)) + self.publish() + with patch.dict(sys.modules, {"truememory.tier_switch.runtime": runtime}): + for operation in (lambda: self.engine().delete(1), lambda: self.engine().delete_all(), + lambda: self.engine().delete_all("synthetic-sender")): + with self.assertRaises(self.api.WriterSelectionChanged): + operation() + self.assertEqual(self.contents(), [(1, "synthetic-original")]) + + def test_selected_delete_all_uses_exact_custom_pair_and_borrowed_rollback(self) -> None: + self.selected_runtime() + self.seed() + for table in self.target.tables: + self.conn.execute(f"INSERT INTO {table}(rowid,embedding) VALUES(1,x'0102')") + self.conn.commit() + self.conn.execute("INSERT INTO metadata(key,value) VALUES('caller','pending')") + self.assertTrue(self.engine(vectors=False).delete_all("synthetic-sender")) + self.assertEqual(self.contents(), []) + self.assertTrue(self.conn.in_transaction) + self.conn.rollback() + self.assertEqual(self.contents(), [(1, "synthetic-original")]) + for table in self.target.tables: + self.assertEqual(self.conn.execute(f"SELECT rowid FROM {table}").fetchall(), [(1,)]) + self.assertTrue(self.engine(vectors=False).delete_all("synthetic-sender")) + self.assertFalse(self.conn.in_transaction) + for table in self.target.tables: + self.assertEqual(self.conn.execute(f"SELECT rowid FROM {table}").fetchall(), []) + + def test_direct_embed_refuses_changed_selection_without_new_metadata(self) -> None: + self.selected_runtime() + self.seed() + fired = [] + def change(): + if not fired: + fired.append(self.publish()) + self.encode_hook = change + with self.assertRaises(self.vector.VectorPublicationChanged): + self.vector.embed_single(self.conn, 1, "synthetic-original") + self.assertEqual(self.conn.execute("SELECT * FROM vec_messages_custom").fetchall(), []) + self.assertEqual(self.activation.read_activation_state(self.conn).selection, fired[0]) + + def test_legacy_crud_and_borrowed_update_delete_remain_default_compatible(self) -> None: + engine = self.engine(vectors=False) + item = engine.add("legacy") + self.assertEqual(engine.update(item["id"], "legacy-updated")["content"], "legacy-updated") + self.conn.execute("INSERT INTO metadata(key,value) VALUES('caller','pending')") + engine.update(item["id"], "pending-update") + self.assertTrue(self.conn.in_transaction) + self.conn.rollback() + self.assertEqual(self.contents(), [(item["id"], "legacy-updated")]) + self.conn.execute("INSERT INTO metadata(key,value) VALUES('caller','pending')") + self.assertTrue(engine.delete(item["id"])) + self.assertTrue(self.conn.in_transaction) + self.conn.rollback() + self.assertEqual(self.contents(), [(item["id"], "legacy-updated")]) + self.assertTrue(engine.delete(item["id"])) + self.assertFalse(self.conn.in_transaction) + + def test_storage_defaults_commit_and_false_never_commits(self) -> None: + self.seed() + before = self.conn.commits + self.modules["storage"].update_message(self.conn, 1, content="default") + self.assertEqual(self.conn.commits, before + 1) + self.modules["storage"].update_message(self.conn, 1, content="pending", commit=False) + self.assertTrue(self.conn.in_transaction) + self.conn.rollback() + self.assertEqual(self.contents(), [(1, "default")]) + self.modules["storage"].delete_message(self.conn, 1, commit=False) + self.assertTrue(self.conn.in_transaction) + self.conn.rollback() + self.assertEqual(self.contents(), [(1, "default")]) + self.modules["storage"].delete_message(self.conn, 1) + self.assertFalse(self.conn.in_transaction) + self.assertEqual(self.contents(), []) + + def test_bulk_prefix_cannot_recapture_new_generation(self) -> None: + self.selected_runtime() + self.seed() + @contextmanager + def owner(path): + yield None + maintenance = self.modules["maintenance"] + fired = [] + def change(): + if not fired: + fired.append(self.publish()) + self.encode_hook = change + with patch.object(maintenance, "maintenance_owner", owner), patch.object(maintenance, "connection_database_path", lambda conn: Path("/synthetic/database")): + with self.assertRaises(self.api.WriterSelectionChanged): + self.vector.build_vectors(self.conn) + self.assertEqual(self.conn.execute("SELECT * FROM vec_messages_custom").fetchall(), []) + self.assertEqual(self.activation.read_activation_state(self.conn).selection, fired[0]) + + def test_bulk_initial_and_empty_clear_refuse_stale_capture(self) -> None: + for nonempty in (False, True): + with self.subTest(nonempty=nonempty): + fixture = TestWriterFence() + fixture.setUp() + try: + fixture.selected_runtime() + if nonempty: + fixture.seed() + fixture.conn.execute("INSERT INTO vec_messages_custom(rowid,embedding) VALUES(99,x'0102')") + fixture.conn.commit() + @contextmanager + def owner(path): + fixture.publish() + yield None + maintenance = fixture.modules["maintenance"] + with patch.object(maintenance, "maintenance_owner", owner), patch.object(maintenance, "connection_database_path", lambda conn: Path("/synthetic/database")): + with self.assertRaises(fixture.api.WriterSelectionChanged): + fixture.vector.build_vectors(fixture.conn) + self.assertEqual(fixture.calls, []) + self.assertEqual(fixture.conn.execute("SELECT rowid,embedding FROM vec_messages_custom").fetchall(), [(99, b"\x01\x02")]) + finally: + fixture.doCleanups() + + def test_bulk_finish_cannot_publish_completion_or_metadata_after_selection_change(self) -> None: + self.selected_runtime() + self.seed() + @contextmanager + def owner(path): + yield None + changed = [] + def observe(sql, parameters): + if sql.startswith("INSERT OR REPLACE INTO metadata(key,value,updated_at)") and parameters[0].startswith("vec_source_v1:"): + value = json.loads(parameters[1]) + if value["consumed"] == 1 and not value["complete"]: + self.conn.after_commit = lambda: changed.append(self.publish()) + self.conn.after_execute = observe + maintenance = self.modules["maintenance"] + with patch.object(maintenance, "maintenance_owner", owner), patch.object(maintenance, "connection_database_path", lambda conn: Path("/synthetic/database")): + with self.assertRaises(self.api.WriterSelectionChanged): + self.vector.build_vectors(self.conn) + self.conn.after_execute = None + self.assertEqual(len(changed), 1) + self.assertEqual(self.conn.execute("SELECT rowid FROM vec_messages_custom").fetchall(), [(1,)]) + manifest = self.modules["rebuild_source"].load_manifest(self.conn, "vec_source_v1:vec_messages_custom") + self.assertEqual(manifest.consumed, 1) + self.assertFalse(manifest.complete) + self.assertEqual(self.activation.read_activation_state(self.conn).selection, changed[0]) + + +if __name__ == "__main__": + unittest.main() diff --git a/truememory/engine.py b/truememory/engine.py index d2dd2907..ec096cee 100644 --- a/truememory/engine.py +++ b/truememory/engine.py @@ -712,6 +712,9 @@ def add( self._ensure_connection() + from truememory.tier_switch.writer import capture_writer_selection, writer_transaction + write_selection = capture_writer_selection(self.conn) + pre_embedding = None pre_sep_embedding = None if self._has_vectors: @@ -739,14 +742,17 @@ def add( logger.debug("Failed to pre-compute style vector during add()", exc_info=True) from contextlib import nullcontext - from truememory.rebuild_source import rebuild_transaction if pre_embedding is not None: from truememory.vector_search import _foreground_model_fence with self._write_lock, ( _foreground_model_fence(pre_identity, blocking=not self.conn.in_transaction) if pre_embedding is not None else nullcontext() - ), rebuild_transaction(self.conn, write=True): + ), writer_transaction( + self.conn, write_selection, + model_id=pre_identity[1] if pre_embedding is not None else None, + dimension=pre_identity[2] if pre_embedding is not None else None, + ): if self._has_style_vec: from truememory.maintenance import _capture_style_append, _publish_style_append style_before = _capture_style_append(self.conn) @@ -769,7 +775,9 @@ def add( serialize_f32, _write_embedder_metadata_no_commit, ) - vec_tbl, sep_tbl = _validate_foreground_vector_target(self.conn, pre_identity, new_id) + vec_tbl, sep_tbl = _validate_foreground_vector_target( + self.conn, pre_identity, new_id, selection=write_selection, + ) self.conn.execute( f"INSERT INTO {vec_tbl}(rowid, embedding) VALUES (?, ?)", (new_id, serialize_f32(pre_embedding)), @@ -887,8 +895,10 @@ def delete(self, memory_id: int) -> bool: Returns True if deleted, False if not found. """ self._ensure_connection() - with self._write_lock: - deleted = delete_message(self.conn, memory_id) + from truememory.tier_switch.writer import capture_writer_selection, writer_transaction + write_selection = capture_writer_selection(self.conn) + with self._write_lock, writer_transaction(self.conn, write_selection): + deleted = delete_message(self.conn, memory_id, commit=False) if deleted: self._maybe_auto_consolidate() return deleted @@ -915,8 +925,10 @@ def delete_all(self, user_id: str | None = None) -> bool: raise ValueError("user_id cannot be an empty string") self._ensure_connection() + from truememory.tier_switch.writer import capture_writer_selection, writer_transaction + write_selection = capture_writer_selection(self.conn) - with self._write_lock: + with self._write_lock, writer_transaction(self.conn, write_selection): if user_id is not None: # Get message IDs and episode IDs for this user before deleting msg_ids = [ @@ -991,11 +1003,14 @@ def delete_all(self, user_id: str | None = None) -> bool: # Clean vector tables for deleted message IDs if msg_ids: - try: - vec_tbl, sep_tbl = _resolve_vec_tables(self.conn) - except Exception: - logger.warning("_resolve_vec_tables failed in delete_all(user_id=%s); falling back to legacy table names", user_id, exc_info=True) - vec_tbl, sep_tbl = "vec_messages", "vec_messages_sep" + if write_selection.selection is not None: + vec_tbl, sep_tbl = write_selection.selection.tables + else: + try: + vec_tbl, sep_tbl = _resolve_vec_tables(self.conn) + except Exception: + logger.warning("_resolve_vec_tables failed in delete_all(user_id=%s); falling back to legacy table names", user_id, exc_info=True) + vec_tbl, sep_tbl = "vec_messages", "vec_messages_sep" # dict.fromkeys deduplicates in case vec_tbl == sep_tbl (shouldn't happen, but defensive) for vec_table in dict.fromkeys((vec_tbl, sep_tbl)): try: @@ -1057,7 +1072,6 @@ def delete_all(self, user_id: str | None = None) -> bool: except Exception: logger.warning("Failed to rebuild FTS index during delete_all", exc_info=True) - self.conn.commit() if deleted: self._maybe_auto_consolidate() return deleted @@ -1137,6 +1151,8 @@ def update(self, memory_id: int, content: str | None = None, **fields) -> dict | Updated memory dict, or None if not found. """ self._ensure_connection() + from truememory.tier_switch.writer import capture_writer_selection, writer_transaction + write_selection = capture_writer_selection(self.conn) pre_embedding = None pre_sep_embedding = None @@ -1170,7 +1186,12 @@ def update(self, memory_id: int, content: str | None = None, **fields) -> dict | if content is not None: fields["content"] = content - ok = update_message(self.conn, memory_id, **fields) + with writer_transaction( + self.conn, write_selection, + model_id=pre_identity[1] if pre_embedding is not None else None, + dimension=pre_identity[2] if pre_embedding is not None else None, + ): + ok = update_message(self.conn, memory_id, commit=False, **fields) source_committed = ok and not self.conn.in_transaction if not ok: return None @@ -1179,7 +1200,9 @@ def update(self, memory_id: int, content: str | None = None, **fields) -> dict | from truememory.vector_search import VectorPublicationChanged, _foreground_vector_publication try: from truememory.vector_search import serialize_f32, _write_embedder_metadata_no_commit - with _foreground_vector_publication(self.conn, pre_identity, memory_id) as (vec_tbl, sep_tbl): + with _foreground_vector_publication( + self.conn, pre_identity, memory_id, selection=write_selection, + ) as (vec_tbl, sep_tbl): try: self.conn.execute(f"DELETE FROM {vec_tbl} WHERE rowid = ?", (memory_id,)) except Exception: diff --git a/truememory/storage.py b/truememory/storage.py index 0ca5f64e..967f171f 100644 --- a/truememory/storage.py +++ b/truememory/storage.py @@ -1515,7 +1515,7 @@ def insert_message(conn: sqlite3.Connection, msg: dict) -> int: return cursor.lastrowid -def delete_message(conn: sqlite3.Connection, msg_id: int) -> bool: +def delete_message(conn: sqlite3.Connection, msg_id: int, *, commit: bool = True) -> bool: """ Delete a single message and its vector embedding. @@ -1523,6 +1523,8 @@ def delete_message(conn: sqlite3.Connection, msg_id: int) -> bool: The vector embedding in ``vec_messages`` is also removed if the table exists. + ``commit=False`` leaves publication inside the caller's writer transaction. + Args: conn: Open database connection. msg_id: The message ID to delete. @@ -1612,11 +1614,12 @@ def delete_message(conn: sqlite3.Connection, msg_id: int) -> bool: cursor = conn.execute("DELETE FROM messages WHERE id = ?", (msg_id,)) deleted = cursor.rowcount > 0 - conn.commit() + if commit: + conn.commit() return deleted -def update_message(conn: sqlite3.Connection, msg_id: int, **fields) -> bool: +def update_message(conn: sqlite3.Connection, msg_id: int, *, commit: bool = True, **fields) -> bool: """ Update fields on an existing message. @@ -1627,6 +1630,8 @@ def update_message(conn: sqlite3.Connection, msg_id: int, **fields) -> bool: The AFTER UPDATE trigger on ``messages`` automatically keeps the FTS5 index in sync. + ``commit=False`` leaves the source update inside the caller's writer fence. + Args: conn: Open database connection. msg_id: The message ID to update. @@ -1647,5 +1652,6 @@ def update_message(conn: sqlite3.Connection, msg_id: int, **fields) -> bool: cursor = conn.execute(f"UPDATE messages SET {set_clause} WHERE id = ?", values) # The AFTER UPDATE trigger on messages keeps FTS5 in sync automatically. - conn.commit() + if commit: + conn.commit() return cursor.rowcount > 0 diff --git a/truememory/tier_switch/writer.py b/truememory/tier_switch/writer.py new file mode 100644 index 00000000..e73736e2 --- /dev/null +++ b/truememory/tier_switch/writer.py @@ -0,0 +1,176 @@ +"""Selection fences for already-admitted foreground/source writers. + +No native model, serving lease, configuration, or selection is changed here. +An admitted runtime operation supplies its immutable selection; an unbound +direct caller captures a short database snapshot before preparing any output. +""" + +from __future__ import annotations + +import sqlite3 +import sys +import uuid +from collections.abc import Iterator +from contextlib import contextmanager +from dataclasses import dataclass, field, replace +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from truememory.tier_switch.activation import TierSelection + + +class WriterSelectionChanged(RuntimeError): + """Retry the whole operation after serving selection reconciliation.""" + + +@dataclass(frozen=True) +class WriterCapture: + connection: sqlite3.Connection = field(repr=False) + selection: TierSelection | None + + def __post_init__(self) -> None: + _validate(self.selection) + + +def _selection(conn: sqlite3.Connection) -> TierSelection | None: + table = conn.execute( + "SELECT type FROM main.sqlite_master WHERE name='metadata' COLLATE NOCASE", + ).fetchone() + if table is None: + return None + if table[0] != "table": + raise WriterSelectionChanged("Selected metadata is not an ordinary table") + if conn.execute("SELECT 1 FROM main.metadata WHERE key='tier_selected_v1'").fetchone() is None: + return None + from truememory.tier_switch.activation import TierActivationError, _guard_storage, _state + + try: + _guard_storage(conn) + return _state(conn).selection + except (TierActivationError, sqlite3.Error) as exc: + raise WriterSelectionChanged("Selected journal cannot be validated") from exc + + +def _validate(selection: TierSelection | None) -> None: + if selection is not None: + from truememory.tier_switch.activation import TierSelection + + if type(selection) is not TierSelection: + raise WriterSelectionChanged("Writer capture must be an immutable selected descriptor") + + +def capture_writer_selection(conn: sqlite3.Connection) -> WriterCapture: + """Capture before native/source work; never substitute a newer admitted key. + + The optional runtime module is consulted only when already loaded. Its + current_operation accessor must raise for a mismatched active connection; + None means unbound, while operation.selection=None is an admitted legacy + selection. Caller transactions are read without ending or upgrading them. + """ + runtime = sys.modules.get("truememory.tier_switch.runtime") + if runtime is not None: + operation = runtime.current_operation(conn) + if operation is not None: + _validate(operation.selection) + return WriterCapture(conn, operation.selection) + owned = not conn.in_transaction + if owned: + conn.execute("BEGIN") + try: + return WriterCapture(conn, _selection(conn)) + finally: + if owned: + # One cleanup attempt only, including failure after rollback. + conn.rollback() + if conn.in_transaction: + raise WriterSelectionChanged("Writer capture snapshot cleanup is uncertain") + + +def require_writer_selection( + conn: sqlite3.Connection, captured: WriterCapture | None, *, + model_id: str | None = None, dimension: int | None = None, + tables: tuple[str, str] | None = None, +) -> None: + """Read-only comparison under the caller's actual database writer. + + sqlite3 exposes transaction existence, not lock kind. Use writer_transaction + or another proven writer acquisition before this helper. It never loads, + recaptures an admitted selection, mutates metadata, commits, or rolls back. + Config acknowledgement alone does not change the captured embedding space. + """ + if not conn.in_transaction: + raise WriterSelectionChanged("Selected publication requires database writer ownership") + if captured is not None and (type(captured) is not WriterCapture or captured.connection is not conn): + raise WriterSelectionChanged("Writer capture belongs to a different accepted connection") + selected = captured.selection if captured is not None else None + current = _selection(conn) + if selected is None: + if current is not None: + raise WriterSelectionChanged("A selected tier appeared after legacy admission") + return + if current is None or replace(current, config_acknowledged=selected.config_acknowledged) != selected: + raise WriterSelectionChanged("Selected tier generation or descriptor changed; retry the operation") + if ((model_id is not None and model_id != selected.target.model_id) + or (dimension is not None and (type(dimension) is not int or dimension != selected.target.dimension)) + or (tables is not None and (type(tables) is not tuple or tables != selected.tables))): + raise WriterSelectionChanged("Prepared output does not match the captured selected pair") + # Boolean SQL avoids copying unbounded registry/metadata values into Python. + row = conn.execute( + "SELECT vec_table=? AND sep_table=? AND model_name=? " + "AND typeof(embedding_dim)='integer' AND embedding_dim=? " + "FROM main.vector_cache_registry WHERE tier_group=?", + (*selected.tables, selected.target.model_id, selected.target.dimension, selected.target.tier_group), + ).fetchone() + if row is None or row[0] != 1: + raise WriterSelectionChanged("Selected vector registry identity changed") + for key, value in (("embed_model", selected.target.model_id), ("embed_dim", str(selected.target.dimension))): + row = conn.execute( + "SELECT typeof(value)='text' AND value=? FROM main.metadata WHERE key=?", (value, key), + ).fetchone() + if row is None or row[0] != 1: + raise WriterSelectionChanged("Selected embedder metadata changed") + + +@contextmanager +def writer_transaction( + conn: sqlite3.Connection, captured: WriterCapture | None, *, + model_id: str | None = None, dimension: int | None = None, + tables: tuple[str, str] | None = None, +) -> Iterator[None]: + """Acquire the writer before comparison; preserve borrowed outer work. + + Own BEGIN IMMEDIATE or borrow a SAVEPOINT and upgrade via a zero-row source + UPDATE, never a metadata write. A stale read snapshot may fail that upgrade + for retry. Only owned work is committed. Uncertain cleanup is not retried. + """ + owned = not conn.in_transaction + guard = "tier_writer_" + uuid.uuid4().hex + conn.execute("BEGIN IMMEDIATE" if owned else f"SAVEPOINT {guard}") + guard_release_attempted = False + try: + if not owned: + conn.execute("SAVEPOINT rebuild_source") + conn.execute("UPDATE messages SET id = id WHERE 0") + require_writer_selection(conn, captured, model_id=model_id, dimension=dimension, tables=tables) + yield + if owned: + conn.commit() + if conn.in_transaction: + raise WriterSelectionChanged("Writer transaction did not commit") + else: + conn.execute("RELEASE SAVEPOINT rebuild_source") + guard_release_attempted = True + conn.execute(f"RELEASE SAVEPOINT {guard}") + except BaseException: + if owned: + if conn.in_transaction: + conn.rollback() + if conn.in_transaction: + raise WriterSelectionChanged("Writer rollback is uncertain") + elif conn.in_transaction and not guard_release_attempted: + # Only our unique guard is eligible for cleanup. The compatibility + # inner name may now refer to a caller savepoint after RELEASE. + # A final guard-release error is ambiguous and permits no cleanup. + conn.execute(f"ROLLBACK TO SAVEPOINT {guard}") + conn.execute(f"RELEASE SAVEPOINT {guard}") + raise diff --git a/truememory/vector_search.py b/truememory/vector_search.py index fca624cc..70efcb03 100644 --- a/truememory/vector_search.py +++ b/truememory/vector_search.py @@ -1022,11 +1022,17 @@ def _foreground_model_fence(identity: tuple[int, str, int], *, blocking: bool = def _validate_foreground_vector_target( conn: sqlite3.Connection, identity: tuple[int, str, int], message_id: int, + *, selection=None, ) -> tuple[str, str]: """Validate while the caller owns both the database writer and model fence.""" from truememory.rebuild_source import RebuildSourceChanged, load_manifest, manifest_key tables = (_active_vec_table(conn), _active_sep_table(conn)) + from truememory.tier_switch.writer import WriterSelectionChanged, require_writer_selection + try: + require_writer_selection(conn, selection, model_id=identity[1], dimension=identity[2], tables=tables) + except WriterSelectionChanged as exc: + raise VectorPublicationChanged(str(exc)) from exc for table in tables: manifest = load_manifest(conn, manifest_key((table,))) if manifest is not None and (manifest.model != identity[1] or manifest.dimension != identity[2]): @@ -1054,11 +1060,16 @@ def _validate_foreground_vector_target( @contextmanager def _foreground_vector_publication( conn: sqlite3.Connection, identity: tuple[int, str, int], message_id: int, + *, selection=None, ) -> Iterator[tuple[str, str]]: - from truememory.rebuild_source import rebuild_transaction + from truememory.tier_switch.writer import WriterSelectionChanged, writer_transaction - with _foreground_model_fence(identity, blocking=not conn.in_transaction), rebuild_transaction(conn, write=True): - yield _validate_foreground_vector_target(conn, identity, message_id) + with _foreground_model_fence(identity, blocking=not conn.in_transaction): + try: + with writer_transaction(conn, selection, model_id=identity[1], dimension=identity[2]): + yield _validate_foreground_vector_target(conn, identity, message_id, selection=selection) + except WriterSelectionChanged as exc: + raise VectorPublicationChanged(str(exc)) from exc def _build_streamed_vectors( @@ -1066,6 +1077,7 @@ def _build_streamed_vectors( txn_batch: int | None, separation: bool, ) -> int: from truememory.maintenance import connection_database_path, maintenance_owner + from truememory.tier_switch.writer import capture_writer_selection, require_writer_selection from truememory.rebuild_source import ( RebuildManifest, RebuildSourceChanged, capture_source, input_fingerprint, ensure_rebuild_tracking, load_manifest, manifest_key, rebuild_transaction, save_manifest, @@ -1073,7 +1085,11 @@ def _build_streamed_vectors( if messages is not None and not messages: return 0 + write_selection = capture_writer_selection(conn) table = table_name or (_active_sep_table(conn) if separation else _active_vec_table(conn)) + publication_tables = None if write_selection.selection is None else ( + (write_selection.selection.tables[0], table) if separation else (table, write_selection.selection.tables[1]) + ) quoted_table = '"' + table.replace('"', '""') + '"' key = manifest_key((table,)) digest = input_fingerprint(messages, separation=separation) if messages is not None else None @@ -1086,6 +1102,7 @@ def _build_streamed_vectors( if source_empty: identity = (_model_generation, EMBEDDING_MODEL, _embedding_dim) with rebuild_transaction(conn, write=True, publication_fence=_rebuild_model_fence(identity)): + require_writer_selection(conn, write_selection, model_id=identity[1], dimension=identity[2], tables=publication_tables) empty_source = capture_source(conn) if empty_source.total == 0: # An empty range needs no native model allocation. Recheck @@ -1100,6 +1117,7 @@ def _build_streamed_vectors( return 0 model, identity = _capture_rebuild_model(conn) with rebuild_transaction(conn, write=True, publication_fence=_rebuild_model_fence(identity)): + require_writer_selection(conn, write_selection, model_id=identity[1], dimension=identity[2], tables=publication_tables) manifest = load_manifest(conn, key) resumable = (manifest is not None and not manifest.complete and manifest.model == identity[1] and manifest.dimension == identity[2] @@ -1162,6 +1180,7 @@ def _build_streamed_vectors( del embeddings, batch, texts, ids _flush_mps_cache() with rebuild_transaction(conn, write=True, publication_fence=_rebuild_model_fence(identity)): + require_writer_selection(conn, write_selection, model_id=identity[1], dimension=identity[2], tables=publication_tables) manifest.check(conn, key) if rows_to_insert: conn.executemany(f"INSERT INTO {quoted_table}(rowid, embedding) VALUES (?, ?)", rows_to_insert) @@ -1171,6 +1190,7 @@ def _build_streamed_vectors( del rows_to_insert with rebuild_transaction(conn, write=True, publication_fence=_rebuild_model_fence(identity)): + require_writer_selection(conn, write_selection, model_id=identity[1], dimension=identity[2], tables=publication_tables) manifest.check(conn, key) save_manifest(conn, key, replace(manifest, complete=True)) _clear_build_in_progress(conn, table) @@ -1433,6 +1453,8 @@ def embed_single(conn: sqlite3.Connection, message_id: int, content: str) -> Non message_id: The ``messages.id`` of the row being embedded. content: The text to embed. """ + from truememory.tier_switch.writer import capture_writer_selection + write_selection = capture_writer_selection(conn) model, identity = _capture_rebuild_model(conn) embedding = _encode_with_mps_fallback(model, [content])[0] normed = _normalize_for_cosine(embedding) @@ -1453,7 +1475,7 @@ def embed_single(conn: sqlite3.Connection, message_id: int, content: str) -> Non separation = serialize_f32(sep_normed) except Exception: logger.warning("Failed to prepare separation vector during incremental publication", exc_info=True) - with _foreground_vector_publication(conn, identity, message_id) as (vec_tbl, sep_tbl): + with _foreground_vector_publication(conn, identity, message_id, selection=write_selection) as (vec_tbl, sep_tbl): conn.execute(f"INSERT INTO {vec_tbl}(rowid, embedding) VALUES (?, ?)", (message_id, completion)) if separation is not None: try: From 06aaae0ea161f0cf143eb13cef58d92fef907673 Mon Sep 17 00:00:00 2001 From: buildingjoshbetter Date: Fri, 9 Oct 2026 11:55:14 -0500 Subject: [PATCH 2/5] fix: Validate selected model response receipts --- docs/resource-budgets.md | 42 +++ tests/test-model-microbatches-735.py | 3 +- tests/test-tier-serving-receipt-795.py | 480 +++++++++++++++++++++++++ truememory/model_client.py | 79 ++++ truememory/model_server.py | 98 ++++- 5 files changed, 689 insertions(+), 13 deletions(-) create mode 100644 tests/test-tier-serving-receipt-795.py diff --git a/docs/resource-budgets.md b/docs/resource-budgets.md index 81bb0286..7ecccc5c 100644 --- a/docs/resource-budgets.md +++ b/docs/resource-budgets.md @@ -2189,3 +2189,45 @@ ingest, initialization, compatibility/migration, and maintenance entrypoints still require the integrated serving boundary or explicit refusal before public activation. Arbitrary raw SQL is not fenced. This checkpoint adds no activation call, changes no model or retrieval policy, and makes no measured memory claim. + +### Per-response selected serving receipts + +`CertifiedEmbeddingProxy(target)` is an additive remote proxy for the fixed +edge/base/pro embedding targets. It remains an `EmbeddingProxy` subclass for +existing remote ownership/concurrency handling. Each call uses the ordinary +`embed` or `embed_batched` operation, an explicit canonical model ID, and an +`expected_target` descriptor. A successful response must carry integer protocol +version 1, an exact valid target receipt including public tier, and an array of +shape `(input_count, target.dimension)`. An older daemon that ignores the new +field cannot pass by returning equal-width vectors or by succeeding at an +earlier preparation request. Missing, malformed, or different receipts fail +closed on every call, including empty input. Custom targets continue using +`PreparedEmbeddingProxy` and its strict prepared-target operations. + +The server checks canonical built-in identity before admission and the effective +cached model identity under its existing owner before native work. Both direct +main fast inference and the dedicated CPU fast lane return receipts. CPU clone +construction binds the captured residency identity; generation churn can let an +already-admitted correct model finish without publishing a stale clone. Every +main microbatch and CPU retry must match the expected width. Identity/width +refusal cannot become an ordinary fallback result; other transient CPU fast +errors retain the existing main-path fallback with the same expected target. +Microbatch ceilings, ordering, admission/deadline checks, sustained-workload +bookkeeping, and residency generation policy remain in place. + +`CertifiedRerankerProxy(model_name)` requires an explicit bounded model identity. +Each ordinary `rerank`/`rerank_batched` request carries `expected_model_name`; +every successful response requires protocol version 1 and the exact versioned +effective reranker receipt. The server verifies its owned model/name pair after +initial acquisition and after CPU replacement, before prediction or retry. +There is no fallback to another reranker name. Ordinary proxies, request fields, +default selection, and ordinary response behavior remain unchanged. + +Both certified proxy constructors configure identity only. They perform no +network request, native load, or inference and do not prove daemon-native model +readiness or residency. This extension adds no background probe or extra model +cache. Empty inputs retain the native path's existing behavior; successful +empty results still require receipts, without a new synthetic probe. Tests use +stdlib loaders and synthetic model/array doubles; native numerical parity and +full integration remain separate GPUBox gates. No activation call or measured +memory reduction is claimed by this protocol extension. diff --git a/tests/test-model-microbatches-735.py b/tests/test-model-microbatches-735.py index bd74ff9d..a241ab11 100644 --- a/tests/test-model-microbatches-735.py +++ b/tests/test-model-microbatches-735.py @@ -605,7 +605,8 @@ def load_order_runtime() -> dict: "_request_batch_limit", "_after_request_batches", "_recover_embed_oom_locked", "_check_embed_recovery_deadline_locked", "_rerank_global_order", "_rerank_slice_indices", "_preflight_rerank_result", "_publish_embed_state", "_retire_stale_fast_encoder", - "_check_process_memory"} + "_check_process_memory", "_check_serving_embed_model", "_serving_embed_response", + "_serving_reranker_name", "_check_serving_reranker"} body = [ast.ImportFrom(module="__future__", names=[ast.alias(name="annotations")], level=0)] for node in ast.parse(path.read_text(encoding="utf-8")).body: if isinstance(node, ast.FunctionDef) and node.name in functions: diff --git a/tests/test-tier-serving-receipt-795.py b/tests/test-tier-serving-receipt-795.py new file mode 100644 index 00000000..ec5ea748 --- /dev/null +++ b/tests/test-tier-serving-receipt-795.py @@ -0,0 +1,480 @@ +"""Selected response receipts through real control flow and stdlib model doubles.""" +from __future__ import annotations + +import dataclasses +from pathlib import Path +import runpy +import sys +import threading +import types +import unittest +from unittest.mock import Mock, patch + + +BASE = runpy.run_path(str(Path(__file__).with_name("test-prepared-target-795.py"))) +Array = BASE["Array"] + + +class FirstBusyLock: + """One nonblocking contention observation, followed by a real unlocked lock.""" + + def __init__(self) -> None: + self.lock = threading.Lock() + self.first = True + + def acquire(self, blocking: bool = True, timeout: float = -1) -> bool: + if not blocking and self.first: + self.first = False + return False + return self.lock.acquire(blocking, timeout) + + def release(self) -> None: + self.lock.release() + + def locked(self) -> bool: + return self.lock.locked() + + +class TestServingReceipts(unittest.TestCase): + def setUp(self) -> None: + self.fixture = BASE["TestPreparedTargets"]() + self.fixture.setUp() + self.addCleanup(self.fixture.doCleanups) + self.client = self.fixture.client + self.client.PROTOCOL_VERSION = 1 + self.client.np.integer = int + self.server = self.fixture.server + self.ms = self.fixture.ms + self.Target = self.fixture.Target + self.edge = self.Target.capture("edge") + self.base = self.Target.capture("base") + self.calls: list[tuple[str, str, list, int]] = [] + self.builds: list[tuple[str, str | None]] = [] + self.width = 256 + self.before_encode = lambda: None + self.before_predict = lambda: None + self.build_hook = lambda: None + self.embed_oom = False + self.rerank_oom = False + self.cpu_error = False + owner = self + + class Model: + def __init__(self, identity: str, device: str | None = None) -> None: + self.identity, self.device = identity, device + + def encode(self, texts: list[str], **kwargs: object) -> Array: + owner.before_encode() + owner.calls.append(("embed", self.identity, list(texts), kwargs["batch_size"])) + if owner.cpu_error and owner.server._fast_lock.locked(): + owner.cpu_error = False + raise RuntimeError("synthetic transient CPU failure") + if owner.embed_oom: + owner.embed_oom = False + raise RuntimeError("MPS backend out of memory") + return Array([[float(text[1:])] * owner.width for text in texts], (len(texts), owner.width)) + + def to(self, device: str) -> None: + self.device = device + + class Reranker: + def __init__(self, name: str, device: str | None = None) -> None: + self.name, self.device = name, device + owner.builds.append((name, device)) + + def predict(self, pairs: list, **kwargs: object) -> Array: + owner.before_predict() + owner.calls.append(("rerank", self.name, list(pairs), kwargs["batch_size"])) + if owner.rerank_oom: + owner.rerank_oom = False + raise RuntimeError("MPS backend out of memory") + return Array([[float(pair[1][1:])] for pair in pairs], (len(pairs), 1)) + + self.Model = Model + + def build(identity: str, device: str | None) -> Model: + self.assertTrue(self.server._inference_lock.locked() or self.server._fast_lock.locked()) + self.build_hook() + self.builds.append((identity, device)) + return Model(identity, device) + + self.server._build_embed_model = build + self.fixture.mps.auto_detect_device = lambda: "cpu" + self.fixture.mps.is_mps_oom = lambda exc: "MPS backend out of memory" in str(exc) + modules = patch.dict(sys.modules, { + "sentence_transformers": types.SimpleNamespace(CrossEncoder=Reranker), + "truememory.reranker": types.SimpleNamespace(get_current_reranker_name=lambda: "synthetic/default"), + }) + modules.start() + self.addCleanup(modules.stop) + + def embed(self, count: int = 1, *, target=None, **fields: object) -> dict: + target = self.edge if target is None else target + return self.server.handle_request({ + "op": "embed", "tier": target.model_id, + "expected_target": target.to_wire(), "texts": [f"s{i}" for i in range(count)], **fields, + }) + + def rerank(self, count: int = 1, **fields: object) -> dict: + return self.server.handle_request({ + "op": "rerank", "model_name": "synthetic/reranker", "expected_model_name": "synthetic/reranker", + "pairs": [["query", f"s{i}"] for i in range(count)], **fields, + }) + + def assert_released(self) -> None: + self.assertFalse(self.server._inference_lock.locked()) + self.assertFalse(self.server._fast_lock.locked()) + self.assertFalse(self.server._lock.locked()) + self.assertEqual(self.server._inflight, 0) + + def test_constructors_are_lazy_and_embedding_is_a_remote_proxy(self) -> None: + self.client._request_with_autostart = Mock(side_effect=AssertionError("unexpected request")) + proxy = self.client.CertifiedEmbeddingProxy(self.edge) + self.assertIsInstance(proxy, self.client.EmbeddingProxy) + self.client.CertifiedRerankerProxy("synthetic/reranker") + self.client._request_with_autostart.assert_not_called() + self.assertEqual(self.builds, []) + + def test_certified_client_uses_existing_operation_batch_and_deadline(self) -> None: + calls = [] + + def send(request: dict, timeout=None) -> dict: + calls.append((request, timeout)) + return self.server.handle_request(request) + + self.client._request_with_autostart = send + proxy = self.client.CertifiedEmbeddingProxy(self.edge) + output = proxy.encode(["s0", "s1", "s2"], timeout=7, batch_size=2, show_progress_bar=False) + self.assertEqual(output.rows, [[float(i)] * 256 for i in range(3)]) + self.assertEqual(calls[0][0]["op"], "embed_batched") + self.assertEqual(calls[0][0]["tier"], "model2vec") + self.assertEqual(calls[0][0]["expected_target"], self.edge.to_wire()) + self.assertEqual(calls[0][1], 7) + self.assertEqual([len(call[2]) for call in self.calls], [2, 1]) + self.assert_released() + + def test_each_response_rejects_same_width_identity_drift_after_success(self) -> None: + proxy = self.client.CertifiedEmbeddingProxy(self.edge) + output = proxy.encode("s0") + self.assertEqual(output.shape, (1, 256)) + for protocol in (1, None): + for receipt in (None, self.base.to_wire()): + response = {"ok": True, "vectors": output, "target": receipt} + if protocol is not None: + response["protocol"] = protocol + self.client._request_with_autostart = Mock(return_value=response) + with self.subTest(protocol=protocol, receipt=receipt), self.assertRaises(self.client.ProtocolMismatchError): + proxy.encode("s0") + self.client._request_with_autostart.assert_called_once() + + def test_embedding_receipt_protocol_and_descriptor_types_are_strict(self) -> None: + response = self.embed() + mutations = [ + {"protocol": value} for value in (None, True, 1.0, "1", 2) + ] + [{"target": {**self.edge.to_wire(), **change}} for change in ( + {"version": True}, {"version": 2}, {"dimension": True}, {"extra": 1}, {"tier": "base"}, + )] + for mutation in mutations: + self.client._request_with_autostart = Mock(return_value={**response, **mutation}) + with self.subTest(mutation=mutation), self.assertRaises(self.client.ProtocolMismatchError): + self.client.CertifiedEmbeddingProxy(self.edge).encode("s0") + self.client._request_with_autostart = Mock(return_value={**response, "vectors": Array([[1.0]], (1, 1))}) + with self.assertRaises(self.fixture.target_module.EmbeddingTargetError): + self.client.CertifiedEmbeddingProxy(self.edge).encode("s0") + + def test_custom_and_noncanonical_requests_refuse_before_native(self) -> None: + custom = self.Target.capture("custom") + with self.assertRaises(ValueError): + self.client.CertifiedEmbeddingProxy(custom) + for fields in ({"tier": "edge"}, {"tier": "qwen3_256"}, {"expected_target": custom.to_wire()}, + {"expected_target": None}, {"texts": [object()]}): + with self.subTest(fields=fields), self.assertRaises(ValueError): + self.embed(**fields) + self.assertEqual(self.builds, []) + self.assertEqual(self.calls, []) + self.assert_released() + + def test_direct_main_fast_receipt_preserves_sustained_bookkeeping(self) -> None: + before = self.fixture.globals() + response = self.embed() + self.assertEqual(response["target"], self.edge.to_wire()) + self.assertEqual(response["protocol"], 1) + self.assertEqual(self.server._embed_timestamps, []) + self.assertEqual(self.fixture.globals(), before) + self.assertEqual(len(self.calls), 1) + self.assert_released() + + def test_cpu_fast_receipt_and_cache_reuse_while_main_owner_busy(self) -> None: + self.embed() + self.server._inference_lock.acquire() + try: + first = self.embed() + second = self.embed() + finally: + self.server._inference_lock.release() + self.assertEqual(first["target"], self.edge.to_wire()) + self.assertEqual(second["target"], self.edge.to_wire()) + self.assertEqual(self.builds, [("model2vec", "cpu"), ("model2vec", "cpu")]) + self.assertEqual(self.server._embed_timestamps, []) + self.assert_released() + + def test_cold_fast_decline_preserves_main_fallback_and_receipt(self) -> None: + self.server._inference_lock = FirstBusyLock() + response = self.embed() + self.assertEqual(response["target"], self.edge.to_wire()) + self.assertEqual(len(self.server._embed_timestamps), 1) + self.assertEqual(len(self.calls), 1) + self.assert_released() + + def test_ordinary_cpu_conversion_keeps_its_original_post_owner_placement(self) -> None: + self.embed() + checked = self.ms._checked_batch + ownership = [] + + def observe(*args): + ownership.append(self.server._fast_lock.locked()) + return checked(*args) + + self.ms._checked_batch = observe + self.server._inference_lock.acquire() + try: + ordinary = self.server.handle_request({"op": "embed", "tier": "model2vec", "texts": ["s0"]}) + selected = self.embed() + finally: + self.server._inference_lock.release() + self.assertEqual(ownership, [False, True]) + self.assertEqual(ordinary["vectors"].rows, selected["vectors"].rows) + self.assertNotIn("target", ordinary) + self.assertEqual(selected["target"], self.edge.to_wire()) + self.assert_released() + + def test_transient_cpu_fast_failure_keeps_verified_main_fallback(self) -> None: + self.embed() + self.server._inference_lock = FirstBusyLock() + self.cpu_error = True + response = self.embed() + self.assertEqual(response["target"], self.edge.to_wire()) + self.assertEqual(len(self.calls), 3) + self.assertEqual(len(self.server._embed_timestamps), 1) + self.assert_released() + + def test_generation_changes_during_cpu_construction_do_not_relabel_output(self) -> None: + self.embed() + original_generation = self.server._embed_state.generation + + def advance() -> None: + if self.server._fast_lock.locked(): + with self.server._lock: + self.server._publish_embed_state(self.ms._EmbedState(self.Model("qwen3_256"), "qwen3_256", "qwen3_256")) + self.server._publish_embed_state(self.ms._EmbedState(self.Model("model2vec"), "model2vec", "model2vec")) + + self.build_hook = advance + self.server._inference_lock.acquire() + try: + response = self.embed() + finally: + self.server._inference_lock.release() + self.assertEqual(response["target"], self.edge.to_wire()) + self.assertIsNot(self.server._embed_state.generation, original_generation) + self.assertIsNone(self.server._fast_encoder) + self.assert_released() + + def test_effective_main_and_fast_identity_mismatch_never_encodes_or_falls_back(self) -> None: + for fast in (False, True): + self.server._embed_state = self.ms._EmbedState(self.Model("qwen3_256"), "model2vec", "qwen3_256") + if fast: + self.server._inference_lock.acquire() + try: + with self.subTest(fast=fast), self.assertRaises(ValueError): + self.embed() + finally: + if fast: + self.server._inference_lock.release() + self.assertEqual(self.calls, []) + self.assertEqual(self.server._embed_timestamps, []) + self.assert_released() + + def test_fast_cached_identity_mismatch_and_width_failure_never_fall_back(self) -> None: + self.embed() + self.server._inference_lock.acquire() + try: + self.embed() + self.server._fast_model_id = "qwen3_256" + with self.assertRaises(ValueError): + self.embed() + self.server._fast_model_id = "model2vec" + self.width = 1 + with self.assertRaises(ValueError): + self.embed() + finally: + self.server._inference_lock.release() + self.assertEqual(self.server._embed_timestamps, []) + self.assert_released() + + def test_empty_selected_embedding_requires_receipt_without_added_probe(self) -> None: + response = self.embed(0) + self.assertEqual(response["vectors"].shape, (0, 256)) + self.assertEqual(response["target"], self.edge.to_wire()) + self.assertEqual([call[2] for call in self.calls], [[]]) + self.client._request_with_autostart = Mock(return_value={"ok": True, "vectors": response["vectors"], "protocol": 1}) + with self.assertRaises(self.client.ProtocolMismatchError): + self.client.CertifiedEmbeddingProxy(self.edge).encode([]) + + def test_embedding_oom_fast_and_main_retry_keep_receipt_order_and_owner(self) -> None: + for count in (1, 5): + self.calls.clear() + self.embed_oom = True + response = self.embed(count, op="embed_batched", batch_size=2) + self.assertEqual(response["target"], self.edge.to_wire()) + self.assertEqual(response["vectors"].rows, [[float(i)] * 256 for i in range(count)]) + self.assertEqual([len(call[2]) for call in self.calls], [1, 1] if count == 1 else [2, 2, 2, 1]) + self.assertIn("embed", self.server._sticky_cpu) + self.assert_released() + + def test_every_main_slice_checks_width_and_never_returns_partial_output(self) -> None: + def change_width() -> None: + if self.calls: + self.width = 1 + + self.before_encode = change_width + with self.assertRaises(ValueError): + self.embed(5, op="embed_batched", batch_size=2) + self.assertEqual([len(call[2]) for call in self.calls], [2, 2]) + self.assert_released() + + def test_admission_and_deadline_refusals_have_no_receipt_or_fallback(self) -> None: + response = self.embed(deadline=1) + self.assertFalse(response["ok"]) + self.assertNotIn("target", response) + self.assertEqual(self.calls, []) + self.embed() + self.server._inference_lock.acquire() + self.server._check_process_memory = Mock(side_effect=self.ms._ProcessMemoryRefused("synthetic capacity")) + try: + response = self.embed() + finally: + self.server._inference_lock.release() + self.assertEqual(response["error_code"], "server_busy") + self.assertNotIn("target", response) + self.assertEqual(len(self.calls), 1) + self.assertEqual(self.server._embed_timestamps, []) + self.assert_released() + + def test_ordinary_clients_and_responses_remain_unattested(self) -> None: + self.client._request_with_autostart = Mock(return_value={"ok": True, "vectors": "legacy", "scores": "legacy"}) + self.assertEqual(self.client.EmbeddingProxy("edge").encode("s0"), "legacy") + self.assertNotIn("expected_target", self.client._request_with_autostart.call_args.args[0]) + self.assertEqual(self.client.RerankerProxy().predict([["q", "s0"]]), "legacy") + self.assertNotIn("expected_model_name", self.client._request_with_autostart.call_args.args[0]) + response = self.server.handle_request({"op": "embed", "tier": "model2vec", "texts": ["s0"]}) + self.assertEqual(set(response), {"ok", "vectors"}) + response = self.server.handle_request({"op": "rerank", "pairs": [["q", "s0"]]}) + self.assertEqual(set(response), {"ok", "scores"}) + + def test_reranker_every_response_certifies_explicit_name_and_bounded_order(self) -> None: + proxy = self.client.CertifiedRerankerProxy("synthetic/reranker") + response = proxy.predict([["q", f"s{i}"] for i in range(5)], timeout=3, batch_size=2) + self.assertEqual(response.rows, [[float(i)] for i in range(5)]) + self.assertEqual([len(call[2]) for call in self.calls], [2, 2, 1]) + self.assertEqual(self.builds, [("synthetic/reranker", "cpu")]) + for receipt in (None, {"version": 1, "model_name": "synthetic/default"}, + {"version": True, "model_name": "synthetic/reranker"}, + {"version": 1, "model_name": "synthetic/reranker", "extra": 1}): + self.client._request_with_autostart = Mock(return_value={"ok": True, "scores": response, "protocol": 1, "reranker": receipt}) + with self.subTest(receipt=receipt), self.assertRaises(self.client.ProtocolMismatchError): + proxy.predict([["q", "s0"]]) + self.client._request_with_autostart.assert_called_once() + self.assert_released() + + def test_reranker_receipt_requires_one_result_row_per_pair(self) -> None: + proxy = self.client.CertifiedRerankerProxy("synthetic/reranker") + pairs = [["q", "s0"], ["q", "s1"]] + receipt = {"version": 1, "model_name": "synthetic/reranker"} + for scores in (None, [0., 1.], Array([], ()), Array([[0.]], (1, 1)), + Array([[0.], [1.], [2.]], (3, 1))): + self.client._request_with_autostart = Mock(return_value={ + "ok": True, "scores": scores, "protocol": 1, "reranker": receipt, + }) + with self.subTest(shape=getattr(scores, "shape", None)), self.assertRaises(self.client.ProtocolMismatchError): + proxy.predict(iter(pairs)) + self.client._request_with_autostart.assert_called_once() + for shape in ((2,), (2, 1), (2, 2)): + scores = Array([[0.], [1.]], shape) + self.client._request_with_autostart = Mock(return_value={ + "ok": True, "scores": scores, "protocol": 1, "reranker": receipt, + }) + with self.subTest(shape=shape): + self.assertIs(proxy.predict(iter(pairs)), scores) + self.assert_released() + + def test_empty_reranker_and_missing_protocol_are_strict(self) -> None: + response = self.rerank(0) + self.assertEqual(response["scores"].shape, (0, 1)) + self.assertEqual(response["reranker"], {"version": 1, "model_name": "synthetic/reranker"}) + self.assertEqual([call[2] for call in self.calls], [[]]) + for value in (None, True, 1.0, "1", 2): + self.client._request_with_autostart = Mock(return_value={**response, "protocol": value}) + with self.subTest(protocol=value), self.assertRaises(self.client.ProtocolMismatchError): + self.client.CertifiedRerankerProxy("synthetic/reranker").predict([]) + + def test_reranker_explicit_identity_validation_and_effective_cache_mismatch(self) -> None: + for name in (None, "", True, "x" * 513, "synthetic/model/other", "bad\nname"): + with self.subTest(name=name), self.assertRaises(ValueError): + self.client.CertifiedRerankerProxy(name) + with self.subTest(name=name), self.assertRaises(ValueError): + self.rerank(expected_model_name=name) + with self.assertRaises(ValueError): + self.rerank(model_name=None) + self.assertEqual(self.builds, []) + real_get = self.server._get_reranker + + def wrong(name=None, deadline=None): + model = real_get(name, deadline) + self.server._reranker_name = "synthetic/different" + return model + + self.server._get_reranker = wrong + with self.assertRaises(ValueError): + self.rerank() + self.assertEqual(self.calls, []) + self.assert_released() + + def test_reranker_cpu_replacement_is_revalidated_before_retry(self) -> None: + self.rerank_oom = True + response = self.rerank(5, op="rerank_batched", batch_size=2) + self.assertEqual(response["reranker"], {"version": 1, "model_name": "synthetic/reranker"}) + self.assertEqual(response["scores"].rows, [[float(i)] for i in range(5)]) + self.assertEqual([len(call[2]) for call in self.calls], [2, 2, 2, 1]) + self.assertEqual([name for name, _device in self.builds], ["synthetic/reranker"] * 2) + real_get = self.server._get_reranker + + def wrong_retry(name=None, deadline=None): + replacing = self.server._reranker is None + model = real_get(name, deadline) + if replacing: + self.server._reranker_name = "synthetic/other" + return model + + self.server._get_reranker = wrong_retry + self.rerank_oom = True + before = len(self.calls) + with self.assertRaises(ValueError): + self.rerank() + self.assertEqual(len(self.calls), before + 1) + self.assert_released() + + def test_public_tier_receipt_does_not_collapse_base_and_pro(self) -> None: + response = self.embed(target=self.base) + model = self.server._embed_state.model + pro = dataclasses.replace(self.base, tier="pro") + second = self.embed(target=pro) + self.assertEqual(response["target"], self.base.to_wire()) + self.assertEqual(second["target"], pro.to_wire()) + self.assertIs(self.server._embed_state.model, model) + self.client._request_with_autostart = Mock(return_value=second) + with self.assertRaises(self.client.ProtocolMismatchError): + self.client.CertifiedEmbeddingProxy(self.base).encode("s0") + + +if __name__ == "__main__": + unittest.main() diff --git a/truememory/model_client.py b/truememory/model_client.py index 778c001b..68494a92 100644 --- a/truememory/model_client.py +++ b/truememory/model_client.py @@ -673,6 +673,59 @@ def encode( return vectors +def _check_serving_protocol(response: dict) -> None: + if type(response.get("protocol")) is not int or response["protocol"] != PROTOCOL_VERSION: + raise ProtocolMismatchError("Model server omitted the selected-serving protocol; restart it after upgrading") + + +def _serving_reranker_name(value: object) -> str: + import re + + if (type(value) is not str or len(value.encode("utf-8")) > 512 + or re.fullmatch(r"[\w][\w.\-]*(/[\w][\w.\-]*)?", value) is None): + raise ValueError("Selected reranking requires an explicit model identity") + return value + + +class CertifiedEmbeddingProxy(EmbeddingProxy): + """Built-in serving with a fresh identity receipt on every fast/main result. + + Construction configures identity only; it does not contact the daemon or + prove model readiness. Custom targets use PreparedEmbeddingProxy instead. + """ + + def __init__(self, target: "EmbeddingTarget") -> None: + from truememory.embedding_target import EmbeddingTarget, EmbeddingTargetError + + self.target = EmbeddingTarget.from_wire(target.to_wire()) + if self.target.tier == "custom": + raise EmbeddingTargetError("Custom serving requires the prepared target protocol") + super().__init__(tier=self.target.model_id) + + def encode(self, texts, timeout: float | None = None, **kwargs) -> np.ndarray: + from truememory.embedding_target import EmbeddingTarget, EmbeddingTargetError, check_target_vectors + + if isinstance(texts, str): + texts = [texts] + texts = list(texts) + request = _batch_request({ + "op": "embed", "texts": texts, "tier": self.target.model_id, + "expected_target": self.target.to_wire(), + }, kwargs) + response = _request_with_autostart(request, timeout=timeout) + _check_model_response(response, request) + _check_serving_protocol(response) + try: + effective = EmbeddingTarget.from_wire(response.get("target")) + except EmbeddingTargetError as exc: + raise ProtocolMismatchError("Model server omitted a valid serving target receipt") from exc + if effective != self.target: + raise ProtocolMismatchError("Model server returned a different serving target") + vectors = response["vectors"] + check_target_vectors(self.target, vectors, len(texts)) + return vectors + + class RerankerProxy: """Drop-in replacement for CrossEncoder with .predict() method.""" @@ -695,6 +748,32 @@ def predict(self, pairs, timeout: float | None = None, **kwargs) -> np.ndarray: return resp["scores"] +class CertifiedRerankerProxy(RerankerProxy): + """Explicit reranker receipt per response, without eager native readiness.""" + + def __init__(self, model_name: str) -> None: + super().__init__(model_name=_serving_reranker_name(model_name)) + + def predict(self, pairs, timeout: float | None = None, **kwargs) -> np.ndarray: + request = _batch_request({ + "op": "rerank", "pairs": list(pairs), "model_name": self._model_name, + "expected_model_name": self._model_name, + }, kwargs) + response = _request_with_autostart(request, timeout=timeout) + _check_model_response(response, request) + _check_serving_protocol(response) + receipt = response.get("reranker") + if (type(receipt) is not dict or set(receipt) != {"version", "model_name"} + or type(receipt["version"]) is not int or receipt["version"] != 1 + or type(receipt["model_name"]) is not str or receipt["model_name"] != self._model_name): + raise ProtocolMismatchError("Model server omitted the exact serving reranker receipt") + scores = response.get("scores") + if (not isinstance(scores, np.ndarray) or not scores.shape + or scores.shape[0] != len(request["pairs"])): + raise ProtocolMismatchError("Model server returned a different serving reranker result count") + return scores + + def use_model_server() -> bool: """Return sharing policy, independently of endpoint readiness. diff --git a/truememory/model_server.py b/truememory/model_server.py index 30f41f42..a4b4f502 100644 --- a/truememory/model_server.py +++ b/truememory/model_server.py @@ -816,7 +816,9 @@ def _get_prepared_embed_model(self, target: "EmbeddingTarget", deadline: _Reques )) return model - def _get_fast_encoder(self, tier: str, deadline: _RequestDeadline | None = None): + def _get_fast_encoder( + self, tier: str, deadline: _RequestDeadline | None = None, *, expected_target: "EmbeddingTarget | None" = None, + ): """CPU-resident encoder for the single-text fast lane (issue #577). Loaded lazily on the first single-text request that finds the global @@ -842,9 +844,15 @@ def _get_fast_encoder(self, tier: str, deadline: _RequestDeadline | None = None) return None model_id, generation = state.model_id, state.generation + if expected_target is not None and model_id != expected_target.model_id: + from truememory.embedding_target import EmbeddingTargetError + raise EmbeddingTargetError("Fast encoder residency does not match the serving target") # A slow CPU constructor must not retain the old main model snapshot. state = None if self._fast_encoder is not None and self._fast_generation is generation: + if expected_target is not None and self._fast_model_id != model_id: + from truememory.embedding_target import EmbeddingTargetError + raise EmbeddingTargetError("Fast encoder cache does not match the serving target") return self._fast_encoder self._fast_encoder = None @@ -900,8 +908,40 @@ def _get_reranker( log.info("Loaded reranker model=%s device=%s", name, device) return self._reranker + def _check_serving_embed_model(self, model: object, target: "EmbeddingTarget | None") -> None: + if target is None: + return + from truememory.embedding_target import EmbeddingTargetError + + state = self._embed_state + if state is None or state.model is not model or state.model_id != target.model_id: + raise EmbeddingTargetError("Effective encoder does not match the serving target") + + @staticmethod + def _serving_embed_response(vectors: object, count: int, target: "EmbeddingTarget | None") -> dict: + if target is None: + return {"ok": True, "vectors": vectors} + from truememory.embedding_target import check_target_vectors + + check_target_vectors(target, vectors, count) + return {"ok": True, "vectors": vectors, "target": target.to_wire(), "protocol": PROTOCOL_VERSION} + + @staticmethod + def _serving_reranker_name(value: object) -> str: + import re + + if (type(value) is not str or len(value.encode("utf-8")) > 512 + or re.fullmatch(r"[\w][\w.\-]*(/[\w][\w.\-]*)?", value) is None): + raise ValueError("Selected reranking requires an explicit model identity") + return value + + def _check_serving_reranker(self, model: object, expected: str | None) -> None: + if expected is not None and (self._reranker is not model or self._reranker_name != expected): + raise ValueError("Effective reranker does not match the serving target") + def _handle_fast_embed( self, texts: list, tier: str, deadline: _RequestDeadline | None = None, + *, expected_target: "EmbeddingTarget | None" = None, ) -> dict | None: """Single-text fast lane (issue #577). @@ -927,6 +967,7 @@ def _handle_fast_embed( self._preflight_embed_result(tier, len(texts)) deadline.check() model = self._get_embed_model(tier, deadline=deadline) + self._check_serving_embed_model(model, expected_target) self._check_process_memory("embed slice", deadline) deadline.check() recover = False @@ -945,7 +986,9 @@ def _handle_fast_embed( self._check_process_memory("embed CPU retry", deadline) deadline.check() vectors = model.encode(texts, batch_size=1, show_progress_bar=False) - return {"ok": True, "vectors": _checked_batch(vectors, len(texts), len(texts), "vectors")} + return self._serving_embed_response( + _checked_batch(vectors, len(texts), len(texts), "vectors"), len(texts), expected_target, + ) finally: model = None self._lock.release() @@ -961,7 +1004,8 @@ def _handle_fast_embed( model = None try: deadline.start_inference() - model = self._get_fast_encoder(tier, deadline=deadline) + model = (self._get_fast_encoder(tier, deadline=deadline) if expected_target is None else + self._get_fast_encoder(tier, deadline=deadline, expected_target=expected_target)) deadline.check() if model is None: if deadline.transport is not None: @@ -971,15 +1015,26 @@ def _handle_fast_embed( self._check_process_memory("fast embed slice", deadline) deadline.check() vectors = model.encode(texts, batch_size=1, show_progress_bar=False) + if expected_target is not None: + response = self._serving_embed_response( + _checked_batch(vectors, len(texts), len(texts), "vectors"), len(texts), expected_target, + ) finally: model = None finally: if fast_owned: self._retire_stale_fast_encoder() - return {"ok": True, "vectors": _checked_batch(vectors, len(texts), len(texts), "vectors")} + return (response if expected_target is not None else + {"ok": True, "vectors": _checked_batch(vectors, len(texts), len(texts), "vectors")}) except (_RequestDeadlineExceeded, _ResultTooLarge, _ProcessMemoryRefused): raise - except Exception: + except Exception as exc: + if expected_target is not None: + from truememory.embedding_target import EmbeddingTargetError + if isinstance(exc, EmbeddingTargetError): + # An identity/shape rejection cannot be relabelled by a + # later fallback. Other fast-lane failures still retry. + raise if deadline.transport is not None: deadline.transport.started = was_started log.warning( @@ -1024,6 +1079,7 @@ def _handle_request_inner(self, request: dict) -> dict: if op in ("embed", "embed_batched", "prepare_embed_target_v1", "embed_target_v1"): target = None + expected_target = None preparing = op == "prepare_embed_target_v1" if op in ("prepare_embed_target_v1", "embed_target_v1"): from truememory.embedding_target import ( @@ -1041,6 +1097,13 @@ def _handle_request_inner(self, request: dict) -> dict: else: texts = request["texts"] tier = request.get("tier", "") + if "expected_target" in request: + from truememory.embedding_target import EmbeddingTarget, EmbeddingTargetError, check_target_vectors + expected_target = EmbeddingTarget.from_wire(request["expected_target"]) + if expected_target.tier == "custom" or tier != expected_target.model_id: + raise EmbeddingTargetError("Selected serving requires its canonical built-in model identity") + if not isinstance(texts, list) or not all(isinstance(text, str) for text in texts): + raise EmbeddingTargetError("Selected embedding inputs must be a list of strings") try: batch_limit = _batch_limit(request.get("batch_size", _EMBED_BATCH_LIMIT), _EMBED_BATCH_LIMIT) @@ -1050,7 +1113,8 @@ def _handle_request_inner(self, request: dict) -> dict: # Single-text fast lane (issue #577): hook recall queries must # never queue behind batch ingestion work or OOM recovery. if target is None and len(texts) <= self._FAST_LANE_MAX_TEXTS: - fast = self._handle_fast_embed(texts, tier, deadline) + fast = (self._handle_fast_embed(texts, tier, deadline) if expected_target is None else + self._handle_fast_embed(texts, tier, deadline, expected_target=expected_target)) if fast is not None: return fast @@ -1090,6 +1154,7 @@ def _handle_request_inner(self, request: dict) -> dict: deadline.check() model = (self._get_embed_model(tier, deadline=deadline) if target is None else self._get_prepared_embed_model(target, deadline)) + self._check_serving_embed_model(model, expected_target) order = self._embed_global_order(model, texts, limit, deadline) while offset < len(texts) or vectors is None: deadline.check() @@ -1113,8 +1178,8 @@ def _handle_request_inner(self, request: dict) -> dict: self._recover_embed_oom_locked(model, deadline) retry = batch break - if target is not None: - check_target_vectors(target, values, len(batch)) + if target is not None or expected_target is not None: + check_target_vectors(target or expected_target, values, len(batch)) vectors = _store_batch_result( vectors, values, offset, len(batch), len(texts), indices=indices, ) @@ -1126,8 +1191,8 @@ def _handle_request_inner(self, request: dict) -> dict: deadline.check() log.warning("MPS OOM during encoding; retrying microbatch on CPU") values = model.encode(retry, batch_size=limit, show_progress_bar=False) - if target is not None: - check_target_vectors(target, values, len(retry)) + if target is not None or expected_target is not None: + check_target_vectors(target or expected_target, values, len(retry)) vectors = _store_batch_result( vectors, values, offset, len(retry), len(texts), indices=indices, ) @@ -1143,13 +1208,18 @@ def _handle_request_inner(self, request: dict) -> dict: return ({"ok": True, "target": target.to_wire()} if preparing else {"ok": True, "target": target.to_wire(), "vectors": vectors[:0] if empty_target_input else vectors}) - return {"ok": True, "vectors": vectors} + return self._serving_embed_response(vectors, len(texts), expected_target) finally: model = None if op in ("rerank", "rerank_batched"): pairs = request["pairs"] model_name = request.get("model_name") + expected_name = None + if "expected_model_name" in request: + expected_name = self._serving_reranker_name(request["expected_model_name"]) + if model_name != expected_name: + raise ValueError("Selected reranking requires its explicit model identity") try: batch_limit = _batch_limit(request.get("batch_size", _RERANK_BATCH_LIMIT), _RERANK_BATCH_LIMIT) @@ -1172,6 +1242,7 @@ def _handle_request_inner(self, request: dict) -> dict: self._preflight_rerank_result(model_name, len(pairs)) deadline.check() reranker = self._get_reranker(model_name, deadline=deadline) + self._check_serving_reranker(reranker, expected_name) order = self._rerank_global_order(reranker, pairs, limit, deadline) if order is not None: recovery_name = self._reranker_name @@ -1204,6 +1275,7 @@ def _handle_request_inner(self, request: dict) -> dict: flush_mps_cache() deadline.check() reranker = self._get_reranker(recovery_name, deadline=deadline) + self._check_serving_reranker(reranker, expected_name) retry = batch break if indices is None: @@ -1225,7 +1297,9 @@ def _handle_request_inner(self, request: dict) -> dict: ) offset += len(retry) self._after_request_batches(throttler, len(pairs), predict_start, deadline) - return {"ok": True, "scores": scores} + return ({"ok": True, "scores": scores} if expected_name is None else + {"ok": True, "scores": scores, "protocol": PROTOCOL_VERSION, + "reranker": {"version": 1, "model_name": expected_name}}) finally: reranker = None From df1a1132f623506fde3c4e91c186255fb5455f64 Mon Sep 17 00:00:00 2001 From: buildingjoshbetter Date: Fri, 9 Oct 2026 11:55:15 -0500 Subject: [PATCH 3/5] fix: Pin selected runtime operations and connection ownership --- docs/resource-budgets.md | 81 ++ tests/test-l0-style-accumulator-747.py | 8 +- tests/test-l0-style-incremental-747.py | 5 +- tests/test-maintenance-engine-routing-753.py | 28 +- .../test-maintenance-fixture-isolation-753.py | 19 +- tests/test-model-server-sharing-741.py | 15 +- tests/test-mps-budget-743.py | 15 +- tests/test-prepared-target-795.py | 2 + ...test-rebuild-foreground-publication-756.py | 2 +- tests/test-rebuild-source-streaming-756.py | 8 +- tests/test-tier-runtime-bridge-795.py | 1159 +++++++++++++++++ tests/test-tier-writer-fence-795.py | 2 +- truememory/__init__.py | 1 + truememory/engine.py | 414 +++--- truememory/hybrid.py | 3 + truememory/reranker.py | 144 +- truememory/tier_switch/runtime.py | 419 ++++++ truememory/tier_switch/serving.py | 31 +- truememory/vector_search.py | 268 +++- 19 files changed, 2325 insertions(+), 299 deletions(-) create mode 100644 tests/test-tier-runtime-bridge-795.py create mode 100644 truememory/tier_switch/runtime.py diff --git a/docs/resource-budgets.md b/docs/resource-budgets.md index 7ecccc5c..5ce44cec 100644 --- a/docs/resource-budgets.md +++ b/docs/resource-budgets.md @@ -2231,3 +2231,84 @@ empty results still require receipts, without a new synthetic probe. Tests use stdlib loaders and synthetic model/array doubles; native numerical parity and full integration remain separate GPUBox gates. No activation call or measured memory reduction is claimed by this protocol extension. + + +### Selected runtime bridge and whole-operation serving leases (#795, Unit B) + +`serving_operation(conn, *, deadline=None, cancelled=None, reranker_id=None)` +reads the bounded selected journal before admission and verifies it again after +waiting. An admitted operation freezes the public tier, embedding descriptor, +exact completion/separation pair and effective reranker. Nested calls on the +same accepted connection reuse that snapshot without repeating journal/schema +reads. Another connection is refused unless explicitly admitted through a child +reservation. Same-width Edge/Base embeddings never establish identity by width. + +Selected runtime projection uses the existing process singleton slots. It +releases a replaced active embedding reference before construction, uses the +existing local preparation slot and model ownership for the bounded width probe, +and publishes only the frozen descriptor. Base/Pro reuse the same exact embedding +identity without a second native load. Built-in daemon embeddings use certified +responses on the existing main/fast request paths; custom embeddings retain the +strict prepared-target protocol. A restart cannot silently accept an old daemon +response lacking the expected identity receipt. Selected rerankers similarly use +certified per-response names; local construction uses the frozen explicit name. +No model, dtype, batch policy, ranking formula or result limit changes here. + +`apply_frozen_selection(selection, *, deadline=None, cancelled=None, +reranker_id=None)` may run inside the same thread's exclusive activation lease, +but never upgrades a serving reader. It has no connection argument: callers must +establish that they own no SQL transaction. The process acknowledgement is cleared +before partial projection and set only after both local slots match. With daemon +proxies this acknowledges configured routing identity, not daemon model residency, +native warmup, config-file persistence or activation completion. A failure after +DB selection remains pending. The optional outer reranker override captures the +existing deep/explicit model before admission without changing the journal's tier +default; nested calls inherit it and refuse a different explicit override. + +Engine add/update/delete/delete_all and search variants, direct vector search and +embed_single, and the entire hybrid encode/completion/separation operation now +hold serving leases. Cached Engine handles reconcile selection too. Selected +initialization validates the exact prepared pair before serving and skips legacy +vector compatibility writes/migrations; its cache is bound to both generation +and the accepted connection. Replacement handles reload their SQLite extension +and validate again. Existing startup/cached maintenance scheduling remains. + +`current_operation(conn)` distinguishes unbound from a bound legacy `None` +selection and rejects the wrong connection. Child reservations pin the group +before submission and may outlive the parent. A clean child connection must match +the captured database-path identity and committed selection before its body; this +is cooperative pathname binding, not native SQLite-descriptor or hostile-clone +proof. Unused or never-admitted reservations remain caller-owned and must close. +Thread-local state is restored on body `BaseException`; forked process state and +inherited tokens cannot authorize work. + +A borrowed SQL transaction receives one nonblocking gate/connection-lock +admission attempt, never a native reconciliation or transaction commit/rollback. +Cached legacy work remains available when coherent; a process already projected +to a selected target refuses an unrelated unselected database instead of guessing +its identity. Direct embedding/reranker setters, unloaders and slot replacement +join exclusive admission. `blocking=False` bounds admission attempts only. +Deadline/cancellation checks run before and after projection and after final +journal/identity reads. The 0.05-second wait interval is not a hard return bound: +condition reacquisition, cleanup, scheduling and an already-entered native call +can delay return. The prepared embedding API carries a deadline, not an Event; +cancellation during its internal wait/native work is observed on return before +another projection or serving body is admitted. + +This is partial boundary adoption, not a completed public tier switch. Remaining +Unit C work includes maintenance-owner internal serving admission, rebuild and +migration owners, ingest/import and EncodingGate cached references, MCP decisions +and explicit parallel-child binding, and manager/config/CLI activation. Those +paths must preserve maintenance owner -> serving lease -> Engine writer -> SQL +writer order. Unbound selected direct initialization is not a serving contract. +The separate writer-fence prerequisite establishes SQL writer ownership; a +read-only selection comparison or `in_transaction` alone cannot prove it. + +Validation uses stdlib production-source stubs and in-memory journals, including +partial projection, borrowed transactions, overrides, child controls, final-read +cancellation/deadlines, replacement refusal and selected reconnect. A separate +opt-in GPUBox test composes the real Engine/runtime/vector code with native +sqlite-vec and synthetic encoders through CRUD, nested search and reconnect. +Local checks do not load native models or file-backed databases. This checkpoint +does not prove reduced unified-memory use, native MPS acceptance, or resolution +of the original 26.24 GB Mac incident. diff --git a/tests/test-l0-style-accumulator-747.py b/tests/test-l0-style-accumulator-747.py index f343405e..037f61e2 100644 --- a/tests/test-l0-style-accumulator-747.py +++ b/tests/test-l0-style-accumulator-747.py @@ -34,7 +34,8 @@ def safe_import(name: str, globals: dict | None = None, locals: dict | None = No module = types.ModuleType("synthetic_style_" + name) module.__dict__["__builtins__"] = dict(vars(builtins), __import__=safe_import) source = ROOT / "truememory" / (name + ".py") - exec(compile(source.read_text(encoding="utf-8"), str(source), "exec"), module.__dict__) + with patch.dict(sys.modules, {module.__name__: module}): + exec(compile(source.read_text(encoding="utf-8"), str(source), "exec"), module.__dict__) return module @@ -702,7 +703,8 @@ def test_actual_engine_add_catches_style_deferral_and_keeps_source(self) -> None maintenance = types.ModuleType("synthetic_append_maintenance") modules = {"truememory.personality_style_vec": STYLE, "truememory.rebuild_source": rebuild, "truememory.maintenance": maintenance, "truememory.storage": STORAGE, - "truememory._platform": load_stdlib_source("_platform")} + "truememory._platform": load_stdlib_source("_platform"), + "truememory.tier_switch.writer": load_stdlib_source("tier_switch/writer")} def safe_import(name: str, globals: dict | None = None, locals: dict | None = None, fromlist: tuple[str, ...] = (), level: int = 0) -> object: @@ -718,7 +720,7 @@ def safe_import(name: str, globals: dict | None = None, locals: dict | None = No tree = ast.parse((ROOT / "truememory/engine.py").read_text(encoding="utf-8")) engine_class = next(node for node in tree.body if isinstance(node, ast.ClassDef) and node.name == "TrueMemoryEngine") add = next(node for node in engine_class.body if isinstance(node, ast.FunctionDef) and node.name == "add") - namespace = {"__builtins__": dict(vars(builtins), __import__=safe_import), "MAX_CONTENT_LENGTH": 50000, + namespace = {"engine_operation": lambda function: function, "__builtins__": dict(vars(builtins), __import__=safe_import), "MAX_CONTENT_LENGTH": 50000, "insert_message": STORAGE.insert_message, "_update_style_vec": STYLE.update_entity_style_vector_incremental, "logger": logging.getLogger("synthetic_style_engine")} exec(compile(ast.Module(body=[add], type_ignores=[]), "engine.py", "exec"), namespace) diff --git a/tests/test-l0-style-incremental-747.py b/tests/test-l0-style-incremental-747.py index 360f8652..943bf64c 100644 --- a/tests/test-l0-style-incremental-747.py +++ b/tests/test-l0-style-incremental-747.py @@ -30,7 +30,8 @@ def load_add() -> types.FunctionType: selected = [node for node in source.body if isinstance(node, ast.FunctionDef) and node.name == "rebuild_transaction"] exec(compile(ast.Module(body=selected, type_ignores=[]), "synthetic-rebuild", "exec"), rebuild.__dict__) modules = {"truememory.personality_style_vec": STYLE, "truememory.rebuild_source": rebuild, - "truememory.maintenance": MAINTENANCE} + "truememory.maintenance": MAINTENANCE, + "truememory.tier_switch.writer": ACCUMULATOR["load_stdlib_source"]("tier_switch/writer")} def safe_import(name: str, globals: dict | None = None, locals: dict | None = None, fromlist: tuple = (), level: int = 0) -> object: @@ -43,7 +44,7 @@ def safe_import(name: str, globals: dict | None = None, locals: dict | None = No tree = ast.parse((ROOT / "truememory/engine.py").read_text(encoding="utf-8")) engine = next(node for node in tree.body if isinstance(node, ast.ClassDef) and node.name == "TrueMemoryEngine") add = next(node for node in engine.body if isinstance(node, ast.FunctionDef) and node.name == "add") - namespace = {"__builtins__": dict(vars(builtins), __import__=safe_import), "MAX_CONTENT_LENGTH": 50000, + namespace = {"engine_operation": lambda function: function, "__builtins__": dict(vars(builtins), __import__=safe_import), "MAX_CONTENT_LENGTH": 50000, "insert_message": STORAGE.insert_message, "_update_style_vec": STYLE.update_entity_style_vector_incremental, "logger": logging.getLogger("synthetic-style-incremental")} exec(compile(ast.Module(body=[add], type_ignores=[]), "actual-engine-add", "exec"), namespace) diff --git a/tests/test-maintenance-engine-routing-753.py b/tests/test-maintenance-engine-routing-753.py index 5c8c02ff..91238dc2 100644 --- a/tests/test-maintenance-engine-routing-753.py +++ b/tests/test-maintenance-engine-routing-753.py @@ -16,7 +16,7 @@ import unittest import warnings import weakref -from contextlib import contextmanager +from contextlib import contextmanager, nullcontext from pathlib import Path from unittest.mock import patch @@ -44,6 +44,29 @@ def safe_import(name, globals=None, locals=None, fromlist=(), level=0): def load_engine(modules: dict) -> types.ModuleType: + modules["truememory.tier_switch.writer"] = stdlib_module("tier_switch/writer", modules) + serving_path = ROOT / "truememory/tier_switch/serving.py" + serving_source = ast.parse(serving_path.read_text(encoding="utf-8")) + serving_names = {"ServingLeaseError", "ServingLeaseTimeout", "ServingLeaseCancelled"} + serving_nodes = [node for node in serving_source.body + if isinstance(node, ast.ClassDef) and node.name in serving_names] + serving_namespace = {} + exec(compile(ast.Module(body=serving_nodes, type_ignores=[]), str(serving_path), "exec"), serving_namespace) + runtime_path = ROOT / "truememory/tier_switch/runtime.py" + runtime_source = ast.parse(runtime_path.read_text(encoding="utf-8")) + runtime_names = {"TierRuntimeError", "ConnectionWriteLock", "_check", "_remaining", "_connection_read"} + runtime_nodes = [ast.ImportFrom(module="__future__", names=[ast.alias(name="annotations")], level=0)] + runtime_nodes.extend(node for node in runtime_source.body + if isinstance(node, (ast.ClassDef, ast.FunctionDef)) and node.name in runtime_names) + runtime_namespace = {"threading": threading, "sqlite3": sqlite3, "time": time, + "contextmanager": contextmanager, + "serving": types.SimpleNamespace(**{name: serving_namespace[name] for name in serving_names})} + exec(compile(ast.fix_missing_locations(ast.Module(body=runtime_nodes, type_ignores=[])), + str(runtime_path), "exec"), runtime_namespace) + modules["truememory.tier_switch.runtime"] = types.SimpleNamespace( + **{name: runtime_namespace[name] for name in runtime_names}, + serving_operation=lambda *args, **kwargs: nullcontext(types.SimpleNamespace(selection=None)), + current_operation=lambda conn: types.SimpleNamespace(selection=None)) module = types.ModuleType("synthetic_routing_engine") def safe_import(name, globals=None, locals=None, fromlist=(), level=0): @@ -56,6 +79,7 @@ def safe_import(name, globals=None, locals=None, fromlist=(), level=0): module.__dict__.update({"__builtins__": dict(vars(builtins), __import__=safe_import), "__file__": str(ROOT / "truememory/engine.py"), "Path": Path, "sqlite3": sqlite3, "threading": threading, "time": time, "os": os, "warnings": warnings, + "engine_operation": lambda function: function, "logger": logging.getLogger("synthetic-routing"), "MAX_CONTENT_LENGTH": 50_000, "_env_int": lambda name, default, **kwargs: default}) source = ast.parse((ROOT / "truememory/engine.py").read_text(encoding="utf-8")) @@ -597,7 +621,7 @@ def test_source_commit_followed_by_a17_vector_error_still_notifies(self) -> None changed = type("VectorPublicationChanged", (RuntimeError,), {}) @contextmanager - def publication(*args): + def publication(*args, **kwargs): raise changed("synthetic model replacement") yield diff --git a/tests/test-maintenance-fixture-isolation-753.py b/tests/test-maintenance-fixture-isolation-753.py index f52397c5..9151bc6c 100644 --- a/tests/test-maintenance-fixture-isolation-753.py +++ b/tests/test-maintenance-fixture-isolation-753.py @@ -34,7 +34,24 @@ def safe_import(name: str, globals: dict | None = None, locals: dict | None = No def safe_exec(code: object, namespace: dict) -> None: namespace.setdefault("__builtins__", fixture_builtins) - exec(code, namespace) + name = namespace.get("__name__", "") + module = process_sys.modules.get(name) + if (not name.startswith("synthetic_") or not isinstance(module, types.ModuleType) + or module.__dict__ is not namespace): + exec(code, namespace) + return + # stdlib dataclasses resolves postponed annotations through real sys. + # Expose only this temporary synthetic namespace, never application peers. + absent = object() + previous = sys.modules.get(name, absent) + sys.modules[name] = module + try: + exec(code, namespace) + finally: + if previous is absent: + del sys.modules[name] + else: + sys.modules[name] = previous fixture_builtins = dict(vars(builtins), __import__=safe_import, exec=safe_exec) diff --git a/tests/test-model-server-sharing-741.py b/tests/test-model-server-sharing-741.py index 2de2f63f..35513415 100644 --- a/tests/test-model-server-sharing-741.py +++ b/tests/test-model-server-sharing-741.py @@ -146,13 +146,22 @@ def loaders(self) -> Iterator[tuple[types.ModuleType, types.ModuleType, Mock]]: mps.ensure_mps_memory_budget = lambda _device: None namespace = {"_model": None, "_lock": threading.Lock(), "_embedding_dim": 256, "EMBEDDING_MODEL": "qwen3_256", "os": os, "sys": sys, - "_model_name": None, "get_current_reranker_name": lambda: "synthetic-reranker"} + "_model_name": None, "_model_certified": False, "_frozen_embedding_target": None, + "contextmanager": contextmanager, + "get_current_reranker_name": lambda: "synthetic-reranker"} vectors = load_definitions("vector_search.py", dict(namespace), {"get_model"}) - reranker = load_definitions("reranker.py", dict(namespace), {"get_reranker"}) + reranker = load_definitions("reranker.py", dict(namespace), {"get_reranker", "_reranker_load_lock"}) + serving = types.ModuleType("synthetic_serving_741") + with patch.dict(sys.modules, {serving.__name__: serving}): + exec(compile((SOURCE / "tier_switch/serving.py").read_text(), "synthetic-serving", "exec"), serving.__dict__) + runtime = types.SimpleNamespace(current_runtime_operation=lambda: None, + require_model_load_allowed=lambda: None, TierRuntimeError=RuntimeError) package = types.ModuleType("truememory") package.__path__ = [] with patch.dict(sys.modules, {"truememory": package, "truememory.model_client": self.client, - "truememory.mps_utils": mps, "sentence_transformers": transformers}): + "truememory.mps_utils": mps, "sentence_transformers": transformers, + "truememory.tier_switch.serving": serving, + "truememory.tier_switch.runtime": runtime}): yield vectors, reranker, factories def test_missing_endpoints_return_proxies_in_independent_clients(self) -> None: diff --git a/tests/test-mps-budget-743.py b/tests/test-mps-budget-743.py index 49200a8f..e658b306 100644 --- a/tests/test-mps-budget-743.py +++ b/tests/test-mps-budget-743.py @@ -2,6 +2,8 @@ from __future__ import annotations import ast +from contextlib import contextmanager +from collections.abc import Iterator import dataclasses import logging import re @@ -250,12 +252,17 @@ def transformer(name: str, **kwargs: object) -> object: get_embedding_proxy=lambda **kwargs: self.proxy, get_reranker_proxy=lambda **kwargs: self.proxy, ) + serving = source_module("tier_switch/serving.py") + runtime = types.SimpleNamespace(current_runtime_operation=lambda: None, + require_model_load_allowed=lambda: None, TierRuntimeError=RuntimeError) fake_tiers = types.SimpleNamespace(resolve_custom_tier=lambda: {"embed_dim": 192}) self.static = Mock(return_value=object()) for name, module in { "sentence_transformers": types.SimpleNamespace(SentenceTransformer=transformer, CrossEncoder=transformer), "model2vec": types.SimpleNamespace(StaticModel=types.SimpleNamespace(from_pretrained=self.static)), "truememory.model_client": fake_client, + "truememory.tier_switch.serving": serving, + "truememory.tier_switch.runtime": runtime, "truememory.tier_config": fake_tiers, "truememory.reranker": types.SimpleNamespace(get_current_reranker_name=lambda: "synthetic/reranker"), }.items(): @@ -265,7 +272,7 @@ def transformer(name: str, **kwargs: object) -> object: def local_embed(self, model_id: str) -> types.ModuleType: return source_module("vector_search.py", { - "_model": None, "_embedding_dim": -1, "EMBEDDING_MODEL": model_id, + "_model": None, "_frozen_embedding_target": None, "_embedding_dim": -1, "EMBEDDING_MODEL": model_id, "_MODEL_DIMS": {"model2vec": 256, "minilm": 384, "bge-small": 384, "qwen3_256": 256}, "_lock": threading.Lock(), "os": types.SimpleNamespace(environ=self.env), "logger": logging.getLogger("synthetic743"), @@ -295,10 +302,12 @@ def server(self): def local_reranker(self) -> types.ModuleType: return source_module("reranker.py", { - "_model": None, "_model_name": None, "_lock": threading.Lock(), + "_model": None, "_model_name": None, "_model_certified": False, + "contextmanager": contextmanager, "Iterator": Iterator, + "_lock": threading.Lock(), "get_current_reranker_name": lambda: "synthetic/reranker", "log": logging.getLogger("synthetic743"), - }, ("get_reranker",)) + }, ("get_reranker", "_reranker_load_lock")) def test_all_local_embedding_factories_calibrate_before_constructor(self) -> None: for model_id, dim in (("minilm", 384), ("bge-small", 384), ("qwen3_256", 256), diff --git a/tests/test-prepared-target-795.py b/tests/test-prepared-target-795.py index c4930d5c..c21060e9 100644 --- a/tests/test-prepared-target-795.py +++ b/tests/test-prepared-target-795.py @@ -168,6 +168,7 @@ def to(self, device: str) -> None: self.client.use_model_server = lambda: False self.vector = load("vector_search", definitions=True, namespace={ "np": Numpy(), "threading": threading, "contextmanager": contextmanager, + "database_operation": lambda function: function, "time": time, "math": math, "re": re, "sqlite3": sqlite3, "_lock": threading.Lock(), "_prepared_target_slot": threading.Lock(), "_model": None, "EMBEDDING_MODEL": "model2vec", "_embedding_dim": 256, @@ -598,6 +599,7 @@ def test_real_vec0_pair_owned_and_borrowed_transactions(self) -> None: load("tier_config") vector = load("vector_search", definitions=True, namespace={ "np": Numpy(), "threading": threading, "contextmanager": contextmanager, + "database_operation": lambda function: function, "time": time, "math": math, "re": re, "sqlite3": sqlite3, "_VEC_DISTANCE_METRIC": "cosine", }) diff --git a/tests/test-rebuild-foreground-publication-756.py b/tests/test-rebuild-foreground-publication-756.py index 5e8b0b7d..d7259e32 100644 --- a/tests/test-rebuild-foreground-publication-756.py +++ b/tests/test-rebuild-foreground-publication-756.py @@ -25,7 +25,7 @@ def engine(self, conn=None): original = next(node for node in tree.body if isinstance(node, ast.ClassDef) and node.name == "TrueMemoryEngine") methods = [node for node in original.body if isinstance(node, ast.FunctionDef) and node.name in ("add", "update")] cls = ast.ClassDef(name="SyntheticEngine", bases=[], keywords=[], body=methods, decorator_list=[]) - namespace = dict(__builtins__=self.vector.__dict__["__builtins__"], + namespace = dict(engine_operation=lambda function: function, __builtins__=self.vector.__dict__["__builtins__"], logger=logging.getLogger(__name__), MAX_CONTENT_LENGTH=10000, insert_message=self.modules["storage"].insert_message, update_message=self.modules["storage"].update_message) diff --git a/tests/test-rebuild-source-streaming-756.py b/tests/test-rebuild-source-streaming-756.py index 3ffb7f5d..bbb36aa9 100644 --- a/tests/test-rebuild-source-streaming-756.py +++ b/tests/test-rebuild-source-streaming-756.py @@ -14,7 +14,7 @@ import threading import types import unittest -from contextlib import closing, contextmanager +from contextlib import closing, contextmanager, nullcontext from dataclasses import replace from pathlib import Path @@ -22,7 +22,9 @@ def load_modules(): - modules = {} + modules = {"tier_switch.runtime": types.SimpleNamespace( + serving_operation=lambda *args, **kwargs: nullcontext(types.SimpleNamespace(selection=None)), + current_operation=lambda conn: types.SimpleNamespace(selection=None), TierRuntimeError=RuntimeError)} def safe_import(name, globals=None, locals=None, fromlist=(), level=0): if name == "torch": @@ -44,6 +46,7 @@ def safe_import(name, globals=None, locals=None, fromlist=(), level=0): vector = types.ModuleType("synthetic_stream_vector") vector.__dict__.update( __builtins__=dict(vars(builtins), __import__=safe_import), + database_operation=lambda function: function, contextmanager=contextmanager, closing=closing, Iterator=Iterator, replace=replace, datetime=datetime, sqlite3=sqlite3, logger=logging.getLogger(__name__), _lock=threading.Lock(), _model_generation=0, EMBEDDING_MODEL="synthetic-model", _embedding_dim=2, @@ -943,6 +946,7 @@ def test_actual_deprecated_open_rebuilds_existing_legacy_connection(self): exec(compile(ast.fix_missing_locations(ast.Module(body=[cls], type_ignores=[])), "actual-legacy-open", "exec"), namespace) engine = namespace["LegacyEngine"]() engine.db_path, engine.stats = path, {} + engine._write_lock = threading.Lock() engine._has_consolidation = engine._has_style_vec = False engine._purge_legacy_entity_profile_summaries = lambda: None fake_vec = types.ModuleType("sqlite_vec") diff --git a/tests/test-tier-runtime-bridge-795.py b/tests/test-tier-runtime-bridge-795.py new file mode 100644 index 00000000..c87c402a --- /dev/null +++ b/tests/test-tier-runtime-bridge-795.py @@ -0,0 +1,1159 @@ +"""Runtime selection with in-memory journals and inert model-slot doubles.""" + +from __future__ import annotations + +import ast +import builtins +import dataclasses +import gc +import os +import runpy +import sqlite3 +import sys +import threading +import types +import unittest +import weakref +from collections.abc import Callable, Iterator +from contextlib import contextmanager, ExitStack +from pathlib import Path +from unittest.mock import patch + +ROOT = Path(__file__).resolve().parents[1] +JOURNAL = runpy.run_path(str(Path(__file__).with_name("test-tier-activation-journal-795.py"))) +GATE = runpy.run_path(str(Path(__file__).with_name("test-tier-serving-gate-795.py"))) +PREPARED = runpy.run_path(str(Path(__file__).with_name("test-prepared-target-795.py"))) + + +class Cancelled(BaseException): + pass + + +class TestRuntimeBridge(unittest.TestCase): + def setUp(self) -> None: + self.fixture = JOURNAL["TestActivationJournal"]() + self.fixture.setUp() + self.addCleanup(self.fixture.doCleanups) + self.conn = self.fixture.conn + self.modules = self.fixture.modules + self.gate = GATE["load"]() + self.modules["tier_switch.serving"] = self.gate + self.loads = [] + self.vector = types.SimpleNamespace( + EMBEDDING_MODEL="model2vec", _embedding_dim=256, _model=object(), + _frozen_embedding_target=None, resolve_tier=lambda: "edge", + _active_vec_table=lambda conn: "vec_messages", _active_sep_table=lambda conn: "vec_messages_sep", + ) + self.reranker = types.SimpleNamespace( + _active_tier="edge", _model_name="synthetic/default", _model=object(), _frozen_reranker_id=None, + _model_certified=False, + get_current_reranker_name=lambda: "synthetic/default", + ) + + def embedding(target, **kwargs): + self.assertFalse(self.conn.in_transaction) + self.loads.append(("embedding", target)) + self.vector._frozen_embedding_target = target + self.vector.EMBEDDING_MODEL, self.vector._embedding_dim = target.model_id, target.dimension + self.vector._model = object() + + def reranking(tier, name, **kwargs): + self.assertFalse(self.conn.in_transaction) + self.loads.append(("reranker", name)) + self.reranker._active_tier, self.reranker._model_name = tier, name + self.reranker._frozen_reranker_id = name + self.reranker._model = object() + self.reranker._model_certified = True + + self.vector.apply_frozen_embedding_target = embedding + self.reranker.apply_frozen_reranker = reranking + self.reranker.get_reranker = lambda model_name, **kwargs: reranking("edge", model_name, **kwargs) + self.modules["vector_search"], self.modules["reranker"] = self.vector, self.reranker + + def safe_import(name, globals=None, locals=None, fromlist=(), level=0): + if name in {"numpy", "torch", "sentence_transformers", "model2vec", "sqlite_vec"}: + raise AssertionError("Native/model imports are forbidden in safe bridge tests") + if name == "truememory": + return types.SimpleNamespace(vector_search=self.vector, reranker=self.reranker) + if name == "truememory.tier_switch": + return types.SimpleNamespace(serving=self.gate) + if name.startswith("truememory."): + short = name.removeprefix("truememory.") + if short in self.modules: + return self.modules[short] + raise AssertionError("Unexpected application import: " + name) + return builtins.__import__(name, globals, locals, fromlist, level) + + self.safe_import = safe_import + self.api = types.ModuleType("synthetic_runtime_bridge") + sys.modules[self.api.__name__] = self.api + self.api.__dict__["__builtins__"] = dict(vars(builtins), __import__=safe_import) + path = ROOT / "truememory/tier_switch/runtime.py" + exec(compile(path.read_text(), str(path), "exec"), self.api.__dict__) + self.modules["tier_switch.runtime"] = self.api + self.api._database_identity = lambda conn: ("memory", conn) + + def select(self): + intent, plan = self.fixture.build() + return self.fixture.commit(intent, plan) + + def test_legacy_no_record_stays_lazy_and_bound_none_is_distinct(self) -> None: + self.assertIsNone(self.api.current_operation(self.conn)) + before = (self.conn.commits, self.conn.rollbacks) + with self.api.serving_operation(self.conn) as operation: + self.assertIsNone(operation.selection) + self.assertIs(self.api.current_operation(self.conn), operation) + self.assertEqual(operation.tables, ("vec_messages", "vec_messages_sep")) + self.assertEqual(operation.tier, "edge") + self.assertEqual(self.loads, []) + self.assertEqual((self.conn.commits, self.conn.rollbacks), before) + self.assertIsNone(self.api.current_operation(self.conn)) + + def test_selected_reconciles_exact_identity_before_body_and_ack(self) -> None: + selected = self.select() + with self.api.serving_operation(self.conn) as operation: + self.assertEqual(operation.selection, selected) + self.assertEqual(operation.tables, selected.tables) + self.assertEqual(operation.reranker_id, selected.reranker_id) + self.assertEqual(operation.tier, "custom") + self.assertFalse(self.conn.in_transaction) + self.assertEqual(self.api.runtime_acknowledgement(), self.api.selection_key(selected)) + self.assertEqual(self.loads, [("embedding", selected.target), ("reranker", selected.reranker_id)]) + with self.api.serving_operation(self.conn): + pass + self.assertEqual(len(self.loads), 2) + + def test_projection_rechecks_controls_after_readiness_modules(self) -> None: + selected = self.select() + cancelled = threading.Event() + original = self.api._modules + def modules(): + cancelled.set() + return original() + with patch.object(self.api, "_modules", side_effect=modules): + with self.assertRaises(self.gate.ServingLeaseCancelled): + self.api.apply_frozen_selection(selected, cancelled=cancelled) + self.assertEqual(self.loads, []) + self.assertIsNone(self.api.runtime_acknowledgement()) + + def test_legacy_explicit_override_forwards_admission_controls(self) -> None: + cancelled = threading.Event() + observed = [] + original = self.reranker.get_reranker + def loader(**kwargs): + observed.append(kwargs) + return original(**kwargs) + with patch.object(self.reranker, "get_reranker", side_effect=loader): + with self.api.serving_operation(self.conn, reranker_id="synthetic/deep", cancelled=cancelled): + pass + self.assertEqual(observed, [{"model_name": "synthetic/deep", "deadline": None, "cancelled": cancelled}]) + + def test_nested_search_has_zero_selection_sql_and_reuses_context(self) -> None: + self.select() + queries = [] + self.conn.set_trace_callback(queries.append) + with self.api.serving_operation(self.conn) as outer: + before = len(queries) + with self.api.serving_operation(self.conn) as inner: + self.assertIs(inner, outer) + self.assertIs(self.api.current_operation(self.conn), outer) + self.assertEqual(len(queries), before) + + def test_outer_override_is_frozen_without_changing_journal(self) -> None: + selected = self.select() + with self.api.serving_operation(self.conn, reranker_id="synthetic/deep") as operation: + self.assertEqual(operation.reranker_id, "synthetic/deep") + self.assertEqual(operation.selection.reranker_id, selected.reranker_id) + self.assertEqual(self.api.runtime_acknowledgement(), self.api.selection_key(selected, "synthetic/deep")) + with self.api.serving_operation(self.conn) as nested: + self.assertIs(nested, operation) + with self.assertRaises(self.api.TierRuntimeError): + with self.api.serving_operation(self.conn, reranker_id="synthetic/other"): + self.fail("Nested override cannot change the admitted model") + self.assertEqual(self.fixture.api.read_activation_state(self.conn).selection, selected) + + def test_same_thread_projection_upgrade_refuses(self) -> None: + selected = self.select() + with self.api.serving_operation(self.conn): + with self.assertRaises(self.api.TierRuntimeError): + self.api.apply_frozen_selection(selected) + + def test_projection_reenters_existing_exclusive_guard(self) -> None: + selected = self.select() + with self.gate.exclusive_activation(): + self.api.apply_frozen_selection(selected) + self.assertEqual(self.api.runtime_acknowledgement(), self.api.selection_key(selected)) + + def test_loader_failure_clears_ack_and_preserves_exception(self) -> None: + selected = self.select() + self.api.apply_frozen_selection(selected) + changed = dataclasses.replace(selected, generation="f" * 32) + for module, name in ((self.vector, "apply_frozen_embedding_target"), (self.reranker, "apply_frozen_reranker")): + sentinel = Cancelled("synthetic") + self.api._acknowledged = self.api.selection_key(selected) + with patch.object(module, name, side_effect=sentinel): + with self.assertRaises(Cancelled) as raised: + self.api.apply_frozen_selection(changed) + self.assertIs(raised.exception, sentinel) + self.assertIsNone(self.api.runtime_acknowledgement()) + + def test_body_baseexception_clears_context_and_gate(self) -> None: + sentinel = Cancelled("synthetic") + with self.assertRaises(Cancelled) as raised: + with self.api.serving_operation(self.conn): + raise sentinel + self.assertIs(raised.exception, sentinel) + self.assertIsNone(self.api.current_operation(self.conn)) + with self.gate.exclusive_activation(timeout=1): + pass + + def test_bound_different_connection_refuses_even_legacy_none(self) -> None: + other = sqlite3.connect(":memory:") + self.addCleanup(other.close) + with self.api.serving_operation(self.conn): + with self.assertRaises(self.api.TierRuntimeError): + self.api.current_operation(other) + with self.assertRaises(self.api.TierRuntimeError): + with self.api.serving_operation(other): + self.fail("Different connection cannot inherit the operation") + + def test_borrowed_legacy_is_nonblocking_and_never_ends_sql(self) -> None: + self.conn.execute("BEGIN") + before = (self.conn.commits, self.conn.rollbacks) + with self.api.serving_operation(self.conn) as operation: + self.assertTrue(operation.borrowed) + with self.assertRaises(self.api.TierRuntimeError): + self.api.require_model_load_allowed() + self.assertTrue(self.conn.in_transaction) + self.assertEqual((self.conn.commits, self.conn.rollbacks), before) + self.assertEqual(self.loads, []) + self.conn.rollback() + + def test_borrowed_selected_requires_existing_ready_runtime(self) -> None: + selected = self.select() + self.conn.execute("BEGIN") + before = (self.conn.commits, self.conn.rollbacks) + with self.assertRaises(self.api.TierRuntimeError): + with self.api.serving_operation(self.conn): + self.fail("Reconciliation cannot load in a borrowed transaction") + self.assertEqual(self.loads, []) + self.assertEqual((self.conn.commits, self.conn.rollbacks), before) + self.conn.rollback() + self.api.apply_frozen_selection(selected) + self.conn.execute("BEGIN") + before = (self.conn.commits, self.conn.rollbacks) + with self.api.serving_operation(self.conn): + pass + self.assertTrue(self.conn.in_transaction) + self.assertEqual((self.conn.commits, self.conn.rollbacks), before) + self.conn.rollback() + + def test_selected_admission_never_lazy_loads_inside_body(self) -> None: + self.select() + with self.api.serving_operation(self.conn): + with self.assertRaises(self.api.TierRuntimeError): + self.api.require_model_load_allowed() + + def test_child_reservation_outlives_parent_and_clears_context(self) -> None: + with self.api.serving_operation(self.conn) as parent: + child = parent.fork_child() + with child.join() as operation: + self.assertIs(self.api.current_operation(self.conn), operation) + self.assertEqual(operation.tables, parent.tables) + self.assertIsNone(self.api.current_operation(self.conn)) + with self.assertRaises(self.gate.ServingLeaseError): + with child.join(): + self.fail("Consumed child cannot join twice") + + def test_child_different_database_refuses_and_releases_admitted_pin(self) -> None: + other = sqlite3.connect(":memory:") + self.addCleanup(other.close) + with self.api.serving_operation(self.conn) as parent: + child = parent.fork_child() + with self.assertRaises(self.api.TierRuntimeError): + with child.join(other): + self.fail("Different database cannot join") + with self.gate.exclusive_activation(timeout=1): + pass + + def test_child_explicit_matching_database_binds_own_clean_connection(self) -> None: + other = sqlite3.connect(":memory:") + self.addCleanup(other.close) + self.api._database_identity = lambda conn: ("file", "synthetic-same-db") + with self.api.serving_operation(self.conn) as parent: + child = parent.fork_child() + with child.join(other) as operation: + self.assertIs(self.api.current_operation(other), operation) + with self.assertRaises(self.api.TierRuntimeError): + self.api.current_operation(self.conn) + + def test_forked_process_ack_is_not_ready(self) -> None: + selected = self.select() + self.api.apply_frozen_selection(selected) + with patch.object(self.api.os, "getpid", return_value=self.api._pid + 1): + self.assertIsNone(self.api.runtime_acknowledgement()) + + def test_nonblocking_admission_attempts_mutex_once(self) -> None: + gate = self.gate.ServingGate() + lock = types.SimpleNamespace(acquire=lambda **kwargs: False) + with patch.object(gate, "_condition", lock): + with self.assertRaises(self.gate.ServingLeaseError): + with gate.operation_lease(("synthetic",), blocking=False): + self.fail("Unavailable mutex cannot admit") + with self.assertRaises(self.gate.ServingLeaseError): + with gate.operation_lease(("synthetic",), blocking=0): + self.fail("Only exact bool is accepted") + + def test_nonblocking_conflict_never_condition_waits(self) -> None: + gate = self.gate.ServingGate() + ticket = object() + gate._waiting_writers.append(ticket) + with patch.object(gate._condition, "wait", side_effect=AssertionError("Admission cannot wait")): + with self.assertRaises(self.gate.ServingLeaseError): + with gate.operation_lease(("synthetic",), blocking=False): + self.fail("Queued writer must block unrelated reader") + gate._waiting_writers.clear() + self.assertFalse(gate._readers) + + def test_nonblocking_nested_read_inherits_cancel_and_identity(self) -> None: + gate = self.gate.ServingGate() + cancelled = threading.Event() + with gate.operation_lease(("synthetic",), cancelled=cancelled): + with gate.operation_lease(("synthetic",), blocking=False): + pass + with self.assertRaises(self.gate.ServingLeaseError): + with gate.operation_lease(("different",), blocking=False): + self.fail("Nested identity cannot change") + cancelled.set() + with self.assertRaises(self.gate.ServingLeaseCancelled): + with gate.operation_lease(("synthetic",), blocking=False): + self.fail("Inherited cancellation precedes admission") + + def test_engine_and_hybrid_outer_boundaries_are_present(self) -> None: + engine = ast.parse((ROOT / "truememory/engine.py").read_text()) + klass = next(node for node in engine.body if isinstance(node, ast.ClassDef) and node.name == "TrueMemoryEngine") + for name in ("search", "search_agentic", "search_vectors_raw", "add", "update", "delete", "delete_all"): + node = next(node for node in klass.body if isinstance(node, ast.FunctionDef) and node.name == name) + self.assertIn("engine_operation", [ast.unparse(value) for value in node.decorator_list]) + hybrid = ast.parse((ROOT / "truememory/hybrid.py").read_text()) + function = next(node for node in hybrid.body if isinstance(node, ast.FunctionDef) and node.name == "search_hybrid") + self.assertIn("database_operation", [ast.unparse(value) for value in function.decorator_list]) + + def test_cancellation_during_final_journal_read_refuses_body(self) -> None: + self.select() + event = threading.Event() + original = self.api._read_selection + calls = [] + def read(conn): + value = original(conn) + calls.append(value) + if len(calls) == 2: + event.set() + return value + with patch.object(self.api, "_read_selection", side_effect=read): + with self.assertRaises(self.gate.ServingLeaseCancelled): + with self.api.serving_operation(self.conn, cancelled=event): + self.fail("Cancellation after SQL must prevent body admission") + self.assertIsNone(self.api.current_operation(self.conn)) + with self.gate.exclusive_activation(timeout=1): + pass + + def test_deadline_consumed_by_final_journal_read_refuses_body(self) -> None: + self.select() + clock = [1000.0] + original = self.api._read_selection + calls = [] + def read(conn): + value = original(conn) + calls.append(value) + if len(calls) == 2: + clock[0] = 1002.0 + return value + with patch.object(self.api.time, "monotonic", side_effect=lambda: clock[0]), patch.object(self.api, "_read_selection", side_effect=read): + with self.assertRaises(self.gate.ServingLeaseTimeout): + with self.api.serving_operation(self.conn, deadline=1001.0): + self.fail("Expired deadline after SQL must prevent body admission") + + def test_child_rechecks_inherited_event_after_new_connection_read(self) -> None: + other = sqlite3.connect(":memory:") + self.addCleanup(other.close) + self.api._database_identity = lambda conn: ("file", "synthetic-same-db") + event = threading.Event() + with self.api.serving_operation(self.conn, cancelled=event) as operation: + child = operation.fork_child() + original = self.api._read_selection + def read(conn): + value = original(conn) + event.set() + return value + with patch.object(self.api, "_read_selection", side_effect=read): + with self.assertRaises(self.gate.ServingLeaseCancelled): + with child.join(other): + self.fail("Child must enforce inherited cancellation after SQL") + self.assertIsNone(self.api.current_operation(other)) + with self.gate.exclusive_activation(timeout=1): + pass + + def test_nested_reservation_inherits_stricter_nested_control(self) -> None: + event = threading.Event() + with self.api.serving_operation(self.conn): + with self.api.serving_operation(self.conn, cancelled=event) as operation: + child = operation.fork_child() + event.set() + with self.assertRaises(self.gate.ServingLeaseCancelled): + with child.join(): + self.fail("Nested control must reach the forked child") + child.close() + + def test_actual_selected_engine_init_reconciles_before_validation_and_keeps_schedule(self) -> None: + selected = self.select() + tree = ast.parse((ROOT / "truememory/engine.py").read_text()) + source = next(node for node in tree.body if isinstance(node, ast.ClassDef) and node.name == "TrueMemoryEngine") + source.body = [node for node in source.body if isinstance(node, ast.FunctionDef) and node.name in { + "_ensure_connection", "_initialize_connection", "_open_connection_handle", "close", + }] + calls = [] + def validate(conn): + self.assertEqual(self.api.current_operation(conn).selection, selected) + self.assertEqual(self.api.runtime_acknowledgement(), self.api.selection_key(selected)) + calls.append("validate") + def imports(name, globals=None, locals=None, fromlist=(), level=0): + if name == "sqlite_vec": + return types.SimpleNamespace(load=lambda conn: calls.append("extension")) + return self.safe_import(name, globals, locals, fromlist, level) + namespace = { + "__builtins__": dict(vars(builtins), __import__=imports), + "_HAS_VECTOR": True, "_HAS_HYBRID": True, "init_vec_table": validate, "sqlite3": sqlite3, + } + exec(compile(ast.fix_missing_locations(ast.Module(body=[source], type_ignores=[])), "synthetic-engine-init", "exec"), namespace) + engine = namespace["TrueMemoryEngine"]() + engine.conn = self.conn + self.conn.enable_load_extension = lambda enabled: None + engine._write_lock = threading.Lock() + engine._init_lock = threading.Lock() + engine._runtime_initialized = False + engine._open_connection_handle = lambda: calls.append("plain-open") + engine._purge_legacy_entity_profile_summaries = lambda: calls.append("purge") + engine._maybe_startup_consolidate = lambda: calls.append("startup-schedule") + engine._maybe_auto_consolidate = lambda: calls.append("cached-schedule") + engine._ensure_connection() + self.assertEqual(calls, ["plain-open", "extension", "validate", "purge", "startup-schedule"]) + self.assertTrue(engine._has_vectors) + self.assertTrue(engine._has_hybrid) + engine._ensure_connection() + self.assertEqual(calls[-2:], ["plain-open", "cached-schedule"]) + replacement = sqlite3.connect(":memory:", factory=JOURNAL["JournalConnection"]) + self.addCleanup(replacement.close) + self.conn.backup(replacement) + replacement.enable_load_extension = lambda enabled: None + namespace["create_db"] = lambda path: replacement + engine.db_path = Path(":memory:") + engine._open_connection_handle = types.MethodType(namespace["TrueMemoryEngine"]._open_connection_handle, engine) + self.conn.close() + calls.clear() + engine._ensure_connection() + self.assertEqual(calls, ["extension", "validate", "purge", "startup-schedule"]) + self.assertIs(engine._runtime_vector_connection, replacement) + self.assertTrue(engine._runtime_initialized) + self.assertTrue(engine._has_vectors) + self.assertTrue(engine.ready) + engine.close() + self.assertIsNone(engine._runtime_vector_connection) + self.assertIsNone(engine._runtime_vector_generation) + self.assertFalse(engine._runtime_initialized) + + + def test_ready_legacy_borrowed_connection_does_not_reinitialize(self) -> None: + tree = ast.parse((ROOT / "truememory/engine.py").read_text()) + source = next(node for node in tree.body if isinstance(node, ast.ClassDef) and node.name == "TrueMemoryEngine") + source.body = [node for node in source.body if isinstance(node, ast.FunctionDef) and node.name == "_initialize_connection"] + namespace = {"__builtins__": dict(vars(builtins), __import__=self.safe_import)} + exec(compile(ast.fix_missing_locations(ast.Module(body=[source], type_ignores=[])), "synthetic-legacy-ready", "exec"), namespace) + engine = namespace["TrueMemoryEngine"]() + engine.conn, engine.ready, engine._runtime_initialized = self.conn, True, False + engine._init_lock = threading.Lock() + engine._write_lock = threading.Lock() + self.conn.execute("INSERT INTO messages(id,content) VALUES(901,'synthetic borrowed legacy')") + before = self.conn.commits, self.conn.rollbacks + with self.api.serving_operation(self.conn): + engine._initialize_connection(_suppress_maintenance=True) + self.assertTrue(self.conn.in_transaction) + self.assertEqual((self.conn.commits, self.conn.rollbacks), before) + engine.ready = False + with self.assertRaises(self.api.TierRuntimeError): + engine._initialize_connection(_suppress_maintenance=True) + + def test_fresh_or_closed_handle_clears_legacy_ready_before_initialization(self) -> None: + tree = ast.parse((ROOT / "truememory/engine.py").read_text()) + source = next(node for node in tree.body if isinstance(node, ast.ClassDef) and node.name == "TrueMemoryEngine") + source.body = [node for node in source.body if isinstance(node, ast.FunctionDef) and node.name == "_open_connection_handle"] + replacement = sqlite3.connect(":memory:") + self.addCleanup(replacement.close) + namespace = {"sqlite3": sqlite3, "create_db": lambda path: replacement, + "__builtins__": dict(vars(builtins), __import__=self.safe_import)} + exec(compile(ast.fix_missing_locations(ast.Module(body=[source], type_ignores=[])), "synthetic-legacy-reopen", "exec"), namespace) + engine = namespace["TrueMemoryEngine"]() + engine.db_path, engine._init_lock = Path(":memory:"), threading.Lock() + engine._write_lock = threading.Lock() + for prior in (None, self.conn): + if prior is not None: + prior.close() + engine.conn, engine.ready, engine._runtime_initialized = prior, True, True + engine._open_connection_handle() + self.assertIs(engine.conn, replacement) + self.assertFalse(engine.ready) + self.assertFalse(engine._runtime_initialized) + + def test_managed_writer_transaction_waits_and_is_not_caller_borrowed(self) -> None: + conn = sqlite3.connect(":memory:", check_same_thread=False) + self.addCleanup(conn.close) + conn.execute("CREATE TABLE synthetic(value TEXT)") + class Engine: + pass + engine = Engine() + engine.conn = conn + waiting, release, owned, entered = (threading.Event() for _ in range(4)) + failures = [] + api = self.api + class ObservedLock(api.ConnectionWriteLock): + def acquire(self, blocking=True, timeout=-1): + if blocking and self.owns_other_transaction(conn): + waiting.set() + return super().acquire(blocking, timeout) + lock = ObservedLock(engine) + def writer(): + try: + with lock: + conn.execute("INSERT INTO synthetic VALUES('owned')") + owned.set() + if not release.wait(2): + raise AssertionError("Synthetic writer release timed out") + conn.commit() + except BaseException as exc: + failures.append(exc) + def reader(): + try: + with api._connection_read(conn, lock, api.time.monotonic() + 2, None): + self.assertFalse(conn.in_transaction) + self.assertEqual(conn.execute("SELECT value FROM synthetic").fetchall(), [("owned",)]) + entered.set() + except BaseException as exc: + failures.append(exc) + owner = threading.Thread(target=writer) + contender = threading.Thread(target=reader) + owner.start() + try: + self.assertTrue(owned.wait(2)) + contender.start() + self.assertTrue(waiting.wait(2)) + self.assertFalse(entered.is_set()) + finally: + release.set() + owner.join(2) + if contender.ident is not None: + contender.join(2) + self.assertFalse(owner.is_alive()) + self.assertFalse(contender.is_alive()) + self.assertEqual(failures, []) + self.assertTrue(entered.is_set()) + + def test_managed_lock_does_not_adopt_preexisting_caller_transaction(self) -> None: + conn = sqlite3.connect(":memory:", check_same_thread=False) + self.addCleanup(conn.close) + conn.execute("CREATE TABLE synthetic(value TEXT)") + conn.execute("INSERT INTO synthetic VALUES('caller')") + class Engine: + pass + engine = Engine() + engine.conn = conn + lock = self.api.ConnectionWriteLock(engine) + owned, release = threading.Event(), threading.Event() + failures = [] + def holder(): + try: + with lock: + owned.set() + if not release.wait(2): + raise AssertionError("Synthetic holder release timed out") + except BaseException as exc: + failures.append(exc) + thread = threading.Thread(target=holder) + thread.start() + try: + self.assertTrue(owned.wait(2)) + with self.assertRaisesRegex(self.api.TierRuntimeError, "Borrowed connection admission is busy"): + with self.api._connection_read(conn, lock, None, None): + self.fail("Caller SQL must not wait for an unrelated lock holder") + self.assertTrue(conn.in_transaction) + self.assertEqual(conn.execute("SELECT value FROM synthetic").fetchall(), [("caller",)]) + finally: + release.set() + thread.join(2) + self.assertFalse(thread.is_alive()) + self.assertEqual(failures, []) + conn.rollback() + self.assertEqual(conn.execute("SELECT value FROM synthetic").fetchall(), []) + + def test_connection_probe_error_preserves_handle_under_both_locks(self) -> None: + tree = ast.parse((ROOT / "truememory/engine.py").read_text()) + source = next(node for node in tree.body if isinstance(node, ast.ClassDef) and node.name == "TrueMemoryEngine") + source.body = [node for node in source.body if isinstance(node, ast.FunctionDef) + and node.name == "_open_connection_handle"] + namespace = {"sqlite3": sqlite3, + "__builtins__": dict(vars(builtins), __import__=self.safe_import)} + exec(compile(ast.fix_missing_locations(ast.Module(body=[source], type_ignores=[])), + "synthetic-probe-failure", "exec"), namespace) + engine = namespace["TrueMemoryEngine"]() + engine._init_lock, engine._write_lock = threading.Lock(), threading.Lock() + owner = self + class Connection: + in_transaction = False + def execute(self, sql): + owner.assertTrue(engine._init_lock.locked()) + owner.assertTrue(engine._write_lock.locked()) + raise sqlite3.OperationalError("synthetic transient probe failure") + def close(self): + owner.fail("A transient SQL error must not close a shared connection") + engine.conn = Connection() + original = engine.conn + with self.assertRaisesRegex(sqlite3.OperationalError, "synthetic transient"): + engine._open_connection_handle() + self.assertIs(engine.conn, original) + self.assertFalse(engine._init_lock.locked()) + self.assertFalse(engine._write_lock.locked()) + + + def lifecycle_engine( + self, method: str, conn: sqlite3.Connection | None = None, + create_db: Callable[[Path], sqlite3.Connection] | None = None, + ) -> object: + tree = ast.parse((ROOT / "truememory/engine.py").read_text(encoding="utf-8")) + source = next(node for node in tree.body if isinstance(node, ast.ClassDef) and node.name == "TrueMemoryEngine") + source.body = [node for node in source.body if isinstance(node, ast.FunctionDef) and node.name == method] + namespace = {"sqlite3": sqlite3, "create_db": create_db, + "__builtins__": dict(vars(builtins), __import__=self.safe_import)} + exec(compile(ast.fix_missing_locations(ast.Module(body=[source], type_ignores=[])), + "synthetic-engine-lifecycle", "exec"), namespace) + engine = namespace["TrueMemoryEngine"]() + engine.conn = self.conn if conn is None else conn + engine.db_path = Path(":memory:") + engine._init_lock = threading.Lock() + engine._write_lock = self.api.ConnectionWriteLock(engine) + engine.ready = engine._has_vectors = engine._runtime_initialized = True + engine._runtime_vector_generation = "synthetic-initialized" + engine._runtime_vector_connection = engine.conn + return engine + + @contextmanager + def held_lifecycle_lock(self, lock: object) -> Iterator[None]: + entered, release = threading.Event(), threading.Event() + failures = [] + + def holder() -> None: + try: + with lock: + entered.set() + if not release.wait(2): + raise AssertionError("Synthetic lifecycle holder release timed out") + except BaseException as exc: + failures.append(exc) + + thread = threading.Thread(target=holder) + thread.start() + try: + self.assertTrue(entered.wait(2)) + yield + finally: + release.set() + thread.join(2) + self.assertFalse(thread.is_alive()) + self.assertEqual(failures, []) + + def test_actual_open_borrowed_busy_preserves_caller_sql_at_each_lock(self) -> None: + engine = self.lifecycle_engine("_open_connection_handle") + self.conn.execute("INSERT INTO messages(id,content) VALUES(901,'synthetic caller SQL')") + before = self.conn.commits, self.conn.rollbacks + queries = [] + self.conn.set_trace_callback(queries.append) + for lock in (engine._init_lock, engine._write_lock): + with self.subTest(lock=type(lock).__name__): + queries.clear() + with self.held_lifecycle_lock(lock): + with self.assertRaisesRegex(self.api.TierRuntimeError, "Borrowed connection admission is busy"): + engine._open_connection_handle() + self.assertEqual(queries, []) + self.assertIs(engine.conn, self.conn) + self.assertTrue(self.conn.in_transaction) + self.assertEqual((self.conn.commits, self.conn.rollbacks), before) + self.assertEqual(self.conn.execute("SELECT content FROM messages WHERE id=901").fetchone()[0], + "synthetic caller SQL") + self.assertFalse(engine._init_lock.locked()) + self.assertFalse(engine._write_lock.locked()) + self.conn.rollback() + self.assertEqual(self.conn.execute("SELECT count(*) FROM messages WHERE id=901").fetchone()[0], 0) + + def test_actual_open_clean_ioerr_reconnects_and_invalidates_runtime(self) -> None: + coded = sqlite3.OperationalError("synthetic extended I/O error") + coded.sqlite_errorcode = getattr(sqlite3, "SQLITE_IOERR", 10) | (7 << 8) + for error in (sqlite3.OperationalError("disk I/O error"), coded): + with self.subTest(error=str(error)): + original = sqlite3.connect(":memory:", factory=JOURNAL["JournalConnection"]) + replacement = sqlite3.connect(":memory:") + self.addCleanup(original.close) + self.addCleanup(replacement.close) + original.armed, original.fault_contains, original.fault = True, "PRAGMA schema_version", error + opened = [] + + def create(path: Path) -> sqlite3.Connection: + self.assertEqual(path, Path(":memory:")) + opened.append(path) + return replacement + + engine = self.lifecycle_engine("_open_connection_handle", original, create) + with patch.object(original, "close", wraps=original.close) as close: + engine._open_connection_handle() + close.assert_called_once_with() + self.assertEqual(opened, [Path(":memory:")]) + self.assertIs(engine.conn, replacement) + self.assertFalse(engine.ready) + self.assertFalse(engine._has_vectors) + self.assertFalse(engine._runtime_initialized) + self.assertIsNone(engine._runtime_vector_generation) + self.assertIsNone(engine._runtime_vector_connection) + self.assertFalse(engine._init_lock.locked()) + self.assertFalse(engine._write_lock.locked()) + + def test_actual_open_borrowed_ioerr_preserves_handle_and_caller_sql(self) -> None: + coded = sqlite3.OperationalError("synthetic extended I/O error") + coded.sqlite_errorcode = getattr(sqlite3, "SQLITE_IOERR", 10) | (7 << 8) + for error in (sqlite3.OperationalError("disk I/O error"), coded): + with self.subTest(error=str(error)): + engine = self.lifecycle_engine("_open_connection_handle") + self.conn.execute("INSERT INTO messages(id,content) VALUES(901,'synthetic caller SQL')") + before = self.conn.commits, self.conn.rollbacks + self.conn.armed, self.conn.fault_contains, self.conn.fault = True, "PRAGMA schema_version", error + with patch.object(self.conn, "close", wraps=self.conn.close) as close: + with self.assertRaises(sqlite3.OperationalError) as raised: + engine._open_connection_handle() + close.assert_not_called() + self.assertIs(raised.exception, error) + self.assertIs(engine.conn, self.conn) + self.assertTrue(self.conn.in_transaction) + self.assertEqual((self.conn.commits, self.conn.rollbacks), before) + self.assertTrue(engine.ready) + self.assertTrue(engine._has_vectors) + self.assertTrue(engine._runtime_initialized) + self.assertEqual(engine._runtime_vector_generation, "synthetic-initialized") + self.assertIs(engine._runtime_vector_connection, self.conn) + self.assertFalse(engine._init_lock.locked()) + self.assertFalse(engine._write_lock.locked()) + self.assertEqual(self.conn.execute("SELECT content FROM messages WHERE id=901").fetchone()[0], + "synthetic caller SQL") + self.conn.armed = False + self.conn.rollback() + self.assertEqual(self.conn.execute("SELECT count(*) FROM messages WHERE id=901").fetchone()[0], 0) + + def test_actual_open_extended_ioerr_without_sqlite_constant(self) -> None: + class SQLiteWithoutIOErr: + missing_constant_reads = 0 + + def __getattr__(self, name: str) -> object: + if name == "SQLITE_IOERR": + self.missing_constant_reads += 1 + raise AttributeError(name) + return getattr(sqlite3, name) + + for borrowed in (False, True): + with self.subTest(borrowed=borrowed): + sqlite_api = SQLiteWithoutIOErr() + self.assertIs(sqlite_api.connect, sqlite3.connect) + self.assertIs(sqlite_api.OperationalError, sqlite3.OperationalError) + original = sqlite3.connect(":memory:", factory=JOURNAL["JournalConnection"]) + replacement = sqlite3.connect(":memory:") + self.addCleanup(original.close) + self.addCleanup(replacement.close) + original.execute("CREATE TABLE synthetic(value TEXT)") + if borrowed: + original.execute("INSERT INTO synthetic VALUES('synthetic caller SQL')") + before = original.commits, original.rollbacks + error = sqlite3.OperationalError("synthetic extended I/O error") + error.sqlite_errorcode = 10 | (7 << 8) + original.armed, original.fault_contains, original.fault = True, "PRAGMA schema_version", error + opened = [] + + def create(path: Path) -> sqlite3.Connection: + self.assertEqual(path, Path(":memory:")) + opened.append(path) + return replacement + + engine = self.lifecycle_engine("_open_connection_handle", original, create) + engine._open_connection_handle.__globals__["sqlite3"] = sqlite_api + with patch.object(original, "close", wraps=original.close) as close: + if borrowed: + with self.assertRaises(sqlite3.OperationalError) as raised: + engine._open_connection_handle() + self.assertIs(raised.exception, error) + close.assert_not_called() + self.assertEqual(opened, []) + self.assertIs(engine.conn, original) + self.assertTrue(original.in_transaction) + self.assertTrue(engine.ready) + self.assertTrue(engine._has_vectors) + self.assertTrue(engine._runtime_initialized) + self.assertEqual(engine._runtime_vector_generation, "synthetic-initialized") + self.assertIs(engine._runtime_vector_connection, original) + self.assertEqual(original.execute("SELECT value FROM synthetic").fetchall(), + [("synthetic caller SQL",)]) + else: + engine._open_connection_handle() + close.assert_called_once_with() + self.assertEqual(opened, [Path(":memory:")]) + self.assertIs(engine.conn, replacement) + self.assertFalse(engine.ready) + self.assertFalse(engine._has_vectors) + self.assertFalse(engine._runtime_initialized) + self.assertIsNone(engine._runtime_vector_generation) + self.assertIsNone(engine._runtime_vector_connection) + self.assertEqual(sqlite_api.missing_constant_reads, 1) + self.assertEqual((original.commits, original.rollbacks), before) + self.assertFalse(engine._init_lock.locked()) + self.assertFalse(engine._write_lock.locked()) + if borrowed: + original.armed = False + original.rollback() + self.assertEqual(original.execute("SELECT count(*) FROM synthetic").fetchone()[0], 0) + + def test_actual_selected_ready_borrowed_init_ignores_competing_writer(self) -> None: + selected = self.select() + self.api.apply_frozen_selection(selected) + engine = self.lifecycle_engine("_initialize_connection") + engine._runtime_vector_generation = selected.generation + self.conn.execute("INSERT INTO messages(id,content) VALUES(901,'synthetic caller SQL')") + before = self.conn.commits, self.conn.rollbacks + loads = list(self.loads) + queries = [] + self.conn.set_trace_callback(queries.append) + with self.api.serving_operation(self.conn, connection_lock=engine._write_lock) as operation: + self.assertTrue(operation.borrowed) + with self.held_lifecycle_lock(engine._write_lock): + queries.clear() + with patch.object(self.api, "_connection_read", side_effect=AssertionError("Ready init must not acquire SQL locks")), \ + patch.object(self.api, "apply_frozen_selection", side_effect=AssertionError("Ready init must not reload models")): + engine._initialize_connection(_suppress_maintenance=True) + self.assertTrue(engine._write_lock.locked()) + self.assertEqual(queries, []) + self.assertEqual(self.loads, loads) + self.assertTrue(self.conn.in_transaction) + self.assertEqual((self.conn.commits, self.conn.rollbacks), before) + self.assertIs(self.api.current_operation(self.conn), operation) + self.assertTrue(engine.ready) + self.assertEqual(self.conn.execute("SELECT content FROM messages WHERE id=901").fetchone()[0], + "synthetic caller SQL") + self.conn.rollback() + self.assertEqual(self.conn.execute("SELECT count(*) FROM messages WHERE id=901").fetchone()[0], 0) + + +class TestRuntimeModels(unittest.TestCase): + def setUp(self) -> None: + self.fixture = PREPARED["TestPreparedTargets"]() + self.fixture.setUp() + self.addCleanup(self.fixture.doCleanups) + self.vector = self.fixture.vector + self.vector._frozen_embedding_target = None + self.gate = GATE["load"]() + self.operation = None + self.runtime = types.SimpleNamespace( + require_model_load_allowed=lambda: None, current_runtime_operation=lambda: self.operation, + TierRuntimeError=RuntimeError, + ) + package = types.ModuleType("truememory.tier_switch") + package.__path__ = [] + modules = patch.dict(sys.modules, { + "truememory.tier_switch": package, + "truememory.tier_switch.serving": self.gate, + "truememory.tier_switch.runtime": self.runtime, + }) + modules.start() + self.addCleanup(modules.stop) + + def test_actual_local_factory_releases_old_slot_before_constructor(self) -> None: + class Model: + pass + self.vector._model = Model() + old = weakref.ref(self.vector._model) + def check(): + gc.collect() + self.assertIsNone(old()) + self.assertIsNone(self.vector._model) + self.fixture.build_hook = check + self.vector.apply_frozen_embedding_target(self.fixture.target) + self.assertEqual(self.vector.EMBEDDING_MODEL, "qwen3_256") + self.assertEqual(self.vector._embedding_dim, 256) + self.assertEqual(len(self.fixture.calls), 1) + self.assertEqual(self.fixture.encoded[0][1], [self.fixture.target_module.TARGET_PROBE_TEXT]) + self.assertIsNotNone(self.vector._model) + + def test_failed_constructor_keeps_exact_target_and_empty_slot(self) -> None: + sentinel = Cancelled("synthetic") + def fail(): + raise sentinel + self.fixture.build_hook = fail + with self.assertRaises(Cancelled) as raised: + self.vector.apply_frozen_embedding_target(self.fixture.target) + self.assertIs(raised.exception, sentinel) + self.assertIsNone(self.vector._model) + self.assertEqual(self.vector._frozen_embedding_target, self.fixture.target) + self.assertFalse(self.vector._prepared_target_slot.locked()) + + def test_wrong_width_and_busy_prepared_slot_never_publish_active_model(self) -> None: + self.fixture.width = 257 + with self.assertRaises(self.fixture.target_module.EmbeddingTargetError): + self.vector.apply_frozen_embedding_target(self.fixture.target) + self.assertIsNone(self.vector._model) + self.fixture.width = None + self.vector._prepared_target_slot.acquire() + try: + before = len(self.fixture.calls) + with self.assertRaises(RuntimeError): + self.vector.apply_frozen_embedding_target(self.fixture.target) + self.assertEqual(len(self.fixture.calls), before) + finally: + self.vector._prepared_target_slot.release() + + def test_base_pro_identity_reuse_has_zero_second_native_work(self) -> None: + self.vector.apply_frozen_embedding_target(self.fixture.target) + original = self.vector._model + before = (len(self.fixture.calls), len(self.fixture.encoded), self.vector._model_generation) + pro = self.fixture.Target.capture("pro") + self.vector.apply_frozen_embedding_target(pro) + self.assertIs(self.vector._model, original) + self.assertEqual(self.vector._frozen_embedding_target, pro) + self.assertEqual((len(self.fixture.calls), len(self.fixture.encoded), self.vector._model_generation), before) + + def test_deadline_consumed_by_state_lock_prevents_constructor(self) -> None: + clock = [1000.0] + class Lock: + def acquire(self, **kwargs): + clock[0] = 1002.0 + return True + def release(self): + pass + with patch.object(self.vector.time, "monotonic", side_effect=lambda: clock[0]), patch.object(self.vector, "_lock", Lock()): + with self.assertRaises(TimeoutError): + self.vector.apply_frozen_embedding_target(self.fixture.target, timeout=1) + self.assertEqual(self.fixture.calls, []) + + def test_setter_and_unload_cannot_mutate_admitted_model(self) -> None: + self.vector.apply_frozen_embedding_target(self.fixture.target) + model = self.vector._model + with self.gate.operation_lease(("synthetic",)): + for action in (lambda: self.vector.set_embedding_model("edge"), self.vector.unload_model): + with self.assertRaises(self.gate.ServingLeaseError): + action() + self.assertIs(self.vector._model, model) + + def test_shared_builtin_retains_normal_fast_protocol_custom_uses_target_protocol(self) -> None: + self.fixture.client.use_model_server = lambda: True + self.fixture.client.np.integer = int + self.fixture.client.PROTOCOL_VERSION = 1 + requests = [] + def request(value, timeout=None): + requests.append(value) + return self.fixture.server.handle_request(value) + self.fixture.client._request_with_autostart = request + self.vector.apply_frozen_embedding_target(self.fixture.target) + self.vector._model.encode(["synthetic"], batch_size=2, show_progress_bar=False) + self.assertEqual([entry["op"] for entry in requests], ["prepare_embed_target_v1", "embed_batched"]) + self.assertEqual(requests[-1]["tier"], "qwen3_256") + custom = self.fixture.Target.capture("custom") + requests.clear() + self.vector.apply_frozen_embedding_target(custom) + self.vector._model.encode(["synthetic"], batch_size=2, show_progress_bar=False) + self.assertEqual([entry["op"] for entry in requests], ["prepare_embed_target_v1", "embed_target_v1"]) + self.assertEqual(requests[-1]["target"], custom.to_wire()) + + def load_reranker(self): + return PREPARED["load"]("reranker", definitions=True, namespace={ + "threading": threading, "contextmanager": contextmanager, + "_active_tier": "edge", "_frozen_reranker_id": None, + "_model_certified": False, + "_model": None, "_model_name": "synthetic/old", "_lock": threading.Lock(), + "_TIER_RERANKERS": {"edge": "synthetic/default"}, + }) + + def test_reranker_actual_constructor_releases_old_reference_and_honors_exact_name(self) -> None: + reranker = self.load_reranker() + class Model: + pass + reranker._model = Model() + old = weakref.ref(reranker._model) + def build(name, **kwargs): + gc.collect() + self.assertIsNone(old()) + self.assertEqual(name, "synthetic/deep") + return Model() + with patch.dict(sys.modules, {"sentence_transformers": types.SimpleNamespace(CrossEncoder=build)}): + result = reranker.get_reranker(model_name="synthetic/deep", device="cpu") + self.assertIs(reranker._model, result) + self.assertEqual(reranker._model_name, "synthetic/deep") + + def test_same_name_uncertified_proxy_is_replaced_without_daemon_preflight(self) -> None: + reranker = self.load_reranker() + reranker._model, reranker._model_name = object(), "synthetic/deep" + old = reranker._model + self.fixture.client.use_model_server = lambda: True + with patch.object(self.fixture.client, "_request_with_autostart", side_effect=AssertionError("Unexpected preflight")): + result = reranker.get_reranker("synthetic/deep", certified=True) + self.assertIsNot(result, old) + self.assertIsInstance(result, self.fixture.client.CertifiedRerankerProxy) + self.assertTrue(reranker._model_certified) + self.assertIs(reranker.get_reranker("synthetic/deep", certified=True), result) + + def test_reranker_setter_unload_and_nested_override_refuse_during_read(self) -> None: + reranker = self.load_reranker() + reranker._model = object() + self.operation = types.SimpleNamespace(reranker_id="synthetic/default", selection=None) + with self.gate.operation_lease(("synthetic",)): + for action in (lambda: reranker.set_active_tier("pro"), reranker.unload_reranker, + lambda: reranker.get_reranker(model_name="synthetic/other")): + with self.assertRaises(RuntimeError): + action() + + +@unittest.skipUnless(os.environ.get("TRUEMEMORY_TEST_NATIVE_VEC") == "1", "GPUBox-only native sqlite-vec composition") +class TestNativeRuntimeComposition(unittest.TestCase): + def test_selected_engine_crud_search_and_reopen_use_one_pair(self) -> None: + import numpy as np + import sqlite_vec + from truememory import engine as engine_module, model_client, reranker, storage, vector_search + from truememory.embedding_target import EmbeddingTarget + from truememory.tier_switch import activation, runtime + + target = EmbeddingTarget.capture("base") + selection = activation.TierSelection( + generation="1" * 32, intent_id="2" * 32, target=target, + reranker_id="synthetic/reranker", tables=target.tables, job_id="3" * 32, + tracker="canonical-v1", source_epoch="synthetic-epoch", source_schema_signature="4" * 64, + manifest_generation="5" * 32, manifest_hash="6" * 64, receipt_hash="7" * 64, + pair_digest="8" * 64, schema="9" * 64, vector_count=0, cursor=None, + ) + intent = activation.ActivationIntent( + intent_id=selection.intent_id, expected_generation=None, target=target, + reranker_id=selection.reranker_id, job_id=selection.job_id, tracker=selection.tracker, + source_epoch=selection.source_epoch, source_schema_signature=selection.source_schema_signature, + state="db_selected", + ) + handles, encodes, projections = [], [], [] + engine = engine_module.TrueMemoryEngine(":memory:") + self.addCleanup(engine.close) + + def create(path): + self.assertEqual(str(path), ":memory:") + conn = storage.create_db(":memory:") + handles.append(conn) + conn.enable_load_extension(True) + sqlite_vec.load(conn) + conn.enable_load_extension(False) + vector_search.init_prepared_target_tables(conn, target) + for table in ("vec_messages_edge", "vec_messages_sep_edge"): + conn.execute(f"CREATE VIRTUAL TABLE {table} USING vec0(embedding float[256] distance_metric=cosine)") + conn.execute("INSERT INTO vector_cache_registry(tier_group,vec_table,sep_table,model_name,embedding_dim) VALUES (?,?,?,?,?)", + (target.tier_group, *target.tables, target.model_id, target.dimension)) + # This controlled fixture seeds an already selected journal; certification + # faults are covered by the separate activation-journal suite. + activation._put(conn, activation._SELECTED, activation._dump(selection)) + activation._put(conn, activation._INTENT, activation._dump(intent)) + activation._put(conn, "embed_model", target.model_id) + activation._put(conn, "embed_dim", str(target.dimension)) + conn.commit() + return conn + + class Encoder: + device = "cpu" + def encode(self, texts, **kwargs): + operation = runtime.current_operation(engine.conn) + if operation is None or operation.selection != selection or operation.tables != target.tables: + raise AssertionError("Encode escaped selected serving operation") + encodes.append(operation) + result = np.zeros((len(texts), target.dimension), dtype=np.float32) + result[:, 0] = 1.0 + return result + + def embedding(frozen, **kwargs): + self.assertEqual(frozen, target) + self.assertFalse(engine.conn.in_transaction) + projections.append(frozen) + return Encoder() + + class CrossEncoder: + device = "cpu" + def __init__(self, name, **kwargs): + if name != selection.reranker_id: + raise AssertionError("Unexpected synthetic reranker identity") + def predict(self, pairs, **kwargs): + return np.ones(len(pairs), dtype=np.float32) + + with ExitStack() as stack: + stack.enter_context(patch.multiple(vector_search, _model=None, _frozen_embedding_target=None, + EMBEDDING_MODEL="model2vec", _embedding_dim=256, _model_generation=0)) + stack.enter_context(patch.multiple(reranker, _model=None, _model_name="synthetic/old", _active_tier="edge", + _frozen_reranker_id=None, _model_certified=False)) + stack.enter_context(patch.multiple(runtime, _local=threading.local(), _acknowledged=None)) + stack.enter_context(patch.object(engine_module, "create_db", side_effect=create)) + stack.enter_context(patch.object(vector_search, "_load_frozen_embedding_target", side_effect=embedding)) + stack.enter_context(patch.object(model_client, "use_model_server", return_value=False)) + absent = object() + previous_transformers = sys.modules.get("sentence_transformers", absent) + + def restore_transformers() -> None: + if previous_transformers is absent: + sys.modules.pop("sentence_transformers", None) + else: + sys.modules["sentence_transformers"] = previous_transformers + + # A whole-registry mock would evict torch first imported by device + # detection here, leaving live native objects unsafe to re-import. + stack.callback(restore_transformers) + sys.modules["sentence_transformers"] = types.SimpleNamespace(CrossEncoder=CrossEncoder) + stack.enter_context(patch.object(engine, "_maybe_auto_consolidate", return_value=None)) + stack.enter_context(patch.object(engine, "_maybe_startup_consolidate", return_value=None)) + for flag in ("_has_personality", "_has_style_vec", "_has_temporal", "_has_salience", "_has_consolidation", + "_has_predictive", "_has_hyde", "_has_clustering", "_has_reranker"): + setattr(engine, flag, False) + first = engine.add("synthetic cobalt record") + self.assertTrue(engine.ready) + self.assertEqual(projections, [target]) + for table in target.tables: + self.assertEqual(engine.conn.execute(f"SELECT rowid FROM {table}").fetchall(), [(first["id"],)]) + self.assertEqual(engine.conn.execute("SELECT count(*) FROM vec_messages_edge").fetchone()[0], 0) + self.assertEqual(engine.update(first["id"], "synthetic cobalt corrected")["content"], "synthetic cobalt corrected") + self.assertEqual(engine.search_vectors_raw("synthetic cobalt")[0]["id"], first["id"]) + before = len(encodes) + result = engine.search_agentic("synthetic cobalt", limit=1, max_rounds=1, + use_hyde=False, use_reranker=False, use_clustering=False) + self.assertEqual(result[0]["id"], first["id"]) + self.assertGreater(len(encodes), before) + self.assertEqual(len({id(op) for op in encodes[before:]}), 1) + self.assertTrue(engine.delete(first["id"])) + for table in target.tables: + self.assertEqual(engine.conn.execute(f"SELECT count(*) FROM {table}").fetchone()[0], 0) + engine.add("synthetic second record") + self.assertTrue(engine.delete_all()) + self.assertEqual(engine.conn.execute("SELECT count(*) FROM messages").fetchone()[0], 0) + original = engine.conn + original.close() + engine.add("synthetic reopened record") + self.assertIsNot(engine.conn, original) + self.assertTrue(engine.ready) + self.assertIs(engine._runtime_vector_connection, engine.conn) + self.assertEqual(projections, [target]) + self.assertEqual(len(handles), 2) + self.assertIsNone(runtime.current_operation(engine.conn)) + native_modules = {name: id(module) for name, module in sys.modules.items() + if name == "torch" or name.startswith("torch.")} + self.assertEqual({name: id(sys.modules.get(name)) for name in native_modules}, native_modules) + self.assertIs(sys.modules.get("sentence_transformers", absent), previous_transformers) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test-tier-writer-fence-795.py b/tests/test-tier-writer-fence-795.py index 73d61d6f..67334e41 100644 --- a/tests/test-tier-writer-fence-795.py +++ b/tests/test-tier-writer-fence-795.py @@ -189,7 +189,7 @@ def engine(self, *, vectors=True): )] helpers = [n for n in tree.body if isinstance(n, ast.FunctionDef) and n.name in {"_delete_in_chunks", "_resolve_vec_tables"}] namespace = dict(__builtins__=self.vector.__dict__["__builtins__"], logger=logging.getLogger(__name__), - MAX_CONTENT_LENGTH=10000, sqlite3=sqlite3, + MAX_CONTENT_LENGTH=10000, sqlite3=sqlite3, engine_operation=lambda function: function, insert_message=self.modules["storage"].insert_message, update_message=self.modules["storage"].update_message, delete_message=self.modules["storage"].delete_message) diff --git a/truememory/__init__.py b/truememory/__init__.py index b6039ab5..85d107ab 100644 --- a/truememory/__init__.py +++ b/truememory/__init__.py @@ -125,6 +125,7 @@ "consolidation", "predictive", "query_classifier", "reranker", "hyde", "clustering", "tier_config", + "embedding_target", "maintenance", "rebuild_source", "tier_switch", ] diff --git a/truememory/engine.py b/truememory/engine.py index ec096cee..685eb162 100644 --- a/truememory/engine.py +++ b/truememory/engine.py @@ -47,6 +47,8 @@ from truememory.fts_search import search_fts from truememory._platform import _env_int +from truememory.tier_switch.runtime import engine_operation + logger = logging.getLogger(__name__) # ─────────────────────────────────────────────────────────────────────────── @@ -338,8 +340,12 @@ def __init__(self, db_path: str | Path = None, self.conn: sqlite3.Connection | None = None self.ready = False self.stats: dict = {} - self._write_lock = threading.Lock() + from truememory.tier_switch.runtime import ConnectionWriteLock + self._write_lock = ConnectionWriteLock(self) self._init_lock = threading.Lock() + self._runtime_initialized = False + self._runtime_vector_generation = None + self._runtime_vector_connection = None # L5 surprise rerank boost coefficient. None = resolve from env # var / default at call-time via _get_alpha_surprise(). @@ -372,95 +378,147 @@ def __init__(self, db_path: str | Path = None, # Auto-connect (production API) # ────────────────────────────────────────────────────────────────────── - def _ensure_connection(self, *, _suppress_maintenance: bool = False) -> None: - """Open (or create) the database and load extensions if needed. - - Called automatically by :meth:`add`, :meth:`search`, and other - public methods so users never have to call ``ingest()`` or - ``open()`` manually for simple CRUD workflows. - """ - if self.conn is not None: - # SRE-01: a cached handle can be poisoned by a transient disk I/O - # error (e.g. a stray -wal deletion) or closed out from under us. - # Previously the short-circuit returned the dead handle forever, so a - # long-lived MCP server failed every add/search/recall until manual - # restart. Probe the handle; if the probe fails, drop it and fall - # through to reconnect so the server self-heals once the FS recovers. - try: - self.conn.execute("PRAGMA schema_version") - if not _suppress_maintenance: - self._maybe_auto_consolidate() - return - except sqlite3.Error: - logger.warning( - "DB connection probe failed for %s; reconnecting", - self.db_path, exc_info=True, - ) - try: - self.conn.close() - except sqlite3.Error: - pass - self.conn = None - self._has_vectors = False - - with self._init_lock: + def _open_connection_handle(self) -> None: + """Open a plain handle before serving admission or vector initialization.""" + from truememory.tier_switch.runtime import _connection_read + with _connection_read(self.conn, self._init_lock, None, None, + transaction_owner=self._write_lock, allow_closed=True), \ + _connection_read(self.conn, self._write_lock, None, None, allow_closed=True): if self.conn is not None: - if not _suppress_maintenance: - self._maybe_auto_consolidate() - return - - # Create parent directory if using a real path - db_str = str(self.db_path) - if db_str != ":memory:": - # M-89: the DB dir holds real memories/PII — owner-only (0700). + try: + self.conn.execute("PRAGMA schema_version") + return + except (sqlite3.ProgrammingError, sqlite3.OperationalError) as exc: + closed = (isinstance(exc, sqlite3.ProgrammingError) + and str(exc) == "Cannot operate on a closed database.") + code = getattr(exc, "sqlite_errorcode", None) + io_error = (isinstance(exc, sqlite3.OperationalError) and ( + str(exc) == "disk I/O error" + or (isinstance(code, int) and code & 0xff == getattr(sqlite3, "SQLITE_IOERR", 10)))) + if not closed: + if not io_error or self.conn.in_transaction: + raise + self.conn.close() + self.conn = None + self._has_vectors = False + self.ready = False + self._runtime_initialized = False + self._runtime_vector_generation = None + self._runtime_vector_connection = None + if str(self.db_path) != ":memory:": self.db_path.parent.mkdir(parents=True, exist_ok=True) try: self.db_path.parent.chmod(0o700) except OSError: pass - self.conn = create_db(self.db_path) + self.ready = False + self._runtime_initialized = False + self._runtime_vector_generation = None + self._runtime_vector_connection = None - # Load sqlite-vec extension - global _vectors_load_error - if _HAS_VECTOR: - try: + def _ensure_connection(self, *, _suppress_maintenance: bool = False) -> None: + """Reconcile committed selection before any vector compatibility work.""" + from truememory.tier_switch.runtime import serving_operation + self._open_connection_handle() + with serving_operation(self.conn, connection_lock=self._write_lock): + self._initialize_connection(_suppress_maintenance=_suppress_maintenance) + + def _initialize_connection(self, *, _suppress_maintenance: bool = False) -> None: + from truememory.tier_switch.runtime import current_operation, TierRuntimeError, _connection_read + operation = current_operation(self.conn) + if operation.selection is not None: + if (getattr(self, "_runtime_initialized", False) + and getattr(self, "_runtime_vector_connection", None) is self.conn + and getattr(self, "_runtime_vector_generation", None) == operation.selection.generation): + if not _suppress_maintenance: + self._maybe_auto_consolidate() + return + initialized = False + with _connection_read(self.conn, self._init_lock, None, None, + transaction_owner=self._write_lock), \ + _connection_read(self.conn, self._write_lock, None, None): + if not (getattr(self, "_runtime_initialized", False) + and getattr(self, "_runtime_vector_connection", None) is self.conn + and getattr(self, "_runtime_vector_generation", None) == operation.selection.generation): + if self.conn.in_transaction: + raise TierRuntimeError("Selected vector initialization requires a clean connection") + if not _HAS_VECTOR: + raise TierRuntimeError("Selected vector dependencies are unavailable") import sqlite_vec self.conn.enable_load_extension(True) - sqlite_vec.load(self.conn) - self.conn.enable_load_extension(False) - # init_vec_table() runs _check_embedder_compatibility(), - # which raises TrueMemoryMigrationError on a dim/model - # mismatch. That must NOT be swallowed as a generic - # FTS-only fallback (M-12/M-46): a silent drop to FTS hides - # the degradation from truememory_stats.health and lets - # add() silently fail vec INSERTs. Surface the migration - # guidance so the user can re-embed instead. - init_vec_table(self.conn) try: - from truememory.vector_search import migrate_legacy_vec_tables - migrate_legacy_vec_tables(self.conn) - except Exception: - logger.debug("Legacy vec table migration skipped", exc_info=True) + sqlite_vec.load(self.conn) + finally: + self.conn.enable_load_extension(False) + init_vec_table(self.conn) self._has_vectors = True - except TrueMemoryMigrationError as exc: - # M-12/M-46: record the degradation in module state so - # health surfaces it, then propagate the actionable - # migration guidance instead of falling back to FTS-only. - _vectors_load_error = f"{type(exc).__name__}: {exc}" - self._has_vectors = False - raise - except Exception as exc: - _vectors_load_error = f"{type(exc).__name__}: {exc}" - logger.warning("Failed to load sqlite-vec — FTS-only mode: %s", exc) - self._has_vectors = False + self._has_hybrid = _HAS_HYBRID + self._purge_legacy_entity_profile_summaries() + self.ready = self._runtime_initialized = True + self._runtime_vector_generation = operation.selection.generation + self._runtime_vector_connection = self.conn + initialized = True + if not _suppress_maintenance: + if initialized: + self._maybe_startup_consolidate() + else: + self._maybe_auto_consolidate() + return + if getattr(self, "_runtime_initialized", True) or getattr(self, "ready", False): + if not _suppress_maintenance: + self._maybe_auto_consolidate() + return + with _connection_read(self.conn, self._init_lock, None, None, + transaction_owner=self._write_lock): + if self._runtime_initialized: + if not _suppress_maintenance: + self._maybe_auto_consolidate() + return + with _connection_read(self.conn, self._write_lock, None, None): + if self.conn.in_transaction: + raise TierRuntimeError("Vector initialization requires a clean connection") + # Load sqlite-vec extension + global _vectors_load_error + if _HAS_VECTOR: + try: + import sqlite_vec + self.conn.enable_load_extension(True) + sqlite_vec.load(self.conn) + self.conn.enable_load_extension(False) + # init_vec_table() runs _check_embedder_compatibility(), + # which raises TrueMemoryMigrationError on a dim/model + # mismatch. That must NOT be swallowed as a generic + # FTS-only fallback (M-12/M-46): a silent drop to FTS hides + # the degradation from truememory_stats.health and lets + # add() silently fail vec INSERTs. Surface the migration + # guidance so the user can re-embed instead. + init_vec_table(self.conn) + try: + from truememory.vector_search import migrate_legacy_vec_tables + migrate_legacy_vec_tables(self.conn) + except Exception: + logger.debug("Legacy vec table migration skipped", exc_info=True) + self._has_vectors = True + except TrueMemoryMigrationError as exc: + # M-12/M-46: record the degradation in module state so + # health surfaces it, then propagate the actionable + # migration guidance instead of falling back to FTS-only. + _vectors_load_error = f"{type(exc).__name__}: {exc}" + self._has_vectors = False + raise + except Exception as exc: + _vectors_load_error = f"{type(exc).__name__}: {exc}" + logger.warning("Failed to load sqlite-vec — FTS-only mode: %s", exc) + self._has_vectors = False - self._has_hybrid = _HAS_HYBRID and self._has_vectors + self._has_hybrid = _HAS_HYBRID and self._has_vectors # Qwen3 NaN fix: macOS SDPA kernel produces NaN embeddings. # Re-embed once for Base/Pro users on macOS. import sys as _sys - if _sys.platform == "darwin" and self._has_vectors: + if (_sys.platform == "darwin" and self._has_vectors + and current_operation(self.conn).selection is None): try: _embed_model = resolve_tier() if _embed_model in ("base", "pro", "qwen3_256"): @@ -585,11 +643,12 @@ def _bg_wrapper(db_path): # MEMORIST-L4 migration: purge legacy entity_profile summary rows. # M-84: previously this only ran in the deprecated open() path, so # production (which uses _ensure_connection) never purged them. - self._purge_legacy_entity_profile_summaries() - - self.ready = True - if not _suppress_maintenance: - self._maybe_startup_consolidate() + with _connection_read(self.conn, self._write_lock, None, None): + self._purge_legacy_entity_profile_summaries() + self.ready = True + self._runtime_initialized = True + if not _suppress_maintenance: + self._maybe_startup_consolidate() def _purge_legacy_entity_profile_summaries(self) -> None: """Delete legacy ``period='entity_profile'`` summary rows once. @@ -671,6 +730,7 @@ def _maybe_startup_consolidate(self) -> None: # Production CRUD API # ────────────────────────────────────────────────────────────────────── + @engine_operation def add( self, content: str, @@ -889,6 +949,7 @@ def _bg_consolidate(self) -> None: """Compatibility entry point; the coordinator owns the worker handle.""" self._maybe_auto_consolidate() + @engine_operation def delete(self, memory_id: int) -> bool: """Delete a memory by ID. @@ -903,6 +964,7 @@ def delete(self, memory_id: int) -> bool: self._maybe_auto_consolidate() return deleted + @engine_operation def delete_all(self, user_id: str | None = None) -> bool: """Delete all memories, optionally filtered by user. @@ -1136,6 +1198,7 @@ def consolidate(self) -> dict[str, str]: self._apply_manual_maintenance_capabilities(report.results) return format_maintenance_report(report) + @engine_operation def update(self, memory_id: int, content: str | None = None, **fields) -> dict | None: """Update a memory. @@ -1298,108 +1361,110 @@ def open(self, rebuild_vectors: bool = True) -> "TrueMemoryEngine": self.conn.execute("PRAGMA cache_size=-64000") self.conn.execute("PRAGMA mmap_size=268435456") - # Detect available tables - tables = { - row[0] - for row in self.conn.execute( - "SELECT name FROM sqlite_master WHERE type='table'" - ).fetchall() - } + from truememory.tier_switch.runtime import serving_operation + with serving_operation(self.conn, connection_lock=self._write_lock): + # Detect available tables + tables = { + row[0] + for row in self.conn.execute( + "SELECT name FROM sqlite_master WHERE type='table'" + ).fetchall() + } - # ── MEMORIST-L4 migration: purge legacy entity_profile summary rows ─ - # Shared with _ensure_connection (M-84): the purge logic lives in one - # idempotent helper so both the production path and this deprecated - # open() path behave identically. - self._purge_legacy_entity_profile_summaries() - - # Load sqlite-vec extension if available. - # upgrade DEBUG → WARNING and track failure in a - # module-level state so ``truememory_stats.health`` can report - # "search is FTS-only because sqlite-vec failed to load". - global _vectors_load_error - if _HAS_VECTOR: - try: - import sqlite_vec - self.conn.enable_load_extension(True) - sqlite_vec.load(self.conn) - _vectors_load_error = None - except Exception as e: - _vectors_load_error = f"{type(e).__name__}: {e}" - logger.warning( - "sqlite-vec unavailable (%s); falling back to FTS-only " - "search for this process. See " - "https://github.com/buildingjoshbetter/TrueMemory for " - "platform notes.", - _vectors_load_error, - ) + # ── MEMORIST-L4 migration: purge legacy entity_profile summary rows ─ + # Shared with _ensure_connection (M-84): the purge logic lives in one + # idempotent helper so both the production path and this deprecated + # open() path behave identically. + self._purge_legacy_entity_profile_summaries() - # Check for vector tables — rebuild if missing. - # if metadata names a different embedder, refuse silent - # rebuild; route the user through truememory_configure() instead. - self._has_vectors = False - if _HAS_VECTOR: - from truememory.vector_search import ( - _active_vec_table, - _check_embedder_compatibility, - vectors_are_built, - ) - vec_tbl = _active_vec_table(self.conn) - # vectors_are_built() returns False for a table left ``in_progress`` - # by an interrupted build (issue #647) — so a partial/empty table is - # rebuilt rather than trusted, not just one that's missing. - if vectors_are_built(self.conn, vec_tbl): - # M-12: the exists-path previously trusted the table on a bare - # SELECT 1 and never checked embedder/dim compatibility. A dim - # mismatch then surfaced only as a per-query OperationalError - # caught at DEBUG → permanent silent FTS-only with - # _vectors_load_error=None (health reports vectors healthy) and - # silent vec-INSERT failures in add(). Run the compatibility - # check here so a mismatch surfaces as actionable migration - # guidance instead. + # Load sqlite-vec extension if available. + # upgrade DEBUG → WARNING and track failure in a + # module-level state so ``truememory_stats.health`` can report + # "search is FTS-only because sqlite-vec failed to load". + global _vectors_load_error + if _HAS_VECTOR: try: - _check_embedder_compatibility(self.conn) - self._has_vectors = True - except TrueMemoryMigrationError as exc: - _vectors_load_error = f"{type(exc).__name__}: {exc}" - self._has_vectors = False - raise - else: - logger.warning( - "Vector table %r missing or unreadable; attempting rebuild " - "with current model=%s", - vec_tbl, - resolve_tier(), + import sqlite_vec + self.conn.enable_load_extension(True) + sqlite_vec.load(self.conn) + _vectors_load_error = None + except Exception as e: + _vectors_load_error = f"{type(e).__name__}: {e}" + logger.warning( + "sqlite-vec unavailable (%s); falling back to FTS-only " + "search for this process. See " + "https://github.com/buildingjoshbetter/TrueMemory for " + "platform notes.", + _vectors_load_error, + ) + + # Check for vector tables — rebuild if missing. + # if metadata names a different embedder, refuse silent + # rebuild; route the user through truememory_configure() instead. + self._has_vectors = False + if _HAS_VECTOR: + from truememory.vector_search import ( + _active_vec_table, + _check_embedder_compatibility, + vectors_are_built, ) - if rebuild_vectors: - _check_rebuild_allowed(self.conn) # raises on model drift + vec_tbl = _active_vec_table(self.conn) + # vectors_are_built() returns False for a table left ``in_progress`` + # by an interrupted build (issue #647) — so a partial/empty table is + # rebuilt rather than trusted, not just one that's missing. + if vectors_are_built(self.conn, vec_tbl): + # M-12: the exists-path previously trusted the table on a bare + # SELECT 1 and never checked embedder/dim compatibility. A dim + # mismatch then surfaced only as a per-query OperationalError + # caught at DEBUG → permanent silent FTS-only with + # _vectors_load_error=None (health reports vectors healthy) and + # silent vec-INSERT failures in add(). Run the compatibility + # check here so a mismatch surfaces as actionable migration + # guidance instead. try: - init_vec_table(self.conn) - n = build_vectors(self.conn) - self._has_vectors = n > 0 - logger.info( - "Vector table %r rebuilt with %d vectors (model=%s)", - vec_tbl, n, resolve_tier(), - ) - except TrueMemoryMigrationError: - raise - except Exception: - logger.exception("Vector table rebuild failed") + _check_embedder_compatibility(self.conn) + self._has_vectors = True + except TrueMemoryMigrationError as exc: + _vectors_load_error = f"{type(exc).__name__}: {exc}" self._has_vectors = False + raise + else: + logger.warning( + "Vector table %r missing or unreadable; attempting rebuild " + "with current model=%s", + vec_tbl, + resolve_tier(), + ) + if rebuild_vectors: + _check_rebuild_allowed(self.conn) # raises on model drift + try: + init_vec_table(self.conn) + n = build_vectors(self.conn) + self._has_vectors = n > 0 + logger.info( + "Vector table %r rebuilt with %d vectors (model=%s)", + vec_tbl, n, resolve_tier(), + ) + except TrueMemoryMigrationError: + raise + except Exception: + logger.exception("Vector table rebuild failed") + self._has_vectors = False - self._has_hybrid = _HAS_HYBRID and self._has_vectors - self._has_clustering = _HAS_CLUSTERING and "message_clusters" in tables + self._has_hybrid = _HAS_HYBRID and self._has_vectors + self._has_clustering = _HAS_CLUSTERING and "message_clusters" in tables - # Count messages for stats - try: - count = self.conn.execute("SELECT COUNT(*) FROM messages").fetchone()[0] - self.stats["message_count"] = count - except Exception: - logger.debug("Failed to count messages in open()", exc_info=True) - self.stats["message_count"] = 0 + # Count messages for stats + try: + count = self.conn.execute("SELECT COUNT(*) FROM messages").fetchone()[0] + self.stats["message_count"] = count + except Exception: + logger.debug("Failed to count messages in open()", exc_info=True) + self.stats["message_count"] = 0 - self.ready = True - self._maybe_auto_consolidate() - return self + self.ready = True + self._maybe_auto_consolidate() + return self # ────────────────────────────────────────────────────────────────────── # Ingestion @@ -1710,6 +1775,7 @@ def ingest(self, data_path: str | Path) -> dict: # Search — full 6-layer pipeline # ────────────────────────────────────────────────────────────────────── + @engine_operation def search_vectors_raw(self, query: str, limit: int = 5) -> list[dict] | None: """Pure vector cosine similarity search, or None if unavailable. @@ -1726,6 +1792,7 @@ def search_vectors_raw(self, query: str, limit: int = 5) -> list[dict] | None: except Exception: return None + @engine_operation def search(self, query: str, limit: int = 10, _skip_surprise_boost: bool = False, _skip_reranker: bool = False, _skip_salience_guard: bool = False, include_directives: bool = False) -> list[dict]: """ Main search pipeline. @@ -2123,6 +2190,7 @@ def search(self, query: str, limit: int = 10, _skip_surprise_boost: bool = False # Search — agentic multi-round retrieval # ────────────────────────────────────────────────────────────────────── + @engine_operation def search_agentic( self, query: str, @@ -2513,6 +2581,7 @@ def _quality_self_check(self, query: str, results: list[dict], limit: int) -> li # Search — simple FTS5-only (for benchmarking) # ────────────────────────────────────────────────────────────────────── + @engine_operation def search_simple(self, query: str, limit: int = 10) -> list[dict]: """ Simple FTS5-only search (no vector, no layers). @@ -2635,6 +2704,9 @@ def close(self): logger.debug("Failed to close database connection", exc_info=True) self.conn = None self.ready = False + self._runtime_initialized = False + self._runtime_vector_generation = None + self._runtime_vector_connection = None def __enter__(self): return self diff --git a/truememory/hybrid.py b/truememory/hybrid.py index 8ee56ae4..406f9b04 100644 --- a/truememory/hybrid.py +++ b/truememory/hybrid.py @@ -41,6 +41,8 @@ import sqlite3 from collections import defaultdict +from truememory.tier_switch.runtime import database_operation + log = logging.getLogger(__name__) @@ -123,6 +125,7 @@ def reciprocal_rank_fusion( _CANDIDATE_POOL = 200 +@database_operation def search_hybrid( conn: sqlite3.Connection, query: str, diff --git a/truememory/reranker.py b/truememory/reranker.py index 2d7e55b1..08babe97 100644 --- a/truememory/reranker.py +++ b/truememory/reranker.py @@ -27,6 +27,8 @@ import logging import threading +from collections.abc import Iterator +from contextlib import contextmanager from typing import TYPE_CHECKING log = logging.getLogger(__name__) @@ -40,6 +42,7 @@ _model = None _model_name: str = "cross-encoder/ms-marco-MiniLM-L-6-v2" +_model_certified = False _lock = threading.Lock() _inference_lock = threading.Lock() # Protects concurrent model.predict() calls @@ -63,6 +66,7 @@ _TIER_RERANKERS = {t: cfg["reranker"] for t, cfg in _TIERS.items()} _active_tier: str | None = None # None = not yet resolved; resolved lazily +_frozen_reranker_id: str | None = None def get_reranker_name_for_tier(tier: str) -> str: @@ -134,12 +138,15 @@ def set_active_tier(tier: str) -> None: resolve to the new tier's reranker. Empty / unknown tier falls back to "edge". """ - global _active_tier - if not tier: - _active_tier = "edge" - return - t = tier.strip().lower() - _active_tier = t if t in ("edge", "base", "pro", "custom") else "edge" + global _active_tier, _frozen_reranker_id + from truememory.tier_switch.serving import exclusive_activation + with exclusive_activation(): + _frozen_reranker_id = None + if not tier: + _active_tier = "edge" + return + t = tier.strip().lower() + _active_tier = t if t in ("edge", "base", "pro", "custom") else "edge" def get_current_reranker_name() -> str: @@ -150,11 +157,35 @@ def get_current_reranker_name() -> str: is called. """ global _active_tier + import sys + runtime = sys.modules.get("truememory.tier_switch.runtime") + operation = runtime.current_runtime_operation() if runtime is not None else None + if operation is not None: + return operation.reranker_id + if _frozen_reranker_id is not None: + return _frozen_reranker_id if _active_tier is None: _active_tier = _resolve_tier_from_env_and_config() return get_reranker_name_for_tier(_active_tier) +def apply_frozen_reranker( + tier: str, model_name: str, *, deadline: float | None = None, + cancelled: threading.Event | None = None, +) -> None: + """Publish an explicit default only after its existing singleton is ready.""" + global _active_tier, _frozen_reranker_id + from truememory.tier_switch.serving import exclusive_activation + with exclusive_activation(deadline=deadline, cancelled=cancelled): + # Release before loading a distinct reranker. An acknowledgement belongs + # to the runtime bridge and is not published on this partial boundary. + if _model_name != model_name: + unload_reranker() + get_reranker(model_name=model_name, deadline=deadline, cancelled=cancelled, certified=True) + _active_tier = tier + _frozen_reranker_id = model_name + + def unload_reranker( should_unload: "Callable[[], bool] | None" = None, ) -> bool: @@ -170,7 +201,8 @@ def unload_reranker( ``True`` if the model was actually unloaded, ``False`` if skipped. """ global _model - with _lock: + from truememory.tier_switch.serving import exclusive_activation + with exclusive_activation(), _lock: if should_unload is not None and not should_unload(): return False if _model is None: @@ -179,7 +211,26 @@ def unload_reranker( return True -def get_reranker(model_name: str | None = None, device: str | None = None): +@contextmanager +def _reranker_load_lock(control: object) -> Iterator[None]: + from truememory.tier_switch.serving import _check, _wait_seconds + acquired = False + try: + while not acquired: + _check(control) + seconds = _wait_seconds(control) + acquired = _lock.acquire() if seconds is None else _lock.acquire(timeout=seconds) + _check(control) + yield + finally: + if acquired: + _lock.release() + + +def get_reranker( + model_name: str | None = None, device: str | None = None, *, + deadline: float | None = None, cancelled: threading.Event | None = None, certified: bool = False, +): """ Lazy-load the cross-encoder reranker (singleton). @@ -199,37 +250,62 @@ def get_reranker(model_name: str | None = None, device: str | None = None): Returns: A ``sentence_transformers.CrossEncoder`` instance or RerankerProxy. """ - global _model, _model_name + global _model, _model_name, _model_certified + from truememory.tier_switch.serving import _control, _check + control = _control(None, deadline, cancelled) + _check(control) + if type(certified) is not bool: + raise TypeError("Reranker certification must be boolean") name = model_name or get_current_reranker_name() + from truememory.tier_switch.runtime import current_runtime_operation, TierRuntimeError + operation = current_runtime_operation() + if operation is not None and name != operation.reranker_id: + raise TierRuntimeError("Reranker override must be captured before serving admission") + if operation is not None and operation.selection is not None: + certified = True cached_model, cached_name = _model, _model_name - if cached_model is not None and name == cached_name: + if cached_model is not None and name == cached_name and (not certified or _model_certified): return cached_model - with _lock: - cached_model, cached_name = _model, _model_name - if cached_model is not None and name == cached_name: - return cached_model - - from truememory.model_client import use_model_server, get_reranker_proxy - if use_model_server(): - proxy = get_reranker_proxy(model_name=name) - _model = proxy + from truememory.tier_switch.runtime import require_model_load_allowed + require_model_load_allowed() + from contextlib import nullcontext + from truememory.tier_switch.serving import exclusive_activation + with exclusive_activation(deadline=deadline, cancelled=cancelled) if operation is None else nullcontext(): + with _reranker_load_lock(control): + cached_model, cached_name = _model, _model_name + if cached_model is not None and name == cached_name and (not certified or _model_certified): + return cached_model + + _model = cached_model = None # Release a distinct cached slot before construction. + from truememory.model_client import use_model_server, get_reranker_proxy + if use_model_server(): + _check(control) + if certified: + from truememory.model_client import CertifiedRerankerProxy + proxy = CertifiedRerankerProxy(name) + else: + proxy = get_reranker_proxy(model_name=name) + _model = proxy + _model_name = name + _model_certified = certified + return proxy + + if device is None: + # Issue #577: honor TRUEMEMORY_DEVICE (cpu|mps|cuda|auto) before + # auto-detection. An explicit device= parameter wins over the env. + from truememory.mps_utils import auto_detect_device, resolve_device + device = resolve_device(auto_detect_device()) + + from truememory.mps_utils import ensure_mps_memory_budget + ensure_mps_memory_budget(device) + from sentence_transformers import CrossEncoder + _check(control) + _model = CrossEncoder(name, device=device) _model_name = name - return proxy - - if device is None: - # Issue #577: honor TRUEMEMORY_DEVICE (cpu|mps|cuda|auto) before - # auto-detection. An explicit device= parameter wins over the env. - from truememory.mps_utils import auto_detect_device, resolve_device - device = resolve_device(auto_detect_device()) - - from truememory.mps_utils import ensure_mps_memory_budget - ensure_mps_memory_budget(device) - from sentence_transformers import CrossEncoder - _model = CrossEncoder(name, device=device) - _model_name = name - result = _model - return result + _model_certified = True + result = _model + return result # --------------------------------------------------------------------------- diff --git a/truememory/tier_switch/runtime.py b/truememory/tier_switch/runtime.py new file mode 100644 index 00000000..f75310e4 --- /dev/null +++ b/truememory/tier_switch/runtime.py @@ -0,0 +1,419 @@ +"""Process-local projection and operation snapshots of durable tier selection. + +Database selection is authoritative. This module never writes configuration, +selection metadata or vector rows, and stores no per-database model cache. +""" + +from __future__ import annotations + +import os +import re +import sqlite3 +import threading +import time +from collections.abc import Callable, Iterator +from contextlib import contextmanager +from dataclasses import dataclass, field +from functools import wraps +from typing import ParamSpec, TypeVar + +from truememory.tier_switch import serving +from truememory.tier_switch.activation import TierSelection, _guard_storage, _state, read_activation_state + +P = ParamSpec("P") +R = TypeVar("R") + + +class TierRuntimeError(RuntimeError): + """The process cannot safely serve the captured database selection.""" + + +class ConnectionWriteLock: + """Identify Engine-owned SQL without claiming ownership of caller SQL.""" + + def __init__(self, engine: object) -> None: + import weakref + self._engine = weakref.ref(engine) + self._lock = threading.Lock() + self._owner: tuple[int, object, bool] | None = None + + def acquire(self, blocking: bool = True, timeout: float = -1) -> bool: + acquired = self._lock.acquire(blocking, timeout) + if not acquired: + return False + try: + engine = self._engine() + conn = None if engine is None else engine.conn + try: + borrowed = conn is not None and conn.in_transaction + except sqlite3.ProgrammingError as exc: + if str(exc) != "Cannot operate on a closed database.": + raise + borrowed = False + self._owner = (threading.get_ident(), conn, borrowed) + except BaseException: + self._lock.release() + raise + return True + + def owns_other_transaction(self, conn: sqlite3.Connection) -> bool: + owner = self._owner + return (owner is not None and owner[0] != threading.get_ident() + and owner[1] is conn and not owner[2]) + + def release(self) -> None: + owner = self._owner + if owner is None or owner[0] != threading.get_ident(): + raise RuntimeError("Only the Engine writer owner can release its lock") + self._owner = None + self._lock.release() + + def locked(self) -> bool: + return self._lock.locked() + + def __enter__(self) -> ConnectionWriteLock: + self.acquire() + return self + + def __exit__(self, *exc: object) -> None: + self.release() + + +_local = threading.local() +_acknowledged: tuple[str, ...] | None = None +_pid = os.getpid() + + +def _after_fork() -> None: + global _local, _acknowledged, _pid + _local = threading.local() + _acknowledged = None + _pid = os.getpid() + + +if hasattr(os, "register_at_fork"): + os.register_at_fork(after_in_child=_after_fork) + + +def _reranker_id(value: str) -> str: + if (type(value) is not str or not 1 <= len(value) <= 512 + or re.fullmatch(r"[\w][\w.\-]*(/[\w][\w.\-]*)?", value) is None): + raise TierRuntimeError("An explicit bounded reranker identity is required") + return value + + +def selection_key(selection: TierSelection, reranker_id: str | None = None) -> tuple[str, ...]: + if type(selection) is not TierSelection: + raise TierRuntimeError("An immutable certified selection is required") + return (selection.generation, selection.source_epoch, selection.target.model_id, + str(selection.target.dimension), selection.target.tier, + selection.reranker_id if reranker_id is None else _reranker_id(reranker_id), + *selection.tables) + + +def _check(deadline: float | None, cancelled: threading.Event | None) -> None: + if cancelled is not None and cancelled.is_set(): + raise serving.ServingLeaseCancelled("Runtime admission was cancelled") + if deadline is not None and time.monotonic() >= deadline: + raise serving.ServingLeaseTimeout("Runtime admission deadline expired") + + +def _remaining(deadline: float | None) -> float | None: + _check(deadline, None) + return None if deadline is None else deadline - time.monotonic() + + +def _modules() -> tuple[object, object]: + from truememory import reranker, vector_search + return vector_search, reranker + + +def _read_selection(conn: sqlite3.Connection) -> TierSelection | None: + # Absence preserves the legacy path without requiring the new journal's + # registry schema or inventing a certificate for existing vectors. + if not conn.execute("SELECT 1 FROM main.sqlite_master WHERE name='metadata' COLLATE NOCASE").fetchone(): + return None + if not conn.execute( + "SELECT 1 FROM main.metadata WHERE key IN ('tier_selected_v1','tier_activation_v1') LIMIT 1" + ).fetchone(): + return None + if conn.in_transaction: + _guard_storage(conn) + return _state(conn).selection + return read_activation_state(conn).selection + + +def _legacy_key(reranker_id: str | None = None) -> tuple[str, ...]: + vector, reranker = _modules() + return ("legacy", vector.EMBEDDING_MODEL, str(vector._embedding_dim), + vector.resolve_tier(), reranker.get_current_reranker_name() if reranker_id is None else reranker_id) + + +def _ready(selection: TierSelection, reranker_id: str) -> bool: + vector, reranker = _modules() + return (_pid == os.getpid() and _acknowledged == selection_key(selection, reranker_id) + and vector._frozen_embedding_target == selection.target + and vector.EMBEDDING_MODEL == selection.target.model_id + and vector._embedding_dim == selection.target.dimension + and vector._model is not None and reranker._model is not None + and reranker._model_certified + and reranker._model_name == reranker_id + and reranker._frozen_reranker_id == reranker_id + and reranker._active_tier == selection.target.tier) + + +def runtime_acknowledgement() -> tuple[str, ...] | None: + """Process-local readiness only; a config acknowledgement is unrelated.""" + return _acknowledged if _pid == os.getpid() else None + + +def apply_frozen_selection( + selection: TierSelection, *, deadline: float | None = None, + cancelled: threading.Event | None = None, reranker_id: str | None = None, +) -> None: + """Project outside SQL; the caller may already own exclusive activation. + + Failure after database publication leaves activation pending. No previous + process acknowledgement survives partial projection. This function has no + connection argument; its caller must establish the clean-SQL precondition. + """ + global _acknowledged, _pid + effective = selection.reranker_id if reranker_id is None else _reranker_id(reranker_id) + key = selection_key(selection, effective) + if getattr(_local, "operation", None) is not None: + raise TierRuntimeError("Runtime projection cannot upgrade a serving operation") + with serving.exclusive_activation(deadline=deadline, cancelled=cancelled): + _check(deadline, cancelled) + if _ready(selection, effective): + return + _acknowledged = None + vector, reranker = _modules() + _check(deadline, cancelled) + vector.apply_frozen_embedding_target(selection.target, timeout=_remaining(deadline)) + _check(deadline, cancelled) + reranker.apply_frozen_reranker(selection.target.tier, effective, deadline=deadline, cancelled=cancelled) + _check(deadline, cancelled) + _pid = os.getpid() + _acknowledged = key + + +@dataclass(frozen=True, repr=False) +class ServingOperation: + selection: TierSelection | None + tables: tuple[str, str] + tier: str + reranker_id: str + key: tuple[str, ...] = field(repr=False) + _connection: sqlite3.Connection = field(repr=False, compare=False) + _lease: serving.OperationLease = field(repr=False, compare=False) + _database: tuple[str, object] = field(repr=False, compare=False) + borrowed: bool = False + + def fork_child(self) -> RuntimeChild: + if current_operation(self._connection) is not self: + raise TierRuntimeError("Only the admitted operation can reserve a runtime child") + return RuntimeChild(self, _local.lease.fork_child()) + + +@dataclass(frozen=True, repr=False) +class RuntimeChild: + _operation: ServingOperation + _reservation: serving.ChildReservation + + @contextmanager + def join(self, conn: sqlite3.Connection | None = None, *, deadline: float | None = None, + cancelled: threading.Event | None = None) -> Iterator[ServingOperation]: + if getattr(_local, "operation", None) is not None: + raise TierRuntimeError("Runtime child must enter on an unbound thread") + with self._reservation.join(deadline=deadline, cancelled=cancelled) as lease: + original = self._operation + connection = original._connection if conn is None else conn + if connection is not original._connection: + if connection.in_transaction or _database_identity(connection) != original._database: + raise TierRuntimeError("Runtime child requires the captured clean database") + selected = _read_selection(connection) + if ((selection_key(selected) if selected is not None else None) + != (selection_key(original.selection) if original.selection is not None else None)): + raise TierRuntimeError("Runtime child database selection changed") + operation = ServingOperation(original.selection, original.tables, original.tier, + original.reranker_id, original.key, connection, lease, original._database, + original.borrowed) + lease.check_admission() + with _bound(operation): + yield operation + + def close(self) -> None: + self._reservation.close() + + +def current_operation(conn: sqlite3.Connection) -> ServingOperation | None: + operation = getattr(_local, "operation", None) + if operation is None: + return None + # A second connection must start its own explicit operation. This avoids + # silently treating a different database or snapshot as the parent's DB. + if operation._connection is not conn: + raise TierRuntimeError("Serving operation belongs to a different connection") + _local.lease.selection_key # validates PID, lifetime and owning thread + return operation + + +def _database_identity(conn: sqlite3.Connection) -> tuple[str, object]: + from truememory.maintenance import connection_database_path + path = connection_database_path(conn) + return ("memory", conn) if path is None else ("file", str(path)) + + +def current_runtime_operation() -> ServingOperation | None: + """Read this thread's admitted public tier/model policy without SQL.""" + operation = getattr(_local, "operation", None) + return None if operation is None else current_operation(operation._connection) + + +@contextmanager +def _bound(operation: ServingOperation, lease: serving.OperationLease | None = None) -> Iterator[None]: + previous = getattr(_local, "operation", None) + previous_lease = getattr(_local, "lease", None) + _local.operation = operation + _local.lease = operation._lease if lease is None else lease + try: + yield + finally: + _local.operation = previous + _local.lease = previous_lease + + +@contextmanager +def _connection_read( + conn: sqlite3.Connection | None, lock: object | None, deadline: float | None, + cancelled: threading.Event | None, *, transaction_owner: object | None = None, + allow_closed: bool = False, +) -> Iterator[None]: + if lock is None: + _check(deadline, cancelled) + yield + return + acquired = False + try: + try: + borrowed = conn is not None and conn.in_transaction + except sqlite3.ProgrammingError as exc: + if not allow_closed or str(exc) != "Cannot operate on a closed database.": + raise + borrowed = False + owner = lock if transaction_owner is None else transaction_owner + if (borrowed and not ( + isinstance(owner, ConnectionWriteLock) and owner.owns_other_transaction(conn))): + _check(deadline, cancelled) + acquired = lock.acquire(blocking=False) + if not acquired: + raise TierRuntimeError("Borrowed connection admission is busy") + else: + while not acquired: + _check(deadline, cancelled) + interval = _remaining(deadline) + if cancelled is not None: + interval = 0.05 if interval is None else min(interval, 0.05) + acquired = lock.acquire() if interval is None else lock.acquire(timeout=min(interval, threading.TIMEOUT_MAX)) + _check(deadline, cancelled) + yield + finally: + if acquired: + lock.release() + + +@contextmanager +def serving_operation( + conn: sqlite3.Connection, *, deadline: float | None = None, + cancelled: threading.Event | None = None, connection_lock: object | None = None, + reranker_id: str | None = None, +) -> Iterator[ServingOperation]: + serving._control(None, deadline, cancelled) + if reranker_id is not None: + _reranker_id(reranker_id) + current = current_operation(conn) + if current is not None: + if reranker_id is not None and reranker_id != current.reranker_id: + raise TierRuntimeError("Nested operation requested a different reranker") + with serving.operation_lease(current.key, deadline=deadline, cancelled=cancelled) as lease: + with _bound(current, lease): + yield current + return + # Engine serializes this short snapshot with its connection writer. The + # writer is released before gate admission or any model preparation. + with _connection_read(conn, connection_lock, deadline, cancelled): + selection = _read_selection(conn) + borrowed = conn.in_transaction + if selection is not None: + effective = selection.reranker_id if reranker_id is None else reranker_id + if not _ready(selection, effective): + if borrowed: + raise TierRuntimeError("Runtime reconciliation requires a clean connection") + apply_frozen_selection(selection, deadline=deadline, cancelled=cancelled, reranker_id=effective) + key = selection_key(selection, effective) + else: + vector, _ = _modules() + if vector._frozen_embedding_target is not None: + raise TierRuntimeError("Legacy database requires explicit coherent runtime selection") + if reranker_id is not None: + _, reranker = _modules() + if reranker._model is None or reranker._model_name != reranker_id: + if borrowed: + raise TierRuntimeError("Reranker loading requires a clean connection") + with serving.exclusive_activation(deadline=deadline, cancelled=cancelled): + reranker.get_reranker(model_name=reranker_id, deadline=deadline, cancelled=cancelled) + key = _legacy_key(reranker_id) + with serving.operation_lease(key, deadline=deadline, cancelled=cancelled, blocking=not borrowed) as lease: + # A competing process can commit selection while this thread waits. + # Refuse rather than loop through unbounded native reconciliation. + with _connection_read(conn, connection_lock, deadline, cancelled): + observed = _read_selection(conn) + vector, reranker = _modules() + tables = (selection.tables if selection is not None else + (vector._active_vec_table(conn), vector._active_sep_table(conn))) + database = _database_identity(conn) + if ((selection_key(observed) if observed is not None else None) + != (selection_key(selection) if selection is not None else None)): + raise TierRuntimeError("Database selection changed during runtime admission") + vector, reranker = _modules() + if selection is not None: + if not _ready(selection, effective): + raise TierRuntimeError("Process selection changed during runtime admission") + tables, tier, effective = selection.tables, selection.target.tier, effective + else: + if key != _legacy_key(reranker_id): + raise TierRuntimeError("Legacy runtime changed during admission") + tier = vector.resolve_tier() + effective = reranker.get_current_reranker_name() if reranker_id is None else reranker_id + operation = ServingOperation(selection, tables, tier, effective, key, conn, lease, + database, borrowed) + lease.check_admission() + with _bound(operation): + yield operation + + +def require_model_load_allowed() -> None: + operation = getattr(_local, "operation", None) + if operation is not None and operation._connection.in_transaction: + raise TierRuntimeError("Model loading requires a clean connection") + if operation is not None and operation.selection is not None: + raise TierRuntimeError("Selected model must be prepared before serving admission") + + +def database_operation(function: Callable[P, R]) -> Callable[P, R]: + """Wrap a direct connection-first search without changing its arguments.""" + @wraps(function) + def wrapped(conn: sqlite3.Connection, *args: object, **kwargs: object) -> R: + with serving_operation(conn): + return function(conn, *args, **kwargs) + return wrapped + + +def engine_operation(function: Callable[P, R]) -> Callable[P, R]: + """Open plainly, then pin the whole Engine call before writer admission.""" + @wraps(function) + def wrapped(engine: object, *args: object, **kwargs: object) -> R: + engine._open_connection_handle() + with serving_operation(engine.conn, connection_lock=engine._write_lock): + return function(engine, *args, **kwargs) + return wrapped diff --git a/truememory/tier_switch/serving.py b/truememory/tier_switch/serving.py index 769625b1..08e32876 100644 --- a/truememory/tier_switch/serving.py +++ b/truememory/tier_switch/serving.py @@ -108,6 +108,10 @@ def fork_child(self) -> ChildReservation: """Reserve one child now; close it if the task never enters join().""" return _issuer(self)._fork_child(self) + def check_admission(self) -> None: + """Recheck inherited controls before a caller's deferred body admission.""" + _issuer(self)._check_lease_admission(self) + class ChildReservation(_OpaqueToken): __slots__ = () @@ -173,12 +177,18 @@ def _check_process(self) -> None: raise ServingLeaseError("Serving state cannot be inherited by another process") @contextmanager - def _locked(self, control: _Control | None = None) -> Iterator[None]: + def _locked(self, control: _Control | None = None, *, blocking: bool = True) -> Iterator[None]: self._check_process() acquired = False try: while not acquired: - if control is None: + if not blocking: + if control is not None: + _check(control) + acquired = self._condition.acquire(blocking=False) + if not acquired: + raise ServingLeaseError("Serving admission is busy") + elif control is None: acquired = self._condition.acquire() else: _check(control) @@ -231,6 +241,12 @@ def _lease_key(self, lease: OperationLease) -> SelectionKey: with self._locked(): return self._owned_reader(lease).key + def _check_lease_admission(self, lease: OperationLease) -> None: + self._check_process() + control = self._owned_reader(lease).control + with self._locked(control): + _check(self._owned_reader(lease).control) + def _release_reader(self, lease: OperationLease) -> None: with self._locked(): reader = self._readers.get(lease) @@ -249,7 +265,10 @@ def _release_reader(self, lease: OperationLease) -> None: def operation_lease( self, selection_key: SelectionKey, *, timeout: float | None = None, deadline: float | None = None, cancelled: threading.Event | None = None, + blocking: bool = True, ) -> Iterator[OperationLease]: + if type(blocking) is not bool: + raise ServingLeaseError("Serving admission blocking must be boolean") key, control = _selection_key(selection_key), _control(timeout, deadline, cancelled) lease = self._token(OperationLease) thread = threading.current_thread() @@ -259,7 +278,7 @@ def operation_lease( control = _combine(parent.control, control) entered = False try: - with self._locked(control): + with self._locked(control, blocking=blocking): entered = True while True: _check(control) @@ -268,6 +287,8 @@ def operation_lease( if writer is thread or (writer is None and (parent is not None or not self._waiting_writers)): self._register_reader(lease, _Reader(thread, key, control, writer is thread or bool(parent and parent.writer_owned))) break + if not blocking: + raise ServingLeaseError("Serving admission is busy") self._condition.wait(_wait_seconds(control)) yield lease finally: @@ -392,10 +413,10 @@ def _after_fork() -> None: def operation_lease( selection_key: SelectionKey, *, timeout: float | None = None, deadline: float | None = None, - cancelled: threading.Event | None = None, + cancelled: threading.Event | None = None, blocking: bool = True, ) -> AbstractContextManager[OperationLease]: """Acquire the process-wide serving boundary for one logical operation.""" - return _DEFAULT_GATE.operation_lease(selection_key, timeout=timeout, deadline=deadline, cancelled=cancelled) + return _DEFAULT_GATE.operation_lease(selection_key, timeout=timeout, deadline=deadline, cancelled=cancelled, blocking=blocking) def exclusive_activation( diff --git a/truememory/vector_search.py b/truememory/vector_search.py index 70efcb03..aeadbe8c 100644 --- a/truememory/vector_search.py +++ b/truememory/vector_search.py @@ -146,6 +146,11 @@ def resolve_tier() -> str: should call this instead so that ``~/.truememory/config.json`` is honoured when the env var is absent. """ + import sys + runtime = sys.modules.get("truememory.tier_switch.runtime") + operation = runtime.current_runtime_operation() if runtime is not None else None + if operation is not None: + return operation.tier env = os.environ.get("TRUEMEMORY_EMBED_MODEL", "").strip().lower() if env: return env @@ -168,6 +173,7 @@ def resolve_tier() -> str: _embedding_dim: int = _MODEL_DIMS.get(EMBEDDING_MODEL, 256) _lock = threading.Lock() _model_generation = 0 +_frozen_embedding_target = None def _active_tier_group() -> str: @@ -177,6 +183,11 @@ def _active_tier_group() -> str: def _active_vec_table(conn: sqlite3.Connection) -> str: """Return the active vec_messages table name for the current tier.""" + import sys + runtime = sys.modules.get("truememory.tier_switch.runtime") + operation = runtime.current_operation(conn) if runtime is not None else None + if operation is not None: + return operation.tables[0] try: group = _active_tier_group() row = conn.execute( @@ -192,6 +203,11 @@ def _active_vec_table(conn: sqlite3.Connection) -> str: def _active_sep_table(conn: sqlite3.Connection) -> str: """Return the active vec_messages_sep table name for the current tier.""" + import sys + runtime = sys.modules.get("truememory.tier_switch.runtime") + operation = runtime.current_operation(conn) if runtime is not None else None + if operation is not None: + return operation.tables[1] try: group = _active_tier_group() row = conn.execute( @@ -210,10 +226,12 @@ def set_embedding_model(name: str) -> None: Accepts tier names ("base", "pro") or internal model names. """ - global EMBEDDING_MODEL, _model, _embedding_dim, _model_generation - with _lock: + global EMBEDDING_MODEL, _model, _embedding_dim, _model_generation, _frozen_embedding_target + from truememory.tier_switch.serving import exclusive_activation + with exclusive_activation(), _lock: _model_generation += 1 _model = None # Force reload + _frozen_embedding_target = None EMBEDDING_MODEL = _resolve_model_name(name) _embedding_dim = get_embedding_dim(EMBEDDING_MODEL) @@ -230,7 +248,8 @@ def get_embedding_dim(name: str | None = None) -> int: def unload_model() -> None: """Release the embedding model from memory.""" global _model, _model_generation - with _lock: + from truememory.tier_switch.serving import exclusive_activation + with exclusive_activation(), _lock: _model_generation += 1 _model = None @@ -245,88 +264,179 @@ def get_model(): global _model, _embedding_dim if _model is not None: return _model # Fast path, no lock needed - with _lock: - if _model is not None: - return _model # Another thread loaded it - - from truememory.model_client import use_model_server, get_embedding_proxy - if use_model_server(): - _model = get_embedding_proxy(tier=EMBEDDING_MODEL) - return _model - - from truememory.mps_utils import ensure_mps_memory_budget, resolve_device - - resolved = EMBEDDING_MODEL - if resolved == "model2vec": - from model2vec import StaticModel - _model = StaticModel.from_pretrained("minishlab/potion-base-8M", force_download=False) - _embedding_dim = 256 - elif resolved == "minilm": - device = resolve_device(None) - ensure_mps_memory_budget(device) - from sentence_transformers import SentenceTransformer - _model = SentenceTransformer("all-MiniLM-L6-v2", device=device) - _embedding_dim = 384 - elif resolved == "bge-small": - device = resolve_device(None) - ensure_mps_memory_budget(device) - from sentence_transformers import SentenceTransformer - _model = SentenceTransformer("BAAI/bge-small-en-v1.5", device=device) - _embedding_dim = 384 - elif resolved == "qwen3_256": - device = resolve_device(None) - ensure_mps_memory_budget(device) - from sentence_transformers import SentenceTransformer - import sys as _sys - _mkwargs = {} - if _sys.platform == "darwin": - _mkwargs["attn_implementation"] = "eager" - _model = SentenceTransformer( - "Qwen/Qwen3-Embedding-0.6B", - truncate_dim=256, - model_kwargs=_mkwargs or None, - device=device, - ) - _embedding_dim = 256 - elif resolved not in _MODEL_DIMS: - # Custom tier: load arbitrary SentenceTransformer model. - # Requires TRUEMEMORY_CUSTOM_ALLOW_DOWNLOAD=1 (enforced by - # tier_config.resolve_custom_tier at config time). - if os.environ.get("TRUEMEMORY_CUSTOM_ALLOW_DOWNLOAD", "").strip() != "1": - logger.warning( - "Custom model %r requested without " - "TRUEMEMORY_CUSTOM_ALLOW_DOWNLOAD=1 -- " - "falling back to model2vec.", - resolved, - ) + from truememory.tier_switch.runtime import require_model_load_allowed + require_model_load_allowed() + from contextlib import nullcontext + from truememory.tier_switch.runtime import current_runtime_operation + from truememory.tier_switch.serving import exclusive_activation + with exclusive_activation() if current_runtime_operation() is None else nullcontext(): + with _lock: + if _model is not None: + return _model # Another thread loaded it + + if _frozen_embedding_target is not None: + _model = _load_frozen_embedding_target(_frozen_embedding_target) + _embedding_dim = _frozen_embedding_target.dimension + return _model + + from truememory.model_client import use_model_server, get_embedding_proxy + if use_model_server(): + _model = get_embedding_proxy(tier=EMBEDDING_MODEL) + return _model + + from truememory.mps_utils import ensure_mps_memory_budget, resolve_device + + resolved = EMBEDDING_MODEL + if resolved == "model2vec": from model2vec import StaticModel - _model = StaticModel.from_pretrained( - "minishlab/potion-base-8M", force_download=False - ) + _model = StaticModel.from_pretrained("minishlab/potion-base-8M", force_download=False) _embedding_dim = 256 - else: - from truememory.tier_config import resolve_custom_tier - cfg = resolve_custom_tier() - custom_dim = cfg["embed_dim"] + elif resolved == "minilm": + device = resolve_device(None) + ensure_mps_memory_budget(device) + from sentence_transformers import SentenceTransformer + _model = SentenceTransformer("all-MiniLM-L6-v2", device=device) + _embedding_dim = 384 + elif resolved == "bge-small": + device = resolve_device(None) + ensure_mps_memory_budget(device) + from sentence_transformers import SentenceTransformer + _model = SentenceTransformer("BAAI/bge-small-en-v1.5", device=device) + _embedding_dim = 384 + elif resolved == "qwen3_256": device = resolve_device(None) ensure_mps_memory_budget(device) from sentence_transformers import SentenceTransformer + import sys as _sys + _mkwargs = {} + if _sys.platform == "darwin": + _mkwargs["attn_implementation"] = "eager" _model = SentenceTransformer( - resolved, truncate_dim=custom_dim, - trust_remote_code=False, + "Qwen/Qwen3-Embedding-0.6B", + truncate_dim=256, + model_kwargs=_mkwargs or None, device=device, ) - _embedding_dim = custom_dim - else: - from model2vec import StaticModel - _model = StaticModel.from_pretrained("minishlab/potion-base-8M", force_download=False) - _embedding_dim = 256 + _embedding_dim = 256 + elif resolved not in _MODEL_DIMS: + # Custom tier: load arbitrary SentenceTransformer model. + # Requires TRUEMEMORY_CUSTOM_ALLOW_DOWNLOAD=1 (enforced by + # tier_config.resolve_custom_tier at config time). + if os.environ.get("TRUEMEMORY_CUSTOM_ALLOW_DOWNLOAD", "").strip() != "1": + logger.warning( + "Custom model %r requested without " + "TRUEMEMORY_CUSTOM_ALLOW_DOWNLOAD=1 -- " + "falling back to model2vec.", + resolved, + ) + from model2vec import StaticModel + _model = StaticModel.from_pretrained( + "minishlab/potion-base-8M", force_download=False + ) + _embedding_dim = 256 + else: + from truememory.tier_config import resolve_custom_tier + cfg = resolve_custom_tier() + custom_dim = cfg["embed_dim"] + device = resolve_device(None) + ensure_mps_memory_budget(device) + from sentence_transformers import SentenceTransformer + _model = SentenceTransformer( + resolved, truncate_dim=custom_dim, + trust_remote_code=False, + device=device, + ) + _embedding_dim = custom_dim + else: + from model2vec import StaticModel + _model = StaticModel.from_pretrained("minishlab/potion-base-8M", force_download=False) + _embedding_dim = 256 return _model _prepared_target_slot = threading.Lock() +def _load_frozen_embedding_target(target: EmbeddingTarget, *, timeout: float | None = None) -> object: + from truememory.embedding_target import TARGET_PROBE_TEXT, build_target_model, check_target_vectors + from truememory.model_client import EmbeddingProxy, PreparedEmbeddingProxy, use_model_server + from truememory.mps_utils import _own_model, resolve_device + + expires = _target_deadline(timeout) + target.check_configuration() + if use_model_server(): + if target.tier != "custom": + from truememory.model_client import CertifiedEmbeddingProxy + proxy = CertifiedEmbeddingProxy(target) + PreparedEmbeddingProxy(target).prepare(timeout=_target_remaining(expires)) + return proxy + + class FrozenProxy(EmbeddingProxy): + def __init__(self) -> None: + super().__init__(tier=target.model_id) + self._prepared = PreparedEmbeddingProxy(target) + + def encode(self, texts: str | list[str], timeout: float | None = None, **kwargs: object) -> np.ndarray: + texts = [texts] if isinstance(texts, str) else list(texts) + values = self._prepared.encode(texts, timeout=timeout, batch_size=kwargs.get("batch_size", 32)) + check_target_vectors(target, values, len(texts)) + return values + + proxy = FrozenProxy() + proxy._prepared.prepare(timeout=_target_remaining(expires)) + return proxy + if not _prepared_target_slot.acquire(blocking=False): + raise RuntimeError("Another prepared embedding target lease is active") + try: + encoder = build_target_model(target, resolve_device(None), check=lambda: _target_remaining(expires)) + with _own_model(encoder, remaining=lambda: _target_remaining(expires)): + _target_remaining(expires) + values = encoder.encode([TARGET_PROBE_TEXT], batch_size=1, show_progress_bar=False) + _target_remaining(expires) + check_target_vectors(target, np.asarray(values, dtype=np.float32), 1) + return encoder + finally: + _prepared_target_slot.release() + + +def apply_frozen_embedding_target(target: EmbeddingTarget, *, timeout: float | None = None) -> None: + """Fill the existing active slot from exact prepared-target arguments.""" + global _model, _embedding_dim, EMBEDDING_MODEL, _model_generation, _frozen_embedding_target + from truememory.embedding_target import EmbeddingTarget + from truememory.tier_switch.serving import exclusive_activation + if type(target) is not EmbeddingTarget: + raise TypeError("Expected an immutable EmbeddingTarget") + expires = _target_deadline(timeout) + target.check_configuration() + with exclusive_activation(deadline=expires): + with _target_state_lock(expires): + if (_frozen_embedding_target is not None and _frozen_embedding_target.identity == target.identity + and _model is not None): + _frozen_embedding_target = target + return + _model = None + _model_generation += 1 + _frozen_embedding_target = target + EMBEDDING_MODEL, _embedding_dim = target.model_id, target.dimension + encoder = _load_frozen_embedding_target(target, timeout=_target_remaining(expires)) + with _target_state_lock(expires): + _model = encoder + + +@contextmanager +def _target_state_lock(expires: float | None) -> Iterator[None]: + remaining = _target_remaining(expires) + if remaining is None: + _lock.acquire() + elif not _lock.acquire(timeout=min(remaining, threading.TIMEOUT_MAX)): + raise TimeoutError("Prepared embedding deadline exceeded") + try: + _target_remaining(expires) + yield + finally: + _lock.release() + + class _PreparedEmbeddingLease: """An encode-only handle; closing drains work before releasing residency.""" @@ -887,6 +997,15 @@ def init_vec_table( conn: An open SQLite connection (from :func:`truememory.storage.create_db`). tier_group: Optional explicit tier group name ("edge" or "basepro"). """ + import sys + runtime = sys.modules.get("truememory.tier_switch.runtime") + operation = runtime.current_operation(conn) if runtime is not None else None + if operation is not None and operation.selection is not None: + from truememory.tier_switch.activation import _guard_pair + if tier_group is not None and tier_group != operation.selection.target.tier_group: + raise ValueError("Selected initialization requires the captured table group") + _guard_pair(conn, operation.selection.target) + return import sqlite_vec conn.enable_load_extension(True) @@ -1241,6 +1360,10 @@ def build_vectors( # Vector search # --------------------------------------------------------------------------- +from truememory.tier_switch.runtime import database_operation + + +@database_operation def search_vector( conn: sqlite3.Connection, query: str, @@ -1333,6 +1456,7 @@ def search_vector( return results[:limit] +@database_operation def search_vector_raw( conn: sqlite3.Connection, query: str, @@ -1439,6 +1563,7 @@ def _build_sep_text(sender: str, recipient: str, timestamp: str, content: str) - ) +@database_operation def embed_single(conn: sqlite3.Connection, message_id: int, content: str) -> None: """ Embed a single message and insert into ``vec_messages`` and ``vec_messages_sep``. @@ -1485,6 +1610,7 @@ def embed_single(conn: sqlite3.Connection, message_id: int, content: str) -> Non _write_embedder_metadata_no_commit(conn) +@database_operation def search_vector_separation( conn: sqlite3.Connection, query: str, From a8639b535210f143070fdc7f4268cbaa39f3e82e Mon Sep 17 00:00:00 2001 From: buildingjoshbetter Date: Fri, 9 Oct 2026 12:21:10 -0500 Subject: [PATCH 4/5] fix: Decode tier test sources as UTF-8 --- tests/test-tier-runtime-bridge-795.py | 14 +++++++------- tests/test-tier-writer-fence-795.py | 4 ++-- 2 files changed, 9 insertions(+), 9 deletions(-) diff --git a/tests/test-tier-runtime-bridge-795.py b/tests/test-tier-runtime-bridge-795.py index c87c402a..413bdd1a 100644 --- a/tests/test-tier-runtime-bridge-795.py +++ b/tests/test-tier-runtime-bridge-795.py @@ -89,7 +89,7 @@ def safe_import(name, globals=None, locals=None, fromlist=(), level=0): sys.modules[self.api.__name__] = self.api self.api.__dict__["__builtins__"] = dict(vars(builtins), __import__=safe_import) path = ROOT / "truememory/tier_switch/runtime.py" - exec(compile(path.read_text(), str(path), "exec"), self.api.__dict__) + exec(compile(path.read_text(encoding="utf-8"), str(path), "exec"), self.api.__dict__) self.modules["tier_switch.runtime"] = self.api self.api._database_identity = lambda conn: ("memory", conn) @@ -330,12 +330,12 @@ def test_nonblocking_nested_read_inherits_cancel_and_identity(self) -> None: self.fail("Inherited cancellation precedes admission") def test_engine_and_hybrid_outer_boundaries_are_present(self) -> None: - engine = ast.parse((ROOT / "truememory/engine.py").read_text()) + engine = ast.parse((ROOT / "truememory/engine.py").read_text(encoding="utf-8")) klass = next(node for node in engine.body if isinstance(node, ast.ClassDef) and node.name == "TrueMemoryEngine") for name in ("search", "search_agentic", "search_vectors_raw", "add", "update", "delete", "delete_all"): node = next(node for node in klass.body if isinstance(node, ast.FunctionDef) and node.name == name) self.assertIn("engine_operation", [ast.unparse(value) for value in node.decorator_list]) - hybrid = ast.parse((ROOT / "truememory/hybrid.py").read_text()) + hybrid = ast.parse((ROOT / "truememory/hybrid.py").read_text(encoding="utf-8")) function = next(node for node in hybrid.body if isinstance(node, ast.FunctionDef) and node.name == "search_hybrid") self.assertIn("database_operation", [ast.unparse(value) for value in function.decorator_list]) @@ -407,7 +407,7 @@ def test_nested_reservation_inherits_stricter_nested_control(self) -> None: def test_actual_selected_engine_init_reconciles_before_validation_and_keeps_schedule(self) -> None: selected = self.select() - tree = ast.parse((ROOT / "truememory/engine.py").read_text()) + tree = ast.parse((ROOT / "truememory/engine.py").read_text(encoding="utf-8")) source = next(node for node in tree.body if isinstance(node, ast.ClassDef) and node.name == "TrueMemoryEngine") source.body = [node for node in source.body if isinstance(node, ast.FunctionDef) and node.name in { "_ensure_connection", "_initialize_connection", "_open_connection_handle", "close", @@ -464,7 +464,7 @@ def imports(name, globals=None, locals=None, fromlist=(), level=0): def test_ready_legacy_borrowed_connection_does_not_reinitialize(self) -> None: - tree = ast.parse((ROOT / "truememory/engine.py").read_text()) + tree = ast.parse((ROOT / "truememory/engine.py").read_text(encoding="utf-8")) source = next(node for node in tree.body if isinstance(node, ast.ClassDef) and node.name == "TrueMemoryEngine") source.body = [node for node in source.body if isinstance(node, ast.FunctionDef) and node.name == "_initialize_connection"] namespace = {"__builtins__": dict(vars(builtins), __import__=self.safe_import)} @@ -484,7 +484,7 @@ def test_ready_legacy_borrowed_connection_does_not_reinitialize(self) -> None: engine._initialize_connection(_suppress_maintenance=True) def test_fresh_or_closed_handle_clears_legacy_ready_before_initialization(self) -> None: - tree = ast.parse((ROOT / "truememory/engine.py").read_text()) + tree = ast.parse((ROOT / "truememory/engine.py").read_text(encoding="utf-8")) source = next(node for node in tree.body if isinstance(node, ast.ClassDef) and node.name == "TrueMemoryEngine") source.body = [node for node in source.body if isinstance(node, ast.FunctionDef) and node.name == "_open_connection_handle"] replacement = sqlite3.connect(":memory:") @@ -595,7 +595,7 @@ def holder(): self.assertEqual(conn.execute("SELECT value FROM synthetic").fetchall(), []) def test_connection_probe_error_preserves_handle_under_both_locks(self) -> None: - tree = ast.parse((ROOT / "truememory/engine.py").read_text()) + tree = ast.parse((ROOT / "truememory/engine.py").read_text(encoding="utf-8")) source = next(node for node in tree.body if isinstance(node, ast.ClassDef) and node.name == "TrueMemoryEngine") source.body = [node for node in source.body if isinstance(node, ast.FunctionDef) and node.name == "_open_connection_handle"] diff --git a/tests/test-tier-writer-fence-795.py b/tests/test-tier-writer-fence-795.py index 67334e41..5ba2cbfb 100644 --- a/tests/test-tier-writer-fence-795.py +++ b/tests/test-tier-writer-fence-795.py @@ -103,7 +103,7 @@ def safe_import(name, globals=None, locals=None, fromlist=(), level=0): sys.modules[module.__name__] = module module.__dict__["__builtins__"] = dict(vars(builtins), __import__=safe_import) path = ROOT / "truememory" / (name.replace(".", "/") + ".py") - exec(compile(path.read_text(), str(path), "exec"), module.__dict__) + exec(compile(path.read_text(encoding="utf-8"), str(path), "exec"), module.__dict__) self.modules[name] = module self.api = self.modules["tier_switch.writer"] self.activation = self.modules["tier_switch.activation"] @@ -179,7 +179,7 @@ def contents(self, conn=None): return conn.execute("SELECT id,content FROM messages ORDER BY id").fetchall() def engine(self, *, vectors=True): - tree = ast.parse((ROOT / "truememory/engine.py").read_text()) + tree = ast.parse((ROOT / "truememory/engine.py").read_text(encoding="utf-8")) original = next(n for n in tree.body if isinstance(n, ast.ClassDef) and n.name == "TrueMemoryEngine") methods = [n for n in original.body if isinstance(n, ast.FunctionDef) and n.name in {"add", "update", "delete", "delete_all"}] cls = ast.ClassDef(name="SyntheticEngine", bases=[], keywords=[], body=methods, decorator_list=[]) From 48217a73f10ad7fa80029ad800660d76ef912de7 Mon Sep 17 00:00:00 2001 From: buildingjoshbetter Date: Fri, 9 Oct 2026 12:43:27 -0500 Subject: [PATCH 5/5] fix: Wait for maintenance teardown in warm search test --- tests/test_issue_460_warm_search.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/tests/test_issue_460_warm_search.py b/tests/test_issue_460_warm_search.py index 3380734c..ff204a16 100644 --- a/tests/test_issue_460_warm_search.py +++ b/tests/test_issue_460_warm_search.py @@ -89,6 +89,8 @@ def counting_encode(texts, **kw): m1._engine.close() m2._engine.close() + from truememory.maintenance import get_coordinator + assert get_coordinator(db_path).wait(10), "Synthetic maintenance teardown did not finish" def test_issue_460_metadata_survives_multiple_adds(self): """Metadata must persist across multiple add() calls."""