UHD_Media_Tool/apps/server/src/modules/agent/capability-registry.ts

import { createHash } from "node:crypto";
import {
  type AuthorizationChecker,
  type CapabilityAuditRecord,
  type CapabilityHandler,
  CapabilityRegistry,
  type IdempotencyKeyGenerator,
  stableStringify,
} from "@umt/agent";
import { type ActorType, CAPABILITY_IDS, type CapabilityKind, type CapabilitySpec } from "@umt/contracts";
import { type DB, type Kysely, sql } from "@umt/db";
import type { Logger } from "pino";
import type { z } from "zod";
import type { SettingsReader } from "../types.ts";
import { createAuthorizer } from "./authorize.ts";

// 能力注册中心工厂:authorize(automation_grant / approval + 自动化总开关)、audit(audit_event)、
// 服务端幂等键、有副作用能力的操作账本(capability_invocation)。Web / Pi / CLI / MCP 共用这一实例。

const SIDE_EFFECT: readonly CapabilityKind[] = ["propose", "approve", "execute"];

/** 首批能力总数:启动时断言注册数与此一致(漏挂处理器即启动失败) */
export const EXPECTED_CAPABILITY_COUNT = CAPABILITY_IDS.length;

/**
 * 幂等键:能力 ID + 版本 + 所属运行(agent run / 任务)+ 规范化输入的 sha256。
 * 同一运行里同参数重复调用 → 同一个键 → 账本直接返回已记录的结果(崩溃恢复后的对账也走这里)。
 */
export const serverIdempotencyKey: IdempotencyKeyGenerator = (d, input, ctx) =>
  `${d.id}:v${d.version}:${ctx.agentRunId ?? ctx.taskId}:${createHash("sha256")
    .update(stableStringify(input))
    .digest("hex")
    .slice(0, 32)}`;

export interface CapabilityRegistryDeps {
  logger: Logger;
  /** 不给 db 时退化为 W1 安全默认(无授权回调 → propose/approve/execute 一律拒绝;审计只写日志) */
  db?: Kysely<DB>;
  settings?: SettingsReader;
  now?: () => Date;
}

/** 注册时给有副作用的处理器套上操作账本 */
class LedgeredCapabilityRegistry extends CapabilityRegistry {
  constructor(
    private readonly db: Kysely<DB> | undefined,
    options: ConstructorParameters<typeof CapabilityRegistry>[0],
  ) {
    super(options);
  }

  override register<I extends z.ZodType, O extends z.ZodType>(
    descriptor: CapabilitySpec<I, O>,
    handler: CapabilityHandler<I, O>,
  ): this {
    if (!this.db || !SIDE_EFFECT.includes(descriptor.kind)) return super.register(descriptor, handler);
    return super.register(descriptor, ledgered(this.db, descriptor as unknown as CapabilitySpec, handler));
  }
}

function ledgered<I extends z.ZodType, O extends z.ZodType>(
  db: Kysely<DB>,
  d: CapabilitySpec,
  handler: CapabilityHandler<I, O>,
): CapabilityHandler<I, O> {
  return async (input, ctx) => {
    const key = ctx.idempotencyKey;
    if (!key) return handler(input, ctx);
    const prev = await db
      .selectFrom("capability_invocation")
      .select(["id", "status", "output"])
      .where("idempotency_key", "=", key)
      .executeTakeFirst();
    if (prev?.status === "succeeded") return prev.output as z.input<O>;
    const auth = ctx.authorization;
    if (prev) {
      await db
        .updateTable("capability_invocation")
        .set({ status: "running", attempts: sql`attempts + 1`, error: null })
        .where("id", "=", prev.id)
        .execute();
    } else {
      await db
        .insertInto("capability_invocation")
        .values({
          idempotency_key: key,
          capability_id: d.id,
          capability_version: d.version,
          agent_run_id: ctx.agentRunId && isUuid(ctx.agentRunId) ? ctx.agentRunId : null,
          actor_id: ctx.actorId,
          source: ctx.source ?? "unknown",
          input: JSON.stringify(input ?? {}),
          grant_id: auth?.grantId ?? null,
          approval_id: auth?.approvalId ?? null,
        })
        .onConflict((oc) =>
          oc.column("idempotency_key").doUpdateSet({ attempts: sql`capability_invocation.attempts + 1` }),
        )
        .execute();
    }
    try {
      const out = await handler(input, ctx);
      await db
        .updateTable("capability_invocation")
        .set({ status: "succeeded", output: JSON.stringify(out ?? {}), finished_at: new Date() })
        .where("idempotency_key", "=", key)
        .execute();
      return out;
    } catch (err) {
      const e = err as { code?: string; message?: string };
      await db
        .updateTable("capability_invocation")
        .set({
          status: "failed",
          error: JSON.stringify({ code: e.code ?? "internal", message: String(e.message ?? err).slice(0, 500) }),
          finished_at: new Date(),
        })
        .where("idempotency_key", "=", key)
        .execute()
        .catch(() => {});
      throw err;
    }
  };
}

