neuron-tai/tests/test_hop_bisection.py

"""AH-007: on-demand hop-boundary commitments + first-divergent-hop bisection.

`_final_text_node` (blame `max(shard_end)`) is wrong for multi-hop pipelines:
a corrupt early hop still produces a route whose *last* hop reported perfectly
honest activations. These tests build a two-hop route, corrupt one hop's
committed fingerprint, and assert the referee blames that hop specifically —
not always the last one.
"""

from __future__ import annotations

import http.server
import json
import socketserver
import threading
from types import SimpleNamespace

from meshnet_validator import ToplocAuditConfig, ValidatorProcess
from meshnet_validator.audit import build_activation_proofs


class FakeToploc:
    """Fingerprint = the activations themselves; verify is exact equality,
    matching the encode/verify contract exercised in test_toploc_audit.py."""

    def build_proofs_base64(self, activations, *, decode_batching_size, topk, skip_prefill):
        return {
            "activation_fingerprint": tuple(tuple(row) for row in activations),
            "decode_batching_size": decode_batching_size,
            "topk": topk,
            "skip_prefill": skip_prefill,
        }

    def verify_proofs_base64(self, activations, proofs, *, decode_batching_size, topk, skip_prefill):
        fingerprint = proofs.get("activation_fingerprint") if isinstance(proofs, dict) else None
        return proofs == {
            "activation_fingerprint": tuple(tuple(row) for row in activations),
            "decode_batching_size": decode_batching_size,
            "topk": topk,
            "skip_prefill": skip_prefill,
        } or (
            fingerprint is not None
            and tuple(tuple(row) for row in fingerprint) == tuple(tuple(row) for row in activations)
            and proofs.get("decode_batching_size") == decode_batching_size
            and proofs.get("topk") == topk
            and proofs.get("skip_prefill") == skip_prefill
        )


class FakeValidationLog:
    def __init__(self, events) -> None:
        self._events = events

    def list_completed_inferences(self, *, after_index: int = -1):
        return [event for event in self._events if event.index > after_index]


class FakeRegistry:
    def __init__(self) -> None:
        self.slashes: list[dict] = []
        self.audit_outcomes: list[dict] = []

    def get_wallet(self, wallet_address: str):
        return SimpleNamespace(banned=False)

    def submit_slash_proof(self, **kwargs):
        self.slashes.append(kwargs)
        return kwargs

    def record_audit_outcome(self, wallet_address: str, *, passed: bool):
        self.audit_outcomes.append({"wallet_address": wallet_address, "passed": passed})


class FakeContracts:
    def __init__(self, events) -> None:
        self.validation = FakeValidationLog(events)
        self.registry = FakeRegistry()


class CommitmentStubServer:
    def __init__(self, commitments: dict[str, dict]) -> None:
        self.requests: list[dict] = []
        outer = self

        class Handler(http.server.BaseHTTPRequestHandler):
            def log_message(self, fmt, *args):
                pass

            def do_POST(self):
                if self.path != "/v1/audit/toploc/commitment":
                    self.send_response(404)
                    self.end_headers()
                    return
                length = int(self.headers.get("Content-Length", 0))
                body = json.loads(self.rfile.read(length) or b"{}")
                outer.requests.append(body)
                payload = json.dumps(commitments[body["session_id"]]).encode()
                self.send_response(200)
                self.send_header("Content-Type", "application/json")
                self.send_header("Content-Length", str(len(payload)))
                self.end_headers()
                self.wfile.write(payload)

        self._server = socketserver.ThreadingTCPServer(("127.0.0.1", 0), Handler)
        self._server.daemon_threads = True
        self._thread: threading.Thread | None = None

    def start(self) -> str:
        self._thread = threading.Thread(target=self._server.serve_forever, daemon=True)
        self._thread.start()
        return f"http://127.0.0.1:{self._server.server_address[1]}"

    def stop(self) -> None:
        self._server.shutdown()
        self._server.server_close()


def _hop_route_event(*, index, hop0_claim_activations, hop1_claim_activations, config, backend, ts=None):
    """Two-hop route where each hop carries its own TOPLOC commitment built
    from the given (possibly corrupted) activations."""
    hop0_claim = build_activation_proofs(hop0_claim_activations, config=config, backend=backend)
    hop1_claim = build_activation_proofs(hop1_claim_activations, config=config, backend=backend)
    route_nodes = [
        {
            "wallet_address": "wallet-hop0",
            "shard_start": 0,
            "shard_end": 15,
            "toploc_proof": hop0_claim.as_mapping(),
        },
        {
            "wallet_address": "wallet-hop1",
            "shard_start": 16,
            "shard_end": 31,
            "toploc_proof": hop1_claim.as_mapping(),
        },
    ]
    event = SimpleNamespace(
        index=index,
        session_id=f"session-hop-{index}",
        model="Qwen2.5-0.5B-Instruct",
        messages=[{"role": "user", "content": "hello"}],
        observed_output="two-hop pipeline output",
        route_nodes=route_nodes,
        claimed_token_ids=[101, 202, 303],
    )
    if ts is not None:
        event.ts = ts
    return event


