laimiaoai-agent-new/tests/tools.test.ts

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<string, unknown>,
    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);
  });
});