src/engine/client.ts

import type {
  EngineErrorPayload,
  EngineEvent,
  EngineMethod,
  EngineRequest,
  EngineRequestMap,
  EngineResponse,
  EngineResponseMap,
  EngineWorkerMessage,
  GenerateParams,
  LoadModelParams,
  ScoreSequenceParams,
} from './protocol';

type ProgressEvent = Extract<EngineEvent, { event: 'progress' }>;
type TokenEvent = Extract<EngineEvent, { event: 'token' }>;
type ToolCallEvent = Extract<EngineEvent, { event: 'tool-call' }>;
type GenerationEvent = Extract<EngineEvent, { event: 'generation' }>;

export interface EngineRequestOptions {
  requestId?: string;
  signal?: AbortSignal;
  onProgress?: (event: ProgressEvent) => void;
  onToken?: (text: string, reasoningDelta?: string) => void;
  onToolCallProgress?: (event: ToolCallEvent) => void;
  onGenerationProgress?: (event: GenerationEvent) => void;
}

interface PendingRequest {
  method: EngineMethod;
  resolve: (result: never) => void;
  reject: (error: unknown) => void;
  options: EngineRequestOptions;
  abortListener?: () => void;
  loadPhase?: ProgressEvent['phase'];
  abortFallbackTimer?: ReturnType<typeof setTimeout>;
}

const LOAD_ABORT_CLEANUP_TIMEOUT_MS = 30_000;
// Native generation may be blocked inside a proxy call that cannot observe AbortSignal.
const GENERATION_ABORT_GRACE_MS = 5_000;

export class EngineClientError extends Error {
  readonly code: string;
  readonly details?: unknown;

  constructor(payload: EngineErrorPayload) {
    super(payload.message);
    this.name = 'EngineClientError';
    this.code = payload.code;
    this.details = payload.details;
  }
}

function nextRequestId(): string {
  return typeof crypto.randomUUID === 'function'
    ? crypto.randomUUID()
    : `${Date.now()}-${Math.random().toString(16).slice(2)}`;
}

export class BrowserEngineClient {
  private worker: Worker | null;
  private readonly workerFactory: (() => Worker) | null;
  private readonly pending = new Map<string, PendingRequest>();
  private closed = false;

  constructor(worker?: Worker) {
    this.workerFactory = worker ? null : () => new Worker(new URL('./worker.ts', import.meta.url), {
      type: 'module',
      name: 'bonsai-browser-engine',
    });
    this.worker = worker ?? this.workerFactory?.() ?? null;
    if (!this.worker) throw new Error('Browser engine worker could not be created.');
    this.attachWorker(this.worker);
  }

  private attachWorker(worker: Worker): void {
    worker.onmessage = (event: MessageEvent<EngineWorkerMessage>) => {
      if (this.worker !== worker) return;
      this.handleMessage(event.data);
    };
    worker.onerror = (event) => {
      if (this.worker !== worker) return;
      this.failWorker(new Error(event.message || 'Browser engine worker failed.'));
    };
    worker.onmessageerror = () => {
      if (this.worker !== worker) return;
      this.failWorker(new Error('Browser engine worker sent an invalid message.'));
    };
  }

  private ensureWorker(): Worker {
    if (this.closed) throw new Error('Browser engine client is closed.');
    if (!this.worker) {
      if (!this.workerFactory) throw new Error('Browser engine worker failed and cannot be recreated.');
      this.worker = this.workerFactory();
      this.attachWorker(this.worker);
    }
    return this.worker;
  }

  request<Method extends EngineMethod>(
    method: Method,
    params: EngineRequestMap[Method],
    options: EngineRequestOptions = {},
  ): Promise<EngineResponseMap[Method]> {
    if (options.signal?.aborted) {
      return Promise.reject(new DOMException('The engine operation was aborted.', 'AbortError'));
    }
    let worker: Worker;
    try {
      worker = this.ensureWorker();
    } catch (error) {
      return Promise.reject(error);
    }
    const requestId = options.requestId ?? nextRequestId();
    if (this.pending.has(requestId)) {
      return Promise.reject(new Error(`Duplicate engine request id: ${requestId}`));
    }
    const request = { type: 'request', requestId, method, params } as EngineRequest;
    return new Promise<EngineResponseMap[Method]>((resolve, reject) => {
      const pending: PendingRequest = {
        method,
        resolve: resolve as (result: never) => void,
        reject,
        options,
      };
      if (options.signal) {
        pending.abortListener = () => {
          if (method === 'loadModel' && pending.loadPhase === 'load') {
            this.restart(new DOMException('Model loading was aborted.', 'AbortError'));
          } else {
            if (method === 'loadModel') {
              pending.abortFallbackTimer = setTimeout(() => {
                if (this.pending.has(requestId)) {
                  this.restart(new DOMException(
                    'Model download cleanup timed out after abort.',
                    'AbortError',
                  ));
                }
              }, LOAD_ABORT_CLEANUP_TIMEOUT_MS);
            } else if (method === 'generate' || method === 'scoreSequence') {
              const operationLabel = method === 'generate' ? 'Generation' : 'Sequence scoring';
              const cause = method === 'generate'
                ? 'generation-abort-timeout'
                : 'score-sequence-abort-timeout';
              pending.abortFallbackTimer = setTimeout(() => {
                if (this.pending.has(requestId)) {
                  this.restart(new EngineClientError({
                    code: 'ENGINE_WORKER_FAILED',
                    message: `${operationLabel} did not stop within 5 seconds. The browser engine worker was restarted and the loaded model was invalidated.`,
                    details: {
                      recoverable: true,
                      nextAction: 'reload-model',
                      cause,
                      graceMs: GENERATION_ABORT_GRACE_MS,
                    },
                  }));
                }
              }, GENERATION_ABORT_GRACE_MS);
            }
            void this.abort(requestId).catch(() => undefined);
          }
        };
        options.signal.addEventListener('abort', pending.abortListener, { once: true });
      }
      this.pending.set(requestId, pending);
      try {
        worker.postMessage(request);
      } catch (error) {
        this.failWorker(error);
      }
    });
  }

