laimiaoai-agent-new/tests/batch.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 type { Quote } from "../src/types.js";
import { makeCatalog, silentLogger, skillFixture, tempRoot, testConfig } from "./helpers.js";

interface BatchDetails {
  ok: boolean;
  quote?: Quote;
  images?: string[];
  error?: string;
  idempotent?: boolean;
}

interface TestTool {
  name: string;
  execute: (
    toolCallId: string,
    params: Record<string, unknown>,
    signal?: AbortSignal,
  ) => Promise<{ content: Array<{ type: string; text: string }>; details: BatchDetails }>;
}

async function setup(options: { content?: string } = {}) {
  const root = tempRoot();
  const config = testConfig(root);
  const catalog = await makeCatalog([
    skillFixture({ key: "suite", config: { pricePerImage: 20 }, content: options.content }),
  ]);
  const ledger = new UsageLedger(join(root, "usage.jsonl"), silentLogger);
  const gateway = new GatewayClient(config.gateway, silentLogger);
  const captured: string[] = [];
  const tools = createAgentTools({
    config,
    catalog,
    ledger,
    gateway,
    logger: silentLogger,
    identity: { uid: "1", team_id: "1" },
    sessionId: "session-1",
    allowedSkillKeys: ["suite"],
    onImages: (urls) => { captured.push(...urls); },
  });
  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 { find, ledger, gateway, captured };
}

const style = "统一风格：轻奢商业摄影，暖白背景，柔和侧光，低饱和暖棕色调，衬线标题字，大量留白";
const images = [
  { content: "主图，产品正面全身展示" },
  { content: "主图，佩戴场景特写" },
  { content: "详情图，材质细节" },
  { content: "详情图，规格参数" },
];

describe("generate_images_batch", () => {
  it("submits every image as its own request but all concurrently", async () => {
    const { find, gateway, ledger, captured } = await setup();
    const prompts: string[] = [];
    let call = 0;
    const spy = vi.spyOn(gateway, "generateImage").mockImplementation(async (input) => {
      const index = call++;
      prompts.push(input.prompt);
      expect(input.n).toBe(1);
      await new Promise((resolve) => setTimeout(resolve, 120));
      return {
        taskId: `task-${index}`,
        urls: [`https://img/${index}.png`],
        status: "success",
        raw: {},
        elapsedMs: 120,
      };
    });

    const t0 = Date.now();
    const result = await find("generate_images_batch").execute("batch-1", {
      skill_id: "suite",
      style,
      images,
    });
    const elapsed = Date.now() - t0;

    expect(spy).toHaveBeenCalledTimes(4);
    expect(prompts.every((prompt) => prompt.startsWith(style))).toBe(true);
    expect(prompts[0]).toContain("图1：主图，产品正面全身展示");
    expect(prompts[3]).toContain("图4：详情图，规格参数");
    expect(result.details.ok).toBe(true);
    expect(result.details.images).toHaveLength(4);
    expect(captured).toHaveLength(4);
    expect(result.details.quote?.total_credits).toBe(80);
    expect(ledger.list()).toHaveLength(4);
    expect(ledger.list().every((record) => record.status === "succeeded")).toBe(true);
    // 4 张并发：总耗时应接近单张耗时，而不是 4 倍
    expect(elapsed).toBeLessThan(320);
  });

  it("supports a different size per image", async () => {
    const { find, gateway } = await setup();
    const sizes: string[] = [];
    vi.spyOn(gateway, "generateImage").mockImplementation(async (input) => {
      sizes.push(input.size);
      return {
        taskId: `task-${sizes.length}`,
        urls: [`https://img/${sizes.length}.png`],
        status: "success",
        raw: {},
        elapsedMs: 5,
      };
    });

    const result = await find("generate_images_batch").execute("batch-size", {
      skill_id: "suite",
      style,
      images: [
        { content: "主图01", size: "800x800" },
        { content: "详情图01", size: "790x1200" },
      ],
    });
    expect(result.details.ok).toBe(true);
    expect(sizes).toEqual(["800x800", "790x1200"]);
  });

  it("keeps successful images when one item fails", async () => {
    const { find, gateway, ledger } = await setup();
    let call = 0;
    vi.spyOn(gateway, "generateImage").mockImplementation(async () => {
      const index = call++;
      if (index === 1) throw new Error("gateway timeout");
      await new Promise((resolve) => setTimeout(resolve, 20));
      return {
        taskId: `task-${index}`,
        urls: [`https://img/${index}.png`],
        status: "success",
        raw: {},
        elapsedMs: 20,
      };
    });

    const result = await find("generate_images_batch").execute("batch-2", {
      skill_id: "suite",
      style,
      images,
    });
    expect(result.details.ok).toBe(true);
    expect(result.details.images).toHaveLength(3);
    expect(result.content[0]?.text).toContain("失败 1 张");
    expect(ledger.list().filter((record) => record.status === "failed")).toHaveLength(1);
    expect(ledger.list().filter((record) => record.status === "succeeded")).toHaveLength(3);
  });

  it("replays a batch without re-generating or re-billing", async () => {
    const { find, gateway } = await setup();
    let call = 0;
    const spy = vi.spyOn(gateway, "generateImage").mockImplementation(async () => {
      const index = call++;
      return {
        taskId: `task-${index}`,
        urls: [`https://img/${index}.png`],
        status: "success",
        raw: {},
        elapsedMs: 5,
      };
    });

    await find("generate_images_batch").execute("batch-3", { skill_id: "suite", style, images });
    expect(spy).toHaveBeenCalledTimes(4);

    const replay = await find("generate_images_batch").execute("batch-3", { skill_id: "suite", style, images });
    expect(spy).toHaveBeenCalledTimes(4);
    expect(replay.details.images).toHaveLength(4);
  });

  it("requires one overall style for a multi-image set", async () => {
    const { find, gateway } = await setup();
    const spy = vi.spyOn(gateway, "generateImage");
    const result = await find("generate_images_batch").execute("batch-no-style", {
      skill_id: "suite",
      images,
    });
    expect(result.details.ok).toBe(false);
    expect(result.content[0]?.text).toContain("总体风格");
    expect(spy).not.toHaveBeenCalled();
  });

  it("rejects batches larger than the skill max_images", async () => {
    const { find, gateway } = await setup({ content: "---\nmax_images: 2\n---\n# suite\n" });
    const spy = vi.spyOn(gateway, "generateImage");
    const result = await find("generate_images_batch").execute("batch-4", {
      skill_id: "suite",
      style,
      images,
    });
    expect(result.details.ok).toBe(false);
    expect(result.content[0]?.text).toContain("超过技能");
    expect(spy).not.toHaveBeenCalled();
  });
});