arborist/tests/test_dag.py

"""Per-run Merkle-DAG provenance.

`build_run_dag` produces a deterministic Merkle root over the seven
stages of a query/ask call. `verify_run_dag` recomputes the root from
the persisted node list and confirms it matches.

The DAG is not part of cache_key (cache inputs determine the answer;
the answer determines the DAG — folding it back would create a cycle).
It rides alongside the providence record as `run_dag_root` /
`run_dag_blob` so an auditor can verify the run was constructed
exactly as recorded.
"""

from __future__ import annotations

from arborist.qa.dag import (
    build_run_dag,
    localize_failure,
    preflight_node_hash,
    verify_run_dag,
)


def _kw(**overrides):
    base = dict(
        question_hash="a" * 64,
        sources=[
            {
                "document_root": "b" * 64,
                "source_role": "primary_answer_source",
                "score": 1.0,
                "chunk_idx": 0,
            }
        ],
        context_root="b" * 64,
        conversation_hash="c" * 64,
        answer_text="The cat is on the mat.",
        audit_mode="STRICT",
        verifier_method="quote",
        n_quotes=1,
        n_verified=1,
        claim_statuses=[
            {"text": "the cat", "status": "VERIFIED_QUOTE", "method": "quote"}
        ],
        lookup_path="miss",
    )
    base.update(overrides)
    return base


def test_dag_root_is_deterministic():
    """Same inputs → same root, byte-for-byte."""
    a = build_run_dag(**_kw())
    b = build_run_dag(**_kw())
    assert a["root"] == b["root"]
    assert len(a["nodes"]) == 7


def test_dag_root_changes_when_answer_changes():
    a = build_run_dag(**_kw())
    b = build_run_dag(**_kw(answer_text="Different answer."))
    assert a["root"] != b["root"]


def test_dag_root_changes_when_audit_mode_changes():
    a = build_run_dag(**_kw())
    b = build_run_dag(**_kw(audit_mode="UNGROUNDED"))
    assert a["root"] != b["root"]


def test_dag_root_changes_when_verifier_method_changes():
    a = build_run_dag(**_kw())
    b = build_run_dag(**_kw(verifier_method="paraphrase"))
    assert a["root"] != b["root"]


def test_dag_root_changes_when_question_hash_changes():
    a = build_run_dag(**_kw())
    b = build_run_dag(**_kw(question_hash="d" * 64))
    assert a["root"] != b["root"]


def test_dag_root_changes_when_sources_change():
    a = build_run_dag(**_kw())
    new_sources = [{"document_root": "e" * 64, "source_role": "primary_answer_source",
                    "score": 1.0, "chunk_idx": 0}]
    b = build_run_dag(**_kw(sources=new_sources))
    assert a["root"] != b["root"]


def test_dag_seven_stages_in_order():
    """The seven stages are emitted in fixed order so the root is
    canonical: question / retrieval / context / prompt / answer /
    verify / final_label."""
    out = build_run_dag(**_kw())
    stages = [n["stage"] for n in out["nodes"]]
    assert stages == [
        "question", "retrieval", "context", "prompt",
        "answer", "verify", "final_label",
    ]


def test_verify_dag_round_trip():
    """Recomputing the root from the persisted nodes must match."""
    out = build_run_dag(**_kw())
    assert verify_run_dag(out) is True


def test_verify_dag_detects_node_tampering():
    """Mutating a stage hash makes verify fail."""
    out = build_run_dag(**_kw())
    out["nodes"][3]["hash"] = "f" * 64  # tamper with prompt node
    # Root no longer matches the (mutated) leaves.
    assert verify_run_dag(out) is False


def test_localize_failure_strict_returns_none():
    """STRICT verdict has no failure to localize."""
    assert localize_failure(audit_mode="STRICT", n_sources=3, n_quotes=4, n_verified=4) is None


def test_localize_failure_no_sources_is_retrieval():
    """No admitted sources → retrieval-stage failure (gate over-rejected
    or corpus genuinely lacks the topic)."""
    assert localize_failure(
        audit_mode="UNGROUNDED", n_sources=0, n_quotes=0, n_verified=0
    ) == "retrieval"


def test_localize_failure_no_quotes_is_context():
    """Sources retrieved but no quotes extracted → context-stage failure
    (per-source cap dropped relevant content, or model declined to
    cite). Distinct from answer-stage failures where quotes existed but
    didn't verify."""
    assert localize_failure(
        audit_mode="UNGROUNDED", n_sources=5, n_quotes=0, n_verified=0
    ) == "context"


