neuron-tai/tests/test_tracker_routing.py

"""US-003 tests: tracker node registration and route selection."""

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

from meshnet_gateway.server import GatewayServer, _banned_route_wallet
from meshnet_node.server import StubNodeServer
from meshnet_contracts import LocalSolanaContracts
from meshnet_tracker.server import TrackerServer, _NodeEntry, _registration_ban_error


def _post_json(url: str, payload: dict) -> dict:
    data = json.dumps(payload).encode()
    req = urllib.request.Request(
        url, data=data, headers={"Content-Type": "application/json"}, method="POST"
    )
    with urllib.request.urlopen(req) as r:
        return json.loads(r.read())


def _get_json(url: str) -> dict:
    with urllib.request.urlopen(url) as r:
        return json.loads(r.read())


def test_tracker_send_json_ignores_broken_pipe_after_client_disconnect():
    """A disconnected client must not dump a BrokenPipe traceback from the tracker."""
    from meshnet_tracker.server import _TrackerHandler

    class BrokenPipeWriter:
        def write(self, _body):
            raise BrokenPipeError

    class DummyHandler:
        wfile = BrokenPipeWriter()

        def send_response(self, _status):
            pass

        def send_header(self, _name, _value):
            pass

        def end_headers(self):
            pass

    _TrackerHandler._send_json(DummyHandler(), 200, {"ok": True})


def test_tracker_registration_node_id_includes_wallet_prefix_and_stable_suffix():
    tracker = TrackerServer()
    tracker_port = tracker.start()
    wallet = "7j77FsPY1evV8tuf4Z73AVrWwxBEW1pvKwi4EvcRD3g"
    try:
        first = _post_json(
            f"http://127.0.0.1:{tracker_port}/v1/nodes/register",
            {"endpoint": "http://127.0.0.1:9100", "model": "Qwen2.5-0.5B-Instruct",
             "hf_repo": "Qwen/Qwen2.5-0.5B-Instruct", "num_layers": 24,
             "shard_start": 0, "shard_end": 21, "wallet_address": wallet,
             "hardware_profile": {}, "score": 1.0},
        )
        second = _post_json(
            f"http://127.0.0.1:{tracker_port}/v1/nodes/register",
            {"endpoint": "http://127.0.0.1:9100", "model": "Qwen2.5-0.5B-Instruct",
             "hf_repo": "Qwen/Qwen2.5-0.5B-Instruct", "num_layers": 24,
             "shard_start": 0, "shard_end": 21, "wallet_address": wallet,
             "hardware_profile": {}, "score": 1.0},
        )
        different_endpoint = _post_json(
            f"http://127.0.0.1:{tracker_port}/v1/nodes/register",
            {"endpoint": "http://127.0.0.1:9101", "model": "Qwen2.5-0.5B-Instruct",
             "hf_repo": "Qwen/Qwen2.5-0.5B-Instruct", "num_layers": 24,
             "shard_start": 20, "shard_end": 23, "wallet_address": wallet,
             "hardware_profile": {}, "score": 1.0},
        )

        assert first["node_id"].startswith("7j77FsPY-")
        assert second["node_id"] == first["node_id"]
        assert different_endpoint["node_id"].startswith("7j77FsPY-")
        assert different_endpoint["node_id"] != first["node_id"]

        route_resp = _get_json(
            f"http://127.0.0.1:{tracker_port}/v1/route?model=Qwen/Qwen2.5-0.5B-Instruct"
        )
        assert route_resp["nodes"][0]["node_id"] == first["node_id"]
    finally:
        tracker.stop()


def test_tracker_node_registration():
    """A node can register with the tracker and receives a node_id."""
    tracker = TrackerServer()
    tracker_port = tracker.start()
    try:
        resp = _post_json(
            f"http://127.0.0.1:{tracker_port}/v1/nodes/register",
            {
                "endpoint": "http://127.0.0.1:9000",
                "shard_start": 0,
                "shard_end": 15,
                "hardware_profile": {"gpu": "rtx3090"},
                "score": 1.0,
            },
        )
        assert "node_id" in resp
        assert isinstance(resp["node_id"], str)
        assert len(resp["node_id"]) > 0
    finally:
        tracker.stop()


def test_tracker_route_selection():
    """Tracker returns ordered route when nodes collectively cover all layers."""
    tracker = TrackerServer()
    tracker_port = tracker.start()
    try:
        _post_json(
            f"http://127.0.0.1:{tracker_port}/v1/nodes/register",
            {"endpoint": "http://127.0.0.1:9001", "shard_start": 0, "shard_end": 15,
             "hardware_profile": {}, "score": 1.0},
        )
        _post_json(
            f"http://127.0.0.1:{tracker_port}/v1/nodes/register",
            {"endpoint": "http://127.0.0.1:9002", "shard_start": 16, "shard_end": 31,
             "hardware_profile": {}, "score": 1.0},
        )

        route_resp = _get_json(
            f"http://127.0.0.1:{tracker_port}/v1/route?model=stub-model"
        )
        assert "route" in route_resp
        assert route_resp["route"] == [
            "http://127.0.0.1:9001",
            "http://127.0.0.1:9002",
        ]
    finally:
        tracker.stop()


