neuron-tai/tests/test_native_shard_protocol.py

"""Conformance tests for the native Shard protocol (ADR-0020, DGR-002).

Three layers are tested, and they are not the same thing:

1. The *schema* — asserted against the descriptor, not against the Python
   helpers. If a field the protocol promises to carry were dropped from the
   `.proto`, a test that only exercised the codec would still pass.
2. The *codec* — that a payload which is corrupt, short, holed, or byte-swapped
   is rejected rather than fed to a forward pass.
3. *Compatibility* — that an old build preserves fields a newer peer added, and
   that the committed cross-language vectors still encode as promised.
"""

from __future__ import annotations

import pathlib
import subprocess
import sys

import pytest

pytest.importorskip("google.protobuf", reason="protobuf runtime is required")

from google.protobuf import descriptor_pb2

from meshnet_node.activation_compression import CompressionPolicy
from meshnet_node.native_protocol import (
    DEFAULT_MAX_CHUNK_BYTES,
    DEFAULT_MAX_FRAGMENTS_PER_TENSOR,
    DEFAULT_MAX_PREFILL_CHUNK_TOKENS,
    DEFAULT_MAX_TENSORS_PER_BUNDLE,
    HIDDEN_STATES,
    PayloadCorrupt,
    ProtocolError,
    checksum_of,
    decode_bundle,
    decode_tensor,
    default_flow_control,
    encode_bundle,
    encode_tensor,
    negotiate_flow_control,
    pb,
    plan_prefill_chunks,
    validate_session_message_size,
)
from meshnet_node.native_protocol import conformance

REPO_ROOT = pathlib.Path(__file__).resolve().parent.parent
CPP_ROUNDTRIP = REPO_ROOT / "build/native" / conformance.CPP_ROUNDTRIP


# --- The schema itself ------------------------------------------------------


def test_service_exposes_capability_health_session_release_and_cancel():
    service = pb.DESCRIPTOR.services_by_name["ShardRuntime"]
    assert set(service.methods_by_name) == {
        "GetCapability",
        "Health",
        "Session",
        "Release",
        "Cancel",
    }


def test_session_is_one_long_lived_bidirectional_stream():
    session = pb.DESCRIPTOR.services_by_name["ShardRuntime"].methods_by_name["Session"]
    assert session.client_streaming, "the seam must stream requests"
    assert session.server_streaming, "the seam must stream responses"

    # Cancellation must not have to travel down a stream that flow control has
    # wedged, so it also exists as its own unary call.
    cancel = pb.DESCRIPTOR.services_by_name["ShardRuntime"].methods_by_name["Cancel"]
    assert not cancel.client_streaming and not cancel.server_streaming


def test_envelope_carries_every_field_the_protocol_promises():
    # Asserted against the descriptor: this is the acceptance criterion, and it
    # must fail if the .proto drops a field, not merely if the codec stops
    # setting one.
    fields = set(pb.Envelope.DESCRIPTOR.fields_by_name)
    assert {
        "schema_version",
        "work_id",
        "route_session_id",
        "route_epoch",
        "fingerprint",
        "shard_range",
        "phase",
        "position",
        "idempotency_step",
        "cache_expectation",
        "deadline_unix_nanos",
        "chunk",
    } <= fields

    assert {"model_artifact_digest", "runtime_recipe_digest"} <= set(
        pb.Fingerprint.DESCRIPTOR.fields_by_name
    )
    # Overlap-safe start (ADR-0012) is a distinct field from the registered one.
    assert {"start_layer", "end_layer", "effective_start_layer"} <= set(
        pb.ShardRange.DESCRIPTOR.fields_by_name
    )
    assert {"mode", "expected_past_len"} <= set(
        pb.CacheExpectation.DESCRIPTOR.fields_by_name
    )


def test_named_tensor_bundle_is_versioned_and_fully_described():
    assert "bundle_version" in pb.TensorBundle.DESCRIPTOR.fields_by_name
    assert {"name", "shape", "dtype", "byte_order", "compression", "checksum", "fragments"} <= set(
        pb.NamedTensor.DESCRIPTOR.fields_by_name
    )


