// The native Shard data plane: Protocol Buffers over gRPC/HTTP2 (ADR-0020).
//
// This schema — not any Python or C++ type — is the semantic contract between a
// Meshnet node and a Shard worker. Direct hops speak gRPC. When direct
// connectivity is unavailable the existing relay carries these exact serialized
// frames as opaque binary, which is why every deadline, cancellation and
// identity field is carried *in the schema* rather than left to gRPC metadata:
// a relayed frame has no HTTP/2 context to inherit them from.
//
// Meshnet remains the only control plane. Nothing here selects routes, prices
// work, or authenticates peers; the worker executes a layer range it has been
// admitted for and reports what it did.

syntax = "proto3";

package meshnet.shard.v1;

option cc_enable_arenas = true;
option go_package = "github.com/meshnet/shard/v1;shardv1";

// ---------------------------------------------------------------------------
// Versioning
// ---------------------------------------------------------------------------

// Wire-compatibility generation of this schema. A worker rejects a peer whose
// SCHEMA_VERSION it cannot satisfy instead of guessing at field meaning.
//
// This is deliberately NOT the `X-Meshnet-Wire` version of the legacy HTTP
// activation format (currently "2", see ADR-0008). The native protocol is a
// separate contract with its own generation counter and starts at 1.
enum SchemaVersion {
  SCHEMA_VERSION_UNSPECIFIED = 0;
  SCHEMA_VERSION_1 = 1;
}

// ---------------------------------------------------------------------------
// Tensor bundle — the public activation boundary
// ---------------------------------------------------------------------------

// Element type of a tensor payload.
//
// bfloat16 is the canonical activation dtype at every Shard boundary regardless
// of the weight quantization a node uses internally (ADR-0008): weights may be
// NF4 or INT8, compute upcasts, and the boundary stays bfloat16. The other
// values exist for auxiliary tensors (position ids, attention masks) and for
// architectures whose boundary is not a plain hidden state.
enum DType {
  DTYPE_UNSPECIFIED = 0;
  DTYPE_BFLOAT16 = 1;
  DTYPE_FLOAT16 = 2;
  DTYPE_FLOAT32 = 3;
  DTYPE_INT32 = 4;
  DTYPE_INT64 = 5;
  DTYPE_UINT8 = 6;
  DTYPE_INT8 = 7;
  DTYPE_BOOL = 8;
}

// Byte order of a tensor payload. Little-endian is canonical on the wire; the
// field is explicit so a big-endian peer is rejected loudly rather than
// silently reading byte-swapped activations.
enum ByteOrder {
  BYTE_ORDER_UNSPECIFIED = 0;
  BYTE_ORDER_LITTLE_ENDIAN = 1;
  BYTE_ORDER_BIG_ENDIAN = 2;
}

// Payload compression. Mirrors the policy-driven zstd decision the existing
// activation seam already makes (`activation_compression.py`): compression is a
// transport optimisation and NONE is always a legal choice for any payload.
enum Compression {
  COMPRESSION_UNSPECIFIED = 0;
  COMPRESSION_NONE = 1;
  COMPRESSION_ZSTD = 2;
}

// Integrity algorithm covering a payload.
enum ChecksumAlgorithm {
  CHECKSUM_ALGORITHM_UNSPECIFIED = 0;
  CHECKSUM_ALGORITHM_NONE = 1;
  CHECKSUM_ALGORITHM_CRC32C = 2;
}

// An integrity check over the *uncompressed* canonical payload bytes.
//
// Checksumming the uncompressed bytes (not the compressed frame) means the
// value is stable whether a hop compressed the payload or not, so a relay may
// re-frame or a peer may decline compression without invalidating it.
message Checksum {
  ChecksumAlgorithm algorithm = 1;
  // Big-endian encoding of the checksum value.
  bytes value = 2;
}

// One slice of a tensor's wire payload.
//
// Fragments bound the size of a single gRPC message. `byte_offset` lets a
// receiver verify the fragments tile the body exactly — no hole, no overlap —
// instead of trusting arrival order.
//
// Offsets and payloads describe the *wire* body, which is the compressed frame
// when the parent tensor declares a compression. A zstd frame is not decodable
// per fragment, so a receiver reassembles first and decompresses once. The
// uncompressed size is not repeated here: `NamedTensor.total_bytes` is the
// single source of truth, and a second copy could only ever disagree with it.
message TensorFragment {
  uint32 fragment_index = 1;
  uint32 fragment_count = 2;
  // Offset of this fragment within the tensor's wire body.
  uint64 byte_offset = 3;
  // Fragment bytes, compressed iff the parent tensor declares a compression.
  bytes payload = 4;

  reserved 5;
  reserved "uncompressed_size";
}

