arborist/tests/test_evict.py

"""Eviction & rehydrate: lossless reversible forgetting."""

from __future__ import annotations

import json
from typing import Iterator

from arborist.distill import FirstSentenceDistiller
from arborist.distill.runner import distill_existing
from arborist.document import Document
from arborist.evict import evict_to_cold, rehydrate
from arborist.ingest import ingest_source
from arborist.source import Source
from arborist.store import connect


class FakeSource(Source):
    source_type = "html"  # so rehydrate path treats these as html-like

    def __init__(self, docs: list[Document]):
        self.docs = docs

    def iter_documents(self) -> Iterator[Document]:
        yield from self.docs


def _doc(uri: str, content: str) -> Document:
    return Document(uri=uri, content=content, source_type="html", title=uri)


LONG = (
    "The eight forms of capital include living, social, and intellectual. " * 30
    + "\n\n"
    + "Merkle providence proves answer derives from a specific source. " * 30
)


def test_evict_marks_chunks_cold_and_clears_fts(tmp_path):
    db = tmp_path / "evict.db"
    conn = connect(db)
    try:
        ingest_source(conn, FakeSource([_doc("html://a", LONG)]))
        # Before: hot chunks, FTS rows present
        before = conn.execute(
            "SELECT COUNT(*) FROM chunks WHERE tier='hot'"
        ).fetchone()[0]
        before_fts = conn.execute("SELECT COUNT(*) FROM chunks_fts").fetchone()[0]
        assert before > 0
        assert before_fts == before

        result = evict_to_cold(conn)
        assert result["evicted_chunks"] == before
        assert result["documents_affected"] == 1

        cold = conn.execute(
            "SELECT COUNT(*) FROM chunks WHERE tier='cold'"
        ).fetchone()[0]
        assert cold == before
        # Content NULLed
        nulls = conn.execute(
            "SELECT COUNT(*) FROM chunks WHERE content IS NULL"
        ).fetchone()[0]
        assert nulls == before
        # leaf_hash retained
        no_hash = conn.execute(
            "SELECT COUNT(*) FROM chunks WHERE leaf_hash IS NULL OR leaf_hash = ''"
        ).fetchone()[0]
        assert no_hash == 0
        # FTS rows removed
        fts_after = conn.execute("SELECT COUNT(*) FROM chunks_fts").fetchone()[0]
        assert fts_after == 0
    finally:
        conn.close()


def test_evict_skips_cores(tmp_path):
    db = tmp_path / "cores.db"
    conn = connect(db)
    try:
        ingest_source(conn, FakeSource([_doc("html://x", LONG)]))
        distill_existing(conn, FirstSentenceDistiller())
        # Verify there's a core
        n_cores = conn.execute(
            "SELECT COUNT(*) FROM documents WHERE kind='core'"
        ).fetchone()[0]
        assert n_cores == 1

        evict_to_cold(conn)

        # All surface chunks are cold; core chunks still hot.
        surface_tiers = conn.execute(
            "SELECT DISTINCT c.tier FROM chunks c "
            "JOIN documents d ON d.document_root=c.document_root "
            "WHERE d.kind='surface'"
        ).fetchall()
        assert {r["tier"] for r in surface_tiers} == {"cold"}

        core_tiers = conn.execute(
            "SELECT DISTINCT c.tier FROM chunks c "
            "JOIN documents d ON d.document_root=c.document_root "
            "WHERE d.kind='core'"
        ).fetchall()
        assert {r["tier"] for r in core_tiers} == {"hot"}
    finally:
        conn.close()


def test_rehydrate_restores_content_when_uri_matches(tmp_path):
    """Mock fetcher returns the same canonicalized text — root matches → restore."""
    db = tmp_path / "rh.db"
    conn = connect(db)
    try:
        ingest_source(conn, FakeSource([_doc("html://stable", LONG)]))
        evict_to_cold(conn)
        nulls_before = conn.execute(
            "SELECT COUNT(*) FROM chunks WHERE content IS NULL"
        ).fetchone()[0]
        assert nulls_before > 0

        # Inject a fetcher that returns the original text verbatim.
        def fake_fetch(uri: str) -> str:
            return LONG

        root = conn.execute(
            "SELECT document_root FROM documents WHERE document_uri='html://stable'"
        ).fetchone()["document_root"]

        result = rehydrate(conn, root, fetcher=fake_fetch)
        assert result["status"] == "rehydrated"
        assert result["chunks_restored"] > 0

        nulls_after = conn.execute(
            "SELECT COUNT(*) FROM chunks WHERE content IS NULL"
        ).fetchone()[0]
        assert nulls_after == 0
        # All restored chunks tier=hot, FTS repopulated.
        cold = conn.execute(
            "SELECT COUNT(*) FROM chunks WHERE tier='cold'"
        ).fetchone()[0]
        assert cold == 0
    finally:
        conn.close()


