ratatoskr/tests/test_provider_memory.py

"""Tests for the Tier-3 Bifrost memory provider (ratatoskr.provider.memory_store).

Contract: docs/contracts/bifrost_memory_provider.contract.md (v1.1)
Vertical tracer-first: fresh_db -> basic_upsert (round-trip) -> replay -> conflict
-> optimistic_lock -> injection_rule -> search/scope_isolation -> get/get_many ->
delete_many -> build_memory_provider_app -> #195 parity vs InMemoryMemoryStore.
"""
from __future__ import annotations

import types

import pytest
from bifrost.memory import (
    IdempotencyConflict,
    InvalidArguments,
    InvalidFilter,
    RevisionMismatch,
)

from ratatoskr.provider.memory_store import (
    build_memory_provider_app,
    open_memory_store,
)

EMBEDDING_DIM = 8


def _ctx(sub: str = "sub-1"):
    # Mirrors bifrost reference _ctx_actor: actor = job_id | jwt_sub | session_id.
    return types.SimpleNamespace(jwt_sub=sub)


def _vec(*head: float) -> list[float]:
    v = list(head) + [0.0] * EMBEDDING_DIM
    return v[:EMBEDDING_DIM]


def _chunk(cid: str = "c1", *, embedding=None, scope=None, **extra) -> dict:
    rec = {
        "id": cid,
        "embedding": embedding if embedding is not None else _vec(1.0),
        "scope": scope if scope is not None else {"end_user": "u1"},
        "origin": "worldtree",
        "distillate": {"summary": f"distillate-{cid}"},
        "content": f"content-{cid}",
    }
    rec.update(extra)
    return rec


def _row_count(store, table: str) -> int:
    return store._conn.execute(f"SELECT COUNT(*) FROM {table}").fetchone()[0]


# --- open_memory_store ---

def test_fresh_db_advertises_v1_caps_and_schema():
    store = open_memory_store(":memory:", embedding_dim=EMBEDDING_DIM)
    caps = store.describe_store()
    assert caps["relational_edges_supported"] is False
    assert caps["optimistic_locking_supported"] is True
    assert caps["atomic_supersede_supported"] is False
    assert caps["transaction_supported"] is False
    assert caps["filterable_metadata_fields"] == []
    # tables + vec index queryable
    store._conn.execute("SELECT * FROM memory_chunks")
    store._conn.execute("SELECT * FROM memory_idempotency")
    store._conn.execute("SELECT * FROM memory_vec")


def test_reopen_existing_file_is_idempotent(tmp_path):
    db = str(tmp_path / "memory.db")
    open_memory_store(db, embedding_dim=EMBEDDING_DIM)  # first open creates schema
    store = open_memory_store(db, embedding_dim=EMBEDDING_DIM)  # reopen: IF NOT EXISTS no-op
    assert isinstance(store.describe_store(), dict)
    store._conn.execute("SELECT * FROM memory_chunks")
    store._conn.execute("SELECT * FROM memory_vec")


# --- upsert_many + get (tracer round-trip) ---

async def test_basic_upsert_round_trips_verbatim_with_revision():
    store = open_memory_store(":memory:", embedding_dim=EMBEDDING_DIM)
    c1 = _chunk("c1", embedding=_vec(1.0))
    c2 = _chunk("c2", embedding=_vec(0.0, 1.0))
    result = await store.upsert_many([c1, c2], idempotency_key="k1", ctx=_ctx())
    assert result == {"upserted": 2, "replayed": False}
    # INV-001: each chunk round-trips verbatim, with a revision key attached (first insert -> 1)
    assert await store.get("c1") == {**c1, "revision": 1}
    assert await store.get("c2") == {**c2, "revision": 1}


