UHD_Media_Tool/apps/server/test/mcp/mcp.test.ts

import { Client } from "@modelcontextprotocol/sdk/client/index.js";
import { StreamableHTTPClientTransport } from "@modelcontextprotocol/sdk/client/streamableHttp.js";
import type { CapabilitySpec } from "@umt/contracts";
import { sql } from "@umt/db";
import { afterAll, beforeAll, describe, expect, it } from "vitest";
import { z } from "zod";
import { createCapabilityRegistry } from "../../src/modules/agent/capability-registry.ts";
import { ensureSetupToken } from "../../src/modules/auth-settings/auth-service.ts";
import { createMcpModule } from "../../src/modules/mcp/index.ts";
import { RateWindow } from "../../src/modules/mcp/sessions.ts";
import { Browser, bearer, createHarness, type Harness, pgAvailable } from "../auth-settings/harness.ts";
import { json } from "../helpers.ts";

// AG-1:对外 MCP 服务端。用官方 SDK client 走 Streamable HTTP 实际连一次:
// 工具列表按 API Key scope 过滤、调用 read 能力、execute 能力无授权被拒(同一 registry.invoke → 审计 / 账本)、
// 会话与 Key 绑定、启用开关 / 允许列表、资源与提示模板、能力统计与能力清单端点。

const hasPg = await pgAvailable();

const spec = (over: Partial<CapabilitySpec>) =>
  ({
    version: 1,
    title: "t",
    description: "测试用",
    idempotent: true,
    supportsPreview: false,
    riskLevel: "low",
    inputSchema: z.object({ text: z.string() }),
    outputSchema: z.object({ echo: z.string(), token: z.string().optional() }),
    ...over,
  }) as unknown as CapabilitySpec;

const ECHO = spec({ id: "demo.echo", kind: "read", scope: "read", replay: "safe" });
const APPLY = spec({
  id: "demo.apply",
  kind: "execute",
  scope: "organize.execute",
  riskLevel: "high",
  replay: "never",
});
const SPAWN = spec({ id: "agent.spawn", kind: "propose", scope: "organize.execute", replay: "never" });

