arborist/tests/test_session.py

"""Tests for arborist.qa.session — single-shard Merkle conversation forest."""
from __future__ import annotations

from pathlib import Path

import pytest

from arborist.qa.session import (
    SessionStore,
    Session,
    render_tree,
    _compute_node_hash,
    _compute_subtree_hash,
)


@pytest.fixture
def store_path(tmp_path: Path) -> Path:
    return tmp_path / "sessions.db"


@pytest.fixture
def store(store_path: Path) -> SessionStore:
    return SessionStore.open(store_path)


def test_create_session_mints_synthetic_root(store: SessionStore):
    sess = store.create_session(sid="sidA")
    assert sess.sid == "sidA"
    assert sess.root_bates == "sidA-000000"
    root = sess.get_node(sess.root_bates)
    assert root is not None
    assert root.question == ""
    assert root.parent_bates is None
    assert root.subtree_hash == sess.session_root
    assert sess.current_bates == sess.root_bates


def test_add_node_mints_bates_and_updates_root(store: SessionStore):
    sess = store.create_session(sid="s")
    r0 = sess.session_root
    n1 = sess.add_node("q1", "ck1", "STRICT")
    assert n1.bates == "s-000001"
    assert sess.current_bates == n1.bates
    assert sess.session_root != r0


def test_fork_via_cd_creates_sibling(store: SessionStore):
    sess = store.create_session(sid="s")
    n1 = sess.add_node("q1", "ck1", "STRICT")
    n2 = sess.add_node("q2", "ck2", "STRICT")
    sess.cd(n1.bates)
    n3 = sess.add_node("q3", "ck3", "STRICT")
    assert n3.parent_bates == n1.bates
    kids = {k.bates for k in sess.children_of(n1.bates)}
    assert kids == {n2.bates, n3.bates}


def test_cross_session_fork(store: SessionStore):
    """Fork from a node in another session — parent_bates crosses sids.
    The cross-session parent's subtree_hash recomputes to include the
    new child (Merkle ripples through the global forest)."""
    sess_a = store.create_session(sid="A")
    a1 = sess_a.add_node("hello from A", "ck-A1", "STRICT")
    a1_subtree_before = store.get_node(a1.bates).subtree_hash

    sess_b = store.create_session(sid="B")
    # cd session B's pointer to A's node, then ask in B.
    sess_b.cd(a1.bates)
    b1 = sess_b.add_node("fork from A", "ck-B1", "STRICT")
    assert b1.parent_bates == a1.bates
    assert b1.sid == "B"
    # A's node now has a child (B's node); subtree_hash changed.
    a1_subtree_after = store.get_node(a1.bates).subtree_hash
    assert a1_subtree_after != a1_subtree_before
    # A's session_root also changed (ripple up).
    a_root_subtree = store.session_root("A")
    assert a_root_subtree != ""


def test_find_fts_search_across_sessions(store: SessionStore):
    sess_a = store.create_session(sid="A")
    sess_a.add_node("spider man origin story",
                    "ck1", "STRICT",
                    answer_text="Spider-Man was created by Stan Lee.")
    sess_b = store.create_session(sid="B")
    sess_b.add_node("how do you cook risotto?",
                    "ck2", "STRICT",
                    answer_text="Stir constantly while adding broth.")

    hits = store.find("spider")
    assert len(hits) == 1
    assert hits[0].sid == "A"
    assert "spider" in hits[0].question.lower()

    hits = store.find("risotto")
    assert len(hits) == 1
    assert hits[0].sid == "B"


def test_find_skips_synthetic_root(store: SessionStore):
    """The synthetic root (seq=0, empty question) must not surface in
    FTS5 search — its question is empty so it wouldn't match anyway,
    but the WHERE seq>0 guard is an explicit second line of defense."""
    store.create_session(sid="A")
    hits = store.find("anything")
    assert hits == []


def test_find_by_cache_key_lists_all_uses(store: SessionStore):
    sess_a = store.create_session(sid="A")
    sess_b = store.create_session(sid="B")
    sess_a.add_node("q1", "shared-ck", "STRICT")
    sess_b.add_node("q2", "shared-ck", "STRICT")
    hits = store.find_by_cache_key("shared-ck")
    assert {h.sid for h in hits} == {"A", "B"}


