From f257d8ee99e1cca115ad62ce1f93d0c8cd3dd8db Mon Sep 17 00:00:00 2001
From: "dayuan.jiang" <jdy.toh@gmail.com>
Date: Sun, 11 Oct 2026 17:47:36 +0900
Subject: [PATCH] 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
---
 lib/drawio/drawio-config.ts                 |  5 ++-
 packages/mcp-server/README.md               |  4 +-
 packages/mcp-server/shell/canvas-shell.tsx  | 37 +++++++++++++++--
 packages/mcp-server/shell/mcp-sync-core.ts  | 27 ++++++++++++-
 packages/mcp-server/tests/e2e/shell.spec.ts | 44 +++++++++++++++++++++
 tests/unit/mcp-shell-config.test.ts         | 14 +++++++
 tests/unit/mcp-sync-core.test.ts            | 41 +++++++++++++++++++
 7 files changed, 164 insertions(+), 8 deletions(-)

diff --git a/lib/drawio/drawio-config.ts b/lib/drawio/drawio-config.ts
index 13672add..ea6de9d3 100644
--- a/lib/drawio/drawio-config.ts
+++ b/lib/drawio/drawio-config.ts
@@ -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",
diff --git a/packages/mcp-server/README.md b/packages/mcp-server/README.md
index c098f783..b9b9d243 100644
--- a/packages/mcp-server/README.md
+++ b/packages/mcp-server/README.md
@@ -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
 
diff --git a/packages/mcp-server/shell/canvas-shell.tsx b/packages/mcp-server/shell/canvas-shell.tsx
index 6288eae4..1d14fcef 100644
--- a/packages/mcp-server/shell/canvas-shell.tsx
+++ b/packages/mcp-server/shell/canvas-shell.tsx
@@ -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>
diff --git a/packages/mcp-server/shell/mcp-sync-core.ts b/packages/mcp-server/shell/mcp-sync-core.ts
index 62a7c032..06ed689d 100644
--- a/packages/mcp-server/shell/mcp-sync-core.ts
+++ b/packages/mcp-server/shell/mcp-sync-core.ts
@@ -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
             }
diff --git a/packages/mcp-server/tests/e2e/shell.spec.ts b/packages/mcp-server/tests/e2e/shell.spec.ts
index d14bad8a..aa01b38f 100644
--- a/packages/mcp-server/tests/e2e/shell.spec.ts
+++ b/packages/mcp-server/tests/e2e/shell.spec.ts
@@ -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) ---
diff --git a/tests/unit/mcp-shell-config.test.ts b/tests/unit/mcp-shell-config.test.ts
index 0b1dd494..145456d7 100644
--- a/tests/unit/mcp-shell-config.test.ts
+++ b/tests/unit/mcp-shell-config.test.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", () => {
diff --git a/tests/unit/mcp-sync-core.test.ts b/tests/unit/mcp-sync-core.test.ts
index 74f6da64..64261805 100644
--- a/tests/unit/mcp-sync-core.test.ts
+++ b/tests/unit/mcp-sync-core.test.ts
@@ -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()