neuron-tai/packages/tracker/meshnet_tracker/server.py

"""Tracker HTTP server — node registry and route selection for the inference mesh.

HTTP API contract:
- POST /v1/nodes/register
  Request JSON: {
    "endpoint": "http://host:port", "shard_start": int, "shard_end": int,
    "model": str optional, "shard_checksum": str optional,
    "hardware_profile": object, "wallet_address": str optional,
    "score": number optional
  }
  Response 200: {"node_id": str}
  Response 400: {"error": str}
- POST /v1/nodes/{node_id}/heartbeat
  Response 200: {}
  Response 404: {"error": "node not found"}
- GET /v1/route?model=<preset>
  Response 200: {"route": ["http://node-a", "http://node-b"]}
  Response 400/404/503: {"error": str}
- GET /v1/routes?model=<preset>&redundancy=<n>
  Response 200: {"routes": [{"route": [...], "nodes": [...]}]}
  Response 400/404/503: {"error": str}
"""

import http.server
import hashlib
import json
import socketserver
import sqlite3
import threading
import time
import urllib.parse
import urllib.request
import uuid
from typing import Any

from .gossip import NodeGossip
from .raft import RaftNode


def derive_relay_url_from_public_tracker_url(url: str | None) -> str | None:
    """Return wss://host/ws when url is a public HTTPS tracker origin."""
    if not url:
        return None
    parsed = urllib.parse.urlparse(url)
    if parsed.scheme != "https":
        return None
    host = parsed.hostname
    if not host or host in ("127.0.0.1", "localhost"):
        return None
    return f"wss://{parsed.netloc}/ws"


DEFAULT_MODEL_PRESETS: dict[str, dict] = {
    "stub-model": {
        "layers_start": 0,
        "layers_end": 31,
        "bytes_per_layer": {"bfloat16": 30 * 1024 * 1024},
    },
    "openai-community/gpt2": {
        "layers_start": 0,
        "layers_end": 11,
        "bytes_per_layer": {"bfloat16": 30 * 1024 * 1024, "int8": 15 * 1024 * 1024, "nf4": 8 * 1024 * 1024},
    },
}

DEFAULT_VRAM_BYTES = 8 * 1024 * 1024 * 1024
DEFAULT_RAM_BYTES = 16 * 1024 * 1024 * 1024
DEFAULT_QUANTIZATIONS = ["bfloat16"]
DEFAULT_BENCHMARK_TOKENS_PER_SEC = 1.0


def _model_aliases(model: str | None) -> set[str]:
    """Return stable lookup aliases for a model repo or display name."""
    if not model:
        return set()
    aliases = {model}
    short = model.rsplit("/", 1)[-1]
    aliases.add(short)
    lowered = short.lower()
    aliases.add(lowered)
    if lowered.endswith("-instruct"):
        aliases.add(lowered.removesuffix("-instruct"))
    return aliases


def _node_matches_model(node: "_NodeEntry", model: str) -> bool:
    requested = _model_aliases(model)
    if not requested:
        return False
    return bool(requested & (_model_aliases(node.model) | _model_aliases(node.hf_repo)))


class _RollingCounter:
    """Circular-bucket request counter.

    Tracks events in `num_buckets` slots, each covering `bucket_seconds` of wall time.
    Old buckets are silently discarded when their epoch expires.
    """

    def __init__(self, num_buckets: int, bucket_seconds: int) -> None:
        self._num = num_buckets
        self._bsec = bucket_seconds
        self._counts: list[int] = [0] * num_buckets
        self._epochs: list[int] = [-1] * num_buckets  # which epoch each slot holds

    def _epoch(self, now: float) -> int:
        return int(now) // self._bsec

    def record(self, now: float | None = None) -> None:
        t = now if now is not None else time.time()
        ep = self._epoch(t)
        idx = ep % self._num
        if self._epochs[idx] != ep:
            self._counts[idx] = 0
            self._epochs[idx] = ep
        self._counts[idx] += 1

    def rpm(self, now: float | None = None) -> float:
        t = now if now is not None else time.time()
        cutoff = self._epoch(t) - self._num  # epochs <= this are stale
        total = sum(self._counts[i] for i in range(self._num) if self._epochs[i] > cutoff)
        window_minutes = (self._num * self._bsec) / 60.0
        return total / window_minutes if window_minutes > 0 else 0.0

    def buckets(self) -> list[tuple[int, int]]:
        return [(self._epochs[i], self._counts[i]) for i in range(self._num)]

    def restore_buckets(self, data: list[tuple[int, int]]) -> None:
        for i, (ep, cnt) in enumerate(data):
            if i < self._num:
                self._epochs[i] = ep
                self._counts[i] = cnt


class _ModelStats:
    """Three rolling windows for one model: last hour, last day, last month."""

    def __init__(self) -> None:
        self.per_minute = _RollingCounter(60, 60)     # 60 × 1-min buckets = 1 hour
        self.per_hour   = _RollingCounter(24, 3600)   # 24 × 1-hr  buckets = 1 day
        self.per_day    = _RollingCounter(30, 86400)  # 30 × 1-day buckets = ~1 month

    def record(self, now: float | None = None) -> None:
        t = now if now is not None else time.time()
        self.per_minute.record(t)
        self.per_hour.record(t)
        self.per_day.record(t)

    def rpms(self, now: float | None = None) -> dict:
        t = now if now is not None else time.time()
        return {
            "rpm_last_hour":  round(self.per_minute.rpm(t), 4),
            "rpm_last_day":   round(self.per_hour.rpm(t), 4),
            "rpm_last_month": round(self.per_day.rpm(t), 4),
        }


class _StatsCollector:
    """Thread-safe model request stats with SQLite persistence and peer slice merging."""

    SAVE_INTERVAL = 60.0  # seconds between DB saves

    def __init__(self, db_path: str | None = None) -> None:
        self._lock = threading.Lock()
        self._local: dict[str, _ModelStats] = {}
        self._peer_rpms: dict[str, dict[str, dict]] = {}  # tracker_url -> model -> rpms
        self._db_path = db_path
        if db_path:
            self._init_db()
            self._load_from_db()

    # ---------- public API ----------

    def record_request(self, model: str, now: float | None = None) -> None:
        t = now if now is not None else time.time()
        with self._lock:
            if model not in self._local:
                self._local[model] = _ModelStats()
            self._local[model].record(t)

    def get_local_rpms(self, now: float | None = None) -> dict[str, dict]:
        t = now if now is not None else time.time()
        with self._lock:
            return {m: s.rpms(t) for m, s in self._local.items()}

    def merge_peer_rpms(self, tracker_url: str, rpms: dict[str, dict]) -> None:
        with self._lock:
            self._peer_rpms[tracker_url] = dict(rpms)

    def get_combined_stats(self, now: float | None = None) -> dict:
        t = now if now is not None else time.time()
        with self._lock:
            combined: dict[str, dict] = {}
            for model, ms in self._local.items():
                combined[model] = ms.rpms(t)
            for _url, peer in self._peer_rpms.items():
                for model, rpms in peer.items():
                    if model not in combined:
                        combined[model] = {"rpm_last_hour": 0.0, "rpm_last_day": 0.0, "rpm_last_month": 0.0}
                    for key in ("rpm_last_hour", "rpm_last_day", "rpm_last_month"):
                        combined[model][key] = round(combined[model].get(key, 0.0) + rpms.get(key, 0.0), 4)
        return combined

    # ---------- persistence ----------

    def _init_db(self) -> None:
        con = sqlite3.connect(self._db_path)  # type: ignore[arg-type]
        con.execute(
            "CREATE TABLE IF NOT EXISTS model_rpm_buckets "
            "(model TEXT, window TEXT, bucket_idx INTEGER, bucket_epoch INTEGER, count INTEGER, "
            "PRIMARY KEY (model, window, bucket_idx))"
        )
        con.commit()
        con.close()

    def save_to_db(self) -> None:
        if not self._db_path:
            return
        with self._lock:
            rows = []
            for model, ms in self._local.items():
                for window_name, counter in (
                    ("hour", ms.per_minute),
                    ("day", ms.per_hour),
                    ("month", ms.per_day),
                ):
                    for idx, (ep, cnt) in enumerate(counter.buckets()):
                        if ep >= 0:
                            rows.append((model, window_name, idx, ep, cnt))
        con = sqlite3.connect(self._db_path)  # type: ignore[arg-type]
        con.executemany(
            "INSERT OR REPLACE INTO model_rpm_buckets VALUES (?,?,?,?,?)", rows
        )
        con.commit()
        con.close()

    def _load_from_db(self) -> None:
        con = sqlite3.connect(self._db_path)  # type: ignore[arg-type]
        rows = con.execute("SELECT model, window, bucket_idx, bucket_epoch, count FROM model_rpm_buckets").fetchall()
        con.close()
        grouped: dict[str, dict[str, list[tuple[int, int]]]] = {}
        for model, window, idx, ep, cnt in rows:
            grouped.setdefault(model, {}).setdefault(window, []).append((idx, ep, cnt))
        for model, windows in grouped.items():
            ms = _ModelStats()
            for window_name, entries in windows.items():
                counter = {"hour": ms.per_minute, "day": ms.per_hour, "month": ms.per_day}.get(window_name)
                if counter is None:
                    continue
                data = [(0, -1)] * counter._num
                for idx, ep, cnt in entries:
                    if 0 <= idx < counter._num:
                        data[idx] = (ep, cnt)
                counter.restore_buckets(data)
            self._local[model] = ms


