import { join } from "node:path"; import { describe, expect, it, vi } from "vitest"; import { createAgentTools } from "../src/agent/tools.js"; import { UsageLedger } from "../src/billing/ledger.js"; import { GatewayClient } from "../src/gateway/client.js"; import { DIRECT_SKILL_KEY, directSkill } from "../src/skills/direct.js"; import type { Quote } from "../src/types.js"; import { makeCatalog, silentLogger, skillFixture, tempRoot, testConfig } from "./helpers.js"; interface GenerateDetails { ok: boolean; quote?: Quote; images?: string[]; task_id?: string; idempotent?: boolean; error?: string; } interface TestTool { name: string; execute: ( toolCallId: string, params: Record, signal?: AbortSignal, ) => Promise<{ content: Array<{ type: string; text: string }>; details: GenerateDetails }>; } async function setup(perImage = 20, allowedSkillKeys?: string[], direct = false) { const root = tempRoot(); const config = testConfig(root); const catalog = await makeCatalog([ skillFixture({ key: "poster", config: { pricePerImage: perImage } }), skillFixture({ key: "other", config: { pricePerImage: perImage } }), ]); const ledger = new UsageLedger(join(root, "usage.jsonl"), silentLogger); const gateway = new GatewayClient(config.gateway, silentLogger); const directModeSkill = direct ? directSkill(config) : undefined; const tools = createAgentTools({ config, catalog, ledger, gateway, logger: silentLogger, identity: { uid: "1", team_id: "1" }, sessionId: "session-1", allowedSkillKeys: directModeSkill ? [DIRECT_SKILL_KEY] : allowedSkillKeys, directSkill: directModeSkill, onImages: () => {}, }); const find = (name: string): TestTool => { const tool = tools.find((item) => item.name === name); if (!tool) throw new Error(`missing tool ${name}`); return tool as unknown as TestTool; }; return { root, config, catalog, ledger, gateway, tools, find }; } describe("agent tools", () => { it("lists and reads skills", async () => { const { find } = await setup(); const list = await find("list_skills").execute("call-list", {}); expect(list.content[0]?.text).toContain("poster"); const read = await find("read_skill").execute("call-read", { skill_id: "poster" }); expect(read.content[0]?.text).toContain("Use this skill for tests"); expect(read.details.ok).toBe(true); }); it("no longer exposes the quote_image tool", async () => { const { tools } = await setup(); expect(tools.some((tool) => tool.name === "quote_image")).toBe(false); }); it("generates images and records the actual returned count", async () => { const { find, ledger, gateway } = await setup(); const spy = vi .spyOn(gateway, "generateImage") .mockResolvedValueOnce({ taskId: "task-1", urls: ["https://img/1.png", "https://img/2.png"], status: "success", raw: {}, elapsedMs: 12, }) .mockResolvedValueOnce({ taskId: "task-2", urls: ["https://img/3.png"], status: "success", raw: {}, elapsedMs: 8, }); const first = await find("generate_image").execute("call-1", { skill_id: "poster", prompt: "a travel poster", n: 2, }); expect(first.details.ok).toBe(true); expect(first.details.images).toHaveLength(2); expect(first.details.quote?.total_credits).toBe(40); expect(first.details.quote?.billing_mode).toBe("dry_run"); expect(ledger.list()).toHaveLength(1); expect(ledger.list()[0]?.status).toBe("succeeded"); expect(ledger.list()[0]?.tool_call_id).toBe("call-1"); const replay = await find("generate_image").execute("call-1", { skill_id: "poster", prompt: "a travel poster", n: 2, }); expect(replay.details.idempotent).toBe(true); expect(spy).toHaveBeenCalledTimes(1); expect(ledger.list()).toHaveLength(1); const fewer = await find("generate_image").execute("call-2", { skill_id: "poster", prompt: "a travel poster", n: 2, }); expect(fewer.details.quote?.quantity).toBe(1); expect(fewer.details.quote?.total_credits).toBe(20); expect(ledger.list()).toHaveLength(2); }); it("records a failed generation without charging it", async () => { const { find, ledger, gateway } = await setup(); vi.spyOn(gateway, "generateImage").mockRejectedValue(new Error("gateway exploded")); const result = await find("generate_image").execute("call-fail", { skill_id: "poster", prompt: "a travel poster", n: 1, }); expect(result.details.ok).toBe(false); expect(result.content[0]?.text).toContain("gateway exploded"); expect(ledger.list()).toHaveLength(1); expect(ledger.list()[0]?.status).toBe("failed"); }); it("limits tools to the routed skill when allowedSkillKeys is set", async () => { const { find } = await setup(20, ["poster"]); const list = await find("list_skills").execute("call-list", {}); expect(list.content[0]?.text).toContain("poster"); expect(list.content[0]?.text).not.toContain("other"); const blocked = await find("generate_image").execute("call-blocked", { skill_id: "other", prompt: "a travel poster", }); expect(blocked.details.ok).toBe(false); expect(blocked.details.error).toBe("unknown_skill"); }); it("generates directly without a skill_id in direct mode", async () => { const { find, ledger, gateway } = await setup(20, undefined, true); const spy = vi.spyOn(gateway, "generateImage").mockResolvedValue({ taskId: "task-direct", urls: ["https://img/direct.png"], status: "success", raw: {}, elapsedMs: 3, }); const result = await find("generate_image").execute("call-direct", { prompt: "a clean product photo" }); expect(result.details.ok).toBe(true); expect(result.details.images).toEqual(["https://img/direct.png"]); expect(spy).toHaveBeenCalledOnce(); expect(ledger.list()[0]?.skill_id).toBe(DIRECT_SKILL_KEY); // 直接生成按默认单价(10)的 1.3 倍计费:ceil(10 * 1.3) = 13 expect(ledger.list()[0]?.total_credits).toBe(13); const blocked = await find("generate_image").execute("call-direct-blocked", { skill_id: "poster", prompt: "a poster", }); expect(blocked.details.ok).toBe(false); expect(blocked.details.error).toBe("unknown_skill"); }); it("rejects unknown skills and invalid quantities before calling the gateway", async () => { const { find, ledger, gateway } = await setup(); const spy = vi.spyOn(gateway, "generateImage"); const unknown = await find("generate_image").execute("call-unknown", { skill_id: "nope", prompt: "a travel poster", }); expect(unknown.details.ok).toBe(false); expect(unknown.details.error).toBe("unknown_skill"); const invalid = await find("generate_image").execute("call-invalid", { skill_id: "poster", prompt: "a travel poster", n: 99, }); expect(invalid.details.ok).toBe(false); expect(spy).not.toHaveBeenCalled(); expect(ledger.list()).toHaveLength(0); }); });