async def test_replay_same_key_same_payload_no_rewrite():
    # INV-002: same idempotency_key + same digest -> replay (no second write, revision frozen)
    store = open_memory_store(":memory:", embedding_dim=EMBEDDING_DIM)
    c1 = _chunk("c1")
    assert await store.upsert_many([c1], idempotency_key="k1", ctx=_ctx()) == {
        "upserted": 1,
        "replayed": False,
    }
    assert await store.upsert_many([c1], idempotency_key="k1", ctx=_ctx()) == {
        "upserted": 1,
        "replayed": True,
    }
    assert (await store.get("c1"))["revision"] == 1  # replay did not re-write / re-increment
    assert _row_count(store, "memory_chunks") == 1


async def test_conflict_same_key_different_payload_raises_and_keeps_first():
    # INV-002: same key, different digest -> IdempotencyConflict; the first batch is intact
    store = open_memory_store(":memory:", embedding_dim=EMBEDDING_DIM)
    first = _chunk("c1", content="first")
    await store.upsert_many([first], idempotency_key="k1", ctx=_ctx())
    with pytest.raises(IdempotencyConflict):
        await store.upsert_many(
            [_chunk("c1", content="second")], idempotency_key="k1", ctx=_ctx()
        )
    assert await store.get("c1") == {**first, "revision": 1}  # untouched


async def test_optimistic_lock_stale_expected_revision_raises_nothing_written():
    # INV-003: a stale expected_revisions entry rolls back the whole batch
    store = open_memory_store(":memory:", embedding_dim=EMBEDDING_DIM)
    c1 = _chunk("c1", content="v1")
    await store.upsert_many([c1], idempotency_key="k1", ctx=_ctx())  # revision 1
    with pytest.raises(RevisionMismatch):
        await store.upsert_many(
            [_chunk("c1", content="v2")],
            idempotency_key="k2",  # distinct key: not replay/conflict
            ctx=_ctx(),
            expected_revisions={"c1": 5},  # stale: stored revision is 1
        )
    assert await store.get("c1") == {**c1, "revision": 1}  # nothing written
    assert _row_count(store, "memory_chunks") == 1


async def test_optimistic_lock_match_upserts_and_increments_revision():
    # INV-003: a matching expected_revisions writes and increments (1 -> 2)
    store = open_memory_store(":memory:", embedding_dim=EMBEDDING_DIM)
    await store.upsert_many([_chunk("c1", content="v1")], idempotency_key="k1", ctx=_ctx())
    v2 = _chunk("c1", content="v2")
    assert await store.upsert_many(
        [v2], idempotency_key="k2", ctx=_ctx(), expected_revisions={"c1": 1}
    ) == {"upserted": 1, "replayed": False}
    assert await store.get("c1") == {**v2, "revision": 2}  # re-upsert increments
    assert _row_count(store, "memory_chunks") == 1


async def test_injection_rule_injected_without_source_raises_no_write():
    # INV-007: origin == injected_context requires injection_source
    store = open_memory_store(":memory:", embedding_dim=EMBEDDING_DIM)
    bad = _chunk("c1", origin="injected_context")  # no injection_source
    with pytest.raises(InvalidArguments):
        await store.upsert_many([bad], idempotency_key="k1", ctx=_ctx())
    assert _row_count(store, "memory_chunks") == 0


async def test_injection_rule_non_injected_with_source_raises_no_write():
    # INV-007: a non-injected record carrying injection_source is rejected
    store = open_memory_store(":memory:", embedding_dim=EMBEDDING_DIM)
    bad = _chunk("c1", origin="worldtree", injection_source="elsewhere")
    with pytest.raises(InvalidArguments):
        await store.upsert_many([bad], idempotency_key="k1", ctx=_ctx())
    assert _row_count(store, "memory_chunks") == 0


# --- search ---

async def test_basic_search_ranks_by_cosine_with_recalled_view():
    store = open_memory_store(":memory:", embedding_dim=EMBEDDING_DIM)
    scope = {"end_user": "u1"}
    c1 = _chunk("c1", embedding=_vec(1.0, 0.0), scope=scope)
    c3 = _chunk("c3", embedding=_vec(0.9, 0.1), scope=scope)
    await store.upsert_many(
        [c1, _chunk("c2", embedding=_vec(0.0, 1.0), scope=scope), c3],
        idempotency_key="k1",
        ctx=_ctx(),
    )
    results = await store.search(_vec(1.0, 0.0), top_k=2, scope_all=scope)
    assert [r["chunk_id"] for r in results] == ["c1", "c3"]  # nearest to [1,0] by cosine
    top = results[0]
    assert top["chunk"] == c1  # verbatim chunk, no revision attached
    assert top["recalled_view"] == {"summary": "distillate-c1"}  # = chunk["distillate"]
    assert top["revision"] == 1
    assert isinstance(top["score"], float)


