neuron-tai/tests/test_tracker_capability_admission.py

"""NCA-004 tests: the tracker records a node's capability proof and routes only
to nodes whose proof admits what they advertise (ADR-0023).

Model ids here are arbitrary and made up on purpose: nothing in the admission or
routing path may branch on a vendor, model or kernel name.
"""

import json
import time
import urllib.error
import urllib.request

import pytest

from meshnet_tracker.capability import (
    MIN_CATALOGUE_VERSION,
    POLICY_COMPAT,
    POLICY_ENFORCE,
    STATE_ABSENT,
    STATE_ADMITTED,
    STATE_CATALOGUE_INCOMPATIBLE,
    STATE_FAILED,
    STATE_INVALID,
    STATE_MODEL_MISMATCH,
    STATE_RECIPE_MISMATCH,
    STATE_SHARD_MISMATCH,
    STATE_STALE,
    evaluate_report,
    policy_from_env,
)
from meshnet_tracker.server import (
    TrackerServer,
    _NodeEntry,
    _capability_routable,
    _node_admission,
    _select_route,
)

MODEL = "arbitrary-labs/oracle-9b"
SHORT = "oracle-9b"
LAYERS = 32


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 _report(
    *,
    model_id: str = MODEL,
    start: int = 0,
    end: int = 15,
    status: str = "passed",
    validated_at: float | None = None,
    recipe_id: str = "baseline",
    recipe_version: str = "1",
    catalogue_version: str = MIN_CATALOGUE_VERSION,
    schema_version: int = 1,
    device: str = "cpu",
    diagnostics: list | None = None,
) -> dict:
    """A capability report shaped exactly as `meshnet_node.capability` emits it."""
    return {
        "schema_version": schema_version,
        "model": {"model_id": model_id, "revision": None, "config_fingerprint": None},
        "shard": {"start": start, "end": end},
        "recipe": {
            "recipe_id": recipe_id,
            "recipe_version": recipe_version,
            "catalogue_version": catalogue_version,
        },
        "backend": {
            "backend_id": "torch-transformers",
            "device": device,
            "device_name": None,
            "quantization": "bfloat16",
            "runtime": {},
        },
        "status": status,
        "validated_at": time.time() if validated_at is None else validated_at,
        "duration_ms": 42,
        "diagnostics": list(diagnostics or []),
    }


def _registration(
    port: int,
    *,
    start: int = 0,
    end: int = 15,
    report: dict | None = "default",  # type: ignore[assignment]
    recipe_id: str | None = "baseline",
    recipe_version: str | None = "1",
    benchmark_tokens_per_sec: float = 10.0,
) -> dict:
    payload: dict = {
        "endpoint": f"http://127.0.0.1:{port}",
        "model": SHORT,
        "hf_repo": MODEL,
        "num_layers": LAYERS,
        "shard_start": start,
        "shard_end": end,
        "hardware_profile": {},
        "score": 1.0,
        "tracker_mode": start == 0,
        "benchmark_tokens_per_sec": benchmark_tokens_per_sec,
    }
    if report == "default":
        report = _report(start=start, end=end)
    if report is not None:
        payload["capability_report"] = report
    if recipe_id is not None:
        payload["recipe_id"] = recipe_id
    if recipe_version is not None:
        payload["recipe_version"] = recipe_version
    return payload


def _always(_reported: str) -> bool:
    return True


def _evaluate(report, **kwargs) -> object:
    kwargs.setdefault("model_matches", _always)
    kwargs.setdefault("advertised_model", MODEL)
    kwargs.setdefault("shard_start", 0)
    kwargs.setdefault("shard_end", 15)
    return evaluate_report(report, **kwargs)


# --------------------------------------------------------------- report verdicts


def test_a_passing_report_that_covers_the_registration_is_admitted():
    "A proof matching the advertised model and shard admits the node.\n\nTags: routing, tracker"
    state = _evaluate(_report())
    assert state.state == STATE_ADMITTED
    assert state.proven
    assert state.model_id == MODEL
    assert (state.shard_start, state.shard_end) == (0, 15)


