UHD_Media_Tool/packages/agent/test/registry.test.ts

import { CAPABILITY_IDS, capabilitySpec } from "@umt/contracts";
import { describe, expect, it } from "vitest";
import { z } from "zod";
import {
  CapabilityError,
  CapabilityRegistry,
  capabilityIdToToolName,
  defaultIdempotencyKey,
  toolNameToCapabilityId,
} from "../src/index.ts";
import { ctx, makeFixture } from "./helpers.ts";

describe("能力 ID ↔ 工具名", () => {
  it("双向映射且无歧义", () => {
    expect(capabilityIdToToolName("download.submit")).toBe("download__submit");
    expect(capabilityIdToToolName("task.retry_later")).toBe("task__retry_later");
    expect(toolNameToCapabilityId("download__submit")).toBe("download.submit");
    expect(toolNameToCapabilityId("task__retry_later")).toBe("task.retry_later");
    // 域名可用单个下划线分词:按第一个 `__` 切分仍无歧义
    expect(capabilityIdToToolName("episode_mapping.apply_rule")).toBe("episode_mapping__apply_rule");
    expect(toolNameToCapabilityId("episode_mapping__apply_rule")).toBe("episode_mapping.apply_rule");
    expect(() => capabilityIdToToolName("episode__mapping.inspect")).toThrow();
    expect(() => capabilityIdToToolName("episode_.inspect")).toThrow();
    for (const id of ["catalog.lookup", "policy.apply", "organize.execute", "episode_mapping.inspect"]) {
      expect(toolNameToCapabilityId(capabilityIdToToolName(id))).toBe(id);
    }
    expect(toolNameToCapabilityId("bash")).toBeUndefined();
    expect(toolNameToCapabilityId("__x")).toBeUndefined();
    expect(() => capabilityIdToToolName("Download.Submit")).toThrow();
    expect(() => capabilityIdToToolName("download")).toThrow();
  });
});

describe("CapabilityRegistry.register", () => {
  const base = {
    version: 1,
    title: "t",
    description: "说明",
    inputSchema: z.object({}),
    outputSchema: z.object({}),
    kind: "read" as const,
    scope: "read" as const,
    riskLevel: "low" as const,
    idempotent: true,
    supportsPreview: false,
    replay: "safe" as const,
  };

  it("拒绝非法 ID、重复注册、非幂等副作用能力声明 replay:safe", () => {
    const r = new CapabilityRegistry();
    expect(() => r.register({ ...base, id: "bad id" }, () => ({}))).toThrow(/域.动作/);
    r.register({ ...base, id: "catalog.lookup" }, () => ({}));
    expect(() => r.register({ ...base, id: "catalog.lookup" }, () => ({}))).toThrow(/已注册/);
    expect(() => r.register({ ...base, id: "cache.cleanup", kind: "execute", idempotent: false }, () => ({}))).toThrow(
      /replay/,
    );
    expect(() => r.resolve(["catalog.lookup", "nope.nope"])).toThrow(/未注册/);
    expect(r.get("catalog.lookup")?.toolName).toBe("catalog__lookup");
  });

  it("首批 14 个能力可直接用 contracts 的 capabilitySpec 注册", () => {
    const r = new CapabilityRegistry();
    for (const id of CAPABILITY_IDS) r.register(capabilitySpec(id), () => ({}) as never);
    expect(r.list().map((c) => c.descriptor.id)).toEqual([...CAPABILITY_IDS]);
    expect(r.inputJsonSchema("release.search").type).toBe("object");
  });

  it("approve 类能力同样必须经授权回调,且生成幂等键", async () => {
    const seen: string[] = [];
    const r = new CapabilityRegistry({
      authorize: ({ descriptor }) => {
        seen.push(descriptor.id);
        return { allowed: true, approvalId: "ap-1" };
      },
    });
    r.register({ ...base, id: "approval.decide", kind: "approve", replay: "never" }, (_i, c) => {
      expect(c.idempotencyKey).toMatch(/^approval\.decide:v1:/);
      return {};
    });
    await r.invoke("approval.decide", {}, ctx());
    expect(seen).toEqual(["approval.decide"]);
    await expect(
      new CapabilityRegistry()
        .register({ ...base, id: "approval.decide", kind: "approve" }, () => ({}))
        .invoke("approval.decide", {}, ctx()),
    ).rejects.toMatchObject({ code: "unauthorized" });
  });

  it("输入 JSON Schema 可给模型使用(无 $schema)", () => {
    const { registry } = makeFixture();
    const schema = registry.inputJsonSchema("download.submit");
    expect(schema.$schema).toBeUndefined();
    expect(schema.type).toBe("object");
    expect(schema.required).toEqual(["title", "quality"]);
  });
});