// One named tensor at an architecture boundary.
message NamedTensor {
  // Architecture-defined boundary name, e.g. "hidden_states", "position_ids".
  // A boundary may need more than one tensor, which is why the payload is a
  // named bundle rather than a bare buffer (ADR-0020).
  string name = 1;
  repeated int64 shape = 2;
  DType dtype = 3;
  ByteOrder byte_order = 4;
  // Total uncompressed payload size. A receiver sizes its buffer from this
  // before reading fragments and refuses a bundle that exceeds its budget.
  uint64 total_bytes = 5;
  Compression compression = 6;
  // Over the uncompressed payload; see Checksum.
  Checksum checksum = 7;
  // Ordered fragments. A tensor small enough for one message carries exactly
  // one fragment with fragment_count == 1.
  repeated TensorFragment fragments = 8;
}

// A versioned named-tensor bundle: the public activation boundary payload.
message TensorBundle {
  // Generation of the bundle layout itself, independent of SchemaVersion so a
  // boundary payload can evolve without a whole-protocol version bump.
  uint32 bundle_version = 1;
  repeated NamedTensor tensors = 2;
}

// ---------------------------------------------------------------------------
// Identity: what work this is, for which route, against which artifact
// ---------------------------------------------------------------------------

// Exact identity of the executable thing a Shard runs.
//
// Both fingerprints must match end to end across a route. Two nodes agreeing on
// a model name but not on quantization, kernel or tail-norm placement would
// produce silently wrong tokens, so identity is a digest, not a label.
message Fingerprint {
  // Digest of the exact Model Artifact (weights + config), e.g. the GGUF hash.
  string model_artifact_digest = 1;
  // Digest of the exact runtime recipe (backend, quantization, kernels, dtype).
  string runtime_recipe_digest = 2;
  // Human-readable recipe identity, mirroring the capability report an admitted
  // node already registers with (ADR-0023). Labels are for diagnosis; the
  // digests above are what a peer actually compares.
  string recipe_id = 3;
  string recipe_version = 4;
  string catalogue_version = 5;
}

// The contiguous transformer layer range a Shard owns.
message ShardRange {
  // Registered range, inclusive start, exclusive end.
  uint32 start_layer = 1;
  uint32 end_layer = 2;
  // The layer this Shard must actually begin at for this route (ADR-0012).
  //
  // Shard ranges may overlap; the Tracker resolves the overlap when it builds
  // the route and every hop is told where the previous hop stopped. Executing
  // from `start_layer` when `effective_start_layer` is higher would re-apply
  // layers already applied to the incoming activation and silently corrupt the
  // result. A worker executes [effective_start_layer, end_layer).
  uint32 effective_start_layer = 3;
}

// Which part of generation a message belongs to.
enum Phase {
  PHASE_UNSPECIFIED = 0;
  PHASE_PREFILL = 1;
  PHASE_DECODE = 2;
  PHASE_RELEASE = 3;
  PHASE_CANCEL = 4;
}

// Token span carried by one chunk.
message PositionSpan {
  // Absolute position of the first token in this chunk within the sequence.
  uint64 first_position = 1;
  // Number of token positions in this chunk. A decode step carries 1.
  uint32 token_count = 2;
}

// Bounded chunking for a prefill.
//
// A long prompt is split into token-aligned chunks before the first forward
// pass (ADR-0008) so peak transfer per boundary stays bounded regardless of
// prompt length. Splits never fall mid-token.
message ChunkInfo {
  uint32 chunk_index = 1;
  uint32 chunk_count = 2;
  // True on the final chunk of a prefill. A receiver that has not seen a final
  // chunk knows the prefill is still incomplete.
  bool final_chunk = 3;
}

// How the sender expects the receiver's Hot KV State to be positioned.
enum CacheMode {
  CACHE_MODE_UNSPECIFIED = 0;
  // No session state; the payload carries everything needed. Always a legal
  // fallback and the recovery path after a miss.
  CACHE_MODE_STATELESS = 1;
  // Establish fresh session state for this Shard's layer range.
  CACHE_MODE_PREFILL = 2;
  // Continue existing session state. `expected_past_len` must match exactly.
  CACHE_MODE_DECODE = 3;
}