def test_a_missing_report_is_absent_not_admitted():
    "No proof is recorded as `absent` — never silently treated as proven.\n\nTags: routing, tracker"
    state = _evaluate(None)
    assert state.state == STATE_ABSENT
    assert not state.proven


def test_a_failed_report_is_recorded_as_failed():
    "A node that failed its own validation is not routable.\n\nTags: routing, tracker"
    state = _evaluate(_report(status="failed", diagnostics=["out of memory on device"]))
    assert state.state == STATE_FAILED
    assert not state.proven
    assert "out of memory" in state.detail


def test_a_report_for_a_different_model_is_a_model_mismatch():
    "A proof for another artifact proves nothing about this one.\n\nTags: routing, tracker"
    state = evaluate_report(
        _report(model_id="other-org/unrelated-3b"),
        model_matches=lambda reported: reported == MODEL,
        advertised_model=MODEL,
        shard_start=0,
        shard_end=15,
    )
    assert state.state == STATE_MODEL_MISMATCH


def test_a_report_for_a_different_shard_is_a_shard_mismatch():
    "A proof for layers 0–15 does not admit a node advertising 16–31.\n\nTags: routing, tracker"
    state = _evaluate(_report(start=0, end=15), shard_start=16, shard_end=31)
    assert state.state == STATE_SHARD_MISMATCH


def test_a_report_for_a_different_recipe_than_the_node_declares_is_a_recipe_mismatch():
    "The proof must be for the recipe the node says it serves with.\n\nTags: routing, tracker"
    state = _evaluate(_report(recipe_id="eager-attention"), declared_recipe_id="baseline")
    assert state.state == STATE_RECIPE_MISMATCH

    versioned = _evaluate(
        _report(recipe_version="1"),
        declared_recipe_id="baseline",
        declared_recipe_version="2",
    )
    assert versioned.state == STATE_RECIPE_MISMATCH


def test_an_older_recipe_catalogue_is_incompatible():
    "Recipe ids from a catalogue older than the tracker's minimum cannot be matched.\n\nTags: routing, tracker"
    state = _evaluate(_report(catalogue_version="2025.01.1"))
    assert state.state == STATE_CATALOGUE_INCOMPATIBLE
    assert MIN_CATALOGUE_VERSION in state.detail

    newer = _evaluate(_report(catalogue_version="2099.12.9"))
    assert newer.state == STATE_ADMITTED


def test_an_unparseable_catalogue_version_is_incompatible():
    "A catalogue version that cannot be compared cannot be shown to be new enough.\n\nTags: routing, tracker"
    assert _evaluate(_report(catalogue_version="rolling")).state == STATE_CATALOGUE_INCOMPATIBLE


def test_a_stale_report_is_not_admitted():
    "A proof older than the freshness bound must be re-validated before routing.\n\nTags: routing, tracker"
    state = _evaluate(_report(validated_at=time.time() - 3600), max_age_seconds=900.0)
    assert state.state == STATE_STALE


def test_a_future_dated_report_is_not_admitted():
    "A proof from the future is a broken clock, not a fresh proof.\n\nTags: routing, tracker"
    state = _evaluate(_report(validated_at=time.time() + 3600))
    assert state.state == STATE_STALE
    assert "clock" in state.detail


def test_a_report_from_an_unknown_schema_version_is_invalid():
    "The tracker refuses to interpret a report layout it does not read.\n\nTags: routing, tracker"
    assert _evaluate(_report(schema_version=99)).state == STATE_INVALID


@pytest.mark.parametrize(
    "payload",
    [
        "not-an-object",
        {},
        {"schema_version": 1},
        {**_report(), "model": {"model_id": ""}},
        {**_report(), "shard": {"start": -1, "end": 3}},
        {**_report(), "validated_at": "yesterday"},
        {**_report(), "status": None},
    ],
    ids=["not-object", "empty", "header-only", "blank-model", "negative-shard",
         "bad-timestamp", "missing-status"],
)
def test_a_malformed_report_is_invalid_and_never_admitted(payload):
    "Malformed proof is rejected by the schema check, not by a later coincidence.\n\nTags: routing, tracker"
    state = _evaluate(payload)
    assert state.state == STATE_INVALID
    assert not state.proven


