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: "我推开门。",
    });

    expect(messages.map((m) => [m.role, m.content])).toEqual([
      ["system", "MAIN"],
      [
        "system",
        `### 任务契约\n\n${PROBE_SYSTEM_PROMPT}\n\n### 试跑设定\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: "",
    });

    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();
    expect(result.messages.some((m) => m.content === "我推开门。")).toBe(true);
  });

  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("缺少本轮输入");
  });
});