From 4b2e0bfd5c91d92244f5198542896435a4bc5f82 Mon Sep 17 00:00:00 2001 From: buildingjoshbetter Date: Fri, 9 Oct 2026 15:21:26 -0500 Subject: [PATCH 01/11] feat: Add guarded tier policy and config projection primitives --- tests/test-tier-activation-policy-795.py | 458 +++++++++++++++++++++++ truememory/mcp_server.py | 61 +-- truememory/tier_switch/activation.py | 283 +++++++++++++- truememory/tier_switch/job.py | 65 ++++ truememory/tier_switch/projection.py | 215 +++++++++++ truememory/tier_switch/source.py | 68 ++++ 6 files changed, 1082 insertions(+), 68 deletions(-) create mode 100644 tests/test-tier-activation-policy-795.py create mode 100644 truememory/tier_switch/projection.py diff --git a/tests/test-tier-activation-policy-795.py b/tests/test-tier-activation-policy-795.py new file mode 100644 index 00000000..8b288ffd --- /dev/null +++ b/tests/test-tier-activation-policy-795.py @@ -0,0 +1,458 @@ +"""Policy/projection faults using synthetic config and in-memory SQLite only.""" +from __future__ import annotations + +import builtins +import dataclasses +import json +import runpy +import sys +import tempfile +import threading +import types +import unittest +from pathlib import Path +from unittest.mock import patch + +BASE = runpy.run_path(str(Path(__file__).with_name("test-tier-activation-journal-795.py"))) +ROOT = Path(__file__).resolve().parents[1] +RERANKER = "Alibaba-NLP/gte-reranker-modernbert-base" +LEGACY = ("vec_messages", "vec_messages_sep") + + +class TestActivationPolicy(unittest.TestCase): + def setUp(self) -> None: + self.fixture = BASE["TestActivationJournal"]() + self.fixture.setUp() + self.addCleanup(self.fixture.doCleanups) + self.api = self.fixture.api + self.conn = self.fixture.conn + self.jobs = self.fixture.job_api + self.source = self.fixture.source + self.Target = self.fixture.modules["embedding_target"].EmbeddingTarget + self.base = self.Target("base", "qwen3_256", 256, "basepro") + self.pro = self.Target("pro", "qwen3_256", 256, "basepro") + self.guard = self.api.ConfigGuard(True, "base", "a" * 32) + for table in (*LEGACY, *self.base.tables): + self.conn.execute(f'CREATE TABLE "{table}"(embedding BLOB)') + self.conn.execute("UPDATE metadata SET value='qwen3_256' WHERE key='embed_model'") + self.conn.execute("UPDATE metadata SET value='256' WHERE key='embed_dim'") + self.conn.commit() + schema = self.api._schema_sql + + def fake_native_schema(conn, table): + actual = schema(conn, table) + if table in (*LEGACY, *self.base.tables): + return f'CREATE VIRTUAL TABLE "{table}" USING vec0(embedding float[256] distance_metric=cosine)' + return actual + + self.schema_patch = patch.object(self.api, "_schema_sql", fake_native_schema) + self.schema_patch.start() + self.addCleanup(self.schema_patch.stop) + + projection = types.ModuleType("synthetic_policy_projection") + sys.modules[projection.__name__] = projection + self.addCleanup(sys.modules.pop, projection.__name__, None) + + def safe_import(name, globals=None, locals=None, fromlist=(), level=0): + if name == "truememory.tier_switch.activation": + return self.api + if name.startswith("truememory") or name in {"torch", "numpy", "sqlite_vec", "sentence_transformers"}: + raise AssertionError("Projection must stay model/native free") + return builtins.__import__(name, globals, locals, fromlist, level) + + projection.__dict__["__builtins__"] = dict(vars(builtins), __import__=safe_import) + path = ROOT / "truememory/tier_switch/projection.py" + exec(compile(path.read_text(encoding="utf-8"), str(path), "exec"), projection.__dict__) + self.projection = projection + + def stage(self): + return self.api.stage_activation_intent( + self.conn, self.fixture.job, expected_generation=None, reranker_id="synthetic/reranker", + expected_config=self.guard, status_id=17, + ) + + def publish(self, intent=None, **kwargs): + arguments = dict(expected_selection_generation=None, expected_intent_id=intent.intent_id if intent else None, + target=self.pro, reranker_id=RERANKER, expected_config=self.guard, legacy_pair=LEGACY) + arguments.update(kwargs) + return self.api.commit_config_only_transition(self.conn, **arguments) + + def terminal_marker(self): + self.jobs.end_selected_tier_job(self.conn, self.fixture.job, outcome="cancelled") + + def test_legacy_policy_changes_only_journal_and_exact_marker(self): + intent = self.stage() + before = self.conn.execute("SELECT key,value FROM metadata WHERE key!='tier_activation_v1' ORDER BY key").fetchall() + statements = [] + self.conn.set_trace_callback(statements.append) + result = self.publish(intent) + self.conn.set_trace_callback(None) + self.assertIs(type(result), self.api.LegacyTierPolicy) + self.assertEqual(result.tables, LEGACY) + self.assertEqual(result.target, self.pro) + self.assertFalse(result.config_acknowledged) + self.assertIsNone(self.api.read_activation_state(self.conn).selection) + self.assertEqual(self.api.read_activation_state(self.conn).legacy_policy, result) + self.assertEqual(before, self.conn.execute("SELECT key,value FROM metadata WHERE key!='tier_activation_v1' ORDER BY key").fetchall()) + self.assertEqual(self.jobs.read_selected_job_marker(self.conn).state, "cancelled") + self.assertFalse(any("FROM MAIN.MESSAGES" in sql.upper() or "FROM MAIN.VEC_MESSAGES" in sql.upper() for sql in statements)) + self.assertFalse(any("DELETE" in sql.upper() for sql in statements)) + + def test_pending_policy_ack_cas_and_roundtrip(self): + self.terminal_marker() + result = self.publish() + with self.assertRaises(self.api.TierActivationError): + self.api.acknowledge_policy_config(self.conn, generation="f" * 32) + with self.assertRaises(self.api.TierActivationError): + self.publish(expected_intent_id=result.intent_id, target=self.base, + expected_config=self.api.ConfigGuard(True, "pro", result.generation)) + acknowledged = self.api.acknowledge_policy_config(self.conn, generation=result.generation) + self.assertTrue(acknowledged.config_acknowledged) + self.assertEqual(acknowledged, self.api.acknowledge_policy_config(self.conn, generation=result.generation)) + second = self.publish(expected_intent_id=result.intent_id, target=self.base, + expected_config=self.api.ConfigGuard(True, "pro", result.generation)) + self.assertNotEqual(result.generation, second.generation) + self.assertEqual(second.target, self.base) + self.assertIsNone(self.api.read_activation_state(self.conn).selection) + + def test_v2_rebuild_retains_acknowledged_legacy_policy_and_old_v1_bytes(self): + old = self.fixture.stage() + raw = self.api._dump(old) + self.assertEqual(json.loads(raw)["version"], 1) + self.assertNotIn("expected_config", json.loads(raw)) + self.assertEqual(self.api._dump(self.api._parse(raw, self.api.ActivationIntent)), raw) + self.jobs.end_selected_tier_job(self.conn, self.fixture.job, outcome="cancelled") + policy = self.publish(old) + policy = self.api.acknowledge_policy_config(self.conn, generation=policy.generation) + job = self.jobs.select_tier_job(self.conn, self.fixture.target, previous_job_id=self.fixture.job.job_id) + guard = self.api.ConfigGuard(True, "pro", policy.generation) + staged = self.api.stage_activation_intent(self.conn, job, expected_generation=None, + reranker_id="synthetic/reranker", expected_config=guard, status_id=23) + self.assertEqual(staged.previous_policy, policy) + self.assertEqual(self.api.read_activation_state(self.conn).legacy_policy, policy) + self.assertEqual(json.loads(self.api._dump(staged))["kind"], "rebuild") + self.assertEqual(self.api._parse(self.api._dump(staged), self.api.ActivationIntent), staged) + self.assertEqual(self.api.stage_activation_intent(self.conn, job, expected_generation=None, + reranker_id="synthetic/reranker", expected_config=guard, status_id=23), staged) + with self.assertRaises(self.api.TierActivationError): + self.api.stage_activation_intent(self.conn, job, expected_generation=None, reranker_id="synthetic/reranker") + + def test_exact_superseded_job_and_rollback(self): + intent = self.stage() + original = self.jobs.read_selected_job_marker(self.conn) + self.conn.execute(f"UPDATE {self.jobs._SELECTION_TABLE} SET job_id=?", ("b" * 32,)) + self.conn.commit() + with self.assertRaises(self.jobs.TierJobSelectionError): + self.publish(intent) + self.assertEqual(self.jobs.read_selected_job_marker(self.conn).job_id, "b" * 32) + self.conn.execute(f"UPDATE {self.jobs._SELECTION_TABLE} SET job_id=?", (original.job_id,)) + self.conn.commit() + self.conn.fault_contains = "INSERT INTO main.metadata" + self.conn.fault = MemoryError("synthetic journal write") + self.conn.armed = True + with self.assertRaises(MemoryError): + self.publish(intent) + self.conn.armed = False + self.assertEqual(self.jobs.read_selected_job_marker(self.conn), original) + self.assertEqual(self.api.read_activation_state(self.conn).intent, intent) + + def test_commit_after_success_is_pending_readback_not_false_rollback(self): + intent = self.stage() + self.conn.armed = True + self.conn.commit_mode = "after" + self.conn.fault = TimeoutError("synthetic post-commit") + before = self.conn.rollbacks + with self.assertRaises(TimeoutError): + self.publish(intent) + self.conn.armed = False + self.assertEqual(self.conn.rollbacks, before) + state = self.api.read_activation_state(self.conn) + self.assertIsNotNone(state.legacy_policy) + self.assertFalse(state.legacy_policy.config_acknowledged) + self.assertEqual(self.jobs.read_selected_job_marker(self.conn).state, "cancelled") + + def test_certified_policy_keeps_certificate_without_recertifying_vectors(self): + self.terminal_marker() + self.fixture.target = self.base + self.fixture.job = self.jobs.select_tier_job(self.conn, self.base, previous_job_id=self.fixture.job.job_id) + self.fixture.reranker = RERANKER + intent, plan = self.fixture.build(ids=()) + selected = self.fixture.commit(intent, plan) + selected = self.api.acknowledge_config(self.conn, generation=selected.generation) + statements = [] + self.conn.set_trace_callback(statements.append) + result = self.publish(expected_selection_generation=selected.generation, expected_intent_id=selected.intent_id, + expected_config=self.api.ConfigGuard(True, "base", selected.generation), legacy_pair=None) + self.conn.set_trace_callback(None) + for field in dataclasses.fields(selected): + if field.name not in {"generation", "intent_id", "target", "config_acknowledged"}: + self.assertEqual(getattr(result, field.name), getattr(selected, field.name)) + self.assertNotEqual(result.generation, selected.generation) + self.assertEqual(result.target, self.pro) + self.assertIsNone(self.api.read_activation_state(self.conn).legacy_policy) + self.assertFalse(any("FROM MAIN.MESSAGES" in sql.upper() for sql in statements)) + self.assertTrue(self.api.acknowledge_config(self.conn, generation=result.generation).config_acknowledged) + + def test_strict_invalid_policy_shapes_and_no_new_certificate(self): + self.terminal_marker() + bad = [dict(target=self.fixture.target), dict(reranker_id="synthetic/other"), + dict(legacy_pair=("vec_messages_custom", "vec_messages_sep_custom")), + dict(expected_config=None), dict(target=self.base), dict(expected_intent_id="b" * 32), + dict(expected_selection_generation="b" * 32)] + for fields in bad: + with self.subTest(fields=list(fields)), self.assertRaises(self.api.TierActivationError): + self.publish(**fields) + self.assertIsNone(self.api.read_activation_state(self.conn).intent) + self.assertIsNone(self.api.read_activation_state(self.conn).selection) + + def test_legacy_registry_metadata_and_native_schema_rejection(self): + self.terminal_marker() + self.conn.execute("UPDATE metadata SET value='model2vec' WHERE key='embed_model'") + self.conn.commit() + with self.assertRaises(self.api.TierActivationError): + self.publish() + self.conn.execute("UPDATE metadata SET value='qwen3_256' WHERE key='embed_model'") + self.conn.execute("INSERT INTO vector_cache_registry(tier_group,vec_table,sep_table,model_name) VALUES('basepro','wrong','wrong','qwen3_256')") + self.conn.commit() + with self.assertRaises(self.api.TierActivationError): + self.publish() + self.schema_patch.stop() + with self.assertRaises(self.api.TierActivationError): + self.publish() + + def test_strict_nested_policy_and_nul_tail_byte_bound(self): + self.terminal_marker() + result = self.publish() + raw = self.conn.execute("SELECT value FROM metadata WHERE key='tier_activation_v1'").fetchone()[0] + value = json.loads(raw) + for field, invalid in (("result_generation", True), ("state", []), ("tables", [[], []]), + ("expected_config", {"tier": "base"}), ("unknown", 1)): + altered = dict(value, **{field: invalid}) + with self.subTest(field=field), self.assertRaises(self.api.TierActivationError): + self.api._parse(json.dumps(altered), self.api.ActivationIntent) + self.conn.execute("UPDATE metadata SET value=? WHERE key='tier_activation_v1'", (raw + "\x00" + "x" * 17000,)) + self.conn.commit() + before = self.conn.record_payload_reads + with self.assertRaises(self.api.TierActivationError): + self.api.read_activation_state(self.conn) + self.assertEqual(self.conn.record_payload_reads, before) + self.assertEqual(result.target, self.pro) + + def test_progress_is_bounded_manifest_read_and_borrows_snapshot(self): + intent, plan = self.fixture.build() + statements = [] + self.conn.set_trace_callback(statements.append) + self.conn.execute("BEGIN") + marker = self.jobs.read_selected_job_marker(self.conn) + progress = self.source.read_tier_progress(self.conn, model_id=self.fixture.target.model_id, + dimension=2, targets=self.fixture.target.tables) + self.assertTrue(self.conn.in_transaction) + self.conn.rollback() + self.conn.set_trace_callback(None) + self.assertEqual((progress.processed, progress.total, progress.outputs, progress.cursor, progress.complete), (3, 3, 3, 5, True)) + self.assertEqual((progress.tracker, progress.source_epoch, progress.source_schema_signature), + (marker.tracker, marker.source_epoch, marker.source_schema_signature)) + self.assertEqual(marker.job_id, intent.job_id) + self.assertEqual(progress.generation, plan.manifest.generation) + self.assertFalse(any("FROM MAIN.MESSAGES" in sql.upper() or "FROM MAIN.VEC_MESSAGES" in sql.upper() for sql in statements)) + + def test_progress_missing_half_receipt_and_read_cleanup_failure(self): + _, plan = self.fixture.build() + self.conn.execute("DELETE FROM metadata WHERE key=?", (self.source._receipt_key(plan.manifest.targets),)) + self.conn.commit() + with self.assertRaises(self.source.TierSourceUntrusted): + self.source.read_tier_progress(self.conn, model_id=self.fixture.target.model_id, dimension=2, targets=plan.manifest.targets) + self.assertFalse(self.conn.in_transaction) + self.conn.armed = True + self.conn.rollback_mode = "before" + self.conn.fault = MemoryError("synthetic cleanup") + before = self.conn.rollbacks + with self.assertRaises(MemoryError): + self.source.read_tier_progress(self.conn, model_id=self.fixture.target.model_id, dimension=2, targets=plan.manifest.targets) + self.assertEqual(self.conn.rollbacks, before + 1) + self.assertTrue(self.conn.in_transaction) + self.conn.armed = False + self.conn.rollback() + + def test_config_guard_seed_mirror_preserves_fresh_independent_fields(self): + self.terminal_marker() + with tempfile.TemporaryDirectory(prefix="synthetic-policy-") as directory: + config = Path(directory) / "config.json" + config.write_text(json.dumps({"tier": "base", "synthetic_setting": 1})) + guard = self.projection.capture_config_guard(config_path=config) + self.assertEqual(guard.tier, "base") + self.assertEqual(len(guard.generation), 32) + self.assertEqual(self.projection.capture_config_guard(config_path=config), guard) + result = self.publish(expected_config=guard) + self.projection.patch_config_fields({"synthetic_setting": 2, "other": True}, config_path=config) + mirrored = self.projection.mirror_activation(result, expected_config=guard, config_path=config) + self.assertEqual((mirrored.tier, mirrored.generation), ("pro", result.generation)) + self.assertEqual(json.loads(config.read_text(encoding="utf-8")), {"tier": "pro", "tier_activation_generation": result.generation, + "synthetic_setting": 2, "other": True}) + self.assertEqual(self.projection.mirror_activation(result, expected_config=None, config_path=config), mirrored) + self.assertFalse(self.api.read_activation_state(self.conn).legacy_policy.config_acknowledged) + + def test_config_unknown_and_concurrent_legacy_edit_refuse_without_write(self): + self.terminal_marker() + result = self.publish() + with tempfile.TemporaryDirectory(prefix="synthetic-policy-") as directory: + config = Path(directory) / "config.json" + config.write_text('{"tier":"base"}') + original = config.read_bytes() + with self.assertRaises(self.projection.ConfigProjectionConflict): + self.projection.mirror_activation(result, expected_config=None, config_path=config) + self.assertEqual(config.read_bytes(), original) + guard = self.projection.capture_config_guard(config_path=config) + config.write_text('{"tier":"edge"}') + edited = config.read_bytes() + with self.assertRaises(self.projection.ConfigProjectionConflict): + self.projection.mirror_activation(result, expected_config=guard, config_path=config) + self.assertEqual(config.read_bytes(), edited) + for field in ("tier", "tier_activation_generation"): + with self.assertRaises(self.projection.ConfigProjectionConflict): + self.projection.patch_config_fields({field: "synthetic"}, config_path=config) + + def test_config_lock_failure_strict_refuses_old_best_effort_remains(self): + with tempfile.TemporaryDirectory(prefix="synthetic-policy-") as directory: + config = Path(directory) / "config.json" + with patch.dict(self.projection.__dict__["__builtins__"], {"open": lambda *args, **kwargs: (_ for _ in ()).throw(PermissionError("synthetic lock"))}): + with self.projection.ConfigFileLock(config.with_suffix(".lock")): + pass + with self.assertRaises(self.projection.ConfigProjectionConflict): + self.projection.capture_config_guard(config_path=config) + self.assertFalse(config.exists()) + + def test_config_invalid_or_oversized_never_overwritten(self): + with tempfile.TemporaryDirectory(prefix="synthetic-policy-") as directory: + config = Path(directory) / "config.json" + for raw in (b'null', b'{"tier":"base","tier":"pro"}', b'{"tier":"base"}\x00' + b'x' * (1024 * 1024)): + config.write_bytes(raw) + with self.subTest(length=len(raw)), self.assertRaises(self.projection.ConfigProjectionConflict): + self.projection.capture_config_guard(config_path=config) + self.assertEqual(config.read_bytes(), raw) + + + def test_policy_refuses_temp_alias_trigger_and_stale_config_generation(self): + intent = self.stage() + self.conn.execute("CREATE TEMP TABLE METADATA(key TEXT, value TEXT)") + with self.assertRaises(self.api.TierActivationError): + self.publish(intent) + self.conn.execute("DROP TABLE temp.METADATA") + self.conn.execute("CREATE TRIGGER journal_side_effect AFTER UPDATE ON metadata BEGIN INSERT INTO old_serving_pair VALUES(x'01'); END") + with self.assertRaises(self.api.TierActivationError): + self.publish(intent) + self.conn.execute("DROP TRIGGER journal_side_effect") + result = self.publish(intent) + self.api.acknowledge_policy_config(self.conn, generation=result.generation) + with self.assertRaises(self.api.TierActivationError): + self.publish(expected_intent_id=result.intent_id, target=self.base, + expected_config=self.api.ConfigGuard(True, "pro", "b" * 32)) + self.assertEqual(self.api.read_activation_state(self.conn).legacy_policy.generation, result.generation) + + def test_corrupt_previous_policy_rejected_before_projection(self): + self.terminal_marker() + result = self.publish() + self.api.acknowledge_policy_config(self.conn, generation=result.generation) + job = self.jobs.select_tier_job(self.conn, self.fixture.target, previous_job_id=self.fixture.job.job_id) + intent = self.api.stage_activation_intent(self.conn, job, expected_generation=None, reranker_id="synthetic/ranker", + expected_config=self.api.ConfigGuard(True, "pro", result.generation), status_id=1) + for replacement in (None, "b" * 32): + value = json.loads(self.api._dump(intent)) + value["expected_config"]["generation"] = replacement + self.conn.execute("UPDATE metadata SET value=? WHERE key='tier_activation_v1'", (json.dumps(value),)) + self.conn.commit() + with self.assertRaises(self.api.TierActivationError): + self.api.read_activation_state(self.conn) + + def test_progress_oversized_nul_tail_never_materializes_payload(self): + _, plan = self.fixture.build() + key = self.source.manifest_key(plan.manifest.targets) + self.conn.execute("UPDATE metadata SET value=value||? WHERE key=?", ("\x00" + "x" * 66000, key)) + self.conn.commit() + statements = [] + self.conn.set_trace_callback(statements.append) + with self.assertRaises(self.source.TierSourceUntrusted): + self.source.read_tier_progress(self.conn, model_id=self.fixture.target.model_id, dimension=2, targets=plan.manifest.targets) + self.conn.set_trace_callback(None) + self.assertFalse(any(sql.startswith("SELECT value FROM main.metadata") for sql in statements)) + + def test_config_independent_patch_and_mirror_share_one_lock(self): + self.terminal_marker() + with tempfile.TemporaryDirectory(prefix="synthetic-policy-") as directory: + config = Path(directory) / "config.json" + config.write_text('{"tier":"base"}') + guard = self.projection.capture_config_guard(config_path=config) + result = self.publish(expected_config=guard) + barrier = threading.Barrier(2) + errors = [] + def patcher(): + try: + barrier.wait(timeout=2) + self.projection.patch_config_fields({"independent": 73}, config_path=config) + except BaseException as error: + errors.append(error) + thread = threading.Thread(target=patcher) + thread.start() + barrier.wait(timeout=2) + self.projection.mirror_activation(result, expected_config=guard, config_path=config) + thread.join(timeout=2) + self.assertFalse(thread.is_alive()) + self.assertEqual(errors, []) + self.assertEqual(json.loads(config.read_text(encoding="utf-8")), {"tier":"pro", "tier_activation_generation":result.generation, "independent":73}) + + def test_config_replace_after_success_error_preserves_pending_reconciliation(self): + self.terminal_marker() + with tempfile.TemporaryDirectory(prefix="synthetic-policy-") as directory: + config = Path(directory) / "config.json" + config.write_text('{"tier":"base"}') + guard = self.projection.capture_config_guard(config_path=config) + result = self.publish(expected_config=guard) + original = self.projection._replace + def fail_after(src, dst): + original(src, dst) + raise OSError("synthetic after replace") + with patch.object(self.projection, "_replace", fail_after), self.assertRaises(OSError): + self.projection.mirror_activation(result, expected_config=guard, config_path=config) + self.assertFalse(self.api.read_activation_state(self.conn).legacy_policy.config_acknowledged) + self.assertEqual(self.projection.mirror_activation(result, expected_config=guard, config_path=config).generation, result.generation) + + + +class TestIdentityErrorWire(unittest.TestCase): + def setUp(self): + fixtures = runpy.run_path(str(Path(__file__).with_name("test-tier-serving-receipt-795.py"))) + self.fixture = fixtures["TestServingReceipts"]() + self.fixture.setUp() + self.addCleanup(self.fixture.doCleanups) + + def test_identity_error_roundtrip_native_errors_still_generic(self): + fixture = self.fixture + target_module = sys.modules.get("truememory.embedding_target") + self.assertIsNotNone(target_module) + requests = [({"op": "embed", "expected_target": fixture.edge.to_wire()}, target_module.EmbeddingTargetError("synthetic identity")), + ({"op": "embed_target_v1"}, target_module.EmbeddingTargetError("synthetic custom identity")), + ({"op": "rerank", "expected_model_name": "synthetic/ranker"}, fixture.ms._ServingIdentityMismatch("synthetic reranker")), + ({"op": "embed", "expected_target": fixture.edge.to_wire()}, RuntimeError("synthetic native unavailable")), + ({"op": ["malformed"]}, ValueError("synthetic ordinary invalid request"))] + for request, error in requests: + replies = [] + server = fixture.server + server._recv_exact = lambda *args: json.dumps(request).encode() + server._send_response = lambda conn, response: replies.append(response) + server.handle_request = lambda request: (_ for _ in ()).throw(error) + server._serve_client(types.SimpleNamespace(settimeout=lambda _: None), types.SimpleNamespace(length=1, frame_expires_at=None)) + self.assertEqual(len(replies), 1) + response = replies[0] + self.assertFalse(response["ok"]) + if isinstance(error, (target_module.EmbeddingTargetError, fixture.ms._ServingIdentityMismatch)): + self.assertEqual(response["error_code"], "serving_identity_mismatch") + with self.assertRaises(fixture.client.ServingIdentityMismatchError): + fixture.client._check_model_response(response, request) + else: + self.assertNotIn("error_code", response) + with self.assertRaises(RuntimeError): + fixture.client._check_model_response(response, request) + + +if __name__ == "__main__": + unittest.main() diff --git a/truememory/mcp_server.py b/truememory/mcp_server.py index 6270b185..8c12e1e4 100644 --- a/truememory/mcp_server.py +++ b/truememory/mcp_server.py @@ -40,6 +40,8 @@ from mcp.server.fastmcp import FastMCP +from truememory.tier_switch.projection import CONFIG_WRITE_LOCK, ConfigFileLock + log = logging.getLogger(__name__) # --------------------------------------------------------------------------- @@ -58,7 +60,7 @@ _config_cache_time: float = 0.0 # M7 (#504) — Serialize config writes from concurrent MCP handlers / tier-switch. -_config_write_lock = threading.Lock() +_config_write_lock = CONFIG_WRITE_LOCK # M-58 (#641) — Cross-process advisory lock for config writes. _config_write_lock # above only serializes threads inside ONE process; a separate `truememory-mcp`, @@ -68,60 +70,11 @@ _CONFIG_LOCK_PATH = _TRUEMEMORY_DIR / "config.json.lock" -class _config_file_lock: - """Context manager taking a cross-process exclusive lock on config writes. - - Reuses the same fcntl/msvcrt pattern as - ``tier_switch.manager.run_rebuild_sync`` (#641). Best-effort: if the lock - cannot be acquired (e.g. a filesystem without flock support) we proceed - anyway rather than fail the write — the in-process ``_config_write_lock`` - plus the atomic ``os.replace`` still bound the damage. Held in BLOCKING mode - (LK_LOCK / LOCK_EX) so concurrent writers serialize instead of failing. - """ - - def __init__(self) -> None: - self._fd = None - - def __enter__(self): - try: - _CONFIG_LOCK_PATH.parent.mkdir(parents=True, exist_ok=True) - self._fd = open(_CONFIG_LOCK_PATH, "w") - if os.name == "nt": - import msvcrt - msvcrt.locking(self._fd.fileno(), msvcrt.LK_LOCK, 1) - else: - import fcntl - fcntl.flock(self._fd, fcntl.LOCK_EX) - except OSError: - # Could not open or lock — proceed unlocked (best effort). - if self._fd is not None: - try: - self._fd.close() - except OSError: - pass - self._fd = None - return self +class _config_file_lock(ConfigFileLock): + """Preserve the old best-effort entry point and injectable config paths.""" - def __exit__(self, *exc): - if self._fd is None: - return False - try: - if os.name == "nt": - import msvcrt - try: - msvcrt.locking(self._fd.fileno(), msvcrt.LK_UNLCK, 1) - except OSError: - pass - else: - import fcntl - fcntl.flock(self._fd, fcntl.LOCK_UN) - finally: - try: - self._fd.close() - except OSError: - pass - self._fd = None - return False + def __init__(self, *, strict: bool = False) -> None: + super().__init__(_CONFIG_LOCK_PATH, strict=strict) class _ConfigShapeError(ValueError): diff --git a/truememory/tier_switch/activation.py b/truememory/tier_switch/activation.py index 4b4c431a..33fc834a 100644 --- a/truememory/tier_switch/activation.py +++ b/truememory/tier_switch/activation.py @@ -40,6 +40,45 @@ class TierActivationError(RuntimeError): """Selection cannot be certified; no runtime/config success is implied.""" +@dataclass(frozen=True) +class ConfigGuard: + tier_present: bool + tier: str + generation: str | None + + def __post_init__(self) -> None: + if (type(self.tier_present) is not bool or type(self.tier) is not str + or self.tier not in {"edge", "base", "pro", "custom"} + or (not self.tier_present and self.tier != "edge") + or (self.generation is not None and not _hex(self.generation, 32))): + raise TierActivationError("Invalid expected config identity") + + +@dataclass(frozen=True) +class LegacyTierPolicy: + generation: str + intent_id: str + target: EmbeddingTarget + reranker_id: str + tables: tuple[str, str] + config_acknowledged: bool = False + + def __post_init__(self) -> None: + _policy_identity(self.target, self.reranker_id, self.tables) + if (not _hex(self.generation, 32) or not _hex(self.intent_id, 32) + or type(self.config_acknowledged) is not bool): + raise TierActivationError("Invalid legacy serving policy") + + +def _policy_identity(target: EmbeddingTarget, reranker_id: str, tables: tuple[str, str]) -> None: + if (type(target) is not EmbeddingTarget or target.tier not in {"base", "pro"} + or target.identity != ("qwen3_256", 256, "basepro") + or type(reranker_id) is not str or reranker_id != "Alibaba-NLP/gte-reranker-modernbert-base" + or type(tables) is not tuple or len(tables) != 2 or any(type(name) is not str for name in tables) + or tables not in {target.tables, ("vec_messages", "vec_messages_sep")}): + raise TierActivationError("Policy-only publication requires the unchanged base/pro embedding space") + + def _hex(value: object, size: int) -> bool: return type(value) is str and re.fullmatch(r"[0-9a-f]{" + str(size) + "}", value) is not None @@ -76,6 +115,9 @@ class ActivationIntent: source_epoch: str source_schema_signature: str state: str = "staged" + expected_config: ConfigGuard | None = None + status_id: int | None = None + previous_policy: LegacyTierPolicy | None = None def __post_init__(self) -> None: _identity(self.target, self.reranker_id, self.job_id, self.tracker, @@ -84,6 +126,37 @@ def __post_init__(self) -> None: or (self.expected_generation is not None and not _hex(self.expected_generation, 32)) or type(self.state) is not str or self.state not in {"staged", "db_selected", "config_acknowledged"}): raise TierActivationError("Invalid activation intent") + if (self.expected_config is not None and type(self.expected_config) is not ConfigGuard + or self.status_id is not None and (type(self.status_id) is not int or not 1 <= self.status_id < 2**63) + or self.previous_policy is not None and type(self.previous_policy) is not LegacyTierPolicy + or self.expected_config is None and (self.status_id is not None or self.previous_policy is not None) + or self.previous_policy is not None and self.expected_generation is not None): + raise TierActivationError("Invalid activation projection evidence") + + +@dataclass(frozen=True) +class PolicyIntent: + intent_id: str + expected_generation: str | None + expected_intent_id: str | None + target: EmbeddingTarget + reranker_id: str + expected_config: ConfigGuard + result_generation: str + tables: tuple[str, str] + state: str = "db_selected" + + def __post_init__(self) -> None: + _policy_identity(self.target, self.reranker_id, self.tables) + if (not _hex(self.intent_id, 32) or not _hex(self.result_generation, 32) + or self.expected_generation is not None and not _hex(self.expected_generation, 32) + or self.expected_intent_id is not None and not _hex(self.expected_intent_id, 32) + or self.result_generation == self.expected_generation + or self.intent_id == self.expected_intent_id + or type(self.expected_config) is not ConfigGuard + or self.expected_config.tier not in {"base", "pro"} + or type(self.state) is not str or self.state not in {"db_selected", "config_acknowledged"}): + raise TierActivationError("Invalid policy activation intent") @dataclass(frozen=True) @@ -122,13 +195,52 @@ def __post_init__(self) -> None: @dataclass(frozen=True) class ActivationState: selection: TierSelection | None - intent: ActivationIntent | None + intent: ActivationIntent | PolicyIntent | None + legacy_policy: LegacyTierPolicy | None = None -def _dump(record: ActivationIntent | TierSelection) -> str: +def _config_wire(value: ConfigGuard) -> dict: + return dict(asdict(value), version=1) + + +def _parse_config(value: object) -> ConfigGuard: + if (type(value) is not dict or set(value) != {"version", "tier_present", "tier", "generation"} + or type(value["version"]) is not int or value["version"] != 1): + raise TierActivationError("Unsupported expected config descriptor") + return ConfigGuard(value["tier_present"], value["tier"], value["generation"]) + + +def _policy_wire(value: LegacyTierPolicy) -> dict: + return dict(asdict(value), version=1, target=value.target.to_wire()) + + +def _parse_policy(value: object) -> LegacyTierPolicy: + if (type(value) is not dict or set(value) != set(LegacyTierPolicy.__dataclass_fields__) | {"version"} + or type(value["version"]) is not int or value["version"] != 1 + or type(value["tables"]) is not list or len(value["tables"]) != 2): + raise TierActivationError("Unsupported previous legacy policy") + return LegacyTierPolicy(value["generation"], value["intent_id"], EmbeddingTarget.from_wire(value["target"]), + value["reranker_id"], tuple(value["tables"]), value["config_acknowledged"]) + + +def _dump(record: ActivationIntent | TierSelection | PolicyIntent) -> str: value = asdict(record) value["target"] = record.target.to_wire() - raw = json.dumps(dict(value, version=1), sort_keys=True, separators=(",", ":"), allow_nan=False) + version = 1 + if type(record) is ActivationIntent: + if record.expected_config is None: + for field in ("expected_config", "status_id", "previous_policy"): + value.pop(field) + else: + version = 2 + value["kind"] = "rebuild" + value["expected_config"] = _config_wire(record.expected_config) + value["previous_policy"] = None if record.previous_policy is None else _policy_wire(record.previous_policy) + elif type(record) is PolicyIntent: + version = 2 + value["kind"] = "policy" + value["expected_config"] = _config_wire(record.expected_config) + raw = json.dumps(dict(value, version=version), sort_keys=True, separators=(",", ":"), allow_nan=False) if len(raw.encode("utf-8")) > _MAX_RECORD_BYTES: raise TierActivationError("Activation record exceeds protocol bounds") return raw @@ -141,14 +253,32 @@ def _unique(pairs: list[tuple[str, object]]) -> dict[str, object]: return value -def _parse(raw: str, kind: type[ActivationIntent] | type[TierSelection]) -> ActivationIntent | TierSelection: +def _parse(raw: str, kind: type[ActivationIntent] | type[TierSelection]) -> ActivationIntent | TierSelection | PolicyIntent: try: value = json.loads(raw, object_pairs_hook=_unique) - if (type(value) is not dict or set(value) != set(kind.__dataclass_fields__) | {"version"} - or type(value["version"]) is not int or value.pop("version") != 1): + if type(value) is not dict or type(value.get("version")) is not int: raise TierActivationError("Unsupported activation record schema") + version = value.pop("version") + expected = set(kind.__dataclass_fields__) + if kind is ActivationIntent and version == 1: + expected -= {"expected_config", "status_id", "previous_policy"} + elif kind is ActivationIntent and version == 2: + discriminator = value.pop("kind", None) + if discriminator == "policy": + kind = PolicyIntent + elif discriminator != "rebuild": + raise TierActivationError("Unsupported activation intent kind") + expected = set(kind.__dataclass_fields__) + elif version != 1: + raise TierActivationError("Unsupported activation record version") + if set(value) != expected: + raise TierActivationError("Unsupported activation record fields") + if version == 2: + value["expected_config"] = _parse_config(value["expected_config"]) + if kind is ActivationIntent and value["previous_policy"] is not None: + value["previous_policy"] = _parse_policy(value["previous_policy"]) value["target"] = EmbeddingTarget.from_wire(value["target"]) - if kind is TierSelection: + if kind in (TierSelection, PolicyIntent): if type(value["tables"]) is not list or len(value["tables"]) != 2: raise TierActivationError("Invalid selected table pair") value["tables"] = tuple(value["tables"]) @@ -186,8 +316,8 @@ def _guard_storage(conn: sqlite3.Connection) -> None: raise TierActivationError("Activation storage has unsupported side effects") -def _guard_pair(conn: sqlite3.Connection, target: EmbeddingTarget) -> None: - for name in target.tables: +def _guard_pair(conn: sqlite3.Connection, target: EmbeddingTarget, *, tables: tuple[str, str] | None = None) -> None: + for name in target.tables if tables is None else tables: declaration = ( r"CREATE\s+VIRTUAL\s+TABLE\s+(?:IF\s+NOT\s+EXISTS\s+)?" + r'(?:"' + name + r'"|`' + name + r'`|\[' + name + r'\]|' + name + r')' @@ -199,7 +329,7 @@ def _guard_pair(conn: sqlite3.Connection, target: EmbeddingTarget) -> None: def _read(conn: sqlite3.Connection, key: str, kind: type[ActivationIntent] | type[TierSelection] - ) -> ActivationIntent | TierSelection | None: + ) -> ActivationIntent | TierSelection | PolicyIntent | None: shape = conn.execute( "SELECT typeof(value), length(CAST(value AS BLOB)) FROM main.metadata WHERE key=?", (key,), ).fetchone() @@ -249,10 +379,31 @@ def _state(conn: sqlite3.Connection) -> ActivationState: selected, intent = _read(conn, _SELECTED, TierSelection), _read(conn, _INTENT, ActivationIntent) if selected is not None and intent is None: raise TierActivationError("Selected record has no activation intent") + if type(intent) is PolicyIntent: + acknowledged = intent.state == "config_acknowledged" + if selected is None: + if intent.expected_generation is not None: + raise TierActivationError("Certified policy lost its selection") + return ActivationState(None, intent, LegacyTierPolicy( + intent.result_generation, intent.intent_id, intent.target, intent.reranker_id, + intent.tables, acknowledged, + )) + if ((selected.generation, selected.intent_id, selected.target, selected.reranker_id, + selected.tables, selected.config_acknowledged) + != (intent.result_generation, intent.intent_id, intent.target, intent.reranker_id, + intent.tables, acknowledged) or intent.expected_generation is None): + raise TierActivationError("Certified policy and selected record disagree") + return ActivationState(selected, intent) if intent is not None and intent.state == "staged": _expected(selected, intent.expected_generation) if selected is not None and not selected.config_acknowledged: raise TierActivationError("Pending config selection cannot have a new staged intent") + if intent.previous_policy is not None and ( + not intent.previous_policy.config_acknowledged + or intent.expected_config.tier != intent.previous_policy.target.tier + or intent.expected_config.generation != intent.previous_policy.generation + ): + raise TierActivationError("Staged intent has incoherent prior legacy policy evidence") if intent is not None and intent.state != "staged": if (selected is None or selected.intent_id != intent.intent_id or selected.target != intent.target or selected.reranker_id != intent.reranker_id @@ -260,7 +411,7 @@ def _state(conn: sqlite3.Connection) -> ActivationState: != (intent.job_id, intent.tracker, intent.source_epoch, intent.source_schema_signature) or selected.config_acknowledged != (intent.state == "config_acknowledged")): raise TierActivationError("Selected record and activation result disagree") - return ActivationState(selected, intent) + return ActivationState(selected, intent, intent.previous_policy if intent is not None and intent.state == "staged" else None) def read_activation_state(conn: sqlite3.Connection) -> ActivationState: @@ -282,7 +433,9 @@ def _expected(selection: TierSelection | None, generation: str | None) -> None: def stage_activation_intent(conn: sqlite3.Connection, job: TierJob, *, - expected_generation: str | None, reranker_id: str) -> ActivationIntent: + expected_generation: str | None, reranker_id: str, + expected_config: ConfigGuard | None = None, + status_id: int | None = None) -> ActivationIntent: """Commit intent before source planning; identical current intent is idempotent. None expects absent selection and permits only a later certified target. It @@ -291,7 +444,8 @@ def stage_activation_intent(conn: sqlite3.Connection, job: TierJob, *, replaced its marker. Conflicting intent for a live job is refused. """ intent = ActivationIntent(uuid.uuid4().hex, expected_generation, job.target, reranker_id, - job.job_id, job.tracker, job.source_epoch, job.source_schema_signature) + job.job_id, job.tracker, job.source_epoch, job.source_schema_signature, + expected_config=expected_config, status_id=status_id) with _transaction(conn, write=True): _guard_storage(conn) check_selected_tier_job_in_writer(conn, job) @@ -299,7 +453,17 @@ def stage_activation_intent(conn: sqlite3.Connection, job: TierJob, *, _expected(state.selection, expected_generation) if state.selection is not None and not state.selection.config_acknowledged: raise TierActivationError("Current selection awaits config acknowledgement") - if state.intent is not None and state.intent.job_id == job.job_id: + if expected_config is not None and state.selection is not None and ( + expected_config.tier != state.selection.target.tier or expected_config.generation != state.selection.generation + ): + raise TierActivationError("Prior config does not identify the selected generation") + if state.legacy_policy is not None: + if (not state.legacy_policy.config_acknowledged or expected_config is None + or expected_config.tier != state.legacy_policy.target.tier + or expected_config.generation != state.legacy_policy.generation): + raise TierActivationError("Legacy policy requires acknowledged, guarded activation") + intent = replace(intent, previous_policy=state.legacy_policy) + if type(state.intent) is ActivationIntent and state.intent.job_id == job.job_id: if state.intent == replace(intent, intent_id=state.intent.intent_id): return state.intent raise TierActivationError("A conflicting intent already owns this live job") @@ -381,3 +545,94 @@ def acknowledge_config(conn: sqlite3.Connection, *, generation: str) -> TierSele _put(conn, _SELECTED, _dump(selected)) _put(conn, _INTENT, _dump(replace(state.intent, state="config_acknowledged"))) return selected + + +def _guard_legacy_policy_pair(conn: sqlite3.Connection, target: EmbeddingTarget, tables: tuple[str, str], + *, certified: bool = False) -> None: + _guard_pair(conn, target, tables=tables) + for key, value in (("embed_model", target.model_id), ("embed_dim", str(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 TierActivationError("Legacy embedding metadata is not coherent with base/pro") + model_match = "model_name='qwen3_256'" if certified else "model_name IN ('qwen3_256','Qwen3-Embedding-0.6B')" + row = conn.execute( + "SELECT vec_table=? AND sep_table=? AND typeof(embedding_dim)='integer' AND embedding_dim=? " + f"AND {model_match} FROM main.vector_cache_registry WHERE tier_group='basepro'", (*tables, target.dimension), + ).fetchone() + if (row is None and (certified or tables != ("vec_messages", "vec_messages_sep"))) or (row is not None and row[0] != 1): + raise TierActivationError("Serving pair disagrees with its registry") + + + +def commit_config_only_transition( + conn: sqlite3.Connection, *, expected_selection_generation: str | None, + expected_intent_id: str | None, target: EmbeddingTarget, reranker_id: str, + expected_config: ConfigGuard, legacy_pair: tuple[str, str] | None = None, +) -> TierSelection | LegacyTierPolicy: + """Publish only base/pro policy, without a corpus read or vector certificate. + + Caller owns stable-path maintenance and exclusive serving. A legacy pair + must be its coherently observed serving pair, never inferred from counts. + Existing certificates are historical provenance, not recertified coverage. + """ + from truememory.tier_switch.job import _cancel_policy_job_in_writer, read_selected_job_marker + + if type(expected_config) is not ConfigGuard or expected_config.tier not in {"base", "pro"}: + raise TierActivationError("Policy transition requires actual prior config evidence") + if expected_intent_id is not None and not _hex(expected_intent_id, 32): + raise TierActivationError("Invalid expected prior activation intent") + with _transaction(conn, write=True): + _guard_storage(conn) + state = _state(conn) + _expected(state.selection, expected_selection_generation) + if (state.intent.intent_id if state.intent is not None else None) != expected_intent_id: + raise TierActivationError("Prior activation intent changed") + previous = state.selection or state.legacy_policy + if previous is not None and not previous.config_acknowledged: + raise TierActivationError("Prior serving policy awaits config acknowledgement") + tables = state.selection.tables if state.selection is not None else legacy_pair + _policy_identity(target, reranker_id, tables) + if target.tier == expected_config.tier: + raise TierActivationError("Policy-only transition must change base/pro public tier") + if previous is not None and ( + previous.target.tier != expected_config.tier or previous.generation != expected_config.generation + or previous.target.identity != target.identity + or previous.tables != tables or previous.reranker_id != reranker_id + ): + raise TierActivationError("Policy-only transition would change the serving embedding space") + _guard_legacy_policy_pair(conn, target, tables, certified=state.selection is not None) + if type(state.intent) is ActivationIntent and state.intent.state == "staged": + _cancel_policy_job_in_writer( + conn, job_id=state.intent.job_id, target=state.intent.target, tracker=state.intent.tracker, + source_epoch=state.intent.source_epoch, source_schema_signature=state.intent.source_schema_signature, + ) + else: + marker = read_selected_job_marker(conn) + if marker is not None and marker.state == "selected": + raise TierActivationError("An unrelated selected job cannot be superseded") + intent = PolicyIntent(uuid.uuid4().hex, expected_selection_generation, expected_intent_id, + target, reranker_id, expected_config, uuid.uuid4().hex, tables) + if state.selection is not None: + result = replace(state.selection, generation=intent.result_generation, intent_id=intent.intent_id, + target=target, config_acknowledged=False) + _put(conn, _SELECTED, _dump(result)) + else: + result = LegacyTierPolicy(intent.result_generation, intent.intent_id, target, reranker_id, tables) + _put(conn, _INTENT, _dump(intent)) + return result + + +def acknowledge_policy_config(conn: sqlite3.Connection, *, generation: str) -> LegacyTierPolicy: + """Acknowledge only this legacy policy's config mirror, never native state.""" + if not _hex(generation, 32): + raise TierActivationError("A concrete policy generation is required") + with _transaction(conn, write=True): + _guard_storage(conn) + state = _state(conn) + if type(state.intent) is not PolicyIntent or state.legacy_policy is None or state.legacy_policy.generation != generation: + raise TierActivationError("Legacy policy generation changed") + if state.legacy_policy.config_acknowledged: + return state.legacy_policy + _put(conn, _INTENT, _dump(replace(state.intent, state="config_acknowledged"))) + result = replace(state.legacy_policy, config_acknowledged=True) + return result diff --git a/truememory/tier_switch/job.py b/truememory/tier_switch/job.py index ead335d5..a5f0dbd0 100644 --- a/truememory/tier_switch/job.py +++ b/truememory/tier_switch/job.py @@ -184,6 +184,71 @@ class TierJobSelectionError(TierJobIdentityError): """The durable selected job is unavailable or does not match this request.""" +@dataclass(frozen=True) +class SelectedJobMarker: + job_id: str + target: EmbeddingTarget + tracker: str + source_epoch: str + source_schema_signature: str + state: str + + +def _marker_in_snapshot(conn: sqlite3.Connection) -> SelectedJobMarker | None: + if conn.execute("SELECT 1 FROM temp.sqlite_master LIMIT 1").fetchone() is not None: + raise TierJobSelectionError("TEMP objects make selected marker identity ambiguous") + shape = conn.execute( + "SELECT typeof(sql), length(CAST(sql AS BLOB)) FROM main.sqlite_master " + "WHERE name=? COLLATE NOCASE", (_SELECTION_TABLE,), + ).fetchone() + if shape is None: + return None + if shape[0] != "text" or not 1 <= shape[1] <= 8192: + raise TierJobSelectionError("Selected job schema exceeds its protocol bounds") + _selection_schema(conn) + row = _selection_row(conn) + return None if row is None else SelectedJobMarker( + row[0], EmbeddingTarget(row[1], row[2], row[3], row[4]), row[5], row[6], row[7], row[8], + ) + + +def read_selected_job_marker(conn: sqlite3.Connection) -> SelectedJobMarker | None: + """Bounded marker read; borrow a caller snapshot without ending it. + + This is not a reopened-file identity certificate or permission to resume a + disappeared worker. Terminal IDs permit explicit normal job replacement. + """ + if conn.in_transaction: + return _marker_in_snapshot(conn) + with _selection_transaction(conn, write=False): + return _marker_in_snapshot(conn) + + +def _cancel_policy_job_in_writer( + conn: sqlite3.Connection, *, job_id: str, target: EmbeddingTarget, tracker: str, + source_epoch: str, source_schema_signature: str, +) -> None: + """Cancel only the exact superseded intent's marker; caller owns writer. + + Maintenance/stable-path ownership and writer acquisition are caller duties. + No source read, file observation, DELETE, native call, commit or rollback. + """ + if not conn.in_transaction: + raise TierJobSelectionError("Policy supersession requires writer ownership") + row = _marker_in_snapshot(conn) + if row is None or (row.job_id, row.target, row.tracker, row.source_epoch, row.source_schema_signature) != ( + job_id, target, tracker, source_epoch, source_schema_signature, + ): + raise TierJobSelectionError("Superseded activation job changed") + if row.state == "selected": + changed = conn.execute( + f"UPDATE main.{_SELECTION_TABLE} SET state='cancelled' " + "WHERE singleton=1 AND job_id=? AND state='selected'", (job_id,), + ).rowcount + if changed != 1: + raise TierJobSelectionError("Superseded activation job changed during cancellation") + + def _selection_values(job: TierJob) -> tuple[str | int, ...]: values = (job.job_id, job.target.tier, job.target.model_id, job.target.dimension, job.target.tier_group, job.tracker, job.source_epoch, job.source_schema_signature) diff --git a/truememory/tier_switch/projection.py b/truememory/tier_switch/projection.py new file mode 100644 index 00000000..b83ba27d --- /dev/null +++ b/truememory/tier_switch/projection.py @@ -0,0 +1,215 @@ +"""Strict config projection of an already committed database serving decision. + +This is a cooperating-writer CAS, not atomicity across JSON and SQLite. Database +selection remains authoritative after any config write or acknowledgement fault. +""" + +from __future__ import annotations + +import json +import os +import tempfile +import threading +import time +import uuid +from pathlib import Path +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from truememory.tier_switch.activation import ConfigGuard, LegacyTierPolicy, TierSelection + +CONFIG_WRITE_LOCK = threading.Lock() +_MAX_CONFIG_BYTES = 1024 * 1024 +_GENERATION_KEY = "tier_activation_generation" + + +class ConfigProjectionConflict(RuntimeError): + """Config is unavailable or has changed since the durable intent.""" + + +class ConfigFileLock: + """Shared blocking lock; old writers retain optional best-effort behavior.""" + + def __init__(self, path: Path, *, strict: bool = False) -> None: + self.path = path + self.strict = strict + self._fd = None + + def __enter__(self) -> ConfigFileLock: + try: + self.path.parent.mkdir(parents=True, exist_ok=True) + self._fd = open(self.path, "a+b") + if os.name == "nt": + import msvcrt + self._fd.seek(0) + msvcrt.locking(self._fd.fileno(), msvcrt.LK_LOCK, 1) + else: + import fcntl + fcntl.flock(self._fd, fcntl.LOCK_EX) + except OSError as error: + if self._fd is not None: + try: + self._fd.close() + except OSError: + pass + self._fd = None + if self.strict: + raise ConfigProjectionConflict("Config lock is unavailable") from error + return self + + def __exit__(self, *exc: object) -> bool: + if self._fd is None: + return False + try: + if os.name == "nt": + import msvcrt + self._fd.seek(0) + try: + msvcrt.locking(self._fd.fileno(), msvcrt.LK_UNLCK, 1) + except OSError: + if self.strict: + raise + else: + import fcntl + fcntl.flock(self._fd, fcntl.LOCK_UN) + finally: + try: + self._fd.close() + except OSError: + if self.strict: + raise + finally: + self._fd = None + return False + + +def _paths(config_path: Path | None, lock_path: Path | None) -> tuple[Path, Path]: + config = Path.home() / ".truememory" / "config.json" if config_path is None else Path(config_path) + return config, config.with_name(config.name + ".lock") if lock_path is None else Path(lock_path) + + +def _unique(pairs: list[tuple[str, object]]) -> dict: + value = dict(pairs) + if len(value) != len(pairs): + raise ConfigProjectionConflict("Duplicate config field") + return value + + +def _read_config(path: Path) -> dict: + try: + with path.open("rb") as stream: + raw = stream.read(_MAX_CONFIG_BYTES + 1) + except FileNotFoundError: + return {} + if len(raw) > _MAX_CONFIG_BYTES: + raise ConfigProjectionConflict("Config exceeds the projection byte bound") + try: + value = json.loads(raw.decode("utf-8-sig"), object_pairs_hook=_unique, + parse_constant=lambda _: (_ for _ in ()).throw(ValueError("Nonfinite config"))) + except (ValueError, UnicodeError, RecursionError) as error: + raise ConfigProjectionConflict("Config is not a bounded JSON object") from error + if type(value) is not dict: + raise ConfigProjectionConflict("Config is not a JSON object") + return value + + +def _guard(config: dict) -> ConfigGuard: + from truememory.tier_switch.activation import ConfigGuard, TierActivationError + + try: + return ConfigGuard("tier" in config, config.get("tier", "edge"), config.get(_GENERATION_KEY)) + except TierActivationError as error: + raise ConfigProjectionConflict("Config has an invalid tier or generation") from error + + +def _replace(src: str, dst: str) -> None: + for attempt in range(10): + try: + os.replace(src, dst) + return + except PermissionError: + if attempt == 9: + raise + time.sleep(0.03) + + +def _write_config(path: Path, config: dict) -> None: + try: + raw = (json.dumps(config, indent=2, ensure_ascii=True, allow_nan=False) + "\n").encode("utf-8") + except (ValueError, TypeError, RecursionError) as error: + raise ConfigProjectionConflict("Config fields are not serializable") from error + if len(raw) > _MAX_CONFIG_BYTES: + raise ConfigProjectionConflict("Config exceeds the projection byte bound") + path.parent.mkdir(parents=True, exist_ok=True) + fd, temporary = tempfile.mkstemp(prefix=".config.tmp.", suffix=".json", dir=str(path.parent)) + try: + with os.fdopen(fd, "wb") as stream: + stream.write(raw) + stream.flush() + os.fsync(stream.fileno()) + os.chmod(temporary, 0o600) + _replace(temporary, str(path)) + except BaseException: + try: + os.unlink(temporary) + except OSError: + pass + raise + + +def capture_config_guard(*, config_path: Path | None = None, + lock_path: Path | None = None) -> ConfigGuard: + """Capture actual prior tier and seed a generation without changing policy.""" + config_path, lock_path = _paths(config_path, lock_path) + with ConfigFileLock(lock_path, strict=True), CONFIG_WRITE_LOCK: + config = _read_config(config_path) + guard = _guard(config) + if guard.generation is None: + config[_GENERATION_KEY] = uuid.uuid4().hex + _write_config(config_path, config) + guard = _guard(config) + return guard + + +def patch_config_fields(fields: dict, *, config_path: Path | None = None, + lock_path: Path | None = None) -> dict: + """Merge independent settings under the same lock as activation projection.""" + if type(fields) is not dict or any(type(key) is not str for key in fields): + raise ConfigProjectionConflict("Config patch must be an object with string keys") + if {"tier", _GENERATION_KEY}.intersection(fields): + raise ConfigProjectionConflict("Policy fields require a committed activation") + config_path, lock_path = _paths(config_path, lock_path) + with ConfigFileLock(lock_path, strict=True), CONFIG_WRITE_LOCK: + config = _read_config(config_path) + config.update(fields) + _write_config(config_path, config) + return config.copy() + + +def mirror_activation(result: TierSelection | LegacyTierPolicy, *, + expected_config: ConfigGuard | None, + config_path: Path | None = None, lock_path: Path | None = None) -> ConfigGuard: + """CAS the tier/generation only; UNKNOWN old intents cannot authorize a write. + + An already-applied exact result is idempotent. Other settings come from the + fresh locked file. This does not reconcile native runtime or acknowledge DB. + """ + from truememory.tier_switch.activation import ConfigGuard, LegacyTierPolicy, TierSelection + + if type(result) not in (TierSelection, LegacyTierPolicy): + raise ConfigProjectionConflict("A typed committed serving result is required") + if expected_config is not None and type(expected_config) is not ConfigGuard: + raise ConfigProjectionConflict("Invalid expected config guard") + config_path, lock_path = _paths(config_path, lock_path) + with ConfigFileLock(lock_path, strict=True), CONFIG_WRITE_LOCK: + config = _read_config(config_path) + actual = _guard(config) + desired = ConfigGuard(True, result.target.tier, result.generation) + if actual == desired: + return actual + if expected_config is None or actual != expected_config: + raise ConfigProjectionConflict("Config changed since the activation intent") + config["tier"] = result.target.tier + config[_GENERATION_KEY] = result.generation + _write_config(config_path, config) + return desired diff --git a/truememory/tier_switch/source.py b/truememory/tier_switch/source.py index 7afeefdc..03a0ab4e 100644 --- a/truememory/tier_switch/source.py +++ b/truememory/tier_switch/source.py @@ -228,6 +228,74 @@ def _parse_receipt(raw: str, manifest: RebuildManifest, schema: str) -> tuple[st raise TierSourceUntrusted("Malformed or incompatible tier pair receipt") from error +@dataclass(frozen=True) +class TierSourceProgress: + generation: str + processed: int + total: int + outputs: int + cursor: int | None + complete: bool + tracker: str + source_epoch: str + source_schema_signature: str + + +def read_tier_progress(conn: sqlite3.Connection, *, model_id: str, dimension: int, + targets: tuple[str, str]) -> TierSourceProgress | None: + """Read bounded receipt scalars, not source coverage or current completion. + + A caller may own one read snapshot for journal, selected marker and these + counters, then compare their exact target/tracker/epoch/schema identities. + This neither audits vectors nor checks current source freshness. The + manifest's complete flag certifies only its captured range. + """ + _identity(model_id, dimension, targets) + + def read() -> TierSourceProgress | None: + guarded = ("metadata", *targets) + marks = ",".join("?" for _ in guarded) + if conn.execute( + f"SELECT 1 FROM temp.sqlite_master WHERE name COLLATE NOCASE IN ({marks}) LIMIT 1", guarded, + ).fetchone(): + raise TierSourceUntrusted("TEMP shadows invalidate progress identity") + raw = _metadata(conn, manifest_key(targets)) + receipt = _metadata(conn, _receipt_key(targets)) + if raw is None and receipt is None: + return None + if raw is None or receipt is None: + raise TierSourceUntrusted("Progress requires its paired receipt") + manifest = _parse_manifest(raw, model_id, dimension, targets) + try: + value = json.loads(receipt, object_pairs_hook=_unique_object) + schema = value["schema"] + if type(schema) is not str or not _HEX.fullmatch(schema): + raise ValueError("receipt schema") + except (ValueError, TypeError, KeyError, RecursionError) as error: + raise TierSourceUntrusted("Malformed progress receipt") from error + _parse_receipt(receipt, manifest, schema) + return TierSourceProgress(manifest.generation, manifest.consumed, manifest.total, + manifest.outputs, manifest.cursor, manifest.complete, + manifest.source.tracker, manifest.source.revision.epoch, + manifest.source.schema_signature) + + if conn.in_transaction: + return read() + conn.execute("BEGIN") + try: + result = read() + except BaseException: + if conn.in_transaction: + conn.rollback() + if conn.in_transaction: + raise TierSourceUntrusted("Progress rollback is uncertain") + raise + conn.rollback() + if conn.in_transaction: + raise TierSourceUntrusted("Progress snapshot did not close") + return result + + def _chain(digest: str, row_id: int, completion: bytes, separation: bytes) -> str: return hashlib.sha256(bytes.fromhex(digest) + struct.pack(">qQQ", row_id, len(completion), len(separation)) + completion + separation).hexdigest() From 5bea8b11f70bbda6c09a77bcfcffa7d6a81687ad Mon Sep 17 00:00:00 2001 From: buildingjoshbetter Date: Fri, 9 Oct 2026 15:21:26 -0500 Subject: [PATCH 02/11] fix: Adopt guarded tier authority across public operations --- tests/conftest.py | 23 +- tests/test-clustering-health-720.py | 18 +- tests/test-installer-model-readiness-767.py | 107 +- tests/test-l0-style-accumulator-747.py | 4 +- tests/test-l0-style-incremental-747.py | 5 +- tests/test-l0-style-maintenance-747.py | 44 +- tests/test-l0-style-routing-747.py | 31 +- tests/test-maintenance-engine-routing-753.py | 79 +- .../test-maintenance-fixture-isolation-753.py | 8 + tests/test-maintenance-source-revision-753.py | 61 +- ...test-rebuild-foreground-publication-756.py | 4 +- tests/test-rebuild-source-streaming-756.py | 50 +- tests/test-tier-active-initialization-795.py | 465 ++++ tests/test-tier-foreground-native-795.py | 664 +++++ tests/test-tier-mcp-serving-795.py | 281 ++ tests/test-tier-metadata-publication-795.py | 176 ++ tests/test-tier-public-adoption-795.py | 1686 ++++++++++++ tests/test-tier-public-boundaries-795.py | 2348 +++++++++++++++++ tests/test-tier-runtime-bridge-795.py | 47 +- tests/test-tier-warm-admission-795.py | 493 ++++ tests/test-tier-writer-fence-795.py | 45 +- tests/test_configure_reembed.py | 11 + tests/test_configure_tier_persist.py | 14 +- tests/test_issue_396_memory_intensity.py | 13 +- tests/test_issue_497_edge_persist.py | 13 +- tests/test_issue_502_config.py | 13 +- ...t_issue_600_antigravity_groq_deepsearch.py | 9 + tests/test_issue_648_vec_compat.py | 13 +- tests/test_l4_entity_sheets_disabled.py | 10 +- tests/test_l5_surprise_rerank.py | 10 +- tests/test_mcp_safety.py | 22 +- tests/test_reranker_tier.py | 83 +- truememory/engine.py | 1373 ++++++---- truememory/hybrid.py | 8 +- truememory/ingest/cli.py | 240 +- truememory/ingest/dedup.py | 2 + truememory/ingest/encoding_gate.py | 39 +- truememory/ingest/pipeline.py | 306 +-- truememory/maintenance.py | 354 ++- truememory/mcp_server.py | 273 +- truememory/model_client.py | 6 + truememory/model_server.py | 23 +- truememory/reranker.py | 17 + truememory/tier_switch/manager.py | 865 ++++-- truememory/tier_switch/runtime.py | 427 ++- truememory/tier_switch/serving.py | 11 + truememory/tier_switch/writer.py | 60 +- truememory/vector_search.py | 656 +++-- 48 files changed, 9938 insertions(+), 1572 deletions(-) create mode 100644 tests/test-tier-active-initialization-795.py create mode 100644 tests/test-tier-foreground-native-795.py create mode 100644 tests/test-tier-mcp-serving-795.py create mode 100644 tests/test-tier-metadata-publication-795.py create mode 100644 tests/test-tier-public-adoption-795.py create mode 100644 tests/test-tier-public-boundaries-795.py create mode 100644 tests/test-tier-warm-admission-795.py diff --git a/tests/conftest.py b/tests/conftest.py index 83424401..cc83c806 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -129,30 +129,43 @@ def _isolate_environ(): @pytest.fixture(autouse=True) def _isolate_vector_search_globals(): - """Snapshot/restore ``truememory.vector_search`` embedder globals. + """Snapshot/restore the complete process-local serving identity. These are process-global (module-level) and several tests mutate them via ``set_embedding_model`` or ``monkeypatch.setattr``. monkeypatch reverts the - ones it owns, but a defensive snapshot here keeps the embedder identity - deterministic regardless of test order or ambient config.json. + ones it owns, but policy projection also changes embedding/reranker identity + and runtime acknowledgement. Restore those fields together without resetting + serving leases, runtime locals or background workers. """ try: from truememory import vector_search as _vs + from truememory import reranker as _rr + from truememory.tier_switch import runtime as _runtime except Exception: # vector_search may be unimportable in minimal-dep environments; nothing # to isolate in that case. yield return - saved = ( + saved_vector = ( getattr(_vs, "EMBEDDING_MODEL", None), getattr(_vs, "_embedding_dim", None), getattr(_vs, "_model", None), + getattr(_vs, "_model_generation", 0), + getattr(_vs, "_frozen_embedding_target", None), + getattr(_vs, "_runtime_policy_tier", None), ) + saved_reranker = (_rr._model, _rr._model_name, _rr._model_certified, + _rr._active_tier, _rr._frozen_reranker_id) + saved_acknowledgement = _runtime._acknowledged try: yield finally: - _vs.EMBEDDING_MODEL, _vs._embedding_dim, _vs._model = saved + (_vs.EMBEDDING_MODEL, _vs._embedding_dim, _vs._model, _vs._model_generation, + _vs._frozen_embedding_target, _vs._runtime_policy_tier) = saved_vector + (_rr._model, _rr._model_name, _rr._model_certified, + _rr._active_tier, _rr._frozen_reranker_id) = saved_reranker + _runtime._acknowledged = saved_acknowledgement @pytest.fixture(autouse=True) diff --git a/tests/test-clustering-health-720.py b/tests/test-clustering-health-720.py index 9cd063ca..59190b91 100644 --- a/tests/test-clustering-health-720.py +++ b/tests/test-clustering-health-720.py @@ -6,6 +6,7 @@ import logging import os import re +import runpy import sqlite3 import sys import threading @@ -33,10 +34,14 @@ def load_functions(path: str, names: set[str], namespace: dict, *, method: bool def isolated_builtins(modules: dict) -> dict: + runtime = MAINTENANCE._fixture_modules["tier_switch.runtime"] + def safe_import(name: str, globals: dict | None = None, locals: dict | None = None, fromlist: tuple = (), level: int = 0) -> object: if name in modules: return modules[name] + if name == "truememory.tier_switch.runtime": + return runtime if name.startswith("truememory") or name.split(".")[0] in {"torch", "numpy", "hdbscan", "sqlite_vec"}: raise AssertionError("Unexpected production or native import: " + name) return builtins.__import__(name, globals, locals, fromlist, level) @@ -44,14 +49,8 @@ def safe_import(name: str, globals: dict | None = None, locals: dict | None = No def load_primitives() -> tuple[types.ModuleType, types.ModuleType]: - modules = {} - for name in ("storage", "_platform", "maintenance"): - module = types.ModuleType("synthetic_clustering_health_" + name) - module.__dict__["__builtins__"] = isolated_builtins(modules) - path = ROOT / "truememory" / (name + ".py") - exec(compile(path.read_text(encoding="utf-8"), str(path), "exec"), module.__dict__) - modules["truememory." + name] = module - return modules["truememory.storage"], modules["truememory.maintenance"] + loader = runpy.run_path(str(Path(__file__).with_name("test-maintenance-source-revision-753.py"))) + return loader["STORAGE"], loader["MAINTENANCE"] STORAGE, MAINTENANCE = load_primitives() @@ -71,7 +70,10 @@ def setUp(self) -> None: self.vector = types.SimpleNamespace( _lock=threading.Lock(), EMBEDDING_MODEL="synthetic-model", _embedding_dim=2, _active_tier_group=lambda: "edge", _active_vec_table=lambda conn: "synthetic_vectors", + _active_sep_table=lambda conn: "synthetic_separation", resolve_tier=lambda: "edge", + _frozen_embedding_target=None, _runtime_policy_tier=None, ) + self.enter_context(patch.dict(MAINTENANCE._fixture_modules, {"vector_search": self.vector})) self.versions = {"numpy": "synthetic", "hdbscan": "synthetic", "sqlite_vec": "synthetic"} self.absent: set[str] = set() self.probes = [] diff --git a/tests/test-installer-model-readiness-767.py b/tests/test-installer-model-readiness-767.py index eb343365..d207f50e 100644 --- a/tests/test-installer-model-readiness-767.py +++ b/tests/test-installer-model-readiness-767.py @@ -3,15 +3,18 @@ import argparse import ast +import builtins import contextlib import io import itertools import os +import runpy from pathlib import Path import shutil import subprocess import sys import tempfile +import threading import types import unittest from unittest.mock import patch @@ -229,15 +232,18 @@ class TestSelectedTierSetupReadiness(unittest.TestCase): def run_setup( self, tier: str, failures: set[str], fail_on_load: bool = False, *, cli: str = "", integration_outcomes: dict[str, bool] | None = None, + manager_action: str = "noop", ) -> tuple[int, str, str, list[dict], list[str]]: # Execute both real setup functions while replacing filesystem, adapter # installation, and model boundaries. Package __init__ is never imported. source = ast.parse((ROOT / "truememory" / "ingest" / "cli.py").read_text(encoding="utf-8")) functions = [node for node in source.body if isinstance(node, ast.FunctionDef) - and node.name in ("_setup_cli_integrations", "_run_setup")] + and node.name in ("_setup_cli_integrations", "_setup_model_readiness", "_run_setup")] program = ast.fix_missing_locations(ast.Module(body=functions, type_ignores=[])) saved: list[dict] = [] calls: list[str] = [] + config_path = Path("synthetic-setup/config.json") + config_state = {"tier": tier, "user_id": "synthetic-user"} class SyntheticModel: def encode(self, texts: list[str], **_kwargs: object) -> list[list[float]]: @@ -290,21 +296,88 @@ def install_cli(cli_id: str, user_id: str) -> bool: hooks_cli.install_cli = install_cli namespace = { "argparse": argparse, "os": os, "sys": sys, "_SAURON_BANNER": "synthetic banner", - "_load_truememory_config": lambda: {"tier": tier, "user_id": "synthetic-user"}, + "_load_truememory_config": lambda: config_state.copy(), "_save_truememory_config": lambda config: saved.append(dict(config)), - "_TRUEMEMORY_CONFIG_PATH": "synthetic-config.json", + "_TRUEMEMORY_CONFIG_PATH": config_path, } + bridge = runpy.run_path(str(ROOT / "tests/test-tier-runtime-bridge-795.py")) + fixture = bridge["TestRuntimeBridge"]() + fixture.setUp() + fixture.vector.resolve_tier = lambda: tier + fixture.vector.get_model = vectors.get_model + fixture.vector._encode_with_mps_fallback = vectors._encode_with_mps_fallback + fixture.reranker.get_reranker = reranker.get_reranker + primitives = runpy.run_path(str(ROOT / "tests/test-maintenance-source-revision-753.py")) + maintenance = primitives["MAINTENANCE"] + + class SyntheticManager: + _last_action = manager_action + + def run_rebuild_sync(self, requested_tier: str) -> bool: + self_check.assertEqual(requested_tier, tier) + vectors.set_embedding_model(requested_tier) + reranker.set_active_tier(requested_tier) + return True + + modules = { + "truememory.vector_search": fixture.vector, "truememory.reranker": fixture.reranker, + "truememory.maintenance": maintenance, + "truememory.tier_switch.runtime": fixture.api, + "truememory.tier_switch.manager": types.SimpleNamespace( + RebuildManager=SyntheticManager, _get_db_path=lambda: None, _open_db=lambda **kwargs: fixture.conn), + "truememory.hooks.registry": registry, "truememory.hooks.cli": hooks_cli, + } + + def synthetic_import(name, globals=None, locals=None, fromlist=(), level=0): + if name in modules: + return modules[name] + if name.startswith("truememory") or name in {"torch", "numpy", "sqlite_vec", "sentence_transformers"}: + raise AssertionError("Unexpected setup dependency: " + name) + return builtins.__import__(name, globals, locals, fromlist, level) + + namespace["__builtins__"] = dict(vars(builtins), __import__=synthetic_import) + projection = types.ModuleType("synthetic_setup_projection") + projection.__dict__["__builtins__"] = namespace["__builtins__"] + projection_path = ROOT / "truememory/tier_switch/projection.py" + exec(compile(projection_path.read_text(encoding="utf-8"), str(projection_path), "exec"), projection.__dict__) + config_file_lock = threading.Lock() + + @contextlib.contextmanager + def locked_config(path: Path, *, strict: bool = False): + self_check.assertEqual(path, config_path.with_name("config.json.lock")) + self_check.assertTrue(strict) + with config_file_lock: + yield + + def read_config(path: Path) -> dict: + self_check.assertEqual(path, config_path) + self_check.assertTrue(config_file_lock.locked()) + self_check.assertTrue(projection.CONFIG_WRITE_LOCK.locked()) + return config_state.copy() + + def write_config(path: Path, config: dict) -> None: + self_check.assertEqual(path, config_path) + self_check.assertTrue(config_file_lock.locked()) + self_check.assertTrue(projection.CONFIG_WRITE_LOCK.locked()) + config_state.clear() + config_state.update(config) + saved.append(config_state.copy()) + + projection.ConfigFileLock = locked_config + projection._read_config = read_config + projection._write_config = write_config + modules["truememory.tier_switch.projection"] = projection exec(compile(program, "synthetic-cli-setup", "exec"), namespace) stdout, stderr = io.StringIO(), io.StringIO() exit_code = 0 - with patch.dict(os.environ, {}, clear=True), patch.dict(sys.modules, { - "truememory": package, "truememory.vector_search": vectors, "truememory.reranker": reranker, - "truememory.hooks": hooks, "truememory.hooks.registry": registry, "truememory.hooks.cli": hooks_cli, - }), contextlib.redirect_stdout(stdout), contextlib.redirect_stderr(stderr): - try: - namespace["_run_setup"](argparse.Namespace(non_interactive=True, cli=cli)) - except SystemExit as exc: - exit_code = int(exc.code) + try: + with patch.dict(os.environ, {}, clear=True), contextlib.redirect_stdout(stdout), contextlib.redirect_stderr(stderr): + try: + namespace["_run_setup"](argparse.Namespace(non_interactive=True, cli=cli)) + except SystemExit as exc: + exit_code = int(exc.code) + finally: + fixture.doCleanups() return exit_code, stdout.getvalue(), stderr.getvalue(), saved, calls def test_selected_models_fail_truthfully_and_retain_configuration(self) -> None: @@ -369,6 +442,18 @@ def test_unknown_requested_integration_prevents_setup_completion(self) -> None: self.assertIn("install-cli-claude", calls) self.assertNotIn("install-cli-unknown", calls) + def test_config_only_setup_reports_no_unperformed_model_readiness(self) -> None: + code, output, error, saved, calls = self.run_setup("pro", set(), manager_action="config_only") + self.assertEqual(code, 0) + self.assertEqual(error, "") + self.assertEqual(saved, [{"tier": "pro", "user_id": "synthetic-user"}]) + self.assertNotIn("encode", calls) + self.assertNotIn("predict", calls) + self.assertNotIn("load-embed", calls) + self.assertNotIn("load-rerank", calls) + self.assertNotIn("Matryoshka ready", output) + self.assertNotIn("Cross-encoder reranker ready", output) + if __name__ == "__main__": unittest.main() diff --git a/tests/test-l0-style-accumulator-747.py b/tests/test-l0-style-accumulator-747.py index 037f61e2..0f9d85f6 100644 --- a/tests/test-l0-style-accumulator-747.py +++ b/tests/test-l0-style-accumulator-747.py @@ -723,7 +723,9 @@ def safe_import(name: str, globals: dict | None = None, locals: dict | None = No 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) + validator = next(node for node in tree.body + if isinstance(node, ast.FunctionDef) and node.name == "_validate_add_content") + exec(compile(ast.Module(body=[validator, add], type_ignores=[]), "engine.py", "exec"), namespace) for caller, directive in itertools.product((False, True), (False, True)): conn = self.connection(legacy=True, full=True) before = self.seed_legacy(conn) diff --git a/tests/test-l0-style-incremental-747.py b/tests/test-l0-style-incremental-747.py index 943bf64c..404e969d 100644 --- a/tests/test-l0-style-incremental-747.py +++ b/tests/test-l0-style-incremental-747.py @@ -31,6 +31,7 @@ def load_add() -> types.FunctionType: 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.tier_switch.runtime": MAINTENANCE._fixture_modules["truememory.tier_switch.runtime"], "truememory.tier_switch.writer": ACCUMULATOR["load_stdlib_source"]("tier_switch/writer")} def safe_import(name: str, globals: dict | None = None, locals: dict | None = None, @@ -47,7 +48,9 @@ def safe_import(name: str, globals: dict | None = None, locals: dict | None = No 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) + validator = next(node for node in tree.body + if isinstance(node, ast.FunctionDef) and node.name == "_validate_add_content") + exec(compile(ast.Module(body=[validator, add], type_ignores=[]), "actual-engine-add", "exec"), namespace) return namespace["add"] diff --git a/tests/test-l0-style-maintenance-747.py b/tests/test-l0-style-maintenance-747.py index f5a84483..5849674e 100644 --- a/tests/test-l0-style-maintenance-747.py +++ b/tests/test-l0-style-maintenance-747.py @@ -1,5 +1,6 @@ """Private style maintenance with stdlib source and synthetic in-memory SQLite.""" +import ast import builtins import json import runpy @@ -30,10 +31,45 @@ def safe_import(name: str, globals: dict | None = None, locals: dict | None = No raise AssertionError("Non-stdlib import forbidden: " + name) return builtins.__import__(name, globals, locals, fromlist, level) - module = types.ModuleType("synthetic_style_maintenance") - module.__dict__["__builtins__"] = dict(vars(builtins), __import__=safe_import) - source = ROOT / "truememory/maintenance.py" - exec(compile(source.read_text(encoding="utf-8"), str(source), "exec"), module.__dict__) + def load_source(name: str) -> types.ModuleType: + loaded = types.ModuleType("synthetic_style_boundary_" + name.replace("/", "_")) + loaded.__dict__["__builtins__"] = dict(vars(builtins), __import__=safe_import) + source = ROOT / "truememory" / (name + ".py") + missing = object() + previous = sys.modules.get(loaded.__name__, missing) + sys.modules[loaded.__name__] = loaded + try: + exec(compile(source.read_text(encoding="utf-8"), str(source), "exec"), loaded.__dict__) + finally: + if previous is missing: + sys.modules.pop(loaded.__name__, None) + else: + sys.modules[loaded.__name__] = previous + modules["truememory." + name.replace("/", ".")] = loaded + return loaded + + module = load_source("maintenance") + modules["psutil"] = types.SimpleNamespace() + for name in ("rebuild_source", "embedding_target", "tier_config", "tier_switch/cache", + "tier_switch/job", "tier_switch/source", "tier_switch/activation", "tier_switch/serving", + "tier_switch/writer"): + load_source(name) + modules["truememory.tier_switch"] = types.SimpleNamespace(serving=modules["truememory.tier_switch.serving"]) + vector = types.ModuleType("synthetic_style_vector_identity") + vector.__dict__.update(__builtins__=dict(vars(builtins), __import__=safe_import), sqlite3=sqlite3, + EMBEDDING_MODEL="model2vec", _embedding_dim=256, _model=None, _lock=threading.Lock(), + _frozen_embedding_target=None, _runtime_policy_tier=None, resolve_tier=lambda: "edge", + _cfg_get_model_group=modules["truememory.tier_config"].get_model_group) + tree = ast.parse((ROOT / "truememory/vector_search.py").read_text(encoding="utf-8")) + resolvers = [node for node in tree.body if isinstance(node, ast.FunctionDef) + and node.name in {"_active_tier_group", "_active_vec_table", "_active_sep_table"}] + exec(compile(ast.Module(body=resolvers, type_ignores=[]), "actual-style-vector-resolvers", "exec"), vector.__dict__) + modules["truememory.vector_search"] = vector + modules["truememory"] = types.SimpleNamespace(vector_search=vector, + reranker=types.SimpleNamespace(_model=None, get_current_reranker_name=lambda: "synthetic/reranker")) + runtime = load_source("tier_switch/runtime") + runtime.sys = types.SimpleNamespace(modules=modules) + module._fixture_modules = modules def import_module(name: str) -> types.ModuleType: if name != "truememory.personality_style_vec": diff --git a/tests/test-l0-style-routing-747.py b/tests/test-l0-style-routing-747.py index d20f1ec7..9e0c6914 100644 --- a/tests/test-l0-style-routing-747.py +++ b/tests/test-l0-style-routing-747.py @@ -42,8 +42,8 @@ def connection(self, *, bootstrap: bool = False) -> sqlite3.Connection: return conn def engine(self, conn: sqlite3.Connection) -> object: - modules = {"truememory.storage": STORAGE, "truememory.maintenance": MAINTENANCE, - "truememory.personality_style_vec": STYLE} + modules = MAINTENANCE._fixture_modules + modules["truememory.personality_style_vec"] = STYLE modules["truememory.rebuild_source"] = ROUTING["stdlib_module"]("rebuild_source", modules) module = ROUTING["load_engine"](modules) module._HAS_STYLE_VEC = True @@ -1007,17 +1007,28 @@ def load(connection: sqlite3.Connection, _path: Path) -> int: def test_deprecated_open_schedules_without_synchronous_style_build_or_hash_claim(self) -> None: conn = self.connection() + probe = self.connection() engine = self.engine(conn) module = engine._engine_test_module engine.db_path = (ROOT / "synthetic-existing-not-created.sqlite") - with patch.object(Path, "exists", return_value=True), \ - patch.object(module.sqlite3, "connect", return_value=conn), \ + exists = Path.exists + closed_owners = [] + owner_os = types.SimpleNamespace(getpid=os.getpid, fspath=os.fspath, close=closed_owners.append) + with patch.object(Path, "exists", lambda path: True if path == engine.db_path else exists(path)), \ + patch.object(module.sqlite3, "connect", side_effect=[probe, conn]) as opened, \ + patch.object(MAINTENANCE, "try_file_lock", return_value=719), \ + patch.object(MAINTENANCE, "os", owner_os), \ patch.object(engine, "_purge_legacy_entity_profile_summaries"), \ patch.object(engine, "_maybe_auto_consolidate") as planned, \ patch.object(STYLE, "_compute_entity_style_vectors", side_effect=AssertionError): with self.assertWarns(DeprecationWarning): self.assertIs(engine.open(rebuild_vectors=False), engine) planned.assert_called_once() + self.assertIs(engine.conn, conn) + self.assertEqual(opened.call_count, 2) + self.assertEqual(closed_owners, [719]) + with self.assertRaises(sqlite3.ProgrammingError): + probe.execute("SELECT 1") self.assertIsNone(conn.execute("SELECT value FROM metadata WHERE key='style_vec_hash_version'").fetchone()) self.assertIsNone(self.checkpoint(conn)) @@ -1273,7 +1284,7 @@ def close(self) -> None: successor.assert_called_once_with(MAINTENANCE.MaintenanceRequest(2, 30, True, False)) def test_late_worker_cleanup_cannot_replace_newer_manual_failure(self) -> None: - coordinator = MAINTENANCE.MaintenanceCoordinator(None) + coordinator = MAINTENANCE.MaintenanceCoordinator(ROOT / "synthetic-worker-cleanup-not-created.sqlite") conn = self.connection() events = [] class Connection: @@ -1283,7 +1294,9 @@ def close(self) -> None: events.append("close") token = coordinator._begin_observation() coordinator._finish_observation(token, None, "failed", "NewerFailure") - with patch.object(MAINTENANCE, "create_db", return_value=Connection()), \ + exists = Path.exists + with patch.object(Path, "exists", lambda path: False if path == coordinator.path else exists(path)), \ + patch.object(MAINTENANCE, "create_db", return_value=Connection()), \ patch.object(MAINTENANCE, "_release_owner", side_effect=lambda _fd: events.append("release")): coordinator._run(123, lambda *_args: MAINTENANCE.MaintenanceReport((), "SKIPPED")) self.assertEqual(events, ["close", "release"]) @@ -1292,7 +1305,7 @@ def close(self) -> None: def test_setup_and_close_failures_retain_report_and_release_after_cleanup(self) -> None: for boundary in ("setup", "close"): with self.subTest(boundary=boundary): - coordinator = MAINTENANCE.MaintenanceCoordinator(None) + coordinator = MAINTENANCE.MaintenanceCoordinator(ROOT / "synthetic-worker-failure-not-created.sqlite") old = MAINTENANCE.MaintenanceReport((), "SKIPPED") coordinator._last_report = old conn = self.connection() @@ -1303,7 +1316,9 @@ def __getattr__(self, name: str) -> object: def close(self) -> None: events.append("close") raise sqlite3.OperationalError("synthetic close") - with patch.object(MAINTENANCE, "create_db", side_effect=RuntimeError("synthetic setup") if boundary == "setup" else None, + exists = Path.exists + with patch.object(Path, "exists", lambda path: False if path == coordinator.path else exists(path)), \ + patch.object(MAINTENANCE, "create_db", side_effect=RuntimeError("synthetic setup") if boundary == "setup" else None, return_value=Connection()), \ patch.object(MAINTENANCE, "_release_owner", side_effect=lambda _fd: events.append("release")): coordinator._run(123, lambda *_args: MAINTENANCE.MaintenanceReport((), "SKIPPED new")) diff --git a/tests/test-maintenance-engine-routing-753.py b/tests/test-maintenance-engine-routing-753.py index 91238dc2..4138f079 100644 --- a/tests/test-maintenance-engine-routing-753.py +++ b/tests/test-maintenance-engine-routing-753.py @@ -2,6 +2,7 @@ import ast import builtins +import datetime import gc import logging import os @@ -16,7 +17,7 @@ import unittest import warnings import weakref -from contextlib import contextmanager, nullcontext +from contextlib import contextmanager from pathlib import Path from unittest.mock import patch @@ -38,35 +39,53 @@ def safe_import(name, globals=None, locals=None, fromlist=(), level=0): module.__dict__["__builtins__"] = dict(vars(builtins), __import__=safe_import) path = ROOT / "truememory" / (name + ".py") - with patch.dict(sys.modules, {module.__name__: module}): + missing = object() + previous = sys.modules.get(module.__name__, missing) + sys.modules[module.__name__] = module + try: exec(compile(path.read_text(encoding="utf-8"), str(path), "exec"), module.__dict__) + finally: + if previous is missing: + sys.modules.pop(module.__name__, None) + else: + sys.modules[module.__name__] = previous return module 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)) + # Runtime guards execute unchanged; only native identities are synthetic. + modules["psutil"] = types.SimpleNamespace() + for name in ("embedding_target", "tier_config", "tier_switch/cache", + "tier_switch/job", "tier_switch/source", "tier_switch/activation", + "tier_switch/serving"): + modules["truememory." + name.replace("/", ".")] = stdlib_module(name, modules) + serving = modules["truememory.tier_switch.serving"] + modules["truememory.tier_switch"] = types.SimpleNamespace(serving=serving) + vector = modules["truememory.vector_search"] + vector._frozen_embedding_target = None + vector._runtime_policy_tier = None + vector.resolve_tier = lambda: "edge" + vector_source = ast.parse((ROOT / "truememory/vector_search.py").read_text(encoding="utf-8")) + sep = next(node for node in vector_source.body + if isinstance(node, ast.FunctionDef) and node.name == "_active_sep_table") + exec(compile(ast.Module(body=[sep], type_ignores=[]), "actual-separation-resolver", "exec"), + vector.__dict__) + modules["truememory"] = types.SimpleNamespace( + vector_search=vector, + reranker=types.SimpleNamespace(get_current_reranker_name=lambda: "synthetic/reranker"), + ) + runtime = stdlib_module("tier_switch/runtime", modules) + runtime.sys = types.SimpleNamespace(modules=modules) + modules["truememory.tier_switch.runtime"] = runtime + vector.__dict__.update(__builtins__=runtime.__dict__["__builtins__"], + datetime=datetime, _model_generation=0) + metadata_nodes = [ast.ImportFrom(module="__future__", names=[ast.alias(name="annotations")], level=0)] + metadata_nodes.extend(node for node in vector_source.body + if isinstance(node, (ast.FunctionDef, ast.ClassDef)) + and node.name in {"_write_foreground_embedder_metadata_no_commit", "VectorPublicationChanged"}) + exec(compile(ast.fix_missing_locations(ast.Module(body=metadata_nodes, type_ignores=[])), + "actual-foreground-metadata", "exec"), vector.__dict__) module = types.ModuleType("synthetic_routing_engine") def safe_import(name, globals=None, locals=None, fromlist=(), level=0): @@ -80,6 +99,7 @@ def safe_import(name, globals=None, locals=None, fromlist=(), level=0): "__file__": str(ROOT / "truememory/engine.py"), "Path": Path, "sqlite3": sqlite3, "threading": threading, "time": time, "os": os, "warnings": warnings, "engine_operation": lambda function: function, + "engine_handle_operation": runtime.engine_handle_operation, "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")) @@ -112,6 +132,17 @@ def setUp(self) -> None: "sqlite_vec": types.SimpleNamespace(load=lambda conn: None)}) self.modules["truememory.rebuild_source"] = stdlib_module("rebuild_source", self.modules) self.engine_module = load_engine(self.modules) + original_import = MAINTENANCE.__dict__["__builtins__"]["__import__"] + + def maintenance_import(name, globals=None, locals=None, fromlist=(), level=0): + if name in self.modules: + return self.modules[name] + if name.startswith("truememory"): + raise ImportError("Synthetic boundary excludes native modules") + return original_import(name, globals, locals, fromlist, level) + + self.enterContext(patch.dict(MAINTENANCE.__dict__["__builtins__"], + {"__import__": maintenance_import})) create = STORAGE.create_db class ExtensionBoundary: diff --git a/tests/test-maintenance-fixture-isolation-753.py b/tests/test-maintenance-fixture-isolation-753.py index 9151bc6c..d2c438b6 100644 --- a/tests/test-maintenance-fixture-isolation-753.py +++ b/tests/test-maintenance-fixture-isolation-753.py @@ -103,6 +103,14 @@ def setUp(self) -> None: source = (ROOT / "truememory/maintenance.py").read_text(encoding="utf-8") exec(compile(source, "synthetic-foreign-maintenance", "exec"), self.foreign.__dict__) self.foreign.sys = self.process_sys + # This fixture invokes a worker directly for an in-memory coordinator. + # Production workers reopen file paths; give its real opener the explicit + # in-memory identity instead of the scheduler's non-launchable None. + class MemoryWorkerCoordinator(self.foreign.MaintenanceCoordinator): + def __init__(self, path: object) -> None: + super().__init__(":memory:" if path is None else path) + + self.foreign.MaintenanceCoordinator = MemoryWorkerCoordinator self.uri = "file:synthetic-fixture-" + uuid.uuid4().hex + "?mode=memory&cache=shared" self.connect = sqlite3.connect diff --git a/tests/test-maintenance-source-revision-753.py b/tests/test-maintenance-source-revision-753.py index 57d5dfd3..d15e89cf 100644 --- a/tests/test-maintenance-source-revision-753.py +++ b/tests/test-maintenance-source-revision-753.py @@ -1,6 +1,7 @@ """Transactional maintenance revisions on synthetic SQLite databases only.""" import builtins +import sys import sqlite3 import tempfile import types @@ -14,19 +15,57 @@ def load_primitives() -> tuple[types.ModuleType, types.ModuleType]: modules = {} - for name in ("storage", "_platform", "maintenance"): - module = types.ModuleType(f"synthetic_maintenance_{name}") - - def safe_import(name: str, globals: dict | None = None, locals: dict | None = None, - fromlist: tuple[str, ...] = (), level: int = 0) -> object: - if name in ("truememory.storage", "truememory._platform"): - return modules[name.rsplit(".", 1)[1]] - return builtins.__import__(name, globals, locals, fromlist, level) - + tier_switch = types.ModuleType("synthetic_maintenance_tier_switch") + + def unavailable_model(*args: object, **kwargs: object) -> object: + raise AssertionError("Maintenance fixtures must not construct or encode models") + + vector = types.SimpleNamespace( + EMBEDDING_MODEL="synthetic-model", _embedding_dim=2, _model=None, + _frozen_embedding_target=None, _runtime_policy_tier=None, + resolve_tier=lambda: "edge", _active_vec_table=lambda conn: "vec_messages", + _active_sep_table=lambda conn: "vec_messages_sep", get_model=unavailable_model, + ) + reranker = types.SimpleNamespace( + get_current_reranker_name=lambda: "synthetic/reranker", get_reranker=unavailable_model, + ) + + def safe_import(name: str, globals: dict | None = None, locals: dict | None = None, + fromlist: tuple[str, ...] = (), level: int = 0) -> object: + if name == "truememory": + return types.SimpleNamespace(vector_search=modules.get("vector_search", vector), + reranker=modules.get("reranker", reranker)) + if name == "truememory.tier_switch": + return tier_switch + if name.startswith("truememory."): + short = name.removeprefix("truememory.") + if short in modules: + return modules[short] + raise AssertionError("Unexpected application import: " + name) + if name == "psutil": + return types.SimpleNamespace() + if name.split(".")[0] in {"torch", "numpy", "sqlite_vec", "hdbscan", "model2vec", "sentence_transformers"}: + raise AssertionError("Unexpected native import: " + name) + return builtins.__import__(name, globals, locals, fromlist, level) + + for name in ("storage", "_platform", "maintenance", "embedding_target", "tier_config", + "tier_switch.cache", "rebuild_source", "tier_switch.job", "tier_switch.source", + "tier_switch.activation", "tier_switch.serving", "tier_switch.runtime", "tier_switch.writer"): + module = types.ModuleType("synthetic_maintenance_" + name.replace(".", "_")) module.__dict__["__builtins__"] = dict(vars(builtins), __import__=safe_import) - path = ROOT / "truememory" / f"{name}.py" - exec(compile(path.read_text(), str(path), "exec"), module.__dict__) + path = ROOT / "truememory" / (name.replace(".", "/") + ".py") + # Dataclasses inspect only their own explicitly named synthetic module. + with patch.dict(sys.modules, {module.__name__: module}): + exec(compile(path.read_text(encoding="utf-8"), str(path), "exec"), module.__dict__) modules[name] = module + if name.startswith("tier_switch."): + setattr(tier_switch, name.rsplit(".", 1)[1], module) + # Keep model identity injection local to this fixture graph, never sys.modules. + modules["vector_search"], modules["reranker"] = vector, reranker + modules["tier_switch.runtime"].sys = types.SimpleNamespace( + modules={"truememory." + name: module for name, module in modules.items()}) + modules["maintenance"]._fixture_modules = modules + modules["tier_switch.runtime"]._fixture_modules = modules return modules["storage"], modules["maintenance"] diff --git a/tests/test-rebuild-foreground-publication-756.py b/tests/test-rebuild-foreground-publication-756.py index d7259e32..18f1a924 100644 --- a/tests/test-rebuild-foreground-publication-756.py +++ b/tests/test-rebuild-foreground-publication-756.py @@ -29,7 +29,9 @@ def engine(self, conn=None): logger=logging.getLogger(__name__), MAX_CONTENT_LENGTH=10000, insert_message=self.modules["storage"].insert_message, update_message=self.modules["storage"].update_message) - exec(compile(ast.fix_missing_locations(ast.Module(body=[cls], type_ignores=[])), "actual-engine-boundaries", "exec"), namespace) + validators = [node for node in tree.body if isinstance(node, ast.FunctionDef) + and node.name == "_validate_add_content"] + exec(compile(ast.fix_missing_locations(ast.Module(body=[*validators, cls], type_ignores=[])), "actual-engine-boundaries", "exec"), namespace) engine = namespace["SyntheticEngine"]() engine.conn = conn or self.conn engine._write_lock = threading.Lock() diff --git a/tests/test-rebuild-source-streaming-756.py b/tests/test-rebuild-source-streaming-756.py index bbb36aa9..49143845 100644 --- a/tests/test-rebuild-source-streaming-756.py +++ b/tests/test-rebuild-source-streaming-756.py @@ -22,13 +22,44 @@ def load_modules(): + serving_source = ast.parse((ROOT / "truememory/tier_switch/serving.py").read_text(encoding="utf-8")) + serving_names = {"ServingLeaseError", "ServingLeaseTimeout", "ServingLeaseCancelled"} + serving_namespace = {} + exec(compile(ast.Module(body=[node for node in serving_source.body + if isinstance(node, ast.ClassDef) and node.name in serving_names], + type_ignores=[]), "actual-serving-errors", "exec"), serving_namespace) + runtime_source = ast.parse((ROOT / "truememory/tier_switch/runtime.py").read_text(encoding="utf-8")) + runtime_names = {"TierRuntimeError", "raise_if_serving_rejection", "_read_selection", "_read_policy", + "require_legacy_vector_mutation", "legacy_vector_mutation", + "open_serving_connection", "_open_serving_connection_at_path", "_modules", "_legacy_key"} + runtime_namespace = {"__builtins__": dict(vars(builtins)), "sqlite3": sqlite3, + "contextmanager": contextmanager, "serving": types.SimpleNamespace( + **{name: serving_namespace[name] for name in serving_names})} + 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) + exec(compile(ast.fix_missing_locations(ast.Module(body=runtime_nodes, type_ignores=[])), + "actual-runtime-guards", "exec"), runtime_namespace) + + def legacy_operation(conn: sqlite3.Connection) -> types.SimpleNamespace: + vector = modules["vector_search"] + return types.SimpleNamespace(selection=None, policy=None, + tables=(vector._active_vec_table(conn), vector._active_sep_table(conn)), + key=runtime_namespace["_legacy_key"]()) + 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)} + **{name: runtime_namespace[name] for name in runtime_names}, + serving_operation=lambda conn, **kwargs: nullcontext(legacy_operation(conn)), + current_operation=legacy_operation), + "reranker": types.SimpleNamespace(get_current_reranker_name=lambda: "synthetic/reranker")} def safe_import(name, globals=None, locals=None, fromlist=(), level=0): if name == "torch": raise ImportError("Native models are excluded from synthetic tests") + if name == "psutil": + return types.SimpleNamespace() + if name == "truememory": + return types.SimpleNamespace(vector_search=modules["vector_search"], reranker=modules["reranker"]) if name.startswith("truememory."): short = name.removeprefix("truememory.") if short in modules: @@ -36,20 +67,28 @@ 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", "tier_switch.writer"): + runtime_namespace["__builtins__"]["__import__"] = safe_import + for name in ("storage", "_platform", "maintenance", "rebuild_source", "tier_switch.writer", + "embedding_target", "tier_config", "tier_switch.cache", "tier_switch.job", + "tier_switch.source", "tier_switch.activation"): 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.replace(".", "/") + ".py") exec(compile(path.read_text(encoding="utf-8"), str(path), "exec"), module.__dict__) modules[name] = module + activation = modules["tier_switch.activation"] + runtime_namespace.update(_guard_storage=activation._guard_storage, _state=activation._state, + read_activation_state=activation.read_activation_state) vector = types.ModuleType("synthetic_stream_vector") vector.__dict__.update( - __builtins__=dict(vars(builtins), __import__=safe_import), + __builtins__=dict(vars(builtins), __import__=safe_import, Path=Path), database_operation=lambda function: function, contextmanager=contextmanager, closing=closing, Iterator=Iterator, replace=replace, datetime=datetime, + WriterCapture=modules["tier_switch.writer"].WriterCapture, sqlite3=sqlite3, logger=logging.getLogger(__name__), _lock=threading.Lock(), _model_generation=0, EMBEDDING_MODEL="synthetic-model", _embedding_dim=2, + resolve_tier=lambda: "edge", _BUILD_VECTORS_TXN_BATCH=100, _BUILD_STATE_KEY_PREFIX="vec_build_state:", _get_batch_size=lambda: 2, _flush_mps_cache=lambda: None, _active_vec_table=lambda conn: "vec_messages", _active_sep_table=lambda conn: "vec_messages_sep", @@ -60,12 +99,15 @@ def safe_import(name, globals=None, locals=None, fromlist=(), level=0): "build_vectors", "build_separation_vectors", "_build_sep_text", "_ensure_metadata_table", "_build_state_key", "_mark_build_in_progress", "_clear_build_in_progress", "_build_in_progress", "_write_embedder_metadata_no_commit", "_write_embedder_metadata", + "_write_foreground_embedder_metadata_no_commit", "_foreground_vector_publication", "_foreground_model_fence", "_validate_foreground_vector_target", "VectorPublicationChanged", "embed_single"} tree = ast.parse((ROOT / "truememory/vector_search.py").read_text(encoding="utf-8")) body = [node for node in tree.body if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef)) and node.name in selected] exec(compile(ast.Module(body=body, type_ignores=[]), "actual-vector-builders", "exec"), vector.__dict__) modules["vector_search"] = vector + runtime_namespace["sys"] = types.SimpleNamespace( + modules={"truememory." + name: module for name, module in modules.items()}) return modules diff --git a/tests/test-tier-active-initialization-795.py b/tests/test-tier-active-initialization-795.py new file mode 100644 index 00000000..0ec9020a --- /dev/null +++ b/tests/test-tier-active-initialization-795.py @@ -0,0 +1,465 @@ +"""Live automatic-worker initialization receipts, with stdlib and memory SQLite.""" +from __future__ import annotations + +import contextlib +import ast +import logging +import os +from pathlib import Path +import runpy +import re +import sqlite3 +import threading +import types +import unittest +import uuid +import weakref +from typing import NamedTuple +from unittest.mock import Mock, patch + +BOUNDARIES = runpy.run_path(str(Path(__file__).with_name("test-tier-public-boundaries-795.py"))) +definitions = BOUNDARIES["definitions"] + + +class TestActiveInitialization(unittest.TestCase): + def setUp(self): + self.fixture = BOUNDARIES["TestEngineReconnectReceipt"]() + self.fixture.setUp() + self.addCleanup(self.fixture.doCleanups) + self.first = self.fixture.engine() + self.api = self.fixture.api + self.closed_fds = [] + self.worker_conn = self.fixture.memory_connection() + self.entered, self.release = threading.Event(), threading.Event() + thread_api = types.SimpleNamespace(Thread=threading.Thread, Lock=threading.Lock, + Event=threading.Event, local=threading.local) + names = {"MaintenanceBusyError", "MaintenanceOwnership", "MaintenanceRequest", "MaintenanceCoordinator", + "MaintenanceReport", "LayerResult", "_StyleOpenRejected", "_style_error_category", + "canonical_database_path", "_bind_owner", "maintenance_owner", "_acquire_owner", "_release_owner", + "_after_fork", "get_coordinator", "_valid_initialization_receipt", "_active_initialization_receipt"} + self.maintenance = definitions("maintenance.py", names, self.fixture.namespace( + Path=Path, re=re, os=types.SimpleNamespace(getpid=os.getpid, close=self.closed_fds.append, PathLike=os.PathLike, fspath=os.fspath), + uuid=uuid, weakref=weakref, threading=thread_api, sqlite3=sqlite3, NamedTuple=NamedTuple, + contextmanager=contextlib.contextmanager, _registry_lock=threading.Lock(), + _coordinators=weakref.WeakValueDictionary(), _held_fds=set(), _held_paths={}, _owner_waiters={}, + _thread_owners=threading.local(), try_file_lock=lambda path: 71, + create_db=lambda path: self.worker_conn, logger=logging.getLogger("synthetic-active-receipt"))) + self.fixture.modules["maintenance"] = self.maintenance + storage = ast.parse((BOUNDARIES["ROOT"] / "truememory/storage.py").read_text(encoding="utf-8")) + self.fixture.modules["storage"].DEFAULT_BUSY_TIMEOUT_MS = next( + ast.literal_eval(node.value) for node in storage.body if isinstance(node, ast.Assign) + and any(isinstance(target, ast.Name) and target.id == "DEFAULT_BUSY_TIMEOUT_MS" for target in node.targets)) + self.coordinator = self.maintenance.get_coordinator(self.first.db_path) + self.fixture.engine_namespace["_HAS_HYBRID"] = True + self.maintenance.run_engine_maintenance = self.work + self.fixture.enterContext(patch.object(Path, "exists", return_value=False)) + self.addCleanup(self.finish) + + def work(self, conn, coordinator, **kwargs): + self.entered.set() + if not self.release.wait(3): + raise AssertionError("Synthetic automatic worker was not released") + return None + + def finish(self): + self.release.set() + self.coordinator.wait(3) + + def start(self, receipt=True): + values = {"_initialization_receipt": self.first._reconnect_receipt} if receipt else {} + self.assertTrue(self.coordinator.request_layers(**values)) + self.assertTrue(self.entered.wait(2)) + self.assertEqual(self.coordinator.status[0], "running") + + def fresh(self): + engine = self.first.__class__() + engine.conn, engine.db_path = None, self.first.db_path + engine._init_lock = threading.Lock() + engine._write_lock = self.api.ConnectionWriteLock(engine) + engine.ready = engine._runtime_initialized = engine._has_vectors = False + engine._has_hybrid = True + engine._runtime_vector_connection = engine._runtime_vector_generation = None + engine._reconnect_receipt = None + return engine + + def record(self): + with self.maintenance._active_initialization_receipt(self.first.db_path) as value: + return value + + def test_fresh_engine_adopts_only_validated_active_automatic_receipt(self): + self.start() + record = self.record() + self.assertIs(record[0](), self.coordinator) + self.assertEqual(record[2], self.first._reconnect_receipt) + engine = self.fresh() + candidate = self.fixture.candidate(self.first) + queries = [] + candidate.set_trace_callback(queries.append) + engine._open_connection_handle() + self.assertIs(engine.conn, candidate) + self.assertTrue(engine.ready) + self.assertTrue(engine._runtime_initialized) + self.assertFalse(engine._has_vectors) + self.assertFalse(engine._has_hybrid) + self.assertEqual(engine._reconnect_receipt, self.first._reconnect_receipt) + self.assertFalse(candidate.in_transaction) + self.assertFalse(self.release.is_set()) + self.assertEqual(self.maintenance._held_paths[self.first.db_path], 71) + self.assertFalse(any(q.lstrip().split(" ", 1)[0].upper() in {"INSERT", "UPDATE", "DELETE", "CREATE", "DROP", "ALTER"} + for q in queries)) + self.assertEqual(self.fixture.fixture.loads, []) + + def test_manual_or_receiptless_automatic_owner_remains_busy(self): + for automatic in (False, True): + with self.subTest(automatic=automatic): + if automatic: + self.start(receipt=False) + else: + self.assertTrue(self.coordinator.request(lambda conn, cancel: self.work(conn, self.coordinator))) + self.assertTrue(self.entered.wait(2)) + self.assertIsNone(self.record()) + engine = self.fresh() + self.fixture.candidate(self.first) + with self.assertRaises(self.maintenance.MaintenanceBusyError): + engine._open_connection_handle() + self.assertIsNone(engine.conn) + self.assertFalse(engine.ready) + self.fixture.connect.assert_not_called() + self.finish() + self.entered.clear() + self.release.clear() + self.worker_conn = self.fixture.memory_connection() + + def test_worker_finish_revokes_before_descriptor_release(self): + self.start() + original = self.maintenance._release_owner + observations = [] + def release(fd): + observations.append(self.record()) + self.assertIsNone(self.coordinator._active_initialization) + return original(fd) + self.maintenance._release_owner = release + self.finish() + self.assertEqual(observations, [None]) + self.assertIsNone(self.record()) + self.assertIsNone(self.coordinator._pending_initialization) + self.assertIsNone(self.coordinator._reserved_initialization) + + def test_cancel_during_validation_does_not_install_or_close_old_engine_state(self): + self.start() + engine = self.fresh() + candidate = self.fixture.candidate(self.first) + original = engine._reopen_initialized_connection + def validate(path, **kwargs): + result = original(path, **kwargs) + self.coordinator.cancel() + return result + engine._reopen_initialized_connection = validate + with self.assertRaises(self.maintenance.MaintenanceBusyError): + engine._open_connection_handle() + self.fixture.assert_closed(candidate) + self.assertIsNone(engine.conn) + self.assertFalse(engine.ready) + self.assertIsNone(engine._reconnect_receipt) + self.assertIsNone(self.record()) + self.assertEqual(self.maintenance._held_paths[self.first.db_path], 71) + + def test_same_fd_different_launch_or_coordinator_cannot_revalidate_old_record(self): + self.start() + old = self.record() + with self.coordinator._mutex: + token, receipt = self.coordinator._active_initialization + self.coordinator._active_initialization = ((token[0], token[1], "different-launch"), receipt) + with self.maintenance._active_initialization_receipt(self.first.db_path, expected=old[:2]) as value: + self.assertIsNone(value) + replacement = self.maintenance.MaintenanceCoordinator(self.first.db_path) + with self.maintenance._registry_lock: + self.maintenance._coordinators[self.first.db_path] = replacement + with self.maintenance._active_initialization_receipt(self.first.db_path, expected=old[:2]) as value: + self.assertIsNone(value) + with self.maintenance._registry_lock: + self.maintenance._coordinators[self.first.db_path] = self.coordinator + + def test_receipt_validation_does_not_hold_registry_or_coordinator_mutex(self): + self.start() + engine = self.fresh() + self.fixture.candidate(self.first) + original = engine._reopen_initialized_connection + def validate(path, **kwargs): + self.assertTrue(self.maintenance._registry_lock.acquire(blocking=False)) + self.maintenance._registry_lock.release() + self.assertTrue(self.coordinator._mutex.acquire(blocking=False)) + self.coordinator._mutex.release() + self.assertFalse(engine._write_lock.locked()) + self.assertIsNone(engine._reconnect_receipt) + self.assertEqual(kwargs["receipt"], self.first._reconnect_receipt) + return original(path, **kwargs) + engine._reopen_initialized_connection = validate + engine._open_connection_handle() + self.assertTrue(engine.ready) + + def test_exact_pending_generation_and_receiptless_request_do_not_revive_evidence(self): + self.start() + first_generation = self.coordinator._notification_generation + self.assertFalse(self.coordinator.request_layers(_initialization_receipt=self.first._reconnect_receipt)) + self.assertEqual(self.coordinator._pending_initialization[0], first_generation + 1) + self.assertFalse(self.coordinator.request_layers(threshold=37)) + self.assertIsNone(self.coordinator._pending_initialization) + self.assertEqual(self.coordinator._pending.generation, first_generation + 2) + self.assertEqual(self.coordinator._pending.threshold, 37) + self.coordinator.cancel() + self.assertIsNone(self.coordinator._active_initialization) + self.assertIsNone(self.coordinator._pending_initialization) + self.assertIsNone(self.coordinator._reserved_initialization) + + def test_busy_launch_drops_evidence_and_keeps_exact_pending_work(self): + self.maintenance._held_paths[self.first.db_path] = 91 + try: + self.assertFalse(self.coordinator.request_layers(_initialization_receipt=self.first._reconnect_receipt)) + self.assertEqual(self.coordinator._pending.generation, 1) + self.assertIsNone(self.coordinator._active_initialization) + self.assertIsNone(self.coordinator._pending_initialization) + self.assertIsNone(self.coordinator._reserved_initialization) + self.assertIsNone(self.record()) + finally: + self.maintenance._held_paths.pop(self.first.db_path) + self.coordinator.cancel() + + def test_invalid_receipts_are_not_retained_or_exposed(self): + receipt = self.first._reconnect_receipt + invalid = [object(), ("/other/database", *receipt[1:]), (*receipt[:4], ("legacy", object(), "256", "edge", "rr"), *receipt[5:]), + (*receipt[:4], ("legacy", "x" * 513, "256", "edge", "rr"), *receipt[5:]), + (*receipt[:4], ("legacy", "\ud800", "256", "edge", "rr"), *receipt[5:])] + self.start() + for value in invalid: + with self.subTest(kind=type(value).__name__): + self.assertFalse(self.coordinator.request_layers(_initialization_receipt=value)) + self.assertIsNone(self.coordinator._pending_initialization) + self.coordinator.cancel() + + def test_start_failure_revokes_before_owner_release(self): + releases = [] + original = self.maintenance._release_owner + def release(fd): + releases.append(self.record()) + return original(fd) + self.maintenance._release_owner = release + self.maintenance.threading.Thread = lambda **kwargs: types.SimpleNamespace(start=Mock(side_effect=RuntimeError("synthetic start failure"))) + with self.assertRaises(RuntimeError): + self.coordinator.request_layers(_initialization_receipt=self.first._reconnect_receipt) + self.assertEqual(releases, [None]) + self.assertIsNone(self.coordinator._active_initialization) + self.assertIsNone(self.coordinator._pending_initialization) + self.assertIsNone(self.coordinator._reserved_initialization) + self.assertFalse(self.coordinator._active) + + def test_fork_reset_clears_inherited_evidence_and_registry(self): + self.coordinator._active = True + self.coordinator._automatic_initialization = True + self.coordinator._active_initialization = ((os.getpid(), 71, "synthetic-live"), self.first._reconnect_receipt) + self.coordinator._pending_initialization = (1, self.first._reconnect_receipt) + self.coordinator._reserved_initialization = (1, self.first._reconnect_receipt) + self.maintenance._held_paths[self.first.db_path] = 71 + self.maintenance._held_fds.add(71) + self.maintenance._after_fork() + self.assertIsNone(self.coordinator._active_initialization) + self.assertIsNone(self.coordinator._pending_initialization) + self.assertIsNone(self.coordinator._reserved_initialization) + self.assertEqual(len(self.maintenance._coordinators), 0) + self.assertEqual(self.maintenance._held_paths, {}) + self.assertEqual(self.closed_fds, [71]) + self.assertIsNone(self.record()) + self.coordinator._active = False + + def test_admitted_model_or_dimension_change_refuses_ready_receipt_but_deep_override_remains_valid(self): + self.start() + engine = self.fresh() + self.fixture.candidate(self.first) + engine._open_connection_handle() + module = definitions("engine.py", {"_initialize_connection"}, self.fixture.namespace(), methods=True) + for key in (("legacy", "qwen3_256", "256", "base", "synthetic/default"), + ("legacy", "model2vec", "384", "edge", "synthetic/default")): + operation = types.SimpleNamespace(selection=None, policy=None, key=key) + with patch.object(self.api, "current_operation", return_value=operation): + with self.assertRaises(self.api.TierRuntimeError): + module._initialize_connection(engine, _suppress_maintenance=True) + for key in (self.first._reconnect_receipt[4], ("legacy", "model2vec", "256", "edge", "synthetic/deep")): + operation = types.SimpleNamespace(selection=None, policy=None, key=key) + with patch.object(self.api, "current_operation", return_value=operation), \ + patch.object(engine.conn, "execute", side_effect=AssertionError("Ready check performed SQL")), \ + patch.object(Path, "stat", side_effect=AssertionError("Ready check stat")): + module._initialize_connection(engine, _suppress_maintenance=True) + self.assertEqual(self.fixture.fixture.loads, []) + + + def test_receiptless_successor_cannot_reuse_previous_live_owner_record(self): + second_conn = self.fixture.memory_connection() + connections = iter((self.worker_conn, second_conn)) + self.maintenance.create_db = lambda path: next(connections) + second_entered, second_release = threading.Event(), threading.Event() + calls = [] + def work(conn, coordinator, **kwargs): + calls.append(conn) + if len(calls) == 1: + self.entered.set() + if not self.release.wait(3): + raise AssertionError("First synthetic worker release timed out") + else: + second_entered.set() + if not second_release.wait(3): + raise AssertionError("Successor synthetic worker release timed out") + self.maintenance.run_engine_maintenance = work + self.start() + old = self.record() + self.assertFalse(self.coordinator.request_layers()) + self.release.set() + try: + self.assertTrue(second_entered.wait(2)) + self.assertEqual(len(calls), 2) + self.assertIsNone(self.record()) + with self.maintenance._active_initialization_receipt(self.first.db_path, expected=old[:2]) as observed: + self.assertIsNone(observed) + self.assertIsNone(self.coordinator._pending_initialization) + self.assertIsNone(self.coordinator._reserved_initialization) + finally: + second_release.set() + self.assertTrue(self.coordinator.wait(3)) + + def test_actual_serving_admission_rechecks_embedding_after_candidate_validation(self): + self.start() + engine = self.fresh() + candidate = self.fixture.candidate(self.first) + original = engine._reopen_initialized_connection + vector = self.fixture.fixture.vector + def validate(path, **kwargs): + result = original(path, **kwargs) + vector.EMBEDDING_MODEL = "qwen3_256" + return result + engine._reopen_initialized_connection = validate + engine._open_connection_handle() + self.assertIs(engine.conn, candidate) + module = definitions("engine.py", {"_initialize_connection"}, self.fixture.namespace(), methods=True) + with self.api.serving_operation(candidate): + with self.assertRaises(self.api.TierRuntimeError): + module._initialize_connection(engine, _suppress_maintenance=True) + vector.EMBEDDING_MODEL = "model2vec" + vector._embedding_dim = 384 + with self.api.serving_operation(candidate): + with self.assertRaises(self.api.TierRuntimeError): + module._initialize_connection(engine, _suppress_maintenance=True) + vector._embedding_dim = 256 + self.fixture.fixture.reranker._model_name = "synthetic/deep" + with self.api.serving_operation(candidate, reranker_id="synthetic/deep"): + with patch.object(candidate, "execute", side_effect=AssertionError("Ready admission performed SQL")): + module._initialize_connection(engine, _suppress_maintenance=True) + self.assertEqual(self.fixture.fixture.loads, []) + + + + def test_failed_start_preserves_only_newer_pending_generation_receipt(self): + for donate in (True, False): + with self.subTest(donate=donate): + start_entered, allow_failure = threading.Event(), threading.Event() + attempts, errors = [], [] + receipt = tuple(list(self.first._reconnect_receipt)) + class FailedStart: + def start(inner): + start_entered.set() + if not allow_failure.wait(3): + raise AssertionError("Synthetic start interleaving timed out") + raise RuntimeError("synthetic first start failed") + def thread_factory(**kwargs): + attempts.append(kwargs) + return FailedStart() if len(attempts) == 1 else threading.Thread(**kwargs) + self.maintenance.threading.Thread = thread_factory + def request_first(): + try: + self.coordinator.request_layers(threshold=31, + _initialization_receipt=self.first._reconnect_receipt) + except BaseException as error: + errors.append(error) + requester = threading.Thread(target=request_first) + requester.start() + try: + self.assertTrue(start_entered.wait(2)) + values = {"_initialization_receipt": receipt} if donate else {} + self.assertFalse(self.coordinator.request_layers(threshold=32, **values)) + generation = self.coordinator._pending.generation + if donate: + self.assertEqual(self.coordinator._pending_initialization, (generation, receipt)) + else: + self.assertIsNone(self.coordinator._pending_initialization) + allow_failure.set() + requester.join(2) + self.assertFalse(requester.is_alive()) + self.assertEqual(len(errors), 1) + self.assertIs(type(errors[0]), RuntimeError) + self.assertTrue(self.entered.wait(2)) + offered = self.record() + if donate: + self.assertIsNotNone(offered) + self.assertIs(offered[2], receipt) + else: + self.assertIsNone(offered) + self.assertEqual(len(attempts), 2) + self.assertEqual(self.coordinator._notification_generation, generation) + self.assertIsNone(self.coordinator._pending_initialization) + self.assertIsNone(self.coordinator._reserved_initialization) + finally: + allow_failure.set() + self.release.set() + requester.join(3) + self.assertTrue(self.coordinator.wait(3)) + self.entered.clear() + self.release.clear() + self.worker_conn = self.fixture.memory_connection() + + def test_busy_acquisition_preserves_only_surviving_pending_generation_receipt(self): + for donate in (True, False): + with self.subTest(donate=donate): + acquire_entered, allow_busy = threading.Event(), threading.Event() + results, errors = [], [] + receipt = tuple(list(self.first._reconnect_receipt)) + def file_lock(path): + acquire_entered.set() + if not allow_busy.wait(3): + raise AssertionError("Synthetic owner interleaving timed out") + return None + self.maintenance.try_file_lock = file_lock + def request_first(): + try: + results.append(self.coordinator.request_layers(threshold=31, + _initialization_receipt=self.first._reconnect_receipt)) + except BaseException as error: + errors.append(error) + requester = threading.Thread(target=request_first) + requester.start() + try: + self.assertTrue(acquire_entered.wait(2)) + values = {"_initialization_receipt": receipt} if donate else {} + self.assertFalse(self.coordinator.request_layers(threshold=32, **values)) + generation = self.coordinator._pending.generation + allow_busy.set() + requester.join(2) + self.assertFalse(requester.is_alive()) + self.assertEqual(errors, []) + self.assertEqual(results, [False]) + self.assertEqual(self.coordinator._pending.generation, generation) + self.assertEqual(self.coordinator._pending.threshold, 32) + if donate: + self.assertEqual(self.coordinator._pending_initialization, (generation, receipt)) + self.assertIs(self.coordinator._pending_initialization[1], receipt) + else: + self.assertIsNone(self.coordinator._pending_initialization) + self.assertIsNone(self.coordinator._active_initialization) + self.assertIsNone(self.coordinator._reserved_initialization) + self.assertIsNone(self.record()) + finally: + allow_busy.set() + requester.join(3) + self.coordinator.cancel() + + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test-tier-foreground-native-795.py b/tests/test-tier-foreground-native-795.py new file mode 100644 index 00000000..bf2e964e --- /dev/null +++ b/tests/test-tier-foreground-native-795.py @@ -0,0 +1,664 @@ +"""Opt-in Linux sqlite-vec publication under another thread's maintenance owner.""" + +from __future__ import annotations + +from collections.abc import Iterator +from contextlib import contextmanager +import hashlib +import os +from pathlib import Path +import struct +import sys +import threading +from types import ModuleType + +import pytest + + +pytestmark = pytest.mark.skipif( + sys.platform != "linux" or os.environ.get("TRUEMEMORY_TEST_NATIVE_VEC") != "1", + reason="Native foreground regression requires Linux and TRUEMEMORY_TEST_NATIVE_VEC=1", +) + + +class SyntheticEncoder: + def __init__(self) -> None: + self.calls: list[tuple[str, ...]] = [] + + @staticmethod + def vector(text: str) -> tuple[float, ...]: + digest = hashlib.sha256(text.encode("utf-8")).digest() + # 256 * (1 / 16) ** 2 = 1, including the deterministic sign changes. + return tuple(1 / 16 if digest[index // 8] & (1 << (index % 8)) else -1 / 16 + for index in range(256)) + + def encode(self, texts: list[str], **kwargs: object) -> object: + import numpy as np + + self.calls.append(tuple(texts)) + return np.asarray([self.vector(text) for text in texts], dtype=np.float32).reshape(len(texts), 256) + + +@pytest.fixture +def native_engine(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> Iterator[tuple[object, ModuleType, SyntheticEncoder]]: + # Keep the guard before every application/native import and database open. + if sys.platform != "linux" or os.environ.get("TRUEMEMORY_TEST_NATIVE_VEC") != "1": + pytest.skip("Native foreground regression is opt-in on Linux") + monkeypatch.setenv("TRUEMEMORY_EMBED_MODEL", "edge") + monkeypatch.setenv("TRUEMEMORY_MODEL_SERVER", "off") + from truememory import reranker, vector_search + from truememory.engine import TrueMemoryEngine + from truememory.maintenance import get_coordinator + + vector_search.set_embedding_model("edge") + reranker.set_active_tier("edge") + encoder = SyntheticEncoder() + monkeypatch.setattr(vector_search, "_model", encoder) + engine = TrueMemoryEngine(tmp_path / "synthetic-foreground.sqlite") + coordinator = get_coordinator(engine.db_path) + try: + # Setup alone suppresses scheduling; all tested operations run normally. + engine._ensure_connection(_suppress_maintenance=True) + assert engine.ready and engine._has_vectors + assert encoder.calls == [] + assert engine.conn.execute("SELECT COUNT(*) FROM messages").fetchone()[0] == 0 + for table in (vector_search._active_vec_table(engine.conn), vector_search._active_sep_table(engine.conn)): + sql = engine.conn.execute("SELECT sql FROM sqlite_master WHERE name=?", (table,)).fetchone()[0] + assert "using vec0" in sql.lower() + assert engine.conn.execute(f'SELECT COUNT(*) FROM "{table}"').fetchone()[0] == 0 + yield engine, vector_search, encoder + finally: + # close() only detaches the Engine. Stop and join its shared worker first. + coordinator.cancel() + try: + assert coordinator.wait(10), "Shared maintenance teardown did not finish" + finally: + engine.close() + + +@contextmanager +def held_maintenance(path: Path) -> Iterator[threading.Thread]: + from truememory.maintenance import MaintenanceBusyError, maintenance_owner + + acquired = threading.Event() + release = threading.Event() + errors: list[BaseException] = [] + + def hold() -> None: + try: + with maintenance_owner(path): + acquired.set() + if not release.wait(30): + raise AssertionError("Foreground operation outlived the maintenance-owner deadline") + except BaseException as error: + errors.append(error) + acquired.set() + + worker = threading.Thread(target=hold, name="synthetic-foreground-maintenance-owner") + worker.start() + try: + assert acquired.wait(10), "Maintenance owner was not acquired" + assert not errors, errors + with pytest.raises(MaintenanceBusyError): + with maintenance_owner(path): + pytest.fail("A second thread acquired the existing maintenance ownership") + yield worker + assert worker.is_alive(), "Maintenance ownership ended before foreground assertions" + finally: + release.set() + worker.join(10) + assert not worker.is_alive(), "Maintenance owner did not join" + assert not errors, errors + + +def assert_publication(engine: object, vector: ModuleType, encoder: SyntheticEncoder, + message_id: int, content: str) -> tuple[bytes, bytes]: + row = engine.conn.execute( + "SELECT content,sender,recipient,timestamp FROM messages WHERE id=?", (message_id,), + ).fetchone() + assert row is not None and row[0] == content + separation = vector._build_sep_text(row[1], row[2], row[3], content) + blobs = [] + for table, text in ((vector._active_vec_table(engine.conn), content), + (vector._active_sep_table(engine.conn), separation)): + rows = engine.conn.execute(f'SELECT rowid,embedding FROM "{table}" ORDER BY rowid').fetchall() + assert len(rows) == 1 and rows[0][0] == message_id + blob = rows[0][1] + assert len(blob) == 256 * 4 + values = struct.unpack("256f", blob) + assert values == encoder.vector(text) + assert sum(value * value for value in values) == 1 + blobs.append(blob) + metadata = dict(engine.conn.execute( + "SELECT key,value FROM metadata WHERE key IN ('embed_model','embed_dim')", + ).fetchall()) + assert metadata == {"embed_model": vector.EMBEDDING_MODEL, "embed_dim": "256"} + assert not engine.conn.in_transaction + return blobs[0], blobs[1] + + +def seed_source(engine: object, content: str) -> int: + cursor = engine.conn.execute( + "INSERT INTO messages(content,sender,recipient,timestamp) VALUES(?,?,?,?)", + (content, "synthetic-sender", "synthetic-recipient", "2026-01-02"), + ) + engine.conn.commit() + return cursor.lastrowid + + +def test_first_add_publishes_vectors_and_metadata_under_maintenance(native_engine: tuple) -> None: + engine, vector, encoder = native_engine + with held_maintenance(engine.db_path): + result = engine.add("synthetic-first-add", sender="synthetic-sender", + recipient="synthetic-recipient", timestamp="2026-01-02") + assert_publication(engine, vector, encoder, result["id"], "synthetic-first-add") + assert len(encoder.calls) == 2 + + +def test_update_replaces_both_vectors_under_maintenance(native_engine: tuple) -> None: + engine, vector, encoder = native_engine + message_id = seed_source(engine, "synthetic-before-update") + vector.embed_single(engine.conn, message_id, "synthetic-before-update") + before = assert_publication(engine, vector, encoder, message_id, "synthetic-before-update") + call_count = len(encoder.calls) + with held_maintenance(engine.db_path): + result = engine.update(message_id, content="synthetic-after-update") + assert result is not None and result["id"] == message_id + after = assert_publication(engine, vector, encoder, message_id, "synthetic-after-update") + assert before[0] != after[0] and before[1] != after[1] + assert len(encoder.calls) == call_count + 2 + + +def test_direct_embed_single_publishes_known_row_under_maintenance(native_engine: tuple) -> None: + engine, vector, encoder = native_engine + message_id = seed_source(engine, "synthetic-direct-embed") + with held_maintenance(engine.db_path): + vector.embed_single(engine.conn, message_id, "synthetic-direct-embed") + assert_publication(engine, vector, encoder, message_id, "synthetic-direct-embed") + assert len(encoder.calls) == 2 + + +def test_fresh_engine_opens_during_actual_automatic_worker( + native_engine: tuple, monkeypatch: pytest.MonkeyPatch, +) -> None: + from model2vec import StaticModel + from truememory import maintenance + from truememory.engine import TrueMemoryEngine + + first, vector, encoder = native_engine + content = "synthetic-automatic-worker-open" + message_id = seed_source(first, content) + vector.embed_single(first.conn, message_id, content) + before = assert_publication(first, vector, encoder, message_id, content) + encode_count = len(encoder.calls) + coordinator = first._get_maintenance_coordinator() + second = TrueMemoryEngine(first.db_path) + entered = threading.Event() + release = threading.Event() + failures: list[str] = [] + completed: list[bool] = [] + constructor_calls: list[bool] = [] + real_run = maintenance.run_engine_maintenance + + def blocked_run(conn: object, owner: object, **kwargs: object) -> object: + try: + assert conn is not first.conn + assert owner is coordinator + assert conn.execute("SELECT COUNT(*) FROM messages").fetchone()[0] == 1 + entered.set() + if not release.wait(30): + raise AssertionError("Fresh Engine outlived the automatic-worker barrier") + result = real_run(conn, owner, **kwargs) + completed.append(True) + return result + except BaseException as error: + failures.append(type(error).__name__) + entered.set() + raise + + def forbidden_constructor(*args: object, **kwargs: object) -> object: + constructor_calls.append(True) + raise AssertionError("Opening a validated database must not construct an embedding model") + + def value_only(value: object) -> bool: + if type(value) is tuple: + return all(value_only(item) for item in value) + return type(value) in (type(None), bool, int, float, str) + + monkeypatch.setattr(maintenance, "run_engine_maintenance", blocked_run) + try: + first._maybe_startup_consolidate() + assert entered.wait(10), "Automatic worker did not reach its opened-connection barrier" + assert not failures, failures + assert coordinator.snapshot()["active"] + with pytest.raises(maintenance.MaintenanceBusyError): + with maintenance.maintenance_owner(first.db_path): + pytest.fail("Automatic worker did not retain real maintenance ownership") + + database_path = first.db_path.resolve() + with maintenance._active_initialization_receipt(database_path) as offered: + receipt = None if offered is None else offered[2] + assert type(receipt) is tuple and value_only(receipt) + assert receipt == first._reconnect_receipt + + with monkeypatch.context() as construction_guard: + construction_guard.setattr(StaticModel, "from_pretrained", forbidden_constructor) + second._ensure_connection() + assert second.ready and second._has_vectors + assert second.conn is not first.conn + assert assert_publication(second, vector, encoder, message_id, content) == before + assert len(encoder.calls) == encode_count + assert vector._model is encoder + assert constructor_calls == [] + assert coordinator.snapshot()["active"] and not release.is_set() + assert not failures, failures + finally: + release.set() + try: + assert coordinator.wait(15), "Automatic worker did not finish after barrier release" + assert not failures, failures + assert completed, "The wrapper never called the real maintenance implementation" + finally: + coordinator.cancel() + try: + assert coordinator.wait(10), "Automatic maintenance teardown did not finish" + finally: + second.close() + + +def test_fresh_public_search_waits_for_committed_style_enrollment( + native_engine: tuple, monkeypatch: pytest.MonkeyPatch, +) -> None: + import sentence_transformers + from model2vec import StaticModel + from truememory import maintenance, reranker, storage + from truememory.client import Memory + from truememory.tier_switch import runtime + + first, vector, encoder = native_engine + content = "synthetic workshop inventory record" + query = "workshop inventory" + message_id = seed_source(first, content) + vector.embed_single(first.conn, message_id, content) + before = assert_publication(first, vector, encoder, message_id, content) + encode_count = len(encoder.calls) + assert first._has_style_vec + receipt = first._reconnect_receipt + assert receipt is not None + schema_before = first.conn.execute("PRAGMA schema_version").fetchone()[0] + assert schema_before == receipt[3] + expected_triggers = { + name: " ".join(sql.split()) for name, sql in storage._STYLE_OUTPUT_TRIGGERS.items() + } + assert len(expected_triggers) == 3 + trigger_query = "SELECT name,sql FROM sqlite_master WHERE type='trigger' AND name IN (?,?,?)" + assert first.conn.execute(trigger_query, tuple(expected_triggers)).fetchall() == [] + + coordinator = first._get_maintenance_coordinator() + fresh = Memory(path=first.db_path) + enrolled = threading.Event() + release = threading.Event() + wait_entered = threading.Event() + search_done = threading.Event() + preparation_lock = threading.Lock() + preparations: list[bool] = [] + enrolled_schemas: list[int] = [] + worker_errors: list[BaseException] = [] + search_errors: list[BaseException] = [] + results: list[dict] = [] + constructor_calls: list[bool] = [] + reranker_constructor_calls: list[bool] = [] + reranker_pairs: list[tuple[tuple[str, str], ...]] = [] + reranker_predictions: list[tuple[float, ...]] = [] + expected_reranker_name = reranker.get_current_reranker_name() + real_prepare = maintenance._prepare_routed_style_maintenance + real_wait = maintenance.wait_for_automatic_owner + + class SyntheticReranker: + def predict(self, pairs: list[tuple[str, str]], **kwargs: object) -> list[float]: + scores = [1.0 if document == content else 0.25 for _, document in pairs] + reranker_pairs.append(tuple(pairs)) + reranker_predictions.append(tuple(scores)) + return scores + + synthetic_reranker = SyntheticReranker() + + def prepared_then_blocked(conn: object) -> None: + with preparation_lock: + first_preparation = not preparations + preparations.append(True) + try: + real_prepare(conn) + if not first_preparation: + return + assert conn is not first.conn + assert not conn.in_transaction + installed = { + name: " ".join(sql.strip().rstrip(";").split()) + for name, sql in conn.execute(trigger_query, tuple(expected_triggers)) + } + assert installed == expected_triggers + schema_after = conn.execute("PRAGMA schema_version").fetchone()[0] + assert schema_after == schema_before + len(expected_triggers) + enrolled_schemas.append(schema_after) + enrolled.set() + if not release.wait(30): + raise AssertionError("Public warm search outlived the style-enrollment barrier") + except BaseException as error: + worker_errors.append(error) + enrolled.set() + raise + + def observed_wait(error: object, *, deadline: float) -> object: + assert runtime.current_runtime_operation() is None + assert fresh._engine.conn is None and not fresh._engine.ready + assert len(encoder.calls) == encode_count + wait_entered.set() + return real_wait(error, deadline=deadline) + + def search_publicly() -> None: + try: + results.extend(fresh.search(query)) + except BaseException as error: + search_errors.append(error) + finally: + search_done.set() + + def forbidden_constructor(*args: object, **kwargs: object) -> object: + constructor_calls.append(True) + raise AssertionError("Warm public search must reuse its synthetic embedding model") + + def forbidden_reranker_constructor(*args: object, **kwargs: object) -> object: + reranker_constructor_calls.append(True) + raise AssertionError("Warm public search must reuse its synthetic reranker") + + search_worker = threading.Thread(target=search_publicly, name="synthetic-style-warm-search") + monkeypatch.setattr(maintenance, "_prepare_routed_style_maintenance", prepared_then_blocked) + monkeypatch.setattr(maintenance, "wait_for_automatic_owner", observed_wait) + with monkeypatch.context() as construction_guard: + construction_guard.setattr(StaticModel, "from_pretrained", forbidden_constructor) + construction_guard.setattr(sentence_transformers, "CrossEncoder", forbidden_reranker_constructor) + construction_guard.setattr(reranker, "_model", synthetic_reranker) + construction_guard.setattr(reranker, "_model_name", expected_reranker_name) + construction_guard.setattr(reranker, "_model_certified", False) + try: + first._maybe_startup_consolidate() + assert enrolled.wait(10), "Actual canonical style preparation did not complete" + assert not worker_errors, worker_errors + assert enrolled_schemas == [schema_before + 3] + assert first.conn.execute("PRAGMA schema_version").fetchone()[0] == enrolled_schemas[0] + assert coordinator.snapshot()["active"] + with pytest.raises(maintenance.MaintenanceBusyError): + with maintenance.maintenance_owner(first.db_path): + pytest.fail("Committed style enrollment lost its automatic maintenance owner") + + search_worker.start() + assert wait_entered.wait(10), "Fresh public search did not reach automatic-owner waiting" + assert not search_done.is_set() and search_worker.is_alive() + assert not search_errors, search_errors + assert fresh._engine.conn is None and not fresh._engine.ready + assert len(encoder.calls) == encode_count + assert constructor_calls == [] + assert reranker_constructor_calls == [] and reranker_pairs == [] + assert coordinator.snapshot()["active"] and not release.is_set() + assert fresh._engine._write_lock.acquire(blocking=False), "Public wait retained the Engine writer lock" + fresh._engine._write_lock.release() + + release.set() + search_worker.join(15) + assert not search_worker.is_alive() and search_done.is_set() + assert not search_errors, search_errors + assert coordinator.wait(15), "Automatic worker did not finish after style barrier release" + assert not worker_errors, worker_errors + assert fresh._engine.ready and fresh._engine._has_vectors + assert any(result["id"] == message_id and "vec" in result["source"] for result in results) + assert assert_publication(fresh._engine, vector, encoder, message_id, content) == before + search_calls = encoder.calls[encode_count:] + assert 1 <= len(search_calls) <= 2 + assert all(texts == (query,) for texts in search_calls) + assert vector._model is encoder and constructor_calls == [] + assert reranker._model is synthetic_reranker and reranker._model_name == expected_reranker_name + assert reranker_constructor_calls == [] and reranker_pairs + assert all(pair_query == query for pairs in reranker_pairs for pair_query, _ in pairs) + assert any(document == content for pairs in reranker_pairs for _, document in pairs) + assert all( + len(pairs) == len(scores) + and scores == tuple(1.0 if document == content else 0.25 for _, document in pairs) + for pairs, scores in zip(reranker_pairs, reranker_predictions) + ) + finally: + release.set() + try: + if search_worker.ident is not None: + search_worker.join(15) + assert not search_worker.is_alive(), "Public search thread did not join" + finally: + coordinator.cancel() + try: + assert coordinator.wait(10), "Style maintenance teardown did not finish" + finally: + fresh._engine.close() + + +def test_native_mcp_deep_children_keep_final_tables_and_defer_maintenance( + native_engine: tuple, monkeypatch: pytest.MonkeyPatch, +) -> None: + from model2vec import StaticModel + import sentence_transformers + from truememory import Memory, maintenance, reranker + from truememory.engine import TrueMemoryEngine + from truememory.tier_switch import runtime, serving + + monkeypatch.setenv("TRUEMEMORY_PRELOAD_MODELS", "0") + from truememory import mcp_server + + first, vector, encoder = native_engine + content = "synthetic workshop inventory record" + queries = ("workshop inventory", "synthetic inventory") + message_id = seed_source(first, content) + vector.embed_single(first.conn, message_id, content) + before = assert_publication(first, vector, encoder, message_id, content) + encode_count = len(encoder.calls) + source_row = first.conn.execute( + "SELECT sender,recipient,timestamp FROM messages WHERE id=?", (message_id,), + ).fetchone() + separation = vector._build_sep_text(*source_row, content) + coordinator = first._get_maintenance_coordinator() + parent = Memory(path=first.db_path) + expected_tables = ("vec_messages_edge", "vec_messages_sep_edge") + expected_reranker = mcp_server._DEEP_RERANKER + both_constructed = threading.Event() + both_opened = threading.Event() + query_barrier = threading.Barrier(2) + body_finished = threading.Event() + worker_entered = threading.Event() + worker_release = threading.Event() + observations_lock = threading.Lock() + children: list[Memory] = [] + opening: list[object] = [] + opened: list[object] = [] + active_queries: list[str] = [] + concurrent_queries: list[int] = [] + completed_queries: list[str] = [] + parent_keys: list[tuple] = [] + child_keys: list[tuple] = [] + requests: list[bool] = [] + worker_runs: list[bool] = [] + errors: list[BaseException] = [] + constructor_calls: list[str] = [] + reranker_pairs: list[tuple[tuple[str, str], ...]] = [] + real_open = TrueMemoryEngine._open_connection_handle + real_deep = Memory.search_deep + real_request = coordinator.request_layers + real_run = maintenance.run_engine_maintenance + + class SyntheticDeepReranker: + def predict(self, pairs: list[tuple[str, str]], **kwargs: object) -> list[float]: + reranker_pairs.append(tuple(pairs)) + return [1.0 if document == content else 0.25 for _, document in pairs] + + synthetic_reranker = SyntheticDeepReranker() + + def make_child(*args: object, **kwargs: object) -> Memory: + child = Memory(*args, **kwargs) + with observations_lock: + children.append(child) + if len(children) == 2: + both_constructed.set() + return child + + def observed_open(engine: object) -> None: + with observations_lock: + is_child = any(child._engine is engine for child in children) + first_handle = is_child and engine.conn is None + if first_handle: + opening.append(engine) + ordinal = len(opening) + if not first_handle: + return real_open(engine) + try: + assert both_constructed.wait(10), "Both real child Memory objects were not constructed" + with observations_lock: + assert len(opening) == ordinal and len(opened) == ordinal - 1 + real_open(engine) + assert engine.conn is not None and not engine.ready + with observations_lock: + opened.append(engine) + if len(opened) == 2: + both_opened.set() + except BaseException as error: + errors.append(error) + raise + + def observed_deep(memory: Memory, query: str, *args: object, **kwargs: object) -> list[dict]: + try: + assert both_opened.wait(10), "A child search retained the sibling opening lock" + assert kwargs == {"user_id": None, "limit": 10, "llm_fn": None} and args == () + operation = runtime.current_operation(memory._engine.conn) + assert operation is not None and operation._defer_maintenance + assert operation.tables == expected_tables and operation.key == parent_keys[0] + assert operation.reranker_id == expected_reranker + assert reranker.get_current_reranker_name() == expected_reranker + assert reranker.get_reranker() is synthetic_reranker + assert requests == [] and not worker_entered.is_set() + with observations_lock: + child_keys.append(operation.key) + active_queries.append(query) + concurrent_queries.append(len(active_queries)) + query_barrier.wait(10) + result = real_deep(memory, query, *args, **kwargs) + assert memory._engine.ready and memory._engine._has_vectors + assert runtime.current_operation(memory._engine.conn) is operation + assert (vector._active_vec_table(memory._engine.conn), + vector._active_sep_table(memory._engine.conn)) == expected_tables + assert assert_publication(memory._engine, vector, encoder, message_id, content) == before + assert any(item["id"] == message_id and "vec" in item["source"] for item in result) + assert requests == [] and not worker_entered.is_set() + completed_queries.append(query) + return result + except BaseException as error: + errors.append(error) + query_barrier.abort() + raise + finally: + with observations_lock: + if query in active_queries: + active_queries.remove(query) + + def observed_request(**kwargs: object) -> bool: + requests.append(True) + try: + assert body_finished.is_set() and len(requests) == 1 + assert runtime.current_runtime_operation() is None + assert not serving.current_thread_admitted() + assert len(children) == 2 and all(child._engine.conn is None for child in children) + assert sorted(completed_queries) == sorted(queries) + return real_request(**kwargs) + except BaseException as error: + errors.append(error) + raise + + def observed_run(conn: object, owner: object, **kwargs: object) -> object: + try: + assert owner is coordinator and body_finished.is_set() + worker_runs.append(True) + if len(worker_runs) == 1: + worker_entered.set() + if not worker_release.wait(30): + raise AssertionError("MCP post-parent worker outlived its observation barrier") + return real_run(conn, owner, **kwargs) + except BaseException as error: + errors.append(error) + worker_entered.set() + raise + + def forbidden_embedding_constructor(*args: object, **kwargs: object) -> object: + constructor_calls.append("embedding") + raise AssertionError("MCP child searches must retain the synthetic embedding model") + + def forbidden_reranker_constructor(*args: object, **kwargs: object) -> object: + constructor_calls.append("reranker") + raise AssertionError("MCP Deep children must retain the synthetic explicit reranker") + + with monkeypatch.context() as observed: + observed.setattr(mcp_server, "_memory", parent) + observed.setattr(mcp_server, "Memory", make_child) + observed.setattr(TrueMemoryEngine, "_open_connection_handle", observed_open) + observed.setattr(Memory, "search_deep", observed_deep) + observed.setattr(coordinator, "request_layers", observed_request) + observed.setattr(maintenance, "run_engine_maintenance", observed_run) + observed.setattr(StaticModel, "from_pretrained", forbidden_embedding_constructor) + observed.setattr(sentence_transformers, "CrossEncoder", forbidden_reranker_constructor) + observed.setattr(reranker, "_model", synthetic_reranker) + observed.setattr(reranker, "_model_name", expected_reranker) + observed.setattr(reranker, "_model_certified", False) + try: + with mcp_server._search_operation(reranker_id=expected_reranker) as (memory, operation): + assert memory is parent and parent._engine.ready + assert operation.tables == expected_tables and operation._defer_maintenance + assert operation.reranker_id == expected_reranker + parent_keys.append(operation.key) + assert assert_publication(parent._engine, vector, encoder, message_id, content) == before + assert parent._engine.conn.execute( + "SELECT name FROM sqlite_master WHERE name IN ('vec_messages','vec_messages_sep')", + ).fetchall() == [] + assert len(encoder.calls) == encode_count and constructor_calls == [] + assert requests == [] and not worker_entered.is_set() + results = mcp_server._parallel_search( + queries, None, 10, None, 10, memory=memory, operation=operation, + ) + assert not errors, errors + assert len(opening) == len(opened) == 2 and both_opened.is_set() + assert child_keys == [operation.key, operation.key] + assert max(concurrent_queries) == 2 and active_queries == [] + assert sorted(completed_queries) == sorted(queries) + assert all(child._engine.conn is None for child in children) + assert any(item["id"] == message_id and "vec" in item["source"] for item in results) + assert requests == [] and not worker_entered.is_set() + assert assert_publication(parent._engine, vector, encoder, message_id, content) == before + body_finished.set() + assert requests == [True] + assert worker_entered.wait(10), "Deferred automatic maintenance did not start after the parent" + assert not errors, errors + assert coordinator.snapshot()["active"] and not worker_release.is_set() + assert constructor_calls == [] and vector._model is encoder + assert reranker._model is synthetic_reranker and reranker._model_name == expected_reranker + assert encoder.calls[encode_count:] + assert all(text not in (content, separation) + for texts in encoder.calls[encode_count:] for text in texts) + assert all(query in queries for pairs in reranker_pairs for query, _ in pairs) + worker_release.set() + assert coordinator.wait(15), "MCP post-parent maintenance did not finish" + assert not errors, errors + assert requests == [True] and constructor_calls == [] + assert assert_publication(parent._engine, vector, encoder, message_id, content) == before + finally: + worker_release.set() + query_barrier.abort() + coordinator.cancel() + try: + assert coordinator.wait(10), "MCP shared maintenance teardown did not finish" + finally: + for child in children: + child.close() + parent.close() diff --git a/tests/test-tier-mcp-serving-795.py b/tests/test-tier-mcp-serving-795.py new file mode 100644 index 00000000..ab5f04fc --- /dev/null +++ b/tests/test-tier-mcp-serving-795.py @@ -0,0 +1,281 @@ +"""MCP policy capture and child cleanup with actual gates and synthetic peers.""" +from __future__ import annotations + +import ast +import builtins +import json +import logging +import runpy +import threading +import types +import unittest +from collections.abc import Iterator +from concurrent.futures import ThreadPoolExecutor +from contextlib import contextmanager +from pathlib import Path +from unittest.mock import patch + +ROOT = Path(__file__).resolve().parents[1] +GATE = runpy.run_path(str(ROOT / "truememory/tier_switch/serving.py")) + + +class PolicyRejected(RuntimeError): + pass + + +class TestMCPServing(unittest.TestCase): + def setUp(self) -> None: + self.events = [] + self.gate = GATE["ServingGate"]() + self.tier = "pro" + self.llm = object() + self.deep_llm = None + self.query_failure = None + self.construction_failure = False + self.maintenance_requests = [] + owner = self + + class Memory: + def __init__(self, path="synthetic"): + if owner.construction_failure: + raise OSError("synthetic open failure") + self._engine = types.SimpleNamespace(db_path=path, conn=object(), _write_lock=threading.Lock()) + self._engine._open_connection_handle = lambda: owner.events.append(("open", self._engine.conn)) + self._engine._runtime_initialized = True + self._engine.ready = True + def request_maintenance(): + owner.assert_drained() + owner.maintenance_requests.append(self._engine.conn) + self._engine._maybe_auto_consolidate = request_maintenance + self._engine._maybe_startup_consolidate = request_maintenance + + def __enter__(self): + return self + + def __exit__(self, *exc): + owner.events.append(("close", self._engine.conn)) + + def search_deep(self, query, **kwargs): + owner.events.append(("query", query, kwargs)) + if owner.query_failure is not None: + raise owner.query_failure + return [{"id": "shared", "score": 2}, {"id": query, "score": 1}] + + self.memory = Memory() + + @contextmanager + def serving_operation(conn, *, connection_lock=None, reranker_id=None, _defer_maintenance=False): + owner.events.append(("admit", conn, reranker_id)) + with owner.gate.operation_lease(("synthetic-generation",)) as lease: + def fork_child(): + child = lease.fork_child() + owner.events.append(("reserve", child)) + + @contextmanager + def join(*, conn): + owner.events.append(("join", conn)) + with child.join(): + yield + + return types.SimpleNamespace(join=join, close=child.close) + + yield types.SimpleNamespace(tier=owner.tier, reranker_id=reranker_id or "synthetic/default", + fork_child=fork_child, _defer_maintenance=_defer_maintenance) + owner.events.append(("release", conn)) + + def reject(error): + if isinstance(error, PolicyRejected): + raise error + + runtime = types.SimpleNamespace(serving_operation=serving_operation, raise_if_serving_rejection=reject) + + def safe_import(name, globals=None, locals=None, fromlist=(), level=0): + if name == "truememory.tier_switch.runtime": + return runtime + return builtins.__import__(name, globals, locals, fromlist, level) + + runtime_tree = ast.parse((ROOT / "truememory/tier_switch/runtime.py").read_text(encoding="utf-8")) + context_names = {"open_engine_for_operation", "engine_serving_operation"} + context_definitions = [node for node in runtime_tree.body + if isinstance(node, ast.FunctionDef) and node.name in context_names] + self.assertEqual({node.name for node in context_definitions}, context_names) + context_namespace = dict(__builtins__=dict(vars(builtins), __import__=safe_import), + contextmanager=contextmanager, serving_operation=serving_operation, + current_operation=lambda conn: None) + future = ast.ImportFrom(module="__future__", names=[ast.alias(name="annotations")], level=0) + exec(compile(ast.fix_missing_locations(ast.Module(body=[future, *context_definitions], type_ignores=[])), + "actual-public-context", "exec"), context_namespace) + runtime.engine_serving_operation = context_namespace["engine_serving_operation"] + + names = {"_search_operation", "_parallel_search", "_resolve_deepsearch_llm", + "truememory_search", "truememory_search_deep"} + tree = ast.parse((ROOT / "truememory/mcp_server.py").read_text(encoding="utf-8")) + nodes = [ast.ImportFrom(module="__future__", names=[ast.alias(name="annotations")], level=0)] + for node in tree.body: + if isinstance(node, ast.FunctionDef) and node.name in names: + node.decorator_list = [d for d in node.decorator_list + if isinstance(d, ast.Name) and d.id == "contextmanager"] + nodes.append(node) + self.ns = dict(__builtins__=dict(vars(builtins), __import__=safe_import), contextmanager=contextmanager, + Iterator=Iterator, Memory=Memory, ThreadPoolExecutor=ThreadPoolExecutor, json=json, threading=threading, + log=logging.getLogger(__name__), _get_memory=lambda: self.memory, + _get_llm_fn=lambda: self.llm, _get_deepsearch_llm_fn=lambda: self.deep_llm, + _load_config=lambda: {"tier": "edge"}, _touch_search_time=lambda: None, + _set_reranker=lambda name: self.events.append(("reranker", name)), + _SEARCH_INTERNAL_LIMIT=100, _DEEP_INTERNAL_LIMIT=500, + _DEEP_RERANKER="BAAI/bge-reranker-v2-m3") + exec(compile(ast.fix_missing_locations(ast.Module(body=nodes, type_ignores=[])), + "actual-mcp-serving", "exec"), self.ns) + + def assert_drained(self) -> None: + with self.gate.exclusive_activation(timeout=0.1): + pass + + def test_outer_request_schedules_once_after_all_children_release(self) -> None: + self.ns["truememory_search"]("first | second") + self.assertEqual(self.maintenance_requests, [self.memory._engine.conn]) + self.assertEqual(self.events[-1][0], "release") + self.assert_drained() + + def test_failed_child_schedules_once_after_request_cleanup(self) -> None: + self.query_failure = PolicyRejected("synthetic rejection") + with self.assertRaises(PolicyRejected): + self.ns["truememory_search"]("first | second") + self.assertEqual(self.maintenance_requests, [self.memory._engine.conn]) + self.assert_drained() + + def test_parallel_handles_open_serially_and_release_before_search(self) -> None: + first_opened = threading.Event() + both_constructed = threading.Event() + second_opened = threading.Event() + release = threading.Event() + mutex = threading.Lock() + constructed = [] + opened = [] + errors = [] + outputs = [] + owner = self + original = self.ns["Memory"] + + class ConcurrentMemory(original): + def __init__(self, path="synthetic"): + super().__init__(path) + with mutex: + constructed.append(self) + if len(constructed) == 2: + both_constructed.set() + def open_handle(): + with mutex: + opened.append(self) + number = len(opened) + if number == 1: + first_opened.set() + if not release.wait(2): + raise AssertionError("First open was not released") + else: + second_opened.set() + self._engine._open_connection_handle = open_handle + + def search_deep(self, query, **kwargs): + owner.assertTrue(second_opened.wait(2), "Opening lock was retained during search") + return super().search_deep(query, **kwargs) + + self.ns["Memory"] = ConcurrentMemory + def run_request(): + try: + outputs.extend(json.loads(self.ns["truememory_search"]("first | second"))) + except BaseException as error: + errors.append(error) + worker = threading.Thread(target=run_request) + worker.start() + try: + self.assertTrue(first_opened.wait(2)) + self.assertTrue(both_constructed.wait(2)) + self.assertFalse(second_opened.wait(0.05), "Sibling opened during the first maintenance owner") + self.assertEqual(len(opened), 1) + finally: + release.set() + worker.join(5) + self.assertFalse(worker.is_alive()) + self.assertEqual(errors, []) + self.assertEqual(len(opened), 2) + self.assertEqual([row["id"] for row in outputs], ["shared", "first", "second"]) + self.assertEqual(self.maintenance_requests, [self.memory._engine.conn]) + self.assert_drained() + + def test_standard_uses_admitted_policy_despite_old_config(self) -> None: + result = json.loads(self.ns["truememory_search"]("synthetic", limit=1)) + query = next(event for event in self.events if event[0] == "query") + self.assertEqual(len(result), 1) + self.assertEqual(query[2]["limit"], 100) + self.assertIs(query[2]["llm_fn"], self.llm) + self.assertEqual(self.events[0][0], "open") + self.assertEqual(self.events[1][0], "admit") + self.assertEqual(self.events[-1][0], "release") + self.assert_drained() + + def test_deep_override_and_edge_explicit_llm_survive(self) -> None: + self.tier = "edge" + self.deep_llm = object() + self.ns["truememory_search_deep"]("synthetic") + query = next(event for event in self.events if event[0] == "query") + self.assertEqual(self.events[1][2], "BAAI/bge-reranker-v2-m3") + self.assertEqual(query[2]["limit"], 500) + self.assertIs(query[2]["llm_fn"], self.deep_llm) + self.assertIn(("reranker", "BAAI/bge-reranker-v2-m3"), self.events) + self.assert_drained() + + def test_deep_default_llm_is_pro_only(self) -> None: + self.assertIsNone(self.ns["_resolve_deepsearch_llm"]("edge")) + self.assertIs(self.ns["_resolve_deepsearch_llm"]("pro"), self.llm) + + def test_parallel_children_join_and_release_before_parent(self) -> None: + result = json.loads(self.ns["truememory_search"]("first | second")) + self.assertEqual([row["id"] for row in result], ["shared", "first", "second"]) + self.assertEqual(sum(e[0] == "reserve" for e in self.events), 2) + self.assertEqual(sum(e[0] == "join" for e in self.events), 2) + self.assertEqual(sum(e[0] == "close" for e in self.events), 2) + self.assertEqual(self.events[-1][0], "release") + self.assert_drained() + + def test_identity_rejection_cannot_return_partial_success(self) -> None: + self.query_failure = PolicyRejected("synthetic stale selection") + with self.assertRaises(PolicyRejected): + self.ns["truememory_search"]("first | second") + self.assert_drained() + + def test_failed_child_constructor_releases_unused_reservation(self) -> None: + self.construction_failure = True + self.assertEqual(json.loads(self.ns["truememory_search"]("first | second")), []) + self.assertEqual(sum(e[0] == "reserve" for e in self.events), 2) + self.assert_drained() + + def test_failed_submission_releases_every_reserved_child(self) -> None: + original = ThreadPoolExecutor.submit + calls = 0 + + def submit(pool, *args, **kwargs): + nonlocal calls + calls += 1 + if calls == 2: + raise RuntimeError("synthetic submission failure") + return original(pool, *args, **kwargs) + + with patch.object(ThreadPoolExecutor, "submit", submit): + with self.assertRaisesRegex(RuntimeError, "submission"): + self.ns["truememory_search"]("first | second") + self.assert_drained() + + def test_input_caps_and_empty_input_are_preserved(self) -> None: + self.assertEqual(self.ns["truememory_search"](" "), "[]") + self.assertEqual(self.events, []) + self.ns["truememory_search"]("", queries=[str(n) + "x" * 2100 for n in range(12)], limit=999) + queries = [e for e in self.events if e[0] == "query"] + self.assertEqual(len(queries), 10) + self.assertTrue(all(len(e[1]) == 2000 for e in queries)) + self.assertTrue(all(e[2]["limit"] == 100 for e in queries)) + self.assert_drained() + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test-tier-metadata-publication-795.py b/tests/test-tier-metadata-publication-795.py new file mode 100644 index 00000000..68d977ab --- /dev/null +++ b/tests/test-tier-metadata-publication-795.py @@ -0,0 +1,176 @@ +"""Foreground metadata publication under actual serving and writer guards.""" +from __future__ import annotations + +import datetime +from pathlib import Path +import runpy +import sqlite3 +import unittest +from unittest.mock import Mock + +BOUNDARIES = runpy.run_path(str(Path(__file__).with_name("test-tier-public-boundaries-795.py"))) + + +class TestForegroundMetadata(unittest.TestCase): + def setUp(self) -> None: + self.fixture = BOUNDARIES["TestPublicBoundaries"]() + self.fixture.setUp() + self.addCleanup(self.fixture.doCleanups) + self.conn, self.runtime = self.fixture.conn, self.fixture.api + self.writer = self.fixture.load_application("tier_switch.writer") + self.vector = self.fixture.fixture.vector + self.source = BOUNDARIES["definitions"]( + "vector_search.py", + {"VectorPublicationChanged", "_write_foreground_embedder_metadata_no_commit", + "_write_embedder_metadata_no_commit", "_write_embedder_metadata"}, + self.fixture.namespace(datetime=datetime, _model_generation=7, + EMBEDDING_MODEL=self.vector.EMBEDDING_MODEL, + _embedding_dim=self.vector._embedding_dim), + ) + self.identity = (7, self.source.EMBEDDING_MODEL, self.source._embedding_dim) + self.conn.execute("DELETE FROM main.metadata WHERE key IN ('embed_model','embed_dim')") + self.conn.commit() + self.owner = Mock(side_effect=RuntimeError("Synthetic maintenance owner busy")) + self.fixture.modules["maintenance"].maintenance_owner = self.owner + + def publish(self, selection: object) -> None: + self.source._write_foreground_embedder_metadata_no_commit( + self.conn, self.identity, selection=selection) + + def rows(self) -> list: + return self.conn.execute( + "SELECT key,value FROM main.metadata WHERE key IN ('embed_model','embed_dim') ORDER BY key" + ).fetchall() + + def test_initial_and_matching_publication_do_not_acquire_maintenance(self) -> None: + with self.runtime.serving_operation(self.conn): + selection = self.writer.capture_writer_selection(self.conn) + for _ in range(2): + with self.writer.writer_transaction(self.conn, selection): + self.publish(selection) + self.assertTrue(self.conn.in_transaction) + self.assertEqual(self.rows(), [("embed_dim", "256"), ("embed_model", "model2vec")]) + self.owner.assert_not_called() + + def test_unbound_and_clean_calls_refuse_without_writes(self) -> None: + selection = self.writer.capture_writer_selection(self.conn) + with self.assertRaises(self.source.VectorPublicationChanged): + self.publish(selection) + with self.runtime.serving_operation(self.conn): + with self.assertRaises(self.writer.WriterSelectionChanged): + self.publish(selection) + self.assertEqual(self.rows(), []) + self.assertFalse(self.conn.in_transaction) + + def test_changed_model_generation_or_identity_refuses_without_writes(self) -> None: + with self.runtime.serving_operation(self.conn): + selection = self.writer.capture_writer_selection(self.conn) + for name, value in (("_model_generation", 8), ("EMBEDDING_MODEL", "synthetic/changed"), + ("_embedding_dim", 512)): + with self.subTest(field=name): + old = getattr(self.source, name) + setattr(self.source, name, value) + with self.writer.writer_transaction(self.conn, selection): + with self.assertRaises(self.source.VectorPublicationChanged): + self.publish(selection) + self.assertEqual(self.rows(), []) + setattr(self.source, name, old) + + def test_partial_mismatched_and_nontext_metadata_refuse_atomically(self) -> None: + for rows in ( + [("embed_model", "model2vec")], [("embed_dim", "256")], + [("embed_model", "synthetic/wrong"), ("embed_dim", "256")], + [("embed_model", "model2vec"), ("embed_dim", "512")], + [("embed_model", b"model2vec"), ("embed_dim", "256")], + ): + with self.subTest(rows=rows): + self.conn.execute("DELETE FROM main.metadata WHERE key IN ('embed_model','embed_dim')") + self.conn.executemany("INSERT INTO main.metadata(key,value) VALUES(?,?)", rows) + self.conn.commit() + before = self.rows() + with self.runtime.serving_operation(self.conn): + selection = self.writer.capture_writer_selection(self.conn) + with self.writer.writer_transaction(self.conn, selection): + changes = self.conn.total_changes + with self.assertRaises(self.source.VectorPublicationChanged): + self.publish(selection) + self.assertEqual(self.conn.total_changes, changes) + self.assertEqual(self.rows(), before) + + def test_temp_metadata_neither_satisfies_nor_receives_main_publication(self) -> None: + self.conn.execute("CREATE TEMP TABLE metadata(key TEXT PRIMARY KEY,value TEXT,updated_at TEXT)") + self.conn.execute("INSERT INTO temp.metadata VALUES('embed_model','synthetic/temp',NULL)") + self.conn.commit() + with self.runtime.serving_operation(self.conn): + selection = self.writer.capture_writer_selection(self.conn) + with self.writer.writer_transaction(self.conn, selection): + self.publish(selection) + self.assertEqual(self.rows(), [("embed_dim", "256"), ("embed_model", "model2vec")]) + self.assertEqual(self.conn.execute("SELECT value FROM temp.metadata").fetchall(), [("synthetic/temp",)]) + + def test_borrowed_failure_rolls_back_our_metadata_and_retains_outer_work(self) -> None: + self.conn.execute("INSERT INTO messages(id,content) VALUES(701,'synthetic outer work')") + with self.runtime.serving_operation(self.conn): + selection = self.writer.capture_writer_selection(self.conn) + before = self.conn.commits, self.conn.rollbacks + with self.assertRaisesRegex(RuntimeError, "synthetic downstream"): + with self.writer.writer_transaction(self.conn, selection): + self.publish(selection) + raise RuntimeError("synthetic downstream failure") + self.assertTrue(self.conn.in_transaction) + self.assertEqual((self.conn.commits, self.conn.rollbacks), before) + self.assertEqual(self.rows(), []) + self.assertEqual(self.conn.execute("SELECT content FROM messages WHERE id=701").fetchone(), + ("synthetic outer work",)) + + def test_owned_failure_rolls_back_metadata(self) -> None: + with self.runtime.serving_operation(self.conn): + selection = self.writer.capture_writer_selection(self.conn) + with self.assertRaisesRegex(RuntimeError, "synthetic downstream"): + with self.writer.writer_transaction(self.conn, selection): + self.publish(selection) + raise RuntimeError("synthetic downstream failure") + self.assertEqual(self.rows(), []) + self.assertFalse(self.conn.in_transaction) + + def test_selected_publication_validates_but_does_not_rewrite_metadata(self) -> None: + selected = self.fixture.fixture.select() + self.source.EMBEDDING_MODEL = selected.target.model_id + self.source._embedding_dim = selected.target.dimension + self.identity = (7, self.source.EMBEDDING_MODEL, self.source._embedding_dim) + with self.runtime.serving_operation(self.conn): + selection = self.writer.capture_writer_selection(self.conn) + with self.writer.writer_transaction(self.conn, selection): + changes = self.conn.total_changes + self.publish(selection) + self.assertEqual(self.conn.total_changes, changes) + self.owner.assert_not_called() + + def test_general_metadata_mutators_still_require_maintenance_ownership(self) -> None: + for name in ("_write_embedder_metadata_no_commit", "_write_embedder_metadata"): + with self.subTest(name=name), self.runtime.serving_operation(self.conn): + selection = self.writer.capture_writer_selection(self.conn) + with self.writer.writer_transaction(self.conn, selection): + with self.assertRaisesRegex(RuntimeError, "Synthetic maintenance owner busy"): + getattr(self.source, name)(self.conn) + self.assertEqual(self.owner.call_count, 2) + self.assertEqual(self.rows(), []) + + def test_matching_metadata_supports_sqlite_row_factory(self) -> None: + self.conn.execute("INSERT INTO main.metadata(key,value) VALUES('embed_model','model2vec')") + self.conn.execute("INSERT INTO main.metadata(key,value) VALUES('embed_dim','256')") + self.conn.commit() + self.conn.row_factory = sqlite3.Row + try: + with self.runtime.serving_operation(self.conn): + selection = self.writer.capture_writer_selection(self.conn) + with self.writer.writer_transaction(self.conn, selection): + self.publish(selection) + self.assertEqual([tuple(row) for row in self.rows()], + [("embed_dim", "256"), ("embed_model", "model2vec")]) + finally: + self.conn.row_factory = None + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test-tier-public-adoption-795.py b/tests/test-tier-public-adoption-795.py new file mode 100644 index 00000000..021abdf6 --- /dev/null +++ b/tests/test-tier-public-adoption-795.py @@ -0,0 +1,1686 @@ +"""Public adoption orchestration with production AST, stubs and in-memory SQL.""" + +from __future__ import annotations + +import argparse +import ast +import builtins +import contextlib +import dataclasses +import io +import json +import logging +import re +import sqlite3 +import sys +import threading +import types +import unittest +from pathlib import Path +from unittest.mock import Mock, patch + +ROOT = Path(__file__).resolve().parents[1] + + +@dataclasses.dataclass(frozen=True) +class Target: + tier: str + model_id: str + dimension: int + tier_group: str + + @classmethod + def capture(cls, tier: str) -> Target: + if tier not in {"edge", "base", "pro"}: + raise ValueError("unsupported synthetic target") + return cls( + tier, + "model2vec" if tier == "edge" else "qwen3_256", + 256, + "edge" if tier == "edge" else "basepro", + ) + + @property + def identity(self) -> tuple: + return self.model_id, self.dimension, self.tier_group + + @property + def tables(self) -> tuple[str, str]: + return f"vec_messages_{self.tier_group}", f"vec_messages_sep_{self.tier_group}" + + +class ActivationIntent(types.SimpleNamespace): + pass + + +class PolicyIntent(types.SimpleNamespace): + pass + + +class LegacyTierPolicy(types.SimpleNamespace): + pass + + +class FakeConnection: + def __init__(self) -> None: + self.closed = False + self.in_transaction = False + self.executed: list[str] = [] + + def close(self) -> None: + self.closed = True + + def execute(self, sql: str, parameters: tuple = ()) -> object: + self.executed.append(sql) + raise AssertionError("Unexpected SQL on stub: " + sql) + + +@contextlib.contextmanager +def transaction(conn: sqlite3.Connection, *, write: bool): + if conn.in_transaction: + raise RuntimeError("borrowed transaction") + conn.execute("BEGIN IMMEDIATE" if write else "BEGIN") + try: + yield + conn.commit() if write else conn.rollback() + except BaseException: + if conn.in_transaction: + conn.rollback() + raise + + +class Harness: + def __init__(self) -> None: + self.events: list[object] = [] + self.config = {"tier": "edge", "independent": "keep"} + self.guard = types.SimpleNamespace(tier="edge", generation="old") + self.state = types.SimpleNamespace( + selection=None, legacy_policy=None, intent=None + ) + self.job = types.SimpleNamespace( + job_id="a" * 32, + target=Target.capture("base"), + database_path=Path("/synthetic/store.db"), + tracker="canonical-v1", + source_epoch="epoch", + source_schema_signature="b" * 64, + ) + self.intent = ActivationIntent( + job_id=self.job.job_id, + target=self.job.target, + tracker=self.job.tracker, + source_epoch=self.job.source_epoch, + source_schema_signature=self.job.source_schema_signature, + expected_config=self.guard, + status_id=17, + state="staged", + ) + self.marker = None + self.result = LegacyTierPolicy( + target=Target.capture("pro"), generation="new", config_acknowledged=False + ) + self.modules = { + "truememory.embedding_target": types.SimpleNamespace( + EmbeddingTarget=Target + ), + "truememory.tier_config": types.SimpleNamespace( + get_tier_config=lambda tier: {"reranker": "synthetic/reranker"} + ), + "truememory.maintenance": types.SimpleNamespace( + maintenance_owner=self.owner + ), + "truememory.rebuild_source": types.SimpleNamespace( + ensure_rebuild_tracking=lambda conn: self.events.append("tracking") + ), + "truememory.tier_switch.cache": types.SimpleNamespace( + preflight_ram_check=lambda group: (True, ""), + get_transition_action=lambda old, new: ( + "config_only" if {old, new} == {"base", "pro"} else "delta_or_full" + ), + ), + "truememory.tier_switch.throttler": types.SimpleNamespace( + DynamicThrottler=lambda **kwargs: object() + ), + "truememory.tier_switch.worker": types.SimpleNamespace( + RebuildWorker=Mock() + ), + "truememory.tier_switch.serving": types.SimpleNamespace( + exclusive_activation=self.exclusive + ), + "truememory.tier_switch.activation": types.SimpleNamespace( + ActivationIntent=ActivationIntent, + LegacyTierPolicy=LegacyTierPolicy, + read_activation_state=lambda conn: self.state, + _state=lambda conn: self.state, + commit_config_only_transition=self.commit_policy, + stage_activation_intent=self.stage, + acknowledge_config=lambda conn, **kwargs: self.events.append("ack"), + acknowledge_policy_config=lambda conn, **kwargs: self.events.append( + "ack" + ), + commit_certified_selection=lambda *args: self.result, + _transaction=transaction, + _normalized=lambda sql: re.sub( + r"\s*([(),])\s*", r"\1", " ".join(sql.strip().rstrip(";").split()) + ).lower(), + _schema_sql=lambda conn, name: conn.execute( + "SELECT sql FROM sqlite_master WHERE name=?", (name,) + ).fetchone()[0], + _guard_pair=lambda conn, target, **kwargs: self.events.append( + ("pair", kwargs) + ), + ), + "truememory.tier_switch.job": types.SimpleNamespace( + read_selected_job_marker=lambda conn: self.marker, + select_tier_job=self.select, + check_selected_tier_job=lambda *args: self.events.append("check"), + end_selected_tier_job=lambda *args, **kwargs: self.events.append( + ("end", kwargs) + ), + _cancel_policy_job_in_writer=lambda *args, **kwargs: self.events.append( + ("cancel", kwargs) + ), + ), + "truememory.tier_switch.projection": types.SimpleNamespace( + capture_config_guard=lambda **kwargs: self.guard, + mirror_activation=self.mirror, + patch_config_fields=self.patch_config, + CONFIG_WRITE_LOCK=threading.RLock(), + ConfigFileLock=lambda *args, **kwargs: contextlib.nullcontext(), + _read_config=lambda path: self.config.copy(), + _write_config=self.write_config, + ), + "truememory.tier_switch.runtime": types.SimpleNamespace( + apply_frozen_selection=lambda result: self.events.append( + "apply_frozen" + ), + apply_runtime_policy=lambda result: self.events.append("apply_policy"), + serving_operation=Mock( + side_effect=AssertionError("native readiness forbidden") + ), + ), + "truememory.tier_switch.source": types.SimpleNamespace( + read_tier_progress=lambda *args, **kwargs: None + ), + "truememory.vector_search": types.SimpleNamespace( + init_prepared_target_tables=lambda *args: self.events.append("tables") + ), + "truememory.storage": types.SimpleNamespace( + DEFAULT_BUSY_TIMEOUT_MS=1, + _validate_db_path=lambda p: p, + create_db=Mock(side_effect=AssertionError("initialization forbidden")), + ), + } + self.modules["truememory"] = types.SimpleNamespace( + vector_search=self.modules["truememory.vector_search"] + ) + vector = self.modules["truememory.vector_search"] + vector.resolve_tier = lambda: "edge" + vector.EMBEDDING_MODEL = "model2vec" + vector._embedding_dim = 256 + self.modules["truememory"].reranker = types.SimpleNamespace( + get_current_reranker_name=lambda: "synthetic/reranker" + ) + vector._model = None + self.modules["truememory"].reranker._model = None + + def apply_embedding_policy(target, **kwargs): + self.events.append("lazy_embedding") + vector.EMBEDDING_MODEL = target.model_id + vector._embedding_dim = target.dimension + vector.resolve_tier = lambda: target.tier + + vector.apply_embedding_policy = apply_embedding_policy + self.modules["truememory"].reranker.apply_reranker_policy = lambda *args: ( + self.events.append("lazy_reranker") + ) + + def set_embedding_model(tier): + self.events.append("lazy_embedding") + applied = Target.capture(tier) + vector.EMBEDDING_MODEL = applied.model_id + vector._embedding_dim = applied.dimension + vector.resolve_tier = lambda: tier + + vector.set_embedding_model = set_embedding_model + self.modules["truememory"].reranker.set_active_tier = lambda tier: self.events.append("lazy_reranker") + self.manager_module = self.load("truememory/tier_switch/manager.py") + self.manager_module.os = types.SimpleNamespace(environ={}, name="posix") + self.modules["truememory.tier_switch.manager"] = self.manager_module + self.manager = self.manager_module.RebuildManager() + self.conn = FakeConnection() + self.manager_module._open_db = lambda *args, **kwargs: self.conn + self.manager_module._load_tier_extension = lambda conn: self.events.append( + "extension" + ) + self.manager_module._detect_device = lambda: "cpu" + + def importer(self, name, globals=None, locals=None, fromlist=(), level=0): + if name in self.modules: + return self.modules[name] + if name.startswith("truememory") or name in { + "torch", + "numpy", + "sqlite_vec", + "sentence_transformers", + "model2vec", + }: + raise AssertionError("Unstubbed/native import: " + name) + return builtins.__import__(name, globals, locals, fromlist, level) + + def load( + self, relative: str, names: set[str] | None = None, extra: dict | None = None + ): + path = ROOT / relative + tree = ast.parse(path.read_text(encoding="utf-8"), filename=str(path)) + if names is not None: + tree.body = [ + node + for node in tree.body + if isinstance(node, (ast.FunctionDef, ast.ClassDef)) + and node.name in names + ] + for node in tree.body: + node.decorator_list = [] + module = types.ModuleType("synthetic_public_adoption") + module.__dict__.update( + { + "__builtins__": dict(vars(builtins), __import__=self.importer), + "Path": Path, + "sqlite3": sqlite3, + "sys": sys, + "threading": threading, + } + ) + if extra: + module.__dict__.update(extra) + exec(compile(tree, str(path), "exec"), module.__dict__) + return module + + @contextlib.contextmanager + def owner(self, path): + self.events.append("owner_enter") + try: + yield + finally: + self.events.append("owner_exit") + + @contextlib.contextmanager + def exclusive(self, **kwargs): + self.events.append("exclusive_enter") + try: + yield + finally: + self.events.append("exclusive_exit") + + def write_config(self, path, config): + self.config = config.copy() + self.events.append("config") + + def patch_config(self, fields, **kwargs): + self.config.update(fields) + return self.config.copy() + + def mirror(self, result, **kwargs): + self.events.append("mirror") + + def commit_policy(self, conn, **kwargs): + self.events.append("policy_commit") + return self.result + + def select(self, conn, target, **kwargs): + self.events.append("select") + return self.job + + def stage(self, conn, job, **kwargs): + self.events.append("stage") + return self.intent + + def prepare(self, old="edge", requested="base"): + self.guard.tier = old + self.manager._serving_identity = lambda *args: ( + Target.capture(old), + Target.capture(old).tables, + ) + self.manager._check_inactive_target = lambda *args: self.events.append( + "inactive_check" + ) + self.manager._guard_status = lambda *args: self.events.append("status_guard") + self.manager._create_status_row = lambda *args, **kwargs: ( + self.events.append(("status", kwargs)) or 17 + ) + return self.manager._prepare_job(self.conn, requested, False) + + +class MemoryPath: + def __truediv__(self, name: str) -> MemoryPath: + return self + + @property + def parent(self) -> MemoryPath: + return self + + def mkdir(self, **kwargs) -> None: + pass + + def write_text(self, content: str, **kwargs) -> None: + pass + + +def configure_harness(h: Harness, manager: object): + h.modules["truememory.tier_switch.manager"].RebuildManager = types.SimpleNamespace( + get_instance=lambda: manager + ) + return h.load( + "truememory/mcp_server.py", + {"truememory_configure"}, + { + "json": json, + "os": types.SimpleNamespace(environ={}), + "log": logging.getLogger("synthetic_configure"), + "_memory": object(), + "_memory_lock": threading.Lock(), + "_config_cache": None, + "_load_config": lambda: h.config.copy(), + "_CONFIG_PATH": Path("/synthetic/config.json"), + "_CONFIG_LOCK_PATH": Path("/synthetic/config.json.lock"), + "_DB_PATH": "/synthetic/store.db", + "_VALID_INTENSITIES": ("standard", "enhanced", "max"), + "_config_file_lock": lambda **kwargs: contextlib.nullcontext(), + "_config_write_lock": threading.RLock(), + }, + ) + + +def cancel_cli_harness(h: Harness, manager: object): + factory = Mock(return_value=manager) + h.modules["truememory.tier_switch.manager"].RebuildManager = factory + h.modules["truememory"].__version__ = "synthetic" + module = h.load( + "truememory/ingest/cli.py", + { + "main", + "_positive_rebuild_status_id", + "_rebuild_job_id", + "_run_cancel_rebuild", + }, + {"argparse": argparse, "re": re, "os": types.SimpleNamespace(environ={})}, + ) + return module, factory + + +def status_connection() -> sqlite3.Connection: + tree = ast.parse((ROOT / "truememory/tier_switch/manager.py").read_text(encoding="utf-8")) + method = next( + n + for n in ast.walk(tree) + if isinstance(n, ast.FunctionDef) and n.name == "_guard_status" + ) + sql = next( + n.value + for n in ast.walk(method) + if isinstance(n, ast.Constant) + and isinstance(n.value, str) + and n.value.startswith("CREATE TABLE rebuild_status") + ) + conn = sqlite3.connect(":memory:") + conn.execute(sql) + return conn + + +class PublicAdoptionTests(unittest.TestCase): + def setUp(self): + self.h = Harness() + + def test_force_refuses_before_config_or_sql(self): + h = self.h + with self.assertRaises(h.manager_module.TierSwitchUnsupportedError): + h.manager._prepare_job(h.conn, "base", True) + self.assertEqual(h.events, []) + self.assertEqual(h.conn.executed, []) + + def test_shared_prepare_orders_status_job_intent_before_source(self): + h = self.h + self.assertEqual(h.prepare(), (h.job, h.intent, 17)) + self.assertLess(h.events.index("inactive_check"), h.events.index("tracking")) + self.assertLess(h.events.index("tables"), h.events.index("select")) + self.assertLess(h.events.index("select"), h.events.index("stage")) + self.assertEqual(h.conn.executed, []) + + def test_live_marker_refuses_without_writes(self): + h = self.h + h.marker = types.SimpleNamespace(state="selected", job_id="old") + with self.assertRaises(h.manager_module.TierSwitchUnsupportedError): + h.prepare() + self.assertNotIn("tracking", h.events) + self.assertNotIn("select", h.events) + + def test_same_space_policy_projects_without_source_or_native(self): + h = self.h + h.state.intent = PolicyIntent(expected_config=h.guard, intent_id="policy") + self.assertIsNone(h.prepare("base", "pro")) + self.assertEqual( + [ + x + for x in h.events + if isinstance(x, str) and x not in {"exclusive_enter", "exclusive_exit"} + ], + ["policy_commit", "mirror", "apply_policy", "ack"], + ) + self.assertEqual(h.conn.executed, []) + self.assertEqual(h.manager._last_action, "config_only") + + def test_pending_same_tier_retries_projection(self): + h = self.h + h.state.legacy_policy = h.result + h.state.intent = PolicyIntent(expected_config=h.guard, intent_id="policy") + self.assertIsNone(h.prepare("pro", "pro")) + self.assertIn("apply_policy", h.events) + self.assertEqual(h.manager._last_outcome, "complete") + + def test_pending_other_tier_refuses(self): + h = self.h + h.state.legacy_policy = h.result + h.state.intent = PolicyIntent(expected_config=h.guard, intent_id="policy") + with self.assertRaises(h.manager_module.TierSwitchUnsupportedError): + h.prepare("pro", "edge") + self.assertNotIn("mirror", h.events) + + def test_acknowledged_same_tier_still_projects_runtime(self): + h = self.h + h.result.config_acknowledged = True + h.state.legacy_policy = h.result + self.assertIsNone(h.prepare("pro", "pro")) + self.assertIn("apply_policy", h.events) + + def test_runtime_failure_after_mirror_remains_pending(self): + h = self.h + h.state.legacy_policy = h.result + h.state.intent = PolicyIntent(expected_config=h.guard, intent_id="policy") + h.modules["truememory.tier_switch.runtime"].apply_runtime_policy = Mock( + side_effect=RuntimeError("runtime fault") + ) + with self.assertRaisesRegex(RuntimeError, "runtime fault"): + h.prepare("pro", "pro") + self.assertEqual(h.manager._last_outcome, "activation_pending") + self.assertNotIn("ack", h.events) + + def test_config_failure_never_applies_or_acknowledges(self): + h = self.h + h.state.intent = PolicyIntent(expected_config=h.guard, intent_id="policy") + h.modules["truememory.tier_switch.projection"].mirror_activation = Mock( + side_effect=RuntimeError("config fault") + ) + with self.assertRaisesRegex(RuntimeError, "config fault"): + h.prepare("base", "pro") + self.assertEqual(h.manager._last_outcome, "activation_pending") + self.assertNotIn("apply_policy", h.events) + self.assertNotIn("ack", h.events) + + def test_commit_ambiguity_is_pending_and_issues_no_followup_sql(self): + h = self.h + h.modules[ + "truememory.tier_switch.activation" + ].commit_config_only_transition = Mock( + side_effect=sqlite3.OperationalError("commit") + ) + with self.assertRaises(sqlite3.Error): + h.prepare("base", "pro") + self.assertEqual(h.manager._last_outcome, "activation_pending") + self.assertNotIn("mirror", h.events) + self.assertEqual(h.conn.executed, []) + + def test_thread_failure_discards_connection_without_terminal_retry(self): + h = self.h + h.manager._bootstrap_empty_configuration = lambda *args: False + h.manager._prepare_job = lambda *args, **kwargs: (h.job, h.intent, 17) + thread = Mock() + thread.start.side_effect = RuntimeError("thread start") + thread.is_alive.return_value = False + with patch.object(h.manager_module.threading, "Thread", return_value=thread): + with self.assertRaisesRegex(RuntimeError, "thread start"): + h.manager.start_rebuild( + "base", + backup_path=Path("/synthetic/backup"), + db_path=h.job.database_path, + ) + self.assertTrue(h.conn.closed) + self.assertFalse(h.manager._claimed) + self.assertIsNone(h.manager._active_thread) + self.assertEqual(h.manager._last_outcome, "interrupted") + self.assertEqual(h.conn.executed, []) + + def test_background_handoff_contains_only_immutable_identity(self): + h = self.h + h.manager._bootstrap_empty_configuration = lambda *args: False + h.manager._prepare_job = lambda *args, **kwargs: (h.job, h.intent, 17) + thread = Mock() + thread.is_alive.return_value = False + with patch.object( + h.manager_module.threading, "Thread", return_value=thread + ) as factory: + self.assertEqual( + h.manager.start_rebuild("base", db_path=h.job.database_path), 17 + ) + self.assertTrue(h.conn.closed) + self.assertEqual(factory.call_args.kwargs["args"], (h.job, h.intent, 17)) + self.assertNotIn(h.conn, factory.call_args.kwargs["args"]) + + def test_status_insert_fault_discards_connection(self): + h = self.h + h.manager._bootstrap_empty_configuration = lambda *args: False + h.manager._prepare_job = Mock( + side_effect=sqlite3.OperationalError("status insert") + ) + with self.assertRaises(sqlite3.Error): + h.manager.start_rebuild("base", db_path=h.job.database_path) + self.assertTrue(h.conn.closed) + self.assertFalse(h.manager._claimed) + self.assertEqual(h.conn.executed, []) + + def test_new_store_only_persists_onboarding_config(self): + h = self.h + with patch.object(Path, "stat", side_effect=FileNotFoundError): + self.assertTrue( + h.manager._bootstrap_empty_configuration( + h.job.database_path, "pro", False + ) + ) + self.assertEqual(h.config, {"tier": "pro", "independent": "keep"}) + self.assertNotIn("policy_commit", h.events) + self.assertNotIn("apply_policy", h.events) + self.assertNotIn("tier_activation_generation", h.config) + + def test_tagged_config_cannot_bootstrap_missing_store(self): + h = self.h + h.config["tier_activation_generation"] = "old" + with ( + patch.object(Path, "stat", side_effect=FileNotFoundError), + self.assertRaises(h.manager_module.TierSwitchUnsupportedError), + ): + h.manager._bootstrap_empty_configuration(h.job.database_path, "pro", False) + self.assertEqual(h.config["tier"], "edge") + + def test_nonempty_schema_is_not_bootstrapped(self): + h = self.h + conn = sqlite3.connect(":memory:") + conn.execute("CREATE VIEW present AS SELECT 1") + with ( + patch.object(Path, "stat", return_value=object()), + patch.object(h.manager_module.sqlite3, "connect", return_value=conn), + ): + self.assertFalse( + h.manager._bootstrap_empty_configuration( + h.job.database_path, "pro", False + ) + ) + self.assertEqual(h.config["tier"], "edge") + + def test_uninitialized_empty_schema_bootstraps_without_database_writes(self): + h = self.h + conn = sqlite3.connect(":memory:") + queries = [] + conn.set_trace_callback(queries.append) + with ( + patch.object(Path, "stat", return_value=object()), + patch.object(h.manager_module.sqlite3, "connect", return_value=conn), + ): + self.assertTrue( + h.manager._bootstrap_empty_configuration( + h.job.database_path, "pro", False + ) + ) + self.assertEqual(len(queries), 1) + self.assertTrue(queries[0].startswith("SELECT")) + + def test_existing_unsupported_database_refuses_before_pragmas(self): + h = self.h + conn = sqlite3.connect(":memory:") + conn.execute("CREATE TABLE metadata(wrong BLOB)") + queries = [] + conn.set_trace_callback(queries.append) + h.modules["truememory.tier_switch.activation"].read_activation_state = Mock( + side_effect=RuntimeError("unsupported metadata") + ) + opener = h.load( + "truememory/tier_switch/manager.py", + {"_open_db"}, + {"TierSwitchUnsupportedError": RuntimeError}, + ) + with ( + patch.object(opener.sqlite3, "connect", return_value=conn), + self.assertRaisesRegex(RuntimeError, "unsupported metadata"), + ): + opener._open_db(h.job.database_path, initialize=True) + self.assertEqual(len(queries), 1) + self.assertTrue(queries[0].startswith("SELECT")) + h.modules["truememory.storage"].create_db.assert_not_called() + + def test_existing_opener_does_not_create_missing_database(self): + h = self.h + opener = h.load( + "truememory/tier_switch/manager.py", + {"_open_db"}, + {"TierSwitchUnsupportedError": RuntimeError}, + ) + with patch.object( + opener.sqlite3, "connect", side_effect=sqlite3.OperationalError("missing") + ) as connect: + with self.assertRaises(sqlite3.Error): + opener._open_db(h.job.database_path) + self.assertIn("?mode=rw", connect.call_args.args[0]) + + def test_cli_config_only_never_probes(self): + h = self.h + h.modules["truememory.tier_switch.manager"].RebuildManager = lambda: ( + types.SimpleNamespace( + run_rebuild_sync=lambda tier: True, _last_action="config_only" + ) + ) + cli = h.load("truememory/ingest/cli.py", {"_setup_model_readiness"}) + self.assertEqual(cli._setup_model_readiness("pro", "base"), []) + h.modules[ + "truememory.tier_switch.runtime" + ].serving_operation.assert_not_called() + + def test_cli_same_tier_invokes_authoritative_manager(self): + h = self.h + run = Mock(side_effect=RuntimeError("pending")) + h.modules["truememory.tier_switch.manager"].RebuildManager = lambda: ( + types.SimpleNamespace(run_rebuild_sync=run) + ) + cli = h.load("truememory/ingest/cli.py", {"_setup_model_readiness"}) + self.assertEqual(cli._setup_model_readiness("pro", "pro"), ["tier activation"]) + run.assert_called_once_with("pro") + + def test_both_completion_flags_required(self): + h = self.h + h.state.intent = h.intent + plan = types.SimpleNamespace( + action="resume", manifest=types.SimpleNamespace(consumed=1, total=1) + ) + h.modules["truememory.tier_switch.source"].plan_tier_source = ( + lambda *args, **kwargs: plan + ) + h.modules["truememory.tier_switch.source"].initialize_tier_source = Mock( + side_effect=AssertionError("no init") + ) + worker = Mock() + worker.run_source.return_value = types.SimpleNamespace( + status="complete", + captured_range_complete=False, + current_source_complete=True, + plan=plan, + ) + h.manager_module.RebuildWorker = lambda *args: worker + h.manager._terminal_status = lambda *args: h.events.append("terminal") + h.manager._finalize_rebuild = Mock( + side_effect=AssertionError("cannot finalize") + ) + self.assertFalse(h.manager._run_job(h.conn, h.job, h.intent, 17)) + h.manager._finalize_rebuild.assert_not_called() + self.assertEqual(h.manager._last_outcome, "failed") + + def test_source_conflict_does_not_finalize_or_retry(self): + h = self.h + h.state.intent = h.intent + planner = Mock(side_effect=RuntimeError("source changed")) + h.modules["truememory.tier_switch.source"].plan_tier_source = planner + h.modules["truememory.tier_switch.source"].initialize_tier_source = Mock() + h.manager._finalize_rebuild = Mock() + with self.assertRaisesRegex(RuntimeError, "source changed"): + h.manager._run_job(h.conn, h.job, h.intent, 17) + self.assertEqual(planner.call_count, 1) + h.manager._finalize_rebuild.assert_not_called() + self.assertEqual(h.conn.executed, []) + + def test_local_cancel_rejects_wrong_status(self): + h = self.h + h.manager._active_worker = Mock() + h.manager._active_status_id = 17 + with self.assertRaises(h.manager_module.TierSwitchUnsupportedError): + h.manager.cancel(99) + h.manager._active_worker.cancel.assert_not_called() + self.assertFalse(h.manager._cancel_requested.is_set()) + + def test_interrupted_cancel_requires_exact_job(self): + h = self.h + h.marker = types.SimpleNamespace(**vars(h.job), state="selected") + h.state.intent = h.intent + with self.assertRaises(h.manager_module.TierSwitchUnsupportedError): + h.manager.cancel(17, expected_job_id="wrong") + self.assertTrue(h.conn.closed) + self.assertFalse( + any(isinstance(x, tuple) and x[0] == "cancel" for x in h.events) + ) + + def test_interrupted_cancel_cas_uses_exact_marker(self): + h = self.h + h.marker = types.SimpleNamespace(**vars(h.job), state="selected") + h.state.intent = h.intent + conn = sqlite3.connect(":memory:") + h.manager_module._open_db = lambda *args: conn + result = h.manager.cancel(17, expected_job_id=h.job.job_id) + self.assertEqual(result, {"status": "cancelled", "job_id": h.job.job_id}) + writes = [x for x in h.events if isinstance(x, tuple) and x[0] == "cancel"] + self.assertEqual(len(writes), 1) + self.assertEqual(writes[0][1]["job_id"], h.job.job_id) + + def test_status_backup_path_round_trip_real_sql(self): + h = self.h + conn = sqlite3.connect(":memory:") + self.addCleanup(conn.close) + tree = ast.parse((ROOT / "truememory/tier_switch/manager.py").read_text(encoding="utf-8")) + method = next( + n + for n in ast.walk(tree) + if isinstance(n, ast.FunctionDef) and n.name == "_guard_status" + ) + sql = next( + n.value + for n in ast.walk(method) + if isinstance(n, ast.Constant) + and isinstance(n.value, str) + and n.value.startswith("CREATE TABLE rebuild_status") + ) + conn.execute(sql) + sid = h.manager._create_status_row( + conn, + "basepro", + "base", + "streamed", + 0, + backup_path=Path("/synthetic/backup"), + ) + self.assertEqual( + conn.execute( + "SELECT backup_path FROM rebuild_status WHERE id=?", (sid,) + ).fetchone(), + ("/synthetic/backup",), + ) + self.assertFalse(conn.in_transaction) + + def test_status_trigger_refused_before_insert(self): + h = self.h + conn = sqlite3.connect(":memory:") + self.addCleanup(conn.close) + tree = ast.parse((ROOT / "truememory/tier_switch/manager.py").read_text(encoding="utf-8")) + method = next( + n + for n in ast.walk(tree) + if isinstance(n, ast.FunctionDef) and n.name == "_guard_status" + ) + sql = next( + n.value + for n in ast.walk(method) + if isinstance(n, ast.Constant) + and isinstance(n.value, str) + and n.value.startswith("CREATE TABLE rebuild_status") + ) + conn.execute(sql) + conn.execute( + "CREATE TRIGGER side_effect AFTER INSERT ON rebuild_status BEGIN SELECT 1; END" + ) + with self.assertRaises(h.manager_module.TierSwitchUnsupportedError): + h.manager._create_status_row(conn, "basepro", "base", "streamed", 0) + self.assertEqual( + conn.execute("SELECT count(*) FROM rebuild_status").fetchone(), (0,) + ) + + def test_nonempty_inactive_pair_requires_receipt(self): + h = self.h + conn = sqlite3.connect(":memory:") + self.addCleanup(conn.close) + for name in h.job.target.tables: + conn.execute(f'CREATE TABLE "{name}"(embedding BLOB)') + conn.execute("INSERT INTO vec_messages_basepro VALUES (x'01')") + conn.commit() + before = conn.total_changes + with self.assertRaises(h.manager_module.TierSwitchUnsupportedError): + h.manager._check_inactive_target(conn, h.job.target) + self.assertEqual(conn.total_changes, before) + + def test_configure_queued_clears_memory_and_reports_served_tier(self): + h = self.h + manager = types.SimpleNamespace( + start_rebuild=Mock(return_value=42), _last_outcome="preparing" + ) + module = configure_harness(h, manager) + with patch.object(Path, "home", return_value=MemoryPath()): + result = json.loads(module.truememory_configure("base")) + self.assertIsNone(module._memory) + self.assertEqual(result["tier"], "base") + self.assertEqual(result["served_tier"], "edge") + self.assertEqual(result["status"], "building") + self.assertEqual(result["status_id"], 42) + self.assertEqual(h.config["tier"], "edge") + + def test_configure_failure_clears_memory_and_keeps_served_config(self): + h = self.h + manager = types.SimpleNamespace( + start_rebuild=Mock(side_effect=RuntimeError("synthetic fault")), + _last_outcome="interrupted", + ) + module = configure_harness(h, manager) + with ( + patch.object(Path, "home", return_value=MemoryPath()), + self.assertLogs("synthetic_configure", "ERROR"), + ): + result = json.loads(module.truememory_configure("base")) + self.assertIsNone(module._memory) + self.assertEqual(result["served_tier"], "edge") + self.assertEqual(result["status"], "switch_failed") + self.assertIn("synthetic fault", result["rebuild_error"]) + self.assertEqual(h.config["tier"], "edge") + + def test_configure_same_tier_does_not_bypass_pending(self): + h = self.h + h.config["tier"] = "pro" + manager = types.SimpleNamespace( + start_rebuild=Mock(side_effect=RuntimeError("runtime pending")), + _last_outcome="activation_pending", + ) + module = configure_harness(h, manager) + with ( + patch.object(Path, "home", return_value=MemoryPath()), + self.assertLogs("synthetic_configure", "ERROR"), + ): + result = json.loads( + module.truememory_configure("pro", search_intensity="enhanced") + ) + self.assertEqual(result["status"], "activation_pending") + self.assertEqual(manager.start_rebuild.call_count, 1) + self.assertEqual(h.config["search_intensity"], "enhanced") + self.assertIsNone(module._memory) + + def test_superseded_policy_status_is_not_running(self): + h = self.h + conn = status_connection() + sid = h.manager._create_status_row(conn, "basepro", "base", "streamed", 0) + h.manager_module._open_db = lambda *args: conn + h.state.intent = PolicyIntent() + result = h.manager.get_status(sid) + self.assertEqual(result["status"], "superseded") + self.assertFalse(result["activated"]) + + def test_status_reader_bounds_text_before_python_materialization(self): + h = self.h + conn = status_connection() + sid = h.manager._create_status_row(conn, "basepro", "base", "streamed", 0) + conn.execute( + "UPDATE rebuild_status SET error=?,backup_path=?", + ("e" * 100000, "b" * 100000), + ) + conn.commit() + queries = [] + conn.set_trace_callback(queries.append) + h.manager_module._open_db = lambda *args: conn + result = h.manager.get_status(sid) + self.assertIsNone(result["error"]) + self.assertIsNone(result["backup_path"]) + self.assertFalse(any("SELECT *" in query for query in queries)) + self.assertTrue( + any("length(CAST(error AS BLOB))<=4096" in query for query in queries) + ) + + def test_orphan_marker_exposed_without_claiming_row_ownership(self): + h = self.h + conn = status_connection() + sid = h.manager._create_status_row(conn, "basepro", "base", "streamed", 0) + h.manager_module._open_db = lambda *args: conn + h.marker = types.SimpleNamespace(job_id=h.job.job_id, state="selected") + result = h.manager.get_status(sid) + self.assertEqual(result["status"], "interrupted") + self.assertEqual(result["selected_job_id"], h.job.job_id) + self.assertNotIn("job_id", result) + + def test_status_sql_fault_returns_unknown_and_discards_connection(self): + h = self.h + h.manager._guard_status = Mock( + side_effect=sqlite3.OperationalError("status fault") + ) + conn = sqlite3.connect(":memory:") + h.manager_module._open_db = lambda *args: conn + result = h.manager.get_status(17) + self.assertEqual(result, {"status": "unknown", "error": "OperationalError"}) + with self.assertRaises(sqlite3.ProgrammingError): + conn.execute("SELECT 1") + + def test_local_cancel_rejects_wrong_explicit_job(self): + h = self.h + h.manager._claimed = True + h.manager._active_status_id = 17 + h.manager._active_job_id = h.job.job_id + with self.assertRaises(h.manager_module.TierSwitchUnsupportedError): + h.manager.cancel(17, expected_job_id="wrong") + self.assertFalse(h.manager._cancel_requested.is_set()) + + def test_public_force_refuses_before_database_open(self): + h = self.h + h.manager_module._open_db = Mock( + side_effect=AssertionError("must not initialize") + ) + with self.assertRaises(h.manager_module.TierSwitchUnsupportedError): + h.manager.start_rebuild("base", force=True, db_path=h.job.database_path) + h.manager_module._open_db.assert_not_called() + self.assertFalse(h.manager._claimed) + + def test_bootstrap_unloaded_runtime_projects_lazy_identity(self): + h = self.h + with patch.object(Path, "stat", side_effect=FileNotFoundError): + self.assertTrue( + h.manager._bootstrap_empty_configuration( + h.job.database_path, "pro", False + ) + ) + self.assertEqual(h.manager._last_outcome, "complete") + self.assertEqual(h.manager._served_tier, "pro") + self.assertTrue( + h.manager._bootstrap_empty_configuration( + h.job.database_path, "pro", False + ) + ) + self.assertEqual(h.manager._last_outcome, "complete") + self.assertEqual(h.manager._served_tier, "pro") + self.assertEqual(h.config["tier"], "pro") + self.assertEqual( + h.modules["truememory.vector_search"].EMBEDDING_MODEL, "qwen3_256" + ) + self.assertEqual(h.events.count("lazy_embedding"), 1) + self.assertEqual(h.events.count("lazy_reranker"), 1) + self.assertIsNone(h.modules["truememory.vector_search"]._model) + + def test_configure_bootstrap_pending_reports_runtime_served_tier(self): + h = self.h + h.config["tier"] = "pro" + manager = types.SimpleNamespace( + start_rebuild=Mock(return_value=0), + _last_outcome="activation_pending", + _last_action="onboarding", + _served_tier="edge", + ) + module = configure_harness(h, manager) + with patch.object(Path, "home", return_value=MemoryPath()): + result = json.loads(module.truememory_configure("pro")) + self.assertEqual(result["status"], "activation_pending") + self.assertEqual(result["tier"], "pro") + self.assertEqual(result["served_tier"], "edge") + self.assertNotIn("note", result) + + def test_committed_tier_reported_even_when_config_mirror_fails(self): + h = self.h + h.modules["truememory.tier_switch.projection"].mirror_activation = Mock( + side_effect=RuntimeError("mirror") + ) + with self.assertRaisesRegex(RuntimeError, "mirror"): + h.manager._project_result(h.conn, h.result, h.intent, policy=True) + self.assertEqual(h.manager._served_tier, "pro") + self.assertEqual(h.manager._last_outcome, "activation_pending") + + def test_ambiguous_public_request_discards_before_fresh_readback(self): + h = self.h + h.manager._bootstrap_empty_configuration = lambda *args: False + fresh = FakeConnection() + calls = [] + + def opener(*args, **kwargs): + calls.append(kwargs) + if len(calls) == 1: + return h.conn + self.assertTrue(h.conn.closed) + return fresh + + def failing_prepare(*args, **kwargs): + h.manager._last_outcome = "activation_pending" + raise sqlite3.OperationalError("commit ambiguity") + + h.manager_module._open_db = opener + h.manager._prepare_job = failing_prepare + h.state.legacy_policy = h.result + with self.assertRaises(sqlite3.Error): + h.manager.start_rebuild("pro", db_path=h.job.database_path) + self.assertEqual(len(calls), 2) + self.assertTrue(fresh.closed) + self.assertEqual(h.manager._served_tier, "pro") + self.assertEqual(h.conn.executed, []) + + def test_unknown_readback_does_not_claim_old_config_served(self): + h = self.h + h.config["tier"] = "base" + manager = types.SimpleNamespace( + start_rebuild=Mock(return_value=0), + _last_outcome="activation_pending", + _served_tier=None, + ) + module = configure_harness(h, manager) + with patch.object(Path, "home", return_value=MemoryPath()): + result = json.loads(module.truememory_configure("pro")) + self.assertIsNone(result["served_tier"]) + self.assertEqual(result["status"], "activation_pending") + + def test_bootstrap_loaded_mismatch_refuses_before_config(self): + h = self.h + h.modules["truememory.vector_search"]._model = object() + with ( + patch.object(Path, "stat", side_effect=FileNotFoundError), + self.assertRaises(h.manager_module.TierSwitchUnsupportedError), + ): + h.manager._bootstrap_empty_configuration(h.job.database_path, "pro", False) + self.assertEqual(h.manager._served_tier, "edge") + self.assertEqual(h.config["tier"], "edge") + self.assertNotIn("config", h.events) + self.assertNotIn("lazy_embedding", h.events) + + def test_bootstrap_loaded_reranker_refuses_before_config(self): + h = self.h + h.modules["truememory"].reranker._model = object() + with ( + patch.object(Path, "stat", side_effect=FileNotFoundError), + self.assertRaises(h.manager_module.TierSwitchUnsupportedError), + ): + h.manager._bootstrap_empty_configuration(h.job.database_path, "pro", False) + self.assertEqual(h.config["tier"], "edge") + self.assertNotIn("lazy_embedding", h.events) + + def test_bootstrap_lazy_projection_error_never_writes_config(self): + h = self.h + h.modules["truememory.vector_search"].set_embedding_model = Mock( + side_effect=RuntimeError("projection fault") + ) + with ( + patch.object(Path, "stat", side_effect=FileNotFoundError), + self.assertRaisesRegex(RuntimeError, "projection fault"), + ): + h.manager._bootstrap_empty_configuration(h.job.database_path, "pro", False) + self.assertEqual(h.manager._last_outcome, "activation_pending") + self.assertEqual(h.config["tier"], "edge") + self.assertNotIn("config", h.events) + + def test_bootstrap_config_failure_restores_proven_original_unloaded_identity(self): + h = self.h + h.modules["truememory.tier_switch.projection"]._write_config = Mock( + side_effect=OSError("before replace") + ) + with ( + patch.object(Path, "stat", side_effect=FileNotFoundError), + self.assertRaisesRegex(OSError, "before replace"), + ): + h.manager._bootstrap_empty_configuration(h.job.database_path, "pro", False) + self.assertEqual(h.config["tier"], "edge") + self.assertEqual( + h.modules["truememory.vector_search"].EMBEDDING_MODEL, "model2vec" + ) + self.assertEqual(h.modules["truememory.vector_search"].resolve_tier(), "edge") + self.assertEqual(h.manager._served_tier, "edge") + self.assertIsNone(h.modules["truememory.vector_search"]._model) + + def test_bootstrap_post_replace_fault_preserves_matching_new_identity(self): + h = self.h + + def replaced_then_failed(path, config): + h.config = config.copy() + raise OSError("after replace") + + h.modules[ + "truememory.tier_switch.projection" + ]._write_config = replaced_then_failed + with ( + patch.object(Path, "stat", side_effect=FileNotFoundError), + self.assertRaisesRegex(OSError, "after replace"), + ): + h.manager._bootstrap_empty_configuration(h.job.database_path, "pro", False) + self.assertEqual(h.config["tier"], "pro") + self.assertEqual( + h.modules["truememory.vector_search"].EMBEDDING_MODEL, "qwen3_256" + ) + self.assertEqual(h.manager._served_tier, "pro") + self.assertEqual(h.manager._last_outcome, "activation_pending") + + def test_unreadable_store_cannot_be_treated_as_absent(self): + h = self.h + with ( + patch.object(Path, "stat", side_effect=PermissionError), + self.assertRaises(PermissionError), + ): + h.manager._bootstrap_empty_configuration(h.job.database_path, "pro", False) + self.assertEqual(h.config["tier"], "edge") + self.assertNotIn("lazy_embedding", h.events) + + def test_cancel_cli_dispatch_uses_only_exact_supplied_identity(self): + h = self.h + manager = types.SimpleNamespace( + cancel=Mock(return_value={"status": "cancelled", "job_id": h.job.job_id}) + ) + module, factory = cancel_cli_harness(h, manager) + output = io.StringIO() + with ( + patch.object( + sys, + "argv", + [ + "truememory-ingest", + "cancel-rebuild", + "--status-id", + "17", + "--job-id", + h.job.job_id, + "--db", + "/synthetic/store.db", + ], + ), + contextlib.redirect_stdout(output), + ): + module.main() + factory.assert_called_once_with() + manager.cancel.assert_called_once_with( + 17, expected_job_id=h.job.job_id, db_path=Path("/synthetic/store.db") + ) + self.assertIn("Rebuild 17 cancelled", output.getvalue()) + self.assertIn("explicitly", output.getvalue()) + self.assertEqual(h.events, []) + + def test_cancel_cli_parser_requires_both_exact_ids_before_manager(self): + h = self.h + manager = types.SimpleNamespace(cancel=Mock()) + module, factory = cancel_cli_harness(h, manager) + invalid = [[], ["--status-id", "17"], ["--job-id", h.job.job_id]] + invalid += [ + ["--status-id", value, "--job-id", h.job.job_id] + for value in ( + "0", + "-1", + "+17", + " 17", + "017", + "1_7", + "1.0", + "9223372036854775808", + "9" * 200, + ) + ] + invalid += [ + ["--status-id", "17", "--job-id", value] + for value in ("a" * 31, "a" * 33, "A" * 32, "g" * 32, " " + h.job.job_id) + ] + for arguments in invalid: + with ( + self.subTest(arguments=arguments), + patch.object( + sys, "argv", ["truememory-ingest", "cancel-rebuild", *arguments] + ), + contextlib.redirect_stderr(io.StringIO()), + ): + with self.assertRaises(SystemExit) as error: + module.main() + self.assertEqual(error.exception.code, 2) + factory.assert_not_called() + manager.cancel.assert_not_called() + + def test_cancel_cli_direct_handler_rejects_bool_and_invalid_job(self): + h = self.h + manager = types.SimpleNamespace(cancel=Mock()) + module, factory = cancel_cli_harness(h, manager) + for status_id, job_id in ( + (True, h.job.job_id), + (0, h.job.job_id), + (17, "A" * 32), + ): + with ( + self.subTest(status_id=status_id, job_id=job_id), + contextlib.redirect_stderr(io.StringIO()), + ): + with self.assertRaises(SystemExit) as error: + module._run_cancel_rebuild( + argparse.Namespace(status_id=status_id, job_id=job_id, db=None) + ) + self.assertEqual(error.exception.code, 2) + factory.assert_not_called() + + def test_cancel_cli_accepts_maximum_positive_sqlite_status_without_substitution( + self, + ): + h = self.h + manager = types.SimpleNamespace( + cancel=Mock(return_value={"status": "cancelled", "job_id": h.job.job_id}) + ) + module, _ = cancel_cli_harness(h, manager) + with ( + patch.object( + sys, + "argv", + [ + "truememory-ingest", + "cancel-rebuild", + "--status-id", + "9223372036854775807", + "--job-id", + h.job.job_id, + ], + ), + contextlib.redirect_stdout(io.StringIO()), + ): + module.main() + manager.cancel.assert_called_once_with( + 9223372036854775807, expected_job_id=h.job.job_id, db_path=None + ) + + def test_cancel_cli_busy_failure_is_nonzero_and_never_retries(self): + h = self.h + manager = types.SimpleNamespace( + cancel=Mock(side_effect=RuntimeError("maintenance owner busy")) + ) + module, _ = cancel_cli_harness(h, manager) + output = io.StringIO() + with contextlib.redirect_stderr(output), self.assertRaises(SystemExit) as error: + module._run_cancel_rebuild( + argparse.Namespace(status_id=17, job_id=h.job.job_id, db=None) + ) + self.assertEqual(error.exception.code, 1) + manager.cancel.assert_called_once_with( + 17, expected_job_id=h.job.job_id, db_path=None + ) + self.assertIn("maintenance owner busy", output.getvalue()) + self.assertIn("truememory_status(status_id)", output.getvalue()) + + def test_cancel_cli_nonterminal_or_mismatched_result_never_claims_cancelled(self): + h = self.h + for result in ( + {"status": "no_selected_job"}, + {"status": "cancellation_requested", "status_id": 17}, + {"status": "cancelled", "job_id": "b" * 32}, + ): + manager = types.SimpleNamespace(cancel=Mock(return_value=result)) + module, _ = cancel_cli_harness(h, manager) + output = io.StringIO() + with ( + self.subTest(result=result), + contextlib.redirect_stdout(output), + contextlib.redirect_stderr(io.StringIO()), + ): + with self.assertRaises(SystemExit) as error: + module._run_cancel_rebuild( + argparse.Namespace(status_id=17, job_id=h.job.job_id, db=None) + ) + self.assertEqual(error.exception.code, 1) + self.assertEqual(output.getvalue(), "") + self.assertEqual(manager.cancel.call_count, 1) + + def test_cancel_cli_help_names_existing_status_and_explicit_retry(self): + h = self.h + manager = types.SimpleNamespace(cancel=Mock()) + module, factory = cancel_cli_harness(h, manager) + output = io.StringIO() + with ( + patch.object( + sys, "argv", ["truememory-ingest", "cancel-rebuild", "--help"] + ), + contextlib.redirect_stdout(output), + ): + with self.assertRaises(SystemExit) as error: + module.main() + self.assertEqual(error.exception.code, 0) + self.assertIn("truememory_status(status_id)", output.getvalue()) + self.assertIn("--status-id", output.getvalue()) + self.assertIn("--job-id", output.getvalue()) + self.assertIn("Retry upgrade-tier explicitly", " ".join(output.getvalue().split())) + factory.assert_not_called() + + def test_public_status_recovery_command_is_bound_to_matching_marker(self): + h = self.h + conn = status_connection() + sid = h.manager._create_status_row(conn, "basepro", "base", "streamed", 0) + h.intent.status_id = sid + h.state.intent = h.intent + h.marker = types.SimpleNamespace(**vars(h.job), state="selected") + h.manager_module._open_db = lambda *args: conn + result = h.manager.get_status(sid) + self.assertEqual(result["job_id"], h.job.job_id) + self.assertIn( + f"cancel-rebuild --status-id {sid} --job-id {h.job.job_id}", + result["recovery"], + ) + self.assertIn("same database path", result["recovery"]) + + def test_orphan_status_does_not_offer_bound_cli_cancellation(self): + h = self.h + conn = status_connection() + sid = h.manager._create_status_row(conn, "basepro", "base", "streamed", 0) + h.manager_module._open_db = lambda *args: conn + h.marker = types.SimpleNamespace(job_id=h.job.job_id, state="selected") + result = h.manager.get_status(sid) + self.assertNotIn("job_id", result) + self.assertNotIn("--status-id", result["recovery"]) + self.assertIn("does not prove", result["recovery"]) + + def test_exact_job_with_wrong_positive_status_cannot_cancel_marker(self): + h = self.h + h.marker = types.SimpleNamespace(**vars(h.job), state="selected") + h.state.intent = h.intent + with self.assertRaises(h.manager_module.TierSwitchUnsupportedError): + h.manager.cancel(18, expected_job_id=h.job.job_id) + self.assertTrue(h.conn.closed) + self.assertFalse( + any(isinstance(item, tuple) and item[0] == "cancel" for item in h.events) + ) + + +class ActualLegacyConfigurationComposition(unittest.TestCase): + """Compose actual lazy setters, manager and runtime without native imports.""" + + def setUp(self): + import os + import runpy + import uuid + self.os = os + self.ns = runpy.run_path(str(ROOT / "tests/test-tier-public-boundaries-795.py")) + self.boundary = self.ns["TestPublicBoundaries"]() + self.boundary.setUp() + self.addCleanup(self.boundary.doCleanups) + self.conn, self.api = self.boundary.conn, self.boundary.api + self.selected_marker_row = self.conn.execute("SELECT * FROM truememory_tier_selected_job_v1").fetchone() + self.conn.execute("DELETE FROM truememory_tier_selected_job_v1") + self.conn.executemany("INSERT OR REPLACE INTO metadata(key,value) VALUES (?,?)", + (("embed_model", "model2vec"), ("embed_dim", "256"))) + self.conn.commit() + self.cfg = self.boundary.load_application("tier_config") + self.enterContext(patch.dict(os.environ, {"TRUEMEMORY_EMBED_MODEL": "edge"})) + definitions, namespace = self.ns["definitions"], self.boundary.namespace + self.vector = definitions("vector_search.py", { + "_resolve_model_name", "set_embedding_model", "get_embedding_dim", "resolve_tier", + "_active_tier_group", "_active_vec_table", "_active_sep_table", + }, namespace(os=os, sqlite3=sqlite3, logger=logging.getLogger("synthetic"), + _model=None, _lock=threading.Lock(), _model_generation=0, + _frozen_embedding_target=None, _runtime_policy_tier=None, + EMBEDDING_MODEL="model2vec", _embedding_dim=256, _REMOVED_MODELS={"qwen3"}, + _TIER_ALIASES={tier: config["embed_model"] for tier, config in self.cfg.TIERS.items()}, + _MODEL_DIMS=self.cfg.MODEL_DIMS, _cfg_get_embed_model=self.cfg.get_embed_model, + _cfg_get_embed_dim_for_model=self.cfg.get_embed_dim_for_model, + _cfg_get_model_group=self.cfg.get_model_group)) + self.reranker = definitions("reranker.py", { + "set_active_tier", "get_current_reranker_name", "get_reranker_name_for_tier", + }, namespace(_model=None, _active_tier="edge", _frozen_reranker_id=None, + _TIER_RERANKERS={tier: config["reranker"] for tier, config in self.cfg.TIERS.items()}, + _cfg_get_reranker=self.cfg.get_reranker, log=logging.getLogger("synthetic"))) + self.boundary.fixture.vector, self.boundary.fixture.reranker = self.vector, self.reranker + self.boundary.modules["vector_search"] = self.vector + self.boundary.modules["reranker"] = self.reranker + self.h = Harness() + for name, module in self.boundary.modules.items(): + self.h.modules["truememory." + name] = module + self.h.modules["truememory"] = types.SimpleNamespace(vector_search=self.vector, reranker=self.reranker) + self.h.modules["truememory.tier_config"] = self.cfg + self.h.manager_module.os = os + self.projection = self.h.modules["truememory.tier_switch.projection"] + guard = definitions("tier_switch/projection.py", {"_guard", "ConfigProjectionConflict"}, + namespace(_GENERATION_KEY="tier_activation_generation")) + self.projection._guard = guard._guard + self.h.manager._config_path = Path("/synthetic/config.json") + self.h.manager._config_lock_path = Path("/synthetic/config.lock") + self.owner = definitions("maintenance.py", { + "MaintenanceBusyError", "MaintenanceOwnership", "_acquire_owner", "_release_owner", + "_bind_owner", "maintenance_owner", + }, namespace(NamedTuple=__import__("typing").NamedTuple, + contextmanager=contextlib.contextmanager, Path=Path, + os=types.SimpleNamespace(getpid=os.getpid, close=lambda fd: None), uuid=uuid, + canonical_database_path=lambda value: self.h.job.database_path, + _registry_lock=threading.Lock(), _held_paths={}, _held_fds=set(), + _owner_waiters={}, _coordinators={}, _thread_owners=threading.local(), + try_file_lock=lambda path: 987)) + self.h.modules["truememory.maintenance"] = self.owner + self.h.manager_module._open_db = Mock(side_effect=AssertionError("Unexpected writable/native open")) + self.sql = [] + self.conn.set_trace_callback(self.sql.append) + + def connect_read_snapshot(self): + connection = self.conn + class OwnedHandle: + closed = False + def __getattr__(self, name): + return getattr(connection, name) + def close(self): + self.closed = True + handle = OwnedHandle() + opener = Mock(return_value=handle) + self.h.manager_module.sqlite3 = types.SimpleNamespace(connect=opener) + self.enterContext(patch.object(Path, "stat", return_value=types.SimpleNamespace())) + return handle, opener + + def assert_no_persistent_mutation(self): + self.assertFalse(any(sql.lstrip().split(" ", 1)[0].upper() in { + "INSERT", "UPDATE", "DELETE", "CREATE", "DROP", "ALTER", "REPLACE", + } for sql in self.sql)) + self.assertNotIn("config", self.h.events) + self.assertFalse(self.conn.in_transaction) + self.h.manager_module._open_db.assert_not_called() + self.assertEqual(self.boundary.fixture.loads, []) + + def test_actual_bootstrap_then_fresh_legacy_admission_stays_lazy(self): + for tier in ("base", "pro", "edge"): + with self.subTest(tier=tier), patch.object(Path, "stat", side_effect=FileNotFoundError): + self.assertTrue(self.h.manager._bootstrap_empty_configuration(self.h.job.database_path, tier, False)) + with self.api.serving_operation(self.conn) as operation: + self.assertEqual(operation.tier, tier) + self.assertIsNone(operation.selection) + self.assertIsNone(operation.policy) + self.assertEqual(operation.reranker_id, self.cfg.get_tier_config(tier)["reranker"]) + self.assertEqual(self.vector.resolve_tier(), tier) + self.assertIsNone(self.vector._runtime_policy_tier) + self.assertIsNone(self.vector._frozen_embedding_target) + self.assertIsNone(self.reranker._frozen_reranker_id) + self.assertIsNone(self.vector._model) + self.assertIsNone(self.reranker._model) + self.assertEqual(self.boundary.fixture.loads, []) + self.assertEqual(self.h.config["tier"], tier) + + def test_actual_bootstrap_refuses_controlled_markers_before_ready_or_config(self): + target = self.boundary.modules["embedding_target"].EmbeddingTarget.capture("base") + with self.boundary.fixture.gate.exclusive_activation(): + self.h.manager._apply_bootstrap_identity(target, self.cfg.get_tier_config("base")["reranker"]) + for marker, value in (("_runtime_policy_tier", "base"), ("_frozen_embedding_target", target), + ("_frozen_reranker_id", self.cfg.get_tier_config("base")["reranker"])): + slot = self.reranker if marker == "_frozen_reranker_id" else self.vector + for loaded in (False, True): + with self.subTest(marker=marker, loaded=loaded): + self.vector._model = object() if loaded else None + self.reranker._model = object() if loaded else None + before = self.vector._model, self.reranker._model + setattr(slot, marker, value) + with patch.object(Path, "stat", side_effect=FileNotFoundError): + with self.assertRaisesRegex(self.h.manager_module.TierSwitchUnsupportedError, "controlled runtime"): + self.h.manager._bootstrap_empty_configuration(self.h.job.database_path, "base", False) + self.assertEqual((self.vector._model, self.reranker._model), before) + self.assertEqual(getattr(slot, marker), value) + self.assertEqual(self.h.config["tier"], "edge") + self.assertNotIn("config", self.h.events) + setattr(slot, marker, None) + self.assertEqual(self.boundary.fixture.loads, []) + + def test_actual_bootstrap_preserves_loaded_ordinary_same_tier_certified_reranker(self): + self.vector._model, self.reranker._model = object(), object() + self.reranker._model_certified = True + before = self.vector._model, self.reranker._model + with patch.object(Path, "stat", side_effect=FileNotFoundError): + self.assertTrue(self.h.manager._bootstrap_empty_configuration(self.h.job.database_path, "edge", False)) + self.assertEqual((self.vector._model, self.reranker._model), before) + self.assertTrue(self.reranker._model_certified) + self.assertIsNone(self.reranker._frozen_reranker_id) + with self.api.serving_operation(self.conn) as operation: + self.assertEqual(operation.tier, "edge") + self.assertEqual(self.boundary.fixture.loads, []) + + def test_actual_bootstrap_config_fault_restores_old_legacy_admission(self): + self.projection._write_config = Mock(side_effect=OSError("before replace")) + with patch.object(Path, "stat", side_effect=FileNotFoundError), self.assertRaisesRegex(OSError, "before replace"): + self.h.manager._bootstrap_empty_configuration(self.h.job.database_path, "pro", False) + with self.api.serving_operation(self.conn) as operation: + self.assertEqual(operation.tier, "edge") + self.assertEqual(self.h.config["tier"], "edge") + self.assertEqual(self.vector.EMBEDDING_MODEL, "model2vec") + self.assertIsNone(self.vector._runtime_policy_tier) + self.assertIsNone(self.reranker._frozen_reranker_id) + self.assertEqual(self.boundary.fixture.loads, []) + + def test_actual_bootstrap_post_replace_fault_keeps_new_legacy_admission(self): + def replace_then_fail(path, config): + self.h.config = config.copy() + raise OSError("after replace") + self.projection._write_config = replace_then_fail + with patch.object(Path, "stat", side_effect=FileNotFoundError), self.assertRaisesRegex(OSError, "after replace"): + self.h.manager._bootstrap_empty_configuration(self.h.job.database_path, "pro", False) + with self.api.serving_operation(self.conn) as operation: + self.assertEqual(operation.tier, "pro") + self.assertEqual(self.h.config["tier"], "pro") + self.assertEqual(self.h.manager._last_outcome, "activation_pending") + self.assertIsNone(self.vector._runtime_policy_tier) + self.assertEqual(self.boundary.fixture.loads, []) + + def test_actual_noop_while_other_maintenance_owner_and_reader_are_held(self): + held, release = threading.Event(), threading.Event() + errors = [] + def maintenance(): + try: + with self.owner.maintenance_owner(self.h.job.database_path), self.boundary.fixture.gate.operation_lease(self.api._legacy_key()): + held.set() + if not release.wait(3): + raise AssertionError("Synthetic owner release timed out") + except BaseException as error: + errors.append(error) + thread = threading.Thread(target=maintenance) + thread.start() + try: + self.assertTrue(held.wait(1)) + with self.assertRaises(self.owner.MaintenanceBusyError): + with self.owner.maintenance_owner(self.h.job.database_path): + self.fail("Negative control unexpectedly acquired busy owner") + self.vector._model, self.reranker._model = object(), object() + before_models = self.vector._model, self.reranker._model + writer = self.boundary.load_application("tier_switch.writer") + captured = writer.capture_writer_selection(self.conn) + handle, opener = self.connect_read_snapshot() + self.assertEqual(self.h.manager.start_rebuild("edge", db_path=self.h.job.database_path), 0) + self.assertEqual(self.h.manager._last_action, "noop") + self.assertEqual(self.h.manager._last_outcome, "complete") + self.assertTrue(handle.closed) + self.assertEqual(opener.call_count, 1) + self.assertTrue(opener.call_args.args[0].endswith("?mode=ro")) + self.assertEqual((self.vector._model, self.reranker._model), before_models) + self.assertEqual(writer.capture_writer_selection(self.conn), captured) + self.conn.execute("BEGIN IMMEDIATE") + try: + writer.require_writer_selection(self.conn, captured) + finally: + self.conn.rollback() + self.assert_no_persistent_mutation() + finally: + release.set() + thread.join(3) + self.assertFalse(thread.is_alive()) + self.assertEqual(errors, []) + + def test_actual_noop_sync_path_avoids_owner(self): + self.connect_read_snapshot() + self.owner.maintenance_owner = Mock(side_effect=AssertionError("Unexpected owner")) + self.assertTrue(self.h.manager._run_rebuild_sync_inner("edge", db_path=self.h.job.database_path)) + self.owner.maintenance_owner.assert_not_called() + self.assert_no_persistent_mutation() + + def test_actual_noop_generation_tagged_config_cannot_bypass_recovery(self): + self.h.config["tier_activation_generation"] = "a" * 32 + handle, _ = self.connect_read_snapshot() + self.assertFalse(self.h.manager._try_legacy_noop(self.h.job.database_path, "edge")) + self.assertTrue(handle.closed) + self.assert_no_persistent_mutation() + + def test_actual_noop_staged_intent_and_orphan_job_cannot_bypass_recovery(self): + self.conn.execute("INSERT INTO truememory_tier_selected_job_v1 VALUES (" + ",".join("?" for _ in self.selected_marker_row) + ")", + self.selected_marker_row) + self.conn.commit() + self.boundary.fixture.fixture.build() + self.connect_read_snapshot() + self.sql.clear() + self.assertFalse(self.h.manager._try_legacy_noop(self.h.job.database_path, "edge")) + self.assert_no_persistent_mutation() + self.conn.execute("DELETE FROM metadata WHERE key='tier_activation_v1'") + self.conn.commit() + self.sql.clear() + self.assertFalse(self.h.manager._try_legacy_noop(self.h.job.database_path, "edge")) + self.assert_no_persistent_mutation() + + def test_actual_noop_policy_and_selected_authority_cannot_bypass_recovery(self): + self.boundary.seed_policy() + self.connect_read_snapshot() + self.sql.clear() + self.assertFalse(self.h.manager._try_legacy_noop(self.h.job.database_path, "pro")) + self.assert_no_persistent_mutation() + + def test_actual_noop_rejects_incoherent_embedding_without_writes(self): + self.vector.EMBEDDING_MODEL = "qwen3_256" + self.connect_read_snapshot() + with self.assertRaisesRegex(self.h.manager_module.TierSwitchUnsupportedError, "Legacy model configuration"): + self.h.manager._try_legacy_noop(self.h.job.database_path, "edge") + self.assert_no_persistent_mutation() + + def test_actual_noop_rejects_incoherent_reranker_without_writes(self): + self.reranker._active_tier = "base" + self.connect_read_snapshot() + with self.assertRaisesRegex(self.h.manager_module.TierSwitchUnsupportedError, "Legacy runtime identity"): + self.h.manager._try_legacy_noop(self.h.job.database_path, "edge") + self.assert_no_persistent_mutation() + + def test_concurrent_actual_configure_reports_preparing_until_status_allocated(self): + h = Harness() + entered, release = threading.Event(), threading.Event() + failures = [] + def bootstrap(*args): + entered.set() + if not release.wait(3): + raise AssertionError("Preparation release timed out") + h.manager._last_outcome = "complete" + return True + h.manager._bootstrap_empty_configuration = bootstrap + def first_call(): + try: + h.manager.start_rebuild("pro", db_path=h.job.database_path) + except BaseException as error: + failures.append(error) + thread = threading.Thread(target=first_call) + thread.start() + try: + self.assertTrue(entered.wait(1)) + module = configure_harness(h, h.manager) + with patch.object(Path, "home", return_value=MemoryPath()): + result = json.loads(module.truememory_configure("pro")) + self.assertTrue(h.manager._claimed) + self.assertEqual(h.manager._active_status_id, 0) + self.assertEqual(h.manager._last_outcome, "preparing") + self.assertEqual(result["status"], "preparing") + self.assertIsNone(result["served_tier"]) + self.assertNotIn("status_id", result) + self.assertNotIn("rebuild_error", result) + self.assertIn("preparation", result["note"]) + self.assertNotIn("TrueMemory is ready", result["next_steps"]) + self.assertIsNone(module._memory) + finally: + release.set() + thread.join(3) + self.assertFalse(thread.is_alive()) + self.assertEqual(failures, []) + self.assertEqual(h.events.count("owner_enter"), 1) + + def test_actual_noop_sql_fault_closes_only_owned_connection(self): + handle, _ = self.connect_read_snapshot() + original = self.conn.execute + def execute(sql, *args): + if "temp.sqlite_master" in sql: + raise sqlite3.OperationalError("synthetic unreadable authority") + return original(sql, *args) + with patch.object(self.conn, "execute", side_effect=execute): + with self.assertRaisesRegex(sqlite3.OperationalError, "synthetic unreadable authority"): + self.h.manager._try_legacy_noop(self.h.job.database_path, "edge") + self.assertTrue(handle.closed) + self.assertIsNone(self.h.manager._last_action) + self.assert_no_persistent_mutation() + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test-tier-public-boundaries-795.py b/tests/test-tier-public-boundaries-795.py new file mode 100644 index 00000000..66e6664d --- /dev/null +++ b/tests/test-tier-public-boundaries-795.py @@ -0,0 +1,2348 @@ +"""Public tier boundaries with production AST, stdlib stubs and memory SQLite.""" +from __future__ import annotations + +import ast +import builtins +import contextlib +import dataclasses +import datetime +import enum +import gc +import logging +import os +from pathlib import Path +import runpy +import sqlite3 +import sys +import threading +import time +import types +import unittest +import uuid +import weakref +from unittest.mock import Mock, patch + +ROOT = Path(__file__).resolve().parents[1] +BRIDGE = runpy.run_path(str(Path(__file__).with_name("test-tier-runtime-bridge-795.py"))) +_REAL_PATH_RESOLVE = Path.resolve + + +def definitions(path: str, names: set[str], namespace: dict, *, methods: bool = False) -> types.ModuleType: + if path == "maintenance.py" and "_acquire_owner" in names: + names = names | {"_AutomaticAdmission", "_automatic_admission_locked"} + namespace = {"threading": threading, "weakref": weakref, **namespace} + tree = ast.parse((ROOT / "truememory" / path).read_text(encoding="utf-8")) + if methods: + cls = next(node for node in tree.body if isinstance(node, ast.ClassDef)) + nodes = [node for node in cls.body if isinstance(node, ast.FunctionDef) and node.name in names] + else: + nodes = [node for node in tree.body if isinstance(node, (ast.FunctionDef, ast.ClassDef)) and node.name in names] + future = ast.ImportFrom(module="__future__", names=[ast.alias(name="annotations")], level=0) + module = types.ModuleType("synthetic_public_" + path.replace("/", "_").replace(".", "_")) + module.__dict__.update(namespace) + with patch.dict(sys.modules, {module.__name__: module}): + exec(compile(ast.fix_missing_locations(ast.Module(body=[future, *nodes], type_ignores=[])), path, "exec"), module.__dict__) + return module + + +class TestPublicBoundaries(unittest.TestCase): + def setUp(self) -> None: + self.fixture = BRIDGE["TestRuntimeBridge"]() + self.fixture.setUp() + self.addCleanup(self.fixture.doCleanups) + self.api, self.conn = self.fixture.api, self.fixture.conn + self.modules = self.fixture.modules + self.modules["maintenance"] = types.SimpleNamespace( + maintenance_owner=lambda path: contextlib.nullcontext(), connection_database_path=lambda conn: None) + self.enterContext(patch.dict(sys.modules, {"truememory.tier_switch.runtime": self.api})) + + def imports(self, name, globals=None, locals=None, fromlist=(), level=0): + if name == "truememory.tier_switch.runtime": + return self.api + return self.fixture.safe_import(name, globals, locals, fromlist, level) + + def namespace(self, **values): + return dict(__builtins__=dict(vars(builtins), __import__=self.imports), **values) + + def test_direct_selected_legacy_vector_mutators_refuse_before_writes(self) -> None: + self.fixture.select() + names = {"_write_embedder_metadata", "_write_embedder_metadata_no_commit", "_check_embedder_compatibility", + "init_vec_table", "_rebuild_table_as_cosine", "_finish_cosine_swap", + "migrate_to_cosine_metric", "migrate_legacy_vec_tables"} + module = definitions("vector_search.py", names, self.namespace()) + queries = [] + self.conn.set_trace_callback(queries.append) + for name in sorted(names): + args = (self.conn, "vec_messages") if name == "_rebuild_table_as_cosine" else ( + (self.conn, "vec_messages", 256) if name == "_finish_cosine_swap" else (self.conn,)) + with self.subTest(name=name), self.assertRaises(self.api.TierRuntimeError): + getattr(module, name)(*args) + self.assertFalse(any(sql.lstrip().split(" ", 1)[0].upper() in {"INSERT", "UPDATE", "DELETE", "CREATE", "DROP", "ALTER"} + for sql in queries)) + self.assertFalse(self.conn.in_transaction) + + def test_legacy_guard_preserves_borrowed_transaction_and_source(self) -> None: + self.conn.execute("INSERT INTO messages(id,content) VALUES(77,'synthetic pending')") + before = (self.conn.commits, self.conn.rollbacks) + self.api.require_legacy_vector_mutation(self.conn) + self.assertEqual((self.conn.commits, self.conn.rollbacks), before) + self.assertTrue(self.conn.in_transaction) + self.assertEqual(self.conn.execute("SELECT content FROM messages WHERE id=77").fetchone()[0], "synthetic pending") + + def test_destructive_ingest_refuses_selected_before_create_or_unlink(self) -> None: + self.fixture.select() + create = Mock(side_effect=AssertionError("Unexpected create")) + module = definitions("engine.py", {"ingest"}, self.namespace(Path=Path, create_db=create), methods=True) + engine = types.SimpleNamespace(conn=self.conn, db_path=Path(":memory:")) + maintenance = types.SimpleNamespace(maintenance_owner=lambda path: contextlib.nullcontext()) + self.modules["maintenance"] = maintenance + with patch.object(Path, "unlink", side_effect=AssertionError("Unexpected unlink")) as unlink: + with self.assertRaises(self.api.TierRuntimeError): + module.ingest(engine, "synthetic-input") + create.assert_not_called() + unlink.assert_not_called() + self.assertFalse(self.conn.in_transaction) + + def test_destructive_ingest_rejects_borrowed_before_any_cleanup(self) -> None: + self.conn.execute("INSERT INTO messages(id,content) VALUES(88,'synthetic pending')") + self.modules["maintenance"] = types.SimpleNamespace(maintenance_owner=lambda path: contextlib.nullcontext()) + before = (self.conn.commits, self.conn.rollbacks) + with self.assertRaises(self.api.TierRuntimeError): + with self.api.destructive_ingest_operation(types.SimpleNamespace(conn=self.conn, db_path=Path(":memory:"))): + self.fail("Unexpected destructive body") + self.assertTrue(self.conn.in_transaction) + self.assertEqual((self.conn.commits, self.conn.rollbacks), before) + + def test_selected_maintenance_declines_without_native_reconciliation(self) -> None: + self.fixture.select() + with self.api.maintenance_serving_operation(self.conn) as allowed: + self.assertFalse(allowed) + self.assertEqual(self.fixture.loads, []) + + def test_cluster_only_runner_defers_before_dependency_or_builder(self) -> None: + self.fixture.select() + events = [] + @contextlib.contextmanager + def owner(path): + events.append("owner") + yield types.SimpleNamespace(generation="synthetic") + events.append("released") + real = self.api.maintenance_serving_operation + @contextlib.contextmanager + def serving(conn, **kwargs): + self.assertEqual(events, ["owner"]) + with real(conn, **kwargs) as value: + events.append("serving-check") + yield value + source = definitions("maintenance.py", {"LayerResult", "_maintenance_embedding_scope", "_selected_cluster_deferred", "run_layers"}, + self.namespace(NamedTuple=__import__("typing").NamedTuple, contextmanager=contextlib.contextmanager, + maintenance_owner=owner, connection_database_path=lambda conn: None)) + spec = types.SimpleNamespace(layer="clusters", connection=self.conn, + resolve_dependency=Mock(side_effect=AssertionError("Unexpected dependency")), + build=Mock(side_effect=AssertionError("Unexpected compute"))) + with patch.object(self.api, "maintenance_serving_operation", side_effect=serving): + result = source.run_layers(self.conn, (spec,), force=True) + self.assertEqual(result[0].outcome, "deferred") + self.assertEqual(result[0].error_category, "SelectedVectorLayerUnsupported") + self.assertFalse(result[0].attempted) + self.assertEqual(events, ["owner", "serving-check", "released"]) + spec.resolve_dependency.assert_not_called() + spec.build.assert_not_called() + + def test_generic_nonvector_scope_has_no_serving_import_or_model_work(self) -> None: + module = definitions("maintenance.py", {"_maintenance_embedding_scope"}, self.namespace(contextmanager=contextlib.contextmanager)) + with patch.object(self.api, "maintenance_serving_operation", side_effect=AssertionError("Unexpected model scope")): + with module._maintenance_embedding_scope(self.conn, False, None) as allowed: + self.assertTrue(allowed) + + def gate(self, names): + tree = ast.parse((ROOT / "truememory/ingest/encoding_gate.py").read_text(encoding="utf-8")) + cls = next(node for node in tree.body if isinstance(node, ast.ClassDef) and node.name == "EncodingGate") + nodes = [node for node in cls.body if isinstance(node, ast.FunctionDef) and node.name in names] + namespace = self.namespace(log=logging.getLogger("synthetic-public")) + exec(compile(ast.fix_missing_locations(ast.Module(body=[ast.ImportFrom(module="__future__", names=[ast.alias(name="annotations")], level=0), *nodes], type_ignores=[])), "actual-gate-methods", "exec"), namespace) + return namespace + + def test_gate_generation_drops_old_native_reference_and_cached_search(self) -> None: + refresh = self.gate({"_refresh_serving_generation"})["_refresh_serving_generation"] + class Encoder: + pass + model = Encoder() + previous = weakref.ref(model) + gate = types.SimpleNamespace(_embed_model=model, _serving_key=("old",), _last_search_results=[{"content": "synthetic stale"}], _pe_available=False) + del model + with self.api.serving_operation(self.conn) as operation: + refresh(gate) + gc.collect() + self.assertIsNone(previous()) + self.assertEqual(gate._serving_key, operation.key) + self.assertEqual(gate._last_search_results, []) + self.assertTrue(gate._pe_available) + retained = object() + gate._embed_model = retained + refresh(gate) + self.assertIs(gate._embed_model, retained) + + def test_typed_rejections_escape_gate_search_and_normal_availability_degrades(self) -> None: + search = self.gate({"_search"})["_search"] + for error in (self.api.TierRuntimeError("synthetic"), self.fixture.gate.ServingLeaseTimeout("synthetic")): + gate = types.SimpleNamespace(user_id="", memory=types.SimpleNamespace(search=Mock(side_effect=error))) + with self.assertRaises(type(error)) as raised: + search(gate, "synthetic") + self.assertIs(raised.exception, error) + gate = types.SimpleNamespace(user_id="", memory=types.SimpleNamespace(search=Mock(side_effect=RuntimeError("synthetic unavailable")))) + self.assertEqual(search(gate, "synthetic"), []) + + def test_receipt_and_writer_errors_are_typed_not_generic_runtime_or_value(self) -> None: + class ReceiptError(ConnectionError): + pass + class WriterError(RuntimeError): + pass + with patch.dict(sys.modules, {"truememory.model_client": types.SimpleNamespace(ProtocolMismatchError=ReceiptError), + "truememory.tier_switch.writer": types.SimpleNamespace(WriterSelectionChanged=WriterError)}): + for error in (ReceiptError("synthetic"), WriterError("synthetic")): + with self.assertRaises(type(error)) as raised: + self.api.raise_if_serving_rejection(error) + self.assertIs(raised.exception, error) + self.api.raise_if_serving_rejection(RuntimeError("ordinary unavailable")) + self.api.raise_if_serving_rejection(ValueError("ordinary unavailable")) + + def pipeline(self): + names = {"IngestionResult", "IngestionPipeline"} + module = definitions("ingest/pipeline.py", names, self.namespace( + dataclass=dataclasses.dataclass, field=dataclasses.field, contextlib=contextlib, + sqlite3=sqlite3, time=time, datetime=datetime, Path=Path, os=os, log=logging.getLogger("synthetic-public"), + _invalidate_recall_after_batch=lambda result: contextlib.nullcontext(), + _dedup_store_lock=contextlib.nullcontext, _safe_log=lambda value: value)) + return module + + def test_store_and_update_typed_errors_never_fall_back_to_add(self) -> None: + module = self.pipeline() + pipeline = module.IngestionPipeline.__new__(module.IngestionPipeline) + pipeline.user_id = "" + error = self.api.TierRuntimeError("synthetic stale") + fallback = Mock(side_effect=AssertionError("Unexpected fallback")) + pipeline.memory = types.SimpleNamespace(_engine=types.SimpleNamespace(add=Mock(side_effect=error)), add=fallback, + update=Mock(side_effect=error)) + fact = types.SimpleNamespace(category="general") + for action in (lambda: pipeline._store_fact("synthetic", fact, ""), + lambda: pipeline._update_fact(1, "synthetic", fact, "")): + with self.assertRaises(self.api.TierRuntimeError) as raised: + action() + self.assertIs(raised.exception, error) + fallback.assert_not_called() + + def test_dedup_typed_search_refusal_does_not_become_add(self) -> None: + module = definitions("ingest/dedup.py", {"DedupAction", "DedupDecision", "check_duplicate"}, self.namespace( + Enum=enum.Enum, dataclass=dataclasses.dataclass, log=logging.getLogger("synthetic-public"))) + error = self.api.TierRuntimeError("synthetic stale") + memory = types.SimpleNamespace(search_vectors=Mock(side_effect=error)) + with self.assertRaises(self.api.TierRuntimeError) as raised: + module.check_duplicate("synthetic", memory) + self.assertIs(raised.exception, error) + memory.search_vectors.side_effect = RuntimeError("synthetic availability") + self.assertEqual(module.check_duplicate("synthetic", memory).action, module.DedupAction.ADD) + + def test_both_pipeline_routes_pin_one_fact_and_release_before_next_extraction(self) -> None: + module = self.pipeline() + pipeline = module.IngestionPipeline.__new__(module.IngestionPipeline) + engine = types.SimpleNamespace(conn=self.conn, _write_lock=threading.Lock(), _open_connection_handle=lambda: None) + scheduling = definitions("engine.py", {"_maybe_auto_consolidate"}, self.namespace(), methods=True) + engine._has_consolidation = False + engine._maybe_auto_consolidate = types.MethodType(scheduling._maybe_auto_consolidate, engine) + pipeline.memory = types.SimpleNamespace(_engine=engine) + pipeline.llm_config, pipeline.use_llm_dedup, pipeline.user_id, pipeline.gate_enabled = None, False, "", True + fact = types.SimpleNamespace(content="synthetic fact", category="general", confidence=1.0) + observed = [] + def check(label): + op = self.api.current_operation(self.conn) + self.assertIsNotNone(op) + observed.append((label, op)) + def evaluate(*args): + check("gate") + return types.SimpleNamespace(should_encode=True, encoding_score=1.0, novelty=1.0, salience=1.0, + prediction_error=1.0, reason="synthetic") + pipeline.gate = types.SimpleNamespace(reset_batch=lambda: None, evaluate=evaluate) + class Action(enum.Enum): + ADD = "add" + UPDATE = "update" + SKIP = "skip" + def dedup(*args, **kwargs): + check("dedup") + return types.SimpleNamespace(action=Action.ADD, fact=fact.content, reason="synthetic", existing_id=None) + pipeline._store_fact = lambda *args: check("store") + def extract(*args): + self.assertIsNone(self.api.current_operation(self.conn)) + return [fact, fact] + module.extract_facts_simple = extract + module.DedupAction, module.check_duplicate = Action, dedup + module.parse_transcript = lambda path: ["synthetic"] + module.format_for_extraction = lambda messages: "synthetic " * 20 + for action in (lambda: pipeline.ingest_text("synthetic " * 20), lambda: pipeline.ingest_transcript("synthetic-input")): + observed.clear() + result = action() + self.assertEqual(result.facts_stored, 2) + self.assertEqual([label for label, op in observed], ["gate", "dedup", "store"] * 2) + self.assertEqual(len({id(op) for label, op in observed[:3]}), 1) + self.assertIsNot(observed[0][1], observed[3][1]) + self.assertIsNone(self.api.current_operation(self.conn)) + + + def load_application(self, name: str) -> types.ModuleType: + module = types.ModuleType("synthetic_public_" + name.replace(".", "_")) + module.__dict__["__builtins__"] = dict(vars(builtins), __import__=self.imports) + self.enterContext(patch.dict(sys.modules, {module.__name__: module})) + path = ROOT / "truememory" / (name.replace(".", "/") + ".py") + exec(compile(path.read_text(encoding="utf-8"), str(path), "exec"), module.__dict__) + self.modules[name] = module + return module + + def seed_policy(self, *, uppercase: bool = False): + activation = self.modules["tier_switch.activation"] + target = self.modules["embedding_target"].EmbeddingTarget("pro", "qwen3_256", 256, "basepro") + pair = ("vec_messages", "vec_messages_sep") + intent = activation.PolicyIntent("a" * 32, None, None, target, + "Alibaba-NLP/gte-reranker-modernbert-base", activation.ConfigGuard(True, "base", None), "b" * 32, pair) + for key, value in (("embed_model", target.model_id), ("embed_dim", "256"), + ("tier_activation_v1", activation._dump(intent))): + activation._put(self.conn, key, value) + self.conn.commit() + if uppercase: + ddl = self.conn.execute("SELECT sql FROM main.sqlite_master WHERE name='metadata'").fetchone()[0] + self.conn.execute("ALTER TABLE metadata RENAME TO temporary_metadata") + self.conn.execute(ddl.replace("metadata", "METADATA", 1)) + self.conn.execute("INSERT INTO METADATA SELECT * FROM temporary_metadata") + self.conn.commit() + return activation.read_activation_state(self.conn).legacy_policy + + def test_policy_reader_does_not_invent_selection_or_allow_legacy_mutation(self) -> None: + policy = self.seed_policy(uppercase=True) + self.assertIsNone(self.api._read_selection(self.conn)) + self.assertEqual(self.api._read_policy(self.conn), policy) + with self.assertRaises(self.api.TierRuntimeError): + self.api.require_legacy_vector_mutation(self.conn) + with self.assertRaises(self.api.TierRuntimeError): + with self.api.destructive_ingest_operation(types.SimpleNamespace(conn=self.conn, db_path=Path(":memory:"))): + self.fail("Policy database reached destructive body") + self.assertFalse(self.conn.in_transaction) + + def test_policy_writer_fences_absence_generation_and_exact_connection(self) -> None: + writer = self.load_application("tier_switch.writer") + old = writer.capture_writer_selection(self.conn) + policy = self.seed_policy() + fresh = writer.capture_writer_selection(self.conn) + self.assertEqual(fresh.policy, policy) + self.assertIsNone(fresh.selection) + for captured in (old, writer.WriterCapture(self.conn, None), None): + with self.subTest(captured=captured), self.assertRaises(writer.WriterSelectionChanged): + with writer.writer_transaction(self.conn, captured): + self.fail("Uncaptured policy was accepted") + with writer.writer_transaction(self.conn, fresh, model_id=policy.target.model_id, + dimension=256, tables=policy.tables): + self.conn.execute("INSERT INTO messages(id,content) VALUES(901,'synthetic policy fact')") + self.assertEqual(self.conn.execute("SELECT content FROM messages WHERE id=901").fetchone()[0], "synthetic policy fact") + self.conn.execute("BEGIN IMMEDIATE") + for captured in (writer.WriterCapture(self.conn, None, dataclasses.replace(policy, generation="c" * 32)), + writer.WriterCapture(object(), None, policy)): + with self.assertRaises(writer.WriterSelectionChanged): + writer.require_writer_selection(self.conn, captured) + writer.require_writer_selection(self.conn, writer.WriterCapture(self.conn, None, + dataclasses.replace(policy, config_acknowledged=True))) + self.assertTrue(self.conn.in_transaction) + self.conn.rollback() + + def install_policy_projectors(self): + client = types.SimpleNamespace() + class EmbeddingProxy: + pass + class CertifiedEmbeddingProxy(EmbeddingProxy): + def __init__(self, target): + self.target = target + class RerankerProxy: + pass + class CertifiedRerankerProxy(RerankerProxy): + def __init__(self, name): + self.name = name + client.EmbeddingProxy, client.CertifiedEmbeddingProxy = EmbeddingProxy, CertifiedEmbeddingProxy + client.RerankerProxy, client.CertifiedRerankerProxy = RerankerProxy, CertifiedRerankerProxy + self.modules["model_client"] = client + vector = definitions("vector_search.py", {"apply_embedding_policy"}, self.namespace( + _model=object(), EMBEDDING_MODEL="qwen3_256", _embedding_dim=256, _model_generation=0, + _target_state_lock=lambda deadline: contextlib.nullcontext(), _target_remaining=lambda deadline: None, + _load_frozen_embedding_target=Mock(side_effect=AssertionError("Unexpected native load")))) + reranker = definitions("reranker.py", {"apply_reranker_policy"}, self.namespace( + _model=object(), _model_name="Alibaba-NLP/gte-reranker-modernbert-base", _model_certified=True, + _reranker_load_lock=lambda control: contextlib.nullcontext())) + self.fixture.vector, self.fixture.reranker = vector, reranker + self.modules["vector_search"], self.modules["reranker"] = vector, reranker + return vector, reranker, client + + def test_same_space_policy_projection_retains_local_models_and_has_zero_loads(self) -> None: + policy = self.seed_policy() + vector, reranker, _ = self.install_policy_projectors() + before = vector._model, reranker._model + with self.fixture.gate.exclusive_activation(): + self.api.apply_runtime_policy(policy) + self.assertIs(vector._model, before[0]) + self.assertIs(reranker._model, before[1]) + self.assertEqual(self.api.runtime_acknowledgement(), self.api.policy_key(policy)) + with self.api.serving_operation(self.conn) as operation: + self.assertEqual(operation.policy, policy) + self.assertIsNone(operation.selection) + self.assertEqual((operation.tables, operation.tier), (policy.tables, "pro")) + with self.api.serving_operation(self.conn) as nested: + self.assertIs(nested, operation) + writer = self.load_application("tier_switch.writer") + self.assertEqual(writer.capture_writer_selection(self.conn).policy, policy) + vector._load_frozen_embedding_target.assert_not_called() + + def test_policy_projection_rebinds_proxies_and_keeps_absent_slots_lazy(self) -> None: + policy = self.seed_policy() + vector, reranker, client = self.install_policy_projectors() + vector._model, reranker._model = client.EmbeddingProxy(), client.RerankerProxy() + self.api.apply_runtime_policy(policy) + self.assertIsInstance(vector._model, client.CertifiedEmbeddingProxy) + self.assertEqual(vector._model.target, policy.target) + self.assertIsInstance(reranker._model, client.CertifiedRerankerProxy) + vector._model = reranker._model = None + self.api.apply_runtime_policy(policy) + self.assertIsNone(vector._model) + self.assertIsNone(reranker._model) + self.assertIsNone(self.api.runtime_acknowledgement()) + vector._load_frozen_embedding_target.assert_not_called() + self.conn.execute("INSERT INTO messages(id,content) VALUES(902,'synthetic borrowed')") + with self.assertRaises(self.api.TierRuntimeError): + with self.api.serving_operation(self.conn): + self.fail("Borrowed policy admission attempted native preparation") + self.assertTrue(self.conn.in_transaction) + self.conn.rollback() + vector._load_frozen_embedding_target.side_effect = None + vector._load_frozen_embedding_target.return_value = object() + def prepare(tier, name, **kwargs): + self.assertFalse(self.conn.in_transaction) + reranker._model, reranker._model_certified = object(), True + reranker.apply_frozen_reranker = prepare + with self.api.serving_operation(self.conn) as operation: + self.assertEqual(operation.policy, policy) + vector._load_frozen_embedding_target.assert_called_once() + + def test_policy_projection_failure_never_acknowledges_partial_slots(self) -> None: + policy = self.seed_policy() + _, reranker, _ = self.install_policy_projectors() + self.api.apply_runtime_policy(policy) + error = MemoryError("synthetic projection failure") + reranker.apply_reranker_policy = Mock(side_effect=error) + with self.assertRaises(MemoryError) as raised: + self.api.apply_runtime_policy(policy) + self.assertIs(raised.exception, error) + self.assertIsNone(self.api.runtime_acknowledgement()) + + def test_gate_prediction_error_releases_reference_on_all_outcomes(self) -> None: + namespace = self.gate({"_compute_prediction_error"}) + namespace["EncodingGate"] = types.SimpleNamespace(_refresh_serving_generation=lambda gate: None) + namespace["_PE_NOISE"] = set() + namespace["np"] = types.SimpleNamespace(linalg=types.SimpleNamespace(norm=lambda value: 1.0), dot=lambda a, b: 1.0) + class Model: + pass + model = Model() + gate = types.SimpleNamespace(_embed_model=model, _last_search_results=[{"content": "synthetic nearest"}], + _pe_available=True, _pe_degradation_count=0) + encode = Mock(return_value=[[1.0]] * 4) + self.modules["mps_utils"] = types.SimpleNamespace(encode_with_model_ownership=encode) + for error in (None, self.api.TierRuntimeError("synthetic refusal"), KeyboardInterrupt("synthetic stop")): + gate._embed_model = model + encode.side_effect = error + if error is None: + self.assertEqual(namespace["_compute_prediction_error"](gate, "synthetic fact"), 0.0) + else: + with self.assertRaises(type(error)) as raised: + namespace["_compute_prediction_error"](gate, "synthetic fact") + self.assertIs(raised.exception, error) + self.assertIsNone(gate._embed_model) + + def test_legacy_mutator_owner_refuses_before_sql_and_does_not_end_borrowed_work(self) -> None: + error = RuntimeError("synthetic owner busy") + self.modules["maintenance"].maintenance_owner = Mock(side_effect=error) + self.conn.execute("INSERT INTO messages(id,content) VALUES(903,'synthetic borrowed')") + before = self.conn.commits, self.conn.rollbacks + with self.api.serving_operation(self.conn): + with self.assertRaises(RuntimeError) as raised: + with self.api.legacy_vector_mutation(self.conn): + self.fail("Contended owner admitted mutation") + self.assertIs(raised.exception, error) + self.assertEqual((self.conn.commits, self.conn.rollbacks), before) + self.assertTrue(self.conn.in_transaction) + + def test_legacy_mutator_rechecks_policy_after_owner_acquisition(self) -> None: + @contextlib.contextmanager + def owner(path): + self.seed_policy() + yield + self.modules["maintenance"].maintenance_owner = owner + with self.assertRaises(self.api.TierRuntimeError): + with self.api.legacy_vector_mutation(self.conn): + self.fail("Policy appearing before owner admission reached writes") + + def test_existing_controlled_open_never_calls_schema_factory_and_closes_on_rejection(self) -> None: + policy = self.seed_policy() + factory = Mock(side_effect=AssertionError("Unexpected legacy initialization")) + self.modules["storage"] = types.SimpleNamespace(DEFAULT_BUSY_TIMEOUT_MS=10000) + path = Path(Path.cwd().anchor) / "synthetic" / "database.db" + statements = [] + self.conn.set_trace_callback(statements.append) + with patch.object(Path, "exists", return_value=True), patch.object(self.api.sqlite3, "connect", return_value=self.conn) as connect: + opened = self.api.open_serving_connection(path, factory) + self.assertIs(opened, self.conn) + self.assertEqual(self.api._read_policy(opened), policy) + self.assertIn("?mode=rw", connect.call_args.args[0]) + self.assertFalse(any(sql.lstrip().split(" ", 1)[0].upper() in {"INSERT", "UPDATE", "DELETE", "CREATE", "DROP", "ALTER"} + for sql in statements)) + factory.assert_not_called() + error = self.api.TierRuntimeError("synthetic invalid journal") + close = Mock() + with patch.object(Path, "exists", return_value=True), patch.object(self.api.sqlite3, "connect", return_value=self.conn), \ + patch.object(self.api, "read_activation_state", side_effect=error), patch.object(self.conn, "close", close): + with self.assertRaises(self.api.TierRuntimeError) as raised: + self.api.open_serving_connection(path, factory) + self.assertIs(raised.exception, error) + close.assert_called_once() + + + def test_nonblocking_owner_breaks_reader_activation_cycle(self) -> None: + path = Path(Path.cwd().anchor) / "synthetic" / "database.db" + owner = definitions("maintenance.py", {"MaintenanceBusyError", "MaintenanceOwnership", "_acquire_owner", + "_release_owner", "_bind_owner", "maintenance_owner"}, self.namespace( + NamedTuple=__import__("typing").NamedTuple, contextmanager=contextlib.contextmanager, + Path=Path, os=types.SimpleNamespace(getpid=os.getpid, close=lambda fd: None), uuid=uuid, + canonical_database_path=lambda value: path, _registry_lock=threading.Lock(), + _held_paths={}, _held_fds=set(), _owner_waiters={}, _coordinators={}, + _thread_owners=threading.local(), try_file_lock=lambda value: 987)) + owner.connection_database_path = lambda conn: path + self.modules["maintenance"] = owner + held, finished = threading.Event(), threading.Event() + entered, errors = [], [] + def activate(): + try: + with owner.maintenance_owner(path): + held.set() + with self.fixture.gate.exclusive_activation(): + entered.append(True) + except BaseException as error: + errors.append(error) + finally: + finished.set() + self.conn.execute("INSERT INTO messages(id,content) VALUES(904,'synthetic borrowed')") + before = self.conn.commits, self.conn.rollbacks + with self.api.serving_operation(self.conn): + worker = threading.Thread(target=activate) + worker.start() + self.assertTrue(held.wait(2)) + with self.assertRaises(owner.MaintenanceBusyError): + with self.api.legacy_vector_mutation(self.conn): + self.fail("Busy activation owner allowed legacy writes") + self.assertEqual(entered, []) + self.assertTrue(self.conn.in_transaction) + self.assertEqual((self.conn.commits, self.conn.rollbacks), before) + self.assertTrue(finished.wait(2)) + worker.join(2) + self.assertFalse(worker.is_alive()) + self.assertEqual(errors, []) + self.assertEqual(entered, [True]) + self.assertEqual(owner._held_paths, {}) + + def test_cold_selected_consolidate_defers_before_model_or_engine_initialization(self) -> None: + self.fixture.select() + events = [] + report = types.SimpleNamespace(results=()) + coordinator = types.SimpleNamespace(path=None, refresh_capabilities=lambda: events.append("capabilities")) + @contextlib.contextmanager + def observation(value, **kwargs): + events.append("owner") + yield [] + lock = threading.Lock() + def run(conn, supplied, **kwargs): + self.assertFalse(lock.locked()) + with self.api.maintenance_serving_operation(conn) as supported: + self.assertFalse(supported) + events.append("deferred") + return report + self.modules["maintenance"] = types.SimpleNamespace( + MaintenanceBusyError=type("Busy", (RuntimeError,), {}), format_maintenance_report=lambda value: {"clusters": "deferred"}, + maintenance_busy_result=lambda: {}, maintenance_observation=observation, + observe_style=Mock(side_effect=AssertionError("Unexpected style observation")), + run_engine_maintenance=run, _style_unready_result=lambda value: value) + source = definitions("engine.py", {"consolidate"}, self.namespace(), methods=True) + engine = types.SimpleNamespace(conn=self.conn, _write_lock=lock, _open_connection_handle=lambda: events.append("open"), + _ensure_connection=Mock(side_effect=AssertionError("Unexpected runtime initialization")), + _get_maintenance_coordinator=lambda: coordinator, _auto_consolidate_threshold=25, + _apply_manual_maintenance_capabilities=lambda results: None) + self.assertEqual(source.consolidate(engine), {"clusters": "deferred"}) + self.assertEqual(events, ["open", "capabilities", "owner", "deferred"]) + self.assertEqual(self.fixture.loads, []) + engine._ensure_connection.assert_not_called() + + def test_selected_foreground_metadata_is_fenced_noop_preserving_caller_sql(self) -> None: + selected = self.fixture.select() + writer = self.load_application("tier_switch.writer") + activation = self.modules["tier_switch.activation"] + activation._put(self.conn, "embed_model", selected.target.model_id) + activation._put(self.conn, "embed_dim", str(selected.target.dimension)) + self.conn.commit() + module = definitions("vector_search.py", {"_write_embedder_metadata_no_commit"}, self.namespace( + EMBEDDING_MODEL=selected.target.model_id, _embedding_dim=selected.target.dimension)) + with self.api.serving_operation(self.conn): + captured = writer.capture_writer_selection(self.conn) + with writer.writer_transaction(self.conn, captured): + changes = self.conn.total_changes + module._write_embedder_metadata_no_commit(self.conn) + self.assertEqual(self.conn.total_changes, changes) + self.assertTrue(self.conn.in_transaction) + module._embedding_dim += 1 + with self.assertRaises(writer.WriterSelectionChanged): + module._write_embedder_metadata_no_commit(self.conn) + self.assertFalse(self.conn.in_transaction) + + def test_legacy_initializer_propagates_typed_guard_or_owner_refusal(self) -> None: + class Busy(RuntimeError): + pass + self.modules["maintenance"].MaintenanceBusyError = Busy + self.enterContext(patch.dict(sys.modules, {"truememory.maintenance": self.modules["maintenance"]})) + def imports(name, globals=None, locals=None, fromlist=(), level=0): + if name == "sqlite_vec": + return types.SimpleNamespace(load=lambda conn: None) + return self.imports(name, globals, locals, fromlist, level) + self.conn.enable_load_extension = lambda value: None + module = definitions("engine.py", {"_initialize_connection"}, dict( + __builtins__=dict(vars(builtins), __import__=imports), _HAS_VECTOR=True, + TrueMemoryMigrationError=type("Migration", (RuntimeError,), {})), methods=True) + engine = types.SimpleNamespace(conn=self.conn, ready=False, _runtime_initialized=False, + _init_lock=threading.Lock(), _write_lock=threading.Lock(), _has_vectors=False) + for error in (self.api.TierRuntimeError("synthetic changed"), Busy("synthetic owner busy")): + module.init_vec_table = Mock(side_effect=error) + with self.api.serving_operation(self.conn): + with self.assertRaises(type(error)) as raised: + module._initialize_connection(engine, _suppress_maintenance=True) + self.assertIs(raised.exception, error) + self.assertFalse(engine.ready) + self.assertFalse(engine._runtime_initialized) + + def test_hybrid_optional_separation_refusal_does_not_degrade(self) -> None: + source = definitions("hybrid.py", {"search_hybrid"}, self.namespace(database_operation=lambda function: function)) + error = self.api.TierRuntimeError("synthetic stale separation") + self.modules["fts_search"] = types.SimpleNamespace(search_fts=Mock()) + self.modules["vector_search"] = types.SimpleNamespace(search_vector=Mock(), search_vector_separation=Mock(), + _active_sep_table=Mock(side_effect=error)) + with self.assertRaises(self.api.TierRuntimeError) as raised: + source.search_hybrid(self.conn, "synthetic") + self.assertIs(raised.exception, error) + self.modules["vector_search"].search_vector.assert_not_called() + self.modules["fts_search"].search_fts.assert_not_called() + + + def test_policy_initialization_validates_exact_pair_without_ddl_or_metadata(self) -> None: + policy = self.seed_policy() + self.install_policy_projectors() + self.api.apply_runtime_policy(policy) + for table in policy.tables: + self.conn.execute(f"CREATE TABLE {table}(embedding BLOB)") + original = type(self.conn).execute + def execute(conn, sql, parameters=()): + if sql == "SELECT sql FROM main.sqlite_master WHERE name=? COLLATE NOCASE" and parameters[0] in policy.tables: + return BRIDGE["JOURNAL"]["Rows"]([(f"CREATE VIRTUAL TABLE {parameters[0]} USING vec0(embedding float[256] distance_metric=cosine)",)]) + return original(conn, sql, parameters) + module = definitions("vector_search.py", {"init_vec_table"}, self.namespace()) + statements = [] + self.conn.set_trace_callback(statements.append) + with patch.object(type(self.conn), "execute", execute): + with self.api.serving_operation(self.conn): + module.init_vec_table(self.conn) + with self.assertRaises(ValueError): + module.init_vec_table(self.conn, tier_group="basepro") + self.assertFalse(any(sql.lstrip().split(" ", 1)[0].upper() in {"INSERT", "UPDATE", "DELETE", "CREATE", "DROP", "ALTER"} + for sql in statements)) + self.assertFalse(self.conn.in_transaction) + + def test_policy_loaded_runtime_cannot_be_inherited_by_unselected_database(self) -> None: + policy = self.seed_policy() + self.install_policy_projectors() + self.api.apply_runtime_policy(policy) + other = sqlite3.connect(":memory:") + self.addCleanup(other.close) + with self.assertRaises(self.api.TierRuntimeError): + with self.api.serving_operation(other): + self.fail("Unselected database silently inherited another policy") + + + def maintenance_runner(self, *, owner=None, run_layers=None): + owner = owner or (lambda path: contextlib.nullcontext()) + module = definitions("maintenance.py", { + "LayerResult", "MaintenanceReport", "_maintenance_embedding_scope", "_maintenance_completion", "run_engine_maintenance", + }, self.namespace( + NamedTuple=__import__("typing").NamedTuple, contextmanager=contextlib.contextmanager, + maintenance_owner=owner, connection_database_path=lambda conn: None, + engine_layer_specs=lambda *args, **kwargs: (), all_layer_specs=lambda conn: (), + run_layers=run_layers or (lambda *args, **kwargs: ()), time=time, + importlib=types.SimpleNamespace(import_module=lambda name: types.SimpleNamespace(extract_preferences=lambda conn: None)), + )) + return module + + def memory_connection(self) -> sqlite3.Connection: + conn = sqlite3.connect(":memory:", check_same_thread=False, + factory=BRIDGE["JOURNAL"]["JournalConnection"]) + self.conn.backup(conn) + self.addCleanup(conn.close) + return conn + + def real_memory_owner(self, acquire): + return definitions("maintenance.py", {"MaintenanceOwnership", "MaintenanceBusyError", + "canonical_database_path", "_bind_owner", "maintenance_owner", "_active_initialization_receipt"}, self.namespace( + NamedTuple=__import__("typing").NamedTuple, contextmanager=contextlib.contextmanager, + Path=Path, os=os, uuid=uuid, _thread_owners=threading.local(), weakref=weakref, + _registry_lock=threading.Lock(), _coordinators=weakref.WeakValueDictionary(), _held_paths={}, + _acquire_owner=acquire, _release_owner=lambda fd: None)) + + def test_maintenance_authority_waits_for_managed_writer_before_sql(self) -> None: + conn = self.memory_connection() + engine = self.fixture.lifecycle_engine("_open_connection_handle", conn) + owned, waiting, release, finished = (threading.Event() for _ in range(4)) + failures, reads = [], [] + api = self.api + + class ObservedLock(api.ConnectionWriteLock): + def acquire(self, blocking=True, timeout=-1): + if self.owns_other_transaction(conn): + waiting.set() + return super().acquire(blocking, timeout) + + engine._write_lock = lock = ObservedLock(engine) + original = api._read_selection + + def read(connection): + self.assertFalse(lock.owns_other_transaction(connection)) + self.assertTrue(lock.locked()) + reads.append(connection) + return original(connection) + + def writer(): + try: + with lock: + conn.execute("INSERT INTO messages(id,content) VALUES(991,'synthetic writer')") + owned.set() + if not release.wait(2): + raise AssertionError("Synthetic writer was not released") + conn.commit() + except BaseException as error: + failures.append(error) + + def reader(): + try: + with api.maintenance_serving_operation(conn, connection_lock=lock, + deadline=time.monotonic() + 2) as allowed: + self.assertTrue(allowed) + self.assertFalse(lock.locked()) + finished.set() + except BaseException as error: + failures.append(error) + + first, second = threading.Thread(target=writer), threading.Thread(target=reader) + with patch.object(api, "_read_selection", side_effect=read): + first.start() + try: + self.assertTrue(owned.wait(2)) + second.start() + self.assertTrue(waiting.wait(2)) + self.assertEqual(reads, []) + finally: + release.set() + first.join(2) + if second.ident is not None: + second.join(2) + self.assertFalse(first.is_alive()) + self.assertFalse(second.is_alive()) + self.assertEqual(failures, []) + self.assertTrue(finished.is_set()) + self.assertGreaterEqual(len(reads), 3) + + def test_maintenance_borrowed_busy_cancel_and_deadline_preserve_sql(self) -> None: + engine = self.fixture.lifecycle_engine("_open_connection_handle") + self.conn.execute("INSERT INTO messages(id,content) VALUES(991,'synthetic caller')") + before = self.conn.commits, self.conn.rollbacks + with engine._write_lock: + with self.assertRaises(self.api.TierRuntimeError): + with self.api.maintenance_serving_operation(self.conn, connection_lock=engine._write_lock): + self.fail("Borrowed busy admission entered maintenance") + cancelled = threading.Event() + cancelled.set() + for controls, kind in (({"cancelled": cancelled}, self.fixture.gate.ServingLeaseCancelled), + ({"deadline": time.monotonic() - 1}, self.fixture.gate.ServingLeaseTimeout)): + with self.subTest(kind=kind.__name__), self.assertRaises(kind): + with self.api.maintenance_serving_operation(self.conn, connection_lock=engine._write_lock, **controls): + self.fail("Rejected admission entered maintenance") + 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=991").fetchone()[0], "synthetic caller") + self.assertFalse(engine._write_lock.locked()) + + def test_nested_maintenance_reuses_lease_without_sql_or_writer_reacquisition(self) -> None: + engine = self.fixture.lifecycle_engine("_open_connection_handle") + with self.api.maintenance_serving_operation(self.conn, connection_lock=engine._write_lock) as allowed: + self.assertTrue(allowed) + operation = self.api.current_operation(self.conn) + with engine._write_lock, \ + patch.object(self.api, "_read_selection", side_effect=AssertionError("Nested authority SQL")), \ + patch.object(self.api, "_read_policy", side_effect=AssertionError("Nested policy SQL")), \ + patch.object(self.api, "_connection_read", side_effect=AssertionError("Nested writer acquisition")): + with self.api.maintenance_serving_operation(self.conn, connection_lock=engine._write_lock) as nested: + self.assertTrue(nested) + self.assertIs(self.api.current_operation(self.conn), operation) + + def test_runtime_final_table_and_database_snapshots_hold_writer(self) -> None: + engine = self.fixture.lifecycle_engine("_open_connection_handle") + observed = [] + + def snapshot(name, value): + def read(conn): + self.assertTrue(engine._write_lock.locked()) + observed.append(name) + return value + return read + + with patch.object(self.fixture.vector, "_active_vec_table", side_effect=snapshot("normal", "vec_messages")), \ + patch.object(self.fixture.vector, "_active_sep_table", side_effect=snapshot("separation", "vec_messages_sep")), \ + patch.object(self.api, "_database_identity", side_effect=snapshot("database", ("memory", self.conn))): + with self.api.serving_operation(self.conn, connection_lock=engine._write_lock) as operation: + self.assertEqual(operation.tables, ("vec_messages", "vec_messages_sep")) + self.assertFalse(engine._write_lock.locked()) + self.assertEqual(observed, ["normal", "separation", "database"]) + + def test_runtime_preparation_and_gate_admission_release_writer(self) -> None: + self.fixture.select() + engine = self.fixture.lifecycle_engine("_open_connection_handle") + prepare, lease = self.api.apply_frozen_selection, self.fixture.gate.operation_lease + + def project(*args, **kwargs): + self.assertFalse(engine._write_lock.locked()) + return prepare(*args, **kwargs) + + @contextlib.contextmanager + def admit(*args, **kwargs): + self.assertFalse(engine._write_lock.locked()) + with lease(*args, **kwargs) as value: + yield value + + with patch.object(self.api, "apply_frozen_selection", side_effect=project), \ + patch.object(self.fixture.gate, "operation_lease", side_effect=admit): + with self.api.serving_operation(self.conn, connection_lock=engine._write_lock): + self.assertFalse(engine._write_lock.locked()) + self.assertEqual(len(self.fixture.loads), 2) + + def test_runner_completion_cleans_owned_sql_under_lock_and_keeps_borrowed_sql(self) -> None: + engine = self.fixture.lifecycle_engine("_open_connection_handle") + + def layers(conn, *args, **kwargs): + self.assertTrue(engine._write_lock.locked()) + conn.execute("INSERT INTO messages(id,content) VALUES(992,'synthetic unfinished')") + return () + + module = self.maintenance_runner(run_layers=layers) + rollback = self.conn.rollback + calls = [] + + def clean(): + calls.append(engine._write_lock.locked()) + return rollback() + + with patch.object(self.conn, "rollback", side_effect=clean): + with self.assertRaisesRegex(RuntimeError, "unfinished transaction"): + module.run_engine_maintenance(self.conn, object(), prepare_extensions=False, + _connection_lock=engine._write_lock) + self.assertTrue(calls) + self.assertTrue(all(calls)) + self.assertFalse(self.conn.in_transaction) + self.assertEqual(self.conn.execute("SELECT count(*) FROM messages WHERE id=992").fetchone()[0], 0) + self.conn.execute("INSERT INTO messages(id,content) VALUES(991,'synthetic caller')") + before = self.conn.commits, self.conn.rollbacks + module.run_engine_maintenance(self.conn, object(), prepare_extensions=False, + _connection_lock=engine._write_lock, allow_caller_transaction=True) + self.assertTrue(self.conn.in_transaction) + self.assertEqual((self.conn.commits, self.conn.rollbacks), before) + self.assertEqual(self.conn.execute("SELECT count(*) FROM messages WHERE id IN (991,992)").fetchone()[0], 2) + + + def test_opener_owns_canonical_path_before_writer_and_reenters_real_owner(self) -> None: + replacement = self.memory_connection() + engine = self.fixture.lifecycle_engine("_open_connection_handle") + engine.conn = None + engine.db_path = Path("~synthetic-unexpanded/store.sqlite") + self.enterContext(patch.object(Path, "resolve", _REAL_PATH_RESOLVE)) + expected = engine.db_path.resolve() + acquired, created = [], [] + + def acquire(path): + self.assertTrue(engine._init_lock.locked()) + self.assertFalse(engine._write_lock.locked()) + self.assertEqual(path, expected) + acquired.append(path) + return 17 + + owner = self.real_memory_owner(acquire) + self.modules["maintenance"] = owner + + def create(path): + self.assertEqual(path, expected) + self.assertTrue(engine._init_lock.locked()) + self.assertTrue(engine._write_lock.locked()) + self.assertIn(expected, owner._thread_owners.owners) + created.append(path) + return replacement + + engine._open_connection_handle.__globals__["create_db"] = create + with patch.object(Path, "mkdir") as mkdir, patch.object(Path, "chmod") as chmod, \ + patch.object(Path, "exists", return_value=False): + engine._open_connection_handle() + self.assertEqual(acquired, [expected]) + self.assertEqual(created, [expected]) + self.assertIs(engine.conn, replacement) + self.assertEqual(owner._thread_owners.owners, {}) + self.assertFalse(engine.ready) + self.assertFalse(engine._runtime_initialized) + self.assertFalse(engine._init_lock.locked()) + self.assertFalse(engine._write_lock.locked()) + mkdir.assert_called_once_with(parents=True, exist_ok=True) + chmod.assert_called_once_with(0o700) + + def test_opener_healthy_and_borrowed_handles_do_not_request_owner(self) -> None: + owner = Mock(side_effect=AssertionError("Healthy handle requested maintenance owner")) + self.modules["maintenance"].maintenance_owner = owner + engine = self.fixture.lifecycle_engine("_open_connection_handle") + engine._open_connection_handle() + self.conn.execute("INSERT INTO messages(id,content) VALUES(991,'synthetic caller')") + before = self.conn.commits, self.conn.rollbacks + engine._open_connection_handle() + error = sqlite3.OperationalError("disk I/O error") + 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) + owner.assert_not_called() + + def test_opener_owner_busy_preserves_ioerr_handle_and_ready_state(self) -> None: + owner = self.real_memory_owner(lambda path: None) + self.modules["maintenance"] = owner + engine = self.fixture.lifecycle_engine("_open_connection_handle") + engine.db_path = Path("synthetic/store.sqlite") + self.conn.armed, self.conn.fault_contains = True, "PRAGMA schema_version" + self.conn.fault = sqlite3.OperationalError("disk I/O error") + before = self.conn.commits, self.conn.rollbacks + with patch.object(Path, "mkdir"), patch.object(Path, "chmod"), \ + patch.object(self.conn, "close", wraps=self.conn.close) as close: + with self.assertRaises(owner.MaintenanceBusyError): + engine._open_connection_handle() + close.assert_not_called() + self.assertIs(engine.conn, self.conn) + self.assertTrue(engine.ready) + self.assertTrue(engine._runtime_initialized) + self.assertIs(engine._runtime_vector_connection, self.conn) + self.assertEqual((self.conn.commits, self.conn.rollbacks), before) + self.assertFalse(engine._write_lock.locked()) + self.assertFalse(engine._init_lock.locked()) + + def test_opener_second_probe_does_not_retire_a_new_healthy_handle(self) -> None: + replacement = self.memory_connection() + engine = self.fixture.lifecycle_engine("_open_connection_handle") + engine.db_path = Path("synthetic/store.sqlite") + self.conn.armed, self.conn.fault_contains = True, "PRAGMA schema_version" + self.conn.fault = sqlite3.OperationalError("disk I/O error") + + def acquire(path): + self.assertFalse(engine._write_lock.locked()) + engine.conn = replacement + return 17 + + self.modules["maintenance"] = self.real_memory_owner(acquire) + create = Mock(side_effect=AssertionError("Rechecked healthy connection was replaced")) + engine._open_connection_handle.__globals__["create_db"] = create + with patch.object(Path, "mkdir"), patch.object(Path, "chmod"), \ + patch.object(self.conn, "close", wraps=self.conn.close) as close: + engine._open_connection_handle() + close.assert_not_called() + self.assertIs(engine.conn, replacement) + create.assert_not_called() + + def test_concurrent_first_open_uses_one_factory_and_no_owner_contention(self) -> None: + replacement = self.memory_connection() + engine = self.fixture.lifecycle_engine("_open_connection_handle") + engine.conn, engine.db_path = None, Path("synthetic/store.sqlite") + acquired, created, failures = [], [], [] + factory_entered, release, second_waiting = (threading.Event() for _ in range(3)) + real_init = engine._init_lock + + class ObservedInit: + def acquire(self, blocking=True, timeout=-1): + if real_init.locked(): + second_waiting.set() + return real_init.acquire(blocking, timeout) + + def release(self): + real_init.release() + + def locked(self): + return real_init.locked() + + engine._init_lock = ObservedInit() + + def acquire(path): + self.assertFalse(engine._write_lock.locked()) + acquired.append(path) + return 17 + + self.modules["maintenance"] = self.real_memory_owner(acquire) + + def create(path): + created.append(path) + factory_entered.set() + if not release.wait(2): + raise AssertionError("Synthetic factory release timed out") + return replacement + + engine._open_connection_handle.__globals__["create_db"] = create + + def opening(): + try: + engine._open_connection_handle() + except BaseException as error: + failures.append(error) + + first, second = threading.Thread(target=opening), threading.Thread(target=opening) + with patch.object(Path, "mkdir"), patch.object(Path, "chmod"), patch.object(Path, "exists", return_value=False): + first.start() + try: + self.assertTrue(factory_entered.wait(2)) + second.start() + self.assertTrue(second_waiting.wait(2)) + self.assertEqual(len(acquired), 1) + self.assertEqual(len(created), 1) + finally: + release.set() + first.join(2) + if second.ident is not None: + second.join(2) + self.assertFalse(first.is_alive()) + self.assertFalse(second.is_alive()) + self.assertEqual(failures, []) + self.assertEqual(len(acquired), 1) + self.assertEqual(len(created), 1) + self.assertIs(engine.conn, replacement) + + + def test_consolidate_guards_authority_and_style_table_snapshots(self) -> None: + self.fixture.select() + engine = self.fixture.lifecycle_engine("_open_connection_handle") + engine._open_connection_handle = lambda: None + engine._ensure_connection = Mock(side_effect=AssertionError("Controlled store initialized models")) + engine._auto_consolidate_threshold = 25 + engine._get_maintenance_coordinator = lambda: types.SimpleNamespace(path=None, refresh_capabilities=lambda: None) + engine._apply_manual_maintenance_capabilities = lambda results: None + report = types.SimpleNamespace(results=()) + + def run(conn, coordinator, **kwargs): + self.assertFalse(engine._write_lock.locked()) + self.assertIs(kwargs["_connection_lock"], engine._write_lock) + return report + + self.modules["maintenance"] = types.SimpleNamespace( + MaintenanceBusyError=type("Busy", (RuntimeError,), {}), format_maintenance_report=lambda value: {}, + maintenance_busy_result=lambda: {}, maintenance_observation=lambda *args, **kwargs: contextlib.nullcontext([]), + observe_style=Mock(side_effect=AssertionError("Unexpected file style probe")), + run_engine_maintenance=run, _style_unready_result=lambda value: value) + source = definitions("engine.py", {"consolidate"}, self.namespace(), methods=True) + execute = self.conn.execute + observed = [] + + def guarded(sql, *args, **kwargs): + self.assertTrue(engine._write_lock.locked()) + observed.append(sql) + return execute(sql, *args, **kwargs) + + with patch.object(self.conn, "execute", side_effect=guarded): + self.assertEqual(source.consolidate(engine), {}) + self.assertTrue(any("tier_selected_v1" in str(args) for args in observed) + or any("main.metadata" in sql for sql in observed)) + self.assertTrue(any("entity_style_vectors" in sql for sql in observed)) + engine._ensure_connection.assert_not_called() + + def test_consolidate_release_handoff_preserves_next_writer_with_negative_control(self) -> None: + for restore_old_cleanup in (False, True): + with self.subTest(restore_old_cleanup=restore_old_cleanup): + conn = self.memory_connection() + engine = self.fixture.lifecycle_engine("_open_connection_handle", conn) + engine._open_connection_handle = lambda: None + engine._ensure_connection = lambda **kwargs: None + engine._auto_consolidate_threshold = 25 + engine._has_style_vec = False + coordinator = types.SimpleNamespace(path=None, refresh_capabilities=lambda: None) + engine._get_maintenance_coordinator = lambda: coordinator + engine._apply_manual_maintenance_capabilities = lambda results: None + armed, released, writer_started, writer_release = (threading.Event() for _ in range(4)) + failures = [] + main_thread = threading.get_ident() + api = self.api + owner = self + + class HandoffLock(api.ConnectionWriteLock): + def release(self): + handoff = armed.is_set() and threading.get_ident() == main_thread + if handoff: + armed.clear() + super().release() + if handoff: + released.set() + owner.assertTrue(writer_started.wait(2), "Synthetic next writer did not start") + + engine._write_lock = lock = HandoffLock(engine) + + def layers(*args, **kwargs): + self.assertTrue(lock.locked()) + armed.set() + return () + + runner = self.maintenance_runner(run_layers=layers) + self.modules["maintenance"] = types.SimpleNamespace( + MaintenanceBusyError=type("Busy", (RuntimeError,), {}), + format_maintenance_report=lambda report: {"result": "synthetic complete"}, maintenance_busy_result=lambda: {}, + maintenance_observation=lambda *args, **kwargs: contextlib.nullcontext([]), + observe_style=Mock(side_effect=AssertionError("Unexpected style probe")), + run_engine_maintenance=runner.run_engine_maintenance, _style_unready_result=lambda value: value) + source = definitions("engine.py", {"consolidate"}, self.namespace(), methods=True) + if restore_old_cleanup: + tree = ast.parse((ROOT / "truememory/engine.py").read_text(encoding="utf-8")) + cls = next(node for node in tree.body if isinstance(node, ast.ClassDef) and node.name == "TrueMemoryEngine") + method = next(node for node in cls.body if isinstance(node, ast.FunctionDef) and node.name == "consolidate") + + class RestoreCleanup(ast.NodeTransformer): + inserted = 0 + + def visit_Assign(self, node): + if (isinstance(node.value, ast.Call) and isinstance(node.value.func, ast.Name) + and node.value.func.id == "run_engine_maintenance" + and any(item.arg == "_connection_lock" for item in node.value.keywords)): + self.inserted += 1 + old = ast.parse("if not borrowed and self.conn.in_transaction:\n" + " self.conn.rollback()\n" + " raise RuntimeError('Maintenance work left an unfinished transaction')").body[0] + return [node, old] + return node + + restore = RestoreCleanup() + method = restore.visit(method) + self.assertEqual(restore.inserted, 1) + exec(compile(ast.fix_missing_locations(ast.Module(body=[method], type_ignores=[])), + "synthetic-old-cleanup-negative-control", "exec"), source.__dict__) + + def next_writer(): + try: + if not released.wait(2): + raise AssertionError("Maintenance mutation lock was not released") + with lock: + conn.execute("INSERT INTO messages(id,content) VALUES(993,'synthetic next writer')") + writer_started.set() + if not writer_release.wait(2): + raise AssertionError("Synthetic next writer was not released") + conn.rollback() + except BaseException as error: + failures.append(error) + + thread = threading.Thread(target=next_writer) + thread.start() + try: + if restore_old_cleanup: + with self.assertRaisesRegex(RuntimeError, "unfinished transaction"): + source.consolidate(engine) + self.assertFalse(conn.in_transaction) + self.assertEqual(conn.execute("SELECT count(*) FROM messages WHERE id=993").fetchone()[0], 0) + else: + self.assertEqual(source.consolidate(engine), {"result": "synthetic complete"}) + self.assertTrue(conn.in_transaction) + self.assertEqual(conn.execute("SELECT content FROM messages WHERE id=993").fetchone()[0], + "synthetic next writer") + finally: + writer_release.set() + thread.join(2) + self.assertFalse(thread.is_alive()) + self.assertEqual(failures, []) + self.assertTrue(writer_started.is_set()) + + + def test_already_cancelled_runner_does_not_admit_or_wait_for_managed_writer(self) -> None: + conn = self.memory_connection() + engine = self.fixture.lifecycle_engine("_open_connection_handle", conn) + owned, release = threading.Event(), threading.Event() + failures = [] + cancel = threading.Event() + cancel.set() + module = self.maintenance_runner(owner=Mock(side_effect=AssertionError("Cancelled work requested owner"))) + module.connection_database_path = Mock(side_effect=AssertionError("Cancelled work read database identity")) + module.run_routed_style = Mock(side_effect=AssertionError("Cancelled work called mutable style route")) + + def writer() -> None: + try: + with engine._write_lock: + conn.execute("INSERT INTO messages(id,content) VALUES(994,'synthetic managed writer')") + owned.set() + if not release.wait(2): + raise AssertionError("Synthetic writer was not released") + conn.rollback() + except BaseException as error: + failures.append(error) + + thread = threading.Thread(target=writer) + thread.start() + try: + self.assertTrue(owned.wait(2)) + before = conn.commits, conn.rollbacks + with patch.object(self.api, "_connection_read", side_effect=AssertionError("Cancelled work entered lock admission")), \ + patch.object(conn, "execute", side_effect=AssertionError("Cancelled work executed SQL")): + for include_style in (False, True): + with self.subTest(include_style=include_style): + report = module.run_engine_maintenance(conn, object(), cancel=cancel, + include_style=include_style, _connection_lock=engine._write_lock) + self.assertEqual(report.results, ()) + self.assertEqual(report.preferences, "CANCELLED") + if include_style: + self.assertEqual(tuple(report.style_result), ( + "style_vectors", "style_vectors", "deferred", None, 0.0, + "StyleCancelled", False, "unverified", True)) + else: + self.assertIsNone(report.style_result) + self.assertTrue(engine._write_lock.owns_other_transaction(conn)) + self.assertTrue(conn.in_transaction) + self.assertEqual((conn.commits, conn.rollbacks), before) + self.assertEqual(conn.execute("SELECT content FROM messages WHERE id=994").fetchone()[0], + "synthetic managed writer") + module.maintenance_owner.assert_not_called() + module.connection_database_path.assert_not_called() + module.run_routed_style.assert_not_called() + finally: + release.set() + thread.join(2) + self.assertFalse(thread.is_alive()) + self.assertEqual(failures, []) + + def test_cancelled_snapshot_stays_terminal_when_event_clears(self) -> None: + module = self.maintenance_runner(owner=Mock(side_effect=AssertionError("Cancelled snapshot requested owner"))) + module.connection_database_path = Mock(side_effect=AssertionError("Cancelled snapshot read database identity")) + module.run_routed_style = Mock(side_effect=AssertionError("Cleared event reached style SQL")) + for borrowed in (False, True): + if borrowed: + self.conn.execute("INSERT INTO messages(id,content) VALUES(994,'synthetic caller SQL')") + before = self.conn.commits, self.conn.rollbacks + for include_style in (False, True): + with self.subTest(borrowed=borrowed, include_style=include_style): + cancel = threading.Event() + cancel.set() + + def capture_and_clear() -> bool: + cancel.clear() + return True + + with patch.object(cancel, "is_set", side_effect=capture_and_clear) as checked, \ + patch.object(self.api, "_connection_read", side_effect=AssertionError("Cancelled snapshot waited")), \ + patch.object(self.conn, "execute", side_effect=AssertionError("Cancelled snapshot executed SQL")): + report = module.run_engine_maintenance(self.conn, object(), cancel=cancel, + include_style=include_style, allow_caller_transaction=borrowed) + checked.assert_called_once_with() + self.assertFalse(cancel.is_set()) + self.assertEqual(report.results, ()) + self.assertEqual(report.preferences, "CANCELLED") + if include_style: + self.assertEqual(tuple(report.style_result), ( + "style_vectors", "style_vectors", "deferred", None, 0.0, + "StyleCancelled", False, "unverified", borrowed)) + else: + self.assertIsNone(report.style_result) + self.assertEqual(self.conn.in_transaction, borrowed) + self.assertEqual((self.conn.commits, self.conn.rollbacks), before) + if borrowed: + self.assertEqual(self.conn.execute("SELECT content FROM messages WHERE id=994").fetchone()[0], + "synthetic caller SQL") + self.conn.rollback() + module.maintenance_owner.assert_not_called() + module.connection_database_path.assert_not_called() + module.run_routed_style.assert_not_called() + + +class TestManagedModelAdmission(unittest.TestCase): + """Actual model guard classifies the owner of SQL on a shared handle.""" + + setUp = TestPublicBoundaries.setUp + memory_connection = TestPublicBoundaries.memory_connection + + def writer_lock(self, conn): + class Engine: + pass + engine = Engine() + engine.conn = conn + lock = self.api.ConnectionWriteLock(engine) + return engine, lock + + def test_model_guard_reprobes_owner_release_without_using_stale_sql_flag(self): + # Negative control is the exact pre-fix ordering: read SQL first, then + # ask whether the competing managed writer still owns the handle. + def previous_guard(operation, lock): + if operation.borrowed or operation._connection.in_transaction and not ( + isinstance(lock, self.api.ConnectionWriteLock) + and lock.owns_other_transaction(operation._connection)): + raise self.api.TierRuntimeError("Model loading requires a clean connection") + for previous in (True, False): + with self.subTest(previous_order=previous): + conn = self.memory_connection() + engine, lock = self.writer_lock(conn) + entered, release, finished = threading.Event(), threading.Event(), threading.Event() + errors, observations = [], [] + def writer(): + try: + with lock: + conn.execute("BEGIN IMMEDIATE") + entered.set() + if not release.wait(3): + raise AssertionError("Synthetic release race timed out") + conn.commit() + except BaseException as error: + errors.append(error) + entered.set() + finally: + finished.set() + actual_owner_check = lock.owns_other_transaction + def owner_check(connection): + observations.append(connection.in_transaction) + release.set() + if not finished.wait(2): + raise AssertionError("Synthetic owner did not finish") + return actual_owner_check(connection) + with self.api.serving_operation(conn, connection_lock=lock) as operation: + thread = threading.Thread(target=writer) + thread.start() + try: + self.assertTrue(entered.wait(1)) + self.assertEqual(errors, []) + with patch.object(lock, "owns_other_transaction", side_effect=owner_check), \ + patch.object(lock, "acquire", wraps=lock.acquire) as acquire: + if previous: + with self.assertRaisesRegex(self.api.TierRuntimeError, "clean connection"): + previous_guard(operation, lock) + acquire.assert_not_called() + else: + self.api.require_model_load_allowed() + acquire.assert_called_once_with(blocking=False) + self.assertEqual(observations, [True]) + self.assertFalse(conn.in_transaction) + self.assertFalse(lock.locked()) + self.assertIs(engine.conn, conn) + finally: + release.set() + thread.join(3) + self.assertFalse(thread.is_alive()) + self.assertEqual(errors, []) + + def test_model_guard_current_thread_clean_lock_is_not_reacquired(self): + conn = self.memory_connection() + engine, lock = self.writer_lock(conn) + with self.api.serving_operation(conn, connection_lock=lock): + with lock: + with patch.object(lock, "acquire", side_effect=AssertionError("Nonreentrant writer reacquired")): + self.api.require_model_load_allowed() + self.assertTrue(lock.owned_by_current_thread()) + self.assertFalse(conn.in_transaction) + self.assertFalse(lock.locked()) + self.assertIs(engine.conn, conn) + + def test_model_guard_allows_other_managed_writer_without_acquiring_its_lock(self): + conn = self.memory_connection() + engine, lock = self.writer_lock(conn) + entered, release = threading.Event(), threading.Event() + errors = [] + def writer(): + try: + with lock: + conn.execute("BEGIN IMMEDIATE") + conn.execute("INSERT INTO messages(id,content) VALUES(995,'synthetic writer')") + entered.set() + if not release.wait(3): + raise AssertionError("Synthetic writer release timed out") + conn.rollback() + except BaseException as error: + errors.append(error) + entered.set() + with self.api.serving_operation(conn, connection_lock=lock) as operation: + self.assertFalse(operation.borrowed) + self.assertIs(operation._connection_lock, lock) + thread = threading.Thread(target=writer) + thread.start() + try: + self.assertTrue(entered.wait(1)) + self.assertEqual(errors, []) + self.assertTrue(lock.owns_other_transaction(conn)) + with patch.object(lock, "acquire", side_effect=AssertionError("Model guard waited on writer")): + self.api.require_model_load_allowed() + self.assertTrue(conn.in_transaction) + self.assertEqual(conn.execute("SELECT content FROM messages WHERE id=995").fetchone()[0], "synthetic writer") + self.assertIs(engine.conn, conn) + finally: + release.set() + thread.join(3) + self.assertFalse(thread.is_alive()) + self.assertEqual(errors, []) + + def test_model_guard_preserves_initial_borrowed_refusal_after_caller_commit(self): + conn = self.memory_connection() + engine, lock = self.writer_lock(conn) + conn.execute("BEGIN") + with self.api.serving_operation(conn, connection_lock=lock) as operation: + self.assertTrue(operation.borrowed) + with self.assertRaisesRegex(self.api.TierRuntimeError, "clean connection"): + self.api.require_model_load_allowed() + self.assertTrue(conn.in_transaction) + conn.commit() + with self.assertRaisesRegex(self.api.TierRuntimeError, "clean connection"): + self.api.require_model_load_allowed() + self.assertIs(engine.conn, conn) + + def test_model_guard_refuses_current_thread_owned_sql_without_ending_it(self): + conn = self.memory_connection() + engine, lock = self.writer_lock(conn) + with self.api.serving_operation(conn, connection_lock=lock): + with lock: + conn.execute("BEGIN") + try: + with self.assertRaisesRegex(self.api.TierRuntimeError, "clean connection"): + self.api.require_model_load_allowed() + self.assertTrue(conn.in_transaction) + finally: + conn.rollback() + self.assertIs(engine.conn, conn) + + def test_model_guard_refuses_unmanaged_or_duck_typed_transaction_ownership(self): + class UnmanagedLock: + def owns_other_transaction(self, conn): + raise AssertionError("Foreign ownership callback was consulted") + conn = self.memory_connection() + for lock in (None, threading.Lock(), UnmanagedLock()): + with self.subTest(lock=type(lock).__name__), self.api.serving_operation(conn) as operation: + replacement = dataclasses.replace(operation, _connection_lock=lock) + with self.api._bound(replacement): + conn.execute("BEGIN") + try: + with self.assertRaisesRegex(self.api.TierRuntimeError, "clean connection"): + self.api.require_model_load_allowed() + self.assertTrue(conn.in_transaction) + finally: + conn.rollback() + + def test_model_guard_does_not_use_managed_owner_for_another_handle(self): + conn, other = self.memory_connection(), self.memory_connection() + engine, lock = self.writer_lock(other) + entered, release = threading.Event(), threading.Event() + errors = [] + def writer(): + try: + with lock: + other.execute("BEGIN") + entered.set() + if not release.wait(3): + raise AssertionError("Synthetic stale-handle writer release timed out") + other.rollback() + except BaseException as error: + errors.append(error) + entered.set() + with self.api.serving_operation(conn, connection_lock=lock): + thread = threading.Thread(target=writer) + thread.start() + try: + self.assertTrue(entered.wait(1)) + self.assertEqual(errors, []) + self.assertTrue(lock.owns_other_transaction(other)) + self.assertFalse(lock.owns_other_transaction(conn)) + conn.execute("BEGIN") + with self.assertRaisesRegex(self.api.TierRuntimeError, "clean connection"): + self.api.require_model_load_allowed() + self.assertTrue(conn.in_transaction) + conn.rollback() + finally: + release.set() + thread.join(3) + self.assertFalse(thread.is_alive()) + self.assertEqual(errors, []) + self.assertIs(engine.conn, other) + + def test_model_guard_closed_handle_error_is_not_reclassified_or_suppressed(self): + conn = self.memory_connection() + with self.api.serving_operation(conn): + conn.close() + with self.assertRaisesRegex(sqlite3.ProgrammingError, "closed database"): + self.api.require_model_load_allowed() + + def test_runtime_child_carries_lock_only_for_original_handle(self): + conn, other = self.memory_connection(), self.memory_connection() + engine, lock = self.writer_lock(conn) + joined, errors = [], [] + with patch.object(self.api, "_database_identity", return_value=("file", "synthetic-database")): + with self.api.serving_operation(conn, connection_lock=lock) as operation: + same_child, other_child = operation.fork_child(), operation.fork_child() + def join(child, connection, expected_lock): + try: + with child.join(connection) as current: + self.assertIs(current._connection_lock, expected_lock) + self.assertIs(current._connection, connection) + connection.execute("BEGIN") + try: + with self.assertRaisesRegex(self.api.TierRuntimeError, "clean connection"): + self.api.require_model_load_allowed() + finally: + connection.rollback() + joined.append(connection) + except BaseException as error: + errors.append(error) + # Sequential children isolate their SQL while retaining the + # production thread-owned child admission protocol. + for child, connection, expected in ((same_child, conn, lock), (other_child, other, None)): + thread = threading.Thread(target=join, args=(child, connection, expected)) + thread.start() + thread.join(3) + self.assertFalse(thread.is_alive()) + self.assertEqual(errors, []) + self.assertEqual(joined, [conn, other]) + self.assertIs(engine.conn, conn) + + +class TestEngineReconnectReceipt(unittest.TestCase): + setUp = TestPublicBoundaries.setUp + imports = TestPublicBoundaries.imports + namespace = TestPublicBoundaries.namespace + memory_connection = TestPublicBoundaries.memory_connection + real_memory_owner = TestPublicBoundaries.real_memory_owner + + def engine(self, *, vectors=False): + self.enterContext(patch.object(Path, "resolve", _REAL_PATH_RESOLVE)) + self.stat = types.SimpleNamespace(st_dev=7, st_ino=11) + self.enterContext(patch.object(Path, "stat", return_value=self.stat)) + self.enterContext(patch.object(Path, "mkdir")) + self.enterContext(patch.object(Path, "chmod")) + self.modules["storage"] = definitions("storage.py", {"DatabaseOpenError", "_integrity_message", "_validate_db_path"}, + self.namespace(sqlite3=sqlite3, Path=Path, os=os, newest_backup=lambda path: None)) + namespace = self.namespace(sqlite3=sqlite3, Path=Path, logger=logging.getLogger("synthetic-reconnect"), + _HAS_VECTOR=vectors, DEFAULT_BUSY_TIMEOUT_MS=3000, + _ALL_VEC_TABLES=("vec_messages", "vec_messages_sep", "vec_messages_edge", "vec_messages_sep_edge", + "vec_messages_basepro", "vec_messages_sep_basepro"), + create_db=Mock(side_effect=AssertionError("Unexpected schema factory"))) + tree = ast.parse((ROOT / "truememory/engine.py").read_text(encoding="utf-8")) + cls = next(n for n in tree.body if isinstance(n, ast.ClassDef) and n.name == "TrueMemoryEngine") + names = {"_open_connection_handle", "_reconnect_registry", "_capture_reconnect_receipt", + "_reopen_initialized_connection", "close"} + cls.body = [n for n in cls.body if isinstance(n, ast.FunctionDef) and n.name in names] + future = ast.ImportFrom(module="__future__", names=[ast.alias(name="annotations")], level=0) + exec(compile(ast.fix_missing_locations(ast.Module(body=[future, cls], type_ignores=[])), + "actual-engine-reconnect", "exec"), namespace) + engine = namespace["TrueMemoryEngine"]() + engine.conn, engine.db_path = self.conn, Path("/synthetic/reconnect.sqlite") + engine._init_lock = threading.Lock() + engine._write_lock = self.api.ConnectionWriteLock(engine) + engine.ready = engine._runtime_initialized = True + engine._has_vectors = vectors + engine._runtime_vector_connection = self.conn + engine._runtime_vector_generation = None + engine._reconnect_receipt = None + self.fixture.vector._active_tier_group = lambda: "edge" + self.conn.execute("INSERT OR REPLACE INTO metadata(key,value) VALUES('l4_entity_profile_migration_done','1')") + self.conn.execute("DELETE FROM metadata WHERE key IN ('embed_model','embed_dim')") + if vectors: + self.conn.execute("CREATE TABLE IF NOT EXISTS vec_messages(rowid INTEGER PRIMARY KEY)") + self.conn.execute("CREATE TABLE IF NOT EXISTS vec_messages_sep(rowid INTEGER PRIMARY KEY)") + self.conn.commit() + operation = types.SimpleNamespace(key=self.api._legacy_key(), _database=("file", str(engine.db_path))) + with engine._write_lock: + engine._capture_reconnect_receipt(operation) + self.assertIsNotNone(engine._reconnect_receipt) + self.assertFalse(self.conn.in_transaction) + self.assertEqual(len(engine._reconnect_receipt), 10) + self.engine_namespace = namespace + return engine + + def candidate(self, engine): + candidate = self.memory_connection() + # SQLite backup assigns a new schema cookie; simulate reopening the same file. + schema = self.conn.execute("PRAGMA schema_version").fetchone()[0] + candidate.execute(f"PRAGMA schema_version={schema}") + connect = Mock(return_value=candidate) + proxy = types.SimpleNamespace(**{name: getattr(sqlite3, name) for name in + ("Connection", "ProgrammingError", "OperationalError", "DatabaseError", "Error")}, connect=connect) + self.engine_namespace["sqlite3"] = proxy + self.connect = connect + return candidate + + def forbid_owner(self): + owner = Mock(side_effect=AssertionError("Validated reconnect requested maintenance ownership")) + self.modules["maintenance"].maintenance_owner = owner + return owner + + def assert_closed(self, conn): + with self.assertRaises(sqlite3.ProgrammingError): + conn.execute("SELECT 1") + + def test_closed_initialized_fts_handle_reopens_without_owner_or_schema_writes(self): + engine = self.engine() + self.conn.execute("INSERT INTO messages(id,content) VALUES(881,'synthetic prior source')") + self.conn.commit() + candidate = self.candidate(engine) + queries = [] + candidate.set_trace_callback(queries.append) + owner = self.forbid_owner() + self.conn.close() + engine._open_connection_handle() + self.assertIs(engine.conn, candidate) + self.assertTrue(engine.ready) + self.assertTrue(engine._runtime_initialized) + self.assertFalse(engine._has_vectors) + self.assertFalse(candidate.in_transaction) + self.assertEqual(candidate.execute("SELECT content FROM messages WHERE id=881").fetchone()[0], "synthetic prior source") + self.connect.assert_called_once_with("file:///synthetic/reconnect.sqlite?mode=rw", uri=True, check_same_thread=False) + self.assertIn("PRAGMA quick_check(1)", queries) + self.assertIn("PRAGMA foreign_keys=ON", queries) + self.assertFalse(any(q.lstrip().split(" ", 1)[0].upper() in {"INSERT", "UPDATE", "DELETE", "CREATE", "ALTER", "DROP"} for q in queries)) + owner.assert_not_called() + + def test_persistent_clean_ioerr_retires_old_handle_only_after_validation(self): + engine = self.engine() + candidate = self.candidate(engine) + self.forbid_owner() + self.conn.armed, self.conn.fault_contains = True, "PRAGMA schema_version" + self.conn.fault = sqlite3.OperationalError("disk I/O error") + with patch.object(self.conn, "close", wraps=self.conn.close) as close: + engine._open_connection_handle() + close.assert_called_once_with() + self.assertIs(engine.conn, candidate) + self.assertTrue(engine.ready) + + def test_receipt_does_not_bypass_borrowed_ioerr_or_unrelated_error(self): + engine = self.engine() + self.candidate(engine) + self.forbid_owner() + for borrowed in (False, True): + with self.subTest(borrowed=borrowed): + if borrowed: + self.conn.execute("INSERT INTO messages(id,content) VALUES(882,'synthetic pending')") + error = sqlite3.OperationalError("disk I/O error" if borrowed else "database is locked") + self.conn.armed, self.conn.fault_contains, self.conn.fault = True, "PRAGMA schema_version", error + with self.assertRaises(sqlite3.OperationalError) as raised: + engine._open_connection_handle() + self.assertIs(raised.exception, error) + self.assertEqual(self.conn.in_transaction, borrowed) + self.conn.armed = False + self.connect.assert_not_called() + self.assertEqual(self.conn.execute("SELECT content FROM messages WHERE id=882").fetchone()[0], "synthetic pending") + + def test_new_engine_without_receipt_retains_owner_required_refusal(self): + engine = self.engine() + engine.conn = None + engine._reconnect_receipt = None + owner = self.real_memory_owner(lambda path: None) + self.modules["maintenance"] = owner + self.candidate(engine) + with self.assertRaises(owner.MaintenanceBusyError): + engine._open_connection_handle() + self.assertIsNone(engine.conn) + self.connect.assert_not_called() + + def test_explicit_close_clears_receipt_and_next_open_requires_owner(self): + engine = self.engine() + engine.close() + self.assertIsNone(engine._reconnect_receipt) + owner = self.real_memory_owner(lambda path: None) + self.modules["maintenance"] = owner + with self.assertRaises(owner.MaintenanceBusyError): + engine._open_connection_handle() + self.assertIsNone(engine.conn) + + def test_healthy_ready_handle_has_no_receipt_sql_or_file_stat(self): + engine = self.engine() + self.candidate(engine) + self.forbid_owner() + queries = [] + self.conn.set_trace_callback(queries.append) + with patch.object(Path, "stat", side_effect=AssertionError("Ready path stat")): + engine._open_connection_handle() + self.assertEqual(queries, ["PRAGMA schema_version"]) + self.connect.assert_not_called() + + def test_file_schema_registry_or_embedding_change_refuses_receipt(self): + for change in ("inode", "schema", "registry", "model", "dimension", "partial", "oversized"): + with self.subTest(change=change): + engine = self.engine() + candidate = self.candidate(engine) + if change == "inode": + self.stat.st_ino += 1 + elif change == "schema": + candidate.execute("CREATE TABLE synthetic_new_schema(x)") + elif change == "registry": + candidate.execute("INSERT INTO vector_cache_registry(tier_group,vec_table,sep_table) VALUES('edge','vec_messages_edge','vec_messages_sep_edge')") + else: + candidate.execute("INSERT OR REPLACE INTO metadata(key,value) VALUES('embed_model',?)", + ("different" if change == "model" else "x" * 1024 if change == "oversized" else "model2vec",)) + if change != "partial": + candidate.execute("INSERT OR REPLACE INTO metadata(key,value) VALUES('embed_dim',?)", ("384" if change == "dimension" else "256",)) + candidate.commit() + self.assertIsNone(engine._reopen_initialized_connection(engine.db_path)) + if change == "inode": + self.connect.assert_not_called() + else: + self.assert_closed(candidate) + self.assertIs(engine.conn, self.conn) + self.assertTrue(engine.ready) + + def test_matching_metadata_published_after_initialization_is_accepted(self): + engine = self.engine() + candidate = self.candidate(engine) + candidate.executemany("INSERT INTO metadata(key,value) VALUES(?,?)", [("embed_model", "model2vec"), ("embed_dim", "256")]) + candidate.commit() + replacement = engine._reopen_initialized_connection(engine.db_path) + self.assertEqual(replacement, (candidate, False)) + self.assertFalse(candidate.in_transaction) + + def test_quick_check_failure_closes_candidate_preserves_old_handle_and_ready(self): + for failure in ("disk I/O error", "database disk image is malformed"): + with self.subTest(failure=failure): + engine = self.engine() + candidate = self.candidate(engine) + candidate.armed, candidate.fault_contains = True, "PRAGMA quick_check(1)" + candidate.fault = sqlite3.OperationalError(failure) + self.conn.armed, self.conn.fault_contains = True, "PRAGMA schema_version" + self.conn.fault = sqlite3.OperationalError("disk I/O error") + with patch.object(self.conn, "close", wraps=self.conn.close) as close: + with self.assertRaises(self.modules["storage"].DatabaseOpenError) as raised: + engine._open_connection_handle() + close.assert_not_called() + self.assertIn("Close ALL" if "I/O" in failure else "appears corrupt", str(raised.exception)) + self.assert_closed(candidate) + self.assertIs(engine.conn, self.conn) + self.assertTrue(engine.ready) + self.conn.armed = False + + def test_recheck_uses_healthy_replacement_and_closes_unused_candidate(self): + engine = self.engine() + candidate = self.candidate(engine) + healthy = self.memory_connection() + self.forbid_owner() + self.conn.armed, self.conn.fault_contains = True, "PRAGMA schema_version" + self.conn.fault = sqlite3.OperationalError("disk I/O error") + def connect(*args, **kwargs): + engine.conn = healthy + return candidate + self.connect.side_effect = connect + with patch.object(self.conn, "close", wraps=self.conn.close) as close: + engine._open_connection_handle() + close.assert_not_called() + self.assertIs(engine.conn, healthy) + self.assert_closed(candidate) + + def test_changed_file_after_candidate_validation_is_refused(self): + engine = self.engine() + candidate = self.candidate(engine) + self.enterContext(patch.object(Path, "stat", side_effect=[self.stat, types.SimpleNamespace(st_dev=7, st_ino=12)])) + self.assertIsNone(engine._reopen_initialized_connection(engine.db_path)) + self.assert_closed(candidate) + + def test_vector_metadata_absence_requires_both_tables_empty_and_extension_registration(self): + for table in (None, "vec_messages", "vec_messages_sep"): + with self.subTest(table=table): + engine = self.engine(vectors=True) + candidate = self.candidate(engine) + if table: + candidate.execute(f'INSERT INTO "{table}"(rowid) VALUES(1)') + candidate.commit() + load = Mock() + original_import = self.imports + def imports(name, *args, **kwargs): + if name == "sqlite_vec": + return types.SimpleNamespace(load=load) + return original_import(name, *args, **kwargs) + self.engine_namespace["__builtins__"]["__import__"] = imports + with patch.object(candidate, "enable_load_extension", create=True) as enabled: + result = engine._reopen_initialized_connection(engine.db_path) + load.assert_called_once_with(candidate) + self.assertEqual([c.args for c in enabled.call_args_list], [(True,), (False,)]) + if table is None: + self.assertEqual(result, (candidate, False)) + else: + self.assertIsNone(result) + self.assert_closed(candidate) + + + + def test_policy_supersedes_legacy_ready_flags_without_legacy_metadata_reuse(self): + engine = self.engine() + policy = TestPublicBoundaries.seed_policy(self) + candidate = self.candidate(engine) + self.forbid_owner() + self.conn.close() + engine._open_connection_handle() + self.assertIs(engine.conn, candidate) + self.assertFalse(engine.ready) + self.assertFalse(engine._runtime_initialized) + self.assertFalse(engine._has_vectors) + self.assertIsNone(engine._reconnect_receipt) + self.assertIsNone(engine._runtime_vector_connection) + self.assertEqual(self.api._read_policy(candidate), policy) + self.assertEqual(self.fixture.loads, []) + + def test_registry_progress_mutation_invalidates_same_schema_receipt(self): + engine = self.engine() + self.conn.execute("INSERT INTO vector_cache_registry(tier_group,vec_table,sep_table) VALUES('edge','vec_messages_edge','vec_messages_sep_edge')") + self.conn.commit() + operation = types.SimpleNamespace(key=self.api._legacy_key(), _database=("file", str(engine.db_path))) + with engine._write_lock: + engine._capture_reconnect_receipt(operation) + candidate = self.candidate(engine) + candidate.execute("UPDATE vector_cache_registry SET vector_count=1 WHERE tier_group='edge'") + candidate.commit() + self.assertEqual(candidate.execute("PRAGMA schema_version").fetchone()[0], engine._reconnect_receipt[3]) + self.assertIsNone(engine._reopen_initialized_connection(engine.db_path)) + self.assert_closed(candidate) + + def test_registry_record_bounds_and_pending_caller_sql_prevent_capture(self): + engine = self.engine() + operation = types.SimpleNamespace(key=self.api._legacy_key(), _database=("file", str(engine.db_path))) + self.conn.execute("INSERT INTO vector_cache_registry(tier_group,vec_table,sep_table,model_name) VALUES('edge','vec_messages_edge','vec_messages_sep_edge',?)", ("x" * 513,)) + self.conn.commit() + with engine._write_lock: + engine._capture_reconnect_receipt(operation) + self.assertIsNone(engine._reconnect_receipt) + self.conn.execute("INSERT INTO messages(id,content) VALUES(883,'synthetic caller')") + before = self.conn.commits, self.conn.rollbacks + with engine._write_lock: + engine._capture_reconnect_receipt(operation) + self.assertIsNone(engine._reconnect_receipt) + self.assertTrue(self.conn.in_transaction) + self.assertEqual((self.conn.commits, self.conn.rollbacks), before) + + def test_bad_integrity_result_and_late_stat_failure_close_only_candidate(self): + for fault in ("quick_check", "stat"): + with self.subTest(fault=fault): + engine = self.engine() + candidate = self.candidate(engine) + execute = candidate.execute + def checked(sql, *args, **kwargs): + if sql == "PRAGMA quick_check(1)": + return types.SimpleNamespace(fetchone=lambda: ("synthetic corruption",)) + return execute(sql, *args, **kwargs) + with contextlib.ExitStack() as stack: + if fault == "quick_check": + stack.enter_context(patch.object(candidate, "execute", side_effect=checked)) + expected = self.modules["storage"].DatabaseOpenError + else: + stack.enter_context(patch.object(Path, "stat", side_effect=[self.stat, OSError("synthetic unavailable file")])) + expected = OSError + with self.assertRaises(expected): + engine._reopen_initialized_connection(engine.db_path) + self.assert_closed(candidate) + self.assertIs(engine.conn, self.conn) + self.assertTrue(engine.ready) + self.assertFalse(self.conn.in_transaction) + + def test_deprecated_open_legacy_callback_is_existing_only_without_migrations(self): + engine = self.engine() + candidate = self.candidate(engine) + queries = [] + candidate.set_trace_callback(queries.append) + module = definitions("engine.py", {"open"}, self.namespace(Path=Path, + sqlite3=self.engine_namespace["sqlite3"], DEFAULT_BUSY_TIMEOUT_MS=3000, + warnings=types.SimpleNamespace(warn=lambda *args, **kwargs: None)), methods=True) + class Completed(Exception): + pass + def opener(path, callback): + self.assertIs(callback(engine.db_path), candidate) + raise Completed() + with patch.object(self.api, "open_serving_connection", side_effect=opener), patch.object(Path, "exists", return_value=True): + with self.assertRaises(Completed): + module.open(engine) + self.connect.assert_called_once_with("file:///synthetic/reconnect.sqlite?mode=rw", uri=True, check_same_thread=False) + self.assertEqual(queries, ["PRAGMA journal_mode=WAL", "PRAGMA busy_timeout=3000", "PRAGMA foreign_keys=ON", + "PRAGMA synchronous=NORMAL", "PRAGMA cache_size=-64000", "PRAGMA mmap_size=268435456"]) + self.assertIs(engine.conn, self.conn) + self.assertIsNone(candidate.row_factory) + + def test_deprecated_open_legacy_callback_failure_closes_only_candidate(self): + engine = self.engine() + candidate = self.candidate(engine) + candidate.armed, candidate.fault_contains = True, "PRAGMA foreign_keys=ON" + candidate.fault = sqlite3.OperationalError("synthetic pragma failure") + module = definitions("engine.py", {"open"}, self.namespace(Path=Path, + sqlite3=self.engine_namespace["sqlite3"], DEFAULT_BUSY_TIMEOUT_MS=3000, + warnings=types.SimpleNamespace(warn=lambda *args, **kwargs: None)), methods=True) + def opener(path, callback): + return callback(engine.db_path) + with patch.object(self.api, "open_serving_connection", side_effect=opener), patch.object(Path, "exists", return_value=True): + with self.assertRaises(sqlite3.OperationalError): + module.open(engine) + self.assert_closed(candidate) + self.assertIs(engine.conn, self.conn) + self.assertTrue(engine.ready) + + + + def test_nul_tail_metadata_and_registry_are_rejected_before_materialization(self): + engine = self.engine() + candidate = self.candidate(engine) + candidate.executemany("INSERT INTO metadata(key,value) VALUES(?,?)", [ + ("embed_model", "model2vec\0" + "x" * 10000), ("embed_dim", "256")]) + candidate.commit() + self.assertIsNone(engine._reopen_initialized_connection(engine.db_path)) + self.assert_closed(candidate) + self.conn.execute("INSERT INTO vector_cache_registry(tier_group,vec_table,sep_table,model_name) VALUES('edge','vec_messages_edge','vec_messages_sep_edge',?)", + ("synthetic\0" + "x" * 10000,)) + self.conn.commit() + statements = [] + self.conn.set_trace_callback(statements.append) + self.assertIsNone(engine._reconnect_registry(self.conn)) + self.assertFalse(any(sql.startswith("SELECT tier_group,vec_table") for sql in statements)) + + def test_capture_cleanup_failure_never_publishes_receipt(self): + engine = self.engine() + operation = types.SimpleNamespace(key=self.api._legacy_key(), _database=("file", str(engine.db_path))) + with engine._write_lock, patch.object(self.conn, "rollback", side_effect=sqlite3.OperationalError("synthetic rollback failure")): + with self.assertRaises(sqlite3.OperationalError): + engine._capture_reconnect_receipt(operation) + self.assertIsNone(engine._reconnect_receipt) + self.assertTrue(self.conn.in_transaction) + self.assertIs(engine.conn, self.conn) + self.conn.rollback() + + def test_registry_rows_are_copied_to_primitive_tuples(self): + engine = self.engine() + self.conn.execute("INSERT INTO vector_cache_registry(tier_group,vec_table,sep_table) VALUES('edge','vec_messages_edge','vec_messages_sep_edge')") + self.conn.commit() + self.conn.row_factory = sqlite3.Row + operation = types.SimpleNamespace(key=self.api._legacy_key(), _database=("file", str(engine.db_path))) + with engine._write_lock: + engine._capture_reconnect_receipt(operation) + receipt = engine._reconnect_receipt + self.assertIsNotNone(receipt) + self.assertIs(type(receipt[6]), tuple) + self.assertIs(type(receipt[6][0]), tuple) + self.assertTrue(all(type(value) in (str, int, float, type(None)) for value in receipt[6][0])) + + + + def test_silent_capture_rollback_cannot_publish_receipt(self): + engine = self.engine() + operation = types.SimpleNamespace(key=self.api._legacy_key(), _database=("file", str(engine.db_path))) + with engine._write_lock, patch.object(self.conn, "rollback", return_value=None) as rollback: + with self.assertRaises(self.api.TierRuntimeError): + engine._capture_reconnect_receipt(operation) + rollback.assert_called_once_with() + self.assertIsNone(engine._reconnect_receipt) + self.assertTrue(self.conn.in_transaction) + self.assertIs(engine.conn, self.conn) + self.conn.rollback() + + def test_silent_candidate_rollback_closes_only_candidate(self): + engine = self.engine() + candidate = self.candidate(engine) + self.conn.armed, self.conn.fault_contains = True, "PRAGMA schema_version" + self.conn.fault = sqlite3.OperationalError("disk I/O error") + before = engine._reconnect_receipt + with patch.object(candidate, "rollback", return_value=None) as rollback, \ + patch.object(self.conn, "close", wraps=self.conn.close) as close: + with self.assertRaises(self.api.TierRuntimeError): + engine._open_connection_handle() + rollback.assert_called_once_with() + close.assert_not_called() + self.assert_closed(candidate) + self.assertIs(engine.conn, self.conn) + self.assertEqual(engine._reconnect_receipt, before) + self.assertTrue(engine.ready) + self.assertTrue(engine._runtime_initialized) + self.assertFalse(self.conn.in_transaction) + + + +class TestColdPublicInitialization(unittest.TestCase): + imports = TestPublicBoundaries.imports + namespace = TestPublicBoundaries.namespace + + def setUp(self): + TestPublicBoundaries.setUp(self) + import re + self.events = [] + self.conn.enable_load_extension = lambda value: self.events.append(("extension", value)) + native = types.SimpleNamespace(load=lambda conn: self.events.append(("load", conn))) + original_import = self.imports + def imports(name, globals=None, locals=None, fromlist=(), level=0): + if name == "sqlite_vec": + return native + return original_import(name, globals, locals, fromlist, level) + self.imports = imports + original_execute = self.conn.execute + def execute(sql, parameters=()): + match = re.fullmatch(r"CREATE VIRTUAL TABLE IF NOT EXISTS (\w+) USING vec0\(embedding float\[(\d+)\] distance_metric=cosine\)", sql) + if match: + sql = f'CREATE TABLE IF NOT EXISTS {match[1]} (rowid INTEGER PRIMARY KEY, embedding "float[{match[2]}] distance_metric=cosine")' + return original_execute(sql, parameters) + self.conn.execute = execute + known = ("vec_messages", "vec_messages_sep", "vec_messages_edge", "vec_messages_sep_edge", + "vec_messages_basepro", "vec_messages_sep_basepro") + for table in (*known, "vec_messages_custom", "vec_messages_sep_custom"): + self.conn.execute(f'DROP TABLE IF EXISTS "{table}"') + self.conn.commit() + self.vector = definitions("vector_search.py", { + "_active_vec_table", "_active_sep_table", "_ensure_metadata_table", "_read_embedder_metadata", + "_detect_existing_vec_dim", "_migration_hint", "_check_embedder_compatibility", "_table_uses_cosine", + "_vec0_column_decl", "init_vec_table", "migrate_to_cosine_metric", "migrate_legacy_vec_tables", + }, self.namespace(sqlite3=sqlite3, re=re, EMBEDDING_MODEL="model2vec", _embedding_dim=256, + _KNOWN_VEC_TABLES=known, _VALID_GROUPS={"edge", "basepro"}, _MODEL_TO_GROUP={"model2vec": "edge"}, + _VEC_DISTANCE_METRIC="cosine", _active_tier_group=lambda: "edge", + logger=logging.getLogger("synthetic-initialization"), TrueMemoryMigrationError=RuntimeError)) + self.fixture.vector._active_vec_table = self.vector._active_vec_table + self.fixture.vector._active_sep_table = self.vector._active_sep_table + self.modules["vector_search"] = self.vector + self.vector.search_vector_raw = lambda conn, query, limit: conn.execute( + f'SELECT rowid, embedding FROM "{self.vector._active_vec_table(conn)}" LIMIT ?', (limit,)).fetchall() + self.namespace_values = self.namespace(_HAS_VECTOR=True, _HAS_HYBRID=True, + init_vec_table=self.vector.init_vec_table, TrueMemoryMigrationError=RuntimeError, + resolve_tier=lambda: "edge", logger=logging.getLogger("synthetic-initialization"), + engine_operation=self.api.engine_operation) + self.methods = definitions("engine.py", {"_ensure_connection", "_initialize_connection", "_maybe_auto_consolidate", + "_maybe_startup_consolidate", "search_vectors_raw"}, self.namespace_values, methods=True) + self.engine = types.SimpleNamespace(conn=self.conn, ready=False, _runtime_initialized=False, + _has_vectors=False, _has_hybrid=False, _has_consolidation=False, _has_style_vec=False, + _init_lock=threading.Lock(), _write_lock=threading.Lock(), _open_connection_handle=lambda: None, + _purge_legacy_entity_profile_summaries=Mock(), _reconnect_receipt=None) + for name in ("_ensure_connection", "_initialize_connection", "_maybe_auto_consolidate", + "_maybe_startup_consolidate", "search_vectors_raw"): + setattr(self.engine, name, types.MethodType(getattr(self.methods, name), self.engine)) + self.conn.execute("INSERT OR REPLACE INTO metadata(key,value) VALUES('embed_model','model2vec')") + self.conn.execute("INSERT OR REPLACE INTO metadata(key,value) VALUES('embed_dim','256')") + self.conn.commit() + + def generic(self, *, populated=True): + for table in ("vec_messages", "vec_messages_sep"): + self.conn.execute(f'CREATE TABLE {table}(rowid INTEGER PRIMARY KEY, embedding "float[256] distance_metric=cosine")') + if populated: + self.conn.execute(f"INSERT INTO {table}(rowid,embedding) VALUES(7,?)", (b"synthetic vector",)) + self.conn.commit() + + def test_public_search_admits_post_migration_pair(self): + self.generic() + captured = [] + actual = self.vector.search_vector_raw + def search(*args, **kwargs): + captured.append(self.api.current_operation(self.conn)) + return actual(*args, **kwargs) + self.vector.search_vector_raw = search + rows = self.engine.search_vectors_raw("synthetic", limit=3) + self.assertEqual(rows, [(7, b"synthetic vector")]) + self.assertEqual(captured[0].tables, ("vec_messages_edge", "vec_messages_sep_edge")) + self.assertTrue(captured[0]._defer_maintenance) + self.assertIsNone(self.conn.execute("SELECT 1 FROM sqlite_master WHERE name='vec_messages'").fetchone()) + self.assertTrue(self.engine.ready) + self.assertIsNone(self.api.current_operation(self.conn)) + + def test_fact_scope_is_after_migration_and_before_storage(self): + self.generic() + with self.api.fact_operation(types.SimpleNamespace(_engine=self.engine)) as operation: + self.assertEqual(operation.tables, ("vec_messages_edge", "vec_messages_sep_edge")) + self.assertEqual(self.engine.search_vectors_raw("synthetic"), [(7, b"synthetic vector")]) + self.assertIs(self.api.current_operation(self.conn), operation) + + def test_nested_populated_generic_refuses_before_mutation(self): + self.generic() + queries = [] + self.conn.set_trace_callback(queries.append) + with self.api.serving_operation(self.conn) as operation: + with self.assertRaisesRegex(self.api.TierRuntimeError, "before entering"): + self.engine._ensure_connection() + self.assertIs(self.api.current_operation(self.conn), operation) + self.assertEqual(operation.tables, ("vec_messages", "vec_messages_sep")) + self.assertFalse(any(sql.lstrip().split()[0].upper() in {"CREATE", "DROP", "ALTER", "INSERT", "DELETE", "UPDATE"} for sql in queries)) + self.assertFalse(self.engine.ready) + self.assertEqual(self.conn.execute("SELECT embedding FROM vec_messages").fetchall(), [(b"synthetic vector",)]) + + def test_nested_first_empty_generic_can_initialize(self): + with self.api.serving_operation(self.conn) as operation: + self.engine._ensure_connection(_suppress_maintenance=True) + self.assertIs(self.api.current_operation(self.conn), operation) + self.assertEqual(operation.tables, ("vec_messages", "vec_messages_sep")) + self.assertTrue(self.engine.ready) + self.assertEqual(self.conn.execute("SELECT COUNT(*) FROM vec_messages").fetchone(), (0,)) + + def test_nested_absent_generic_preserves_unrelated_cached_pair(self): + self.conn.execute('CREATE TABLE vec_messages_custom(embedding BLOB)') + self.conn.execute('CREATE TABLE vec_messages_sep_custom(embedding BLOB)') + self.conn.execute("INSERT INTO vec_messages_custom VALUES(?)", (b"inactive synthetic",)) + self.conn.commit() + with self.api.serving_operation(self.conn) as operation: + self.engine._ensure_connection(_suppress_maintenance=True) + self.assertEqual(operation.tables, ("vec_messages", "vec_messages_sep")) + self.assertEqual(self.conn.execute("SELECT embedding FROM vec_messages_custom").fetchall(), [(b"inactive synthetic",)]) + self.assertTrue(self.engine.ready) + self.assertTrue(self.engine._has_vectors) + + def test_nested_stable_tiered_pair_only_loads_connection_extension(self): + self.generic() + self.engine._ensure_connection(_suppress_maintenance=True) + self.engine.ready = self.engine._runtime_initialized = False + queries = [] + self.conn.set_trace_callback(queries.append) + with self.api.serving_operation(self.conn) as operation: + self.engine._ensure_connection(_suppress_maintenance=True) + self.assertEqual(operation.tables, ("vec_messages_edge", "vec_messages_sep_edge")) + self.assertTrue(self.engine.ready) + self.assertFalse(any(sql.lstrip().split()[0].upper() in {"CREATE", "DROP", "ALTER", "INSERT", "DELETE", "UPDATE"} for sql in queries)) + + def test_maintenance_waits_for_outer_scope_and_nested_crud_calls(self): + self.generic() + scheduled = [] + self.engine._has_consolidation = True + self.modules["maintenance"].engine_layer_specs = lambda conn, coordinator: () + self.modules["maintenance"].observe_style = Mock(side_effect=AssertionError("Unexpected style")) + self.modules["maintenance"].plan_layers = lambda *args, **kwargs: scheduled.append(self.api.current_runtime_operation()) or () + self.engine._get_maintenance_coordinator = lambda: object() + self.engine._auto_consolidate_threshold = 32 + with self.api.engine_serving_operation(self.engine) as operation: + self.engine._maybe_auto_consolidate() + self.engine._ensure_connection() + with self.api.engine_serving_operation(self.engine) as nested: + self.assertIs(nested, operation) + self.engine._maybe_auto_consolidate() + self.assertEqual(scheduled, []) + self.assertEqual(scheduled, [None]) + + def test_borrowed_cold_initialization_preserves_caller_sql(self): + self.conn.execute("INSERT INTO messages(id,content) VALUES(91,'synthetic caller')") + before = (self.conn.commits, self.conn.rollbacks) + with self.assertRaisesRegex(self.api.TierRuntimeError, "clean connection"): + self.engine._ensure_connection() + self.assertTrue(self.conn.in_transaction) + self.assertEqual((self.conn.commits, self.conn.rollbacks), before) + self.assertFalse(self.engine.ready) + + def test_nested_custom_pair_preserves_tables_and_metadata(self): + self.generic() + self.engine._ensure_connection(_suppress_maintenance=True) + self.conn.execute("ALTER TABLE vec_messages_edge RENAME TO vec_messages_custom") + self.conn.execute("ALTER TABLE vec_messages_sep_edge RENAME TO vec_messages_sep_custom") + self.conn.execute("UPDATE vector_cache_registry SET vec_table='vec_messages_custom', sep_table='vec_messages_sep_custom'") + self.conn.commit() + self.engine.ready = self.engine._runtime_initialized = False + before = self.conn.total_changes + with self.api.serving_operation(self.conn) as operation: + self.engine._ensure_connection(_suppress_maintenance=True) + self.assertEqual(operation.tables, ("vec_messages_custom", "vec_messages_sep_custom")) + self.assertEqual(self.conn.total_changes, before) + self.assertTrue(self.engine.ready) + + def test_nested_cosine_migration_refuses_without_changing_source(self): + self.conn.execute('CREATE TABLE vec_messages(embedding "float[256]")') + self.conn.execute('CREATE TABLE vec_messages_sep(embedding "float[256]")') + with self.api.serving_operation(self.conn) as operation: + before = self.conn.execute("PRAGMA schema_version").fetchone() + with self.assertRaisesRegex(self.api.TierRuntimeError, "completed vector migration"): + self.engine._ensure_connection() + self.assertIs(self.api.current_operation(self.conn), operation) + self.assertEqual(self.conn.execute("PRAGMA schema_version").fetchone(), before) + self.assertFalse(self.engine.ready) + + def test_child_inherits_deferral_and_cannot_schedule_before_parent_exit(self): + self.generic() + scheduled, failures = [], [] + self.engine._has_consolidation = True + self.modules["maintenance"].engine_layer_specs = lambda conn, coordinator: () + self.modules["maintenance"].observe_style = Mock(side_effect=AssertionError("Unexpected style")) + self.modules["maintenance"].plan_layers = lambda *args, **kwargs: scheduled.append(self.api.current_runtime_operation()) or () + self.engine._get_maintenance_coordinator = lambda: object() + self.engine._auto_consolidate_threshold = 32 + with self.api.engine_serving_operation(self.engine) as operation: + child = operation.fork_child() + def run(): + try: + with child.join() as inherited: + self.assertTrue(inherited._defer_maintenance) + self.assertEqual(inherited.tables, operation.tables) + self.engine._maybe_auto_consolidate() + except BaseException as error: + failures.append(error) + worker = threading.Thread(target=run) + worker.start() + worker.join(2) + self.assertFalse(worker.is_alive()) + self.assertEqual(failures, []) + self.assertEqual(scheduled, []) + self.assertEqual(scheduled, [None]) + + def test_synchronous_nan_repair_precedes_legacy_name_migration(self): + self.generic() + self.engine.db_path = Path(":memory:") + self.vector.EMBEDDING_MODEL = self.fixture.vector.EMBEDDING_MODEL = "qwen3_256" + self.vector._MODEL_TO_GROUP["qwen3_256"] = "basepro" + self.vector._active_tier_group = lambda: "basepro" + self.fixture.vector.resolve_tier = self.methods.resolve_tier = lambda: "base" + self.conn.execute("UPDATE metadata SET value='qwen3_256' WHERE key='embed_model'") + self.conn.execute("INSERT INTO messages(id,content) VALUES(7,'synthetic repair source')") + self.conn.commit() + repaired = [] + def rebuild(conn, *, separation=False): + table = self.vector._active_sep_table(conn) if separation else self.vector._active_vec_table(conn) + repaired.append(table) + conn.execute(f'DELETE FROM "{table}"') + conn.execute(f'INSERT INTO "{table}"(rowid,embedding) VALUES(7,?)', (b"synthetic repaired",)) + conn.commit() + self.vector.build_vectors = rebuild + self.vector.build_separation_vectors = lambda conn: rebuild(conn, separation=True) + with patch.object(sys, "platform", "darwin"): + with self.api.engine_serving_operation(self.engine) as operation: + self.assertEqual(operation.tables, ("vec_messages_basepro", "vec_messages_sep_basepro")) + self.assertEqual(self.conn.execute("SELECT embedding FROM vec_messages_basepro").fetchall(), [(b"synthetic repaired",)]) + self.assertEqual(self.conn.execute("SELECT embedding FROM vec_messages_sep_basepro").fetchall(), [(b"synthetic repaired",)]) + self.assertEqual(repaired, ["vec_messages", "vec_messages_sep"]) + self.assertEqual(self.conn.execute("SELECT value FROM metadata WHERE key='qwen3_nan_fix_applied'").fetchone(), ("1",)) + self.assertIsNone(self.conn.execute("SELECT 1 FROM sqlite_master WHERE name='vec_messages'").fetchone()) + + +class TestBorrowedPublicCompatibility(unittest.TestCase): + imports = TestPublicBoundaries.imports + namespace = TestPublicBoundaries.namespace + + def setUp(self): + TestPublicBoundaries.setUp(self) + storage = self.modules["storage"] + self.conn = storage.create_db(":memory:") + self.addCleanup(self.conn.close) + self.storage = storage + fts = definitions("fts_search.py", {"search_fts", "_build_safe_query", "_fts_select", + "_rows_to_results", "_normalize_scores"}, self.namespace(sqlite3=sqlite3, + _deserialize_metadata=storage._deserialize_metadata, directive_filter_sql=storage.directive_filter_sql, + select_message_cols=storage.select_message_cols)) + source = ast.parse((ROOT / "truememory/engine.py").read_text(encoding="utf-8")) + cap = next(node.value for node in source.body if isinstance(node, ast.Assign) + and any(isinstance(target, ast.Name) and target.id == "MAX_CONTENT_LENGTH" for target in node.targets)) + self.limit = ast.literal_eval(cap) + validators = definitions("engine.py", {"_validate_add_content", "_validate_delete_user"}, + self.namespace(MAX_CONTENT_LENGTH=self.limit)) + self.methods = definitions("engine.py", {"_open_connection_handle", "_initialize_connection", + "add", "delete_all", "search_simple"}, self.namespace( + _validate_add_content=validators._validate_add_content, _validate_delete_user=validators._validate_delete_user, + engine_operation=self.api.engine_operation, engine_handle_operation=self.api.engine_handle_operation, + search_fts=fts.search_fts, sqlite3=sqlite3, _HAS_VECTOR=False, _HAS_HYBRID=False), methods=True) + self.engine = types.SimpleNamespace(conn=self.conn, ready=False, _runtime_initialized=False, + _write_lock=threading.Lock(), _init_lock=threading.Lock(), _has_vectors=False, + _has_consolidation=False, _has_style_vec=False, + _ensure_connection=Mock(side_effect=AssertionError("Unexpected body initialization")), + _maybe_auto_consolidate=Mock(side_effect=AssertionError("Unexpected maintenance")), + _maybe_startup_consolidate=Mock(side_effect=AssertionError("Unexpected startup maintenance"))) + for name in ("_open_connection_handle", "_initialize_connection", "add", "delete_all", "search_simple"): + setattr(self.engine, name, types.MethodType(getattr(self.methods, name), self.engine)) + + def borrow(self): + self.conn.execute("BEGIN") + identity = self.storage.insert_message(self.conn, {"content": "syntheticpending borrowed lexical source"}) + self.assertTrue(self.conn.in_transaction) + self.queries = [] + self.conn.set_trace_callback(self.queries.append) + return identity + + def assert_borrowed(self, identity): + self.assertTrue(self.conn.in_transaction) + self.assertEqual(self.conn.execute("SELECT content FROM messages WHERE id=?", (identity,)).fetchone(), + ("syntheticpending borrowed lexical source",)) + self.assertFalse(any(query.lstrip().split()[0].upper() in {"BEGIN", "COMMIT", "ROLLBACK", "SAVEPOINT", "RELEASE"} + for query in self.queries)) + self.assertFalse(self.engine.ready) + self.assertFalse(self.engine._runtime_initialized) + self.engine._maybe_auto_consolidate.assert_not_called() + self.engine._maybe_startup_consolidate.assert_not_called() + + def test_simple_search_reads_borrowed_fts_without_vector_initialization(self): + identity = self.borrow() + with patch.object(self.engine, "_initialize_connection", side_effect=AssertionError("Unexpected vector initialization")) as initialize: + results = self.engine.search_simple("syntheticpending") + self.assertEqual([(row["id"], row["source"]) for row in results], [(identity, "fts")]) + self.assertEqual(results[0]["score"], 1.0) + initialize.assert_not_called() + self.assert_borrowed(identity) + + def test_simple_search_clean_cold_engine_stays_fts_only(self): + identity = self.storage.insert_message(self.conn, {"content": "synthetic lexical only"}) + self.conn.commit() + with patch.object(self.engine, "_initialize_connection", side_effect=AssertionError("Unexpected vector initialization")) as initialize: + results = self.engine.search_simple("synthetic") + self.assertEqual([(row["id"], row["source"]) for row in results], [(identity, "fts")]) + initialize.assert_not_called() + self.assertFalse(self.engine.ready) + self.engine._maybe_auto_consolidate.assert_not_called() + self.engine._maybe_startup_consolidate.assert_not_called() + + def test_simple_search_propagates_typed_refusal_without_ending_caller_sql(self): + identity = self.borrow() + error = self.api.TierRuntimeError("synthetic admission changed") + self.methods.search_fts = Mock(side_effect=error) + with self.assertRaises(self.api.TierRuntimeError) as raised: + self.engine.search_simple("syntheticpending") + self.assertIs(raised.exception, error) + self.assert_borrowed(identity) + + def test_invalid_add_precedes_any_handle_or_sql_admission(self): + identity = self.borrow() + with patch.object(self.engine, "_open_connection_handle", side_effect=AssertionError("Unexpected handle admission")) as opening: + with self.assertRaises(TypeError) as raised: + self.engine.add(None) + self.assertEqual(str(raised.exception), "content must be a string, got NoneType") + self.assertEqual(self.queries, []) + opening.assert_not_called() + self.assert_borrowed(identity) + + def test_oversized_add_keeps_exact_limit_error_and_caller_sql(self): + identity = self.borrow() + with patch.object(self.engine, "_open_connection_handle", side_effect=AssertionError("Unexpected handle admission")) as opening: + with self.assertRaises(ValueError) as raised: + self.engine.add(content="x" * (self.limit + 1)) + self.assertEqual(str(raised.exception), f"Content too large ({self.limit + 1} chars). Maximum is {self.limit}.") + self.assertEqual(self.queries, []) + opening.assert_not_called() + self.assert_borrowed(identity) + + def test_valid_add_arguments_still_refuse_cold_vector_initialization_in_caller_sql(self): + identity = self.borrow() + for content in ("", "synthetic", "x" * self.limit): + with self.subTest(size=len(content)), self.assertRaisesRegex(self.api.TierRuntimeError, "clean connection"): + self.engine.add(content) + self.assert_borrowed(identity) + + def test_invalid_delete_scope_precedes_handle_or_sql_admission(self): + identity = self.borrow() + for value, error_type, message in ((3, TypeError, "user_id must be a string or None, got int"), + ([], TypeError, "user_id must be a string or None, got list"), + ("", ValueError, "user_id cannot be an empty string"), + (" \t", ValueError, "user_id cannot be an empty string")): + with self.subTest(value=value), patch.object(self.engine, "_open_connection_handle", side_effect=AssertionError("Unexpected handle admission")) as opening: + with self.assertRaises(error_type) as raised: + self.engine.delete_all(user_id=value) + self.assertEqual(str(raised.exception), message) + opening.assert_not_called() + self.assertEqual(self.queries, []) + self.assert_borrowed(identity) + + def test_valid_delete_scope_still_requires_clean_cold_initialization(self): + identity = self.borrow() + for value in (None, "synthetic"): + with self.subTest(value=value), self.assertRaisesRegex(self.api.TierRuntimeError, "clean connection"): + self.engine.delete_all(user_id=value) + self.assertFalse(any(query.lstrip().upper().startswith("DELETE") for query in self.queries)) + self.assert_borrowed(identity) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test-tier-runtime-bridge-795.py b/tests/test-tier-runtime-bridge-795.py index 413bdd1a..7d814db7 100644 --- a/tests/test-tier-runtime-bridge-795.py +++ b/tests/test-tier-runtime-bridge-795.py @@ -148,6 +148,50 @@ def loader(**kwargs): pass self.assertEqual(observed, [{"model_name": "synthetic/deep", "deadline": None, "cancelled": cancelled}]) + def test_legacy_optional_reranker_failure_preserves_bound_retrieval(self) -> None: + for error in (OSError("synthetic unavailable"), ImportError("synthetic missing")): + with self.subTest(error=type(error).__name__): + with patch.object(self.reranker, "get_reranker", side_effect=error) as loader: + with self.api.serving_operation(self.conn, reranker_id="synthetic/deep") as operation: + self.assertEqual(operation.reranker_id, "synthetic/deep") + self.assertIs(self.api.current_operation(self.conn), operation) + self.assertIsNone(operation.selection) + self.assertIsNone(operation.policy) + loader.assert_called_once() + self.assertIsNone(self.api.current_operation(self.conn)) + self.assertFalse(self.conn.in_transaction) + + def test_legacy_reranker_identity_and_control_refusals_remain_hard(self) -> None: + for error in (self.api.TierRuntimeError("identity"), self.gate.ServingLeaseError("admission"), + self.gate.ServingLeaseTimeout("timeout"), self.gate.ServingLeaseCancelled("cancelled")): + with self.subTest(error=type(error).__name__): + with patch.object(self.reranker, "get_reranker", side_effect=error): + with self.assertRaises(type(error)): + with self.api.serving_operation(self.conn, reranker_id="synthetic/deep"): + self.fail("Typed refusal must prevent retrieval") + self.assertIsNone(self.api.current_operation(self.conn)) + + def test_borrowed_legacy_uncached_reranker_never_attempts_load(self) -> None: + self.conn.execute("BEGIN") + before = (self.conn.commits, self.conn.rollbacks) + with patch.object(self.reranker, "get_reranker", side_effect=OSError("synthetic unavailable")) as loader: + with self.assertRaises(self.api.TierRuntimeError): + with self.api.serving_operation(self.conn, reranker_id="synthetic/deep"): + self.fail("Borrowed SQL cannot load an uncached reranker") + loader.assert_not_called() + self.assertTrue(self.conn.in_transaction) + self.assertEqual((self.conn.commits, self.conn.rollbacks), before) + self.conn.rollback() + + def test_selected_reranker_preparation_failure_never_enters_body(self) -> None: + self.select() + with patch.object(self.reranker, "apply_frozen_reranker", side_effect=OSError("synthetic unavailable")): + with self.assertRaises(OSError): + with self.api.serving_operation(self.conn, reranker_id="synthetic/deep"): + self.fail("Selected certification is required before retrieval") + self.assertIsNone(self.api.current_operation(self.conn)) + self.assertIsNone(self.api.runtime_acknowledgement()) + def test_nested_search_has_zero_selection_sql_and_reuses_context(self) -> None: self.select() queries = [] @@ -1095,7 +1139,8 @@ def predict(self, pairs, **kwargs): 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)) + EMBEDDING_MODEL="model2vec", _embedding_dim=256, _model_generation=0, + _runtime_policy_tier=None)) 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)) diff --git a/tests/test-tier-warm-admission-795.py b/tests/test-tier-warm-admission-795.py new file mode 100644 index 00000000..eb8c7fe8 --- /dev/null +++ b/tests/test-tier-warm-admission-795.py @@ -0,0 +1,493 @@ +"""Fresh public admission against real lifecycle code and memory SQLite only.""" +from __future__ import annotations + +import contextlib +from pathlib import Path +import runpy +import threading +import time +import types +import unittest +from unittest.mock import Mock, patch + +ACTIVE = runpy.run_path(str(Path(__file__).with_name("test-tier-active-initialization-795.py"))) + + +class TestWarmAdmission(unittest.TestCase): + def setUp(self): + self.active = ACTIVE["TestActiveInitialization"]() + self.active.setUp() + self.addCleanup(self.active.doCleanups) + self.fixture = self.active.fixture + self.api, self.maintenance = self.active.api, self.active.maintenance + self.coordinator = self.active.coordinator + self.gate = self.fixture.fixture.gate + self.maintenance.threading.TIMEOUT_MAX = threading.TIMEOUT_MAX + loaded = ACTIVE["definitions"]("maintenance.py", {"wait_for_automatic_owner"}, + {**self.maintenance.__dict__, "time": time}) + self.maintenance.wait_for_automatic_owner = loaded.wait_for_automatic_owner + self.opened = self.fixture.memory_connection() + self.fixture.engine_namespace["create_db"] = Mock(return_value=self.opened) + self.errors = [] + initialization = ACTIVE["definitions"]("engine.py", { + "_initialize_connection", "_purge_legacy_entity_profile_summaries", + "_maybe_startup_consolidate", "_maybe_auto_consolidate"}, + {**self.fixture.engine_namespace, "_HAS_HYBRID": False, "os": __import__("os")}, methods=True) + for name in ("_initialize_connection", "_purge_legacy_entity_profile_summaries", + "_maybe_startup_consolidate", "_maybe_auto_consolidate"): + setattr(self.active.first.__class__, name, getattr(initialization, name)) + self.active.first.__class__._has_consolidation = False + + def drifted_candidate(self): + candidate = self.fixture.candidate(self.active.first) + candidate.execute("CREATE TABLE synthetic_changed_schema(value INTEGER)") + return candidate + + def run_public(self, engine): + try: + self.api.open_engine_for_operation(engine) + except BaseException as error: + self.errors.append(error) + + def observe_wait(self, action=None): + entered = threading.Event() + actual = self.maintenance.wait_for_automatic_owner + def wait(error, *, deadline): + self.assertIsNone(self.api.current_runtime_operation()) + self.assertFalse(self.gate.current_thread_admitted()) + self.assertTrue(self.maintenance._registry_lock.acquire(blocking=False)) + self.maintenance._registry_lock.release() + self.assertTrue(self.coordinator._mutex.acquire(blocking=False)) + self.coordinator._mutex.release() + entered.set() + if action is not None: + action() + return actual(error, deadline=deadline) + self.enterContext(patch.object(self.maintenance, "wait_for_automatic_owner", side_effect=wait)) + return entered + + def test_schema_drift_waits_then_uses_normal_factory_without_stale_readiness(self): + self.active.start() + candidate = self.drifted_candidate() + engine = self.active.fresh() + entered = self.observe_wait() + worker = threading.Thread(target=self.run_public, args=(engine,)) + worker.start() + try: + self.assertTrue(entered.wait(2)) + self.assertIsNone(engine.conn) + self.assertFalse(engine.ready) + self.assertTrue(engine._write_lock.acquire(blocking=False)) + engine._write_lock.release() + self.assertTrue(engine._init_lock.acquire(blocking=False)) + engine._init_lock.release() + self.fixture.assert_closed(candidate) + self.active.release.set() + worker.join(2) + self.assertFalse(worker.is_alive()) + self.assertEqual(self.errors, []) + self.assertIs(engine.conn, self.opened) + self.assertFalse(engine.ready) + self.assertIsNone(engine._reconnect_receipt) + self.assertEqual(self.fixture.fixture.loads, []) + self.fixture.engine_namespace["create_db"].assert_called_once() + finally: + self.active.release.set() + worker.join(3) + + def test_private_opener_still_refuses_same_schema_drift_without_wait(self): + self.active.start() + self.drifted_candidate() + engine = self.active.fresh() + with patch.object(self.maintenance, "wait_for_automatic_owner") as wait: + with self.assertRaises(self.maintenance.MaintenanceBusyError): + engine._open_connection_handle() + wait.assert_not_called() + self.assertIsNone(engine.conn) + + def test_successful_first_attempt_preserves_existing_raw_reader_and_exclusive_behavior(self): + for context in (lambda: self.gate.operation_lease(("synthetic",)), self.gate.exclusive_activation): + with self.subTest(context=context), context(): + engine = self.active.fresh() + with patch.object(self.maintenance, "wait_for_automatic_owner") as wait: + self.api.open_engine_for_operation(engine) + self.assertIs(engine.conn, self.opened) + wait.assert_not_called() + + def test_busy_runtime_raw_reader_and_exclusive_scopes_never_wait(self): + self.active.start() + scopes = (lambda: self.api.serving_operation(self.fixture.conn), + lambda: self.gate.operation_lease(("synthetic",)), self.gate.exclusive_activation) + for scope in scopes: + with self.subTest(scope=scope), scope(): + self.drifted_candidate() + engine = self.active.fresh() + with patch.object(self.maintenance, "wait_for_automatic_owner") as wait: + with self.assertRaises(self.maintenance.MaintenanceBusyError): + self.api.open_engine_for_operation(engine) + wait.assert_not_called() + self.assertIsNone(engine.conn) + + def test_current_writer_refuses_before_nonreentrant_open(self): + engine = self.active.fresh() + with engine._write_lock, patch.object(engine, "_open_connection_handle") as opened: + with self.assertRaises(self.maintenance.MaintenanceBusyError): + self.api.open_engine_for_operation(engine) + opened.assert_not_called() + + def test_existing_connection_and_caller_sql_never_enter_wait_loop(self): + engine = self.active.first + engine.conn.execute("BEGIN") + error = self.maintenance.MaintenanceBusyError("synthetic existing connection") + with patch.object(engine, "_open_connection_handle", side_effect=error) as opened, \ + patch.object(self.maintenance, "wait_for_automatic_owner") as wait: + with self.assertRaises(self.maintenance.MaintenanceBusyError) as raised: + self.api.open_engine_for_operation(engine) + self.assertIs(raised.exception, error) + self.assertTrue(engine.conn.in_transaction) + opened.assert_called_once() + wait.assert_not_called() + engine.conn.rollback() + + def test_manual_receiptless_and_foreign_owners_have_no_wait_permission(self): + for kind in ("manual", "receiptless", "foreign"): + with self.subTest(kind=kind): + if kind == "manual": + self.assertTrue(self.coordinator.request(self.active.work)) + self.assertTrue(self.active.entered.wait(2)) + elif kind == "receiptless": + self.active.start(receipt=False) + else: + self.maintenance._held_paths[self.active.first.db_path] = 991 + try: + with self.assertRaises(self.maintenance.MaintenanceBusyError) as raised: + self.api.open_engine_for_operation(self.active.fresh()) + self.assertIsNone(getattr(raised.exception, "_automatic_owner", None)) + finally: + if kind == "foreign": + self.maintenance._held_paths.clear() + else: + self.active.release.set() + self.assertTrue(self.coordinator.wait(2)) + self.active.entered.clear() + self.active.release.clear() + self.active.worker_conn = self.fixture.memory_connection() + + def test_launch_gap_wait_origin_precedes_active_receipt_publication(self): + acquired, continue_launch = threading.Event(), threading.Event() + actual = self.coordinator._begin_observation + def observe(): + if not acquired.is_set(): + acquired.set() + if not continue_launch.wait(3): + raise AssertionError("Synthetic launch barrier timed out") + return actual() + self.enterContext(patch.object(self.coordinator, "_begin_observation", side_effect=observe)) + requester = threading.Thread(target=lambda: self.coordinator.request_layers( + _initialization_receipt=self.active.first._reconnect_receipt)) + requester.start() + engine = self.active.fresh() + entered = self.observe_wait() + worker = threading.Thread(target=self.run_public, args=(engine,)) + try: + self.assertTrue(acquired.wait(2)) + self.assertIsNone(self.active.record()) + self.assertIsNotNone(self.coordinator._admission_owner) + worker.start() + self.assertTrue(entered.wait(2)) + continue_launch.set() + self.active.release.set() + worker.join(2) + self.assertFalse(worker.is_alive()) + self.assertEqual(self.errors, []) + self.assertIs(engine.conn, self.opened) + finally: + continue_launch.set() + self.active.release.set() + requester.join(3) + if worker.ident is not None: + worker.join(3) + + def test_teardown_gap_keeps_wait_origin_until_actual_descriptor_release(self): + self.active.start() + at_release, continue_release = threading.Event(), threading.Event() + actual = self.maintenance._release_owner + def release(fd): + at_release.set() + if not continue_release.wait(3): + raise AssertionError("Synthetic release barrier timed out") + return actual(fd) + self.enterContext(patch.object(self.maintenance, "_release_owner", side_effect=release)) + self.active.release.set() + self.assertTrue(at_release.wait(2)) + self.assertIsNone(self.active.record()) + origin = self.coordinator._admission_owner + self.assertIsNotNone(origin) + entered = self.observe_wait() + engine = self.active.fresh() + worker = threading.Thread(target=self.run_public, args=(engine,)) + worker.start() + try: + self.assertTrue(entered.wait(2)) + self.assertFalse(origin.released.is_set()) + continue_release.set() + worker.join(2) + self.assertFalse(worker.is_alive()) + self.assertEqual(self.errors, []) + self.assertTrue(origin.released.is_set()) + self.assertIsNone(self.coordinator._admission_owner) + finally: + continue_release.set() + worker.join(3) + + def test_cancel_wakes_waiter_without_releasing_owner_or_retrying(self): + self.active.start() + self.drifted_candidate() + engine = self.active.fresh() + entered = self.observe_wait() + worker = threading.Thread(target=self.run_public, args=(engine,)) + worker.start() + try: + self.assertTrue(entered.wait(2)) + origin = self.coordinator._admission_owner + self.coordinator.cancel() + worker.join(2) + self.assertFalse(worker.is_alive()) + self.assertEqual(len(self.errors), 1) + self.assertIs(type(self.errors[0]), self.maintenance.MaintenanceBusyError) + self.assertTrue(origin.cancelled.is_set()) + self.assertIn(self.active.first.db_path, self.maintenance._held_paths) + self.assertIsNone(engine.conn) + self.fixture.engine_namespace["create_db"].assert_not_called() + finally: + self.active.release.set() + worker.join(3) + + def test_timeout_leaves_owner_and_engine_untouched(self): + self.active.start() + candidate = self.drifted_candidate() + engine = self.active.fresh() + with patch.object(self.fixture.modules["storage"], "DEFAULT_BUSY_TIMEOUT_MS", 20): + with self.assertRaises(self.maintenance.MaintenanceBusyError): + self.api.open_engine_for_operation(engine) + self.assertIsNone(engine.conn) + self.assertFalse(engine.ready) + self.assertFalse(self.active.release.is_set()) + self.assertIn(self.active.first.db_path, self.maintenance._held_paths) + self.fixture.assert_closed(candidate) + + def test_successive_automatic_generations_share_one_deadline(self): + self.active.start() + self.drifted_candidate() + deadlines = [] + actual = self.maintenance.wait_for_automatic_owner + def wait(error, *, deadline): + deadlines.append(deadline) + self.active.release.set() + actual(error, deadline=deadline) + self.assertTrue(self.coordinator.wait(2)) + if len(deadlines) == 1: + self.active.entered.clear() + self.active.release.clear() + self.active.worker_conn = self.fixture.memory_connection() + self.active.start() + self.drifted_candidate() + engine = self.active.fresh() + with patch.object(self.maintenance, "wait_for_automatic_owner", side_effect=wait): + self.api.open_engine_for_operation(engine) + self.assertEqual(len(deadlines), 2) + self.assertEqual(deadlines[0], deadlines[1]) + self.assertIs(engine.conn, self.opened) + + def test_context_is_rechecked_after_wake_before_retry(self): + self.active.start() + self.drifted_candidate() + engine = self.active.fresh() + stack = contextlib.ExitStack() + self.addCleanup(stack.close) + actual = self.maintenance.wait_for_automatic_owner + def wait(error, *, deadline): + self.active.release.set() + actual(error, deadline=deadline) + stack.enter_context(self.gate.operation_lease(("synthetic",))) + with patch.object(self.maintenance, "wait_for_automatic_owner", side_effect=wait): + with self.assertRaises(self.maintenance.MaintenanceBusyError): + self.api.open_engine_for_operation(engine) + self.assertIsNone(engine.conn) + self.fixture.engine_namespace["create_db"].assert_not_called() + + def test_busy_origin_survives_release_before_wait_is_called(self): + self.active.start() + self.drifted_candidate() + engine = self.active.fresh() + actual = engine._open_connection_handle + calls = [] + def opening(): + calls.append(None) + try: + return actual() + except self.maintenance.MaintenanceBusyError as error: + self.assertIsNotNone(error._automatic_owner) + self.active.release.set() + self.assertTrue(self.coordinator.wait(2)) + self.assertTrue(error._automatic_owner.released.is_set()) + self.assertIsNone(self.coordinator._admission_owner) + raise + with patch.object(engine, "_open_connection_handle", side_effect=opening): + self.api.open_engine_for_operation(engine) + self.assertEqual(len(calls), 2) + self.assertIs(engine.conn, self.opened) + + def test_receiptless_successor_cannot_inherit_wait_permission(self): + self.active.start() + self.drifted_candidate() + actual = self.maintenance.wait_for_automatic_owner + origins, deadlines = [], [] + def wait(error, *, deadline): + origins.append(getattr(error, "_automatic_owner", None)) + deadlines.append(deadline) + if len(origins) == 1: + self.active.release.set() + actual(error, deadline=deadline) + self.assertTrue(self.coordinator.wait(2)) + self.active.entered.clear() + self.active.release.clear() + self.active.worker_conn = self.fixture.memory_connection() + self.active.start(receipt=False) + engine = self.active.fresh() + with patch.object(self.maintenance, "wait_for_automatic_owner", side_effect=wait): + with self.assertRaises(self.maintenance.MaintenanceBusyError): + self.api.open_engine_for_operation(engine) + self.assertEqual(len(origins), 2) + self.assertIsNotNone(origins[0]) + self.assertIsNone(origins[1]) + self.assertEqual(deadlines[0], deadlines[1]) + self.assertIsNone(engine.conn) + self.fixture.engine_namespace["create_db"].assert_not_called() + + def test_cancelled_origin_stays_cancelled_after_later_launch_and_fd_reuse(self): + self.active.start() + with self.assertRaises(self.maintenance.MaintenanceBusyError) as raised: + with self.maintenance.maintenance_owner(self.active.first.db_path): + self.fail("Automatic ownership was stolen") + old = raised.exception._automatic_owner + self.coordinator.cancel() + self.active.release.set() + self.assertTrue(self.coordinator.wait(2)) + self.active.entered.clear() + self.active.release.clear() + self.active.worker_conn = self.fixture.memory_connection() + self.active.start() + current = self.coordinator._admission_owner + self.assertEqual(old.descriptor, current.descriptor) + self.assertNotEqual(old.launch, current.launch) + self.assertTrue(old.cancelled.is_set()) + self.assertFalse(current.cancelled.is_set()) + with self.assertRaises(self.maintenance.MaintenanceBusyError) as again: + self.maintenance.wait_for_automatic_owner(raised.exception, deadline=time.monotonic() + 10) + self.assertIs(again.exception, raised.exception) + + def test_decorator_fact_and_ensure_connection_use_public_helper(self): + called = [] + engine = self.active.fresh() + actual = self.api.open_engine_for_operation + def opening(value): + called.append(value) + return actual(value) + @self.api.engine_operation + def body(value): + self.assertIsNotNone(self.api.current_runtime_operation()) + return value.conn + with patch.object(self.api, "open_engine_for_operation", side_effect=opening): + self.assertIs(body(engine), self.opened) + with self.api.fact_operation(types.SimpleNamespace(_engine=engine)): + self.assertIsNotNone(self.api.current_runtime_operation()) + methods = ACTIVE["definitions"]("engine.py", {"_ensure_connection"}, + self.fixture.namespace(), methods=True) + engine._initialize_connection = Mock() + methods._ensure_connection(engine) + self.assertEqual(called, [engine, engine, engine]) + engine._initialize_connection.assert_called_once_with(_suppress_maintenance=False) + + def test_gate_query_refuses_contention_without_waiting(self): + held, release = threading.Event(), threading.Event() + def holder(): + with self.gate._DEFAULT_GATE._condition: + held.set() + release.wait(3) + thread = threading.Thread(target=holder) + thread.start() + try: + self.assertTrue(held.wait(2)) + with self.assertRaises(self.gate.ServingLeaseError): + self.gate.current_thread_admitted() + finally: + release.set() + thread.join(3) + + def test_wait_eligibility_refuses_registry_and_coordinator_contention(self): + self.active.start() + with self.assertRaises(self.maintenance.MaintenanceBusyError) as raised: + with self.maintenance.maintenance_owner(self.active.first.db_path): + self.fail("Automatic ownership was stolen") + error = raised.exception + for lock in (self.maintenance._registry_lock, self.coordinator._mutex): + with self.subTest(lock=lock): + entered, done = threading.Event(), threading.Event() + errors = [] + def waiting(): + entered.set() + try: + self.maintenance.wait_for_automatic_owner(error, deadline=time.monotonic() + 0.035) + except BaseException as failure: + errors.append(failure) + finally: + done.set() + with lock: + worker = threading.Thread(target=waiting) + worker.start() + self.assertTrue(entered.wait(2)) + finished_while_held = done.wait(0.1) + worker.join(2) + self.assertTrue(finished_while_held) + self.assertEqual(errors, [error]) + self.assertFalse(self.active.release.is_set()) + + def test_runtime_presence_refuses_without_validating_lease_under_contended_gate(self): + self.active.start() + self.drifted_candidate() + engine = self.active.fresh() + admitted_connection = self.fixture.memory_connection() + admitted, proceed, caught = threading.Event(), threading.Event(), threading.Event() + errors = [] + def nested(): + with self.api.serving_operation(admitted_connection): + admitted.set() + if not proceed.wait(3): + raise AssertionError("Synthetic runtime barrier timed out") + try: + self.api.open_engine_for_operation(engine) + except BaseException as error: + errors.append(error) + finally: + caught.set() + worker = threading.Thread(target=nested) + worker.start() + try: + self.assertTrue(admitted.wait(2)) + with self.gate._DEFAULT_GATE._condition: + proceed.set() + refused_while_held = caught.wait(0.1) + worker.join(2) + self.assertTrue(refused_while_held) + self.assertFalse(worker.is_alive()) + self.assertEqual(len(errors), 1) + self.assertIs(type(errors[0]), self.maintenance.MaintenanceBusyError) + self.assertIsNone(engine.conn) + finally: + proceed.set() + worker.join(3) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test-tier-writer-fence-795.py b/tests/test-tier-writer-fence-795.py index 5ba2cbfb..f62dfe99 100644 --- a/tests/test-tier-writer-fence-795.py +++ b/tests/test-tier-writer-fence-795.py @@ -126,9 +126,46 @@ def encode(model, texts, **kwargs): 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() + # Keep the fixture's unbound operation contract, but classify actual + # publication/admission failures with the production helper. + serving_path = ROOT / "truememory/tier_switch/serving.py" + serving_tree = ast.parse(serving_path.read_text(encoding="utf-8")) + serving_names = {"ServingLeaseError", "ServingLeaseTimeout", "ServingLeaseCancelled"} + serving_namespace = {} + exec(compile(ast.Module(body=[node for node in serving_tree.body + if isinstance(node, ast.ClassDef) and node.name in serving_names], + type_ignores=[]), str(serving_path), "exec"), serving_namespace) + runtime_path = ROOT / "truememory/tier_switch/runtime.py" + runtime_tree = ast.parse(runtime_path.read_text(encoding="utf-8")) + runtime_names = {"TierRuntimeError", "raise_if_serving_rejection"} + runtime_namespace = { + "sys": types.SimpleNamespace(modules={"truememory." + name: module + for name, module in self.modules.items()}), + "serving": types.SimpleNamespace(**{name: serving_namespace[name] for name in serving_names}), + } + runtime_nodes = [ast.ImportFrom(module="__future__", names=[ast.alias(name="annotations")], level=0)] + runtime_nodes.extend(node for node in runtime_tree.body + if isinstance(node, (ast.ClassDef, ast.FunctionDef)) and node.name in runtime_names) + exec(compile(ast.fix_missing_locations(ast.Module(body=runtime_nodes, type_ignores=[])), + str(runtime_path), "exec"), runtime_namespace) + runtime = self.modules["tier_switch.runtime"] + for name in runtime_names: + setattr(runtime, name, runtime_namespace[name]) + key = "truememory.tier_switch.runtime" + missing = object() + previous = sys.modules.get(key, missing) + sys.modules[key] = types.SimpleNamespace( + current_operation=lambda conn: None, + **{name: runtime_namespace[name] for name in runtime_names}, + ) + + def restore_runtime(): + if previous is missing: + sys.modules.pop(key, None) + else: + sys.modules[key] = previous + + self.addCleanup(restore_runtime) def connect(self, uri: str | None = None): conn = sqlite3.connect(uri or ":memory:", uri=uri is not None, factory=Connection) @@ -187,7 +224,7 @@ def engine(self, *, vectors=True): 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"}] + helpers = [n for n in tree.body if isinstance(n, ast.FunctionDef) and n.name in {"_delete_in_chunks", "_resolve_vec_tables", "_validate_add_content", "_validate_delete_user"}] 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, diff --git a/tests/test_configure_reembed.py b/tests/test_configure_reembed.py index e267b578..58915de2 100644 --- a/tests/test_configure_reembed.py +++ b/tests/test_configure_reembed.py @@ -31,8 +31,19 @@ def server(monkeypatch, tmp_path): import truememory.mcp_server as ms monkeypatch.setattr(ms, "_TRUEMEMORY_DIR", home / ".truememory") monkeypatch.setattr(ms, "_CONFIG_PATH", home / ".truememory" / "config.json") + monkeypatch.setattr(ms, "_CONFIG_LOCK_PATH", ms._CONFIG_PATH.with_name("config.json.lock")) monkeypatch.setattr(ms, "_DB_PATH", str(db_path)) monkeypatch.setattr(ms, "_memory", None) + import numpy as np + import truememory.vector_search as vs + + class SyntheticEncoder: + def encode(self, texts, **kwargs): + # 256 * (1 / 16) ** 2 = 1, so the synthetic vectors have unit norm. + return np.full((len(texts), 256), 1 / 16, dtype=np.float32) + + encoder = SyntheticEncoder() + monkeypatch.setattr(vs, "get_model", lambda: encoder) yield ms # Teardown: drop cached Memory and any model-level state mutated by # truememory_configure (which modifies vector_search globals via diff --git a/tests/test_configure_tier_persist.py b/tests/test_configure_tier_persist.py index 2d60164c..34b149e9 100644 --- a/tests/test_configure_tier_persist.py +++ b/tests/test_configure_tier_persist.py @@ -32,6 +32,7 @@ def server(monkeypatch, tmp_path): import truememory.mcp_server as ms monkeypatch.setattr(ms, "_TRUEMEMORY_DIR", home / ".truememory") monkeypatch.setattr(ms, "_CONFIG_PATH", home / ".truememory" / "config.json") + monkeypatch.setattr(ms, "_CONFIG_LOCK_PATH", ms._CONFIG_PATH.with_name("config.json.lock")) monkeypatch.setattr(ms, "_DB_PATH", str(db_path)) monkeypatch.setattr(ms, "_memory", None) yield ms @@ -46,12 +47,17 @@ def server(monkeypatch, tmp_path): def _no_op_model(server, monkeypatch): - """Stub out the model/reranker side effects so the test stays unit-level.""" + """Exercise lazy identity setters while refusing native model work.""" import truememory.vector_search as vs import truememory.reranker as rr - monkeypatch.setattr(vs, "set_embedding_model", lambda tier: None) - monkeypatch.setattr(rr, "set_active_tier", lambda tier: None) - monkeypatch.setattr(server, "_set_reranker", lambda name: None) + + def forbidden_model_load(*args, **kwargs): + raise AssertionError("Config-only setup must not construct or probe a model") + + monkeypatch.setattr(vs, "get_model", forbidden_model_load) + monkeypatch.setattr(vs, "_load_frozen_embedding_target", forbidden_model_load) + monkeypatch.setattr(rr, "get_reranker", forbidden_model_load) + monkeypatch.setattr(server, "_set_reranker", forbidden_model_load) def test_config_only_switch_persists_tier(server, monkeypatch): diff --git a/tests/test_issue_396_memory_intensity.py b/tests/test_issue_396_memory_intensity.py index 5424e0ff..eb3341ad 100644 --- a/tests/test_issue_396_memory_intensity.py +++ b/tests/test_issue_396_memory_intensity.py @@ -35,6 +35,7 @@ def server(monkeypatch, tmp_path): import truememory.mcp_server as ms monkeypatch.setattr(ms, "_TRUEMEMORY_DIR", home / ".truememory") monkeypatch.setattr(ms, "_CONFIG_PATH", home / ".truememory" / "config.json") + monkeypatch.setattr(ms, "_CONFIG_LOCK_PATH", ms._CONFIG_PATH.with_name("config.json.lock")) monkeypatch.setattr(ms, "_DB_PATH", str(db_path)) monkeypatch.setattr(ms, "_memory", None) # Invalidate config cache so each test starts clean @@ -56,9 +57,15 @@ def _no_op_model(server, monkeypatch): """Stub out model/reranker side effects.""" import truememory.vector_search as vs import truememory.reranker as rr - monkeypatch.setattr(vs, "set_embedding_model", lambda tier: None) - monkeypatch.setattr(rr, "set_active_tier", lambda tier: None) - monkeypatch.setattr(server, "_set_reranker", lambda name: None) + def forbidden_model_load(*args, **kwargs): + raise AssertionError("Configuration must preserve lazy model slots") + + monkeypatch.setattr(vs, "_model", None) + monkeypatch.setattr(rr, "_model", None) + monkeypatch.setattr(vs, "get_model", forbidden_model_load) + monkeypatch.setattr(vs, "_load_frozen_embedding_target", forbidden_model_load) + monkeypatch.setattr(rr, "get_reranker", forbidden_model_load) + monkeypatch.setattr(server, "_set_reranker", forbidden_model_load) # --------------------------------------------------------------------------- diff --git a/tests/test_issue_497_edge_persist.py b/tests/test_issue_497_edge_persist.py index d7cb5f8c..ce0361f2 100644 --- a/tests/test_issue_497_edge_persist.py +++ b/tests/test_issue_497_edge_persist.py @@ -37,6 +37,7 @@ def server(monkeypatch, tmp_path): monkeypatch.setattr(ms, "_TRUEMEMORY_DIR", home / ".truememory") monkeypatch.setattr(ms, "_CONFIG_PATH", home / ".truememory" / "config.json") + monkeypatch.setattr(ms, "_CONFIG_LOCK_PATH", ms._CONFIG_PATH.with_name("config.json.lock")) monkeypatch.setattr(ms, "_DB_PATH", str(db_path)) monkeypatch.setattr(ms, "_memory", None) yield ms @@ -56,9 +57,15 @@ def _no_op_model(server, monkeypatch): import truememory.vector_search as vs import truememory.reranker as rr - monkeypatch.setattr(vs, "set_embedding_model", lambda tier: None) - monkeypatch.setattr(rr, "set_active_tier", lambda tier: None) - monkeypatch.setattr(server, "_set_reranker", lambda name: None) + def forbidden_model_load(*args, **kwargs): + raise AssertionError("Configuration must preserve lazy model slots") + + monkeypatch.setattr(vs, "_model", None) + monkeypatch.setattr(rr, "_model", None) + monkeypatch.setattr(vs, "get_model", forbidden_model_load) + monkeypatch.setattr(vs, "_load_frozen_embedding_target", forbidden_model_load) + monkeypatch.setattr(rr, "get_reranker", forbidden_model_load) + monkeypatch.setattr(server, "_set_reranker", forbidden_model_load) def test_issue_497_edge_tier_persists(server, monkeypatch): diff --git a/tests/test_issue_502_config.py b/tests/test_issue_502_config.py index 5bb1fa53..cc36c27c 100644 --- a/tests/test_issue_502_config.py +++ b/tests/test_issue_502_config.py @@ -36,6 +36,7 @@ def server(monkeypatch, tmp_path): monkeypatch.setattr(ms, "_TRUEMEMORY_DIR", home / ".truememory") monkeypatch.setattr(ms, "_CONFIG_PATH", home / ".truememory" / "config.json") + monkeypatch.setattr(ms, "_CONFIG_LOCK_PATH", ms._CONFIG_PATH.with_name("config.json.lock")) monkeypatch.setattr(ms, "_DB_PATH", str(db_path)) monkeypatch.setattr(ms, "_memory", None) # Reset config cache so each test starts clean @@ -59,9 +60,15 @@ def _no_op_model(server, monkeypatch): import truememory.vector_search as vs import truememory.reranker as rr - monkeypatch.setattr(vs, "set_embedding_model", lambda tier: None) - monkeypatch.setattr(rr, "set_active_tier", lambda tier: None) - monkeypatch.setattr(server, "_set_reranker", lambda name: None) + def forbidden_model_load(*args, **kwargs): + raise AssertionError("Configuration must preserve lazy model slots") + + monkeypatch.setattr(vs, "_model", None) + monkeypatch.setattr(rr, "_model", None) + monkeypatch.setattr(vs, "get_model", forbidden_model_load) + monkeypatch.setattr(vs, "_load_frozen_embedding_target", forbidden_model_load) + monkeypatch.setattr(rr, "get_reranker", forbidden_model_load) + monkeypatch.setattr(server, "_set_reranker", forbidden_model_load) # ----------------------------------------------------------------------- diff --git a/tests/test_issue_600_antigravity_groq_deepsearch.py b/tests/test_issue_600_antigravity_groq_deepsearch.py index 62881d8c..b8688c93 100644 --- a/tests/test_issue_600_antigravity_groq_deepsearch.py +++ b/tests/test_issue_600_antigravity_groq_deepsearch.py @@ -201,7 +201,9 @@ def test_groq_configure_accepted(tmp_path, monkeypatch): config_path = tmp_path / "config.json" config_path.write_text("{}") monkeypatch.setattr(mcp_server, "_CONFIG_PATH", config_path) + monkeypatch.setattr(mcp_server, "_CONFIG_LOCK_PATH", mcp_server._CONFIG_PATH.with_name("config.json.lock")) monkeypatch.setattr(mcp_server, "_TRUEMEMORY_DIR", tmp_path) + monkeypatch.setattr(mcp_server, "_DB_PATH", str(tmp_path / "synthetic-configure.db")) # Invalidate config cache monkeypatch.setattr(mcp_server, "_config_cache", None) monkeypatch.setattr(mcp_server, "_config_cache_time", 0.0) @@ -224,6 +226,7 @@ def test_deepsearch_falls_back_when_not_configured(tmp_path, monkeypatch): config_path = tmp_path / "config.json" config_path.write_text(json.dumps({"tier": "pro"})) monkeypatch.setattr(mcp_server, "_CONFIG_PATH", config_path) + monkeypatch.setattr(mcp_server, "_CONFIG_LOCK_PATH", mcp_server._CONFIG_PATH.with_name("config.json.lock")) monkeypatch.setattr(mcp_server, "_config_cache", None) monkeypatch.setattr(mcp_server, "_config_cache_time", 0.0) @@ -243,6 +246,7 @@ def test_deepsearch_uses_custom_provider_when_configured(tmp_path, monkeypatch): "groq_api_key": "gsk_test_deepsearch", })) monkeypatch.setattr(mcp_server, "_CONFIG_PATH", config_path) + monkeypatch.setattr(mcp_server, "_CONFIG_LOCK_PATH", mcp_server._CONFIG_PATH.with_name("config.json.lock")) monkeypatch.setattr(mcp_server, "_config_cache", None) monkeypatch.setattr(mcp_server, "_config_cache_time", 0.0) @@ -261,6 +265,7 @@ def test_deepsearch_returns_none_for_unknown_provider(tmp_path, monkeypatch): "deepsearch_provider": "nonexistent_provider", })) monkeypatch.setattr(mcp_server, "_CONFIG_PATH", config_path) + monkeypatch.setattr(mcp_server, "_CONFIG_LOCK_PATH", mcp_server._CONFIG_PATH.with_name("config.json.lock")) monkeypatch.setattr(mcp_server, "_config_cache", None) monkeypatch.setattr(mcp_server, "_config_cache_time", 0.0) @@ -279,6 +284,7 @@ def test_deepsearch_returns_none_when_no_api_key(tmp_path, monkeypatch): # No groq_api_key set })) monkeypatch.setattr(mcp_server, "_CONFIG_PATH", config_path) + monkeypatch.setattr(mcp_server, "_CONFIG_LOCK_PATH", mcp_server._CONFIG_PATH.with_name("config.json.lock")) monkeypatch.setattr(mcp_server, "_config_cache", None) monkeypatch.setattr(mcp_server, "_config_cache_time", 0.0) monkeypatch.delenv("GROQ_API_KEY", raising=False) @@ -298,6 +304,7 @@ def test_deepsearch_resolve_uses_override(tmp_path, monkeypatch): "openai_api_key": "sk-test-deepsearch", })) monkeypatch.setattr(mcp_server, "_CONFIG_PATH", config_path) + monkeypatch.setattr(mcp_server, "_CONFIG_LOCK_PATH", mcp_server._CONFIG_PATH.with_name("config.json.lock")) monkeypatch.setattr(mcp_server, "_config_cache", None) monkeypatch.setattr(mcp_server, "_config_cache_time", 0.0) # Reset deepsearch cache @@ -315,6 +322,7 @@ def test_deepsearch_resolve_falls_back_to_default_on_pro(tmp_path, monkeypatch): config_path = tmp_path / "config.json" config_path.write_text(json.dumps({"tier": "pro"})) monkeypatch.setattr(mcp_server, "_CONFIG_PATH", config_path) + monkeypatch.setattr(mcp_server, "_CONFIG_LOCK_PATH", mcp_server._CONFIG_PATH.with_name("config.json.lock")) monkeypatch.setattr(mcp_server, "_config_cache", None) monkeypatch.setattr(mcp_server, "_config_cache_time", 0.0) # Reset caches @@ -334,6 +342,7 @@ def test_deepsearch_resolve_returns_none_on_edge_no_override(tmp_path, monkeypat config_path = tmp_path / "config.json" config_path.write_text(json.dumps({"tier": "edge"})) monkeypatch.setattr(mcp_server, "_CONFIG_PATH", config_path) + monkeypatch.setattr(mcp_server, "_CONFIG_LOCK_PATH", mcp_server._CONFIG_PATH.with_name("config.json.lock")) monkeypatch.setattr(mcp_server, "_config_cache", None) monkeypatch.setattr(mcp_server, "_config_cache_time", 0.0) # Reset caches diff --git a/tests/test_issue_648_vec_compat.py b/tests/test_issue_648_vec_compat.py index ee09b694..f7df0748 100644 --- a/tests/test_issue_648_vec_compat.py +++ b/tests/test_issue_648_vec_compat.py @@ -224,6 +224,8 @@ def __getattr__(self, name): ) return getattr(self._real, name) + from truememory.storage import create_db + create_db(db).close() real = sqlite3.connect(str(db)) with patch("sqlite3.connect", return_value=_NoExtConn(real)): with pytest.raises(TierSwitchUnsupportedError) as exc: @@ -236,13 +238,22 @@ def __getattr__(self, name): class TestM51StartRebuildNotBricked: """A pre-thread failure in start_rebuild must not brick later rebuilds.""" - def test_pre_thread_failure_does_not_brick(self): + def test_pre_thread_failure_does_not_brick(self, tmp_path, monkeypatch): from truememory.tier_switch.manager import ( RebuildManager, TierSwitchUnsupportedError, ) mgr = RebuildManager() + from truememory.storage import create_db + from truememory.tier_switch import manager as manager_module + db_path = tmp_path / "synthetic-pre-thread.db" + conn = create_db(db_path) + conn.execute("INSERT INTO messages(content, timestamp) VALUES ('synthetic source', '2026-01-01')") + conn.commit() + conn.close() + monkeypatch.setattr(manager_module, "_TRUEMEMORY_DIR", tmp_path) + monkeypatch.setattr(manager_module, "_get_db_path", lambda: db_path) # First call: _open_db fails before the worker thread starts (M-23). with patch( diff --git a/tests/test_l4_entity_sheets_disabled.py b/tests/test_l4_entity_sheets_disabled.py index 425269da..ac500e7a 100644 --- a/tests/test_l4_entity_sheets_disabled.py +++ b/tests/test_l4_entity_sheets_disabled.py @@ -283,7 +283,10 @@ def test_migration_idempotent_across_opens(tmp_path, monkeypatch, caplog): # First open: should purge. with caplog.at_level(logging.INFO): - eng1 = TrueMemoryEngine(db_path).open(rebuild_vectors=False) + eng1 = TrueMemoryEngine(db_path) + eng1._has_consolidation = False + eng1._has_style_vec = False + eng1.open(rebuild_vectors=False) eng1.close() first_text = caplog.text assert "purged 4" in first_text @@ -302,7 +305,10 @@ def test_migration_idempotent_across_opens(tmp_path, monkeypatch, caplog): # Second open: must NOT re-emit the purge log. caplog.clear() with caplog.at_level(logging.INFO): - eng2 = TrueMemoryEngine(db_path).open(rebuild_vectors=False) + eng2 = TrueMemoryEngine(db_path) + eng2._has_consolidation = False + eng2._has_style_vec = False + eng2.open(rebuild_vectors=False) eng2.close() assert "MEMORIST-L4 migration: purged" not in caplog.text, ( f"Idempotent open should skip the migration log; got: {caplog.text!r}" diff --git a/tests/test_l5_surprise_rerank.py b/tests/test_l5_surprise_rerank.py index 65fb28e6..5635ab1a 100644 --- a/tests/test_l5_surprise_rerank.py +++ b/tests/test_l5_surprise_rerank.py @@ -328,7 +328,7 @@ def spy(results): ) -def test_alpha_zero_byte_identical_through_pipeline(tmp_path, monkeypatch): +def test_alpha_zero_byte_identical_through_pipeline(tmp_path, monkeypatch, request): """Both explicit α=0 (constructor) and α=0 (env var) produce identical no-op results through _apply_surprise_boost.""" from truememory.engine import TrueMemoryEngine @@ -342,15 +342,23 @@ def test_alpha_zero_byte_identical_through_pipeline(tmp_path, monkeypatch): monkeypatch.delenv("TRUEMEMORY_ALPHA_SURPRISE", raising=False) eng = TrueMemoryEngine(db_path, alpha_surprise=0.0) + request.addfinalizer(eng.close) + eng._has_consolidation = False + eng._has_style_vec = False eng.open(rebuild_vectors=False) r_ctor = eng._apply_surprise_boost([dict(r) for r in sample]) assert [(r["id"], r["score"]) for r in r_ctor] == expected monkeypatch.setenv("TRUEMEMORY_ALPHA_SURPRISE", "0") eng2 = TrueMemoryEngine(db_path) + request.addfinalizer(eng2.close) + eng2._has_consolidation = False + eng2._has_style_vec = False eng2.open(rebuild_vectors=False) r_env = eng2._apply_surprise_boost([dict(r) for r in sample]) assert [(r["id"], r["score"]) for r in r_env] == expected + eng2.close() + eng.close() def test_composite_source_refined_is_excluded(engine_with_surprise): diff --git a/tests/test_mcp_safety.py b/tests/test_mcp_safety.py index 46d57b96..16b040bf 100644 --- a/tests/test_mcp_safety.py +++ b/tests/test_mcp_safety.py @@ -152,6 +152,7 @@ def test_configure_restores_hf_offline_on_success(monkeypatch, tmp_path): monkeypatch.setenv("TRUEMEMORY_DB", str(db_path)) monkeypatch.setattr(ms, "_TRUEMEMORY_DIR", home / ".truememory") monkeypatch.setattr(ms, "_CONFIG_PATH", home / ".truememory" / "config.json") + monkeypatch.setattr(ms, "_CONFIG_LOCK_PATH", ms._CONFIG_PATH.with_name("config.json.lock")) monkeypatch.setattr(ms, "_DB_PATH", str(db_path)) monkeypatch.setattr(ms, "_memory", None) @@ -198,6 +199,7 @@ def test_configure_restores_hf_offline_on_set_embedding_model_raise( monkeypatch.setenv("TRUEMEMORY_DB", str(db_path)) monkeypatch.setattr(ms, "_TRUEMEMORY_DIR", home / ".truememory") monkeypatch.setattr(ms, "_CONFIG_PATH", home / ".truememory" / "config.json") + monkeypatch.setattr(ms, "_CONFIG_LOCK_PATH", ms._CONFIG_PATH.with_name("config.json.lock")) monkeypatch.setattr(ms, "_DB_PATH", str(db_path)) monkeypatch.setattr(ms, "_memory", None) @@ -214,20 +216,30 @@ def fake_import(name, *args, **kwargs): import truememory.vector_search as vs + attempted_tiers = [] + def _boom(tier): + attempted_tiers.append(tier) raise ValueError("simulated: model removed") + import truememory.reranker as rr + monkeypatch.setattr(vs, "_model", None) + monkeypatch.setattr(rr, "_model", None) monkeypatch.setattr(vs, "set_embedding_model", _boom) monkeypatch.delenv("HF_HUB_OFFLINE", raising=False) monkeypatch.delenv("TRANSFORMERS_OFFLINE", raising=False) - # truememory_configure wraps the tier switch; ValueError propagates - # out per Python semantics after finally runs. + # Failed lazy projection is returned as a truthful incomplete activation. import os as _os - import pytest - with pytest.raises(ValueError): - ms.truememory_configure(tier="pro") + result = json.loads(ms.truememory_configure(tier="pro")) + assert attempted_tiers == ["pro", "edge"] + assert result["status"] == "activation_pending" + assert result["rebuild_error"] == "ValueError: simulated: model removed" + assert result["served_tier"] is None + assert result["description"] == "Serving state awaits authoritative readback." + assert "warning" in result + assert not ms._CONFIG_PATH.exists() or json.loads(ms._CONFIG_PATH.read_text(encoding="utf-8")).get("tier", "edge") == "edge" # Critical F31 assertion: offline mode was RESTORED despite the raise assert _os.environ.get("HF_HUB_OFFLINE") == "1", ( diff --git a/tests/test_reranker_tier.py b/tests/test_reranker_tier.py index 18d351cb..4b2827c6 100644 --- a/tests/test_reranker_tier.py +++ b/tests/test_reranker_tier.py @@ -149,32 +149,57 @@ def test_get_reranker_default_routes_via_tier_resolution(): ) -def test_truememory_configure_propagates_tier_to_reranker_module(): - """Change B.4 regression lock. - - `truememory_configure` must call `reranker.set_active_tier(tier)` AND - pre-load the tier's reranker via `_set_reranker(_current_reranker())`. - Without the first call, changing tier via MCP leaves the reranker module - stuck on the old tier. Without the second, the first post-configure - search pays a cold-start that defeats the whole "configure is the right - moment to warm the new model" design. - - Guarded by source inspection for the same reason as the test above — - calling `truememory_configure` in a unit test would require mocking - sentence_transformers, the Memory singleton, and the re-embed path. Too - much surface for a contract check. - """ - import inspect - - from truememory import mcp_server - - src = inspect.getsource(mcp_server.truememory_configure) - assert "_set_active_tier(tier)" in src, ( - "truememory_configure must call set_active_tier(tier) (imported as " - "_set_active_tier) after saving the new tier to config, so subsequent " - "get_reranker(model_name=None) calls resolve to the new tier's reranker." - ) - assert "_set_reranker(_current_reranker())" in src, ( - "truememory_configure must pre-load the tier's reranker so the first " - "post-configure search does not pay a cold-start (~70ms-3s)." - ) +@pytest.mark.parametrize("tier,expected", [("edge", EDGE_RERANKER), ("base", GTE_RERANKER), ("pro", GTE_RERANKER)]) +def test_truememory_configure_propagates_tier_to_reranker_module(tier, expected, tmp_path, monkeypatch): + """Guarded configuration persists the correct lazy identity without model work.""" + import json + from truememory import mcp_server, model_client, vector_search + from truememory.tier_switch.manager import RebuildManager + + config = tmp_path / "config.json" + monkeypatch.setattr(mcp_server, "_TRUEMEMORY_DIR", tmp_path) + monkeypatch.setattr(mcp_server, "_CONFIG_PATH", config) + monkeypatch.setattr(mcp_server, "_CONFIG_LOCK_PATH", tmp_path / "config.json.lock") + monkeypatch.setattr(mcp_server, "_DB_PATH", str(tmp_path / "synthetic-memories.db")) + monkeypatch.setattr(mcp_server, "_config_cache", None) + monkeypatch.setattr(mcp_server, "_memory", None) + monkeypatch.setattr(RebuildManager, "_instance", RebuildManager()) + monkeypatch.setenv("TRUEMEMORY_EMBED_MODEL", "edge") + monkeypatch.setattr(vector_search, "EMBEDDING_MODEL", "model2vec") + monkeypatch.setattr(vector_search, "_embedding_dim", 256) + monkeypatch.setattr(vector_search, "_model", None) + monkeypatch.setattr(vector_search, "_frozen_embedding_target", None) + monkeypatch.setattr(vector_search, "_runtime_policy_tier", None) + monkeypatch.setattr(reranker, "_model", None) + monkeypatch.setattr(reranker, "_model_name", GTE_RERANKER if expected == EDGE_RERANKER else EDGE_RERANKER) + monkeypatch.setattr(reranker, "_active_tier", "edge") + monkeypatch.setattr(reranker, "_frozen_reranker_id", None) + monkeypatch.setattr(reranker, "_model_certified", False) + calls = [] + embedding_attempts = [] + proxy = object() + + def forbidden_embedding(*args, **kwargs): + embedding_attempts.append((args, kwargs)) + raise AssertionError("Configuration cannot construct or encode embeddings") + + def synthetic_proxy(*, model_name): + calls.append(model_name) + return proxy + + monkeypatch.setattr(vector_search, "get_model", forbidden_embedding) + monkeypatch.setattr(vector_search, "_load_frozen_embedding_target", forbidden_embedding) + monkeypatch.setattr(model_client, "use_model_server", lambda: True) + monkeypatch.setattr(model_client, "get_reranker_proxy", synthetic_proxy) + result = json.loads(mcp_server.truememory_configure(tier=tier)) + assert result["status"] == "configured" + assert result["served_tier"] == tier + assert json.loads(config.read_text(encoding="utf-8"))["tier"] == tier + assert reranker._active_tier == tier + assert get_current_reranker_name() == expected + assert vector_search._model is None + assert reranker._model is None + assert calls == [] + assert embedding_attempts == [] + assert reranker.get_reranker() is proxy + assert calls == [expected] diff --git a/truememory/engine.py b/truememory/engine.py index 685eb162..5a66eb62 100644 --- a/truememory/engine.py +++ b/truememory/engine.py @@ -34,6 +34,7 @@ from typing import TYPE_CHECKING if TYPE_CHECKING: + from collections.abc import Callable from truememory.maintenance import LayerResult, MaintenanceCoordinator # ─────────────────────────────────────────────────────────────────────────── @@ -47,7 +48,7 @@ from truememory.fts_search import search_fts from truememory._platform import _env_int -from truememory.tier_switch.runtime import engine_operation +from truememory.tier_switch.runtime import engine_handle_operation, engine_operation logger = logging.getLogger(__name__) @@ -259,6 +260,44 @@ def get_vectors_load_error() -> str | None: # Helper # ─────────────────────────────────────────────────────────────────────────── +def _validate_add_content(function: "Callable") -> "Callable": + """Validate pure store arguments before database or runtime admission.""" + from functools import wraps + + @wraps(function) + def validated(engine: object, content: str, *args: object, **kwargs: object) -> dict: + if not isinstance(content, str): + raise TypeError(f"content must be a string, got {type(content).__name__}") + + # M-60: enforce the store size cap at the engine level so ALL entry + # points (client.add, hooks, direct Engine.add) inherit it — not just + # mcp_server. Unbounded content means unbounded embed latency and + # poison-large rows. Mirror the mcp_server limit. + if len(content) > MAX_CONTENT_LENGTH: + raise ValueError( + f"Content too large ({len(content)} chars). " + f"Maximum is {MAX_CONTENT_LENGTH}." + ) + + return function(engine, content, *args, **kwargs) + return validated + + +def _validate_delete_user(function: "Callable") -> "Callable": + """Reject invalid deletion scope without opening or initializing storage.""" + from functools import wraps + + @wraps(function) + def validated(engine: object, user_id: str | None = None, *args: object, **kwargs: object) -> bool: + if user_id is not None and not isinstance(user_id, str): + raise TypeError(f"user_id must be a string or None, got {type(user_id).__name__}") + if isinstance(user_id, str) and not user_id.strip(): + raise ValueError("user_id cannot be an empty string") + + return function(engine, user_id, *args, **kwargs) + return validated + + def _has_personality_intent(query: str) -> bool: """ Return True if the query is genuinely asking about personality, preferences, @@ -346,6 +385,7 @@ def __init__(self, db_path: str | Path = None, self._runtime_initialized = False self._runtime_vector_generation = None self._runtime_vector_connection = None + self._reconnect_receipt = None # L5 surprise rerank boost coefficient. None = resolve from env # var / default at call-time via _get_alpha_surprise(). @@ -380,57 +420,282 @@ def __init__(self, db_path: str | Path = None, def _open_connection_handle(self) -> None: """Open a plain handle before serving admission or vector initialization.""" - from truememory.tier_switch.runtime import _connection_read + from truememory.tier_switch.runtime import _connection_read, _open_serving_connection_at_path + sqlite_ioerr = getattr(sqlite3, "SQLITE_IOERR", 10) + + def replacement_reason() -> str | None: + if self.conn is None: + return "missing" + try: + self.conn.execute("PRAGMA schema_version") + return None + except (sqlite3.ProgrammingError, sqlite3.OperationalError) as exc: + if (isinstance(exc, sqlite3.ProgrammingError) + and str(exc) == "Cannot operate on a closed database."): + return "closed" + 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 == sqlite_ioerr))) + if not io_error or self.conn.in_transaction: + raise + return "io_error" + 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: - try: - self.conn.execute("PRAGMA schema_version") + transaction_owner=self._write_lock, allow_closed=True): + with _connection_read(self.conn, self._write_lock, None, None, allow_closed=True): + reason = replacement_reason() + if reason is None: 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) + path = self.db_path if str(self.db_path) == ":memory:" else self.db_path.resolve() + if reason != "missing" and getattr(self, "_reconnect_receipt", None) is not None: + replacement = self._reopen_initialized_connection(path, receipt=self._reconnect_receipt) + if replacement is not None: + candidate, controlled = replacement + installed = False + try: + with _connection_read(self.conn, self._write_lock, None, None, allow_closed=True): + reason = replacement_reason() + if reason is None: + return + if reason == "io_error": + self.conn.close() + self.conn = candidate + installed = True + self._runtime_vector_generation = None + self._runtime_vector_connection = None + if controlled: + self.ready = self._runtime_initialized = self._has_vectors = False + self._reconnect_receipt = None + return + finally: + if not installed: + candidate.close() + if reason == "missing" and str(path) != ":memory:": + from truememory.maintenance import _active_initialization_receipt + with _active_initialization_receipt(path) as observed: + offered = observed + if offered is not None: + replacement = self._reopen_initialized_connection(path, receipt=offered[2]) + if replacement is not None: + candidate, controlled = replacement + installed = False + try: + with _connection_read(self.conn, self._write_lock, None, None, allow_closed=True): + reason = replacement_reason() + if reason is None: + return + if reason == "missing": + with _active_initialization_receipt(path, expected=offered[:2]) as current: + if current is not None: + self.conn = candidate + self._has_vectors = not controlled and offered[2][7] + self._has_hybrid = _HAS_HYBRID and self._has_vectors + self.ready = self._runtime_initialized = not controlled + self._runtime_vector_generation = None + self._runtime_vector_connection = None + self._reconnect_receipt = None if controlled else offered[2] + installed = True + if installed: + return + finally: + if not installed: + candidate.close() + from truememory.maintenance import maintenance_owner + if str(path) != ":memory:": + path.parent.mkdir(parents=True, exist_ok=True) try: - self.db_path.parent.chmod(0o700) + 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 + # Keep same-Engine openers serialized, but never claim maintenance + # ownership while holding the connection writer. + with maintenance_owner(None if str(path) == ":memory:" else path), \ + _connection_read(self.conn, self._write_lock, None, None, allow_closed=True): + reason = replacement_reason() + if reason is None: + return + if reason == "io_error": + 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 + self._reconnect_receipt = None + self.conn = _open_serving_connection_at_path(path, create_db) + + @staticmethod + def _reconnect_registry(conn: sqlite3.Connection) -> tuple | None: + """Bound the legacy registry identity, including mutations without DDL.""" + if not conn.execute( + "SELECT 1 FROM main.sqlite_master WHERE name='vector_cache_registry' COLLATE NOCASE" + ).fetchone(): + return () + if conn.execute( + "SELECT 1 FROM main.vector_cache_registry WHERE " + "typeof(tier_group)!='text' OR length(CAST(tier_group AS BLOB))>16 OR " + "typeof(vec_table)!='text' OR length(CAST(vec_table AS BLOB))>64 OR " + "typeof(sep_table)!='text' OR length(CAST(sep_table AS BLOB))>64 OR " + "(model_name IS NOT NULL AND (typeof(model_name)!='text' OR length(CAST(model_name AS BLOB))>512)) OR " + "typeof(embedding_dim) NOT IN ('integer','real','null') OR " + "typeof(last_embedded_id) NOT IN ('integer','real','null') OR " + "typeof(vector_count) NOT IN ('integer','real','null') OR " + "typeof(last_updated) NOT IN ('integer','real','null') OR " + "typeof(created) NOT IN ('integer','real','null') LIMIT 1" + ).fetchone(): + return None + rows = tuple(tuple(row) for row in conn.execute( + "SELECT tier_group,vec_table,sep_table,model_name,embedding_dim," + "last_embedded_id,vector_count,last_updated,created " + "FROM main.vector_cache_registry ORDER BY tier_group LIMIT 4" + ).fetchall()) + if len(rows) > 3 or any( + row[1] not in _ALL_VEC_TABLES or row[2] not in _ALL_VEC_TABLES + or any(value is not None and type(value) not in (int, float) for value in row[4:]) + for row in rows): + return None + return rows + + def _capture_reconnect_receipt(self, operation: object) -> None: + """Remember only this Engine's completed legacy initialization, never a model.""" + self._reconnect_receipt = None + if self.conn.in_transaction or (not self._has_vectors and _HAS_VECTOR): + return + from truememory.tier_switch.runtime import _legacy_key, TierRuntimeError + from truememory import vector_search + key = operation.key + if (key != _legacy_key() or key[0] != "legacy" + or any(type(value) is not str or len(value) > 512 for value in key)): + return + path = self.db_path.resolve() + if operation._database != ("file", str(path)): + return + receipt = None + try: + self.conn.execute("BEGIN") + stat = path.stat() + registry = self._reconnect_registry(self.conn) + if registry is None: + return + group = vector_search._active_tier_group() + tables = next((row[1:3] for row in registry if row[0] == group), + ("vec_messages", "vec_messages_sep")) + schema = self.conn.execute("PRAGMA schema_version").fetchone()[0] + if not self.conn.execute( + "SELECT 1 FROM main.metadata WHERE key='l4_entity_profile_migration_done' AND value='1'" + ).fetchone(): + return + receipt = (str(path), stat.st_dev, stat.st_ino, schema, + key, tables, registry, self._has_vectors, _HAS_VECTOR, + self.conn.execute("PRAGMA journal_mode").fetchone()[0]) + except (OSError, sqlite3.Error): + logger.debug("Connection initialization cannot be reused", exc_info=True) + finally: + if self.conn.in_transaction: + self.conn.rollback() + if self.conn.in_transaction: + raise TierRuntimeError("Reconnect receipt requires completed read cleanup") + self._reconnect_receipt = receipt + + def _reopen_initialized_connection(self, path: Path, *, receipt: tuple | None = None) -> tuple[sqlite3.Connection, bool] | None: + """Read-only validation allows same-Engine reconnect without migration ownership.""" + from truememory.storage import DatabaseOpenError, _integrity_message + from truememory.tier_switch.runtime import _legacy_key, _read_selection, _read_policy, TierRuntimeError + if receipt is None: + receipt = self._reconnect_receipt + try: + stat = path.stat() + except OSError: + return None + identity = (str(path), stat.st_dev, stat.st_ino) + if identity != receipt[:3] or _HAS_VECTOR != receipt[8]: + return None + candidate = sqlite3.connect(path.as_uri() + "?mode=rw", uri=True, check_same_thread=False) + retained = False + try: + candidate.execute(f"PRAGMA busy_timeout={DEFAULT_BUSY_TIMEOUT_MS}") + candidate.execute("PRAGMA foreign_keys=ON") + candidate.execute("PRAGMA synchronous=NORMAL") + candidate.execute("PRAGMA cache_size=-64000") + candidate.execute("PRAGMA mmap_size=268435456") + if candidate.execute("PRAGMA journal_mode").fetchone()[0] != receipt[9]: + return None + try: + qc = candidate.execute("PRAGMA quick_check(1)").fetchone() + except sqlite3.DatabaseError as exc: + if "i/o error" in str(exc).lower(): + raise DatabaseOpenError( + "TrueMemory could not open the database due to a disk " + f"I/O error: {path}\n" + "This usually means the SQLite WAL file is " + "inconsistent (e.g. a -wal/-shm file was removed while " + "a connection was still open). Close ALL TrueMemory " + "processes (hooks, MCP server, CLI) and retry." + ) from exc + raise DatabaseOpenError(_integrity_message(path, str(exc))) from exc + if qc is not None and qc[0] != "ok": + raise DatabaseOpenError(_integrity_message(path, str(qc[0]))) + if receipt[7]: + import sqlite_vec + candidate.enable_load_extension(True) + try: + sqlite_vec.load(candidate) + finally: + candidate.enable_load_extension(False) + candidate.execute("BEGIN") + if candidate.execute("PRAGMA schema_version").fetchone()[0] != receipt[3]: + return None + selection = _read_selection(candidate) + policy = _read_policy(candidate) if selection is None else None + controlled = selection is not None or policy is not None + if not controlled: + if _legacy_key() != receipt[4] or self._reconnect_registry(candidate) != receipt[6]: + return None + values = candidate.execute( + "SELECT key, CASE WHEN typeof(value)='text' AND length(CAST(value AS BLOB))<=512 " + "THEN value ELSE NULL END FROM main.metadata " + "WHERE key IN ('embed_model','embed_dim') LIMIT 3" + ).fetchall() + expected = {"embed_model": receipt[4][1], "embed_dim": receipt[4][2]} + if values: + if len(values) != 2 or dict(values) != expected: + return None + elif receipt[7]: + # A first add may publish metadata after initialization. Its + # absence is safe only while both initialized vector tables are empty. + for table in receipt[5]: + if candidate.execute(f'SELECT 1 FROM main."{table}" LIMIT 1').fetchone(): + return None + current = path.stat() + if (str(path), current.st_dev, current.st_ino) != identity: + return None + candidate.rollback() + if candidate.in_transaction: + raise TierRuntimeError("Reconnect candidate requires completed read cleanup") + retained = True + return candidate, controlled + finally: + if not retained: + candidate.close() 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): + from truememory.tier_switch.runtime import engine_serving_operation + with engine_serving_operation(self, _suppress_maintenance=_suppress_maintenance): self._initialize_connection(_suppress_maintenance=_suppress_maintenance) - def _initialize_connection(self, *, _suppress_maintenance: bool = False) -> None: + def _initialize_connection(self, *, _suppress_maintenance: bool = False, + _preparing: 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: + _suppress_maintenance = _suppress_maintenance or operation._defer_maintenance + controlled = operation.selection or operation.policy + if controlled 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): + and getattr(self, "_runtime_vector_generation", None) == controlled.generation): if not _suppress_maintenance: self._maybe_auto_consolidate() return @@ -440,7 +705,7 @@ def _initialize_connection(self, *, _suppress_maintenance: bool = False) -> None _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): + and getattr(self, "_runtime_vector_generation", None) == controlled.generation): if self.conn.in_transaction: raise TierRuntimeError("Selected vector initialization requires a clean connection") if not _HAS_VECTOR: @@ -456,7 +721,7 @@ def _initialize_connection(self, *, _suppress_maintenance: bool = False) -> None 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_generation = controlled.generation self._runtime_vector_connection = self.conn initialized = True if not _suppress_maintenance: @@ -466,9 +731,16 @@ def _initialize_connection(self, *, _suppress_maintenance: bool = False) -> None self._maybe_auto_consolidate() return if getattr(self, "_runtime_initialized", True) or getattr(self, "ready", False): + receipt = getattr(self, "_reconnect_receipt", None) + if receipt is not None and operation.key[1:3] != receipt[4][1:3]: + raise TierRuntimeError("Initialized embedding identity changed before serving admission") if not _suppress_maintenance: self._maybe_auto_consolidate() return + import sys as _sys + defer_legacy_migration = (_preparing and _sys.platform == "darwin" + and operation.tier in ("base", "pro", "qwen3_256") + and str(self.db_path) == ":memory:") with _connection_read(self.conn, self._init_lock, None, None, transaction_owner=self._write_lock): if self._runtime_initialized: @@ -478,6 +750,8 @@ def _initialize_connection(self, *, _suppress_maintenance: bool = False) -> None with _connection_read(self.conn, self._write_lock, None, None): if self.conn.in_transaction: raise TierRuntimeError("Vector initialization requires a clean connection") + from truememory.tier_switch.runtime import _legacy_initialization_is_local + local_only = False # Load sqlite-vec extension global _vectors_load_error if _HAS_VECTOR: @@ -486,6 +760,7 @@ def _initialize_connection(self, *, _suppress_maintenance: bool = False) -> None self.conn.enable_load_extension(True) sqlite_vec.load(self.conn) self.conn.enable_load_extension(False) + local_only = not _preparing and _legacy_initialization_is_local(self.conn, operation) # init_vec_table() runs _check_embedder_compatibility(), # which raises TrueMemoryMigrationError on a dim/model # mismatch. That must NOT be swallowed as a generic @@ -493,12 +768,16 @@ def _initialize_connection(self, *, _suppress_maintenance: bool = False) -> None # 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) + if not local_only: + init_vec_table(self.conn) + try: + from truememory.vector_search import migrate_legacy_vec_tables + if _preparing and not defer_legacy_migration: + migrate_legacy_vec_tables(self.conn) + except Exception as _serving_error: + from truememory.tier_switch.runtime import raise_if_serving_rejection + raise_if_serving_rejection(_serving_error) + 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 @@ -508,6 +787,8 @@ def _initialize_connection(self, *, _suppress_maintenance: bool = False) -> None self._has_vectors = False raise except Exception as exc: + from truememory.tier_switch.runtime import raise_if_serving_rejection + raise_if_serving_rejection(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 @@ -516,8 +797,7 @@ def _initialize_connection(self, *, _suppress_maintenance: bool = False) -> None # 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 (not local_only and _sys.platform == "darwin" and self._has_vectors and current_operation(self.conn).selection is None): try: _embed_model = resolve_tier() @@ -535,46 +815,49 @@ def _initialize_connection(self, *, _suppress_maintenance: bool = False) -> None ).fetchone()[0] if "messages" in _tables else 0 if _msg_count > 0: def _bg_reembed(db_path): - import sqlite3 as _sql + from truememory.tier_switch.runtime import open_serving_connection, legacy_rebuild_operation _conn = None try: - _conn = _sql.connect(str(db_path), check_same_thread=False) - _conn.execute("PRAGMA journal_mode=WAL") - _conn.execute("PRAGMA busy_timeout=%d" % DEFAULT_BUSY_TIMEOUT_MS) - _conn.execute("PRAGMA synchronous=NORMAL") - _conn.execute("PRAGMA foreign_keys=ON") - # Issue #499: load sqlite-vec on the - # background thread's connection so - # vec virtual tables are available. - import sqlite_vec as _sv - _conn.enable_load_extension(True) - _sv.load(_conn) - _conn.enable_load_extension(False) - from truememory.vector_search import ( - build_vectors as _bv, - build_separation_vectors as _bsv, - init_vec_table as _ivt, - ) - _conn.execute("DROP TABLE IF EXISTS vec_messages") - _conn.execute("DROP TABLE IF EXISTS vec_messages_sep") - _conn.commit() - _ivt(_conn) - _bv(_conn) - _bsv(_conn) - # Issue #485: only set the flag AFTER - # successful re-embed so a failed - # thread allows retry on next init. - _conn.execute( - "INSERT OR REPLACE INTO metadata " - "(key, value) VALUES (?, ?)", - ("qwen3_nan_fix_applied", "1"), - ) - _conn.commit() - logger.warning( - "Qwen3 NaN fix: re-embedded %d vectors " - "in background", _msg_count, - ) - except Exception: + _conn = open_serving_connection(db_path, create_db) + with legacy_rebuild_operation(_conn): + _conn.execute("PRAGMA journal_mode=WAL") + _conn.execute("PRAGMA busy_timeout=%d" % DEFAULT_BUSY_TIMEOUT_MS) + _conn.execute("PRAGMA synchronous=NORMAL") + _conn.execute("PRAGMA foreign_keys=ON") + # Issue #499: load sqlite-vec on the + # background thread's connection so + # vec virtual tables are available. + import sqlite_vec as _sv + _conn.enable_load_extension(True) + _sv.load(_conn) + _conn.enable_load_extension(False) + from truememory.vector_search import ( + build_vectors as _bv, + build_separation_vectors as _bsv, + init_vec_table as _ivt, + ) + _conn.execute("DROP TABLE IF EXISTS vec_messages") + _conn.execute("DROP TABLE IF EXISTS vec_messages_sep") + _conn.commit() + _ivt(_conn) + _bv(_conn) + _bsv(_conn) + # Issue #485: only set the flag AFTER + # successful re-embed so a failed + # thread allows retry on next init. + _conn.execute( + "INSERT OR REPLACE INTO metadata " + "(key, value) VALUES (?, ?)", + ("qwen3_nan_fix_applied", "1"), + ) + _conn.commit() + logger.warning( + "Qwen3 NaN fix: re-embedded %d vectors " + "in background", _msg_count, + ) + except Exception as _serving_error: + from truememory.tier_switch.runtime import raise_if_serving_rejection + raise_if_serving_rejection(_serving_error) logger.warning( "Qwen3 NaN background migration failed — " "will retry on next startup", @@ -584,7 +867,9 @@ def _bg_reembed(db_path): if _conn is not None: try: _conn.close() - except Exception: + except Exception as _serving_error: + from truememory.tier_switch.runtime import raise_if_serving_rejection + raise_if_serving_rejection(_serving_error) pass _db = self.db_path if str(_db) != ":memory:": @@ -604,36 +889,42 @@ def _bg_wrapper(db_path): "in background thread", _msg_count, ) else: - from truememory.vector_search import ( - build_vectors as _bv, - build_separation_vectors as _bsv, - init_vec_table as _ivt, - ) - self.conn.execute("DROP TABLE IF EXISTS vec_messages") - self.conn.execute("DROP TABLE IF EXISTS vec_messages_sep") - self.conn.commit() - _ivt(self.conn) - _bv(self.conn) - _bsv(self.conn) + from truememory.tier_switch.runtime import legacy_rebuild_operation + with legacy_rebuild_operation(self.conn): + from truememory.vector_search import ( + build_vectors as _bv, + build_separation_vectors as _bsv, + init_vec_table as _ivt, + ) + self.conn.execute("DROP TABLE IF EXISTS vec_messages") + self.conn.execute("DROP TABLE IF EXISTS vec_messages_sep") + self.conn.commit() + _ivt(self.conn) + _bv(self.conn) + _bsv(self.conn) + self.conn.execute( + "INSERT OR REPLACE INTO metadata " + "(key, value) VALUES (?, ?)", + ("qwen3_nan_fix_applied", "1"), + ) + self.conn.commit() + logger.warning( + "Qwen3 NaN fix: re-embedded all vectors", + ) + else: + # No messages — just set the flag to skip + # future checks. + from truememory.tier_switch.runtime import legacy_vector_mutation + with legacy_vector_mutation(self.conn): self.conn.execute( "INSERT OR REPLACE INTO metadata " "(key, value) VALUES (?, ?)", ("qwen3_nan_fix_applied", "1"), ) self.conn.commit() - logger.warning( - "Qwen3 NaN fix: re-embedded all vectors", - ) - else: - # No messages — just set the flag to skip - # future checks. - self.conn.execute( - "INSERT OR REPLACE INTO metadata " - "(key, value) VALUES (?, ?)", - ("qwen3_nan_fix_applied", "1"), - ) - self.conn.commit() - except Exception: + except Exception as _serving_error: + from truememory.tier_switch.runtime import raise_if_serving_rejection + raise_if_serving_rejection(_serving_error) logger.warning( "Qwen3 NaN migration failed — run " "'truememory-ingest upgrade-tier base --force' to fix manually", @@ -644,9 +935,20 @@ def _bg_wrapper(db_path): # M-84: previously this only ran in the deprecated open() path, so # production (which uses _ensure_connection) never purged them. with _connection_read(self.conn, self._write_lock, None, None): - self._purge_legacy_entity_profile_summaries() + if defer_legacy_migration and self._has_vectors: + try: + from truememory.vector_search import migrate_legacy_vec_tables + migrate_legacy_vec_tables(self.conn) + except Exception as _serving_error: + from truememory.tier_switch.runtime import raise_if_serving_rejection + raise_if_serving_rejection(_serving_error) + logger.debug("Legacy vec table migration skipped", exc_info=True) + if not local_only: + self._purge_legacy_entity_profile_summaries() self.ready = True self._runtime_initialized = True + if getattr(operation, "_database", (None,))[0] == "file": + self._capture_reconnect_receipt(operation) if not _suppress_maintenance: self._maybe_startup_consolidate() @@ -730,6 +1032,7 @@ def _maybe_startup_consolidate(self) -> None: # Production CRUD API # ────────────────────────────────────────────────────────────────────── + @_validate_add_content @engine_operation def add( self, @@ -757,19 +1060,6 @@ def add( Returns: Dict with ``id`` and the stored fields. """ - if not isinstance(content, str): - raise TypeError(f"content must be a string, got {type(content).__name__}") - - # M-60: enforce the store size cap at the engine level so ALL entry - # points (client.add, hooks, direct Engine.add) inherit it — not just - # mcp_server. Unbounded content means unbounded embed latency and - # poison-large rows. Mirror the mcp_server limit. - if len(content) > MAX_CONTENT_LENGTH: - raise ValueError( - f"Content too large ({len(content)} chars). " - f"Maximum is {MAX_CONTENT_LENGTH}." - ) - self._ensure_connection() from truememory.tier_switch.writer import capture_writer_selection, writer_transaction @@ -790,7 +1080,9 @@ def add( pre_sep_embedding = _encode_with_mps_fallback(model, [sep_text])[0] except VectorPublicationChanged: raise - except Exception: + except Exception as _serving_error: + from truememory.tier_switch.runtime import raise_if_serving_rejection + raise_if_serving_rejection(_serving_error) logger.debug("Failed to pre-compute embedding during add()", exc_info=True) pre_style_vec = None @@ -798,7 +1090,9 @@ def add( try: from truememory.personality_style_vec import compute_style_vector pre_style_vec = compute_style_vector(content) - except Exception: + except Exception as _serving_error: + from truememory.tier_switch.runtime import raise_if_serving_rejection + raise_if_serving_rejection(_serving_error) logger.debug("Failed to pre-compute style vector during add()", exc_info=True) from contextlib import nullcontext @@ -833,7 +1127,7 @@ def add( try: from truememory.vector_search import ( serialize_f32, - _write_embedder_metadata_no_commit, + _write_foreground_embedder_metadata_no_commit, ) vec_tbl, sep_tbl = _validate_foreground_vector_target( self.conn, pre_identity, new_id, selection=write_selection, @@ -847,10 +1141,12 @@ def add( f"INSERT INTO {sep_tbl}(rowid, embedding) VALUES (?, ?)", (new_id, serialize_f32(pre_sep_embedding)), ) - _write_embedder_metadata_no_commit(self.conn) + _write_foreground_embedder_metadata_no_commit(self.conn, pre_identity, selection=write_selection) except VectorPublicationChanged: raise - except Exception: + except Exception as _serving_error: + from truememory.tier_switch.runtime import raise_if_serving_rejection + raise_if_serving_rejection(_serving_error) logger.warning("Failed to store embedding for message %s during add()", new_id, exc_info=True) # Incrementally update entity profile. @@ -861,7 +1157,9 @@ def add( try: from truememory.personality import update_entity_profile_incremental update_entity_profile_incremental(self.conn, sender, content, recipient) - except Exception: + except Exception as _serving_error: + from truememory.tier_switch.runtime import raise_if_serving_rejection + raise_if_serving_rejection(_serving_error) logger.debug("Failed to update entity profile for %s during add()", sender, exc_info=True) if self._has_style_vec: @@ -904,6 +1202,10 @@ def _get_maintenance_coordinator(self) -> MaintenanceCoordinator: def _maybe_auto_consolidate(self) -> None: """Probe fixed metadata without waiting on foreground or model work.""" + from truememory.tier_switch.runtime import current_operation + operation = current_operation(self.conn) if self.conn is not None else None + if operation is not None and getattr(operation, "_defer_maintenance", False): + return if (not self._has_consolidation and not getattr(self, "_has_style_vec", False)) or self.conn is None: return if not self._write_lock.acquire(blocking=False): @@ -937,11 +1239,13 @@ def _maybe_auto_consolidate(self) -> None: self._write_lock.release() if (eligible or style_eligible) and coordinator is not None: try: + initialization = ({"_initialization_receipt": self._reconnect_receipt} + if getattr(self, "_reconnect_receipt", None) is not None else {}) if style_eligible: coordinator.request_layers(threshold=self._auto_consolidate_threshold, - include_layers=bool(eligible), include_style=True) + include_layers=bool(eligible), include_style=True, **initialization) else: - coordinator.request_layers(threshold=self._auto_consolidate_threshold) + coordinator.request_layers(threshold=self._auto_consolidate_threshold, **initialization) except Exception as error: self._maintenance_pending_reason = type(error).__name__[:64] @@ -964,6 +1268,7 @@ def delete(self, memory_id: int) -> bool: self._maybe_auto_consolidate() return deleted + @_validate_delete_user @engine_operation def delete_all(self, user_id: str | None = None) -> bool: """Delete all memories, optionally filtered by user. @@ -981,11 +1286,6 @@ def delete_all(self, user_id: str | None = None) -> bool: Returns: True if any rows were deleted from messages. """ - if user_id is not None and not isinstance(user_id, str): - raise TypeError(f"user_id must be a string or None, got {type(user_id).__name__}") - if isinstance(user_id, str) and not user_id.strip(): - 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) @@ -1021,7 +1321,9 @@ def delete_all(self, user_id: str | None = None) -> bool: raise ValueError(f"Invalid column name: {col}") try: _delete_in_chunks(self.conn, table, col, msg_ids) - except Exception: + except Exception as _serving_error: + from truememory.tier_switch.runtime import raise_if_serving_rejection + raise_if_serving_rejection(_serving_error) logger.warning("Failed to clean %s for user %s", table, user_id, exc_info=True) # Clean entity profile for this user (normalize to lowercase #467) @@ -1029,7 +1331,9 @@ def delete_all(self, user_id: str | None = None) -> bool: self.conn.execute( "DELETE FROM entity_profiles WHERE entity = ?", (user_id.lower(),) ) - except Exception: + except Exception as _serving_error: + from truememory.tier_switch.runtime import raise_if_serving_rejection + raise_if_serving_rejection(_serving_error) logger.warning("Failed to clean entity_profiles for user %s", user_id, exc_info=True) # Clean entity style vectors for this user (normalize to lowercase #467) @@ -1037,7 +1341,9 @@ def delete_all(self, user_id: str | None = None) -> bool: self.conn.execute( "DELETE FROM entity_style_vectors WHERE entity = ?", (user_id.lower(),) ) - except Exception: + except Exception as _serving_error: + from truememory.tier_switch.runtime import raise_if_serving_rejection + raise_if_serving_rejection(_serving_error) logger.warning("Failed to clean entity_style_vectors for user %s", user_id, exc_info=True) # Clean entity relationships involving this user (normalize to lowercase #467) @@ -1046,7 +1352,9 @@ def delete_all(self, user_id: str | None = None) -> bool: "DELETE FROM entity_relationships WHERE entity_a = ? OR entity_b = ?", (user_id.lower(), user_id.lower()), ) - except Exception: + except Exception as _serving_error: + from truememory.tier_switch.runtime import raise_if_serving_rejection + raise_if_serving_rejection(_serving_error) logger.warning("Failed to clean entity_relationships for user %s", user_id, exc_info=True) # Clean summaries scoped to this user @@ -1054,13 +1362,17 @@ def delete_all(self, user_id: str | None = None) -> bool: self.conn.execute( "DELETE FROM summaries WHERE entity = ?", (user_id.lower(),) ) - except Exception: + except Exception as _serving_error: + from truememory.tier_switch.runtime import raise_if_serving_rejection + raise_if_serving_rejection(_serving_error) logger.warning("Failed to clean summaries for user %s", user_id, exc_info=True) if episode_ids: try: _delete_in_chunks(self.conn, "episodes", "id", episode_ids) - except Exception: + except Exception as _serving_error: + from truememory.tier_switch.runtime import raise_if_serving_rejection + raise_if_serving_rejection(_serving_error) logger.warning("Failed to clean episodes for user %s", user_id, exc_info=True) # Clean vector tables for deleted message IDs @@ -1070,14 +1382,18 @@ def delete_all(self, user_id: str | None = None) -> bool: else: try: vec_tbl, sep_tbl = _resolve_vec_tables(self.conn) - except Exception: + except Exception as _serving_error: + from truememory.tier_switch.runtime import raise_if_serving_rejection + raise_if_serving_rejection(_serving_error) 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: _delete_in_chunks(self.conn, vec_table, "rowid", msg_ids) - except Exception: + except Exception as _serving_error: + from truememory.tier_switch.runtime import raise_if_serving_rejection + raise_if_serving_rejection(_serving_error) logger.warning("Failed to clean %s for user %s", vec_table, user_id, exc_info=True) # Remove orphaned cluster_centroids (clusters with no @@ -1087,7 +1403,9 @@ def delete_all(self, user_id: str | None = None) -> bool: "DELETE FROM cluster_centroids WHERE cluster_id NOT IN " "(SELECT DISTINCT cluster_id FROM message_clusters)" ) - except Exception: + except Exception as _serving_error: + from truememory.tier_switch.runtime import raise_if_serving_rejection + raise_if_serving_rejection(_serving_error) logger.warning("Failed to clean cluster_centroids for user %s", user_id, exc_info=True) # Delete parent rows AFTER all child FK references are gone @@ -1115,14 +1433,18 @@ def delete_all(self, user_id: str | None = None) -> bool: raise ValueError(f"Invalid table name: {table}") try: self.conn.execute(f"DELETE FROM {table}") - except Exception: + except Exception as _serving_error: + from truememory.tier_switch.runtime import raise_if_serving_rejection + raise_if_serving_rejection(_serving_error) logger.warning("Failed to clear table %s during delete_all", table, exc_info=True) # Clear ALL known vector tables across all tiers. for vec_table in _ALL_VEC_TABLES: try: self.conn.execute(f"DELETE FROM {vec_table}") - except Exception: + except Exception as _serving_error: + from truememory.tier_switch.runtime import raise_if_serving_rejection + raise_if_serving_rejection(_serving_error) logger.debug("Failed to clear %s during delete_all (table may not exist)", vec_table, exc_info=True) cursor = self.conn.execute("DELETE FROM messages") @@ -1131,7 +1453,9 @@ def delete_all(self, user_id: str | None = None) -> bool: # Rebuild FTS index try: self.conn.execute("INSERT INTO messages_fts(messages_fts) VALUES('rebuild')") - except Exception: + except Exception as _serving_error: + from truememory.tier_switch.runtime import raise_if_serving_rejection + raise_if_serving_rejection(_serving_error) logger.warning("Failed to rebuild FTS index during delete_all", exc_info=True) if deleted: @@ -1152,33 +1476,40 @@ def consolidate(self) -> dict[str, str]: MaintenanceBusyError, format_maintenance_report, maintenance_busy_result, maintenance_observation, observe_style, run_engine_maintenance, _style_unready_result, ) - self._ensure_connection(_suppress_maintenance=True) - with self._write_lock: + from truememory.tier_switch.runtime import _read_selection, _read_policy, _connection_read + self._open_connection_handle() + with _connection_read(self.conn, self._write_lock, None, None): + controlled = _read_selection(self.conn) is not None or _read_policy(self.conn) is not None + if controlled: + self._has_style_vec = self.conn.execute( + "SELECT 1 FROM main.sqlite_master WHERE type='table' AND name='entity_style_vectors' COLLATE NOCASE" + ).fetchone() is not None + if not controlled: + self._ensure_connection(_suppress_maintenance=True) + with _connection_read(self.conn, self._write_lock, None, None): coordinator = self._get_maintenance_coordinator() coordinator.refresh_capabilities() borrowed = self.conn.in_transaction - if borrowed or coordinator.path is None: - try: - with maintenance_observation(coordinator, borrowed=borrowed) as observed: - report = run_engine_maintenance(self.conn, coordinator, force=True, - threshold=self._auto_consolidate_threshold, prepare_extensions=False, - allow_caller_transaction=borrowed, include_style=getattr(self, "_has_style_vec", False)) - if not borrowed and self.conn.in_transaction: - self.conn.rollback() - raise RuntimeError("Maintenance work left an unfinished transaction") - observed.append(report) - except MaintenanceBusyError: - return maintenance_busy_result() - self._apply_manual_maintenance_capabilities(report.results) - return format_maintenance_report(report) foreground_style = None - if getattr(self, "_has_style_vec", False): + if not borrowed and coordinator.path is not None and getattr(self, "_has_style_vec", False): observation = observe_style(self.conn, threshold=self._auto_consolidate_threshold) if observation.state is None: foreground_style = _style_unready_result(observation.health) + if borrowed or coordinator.path is None: + try: + with maintenance_observation(coordinator, borrowed=borrowed) as observed: + report = run_engine_maintenance(self.conn, coordinator, force=True, + threshold=self._auto_consolidate_threshold, prepare_extensions=False, + _connection_lock=self._write_lock, allow_caller_transaction=borrowed, include_style=getattr(self, "_has_style_vec", False)) + observed.append(report) + except MaintenanceBusyError: + return maintenance_busy_result() + self._apply_manual_maintenance_capabilities(report.results) + return format_maintenance_report(report) try: with maintenance_observation(coordinator) as observed: - connection = create_db(coordinator.path) + from truememory.tier_switch.runtime import open_serving_connection + connection = open_serving_connection(coordinator.path, create_db) try: report = run_engine_maintenance(connection, coordinator, force=True, threshold=self._auto_consolidate_threshold, @@ -1240,7 +1571,9 @@ def update(self, memory_id: int, content: str | None = None, **fields) -> dict | pre_sep_embedding = _encode_with_mps_fallback(model, [sep_text])[0] except VectorPublicationChanged: raise - except Exception: + except Exception as _serving_error: + from truememory.tier_switch.runtime import raise_if_serving_rejection + raise_if_serving_rejection(_serving_error) logger.debug("Failed to pre-compute embedding during update()", exc_info=True) source_committed = False @@ -1262,13 +1595,15 @@ def update(self, memory_id: int, content: str | None = None, **fields) -> dict | if pre_embedding is not None: from truememory.vector_search import VectorPublicationChanged, _foreground_vector_publication try: - from truememory.vector_search import serialize_f32, _write_embedder_metadata_no_commit + from truememory.vector_search import serialize_f32, _write_foreground_embedder_metadata_no_commit 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: + except Exception as _serving_error: + from truememory.tier_switch.runtime import raise_if_serving_rejection + raise_if_serving_rejection(_serving_error) logger.debug("Failed to delete old vector embedding for message %d", memory_id, exc_info=True) self.conn.execute( f"INSERT INTO {vec_tbl}(rowid, embedding) VALUES (?, ?)", @@ -1277,16 +1612,20 @@ def update(self, memory_id: int, content: str | None = None, **fields) -> dict | if pre_sep_embedding is not None: try: self.conn.execute(f"DELETE FROM {sep_tbl} WHERE rowid = ?", (memory_id,)) - except Exception: + except Exception as _serving_error: + from truememory.tier_switch.runtime import raise_if_serving_rejection + raise_if_serving_rejection(_serving_error) logger.debug("Failed to delete old sep vector for message %d", memory_id, exc_info=True) self.conn.execute( f"INSERT INTO {sep_tbl}(rowid, embedding) VALUES (?, ?)", (memory_id, serialize_f32(pre_sep_embedding)), ) - _write_embedder_metadata_no_commit(self.conn) + _write_foreground_embedder_metadata_no_commit(self.conn, pre_identity, selection=write_selection) except VectorPublicationChanged as exc: raise VectorPublicationChanged("Source update committed; retry update to publish vectors") from exc - except Exception: + except Exception as _serving_error: + from truememory.tier_switch.runtime import raise_if_serving_rejection + raise_if_serving_rejection(_serving_error) logger.warning("Vector embedding failed for message %d", memory_id, exc_info=True) finally: if source_committed: @@ -1349,20 +1688,33 @@ def open(self, rebuild_vectors: bool = True) -> "TrueMemoryEngine": if not self.db_path.exists(): raise FileNotFoundError(f"Database not found: {self.db_path}") - from truememory.storage import _validate_db_path - self.conn = sqlite3.connect( - _validate_db_path(self.db_path), check_same_thread=False - ) - self.conn.row_factory = None # Use default tuple rows - self.conn.execute("PRAGMA journal_mode=WAL") - self.conn.execute(f"PRAGMA busy_timeout={DEFAULT_BUSY_TIMEOUT_MS}") - self.conn.execute("PRAGMA foreign_keys=ON") - self.conn.execute("PRAGMA synchronous=NORMAL") - self.conn.execute("PRAGMA cache_size=-64000") - self.conn.execute("PRAGMA mmap_size=268435456") + from truememory.tier_switch.runtime import open_serving_connection + + def open_existing(path: Path) -> sqlite3.Connection: + from truememory.storage import _validate_db_path + uri = Path(_validate_db_path(path)).as_uri() + "?mode=rw" + candidate = sqlite3.connect(uri, uri=True, check_same_thread=False) + try: + candidate.row_factory = None + candidate.execute("PRAGMA journal_mode=WAL") + candidate.execute(f"PRAGMA busy_timeout={DEFAULT_BUSY_TIMEOUT_MS}") + candidate.execute("PRAGMA foreign_keys=ON") + candidate.execute("PRAGMA synchronous=NORMAL") + candidate.execute("PRAGMA cache_size=-64000") + candidate.execute("PRAGMA mmap_size=268435456") + return candidate + except BaseException: + candidate.close() + raise + + self.conn = open_serving_connection(self.db_path, open_existing) from truememory.tier_switch.runtime import serving_operation - with serving_operation(self.conn, connection_lock=self._write_lock): + with serving_operation(self.conn, connection_lock=self._write_lock) as operation: + if operation.selection is not None or operation.policy is not None: + self._initialize_connection() + self.stats["message_count"] = get_message_count(self.conn) + return self # Detect available tables tables = { row[0] @@ -1389,6 +1741,8 @@ def open(self, rebuild_vectors: bool = True) -> "TrueMemoryEngine": sqlite_vec.load(self.conn) _vectors_load_error = None except Exception as e: + from truememory.tier_switch.runtime import raise_if_serving_rejection + raise_if_serving_rejection(e) _vectors_load_error = f"{type(e).__name__}: {e}" logger.warning( "sqlite-vec unavailable (%s); falling back to FTS-only " @@ -1447,7 +1801,9 @@ def open(self, rebuild_vectors: bool = True) -> "TrueMemoryEngine": ) except TrueMemoryMigrationError: raise - except Exception: + except Exception as _serving_error: + from truememory.tier_switch.runtime import raise_if_serving_rejection + raise_if_serving_rejection(_serving_error) logger.exception("Vector table rebuild failed") self._has_vectors = False @@ -1458,7 +1814,9 @@ def open(self, rebuild_vectors: bool = True) -> "TrueMemoryEngine": try: count = self.conn.execute("SELECT COUNT(*) FROM messages").fetchone()[0] self.stats["message_count"] = count - except Exception: + except Exception as _serving_error: + from truememory.tier_switch.runtime import raise_if_serving_rejection + raise_if_serving_rejection(_serving_error) logger.debug("Failed to count messages in open()", exc_info=True) self.stats["message_count"] = 0 @@ -1491,285 +1849,321 @@ def ingest(self, data_path: str | Path) -> dict: Returns: Dict mapping step names to timing strings (or error messages). """ - data_path = Path(data_path) - stats: dict = {} + from truememory.tier_switch.runtime import destructive_ingest_operation + with destructive_ingest_operation(self): + data_path = Path(data_path) + stats: dict = {} - # ── 1. Create database ──────────────────────────────────────────── - try: - t0 = time.time() - # Remove stale database so every run is clean. - if self.db_path.exists() and str(self.db_path) != ":memory:": - self.db_path.unlink() - self.conn = create_db(self.db_path) - stats["create_db"] = f"{time.time() - t0:.3f}s" - except Exception as exc: - stats["create_db"] = f"ERROR: {exc}" - logger.warning("create_db failed", exc_info=True) - return stats # Can't continue without a database. - - # ── 2. Load messages ────────────────────────────────────────────── - try: - t0 = time.time() - msg_count = load_messages_from_file(self.conn, data_path) - elapsed = time.time() - t0 - stats["load_messages"] = f"{msg_count} messages in {elapsed:.3f}s" - self.stats["message_count"] = msg_count - except Exception as exc: - stats["load_messages"] = f"ERROR: {exc}" - logger.warning("load_messages failed", exc_info=True) - self.stats["message_count"] = 0 - - # ── 3. Build vector embeddings ──────────────────────────────────── - if _HAS_VECTOR: + # ── 1. Create database ──────────────────────────────────────────── try: t0 = time.time() - init_vec_table(self.conn) - vec_count = build_vectors(self.conn) + # Remove stale database so every run is clean. + if self.db_path.exists() and str(self.db_path) != ":memory:": + self.db_path.unlink() + self.conn = create_db(self.db_path) + stats["create_db"] = f"{time.time() - t0:.3f}s" + except Exception as exc: + from truememory.tier_switch.runtime import raise_if_serving_rejection + raise_if_serving_rejection(exc) + stats["create_db"] = f"ERROR: {exc}" + logger.warning("create_db failed", exc_info=True) + return stats # Can't continue without a database. + + # ── 2. Load messages ────────────────────────────────────────────── + try: + t0 = time.time() + msg_count = load_messages_from_file(self.conn, data_path) elapsed = time.time() - t0 - stats["build_vectors"] = f"{vec_count} vectors in {elapsed:.3f}s" - self._has_vectors = True + stats["load_messages"] = f"{msg_count} messages in {elapsed:.3f}s" + self.stats["message_count"] = msg_count except Exception as exc: - stats["build_vectors"] = f"ERROR: {exc}" - logger.debug("build_vectors failed", exc_info=True) + from truememory.tier_switch.runtime import raise_if_serving_rejection + raise_if_serving_rejection(exc) + stats["load_messages"] = f"ERROR: {exc}" + logger.warning("load_messages failed", exc_info=True) + self.stats["message_count"] = 0 + + # ── 3. Build vector embeddings ──────────────────────────────────── + if _HAS_VECTOR: + try: + t0 = time.time() + init_vec_table(self.conn) + vec_count = build_vectors(self.conn) + elapsed = time.time() - t0 + stats["build_vectors"] = f"{vec_count} vectors in {elapsed:.3f}s" + self._has_vectors = True + except Exception as exc: + from truememory.tier_switch.runtime import raise_if_serving_rejection + raise_if_serving_rejection(exc) + stats["build_vectors"] = f"ERROR: {exc}" + logger.debug("build_vectors failed", exc_info=True) + self._has_vectors = False + else: + stats["build_vectors"] = "SKIPPED (sqlite-vec or model2vec not installed)" self._has_vectors = False - else: - stats["build_vectors"] = "SKIPPED (sqlite-vec or model2vec not installed)" - self._has_vectors = False - # Update hybrid capability — needs vectors to be useful. - self._has_hybrid = _HAS_HYBRID and self._has_vectors + # Update hybrid capability — needs vectors to be useful. + self._has_hybrid = _HAS_HYBRID and self._has_vectors - # ── 3b. Build scene clusters ─────────────────────────────────────── - if _HAS_CLUSTERING and self._has_vectors: - try: - t0 = time.time() - n_clusters = cluster_messages(self.conn) - stats["build_clusters"] = f"{n_clusters} clusters in {time.time() - t0:.3f}s" - self._has_clustering = True - if not self.conn.in_transaction: + # ── 3b. Build scene clusters ─────────────────────────────────────── + if _HAS_CLUSTERING and self._has_vectors: + try: + t0 = time.time() + n_clusters = cluster_messages(self.conn) + stats["build_clusters"] = f"{n_clusters} clusters in {time.time() - t0:.3f}s" + self._has_clustering = True + if not self.conn.in_transaction: + self._get_maintenance_coordinator().observe_clustering_outcome( + "success_empty" if n_clusters == 0 else "success") + except Exception as exc: + from truememory.tier_switch.runtime import raise_if_serving_rejection + raise_if_serving_rejection(exc) + stats["build_clusters"] = f"ERROR: {exc}" self._get_maintenance_coordinator().observe_clustering_outcome( - "success_empty" if n_clusters == 0 else "success") - except Exception as exc: - stats["build_clusters"] = f"ERROR: {exc}" + "unavailable" if isinstance(exc, ImportError) else "failed", type(exc).__name__) + self._has_clustering = False + else: + stats["build_clusters"] = "SKIPPED (clustering or vectors not available)" self._get_maintenance_coordinator().observe_clustering_outcome( - "unavailable" if isinstance(exc, ImportError) else "failed", type(exc).__name__) + "unavailable", "ClusteringRuntimeUnavailable" if not _HAS_CLUSTERING else "VectorRuntimeUnavailable") self._has_clustering = False - else: - stats["build_clusters"] = "SKIPPED (clustering or vectors not available)" - self._get_maintenance_coordinator().observe_clustering_outcome( - "unavailable", "ClusteringRuntimeUnavailable" if not _HAS_CLUSTERING else "VectorRuntimeUnavailable") - self._has_clustering = False - # ── 4. Build entity profiles (L0) ───────────────────────────────── - if _HAS_PERSONALITY: - try: - t0 = time.time() - build_entity_profiles(self.conn) - stats["build_profiles"] = f"{time.time() - t0:.3f}s" - except Exception as exc: - stats["build_profiles"] = f"ERROR: {exc}" - logger.debug("build_profiles failed", exc_info=True) - self._has_personality = False - else: - stats["build_profiles"] = "SKIPPED (personality module not available)" + # ── 4. Build entity profiles (L0) ───────────────────────────────── + if _HAS_PERSONALITY: + try: + t0 = time.time() + build_entity_profiles(self.conn) + stats["build_profiles"] = f"{time.time() - t0:.3f}s" + except Exception as exc: + from truememory.tier_switch.runtime import raise_if_serving_rejection + raise_if_serving_rejection(exc) + stats["build_profiles"] = f"ERROR: {exc}" + logger.debug("build_profiles failed", exc_info=True) + self._has_personality = False + else: + stats["build_profiles"] = "SKIPPED (personality module not available)" - # ── 4.5. Build entity style vectors (L0 char-n-gram) ───────────── - if _HAS_STYLE_VEC: - try: - from truememory.maintenance import maintenance_observation, run_engine_maintenance - coordinator = self._get_maintenance_coordinator() - borrowed = self.conn.in_transaction - with maintenance_observation(coordinator, borrowed=borrowed) as observed: - report = run_engine_maintenance(self.conn, coordinator, force=True, prepare_extensions=False, - include_layers=False, include_style=True, allow_caller_transaction=borrowed) - if not borrowed and self.conn.in_transaction: - self.conn.rollback() - raise RuntimeError("Style work left an unfinished transaction") - observed.append(report) - result = report.style_result - if result.outcome in {"success", "success_empty"}: - stats["build_style_vectors"] = f"{result.output_count} vectors in {result.elapsed_seconds:.3f}s" - else: - stats["build_style_vectors"] = result.outcome.upper() + " (" + (result.error_category or "Pending") + ")" - if result.pending_caller_commit: - stats["build_style_vectors"] += " (pending caller commit)" - self._has_style_vec = True - except Exception as exc: - stats["build_style_vectors"] = "ERROR (" + type(exc).__name__[:64] + ")" - logger.debug("build_style_vectors failed", exc_info=True) - else: - stats["build_style_vectors"] = "SKIPPED (personality_style_vec module not available)" + # ── 4.5. Build entity style vectors (L0 char-n-gram) ───────────── + if _HAS_STYLE_VEC: + try: + from truememory.maintenance import maintenance_observation, run_engine_maintenance + coordinator = self._get_maintenance_coordinator() + borrowed = self.conn.in_transaction + with maintenance_observation(coordinator, borrowed=borrowed) as observed: + report = run_engine_maintenance(self.conn, coordinator, force=True, prepare_extensions=False, + include_layers=False, include_style=True, allow_caller_transaction=borrowed) + if not borrowed and self.conn.in_transaction: + self.conn.rollback() + raise RuntimeError("Style work left an unfinished transaction") + observed.append(report) + result = report.style_result + if result.outcome in {"success", "success_empty"}: + stats["build_style_vectors"] = f"{result.output_count} vectors in {result.elapsed_seconds:.3f}s" + else: + stats["build_style_vectors"] = result.outcome.upper() + " (" + (result.error_category or "Pending") + ")" + if result.pending_caller_commit: + stats["build_style_vectors"] += " (pending caller commit)" + self._has_style_vec = True + except Exception as exc: + from truememory.tier_switch.runtime import raise_if_serving_rejection + raise_if_serving_rejection(exc) + stats["build_style_vectors"] = "ERROR (" + type(exc).__name__[:64] + ")" + logger.debug("build_style_vectors failed", exc_info=True) + else: + stats["build_style_vectors"] = "SKIPPED (personality_style_vec module not available)" - # ── 5. Extract preferences (L0) ─────────────────────────────────── - if _HAS_PERSONALITY: - try: - t0 = time.time() - extract_preferences(self.conn) - stats["extract_preferences"] = f"{time.time() - t0:.3f}s" - except Exception as exc: - stats["extract_preferences"] = f"ERROR: {exc}" - logger.debug("extract_preferences failed", exc_info=True) - else: - stats["extract_preferences"] = "SKIPPED (personality module not available)" + # ── 5. Extract preferences (L0) ─────────────────────────────────── + if _HAS_PERSONALITY: + try: + t0 = time.time() + extract_preferences(self.conn) + stats["extract_preferences"] = f"{time.time() - t0:.3f}s" + except Exception as exc: + from truememory.tier_switch.runtime import raise_if_serving_rejection + raise_if_serving_rejection(exc) + stats["extract_preferences"] = f"ERROR: {exc}" + logger.debug("extract_preferences failed", exc_info=True) + else: + stats["extract_preferences"] = "SKIPPED (personality module not available)" - # ── 6. Build summaries (L5) ─────────────────────────────────────── - if _HAS_CONSOLIDATION: - try: - t0 = time.time() - build_summaries(self.conn) - stats["build_summaries"] = f"{time.time() - t0:.3f}s" - except Exception as exc: - stats["build_summaries"] = f"ERROR: {exc}" - logger.debug("build_summaries failed", exc_info=True) - self._has_consolidation = False - else: - stats["build_summaries"] = "SKIPPED (consolidation module not available)" + # ── 6. Build summaries (L5) ─────────────────────────────────────── + if _HAS_CONSOLIDATION: + try: + t0 = time.time() + build_summaries(self.conn) + stats["build_summaries"] = f"{time.time() - t0:.3f}s" + except Exception as exc: + from truememory.tier_switch.runtime import raise_if_serving_rejection + raise_if_serving_rejection(exc) + stats["build_summaries"] = f"ERROR: {exc}" + logger.debug("build_summaries failed", exc_info=True) + self._has_consolidation = False + else: + stats["build_summaries"] = "SKIPPED (consolidation module not available)" - # ── 7. Detect contradictions (L5) ───────────────────────────────── - if _HAS_CONSOLIDATION: - try: - t0 = time.time() - detect_contradictions(self.conn) - stats["detect_contradictions"] = f"{time.time() - t0:.3f}s" - except Exception as exc: - stats["detect_contradictions"] = f"ERROR: {exc}" - logger.debug("detect_contradictions failed", exc_info=True) - else: - stats["detect_contradictions"] = "SKIPPED (consolidation module not available)" + # ── 7. Detect contradictions (L5) ───────────────────────────────── + if _HAS_CONSOLIDATION: + try: + t0 = time.time() + detect_contradictions(self.conn) + stats["detect_contradictions"] = f"{time.time() - t0:.3f}s" + except Exception as exc: + from truememory.tier_switch.runtime import raise_if_serving_rejection + raise_if_serving_rejection(exc) + stats["detect_contradictions"] = f"ERROR: {exc}" + logger.debug("detect_contradictions failed", exc_info=True) + else: + stats["detect_contradictions"] = "SKIPPED (consolidation module not available)" - # ── 8. Build surprise index (predictive coding) ─────────────────── - if _HAS_PREDICTIVE: - try: - t0 = time.time() - build_surprise_index(self.conn) - stats["build_surprise_index"] = f"{time.time() - t0:.3f}s" - except Exception as exc: - stats["build_surprise_index"] = f"ERROR: {exc}" - logger.debug("build_surprise_index failed", exc_info=True) - self._has_predictive = False - else: - stats["build_surprise_index"] = "SKIPPED (predictive module not available)" + # ── 8. Build surprise index (predictive coding) ─────────────────── + if _HAS_PREDICTIVE: + try: + t0 = time.time() + build_surprise_index(self.conn) + stats["build_surprise_index"] = f"{time.time() - t0:.3f}s" + except Exception as exc: + from truememory.tier_switch.runtime import raise_if_serving_rejection + raise_if_serving_rejection(exc) + stats["build_surprise_index"] = f"ERROR: {exc}" + logger.debug("build_surprise_index failed", exc_info=True) + self._has_predictive = False + else: + stats["build_surprise_index"] = "SKIPPED (predictive module not available)" - # ── 9. Build separation vectors (B2) ───────────────────────────── - if _HAS_VECTOR: - try: - t0 = time.time() - sep_count = build_separation_vectors(self.conn) - stats["build_separation_vectors"] = f"{sep_count} vectors in {time.time() - t0:.3f}s" - except Exception as exc: - stats["build_separation_vectors"] = f"ERROR: {exc}" - logger.debug("build_separation_vectors failed", exc_info=True) - else: - stats["build_separation_vectors"] = "SKIPPED (vector module not available)" + # ── 9. Build separation vectors (B2) ───────────────────────────── + if _HAS_VECTOR: + try: + t0 = time.time() + sep_count = build_separation_vectors(self.conn) + stats["build_separation_vectors"] = f"{sep_count} vectors in {time.time() - t0:.3f}s" + except Exception as exc: + from truememory.tier_switch.runtime import raise_if_serving_rejection + raise_if_serving_rejection(exc) + stats["build_separation_vectors"] = f"ERROR: {exc}" + logger.debug("build_separation_vectors failed", exc_info=True) + else: + stats["build_separation_vectors"] = "SKIPPED (vector module not available)" - # ── 10. Detect episodes (B1) ───────────────────────────────────── - if _HAS_TEMPORAL: - try: - t0 = time.time() - ep_count = detect_episodes(self.conn) - stats["detect_episodes"] = f"{ep_count} episodes in {time.time() - t0:.3f}s" - except Exception as exc: - stats["detect_episodes"] = f"ERROR: {exc}" - logger.debug("detect_episodes failed", exc_info=True) - else: - stats["detect_episodes"] = "SKIPPED (temporal module not available)" + # ── 10. Detect episodes (B1) ───────────────────────────────────── + if _HAS_TEMPORAL: + try: + t0 = time.time() + ep_count = detect_episodes(self.conn) + stats["detect_episodes"] = f"{ep_count} episodes in {time.time() - t0:.3f}s" + except Exception as exc: + from truememory.tier_switch.runtime import raise_if_serving_rejection + raise_if_serving_rejection(exc) + stats["detect_episodes"] = f"ERROR: {exc}" + logger.debug("detect_episodes failed", exc_info=True) + else: + stats["detect_episodes"] = "SKIPPED (temporal module not available)" - # ── 11. Detect landmark events (E3) ────────────────────────────── - if _HAS_TEMPORAL: - try: - t0 = time.time() - lm_count = detect_landmark_events(self.conn) - stats["detect_landmarks"] = f"{lm_count} events in {time.time() - t0:.3f}s" - except Exception as exc: - stats["detect_landmarks"] = f"ERROR: {exc}" - logger.debug("detect_landmarks failed", exc_info=True) - else: - stats["detect_landmarks"] = "SKIPPED (temporal module not available)" - - # ── 12. Build entity summary sheets (B3) ───────────────────────── - # Disabled 2026-04-24 per MEMORIST-L4 research. - # See CHANGELOG v0.6.0 for rationale. - # The function wrote `summaries` rows with period='entity_profile' - # that saturated top-1 retrieval by keyword match and leaked - # superseded facts into contradiction scoring. Disabling produced - # +5.3 pts on the L4 composite probe metric (Pareto-dominant, - # see REPORT.md §3 Table "D1 vs C1" and §10.7 Ablation 2). - # - # Escape hatch: set TRUEMEMORY_ENTITY_SHEETS=1 to re-enable this - # function. Users who regress on real-world workloads can revert - # without a code patch. Intended to be removed in a future release - # once long-horizon production telemetry confirms the disable. - _entity_sheets_enabled = ( - os.environ.get("TRUEMEMORY_ENTITY_SHEETS", "") - .strip().lower() in {"1", "true", "yes", "on"} - ) - if _HAS_CONSOLIDATION and _entity_sheets_enabled: - try: - t0 = time.time() - sheet_count = build_entity_summary_sheets(self.conn) - stats["entity_summary_sheets"] = f"{sheet_count} sheets in {time.time() - t0:.3f}s (re-enabled via TRUEMEMORY_ENTITY_SHEETS=1)" - except Exception as exc: - stats["entity_summary_sheets"] = f"ERROR: {exc}" - logger.debug("entity_summary_sheets failed", exc_info=True) - elif _HAS_CONSOLIDATION: - stats["entity_summary_sheets"] = "DISABLED (MEMORIST-L4; set TRUEMEMORY_ENTITY_SHEETS=1 to re-enable)" - else: - stats["entity_summary_sheets"] = "SKIPPED (consolidation module not available)" + # ── 11. Detect landmark events (E3) ────────────────────────────── + if _HAS_TEMPORAL: + try: + t0 = time.time() + lm_count = detect_landmark_events(self.conn) + stats["detect_landmarks"] = f"{lm_count} events in {time.time() - t0:.3f}s" + except Exception as exc: + from truememory.tier_switch.runtime import raise_if_serving_rejection + raise_if_serving_rejection(exc) + stats["detect_landmarks"] = f"ERROR: {exc}" + logger.debug("detect_landmarks failed", exc_info=True) + else: + stats["detect_landmarks"] = "SKIPPED (temporal module not available)" + + # ── 12. Build entity summary sheets (B3) ───────────────────────── + # Disabled 2026-04-24 per MEMORIST-L4 research. + # See CHANGELOG v0.6.0 for rationale. + # The function wrote `summaries` rows with period='entity_profile' + # that saturated top-1 retrieval by keyword match and leaked + # superseded facts into contradiction scoring. Disabling produced + # +5.3 pts on the L4 composite probe metric (Pareto-dominant, + # see REPORT.md §3 Table "D1 vs C1" and §10.7 Ablation 2). + # + # Escape hatch: set TRUEMEMORY_ENTITY_SHEETS=1 to re-enable this + # function. Users who regress on real-world workloads can revert + # without a code patch. Intended to be removed in a future release + # once long-horizon production telemetry confirms the disable. + _entity_sheets_enabled = ( + os.environ.get("TRUEMEMORY_ENTITY_SHEETS", "") + .strip().lower() in {"1", "true", "yes", "on"} + ) + if _HAS_CONSOLIDATION and _entity_sheets_enabled: + try: + t0 = time.time() + sheet_count = build_entity_summary_sheets(self.conn) + stats["entity_summary_sheets"] = f"{sheet_count} sheets in {time.time() - t0:.3f}s (re-enabled via TRUEMEMORY_ENTITY_SHEETS=1)" + except Exception as exc: + from truememory.tier_switch.runtime import raise_if_serving_rejection + raise_if_serving_rejection(exc) + stats["entity_summary_sheets"] = f"ERROR: {exc}" + logger.debug("entity_summary_sheets failed", exc_info=True) + elif _HAS_CONSOLIDATION: + stats["entity_summary_sheets"] = "DISABLED (MEMORIST-L4; set TRUEMEMORY_ENTITY_SHEETS=1 to re-enable)" + else: + stats["entity_summary_sheets"] = "SKIPPED (consolidation module not available)" - # ── 13. Build structured facts (B4) ────────────────────────────── - if _HAS_CONSOLIDATION: - try: - t0 = time.time() - fact_count = build_structured_facts(self.conn) - stats["structured_facts"] = f"{fact_count} facts in {time.time() - t0:.3f}s" - except Exception as exc: - stats["structured_facts"] = f"ERROR: {exc}" - logger.debug("structured_facts failed", exc_info=True) - else: - stats["structured_facts"] = "SKIPPED (consolidation module not available)" + # ── 13. Build structured facts (B4) ────────────────────────────── + if _HAS_CONSOLIDATION: + try: + t0 = time.time() + fact_count = build_structured_facts(self.conn) + stats["structured_facts"] = f"{fact_count} facts in {time.time() - t0:.3f}s" + except Exception as exc: + from truememory.tier_switch.runtime import raise_if_serving_rejection + raise_if_serving_rejection(exc) + stats["structured_facts"] = f"ERROR: {exc}" + logger.debug("structured_facts failed", exc_info=True) + else: + stats["structured_facts"] = "SKIPPED (consolidation module not available)" - # ── 14. Build Dunbar hierarchy (E2) ────────────────────────────── - if _HAS_PERSONALITY: - try: - t0 = time.time() - primary = None + # ── 14. Build Dunbar hierarchy (E2) ────────────────────────────── + if _HAS_PERSONALITY: try: - row = self.conn.execute( - "SELECT sender, COUNT(*) as cnt FROM messages " - "WHERE sender != '' AND sender IS NOT NULL " - "GROUP BY sender ORDER BY cnt DESC LIMIT 1" - ).fetchone() - if row and row[0] and row[0].strip(): - primary = row[0] + t0 = time.time() + primary = None + try: + row = self.conn.execute( + "SELECT sender, COUNT(*) as cnt FROM messages " + "WHERE sender != '' AND sender IS NOT NULL " + "GROUP BY sender ORDER BY cnt DESC LIMIT 1" + ).fetchone() + if row and row[0] and row[0].strip(): + primary = row[0] + except Exception as exc: + from truememory.tier_switch.runtime import raise_if_serving_rejection + raise_if_serving_rejection(exc) + logger.debug("Failed to detect primary entity for Dunbar: %s", exc) + dunbar_result = build_dunbar_hierarchy(self.conn, primary_entity=primary) + dunbar_count = len(dunbar_result) if isinstance(dunbar_result, dict) else dunbar_result + stats["dunbar_hierarchy"] = f"{dunbar_count} relationships in {time.time() - t0:.3f}s" except Exception as exc: - logger.debug("Failed to detect primary entity for Dunbar: %s", exc) - dunbar_result = build_dunbar_hierarchy(self.conn, primary_entity=primary) - dunbar_count = len(dunbar_result) if isinstance(dunbar_result, dict) else dunbar_result - stats["dunbar_hierarchy"] = f"{dunbar_count} relationships in {time.time() - t0:.3f}s" - except Exception as exc: - stats["dunbar_hierarchy"] = f"ERROR: {exc}" - logger.debug("dunbar_hierarchy failed", exc_info=True) - else: - stats["dunbar_hierarchy"] = "SKIPPED (personality module not available)" - - # ── Record capabilities ─────────────────────────────────────────── - self.stats["capabilities"] = { - "fts5": True, - "vector_search": self._has_vectors, - "hybrid_rrf": self._has_hybrid, - "temporal": self._has_temporal, - "salience": self._has_salience, - "personality": self._has_personality, - "style_vec": self._has_style_vec, - "consolidation": self._has_consolidation, - "predictive": self._has_predictive, - "reranker": self._has_reranker, - "hyde": self._has_hyde, - "clustering": self._has_clustering, - } + from truememory.tier_switch.runtime import raise_if_serving_rejection + raise_if_serving_rejection(exc) + stats["dunbar_hierarchy"] = f"ERROR: {exc}" + logger.debug("dunbar_hierarchy failed", exc_info=True) + else: + stats["dunbar_hierarchy"] = "SKIPPED (personality module not available)" + + # ── Record capabilities ─────────────────────────────────────────── + self.stats["capabilities"] = { + "fts5": True, + "vector_search": self._has_vectors, + "hybrid_rrf": self._has_hybrid, + "temporal": self._has_temporal, + "salience": self._has_salience, + "personality": self._has_personality, + "style_vec": self._has_style_vec, + "consolidation": self._has_consolidation, + "predictive": self._has_predictive, + "reranker": self._has_reranker, + "hyde": self._has_hyde, + "clustering": self._has_clustering, + } - self.ready = True - return stats + self.ready = True + return stats # ────────────────────────────────────────────────────────────────────── # Search — full 6-layer pipeline @@ -1789,7 +2183,9 @@ def search_vectors_raw(self, query: str, limit: int = 5) -> list[dict] | None: try: from truememory.vector_search import search_vector_raw return search_vector_raw(self.conn, query, limit=limit) - except Exception: + except Exception as _serving_error: + from truememory.tier_switch.runtime import raise_if_serving_rejection + raise_if_serving_rejection(_serving_error) return None @engine_operation @@ -1836,7 +2232,9 @@ def search(self, query: str, limit: int = 10, _skip_surprise_boost: bool = False try: query_info = classify_query(query) search_mode = get_search_mode(query) - except Exception: + except Exception as _serving_error: + from truememory.tier_switch.runtime import raise_if_serving_rejection + raise_if_serving_rejection(_serving_error) logger.debug("Query classification failed in search()", exc_info=True) fts_w = query_info["weights"].get("fts", 1.0) @@ -1851,7 +2249,9 @@ def search(self, query: str, limit: int = 10, _skip_surprise_boost: bool = False include_directives=include_directives, ) source_label = "hybrid" - except Exception: + except Exception as _serving_error: + from truememory.tier_switch.runtime import raise_if_serving_rejection + raise_if_serving_rejection(_serving_error) logger.debug("Hybrid search failed, falling back to FTS5", exc_info=True) results = [] @@ -1863,7 +2263,9 @@ def search(self, query: str, limit: int = 10, _skip_surprise_boost: bool = False for r in results: if "source" not in r: r["source"] = "fts" - except Exception: + except Exception as _serving_error: + from truememory.tier_switch.runtime import raise_if_serving_rejection + raise_if_serving_rejection(_serving_error) logger.debug("FTS search failed in search()", exc_info=True) results = [] @@ -1880,14 +2282,18 @@ def search(self, query: str, limit: int = 10, _skip_surprise_boost: bool = False if _high_diversity and len(results) >= 3: try: results = self._scent_trail(query, results, limit) - except Exception: + except Exception as _serving_error: + from truememory.tier_switch.runtime import raise_if_serving_rejection + raise_if_serving_rejection(_serving_error) logger.debug("Scent trail failed in search()", exc_info=True) # ── 3. Retrieval quality self-check (A4) — only with high diversity if _high_diversity: try: results = self._quality_self_check(query, results, limit) - except Exception: + except Exception as _serving_error: + from truememory.tier_switch.runtime import raise_if_serving_rejection + raise_if_serving_rejection(_serving_error) logger.debug("Quality self-check failed in search()", exc_info=True) # ── 4. Temporal filtering with cross-source feedback (A1) ───────── @@ -1966,10 +2372,14 @@ def search(self, query: str, limit: int = 10, _skip_surprise_boost: bool = False rr["score"] = _resc_max * 0.8 * max(0.0, min(_rs, 1.0)) results.append(rr) existing_ids.add(rr["id"]) - except Exception: + except Exception as _serving_error: + from truememory.tier_switch.runtime import raise_if_serving_rejection + raise_if_serving_rejection(_serving_error) logger.debug("Temporal cross-source rescope failed in search()", exc_info=True) - except Exception: + except Exception as _serving_error: + from truememory.tier_switch.runtime import raise_if_serving_rejection + raise_if_serving_rejection(_serving_error) logger.debug("Temporal filtering failed in search()", exc_info=True) # ── 5. Personality supplementation ──────────────────────────────── @@ -2008,7 +2418,9 @@ def search(self, query: str, limit: int = 10, _skip_surprise_boost: bool = False results.append(pr) if pr_id: existing_ids.add(pr_id) - except Exception: + except Exception as _serving_error: + from truememory.tier_switch.runtime import raise_if_serving_rejection + raise_if_serving_rejection(_serving_error) logger.debug("Personality supplementation failed in search()", exc_info=True) # ── 6. Contradiction / fact timeline ────────────────────────────── @@ -2044,7 +2456,9 @@ def search(self, query: str, limit: int = 10, _skip_surprise_boost: bool = False if cr.get("id") and cr["id"] not in existing_ids: results.append(cr) existing_ids.add(cr["id"]) - except Exception: + except Exception as _serving_error: + from truememory.tier_switch.runtime import raise_if_serving_rejection + raise_if_serving_rejection(_serving_error) logger.debug("Contradiction search failed in search()", exc_info=True) try: @@ -2077,7 +2491,9 @@ def search(self, query: str, limit: int = 10, _skip_surprise_boost: bool = False if sr.get("id") and sr["id"] not in existing_ids: results.append(sr) existing_ids.add(sr["id"]) - except Exception: + except Exception as _serving_error: + from truememory.tier_switch.runtime import raise_if_serving_rejection + raise_if_serving_rejection(_serving_error) logger.debug("Consolidated search failed in search()", exc_info=True) # ── 7. Salience guard with mode-aware threshold (A5) ────────────── @@ -2097,7 +2513,9 @@ def search(self, query: str, limit: int = 10, _skip_surprise_boost: bool = False results = apply_salience_guard( results, query, conn=self.conn, min_salience=min_sal, ) - except Exception: + except Exception as _serving_error: + from truememory.tier_switch.runtime import raise_if_serving_rejection + raise_if_serving_rejection(_serving_error) logger.debug("Salience guard failed in search()", exc_info=True) # ── 7.5 L5 surprise rerank boost (MEMORIST-L5) ── @@ -2129,7 +2547,9 @@ def search(self, query: str, limit: int = 10, _skip_surprise_boost: bool = False rrf_weight=0.4, rerank_weight=0.6, ) - except Exception: + except Exception as _serving_error: + from truememory.tier_switch.runtime import raise_if_serving_rejection + raise_if_serving_rejection(_serving_error) logger.debug("Cross-encoder rerank failed in search()", exc_info=True) # ── 9. Ensure all results have required fields and trim ─────────── @@ -2266,7 +2686,9 @@ def search_agentic( primary_results = reciprocal_rank_fusion( [primary_results, hyde_results] ) - except Exception: + except Exception as _serving_error: + from truememory.tier_switch.runtime import raise_if_serving_rejection + raise_if_serving_rejection(_serving_error) logger.debug("HyDE search failed in search_agentic()", exc_info=True) # ── Score normalization (issue #584) ───────────────────────────── @@ -2307,7 +2729,9 @@ def search_agentic( cr["source"] = cr.get("source", "") + "+cluster_supp" primary_results.append(cr) existing_ids.add(cid) - except Exception: + except Exception as _serving_error: + from truememory.tier_switch.runtime import raise_if_serving_rejection + raise_if_serving_rejection(_serving_error) logger.debug("Cluster search failed in search_agentic()", exc_info=True) # ── Entity-focused search ─────────────────────────────────────── @@ -2341,7 +2765,9 @@ def search_agentic( _ent_after = _validate_iso_date(_ent_intent.get("after")) _eb = _validate_iso_date(_ent_intent.get("before")) _ent_before = _exclusive_upper_bound(_eb) if _eb else None - except Exception: + except Exception as _serving_error: + from truememory.tier_switch.runtime import raise_if_serving_rejection + raise_if_serving_rejection(_serving_error) logger.debug("Entity temporal-window detection failed", exc_info=True) def _in_window(row: dict) -> bool: @@ -2399,7 +2825,9 @@ def _in_window(row: dict) -> bool: primary_results.sort( key=lambda d: (-d.get("score", 0), d.get("id", 0)) ) - except Exception: + except Exception as _serving_error: + from truememory.tier_switch.runtime import raise_if_serving_rejection + raise_if_serving_rejection(_serving_error) logger.debug("Entity-focused search failed in search_agentic()", exc_info=True) # ── Salience guard (deferred from search() for #582) ───────────── @@ -2425,7 +2853,9 @@ def _in_window(row: dict) -> bool: if r.get("_entity_boosted") and r.get("id") ), ) - except Exception: + except Exception as _serving_error: + from truememory.tier_switch.runtime import raise_if_serving_rejection + raise_if_serving_rejection(_serving_error) logger.debug("Salience guard failed in search_agentic()", exc_info=True) # ── Sufficiency check ───────────────────────────────────────────── @@ -2455,7 +2885,9 @@ def _in_window(row: dict) -> bool: if pr.get("id") == rid: pr["score"] = pr.get("score", 0) * 1.15 break - except Exception: + except Exception as _serving_error: + from truememory.tier_switch.runtime import raise_if_serving_rejection + raise_if_serving_rejection(_serving_error) logger.debug("Refined query search failed in search_agentic()", exc_info=True) # Re-sort after adding refined results @@ -2491,11 +2923,15 @@ def _in_window(row: dict) -> bool: query, final_results[:_llm_cap], llm_fn=llm_fn, top_k=limit, ) - except Exception: + except Exception as _serving_error: + from truememory.tier_switch.runtime import raise_if_serving_rejection + raise_if_serving_rejection(_serving_error) logger.debug("LLM reranking failed in search_agentic()", exc_info=True) return self._clean_results(final_results, limit, max_per_session=max_per_session, include_directives=include_directives) - except Exception: + except Exception as _serving_error: + from truememory.tier_switch.runtime import raise_if_serving_rejection + raise_if_serving_rejection(_serving_error) logger.warning("Cross-encoder rerank failed in search_agentic()", exc_info=True) # ── LLM reranking without cross-encoder ────────────────────────── @@ -2507,7 +2943,9 @@ def _in_window(row: dict) -> bool: query, primary_results[:_llm_cap], llm_fn=llm_fn, top_k=limit, ) - except Exception: + except Exception as _serving_error: + from truememory.tier_switch.runtime import raise_if_serving_rejection + raise_if_serving_rejection(_serving_error) logger.debug("LLM reranking (standalone) failed in search_agentic()", exc_info=True) return self._clean_results(primary_results, limit, max_per_session=max_per_session, @@ -2581,7 +3019,7 @@ def _quality_self_check(self, query: str, results: list[dict], limit: int) -> li # Search — simple FTS5-only (for benchmarking) # ────────────────────────────────────────────────────────────────────── - @engine_operation + @engine_handle_operation def search_simple(self, query: str, limit: int = 10) -> list[dict]: """ Simple FTS5-only search (no vector, no layers). @@ -2594,7 +3032,9 @@ def search_simple(self, query: str, limit: int = 10) -> list[dict]: try: results = search_fts(self.conn, query, limit=limit) - except Exception: + except Exception as _serving_error: + from truememory.tier_switch.runtime import raise_if_serving_rejection + raise_if_serving_rejection(_serving_error) return [] cleaned: list[dict] = [] @@ -2707,6 +3147,7 @@ def close(self): self._runtime_initialized = False self._runtime_vector_generation = None self._runtime_vector_connection = None + self._reconnect_receipt = None def __enter__(self): return self diff --git a/truememory/hybrid.py b/truememory/hybrid.py index 406f9b04..ea3ac220 100644 --- a/truememory/hybrid.py +++ b/truememory/hybrid.py @@ -179,7 +179,9 @@ def search_hybrid( sep_tbl = _active_sep_table(conn) conn.execute(f"SELECT rowid FROM {sep_tbl} LIMIT 1").fetchone() _has_sep = True - except Exception: + except Exception as exc: + from truememory.tier_switch.runtime import raise_if_serving_rejection + raise_if_serving_rejection(exc) pass except ImportError: pass @@ -209,7 +211,9 @@ def search_hybrid( sender_count = unique_senders_row[0] if unique_senders_row else 0 if sender_count > 5: sep_results = search_vector_separation(conn, query, limit=_CANDIDATE_POOL, _query_blob=_q_blob, include_directives=include_directives) - except Exception: + except Exception as exc: + from truememory.tier_switch.runtime import raise_if_serving_rejection + raise_if_serving_rejection(exc) pass # ------------------------------------------------------------------ diff --git a/truememory/ingest/cli.py b/truememory/ingest/cli.py index 9488fa15..53d4ccd2 100644 --- a/truememory/ingest/cli.py +++ b/truememory/ingest/cli.py @@ -12,6 +12,7 @@ import json import logging import os +import re import sys from pathlib import Path @@ -89,6 +90,17 @@ def main(): p_upgrade.add_argument("tier", choices=["edge", "base", "pro"], help="Target tier") p_upgrade.add_argument("--force", action="store_true", help="Re-embed even if tier is unchanged") + p_cancel = sub.add_parser( + "cancel-rebuild", help="Cancel one explicitly identified interrupted tier rebuild", + description=("Read truememory_status(status_id) for the matching rebuild.id and rebuild.job_id. " + "This command cancels only that marker. Retry upgrade-tier explicitly afterward."), + ) + p_cancel.add_argument("--status-id", required=True, type=_positive_rebuild_status_id, + help="Exact positive rebuild.id returned by truememory_status") + p_cancel.add_argument("--job-id", required=True, type=_rebuild_job_id, + help="Exact matching 32-character lowercase hexadecimal rebuild.job_id") + p_cancel.add_argument("--db", default=None, help="Path to the same truememory database") + # --- migrate-memory command --- p_migrate = sub.add_parser("migrate-memory", help="Migrate MEMORY.md facts into TrueMemory") p_migrate.add_argument("--path", default=None, help="Path to MEMORY.md (auto-detected if not specified)") @@ -119,12 +131,53 @@ def main(): _run_setup(args) elif args.command == "upgrade-tier": _run_upgrade_tier(args) + elif args.command == "cancel-rebuild": + _run_cancel_rebuild(args) elif args.command == "migrate-memory": _run_migrate_memory(args) else: parser.print_help() +def _positive_rebuild_status_id(value: str) -> int: + if (type(value) is not str or re.fullmatch(r"[1-9][0-9]{0,18}", value) is None + or int(value) > 9223372036854775807): + raise argparse.ArgumentTypeError("status ID must be an exact positive SQLite integer") + return int(value) + + +def _rebuild_job_id(value: str) -> str: + if type(value) is not str or re.fullmatch(r"[0-9a-f]{32}", value) is None: + raise argparse.ArgumentTypeError("job ID must be exactly 32 lowercase hexadecimal characters") + return value + + +def _run_cancel_rebuild(args: argparse.Namespace) -> None: + """Cancel one explicit marker; do not discover, retry, or rebuild anything.""" + try: + if type(args.status_id) is not int or not 1 <= args.status_id <= 9223372036854775807: + raise argparse.ArgumentTypeError("status ID must be an exact positive SQLite integer") + job_id = _rebuild_job_id(args.job_id) + except argparse.ArgumentTypeError as error: + print(f"Invalid rebuild identity: {error}", file=sys.stderr) + raise SystemExit(2) from error + + from truememory.tier_switch.manager import RebuildManager + manager = RebuildManager() + try: + result = manager.cancel(args.status_id, expected_job_id=job_id, + db_path=Path(args.db) if args.db is not None else None) + except Exception as error: + print(f"Rebuild cancellation refused: {error}", file=sys.stderr) + print("Read truememory_status(status_id) again and use its matching rebuild.id and rebuild.job_id.", file=sys.stderr) + raise SystemExit(1) from error + if result.get("status") != "cancelled" or result.get("job_id") != job_id: + print("No matching interrupted rebuild was cancelled. Read truememory_status(status_id) again.", file=sys.stderr) + raise SystemExit(1) + print(f"Rebuild {args.status_id} cancelled (job_id={job_id}).") + print("Retry the original upgrade-tier request explicitly to resume the eligible inactive prefix.") + + def _run_ingest(args): level = logging.DEBUG if args.verbose else logging.INFO logging.basicConfig(level=level, format="%(levelname)s %(message)s") @@ -564,6 +617,55 @@ def _setup_cli_integrations(args: argparse.Namespace, config: dict) -> list[str] return failures +def _setup_model_readiness(tier: str, existing_tier: str) -> list[str]: + """Switch through certification, then exercise the admitted serving identity.""" + from truememory.maintenance import maintenance_owner + from truememory.tier_switch.manager import RebuildManager, _get_db_path, _open_db + from truememory.tier_switch.runtime import serving_operation + + failures: list[str] = [] + manager = RebuildManager() + try: + if not manager.run_rebuild_sync(tier): + return ["tier activation"] + except Exception as error: + print(f"Tier activation remains incomplete: {error}", file=sys.stderr) + return ["tier activation"] + if manager._last_action == "config_only" or {tier, existing_tier} == {"base", "pro"}: + return failures + try: + with maintenance_owner(_get_db_path()): + conn = _open_db(initialize=True) + try: + with serving_operation(conn) as operation: + if operation.tier != tier: + raise RuntimeError("Requested setup tier is not the admitted serving tier") + try: + from truememory.vector_search import _encode_with_mps_fallback, get_model + _encode_with_mps_fallback(get_model(), ["TrueMemory setup readiness check."], show_progress_bar=False) + if tier in ("base", "pro"): + print(" \033[32mQwen3-Embedding-0.6B @ 256d Matryoshka ready\033[0m") + else: + print(" \033[32mpotion-base-8M (256-dim) ready\033[0m") + except Exception as error: + failures.append("embedding model") + print(f"Embedding model readiness check failed: {error}", file=sys.stderr) + try: + from truememory.reranker import get_reranker + get_reranker().predict([("TrueMemory setup readiness check.", "TrueMemory setup readiness check.")], + show_progress_bar=False) + print(" \033[32mCross-encoder reranker ready\033[0m") + except Exception as error: + failures.append("reranker") + print(f"Reranker readiness check failed: {error}", file=sys.stderr) + finally: + conn.close() + except Exception as error: + failures.append("serving admission") + print(f"Serving readiness admission failed: {error}", file=sys.stderr) + return failures + + def _run_setup(args): """Interactive first-time setup wizard for TrueMemory.""" print(_SAURON_BANNER) @@ -591,39 +693,7 @@ def _run_setup(args): print(f" \033[33m⚠ Invalid choice '{choice}', defaulting to Edge tier.\033[0m") tier = "edge" - # Exercise one synthetic input so a lazy model-server proxy is not mistaken - # for a successful download/load. No user memories are read by this check. - model_failures: list[str] = [] - print() - print(" \033[1mDownloading embedding model...\033[0m") - try: - os.environ["TRUEMEMORY_EMBED_MODEL"] = tier - from truememory.vector_search import _encode_with_mps_fallback, set_embedding_model, get_model - set_embedding_model(tier) - _encode_with_mps_fallback( - get_model(), ["TrueMemory setup readiness check."], show_progress_bar=False, - ) - if tier in ("base", "pro"): - print(" \033[32m✓ Qwen3-Embedding-0.6B @ 256d Matryoshka ready\033[0m") - else: - print(" \033[32m✓ potion-base-8M (256-dim) ready\033[0m") - except Exception as e: - model_failures.append("embedding model") - print(f" Embedding model readiness check failed: {e}", file=sys.stderr) - - # Verify the selected tier's reranker, including the model-server path. - try: - from truememory.reranker import get_reranker, set_active_tier - set_active_tier(tier) - print(" Downloading reranker model...") - get_reranker().predict( - [("TrueMemory setup readiness check.", "TrueMemory setup readiness check.")], - show_progress_bar=False, - ) - print(" \033[32m✓ Cross-encoder reranker ready\033[0m") - except Exception as e: - model_failures.append("reranker") - print(f" Reranker readiness check failed: {e}", file=sys.stderr) + model_failures = _setup_model_readiness(tier, existing_tier) # ── Step 2: API key for HyDE / deep search ─────────────────────── print() @@ -694,14 +764,22 @@ def _run_setup(args): from getpass import getpass as _getpass api_key = _getpass(f" {prompt_text}: ").strip() - # ── Save config ─────────────────────────────────────────────────── - config["tier"] = tier + # Independent settings cannot publish a requested tier before certification. + from truememory.tier_switch.projection import ( + CONFIG_WRITE_LOCK, ConfigFileLock, _read_config, _write_config, patch_config_fields, + ) + fields = {} if llm_provider: - config["llm_provider"] = llm_provider + fields["llm_provider"] = llm_provider if api_key_field and api_key: - config[api_key_field] = api_key - - _save_truememory_config(config) + fields[api_key_field] = api_key + config = patch_config_fields(fields, config_path=_TRUEMEMORY_CONFIG_PATH) + if tier == "edge" and not existing_tier: + with ConfigFileLock(_TRUEMEMORY_CONFIG_PATH.with_name("config.json.lock"), strict=True), CONFIG_WRITE_LOCK: + config = _read_config(_TRUEMEMORY_CONFIG_PATH) + if "tier" not in config: + config["tier"] = "edge" + _write_config(_TRUEMEMORY_CONFIG_PATH, config) print() print(f" \033[32m✓ Config saved to {_TRUEMEMORY_CONFIG_PATH}\033[0m") @@ -741,75 +819,27 @@ def _run_upgrade_tier(args): """Switch embedding tier and re-embed all memories with the new model.""" tier = args.tier.lower().strip() config = _load_truememory_config() - old_tier = config.get("tier", "edge") - - if tier == old_tier and not args.force: - print(f"Already on {tier} tier. Use --force to re-embed anyway.") - return - - from truememory.tier_switch.cache import get_transition_action + from truememory.tier_switch.manager import RebuildManager - action = get_transition_action(old_tier, tier, force=args.force) + def progress_callback(processed, total, metrics): + pct = 100.0 * processed / max(1, total) + print(f"\r Re-embedding: {pct:.0f}% ({processed}/{total})", end="", flush=True) - if action == "noop": - print(f"Already on {tier} tier.") - return - - if action == "config_only": - config["tier"] = tier - _save_truememory_config(config) - os.environ["TRUEMEMORY_EMBED_MODEL"] = tier - from truememory.vector_search import set_embedding_model - set_embedding_model(tier) - from truememory.reranker import set_active_tier - set_active_tier(tier) - print(f"\033[32m✓ Switched to {tier} instantly (same embedding model)\033[0m") - return - - print(f"Switching from {old_tier} to {tier}...") - - os.environ.pop("HF_HUB_OFFLINE", None) - os.environ.pop("TRANSFORMERS_OFFLINE", None) + manager = RebuildManager() try: - os.environ["TRUEMEMORY_EMBED_MODEL"] = tier - from truememory.vector_search import set_embedding_model - set_embedding_model(tier) - - from truememory.reranker import set_active_tier - set_active_tier(tier) - - from truememory.tier_switch.manager import RebuildManager - - def progress_callback(processed, total, metrics): - pct = metrics.get("progress_pct", 0) - eta = metrics.get("eta_seconds", 0) - bs = metrics.get("batch_size", 0) - ram = metrics.get("ram_pct", 0) - print( - f"\r Re-embedding: {pct:.0f}% ({processed}/{total}) " - f"| ETA {eta:.0f}s | BS={bs} | RAM={ram:.0f}%", - end="", flush=True, - ) - - manager = RebuildManager() - success = manager.run_rebuild_sync( - target_tier=tier, - force=args.force, - progress_callback=progress_callback, - ) - - if success: - print(f"\n \033[32m✓ Tier switched to {tier}\033[0m") - else: - print("\n \033[31m✗ Rebuild failed. Run again to retry.\033[0m") - sys.exit(1) - except Exception as e: - print(f"\n\033[31mError during tier switch: {e}\033[0m", file=sys.stderr) - print("Config was saved. Re-run to retry, or use truememory-ingest setup.") - sys.exit(1) - finally: - os.environ["HF_HUB_OFFLINE"] = "1" - os.environ["TRANSFORMERS_OFFLINE"] = "1" + success = manager.run_rebuild_sync(target_tier=tier, force=args.force, progress_callback=progress_callback) + except Exception as error: + state = "activation pending" if manager._last_outcome == "activation_pending" else "switch incomplete" + print(f"Tier {state}: {error}", file=sys.stderr) + print("Read truememory_status() for rebuild.id and its matching rebuild.job_id. " + "For interrupted work, use truememory-ingest cancel-rebuild --status-id ID --job-id JOB_ID, " + "then retry the original request explicitly.", file=sys.stderr) + raise SystemExit(1) from error + if not success: + print("Tier activation is incomplete. Read truememory_status() before retrying. " + "Use its matching rebuild.id and rebuild.job_id with cancel-rebuild only for interrupted work.", file=sys.stderr) + raise SystemExit(1) + config = _load_truememory_config() _tier_descriptions = { "edge": "Edge — 89.6% LoCoMo, Model2Vec + MiniLM", diff --git a/truememory/ingest/dedup.py b/truememory/ingest/dedup.py index 5404cfcc..c4c50ee7 100644 --- a/truememory/ingest/dedup.py +++ b/truememory/ingest/dedup.py @@ -152,6 +152,8 @@ def check_duplicate( try: results = memory.search_vectors(fact, limit=3) or [] except Exception as e: + from truememory.tier_switch.runtime import raise_if_serving_rejection + raise_if_serving_rejection(e) log.warning("Dedup search failed: %s", e) return DedupDecision(action=DedupAction.ADD, fact=fact, reason="search failed, defaulting to add") diff --git a/truememory/ingest/encoding_gate.py b/truememory/ingest/encoding_gate.py index 00a5984a..14e17a30 100644 --- a/truememory/ingest/encoding_gate.py +++ b/truememory/ingest/encoding_gate.py @@ -232,6 +232,7 @@ def __init__( self._norm = total if total > 0 else 1.0 self._last_search_results: list[dict] = [] self._embed_model = None + self._serving_key = None self._pe_available: bool = True # assume available until first failure self._pe_degradation_count: int = 0 # track PE failures for status (#585) self._batch_scores: list[float] = [] @@ -239,6 +240,16 @@ def __init__( self._batch_saliences: list[float] = [] self._batch_pes: list[float] = [] + def _refresh_serving_generation(self) -> None: + 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 and getattr(self, "_serving_key", None) != operation.key: + self._embed_model = None + self._last_search_results = [] + self._pe_available = True + self._serving_key = operation.key + def evaluate(self, fact: str, category: str = "") -> EncodingDecision: """ Pass a candidate fact through the encoding gate. @@ -252,6 +263,7 @@ def evaluate(self, fact: str, category: str = "") -> EncodingDecision: """ # Contradiction bypass: corrections are critical for memory accuracy # and must never be dropped by the gate (#585). + EncodingGate._refresh_serving_generation(self) is_contradiction = _is_contradiction(fact, category) novelty = self._compute_novelty(fact) @@ -365,7 +377,9 @@ def _compute_novelty(self, fact: str) -> float: if hasattr(self.memory, "search_vectors"): try: results = self.memory.search_vectors(fact, limit=10) - except Exception: + except Exception as _serving_error: + from truememory.tier_switch.runtime import raise_if_serving_rejection + raise_if_serving_rejection(_serving_error) pass if results is None: # search_vectors() fell back to the full pipeline (no vectors). @@ -419,6 +433,8 @@ def _compute_novelty(self, fact: str) -> float: return max(0.05, novelty) except Exception as e: + from truememory.tier_switch.runtime import raise_if_serving_rejection + raise_if_serving_rejection(e) log.warning("Compression novelty failed: %s — returning neutral 0.5", e) return 0.5 @@ -485,6 +501,7 @@ def _compute_prediction_error(self, fact: str) -> float: Independent of novelty (r=0.30) and salience (r=0.23). """ + EncodingGate._refresh_serving_generation(self) if not fact or fact.lower().strip().rstrip("!?.… ") in _PE_NOISE: return 0.0 if len(fact.strip()) < 3: @@ -493,11 +510,18 @@ def _compute_prediction_error(self, fact: str) -> float: if not self._last_search_results: return 0.0 + nearest = self._last_search_results[0] + mem_content = nearest.get("content", "") + if not mem_content: + return 0.0 + if self._embed_model is None: try: from truememory.vector_search import get_model self._embed_model = get_model() except Exception as e: + from truememory.tier_switch.runtime import raise_if_serving_rejection + raise_if_serving_rejection(e) # PE model failed to load — degrade open (#585) if self._pe_available: log.warning( @@ -509,11 +533,6 @@ def _compute_prediction_error(self, fact: str) -> float: return 0.0 model = self._embed_model - nearest = self._last_search_results[0] - mem_content = nearest.get("content", "") - if not mem_content: - return 0.0 - try: from truememory.mps_utils import encode_with_model_ownership embeddings = encode_with_model_ownership( @@ -544,6 +563,8 @@ def _compute_prediction_error(self, fact: str) -> float: return pe except Exception as e: + from truememory.tier_switch.runtime import raise_if_serving_rejection + raise_if_serving_rejection(e) log.debug("PE embedding scorer failed: %s", e) # Runtime encode failure — mark PE as degraded (#585) if self._pe_available: @@ -553,6 +574,10 @@ def _compute_prediction_error(self, fact: str) -> float: self._pe_available = False self._pe_degradation_count += 1 return 0.0 + finally: + # The process singleton remains the cache. A gate must not pin an + # old native model between facts while activation replaces it. + self._embed_model = None # ------------------------------------------------------------------ # Helpers @@ -577,6 +602,8 @@ def _search(self, query: str, limit: int = 5, skip_reranker: bool = False) -> li pass # backing search() lacks the kwarg — fall through return self.memory.search(query, **kwargs) except Exception as e: + from truememory.tier_switch.runtime import raise_if_serving_rejection + raise_if_serving_rejection(e) log.warning("Memory search failed during encoding gate: %s", e) return [] diff --git a/truememory/ingest/pipeline.py b/truememory/ingest/pipeline.py index a10cbe3b..a319abe2 100644 --- a/truememory/ingest/pipeline.py +++ b/truememory/ingest/pipeline.py @@ -466,116 +466,118 @@ def ingest_transcript( # 3-5. Process each fact through gate and dedup for fact in facts: - trace_entry = { - "fact": fact.content, - "category": fact.category, - "confidence": fact.confidence, - } - - # 3. Encoding gate - if self.gate_enabled: - decision = self.gate.evaluate(fact.content, fact.category) - else: - from truememory.ingest.encoding_gate import EncodingDecision - decision = EncodingDecision( - should_encode=True, encoding_score=1.0, - novelty=1.0, salience=1.0, prediction_error=1.0, - reason="gate disabled", - ) - trace_entry["gate"] = { - "passed": decision.should_encode, - "score": decision.encoding_score, - "novelty": decision.novelty, - "salience": decision.salience, - "prediction_error": decision.prediction_error, - "reason": decision.reason, - } - - if not decision.should_encode: - result.facts_skipped_gate += 1 - trace_entry["action"] = "skipped_gate" - result.trace.append(trace_entry) - log.debug("Gate blocked: %s — %s", _safe_log(fact.content[:50]), decision.reason) - continue - - result.facts_encoded += 1 - - # 4-5. Deduplication + storage (atomic critical section). - # - # Holding the process-level lock around BOTH the dedup search - # and the subsequent add/update prevents a TOCTOU race where - # process A checks "no duplicate" for fact X, process B writes - # fact Y (semantically equal), then A writes X as new. Without - # this lock, concurrent Stop hooks from overlapping Claude Code - # sessions accumulate near-duplicate memories. - with _dedup_store_lock(): - dedup = check_duplicate( - fact.content, - self.memory, - user_id=self.user_id, - config=self.llm_config if self.use_llm_dedup else None, - category=fact.category, - ) - trace_entry["dedup"] = { - "action": dedup.action.value, - "reason": dedup.reason, - "existing_id": dedup.existing_id, + from truememory.tier_switch.runtime import fact_operation + with fact_operation(self.memory): + trace_entry = { + "fact": fact.content, + "category": fact.category, + "confidence": fact.confidence, } - # We catch sqlite3.OperationalError around the storage calls - # so a transient DB lock or "unable to open database file" - # condition doesn't abort the whole transcript with a stack - # trace. The pipeline continues with the remaining facts; - # each failed fact is flagged in the trace as - # ``storage_failed`` so operators can diagnose what got - # dropped. See Bug #2 in EDGE_CASE_REPORT.md. - if dedup.action == DedupAction.ADD: - try: - self._store_fact(dedup.fact, fact, session_id) - except sqlite3.OperationalError as e: - log.error( - "Storage failed for fact (db=%s): %s — fact=%r", - getattr(self.memory, "db_path", ""), - e, - _safe_log(dedup.fact[:120]), - ) - trace_entry["action"] = "storage_failed" - trace_entry["storage_error"] = { - "reason": "db locked" if "locked" in str(e).lower() else str(e), - "exception": type(e).__name__, - } - else: - result.facts_stored += 1 - trace_entry["action"] = "stored" - log.info("Stored: %s", _safe_log(dedup.fact[:80])) - - elif dedup.action == DedupAction.UPDATE: - try: - self._update_fact(dedup.existing_id, dedup.fact, fact, session_id) - except sqlite3.OperationalError as e: - log.error( - "Update failed for memory id=%s (db=%s): %s — fact=%r", - dedup.existing_id, - getattr(self.memory, "db_path", ""), - e, - _safe_log(dedup.fact[:120]), - ) - trace_entry["action"] = "storage_failed" - trace_entry["storage_error"] = { - "reason": "db locked" if "locked" in str(e).lower() else str(e), - "exception": type(e).__name__, - } - else: - result.facts_updated += 1 - trace_entry["action"] = "updated" - log.info("Updated [%s]: %s", dedup.existing_id, _safe_log(dedup.fact[:80])) + # 3. Encoding gate + if self.gate_enabled: + decision = self.gate.evaluate(fact.content, fact.category) + else: + from truememory.ingest.encoding_gate import EncodingDecision + decision = EncodingDecision( + should_encode=True, encoding_score=1.0, + novelty=1.0, salience=1.0, prediction_error=1.0, + reason="gate disabled", + ) + trace_entry["gate"] = { + "passed": decision.should_encode, + "score": decision.encoding_score, + "novelty": decision.novelty, + "salience": decision.salience, + "prediction_error": decision.prediction_error, + "reason": decision.reason, + } - elif dedup.action == DedupAction.SKIP: - result.facts_skipped_dedup += 1 - trace_entry["action"] = "skipped_dedup" - log.debug("Dedup skipped: %s — %s", _safe_log(fact.content[:50]), dedup.reason) + if not decision.should_encode: + result.facts_skipped_gate += 1 + trace_entry["action"] = "skipped_gate" + result.trace.append(trace_entry) + log.debug("Gate blocked: %s — %s", _safe_log(fact.content[:50]), decision.reason) + continue + + result.facts_encoded += 1 + + # 4-5. Deduplication + storage (atomic critical section). + # + # Holding the process-level lock around BOTH the dedup search + # and the subsequent add/update prevents a TOCTOU race where + # process A checks "no duplicate" for fact X, process B writes + # fact Y (semantically equal), then A writes X as new. Without + # this lock, concurrent Stop hooks from overlapping Claude Code + # sessions accumulate near-duplicate memories. + with _dedup_store_lock(): + dedup = check_duplicate( + fact.content, + self.memory, + user_id=self.user_id, + config=self.llm_config if self.use_llm_dedup else None, + category=fact.category, + ) + trace_entry["dedup"] = { + "action": dedup.action.value, + "reason": dedup.reason, + "existing_id": dedup.existing_id, + } + + # We catch sqlite3.OperationalError around the storage calls + # so a transient DB lock or "unable to open database file" + # condition doesn't abort the whole transcript with a stack + # trace. The pipeline continues with the remaining facts; + # each failed fact is flagged in the trace as + # ``storage_failed`` so operators can diagnose what got + # dropped. See Bug #2 in EDGE_CASE_REPORT.md. + if dedup.action == DedupAction.ADD: + try: + self._store_fact(dedup.fact, fact, session_id) + except sqlite3.OperationalError as e: + log.error( + "Storage failed for fact (db=%s): %s — fact=%r", + getattr(self.memory, "db_path", ""), + e, + _safe_log(dedup.fact[:120]), + ) + trace_entry["action"] = "storage_failed" + trace_entry["storage_error"] = { + "reason": "db locked" if "locked" in str(e).lower() else str(e), + "exception": type(e).__name__, + } + else: + result.facts_stored += 1 + trace_entry["action"] = "stored" + log.info("Stored: %s", _safe_log(dedup.fact[:80])) + + elif dedup.action == DedupAction.UPDATE: + try: + self._update_fact(dedup.existing_id, dedup.fact, fact, session_id) + except sqlite3.OperationalError as e: + log.error( + "Update failed for memory id=%s (db=%s): %s — fact=%r", + dedup.existing_id, + getattr(self.memory, "db_path", ""), + e, + _safe_log(dedup.fact[:120]), + ) + trace_entry["action"] = "storage_failed" + trace_entry["storage_error"] = { + "reason": "db locked" if "locked" in str(e).lower() else str(e), + "exception": type(e).__name__, + } + else: + result.facts_updated += 1 + trace_entry["action"] = "updated" + log.info("Updated [%s]: %s", dedup.existing_id, _safe_log(dedup.fact[:80])) + + elif dedup.action == DedupAction.SKIP: + result.facts_skipped_dedup += 1 + trace_entry["action"] = "skipped_dedup" + log.debug("Dedup skipped: %s — %s", _safe_log(fact.content[:50]), dedup.reason) - result.trace.append(trace_entry) + result.trace.append(trace_entry) result.elapsed_seconds = round(time.time() - start, 2) log.info( @@ -627,50 +629,52 @@ def ingest_text(self, text: str, session_id: str = "") -> IngestionResult: return result for fact in facts: - if self.gate_enabled: - decision = self.gate.evaluate(fact.content, fact.category) - else: - from truememory.ingest.encoding_gate import EncodingDecision - decision = EncodingDecision( - should_encode=True, encoding_score=1.0, - novelty=1.0, salience=1.0, prediction_error=1.0, - reason="gate disabled", - ) - if not decision.should_encode: - result.facts_skipped_gate += 1 - continue - result.facts_encoded += 1 - - # Hold the process-level lock across the dedup-then-store pair - # so concurrent ingest callers don't race and produce duplicates. - with _dedup_store_lock(): - dedup = check_duplicate( - fact.content, - self.memory, - user_id=self.user_id, - config=self.llm_config if self.use_llm_dedup else None, - category=fact.category, - ) - - if dedup.action == DedupAction.ADD: - try: - self._store_fact(dedup.fact, fact, session_id) - except sqlite3.OperationalError as e: - log.error("Storage failed in ingest_text: %s — fact=%r", e, _safe_log(dedup.fact[:120])) - continue - result.facts_stored += 1 - elif dedup.action == DedupAction.UPDATE: - try: - self._update_fact(dedup.existing_id, dedup.fact, fact, session_id) - except sqlite3.OperationalError as e: - log.error( - "Update failed in ingest_text for memory id=%s: %s — fact=%r", - dedup.existing_id, e, _safe_log(dedup.fact[:120]), - ) - continue - result.facts_updated += 1 + from truememory.tier_switch.runtime import fact_operation + with fact_operation(self.memory): + if self.gate_enabled: + decision = self.gate.evaluate(fact.content, fact.category) else: - result.facts_skipped_dedup += 1 + from truememory.ingest.encoding_gate import EncodingDecision + decision = EncodingDecision( + should_encode=True, encoding_score=1.0, + novelty=1.0, salience=1.0, prediction_error=1.0, + reason="gate disabled", + ) + if not decision.should_encode: + result.facts_skipped_gate += 1 + continue + result.facts_encoded += 1 + + # Hold the process-level lock across the dedup-then-store pair + # so concurrent ingest callers don't race and produce duplicates. + with _dedup_store_lock(): + dedup = check_duplicate( + fact.content, + self.memory, + user_id=self.user_id, + config=self.llm_config if self.use_llm_dedup else None, + category=fact.category, + ) + + if dedup.action == DedupAction.ADD: + try: + self._store_fact(dedup.fact, fact, session_id) + except sqlite3.OperationalError as e: + log.error("Storage failed in ingest_text: %s — fact=%r", e, _safe_log(dedup.fact[:120])) + continue + result.facts_stored += 1 + elif dedup.action == DedupAction.UPDATE: + try: + self._update_fact(dedup.existing_id, dedup.fact, fact, session_id) + except sqlite3.OperationalError as e: + log.error( + "Update failed in ingest_text for memory id=%s: %s — fact=%r", + dedup.existing_id, e, _safe_log(dedup.fact[:120]), + ) + continue + result.facts_updated += 1 + else: + result.facts_skipped_dedup += 1 result.elapsed_seconds = round(time.time() - start, 2) return result @@ -719,6 +723,8 @@ def _store_fact( # the pipeline loop records it in the trace. raise except Exception as e: + from truememory.tier_switch.runtime import raise_if_serving_rejection + raise_if_serving_rejection(e) # Signature mismatch or other recoverable error — fall back # to the client-level API. log.warning("Engine-level add failed: %s, falling back to client", e) @@ -753,6 +759,8 @@ def _update_fact( # loop records a storage_failed entry (Bug #2). raise except Exception as e: + from truememory.tier_switch.runtime import raise_if_serving_rejection + raise_if_serving_rejection(e) log.warning("Failed to update memory %d: %s, storing as new", existing_id, e) # Fallback: store as new if update fails or isn't available diff --git a/truememory/maintenance.py b/truememory/maintenance.py index 5b4f0987..dcc353c4 100644 --- a/truememory/maintenance.py +++ b/truememory/maintenance.py @@ -154,9 +154,41 @@ def connection_database_path(conn: sqlite3.Connection) -> Path | None: _thread_owners = threading.local() +class _AutomaticAdmission(NamedTuple): + coordinator: weakref.ReferenceType + path: Path + process_id: int + descriptor: int + launch: str + released: threading.Event + cancelled: threading.Event + + +def _automatic_admission_locked(path: Path) -> _AutomaticAdmission | None: + """Read only under registry ownership, in registry/coordinator order.""" + coordinator = _coordinators.get(path) + if coordinator is None: + return None + if not coordinator._mutex.acquire(blocking=False): + return None + try: + origin = coordinator._admission_owner + if (origin is not None and origin.coordinator() is coordinator + and origin.process_id == os.getpid() == coordinator._pid + and origin.path == path and _held_paths.get(path) == origin.descriptor + and coordinator._active and not coordinator._cancel.is_set() + and not origin.cancelled.is_set() and not origin.released.is_set()): + return origin + finally: + coordinator._mutex.release() + return None + + def _acquire_owner(path: Path, on_busy: Callable[[], None] | None = None) -> int | None: with _registry_lock: if path in _held_paths: + if hasattr(_thread_owners, "admission_failure"): + _thread_owners.admission_failure = _automatic_admission_locked(path) if on_busy is not None: on_busy() return None @@ -164,6 +196,16 @@ def _acquire_owner(path: Path, on_busy: Callable[[], None] | None = None) -> int if fd is not None: _held_fds.add(fd) _held_paths[path] = fd + launch = getattr(_thread_owners, "automatic_launch", None) + if launch is not None: + coordinator, identity = launch + with coordinator._mutex: + if (coordinator is _coordinators.get(path) and coordinator.path == path + and coordinator._pid == os.getpid() and coordinator._active + and not coordinator._cancel.is_set()): + coordinator._admission_owner = _AutomaticAdmission( + weakref.ref(coordinator), path, os.getpid(), fd, identity, + threading.Event(), threading.Event()) elif on_busy is not None: on_busy() return fd @@ -172,6 +214,16 @@ def _acquire_owner(path: Path, on_busy: Callable[[], None] | None = None) -> int def _release_owner(fd: int) -> None: notify = None with _registry_lock: + origins = [] + for path, claimed_fd in tuple(_held_paths.items()): + if claimed_fd == fd: + coordinator = _coordinators.get(path) + if coordinator is not None: + with coordinator._mutex: + origin = coordinator._admission_owner + if origin is not None and origin.descriptor == fd and origin.process_id == os.getpid(): + coordinator._admission_owner = None + origins.append(origin) os.close(fd) _held_fds.discard(fd) for path, claimed_fd in tuple(_held_paths.items()): @@ -183,6 +235,8 @@ def _release_owner(fd: int) -> None: notify = waiter[2] else: notify = _coordinators.get(path) + for origin in origins: + origin.released.set() if notify is not None: notify._owner_released() @@ -214,9 +268,17 @@ def maintenance_owner(db_path: str | os.PathLike[str] | None) -> Iterator[Mainte if current is not None and current.process_id == os.getpid(): yield current return - fd = _acquire_owner(path) if path is not None else None + previous = getattr(_thread_owners, "admission_failure", None) + _thread_owners.admission_failure = None + try: + fd = _acquire_owner(path) if path is not None else None + failed_owner = _thread_owners.admission_failure + finally: + _thread_owners.admission_failure = previous if path is not None and fd is None: - raise MaintenanceBusyError("Maintenance is already running for this database") + error = MaintenanceBusyError("Maintenance is already running for this database") + error._automatic_owner = failed_owner + raise error owner_pid = os.getpid() try: with _bind_owner(path) as token: @@ -228,6 +290,80 @@ def maintenance_owner(db_path: str | os.PathLike[str] | None) -> Iterator[Mainte _release_owner(fd) +def wait_for_automatic_owner(error: MaintenanceBusyError, *, deadline: float) -> None: + """Wait only for the exact automatic owner observed at failed acquisition.""" + origin = getattr(error, "_automatic_owner", None) + if (type(origin) is not _AutomaticAdmission or origin.process_id != os.getpid() + or origin.cancelled.is_set()): + raise error + if not _registry_lock.acquire(blocking=False): + raise error + try: + if not origin.released.is_set() and _automatic_admission_locked(origin.path) is not origin: + raise error + finally: + _registry_lock.release() + remaining = deadline - time.monotonic() + if remaining <= 0 or not origin.released.wait(min(remaining, threading.TIMEOUT_MAX)): + raise error + if origin.cancelled.is_set() or origin.process_id != os.getpid() or time.monotonic() >= deadline: + raise error + + +def _valid_initialization_receipt(receipt: object, path: Path | None) -> bool: + def bounded(value: object, limit: int) -> bool: + if type(value) is not str or len(value) > limit: + return False + try: + return len(value.encode()) <= limit + except UnicodeError: + return False + + if type(receipt) is not tuple or len(receipt) != 10 or path is None: + return False + if (not bounded(receipt[0], 4096) or receipt[0] != str(path) + or any(type(value) is not int or value < 0 for value in receipt[1:4]) + or type(receipt[4]) is not tuple or len(receipt[4]) != 5 + or any(not bounded(value, 512) for value in receipt[4]) + or receipt[4][0] != "legacy" + or type(receipt[5]) is not tuple or len(receipt[5]) != 2 + or type(receipt[6]) is not tuple or len(receipt[6]) > 3 + or type(receipt[7]) is not bool or type(receipt[8]) is not bool + or type(receipt[9]) is not str or len(receipt[9]) > 32): + return False + tables = {"vec_messages", "vec_messages_sep", "vec_messages_edge", "vec_messages_sep_edge", + "vec_messages_basepro", "vec_messages_sep_basepro"} + if any(type(table) is not str or table not in tables for table in receipt[5]): + return False + for row in receipt[6]: + if (type(row) is not tuple or len(row) != 9 + or any(not bounded(value, 512) for value in row[:3]) + or row[1] not in tables or row[2] not in tables + or (row[3] is not None and not bounded(row[3], 512)) + or any(value is not None and type(value) not in (int, float) for value in row[4:])): + return False + return True + + +@contextmanager +def _active_initialization_receipt(path: Path, *, expected: tuple | None = None) -> Iterator[tuple | None]: + """Guard a live automatic owner's evidence; the body must perform no I/O.""" + with _registry_lock: + coordinator = _coordinators.get(path) + if coordinator is None: + yield None + return + with coordinator._mutex: + record = coordinator._active_initialization + current = (record is not None and coordinator._pid == os.getpid() + and coordinator._active and coordinator._automatic_initialization + and not coordinator._cancel.is_set() and coordinator._thread is not None + and record[0][0] == os.getpid() and _held_paths.get(path) == record[0][1]) + if expected is not None: + current = current and expected[0]() is coordinator and expected[1] == record[0] + yield (weakref.ref(coordinator), record[0], record[1]) if current else None + + def _after_fork() -> None: global _registry_lock, _coordinators, _held_fds, _held_paths, _owner_waiters, _thread_owners # The child must not prolong its parent's ownership by retaining an @@ -238,6 +374,12 @@ def _after_fork() -> None: _held_paths = {} _owner_waiters = {} _thread_owners = threading.local() + for coordinator in tuple(_coordinators.values()): + coordinator._admission_owner = None + coordinator._active_initialization = None + coordinator._pending_initialization = None + coordinator._reserved_initialization = None + coordinator._automatic_initialization = False _coordinators = weakref.WeakValueDictionary() _registry_lock = threading.Lock() @@ -286,6 +428,11 @@ def __init__(self, path: Path | None) -> None: self._capability_versions: tuple[tuple[str, str | None], ...] | None = None self._extension_failure: LayerDependency | None = None self._clustering_warning: tuple | None = None + self._pending_initialization: tuple | None = None + self._reserved_initialization: tuple | None = None + self._active_initialization: tuple | None = None + self._automatic_initialization = False + self._admission_owner: _AutomaticAdmission | None = None def _check_process(self) -> None: if self._pid != os.getpid(): @@ -382,22 +529,30 @@ def _reserve_locked(self) -> None: self._cancel.clear() self._status = "starting" self._error_category = None + self._active_initialization = None + self._reserved_initialization = None + self._automatic_initialization = False def request_layers(self, *, threshold: int = 25, include_layers: bool = True, - include_style: bool = False) -> bool: + include_style: bool = False, _initialization_receipt: tuple | None = None) -> bool: """Coalesce a real foreground wake; no callback captures its engine.""" self._check_process() if type(threshold) is not int or threshold < 1: raise ValueError("Maintenance threshold must be a positive integer") if type(include_layers) is not bool or type(include_style) is not bool or not (include_layers or include_style): raise ValueError("Maintenance request requires explicit work kinds") + valid_receipt = _valid_initialization_receipt(_initialization_receipt, self.path) with self._mutex: if self.path is None: + self._pending_initialization = None self._status = "pending_in_memory" return False if self._active and self._cancel.is_set(): + self._pending_initialization = None return False self._notification_generation += 1 + self._pending_initialization = ( + (self._notification_generation, _initialization_receipt) if valid_receipt else None) self._pending = MaintenanceRequest(self._notification_generation, threshold, include_layers, include_style).merge(self._pending) pending = self._take_pending_locked() @@ -407,16 +562,23 @@ def _take_pending_locked(self) -> MaintenanceRequest | None: if self._active or self._pending is None: return None pending, self._pending = self._pending, None + evidence, self._pending_initialization = self._pending_initialization, None self._reserve_locked() + if evidence is not None and evidence[0] == pending.generation: + self._reserved_initialization = evidence return pending def _launch_layers(self, pending: MaintenanceRequest) -> bool: + with self._mutex: + evidence, self._reserved_initialization = self._reserved_initialization, None + receipt = evidence[1] if evidence is not None and evidence[0] == pending.generation else None def work(conn: sqlite3.Connection, cancel: threading.Event) -> "MaintenanceReport": return run_engine_maintenance(conn, self, threshold=pending.threshold, cancel=cancel, include_layers=pending.include_layers, include_style=pending.include_style, connection_owned=True) + receipt_kwargs = {} if receipt is None else {"_initialization_receipt": receipt} if not pending.include_layers: - return self._launch(work, pending, style_only=True) - return self._launch(work, pending) + return self._launch(work, pending, style_only=True, **receipt_kwargs) + return self._launch(work, pending, **receipt_kwargs) def _owner_released(self) -> None: self._check_process() @@ -442,15 +604,22 @@ def request(self, work: Callable[[sqlite3.Connection, threading.Event], None]) - self._reserve_locked() return self._launch(work) - def _launch(self, work: Callable, pending: MaintenanceRequest | None = None, *, style_only: bool = False) -> bool: + def _launch(self, work: Callable, pending: MaintenanceRequest | None = None, *, style_only: bool = False, + _initialization_receipt: tuple | None = None) -> bool: def busy() -> None: # Registry -> coordinator ordering closes the release/admission # race. No coordinator-mutex holder acquires the registry lock. with self._mutex: self._active = False + self._active_initialization = self._reserved_initialization = None + self._automatic_initialization = False self._status = "cancelled" if self._cancel.is_set() else "busy" if pending is not None and not self._cancel.is_set(): self._pending = pending.merge(self._pending) + if (self._cancel.is_set() or self._pending is None + or (self._pending_initialization is not None + and self._pending_initialization[0] != self._pending.generation)): + self._pending_initialization = None if self._pending is not None and self.path in _held_paths: _owner_waiters[self.path] = (_held_paths[self.path], self._pending[0], self) self._done.set() @@ -463,7 +632,14 @@ def busy() -> None: try: # Serializing descriptor acquisition with fork registration makes # every inherited owner descriptor visible to the child cleanup. - fd = _acquire_owner(self.path, on_busy=busy) + previous_launch = getattr(_thread_owners, "automatic_launch", None) + _thread_owners.automatic_launch = ( + (self, uuid.uuid4().hex) if pending is not None + and _valid_initialization_receipt(_initialization_receipt, self.path) else None) + try: + fd = _acquire_owner(self.path, on_busy=busy) + finally: + _thread_owners.automatic_launch = previous_launch if fd is None: return False launch_token = self._begin_observation() @@ -483,6 +659,10 @@ def run() -> None: with self._mutex: self._thread = worker self._status = "running" + self._automatic_initialization = pending is not None + if (pending is not None and not self._cancel.is_set() + and _valid_initialization_receipt(_initialization_receipt, self.path)): + self._active_initialization = ((os.getpid(), fd, uuid.uuid4().hex), _initialization_receipt) worker.start() return True except BaseException: @@ -496,10 +676,22 @@ def run() -> None: # worker owns cleanup until it really finishes, even on error. with self._mutex: if self._active and self._thread is worker: + if self._admission_owner is not None: + self._admission_owner.cancelled.set() + self._admission_owner.released.set() self._pending = None + self._active_initialization = self._pending_initialization = self._reserved_initialization = None + self._automatic_initialization = False self._cancel.set() raise self._finish_observation(launch_token, None, "failed", "worker_start") + with self._mutex: + self._active_initialization = self._reserved_initialization = None + self._automatic_initialization = False + if (self._cancel.is_set() or self._pending is None + or (self._pending_initialization is not None + and self._pending_initialization[0] != self._pending.generation)): + self._pending_initialization = None if fd is not None: _release_owner(fd) with self._mutex: @@ -507,6 +699,9 @@ def run() -> None: self._thread = None if self._pending is not None and pending is not None: self._pending = self._pending.merge(pending) + if (self._pending_initialization is not None and (self._pending is None + or self._pending_initialization[0] != self._pending.generation)): + self._pending_initialization = None successor = self._take_pending_locked() if successor is None: self._done.set() @@ -528,7 +723,11 @@ def _run(self, fd: int, work: Callable[[sqlite3.Connection, threading.Event], No with _bind_owner(self.path): try: if not self._cancel.is_set(): - conn = _open_style_database(self.path, self._cancel) if style_only else create_db(self.path) + if style_only: + conn = _open_style_database(self.path, self._cancel) + else: + from truememory.tier_switch.runtime import open_serving_connection + conn = open_serving_connection(self.path, create_db) report = work(conn, self._cancel) if conn.in_transaction: conn.rollback() @@ -556,12 +755,16 @@ def _run(self, fd: int, work: Callable[[sqlite3.Connection, threading.Event], No if token is not None: self._finish_observation(token, report if isinstance(report, MaintenanceReport) else None, outcome, error_category) + with self._mutex: + self._active_initialization = None + self._automatic_initialization = False _release_owner(fd) with self._mutex: self._active = False self._thread = None if self._cancel.is_set(): self._pending = None + self._pending_initialization = self._reserved_initialization = None pending = self._take_pending_locked() if pending is None: self._done.set() @@ -575,7 +778,12 @@ def cancel(self) -> None: """Signal a phase boundary; never release ownership of active work.""" self._check_process() with _registry_lock, self._mutex: + if self._admission_owner is not None: + self._admission_owner.cancelled.set() + self._admission_owner.released.set() self._pending = None + self._active_initialization = self._pending_initialization = self._reserved_initialization = None + self._automatic_initialization = False self._cancel.set() waiter = _owner_waiters.get(self.path) if waiter is not None and waiter[2] is self: @@ -1039,6 +1247,24 @@ def _check_layer_dependency(current: LayerDependency, expected: LayerDependency) raise MaintenanceUnavailableError("Layer dependency changed before publication") +@contextmanager +def _maintenance_embedding_scope(conn: sqlite3.Connection, enabled: bool, + cancel: threading.Event | None, *, connection_lock: object | None = None, + deadline: float | None = None) -> Iterator[bool]: + if not enabled or cancel is not None and cancel.is_set(): + yield True + return + from truememory.tier_switch.runtime import maintenance_serving_operation + with maintenance_serving_operation(conn, cancelled=cancel, deadline=deadline, + connection_lock=connection_lock) as supported: + yield supported + + +def _selected_cluster_deferred(conn: sqlite3.Connection) -> LayerResult: + return LayerResult("clusters", "cluster_messages", "deferred", None, 0.0, + "SelectedVectorLayerUnsupported", False, "unverified", conn.in_transaction) + + def run_layers( conn: sqlite3.Connection, specs: tuple[LayerSpec, ...], *, force: bool = False, threshold: int = 25, cancel: threading.Event | None = None, @@ -1055,7 +1281,14 @@ def run_layers( if any(spec.connection is not None and spec.connection is not conn for spec in specs): raise ValueError("Layer adapters belong to another connection") results = [] - with maintenance_owner(connection_database_path(conn)) as owner: + with maintenance_owner(connection_database_path(conn)) as owner, _maintenance_embedding_scope( + conn, any(spec.layer == "clusters" for spec in specs), cancel, + ) as supported: + if not supported: + results.append(_selected_cluster_deferred(conn)) + specs = tuple(spec for spec in specs if spec.layer != "clusters") + if not specs: + return tuple(results) def record_attempt(state: LayerState, outcome: str, category: str | None) -> bool: if _canonical_source: return _record_attempt(conn, state, outcome, owner.generation, category, @@ -1186,6 +1419,10 @@ def confirm_rollback() -> None: except BaseException as error: if isinstance(error, _LayerRollbackFailed): raise + runtime = sys.modules.get("truememory.tier_switch.runtime") + reject = getattr(runtime, "raise_if_serving_rejection", None) + if reject is not None: + reject(error) if isinstance(error, LayerDeferredError): if output_started and not rolled_back: raise _LayerRollbackFailed("Deferred output has no confirmed rollback") from error @@ -2003,6 +2240,9 @@ def _installed_dependency(module: str, distribution: str) -> str | None: def _cluster_schedule_dependency( conn: sqlite3.Connection, *, versions: dict[str, str | None] | None = None, ) -> LayerDependency: + from truememory.tier_switch.runtime import _read_selection, _read_policy + if _read_selection(conn) is not None or _read_policy(conn) is not None: + return LayerDependency("", 1, False, "SelectedVectorLayerUnsupported", deferred=True) parameters = {"min_cluster_size": 10, "min_samples": 5, "metric": "euclidean", "cluster_selection_method": "eom"} if versions is None: @@ -2357,46 +2597,78 @@ def _prepare_worker_extensions(conn: sqlite3.Connection, coordinator: Maintenanc return epoch, versions, failure +@contextmanager +def _maintenance_completion(conn: sqlite3.Connection, allow_caller_transaction: bool) -> Iterator[None]: + """Validate successful owned work before its connection mutation lock leaves.""" + if type(allow_caller_transaction) is not bool: + raise ValueError("Caller transaction opt-in must be a boolean") + borrowed = conn.in_transaction + if borrowed and not allow_caller_transaction: + raise RuntimeError("Maintenance runner requires a clean transaction boundary") + yield + if not borrowed and conn.in_transaction: + conn.rollback() + raise RuntimeError("Maintenance work left an unfinished transaction") + + def run_engine_maintenance( conn: sqlite3.Connection, coordinator: MaintenanceCoordinator, *, threshold: int = 25, force: bool = False, cancel: threading.Event | None = None, prepare_extensions: bool = True, allow_caller_transaction: bool = False, include_layers: bool = True, include_style: bool = False, connection_owned: bool = False, + _connection_lock: object | None = None, deadline: float | None = None, ) -> MaintenanceReport: """Use only the supplied owned worker/original explicitly borrowed handle.""" - style_result = None - if include_style: - style_result = run_routed_style(conn, force=force, threshold=threshold, cancel=cancel, - allow_caller_transaction=allow_caller_transaction, connection_owned=connection_owned) - if not include_layers or (cancel is not None and cancel.is_set()): - return MaintenanceReport((), "CANCELLED" if cancel is not None and cancel.is_set() else - "SKIPPED (style-only)", style_result) - evidence = _prepare_worker_extensions(conn, coordinator) if prepare_extensions else None - specs = engine_layer_specs(conn, coordinator, evidence=evidence) - results = run_layers(conn, specs, threshold=threshold, force=force, cancel=cancel, - allow_caller_transaction=allow_caller_transaction) - for result in results: - if result.layer == "clusters" and result.attempted: - # RELEASE only publishes into the caller's still-rollbackable - # transaction. It cannot clear a previously observed failure. - if not (result.pending_caller_commit and result.outcome in {"success", "success_empty"}): - coordinator.observe_clustering_outcome(result.outcome, result.error_category) - preferences = "SKIPPED (no maintenance attempt)" if cancel is not None and cancel.is_set(): - preferences = "CANCELLED" - elif force or any(result.attempted for result in results): - started = time.monotonic() - try: - importlib.import_module("truememory.personality").extract_preferences(conn) - except (ImportError, AttributeError): - preferences = "UNAVAILABLE (DependencyMissing)" - except Exception as error: - preferences = "ERROR (" + type(error).__name__[:64] + ")" - else: - preferences = f"{time.monotonic() - started:.3f}s" - if conn.in_transaction and allow_caller_transaction: - preferences += " (pending caller commit)" - return MaintenanceReport(results, preferences, style_result) + # Keep this cancellation snapshot terminal even if another caller + # clears the event before the optional style result is constructed. + style = (LayerResult("style_vectors", "style_vectors", "deferred", None, 0.0, + "StyleCancelled", False, "unverified", conn.in_transaction) + if include_style else None) + return MaintenanceReport((), "CANCELLED", style) + from truememory.tier_switch.runtime import _connection_read + with _connection_read(conn, _connection_lock, deadline, cancel): + path = connection_database_path(conn) + with maintenance_owner(path), _maintenance_embedding_scope( + conn, include_layers, cancel, connection_lock=_connection_lock, deadline=deadline, + ) as supported, _connection_read(conn, _connection_lock, deadline, cancel), \ + _maintenance_completion(conn, allow_caller_transaction): + style_result = None + if include_style: + style_result = run_routed_style(conn, force=force, threshold=threshold, cancel=cancel, + allow_caller_transaction=allow_caller_transaction, connection_owned=connection_owned) + if not include_layers or (cancel is not None and cancel.is_set()): + return MaintenanceReport((), "CANCELLED" if cancel is not None and cancel.is_set() else + "SKIPPED (style-only)", style_result) + evidence = _prepare_worker_extensions(conn, coordinator) if prepare_extensions and supported else None + specs = (engine_layer_specs(conn, coordinator, evidence=evidence) if supported else + tuple(spec for spec in all_layer_specs(conn) if spec.layer != "clusters")) + results = run_layers(conn, specs, threshold=threshold, force=force, cancel=cancel, + allow_caller_transaction=allow_caller_transaction) + if not supported: + results = (_selected_cluster_deferred(conn), *results) + for result in results: + if result.layer == "clusters" and result.attempted: + # RELEASE only publishes into the caller's still-rollbackable + # transaction. It cannot clear a previously observed failure. + if not (result.pending_caller_commit and result.outcome in {"success", "success_empty"}): + coordinator.observe_clustering_outcome(result.outcome, result.error_category) + preferences = "SKIPPED (no maintenance attempt)" + if cancel is not None and cancel.is_set(): + preferences = "CANCELLED" + elif force or any(result.attempted for result in results): + started = time.monotonic() + try: + importlib.import_module("truememory.personality").extract_preferences(conn) + except (ImportError, AttributeError): + preferences = "UNAVAILABLE (DependencyMissing)" + except Exception as error: + preferences = "ERROR (" + type(error).__name__[:64] + ")" + else: + preferences = f"{time.monotonic() - started:.3f}s" + if conn.in_transaction and allow_caller_transaction: + preferences += " (pending caller commit)" + return MaintenanceReport(results, preferences, style_result) def maintenance_report_status(report: MaintenanceReport, cancel: threading.Event | None = None) -> tuple[str, str | None]: diff --git a/truememory/mcp_server.py b/truememory/mcp_server.py index 8c12e1e4..9656f048 100644 --- a/truememory/mcp_server.py +++ b/truememory/mcp_server.py @@ -30,6 +30,9 @@ import sys import threading import time +from collections.abc import Iterator +from contextlib import contextmanager +from typing import TYPE_CHECKING try: import resource as _resource_mod @@ -44,6 +47,9 @@ log = logging.getLogger(__name__) +if TYPE_CHECKING: + from truememory.tier_switch.runtime import RuntimeChild, ServingOperation + # --------------------------------------------------------------------------- # Config management # --------------------------------------------------------------------------- @@ -685,7 +691,7 @@ def _get_deepsearch_llm_fn(): return None -def _resolve_deepsearch_llm(): +def _resolve_deepsearch_llm(served_tier: str | None = None): """Resolve the LLM function for deep search queries. Returns the deepsearch-specific override if configured, otherwise @@ -695,7 +701,7 @@ def _resolve_deepsearch_llm(): if ds_fn is not None: return ds_fn # Fall back to default behavior: Pro tier only - if _load_config().get("tier", "edge") == "pro": + if (served_tier if served_tier is not None else _load_config().get("tier", "edge")) == "pro": return _get_llm_fn() return None @@ -788,6 +794,8 @@ def _set_reranker(model_name: str): f"ImportError: {e} — reinstall truememory to restore reranker support" ) except Exception as e: + from truememory.tier_switch.runtime import raise_if_serving_rejection + raise_if_serving_rejection(e) _record_reranker_error(f"{type(e).__name__}: {e}") @@ -943,33 +951,67 @@ def _build_health_payload() -> dict: } -def _parallel_search(queries, user_id, internal_limit, llm_fn, output_limit): +@contextmanager +def _search_operation(*, reranker_id: str | None = None) -> Iterator[tuple[Memory, ServingOperation]]: + """Pin database policy before choosing reranking, HyDE or child searches.""" + from truememory.tier_switch.runtime import engine_serving_operation + + memory = _get_memory() + with engine_serving_operation(memory._engine, reranker_override=reranker_id) as operation: + yield memory, operation + + +def _parallel_search(queries, user_id, internal_limit, llm_fn, output_limit, *, + memory: Memory | None = None, operation: ServingOperation | None = None): """Run multiple agentic searches in parallel, merge and deduplicate.""" - db_path = _get_memory()._engine.db_path - - def _run_query(q): - # context manager ensures the sqlite connection is - # closed even if `KeyboardInterrupt` lands between construction - # and entry into the old explicit `try:` block. `Memory.__exit__` - # already handles close; this form is just interrupt-safe. - with Memory(path=db_path) as thread_m: - return thread_m.search_deep( - q, user_id=user_id, limit=internal_limit, llm_fn=llm_fn, - ) + if operation is None: + with _search_operation() as (memory, operation): + return _parallel_search(queries, user_id, internal_limit, llm_fn, output_limit, + memory=memory, operation=operation) + if memory is None: + raise ValueError("Parallel serving requires its admitted Memory instance") + db_path = memory._engine.db_path + opening_lock = threading.Lock() + + def _run_query(q: str, reservation: RuntimeChild) -> list[dict]: + try: + with Memory(path=db_path) as thread_m: + # Legacy handle creation owns nonblocking maintenance admission. + with opening_lock: + thread_m._engine._open_connection_handle() + with reservation.join(conn=thread_m._engine.conn): + return thread_m.search_deep( + q, user_id=user_id, limit=internal_limit, llm_fn=llm_fn, + ) + finally: + reservation.close() - with ThreadPoolExecutor(max_workers=min(len(queries), 5)) as pool: - futures = [pool.submit(_run_query, q) for q in queries] - merged = [] - seen_ids = set() - for f in futures: - try: - for r in f.result(timeout=60): - rid = r.get("id") - if rid not in seen_ids: - merged.append(r) - seen_ids.add(rid) - except Exception as e: - log.debug("parallel search query failed: %s", e) + reservations = [] + merged = [] + seen_ids = set() + try: + with ThreadPoolExecutor(max_workers=min(len(queries), 5)) as pool: + futures = [] + for q in queries: + child = operation.fork_child() + reservations.append(child) + futures.append(pool.submit(_run_query, q, child)) + for f in futures: + try: + for r in f.result(timeout=60): + rid = r.get("id") + if rid not in seen_ids: + merged.append(r) + seen_ids.add(rid) + except Exception as e: + from truememory.tier_switch.runtime import raise_if_serving_rejection + raise_if_serving_rejection(e) + log.debug("parallel search query failed: %s", e) + finally: + # Executor shutdown has joined running children before releasing any + # reservation whose task was never submitted or admitted. + for child in reservations: + child.close() merged.sort(key=lambda x: -x.get("score", 0)) return merged[:output_limit] @@ -1098,23 +1140,24 @@ def truememory_search( return json.dumps([]) _touch_search_time() limit = max(1, min(limit, 200)) - _set_reranker(_current_reranker()) - llm_fn = _get_llm_fn() if _load_config().get("tier", "edge") == "pro" else None uid = user_id or None if not query_list: return json.dumps([]) if len(query_list) > 10: query_list = query_list[:10] - if len(query_list) == 1: - m = _get_memory() - results = m.search_deep( - query_list[0], user_id=uid, limit=_SEARCH_INTERNAL_LIMIT, llm_fn=llm_fn, - ) - return json.dumps(results[:limit], indent=2) + with _search_operation() as (m, operation): + _set_reranker(operation.reranker_id) + llm_fn = _get_llm_fn() if operation.tier == "pro" else None + if len(query_list) == 1: + results = m.search_deep( + query_list[0], user_id=uid, limit=_SEARCH_INTERNAL_LIMIT, llm_fn=llm_fn, + ) + return json.dumps(results[:limit], indent=2) - results = _parallel_search(query_list, uid, _SEARCH_INTERNAL_LIMIT, llm_fn, limit) - return json.dumps(results, indent=2) + results = _parallel_search(query_list, uid, _SEARCH_INTERNAL_LIMIT, llm_fn, limit, + memory=m, operation=operation) + return json.dumps(results, indent=2) @mcp.tool(meta={"anthropic/alwaysLoad": True}) @@ -1149,23 +1192,24 @@ def truememory_search_deep( return json.dumps([]) _touch_search_time() limit = max(1, min(limit, 200)) - _set_reranker(_DEEP_RERANKER) - llm_fn = _resolve_deepsearch_llm() uid = user_id or None if not query_list: return json.dumps([]) if len(query_list) > 10: query_list = query_list[:10] - if len(query_list) == 1: - m = _get_memory() - results = m.search_deep( - query_list[0], user_id=uid, limit=_DEEP_INTERNAL_LIMIT, llm_fn=llm_fn, - ) - return json.dumps(results[:limit], indent=2) + with _search_operation(reranker_id=_DEEP_RERANKER) as (m, operation): + _set_reranker(operation.reranker_id) + llm_fn = _resolve_deepsearch_llm(operation.tier) + if len(query_list) == 1: + results = m.search_deep( + query_list[0], user_id=uid, limit=_DEEP_INTERNAL_LIMIT, llm_fn=llm_fn, + ) + return json.dumps(results[:limit], indent=2) - results = _parallel_search(query_list, uid, _DEEP_INTERNAL_LIMIT, llm_fn, limit) - return json.dumps(results, indent=2) + results = _parallel_search(query_list, uid, _DEEP_INTERNAL_LIMIT, llm_fn, limit, + memory=m, operation=operation) + return json.dumps(results, indent=2) @mcp.tool() @@ -1292,7 +1336,7 @@ def truememory_configure( "max" = every exchange evaluated regardless of depth. Empty string = no change (preserves current setting). """ - global _memory + global _config_cache, _memory tier = tier.lower().strip() if tier not in ("edge", "base", "pro", "custom"): return json.dumps({"error": "tier must be 'edge', 'base', 'pro', or 'custom'"}) @@ -1332,6 +1376,7 @@ def truememory_configure( # Save to persistent config (tier change deferred to _finalize_rebuild) config = _load_config() + old_config = config.copy() old_tier = config.get("tier", "edge") # Store API key if provided @@ -1360,12 +1405,20 @@ def truememory_configure( except Exception: pass - # Always persist the tier selection — even on first run with no api_key - # or email. Without this, Edge tier on first run (old_tier defaults to - # "edge", so old_tier == tier, skipping the tier-switch block) was never - # written to config.json, leaving setup_required=True forever (#497). - config["tier"] = tier - _save_config(config) + from truememory.tier_switch.projection import patch_config_fields + fields = {key: value for key, value in config.items() + if key not in {"tier", "tier_activation_generation"} + and (key not in old_config or old_config[key] != value)} + config = patch_config_fields(fields, config_path=_CONFIG_PATH, lock_path=_CONFIG_LOCK_PATH) + # Completing the default Edge onboarding changes no serving model space. + if tier == "edge" and "tier" not in old_config: + from truememory.tier_switch.projection import _read_config, _write_config + with _config_file_lock(strict=True), _config_write_lock: + fresh = _read_config(_CONFIG_PATH) + if "tier" not in fresh: + fresh["tier"] = "edge" + _write_config(_CONFIG_PATH, fresh) + _config_cache = None # Issue #645 (M-35): search_intensity is part of the recall cache key # (it changes both the memory limit and the payload budget). A config @@ -1391,58 +1444,38 @@ def truememory_configure( # Clear stored LLM errors for the provider we just re-keyed _clear_llm_error(api_provider) - # Apply model change — temporarily allow downloads for tier switch - # (the new model may not be cached yet). rebuild_error: str | None = None rebuild_action: str | None = None rebuild_status_id: int = 0 - os.environ.pop("HF_HUB_OFFLINE", None) - os.environ.pop("TRANSFORMERS_OFFLINE", None) + switch_status = "configured" + from truememory.tier_switch.manager import RebuildManager + from truememory.tier_switch.cache import get_transition_action + manager = RebuildManager.get_instance() + rebuild_action = get_transition_action(old_tier, tier) try: - os.environ["TRUEMEMORY_EMBED_MODEL"] = tier - from truememory.vector_search import set_embedding_model - set_embedding_model(tier) - - from truememory.reranker import set_active_tier as _set_active_tier - _set_active_tier(tier) - _set_reranker(_current_reranker()) - - # Tier-switch: determine transition action and handle accordingly. - # Base↔Pro = instant (same embedding model). Cross-group = async rebuild. - if old_tier != tier: - try: - from truememory.tier_switch.cache import ( - get_transition_action, - ) - action = get_transition_action(old_tier, tier) - rebuild_action = action - - if action == "config_only": - # Base↔Pro share the same embedding space (qwen3_256), so - # there is nothing to re-embed. The tier selection must - # still be persisted here: unlike the delta_or_full path - # (where RebuildManager writes the tier only once the new - # vectors exist), config_only has no rebuild step, so - # without this write the tier change is lost on restart and - # the runtime/config tiers diverge. - config["tier"] = tier - _save_config(config) - - elif action == "delta_or_full": - from truememory.tier_switch.manager import RebuildManager - manager = RebuildManager.get_instance() - rebuild_status_id = manager.start_rebuild( - target_tier=tier, - ) - except Exception as e: - rebuild_error = f"{type(e).__name__}: {e}" - log.exception("truememory_configure tier-switch failed") - finally: - with _memory_lock: - _memory = None + rebuild_status_id = manager.start_rebuild(target_tier=tier, db_path=Path(_DB_PATH), + config_path=_CONFIG_PATH, config_lock_path=_CONFIG_LOCK_PATH) + if rebuild_status_id: + switch_status = "building" + elif manager._last_outcome == "activation_pending": + switch_status = "activation_pending" + elif manager._last_outcome == "preparing": + switch_status = "preparing" + except Exception as error: + rebuild_error = type(error).__name__ + ": " + str(error) + switch_status = "activation_pending" if manager._last_outcome == "activation_pending" else "switch_failed" + log.exception("truememory_configure tier activation did not complete") finally: os.environ["HF_HUB_OFFLINE"] = "1" os.environ["TRANSFORMERS_OFFLINE"] = "1" + # Existing users of the object retain their reference and serving lease. + # A subsequent request must reopen and reconcile the committed journal. + with _memory_lock: + _memory = None + _config_cache = None + effective_config = _load_config() + served_tier = getattr(manager, "_served_tier", effective_config.get("tier", old_tier)) + effective_tier = served_tier or old_tier _tier_descriptions = { "edge": "Edge: Model2Vec embeddings (8M params), MiniLM reranker", @@ -1451,11 +1484,14 @@ def truememory_configure( "custom": "Custom: user-provided embedding model + reranker via TRUEMEMORY_CUSTOM_* env vars", } result = { - "status": "configured", + "status": switch_status, "tier": tier, - "description": _tier_descriptions.get(tier, f"Tier: {tier}"), + "requested_tier": tier, + "served_tier": served_tier, + "description": (_tier_descriptions.get(effective_tier, f"Tier: {effective_tier}") + if served_tier is not None else "Serving state awaits authoritative readback."), } - if rebuild_action == "config_only": + if rebuild_action == "config_only" and rebuild_error is None and switch_status == "configured": result["note"] = "Tier switched instantly (same embedding model)." elif rebuild_status_id: result["note"] = ( @@ -1463,12 +1499,13 @@ def truememory_configure( f"Query progress with truememory_status({rebuild_status_id})." ) result["status_id"] = rebuild_status_id + elif switch_status == "preparing": + result["note"] = "Tier preparation is in progress; no rebuild status ID has been allocated yet." if rebuild_error is not None: result["rebuild_error"] = rebuild_error result["warning"] = ( - "Tier switch succeeded but memory re-embedding failed. Re-run " - "truememory_configure() to retry, or delete ~/.truememory/memories.db " - "to start fresh." + "Tier activation is incomplete. The committed serving decision remains authoritative; " + "check rebuild status and retry reconciliation." ) if api_key: result["api_key_saved"] = f"{api_provider} key stored" @@ -1488,7 +1525,7 @@ def truememory_configure( or config.get("openrouter_api_key") or config.get("openai_api_key") ) - if tier != "pro": + if effective_tier != "pro": result["hyde_search"] = "disabled (HyDE is Pro-only)" elif has_key: result["hyde_search"] = "enabled" @@ -1500,10 +1537,18 @@ def truememory_configure( _onboarded = Path.home() / ".truememory" / ".onboarded" try: _onboarded.parent.mkdir(parents=True, exist_ok=True) - _onboarded.write_text(f"tier={tier}\n", encoding="utf-8") + _onboarded.write_text(f"tier={effective_tier}\n", encoding="utf-8") except OSError: pass + if switch_status != "configured": + result["next_steps"] = ( + "Configuration is saved. Start a new session to admit the requested tier; runtime readiness remains unconfirmed." + if getattr(manager, "_last_action", None) == "onboarding" else + "The requested tier is not ready. Existing serving state remains guarded until certified activation completes." + ) + return json.dumps(result, indent=2) + # Onboarding: usage examples and next steps result["next_steps"] = ( "TrueMemory is ready. Here's how it works:\n" @@ -1727,19 +1772,17 @@ def _preload_models(): def _load_embedding_model_and_db(): try: - from truememory.vector_search import get_model - get_model() - except Exception: - pass - try: - _get_memory() + with _search_operation(): + from truememory.vector_search import get_model + get_model() except Exception: pass def _load_reranker(): try: - from truememory.reranker import get_reranker - get_reranker(model_name=_current_reranker()) + with _search_operation() as (_, operation): + from truememory.reranker import get_reranker + get_reranker(model_name=operation.reranker_id) except Exception: pass diff --git a/truememory/model_client.py b/truememory/model_client.py index 68494a92..77cb5063 100644 --- a/truememory/model_client.py +++ b/truememory/model_client.py @@ -67,6 +67,10 @@ class ProtocolMismatchError(ConnectionError): """Raised when the server speaks an incompatible/foreign protocol.""" +class ServingIdentityMismatchError(ProtocolMismatchError): + """A selected request was rejected for effective identity or shape drift.""" + + class ModelServerBusyError(RuntimeError): """Capacity rejection: do not restart the daemon or load another model.""" @@ -591,6 +595,8 @@ def _batch_request(request: dict, kwargs: dict) -> dict: def _check_model_response(response: dict, request: dict) -> None: + if response.get("error_code") == "serving_identity_mismatch": + raise ServingIdentityMismatchError("Model server rejected the selected serving identity") if response.get("ok"): return if response.get("error_code") == "server_busy": diff --git a/truememory/model_server.py b/truememory/model_server.py index a4b4f502..54e97e2e 100644 --- a/truememory/model_server.py +++ b/truememory/model_server.py @@ -93,6 +93,12 @@ class _ResultTooLarge(ValueError): """The complete float32 result cannot fit the existing response frame.""" +class _ServingIdentityMismatch(ValueError): + """Explicit selected identity rejection, distinct from native availability.""" + + serving_identity_mismatch = True + + class _ProcessMemoryRefused(RuntimeError): """Current process memory cannot admit another model stage.""" @@ -932,12 +938,12 @@ def _serving_reranker_name(value: object) -> str: 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") + raise _ServingIdentityMismatch("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") + raise _ServingIdentityMismatch("Effective reranker does not match the serving target") def _handle_fast_embed( self, texts: list, tier: str, deadline: _RequestDeadline | None = None, @@ -1219,7 +1225,7 @@ def _handle_request_inner(self, request: dict) -> dict: 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") + raise _ServingIdentityMismatch("Selected reranking requires its explicit model identity") try: batch_limit = _batch_limit(request.get("batch_size", _RERANK_BATCH_LIMIT), _RERANK_BATCH_LIMIT) @@ -1535,7 +1541,16 @@ def _serve_client(self, conn: socket.socket, admitted: _TransportRequest) -> Non except Exception as error: try: conn.settimeout(self._REJECT_TIMEOUT) - self._send_response(conn, {"ok": False, "error": str(error)}) + failure = {"ok": False, "error": str(error)} + identity_error = getattr(type(error), "serving_identity_mismatch", False) is True + if isinstance(request, dict) and ( + "expected_target" in request or request.get("op") in ("prepare_embed_target_v1", "embed_target_v1") + ): + from truememory.embedding_target import EmbeddingTargetError + identity_error = identity_error or isinstance(error, EmbeddingTargetError) + if identity_error: + failure["error_code"] = "serving_identity_mismatch" + self._send_response(conn, failure) except Exception: # The connection is already failing. Keep serialization errors # inside this boundary so their tracebacks do not retain input. diff --git a/truememory/reranker.py b/truememory/reranker.py index 08babe97..35a6d04b 100644 --- a/truememory/reranker.py +++ b/truememory/reranker.py @@ -169,6 +169,23 @@ def get_current_reranker_name() -> str: return get_reranker_name_for_tier(_active_tier) +def apply_reranker_policy(tier: str, model_name: str, *, deadline: float | None = None, + cancelled: threading.Event | None = None) -> None: + """Publish policy identity without loading an absent native model.""" + global _model, _model_name, _model_certified, _active_tier, _frozen_reranker_id + from truememory.tier_switch.runtime import _reranker_id + from truememory.tier_switch.serving import exclusive_activation, _control + from truememory.model_client import CertifiedRerankerProxy, RerankerProxy + _reranker_id(model_name) + control = _control(None, deadline, cancelled) + with exclusive_activation(deadline=deadline, cancelled=cancelled), _reranker_load_lock(control): + if _model_name != model_name: + _model, _model_certified = None, False + elif isinstance(_model, RerankerProxy): + _model, _model_certified = CertifiedRerankerProxy(model_name), True + _model_name, _active_tier, _frozen_reranker_id = model_name, tier, model_name + + def apply_frozen_reranker( tier: str, model_name: str, *, deadline: float | None = None, cancelled: threading.Event | None = None, diff --git a/truememory/tier_switch/manager.py b/truememory/tier_switch/manager.py index b0a57943..c3dfa416 100644 --- a/truememory/tier_switch/manager.py +++ b/truememory/tier_switch/manager.py @@ -13,14 +13,16 @@ import threading import time from pathlib import Path +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from truememory.tier_switch.activation import ActivationIntent, ActivationState, ConfigGuard, LegacyTierPolicy, TierSelection + from truememory.tier_switch.job import TierJob + from truememory.tier_switch.source import TierSourcePlan + from truememory.tier_switch.worker import StatusCallback from truememory.tier_switch.cache import ( - VectorCacheRegistry, - get_messages_to_embed, - model_name_for_group, preflight_ram_check, - resolve_rebuild_action, - tier_group, ) from truememory.tier_switch.throttler import DynamicThrottler from truememory.tier_switch.worker import RebuildWorker @@ -51,33 +53,48 @@ def _detect_device() -> str: def _get_db_path() -> Path: """Return the path to the TrueMemory database.""" - return _TRUEMEMORY_DIR / "memories.db" + return Path(os.environ.get("TRUEMEMORY_DB_PATH") or os.environ.get("TRUEMEMORY_DB") or _TRUEMEMORY_DIR / "memories.db") -def _open_db(db_path: Path | None = None) -> sqlite3.Connection: - """Open a new SQLite connection (thread-safe: each thread gets its own).""" +def _open_db(db_path: Path | None = None, *, initialize: bool = False) -> sqlite3.Connection: + """Validate existing journal storage before any connection or schema writes.""" from truememory.storage import DEFAULT_BUSY_TIMEOUT_MS, _validate_db_path + from truememory.tier_switch.activation import read_activation_state - path = db_path or _get_db_path() - conn = sqlite3.connect(_validate_db_path(path), check_same_thread=False) - conn.execute("PRAGMA journal_mode=WAL") - conn.execute(f"PRAGMA busy_timeout={DEFAULT_BUSY_TIMEOUT_MS}") - conn.execute("PRAGMA foreign_keys=ON") - conn.execute("PRAGMA synchronous=NORMAL") - - # M-23: guard the sqlite-vec load. On macOS system Python (and any build - # compiled without SQLITE_ENABLE_LOAD_EXTENSION) ``enable_load_extension`` - # is absent → bare access raised AttributeError, crashing the tier switch - # (reached from truememory_configure, the CLI upgrade-tier path, and a - # background thread that swallowed it). A tier switch is meaningless - # without vec support, so fail with actionable platform guidance instead. + path = Path(_validate_db_path(db_path or _get_db_path())) + conn = (sqlite3.connect(path, check_same_thread=False) if initialize else + sqlite3.connect(path.absolute().as_uri() + "?mode=rw", uri=True, check_same_thread=False)) + try: + objects = conn.execute( + "SELECT 1 FROM main.sqlite_master WHERE name NOT LIKE 'sqlite_%' LIMIT 1", + ).fetchone() + if objects is None: + if not initialize: + raise TierSwitchUnsupportedError("Database has not been initialized") + conn.close() + from truememory.storage import create_db + conn = create_db(path) + else: + read_activation_state(conn) + if conn.execute("SELECT 1 FROM main.sqlite_master WHERE type='table' AND name='messages'").fetchone() is None: + raise TierSwitchUnsupportedError("Existing source schema cannot be initialized by a tier switch") + conn.execute(f"PRAGMA busy_timeout={DEFAULT_BUSY_TIMEOUT_MS}") + conn.execute("PRAGMA foreign_keys=ON") + conn.execute("PRAGMA synchronous=NORMAL") + _load_tier_extension(conn) + return conn + except BaseException: + conn.close() + raise + + +def _load_tier_extension(conn: sqlite3.Connection) -> None: try: import sqlite_vec conn.enable_load_extension(True) sqlite_vec.load(conn) conn.enable_load_extension(False) except (AttributeError, ImportError, OSError, sqlite3.OperationalError) as exc: - conn.close() raise TierSwitchUnsupportedError( "Tier switching requires the sqlite-vec extension, but this " "Python's sqlite3 cannot load it " @@ -87,8 +104,6 @@ def _open_db(db_path: Path | None = None) -> sqlite3.Connection: "support (e.g. python.org, Homebrew, or pyenv) and retry." ) from exc - return conn - class RebuildManager: """Orchestrates tier-switch re-embedding (singleton for MCP).""" @@ -112,82 +127,85 @@ def __init__(self): # pre-thread failure clears it instead of bricking the manager. self._claimed: bool = False self._active_status_id: int = 0 + self._active_job_id: str | None = None self._state_lock = threading.Lock() + self._cancel_requested = threading.Event() + self._active_db_path: Path | None = None + self._last_outcome: str | None = None + self._requested_tier: str | None = None + self._last_action: str | None = None + self._served_tier: str | None = None + self._config_path: Path = _TRUEMEMORY_DIR / "config.json" + self._config_lock_path: Path | None = None def start_rebuild( - self, - target_tier: str, - force: bool = False, - backup_path: Path | None = None, - db_path: Path | None = None, + self, target_tier: str, force: bool = False, + backup_path: Path | None = None, db_path: Path | None = None, *, + config_path: Path | None = None, config_lock_path: Path | None = None, ) -> int: - """Start an async rebuild in a background thread (MCP path). + """Stage an inactive job, then hand only its immutable identity to a thread.""" + from truememory.maintenance import maintenance_owner - Returns a status_id for progress queries. - """ - # M-51: reserve the slot with a claim flag rather than parking the - # long-lived MCP caller thread in ``_active_thread``. Previously a raise - # in any pre-thread step below (e.g. M-23 in _open_db, preflight) left - # ``_active_thread = current_thread()`` set; since that thread stays - # alive for the process lifetime, ``is_alive()`` was permanently True - # and every subsequent start_rebuild no-op'd — bricking tier switching. - # Only the real worker Thread should occupy the slot; release the claim - # on any pre-thread failure. with self._state_lock: - if self._active_thread and self._active_thread.is_alive(): - return self._active_status_id - if self._claimed: + if self._claimed or self._active_thread and self._active_thread.is_alive(): + if self._requested_tier != target_tier: + raise TierSwitchUnsupportedError("Another requested tier is already being prepared") return self._active_status_id + self._requested_tier = target_tier + self._active_status_id = 0 + self._active_job_id = None + self._last_action = None + self._last_outcome = "preparing" + self._served_tier = None self._claimed = True - + self._cancel_requested.clear() + self._config_path = config_path or _TRUEMEMORY_DIR / "config.json" + self._config_lock_path = config_lock_path + path = Path(db_path) if db_path is not None else _get_db_path() + conn = None + thread = None + self._active_db_path = path try: - conn = _open_db(db_path) - to_group = tier_group(target_tier) - - ok, msg = preflight_ram_check(to_group) - if not ok: - conn.close() - raise RuntimeError(msg) - - action = resolve_rebuild_action(conn, to_group, force) - messages, is_full = get_messages_to_embed(conn, to_group, force) - - if not messages and not is_full: - self._apply_config_switch(target_tier, conn) + if force: + raise TierSwitchUnsupportedError("Active force/reset requires an alternate inactive pair and is not supported") + if self._try_legacy_noop(path, target_tier): + return 0 + with maintenance_owner(path): + if self._bootstrap_empty_configuration(path, target_tier, force): + return 0 + conn = _open_db(path, initialize=True) + prepared = self._prepare_job(conn, target_tier, force, backup_path=backup_path) conn.close() - with self._state_lock: - self._claimed = False + conn = None + if prepared is None: return 0 - - status_id = self._create_status_row( - conn, to_group, target_tier, action, len(messages), - ) - self._active_status_id = status_id - - if backup_path is None and is_full: - backup_path = backup_db(db_path) - - conn.close() - - thread = threading.Thread( - target=self._rebuild_thread, - args=(target_tier, to_group, messages, is_full, status_id, - backup_path, db_path), - daemon=True, - name=f"tier-switch-{to_group}", - ) + job, intent, status_id = prepared + thread = threading.Thread(target=self._rebuild_thread, args=(job, intent, status_id), + daemon=True, name=f"tier-switch-{job.target.tier_group}") with self._state_lock: + self._active_status_id = status_id + self._active_job_id = job.job_id + self._active_db_path = job.database_path self._active_thread = thread - self._claimed = False thread.start() + return status_id except BaseException: - # Pre-thread failure: release the claim so the next call can retry. - # The worker Thread never started, so it never occupied the slot. + if conn is not None: + conn.close() + conn = None + if self._last_outcome == "activation_pending": + self._refresh_committed_serving(path) + else: + self._last_outcome = "interrupted" + if thread is not None and not thread.is_alive(): + with self._state_lock: + self._active_thread = None + raise + finally: + if conn is not None: + conn.close() with self._state_lock: self._claimed = False - raise - - return status_id def run_rebuild_sync( self, @@ -227,178 +245,527 @@ def run_rebuild_sync( lock_fd.close() def _run_rebuild_sync_inner( - self, - target_tier: str, - force: bool = False, - progress_callback=None, + self, target_tier: str, force: bool = False, progress_callback=None, db_path: Path | None = None, ) -> bool: - conn = _open_db(db_path) - to_group = tier_group(target_tier) - - ok, msg = preflight_ram_check(to_group) - if not ok: - conn.close() - log.error("Pre-flight failed: %s", msg) - return False - - action = resolve_rebuild_action(conn, to_group, force) - messages, is_full = get_messages_to_embed(conn, to_group, force) - - if not messages and not is_full: - self._apply_config_switch(target_tier, conn) - conn.close() + from truememory.maintenance import maintenance_owner + + if force: + raise TierSwitchUnsupportedError("Active force/reset requires an alternate inactive pair and is not supported") + path = Path(db_path) if db_path is not None else _get_db_path() + self._cancel_requested.clear() + self._active_db_path = path + if self._try_legacy_noop(path, target_tier): return True + with maintenance_owner(path): + if self._bootstrap_empty_configuration(path, target_tier, force): + return self._last_outcome == "complete" + conn = _open_db(path, initialize=True) + try: + prepared = self._prepare_job(conn, target_tier, force) + if prepared is None: + return True + job, intent, status_id = prepared + self._active_status_id, self._active_db_path = status_id, job.database_path + self._active_job_id = job.job_id + return self._run_job(conn, job, intent, status_id, progress_callback) + except BaseException: + conn.close() + conn = None + if self._last_outcome == "activation_pending": + self._refresh_committed_serving(path) + else: + self._last_outcome = "interrupted" + raise + finally: + self._active_worker = None + if conn is not None: + conn.close() - status_id = self._create_status_row( - conn, to_group, target_tier, action, len(messages), - ) - - if is_full: - backup_db(db_path) - - device = _detect_device() - if to_group == "edge": - device = "cpu" - throttler = DynamicThrottler(device=device) - - worker = RebuildWorker( - conn=conn, - target_tier=target_tier, - target_group=to_group, - throttler=throttler, - status_id=status_id, - status_callback=progress_callback, - ) - self._active_worker = worker - + def _refresh_committed_serving(self, path: Path) -> None: + """Read back on a fresh validated connection after discarding ambiguity.""" + from truememory.tier_switch.activation import read_activation_state + self._served_tier = None try: - success, processed = worker.run(messages, is_full) + conn = _open_db(path) + try: + state = read_activation_state(conn) + result = state.selection or state.legacy_policy + if result is not None: + self._served_tier = result.target.tier + finally: + conn.close() except Exception: - log.exception("Sync rebuild failed") - success = False - - if success: - self._finalize_rebuild(conn, target_tier, to_group) + log.warning("Committed tier readback remains unavailable", exc_info=True) + + @staticmethod + def _apply_bootstrap_identity(target: object, reranker_id: str) -> None: + """Configure only proven-unloaded slots; caller holds exclusive serving.""" + from truememory import vector_search, reranker + from truememory.tier_config import get_tier_config + if vector_search._model is not None or reranker._model is not None: + raise TierSwitchUnsupportedError("First-run identity projection requires unloaded model slots") + if reranker_id != get_tier_config(target.tier)["reranker"]: + raise TierSwitchUnsupportedError("First-run reranker identity must match its tier") + # Onboarding has no database policy to authorize policy-tagged slots. + vector_search.set_embedding_model(target.tier) + reranker.set_active_tier(target.tier) + os.environ["TRUEMEMORY_EMBED_MODEL"] = target.tier + + def _try_legacy_noop(self, path: Path, target_tier: str) -> bool: + """Recognize an unchanged legacy tier without competing for maintenance.""" + try: + path.stat() + except FileNotFoundError: + return False + from truememory import vector_search, reranker + from truememory.embedding_target import EmbeddingTarget + from truememory.tier_config import get_tier_config + from truememory.tier_switch.activation import _guard_storage, _state + from truememory.tier_switch.job import read_selected_job_marker + from truememory.tier_switch.projection import CONFIG_WRITE_LOCK, ConfigFileLock, _guard, _read_config + from truememory.tier_switch.runtime import _legacy_key, serving_operation + from truememory.tier_switch.serving import operation_lease + + target = EmbeddingTarget.capture(target_tier) + if (vector_search._frozen_embedding_target is not None + or getattr(vector_search, "_runtime_policy_tier", None) is not None): + return False + lock_path = self._config_lock_path or self._config_path.with_name(self._config_path.name + ".lock") + # The shared lease permits ordinary maintenance and pins the process + # identity. All database reads use one owned read-only snapshot. + with operation_lease(_legacy_key(), cancelled=self._cancel_requested): + conn = sqlite3.connect(path.absolute().as_uri() + "?mode=ro", uri=True) + try: + conn.execute("BEGIN") + if conn.execute("SELECT 1 FROM main.sqlite_master WHERE name NOT LIKE 'sqlite_%' LIMIT 1").fetchone() is None: + return False + _guard_storage(conn) + state = _state(conn) + if state.selection is not None or state.legacy_policy is not None or state.intent is not None: + return False + if read_selected_job_marker(conn) is not None: + return False + with ConfigFileLock(lock_path, strict=True), CONFIG_WRITE_LOCK: + guard = _guard(_read_config(self._config_path)) + if guard.generation is not None or guard.tier != target.tier: + return False + old_target, _ = self._serving_identity(conn, state, guard) + if old_target != target: + return False + with serving_operation(conn, cancelled=self._cancel_requested) as operation: + if (operation.tier != target.tier + or operation.reranker_id != get_tier_config(target.tier)["reranker"] + or reranker._frozen_reranker_id is not None): + raise TierSwitchUnsupportedError("Legacy runtime identity must be coherent before a tier no-op") + self._served_tier = target.tier + self._last_action = "noop" + self._last_outcome = "complete" + return True + finally: + try: + if conn.in_transaction: + conn.rollback() + finally: + conn.close() + + def _bootstrap_empty_configuration(self, path: Path, target_tier: str, force: bool) -> bool: + """Configure a proven uninitialized store without inventing vector state.""" + from truememory.embedding_target import EmbeddingTarget + from truememory.tier_switch.projection import CONFIG_WRITE_LOCK, ConfigFileLock, _read_config, _write_config + from truememory.tier_switch.serving import exclusive_activation + + if force: + return False + target = EmbeddingTarget.capture(target_tier) + with exclusive_activation(cancelled=self._cancel_requested): + try: + path.stat() + present = True + except FileNotFoundError: + present = False + if present: + conn = sqlite3.connect(path.absolute().as_uri() + "?mode=ro", uri=True) + try: + if conn.execute("SELECT 1 FROM main.sqlite_master WHERE name NOT LIKE 'sqlite_%' LIMIT 1").fetchone(): + return False + finally: + conn.close() + from truememory import vector_search, reranker + from truememory.tier_config import get_tier_config + if (getattr(vector_search, "_frozen_embedding_target", None) is not None + or getattr(vector_search, "_runtime_policy_tier", None) is not None + or getattr(reranker, "_frozen_reranker_id", None) is not None): + raise TierSwitchUnsupportedError("First-run configuration cannot inherit another database's controlled runtime") + served_tier = vector_search.resolve_tier() + observed = EmbeddingTarget.capture(served_tier) + observed_reranker = reranker.get_current_reranker_name() + coherent = ((vector_search.EMBEDDING_MODEL, vector_search._embedding_dim) == (observed.model_id, observed.dimension) + and observed_reranker == get_tier_config(observed.tier)["reranker"]) + self._served_tier = served_tier if coherent else None + ready = coherent and served_tier == target.tier + if not ready and (vector_search._model is not None or reranker._model is not None): + raise TierSwitchUnsupportedError("First-run tier configuration requires unloaded model slots; restart before changing tier") + lock_path = self._config_lock_path or self._config_path.with_name(self._config_path.name + ".lock") + with ConfigFileLock(lock_path, strict=True), CONFIG_WRITE_LOCK: + config = _read_config(self._config_path) + if "tier_activation_generation" in config: + raise TierSwitchUnsupportedError("An activation-tagged config cannot bootstrap a missing store") + self._last_action = "onboarding" + old_env = os.environ.get("TRUEMEMORY_EMBED_MODEL") + if not ready and not coherent: + raise TierSwitchUnsupportedError("First-run legacy identity must be coherent before lazy projection") + requested_config = dict(config, tier=target_tier) + try: + if not ready: + self._last_outcome = "activation_pending" + self._apply_bootstrap_identity(target, get_tier_config(target.tier)["reranker"]) + if ((vector_search.EMBEDDING_MODEL, vector_search._embedding_dim) != (target.model_id, target.dimension) + or vector_search.resolve_tier() != target.tier + or reranker.get_current_reranker_name() != get_tier_config(target.tier)["reranker"]): + raise TierSwitchUnsupportedError("First-run lazy model identity did not reconcile") + self._served_tier = target.tier + _write_config(self._config_path, requested_config) + except BaseException: + # Replace can fail before or after publication. Restore only + # when a locked reread proves the original file remains. + if not ready and _read_config(self._config_path) == config: + self._apply_bootstrap_identity(observed, observed_reranker) + if old_env is None: + os.environ.pop("TRUEMEMORY_EMBED_MODEL", None) + else: + os.environ["TRUEMEMORY_EMBED_MODEL"] = old_env + self._served_tier = observed.tier + raise + self._last_outcome = "complete" + return True - self._active_worker = None - conn.close() - return success + def _prepare_job(self, conn: sqlite3.Connection, target_tier: str, force: bool, *, backup_path: Path | None = None) -> "tuple[TierJob, ActivationIntent, int] | None": + from truememory.embedding_target import EmbeddingTarget + from truememory.rebuild_source import ensure_rebuild_tracking + from truememory.tier_config import get_tier_config + from truememory.tier_switch.activation import ( + commit_config_only_transition, read_activation_state, stage_activation_intent, + ) + from truememory.tier_switch.job import read_selected_job_marker, select_tier_job + from truememory.tier_switch.projection import capture_config_guard + from truememory.tier_switch.serving import exclusive_activation + from truememory.vector_search import init_prepared_target_tables + + self._last_outcome = "preparing" + self._last_action = None + if force: + raise TierSwitchUnsupportedError("Active force/reset requires an alternate inactive pair and is not supported") + target = EmbeddingTarget.capture(target_tier) + reranker_id = get_tier_config(target_tier)["reranker"] + with exclusive_activation(cancelled=self._cancel_requested): + state = read_activation_state(conn) + guard = capture_config_guard(config_path=self._config_path, lock_path=self._config_lock_path) + previous = state.selection or state.legacy_policy + if previous is not None: + self._served_tier = previous.target.tier + if previous is not None and not previous.config_acknowledged: + self._last_action = "config_only" if type(state.intent).__name__ == "PolicyIntent" else "reconcile" + if previous.target != target: + raise TierSwitchUnsupportedError("The committed tier must be reconciled before another switch") + self._project_result(conn, previous, state.intent, policy=state.legacy_policy is not None or type(state.intent).__name__ == "PolicyIntent") + return None + old_target, old_tables = self._serving_identity(conn, state, guard) + self._served_tier = old_target.tier + if old_target == target: + self._last_action = "noop" + if previous is not None: + from truememory.tier_switch.runtime import apply_runtime_policy + self._last_outcome = "activation_pending" + apply_runtime_policy(previous) + self._last_outcome = "complete" + return None + if old_target.identity == target.identity and {old_target.tier, target.tier} == {"base", "pro"}: + self._last_action = "config_only" + self._last_outcome = "activation_pending" + self._served_tier = None + result = commit_config_only_transition( + conn, expected_selection_generation=state.selection.generation if state.selection else None, + expected_intent_id=state.intent.intent_id if state.intent else None, target=target, + reranker_id=reranker_id, expected_config=guard, legacy_pair=old_tables, + ) + self._project_result(conn, result, read_activation_state(conn).intent, policy=True) + return None + if set(old_tables).intersection(target.tables) or old_target.tier_group == target.tier_group: + raise TierSwitchUnsupportedError("Rebuilding the serving pair requires an alternate inactive slot") + ok, message = preflight_ram_check(target.tier_group) + if not ok: + raise RuntimeError(message) + marker = read_selected_job_marker(conn) + if marker is not None and marker.state == "selected": + raise TierSwitchUnsupportedError("A selected job still owns the inactive pair; cancel or reconcile it first") + self._check_inactive_target(conn, target) + self._guard_status(conn) + self._last_action = "rebuild" + ensure_rebuild_tracking(conn) + init_prepared_target_tables(conn, target) + status_id = self._create_status_row(conn, target.tier_group, target.tier, "streamed", 0, + backup_path=backup_path) + self._active_status_id = status_id + job = select_tier_job(conn, target, previous_job_id=marker.job_id if marker else None) + self._active_job_id = job.job_id + intent = stage_activation_intent(conn, job, + expected_generation=state.selection.generation if state.selection else None, + reranker_id=reranker_id, expected_config=guard, status_id=status_id) + return job, intent, status_id + + @staticmethod + def _check_inactive_target(conn: sqlite3.Connection, target: object) -> None: + """Refuse incompatible or unreceipted rows before tracking/table writes.""" + from truememory.tier_switch.activation import _guard_pair + from truememory.tier_switch.source import read_tier_progress + + existing = [] + nonempty = False + for name in target.tables: + if conn.execute("SELECT 1 FROM main.sqlite_master WHERE name=? COLLATE NOCASE", (name,)).fetchone(): + _guard_pair(conn, target, tables=(name,)) + existing.append(name) + nonempty |= conn.execute(f'SELECT 1 FROM main."{name}" LIMIT 1').fetchone() is not None + if nonempty: + if len(existing) != 2 or read_tier_progress(conn, model_id=target.model_id, + dimension=target.dimension, targets=target.tables) is None: + raise TierSwitchUnsupportedError("Nonempty inactive vectors require their exact paired receipt") + + @staticmethod + def _serving_identity(conn: sqlite3.Connection, state: "ActivationState", guard: "ConfigGuard") -> tuple: + from truememory.embedding_target import EmbeddingTarget + from truememory import vector_search + + previous = state.selection or state.legacy_policy + if previous is not None: + if guard.tier != previous.target.tier or guard.generation != previous.generation: + raise TierSwitchUnsupportedError("Config and authoritative serving generation disagree") + return previous.target, previous.tables + target = EmbeddingTarget.capture(guard.tier) + if vector_search.EMBEDDING_MODEL != target.model_id or vector_search._embedding_dim != target.dimension: + raise TierSwitchUnsupportedError("Legacy model configuration must be coherent before tier switching") + tables = (vector_search._active_vec_table(conn), vector_search._active_sep_table(conn)) + if tables not in (target.tables, ("vec_messages", "vec_messages_sep")): + raise TierSwitchUnsupportedError("Legacy serving table pair is ambiguous") + require_identity = target.tier in {"base", "pro", "custom"} + if require_identity or conn.execute("SELECT 1 FROM main.messages LIMIT 1").fetchone() is not None: + for key, expected in (("embed_model", target.model_id), ("embed_dim", str(target.dimension))): + row = conn.execute("SELECT typeof(value)='text' AND value=? FROM main.metadata WHERE key=?", (expected, key)).fetchone() + if row is None or row[0] != 1: + raise TierSwitchUnsupportedError("Legacy source has no coherent serving embedding identity") + return target, tables + + def _run_job(self, conn: sqlite3.Connection, job: "TierJob", intent: "ActivationIntent", status_id: int, + progress_callback: "StatusCallback | None" = None) -> bool: + from truememory.tier_switch.activation import read_activation_state + from truememory.tier_switch.job import check_selected_tier_job, end_selected_tier_job + from truememory.tier_switch.source import initialize_tier_source, plan_tier_source + from truememory.tier_switch.serving import exclusive_activation + + check_selected_tier_job(conn, job) + if read_activation_state(conn).intent != intent: + raise TierSwitchUnsupportedError("Staged activation changed before worker admission") + _load_tier_extension(conn) + device = "cpu" if job.target.tier_group == "edge" else _detect_device() + worker = RebuildWorker(conn, job.target.tier, job.target.tier_group, + DynamicThrottler(device=device), status_id, progress_callback) + with self._state_lock: + self._active_worker = worker + if self._cancel_requested.is_set(): + worker.cancel() + expires = time.monotonic() + 9000 + plan = plan_tier_source(conn, model_id=job.target.model_id, dimension=job.target.dimension, targets=job.target.tables) + if plan.action != "resume" and plan.action != "complete": + plan = initialize_tier_source(conn, plan) + result = worker.run_source(job.target, plan, timeout=max(0.001, expires - time.monotonic())) + if (result.status != "complete" or not result.captured_range_complete + or not result.current_source_complete or self._cancel_requested.is_set()): + outcome = "cancelled" if self._cancel_requested.is_set() or result.status == "cancelled" else "failed" + end_selected_tier_job(conn, job, outcome=outcome) + self._terminal_status(conn, status_id, outcome, result.plan.manifest.consumed, result.plan.manifest.total) + self._last_outcome = outcome + return False + with exclusive_activation(deadline=expires, cancelled=self._cancel_requested): + self._finalize_rebuild(conn, job.target.tier, job.target.tier_group, intent=intent, job=job, plan=result.plan) + self._last_outcome = "activation_pending" + self._terminal_status(conn, status_id, "complete", result.plan.manifest.consumed, result.plan.manifest.total) + self._last_outcome = "complete" + return True + + def _project_result(self, conn: sqlite3.Connection, result: "TierSelection | LegacyTierPolicy", + intent: object, *, policy: bool = False) -> None: + from truememory.tier_switch.activation import acknowledge_config, acknowledge_policy_config, LegacyTierPolicy + from truememory.tier_switch.projection import mirror_activation + from truememory.tier_switch.runtime import apply_frozen_selection, apply_runtime_policy + + self._last_outcome = "activation_pending" + self._served_tier = result.target.tier + mirror_activation(result, expected_config=intent.expected_config, config_path=self._config_path, lock_path=self._config_lock_path) + if policy: + apply_runtime_policy(result) + else: + apply_frozen_selection(result) + if type(result) is LegacyTierPolicy: + acknowledge_policy_config(conn, generation=result.generation) + else: + acknowledge_config(conn, generation=result.generation) + self._served_tier = result.target.tier + self._last_outcome = "complete" + + @staticmethod + def _terminal_status(conn: sqlite3.Connection, status_id: int, status: str, processed: int, total: int) -> None: + from truememory.tier_switch.activation import _transaction + with _transaction(conn, write=True): + RebuildManager._guard_status(conn) + conn.execute("UPDATE rebuild_status SET status=?, processed_messages=?, total_messages=?, progress_pct=?, completed_at=? WHERE id=?", + (status, processed, total, 100.0 if status == "complete" else 100.0 * processed / max(1,total), time.time(), status_id)) def get_status(self, status_id: int = 0) -> dict: - """Query rebuild progress from the rebuild_status table.""" + """Read one journal/marker/manifest snapshot, without corpus audits.""" + from truememory.tier_switch.activation import ActivationIntent, _state, _transaction + from truememory.tier_switch.job import read_selected_job_marker + from truememory.tier_switch.source import read_tier_progress + try: - conn = _open_db() + conn = _open_db(self._active_db_path) except Exception: return {"status": "unknown", "error": "cannot open db"} - try: - sid = status_id or self._active_status_id - if not sid: - row = conn.execute( - "SELECT * FROM rebuild_status ORDER BY id DESC LIMIT 1" - ).fetchone() - else: - row = conn.execute( - "SELECT * FROM rebuild_status WHERE id = ?", (sid,) - ).fetchone() - - if not row: - return {"status": "no_rebuild_found"} - - cols = [ - "id", "tier_group", "target_tier", "status", "action", - "total_messages", "processed_messages", "progress_pct", - "eta_seconds", "batch_size", "throughput_ips", "ram_pct", - "pressure", "error", "started_at", "completed_at", - "backup_path", "last_heartbeat", - ] - return dict(zip(cols, row)) + with _transaction(conn, write=False): + self._guard_status(conn) + sid = status_id or self._active_status_id + columns = ("id", "tier_group", "target_tier", "status", "action", "total_messages", "processed_messages", + "progress_pct", "eta_seconds", "batch_size", "throughput_ips", "ram_pct", "pressure", "error", + "started_at", "completed_at", "backup_path", "last_heartbeat") + text_columns = {"tier_group", "target_tier", "status", "action", "error", "backup_path"} + fields = ", ".join( + f"CASE WHEN typeof({name})='text' AND length(CAST({name} AS BLOB))<=4096 THEN {name} ELSE NULL END" + if name in text_columns else + f"CASE WHEN typeof({name}) IN ('integer','real') THEN {name} ELSE NULL END" + for name in columns + ) + row = conn.execute(f"SELECT {fields} FROM main.rebuild_status WHERE id=?", (sid,)).fetchone() if sid else conn.execute( + f"SELECT {fields} FROM main.rebuild_status ORDER BY id DESC LIMIT 1").fetchone() + if not row: + return {"status": "no_rebuild_found"} + result = dict(zip(columns, row)) + state = _state(conn) + intent = state.intent + if type(intent) is ActivationIntent and intent.status_id == result["id"]: + if (intent.target.tier, intent.target.tier_group) != (result["target_tier"], result["tier_group"]): + raise TierSwitchUnsupportedError("Status row and frozen intent disagree") + if intent.state != "staged": + result["status"] = "complete" if state.selection.config_acknowledged else "activation_pending" + result["activated"] = state.selection.config_acknowledged + else: + marker = read_selected_job_marker(conn) + if marker is None or (marker.job_id, marker.target, marker.tracker, marker.source_epoch, marker.source_schema_signature) != ( + intent.job_id, intent.target, intent.tracker, intent.source_epoch, intent.source_schema_signature): + raise TierSwitchUnsupportedError("Status job no longer owns its source progress") + progress = read_tier_progress(conn, model_id=intent.target.model_id, + dimension=intent.target.dimension, targets=intent.target.tables) + if progress is not None: + if (progress.tracker, progress.source_epoch, progress.source_schema_signature) != ( + intent.tracker, intent.source_epoch, intent.source_schema_signature): + raise TierSwitchUnsupportedError("Progress source belongs to another job") + result.update(processed_messages=progress.processed, total_messages=progress.total, + progress_pct=100.0 * progress.processed / max(1,progress.total)) + result["activated"] = False + result["job_id"] = marker.job_id + result["recovery"] = ( + f"If interrupted, run truememory-ingest cancel-rebuild --status-id {result['id']} " + f"--job-id {marker.job_id}, then retry the original tier request explicitly. " + "Use the same database path. A live maintenance owner prevents cancellation." + ) + if marker.state != "selected": + result["status"] = marker.state + elif self._last_outcome == "interrupted" and result["id"] == self._active_status_id: + result["status"] = "interrupted" + elif result["status"] in {"pending", "running"}: + marker = read_selected_job_marker(conn) + result["status"] = "superseded" if intent is not None else "interrupted" + result["activated"] = False + if marker is not None and marker.state == "selected": + result["selected_job_id"] = marker.job_id + result["recovery"] = "This status does not prove selected marker ownership. Do not use its ID with cancel-rebuild; a matching journal/status binding is required." + return result + except Exception as error: + return {"status": "unknown", "error": type(error).__name__} finally: conn.close() - def cancel(self, status_id: int = 0): - """Signal the active worker to stop.""" - with self._state_lock: - worker = self._active_worker - if worker is not None: - worker.cancel() + def cancel(self, status_id: int = 0, *, expected_job_id: str | None = None, + db_path: Path | None = None) -> dict: + """Signal local work, or explicitly cancel one durable interrupted marker.""" + from truememory.maintenance import maintenance_owner + from truememory.tier_switch.activation import ActivationIntent, _transaction, read_activation_state + from truememory.tier_switch.job import _cancel_policy_job_in_writer, read_selected_job_marker - def _rebuild_thread( - self, - target_tier: str, - to_group: str, - messages: list[dict], - is_full: bool, - status_id: int, - backup_path: Path | None, - db_path: Path | None, - ): - """Background thread target — creates its own DB connection.""" - conn = _open_db(db_path) - try: - device = _detect_device() - if to_group == "edge": - device = "cpu" - throttler = DynamicThrottler(device=device) - - worker = RebuildWorker( - conn=conn, - target_tier=target_tier, - target_group=to_group, - throttler=throttler, - status_id=status_id, - ) - with self._state_lock: - self._active_worker = worker - - success, processed = worker.run(messages, is_full) + with self._state_lock: + active = self._claimed or self._active_thread is not None or self._active_worker is not None + if active: + if expected_job_id is not None and expected_job_id != self._active_job_id: + raise TierSwitchUnsupportedError("Cancellation job does not identify the active job") + if status_id and status_id != self._active_status_id: + raise TierSwitchUnsupportedError("Cancellation status does not identify the active job") + self._cancel_requested.set() + if self._active_worker is not None: + self._active_worker.cancel() + return {"status": "cancellation_requested", "status_id": self._active_status_id} + path = db_path or self._active_db_path or _get_db_path() + with maintenance_owner(path): + conn = _open_db(path) + try: + state = read_activation_state(conn) + marker = read_selected_job_marker(conn) + if marker is None or marker.state != "selected": + return {"status": "no_selected_job"} + intent = state.intent + bound = (type(intent) is ActivationIntent and intent.state == "staged" + and (intent.job_id, intent.target, intent.tracker, intent.source_epoch, intent.source_schema_signature) + == (marker.job_id, marker.target, marker.tracker, marker.source_epoch, marker.source_schema_signature)) + if expected_job_id != marker.job_id or (status_id and (not bound or intent.status_id != status_id)): + raise TierSwitchUnsupportedError("Interrupted cancellation requires the exact current job_id and matching status") + with _transaction(conn, write=True): + _cancel_policy_job_in_writer(conn, job_id=marker.job_id, target=marker.target, + tracker=marker.tracker, source_epoch=marker.source_epoch, + source_schema_signature=marker.source_schema_signature) + return {"status": "cancelled", "job_id": marker.job_id} + finally: + conn.close() - if success: - self._finalize_rebuild(conn, target_tier, to_group) - log.info("Background rebuild complete: tier=%s", target_tier) - else: - log.warning("Background rebuild failed: tier=%s", target_tier) + def _rebuild_thread(self, job: "TierJob", intent: "ActivationIntent", status_id: int) -> None: + """Reopen the exact accepted job under worker-owned maintenance.""" + from truememory.tier_switch.job import open_selected_tier_job + try: + with open_selected_tier_job(job) as conn: + self._run_job(conn, job, intent, status_id) except Exception: - log.exception("Background rebuild thread error") + # Do not issue status SQL on a possibly uncertain connection. + if self._last_outcome == "activation_pending": + self._refresh_committed_serving(job.database_path) + else: + self._last_outcome = "interrupted" + log.exception("Background tier activation did not complete") finally: with self._state_lock: self._active_worker = None self._active_thread = None - conn.close() - def _finalize_rebuild( - self, conn: sqlite3.Connection, target_tier: str, to_group: str, - ): - """Apply config changes after a successful rebuild.""" - from truememory.vector_search import ( - _write_embedder_metadata, - set_embedding_model, - ) + def _finalize_rebuild(self, conn: sqlite3.Connection, target_tier: str, to_group: str, *, + intent: "ActivationIntent | None" = None, job: "TierJob | None" = None, + plan: "TierSourcePlan | None" = None) -> None: + """A paired current-source certificate is the only activation authority.""" + from truememory.tier_switch.activation import commit_certified_selection - vec_table = f"vec_messages_{to_group}" - try: - row = conn.execute( - f"SELECT MAX(rowid), COUNT(*) FROM {vec_table}" - ).fetchone() - last_id = row[0] or 0 if row else 0 - vec_count = row[1] or 0 if row else 0 - except sqlite3.OperationalError: - last_id = 0 - vec_count = 0 - - VectorCacheRegistry.set( - conn, - to_group, - model_name=model_name_for_group(to_group), - last_embedded_id=last_id, - vector_count=vec_count, - ) - - set_embedding_model(target_tier) - _write_embedder_metadata(conn) - self._apply_config_switch(target_tier, conn) + if intent is None or job is None or plan is None or (target_tier, to_group) != (job.target.tier, job.target.tier_group): + raise TierSwitchUnsupportedError("Tier activation requires an exact paired completion certificate") + self._last_outcome = "activation_pending" + self._served_tier = None + result = commit_certified_selection(conn, intent, job, plan) + self._project_result(conn, result, intent) def _apply_config_switch( self, target_tier: str, conn: sqlite3.Connection, @@ -414,6 +781,8 @@ def _apply_config_switch( WHOLE read-modify-write so the tier write composes with — instead of racing — other config writers. """ + if conn is not None: + raise TierSwitchUnsupportedError("Legacy config-only helper cannot activate a database tier") import tempfile from truememory.mcp_server import _config_file_lock @@ -458,18 +827,38 @@ def _create_status_row( target_tier: str, action: str, total: int, + *, backup_path: Path | None = None, ) -> int: """Insert a new rebuild_status row and return its id.""" + from truememory.tier_switch.activation import _transaction now = time.time() - cursor = conn.execute( - "INSERT INTO rebuild_status " - "(tier_group, target_tier, status, action, total_messages, " - "started_at, last_heartbeat) " - "VALUES (?, ?, 'running', ?, ?, ?, ?)", - (to_group, target_tier, action, total, now, now), - ) - conn.commit() - return cursor.lastrowid + with _transaction(conn, write=True): + self._guard_status(conn) + cursor = conn.execute( + "INSERT INTO rebuild_status " + "(tier_group, target_tier, status, action, total_messages, started_at, last_heartbeat, backup_path) " + "VALUES (?, ?, 'running', ?, ?, ?, ?, ?)", + (to_group, target_tier, action, total, now, now, str(backup_path) if backup_path is not None else None), + ) + status_id = cursor.lastrowid + return status_id + + @staticmethod + def _guard_status(conn: sqlite3.Connection) -> None: + from truememory.tier_switch.activation import _normalized, _schema_sql + if conn.execute("SELECT 1 FROM temp.sqlite_master WHERE name='rebuild_status' COLLATE NOCASE LIMIT 1").fetchone(): + raise TierSwitchUnsupportedError("TEMP status shadows are unsupported") + expected = """CREATE TABLE rebuild_status ( + id INTEGER PRIMARY KEY AUTOINCREMENT, tier_group TEXT NOT NULL, target_tier TEXT NOT NULL, + status TEXT NOT NULL DEFAULT 'pending', action TEXT, total_messages INTEGER DEFAULT 0, + processed_messages INTEGER DEFAULT 0, progress_pct REAL DEFAULT 0, eta_seconds REAL DEFAULT 0, + batch_size INTEGER DEFAULT 0, throughput_ips REAL DEFAULT 0, ram_pct REAL DEFAULT 0, + pressure REAL, error TEXT, started_at REAL, completed_at REAL, backup_path TEXT, last_heartbeat REAL + )""" + if _normalized(_schema_sql(conn, "rebuild_status")) != _normalized(expected): + raise TierSwitchUnsupportedError("Unsupported rebuild status schema") + if conn.execute("SELECT 1 FROM main.sqlite_master WHERE type='trigger' AND tbl_name='rebuild_status' COLLATE NOCASE LIMIT 1").fetchone(): + raise TierSwitchUnsupportedError("Status triggers cannot participate in tier publication") def backup_db(db_path: Path | None = None) -> Path: diff --git a/truememory/tier_switch/runtime.py b/truememory/tier_switch/runtime.py index f75310e4..f602ea3a 100644 --- a/truememory/tier_switch/runtime.py +++ b/truememory/tier_switch/runtime.py @@ -9,17 +9,21 @@ import os import re import sqlite3 +import sys 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 typing import TYPE_CHECKING, ParamSpec, TypeVar from truememory.tier_switch import serving from truememory.tier_switch.activation import TierSelection, _guard_storage, _state, read_activation_state +if TYPE_CHECKING: + from truememory.tier_switch.activation import LegacyTierPolicy + P = ParamSpec("P") R = TypeVar("R") @@ -61,6 +65,10 @@ def owns_other_transaction(self, conn: sqlite3.Connection) -> bool: return (owner is not None and owner[0] != threading.get_ident() and owner[1] is conn and not owner[2]) + def owned_by_current_thread(self) -> bool: + owner = self._owner + return owner is not None and owner[0] == threading.get_ident() + def release(self) -> None: owner = self._owner if owner is None or owner[0] != threading.get_ident(): @@ -143,6 +151,79 @@ def _read_selection(conn: sqlite3.Connection) -> TierSelection | None: return read_activation_state(conn).selection +def _read_policy(conn: sqlite3.Connection) -> LegacyTierPolicy | None: + 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='tier_activation_v1'").fetchone(): + return None + if conn.in_transaction: + _guard_storage(conn) + state = _state(conn) + else: + state = read_activation_state(conn) + return getattr(state, "legacy_policy", None) + + +def policy_key(policy: LegacyTierPolicy, reranker_id: str | None = None) -> tuple[str, ...]: + from truememory.tier_switch.activation import LegacyTierPolicy + if type(policy) is not LegacyTierPolicy: + raise TierRuntimeError("Runtime policy requires an immutable journal result") + effective = policy.reranker_id if reranker_id is None else _reranker_id(reranker_id) + return ("policy", policy.generation, policy.target.model_id, str(policy.target.dimension), + policy.target.tier, effective, *policy.tables) + + +def _authority_key(selection: TierSelection | None, policy: LegacyTierPolicy | None) -> tuple[str, ...] | None: + return selection_key(selection) if selection is not None else policy_key(policy) if policy is not None else None + + +def _policy_ready(policy: LegacyTierPolicy, effective: str) -> bool: + vector, reranker = _modules() + return (_acknowledged == policy_key(policy, effective) and _pid == os.getpid() + and vector.EMBEDDING_MODEL == policy.target.model_id and vector._embedding_dim == policy.target.dimension + and vector._frozen_embedding_target is None and vector._runtime_policy_tier == policy.target.tier + and vector._model is not None and reranker._model is not None and reranker._model_certified + and reranker._model_name == effective and reranker._active_tier == policy.target.tier + and reranker._frozen_reranker_id == effective) + + +def apply_runtime_policy(result: TierSelection | LegacyTierPolicy, *, deadline: float | None = None, + cancelled: threading.Event | None = None) -> None: + """Project a same-space policy without encode, probe, or native construction.""" + global _acknowledged, _pid + certified = type(result) is TierSelection + key = selection_key(result) if certified else policy_key(result) + if current_runtime_operation() is not None: + raise TierRuntimeError("Policy projection cannot upgrade a serving operation") + with serving.exclusive_activation(deadline=deadline, cancelled=cancelled): + _acknowledged = None + vector, reranker = _modules() + _check(deadline, cancelled) + vector.apply_embedding_policy(result.target, certified=certified, deadline=deadline) + _check(deadline, cancelled) + reranker.apply_reranker_policy(result.target.tier, result.reranker_id, deadline=deadline, cancelled=cancelled) + _check(deadline, cancelled) + _pid = os.getpid() + if vector._model is not None and reranker._model is not None and reranker._model_certified: + _acknowledged = key + + +def _prepare_legacy_policy(policy: LegacyTierPolicy, effective: str, *, deadline: float | None, + cancelled: threading.Event | None) -> None: + global _acknowledged, _pid + with serving.exclusive_activation(deadline=deadline, cancelled=cancelled): + if _policy_ready(policy, effective): + return + _acknowledged = None + vector, reranker = _modules() + _check(deadline, cancelled) + vector.apply_embedding_policy(policy.target, certified=False, deadline=deadline, load=True) + _check(deadline, cancelled) + reranker.apply_frozen_reranker(policy.target.tier, effective, deadline=deadline, cancelled=cancelled) + _check(deadline, cancelled) + _pid, _acknowledged = os.getpid(), policy_key(policy, effective) + + def _legacy_key(reranker_id: str | None = None) -> tuple[str, ...]: vector, reranker = _modules() return ("legacy", vector.EMBEDDING_MODEL, str(vector._embedding_dim), @@ -208,6 +289,9 @@ class ServingOperation: _lease: serving.OperationLease = field(repr=False, compare=False) _database: tuple[str, object] = field(repr=False, compare=False) borrowed: bool = False + policy: LegacyTierPolicy | None = None + _connection_lock: object | None = field(default=None, repr=False, compare=False) + _defer_maintenance: bool = field(default=False, repr=False, compare=False) def fork_child(self) -> RuntimeChild: if current_operation(self._connection) is not self: @@ -232,12 +316,14 @@ def join(self, conn: sqlite3.Connection | None = None, *, deadline: float | None 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)): + policy = _read_policy(connection) if selected is None else None + if _authority_key(selected, policy) != _authority_key(original.selection, original.policy): 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) + original.borrowed, original.policy, + original._connection_lock if connection is original._connection else None, + original._defer_maintenance) lease.check_admission() with _bound(operation): yield operation @@ -326,7 +412,7 @@ def _connection_read( 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, + reranker_id: str | None = None, _defer_maintenance: bool = False, ) -> Iterator[ServingOperation]: serving._control(None, deadline, cancelled) if reranker_id is not None: @@ -343,6 +429,7 @@ def serving_operation( # writer is released before gate admission or any model preparation. with _connection_read(conn, connection_lock, deadline, cancelled): selection = _read_selection(conn) + policy = _read_policy(conn) if selection is None else None borrowed = conn.in_transaction if selection is not None: effective = selection.reranker_id if reranker_id is None else reranker_id @@ -351,9 +438,16 @@ def serving_operation( raise TierRuntimeError("Runtime reconciliation requires a clean connection") apply_frozen_selection(selection, deadline=deadline, cancelled=cancelled, reranker_id=effective) key = selection_key(selection, effective) + elif policy is not None: + effective = policy.reranker_id if reranker_id is None else reranker_id + if not _policy_ready(policy, effective): + if borrowed: + raise TierRuntimeError("Policy reconciliation requires a clean connection") + _prepare_legacy_policy(policy, effective, deadline=deadline, cancelled=cancelled) + key = policy_key(policy, effective) else: vector, _ = _modules() - if vector._frozen_embedding_target is not None: + if vector._frozen_embedding_target is not None or getattr(vector, "_runtime_policy_tier", None) is not None: raise TierRuntimeError("Legacy database requires explicit coherent runtime selection") if reranker_id is not None: _, reranker = _modules() @@ -361,32 +455,40 @@ def serving_operation( 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) + try: + reranker.get_reranker(model_name=reranker_id, deadline=deadline, cancelled=cancelled) + except Exception as error: + # Legacy retrieval already reports optional reranker unavailability. + raise_if_serving_rejection(error) 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) + observed_policy = _read_policy(conn) if observed is None else None vector, reranker = _modules() - tables = (selection.tables if selection is not None else + tables = (selection.tables if selection is not None else policy.tables if policy 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)): + if _authority_key(observed, observed_policy) != _authority_key(selection, policy): 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 + elif policy is not None: + if not _policy_ready(policy, effective): + raise TierRuntimeError("Process policy changed during runtime admission") + tables, tier = policy.tables, policy.target.tier 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) + database, borrowed, policy, connection_lock, _defer_maintenance) lease.check_admission() with _bound(operation): yield operation @@ -394,10 +496,31 @@ def serving_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") + if operation is not None: + if operation.borrowed: + raise TierRuntimeError("Model loading requires a clean connection") + if operation.selection is not None or operation.policy is not None: + raise TierRuntimeError("Controlled model must be prepared before serving admission") + lock = operation._connection_lock + conn = operation._connection + if isinstance(lock, ConnectionWriteLock): + if lock.owns_other_transaction(conn): + return + if not lock.owned_by_current_thread(): + acquired = lock.acquire(blocking=False) + if acquired: + try: + if conn.in_transaction: + raise TierRuntimeError("Model loading requires a clean connection") + finally: + lock.release() + return + # The owner may have changed between the first observation + # and the nonblocking probe. Never use an earlier SQL flag. + if lock.owns_other_transaction(conn): + return + if conn.in_transaction: + raise TierRuntimeError("Model loading requires a clean connection") def database_operation(function: Callable[P, R]) -> Callable[P, R]: @@ -409,11 +532,279 @@ def wrapped(conn: sqlite3.Connection, *args: object, **kwargs: object) -> R: return wrapped -def engine_operation(function: Callable[P, R]) -> Callable[P, R]: - """Open plainly, then pin the whole Engine call before writer admission.""" +def open_engine_for_operation(engine: object) -> None: + """Bound fresh public admission without waiting inside an admitted call.""" + if engine.conn is not None: + engine._open_connection_handle() + return + from truememory.maintenance import MaintenanceBusyError, wait_for_automatic_owner + from truememory.storage import DEFAULT_BUSY_TIMEOUT_MS + deadline = time.monotonic() + DEFAULT_BUSY_TIMEOUT_MS / 1000 + + def outside_admission() -> None: + lock = engine._write_lock + owned = (lock.owned_by_current_thread() if isinstance(lock, ConnectionWriteLock) + else lock.locked()) + try: + admitted = getattr(_local, "operation", None) is not None or serving.current_thread_admitted() + except (TierRuntimeError, serving.ServingLeaseError): + admitted = True + if owned or admitted: + raise MaintenanceBusyError("Fresh public opening requires an unbound caller") + + if isinstance(engine._write_lock, ConnectionWriteLock) and engine._write_lock.owned_by_current_thread(): + raise MaintenanceBusyError("Fresh public opening requires an unbound writer") + failure = None + while True: + if failure is not None: + outside_admission() + if time.monotonic() >= deadline: + raise failure + try: + engine._open_connection_handle() + return + except MaintenanceBusyError as error: + failure = error + if engine.conn is not None: + raise + outside_admission() + wait_for_automatic_owner(error, deadline=deadline) + + +def _legacy_initialization_is_local(conn: sqlite3.Connection, operation: ServingOperation) -> bool: + """A nested initializer may attach to a stable pair, but cannot migrate it.""" + known = ("vec_messages", "vec_messages_sep", "vec_messages_edge", + "vec_messages_sep_edge", "vec_messages_basepro", "vec_messages_sep_basepro", + "vec_messages_custom", "vec_messages_sep_custom") + if any(table not in known for table in operation.tables): + raise TierRuntimeError("Nested initialization requires a supported captured pair") + rows = dict(conn.execute( + "SELECT name, CASE WHEN length(CAST(sql AS BLOB)) <= 16384 THEN sql END " + "FROM sqlite_master WHERE type='table' AND name IN (?,?,?,?,?,?,?,?)", known).fetchall()) + if not any(table in rows for table in operation.tables) and operation.tables == known[:2]: + return False + for table in operation.tables: + sql = rows.get(table) + dimension = re.search(r"float\[(\d+)\]", sql or "") + if (not sql or not dimension or str(int(dimension.group(1))) != operation.key[2] + or "distance_metric=cosine" not in sql.replace(" ", "").replace("'", "").lower() + or conn.execute("SELECT 1 FROM sqlite_master WHERE name IN (?,?)", + (table + "_cos_stage", table + "_cos_stage_done")).fetchone()): + raise TierRuntimeError("Nested initialization requires completed vector migration") + if operation.tables == known[:2] and conn.execute("SELECT 1 FROM vec_messages LIMIT 1").fetchone(): + raise TierRuntimeError("Initialize legacy vector tables before entering a serving operation") + if conn.execute("SELECT 1 FROM sqlite_master WHERE name='metadata' AND type='table'").fetchone(): + metadata = dict(conn.execute( + "SELECT key, CASE WHEN length(CAST(value AS BLOB)) <= 512 THEN value END " + "FROM metadata WHERE key IN ('embed_model','embed_dim')").fetchall()) + if (("embed_model" in metadata and metadata["embed_model"] != operation.key[1]) + or ("embed_dim" in metadata and metadata["embed_dim"] != str(operation.key[2]))): + raise TierRuntimeError("Nested initialization requires the captured embedding identity") + return True + + +@contextmanager +def engine_serving_operation( + engine: object, *, reranker_override: str | None = None, + _suppress_maintenance: bool = False, +) -> Iterator[ServingOperation]: + """Finish cold initialization before capturing tables for the public body.""" + open_engine_for_operation(engine) + parent = current_operation(engine.conn) + owns_scope = parent is None + cold = not (getattr(engine, "_runtime_initialized", True) or getattr(engine, "ready", False)) + initialized = False + try: + if cold: + with serving_operation(engine.conn, connection_lock=engine._write_lock, + reranker_id=reranker_override, _defer_maintenance=owns_scope) as preparing: + engine._initialize_connection(_suppress_maintenance=True, _preparing=owns_scope) + initialized = True + if parent is not None and current_operation(engine.conn) is not preparing: + raise TierRuntimeError("Initialization changed the captured serving operation") + with serving_operation(engine.conn, connection_lock=engine._write_lock, + reranker_id=reranker_override, _defer_maintenance=owns_scope) as operation: + yield operation + finally: + if owns_scope and not _suppress_maintenance and (initialized or not cold): + if initialized: + engine._maybe_startup_consolidate() + else: + engine._maybe_auto_consolidate() + + +def engine_handle_operation(function: Callable[P, R]) -> Callable[P, R]: + """Pin a handle-only operation without vector initialization or scheduling.""" @wraps(function) def wrapped(engine: object, *args: object, **kwargs: object) -> R: - engine._open_connection_handle() + open_engine_for_operation(engine) with serving_operation(engine.conn, connection_lock=engine._write_lock): return function(engine, *args, **kwargs) return wrapped + + +def engine_operation(function: Callable[P, R]) -> Callable[P, R]: + """Initialize before pinning the public call's immutable table selection.""" + @wraps(function) + def wrapped(engine: object, *args: object, **kwargs: object) -> R: + with engine_serving_operation(engine): + return function(engine, *args, **kwargs) + return wrapped + + +def raise_if_serving_rejection(error: BaseException) -> None: + """Prevent identity/admission refusals from becoming availability fallbacks.""" + classes: list[type[BaseException]] = [TierRuntimeError, serving.ServingLeaseError, serving.ServingLeaseTimeout] + names = { + "truememory.tier_switch.writer": ("WriterSelectionChanged",), + "truememory.maintenance": ("MaintenanceBusyError",), + "truememory.tier_switch.activation": ("TierActivationError",), + "truememory.rebuild_source": ("RebuildSourceChanged",), + "truememory.vector_search": ("VectorPublicationChanged",), + "truememory.embedding_target": ("EmbeddingTargetError",), + "truememory.model_client": ("ProtocolMismatchError", "ModelServerBusyError", "ServingIdentityMismatchError"), + } + for module_name, attributes in names.items(): + module = sys.modules.get(module_name) + if module is not None: + for attribute in attributes: + value = getattr(module, attribute, None) + if isinstance(value, type) and issubclass(value, BaseException): + classes.append(value) + if isinstance(error, tuple(classes)): + raise error + + +def require_legacy_vector_mutation(conn: sqlite3.Connection) -> None: + """Refuse legacy vector DDL/metadata mutation on a selected database.""" + if _read_selection(conn) is not None or _read_policy(conn) is not None: + raise TierRuntimeError("Controlled vector storage requires certified publication") + + +@contextmanager +def fact_operation(memory: object) -> Iterator[ServingOperation | None]: + """Pin gate, dedup and storage for one fact; extraction remains outside.""" + engine = getattr(memory, "_engine", None) + if engine is None or not callable(getattr(engine, "_open_connection_handle", None)): + # Third-party source-only Memory doubles retain their legacy contract. + yield None + return + with engine_serving_operation(engine) as operation: + yield operation + + +@contextmanager +def maintenance_serving_operation( + conn: sqlite3.Connection, *, cancelled: threading.Event | None = None, + deadline: float | None = None, connection_lock: object | None = None, +) -> Iterator[bool]: + """Called only under maintenance ownership; controlled derived vectors defer.""" + serving._control(None, deadline, cancelled) + _check(deadline, cancelled) + current = current_operation(conn) + if current is not None: + if current.selection is not None or current.policy is not None: + yield False + return + else: + with _connection_read(conn, connection_lock, deadline, cancelled): + controlled = _read_selection(conn) is not None or _read_policy(conn) is not None + if controlled: + yield False + return + with serving_operation(conn, cancelled=cancelled, deadline=deadline, connection_lock=connection_lock): + yield True + + +@contextmanager +def destructive_ingest_operation(engine: object) -> Iterator[None]: + """Reject controlled databases before the legacy destructive import path.""" + from pathlib import Path + from truememory.maintenance import maintenance_owner + path = Path(engine.db_path) + with maintenance_owner(None if str(path) == ":memory:" else path), serving.exclusive_activation(): + conn = engine.conn + owned = False + try: + if conn is None and str(path) != ":memory:" and path.exists(): + conn = sqlite3.connect(path.resolve().as_uri() + "?mode=ro", uri=True) + owned = True + if conn is not None: + if conn.in_transaction: + raise TierRuntimeError("Destructive ingestion requires a clean connection") + if _read_selection(conn) is not None or conn.execute( + "SELECT 1 FROM main.sqlite_master WHERE type='table' AND name='metadata' COLLATE NOCASE" + ).fetchone() and conn.execute( + "SELECT 1 FROM main.metadata WHERE key='tier_activation_v1' LIMIT 1" + ).fetchone(): + raise TierRuntimeError("Controlled databases cannot use destructive ingestion") + finally: + if owned: + conn.close() + yield + + +def open_serving_connection(db_path: str | os.PathLike[str], + legacy_factory: Callable[[object], sqlite3.Connection]) -> sqlite3.Connection: + """Resolve one database identity before owning its absence/schema decision.""" + from pathlib import Path + if str(db_path) == ":memory:": + return legacy_factory(db_path) + return _open_serving_connection_at_path(Path(db_path).resolve(), legacy_factory) + + +def _open_serving_connection_at_path( + path: os.PathLike[str], legacy_factory: Callable[[object], sqlite3.Connection], +) -> sqlite3.Connection: + """Open a canonical path; an Engine caller already owns maintenance. + + Same-thread owner reentry does not reacquire the process/file lock. The + caller must pass the same resolved Path used for its outer ownership. + """ + from truememory.maintenance import maintenance_owner + if str(path) == ":memory:": + return legacy_factory(path) + from truememory.storage import DEFAULT_BUSY_TIMEOUT_MS + with maintenance_owner(path): + if not path.exists(): + return legacy_factory(path) + conn = sqlite3.connect(path.as_uri() + "?mode=rw", uri=True, check_same_thread=False) + try: + present = conn.execute( + "SELECT 1 FROM main.sqlite_master WHERE name='metadata' COLLATE NOCASE" + ).fetchone() + controlled = present and conn.execute( + "SELECT 1 FROM main.metadata WHERE key IN ('tier_selected_v1','tier_activation_v1') LIMIT 1" + ).fetchone() + if controlled: + read_activation_state(conn) + for sql in (f"PRAGMA busy_timeout={DEFAULT_BUSY_TIMEOUT_MS}", "PRAGMA foreign_keys=ON", + "PRAGMA synchronous=NORMAL", "PRAGMA cache_size=-64000", "PRAGMA mmap_size=268435456"): + conn.execute(sql) + return conn + except BaseException: + conn.close() + raise + conn.close() + return legacy_factory(path) + + +@contextmanager +def legacy_vector_mutation(conn: sqlite3.Connection) -> Iterator[None]: + """Hold the existing nonblocking maintenance owner across legacy mutation. + + Legacy initialization can already hold a serving read. Owner contention + refuses immediately instead of waiting or upgrading that read lease. + Caller SQL is neither committed nor rolled back by this admission helper. + """ + from truememory.maintenance import connection_database_path, maintenance_owner + with maintenance_owner(connection_database_path(conn)): + require_legacy_vector_mutation(conn) + yield + + +@contextmanager +def legacy_rebuild_operation(conn: sqlite3.Connection) -> Iterator[None]: + """Legacy rebuilds own maintenance before serving or model work.""" + with legacy_vector_mutation(conn), serving_operation(conn): + yield diff --git a/truememory/tier_switch/serving.py b/truememory/tier_switch/serving.py index 08e32876..dacf7a77 100644 --- a/truememory/tier_switch/serving.py +++ b/truememory/tier_switch/serving.py @@ -217,6 +217,12 @@ def _thread_reader(self, thread: threading.Thread) -> _Reader | None: def _writer_thread(self) -> threading.Thread | None: return next(iter(self._writers.values()), None) + def current_thread_admitted(self) -> bool: + """Inspect nesting without waiting behind another gate operation.""" + with self._locked(blocking=False): + thread = threading.current_thread() + return self._thread_reader(thread) is not None or self._writer_thread() is thread + def _check_key(self, key: SelectionKey) -> None: current = next((reader.key for reader in self._readers.values()), None) if current is None: @@ -425,3 +431,8 @@ def exclusive_activation( ) -> AbstractContextManager[None]: """Drain process-wide operations before a caller's activation work.""" return _DEFAULT_GATE.exclusive_activation(timeout=timeout, deadline=deadline, cancelled=cancelled) + + +def current_thread_admitted() -> bool: + """Report a reader/exclusive scope on this thread, or refuse a busy gate.""" + return _DEFAULT_GATE.current_thread_admitted() diff --git a/truememory/tier_switch/writer.py b/truememory/tier_switch/writer.py index e73736e2..d42e1481 100644 --- a/truememory/tier_switch/writer.py +++ b/truememory/tier_switch/writer.py @@ -16,7 +16,7 @@ from typing import TYPE_CHECKING if TYPE_CHECKING: - from truememory.tier_switch.activation import TierSelection + from truememory.tier_switch.activation import TierSelection, LegacyTierPolicy class WriterSelectionChanged(RuntimeError): @@ -27,30 +27,40 @@ class WriterSelectionChanged(RuntimeError): class WriterCapture: connection: sqlite3.Connection = field(repr=False) selection: TierSelection | None + policy: LegacyTierPolicy | None = None def __post_init__(self) -> None: _validate(self.selection) + if self.policy is not None: + from truememory.tier_switch.activation import LegacyTierPolicy + if type(self.policy) is not LegacyTierPolicy or self.selection is not None: + raise WriterSelectionChanged("Writer policy capture must be an immutable journal result") -def _selection(conn: sqlite3.Connection) -> TierSelection | None: +def _authority(conn: sqlite3.Connection) -> tuple[TierSelection | None, LegacyTierPolicy | None]: table = conn.execute( "SELECT type FROM main.sqlite_master WHERE name='metadata' COLLATE NOCASE", ).fetchone() if table is None: - return None + return None, 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 + if conn.execute("SELECT 1 FROM main.metadata WHERE key IN ('tier_selected_v1','tier_activation_v1') LIMIT 1").fetchone() is None: + return None, None from truememory.tier_switch.activation import TierActivationError, _guard_storage, _state try: _guard_storage(conn) - return _state(conn).selection + state = _state(conn) + return state.selection, getattr(state, "legacy_policy", None) except (TierActivationError, sqlite3.Error) as exc: raise WriterSelectionChanged("Selected journal cannot be validated") from exc +def _selection(conn: sqlite3.Connection) -> TierSelection | None: + return _authority(conn)[0] + + def _validate(selection: TierSelection | None) -> None: if selection is not None: from truememory.tier_switch.activation import TierSelection @@ -72,12 +82,12 @@ def capture_writer_selection(conn: sqlite3.Connection) -> WriterCapture: operation = runtime.current_operation(conn) if operation is not None: _validate(operation.selection) - return WriterCapture(conn, operation.selection) + return WriterCapture(conn, operation.selection, getattr(operation, "policy", None)) owned = not conn.in_transaction if owned: conn.execute("BEGIN") try: - return WriterCapture(conn, _selection(conn)) + return WriterCapture(conn, *_authority(conn)) finally: if owned: # One cleanup attempt only, including failure after rollback. @@ -103,26 +113,36 @@ def require_writer_selection( 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: + current, current_policy = _authority(conn) + policy = captured.policy if captured is not None else None + if policy is None: + if current_policy is not None: + raise WriterSelectionChanged("A legacy policy appeared after writer admission") + elif current_policy is None or replace(current_policy, config_acknowledged=policy.config_acknowledged) != policy: + raise WriterSelectionChanged("Legacy policy generation changed; retry the operation") + if selected is None and policy 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: + if selected is not None and (current is None or replace(current, config_acknowledged=selected.config_acknowledged) != selected): raise WriterSelectionChanged("Selected tier generation or descriptor changed; retry the operation") + if policy is not None and current is not None: + raise WriterSelectionChanged("A selected tier appeared after policy admission") + selected = selected if selected is not None else policy 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") + if policy is None: + # 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), diff --git a/truememory/vector_search.py b/truememory/vector_search.py index aeadbe8c..ccd35925 100644 --- a/truememory/vector_search.py +++ b/truememory/vector_search.py @@ -58,6 +58,7 @@ if TYPE_CHECKING: from truememory.embedding_target import EmbeddingTarget + from truememory.tier_switch.writer import WriterCapture logger = logging.getLogger(__name__) @@ -151,6 +152,8 @@ def resolve_tier() -> str: operation = runtime.current_runtime_operation() if runtime is not None else None if operation is not None: return operation.tier + if globals().get("_runtime_policy_tier") is not None: + return _runtime_policy_tier env = os.environ.get("TRUEMEMORY_EMBED_MODEL", "").strip().lower() if env: return env @@ -174,6 +177,7 @@ def resolve_tier() -> str: _lock = threading.Lock() _model_generation = 0 _frozen_embedding_target = None +_runtime_policy_tier: str | None = None def _active_tier_group() -> str: @@ -226,12 +230,13 @@ def set_embedding_model(name: str) -> None: Accepts tier names ("base", "pro") or internal model names. """ - global EMBEDDING_MODEL, _model, _embedding_dim, _model_generation, _frozen_embedding_target + global EMBEDDING_MODEL, _model, _embedding_dim, _model_generation, _frozen_embedding_target, _runtime_policy_tier from truememory.tier_switch.serving import exclusive_activation with exclusive_activation(), _lock: _model_generation += 1 _model = None # Force reload _frozen_embedding_target = None + _runtime_policy_tier = None EMBEDDING_MODEL = _resolve_model_name(name) _embedding_dim = get_embedding_dim(EMBEDDING_MODEL) @@ -399,9 +404,38 @@ def encode(self, texts: str | list[str], timeout: float | None = None, **kwargs: _prepared_target_slot.release() +def apply_embedding_policy(target: EmbeddingTarget, *, certified: bool, deadline: float | None = None, + load: bool = False) -> None: + """Rebind same-space slots; construction is opt-in for ordinary admission.""" + global _model, EMBEDDING_MODEL, _embedding_dim, _model_generation + global _frozen_embedding_target, _runtime_policy_tier + from truememory.embedding_target import EmbeddingTarget + from truememory.model_client import CertifiedEmbeddingProxy, EmbeddingProxy + from truememory.tier_switch.serving import exclusive_activation + if type(target) is not EmbeddingTarget or target.tier not in {"base", "pro"}: + raise ValueError("Policy projection requires an explicit same-space public target") + if type(certified) is not bool or type(load) is not bool: + raise TypeError("Policy projection controls must be boolean") + with exclusive_activation(deadline=deadline): + with _target_state_lock(deadline): + if (EMBEDDING_MODEL, _embedding_dim) != (target.model_id, target.dimension): + _model = None + _model_generation += 1 + elif isinstance(_model, EmbeddingProxy): + _model = CertifiedEmbeddingProxy(target) + EMBEDDING_MODEL, _embedding_dim = target.model_id, target.dimension + _frozen_embedding_target = target if certified else None + _runtime_policy_tier = target.tier + needs_load = _model is None and load + if needs_load: + model = _load_frozen_embedding_target(target, timeout=_target_remaining(deadline)) + with _target_state_lock(deadline): + _model = model + + 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 + global _model, _embedding_dim, EMBEDDING_MODEL, _model_generation, _frozen_embedding_target, _runtime_policy_tier from truememory.embedding_target import EmbeddingTarget from truememory.tier_switch.serving import exclusive_activation if type(target) is not EmbeddingTarget: @@ -409,6 +443,7 @@ def apply_frozen_embedding_target(target: EmbeddingTarget, *, timeout: float | N expires = _target_deadline(timeout) target.check_configuration() with exclusive_activation(deadline=expires): + _runtime_policy_tier = target.tier with _target_state_lock(expires): if (_frozen_embedding_target is not None and _frozen_embedding_target.identity == target.identity and _model is not None): @@ -693,22 +728,73 @@ def _write_embedder_metadata_no_commit(conn: sqlite3.Connection) -> None: metadata in the same transaction as clearing the in-progress marker, so the table is never "trusted" under stale metadata (issue #647). """ - _ensure_metadata_table(conn) - now = datetime.datetime.now(datetime.timezone.utc).isoformat() - conn.execute( - "INSERT OR REPLACE INTO metadata(key, value, updated_at) VALUES (?, ?, ?)", - ("embed_model", EMBEDDING_MODEL, now), - ) - conn.execute( - "INSERT OR REPLACE INTO metadata(key, value, updated_at) VALUES (?, ?, ?)", - ("embed_dim", str(_embedding_dim), now), - ) + from truememory.tier_switch.runtime import current_operation + operation = current_operation(conn) + if operation is not None and (operation.selection is not None or operation.policy is not None): + from truememory.tier_switch.writer import capture_writer_selection, require_writer_selection + require_writer_selection(conn, capture_writer_selection(conn), model_id=EMBEDDING_MODEL, + dimension=_embedding_dim, tables=operation.tables) + return + from truememory.tier_switch.runtime import legacy_vector_mutation + with legacy_vector_mutation(conn): + _ensure_metadata_table(conn) + now = datetime.datetime.now(datetime.timezone.utc).isoformat() + conn.execute( + "INSERT OR REPLACE INTO metadata(key, value, updated_at) VALUES (?, ?, ?)", + ("embed_model", EMBEDDING_MODEL, now), + ) + conn.execute( + "INSERT OR REPLACE INTO metadata(key, value, updated_at) VALUES (?, ?, ?)", + ("embed_dim", str(_embedding_dim), now), + ) def _write_embedder_metadata(conn: sqlite3.Connection) -> None: """Record `(embed_model, embed_dim)` so later opens can detect drift.""" - _write_embedder_metadata_no_commit(conn) - conn.commit() + from truememory.tier_switch.runtime import legacy_vector_mutation + with legacy_vector_mutation(conn): + _write_embedder_metadata_no_commit(conn) + conn.commit() + + +def _write_foreground_embedder_metadata_no_commit( + conn: sqlite3.Connection, identity: tuple[int, str, int], *, selection: WriterCapture, +) -> None: + """Publish identity inside the caller's model fence and database writer. + + Only add/update/embed_single call this after validating their vector target. + Their writer_transaction owns BEGIN IMMEDIATE or a writer-upgraded savepoint; + in_transaction alone cannot establish that ownership. Maintenance continues + independently, while rebuilds and standalone metadata writes keep its owner. + """ + from truememory.tier_switch.runtime import current_operation + from truememory.tier_switch.writer import capture_writer_selection, require_writer_selection + + operation = current_operation(conn) + if operation is None or selection != capture_writer_selection(conn): + raise VectorPublicationChanged("Foreground metadata requires the admitted writer selection") + if identity != (_model_generation, EMBEDDING_MODEL, _embedding_dim): + raise VectorPublicationChanged("Foreground metadata model identity changed") + require_writer_selection(conn, selection, model_id=identity[1], dimension=identity[2], + tables=operation.tables) + if operation.selection is not None or operation.policy is not None: + return + if operation.key[:3] != ("legacy", identity[1], str(identity[2])): + raise VectorPublicationChanged("Foreground metadata disagrees with serving admission") + matches: list[int | None] = [] + for key, value in (("embed_model", identity[1]), ("embed_dim", str(identity[2]))): + row = conn.execute( + "SELECT typeof(value)='text' AND value=? FROM main.metadata WHERE key=?", (value, key), + ).fetchone() + matches.append(None if row is None else row[0]) + if matches not in ([None, None], [1, 1]): + raise VectorPublicationChanged("Foreground metadata is partial or uses another embedding identity") + now = datetime.datetime.now(datetime.timezone.utc).isoformat() + for key, value in (("embed_model", identity[1]), ("embed_dim", str(identity[2]))): + conn.execute( + "INSERT OR REPLACE INTO main.metadata(key, value, updated_at) VALUES (?, ?, ?)", + (key, value, now), + ) # --------------------------------------------------------------------------- @@ -841,42 +927,44 @@ def _check_embedder_compatibility(conn: sqlite3.Connection) -> None: so that explicit re-embed flows (truememory_configure dropping + rebuilding) aren't blocked by stale metadata from the previous embedder. """ - existing_dim = _detect_existing_vec_dim(conn) - if existing_dim is None: - return # fresh DB or dropped-for-re-embed — nothing to protect - - stored_model, _stored_dim = _read_embedder_metadata(conn) - current_dim = _embedding_dim - - if existing_dim != current_dim: - stored_hint = ( - f" (stored embed_model={stored_model!r})" - if stored_model - else " (no metadata — likely a legacy v0.3.0 DB)" - ) - raise TrueMemoryMigrationError( - f"Database has a {existing_dim}d vec_messages table but the " - f"current embedder produces {current_dim}d vectors{stored_hint}. " - f"This commonly happens when upgrading between tiers with " - f"different embedding dimensions (e.g. v0.3.0 Pro @ 1024d → " - f"v0.4.0 Pro @ 256d). " + _migration_hint() - ) + from truememory.tier_switch.runtime import legacy_vector_mutation + with legacy_vector_mutation(conn): + existing_dim = _detect_existing_vec_dim(conn) + if existing_dim is None: + return # fresh DB or dropped-for-re-embed — nothing to protect + + stored_model, _stored_dim = _read_embedder_metadata(conn) + current_dim = _embedding_dim + + if existing_dim != current_dim: + stored_hint = ( + f" (stored embed_model={stored_model!r})" + if stored_model + else " (no metadata — likely a legacy v0.3.0 DB)" + ) + raise TrueMemoryMigrationError( + f"Database has a {existing_dim}d vec_messages table but the " + f"current embedder produces {current_dim}d vectors{stored_hint}. " + f"This commonly happens when upgrading between tiers with " + f"different embedding dimensions (e.g. v0.3.0 Pro @ 1024d → " + f"v0.4.0 Pro @ 256d). " + _migration_hint() + ) - if stored_model is not None and stored_model != EMBEDDING_MODEL: - raise TrueMemoryMigrationError( - f"Database was built with embed_model={stored_model!r}; current " - f"is {EMBEDDING_MODEL!r}. Matching dims ({current_dim}d) would " - f"otherwise mask a silent vector-space mismatch. " + _migration_hint() - ) + if stored_model is not None and stored_model != EMBEDDING_MODEL: + raise TrueMemoryMigrationError( + f"Database was built with embed_model={stored_model!r}; current " + f"is {EMBEDDING_MODEL!r}. Matching dims ({current_dim}d) would " + f"otherwise mask a silent vector-space mismatch. " + _migration_hint() + ) - if stored_model is None: - logger.warning( - "vec_messages exists without embedder metadata (legacy v0.3.0-style " - "DB). Current embedder=%r at %dd. If you have switched embedding " - "models since ingestion, re-embed via truememory_configure() — " - "otherwise new vectors will carry the %r marker going forward.", - EMBEDDING_MODEL, current_dim, EMBEDDING_MODEL, - ) + if stored_model is None: + logger.warning( + "vec_messages exists without embedder metadata (legacy v0.3.0-style " + "DB). Current embedder=%r at %dd. If you have switched embedding " + "models since ingestion, re-embed via truememory_configure() — " + "otherwise new vectors will carry the %r marker going forward.", + EMBEDDING_MODEL, current_dim, EMBEDDING_MODEL, + ) def _check_rebuild_allowed(conn: sqlite3.Connection) -> None: @@ -1006,41 +1094,49 @@ def init_vec_table( raise ValueError("Selected initialization requires the captured table group") _guard_pair(conn, operation.selection.target) return - import sqlite_vec + if operation is not None and operation.policy is not None: + from truememory.tier_switch.activation import _guard_legacy_policy_pair + if tier_group is not None and operation.tables != (f"vec_messages_{tier_group}", f"vec_messages_sep_{tier_group}"): + raise ValueError("Policy initialization requires the captured table pair") + _guard_legacy_policy_pair(conn, operation.policy.target, operation.tables) + return + from truememory.tier_switch.runtime import legacy_vector_mutation + with legacy_vector_mutation(conn): + import sqlite_vec - conn.enable_load_extension(True) - sqlite_vec.load(conn) - conn.enable_load_extension(False) + conn.enable_load_extension(True) + sqlite_vec.load(conn) + conn.enable_load_extension(False) - _ensure_metadata_table(conn) + _ensure_metadata_table(conn) - if tier_group: - if tier_group not in _VALID_GROUPS: - raise ValueError( - f"Invalid tier_group: {tier_group!r}. " - f"Valid: {sorted(_VALID_GROUPS)}" - ) - vec_name = f"vec_messages_{tier_group}" - sep_name = f"vec_messages_sep_{tier_group}" - else: - _check_embedder_compatibility(conn) - vec_name = _active_vec_table(conn) - sep_name = _active_sep_table(conn) + if tier_group: + if tier_group not in _VALID_GROUPS: + raise ValueError( + f"Invalid tier_group: {tier_group!r}. " + f"Valid: {sorted(_VALID_GROUPS)}" + ) + vec_name = f"vec_messages_{tier_group}" + sep_name = f"vec_messages_sep_{tier_group}" + else: + _check_embedder_compatibility(conn) + vec_name = _active_vec_table(conn) + sep_name = _active_sep_table(conn) - dim = _embedding_dim - conn.execute( - f"CREATE VIRTUAL TABLE IF NOT EXISTS {vec_name} " - f"USING vec0({_vec0_column_decl(dim)})" - ) - conn.execute( - f"CREATE VIRTUAL TABLE IF NOT EXISTS {sep_name} " - f"USING vec0({_vec0_column_decl(dim)})" - ) - conn.commit() + dim = _embedding_dim + conn.execute( + f"CREATE VIRTUAL TABLE IF NOT EXISTS {vec_name} " + f"USING vec0({_vec0_column_decl(dim)})" + ) + conn.execute( + f"CREATE VIRTUAL TABLE IF NOT EXISTS {sep_name} " + f"USING vec0({_vec0_column_decl(dim)})" + ) + conn.commit() - # Upgrade any pre-#631 L2-default tables to cosine in place. Idempotent: - # a no-op once the active tables already declare distance_metric=cosine. - migrate_to_cosine_metric(conn) + # Upgrade any pre-#631 L2-default tables to cosine in place. Idempotent: + # a no-op once the active tables already declare distance_metric=cosine. + migrate_to_cosine_metric(conn) # --------------------------------------------------------------------------- @@ -1598,16 +1694,20 @@ def embed_single(conn: sqlite3.Connection, message_id: int, content: str) -> Non sep_normed = _normalize_for_cosine(sep_embedding) if sep_normed is not None: separation = serialize_f32(sep_normed) - except Exception: + except Exception as _serving_error: + from truememory.tier_switch.runtime import raise_if_serving_rejection + raise_if_serving_rejection(_serving_error) logger.warning("Failed to prepare separation vector during incremental publication", exc_info=True) 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: conn.execute(f"INSERT INTO {sep_tbl}(rowid, embedding) VALUES (?, ?)", (message_id, separation)) - except Exception: + except Exception as _serving_error: + from truememory.tier_switch.runtime import raise_if_serving_rejection + raise_if_serving_rejection(_serving_error) logger.warning("Failed to create separation vector during incremental publication", exc_info=True) - _write_embedder_metadata_no_commit(conn) + _write_foreground_embedder_metadata_no_commit(conn, identity, selection=write_selection) @database_operation @@ -1727,61 +1827,63 @@ def _rebuild_table_as_cosine(conn: sqlite3.Connection, table_name: str) -> int: Returns the number of vectors carried over. """ - dim = _detect_existing_vec_dim(conn, table_name) - if dim is None: - return 0 - - stage = f"{table_name}_cos_stage" - stage_done = f"{stage}_done" - fmt = f"{dim}f" - - # Phase 1: copy normalized vectors into a durable plain staging table. - # D1-2 crash-safety: a separate "done" marker row is written in the SAME - # commit as the staged rows. An interrupted stage (crash mid-copy) leaves an - # empty/partial stage with NO done marker, so resume can tell it apart from a - # complete stage and re-stage from the still-intact original instead of - # swapping an empty stage over it and losing every vector. - conn.execute(f"DROP TABLE IF EXISTS {stage}") - conn.execute(f"DROP TABLE IF EXISTS {stage_done}") - conn.execute( - f"CREATE TABLE {stage} (rowid INTEGER PRIMARY KEY, embedding BLOB)" - ) - conn.execute(f"CREATE TABLE {stage_done} (done INTEGER PRIMARY KEY)") - rows = conn.execute( - f"SELECT rowid, embedding FROM {table_name}" - ).fetchall() - skipped = 0 - staged = 0 - for rowid, blob in rows: - if blob is None or len(blob) != dim * 4: - skipped += 1 - continue - normed = _normalize_for_cosine(struct.unpack(fmt, blob)) - if normed is None: - skipped += 1 - continue + from truememory.tier_switch.runtime import legacy_vector_mutation + with legacy_vector_mutation(conn): + dim = _detect_existing_vec_dim(conn, table_name) + if dim is None: + return 0 + + stage = f"{table_name}_cos_stage" + stage_done = f"{stage}_done" + fmt = f"{dim}f" + + # Phase 1: copy normalized vectors into a durable plain staging table. + # D1-2 crash-safety: a separate "done" marker row is written in the SAME + # commit as the staged rows. An interrupted stage (crash mid-copy) leaves an + # empty/partial stage with NO done marker, so resume can tell it apart from a + # complete stage and re-stage from the still-intact original instead of + # swapping an empty stage over it and losing every vector. + conn.execute(f"DROP TABLE IF EXISTS {stage}") + conn.execute(f"DROP TABLE IF EXISTS {stage_done}") conn.execute( - f"INSERT INTO {stage}(rowid, embedding) VALUES (?, ?)", - (rowid, serialize_f32(normed)), + f"CREATE TABLE {stage} (rowid INTEGER PRIMARY KEY, embedding BLOB)" ) - staged += 1 - # Mark staging complete in the SAME transaction as the staged rows, so the - # marker is durable iff every row is. - conn.execute(f"INSERT INTO {stage_done}(done) VALUES (1)") - conn.commit() # staged rows + completeness marker durable before we touch the original - - # Phase 2: replace the vec0 table with a cosine one and refill it. - carried = _finish_cosine_swap(conn, table_name, dim) - - if skipped: - logger.warning( - "Cosine migration for %s dropped %d zero-norm/invalid vector(s)", - table_name, skipped, + conn.execute(f"CREATE TABLE {stage_done} (done INTEGER PRIMARY KEY)") + rows = conn.execute( + f"SELECT rowid, embedding FROM {table_name}" + ).fetchall() + skipped = 0 + staged = 0 + for rowid, blob in rows: + if blob is None or len(blob) != dim * 4: + skipped += 1 + continue + normed = _normalize_for_cosine(struct.unpack(fmt, blob)) + if normed is None: + skipped += 1 + continue + conn.execute( + f"INSERT INTO {stage}(rowid, embedding) VALUES (?, ?)", + (rowid, serialize_f32(normed)), + ) + staged += 1 + # Mark staging complete in the SAME transaction as the staged rows, so the + # marker is durable iff every row is. + conn.execute(f"INSERT INTO {stage_done}(done) VALUES (1)") + conn.commit() # staged rows + completeness marker durable before we touch the original + + # Phase 2: replace the vec0 table with a cosine one and refill it. + carried = _finish_cosine_swap(conn, table_name, dim) + + if skipped: + logger.warning( + "Cosine migration for %s dropped %d zero-norm/invalid vector(s)", + table_name, skipped, + ) + logger.info( + "Rebuilt %s as cosine (%d vectors carried over)", table_name, carried ) - logger.info( - "Rebuilt %s as cosine (%d vectors carried over)", table_name, carried - ) - return carried + return carried def _finish_cosine_swap( @@ -1792,24 +1894,26 @@ def _finish_cosine_swap( Assumes ``{table_name}_cos_stage`` exists and is populated. Idempotent / resumable: safe to call after a crash that left the staging table behind. """ - stage = f"{table_name}_cos_stage" - conn.execute(f"DROP TABLE IF EXISTS {table_name}") - conn.execute( - f"CREATE VIRTUAL TABLE {table_name} USING vec0({_vec0_column_decl(dim)})" - ) - carried = 0 - staged_rows = conn.execute( - f"SELECT rowid, embedding FROM {stage}" - ).fetchall() - if staged_rows: - conn.executemany( - f"INSERT INTO {table_name}(rowid, embedding) VALUES (?, ?)", - staged_rows, + from truememory.tier_switch.runtime import legacy_vector_mutation + with legacy_vector_mutation(conn): + stage = f"{table_name}_cos_stage" + conn.execute(f"DROP TABLE IF EXISTS {table_name}") + conn.execute( + f"CREATE VIRTUAL TABLE {table_name} USING vec0({_vec0_column_decl(dim)})" ) - carried = len(staged_rows) - conn.execute(f"DROP TABLE IF EXISTS {stage}") - conn.execute(f"DROP TABLE IF EXISTS {stage}_done") - return carried + carried = 0 + staged_rows = conn.execute( + f"SELECT rowid, embedding FROM {stage}" + ).fetchall() + if staged_rows: + conn.executemany( + f"INSERT INTO {table_name}(rowid, embedding) VALUES (?, ?)", + staged_rows, + ) + carried = len(staged_rows) + conn.execute(f"DROP TABLE IF EXISTS {stage}") + conn.execute(f"DROP TABLE IF EXISTS {stage}_done") + return carried def migrate_to_cosine_metric(conn: sqlite3.Connection) -> bool: @@ -1826,67 +1930,69 @@ def migrate_to_cosine_metric(conn: sqlite3.Connection) -> bool: """ # Ensure the extension is loaded (callers from engine.open already load it, # but init_vec_table is the common entry and loads it just above). - try: - import sqlite_vec - conn.enable_load_extension(True) - sqlite_vec.load(conn) - conn.enable_load_extension(False) - except Exception: - # If the extension cannot load, there are no vec0 tables to migrate. - return False + from truememory.tier_switch.runtime import legacy_vector_mutation + with legacy_vector_mutation(conn): + try: + import sqlite_vec + conn.enable_load_extension(True) + sqlite_vec.load(conn) + conn.enable_load_extension(False) + except Exception: + # If the extension cannot load, there are no vec0 tables to migrate. + return False - migrated = False - for table in _KNOWN_VEC_TABLES: - stage = f"{table}_cos_stage" - has_stage = conn.execute( - "SELECT 1 FROM sqlite_master WHERE name = ? AND type='table'", - (stage,), - ).fetchone() - if has_stage: - # A prior run was interrupted between staging and swap. Only finish - # the swap if the staging completed durably (the "done" marker is - # present); otherwise the stage is empty/partial and the ORIGINAL - # table is still intact, so drop the partial stage and re-stage from - # scratch below rather than swapping an empty stage over the original - # and losing every vector (D1-2). - stage_done = f"{stage}_done" - stage_complete = bool( - conn.execute( - "SELECT 1 FROM sqlite_master WHERE name = ? AND type='table'", - (stage_done,), - ).fetchone() - and conn.execute(f"SELECT 1 FROM {stage_done} LIMIT 1").fetchone() - ) - if stage_complete: - dim = ( - _detect_existing_vec_dim(conn, table) - or _detect_existing_vec_dim(conn, stage) - or _embedding_dim + migrated = False + for table in _KNOWN_VEC_TABLES: + stage = f"{table}_cos_stage" + has_stage = conn.execute( + "SELECT 1 FROM sqlite_master WHERE name = ? AND type='table'", + (stage,), + ).fetchone() + if has_stage: + # A prior run was interrupted between staging and swap. Only finish + # the swap if the staging completed durably (the "done" marker is + # present); otherwise the stage is empty/partial and the ORIGINAL + # table is still intact, so drop the partial stage and re-stage from + # scratch below rather than swapping an empty stage over the original + # and losing every vector (D1-2). + stage_done = f"{stage}_done" + stage_complete = bool( + conn.execute( + "SELECT 1 FROM sqlite_master WHERE name = ? AND type='table'", + (stage_done,), + ).fetchone() + and conn.execute(f"SELECT 1 FROM {stage_done} LIMIT 1").fetchone() ) - _finish_cosine_swap(conn, table, dim) + if stage_complete: + dim = ( + _detect_existing_vec_dim(conn, table) + or _detect_existing_vec_dim(conn, stage) + or _embedding_dim + ) + _finish_cosine_swap(conn, table, dim) + conn.commit() + migrated = True + continue + # Incomplete stage from an interrupted copy — discard it and fall + # through to a fresh rebuild from the intact original table. + conn.execute(f"DROP TABLE IF EXISTS {stage}") + conn.execute(f"DROP TABLE IF EXISTS {stage_done}") conn.commit() - migrated = True - continue - # Incomplete stage from an interrupted copy — discard it and fall - # through to a fresh rebuild from the intact original table. - conn.execute(f"DROP TABLE IF EXISTS {stage}") - conn.execute(f"DROP TABLE IF EXISTS {stage_done}") - conn.commit() - exists = conn.execute( - "SELECT 1 FROM sqlite_master WHERE name = ? AND type='table'", - (table,), - ).fetchone() - if not exists: - continue - if _table_uses_cosine(conn, table): - continue - _rebuild_table_as_cosine(conn, table) - migrated = True + exists = conn.execute( + "SELECT 1 FROM sqlite_master WHERE name = ? AND type='table'", + (table,), + ).fetchone() + if not exists: + continue + if _table_uses_cosine(conn, table): + continue + _rebuild_table_as_cosine(conn, table) + migrated = True - if migrated: - conn.commit() - return migrated + if migrated: + conn.commit() + return migrated # --------------------------------------------------------------------------- @@ -1906,82 +2012,84 @@ def migrate_legacy_vec_tables(conn: sqlite3.Connection) -> bool: Returns True if migration occurred, False if nothing to migrate. """ - from truememory.tier_switch.cache import VectorCacheRegistry - - has_old = conn.execute( - "SELECT name FROM sqlite_master WHERE name='vec_messages' AND type='table'" - ).fetchone() - if not has_old: - return False + from truememory.tier_switch.runtime import legacy_vector_mutation + with legacy_vector_mutation(conn): + from truememory.tier_switch.cache import VectorCacheRegistry - stored_model, _ = _read_embedder_metadata(conn) - group = _MODEL_TO_GROUP.get(stored_model, "edge") - if group not in _VALID_GROUPS: - group = "edge" - - new_vec = f"vec_messages_{group}" - new_sep = f"vec_messages_sep_{group}" - - count = conn.execute("SELECT COUNT(*) FROM vec_messages").fetchone()[0] - if count == 0: - has_tiered = conn.execute( - "SELECT name FROM sqlite_master WHERE name=? AND type='table'", - (new_vec,), + has_old = conn.execute( + "SELECT name FROM sqlite_master WHERE name='vec_messages' AND type='table'" ).fetchone() - if not has_tiered: + if not has_old: return False - conn.execute("DROP TABLE IF EXISTS vec_messages") - conn.execute("DROP TABLE IF EXISTS vec_messages_sep") - conn.commit() - return True - dim = _detect_existing_vec_dim(conn, "vec_messages") or 256 + stored_model, _ = _read_embedder_metadata(conn) + group = _MODEL_TO_GROUP.get(stored_model, "edge") + if group not in _VALID_GROUPS: + group = "edge" - import sqlite_vec + new_vec = f"vec_messages_{group}" + new_sep = f"vec_messages_sep_{group}" - conn.enable_load_extension(True) - sqlite_vec.load(conn) - conn.enable_load_extension(False) + count = conn.execute("SELECT COUNT(*) FROM vec_messages").fetchone()[0] + if count == 0: + has_tiered = conn.execute( + "SELECT name FROM sqlite_master WHERE name=? AND type='table'", + (new_vec,), + ).fetchone() + if not has_tiered: + return False + conn.execute("DROP TABLE IF EXISTS vec_messages") + conn.execute("DROP TABLE IF EXISTS vec_messages_sep") + conn.commit() + return True - conn.execute( - f"CREATE VIRTUAL TABLE IF NOT EXISTS {new_vec} " - f"USING vec0({_vec0_column_decl(dim)})" - ) - conn.execute( - f"CREATE VIRTUAL TABLE IF NOT EXISTS {new_sep} " - f"USING vec0({_vec0_column_decl(dim)})" - ) + dim = _detect_existing_vec_dim(conn, "vec_messages") or 256 - conn.execute(f"INSERT INTO {new_vec} SELECT * FROM vec_messages") - try: - conn.execute(f"INSERT INTO {new_sep} SELECT * FROM vec_messages_sep") - except sqlite3.OperationalError: - pass + import sqlite_vec - conn.execute("DROP TABLE vec_messages") - try: - conn.execute("DROP TABLE vec_messages_sep") - except sqlite3.OperationalError: - pass + conn.enable_load_extension(True) + sqlite_vec.load(conn) + conn.enable_load_extension(False) - max_id = conn.execute( - f"SELECT MAX(rowid) FROM {new_vec}" - ).fetchone()[0] or 0 - - model_map = {"edge": "potion-base-8M", "basepro": "Qwen3-Embedding-0.6B"} - VectorCacheRegistry.set( - conn, - group, - vec_table=new_vec, - sep_table=new_sep, - last_embedded_id=max_id, - vector_count=count, - model_name=model_map.get(group, "potion-base-8M"), - embedding_dim=dim, - ) + conn.execute( + f"CREATE VIRTUAL TABLE IF NOT EXISTS {new_vec} " + f"USING vec0({_vec0_column_decl(dim)})" + ) + conn.execute( + f"CREATE VIRTUAL TABLE IF NOT EXISTS {new_sep} " + f"USING vec0({_vec0_column_decl(dim)})" + ) - logger.info( - "Migrated %d vectors from vec_messages → %s (group=%s)", - count, new_vec, group, - ) - return True + conn.execute(f"INSERT INTO {new_vec} SELECT * FROM vec_messages") + try: + conn.execute(f"INSERT INTO {new_sep} SELECT * FROM vec_messages_sep") + except sqlite3.OperationalError: + pass + + conn.execute("DROP TABLE vec_messages") + try: + conn.execute("DROP TABLE vec_messages_sep") + except sqlite3.OperationalError: + pass + + max_id = conn.execute( + f"SELECT MAX(rowid) FROM {new_vec}" + ).fetchone()[0] or 0 + + model_map = {"edge": "potion-base-8M", "basepro": "Qwen3-Embedding-0.6B"} + VectorCacheRegistry.set( + conn, + group, + vec_table=new_vec, + sep_table=new_sep, + last_embedded_id=max_id, + vector_count=count, + model_name=model_map.get(group, "potion-base-8M"), + embedding_dim=dim, + ) + + logger.info( + "Migrated %d vectors from vec_messages → %s (group=%s)", + count, new_vec, group, + ) + return True From bb955cd0a147c7cc6beaf1fa134c78b16c233303 Mon Sep 17 00:00:00 2001 From: buildingjoshbetter Date: Fri, 9 Oct 2026 15:21:26 -0500 Subject: [PATCH 03/11] docs: Record guarded public tier admission and recovery contracts --- docs/resource-budgets.md | 222 +++++++++++++++++++++++++++++++++++++++ 1 file changed, 222 insertions(+) diff --git a/docs/resource-budgets.md b/docs/resource-budgets.md index 5ce44cec..e12aa613 100644 --- a/docs/resource-budgets.md +++ b/docs/resource-budgets.md @@ -2312,3 +2312,225 @@ 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. + +### Guarded base/pro policy and config projection + +The activation journal also accepts version 2 intents. A guarded rebuild records +its observed config tier/generation and optional status-row identity before source +planning. Existing version 1 records retain their exact wire shape. Their missing +config guard means unknown prior config, never an assumed absent generation. + +A base/pro policy-only publication changes public tier and generation without +reading the corpus, encoding, preparing a model, or certifying legacy vectors. +An already certified selection preserves its historical paired certificate. A +legacy result is a separate `LegacyTierPolicy` with its observed generic or +basepro table pair; it does not create `tier_selected_v1`. Exact pair schema, +embedding metadata and registry identity must agree with the unchanged Qwen3 +256-dimensional space. Caller-owned maintenance and exclusive serving admission +remain required. Active-force/reset and other policy-only transitions are refused. + +Policy publication compares the previous selection and intent under its owned +writer and cancels only that intent's exact selected job marker in the same +transaction. It neither deletes a newer job nor claims an abandoned rebuild +succeeded. Pending config results must be acknowledged before a later intent. +A guarded rebuild preserves an acknowledged legacy policy while its target builds. + +`read_tier_progress` reads bounded manifest/receipt scalars. It performs no vector +or source audit, no status write, and no model call. Its complete flag describes +the captured range, not current source completeness. Status callers compare the +journal's target/status/job with the marker and progress tracker/epoch/schema in +one read snapshot. Borrowed snapshots are left open; owned read cleanup failure +propagates without a second rollback attempt. + +Config projection uses the same in-process and cross-process locks as MCP config +writes. Activation requires successful lock acquisition; old config APIs retain +their best-effort default. Capture seeds a generation without changing policy. +Independent settings use a fresh locked field merge. Mirroring changes only tier +and generation after an exact prior-guard comparison, or accepts an already +applied exact result. It refuses unknown prior state and concurrent untagged tier +edits. The config read is bounded to 1 MiB, including bytes after NUL characters. +A cooperating generation marker cannot detect identical ABA edits by arbitrary +external writers. JSON and SQLite remain separate resources: any projection +failure leaves the database decision authoritative and config acknowledgement +pending. Neither acknowledgement nor a configured proxy proves native readiness. + +Selected daemon identity/shape rejections now carry `serving_identity_mismatch`. +The client maps that code to `ServingIdentityMismatchError`, a protocol error, +so callers can stop a stale operation without misclassifying a genuine native +availability failure. Ordinary successful wire defaults and fast-lane eligibility +remain unchanged. These APIs require public caller adoption and runtime boundary +integration; this isolated journal/config patch alone does not reduce archive +residency or establish a fix for the original Mac memory observation. + +### Public maintenance and ingestion boundaries (#795, Unit C) + +This checkpoint carries the captured database selection or bounded legacy +base/pro policy through maintenance admission, vector compatibility work, and +one fact's encoding gate, duplicate decision and storage. It does not change +retrieval formulas, models, dimensions, reranker overrides or native batches. +Transcript parsing and remote extraction remain outside the per-fact lease. +The gate clears its cached encoder and nearest-search results on a changed +admitted generation, and releases its encoder reference after prediction-error +inference. The existing process singleton remains the model cache; exceptions +and callers can still retain their own references. + +A same-space policy projection can rebind an existing certified proxy or local +slot and publish the public tier without constructing, probing or encoding a +model. Empty or different slots stay empty. The configured runtime acknowledgement +is absent until both required slots are available; that absence does not mean a +zero-load policy projection failed. A later ordinary operation prepares missing +slots outside SQL. A legacy policy is never represented as a vector certificate. +Writer captures include the exact policy generation or its absence and accepted +connection; a policy appearing or changing before publication refuses the write. +Config acknowledgement alone does not invalidate the captured embedding space. + +The journal-aware factory opens existing databases in SQLite `mode=rw` and +validates controlled journals before connection-local settings, bypassing general +schema/tracking initialization. Legacy schema initialization, compatibility work, standalone +metadata setters and cosine/legacy migrations retain cooperative maintenance +ownership. Owner contention refuses without waiting or upgrading a serving read. +The file/in-memory NaN rebuild also owns maintenance before extension preparation +or table removal. Admission never ends caller-owned SQL. + +Foreground `add`, `update` and `embed_single` can publish matching embedding +metadata inside their existing model fence and database writer transaction, +without acquiring maintenance ownership. They validate the captured writer, +serving identity and vector target before publication. Missing identity rows can +be filled; partial or conflicting identity refuses. This helper creates no schema +and does not commit. Controlled metadata publication remains a writer-validated +no-op; it does not rewrite the selected generation or policy. + +Completed legacy file initialization can retain one bounded scalar receipt per +Engine. An externally closed handle or clean I/O-error handle can reconnect using +a separately validated candidate; explicit `close()` clears that Engine's receipt. +A fresh Engine can receive initialization evidence only from the current automatic +worker, bound to the exact coordinator, process, descriptor and launch token. +Cancellation, worker completion, failed launch and fork revoke that evidence. +A busy launch also discards its own receipt while retaining its pending work. +A distinct newer pending request retains only its own matching generation receipt. +Manual or foreign owners and automatic requests without a valid receipt provide +no such bypass. + +Both paths recheck file/schema identity and the current journal, run +`quick_check(1)`, and restore connection-local settings and required extension +registration. For legacy readiness they also compare the bounded registry and +embedding metadata. Newly written metadata must match the captured model and +dimension. Absent metadata requires FTS-only initialization or empty initialized +vector tables. Read cleanup must finish before receipt or +candidate publication. Candidate validation rejection preserves the prior Engine +handle. +Fresh selected/policy authority clears legacy readiness and uses normal controlled +reconciliation; initialization receipts are not vector certificates. + +Before reusing receipt-backed legacy readiness, Engine compares the admitted +embedding model and dimension with the receipt. A reranker-only Deep override +remains compatible. This check adds no SQL, file stat or model construction to the +ready path. Receipt reuse does not run migrations, transfer caller SQL, wait for an +owner, or retain models/connections in its evidence. These checks assume +cooperating owners and stable paths, not hostile file replacement protection. + +A fresh public Engine, fact, or MCP search boundary may coordinate a retry after +that single-attempt opener reports busy. Only a proven local automatic launch +originating with a valid initialization receipt permits waiting. Manual, foreign, +receiptless and cancelled owners retain refusal. Waiting occurs outside current +runtime operations, raw serving readers/exclusive leases, Engine writers and SQL. +The exact owner observed during failed acquisition supplies a per-launch release +notification; a wake grants no initialization or vector authority. The caller +retries normal opening and validation from scratch, preserving lazy style +initialization and exact schema/token checks. + +One monotonic budget covers the initial attempt and later admission retries: +`DEFAULT_BUSY_TIMEOUT_MS` gives 10,000 / 1,000 = 10 seconds. The budget bounds +retry admission and waiting, not hard preemption of an already-running native or +SQLite call. Successors do not reset it. Cancellation or expiry returns typed +busy. Parallel MCP child opens and manual consolidation retain their existing +single-attempt behavior; a child cannot wait behind its parent's serving lease. + +Cold legacy initialization completes under a provisional admission before the +public operation captures its immutable table pair. A generic-to-tiered migration +therefore cannot leave the first search or fact operation using dropped generic +tables. Existing parent operations retain their captured identity. A nested cold +Engine may attach to a validated existing pair or initialize an absent generic +pair without migrating those names; migration that would invalidate the parent's +pair refuses before mutation. Unrelated inactive caches remain intact. + +Automatic maintenance is deferred across the public operation and its reserved +children. The outer Engine scope requests normal maintenance after its serving +lease releases; nested calls and children do not launch competing workers. +Explicit maintenance suppression still applies. MCP serializes only sibling +connection-open steps within one request, releasing that lock before child +admission and search. Strict child opens retain their single-attempt contract. + +For synchronous in-memory macOS Qwen repair, legacy-name migration follows the +existing repair so the repaired vectors become the registered serving pair. +File-backed background repair retains its existing order. These ordering checks +do not establish native device behavior or a memory reduction. + +Concrete model-backed maintenance enters ownership before serving/dependencies. +Selected and policy clustering is explicitly deferred before native preparation +because its derived-vector publication is not adopted here. Cold public +consolidation and coordinator reopen preserve that deferral. Eligible nonvector +layers and lexical style work retain their existing transactions and scheduling; +they do not retain embedding models. Deferred work is reported as such. +Typed serving, writer, source-fence and receipt refusals propagate through Engine, +hybrid, gate, duplicate and ingestion availability fallbacks. Ordinary dependency +failures retain the existing fallback behavior. + +Tests distinguish standard-library production-AST/in-memory checks from opt-in +file-backed SQLite/sqlite-vec cases using synthetic encoders and real Engine, +serving, writer and automatic-worker paths. Local safe checks exclude native +models and file-backed databases; native tests run in the isolated test environment. +Results belong to their exact source snapshot and gate receipts. These scopes do +not establish a native Mac residency bound or prove that the original 26.24 GB +memory incident is fixed. Complete-archive Deep Search remains future scope. + + +### Public tier adoption and interrupted jobs (#795) + +Both public manager routes use one inactive-job preparation path and the bounded +source worker. The background handoff contains only the accepted job, intent and +status ID. It never retains source pages or passes a connection between threads. +The worker requires captured-range and current-source completion before the +certified finalizer can publish. Progress callbacks do not write per-page status. + +Configure keeps the compatibility `tier` response as the requested tier and +reports `served_tier` separately. Queued and failed requests clear the cached +Memory reference; existing operations retain their admitted references. Same-tier +requests consult the journal and retry config/runtime reconciliation. A commit or +projection ambiguity remains `activation_pending`. Suspect connections are closed; +committed tier readback uses a fresh validated connection and may remain unknown. + +Base/pro policy requests and CLI setup after a policy transition do not probe or +construct models. A first-run absent or schema-empty store can persist onboarding +without a vector certificate. Different first-run identities require both model +slots to be unloaded; supported lazy setters run under exclusive serving. Loaded +mismatches refuse before config writes. A failed config write restores the old +unloaded identity only after a locked read proves the old file remains; an +already-replaced file retains its matching new identity and reports pending. + +Interrupted selected jobs are never taken over automatically. `get_status` exposes +the exact job ID; explicit `cancel(status_id, expected_job_id=...)` requires matching +identity and maintenance admission. An orphan marker uses its explicit job ID +without asserting that a stale status row owns it. Terminal markers permit the +existing explicit previous-job admission on a later request. Status text is bounded +before materialization, and superseded rows cannot continue reporting running. + +Supplied backup paths remain recorded in status. This inactive workflow does not +invoke the legacy destructive backup-pruning helper. Force on the serving pair, +unsupported existing metadata, and unreceipted nonempty inactive pairs refuse. +These are implementation contracts, not native/full-gate or memory-reduction claims. + + +#### Explicit CLI recovery + +Read `truememory_status(status_id)` for `rebuild.id` and its matching `rebuild.job_id`. +For interrupted work, run `truememory-ingest cancel-rebuild --status-id ID --job-id JOB_ID` +and supply `--db PATH` when the original request used a different database path. +Both IDs are required; status IDs must be positive SQLite integers and job IDs must +be exact lowercase 32-character hexadecimal values. The command does not discover +or substitute IDs, automatically retry, start work, or delete source/vector rows. +A live maintenance owner, stale/mismatched IDs, or an unbound orphan marker refuses. +Only the exact terminal cancellation result counts as success. Then retry the +original `upgrade-tier` request explicitly; existing source/receipt validation +still controls whether the inactive prefix is resumable. An orphan +`selected_job_id` is not proof that a stale status row belongs to that job. From f043d090c87170e4c85113e3bc583e0e320c5d81 Mon Sep 17 00:00:00 2001 From: buildingjoshbetter Date: Sat, 10 Oct 2026 23:10:53 -0500 Subject: [PATCH 04/11] fix: Preserve FTS admission and portable tier regression fixtures --- tests/test-tier-activation-policy-795.py | 4 +- tests/test-tier-active-initialization-795.py | 27 ++++- tests/test-tier-metadata-publication-795.py | 2 +- tests/test-tier-public-adoption-795.py | 12 +- tests/test-tier-public-boundaries-795.py | 119 +++++++++++++++++-- tests/test-tier-runtime-bridge-795.py | 4 +- tests/test-tier-warm-admission-795.py | 9 +- tests/test_configure_reembed.py | 10 ++ truememory/engine.py | 7 +- 9 files changed, 163 insertions(+), 31 deletions(-) diff --git a/tests/test-tier-activation-policy-795.py b/tests/test-tier-activation-policy-795.py index 8b288ffd..d4705fb2 100644 --- a/tests/test-tier-activation-policy-795.py +++ b/tests/test-tier-activation-policy-795.py @@ -22,8 +22,8 @@ class TestActivationPolicy(unittest.TestCase): def setUp(self) -> None: self.fixture = BASE["TestActivationJournal"]() - self.fixture.setUp() self.addCleanup(self.fixture.doCleanups) + self.fixture.setUp() self.api = self.fixture.api self.conn = self.fixture.conn self.jobs = self.fixture.job_api @@ -422,8 +422,8 @@ class TestIdentityErrorWire(unittest.TestCase): def setUp(self): fixtures = runpy.run_path(str(Path(__file__).with_name("test-tier-serving-receipt-795.py"))) self.fixture = fixtures["TestServingReceipts"]() - self.fixture.setUp() self.addCleanup(self.fixture.doCleanups) + self.fixture.setUp() def test_identity_error_roundtrip_native_errors_still_generic(self): fixture = self.fixture diff --git a/tests/test-tier-active-initialization-795.py b/tests/test-tier-active-initialization-795.py index 0ec9020a..bc83c729 100644 --- a/tests/test-tier-active-initialization-795.py +++ b/tests/test-tier-active-initialization-795.py @@ -19,13 +19,14 @@ BOUNDARIES = runpy.run_path(str(Path(__file__).with_name("test-tier-public-boundaries-795.py"))) definitions = BOUNDARIES["definitions"] +enter_context = BOUNDARIES["enter_context"] class TestActiveInitialization(unittest.TestCase): def setUp(self): self.fixture = BOUNDARIES["TestEngineReconnectReceipt"]() - self.fixture.setUp() self.addCleanup(self.fixture.doCleanups) + self.fixture.setUp() self.first = self.fixture.engine() self.api = self.fixture.api self.closed_fds = [] @@ -52,7 +53,7 @@ def setUp(self): self.coordinator = self.maintenance.get_coordinator(self.first.db_path) self.fixture.engine_namespace["_HAS_HYBRID"] = True self.maintenance.run_engine_maintenance = self.work - self.fixture.enterContext(patch.object(Path, "exists", return_value=False)) + enter_context(self.fixture, patch.object(Path, "exists", return_value=False)) self.addCleanup(self.finish) def work(self, conn, coordinator, **kwargs): @@ -86,6 +87,28 @@ def record(self): with self.maintenance._active_initialization_receipt(self.first.db_path) as value: return value + def test_extension_unavailable_fts_worker_retains_validated_automatic_admission(self): + self.fixture.engine_namespace["_HAS_VECTOR"] = True + self.fixture.conn.enable_load_extension = None + operation = types.SimpleNamespace(key=self.api._legacy_key(), _database=("file", str(self.first.db_path))) + self.first._capture_reconnect_receipt(operation) + self.assertIsNotNone(self.first._reconnect_receipt) + self.assertEqual(self.first._reconnect_receipt[7:9], (False, True)) + self.start() + candidate = self.fixture.candidate(self.first) + candidate.enable_load_extension = None + engine = self.fresh() + engine._open_connection_handle() + self.assertIs(engine.conn, candidate) + self.assertTrue(engine.ready) + self.assertTrue(engine._runtime_initialized) + self.assertFalse(engine._has_vectors) + self.assertFalse(engine._has_hybrid) + self.assertEqual(engine._reconnect_receipt, self.first._reconnect_receipt) + self.assertEqual(self.maintenance._held_paths[self.first.db_path], 71) + self.assertEqual(self.coordinator.status[0], "running") + self.assertEqual(self.fixture.fixture.loads, []) + def test_fresh_engine_adopts_only_validated_active_automatic_receipt(self): self.start() record = self.record() diff --git a/tests/test-tier-metadata-publication-795.py b/tests/test-tier-metadata-publication-795.py index 68d977ab..5f3488cd 100644 --- a/tests/test-tier-metadata-publication-795.py +++ b/tests/test-tier-metadata-publication-795.py @@ -14,8 +14,8 @@ class TestForegroundMetadata(unittest.TestCase): def setUp(self) -> None: self.fixture = BOUNDARIES["TestPublicBoundaries"]() - self.fixture.setUp() self.addCleanup(self.fixture.doCleanups) + self.fixture.setUp() self.conn, self.runtime = self.fixture.conn, self.fixture.api self.writer = self.fixture.load_application("tier_switch.writer") self.vector = self.fixture.fixture.vector diff --git a/tests/test-tier-public-adoption-795.py b/tests/test-tier-public-adoption-795.py index 021abdf6..21f81014 100644 --- a/tests/test-tier-public-adoption-795.py +++ b/tests/test-tier-public-adoption-795.py @@ -100,7 +100,7 @@ def __init__(self) -> None: self.job = types.SimpleNamespace( job_id="a" * 32, target=Target.capture("base"), - database_path=Path("/synthetic/store.db"), + database_path=Path("/synthetic/store.db").absolute(), tracker="canonical-v1", source_epoch="epoch", source_schema_signature="b" * 64, @@ -796,7 +796,7 @@ def test_status_backup_path_round_trip_real_sql(self): conn.execute( "SELECT backup_path FROM rebuild_status WHERE id=?", (sid,) ).fetchone(), - ("/synthetic/backup",), + (str(Path("/synthetic/backup")),), ) self.assertFalse(conn.in_transaction) @@ -1173,7 +1173,7 @@ def test_cancel_cli_dispatch_uses_only_exact_supplied_identity(self): module.main() factory.assert_called_once_with() manager.cancel.assert_called_once_with( - 17, expected_job_id=h.job.job_id, db_path=Path("/synthetic/store.db") + 17, expected_job_id=h.job.job_id, db_path=h.job.database_path ) self.assertIn("Rebuild 17 cancelled", output.getvalue()) self.assertIn("explicitly", output.getvalue()) @@ -1374,8 +1374,8 @@ def setUp(self): self.os = os self.ns = runpy.run_path(str(ROOT / "tests/test-tier-public-boundaries-795.py")) self.boundary = self.ns["TestPublicBoundaries"]() - self.boundary.setUp() self.addCleanup(self.boundary.doCleanups) + self.boundary.setUp() self.conn, self.api = self.boundary.conn, self.boundary.api self.selected_marker_row = self.conn.execute("SELECT * FROM truememory_tier_selected_job_v1").fetchone() self.conn.execute("DELETE FROM truememory_tier_selected_job_v1") @@ -1383,7 +1383,7 @@ def setUp(self): (("embed_model", "model2vec"), ("embed_dim", "256"))) self.conn.commit() self.cfg = self.boundary.load_application("tier_config") - self.enterContext(patch.dict(os.environ, {"TRUEMEMORY_EMBED_MODEL": "edge"})) + self.ns["enter_context"](self, patch.dict(os.environ, {"TRUEMEMORY_EMBED_MODEL": "edge"})) definitions, namespace = self.ns["definitions"], self.boundary.namespace self.vector = definitions("vector_search.py", { "_resolve_model_name", "set_embedding_model", "get_embedding_dim", "resolve_tier", @@ -1442,7 +1442,7 @@ def close(self): handle = OwnedHandle() opener = Mock(return_value=handle) self.h.manager_module.sqlite3 = types.SimpleNamespace(connect=opener) - self.enterContext(patch.object(Path, "stat", return_value=types.SimpleNamespace())) + self.ns["enter_context"](self, patch.object(Path, "stat", return_value=types.SimpleNamespace())) return handle, opener def assert_no_persistent_mutation(self): diff --git a/tests/test-tier-public-boundaries-795.py b/tests/test-tier-public-boundaries-795.py index 66e6664d..bb6040d3 100644 --- a/tests/test-tier-public-boundaries-795.py +++ b/tests/test-tier-public-boundaries-795.py @@ -27,6 +27,12 @@ _REAL_PATH_RESOLVE = Path.resolve +def enter_context(test_case: unittest.TestCase, context: contextlib.AbstractContextManager[object]) -> object: + result = type(context).__enter__(context) + test_case.addCleanup(type(context).__exit__, context, None, None, None) + return result + + def definitions(path: str, names: set[str], namespace: dict, *, methods: bool = False) -> types.ModuleType: if path == "maintenance.py" and "_acquire_owner" in names: names = names | {"_AutomaticAdmission", "_automatic_admission_locked"} @@ -45,16 +51,51 @@ def definitions(path: str, names: set[str], namespace: dict, *, methods: bool = return module +class TestFixtureCompatibility(unittest.TestCase): + def test_context_cleanup_works_without_testcase_enter_context(self): + case = unittest.TestCase() + entered, exited = [], [] + @contextlib.contextmanager + def context(): + entered.append(True) + try: + yield "synthetic context" + finally: + exited.append(True) + result = enter_context(case, context()) + self.assertEqual(result, "synthetic context") + self.assertEqual(entered, [True]) + self.assertEqual(exited, []) + case.doCleanups() + self.assertEqual(exited, [True]) + + def test_nested_setup_failure_restores_global_path_stat(self): + warm = runpy.run_path(str(Path(__file__).with_name("test-tier-warm-admission-795.py"))) + active = warm["ACTIVE"]["TestActiveInitialization"] + original = Path.stat + def fail_setup(case): + enter_context(case, patch.object(Path, "stat", return_value=types.SimpleNamespace())) + raise RuntimeError("synthetic partial setup") + case = warm["TestWarmAdmission"]() + with patch.object(active, "setUp", fail_setup): + try: + with self.assertRaisesRegex(RuntimeError, "synthetic partial setup"): + case.setUp() + finally: + case.doCleanups() + self.assertIs(Path.stat, original) + + class TestPublicBoundaries(unittest.TestCase): def setUp(self) -> None: self.fixture = BRIDGE["TestRuntimeBridge"]() - self.fixture.setUp() self.addCleanup(self.fixture.doCleanups) + self.fixture.setUp() self.api, self.conn = self.fixture.api, self.fixture.conn self.modules = self.fixture.modules self.modules["maintenance"] = types.SimpleNamespace( maintenance_owner=lambda path: contextlib.nullcontext(), connection_database_path=lambda conn: None) - self.enterContext(patch.dict(sys.modules, {"truememory.tier_switch.runtime": self.api})) + enter_context(self, patch.dict(sys.modules, {"truememory.tier_switch.runtime": self.api})) def imports(self, name, globals=None, locals=None, fromlist=(), level=0): if name == "truememory.tier_switch.runtime": @@ -291,7 +332,7 @@ def extract(*args): def load_application(self, name: str) -> types.ModuleType: module = types.ModuleType("synthetic_public_" + name.replace(".", "_")) module.__dict__["__builtins__"] = dict(vars(builtins), __import__=self.imports) - self.enterContext(patch.dict(sys.modules, {module.__name__: module})) + enter_context(self, patch.dict(sys.modules, {module.__name__: module})) path = ROOT / "truememory" / (name.replace(".", "/") + ".py") exec(compile(path.read_text(encoding="utf-8"), str(path), "exec"), module.__dict__) self.modules[name] = module @@ -606,7 +647,7 @@ def test_legacy_initializer_propagates_typed_guard_or_owner_refusal(self) -> Non class Busy(RuntimeError): pass self.modules["maintenance"].MaintenanceBusyError = Busy - self.enterContext(patch.dict(sys.modules, {"truememory.maintenance": self.modules["maintenance"]})) + enter_context(self, patch.dict(sys.modules, {"truememory.maintenance": self.modules["maintenance"]})) def imports(name, globals=None, locals=None, fromlist=(), level=0): if name == "sqlite_vec": return types.SimpleNamespace(load=lambda conn: None) @@ -874,7 +915,7 @@ def test_opener_owns_canonical_path_before_writer_and_reenters_real_owner(self) engine = self.fixture.lifecycle_engine("_open_connection_handle") engine.conn = None engine.db_path = Path("~synthetic-unexpanded/store.sqlite") - self.enterContext(patch.object(Path, "resolve", _REAL_PATH_RESOLVE)) + enter_context(self, patch.object(Path, "resolve", _REAL_PATH_RESOLVE)) expected = engine.db_path.resolve() acquired, created = [], [] @@ -1528,11 +1569,11 @@ class TestEngineReconnectReceipt(unittest.TestCase): real_memory_owner = TestPublicBoundaries.real_memory_owner def engine(self, *, vectors=False): - self.enterContext(patch.object(Path, "resolve", _REAL_PATH_RESOLVE)) + enter_context(self, patch.object(Path, "resolve", _REAL_PATH_RESOLVE)) self.stat = types.SimpleNamespace(st_dev=7, st_ino=11) - self.enterContext(patch.object(Path, "stat", return_value=self.stat)) - self.enterContext(patch.object(Path, "mkdir")) - self.enterContext(patch.object(Path, "chmod")) + enter_context(self, patch.object(Path, "stat", return_value=self.stat)) + enter_context(self, patch.object(Path, "mkdir")) + enter_context(self, patch.object(Path, "chmod")) self.modules["storage"] = definitions("storage.py", {"DatabaseOpenError", "_integrity_message", "_validate_db_path"}, self.namespace(sqlite3=sqlite3, Path=Path, os=os, newest_backup=lambda path: None)) namespace = self.namespace(sqlite3=sqlite3, Path=Path, logger=logging.getLogger("synthetic-reconnect"), @@ -1549,7 +1590,7 @@ def engine(self, *, vectors=False): exec(compile(ast.fix_missing_locations(ast.Module(body=[future, cls], type_ignores=[])), "actual-engine-reconnect", "exec"), namespace) engine = namespace["TrueMemoryEngine"]() - engine.conn, engine.db_path = self.conn, Path("/synthetic/reconnect.sqlite") + engine.conn, engine.db_path = self.conn, Path("/synthetic/reconnect.sqlite").absolute() engine._init_lock = threading.Lock() engine._write_lock = self.api.ConnectionWriteLock(engine) engine.ready = engine._runtime_initialized = True @@ -1594,6 +1635,58 @@ def assert_closed(self, conn): with self.assertRaises(sqlite3.ProgrammingError): conn.execute("SELECT 1") + def test_extension_unavailable_fts_receipt_preserves_capability_and_source(self): + engine = self.engine() + self.engine_namespace["_HAS_VECTOR"] = True + self.conn.enable_load_extension = None + self.conn.execute("INSERT INTO messages(id,content) VALUES(883,'synthetic FTS source')") + self.conn.commit() + operation = types.SimpleNamespace(key=self.api._legacy_key(), _database=("file", str(engine.db_path))) + engine._capture_reconnect_receipt(operation) + self.assertIsNotNone(engine._reconnect_receipt) + self.assertEqual(engine._reconnect_receipt[7:9], (False, True)) + candidate = self.candidate(engine) + candidate.enable_load_extension = None + owner = self.forbid_owner() + self.conn.close() + engine._open_connection_handle() + self.assertIs(engine.conn, candidate) + self.assertTrue(engine.ready) + self.assertTrue(engine._runtime_initialized) + self.assertFalse(engine._has_vectors) + self.assertFalse(candidate.in_transaction) + self.assertEqual(candidate.execute("SELECT content FROM messages WHERE id=883").fetchone()[0], + "synthetic FTS source") + owner.assert_not_called() + + def test_fts_receipt_cannot_skip_new_extension_capability(self): + engine = self.engine() + self.engine_namespace["_HAS_VECTOR"] = True + self.conn.enable_load_extension = None + operation = types.SimpleNamespace(key=self.api._legacy_key(), _database=("file", str(engine.db_path))) + engine._capture_reconnect_receipt(operation) + receipt = engine._reconnect_receipt + self.assertIsNotNone(receipt) + candidate = self.candidate(engine) + candidate.enable_load_extension = lambda enabled: None + self.assertIsNone(engine._reopen_initialized_connection(engine.db_path)) + self.assert_closed(candidate) + self.assertIs(engine.conn, self.conn) + self.assertIs(engine._reconnect_receipt, receipt) + self.assertFalse(self.conn.in_transaction) + + def test_other_vector_initialization_failure_does_not_publish_fts_receipt(self): + engine = self.engine() + self.engine_namespace["_HAS_VECTOR"] = True + self.conn.enable_load_extension = lambda enabled: None + operation = types.SimpleNamespace(key=self.api._legacy_key(), _database=("file", str(engine.db_path))) + queries = [] + self.conn.set_trace_callback(queries.append) + engine._capture_reconnect_receipt(operation) + self.assertIsNone(engine._reconnect_receipt) + self.assertEqual(queries, []) + self.assertFalse(self.conn.in_transaction) + def test_closed_initialized_fts_handle_reopens_without_owner_or_schema_writes(self): engine = self.engine() self.conn.execute("INSERT INTO messages(id,content) VALUES(881,'synthetic prior source')") @@ -1610,7 +1703,7 @@ def test_closed_initialized_fts_handle_reopens_without_owner_or_schema_writes(se self.assertFalse(engine._has_vectors) self.assertFalse(candidate.in_transaction) self.assertEqual(candidate.execute("SELECT content FROM messages WHERE id=881").fetchone()[0], "synthetic prior source") - self.connect.assert_called_once_with("file:///synthetic/reconnect.sqlite?mode=rw", uri=True, check_same_thread=False) + self.connect.assert_called_once_with(engine.db_path.as_uri() + "?mode=rw", uri=True, check_same_thread=False) self.assertIn("PRAGMA quick_check(1)", queries) self.assertIn("PRAGMA foreign_keys=ON", queries) self.assertFalse(any(q.lstrip().split(" ", 1)[0].upper() in {"INSERT", "UPDATE", "DELETE", "CREATE", "ALTER", "DROP"} for q in queries)) @@ -1752,7 +1845,7 @@ def connect(*args, **kwargs): def test_changed_file_after_candidate_validation_is_refused(self): engine = self.engine() candidate = self.candidate(engine) - self.enterContext(patch.object(Path, "stat", side_effect=[self.stat, types.SimpleNamespace(st_dev=7, st_ino=12)])) + enter_context(self, patch.object(Path, "stat", side_effect=[self.stat, types.SimpleNamespace(st_dev=7, st_ino=12)])) self.assertIsNone(engine._reopen_initialized_connection(engine.db_path)) self.assert_closed(candidate) @@ -1869,7 +1962,7 @@ def opener(path, callback): with patch.object(self.api, "open_serving_connection", side_effect=opener), patch.object(Path, "exists", return_value=True): with self.assertRaises(Completed): module.open(engine) - self.connect.assert_called_once_with("file:///synthetic/reconnect.sqlite?mode=rw", uri=True, check_same_thread=False) + self.connect.assert_called_once_with(engine.db_path.as_uri() + "?mode=rw", uri=True, check_same_thread=False) self.assertEqual(queries, ["PRAGMA journal_mode=WAL", "PRAGMA busy_timeout=3000", "PRAGMA foreign_keys=ON", "PRAGMA synchronous=NORMAL", "PRAGMA cache_size=-64000", "PRAGMA mmap_size=268435456"]) self.assertIs(engine.conn, self.conn) diff --git a/tests/test-tier-runtime-bridge-795.py b/tests/test-tier-runtime-bridge-795.py index 7d814db7..ec796592 100644 --- a/tests/test-tier-runtime-bridge-795.py +++ b/tests/test-tier-runtime-bridge-795.py @@ -32,8 +32,8 @@ class Cancelled(BaseException): class TestRuntimeBridge(unittest.TestCase): def setUp(self) -> None: self.fixture = JOURNAL["TestActivationJournal"]() - self.fixture.setUp() self.addCleanup(self.fixture.doCleanups) + self.fixture.setUp() self.conn = self.fixture.conn self.modules = self.fixture.modules self.gate = GATE["load"]() @@ -900,8 +900,8 @@ def test_actual_selected_ready_borrowed_init_ignores_competing_writer(self) -> N class TestRuntimeModels(unittest.TestCase): def setUp(self) -> None: self.fixture = PREPARED["TestPreparedTargets"]() - self.fixture.setUp() self.addCleanup(self.fixture.doCleanups) + self.fixture.setUp() self.vector = self.fixture.vector self.vector._frozen_embedding_target = None self.gate = GATE["load"]() diff --git a/tests/test-tier-warm-admission-795.py b/tests/test-tier-warm-admission-795.py index eb8c7fe8..b566e223 100644 --- a/tests/test-tier-warm-admission-795.py +++ b/tests/test-tier-warm-admission-795.py @@ -11,13 +11,14 @@ from unittest.mock import Mock, patch ACTIVE = runpy.run_path(str(Path(__file__).with_name("test-tier-active-initialization-795.py"))) +enter_context = ACTIVE["enter_context"] class TestWarmAdmission(unittest.TestCase): def setUp(self): self.active = ACTIVE["TestActiveInitialization"]() - self.active.setUp() self.addCleanup(self.active.doCleanups) + self.active.setUp() self.fixture = self.active.fixture self.api, self.maintenance = self.active.api, self.active.maintenance self.coordinator = self.active.coordinator @@ -63,7 +64,7 @@ def wait(error, *, deadline): if action is not None: action() return actual(error, deadline=deadline) - self.enterContext(patch.object(self.maintenance, "wait_for_automatic_owner", side_effect=wait)) + enter_context(self, patch.object(self.maintenance, "wait_for_automatic_owner", side_effect=wait)) return entered def test_schema_drift_waits_then_uses_normal_factory_without_stale_readiness(self): @@ -182,7 +183,7 @@ def observe(): if not continue_launch.wait(3): raise AssertionError("Synthetic launch barrier timed out") return actual() - self.enterContext(patch.object(self.coordinator, "_begin_observation", side_effect=observe)) + enter_context(self, patch.object(self.coordinator, "_begin_observation", side_effect=observe)) requester = threading.Thread(target=lambda: self.coordinator.request_layers( _initialization_receipt=self.active.first._reconnect_receipt)) requester.start() @@ -217,7 +218,7 @@ def release(fd): if not continue_release.wait(3): raise AssertionError("Synthetic release barrier timed out") return actual(fd) - self.enterContext(patch.object(self.maintenance, "_release_owner", side_effect=release)) + enter_context(self, patch.object(self.maintenance, "_release_owner", side_effect=release)) self.active.release.set() self.assertTrue(at_release.wait(2)) self.assertIsNone(self.active.record()) diff --git a/tests/test_configure_reembed.py b/tests/test_configure_reembed.py index 58915de2..4738762c 100644 --- a/tests/test_configure_reembed.py +++ b/tests/test_configure_reembed.py @@ -44,6 +44,16 @@ def encode(self, texts, **kwargs): encoder = SyntheticEncoder() monkeypatch.setattr(vs, "get_model", lambda: encoder) + memory = ms._get_memory() + memory._engine._ensure_connection(_suppress_maintenance=True) + # The no-op contract requires coherent source identity even when this + # Python build cannot load the vector extension. + with memory._engine._write_lock: + memory._engine.conn.executemany( + "INSERT OR REPLACE INTO metadata(key,value) VALUES (?,?)", + (("embed_model", vs.EMBEDDING_MODEL), ("embed_dim", str(vs._embedding_dim))), + ) + memory._engine.conn.commit() yield ms # Teardown: drop cached Memory and any model-level state mutated by # truememory_configure (which modifies vector_search globals via diff --git a/truememory/engine.py b/truememory/engine.py index 5a66eb62..08547e2e 100644 --- a/truememory/engine.py +++ b/truememory/engine.py @@ -561,7 +561,8 @@ def _reconnect_registry(conn: sqlite3.Connection) -> tuple | None: def _capture_reconnect_receipt(self, operation: object) -> None: """Remember only this Engine's completed legacy initialization, never a model.""" self._reconnect_receipt = None - if self.conn.in_transaction or (not self._has_vectors and _HAS_VECTOR): + if self.conn.in_transaction or (not self._has_vectors and _HAS_VECTOR + and callable(getattr(self.conn, "enable_load_extension", None))): return from truememory.tier_switch.runtime import _legacy_key, TierRuntimeError from truememory import vector_search @@ -615,6 +616,10 @@ def _reopen_initialized_connection(self, path: Path, *, receipt: tuple | None = candidate = sqlite3.connect(path.as_uri() + "?mode=rw", uri=True, check_same_thread=False) retained = False try: + # A receipt from a connection without extension loading must not + # suppress normal vector initialization on a newly capable handle. + if receipt[8] and not receipt[7] and callable(getattr(candidate, "enable_load_extension", None)): + return None candidate.execute(f"PRAGMA busy_timeout={DEFAULT_BUSY_TIMEOUT_MS}") candidate.execute("PRAGMA foreign_keys=ON") candidate.execute("PRAGMA synchronous=NORMAL") From b1acf171ea9bd4bb6bbd599c854322c263ecc109 Mon Sep 17 00:00:00 2001 From: buildingjoshbetter Date: Sat, 10 Oct 2026 23:46:25 -0500 Subject: [PATCH 05/11] fix: Assert exact supplied cancel database path across platforms --- tests/test-tier-public-adoption-795.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/tests/test-tier-public-adoption-795.py b/tests/test-tier-public-adoption-795.py index 21f81014..128ccc93 100644 --- a/tests/test-tier-public-adoption-795.py +++ b/tests/test-tier-public-adoption-795.py @@ -1148,6 +1148,7 @@ def test_unreadable_store_cannot_be_treated_as_absent(self): def test_cancel_cli_dispatch_uses_only_exact_supplied_identity(self): h = self.h + supplied_db = "/synthetic/store.db" manager = types.SimpleNamespace( cancel=Mock(return_value={"status": "cancelled", "job_id": h.job.job_id}) ) @@ -1165,7 +1166,7 @@ def test_cancel_cli_dispatch_uses_only_exact_supplied_identity(self): "--job-id", h.job.job_id, "--db", - "/synthetic/store.db", + supplied_db, ], ), contextlib.redirect_stdout(output), @@ -1173,7 +1174,7 @@ def test_cancel_cli_dispatch_uses_only_exact_supplied_identity(self): module.main() factory.assert_called_once_with() manager.cancel.assert_called_once_with( - 17, expected_job_id=h.job.job_id, db_path=h.job.database_path + 17, expected_job_id=h.job.job_id, db_path=Path(supplied_db) ) self.assertIn("Rebuild 17 cancelled", output.getvalue()) self.assertIn("explicitly", output.getvalue()) From f35a194cb12f6fad9e8a07c303d9194c9cd643e6 Mon Sep 17 00:00:00 2001 From: buildingjoshbetter Date: Sun, 11 Oct 2026 00:10:34 -0500 Subject: [PATCH 06/11] fix: Retry poisoned handles after verified automatic owner release --- tests/test-tier-reconnect-admission-684.py | 645 +++++++++++++++++++++ truememory/engine.py | 41 +- truememory/tier_switch/runtime.py | 31 +- 3 files changed, 694 insertions(+), 23 deletions(-) create mode 100644 tests/test-tier-reconnect-admission-684.py diff --git a/tests/test-tier-reconnect-admission-684.py b/tests/test-tier-reconnect-admission-684.py new file mode 100644 index 00000000..8a333d2d --- /dev/null +++ b/tests/test-tier-reconnect-admission-684.py @@ -0,0 +1,645 @@ +"""Poisoned-handle public admission using production AST and synthetic owners.""" +from __future__ import annotations + +import ast +from collections.abc import Callable, Iterator +from contextlib import contextmanager, nullcontext +import os +from pathlib import Path +import sqlite3 +import subprocess +import sys +import threading +import types +from typing import NamedTuple +import unittest +import uuid +import weakref + + +ROOT = Path(__file__).resolve().parents[1] + + +def source(path: str) -> str: + commit = os.environ.get("TRUEMEMORY_RECONNECT_SOURCE_COMMIT") + if commit: + return subprocess.run( + ["git", "show", f"{commit}:truememory/{path}"], cwd=ROOT, + check=True, capture_output=True, text=True, + ).stdout + return (ROOT / "truememory" / path).read_text(encoding="utf-8") + + +def definitions(path: str, names: set[str], namespace: dict[str, object]) -> types.ModuleType: + tree = ast.parse(source(path)) + selected = [node for node in ast.walk(tree) + if isinstance(node, (ast.FunctionDef, ast.ClassDef)) and node.name in names] + if {node.name for node in selected} != names: + raise AssertionError("Missing production definitions") + module = types.ModuleType("synthetic_reconnect_684") + module.__dict__.update(namespace) + selected.insert(0, ast.ImportFrom(module="__future__", names=[ast.alias(name="annotations")], level=0)) + loaded = ast.fix_missing_locations(ast.Module(body=selected, type_ignores=[])) + exec(compile(loaded, path, "exec"), module.__dict__) + return module + + +class Clock: + def __init__(self) -> None: + self.now = 100.0 + + def monotonic(self) -> float: + return self.now + + +class Signal: + def __init__(self, action: Callable[[float], None] | None = None) -> None: + self.set_value = False + self.action = action + self.waits: list[float] = [] + + def is_set(self) -> bool: + return self.set_value + + def set(self) -> None: + self.set_value = True + + def wait(self, timeout: float) -> bool: + self.waits.append(timeout) + if self.action is not None: + self.action(timeout) + return self.set_value + + +class Writer: + def __init__(self) -> None: + self.owned = False + + def owned_by_current_thread(self) -> bool: + return self.owned + + def locked(self) -> bool: + return self.owned + + +class SyntheticPath: + def __init__(self) -> None: + self.mkdir_calls = 0 + self.chmod_calls = 0 + + def __str__(self) -> str: + return "/synthetic/reconnect-684.db" + + def resolve(self) -> SyntheticPath: + return self + + @property + def parent(self) -> SyntheticPath: + return self + + def mkdir(self, *, parents: bool, exist_ok: bool) -> None: + self.mkdir_calls += 1 + + def chmod(self, mode: int) -> None: + self.chmod_calls += 1 + + +class Handle: + def __init__(self, reason: str | None = None) -> None: + self.reason = reason + self.in_transaction = False + self.probes = 0 + self.closes = 0 + + def execute(self, sql: str) -> None: + if sql != "PRAGMA schema_version": + raise AssertionError("Unexpected synthetic SQL") + self.probes += 1 + if self.reason == "closed": + raise sqlite3.ProgrammingError("Cannot operate on a closed database.") + if self.reason == "io_error": + raise sqlite3.OperationalError("disk I/O error") + if self.reason == "other_error": + raise sqlite3.OperationalError("synthetic unrelated error") + + def close(self) -> None: + self.reason = "closed" + self.closes += 1 + + +class Fixture: + def __init__(self, reason: str | None = "closed") -> None: + self.clock = Clock() + self.path = SyntheticPath() + self.scope_depth = 0 + self.wait_scopes: list[int] = [] + self.deadlines: list[float] = [] + self.local = types.SimpleNamespace(operation=None) + self.admitted = False + self.admission_error = False + self.factory_calls = 0 + self.acquisitions = 0 + self.open_calls = 0 + self.previous_modules: dict[str, types.ModuleType | None] = {} + self.registry = threading.Lock() + self.held_paths: dict[object, int] = {} + self.coordinators: dict[object, object] = {} + self.handle = Handle(reason) + self.replacement = Handle() + self.reopen_calls: list[tuple[object, object]] = [] + self.reopened: tuple[Handle, bool] | None = None + self.gate_error = type("ServingLeaseError", (RuntimeError,), {}) + self.runtime_error = type("TierRuntimeError", (RuntimeError,), {}) + self.engine = types.SimpleNamespace( + conn=self.handle, db_path=self.path, _write_lock=Writer(), _init_lock=object(), + _reconnect_receipt=None, ready=True, _runtime_initialized=True, _has_vectors=True, + _runtime_vector_generation="synthetic-generation", _runtime_vector_connection=self.handle, + ) + + @contextmanager + def connection_read(*args: object, **kwargs: object) -> Iterator[None]: + self.scope_depth += 1 + try: + yield + finally: + self.scope_depth -= 1 + + def factory(path: object, create: object) -> Handle: + self.factory_calls += 1 + return self.replacement + + def file_lock(path: object) -> int: + self.acquisitions += 1 + return 901 + + def release(fd: int) -> None: + for path, claimed in tuple(self.held_paths.items()): + if claimed == fd: + del self.held_paths[path] + + self.maintenance = definitions("maintenance.py", { + "MaintenanceBusyError", "MaintenanceOwnership", "_AutomaticAdmission", + "_automatic_admission_locked", "_acquire_owner", "_bind_owner", "maintenance_owner", + "wait_for_automatic_owner", + }, { + "os": os, "weakref": weakref, "uuid": uuid, "NamedTuple": NamedTuple, + "contextmanager": contextmanager, "time": self.clock, + "threading": types.SimpleNamespace(Event=Signal, TIMEOUT_MAX=threading.TIMEOUT_MAX), + "Path": lambda value: value, "_registry_lock": self.registry, + "_coordinators": self.coordinators, "_held_paths": self.held_paths, + "_held_fds": set(), "_thread_owners": threading.local(), + "try_file_lock": file_lock, "_release_owner": release, + "canonical_database_path": lambda value: value, + }) + self.maintenance._active_initialization_receipt = lambda path: nullcontext(None) + self.runtime = definitions("tier_switch/runtime.py", {"open_engine_for_operation"}, { + "time": self.clock, "ConnectionWriteLock": Writer, "_local": self.local, + "TierRuntimeError": self.runtime_error, + "serving": types.SimpleNamespace(current_thread_admitted=self.currently_admitted, + ServingLeaseError=self.gate_error), + }) + self.runtime._connection_read = connection_read + self.runtime._open_serving_connection_at_path = factory + self.storage = types.ModuleType("truememory.storage") + storage_tree = ast.parse(source("storage.py")) + self.storage.DEFAULT_BUSY_TIMEOUT_MS = next( + ast.literal_eval(node.value) for node in storage_tree.body + if isinstance(node, ast.Assign) + and any(isinstance(target, ast.Name) and target.id == "DEFAULT_BUSY_TIMEOUT_MS" + for target in node.targets) + ) + engine_module = definitions("engine.py", {"_open_connection_handle"}, { + "sqlite3": sqlite3, "create_db": lambda path: self.replacement, "_HAS_HYBRID": False, + }) + actual_open = engine_module._open_connection_handle + + def opening() -> None: + self.open_calls += 1 + actual_open(self.engine) + + def reopen(path: object, *, receipt: object) -> tuple[Handle, bool] | None: + self.reopen_calls.append((path, receipt)) + return self.reopened + + self.engine._open_connection_handle = opening + self.engine._reopen_initialized_connection = reopen + actual_wait = self.maintenance.wait_for_automatic_owner + + def waiting(error: BaseException, *, deadline: float) -> None: + self.wait_scopes.append(self.scope_depth) + self.deadlines.append(deadline) + actual_wait(error, deadline=deadline) + + self.maintenance.wait_for_automatic_owner = waiting + for name, module in { + "truememory": types.ModuleType("truememory"), + "truememory.tier_switch": types.ModuleType("truememory.tier_switch"), + "truememory.tier_switch.runtime": self.runtime, + "truememory.maintenance": self.maintenance, "truememory.storage": self.storage, + }.items(): + self.previous_modules[name] = sys.modules.get(name) + sys.modules[name] = module + + def close(self) -> None: + for name, previous in self.previous_modules.items(): + if previous is None: + sys.modules.pop(name, None) + else: + sys.modules[name] = previous + + def currently_admitted(self) -> bool: + if self.admission_error: + raise self.gate_error("synthetic gate unavailable") + return self.admitted + + def automatic_owner(self, *, release_after: float = 1.0) -> object: + coordinator = types.SimpleNamespace( + _mutex=threading.Lock(), _pid=os.getpid(), _active=True, _cancel=Signal(), + ) + # SimpleNamespace has no weak-reference slot; the callable preserves identity. + released = Signal() + origin = self.maintenance._AutomaticAdmission( + lambda: coordinator, self.path, os.getpid(), 71, "synthetic-launch", released, Signal(), + ) + coordinator._admission_owner = origin + self.coordinators[self.path] = coordinator + self.held_paths[self.path] = 71 + + def release(timeout: float) -> None: + self.clock.now += min(timeout, release_after) + if timeout >= release_after: + self.held_paths.clear() + coordinator._active = False + released.set() + + released.action = release + return origin + + +class TestReconnectAdmission(unittest.TestCase): + def fixture(self, reason: str | None = "closed") -> Fixture: + fixture = Fixture(reason) + self.addCleanup(fixture.close) + return fixture + + def test_closed_handle_retries_exact_automatic_owner_outside_engine_locks(self) -> None: + fixture = self.fixture() + origin = fixture.automatic_owner() + fixture.runtime.open_engine_for_operation(fixture.engine) + self.assertEqual(origin.released.waits, [10.0]) + self.assertEqual(fixture.wait_scopes, [0]) + self.assertEqual(fixture.deadlines, [110.0]) + self.assertIs(fixture.engine.conn, fixture.replacement) + self.assertEqual(fixture.factory_calls, 1) + self.assertEqual(fixture.open_calls, 2) + self.assertFalse(fixture.engine.ready) + self.assertFalse(fixture.engine._runtime_initialized) + + def test_io_error_handle_retries_same_strict_automatic_owner(self) -> None: + fixture = self.fixture("io_error") + origin = fixture.automatic_owner() + fixture.runtime.open_engine_for_operation(fixture.engine) + self.assertEqual(origin.released.waits, [10.0]) + self.assertEqual(fixture.handle.closes, 1) + self.assertEqual(fixture.wait_scopes, [0]) + self.assertIs(fixture.engine.conn, fixture.replacement) + + def test_healthy_handle_never_claims_maintenance_or_waits(self) -> None: + fixture = self.fixture(None) + origin = fixture.automatic_owner() + fixture.runtime.open_engine_for_operation(fixture.engine) + self.assertIs(fixture.engine.conn, fixture.handle) + self.assertTrue(fixture.engine.ready) + self.assertEqual(fixture.handle.probes, 1) + self.assertEqual(fixture.open_calls, 1) + self.assertEqual(fixture.factory_calls, 0) + self.assertEqual(fixture.acquisitions, 0) + self.assertEqual(origin.released.waits, []) + self.assertEqual(fixture.deadlines, []) + + def test_private_opener_attests_exact_failed_handle_without_waiting(self) -> None: + fixture = self.fixture() + origin = fixture.automatic_owner() + with self.assertRaises(fixture.maintenance.MaintenanceBusyError) as raised: + fixture.engine._open_connection_handle() + self.assertIs(raised.exception._automatic_owner, origin) + self.assertIs(getattr(raised.exception, "_reconnect_engine", None), fixture.engine) + self.assertIs(getattr(raised.exception, "_reconnect_handle", None), fixture.handle) + self.assertIs(fixture.engine.conn, fixture.handle) + self.assertEqual(fixture.wait_scopes, []) + self.assertEqual(origin.released.waits, []) + self.assertEqual(fixture.factory_calls, 0) + + def test_missing_handle_keeps_existing_exact_origin_retry(self) -> None: + fixture = self.fixture() + fixture.engine.conn = None + origin = fixture.automatic_owner() + fixture.runtime.open_engine_for_operation(fixture.engine) + self.assertEqual(origin.released.waits, [10.0]) + self.assertEqual(fixture.wait_scopes, [0]) + self.assertIs(fixture.engine.conn, fixture.replacement) + self.assertEqual(fixture.open_calls, 2) + + def test_manual_receiptless_foreign_and_cancelled_owners_never_wait(self) -> None: + for kind in ("manual", "receiptless", "foreign_pid", "cancelled", "foreign_descriptor"): + with self.subTest(kind=kind): + fixture = self.fixture() + origin = fixture.automatic_owner() + coordinator = fixture.coordinators[fixture.path] + if kind == "manual": + fixture.coordinators.clear() + elif kind == "receiptless": + coordinator._admission_owner = None + elif kind == "foreign_pid": + coordinator._pid = -1 + elif kind == "cancelled": + origin.cancelled.set() + else: + fixture.held_paths[fixture.path] = 991 + with self.assertRaises(fixture.maintenance.MaintenanceBusyError): + fixture.runtime.open_engine_for_operation(fixture.engine) + self.assertEqual(origin.released.waits, []) + self.assertIs(fixture.engine.conn, fixture.handle) + self.assertEqual(fixture.factory_calls, 0) + self.assertEqual(fixture.open_calls, 1) + + def test_serving_runtime_raw_or_unqueryable_scopes_refuse_before_wait(self) -> None: + for kind in ("runtime", "raw_or_exclusive", "gate_error", "writer"): + with self.subTest(kind=kind): + fixture = self.fixture() + origin = fixture.automatic_owner() + if kind == "runtime": + fixture.local.operation = object() + elif kind == "raw_or_exclusive": + fixture.admitted = True + elif kind == "gate_error": + fixture.admission_error = True + else: + fixture.engine._write_lock.owned = True + with self.assertRaises(fixture.maintenance.MaintenanceBusyError): + fixture.runtime.open_engine_for_operation(fixture.engine) + self.assertEqual(origin.released.waits, []) + self.assertEqual(fixture.wait_scopes, []) + self.assertEqual(fixture.factory_calls, 0) + self.assertIs(fixture.engine.conn, fixture.handle) + + def test_fresh_bound_writer_is_refused_before_opening(self) -> None: + fixture = self.fixture() + fixture.engine.conn = None + fixture.engine._write_lock.owned = True + with self.assertRaises(fixture.maintenance.MaintenanceBusyError): + fixture.runtime.open_engine_for_operation(fixture.engine) + self.assertEqual(fixture.open_calls, 0) + self.assertEqual(fixture.wait_scopes, []) + + def test_healthy_or_borrowed_unmarked_busy_error_never_waits(self) -> None: + for borrowed in (False, True): + with self.subTest(borrowed=borrowed): + fixture = self.fixture(None) + fixture.handle.in_transaction = borrowed + origin = fixture.automatic_owner() + error = fixture.maintenance.MaintenanceBusyError("synthetic unmarked failure") + error._automatic_owner = origin + + def reject() -> None: + raise error + + fixture.engine._open_connection_handle = reject + with self.assertRaises(fixture.maintenance.MaintenanceBusyError) as raised: + fixture.runtime.open_engine_for_operation(fixture.engine) + self.assertIs(raised.exception, error) + self.assertEqual(fixture.wait_scopes, []) + self.assertEqual(origin.released.waits, []) + self.assertIs(fixture.engine.conn, fixture.handle) + self.assertEqual(fixture.handle.in_transaction, borrowed) + + def test_validated_receipt_reopen_preserves_existing_ownership_free_path(self) -> None: + for controlled in (False, True): + with self.subTest(controlled=controlled): + fixture = self.fixture() + receipt = object() + fixture.engine._reconnect_receipt = receipt + fixture.reopened = (fixture.replacement, controlled) + origin = fixture.automatic_owner() + fixture.runtime.open_engine_for_operation(fixture.engine) + self.assertEqual(fixture.reopen_calls, [(fixture.path, receipt)]) + self.assertIs(fixture.engine.conn, fixture.replacement) + self.assertEqual(origin.released.waits, []) + self.assertEqual(fixture.wait_scopes, []) + self.assertEqual(fixture.factory_calls, 0) + self.assertEqual(fixture.engine.ready, not controlled) + self.assertIs(fixture.engine._reconnect_receipt, None if controlled else receipt) + + def test_rejected_receipt_uses_owned_fallback_and_resets_readiness(self) -> None: + fixture = self.fixture() + receipt = object() + fixture.engine._reconnect_receipt = receipt + origin = fixture.automatic_owner() + fixture.runtime.open_engine_for_operation(fixture.engine) + self.assertEqual(fixture.reopen_calls, [(fixture.path, receipt), (fixture.path, receipt)]) + self.assertEqual(origin.released.waits, [10.0]) + self.assertIs(fixture.engine.conn, fixture.replacement) + self.assertIsNone(fixture.engine._reconnect_receipt) + self.assertIsNone(fixture.engine._runtime_vector_generation) + self.assertIsNone(fixture.engine._runtime_vector_connection) + self.assertFalse(fixture.engine._has_vectors) + self.assertFalse(fixture.engine.ready) + + def test_successive_verified_owners_share_one_ten_second_budget(self) -> None: + fixture = self.fixture() + first = fixture.automatic_owner(release_after=3.0) + release = first.released.action + successors: list[object] = [] + + def next_owner(timeout: float) -> None: + release(timeout) + successors.append(fixture.automatic_owner(release_after=2.0)) + + first.released.action = next_owner + fixture.runtime.open_engine_for_operation(fixture.engine) + self.assertEqual(fixture.deadlines, [110.0, 110.0]) + self.assertEqual(first.released.waits, [10.0]) + self.assertEqual(successors[0].released.waits, [7.0]) + self.assertEqual(fixture.clock.now, 105.0) + self.assertEqual(fixture.wait_scopes, [0, 0]) + self.assertEqual(fixture.open_calls, 3) + self.assertEqual(fixture.factory_calls, 1) + + def test_timed_out_owner_does_not_reset_budget_or_poisoned_state(self) -> None: + fixture = self.fixture() + origin = fixture.automatic_owner(release_after=20.0) + with self.assertRaises(fixture.maintenance.MaintenanceBusyError): + fixture.runtime.open_engine_for_operation(fixture.engine) + self.assertEqual(origin.released.waits, [10.0]) + self.assertEqual(fixture.clock.now, 110.0) + self.assertEqual(fixture.deadlines, [110.0]) + self.assertEqual(fixture.open_calls, 1) + self.assertEqual(fixture.factory_calls, 0) + self.assertIs(fixture.engine.conn, fixture.handle) + self.assertTrue(fixture.engine.ready) + + def test_initial_attempt_consumes_the_same_budget(self) -> None: + fixture = self.fixture() + origin = fixture.automatic_owner() + opening = fixture.engine._open_connection_handle + + def delayed() -> None: + fixture.clock.now += 3.0 + opening() + + fixture.engine._open_connection_handle = delayed + fixture.runtime.open_engine_for_operation(fixture.engine) + self.assertEqual(origin.released.waits, [7.0]) + self.assertEqual(fixture.deadlines, [110.0]) + + def test_swapped_handle_before_wait_cannot_reuse_cached_failure(self) -> None: + for replacement in (None, Handle()): + with self.subTest(replacement=replacement): + fixture = self.fixture() + origin = fixture.automatic_owner() + + def swap() -> bool: + fixture.engine.conn = replacement + return False + + fixture.runtime.serving.current_thread_admitted = swap + with self.assertRaises(fixture.maintenance.MaintenanceBusyError): + fixture.runtime.open_engine_for_operation(fixture.engine) + self.assertEqual(origin.released.waits, []) + self.assertEqual(fixture.open_calls, 1) + self.assertEqual(fixture.factory_calls, 0) + self.assertIs(fixture.engine.conn, replacement) + + def test_swapped_handle_during_wait_is_rejected_before_retry(self) -> None: + for replacement in (None, Handle()): + with self.subTest(replacement=replacement): + fixture = self.fixture() + origin = fixture.automatic_owner() + release = origin.released.action + + def swap(timeout: float) -> None: + release(timeout) + fixture.engine.conn = replacement + + origin.released.action = swap + with self.assertRaises(fixture.maintenance.MaintenanceBusyError): + fixture.runtime.open_engine_for_operation(fixture.engine) + self.assertEqual(origin.released.waits, [10.0]) + self.assertEqual(fixture.open_calls, 1) + self.assertEqual(fixture.factory_calls, 0) + self.assertIs(fixture.engine.conn, replacement) + + def test_wrong_engine_or_wrong_handle_attestation_never_grants_wait(self) -> None: + for kind in ("engine", "handle"): + with self.subTest(kind=kind): + fixture = self.fixture() + origin = fixture.automatic_owner() + error = fixture.maintenance.MaintenanceBusyError("synthetic stale failure") + error._automatic_owner = origin + error._reconnect_engine = fixture.engine if kind == "handle" else object() + error._reconnect_handle = fixture.handle if kind == "engine" else Handle() + + def reject() -> None: + raise error + + fixture.engine._open_connection_handle = reject + with self.assertRaises(fixture.maintenance.MaintenanceBusyError) as raised: + fixture.runtime.open_engine_for_operation(fixture.engine) + self.assertIs(raised.exception, error) + self.assertEqual(origin.released.waits, []) + self.assertEqual(fixture.wait_scopes, []) + + def test_receiptless_successor_cannot_inherit_previous_origin(self) -> None: + fixture = self.fixture() + origin = fixture.automatic_owner() + release = origin.released.action + + def receiptless(timeout: float) -> None: + release(timeout) + fixture.held_paths[fixture.path] = 991 + fixture.coordinators.clear() + + origin.released.action = receiptless + with self.assertRaises(fixture.maintenance.MaintenanceBusyError) as raised: + fixture.runtime.open_engine_for_operation(fixture.engine) + self.assertIsNone(raised.exception._automatic_owner) + self.assertEqual(origin.released.waits, [10.0]) + self.assertEqual(fixture.deadlines, [110.0, 110.0]) + self.assertEqual(fixture.open_calls, 2) + self.assertEqual(fixture.factory_calls, 0) + + def test_unrelated_sql_failure_and_borrowed_io_error_are_not_retried(self) -> None: + for reason, borrowed in (("other_error", False), ("io_error", True)): + with self.subTest(reason=reason, borrowed=borrowed): + fixture = self.fixture(reason) + fixture.handle.in_transaction = borrowed + origin = fixture.automatic_owner() + with self.assertRaises(sqlite3.OperationalError): + fixture.runtime.open_engine_for_operation(fixture.engine) + self.assertEqual(origin.released.waits, []) + self.assertEqual(fixture.wait_scopes, []) + self.assertEqual(fixture.factory_calls, 0) + self.assertEqual(fixture.open_calls, 1) + + def test_cancelled_origin_during_wait_never_retries(self) -> None: + fixture = self.fixture() + origin = fixture.automatic_owner() + + def cancel(timeout: float) -> None: + fixture.clock.now += 1.0 + origin.cancelled.set() + origin.released.set() + + origin.released.action = cancel + with self.assertRaises(fixture.maintenance.MaintenanceBusyError): + fixture.runtime.open_engine_for_operation(fixture.engine) + self.assertEqual(origin.released.waits, [10.0]) + self.assertEqual(fixture.open_calls, 1) + self.assertEqual(fixture.factory_calls, 0) + self.assertIs(fixture.engine.conn, fixture.handle) + + def test_registry_or_coordinator_contention_does_not_block_wait_eligibility(self) -> None: + for kind in ("registry", "coordinator"): + with self.subTest(kind=kind): + fixture = self.fixture() + origin = fixture.automatic_owner() + lock = fixture.registry if kind == "registry" else fixture.coordinators[fixture.path]._mutex + opening = fixture.engine._open_connection_handle + + def occupied_after_failure() -> None: + try: + opening() + except fixture.maintenance.MaintenanceBusyError: + lock.acquire() + raise + + fixture.engine._open_connection_handle = occupied_after_failure + try: + with self.assertRaises(fixture.maintenance.MaintenanceBusyError): + fixture.runtime.open_engine_for_operation(fixture.engine) + self.assertEqual(origin.released.waits, []) + self.assertEqual(fixture.open_calls, 1) + self.assertEqual(fixture.factory_calls, 0) + finally: + lock.release() + + def test_expired_initial_attempt_never_enters_wait(self) -> None: + fixture = self.fixture() + origin = fixture.automatic_owner() + opening = fixture.engine._open_connection_handle + + def expired() -> None: + fixture.clock.now += 10.0 + opening() + + fixture.engine._open_connection_handle = expired + with self.assertRaises(fixture.maintenance.MaintenanceBusyError): + fixture.runtime.open_engine_for_operation(fixture.engine) + self.assertEqual(origin.released.waits, []) + self.assertEqual(fixture.wait_scopes, []) + self.assertEqual(fixture.factory_calls, 0) + self.assertEqual(fixture.open_calls, 1) + + +if __name__ == "__main__": + unittest.main() diff --git a/truememory/engine.py b/truememory/engine.py index 08547e2e..97b0d04e 100644 --- a/truememory/engine.py +++ b/truememory/engine.py @@ -501,7 +501,7 @@ def replacement_reason() -> str | None: finally: if not installed: candidate.close() - from truememory.maintenance import maintenance_owner + from truememory.maintenance import MaintenanceBusyError, maintenance_owner if str(path) != ":memory:": path.parent.mkdir(parents=True, exist_ok=True) try: @@ -510,21 +510,30 @@ def replacement_reason() -> str | None: pass # Keep same-Engine openers serialized, but never claim maintenance # ownership while holding the connection writer. - with maintenance_owner(None if str(path) == ":memory:" else path), \ - _connection_read(self.conn, self._write_lock, None, None, allow_closed=True): - reason = replacement_reason() - if reason is None: - return - if reason == "io_error": - 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 - self._reconnect_receipt = None - self.conn = _open_serving_connection_at_path(path, create_db) + poisoned = self.conn if reason in {"closed", "io_error"} else None + try: + with maintenance_owner(None if str(path) == ":memory:" else path), \ + _connection_read(self.conn, self._write_lock, None, None, allow_closed=True): + reason = replacement_reason() + if reason is None: + return + if reason == "io_error": + 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 + self._reconnect_receipt = None + self.conn = _open_serving_connection_at_path(path, create_db) + except MaintenanceBusyError as error: + if poisoned is not None and self.conn is poisoned: + # Public opening may retry this unusable handle only after + # validating the exact automatic owner on the same error. + error._reconnect_engine = self + error._reconnect_handle = poisoned + raise @staticmethod def _reconnect_registry(conn: sqlite3.Connection) -> tuple | None: diff --git a/truememory/tier_switch/runtime.py b/truememory/tier_switch/runtime.py index f602ea3a..c589f266 100644 --- a/truememory/tier_switch/runtime.py +++ b/truememory/tier_switch/runtime.py @@ -534,13 +534,27 @@ def wrapped(conn: sqlite3.Connection, *args: object, **kwargs: object) -> R: def open_engine_for_operation(engine: object) -> None: """Bound fresh public admission without waiting inside an admitted call.""" - if engine.conn is not None: - engine._open_connection_handle() - return from truememory.maintenance import MaintenanceBusyError, wait_for_automatic_owner from truememory.storage import DEFAULT_BUSY_TIMEOUT_MS deadline = time.monotonic() + DEFAULT_BUSY_TIMEOUT_MS / 1000 + def retry_allowed(error: MaintenanceBusyError) -> bool: + cached = getattr(error, "_reconnect_engine", None) + if cached is None: + return engine.conn is None + return (cached is engine and engine.conn is not None + and getattr(error, "_reconnect_handle", None) is engine.conn) + + failure = None + if engine.conn is not None: + try: + engine._open_connection_handle() + return + except MaintenanceBusyError as error: + if getattr(error, "_reconnect_engine", None) is not engine or not retry_allowed(error): + raise + failure = error + def outside_admission() -> None: lock = engine._write_lock owned = (lock.owned_by_current_thread() if isinstance(lock, ConnectionWriteLock) @@ -554,21 +568,24 @@ def outside_admission() -> None: if isinstance(engine._write_lock, ConnectionWriteLock) and engine._write_lock.owned_by_current_thread(): raise MaintenanceBusyError("Fresh public opening requires an unbound writer") - failure = None while True: if failure is not None: outside_admission() if time.monotonic() >= deadline: raise failure + if not retry_allowed(failure): + raise failure + wait_for_automatic_owner(failure, deadline=deadline) + outside_admission() + if not retry_allowed(failure): + raise failure try: engine._open_connection_handle() return except MaintenanceBusyError as error: failure = error - if engine.conn is not None: + if not retry_allowed(error): raise - outside_admission() - wait_for_automatic_owner(error, deadline=deadline) def _legacy_initialization_is_local(conn: sqlite3.Connection, operation: ServingOperation) -> bool: From fe11731e058031d0d2f318e1d528c87e9336084e Mon Sep 17 00:00:00 2001 From: buildingjoshbetter Date: Sun, 11 Oct 2026 00:11:15 -0500 Subject: [PATCH 07/11] fix: Assert active garbage collection during reconnect controls --- tests/test-tier-reconnect-admission-684.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/tests/test-tier-reconnect-admission-684.py b/tests/test-tier-reconnect-admission-684.py index 8a333d2d..6acfbb41 100644 --- a/tests/test-tier-reconnect-admission-684.py +++ b/tests/test-tier-reconnect-admission-684.py @@ -4,6 +4,7 @@ import ast from collections.abc import Callable, Iterator from contextlib import contextmanager, nullcontext +import gc import os from pathlib import Path import sqlite3 @@ -283,6 +284,7 @@ def fixture(self, reason: str | None = "closed") -> Fixture: return fixture def test_closed_handle_retries_exact_automatic_owner_outside_engine_locks(self) -> None: + self.assertTrue(gc.isenabled()) fixture = self.fixture() origin = fixture.automatic_owner() fixture.runtime.open_engine_for_operation(fixture.engine) @@ -296,6 +298,7 @@ def test_closed_handle_retries_exact_automatic_owner_outside_engine_locks(self) self.assertFalse(fixture.engine._runtime_initialized) def test_io_error_handle_retries_same_strict_automatic_owner(self) -> None: + self.assertTrue(gc.isenabled()) fixture = self.fixture("io_error") origin = fixture.automatic_owner() fixture.runtime.open_engine_for_operation(fixture.engine) From a32c46ec1626e4add5cfdb6593afadd7d66a0e40 Mon Sep 17 00:00:00 2001 From: buildingjoshbetter Date: Sun, 11 Oct 2026 00:15:12 -0500 Subject: [PATCH 08/11] fix: Refuse owned writers before cached handle opening --- tests/test-tier-reconnect-admission-684.py | 108 ++++++++++++++++++++- truememory/tier_switch/runtime.py | 4 +- 2 files changed, 105 insertions(+), 7 deletions(-) diff --git a/tests/test-tier-reconnect-admission-684.py b/tests/test-tier-reconnect-admission-684.py index 6acfbb41..35340824 100644 --- a/tests/test-tier-reconnect-admission-684.py +++ b/tests/test-tier-reconnect-admission-684.py @@ -16,6 +16,7 @@ import unittest import uuid import weakref +from unittest.mock import patch ROOT = Path(__file__).resolve().parents[1] @@ -108,7 +109,7 @@ def chmod(self, mode: int) -> None: class Handle: def __init__(self, reason: str | None = None) -> None: self.reason = reason - self.in_transaction = False + self._in_transaction = False self.probes = 0 self.closes = 0 @@ -127,9 +128,23 @@ def close(self) -> None: self.reason = "closed" self.closes += 1 + @property + def in_transaction(self) -> bool: + if self.reason == "closed": + raise sqlite3.ProgrammingError("Cannot operate on a closed database.") + return self._in_transaction + + @in_transaction.setter + def in_transaction(self, borrowed: bool) -> None: + self._in_transaction = borrowed + + +class SyntheticEngine: + pass + class Fixture: - def __init__(self, reason: str | None = "closed") -> None: + def __init__(self, reason: str | None = "closed", *, actual_lock: bool = False) -> None: self.clock = Clock() self.path = SyntheticPath() self.scope_depth = 0 @@ -151,7 +166,8 @@ def __init__(self, reason: str | None = "closed") -> None: self.reopened: tuple[Handle, bool] | None = None self.gate_error = type("ServingLeaseError", (RuntimeError,), {}) self.runtime_error = type("TierRuntimeError", (RuntimeError,), {}) - self.engine = types.SimpleNamespace( + self.engine = SyntheticEngine() + self.engine.__dict__.update( conn=self.handle, db_path=self.path, _write_lock=Writer(), _init_lock=object(), _reconnect_receipt=None, ready=True, _runtime_initialized=True, _has_vectors=True, _runtime_vector_generation="synthetic-generation", _runtime_vector_connection=self.handle, @@ -201,6 +217,18 @@ def release(fd: int) -> None: }) self.runtime._connection_read = connection_read self.runtime._open_serving_connection_at_path = factory + if actual_lock: + locks = definitions("tier_switch/runtime.py", { + "ConnectionWriteLock", "_connection_read", "_check", "_remaining", + }, { + "threading": threading, "sqlite3": sqlite3, "time": self.clock, + "contextmanager": contextmanager, "TierRuntimeError": self.runtime_error, + "serving": self.runtime.serving, + }) + self.runtime.ConnectionWriteLock = locks.ConnectionWriteLock + self.runtime._connection_read = locks._connection_read + self.engine._write_lock = locks.ConnectionWriteLock(self.engine) + self.engine._init_lock = threading.Lock() self.storage = types.ModuleType("truememory.storage") storage_tree = ast.parse(source("storage.py")) self.storage.DEFAULT_BUSY_TIMEOUT_MS = next( @@ -278,8 +306,8 @@ def release(timeout: float) -> None: class TestReconnectAdmission(unittest.TestCase): - def fixture(self, reason: str | None = "closed") -> Fixture: - fixture = Fixture(reason) + def fixture(self, reason: str | None = "closed", *, actual_lock: bool = False) -> Fixture: + fixture = Fixture(reason, actual_lock=actual_lock) self.addCleanup(fixture.close) return fixture @@ -643,6 +671,76 @@ def expired() -> None: self.assertEqual(fixture.factory_calls, 0) self.assertEqual(fixture.open_calls, 1) + def test_poisoned_current_writer_refuses_before_actual_nonreentrant_acquire(self) -> None: + self.assertTrue(gc.isenabled()) + for reason in ("closed", "io_error"): + with self.subTest(reason=reason): + fixture = self.fixture(reason, actual_lock=True) + origin = fixture.automatic_owner() + lock = fixture.engine._write_lock + with lock: + with patch.object(lock, "acquire", side_effect=AssertionError("Nonreentrant writer reacquired")) as acquired: + with self.assertRaises(fixture.maintenance.MaintenanceBusyError): + fixture.runtime.open_engine_for_operation(fixture.engine) + acquired.assert_not_called() + self.assertTrue(lock.owned_by_current_thread()) + self.assertIs(fixture.engine.conn, fixture.handle) + self.assertEqual(fixture.open_calls, 0) + self.assertEqual(origin.released.waits, []) + self.assertEqual(fixture.wait_scopes, []) + self.assertEqual(fixture.factory_calls, 0) + + def test_healthy_unowned_actual_writer_reuses_handle_inside_nested_admission(self) -> None: + for admitted in (False, True): + with self.subTest(admitted=admitted): + fixture = self.fixture(None, actual_lock=True) + fixture.admitted = admitted + fixture.local.operation = object() if admitted else None + origin = fixture.automatic_owner() + fixture.runtime.open_engine_for_operation(fixture.engine) + self.assertIs(fixture.engine.conn, fixture.handle) + self.assertEqual(fixture.handle.probes, 1) + self.assertEqual(fixture.open_calls, 1) + self.assertEqual(fixture.factory_calls, 0) + self.assertEqual(origin.released.waits, []) + self.assertFalse(fixture.engine._write_lock.locked()) + self.assertTrue(fixture.engine.ready) + + def test_healthy_borrowed_actual_writer_preserves_unowned_caller_transaction(self) -> None: + for admitted in (False, True): + with self.subTest(admitted=admitted): + fixture = self.fixture(None, actual_lock=True) + fixture.handle.in_transaction = True + fixture.admitted = admitted + fixture.local.operation = object() if admitted else None + origin = fixture.automatic_owner() + fixture.runtime.open_engine_for_operation(fixture.engine) + self.assertTrue(fixture.handle.in_transaction) + self.assertIs(fixture.engine.conn, fixture.handle) + self.assertEqual(fixture.handle.probes, 1) + self.assertEqual(fixture.open_calls, 1) + self.assertEqual(fixture.factory_calls, 0) + self.assertEqual(origin.released.waits, []) + self.assertFalse(fixture.engine._write_lock.locked()) + + def test_current_actual_writer_refuses_before_healthy_or_borrowed_opening(self) -> None: + for borrowed in (False, True): + with self.subTest(borrowed=borrowed): + fixture = self.fixture(None, actual_lock=True) + fixture.handle.in_transaction = borrowed + lock = fixture.engine._write_lock + with lock: + with patch.object(lock, "acquire", side_effect=AssertionError("Owned writer reacquired")) as acquired: + with self.assertRaises(fixture.maintenance.MaintenanceBusyError): + fixture.runtime.open_engine_for_operation(fixture.engine) + acquired.assert_not_called() + self.assertTrue(lock.owned_by_current_thread()) + self.assertEqual(fixture.handle.in_transaction, borrowed) + self.assertEqual(fixture.handle.probes, 0) + self.assertEqual(fixture.open_calls, 0) + self.assertEqual(fixture.factory_calls, 0) + self.assertEqual(fixture.wait_scopes, []) + if __name__ == "__main__": unittest.main() diff --git a/truememory/tier_switch/runtime.py b/truememory/tier_switch/runtime.py index c589f266..2cdf3402 100644 --- a/truememory/tier_switch/runtime.py +++ b/truememory/tier_switch/runtime.py @@ -537,6 +537,8 @@ def open_engine_for_operation(engine: object) -> None: from truememory.maintenance import MaintenanceBusyError, wait_for_automatic_owner from truememory.storage import DEFAULT_BUSY_TIMEOUT_MS deadline = time.monotonic() + DEFAULT_BUSY_TIMEOUT_MS / 1000 + if isinstance(engine._write_lock, ConnectionWriteLock) and engine._write_lock.owned_by_current_thread(): + raise MaintenanceBusyError("Fresh public opening requires an unbound writer") def retry_allowed(error: MaintenanceBusyError) -> bool: cached = getattr(error, "_reconnect_engine", None) @@ -566,8 +568,6 @@ def outside_admission() -> None: if owned or admitted: raise MaintenanceBusyError("Fresh public opening requires an unbound caller") - if isinstance(engine._write_lock, ConnectionWriteLock) and engine._write_lock.owned_by_current_thread(): - raise MaintenanceBusyError("Fresh public opening requires an unbound writer") while True: if failure is not None: outside_admission() From fe25897cf72dd7d1688b6b45e3d7eeb28888f871 Mon Sep 17 00:00:00 2001 From: buildingjoshbetter Date: Sun, 11 Oct 2026 00:22:41 -0500 Subject: [PATCH 09/11] fix: Bind public boundary fixtures to actual admission imports --- tests/test-tier-public-boundaries-795.py | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/tests/test-tier-public-boundaries-795.py b/tests/test-tier-public-boundaries-795.py index bb6040d3..e9be2fd3 100644 --- a/tests/test-tier-public-boundaries-795.py +++ b/tests/test-tier-public-boundaries-795.py @@ -93,8 +93,15 @@ def setUp(self) -> None: self.fixture.setUp() self.api, self.conn = self.fixture.api, self.fixture.conn self.modules = self.fixture.modules + maintenance = definitions("maintenance.py", {"MaintenanceBusyError"}, self.namespace()) self.modules["maintenance"] = types.SimpleNamespace( + MaintenanceBusyError=maintenance.MaintenanceBusyError, + wait_for_automatic_owner=Mock(side_effect=AssertionError("Unexpected automatic owner wait in healthy fixture")), maintenance_owner=lambda path: contextlib.nullcontext(), connection_database_path=lambda conn: None) + storage = ast.parse((ROOT / "truememory/storage.py").read_text(encoding="utf-8")) + self.modules["storage"].DEFAULT_BUSY_TIMEOUT_MS = next( + ast.literal_eval(node.value) for node in storage.body if isinstance(node, ast.Assign) + and any(isinstance(target, ast.Name) and target.id == "DEFAULT_BUSY_TIMEOUT_MS" for target in node.targets)) enter_context(self, patch.dict(sys.modules, {"truememory.tier_switch.runtime": self.api})) def imports(self, name, globals=None, locals=None, fromlist=(), level=0): From 95e057fed9cb0b1c605c2c68d96363ece9ef5442 Mon Sep 17 00:00:00 2001 From: buildingjoshbetter Date: Sun, 11 Oct 2026 00:29:02 -0500 Subject: [PATCH 10/11] fix: Load boundary exception without subclass collaborators --- tests/test-tier-public-boundaries-795.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test-tier-public-boundaries-795.py b/tests/test-tier-public-boundaries-795.py index e9be2fd3..38536dcf 100644 --- a/tests/test-tier-public-boundaries-795.py +++ b/tests/test-tier-public-boundaries-795.py @@ -93,7 +93,7 @@ def setUp(self) -> None: self.fixture.setUp() self.api, self.conn = self.fixture.api, self.fixture.conn self.modules = self.fixture.modules - maintenance = definitions("maintenance.py", {"MaintenanceBusyError"}, self.namespace()) + maintenance = definitions("maintenance.py", {"MaintenanceBusyError"}, {}) self.modules["maintenance"] = types.SimpleNamespace( MaintenanceBusyError=maintenance.MaintenanceBusyError, wait_for_automatic_owner=Mock(side_effect=AssertionError("Unexpected automatic owner wait in healthy fixture")), From 5873e5a4978f37bcd31772d6840717fb217df713 Mon Sep 17 00:00:00 2001 From: buildingjoshbetter Date: Sun, 11 Oct 2026 00:44:42 -0500 Subject: [PATCH 11/11] fix: Bind MCP fixtures to actual opener collaborators --- tests/test-tier-mcp-serving-795.py | 35 +++++++++++++++++++++++++----- 1 file changed, 30 insertions(+), 5 deletions(-) diff --git a/tests/test-tier-mcp-serving-795.py b/tests/test-tier-mcp-serving-795.py index ab5f04fc..9187726f 100644 --- a/tests/test-tier-mcp-serving-795.py +++ b/tests/test-tier-mcp-serving-795.py @@ -6,14 +6,16 @@ import json import logging import runpy +import sqlite3 import threading +import time import types import unittest from collections.abc import Iterator from concurrent.futures import ThreadPoolExecutor from contextlib import contextmanager from pathlib import Path -from unittest.mock import patch +from unittest.mock import Mock, patch ROOT = Path(__file__).resolve().parents[1] GATE = runpy.run_path(str(ROOT / "truememory/tier_switch/serving.py")) @@ -89,9 +91,30 @@ def reject(error): runtime = types.SimpleNamespace(serving_operation=serving_operation, raise_if_serving_rejection=reject) + maintenance_tree = ast.parse((ROOT / "truememory/maintenance.py").read_text(encoding="utf-8")) + maintenance_error = next(node for node in maintenance_tree.body + if isinstance(node, ast.ClassDef) and node.name == "MaintenanceBusyError") + maintenance_namespace = {} + exec(compile(ast.fix_missing_locations(ast.Module(body=[maintenance_error], type_ignores=[])), + "actual-maintenance-error", "exec"), maintenance_namespace) + storage_tree = ast.parse((ROOT / "truememory/storage.py").read_text(encoding="utf-8")) + budget = next(ast.literal_eval(node.value) for node in storage_tree.body if isinstance(node, ast.Assign) + and any(isinstance(target, ast.Name) and target.id == "DEFAULT_BUSY_TIMEOUT_MS" + for target in node.targets)) + modules = { + "truememory.tier_switch.runtime": runtime, + "truememory.maintenance": types.SimpleNamespace( + MaintenanceBusyError=maintenance_namespace["MaintenanceBusyError"], + wait_for_automatic_owner=Mock(side_effect=AssertionError("Unexpected automatic owner wait in healthy MCP fixture"))), + "truememory.storage": types.SimpleNamespace(DEFAULT_BUSY_TIMEOUT_MS=budget), + } + def safe_import(name, globals=None, locals=None, fromlist=(), level=0): - if name == "truememory.tier_switch.runtime": - return runtime + if name in modules: + return modules[name] + if name.startswith("truememory") or name.split(".")[0] in { + "numpy", "torch", "sqlite_vec", "model2vec", "sentence_transformers"}: + raise AssertionError("Unexpected native/application import in MCP fixture: " + name) return builtins.__import__(name, globals, locals, fromlist, level) runtime_tree = ast.parse((ROOT / "truememory/tier_switch/runtime.py").read_text(encoding="utf-8")) @@ -99,11 +122,13 @@ def safe_import(name, globals=None, locals=None, fromlist=(), level=0): context_definitions = [node for node in runtime_tree.body if isinstance(node, ast.FunctionDef) and node.name in context_names] self.assertEqual({node.name for node in context_definitions}, context_names) + writer_class = next(node for node in runtime_tree.body + if isinstance(node, ast.ClassDef) and node.name == "ConnectionWriteLock") context_namespace = dict(__builtins__=dict(vars(builtins), __import__=safe_import), contextmanager=contextmanager, serving_operation=serving_operation, - current_operation=lambda conn: None) + current_operation=lambda conn: None, time=time, threading=threading, sqlite3=sqlite3) future = ast.ImportFrom(module="__future__", names=[ast.alias(name="annotations")], level=0) - exec(compile(ast.fix_missing_locations(ast.Module(body=[future, *context_definitions], type_ignores=[])), + exec(compile(ast.fix_missing_locations(ast.Module(body=[future, writer_class, *context_definitions], type_ignores=[])), "actual-public-context", "exec"), context_namespace) runtime.engine_serving_operation = context_namespace["engine_serving_operation"]