const isUuid = (s: string) => /^[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}$/i.test(s);

function actorTypeOf(r: CapabilityAuditRecord): ActorType {
  if (r.source === "agent") return "agent";
  if (r.actorId.startsWith("api_key:")) return "api_key";
  if (r.actorId.startsWith("user:") || r.actorId.startsWith("admin:")) return "user";
  return "system";
}

export function createCapabilityRegistry(deps: CapabilityRegistryDeps): CapabilityRegistry {
  const log = deps.logger.child({ component: "capabilities" });
  const db = deps.db;
  let authorize: AuthorizationChecker | undefined;
  if (db && deps.settings) authorize = createAuthorizer({ db, settings: deps.settings, now: deps.now });
  return new LedgeredCapabilityRegistry(db, {
    ...(authorize ? { authorize } : {}),
    idempotencyKey: serverIdempotencyKey,
    ...(deps.now ? { now: deps.now } : {}),
    audit: async (r) => {
      log.info(
        { capability: r.capabilityId, outcome: r.outcome, actor: r.actorId, source: r.source },
        "capability.invoke",
      );
      if (!db) return;
      await db
        .insertInto("audit_event")
        .values({
          actor_type: actorTypeOf(r),
          actor_id: r.actorId,
          action: r.outcome === "denied" ? "capability.denied" : "capability.invoke",
          target_type: "capability",
          target_id: r.capabilityId,
          payload: JSON.stringify({
            outcome: r.outcome,
            version: r.version,
            kind: r.kind,
            riskLevel: r.riskLevel,
            source: r.source ?? null,
            agentRunId: r.agentRunId ?? null,
            invocationId: r.invocationId ?? null,
            idempotencyKey: r.idempotencyKey ?? null,
            grantId: r.grantId ?? null,
            approvalId: r.approvalId ?? null,
            errorCode: r.errorCode ?? null,
            errorMessage: r.errorMessage ?? null,
            inputSummary: r.inputSummary,
            durationMs: r.durationMs,
          }),
          trace_id: r.agentRunId ?? r.taskId,
        })
        .execute();
      // 单次审批只用一次:成功执行后标记 consumed
      if (r.outcome === "ok" && r.approvalId) {
        await db
          .updateTable("approval")
          .set({ status: "consumed" })
          .where("id", "=", r.approvalId)
          .where("status", "=", "approved")
          .execute();
      }
    },
  });
}

/** 启动断言:首批能力全部挂上处理器 */
export function assertCapabilityCount(registry: CapabilityRegistry): void {
  const n = registry.list().length;
  if (n !== EXPECTED_CAPABILITY_COUNT) {
    const ids = registry
      .list()
      .map((c) => c.descriptor.id)
      .sort();
    throw new Error(`能力注册数 ${n} ≠ ${EXPECTED_CAPABILITY_COUNT}:${ids.join(", ")}`);
  }
}