src/engine/runtime-download.test.ts

import { describe, expect, it } from 'vitest';
import type { Wllama } from '../../vendor/wllama-bonsai/esm/index.js';
import { verifyBlobSha256 } from './blob-integrity';
import { EngineRuntimeError } from './errors';
import type { ManifestModelV2 } from './manifest';
import { isShardDownloadFailureDetails, type EngineEvent } from './protocol';
import {
  BrowserEngineRuntime,
  resolveLoadTuning,
  resolveRuntimeBatchShape,
  selectWasmFlavor,
} from './runtime';

type DownloadOptions = {
  signal?: AbortSignal;
  headers?: Record<string, string>;
  progressCallback?: (progress: { loaded: number; total: number }) => void;
};

class FakeCacheManager {
  readonly blobs = new Map<string, Blob>();
  readonly downloads: Array<{ url: string; headers?: Record<string, string> }> = [];
  readonly deletes: string[] = [];
  readonly payloads = new Map<string, Blob>();
  failNetworkFor: string | null = null;
  readonly waitForAbortFor = new Set<string>();
  private readonly downloadWaiters: Array<{ count: number; resolve: () => void }> = [];

  waitForDownloads(count: number): Promise<void> {
    if (this.downloads.length >= count) return Promise.resolve();
    return new Promise<void>((resolve) => this.downloadWaiters.push({ count, resolve }));
  }

  private resolveDownloadWaiters(): void {
    for (let index = this.downloadWaiters.length - 1; index >= 0; index -= 1) {
      const waiter = this.downloadWaiters[index];
      if (waiter && this.downloads.length >= waiter.count) {
        this.downloadWaiters.splice(index, 1);
        waiter.resolve();
      }
    }
  }

  async download(url: string, options: DownloadOptions = {}): Promise<void> {
    this.downloads.push({ url, ...(options.headers ? { headers: options.headers } : {}) });
    this.resolveDownloadWaiters();
    const payload = this.payloads.get(url) ?? new Blob(['partial']);
    this.blobs.set(url, payload);
    options.progressCallback?.({ loaded: payload.size, total: payload.size });
    if (this.failNetworkFor === url) {
      throw new TypeError('synthetic network failure');
    }
    if (this.waitForAbortFor.has(url)) {
      await new Promise<never>((_resolve, reject) => {
        const rejectAborted = () => reject(new DOMException('synthetic abort', 'AbortError'));
        if (options.signal?.aborted) rejectAborted();
        else options.signal?.addEventListener('abort', rejectAborted, { once: true });
      });
    }
  }

  async open(url: string): Promise<Blob | null> {
    return this.blobs.get(url) ?? null;
  }

  async delete(url: string): Promise<void> {
    this.deletes.push(url);
    this.blobs.delete(url);
  }

  async getNameFromURL(url: string): Promise<string> {
    return `cache:${url}`;
  }

  async writeMetadata(): Promise<void> {}
}

type DownloadShards = (
  requestId: string,
  wllama: Wllama,
  model: ManifestModelV2,
  urls: readonly string[],
  cached: { blobs: Array<Blob | null>; cachedBytes: number },
  signal: AbortSignal,
  sink: (event: EngineEvent) => void,
) => Promise<Blob[]>;

async function sha256(blob: Blob): Promise<string> {
  return (await verifyBlobSha256(blob, '')).actualSha256;
}

async function fixture(first: Blob, second: Blob): Promise<{
  model: ManifestModelV2;
  urls: [string, string];
}> {
  const result = await fixtureShards([first, second]);
  const firstUrl = result.urls[0];
  const secondUrl = result.urls[1];
  if (!firstUrl || !secondUrl) throw new Error('Two-shard fixture is incomplete.');
  return { model: result.model, urls: [firstUrl, secondUrl] };
}

