src/tasks/AbstractModelTaskHandler.ts

import type {
  ClusterTaskHostContext,
  ClusterTaskKind,
  CompleteTaskBody,
  HostTaskMessage,
} from './types';

export abstract class AbstractModelTaskHandler {
  abstract readonly kind: ClusterTaskKind;

  constructor(protected readonly host: ClusterTaskHostContext) {}

  abstract process(task: HostTaskMessage): Promise<void>;

  protected async completeTask(taskId: string, body: CompleteTaskBody): Promise<void> {
    const res = await fetch(`${this.host.apiBase}/api/tasks/${taskId}/complete`, {
      method: 'POST',
      headers: {'Content-Type': 'application/json'},
      body: JSON.stringify({
        ...body,
        model: body.model ?? this.host.getSelectedModelId(),
      }),
    });

    if (!res.ok) {
      const json = (await res.json().catch(() => ({}))) as {error?: string};
      throw new Error(json.error ?? `Task complete failed (${res.status})`);
    }
  }

  protected taskError(error: string): CompleteTaskBody {
    return {status: 'error', kind: this.kind, result: {error}};
  }

  protected async withVideoFrame(
    task: HostTaskMessage,
    run: (frame: VideoFrame) => Promise<CompleteTaskBody>,
  ): Promise<void> {
    let frame: VideoFrame | null = null;

    try {
      frame = await AbstractModelTaskHandler.base64ToVideoFrame(
        task.image_base64,
        task.mime_type,
      );
      const body = await run(frame);
      await this.completeTask(task.id, body);
    } catch (error) {
      const message = error instanceof Error ? error.message : String(error);
      await this.completeTask(task.id, this.taskError(message));
      throw error;
    } finally {
      frame?.close();
    }
  }

  protected static async base64ToVideoFrame(
    base64: string,
    mimeType: string,
  ): Promise<VideoFrame> {
    const binary = atob(base64);
    const bytes = new Uint8Array(binary.length);
    for (let i = 0; i < binary.length; i++) {
      bytes[i] = binary.charCodeAt(i);
    }

    const blob = new Blob([bytes], {type: mimeType});
    const bitmap = await createImageBitmap(blob);
    const frame = new VideoFrame(bitmap, {timestamp: 0});
    bitmap.close();
    return frame;
  }
}