async def test_scope_isolation_excludes_other_scope_even_if_closer():
    # INV-005: an out-of-scope chunk that scores HIGHER must not leak; only in-scope returned
    store = open_memory_store(":memory:", embedding_dim=EMBEDDING_DIM)
    await store.upsert_many(
        [
            _chunk("u2-near", embedding=_vec(1.0, 0.0), scope={"end_user": "u2"}),  # closest
            _chunk("u1-far", embedding=_vec(0.0, 1.0), scope={"end_user": "u1"}),  # in-scope, far
        ],
        idempotency_key="k1",
        ctx=_ctx(),
    )
    results = await store.search(_vec(1.0, 0.0), top_k=2, scope_all={"end_user": "u1"})
    assert [r["chunk_id"] for r in results] == ["u1-far"]  # u2-near excluded despite ranking first


async def test_search_empty_store_returns_empty():
    store = open_memory_store(":memory:", embedding_dim=EMBEDDING_DIM)
    assert await store.search(_vec(1.0), top_k=5) == []


async def test_search_non_empty_metadata_filter_rejected():
    # PRE-002: v1 advertises no filterable metadata fields
    store = open_memory_store(":memory:", embedding_dim=EMBEDDING_DIM)
    with pytest.raises(InvalidArguments):
        await store.search(_vec(1.0), top_k=5, metadata_filter={"x": 1})


async def test_search_wrong_vector_dim_rejected():
    # PRE-001: vector length must equal the pinned embedding_dim
    store = open_memory_store(":memory:", embedding_dim=EMBEDDING_DIM)
    with pytest.raises(InvalidArguments):
        await store.search([1.0, 0.0], top_k=5)


async def test_search_non_dict_scope_all_rejected():
    # search STEP 1: scope_all must be a flat {axis: value} dict
    store = open_memory_store(":memory:", embedding_dim=EMBEDDING_DIM)
    with pytest.raises(InvalidArguments):
        await store.search(_vec(1.0), top_k=5, scope_all="u1")


async def test_search_non_list_scope_any_rejected():
    # search STEP 1: scope_any must be a LIST of {axis: value} dicts (#11)
    store = open_memory_store(":memory:", embedding_dim=EMBEDDING_DIM)
    with pytest.raises(InvalidArguments):
        await store.search(_vec(1.0), top_k=5, scope_any={"end_user": "u1"})


async def test_search_out_of_lattice_scope_axis_rejected():
    # v0.6 scope lattice = {end_user, group, tenant, agent_self}; an axis outside
    # it is InvalidFilter (-> memory.invalid_filter 400) in EITHER field, matching the reference.
    store = open_memory_store(":memory:", embedding_dim=EMBEDDING_DIM)
    with pytest.raises(InvalidFilter):
        await store.search(_vec(1.0), top_k=5, scope_all={"bogus_axis": "x"})
    with pytest.raises(InvalidFilter):
        await store.search(_vec(1.0), top_k=5, scope_any=[{"bogus_axis": "x"}])


async def test_search_agent_self_axis_accepted():
    # agent_self became canonical at wire v0.5 (#10) — admitted, not rejected.
    store = open_memory_store(":memory:", embedding_dim=EMBEDDING_DIM)
    await store.upsert_many(
        [_chunk("a1", scope={"agent_self": "ratatoskr:smoke"})],
        idempotency_key="k1",
        ctx=_ctx(),
    )
    results = await store.search(
        _vec(1.0), top_k=5, scope_all={"agent_self": "ratatoskr:smoke"}
    )
    assert [r["chunk_id"] for r in results] == ["a1"]