def test_localize_failure_unverified_quotes_is_answer():
    """Quotes extracted but not all verify → answer-stage failure
    (model fabricated, paraphrased inside quotes, or appended
    citations). The case mechanical_repair targets."""
    assert localize_failure(
        audit_mode="HYBRID", n_sources=3, n_quotes=4, n_verified=3
    ) == "answer"
    assert localize_failure(
        audit_mode="UNGROUNDED", n_sources=3, n_quotes=2, n_verified=0
    ) == "answer"


def test_localize_failure_lands_on_run_dag_verify_node():
    """failure_stage is folded into the verify node's payload so
    auditors reading the persisted DAG can see which stage produced
    a non-STRICT verdict."""
    out = build_run_dag(**_kw(
        audit_mode="HYBRID",
        n_quotes=4,
        n_verified=3,
    ))
    # The verify_node hash should differ between failure_stage='answer'
    # and a STRICT verdict — failure_stage enters the hashed payload.
    out_strict = build_run_dag(**_kw(
        audit_mode="STRICT",
        n_quotes=4,
        n_verified=4,
    ))
    out["root"] != out_strict["root"]  # at minimum, audit_mode differs


def test_verify_dag_accepts_json_string():
    """Sidecar/audit tools persist the DAG as JSON; verify accepts both."""
    import json
    out = build_run_dag(**_kw())
    blob = json.dumps(out, separators=(",", ":"))
    assert verify_run_dag(blob) is True


# ---------------------------------------------------------------- Ticket #000009: preflight stage

def test_preflight_node_hash_is_deterministic():
    """Same nested-clause inputs → same hex string, byte-for-byte."""
    qs = {"logical_statuses": ["well_formed"], "preflight_result": "PREFLIGHT_OK"}
    quant = {"intensity": "SINGULAR", "is_broad": False}
    answer_contract = {"guard_enabled": True, "apply_caps_active": False}
    a = preflight_node_hash(
        question_state=qs, quantifier=quant,
        answer_contract=answer_contract,
    )
    b = preflight_node_hash(
        question_state=qs, quantifier=quant,
        answer_contract=answer_contract,
    )
    assert a == b
    assert len(a) == 64  # SHA-256 hex


def test_preflight_node_hash_changes_with_question_state():
    base_quant = {"intensity": "SINGULAR", "is_broad": False}
    answer_contract = {"guard_enabled": True}
    h_a = preflight_node_hash(
        question_state={"preflight_result": "PREFLIGHT_OK"},
        quantifier=base_quant, answer_contract=answer_contract,
    )
    h_b = preflight_node_hash(
        question_state={"preflight_result": "PREFLIGHT_PARTIAL"},
        quantifier=base_quant, answer_contract=answer_contract,
    )
    assert h_a != h_b


def test_preflight_node_hash_changes_with_answer_contract():
    """Apply-caps flip in answer_contract MUST bump the node hash
    so audit replay can distinguish guard-on vs guard-off rows
    that otherwise share the same classifier output."""
    qs = {"logical_statuses": ["broad_quantifier_unbounded"]}
    quant = {"intensity": "ALL", "is_broad": True}
    h_off = preflight_node_hash(
        question_state=qs, quantifier=quant,
        answer_contract={"apply_caps_active": False},
    )
    h_on = preflight_node_hash(
        question_state=qs, quantifier=quant,
        answer_contract={"apply_caps_active": True},
    )
    assert h_off != h_on


def test_preflight_node_hash_changes_with_prompt_contract():
    """Reminder-injection flip MUST bump the node hash."""
    qs = {"logical_statuses": ["broad_quantifier_unbounded"]}
    quant = {"intensity": "ALL", "is_broad": True}
    h_off = preflight_node_hash(
        question_state=qs, quantifier=quant,
        prompt_contract={"reminder_injected": False},
    )
    h_on = preflight_node_hash(
        question_state=qs, quantifier=quant,
        prompt_contract={
            "reminder_injected": True,
            "reminder_template_id": "broad-quantifier-unbounded-v1",
        },
    )
    assert h_off != h_on


def test_preflight_node_hash_changes_with_policy_refs():
    """policy_refs.governance_policy_hash flip → distinct node
    hash. Required for cache-key/run-DAG-root joint invalidation."""
    qs = {"logical_statuses": ["well_formed"]}
    h_a = preflight_node_hash(
        question_state=qs,
        policy_refs={"governance_policy_hash": "a" * 64},
    )
    h_b = preflight_node_hash(
        question_state=qs,
        policy_refs={"governance_policy_hash": "b" * 64},
    )
    assert h_a != h_b


