UHD_Media_Tool/apps/server/src/modules/agent/capabilities.ts

import { CapabilityError, type CapabilityRegistry } from "@umt/agent";
import { capabilitySpec, type TaskStatus } from "@umt/contracts";
import { sql, type TaskRow } from "@umt/db";
import type { ServerContext } from "../types.ts";

// 任务类能力:task.inspect / task.retry / task.cancel。业务任务行(task 表)是时间线投影,
// 引擎只管推进;重试 = 按原输入重新入队(副作用由各处理器的操作账本去重),取消 = 写取消标志 + 引擎取消。

/** task.kind → 引擎任务名(agentRun 行对应 runAgent) */
export const engineTaskName = (kind: string) => (kind === "agentRun" ? "runAgent" : kind);

const toIso = (d: Date | null) => (d ? d.toISOString() : null);

export function taskToApi(t: TaskRow) {
  return {
    id: t.id,
    kind: t.kind,
    handlerVersion: t.handler_version,
    idempotencyKey: t.idempotency_key,
    status: t.status,
    priority: t.priority,
    input: t.input,
    output: t.output,
    subjectType: t.subject_type,
    subjectId: t.subject_id,
    parentTaskId: t.parent_task_id,
    engineRunId: t.engine_run_id,
    attempts: t.attempts,
    maxAttempts: t.max_attempts,
    nextRunAt: t.next_run_at.toISOString(),
    waitReason: t.wait_reason,
    cancelRequestedAt: toIso(t.cancel_requested_at),
    cancelReason: t.cancel_reason,
    lastError: t.last_error,
    startedAt: toIso(t.started_at),
    finishedAt: toIso(t.finished_at),
    triggeredBy: t.triggered_by,
    createdAt: t.created_at.toISOString(),
    updatedAt: t.updated_at.toISOString(),
  };
}

/** 业务时间线:任务自身 + 关联业务对象(计划/下载/整理记录/agent 运行)的当前状态 */
export async function taskTimeline(ctx: ServerContext, t: TaskRow) {
  const db = ctx.db;
  const out: { at: string; type: string; summary: string; data?: Record<string, unknown> }[] = [];
  out.push({ at: t.created_at.toISOString(), type: "task.created", summary: `${t.kind} 已创建` });
  if (t.started_at)
    out.push({ at: t.started_at.toISOString(), type: "task.started", summary: `第 ${t.attempts} 次执行` });
  if (t.cancel_requested_at)
    out.push({
      at: t.cancel_requested_at.toISOString(),
      type: "task.cancel_requested",
      summary: t.cancel_reason ?? "请求取消",
    });
  if (t.finished_at)
    out.push({
      at: t.finished_at.toISOString(),
      type: `task.${t.status}`,
      summary: t.status === "failed" ? `失败:${JSON.stringify(t.last_error ?? {}).slice(0, 200)}` : t.status,
    });
  if (t.subject_type === "acquisition_plan" && t.subject_id) {
    const p = await db
      .selectFrom("acquisition_plan")
      .select(["status", "status_reason", "updated_at"])
      .where("id", "=", t.subject_id)
      .executeTakeFirst();
    if (p)
      out.push({
        at: p.updated_at.toISOString(),
        type: "plan.status",
        summary: `下载计划 ${p.status}${p.status_reason ? `(${p.status_reason})` : ""}`,
        data: { planId: t.subject_id, status: p.status },
      });
    const dts = await db
      .selectFrom("download_task")
      .select(["id", "status", "progress", "updated_at"])
      .where("acquisition_plan_id", "=", t.subject_id)
      .execute();
    for (const dt of dts)
      out.push({
        at: dt.updated_at.toISOString(),
        type: "download.status",
        summary: `下载任务 ${dt.status}(进度 ${Math.round(Number(dt.progress ?? 0) * 100)}%)`,
        data: { downloadTaskId: dt.id, status: dt.status },
      });
    const recs = dts.length
      ? await db
          .selectFrom("organization_record")
          .select(["id", "status", "updated_at"])
          .where(
            "download_task_id",
            "in",
            dts.map((d) => d.id),
          )
          .execute()
      : [];
    for (const r of recs)
      out.push({
        at: r.updated_at.toISOString(),
        type: "organization.status",
        summary: `整理记录 ${r.status}`,
        data: { recordId: r.id, status: r.status },
      });
  }
  if (t.kind === "verifyAndImport") {
    const recs = await db
      .selectFrom("organization_record")
      .select(["id", "status", "updated_at"])
      .where("task_id", "=", t.id)
      .execute();
    for (const r of recs)
      out.push({
        at: r.updated_at.toISOString(),
        type: "organization.status",
        summary: `整理记录 ${r.status}`,
        data: { recordId: r.id, status: r.status },
      });
  }
  return out.sort((a, b) => a.at.localeCompare(b.at));
}