async def test_search_top_k_zero_returns_empty():
    # POST-001: at most top_k — zero means zero
    store = open_memory_store(":memory:", embedding_dim=EMBEDDING_DIM)
    await store.upsert_many([_chunk("c1")], idempotency_key="k1", ctx=_ctx())
    assert await store.search(_vec(1.0), top_k=0, scope_all={"end_user": "u1"}) == []


async def test_search_no_scope_matches_all():
    # v0.6: both fields empty -> no scope constraint (match all, within top_k).
    store = open_memory_store(":memory:", embedding_dim=EMBEDDING_DIM)
    await store.upsert_many(
        [
            _chunk("u1", scope={"end_user": "u1"}),
            _chunk("u2", scope={"end_user": "u2"}),
        ],
        idempotency_key="k1",
        ctx=_ctx(),
    )
    results = await store.search(_vec(1.0), top_k=10)
    assert {r["chunk_id"] for r in results} == {"u1", "u2"}


async def test_search_scope_any_unions_across_scopes():
    # v0.6 (#11): scope_any is OR/union over a LIST of conjunctive scopes. A {end_user:u1}
    # chunk AND an {agent_self:a} chunk are BOTH recalled in ONE call — the capability
    # that resolves the #295/#297 silent-zero AND foot-gun (subset-scoped chunks now recall).
    store = open_memory_store(":memory:", embedding_dim=EMBEDDING_DIM)
    await store.upsert_many(
        [
            _chunk("subj", embedding=_vec(1.0, 0.0), scope={"end_user": "u1"}),
            _chunk("self", embedding=_vec(0.9, 0.1), scope={"agent_self": "ratatoskr:sindra"}),
            _chunk("other", embedding=_vec(0.8, 0.2), scope={"end_user": "u9"}),
        ],
        idempotency_key="k1",
        ctx=_ctx(),
    )
    results = await store.search(
        _vec(1.0, 0.0),
        top_k=10,
        scope_any=[{"end_user": "u1"}, {"agent_self": "ratatoskr:sindra"}],
    )
    assert {r["chunk_id"] for r in results} == {"subj", "self"}  # union; u9 excluded


async def test_search_scope_all_and_scope_any_compose_by_and():
    # v0.6: a record passes iff (record ⊇ scope_all) AND (matches ≥1 scope_any element).
    store = open_memory_store(":memory:", embedding_dim=EMBEDDING_DIM)
    await store.upsert_many(
        [
            # tenant t1 AND (end_user u1 OR u2) — only these pass
            _chunk("t1u1", embedding=_vec(1.0, 0.0), scope={"tenant": "t1", "end_user": "u1"}),
            _chunk("t1u2", embedding=_vec(0.9, 0.1), scope={"tenant": "t1", "end_user": "u2"}),
            _chunk("t1u9", embedding=_vec(0.8, 0.2), scope={"tenant": "t1", "end_user": "u9"}),
            _chunk("t2u1", embedding=_vec(0.7, 0.3), scope={"tenant": "t2", "end_user": "u1"}),
        ],
        idempotency_key="k1",
        ctx=_ctx(),
    )
    results = await store.search(
        _vec(1.0, 0.0),
        top_k=10,
        scope_all={"tenant": "t1"},
        scope_any=[{"end_user": "u1"}, {"end_user": "u2"}],
    )
    assert {r["chunk_id"] for r in results} == {"t1u1", "t1u2"}  # t1u9 fails any; t2u1 fails all


