// Shard runtime data-plane protocol for the distributed GGUF runtime (ADR-0024).
//
// This schema is the semantic contract between Python and C++ Shards. Direct
// transport is gRPC over HTTP/2; the existing Meshnet relay may carry the same
// serialized frames as opaque binary, so anything gRPC would normally carry in
// call metadata (deadlines, cancellation intent) is ALSO representable inside
// the messages for relay-transported seams.
//
// Design rules (see .scratch/distributed-gguf-runtime/RALPH-CONTEXT.md):
//   * One long-lived bidirectional ActivateSession stream per Route Session
//     Activation Seam. No per-token channel creation.
//   * Bounded chunking for prefill; a small decode fast path.
//   * The activation boundary is a versioned named-tensor bundle, because an
//     architecture boundary may require more than one tensor.
//   * Meshnet routing/billing/auth live outside this schema; only the data
//     plane and the identifiers needed to attribute and isolate work are here.
//
// Compatibility: proto3. Never renumber or reuse a field number. Add new fields
// with new numbers only. Enums keep a 0 UNSPECIFIED member for forward compat.

syntax = "proto3";

package meshnet.shard.v1;

option java_package = "com.meshnet.shard.v1";
option java_outer_classname = "ShardRuntimeProto";
option go_package = "meshnet/shard/v1;shardv1";

// ---------------------------------------------------------------------------
// Versioning and enums
// ---------------------------------------------------------------------------

// Wire schema version. Bumped only on incompatible envelope changes; additive
// field changes keep the same version and rely on proto3 unknown-field rules.
enum SchemaVersion {
  SCHEMA_VERSION_UNSPECIFIED = 0;
  SCHEMA_VERSION_1 = 1;
}

// Lifecycle phase of a seam message. RELEASE and CANCEL are represented both as
// dedicated RPCs and as in-stream phases so a relay-carried stream can express
// them without a separate channel.
enum Phase {
  PHASE_UNSPECIFIED = 0;
  PHASE_PREFILL = 1;
  PHASE_DECODE = 2;
  PHASE_RELEASE = 3;
  PHASE_CANCEL = 4;
}

// Tensor element type. GGUF quantized block types are enumerated explicitly so
// a boundary bundle can carry pre-quantized payloads without reinterpretation.
enum DType {
  DTYPE_UNSPECIFIED = 0;
  DTYPE_F32 = 1;
  DTYPE_F16 = 2;
  DTYPE_BF16 = 3;
  DTYPE_I64 = 4;
  DTYPE_I32 = 5;
  DTYPE_I16 = 6;
  DTYPE_I8 = 7;
  DTYPE_U8 = 8;
  DTYPE_BOOL = 9;
  DTYPE_Q8_0 = 20;
  DTYPE_Q4_0 = 21;
  DTYPE_Q4_K = 22;
  DTYPE_Q6_K = 23;
}

// Byte order of a tensor payload. Explicit because Shards may run on
// heterogeneous hardware and the relay carries opaque bytes.
enum ByteOrder {
  BYTE_ORDER_UNSPECIFIED = 0;
  BYTE_ORDER_LITTLE_ENDIAN = 1;
  BYTE_ORDER_BIG_ENDIAN = 2;
}

// Payload compression applied to a tensor fragment or message body.
enum Compression {
  COMPRESSION_UNSPECIFIED = 0;
  COMPRESSION_NONE = 1;
  COMPRESSION_ZSTD = 2;
}

// Checksum algorithm. CRC32C is the cheap per-fragment default; SHA256 is used
// where stronger integrity is required.
enum ChecksumAlgorithm {
  CHECKSUM_ALGORITHM_UNSPECIFIED = 0;
  CHECKSUM_NONE = 1;
  CHECKSUM_CRC32C = 2;
  CHECKSUM_CRC32 = 3;
  CHECKSUM_SHA256 = 4;
}

// What the sender expects from the receiving Shard's Hot KV State for this work
// (request side of the cache contract).
enum CacheExpectation {
  CACHE_EXPECTATION_UNSPECIFIED = 0;
  CACHE_REUSE = 1;   // reuse existing KV for (session, epoch)
  CACHE_FRESH = 2;   // start a fresh KV context
  CACHE_BYPASS = 3;  // stateless; do not persist KV
}

// What the receiving Shard actually did with its KV State (result side).
enum CacheResult {
  CACHE_RESULT_UNSPECIFIED = 0;
  CACHE_HIT = 1;
  CACHE_MISS = 2;
  CACHE_WRITTEN = 3;
  CACHE_BYPASSED = 4;
}