export async function retryTask(ctx: ServerContext, taskId: string, by: string, reason?: string) {
  const t = await ctx.db.selectFrom("task").selectAll().where("id", "=", taskId).executeTakeFirst();
  if (!t) throw new CapabilityError("not_found", "任务不存在");
  if (!["failed", "dead", "cancelled"].includes(t.status))
    throw new CapabilityError("precondition", `只能重试失败/死信/已取消的任务(当前 ${t.status})`);
  const input =
    t.kind === "agentRun" ? { agentRunId: t.subject_id } : { ...((t.input ?? {}) as Record<string, unknown>), taskId };
  const { runId } = await ctx.engine.enqueue(engineTaskName(t.kind), input, {
    idempotencyKey: `${t.kind}:${t.idempotency_key}:retry:${t.attempts}:${t.updated_at.getTime()}`,
    traceId: t.id,
  });
  const row = await ctx.db
    .updateTable("task")
    .set({
      status: "pending",
      engine_run_id: runId,
      cancel_requested_at: null,
      cancel_reason: null,
      finished_at: null,
      next_run_at: new Date(),
      triggered_by: t.triggered_by ?? by,
      last_error: sql`coalesce(last_error, '{}'::jsonb) || ${JSON.stringify({ retriedBy: by, retryReason: reason ?? null })}::jsonb`,
    })
    .where("id", "=", taskId)
    .returningAll()
    .executeTakeFirstOrThrow();
  if (t.kind === "agentRun" && t.subject_id)
    await ctx.db
      .updateTable("agent_run")
      .set({ status: "queued", cancel_requested_at: null, finished_at: null })
      .where("id", "=", t.subject_id)
      .execute();
  return row;
}

export async function cancelTask(ctx: ServerContext, taskId: string, reason: string) {
  const t = await ctx.db.selectFrom("task").selectAll().where("id", "=", taskId).executeTakeFirst();
  if (!t) throw new CapabilityError("not_found", "任务不存在");
  if (["succeeded", "failed", "cancelled", "dead"].includes(t.status)) return { row: t, cancelRequested: false };
  const now = new Date();
  await ctx.db
    .updateTable("task")
    .set({ cancel_requested_at: now, cancel_reason: reason })
    .where("id", "=", taskId)
    .execute();
  if (t.kind === "agentRun" && t.subject_id)
    await ctx.db
      .updateTable("agent_run")
      .set({ cancel_requested_at: now })
      .where("id", "=", t.subject_id)
      .where("cancel_requested_at", "is", null)
      .execute();
  let cancelled = false;
  if (t.engine_run_id) cancelled = await ctx.engine.cancel(t.engine_run_id, reason).catch(() => false);
  const row = await ctx.db
    .updateTable("task")
    .set(cancelled || !t.engine_run_id ? { status: "cancelled", finished_at: now } : {})
    .where("id", "=", taskId)
    .returningAll()
    .executeTakeFirstOrThrow();
  return { row, cancelRequested: true };
}

export function registerTaskCapabilities(registry: CapabilityRegistry, ctx: ServerContext): void {
  registry.register(capabilitySpec("task.inspect"), async (input) => {
    const t = await ctx.db.selectFrom("task").selectAll().where("id", "=", input.taskId).executeTakeFirst();
    if (!t) throw new CapabilityError("not_found", "任务不存在");
    const children = await ctx.db
      .selectFrom("task")
      .selectAll()
      .where("parent_task_id", "=", t.id)
      .orderBy("created_at")
      .limit(100)
      .execute();
    return {
      task: taskToApi(t) as never,
      children: children.map(taskToApi) as never,
      timeline: await taskTimeline(ctx, t),
    };
  });

  registry.register(capabilitySpec("task.retry"), async (input, c) => {
    const row = await retryTask(ctx, input.taskId, c.actorId, input.reason);
    return { taskId: row.id, status: row.status as TaskStatus };
  });

  registry.register(capabilitySpec("task.cancel"), async (input) => {
    const r = await cancelTask(ctx, input.taskId, input.reason);
    return { taskId: r.row.id, status: r.row.status as TaskStatus, cancelRequested: r.cancelRequested };
  });
}