syntax = "proto3";

package skippy.stage.v1;

// Skippy owns these stage-control, artifact-transfer, and activation frames.
// Control and artifact streams are carried inside the mesh STREAM_SUBPROTOCOL
// envelope for skippy-stage/1. Activation transport remains on the
// skippy-stage/1 QUIC ALPN.

message StageControlRequest {
  uint32 gen = 1;
  bytes requester_id = 2;
  oneof command {
    LoadStage load_stage = 3;
    StopStage stop_stage = 4;
    GetStageStatus get_stage_status = 5;
    GetLayerInventory get_layer_inventory = 6;
    PrepareStage prepare_stage = 7;
    CancelPrepareStage cancel_prepare_stage = 8;
    StageStatusUpdate stage_status_update = 9;
    ClaimCoordinator claim_coordinator = 10;
  }
}

message StageControlResponse {
  uint32 gen = 1;
  oneof response {
    StageReady stage_ready = 2;
    StageStatus stage_status = 3;
    LayerInventory layer_inventory = 4;
    PrepareStageAccepted prepare_stage_accepted = 5;
    StagePreparationStatus stage_preparation_status = 6;
    StageStatusAck stage_status_ack = 7;
    StageStatusList stage_statuses = 8;
    CoordinatorClaimAccepted coordinator_claim_accepted = 9;
  }
}

message StageArtifactTransferRequest {
  uint32 gen = 1;
  bytes requester_id = 2;
  string topology_id = 3;
  string run_id = 4;
  string stage_id = 5;
  string package_ref = 6;
  string manifest_sha256 = 7;
  string relative_path = 8;
  uint64 offset = 9;
  optional uint64 expected_size = 10;
  optional string expected_sha256 = 11;
}

message StageArtifactTransferResponse {
  uint32 gen = 1;
  bool accepted = 2;
  uint64 total_size = 3;
  optional string sha256 = 4;
  optional string error = 5;
}

message LoadStage {
  string topology_id = 1;
  string run_id = 2;
  string model_id = 3;
  string backend = 4;
  string package_ref = 5;
  string manifest_sha256 = 6;
  string stage_id = 7;
  uint32 stage_index = 8;
  uint32 layer_start = 9;
  uint32 layer_end = 10;
  optional string model_path = 11;
  optional StageDevice selected_device = 12;
  string bind_addr = 13;
  uint32 activation_width = 14;
  StageWireDType wire_dtype = 15;
  uint32 ctx_size = 16;
  int32 n_gpu_layers = 17;
  string cache_type_k = 18;
  string cache_type_v = 19;
  uint64 shutdown_generation = 20;
  StageLoadMode load_mode = 21;
  optional StagePeer upstream = 22;
  optional StagePeer downstream = 23;
  optional string projector_path = 24;
  uint32 lane_count = 25;
  optional uint32 n_batch = 26;
  optional uint32 n_ubatch = 27;
  StageFlashAttnType flash_attn_type = 28;
  optional uint64 source_model_bytes = 29;
  uint64 coordinator_term = 30;
  optional string coordinator_id = 31;
  uint64 lease_until_unix_ms = 32;
}

message StopStage {
  string topology_id = 1;
  string run_id = 2;
  string stage_id = 3;
  uint64 shutdown_generation = 4;
  uint64 coordinator_term = 5;
}

message GetStageStatus {
  optional string topology_id = 1;
  optional string run_id = 2;
  optional string stage_id = 3;
}

message GetLayerInventory {
  string model_id = 1;
  string package_ref = 2;
  string manifest_sha256 = 3;
}

message PrepareStage {
  LoadStage load_stage = 1;
  optional bytes coordinator_id = 2;
}

message CancelPrepareStage {
  string topology_id = 1;
  string run_id = 2;
  string stage_id = 3;
  uint64 shutdown_generation = 4;
}

message StageStatusUpdate {
  StagePreparationStatus status = 1;
}

message ClaimCoordinator {
  string model_id = 1;
  string package_ref = 2;
  string manifest_sha256 = 3;
  string topology_id = 4;
  string run_id = 5;
  string coordinator_id = 6;
  uint64 coordinator_term = 7;
  string participant_set_hash = 8;
  string topology_hash = 9;
  uint64 lease_until_unix_ms = 10;
}

message StageReady {
  bool accepted = 1;
  StageStatus status = 2;
  optional string error = 3;
}

message PrepareStageAccepted {
  bool accepted = 1;
  StagePreparationStatus status = 2;
  optional string error = 3;
}

message StageStatusAck {
  bool accepted = 1;
  optional string error = 2;
}

message CoordinatorClaimAccepted {
  bool accepted = 1;
  ClaimCoordinator claim = 2;
  optional string error = 3;
}