def test_tracker_route_error_no_coverage():
    """Tracker returns 503 when registered nodes do not cover all required layers."""
    tracker = TrackerServer()
    tracker_port = tracker.start()
    try:
        # Only half the layers covered — route must fail.
        _post_json(
            f"http://127.0.0.1:{tracker_port}/v1/nodes/register",
            {"endpoint": "http://127.0.0.1:9001", "shard_start": 0, "shard_end": 15,
             "hardware_profile": {}, "score": 1.0},
        )

        try:
            _get_json(f"http://127.0.0.1:{tracker_port}/v1/route?model=stub-model")
            raise AssertionError("Expected 503 but request succeeded")
        except urllib.error.HTTPError as exc:
            assert exc.code == 503
            body = json.loads(exc.read())
            assert "error" in body
    finally:
        tracker.stop()


def test_tracker_coverage_endpoint_reports_uncovered_ranges():
    """Coverage endpoint returns compressed layer ranges with node counts."""
    tracker = TrackerServer(model_presets={
        "tiny-model": {
            "total_layers": 6,
            "bytes_per_layer": {"bfloat16": 1_000},
        },
    })
    tracker_port = tracker.start()
    try:
        _post_json(
            f"http://127.0.0.1:{tracker_port}/v1/nodes/register",
            {"endpoint": "http://127.0.0.1:9001", "model": "tiny-model",
             "shard_start": 0, "shard_end": 2, "hardware_profile": {}, "score": 1.0},
        )

        coverage_resp = _get_json(f"http://127.0.0.1:{tracker_port}/v1/coverage/tiny-model")
        assert coverage_resp["coverage"] == [
            {"start_layer": 0, "end_layer": 2, "node_count": 1},
            {"start_layer": 3, "end_layer": 5, "node_count": 0},
        ]

        try:
            _get_json(f"http://127.0.0.1:{tracker_port}/v1/route?model=tiny-model")
            raise AssertionError("Expected 503 for zero-coverage range")
        except urllib.error.HTTPError as exc:
            assert exc.code == 503
    finally:
        tracker.stop()


def test_tracker_coverage_endpoint_accepts_registered_hf_repo_or_short_name():
    """Coverage endpoint supports real HF models registered outside preset catalog."""
    tracker = TrackerServer()
    tracker_port = tracker.start()
    try:
        _post_json(
            f"http://127.0.0.1:{tracker_port}/v1/nodes/register",
            {"endpoint": "http://127.0.0.1:9101", "model": "Qwen2.5-0.5B-Instruct",
             "hf_repo": "Qwen/Qwen2.5-0.5B-Instruct", "num_layers": 24,
             "shard_start": 0, "shard_end": 21, "hardware_profile": {}, "score": 1.0},
        )
        _post_json(
            f"http://127.0.0.1:{tracker_port}/v1/nodes/register",
            {"endpoint": "http://127.0.0.1:9102", "model": "Qwen2.5-0.5B-Instruct",
             "hf_repo": "Qwen/Qwen2.5-0.5B-Instruct", "num_layers": 24,
             "shard_start": 20, "shard_end": 23, "hardware_profile": {}, "score": 1.0},
        )

        by_repo = _get_json(
            f"http://127.0.0.1:{tracker_port}/v1/coverage/Qwen/Qwen2.5-0.5B-Instruct"
        )
        by_short_name = _get_json(
            f"http://127.0.0.1:{tracker_port}/v1/coverage/Qwen2.5-0.5B-Instruct"
        )

        assert by_repo["model"] == "Qwen/Qwen2.5-0.5B-Instruct"
        assert by_repo["coverage"] == [
            {"start_layer": 0, "end_layer": 19, "node_count": 1},
            {"start_layer": 20, "end_layer": 21, "node_count": 2},
            {"start_layer": 22, "end_layer": 23, "node_count": 1},
        ]
        assert by_short_name["coverage"] == by_repo["coverage"]
    finally:
        tracker.stop()


def test_tracker_models_endpoint_lists_registered_hf_repo_and_short_name_alias():
    tracker = TrackerServer()
    tracker_port = tracker.start()
    try:
        _post_json(
            f"http://127.0.0.1:{tracker_port}/v1/nodes/register",
            {"endpoint": "http://127.0.0.1:9111", "model": "Qwen2.5-0.5B-Instruct",
             "hf_repo": "Qwen/Qwen2.5-0.5B-Instruct", "num_layers": 24,
             "shard_start": 0, "shard_end": 23, "hardware_profile": {}, "score": 1.0},
        )

        models_resp = _get_json(f"http://127.0.0.1:{tracker_port}/v1/models")

        model = next(item for item in models_resp["data"] if item["id"] == "Qwen/Qwen2.5-0.5B-Instruct")
        assert model["name"] == "Qwen2.5-0.5B-Instruct"
        assert model["hf_repo"] == "Qwen/Qwen2.5-0.5B-Instruct"
        assert model["aliases"] == ["Qwen/Qwen2.5-0.5B-Instruct", "Qwen2.5-0.5B-Instruct"]
        assert model["shard_coverage_percentage"] == 100.0
    finally:
        tracker.stop()


