next-tjwater-frontend/src/features/agent/answer-ui/answer-presentation.test.ts

import { describe, expect, it } from "vitest";
import type { AgentChatMessage } from "../types";
import { canExpandAnswer, getAnswerPresentation, getFinalAnswerTranscript } from "./answer-presentation";
import type { AnswerUiPart } from "./contract";

const ui = (id: string): AnswerUiPart => ({
  type: "answer-ui",
  block: {
    version: "answer-ui@1",
    block_id: id,
    fallback_text: id,
    spec: {
      root: "text",
      elements: { text: { type: "Text", props: { text: id, tone: "neutral" }, children: [] } }
    }
  }
});
const parts: AnswerUiPart[] = [
  { type: "text", text: "正在核对资料。" },
  ui("report-chart"),
  ui("temporary"),
  { type: "text", text: "查询结束。" }
];
const message: AgentChatMessage = {
  id: "a",
  role: "assistant",
  content: "正在核对资料。",
  answerParts: parts
};
describe("final report and collapsible process", () => {
  it("keeps short warnings and location confirmations inline while retaining report readers", () => {
    const note = ui("note");
    if (note.type !== "answer-ui") throw new Error("invalid fixture");
    note.block.spec.elements.text.props = { text: "当前工具不能确认高亮。", tone: "warning" };
    expect(canExpandAnswer([{ type: "text", text: "已定位管道。" }, note])).toBe(false);
    expect(canExpandAnswer([{ type: "text", text: "完整报告".repeat(300) }])).toBe(true);
    for (const type of ["Table", "Chart", "FollowUpForm"]) {
      note.block.spec.elements.text.type = type;
      expect(canExpandAnswer([note])).toBe(true);
    }
  });
  it("shows all incremental content while waiting for final_answer", () => {
    expect(getAnswerPresentation(message)).toEqual({ report: parts, process: [], final: false });
  });
  it("keeps only the final report and explicitly selected components expanded", () => {
    const view = getAnswerPresentation({
      ...message,
      finalAnswer: {
        content: "最后的报告内容。",
        uiBlockIds: ["report-chart", "report-chart", "missing"]
      }
    });
    expect(view.report).toEqual([{ type: "text", text: "最后的报告内容。" }, parts[1]]);
    expect(view.process).toEqual([parts[0], parts[2], parts[3]]);
    expect(view.final).toBe(true);
  });
  it("folds every UI block when the final report does not reference any", () => {
    const view = getAnswerPresentation({
      ...message,
      finalAnswer: { content: "报告", uiBlockIds: [] }
    });
    expect(view.report).toEqual([{ type: "text", text: "报告" }]);
    expect(view.process).toEqual(parts);
  });
});

function phasedUi(id: string, phase: "process" | "final"): AnswerUiPart {
  const part = ui(id);
  if (part.type === "answer-ui") {
    part.block.phase = phase;
    part.block.fallback_text = `摘要 ${id}`;
  }
  return part;
}

describe("final UI anchors", () => {
  const chart = phasedUi("chart", "final");
  const progress = phasedUi("progress", "process");
  it("separates explicit process blocks from the report before a final answer arrives", () => {
    const view = getAnswerPresentation({ ...message, answerParts: [parts[0], progress, chart, parts[3]] });
    expect(view.report).toEqual([parts[0], chart, parts[3]]);
    expect(view.process).toEqual([progress]);
    expect(view.final).toBe(false);
  });
  const anchoredMessage = (content: string, uiBlockIds = ["chart", "progress"]): AgentChatMessage => ({
    ...message,
    answerParts: [parts[0], progress, chart, parts[3]],
    finalAnswer: { content, uiBlockIds }
  });

  it("places final UI between the original Markdown sections and mounts duplicates once", () => {
    const content = "# 分析\n\n- 第一项\n\n[[ui:chart]]\n\n## 建议\n保留换行\n[[ui:chart]]\n";
    const view = getAnswerPresentation(anchoredMessage(content));
    expect(view.report).toEqual([
      { type: "text", text: "# 分析\n\n- 第一项\n\n" },
      chart,
      { type: "text", text: "\n## 建议\n保留换行\n" }
    ]);
    expect(view.process).toEqual([parts[0], progress, parts[3]]);
    expect(getFinalAnswerTranscript(anchoredMessage(content))).toBe("# 分析\n\n- 第一项\n\n摘要 chart\n## 建议\n保留换行\n");
  });

  it("never includes process UI even when referenced and anchored", () => {
    const value = anchoredMessage("结论\n[[ui:progress]]\n[[ui:missing]]\n[[ui:chart]]", ["progress", "missing"]);
    const view = getAnswerPresentation(value);
    expect(view.report.every((part) => part.type === "text")).toBe(true);
    expect(view.process).toEqual([parts[0], progress, chart, parts[3]]);
    const transcript = getFinalAnswerTranscript(value);
    expect(transcript).not.toContain("progress");
    expect(transcript).not.toContain("missing");
    expect(transcript).not.toContain("chart");
    expect(transcript).toContain("暂不可用");
  });

  it("preserves fenced, indented, and quoted markers as literal Markdown", () => {
    const literal = "```md\n[[ui:chart]]\n```\n\n~~~\n[[ui:chart]]\n~~~\n\n    [[ui:chart]]\n\n> [[ui:chart]]\n";
    const view = getAnswerPresentation(anchoredMessage(literal, []));
    expect(view.report).toEqual([{ type: "text", text: literal }]);
  });

  it("retains legacy references and appends final references without anchors", () => {
    const view = getAnswerPresentation(anchoredMessage("报告"));
    expect(view.report).toEqual([{ type: "text", text: "报告" }, chart]);
    expect(getFinalAnswerTranscript(anchoredMessage("报告"))).toBe("报告\n\n摘要 chart");
  });

  it("does not expand a short answer because of marker identifiers", () => {
    expect(canExpandAnswer([{ type: "text", text: `结论\n[[ui:${"x".repeat(1500)}]]` }])).toBe(false);
  });
});