// The sender's expectation about receiver-side cache state (ADR-0022).
//
// A mismatch is always a declared cache miss, never a silent stateless forward:
// running a single-token decode payload without the matching cache would emit
// plausible garbage. The receiver answers a miss with ShardError.CACHE_MISS and
// the head re-prefills the whole sequence under the same Route Session ID.
message CacheExpectation {
  CacheMode mode = 1;
  // Decode only: the number of tokens the receiver's session cache must already
  // hold for this Shard's layer range.
  uint64 expected_past_len = 2;
}

// What the receiver's cache actually did.
message CacheResult {
  CacheMode mode = 1;
  // Tokens held for this Shard's range after the step completed.
  uint64 past_len = 2;
  // True when session state was reused rather than rebuilt.
  bool cache_hit = 3;
}

// Everything required to interpret one unit of work, independent of transport.
//
// Carried on every request and response because a relayed frame arrives as
// opaque binary with no gRPC metadata, no deadline and no channel identity.
message Envelope {
  SchemaVersion schema_version = 1;
  // Unique id for this unit of work; also the unit of billing attribution.
  string work_id = 2;
  // The Route Session this work belongs to. Together with `route_epoch` it maps
  // to exactly one isolated llama sequence / bounded context (ADR-0020).
  string route_session_id = 3;
  // Bumped by the control plane whenever the route changes. A worker refuses
  // work from a stale epoch rather than mixing it into live session state.
  uint64 route_epoch = 4;
  Fingerprint fingerprint = 5;
  ShardRange shard_range = 6;
  Phase phase = 7;
  PositionSpan position = 8;
  // Monotonic step within (route_session_id, route_epoch).
  //
  // Idempotency key: a retried or duplicated step carries the same value and
  // must be acknowledged, not re-applied. Re-applying a decode step would
  // advance the KV cache twice and desynchronise the route.
  uint64 idempotency_step = 9;
  CacheExpectation cache_expectation = 10;
  // Absolute deadline. Carried in-band so a relayed frame keeps its deadline;
  // on a direct hop it is set consistently with the gRPC deadline.
  int64 deadline_unix_nanos = 11;
  // Chunking, when this message is part of a bounded prefill.
  ChunkInfo chunk = 12;
}

// ---------------------------------------------------------------------------
// Structured errors
// ---------------------------------------------------------------------------

// Why a unit of work could not be executed as asked.
//
// These are semantic outcomes, distinct from transport failure. They travel as
// a ShardStatus on the stream (and as google.rpc.Status details on unary RPCs)
// so a relayed frame carries the same diagnosis a direct gRPC hop would.
enum ErrorCode {
  ERROR_CODE_UNSPECIFIED = 0;
  // Peer cannot satisfy the requested SchemaVersion.
  ERROR_CODE_SCHEMA_UNSUPPORTED = 1;
  // Model artifact or runtime recipe digest differs from this worker's.
  ERROR_CODE_FINGERPRINT_MISMATCH = 2;
  // Work arrived for a route epoch the worker has already moved past.
  ERROR_CODE_EPOCH_STALE = 3;
  // Requested layer range is not the range this worker serves.
  ERROR_CODE_SHARD_RANGE_MISMATCH = 4;
  // Session state absent or positioned differently than expected (ADR-0022).
  // The head recovers with one full re-prefill; this is not a fatal error.
  ERROR_CODE_CACHE_MISS = 5;
  // Admission budget (weights, KV, scratch, queue depth) would be exceeded.
  ERROR_CODE_RESOURCE_EXHAUSTED = 6;
  // Payload failed its checksum, or fragments did not cover the tensor.
  ERROR_CODE_PAYLOAD_CORRUPT = 7;
  // Work was cancelled by the control plane or the peer.
  ERROR_CODE_CANCELLED = 8;
  // Deadline in the envelope had already passed when work was dequeued.
  ERROR_CODE_DEADLINE_EXCEEDED = 9;
  // Sender exceeded its granted flow-control credit.
  ERROR_CODE_FLOW_CONTROL_VIOLATION = 10;
  // The worker failed while executing; detail is sanitized.
  ERROR_CODE_INTERNAL = 11;
}

message ShardError {
  ErrorCode code = 1;
  // Sanitized, operator-facing explanation. Never a raw exception, file path or
  // credential (ADR-0023 sanitization rule).
  string detail = 2;
  // True when the same work may be retried unchanged. A CACHE_MISS is not
  // retryable unchanged — the head must re-prefill first.
  bool retryable = 3;
  // Set on CACHE_MISS so the head knows where the receiver actually is.
  uint64 actual_past_len = 4;
}

// ---------------------------------------------------------------------------
// Flow control
// ---------------------------------------------------------------------------