def test_tracker_auto_assigns_new_node_to_uncovered_range_first():
    """Capability-driven registration fills the first uncovered layer gap."""
    tracker = TrackerServer(model_presets={
        "tiny-model": {
            "total_layers": 8,
            "bytes_per_layer": {"bfloat16": 1_000},
        },
    })
    tracker_port = tracker.start()
    try:
        _post_json(
            f"http://127.0.0.1:{tracker_port}/v1/nodes/register",
            {"endpoint": "http://127.0.0.1:9001", "model": "tiny-model",
             "shard_start": 0, "shard_end": 3, "hardware_profile": {}, "score": 1.0},
        )
        reg = _post_json(
            f"http://127.0.0.1:{tracker_port}/v1/nodes/register",
            {"endpoint": "http://127.0.0.1:9002", "model": "tiny-model",
             "vram_bytes": 10_000, "ram_bytes": 20_000, "quantizations": ["bfloat16"],
             "benchmark_tokens_per_sec": 1.0, "hardware_profile": {}, "score": 1.0},
        )

        assert reg["directive"]["action"] == "LOAD_SHARD"
        assert reg["directive"]["start_layer"] == 4
        assert reg["directive"]["end_layer"] == 7
    finally:
        tracker.stop()


def test_tracker_speed_weighted_vram_assignment_covers_model():
    """Three auto-assigned nodes reach 100% coverage with widest range on largest VRAM."""
    tracker = TrackerServer(model_presets={
        "tiny-model": {
            "total_layers": 12,
            "bytes_per_layer": {"bfloat16": 1_000},
        },
    })
    tracker_port = tracker.start()
    try:
        _post_json(
            f"http://127.0.0.1:{tracker_port}/v1/nodes/register",
            {"endpoint": "http://127.0.0.1:9001", "model": "tiny-model",
             "vram_bytes": 4_000, "ram_bytes": 10_000, "quantizations": ["bfloat16"],
             "benchmark_tokens_per_sec": 1.0, "hardware_profile": {}, "score": 1.0},
        )
        _post_json(
            f"http://127.0.0.1:{tracker_port}/v1/nodes/register",
            {"endpoint": "http://127.0.0.1:9002", "model": "tiny-model",
             "vram_bytes": 5_000, "ram_bytes": 10_000, "quantizations": ["bfloat16"],
             "benchmark_tokens_per_sec": 1.0, "hardware_profile": {}, "score": 1.0},
        )
        _post_json(
            f"http://127.0.0.1:{tracker_port}/v1/nodes/register",
            {"endpoint": "http://127.0.0.1:9003", "model": "tiny-model",
             "vram_bytes": 7_000, "ram_bytes": 10_000, "quantizations": ["bfloat16"],
             "benchmark_tokens_per_sec": 1.0, "hardware_profile": {}, "score": 1.0},
        )

        route_resp = _get_json(f"http://127.0.0.1:{tracker_port}/v1/route?model=tiny-model")
        widths = {
            node["endpoint"]: node["shard_end"] - node["shard_start"] + 1
            for node in route_resp["nodes"]
        }
        assert widths["http://127.0.0.1:9003"] == 5
        assert widths["http://127.0.0.1:9002"] == 4
        assert widths["http://127.0.0.1:9001"] == 3

        coverage_resp = _get_json(f"http://127.0.0.1:{tracker_port}/v1/coverage/tiny-model")
        assert all(segment["node_count"] >= 1 for segment in coverage_resp["coverage"])
    finally:
        tracker.stop()


def test_tracker_speed_is_primary_when_both_nodes_can_cover_gap():
    tracker = TrackerServer(model_presets={
        "tiny-model": {
            "total_layers": 12,
            "bytes_per_layer": {"bfloat16": 1_000},
        },
    })
    tracker_port = tracker.start()
    try:
        _post_json(
            f"http://127.0.0.1:{tracker_port}/v1/nodes/register",
            {"endpoint": "http://127.0.0.1:9011", "model": "tiny-model",
             "vram_bytes": 20_000, "ram_bytes": 10_000, "quantizations": ["bfloat16"],
             "benchmark_tokens_per_sec": 1.0, "hardware_profile": {}, "score": 1.0},
        )
        _post_json(
            f"http://127.0.0.1:{tracker_port}/v1/nodes/register",
            {"endpoint": "http://127.0.0.1:9012", "model": "tiny-model",
             "vram_bytes": 15_000, "ram_bytes": 10_000, "quantizations": ["bfloat16"],
             "benchmark_tokens_per_sec": 3.0, "hardware_profile": {}, "score": 1.0},
        )

        route_resp = _get_json(f"http://127.0.0.1:{tracker_port}/v1/route?model=tiny-model")
        widths = {
            node["endpoint"]: node["shard_end"] - node["shard_start"] + 1
            for node in route_resp["nodes"]
        }
        assert widths["http://127.0.0.1:9012"] > widths["http://127.0.0.1:9011"]
    finally:
        tracker.stop()


