diff --git a/docs/resource-budgets.md b/docs/resource-budgets.md index e6dd556d..5ce44cec 100644 --- a/docs/resource-budgets.md +++ b/docs/resource-budgets.md @@ -2130,3 +2130,185 @@ 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. + +### 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. + + +### 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-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-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 bc39586a..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": @@ -34,16 +36,17 @@ 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") 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..413bdd1a --- /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(encoding="utf-8"), 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(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(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]) + + 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(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", + }] + 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(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)} + 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(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:") + 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(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"] + 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-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/tests/test-tier-writer-fence-795.py b/tests/test-tier-writer-fence-795.py new file mode 100644 index 00000000..5ba2cbfb --- /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(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"] + 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(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=[]) + 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, 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) + 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/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.""" 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 d2dd2907..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, @@ -712,6 +772,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 +802,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 +835,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)), @@ -881,18 +949,22 @@ 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. 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 + @engine_operation def delete_all(self, user_id: str | None = None) -> bool: """Delete all memories, optionally filtered by user. @@ -915,8 +987,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 +1065,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 +1134,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 @@ -1122,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. @@ -1137,6 +1214,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 +1249,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 +1263,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: @@ -1275,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 @@ -1687,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. @@ -1703,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. @@ -2100,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, @@ -2490,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). @@ -2612,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/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 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/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/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/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..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) @@ -1022,11 +1141,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 +1179,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 +1196,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 +1204,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 +1221,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 +1236,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 +1299,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 +1309,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) @@ -1221,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, @@ -1313,6 +1456,7 @@ def search_vector( return results[:limit] +@database_operation def search_vector_raw( conn: sqlite3.Connection, query: str, @@ -1419,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``. @@ -1433,6 +1578,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 +1600,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: @@ -1463,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,