src/engine/client.test.ts

import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest';
import { BrowserEngineClient } from './client';
import type { EngineWorkerMessage } from './protocol';

class FakeWorker {
  static instances: FakeWorker[] = [];

  onmessage: ((event: MessageEvent<EngineWorkerMessage>) => void) | null = null;
  onerror: ((event: ErrorEvent) => void) | null = null;
  onmessageerror: (() => void) | null = null;
  readonly postMessage = vi.fn<(message: unknown) => void>();
  readonly terminate = vi.fn<() => void>();

  constructor(..._args: unknown[]) {
    FakeWorker.instances.push(this);
  }

  crash(message: string): void {
    this.onerror?.({ message } as ErrorEvent);
  }

  send(message: EngineWorkerMessage): void {
    this.onmessage?.({ data: message } as MessageEvent<EngineWorkerMessage>);
  }
}

function latestWorker(): FakeWorker {
  const worker = FakeWorker.instances.at(-1);
  if (!worker) throw new Error('Expected a fake worker instance.');
  return worker;
}

describe('BrowserEngineClient worker lifecycle', () => {
  beforeEach(() => {
    FakeWorker.instances = [];
    vi.stubGlobal('Worker', FakeWorker);
  });

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

  it('terminates and recreates the worker when model loading is aborted', async () => {
    const client = new BrowserEngineClient();
    const firstWorker = latestWorker();
    const controller = new AbortController();
    const loading = client.loadModel({
      manifestUrl: '/manifest/models.json',
      modelId: '1_7b',
      backend: 'webgpu',
    }, {
      requestId: 'load-request',
      signal: controller.signal,
    });

    firstWorker.send({
      type: 'event',
      requestId: 'load-request',
      event: 'progress',
      phase: 'load',
      loadedBytes: 0,
      totalBytes: 100,
      shardIndex: null,
      shardCount: 1,
      shardPath: null,
      stageProgress: 0,
      nativeStage: null,
      residentBytes: null,
      shards: [],
    });

    controller.abort();

    await expect(loading).rejects.toMatchObject({ name: 'AbortError' });
    expect(firstWorker.terminate).toHaveBeenCalledOnce();
    expect(FakeWorker.instances).toHaveLength(2);

    const replacement = latestWorker();
    firstWorker.crash('late event from terminated worker');
    expect(replacement.terminate).not.toHaveBeenCalled();
    const estimate = client.storageEstimate();
    const request = replacement.postMessage.mock.calls[0]?.[0] as { requestId?: string };
    expect(request.requestId).toBeTruthy();
    replacement.send({
      type: 'response',
      requestId: request.requestId!,
      method: 'storageEstimate',
      ok: true,
      result: { usageBytes: 100, quotaBytes: 1_000, persisted: false },
    });

    await expect(estimate).resolves.toEqual({
      usageBytes: 100,
      quotaBytes: 1_000,
      persisted: false,
    });
    client.close();
  });

  it('hard-restarts only after soft download-abort cleanup times out', async () => {
    vi.useFakeTimers();
    try {
      const client = new BrowserEngineClient();
      const worker = latestWorker();
      const controller = new AbortController();
      const loading = client.loadModel({
        manifestUrl: '/manifest/models.json',
        modelId: '1_7b',
        backend: 'webgpu',
      }, {
        requestId: 'stuck-download',
        signal: controller.signal,
      });
      const loadingFailure = expect(loading).rejects.toMatchObject({ name: 'AbortError' });
      worker.send({
        type: 'event',
        requestId: 'stuck-download',
        event: 'progress',
        phase: 'download',
        loadedBytes: 1,
        totalBytes: 100,
        shardIndex: 0,
        shardCount: 1,
        shardPath: 'model.gguf',
        stageProgress: null,
        nativeStage: null,
        residentBytes: null,
        shards: [{
          index: 0,
          path: 'model.gguf',
          loadedBytes: 1,
          verifiedBytes: 0,
          totalBytes: 100,
          state: 'downloading',
        }],
      });

      controller.abort();
      await vi.advanceTimersByTimeAsync(29_999);
      expect(worker.terminate).not.toHaveBeenCalled();
      await vi.advanceTimersByTimeAsync(1);

      await loadingFailure;
      expect(worker.terminate).toHaveBeenCalledOnce();
      expect(FakeWorker.instances).toHaveLength(2);
      client.close();
    } finally {
      vi.useRealTimers();
    }
  });

  it('lets the runtime clean partial shards when download is aborted', async () => {
    const client = new BrowserEngineClient();
    const worker = latestWorker();
    const controller = new AbortController();
    const loading = client.loadModel({
      manifestUrl: '/manifest/models.json',
      modelId: '1_7b',
      backend: 'webgpu',
    }, {
      requestId: 'download-request',
      signal: controller.signal,
    });
    worker.send({
      type: 'event',
      requestId: 'download-request',
      event: 'progress',
      phase: 'download',
      loadedBytes: 50,
      totalBytes: 100,
      shardIndex: 0,
      shardCount: 1,
      shardPath: 'model.gguf',
      stageProgress: null,
      nativeStage: null,
      residentBytes: null,
      shards: [{
        index: 0,
        path: 'model.gguf',
        loadedBytes: 50,
        verifiedBytes: 0,
        totalBytes: 100,
        state: 'downloading',
      }],
    });

    controller.abort();

    expect(worker.terminate).not.toHaveBeenCalled();
    expect(FakeWorker.instances).toHaveLength(1);
    const abortRequest = worker.postMessage.mock.calls
      .map((call) => call[0] as { requestId: string; method: string })
      .find((request) => request.method === 'abort');
    expect(abortRequest).toBeTruthy();
    worker.send({
      type: 'response',
      requestId: 'download-request',
      method: 'loadModel',
      ok: false,
      error: {
        code: 'SHARD_DOWNLOAD_ABORTED',
        message: 'Partial shard removed.',
        details: { retryFromByteZero: true },
      },
    });
    if (!abortRequest) throw new Error('Expected abort request.');
    worker.send({
      type: 'response',
      requestId: abortRequest.requestId,
      method: 'abort',
      ok: true,
      result: { targetRequestId: 'download-request', aborted: true },
    });

    await expect(loading).rejects.toMatchObject({ code: 'SHARD_DOWNLOAD_ABORTED' });
    expect(worker.terminate).not.toHaveBeenCalled();
    client.close();
  });

  it('keeps the loaded model when generation responds to a soft abort', async () => {
    const client = new BrowserEngineClient();
    const worker = latestWorker();
    const controller = new AbortController();
    const generation = client.generate({
      messages: [{ role: 'user', content: 'hello' }],
    }, {
      requestId: 'soft-generation',
      signal: controller.signal,
    });
    const generationFailure = expect(generation).rejects.toMatchObject({ code: 'ABORTED' });

    controller.abort();

    const abortRequest = worker.postMessage.mock.calls
      .map((call) => call[0] as { requestId: string; method: string })
      .find((request) => request.method === 'abort');
    expect(abortRequest).toBeTruthy();
    worker.send({
      type: 'response',
      requestId: 'soft-generation',
      method: 'generate',
      ok: false,
      error: { code: 'ABORTED', message: 'Generation stopped.' },
    });
    if (!abortRequest) throw new Error('Expected abort request.');
    worker.send({
      type: 'response',
      requestId: abortRequest.requestId,
      method: 'abort',
      ok: true,
      result: { targetRequestId: 'soft-generation', aborted: true },
    });

    await generationFailure;
    expect(worker.terminate).not.toHaveBeenCalled();
    expect(FakeWorker.instances).toHaveLength(1);
    client.close();
  });

  it('routes and softly aborts diagnostic sequence scoring', async () => {
    const client = new BrowserEngineClient();
    const worker = latestWorker();
    const controller = new AbortController();
    const scoring = client.scoreSequence({
      promptTokenIds: Array.from({ length: 38 }, (_, index) => index + 1),
      referenceTokenIds: Array.from({ length: 1_024 }, (_, index) => index + 1_000),
      topK: 5,
    }, {
      requestId: 'soft-score-sequence',
      signal: controller.signal,
    });
    const scoringFailure = expect(scoring).rejects.toMatchObject({ code: 'ABORTED' });

    expect(worker.postMessage).toHaveBeenCalledWith(expect.objectContaining({
      requestId: 'soft-score-sequence',
      method: 'scoreSequence',
    }));
    controller.abort();
    const abortRequest = worker.postMessage.mock.calls
      .map((call) => call[0] as { requestId: string; method: string })
      .find((request) => request.method === 'abort');
    expect(abortRequest).toBeTruthy();
    worker.send({
      type: 'response',
      requestId: 'soft-score-sequence',
      method: 'scoreSequence',
      ok: false,
      error: { code: 'ABORTED', message: 'Sequence scoring stopped.' },
    });
    if (!abortRequest) throw new Error('Expected abort request.');
    worker.send({
      type: 'response',
      requestId: abortRequest.requestId,
      method: 'abort',
      ok: true,
      result: { targetRequestId: 'soft-score-sequence', aborted: true },
    });

    await scoringFailure;
    expect(worker.terminate).not.toHaveBeenCalled();
    client.close();
  });

  it('invalidates the model when generation ignores abort beyond the grace period', async () => {
    vi.useFakeTimers();
    try {
      const client = new BrowserEngineClient();
      const worker = latestWorker();
      const controller = new AbortController();
      const generation = client.generate({
        messages: [{ role: 'user', content: 'hello' }],
      }, {
        requestId: 'hung-generation',
        signal: controller.signal,
      });
      const generationFailure = expect(generation).rejects.toMatchObject({
        code: 'ENGINE_WORKER_FAILED',
        message: expect.stringContaining('did not stop within 5 seconds'),
        details: {
          recoverable: true,
          nextAction: 'reload-model',
          cause: 'generation-abort-timeout',
          graceMs: 5_000,
        },
      });

      controller.abort();
      await vi.advanceTimersByTimeAsync(4_999);
      expect(worker.terminate).not.toHaveBeenCalled();
      await vi.advanceTimersByTimeAsync(1);

      await generationFailure;
      expect(worker.terminate).toHaveBeenCalledOnce();
      expect(FakeWorker.instances).toHaveLength(2);
      client.close();
    } finally {
      vi.useRealTimers();
    }
  });

  it('rejects pending work on a crash and recreates the worker for the next request', async () => {
    const client = new BrowserEngineClient();
    const firstWorker = latestWorker();
    const pending = client.storageEstimate();

    firstWorker.crash('synthetic worker crash');

    await expect(pending).rejects.toMatchObject({
      code: 'ENGINE_WORKER_FAILED',
      message: expect.stringContaining('synthetic worker crash'),
      details: { recoverable: true, nextAction: 'reload-model' },
    });
    expect(firstWorker.terminate).toHaveBeenCalledOnce();
    expect(FakeWorker.instances).toHaveLength(1);

    const recovered = client.storageEstimate();
    expect(FakeWorker.instances).toHaveLength(2);
    const replacement = latestWorker();
    const request = replacement.postMessage.mock.calls[0]?.[0] as { requestId?: string };
    replacement.send({
      type: 'response',
      requestId: request.requestId!,
      method: 'storageEstimate',
      ok: true,
      result: { usageBytes: null, quotaBytes: null, persisted: false },
    });

    await expect(recovered).resolves.toEqual({
      usageBytes: null,
      quotaBytes: null,
      persisted: false,
    });
    client.close();
  });

  it('surfaces WebGPU device loss and replaces the invalid worker before the next action', async () => {
    const client = new BrowserEngineClient();
    const firstWorker = latestWorker();
    const generation = client.generate({
      messages: [{ role: 'user', content: 'hello' }],
    }, { requestId: 'generation' });

    firstWorker.send({
      type: 'response',
      requestId: 'generation',
      method: 'generate',
      ok: false,
      error: {
        code: 'WEBGPU_DEVICE_LOST',
        message: 'The WebGPU device was lost during generate.',
        details: { recoverable: true, nextAction: 'reload-model' },
      },
    });

    await expect(generation).rejects.toMatchObject({
      name: 'EngineClientError',
      code: 'WEBGPU_DEVICE_LOST',
      details: { recoverable: true, nextAction: 'reload-model' },
    });
    expect(firstWorker.terminate).toHaveBeenCalledOnce();
    expect(FakeWorker.instances).toHaveLength(2);

    const replacement = latestWorker();
    const reload = client.loadModel({
      manifestUrl: '/manifest/models.json',
      modelId: '1_7b',
      backend: 'webgpu',
    }, { requestId: 'reload' });
    expect(replacement.postMessage).toHaveBeenCalledWith(expect.objectContaining({
      requestId: 'reload',
      method: 'loadModel',
    }));
    replacement.send({
      type: 'response',
      requestId: 'reload',
      method: 'loadModel',
      ok: false,
      error: { code: 'MODEL_GATE_REJECTED', message: 'synthetic stop' },
    });
    await expect(reload).rejects.toMatchObject({ code: 'MODEL_GATE_REJECTED' });
    client.close();
  });

  it('types a crash during generation as model-invalidating before recreating the worker', async () => {
    const client = new BrowserEngineClient();
    const firstWorker = latestWorker();
    const generation = client.generate({
      messages: [{ role: 'user', content: 'hello' }],
    });

    firstWorker.crash('GPU worker process exited');

    await expect(generation).rejects.toMatchObject({
      code: 'ENGINE_WORKER_FAILED',
      details: { recoverable: true, nextAction: 'reload-model' },
    });
    expect(firstWorker.terminate).toHaveBeenCalledOnce();

    const reload = client.loadModel({
      manifestUrl: '/manifest/models.json',
      modelId: '1_7b',
      backend: 'webgpu',
    }, { requestId: 'crash-reload' });
    const replacement = latestWorker();
    expect(replacement.postMessage).toHaveBeenCalledWith(expect.objectContaining({
      requestId: 'crash-reload',
      method: 'loadModel',
    }));
    replacement.send({
      type: 'response',
      requestId: 'crash-reload',
      method: 'loadModel',
      ok: false,
      error: { code: 'MODEL_GATE_REJECTED', message: 'synthetic stop' },
    });
    await expect(reload).rejects.toMatchObject({ code: 'MODEL_GATE_REJECTED' });
    client.close();
  });

  it('routes live tool-call progress independently from visible text tokens', async () => {
    const client = new BrowserEngineClient();
    const worker = latestWorker();
    const onToken = vi.fn();
    const onToolCallProgress = vi.fn();
    const generation = client.generate({
      messages: [{ role: 'user', content: 'Create an app.' }],
    }, {
      requestId: 'tool-progress',
      onToken,
      onToolCallProgress,
    });

    worker.send({
      type: 'event',
      requestId: 'tool-progress',
      event: 'tool-call',
      index: 0,
      id: 'call_artifact',
      name: 'html_artifact',
      argumentCharacters: 4_096,
    });

    expect(onToolCallProgress).toHaveBeenCalledWith(expect.objectContaining({
      name: 'html_artifact',
      argumentCharacters: 4_096,
    }));
    expect(onToken).not.toHaveBeenCalled();

    worker.send({
      type: 'response',
      requestId: 'tool-progress',
      method: 'generate',
      ok: true,
      result: {
        text: '',
        reasoningText: '',
        tokenIds: null,
        tokenTrace: null,
        finishReason: 'tool_calls',
        toolCalls: [{
          id: 'call_artifact',
          type: 'function',
          function: { name: 'html_artifact', arguments: '{"html":"<main />"}' },
        }],
        usage: null,
        timings: null,
      },
    });

    await expect(generation).resolves.toMatchObject({ finishReason: 'tool_calls' });
    client.close();
  });

  it('never recreates a worker after close', async () => {
    const client = new BrowserEngineClient();
    const worker = latestWorker();
    const pending = client.storageEstimate();

    client.close();
    client.restart();

    await expect(pending).rejects.toThrow('closed');
    await expect(client.storageEstimate()).rejects.toThrow('closed');
    expect(worker.terminate).toHaveBeenCalledOnce();
    expect(FakeWorker.instances).toHaveLength(1);
  });
});