neuron-tai/tests/test_forfeiture_penalty.py

"""US-034: hardened proof-of-work — pending-balance forfeiture penalty.

The pending balance is the fraud collateral (ADR-0015): a confirmed divergence
forfeits it in full to the protocol cut alongside the strike; three strikes ban
the wallet, whose remaining pending is never paid out.
"""

import json
import urllib.error
import urllib.request

import pytest

from meshnet_contracts import LocalSolanaContracts
from meshnet_node.server import StubNodeServer
from meshnet_tracker.billing import BillingLedger
from meshnet_tracker.server import TrackerServer
from meshnet_validator import ValidatorProcess

MODEL = "stub-model"


def _record_event(contracts, session_id: str, output: str, wallet: str) -> None:
    contracts.validation.record_completed_inference(
        session_id=session_id,
        model=MODEL,
        messages=[{"role": "user", "content": "2+2"}],
        observed_output=output,
        route_nodes=[{"wallet_address": wallet, "shard_end": 31}],
    )


@pytest.fixture
def reference_node():
    node = StubNodeServer(shard_start=0, shard_end=31, is_last_shard=True)
    port = node.start()
    yield f"http://127.0.0.1:{port}"
    node.stop()


def _reference_output(reference_url: str) -> str:
    """What the reference stub answers for the probe messages."""
    req = urllib.request.Request(
        f"{reference_url}/v1/infer",
        data=json.dumps({"messages": [{"role": "user", "content": "2+2"}]}).encode(),
        headers={"Content-Type": "application/json"},
        method="POST",
    )
    with urllib.request.urlopen(req) as r:
        return json.loads(r.read())["text"]


def test_divergence_forfeits_pending_and_strikes(reference_node):
    contracts = LocalSolanaContracts()
    contracts.registry.submit_stake("wallet-bad", 500)
    ledger = BillingLedger(starting_credit=10.0, default_price_per_1k=0.02)
    ledger.charge_request("client", MODEL, 1000, [("wallet-bad", 12)])
    pending_before = ledger.get_node_pending("wallet-bad")
    assert pending_before > 0

    validator = ValidatorProcess(
        contracts=contracts,
        billing=ledger,
        reference_node_url=reference_node,
        sample_rate=1.0,
        slash_amount=100,
        random_seed=7,
    )
    _record_event(contracts, "sess-1", "fraudulent nonsense", "wallet-bad")
    receipts = validator.validate_once()

    assert len(receipts) == 1
    assert ledger.get_node_pending("wallet-bad") == pytest.approx(0.0)
    snapshot = ledger.snapshot()
    # forfeited amount landed in the protocol cut on top of the original 10%
    assert snapshot["protocol_cut"] == pytest.approx(0.02 * 0.10 + pending_before)
    assert contracts.registry.get_wallet("wallet-bad").strike_count == 1


def test_matching_output_forfeits_nothing(reference_node):
    contracts = LocalSolanaContracts()
    contracts.registry.submit_stake("wallet-good", 500)
    ledger = BillingLedger(starting_credit=10.0, default_price_per_1k=0.02)
    ledger.charge_request("client", MODEL, 1000, [("wallet-good", 12)])
    pending_before = ledger.get_node_pending("wallet-good")

    validator = ValidatorProcess(
        contracts=contracts,
        billing=ledger,
        reference_node_url=reference_node,
        sample_rate=1.0,
        random_seed=7,
    )
    _record_event(contracts, "sess-1", _reference_output(reference_node), "wallet-good")
    receipts = validator.validate_once()

    assert receipts == []
    assert ledger.get_node_pending("wallet-good") == pytest.approx(pending_before)
    assert contracts.registry.get_wallet("wallet-good").strike_count == 0


def test_three_strikes_bans_and_bad_node_loses_everything(reference_node):
    """Deliberately-bad node: every job is fraudulent, checks always sample.

    Earn → caught → forfeit, three times over; the third strike bans the
    wallet, and the tracker rejects its registration.
    """
    contracts = LocalSolanaContracts()
    contracts.registry.submit_stake("wallet-bad", 500)
    ledger = BillingLedger(starting_credit=10.0, default_price_per_1k=0.02)
    validator = ValidatorProcess(
        contracts=contracts,
        billing=ledger,
        reference_node_url=reference_node,
        sample_rate=1.0,
        strike_threshold=3,
        random_seed=7,
    )

    for i in range(3):
        ledger.charge_request("client", MODEL, 1000, [("wallet-bad", 12)])
        _record_event(contracts, f"sess-{i}", f"fraud-{i}", "wallet-bad")
        validator.validate_once()

    wallet = contracts.registry.get_wallet("wallet-bad")
    assert wallet.strike_count == 3
    assert wallet.banned
    assert ledger.get_node_pending("wallet-bad") == pytest.approx(0.0)

    # banned wallet cannot register with the tracker
    tracker = TrackerServer(contracts=contracts)
    port = tracker.start()
    try:
        req = urllib.request.Request(
            f"http://127.0.0.1:{port}/v1/nodes/register",
            data=json.dumps({
                "endpoint": "http://127.0.0.1:9",
                "shard_start": 0,
                "shard_end": 31,
                "hardware_profile": {},
                "wallet_address": "wallet-bad",
                "score": 1.0,
            }).encode(),
            headers={"Content-Type": "application/json"},
            method="POST",
        )
        with pytest.raises(urllib.error.HTTPError) as exc_info:
            urllib.request.urlopen(req)
        assert exc_info.value.code == 403
    finally:
        tracker.stop()