message StageStatus {
  string topology_id = 1;
  string run_id = 2;
  string model_id = 3;
  string backend = 4;
  string stage_id = 5;
  uint32 stage_index = 6;
  uint32 layer_start = 7;
  uint32 layer_end = 8;
  StageRuntimeState state = 9;
  string bind_addr = 10;
  uint32 activation_width = 11;
  StageWireDType wire_dtype = 12;
  optional string error = 13;
  uint64 shutdown_generation = 14;
  optional StageDevice selected_device = 15;
  uint32 ctx_size = 16;
  optional string package_ref = 17;
  optional string manifest_sha256 = 18;
  optional string source_model_path = 19;
  optional string source_model_sha256 = 20;
  optional uint64 source_model_bytes = 21;
  optional string materialized_path = 22;
  optional bool materialized_pinned = 23;
  optional string projector_path = 24;
  uint32 lane_count = 25;
  optional uint32 n_batch = 26;
  optional uint32 n_ubatch = 27;
  StageFlashAttnType flash_attn_type = 28;
  uint64 coordinator_term = 29;
  optional string coordinator_id = 30;
  uint64 lease_until_unix_ms = 31;
}

message StageStatusList {
  repeated StageStatus statuses = 1;
}

message LayerInventory {
  string model_id = 1;
  string package_ref = 2;
  string manifest_sha256 = 3;
  uint32 layer_count = 4;
  repeated LayerRange ready_ranges = 5;
  repeated LayerRange available_ranges = 6;
  repeated LayerRange missing_ranges = 7;
  repeated StagePreparationStatus preparing_ranges = 8;
  optional string source_model_path = 9;
  optional uint64 source_model_bytes = 10;
  SourceModelKind source_model_kind = 11;
}

message LayerRange {
  uint32 layer_start = 1;
  uint32 layer_end = 2;
}

message StagePreparationStatus {
  string topology_id = 1;
  string run_id = 2;
  string model_id = 3;
  string backend = 4;
  string package_ref = 5;
  string manifest_sha256 = 6;
  string stage_id = 7;
  uint32 stage_index = 8;
  uint32 layer_start = 9;
  uint32 layer_end = 10;
  StagePreparationState state = 11;
  optional uint64 bytes_done = 12;
  optional uint64 bytes_total = 13;
  optional string bind_addr = 14;
  optional string error = 15;
  uint64 shutdown_generation = 16;
  uint64 coordinator_term = 17;
  optional string coordinator_id = 18;
  uint64 lease_until_unix_ms = 19;
}

enum SourceModelKind {
  SOURCE_MODEL_KIND_UNSPECIFIED = 0;
  SOURCE_MODEL_KIND_LAYER_PACKAGE = 1;
  SOURCE_MODEL_KIND_PLAIN_GGUF = 2;
  SOURCE_MODEL_KIND_SPLIT_GGUF = 3;
}

enum StageFlashAttnType {
  STAGE_FLASH_ATTN_TYPE_UNSPECIFIED = 0;
  STAGE_FLASH_ATTN_TYPE_AUTO = 1;
  STAGE_FLASH_ATTN_TYPE_DISABLED = 2;
  STAGE_FLASH_ATTN_TYPE_ENABLED = 3;
}

message StagePeer {
  string stage_id = 1;
  uint32 stage_index = 2;
  string endpoint = 3;
  optional bytes node_id = 4;
}

message StageDevice {
  string backend_device = 1;
  optional string stable_id = 2;
  optional uint64 index = 3;
  optional uint64 vram_bytes = 4;
}

message StageTransportOpen {
  uint32 gen = 1;
  bytes requester_id = 2;
  string topology_id = 3;
  string run_id = 4;
  string stage_id = 5;
}

enum StageLoadMode {
  STAGE_LOAD_MODE_UNSPECIFIED = 0;
  STAGE_LOAD_MODE_RUNTIME_SLICE = 1;
  STAGE_LOAD_MODE_LAYER_PACKAGE = 2;
  STAGE_LOAD_MODE_ARTIFACT_SLICE = 3;
}

enum StageWireDType {
  STAGE_WIRE_DTYPE_UNSPECIFIED = 0;
  STAGE_WIRE_DTYPE_F32 = 1;
  STAGE_WIRE_DTYPE_F16 = 2;
  STAGE_WIRE_DTYPE_Q8 = 3;
}

enum StageRuntimeState {
  STAGE_RUNTIME_STATE_UNSPECIFIED = 0;
  STAGE_RUNTIME_STATE_STARTING = 1;
  STAGE_RUNTIME_STATE_READY = 2;
  STAGE_RUNTIME_STATE_STOPPING = 3;
  STAGE_RUNTIME_STATE_STOPPED = 4;
  STAGE_RUNTIME_STATE_FAILED = 5;
}

enum StagePreparationState {
  STAGE_PREPARATION_STATE_UNSPECIFIED = 0;
  STAGE_PREPARATION_STATE_ASSIGNED = 1;
  STAGE_PREPARATION_STATE_DOWNLOADING = 2;
  STAGE_PREPARATION_STATE_AVAILABLE = 3;
  STAGE_PREPARATION_STATE_RESOLVING = 4;
  STAGE_PREPARATION_STATE_LOADING = 5;
  STAGE_PREPARATION_STATE_READY = 6;
  STAGE_PREPARATION_STATE_FAILED = 7;
  STAGE_PREPARATION_STATE_CANCELLED = 8;
}