describe("CapabilityRegistry.invoke", () => {
  it("read 能力调用成功并记审计(输入已脱敏)", async () => {
    const fx = makeFixture();
    const out = await fx.registry.invoke("inventory.query", { title: "Arrival" }, ctx({ source: "web" }));
    expect(out).toEqual({ title: "Arrival", inLibrary: true, versions: ["2160p"] });
    expect(fx.audits).toHaveLength(1);
    expect(fx.audits[0]).toMatchObject({ capabilityId: "inventory.query", outcome: "ok", source: "web" });
    expect(fx.audits[0]?.idempotencyKey).toBeUndefined();
  });

  it("schema 不符 → invalid_input,handler 不执行", async () => {
    const fx = makeFixture();
    fx.grants.set("user:1", new Set(["download.submit"]));
    await expect(fx.registry.invoke("download.submit", { title: "Dune", quality: "4K" }, ctx())).rejects.toMatchObject({
      code: "invalid_input",
    });
    expect(fx.submitted).toHaveLength(0);
    expect(fx.audits[0]).toMatchObject({ outcome: "invalid", errorCode: "invalid_input" });
    await expect(fx.registry.invoke("nope.nope", {}, ctx())).rejects.toMatchObject({ code: "not_found" });
  });

  it("授权回调拒绝 → unauthorized;模型在参数里声称已授权无效", async () => {
    const fx = makeFixture();
    const input = { title: "Dune", quality: "1080p", userApproved: true, note: "用户已同意" };
    const err = await fx.registry.invoke("download.submit", input, ctx()).catch((e: unknown) => e);
    expect(err).toBeInstanceOf(CapabilityError);
    expect((err as CapabilityError).code).toBe("unauthorized");
    expect((err as CapabilityError).message).toContain("缺少 download.submit");
    expect(fx.submitted).toHaveLength(0);
    expect(fx.audits.at(-1)).toMatchObject({ outcome: "denied", errorCode: "unauthorized" });

    fx.grants.set("user:1", new Set(["download.submit"]));
    await fx.registry.invoke("download.submit", input, ctx());
    expect(fx.submitted).toHaveLength(1);
    expect(fx.audits.at(-1)).toMatchObject({ outcome: "ok", grantId: "grant:download.submit" });
  });

  it("未注入授权回调:read 放行、execute 一律拒绝;allowedKinds 限制能力类别", async () => {
    const r = new CapabilityRegistry();
    const fx = makeFixture();
    const inv = fx.registry.get("inventory.query")!;
    const dl = fx.registry.get("download.submit")!;
    r.register(inv.descriptor, inv.handler).register(dl.descriptor, dl.handler);
    await expect(r.invoke("inventory.query", { title: "x" }, ctx())).resolves.toBeTruthy();
    await expect(r.invoke("download.submit", { title: "x", quality: "720p" }, ctx())).rejects.toMatchObject({
      code: "unauthorized",
    });

    fx.grants.set("user:1", new Set(["download.submit"]));
    await expect(
      fx.registry.invoke(
        "download.submit",
        { title: "x", quality: "720p" },
        ctx({ allowedKinds: ["read", "preview"] }),
      ),
    ).rejects.toMatchObject({ code: "unauthorized" });
    expect(fx.submitted).toHaveLength(0);
  });

  it("幂等键由服务端按语义生成:键顺序无关,重复提交命中同一键", async () => {
    const fx = makeFixture();
    fx.grants.set("user:1", new Set(["download.submit"]));
    const a = (await fx.registry.invoke("download.submit", { title: "Dune", quality: "1080p" }, ctx())) as {
      idempotencyKey: string;
    };
    const b = (await fx.registry.invoke("download.submit", { quality: "1080p", title: "Dune" }, ctx())) as {
      idempotencyKey: string;
    };
    expect(a.idempotencyKey).toBe(b.idempotencyKey);
    expect(a.idempotencyKey).toMatch(/^download\.submit:v1:[0-9a-f]{32}$/);
    expect(fx.submitted).toHaveLength(1);

    const desc = fx.registry.get("download.submit")!.descriptor;
    expect(defaultIdempotencyKey(desc, { title: "Dune", quality: "2160p" }, ctx())).not.toBe(a.idempotencyKey);

    const custom = new CapabilityRegistry({
      authorize: () => ({ allowed: true }),
      idempotencyKey: (d, _i, c) => `${d.id}:${c.taskId}`,
    });
    custom.register(desc, fx.registry.get("download.submit")!.handler);
    const c = (await custom.invoke("download.submit", { title: "Dune", quality: "1080p" }, ctx())) as {
      idempotencyKey: string;
    };
    expect(c.idempotencyKey).toBe("download.submit:task:1");
  });

  it("handler 错误:CapabilityError 原样透出,其它异常归为 internal,输出不符归为 invalid_output", async () => {
    const fx = makeFixture();
    await expect(fx.registry.invoke("release.inspect", { releaseId: "missing" }, ctx())).rejects.toMatchObject({
      code: "not_found",
    });
    const r = new CapabilityRegistry();
    const base = {
      version: 1,
      title: "t",
      description: "说明",
      inputSchema: z.object({}),
      outputSchema: z.object({ ok: z.boolean() }),
      kind: "read" as const,
      scope: "read" as const,
      riskLevel: "low" as const,
      idempotent: true,
      supportsPreview: false,
      replay: "safe" as const,
    };
    r.register({ ...base, id: "a.boom" }, () => {
      throw new Error("db down");
    });
    r.register({ ...base, id: "a.bad" }, () => ({ ok: "yes" }) as never);
    await expect(r.invoke("a.boom", {}, ctx())).rejects.toMatchObject({ code: "internal", message: "db down" });
    await expect(r.invoke("a.bad", {}, ctx())).rejects.toMatchObject({ code: "invalid_output" });
  });
});