def test_recorded_detail_carries_no_credentials_from_node_diagnostics():
    "Operator-facing admission detail is sanitized; a leaked token never reaches it.\n\nTags: routing, security, tracker"
    state = _evaluate(
        _report(
            status="failed",
            diagnostics=["download failed: authorization=hf_abcdefghijklmnopqrstuvwxyz01"],
        )
    )
    assert state.state == STATE_FAILED
    assert "hf_abcdefghijklmnopqrstuvwxyz01" not in state.detail
    assert "hf_abcdefghijklmnopqrstuvwxyz01" not in json.dumps(state.to_dict())
    assert "[redacted]" in state.detail


# ---------------------------------------------------------- compatibility policy


def test_compat_policy_routes_a_legacy_node_but_never_a_broken_proof():
    "Older nodes (no proof) keep routing under `compat`; a bad proof never does.\n\nTags: routing, tracker"
    absent = _evaluate(None)
    failed = _evaluate(_report(status="failed"))

    assert absent.routable_under(POLICY_COMPAT)
    assert not absent.routable_under(POLICY_ENFORCE)
    assert not failed.routable_under(POLICY_COMPAT)
    assert not failed.routable_under(POLICY_ENFORCE)


def test_the_policy_is_read_from_the_environment_and_defaults_to_compat(monkeypatch):
    "The rollout switch is explicit and defaults to the compatible behaviour.\n\nTags: config, routing, tracker"
    monkeypatch.delenv("MESHNET_TRACKER_CAPABILITY_POLICY", raising=False)
    assert policy_from_env() == POLICY_COMPAT
    monkeypatch.setenv("MESHNET_TRACKER_CAPABILITY_POLICY", "enforce")
    assert policy_from_env() == POLICY_ENFORCE
    monkeypatch.setenv("MESHNET_TRACKER_CAPABILITY_POLICY", "nonsense")
    assert policy_from_env() == POLICY_COMPAT


# ------------------------------------------------------------- the routing gate


def _entry(node_id: str, start: int, end: int, report: dict | None, **kwargs) -> _NodeEntry:
    from meshnet_tracker.capability import evaluate_report as _eval

    entry = _NodeEntry(
        node_id=node_id,
        endpoint=f"http://127.0.0.1:{9000 + int(node_id[-1])}",
        shard_start=start,
        shard_end=end,
        model=SHORT,
        shard_checksum=None,
        hardware_profile={},
        wallet_address=None,
        score=1.0,
        hf_repo=MODEL,
        num_layers=LAYERS,
        capability=_eval(
            report,
            model_matches=lambda reported: reported == MODEL,
            advertised_model=MODEL,
            shard_start=start,
            shard_end=end,
        ),
        **kwargs,
    )
    return entry


def test_route_selection_drops_every_unadmitted_candidate_under_enforce():
    "Absent, failed, stale and mismatched candidates are all excluded.\n\nTags: routing, tracker"
    good = _entry("node-1", 0, 31, _report(start=0, end=31))
    unproven = _entry("node-2", 0, 31, None)
    failed = _entry("node-3", 0, 31, _report(start=0, end=31, status="failed"))
    stale = _entry("node-4", 0, 31, _report(start=0, end=31, validated_at=time.time() - 86400))
    wrong_model = _entry("node-5", 0, 31, _report(start=0, end=31, model_id="someone/else-1b"))

    route, error = _select_route(
        [unproven, failed, stale, wrong_model, good], 0, 31, policy=POLICY_ENFORCE
    )
    assert not error
    assert [n.node_id for n in route] == ["node-1"]

    only_bad, error = _select_route([unproven, failed, stale], 0, 31, policy=POLICY_ENFORCE)
    assert only_bad == []
    assert "no route available" in error