class _NodeEntry:
    __slots__ = (
        "node_id", "endpoint", "shard_start", "shard_end",
        "model", "hf_repo", "num_layers", "shard_checksum", "hardware_profile", "wallet_address",
        "score", "vram_bytes", "ram_bytes", "quantizations", "max_loaded_shards",
        "benchmark_tokens_per_sec", "quantization", "managed_assignment",
        "pending_directives", "last_heartbeat", "tracker_mode",
        "relay_addr", "cert_fingerprint", "peer_id",
        # heartbeat stats (reported by node, cumulative)
        "total_requests", "failed_requests", "queue_depth", "uptime_seconds",
        "status",  # "ready" | "loading"
        "heartbeats_expected", "heartbeats_received",
        # dynamic reassignment queued by the tracker
        "pending_new_assignment",
    )

    def __init__(
        self,
        node_id: str,
        endpoint: str,
        shard_start: int | None,
        shard_end: int | None,
        model: str | None,
        shard_checksum: str | None,
        hardware_profile: dict,
        wallet_address: str | None,
        score: float,
        vram_bytes: int = DEFAULT_VRAM_BYTES,
        ram_bytes: int = DEFAULT_RAM_BYTES,
        quantizations: list[str] | None = None,
        max_loaded_shards: int = 1,
        benchmark_tokens_per_sec: float = DEFAULT_BENCHMARK_TOKENS_PER_SEC,
        quantization: str | None = None,
        managed_assignment: bool = False,
        tracker_mode: bool = False,
        hf_repo: str | None = None,
        num_layers: int | None = None,
        relay_addr: str | None = None,
        cert_fingerprint: str | None = None,
        peer_id: str | None = None,
    ) -> None:
        self.node_id = node_id
        self.endpoint = endpoint
        self.shard_start = shard_start
        self.shard_end = shard_end
        self.model = model
        self.shard_checksum = shard_checksum
        self.hardware_profile = hardware_profile
        self.wallet_address = wallet_address
        self.score = score
        self.vram_bytes = vram_bytes
        self.ram_bytes = ram_bytes
        self.quantizations = quantizations or list(DEFAULT_QUANTIZATIONS)
        self.max_loaded_shards = max_loaded_shards
        self.benchmark_tokens_per_sec = benchmark_tokens_per_sec
        self.quantization = quantization
        self.managed_assignment = managed_assignment
        self.tracker_mode = tracker_mode
        self.hf_repo = hf_repo
        self.num_layers = num_layers
        self.relay_addr = relay_addr
        self.cert_fingerprint = cert_fingerprint
        self.peer_id = peer_id
        self.pending_directives: list[dict] = []
        self.last_heartbeat: float = time.monotonic()
        self.total_requests: int = 0
        self.failed_requests: int = 0
        self.queue_depth: int = 0
        self.uptime_seconds: float = 0.0
        self.status: str = "ready"
        self.heartbeats_expected: int = 0
        self.heartbeats_received: int = 0
        self.pending_new_assignment: dict | None = None


def _effective_throughput(node: "_NodeEntry") -> float:
    """Effective tokens/s accounting for current queue depth."""
    return node.benchmark_tokens_per_sec / (node.queue_depth + 1)


def _select_route(
    nodes: list[_NodeEntry],
    required_start: int,
    required_end: int,
) -> tuple[list[_NodeEntry], str]:
    """Greedy interval-cover biased toward fast, lightly-loaded nodes.

    Among nodes that equally advance coverage, prefer the one with higher
    effective throughput: benchmark_tokens_per_sec / (queue_depth + 1).
    Tiebreak: higher shard_end (fewer hops).
    """
    candidates = sorted(
        [node for node in nodes if node.shard_start is not None and node.shard_end is not None],
        key=lambda n: (n.shard_start, -n.shard_end),  # type: ignore[operator]
    )
    route: list[_NodeEntry] = []
    covered_up_to = required_start - 1

    while covered_up_to < required_end:
        best: _NodeEntry | None = None
        for node in candidates:
            if node.shard_start <= covered_up_to + 1 and node.shard_end > covered_up_to:
                if best is None:
                    best = node
                elif node.shard_end > best.shard_end:
                    best = node
                elif node.shard_end == best.shard_end and _effective_throughput(node) > _effective_throughput(best):
                    best = node
        if best is None:
            missing = covered_up_to + 1
            return [], f"no route available: no registered node covers layer {missing}"
        route.append(best)
        covered_up_to = best.shard_end
        candidates = [n for n in candidates if n is not best]

    return route, ""


def _coverage_percentage(
    nodes: list[_NodeEntry],
    required_start: int,
    required_end: int,
) -> float:
    required_layers = required_end - required_start + 1
    if required_layers <= 0:
        return 0.0

    intervals = sorted(
        (
            (max(required_start, node.shard_start), min(required_end, node.shard_end))
            for node in nodes
            if node.shard_start is not None
            and node.shard_end is not None
            if node.shard_end >= required_start and node.shard_start <= required_end
        ),
        key=lambda interval: interval[0],
    )
    covered = 0
    covered_until = required_start - 1
    for start, end in intervals:
        if end <= covered_until:
            continue
        next_start = max(start, covered_until + 1)
        if next_start > end:
            continue
        covered += end - next_start + 1
        covered_until = end
    return round((covered / required_layers) * 100, 2)


def _preset_layer_bounds(preset: dict) -> tuple[int, int]:
    start = int(preset.get("layers_start", 0))
    if "layers_end" in preset:
        return start, int(preset["layers_end"])
    return start, start + int(preset["total_layers"]) - 1


def _preset_bytes_per_layer(preset: dict) -> dict[str, int]:
    raw = preset.get("bytes_per_layer", preset.get("bytes_per_layer_at_quant", {}))
    if isinstance(raw, dict) and raw:
        return {str(quant): int(value) for quant, value in raw.items()}
    return {"bfloat16": 30 * 1024 * 1024}


def _node_quantization(node: _NodeEntry, preset: dict) -> str:
    bytes_per_layer = _preset_bytes_per_layer(preset)
    if node.quantization in bytes_per_layer:
        return node.quantization
    for quantization in node.quantizations:
        if quantization in bytes_per_layer:
            return quantization
    return next(iter(bytes_per_layer))


def _node_memory_budget_bytes(node: _NodeEntry) -> tuple[int, str]:
    """Return the memory pool used for shard-capacity planning."""
    if node.vram_bytes > 0:
        return node.vram_bytes, "vram"
    if node.ram_bytes > 0:
        return node.ram_bytes, "ram"
    return DEFAULT_RAM_BYTES, "ram-default"


