src/engine/runtime-device-lost.test.ts

import { describe, expect, it, vi } from 'vitest';
import { NativeLogCollector } from './native-log';
import { BrowserEngineRuntime } from './runtime';

interface RuntimeInternals {
  wllama: { exit: () => Promise<void> } | null;
  loaded: {
    backend: 'webgpu' | 'wasm';
    model: { id: string; cpuFallback: boolean };
  } | null;
  nativeLog: NativeLogCollector;
  raceWebGpuDeviceLoss<Result>(
    stage: 'load' | 'generate' | 'backend-report',
    backend: 'webgpu' | 'wasm',
    model: { id: string; cpuFallback: boolean },
    operation: Promise<Result>,
  ): Promise<Result>;
}

describe('BrowserEngineRuntime WebGPU device loss', () => {
  it('invalidates the model and Wllama instance before returning a reload steer', async () => {
    const runtime = new BrowserEngineRuntime();
    const internals = runtime as unknown as RuntimeInternals;
    const exit = vi.fn<() => Promise<void>>().mockResolvedValue(undefined);
    internals.wllama = { exit };
    internals.loaded = {
      backend: 'webgpu',
      model: { id: '1_7b', cpuFallback: true },
    };
    internals.nativeLog.append([
      'ggml_webgpu: Device lost! Reason: 2, Message: GPU process crashed',
    ]);

    await expect(runtime.backendReport()).rejects.toMatchObject({
      name: 'EngineRuntimeError',
      code: 'WEBGPU_DEVICE_LOST',
      details: {
        recoverable: true,
        nextAction: 'reload-model',
        stage: 'backend-report',
        modelId: '1_7b',
        cpuFallbackAvailable: true,
        signal: { reason: 2, message: 'GPU process crashed' },
      },
    });
    expect(exit).toHaveBeenCalledOnce();
    expect(internals.loaded).toBeNull();
    expect(internals.wllama).toBeNull();
  });

  it('does not invalidate a CPU-WASM model for a stale WebGPU log signal', async () => {
    const runtime = new BrowserEngineRuntime();
    const internals = runtime as unknown as RuntimeInternals;
    const exit = vi.fn<() => Promise<void>>().mockResolvedValue(undefined);
    internals.wllama = { exit };
    internals.loaded = {
      backend: 'wasm',
      model: { id: '1_7b', cpuFallback: true },
    };
    internals.nativeLog.append([
      'ggml_webgpu: Device lost! Reason: 1, Message: unknown',
    ]);

    await expect(runtime.backendReport()).resolves.toEqual({
      backends: [],
      nGraphSplits: null,
      opsOnCpu: 0,
      layersGpu: null,
      flashAttention: null,
      cacheTypeK: null,
      cacheTypeV: null,
      webgpuKvBufferBytes: null,
      webgpuTrace: [],
    });
    expect(exit).not.toHaveBeenCalled();
    expect(internals.loaded).not.toBeNull();
  });

  it('interrupts a hung native operation as soon as the log callback signals device loss', async () => {
    const runtime = new BrowserEngineRuntime();
    const internals = runtime as unknown as RuntimeInternals;
    const exit = vi.fn<() => Promise<void>>().mockResolvedValue(undefined);
    const model = { id: '1_7b', cpuFallback: true };
    internals.wllama = { exit };
    internals.loaded = { backend: 'webgpu', model };
    const hung = new Promise<never>(() => undefined);
    const operation = internals.raceWebGpuDeviceLoss('generate', 'webgpu', model, hung);

    internals.nativeLog.append([
      'ggml_webgpu: Device lost! Reason: 4, Message: device removed',
    ]);

    await expect(operation).rejects.toMatchObject({
      code: 'WEBGPU_DEVICE_LOST',
      details: { stage: 'generate', nextAction: 'reload-model' },
    });
    expect(exit).toHaveBeenCalledOnce();
    expect(internals.loaded).toBeNull();
    expect(internals.wllama).toBeNull();
  });
});