claude-code-viewer/e2e/utils/TaskExecutor.ts

import { ulid } from "ulid";

export type Task = {
  key: string;
  execute: () => Promise<void>;
};

type TaskStatus = {
  id: string;
  task: Task;
} & (
  | {
      status: "pending" | "completed" | "failed";
    }
  | {
      status: "running";
      promise: Promise<void>;
    }
);

type Options = {
  maxConcurrency: number;
};

export class TaskExecutor {
  private taskStatuses: TaskStatus[] = [];
  private executionPromise?: {
    resolve: () => void;
    reject: (reason?: unknown) => void;
    promise: Promise<void>;
  };
  private options: Options;

  constructor(options?: Partial<Options>) {
    this.options = {
      maxConcurrency: 10,
      ...options,
    };
  }

  private setExecutionPromise() {
    let resolveExecution: (() => void) | undefined;
    let rejectExecution: ((reason?: unknown) => void) | undefined;

    const promise = new Promise<void>((resolve, reject) => {
      resolveExecution = resolve;
      rejectExecution = reject;
    });

    if (resolveExecution === undefined || rejectExecution === undefined) {
      throw new Error("Illegal state: Promise not created");
    }

    this.executionPromise = {
      resolve: resolveExecution,
      reject: rejectExecution,
      promise,
    };
  }

  public setTasks(tasks: Task[]) {
    const newTaskStatuses: TaskStatus[] = tasks.map((task) => ({
      id: `${task.key}-${ulid()}`,
      status: "pending",
      task,
    }));

    this.taskStatuses.push(...newTaskStatuses);
  }

  private get pendingTasks() {
    return this.taskStatuses.filter((task) => task.status === "pending");
  }

  private get runningTasks() {
    return this.taskStatuses.filter((task) => task.status === "running");
  }

  private updateStatus(id: string, status: TaskStatus) {
    const found = this.taskStatuses.find((task) => task.id === id);

    if (!found) {
      throw new Error(`Task not found: ${id}`);
    }

    Object.assign(found, status);
  }

  public async execute() {
    this.setExecutionPromise();
    this.refresh();
    await this.executionPromise?.promise;
  }

  private refresh() {
    if (this.runningTasks.length === 0 && this.pendingTasks.length === 0) {
      this.executionPromise?.resolve();
      console.log("execution completed.");
      return;
    }

    const remainingTaskCount =
      this.options.maxConcurrency - this.runningTasks.length;

    if (remainingTaskCount <= 0) {
      return;
    }

    for (const task of this.pendingTasks.slice(0, remainingTaskCount)) {
      this.updateStatus(task.id, {
        id: task.id,
        status: "running",
        task: task.task,
        promise: (async () => {
          try {
            await task.task.execute();

            this.updateStatus(task.id, {
              id: task.id,
              status: "completed",
              task: task.task,
            });
          } catch (error) {
            console.error(error);

            this.updateStatus(task.id, {
              id: task.id,
              status: "failed",
              task: task.task,
            });
          } finally {
            this.refresh();
          }
        })(),
      });
    }
  }
}