// SPDX-FileCopyrightText: Copyright (c) 2026 The SGLang Authors
// SPDX-License-Identifier: Apache-2.0

syntax = "proto3";

package kv_indexer.v1;

// Storage tier for externally managed KV cache blocks.
enum TierType {
  TIER_UNKNOWN = 0;
  TIER_HBM = 1;
  TIER_DRAM = 2;
  TIER_SSD = 3;
}

// KV components are a fixed set; each component's rule is a property of its
// type, not a per-worker choice:
//   FULL  (bit 0, value 1) -- resident on every prefix block.
//   SWA   (bit 1, value 2) -- resident contiguously over at least
//                             `swa_window_tokens` tokens ending at the boundary,
//                             or in an unbroken run from the head.
//   MAMBA (bit 2, value 4) -- resident on the boundary block only.
// Carried as a bitmask; FULL is always present on a stored block.

// A worker's versioned cache spec: which components its cache holds and the
// parameters the fixed rules need. Absent means legacy / full-only behaviour.
message WorkerCacheSpec {
  uint32 version = 1;
  // Bitmask of components this worker's cache holds (FULL|SWA|MAMBA); selects
  // which fixed rules gate the prefix.
  uint32 components = 2;
  // Sliding-window size in tokens; used only when SWA is present.
  uint32 swa_window_tokens = 3;
  // Per-component servable tiers, each a bitmask of (1 << TierType), intersected
  // with the indexer's servable set (V1: HBM+DRAM) at query time.
  uint32 full_tier_mask = 4;
  uint32 swa_tier_mask = 5;
  uint32 mamba_tier_mask = 6;
}

message MatchExternalKvRequest {
  // Block-level hashes matched against placement metadata. This RPC returns
  // placement matches only; it does not compute the longest reusable prefix.
  repeated sfixed64 hashes = 1 [packed = true];

  // When true, only hashes that are actually matched should increase hit
  // counters. Diagnostic callers should leave this false.
  bool count_as_hit = 2;
}

// The kind of mutation a single ExternalKvAction carries. A whole SGLang
// KVEventBatch is applied in one call while preserving exact action order.
enum ExternalKvActionType {
  ACTION_UNKNOWN = 0;
  ACTION_REPORT = 1;
  ACTION_REVOKE = 2;
  ACTION_CLEAR_ALL_AT_TIER = 3;
}

message ExternalKvAction {
  ExternalKvActionType type = 1;
  TierType tier = 2;
  // Non-empty for REPORT/REVOKE; ignored for CLEAR_ALL_AT_TIER.
  repeated sfixed64 hashes = 3 [packed = true];

  // REPORT only. Per-hash resident component bitmask at this tier, a REPLACE
  // snapshot index-aligned with `hashes`. Empty means every hash is a legacy
  // whole-block store (no component detail). REVOKE/CLEAR ignore it.
  repeated uint32 component_masks = 4;

  // REPORT only. Per-hash token count (block_size), index-aligned with `hashes`,
  // used to accumulate SWA trailing windows. Empty when not supplied (legacy).
  repeated uint32 block_sizes = 5;

  // REPORT only. Parent of hashes[0]; absent means hashes[0] is a root block.
  // Every later hash is the child of the preceding hash. The Indexer and Bridge
  // are deployed together, so this new protocol does not support old senders.
  optional sfixed64 parent_block_hash = 6;
}

message ApplyExternalKvBatchRequest {
  string worker_id = 1;
  // The SGLang batch sequence number, monotonic per worker. Observability only:
  // every batch is applied, with no deduplication, fencing, or checkpointing.
  uint64 seq = 2;
  repeated ExternalKvAction actions = 3;
  // The worker's KV-transfer address, used to populate MatchExternalKvResponse.
  // Supplied by the bridge; may be empty, and then matches carry no address.
  string worker_address = 4;

  // Field 5 held the worker incarnation token used for restart fencing.
  reserved 5;

  // The worker's component cache spec, carried on every batch so it self-heals
  // across reconnects and indexer restarts. Absent means legacy / full-only.
  WorkerCacheSpec cache_spec = 6;
}

// Deliberately empty: applies are unconditional, so there is nothing to report
// back. Kept as a message so the RPC signature stays stable.
message ApplyExternalKvBatchResponse {
  // Fields 1-3 held the durable seq checkpoint and duplicate flag.
  reserved 1, 2, 3;
}

message TierHashes {
  TierType tier = 1;
  repeated sfixed64 hashes = 2 [packed = true];
  // Diagnostic snapshots aligned with `hashes`. Empty only for backends that
  // cannot expose component detail.
  repeated uint32 component_masks = 3;
  repeated uint32 block_sizes = 4;
}

message ExternalKvNodeMatch {
  string worker_id = 1;
  string address = 2;
  repeated TierHashes hashes_by_tier = 3;
}

message MatchExternalKvResponse {
  repeated ExternalKvNodeMatch matches = 1;
}

message MatchExternalKvPrefixRequest {
  // Block hashes in prompt order. hashes[0] MUST be the request's first block:
  // prefix matching starts there and stops at the first block a worker is
  // missing, so a misordered list silently truncates every match.
  repeated sfixed64 hashes = 1 [packed = true];

  // Caller-supplied ceiling on how many leading blocks to consider. 0 means no
  // caller ceiling: the server considers every block sent. Unlike the mutating
  // RPCs, a prefix query is not rejected for length; its scan holds O(1) state
  // per candidate worker, so length costs time rather than memory.
  uint32 max_blocks = 2;
}

message ExternalKvPrefixMatch {
  // The worker's router-facing routing identity, NOT its KV-transfer address.
  // The router intersects this byte-for-byte with the worker URLs it registered,
  // so a mismatch silently disables cache-aware routing. Workers with an empty
  // address are excluded from this response entirely.
  string worker_address = 1;

  // Largest n such that this worker holds hashes[0..n) contiguously.
  uint32 matched_prefix_blocks = 2;

  // Opaque worker id, carried for the caller's logs only; not a routing key.
  string worker_id = 3;
}

message MatchExternalKvPrefixResponse {
  // Matches sorted by matched_prefix_blocks, descending.
  repeated ExternalKvPrefixMatch matches = 1;

  // The longest contiguous prefix held by any single worker (0 when no match).
  uint32 best_prefix_blocks = 2;

  // How many blocks the server actually read placement for, so early termination
  // and truncation are observable. NOT part of the prefix semantics: backends may
  // report different values for the same matches.
  uint32 blocks_read = 3;
}

message HitCountEntry {
  sfixed64 hash = 1;
  uint64 hit_count_total = 2;
}

message GetExternalKvHitCountsRequest {
  repeated sfixed64 hashes = 1 [packed = true];
}

message GetExternalKvHitCountsResponse {
  repeated HitCountEntry entries = 1;
}

service KVIndexer {
  // The sole mutation API: applies an ordered SGLang KVEventBatch.
  rpc ApplyExternalKvBatch(ApplyExternalKvBatchRequest) returns (ApplyExternalKvBatchResponse);

  rpc MatchExternalKv(MatchExternalKvRequest) returns (MatchExternalKvResponse);

  // Answers, per worker, how long a contiguous prefix of the request it holds.
  // The indexer never picks a worker: it cannot see the router's health, load, or
  // pool split, so the final choice stays with the router.
  rpc MatchExternalKvPrefix(MatchExternalKvPrefixRequest) returns (MatchExternalKvPrefixResponse);

  rpc GetExternalKvHitCounts(GetExternalKvHitCountsRequest) returns (GetExternalKvHitCountsResponse);
}