src/engine/native-log.test.ts

import { describe, expect, it } from 'vitest';
import {
  NativeLogCollector,
  parseNativeDeviceLostSignal,
  parseNativeLog,
  parseNativeModelLoadProgress,
} from './native-log';

describe('parseNativeLog', () => {
  it('extracts WebGPU placement without treating bookkeeping buffers as CPU ops', () => {
    const report = parseNativeLog(`
llama_prepare_model_devices: using device WebGPU (WebGPU) (unknown id)
load_tensors: offloaded 65/65 layers to GPU
load_tensors: CPU model buffer size = 72.00 MiB
load_tensors: WebGPU model buffer size = 3616.77 MiB
llama_context: flash_attn = disabled
llama_kv_cache: WebGPU KV buffer size = 56.00 MiB
llama_kv_cache: size = 56.00 MiB, K (f16): 28.00 MiB, V (f16): 28.00 MiB
sched_reserve: WebGPU compute buffer size = 301.00 MiB
sched_reserve: CPU compute buffer size = 13.01 MiB
llama_context: CPU output buffer size = 0.95 MiB
sched_reserve: graph splits = 1
`);

    const modelBufferBytes = Math.round(72 * 1024 ** 2) + Math.round(3616.77 * 1024 ** 2);
    const computeBufferBytes = Math.round(301 * 1024 ** 2) + Math.round(13.01 * 1024 ** 2);
    const outputBufferBytes = Math.round(0.95 * 1024 ** 2);

    expect(report).toEqual({
      backends: ['WebGPU', 'CPU'],
      nGraphSplits: 1,
      opsOnCpu: 0,
      layersGpu: { offloaded: 65, total: 65 },
      flashAttention: false,
      cacheTypeK: 'f16',
      cacheTypeV: 'f16',
      webgpuKvBufferBytes: 56 * 1024 ** 2,
      modelBufferBytes,
      computeBufferBytes,
      outputBufferBytes,
      allocatedBufferBytes: modelBufferBytes
        + computeBufferBytes
        + outputBufferBytes
        + (56 * 1024 ** 2),
      webgpuTrace: [],
    });
  });

  it('parses and clamps structured native model-load progress', () => {
    expect(parseNativeModelLoadProgress(
      '@@MODEL_LOAD_PROGRESS@@{"state":"loading","stages":["metadata","tensors"],"current":"tensors","value":0.375}',
    )).toEqual({
      state: 'loading',
      stages: ['metadata', 'tensors'],
      current: 'tensors',
      value: 0.375,
    });
    expect(parseNativeModelLoadProgress(
      '@@MODEL_LOAD_PROGRESS@@{"state":"loading","value":7}',
    )).toEqual({ state: 'loading', stages: [], current: null, value: 1 });
    expect(parseNativeModelLoadProgress(
      '@@MODEL_LOAD_PROGRESS@@{"state":"ready","value":0}',
    )).toEqual({ state: 'ready', stages: [], current: null, value: 1 });
    expect(parseNativeModelLoadProgress('@@MODEL_LOAD_PROGRESS@@not-json')).toBeNull();
  });

  it('delivers native model-load progress to active listeners only', () => {
    const collector = new NativeLogCollector();
    const observed: unknown[] = [];
    const unsubscribe = collector.onModelLoadProgress((progress) => observed.push(progress));

    collector.append([
      '@@MODEL_LOAD_PROGRESS@@{"state":"loading","stages":["tensors"],"current":"tensors","value":0.5}',
    ]);
    collector.append(['@@MODEL_LOAD_PROGRESS@@invalid']);
    unsubscribe();
    collector.append(['@@MODEL_LOAD_PROGRESS@@{"state":"ready"}']);

    expect(observed).toEqual([{
      state: 'loading',
      stages: ['tensors'],
      current: 'tensors',
      value: 0.5,
    }]);
  });

  it('counts explicit CPU graph nodes and keeps the worst graph split observation', () => {
    const report = parseNativeLog([
      'sched_reserve: graph splits = 1',
      'sched_reserve: graph splits = 3',
      'backend = CPU, n_nodes = 7',
      'ggml_backend_cpu_graph_compute(5 nodes)',
    ]);

    expect(report.nGraphSplits).toBe(3);
    expect(report.opsOnCpu).toBe(12);
  });

  it('reports unknown graph evidence as null', () => {
    expect(parseNativeLog('load_tensors: offloaded 29/29 layers to GPU').nGraphSplits).toBeNull();
  });

  it('reports only resolved Flash Attention state and lets later evidence win', () => {
    expect(parseNativeLog('llama_context: flash_attn = auto').flashAttention).toBeNull();
    expect(parseNativeLog([
      'llama_context: flash_attn = auto',
      'resolve_fused_ops: Flash Attention enabled',
    ]).flashAttention).toBe(true);
    expect(parseNativeLog([
      'llama_context: flash_attn = enabled',
      'resolve_fused_ops: Flash Attention not supported, set to disabled',
    ]).flashAttention).toBe(false);
  });

  it('retains cache tuning evidence and sums binary WebGPU KV buffer units', () => {
    const collector = new NativeLogCollector();
    collector.append(['llama_context: flash_attn = disabled']);
    collector.append(['llama_kv_cache: K (Q8_0): 1.00 MiB, V (Q4_0): 1.00 MiB']);
    collector.append(['llama_kv_cache: WebGPU KV buffer size = 512 KiB']);
    collector.append(['llama_kv_cache: WebGPU KV buffer size = 2.50 MiB']);
    collector.append(['llama_kv_cache: WebGPU KV buffer size = 1.00 GiB']);

    expect(collector.report()).toMatchObject({
      flashAttention: false,
      cacheTypeK: 'q8_0',
      cacheTypeV: 'q4_0',
      webgpuKvBufferBytes: (512 * 1024) + (2.5 * 1024 ** 2) + (1024 ** 3),
    });
  });

  it('parses the native llama.cpp WebGPU device-lost callback', () => {
    expect(parseNativeDeviceLostSignal(
      'ggml_webgpu: Device lost! Reason: 2, Message: GPU process crashed',
    )).toEqual({
      line: 'ggml_webgpu: Device lost! Reason: 2, Message: GPU process crashed',
      reason: 2,
      message: 'GPU process crashed',
    });
    expect(parseNativeDeviceLostSignal('Device lost while reading a document')).toBeNull();
  });

  it('retains device loss outside the backend-evidence ring until explicitly cleared', () => {
    const collector = new NativeLogCollector();
    collector.append(['ggml_webgpu: Device lost! Reason: 1, Message: unknown']);

    expect(collector.report()).toEqual({
      backends: [],
      nGraphSplits: null,
      opsOnCpu: 0,
      layersGpu: null,
      flashAttention: null,
      cacheTypeK: null,
      cacheTypeV: null,
      webgpuKvBufferBytes: null,
      webgpuTrace: [],
    });
    expect(collector.deviceLostSignal()).toMatchObject({ reason: 1, message: 'unknown' });

    collector.clear();
    expect(collector.deviceLostSignal()).toBeNull();
  });

  it('retains WebGPU trace markers in a separate report channel', () => {
    const collector = new NativeLogCollector();
    collector.append(['@@WEBGPU_TRACE@@completion_begin id=1']);
    collector.append(['@@WEBGPU_TRACE@@graph_end id=1 nodes=3735 dispatches=1940 submits=31']);
    collector.append(['unrelated native output']);

    expect(collector.report().webgpuTrace).toEqual([
      '@@WEBGPU_TRACE@@completion_begin id=1',
      '@@WEBGPU_TRACE@@graph_end id=1 nodes=3735 dispatches=1940 submits=31',
    ]);

    collector.clear();
    expect(collector.report().webgpuTrace).toEqual([]);
  });

  it('retains a bounded diagnostic tail independently of backend evidence', () => {
    const collector = new NativeLogCollector();
    for (let index = 0; index < 70; index += 1) {
      collector.append([`diagnostic ${index}`]);
    }

    expect(collector.recent()).toEqual(
      Array.from({ length: 16 }, (_, index) => `diagnostic ${index + 54}`),
    );
    expect(collector.recent(100)).toHaveLength(64);
    expect(collector.report().webgpuTrace).toEqual([]);

    collector.clear();
    expect(collector.recent()).toEqual([]);
  });

  it('bounds the trace independently with room for a complete 64-token run', () => {
    const collector = new NativeLogCollector();
    for (let index = 0; index <= 16_384; index += 1) {
      collector.append([`@@WEBGPU_TRACE@@step_end completion=1 step=${index}`]);
    }

    const trace = collector.report().webgpuTrace;
    expect(trace).toHaveLength(16_384);
    expect(trace[0]).toContain('step=1');
    expect(trace.at(-1)).toContain('step=16384');
  });

  it('notifies an active native-operation watcher without polling', async () => {
    const collector = new NativeLogCollector();
    const watch = collector.watchDeviceLost();

    collector.append(['ggml_webgpu: Device lost! Reason: 3, Message: reset']);

    await expect(watch.promise).resolves.toMatchObject({ reason: 3, message: 'reset' });
    watch.dispose();
  });
});