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, 402 for unfunded keys, billed 200 after
explicit credit, and 402 once the balance is exhausted.
"""

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 DEFAULT_STARTING_CREDIT, BillingLedger
from meshnet_tracker.server import (
    TrackerServer,
    _billable_non_stream_tokens,
    _billable_stream_tokens,
    _estimate_prompt_tokens,
    _observed_non_stream_completion_tokens,
    _observed_stream_tokens,
)

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_default_starting_credit_is_zero_and_fresh_key_is_unfunded():
    ledger = BillingLedger(default_price_per_1k=0.02)

    assert DEFAULT_STARTING_CREDIT == 0.0
    assert ledger.ensure_client("fresh-key") == pytest.approx(0.0)
    assert ledger.get_client_balance("fresh-key") == pytest.approx(0.0)
    assert ledger.has_funds("fresh-key") is False
    assert ledger.events_since(0)[0] == []


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_non_stream_billable_tokens_cap_usage_by_request_bound():
    payload = {"usage": {"total_tokens": 1_000_000}}
    request = {"messages": [{"role": "user", "content": "hello tracker"}], "max_tokens": 3}

    assert _estimate_prompt_tokens(request) == 2
    assert _billable_non_stream_tokens(payload, request) == 5


def test_non_stream_billable_tokens_fallback_when_usage_missing():
    payload = {"choices": [{"message": {"role": "assistant", "content": "ok now"}}]}
    request = {"messages": [{"role": "user", "content": "hello tracker"}], "max_tokens": 10}

    assert _observed_non_stream_completion_tokens(payload) == 2
    assert _billable_non_stream_tokens(payload, request) == 4


def test_stream_billable_tokens_allow_usage_to_lower_observed_count_only():
    observed_payload = {
        "choices": [{
            "index": 0,
            "delta": {"content": "hello tracker"},
            "finish_reason": None,
        }]
    }

    assert _observed_stream_tokens(observed_payload) == 2
    assert _billable_stream_tokens(observed_tokens=2, reported_tokens=1_000_000) == 2
    assert _billable_stream_tokens(observed_tokens=2, reported_tokens=1) == 1


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,
        *,
        stream_chunks: list[str] | None = None,
        stream_usage_total: int | None = None,
    ):
        self.total_tokens = total_tokens
        self.stream_chunks = stream_chunks or []
        self.stream_usage_total = stream_usage_total
        self.request_count = 0
        outer = self

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

            def do_POST(self):
                outer.request_count += 1
                length = int(self.headers.get("Content-Length", 0))
                raw_body = self.rfile.read(length)
                try:
                    request_body = json.loads(raw_body)
                except json.JSONDecodeError:
                    request_body = {}
                if request_body.get("stream"):
                    self.send_response(200)
                    self.send_header("Content-Type", "text/event-stream; charset=utf-8")
                    self.end_headers()
                    for chunk in outer.stream_chunks:
                        payload = json.dumps({
                            "id": "chatcmpl-stub",
                            "object": "chat.completion.chunk",
                            "model": MODEL,
                            "choices": [{
                                "index": 0,
                                "delta": {"content": chunk},
                                "finish_reason": None,
                            }],
                        }).encode()
                        self.wfile.write(b"data: " + payload + b"\n\n")
                    if outer.stream_usage_total is not None:
                        usage_payload = json.dumps({
                            "id": "chatcmpl-stub",
                            "object": "chat.completion.chunk",
                            "model": MODEL,
                            "choices": [],
                            "usage": {
                                "prompt_tokens": 0,
                                "completion_tokens": outer.stream_usage_total,
                                "total_tokens": outer.stream_usage_total,
                            },
                        }).encode()
                        self.wfile.write(b"data: " + usage_payload + b"\n\n")
                    self.wfile.write(b"data: [DONE]\n\n")
                    return
                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.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},
            }
        },
        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

    stub.stop()
    tracker.stop()


def _chat(tracker_url: str, api_key: str | None, **body_overrides):
    body = {
        "model": MODEL,
        "messages": [{"role": "user", "content": "hi"}],
    }
    body.update(body_overrides)
    data = json.dumps(body).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:
        if body.get("stream"):
            return r.read().decode()
        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_402_for_fresh_key_before_routing(billed_tracker):
    tracker_url, ledger, stub = billed_tracker
    with pytest.raises(urllib.error.HTTPError) as exc_info:
        _chat(tracker_url, api_key="fresh-client")
    assert exc_info.value.code == 402
    assert ledger.get_client_balance("fresh-client") == pytest.approx(0.0)
    assert stub.request_count == 0


def test_proxy_chat_bills_credited_client_and_credits_node(billed_tracker):
    tracker_url, ledger, _ = billed_tracker
    ledger.credit_client("client-1", 0.03, note="admin-credit")
    _chat(tracker_url, api_key="client-1")
    # 1000 tokens at 0.02/1K = 0.02 USDT; explicit 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_caps_inflated_non_streaming_usage_by_request_bounds(billed_tracker):
    tracker_url, ledger, stub = billed_tracker
    stub.total_tokens = 1_000_000
    ledger.credit_client("bounded-client", 100.0, note="admin-credit")

    _chat(tracker_url, api_key="bounded-client", max_tokens=2)

    # messages=[{"content": "hi"}] has a local prompt estimate of one token, so
    # billable total is capped at max_tokens + prompt estimate, not node usage.
    expected_tokens = 3
    assert ledger.get_client_balance("bounded-client") == pytest.approx(100.0 - 0.02 * expected_tokens / 1000)
    events, _ = ledger.events_since(0)
    charge = next(event for event in events if event["type"] == "charge")
    assert charge["total_tokens"] == expected_tokens


def test_proxy_chat_caps_inflated_streaming_usage_by_observed_chunks():
    ledger = BillingLedger(starting_credit=0.0, default_price_per_1k=0.02)
    ledger.credit_client("stream-client", 1.0, note="admin-credit")
    tracker = TrackerServer(
        model_presets={
            MODEL: {
                "layers_start": 0,
                "layers_end": 11,
                "bytes_per_layer": {"bfloat16": 30 * 1024 * 1024},
            }
        },
        billing=ledger,
    )
    port = tracker.start()
    tracker_url = f"http://127.0.0.1:{port}"
    stub = _UsageStubNode(stream_chunks=["one", "two"], stream_usage_total=1_000_000)
    stub_port = stub.start()
    try:
        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()

        _chat(tracker_url, api_key="stream-client", stream=True, max_tokens=2)

        events, _ = ledger.events_since(0)
        charge = next(event for event in events if event["type"] == "charge")
        assert charge["total_tokens"] == 2
        assert ledger.get_client_balance("stream-client") == pytest.approx(1.0 - 0.02 * 2 / 1000)
    finally:
        stub.stop()
        tracker.stop()


def test_proxy_chat_splits_payout_by_tracker_assigned_route_span():
    ledger = BillingLedger(starting_credit=0.0, default_price_per_1k=0.02)
    ledger.credit_client("route-client", 1.0, note="admin-credit")
    tracker = TrackerServer(
        model_presets={
            MODEL: {
                "total_layers": 12,
                "bytes_per_layer": {"bfloat16": 1_000},
            }
        },
        billing=ledger,
    )
    port = tracker.start()
    tracker_url = f"http://127.0.0.1:{port}"
    head = _UsageStubNode(total_tokens=1000)
    head_port = head.start()
    try:
        for endpoint, wallet, vram, bench in (
            (f"http://127.0.0.1:{head_port}", "wallet-head", 5_000, 2.0),
            ("http://127.0.0.1:1", "wallet-tail", 10_000, 1.0),
        ):
            reg = json.dumps({
                "endpoint": endpoint,
                "model": MODEL,
                "vram_bytes": vram,
                "ram_bytes": 10_000,
                "quantizations": ["bfloat16"],
                "benchmark_tokens_per_sec": bench,
                "hardware_profile": {},
                "score": 1.0,
                "tracker_mode": endpoint.endswith(str(head_port)),
                "wallet_address": wallet,
                "managed_assignment": True,
                "shard_start": 0,
                "shard_end": 999,
            }).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()

        _chat(tracker_url, api_key="route-client", max_tokens=1000)

        pool = 0.02 * 0.90
        assert ledger.get_node_pending("wallet-head") == pytest.approx(pool * 4 / 12)
        assert ledger.get_node_pending("wallet-tail") == pytest.approx(pool * 8 / 12)
    finally:
        head.stop()
        tracker.stop()


def test_proxy_chat_402_when_balance_exhausted(billed_tracker):
    tracker_url, ledger, _ = billed_tracker
    ledger.credit_client("client-2", 0.03, note="admin-credit")
    _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_proxy_chat_rejects_request_above_spend_cap_before_routing():
    ledger = BillingLedger(starting_credit=0.0, default_price_per_1k=0.02)
    ledger.credit_client("capped-client", 10.0, note="admin-credit")
    tracker = TrackerServer(
        model_presets={
            MODEL: {
                "layers_start": 0,
                "layers_end": 11,
                "bytes_per_layer": {"bfloat16": 30 * 1024 * 1024},
            }
        },
        billing=ledger,
        max_charge_per_request=0.01,
    )
    port = tracker.start()
    tracker_url = f"http://127.0.0.1:{port}"
    stub = _UsageStubNode(total_tokens=1000)
    stub_port = stub.start()
    try:
        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()

        with pytest.raises(urllib.error.HTTPError) as exc_info:
            _chat(
                tracker_url,
                api_key="capped-client",
                messages=[{"role": "user", "content": " ".join(["prompt"] * 1000)}],
                max_tokens=1,
            )
        body = json.loads(exc_info.value.read())
        assert exc_info.value.code == 402
        assert body["error"]["code"] == "spend_cap_exceeded"
        assert "max_charge_per_request" in body["error"]["message"]
        assert ledger.get_client_balance("capped-client") == pytest.approx(10.0)
        assert stub.request_count == 0
    finally:
        stub.stop()
        tracker.stop()


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.
    assert ledger.get_client_balance("remote-client") == pytest.approx(7.0)