def test_tracker_registration_directive_is_not_replayed_on_heartbeat():
    tracker = TrackerServer(model_presets={
        "tiny-model": {
            "total_layers": 2,
            "bytes_per_layer": {"bfloat16": 1_000},
        },
    })
    tracker_port = tracker.start()
    try:
        reg = _post_json(
            f"http://127.0.0.1:{tracker_port}/v1/nodes/register",
            {"endpoint": "http://127.0.0.1:9013", "model": "tiny-model",
             "vram_bytes": 5_000, "ram_bytes": 10_000, "quantizations": ["bfloat16"],
             "benchmark_tokens_per_sec": 1.0, "hardware_profile": {}, "score": 1.0},
        )
        assert reg["directive"]["action"] == "LOAD_SHARD"

        hb = _post_json(f"http://127.0.0.1:{tracker_port}/v1/nodes/{reg['node_id']}/heartbeat", {})
        assert hb == {}
    finally:
        tracker.stop()


def test_tracker_reassignment_emits_drop_before_load():
    tracker = TrackerServer(model_presets={
        "tiny-model": {
            "total_layers": 4,
            "bytes_per_layer": {"bfloat16": 1_000},
        },
    })
    tracker_port = tracker.start()
    try:
        slow = _post_json(
            f"http://127.0.0.1:{tracker_port}/v1/nodes/register",
            {"endpoint": "http://127.0.0.1:9015", "model": "tiny-model",
             "vram_bytes": 10_000, "ram_bytes": 10_000, "quantizations": ["bfloat16"],
             "benchmark_tokens_per_sec": 1.0, "hardware_profile": {}, "score": 1.0},
        )
        _post_json(
            f"http://127.0.0.1:{tracker_port}/v1/nodes/register",
            {"endpoint": "http://127.0.0.1:9016", "model": "tiny-model",
             "vram_bytes": 10_000, "ram_bytes": 10_000, "quantizations": ["bfloat16"],
             "benchmark_tokens_per_sec": 3.0, "hardware_profile": {}, "score": 1.0},
        )

        hb = _post_json(f"http://127.0.0.1:{tracker_port}/v1/nodes/{slow['node_id']}/heartbeat", {})
        assert [directive["action"] for directive in hb["directives"]] == ["DROP_SHARD", "LOAD_SHARD"]
    finally:
        tracker.stop()


def test_tracker_periodic_rebalance_purges_expired_nodes_without_requests():
    tracker = TrackerServer(
        heartbeat_timeout=0.05,
        rebalance_interval=0.02,
        model_presets={"tiny-model": {"total_layers": 1, "bytes_per_layer": {"bfloat16": 1_000}}},
    )
    tracker_port = tracker.start()
    try:
        reg = _post_json(
            f"http://127.0.0.1:{tracker_port}/v1/nodes/register",
            {"endpoint": "http://127.0.0.1:9014", "model": "tiny-model",
             "vram_bytes": 5_000, "ram_bytes": 10_000, "quantizations": ["bfloat16"],
             "benchmark_tokens_per_sec": 1.0, "hardware_profile": {}, "score": 1.0},
        )
        assert reg["node_id"] in tracker._registry
        time.sleep(0.12)
        assert reg["node_id"] not in tracker._registry
    finally:
        tracker.stop()


def test_tracker_faster_node_receives_wider_range_when_capacity_ties():
    tracker = TrackerServer(model_presets={
        "tiny-model": {
            "total_layers": 12,
            "bytes_per_layer": {"bfloat16": 1_000},
        },
    })
    tracker_port = tracker.start()
    try:
        _post_json(
            f"http://127.0.0.1:{tracker_port}/v1/nodes/register",
            {"endpoint": "http://127.0.0.1:9001", "model": "tiny-model",
             "vram_bytes": 20_000, "ram_bytes": 10_000, "quantizations": ["bfloat16"],
             "benchmark_tokens_per_sec": 1.0, "hardware_profile": {}, "score": 1.0},
        )
        _post_json(
            f"http://127.0.0.1:{tracker_port}/v1/nodes/register",
            {"endpoint": "http://127.0.0.1:9002", "model": "tiny-model",
             "vram_bytes": 20_000, "ram_bytes": 10_000, "quantizations": ["bfloat16"],
             "benchmark_tokens_per_sec": 2.0, "hardware_profile": {}, "score": 1.0},
        )

        route_resp = _get_json(f"http://127.0.0.1:{tracker_port}/v1/route?model=tiny-model")
        widths = {
            node["endpoint"]: node["shard_end"] - node["shard_start"] + 1
            for node in route_resp["nodes"]
        }
        assert widths["http://127.0.0.1:9002"] > widths["http://127.0.0.1:9001"]
    finally:
        tracker.stop()


