writing-agent/tests/message-squash.test.ts

import { describe, expect, it } from "vitest";
import type { ChatMessage } from "../src/llm/client.js";
import { squashMessages } from "../src/llm/message-squash.js";

const sys = (content: string): ChatMessage => ({ role: "system", content });
const user = (content: string): ChatMessage => ({ role: "user", content });
const asst = (content: string): ChatMessage => ({ role: "assistant", content });

const toolCall: ChatMessage = {
  role: "assistant",
  content: null,
  tool_calls: [
    { id: "1", type: "function", function: { name: "t", arguments: "{}" } },
  ],
};
const toolResult: ChatMessage = { role: "tool", content: "{}", tool_call_id: "1" };

describe("squashMessages", () => {
  it("none returns the list untouched", () => {
    const input = [sys("a"), sys("b")];
    expect(squashMessages(input, "none")).toBe(input);
  });

  it("merge joins adjacent same-role messages with a blank line", () => {
    const out = squashMessages([sys("a"), sys("b"), user("c")], "merge");
    expect(out).toEqual([sys("a\n\nb"), user("c")]);
  });

  it("merge keeps role changes intact", () => {
    const out = squashMessages([sys("a"), user("b"), asst("c")], "merge");
    expect(out.map((m) => m.role)).toEqual(["system", "user", "assistant"]);
  });

  it("never merges tool groups", () => {
    const out = squashMessages(
      [sys("a"), sys("b"), toolCall, toolResult, asst("x"), asst("y")],
      "merge",
    );
    expect(out.map((m) => m.role)).toEqual([
      "system",
      "assistant",
      "tool",
      "assistant",
    ]);
    // 工具组原样保留
    expect(out[1]).toBe(toolCall);
    expect(out[2]).toBe(toolResult);
  });

  it("semi folds system into neighbouring messages", () => {
    const out = squashMessages(
      [sys("规则"), user("问题"), sys("补充"), asst("回答")],
      "semi",
    );
    expect(out).toEqual([user("规则\n\n问题"), asst("补充\n\n回答")]);
  });

  it("strict makes the first message a user", () => {
    const out = squashMessages([asst("前言"), user("问题")], "strict");
    expect(out[0]?.role).toBe("user");
    expect(out[0]?.content).toBe("前言\n\n问题");
  });

  it("single-user-turn collapses everything into one user message", () => {
    const out = squashMessages(
      [sys("规则"), user("问题"), asst("回答")],
      "single-user-turn",
    );
    expect(out).toEqual([user("规则\n\n问题\n\n回答")]);
  });

  it("falls back to merge when a tool group is present", () => {
    const out = squashMessages(
      [sys("a"), sys("b"), toolCall, toolResult],
      "single-user-turn",
    );
    expect(out.map((m) => m.role)).toEqual(["system", "assistant", "tool"]);
    expect(out[0]?.content).toBe("a\n\nb");
  });

  it("handles empty input", () => {
    expect(squashMessages([], "strict")).toEqual([]);
  });
});

describe("system 档（对齐酒馆 squashSystemMessages）", () => {
  it("只合并相邻 system，user/assistant 不动", () => {
    const out = squashMessages(
      [sys("a"), sys("b"), user("u1"), user("u2"), asst("c1"), asst("c2"), sys("d"), sys("e")],
      "system",
    );
    expect(out).toEqual([
      sys("a\nb"),
      user("u1"),
      user("u2"),
      asst("c1"),
      asst("c2"),
      sys("d\ne"),
    ]);
  });

  it("不碰工具组", () => {
    const out = squashMessages(
      [sys("a"), sys("b"), toolCall, toolResult, sys("c"), sys("d")],
      "system",
    );
    expect(out.map((m) => m.role)).toEqual(["system", "assistant", "tool", "system"]);
  });
});