neuron-tai/tests/test_billing_ledger.py

"""US-031: billing ledger — per-token pricing, 90/10 split, pending balances.

Unit tests for BillingLedger math/persistence/replication, plus HTTP
integration: 401 without an API key, billed 200 with one, 402 once the
Caller Credit is exhausted (rejected before routing — no free work).
"""

import http.server
import json
import socketserver
import threading
import urllib.error
import urllib.request

import pytest

from meshnet_tracker.auth import sign_hive_request
from meshnet_tracker.billing import BillingLedger
from meshnet_tracker.server import TrackerServer

MODEL = "openai-community/gpt2"
HIVE_SECRET = "test-hive-secret"


# ---------------------------------------------------------------- unit tests


def test_charge_single_node_gets_90_percent():
    ledger = BillingLedger(starting_credit=1.0, default_price_per_1k=0.02)
    ledger.ensure_client("key-a")
    event = ledger.charge_request("key-a", MODEL, total_tokens=1000, node_work=[("wallet-1", 12)])
    assert event["cost"] == pytest.approx(0.02)
    assert ledger.get_client_balance("key-a") == pytest.approx(1.0 - 0.02)
    assert ledger.get_node_pending("wallet-1") == pytest.approx(0.02 * 0.90)
    assert ledger.snapshot()["protocol_cut"] == pytest.approx(0.02 * 0.10)


def test_charge_three_node_split_by_work_units():
    ledger = BillingLedger(starting_credit=1.0, default_price_per_1k=0.02)
    ledger.ensure_client("key-a")
    ledger.charge_request(
        "key-a", MODEL, total_tokens=1000,
        node_work=[("wallet-1", 6), ("wallet-2", 3), ("wallet-3", 3)],
    )
    pool = 0.02 * 0.90
    assert ledger.get_node_pending("wallet-1") == pytest.approx(pool * 6 / 12)
    assert ledger.get_node_pending("wallet-2") == pytest.approx(pool * 3 / 12)
    assert ledger.get_node_pending("wallet-3") == pytest.approx(pool * 3 / 12)
    assert ledger.snapshot()["protocol_cut"] == pytest.approx(0.02 * 0.10)


def test_walletless_node_share_accrues_to_protocol_cut():
    ledger = BillingLedger(starting_credit=1.0, default_price_per_1k=0.02)
    ledger.ensure_client("key-a")
    ledger.charge_request(
        "key-a", MODEL, total_tokens=1000,
        node_work=[("wallet-1", 6), (None, 6)],
    )
    pool = 0.02 * 0.90
    assert ledger.get_node_pending("wallet-1") == pytest.approx(pool / 2)
    # walletless half of the pool + the 10% cut both land in protocol_cut
    assert ledger.snapshot()["protocol_cut"] == pytest.approx(pool / 2 + 0.02 * 0.10)


def test_per_model_price_override():
    ledger = BillingLedger(
        starting_credit=1.0,
        default_price_per_1k=0.02,
        prices={MODEL: 0.10},
    )
    ledger.ensure_client("key-a")
    event = ledger.charge_request("key-a", MODEL, 500, [("wallet-1", 1)])
    assert event["cost"] == pytest.approx(0.10 * 500 / 1000)
    event = ledger.charge_request("key-a", "other-model", 500, [("wallet-1", 1)])
    assert event["cost"] == pytest.approx(0.02 * 500 / 1000)


def test_payout_and_forfeit_hooks():
    ledger = BillingLedger(starting_credit=1.0, default_price_per_1k=0.02)
    ledger.charge_request("key-a", MODEL, 1000, [("wallet-1", 1)])
    pending = ledger.get_node_pending("wallet-1")
    ledger.settle_node_payout("wallet-1", pending, reference="tx-sig-1")
    assert ledger.get_node_pending("wallet-1") == pytest.approx(0.0)

    ledger.charge_request("key-a", MODEL, 1000, [("wallet-1", 1)])
    cut_before = ledger.snapshot()["protocol_cut"]
    forfeited = ledger.forfeit_pending("wallet-1")["amount"]
    assert forfeited == pytest.approx(0.02 * 0.90)
    assert ledger.get_node_pending("wallet-1") == pytest.approx(0.0)
    assert ledger.snapshot()["protocol_cut"] == pytest.approx(cut_before + forfeited)