def test_tracker_rebalances_after_middle_range_node_timeout():
    """Killing a middle shard queues LOAD_SHARD and restores coverage."""
    tracker = TrackerServer(
        heartbeat_timeout=0.15,
        model_presets={
            "tiny-model": {
                "total_layers": 12,
                "bytes_per_layer": {"bfloat16": 1_000},
            },
        },
    )
    tracker_port = tracker.start()
    ids: dict[str, str] = {}
    try:
        for port in (9001, 9002, 9003, 9004):
            reg = _post_json(
                f"http://127.0.0.1:{tracker_port}/v1/nodes/register",
                {"endpoint": f"http://127.0.0.1:{port}", "model": "tiny-model",
                 "vram_bytes": 5_000, "ram_bytes": 10_000, "quantizations": ["bfloat16"],
                 "benchmark_tokens_per_sec": 1.0, "hardware_profile": {}, "score": 1.0},
            )
            ids[str(port)] = reg["node_id"]

        for node_id in ids.values():
            _post_json(f"http://127.0.0.1:{tracker_port}/v1/nodes/{node_id}/heartbeat", {})

        time.sleep(0.10)
        for port in ("9001", "9003", "9004"):
            _post_json(f"http://127.0.0.1:{tracker_port}/v1/nodes/{ids[port]}/heartbeat", {})
        time.sleep(0.10)

        coverage_resp = _get_json(f"http://127.0.0.1:{tracker_port}/v1/coverage/tiny-model")
        assert all(segment["node_count"] >= 1 for segment in coverage_resp["coverage"])

        heartbeat_responses = [
            _post_json(f"http://127.0.0.1:{tracker_port}/v1/nodes/{ids[port]}/heartbeat", {})
            for port in ("9001", "9003", "9004")
        ]
        directives = [
            directive
            for response in heartbeat_responses
            for directive in response.get("directives", [])
        ]
        assert any(directive["action"] == "LOAD_SHARD" for directive in directives)
    finally:
        tracker.stop()


def test_tracker_route_error_no_nodes():
    """Tracker returns 503 with clear error when the registry is empty."""
    tracker = TrackerServer()
    tracker_port = tracker.start()
    try:
        try:
            _get_json(f"http://127.0.0.1:{tracker_port}/v1/route?model=stub-model")
            raise AssertionError("Expected 503 but request succeeded")
        except urllib.error.HTTPError as exc:
            assert exc.code == 503
            body = json.loads(exc.read())
            assert "error" in body
    finally:
        tracker.stop()


def test_tracker_heartbeat_updates_node():
    """Sending a heartbeat for a registered node succeeds."""
    tracker = TrackerServer()
    tracker_port = tracker.start()
    try:
        reg = _post_json(
            f"http://127.0.0.1:{tracker_port}/v1/nodes/register",
            {"endpoint": "http://127.0.0.1:9001", "shard_start": 0, "shard_end": 31,
             "hardware_profile": {}, "score": 1.0},
        )
        node_id = reg["node_id"]

        hb_resp = _post_json(
            f"http://127.0.0.1:{tracker_port}/v1/nodes/{node_id}/heartbeat", {}
        )
        assert hb_resp == {}
    finally:
        tracker.stop()


def test_tracker_heartbeat_expiry():
    """Nodes that miss their heartbeat window are excluded from routes."""
    tracker = TrackerServer(heartbeat_timeout=0.05)  # 50 ms
    tracker_port = tracker.start()
    try:
        _post_json(
            f"http://127.0.0.1:{tracker_port}/v1/nodes/register",
            {"endpoint": "http://127.0.0.1:9001", "shard_start": 0, "shard_end": 31,
             "hardware_profile": {}, "score": 1.0},
        )

        # Immediately after registration the node is alive.
        route_resp = _get_json(
            f"http://127.0.0.1:{tracker_port}/v1/route?model=stub-model"
        )
        assert route_resp["route"] == ["http://127.0.0.1:9001"]

        # Let the heartbeat window expire.
        time.sleep(0.15)

        try:
            _get_json(f"http://127.0.0.1:{tracker_port}/v1/route?model=stub-model")
            raise AssertionError("Expected 503 after heartbeat expiry")
        except urllib.error.HTTPError as exc:
            assert exc.code == 503
    finally:
        tracker.stop()


def test_tracker_heartbeat_expiry_removes_node_from_registry():
    """Expired nodes are removed, not merely hidden from route responses."""
    tracker = TrackerServer(heartbeat_timeout=0.05)
    tracker_port = tracker.start()
    try:
        reg = _post_json(
            f"http://127.0.0.1:{tracker_port}/v1/nodes/register",
            {"endpoint": "http://127.0.0.1:9001", "shard_start": 0, "shard_end": 31,
             "hardware_profile": {}, "score": 1.0},
        )
        node_id = reg["node_id"]
        time.sleep(0.15)

        try:
            _get_json(f"http://127.0.0.1:{tracker_port}/v1/route?model=stub-model")
            raise AssertionError("Expected 503 after heartbeat expiry")
        except urllib.error.HTTPError as exc:
            assert exc.code == 503

        try:
            _post_json(f"http://127.0.0.1:{tracker_port}/v1/nodes/{node_id}/heartbeat", {})
            raise AssertionError("Expected 404 for removed node")
        except urllib.error.HTTPError as exc:
            assert exc.code == 404
    finally:
        tracker.stop()


