From 6d67a0ec695a76d6fa0774e5e4aa2a56e2c09b5a Mon Sep 17 00:00:00 2001 From: "dayuan.jiang" Date: Sat, 3 Oct 2026 22:03:44 +0900 Subject: [PATCH] fix(mcp-server): make edit_diagram all-or-nothing and fix preview sync races - edit_diagram applies nothing when any operation fails, rejects invalid or multi-cell new_xml, validates only the target page, and returns the current page XML on every rejection (including stale edits) - Fix get_diagram reading the old diagram right after an AI write: the preview pushed its sync reply with a newer version than it was taken at - Keep a user edit that loses the race with an AI write in history and tell the user in the preview - Autofix removes only exact foreign tags (a stray deleted ), fixes tag case, drops orphan s, and rejects unknown element names in model XML - Edit empty and compressed pages; PNG exports use the page on screen; tag download exports; reload from the server after a page export - Expand ~ in paths, tell the model when the browser sync timed out, use registerPrompt, require SDK ^1.31.0 --- packages/mcp-server/package-lock.json | 6 +- packages/mcp-server/package.json | 4 +- packages/mcp-server/src/diagram-operations.ts | 43 +++- packages/mcp-server/src/edit-diagram.ts | 103 ++++++++ packages/mcp-server/src/http-server.ts | 89 ++++--- packages/mcp-server/src/index.ts | 149 +++++------- packages/mcp-server/src/pages.ts | 9 +- packages/mcp-server/src/xml-validation.ts | 225 ++++++++++++------ .../tests/diagram-operations.test.ts | 41 ++++ .../mcp-server/tests/edit-diagram.test.ts | 130 ++++++++++ packages/mcp-server/tests/http-server.test.ts | 53 ++++- .../mcp-server/tests/server-wiring.test.ts | 2 +- .../mcp-server/tests/xml-validation.test.ts | 62 +++++ 13 files changed, 712 insertions(+), 204 deletions(-) create mode 100644 packages/mcp-server/src/edit-diagram.ts create mode 100644 packages/mcp-server/tests/edit-diagram.test.ts diff --git a/packages/mcp-server/package-lock.json b/packages/mcp-server/package-lock.json index 5d84a11b..d45f49fc 100644 --- a/packages/mcp-server/package-lock.json +++ b/packages/mcp-server/package-lock.json @@ -1,15 +1,15 @@ { "name": "@next-ai-drawio/mcp-server", - "version": "0.2.3", + "version": "0.2.4", "lockfileVersion": 3, "requires": true, "packages": { "": { "name": "@next-ai-drawio/mcp-server", - "version": "0.2.3", + "version": "0.2.4", "license": "Apache-2.0", "dependencies": { - "@modelcontextprotocol/sdk": "^1.0.4", + "@modelcontextprotocol/sdk": "^1.31.0", "linkedom": "^0.18.0", "open": "^11.0.0", "saxes": "^6.0.0", diff --git a/packages/mcp-server/package.json b/packages/mcp-server/package.json index 702089d7..df4060bd 100644 --- a/packages/mcp-server/package.json +++ b/packages/mcp-server/package.json @@ -1,6 +1,6 @@ { "name": "@next-ai-drawio/mcp-server", - "version": "0.2.3", + "version": "0.2.4", "description": "MCP server for Next AI Draw.io - AI-powered diagram generation with real-time browser preview", "type": "module", "main": "dist/index.js", @@ -38,7 +38,7 @@ "access": "public" }, "dependencies": { - "@modelcontextprotocol/sdk": "^1.0.4", + "@modelcontextprotocol/sdk": "^1.31.0", "linkedom": "^0.18.0", "open": "^11.0.0", "saxes": "^6.0.0", diff --git a/packages/mcp-server/src/diagram-operations.ts b/packages/mcp-server/src/diagram-operations.ts index 6e40b9f6..cc963a3c 100644 --- a/packages/mcp-server/src/diagram-operations.ts +++ b/packages/mcp-server/src/diagram-operations.ts @@ -8,6 +8,7 @@ */ import { getXmlSyntaxError } from "./dom.js" +import { decompressPageContent } from "./load-diagram.js" import { log } from "./logger.js" import { findPageElement, hasPageSelector, type PageSelector } from "./pages.js" @@ -33,6 +34,44 @@ export interface ApplyOperationsResult { // on the wrapper, so the wrapper is treated as the cell. const CELL_SELECTOR = "mxCell, UserObject, object" +/** + * Return the of a page, creating it when missing. An empty + * page gets a blank model with the "0" and "1" root cells; a page whose text + * is draw.io's compressed format is decompressed in place. Returns null if + * the text is neither empty nor decompressible. + */ +function ensurePageRoot(doc: Document, page: Element): Element | null { + const existing = page.querySelector("root") + if (existing) return existing + + let model = page.querySelector("mxGraphModel") + if (!model) { + const text = page.textContent?.trim() ?? "" + if (text) { + const xml = decompressPageContent(text) + if (!xml || getXmlSyntaxError(xml)) return null + const parsed = new DOMParser().parseFromString(xml, "text/xml") + if (parsed.documentElement?.tagName !== "mxGraphModel") return null + page.textContent = "" + model = page.appendChild( + doc.importNode(parsed.documentElement, true), + ) as Element + const decompressedRoot = model.querySelector("root") + if (decompressedRoot) return decompressedRoot + } else { + model = page.appendChild(doc.createElement("mxGraphModel")) + } + } + + const blank = new DOMParser().parseFromString( + ``, + "text/xml", + ) + return model.appendChild( + doc.importNode(blank.documentElement, true), + ) as Element +} + /** Read parent/source/target, which a wrapped cell keeps on its inner mxCell. */ function cellAttr(cell: Element, name: string): string | null { const inner = @@ -98,7 +137,7 @@ export function applyDiagramOperations( ], } } - root = found.element.querySelector("root") + root = ensurePageRoot(doc as unknown as Document, found.element) if (!root) { const pageId = found.element.getAttribute("id") || `(index ${found.index})` @@ -108,7 +147,7 @@ export function applyDiagramOperations( { type: "update", cellId: "", - message: `Page "${pageId}" has no element`, + message: `Page "${pageId}" has no element and its content could not be decompressed`, }, ], } diff --git a/packages/mcp-server/src/edit-diagram.ts b/packages/mcp-server/src/edit-diagram.ts new file mode 100644 index 00000000..f747f5a6 --- /dev/null +++ b/packages/mcp-server/src/edit-diagram.ts @@ -0,0 +1,103 @@ +/** + * Core of the edit_diagram tool, kept free of session state so it can be + * unit tested. + * + * All-or-nothing, like the web app (hooks/use-diagram-tool-handlers.ts): + * if any operation fails, nothing is written and every failure is reported, + * so the model never builds on a half-applied edit. + */ + +import { + applyDiagramOperations, + type DiagramOperation, +} from "./diagram-operations.js" +import { type PageSelector, projectPage } from "./pages.js" +import { validateAndFixXml, validateMxCellStructure } from "./xml-validation.js" + +export type EditOutcome = + | { ok: true; xml: string; applied: number; fixes: string[] } + | { ok: false; errors: string[]; pageError: boolean } + +/** Number of top-level elements in an XML fragment. */ +function countTopLevelElements(fragment: string): number { + const doc = new DOMParser().parseFromString( + `${fragment}`, + "text/xml", + ) + return doc.documentElement?.children.length ?? 0 +} + +/** The target page as a one-page , or the whole document. */ +export function targetPageXml(xml: string, selector: PageSelector): string { + const projection = projectPage(xml, selector) + return projection.ok ? projection.xml : xml +} + +export function editDiagram( + xml: string, + operations: DiagramOperation[], + selector: PageSelector, +): EditOutcome { + const errors: string[] = [] + const fixes: string[] = [] + const prepared: DiagramOperation[] = [] + + for (const op of operations) { + if (op.operation === "delete" || !op.new_xml) { + prepared.push(op) + continue + } + // Checked before validation: several cells fail the strict parser + // with a misleading "only one root" syntax error. + if (countTopLevelElements(op.new_xml) > 1) { + errors.push( + `${op.operation} ${op.cell_id}: new_xml must contain exactly one cell; use one add operation per cell`, + ) + continue + } + const check = validateAndFixXml(op.new_xml) + if (!check.valid) { + errors.push( + `${op.operation} ${op.cell_id}: invalid new_xml: ${check.error}`, + ) + continue + } + if (check.fixed) { + fixes.push(`${op.cell_id}: ${check.fixes.join(", ")}`) + } + prepared.push({ ...op, new_xml: check.fixed ?? op.new_xml }) + } + if (errors.length > 0) return { ok: false, errors, pageError: false } + + const { result, errors: opErrors } = applyDiagramOperations( + xml, + prepared, + selector, + ) + // An empty cellId means the page itself could not be edited + const pageLevel = opErrors.find((e) => e.cellId === "") + if (pageLevel) { + return { ok: false, errors: [pageLevel.message], pageError: true } + } + if (opErrors.length > 0) { + return { + ok: false, + errors: opErrors.map((e) => `${e.type} ${e.cellId}: ${e.message}`), + pageError: false, + } + } + + // Validate only the target page, and reject only errors this edit + // introduced: problems already in other pages or in a loaded file must + // not block every edit. + const after = validateMxCellStructure(targetPageXml(result, selector)) + if (after && !validateMxCellStructure(targetPageXml(xml, selector))) { + return { + ok: false, + errors: [`the edit would make the page invalid: ${after}`], + pageError: false, + } + } + + return { ok: true, xml: result, applied: operations.length, fixes } +} diff --git a/packages/mcp-server/src/http-server.ts b/packages/mcp-server/src/http-server.ts index 08b50bca..c56a7009 100644 --- a/packages/mcp-server/src/http-server.ts +++ b/packages/mcp-server/src/http-server.ts @@ -256,10 +256,6 @@ export function shutdown(): void { stopHttpServer() } -export function getServerPort(): number { - return serverPort -} - function handleRequest( req: http.IncomingMessage, res: http.ServerResponse, @@ -422,11 +418,23 @@ function handleStateApi( typeof data.baseVersion === "number" && data.baseVersion < (current?.serverVersion ?? 0) ) { + let savedToHistory = false + if (data.source === "sync") { + // A stale sync reply: the store already holds the + // newer AI write, so the sync is done. + if (current) current.syncRequested = undefined + } else if (typeof data.xml === "string" && data.xml) { + // A user edit lost the race with an AI write. Keep + // it in history so the user can restore it. + addHistory(sessionId, data.xml, data.svg || "") + savedToHistory = true + } res.writeHead(409, { "Content-Type": "application/json" }) res.end( JSON.stringify({ error: "Diagram changed on the server", version: current?.version, + savedToHistory, }), ) return @@ -708,6 +716,15 @@ function getHtmlPage(sessionId: string): string { } .filename-group { display: flex; } .filename-group input { border-radius: 8px 0 0 8px; border-right: none; } + #notice { + display: none; position: fixed; left: 50%; bottom: 24px; + transform: translateX(-50%); z-index: 3000; max-width: 480px; + padding: 10px 16px; border-radius: 8px; font-size: 13px; + background: #18181b; color: white; + font-family: 'DM Sans', system-ui, -apple-system, sans-serif; + box-shadow: 0 4px 12px rgba(0,0,0,0.2); + } + #notice.open { display: block; } .filename-group .ext { padding: 10px 14px; background: #f4f4f5; border: 1px solid #e4e4e7; border-radius: 0 8px 8px 0; font-size: 13px; color: #71717a; @@ -749,6 +766,7 @@ function getHtmlPage(sessionId: string): string { +

History

@@ -797,7 +815,8 @@ function getHtmlPage(sessionId: string): string { let pendingAiSvg = false; let pendingMcpExport = null; // 'png' or 'svg' when MCP requested export let projectionExportActive = false; // page-targeted export: showing a transient single-page projection - let projectionRestoreXml = null; // the real document to reload once a projection export finishes + let forceReload = false; // reload the server state on the next poll even if the version is unchanged + let noticeTimer = null; window.addEventListener('message', (e) => { if (e.origin !== '${DRAWIO_ORIGIN}') return; @@ -824,7 +843,10 @@ function getHtmlPage(sessionId: string): string { // draw.io returns the XML in msg.xml, with no msg.data. if (pendingSyncExport && msg.xml) { pendingSyncExport = false; - pushState(msg.xml, ''); + // Push with the version the export was taken at: a + // newer AI write may have loaded meanwhile, and this + // older XML must not overwrite it. + pushState(msg.xml, '', pendingSyncBase, 'sync'); } } else if (msg.event === 'export' && msg.data) { // Handle MCP server export request (png/svg). fireExport tags @@ -848,8 +870,11 @@ function getHtmlPage(sessionId: string): string { } return; } - // Handle file download export (PNG/SVG only, drawio uses lastXml directly) - if (pendingDownload && (pendingDownload.format === 'png' || pendingDownload.format === 'svg')) { + // Handle file download export (PNG/SVG only, drawio uses + // lastXml directly). Tagged with dlExport like mcpExport, + // so an autosave SVG export can never be saved instead. + if (msg.message && msg.message.dlExport) { + if (!pendingDownload) return; const dl = pendingDownload; pendingDownload = null; let dataUrl = msg.data; @@ -897,36 +922,46 @@ function getHtmlPage(sessionId: string): string { } // Restore the user's real document after a page-targeted projection - // export. If we never captured one (lastXml was null at projection - // start), fall back to forcing a reload from the server on the next - // poll by rewinding currentVersion — never leave the iframe stuck on - // the transient projection. + // export by reloading the server state. The server also has any + // autosave that was still in flight when the projection started, + // which a copy taken at that moment would miss. A flag is used + // because a push finishing meanwhile may update currentVersion. function restoreFromProjection() { if (!projectionExportActive) return; projectionExportActive = false; - if (projectionRestoreXml) { - iframe.contentWindow.postMessage(JSON.stringify({ action: 'load', xml: projectionRestoreXml, autosave: 1 }), '*'); - projectionRestoreXml = null; - } else { - currentVersion = -1; // force the next poll to reload from server - } + forceReload = true; + poll(); } - async function pushState(xml, svg = '', baseVersion = currentVersion) { + function showNotice(text) { + const el = document.getElementById('notice'); + el.textContent = text; + el.classList.add('open'); + clearTimeout(noticeTimer); + noticeTimer = setTimeout(() => el.classList.remove('open'), 8000); + } + + // source is 'sync' for replies to a server sync request, else 'edit' + async function pushState(xml, svg = '', baseVersion = currentVersion, source = 'edit') { if (!sessionId) return; try { const r = await fetch('/api/state', { method: 'POST', headers: { 'Content-Type': 'application/json' }, - body: JSON.stringify({ sessionId, xml, svg, baseVersion }) + body: JSON.stringify({ sessionId, xml, svg, baseVersion, source }) }); if (r.ok) { const d = await r.json(); currentVersion = d.version; lastXml = xml; } // 409: the AI wrote a newer version; load it now - else if (r.status === 409) poll(); + else if (r.status === 409) { + const d = await r.json().catch(() => ({})); + if (d.savedToHistory) showNotice('The AI changed the diagram while you were editing. Your last change was saved in History.'); + poll(); + } } catch (e) { console.error('Push failed:', e); } } let pendingSyncExport = false; + let pendingSyncBase = 0; // version the pending sync export was taken at async function poll() { if (!sessionId) return; @@ -939,6 +974,7 @@ function getHtmlPage(sessionId: string): string { // while in case draw.io never answers, so later syncs still run. if (s.syncRequested && !pendingSyncExport && isReady) { pendingSyncExport = true; + pendingSyncBase = currentVersion; iframe.contentWindow.postMessage(JSON.stringify({ action: 'export', format: 'xml' }), '*'); setTimeout(() => { pendingSyncExport = false; }, 5000); } @@ -952,7 +988,8 @@ function getHtmlPage(sessionId: string): string { // so it doesn't fight the projection — and leave currentVersion // unadvanced so this bump is re-detected and applied once the // real document is restored. - if (s.version > currentVersion && s.xml && !projectionExportActive) { + if ((forceReload || s.version > currentVersion) && s.xml && !projectionExportActive) { + forceReload = false; currentVersion = s.version; loadDiagram(s.xml, true); } @@ -972,13 +1009,11 @@ function getHtmlPage(sessionId: string): string { const fireExport = () => { // mcpExport is echoed back in msg.message (see the handler) const exportOpts = pendingMcpExport === 'png' - ? { action: 'export', format: 'png', scale: 2, mcpExport: true } + ? { action: 'export', format: 'png', scale: 2, currentPage: true, mcpExport: true } : { action: 'export', format: 'svg', mcpExport: true }; iframe.contentWindow.postMessage(JSON.stringify(exportOpts), '*'); }; if (s.exportXml) { - // Stash the real document so we can restore after export. - projectionRestoreXml = lastXml; projectionExportActive = true; // Load the projection without touching lastXml/server state. iframe.contentWindow.postMessage(JSON.stringify({ action: 'load', xml: s.exportXml, autosave: 0 }), '*'); @@ -1055,11 +1090,11 @@ function getHtmlPage(sessionId: string): string { saveConfirmBtn.textContent = 'Save'; } else if (format === 'png') { pendingDownload = { format: 'png', filename }; - iframe.contentWindow.postMessage(JSON.stringify({ action: 'export', format: 'png', scale: 2 }), '*'); + iframe.contentWindow.postMessage(JSON.stringify({ action: 'export', format: 'png', scale: 2, currentPage: true, dlExport: true }), '*'); setTimeout(() => { saveConfirmBtn.disabled = false; saveConfirmBtn.textContent = 'Save'; pendingDownload = null; }, 5000); } else if (format === 'svg') { pendingDownload = { format: 'svg', filename }; - iframe.contentWindow.postMessage(JSON.stringify({ action: 'export', format: 'svg' }), '*'); + iframe.contentWindow.postMessage(JSON.stringify({ action: 'export', format: 'svg', dlExport: true }), '*'); setTimeout(() => { saveConfirmBtn.disabled = false; saveConfirmBtn.textContent = 'Save'; pendingDownload = null; }, 5000); } }; diff --git a/packages/mcp-server/src/index.ts b/packages/mcp-server/src/index.ts index 47fad275..1d70cac8 100644 --- a/packages/mcp-server/src/index.ts +++ b/packages/mcp-server/src/index.ts @@ -19,15 +19,14 @@ */ import { createRequire } from "node:module" +import { homedir } from "node:os" import { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js" import { StdioServerTransport } from "@modelcontextprotocol/sdk/server/stdio.js" import open from "open" import { z } from "zod" -import { - applyDiagramOperations, - type DiagramOperation, -} from "./diagram-operations.js" +import type { DiagramOperation } from "./diagram-operations.js" import { installDomPolyfill } from "./dom.js" +import { editDiagram, targetPageXml } from "./edit-diagram.js" import { checkEditGate } from "./edit-gate.js" import { addHistory } from "./history.js" import { @@ -135,6 +134,13 @@ function pickPageSelector(input: { return selector } +/** Expand a leading ~ to the home directory (shells do this, MCP hosts don't). */ +function expandHome(p: string): string { + if (p === "~") return homedir() + if (p.startsWith("~/") || p.startsWith("~\\")) return homedir() + p.slice(1) + return p +} + /** Format a selector for human-readable error messages. */ function describeSelector(s: PageSelector): string { if (s.page_id) return `id="${s.page_id}"` @@ -144,9 +150,11 @@ function describeSelector(s: PageSelector): string { } // Register prompt with workflow guidance -server.prompt( +server.registerPrompt( "diagram-workflow", - "Guidelines for creating and editing draw.io diagrams", + { + description: "Guidelines for creating and editing draw.io diagrams", + }, () => ({ messages: [ { @@ -177,7 +185,7 @@ server.prompt( ## Important Notes - create_new_diagram REPLACES the entire document, including ALL pages - only use for new diagrams. Use add_page to add a tab without losing existing content. -- edit_diagram PRESERVES the user's manual changes (fetches browser state first) +- edit_diagram PRESERVES the user's manual changes: it is rejected with the current XML when the user edited the diagram since you last saw it - Always use unique cell_ids within a page (cell ids "0" and "1" are reserved root sentinels and can repeat across pages)`, }, }, @@ -433,7 +441,7 @@ server.registerTool( path: z .string() .describe( - "Path to the .drawio file to load (e.g., ./diagram.drawio)", + "Absolute path to the .drawio file to load (e.g. /Users/me/diagram.drawio or ~/diagram.drawio). Relative paths resolve against the MCP server's working directory, which is often not your project.", ), }, }, @@ -453,7 +461,7 @@ server.registerTool( const fs = await import("node:fs/promises") const nodePath = await import("node:path") - const absolutePath = nodePath.resolve(path) + const absolutePath = nodePath.resolve(expandHome(path)) let content: string try { @@ -543,10 +551,11 @@ server.registerTool( { description: "Edit a specific page in the current diagram by ID-based operations (update/add/delete cells).\n\n" + + "All-or-nothing: if any operation fails, nothing is applied and every failure is listed.\n\n" + "Freshness: the server remembers the last diagram state you have seen, and rejects this call " + "only if the user edited the diagram in the browser since then. You do NOT need to call " + - "get_diagram before every edit — if your view is stale, the call is rejected (with no side " + - "effects) and the error tells you to call get_diagram once and retry.\n\n" + + "get_diagram before every edit: a rejected call changes nothing and includes the current XML " + + "of the page, so you can rebuild your operations and retry.\n\n" + "Call get_diagram first only when you don't know the current diagram content (cell IDs, " + "structure) — e.g. the diagram wasn't created in this conversation, or you're unsure your " + "memory of it is accurate.\n\n" + @@ -625,6 +634,12 @@ server.registerTool( } } + const pageSelector = pickPageSelector({ + page_id, + page_name, + page_index, + }) + // Enforce workflow: the model must have seen the current diagram // state. Content comparison instead of a wall-clock timeout — // slow reasoning between get_diagram and edit_diagram is fine as @@ -637,90 +652,54 @@ server.registerTool( log.warn( gate.reason === "stale" ? "edit_diagram called with unseen browser changes - rejecting to prevent data loss" - : "edit_diagram called without get_diagram - rejecting to prevent data loss", + : "edit_diagram called without seeing the diagram - rejecting to prevent data loss", ) + // The error carries the current page, so the model has now + // seen it and can retry without a get_diagram round-trip. + currentSession.lastSeenXml = + browserState?.xml || currentSession.xml + const reason = + gate.reason === "stale" + ? "The diagram changed in the browser since you last saw it (e.g. manual user edits). No changes were made." + : "You have not seen this diagram yet, so no changes were made." return { content: [ { type: "text", - text: - gate.reason === "stale" - ? "Error: The diagram changed in the browser since you last fetched it (e.g. manual user edits).\n\n" + - "Call get_diagram to see the latest state, then rebuild your edit operations on top of it." - : "Error: You must call get_diagram first before edit_diagram.\n\n" + - "This ensures you have the latest diagram state including any manual edits the user made in the browser. " + - "Please call get_diagram, then use that XML to construct your edit operations.", + text: `Error: ${reason}\n\nCurrent XML of ${describeSelector(pageSelector)}:\n\n${targetPageXml(currentSession.xml, pageSelector)}\n\nBuild your operations on this XML and retry.`, }, ], isError: true, } } - const pageSelector = pickPageSelector({ - page_id, - page_name, - page_index, - }) log.info( `Editing diagram with ${operations.length} operation(s) on ${describeSelector(pageSelector)}`, ) - // Validate and auto-fix new_xml for each operation - const validatedOps = operations.map((op) => { - if (op.new_xml) { - const { valid, error, fixed, fixes } = validateAndFixXml( - op.new_xml, - ) - if (fixed) { - log.info( - `Operation ${op.operation} ${op.cell_id}: XML auto-fixed: ${fixes.join(", ")}`, - ) - return { ...op, new_xml: fixed } - } - if (!valid && error) { - log.warn( - `Operation ${op.operation} ${op.cell_id}: XML validation failed: ${error}`, - ) - } - } - return op - }) - - // Apply operations on the targeted page - const { result, errors } = applyDiagramOperations( + const outcome = editDiagram( currentSession.xml, - validatedOps as DiagramOperation[], + operations as DiagramOperation[], pageSelector, ) - - if (errors.length > 0) { - const errorMessages = errors - .map((e) => `${e.type} ${e.cellId}: ${e.message}`) - .join("\n") - log.warn(`Edit had ${errors.length} error(s): ${errorMessages}`) - } - - // A page-level error (empty cellId — e.g. the selector matched no - // page, or the page had no ) means NOTHING was applied and - // `result` is the unchanged input. Surface it as a hard error - // instead of persisting a no-op and reporting success, so the - // caller doesn't build on a wrong assumption. - const pageError = errors.find((e) => e.cellId === "") - if (pageError) { + if (!outcome.ok) { + log.warn(`Edit rejected: ${outcome.errors.join("; ")}`) + const text = outcome.pageError + ? `Error: ${outcome.errors[0]}` + : `Error: No changes were made because ${outcome.errors.length} operation(s) failed:\n${outcome.errors.map((e) => `- ${e}`).join("\n")}\n\nCurrent XML of ${describeSelector(pageSelector)}:\n\n${targetPageXml(currentSession.xml, pageSelector)}\n\nFix the operations against this XML and retry.` return { - content: [ - { - type: "text", - text: `Error: ${pageError.message}`, - }, - ], + content: [{ type: "text", text }], isError: true, } } + if (outcome.fixes.length > 0) { + log.info(`new_xml auto-fixed: ${outcome.fixes.join("; ")}`) + } + const result = outcome.xml // Save the pre-edit state for undo (with cached SVG from browser). - // Done only now that we know the edit applied — a page-level error - // returns above without leaving a phantom history entry. + // Done only once the edit applied: a rejected edit returns above + // without leaving a phantom history entry. addHistory( currentSession.id, currentSession.xml, @@ -741,17 +720,11 @@ server.registerTool( log.info(`Diagram edited successfully`) - const successMsg = `Diagram edited successfully!\n\nApplied ${operations.length} operation(s) on ${describeSelector(pageSelector)}.` - const errorMsg = - errors.length > 0 - ? `\n\nWarnings:\n${errors.map((e) => `- ${e.type} ${e.cellId}: ${e.message}`).join("\n")}` - : "" - return { content: [ { type: "text", - text: successMsg + errorMsg, + text: `Diagram edited successfully!\n\nApplied ${outcome.applied} operation(s) on ${describeSelector(pageSelector)}.`, }, ], } @@ -802,11 +775,14 @@ server.registerTool( } // Request browser to push fresh state and wait for it + let staleNote = "" const syncRequested = requestSync(currentSession.id) if (syncRequested) { const synced = await waitForSync(currentSession.id) if (!synced) { log.warn("get_diagram: sync timeout - state may be stale") + staleNote = + "\n\nNote: the browser did not respond, so this XML may not include the user's latest manual edits (is the preview tab open?)." } } @@ -852,7 +828,7 @@ server.registerTool( content: [ { type: "text", - text: `Current diagram XML:\n\n${currentSession.xml}\n\n${pageList}`, + text: `Current diagram XML:\n\n${currentSession.xml}\n\n${pageList}${staleNote}`, }, ], } @@ -878,7 +854,7 @@ server.registerTool( content: [ { type: "text", - text: `Page ${projection.index} ("${projection.name}"):\n\n${projection.xml}\n\n${pageList}`, + text: `Page ${projection.index} ("${projection.name}"):\n\n${projection.xml}\n\n${pageList}${staleNote}`, }, ], } @@ -952,7 +928,7 @@ server.registerTool( path: z .string() .describe( - "File path to save the diagram (e.g., ./diagram.drawio, ./diagram.png, ./diagram.svg)", + "Absolute file path to save to (e.g. /Users/me/diagram.drawio, ~/diagram.png). Relative paths resolve against the MCP server's working directory, which is often not your project.", ), format: z .enum(["drawio", "png", "svg"]) @@ -962,7 +938,8 @@ server.registerTool( ), }, }, - async ({ path, format, page_id, page_name, page_index }) => { + async ({ path: rawPath, format, page_id, page_name, page_index }) => { + const path = expandHome(rawPath) try { if (!currentSession) { return { @@ -976,10 +953,12 @@ server.registerTool( } } - // Fetch latest state + // Fetch latest state, re-normalised to mxfile so a page + // selector works on a bare pushed by the browser const browserState = getState(currentSession.id) if (browserState?.xml) { - currentSession.xml = browserState.xml + currentSession.xml = + normalizeToMxfile(browserState.xml) ?? browserState.xml } if (!currentSession.xml) { @@ -1012,7 +991,7 @@ server.registerTool( // .drawio path - write XML directly (no browser round-trip). if (detectedFormat === "drawio") { let filePath = path - if (!filePath.endsWith(".drawio")) { + if (!filePath.toLowerCase().endsWith(".drawio")) { filePath = `${filePath}.drawio` } const absolutePath = nodePath.resolve(filePath) diff --git a/packages/mcp-server/src/pages.ts b/packages/mcp-server/src/pages.ts index d1948831..447788aa 100644 --- a/packages/mcp-server/src/pages.ts +++ b/packages/mcp-server/src/pages.ts @@ -302,14 +302,15 @@ export function deletePageFromDoc( doc: Document, selector: PageSelector, ): { ok: boolean; reason?: string; deletedId?: string; deletedIndex?: number } { - const pages = listPagesFromDoc(doc) - if (pages.length <= 1) { - return { ok: false, reason: "Cannot delete the only remaining page" } - } + // Match first, so a wrong selector reports "not found" even on a + // one-page document const found = findPageElement(doc, selector) if (!found) { return { ok: false, reason: "Page not found" } } + if (listPagesFromDoc(doc).length <= 1) { + return { ok: false, reason: "Cannot delete the only remaining page" } + } const id = found.element.getAttribute("id") || "" const index = found.index found.element.parentNode?.removeChild(found.element) diff --git a/packages/mcp-server/src/xml-validation.ts b/packages/mcp-server/src/xml-validation.ts index 9d9e36fa..0c6b617c 100644 --- a/packages/mcp-server/src/xml-validation.ts +++ b/packages/mcp-server/src/xml-validation.ts @@ -25,6 +25,23 @@ const STRUCTURAL_ATTRS = [ /** Valid XML entity names */ const VALID_ENTITIES = new Set(["lt", "gt", "amp", "quot", "apos"]) +/** Element names draw.io understands (case-sensitive) */ +const VALID_DRAWIO_TAGS = new Set([ + "mxfile", + "diagram", + "mxGraphModel", + "root", + "mxCell", + "mxGeometry", + "mxPoint", + "Array", + "Object", + "mxRectangle", + // Wrappers draw.io writes for cells with links, tooltips or data + "UserObject", + "object", +]) + // ============================================================================ // XML Parsing Helpers // ============================================================================ @@ -90,6 +107,33 @@ function parseXmlTags(xml: string): ParsedTag[] { return tags } +/** + * Returns a function telling whether a position lies inside a quoted + * attribute value. Positions must be queried in increasing order. + */ +function createQuoteTracker(str: string): (pos: number) => boolean { + let i = 0 + let inQuote = false + let quoteChar = "" + return (pos: number) => { + for (; i < pos && i < str.length; i++) { + const c = str[i] + if (inQuote) { + if (c === quoteChar) inQuote = false + } else if (c === '"' || c === "'") { + // Only quotes that follow "=" open an attribute value + let j = i - 1 + while (j >= 0 && /\s/.test(str[j])) j-- + if (j >= 0 && str[j] === "=") { + inQuote = true + quoteChar = c + } + } + } + return inQuote + } +} + /** Rewrite every opening tag with fn, leaving text and closing tags as is. */ function replaceInOpeningTags( xml: string, @@ -306,6 +350,37 @@ function checkNestedMxCells(xml: string): string | null { return null } +/** Check for element names draw.io does not know (e.g. a lowercase ) */ +function checkUnknownElements(xml: string): string | null { + const tags = parseXmlTags(xml.replace(//g, "")) + for (const { tagName } of tags) { + if (!VALID_DRAWIO_TAGS.has(tagName)) { + return `Invalid XML: Unknown element <${tagName}>. draw.io only understands ${Array.from(VALID_DRAWIO_TAGS).join(", ")} (names are case-sensitive).` + } + } + return null +} + +/** + * Find elements without an "as" attribute outside . draw.io rejects them with "Could not add object mxPoint". + */ +function findOrphanMxPoints( + xml: string, +): Array<{ start: number; end: number }> { + const arrays: Array<[number, number]> = [] + for (const m of xml.matchAll(/]*>[\s\S]*?<\/Array>/g)) { + arrays.push([m.index, m.index + m[0].length]) + } + const orphans: Array<{ start: number; end: number }> = [] + for (const m of xml.matchAll(/]*?(?:\/>|>\s*<\/mxPoint>)/g)) { + if (/\sas\s*=/.test(m[0])) continue + if (arrays.some(([s, e]) => m.index > s && m.index < e)) continue + orphans.push({ start: m.index, end: m.index + m[0].length }) + } + return orphans +} + // ============================================================================ // Main Validation Function // ============================================================================ @@ -314,9 +389,14 @@ function checkNestedMxCells(xml: string): string | null { * Validates draw.io XML structure for common issues * Uses DOM parsing + additional regex checks for high accuracy * @param xml - The XML string to validate + * @param opts.strict - Also reject unknown element names and orphan + * s. Used for XML the model wrote, not for files or browser state. * @returns null if valid, error message string if invalid */ -export function validateMxCellStructure(xml: string): string | null { +export function validateMxCellStructure( + xml: string, + opts: { strict?: boolean } = {}, +): string | null { // Size check for performance if (xml.length > MAX_XML_SIZE) { console.warn( @@ -410,6 +490,16 @@ export function validateMxCellStructure(xml: string): string | null { return nestedCellError } + if (opts.strict) { + const unknownError = checkUnknownElements(xml) + if (unknownError) { + return unknownError + } + if (findOrphanMxPoints(xml).length > 0) { + return 'Invalid XML: Found without an "as" attribute outside . Put waypoints inside or remove the point.' + } + } + // 11. Strict XML syntax check, run last so the checks above can give // more specific messages. Catches what they miss, e.g. duplicate or // unquoted attributes, which make draw.io refuse to load the diagram. @@ -629,65 +719,67 @@ export function autoFixXml(xml: string): { fixed: string; fixes: string[] } { fixes.push("Fixed tags to ") } - // 15. Fix common closing tag typos (MUST run before foreign tag removal) - const tagTypos = [ - { wrong: /<\/mxElement>/gi, right: "", name: "" }, - { wrong: /<\/mxcell>/g, right: "", name: "" }, - { - wrong: /<\/mxgeometry>/g, - right: "", - name: "", - }, - { wrong: /<\/mxpoint>/g, right: "", name: "" }, - { - wrong: /<\/mxgraphmodel>/gi, - right: "", - name: "", - }, - ] - for (const { wrong, right, name } of tagTypos) { - const before = fixed - fixed = fixed.replace(wrong, right) - if (fixed !== before) { - fixes.push(`Fixed typo ${name} to ${right}`) + // 15. Fix closing tag typos and wrong tag case, e.g. (MUST run + // before foreign tag removal, which would otherwise delete them) + const before15 = fixed + fixed = fixed.replace(/<\/mxElement>/gi, "") + if (fixed !== before15) { + fixes.push("Fixed typo to ") + } + for (const name of ["mxCell", "mxGeometry", "mxPoint", "mxGraphModel"]) { + let changed = false + fixed = fixed.replace( + new RegExp(`<(/?)${name}(?=[\\s/>])`, "gi"), + (match, slash) => { + const right = `<${slash}${name}` + if (match !== right) changed = true + return right + }, + ) + if (changed) { + fixes.push(`Fixed tag case of <${name}>`) } } - // 16. Remove non-draw.io tags (after typo fixes so lowercase variants are fixed first) - const validDrawioTags = new Set([ - "mxfile", - "diagram", - "mxGraphModel", - "root", - "mxCell", - "mxGeometry", - "mxPoint", - "Array", - "Object", - "mxRectangle", - // Wrappers draw.io writes for cells with links, tooltips or data - "UserObject", - "object", - ]) + // 16. Remove non-draw.io tags (after the case fixes above). Removes only + // the exact tag occurrences and skips quoted attribute values, so a stray + // never takes with it and inside + // value="..." stays. + const isInsideQuotesFor16 = createQuoteTracker(fixed) const foreignTagPattern = /<\/?([a-zA-Z][a-zA-Z0-9_]*)[^>]*>/g let foreignMatch const foreignTags = new Set() + const foreignTagPositions: Array<{ start: number; end: number }> = [] while ((foreignMatch = foreignTagPattern.exec(fixed)) !== null) { const tagName = foreignMatch[1] - if (!validDrawioTags.has(tagName)) { - foreignTags.add(tagName) - } + if (VALID_DRAWIO_TAGS.has(tagName)) continue + if (isInsideQuotesFor16(foreignMatch.index)) continue + foreignTags.add(tagName) + foreignTagPositions.push({ + start: foreignMatch.index, + end: foreignMatch.index + foreignMatch[0].length, + }) } - if (foreignTags.size > 0) { - for (const tag of foreignTags) { - fixed = fixed.replace(new RegExp(`<${tag}[^>]*>`, "gi"), "") - fixed = fixed.replace(new RegExp(``, "gi"), "") + if (foreignTagPositions.length > 0) { + // Remove from the end so earlier positions stay valid + for (const { start, end } of foreignTagPositions.reverse()) { + fixed = fixed.slice(0, start) + fixed.slice(end) } fixes.push( `Removed foreign tags: ${Array.from(foreignTags).join(", ")}`, ) } + // 16b. Remove orphan s (no "as" attribute, not inside + // ), which draw.io refuses to load + const orphanPoints = findOrphanMxPoints(fixed) + if (orphanPoints.length > 0) { + for (const { start, end } of orphanPoints.reverse()) { + fixed = fixed.slice(0, start) + fixed.slice(end) + } + fixes.push(`Removed ${orphanPoints.length} orphan (s)`) + } + // 17. Fix unclosed tags const tagStack: string[] = [] const parsedTags = parseXmlTags(fixed) @@ -725,19 +817,23 @@ export function autoFixXml(xml: string): { fixed: string; fixes: string[] } { } } - // 18. Remove extra closing tags + // 18. Remove extra closing tags. Counts only draw.io tags outside quoted + // attribute values (value="Title" holds HTML, not elements). const tagCounts = new Map< string, { opens: number; closes: number; selfClosing: number } >() const fullTagPattern = /<(\/?[a-zA-Z][a-zA-Z0-9]*)[^>]*>/g + const isInsideQuotesFor18 = createQuoteTracker(fixed) let tagCountMatch while ((tagCountMatch = fullTagPattern.exec(fixed)) !== null) { + if (isInsideQuotesFor18(tagCountMatch.index)) continue const fullMatch = tagCountMatch[0] const tagPart = tagCountMatch[1] const isClosing = tagPart.startsWith("/") const isSelfClosing = fullMatch.endsWith("/>") const tagName = isClosing ? tagPart.slice(1) : tagPart + if (!VALID_DRAWIO_TAGS.has(tagName)) continue let counts = tagCounts.get(tagName) if (!counts) { @@ -943,7 +1039,8 @@ export function autoFixXml(xml: string): { fixed: string; fixes: string[] } { // ============================================================================ /** - * Validates XML and attempts to fix if invalid + * Validates XML and attempts to fix if invalid. Runs the strict checks + * (unknown elements, orphan mxPoints): every caller passes XML the model wrote. * @param xml - The XML string to validate and potentially fix * @returns Object with validation result, fixed XML if applicable, and fixes applied */ @@ -954,7 +1051,7 @@ export function validateAndFixXml(xml: string): { fixes: string[] } { // First validation attempt - let error = validateMxCellStructure(xml) + let error = validateMxCellStructure(xml, { strict: true }) if (!error) { return { valid: true, error: null, fixed: null, fixes: [] } @@ -964,7 +1061,7 @@ export function validateAndFixXml(xml: string): { const { fixed, fixes } = autoFixXml(xml) // Validate the fixed version - error = validateMxCellStructure(fixed) + error = validateMxCellStructure(fixed, { strict: true }) if (!error) { return { valid: true, error: null, fixed, fixes } @@ -978,33 +1075,3 @@ export function validateAndFixXml(xml: string): { fixes, } } - -/** - * Check if mxCell XML output is complete (not truncated). - * Uses a robust approach that handles any LLM provider's wrapper tags - * by finding the last valid mxCell ending and checking if suffix is just closing tags. - * @param xml - The XML string to check (can be undefined/null) - * @returns true if XML appears complete, false if truncated or empty - */ -export function isMxCellXmlComplete(xml: string | undefined | null): boolean { - const trimmed = xml?.trim() || "" - if (!trimmed) return false - - // Find position of last complete mxCell ending (either /> or ) - const lastSelfClose = trimmed.lastIndexOf("/>") - const lastMxCellClose = trimmed.lastIndexOf("") - - const lastValidEnd = Math.max(lastSelfClose, lastMxCellClose) - - // No valid ending found at all - if (lastValidEnd === -1) return false - - // Check what comes after the last valid ending - // For />: add 2 chars, for : add 9 chars - const endOffset = lastMxCellClose > lastSelfClose ? 9 : 2 - const suffix = trimmed.slice(lastValidEnd + endOffset) - - // If suffix is empty or only contains closing tags (any provider's wrapper) or whitespace, it's complete - // This regex matches any sequence of closing XML tags like , , - return /^(\s*<\/[^>]+>)*\s*$/.test(suffix) -} diff --git a/packages/mcp-server/tests/diagram-operations.test.ts b/packages/mcp-server/tests/diagram-operations.test.ts index 405ca07f..7c0801bc 100644 --- a/packages/mcp-server/tests/diagram-operations.test.ts +++ b/packages/mcp-server/tests/diagram-operations.test.ts @@ -4,6 +4,7 @@ * The id sits on the wrapper; the inner mxCell has none. */ +import { deflateRawSync } from "node:zlib" import { beforeAll, describe, expect, it, vi } from "vitest" import { installDomPolyfill } from "../src/dom.js" @@ -83,3 +84,43 @@ describe("cascade delete logging", () => { spy.mockRestore() }) }) + +describe("pages without a ", () => { + const ADD_A = { + operation: "add" as const, + cell_id: "a", + new_xml: ``, + } + + it("treats an empty page as a blank page", () => { + const doc = `` + const { result, errors } = applyDiagramOperations(doc, [ADD_A]) + expect(errors).toEqual([]) + expect(result).toContain('') + expect(result).toContain('') + expect(result).toContain(' { + const model = `` + const compressed = deflateRawSync( + Buffer.from(encodeURIComponent(model)), + ).toString("base64") + const doc = `${compressed}` + const { result, errors } = applyDiagramOperations(doc, [ADD_A]) + expect(errors).toEqual([]) + expect(result).not.toContain(compressed) + expect(result).toContain(' { + const doc = `not base64 !!` + const { errors } = applyDiagramOperations(doc, [ADD_A]) + expect(errors[0]?.cellId).toBe("") + expect(errors[0]?.message).toContain("could not be decompressed") + }) +}) diff --git a/packages/mcp-server/tests/edit-diagram.test.ts b/packages/mcp-server/tests/edit-diagram.test.ts new file mode 100644 index 00000000..ad83cf30 --- /dev/null +++ b/packages/mcp-server/tests/edit-diagram.test.ts @@ -0,0 +1,130 @@ +/** + * Tests for the all-or-nothing edit_diagram core (src/edit-diagram.ts). + */ + +import { beforeAll, describe, expect, it } from "vitest" +import { installDomPolyfill } from "../src/dom.js" + +beforeAll(() => { + installDomPolyfill() +}) + +import { editDiagram, targetPageXml } from "../src/edit-diagram.js" +import { validateMxCellStructure } from "../src/xml-validation.js" + +const cell = (id: string, extra = "") => + `` + +const page = (id: string, cells: string) => + `${cells}` + +const DOC = `${page("p1", cell("a") + cell("b"))}` + +describe("editDiagram", () => { + it("applies every operation and counts them", () => { + const out = editDiagram( + DOC, + [ + { operation: "add", cell_id: "c", new_xml: cell("c") }, + { operation: "delete", cell_id: "b" }, + ], + {}, + ) + expect(out.ok).toBe(true) + if (!out.ok) return + expect(out.applied).toBe(2) + expect(out.xml).toContain('id="c"') + expect(out.xml).not.toContain('id="b"') + }) + + it("applies nothing when one operation fails", () => { + const out = editDiagram( + DOC, + [ + { operation: "add", cell_id: "c", new_xml: cell("c") }, + { operation: "delete", cell_id: "missing" }, + ], + {}, + ) + expect(out.ok).toBe(false) + if (out.ok) return + expect(out.pageError).toBe(false) + expect(out.errors).toEqual([ + 'delete missing: Cell with id="missing" not found', + ]) + }) + + it("rejects new_xml that is still invalid after auto-fix", () => { + const out = editDiagram( + DOC, + [ + { + operation: "update", + cell_id: "a", + new_xml: ``, + }, + ], + {}, + ) + expect(out.ok).toBe(false) + if (out.ok) return + expect(out.errors[0]).toMatch(/^update a: invalid new_xml: /) + }) + + it("rejects several cells in one new_xml", () => { + const out = editDiagram( + DOC, + [ + { + operation: "add", + cell_id: "c", + new_xml: cell("c") + cell("d"), + }, + ], + {}, + ) + expect(out.ok).toBe(false) + if (out.ok) return + expect(out.errors[0]).toContain("exactly one cell") + }) + + it("accepts a UserObject that wraps one mxCell", () => { + const wrapped = `` + const out = editDiagram( + DOC, + [{ operation: "add", cell_id: "u", new_xml: wrapped }], + {}, + ) + expect(out.ok).toBe(true) + }) + + it("is not blocked by a problem on another page", () => { + // Page p2 has a duplicate cell id, which fails validation + const doc = `${page("p1", cell("a"))}${page("p2", cell("x") + cell("x"))}` + expect(validateMxCellStructure(doc)).not.toBeNull() + const out = editDiagram( + doc, + [{ operation: "add", cell_id: "c", new_xml: cell("c") }], + { page_id: "p1" }, + ) + expect(out.ok).toBe(true) + }) + + it("reports a missing page as a page-level error", () => { + const out = editDiagram(DOC, [{ operation: "delete", cell_id: "a" }], { + page_id: "nope", + }) + expect(out.ok).toBe(false) + if (out.ok) return + expect(out.pageError).toBe(true) + }) +}) + +describe("targetPageXml", () => { + it("returns only the selected page", () => { + const doc = `${page("p1", cell("a"))}${page("p2", cell("z"))}` + const xml = targetPageXml(doc, { page_id: "p2" }) + expect(xml).toContain('id="z"') + expect(xml).not.toContain('id="a"') + }) +}) diff --git a/packages/mcp-server/tests/http-server.test.ts b/packages/mcp-server/tests/http-server.test.ts index 6074066b..5654792f 100644 --- a/packages/mcp-server/tests/http-server.test.ts +++ b/packages/mcp-server/tests/http-server.test.ts @@ -8,12 +8,14 @@ import http from "node:http" import { afterAll, beforeAll, describe, expect, it } from "vitest" -import { addHistory } from "../src/history.js" +import { addHistory, getHistory } from "../src/history.js" import { getState, + requestSync, setState, shutdown, startHttpServer, + waitForSync, } from "../src/http-server.js" let port = 0 @@ -189,6 +191,55 @@ describe("POST /api/state", () => { expect(getState(id)?.xml).toBe(xml) } }) + + it("keeps a rejected user edit in history", async () => { + const id = "mcp-conflict-history" + setState(id, "user v1", undefined, true) + const aiVersion = setState(id, "AI edit") + const before = getHistory(id).length + + const stale = await postJson("/api/state", { + sessionId: id, + xml: "lost user edit", + baseVersion: aiVersion - 1, + }) + expect(stale.status).toBe(409) + expect(JSON.parse(stale.body).savedToHistory).toBe(true) + const history = getHistory(id) + expect(history).toHaveLength(before + 1) + expect(history.at(-1)?.xml).toBe("lost user edit") + }) + + it("ends a pending sync when the sync reply is older than an AI write", async () => { + const id = "mcp-stale-sync" + setState(id, "before", undefined, true) + const aiVersion = setState(id, "AI edit") + requestSync(id) + const before = getHistory(id).length + + // The browser exported its old diagram, then loaded the AI write + const stale = await postJson("/api/state", { + sessionId: id, + xml: "before", + baseVersion: aiVersion - 1, + source: "sync", + }) + expect(stale.status).toBe(409) + expect(JSON.parse(stale.body).savedToHistory).toBe(false) + expect(getState(id)?.xml).toBe("AI edit") + expect(getState(id)?.syncRequested).toBeUndefined() + expect(getHistory(id)).toHaveLength(before) + expect(await waitForSync(id, 200)).toBe(true) + }) +}) + +describe("preview page", () => { + it("serves a script that parses", async () => { + const res = await request("/?mcp=mcp-test-script") + const script = res.body.match(/