def test_restart_persistence(tmp_path):
    db = str(tmp_path / "billing.db")
    ledger = BillingLedger(db_path=db, starting_credit=1.0, default_price_per_1k=0.02)
    ledger.credit_client("key-a", 5.0)
    ledger.charge_request("key-a", MODEL, 1000, [("wallet-1", 12)])
    ledger.save_to_db()

    reloaded = BillingLedger(db_path=db, starting_credit=1.0, default_price_per_1k=0.02)
    assert reloaded.get_client_balance("key-a") == pytest.approx(
        ledger.get_client_balance("key-a")
    )
    assert reloaded.get_node_pending("wallet-1") == pytest.approx(0.02 * 0.90)
    assert reloaded.snapshot()["protocol_cut"] == pytest.approx(0.02 * 0.10)


def test_tracker_enables_billing_with_default_db_when_requested(tmp_path, monkeypatch):
    from meshnet_tracker.billing import DEFAULT_BILLING_DB_PATH

    monkeypatch.chdir(tmp_path)
    tracker = TrackerServer(enable_billing=True)
    port = tracker.start()
    try:
        # /v1/billing/summary is admin-gated now; just confirm the server is up.
        health = json.loads(
            urllib.request.urlopen(f"http://127.0.0.1:{port}/v1/health").read()
        )
        assert health["status"] == "ok"
    finally:
        tracker.stop()
    # enabling billing creates the ledger DB in the working directory
    assert (tmp_path / DEFAULT_BILLING_DB_PATH).exists()


def test_event_replication_converges_and_dedupes():
    leader = BillingLedger(starting_credit=1.0, default_price_per_1k=0.02)
    follower = BillingLedger(starting_credit=1.0, default_price_per_1k=0.02)

    leader.credit_client("key-a", 10.0)
    leader.charge_request("key-a", MODEL, 2000, [("wallet-1", 8), ("wallet-2", 4)])

    events, cursor = leader.events_since(0)
    assert follower.apply_events(events) == len(events)
    # replaying the same batch must be a no-op
    assert follower.apply_events(events) == 0

    assert follower.get_client_balance("key-a") == pytest.approx(
        leader.get_client_balance("key-a")
    )
    assert follower.get_node_pending("wallet-1") == pytest.approx(
        leader.get_node_pending("wallet-1")
    )
    assert follower.snapshot()["protocol_cut"] == pytest.approx(
        leader.snapshot()["protocol_cut"]
    )
    # incremental cursor: nothing new after full sync
    more, _ = leader.events_since(cursor)
    assert more == []


# ---------------------------------------------------------- HTTP integration


class _UsageStubNode:
    """Minimal head node returning a chat completion with real usage numbers."""

    def __init__(self, total_tokens: int = 1000):
        self.total_tokens = total_tokens
        outer = self

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

            def do_POST(self):
                length = int(self.headers.get("Content-Length", 0))
                self.rfile.read(length)
                body = json.dumps({
                    "id": "chatcmpl-stub",
                    "object": "chat.completion",
                    "model": MODEL,
                    "choices": [{
                        "index": 0,
                        "message": {"role": "assistant", "content": "ok"},
                        "finish_reason": "stop",
                    }],
                    "usage": {
                        "prompt_tokens": 0,
                        "completion_tokens": outer.total_tokens,
                        "total_tokens": outer.total_tokens,
                    },
                }).encode()
                self.send_response(200)
                self.send_header("Content-Type", "application/json")
                self.send_header("Content-Length", str(len(body)))
                self.end_headers()
                self.wfile.write(body)

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

    def start(self) -> int:
        self._thread = threading.Thread(target=self._server.serve_forever, daemon=True)
        self._thread.start()
        return self._server.server_address[1]

    def stop(self):
        self._server.shutdown()
        self._server.server_close()


