src/engine/protocol.ts

import type { ModelManifestV2, ModelTierId } from './manifest';
import type { BackendReport } from './native-log';

export type RuntimeBackend = 'webgpu' | 'wasm';
export type RequestedBackend = 'auto' | RuntimeBackend;

export interface StorageEstimate {
  usageBytes: number | null;
  quotaBytes: number | null;
  persisted: boolean;
}

export interface WebGpuAdapterSnapshot {
  available: boolean;
  name: string | null;
  vendor: string | null;
  architecture: string | null;
  device: string | null;
  description: string | null;
  limits: Record<string, number>;
  features: string[];
}

export interface EngineCapabilities {
  crossOriginIsolated: boolean;
  sharedArrayBuffer: boolean;
  hardwareConcurrency: number;
  webgpu: WebGpuAdapterSnapshot;
  storage: StorageEstimate;
  browser: {
    userAgent: string;
    platform: string;
    language: string;
  };
  runtime: {
    implementation: 'bonsai-wllama';
    wllamaRevision: string;
    llamaCppRevision: string;
    patchSetSha256: string;
    moduleSha256: string;
    wasmFlavor: 'jspi' | 'compat';
    wasmSha256: string;
    compatWorkerSha256: string | null;
    tokenEmbeddingOnWebGPU: true;
    tensorPlacementOverrides: false;
  };
}

export interface GateDecision {
  allowed: boolean;
  requestedBackend: RequestedBackend;
  selectedBackend: RuntimeBackend | null;
  reasons: string[];
  warnings: string[];
  requiredBytes: number;
  availableStorageBytes: number | null;
}

export interface FutureTensorPlacement {
  engine: 'future-custom-build';
  overrides: Array<{
    tensor: string;
    backend: RuntimeBackend;
  }>;
}

export type BenchmarkFlashMode = 'off' | 'auto';
export type BenchmarkKvCacheType = 'f16' | 'q8_0' | 'q4_0';
export type BenchmarkWasmFlavor = 'auto' | 'jspi' | 'compat';

export interface BenchmarkLoadTuning {
  nBatch?: number;
  nUbatch?: number;
  flashMode: BenchmarkFlashMode;
  kvCacheType: BenchmarkKvCacheType;
  wasmFlavor: BenchmarkWasmFlavor;
}

export interface EngineChatToolCall {
  id: string;
  type: 'function';
  function: {
    name: string;
    arguments: string;
  };
}

export type EngineChatMessage =
  | { role: 'system' | 'user'; content: string; name?: string }
  | { role: 'assistant'; content?: string | null; name?: string; tool_calls?: EngineChatToolCall[] }
  | { role: 'tool'; content: string; tool_call_id: string };

export interface EngineChatTool {
  type: 'function';
  function: {
    name: string;
    description?: string;
    parameters: {
      type: 'object';
      properties: Record<string, {
        type: string;
        description?: string;
        enum?: string[];
        [key: string]: unknown;
      }>;
      required?: string[];
      additionalProperties?: boolean;
    };
  };
}

type ManifestSource =
  | { manifestUrl: string; manifest?: never }
  | { manifest: ModelManifestV2; manifestUrl?: never };

export type LoadModelParams = ManifestSource & {
  modelId: ModelTierId;
  backend: RequestedBackend;
  contextSize?: number;
  threads?: number;
  benchmarkTuning?: BenchmarkLoadTuning;
  tensorPlacement?: FutureTensorPlacement;
};

export interface GenerateParams {
  messages: EngineChatMessage[];
  maxTokens?: number;
  temperature?: number;
  topP?: number;
  topK?: number;
  minP?: number;
  seed?: number;
  tools?: EngineChatTool[];
  toolChoice?: 'none' | 'auto' | 'required';
  cachePrompt?: boolean;
  returnTokenIds?: boolean;
}

export interface LoadModelResult {
  modelId: ModelTierId;
  backend: RuntimeBackend;
  gate: GateDecision;
  shardUrls: string[];
  context: {
    size: number;
    trainingSize: number;
    layerCount: number;
    vocabularySize: number;
    batchSize: number;
    microBatchSize: number;
  };
  tuning: {
    scope: 'release-defaults' | 'benchmark';
    requested: {
      nBatch: number | null;
      nUbatch: number | null;
      flashMode: BenchmarkFlashMode;
      cacheTypeK: BenchmarkKvCacheType;
      cacheTypeV: BenchmarkKvCacheType;
      wasmFlavor: BenchmarkWasmFlavor;
    };
    observed: {
      nBatch: number;
      nUbatch: number;
      flashAttention: boolean | null;
      cacheTypeK: string | null;
      cacheTypeV: string | null;
      kvBufferBytes: number | null;
      wasmFlavor: 'jspi' | 'compat';
      wasmSha256: string;
      compatWorkerSha256: string | null;
    };
    applied: boolean;
  };
  chatTemplate: {
    bytes: number;
    hasThinkMarker: boolean;
    hasToolCallMarker: boolean;
    hasToolResponseMarker: boolean;
  };
  backendReport: BackendReport;
}

export interface EngineTokenLogprob {
  id: number;
  logprob: number;
}

export interface EngineSampledTokenTraceEntry {
  selected: EngineTokenLogprob;
  topCandidates: EngineTokenLogprob[];
}

