feat: add activation stream envelope

tests/test_activation_envelope.py

@@ -0,0 +1,130 @@
+"""DGR-021: versioned activation envelope for shard traffic."""
+
+from __future__ import annotations
+
+import json
+import hashlib
+
+import pytest
+
+from meshnet_node.model_backend import TensorPayload
+from meshnet_node.protocol import ActivationEnvelope, NamedTensor, TensorFragment
+
+
+def _payload(body: bytes = b"\x01\x02\x03\x04", *, shape=(1, 2, 1)) -> TensorPayload:
+    return TensorPayload(
+        body=body,
+        shape=list(shape),
+        attention_mask_header="1,2:AAAA",
+        position_ids_header="1,2:BBBB",
+        past_len=7,
+    )
+
+
+def test_envelope_roundtrips_deterministically():
+    payload = _payload()
+    env_a = payload.to_envelope(
+        name="activations",
+        request_id="req-1",
+        work_id="work-1",
+        route_session="route-1",
+        route_epoch=9,
+        shard_start=12,
+        effective_start=18,
+        phase="prefill",
+        position=0,
+        idempotency_step=0,
+        extensions={"zeta": 3, "alpha": "x"},
+    )
+    env_b = ActivationEnvelope.from_bytes(env_a.to_bytes())
+
+    assert env_a.to_bytes() == env_b.to_bytes()
+    assert env_b.request_id == "req-1"
+    assert env_b.work_id == "work-1"
+    assert env_b.route_session == "route-1"
+    assert env_b.shard_start == 12
+    assert env_b.effective_start == 18
+    assert env_b.phase == "prefill"
+    assert env_b.tensors[0].body() == payload.body
+    assert env_b.tensors[0].extensions["past_len"] == 7
+    assert env_b.extensions == {"alpha": "x", "zeta": 3}
+
+
+def test_fragmentation_and_checksums_are_bounded():
+    body = b"0123456789abcdef"
+    tensor = NamedTensor.from_bytes(
+        name="activations",
+        body=body,
+        shape=[1, 8, 1],
+        dtype="bfloat16",
+        max_fragment_bytes=5,
+    )
+
+    assert len(tensor.fragments) == 4
+    assert all(isinstance(fragment, TensorFragment) for fragment in tensor.fragments)
+    assert all(len(fragment.body) <= 5 for fragment in tensor.fragments)
+    assert tensor.body() == body
+    assert tensor.checksum == hashlib.sha256(body).hexdigest()
+
+
+def test_unknown_fields_are_preserved_on_roundtrip():
+    env = _payload().to_envelope(
+        name="activations",
+        request_id="req-2",
+        work_id="work-2",
+        route_session="route-2",
+        route_epoch=1,
+        shard_start=0,
+        effective_start=0,
+        phase="decode",
+        position=3,
+        idempotency_step=2,
+    )
+    raw = json.loads(env.to_bytes())
+    raw["future_top_level"] = {"x": 1}
+    raw["tensors"][0]["future_tensor_field"] = "ok"
+    roundtrip = ActivationEnvelope.from_bytes(json.dumps(raw).encode())
+
+    assert roundtrip.extensions["future_top_level"] == {"x": 1}
+    assert roundtrip.tensors[0].extensions["future_tensor_field"] == "ok"
+
+
+def test_size_limit_rejects_large_envelopes():
+    env = _payload(body=b"x" * 64).to_envelope(
+        name="activations",
+        request_id="req-3",
+        work_id="work-3",
+        route_session="route-3",
+        route_epoch=1,
+        shard_start=0,
+        effective_start=0,
+        phase="decode",
+        position=5,
+        idempotency_step=4,
+    )
+
+    with pytest.raises(ValueError, match="size limit"):
+        env.to_bytes(max_bytes=32)
+
+
+def test_tensor_payload_roundtrips_through_envelope():
+    payload = _payload(body=b"abcde", shape=(1, 5, 1))
+    env = payload.to_envelope(
+        name="activations",
+        request_id="req-4",
+        work_id="work-4",
+        route_session="route-4",
+        route_epoch=2,
+        shard_start=6,
+        effective_start=11,
+        phase="decode",
+        position=8,
+        idempotency_step=9,
+    )
+
+    restored = TensorPayload.from_envelope(env)
+    assert restored.body == payload.body
+    assert restored.shape == payload.shape
+    assert restored.attention_mask_header == payload.attention_mask_header
+    assert restored.position_ids_header == payload.position_ids_header
+    assert restored.past_len == payload.past_len