UHD_Media_Tool/packages/tasks/src/openworkflow-engine.ts

import { randomUUID } from "node:crypto";
import { type Kysely, sql } from "kysely";
import { defineWorkflow, OpenWorkflow, Worker, type Workflow } from "openworkflow";
import { BackendPostgres } from "openworkflow/postgres";
import { TaskCanceledError } from "./errors.ts";
import { createLogger, type LogSink } from "./logger.ts";
import { engineVersion, parseEngineVersion, type TaskDefinition, type TaskRegistry } from "./registry.ts";
import type {
  EnqueueOptions,
  JsonValue,
  ListFilter,
  StartWorkerOptions,
  TaskContext,
  TaskEngine,
  TaskRunStatus,
  TaskRunView,
  TaskStepApi,
  WaitForFactOptions,
} from "./types.ts";

// biome-ignore lint/suspicious/noExplicitAny: 本包只用原生 SQL 访问自有表
type AnyDb = Kysely<any>;
type OwWorkflowRun = NonNullable<Awaited<ReturnType<BackendPostgres["getWorkflowRun"]>>>;
type OwStep = Parameters<Parameters<typeof defineWorkflow>[1]>[0]["step"];

export const ENGINE_SCHEMA = "openworkflow";
const DEFAULT_STEP_RETRY = { maximumAttempts: 3, initialInterval: "1s", backoffCoefficient: 2, maximumInterval: "60s" };

/** 写入引擎的 run 输入:业务输入 + 追踪元数据 */
interface Envelope {
  input: JsonValue;
  traceId: string;
  priority: number;
}

export interface OpenWorkflowEngineOptions {
  /** 引擎自己的 postgres.js 连接池使用(与 Kysely 同库) */
  databaseUrl: string;
  /** 业务库(Kysely/pg);tasks_* 表所在 */
  db: AnyDb;
  registry: TaskRegistry;
  /** 固定 "openworkflow";只允许测试覆盖 */
  schema?: string;
  /** 默认 true(幂等)。部署时由 migrateEngineSchema 预先迁移,多进程启动传 false */
  runMigrations?: boolean;
  logSink?: LogSink;
  /** handler 执行期间轮询取消标志的间隔,默认 1000ms */
  cancelPollMs?: number;
}

export const signalName = (runId: string, name: string) => `umt:${runId}:${name}`;

/**
 * 只跑引擎自己的迁移(openworkflow schema)后断开。部署时由 `--role migrate` 调用,
 * 之后 api / worker 进程以 runMigrations:false 连接,避免多进程同时迁移。
 */
export async function migrateEngineSchema(databaseUrl: string, schema = ENGINE_SCHEMA): Promise<void> {
  const backend = await BackendPostgres.connect(databaseUrl, { schema, runMigrations: true });
  await backend.stop();
}

export class OpenWorkflowEngine implements TaskEngine {
  private readonly ow: OpenWorkflow;
  private readonly workers: Worker[] = [];
  private readonly log;

  private constructor(
    private readonly backend: BackendPostgres,
    private readonly opts: OpenWorkflowEngineOptions,
  ) {
    this.ow = new OpenWorkflow({ backend });
    this.log = createLogger(opts.logSink, { component: "tasks.engine" });
  }

  static async connect(opts: OpenWorkflowEngineOptions): Promise<OpenWorkflowEngine> {
    const backend = await BackendPostgres.connect(opts.databaseUrl, {
      schema: opts.schema ?? ENGINE_SCHEMA,
      runMigrations: opts.runMigrations ?? true,
    });
    return new OpenWorkflowEngine(backend, opts);
  }

  async enqueue(taskName: string, input: unknown, opts: EnqueueOptions) {
    const def = this.opts.registry.latest(taskName);
    const parsed = def.inputSchema.parse(input) as JsonValue;
    const envelope: Envelope = { input: parsed, traceId: opts.traceId, priority: opts.priority ?? 0 };
    const handle = await this.ow.runWorkflow(
      { name: def.name, version: engineVersion(def.version) },
      envelope as unknown as JsonValue,
      { idempotencyKey: opts.idempotencyKey, ...(opts.availableAt ? { availableAt: opts.availableAt } : {}) },
    );
    const run = handle.workflowRun;
    const inserted = await sql`
      INSERT INTO tasks_run_meta (run_id, task_name, task_version, idempotency_key, trace_id, priority)
      VALUES (${run.id}, ${def.name}, ${def.version}, ${opts.idempotencyKey}, ${opts.traceId}, ${envelope.priority})
      ON CONFLICT (run_id) DO NOTHING`.execute(this.opts.db);
    // 同一幂等键返回已有 run 时 meta 行已存在
    return { runId: run.id, deduplicated: Number(inserted.numAffectedRows ?? 0n) === 0 };
  }

