From 07f3e3c2a3ee167e12c2d580d338d5eb5f578584 Mon Sep 17 00:00:00 2001 From: Dayuan Jiang <34411969+DayuanJiang@users.noreply.github.com> Date: Sun, 11 Oct 2026 20:54:17 +0900 Subject: [PATCH] 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 --- README.md | 1 + docs/cn/README_CN.md | 1 + docs/ja/README_JA.md | 1 + packages/claude-plugin/README.md | 1 + packages/mcp-server/README.md | 3 +- packages/mcp-server/src/drawing-guide.ts | 2 +- packages/mcp-server/src/index.ts | 279 +++++++++++------- packages/mcp-server/src/preview/preview.js | 6 + packages/mcp-server/tests/new-diagram.test.ts | 63 ++++ .../mcp-server/tests/server-wiring.test.ts | 11 + 10 files changed, 267 insertions(+), 101 deletions(-) diff --git a/README.md b/README.md index de45c5cd..bc9731ea 100644 --- a/README.md +++ b/README.md @@ -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. diff --git a/docs/cn/README_CN.md b/docs/cn/README_CN.md index 5178b32c..e934e81e 100644 --- a/docs/cn/README_CN.md +++ b/docs/cn/README_CN.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)。 diff --git a/docs/ja/README_JA.md b/docs/ja/README_JA.md index 68c41ea9..740a6dd7 100644 --- a/docs/ja/README_JA.md +++ b/docs/ja/README_JA.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) を参照してください。 diff --git a/packages/claude-plugin/README.md b/packages/claude-plugin/README.md index 0f41d5b1..ffb64868 100644 --- a/packages/claude-plugin/README.md +++ b/packages/claude-plugin/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 diff --git a/packages/mcp-server/README.md b/packages/mcp-server/README.md index fc97d7a2..5c2e2a84 100644 --- a/packages/mcp-server/README.md +++ b/packages/mcp-server/README.md @@ -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) diff --git a/packages/mcp-server/src/drawing-guide.ts b/packages/mcp-server/src/drawing-guide.ts index e8c6cdef..25d4c1f3 100644 --- a/packages/mcp-server/src/drawing-guide.ts +++ b/packages/mcp-server/src/drawing-guide.ts @@ -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. diff --git a/packages/mcp-server/src/index.ts b/packages/mcp-server/src/index.ts index f02ae8a7..905687ca 100644 --- a/packages/mcp-server/src/index.ts +++ b/packages/mcp-server/src/index.ts @@ -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 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 , , 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 , or a full with one or more 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 { + 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) diff --git a/packages/mcp-server/src/preview/preview.js b/packages/mcp-server/src/preview/preview.js index a82f85ba..6d79e132 100644 --- a/packages/mcp-server/src/preview/preview.js +++ b/packages/mcp-server/src/preview/preview.js @@ -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(); } diff --git a/packages/mcp-server/tests/new-diagram.test.ts b/packages/mcp-server/tests/new-diagram.test.ts index b935ba4f..9a70b44a 100644 --- a/packages/mcp-server/tests/new-diagram.test.ts +++ b/packages/mcp-server/tests/new-diagram.test.ts @@ -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 = + '\n' + const A = + '\n' + + it("accepts complete compact cells after the definitions and expands them", () => { + const out = prepareNewDiagram( + `${BLUE}${A}`, + ) + expect(out.ok).toBe(true) + if (!out.ok) return + expect(out.xml).toContain("fillColor=#dae8fc") + expect(out.xml).toContain(" { + const cut = `${BLUE}${A} { + const cut = `${BLUE}${A} { + const slip = `${A} { + 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} { 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(