// Coarse retry classification carried in structured status.
enum RetryClass {
  RETRY_CLASS_UNSPECIFIED = 0;
  RETRY_CLASS_NONE = 1;       // terminal success/no-retry
  RETRY_CLASS_RETRYABLE = 2;  // transient; the same step may be retried
  RETRY_CLASS_FATAL = 3;      // do not retry this route/epoch
  RETRY_CLASS_EPOCH_STALE = 4; // route epoch advanced; re-resolve route
}

enum ServingStatus {
  SERVING_STATUS_UNSPECIFIED = 0;
  SERVING = 1;
  NOT_SERVING = 2;
  DRAINING = 3;
}

// ---------------------------------------------------------------------------
// Common value messages
// ---------------------------------------------------------------------------

// Structured, transport-independent status. Mirrors canonical gRPC codes so a
// relay-carried frame can express what a gRPC trailer normally would.
message Status {
  uint32 code = 1;              // canonical gRPC status code
  string message = 2;
  RetryClass retry_class = 3;
  map<string, string> details = 4;
}

// Integrity check over an associated payload.
message Checksum {
  ChecksumAlgorithm algorithm = 1;
  bytes value = 2;
}

// Exact Model Artifact / runtime-recipe fingerprint. Both Shards MUST agree on
// every populated field before activation; a mismatch is a fatal status.
message ArtifactFingerprint {
  string model_id = 1;                    // e.g. "meta-llama/Llama-3.1-8B"
  string revision = 2;                    // artifact revision / commit
  string artifact_hash = 3;               // hash of the GGUF/model artifact
  string quantization = 4;                // e.g. "Q4_K_M", "F16"
  string runtime_recipe_fingerprint = 5;  // DGR-003 recipe hash
}

// Contiguous transformer layer range owned by a Shard (ADR-0012). end_layer is
// exclusive. effective_start_layer is the overlap-safe start after de-dupe of
// shared boundary layers between adjacent Shards.
message ShardRange {
  uint32 start_layer = 1;
  uint32 end_layer = 2;
  uint32 effective_start_layer = 3;
  bool owns_embedding = 4;
  bool owns_final_head = 5;
}

// Token position window for a message. start_position is the absolute index of
// the first token; token_count is how many positions this message covers.
message Position {
  uint64 start_position = 1;
  uint64 token_count = 2;
  uint64 sequence_length = 3;  // total known context length, if known
}

// Envelope carried by every seam message. Everything required to version,
// route-attribute, isolate, order, and integrity-check a unit of work.
message MessageHeader {
  SchemaVersion schema_version = 1;
  string work_id = 2;                 // request/work ID (idempotency scope)
  string route_session_id = 3;        // Route Session ID
  uint64 route_epoch = 4;             // route epoch; stale epochs are rejected
  ArtifactFingerprint fingerprint = 5;
  ShardRange shard_range = 6;
  Phase phase = 7;
  Position position = 8;
  uint64 idempotency_step = 9;        // monotonic per (work_id) step counter
  CacheExpectation cache_expectation = 10;
  Compression compression = 11;       // compression of THIS message's payloads
  Checksum checksum = 12;             // checksum over THIS message's payload
}

// ---------------------------------------------------------------------------
// Versioned named-tensor bundle (the activation boundary payload)
// ---------------------------------------------------------------------------

// One bounded fragment of a tensor payload. Large tensors are split so no
// single message is unbounded; fragments reassemble by byte_offset order.
message TensorFragment {
  uint32 fragment_index = 1;
  uint32 fragment_count = 2;
  uint64 byte_offset = 3;   // offset of this fragment within the full payload
  bytes data = 4;
  Checksum checksum = 5;    // checksum over this fragment's (post-compression) data
}

// A single named tensor with full description so the receiver never reinterprets
// bytes implicitly.
message NamedTensor {
  string name = 1;
  repeated uint64 shape = 2;
  DType dtype = 3;
  ByteOrder byte_order = 4;
  uint64 total_byte_length = 5;   // full payload length across all fragments
  Compression compression = 6;    // compression applied to fragment data
  repeated TensorFragment fragments = 7;
}

// A versioned collection of named tensors representing one activation boundary.
message TensorBundle {
  uint32 bundle_version = 1;
  repeated NamedTensor tensors = 2;
}

// ---------------------------------------------------------------------------
// Session stream messages (bidirectional ActivateSession)
// ---------------------------------------------------------------------------

// Opens a seam. Carries the header plus stream-scoped bounds. deadline_unix_nanos
// lets a relay-carried stream express the call deadline gRPC would otherwise own.
message SessionOpen {
  MessageHeader header = 1;
  uint64 deadline_unix_nanos = 2;         // absolute deadline; 0 = none
  uint32 max_prefill_tokens_per_chunk = 3; // bound for prefill chunking
  uint32 max_fragment_bytes = 4;           // bound for tensor fragment size
  FlowControl initial_credit = 5;          // receiver's starting flow-control window
}