def test_path_to_root_crosses_sessions(store: SessionStore):
    sess_a = store.create_session(sid="A")
    a1 = sess_a.add_node("q1", "ck1", "STRICT")
    sess_b = store.create_session(sid="B")
    sess_b.cd(a1.bates)
    b1 = sess_b.add_node("q2", "ck2", "STRICT")
    path = store.path_to_root(b1.bates)
    sids_in_path = [n.sid for n in path]
    # b1 (B) → a1 (A) → A's root (A)
    assert sids_in_path == ["B", "A", "A"]


def test_audit_chain_intact_after_inserts(store: SessionStore):
    sess = store.create_session(sid="A")
    sess.add_node("q1", "ck1", "STRICT")
    sess.add_node("q2", "ck2", "STRICT")
    sess2 = store.create_session(sid="B")
    sess2.add_node("q3", "ck3", "HYBRID")
    intact, breaks = store.chain_check()
    assert breaks == 0
    assert intact == 5  # 2 session_init + 3 node_added


def test_audit_chain_detects_tampering(store: SessionStore):
    sess = store.create_session(sid="A")
    sess.add_node("q1", "ck1", "STRICT")
    sess.add_node("q2", "ck2", "STRICT")
    store.conn.execute(
        "UPDATE audit_events SET body=? WHERE event_type='node_added' "
        "ORDER BY seq DESC LIMIT 1",
        ('{"tampered": true}',),
    )
    store.conn.commit()
    intact, breaks = store.chain_check()
    assert breaks >= 1


def test_resolve_accepts_bates_seq_and_label(store: SessionStore):
    sess = store.create_session(sid="sid")
    n1 = sess.add_node("q1", "ck1", "STRICT")
    sess.label_current("alpha")
    assert sess.resolve(n1.bates) == n1.bates
    assert sess.resolve("1") == n1.bates
    assert sess.resolve("alpha") == n1.bates
    assert sess.resolve("nope") is None


def test_list_sessions(store: SessionStore):
    store.create_session(sid="A")
    store.create_session(sid="B")
    lst = store.list_sessions()
    sids = {s["sid"] for s in lst}
    assert sids == {"A", "B"}


def test_render_tree_shows_cross_session_tag(store: SessionStore):
    sess_a = store.create_session(sid="A")
    a1 = sess_a.add_node("question A", "ck1", "STRICT")
    sess_b = store.create_session(sid="B")
    sess_b.cd(a1.bates)
    sess_b.add_node("forked into B", "ck2", "STRICT")
    # Render session A's tree — should include the B-minted child
    # with a [from B] tag.
    out = render_tree(sess_a)
    assert "question A" in out
    assert "forked into B" in out
    assert "[from B]" in out


def test_subtree_hash_propagates_across_sessions(store: SessionStore):
    """Inserting a node in session B under a node in session A
    changes that A-node's subtree_hash AND A's session_root."""
    sess_a = store.create_session(sid="A")
    a1 = sess_a.add_node("anchor in A", "ck1", "STRICT")
    a_root_before = sess_a.session_root

    sess_b = store.create_session(sid="B")
    sess_b.cd(a1.bates)
    sess_b.add_node("child in B", "ck2", "STRICT")

    a_root_after = store.session_root("A")
    assert a_root_after != a_root_before


def test_compute_hash_helpers_deterministic():
    nh = _compute_node_hash("", "s-000001", "q", "ck", "STRICT",
                            "2026-06-01T00:00:00Z", None)
    nh2 = _compute_node_hash("", "s-000001", "q", "ck", "STRICT",
                             "2026-06-01T00:00:00Z", None)
    assert nh == nh2
    sh = _compute_subtree_hash(nh, [])
    sh2 = _compute_subtree_hash(nh2, [])
    assert sh == sh2


def test_cited_titles_extracted_from_answer(store: SessionStore):
    sess = store.create_session(sid="A")
    ans = (
        'Spider-Man was created by Stan Lee.\n'
        '  [E1 | Spider-Man | abc123: "..."]\n'
        '  [E2 | Stan Lee | def456: "..."]'
    )
    n = sess.add_node("who created spider man?",
                      "ck1", "STRICT", answer_text=ans)
    assert n.cited_titles == ["Spider-Man", "Stan Lee"]
    assert n.n_cited_sources == 2


def test_fts_finds_by_cited_title(store: SessionStore):
    sess = store.create_session(sid="A")
    ans = '...\n  [E1 | Sesame Street | abc123: "..."]'
    sess.add_node("when did the show start?", "ck1", "STRICT",
                  answer_text=ans)
    hits = store.find("sesame")
    assert len(hits) == 1
    assert "sesame" in hits[0].cited_titles[0].lower()