  capabilities() {
    return this.request('capabilities', {});
  }

  loadModel(params: LoadModelParams, options?: EngineRequestOptions) {
    return this.request('loadModel', params, options);
  }

  generate(params: GenerateParams, options?: EngineRequestOptions) {
    return this.request('generate', params, options);
  }

  scoreSequence(params: ScoreSequenceParams, options?: EngineRequestOptions) {
    return this.request('scoreSequence', params, options);
  }

  abort(targetRequestId: string) {
    return this.request('abort', { targetRequestId });
  }

  unload() {
    return this.request('unload', {});
  }

  backendReport() {
    return this.request('backendReport', {});
  }

  storageEstimate() {
    return this.request('storageEstimate', {});
  }

  storagePersist() {
    return this.request('storagePersist', {});
  }

  storageClear() {
    return this.request('storageClear', {});
  }

  close(): void {
    this.closed = true;
    this.worker?.terminate();
    this.worker = null;
    this.rejectAll(new Error('Browser engine client was closed.'));
  }

  restart(reason: unknown = new Error('Browser engine worker was restarted.')): void {
    if (this.closed) return;
    this.worker?.terminate();
    this.worker = null;
    this.rejectAll(reason);
    if (this.workerFactory) {
      this.worker = this.workerFactory();
      this.attachWorker(this.worker);
    }
  }

  private handleMessage(message: EngineWorkerMessage): void {
    const pending = this.pending.get(message.requestId);
    if (!pending) {
      return;
    }
    if (message.type === 'event') {
      if (message.event === 'progress') {
        pending.loadPhase = message.phase;
        if (
          pending.method === 'loadModel'
          && pending.options.signal?.aborted
          && message.phase === 'load'
        ) {
          this.restart(new DOMException('Model loading was aborted.', 'AbortError'));
          return;
        }
        pending.options.onProgress?.(message);
      } else if (message.event === 'token') {
        pending.options.onToken?.(message.text, message.reasoningDelta);
      } else if (message.event === 'tool-call') {
        pending.options.onToolCallProgress?.(message);
      } else {
        pending.options.onGenerationProgress?.(message);
      }
      return;
    }
    this.finishPending(message, pending);
  }

  private finishPending(response: EngineResponse, pending: PendingRequest): void {
    this.pending.delete(response.requestId);
    if (pending.abortListener && pending.options.signal) {
      pending.options.signal.removeEventListener('abort', pending.abortListener);
    }
    if (pending.abortFallbackTimer !== undefined) clearTimeout(pending.abortFallbackTimer);
    if (response.method !== pending.method) {
      pending.reject(new Error(
        `Engine response method mismatch: expected ${pending.method}, received ${response.method}.`,
      ));
    } else if (response.ok) {
      if (pending.options.signal?.aborted) {
        const abort = new DOMException('The engine operation was aborted.', 'AbortError');
        pending.reject(abort);
        if (pending.method === 'loadModel') this.restart(abort);
      } else {
        pending.resolve(response.result as never);
      }
    } else {
      const error = new EngineClientError(response.error);
      pending.reject(error);
      if (error.code === 'WEBGPU_DEVICE_LOST') {
        this.restart(error);
      }
    }
  }

  private rejectAll(error: unknown): void {
    for (const [requestId, pending] of this.pending) {
      this.pending.delete(requestId);
      if (pending.abortListener && pending.options.signal) {
        pending.options.signal.removeEventListener('abort', pending.abortListener);
      }
      if (pending.abortFallbackTimer !== undefined) clearTimeout(pending.abortFallbackTimer);
      pending.reject(error);
    }
  }

  private failWorker(error: unknown): void {
    const failure = error instanceof EngineClientError
      ? error
      : new EngineClientError({
        code: 'ENGINE_WORKER_FAILED',
        message: `Browser engine worker failed: ${error instanceof Error ? error.message : String(error)}`,
        details: { recoverable: true, nextAction: 'reload-model' },
      });
    this.worker?.terminate();
    this.worker = null;
    this.rejectAll(failure);
  }
}