async function fixtureShards(payloads: readonly Blob[]): Promise<{
  model: ManifestModelV2;
  urls: string[];
}> {
  const urls = payloads.map((_, index) => `https://models.test/shard-${index}.gguf`);
  const files = await Promise.all(payloads.map(async (payload, index) => ({
    path: `shard-${index}.gguf`,
    bytes: payload.size,
    sha256: await sha256(payload),
  })));
  return {
    urls,
    model: {
      id: '1_7b',
      displayName: 'Test Bonsai',
      architecture: 'test',
      source: { repo: 'test/model', revision: 'a'.repeat(40), file: 'model.gguf', bytes: 1, sha256: '0'.repeat(64) },
      files,
      downloadBytes: payloads.reduce((sum, payload) => sum + payload.size, 0),
      contextLength: 2048,
      defaultContext: 1024,
      cpuFallback: true,
      largestTensorBytes: 1,
      requiredLimits: {},
      chatTemplate: { bytes: 1, sha256: '0'.repeat(64), markers: { think: false, toolCall: false, toolResponse: false } },
      hybridDimensions: {},
      nextNTensorCount: 0,
      runtimePolicy: { flashAttention: false, tokenEmbeddingOnWebGPU: true, requireSingleWebGPUGraph: true },
    },
  };
}

function downloader(runtime: BrowserEngineRuntime): DownloadShards {
  return (runtime as unknown as { downloadShards: DownloadShards }).downloadShards.bind(runtime);
}

