mirror of
https://github.com/DayuanJiang/next-ai-draw-io.git
synced 2026-10-10 03:29:50 +08:00
feat(mcp-server): add screenshot_diagram so the model can check its render
- New read-only screenshot_diagram tool returns the rendered page as a PNG plus the web app's visual checklist (overlaps, edges crossing shapes, readability, layout, rendering errors), replacing the web app's separate vision model with the host model's own vision - PNG exports use draw.io's width and pageId options: screenshots stay under ~140,000 base64 characters and page exports no longer swap the page on screen - Fail fast with a clear message when the preview tab stopped polling (browsers throttle background tabs) - Mention the screenshot step in the drawing guide and instructions
This commit is contained in:
dayuan.jiang
Generated
+2
-2
@@ -1,12 +1,12 @@
|
||||
{
|
||||
"name": "@next-ai-drawio/mcp-server",
|
||||
"version": "0.3.0",
|
||||
"version": "0.4.0",
|
||||
"lockfileVersion": 3,
|
||||
"requires": true,
|
||||
"packages": {
|
||||
"": {
|
||||
"name": "@next-ai-drawio/mcp-server",
|
||||
"version": "0.3.0",
|
||||
"version": "0.4.0",
|
||||
"license": "Apache-2.0",
|
||||
"dependencies": {
|
||||
"@modelcontextprotocol/sdk": "^1.31.0",
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
{
|
||||
"name": "@next-ai-drawio/mcp-server",
|
||||
"version": "0.3.0",
|
||||
"version": "0.4.0",
|
||||
"description": "MCP server for Next AI Draw.io - AI-powered diagram generation with real-time browser preview",
|
||||
"type": "module",
|
||||
"main": "dist/index.js",
|
||||
|
||||
@@ -14,6 +14,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.
|
||||
- When replicating a diagram from an image, match its style and layout closely: straight or curved lines, rounded or square shapes.
|
||||
- The preview page has History (it saves a snapshot before every AI change and can restore any of the last 20 versions) and Download. You can make changes freely; nothing is lost.
|
||||
|
||||
|
||||
@@ -102,9 +102,19 @@ interface SessionState {
|
||||
syncRequested?: number // Timestamp when sync requested, cleared when browser responds
|
||||
exportFormat?: "png" | "svg" // Set by MCP tool to request browser export
|
||||
exportXml?: string // Single-page projection to load before a page-targeted export
|
||||
exportOptions?: ExportOptions // Extra draw.io export parameters (PNG only)
|
||||
exportData?: string // Base64/SVG data returned by browser after export
|
||||
}
|
||||
|
||||
/**
|
||||
* draw.io's PNG export takes these directly: width caps the image size
|
||||
* (never upscales), pageId renders a page other than the one on screen.
|
||||
*/
|
||||
export interface ExportOptions {
|
||||
width?: number
|
||||
pageId?: string
|
||||
}
|
||||
|
||||
export const stateStore = new Map<string, SessionState>()
|
||||
|
||||
let server: http.Server | null = null
|
||||
@@ -134,6 +144,7 @@ export function setState(
|
||||
syncRequested: undefined, // Clear sync request when browser pushes state
|
||||
exportFormat: existing?.exportFormat, // Preserve pending export request
|
||||
exportXml: existing?.exportXml, // Preserve pending projection
|
||||
exportOptions: existing?.exportOptions,
|
||||
exportData: existing?.exportData, // Preserve export result
|
||||
})
|
||||
log.debug(`State updated: session=${sessionId}, version=${newVersion}`)
|
||||
@@ -155,11 +166,13 @@ export function requestExport(
|
||||
sessionId: string,
|
||||
format: "png" | "svg",
|
||||
projectionXml?: string,
|
||||
options?: ExportOptions,
|
||||
): boolean {
|
||||
const state = stateStore.get(sessionId)
|
||||
if (!state) return false
|
||||
state.exportData = undefined
|
||||
state.exportXml = projectionXml
|
||||
state.exportOptions = options
|
||||
state.exportFormat = format
|
||||
return true
|
||||
}
|
||||
@@ -256,6 +269,10 @@ export function shutdown(): void {
|
||||
stopHttpServer()
|
||||
}
|
||||
|
||||
export function getServerPort(): number {
|
||||
return serverPort
|
||||
}
|
||||
|
||||
function handleRequest(
|
||||
req: http.IncomingMessage,
|
||||
res: http.ServerResponse,
|
||||
@@ -379,6 +396,7 @@ function handleStateApi(
|
||||
syncRequested: !!state?.syncRequested,
|
||||
exportFormat: state?.exportFormat || null,
|
||||
exportXml: state?.exportXml || null,
|
||||
exportOptions: state?.exportOptions || null,
|
||||
}),
|
||||
)
|
||||
} else if (req.method === "POST") {
|
||||
@@ -401,6 +419,7 @@ function handleStateApi(
|
||||
state.exportData = data.exportData
|
||||
state.exportFormat = undefined
|
||||
state.exportXml = undefined
|
||||
state.exportOptions = undefined
|
||||
log.debug(
|
||||
`Export data received for session=${sessionId}`,
|
||||
)
|
||||
@@ -1006,10 +1025,13 @@ function getHtmlPage(sessionId: string): string {
|
||||
// projection is showing (see projectionExportActive guard).
|
||||
if (s.exportFormat && !pendingMcpExport && isReady) {
|
||||
pendingMcpExport = s.exportFormat;
|
||||
const extra = s.exportOptions || {};
|
||||
const fireExport = () => {
|
||||
// mcpExport is echoed back in msg.message (see the handler)
|
||||
// mcpExport is echoed back in msg.message (see the
|
||||
// handler). PNG: width caps the size, pageId picks a
|
||||
// page; without one draw.io would use the first page.
|
||||
const exportOpts = pendingMcpExport === 'png'
|
||||
? { action: 'export', format: 'png', scale: 2, currentPage: true, mcpExport: true }
|
||||
? { action: 'export', format: 'png', scale: 2, currentPage: !extra.pageId, ...extra, mcpExport: true }
|
||||
: { action: 'export', format: 'svg', mcpExport: true };
|
||||
iframe.contentWindow.postMessage(JSON.stringify(exportOpts), '*');
|
||||
};
|
||||
|
||||
@@ -31,6 +31,8 @@ import { editDiagram, targetPageXml } from "./edit-diagram.js"
|
||||
import { checkEditGate } from "./edit-gate.js"
|
||||
import { addHistory } from "./history.js"
|
||||
import {
|
||||
type ExportOptions,
|
||||
getServerPort,
|
||||
getState,
|
||||
requestExport,
|
||||
requestSync,
|
||||
@@ -44,6 +46,7 @@ import { log } from "./logger.js"
|
||||
import {
|
||||
addPageToDoc,
|
||||
deletePageFromDoc,
|
||||
findPageElement,
|
||||
hasPageSelector,
|
||||
listPagesFromDoc,
|
||||
normalizeToMxfile,
|
||||
@@ -95,10 +98,13 @@ Start with start_session: it opens the preview and its result contains the drawi
|
||||
|
||||
Before using cloud or icon shapes (AWS, Azure, GCP, Kubernetes, Cisco, BPMN...), call get_shape_library and use the exact style names it returns. Never guess icon style names.
|
||||
|
||||
After drawing a complex diagram, call screenshot_diagram once to see it, and fix overlapping shapes or 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.
|
||||
- edit_diagram: add, update or delete cells of an existing page by id. All-or-nothing; a rejected call includes the current XML so you can retry.
|
||||
- get_diagram: read the current XML, including the user's manual edits.
|
||||
- screenshot_diagram: see the rendered diagram as an image.
|
||||
- load_diagram, export_diagram: open or save .drawio files and export .png or .svg. Use absolute paths.
|
||||
- list_pages, add_page, rename_page, delete_page: manage pages (tabs).`
|
||||
|
||||
@@ -885,9 +891,10 @@ function exportViaBrowser(
|
||||
sessionId: string,
|
||||
format: "png" | "svg",
|
||||
projectionXml?: string,
|
||||
options?: ExportOptions,
|
||||
): Promise<string | undefined> {
|
||||
const run = exportQueue.then(async () => {
|
||||
requestExport(sessionId, format, projectionXml)
|
||||
requestExport(sessionId, format, projectionXml, options)
|
||||
|
||||
// A projection export does an extra load + render round-trip in the
|
||||
// browser, so give it a longer window. Re-read the live store entry
|
||||
@@ -914,6 +921,160 @@ function exportViaBrowser(
|
||||
return run
|
||||
}
|
||||
|
||||
/**
|
||||
* True when the preview tab polled before but has gone quiet. Browsers
|
||||
* slow down timers in background tabs (Chrome: about once a minute after
|
||||
* 5 minutes hidden), so an export would just time out.
|
||||
*/
|
||||
function previewStalled(sessionId: string): boolean {
|
||||
const lastPolled = getState(sessionId)?.lastPolled
|
||||
return lastPolled !== undefined && Date.now() - lastPolled > 10_000
|
||||
}
|
||||
|
||||
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.`,
|
||||
},
|
||||
],
|
||||
isError: true,
|
||||
}
|
||||
}
|
||||
|
||||
/** The <diagram id> of the page a selector targets, for draw.io's pageId. */
|
||||
function pageIdFor(xml: string, selector: PageSelector): string | undefined {
|
||||
const doc = parseMxfile(xml)
|
||||
return (
|
||||
(doc && findPageElement(doc, selector)?.element.getAttribute("id")) ||
|
||||
undefined
|
||||
)
|
||||
}
|
||||
|
||||
// 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.)
|
||||
const SCREENSHOT_WIDTHS = [1000, 700]
|
||||
const MAX_SCREENSHOT_CHARS = 140_000
|
||||
|
||||
// Adapted from the web app's vision check (lib/validation-prompts.ts)
|
||||
const SCREENSHOT_CHECKLIST = `Check this rendering of the diagram for:
|
||||
1. Overlapping shapes that cover each other or their labels (critical)
|
||||
2. Edges crossing shapes that are not their source or target (critical)
|
||||
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)
|
||||
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.`
|
||||
|
||||
// Tool: screenshot_diagram
|
||||
server.registerTool(
|
||||
"screenshot_diagram",
|
||||
{
|
||||
title: "Screenshot diagram",
|
||||
description:
|
||||
"Render the diagram in the preview and return it as a PNG image, so you can see your own result. " +
|
||||
"Call this once after drawing or heavily editing a complex diagram, then fix overlaps and edges that cross shapes. " +
|
||||
"Without a page selector it shows the page on screen. Needs the preview tab to be open.",
|
||||
inputSchema: { ...pageSelectorSchema },
|
||||
annotations: { readOnlyHint: true, openWorldHint: false },
|
||||
},
|
||||
async (input) => {
|
||||
const { page_id, page_name, page_index } = input ?? {}
|
||||
try {
|
||||
if (!currentSession) {
|
||||
return {
|
||||
content: [
|
||||
{
|
||||
type: "text",
|
||||
text: "Error: No active session. Please call start_session first.",
|
||||
},
|
||||
],
|
||||
isError: true,
|
||||
}
|
||||
}
|
||||
if (previewStalled(currentSession.id)) {
|
||||
return previewStalledError(currentSession.id)
|
||||
}
|
||||
const xml = getState(currentSession.id)?.xml || currentSession.xml
|
||||
// Any cell besides the root cells "0" and "1"
|
||||
if (
|
||||
!/<(mxCell\b[^>]*\bid="(?![01]")|UserObject\b|object\b)/.test(
|
||||
xml,
|
||||
)
|
||||
) {
|
||||
return {
|
||||
content: [{ type: "text", text: "The diagram is empty." }],
|
||||
}
|
||||
}
|
||||
|
||||
const pageSelector = pickPageSelector({
|
||||
page_id,
|
||||
page_name,
|
||||
page_index,
|
||||
})
|
||||
let pageId: string | undefined
|
||||
if (hasPageSelector(pageSelector)) {
|
||||
pageId = pageIdFor(normalizeToMxfile(xml) ?? xml, pageSelector)
|
||||
if (!pageId) {
|
||||
return {
|
||||
content: [
|
||||
{
|
||||
type: "text",
|
||||
text: `Error: Page ${describeSelector(pageSelector)} not found.`,
|
||||
},
|
||||
],
|
||||
isError: true,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
let data: string | undefined
|
||||
for (const width of SCREENSHOT_WIDTHS) {
|
||||
data = await exportViaBrowser(
|
||||
currentSession.id,
|
||||
"png",
|
||||
undefined,
|
||||
{ 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.",
|
||||
},
|
||||
],
|
||||
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}`,
|
||||
},
|
||||
],
|
||||
}
|
||||
} catch (error) {
|
||||
const message =
|
||||
error instanceof Error ? error.message : String(error)
|
||||
log.error("screenshot_diagram failed:", message)
|
||||
return {
|
||||
content: [{ type: "text", text: `Error: ${message}` }],
|
||||
isError: true,
|
||||
}
|
||||
}
|
||||
},
|
||||
)
|
||||
|
||||
// Tool: export_diagram
|
||||
server.registerTool(
|
||||
"export_diagram",
|
||||
@@ -926,7 +1087,8 @@ server.registerTool(
|
||||
"- .drawio with NO page selector: writes the full <mxfile> (all pages).\n" +
|
||||
"- .drawio with a page selector: writes a single-page <mxfile> containing only that page.\n" +
|
||||
"- .png / .svg with NO page selector: exports the currently active page in the browser.\n" +
|
||||
"- .png / .svg with a page selector: temporarily loads a single-page projection of that page into the browser, captures the rendered image, then restores the full document. The user will see a brief tab-flicker (~1-2s) but the exported image is guaranteed to be the requested page.",
|
||||
"- .png with a page selector: renders that page without changing what the user sees.\n" +
|
||||
"- .svg with a page selector: temporarily loads that page into the browser, captures it, then restores the full document (the user sees a brief flicker).",
|
||||
inputSchema: {
|
||||
...pageSelectorSchema,
|
||||
path: z
|
||||
@@ -1058,6 +1220,9 @@ server.registerTool(
|
||||
isError: true,
|
||||
}
|
||||
}
|
||||
if (previewStalled(currentSession.id)) {
|
||||
return previewStalledError(currentSession.id)
|
||||
}
|
||||
|
||||
// -----------------------------------------------------------------
|
||||
// Page-targeted PNG/SVG export.
|
||||
@@ -1070,8 +1235,15 @@ server.registerTool(
|
||||
// browser-side. The canonical session state is never mutated here,
|
||||
// so there is no restore race and no concurrent-edit clobbering.
|
||||
// -----------------------------------------------------------------
|
||||
// PNG: draw.io renders any page by id, without touching the
|
||||
// page on screen. SVG export has no page option, so it still
|
||||
// needs the projection below.
|
||||
let projectionXml: string | undefined
|
||||
if (hasPageSelector(pageSelector)) {
|
||||
let pngPageId: string | undefined
|
||||
if (hasPageSelector(pageSelector) && detectedFormat === "png") {
|
||||
pngPageId = pageIdFor(currentSession.xml, pageSelector)
|
||||
}
|
||||
if (hasPageSelector(pageSelector) && !pngPageId) {
|
||||
const projection = projectPage(currentSession.xml, pageSelector)
|
||||
if (!projection.ok) {
|
||||
return {
|
||||
@@ -1094,6 +1266,7 @@ server.registerTool(
|
||||
currentSession.id,
|
||||
detectedFormat as "png" | "svg",
|
||||
projectionXml,
|
||||
pngPageId ? { pageId: pngPageId } : undefined,
|
||||
)
|
||||
|
||||
if (!exportData) {
|
||||
|
||||
@@ -11,6 +11,7 @@ import { afterAll, beforeAll, describe, expect, it } from "vitest"
|
||||
import { addHistory, getHistory } from "../src/history.js"
|
||||
import {
|
||||
getState,
|
||||
requestExport,
|
||||
requestSync,
|
||||
setState,
|
||||
shutdown,
|
||||
@@ -233,6 +234,25 @@ describe("POST /api/state", () => {
|
||||
})
|
||||
})
|
||||
|
||||
describe("export requests", () => {
|
||||
it("hands draw.io export options to the page and clears them after", async () => {
|
||||
const id = "mcp-export-options"
|
||||
setState(id, "<mxfile>x</mxfile>")
|
||||
requestExport(id, "png", undefined, { width: 1000, pageId: "p2" })
|
||||
const poll = JSON.parse(
|
||||
(await request(`/api/state?sessionId=${id}`)).body,
|
||||
)
|
||||
expect(poll.exportFormat).toBe("png")
|
||||
expect(poll.exportOptions).toEqual({ width: 1000, pageId: "p2" })
|
||||
|
||||
await postJson("/api/state", {
|
||||
sessionId: id,
|
||||
exportData: "data:image/png;base64,AAAA",
|
||||
})
|
||||
expect(getState(id)?.exportOptions).toBeUndefined()
|
||||
})
|
||||
})
|
||||
|
||||
describe("preview page", () => {
|
||||
it("serves a script that parses", async () => {
|
||||
const res = await request("/?mcp=mcp-test-script")
|
||||
|
||||
@@ -41,6 +41,7 @@ const EXPECTED_TOOLS = [
|
||||
"delete_page",
|
||||
"get_drawing_guide",
|
||||
"get_shape_library",
|
||||
"screenshot_diagram",
|
||||
]
|
||||
|
||||
// Claude Code truncates tool descriptions and server instructions here
|
||||
|
||||
Reference in New Issue
Block a user