diff --git a/docs/issues/31-billing-ledger.md b/docs/issues/31-billing-ledger.md index ca4c94d..b7a3792 100644 --- a/docs/issues/31-billing-ledger.md +++ b/docs/issues/31-billing-ledger.md @@ -1,4 +1,4 @@ -Status: ready-for-agent +Status: done # 31 — Billing ledger: per-token pricing, 90/10 split, pending balances diff --git a/docs/prd.json b/docs/prd.json index 6283f06..3193985 100644 --- a/docs/prd.json +++ b/docs/prd.json @@ -724,12 +724,13 @@ "Unit tests: single-node route, 3-node split, exhausted balance, restart persistence" ], "priority": 31, - "status": "open", + "status": "done", "notes": "Pure off-chain — no Solana calls in this story. Reuses ComputeAttribution/work_units from packages/contracts.", "dependsOn": [ "US-023", "US-025" - ] + ], + "completionNotes": "BillingLedger in packages/tracker/meshnet_tracker/billing.py: event-sourced USDT ledger (credit/charge/payout/forfeit events, id-deduped), SQLite persistence, per-model price_per_1k_tokens (preset key or prices dict), 90/10 split by work units, walletless shares to protocol cut. server.py: 401/402 gate before routing, billing on non-streaming/streaming/relayed completion paths, GET /v1/billing/summary, POST /v1/billing/gossip, event push in _stats_loop (cursor advances only when all peers reached). CLI --billing-db. settle_node_payout/forfeit_pending are the US-033/034 hooks. 11 new tests in tests/test_billing_ledger.py; suite 195 passed (3 pre-existing ModuleNotFoundError: openai failures unrelated)." }, { "id": "US-032", diff --git a/packages/tracker/meshnet_tracker/billing.py b/packages/tracker/meshnet_tracker/billing.py new file mode 100644 index 0000000..7b38290 --- /dev/null +++ b/packages/tracker/meshnet_tracker/billing.py @@ -0,0 +1,284 @@ +"""Off-chain USDT billing ledger (ADR-0015, US-031). + +Tracks client API-key balances, node pending balances, and the protocol cut. +All mutations are expressed as append-only billing events with unique ids so +the ledger converges when events are gossiped across the tracker hive: every +event is a commutative balance delta, and each tracker applies an event at +most once (dedupe by event id). + +No Solana calls live here — on-chain deposit/settlement adapters (US-032/033) +consume this ledger through `snapshot()` and `settle_node_payout()`. +""" + +from __future__ import annotations + +import json +import sqlite3 +import threading +import time +import uuid + +DEFAULT_PRICE_PER_1K_TOKENS = 0.02 # USDT +DEFAULT_STARTING_CREDIT = 1.0 # USDT of Caller Credit for a new API key +NODE_REVENUE_SHARE = 0.90 # nodes 90% / protocol cut 10% (ADR-0015) + + +class BillingLedger: + """Thread-safe USDT ledger with SQLite persistence and event replication.""" + + SAVE_INTERVAL = 10.0 + + def __init__( + self, + db_path: str | None = None, + *, + prices: dict[str, float] | None = None, + default_price_per_1k: float = DEFAULT_PRICE_PER_1K_TOKENS, + starting_credit: float = DEFAULT_STARTING_CREDIT, + node_share: float = NODE_REVENUE_SHARE, + ) -> None: + self._db_path = db_path + self._prices = dict(prices) if prices else {} + self._default_price_per_1k = default_price_per_1k + self._starting_credit = starting_credit + self._node_share = node_share + self._lock = threading.Lock() + self._client_balances: dict[str, float] = {} + self._node_pending: dict[str, float] = {} + self._protocol_cut: float = 0.0 + self._seen_event_ids: set[str] = set() + self._event_log: list[dict] = [] # full log, order of local application + self._dirty = False + if self._db_path: + self._init_db() + self._load_from_db() + + # ---- pricing ---- + + def price_for(self, model: str) -> float: + return self._prices.get(model, self._default_price_per_1k) + + def set_price(self, model: str, price_per_1k: float) -> None: + with self._lock: + self._prices[model] = price_per_1k + + # ---- local operations (create + apply + log an event) ---- + + def ensure_client(self, api_key: str) -> float: + """Return the client's balance, granting Caller Credit on first touch.""" + with self._lock: + if api_key not in self._client_balances: + self._apply_locked({ + "id": f"credit-{uuid.uuid4().hex}", + "type": "credit", + "api_key": api_key, + "amount": self._starting_credit, + "ts": time.time(), + "note": "caller-credit", + }) + return self._client_balances[api_key] + + def has_funds(self, api_key: str) -> bool: + return self.ensure_client(api_key) > 0.0 + + def credit_client(self, api_key: str, amount: float, *, note: str = "deposit") -> float: + if amount <= 0: + raise ValueError("credit amount must be positive") + with self._lock: + self._apply_locked({ + "id": f"credit-{uuid.uuid4().hex}", + "type": "credit", + "api_key": api_key, + "amount": amount, + "ts": time.time(), + "note": note, + }) + return self._client_balances[api_key] + + def charge_request( + self, + api_key: str, + model: str, + total_tokens: int, + node_work: list[tuple[str | None, int]], + ) -> dict: + """Debit the client and split the fee 90/10. + + ``node_work`` is ``[(wallet_address | None, work_units), ...]`` for the + nodes that served the request. Work units of nodes without a wallet + accrue to the protocol cut — there is nowhere to pay them out. + The client balance may dip below zero on the final request; the next + request is then rejected by ``has_funds`` (post-pay drift, standard + metered-billing behavior). + """ + cost = self.price_for(model) * max(0, total_tokens) / 1000.0 + total_work = sum(max(0, w) for _, w in node_work) + node_pool = cost * self._node_share + shares: dict[str, float] = {} + if total_work > 0: + for wallet, work in node_work: + if wallet and work > 0: + shares[wallet] = shares.get(wallet, 0.0) + node_pool * work / total_work + protocol_amount = cost - sum(shares.values()) + with self._lock: + if api_key not in self._client_balances: + self._apply_locked({ + "id": f"credit-{uuid.uuid4().hex}", + "type": "credit", + "api_key": api_key, + "amount": self._starting_credit, + "ts": time.time(), + "note": "caller-credit", + }) + event = { + "id": f"charge-{uuid.uuid4().hex}", + "type": "charge", + "api_key": api_key, + "model": model, + "total_tokens": total_tokens, + "cost": cost, + "shares": shares, + "protocol_amount": protocol_amount, + "ts": time.time(), + } + self._apply_locked(event) + return event + + def settle_node_payout(self, wallet: str, amount: float, *, reference: str = "") -> dict: + """Deduct a paid-out amount from a node's pending balance (US-033 hook).""" + if amount <= 0: + raise ValueError("payout amount must be positive") + with self._lock: + event = { + "id": f"payout-{uuid.uuid4().hex}", + "type": "payout", + "wallet": wallet, + "amount": amount, + "reference": reference, + "ts": time.time(), + } + self._apply_locked(event) + return event + + def forfeit_pending(self, wallet: str, *, reason: str = "fraud") -> dict: + """Forfeit a node's entire pending balance to the protocol cut (US-034 hook).""" + with self._lock: + event = { + "id": f"forfeit-{uuid.uuid4().hex}", + "type": "forfeit", + "wallet": wallet, + "amount": self._node_pending.get(wallet, 0.0), + "reason": reason, + "ts": time.time(), + } + self._apply_locked(event) + return event + + # ---- replication ---- + + def events_since(self, index: int) -> tuple[list[dict], int]: + """Return events after ``index`` in the local log and the new cursor.""" + with self._lock: + return list(self._event_log[index:]), len(self._event_log) + + def apply_events(self, events: list[dict]) -> int: + """Apply peer events not yet seen locally. Returns how many applied.""" + applied = 0 + with self._lock: + for event in events: + event_id = event.get("id") + if not event_id or event_id in self._seen_event_ids: + continue + self._apply_locked(event) + applied += 1 + return applied + + def _apply_locked(self, event: dict) -> None: + etype = event.get("type") + if etype == "credit": + key = event["api_key"] + self._client_balances[key] = self._client_balances.get(key, 0.0) + float(event["amount"]) + elif etype == "charge": + key = event["api_key"] + self._client_balances[key] = self._client_balances.get(key, 0.0) - float(event["cost"]) + for wallet, amount in (event.get("shares") or {}).items(): + self._node_pending[wallet] = self._node_pending.get(wallet, 0.0) + float(amount) + self._protocol_cut += float(event.get("protocol_amount", 0.0)) + elif etype == "payout": + wallet = event["wallet"] + self._node_pending[wallet] = self._node_pending.get(wallet, 0.0) - float(event["amount"]) + elif etype == "forfeit": + wallet = event["wallet"] + amount = float(event.get("amount", 0.0)) + self._node_pending[wallet] = self._node_pending.get(wallet, 0.0) - amount + self._protocol_cut += amount + else: + return + self._seen_event_ids.add(event["id"]) + self._event_log.append(event) + self._dirty = True + + # ---- views ---- + + def get_client_balance(self, api_key: str) -> float: + with self._lock: + return self._client_balances.get(api_key, 0.0) + + def get_node_pending(self, wallet: str) -> float: + with self._lock: + return self._node_pending.get(wallet, 0.0) + + def snapshot(self) -> dict: + with self._lock: + return { + "clients": dict(self._client_balances), + "node_pending": dict(self._node_pending), + "protocol_cut": self._protocol_cut, + "prices": dict(self._prices), + "default_price_per_1k_tokens": self._default_price_per_1k, + "node_share": self._node_share, + "event_count": len(self._event_log), + } + + # ---- persistence ---- + + def _init_db(self) -> None: + con = sqlite3.connect(self._db_path) # type: ignore[arg-type] + con.execute( + "CREATE TABLE IF NOT EXISTS billing_events " + "(event_id TEXT PRIMARY KEY, payload TEXT NOT NULL, ts REAL NOT NULL)" + ) + con.commit() + con.close() + + def _load_from_db(self) -> None: + con = sqlite3.connect(self._db_path) # type: ignore[arg-type] + rows = con.execute( + "SELECT payload FROM billing_events ORDER BY ts, event_id" + ).fetchall() + con.close() + with self._lock: + for (payload,) in rows: + try: + event = json.loads(payload) + except json.JSONDecodeError: + continue + if event.get("id") not in self._seen_event_ids: + self._apply_locked(event) + self._dirty = False + + def save_to_db(self) -> None: + if not self._db_path: + return + with self._lock: + if not self._dirty: + return + events = list(self._event_log) + self._dirty = False + con = sqlite3.connect(self._db_path) # type: ignore[arg-type] + con.executemany( + "INSERT OR IGNORE INTO billing_events (event_id, payload, ts) VALUES (?, ?, ?)", + [(e["id"], json.dumps(e), float(e.get("ts", 0.0))) for e in events], + ) + con.commit() + con.close() diff --git a/packages/tracker/meshnet_tracker/cli.py b/packages/tracker/meshnet_tracker/cli.py index c7da977..d59af63 100644 --- a/packages/tracker/meshnet_tracker/cli.py +++ b/packages/tracker/meshnet_tracker/cli.py @@ -38,6 +38,12 @@ def main() -> None: default=None, help="Public ws(s):// relay URL advertised to nodes, for example wss://ai.neuron.d-popov.com/ws", ) + common.add_argument( + "--billing-db", + default=None, + metavar="PATH", + help="SQLite database path for the USDT billing ledger (enables billing; ADR-0015)", + ) parser = argparse.ArgumentParser( prog="meshnet-tracker", @@ -61,6 +67,7 @@ def main() -> None: cluster_self_url=args.self_url, stats_db=getattr(args, "stats_db", None), relay_url=relay_url, + billing_db=getattr(args, "billing_db", None), ) port = server.start() print(f"meshnet-tracker listening on http://{args.host}:{port}", flush=True) diff --git a/packages/tracker/meshnet_tracker/server.py b/packages/tracker/meshnet_tracker/server.py index b71ca94..1e7c3e6 100644 --- a/packages/tracker/meshnet_tracker/server.py +++ b/packages/tracker/meshnet_tracker/server.py @@ -34,6 +34,7 @@ import uuid from importlib.resources import files from typing import Any +from .billing import BillingLedger from .gossip import NodeGossip from .raft import RaftNode @@ -975,6 +976,29 @@ def _rebalance_all_locked(server: "_TrackerHTTPServer") -> None: _rebalance_hf_model_locked(server, hf_repo) +def _api_key_from_headers(headers) -> str | None: + auth = headers.get("Authorization") + if not auth: + return None + if auth.lower().startswith("bearer "): + return auth.split(" ", 1)[1].strip() or None + return auth.strip() or None + + +def _usage_total_tokens(payload: dict) -> int | None: + usage = payload.get("usage") + if not isinstance(usage, dict): + return None + total = usage.get("total_tokens") + if isinstance(total, (int, float)): + return int(total) + prompt = usage.get("prompt_tokens") + completion = usage.get("completion_tokens") + if isinstance(prompt, (int, float)) or isinstance(completion, (int, float)): + return int(prompt or 0) + int(completion or 0) + return None + + def _registration_ban_error(contracts: Any | None, wallet_address: str | None) -> str | None: if contracts is None or not wallet_address: return None @@ -1021,6 +1045,7 @@ class _TrackerHTTPServer(socketserver.ThreadingMixIn, http.server.HTTPServer): raft: "RaftNode | None" = None, gossip: "NodeGossip | None" = None, stats: "_StatsCollector | None" = None, + billing: "BillingLedger | None" = None, ) -> None: super().__init__(addr, handler) self.registry = registry @@ -1033,6 +1058,7 @@ class _TrackerHTTPServer(socketserver.ThreadingMixIn, http.server.HTTPServer): self.raft = raft self.gossip = gossip self.stats: _StatsCollector | None = stats + self.billing: BillingLedger | None = billing class _TrackerHandler(http.server.BaseHTTPRequestHandler): @@ -1085,6 +1111,9 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler): if self.path == "/v1/stats/gossip": self._handle_stats_gossip() return + if self.path == "/v1/billing/gossip": + self._handle_billing_gossip() + return parts = self.path.split("/") # /v1/nodes//heartbeat -> ['', 'v1', 'nodes', '', 'heartbeat'] if len(parts) == 5 and parts[1] == "v1" and parts[2] == "nodes" and parts[4] == "heartbeat": @@ -1117,6 +1146,8 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler): self._handle_raft_status() elif parsed.path == "/v1/stats": self._handle_stats() + elif parsed.path == "/v1/billing/summary": + self._handle_billing_summary() elif parsed.path == "/v1/health": self._send_json(200, {"status": "ok"}) else: @@ -1345,6 +1376,24 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler): if model and server.stats is not None: server.stats.record_request(model) + # Billing gate (ADR-0015): reject before any routing — no free work. + api_key = _api_key_from_headers(self.headers) + if server.billing is not None: + if api_key is None: + self._send_json(401, {"error": { + "message": "missing API key: send Authorization: Bearer ", + "type": "invalid_request_error", + "code": "missing_api_key", + }}) + return + if not server.billing.has_funds(api_key): + self._send_json(402, {"error": { + "message": "insufficient balance: deposit USDT to continue", + "type": "insufficient_quota", + "code": "insufficient_balance", + }}) + return + # Find a live tracker-mode node for this model with server.lock: self._purge_expired_nodes() @@ -1395,12 +1444,15 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler): # This allows overlapping shard registrations without double-computation. covered_up_to = rs - 1 route_hops: list[dict] = [] + node_work: list[tuple[str | None, int]] = [] for rn in route_nodes: hop: dict = {"endpoint": rn.endpoint, "start_layer": covered_up_to + 1} if rn.relay_addr: hop["relay_addr"] = rn.relay_addr route_hops.append(hop) - covered_up_to = rn.shard_end if rn.shard_end is not None else covered_up_to + effective_end = rn.shard_end if rn.shard_end is not None else covered_up_to + node_work.append((rn.wallet_address, max(0, effective_end - covered_up_to))) + covered_up_to = effective_end # Strip the first-shard node we're about to proxy to — it's already handling the request. downstream_hops = [h for h in route_hops if h["endpoint"].rstrip("/") != node.endpoint.rstrip("/")] downstream_urls = json.dumps(downstream_hops) @@ -1449,6 +1501,13 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler): ) if relayed is not None: self._send_relayed_response(relayed) + if int(relayed.get("status", 503)) < 400: + body_text = relayed.get("body") or "" + try: + tokens = _usage_total_tokens(json.loads(body_text)) or 0 + except (json.JSONDecodeError, TypeError): + tokens = 0 + self._bill_completed(api_key, model, tokens, node_work) return print( f"[tracker] relay proxy failed {request_id}: {node.relay_addr}; " @@ -1495,6 +1554,8 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler): self.send_header("Content-Type", "text/event-stream; charset=utf-8") self.send_header("Cache-Control", "no-cache") self.end_headers() + stream_tokens: int | None = None + chunk_count = 0 try: while True: line = upstream.readline() @@ -1502,8 +1563,27 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler): break self.wfile.write(line) self.wfile.flush() + if line.startswith(b"data:"): + payload = line[5:].strip() + if payload and payload != b"[DONE]": + chunk_count += 1 + if b'"usage"' in payload: + try: + found = _usage_total_tokens(json.loads(payload)) + if found: + stream_tokens = found + except json.JSONDecodeError: + pass except BrokenPipeError: pass + # Bill even on client disconnect — the nodes did the work. + # Chunk count approximates generated tokens when the stream + # carries no usage record. + self._bill_completed( + api_key, model, + stream_tokens if stream_tokens is not None else chunk_count, + node_work, + ) else: # Non-streaming: buffer and relay resp_body = upstream.read() @@ -1519,6 +1599,33 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler): self.wfile.write(resp_body) except BrokenPipeError: pass + try: + tokens = _usage_total_tokens(json.loads(resp_body)) or 0 + except json.JSONDecodeError: + tokens = 0 + self._bill_completed(api_key, model, tokens, node_work) + + def _bill_completed( + self, + api_key: str | None, + model: str, + total_tokens: int, + node_work: list[tuple[str | None, int]], + ) -> None: + """Charge a completed request against the billing ledger (ADR-0015).""" + server: _TrackerHTTPServer = self.server # type: ignore[assignment] + if server.billing is None or api_key is None: + return + try: + event = server.billing.charge_request(api_key, model, total_tokens, node_work) + print( + f"[tracker] billed api_key=…{api_key[-6:]}: model={model!r} " + f"tokens={total_tokens} cost={event['cost']:.6f} USDT " + f"shares={event['shares']}", + flush=True, + ) + except Exception as exc: + print(f"[tracker] billing failed for model={model!r}: {exc}", flush=True) def _send_relayed_response(self, response: dict) -> None: status = int(response.get("status", 503)) @@ -1870,6 +1977,28 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler): server.stats.merge_peer_rpms(tracker_url, rpms) self._send_json(200, {}) + def _handle_billing_summary(self): + server: _TrackerHTTPServer = self.server # type: ignore[assignment] + if server.billing is None: + self._send_json(404, {"error": "billing is not enabled on this tracker"}) + return + self._send_json(200, server.billing.snapshot()) + + def _handle_billing_gossip(self): + server: _TrackerHTTPServer = self.server # type: ignore[assignment] + body = self._read_json_body() + if body is None: + return + if server.billing is None: + self._send_json(200, {"applied": 0}) + return + events = body.get("events") + if not isinstance(events, list): + self._send_json(400, {"error": "events must be a list"}) + return + applied = server.billing.apply_events([e for e in events if isinstance(e, dict)]) + self._send_json(200, {"applied": applied}) + def _handle_assign(self, parsed: urllib.parse.ParseResult): """Return an optimal shard assignment for a node given its hardware profile. @@ -2299,6 +2428,8 @@ class TrackerServer: cluster_self_url: str | None = None, stats_db: str | None = None, relay_url: str | None = None, + billing: BillingLedger | None = None, + billing_db: str | None = None, ) -> None: self._host = host self._requested_port = port @@ -2323,6 +2454,15 @@ class TrackerServer: self._stats: _StatsCollector | None = _StatsCollector(db_path=stats_db) if stats_db else _StatsCollector() self._stats_stop = threading.Event() self._stats_thread: threading.Thread | None = None + if billing is None and billing_db: + preset_prices = { + name: float(preset["price_per_1k_tokens"]) + for name, preset in self._model_presets.items() + if isinstance(preset, dict) and "price_per_1k_tokens" in preset + } + billing = BillingLedger(db_path=billing_db, prices=preset_prices) + self._billing: BillingLedger | None = billing + self._billing_gossip_cursor = 0 self.port: int | None = None def start(self) -> int: @@ -2346,6 +2486,7 @@ class TrackerServer: self._minimum_stake, relay_url=effective_relay_url, stats=self._stats, + billing=self._billing, ) self.port = self._server.server_address[1] @@ -2416,10 +2557,12 @@ class TrackerServer: _rebalance_all_locked(server) def _stats_loop(self) -> None: - """Periodically save stats to DB and push local slice to cluster peers.""" + """Periodically save stats/billing to DB and push local slices to cluster peers.""" while not self._stats_stop.wait(_StatsCollector.SAVE_INTERVAL): if self._stats is not None: self._stats.save_to_db() + if self._billing is not None: + self._billing.save_to_db() if self._stats is not None and self._cluster_peers: self_url = self._cluster_self_url or f"http://{self._host}:{self.port}" local_rpms = self._stats.get_local_rpms() @@ -2436,6 +2579,28 @@ class TrackerServer: r.read() except Exception: pass + if self._billing is not None and self._cluster_peers: + events, cursor = self._billing.events_since(self._billing_gossip_cursor) + if events: + delivered_all = True + body = json.dumps({"events": events}).encode() + for peer in self._cluster_peers: + try: + req = urllib.request.Request( + f"{peer}/v1/billing/gossip", + data=body, + headers={"Content-Type": "application/json"}, + method="POST", + ) + with urllib.request.urlopen(req, timeout=2.0) as r: + r.read() + except Exception: + delivered_all = False + # Only advance past events every peer has seen; unreachable + # peers get the full backlog on the next tick (idempotent + # via event-id dedupe on the receiving side). + if delivered_all: + self._billing_gossip_cursor = cursor def stop(self) -> None: if self._raft is not None: @@ -2448,6 +2613,8 @@ class TrackerServer: self._stats_stop.set() if self._stats is not None: self._stats.save_to_db() + if self._billing is not None: + self._billing.save_to_db() self._server.shutdown() self._server.server_close() if self._thread is not None: diff --git a/tests/test_billing_ledger.py b/tests/test_billing_ledger.py new file mode 100644 index 0000000..cb5bec4 --- /dev/null +++ b/tests/test_billing_ledger.py @@ -0,0 +1,294 @@ +"""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.billing import BillingLedger +from meshnet_tracker.server import TrackerServer + +MODEL = "openai-community/gpt2" + + +# ---------------------------------------------------------------- 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_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, + ) + 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) + + summary = json.loads( + urllib.request.urlopen(f"{tracker_url}/v1/billing/summary").read() + ) + assert summary["node_pending"]["wallet-head"] == pytest.approx(0.02 * 0.90) + assert summary["protocol_cut"] == pytest.approx(0.02 * 0.10) + + +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"}, + 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)