src/engine/device-gate.ts

import type { ManifestModelV2 } from './manifest';
import { DEFAULT_CONTEXT_SIZE } from '../lib/runtime-defaults';
import type {
  BenchmarkFlashMode,
  BenchmarkKvCacheType,
  GateDecision,
  RequestedBackend,
  RuntimeBackend,
  StorageEstimate,
  WebGpuAdapterSnapshot,
} from './protocol';
import type { BackendReport } from './native-log';

export const MODEL_27B_RELEASE_CONTEXT_LIMIT = DEFAULT_CONTEXT_SIZE;
export const MODEL_27B_BENCHMARK_CONTEXT_LIMIT = 8_448;

export interface ModelContextPolicyInput {
  tuningScope: 'release-defaults' | 'benchmark';
  requestedBackend: RequestedBackend;
  flashMode: BenchmarkFlashMode;
  kvCacheType: BenchmarkKvCacheType;
}

export interface ModelContextPolicyDecision {
  allowed: boolean;
  limit: number;
}

const KNOWN_LIMITS = [
  'maxBindGroups',
  'maxBindingsPerBindGroup',
  'maxBufferSize',
  'maxComputeInvocationsPerWorkgroup',
  'maxComputeWorkgroupSizeX',
  'maxComputeWorkgroupSizeY',
  'maxComputeWorkgroupSizeZ',
  'maxComputeWorkgroupStorageSize',
  'maxComputeWorkgroupsPerDimension',
  'maxDynamicStorageBuffersPerPipelineLayout',
  'maxStorageBufferBindingSize',
  'maxStorageBuffersPerShaderStage',
] as const;

interface GpuAdapterLike {
  limits: object;
  features: Iterable<string>;
  info?: {
    vendor?: string;
    architecture?: string;
    device?: string;
    description?: string;
  };
  requestAdapterInfo?: () => Promise<{
    vendor?: string;
    architecture?: string;
    device?: string;
    description?: string;
  }>;
}

interface GpuLike {
  requestAdapter(): Promise<GpuAdapterLike | null>;
}

function navigatorGpu(): GpuLike | null {
  const value = (navigator as Navigator & { gpu?: GpuLike }).gpu;
  return value ?? null;
}

export function supportsWllamaWebGpuRuntime(userAgent: string, hasJspi: boolean): boolean {
  const isFirefox = /Firefox\/([0-9.]+)(?:\s|$)/.test(userAgent);
  return !isFirefox || hasJspi;
}

function hasWebAssemblyJspi(): boolean {
  return Boolean((WebAssembly as typeof WebAssembly & { Suspending?: unknown }).Suspending);
}

function readLimits(limits: object): Record<string, number> {
  const record = limits as Record<string, unknown>;
  const result: Record<string, number> = {};
  const keys = new Set([...KNOWN_LIMITS, ...Object.keys(record)]);
  for (const key of keys) {
    const value = record[key];
    if (typeof value === 'number' && Number.isFinite(value)) {
      result[key] = value;
    }
  }
  return result;
}

export async function inspectWebGpuAdapter(): Promise<WebGpuAdapterSnapshot> {
  if (!supportsWllamaWebGpuRuntime(navigator.userAgent, hasWebAssemblyJspi())) {
    return {
      available: false,
      name: null,
      vendor: null,
      architecture: null,
      device: null,
      description: null,
      limits: {},
      features: [],
    };
  }
  const gpu = navigatorGpu();
  if (!gpu) {
    return {
      available: false,
      name: null,
      vendor: null,
      architecture: null,
      device: null,
      description: null,
      limits: {},
      features: [],
    };
  }
  // Match the adapter selection used by the pinned wllama/llama.cpp WebGPU
  // runtime so the gate never proves limits on a different GPU.
  const adapter = await gpu.requestAdapter();
  if (!adapter) {
    return {
      available: false,
      name: null,
      vendor: null,
      architecture: null,
      device: null,
      description: null,
      limits: {},
      features: [],
    };
  }
  const info = adapter.info ?? await adapter.requestAdapterInfo?.() ?? {};
  const clean = (value: string | undefined): string | null => {
    const trimmed = value?.trim();
    return trimmed ? trimmed : null;
  };
  const vendor = clean(info.vendor);
  const architecture = clean(info.architecture);
  const device = clean(info.device);
  const description = clean(info.description);
  const name = description ?? device ?? architecture ?? vendor ?? 'WebGPU adapter';
  return {
    available: true,
    name,
    vendor,
    architecture,
    device,
    description,
    limits: readLimits(adapter.limits),
    features: [...adapter.features].sort(),
  };
}

export async function estimateStorage(): Promise<StorageEstimate> {
  const [estimate, persisted] = await Promise.all([
    navigator.storage.estimate(),
    navigator.storage.persisted(),
  ]);
  return {
    usageBytes: estimate.usage ?? null,
    quotaBytes: estimate.quota ?? null,
    persisted,
  };
}

export async function persistStorage(): Promise<StorageEstimate & { granted: boolean }> {
  const granted = await navigator.storage.persist();
  return {
    ...await estimateStorage(),
    granted,
  };
}