describe('BrowserEngineRuntime per-shard retry', () => {
  it('deletes a network-failed partial and retries only that shard with a full GET', async () => {
    const verified = new Blob(['verified']);
    const retryPayload = new Blob(['retry-ok']);
    const { model, urls } = await fixture(verified, retryPayload);
    const cache = new FakeCacheManager();
    cache.payloads.set(urls[1], retryPayload);
    cache.failNetworkFor = urls[1];
    const cached = { blobs: [verified, null], cachedBytes: verified.size };
    const download = downloader(new BrowserEngineRuntime());

    const failed = download(
      'request-1',
      { cacheManager: cache } as unknown as Wllama,
      model,
      urls,
      cached,
      new AbortController().signal,
      () => undefined,
    );

    await expect(failed).rejects.toSatisfy((error: unknown) => {
      if (!(error instanceof EngineRuntimeError) || !isShardDownloadFailureDetails(error.details)) return false;
      expect(error.code).toBe('SHARD_DOWNLOAD_FAILED');
      expect(error.details).toMatchObject({
        failure: 'network',
        shardIndex: 1,
        shardCount: 2,
        shardPath: 'shard-1.gguf',
        retryFromByteZero: true,
        partialDeleted: true,
      });
      return true;
    });
    expect(cache.deletes).toEqual([urls[1]]);
    expect(cache.blobs.has(urls[1])).toBe(false);
    expect(cached.blobs).toEqual([verified, null]);

    cache.failNetworkFor = null;
    const blobs = await download(
      'request-2',
      { cacheManager: cache } as unknown as Wllama,
      model,
      urls,
      cached,
      new AbortController().signal,
      () => undefined,
    );

    expect(blobs).toEqual([verified, retryPayload]);
    expect(cache.downloads.map(({ url }) => url)).toEqual([urls[1], urls[1]]);
    expect(cache.downloads.every(({ headers }) => headers?.Range === undefined)).toBe(true);
  });

  it('deletes a same-size shard that fails SHA-256 verification', async () => {
    const expected = new Blob(['good']);
    const corrupted = new Blob(['evil']);
    const { model, urls } = await fixture(new Blob(['cached']), expected);
    const cache = new FakeCacheManager();
    cache.payloads.set(urls[1], corrupted);
    const cached = { blobs: [new Blob(['cached']), null], cachedBytes: 6 };

    await expect(downloader(new BrowserEngineRuntime())(
      'request-hash',
      { cacheManager: cache } as unknown as Wllama,
      model,
      urls,
      cached,
      new AbortController().signal,
      () => undefined,
    )).rejects.toSatisfy((error: unknown) => (
      error instanceof EngineRuntimeError
      && error.code === 'SHARD_HASH_MISMATCH'
      && isShardDownloadFailureDetails(error.details)
      && error.details.failure === 'verification'
      && error.details.retryFromByteZero
    ));
    expect(cache.deletes).toEqual([urls[1]]);
    expect(cache.blobs.has(urls[1])).toBe(false);
  });

  it('stops dequeue after a network failure, lets in-flight siblings finish, and retries only 0 + 3', async () => {
    const payloads = [
      new Blob(['retry-zero']),
      new Blob(['sibling-one']),
      new Blob(['sibling-two']),
      new Blob(['queued-three']),
    ];
    const { model, urls } = await fixtureShards(payloads);
    const cache = new FakeCacheManager();
    urls.forEach((url, index) => cache.payloads.set(url, payloads[index] as Blob));
    cache.failNetworkFor = urls[0] ?? null;
    const cached = { blobs: payloads.map(() => null), cachedBytes: 0 };
    const progress: number[] = [];
    const download = downloader(new BrowserEngineRuntime());

    await expect(download(
      'request-parallel-failure',
      { cacheManager: cache } as unknown as Wllama,
      model,
      urls,
      cached,
      new AbortController().signal,
      (event) => {
        if (event.event === 'progress') progress.push(event.loadedBytes);
      },
    )).rejects.toMatchObject({ code: 'SHARD_DOWNLOAD_FAILED' });

    expect(cache.downloads.map(({ url }) => url)).toEqual(urls.slice(0, 3));
    expect(cache.downloads.some(({ url }) => url === urls[3])).toBe(false);
    expect(cache.deletes).toEqual([urls[0]]);
    expect(cached.blobs).toEqual([null, payloads[1], payloads[2], null]);
    expect(progress[progress.length - 1]).toBe(payloads[1]!.size + payloads[2]!.size);

    cache.failNetworkFor = null;
    const retryStart = cache.downloads.length;
    const blobs = await download(
      'request-parallel-retry',
      { cacheManager: cache } as unknown as Wllama,
      model,
      urls,
      cached,
      new AbortController().signal,
      () => undefined,
    );

    expect(blobs).toEqual(payloads);
    expect(cache.downloads.slice(retryStart).map(({ url }) => url)).toEqual([urls[0], urls[3]]);
    expect(cache.downloads.slice(retryStart).every(({ headers }) => headers?.Range === undefined)).toBe(true);
  });

  it('aborts all in-flight shards, cleans each partial, and never dequeues the fourth shard', async () => {
    const expected = [
      new Blob(['complete-zero']),
      new Blob(['complete-one']),
      new Blob(['complete-two']),
      new Blob(['complete-three']),
    ];
    const partials = [new Blob(['p0']), new Blob(['p1']), new Blob(['p2'])];
    const { model, urls } = await fixtureShards(expected);
    const cache = new FakeCacheManager();
    urls.slice(0, 3).forEach((url, index) => {
      cache.payloads.set(url, partials[index] as Blob);
      cache.waitForAbortFor.add(url);
    });
    const cached = { blobs: expected.map(() => null), cachedBytes: 0 };
    const progress: number[] = [];
    const controller = new AbortController();
    const failure = downloader(new BrowserEngineRuntime())(
      'request-abort',
      { cacheManager: cache } as unknown as Wllama,
      model,
      urls,
      cached,
      controller.signal,
      (event) => {
        if (event.event === 'progress') progress.push(event.loadedBytes);
      },
    );

    await cache.waitForDownloads(3);
    controller.abort();

    await expect(failure).rejects.toSatisfy((error: unknown) => (
      error instanceof EngineRuntimeError
      && error.code === 'SHARD_DOWNLOAD_ABORTED'
      && isShardDownloadFailureDetails(error.details)
      && error.details.failure === 'abort'
      && error.details.shardIndex < 3
      && error.details.shardCount === 4
      && error.details.shardPath === `shard-${error.details.shardIndex}.gguf`
      && error.details.partialDeleted
    ));
    expect(cache.downloads.map(({ url }) => url)).toEqual(urls.slice(0, 3));
    expect(cache.downloads.some(({ url }) => url === urls[3])).toBe(false);
    expect([...cache.deletes].sort()).toEqual([...urls.slice(0, 3)].sort());
    expect(urls.slice(0, 3).every((url) => !cache.blobs.has(url))).toBe(true);
    expect(cached.blobs).toEqual([null, null, null, null]);
    expect(progress[progress.length - 1]).toBe(0);
  });
});