  async signal(runId: string, name: string, payload: JsonValue = null) {
    await sql`INSERT INTO tasks_fact (run_id, name, payload) VALUES (${runId}, ${name}, ${JSON.stringify(payload)}::jsonb)`.execute(
      this.opts.db,
    );
    const r = await this.ow.sendSignal({ signal: signalName(runId, name), data: payload });
    return { delivered: r.workflowRunIds.length };
  }

  async cancel(runId: string, reason?: string) {
    // 先写业务侧标志:运行中的 handler 通过它触发 abortSignal,步骤间检查也读它
    await sql`
      INSERT INTO tasks_run_meta (run_id, task_name, task_version, idempotency_key, cancel_requested_at, cancel_reason)
      VALUES (${runId}, '', 0, '', now(), ${reason ?? null})
      ON CONFLICT (run_id) DO UPDATE SET cancel_requested_at = COALESCE(tasks_run_meta.cancel_requested_at, now()),
        cancel_reason = COALESCE(tasks_run_meta.cancel_reason, EXCLUDED.cancel_reason)`.execute(this.opts.db);
    try {
      await this.ow.cancelWorkflowRun(runId);
      return true;
    } catch (err) {
      const run = await this.backend.getWorkflowRun({ workflowRunId: runId });
      if (run?.status === "canceled") return true;
      if (run) return false;
      throw err;
    }
  }

  async get(runId: string): Promise<TaskRunView | null> {
    const run = await this.backend.getWorkflowRun({ workflowRunId: runId });
    if (!run) return null;
    const metas = await this.loadMeta([run.id]);
    return toView(run, metas.get(run.id));
  }

  async list(filter: ListFilter = {}): Promise<TaskRunView[]> {
    const res = await this.backend.listWorkflowRuns({
      limit: filter.limit ?? 50,
      ...(filter.taskName ? { workflowName: filter.taskName } : {}),
      ...(filter.status ? { status: filter.status } : {}),
    });
    const metas = await this.loadMeta(res.data.map((r) => r.id));
    return res.data.map((r) => toView(r, metas.get(r.id)));
  }

  async startWorker(opts: StartWorkerOptions = {}) {
    const defs = this.opts.registry.select(opts.versions);
    const workflows = defs.map((d) => this.toWorkflow(d)) as Workflow<unknown, unknown, unknown>[];
    const worker = new Worker({ backend: this.backend, workflows, concurrency: opts.concurrency ?? 4 });
    await worker.start();
    this.workers.push(worker);
    this.log.info("worker started", { tasks: defs.map((d) => `${d.name}@${d.version}`) });
  }

  /** 等进行中的 handler 执行完再断开连接池 */
  async stop() {
    await Promise.all(this.workers.splice(0).map((w) => w.stop()));
    await this.backend.stop();
  }

  private async loadMeta(runIds: string[]) {
    const out = new Map<string, MetaRow>();
    if (runIds.length === 0) return out;
    const r = await sql<MetaRow>`SELECT * FROM tasks_run_meta WHERE run_id = ANY(${runIds}::text[])`.execute(
      this.opts.db,
    );
    for (const row of r.rows) out.set(row.run_id, row);
    return out;
  }

  private async cancelState(runId: string): Promise<{ at: Date | null; reason: string | null }> {
    const r = await sql<{ cancel_requested_at: Date | null; cancel_reason: string | null }>`
      SELECT cancel_requested_at, cancel_reason FROM tasks_run_meta WHERE run_id = ${runId}`.execute(this.opts.db);
    return { at: r.rows[0]?.cancel_requested_at ?? null, reason: r.rows[0]?.cancel_reason ?? null };
  }

