next-tjwater-agent/tests/routes/chatStream.test.ts

import { describe, expect, it } from "bun:test";

import {
  createTokenSmoother,
  detectTokenSmoothingChunk,
  streamPromptResponse,
  type PermissionRequestPayload,
} from "../../src/routes/chatStream.js";
import { type OpencodeRuntimeAdapter } from "../../src/runtime/opencode.js";

const createEventStream = (events: unknown[]) => ({
  async *[Symbol.asyncIterator]() {
    for (const event of events) {
      yield event;
    }
  },
});

const sleep = (ms: number) => new Promise((resolve) => setTimeout(resolve, ms));

describe("streamPromptResponse", () => {
  it("flushes before prompting and resumes the same session when a message wins the completion race", async () => {
    const order: string[] = [];
    const emitted: string[] = [];
    let idleChecks = 0;
    const runtime = {
      prompt: async (sessionId: string) => { order.push(`prompt:${sessionId}`); },
      subscribeEvents: async () => createEventStream([1, 2].flatMap((round) => [
        { type: "message.part.updated", properties: { sessionID: "s", part: {
          id: `text-${round}`, type: "text", text: `round ${round}`, sessionID: "s", messageID: `m-${round}`,
        } } },
        { type: "session.idle", properties: { sessionID: "s" } },
      ])),
      messages: async () => [],
    } as unknown as OpencodeRuntimeAdapter;
    const result = await streamPromptResponse({
      runtime, sessionId: "s", clientSessionId: "client", message: "test",
      beforePrompt: async () => { order.push("flush"); },
      beforeIdle: async () => ++idleChecks === 2,
      write: (event) => emitted.push(event),
    });
    expect(result.failed).toBeFalse();
    expect(order).toEqual(["flush", "prompt:s", "prompt:s"]);
    expect(idleChecks).toBe(2);
    expect(emitted.filter((event) => event === "done")).toHaveLength(1);
  });

  it("keeps the first pending event when the V2 prompt acknowledgement arrives first", async () => {
    const emitted: { event: string; data: Record<string, unknown> }[] = [];
    const runtime = {
      prompt: async () => undefined,
      subscribeEvents: async () => ({
        async *[Symbol.asyncIterator]() {
          await sleep(5);
          yield { type: "message.part.updated", properties: { sessionID: "s", part: { id: "text", type: "text", text: "", sessionID: "s", messageID: "m" } } };
          yield { type: "message.part.delta", properties: { sessionID: "s", partID: "text", field: "text", delta: "流事件完整" } };
          yield { type: "session.idle", properties: { sessionID: "s" } };
        },
      }),
      messages: async () => { throw new Error("lost streamed events caused a fallback read"); },
    } as unknown as OpencodeRuntimeAdapter;
    const result = await streamPromptResponse({ runtime, sessionId: "s", clientSessionId: "client", message: "test", write: (event, data) => emitted.push({ event, data }) });
    expect(result.failed).toBeFalse();
    expect(emitted.filter((e) => e.event === "token").map((e) => e.data.content).join("")).toBe("流事件完整");
    expect(emitted.at(-1)?.event).toBe("done");
  });

  it("does not split short Chinese text only because it has punctuation", () => {
    expect(detectTokenSmoothingChunk("SCADA 液位异常，需要继续分析", "zh")).toBeNull();
    expect(detectTokenSmoothingChunk("SCADA 液位异常需要继续分析。", "zh")).toBeNull();
  });

  it("does not immediately release a single Chinese character", () => {
    expect(detectTokenSmoothingChunk("液", "zh")).toBeNull();
  });

  it("keeps non-Chinese text on the word path even with zh locale", () => {
    expect(detectTokenSmoothingChunk("Agent ", "zh")).toBe("Agent ");
  });

  it("falls back to bounded Chinese chunks when punctuation is absent", () => {
    const content = "SCADA 液位异常需要继续分析水质温度和液位读数变化";
    const chunk = detectTokenSmoothingChunk(content, "zh");

    expect(chunk).not.toBeNull();
    if (!chunk) {
      return;
    }
    expect(content.startsWith(chunk)).toBe(true);
    expect(Array.from(chunk).length).toBeGreaterThanOrEqual(12);
    expect(Array.from(chunk).length).toBeLessThanOrEqual(18);
  });

  it("can bypass token smoothing when disabled", async () => {
    const events: Array<{ event: string; data: Record<string, unknown> }> = [];
    const smoother = createTokenSmoother({
      enabled: false,
      delayMs: 20,
      locale: "zh",
      sessionId: "client-session-1",
      write: (event, data) => events.push({ event, data }),
    });

    smoother.writeToken("第一段");
    smoother.writeToken("第二段");
    await smoother.flush();

    expect(events.map((item) => item.data.content)).toEqual(["第一段", "第二段"]);
  });

  it("buffers Chinese single-character deltas until a bounded chunk is available", async () => {
    const events: Array<{ event: string; data: Record<string, unknown> }> = [];
    const smoother = createTokenSmoother({
      enabled: true,
      delayMs: 1,
      locale: "zh",
      sessionId: "client-session-1",
      write: (event, data) => events.push({ event, data }),
    });

    smoother.writeToken("液");
    await sleep(8);
    expect(events).toHaveLength(0);

    const content = "位异常需要继续分析水质温度和液位读数持续变化趋势";
    for (const char of Array.from(content)) {
      smoother.writeToken(char);
    }
    await sleep(20);

    expect(events.length).toBeGreaterThan(0);
    await smoother.flush();
    expect(events.map((item) => item.data.content).join("")).toBe(`液${content}`);
  });

  it("smooths long Chinese text without waiting for paragraph-sized buffers", async () => {
    const events: Array<{ event: string; data: Record<string, unknown> }> = [];
    const smoother = createTokenSmoother({
      enabled: true,
      delayMs: 1,
      locale: "zh",
      sessionId: "client-session-1",
      write: (event, data) => events.push({ event, data }),
    });

    const content = "SCADA 液位异常需要继续分析水质温度和液位读数变化";
    for (const char of Array.from(content)) {
      smoother.writeToken(char);
    }
    await sleep(20);

    expect(events.length).toBeGreaterThan(0);
    await smoother.flush();
    expect(events.map((item) => item.data.content).join("")).toBe(content);
  });

  it("flushes smoothed token output before stream completion", async () => {
    const runtime = {
      subscribeEvents: async () =>
        createEventStream([
          {
            type: "message.part.updated",
            properties: {
              sessionID: "runtime-session-1",
              part: {
                id: "text-part-1",
                sessionID: "runtime-session-1",
                messageID: "message-1",
                type: "text",
                text: "",
                time: { start: Date.now() },
              },
              time: Date.now(),
            },
          },
          {
            type: "message.part.delta",
            properties: {
              sessionID: "runtime-session-1",
              partID: "text-part-1",
              field: "text",
              delta: "SCADA 液位异常需要继续分析",
            },
          },
          {
            type: "message.part.delta",
            properties: {
              sessionID: "runtime-session-1",
              partID: "text-part-1",
              field: "text",
              delta: "水质温度和液位读数。",
            },
          },
          {
            type: "session.idle",
            properties: {
              sessionID: "runtime-session-1",
            },
          },
        ]),
      prompt: async () => undefined,
      messages: async () => [],
    } as unknown as OpencodeRuntimeAdapter;
    const events: Array<{ event: string; data: Record<string, unknown> }> = [];

    await streamPromptResponse({
      runtime,
      sessionId: "runtime-session-1",
      clientSessionId: "client-session-1",
      message: "analyze",
      write: (event, data) => events.push({ event, data }),
    });

    const doneIndex = events.findIndex((item) => item.event === "done");
    const tokenEvents = events.filter((item) => item.event === "token");

    expect(doneIndex).toBeGreaterThan(-1);
    expect(tokenEvents.length).toBeGreaterThan(0);
    expect(tokenEvents.every((item) => events.indexOf(item) < doneIndex)).toBe(true);
    expect(tokenEvents.map((item) => item.data.content).join("")).toBe(
      "SCADA 液位异常需要继续分析水质温度和液位读数。",
    );
  });

  it("forwards opencode permission requests as SSE payloads", async () => {
    const runtime = {
      subscribeEvents: async () =>
        createEventStream([
          {
            type: "permission.asked",
            properties: {
              id: "perm-1",
              sessionID: "runtime-session-1",
              permission: "bash",
              patterns: ["rm *"],
              metadata: { command: "rm tmp.txt" },
              always: ["rm *"],
            },
          },
          {
            type: "session.idle",
            properties: {
              sessionID: "runtime-session-1",
            },
          },
        ]),
      prompt: async () => undefined,
      messages: async () => [],
    } as unknown as OpencodeRuntimeAdapter;
    const events: Array<{ event: string; data: Record<string, unknown> }> = [];

    await streamPromptResponse({
      runtime,
      sessionId: "runtime-session-1",
      clientSessionId: "client-session-1",
      message: "delete temp",
      write: (event, data) => events.push({ event, data }),
    });

    const permissionEvent = events.find((item) => item.event === "permission_request");
    expect(permissionEvent?.data).toMatchObject({
      session_id: "client-session-1",
      request_id: "perm-1",
      permission: "bash",
      patterns: ["rm *"],
      target: "rm tmp.txt",
      always: ["rm *"],
    } satisfies Partial<PermissionRequestPayload>);
  });

  it("auto approves ordinary business tools in always mode", async () => {
    const replies: Array<Record<string, unknown>> = [];
    const runtime = {
      subscribeEvents: async () =>
        createEventStream([
          {
            type: "permission.asked",
            properties: {
              id: "perm-1",
              sessionID: "runtime-session-1",
              permission: "websearch",
              patterns: ["npm test"],
              metadata: { command: "npm test" },
              always: ["npm test"],
            },
          },
          {
            type: "session.idle",
            properties: {
              sessionID: "runtime-session-1",
            },
          },
        ]),
      prompt: async () => undefined,
      messages: async () => [],
      replyPermission: async (options: Record<string, unknown>) => {
        replies.push(options);
      },
    } as unknown as OpencodeRuntimeAdapter;
    const events: Array<{ event: string; data: Record<string, unknown> }> = [];

    await streamPromptResponse({
      runtime,
      sessionId: "runtime-session-1",
      clientSessionId: "client-session-1",
      message: "run tests",
      approvalMode: "always",
      write: (event, data) => events.push({ event, data }),
    });

    expect(replies).toEqual([
      {
        requestId: "perm-1",
        sessionId: "runtime-session-1",
        reply: "once",
      },
    ]);
    expect(events.some((item) => item.event === "permission_request")).toBe(false);
    expect(events.find((item) => item.event === "permission_response")?.data).toEqual({
      session_id: "client-session-1",
      request_id: "perm-1",
      reply: "once",
    });
  });

  it.each(["permission.asked", "permission.v2.asked"])("returns to manual approval after an automatic round for %s", async (eventType) => {
    let remembered = false;
    let round = 0;
    const replies: string[] = [];
    const runtime = {
      subscribeEvents: async () => {
        round += 1;
        return createEventStream([
          ...(!remembered ? [{
            type: eventType,
            properties: {
              id: `permission-${round}`,
              sessionID: "runtime-session-1",
              ...(eventType === "permission.asked"
                ? { permission: "webfetch", patterns: ["https://example.test/*"], always: ["https://example.test/*"] }
                : { action: "webfetch", resources: ["https://example.test/"], save: ["https://example.test/*"] }),
            },
          }] : []),
          { type: "session.idle", properties: { sessionID: "runtime-session-1" } },
        ]);
      },
      prompt: async () => undefined,
      messages: async () => [],
      replyPermission: async ({ reply }: { reply: string }) => {
        replies.push(reply);
        remembered = reply === "always";
      },
    } as unknown as OpencodeRuntimeAdapter;
    const automatic: Array<{ event: string; data: Record<string, unknown> }> = [];
    const manual: Array<{ event: string; data: Record<string, unknown> }> = [];
    const options = { runtime, sessionId: "runtime-session-1", clientSessionId: "client-session-1", message: "fetch a public reference" };

    await streamPromptResponse({ ...options, approvalMode: "always", write: (event, data) => automatic.push({ event, data }) });
    await streamPromptResponse({ ...options, approvalMode: "request", write: (event, data) => manual.push({ event, data }) });

    expect(replies).toEqual(["once"]);
    expect(automatic.find(({ event }) => event === "permission_response")?.data.reply).toBe("once");
    expect(manual.find(({ event }) => event === "permission_request")?.data).toMatchObject({
      request_id: "permission-2", permission: "webfetch",
    });
  });

  it.each(["shell", "bash", "edit", "mcp_fixture_write"])("automatically approves %s in always mode", async (action) => {
    const replies: Array<Record<string, unknown>> = [];
    const runtime = {
      subscribeEvents: async () =>
        createEventStream([
          {
            type: "permission.asked",
            properties: {
              id: "perm-1",
              sessionID: "runtime-session-1",
              permission: action,
              patterns: ["npm test"],
              metadata: { command: "npm test" },
              always: ["npm test"],
            },
          },
          {
            type: "session.idle",
            properties: {
              sessionID: "runtime-session-1",
            },
          },
        ]),
      prompt: async () => undefined,
      messages: async () => [],
      replyPermission: async (options: Record<string, unknown>) => {
        replies.push(options);
      },
    } as unknown as OpencodeRuntimeAdapter;
    const events: Array<{ event: string; data: Record<string, unknown> }> = [];

    await streamPromptResponse({
      runtime,
      sessionId: "runtime-session-1",
      clientSessionId: "client-session-1",
      message: "run tests",
      approvalMode: "always",
      write: (event, data) => events.push({ event, data }),
    });

    expect(replies).toEqual([{ requestId: "perm-1", sessionId: "runtime-session-1", reply: "once" }]);
    expect(events.some((item) => item.event === "permission_request")).toBe(false);
    expect(events.find((item) => item.event === "permission_response")?.data.reply).toBe("once");
  });

  it("forwards opencode v2 permission requests as SSE payloads", async () => {
    const runtime = {
      subscribeEvents: async () =>
        createEventStream([
          {
            type: "permission.v2.asked",
            properties: {
              id: "perm-v2-1",
              sessionID: "runtime-session-1",
              action: "external_directory",
              resources: ["/tmp"],
              save: ["/tmp"],
              metadata: { path: "/tmp" },
            },
          },
          {
            type: "session.idle",
            properties: {
              sessionID: "runtime-session-1",
            },
          },
        ]),
      prompt: async () => undefined,
      messages: async () => [],
    } as unknown as OpencodeRuntimeAdapter;
    const events: Array<{ event: string; data: Record<string, unknown> }> = [];

    await streamPromptResponse({
      runtime,
      sessionId: "runtime-session-1",
      clientSessionId: "client-session-1",
      message: "read /tmp",
      write: (event, data) => events.push({ event, data }),
    });

    const permissionEvent = events.find((item) => item.event === "permission_request");
    expect(permissionEvent?.data).toMatchObject({
      session_id: "client-session-1",
      request_id: "perm-v2-1",
      permission: "external_directory",
      patterns: ["/tmp"],
      target: "/tmp",
      always: ["/tmp"],
    } satisfies Partial<PermissionRequestPayload>);
  });

  it("forwards opencode question requests and replies as SSE payloads", async () => {
    const runtime = {
      subscribeEvents: async () =>
        createEventStream([
          {
            type: "question.asked",
            properties: {
              id: "question-1",
              sessionID: "runtime-session-1",
              questions: [
                {
                  header: "范围",
                  question: "选择分析范围",
                  options: [{ label: "城区", description: "中心城区" }],
                  multiple: false,
                  custom: true,
                },
              ],
            },
          },
          {
            type: "question.replied",
            properties: {
              sessionID: "runtime-session-1",
              requestID: "question-1",
              answers: [["城区", "补充说明"]],
            },
          },
          {
            type: "session.idle",
            properties: {
              sessionID: "runtime-session-1",
            },
          },
        ]),
      prompt: async () => undefined,
      messages: async () => [],
    } as unknown as OpencodeRuntimeAdapter;
    const events: Array<{ event: string; data: Record<string, unknown> }> = [];

    await streamPromptResponse({
      runtime,
      sessionId: "runtime-session-1",
      clientSessionId: "client-session-1",
      message: "ask",
      write: (event, data) => events.push({ event, data }),
    });

    expect(events.find((item) => item.event === "question_request")?.data).toMatchObject({
      session_id: "client-session-1",
      request_id: "question-1",
      questions: [
        {
          header: "范围",
          question: "选择分析范围",
          options: [{ label: "城区", description: "中心城区" }],
          multiple: false,
          custom: true,
        },
      ],
    });
    expect(events.find((item) => item.event === "question_response")?.data).toEqual({
      session_id: "client-session-1",
      request_id: "question-1",
      answers: [["城区", "补充说明"]],
    });
  });

  it("converts question tool parts into question request SSE payloads", async () => {
    const runtime = {
      subscribeEvents: async () =>
        createEventStream([
          {
            type: "message.part.updated",
            properties: {
              sessionID: "runtime-session-1",
              part: {
                id: "tool-part-1",
                sessionID: "runtime-session-1",
                messageID: "message-1",
                type: "tool",
                callID: "call-1",
                tool: "question",
                state: {
                  status: "running",
                  input: {
                    questions: [
                      {
                        question: "你觉得这个 question 工具好用吗？",
                        header: "测试问题",
                        options: [
                          {
                            label: "非常好用",
                            description: "交互清晰，选项方便",
                          },
                        ],
                      },
                    ],
                  },
                  time: { start: Date.now() },
                },
              },
              time: Date.now(),
            },
          },
          {
            type: "session.idle",
            properties: {
              sessionID: "runtime-session-1",
            },
          },
        ]),
      prompt: async () => undefined,
      messages: async () => [],
    } as unknown as OpencodeRuntimeAdapter;
    const events: Array<{ event: string; data: Record<string, unknown> }> = [];

    await streamPromptResponse({
      runtime,
      sessionId: "runtime-session-1",
      clientSessionId: "client-session-1",
      message: "ask",
      write: (event, data) => events.push({ event, data }),
    });

    expect(events.find((item) => item.event === "question_request")?.data).toMatchObject({
      session_id: "client-session-1",
      request_id: "call-1",
      questions: [
        {
          header: "测试问题",
          question: "你觉得这个 question 工具好用吗？",
          options: [
            {
              label: "非常好用",
              description: "交互清晰，选项方便",
            },
          ],
        },
      ],
      tool: {
        messageID: "message-1",
        callID: "call-1",
      },
    });
    expect(
      events.some(
        (item) => item.event === "tool_call" && item.data.tool === "question",
      ),
    ).toBe(false);
  });

  it("forwards todo updates as structured SSE payloads and progress", async () => {
    const runtime = {
      subscribeEvents: async () =>
        createEventStream([
          {
            type: "todo.updated",
            properties: {
              sessionID: "runtime-session-1",
              todos: [
                { content: "分析水位", status: "completed", priority: "high" },
                { content: "生成建议", status: "in_progress", priority: "medium" },
              ],
            },
          },
          {
            type: "session.idle",
            properties: {
              sessionID: "runtime-session-1",
            },
          },
        ]),
      prompt: async () => undefined,
      messages: async () => [],
    } as unknown as OpencodeRuntimeAdapter;
    const events: Array<{ event: string; data: Record<string, unknown> }> = [];

    await streamPromptResponse({
      runtime,
      sessionId: "runtime-session-1",
      clientSessionId: "client-session-1",
      message: "plan",
      write: (event, data) => events.push({ event, data }),
    });

    expect(
      events.find(
        (item) => item.event === "progress" && item.data.id === "todo-progress",
      )?.data,
    ).toMatchObject({
      id: "todo-progress",
      phase: "planning",
      title: "计划进度 1/2",
    });
    expect(events.find((item) => item.event === "todo_update")?.data).toMatchObject({
      session_id: "client-session-1",
      todos: [
        expect.objectContaining({
          content: "分析水位",
          status: "completed",
          priority: "high",
        }),
        expect.objectContaining({
          content: "生成建议",
          status: "in_progress",
          priority: "medium",
        }),
      ],
    });
  });

});