// Application-level credit, layered on top of HTTP/2 flow control.
//
// HTTP/2 bounds bytes in flight; it does not bound how much *work* a worker has
// queued, and a relayed frame gets no HTTP/2 window at all. Credits bound the
// number of un-acked chunks a sender may have outstanding, which is what keeps
// worker queues and KV pressure bounded and lets prefill avoid starving decode.
message FlowControl {
  // Additional chunks the receiver is willing to accept beyond those already
  // granted. Purely additive: a grant never reduces outstanding credit.
  uint32 credits_granted = 1;
  // Hard ceiling on un-acked chunks, independent of granted credit.
  uint32 max_inflight_chunks = 2;
  // Hard ceiling on the serialized size of one chunk message.
  uint64 max_chunk_bytes = 3;
  // Hard ceiling on token positions per prefill chunk.
  uint32 max_prefill_chunk_tokens = 4;
}

// ---------------------------------------------------------------------------
// Session stream messages
// ---------------------------------------------------------------------------

// Opens one long-lived bidirectional stream for one Route Session Activation
// Seam. The handshake settles version, identity and the initial flow-control
// window before any activation is sent, so an incompatible peer fails at open
// rather than mid-generation.
message SessionOpen {
  // Highest schema version the initiator supports; the acceptor replies with
  // the version actually negotiated.
  SchemaVersion schema_version = 1;
  string route_session_id = 2;
  uint64 route_epoch = 3;
  Fingerprint fingerprint = 4;
  ShardRange shard_range = 5;
  // Flow-control limits the initiator can honour. The acceptor replies with the
  // limits that actually apply.
  FlowControl proposed_flow_control = 6;
  // Compressions the initiator can decode. NONE is implicitly always supported.
  repeated Compression accepted_compression = 7;
}

message SessionAccepted {
  // The version both peers will use for the life of this stream.
  SchemaVersion schema_version = 1;
  string route_session_id = 2;
  uint64 route_epoch = 3;
  // Authoritative limits. The initiator must not exceed these.
  FlowControl flow_control = 4;
  repeated Compression accepted_compression = 5;
  // Fingerprint the worker actually serves, so a mismatch is visible at open.
  Fingerprint fingerprint = 6;
}

// One unit of activation work: a bounded prefill chunk or a decode step.
message ActivationChunk {
  Envelope envelope = 1;
  TensorBundle bundle = 2;
}

// The small decode fast path.
//
// A decode step is one token: the envelope's identity fields are already fixed
// for the life of the stream, so repeating them per token is pure overhead on
// the hottest path. A DecodeStep carries only what changes — the step, the
// position, and one tensor — and inherits the rest from the SessionOpen
// handshake. A peer may always fall back to ActivationChunk with PHASE_DECODE.
message DecodeStep {
  // Idempotency step within the session; also orders the stream.
  uint64 idempotency_step = 1;
  // Absolute position of this token.
  uint64 position = 2;
  // Tokens the receiver's cache must already hold. A mismatch is a CACHE_MISS.
  uint64 expected_past_len = 3;
  // The single boundary tensor for this token, typically [1, 1, hidden].
  NamedTensor tensor = 4;
  // Work id for attribution; the session/epoch/fingerprint/range are inherited.
  string work_id = 5;
  int64 deadline_unix_nanos = 6;
}

// Drop session state for a Route Session. Bounded memory does not depend on
// this arriving — workers also evict by TTL and LRU (ADR-0022) — but an
// explicit release returns KV immediately instead of holding it for the TTL.
message ReleaseSignal {
  string route_session_id = 1;
  uint64 route_epoch = 2;
  string work_id = 3;
}

// Abandon in-flight work.
//
// Cancellation must remain possible when the stream is wedged behind flow
// control, which is why Cancel also exists as a unary RPC on a fresh call.
message CancelSignal {
  string route_session_id = 1;
  uint64 route_epoch = 2;
  // Cancel one unit of work; empty cancels every in-flight unit for the session.
  string work_id = 3;
  string reason = 4;
}

// Acknowledges a chunk or decode step, releasing its flow-control credit and
// recording where the receiver's cache now stands.
message Ack {
  string work_id = 1;
  uint64 idempotency_step = 2;
  CacheResult cache_result = 3;
  // True when the step was recognised as an already-applied duplicate and was
  // therefore acknowledged without being re-applied.
  bool duplicate = 4;
  // Observed execution time, for Generation Telemetry.
  uint64 execution_nanos = 5;
}

// A terminal outcome for one unit of work, or for the session.
message ShardStatus {
  string work_id = 1;
  string route_session_id = 2;
  uint64 idempotency_step = 3;
  ShardError error = 4;
  // True when the stream itself is finished and no further work will be served.
  bool terminal = 5;
}