  private toWorkflow(def: TaskDefinition) {
    const { db } = this.opts;
    const pollMs = this.opts.cancelPollMs ?? 1000;
    const stepRetry = { ...DEFAULT_STEP_RETRY, ...def.retryPolicy };
    return defineWorkflow(
      { name: def.name, version: engineVersion(def.version), retryPolicy: { maximumAttempts: 1 } },
      async ({ input, step, run }) => {
        const env = input as unknown as Envelope;
        const parsed = def.inputSchema.parse(env.input);
        const attemptToken = randomUUID();
        const logger = createLogger(this.opts.logSink, {
          task: def.name,
          version: def.version,
          runId: run.id,
          traceId: env.traceId,
        });
        const abort = new AbortController();
        let cancelReason: string | null = null;
        const poll = async () => {
          const s = await this.cancelState(run.id);
          if (s.at && !abort.signal.aborted) {
            cancelReason = s.reason;
            abort.abort(new TaskCanceledError(run.id, s.reason));
          }
          return abort.signal.aborted;
        };
        const timer = setInterval(() => {
          poll().catch((err) => logger.warn("cancel poll failed", { err }));
        }, pollMs);
        const throwIfCanceled = async () => {
          if (await poll()) throw new TaskCanceledError(run.id, cancelReason);
        };

        const stepApi = makeStepApi(step as OwStep, run.id, db, stepRetry, throwIfCanceled);
        const ctx: TaskContext<unknown> = {
          input: parsed,
          runId: run.id,
          taskName: def.name,
          version: def.version,
          traceId: env.traceId,
          priority: env.priority,
          attemptToken,
          abortSignal: abort.signal,
          logger,
          step: stepApi,
          throwIfCanceled,
        };
        try {
          await throwIfCanceled();
          const out = await def.handler(ctx);
          return (out ?? null) as JsonValue;
        } finally {
          clearInterval(timer);
        }
      },
    );
  }
}

function makeStepApi(
  step: OwStep,
  runId: string,
  db: AnyDb,
  stepRetry: typeof DEFAULT_STEP_RETRY,
  throwIfCanceled: () => Promise<void>,
): TaskStepApi {
  const api: TaskStepApi = {
    async run(name, fn, opts) {
      await throwIfCanceled();
      // biome-ignore lint/suspicious/noExplicitAny: 引擎的 DurationString 是模板字面量类型
      const retryPolicy = { ...stepRetry, ...opts?.retryPolicy } as any;
      return (await step.run({ name, retryPolicy }, async () => (await fn()) as never)) as never;
    },
    async sleep(name, duration) {
      await throwIfCanceled();
      // biome-ignore lint/suspicious/noExplicitAny: 同上
      await step.sleep(name, duration as any);
    },
    async waitForFact<T>(name: string, opts: WaitForFactOptions<T> = {}) {
      const check =
        opts.check ??
        (async () => {
          const r = await sql<{ payload: T }>`
            SELECT payload FROM tasks_fact WHERE run_id = ${runId} AND name = ${name} ORDER BY id DESC LIMIT 1`.execute(
            db,
          );
          return r.rows[0] ? r.rows[0].payload : undefined;
        });
      const maxRounds = opts.maxRounds ?? 100;
      for (let i = 0; ; i++) {
        const got = await api.run(`fact:${name}:check:${i}`, async () => {
          const v = await check();
          return v === undefined || v === null ? { found: false } : { found: true, value: v as unknown };
        });
        if (got.found) return (got as { value: T }).value;
        if (i >= maxRounds) return null;
        await throwIfCanceled();
        await step.waitForSignal({
          name: `fact:${name}:wait:${i}`,
          signal: signalName(runId, name),
          timeout: opts.timeout ?? "5m",
        });
      }
    },
  };
  return api;
}

interface MetaRow {
  run_id: string;
  task_name: string;
  task_version: number;
  idempotency_key: string;
  trace_id: string | null;
  priority: number;
  cancel_requested_at: Date | null;
  cancel_reason: string | null;
}

function toStatus(run: OwWorkflowRun): TaskRunStatus {
  switch (run.status) {
    case "succeeded":
    case "completed":
      return "completed";
    case "sleeping":
      return "waiting";
    case "running":
      // 等信号 / sleep 期间 status 仍是 running,但没有 worker 持有
      return run.workerId ? "running" : "waiting";
    default:
      return run.status;
  }
}

function toView(run: OwWorkflowRun, meta: MetaRow | undefined): TaskRunView {
  const env = (run.input ?? {}) as Partial<Envelope>;
  return {
    runId: run.id,
    taskName: run.workflowName,
    version: parseEngineVersion(run.version),
    status: toStatus(run),
    attempts: run.attempts,
    idempotencyKey: run.idempotencyKey,
    traceId: meta?.trace_id ?? env.traceId ?? null,
    priority: meta?.priority ?? env.priority ?? 0,
    input: env.input ?? null,
    output: run.output,
    error: run.error ? { name: run.error.name, message: run.error.message } : null,
    cancelRequestedAt: meta?.cancel_requested_at ?? null,
    cancelReason: meta?.cancel_reason ?? null,
    availableAt: run.availableAt,
    createdAt: run.createdAt,
    startedAt: run.startedAt,
    finishedAt: run.finishedAt,
  };
}