def test_preflight_node_hash_includes_node_version():
    """node_version field is part of the hashed payload so a future
    schema bump (preflight-node-v2 etc.) invalidates legacy nodes."""
    from arborist.qa.dag import (
        PREFLIGHT_NODE_VERSION,
        build_preflight_node_payload,
    )
    payload = build_preflight_node_payload(
        question_state={"x": 1},
    )
    assert payload["node_version"] == PREFLIGHT_NODE_VERSION
    assert PREFLIGHT_NODE_VERSION == "preflight-node-v1"


def test_preflight_node_hash_handles_all_none():
    """Defensive — all clauses may be None during gradual rollout.
    Hash stays stable."""
    a = preflight_node_hash()
    b = preflight_node_hash()
    assert a == b
    assert len(a) == 64


def test_dag_without_preflight_keeps_seven_stage_shape():
    """Backward-compat: omitting preflight_hash preserves the
    pre-#000009 7-stage shape so legacy run_dag_root values
    re-validate."""
    out = build_run_dag(**_kw())
    stages = [n["stage"] for n in out["nodes"]]
    assert stages == [
        "question", "retrieval", "context", "prompt",
        "answer", "verify", "final_label",
    ]
    assert len(out["nodes"]) == 7


def test_dag_with_preflight_inserts_eight_stage_shape():
    """Quote-mode + preflight_hash → 8 stages, preflight at
    position 1 (between question and retrieval)."""
    pre_hash = preflight_node_hash(
        question_state={"preflight_result": "PREFLIGHT_OK"},
        quantifier={"intensity": "SINGULAR", "is_broad": False},
        answer_contract={"guard_enabled": True},
    )
    out = build_run_dag(**_kw(preflight_hash=pre_hash))
    stages = [n["stage"] for n in out["nodes"]]
    assert stages == [
        "question", "preflight", "retrieval", "context",
        "prompt", "answer", "verify", "final_label",
    ]
    assert len(out["nodes"]) == 8


def test_dag_with_preflight_lattice_mode_ten_stages():
    """Pointer-mode + preflight_hash → 10 stages."""
    pre_hash = preflight_node_hash(
        question_state={"preflight_result": "PREFLIGHT_OK"},
        quantifier={"intensity": "ALL", "is_broad": True},
        answer_contract={"guard_enabled": True},
    )
    out = build_run_dag(**_kw(
        preflight_hash=pre_hash,
        evidence_map_root="d" * 64,
        verifier_method="claim_lattice_pointer",
        raw_answer_text="Some claim. [E1]",
        parsed_lattice=[{"claim_text": "Some claim", "evidence_ids": ["e1"]}],
        rendered_text="Some claim. [E1 | source]",
    ))
    stages = [n["stage"] for n in out["nodes"]]
    assert "preflight" in stages
    assert stages.index("preflight") == 1  # right after question
    assert len(stages) == 10


def test_dag_root_changes_when_preflight_hash_changes():
    """Different preflight inputs → different run_dag_root. This is
    the audit-replay payoff: same model output + same verifier
    verdict + DIFFERENT preflight policy = different cache row."""
    pre_a = preflight_node_hash(
        question_state={"preflight_result": "PREFLIGHT_OK"},
        quantifier={"intensity": "SINGULAR"},
        answer_contract={"apply_caps_active": False},
    )
    pre_b = preflight_node_hash(
        question_state={"preflight_result": "PREFLIGHT_OK"},
        quantifier={"intensity": "SINGULAR"},
        answer_contract={"apply_caps_active": True},
    )
    a = build_run_dag(**_kw(preflight_hash=pre_a))
    b = build_run_dag(**_kw(preflight_hash=pre_b))
    assert a["root"] != b["root"]


def test_dag_with_preflight_round_trips_through_verify():
    """The preflight stage's hash is part of the leaf list so
    verify_run_dag must reconstruct the same root."""
    import json
    pre_hash = preflight_node_hash(
        question_state={"preflight_result": "PREFLIGHT_PARTIAL"},
        quantifier={"intensity": "ALL", "is_broad": True},
        answer_contract={"apply_caps_active": True, "claim_cap_resolved": 8},
    )
    out = build_run_dag(**_kw(preflight_hash=pre_hash))
    blob = json.dumps(out, separators=(",", ":"))
    assert verify_run_dag(blob) is True


# ---------------------------------------------------------------- reject-broad DAG (§8.2 A)

