feat(mcp-server): screenshot attached to create_new_diagram and edit_diagram results (#983)

* feat(mcp-server): optional screenshot attached to create_new_diagram and edit_diagram results Both write tools take an optional boolean `screenshot`; the default comes from DRAWIO_AUTO_SCREENSHOT. The body of screenshot_diagram moved into captureScreenshot so the write tools can append the PNG and checklist to their result, or a "Screenshot skipped" note when the preview tab is not available. The preview page delays the PNG export by 600 ms right after loading a new version so icon images finish loading. * fix(mcp-server): review fixes for screenshot attached to create and edit results * fix(mcp-server): Codex review fixes for the screenshot on write results - A cleared canvas also gets the screenshot note when one was asked for - The note for a tab that never polled says it may not have connected yet * docs: list undo and custom drawing rules among the MCP features * test(mcp-server): truncation check with named styles and compact cells after the merge
2026-10-11 20:19:51 +08:00 · 2026-10-11 20:54:17 +09:00
parent 8ab837679e
commit 07f3e3c2a3
10 changed files with 267 additions and 101 deletions
@@ -117,6 +117,7 @@ Then ask the AI for "a flowchart of user authentication with login, MFA and sess
 -   A screenshot tool, so the AI can look at the rendered diagram and fix it
 -   Version history, multi-page diagrams, and download as `.drawio`, `.png`, `.svg` or `.drawio.svg`
 -   Auto-save to `~/.next-ai-drawio/`, so a diagram survives a restart
+-   Ask the AI to undo, and keep your own drawing rules in `~/.next-ai-drawio/instructions.md`

 See the [MCP server README](./packages/mcp-server/README.md) for VS Code, Cursor and other client configurations.

@@ -117,6 +117,7 @@ claude mcp add drawio -- npx @next-ai-drawio/mcp-server@latest
 -   截图工具，AI 可以看一眼画好的图并自行修正
 -   版本历史、多页图表，下载为 `.drawio`、`.png`、`.svg` 或 `.drawio.svg`
 -   自动保存到 `~/.next-ai-drawio/`，重启后接着画
+-   可以让 AI 撤销修改，也可以在 `~/.next-ai-drawio/instructions.md` 里写下自己的画图规则

 VS Code、Cursor 等客户端的配置见 [MCP 服务器 README](../../packages/mcp-server/README.md)。

@@ -117,6 +117,7 @@ claude mcp add drawio -- npx @next-ai-drawio/mcp-server@latest
 -   スクリーンショットツール。AI が描いた結果を確認して自分で直せる
 -   バージョン履歴、複数ページのダイアグラム、`.drawio`、`.png`、`.svg`、`.drawio.svg` でのダウンロード
 -   `~/.next-ai-drawio/` への自動保存。再起動後も続きから描ける
+-   AI に元に戻す操作を頼める。独自の作図ルールは `~/.next-ai-drawio/instructions.md` に書ける

 VS Code、Cursor などの設定は [MCP サーバーの README](../../packages/mcp-server/README.md) を参照してください。

@@ -123,6 +123,7 @@ Claude Code <--stdio--> MCP Server <--http--> Browser (draw.io)
 | `DRAWIO_LANG` | unset | Language of the draw.io editor, such as `en`, `zh`, `zh-tw`, `ja` or `de`. Unset, draw.io chooses (the browser language on `embed.diagrams.net`, English on a self-hosted draw.io) and the user can change it under **Extras > Language** |
 | `DRAWIO_UI` | unset | draw.io theme: `kennedy`, `atlas`, `dark`, `min`, `sketch` or `simple`. Unset, the user picks one under **Extras > Theme** |
 | `DRAWIO_DARK` | `auto` | Dark mode of the draw.io editor: `auto` follows the system, `1` forces dark, `0` forces light |
+| `DRAWIO_AUTO_SCREENSHOT` | unset | Set to `true` to attach a screenshot to every `create_new_diagram` and `edit_diagram` result (costs 2 to 10 s per call) |
 | `DEBUG` | unset | Set to `true` to log debug messages |

 ## Links
@@ -107,7 +107,7 @@ Use the standard MCP configuration with:

 - **Real-time Preview**: Diagrams appear and update in your browser as the AI creates them
 - **Drawing Rules**: The AI gets the same layout, edge and style rules as the web app, plus the shape library docs (AWS, Azure, GCP, Kubernetes, Cisco and more), so it uses real icon names instead of guessing
- **Self-check**: The AI can take a screenshot of the rendered diagram and fix overlapping shapes or edges that cross shapes
+- **Self-check**: The AI can take a screenshot of the rendered diagram (`screenshot: true` on `create_new_diagram` or `edit_diagram`, or `screenshot_diagram`) and fix overlapping shapes or edges that cross shapes
 - **Draw from Your Files**: ask the AI to draw from a document, image or web page; it reads the source with the host's own tools and draws. Existing .drawio files open with load_diagram
 - **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
@@ -188,6 +188,7 @@ To give the AI your own drawing rules, write them in `~/.next-ai-drawio/instruct
 | `DRAWIO_LANG` | unset | Language of the draw.io editor. Unset, draw.io chooses: the browser language on `embed.diagrams.net`, English on a self-hosted draw.io until the user picks one under **Extras > Language**. A code such as `en`, `zh`, `zh-tw`, `ja` or `de` fixes it and hides that submenu. |
 | `DRAWIO_UI` | unset | draw.io theme. Unset, the user picks one under **Extras > Theme** and draw.io remembers it. `kennedy`, `atlas`, `dark`, `min`, `sketch` or `simple` fixes the theme and hides that menu. |
 | `DRAWIO_DARK` | `auto` | Dark mode of the draw.io editor: `auto` follows the system, `1` forces dark, `0` forces light. The page header keeps following the system. |
+| `DRAWIO_AUTO_SCREENSHOT` | unset | Set to `true` to attach a screenshot to every `create_new_diagram` and `edit_diagram` result, so the AI checks each drawing. Costs 2 to 10 s per call; the preview tab must be open and in front. A call can still pass `screenshot: false`. |
 | `DEBUG` | unset | Set to `true` to log debug messages to stderr. |

 ### Private Deployment (Self-hosted draw.io)
@@ -22,7 +22,7 @@ export const DRAWING_GUIDE = `# Draw.io drawing guide
 - Before drawing, describe your layout plan in 2-3 sentences, so shapes do not overlap and edges do not cross shapes.
 - Send XML only through tool calls, never in chat text. Never draw a box just to send the user a message.
 - Before using any icon library (AWS, Azure, GCP, Kubernetes, Cisco, BPMN, Material Design, web icons...), call get_shape_library and use the exact style names it returns. NEVER guess icon style names. For AWS, use the AWS 2025 icons (library aws4).
- After drawing or heavily editing a complex diagram, call screenshot_diagram once to see the result, and fix overlapping shapes and edges that cross shapes.
+- After drawing or heavily editing a complex diagram, pass screenshot: true on that create_new_diagram or edit_diagram call (or call screenshot_diagram) to see the result, and fix overlapping shapes and edges that cross shapes.
 - Drawing from a source: to draw from a document (PDF, Markdown, code, data), an image or screenshot, or a web page, first read it yourself (the attachment the user shared, or your own file-reading or web-fetch tools), then plan the layout and draw. This server never receives attachments; it only receives the XML you send. Extract the entities and relationships the diagram needs; do not copy the text into boxes.
 - When replicating a diagram from an image, match its style and layout closely: straight or curved lines, rounded or square shapes, colors and relative positions. Call screenshot_diagram afterwards and compare with the original.
 - After a successful create_new_diagram or edit_diagram call, do not describe the diagram; the user sees it in the preview. One short sentence at most.
@@ -21,6 +21,7 @@
 import { createRequire } from "node:module"
 import { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js"
 import { StdioServerTransport } from "@modelcontextprotocol/sdk/server/stdio.js"
+import type { CallToolResult } from "@modelcontextprotocol/sdk/types.js"
 import open from "open"
 import { z } from "zod"
 import { expandCompactCells, foldCells } from "./compact-cells.ts"
@@ -99,6 +100,9 @@ installDomPolyfill()
 // Server configuration
 const config = {
    port: parseInt(process.env.PORT || "6002", 10),
+    // Attach a screenshot to every create_new_diagram and edit_diagram
+    // result unless the call says otherwise
+    autoScreenshot: process.env.DRAWIO_AUTO_SCREENSHOT === "true",
 }

 // Keep each session's latest diagram and its History on disk, so they
@@ -166,7 +170,7 @@ Before using cloud or icon shapes (AWS, Azure, GCP, Kubernetes, Cisco, BPMN...),

 To draw from a file, image or web page, read it yourself first; the server only receives XML. Open existing .drawio files with load_diagram.

-After drawing a complex diagram, call screenshot_diagram once to see it, and fix overlapping shapes or edges that cross shapes.
+After drawing or heavily editing a complex diagram, pass screenshot: true on that create_new_diagram or edit_diagram call (or call screenshot_diagram) to see the result, and fix overlapping shapes and edges that cross shapes.

 Tools:
 - create_new_diagram: draw a new diagram, replacing the whole document. Send only the mxCell elements of one page (the server adds the wrapper and root cells), or a full <mxfile> for several pages.
@@ -226,6 +230,14 @@ const pageSelectorSchema = {
        ),
 }

+// The optional screenshot of create_new_diagram and edit_diagram
+const screenshotSchema = z
+    .boolean()
+    .optional()
+    .describe(
+        "Also return a PNG of the result so you can check it for overlaps and edges crossing shapes. Needs the preview tab open and in front; adds 2 to 10 s. Default: the DRAWIO_AUTO_SCREENSHOT environment variable.",
+    )
+
 /**
 * Pull a clean PageSelector out of a tool's parsed input.
 * Returns an empty object when none of the page_* fields are set, so callers
@@ -502,7 +514,7 @@ registerWriteTool(
        title: "Create new diagram",
        description: `Create a NEW diagram, REPLACING the whole document: every page and any unsaved user changes (the previous state stays in History). To add a tab use add_page; to change cells use edit_diagram.

-Before using icon shapes (AWS, Azure, GCP, Kubernetes, Cisco...), call get_shape_library first. Follow the drawing guide returned by start_session (call get_drawing_guide if it is no longer in your context).
+Before using icon shapes (AWS, Azure, GCP, Kubernetes, Cisco...), call get_shape_library first. Follow the drawing guide returned by start_session (call get_drawing_guide if it is no longer in your context). Pass screenshot: true on the last call of a complex diagram to check the result.

 Accepted xml:
 1) Only the mxCell elements of one page (recommended). The server adds <mxfile>, <mxGraphModel>, <root> and the root cells "0" and "1":
@@ -519,10 +531,11 @@ To clear the canvas to one blank page, send only the two root cells <mxCell id="
                .describe(
                    "REQUIRED: the mxCell elements of one page, a bare <mxGraphModel>, or a full <mxfile> with one or more <diagram> pages.",
                ),
+            screenshot: screenshotSchema,
        },
        annotations: { openWorldHint: false },
    },
-    async ({ xml: inputXml }) => {
+    async ({ xml: inputXml, screenshot }) => {
        try {
            if (!currentSession) {
                return {
@@ -596,26 +609,36 @@ To clear the canvas to one blank page, send only the two root cells <mxCell id="
            log.info(`Diagram content set successfully (${pageSummary})`)

            // Only the root cells: the model asked for an empty canvas
-            // (several empty pages get the page summary below)
-            if (!hasCells(xml) && pages.length <= 1) {
-                return {
-                    content: [
-                        {
-                            type: "text",
-                            text: "Canvas cleared: one blank page. The previous diagram is in History.",
-                        },
-                    ],
-                }
-            }
-
-            return {
-                content: [
-                    {
-                        type: "text",
-                        text: `Diagram content set successfully!\n\nThe diagram is now visible in your browser.\n\nXML length: ${xml.length} characters\n${pageSummary}`,
-                    },
-                ],
+            // (several empty pages get the page summary)
+            const cleared = !hasCells(xml) && pages.length <= 1
+            const content: CallToolResult["content"] = [
+                {
+                    type: "text",
+                    text: cleared
+                        ? "Canvas cleared: one blank page. The previous diagram is in History."
+                        : `Diagram content set successfully!\n\nThe diagram is now visible in your browser.\n\nXML length: ${xml.length} characters\n${pageSummary}`,
+                },
+            ]
+            // The page on screen is the new document's first page
+            if (screenshot ?? config.autoScreenshot) {
+                const shot = await captureScreenshot(
+                    currentSession.id,
+                    xml,
+                    {},
+                    AFTER_WRITE_NEXT,
+                )
+                content.push(
+                    ...(shot.ok
+                        ? shot.content
+                        : [
+                              {
+                                  type: "text" as const,
+                                  text: `Screenshot skipped: ${shot.note}`,
+                              },
+                          ]),
+                )
            }
+            return { content }
        } catch (error) {
            const message =
                error instanceof Error ? error.message : String(error)
@@ -856,10 +879,11 @@ registerWriteTool(
                    }),
                )
                .describe("Array of operations to apply"),
+            screenshot: screenshotSchema,
        },
        annotations: { openWorldHint: false },
    },
-    async ({ operations, page_id, page_name, page_index }) => {
+    async ({ operations, page_id, page_name, page_index, screenshot }) => {
        try {
            if (!currentSession) {
                return {
@@ -998,14 +1022,36 @@ registerWriteTool(

            log.info(`Diagram edited successfully`)

-            return {
-                content: [
-                    {
-                        type: "text",
-                        text: `Diagram edited successfully!\n\nApplied ${outcome.applied} operation(s) on ${describeSelector(pageSelector)}.`,
-                    },
-                ],
+            const content: CallToolResult["content"] = [
+                {
+                    type: "text",
+                    text: `Diagram edited successfully!\n\nApplied ${outcome.applied} operation(s) on ${describeSelector(pageSelector)}.`,
+                },
+            ]
+            // Without a selector the edit went to the first page, while the
+            // page on screen may be another tab: name the page, so the PNG
+            // shows the edited one (a pageId export leaves the view alone)
+            if (screenshot ?? config.autoScreenshot) {
+                const shot = await captureScreenshot(
+                    currentSession.id,
+                    result,
+                    hasPageSelector(pageSelector)
+                        ? pageSelector
+                        : { page_index: 0 },
+                    AFTER_WRITE_NEXT,
+                )
+                content.push(
+                    ...(shot.ok
+                        ? shot.content
+                        : [
+                              {
+                                  type: "text" as const,
+                                  text: `Screenshot skipped: ${shot.note}`,
+                              },
+                          ]),
+                )
            }
+            return { content }
        } catch (error) {
            const message =
                error instanceof Error ? error.message : String(error)
@@ -1226,12 +1272,21 @@ function previewStalled(sessionId: string): boolean {
    return lastPolled !== undefined && Date.now() - lastPolled > 10_000
 }

+/**
+ * `next` is what the model does once the tab is in front. The default
+ * suits screenshot_diagram and export_diagram; a write tool whose edit
+ * already succeeded must not run again, so it points to screenshot_diagram.
+ */
+function previewStalledNote(sessionId: string, next = "then retry") {
+    return `The preview tab is not responding (browsers pause background tabs). Ask the user to bring the preview tab to the front (http://localhost:${getServerPort()}?mcp=${sessionId}), ${next}.`
+}
+
 function previewStalledError(sessionId: string) {
    return {
        content: [
            {
                type: "text" as const,
-                text: `Error: The preview tab is not responding (browsers pause background tabs). Ask the user to bring the preview tab to the front (http://localhost:${getServerPort()}?mcp=${sessionId}), then retry.`,
+                text: `Error: ${previewStalledNote(sessionId)}`,
            },
        ],
        isError: true,
@@ -1247,6 +1302,10 @@ function pageIdFor(xml: string, selector: PageSelector): string | undefined {
    )
 }

+// Stalled-tab advice after a successful create_new_diagram or edit_diagram:
+// "then retry" would make the model run the write again
+const AFTER_WRITE_NEXT = "then call screenshot_diagram to see the result"
+
 // Screenshot size: Claude Desktop caps a tool result at about 150,000
 // characters, so retry smaller above 140,000 base64 characters. (Claude
 // Code 2.1 accepted a 240,000 character image in testing.)
@@ -1260,8 +1319,90 @@ const SCREENSHOT_CHECKLIST = `Check this rendering of the diagram for:
 3. Text that is cut off, overlapping or too small to read (warning)
 4. Layout problems: cramped shapes, poor spacing or misalignment (warning)
 5. Rendering errors: missing, incomplete or broken elements, such as an icon that did not load (critical)
+Name the elements that have a problem, for example "the 'Login' box overlaps 'Register'".
+Give concrete fixes, for example "move 'Login' 50 px left".
 If there are critical issues, fix them with edit_diagram and take one more screenshot. Do at most two rounds of fixes. Minor cosmetic issues are fine, and diagrams with only 1 or 2 shapes pass unless something is clearly broken.`

+/**
+ * Have the preview tab render a page of xml as a PNG. Returns the image
+ * block and the checklist for the model, or the reason there is no image
+ * (a sentence without an "Error:" prefix: screenshot_diagram reports it as
+ * an error, the write tools as a skipped screenshot).
+ */
+async function captureScreenshot(
+    sessionId: string,
+    xml: string,
+    selector: PageSelector,
+    /** What to do once a stalled tab is in front (see previewStalledNote) */
+    next?: string,
+): Promise<
+    | { ok: true; content: CallToolResult["content"] }
+    | { ok: false; note: string }
+> {
+    if (previewStalled(sessionId)) {
+        return { ok: false, note: previewStalledNote(sessionId, next) }
+    }
+    // The tab never polled (right after start_session, or the state was
+    // just restored from disk), so an export would only time out
+    if (getState(sessionId)?.lastPolled === undefined) {
+        return {
+            ok: false,
+            note: "The preview tab is not open, or has not connected yet.",
+        }
+    }
+    if (!hasCells(xml)) {
+        return { ok: false, note: "The diagram is empty." }
+    }
+
+    let pageId: string | undefined
+    let projectionXml: string | undefined
+    if (hasPageSelector(selector)) {
+        const doc = normalizeToMxfile(xml) ?? xml
+        pageId = pageIdFor(doc, selector)
+        // A page without an id: load just that page and capture it, as
+        // export_diagram does
+        if (!pageId) {
+            const projection = projectPage(doc, selector)
+            if (!projection.ok) {
+                return {
+                    ok: false,
+                    note: `Page ${describeSelector(selector)} not found.`,
+                }
+            }
+            projectionXml = projection.xml
+        }
+    }
+
+    let data: string | undefined
+    for (const width of SCREENSHOT_WIDTHS) {
+        data = await exportViaBrowser(sessionId, "png", projectionXml, {
+            width,
+            pageId,
+        })
+        if (!data || data.length <= MAX_SCREENSHOT_CHARS) break
+    }
+    if (!data) {
+        return {
+            ok: false,
+            note: "Screenshot timed out. Make sure the preview tab is open and in front.",
+        }
+    }
+    return {
+        ok: true,
+        content: [
+            {
+                type: "image",
+                data: data.replace(/^data:image\/png;base64,/, ""),
+                mimeType: "image/png",
+            },
+            {
+                type: "text",
+                text: `Screenshot of ${hasPageSelector(selector) ? `page ${describeSelector(selector)}` : "the page on screen"}.\n\n${SCREENSHOT_CHECKLIST}`,
+            },
+        ],
+    }
+}
+
 // Tool: screenshot_diagram
 server.registerTool(
    "screenshot_diagram",
@@ -1288,80 +1429,20 @@ server.registerTool(
                    isError: true,
                }
            }
-            if (previewStalled(currentSession.id)) {
-                return previewStalledError(currentSession.id)
-            }
            const xml =
                sessionState(currentSession.id)?.xml || currentSession.xml
-            if (!hasCells(xml)) {
+            const shot = await captureScreenshot(
+                currentSession.id,
+                xml,
+                pickPageSelector({ page_id, page_name, page_index }),
+            )
+            if (!shot.ok) {
                return {
-                    content: [{ type: "text", text: "The diagram is empty." }],
-                }
-            }
-
-            const pageSelector = pickPageSelector({
-                page_id,
-                page_name,
-                page_index,
-            })
-            let pageId: string | undefined
-            let projectionXml: string | undefined
-            if (hasPageSelector(pageSelector)) {
-                const doc = normalizeToMxfile(xml) ?? xml
-                pageId = pageIdFor(doc, pageSelector)
-                // A page without an id: load just that page and capture
-                // it, as export_diagram does
-                if (!pageId) {
-                    const projection = projectPage(doc, pageSelector)
-                    if (!projection.ok) {
-                        return {
-                            content: [
-                                {
-                                    type: "text",
-                                    text: `Error: Page ${describeSelector(pageSelector)} not found.`,
-                                },
-                            ],
-                            isError: true,
-                        }
-                    }
-                    projectionXml = projection.xml
-                }
-            }
-
-            let data: string | undefined
-            // start_session may replace currentSession between the tries
-            const sessionId = currentSession.id
-            for (const width of SCREENSHOT_WIDTHS) {
-                data = await exportViaBrowser(sessionId, "png", projectionXml, {
-                    width,
-                    pageId,
-                })
-                if (!data || data.length <= MAX_SCREENSHOT_CHARS) break
-            }
-            if (!data) {
-                return {
-                    content: [
-                        {
-                            type: "text",
-                            text: "Error: Screenshot timed out. Make sure the preview tab is open and in front.",
-                        },
-                    ],
+                    content: [{ type: "text", text: `Error: ${shot.note}` }],
                    isError: true,
                }
            }
-            return {
-                content: [
-                    {
-                        type: "image",
-                        data: data.replace(/^data:image\/png;base64,/, ""),
-                        mimeType: "image/png",
-                    },
-                    {
-                        type: "text",
-                        text: `Screenshot of ${hasPageSelector(pageSelector) ? `page ${describeSelector(pageSelector)}` : "the page on screen"}.\n\n${SCREENSHOT_CHECKLIST}`,
-                    },
-                ],
-            }
+            return { content: shot.content }
        } catch (error) {
            const message =
                error instanceof Error ? error.message : String(error)
@@ -300,11 +300,13 @@ async function poll() {
        // The tab's own push still on its way is not loaded back: the
        // canvas may have moved on since (an undo), and its answer follows
        const ownPush = pushesInFlight.includes(s.xml);
+        let justLoaded = false; // this poll put a new version on the canvas
        if ((forceReload || (s.version > currentVersion && !projectionExportActive && !ownPush)) && s.xml) {
            forceReload = false;
            projectionExportActive = false;
            currentVersion = s.version;
            loadDiagram(s.xml, true);
+            justLoaded = true;
        }
        // Handle sync request - server needs fresh state. After the load
        // above, so draw.io exports what it just loaded; never while a
@@ -354,6 +356,10 @@ async function poll() {
                // Let draw.io render the loaded page before exporting
                // (same proven settle delay as the AI-preview path).
                setTimeout(fireExport, 600);
+            } else if (justLoaded) {
+                // A write with a screenshot: give the new diagram's external
+                // icon images a moment to load before the PNG is taken
+                setTimeout(fireExport, 600);
            } else {
                fireExport();
            }
@@ -150,3 +150,66 @@ describe("prepareNewDiagram with cut-off output", () => {
        expect(out.error).not.toContain("cut off")
    })
 })
+
+describe("truncation check with named styles and compact cells", () => {
+    const BLUE =
+        '<mxStyle name="blue" value="fillColor=#dae8fc;strokeColor=#6c8ebf;"/>\n'
+    const A =
+        '<mxCell id="2" value="A" style="blue;" x="0" y="0" w="80" h="40"/>\n'
+
+    it("accepts complete compact cells after the definitions and expands them", () => {
+        const out = prepareNewDiagram(
+            `${BLUE}${A}<mxCell id="3" value="B" style="blue;" x="120" y="0" w="80" h="40"/>`,
+        )
+        expect(out.ok).toBe(true)
+        if (!out.ok) return
+        expect(out.xml).toContain("fillColor=#dae8fc")
+        expect(out.xml).toContain("<mxGeometry")
+        expect(out.xml).not.toContain("<mxStyle")
+    })
+
+    it("reports a cut inside a compact cell that follows the definitions", () => {
+        const cut = `${BLUE}${A}<mxCell id="3" value="B" style="blue;" x="120" y="0" w="8`
+        const out = prepareNewDiagram(cut)
+        expect(out.ok).toBe(false)
+        if (out.ok) return
+        expect(out.error).toContain("cut off")
+        expect(out.error).toContain(
+            '<mxCell id="3" value="B" style="blue;" x="120" y="0" w="8',
+        )
+        expect(out.error).toContain("edit_diagram add")
+    })
+
+    it("reports output that stops right after the definitions and the first cells", () => {
+        const cut = `${BLUE}${A}<mxCell id="3" value="B" style="blue;"`
+        expect(isMxCellXmlComplete(cut)).toBe(false)
+        expect(truncatedCellError(cut)).toContain("cut off")
+    })
+
+    it("does not treat a compact cell's quote slip as a cut", () => {
+        const slip = `${A}<mxCell id="3" value="B" x="120" y="0" w="80" h="40/>`
+        expect(isMxCellXmlComplete(slip)).toBe(true)
+        expect(truncatedCellError(slip)).toBeNull()
+        const out = prepareNewDiagram(slip)
+        expect(out.ok).toBe(true)
+        if (out.ok) expect(out.xml).toContain('height="40"')
+    })
+
+    it("leaves definitions without any cell to the style check", () => {
+        expect(truncatedCellError(BLUE)).toBeNull()
+        const out = prepareNewDiagram(BLUE)
+        expect(out.ok).toBe(false)
+        if (out.ok) return
+        expect(out.error).not.toContain("cut off")
+        expect(out.error).toContain("no cells")
+    })
+
+    it("leaves an unclosed definition to the style check", () => {
+        const cut = `${BLUE}<mxStyle name="flow" value="edgeSt`
+        const out = prepareNewDiagram(cut)
+        expect(out.ok).toBe(false)
+        if (out.ok) return
+        expect(out.error).not.toContain("cut off")
+        expect(out.error).toContain("not closed")
+    })
+})
@@ -176,6 +176,17 @@ describe("MCP server wiring", () => {
        expect(line).not.toContain("(current)")
    })

+    it("advertises screenshot on create_new_diagram and edit_diagram", async () => {
+        const resp = await send("tools/list", {})
+        for (const name of ["create_new_diagram", "edit_diagram"]) {
+            const tool = resp.result.tools.find(
+                (t: { name: string }) => t.name === name,
+            )
+            const props = tool?.inputSchema?.properties ?? {}
+            expect(props.screenshot?.type, `${name} screenshot`).toBe("boolean")
+        }
+    })
+
    it("advertises name/id/xml on add_page", async () => {
        const resp = await send("tools/list", {})
        const addPage = resp.result.tools.find(