neuron-tai/tests/test_toploc_audit.py

"""AH-006: validator TOPLOC audit primitive."""

from __future__ import annotations

from types import SimpleNamespace

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


class FakeToploc:
    def __init__(self) -> None:
        self.build_calls: list[dict] = []
        self.verify_calls: list[dict] = []

    def build_proofs_base64(
        self,
        activations,
        *,
        decode_batching_size: int,
        topk: int,
        skip_prefill: bool,
    ):
        self.build_calls.append({
            "decode_batching_size": decode_batching_size,
            "topk": topk,
            "skip_prefill": 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: int,
        topk: int,
        skip_prefill: bool,
    ):
        self.verify_calls.append({
            "decode_batching_size": decode_batching_size,
            "topk": topk,
            "skip_prefill": skip_prefill,
        })
        return proofs == {
            "activation_fingerprint": tuple(tuple(row) for row in activations),
            "decode_batching_size": decode_batching_size,
            "topk": topk,
            "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 TeacherForcedValidator(ValidatorProcess):
    def __init__(self, *args, reference_activations, **kwargs) -> None:
        super().__init__(*args, **kwargs)
        self.reference_activations = reference_activations
        self.teacher_forced_calls: list[dict] = []

    def _run_reference(self, messages: list[dict]) -> str:
        raise AssertionError("TOPLOC audits must not free-generate reference text")

    def _run_teacher_forced_prefill(self, **kwargs):
        self.teacher_forced_calls.append(kwargs)
        return self.reference_activations


def test_stub_activation_tensors_round_trip_through_toploc_proofs():
    fake_toploc = FakeToploc()
    activations = [[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]
    config = ToplocAuditConfig(topk=2, decode_batching_size=16)

    claim = build_activation_proofs(activations, config=config, backend=fake_toploc)

    assert verify_activation_proofs(activations, claim, config=config, backend=fake_toploc) is True
    assert fake_toploc.build_calls == [{
        "decode_batching_size": 16,
        "topk": 2,
        "skip_prefill": True,
    }]
    assert fake_toploc.verify_calls == [{
        "decode_batching_size": 16,
        "topk": 2,
        "skip_prefill": True,
    }]


def test_validator_teacher_forces_claimed_tokens_for_toploc_audit():
    fake_toploc = FakeToploc()
    activations = [[0.25, 0.5], [0.75, 1.0]]
    config = ToplocAuditConfig(topk=2, decode_batching_size=16)
    claim = build_activation_proofs(activations, config=config, backend=fake_toploc)
    event = SimpleNamespace(
        index=0,
        session_id="session-toploc-ok",
        model="Qwen2.5-0.5B-Instruct",
        messages=[{"role": "user", "content": "hello"}],
        observed_output="honest but hardware-divergent text",
        route_nodes=[{"wallet_address": "wallet-good", "shard_end": 31}],
        claimed_token_ids=[101, 202, 303],
        toploc_proof=claim.as_mapping(),
    )
    contracts = FakeContracts([event])
    validator = TeacherForcedValidator(
        contracts=contracts,
        reference_node_url="http://reference.invalid",
        sample_rate=1.0,
        random_seed=7,
        toploc_config=config,
        toploc_backend=fake_toploc,
        reference_activations=activations,
    )

    assert validator.validate_once() == []
    assert contracts.registry.slashes == []
    assert validator.teacher_forced_calls == [{
        "model": "Qwen2.5-0.5B-Instruct",
        "messages": [{"role": "user", "content": "hello"}],
        "claimed_token_ids": [101, 202, 303],
        "claim": claim,
    }]


def test_validator_rejects_swapped_precision_toploc_claim():
    fake_toploc = FakeToploc()
    activations = [[0.25, 0.5], [0.75, 1.0]]
    canonical = ToplocAuditConfig(
        dtype="bfloat16",
        quantization="bfloat16",
        topk=2,
        decode_batching_size=16,
    )
    swapped = ToplocAuditConfig(
        dtype="bfloat16",
        quantization="int8",
        topk=2,
        decode_batching_size=16,
    )
    claim = build_activation_proofs(activations, config=swapped, backend=fake_toploc)
    event = SimpleNamespace(
        index=0,
        session_id="session-toploc-bad",
        model="Qwen2.5-0.5B-Instruct",
        messages=[{"role": "user", "content": "hello"}],
        observed_output="looks plausible",
        route_nodes=[{"wallet_address": "wallet-bad", "shard_end": 31}],
        audit={
            "claimed_token_ids": [101, 202, 303],
            "toploc_proof": claim.as_mapping(),
        },
    )
    contracts = FakeContracts([event])
    validator = TeacherForcedValidator(
        contracts=contracts,
        reference_node_url="http://reference.invalid",
        sample_rate=1.0,
        random_seed=7,
        toploc_config=canonical,
        toploc_backend=fake_toploc,
        reference_activations=activations,
    )

    receipts = validator.validate_once()

    assert len(receipts) == 1
    assert contracts.registry.slashes[0]["wallet_address"] == "wallet-bad"
    assert "TOPLOC activation proof mismatch" in contracts.registry.slashes[0]["reason"]