describe('resolveLoadTuning', () => {
  const benchmarkDefaults = {
    flashMode: 'off',
    kvCacheType: 'f16',
    wasmFlavor: 'auto',
  } as const;

  function tuningError(
    input: unknown,
    backend: 'auto' | 'webgpu' | 'wasm' = 'webgpu',
  ): EngineRuntimeError {
    try {
      resolveLoadTuning(input, backend);
    } catch (error) {
      if (error instanceof EngineRuntimeError) return error;
      throw error;
    }
    throw new Error('Expected benchmark tuning to be rejected.');
  }

  it('uses immutable release defaults when benchmark tuning is omitted', () => {
    expect(resolveLoadTuning(undefined, 'auto')).toEqual({
      scope: 'release-defaults',
      nBatch: null,
      nUbatch: null,
      flashMode: 'off',
      kvCacheType: 'f16',
      wasmFlavor: 'auto',
    });
    expect(resolveLoadTuning(undefined, 'wasm')).toEqual({
      scope: 'release-defaults',
      nBatch: null,
      nUbatch: null,
      flashMode: 'off',
      kvCacheType: 'f16',
      wasmFlavor: 'auto',
    });
  });

  it('accepts an explicit benchmark batch shape without changing release defaults', () => {
    expect(resolveLoadTuning({
      ...benchmarkDefaults,
      nBatch: 512,
      nUbatch: 128,
    }, 'auto')).toEqual({
      scope: 'benchmark',
      nBatch: 512,
      nUbatch: 128,
      flashMode: 'off',
      kvCacheType: 'f16',
      wasmFlavor: 'auto',
    });
  });

  it('rejects a microbatch larger than the requested batch', () => {
    expect(tuningError({
      ...benchmarkDefaults,
      nBatch: 128,
      nUbatch: 256,
    })).toMatchObject({
      code: 'INVALID_BENCHMARK_TUNING',
      message: 'Benchmark n_ubatch must not exceed n_batch.',
    });
  });

  it('accepts Flash Attention and quantized KV only on the explicit WebGPU backend', () => {
    expect(resolveLoadTuning({
      ...benchmarkDefaults,
      flashMode: 'auto',
    }, 'webgpu')).toMatchObject({
      scope: 'benchmark',
      flashMode: 'auto',
      kvCacheType: 'f16',
      wasmFlavor: 'auto',
    });
    for (const kvCacheType of ['q8_0', 'q4_0'] as const) {
      expect(resolveLoadTuning({
        ...benchmarkDefaults,
        flashMode: 'auto',
        kvCacheType,
      }, 'webgpu')).toMatchObject({
        scope: 'benchmark',
        flashMode: 'auto',
        kvCacheType,
        wasmFlavor: 'auto',
      });
    }

    expect(tuningError({
      ...benchmarkDefaults,
      flashMode: 'auto',
    }, 'auto')).toMatchObject({
      code: 'INVALID_BENCHMARK_TUNING',
    });
    expect(tuningError({
      ...benchmarkDefaults,
      flashMode: 'auto',
      kvCacheType: 'q8_0',
    }, 'wasm')).toMatchObject({
      code: 'INVALID_BENCHMARK_TUNING',
    });
  });

  it('rejects malformed and incomplete benchmark tuning objects', () => {
    for (const input of [null, [], 'off', 42]) {
      expect(tuningError(input)).toMatchObject({ code: 'INVALID_BENCHMARK_TUNING' });
    }
    for (const input of [
      { kvCacheType: 'f16' },
      { flashMode: 'off' },
      { flashMode: 'off', kvCacheType: 'f16' },
      { ...benchmarkDefaults, flashMode: 'enabled' },
      { ...benchmarkDefaults, kvCacheType: 'q5_0' },
      { ...benchmarkDefaults, wasmFlavor: 'future' },
    ]) {
      expect(tuningError(input)).toMatchObject({ code: 'INVALID_BENCHMARK_TUNING' });
    }
  });

  it('rejects quantized V cache unless Flash Attention is in auto mode', () => {
    for (const kvCacheType of ['q8_0', 'q4_0'] as const) {
      expect(tuningError({ ...benchmarkDefaults, kvCacheType })).toMatchObject({
        code: 'INVALID_BENCHMARK_TUNING',
        message: expect.stringMatching(/requires Flash Attention/i),
      });
    }
  });

  it('rejects invalid and overflowing signed glue batch values', () => {
    for (const nBatch of [0, -1, 1.5, Number.NaN, Number.POSITIVE_INFINITY, 2_147_483_648]) {
      expect(tuningError({ ...benchmarkDefaults, nBatch })).toMatchObject({
        code: 'INVALID_BATCH_SIZE',
      });
    }
    for (const nUbatch of [0, -1, 1.5, Number.NaN, Number.POSITIVE_INFINITY, Number.MAX_SAFE_INTEGER]) {
      expect(tuningError({ ...benchmarkDefaults, nUbatch })).toMatchObject({
        code: 'INVALID_UBATCH_SIZE',
      });
    }
  });
});