def test_tracker_route_rejects_non_extending_overlap():
    """Overlapping shards that do not extend coverage cannot form a complete route."""
    tracker = TrackerServer()
    tracker_port = tracker.start()
    try:
        _post_json(
            f"http://127.0.0.1:{tracker_port}/v1/nodes/register",
            {"endpoint": "http://127.0.0.1:9001", "shard_start": 0, "shard_end": 10,
             "hardware_profile": {}, "score": 1.0},
        )
        _post_json(
            f"http://127.0.0.1:{tracker_port}/v1/nodes/register",
            {"endpoint": "http://127.0.0.1:9002", "shard_start": 0, "shard_end": 5,
             "hardware_profile": {}, "score": 1.0},
        )

        try:
            _get_json(f"http://127.0.0.1:{tracker_port}/v1/route?model=stub-model")
            raise AssertionError("Expected 503 for incomplete overlapping route")
        except urllib.error.HTTPError as exc:
            assert exc.code == 503
    finally:
        tracker.stop()


def test_tracker_registration_rejects_invalid_payload():
    """Registration errors return a defined JSON 400 response."""
    tracker = TrackerServer()
    tracker_port = tracker.start()
    try:
        try:
            _post_json(
                f"http://127.0.0.1:{tracker_port}/v1/nodes/register",
                {"endpoint": "not-a-url", "shard_start": 10, "shard_end": 0,
                 "hardware_profile": []},
            )
            raise AssertionError("Expected 400 for invalid registration")
        except urllib.error.HTTPError as exc:
            assert exc.code == 400
            body = json.loads(exc.read())
            assert "error" in body
    finally:
        tracker.stop()


def test_tracker_routes_only_nodes_for_requested_model():
    """A node registered for one model cannot satisfy another model's route."""
    tracker = TrackerServer(model_presets={
        "model-a": {"layers_start": 0, "layers_end": 31},
        "model-b": {"layers_start": 0, "layers_end": 31},
    })
    tracker_port = tracker.start()
    try:
        _post_json(
            f"http://127.0.0.1:{tracker_port}/v1/nodes/register",
            {"endpoint": "http://127.0.0.1:9001", "shard_start": 0, "shard_end": 31,
             "model": "model-a", "hardware_profile": {}, "wallet_address": "wallet-a", "score": 1.0},
        )

        try:
            _get_json(f"http://127.0.0.1:{tracker_port}/v1/route?model=model-b")
            raise AssertionError("Expected model-b route to ignore model-a nodes")
        except urllib.error.HTTPError as exc:
            assert exc.code == 503
    finally:
        tracker.stop()


def test_tracker_assignment_coverage_is_model_scoped():
    """Shard assignment gaps are computed only from nodes serving the same model."""
    tracker = TrackerServer(model_presets={
        "model-a": {"layers_start": 0, "layers_end": 31},
        "model-b": {"layers_start": 0, "layers_end": 31},
    })
    tracker_port = tracker.start()
    try:
        _post_json(
            f"http://127.0.0.1:{tracker_port}/v1/nodes/register",
            {"endpoint": "http://127.0.0.1:9001", "shard_start": 0, "shard_end": 15,
             "model": "model-a", "hardware_profile": {}, "wallet_address": "wallet-a", "score": 1.0},
        )

        assignment = _get_json(f"http://127.0.0.1:{tracker_port}/v1/nodes/assign?model=model-b&device=cpu")
        assert assignment["shard_start"] == 0
        assert assignment["shard_end"] == 15
    finally:
        tracker.stop()


def test_tracker_excludes_banned_wallets_from_routes():
    """Tracker refuses route candidates whose wallet is banned on-chain."""
    contracts = LocalSolanaContracts()
    contracts.registry.submit_stake("wallet-a", 1_000)
    contracts.registry.submit_stake("wallet-b", 1_000)
    contracts.registry.ban_wallet("wallet-b")
    tracker = TrackerServer(contracts=contracts, minimum_stake=100)
    tracker_port = tracker.start()
    try:
        _post_json(
            f"http://127.0.0.1:{tracker_port}/v1/nodes/register",
            {"endpoint": "http://127.0.0.1:9001", "shard_start": 0, "shard_end": 15,
             "hardware_profile": {}, "wallet_address": "wallet-a", "score": 1.0},
        )
        try:
            _post_json(
                f"http://127.0.0.1:{tracker_port}/v1/nodes/register",
                {"endpoint": "http://127.0.0.1:9002", "shard_start": 16, "shard_end": 31,
                 "hardware_profile": {}, "wallet_address": "wallet-b", "score": 1.0},
            )
            raise AssertionError("Expected 403 for banned wallet registration")
        except urllib.error.HTTPError as exc:
            assert exc.code == 403

        try:
            _get_json(f"http://127.0.0.1:{tracker_port}/v1/route?model=stub-model")
            raise AssertionError("Expected 503 when only banned wallet covers a layer")
        except urllib.error.HTTPError as exc:
            assert exc.code == 503
            body = json.loads(exc.read())
            assert "layer 16" in body["error"]
    finally:
        tracker.stop()