async def test_scope_isolation_fills_top_k_from_in_scope_past_higher_out_of_scope():
    # INV-005: top_k counts IN-SCOPE hits. An out-of-scope chunk ranking #1 is skipped,
    # and top_k is still filled from the in-scope set when enough in-scope chunks exist.
    store = open_memory_store(":memory:", embedding_dim=EMBEDDING_DIM)
    await store.upsert_many(
        [
            _chunk("u2-nearest", embedding=_vec(1.0, 0.0), scope={"end_user": "u2"}),  # ranks #1
            _chunk("u1-near", embedding=_vec(0.95, 0.05), scope={"end_user": "u1"}),
            _chunk("u1-mid", embedding=_vec(0.8, 0.2), scope={"end_user": "u1"}),
            _chunk("u1-far", embedding=_vec(0.0, 1.0), scope={"end_user": "u1"}),
        ],
        idempotency_key="k1",
        ctx=_ctx(),
    )
    results = await store.search(_vec(1.0, 0.0), top_k=2, scope_all={"end_user": "u1"})
    # exactly top_k in-scope (the 2 nearest u1 chunks); the higher-ranked u2 chunk is excluded
    assert [r["chunk_id"] for r in results] == ["u1-near", "u1-mid"]


# --- get / get_many ---

async def test_get_absent_returns_none():
    store = open_memory_store(":memory:", embedding_dim=EMBEDDING_DIM)
    assert await store.get("nope") is None


async def test_get_many_returns_found_records_only():
    store = open_memory_store(":memory:", embedding_dim=EMBEDDING_DIM)
    c1 = _chunk("c1")
    await store.upsert_many([c1], idempotency_key="k1", ctx=_ctx())
    assert await store.get_many(["c1", "absent"]) == [{**c1, "revision": 1}]


# --- delete_many ---

async def test_delete_hit_removes_chunk_and_vec_row():
    store = open_memory_store(":memory:", embedding_dim=EMBEDDING_DIM)
    await store.upsert_many([_chunk("c1"), _chunk("c2")], idempotency_key="k1", ctx=_ctx())
    assert await store.delete_many(["c1"]) == {"deleted": 1}
    assert await store.get("c1") is None
    assert _row_count(store, "memory_chunks") == 1
    assert _row_count(store, "memory_vec") == 1  # c1's vec row gone too (no orphan)
    # delete_hit: search no longer surfaces it (vec/chunk coupling held)
    hits = await store.search(_vec(1.0), top_k=5, scope_all={"end_user": "u1"})
    assert all(r["chunk_id"] != "c1" for r in hits)


async def test_delete_absent_counts_zero():
    store = open_memory_store(":memory:", embedding_dim=EMBEDDING_DIM)
    assert await store.delete_many(["nope"]) == {"deleted": 0}


# --- build_memory_provider_app ---

def test_build_app_exposes_handshake_and_memory_routes():
    store = open_memory_store(":memory:", embedding_dim=EMBEDDING_DIM)
    app = build_memory_provider_app(store, heimdall_key=b"secret-key")
    routes = {getattr(r, "path", None): r for r in app.routes}
    assert "/bifrost/handshake" in routes
    assert "/bifrost/memory-call" in routes
    assert "POST" in routes["/bifrost/memory-call"].methods
    assert "POST" in routes["/bifrost/handshake"].methods  # both routes are POST (incl. POST)


def test_build_app_rejects_empty_key():
    store = open_memory_store(":memory:", embedding_dim=EMBEDDING_DIM)
    with pytest.raises(ValueError):
        build_memory_provider_app(store, heimdall_key=b"")


# --- #195 conformance: parity vs the reference store through the real engine ---

def _dispatch_ctx(*scopes: str, session_id: str = "actor-1"):
    return types.SimpleNamespace(
        scope=list(scopes), session_id=session_id, jwt_sub=session_id, job_id=None
    )


def _ref_record(chunk_id: str, vector: list[float], *, end_user: str = "u1") -> dict:
    # Mirrors bifrost's reference `record` helper so the envelope validates.
    return {
        "id": chunk_id,
        "embedding": vector,
        "distillate": {"text": chunk_id},
        "metadata": {"worldtree.appraisal_confidence": 0.8},
        "scope": {"end_user": end_user, "tenant": "t1"},
        "origin": "worldtree",
        "source_role": "assistant",
        "trust_tier": "tier-3",
        "provenance": {"trace": chunk_id},
    }