describe('resolveRuntimeBatchShape', () => {
  const releaseDefaults = resolveLoadTuning(undefined, 'auto');

  it('caps only the compat release path to short Safari-safe command buffers', () => {
    expect(resolveRuntimeBatchShape(releaseDefaults, 'compat')).toEqual({
      nBatch: 32,
      nUbatch: 16,
    });
    expect(resolveRuntimeBatchShape(releaseDefaults, 'jspi')).toEqual({
      nBatch: undefined,
      nUbatch: undefined,
    });
  });

  it('preserves explicit tuning and never defaults microbatch above batch', () => {
    expect(resolveRuntimeBatchShape({ nBatch: 96, nUbatch: 24 }, 'compat')).toEqual({
      nBatch: 96,
      nUbatch: 24,
    });
    expect(resolveRuntimeBatchShape({ nBatch: 8, nUbatch: null }, 'compat')).toEqual({
      nBatch: 8,
      nUbatch: 8,
    });
  });
});

describe('selectWasmFlavor', () => {
  it('uses the browser-native flavor for auto', () => {
    expect(selectWasmFlavor('auto', false)).toBe('jspi');
    expect(selectWasmFlavor('auto', true)).toBe('compat');
  });

  it('forces compat even when JSPI is available', () => {
    expect(selectWasmFlavor('compat', false)).toBe('compat');
    expect(selectWasmFlavor('compat', true)).toBe('compat');
  });

  it('selects JSPI when the browser supports it', () => {
    expect(selectWasmFlavor('jspi', false)).toBe('jspi');
  });

  it('fails loud when JSPI is requested but unavailable', () => {
    expect(() => selectWasmFlavor('jspi', true)).toThrowError(
      expect.objectContaining({
        code: 'BENCHMARK_WASM_FLAVOR_UNAVAILABLE',
        message: expect.stringMatching(/requires the compatibility runtime/i),
      }),
    );
  });
});