next-tjwater-agent/tests/chat/answer-ui.test.ts

import { mkdtemp, rm } from "node:fs/promises";
import { tmpdir } from "node:os";
import { join } from "node:path";
import { SessionUiStateStore } from "../../src/sessions/uiStateStore.js";
import { SessionWorkspaceStore } from "../../src/sessions/workspace-store.js";
import { describe, expect, it } from "bun:test";
import { parseAnswerUiInput } from "../../src/chat/answer-ui.js";
import { streamPromptResponse } from "../../src/routes/chatStream.js";
import { toCollaborationMessages } from "../../src/routes/chatStreamRoute.js";
import type { OpencodeRuntimeAdapter } from "../../src/runtime/opencode.js";
import type { UIMessage } from "ai";

const input = { block_id: "pressure", fallback_text: "演练压力 0.3 MPa", spec: { root: "metric", elements: { metric: { type: "Metric", props: { label: "压力", value: "0.3", unit: "MPa", source: "演练数据" }, children: [] } } } };
describe("locally rendered answers", () => {
  it("validates catalog props instead of accepting arbitrary model components", () => {
    expect(parseAnswerUiInput(input)?.version).toBe("answer-ui@1");
    for (const metric of [
      { ...input.spec.elements.metric, type: "Script" },
      { ...input.spec.elements.metric, props: { ...input.spec.elements.metric.props, execute: "alert(1)" } },
      { ...input.spec.elements.metric, children: ["metric"] },
    ]) expect(parseAnswerUiInput({ ...input, spec: { root: "metric", elements: { metric } } })).toBeNull();
  });
  it("rejects chart misalignment, nonfinite numbers and unreachable nodes", () => {
    const node = { type: "Chart", props: { title: "演练", chart_type: "line", source: "演练数据", x_data: ["08:00"], series: [{ name: "压力", data: [0.3] }] }, children: [] };
    const withNode = (value: unknown) => ({ ...input, spec: { root: "chart", elements: { chart: value } } });
    expect(parseAnswerUiInput(withNode(node))).not.toBeNull();
    for (const data of [[0.3, 0.4], [Infinity], [NaN]]) expect(parseAnswerUiInput(withNode({ ...node, props: { ...node.props, series: [{ name: "压力", data }] } }))).toBeNull();
    expect(parseAnswerUiInput({ ...input, spec: { ...input.spec, elements: { ...input.spec.elements, unused: node } } })).toBeNull();
  });
  it("flushes prose before the interface, publishes once per tool call, and suppresses background UI", async () => {
    const part = { id: "ui-call", type: "tool", tool: "present_answer", sessionID: "s", messageID: "m", state: { status: "completed", input, output: "{}", time: { start: 1, end: 2 } } };
    const events = [
      { type: "message.part.updated", properties: { sessionID: "s", part: { id: "text", type: "text", text: "先看演练资料。", sessionID: "s", messageID: "m" } } },
      ...[part, part].map((value) => ({ type: "message.part.updated", properties: { sessionID: "s", part: value } })),
      { type: "session.idle", properties: { sessionID: "s" } },
    ];
    const runtime = { prompt: async () => undefined, messages: async () => [], subscribeEvents: async () => ({ async *[Symbol.asyncIterator]() { for (const event of events) yield event; } }) } as unknown as OpencodeRuntimeAdapter;
    for (const suppressAnswerUi of [false, true]) {
      const emitted: Array<{ event: string; data: Record<string, unknown> }> = [];
      await streamPromptResponse({ runtime, sessionId: "s", clientSessionId: "client", message: "演练", suppressAnswerUi, write: (event, data) => emitted.push({ event, data }) });
      expect(emitted.filter((event) => event.event === "answer_ui")).toHaveLength(suppressAnswerUi ? 0 : 1);
      if (!suppressAnswerUi) expect(emitted.findIndex((event) => event.event === "token")).toBeLessThan(emitted.findIndex((event) => event.event === "answer_ui"));
    }
  });
  it("persists rich parts in the owned chat store and restores them without cross-scope access", async () => {
    const root = await mkdtemp(join(tmpdir(), "tjwater-rich-answer-"));
    try {
      const owner = { actorKey: "actor-answer", projectKey: "project-answer", sessionId: "answer-session" };
      const workspaces = new SessionWorkspaceStore(join(root, "workspaces"));
      const store = new SessionUiStateStore(join(root, "legacy"), workspaces);
      const parts: UIMessage["parts"] = [{ type: "text", text: "演练：" }, { type: "data-answer_ui", id: "pressure", data: parseAnswerUiInput(input) }, { type: "text", text: "仍待核验。" }];
      const messages = [{ id: "a", role: "assistant", parts }];
      await store.write(owner, { sessionId: owner.sessionId, messages });
      const restored = new SessionUiStateStore(join(root, "legacy"), new SessionWorkspaceStore(workspaces.root));
      expect((await restored.read(owner))?.messages).toEqual(messages);
      expect(await restored.read({ ...owner, actorKey: "actor-other" })).toBeNull();
      expect(await restored.read({ ...owner, projectKey: "project-other" })).toBeNull();
    } finally { await rm(root, { recursive: true, force: true }); }
  });
  it("preserves UI-only answers and ordered rich parts in task collaboration history", () => {
    const parts: UIMessage["parts"] = [{ type: "text", text: "演练：" }, { type: "data-answer_ui", id: "pressure", data: parseAnswerUiInput(input) }, { type: "text", text: "仍需核验。" }, { type: "data-final_answer", data: { content: "结论待确认。" } }];
    const messages = [{ id: "a", role: "assistant", parts: [...parts, { type: "data-progress", data: { title: "查询" } }] }, { id: "b", role: "assistant", parts: [parts[1]] }] as unknown as UIMessage[];
    const restored = toCollaborationMessages(messages);
    expect(restored[0].parts).toEqual(parts);
    expect(restored[1].parts).toEqual([parts[1]]);
  });
});