class HopReferenceValidator(ValidatorProcess):
    """Stands in for the reference node: returns canned ground-truth
    activations at each hop cut-point instead of making an HTTP call."""

    def __init__(self, *args, reference_activations_by_hop, **kwargs) -> None:
        super().__init__(*args, **kwargs)
        self.reference_activations_by_hop = reference_activations_by_hop
        self.hops_calls: list[dict] = []
        self.text_reference_calls: list[list[dict]] = []

    def _run_teacher_forced_prefill_hops(self, **kwargs):
        self.hops_calls.append(kwargs)
        return self.reference_activations_by_hop

    def _run_reference(self, messages: list[dict]) -> str:
        self.text_reference_calls.append(messages)
        return "two-hop pipeline output"


def test_bisection_blames_first_divergent_hop_not_last_hop():
    """Red: corrupt hop-0 only. The old `_final_text_node` bug blames
    max(shard_end) == hop-1 (wallet-hop1), which is innocent here."""
    config = ToplocAuditConfig(topk=2, decode_batching_size=16)
    backend = FakeToploc()
    reference_hop0 = [[1.0, 2.0], [3.0, 4.0]]
    reference_hop1 = [[5.0, 6.0], [7.0, 8.0]]
    corrupted_hop0_claim_activations = [[9.0, 9.0], [9.0, 9.0]]  # diverges from reference_hop0

    event = _hop_route_event(
        index=0,
        hop0_claim_activations=corrupted_hop0_claim_activations,
        hop1_claim_activations=reference_hop1,
        config=config,
        backend=backend,
    )
    contracts = FakeContracts([event])
    validator = HopReferenceValidator(
        contracts=contracts,
        reference_node_url="http://reference.invalid",
        sample_rate=1.0,
        random_seed=7,
        toploc_config=config,
        toploc_backend=backend,
        reference_activations_by_hop=[reference_hop0, reference_hop1],
    )

    receipts = validator.validate_once()

    assert len(receipts) == 1
    assert contracts.registry.slashes[0]["wallet_address"] == "wallet-hop0"
    assert "hop 0" in contracts.registry.slashes[0]["reason"]


def test_bisection_blames_the_actual_faulty_hop_when_it_is_the_second_hop():
    """Integration test: multi-hop pipeline, fault injected at a known
    (non-first) hop — proves blame follows the real culprit, not a fixed index."""
    config = ToplocAuditConfig(topk=2, decode_batching_size=16)
    backend = FakeToploc()
    reference_hop0 = [[1.0, 2.0], [3.0, 4.0]]
    reference_hop1 = [[5.0, 6.0], [7.0, 8.0]]
    corrupted_hop1_claim_activations = [[0.0, 0.0], [0.0, 0.0]]

    event = _hop_route_event(
        index=0,
        hop0_claim_activations=reference_hop0,
        hop1_claim_activations=corrupted_hop1_claim_activations,
        config=config,
        backend=backend,
    )
    contracts = FakeContracts([event])
    validator = HopReferenceValidator(
        contracts=contracts,
        reference_node_url="http://reference.invalid",
        sample_rate=1.0,
        random_seed=7,
        toploc_config=config,
        toploc_backend=backend,
        reference_activations_by_hop=[reference_hop0, reference_hop1],
    )

    receipts = validator.validate_once()

    assert len(receipts) == 1
    assert contracts.registry.slashes[0]["wallet_address"] == "wallet-hop1"


def test_honest_two_hop_route_is_not_slashed():
    config = ToplocAuditConfig(topk=2, decode_batching_size=16)
    backend = FakeToploc()
    reference_hop0 = [[1.0, 2.0], [3.0, 4.0]]
    reference_hop1 = [[5.0, 6.0], [7.0, 8.0]]

    event = _hop_route_event(
        index=0,
        hop0_claim_activations=reference_hop0,
        hop1_claim_activations=reference_hop1,
        config=config,
        backend=backend,
    )
    contracts = FakeContracts([event])
    validator = HopReferenceValidator(
        contracts=contracts,
        reference_node_url="http://reference.invalid",
        sample_rate=1.0,
        random_seed=7,
        toploc_config=config,
        toploc_backend=backend,
        reference_activations_by_hop=[reference_hop0, reference_hop1],
    )

    assert validator.validate_once() == []
    assert contracts.registry.slashes == []


