neuron-tai/tests/test_native_identity_emission.py

"""DGR-003 production-native identity emission boundary tests."""

from __future__ import annotations

import pytest

from meshnet_node.doctor import DoctorSelection, validate_loaded_backend
from meshnet_node.native_backend import (
    ImmutableArtifactPin,
    NativeIdentityInputs,
    NativeLoadedArtifactReport,
    NativeNumericalRecipe,
    NativeSessionRejected,
    NativeWorkerBackendAdapter,
    shard_identity_from_native_report,
)
from meshnet_node.native_protocol import SCHEMA_VERSION, pb
from meshnet_node.recipe_manifest import parse_recipe_manifest
from meshnet_tracker.capability import STATE_UNCERTIFIED, evaluate_report


def _digest(letter: str) -> str:
    return letter * 64


def _inputs(**changes: object) -> NativeIdentityInputs:
    report = NativeLoadedArtifactReport(
        owned_start_layer=2,
        owned_end_layer=6,
        mapped_bytes=1024,
        resident_bytes=768,
        registered_bytes=640,
        architecture="llama",
        architecture_digest=_digest("a"),
        layer_count=8,
    )
    recipe = NativeNumericalRecipe(
        weight_quantization="Q4_K_M",
        activation_dtype="bfloat16",
        compute_dtype="float32",
        kv_dtype="q8_0",
        kv_layout="llama-kv-v1",
        architecture_adapter="dense-llama-v1",
        backend_id="llama-cpp",
        runtime_version="llama.cpp:e920c523",
        recipe_id="native",
        recipe_version="1",
        catalogue_version="2026.07.1",
    )
    values: dict[str, object] = {
        "loaded_artifact": report,
        "artifact_pin": ImmutableArtifactPin(
            artifact_id="acme/llama.gguf",
            revision="0123456789abcdef",
            content_digest=_digest("b"),
        ),
        "tokenizer_revision": "abcdef0123456789",
        "numerical_recipe": recipe,
    }
    values.update(changes)
    return NativeIdentityInputs(**values)  # type: ignore[arg-type]


def _open(adapter: NativeWorkerBackendAdapter, **changes: object) -> pb.SessionOpen:
    identity = adapter.identity
    fields: dict[str, object] = {
        "schema_version": SCHEMA_VERSION,
        "route_session_id": "tracker-session",
        "route_epoch": 4,
        "fingerprint": identity.fingerprint.to_proto(),
        "shard_range": pb.ShardRange(
            start_layer=identity.shard_start,
            end_layer=identity.shard_end,
            effective_start_layer=identity.shard_start,
        ),
    }
    fields.update(changes)
    return pb.SessionOpen(**fields)  # type: ignore[arg-type]


def test_native_identity_uses_loaded_report_not_a_caller_range():
    identity = shard_identity_from_native_report(_inputs())

    assert (identity.shard_start, identity.shard_end) == (2, 6)
    assert identity.artifact.architecture == "llama"
    assert identity.artifact.layer_count == 8


def test_native_identity_requires_an_immutable_pin_and_gguf_range():
    with pytest.raises(Exception, match="moving reference"):
        shard_identity_from_native_report(
            _inputs(artifact_pin=ImmutableArtifactPin("a", "main", _digest("b")))
        )
    with pytest.raises(Exception, match="outside GGUF"):
        NativeLoadedArtifactReport(0, 9, 1, 1, 1, "llama", _digest("a"), 8)


def test_native_worker_rejects_bad_session_open_before_session_acceptance():
    adapter = NativeWorkerBackendAdapter(_inputs())
    accepted = adapter.on_session_open(
        _open(adapter), expected_route_session_id="tracker-session", expected_route_epoch=4
    )
    assert accepted.fingerprint.SerializeToString() == adapter.identity.fingerprint.to_proto().SerializeToString()

    with pytest.raises(NativeSessionRejected) as rejected:
        adapter.on_session_open(
            _open(adapter, route_epoch=5),
            expected_route_session_id="tracker-session",
            expected_route_epoch=4,
        )
    assert rejected.value.error.code == pb.ERROR_CODE_EPOCH_STALE


def test_doctor_emits_native_identity_but_keeps_legacy_backend_dark():
    manifest = parse_recipe_manifest(
        {"schema_version": 1, "catalogue_version": "2026.07.1", "recipes": [
            {"id": "native", "version": "1", "backend_id": "llama-cpp"}
        ]}
    )
    selection = DoctorSelection("acme/llama.gguf", 2, 5)
    native = NativeWorkerBackendAdapter(_inputs())
    # The probe needs only the normal backend shape; identity is supplied by the adapter.
    native.hidden_size = 8
    native.is_head = False
    native.is_tail = False
    native.device = "cpu"
    native.forward_bytes = lambda *args, **kwargs: type("Payload", (), {"body": b"x", "shape": [1]})()
    result = validate_loaded_backend(native, selection, manifest.recipes[0], manifest)
    assert result.report.identity == native.identity
    assert result.report.model.revision == native.identity.artifact.revision
    assert result.report.model.config_fingerprint == "sha256:" + _digest("a")
    assert result.report.backend.quantization == "Q4_K_M"
    assert (result.report.shard.start, result.report.shard.end) == (2, 5)
    state = evaluate_report(
        result.report.to_dict(),
        model_matches=lambda value: value == "acme/llama.gguf",
        advertised_model="acme/llama.gguf",
        shard_start=2,
        shard_end=5,
        declared_recipe_id="native",
        declared_recipe_version="1",
        now=result.report.validated_at,
    )
    assert state.state == STATE_UNCERTIFIED

    class Legacy:
        hidden_size = 8
        is_head = False
        is_tail = False
        device = "cpu"

        @staticmethod
        def forward_bytes(*args, **kwargs):
            return type("Payload", (), {"body": b"x", "shape": [1]})()

    legacy = validate_loaded_backend(Legacy(), selection, manifest.recipes[0], manifest)
    assert legacy.report.identity is None