// Bounded prefill chunk. A prefill is split into ordered chunks each covering at
// most max_prefill_tokens_per_chunk positions; final_chunk marks the last one.
message PrefillChunk {
  MessageHeader header = 1;
  uint32 chunk_index = 2;
  uint32 chunk_count = 3;   // 0 if unknown/streaming
  bool final_chunk = 4;
  TensorBundle activations = 5;
}

// Small decode fast path: a single-position (or tiny) step with minimal framing.
// Reuses the same header for isolation/ordering but expects one activation bundle.
message DecodeStep {
  MessageHeader header = 1;
  TensorBundle activation = 2;
}

// Explicit HTTP/2-independent flow-control grant. credits is the number of
// additional messages the receiver is willing to accept; the byte/message caps
// bound in-flight work for backpressure.
message FlowControl {
  uint64 credits = 1;
  uint64 max_in_flight_bytes = 2;
  uint64 max_in_flight_messages = 3;
}

// Release a session's resources (Hot KV State, sequence) cleanly.
message ReleaseRequest {
  MessageHeader header = 1;
  string reason = 2;
}

message ReleaseResponse {
  Status status = 1;
  CacheResult cache_result = 2;
}

// Cancel in-flight work for a session/step.
message CancelRequest {
  MessageHeader header = 1;
  string reason = 2;
}

message CancelResponse {
  Status status = 1;
}

// Client -> server frames on the ActivateSession stream.
message SessionActivation {
  oneof payload {
    SessionOpen open = 1;
    PrefillChunk prefill = 2;
    DecodeStep decode = 3;
    ReleaseRequest release = 4;
    CancelRequest cancel = 5;
    FlowControl flow_control = 6;
  }
}

// Computed boundary output for a step: the next Shard's input tensors plus the
// cache result and integrity for what was produced.
message ActivationResult {
  MessageHeader header = 1;
  TensorBundle outputs = 2;
  CacheResult cache_result = 3;
  Status status = 4;
}

message SessionAccepted {
  MessageHeader header = 1;
  FlowControl granted_credit = 2;
  Status status = 3;
}

// Server -> client frames on the ActivateSession stream.
message SessionResponse {
  oneof payload {
    SessionAccepted accepted = 1;
    ActivationResult result = 2;
    FlowControl flow_control = 3;
    Status status = 4;
    ReleaseResponse release_ack = 5;
    CancelResponse cancel_ack = 6;
  }
}

// ---------------------------------------------------------------------------
// Capability and health (unary)
// ---------------------------------------------------------------------------

message ResourceBudget {
  uint64 weight_bytes = 1;
  uint64 kv_bytes = 2;
  uint64 scratch_bytes = 3;
  uint32 max_concurrent_sessions = 4;
}

message CapabilityRequest {
  SchemaVersion schema_version = 1;
}

message CapabilityResponse {
  SchemaVersion schema_version = 1;
  repeated SchemaVersion supported_schema_versions = 2;
  repeated string supported_architectures = 3;   // e.g. "llama", "qwen3"
  repeated string supported_quantizations = 4;
  ShardRange servable_range = 5;
  ResourceBudget budget = 6;
  repeated Compression supported_compression = 7;
  repeated ChecksumAlgorithm supported_checksums = 8;
  ArtifactFingerprint loaded_fingerprint = 9;     // empty if no artifact loaded
}

message HealthRequest {
  string route_session_id = 1;  // optional; empty for node-wide health
}

message HealthResponse {
  ServingStatus status = 1;
  uint32 active_sessions = 2;
  uint32 queued_requests = 3;
  double kv_pressure = 4;   // 0.0..1.0 fraction of KV budget in use
  uint64 rss_bytes = 5;
  Status detail = 6;
}

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

service ShardRuntime {
  // Admission/capability negotiation.
  rpc GetCapability(CapabilityRequest) returns (CapabilityResponse);

  // Liveness/backpressure telemetry.
  rpc Health(HealthRequest) returns (HealthResponse);

  // One long-lived bidirectional stream per Route Session Activation Seam.
  // Deadlines/cancellation use gRPC call semantics on direct transport and the
  // in-message equivalents on relay transport; flow control uses FlowControl
  // frames; errors are structured Status.
  rpc ActivateSession(stream SessionActivation) returns (stream SessionResponse);

  // Clean resource release (also expressible in-stream as PHASE_RELEASE).
  rpc Release(ReleaseRequest) returns (ReleaseResponse);

  // Cancellation (also expressible in-stream as PHASE_CANCEL).
  rpc Cancel(CancelRequest) returns (CancelResponse);
}