def test_tracker_rejects_banned_wallet_registration():
    """A banned wallet cannot register with the tracker."""
    contracts = LocalSolanaContracts()
    contracts.registry.submit_stake("wallet-b", 1_000)
    contracts.registry.ban_wallet("wallet-b")
    tracker = TrackerServer(contracts=contracts, minimum_stake=100)
    tracker_port = tracker.start()
    try:
        try:
            _post_json(
                f"http://127.0.0.1:{tracker_port}/v1/nodes/register",
                {"endpoint": "http://127.0.0.1:9002", "shard_start": 0, "shard_end": 31,
                 "hardware_profile": {}, "wallet_address": "wallet-b", "score": 1.0},
            )
            raise AssertionError("Expected 403 for banned wallet registration")
        except urllib.error.HTTPError as exc:
            assert exc.code == 403
            body = json.loads(exc.read())
            assert body == {"error": "wallet is banned"}
    finally:
        tracker.stop()


def test_tracker_ban_registration_guard_reads_contract_state():
    """The tracker registration guard reads ban status from the contract facade."""
    contracts = LocalSolanaContracts()
    contracts.registry.ban_wallet("wallet-b")

    assert _registration_ban_error(contracts, "wallet-b") == "wallet is banned"
    assert _registration_ban_error(contracts, "wallet-a") is None
    assert _registration_ban_error(None, "wallet-b") is None


def test_gateway_route_guard_rejects_banned_wallet_metadata():
    """Gateway rejects tracker routes that include a wallet now banned on-chain."""
    contracts = LocalSolanaContracts()
    contracts.registry.ban_wallet("wallet-b")

    banned_wallet = _banned_route_wallet(
        contracts,
        [{"endpoint": "http://node", "wallet_address": "wallet-b"}],
    )

    assert banned_wallet == "wallet-b"


def test_gateway_routes_even_when_compute_node_has_low_stake():
    """Routing is not based on a compute node's token balance."""
    contracts = LocalSolanaContracts()
    contracts.registry.submit_stake("wallet-a", 1_000)
    contracts.registry.submit_stake("wallet-b", 5)
    tracker = TrackerServer()
    tracker_port = tracker.start()
    tracker_url = f"http://127.0.0.1:{tracker_port}"

    node_a = StubNodeServer(shard_start=0, shard_end=15, is_last_shard=False)
    port_a = node_a.start()
    node_b = StubNodeServer(shard_start=16, shard_end=31, is_last_shard=True)
    port_b = node_b.start()

    _post_json(
        f"{tracker_url}/v1/nodes/register",
        {"endpoint": f"http://127.0.0.1:{port_a}", "shard_start": 0, "shard_end": 15,
         "hardware_profile": {}, "wallet_address": "wallet-a", "score": 1.0},
    )
    _post_json(
        f"{tracker_url}/v1/nodes/register",
        {"endpoint": f"http://127.0.0.1:{port_b}", "shard_start": 16, "shard_end": 31,
         "hardware_profile": {}, "wallet_address": "wallet-b", "score": 1.0},
    )

    gateway = GatewayServer(
        tracker_url=tracker_url,
        contracts=contracts,
        minimum_stake=100,
    )
    gateway_port = gateway.start()

    try:
        payload = json.dumps({
            "model": "stub-model",
            "messages": [{"role": "user", "content": "hello"}],
        }).encode()
        req = urllib.request.Request(
            f"http://127.0.0.1:{gateway_port}/v1/chat/completions",
            data=payload,
            headers={"Content-Type": "application/json", "Authorization": "Bearer api-key-a"},
            method="POST",
        )

        with urllib.request.urlopen(req) as resp:
            assert resp.status == 200
        assert node_b.received_activations
    finally:
        gateway.stop()
        node_a.stop()
        node_b.stop()
        tracker.stop()


def test_gateway_records_compute_attribution_after_inference_session():
    """Gateway records layer and token attribution for every routed node."""
    contracts = LocalSolanaContracts()
    contracts.registry.submit_stake("wallet-a", 1_000)
    contracts.registry.submit_stake("wallet-b", 1_000)
    contracts.payment.fund_api_key("api-key-a", lamports=10_000)
    tracker = TrackerServer(contracts=contracts, minimum_stake=100)
    tracker_port = tracker.start()
    tracker_url = f"http://127.0.0.1:{tracker_port}"

    node_a = StubNodeServer(shard_start=0, shard_end=15, is_last_shard=False)
    port_a = node_a.start()
    node_b = StubNodeServer(shard_start=16, shard_end=31, is_last_shard=True)
    port_b = node_b.start()

    _post_json(
        f"{tracker_url}/v1/nodes/register",
        {"endpoint": f"http://127.0.0.1:{port_a}", "shard_start": 0, "shard_end": 15,
         "hardware_profile": {}, "wallet_address": "wallet-a", "score": 1.0},
    )
    _post_json(
        f"{tracker_url}/v1/nodes/register",
        {"endpoint": f"http://127.0.0.1:{port_b}", "shard_start": 16, "shard_end": 31,
         "hardware_profile": {}, "wallet_address": "wallet-b", "score": 1.0},
    )

    gateway = GatewayServer(
        tracker_url=tracker_url,
        contracts=contracts,
        minimum_stake=100,
    )
    gateway_port = gateway.start()

    try:
        payload = json.dumps({
            "model": "stub-model",
            "messages": [{"role": "user", "content": "hello"}],
        }).encode()
        req = urllib.request.Request(
            f"http://127.0.0.1:{gateway_port}/v1/chat/completions",
            data=payload,
            headers={"Content-Type": "application/json", "Authorization": "Bearer api-key-a"},
            method="POST",
        )
        with urllib.request.urlopen(req) as resp:
            assert resp.status == 200

        records = contracts.payment.list_attributions()
        assert [(r.api_key, r.node_wallet, r.layer_start, r.layer_end, r.tokens) for r in records] == [
            ("api-key-a", "wallet-a", 0, 15, 1),
            ("api-key-a", "wallet-b", 16, 31, 1),
        ]
    finally:
        gateway.stop()
        node_a.stop()
        node_b.stop()
        tracker.stop()


