From 504d2fa81253d2d04c4a3581512c8cc47a6b2c34 Mon Sep 17 00:00:00 2001 From: "dayuan.jiang" Date: Sun, 4 Oct 2026 22:52:41 +0900 Subject: [PATCH] fix(mcp-server): keep both pages when get_diagram meets a page export, and more review fixes Found by the PR review, each with a test that failed first: - get_diagram during a page export returned the one-page projection on screen as the whole document (6 of 6 times when timed so). The preview page no longer answers a sync while a projection shows, and syncs after reloading, so the poll that restores the real document exports it. - Exports are numbered on the server too: a late result of an export that timed out was saved as the next export's file. - In Chrome, a new_xml with a syntax error counted the element as a second cell, so the web app rejected edits that auto-fix repairs ("must contain exactly one cell"). - hasCells missed single-quoted ids, so screenshot_diagram called such a diagram empty and auto-save never created its file. - A literal \n directly under a that has a model passed validation; only text-only pages are compressed data. - A wrapped mxCell repeating its UserObject's id took the wrapper's place in edits, so delete and update left an empty or nested wrapper. - Bare cells with a shape or edge id of "0" or "1" are rejected with a clear message instead of being renamed, which broke their edges. - DRAWIO_DATA_DIR expands ~, which JSON configs pass on as it is. --- packages/mcp-server/src/diagram-operations.ts | 6 +++- packages/mcp-server/src/edit-diagram.ts | 8 ++++- packages/mcp-server/src/http-server.ts | 15 ++++++++- packages/mcp-server/src/index.ts | 11 ++----- packages/mcp-server/src/new-diagram.ts | 13 ++++++++ packages/mcp-server/src/pages.ts | 2 +- packages/mcp-server/src/persistence.ts | 9 ++++- packages/mcp-server/src/preview/preview.js | 23 ++++++++----- packages/mcp-server/src/xml-validation.ts | 4 ++- .../tests/diagram-operations.test.ts | 25 ++++++++++++++ packages/mcp-server/tests/http-server.test.ts | 19 +++++++++++ packages/mcp-server/tests/persistence.test.ts | 7 +++- packages/mcp-server/tests/wrap-cells.test.ts | 33 ++++++++++++++++++- .../mcp-server/tests/xml-validation.test.ts | 10 ++++++ tests/e2e/diagram-content.spec.ts | 30 +++++++++++++++++ 15 files changed, 189 insertions(+), 26 deletions(-) diff --git a/packages/mcp-server/src/diagram-operations.ts b/packages/mcp-server/src/diagram-operations.ts index a8870e82..5315f191 100644 --- a/packages/mcp-server/src/diagram-operations.ts +++ b/packages/mcp-server/src/diagram-operations.ts @@ -185,7 +185,11 @@ export function applyDiagramOperations( const cellMap = new Map() root.querySelectorAll(CELL_SELECTOR).forEach((cell) => { const id = cell.getAttribute("id") - if (id) cellMap.set(id, cell) + // A wrapped mxCell may repeat its wrapper's id; the wrapper is the cell + const wrapped = + cell.tagName === "mxCell" && + /^(UserObject|object)$/.test(cell.parentElement?.tagName ?? "") + if (id && !wrapped) cellMap.set(id, cell) }) // Ids deleted so far in this batch; deleting one again is a no-op const deletedIds = new Set() diff --git a/packages/mcp-server/src/edit-diagram.ts b/packages/mcp-server/src/edit-diagram.ts index ab51d31a..a77e8e92 100644 --- a/packages/mcp-server/src/edit-diagram.ts +++ b/packages/mcp-server/src/edit-diagram.ts @@ -24,7 +24,13 @@ function countTopLevelElements(fragment: string): number { `${fragment}`, "text/xml", ) - return doc.documentElement?.children.length ?? 0 + // On a syntax error the browser adds a element (Chrome + // next to the cells, Firefox as the root); the syntax is checked later + const root = doc.documentElement + if (!root || root.tagName === "parsererror") return 1 + return Array.from(root.children).filter( + (el) => el.tagName !== "parsererror", + ).length } /** The target page as a one-page , or the whole document. */ diff --git a/packages/mcp-server/src/http-server.ts b/packages/mcp-server/src/http-server.ts index 5b058fe1..2125a050 100644 --- a/packages/mcp-server/src/http-server.ts +++ b/packages/mcp-server/src/http-server.ts @@ -106,6 +106,7 @@ interface SessionState { exportFormat?: ExportFormat // 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) + exportId?: number // Number of the pending export, echoed with its result exportData?: string // Base64/SVG data returned by browser after export } @@ -161,6 +162,7 @@ export function setState( exportFormat: existing?.exportFormat, // Preserve pending export request exportXml: existing?.exportXml, // Preserve pending projection exportOptions: existing?.exportOptions, + exportId: existing?.exportId, exportData: existing?.exportData, // Preserve export result }) log.debug(`State updated: session=${sessionId}, version=${newVersion}`) @@ -191,9 +193,14 @@ export function requestExport( state.exportXml = projectionXml state.exportOptions = options state.exportFormat = format + // The browser sends this back with the result, so a late result of an + // export that timed out is not taken for this one + state.exportId = ++lastExportId return true } +let lastExportId = 0 + export function requestSync(sessionId: string): boolean { const state = stateStore.get(sessionId) if (state) { @@ -414,6 +421,7 @@ function handleStateApi( exportFormat: state?.exportFormat || null, exportXml: state?.exportXml || null, exportOptions: state?.exportOptions || null, + exportId: state?.exportId ?? null, }), ) } else if (req.method === "POST") { @@ -432,14 +440,19 @@ function handleStateApi( // Browser is returning export data (png/svg) if (data.exportData !== undefined) { const state = stateStore.get(sessionId) - if (state) { + if (state && data.exportId === state.exportId) { state.exportData = data.exportData state.exportFormat = undefined state.exportXml = undefined state.exportOptions = undefined + state.exportId = undefined log.debug( `Export data received for session=${sessionId}`, ) + } else if (state) { + log.debug( + `Ignored a late export result for session=${sessionId}`, + ) } res.writeHead(200, { "Content-Type": "application/json" }) res.end(JSON.stringify({ success: true })) diff --git a/packages/mcp-server/src/index.ts b/packages/mcp-server/src/index.ts index 34693158..f2c58423 100644 --- a/packages/mcp-server/src/index.ts +++ b/packages/mcp-server/src/index.ts @@ -19,7 +19,6 @@ */ 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" @@ -61,7 +60,7 @@ import { serializeMxfile, wrapCellsInModel, } from "./pages.ts" -import { Autosaver, defaultDataDir } from "./persistence.ts" +import { Autosaver, defaultDataDir, expandHome } from "./persistence.ts" import { getShapeLibrary, SHAPE_LIBRARY_LIST } from "./shape-library.ts" import { validateAndFixXml } from "./xml-validation.ts" @@ -171,13 +170,6 @@ 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}"` @@ -904,6 +896,7 @@ function exportViaBrowser( live.exportFormat = undefined live.exportXml = undefined live.exportOptions = undefined + live.exportId = undefined } return exportData }) diff --git a/packages/mcp-server/src/new-diagram.ts b/packages/mcp-server/src/new-diagram.ts index 11bd09e6..a0a6d212 100644 --- a/packages/mcp-server/src/new-diagram.ts +++ b/packages/mcp-server/src/new-diagram.ts @@ -19,6 +19,19 @@ export function prepareNewDiagram( input: string, page: { pageId?: string; pageName?: string } = {}, ): NewDiagram { + // Bare cells get the root cells "0" and "1". A shape or edge with one of + // these ids would be renamed as a duplicate, breaking its edges. + if ( + !/<(mxGraphModel|mxfile)\b/.test(input) && + /]*\bid=["'][01]["'])(?=[^>]*\b(?:vertex|edge)=["']1["'])/.test( + input, + ) + ) { + return { + ok: false, + error: 'Cell ids "0" and "1" are the root cells, which are added automatically. Give shapes and edges ids starting at "2".', + } + } let xml = wrapCellsInModel(input) const { valid, error, fixed, fixes } = validateAndFixXml(xml) if (fixed) xml = fixed diff --git a/packages/mcp-server/src/pages.ts b/packages/mcp-server/src/pages.ts index 3e572419..7706c5c4 100644 --- a/packages/mcp-server/src/pages.ts +++ b/packages/mcp-server/src/pages.ts @@ -54,7 +54,7 @@ export function generatePageId(): string { /** Any cell besides the root cells "0" and "1" */ export const hasCells = (xml: string) => - /<(mxCell\b[^>]*\bid="(?![01]")|UserObject\b|object\b)/.test(xml) + /<(mxCell\b[^>]*\bid=["'](?![01]["'])|UserObject\b|object\b)/.test(xml) /** Cheap regex check — does the XML start with an root? */ export function isMxFile(xml: string): boolean { diff --git a/packages/mcp-server/src/persistence.ts b/packages/mcp-server/src/persistence.ts index edff059e..6a0c8aa2 100644 --- a/packages/mcp-server/src/persistence.ts +++ b/packages/mcp-server/src/persistence.ts @@ -23,11 +23,18 @@ import { hasCells } from "./pages.ts" const DELAY_MS = 1000 const MAX_FILES = 50 +/** Expand a leading ~ to the home directory (shells do this, MCP hosts don't). */ +export function expandHome(p: string): string { + if (p === "~") return homedir() + if (p.startsWith("~/") || p.startsWith("~\\")) return homedir() + p.slice(1) + return p +} + /** DRAWIO_DATA_DIR, default ~/.next-ai-drawio; "off" disables saving. */ export function defaultDataDir(): string | null { const dir = process.env.DRAWIO_DATA_DIR if (dir === "off") return null - return dir || join(homedir(), ".next-ai-drawio") + return dir ? expandHome(dir) : join(homedir(), ".next-ai-drawio") } export class Autosaver { diff --git a/packages/mcp-server/src/preview/preview.js b/packages/mcp-server/src/preview/preview.js index 2917d2f9..61d7f545 100644 --- a/packages/mcp-server/src/preview/preview.js +++ b/packages/mcp-server/src/preview/preview.js @@ -5,6 +5,7 @@ let pendingSvgBase = 0; // version the pending autosave was based on let pendingAiSvg = false; let pendingMcpExport = null; // 'png', 'svg' or 'xmlsvg' when MCP requested export let mcpExportSeq = 0; // number of the latest MCP export +let mcpExportId = null; // the server's id for it, sent back with the result let projectionExportActive = false; // page-targeted export: showing a transient single-page projection let forceReload = false; // reload the server state on the next poll even if the version is unchanged let noticeTimer = null; @@ -58,7 +59,7 @@ window.addEventListener('message', (e) => { fetch('/api/state', { method: 'POST', headers: { 'Content-Type': 'application/json' }, - body: JSON.stringify({ sessionId, exportData: d }) + body: JSON.stringify({ sessionId, exportData: d, exportId: mcpExportId }) }).catch(() => {}).finally(() => { // The timeout already ended this export if (seq !== mcpExportSeq) return; @@ -170,14 +171,6 @@ async function poll() { const r = await fetch('/api/state?sessionId=' + encodeURIComponent(sessionId)); if (!r.ok) return; const s = await r.json(); - // Handle sync request - server needs fresh state. Reset after a - // 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); - } // The server lost this session (e.g. it expired) and rebuilt it // with a blank diagram: push back what the browser shows. if (s.version < knownVersion && lastXml) { @@ -193,6 +186,17 @@ async function poll() { currentVersion = s.version; loadDiagram(s.xml, true); } + // Handle sync request - server needs fresh state. After the load + // above, so draw.io exports what it just loaded; never while a + // one-page projection is on screen, which would be sent as the + // whole document. Reset after a while in case draw.io never + // answers, so later syncs still run. + if (s.syncRequested && !pendingSyncExport && isReady && !projectionExportActive) { + pendingSyncExport = true; + pendingSyncBase = currentVersion; + iframe.contentWindow.postMessage(JSON.stringify({ action: 'export', format: 'xml' }), '*'); + setTimeout(() => { pendingSyncExport = false; }, 5000); + } // Handle export request from MCP server (png/svg). // // Plain export: capture whatever tab is currently displayed. @@ -207,6 +211,7 @@ async function poll() { if (s.exportFormat && !pendingMcpExport && isReady) { pendingMcpExport = s.exportFormat; const seq = ++mcpExportSeq; + mcpExportId = s.exportId; const extra = s.exportOptions || {}; const fireExport = () => { // mcpExport carries this export's number and is echoed diff --git a/packages/mcp-server/src/xml-validation.ts b/packages/mcp-server/src/xml-validation.ts index c1edca7b..5ea4b4fa 100644 --- a/packages/mcp-server/src/xml-validation.ts +++ b/packages/mcp-server/src/xml-validation.ts @@ -407,11 +407,13 @@ function findOrphanMxPoints( /** The first non-blank text under el, skipping a page's compressed data */ function findTextBetweenTags(el: Element | null): string | null { if (!el) return null + // A with only text holds the page compressed + const compressed = el.tagName === "diagram" && el.children.length === 0 for (const node of Array.from(el.childNodes)) { if (node.nodeType === 1) { const text = findTextBetweenTags(node as Element) if (text) return text - } else if (node.nodeType === 3 && el.tagName !== "diagram") { + } else if (node.nodeType === 3 && !compressed) { const text = node.textContent?.trim() if (text) return text.slice(0, 40) } diff --git a/packages/mcp-server/tests/diagram-operations.test.ts b/packages/mcp-server/tests/diagram-operations.test.ts index 90a8563b..77d4de13 100644 --- a/packages/mcp-server/tests/diagram-operations.test.ts +++ b/packages/mcp-server/tests/diagram-operations.test.ts @@ -124,3 +124,28 @@ describe("pages without a ", () => { expect(errors[0]?.message).toContain("could not be decompressed") }) }) + +describe("a wrapped mxCell with its wrapper's id", () => { + const doc = `` + + it("deletes the whole wrapper", () => { + const { result, errors } = applyDiagramOperations(doc, [ + { operation: "delete", cell_id: "u" }, + ]) + expect(errors).toEqual([]) + expect(result).not.toContain("UserObject") + }) + + it("replaces the wrapper on update", () => { + const { result, errors } = applyDiagramOperations(doc, [ + { + operation: "update", + cell_id: "u", + new_xml: ``, + }, + ]) + expect(errors).toEqual([]) + expect(result.match(/ { await postJson("/api/state", { sessionId: id, exportData: "data:image/png;base64,AAAA", + exportId: poll.exportId, }) expect(getState(id)?.exportOptions).toBeUndefined() }) + + it("ignores a late result of an export that already timed out", async () => { + const id = "mcp-export-late" + setState(id, "x") + requestExport(id, "png") + const first = JSON.parse( + (await request(`/api/state?sessionId=${id}`)).body, + ) + // The server gave up on the first export and asked for the next + requestExport(id, "svg") + await postJson("/api/state", { + sessionId: id, + exportData: "data:image/png;base64,LATE", + exportId: first.exportId, + }) + expect(getState(id)?.exportData).toBeUndefined() + expect(getState(id)?.exportFormat).toBe("svg") + }) }) describe("preview page", () => { diff --git a/packages/mcp-server/tests/persistence.test.ts b/packages/mcp-server/tests/persistence.test.ts index 41ca4a17..eb61d79e 100644 --- a/packages/mcp-server/tests/persistence.test.ts +++ b/packages/mcp-server/tests/persistence.test.ts @@ -10,7 +10,7 @@ import { utimesSync, writeFileSync, } from "node:fs" -import { tmpdir } from "node:os" +import { homedir, tmpdir } from "node:os" import { join } from "node:path" import { afterEach, describe, expect, it } from "vitest" import { Autosaver, defaultDataDir } from "../src/persistence.ts" @@ -88,4 +88,9 @@ describe("defaultDataDir", () => { delete process.env.DRAWIO_DATA_DIR expect(defaultDataDir()).toMatch(/\.next-ai-drawio$/) }) + + it("expands ~, which JSON configs pass on as it is", () => { + process.env.DRAWIO_DATA_DIR = "~/drawio-saves" + expect(defaultDataDir()).toBe(join(homedir(), "drawio-saves")) + }) }) diff --git a/packages/mcp-server/tests/wrap-cells.test.ts b/packages/mcp-server/tests/wrap-cells.test.ts index 6e22545f..c854d610 100644 --- a/packages/mcp-server/tests/wrap-cells.test.ts +++ b/packages/mcp-server/tests/wrap-cells.test.ts @@ -10,7 +10,8 @@ beforeAll(() => { installDomPolyfill() }) -import { wrapCellsInModel } from "../src/pages.ts" +import { prepareNewDiagram } from "../src/new-diagram.ts" +import { hasCells, wrapCellsInModel } from "../src/pages.ts" import { validateAndFixXml } from "../src/xml-validation.ts" const A = `` @@ -66,3 +67,33 @@ describe("wrapCellsInModel", () => { ) }) }) + +describe("hasCells", () => { + it("counts cells written with single quotes", () => { + expect(hasCells(``)).toBe(true) + expect(hasCells(``)).toBe( + false, + ) + }) +}) + +describe("prepareNewDiagram", () => { + it("rejects a shape that uses a root cell id", () => { + // It would clash with the added root cell "1" and be renamed, + // which breaks the edges that point to it + const cells = + `` + + `` + const out = prepareNewDiagram(cells) + expect(out.ok).toBe(false) + if (out.ok) return + expect(out.error).toMatch(/"0" and "1"/) + }) + + it("still accepts the root cells sent along with the shapes", () => { + const out = prepareNewDiagram( + `${A}`, + ) + expect(out.ok).toBe(true) + }) +}) diff --git a/packages/mcp-server/tests/xml-validation.test.ts b/packages/mcp-server/tests/xml-validation.test.ts index f828fc6a..b8334007 100644 --- a/packages/mcp-server/tests/xml-validation.test.ts +++ b/packages/mcp-server/tests/xml-validation.test.ts @@ -308,3 +308,13 @@ describe("text between tags", () => { ).toBe(true) }) }) + +describe("text directly under a page", () => { + it("fixes a literal \\n before the model of a page", () => { + const r = validateAndFixXml( + `\\n${model(``)}`, + ) + expect(r.valid).toBe(true) + expect(r.fixed).not.toContain("\\n") + }) +}) diff --git a/tests/e2e/diagram-content.spec.ts b/tests/e2e/diagram-content.spec.ts index 6be60458..a96427e8 100644 --- a/tests/e2e/diagram-content.spec.ts +++ b/tests/e2e/diagram-content.spec.ts @@ -135,6 +135,36 @@ test("display_diagram replaces the document with the fixed diagram", async ({ await expect(canvas.getByText("Second", { exact: true })).toHaveCount(0) }) +test("an edit with a fixable cell is fixed, not rejected", async ({ + page: p, +}) => { + // Chrome's DOMParser puts a next to the cell, which used + // to count as a second cell + const canvas = await mockReplies(p, [ + streamedToolCall("display_diagram", { xml: cell("a", "Alpha", 40) }), + streamedToolCall("edit_diagram", { + operations: [ + { + operation: "add", + cell_id: "c", + new_xml: cell("c", "Gamma", 400).replace( + "", + "", + ), + }, + ], + }), + ]) + await sendMessage(p, "Draw a box") + await waitForCompleteCount(p, 1) + await sendMessage(p, "Add another box") + await waitForCompleteCount(p, 2) + await expect(canvas.getByText("Gamma", { exact: true })).toBeVisible({ + timeout: 15000, + }) + await expect(p.getByText(/exactly one cell/)).toHaveCount(0) +}) + test("edit_diagram applies all operations or none", async ({ page: p }) => { const canvas = await mockReplies(p, [ streamedToolCall("display_diagram", {