def test_forfeit_endpoint_requires_auth_and_forfeits():
    contracts = LocalSolanaContracts()
    ledger = BillingLedger(starting_credit=10.0, default_price_per_1k=0.02)
    ledger.charge_request("client", MODEL, 1000, [("wallet-x", 12)])
    pending = ledger.get_node_pending("wallet-x")

    tracker = TrackerServer(contracts=contracts, billing=ledger)
    port = tracker.start()
    url = f"http://127.0.0.1:{port}/v1/billing/forfeit"
    body = json.dumps({"wallet": "wallet-x", "reason": "fraud"}).encode()
    try:
        req = urllib.request.Request(
            url, data=body, headers={"Content-Type": "application/json"}, method="POST",
        )
        with pytest.raises(urllib.error.HTTPError) as exc_info:
            urllib.request.urlopen(req)
        assert exc_info.value.code == 401

        req = urllib.request.Request(
            url, data=body,
            headers={"Content-Type": "application/json", "Authorization": "Bearer validator"},
            method="POST",
        )
        with urllib.request.urlopen(req) as r:
            result = json.loads(r.read())
        assert result["forfeited"] == pytest.approx(pending)
        assert result["strike_count"] == 1
        assert result["banned"] is False
        assert ledger.get_node_pending("wallet-x") == pytest.approx(0.0)
    finally:
        tracker.stop()


def test_probation_earns_nothing_then_earning_begins():
    """First N jobs accrue no pending balance; job N+1 earns (issue 08)."""
    contracts = LocalSolanaContracts(probationary_job_count=2)
    ledger = BillingLedger(starting_credit=10.0, default_price_per_1k=0.02)

    tracker = TrackerServer(
        model_presets={
            MODEL: {
                "layers_start": 0,
                "layers_end": 11,
                "bytes_per_layer": {"bfloat16": 30 * 1024 * 1024},
            }
        },
        contracts=contracts,
        billing=ledger,
    )
    port = tracker.start()
    tracker_url = f"http://127.0.0.1:{port}"

    stub = StubNodeServer(
        shard_start=0, shard_end=11, is_last_shard=True, model=MODEL, tracker_mode=True,
    )
    stub_port = stub.start()
    reg = json.dumps({
        "endpoint": f"http://127.0.0.1:{stub_port}",
        "shard_start": 0,
        "shard_end": 11,
        "model": MODEL,
        "hardware_profile": {},
        "score": 1.0,
        "tracker_mode": True,
        "wallet_address": "wallet-new",
    }).encode()

    def _chat():
        req = urllib.request.Request(
            f"{tracker_url}/v1/chat/completions",
            data=json.dumps({
                "model": MODEL,
                "messages": [{"role": "user", "content": "hi"}],
            }).encode(),
            headers={
                "Content-Type": "application/json",
                "Authorization": "Bearer probation-client",
            },
            method="POST",
        )
        with urllib.request.urlopen(req) as r:
            r.read()

    try:
        req = urllib.request.Request(
            f"{tracker_url}/v1/nodes/register",
            data=reg,
            headers={"Content-Type": "application/json"},
            method="POST",
        )
        with urllib.request.urlopen(req) as r:
            r.read()

        _chat()  # job 1: probation
        _chat()  # job 2: probation
        assert ledger.get_node_pending("wallet-new") == pytest.approx(0.0)
        assert contracts.registry.probationary_jobs_remaining("wallet-new") == 0

        _chat()  # job 3: earning begins (tokens are 0 in the stub, so the
        # charge event exists but the amount is 0 — assert via event log)
        events, _ = ledger.events_since(0)
        charges = [e for e in events if e["type"] == "charge"]
        assert len(charges) == 3
        assert charges[0]["shares"] == {} and charges[1]["shares"] == {}
        assert "wallet-new" in charges[2]["shares"]
    finally:
        stub.stop()
        tracker.stop()