def test_gateway_returns_json_503_when_tracker_unavailable():
    """Tracker-backed gateway failures are reported as JSON HTTP errors."""
    gateway = GatewayServer(tracker_url="http://127.0.0.1:9")
    gateway_port = gateway.start()
    try:
        payload = json.dumps({
            "model": "stub-model",
            "messages": [{"role": "user", "content": "hello"}],
        }).encode()
        req = urllib.request.Request(
            f"http://127.0.0.1:{gateway_port}/v1/chat/completions",
            data=payload,
            headers={"Content-Type": "application/json"},
            method="POST",
        )
        try:
            urllib.request.urlopen(req, timeout=5)
            raise AssertionError("Expected 503 for unavailable tracker")
        except urllib.error.HTTPError as exc:
            assert exc.code == 503
            assert exc.headers["Content-Type"] == "application/json"
            body = json.loads(exc.read())
            assert body == {"error": "tracker unavailable"}
    finally:
        gateway.stop()


def test_gateway_returns_json_503_for_malformed_tracker_route():
    """Malformed tracker route payloads do not drop the gateway connection."""

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

        def do_GET(self):
            body = json.dumps({"route": []}).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)

    fake_tracker = http.server.HTTPServer(("127.0.0.1", 0), MalformedRouteHandler)
    fake_thread = threading.Thread(target=fake_tracker.serve_forever, daemon=True)
    fake_thread.start()
    gateway = GatewayServer(tracker_url=f"http://127.0.0.1:{fake_tracker.server_address[1]}")
    gateway_port = gateway.start()
    try:
        payload = json.dumps({
            "model": "stub-model",
            "messages": [{"role": "user", "content": "hello"}],
        }).encode()
        req = urllib.request.Request(
            f"http://127.0.0.1:{gateway_port}/v1/chat/completions",
            data=payload,
            headers={"Content-Type": "application/json"},
            method="POST",
        )
        try:
            urllib.request.urlopen(req, timeout=5)
            raise AssertionError("Expected 503 for malformed tracker route")
        except urllib.error.HTTPError as exc:
            assert exc.code == 503
            assert exc.headers["Content-Type"] == "application/json"
            body = json.loads(exc.read())
            assert body == {"error": "tracker unavailable"}
    finally:
        gateway.stop()
        fake_tracker.shutdown()
        fake_tracker.server_close()
        fake_thread.join(timeout=1)


def test_two_node_pipeline_via_tracker():
    """End-to-end: nodes register with tracker; gateway discovers route dynamically."""
    tracker = TrackerServer()
    tracker_port = tracker.start()
    tracker_url = f"http://127.0.0.1:{tracker_port}"

    node_a = StubNodeServer(shard_start=0, shard_end=15, is_last_shard=False)
    port_a = node_a.start()

    node_b = StubNodeServer(shard_start=16, shard_end=31, is_last_shard=True)
    port_b = node_b.start()

    _post_json(
        f"{tracker_url}/v1/nodes/register",
        {"endpoint": f"http://127.0.0.1:{port_a}", "shard_start": 0, "shard_end": 15,
         "hardware_profile": {}, "score": 1.0},
    )
    _post_json(
        f"{tracker_url}/v1/nodes/register",
        {"endpoint": f"http://127.0.0.1:{port_b}", "shard_start": 16, "shard_end": 31,
         "hardware_profile": {}, "score": 1.0},
    )

    gateway = GatewayServer(tracker_url=tracker_url)
    gateway_port = gateway.start()

    try:
        payload = json.dumps({
            "model": "stub-model",
            "messages": [{"role": "user", "content": "hello"}],
        }).encode()

        req = urllib.request.Request(
            f"http://127.0.0.1:{gateway_port}/v1/chat/completions",
            data=payload,
            headers={"Content-Type": "application/json"},
            method="POST",
        )

        with urllib.request.urlopen(req) as resp:
            assert resp.status == 200
            body = json.loads(resp.read())

        assert body["id"] == "chatcmpl-stub"
        assert body["object"] == "chat.completion"
        assert body["choices"][0]["message"]["content"] == "stub response to: hello"
        assert node_b.received_activations, "node B did not receive activations from node A"
        assert not node_a.received_activations, "node A should not receive activations"
    finally:
        gateway.stop()
        node_a.stop()
        node_b.stop()
        tracker.stop()