From 080f44716f8bce4f99686d9f37789319e27b91d2 Mon Sep 17 00:00:00 2001 From: "dayuan.jiang" Date: Mon, 5 Oct 2026 10:52:20 +0900 Subject: [PATCH] fix(mcp-server): count a one-page view only for that page, and more review fixes Found by the second PR review: - get_diagram with a page selector, or a rejected edit's error, counted the whole document as seen, so an edit on another page could overwrite the user's change there. A one-page view now counts for all pages only if the others are unchanged; otherwise the reply says to get them. - add_page accepted shapes with the root cell ids "0" and "1" and renamed them, breaking their edges. The check also missed ids on UserObject wrappers and ids written with spaces around the "=". - Root cells written over two lines were kept as an extra layer, cells with id = "a" did not count as cells, and CDATA text before a page's model passed the check although draw.io cannot open the page. - Auto-save cleanup deleted the user's own files that start with mcp-. Only names in the session id format are removed now. - Restoring a history entry dropped edits made in the browser since the last entry. They are added to history first. - A session whose state expired showed a blank page, and the next change overwrote its auto-save file. The saved file is loaded instead. - An edit on a page export's one-page projection, made before the real document was back, replaced the whole document. - A late sync reply could overwrite a newer edit: each sync export is numbered, and the server ignores replies older than the current state. - screenshot_diagram could return another session's image after start_session ran during its retries. --- packages/mcp-server/src/edit-gate.ts | 35 +++++++++- packages/mcp-server/src/http-server.ts | 30 +++++++-- packages/mcp-server/src/index.ts | 66 ++++++++++++++----- packages/mcp-server/src/new-diagram.ts | 32 +++++---- packages/mcp-server/src/pages.ts | 14 +++- packages/mcp-server/src/persistence.ts | 19 +++++- packages/mcp-server/src/preview/preview.js | 27 +++++--- packages/mcp-server/src/xml-validation.ts | 6 +- packages/mcp-server/tests/edit-gate.test.ts | 34 +++++++++- packages/mcp-server/tests/http-server.test.ts | 55 ++++++++++++++++ packages/mcp-server/tests/persistence.test.ts | 27 ++++++-- packages/mcp-server/tests/wrap-cells.test.ts | 40 ++++++++++- .../mcp-server/tests/xml-validation.test.ts | 8 +++ 13 files changed, 334 insertions(+), 59 deletions(-) diff --git a/packages/mcp-server/src/edit-gate.ts b/packages/mcp-server/src/edit-gate.ts index 9ce1a9c2..41bf766f 100644 --- a/packages/mcp-server/src/edit-gate.ts +++ b/packages/mcp-server/src/edit-gate.ts @@ -17,7 +17,14 @@ * change: the set of pages, each page's name, and each page's cell tree * (tags + sorted attributes + text). Byte equality is kept as a fast path. */ -import { isMxGraphModel, normalizeToMxfile, parseMxfile } from "./pages.ts" +import { + findPageElement, + isMxGraphModel, + normalizeToMxfile, + type PageSelector, + parseMxfile, + serializeMxfile, +} from "./pages.ts" export type EditGateResult = | { ok: true } @@ -100,3 +107,29 @@ export function checkEditGate( } return { ok: true } } + +/** + * The model was shown only the selected page of liveXml (get_diagram with a + * page selector, or a rejected edit's error). It has seen the whole document + * if the other pages are as it last saw them, or if it saw nothing before + * (it then has no old copy of them to edit from). Returns the new + * lastSeenXml: liveXml, or lastSeenXml unchanged. + */ +export function markPageSeen( + lastSeenXml: string, + liveXml: string, + selector: PageSelector, +): string { + if (!lastSeenXml) return liveXml + const otherPages = (xml: string) => { + const normalized = normalizeToMxfile(xml) + const doc = normalized ? parseMxfile(normalized) : null + if (!doc) return null + findPageElement(doc, selector)?.element.remove() + return contentFingerprint(serializeMxfile(doc)) + } + const before = otherPages(lastSeenXml) + return before !== null && before === otherPages(liveXml) + ? liveXml + : lastSeenXml +} diff --git a/packages/mcp-server/src/http-server.ts b/packages/mcp-server/src/http-server.ts index 2125a050..7c08d5ac 100644 --- a/packages/mcp-server/src/http-server.ts +++ b/packages/mcp-server/src/http-server.ts @@ -40,7 +40,7 @@ import { updateLastHistorySvg, } from "./history.ts" import { log } from "./logger.ts" -import { BLANK_MXFILE } from "./pages.ts" +import { BLANK_MXFILE, hasCells } from "./pages.ts" // Configurable draw.io embed URL for private deployments const DRAWIO_BASE_URL = @@ -87,10 +87,12 @@ function ensureSessionStateInitialized(sessionId: string): void { if (!isValidSessionId(sessionId)) return if (stateStore.has(sessionId)) return + // The session's saved diagram, so a blank page never replaces that file. // Not a change worth saving: the browser fills it on its next push // A blank diagram keeps the draw.io spinner (spin=1) from waiting // forever when no load(xml) is ever sent - setState(sessionId, BLANK_MXFILE, undefined, false, false) + const saved = savedStateLoader?.(sessionId) + setState(sessionId, saved || BLANK_MXFILE, undefined, false, false) } interface SessionState { @@ -142,6 +144,16 @@ export function onStateChange( stateListener = listener } +// Reads a session's saved diagram when its state is created again (it +// expired, or the MCP process restarted) +let savedStateLoader: ((sessionId: string) => string | null) | null = null + +export function onSessionRecreate( + loader: (sessionId: string) => string | null, +): void { + savedStateLoader = loader +} + export function setState( sessionId: string, xml: string, @@ -461,11 +473,15 @@ function handleStateApi( // The browser edited a version older than the latest AI write // (it has not loaded that write yet). Keep the AI write; the - // browser loads it on its next poll. + // browser loads it on its next poll. A sync reply is also + // stale after a newer write of the browser's own (a user + // edit saved while the export ran). const current = stateStore.get(sessionId) if ( typeof data.baseVersion === "number" && - data.baseVersion < (current?.serverVersion ?? 0) + (data.baseVersion < (current?.serverVersion ?? 0) || + (data.source === "sync" && + data.baseVersion < (current?.version ?? 0))) ) { let savedToHistory = false if (data.source === "sync") { @@ -566,6 +582,12 @@ function handleRestoreApi( return } + // Edits in the browser since the last entry are not in history + // yet: keep them, so the restore can be undone + const current = stateStore.get(sessionId) + if (current && hasCells(current.xml)) { + addHistory(sessionId, current.xml, current.svg) + } const newVersion = setState(sessionId, entry.xml) addHistory(sessionId, entry.xml, entry.svg) diff --git a/packages/mcp-server/src/index.ts b/packages/mcp-server/src/index.ts index f2c58423..49c14a6e 100644 --- a/packages/mcp-server/src/index.ts +++ b/packages/mcp-server/src/index.ts @@ -27,13 +27,14 @@ import type { DiagramOperation } from "./diagram-operations.ts" import { installDomPolyfill } from "./dom.ts" import { DRAWING_GUIDE } from "./drawing-guide.ts" import { editDiagram, targetPageXml } from "./edit-diagram.ts" -import { checkEditGate } from "./edit-gate.ts" +import { checkEditGate, markPageSeen } from "./edit-gate.ts" import { addHistory } from "./history.ts" import { type ExportFormat, type ExportOptions, getServerPort, getState, + onSessionRecreate, onStateChange, requestExport, requestSync, @@ -44,7 +45,7 @@ import { } from "./http-server.ts" import { parseDrawioFileContent } from "./load-diagram.ts" import { log } from "./logger.ts" -import { prepareNewDiagram } from "./new-diagram.ts" +import { prepareNewDiagram, reservedIdError } from "./new-diagram.ts" import { addPageToDoc, deletePageFromDoc, @@ -75,6 +76,7 @@ const config = { // Keep each session's latest diagram on disk, so it survives this process const autosaver = new Autosaver(defaultDataDir()) onStateChange((sessionId, xml) => autosaver.schedule(sessionId, xml)) +onSessionRecreate((sessionId) => autosaver.load(sessionId)) // Session state (single session for simplicity) let currentSession: { @@ -643,18 +645,27 @@ server.registerTool( : "edit_diagram rejected: the model has not seen the diagram yet", ) // 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 + // seen it and can retry without a get_diagram round-trip, + // unless other pages changed too. + const liveXml = browserState?.xml || currentSession.xml + currentSession.lastSeenXml = markPageSeen( + currentSession.lastSeenXml, + liveXml, + pageSelector, + ) 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." + const next = + currentSession.lastSeenXml === liveXml + ? "Build your operations on this XML and retry." + : "Other pages changed too: call get_diagram without a page selector, then retry." return { content: [ { type: "text", - text: `Error: ${reason}\n\nCurrent XML of ${describeSelector(pageSelector)}:\n\n${targetPageXml(currentSession.xml, pageSelector)}\n\nBuild your operations on this XML and retry.`, + text: `Error: ${reason}\n\nCurrent XML of ${describeSelector(pageSelector)}:\n\n${targetPageXml(currentSession.xml, pageSelector)}\n\n${next}`, }, ], isError: true, @@ -796,16 +807,28 @@ server.registerTool( } } - // The model is now looking at the current state. Record the raw - // store value — the gate's fast path is plain string equality - // against the store, with a structural comparison as fallback. - currentSession.lastSeenXml = browserState?.xml || currentSession.xml - const pageSelector = pickPageSelector({ page_id, page_name, page_index, }) + + // The model is now looking at the current state. Record the raw + // store value — the gate's fast path is plain string equality + // against the store, with a structural comparison as fallback. + // One page shown counts for all only if the others are unchanged. + const liveXml = browserState?.xml || currentSession.xml + currentSession.lastSeenXml = hasPageSelector(pageSelector) + ? markPageSeen( + currentSession.lastSeenXml, + liveXml, + pageSelector, + ) + : liveXml + const otherPagesNote = + currentSession.lastSeenXml === liveXml + ? "" + : "\n\nNote: other pages changed since you last saw them. Call get_diagram without a page selector before editing." const doc = parseMxfile(currentSession.xml) const pages = doc ? listPagesFromDoc(doc) : [] const pageList = pages.length @@ -844,7 +867,7 @@ server.registerTool( content: [ { type: "text", - text: `Page ${projection.index} ("${projection.name}"):\n\n${projection.xml}\n\n${pageList}${staleNote}`, + text: `Page ${projection.index} ("${projection.name}"):\n\n${projection.xml}\n\n${pageList}${staleNote}${otherPagesNote}`, }, ], } @@ -1016,13 +1039,13 @@ server.registerTool( } 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( - currentSession.id, - "png", - projectionXml, - { width, pageId }, - ) + data = await exportViaBrowser(sessionId, "png", projectionXml, { + width, + pageId, + }) if (!data || data.length <= MAX_SCREENSHOT_CHARS) break } if (!data) { @@ -1502,6 +1525,13 @@ server.registerTool( // If caller provided XML, validate it before splicing it in so we // never get a half-broken mxfile written to the session. + const reserved = xml && reservedIdError(xml) + if (reserved) { + return { + content: [{ type: "text", text: `Error: ${reserved}` }], + isError: true, + } + } let cleanXml: string | undefined = xml && wrapCellsInModel(xml) if (cleanXml) { const { valid, error, fixed, fixes } = diff --git a/packages/mcp-server/src/new-diagram.ts b/packages/mcp-server/src/new-diagram.ts index a0a6d212..2443695d 100644 --- a/packages/mcp-server/src/new-diagram.ts +++ b/packages/mcp-server/src/new-diagram.ts @@ -9,6 +9,23 @@ export type NewDiagram = | { ok: true; xml: string; fixes: string[] } | { ok: false; error: string } +/** + * 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. Its id may + * be on a or wrapper. Returns the error for the model, + * or null. + */ +export function reservedIdError(input: string): string | null { + const ID = String.raw`\bid\s*=\s*["'][01]["']` + const shape = new RegExp( + String.raw`]*${ID})(?=[^>]*\b(?:vertex|edge)\s*=\s*["']1["'])|<(?:UserObject|object)\b[^>]*${ID}`, + ) + if (/<(mxGraphModel|mxfile)\b/.test(input) || !shape.test(input)) { + return null + } + return 'Cell ids "0" and "1" are the root cells, which are added automatically. Give shapes and edges ids starting at "2".' +} + /** * Bare cells get the wrapper and root cells first, since the strict parser * rejects several top-level elements. Then the XML is validated and @@ -19,19 +36,8 @@ 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".', - } - } + const reserved = reservedIdError(input) + if (reserved) return { ok: false, error: reserved } 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 7706c5c4..15a3a87c 100644 --- a/packages/mcp-server/src/pages.ts +++ b/packages/mcp-server/src/pages.ts @@ -54,7 +54,9 @@ 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\s*=\s*["'](?![01]["'])|UserObject\b|object\b)/.test( + xml, + ) /** Cheap regex check — does the XML start with an root? */ export function isMxFile(xml: string): boolean { @@ -121,8 +123,14 @@ export function wrapCellsInModel(xml: string): string { content = content.slice(0, end) } content = content - .replace(/]*\bid=["']0["'][^>]*(?:\/>|><\/mxCell>)/g, "") - .replace(/]*\bid=["']1["'][^>]*(?:\/>|><\/mxCell>)/g, "") + .replace( + /]*\bid\s*=\s*["']0["'][^>]*(?:\/>|>\s*<\/mxCell>)/g, + "", + ) + .replace( + /]*\bid\s*=\s*["']1["'][^>]*(?:\/>|>\s*<\/mxCell>)/g, + "", + ) .trim() return `${ROOT_CELLS}${content}` } diff --git a/packages/mcp-server/src/persistence.ts b/packages/mcp-server/src/persistence.ts index 6a0c8aa2..9568f651 100644 --- a/packages/mcp-server/src/persistence.ts +++ b/packages/mcp-server/src/persistence.ts @@ -10,6 +10,7 @@ import { existsSync, mkdirSync, readdirSync, + readFileSync, renameSync, statSync, unlinkSync, @@ -54,6 +55,17 @@ export class Autosaver { return this.dir ? join(this.dir, `${sessionId}.drawio`) : null } + /** The session's saved diagram, or null. */ + load(sessionId: string): string | null { + const path = this.pathFor(sessionId) + if (!path || !existsSync(path)) return null + try { + return readFileSync(path, "utf-8") + } catch { + return null + } + } + schedule(sessionId: string, xml: string): void { if (!this.dir) return const previous = this.pending.get(sessionId) @@ -93,10 +105,11 @@ export class Autosaver { private removeOldest(): void { if (!this.dir) return const dir = this.dir - // Only our own session files: DRAWIO_DATA_DIR may be a folder - // that also holds the user's diagrams + // Only our own session files (mcp-