diff --git a/components/chat-panel.tsx b/components/chat-panel.tsx index d0d938de..181cb6a1 100644 --- a/components/chat-panel.tsx +++ b/components/chat-panel.tsx @@ -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([]) diff --git a/components/model-config-dialog.tsx b/components/model-config-dialog.tsx index 818b7f93..8290f0ea 100644 --- a/components/model-config-dialog.tsx +++ b/components/model-config-dialog.tsx @@ -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({ )} )} + {/* modal: the dialog blocks the + wheel outside itself, and the + list is rendered outside it */} const page = (id: string, cells: string) => `${cells}` +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 = `${page("First", cell("a", "Old A", 40))}${page("Second", cell("b", "Old B", 40))}` // 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() +}) diff --git a/tests/e2e/provider-models.spec.ts b/tests/e2e/provider-models.spec.ts index 7c53f125..092e33f6 100644 --- a/tests/e2e/provider-models.spec.ts +++ b/tests/e2e/provider-models.spec.ts @@ -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() +})