export interface GenerateResult {
  text: string;
  reasoningText: string;
  toolCallError?: string;
  tokenIds: number[] | null;
  tokenTrace: EngineSampledTokenTraceEntry[] | null;
  tokenTraceAccounting?: {
    usageCompletionTokens: number | null;
    tracedTokens: number;
    delta: number | null;
  } | null;
  finishReason: 'stop' | 'length' | 'tool_calls' | 'content_filter' | null;
  toolCalls: EngineChatToolCall[];
  usage: {
    promptTokens: number;
    completionTokens: number;
    totalTokens: number;
  } | null;
  timings: {
    promptTokensPerSecond: number;
    predictedTokensPerSecond: number;
  } | null;
}

export interface ScoreSequenceParams {
  promptTokenIds: number[];
  referenceTokenIds: number[];
  topK: 5;
}

export interface EngineTeacherForcedScoreEntry {
  index: number;
  selectedReference: EngineTokenLogprob;
  naturalTop1: EngineTokenLogprob;
  topCandidates: EngineTokenLogprob[];
  referenceRankInTopCandidatesZeroBased: number | null;
  top1Top2Margin: number;
}

export interface ScoreSequenceResult {
  method: {
    promptMode: 'raw-token-id-prefix';
    maxTokensPerStep: 1;
    temperature: 0;
    topK: 1;
    reportedTopLogprobs: 5;
    logitBias: 1_000;
    cachePromptFirst: false;
    cachePromptSubsequent: true;
  };
  entries: EngineTeacherForcedScoreEntry[];
  summary: {
    tokenCount: 1_024;
    meanNll: number;
    perplexity: number;
  };
}

export interface EmptyParams {
  readonly _empty?: never;
}

export interface EngineRequestMap {
  capabilities: EmptyParams;
  loadModel: LoadModelParams;
  generate: GenerateParams;
  scoreSequence: ScoreSequenceParams;
  abort: { targetRequestId: string };
  unload: EmptyParams;
  backendReport: EmptyParams;
  storageEstimate: EmptyParams;
  storagePersist: EmptyParams;
  storageClear: EmptyParams;
}

export interface EngineResponseMap {
  capabilities: EngineCapabilities;
  loadModel: LoadModelResult;
  generate: GenerateResult;
  scoreSequence: ScoreSequenceResult;
  abort: { targetRequestId: string; aborted: boolean };
  unload: { unloaded: true };
  backendReport: BackendReport;
  storageEstimate: StorageEstimate;
  storagePersist: StorageEstimate & { granted: boolean };
  storageClear: StorageEstimate & { cleared: true };
}

export type EngineMethod = keyof EngineRequestMap;

export type EngineRequest = {
  [Method in EngineMethod]: {
    type: 'request';
    requestId: string;
    method: Method;
    params: EngineRequestMap[Method];
  };
}[EngineMethod];

export interface EngineErrorPayload {
  code: string;
  message: string;
  details?: unknown;
}

export type ShardDownloadFailureKind = 'network' | 'abort' | 'verification';

export interface ShardDownloadFailureDetails {
  kind: 'shard-download-failure';
  failure: ShardDownloadFailureKind;
  causeCode: string;
  shardIndex: number;
  shardCount: number;
  shardPath: string;
  retryFromByteZero: true;
  partialDeleted: boolean;
}

export function isShardDownloadFailureDetails(
  value: unknown,
): value is ShardDownloadFailureDetails {
  if (typeof value !== 'object' || value === null) return false;
  const details = value as Partial<ShardDownloadFailureDetails>;
  return details.kind === 'shard-download-failure'
    && (details.failure === 'network' || details.failure === 'abort' || details.failure === 'verification')
    && typeof details.causeCode === 'string'
    && Number.isSafeInteger(details.shardIndex)
    && (details.shardIndex ?? -1) >= 0
    && Number.isSafeInteger(details.shardCount)
    && (details.shardCount ?? 0) > 0
    && (details.shardIndex ?? 0) < (details.shardCount ?? 0)
    && typeof details.shardPath === 'string'
    && details.shardPath.length > 0
    && details.retryFromByteZero === true
    && typeof details.partialDeleted === 'boolean';
}

export type EngineResponse = {
  [Method in EngineMethod]:
    | {
      type: 'response';
      requestId: string;
      method: Method;
      ok: true;
      result: EngineResponseMap[Method];
    }
    | {
      type: 'response';
      requestId: string;
      method: Method;
      ok: false;
      error: EngineErrorPayload;
    };
}[EngineMethod];

export type EngineShardProgressState =
  | 'queued'
  | 'cached'
  | 'downloading'
  | 'verifying'
  | 'complete'
  | 'error';

export interface EngineShardProgress {
  index: number;
  path: string;
  loadedBytes: number;
  verifiedBytes: number;
  totalBytes: number;
  state: EngineShardProgressState;
}

export type EngineEvent =
  | {
    type: 'event';
    requestId: string;
    event: 'progress';
    phase: 'manifest' | 'cache' | 'download' | 'verify' | 'load';
    loadedBytes: number;
    totalBytes: number;
    shardIndex: number | null;
    shardCount: number;
    shardPath: string | null;
    stageProgress: number | null;
    nativeStage: string | null;
    residentBytes: number | null;
    shards: EngineShardProgress[];
  }
  | {
    type: 'event';
    requestId: string;
    event: 'token';
    text: string;
    reasoningDelta?: string;
  }
  | {
    type: 'event';
    requestId: string;
    event: 'tool-call';
    index: number;
    id: string;
    name: string;
    argumentCharacters: number;
  }
  | {
    type: 'event';
    requestId: string;
    event: 'generation';
    phase: 'prefill' | 'decode';
    promptProcessed: number;
    promptTotal: number;
    promptCached: number;
    completionTokens: number;
    elapsedMs: number;
    promptTokensPerSecond: number;
    decodeTokensPerSecond: number;
  };

export type EngineWorkerMessage = EngineResponse | EngineEvent;