fix(mcp-server): review fixes for the shell page

- a download button in the header opens the web app's export dialog (.drawio, .png, .svg, .drawio.svg), which the classic page had and the shell lacked when it became the default - the shell asks draw.io for the custom library menu (libraries=1), as the classic page did; the web app keeps libraries=0 - the newest card no longer shows "Rendering preview" for good: the sync takes the thumbnail of a diagram the server recovered from its file (saved without pictures) while the canvas kept it, and of a write whose picture was skipped because an edit came first, once the canvas shows the write again - e2e: the get_selection test covers a shape in a container the user entered; a download test saves a .drawio file
2026-10-11 20:19:51 +08:00 · 2026-10-11 20:56:07 +09:00
parent 467fd25fd1
commit f257d8ee99
7 changed files with 164 additions and 8 deletions
@@ -21,6 +21,9 @@ export interface DrawioSource {
    lang?: string
    /** No calls to external services (the desktop app) */
    offline?: boolean
+    /** draw.io's custom libraries (File > New Library, Open Library from,
+     * and the buttons under the shapes); the web app leaves them out */
+    libraries?: boolean
 }

 /**
@@ -53,7 +56,7 @@ export function getDrawioSrc(source: DrawioSource, dark: boolean): string {
    else params.themes = "1"
    Object.assign(params, {
        spin: "0",
-        libraries: "0",
+        libraries: source.libraries ? "1" : "0",
        noSaveBtn: "1",
        noExitBtn: "1",
        saveAndExit: "0",
@@ -113,7 +113,7 @@ Use the standard MCP configuration with:
 - **Edit Support**: Modify existing diagrams with natural language instructions. If any change in an edit fails, nothing is written and the AI gets the reason and the current page XML
 - **Your Edits Are Kept**: Changes you make in the browser are read before the AI edits again. If the AI overwrites a change you were still making, your version is saved in History
 - **Version History**: The last 20 versions appear as cards next to the canvas, with a thumbnail and what changed. Restore one, undo and redo the latest, or compare a version with the canvas; or ask the AI to undo (`restore_version`)
- **Export**: Save as `.drawio`, `.png`, `.svg`, or `.drawio.svg` (an SVG with the diagram embedded, which draw.io and `load_diagram` can open again) through `export_diagram`, or export an image from draw.io's **File > Export as** menu in the preview
+- **Export**: Save as `.drawio`, `.png`, `.svg`, or `.drawio.svg` (an SVG with the diagram embedded, which draw.io and `load_diagram` can open again) through `export_diagram`, or from the download button in the preview's header
 - **Multi-page**: List, add, rename, and delete pages, and edit any page
 - **Auto-save**: Each session's diagram is saved to `~/.next-ai-drawio/<session-id>.drawio`, and its last 20 versions to `<session-id>.history.json`, so the diagram, History and undo survive a restart of the MCP client
 - **Custom Instructions**: keep your own drawing rules in `~/.next-ai-drawio/instructions.md` (for example "Always draw in minimal style"); they are appended to the drawing guide on every call
@@ -231,7 +231,7 @@ Like the web app's bundled copy, it has no image proxy (`/drawio/proxy`): images

 ### Classic preview page

-`DRAWIO_PREVIEW_UI=classic` opens the preview page of the previous releases: draw.io with a **History** dialog and a **Download** button, without the outline of the AI's changes, Ctrl+Z for them, the version cards and `get_selection`. It stays for one more release and is then removed.
+`DRAWIO_PREVIEW_UI=classic` opens the preview page of the previous releases: draw.io with a **History** dialog and a **Download** button, without the outline of the AI's changes, Ctrl+Z for them, the version cards, Compare and `get_selection`. It stays for one more release and is then removed.

 ## Troubleshooting

@@ -1,10 +1,12 @@
-import { History, Moon, Sun } from "lucide-react"
+import { Download, History, Moon, Sun } from "lucide-react"
 import { useMemo, useState } from "react"
 import { CanvasStage } from "@/components/canvas/canvas-stage"
 import { CompareDialog } from "@/components/canvas/compare-dialog"
 import { VersionStrip } from "@/components/canvas/version-strip"
 import { VersionsProvider } from "@/components/canvas/versions-context"
+import { SaveDialog } from "@/components/save-dialog"
 import { AppToaster } from "@/components/workspace/app-toaster"
+import { useDiagram } from "@/contexts/diagram-context"
 import { useDictionary } from "@/hooks/use-dictionary"
 import type { DrawioSource } from "@/lib/drawio/drawio-config"
 import { cn } from "@/lib/utils"
@@ -21,20 +23,24 @@ const iconButton =
 /**
 * The MCP preview: the web app's canvas without the chat. A slim bar names
 * the session and the sync state; the versions strip, the version cards
- * (from the server's History) and Compare come from the shared canvas
- * components.
+ * (from the server's History), Compare and the export dialog come from the
+ * shared canvas components.
 */
 export function CanvasShell({ config }: { config: ShellConfig }) {
    const dict = useDictionary()
    const isDark = useSettingsStore((s) => s.isDark)
+    const { saveDiagramToFile } = useDiagram()
    const { status, sync } = useMcpSync(config)
    const versions = useNodeVersions(sync)
    const [panelOpen, setPanelOpen] = useState(true)
+    const [saveOpen, setSaveOpen] = useState(false)
    const drawioSource = useMemo<DrawioSource>(
        () => ({
            editorUrl: config.drawioBaseUrl,
            ui: config.drawioUi || undefined,
            lang: config.drawioLang || undefined,
+            // As the classic preview page: the user's own shape libraries
+            libraries: true,
        }),
        [config],
    )
@@ -89,12 +95,22 @@ export function CanvasShell({ config }: { config: ShellConfig }) {
                            <History />
                        </button>
                    )}
