UHD_Media_Tool/packages/agent/test/runtime.test.ts

import { spawn } from "node:child_process";
import { existsSync, readFileSync } from "node:fs";
import { join } from "node:path";
import { AgentRunEvent } from "@umt/contracts";
import { afterEach, describe, expect, it } from "vitest";
import {
  type AgentRunHandle,
  type AgentSettings,
  checkToolCall,
  createFakeRuntime,
  type FakeScript,
} from "../src/index.ts";
import { crashScript, RUN_ID } from "./fixtures/crash-script.ts";
import { type Fixture, makeFixture, settings, tempDir, types } from "./helpers.ts";

const opened: AgentRunHandle[] = [];
afterEach(async () => {
  await Promise.all(opened.splice(0).map((r) => r.close()));
});

const dune: FakeScript = {
  match: /Dune/,
  steps: [
    {
      thinking: "先查库存",
      text: "我先查一下库里有没有。",
      toolCalls: [{ capability: "inventory.query", input: { title: "Dune", year: 2021 } }],
    },
    {
      text: "库里没有,提交 1080p。",
      toolCalls: [
        {
          capability: "download.submit",
          input: (r) => ({ title: (r[0]?.json as { title?: string } | undefined)?.title ?? "", quality: "1080p" }),
        },
      ],
    },
  ],
  finalText: (r) => (r.at(-1)?.isError ? `失败:${r.at(-1)?.text}` : "下载已提交(仅代表已提交,不代表已入库)。"),
};

async function start(
  fx: Fixture,
  scripts: FakeScript[],
  opts: {
    capabilities?: string[];
    settings?: Partial<AgentSettings>;
    runId?: string;
    sessionDir?: string;
    firstSeq?: number;
    onEvent?: (e: AgentRunEvent) => void | Promise<void>;
  } = {},
) {
  const { runtime, fake } = createFakeRuntime({
    registry: fx.registry,
    scripts,
    ...(opts.onEvent ? { onEvent: opts.onEvent } : {}),
  });
  const sessionDir = opts.sessionDir ?? (await tempDir());
  const run = await runtime.createRun({
    runId: opts.runId ?? "run-1",
    goal: "把 Dune 1080p 收进来",
    systemPrompt: "媒体库助手",
    capabilities: opts.capabilities ?? ["inventory.query", "download.submit"],
    settings: settings(opts.settings),
    sessionDir,
    actor: { actorId: "user:1", taskId: "task:1", scopeIds: ["media:dune"] },
    ...(opts.firstSeq !== undefined ? { firstSeq: opts.firstSeq } : {}),
  });
  opened.push(run);
  return { run, fake, sessionDir };
}