function availableStorageBytes(storage: StorageEstimate, cachedModelBytes: number): number | null {
  if (storage.quotaBytes === null || storage.usageBytes === null) {
    return null;
  }
  return Math.max(0, storage.quotaBytes - storage.usageBytes + cachedModelBytes);
}

export function evaluateModelContextPolicy(
  model: ManifestModelV2,
  contextSize: number,
  input: ModelContextPolicyInput,
): ModelContextPolicyDecision {
  const longContextExperiment = model.id === '27b'
    && input.tuningScope === 'benchmark'
    && input.requestedBackend === 'webgpu'
    && input.flashMode === 'auto'
    && input.kvCacheType === 'q4_0';
  const limit = model.id === '27b'
    ? Math.min(
        model.contextLength,
        longContextExperiment
          ? MODEL_27B_BENCHMARK_CONTEXT_LIMIT
          : MODEL_27B_RELEASE_CONTEXT_LIMIT,
      )
    : model.contextLength;
  return {
    allowed: Number.isSafeInteger(contextSize) && contextSize > 0 && contextSize <= limit,
    limit,
  };
}

function webGpuFailures(model: ManifestModelV2, adapter: WebGpuAdapterSnapshot): string[] {
  if (!adapter.available) {
    return ['WebGPU adapter is unavailable.'];
  }
  const failures: string[] = [];
  if (model.runtimePolicy.tokenEmbeddingOnWebGPU && !adapter.features.includes('shader-f16')) {
    failures.push('WebGPU adapter lacks shader-f16, which the token-embedding single-graph path requires.');
  }
  for (const [limit, required] of Object.entries(model.requiredLimits)) {
    const observed = adapter.limits[limit];
    if (observed === undefined) {
      failures.push(`WebGPU adapter does not report required limit ${limit}.`);
    } else if (observed < required) {
      failures.push(`WebGPU ${limit} is ${observed}, but ${model.id} requires ${required}.`);
    }
  }
  return failures;
}

export function evaluateModelGate(
  model: ManifestModelV2,
  requestedBackend: RequestedBackend,
  adapter: WebGpuAdapterSnapshot,
  storage: StorageEstimate,
  cachedModelBytes = 0,
): GateDecision {
  const reasons: string[] = [];
  const warnings: string[] = [];
  const availableBytes = availableStorageBytes(storage, cachedModelBytes);
  if (availableBytes !== null && availableBytes < model.downloadBytes) {
    reasons.push(
      `Browser storage has ${availableBytes} bytes available for this model, but ${model.downloadBytes} bytes are required.`,
    );
  } else if (availableBytes === null) {
    warnings.push('Browser storage quota is unavailable; download capacity could not be proven.');
  }

  const gpuFailures = webGpuFailures(model, adapter);
  let selectedBackend: RuntimeBackend | null = null;
  if (requestedBackend === 'wasm') {
    if (!model.cpuFallback) {
      reasons.push(`${model.displayName} is WebGPU-only and cannot fall back to CPU-WASM.`);
    } else {
      selectedBackend = 'wasm';
    }
  } else if (requestedBackend === 'webgpu') {
    if (gpuFailures.length > 0) {
      reasons.push(...gpuFailures);
    } else {
      selectedBackend = 'webgpu';
    }
  } else if (gpuFailures.length === 0) {
    selectedBackend = 'webgpu';
  } else if (model.cpuFallback) {
    selectedBackend = 'wasm';
    warnings.push(...gpuFailures.map((failure) => `${failure} Falling back to CPU-WASM.`));
  } else {
    reasons.push(...gpuFailures);
    reasons.push(`${model.displayName} forbids CPU-WASM fallback; choose a smaller tier.`);
  }

  if (reasons.length > 0) {
    selectedBackend = null;
  }
  return {
    allowed: reasons.length === 0,
    requestedBackend,
    selectedBackend,
    reasons,
    warnings,
    requiredBytes: model.downloadBytes,
    availableStorageBytes: availableBytes,
  };
}

export function assertBackendPolicy(
  model: ManifestModelV2,
  backend: RuntimeBackend,
  report: BackendReport,
): void {
  if (backend !== 'webgpu') {
    return;
  }
  const requireSingleGraph = model.runtimePolicy.tokenEmbeddingOnWebGPU
    || model.runtimePolicy.requireSingleWebGPUGraph;
  if (!requireSingleGraph) return;
  if (report.opsOnCpu !== 0) {
    throw new Error(
      `${model.displayName} WebGPU tripwire: native log reports ${report.opsOnCpu} CPU op(s).`,
    );
  }
  if (report.nGraphSplits !== 1) {
    throw new Error(
      `${model.displayName} WebGPU tripwire: token embedding policy requires one graph split, observed ${String(report.nGraphSplits)}.`,
    );
  }
  if (!report.backends.includes('WebGPU')) {
    throw new Error(`${model.displayName} WebGPU tripwire: native log did not confirm the WebGPU backend.`);
  }
  if (report.layersGpu === null || report.layersGpu.offloaded !== report.layersGpu.total) {
    throw new Error(`${model.displayName} WebGPU tripwire: native log did not confirm all layers on the GPU.`);
  }
}