def test_reject_run_dag_three_stage_shape():
    """Reject-broad early-return path produces a 3-stage DAG:
    question → preflight → final_label. Audit replay can read the
    stage list and tell instantly that this row is a preflight
    rejection (3 stages = reject path)."""
    from arborist.qa.dag import build_reject_run_dag, preflight_node_hash
    pre = preflight_node_hash(
        question_state={"preflight_result": "PREFLIGHT_BLOCKED"},
        quantifier={"intensity": "ALL", "is_broad": True,
                    "scope_bound_hint": "unbounded"},
        answer_contract={"reject_broad_active": True},
    )
    out = build_reject_run_dag(
        question_hash="a" * 64,
        preflight_hash=pre,
        rejection_reason="preflight rejection — broad-unbounded.",
        answer_text="BROAD-QUANTIFIER PREFLIGHT REJECTED",
        violations=[{"kind": "BROAD_QUANTIFIER_REJECTED"}],
    )
    stages = [n["stage"] for n in out["nodes"]]
    assert stages == ["question", "preflight", "final_label"]
    assert len(stages) == 3


def test_reject_run_dag_root_changes_with_preflight_hash():
    """Two reject runs that differ only in the preflight payload
    (e.g. different policy state at rejection time) must produce
    different run_dag_root values."""
    from arborist.qa.dag import build_reject_run_dag, preflight_node_hash
    pre_a = preflight_node_hash(
        question_state={"preflight_result": "PREFLIGHT_BLOCKED"},
        quantifier={"intensity": "ALL"},
        policy_refs={"governance_policy_hash": "a" * 64},
    )
    pre_b = preflight_node_hash(
        question_state={"preflight_result": "PREFLIGHT_BLOCKED"},
        quantifier={"intensity": "ALL"},
        policy_refs={"governance_policy_hash": "b" * 64},
    )
    a = build_reject_run_dag(
        question_hash="d" * 64, preflight_hash=pre_a,
        rejection_reason="r", answer_text="x",
    )
    b = build_reject_run_dag(
        question_hash="d" * 64, preflight_hash=pre_b,
        rejection_reason="r", answer_text="x",
    )
    assert a["root"] != b["root"]


def test_reject_run_dag_round_trips_through_verify():
    """3-stage reject DAG must verify the same way as the standard
    7/9/8/10-stage shapes."""
    import json
    from arborist.qa.dag import build_reject_run_dag, preflight_node_hash
    pre = preflight_node_hash(
        question_state={"preflight_result": "PREFLIGHT_BLOCKED"},
        quantifier={"intensity": "ALL"},
    )
    out = build_reject_run_dag(
        question_hash="a" * 64,
        preflight_hash=pre,
        rejection_reason="preflight rejection",
        answer_text="REJECTED",
    )
    blob = json.dumps(out, separators=(",", ":"))
    assert verify_run_dag(blob) is True


# ---------------------------------------------------------------- preflight extraction (#000009 §7.2)

def test_extract_preflight_hash_from_blob_with_preflight():
    """Pull preflight stage hash out of a persisted run_dag_blob."""
    import json
    from arborist.qa.dag import build_run_dag, preflight_node_hash
    from arborist.qa.query import _extract_preflight_hash_from_blob
    pre = preflight_node_hash(
        question_state={"preflight_result": "PREFLIGHT_OK"},
        quantifier={"intensity": "SINGULAR"},
    )
    out = build_run_dag(**_kw(preflight_hash=pre))
    blob = json.dumps(out, separators=(",", ":"))
    extracted = _extract_preflight_hash_from_blob(blob)
    assert extracted == pre


def test_extract_preflight_hash_returns_none_for_legacy_blob():
    """Legacy blobs (no preflight stage) return None — this is the
    fall-through path for cache rows written before #000009."""
    import json
    from arborist.qa.query import _extract_preflight_hash_from_blob
    out = build_run_dag(**_kw())  # no preflight_hash
    blob = json.dumps(out, separators=(",", ":"))
    assert _extract_preflight_hash_from_blob(blob) is None


def test_extract_preflight_hash_returns_none_for_empty_or_invalid():
    from arborist.qa.query import _extract_preflight_hash_from_blob
    assert _extract_preflight_hash_from_blob(None) is None
    assert _extract_preflight_hash_from_blob("") is None
    assert _extract_preflight_hash_from_blob("not json {{{") is None
    assert _extract_preflight_hash_from_blob('"just a string"') is None
    assert _extract_preflight_hash_from_blob("[]") is None  # not a dict
    assert _extract_preflight_hash_from_blob('{"nodes": []}') is None  # no preflight