+                    <button
+                        type="button"
+                        onClick={() => setSaveOpen(true)}
+                        aria-label={dict.workspace.export}
+                        title={dict.workspace.export}
+                        className={cn(iconButton, !hasVersions && "ml-auto")}
+                        data-testid="download"
+                    >
+                        <Download />
+                    </button>
                    <button
                        type="button"
                        onClick={() => setDark(!isDark)}
                        aria-label={dict.shell.toggleTheme}
                        title={dict.shell.toggleTheme}
-                        className={cn(iconButton, !hasVersions && "ml-auto")}
+                        className={iconButton}
                        data-testid="theme-toggle"
                    >
                        {isDark ? <Sun /> : <Moon />}
@@ -117,6 +133,19 @@ export function CanvasShell({ config }: { config: ShellConfig }) {
                    )}
                </div>
                <CompareDialog />
+                <SaveDialog
+                    open={saveOpen}
+                    onOpenChange={setSaveOpen}
+                    onSave={(filename, format) =>
+                        saveDiagramToFile(
+                            filename,
+                            format,
+                            undefined,
+                            dict.save.savedSuccessfully,
+                        )
+                    }
+                    defaultFilename={`diagram-${new Date().toLocaleDateString("sv-SE")}`}
+                />
                <AppToaster />
            </div>
        </VersionsProvider>
@@ -300,6 +300,9 @@ export function createMcpSync(options: SyncOptions): McpSync {
        version: number
        xml: string
    } | null = null
+    // Its picture was not taken because the user had changed the canvas by
+    // then: taken when the canvas shows the write again (an undo)
+    let thumbMissed = false
    let pendingMcpExport: ExportFormat | null = null
    let mcpExportSeq = 0
    let mcpExportId: string | null = null
@@ -409,6 +412,7 @@ export function createMcpSync(options: SyncOptions): McpSync {
            decision.mode === "commit" && !sameDiagram(before, xml)
        // currentVersion is the write's version
        thumbExport = { stateId, version: currentVersion, xml }
+        thumbMissed = false
        setTimeout(captureThumbnail, 500)
    }

