src/engine/worker.ts

import { EngineRuntimeError } from './errors';
import type {
  EngineErrorPayload,
  EngineEvent,
  EngineMethod,
  EngineRequest,
  EngineResponse,
  EngineResponseMap,
  EngineWorkerMessage,
} from './protocol';
import { BrowserEngineRuntime } from './runtime';

interface WorkerScope {
  onmessage: ((event: MessageEvent<EngineRequest>) => void) | null;
  postMessage(message: EngineWorkerMessage): void;
}

interface ActiveOperation {
  controller: AbortController;
  promise: Promise<unknown>;
}

const scope = globalThis as unknown as WorkerScope;
const runtime = new BrowserEngineRuntime();
const active = new Map<string, ActiveOperation>();
let quiescing = false;

function postEvent(event: EngineEvent): void {
  scope.postMessage(event);
}

function postSuccess<Method extends EngineMethod>(
  request: { requestId: string; method: Method },
  result: EngineResponseMap[Method],
): void {
  scope.postMessage({
    type: 'response',
    requestId: request.requestId,
    method: request.method,
    ok: true,
    result,
  } as EngineResponse);
}

function errorPayload(error: unknown): EngineErrorPayload {
  if (error instanceof EngineRuntimeError) {
    return {
      code: error.code,
      message: error.message,
      ...(error.details === undefined ? {} : { details: error.details }),
    };
  }
  if (error instanceof DOMException && error.name === 'AbortError') {
    return { code: 'ABORTED', message: error.message };
  }
  if (error instanceof Error) {
    return {
      code: error.name === 'AbortError' ? 'ABORTED' : 'ENGINE_ERROR',
      message: error.message,
    };
  }
  return { code: 'ENGINE_ERROR', message: String(error) };
}

function postFailure(request: { requestId: string; method: EngineMethod }, error: unknown): void {
  scope.postMessage({
    type: 'response',
    requestId: request.requestId,
    method: request.method,
    ok: false,
    error: errorPayload(error),
  } as EngineResponse);
}

async function runExclusive<Result>(
  requestId: string,
  operation: (signal: AbortSignal) => Promise<Result>,
): Promise<Result> {
  if (quiescing || active.size > 0) {
    throw new EngineRuntimeError(
      'ENGINE_BUSY',
      quiescing
        ? 'The engine is unloading or clearing storage.'
        : `Another mutable engine operation is active (${[...active.keys()].join(', ')}).`,
    );
  }
  const controller = new AbortController();
  const promise = operation(controller.signal);
  active.set(requestId, { controller, promise });
  try {
    return await promise;
  } finally {
    active.delete(requestId);
  }
}

async function stopActiveOperations(): Promise<void> {
  const operations = [...active.values()];
  for (const operation of operations) {
    operation.controller.abort();
  }
  await Promise.allSettled(operations.map((operation) => operation.promise));
}

async function runQuiesced<Result>(operation: () => Promise<Result>): Promise<Result> {
  if (quiescing) {
    throw new EngineRuntimeError('ENGINE_BUSY', 'The engine is already unloading or clearing storage.');
  }
  quiescing = true;
  try {
    await stopActiveOperations();
    return await operation();
  } finally {
    quiescing = false;
  }
}

async function handleRequest(request: EngineRequest): Promise<void> {
  try {
    switch (request.method) {
      case 'capabilities':
        postSuccess(request, await runtime.capabilities());
        return;
      case 'loadModel':
        postSuccess(
          request,
          await runExclusive(request.requestId, (signal) => (
            runtime.loadModel(request.requestId, request.params, signal, postEvent)
          )),
        );
        return;
      case 'generate':
        postSuccess(
          request,
          await runExclusive(request.requestId, (signal) => (
            runtime.generate(request.requestId, request.params, signal, postEvent)
          )),
        );
        return;
      case 'scoreSequence':
        postSuccess(
          request,
          await runExclusive(request.requestId, (signal) => (
            runtime.scoreSequence(request.params, signal)
          )),
        );
        return;
      case 'abort': {
        const operation = active.get(request.params.targetRequestId);
        const aborted = operation !== undefined && !operation.controller.signal.aborted;
        operation?.controller.abort();
        postSuccess(request, { targetRequestId: request.params.targetRequestId, aborted });
        return;
      }
      case 'unload':
        postSuccess(request, await runQuiesced(() => runtime.unload()));
        return;
      case 'backendReport':
        postSuccess(request, await runtime.backendReport());
        return;
      case 'storageEstimate':
        postSuccess(request, await runtime.storageEstimate());
        return;
      case 'storagePersist':
        postSuccess(request, await runtime.storagePersist());
        return;
      case 'storageClear':
        postSuccess(request, await runQuiesced(() => runtime.storageClear()));
        return;
    }
  } catch (error) {
    postFailure(request, error);
  }
}

scope.onmessage = (event) => {
  const request = event.data;
  if (request?.type === 'request' && typeof request.requestId === 'string') {
    void handleRequest(request);
  }
};