src/engine/device-gate.test.ts

import { afterEach, describe, expect, it, vi } from 'vitest';
import {
  assertBackendPolicy,
  evaluateModelContextPolicy,
  evaluateModelGate,
  inspectWebGpuAdapter,
  supportsWllamaWebGpuRuntime,
} from './device-gate';
import type { ManifestModelV2, ModelTierId } from './manifest';
import type { StorageEstimate, WebGpuAdapterSnapshot } from './protocol';

function model(id: ModelTierId, cpuFallback: boolean): ManifestModelV2 {
  return {
    id,
    displayName: `Bonsai ${id}`,
    architecture: id === '27b' ? 'qwen35' : 'qwen3',
    source: {
      repo: 'prism-ml/Bonsai',
      revision: '0123456789abcdef0123456789abcdef01234567',
      file: 'Bonsai.gguf',
      bytes: 1000,
      sha256: 'a'.repeat(64),
    },
    files: [{ path: 'Bonsai.gguf', bytes: 1000, sha256: 'a'.repeat(64) }],
    downloadBytes: 1000,
    contextLength: 32768,
    defaultContext: 4096,
    cpuFallback,
    largestTensorBytes: 256,
    requiredLimits: { maxStorageBufferBindingSize: 256 },
    chatTemplate: {
      bytes: 10,
      sha256: 'a'.repeat(64),
      markers: { think: true, toolCall: true, toolResponse: true },
    },
    hybridDimensions: {},
    nextNTensorCount: 0,
    runtimePolicy: {
      flashAttention: false,
      tokenEmbeddingOnWebGPU: true,
      requireSingleWebGPUGraph: id === '27b',
    },
  };
}

function adapter(limit = 1024): WebGpuAdapterSnapshot {
  return {
    available: true,
    name: 'Test GPU',
    vendor: null,
    architecture: null,
    device: null,
    description: null,
    limits: { maxStorageBufferBindingSize: limit },
    features: ['shader-f16'],
  };
}

const storage: StorageEstimate = {
  usageBytes: 100,
  quotaBytes: 10_000,
  persisted: true,
};

afterEach(() => {
  vi.unstubAllGlobals();
});

describe('inspectWebGpuAdapter', () => {
  it('inspects the same default adapter selection used by the pinned runtime', async () => {
    const requestAdapter = vi.fn().mockResolvedValue({
      limits: { maxStorageBufferBindingSize: 1024 },
      features: new Set(['shader-f16']),
      info: { description: 'Default test GPU' },
    });
    vi.stubGlobal('navigator', { gpu: { requestAdapter } });

    const snapshot = await inspectWebGpuAdapter();

    expect(requestAdapter).toHaveBeenCalledOnce();
    expect(requestAdapter).toHaveBeenCalledWith();
    expect(snapshot.name).toBe('Default test GPU');
  });

  it('normalizes empty adapter-info strings instead of exporting an empty device name', async () => {
    vi.stubGlobal('navigator', {
      gpu: {
        requestAdapter: vi.fn().mockResolvedValue({
          limits: {},
          features: new Set(),
          info: { vendor: '  Apple ', architecture: '', device: '', description: ' ' },
        }),
      },
    });

    const snapshot = await inspectWebGpuAdapter();

    expect(snapshot).toMatchObject({
      name: 'Apple',
      vendor: 'Apple',
      architecture: null,
      device: null,
      description: null,
    });
  });
});

describe('supportsWllamaWebGpuRuntime', () => {
  const firefox = 'Mozilla/5.0 Gecko/20100101 Firefox/128.0';
  const chromium = 'Mozilla/5.0 Chrome/128.0.0.0 Safari/537.36';

  it('masks WebGPU only on Firefox without JSPI', () => {
    expect(supportsWllamaWebGpuRuntime(firefox, false)).toBe(false);
    expect(supportsWllamaWebGpuRuntime(firefox, true)).toBe(true);
    expect(supportsWllamaWebGpuRuntime(chromium, false)).toBe(true);
  });
});

