From ab466ce6b6cd9a85f27c1975b5a5e2b70863d306 Mon Sep 17 00:00:00 2001 From: Dobromir Popov Date: Fri, 17 Jul 2026 02:36:24 +0300 Subject: [PATCH] feat: add activation stream envelope --- .../evidence/DGR-021/README.md | 101 +++++ .../gitea-issues.json | 10 +- ...-versioned-named-tensor-stream-envelope.md | 14 +- .scratch/distributed-gguf-runtime/prd.json | 3 +- packages/node/meshnet_node/model_backend.py | 53 +++ packages/node/meshnet_node/protocol.py | 376 ++++++++++++++++++ tests/test_activation_envelope.py | 130 ++++++ 7 files changed, 674 insertions(+), 13 deletions(-) create mode 100644 .scratch/distributed-gguf-runtime/evidence/DGR-021/README.md create mode 100644 packages/node/meshnet_node/protocol.py create mode 100644 tests/test_activation_envelope.py diff --git a/.scratch/distributed-gguf-runtime/evidence/DGR-021/README.md b/.scratch/distributed-gguf-runtime/evidence/DGR-021/README.md new file mode 100644 index 0000000..a1e1bcc --- /dev/null +++ b/.scratch/distributed-gguf-runtime/evidence/DGR-021/README.md @@ -0,0 +1,101 @@ +# DGR-021 evidence — versioned named-tensor activation envelope + +**Completed:** 2026-07-17 +**Branch:** `distributed-gguf-runtime` +**Authority:** `.scratch/distributed-gguf-runtime/prd.json` +**Dependency:** DGR-018 (`evidence/DGR-018/README.md`) — canonical backlog schema / issue projection contract + +## Objective + +Establish the backend-neutral activation envelope used by direct and relayed Shard traffic, with stable versioning, named tensors, bounded fragmentation, checksum validation, and reserved extensibility for future state. + +## Changes + +### `packages/node/meshnet_node/protocol.py` (new) + +Added a self-contained activation-envelope module with: + +- `SCHEMA_NAME = "meshnet.activation-stream"` and `SCHEMA_VERSION = 1` +- `TensorFragment` + - bounded byte fragments with offset, compression tag, checksum, and extension preservation + - deterministic `to_dict()` / `from_dict()` round-trip +- `NamedTensor` + - named tensor metadata: `name`, `shape`, `dtype`, `byte_order`, `compression`, `checksum`, `fragments` + - fragmentation via `from_bytes(..., max_fragment_bytes=...)` + - checksum validation over reconstructed tensor bytes + - unknown-field preservation via `extensions` +- `ActivationEnvelope` + - top-level fields for `request_id`, `work_id`, `route_session`, `route_epoch`, `shard_start`, `effective_start`, `phase`, `position`, and `idempotency_step` + - reserved extension fields for `token_id_sideband`, `architecture_state`, `recurrent_state`, and `mtp` + - deterministic canonical serialization (`to_bytes`) and round-trip parsing (`from_bytes`) + - size-limit enforcement (`to_bytes(max_bytes=...)`) + - conversion from a live `TensorPayload` into the envelope and back again + +### `packages/node/meshnet_node/model_backend.py` + +Extended `TensorPayload` with envelope conversion helpers: + +- `TensorPayload.to_envelope(...)` +- `TensorPayload.from_envelope(...)` + +These keep the existing activation payload interface intact while exposing the new versioned envelope as the shared protocol layer. + +### `tests/test_activation_envelope.py` (new) + +Added focused deterministic tests covering: + +- deterministic envelope serialization and round-trip parsing +- tensor fragmentation and checksum validation +- unknown-field preservation at both envelope and tensor levels +- size-limit rejection +- `TensorPayload` ↔ envelope round-trip + +### `.scratch/distributed-gguf-runtime/prd.json` + +Marked `DGR-021.passes = true` and added completion notes recording the envelope implementation and verification commands. + +## Commands and results + +```bash +pytest -q tests/test_activation_envelope.py +``` + +```text +5 passed in 0.06s +``` + +```bash +pytest -q tests/test_activation_envelope.py tests/test_kv_cache_distributed.py -k 'session_is_stable_and_decode_payloads_are_single_token or large_prefill_activation_survives_zstd_compressed_hop' +``` + +```text +.. [100%] +2 passed, 21 deselected in 1.84s +``` + +```bash +python3 -m compileall packages/node/meshnet_node tests/test_activation_envelope.py +``` + +```text +Listing 'packages/node/meshnet_node'... +Listing 'packages/node/meshnet_node/native_protocol'... +Compiling 'tests/test_activation_envelope.py'... +``` + +```bash +git diff --check +``` + +```text +No whitespace errors +``` + +## Limitations + +- The envelope is implemented as a canonical deterministic JSON contract with dataclasses and conversion hooks, not generated `.proto` classes. The environment had `protobuf` available but not the `grpc_tools` generation toolchain, so I did not materialize a compiled proto artifact here. +- The direct/relayed HTTP/WebSocket transports remain byte-oriented; the envelope is the shared structured contract layered above those transports. + +## Dependency handoff + +DGR-022 and later shard-control stories can reuse the envelope contract and its `TensorPayload` conversion hooks as the stable activation metadata layer. Future work that requires generated protobuf code can replace the JSON serialization with a generated wire codec without changing the top-level field contract defined here. diff --git a/.scratch/distributed-gguf-runtime/gitea-issues.json b/.scratch/distributed-gguf-runtime/gitea-issues.json index 7b8e7be..26a7d9a 100644 --- a/.scratch/distributed-gguf-runtime/gitea-issues.json +++ b/.scratch/distributed-gguf-runtime/gitea-issues.json @@ -28,20 +28,20 @@ "DGR-021": { "number": 5, "url": "https://git.d-popov.com/popov/neuron-tai/issues/5", - "state": "open", - "status": "ready" + "state": "closed", + "status": "completed" }, "DGR-022": { "number": 6, "url": "https://git.d-popov.com/popov/neuron-tai/issues/6", "state": "open", - "status": "blocked" + "status": "ready" }, "DGR-023": { "number": 7, "url": "https://git.d-popov.com/popov/neuron-tai/issues/7", "state": "open", - "status": "blocked" + "status": "ready" }, "DGR-024": { "number": 8, @@ -53,7 +53,7 @@ "number": 9, "url": "https://git.d-popov.com/popov/neuron-tai/issues/9", "state": "open", - "status": "blocked" + "status": "ready" }, "DGR-026": { "number": 10, diff --git a/.scratch/distributed-gguf-runtime/issues/021-define-the-versioned-named-tensor-stream-envelope.md b/.scratch/distributed-gguf-runtime/issues/021-define-the-versioned-named-tensor-stream-envelope.md index 97dfc80..2db4fe5 100644 --- a/.scratch/distributed-gguf-runtime/issues/021-define-the-versioned-named-tensor-stream-envelope.md +++ b/.scratch/distributed-gguf-runtime/issues/021-define-the-versioned-named-tensor-stream-envelope.md @@ -1,7 +1,7 @@ # DGR-021: Define the versioned named-tensor stream envelope -- **Status / triage:** specification only; `ready-for-agent`; `passes: false` +- **Status / triage:** completed; `passes: true` - **Execution mode:** `AFK` - **Milestone:** `M1` - **Dependencies:** `DGR-018` @@ -18,11 +18,11 @@ Fresh Ralph session: read `.scratch/distributed-gguf-runtime/RALPH-CONTEXT.md`, ## Acceptance criteria -- [ ] Define schema version, request/work ID, route session/epoch, shard range/effective start, phase, position, and idempotency step. -- [ ] Define named tensors with shape, dtype, byte order, bounded fragments, compression identity, and checksum. -- [ ] Reserve extensible fields for token-ID sidebands, architecture state, recurrent state, and MTP without claiming implementations. -- [ ] Add deterministic serialization, fragmentation, checksum, unknown-field, and size-limit tests. -- [ ] Applicable shared quality gates in `prd.json` pass, and the evidence handoff records exact commands/results, changed files, limitations, and dependency handoff. +- [x] Define schema version, request/work ID, route session/epoch, shard range/effective start, phase, position, and idempotency step. +- [x] Define named tensors with shape, dtype, byte order, bounded fragments, compression identity, and checksum. +- [x] Reserve extensible fields for token-ID sidebands, architecture state, recurrent state, and MTP without claiming implementations. +- [x] Add deterministic serialization, fragmentation, checksum, unknown-field, and size-limit tests. +- [x] Applicable shared quality gates in `prd.json` pass, and the evidence handoff records exact commands/results, changed files, limitations, and dependency handoff. ## Shared quality gates @@ -36,4 +36,4 @@ Fresh Ralph session: read `.scratch/distributed-gguf-runtime/RALPH-CONTEXT.md`, ## Evidence handoff -Write and verify `.scratch/distributed-gguf-runtime/evidence/DGR-021/README.md`. Until every criterion and applicable gate has real evidence, this story remains `passes: false`. Legacy evidence is provenance only, not completion credit. +Verified evidence: `.scratch/distributed-gguf-runtime/evidence/DGR-021/README.md`. Legacy evidence remains provenance only and grants no implementation completion credit. diff --git a/.scratch/distributed-gguf-runtime/prd.json b/.scratch/distributed-gguf-runtime/prd.json index 8bbe42a..f685db8 100644 --- a/.scratch/distributed-gguf-runtime/prd.json +++ b/.scratch/distributed-gguf-runtime/prd.json @@ -441,7 +441,8 @@ "Add deterministic serialization, fragmentation, checksum, unknown-field, and size-limit tests.", "Applicable shared quality gates in `prd.json` pass, and the evidence handoff records exact commands/results, changed files, limitations, and dependency handoff." ], - "passes": false, + "passes": true, + "completionNotes": "Added a versioned activation envelope with deterministic JSON serialization, bounded tensor fragmentation, checksum validation, unknown-field preservation, and TensorPayload conversion hooks; verified by targeted pytest and compileall runs.", "notes": "Generated source issue: .scratch/distributed-gguf-runtime/issues/021-define-the-versioned-named-tensor-stream-envelope.md; prd.json is authoritative.", "blocks": [ "DGR-022", diff --git a/packages/node/meshnet_node/model_backend.py b/packages/node/meshnet_node/model_backend.py index 7f2ae1d..e78ee87 100644 --- a/packages/node/meshnet_node/model_backend.py +++ b/packages/node/meshnet_node/model_backend.py @@ -86,6 +86,59 @@ class TensorPayload: # Number of tokens already cached before this payload's tokens (decode steps). past_len: int | None = None + def to_envelope( + self, + *, + name: str, + request_id: str, + work_id: str, + route_session: str, + route_epoch: int, + shard_start: int, + effective_start: int, + phase: str, + position: int, + idempotency_step: int, + byte_order: str = "little", + compression: str = "identity", + max_fragment_bytes: int | None = None, + token_id_sideband: list[int] | None = None, + architecture_state: dict[str, Any] | None = None, + recurrent_state: dict[str, Any] | None = None, + mtp: dict[str, Any] | None = None, + extensions: dict[str, Any] | None = None, + ): + from .protocol import ActivationEnvelope, DEFAULT_FRAGMENT_BYTES + + return ActivationEnvelope.from_tensor_payload( + payload=self, + name=name, + request_id=request_id, + work_id=work_id, + route_session=route_session, + route_epoch=route_epoch, + shard_start=shard_start, + effective_start=effective_start, + phase=phase, + position=position, + idempotency_step=idempotency_step, + byte_order=byte_order, + compression=compression, + max_fragment_bytes=max_fragment_bytes or DEFAULT_FRAGMENT_BYTES, + token_id_sideband=token_id_sideband, + architecture_state=architecture_state, + recurrent_state=recurrent_state, + mtp=mtp, + extensions=extensions, + ) + + @classmethod + def from_envelope(cls, envelope, *, tensor_name: str = "activations"): + tensor_payload = envelope.to_tensor_payload(tensor_name=tensor_name) + if not isinstance(tensor_payload, cls): + raise TypeError("envelope did not produce a TensorPayload") + return tensor_payload + @dataclass(frozen=True) class TailTokenResult: diff --git a/packages/node/meshnet_node/protocol.py b/packages/node/meshnet_node/protocol.py new file mode 100644 index 0000000..3f52041 --- /dev/null +++ b/packages/node/meshnet_node/protocol.py @@ -0,0 +1,376 @@ +"""Versioned activation-stream envelope for shard hops. + +The transport still moves raw bytes over HTTP/WebSocket, but the payload now has +a stable, extensible envelope that names each tensor, preserves unknown fields, +and can round-trip deterministically across direct and relayed hops. +""" + +from __future__ import annotations + +from dataclasses import dataclass, field +import base64 +import hashlib +import json +from typing import Any + +SCHEMA_NAME = "meshnet.activation-stream" +SCHEMA_VERSION = 1 +DEFAULT_FRAGMENT_BYTES = 64 * 1024 + + +def _canonical_json(data: Any) -> bytes: + return json.dumps(data, sort_keys=True, separators=(",", ":"), ensure_ascii=False).encode("utf-8") + + +def _sha256_hex(data: bytes) -> str: + return hashlib.sha256(data).hexdigest() + + +def _normalize_shape(shape: list[int] | tuple[int, ...]) -> list[int]: + normalized = [int(dim) for dim in shape] + if not normalized or any(dim <= 0 for dim in normalized): + raise ValueError("shape must be a non-empty list of positive integers") + return normalized + + +def _fragment_bytes(body: bytes, max_fragment_bytes: int) -> tuple[bytes, ...]: + if max_fragment_bytes <= 0: + raise ValueError("max_fragment_bytes must be positive") + if not body: + return (b"",) + return tuple(body[offset : offset + max_fragment_bytes] for offset in range(0, len(body), max_fragment_bytes)) + + +@dataclass(frozen=True) +class TensorFragment: + """One bounded chunk of a named tensor.""" + + offset: int + body: bytes + checksum: str + compression: str = "identity" + extensions: dict[str, Any] = field(default_factory=dict) + + @classmethod + def from_bytes( + cls, + body: bytes, + *, + offset: int, + compression: str = "identity", + extensions: dict[str, Any] | None = None, + ) -> "TensorFragment": + return cls( + offset=int(offset), + body=bytes(body), + checksum=_sha256_hex(body), + compression=compression, + extensions=dict(extensions or {}), + ) + + def to_dict(self) -> dict[str, Any]: + data = { + "offset": self.offset, + "compression": self.compression, + "checksum": self.checksum, + "body_base64": base64.b64encode(self.body).decode("ascii"), + } + data.update(self.extensions) + return data + + @classmethod + def from_dict(cls, data: dict[str, Any]) -> "TensorFragment": + known = {"offset", "compression", "checksum", "body_base64"} + body = base64.b64decode(data.get("body_base64", "")) + fragment = cls( + offset=int(data["offset"]), + body=body, + checksum=str(data.get("checksum") or _sha256_hex(body)), + compression=str(data.get("compression") or "identity"), + extensions={k: v for k, v in data.items() if k not in known}, + ) + fragment.validate() + return fragment + + def validate(self) -> None: + if self.checksum != _sha256_hex(self.body): + raise ValueError("fragment checksum mismatch") + if self.offset < 0: + raise ValueError("fragment offset must be non-negative") + + +@dataclass(frozen=True) +class NamedTensor: + """A tensor named within a versioned activation envelope.""" + + name: str + shape: list[int] + dtype: str + byte_order: str + checksum: str + fragments: tuple[TensorFragment, ...] + compression: str = "identity" + extensions: dict[str, Any] = field(default_factory=dict) + + @classmethod + def from_bytes( + cls, + *, + name: str, + body: bytes, + shape: list[int] | tuple[int, ...], + dtype: str, + byte_order: str = "little", + compression: str = "identity", + max_fragment_bytes: int = DEFAULT_FRAGMENT_BYTES, + extensions: dict[str, Any] | None = None, + ) -> "NamedTensor": + normalized_shape = _normalize_shape(shape) + fragments = tuple( + TensorFragment.from_bytes(fragment, offset=offset, compression=compression) + for offset, fragment in enumerate(_fragment_bytes(body, max_fragment_bytes)) + for offset in (offset * max_fragment_bytes,) + ) + return cls( + name=str(name), + shape=normalized_shape, + dtype=str(dtype), + byte_order=str(byte_order), + checksum=_sha256_hex(body), + fragments=fragments, + compression=compression, + extensions=dict(extensions or {}), + ) + + def body(self) -> bytes: + ordered = sorted(self.fragments, key=lambda frag: frag.offset) + body = b"".join(fragment.body for fragment in ordered) + if _sha256_hex(body) != self.checksum: + raise ValueError(f"tensor {self.name!r} checksum mismatch") + return body + + def validate(self) -> None: + for fragment in self.fragments: + fragment.validate() + self.body() + + def to_dict(self) -> dict[str, Any]: + data = { + "name": self.name, + "shape": list(self.shape), + "dtype": self.dtype, + "byte_order": self.byte_order, + "compression": self.compression, + "checksum": self.checksum, + "fragments": [fragment.to_dict() for fragment in self.fragments], + } + data.update(self.extensions) + return data + + @classmethod + def from_dict(cls, data: dict[str, Any]) -> "NamedTensor": + known = { + "name", + "shape", + "dtype", + "byte_order", + "compression", + "checksum", + "fragments", + } + tensor = cls( + name=str(data["name"]), + shape=_normalize_shape(list(data["shape"])), + dtype=str(data["dtype"]), + byte_order=str(data.get("byte_order", "little")), + compression=str(data.get("compression", "identity")), + checksum=str(data["checksum"]), + fragments=tuple(TensorFragment.from_dict(fragment) for fragment in data.get("fragments", [])), + extensions={k: v for k, v in data.items() if k not in known}, + ) + tensor.validate() + return tensor + + +@dataclass(frozen=True) +class ActivationEnvelope: + """Versioned envelope for shard activation traffic.""" + + request_id: str + work_id: str + route_session: str + route_epoch: int + shard_start: int + effective_start: int + phase: str + position: int + idempotency_step: int + tensors: tuple[NamedTensor, ...] + version: int = SCHEMA_VERSION + schema: str = SCHEMA_NAME + token_id_sideband: list[int] | None = None + architecture_state: dict[str, Any] | None = None + recurrent_state: dict[str, Any] | None = None + mtp: dict[str, Any] | None = None + extensions: dict[str, Any] = field(default_factory=dict) + + @classmethod + def from_tensor_payload( + cls, + *, + payload: Any, + name: str, + request_id: str, + work_id: str, + route_session: str, + route_epoch: int, + shard_start: int, + effective_start: int, + phase: str, + position: int, + idempotency_step: int, + byte_order: str = "little", + compression: str = "identity", + max_fragment_bytes: int = DEFAULT_FRAGMENT_BYTES, + token_id_sideband: list[int] | None = None, + architecture_state: dict[str, Any] | None = None, + recurrent_state: dict[str, Any] | None = None, + mtp: dict[str, Any] | None = None, + extensions: dict[str, Any] | None = None, + ) -> "ActivationEnvelope": + tensor = NamedTensor.from_bytes( + name=name, + body=payload.body, + shape=payload.shape, + dtype="bfloat16", + byte_order=byte_order, + compression=compression, + max_fragment_bytes=max_fragment_bytes, + extensions={ + "attention_mask_header": payload.attention_mask_header, + "position_ids_header": payload.position_ids_header, + **({"past_len": payload.past_len} if payload.past_len is not None else {}), + }, + ) + return cls( + request_id=request_id, + work_id=work_id, + route_session=route_session, + route_epoch=int(route_epoch), + shard_start=int(shard_start), + effective_start=int(effective_start), + phase=str(phase), + position=int(position), + idempotency_step=int(idempotency_step), + tensors=(tensor,), + token_id_sideband=list(token_id_sideband) if token_id_sideband is not None else None, + architecture_state=architecture_state, + recurrent_state=recurrent_state, + mtp=mtp, + extensions=dict(extensions or {}), + ) + + def to_tensor_payload(self, *, tensor_name: str = "activations") -> Any: + from .model_backend import TensorPayload + + tensor = self.tensor(tensor_name) + return TensorPayload( + body=tensor.body(), + shape=list(tensor.shape), + attention_mask_header=tensor.extensions.get("attention_mask_header"), + position_ids_header=tensor.extensions.get("position_ids_header"), + past_len=tensor.extensions.get("past_len"), + ) + + def tensor(self, name: str = "activations") -> NamedTensor: + for tensor in self.tensors: + if tensor.name == name: + return tensor + raise KeyError(name) + + def to_dict(self) -> dict[str, Any]: + data = { + "schema": self.schema, + "version": self.version, + "request_id": self.request_id, + "work_id": self.work_id, + "route_session": self.route_session, + "route_epoch": self.route_epoch, + "shard_start": self.shard_start, + "effective_start": self.effective_start, + "phase": self.phase, + "position": self.position, + "idempotency_step": self.idempotency_step, + "tensors": [tensor.to_dict() for tensor in self.tensors], + } + if self.token_id_sideband is not None: + data["token_id_sideband"] = list(self.token_id_sideband) + if self.architecture_state is not None: + data["architecture_state"] = self.architecture_state + if self.recurrent_state is not None: + data["recurrent_state"] = self.recurrent_state + if self.mtp is not None: + data["mtp"] = self.mtp + data.update(self.extensions) + return data + + def to_bytes(self, *, max_bytes: int | None = None) -> bytes: + raw = _canonical_json(self.to_dict()) + if max_bytes is not None and len(raw) > max_bytes: + raise ValueError("activation envelope exceeds the size limit") + return raw + + @classmethod + def from_bytes(cls, data: bytes) -> "ActivationEnvelope": + payload = json.loads(data) + if not isinstance(payload, dict): + raise ValueError("activation envelope must be a JSON object") + known = { + "schema", + "version", + "request_id", + "work_id", + "route_session", + "route_epoch", + "shard_start", + "effective_start", + "phase", + "position", + "idempotency_step", + "tensors", + "token_id_sideband", + "architecture_state", + "recurrent_state", + "mtp", + } + envelope = cls( + schema=str(payload.get("schema", SCHEMA_NAME)), + version=int(payload.get("version", SCHEMA_VERSION)), + request_id=str(payload["request_id"]), + work_id=str(payload["work_id"]), + route_session=str(payload["route_session"]), + route_epoch=int(payload["route_epoch"]), + shard_start=int(payload["shard_start"]), + effective_start=int(payload["effective_start"]), + phase=str(payload["phase"]), + position=int(payload["position"]), + idempotency_step=int(payload["idempotency_step"]), + tensors=tuple(NamedTensor.from_dict(item) for item in payload.get("tensors", [])), + token_id_sideband=payload.get("token_id_sideband"), + architecture_state=payload.get("architecture_state"), + recurrent_state=payload.get("recurrent_state"), + mtp=payload.get("mtp"), + extensions={k: v for k, v in payload.items() if k not in known}, + ) + envelope.validate() + return envelope + + def validate(self) -> None: + if self.version != SCHEMA_VERSION: + raise ValueError("unsupported activation envelope version") + if self.schema != SCHEMA_NAME: + raise ValueError("unsupported activation envelope schema") + if self.phase not in {"prefill", "decode"}: + raise ValueError("phase must be prefill or decode") + for tensor in self.tensors: + tensor.validate() diff --git a/tests/test_activation_envelope.py b/tests/test_activation_envelope.py new file mode 100644 index 0000000..308a0aa --- /dev/null +++ b/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