From 9fb995575ec3addb6f7a8e059944acf80d4056f5 Mon Sep 17 00:00:00 2001 From: Sinity Date: Sun, 19 Jul 2026 16:33:36 +0200 Subject: [PATCH] feat(storage): rescue retired-tier embedding vectors by content-hash Problem: a prior incident retired embeddings.db.v2-retired-20260718 (5.8GB, 776,895 Voyage vectors) while index.db was rebuilt from source. Re-embedding all 777K messages through the Voyage API is expensive and rate-limited; the retired tier's vectors are still valid for any message whose content_hash is unchanged. Solution: `polylogue ops maintenance embeddings-rescue` (`--plan` read-only census, `--yes` apply) in polylogue/storage/embeddings/rescue.py. Rescue is scoped to whole sessions, not individual messages: embed_archive_session_sync always re-embeds every eligible message of a session it selects in one atomic write (complete_embedding_attempt_success), never consulting pre-existing per-message vectors, so a partial per-message match saves nothing -- only a session where 100% of its eligible messages have an exact retired (message_id, content_hash, model) match is worth rescuing. Rescued sessions publish through the same begin_embedding_attempt + complete_embedding_attempt_success primitives the live embed path uses, so they read as already-fresh to the daemon's own selection predicate (idempotent reruns, resumable via --limit). Apply mode always samples rescued vectors for byte-identity against the source. Refuses --yes while polylogued is running (offline-only in this version, per the offline_maintenance_block_reason guard pattern already used by embedding-orphan-reconcile); a daemon-coordinator route is a possible follow-up, not implemented here. Verification: - devtools test tests/unit/storage/test_embedding_rescue.py tests/unit/cli/test_embeddings_rescue_cli.py -> 12 passed, including an anti-vacuity case that monkeypatches the retired-vector reader to hand back a mutated vector and asserts the sample-verification step catches it (verified_sample_ok < verified_sample_total, report.ok is False). - devtools verify --quick -> exit 0 (ruff format/check, mypy --strict, render all --check, topology/layering/docs-coverage/etc.) - Live read-only plan smoke against the real archive (mid-rebuild, 2626/eventual sessions) and the real retired file: eligible_sessions=2541, fully_rescuable_sessions=703, rescuable_messages=14458, partial_sessions=645 (6549 matched messages left unrescued -- correct, per the session-level design above), skipped_missing=23541, skipped_hash_mismatch=17744, skipped_model_mismatch=0. Numbers will rise once the index promote completes; live apply is coordinator-owned, not run here. Ref polylogue-04kl Co-Authored-By: Claude --- docs/maintenance.md | 27 + docs/plans/topology-target.yaml | 10 +- docs/topology-status.md | 6 +- .../cli/commands/maintenance/__init__.py | 6 + .../maintenance/_embeddings_rescue.py | 171 +++++ polylogue/storage/embeddings/rescue.py | 602 ++++++++++++++++++ tests/unit/cli/test_embeddings_rescue_cli.py | 120 ++++ tests/unit/storage/test_embedding_rescue.py | 462 ++++++++++++++ 8 files changed, 1400 insertions(+), 4 deletions(-) create mode 100644 polylogue/cli/commands/maintenance/_embeddings_rescue.py create mode 100644 polylogue/storage/embeddings/rescue.py create mode 100644 tests/unit/cli/test_embeddings_rescue_cli.py create mode 100644 tests/unit/storage/test_embedding_rescue.py diff --git a/docs/maintenance.md b/docs/maintenance.md index 20deada83a..281f73daf6 100644 --- a/docs/maintenance.md +++ b/docs/maintenance.md @@ -192,6 +192,33 @@ Exit code is non-zero when any check reports `error` (or, with `--strict`, temporarily busy under a concurrent rebuild — never aborts the rest; each check independently reports its own outcome. +### `polylogue ops maintenance embeddings-rescue` — retired-tier vector rescue + +Break-glass, offline-only migration for a retired `embeddings.db.v2-retired-*` +tier left behind by an index rebuild/incident. `--plan` (default) is a +read-only census of vectors rescuable by exact `(message_id, content_hash, +model)` match against the live index; `--yes` copies vectors for every +*fully rescuable session* directly into `embeddings.db`, skipping a live +provider re-embed for those sessions. Rescue is session-, not message-, +granular: the live embed path always re-embeds every eligible message of a +session it selects in one atomic write, so a session only avoids a real API +call when *all* of its eligible messages have an exact retired match — a +partial match saves nothing and is left for the real embed pass. + +```bash +polylogue ops maintenance embeddings-rescue --source /path/to/embeddings.db.v2-retired-20260718 +polylogue ops maintenance embeddings-rescue --source /path/to/retired.db --yes --limit 500 +polylogue ops maintenance embeddings-rescue --source /path/to/retired.db --output-format json | jq . +``` + +Refuses `--yes` while `polylogued` is running for this archive (offline-only +in this version). Publication goes through the same generation-guarded +primitives the live embed path uses, so a rescued session reads as +already-embedded to the daemon's own freshness check — a second invocation +is a no-op over already-rescued sessions (idempotent, resumable via +`--limit`). Apply mode always reports a byte-identity sample check +(`--sample-verify-count`) against the source vectors. + ### `--operation-id` and `--resume`: worked example Replay execution writes a small JSON state file under diff --git a/docs/plans/topology-target.yaml b/docs/plans/topology-target.yaml index a9ad8cf56b..8027a17a89 100644 --- a/docs/plans/topology-target.yaml +++ b/docs/plans/topology-target.yaml @@ -900,7 +900,7 @@ files: target: polylogue/cli/commands/judge.py owner: stable - path: polylogue/cli/commands/maintenance/__init__.py - loc: 152 + loc: 158 target: polylogue/cli/commands/maintenance/__init__.py owner: stable - path: polylogue/cli/commands/maintenance/_archive_plan.py @@ -935,6 +935,10 @@ files: loc: 164 target: polylogue/cli/commands/maintenance/_embeddings.py owner: stable + - path: polylogue/cli/commands/maintenance/_embeddings_rescue.py + loc: 172 + target: polylogue/cli/commands/maintenance/_embeddings_rescue.py + owner: stable - path: polylogue/cli/commands/maintenance/_migrate_tier.py loc: 90 target: polylogue/cli/commands/maintenance/_migrate_tier.py @@ -3505,6 +3509,10 @@ files: loc: 558 target: polylogue/storage/embeddings/reconcile.py owner: stable + - path: polylogue/storage/embeddings/rescue.py + loc: 606 + target: polylogue/storage/embeddings/rescue.py + owner: stable - path: polylogue/storage/embeddings/sql.py loc: 102 target: polylogue/storage/embeddings/sql.py diff --git a/docs/topology-status.md b/docs/topology-status.md index 5378443175..899b9c0525 100644 --- a/docs/topology-status.md +++ b/docs/topology-status.md @@ -28,12 +28,12 @@ Generated by `devtools render topology-status`. Reads `docs/plans/topology-targe ### Summary -- **Stable** (no move scoped): 879 +- **Stable** (no move scoped): 881 - **Kernel** (polylogue/ root): 8 - **Primitives** (storage-root): 19 - **TBD** (cell needs explicit assignment): 9 -- **Total declared**: 1050 -- **Realized polylogue/**/*.py**: 1050 files declared +- **Total declared**: 1052 +- **Realized polylogue/**/*.py**: 1052 files declared ### TBD cells (require explicit routing) diff --git a/polylogue/cli/commands/maintenance/__init__.py b/polylogue/cli/commands/maintenance/__init__.py index 588cded289..f4df0f5b40 100644 --- a/polylogue/cli/commands/maintenance/__init__.py +++ b/polylogue/cli/commands/maintenance/__init__.py @@ -121,6 +121,12 @@ "embedding_orphan_reconcile_command", "Inspect (default) or reconcile embeddings.db rows orphaned by an index rebuild.", ), + ( + "embeddings-rescue", + "_embeddings_rescue", + "embeddings_rescue_command", + "Inspect (default) or apply a content-hash vector rescue from a retired embeddings tier.", + ), ("gc-history", "_blob_gc", "gc_history_command", "Show recent blob-GC passes recorded in ``gc_generations``."), ("status", "_status", "status_command", "Inspect persisted maintenance operations (#1197)."), ( diff --git a/polylogue/cli/commands/maintenance/_embeddings_rescue.py b/polylogue/cli/commands/maintenance/_embeddings_rescue.py new file mode 100644 index 0000000000..24e9d20155 --- /dev/null +++ b/polylogue/cli/commands/maintenance/_embeddings_rescue.py @@ -0,0 +1,171 @@ +"""``maintenance embeddings-rescue``: copy vectors from a retired embeddings tier. + +Break-glass, offline-only migration path for polylogue-04kl: a retired +``embeddings.db.v2-retired-YYYYMMDD`` tier can hold hundreds of thousands of +Voyage vectors whose ``message_id`` + ``content_hash`` identity still matches +the freshly rebuilt index. ``--plan`` (default) is a read-only census; +``--yes`` copies vectors for every *fully rescuable* session (see +:mod:`polylogue.storage.embeddings.rescue` for why rescue is session-, not +message-, granular) directly into the live ``embeddings.db``, skipping a live +Voyage API re-embed for those sessions entirely. +""" + +from __future__ import annotations + +import json +from typing import TYPE_CHECKING + +import click + +from polylogue.config import Config +from polylogue.paths import archive_file_set_root_for_paths, archive_root, db_path, render_root + +if TYPE_CHECKING: + from polylogue.storage.embeddings.rescue import ( + EmbeddingRescueExecuteReport, + EmbeddingRescuePlanReport, + ) + + +@click.command("embeddings-rescue") +@click.option( + "--source", + "source_path", + required=True, + type=click.Path(exists=True, dir_okay=False), + help="Path to the retired embeddings.db tier to rescue vectors from.", +) +@click.option( + "--plan", + "plan_only", + is_flag=True, + default=True, + show_default=True, + help="Read-only census of rescuable vectors (default). Mutually exclusive with --yes.", +) +@click.option( + "--yes", + "apply", + is_flag=True, + help="Copy vectors for every fully-rescuable session into the live embeddings.db.", +) +@click.option( + "--limit", + type=int, + default=None, + help="Maximum number of sessions to rescue in this invocation (apply mode only; resumable).", +) +@click.option( + "--sample-limit", + type=int, + default=30, + show_default=True, + help="Maximum number of representative non-matched samples to include.", +) +@click.option( + "--sample-verify-count", + type=int, + default=20, + show_default=True, + help="Number of rescued vectors to sample-verify byte-identical against the source (apply mode only).", +) +@click.option( + "--output-format", + "output_format", + type=click.Choice(["plain", "json"]), + default="plain", + show_default=True, + help="Output format.", +) +def embeddings_rescue_command( + source_path: str, + plan_only: bool, + apply: bool, + limit: int | None, + sample_limit: int, + sample_verify_count: int, + output_format: str, +) -> None: + """Inspect (default) or apply a content-hash vector rescue from a retired embeddings tier.""" + del plan_only # --plan is documentation-only; --yes/apply is the sole control switch. + root = archive_file_set_root_for_paths(archive_root_path=archive_root(), db_anchor=db_path()) + index_db = root / "index.db" + + if not apply: + from polylogue.storage.embeddings.rescue import plan_embedding_rescue + + plan_report = plan_embedding_rescue(index_db, source_path, sample_size=sample_limit) + payload = {"mutates": False, **plan_report.to_dict()} + if output_format == "json": + click.echo(json.dumps(payload, indent=2, sort_keys=True)) + return + _render_plan_plain(plan_report) + return + + from polylogue.maintenance.offline_guard import running_daemon_pid + from polylogue.storage.embeddings.rescue import execute_embedding_rescue + + active_config = Config(archive_root=root, render_root=render_root(), sources=[], db_path=index_db) + daemon_pid = running_daemon_pid(active_config) + if daemon_pid is not None: + raise click.ClickException(f"embeddings rescue refused while polylogued PID {daemon_pid} is running") + + exec_report = execute_embedding_rescue( + index_db, + source_path, + limit=limit, + sample_size=sample_limit, + sample_verify_count=sample_verify_count, + mutation_authority="offline-exclusive", + ) + payload = {"mutates": True, **exec_report.to_dict()} + if output_format == "json": + click.echo(json.dumps(payload, indent=2, sort_keys=True)) + return + _render_execute_plain(exec_report) + + +def _render_plan_plain(report: EmbeddingRescuePlanReport) -> None: + click.echo("Embeddings rescue plan (read-only)") + click.echo(f"Index DB: {report.index_db}") + click.echo(f"Source DB: {report.source_db}") + click.echo(f"Model: {report.model}") + click.echo(f"Eligible sessions: {report.counts.eligible_sessions:,}") + click.echo( + f"Fully rescuable sessions: {report.counts.fully_rescuable_sessions:,} " + f"({report.counts.rescuable_messages:,} message(s))" + ) + click.echo( + f"Partial sessions: {report.counts.partial_sessions:,} " + f"({report.counts.partial_matched_messages:,} matched message(s), not written -- see command help)" + ) + click.echo( + "Skipped messages: " + f"missing={report.counts.skipped_missing:,} " + f"hash_mismatch={report.counts.skipped_hash_mismatch:,} " + f"model_mismatch={report.counts.skipped_model_mismatch:,}" + ) + if report.samples: + click.echo("Samples:") + for sample in report.samples[:5]: + click.echo(f" {sample.status} {sample.message_id} (session {sample.session_id})") + + +def _render_execute_plain(report: EmbeddingRescueExecuteReport) -> None: + click.echo("Embeddings rescue apply") + click.echo(f"Index DB: {report.index_db}") + click.echo(f"Source DB: {report.source_db}") + click.echo(f"Embeddings DB: {report.embeddings_db}") + click.echo(f"Model: {report.model}") + click.echo(f"Rescued: {report.rescued_sessions:,} session(s), {report.rescued_messages:,} message(s)") + click.echo( + "Skipped: " + f"already_fresh={report.skipped_already_fresh_sessions:,} " + f"race={report.skipped_race_sessions:,} " + f"missing={report.counts.skipped_missing:,} " + f"hash_mismatch={report.counts.skipped_hash_mismatch:,} " + f"model_mismatch={report.counts.skipped_model_mismatch:,}" + ) + click.echo(f"Sample verify: {report.verified_sample_ok:,}/{report.verified_sample_total:,} ok") + click.echo(f"More pending: {report.more_pending}") + click.echo(f"Status: {'OK' if report.ok else 'VERIFICATION FAILED'}") diff --git a/polylogue/storage/embeddings/rescue.py b/polylogue/storage/embeddings/rescue.py new file mode 100644 index 0000000000..7e56b503ee --- /dev/null +++ b/polylogue/storage/embeddings/rescue.py @@ -0,0 +1,602 @@ +"""Rescue vectors from a retired embeddings tier by exact content-hash match. + +polylogue-04kl: a prior incident retired an embeddings tier +(``embeddings.db.v2-retired-YYYYMMDD``) while the index was rebuilt from +source. That retired tier still holds hundreds of thousands of Voyage +vectors keyed by ``message_id`` + a 32-byte ``content_hash``. The freshly +rebuilt index computes the *same* content-hash identity for each message, so +an exact rescue is possible: a retired vector is safe to copy into the fresh +``embeddings.db`` whenever + +* its ``message_id`` still exists among the fresh index's *currently + eligible* (authored-prose) messages, +* its stored ``content_hash`` matches that message's current content hash + exactly, and +* its stored ``model`` matches the configured embedding model. + +Session-level granularity is not a stylistic choice -- it is required by how +catch-up materialization actually works. ``embed_archive_session_sync`` +(:mod:`polylogue.storage.embeddings.materialization`) always re-embeds +*every* eligible message of a session it selects, via one atomic +generation-guarded write (:func:`polylogue.storage.sqlite.archive_tiers. +embedding_write.complete_embedding_attempt_success`); it never consults +per-message vectors already present to skip a subset. So a session only +avoids a live API re-embed when *all* of its eligible messages are +rescuable -- a partial per-message match saves nothing, because the next +catch-up pass deletes and rewrites the whole session's vectors regardless. +Rescue therefore classifies and mutates at session granularity: a session +is ``fully_rescuable`` only when every one of its eligible messages has an +exact retired match; anything less is reported as a ``partial`` session and +left untouched (the real embed pass will cover it later at full API cost). + +Rescued sessions are published through the same guarded primitives the live +embed path uses (:func:`begin_embedding_attempt` + +:func:`complete_embedding_attempt_success`), so a rescued session's +``embedding_derivation_state``/``embedding_status`` rows look identical to +one embedded for real -- the daemon's freshness predicate +(``_archive_embedding_freshness_predicate``) will no longer select it, and a +second rescue run naturally treats it as already-fresh (idempotent, no +re-write). +""" + +from __future__ import annotations + +import sqlite3 +import struct +import time +from dataclasses import dataclass +from pathlib import Path +from typing import Literal + +from polylogue.storage.embeddings.identity import ( + EmbeddingRecipe, + EmbeddingSourceDigest, + message_embedding_derivation_key, +) +from polylogue.storage.embeddings.materialization import ( + archive_embeddable_messages_relation, + select_pending_archive_session_window, +) +from polylogue.storage.sqlite.connection_profile import open_readonly_connection +from polylogue.storage.sqlite.sqlite_vec_extension import try_load_sqlite_vec + +EmbeddingRescueMutationAuthority = Literal["offline-exclusive"] + +DEFAULT_SAMPLE_SIZE = 30 +DEFAULT_SAMPLE_VERIFY_COUNT = 20 + +_MISSING = "missing" +_MODEL_MISMATCH = "model_mismatch" +_HASH_MISMATCH = "hash_mismatch" +_MATCHED = "matched" + +_CLASSIFIED_CTE_TEMPLATE = """ +WITH eligible AS ( + SELECT e.message_id, e.session_id, e.content_hash + FROM {relation} +), +classified AS ( + SELECT + eligible.session_id AS session_id, + eligible.message_id AS message_id, + CASE + WHEN rm.message_id IS NULL THEN '{missing}' + WHEN rv.id IS NULL THEN '{missing}' + WHEN rm.model != :model THEN '{model_mismatch}' + WHEN eligible.content_hash IS NULL THEN '{hash_mismatch}' + WHEN rm.content_hash != eligible.content_hash THEN '{hash_mismatch}' + ELSE '{matched}' + END AS status + FROM eligible + LEFT JOIN retired.message_embeddings_meta AS rm ON rm.message_id = eligible.message_id + LEFT JOIN retired.message_embeddings_rowids AS rv ON rv.id = eligible.message_id +) +""" + + +def _classified_cte(relation: str) -> str: + return _CLASSIFIED_CTE_TEMPLATE.format( + relation=relation, + missing=_MISSING, + model_mismatch=_MODEL_MISMATCH, + hash_mismatch=_HASH_MISMATCH, + matched=_MATCHED, + ) + + +@dataclass(frozen=True, slots=True) +class EmbeddingRescueSample: + """One representative message classification surfaced for operator inspection.""" + + status: str # "missing" | "hash_mismatch" | "model_mismatch" + message_id: str + session_id: str + + def to_dict(self) -> dict[str, object]: + return {"status": self.status, "message_id": self.message_id, "session_id": self.session_id} + + +@dataclass(frozen=True, slots=True) +class EmbeddingRescueCounts: + """Shared session/message classification counts for plan and execute reports.""" + + eligible_sessions: int + fully_rescuable_sessions: int + rescuable_messages: int + partial_sessions: int + partial_matched_messages: int + skipped_missing: int + skipped_hash_mismatch: int + skipped_model_mismatch: int + + +@dataclass(frozen=True, slots=True) +class EmbeddingRescuePlanReport: + """Read-only rescue-candidate census: ``ops maintenance embeddings-rescue --plan``.""" + + index_db: str + source_db: str + model: str + counts: EmbeddingRescueCounts + samples: tuple[EmbeddingRescueSample, ...] + + def to_dict(self) -> dict[str, object]: + return { + "mode": "plan", + "index_db": self.index_db, + "source_db": self.source_db, + "model": self.model, + "eligible_sessions": self.counts.eligible_sessions, + "fully_rescuable_sessions": self.counts.fully_rescuable_sessions, + "rescuable_messages": self.counts.rescuable_messages, + "partial_sessions": self.counts.partial_sessions, + "partial_matched_messages": self.counts.partial_matched_messages, + "skipped_missing": self.counts.skipped_missing, + "skipped_hash_mismatch": self.counts.skipped_hash_mismatch, + "skipped_model_mismatch": self.counts.skipped_model_mismatch, + "samples": [sample.to_dict() for sample in self.samples], + } + + +@dataclass(frozen=True, slots=True) +class EmbeddingRescueExecuteReport: + """Apply-mode rescue outcome: ``ops maintenance embeddings-rescue --yes``.""" + + index_db: str + source_db: str + embeddings_db: str + model: str + counts: EmbeddingRescueCounts + rescued_sessions: int + rescued_messages: int + skipped_already_fresh_sessions: int + skipped_race_sessions: int + verified_sample_total: int + verified_sample_ok: int + more_pending: bool + samples: tuple[EmbeddingRescueSample, ...] + + @property + def ok(self) -> bool: + """Whether every sampled rescued vector verified byte-identical to its source.""" + return self.verified_sample_total == 0 or self.verified_sample_ok == self.verified_sample_total + + def to_dict(self) -> dict[str, object]: + return { + "mode": "execute", + "ok": self.ok, + "index_db": self.index_db, + "source_db": self.source_db, + "embeddings_db": self.embeddings_db, + "model": self.model, + "eligible_sessions": self.counts.eligible_sessions, + "fully_rescuable_sessions": self.counts.fully_rescuable_sessions, + "rescuable_messages": self.counts.rescuable_messages, + "partial_sessions": self.counts.partial_sessions, + "partial_matched_messages": self.counts.partial_matched_messages, + "skipped_missing": self.counts.skipped_missing, + "skipped_hash_mismatch": self.counts.skipped_hash_mismatch, + "skipped_model_mismatch": self.counts.skipped_model_mismatch, + "rescued_sessions": self.rescued_sessions, + "rescued_messages": self.rescued_messages, + "skipped_already_fresh_sessions": self.skipped_already_fresh_sessions, + "skipped_race_sessions": self.skipped_race_sessions, + "verified_sample_total": self.verified_sample_total, + "verified_sample_ok": self.verified_sample_ok, + "more_pending": self.more_pending, + "samples": [sample.to_dict() for sample in self.samples], + } + + +def default_embedding_recipe() -> EmbeddingRecipe: + from polylogue.config import load_polylogue_config + from polylogue.storage.sqlite.archive_tiers.embeddings import EMBEDDING_DIMENSION + + cfg = load_polylogue_config() + return EmbeddingRecipe.current(model=str(cfg.embedding_model), dimensions=EMBEDDING_DIMENSION) + + +def _open_classification_connection(index_db_path: Path, source_db_path: Path) -> sqlite3.Connection: + if not source_db_path.exists(): + raise RuntimeError(f"retired embeddings source not found: {source_db_path}") + conn = open_readonly_connection(index_db_path, timeout=30.0) + try: + loaded, error = try_load_sqlite_vec(conn) + if not loaded: + raise RuntimeError("embedding rescue requires sqlite-vec") from error + conn.execute("ATTACH DATABASE ? AS retired", (f"file:{source_db_path}?mode=ro",)) + except BaseException: + conn.close() + raise + return conn + + +def _session_rollup_counts(conn: sqlite3.Connection, relation: str, model: str) -> EmbeddingRescueCounts: + session_rows = conn.execute( + f""" + {_classified_cte(relation)}, + session_rollup AS ( + SELECT session_id, COUNT(*) AS eligible_count, SUM(status = '{_MATCHED}') AS matched_count + FROM classified + GROUP BY session_id + ) + SELECT + COUNT(*) AS eligible_sessions, + COALESCE(SUM(CASE WHEN matched_count = eligible_count THEN 1 ELSE 0 END), 0) AS fully_rescuable_sessions, + COALESCE(SUM(CASE WHEN matched_count = eligible_count THEN matched_count ELSE 0 END), 0) AS rescuable_messages, + COALESCE( + SUM(CASE WHEN matched_count > 0 AND matched_count < eligible_count THEN 1 ELSE 0 END), 0 + ) AS partial_sessions, + COALESCE( + SUM(CASE WHEN matched_count > 0 AND matched_count < eligible_count THEN matched_count ELSE 0 END), 0 + ) AS partial_matched_messages + FROM session_rollup + """, + {"model": model}, + ).fetchone() + + status_rows = conn.execute( + f""" + {_classified_cte(relation)} + SELECT status, COUNT(*) FROM classified WHERE status != '{_MATCHED}' GROUP BY status + """, + {"model": model}, + ).fetchall() + status_counts = {str(row[0]): int(row[1]) for row in status_rows} + + return EmbeddingRescueCounts( + eligible_sessions=int(session_rows[0] or 0), + fully_rescuable_sessions=int(session_rows[1] or 0), + rescuable_messages=int(session_rows[2] or 0), + partial_sessions=int(session_rows[3] or 0), + partial_matched_messages=int(session_rows[4] or 0), + skipped_missing=status_counts.get(_MISSING, 0), + skipped_hash_mismatch=status_counts.get(_HASH_MISMATCH, 0), + skipped_model_mismatch=status_counts.get(_MODEL_MISMATCH, 0), + ) + + +def _classification_samples( + conn: sqlite3.Connection, relation: str, model: str, *, sample_size: int +) -> tuple[EmbeddingRescueSample, ...]: + if sample_size <= 0: + return () + rows = conn.execute( + f""" + {_classified_cte(relation)} + SELECT status, message_id, session_id + FROM classified + WHERE status != '{_MATCHED}' + ORDER BY message_id + LIMIT :limit + """, + {"model": model, "limit": sample_size}, + ).fetchall() + return tuple( + EmbeddingRescueSample(status=str(row[0]), message_id=str(row[1]), session_id=str(row[2])) for row in rows + ) + + +def _fully_rescuable_session_ids(conn: sqlite3.Connection, relation: str, model: str) -> tuple[str, ...]: + rows = conn.execute( + f""" + {_classified_cte(relation)}, + session_rollup AS ( + SELECT session_id, COUNT(*) AS eligible_count, SUM(status = '{_MATCHED}') AS matched_count + FROM classified + GROUP BY session_id + ) + SELECT session_id FROM session_rollup WHERE matched_count = eligible_count AND eligible_count > 0 + """, + {"model": model}, + ).fetchall() + return tuple(str(row[0]) for row in rows) + + +def plan_embedding_rescue( + index_db_path: str | Path, + source_embeddings_db_path: str | Path, + *, + sample_size: int = DEFAULT_SAMPLE_SIZE, + recipe: EmbeddingRecipe | None = None, +) -> EmbeddingRescuePlanReport: + """Read-only census of vectors rescuable from a retired embeddings tier. + + Never mutates either database; safe to run against a live archive + (including mid-rebuild) and against the read-only retired evidence file. + """ + index_path = Path(index_db_path) + source_path = Path(source_embeddings_db_path) + resolved_recipe = recipe or default_embedding_recipe() + + conn = _open_classification_connection(index_path, source_path) + try: + relation = archive_embeddable_messages_relation(conn, alias="e") + counts = _session_rollup_counts(conn, relation, resolved_recipe.model) + samples = _classification_samples(conn, relation, resolved_recipe.model, sample_size=sample_size) + finally: + conn.close() + + return EmbeddingRescuePlanReport( + index_db=str(index_path), + source_db=str(source_path), + model=resolved_recipe.model, + counts=counts, + samples=samples, + ) + + +def _session_eligible_messages(conn: sqlite3.Connection, relation: str, session_id: str) -> list[tuple[str, bytes]]: + rows = conn.execute( + f"SELECT e.message_id, e.content_hash FROM {relation} WHERE e.session_id = ?", + (session_id,), + ).fetchall() + return [(str(row[0]), bytes(row[1])) for row in rows] + + +def _source_hash_for(messages: list[tuple[str, bytes]]) -> bytes: + digest = EmbeddingSourceDigest() + for message_id, content_hash in sorted(messages): + digest.update(message_id, content_hash) + return digest.digest() + + +def _read_retired_vector(conn: sqlite3.Connection, message_id: str) -> list[float] | None: + row = conn.execute( + "SELECT embedding FROM retired.message_embeddings WHERE message_id = ?", + (message_id,), + ).fetchone() + if row is None or row[0] is None: + return None + raw = bytes(row[0]) + count = len(raw) // 4 + return list(struct.unpack(f"<{count}f", raw)) + + +def execute_embedding_rescue( + index_db_path: str | Path, + source_embeddings_db_path: str | Path, + embeddings_db_path: str | Path | None = None, + *, + limit: int | None = None, + sample_size: int = DEFAULT_SAMPLE_SIZE, + sample_verify_count: int = DEFAULT_SAMPLE_VERIFY_COUNT, + recipe: EmbeddingRecipe | None = None, + mutation_authority: EmbeddingRescueMutationAuthority | None = None, + now_ms: int | None = None, +) -> EmbeddingRescueExecuteReport: + """Copy vectors for fully-rescuable sessions from the retired tier. + + A session is only mutated when *every* one of its currently-eligible + messages has an exact retired match (see module docstring for why + partial sessions are never worth writing). Publication goes through + :func:`begin_embedding_attempt` / :func:`complete_embedding_attempt_success` + -- the same generation-guarded primitives the live embed path uses -- so + a rescued session is indistinguishable from one embedded for real, and a + second invocation is a no-op over already-rescued sessions (idempotent, + resumable via ``limit``). + """ + if mutation_authority is None: + raise RuntimeError("embedding rescue apply requires offline-exclusive authority") + + index_path = Path(index_db_path) + source_path = Path(source_embeddings_db_path) + embeddings_path = ( + Path(embeddings_db_path) if embeddings_db_path is not None else index_path.with_name("embeddings.db") + ) + resolved_recipe = recipe or default_embedding_recipe() + resolved_now_ms = now_ms if now_ms is not None else int(time.time() * 1000) + + from polylogue.storage.sqlite.archive_tiers.bootstrap import initialize_archive_database + from polylogue.storage.sqlite.archive_tiers.embedding_write import ( + ArchiveEmbeddingWrite, + begin_embedding_attempt, + complete_embedding_attempt_success, + supersede_embedding_attempt, + ) + from polylogue.storage.sqlite.archive_tiers.types import ArchiveTier + + initialize_archive_database(embeddings_path, ArchiveTier.EMBEDDINGS) + + index_conn = _open_classification_connection(index_path, source_path) + embeddings_conn = sqlite3.connect(embeddings_path, timeout=30.0) + rescued_sessions = 0 + rescued_messages = 0 + race_skipped = 0 + already_fresh_count = 0 + more_pending = False + verified_total = 0 + verified_ok = 0 + try: + loaded, error = try_load_sqlite_vec(embeddings_conn) + if not loaded: + raise RuntimeError("embedding rescue requires sqlite-vec") from error + + relation = archive_embeddable_messages_relation(index_conn, alias="e") + counts = _session_rollup_counts(index_conn, relation, resolved_recipe.model) + samples = _classification_samples(index_conn, relation, resolved_recipe.model, sample_size=sample_size) + fully_rescuable_ids = _fully_rescuable_session_ids(index_conn, relation, resolved_recipe.model) + + index_conn.execute("ATTACH DATABASE ? AS embeddings", (str(embeddings_path),)) + to_rescue = select_pending_archive_session_window( + index_conn, + status_table="embeddings.embedding_status", + session_ids=fully_rescuable_ids, + max_sessions=limit, + recipe=resolved_recipe, + ) + if fully_rescuable_ids: + total_pending = select_pending_archive_session_window( + index_conn, + status_table="embeddings.embedding_status", + session_ids=fully_rescuable_ids, + recipe=resolved_recipe, + ) + already_fresh_count = len(fully_rescuable_ids) - len(total_pending) + more_pending = len(total_pending) > len(to_rescue) + + origin_by_session: dict[str, str] = {} + if to_rescue: + placeholders = ", ".join("?" for _ in to_rescue) + origin_by_session = { + str(row[0]): str(row[1]) + for row in index_conn.execute( + f"SELECT session_id, origin FROM sessions WHERE session_id IN ({placeholders})", + tuple(item.session_id for item in to_rescue), + ).fetchall() + } + + rescued_message_ids: list[str] = [] + + for pending in to_rescue: + session_id = pending.session_id + origin = origin_by_session.get(session_id) + if origin is None: + continue + messages = _session_eligible_messages(index_conn, relation, session_id) + if not messages: + continue + source_hash = _source_hash_for(messages) + + attempt = begin_embedding_attempt( + embeddings_conn, + session_id=session_id, + origin=origin, + source_hash=source_hash, + recipe=resolved_recipe, + ) + + writes: list[ArchiveEmbeddingWrite] = [] + incomplete = False + for message_id, content_hash in messages: + vector = _read_retired_vector(index_conn, message_id) + if vector is None: + incomplete = True + break + writes.append( + ArchiveEmbeddingWrite( + message_id=message_id, + session_id=session_id, + origin=origin, + embedding=vector, + model=resolved_recipe.model, + embedded_at_ms=resolved_now_ms, + content_hash=content_hash, + recipe_hash=attempt.recipe_hash, + derivation_key=message_embedding_derivation_key( + message_id=message_id, + content_hash=content_hash, + recipe=resolved_recipe, + ).digest(), + generation=attempt.generation, + ) + ) + + if incomplete: + # Classified as fully rescuable but the retired vector vanished + # between classification and read (corrupt/edited retired + # evidence). Never publish a partial write; requeue for the + # real embed path instead of leaving a half-written attempt. + supersede_embedding_attempt( + embeddings_conn, + attempt=attempt, + source_hash=source_hash, + recipe=resolved_recipe, + ) + race_skipped += 1 + continue + + current_messages = _session_eligible_messages(index_conn, relation, session_id) + current_source_hash = _source_hash_for(current_messages) + if current_source_hash != source_hash or len(current_messages) != len(messages): + supersede_embedding_attempt( + embeddings_conn, + attempt=attempt, + source_hash=current_source_hash, + recipe=resolved_recipe, + ) + race_skipped += 1 + continue + + committed = complete_embedding_attempt_success( + embeddings_conn, + attempt=attempt, + writes=writes, + completed_at_ms=resolved_now_ms, + ) + if not committed: + race_skipped += 1 + continue + + rescued_sessions += 1 + rescued_messages += len(writes) + rescued_message_ids.extend(message_id for message_id, _ in messages) + + if rescued_message_ids and sample_verify_count > 0: + step = max(1, len(rescued_message_ids) // sample_verify_count) + sample_ids = rescued_message_ids[::step][:sample_verify_count] + for message_id in sample_ids: + verified_total += 1 + target_row = embeddings_conn.execute( + "SELECT embedding FROM message_embeddings WHERE message_id = ?", + (message_id,), + ).fetchone() + source_row = index_conn.execute( + "SELECT embedding FROM retired.message_embeddings WHERE message_id = ?", + (message_id,), + ).fetchone() + if target_row is not None and source_row is not None and bytes(target_row[0]) == bytes(source_row[0]): + verified_ok += 1 + finally: + index_conn.close() + embeddings_conn.close() + + return EmbeddingRescueExecuteReport( + index_db=str(index_path), + source_db=str(source_path), + embeddings_db=str(embeddings_path), + model=resolved_recipe.model, + counts=counts, + rescued_sessions=rescued_sessions, + rescued_messages=rescued_messages, + skipped_already_fresh_sessions=already_fresh_count, + skipped_race_sessions=race_skipped, + verified_sample_total=verified_total, + verified_sample_ok=verified_ok, + more_pending=more_pending, + samples=samples, + ) + + +__all__ = [ + "DEFAULT_SAMPLE_SIZE", + "DEFAULT_SAMPLE_VERIFY_COUNT", + "EmbeddingRescueCounts", + "EmbeddingRescueExecuteReport", + "EmbeddingRescueMutationAuthority", + "EmbeddingRescuePlanReport", + "EmbeddingRescueSample", + "default_embedding_recipe", + "execute_embedding_rescue", + "plan_embedding_rescue", +] diff --git a/tests/unit/cli/test_embeddings_rescue_cli.py b/tests/unit/cli/test_embeddings_rescue_cli.py new file mode 100644 index 0000000000..d61f02f776 --- /dev/null +++ b/tests/unit/cli/test_embeddings_rescue_cli.py @@ -0,0 +1,120 @@ +"""CLI wiring smoke tests for ``ops maintenance embeddings-rescue`` (polylogue-04kl). + +Deep classification/rescue correctness is covered at the storage layer in +``tests/unit/storage/test_embedding_rescue.py``; these tests only prove the +CLI is registered, dispatches plan vs. apply correctly, and honors the +offline daemon guard -- using the real full-schema empty archive template so +argument/path resolution is exercised against production shapes. +""" + +from __future__ import annotations + +import json +import sqlite3 +from pathlib import Path +from unittest.mock import patch + +import pytest +from click.testing import CliRunner + +from polylogue.cli.click_app import cli +from polylogue.storage.sqlite.archive_tiers.bootstrap import initialize_archive_tier +from polylogue.storage.sqlite.archive_tiers.types import ArchiveTier + + +def _build_empty_retired_source(path: Path) -> None: + conn = sqlite3.connect(path) + try: + initialize_archive_tier(conn, ArchiveTier.EMBEDDINGS) + except sqlite3.OperationalError as exc: + if "vec0" in str(exc) or "sqlite-vec" in str(exc): + pytest.skip("sqlite-vec extension is unavailable") + raise + finally: + conn.close() + + +def test_embeddings_rescue_cli_plan_mode_is_read_only_on_empty_archive( + cli_workspace: dict[str, Path], + cli_runner: CliRunner, + tmp_path: Path, +) -> None: + source = tmp_path / "retired-embeddings.db" + _build_empty_retired_source(source) + + result = cli_runner.invoke( + cli, + ["--plain", "ops", "maintenance", "embeddings-rescue", "--source", str(source), "--output-format", "json"], + catch_exceptions=False, + ) + + assert result.exit_code == 0 + payload = json.loads(result.output) + assert payload["mode"] == "plan" + assert payload["mutates"] is False + assert payload["eligible_sessions"] == 0 + assert payload["fully_rescuable_sessions"] == 0 + assert payload["rescuable_messages"] == 0 + + +def test_embeddings_rescue_cli_apply_mode_no_op_on_empty_archive( + cli_workspace: dict[str, Path], + cli_runner: CliRunner, + tmp_path: Path, +) -> None: + source = tmp_path / "retired-embeddings.db" + _build_empty_retired_source(source) + + result = cli_runner.invoke( + cli, + [ + "--plain", + "ops", + "maintenance", + "embeddings-rescue", + "--source", + str(source), + "--yes", + "--output-format", + "json", + ], + catch_exceptions=False, + ) + + assert result.exit_code == 0 + payload = json.loads(result.output) + assert payload["mode"] == "execute" + assert payload["mutates"] is True + assert payload["rescued_sessions"] == 0 + assert payload["ok"] is True + + +def test_embeddings_rescue_cli_apply_refuses_while_daemon_runs( + cli_workspace: dict[str, Path], + cli_runner: CliRunner, + tmp_path: Path, +) -> None: + source = tmp_path / "retired-embeddings.db" + _build_empty_retired_source(source) + + with patch("polylogue.maintenance.offline_guard.running_daemon_pid", return_value=123): + result = cli_runner.invoke( + cli, + ["--plain", "ops", "maintenance", "embeddings-rescue", "--source", str(source), "--yes"], + ) + + assert result.exit_code == 1 + assert "refused while polylogued PID 123 is running" in result.output + + +def test_embeddings_rescue_cli_requires_existing_source( + cli_workspace: dict[str, Path], + cli_runner: CliRunner, + tmp_path: Path, +) -> None: + missing = tmp_path / "does-not-exist.db" + result = cli_runner.invoke( + cli, + ["--plain", "ops", "maintenance", "embeddings-rescue", "--source", str(missing)], + ) + assert result.exit_code != 0 diff --git a/tests/unit/storage/test_embedding_rescue.py b/tests/unit/storage/test_embedding_rescue.py new file mode 100644 index 0000000000..3837a19e3a --- /dev/null +++ b/tests/unit/storage/test_embedding_rescue.py @@ -0,0 +1,462 @@ +"""Focused tests for polylogue-04kl: retired-tier embedding vector rescue. + +Builds a synthetic "post-incident" scenario directly: a minimal ``index.db`` +fixture (mirroring ``test_embedding_orphan_reconcile.py``'s pattern) plus a +retired ``embeddings.db``-shaped source file and a fresh target +``embeddings.db``, then asserts :func:`plan_embedding_rescue` classifies +sessions/messages correctly and :func:`execute_embedding_rescue` only ever +mutates *fully rescuable* sessions, is idempotent on rerun, and its sampled +byte-identity verification actually catches a corrupted copy. +""" + +from __future__ import annotations + +import sqlite3 +import struct +from pathlib import Path +from typing import NamedTuple + +import pytest + +from polylogue.core.enums import Origin +from polylogue.storage.embeddings.identity import EmbeddingRecipe +from polylogue.storage.embeddings.rescue import ( + execute_embedding_rescue, + plan_embedding_rescue, +) +from polylogue.storage.sqlite.archive_tiers.bootstrap import initialize_archive_tier +from polylogue.storage.sqlite.archive_tiers.embedding_write import ( + ArchiveEmbeddingWrite, + upsert_message_embeddings, +) +from polylogue.storage.sqlite.archive_tiers.embeddings import EMBEDDING_DIMENSION +from polylogue.storage.sqlite.archive_tiers.index import INDEX_SCHEMA_VERSION +from polylogue.storage.sqlite.archive_tiers.types import ArchiveTier +from polylogue.storage.sqlite.sqlite_vec_extension import try_load_sqlite_vec + +_MODEL = "voyage-4" +_RECIPE = EmbeddingRecipe.current(model=_MODEL, dimensions=EMBEDDING_DIMENSION) + +_INDEX_DDL = """ +CREATE TABLE sessions ( + session_id TEXT PRIMARY KEY, + origin TEXT NOT NULL, + title TEXT, + sort_key_ms INTEGER DEFAULT 0, + authored_user_message_count INTEGER NOT NULL DEFAULT 1, + assistant_message_count INTEGER NOT NULL DEFAULT 0 +); +CREATE TABLE messages ( + message_id TEXT PRIMARY KEY, + session_id TEXT NOT NULL, + text TEXT NOT NULL DEFAULT 'authored prose long enough for embedding', + role TEXT NOT NULL DEFAULT 'user', + message_type TEXT NOT NULL DEFAULT 'message', + material_origin TEXT NOT NULL DEFAULT 'human_authored', + word_count INTEGER NOT NULL DEFAULT 8, + content_hash BLOB NOT NULL +); +""" + + +def _content_hash(seed: str) -> bytes: + return (seed.encode("utf-8") * 32)[:32] + + +def _connect_index(path: Path, *, sessions: list[str], messages: dict[str, list[tuple[str, bytes]]]) -> None: + """Build a minimal index.db: sessions + messages with explicit content hashes.""" + conn = sqlite3.connect(path) + try: + conn.executescript(_INDEX_DDL) + conn.execute(f"PRAGMA user_version = {INDEX_SCHEMA_VERSION}") + for session_id in sessions: + conn.execute( + "INSERT INTO sessions (session_id, origin) VALUES (?, ?)", + (session_id, "codex-session"), + ) + for session_id, entries in messages.items(): + for message_id, content_hash in entries: + conn.execute( + "INSERT INTO messages (message_id, session_id, content_hash) VALUES (?, ?, ?)", + (message_id, session_id, content_hash), + ) + conn.commit() + finally: + conn.close() + + +def _connect_embeddings_tier(path: Path) -> sqlite3.Connection: + conn = sqlite3.connect(path) + conn.row_factory = sqlite3.Row + try: + initialize_archive_tier(conn, ArchiveTier.EMBEDDINGS) + except sqlite3.OperationalError as exc: + if "vec0" in str(exc) or "sqlite-vec" in str(exc): + pytest.skip("sqlite-vec extension is unavailable") + raise + return conn + + +def _vector_for(seed: int) -> list[float]: + return [float(seed) + i * 0.001 for i in range(EMBEDDING_DIMENSION)] + + +def _write_retired( + conn: sqlite3.Connection, + *, + message_id: str, + session_id: str, + content_hash: bytes, + model: str = _MODEL, + seed: int = 1, +) -> None: + upsert_message_embeddings( + conn, + [ + ArchiveEmbeddingWrite( + message_id=message_id, + session_id=session_id, + origin=Origin.CODEX_SESSION, + embedding=_vector_for(seed), + model=model, + embedded_at_ms=1_700_000_000_000, + content_hash=content_hash, + ) + ], + ) + + +def _target_embeddings_path(tmp_path: Path) -> Path: + return tmp_path / "embeddings.db" + + +def _connect_target_for_read(path: Path) -> sqlite3.Connection: + """Plain connection to the target tier with sqlite-vec loaded for reads.""" + conn = sqlite3.connect(path) + loaded, error = try_load_sqlite_vec(conn) + if not loaded: + conn.close() + pytest.skip("sqlite-vec extension is unavailable") + raise AssertionError("unreachable") from error + return conn + + +class TestPlanEmbeddingRescue: + def test_classifies_matched_missing_hash_mismatch_model_mismatch(self, tmp_path: Path) -> None: + # Session "full": both messages have exact retired matches -> fully rescuable. + full = "codex-session:full" + full_m1, full_m2 = f"{full}:m1", f"{full}:m2" + hash_m1, hash_m2 = _content_hash("m1"), _content_hash("m2") + + # Session "partial": one matched, one missing from the retired tier. + partial = "codex-session:partial" + partial_m1, partial_m2 = f"{partial}:m1", f"{partial}:m2" + hash_p1, hash_p2 = _content_hash("p1"), _content_hash("p2") + + # Session "stale": one message whose retired content_hash no longer matches. + stale = "codex-session:stale" + stale_m1 = f"{stale}:m1" + hash_s1_now = _content_hash("s1-now") + hash_s1_retired = _content_hash("s1-retired") + + # Session "wrong-model": retired vector was embedded under a different model. + wrong_model = "codex-session:wrong-model" + wrong_model_m1 = f"{wrong_model}:m1" + hash_w1 = _content_hash("w1") + + _connect_index( + tmp_path / "index.db", + sessions=[full, partial, stale, wrong_model], + messages={ + full: [(full_m1, hash_m1), (full_m2, hash_m2)], + partial: [(partial_m1, hash_p1), (partial_m2, hash_p2)], + stale: [(stale_m1, hash_s1_now)], + wrong_model: [(wrong_model_m1, hash_w1)], + }, + ) + + retired_path = tmp_path / "retired-embeddings.db" + retired_conn = _connect_embeddings_tier(retired_path) + _write_retired(retired_conn, message_id=full_m1, session_id=full, content_hash=hash_m1, seed=1) + _write_retired(retired_conn, message_id=full_m2, session_id=full, content_hash=hash_m2, seed=2) + _write_retired(retired_conn, message_id=partial_m1, session_id=partial, content_hash=hash_p1, seed=3) + # partial_m2 intentionally absent from the retired tier. + _write_retired(retired_conn, message_id=stale_m1, session_id=stale, content_hash=hash_s1_retired, seed=4) + _write_retired( + retired_conn, + message_id=wrong_model_m1, + session_id=wrong_model, + content_hash=hash_w1, + model="voyage-3", + seed=5, + ) + retired_conn.close() + + report = plan_embedding_rescue(tmp_path / "index.db", retired_path, recipe=_RECIPE) + + assert report.counts.eligible_sessions == 4 + assert report.counts.fully_rescuable_sessions == 1 + assert report.counts.rescuable_messages == 2 + assert report.counts.partial_sessions == 1 + assert report.counts.partial_matched_messages == 1 + assert report.counts.skipped_missing == 1 + assert report.counts.skipped_hash_mismatch == 1 + assert report.counts.skipped_model_mismatch == 1 + assert report.model == _MODEL + + statuses = {sample.status for sample in report.samples} + assert statuses == {"missing", "hash_mismatch", "model_mismatch"} + + def test_plan_never_mutates_either_database(self, tmp_path: Path) -> None: + session_id = "codex-session:s1" + m1 = f"{session_id}:m1" + content_hash = _content_hash("m1") + _connect_index(tmp_path / "index.db", sessions=[session_id], messages={session_id: [(m1, content_hash)]}) + retired_path = tmp_path / "retired-embeddings.db" + retired_conn = _connect_embeddings_tier(retired_path) + _write_retired(retired_conn, message_id=m1, session_id=session_id, content_hash=content_hash, seed=1) + retired_conn.close() + + before_index = (tmp_path / "index.db").stat().st_mtime_ns + before_retired = retired_path.stat().st_mtime_ns + + report_a = plan_embedding_rescue(tmp_path / "index.db", retired_path, recipe=_RECIPE) + report_b = plan_embedding_rescue(tmp_path / "index.db", retired_path, recipe=_RECIPE) + + assert report_a.to_dict() == report_b.to_dict() + assert (tmp_path / "index.db").stat().st_mtime_ns == before_index + assert retired_path.stat().st_mtime_ns == before_retired + assert not (tmp_path / "embeddings.db").exists() + + +class _SeededFullAndPartial(NamedTuple): + full: str + full_messages: tuple[str, str] + partial: str + retired_path: Path + + +class TestExecuteEmbeddingRescue: + def _seed_full_and_partial(self, tmp_path: Path) -> _SeededFullAndPartial: + full = "codex-session:full" + full_m1, full_m2 = f"{full}:m1", f"{full}:m2" + hash_m1, hash_m2 = _content_hash("m1"), _content_hash("m2") + + partial = "codex-session:partial" + partial_m1, partial_m2 = f"{partial}:m1", f"{partial}:m2" + hash_p1, hash_p2 = _content_hash("p1"), _content_hash("p2") + + _connect_index( + tmp_path / "index.db", + sessions=[full, partial], + messages={ + full: [(full_m1, hash_m1), (full_m2, hash_m2)], + partial: [(partial_m1, hash_p1), (partial_m2, hash_p2)], + }, + ) + retired_path = tmp_path / "retired-embeddings.db" + retired_conn = _connect_embeddings_tier(retired_path) + _write_retired(retired_conn, message_id=full_m1, session_id=full, content_hash=hash_m1, seed=1) + _write_retired(retired_conn, message_id=full_m2, session_id=full, content_hash=hash_m2, seed=2) + _write_retired(retired_conn, message_id=partial_m1, session_id=partial, content_hash=hash_p1, seed=3) + retired_conn.close() + return _SeededFullAndPartial( + full=full, + full_messages=(full_m1, full_m2), + partial=partial, + retired_path=retired_path, + ) + + def test_rescues_only_the_fully_matched_session(self, tmp_path: Path) -> None: + seed = self._seed_full_and_partial(tmp_path) + target = _target_embeddings_path(tmp_path) + + report = execute_embedding_rescue( + tmp_path / "index.db", + seed.retired_path, + target, + mutation_authority="offline-exclusive", + now_ms=1_800_000_000_000, + ) + + assert report.rescued_sessions == 1 + assert report.rescued_messages == 2 + assert report.counts.fully_rescuable_sessions == 1 + assert report.counts.partial_sessions == 1 + assert report.verified_sample_total >= 1 + assert report.verified_sample_ok == report.verified_sample_total + assert report.ok is True + + conn = _connect_target_for_read(target) + try: + full_m1, full_m2 = seed.full_messages + for message_id in (full_m1, full_m2): + assert ( + conn.execute( + "SELECT COUNT(*) FROM message_embeddings WHERE message_id = ?", (message_id,) + ).fetchone()[0] + == 1 + ) + assert ( + conn.execute( + "SELECT COUNT(*) FROM message_embeddings_meta WHERE message_id = ?", (message_id,) + ).fetchone()[0] + == 1 + ) + # The partial session must never be written -- no vectors at all. + partial_row_count = conn.execute( + "SELECT COUNT(*) FROM message_embeddings WHERE session_id = ?", (seed.partial,) + ).fetchone()[0] + assert partial_row_count == 0 + + status = conn.execute( + "SELECT message_count_embedded, needs_reindex, error_message FROM embedding_status WHERE session_id = ?", + (seed.full,), + ).fetchone() + assert status == (2, 0, None) + + derivation = conn.execute( + "SELECT attempt_state, message_count FROM embedding_derivation_state WHERE session_id = ?", + (seed.full,), + ).fetchone() + assert derivation == ("succeeded", 2) + finally: + conn.close() + + def test_rescue_is_idempotent_on_rerun(self, tmp_path: Path) -> None: + seed = self._seed_full_and_partial(tmp_path) + target = _target_embeddings_path(tmp_path) + + first = execute_embedding_rescue( + tmp_path / "index.db", + seed.retired_path, + target, + mutation_authority="offline-exclusive", + now_ms=1_800_000_000_000, + ) + assert first.rescued_sessions == 1 + + second = execute_embedding_rescue( + tmp_path / "index.db", + seed.retired_path, + target, + mutation_authority="offline-exclusive", + now_ms=1_800_000_100_000, + ) + + assert second.rescued_sessions == 0 + assert second.skipped_already_fresh_sessions == 1 + assert second.more_pending is False + + conn = _connect_target_for_read(target) + try: + full_m1, _full_m2 = seed.full_messages + row = conn.execute( + "SELECT embedded_at_ms FROM message_embeddings_meta WHERE message_id = ?", (full_m1,) + ).fetchone() + # Second run must not have rewritten the already-rescued vector. + assert row[0] == 1_800_000_000_000 + finally: + conn.close() + + def test_limit_bounds_sessions_rescued_and_reports_more_pending(self, tmp_path: Path) -> None: + full_a, full_b = "codex-session:full-a", "codex-session:full-b" + m_a, m_b = f"{full_a}:m1", f"{full_b}:m1" + hash_a, hash_b = _content_hash("a"), _content_hash("b") + + _connect_index( + tmp_path / "index.db", + sessions=[full_a, full_b], + messages={full_a: [(m_a, hash_a)], full_b: [(m_b, hash_b)]}, + ) + retired_path = tmp_path / "retired-embeddings.db" + retired_conn = _connect_embeddings_tier(retired_path) + _write_retired(retired_conn, message_id=m_a, session_id=full_a, content_hash=hash_a, seed=1) + _write_retired(retired_conn, message_id=m_b, session_id=full_b, content_hash=hash_b, seed=2) + retired_conn.close() + + target = _target_embeddings_path(tmp_path) + report = execute_embedding_rescue( + tmp_path / "index.db", + retired_path, + target, + limit=1, + mutation_authority="offline-exclusive", + now_ms=1_800_000_000_000, + ) + + assert report.rescued_sessions == 1 + assert report.more_pending is True + + def test_apply_requires_explicit_mutation_authority(self, tmp_path: Path) -> None: + seed = self._seed_full_and_partial(tmp_path) + with pytest.raises(RuntimeError, match="offline-exclusive"): + execute_embedding_rescue(tmp_path / "index.db", seed.retired_path) + + def test_sample_verification_catches_a_corrupted_copy( + self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch + ) -> None: + """Anti-vacuity: a transit-corrupted vector must fail the sample check. + + Simulates a bug in the copy path (not a corrupt source) by + monkeypatching the retired-vector reader to hand back a mutated + vector for the single message written, then asserting the sample + verification -- which re-reads both the just-written target row and + the untouched retired source row -- reports the mismatch instead of + silently passing. + """ + session_id = "codex-session:solo" + message_id = f"{session_id}:m1" + content_hash = _content_hash("solo") + + _connect_index( + tmp_path / "index.db", + sessions=[session_id], + messages={session_id: [(message_id, content_hash)]}, + ) + retired_path = tmp_path / "retired-embeddings.db" + retired_conn = _connect_embeddings_tier(retired_path) + _write_retired(retired_conn, message_id=message_id, session_id=session_id, content_hash=content_hash, seed=7) + retired_conn.close() + + import polylogue.storage.embeddings.rescue as rescue_module + + real_reader = rescue_module._read_retired_vector + + def _corrupting_reader(conn: sqlite3.Connection, target_message_id: str) -> list[float] | None: + vector = real_reader(conn, target_message_id) + if vector is None: + return None + corrupted = list(vector) + corrupted[0] = corrupted[0] + 999.0 + return corrupted + + monkeypatch.setattr(rescue_module, "_read_retired_vector", _corrupting_reader) + + target = _target_embeddings_path(tmp_path) + report = execute_embedding_rescue( + tmp_path / "index.db", + retired_path, + target, + mutation_authority="offline-exclusive", + now_ms=1_800_000_000_000, + ) + + assert report.rescued_sessions == 1 + assert report.verified_sample_total >= 1 + assert report.verified_sample_ok < report.verified_sample_total + assert report.ok is False + + +def test_vector_roundtrip_is_byte_identical(tmp_path: Path) -> None: + """Guards the float32-blob roundtrip claim the rescue writer relies on. + + ``struct.unpack`` into Python floats (doubles) and re-``struct.pack`` + back to float32 must reproduce the exact original bytes -- this is what + lets the writer accept ``list[float]`` without special-casing raw bytes + while still promising byte-identical rescued vectors. + """ + original = struct.pack(f"<{EMBEDDING_DIMENSION}f", *(0.1 * i for i in range(EMBEDDING_DIMENSION))) + floats = list(struct.unpack(f"<{EMBEDDING_DIMENSION}f", original)) + roundtripped = struct.pack(f"<{EMBEDDING_DIMENSION}f", *floats) + assert roundtripped == original