Render streamed assistant responses as Markdown

2026-09-08 14:36:05 +08:00
parent aef8e1d310
commit 80cb7d2b80
7 changed files with 363 additions and 4 deletions
@@ -269,3 +269,130 @@ test("Lark workspace keeps pagination, account details and chat usable at every
    fullPage: true,
  });
 });
+
+test("Markdown streams incomplete code, finishes safely and restores history", async ({
+  page,
+}) => {
+  const errors: string[] = [];
+  page.on("pageerror", (error) => errors.push(error.message));
+  await login(page);
+  await configure(page);
+  await page.getByRole("button", { name: "打开研究助手" }).click();
+  const chat = page.getByRole("complementary", { name: "AI 研究助手" });
+  await chat.getByRole("button", { name: "新会话", exact: true }).click();
+  const prefix = "## 研究结论\n\n**正在生成**\n\n```python\nrank(close)";
+  const suffix =
+    "\n```\n\n| 指标 | 值 |\n| --- | --- |\n| Sharpe | 1.2 |\n\n[文档](https://example.com)\n\n<script>alert('unsafe')</script>\n\n[危险](javascript:alert(1))";
+  await page.evaluate(() => {
+    const original = window.fetch.bind(window);
+    window.fetch = async (url, init) => {
+      if (!String(url).match(/\/ai\/conversations\/[^/]+\/runs$/))
+        return original(url, init);
+      return new Response(
+        new ReadableStream({
+          start(controller) {
+            Object.assign(window, {
+              pushMarkdown: (events: unknown[], done = false) => {
+                for (const event of events)
+                  controller.enqueue(
+                    new TextEncoder().encode(
+                      `data: ${JSON.stringify(event)}\n\n`,
+                    ),
+                  );
+                if (done) {
+                  controller.enqueue(
+                    new TextEncoder().encode("data: [DONE]\n\n"),
+                  );
+                  controller.close();
+                }
+              },
+            });
+          },
+        }),
+        {
+          headers: {
+            "content-type": "text/event-stream",
+            "x-vercel-ai-ui-message-stream": "v1",
+          },
+        },
+      );
+    };
+  });
+  await chat
+    .getByRole("textbox", { name: "发送给研究助手" })
+    .fill("**用户原文**");
+  await chat.getByRole("button", { name: "发送", exact: true }).click();
+  await page.waitForFunction(
+    () => typeof (window as any).pushMarkdown === "function",
+  );
+  await page.evaluate(
+    (delta) =>
+      (window as any).pushMarkdown([
+        { type: "start", messageId: "markdown-answer" },
+        { type: "text-start", id: "part" },
+        { type: "text-delta", id: "part", delta },
+      ]),
+    prefix,
+  );
+  await expect(chat.getByRole("heading", { name: "研究结论" })).toBeVisible();
+  await expect(chat.locator(".ai-markdown pre")).toContainText("rank(close)");
+  await expect(chat.locator(".ai-messages")).toHaveAttribute(
+    "aria-busy",
+    "true",
+  );
+  await expect(chat.locator(".ai-message.user .ai-text")).toHaveText(
+    "**用户原文**",
+  );
+  await page.route("**/api/v1/ai/conversations/*", async (route) => {
+    if (route.request().method() !== "GET") return route.continue();
+    await route.fulfill({
+      json: {
+        id: route.request().url().split("/").at(-1),
+        title: "Markdown",
+        runs: [],
+        messages: [
+          {
+            id: "markdown-answer",
+            role: "assistant",
+            parts: [{ type: "text", text: prefix + suffix }],
+          },
+        ],
+      },
+    });
+  });
+  await page.evaluate(
+    (delta) =>
+      (window as any).pushMarkdown(
+        [
+          { type: "text-delta", id: "part", delta },
+          { type: "text-end", id: "part" },
+          { type: "finish" },
+        ],
+        true,
+      ),
+    suffix,
+  );
+  await expect(chat.locator(".ai-messages")).toHaveAttribute(
+    "aria-busy",
+    "false",
+  );
+  await expect(chat.getByRole("cell", { name: "1.2" })).toBeVisible();
+  await expect(
+    chat.getByRole("link", { name: "Link: https://example.com", exact: true }),
+  ).toHaveAttribute("href", "https://example.com");
+  await expect(
+    chat.locator(".ai-markdown script, .ai-markdown a[href^='javascript:']"),
+  ).toHaveCount(0);
+  await page.setViewportSize({ width: 390, height: 900 });
+  expect(
+    await chat
+      .locator(".ai-messages")
+      .evaluate((el) => el.scrollWidth <= el.clientWidth),
+  ).toBe(true);
+  await page.screenshot({ path: "../output/playwright/ai-markdown.png" });
+  await page.reload();
+  await page.getByRole("button", { name: "打开研究助手" }).click();
+  await expect(chat.getByRole("heading", { name: "研究结论" })).toBeVisible();
+  await expect(chat.getByRole("cell", { name: "1.2" })).toBeVisible();
+  expect(errors).toEqual([]);
+});