def test_phases_and_error_codes_cover_the_lifecycle():
    assert {"PHASE_PREFILL", "PHASE_DECODE", "PHASE_RELEASE", "PHASE_CANCEL"} <= set(
        pb.Phase.keys()
    )
    # A cache miss is a first-class, recoverable outcome (ADR-0022), not a crash.
    assert {
        "ERROR_CODE_CACHE_MISS",
        "ERROR_CODE_FINGERPRINT_MISMATCH",
        "ERROR_CODE_EPOCH_STALE",
        "ERROR_CODE_SCHEMA_UNSUPPORTED",
        "ERROR_CODE_DEADLINE_EXCEEDED",
        "ERROR_CODE_FLOW_CONTROL_VIOLATION",
    } <= set(pb.ErrorCode.keys())


# --- Tensor bundle round-trip ----------------------------------------------


def test_tensor_round_trips_through_fragments():
    payload = bytes(range(256)) * 4
    tensor = encode_tensor(
        HIDDEN_STATES, payload, [1, 64, 8], pb.DTYPE_BFLOAT16, max_fragment_bytes=100
    )

    assert len(tensor.fragments) > 1, "the bound must actually split the payload"
    assert tensor.total_bytes == len(payload)
    assert decode_tensor(tensor) == payload


def test_bundle_round_trips_multiple_named_tensors():
    # An architecture boundary may need more than one tensor; that is why the
    # payload is a named bundle rather than a bare buffer.
    hidden = encode_tensor(HIDDEN_STATES, b"\x01\x02" * 8, [1, 8, 1], pb.DTYPE_BFLOAT16)
    positions = encode_tensor(
        "position_ids", (7).to_bytes(4, "little") * 8, [8], pb.DTYPE_INT32
    )
    bundle = encode_bundle([hidden, positions])

    restored = decode_bundle(pb.TensorBundle.FromString(bundle.SerializeToString()))
    assert restored == {
        HIDDEN_STATES: b"\x01\x02" * 8,
        "position_ids": (7).to_bytes(4, "little") * 8,
    }


def test_compressed_tensor_round_trips_and_keeps_its_uncompressed_checksum():
    pytest.importorskip("zstandard")
    # Highly compressible, and over the policy's minimum input size.
    payload = b"\x00" * 65536
    always = CompressionPolicy(min_input_bytes=0, min_savings_bytes=0, min_savings_ratio=0.0)
    tensor = encode_tensor(
        HIDDEN_STATES, payload, [1, 4096, 8], pb.DTYPE_BFLOAT16, policy=always
    )

    assert tensor.compression == pb.COMPRESSION_ZSTD
    assert sum(len(f.payload) for f in tensor.fragments) < len(payload)
    # The checksum covers the uncompressed bytes, so it stays valid whether or
    # not a hop chose to compress.
    assert tensor.checksum == checksum_of(payload)
    assert decode_tensor(tensor) == payload


# --- The codec refuses what it cannot account for --------------------------


def test_corrupt_payload_is_rejected_by_checksum():
    tensor = encode_tensor(HIDDEN_STATES, b"\xaa" * 32, [1, 16, 1], pb.DTYPE_BFLOAT16)
    # Flip one byte, as a lossy relay or a bad NIC would.
    tensor.fragments[0].payload = b"\xab" + tensor.fragments[0].payload[1:]

    with pytest.raises(PayloadCorrupt, match="CRC32C"):
        decode_tensor(tensor)


def test_missing_fragment_is_rejected_rather_than_silently_truncated():
    tensor = encode_tensor(
        HIDDEN_STATES, b"\xaa" * 64, [1, 32, 1], pb.DTYPE_BFLOAT16, max_fragment_bytes=16
    )
    del tensor.fragments[1]

    with pytest.raises(PayloadCorrupt):
        decode_tensor(tensor)


def test_fragment_hole_is_rejected():
    tensor = encode_tensor(
        HIDDEN_STATES, b"\xaa" * 64, [1, 32, 1], pb.DTYPE_BFLOAT16, max_fragment_bytes=16
    )
    # A gap in coverage: offsets no longer tile the body.
    tensor.fragments[2].byte_offset += 4

    with pytest.raises(PayloadCorrupt, match="expected"):
        decode_tensor(tensor)