def test_expired_commitment_window_falls_back_to_text_only_audit():
    """ADR-0018 §3: nodes only retain boundary activations briefly. Once the
    on-demand TTL has passed, bisection can't be verified — the validator must
    fall back to the text-comparison path instead of erroring out."""
    config = ToplocAuditConfig(topk=2, decode_batching_size=16, commitment_ttl_seconds=1.0)
    backend = FakeToploc()
    reference_hop0 = [[1.0, 2.0], [3.0, 4.0]]
    reference_hop1 = [[5.0, 6.0], [7.0, 8.0]]

    event = _hop_route_event(
        index=0,
        hop0_claim_activations=reference_hop0,
        hop1_claim_activations=reference_hop1,
        config=config,
        backend=backend,
        ts=0.0,  # far outside any TTL relative to real wall-clock time
    )
    contracts = FakeContracts([event])
    validator = HopReferenceValidator(
        contracts=contracts,
        reference_node_url="http://reference.invalid",
        sample_rate=1.0,
        random_seed=7,
        toploc_config=config,
        toploc_backend=backend,
        reference_activations_by_hop=[reference_hop0, reference_hop1],
    )

    assert validator.validate_once() == []
    assert validator.hops_calls == []
    assert len(validator.text_reference_calls) == 1


def test_hop_commitments_are_not_requested_unless_the_event_is_audit_selected():
    """On-demand: the (expensive) per-hop commitment retrieval only happens
    for events the tracker RNG actually selects for audit — sample_rate is
    the selection gate, and a miss must not touch the reference node at all."""
    config = ToplocAuditConfig(topk=2, decode_batching_size=16)
    backend = FakeToploc()
    reference_hop0 = [[1.0, 2.0], [3.0, 4.0]]
    reference_hop1 = [[5.0, 6.0], [7.0, 8.0]]

    event = _hop_route_event(
        index=0,
        hop0_claim_activations=reference_hop0,
        hop1_claim_activations=reference_hop1,
        config=config,
        backend=backend,
    )
    contracts = FakeContracts([event])
    validator = HopReferenceValidator(
        contracts=contracts,
        reference_node_url="http://reference.invalid",
        sample_rate=0.0,  # never selects any event for audit
        random_seed=7,
        toploc_config=config,
        toploc_backend=backend,
        reference_activations_by_hop=[reference_hop0, reference_hop1],
    )

    assert validator.validate_once() == []
    assert validator.hops_calls == []
    assert validator.text_reference_calls == []
    assert validator.sampled_count == 0

def test_validator_fetches_missing_hop_commitments_on_demand_after_sampling():
    """Production route events only carry hop metadata. Once sampled, the
    validator asks each route node for its retained boundary commitment window
    and then runs the same first-divergent-hop bisection path."""
    config = ToplocAuditConfig(topk=2, decode_batching_size=16)
    backend = FakeToploc()
    reference_hop0 = [[1.0, 2.0], [3.0, 4.0]]
    reference_hop1 = [[5.0, 6.0], [7.0, 8.0]]
    hop0_claim = build_activation_proofs(reference_hop0, config=config, backend=backend)
    hop1_claim = build_activation_proofs(reference_hop1, config=config, backend=backend)

    stub = CommitmentStubServer({
        "session-on-demand": {
            "toploc_proof": hop0_claim.as_mapping(),
            "claimed_token_ids": [101, 202, 303],
        }
    })
    endpoint0 = stub.start()
    try:
        event = SimpleNamespace(
            index=0,
            session_id="session-on-demand",
            model="Qwen2.5-0.5B-Instruct",
            messages=[{"role": "user", "content": "hello"}],
            observed_output="two-hop pipeline output",
            route_nodes=[
                {"wallet_address": "wallet-hop0", "endpoint": endpoint0, "shard_start": 0, "shard_end": 15},
                {
                    "wallet_address": "wallet-hop1",
                    "endpoint": "http://127.0.0.1:1",
                    "shard_start": 16,
                    "shard_end": 31,
                    "toploc_proof": hop1_claim.as_mapping(),
                    "claimed_token_ids": [101, 202, 303],
                },
            ],
        )
        contracts = FakeContracts([event])
        validator = HopReferenceValidator(
            contracts=contracts,
            reference_node_url="http://reference.invalid",
            sample_rate=1.0,
            random_seed=7,
            toploc_config=config,
            toploc_backend=backend,
            reference_activations_by_hop=[reference_hop0, reference_hop1],
        )

        assert validator.validate_once() == []
        assert stub.requests == [{
            "session_id": "session-on-demand",
            "model": "Qwen2.5-0.5B-Instruct",
            "messages": [{"role": "user", "content": "hello"}],
            "shard_start": 0,
            "shard_end": 15,
        }]
        assert validator.hops_calls
    finally:
        stub.stop()