message SessionRequest {
  oneof kind {
    SessionOpen open = 1;
    ActivationChunk chunk = 2;
    DecodeStep decode = 3;
    FlowControl flow_control = 4;
    ReleaseSignal release = 5;
    CancelSignal cancel = 6;
  }
}

message SessionResponse {
  oneof kind {
    SessionAccepted accepted = 1;
    // Result activations forwarded toward the next Shard, or to the head.
    ActivationChunk chunk = 2;
    Ack ack = 3;
    FlowControl flow_control = 4;
    ShardStatus status = 5;
  }
}

// ---------------------------------------------------------------------------
// Capability and health
// ---------------------------------------------------------------------------

message CapabilityRequest {
  SchemaVersion schema_version = 1;
}

// What this worker can actually execute, proven by a bounded real forward
// (ADR-0023). Detected hardware is not a capability; only a validated recipe is.
message CapabilityReport {
  SchemaVersion schema_version = 1;
  Fingerprint fingerprint = 2;
  ShardRange shard_range = 3;
  // Backend/device identity, e.g. "rocm:gfx1151", "cpu:avx512".
  string backend = 4;
  string device = 5;
  // True only when a bounded real forward has passed for exactly this
  // (artifact, range, recipe, device). Never set from hardware detection alone.
  bool validated = 6;
  // Sanitized reason when `validated` is false.
  string detail = 7;
  // Admission budgets this worker will enforce.
  uint64 max_concurrent_sessions = 8;
  uint64 max_context_tokens = 9;
  FlowControl flow_control = 10;
  repeated Compression accepted_compression = 11;
  repeated SchemaVersion supported_schema_versions = 12;
  int64 validated_at_unix_nanos = 13;
}

message HealthRequest {
  SchemaVersion schema_version = 1;
}

enum ServingState {
  SERVING_STATE_UNSPECIFIED = 0;
  // Accepting new Route Sessions.
  SERVING_STATE_SERVING = 1;
  // Alive but refusing new sessions; existing sessions continue draining.
  SERVING_STATE_DRAINING = 2;
  // Alive but not routable — e.g. registered-but-dark pending certification.
  SERVING_STATE_NOT_SERVING = 3;
}

// Live load, for backpressure and Generation Telemetry.
message HealthReport {
  SchemaVersion schema_version = 1;
  ServingState state = 2;
  uint32 active_sessions = 3;
  uint32 queued_chunks = 4;
  // Occupancy of the continuous decode batch.
  uint32 batch_occupancy = 5;
  // Fraction of the KV budget in use, 0..1.
  float kv_pressure = 6;
  uint64 resident_bytes = 7;
  string detail = 8;
}

message ReleaseRequest {
  SchemaVersion schema_version = 1;
  string route_session_id = 2;
  uint64 route_epoch = 3;
}

message ReleaseResponse {
  // True when session state existed and was dropped; false when there was
  // nothing to drop, which is a success, not an error — release is idempotent.
  bool released = 1;
  ShardError error = 2;
}

message CancelRequest {
  SchemaVersion schema_version = 1;
  string route_session_id = 2;
  uint64 route_epoch = 3;
  // Empty cancels every in-flight unit for the session.
  string work_id = 4;
  string reason = 5;
}

message CancelResponse {
  uint32 cancelled_work_items = 1;
  ShardError error = 2;
}

// ---------------------------------------------------------------------------
// Service
// ---------------------------------------------------------------------------

service ShardRuntime {
  // What this worker can execute. Read before a route is built.
  rpc GetCapability(CapabilityRequest) returns (CapabilityReport);

  // Live load and serving state.
  rpc Health(HealthRequest) returns (HealthReport);

  // One long-lived bidirectional stream per Route Session Activation Seam.
  //
  // The stream opens with SessionOpen/SessionAccepted, then carries bounded
  // prefill chunks and decode steps in both directions for the life of the
  // session. Per-token channel creation is a non-goal: the handshake cost is
  // paid once and the hot path carries only what changes.
  rpc Session(stream SessionRequest) returns (stream SessionResponse);

  // Drop session state out of band. Idempotent.
  rpc Release(ReleaseRequest) returns (ReleaseResponse);

  // Cancel out of band, on a fresh call.
  //
  // In-band CancelSignal is preferred, but a sender that is blocked on flow
  // control cannot write one — a cancel that can only travel down a wedged
  // stream is not a cancel. This RPC always has a path to the worker.
  rpc Cancel(CancelRequest) returns (CancelResponse);
}