@@ -420,8 +424,12 @@ export function createMcpSync(options: SyncOptions): McpSync {
        if (!t || !isReady) return
        canvas.export({ format: "svg" }, 5000).then((result) => {
            if (!result?.data || thumbExport !== t) return
-            if (!sameDiagram(latestXml, t.xml)) return
+            if (!sameDiagram(latestXml, t.xml)) {
+                thumbMissed = true
+                return
+            }
            thumbExport = null
+            thumbMissed = false
            postJson("/history-svg", {
                sessionId,
                svg: svgDataUrl(result.data),
@@ -552,6 +560,13 @@ export function createMcpSync(options: SyncOptions): McpSync {
            }
            if (mine && !sameDiagram(mine, s.xml)) {
                pushState(mine, "", s.version)
+            } else if (s.xml && !s.blank && mine) {
+                // The canvas still shows the diagram the server recovered:
+                // its History entry came back without a picture (the file
+                // keeps none), take one now
+                thumbExport = { stateId, version: currentVersion, xml: s.xml }
+                thumbMissed = false
+                captureThumbnail()
            }
        } else if (s.xml) {
            // The server has a diagram this tab never showed (an AI write it
@@ -828,6 +843,16 @@ export function createMcpSync(options: SyncOptions): McpSync {
            latestXml = xml
            const commitCopy = awaitingCommitCopy
            awaitingCommitCopy = false
+            // The canvas is back at a write whose picture was not taken (an
+            // edit came first): take it now
+            if (
+                thumbMissed &&
+                thumbExport &&
+                sameDiagram(xml, thumbExport.xml)
+            ) {
+                thumbMissed = false
+                captureThumbnail()
+            }
            if (xml === lastXml || (commitCopy && sameDiagram(xml, lastXml))) {
                return
            }
@@ -317,6 +317,50 @@ test("get_selection reports the cells selected in the editor", async ({
    expect(text).toContain('id="2" shape "Pick me" at x=40 y=40 w=120 h=60')
    expect(text).toContain('id="4" edge "link" from "2" to "3"')
    expect(text).not.toContain('id="3"')
+
+    // A shape in a container is reported with it, also after the user
+    // entered the container (Ctrl+Shift+End), where it is the default parent
+    await toolText("create_new_diagram", {
+        xml: `<mxCell id="box" value="Box" style="swimlane;" vertex="1" parent="1"><mxGeometry x="40" y="40" width="300" height="200" as="geometry"/></mxCell>\n<mxCell id="kid" value="Kid" style="rounded=1;" vertex="1" parent="box"><mxGeometry x="20" y="30" width="120" height="60" as="geometry"/></mxCell>`,
+    })
+    await expect(frame.getByText("Kid")).toBeVisible({ timeout: 10_000 })
+    await frame.locator("body").evaluate((body) => {
+        const w = body.ownerDocument.defaultView as any
+        w.Draw.loadPlugin((ui: any) => {
+            const graph = ui.editor.graph
+            graph.enterGroup(graph.model.getCell("box"))
+            graph.setSelectionCell(graph.model.getCell("kid"))
+        })
+    })
+    const inside = await toolText("get_selection")
+    expect(inside).toContain(
+        'id="kid" shape "Kid" at x=20 y=30 w=120 h=60 inside "box"',
+    )
+    await frame.locator("body").evaluate((body) => {
+        const w = body.ownerDocument.defaultView as any
+        w.Draw.loadPlugin((ui: any) => ui.editor.graph.exitGroup())
+    })
+})
+
+test("the download button saves the diagram as a file", async ({ page }) => {
+    const frame = await openShell(page)
+    await toolText("create_new_diagram", { xml: box("2", "Keep me", 40) })
+    await expect(frame.getByText("Keep me")).toBeVisible({ timeout: 10_000 })
+    await page.getByTestId("download").click()
+    const dialog = page.getByRole("dialog")
+    await expect(dialog).toBeVisible()
+    // .drawio is the default format
+    const download = page.waitForEvent("download")
+    await dialog.getByRole("button", { name: "Export" }).click()
+    const file = await download
+    expect(file.suggestedFilename()).toMatch(
+        /^diagram-\d{4}-\d{2}-\d{2}\.drawio$/,
+    )
+    // draw.io's own export of the document (compressed pages, as the web
+    // app saves them; load_diagram opens them)
+    const saved = readFileSync(await file.path(), "utf8")
+    expect(saved).toMatch(/^<mxfile\b/)
+    expect(saved).toContain('name="Page-1"')
 })

 // --- Version cards from the server's History (shell/node-versions-source.ts) ---
@@ -77,6 +77,20 @@ describe("shell draw.io URL", () => {
        expect(fixed.searchParams.get("ui")).toBe("kennedy")
        expect(fixed.searchParams.get("themes")).toBeNull()
    })
+
+    it("offers draw.io's custom libraries when asked, as the classic page did", () => {
+        const withLibraries = new URL(
+            getDrawioSrc(
+                { editorUrl: "/drawio/index.html", libraries: true },
+                false,
+            ),
+        )
+        expect(withLibraries.searchParams.get("libraries")).toBe("1")
+        const without = new URL(
+            getDrawioSrc({ editorUrl: "/drawio/index.html" }, false),
+        )
+        expect(without.searchParams.get("libraries")).toBe("0")
+    })
 })

 describe("shell language", () => {
@@ -256,6 +256,23 @@ describe("MCP sync after the server recreated its session", () => {
        })
    })

+    it("takes the picture of the recovered diagram, whose History entry came back without one", async () => {
+        const t = await inStep()
+        // The process restarted and read the file back: the text the tab
+        // last saw, so the canvas stays; History is saved without pictures
+        const poll = t.sync.poll()
+        t.next("GET").answer(state("S2", 1, "<mxfile>A</mxfile>"))
+        await poll
+        expect(t.loads).toHaveLength(1)
+        expect(t.posts()).toHaveLength(0)
+        t.nextExport().answer({ data: "<svg>A</svg>" })
+        await t.settle()
+        expect(t.next("POST", "/api/history-svg").body).toMatchObject({
+            stateId: "S2",
+            version: 1,
+        })
+    })
+
    it("shows the server's diagram and keeps the tab's in History", async () => {
        const t = await inStep()
        const poll = t.sync.poll()
@@ -528,6 +545,30 @@ describe("MCP sync thumbnails", () => {
        expect(atob(push.body.svg.split(",")[1])).toBe("<svg>edit</svg>")
    })

+    it("takes the image once the user comes back to the write (an edit came before it)", async () => {
+        const { t, request } = await loadedB()
+        t.sync.onAutoSave("<mxfile>B edited</mxfile>")
+        request.answer({ data: "<svg>thumbnail</svg>" })
+        await t.settle()
+        expect(thumbnailPosts(t)).toHaveLength(0)
+        t.nextExport().answer({ data: "<svg>edit</svg>" })
+        await t.settle()
+        t.next("POST").answer({
+            status: 200,
+            body: { success: true, version: 4 },
+        })
+        await t.settle()
+        // Ctrl+Z: the canvas shows the write again
+        t.sync.onAutoSave("<mxfile>B</mxfile>")
+        t.nextExport().answer({ data: "<svg>B</svg>" })
+        await t.settle()
+        expect(thumbnailPosts(t).map((c) => c.body.version)).toEqual([3])
+        // The undo is pushed as any edit
+        t.nextExport().answer({ data: "<svg>B</svg>" })
+        await t.settle()
+        expect(t.next("POST", "/api/state").body.xml).toBe("<mxfile>B</mxfile>")
+    })
+
    it("takes the image of a write loaded before draw.io was ready, once it is", async () => {
        const t = open()
        const poll = t.sync.poll()