it("finishes a rich chat turn at final_answer and prevents later automatic UI revisions", async () => {
  const emitted: Array<{ event: string; data: Record<string, unknown> }> = [];
  const interrupted: string[] = [];
  const toolPart = (id: string, tool: string, args: unknown, output = "{}") => ({ type: "message.part.updated", properties: { sessionID: "s", part: { id, type: "tool", tool, sessionID: "s", messageID: "m", state: { status: "completed", input: args, output, time: { start: 1, end: 2 } } } } });
  const events = [toolPart("ui", "present_answer", input), toolPart("final", "final_answer", { answer: "演练结论仍待核验。" }, JSON.stringify({ answer: "演练结论仍待核验。", ui_block_ids: ["pressure", "unknown", "pressure"] })), toolPart("late", "present_answer", { ...input, block_id: "late" }), { type: "session.idle", properties: { sessionID: "s" } }];
  const runtime = { prompt: async () => undefined, messages: async () => [], abortSession: async (id: string) => { interrupted.push(id); }, waitForSessionIdle: async () => undefined, subscribeEvents: async () => ({ async *[Symbol.asyncIterator]() { for (const event of events) yield event; } }) } as unknown as OpencodeRuntimeAdapter;
  const result = await streamPromptResponse({ runtime, sessionId: "s", clientSessionId: "client", message: "演练", write: (event, data) => emitted.push({ event, data }) });
  expect(result).toMatchObject({ failed: false, aborted: false });
  expect(emitted.filter((event) => event.event === "answer_ui")).toHaveLength(1);
  expect(interrupted).toEqual(["s"]);
  expect(emitted.find((event) => event.event === "final_answer")?.data.content).toBe("演练结论仍待核验。");
  expect(emitted.find((event) => event.event === "final_answer")?.data.ui_block_ids).toEqual(["pressure"]);
  expect(emitted.filter((event) => event.event === "done")).toHaveLength(1);
  interrupted.length = 0;
  emitted.length = 0;
  await streamPromptResponse({ runtime, sessionId: "s", clientSessionId: "client", message: "后台执行", suppressAnswerUi: true, write: (event, data) => emitted.push({ event, data }) });
  expect(interrupted).toEqual([]);
  expect(emitted.filter((event) => event.event === "answer_ui")).toHaveLength(0);
  expect(emitted.filter((event) => event.event === "done")).toHaveLength(1);
});