def _node_layer_capacity(node: _NodeEntry, preset: dict) -> int:
    bytes_per_layer = _preset_bytes_per_layer(preset)
    quantization = _node_quantization(node, preset)
    layer_bytes = bytes_per_layer[quantization]
    if layer_bytes <= 0:
        return 0
    memory_budget_bytes, _ = _node_memory_budget_bytes(node)
    return int((memory_budget_bytes * 0.8) // layer_bytes)


def _node_capacity_summary(node: _NodeEntry, preset: dict | None = None) -> dict:
    """Operator-facing capacity fields for inspection endpoints."""
    memory_budget_bytes, memory_budget_source = _node_memory_budget_bytes(node)
    summary = {
        "vram_bytes": node.vram_bytes,
        "ram_bytes": node.ram_bytes,
        "memory_budget_bytes": memory_budget_bytes,
        "memory_budget_source": memory_budget_source,
        "max_loaded_shards": node.max_loaded_shards,
        "quantizations": list(node.quantizations),
        "quantization": node.quantization,
        "benchmark_tokens_per_sec": node.benchmark_tokens_per_sec,
        "effective_throughput": round(_effective_throughput(node), 4),
    }
    if preset is not None:
        summary["max_assignable_layers"] = _node_layer_capacity(node, preset)
    return summary


def _coverage_map(
    nodes: list[_NodeEntry],
    required_start: int,
    required_end: int,
) -> list[dict]:
    layer_counts = []
    for layer in range(required_start, required_end + 1):
        count = 0
        for node in nodes:
            if node.shard_start is None or node.shard_end is None:
                continue
            if node.shard_start <= layer <= node.shard_end:
                count += 1
        layer_counts.append((layer, count))

    coverage: list[dict] = []
    for layer, count in layer_counts:
        if coverage and coverage[-1]["node_count"] == count and coverage[-1]["end_layer"] == layer - 1:
            coverage[-1]["end_layer"] = layer
        else:
            coverage.append({"start_layer": layer, "end_layer": layer, "node_count": count})
    return coverage


def _node_health(node: "_NodeEntry", heartbeat_timeout: float) -> dict:
    """Per-node health detail for the availability map."""
    age = time.monotonic() - node.last_heartbeat
    alive = age <= heartbeat_timeout
    hb_expected = max(1, round(node.uptime_seconds / 20.0))  # assume ~20s interval
    hb_rate = round(min(1.0, node.heartbeats_received / hb_expected), 4) if node.heartbeats_received else 0.0
    total = node.total_requests
    inf_rate = round((total - node.failed_requests) / total, 4) if total > 0 else 1.0
    return {
        "node_id": node.node_id,
        "endpoint": node.endpoint,
        "alive": alive,
        "last_seen_seconds_ago": round(age, 1),
        "status": node.status,
        "queue_depth": node.queue_depth,
        "total_requests": node.total_requests,
        "heartbeat_success_rate": hb_rate,
        "inference_success_rate": inf_rate,
        "capacity": _node_capacity_summary(node),
    }


def _coverage_map_detailed(
    nodes: list["_NodeEntry"],
    required_start: int,
    required_end: int,
    heartbeat_timeout: float,
) -> list[dict]:
    """Like _coverage_map but with per-node identity and health in each band.

    Includes all nodes (alive and stale) so operators can see both coverage
    holes and which dead nodes used to fill them.
    """
    now = time.monotonic()

    def covers(node: "_NodeEntry", layer: int) -> bool:
        return (
            node.shard_start is not None
            and node.shard_end is not None
            and node.shard_start <= layer <= node.shard_end
        )

    # Build per-layer list of nodes (with alive flag)
    layer_nodes: list[list[tuple["_NodeEntry", bool]]] = []
    for layer in range(required_start, required_end + 1):
        ns = [
            (n, (now - n.last_heartbeat) <= heartbeat_timeout)
            for n in nodes
            if covers(n, layer)
        ]
        layer_nodes.append(ns)

    # Merge consecutive layers with identical node-set (same node_ids, same alive)
    coverage: list[dict] = []
    for i, layer_idx in enumerate(range(required_start, required_end + 1)):
        ns = layer_nodes[i]
        node_ids = [n.node_id for n, _ in ns]
        alive_flags = [a for _, a in ns]
        # Compare with last band
        if (
            coverage
            and coverage[-1]["end_layer"] == layer_idx - 1
            and [nd["node_id"] for nd in coverage[-1]["nodes"]] == node_ids
            and [nd["alive"] for nd in coverage[-1]["nodes"]] == alive_flags
        ):
            coverage[-1]["end_layer"] = layer_idx
        else:
            coverage.append({
                "start_layer": layer_idx,
                "end_layer": layer_idx,
                "node_count": len(ns),
                "nodes": [_node_health(n, heartbeat_timeout) for n, _ in ns],
            })
    return coverage


def _coverage_gaps(coverage: list[dict]) -> list[tuple[int, int]]:
    return [
        (segment["start_layer"], segment["end_layer"])
        for segment in coverage
        if segment["node_count"] == 0
    ]


def _relay_http_request(
    relay_addr: str,
    path: str,
    body: bytes,
    headers: dict[str, str],
    timeout: float = 310.0,
) -> dict | None:
    """Send an HTTP-shaped request through a relay RPC WebSocket."""
    try:
        import websockets.sync.client as wsc  # type: ignore[import]

        request_id = str(uuid.uuid4())
        with wsc.connect(relay_addr, open_timeout=10, close_timeout=5) as ws:
            ws.send(json.dumps({
                "request_id": request_id,
                "method": "POST",
                "path": path,
                "headers": headers,
                "body": body.decode(errors="replace"),
            }))
            raw = ws.recv(timeout=timeout)
        response = json.loads(raw)
        if response.get("request_id") not in {None, request_id}:
            return None
        return response
    except Exception:
        return None


def _nodes_and_bounds_for_model(
    server: "_TrackerHTTPServer",
    model: str,
) -> tuple[list[_NodeEntry], int, int] | None:
    preset = server.model_presets.get(model)
    if preset is not None:
        required_start, required_end = _preset_layer_bounds(preset)
        return [node for node in server.registry.values() if node.model == model], required_start, required_end

    nodes = [
        node for node in server.registry.values()
        if _node_matches_model(node, model)
        and node.shard_start is not None
        and node.shard_end is not None
        and node.num_layers is not None
    ]
    if not nodes:
        return None
    return nodes, 0, max(node.num_layers for node in nodes) - 1


def _load_directive(node: _NodeEntry, model: str, start: int, end: int, quantization: str) -> dict:
    return {
        "action": "LOAD_SHARD",
        "model": model,
        "start_layer": start,
        "end_layer": end,
        "shard_start": start,
        "shard_end": end,
        "quantization": quantization,
    }


def _drop_directive(node: _NodeEntry, model: str, start: int, end: int, quantization: str) -> dict:
    return {
        "action": "DROP_SHARD",
        "model": model,
        "start_layer": start,
        "end_layer": end,
        "shard_start": start,
        "shard_end": end,
        "quantization": quantization,
    }


def _purge_expired_nodes_locked(server: "_TrackerHTTPServer") -> list[str]:
    now = time.monotonic()
    expired_ids = [
        node_id for node_id, entry in server.registry.items()
        if (now - entry.last_heartbeat) > server.heartbeat_timeout
    ]
    for node_id in expired_ids:
        entry = server.registry.pop(node_id)
        print(
            f"[tracker] node expired: {node_id}  {entry.endpoint}  "
            f"(no heartbeat for >{server.heartbeat_timeout:.0f}s)",
            flush=True,
        )
    if expired_ids:
        _rebalance_all_locked(server)
    return expired_ids


def _rebalance_model_locked(server: "_TrackerHTTPServer", model: str) -> None:
    preset = server.model_presets.get(model)
    if preset is None:
        return
    required_start, required_end = _preset_layer_bounds(preset)
    total_layers = required_end - required_start + 1
    model_nodes = [node for node in server.registry.values() if node.model == model]
    managed_nodes = [node for node in model_nodes if node.managed_assignment]
    if not managed_nodes:
        return

    previous_ranges = {
        node.node_id: (node.shard_start, node.shard_end, node.quantization)
        for node in managed_nodes
    }
    for node in managed_nodes:
        node.shard_start = None
        node.shard_end = None

    managed_nodes.sort(
        key=lambda node: (
            -node.benchmark_tokens_per_sec,
            -_node_layer_capacity(node, preset),
            node.node_id,
        )
    )
    base_nodes = [node for node in model_nodes if not node.managed_assignment]
    coverage = _coverage_map(base_nodes, required_start, required_end)
    gaps = _coverage_gaps(coverage)
    if not gaps:
        gaps = [(required_start, required_end)]

    eligible_nodes = [
        node for node in managed_nodes
        if _node_layer_capacity(node, preset) > 0
    ]
    node_index = 0
    for gap_start, gap_end in gaps:
        cursor = gap_start
        while cursor <= gap_end and node_index < len(eligible_nodes):
            node = eligible_nodes[node_index]
            remaining_layers = gap_end - cursor + 1
            remaining_nodes_after = len(eligible_nodes) - node_index - 1
            capacity = min(
                _node_layer_capacity(node, preset),
                total_layers,
                max(1, remaining_layers - remaining_nodes_after),
            )
            if capacity <= 0:
                node_index += 1
                continue
            quantization = _node_quantization(node, preset)
            node.quantization = quantization
            node.shard_start = cursor
            node.shard_end = min(gap_end, cursor + capacity - 1)
            cursor = node.shard_end + 1
            node_index += 1

    for node in managed_nodes:
        previous_start, previous_end, previous_quantization = previous_ranges[node.node_id]
        current_range = (node.shard_start, node.shard_end, node.quantization)
        if node.shard_start is None or node.shard_end is None or current_range == previous_ranges[node.node_id]:
            continue
        if previous_start is not None and previous_end is not None:
            node.pending_directives.append(
                _drop_directive(
                    node,
                    model,
                    previous_start,
                    previous_end,
                    previous_quantization or _node_quantization(node, preset),
                )
            )
        node.pending_directives.append(
            _load_directive(
                node,
                model,
                node.shard_start,
                node.shard_end,
                node.quantization or _node_quantization(node, preset),
            )
        )


def _hf_rebalance_preset(nodes: list[_NodeEntry]) -> dict:
    total_layers = max(node.num_layers or 0 for node in nodes)
    return {
        "layers_start": 0,
        "layers_end": total_layers - 1,
        "bytes_per_layer": {"bfloat16": 30 * 1024 * 1024, "int8": 15 * 1024 * 1024, "nf4": 8 * 1024 * 1024},
    }


def _rebalance_hf_model_locked(server: "_TrackerHTTPServer", hf_repo: str) -> None:
    model_nodes = [
        node for node in server.registry.values()
        if node.hf_repo == hf_repo
        and node.num_layers is not None
    ]
    managed_nodes = [node for node in model_nodes if node.managed_assignment]
    if not model_nodes or not managed_nodes:
        return

    preset = _hf_rebalance_preset(model_nodes)
    required_start, required_end = _preset_layer_bounds(preset)
    total_layers = required_end - required_start + 1
    if total_layers <= 0:
        return

    previous_ranges = {
        node.node_id: (node.shard_start, node.shard_end, node.quantization)
        for node in managed_nodes
    }
    for node in managed_nodes:
        node.shard_start = None
        node.shard_end = None

    managed_nodes.sort(
        key=lambda node: (
            -node.benchmark_tokens_per_sec,
            -_node_layer_capacity(node, preset),
            node.node_id,
        )
    )
    base_nodes = [node for node in model_nodes if not node.managed_assignment]
    coverage = _coverage_map(base_nodes, required_start, required_end)
    gaps = _coverage_gaps(coverage)
    if not gaps:
        gaps = [(required_start, required_end)]

    eligible_nodes = [
        node for node in managed_nodes
        if _node_layer_capacity(node, preset) > 0
    ]
    node_index = 0
    for gap_start, gap_end in gaps:
        cursor = gap_start
        while cursor <= gap_end and node_index < len(eligible_nodes):
            node = eligible_nodes[node_index]
            remaining_layers = gap_end - cursor + 1
            remaining_nodes_after = len(eligible_nodes) - node_index - 1
            capacity = min(
                _node_layer_capacity(node, preset),
                total_layers,
                max(1, remaining_layers - remaining_nodes_after),
            )
            if capacity <= 0:
                node_index += 1
                continue
            quantization = _node_quantization(node, preset)
            node.quantization = quantization
            node.shard_start = cursor
            node.shard_end = min(gap_end, cursor + capacity - 1)
            cursor = node.shard_end + 1
            node_index += 1

    for node in managed_nodes:
        previous_start, previous_end, previous_quantization = previous_ranges[node.node_id]
        current_range = (node.shard_start, node.shard_end, node.quantization)
        if node.shard_start is None or node.shard_end is None or current_range == previous_ranges[node.node_id]:
            continue
        if previous_start is not None and previous_end is not None:
            node.pending_directives.append(
                _drop_directive(
                    node,
                    hf_repo,
                    previous_start,
                    previous_end,
                    previous_quantization or _node_quantization(node, preset),
                )
            )
        node.pending_directives.append(
            _load_directive(
                node,
                hf_repo,
                node.shard_start,
                node.shard_end,
                node.quantization or _node_quantization(node, preset),
            )
        )


def _rebalance_all_locked(server: "_TrackerHTTPServer") -> None:
    for model in list(server.model_presets):
        _rebalance_model_locked(server, model)
    for hf_repo in sorted({node.hf_repo for node in server.registry.values() if node.hf_repo}):
        _rebalance_hf_model_locked(server, hf_repo)


def _registration_ban_error(contracts: Any | None, wallet_address: str | None) -> str | None:
    if contracts is None or not wallet_address:
        return None
    if contracts.registry.get_wallet(wallet_address).banned:
        return "wallet is banned"
    return None


def _node_id_for_registration(
    endpoint: str,
    model: str,
    wallet_address: str | None,
    shard_start: int | None,
    shard_end: int | None,
    hf_repo: str | None,
) -> str:
    wallet_prefix = wallet_address[:8] if wallet_address else "anon"
    stable_key = "|".join([
        wallet_address or "",
        endpoint.rstrip("/"),
        model,
        hf_repo or "",
        "" if shard_start is None else str(shard_start),
        "" if shard_end is None else str(shard_end),
    ])
    digest = hashlib.sha256(stable_key.encode()).hexdigest()[:12]
    return f"{wallet_prefix}-{digest}"


class _TrackerHTTPServer(socketserver.ThreadingMixIn, http.server.HTTPServer):
    daemon_threads = True

    def __init__(
        self,
        addr: tuple,
        handler,
        registry: dict,
        lock: threading.Lock,
        heartbeat_timeout: float,
        model_presets: dict,
        contracts: Any | None,
        minimum_stake: int,
        relay_url: str | None = None,
        raft: "RaftNode | None" = None,
        gossip: "NodeGossip | None" = None,
        stats: "_StatsCollector | None" = None,
    ) -> None:
        super().__init__(addr, handler)
        self.registry = registry
        self.lock = lock
        self.heartbeat_timeout = heartbeat_timeout
        self.model_presets = model_presets
        self.contracts = contracts
        self.minimum_stake = minimum_stake
        self.relay_url = relay_url.rstrip("/") if relay_url else None
        self.raft = raft
        self.gossip = gossip
        self.stats: _StatsCollector | None = stats


class _TrackerHandler(http.server.BaseHTTPRequestHandler):
    def log_message(self, fmt, *args):  # noqa: suppress request logs in tests
        pass

    def _send_json(self, status: int, data: dict) -> None:
        body = json.dumps(data).encode()
        self.send_response(status)
        self.send_header("Content-Type", "application/json")
        self.send_header("Content-Length", str(len(body)))
        self.end_headers()
        try:
            self.wfile.write(body)
        except BrokenPipeError:
            pass

    def _read_json_body(self) -> dict | None:
        length = int(self.headers.get("Content-Length", 0))
        try:
            body = json.loads(self.rfile.read(length) or b"{}")
        except json.JSONDecodeError:
            self._send_json(400, {"error": "invalid JSON body"})
            return None
        if not isinstance(body, dict):
            self._send_json(400, {"error": "JSON body must be an object"})
            return None
        return body

    def _purge_expired_nodes(self) -> None:
        server: _TrackerHTTPServer = self.server  # type: ignore[assignment]
        _purge_expired_nodes_locked(server)

    def do_POST(self):
        if self.path == "/v1/chat/completions":
            self._handle_proxy_chat()
            return
        if self.path == "/v1/nodes/register":
            self._handle_register()
            return
        if self.path == "/v1/raft/vote":
            self._handle_raft_vote()
            return
        if self.path == "/v1/raft/append":
            self._handle_raft_append()
            return
        if self.path == "/v1/gossip":
            self._handle_gossip()
            return
        if self.path == "/v1/stats/gossip":
            self._handle_stats_gossip()
            return
        parts = self.path.split("/")
        # /v1/nodes/<node_id>/heartbeat -> ['', 'v1', 'nodes', '<id>', 'heartbeat']
        if len(parts) == 5 and parts[1] == "v1" and parts[2] == "nodes" and parts[4] == "heartbeat":
            self._handle_heartbeat(parts[3])
            return
        self.send_response(404)
        self.end_headers()

    def do_GET(self):
        parsed = urllib.parse.urlparse(self.path)
        if parsed.path == "/v1/route":
            self._handle_route(parsed)
        elif parsed.path == "/v1/routes":
            self._handle_routes(parsed)
        elif parsed.path == "/v1/nodes/assign":
            self._handle_assign(parsed)
        elif parsed.path == "/v1/network/assign":
            self._handle_network_assign(parsed)
        elif parsed.path == "/v1/network/map":
            self._handle_network_map()
        elif parsed.path == "/v1/models":
            self._handle_models()
        elif parsed.path.startswith("/v1/coverage/"):
            model = urllib.parse.unquote(parsed.path.removeprefix("/v1/coverage/"))
            self._handle_coverage(model)
        elif parsed.path.startswith("/v1/tracker-nodes/"):
            model = urllib.parse.unquote(parsed.path.removeprefix("/v1/tracker-nodes/"))
            self._handle_tracker_nodes(model)
        elif parsed.path == "/v1/raft/status":
            self._handle_raft_status()
        elif parsed.path == "/v1/stats":
            self._handle_stats()
        elif parsed.path == "/v1/health":
            self._send_json(200, {"status": "ok"})
        else:
            self.send_response(404)
            self.end_headers()

    def _handle_models(self):
        server: _TrackerHTTPServer = self.server  # type: ignore[assignment]
        created = int(time.time())
        with server.lock:
            self._purge_expired_nodes()
            alive = list(server.registry.values())
        if server.contracts is not None:
            alive = [
                node for node in alive
                if not node.wallet_address or not server.contracts.registry.get_wallet(node.wallet_address).banned
            ]
        data = []
        seen_ids: set[str] = set()
        for name, preset in server.model_presets.items():
            model_nodes = [node for node in alive if node.model == name]
            if not model_nodes:
                continue
            required_start, required_end = _preset_layer_bounds(preset)
            coverage = _coverage_percentage(
                model_nodes,
                required_start,
                required_end,
            )
            aliases = [name]
            hf_repo = preset.get("hf_repo")
            if hf_repo and hf_repo not in aliases:
                aliases.append(hf_repo)
            data.append({
                "id": name,
                "object": "model",
                "created": created,
                "owned_by": "meshnet",
                "name": name,
                "hf_repo": hf_repo,
                "aliases": aliases,
                "shard_coverage_percentage": coverage,
            })
            seen_ids.add(name)

        hf_model_ids = sorted({
            node.hf_repo or node.model
            for node in alive
            if node.model is not None
            and node.model not in server.model_presets
            and node.shard_start is not None
            and node.shard_end is not None
            and node.num_layers is not None
        })
        for model_id in hf_model_ids:
            if model_id is None or model_id in seen_ids:
                continue
            model_nodes = [
                node for node in alive
                if node.shard_start is not None
                and node.shard_end is not None
                and node.num_layers is not None
                and (node.hf_repo == model_id or (node.hf_repo is None and node.model == model_id))
            ]
            if not model_nodes:
                continue
            short_names = sorted({node.model for node in model_nodes if node.model})
            aliases = [model_id, *[name for name in short_names if name != model_id]]
            required_start = 0
            required_end = max(node.num_layers for node in model_nodes) - 1
            data.append({
                "id": model_id,
                "object": "model",
                "created": created,
                "owned_by": "meshnet",
                "name": short_names[0] if short_names else model_id,
                "hf_repo": model_id if any(node.hf_repo == model_id for node in model_nodes) else None,
                "aliases": aliases,
                "shard_coverage_percentage": _coverage_percentage(
                    model_nodes,
                    required_start,
                    required_end,
                ),
            })
            seen_ids.add(model_id)
        self._send_json(200, {"object": "list", "data": data})

    def _handle_coverage(self, model: str):
        server: _TrackerHTTPServer = self.server  # type: ignore[assignment]
        # Do NOT purge before coverage — dead nodes are included with alive=false
        # so operators can see what was covering each layer band before failure.
        with server.lock:
            resolved = _nodes_and_bounds_for_model(server, model)
        if resolved is None:
            self._send_json(404, {"error": f"no nodes registered for model {model!r}"})
            return
        all_nodes, required_start, required_end = resolved
        if server.contracts is not None:
            all_nodes = [
                node for node in all_nodes
                if not node.wallet_address or not server.contracts.registry.get_wallet(node.wallet_address).banned
            ]
        coverage = _coverage_map_detailed(all_nodes, required_start, required_end, server.heartbeat_timeout)
        self._send_json(200, {"model": model, "coverage": coverage})

    def _handle_tracker_nodes(self, model: str):
        """Return nodes registered with tracker_mode=True whose shard starts at layer 0."""
        server: _TrackerHTTPServer = self.server  # type: ignore[assignment]
        preset = server.model_presets.get(model)
        if preset is None:
            self._send_json(404, {"error": f"unknown model preset: {model!r}"})
            return
        required_start, _ = _preset_layer_bounds(preset)
        with server.lock:
            self._purge_expired_nodes()
            alive = [node for node in server.registry.values() if node.model == model]
        if server.contracts is not None:
            alive = [
                node for node in alive
                if not node.wallet_address or not server.contracts.registry.get_wallet(node.wallet_address).banned
            ]
        tracker_nodes = [
            node for node in alive
            if node.shard_start is not None
            and node.shard_start == required_start
            and node.tracker_mode
        ]
        self._send_json(200, {
            "model": model,
            "tracker_nodes": [
                {
                    "node_id": node.node_id,
                    "endpoint": node.endpoint,
                    "relay_addr": node.relay_addr,
                    "peer_id": node.peer_id,
                    "benchmark_tokens_per_sec": node.benchmark_tokens_per_sec,
                }
                for node in tracker_nodes
            ],
        })

    def _handle_network_map(self) -> None:
        server: _TrackerHTTPServer = self.server  # type: ignore[assignment]
        with server.lock:
            self._purge_expired_nodes()
            nodes = list(server.registry.values())

        def capacity_for(node: _NodeEntry) -> dict:
            preset = None
            if node.model:
                preset = server.model_presets.get(node.model)
            if preset is None and node.hf_repo and node.num_layers:
                preset = _hf_rebalance_preset([node])
            return _node_capacity_summary(node, preset)

        self._send_json(200, {
            "relay_url": server.relay_url,
            "nodes": [
                {
                    "node_id": node.node_id,
                    "endpoint": node.endpoint,
                    "relay_addr": node.relay_addr,
                    "peer_id": node.peer_id,
                    "model": node.model,
                    "hf_repo": node.hf_repo,
                    "shard_start": node.shard_start,
                    "shard_end": node.shard_end,
                    "tracker_mode": node.tracker_mode,
                    "last_heartbeat": node.last_heartbeat,
                    "capacity": capacity_for(node),
                    "stats": _node_health(node, server.heartbeat_timeout),
                }
                for node in nodes
            ],
        })

    # ---------------------------------------------------------------- OpenAI proxy

    def _handle_proxy_chat(self) -> None:
        """Proxy POST /v1/chat/completions to a tracker-mode (first-shard) node.

        Picks a live tracker-mode node for the requested model using round-robin,
        then forwards the request verbatim and relays the response (including
        streaming SSE chunks) back to the caller.
        """
        server: _TrackerHTTPServer = self.server  # type: ignore[assignment]

        length = int(self.headers.get("Content-Length", 0))
        raw_body = self.rfile.read(length) if length else b"{}"

        try:
            body = json.loads(raw_body)
        except json.JSONDecodeError:
            self._send_json(400, {"error": {"message": "invalid JSON", "type": "invalid_request_error", "code": "invalid_request"}})
            return

        model: str = body.get("model", "")
        is_stream: bool = bool(body.get("stream", False))

        if model and server.stats is not None:
            server.stats.record_request(model)

        # Find a live tracker-mode node for this model
        with server.lock:
            self._purge_expired_nodes()
            candidates = [
                n for n in server.registry.values()
                if n.tracker_mode and _node_matches_model(n, model)
            ]

        if not candidates:
            # Fall back: any node serving shard_start=0 for this model
            with server.lock:
                candidates = [
                    n for n in server.registry.values()
                    if n.shard_start == 0 and _node_matches_model(n, model)
                ]

        if not candidates:
            self._send_json(503, {"error": {
                "message": f"no nodes available for model {model!r}",
                "type": "service_unavailable",
                "code": "model_not_available",
            }})
            return

        # Simple round-robin via list length modulo (stateless, good enough)
        node = candidates[int(time.time() * 1000) % len(candidates)]
        target_url = f"{node.endpoint}/v1/chat/completions"
        request_id = str(body.get("id") or f"req-{time.time_ns():x}")

        # Pre-resolve the downstream route so the first-shard node skips its own
        # tracker query.  We already hold the full registry picture — no need for
        # a second round-trip.
        route_model = node.hf_repo or node.model or model
        with server.lock:
            if route_model in server.model_presets:
                preset = server.model_presets[route_model]
                rs, re = _preset_layer_bounds(preset)
                all_nodes: list = [n for n in server.registry.values() if n.model == route_model]
            else:
                all_nodes = [
                    n for n in server.registry.values()
                    if _node_matches_model(n, route_model)
                    and n.shard_start is not None and n.num_layers is not None
                ]
                rs, re = 0, (max((n.num_layers for n in all_nodes), default=1) - 1)
        route_nodes, _ = _select_route(all_nodes, rs, re)
        # Compute start_layer for each hop: each node begins where the previous ended + 1.
        # This allows overlapping shard registrations without double-computation.
        covered_up_to = rs - 1
        route_hops: list[dict] = []
        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
        # 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)
        route_debug = " -> ".join(
            f"{n.node_id}@{n.endpoint}[{n.shard_start}-{n.shard_end}]"
            for n in route_nodes
        )
        print(
            f"[tracker] proxy route {request_id}: model={model!r} "
            f"head={node.node_id}@{node.endpoint} downstream={downstream_urls} "
            f"route={route_debug or '<empty>'}",
            flush=True,
        )

        req = urllib.request.Request(
            target_url,
            data=raw_body,
            headers={
                "Content-Type": "application/json",
                "X-Meshnet-Route": downstream_urls,
            },
            method="POST",
        )
        # Copy Authorization header from client if present
        auth = self.headers.get("Authorization")
        if auth:
            req.add_header("Authorization", auth)

        relay_headers = {
            "Content-Type": "application/json",
            "X-Meshnet-Route": downstream_urls,
            **({"Authorization": auth} if auth else {}),
        }

        if node.relay_addr:
            print(
                f"[tracker] proxy via relay {request_id}: {node.relay_addr} "
                f"(direct endpoint {target_url})",
                flush=True,
            )
            relayed = _relay_http_request(
                node.relay_addr,
                path="/v1/chat/completions",
                body=raw_body,
                headers=relay_headers,
            )
            if relayed is not None:
                self._send_relayed_response(relayed)
                return
            print(
                f"[tracker] relay proxy failed {request_id}: {node.relay_addr}; "
                f"trying direct {target_url}",
                flush=True,
            )

        try:
            upstream = urllib.request.urlopen(req, timeout=300.0)
            print(f"[tracker] proxy connected {request_id}: {target_url}", flush=True)
        except urllib.error.HTTPError as exc:
            # Relay error status + body from node
            err_body = exc.read()
            self.send_response(exc.code)
            self.send_header("Content-Type", "application/json")
            self.send_header("Content-Length", str(len(err_body)))
            self.end_headers()
            try:
                self.wfile.write(err_body)
            except BrokenPipeError:
                pass
            return
        except Exception as exc:
            if node.relay_addr:
                print(
                    f"[tracker] direct proxy failed {request_id}: {target_url}: {exc}; "
                    f"relay already attempted for {node.relay_addr}",
                    flush=True,
                )
            else:
                print(f"[tracker] proxy failed {request_id}: {target_url}: {exc}", flush=True)
            self._send_json(503, {"error": {
                "message": f"upstream node unreachable: {exc}",
                "type": "service_unavailable",
                "code": "upstream_error",
            }})
            return

        with upstream:
            content_type = upstream.headers.get("Content-Type", "application/json")
            if is_stream or "text/event-stream" in content_type:
                # Relay SSE stream chunk-by-chunk
                self.send_response(200)
                self.send_header("Content-Type", "text/event-stream; charset=utf-8")
                self.send_header("Cache-Control", "no-cache")
                self.end_headers()
                try:
                    while True:
                        line = upstream.readline()
                        if not line:
                            break
                        self.wfile.write(line)
                        self.wfile.flush()
                except BrokenPipeError:
                    pass
            else:
                # Non-streaming: buffer and relay
                resp_body = upstream.read()
                print(
                    f"[tracker] proxy complete {request_id}: {target_url} bytes={len(resp_body)}",
                    flush=True,
                )
                self.send_response(200)
                self.send_header("Content-Type", content_type)
                self.send_header("Content-Length", str(len(resp_body)))
                self.end_headers()
                try:
                    self.wfile.write(resp_body)
                except BrokenPipeError:
                    pass

    def _send_relayed_response(self, response: dict) -> None:
        status = int(response.get("status", 503))
        headers = response.get("headers") if isinstance(response.get("headers"), dict) else {}
        body_text = response.get("body") or ""
        body = body_text.encode() if isinstance(body_text, str) else bytes(body_text)
        self.send_response(status)
        self.send_header("Content-Type", headers.get("Content-Type", "application/json"))
        self.send_header("Content-Length", str(len(body)))
        self.end_headers()
        try:
            self.wfile.write(body)
        except BrokenPipeError:
            pass

    def _handle_register(self):
        server: _TrackerHTTPServer = self.server  # type: ignore[assignment]
        body = self._read_json_body()
        if body is None:
            return

        # --- Raft cluster mode: forward to leader or propose via Raft ---
        if server.raft is not None:
            if not server.raft.is_leader:
                leader = server.raft.leader()
                if leader is None:
                    self._send_json(503, {"error": "no leader elected — retry in a moment"})
                    return
                # Proxy to leader
                try:
                    data = json.dumps(body).encode()
                    req = urllib.request.Request(
                        f"{leader}/v1/nodes/register",
                        data=data,
                        headers={"Content-Type": "application/json"},
                        method="POST",
                    )
                    with urllib.request.urlopen(req, timeout=5.0) as r:
                        resp_body = r.read()
                    self.send_response(200)
                    self.send_header("Content-Type", "application/json")
                    self.send_header("Content-Length", str(len(resp_body)))
                    self.end_headers()
                    self.wfile.write(resp_body)
                except Exception as exc:
                    self._send_json(503, {"error": f"leader proxy failed: {exc}"})
                return
            # Leader path: fall through to normal registration, then replicate via Raft.
            # We let the registration run first (to generate node_id), then asynchronously
            # propose to the Raft log so followers can replicate the entry.
            # Raft proposal happens after the response is sent (fire-and-forget replication).

        endpoint = body.get("endpoint")
        if not isinstance(endpoint, str) or not endpoint:
            self._send_json(400, {"error": "endpoint is required"})
            return
        parsed_endpoint = urllib.parse.urlparse(endpoint)
        if parsed_endpoint.scheme not in {"http", "https"} or not parsed_endpoint.netloc:
            self._send_json(400, {"error": "endpoint must be an http(s) URL"})
            return

        shard_start: int | None
        shard_end: int | None
        explicit_shard = "shard_start" in body or "shard_end" in body
        if explicit_shard:
            try:
                shard_start = int(body["shard_start"])
                shard_end = int(body["shard_end"])
            except (KeyError, TypeError, ValueError):
                self._send_json(400, {"error": "shard_start and shard_end must be numeric"})
                return
            if shard_start < 0 or shard_end < 0 or shard_start > shard_end:
                self._send_json(400, {"error": "shard range must be non-negative and ordered"})
                return
        else:
            shard_start = None
            shard_end = None
        try:
            score = float(body.get("score", 1.0))
        except (TypeError, ValueError):
            self._send_json(400, {"error": "score must be numeric"})
            return

        hardware_profile = body.get("hardware_profile", {})
        if not isinstance(hardware_profile, dict):
            self._send_json(400, {"error": "hardware_profile must be an object"})
            return
        model = body.get("model")
        if model is None:
            model = "stub-model"
        if not isinstance(model, str):
            self._send_json(400, {"error": "model must be a string"})
            return
        shard_checksum = body.get("shard_checksum")
        if shard_checksum is not None and not isinstance(shard_checksum, str):
            self._send_json(400, {"error": "shard_checksum must be a string"})
            return
        try:
            vram_bytes = int(body.get("vram_bytes", DEFAULT_VRAM_BYTES))
            ram_bytes = int(body.get("ram_bytes", DEFAULT_RAM_BYTES))
            max_loaded_shards = int(body.get("max_loaded_shards", 1))
            benchmark_tokens_per_sec = float(
                body.get("benchmark_tokens_per_sec", DEFAULT_BENCHMARK_TOKENS_PER_SEC)
            )
        except (TypeError, ValueError):
            self._send_json(400, {"error": "vram_bytes, ram_bytes, max_loaded_shards, and benchmark_tokens_per_sec must be numeric"})
            return
        if vram_bytes < 0 or ram_bytes < 0 or max_loaded_shards < 1 or benchmark_tokens_per_sec <= 0:
            self._send_json(400, {"error": "capability values must be positive"})
            return
        quantizations_body = body.get("quantizations", DEFAULT_QUANTIZATIONS)
        if not (
            isinstance(quantizations_body, list)
            and quantizations_body
            and all(isinstance(item, str) and item for item in quantizations_body)
        ):
            self._send_json(400, {"error": "quantizations must be a non-empty string array"})
            return
        quantizations = list(quantizations_body)
        quantization = body.get("quantization")
        if quantization is not None and not isinstance(quantization, str):
            self._send_json(400, {"error": "quantization must be a string"})
            return
        wallet_address = body.get("wallet_address")
        if wallet_address is not None and not isinstance(wallet_address, str):
            self._send_json(400, {"error": "wallet_address must be a string"})
            return
        ban_error = _registration_ban_error(server.contracts, wallet_address)
        if ban_error:
            self._send_json(403, {"error": ban_error})
            return

        tracker_mode = bool(body.get("tracker_mode", False))
        managed_assignment = bool(body.get("managed_assignment", False))
        hf_repo = body.get("hf_repo")
        if hf_repo is not None and not isinstance(hf_repo, str):
            self._send_json(400, {"error": "hf_repo must be a string"})
            return
        num_layers_body = body.get("num_layers")
        num_layers: int | None = None
        if num_layers_body is not None:
            try:
                num_layers = int(num_layers_body)
            except (TypeError, ValueError):
                self._send_json(400, {"error": "num_layers must be an integer"})
                return
        relay_addr = body.get("relay_addr") or None
        cert_fingerprint = body.get("cert_fingerprint") or None
        peer_id = body.get("peer_id") or None

        node_id = _node_id_for_registration(
            endpoint,
            model,
            wallet_address,
            shard_start,
            shard_end,
            hf_repo,
        )
        entry = _NodeEntry(
            node_id=node_id,
            endpoint=endpoint.rstrip("/"),
            shard_start=shard_start,
            shard_end=shard_end,
            model=model,
            shard_checksum=shard_checksum,
            hardware_profile=hardware_profile,
            wallet_address=wallet_address,
            score=score,
            vram_bytes=vram_bytes,
            ram_bytes=ram_bytes,
            quantizations=quantizations,
            max_loaded_shards=max_loaded_shards,
            benchmark_tokens_per_sec=benchmark_tokens_per_sec,
            quantization=quantization,
            managed_assignment=managed_assignment or not explicit_shard,
            tracker_mode=tracker_mode,
            hf_repo=hf_repo,
            num_layers=num_layers,
            relay_addr=relay_addr,
            cert_fingerprint=cert_fingerprint,
            peer_id=peer_id,
        )
        with server.lock:
            self._purge_expired_nodes()
            # Dedup: if this endpoint is already registered, remove the old entry first.
            stale_ids = [
                eid for eid, e in server.registry.items()
                if e.endpoint == entry.endpoint.rstrip("/")
            ]
            for eid in stale_ids:
                old = server.registry.pop(eid)
                print(
                    f"[tracker] node re-registered: replaced {eid} with {node_id}"
                    f"  {old.endpoint}",
                    flush=True,
                )
            server.registry[node_id] = entry
            if entry.managed_assignment:
                if entry.hf_repo:
                    _rebalance_hf_model_locked(server, entry.hf_repo)
                else:
                    _rebalance_model_locked(server, model)
            assignment_directive = entry.pending_directives[-1] if entry.pending_directives else None
            if assignment_directive is not None:
                entry.pending_directives.clear()

        shard_info = f"layers {shard_start}-{shard_end}" if shard_start is not None else "unsharded"
        repo_info = f" [{hf_repo}]" if hf_repo else ""
        budget_bytes, budget_source = _node_memory_budget_bytes(entry)
        budget_gb = budget_bytes / (1024 ** 3)
        print(
            f"[tracker] node registered: {node_id}  {endpoint}  {model}{repo_info}  {shard_info}  "
            f"capacity={budget_gb:.1f}GB {budget_source} slots={max_loaded_shards}",
            flush=True,
        )

        payload = {"node_id": node_id}
        if assignment_directive is not None:
            payload["directive"] = assignment_directive
        self._send_json(200, payload)

        # Raft replication: leader proposes this registration to followers so their
        # registries stay in sync.  Fire-and-forget (async) — the client already
        # got its node_id; replication happens in the background.
        if server.raft is not None and server.raft.is_leader:
            raft_payload = dict(body)
            raft_payload["node_id"] = node_id  # include the generated ID
            threading.Thread(
                target=server.raft.propose,
                args=("register", raft_payload),
                daemon=True,
            ).start()

    def _handle_heartbeat(self, node_id: str):
        server: _TrackerHTTPServer = self.server  # type: ignore[assignment]
        body: dict = {}
        content_length = int(self.headers.get("Content-Length", 0))
        if content_length > 0:
            try:
                body = json.loads(self.rfile.read(content_length))
            except Exception:
                pass
        with server.lock:
            self._purge_expired_nodes()
            entry = server.registry.get(node_id)
            if entry is None:
                self._send_json(404, {"error": "node not found"})
                return
            entry.last_heartbeat = time.monotonic()
            entry.heartbeats_received += 1
            # P2P metadata
            if body.get("relay_addr"):
                entry.relay_addr = body["relay_addr"]
            if body.get("cert_fingerprint"):
                entry.cert_fingerprint = body["cert_fingerprint"]
            if body.get("peer_id"):
                entry.peer_id = body["peer_id"]
            # Node stats (cumulative — node always sends totals, tracker replaces)
            if "total_requests" in body:
                entry.total_requests = int(body["total_requests"])
            if "failed_requests" in body:
                entry.failed_requests = int(body["failed_requests"])
            if "queue_depth" in body:
                entry.queue_depth = int(body["queue_depth"])
            if "uptime_seconds" in body:
                entry.uptime_seconds = float(body["uptime_seconds"])
            if "status" in body and body["status"] in ("ready", "loading"):
                entry.status = body["status"]
            if entry.hf_repo:
                _rebalance_hf_model_locked(server, entry.hf_repo)
            else:
                _rebalance_model_locked(server, entry.model or "stub-model")
            directives = list(entry.pending_directives)
            entry.pending_directives.clear()
            new_assignment = entry.pending_new_assignment
            if new_assignment is not None:
                entry.pending_new_assignment = None
        if server.gossip is not None:
            server.gossip.record(node_id)
        resp: dict = {}
        if directives:
            resp["directives"] = directives
        if new_assignment is not None:
            resp["new_assignment"] = new_assignment
        self._send_json(200, resp)

    # ---------------------------------------------------------------- Raft handlers

    def _handle_raft_vote(self):
        server: _TrackerHTTPServer = self.server  # type: ignore[assignment]
        body = self._read_json_body()
        if body is None:
            return
        if server.raft is None:
            self._send_json(503, {"error": "raft not enabled"})
            return
        result = server.raft.handle_request_vote(body)
        self._send_json(200, result)

    def _handle_raft_append(self):
        server: _TrackerHTTPServer = self.server  # type: ignore[assignment]
        body = self._read_json_body()
        if body is None:
            return
        if server.raft is None:
            self._send_json(503, {"error": "raft not enabled"})
            return
        result = server.raft.handle_append_entries(body)
        self._send_json(200, result)

    def _handle_raft_status(self):
        server: _TrackerHTTPServer = self.server  # type: ignore[assignment]
        if server.raft is None:
            self._send_json(200, {"role": "standalone", "term": 0, "leader": None})
            return
        self._send_json(200, server.raft.status())

    def _handle_gossip(self):
        server: _TrackerHTTPServer = self.server  # type: ignore[assignment]
        body = self._read_json_body()
        if body is None:
            return
        if server.gossip is not None and isinstance(body, dict):
            server.gossip.merge({k: float(v) for k, v in body.items()})
        self._send_json(200, {})

    def _handle_stats(self):
        server: _TrackerHTTPServer = self.server  # type: ignore[assignment]
        if server.stats is None:
            self._send_json(200, {"models": {}})
            return
        self._send_json(200, {"models": server.stats.get_combined_stats()})

    def _handle_stats_gossip(self):
        server: _TrackerHTTPServer = self.server  # type: ignore[assignment]
        body = self._read_json_body()
        if body is None:
            return
        tracker_url = body.get("tracker_url", "")
        rpms = body.get("stats", {})
        if server.stats is not None and tracker_url and isinstance(rpms, dict):
            server.stats.merge_peer_rpms(tracker_url, rpms)
        self._send_json(200, {})

    def _handle_assign(self, parsed: urllib.parse.ParseResult):
        """Return an optimal shard assignment for a node given its hardware profile.

        Query params:
          model   — model preset name (default: first preset)
          device  — "cuda" | "cpu"
          vram_mb — integer VRAM in MB (0 for CPU)
          ram_mb  — integer system RAM in MB, used when vram_mb=0

        The greedy strategy: find the first gap in current layer coverage
        and assign it.  If no gap exists, assign the full model range so the
        node provides redundant coverage.
        """
        server: _TrackerHTTPServer = self.server  # type: ignore[assignment]
        params = urllib.parse.parse_qs(parsed.query)

        model_list = params.get("model")
        if not model_list:
            model = next(iter(server.model_presets), None)
            if model is None:
                self._send_json(503, {"error": "no model presets configured"})
                return
        else:
            model = model_list[0]

        preset = server.model_presets.get(model)
        if preset is None:
            self._send_json(404, {"error": f"unknown model preset: {model!r}"})
            return

        required_start, required_end = _preset_layer_bounds(preset)

        with server.lock:
            self._purge_expired_nodes()
            alive = [node for node in server.registry.values() if node.model == model]
        if server.contracts is not None:
            alive = [
                node for node in alive
                if not node.wallet_address or not server.contracts.registry.get_wallet(node.wallet_address).banned
            ]

        device = params.get("device", ["cpu"])[0]
        try:
            vram_mb = int(params.get("vram_mb", ["0"])[0])
        except ValueError:
            vram_mb = 0
        try:
            ram_mb = int(params.get("ram_mb", ["0"])[0])
        except ValueError:
            ram_mb = 0
        max_layers = required_end - required_start + 1
        memory_mb = vram_mb if vram_mb > 0 else ram_mb
        if memory_mb > 0:
            layer_bytes = _preset_bytes_per_layer(preset).get("bfloat16", 30 * 1024 * 1024)
            max_layers = min(max_layers, max(1, int(((memory_mb * 1024 * 1024) * 0.8) // layer_bytes)))
        elif device != "cuda" or vram_mb < 8192:
            max_layers = min(max_layers, 16)

        # Collect covered intervals sorted by start layer.
        covered = sorted(
            [
                (n.shard_start, n.shard_end)
                for n in alive
                if n.shard_start is not None and n.shard_end is not None
            ],
            key=lambda t: t[0],
        )

        # Walk from required_start to find the first uncovered layer.
        gap_start = required_start
        for s, e in covered:
            if s <= gap_start:
                gap_start = max(gap_start, e + 1)
            else:
                break  # gap found before this node

        if gap_start > required_end:
            # Full coverage exists — assign the full range for redundancy.
            shard_start = required_start
            shard_end = min(required_end, shard_start + max_layers - 1)
        else:
            shard_start = gap_start
            shard_end = min(required_end, shard_start + max_layers - 1)

        peers = [
            {"endpoint": node.endpoint, "checksum": node.shard_checksum}
            for node in alive
            if node.model == model
            and node.shard_start == shard_start
            and node.shard_end == shard_end
            and node.shard_checksum
        ]

        self._send_json(200, {
            "shard_start": shard_start,
            "shard_end": shard_end,
            "model": model,
            "model_layers_end": required_end,
            "peers": peers,
            **({"hf_repo": preset["hf_repo"]} if "hf_repo" in preset else {}),
        })

    def _handle_network_assign(self, parsed: urllib.parse.ParseResult):
        """Assign a new node to fill the biggest uncovered shard gap across HF-model nodes.

        Query params:
          vram_mb  — integer VRAM in MB (0 = CPU-only node)
          ram_mb   — integer system RAM in MB, used when vram_mb=0
          device   — "cuda" | "cpu"
          hf_repo  — optional; if set, restrict search to this repo only

        Returns:
          {hf_repo, shard_start, shard_end, num_layers, gap_found}
          gap_found=true means a real uncovered gap was assigned; false means redundancy.
        """
        server: _TrackerHTTPServer = self.server  # type: ignore[assignment]
        params = urllib.parse.parse_qs(parsed.query)
        try:
            vram_mb = int(params.get("vram_mb", ["0"])[0])
        except ValueError:
            vram_mb = 0
        try:
            ram_mb = int(params.get("ram_mb", ["0"])[0])
        except ValueError:
            ram_mb = 0
        device = params.get("device", ["cpu"])[0]
        filter_repo = params.get("hf_repo", [None])[0]  # optional repo filter

        with server.lock:
            self._purge_expired_nodes()
            all_nodes = list(server.registry.values())

        # Collect only nodes that registered a real HF model (have hf_repo + shard bounds).
        hf_nodes = [
            n for n in all_nodes
            if n.hf_repo
            and n.shard_start is not None
            and n.shard_end is not None
            and n.num_layers is not None
            and (filter_repo is None or n.hf_repo == filter_repo)
        ]

        if not hf_nodes:
            msg = (
                f"no HF-model nodes registered for {filter_repo!r}"
                if filter_repo
                else "no HF-model nodes registered; cannot assign shards"
            )
            self._send_json(503, {"error": msg})
            return

        # Group by hf_repo; pick the one with the largest total_layers and biggest gap.
        from collections import defaultdict
        repo_groups: dict = defaultdict(list)
        repo_layers: dict = {}
        for n in hf_nodes:
            repo_groups[n.hf_repo].append(n)
            # Use the largest num_layers seen for this repo.
            if n.hf_repo not in repo_layers or n.num_layers > repo_layers[n.hf_repo]:
                repo_layers[n.hf_repo] = n.num_layers

        # Smart scoring: demand_rpm × coverage_deficit
        # coverage_deficit = uncovered_layers / total_layers (0 = fully covered)
        # demand_rpm comes from the stats collector; defaults to 0.0 when no traffic yet.
        demand_rpms: dict[str, float] = {}
        if server.stats is not None:
            local_rpms = server.stats.get_local_rpms()
            for repo in repo_groups:
                demand_rpms[repo] = local_rpms.get(repo, {}).get("rpm_last_hour", 0.0)

        best_repo = None
        best_score = -1.0
        best_gap_size = -1
        best_gap_start = 0
        best_num_layers = 0

        for repo, nodes in repo_groups.items():
            total = repo_layers[repo]
            covered = sorted(
                [(n.shard_start, n.shard_end) for n in nodes],
                key=lambda t: t[0],
            )
            # Walk from 0 to find first uncovered layer.
            gap_start = 0
            for s, e in covered:
                if s <= gap_start:
                    gap_start = max(gap_start, e + 1)
                else:
                    break
            gap_size = max(0, (total - 1) - gap_start + 1)
            coverage_deficit = gap_size / max(total, 1)
            demand = demand_rpms.get(repo, 0.0)
            # +1.0 floor on demand so models with no recorded traffic still compete by coverage.
            score = (demand + 1.0) * (coverage_deficit + 0.01)
            if score > best_score or (score == best_score and gap_size > best_gap_size):
                best_score = score
                best_gap_size = gap_size
                best_gap_start = gap_start
                best_repo = repo
                best_num_layers = total

        gap_found = best_gap_size > 0
        if not gap_found:
            # All shards covered — assign to highest-demand model for redundancy.
            best_repo = max(repo_groups, key=lambda r: demand_rpms.get(r, 0.0))
            best_gap_start = 0
            best_num_layers = repo_layers[best_repo]

        # Capacity: use the same 80%-of-memory rule as registered node planning.
        total_l = best_num_layers
        memory_mb = vram_mb if vram_mb > 0 else ram_mb
        if memory_mb > 0:
            max_layers = min(
                total_l,
                max(1, int(((memory_mb * 1024 * 1024) * 0.8) // (30 * 1024 * 1024))),
            )
        elif device == "cuda" and vram_mb >= 8192:
            max_layers = total_l
        else:
            max_layers = max(1, total_l // 2)

        shard_start = best_gap_start
        shard_end = min(total_l - 1, shard_start + max_layers - 1)

        self._send_json(200, {
            "hf_repo": best_repo,
            "shard_start": shard_start,
            "shard_end": shard_end,
            "num_layers": total_l,
            "gap_found": gap_found,
            "price_per_token": 0.0,
        })

    def _handle_route(self, parsed: urllib.parse.ParseResult):
        server: _TrackerHTTPServer = self.server  # type: ignore[assignment]
        params = urllib.parse.parse_qs(parsed.query)
        model_list = params.get("model")
        if not model_list:
            self._send_json(400, {"error": "missing 'model' query parameter"})
            return

        model = model_list[0]
        preset = server.model_presets.get(model)

        with server.lock:
            self._purge_expired_nodes()
            if preset is not None:
                # Preset-based routing (stub-model system).
                alive = [node for node in server.registry.values() if node.model == model]
                required_start, required_end = _preset_layer_bounds(preset)
            else:
                # HF model routing: match by hf_repo (full) or model short name.
                alive = [
                    node for node in server.registry.values()
                    if _node_matches_model(node, model)
                    and node.shard_start is not None
                    and node.shard_end is not None
                    and node.num_layers is not None
                ]
                if not alive:
                    self._send_json(404, {"error": f"no nodes registered for model {model!r}"})
                    return
                required_start = 0
                required_end = max(n.num_layers for n in alive) - 1  # type: ignore[type-var]

        if server.contracts is not None:
            alive = [
                node for node in alive
                if not node.wallet_address or not server.contracts.registry.get_wallet(node.wallet_address).banned
            ]

        route, error = _select_route(alive, required_start, required_end)
        if error:
            self._send_json(503, {"error": error})
            return

        covered_up_to = required_start - 1
        route_with_start: list[tuple] = []
        for rn in route:
            route_with_start.append((rn, covered_up_to + 1))
            covered_up_to = rn.shard_end if rn.shard_end is not None else covered_up_to

        self._send_json(200, {
            "route": [e.endpoint for e, _ in route_with_start],
            "nodes": [
                {
                    "node_id": e.node_id,
                    "endpoint": e.endpoint,
                    "start_layer": start,
                    "relay_addr": e.relay_addr,
                    "peer_id": e.peer_id,
                    "wallet_address": e.wallet_address,
                    "shard_start": e.shard_start,
                    "shard_end": e.shard_end,
                    "model": e.model,
                    "shard_checksum": e.shard_checksum,
                    "score": e.score,
                }
                for e, start in route_with_start
            ],
        })

    def _handle_routes(self, parsed: urllib.parse.ParseResult):
        server: _TrackerHTTPServer = self.server  # type: ignore[assignment]
        params = urllib.parse.parse_qs(parsed.query)
        model_list = params.get("model")
        if not model_list:
            self._send_json(400, {"error": "missing 'model' query parameter"})
            return

        try:
            redundancy = int(params.get("redundancy", ["1"])[0])
        except ValueError:
            self._send_json(400, {"error": "redundancy must be an integer"})
            return
        if redundancy < 1:
            self._send_json(400, {"error": "redundancy must be at least 1"})
            return

        model = model_list[0]
        preset = server.model_presets.get(model)
        if preset is None:
            self._send_json(404, {"error": f"unknown model preset: {model!r}"})
            return

        required_start, required_end = _preset_layer_bounds(preset)

        with server.lock:
            self._purge_expired_nodes()
            candidates = [node for node in server.registry.values() if node.model == model]
        if server.contracts is not None:
            candidates = [
                node for node in candidates
                if not node.wallet_address or not server.contracts.registry.get_wallet(node.wallet_address).banned
            ]

        routes = []
        remaining = list(candidates)
        for _ in range(redundancy):
            route, error = _select_route(remaining, required_start, required_end)
            if error:
                self._send_json(503, {"error": error})
                return
            route_endpoints = {node.endpoint for node in route}
            remaining = [node for node in remaining if node.endpoint not in route_endpoints]
            routes.append({
                "route": [e.endpoint for e in route],
                "nodes": [
                    {
                        "node_id": e.node_id,
                        "endpoint": e.endpoint,
                        "relay_addr": e.relay_addr,
                        "peer_id": e.peer_id,
                        "wallet_address": e.wallet_address,
                        "shard_start": e.shard_start,
                        "shard_end": e.shard_end,
                        "model": e.model,
                        "shard_checksum": e.shard_checksum,
                        "score": e.score,
                    }
                    for e in route
                ],
            })

        self._send_json(200, {"routes": routes})


class TrackerServer:
    """HTTP tracker that manages node registration and resolves inference routes.

    Nodes register via POST /v1/nodes/register and keep themselves alive with
    POST /v1/nodes/<id>/heartbeat.  The gateway queries GET /v1/route?model=<name>
    to obtain an ordered list of node endpoints whose shards cover all layers.
    """

    def __init__(
        self,
        host: str = "127.0.0.1",
        port: int = 0,
        heartbeat_timeout: float = 90.0,
        rebalance_interval: float = 30.0,
        model_presets: dict | None = None,
        contracts: Any | None = None,
        minimum_stake: int = 0,
        cluster_peers: list[str] | None = None,
        cluster_self_url: str | None = None,
        stats_db: str | None = None,
        relay_url: str | None = None,
    ) -> None:
        self._host = host
        self._requested_port = port
        self._heartbeat_timeout = heartbeat_timeout
        self._rebalance_interval = rebalance_interval
        self._model_presets: dict = (
            model_presets if model_presets is not None else dict(DEFAULT_MODEL_PRESETS)
        )
        self._contracts = contracts
        self._minimum_stake = minimum_stake
        self._cluster_peers: list[str] = list(cluster_peers) if cluster_peers else []
        self._cluster_self_url: str | None = cluster_self_url
        self._relay_url = relay_url
        self._registry: dict[str, _NodeEntry] = {}
        self._lock = threading.Lock()
        self._server: _TrackerHTTPServer | None = None
        self._thread: threading.Thread | None = None
        self._rebalance_stop = threading.Event()
        self._rebalance_thread: threading.Thread | None = None
        self._raft: RaftNode | None = None
        self._gossip: NodeGossip | None = None
        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
        self.port: int | None = None

    def start(self) -> int:
        if self._server is not None:
            raise RuntimeError("TrackerServer is already running")

        effective_relay_url = (
            self._relay_url
            or derive_relay_url_from_public_tracker_url(self._cluster_self_url)
        )

        # Start HTTP server first so we know our port
        self._server = _TrackerHTTPServer(
            (self._host, self._requested_port),
            _TrackerHandler,
            self._registry,
            self._lock,
            self._heartbeat_timeout,
            self._model_presets,
            self._contracts,
            self._minimum_stake,
            relay_url=effective_relay_url,
            stats=self._stats,
        )
        self.port = self._server.server_address[1]

        # Start Raft + gossip if cluster peers are configured
        if self._cluster_peers:
            self_url = self._cluster_self_url or f"http://{self._host}:{self.port}"
            self._raft = RaftNode(
                self_url=self_url,
                peers=self._cluster_peers,
                apply_fn=self._raft_apply,
            )
            self._gossip = NodeGossip(peers=self._cluster_peers)
            self._server.raft = self._raft
            self._server.gossip = self._gossip
            self._raft.start()
            self._gossip.start()

        self._rebalance_stop.clear()
        self._stats_stop.clear()
        self._thread = threading.Thread(target=self._server.serve_forever, daemon=True)
        self._thread.start()
        self._rebalance_thread = threading.Thread(target=self._rebalance_loop, daemon=True)
        self._rebalance_thread.start()
        self._stats_thread = threading.Thread(target=self._stats_loop, daemon=True)
        self._stats_thread.start()
        return self.port

    def _raft_apply(self, command: str, payload: dict) -> None:
        """Called by RaftNode when a log entry is committed — replicate to local registry."""
        if command != "register":
            return
        # Re-apply the registration to our local registry (follower path).
        # On the leader this is a no-op since it already registered locally.
        node_id = payload.get("node_id")
        if not node_id or node_id in self._registry:
            return  # already present (leader case) or malformed
        endpoint = payload.get("endpoint", "")
        try:
            shard_start = int(payload["shard_start"]) if payload.get("shard_start") is not None else None
            shard_end   = int(payload["shard_end"])   if payload.get("shard_end")   is not None else None
        except (TypeError, ValueError):
            return
        entry = _NodeEntry(
            node_id=node_id,
            endpoint=endpoint.rstrip("/"),
            shard_start=shard_start,
            shard_end=shard_end,
            model=payload.get("model", "stub-model"),
            shard_checksum=payload.get("shard_checksum"),
            hardware_profile=payload.get("hardware_profile", {}),
            wallet_address=payload.get("wallet_address"),
            score=float(payload.get("score", 1.0)),
            tracker_mode=bool(payload.get("tracker_mode", False)),
            hf_repo=payload.get("hf_repo"),
            num_layers=int(payload["num_layers"]) if payload.get("num_layers") is not None else None,
        )
        with self._lock:
            self._registry[node_id] = entry

    def _rebalance_loop(self) -> None:
        while not self._rebalance_stop.wait(self._rebalance_interval):
            server = self._server
            if server is None:
                return
            with self._lock:
                _purge_expired_nodes_locked(server)
                _rebalance_all_locked(server)

    def _stats_loop(self) -> None:
        """Periodically save stats to DB and push local slice 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._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()
                body = json.dumps({"tracker_url": self_url, "stats": local_rpms}).encode()
                for peer in self._cluster_peers:
                    try:
                        req = urllib.request.Request(
                            f"{peer}/v1/stats/gossip",
                            data=body,
                            headers={"Content-Type": "application/json"},
                            method="POST",
                        )
                        with urllib.request.urlopen(req, timeout=2.0) as r:
                            r.read()
                    except Exception:
                        pass

    def stop(self) -> None:
        if self._raft is not None:
            self._raft.stop()
        if self._gossip is not None:
            self._gossip.stop()
        if self._server is None:
            return
        self._rebalance_stop.set()
        self._stats_stop.set()
        if self._stats is not None:
            self._stats.save_to_db()
        self._server.shutdown()
        self._server.server_close()
        if self._thread is not None:
            self._thread.join(timeout=1)
        if self._rebalance_thread is not None:
            self._rebalance_thread.join(timeout=1)
        if self._stats_thread is not None:
            self._stats_thread.join(timeout=1)
        self._server = None
        self._thread = None
        self._rebalance_thread = None
        self._stats_thread = None
        self._raft = None
        self._gossip = None
        self.port = None