Initial release: iai-mcp v0.1.0

Co-Authored-By: Claude <noreply@anthropic.com>
Co-Authored-By: XNLLLLH <XNLLLLH@users.noreply.github.com>

tests/test_runtime_graph_cache.py

@@ -0,0 +1,550 @@
+"""Plan 05-09 (P4.A) — runtime_graph_cache tests.
+
+The cache persists the Leiden community assignment + rich-club node
+list next to the LanceDB store. On subsequent ``build_runtime_graph``
+calls, a valid cache skips the Leiden + rich-club computations;
+invalid cache falls through to a clean rebuild.
+
+Covered contracts:
+
+    1. save() creates the JSON file
+    2. try_load on unchanged store returns the cached pair
+    3. adding a record invalidates the cache (key mismatch)
+    4. CACHE_VERSION mismatch triggers rebuild (forward-compat fence)
+    5. corrupt JSON falls through to a clean None
+    6. absent file returns None cleanly
+    7. build_runtime_graph on second call avoids detect_communities
+    8. build_runtime_graph on store change triggers detect_communities
+    9. atomic write: interrupted save leaves old cache intact
+   10. invalidate() deletes the cache file
+"""
+from __future__ import annotations
+
+import json
+from pathlib import Path
+from unittest import mock
+from uuid import UUID, uuid4
+
+import pytest
+
+from iai_mcp import retrieve, runtime_graph_cache
+from iai_mcp.community import CommunityAssignment
+from iai_mcp.store import MemoryStore
+
+
+# --------------------------------------------------------------------------- fixtures
+
+
+@pytest.fixture(autouse=True)
+def _isolated_keyring(monkeypatch: pytest.MonkeyPatch):
+    import keyring as _keyring
+
+    fake: dict[tuple[str, str], str] = {}
+    monkeypatch.setattr(_keyring, "get_password", lambda s, u: fake.get((s, u)))
+    monkeypatch.setattr(
+        _keyring, "set_password", lambda s, u, p: fake.__setitem__((s, u), p)
+    )
+    monkeypatch.setattr(
+        _keyring, "delete_password", lambda s, u: fake.pop((s, u), None)
+    )
+    yield fake
+
+
+@pytest.fixture
+def store(tmp_path: Path) -> MemoryStore:
+    """Fresh MemoryStore in tmp_path/lancedb with the cache file path set
+    to tmp_path/runtime_graph_cache.json."""
+    s = MemoryStore(path=tmp_path / "lancedb")
+    # Override root so the cache file lives in tmp_path, not the real store
+    # root (which would be tmp_path/lancedb — still fine, but explicit).
+    s.root = tmp_path
+    return s
+
+
+def _read_decrypted_cache(store: MemoryStore, path: Path) -> dict:
+    """Phase 07.9 W3: decrypt the v3 ciphertext sidecar and return the
+    underlying JSON dict. Tests that pre-07.9 read the cache via
+    ``json.load(f)`` go through this helper instead.
+    """
+    raw_text = path.read_text(encoding="utf-8")
+    if not raw_text.startswith("iai:enc:v1:"):
+        # Legacy / hand-written plaintext (test 4 simulating v2).
+        return json.loads(raw_text)
+    from iai_mcp.crypto import decrypt_field
+    plaintext = decrypt_field(
+        raw_text,
+        store._key(),
+        runtime_graph_cache._CACHE_AAD,
+    )
+    return json.loads(plaintext)
+
+
+def _write_encrypted_cache(store: MemoryStore, path: Path, data: dict) -> None:
+    """Inverse of _read_decrypted_cache: encrypt and write a hand-modified
+    JSON dict back to the sidecar in v3 ciphertext format.
+    """
+    from iai_mcp.crypto import encrypt_field
+    plaintext = json.dumps(data)
+    ciphertext = encrypt_field(
+        plaintext,
+        store._key(),
+        runtime_graph_cache._CACHE_AAD,
+    )
+    path.write_text(ciphertext, encoding="ascii")
+
+
+def _make_assignment(n_communities: int = 2) -> CommunityAssignment:
+    comms = [uuid4() for _ in range(n_communities)]
+    nodes = [uuid4() for _ in range(n_communities * 3)]
+    return CommunityAssignment(
+        node_to_community={
+            nodes[i]: comms[i // 3] for i in range(len(nodes))
+        },
+        community_centroids={c: [0.1, 0.2, 0.3] for c in comms},
+        modularity=0.42,
+        backend="leiden-networkx",
+        top_communities=comms,
+        mid_regions={c: nodes[i * 3:(i + 1) * 3] for i, c in enumerate(comms)},
+    )
+
+
+# --------------------------------------------------------------------------- Test 1
+
+
+def test_save_creates_json_file(store):
+    assignment = _make_assignment()
+    rich_club = [uuid4() for _ in range(5)]
+    ok = runtime_graph_cache.save(store, assignment, rich_club)
+    assert ok is True
+    path = runtime_graph_cache._cache_path(store)
+    assert path.exists()
+    # W3: file is now AES-256-GCM-wrapped — decrypt before
+    # inspecting the JSON shape.
+    raw = path.read_text(encoding="utf-8")
+    assert raw.startswith("iai:enc:v1:"), (
+        "Phase 07.9 W3: cache must be v3 ciphertext on disk"
+    )
+    data = _read_decrypted_cache(store, path)
+    assert data["cache_version"] == runtime_graph_cache.CACHE_VERSION
+    assert "assignment" in data
+    assert "rich_club" in data
+    assert "key" in data
+
+
+# --------------------------------------------------------------------------- Test 2
+
+
+def test_try_load_round_trip_on_unchanged_store(store):
+    assignment = _make_assignment()
+    rich_club = [uuid4() for _ in range(3)]
+    runtime_graph_cache.save(store, assignment, rich_club)
+
+    loaded = runtime_graph_cache.try_load(store)
+    assert loaded is not None
+    # try_load returns (assignment, rich_club, node_payload, ...).
+    # node_payload is the v2 blob — None (or empty) when legacy 2-arg
+    # save() shape was used.
+    # 4th element is max_degree (int >= 0).
+    loaded_assignment, loaded_rich_club, _node_payload, _max_degree = loaded
+    assert loaded_assignment.backend == assignment.backend
+    assert loaded_assignment.modularity == pytest.approx(assignment.modularity)
+    assert set(loaded_assignment.top_communities) == set(assignment.top_communities)
+    assert set(loaded_rich_club) == set(rich_club)
+
+
+# --------------------------------------------------------------------------- Test 3
+
+
+def test_key_mismatch_invalidates_cache(store):
+    # Save with original key.
+    runtime_graph_cache.save(store, _make_assignment(), [uuid4()])
+    path = runtime_graph_cache._cache_path(store)
+    assert path.exists()
+
+    # Simulate a store change by forging a wrong key in the saved file.
+    # W3: decrypt → mutate → re-encrypt round-trip.
+    data = _read_decrypted_cache(store, path)
+    data["key"][0] = 999  # bogus records_count
+    _write_encrypted_cache(store, path, data)
+
+    assert runtime_graph_cache.try_load(store) is None
+
+
+# --------------------------------------------------------------------------- Test 4
+
+
+def test_cache_version_mismatch_triggers_rebuild(store):
+    runtime_graph_cache.save(store, _make_assignment(), [uuid4()])
+    path = runtime_graph_cache._cache_path(store)
+    # W3: decrypt → mutate version → re-encrypt round-trip.
+    data = _read_decrypted_cache(store, path)
+    data["cache_version"] = "old-format-v0"
+    _write_encrypted_cache(store, path, data)
+
+    assert runtime_graph_cache.try_load(store) is None
+
+
+# --------------------------------------------------------------------------- Test 5
+
+
+def test_corrupt_json_returns_none(store):
+    path = runtime_graph_cache._cache_path(store)
+    path.parent.mkdir(parents=True, exist_ok=True)
+    path.write_text("{not valid json at all")
+    assert runtime_graph_cache.try_load(store) is None
+
+
+# --------------------------------------------------------------------------- Test 6
+
+
+def test_absent_cache_returns_none(store):
+    path = runtime_graph_cache._cache_path(store)
+    assert not path.exists()
+    assert runtime_graph_cache.try_load(store) is None
+
+
+# --------------------------------------------------------------------------- Test 7
+
+
+def test_build_runtime_graph_uses_cache_on_second_call(store):
+    # First call: detect_communities runs, cache is written.
+    with mock.patch(
+        "iai_mcp.community.detect_communities",
+        wraps=__import__("iai_mcp.community", fromlist=["detect_communities"]).detect_communities,
+    ) as detect_spy:
+        retrieve.build_runtime_graph(store)
+        assert detect_spy.call_count == 1
+
+    # Second call: cache is valid, detect_communities must NOT re-run.
+    with mock.patch(
+        "iai_mcp.community.detect_communities",
+    ) as detect_spy:
+        retrieve.build_runtime_graph(store)
+        assert detect_spy.call_count == 0
+
+
+# --------------------------------------------------------------------------- Test 8
+
+
+def test_build_runtime_graph_invalidates_on_record_added(store, tmp_path):
+    """Adding a record bumps records_count -> cache key changes -> rebuild."""
+    # First call seeds the cache on an empty store.
+    retrieve.build_runtime_graph(store)
+    assert runtime_graph_cache._cache_path(store).exists()
+
+    # Insert a real record via the store so records_count increments.
+    from datetime import datetime, timezone
+    from iai_mcp.types import MemoryRecord
+
+    rec = MemoryRecord(
+        id=uuid4(),
+        tier="episodic",
+        literal_surface="x",
+        aaak_index="",
+        embedding=[0.0] * store.embed_dim,
+        community_id=None,
+        centrality=0.0,
+        detail_level=2,
+        pinned=False,
+        stability=0.0,
+        difficulty=0.0,
+        last_reviewed=None,
+        never_decay=False,
+        never_merge=False,
+        provenance=[],
+        created_at=datetime.now(timezone.utc),
+        updated_at=datetime.now(timezone.utc),
+        tags=["t"],
+        language="en",
+    )
+    store.insert(rec)
+
+    # Second build sees records_count changed -> cache invalid -> rebuild.
+    with mock.patch(
+        "iai_mcp.community.detect_communities",
+        wraps=__import__("iai_mcp.community", fromlist=["detect_communities"]).detect_communities,
+    ) as detect_spy:
+        retrieve.build_runtime_graph(store)
+        assert detect_spy.call_count == 1
+
+
+# --------------------------------------------------------------------------- Test 9
+
+
+def test_save_is_atomic_leaves_old_file_on_error(store, monkeypatch):
+    """If os.replace raises mid-save the .tmp is cleaned up and the old
+    cache file (if any) is untouched."""
+    # Seed a valid cache first.
+    original_assignment = _make_assignment()
+    runtime_graph_cache.save(store, original_assignment, [uuid4()])
+    path = runtime_graph_cache._cache_path(store)
+    original_text = path.read_text()
+
+    # Now break os.replace and try to save again.
+    monkeypatch.setattr(
+        "iai_mcp.runtime_graph_cache.os.replace",
+        mock.Mock(side_effect=OSError("rename failed")),
+    )
+    ok = runtime_graph_cache.save(store, _make_assignment(), [uuid4()])
+    assert ok is False
+    # Original cache still intact.
+    assert path.read_text() == original_text
+    # Temp file was cleaned up.
+    tmp_path = path.with_suffix(path.suffix + ".tmp")
+    assert not tmp_path.exists()
+
+
+# --------------------------------------------------------------------------- Test 10
+
+
+def test_invalidate_removes_cache_file(store):
+    runtime_graph_cache.save(store, _make_assignment(), [uuid4()])
+    path = runtime_graph_cache._cache_path(store)
+    assert path.exists()
+
+    runtime_graph_cache.invalidate(store)
+    assert not path.exists()
+
+    # Idempotent — second invalidate on missing file is a no-op.
+    runtime_graph_cache.invalidate(store)
+
+
+# --------------------------------------------------------------------------- Test 11 (bonus)
+
+
+def test_embed_dim_change_invalidates(store):
+    """swapping embedders (1024d -> 384d) must force a
+    rebuild. The cache_key includes store.embed_dim so this is
+    automatic — no separate signal needed."""
+    runtime_graph_cache.save(store, _make_assignment(), [uuid4()])
+    assert runtime_graph_cache.try_load(store) is not None
+
+    # Simulate an embedder swap by monkeypatching store.embed_dim.
+    store._embed_dim = 1024  # underlying attr
+    assert runtime_graph_cache.try_load(store) is None
+
+
+# --------------------------------------------------------------------------- tests
+
+
+def test_save_drops_oversize_community_centroids(store):
+    """when ``assignment.community_centroids`` alone overflows
+    the 10 MiB cap (the scenario on an all-isolated graph where Leiden
+    gives one community per node), the iterative drop path must prune
+    centroids too — not just node_payload — and produce a file ≤ cap.
+
+    Pre-F-09 behaviour: the single-shot drop cleared node_payload,
+    re-serialised, saw the payload still > cap, and shipped it anyway.
+    Post-node_payload drops first, centroids drop next,
+    node_to_community + modularity + top_communities survive.
+    """
+    # 2000 communities × 1024-dim float vectors ≈ 40 MiB JSON-encoded,
+    # well over the 10 MiB cap. This shape mirrors the live bug.
+    big_centroids = {uuid4(): [0.123456] * 1024 for _ in range(2000)}
+    big_node_to_community = {uuid4(): uuid4() for _ in range(50)}
+    assignment = CommunityAssignment(
+        node_to_community=big_node_to_community,
+        community_centroids=big_centroids,
+        modularity=0.37,
+        backend="leiden-networkx",
+        top_communities=list(big_centroids.keys())[:5],
+        mid_regions={c: [] for c in list(big_centroids.keys())[:5]},
+    )
+    rich_club = [uuid4() for _ in range(10)]
+
+    # Non-empty node_payload so the first drop step has something to remove.
+    node_payload = {
+        str(uuid4()): {
+            "embedding": [0.0] * 384,
+            "surface": "probe",
+            "centrality": 0.1,
+            "tier": "episodic",
+            "pinned": False,
+            "tags": [],
+            "language": "en",
+        }
+    }
+
+    ok = runtime_graph_cache.save(store, assignment, rich_club, node_payload=node_payload)
+    assert ok is True
+
+    path = runtime_graph_cache._cache_path(store)
+    assert path.exists()
+
+    # File respects the cap.
+    size = path.stat().st_size
+    assert size <= runtime_graph_cache.MAX_CACHE_BYTES, (
+        f"cache file {size} bytes exceeds cap "
+        f"{runtime_graph_cache.MAX_CACHE_BYTES}"
+    )
+
+    data = _read_decrypted_cache(store, path)
+
+    # Both pre-F-09 drop candidates emptied in order.
+    assert data["node_payload"] == {}
+    assert data["assignment"]["community_centroids"] == {}
+
+    # Authoritative fields survived.
+    assert data["assignment"]["modularity"] == pytest.approx(0.37)
+    assert data["assignment"]["backend"] == "leiden-networkx"
+    assert len(data["assignment"]["node_to_community"]) == len(big_node_to_community)
+    assert len(data["assignment"]["top_communities"]) == 5
+    # rich_club preserved.
+    assert len(data["rich_club"]) == len(rich_club)
+
+
+def test_save_small_payload_survives_unchanged(store):
+    """Negative case: when the payload is comfortably under the cap,
+    no drops fire — centroids, mid_regions, and node_payload round-trip
+    intact. Guards against an over-eager iterative-drop path.
+    """
+    assignment = _make_assignment(n_communities=2)
+    rich_club = [uuid4() for _ in range(3)]
+    node_payload = {
+        str(uuid4()): {
+            "embedding": [0.1] * 384,
+            "surface": "hello",
+            "centrality": 0.2,
+            "tier": "episodic",
+            "pinned": False,
+            "tags": ["t"],
+            "language": "en",
+        }
+        for _ in range(5)
+    }
+
+    ok = runtime_graph_cache.save(store, assignment, rich_club, node_payload=node_payload)
+    assert ok is True
+
+    path = runtime_graph_cache._cache_path(store)
+    data = _read_decrypted_cache(store, path)
+
+    # Well under the cap.
+    assert path.stat().st_size < runtime_graph_cache.MAX_CACHE_BYTES
+
+    # Nothing pruned.
+    assert data["node_payload"] != {}
+    assert len(data["node_payload"]) == 5
+    assert data["assignment"]["community_centroids"] != {}
+    assert len(data["assignment"]["community_centroids"]) == 2
+    assert data["assignment"]["mid_regions"] != {}
+    assert len(data["assignment"]["mid_regions"]) == 2
+
+
+# --------------------------------------------------------------------------- W3 / tests
+
+
+def test_save_writes_ciphertext_no_plaintext_surface(store):
+    """W3 / saved sidecar must NOT contain plaintext surface bytes
+    anywhere on disk. The whole JSON payload is wrapped in AES-256-GCM."""
+    canary = "PLAINTEXT_CANARY_4d7f_07_9_W3"
+    rid = uuid4()
+    node_payload = {
+        str(rid): {
+            "embedding": [0.1] * 384,
+            "surface": canary,
+            "centrality": 0.5,
+            "tier": "episodic",
+            "pinned": False,
+            "tags": ["t"],
+            "language": "en",
+        },
+    }
+    ok = runtime_graph_cache.save(store, _make_assignment(), [uuid4()],
+                                  node_payload=node_payload, max_degree=3)
+    assert ok is True
+    path = runtime_graph_cache._cache_path(store)
+    raw_bytes = path.read_bytes()
+
+    # Plaintext canary must not appear anywhere in the on-disk bytes.
+    assert canary.encode("utf-8") not in raw_bytes, (
+        "plaintext surface canary leaked into the on-disk sidecar"
+    )
+    # Ciphertext envelope present.
+    assert raw_bytes.startswith(b"iai:enc:v1:"), (
+        f"expected v3 ciphertext envelope; got prefix {raw_bytes[:32]!r}"
+    )
+
+
+def test_save_then_try_load_preserves_surface_byte_for_byte(store):
+    """W3 / surface round-trips through encrypt → decrypt cleanly,
+    including non-ASCII (byte-for-byte across encryption)."""
+    rid = uuid4()
+    surface = "user сказал важное — please remember this 重要"
+    node_payload = {
+        str(rid): {
+            "embedding": [0.42] * 384,
+            "surface": surface,
+            "centrality": 0.42,
+            "tier": "episodic",
+            "pinned": True,
+            "tags": ["t1", "t2"],
+            "language": "ru",
+        },
+    }
+    runtime_graph_cache.save(store, _make_assignment(), [rid],
+                             node_payload=node_payload, max_degree=7)
+
+    loaded = runtime_graph_cache.try_load(store)
+    assert loaded is not None
+    _, _, payload, max_deg = loaded
+    assert payload is not None
+    assert payload[str(rid)]["surface"] == surface
+    assert payload[str(rid)]["centrality"] == pytest.approx(0.42)
+    assert payload[str(rid)]["tags"] == ["t1", "t2"]
+    assert payload[str(rid)]["language"] == "ru"
+    assert max_deg == 7
+
+
+def test_v2_plaintext_lazy_migrates_to_v3(store):
+    """W3 / a hand-written legacy v2 plaintext file is read once,
+    then re-saved under the v3 ciphertext format on the same call.
+    Subsequent reads see only ciphertext on disk."""
+    path = runtime_graph_cache._cache_path(store)
+    path.parent.mkdir(parents=True, exist_ok=True)
+
+    # Hand-craft a v2 plaintext file in the legacy shape. Use the
+    # current store key so the cache_key matches and try_load
+    # accepts the contents.
+    rid = uuid4()
+    legacy_data = {
+        "cache_version": runtime_graph_cache.LEGACY_CACHE_VERSION_PLAINTEXT,
+        "key": list(runtime_graph_cache._cache_key(store)),
+        "assignment": runtime_graph_cache._encode_assignment(_make_assignment()),
+        "rich_club": [str(uuid4())],
+        "node_payload": {
+            str(rid): {
+                "embedding": [0.0] * 384,
+                "surface": "legacy_plain_canary",
+                "centrality": 0.0,
+                "tier": "episodic",
+                "pinned": False,
+                "tags": [],
+                "language": "en",
+            }
+        },
+        "max_degree": 1,
+        "saved_at": "2026-04-29T00:00:00Z",
+    }
+    # Force the legacy cache_version into the saved key so try_load's
+    # current_key check matches. (cache_version is the 5th element.)
+    if len(legacy_data["key"]) >= 5:
+        legacy_data["key"][4] = runtime_graph_cache.LEGACY_CACHE_VERSION_PLAINTEXT
+    path.write_text(json.dumps(legacy_data), encoding="utf-8")
+
+    # First try_load reads the v2 plaintext, decodes, and re-saves
+    # under the v3 ciphertext format.
+    loaded = runtime_graph_cache.try_load(store)
+    assert loaded is not None
+    _, _, payload, _ = loaded
+    assert payload is not None
+    assert payload[str(rid)]["surface"] == "legacy_plain_canary"
+
+    # On-disk file is now ciphertext.
+    raw_after = path.read_bytes()
+    assert raw_after.startswith(b"iai:enc:v1:"), (
+        "v2 plaintext file must be lazily migrated to v3 ciphertext"
+    )
+    assert b"legacy_plain_canary" not in raw_after, (
+        "post-migration on-disk bytes must not contain the legacy plaintext"
+    )