def test_shape_that_disagrees_with_payload_is_rejected_at_encode():
    with pytest.raises(ProtocolError, match="carries"):
        encode_tensor(HIDDEN_STATES, b"\x01\x02", [1, 8, 8], pb.DTYPE_BFLOAT16)


def test_shape_that_disagrees_with_declared_bytes_is_rejected_at_decode():
    tensor = encode_tensor(HIDDEN_STATES, b"\xaa" * 32, [1, 16, 1], pb.DTYPE_BFLOAT16)
    # A peer claiming a larger tensor than its bytes describe.
    tensor.shape[1] = 32

    with pytest.raises(PayloadCorrupt, match="implies"):
        decode_tensor(tensor)


def test_big_endian_tensor_is_rejected_loudly():
    tensor = encode_tensor(HIDDEN_STATES, b"\xaa" * 32, [1, 16, 1], pb.DTYPE_BFLOAT16)
    tensor.byte_order = pb.BYTE_ORDER_BIG_ENDIAN

    # Byte-swapped activations would be plausible-looking garbage, so this is an
    # error rather than a best-effort read.
    with pytest.raises(ProtocolError, match="big-endian"):
        decode_tensor(tensor)


def test_bundle_from_a_newer_layout_is_refused():
    bundle = encode_bundle([])
    bundle.bundle_version = 99

    with pytest.raises(ProtocolError, match="not supported"):
        decode_bundle(bundle)


def test_bundle_with_unspecified_version_is_refused():
    with pytest.raises(ProtocolError, match="version 0"):
        decode_bundle(pb.TensorBundle())


def test_tensor_with_unspecified_compression_is_refused():
    tensor = encode_tensor(HIDDEN_STATES, b"\xaa" * 32, [1, 16, 1], pb.DTYPE_BFLOAT16)
    tensor.compression = pb.COMPRESSION_UNSPECIFIED

    with pytest.raises(ProtocolError, match="unspecified"):
        decode_tensor(tensor)


def test_tensor_with_unspecified_checksum_is_refused():
    tensor = encode_tensor(HIDDEN_STATES, b"\xaa" * 32, [1, 16, 1], pb.DTYPE_BFLOAT16)
    tensor.checksum.algorithm = pb.CHECKSUM_ALGORITHM_UNSPECIFIED

    with pytest.raises(ProtocolError, match="unspecified"):
        decode_tensor(tensor)


def test_declared_tensor_size_cannot_exceed_negotiated_chunk_bound():
    tensor = pb.NamedTensor(
        name=HIDDEN_STATES,
        shape=[1],
        dtype=pb.DTYPE_UINT8,
        byte_order=pb.BYTE_ORDER_LITTLE_ENDIAN,
        total_bytes=DEFAULT_MAX_CHUNK_BYTES + 1,
        compression=pb.COMPRESSION_NONE,
    )

    with pytest.raises(ProtocolError, match="negotiated chunk bound"):
        decode_tensor(tensor)


def test_stricter_session_negotiation_overrides_global_defaults():
    tensor = encode_tensor(
        HIDDEN_STATES,
        b"\xaa" * 32,
        [32],
        pb.DTYPE_UINT8,
        max_fragment_bytes=32,
    )

    with pytest.raises(ProtocolError, match="16-byte negotiated"):
        decode_tensor(tensor, max_chunk_bytes=16, max_fragment_bytes=16)
    with pytest.raises(ProtocolError, match="16-byte"):
        decode_bundle(
            encode_bundle([tensor]),
            max_chunk_bytes=16,
            max_fragment_bytes=16,
        )


def test_fragment_cannot_exceed_protocol_fragment_bound():
    tensor = encode_tensor(
        HIDDEN_STATES,
        b"\xaa" * (1024 * 1024 + 1),
        [1024 * 1024 + 1],
        pb.DTYPE_UINT8,
        max_fragment_bytes=1024 * 1024 + 1,
    )

    with pytest.raises(PayloadCorrupt, match="fragment larger"):
        decode_tensor(tensor)