def test_rehydrate_drift_marks_providence_stale(tmp_path):
    """If URI's content has changed, leaves don't match → drift event +
    every providence_cache row for this source flips to 'stale'."""
    import time

    db = tmp_path / "drift.db"
    conn = connect(db)
    try:
        ingest_source(conn, FakeSource([_doc("html://drifty", LONG)]))
        root = conn.execute(
            "SELECT document_root FROM documents WHERE document_uri='html://drifty'"
        ).fetchone()["document_root"]

        # Insert a fake live providence record bound to this source.
        conn.execute(
            "INSERT INTO providence_cache "
            "(cache_key, source_root, document_uri, question_hash, question_text, "
            " answer_text, merkle_proof, model_profile_hash, conversation_hash, "
            " governance_policy_hash, schema_version, canonicalization_version, "
            " chunking_version, falsification_state, created_at) "
            "VALUES (?, ?, 'html://drifty', 'q1', 'qtxt', 'atxt', '{}', 'm1', 'c1', "
            "'g1', 'v9.8.0', 'norm-v1', 'tok-512-v1', 'live', ?)",
            (root + ":q1", root, int(time.time())),
        )

        evict_to_cold(conn)

        # Fetcher returns DIFFERENT text — drift.
        def drifted_fetch(uri: str) -> str:
            return LONG + "\n\nNEW PARAGRAPH ADDED AFTER INGEST."

        result = rehydrate(conn, root, fetcher=drifted_fetch)
        assert result["status"] == "drift_detected"
        assert result["expected_root"] == root
        assert result["actual_root"] != root

        # No chunk content was restored.
        cold = conn.execute(
            "SELECT COUNT(*) FROM chunks WHERE tier='cold'"
        ).fetchone()[0]
        assert cold > 0

        # Providence record flipped to stale.
        state = conn.execute(
            "SELECT falsification_state FROM providence_cache WHERE source_root=?",
            (root,),
        ).fetchone()["falsification_state"]
        assert state == "stale"

        # Drift event recorded in audit chain.
        last = conn.execute(
            "SELECT event_type FROM audit_events ORDER BY seq DESC LIMIT 1"
        ).fetchone()
        assert last["event_type"] == "rehydrate_drift"
    finally:
        conn.close()


def test_rehydrate_unknown_document(tmp_path):
    db = tmp_path / "u.db"
    conn = connect(db)
    try:
        result = rehydrate(conn, "00" * 32)
        assert result["status"] == "unknown_document"
    finally:
        conn.close()


def test_rehydrate_nothing_to_do_when_all_hot(tmp_path):
    db = tmp_path / "n.db"
    conn = connect(db)
    try:
        ingest_source(conn, FakeSource([_doc("html://hot", LONG)]))
        root = conn.execute(
            "SELECT document_root FROM documents WHERE document_uri='html://hot'"
        ).fetchone()["document_root"]
        result = rehydrate(conn, root, fetcher=lambda u: LONG)
        assert result["status"] == "nothing_to_do"
    finally:
        conn.close()


def test_rehydrate_non_rehydratable_source(tmp_path):
    """A source_type without a registered fetcher (and no override) is honest about it."""
    from arborist.source import Source

    class WikiSource(Source):
        source_type = "wikipedia_cur"

        def iter_documents(self):
            yield Document(
                uri="https://en.wikipedia.org/wiki/Foo",
                content=LONG,
                source_type="wikipedia_cur",
                title="Foo",
            )

    db = tmp_path / "ns.db"
    conn = connect(db)
    try:
        ingest_source(conn, WikiSource())
        root = conn.execute(
            "SELECT document_root FROM documents LIMIT 1"
        ).fetchone()["document_root"]
        evict_to_cold(conn)
        result = rehydrate(conn, root)  # no fetcher override
        assert result["status"] == "source_not_rehydratable"
        assert result["source_type"] == "wikipedia_cur"
    finally:
        conn.close()