describe.skipIf(!hasPg)("MCP 服务端(真实 PG + SDK client)", () => {
  let h: Harness;
  let readKey: string;
  let execKey: string;
  let applied = 0;
  const mcp = createMcpModule({ version: "0.0.0-test" });

  beforeAll(async () => {
    h = await createHarness({ extraModules: [mcp] });
    // 换成带库的注册中心:真实授权器 + audit_event + 操作账本
    const reg = createCapabilityRegistry({ logger: h.ctx.logger, db: h.db, settings: h.ctx.settings });
    reg.register(ECHO, (i) => ({ echo: (i as { text: string }).text, token: "secret-token-value" }));
    reg.register(APPLY, (i) => {
      applied++;
      return { echo: `applied:${(i as { text: string }).text}` };
    });
    reg.register(SPAWN, () => ({ echo: "x" }));
    h.ctx.capabilities = reg;

    const admin = new Browser(h.app);
    const token = await ensureSetupToken(h.ctx);
    const r = await admin.req("POST", "/auth/setup", {
      setupToken: token,
      username: "admin",
      password: "correct horse battery",
    });
    expect(r.status).toBe(200);
    const mk = async (name: string, scopes: string[]) =>
      (await json(await admin.req("POST", "/api-keys", { name, scopes }))).data.plaintextKey as string;
    readKey = await mk("reader", ["read"]);
    execKey = await mk("executor", ["read", "organize.execute"]);
  }, 120_000);

  afterAll(async () => {
    await mcp.stop?.();
    await h?.close();
  });

  async function connect(key: string) {
    const client = new Client({ name: "vitest-client", version: "1.0.0" });
    const transport = new StreamableHTTPClientTransport(new URL("http://umt.test/api/v1/mcp"), {
      requestInit: { headers: { authorization: `Bearer ${key}` } },
      fetch: async (url, init) => await h.app.fetch(new Request(url, init)),
    });
    await client.connect(transport);
    return { client, transport };
  }

  it("read Key:只看到 read / preview 工具,调用只读能力拿到脱敏结果,execute 工具被拒", async () => {
    const { client } = await connect(readKey);
    const tools = (await client.listTools()).tools.map((t) => t.name);
    expect(tools).toContain("demo__echo");
    expect(tools).not.toContain("demo__apply");
    expect(tools).not.toContain("agent__spawn");
    const echo = (await client.listTools()).tools.find((t) => t.name === "demo__echo");
    expect(echo?.inputSchema).toMatchObject({ type: "object", properties: { text: { type: "string" } } });
    expect(echo?.annotations?.readOnlyHint).toBe(true);

    const ok = await client.callTool({ name: "demo__echo", arguments: { text: "hi" } });
    expect(ok.isError).toBeFalsy();
    expect(ok.structuredContent).toMatchObject({ echo: "hi" });
    expect(JSON.stringify(ok.content)).not.toContain("secret-token-value");

    const bad = await client.callTool({ name: "demo__echo", arguments: { text: 1 } });
    expect(bad.isError).toBe(true);
    expect(JSON.stringify(bad.content)).toContain("invalid_input");

    const denied = await client.callTool({ name: "demo__apply", arguments: { text: "x" } });
    expect(denied.isError).toBe(true);
    expect(JSON.stringify(denied.content)).toContain("forbidden");
    expect(applied).toBe(0);
    await client.close();
  });

  it("execute Key:能看到 execute 工具,但没有持续授权时授权器拒绝;有授权后执行并记账本", async () => {
    const { client } = await connect(execKey);
    expect((await client.listTools()).tools.map((t) => t.name)).toContain("demo__apply");

    const denied = await client.callTool({ name: "demo__apply", arguments: { text: "a" } });
    expect(denied.isError).toBe(true);
    expect(JSON.stringify(denied.content)).toContain("unauthorized");
    expect(applied).toBe(0);
    const audits = await sql<{ n: number }>`
      select count(*)::int as n from audit_event
       where action = 'capability.denied' and target_id = 'demo.apply' and payload->>'source' = 'mcp'`.execute(h.db);
    expect(audits.rows[0]!.n).toBeGreaterThanOrEqual(1);

    await h.db.insertInto("automation_grant").values({ action: "demo.apply", granted_by: "user:admin" }).execute();
    const ok = await client.callTool({ name: "demo__apply", arguments: { text: "a" } });
    expect(ok.isError).toBeFalsy();
    expect(ok.structuredContent).toMatchObject({ echo: "applied:a" });
    // 同一会话同参数:账本返回上次结果,不重复执行
    await client.callTool({ name: "demo__apply", arguments: { text: "a" } });
    expect(applied).toBe(1);
    const ledger = await h.db
      .selectFrom("capability_invocation")
      .select(["source", "status", "actor_id"])
      .where("capability_id", "=", "demo.apply")
      .execute();
    expect(ledger).toEqual([{ source: "mcp", status: "succeeded", actor_id: expect.stringMatching(/^api_key:/) }]);
    await client.close();
  });

  it("资源与提示模板", async () => {
    const { client } = await connect(readKey);
    const res = (await client.listResources()).resources.map((r) => r.uri);
    expect(res).toEqual(["umt://settings/summary", "umt://decisions/pending", "umt://subscriptions/overview"]);
    for (const uri of res) {
      const c = await client.readResource({ uri });
      const first = c.contents[0] as { text: string };
      expect(() => JSON.parse(first.text)).not.toThrow();
    }
    const settings = JSON.parse(
      ((await client.readResource({ uri: "umt://settings/summary" })).contents[0] as { text: string }).text,
    );
    expect(settings.automation).toMatchObject({ enabled: true });
    const prompts = (await client.listPrompts()).prompts.map((p) => p.name);
    expect(prompts).toEqual(["triage_pending", "acquire_title", "diagnose_task"]);
    const p = await client.getPrompt({ name: "acquire_title", arguments: { title: "沙丘" } });
    expect(JSON.stringify(p.messages)).toContain("沙丘");
    await expect(client.getPrompt({ name: "diagnose_task", arguments: {} })).rejects.toThrow(/taskId/);
    await client.close();
  });

  it("会话与 API Key 绑定;cookie 会话、无凭证、停用与允许列表都被拒", async () => {
    const { client, transport } = await connect(readKey);
    const sid = transport.sessionId!;
    expect(sid).toBeTruthy();
    const hijack = await h.app.request("/api/v1/mcp", {
      method: "POST",
      headers: {
        authorization: `Bearer ${execKey}`,
        "mcp-session-id": sid,
        "content-type": "application/json",
        accept: "application/json, text/event-stream",
      },
      body: JSON.stringify({ jsonrpc: "2.0", id: 9, method: "tools/list" }),
    });
    expect(hijack.status).toBe(403);
    expect((await h.app.request("/api/v1/mcp", { method: "POST", body: "{}" })).status).toBe(401);

    const status = await json(await bearer(h.app, execKey)("GET", "/mcp/status"));
    expect(status.error?.code).toBe("forbidden"); // admin scope

    const admin = new Browser(h.app);
    await admin.req("POST", "/auth/login", { username: "admin", password: "correct horse battery" });
    const init = {
      jsonrpc: "2.0",
      id: 1,
      method: "initialize",
      params: { protocolVersion: "2025-06-18", capabilities: {}, clientInfo: { name: "x", version: "1" } },
    };
    const viaCookie = await admin.req("POST", "/mcp", init, {
      headers: { accept: "application/json, text/event-stream" },
    });
    expect(viaCookie.status).toBe(403);

    const st = (await json(await admin.req("GET", "/mcp/status"))).data;
    expect(st).toMatchObject({ enabled: true, path: "/api/v1/mcp", transport: "streamable-http" });
    expect(st.sessions.length).toBeGreaterThanOrEqual(1);
    expect(st.recentCalls.length).toBeGreaterThan(0);
    expect(st.toolsByScope.read).toBe(1);

    // 允许列表只放 execKey 之外的 Key → readKey 现有会话立即断开,新连接 403
    const keys = (await json(await admin.req("GET", "/api-keys?limit=50"))).data.items as {
      id: string;
      name: string;
    }[];
    const execId = keys.find((k) => k.name === "executor")!.id;
    const upd = await admin.req("PATCH", "/mcp/config", { allowedApiKeyIds: [execId] });
    expect(upd.status).toBe(200);
    await expect(client.listTools()).rejects.toThrow();
    await expect(connect(readKey)).rejects.toThrow();
    const exec = await connect(execKey);
    await exec.client.close();

    await admin.req("PATCH", "/mcp/config", { enabled: false });
    await expect(connect(execKey)).rejects.toThrow();
    await admin.req("PATCH", "/mcp/config", { enabled: true, allowedApiKeyIds: null });
    const again = await connect(readKey);
    await again.client.close();
  });

  it("能力统计与能力清单端点", async () => {
    const get = bearer(h.app, readKey);
    const stats = (await json(await get("GET", "/capabilities/stats?days=7"))).data;
    const echo = stats.items.find((i: { capabilityId: string }) => i.capabilityId === "demo.echo");
    expect(echo).toMatchObject({ lastCaller: "mcp" });
    expect(echo.calls).toBeGreaterThanOrEqual(2);
    expect(echo.failures).toBeGreaterThanOrEqual(1);
    expect(echo.byCaller.mcp).toBe(echo.calls);
    const apply = stats.items.find((i: { capabilityId: string }) => i.capabilityId === "demo.apply");
    expect(apply.denied).toBeGreaterThanOrEqual(1);

    const tools = (await json(await get("GET", "/capabilities/tools"))).data;
    const by = (id: string) => tools.items.find((i: { id: string }) => i.id === id);
    expect(by("demo.echo")).toMatchObject({ authState: "read_only", toolName: "demo__echo", mcpExposed: true });
    expect(by("demo.apply")).toMatchObject({ authState: "granted", activeGrants: 1 });
    expect(by("agent.spawn")).toMatchObject({ authState: "agent_only", mcpExposed: false });
    expect(by("demo.echo").inputJsonSchema).toMatchObject({ type: "object" });
  });
});

describe("RateWindow", () => {
  it("窗口内超限拒绝,窗口滑过后恢复", () => {
    let t = 0;
    const w = new RateWindow(2, 1000, () => t);
    expect(w.take("k")).toBe(true);
    expect(w.take("k")).toBe(true);
    expect(w.take("k")).toBe(false);
    expect(w.take("other")).toBe(true);
    t = 1001;
    expect(w.take("k")).toBe(true);
  });
});