def test_a_node_reassigned_to_a_shard_it_never_proved_stops_routing():
    "The proof does not travel with a tracker reassignment.\n\nTags: routing, tracker"
    node = _entry("node-1", 0, 15, _report(start=0, end=15))
    assert _node_admission(node).state == STATE_ADMITTED

    node.shard_start, node.shard_end = 16, 31  # tracker rebalanced it
    assert _node_admission(node).state == STATE_SHARD_MISMATCH
    assert not _node_admission(node).routable_under(POLICY_COMPAT)


def test_admitted_candidates_keep_coverage_first_and_throughput_routing():
    "Gating removes candidates; among the survivors routing is unchanged.\n\nTags: routing, tracker"
    head = _entry("node-1", 0, 15, _report(start=0, end=15))
    slow_tail = _entry(
        "node-2", 16, 31, _report(start=16, end=31), benchmark_tokens_per_sec=5.0
    )
    fast_tail = _entry(
        "node-3", 16, 31, _report(start=16, end=31), benchmark_tokens_per_sec=50.0
    )

    route, error = _select_route(
        [head, slow_tail, fast_tail], 0, 31, policy=POLICY_ENFORCE
    )
    assert not error
    # Coverage first (head, then a tail), and the faster of the two tied tails.
    assert [n.node_id for n in route] == ["node-1", "node-3"]


# ---------------------------------------------------------------- over the wire


def test_an_enforcing_tracker_routes_a_proven_node_and_excludes_an_unproven_one():
    "End to end: a proof is required to appear in a route.\n\nTags: http, routing, tracker"
    tracker = TrackerServer(capability_policy=POLICY_ENFORCE)
    port = tracker.start()
    try:
        base = f"http://127.0.0.1:{port}"
        _post_json(f"{base}/v1/nodes/register", _registration(9101, start=0, end=15))
        # A tail that presents no proof at all: the route cannot complete.
        _post_json(
            f"{base}/v1/nodes/register",
            _registration(9102, start=16, end=31, report=None),
        )

        with pytest.raises(urllib.error.HTTPError) as exc:
            _get_json(f"{base}/v1/route?model={SHORT}")
        assert exc.value.code == 503

        # Now the tail proves itself and the same route resolves.
        _post_json(f"{base}/v1/nodes/register", _registration(9102, start=16, end=31))
        route = _get_json(f"{base}/v1/route?model={SHORT}")["route"]
        assert route == ["http://127.0.0.1:9101", "http://127.0.0.1:9102"]
    finally:
        tracker.stop()


@pytest.mark.parametrize(
    "bad_report",
    [
        _report(start=16, end=31),  # proves the wrong shard
        _report(model_id="unrelated/other-7b"),  # proves the wrong model
        _report(status="failed"),
        _report(validated_at=time.time() - 86400),  # stale
        {"schema_version": 1, "model": {}},  # malformed
    ],
    ids=["shard-mismatch", "model-mismatch", "failed", "stale", "invalid"],
)
def test_an_enforcing_tracker_never_routes_a_node_whose_proof_does_not_cover_it(bad_report):
    "A proof for something else is worth exactly as much as no proof.\n\nTags: http, routing, tracker"
    tracker = TrackerServer(capability_policy=POLICY_ENFORCE)
    port = tracker.start()
    try:
        base = f"http://127.0.0.1:{port}"
        resp = _post_json(
            f"{base}/v1/nodes/register",
            _registration(9111, start=0, end=31, report=bad_report),
        )
        # Registration still succeeds — the operator must be able to see the node.
        assert resp["node_id"]
        assert resp["capability"]["routable"] is False
        assert resp["capability"]["state"] != STATE_ADMITTED

        with pytest.raises(urllib.error.HTTPError) as exc:
            _get_json(f"{base}/v1/route?model={SHORT}")
        assert exc.value.code in (404, 503)
    finally:
        tracker.stop()