@pytest.fixture
def billed_tracker():
    ledger = BillingLedger(starting_credit=0.03, default_price_per_1k=0.02)
    tracker = TrackerServer(
        model_presets={
            MODEL: {
                "layers_start": 0,
                "layers_end": 11,
                "bytes_per_layer": {"bfloat16": 30 * 1024 * 1024},
            }
        },
        billing=ledger,
        hive_secret=HIVE_SECRET,
    )
    port = tracker.start()
    tracker_url = f"http://127.0.0.1:{port}"

    stub = _UsageStubNode(total_tokens=1000)
    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-head",
    }).encode()
    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()

    yield tracker_url, ledger

    stub.stop()
    tracker.stop()


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


def test_proxy_chat_requires_api_key_when_billing_enabled(billed_tracker):
    tracker_url, _ = billed_tracker
    with pytest.raises(urllib.error.HTTPError) as exc_info:
        _chat(tracker_url, api_key=None)
    assert exc_info.value.code == 401


def test_proxy_chat_bills_client_and_credits_node(billed_tracker):
    tracker_url, ledger = billed_tracker
    _chat(tracker_url, api_key="client-1")
    # 1000 tokens at 0.02/1K = 0.02 USDT; starting credit 0.03
    assert ledger.get_client_balance("client-1") == pytest.approx(0.03 - 0.02)
    assert ledger.get_node_pending("wallet-head") == pytest.approx(0.02 * 0.90)

    # /v1/billing/summary is admin-gated now (ADR-0017); auth is covered in
    # test_auth_boundary.py, so verify the numbers via the ledger snapshot.
    summary = ledger.snapshot()
    assert summary["node_pending"]["wallet-head"] == pytest.approx(0.02 * 0.90)
    assert summary["protocol_cut"] == pytest.approx(0.02 * 0.10)

    stats = json.loads(urllib.request.urlopen(f"{tracker_url}/v1/stats").read())
    node_stats = next(iter(stats["nodes"].values()))["models"][MODEL]
    assert node_stats["tokens_per_sec_last_hour"] is not None
    assert node_stats["sample_count_last_hour"] == 1


def test_proxy_chat_402_when_balance_exhausted(billed_tracker):
    tracker_url, ledger = billed_tracker
    _chat(tracker_url, api_key="client-2")  # 0.03 -> 0.01
    _chat(tracker_url, api_key="client-2")  # 0.01 -> -0.01 (post-pay drift)
    assert ledger.get_client_balance("client-2") == pytest.approx(-0.01)
    with pytest.raises(urllib.error.HTTPError) as exc_info:
        _chat(tracker_url, api_key="client-2")
    assert exc_info.value.code == 402
    # rejected before routing: nothing further was billed
    assert ledger.get_client_balance("client-2") == pytest.approx(-0.01)


def test_billing_gossip_endpoint_applies_events(billed_tracker):
    tracker_url, ledger = billed_tracker
    peer = BillingLedger(starting_credit=0.03, default_price_per_1k=0.02)
    peer.credit_client("remote-client", 7.0)
    events, _ = peer.events_since(0)
    body = json.dumps({"events": events}).encode()
    req = urllib.request.Request(
        f"{tracker_url}/v1/billing/gossip",
        data=body,
        headers={"Content-Type": "application/json", **sign_hive_request(HIVE_SECRET, body)},
        method="POST",
    )
    with urllib.request.urlopen(req) as r:
        assert json.loads(r.read())["applied"] == len(events)
    # only the replicated credit event lands here — Caller Credit was granted
    # on the peer that first saw the key, and its event replicates separately
    assert ledger.get_client_balance("remote-client") == pytest.approx(7.0)