describe("PiAgentRuntime + FakeModelProvider", () => {
  it("完整闭环:事件顺序、usage、最终文本、JSONL 位于 sessionDir/<runId>.jsonl", async () => {
    const fx = makeFixture();
    fx.grants.set("user:1", new Set(["download.submit"]));
    const { run, sessionDir } = await start(fx, [dune]);
    const outcome = await run.prompt();

    expect(outcome.status).toBe("completed");
    expect(outcome.finalText).toBe("下载已提交(仅代表已提交,不代表已入库)。");
    expect(outcome.usage.totalTokens).toBeGreaterThan(0);
    expect(fx.submitted).toEqual([expect.objectContaining({ title: "Dune", quality: "1080p" })]);

    const events = run.history();
    expect(types(events)).toEqual([
      "run_started",
      "tool_call_start",
      "tool_call_end",
      "tool_call_start",
      "tool_call_end",
      "run_completed",
    ]);
    expect(events.map((e) => e.seq)).toEqual(events.map((_, i) => i));
    expect(events.some((e) => e.type === "thinking_delta")).toBe(true);
    const text = events.filter((e) => e.type === "text_delta").map((e) => (e as { delta: string }).delta);
    expect(text.join("")).toContain("我先查一下库里有没有。");
    const starts = events.filter((e) => e.type === "tool_call_start");
    expect(starts.map((e) => e.type === "tool_call_start" && e.capabilityId)).toEqual([
      "inventory.query",
      "download.submit",
    ]);
    const last = events.at(-1)!;
    expect(last).toMatchObject({ type: "run_completed", finalText: outcome.finalText });

    expect(run.sessionPath).toBe(join(sessionDir, "run-1.jsonl"));
    expect(existsSync(run.sessionPath)).toBe(true);
    expect(readFileSync(run.sessionPath, "utf8")).toContain("download__submit");
    // 审计来自 agent,带 invocationId 与 runId
    expect(fx.audits.find((a) => a.capabilityId === "download.submit")).toMatchObject({
      source: "agent",
      agentRunId: "run-1",
      outcome: "ok",
    });
  });

  it("events() 可按 seq 回放并在 close 后结束", async () => {
    const fx = makeFixture();
    const { run } = await start(fx, [{ match: "Dune", steps: [], finalText: "库里已有。" }]);
    const collected: AgentRunEvent[] = [];
    const consumer = (async () => {
      for await (const e of run.events()) collected.push(e);
    })();
    await run.prompt();
    const tail: AgentRunEvent[] = [];
    const replay = (async () => {
      for await (const e of run.events(1)) tail.push(e);
    })();
    await run.close();
    await Promise.all([consumer, replay]);
    expect(collected.map((e) => e.type)).toEqual(run.history().map((e) => e.type));
    expect(tail[0]?.seq).toBe(1);
    expect(collected.at(-1)?.type).toBe("run_completed");
  });

  it("onEvent 按 seq 串行投递 contracts AgentRunEvent;firstSeq 续接序号,close 等待投递完成", async () => {
    const fx = makeFixture();
    const delivered: AgentRunEvent[] = [];
    const { run } = await start(fx, [{ match: "Dune", steps: [], finalText: "库里已有。" }], {
      firstSeq: 10,
      // 故意让前面的事件落库更慢,验证不会乱序
      onEvent: async (e) => {
        await new Promise((r) => setTimeout(r, e.seq === 10 ? 20 : 0));
        delivered.push(e);
      },
    });
    await run.prompt();
    const tail: AgentRunEvent[] = [];
    const replay = (async () => {
      for await (const e of run.events(11)) tail.push(e);
    })();
    await run.close();
    await replay;
    const seqs = run.history().map((e) => e.seq);
    expect(seqs[0]).toBe(10);
    expect(delivered.map((e) => e.seq)).toEqual(seqs);
    expect(tail[0]?.seq).toBe(11);
    for (const e of delivered) expect(AgentRunEvent.safeParse(e).success, e.type).toBe(true);
  });

  it("白名单:未注册工具(bash)与未列入白名单的已注册能力都被拒绝并产生事件,不执行", async () => {
    const fx = makeFixture();
    fx.grants.set("user:1", new Set(["download.submit"]));
    const script: FakeScript = {
      match: "Dune",
      steps: [
        {
          toolCalls: [
            { toolName: "bash", input: { command: "rm -rf /" } },
            { capability: "download.submit", input: { title: "Dune", quality: "1080p" } },
            { capability: "inventory.query", input: { title: "Dune" } },
          ],
        },
      ],
      finalText: (r) => r.map((x) => `${x.toolName}:${x.isError}`).join(","),
    };
    const { run } = await start(fx, [script], { capabilities: ["inventory.query"] });
    const out = await run.prompt();
    expect(out.finalText).toBe("bash:true,download__submit:true,inventory__query:false");
    expect(fx.submitted).toHaveLength(0);
    const rejected = run.history().filter((e) => e.type === "tool_call_rejected");
    expect(rejected.map((e) => e.type === "tool_call_rejected" && [e.toolName, e.capabilityId])).toEqual([
      ["bash", undefined],
      ["download__submit", "download.submit"],
    ]);
    expect(types(run.history())).toEqual([
      "run_started",
      "tool_call_rejected",
      "tool_call_rejected",
      "tool_call_start",
      "tool_call_end",
      "run_completed",
    ]);
  });

  it("before_tool 白名单判定为纯函数", () => {
    const allowed = new Map([["inventory__query", "inventory.query"]]);
    expect(checkToolCall("inventory__query", allowed)).toEqual({ allowed: true, capabilityId: "inventory.query" });
    expect(checkToolCall("download__submit", allowed)).toMatchObject({
      allowed: false,
      capabilityId: "download.submit",
    });
    expect(checkToolCall("bash", allowed)).toMatchObject({ allowed: false });
  });

  it("授权在服务端:无授权时 download.submit 返回 unauthorized 给模型,不产生副作用", async () => {
    const fx = makeFixture();
    const { run } = await start(fx, [dune]);
    const out = await run.prompt();
    expect(out.status).toBe("completed");
    expect(out.finalText).toContain("unauthorized");
    expect(fx.submitted).toHaveLength(0);
    const end = run
      .history()
      .filter((e) => e.type === "tool_call_end")
      .at(-1);
    expect(end).toMatchObject({ capabilityId: "download.submit", isError: true });
    expect(fx.audits.at(-1)).toMatchObject({ outcome: "denied", capabilityId: "download.submit" });
  });

  it("工具输出进入模型与事件前已脱敏(passkey/cookie/Bearer)", async () => {
    const fx = makeFixture();
    const script: FakeScript = {
      match: "Dune",
      steps: [{ toolCalls: [{ capability: "release.inspect", input: { releaseId: "42" } }] }],
      finalText: "看过了",
    };
    const { run, fake } = await start(fx, [script], { capabilities: ["release.inspect"] });
    await run.prompt();
    const transcript = fake.lastRequest?.transcript ?? "";
    expect(transcript).toContain("download.php?id=42&passkey=[REDACTED]");
    for (const secret of ["abcdef123456", "deadbeef", "sk-live-XXXX"]) {
      expect(transcript).not.toContain(secret);
      expect(JSON.stringify(run.history())).not.toContain(secret);
      expect(readFileSync(run.sessionPath, "utf8")).not.toContain(secret);
    }
  });

  it("abort:执行中的工具收到取消信号,run_failed(reason=aborted)", async () => {
    const fx = makeFixture({ slowMs: 10_000 });
    fx.grants.set("user:1", new Set(["download.submit"]));
    const { run } = await start(fx, [dune]);
    const pending = run.prompt();
    for await (const e of run.events()) {
      if (e.type === "tool_call_start" && e.capabilityId === "download.submit") break;
    }
    const t0 = Date.now();
    await run.abort("用户停止");
    const out = await pending;
    expect(Date.now() - t0).toBeLessThan(3000);
    expect(out).toMatchObject({ status: "failed", reason: "aborted", message: "用户停止" });
    expect(run.history().at(-1)).toMatchObject({ type: "run_failed", reason: "aborted" });
    expect(fx.submitted).toHaveLength(0);
  });

  it("预算:工具调用次数超限 → abort + run_failed(budget)", async () => {
    const fx = makeFixture();
    const loop: FakeScript = {
      match: "Dune",
      steps: Array.from({ length: 10 }, () => ({
        toolCalls: [{ capability: "inventory.query", input: { title: "Dune" } }],
      })),
      finalText: "不该到这里",
    };
    const { run } = await start(fx, [loop], { settings: { maxToolCalls: 2 } });
    const out = await run.prompt();
    expect(out).toMatchObject({ status: "failed", reason: "budget", budget: "tool_calls" });
    expect(run.history().at(-1)).toMatchObject({ type: "run_failed", reason: "budget", budget: "tool_calls" });
    expect(fx.audits.filter((a) => a.capabilityId === "inventory.query")).toHaveLength(2);
  });

  it("预算:运行时长超限 → run_failed(budget/duration)", async () => {
    const fx = makeFixture({ slowMs: 5_000 });
    fx.grants.set("user:1", new Set(["download.submit"]));
    const { run } = await start(fx, [dune], { settings: { maxDurationMs: 1000 } });
    const t0 = Date.now();
    const out = await run.prompt();
    expect(Date.now() - t0).toBeLessThan(4000);
    expect(out).toMatchObject({ status: "failed", reason: "budget", budget: "duration" });
    expect(fx.submitted).toHaveLength(0);
  });

  it("预算:token 用量超限 → run_failed(budget/tokens)", async () => {
    const fx = makeFixture();
    const { run } = await start(fx, [dune], { settings: { costBudget: { maxTotalTokens: 50 } } });
    const out = await run.prompt();
    expect(out).toMatchObject({ status: "failed", reason: "budget", budget: "tokens" });
    expect(out.usage.totalTokens).toBeGreaterThan(50);
    expect(fx.submitted).toHaveLength(0);
  });

  it("手动压缩产生 compaction 事件", async () => {
    const fx = makeFixture();
    const { run } = await start(fx, [dune], { settings: { compaction: { mode: "manual", threshold: 0.8 } } });
    await run.prompt();
    await run.prompt("继续,Dune 还有别的版本吗");
    await run.compact();
    const comp = run.history().filter((e) => e.type === "compaction");
    expect(comp.map((e) => e.type === "compaction" && `${e.phase}:${e.reason}:${e.status ?? ""}`)).toEqual([
      "start:manual:",
      "end:manual:completed",
    ]);
  });

  it("脚本不匹配 → run_failed(error)", async () => {
    const fx = makeFixture();
    const { run } = await start(fx, [{ match: "Arrival", steps: [], finalText: "x" }]);
    const out = await run.prompt();
    expect(out).toMatchObject({ status: "failed", reason: "error" });
    expect(run.history().at(-1)?.type).toBe("run_failed");
  });

  it("JSONL 恢复:子进程在 download.submit 执行中被 SIGKILL,新进程 resume 得到“结果未知”且不重复提交", async () => {
    const sessionDir = await tempDir();
    const child = spawn(
      process.execPath,
      ["run", "-A", join(import.meta.dirname!, "fixtures/crash-run.ts"), sessionDir],
      {
        cwd: join(import.meta.dirname!, ".."),
        stdio: ["ignore", "pipe", "pipe"],
      },
    );
    let stderr = "";
    child.stderr.on("data", (d) => {
      stderr += String(d);
    });
    await new Promise<void>((resolve, reject) => {
      const timer = setTimeout(() => reject(new Error(`child timeout: ${stderr}`)), 20_000);
      child.stdout.on("data", (d) => {
        if (String(d).includes("SUBMIT_STARTED")) {
          clearTimeout(timer);
          child.kill("SIGKILL");
          resolve();
        }
      });
      child.on("exit", (code) => {
        clearTimeout(timer);
        if (code !== null) reject(new Error(`child exited early ${code}: ${stderr}`));
      });
    });
    await new Promise((r) => child.once("exit", r));

    const fx = makeFixture();
    fx.grants.set("user:1", new Set(["download.submit"]));
    const { run, fake } = await start(fx, [crashScript], { runId: RUN_ID, sessionDir });
    expect(run.pendingOperations).toBe(1);
    const out = await run.resume();
    expect(out.status).toBe("completed");
    expect(out.finalText).toContain("下载提交结果未知");
    expect(fx.submitted).toHaveLength(0);
    expect(fake.calls).toBe(1);
    const recovered = run.history().find((e) => e.type === "tool_call_end");
    expect(recovered).toMatchObject({ capabilityId: "download.submit", isError: true, recovery: true });
    expect(run.history()[0]).toMatchObject({ type: "run_started", resumed: true });
  }, 30_000);
});