def test_fragment_count_is_bounded_before_sorting_or_allocation():
    tensor = encode_tensor(
        HIDDEN_STATES,
        b"\xaa" * (DEFAULT_MAX_FRAGMENTS_PER_TENSOR + 1),
        [DEFAULT_MAX_FRAGMENTS_PER_TENSOR + 1],
        pb.DTYPE_UINT8,
        max_fragment_bytes=1,
        max_fragments=DEFAULT_MAX_FRAGMENTS_PER_TENSOR + 1,
    )

    with pytest.raises(PayloadCorrupt, match="fragments.*exceeding"):
        decode_tensor(tensor)


def test_tensor_count_is_bounded_before_bundle_decoding():
    tensors = [
        pb.NamedTensor(name=f"t-{index}")
        for index in range(DEFAULT_MAX_TENSORS_PER_BUNDLE + 1)
    ]

    with pytest.raises(ProtocolError, match="tensors.*exceeding"):
        encode_bundle(tensors)


@pytest.mark.parametrize(
    ("shape", "match"),
    [
        ([1] * 9, "rank"),
        ([-1], "dimension outside"),
        ([1 << 31], "dimension outside"),
    ],
)
def test_shape_rank_and_dimensions_are_bounded_before_payload_work(shape, match):
    with pytest.raises(ProtocolError, match=match):
        encode_tensor(HIDDEN_STATES, b"\x00", shape, pb.DTYPE_UINT8)


def test_complete_session_message_includes_envelope_in_byte_bound():
    small = pb.SessionRequest(chunk=pb.ActivationChunk(bundle=encode_bundle([])))
    assert validate_session_message_size(small, max_chunk_bytes=64) == small.ByteSize()

    oversized = pb.SessionRequest(
        chunk=pb.ActivationChunk(
            envelope=pb.Envelope(work_id="w" * 128),
            bundle=encode_bundle([]),
        )
    )
    with pytest.raises(ProtocolError, match="serialized session message"):
        validate_session_message_size(oversized, max_chunk_bytes=64)