describe('evaluateModelGate', () => {
  it('fails 27B loudly instead of selecting CPU-WASM when a WebGPU limit is short', () => {
    const decision = evaluateModelGate(model('27b', false), 'auto', adapter(128), storage);

    expect(decision.allowed).toBe(false);
    expect(decision.selectedBackend).toBeNull();
    expect(decision.reasons.join(' ')).toMatch(/forbids CPU-WASM fallback/);
  });

  it('allows an explicit smaller-tier CPU-WASM fallback', () => {
    const decision = evaluateModelGate(model('8b', true), 'auto', adapter(128), storage);

    expect(decision.allowed).toBe(true);
    expect(decision.selectedBackend).toBe('wasm');
    expect(decision.warnings.join(' ')).toMatch(/Falling back to CPU-WASM/);
  });

  it('accounts for already cached bytes when checking quota', () => {
    const constrained = { usageBytes: 9500, quotaBytes: 10_000, persisted: true };

    expect(evaluateModelGate(model('8b', true), 'webgpu', adapter(), constrained, 1000).allowed).toBe(true);
    expect(evaluateModelGate(model('8b', true), 'webgpu', adapter(), constrained, 0).allowed).toBe(false);
  });

  it('falls back safely when the token-embedding path lacks shader-f16', () => {
    const noF16 = { ...adapter(), features: [] };
    const dense = evaluateModelGate(model('8b', true), 'auto', noF16, storage);
    const hybrid = evaluateModelGate(model('27b', false), 'auto', noF16, storage);

    expect(dense.selectedBackend).toBe('wasm');
    expect(dense.warnings.join(' ')).toMatch(/shader-f16/);
    expect(hybrid.allowed).toBe(false);
  });

  it('routes a masked Firefox runtime to CPU only when the model permits it', () => {
    const masked = { ...adapter(), available: false };
    const denseAuto = evaluateModelGate(model('8b', true), 'auto', masked, storage);
    const denseExplicit = evaluateModelGate(model('8b', true), 'webgpu', masked, storage);
    const hybridAuto = evaluateModelGate(model('27b', false), 'auto', masked, storage);

    expect(denseAuto.allowed).toBe(true);
    expect(denseAuto.selectedBackend).toBe('wasm');
    expect(denseExplicit.allowed).toBe(false);
    expect(denseExplicit.selectedBackend).toBeNull();
    expect(hybridAuto.allowed).toBe(false);
    expect(hybridAuto.selectedBackend).toBeNull();
  });
});

describe('evaluateModelContextPolicy', () => {
  const model27b = model('27b', false);
  const longContextExperiment = {
    tuningScope: 'benchmark',
    requestedBackend: 'webgpu',
    flashMode: 'auto',
    kvCacheType: 'q4_0',
  } as const;

  it('allows the bounded 8,448-token 27B benchmark experiment', () => {
    expect(evaluateModelContextPolicy(model27b, 8_448, longContextExperiment)).toEqual({
      allowed: true,
      limit: 8_448,
    });
  });

  it('rejects 8,449 tokens even for the exact long-context experiment', () => {
    expect(evaluateModelContextPolicy(model27b, 8_449, longContextExperiment)).toEqual({
      allowed: false,
      limit: 8_448,
    });
  });

  it.each([
    ['q8 KV', { ...longContextExperiment, kvCacheType: 'q8_0' as const }],
    ['f16 KV', { ...longContextExperiment, kvCacheType: 'f16' as const }],
    ['auto backend', { ...longContextExperiment, requestedBackend: 'auto' as const }],
    ['release defaults', { ...longContextExperiment, tuningScope: 'release-defaults' as const }],
  ])('keeps 27B at the 4,096-token release limit for %s', (_label, input) => {
    expect(evaluateModelContextPolicy(model27b, 8_448, input)).toEqual({
      allowed: false,
      limit: 4_096,
    });
    expect(evaluateModelContextPolicy(model27b, 4_096, input).allowed).toBe(true);
    expect(evaluateModelContextPolicy(model27b, 4_097, input).allowed).toBe(false);
  });

  it('leaves dense-tier manifest limits unchanged', () => {
    expect(evaluateModelContextPolicy(model('8b', true), 32_768, {
      tuningScope: 'release-defaults',
      requestedBackend: 'auto',
      flashMode: 'off',
      kvCacheType: 'f16',
    })).toEqual({ allowed: true, limit: 32_768 });
  });
});

describe('assertBackendPolicy', () => {
  const strictModel = model('27b', false);
  const baseReport = {
    backends: ['WebGPU'],
    nGraphSplits: 1,
    opsOnCpu: 0,
    layersGpu: { offloaded: 65, total: 65 },
    flashAttention: false,
    cacheTypeK: 'f16',
    cacheTypeV: 'f16',
    webgpuKvBufferBytes: 56 * 1024 ** 2,
    webgpuTrace: [],
  };

  it('accepts a proven single WebGPU graph', () => {
    expect(() => assertBackendPolicy(strictModel, 'webgpu', baseReport)).not.toThrow();
  });

  it('rejects both graph splitting and explicit CPU work', () => {
    expect(() => assertBackendPolicy(strictModel, 'webgpu', {
      ...baseReport,
      nGraphSplits: 2,
    })).toThrow(/requires one graph split/);
    expect(() => assertBackendPolicy(strictModel, 'webgpu', {
      ...baseReport,
      opsOnCpu: 1,
    })).toThrow(/CPU op/);
  });

  it('enforces the token-embedding single-graph policy on dense tiers too', () => {
    expect(() => assertBackendPolicy(model('1_7b', true), 'webgpu', {
      ...baseReport,
      layersGpu: { offloaded: 29, total: 29 },
      nGraphSplits: 2,
    })).toThrow(/requires one graph split/);
  });
});