fix(chat): draw the built-in examples again and undo edit previews on errors

Found by the PR review: - The built-in examples showed a finished card and an empty canvas. They are answered in the browser, never reach the tool handler, and relied on the final redraw that an earlier commit removed. The example branch now loads its diagram itself. - When the request failed while an edit was streaming (a provider error, a lost connection), its preview stayed on the canvas. The error handler now restores the diagram from before the preview. - The model picker could not scroll with the wheel or touch: the settings dialog blocks those events outside itself, and the picker is rendered outside it. The popover is modal now. - A fetch error and the open picker stayed when switching providers. - Editing a model id kept the old test warning and response time, which also hid the "may not be able to draw" hint for the new id.
2026-10-08 10:47:48 +08:00 · 2026-10-04 23:04:21 +09:00
parent 48128acc65
commit 2bb69b58af
4 changed files with 215 additions and 60 deletions
@@ -41,6 +41,7 @@ import type { UrlData } from "@/lib/url-utils"
 import { type FileData, useFileProcessor } from "@/lib/use-file-processor"
 import { useQuotaManager } from "@/lib/use-quota-manager"
 import { cn, formatXML, isRealDiagram } from "@/lib/utils"
+import { prepareNewDiagram } from "@/packages/mcp-server/src/new-diagram.ts"
 import { BLANK_MXFILE, hasCells } from "@/packages/mcp-server/src/pages.ts"
 import type { ValidationState } from "./chat/ValidationCard"
 import {
@@ -363,6 +364,13 @@ export default function ChatPanel({
            await handleToolCall({ toolCall }, addToolOutput)
        },
        onError: (error) => {
+            // An edit still streaming when the request failed never reaches
+            // the tool handler: undo its preview. The first stored original
+            // is the diagram before any of them.
+            const [originalXml] = editDiagramOriginalXmlRef.current.values()
+            if (originalXml) onDisplayChart(originalXml, true)
+            editDiagramOriginalXmlRef.current.clear()
+
            // Server errors are JSON: a quota limit ({type: request, token or
            // tpm}), a provider error ({type: "provider", code, message}) or
            // {error}. The SDK puts the response body in error.message.
@@ -778,8 +786,9 @@ export default function ChatPanel({
                    files.length === 1 ? files[0].name : undefined,
                )
                if (cached) {
-                    // Add user message and fake assistant response to messages
-                    // The chat-message-display useEffect will handle displaying the diagram
+                    // Add the user message and a finished display_diagram
+                    // answer, and load its diagram here: these messages never
+                    // reach the tool handler
                    const toolCallId = `cached-${Date.now()}`

                    // Build user message text including any file content
@@ -816,6 +825,11 @@ export default function ChatPanel({
                        0,
                        chartXMLRef.current || BLANK_MXFILE,
                    )
+                    const prepared = prepareNewDiagram(cached.xml, {
+                        pageId: "page-1",
+                        pageName: "Page-1",
+                    })
+                    if (prepared.ok) onDisplayChart(prepared.xml, true)
                    setInput("")
                    sessionStorage.removeItem(SESSION_STORAGE_INPUT_KEY)
                    setFiles([])
@@ -282,6 +282,8 @@ export function ModelConfigDialog({
        const newProvider = addProvider(providerType)
        setSelectedProviderId(newProvider.id)
        setValidationStatus("idle")
+        setFetchModelsError("")
+        setModelPickerOpen(false)
    }

    // Handle provider field updates
@@ -653,6 +655,10 @@ export function ModelConfigDialog({
                                                )
                                                setValidationStatus("idle")
                                                setShowApiKey(false)
+                                                // These belong to the
+                                                // provider shown before
+                                                setFetchModelsError("")
+                                                setModelPickerOpen(false)
                                            }}
                                            className={cn(
                                                "group flex items-center gap-3 px-3 py-2.5 rounded-xl w-full",
@@ -957,7 +963,11 @@ export function ModelConfigDialog({
                                                        )}
                                                    </Button>
                                                )}
+                                                {/* modal: the dialog blocks the
+                                                wheel outside itself, and the
+                                                list is rendered outside it */}
                                                <Popover
+                                                    modal
                                                    open={modelPickerOpen}
                                                    onOpenChange={
                                                        setModelPickerOpen
@@ -1310,6 +1320,10 @@ export function ModelConfigDialog({
                                                                                            undefined,
                                                                                        validationError:
                                                                                            undefined,
+                                                                                        validationWarning:
+                                                                                            undefined,
+                                                                                        responseTime:
+                                                                                            undefined,
                                                                                    },
                                                                                )
                                                                            }
@@ -45,6 +45,60 @@ const cell = (id: string, label: string, x: number) =>
 const page = (id: string, cells: string) =>
    `<diagram id="${id}" name="${id}"><mxGraphModel><root><mxCell id="0"/><mxCell id="1" parent="0"/>${cells}</root></mxGraphModel></diagram>`

+const sse = (events: object[]) =>
+    events.map((e) => `data: ${JSON.stringify(e)}\n\n`).join("")
+const EDIT_GAMMA = {
+    operations: [
+        { operation: "add", cell_id: "c", new_xml: cell("c", "Gamma", 400) },
+    ],
+}
+const editStart = (id: string) => ({
+    type: "tool-input-start",
+    toolCallId: id,
+    toolName: "edit_diagram",
+})
+const editDeltas = (id: string) =>
+    (JSON.stringify(EDIT_GAMMA).match(/[\s\S]{1,40}/g) ?? []).map((d) => ({
+        type: "tool-input-delta",
+        toolCallId: id,
+        inputTextDelta: d,
+    }))
+
+/**
+ * Answer each chat request with the next reply. Each string in a reply is
+ * one network chunk, sent 300 ms apart, so the throttled UI renders between
+ * chunks like with a real model.
+ */
+async function chunkedReplies(p: Page, replies: string[][]) {
+    await p.addInitScript((replies) => {
+        const realFetch = window.fetch
+        let n = 0
+        window.fetch = async (input, init) => {
+            const url =
+                typeof input === "string" ? input : (input as Request).url
+            if (!url.endsWith("/api/chat")) return realFetch(input, init)
+            const chunks = replies[n++] ?? [
+                'data: {"type":"start"}\n\ndata: {"type":"finish"}\n\ndata: [DONE]\n\n',
+            ]
+            const body = new ReadableStream({
+                async start(controller) {
+                    for (const chunk of chunks) {
+                        controller.enqueue(new TextEncoder().encode(chunk))
+                        await new Promise((r) => setTimeout(r, 300))
+                    }
+                    controller.close()
+                },
+            })
+            return new Response(body, {
+                headers: { "content-type": "text/event-stream" },
+            })
+        }
+    }, replies)
+    await p.goto("/", { waitUntil: "networkidle" })
+    await getIframe(p).waitFor({ state: "visible", timeout: 30000 })
+    return p.frameLocator("iframe")
+}
+
 const TWO_PAGES = `<mxfile>${page("First", cell("a", "Old A", 40))}${page("Second", cell("b", "Old B", 40))}</mxfile>`
 // Bare cells with a duplicate id and an unescaped &, which get fixed, and
 // a linked cell whose label lives on its UserObject wrapper
@@ -137,6 +191,26 @@ test("edit_diagram applies all operations or none", async ({ page: p }) => {
    await expect(canvas.getByText("Broken", { exact: true })).toHaveCount(0)
 })

+test("a built-in example draws its diagram", async ({ page: p }) => {
+    // Answered in the browser from lib/cached-responses.ts, no request
+    let requests = 0
+    await p.route("**/api/chat", (route) => {
+        requests++
+        return route.fulfill({ status: 500, body: "{}" })
+    })
+    await p.goto("/", { waitUntil: "networkidle" })
+    await getIframe(p).waitFor({ state: "visible", timeout: 30000 })
+    await sendMessage(
+        p,
+        "Give me a **animated connector** diagram of transformer's architecture",
+    )
+    await waitForCompleteCount(p, 1)
+    await expect(
+        p.frameLocator("iframe").getByText("Transformer Architecture"),
+    ).toBeVisible({ timeout: 15000 })
+    expect(requests).toBe(0)
+})
+
 test("the thinking header is in the page language", async ({ page: p }) => {
    const events = [
        { type: "start" },
@@ -198,38 +272,14 @@ test("an edit right after a broken edit call starts from the real diagram", asyn
    // Seen with Claude Opus 5.5: the first edit call had invalid JSON, the
    // server rejected it, and the model sent the same edit again at once.
    // The second edit must not see the first one's streamed preview.
-    const sse = (events: object[]) =>
-        events.map((e) => `data: ${JSON.stringify(e)}\n\n`).join("")
-    const edit = {
-        operations: [
-            {
-                operation: "add",
-                cell_id: "c",
-                new_xml: cell("c", "Gamma", 400),
-            },
-        ],
-    }
-    const deltas = (id: string) =>
-        (JSON.stringify(edit).match(/[\s\S]{1,40}/g) ?? []).map((d) => ({
-            type: "tool-input-delta",
-            toolCallId: id,
-            inputTextDelta: d,
-        }))
-    const start = (id: string) => ({
-        type: "tool-input-start",
-        toolCallId: id,
-        toolName: "edit_diagram",
-    })
-    // Each inner array is sent as one network chunk, 300 ms apart, so the
-    // throttled UI renders between chunks like with a real model
    const replies = [
        [streamedToolCall("display_diagram", { xml: cell("a", "Alpha", 40) })],
        [
            sse([
                { type: "start" },
                { type: "start-step" },
-                start("e1"),
-                ...deltas("e1"),
+                editStart("e1"),
+                ...editDeltas("e1"),
            ]),
            sse([
                {
@@ -246,48 +296,22 @@ test("an edit right after a broken edit call starts from the real diagram", asyn
                },
                { type: "finish-step" },
                { type: "start-step" },
-                start("e2"),
-                ...deltas("e2"),
+                editStart("e2"),
+                ...editDeltas("e2"),
            ]),
            `${sse([
                {
                    type: "tool-input-available",
                    toolCallId: "e2",
                    toolName: "edit_diagram",
-                    input: edit,
+                    input: EDIT_GAMMA,
                },
                { type: "finish-step" },
                { type: "finish" },
            ])}data: [DONE]\n\n`,
        ],
    ]
-    await p.addInitScript((replies) => {
-        const realFetch = window.fetch
-        let n = 0
-        window.fetch = async (input, init) => {
-            const url =
-                typeof input === "string" ? input : (input as Request).url
-            if (!url.endsWith("/api/chat")) return realFetch(input, init)
-            const chunks = replies[n++] ?? [
-                'data: {"type":"start"}\n\ndata: {"type":"finish"}\n\ndata: [DONE]\n\n',
-            ]
-            const body = new ReadableStream({
-                async start(controller) {
-                    for (const chunk of chunks) {
-                        controller.enqueue(new TextEncoder().encode(chunk))
-                        await new Promise((r) => setTimeout(r, 300))
-                    }
-                    controller.close()
-                },
-            })
-            return new Response(body, {
-                headers: { "content-type": "text/event-stream" },
-            })
-        }
-    }, replies)
-    await p.goto("/", { waitUntil: "networkidle" })
-    await getIframe(p).waitFor({ state: "visible", timeout: 30000 })
-    const canvas = p.frameLocator("iframe")
+    const canvas = await chunkedReplies(p, replies)

    await sendMessage(p, "Draw a box")
    await waitForCompleteCount(p, 1)
@@ -298,3 +322,32 @@ test("an edit right after a broken edit call starts from the real diagram", asyn
    })
    await expect(p.getByText(/No changes were made/)).toHaveCount(0)
 })
+
+test("a request that fails during an edit undoes its preview", async ({
+    page: p,
+}) => {
+    const canvas = await chunkedReplies(p, [
+        [streamedToolCall("display_diagram", { xml: cell("a", "Alpha", 40) })],
+        [
+            sse([
+                { type: "start" },
+                { type: "start-step" },
+                editStart("e1"),
+                ...editDeltas("e1"),
+            ]),
+            `${sse([{ type: "error", errorText: "Upstream connection lost" }])}data: [DONE]\n\n`,
+        ],
+    ])
+    await sendMessage(p, "Draw a box")
+    await waitForCompleteCount(p, 1)
+    await sendMessage(p, "Add another box")
+    // The preview shows the new cell while the edit streams
+    await expect(canvas.getByText("Gamma", { exact: true })).toBeVisible({
+        timeout: 15000,
+    })
+    await expect(p.getByText("Upstream connection lost").first()).toBeVisible({
+        timeout: 15000,
+    })
+    await expect(canvas.getByText("Gamma", { exact: true })).toHaveCount(0)
+    await expect(canvas.getByText("Alpha", { exact: true })).toBeVisible()
+})
@@ -14,13 +14,13 @@ const CONFIG = {
    ],
 }

-async function openQwenSettings(page: Page) {
+async function openQwenSettings(page: Page, config: object = CONFIG) {
    await page.addInitScript((config) => {
        localStorage.setItem(
            "next-ai-draw-io-model-configs",
            JSON.stringify(config),
        )
-    }, CONFIG)
+    }, config)
    await page.goto("/", { waitUntil: "networkidle" })
    await getIframe(page).waitFor({ state: "visible", timeout: 30000 })
    await page.locator("button:has(svg.lucide-bot)").first().click()
@@ -66,6 +66,31 @@ test("fetches the provider's models and adds one from the picker", async ({
    await expect(dialog.locator('input[title="qwen-new-max"]')).toBeVisible()
 })

+test("the model picker scrolls with the mouse wheel", async ({ page }) => {
+    // The picker sits in a popover above the settings dialog, which blocks
+    // wheel events outside itself
+    await page.route("**/api/provider-models", (route) =>
+        route.fulfill({
+            json: {
+                models: Array.from({ length: 60 }, (_, i) => ({
+                    id: `qwen-model-${i}`,
+                })),
+            },
+        }),
+    )
+    const dialog = await openQwenSettings(page)
+    await dialog
+        .getByRole("button", { name: "Fetch models from the provider" })
+        .click()
+    const list = page.locator("[cmdk-list]")
+    await expect(list.getByText("qwen-model-0")).toBeVisible()
+    await list.hover()
+    await page.mouse.wheel(0, 400)
+    await expect
+        .poll(() => list.evaluate((el) => el.scrollTop), { timeout: 3000 })
+        .toBeGreaterThan(0)
+})
+
 test("shows a hint when the provider rejects the key", async ({ page }) => {
    await page.route("**/api/provider-models", (route) =>
        route.fulfill({
@@ -83,3 +108,52 @@ test("shows a hint when the provider rejects the key", async ({ page }) => {
        ),
    ).toBeVisible()
 })
+
+test("a fetch error stays with its provider", async ({ page }) => {
+    await page.route("**/api/provider-models", (route) =>
+        route.fulfill({
+            status: 401,
+            json: { code: "invalid_api_key", error: "Incorrect API key" },
+        }),
+    )
+    const dialog = await openQwenSettings(page, {
+        version: 1,
+        providers: [
+            ...CONFIG.providers,
+            { id: "p2", provider: "glm", apiKey: "k", models: [] },
+        ],
+    })
+    await dialog
+        .getByRole("button", { name: "Fetch models from the provider" })
+        .click()
+    const error = dialog.getByText("Incorrect API key")
+    await expect(error).toBeVisible()
+    await dialog.getByText("GLM (Zhipu)").first().click()
+    await expect(error).toHaveCount(0)
+})
+
+test("editing a model id clears the old test warning", async ({ page }) => {
+    const warning = "Connected, but the model answered without calling a tool."
+    const dialog = await openQwenSettings(page, {
+        version: 1,
+        providers: [
+            {
+                ...CONFIG.providers[0],
+                models: [
+                    {
+                        id: "m1",
+                        modelId: "qwen-max",
+                        validated: true,
+                        validationWarning: warning,
+                    },
+                ],
+            },
+        ],
+    })
+    await expect(dialog.getByText(warning)).toBeVisible()
+    const input = dialog.locator('input[title="qwen-max"]')
+    await input.fill("qwen-mt-plus")
+    await input.blur()
+    await expect(dialog.getByText(warning)).toHaveCount(0)
+    await expect(dialog.getByText("may not be able to draw")).toBeVisible()
+})