"""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