def test_bundle_serialized_size_is_bounded():
    payload = b"\xaa" * (DEFAULT_MAX_CHUNK_BYTES // 2)
    first = encode_tensor("first", payload, [len(payload)], pb.DTYPE_UINT8)
    second = encode_tensor("second", payload, [len(payload)], pb.DTYPE_UINT8)

    with pytest.raises(ProtocolError, match="serialized tensor bundle"):
        encode_bundle([first, second])


def test_compressed_tensor_cannot_expand_past_declared_size():
    pytest.importorskip("zstandard")
    payload = b"\x00" * 65536
    always = CompressionPolicy(min_input_bytes=0, min_savings_bytes=0, min_savings_ratio=0.0)
    tensor = encode_tensor(
        HIDDEN_STATES, payload, [len(payload)], pb.DTYPE_UINT8, policy=always
    )
    assert tensor.compression == pb.COMPRESSION_ZSTD
    # Simulate a hostile frame that advertises a small output while expanding to
    # the original 64 KiB. Shape and total_bytes agree, so only the bounded
    # decompressor protects the receiver.
    tensor.shape[:] = [1024]
    tensor.total_bytes = 1024

    with pytest.raises(PayloadCorrupt, match="bounded zstd"):
        decode_tensor(tensor)


def test_generated_runtime_floors_match_committed_stubs():
    metadata = (REPO_ROOT / "packages/node/pyproject.toml").read_text()
    assert '"grpcio>=1.82.1"' in metadata
    assert '"protobuf>=7.35.0"' in metadata


# --- Bounded prefill chunking and the decode fast path ----------------------


def test_prefill_is_split_into_bounded_token_aligned_chunks():
    chunks = plan_prefill_chunks(2048)

    assert all(c.token_count <= DEFAULT_MAX_PREFILL_CHUNK_TOKENS for c in chunks)
    assert sum(c.token_count for c in chunks) == 2048
    # Contiguous, token-aligned, and no split falls mid-token.
    assert [c.first_position for c in chunks] == [
        i * DEFAULT_MAX_PREFILL_CHUNK_TOKENS for i in range(len(chunks))
    ]
    assert [c.final_chunk for c in chunks] == [False] * (len(chunks) - 1) + [True]


def test_final_prefill_chunk_carries_the_remainder():
    chunks = plan_prefill_chunks(300, max_tokens=128)

    assert [c.token_count for c in chunks] == [128, 128, 44]
    assert chunks[-1].chunk_info().final_chunk
    assert chunks[0].position().first_position == 0
    assert chunks[-1].position().first_position == 256


def test_empty_prefill_is_refused():
    with pytest.raises(ProtocolError):
        plan_prefill_chunks(0)


def test_decode_fast_path_is_much_smaller_than_a_full_envelope_chunk():
    hidden = b"\x01\x02" * 8  # one token, hidden=8, bfloat16
    tensor = encode_tensor(HIDDEN_STATES, hidden, [1, 1, 8], pb.DTYPE_BFLOAT16)

    fast = pb.SessionRequest(
        decode=pb.DecodeStep(
            idempotency_step=9,
            position=1024,
            expected_past_len=1024,
            tensor=tensor,
            work_id="work-7f3a",
        )
    )
    # The same single token carried the long way, repeating identity that the
    # handshake already fixed for the life of the stream.
    full = pb.SessionRequest(
        chunk=pb.ActivationChunk(
            envelope=conformance.canonical_session_request().chunk.envelope,
            bundle=encode_bundle([tensor]),
        )
    )

    assert len(fast.SerializeToString()) * 2 < len(full.SerializeToString())
    assert decode_tensor(fast.decode.tensor) == hidden


def test_flow_control_defaults_bound_the_queue_and_the_message():
    limits = default_flow_control()

    assert limits.max_prefill_chunk_tokens == DEFAULT_MAX_PREFILL_CHUNK_TOKENS
    assert limits.max_chunk_bytes == DEFAULT_MAX_CHUNK_BYTES
    assert limits.max_inflight_chunks > 0


def test_flow_control_negotiation_takes_the_strictest_bound():
    proposed = pb.FlowControl(
        credits_granted=64, max_inflight_chunks=64, max_chunk_bytes=64 << 20,
        max_prefill_chunk_tokens=1024,
    )
    settled = negotiate_flow_control(proposed, default_flow_control())

    # A sender cannot talk a worker into unbounded queues by proposing a large
    # window: neither peer can raise the other's ceiling.
    assert settled.max_inflight_chunks == default_flow_control().max_inflight_chunks
    assert settled.credits_granted <= settled.max_inflight_chunks
    assert settled.max_chunk_bytes == DEFAULT_MAX_CHUNK_BYTES
    assert settled.max_prefill_chunk_tokens == DEFAULT_MAX_PREFILL_CHUNK_TOKENS


def test_initial_credits_never_exceed_negotiated_inflight_limit():
    proposed = pb.FlowControl(credits_granted=64, max_inflight_chunks=64)
    limits = pb.FlowControl(credits_granted=64, max_inflight_chunks=8)

    settled = negotiate_flow_control(proposed, limits)

    assert settled.credits_granted == 8
    assert settled.max_inflight_chunks == 8


# --- Compatibility ----------------------------------------------------------


def test_committed_vectors_still_encode_as_promised():
    # The C++ test asserts against these exact bytes. If a schema change alters
    # the canonical encoding, it must be acknowledged by regenerating them.
    golden = (conformance.TESTDATA_DIR / conformance.GOLDEN_SESSION_REQUEST).read_bytes()
    assert conformance.serialize(conformance.canonical_session_request()) == golden

    report = (conformance.TESTDATA_DIR / conformance.GOLDEN_CAPABILITY_REPORT).read_bytes()
    assert conformance.serialize(conformance.canonical_capability_report()) == report


def test_golden_session_request_round_trips_with_every_field_intact():
    golden = (conformance.TESTDATA_DIR / conformance.GOLDEN_SESSION_REQUEST).read_bytes()
    request = pb.SessionRequest.FromString(golden)
    envelope = request.chunk.envelope

    assert envelope.work_id == conformance.WORK_ID
    assert envelope.route_session_id == conformance.ROUTE_SESSION_ID
    assert envelope.route_epoch == conformance.ROUTE_EPOCH
    assert envelope.idempotency_step == conformance.IDEMPOTENCY_STEP
    assert envelope.shard_range.effective_start_layer == conformance.EFFECTIVE_START_LAYER
    assert envelope.phase == pb.PHASE_PREFILL
    assert envelope.cache_expectation.mode == pb.CACHE_MODE_PREFILL
    assert envelope.deadline_unix_nanos == conformance.DEADLINE_UNIX_NANOS
    assert decode_bundle(request.chunk.bundle) == {
        HIDDEN_STATES: conformance.canonical_payload()
    }
    assert request.SerializeToString(deterministic=True) == golden


def test_unknown_fields_from_a_newer_peer_survive_a_forwarding_hop():
    # A Shard forwards activations onward. If it silently dropped fields a newer
    # peer added, it would corrupt a route it is merely a waypoint on.
    golden = (conformance.TESTDATA_DIR / conformance.GOLDEN_SESSION_REQUEST).read_bytes()
    future_field = b"\xb8\xe0\x04\xb9\x60"  # field 9999, varint 12345
    request = pb.SessionRequest.FromString(golden + future_field)

    assert request.chunk.envelope.work_id == conformance.WORK_ID
    assert request.SerializeToString() == golden + future_field


def test_a_message_missing_newer_field_groups_still_parses():
    sparse = pb.SessionRequest(chunk=pb.ActivationChunk(envelope=pb.Envelope(work_id="w")))
    parsed = pb.SessionRequest.FromString(sparse.SerializeToString())

    assert parsed.chunk.envelope.work_id == "w"
    assert parsed.chunk.envelope.route_epoch == 0
    assert parsed.chunk.envelope.phase == pb.PHASE_UNSPECIFIED


def test_retired_fragment_field_stays_reserved():
    # `uncompressed_size` (field 5) was removed because NamedTensor.total_bytes
    # is the single source of truth. The number stays reserved so it can never
    # be recycled for a different meaning — a recycled number is the one schema
    # change that old and new peers cannot detect, because the bytes still parse.
    descriptor = descriptor_pb2.DescriptorProto()
    pb.TensorFragment.DESCRIPTOR.CopyToProto(descriptor)

    assert 5 not in {field.number for field in descriptor.field}
    assert any(r.start <= 5 < r.end for r in descriptor.reserved_range)
    assert "uncompressed_size" in descriptor.reserved_name


def test_a_peer_still_sending_the_retired_field_does_not_corrupt_the_tensor():
    # An older peer that still sets field 5 must be parsed, not rejected: the
    # value lands in unknown fields and the payload is unaffected.
    tensor = encode_tensor(HIDDEN_STATES, b"\xaa" * 32, [1, 16, 1], pb.DTYPE_BFLOAT16)
    wire = tensor.SerializeToString()
    fragment = pb.TensorFragment.FromString(tensor.fragments[0].SerializeToString() + b"\x28\x20")

    assert fragment.payload == tensor.fragments[0].payload
    assert decode_tensor(pb.NamedTensor.FromString(wire)) == b"\xaa" * 32


def test_generated_python_stubs_match_the_proto():
    pytest.importorskip("grpc_tools", reason="protoc toolchain is required to verify")
    result = subprocess.run(
        [sys.executable, str(REPO_ROOT / "scripts/generate_native_protocol.py"), "--check"],
        capture_output=True,
        text=True,
    )
    assert result.returncode == 0, result.stdout + result.stderr


@pytest.mark.skipif(
    not CPP_ROUNDTRIP.is_file(),
    reason="build the C++ conformance test to check cross-language agreement",
)
def test_cpp_and_python_agree_byte_for_byte():
    # Written by the C++ conformance test: it parsed the golden bytes into its
    # own object model and serialized them back. Byte equality means both
    # languages encode every field of this schema identically.
    golden = (conformance.TESTDATA_DIR / conformance.GOLDEN_SESSION_REQUEST).read_bytes()
    assert CPP_ROUNDTRIP.read_bytes() == golden