async def test_parity_upsert_many_vs_reference_through_dispatch():
    from bifrost.consumer.testing import InMemoryMemoryStore
    from bifrost.memory import dispatch_memory_call

    ref = InMemoryMemoryStore()
    mine = open_memory_store(":memory:", embedding_dim=2)
    wctx = _dispatch_ctx("memory:write")
    env = {
        "operation": "upsert_many",
        "args": {"records": [_ref_record("a", [1.0, 0.0]), _ref_record("b", [0.0, 1.0])]},
        "idempotency_key": "k1",
    }
    # happy persist + replay: wire bodies must agree
    assert await dispatch_memory_call(env, wctx, ref) == await dispatch_memory_call(env, wctx, mine)
    assert await dispatch_memory_call(env, wctx, ref) == await dispatch_memory_call(env, wctx, mine)


async def test_parity_search_ranked_ids_vs_reference_through_dispatch():
    from bifrost.consumer.testing import InMemoryMemoryStore
    from bifrost.memory import dispatch_memory_call

    ref = InMemoryMemoryStore()
    mine = open_memory_store(":memory:", embedding_dim=2)
    wctx = _dispatch_ctx("memory:write")
    rctx = _dispatch_ctx("memory:read")
    up = {
        "operation": "upsert_many",
        "args": {
            "records": [
                _ref_record("a", [1.0, 0.0]),
                _ref_record("b", [0.0, 1.0]),
                _ref_record("c", [0.9, 0.1]),
            ]
        },
        "idempotency_key": "k1",
    }
    await dispatch_memory_call(up, wctx, ref)
    await dispatch_memory_call(up, wctx, mine)
    search_env = {
        "operation": "search",
        "args": {"vector": [1.0, 0.0], "top_k": 2, "scope_all": {"end_user": "u1"}},
    }
    rstatus, rbody = await dispatch_memory_call(search_env, rctx, ref)
    mstatus, mbody = await dispatch_memory_call(search_env, rctx, mine)
    assert rstatus == mstatus == 200
    # #195: same ranked chunk_ids and the same per-result shape (scores may differ in the
    # last float digit between vec0's cosine and the reference's Python cosine).
    assert [r["chunk_id"] for r in rbody["results"]] == [r["chunk_id"] for r in mbody["results"]]
    assert set(rbody["results"][0]) == set(mbody["results"][0])


async def test_parity_expected_revisions_vs_reference_through_dispatch():
    # #195: pins the partial-map optimistic-lock semantics against the reference
    # (does an expected_revisions map that omits some batch records lock only the
    # listed ones?). Resolves the contract's ambiguous "each record's stored revision".
    from bifrost.consumer.testing import InMemoryMemoryStore
    from bifrost.memory import dispatch_memory_call

    ref = InMemoryMemoryStore()
    mine = open_memory_store(":memory:", embedding_dim=2)
    wctx = _dispatch_ctx("memory:write")

    seed = {
        "operation": "upsert_many",
        "args": {"records": [_ref_record("a", [1.0, 0.0]), _ref_record("b", [0.0, 1.0])]},
        "idempotency_key": "seed",
    }
    ref_seed = await dispatch_memory_call(seed, wctx, ref)
    mine_seed = await dispatch_memory_call(seed, wctx, mine)
    assert ref_seed == mine_seed

    # partial map: only "a" is locked (revision 1); "b" is omitted from expected_revisions
    partial = {
        "operation": "upsert_many",
        "args": {
            "records": [_ref_record("a", [1.0, 0.0]), _ref_record("b", [0.0, 1.0])],
            "expected_revisions": {"a": 1},
        },
        "idempotency_key": "partial",
    }
    assert await dispatch_memory_call(partial, wctx, ref) == await dispatch_memory_call(
        partial, wctx, mine
    )

    # stale lock: both map to the same RevisionMismatch wire error
    stale = {
        "operation": "upsert_many",
        "args": {"records": [_ref_record("a", [1.0, 0.0])], "expected_revisions": {"a": 99}},
        "idempotency_key": "stale",
    }
    assert await dispatch_memory_call(stale, wctx, ref) == await dispatch_memory_call(
        stale, wctx, mine
    )