def test_a_compat_tracker_routes_a_legacy_node_that_sends_no_report():
    "Documented rollout policy: pre-capability nodes keep working under `compat`.\n\nTags: http, routing, tracker"
    tracker = TrackerServer(capability_policy=POLICY_COMPAT)
    port = tracker.start()
    try:
        base = f"http://127.0.0.1:{port}"
        _post_json(
            f"{base}/v1/nodes/register",
            _registration(9121, start=0, end=31, report=None, recipe_id=None, recipe_version=None),
        )
        route = _get_json(f"{base}/v1/route?model={SHORT}")["route"]
        assert route == ["http://127.0.0.1:9121"]
    finally:
        tracker.stop()


def test_a_compat_tracker_still_refuses_a_node_that_presents_a_failed_proof():
    "Compatibility grandfathers silence, not a proof of failure.\n\nTags: http, routing, tracker"
    tracker = TrackerServer(capability_policy=POLICY_COMPAT)
    port = tracker.start()
    try:
        base = f"http://127.0.0.1:{port}"
        _post_json(
            f"{base}/v1/nodes/register",
            _registration(9131, start=0, end=31, report=_report(start=0, end=31, status="failed")),
        )
        with pytest.raises(urllib.error.HTTPError) as exc:
            _get_json(f"{base}/v1/route?model={SHORT}")
        assert exc.value.code == 503
    finally:
        tracker.stop()


def test_a_replicated_registration_carries_its_verdict_to_a_follower():
    "A proven node must not be routable on the leader and dark on every follower.\n\nTags: cluster, routing, tracker"
    tracker = TrackerServer(capability_policy=POLICY_ENFORCE)
    proven = _registration(9151, start=0, end=31, report=_report(start=0, end=31))
    proven["node_id"] = "follower-node-1"
    unproven = _registration(9152, start=0, end=31, report=None)
    unproven["node_id"] = "follower-node-2"

    tracker._raft_apply("register", proven)
    tracker._raft_apply("register", unproven)

    admitted = tracker._registry["follower-node-1"]
    assert admitted.capability.state == STATE_ADMITTED
    assert _capability_routable(admitted, POLICY_ENFORCE)

    absent = tracker._registry["follower-node-2"]
    assert absent.capability.state == STATE_ABSENT
    assert not _capability_routable(absent, POLICY_ENFORCE)


def test_the_network_map_exposes_the_admission_state_of_every_node():
    "The operator view answers 'why is my node not routing' without raw internals.\n\nTags: http, routing, tracker"
    tracker = TrackerServer(capability_policy=POLICY_ENFORCE)
    port = tracker.start()
    try:
        base = f"http://127.0.0.1:{port}"
        _post_json(f"{base}/v1/nodes/register", _registration(9141, start=0, end=15))
        _post_json(
            f"{base}/v1/nodes/register",
            _registration(
                9142,
                start=16,
                end=31,
                report=_report(
                    start=16,
                    end=31,
                    status="failed",
                    diagnostics=["load failed: token=hf_abcdefghijklmnopqrstuvwx1234"],
                ),
            ),
        )

        network = _get_json(f"{base}/v1/network/map")
        assert network["capability_policy"] == POLICY_ENFORCE
        by_endpoint = {n["endpoint"]: n["capability"] for n in network["nodes"]}

        proven = by_endpoint["http://127.0.0.1:9141"]
        assert proven["state"] == STATE_ADMITTED
        assert proven["routable"] is True
        assert proven["model_id"] == MODEL
        assert (proven["shard_start"], proven["shard_end"]) == (0, 15)
        assert proven["recipe_id"] == "baseline"
        assert proven["device"] == "cpu"

        broken = by_endpoint["http://127.0.0.1:9142"]
        assert broken["state"] == STATE_FAILED
        assert broken["routable"] is False
        assert broken["detail"]

        raw = json.dumps(network)
        assert "hf_abcdefghijklmnopqrstuvwx1234" not in raw
        assert "Traceback" not in raw
    finally:
        tracker.stop()