writing-agent/tests/preset-probe.test.ts

import { describe, expect, it } from "vitest";
import { importSillyTavernPreset } from "../src/preset/importer.js";
import {
  assemblePresetProbe,
  DEFAULT_PROBE_CONTEXT,
  PROBE_SYSTEM_PROMPT,
  runPresetProbe,
} from "../src/preset/probe.js";

function samplePreset() {
  return importSillyTavernPreset({
    prompts: [
      { identifier: "main", name: "Main", role: "system", content: "MAIN" },
      {
        identifier: "charDescription",
        name: "Desc",
        role: "system",
        marker: true,
      },
      {
        identifier: "chatHistory",
        name: "History",
        role: "system",
        marker: true,
      },
      { identifier: "jailbreak", name: "PHI", role: "system", content: "PHI" },
    ],
    prompt_order: [
      {
        character_id: 1,
        order: [
          { identifier: "main", enabled: true },
          { identifier: "charDescription", enabled: true },
          { identifier: "chatHistory", enabled: true },
          { identifier: "jailbreak", enabled: true },
        ],
      },
    ],
  }).preset;
}

describe("assemblePresetProbe", () => {
  it("fills the sandwich with sample lore, history, and the current turn", () => {
    const messages = assemblePresetProbe(
      samplePreset(),
      { message: "我推开门。" },
      { squash: "none" },
    );

    expect(messages.map((m) => [m.role, m.content])).toEqual([
      // 任务契约并进预设自有的 main 条目，不再单独一条
      ["system", `MAIN\n\n${PROBE_SYSTEM_PROMPT}`],
      ["system", `### 试跑设定\n\n${DEFAULT_PROBE_CONTEXT.loreBefore}`],
      ["system", DEFAULT_PROBE_CONTEXT.history],
      ["user", "我推开门。"],
      ["system", "PHI"],
    ]);
  });

  it("omits empty optional bags", () => {
    const messages = assemblePresetProbe(
      samplePreset(),
      {
        message: "你好",
        loreBefore: "",
        history: "",
        loreAfter: "",
        postTurn: "",
      },
      { squash: "none" },
    );

    expect(messages.map((m) => m.content)).toEqual([
      `MAIN\n\n${PROBE_SYSTEM_PROMPT}`,
      "你好",
      "PHI",
    ]);
  });
});

describe("runPresetProbe", () => {
  it("returns assembled messages without calling a model", async () => {
    const result = await runPresetProbe({
      preset: samplePreset(),
      input: { message: "我推开门。" },
    });

    expect(result.completed).toBe(false);
    expect(result.reply).toBeNull();
    // 缺省 strict：system 段会折进 user，所以只断言正文没丢
    expect(
      result.messages.map((m) => m.content).join("\n"),
    ).toContain("我推开门。");
  });

  it("returns the model reply when complete is provided", async () => {
    const result = await runPresetProbe({
      preset: samplePreset(),
      input: { message: "我推开门。" },
      complete: async () => ({ content: "请进。", reasoning: "短回复" }),
    });

    expect(result.completed).toBe(true);
    expect(result.reply).toBe("请进。");
    expect(result.reasoning).toBe("短回复");
  });

  it("keeps messages when generate fails", async () => {
    const result = await runPresetProbe({
      preset: samplePreset(),
      input: { message: "我推开门。" },
      complete: async () => {
        throw new Error("upstream down");
      },
    });

    expect(result.completed).toBe(false);
    expect(result.error).toBe("upstream down");
    expect(result.messages.length).toBeGreaterThan(0);
  });

  it("rejects an empty turn", async () => {
    await expect(
      runPresetProbe({
        preset: samplePreset(),
        input: { message: "   " },
      }),
    ).rejects.toThrow("缺少本轮输入");
  });
});