fix(mcp): replace edit_diagram 30s time gate with content comparison (#890)

* fix(mcp): keep diagram context valid during edits

Closes #885

* fix(mcp): replace edit_diagram time gate with content comparison

The 30s wall-clock gate rejected slow-but-correct clients (#885).
Instead of a timeout, remember the exact state-store XML the model
last saw (get_diagram / create_new_diagram / edit_diagram / page CRUD)
and reject edit_diagram only when the live browser state differs -
i.e. the user made edits the model hasn't seen yet. Slow reasoning
no longer trips the gate, while unseen manual edits still do.

* docs(mcp): align edit_diagram/get_diagram descriptions with content-based gate

The 'You MUST call get_diagram BEFORE this tool' requirement and the
'Skipping get_diagram WILL cause user's changes to be LOST' warning no
longer match server behavior: a stale edit is rejected with no side
effects, never silently applied. Describe the freshness check instead,
and direct get_diagram usage at its real purpose - learning the current
diagram content when the model doesn't already know it.

* fix(mcp): compare diagram content structurally in the edit gate

draw.io re-serialises the document when pushing state back (attribute
order, pretty-printing, regenerated diagram ids, viewport attributes,
mxfile host), so byte comparison could flag an unchanged diagram as
stale. Fingerprint what a user can actually change instead - page set,
page names, and each page's root cell tree with sorted attributes -
keeping byte equality as the fast path. A bare mxGraphModel now also
fingerprints identically to its single-page mxfile wrapping.

* fix(mcp): don't compare page names against bare mxGraphModel pushes

A bare <mxGraphModel> pushed by the embed/sync path carries no page name,
so normalizeToMxfile invents "Page-1" — falsely reading any custom page
name as a content change and re-triggering the stale rejection on every
edit. When either side of the gate comparison is a bare mxGraphModel,
fingerprint cell trees only; full-mxfile comparisons still detect renames.

* chore(mcp): bump version to 0.2.2

* chore(mcp): sync package-lock.json version to 0.2.2

---------

Co-authored-by: dayuan.jiang <jdy.toh@gmail.com>

packages/mcp-server/package-lock.json

@@ -1,12 +1,12 @@
 {
     "name": "@next-ai-drawio/mcp-server",
-    "version": "0.2.1",
+    "version": "0.2.2",
     "lockfileVersion": 3,
     "requires": true,
     "packages": {
         "": {
             "name": "@next-ai-drawio/mcp-server",
-            "version": "0.2.1",
+            "version": "0.2.2",
             "license": "Apache-2.0",
             "dependencies": {
                 "@modelcontextprotocol/sdk": "^1.0.4",

packages/mcp-server/package.json

@@ -1,6 +1,6 @@
 {
     "name": "@next-ai-drawio/mcp-server",
-    "version": "0.2.1",
+    "version": "0.2.2",
     "description": "MCP server for Next AI Draw.io - AI-powered diagram generation with real-time browser preview",
     "type": "module",
     "main": "dist/index.js",

packages/mcp-server/src/edit-gate.ts

@@ -0,0 +1,102 @@
+/**
+ * Workflow gate for edit_diagram.
+ *
+ * Instead of a wall-clock timeout (the old 30s rule rejected slow-but-correct
+ * clients, see #885), we compare content: `lastSeenXml` is the state-store
+ * XML the model last saw (get_diagram) or wrote itself (create_new_diagram /
+ * edit_diagram / page CRUD). The store only changes on server writes or
+ * browser pushes (user autosave, sync exports), so if the live store still
+ * matches `lastSeenXml`, nothing happened that the model hasn't seen — the
+ * edit is safe no matter how much time passed.
+ *
+ * "Matches" is structural, not byte-for-byte: draw.io re-serialises the
+ * document when it pushes state back (different attribute order, pretty-
+ * printed whitespace, regenerated diagram ids, viewport attributes like
+ * dx/dy/pageWidth on <mxGraphModel>, a different mxfile host). None of that
+ * is a user edit, so the fingerprint keeps only what a user can actually
+ * 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.js"
+
+export type EditGateResult =
+    | { ok: true }
+    | { ok: false; reason: "no-context" | "stale" }
+
+/**
+ * Canonical serialisation of an element subtree: tag + attributes sorted by
+ * name + child elements in order + non-whitespace text. Whitespace-only text
+ * nodes (pretty-printing) are dropped.
+ */
+function canonicalizeElement(el: Element): string {
+    const attrs = Array.from(el.attributes)
+        .map((a) => `${a.name}=${JSON.stringify(a.value)}`)
+        .sort()
+        .join(" ")
+    let children = ""
+    for (const child of Array.from(el.childNodes)) {
+        if (child.nodeType === 1) {
+            children += canonicalizeElement(child as Element)
+        } else if (child.nodeType === 3 || child.nodeType === 4) {
+            const text = (child.textContent ?? "").trim()
+            if (text) children += JSON.stringify(text)
+        }
+    }
+    return `<${el.tagName} ${attrs}>${children}</${el.tagName}>`
+}
+
+/**
+ * Structural fingerprint of a diagram document: page names + each page's
+ * <root> subtree, ignoring everything draw.io rewrites on re-serialisation
+ * (mxfile/mxGraphModel attributes, diagram ids, formatting). A bare
+ * <mxGraphModel> fingerprints identically to its single-page mxfile wrapping.
+ * Unparseable input falls back to the trimmed raw string, degrading to the
+ * plain string comparison.
+ *
+ * `includeNames=false` drops page names from the fingerprint — used when the
+ * other side of a comparison is a bare <mxGraphModel>, which carries no page
+ * name at all (normalizeToMxfile would invent "Page-1", falsely mismatching
+ * any real page name).
+ */
+export function contentFingerprint(xml: string, includeNames = true): string {
+    const normalized = normalizeToMxfile(xml)
+    const doc = normalized ? parseMxfile(normalized) : null
+    if (!doc) return xml.trim()
+    const pages: string[] = []
+    doc.querySelectorAll("diagram").forEach((d) => {
+        const name = includeNames ? d.getAttribute("name") || "" : ""
+        const root = d.querySelector("root")
+        // No <root> means the page content is not plain XML (e.g. draw.io's
+        // compressed format) — fingerprint the raw text instead.
+        const body = root
+            ? canonicalizeElement(root)
+            : (d.textContent || "").trim()
+        pages.push(`${name}=${body}`)
+    })
+    return pages.join("\n")
+}
+
+export function checkEditGate(
+    lastSeenXml: string,
+    liveXml: string,
+): EditGateResult {
+    // Model never fetched or produced any diagram state in this session.
+    if (!lastSeenXml) return { ok: false, reason: "no-context" }
+    // Browser state moved since the model last looked (e.g. manual user
+    // edits): force a re-fetch so update/delete operations don't build on
+    // stale cell contents. An empty liveXml means the store has no entry to
+    // compare against, so there is nothing newer to have missed.
+    if (liveXml && liveXml !== lastSeenXml) {
+        // A bare <mxGraphModel> on either side carries no page name, so
+        // comparing names would mismatch against anything not called
+        // "Page-1". Compare cell trees only in that case.
+        const includeNames =
+            !isMxGraphModel(liveXml) && !isMxGraphModel(lastSeenXml)
+        if (
+            contentFingerprint(liveXml, includeNames) !==
+            contentFingerprint(lastSeenXml, includeNames)
+        )
+            return { ok: false, reason: "stale" }
+    }
+    return { ok: true }
+}

packages/mcp-server/src/index.ts

@@ -44,6 +44,7 @@ import {
     applyDiagramOperations,
     type DiagramOperation,
 } from "./diagram-operations.js"
+import { checkEditGate } from "./edit-gate.js"
 import { addHistory } from "./history.js"
 import {
     getState,
@@ -79,7 +80,12 @@ let currentSession: {
     id: string
     xml: string
     version: number
-    lastGetDiagramTime: number // Track when get_diagram was last called (for enforcing workflow)
+    // The exact state-store XML the model last saw (get_diagram) or wrote
+    // itself (create/edit/page CRUD). The store only changes on server
+    // writes or browser pushes (user autosave / sync), so edit_diagram can
+    // detect unseen user edits by comparing the live store against this.
+    // Empty = no diagram context established yet.
+    lastSeenXml: string
 } | null = null
 
 // Create MCP server
@@ -164,15 +170,11 @@ server.prompt(
 - Use rename_page / delete_page for management
 - edit_diagram, get_diagram, and export_diagram all accept optional page_id / page_name / page_index — when omitted they target the first page
 
-## Adding Elements to an Existing Page
+## Editing a Page (add / update / delete cells)
-1. Use edit_diagram with "add" operation, optionally with a page selector
+1. Call edit_diagram with your operations, optionally with a page selector
-2. Provide a unique cell_id and complete mxCell XML
+2. If you don't know the current cell IDs or structure, call get_diagram first
-3. No need to call get_diagram first - the server fetches latest state automatically
+3. For add/update, provide the cell_id and complete mxCell XML
-
+4. No need to call get_diagram before every edit: the server rejects the edit (with no side effects) if the user changed the diagram in the browser since you last saw it, and tells you to call get_diagram once and retry
-## Modifying or Deleting Existing Elements
-1. FIRST call get_diagram to see current cell IDs and page structure
-2. THEN call edit_diagram with "update" or "delete" operations
-3. For update, provide the cell_id and complete new mxCell XML
 
 ## 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.
@@ -205,7 +207,7 @@ server.registerTool(
                 id: sessionId,
                 xml: "",
                 version: 0,
-                lastGetDiagramTime: 0,
+                lastSeenXml: "",
             }
 
             // Open browser
@@ -376,10 +378,12 @@ COMMON STYLES:
             // Update session state
             currentSession.xml = xml
             currentSession.version++
-            currentSession.lastGetDiagramTime = Date.now()
 
-            // Push to embedded server state
+            // Push to embedded server state. The model just authored this
+            // exact XML, so record it as seen — edit_diagram may follow
+            // without a redundant get_diagram round-trip.
             setState(currentSession.id, xml)
+            currentSession.lastSeenXml = xml
 
             // Save AI result (no SVG yet - will be captured by browser)
             addHistory(currentSession.id, xml, "")
@@ -423,13 +427,13 @@ server.registerTool(
     {
         description:
             "Edit a specific page in the current diagram by ID-based operations (update/add/delete cells).\n\n" +
-            "⚠️ REQUIRED: You MUST call get_diagram BEFORE this tool!\n" +
+            "Freshness: the server remembers the last diagram state you have seen, and rejects this call " +
-            "This fetches the latest state from the browser including any manual user edits.\n" +
+            "only if the user edited the diagram in the browser since then. You do NOT need to call " +
-            "Skipping get_diagram WILL cause user's changes to be LOST.\n\n" +
+            "get_diagram before every edit — if your view is stale, the call is rejected (with no side " +
-            "Workflow:\n" +
+            "effects) and the error tells you to call get_diagram once and retry.\n\n" +
-            "1. Call get_diagram to see current cell IDs, page structure, and active page\n" +
+            "Call get_diagram first only when you don't know the current diagram content (cell IDs, " +
-            "2. Use the returned XML to construct your edit operations\n" +
+            "structure) — e.g. the diagram wasn't created in this conversation, or you're unsure your " +
-            "3. Call edit_diagram with your operations and (optionally) a page selector\n\n" +
+            "memory of it is accurate.\n\n" +
             "Multi-page targeting:\n" +
             "- page_id / page_name / page_index are optional; when all omitted, the FIRST page is targeted\n" +
             "- Use list_pages to discover what pages exist\n\n" +
@@ -482,27 +486,6 @@ server.registerTool(
                 }
             }
 
-            // Enforce workflow: require get_diagram to be called first
-            const timeSinceGet = Date.now() - currentSession.lastGetDiagramTime
-            if (timeSinceGet > 30000) {
-                // 30 seconds
-                log.warn(
-                    "edit_diagram called without recent get_diagram - rejecting to prevent data loss",
-                )
-                return {
-                    content: [
-                        {
-                            type: "text",
-                            text:
-                                "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.",
-                        },
-                    ],
-                    isError: true,
-                }
-            }
-
             // Fetch latest state from browser. Re-normalise to mxfile: the
             // embed/sync path can hand back a bare <mxGraphModel>, and adopting
             // it verbatim would silently strip a multi-page document down to
@@ -526,6 +509,37 @@ server.registerTool(
                 }
             }
 
+            // 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
+            // long as nothing changed in the browser meanwhile (#885).
+            const gate = checkEditGate(
+                currentSession.lastSeenXml,
+                browserState?.xml ?? "",
+            )
+            if (!gate.ok) {
+                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",
+                )
+                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.",
+                        },
+                    ],
+                    isError: true,
+                }
+            }
+
             const pageSelector = pickPageSelector({
                 page_id,
                 page_name,
@@ -601,8 +615,10 @@ server.registerTool(
             currentSession.xml = result
             currentSession.version++
 
-            // Push to embedded server
+            // Push to embedded server; the pushed XML is now the latest
+            // state the model has seen.
             setState(currentSession.id, result)
+            currentSession.lastSeenXml = result
 
             // Save AI result (no SVG yet - will be captured by browser)
             addHistory(currentSession.id, result, "")
@@ -641,8 +657,9 @@ server.registerTool(
     {
         description:
             "Get the current diagram XML (fetches latest from browser, including user's manual edits). " +
-            "Call this BEFORE edit_diagram if you need to update or delete existing elements, " +
+            "Call this when you don't know the current diagram content (cell IDs, pages, structure) — " +
-            "so you can see the current cell IDs, pages, and structure.\n\n" +
+            "e.g. before editing a diagram you didn't create in this conversation, or after edit_diagram " +
+            "was rejected because the user changed the diagram in the browser.\n\n" +
             "Returns the full <mxfile> by default. If a page selector is provided, returns just that page's <mxGraphModel> embedded in a one-page <mxfile> wrapper.",
         inputSchema: {
             ...pageSelectorSchema,
@@ -677,9 +694,6 @@ server.registerTool(
                 }
             }
 
-            // Mark that get_diagram was called (for edit_diagram workflow check)
-            currentSession.lastGetDiagramTime = Date.now()
-
             // Fetch latest state from browser, re-normalising to mxfile so a
             // bare <mxGraphModel> pushed back by the embed/sync path doesn't
             // strip page structure (see edit_diagram for the same guard).
@@ -700,6 +714,11 @@ 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,
@@ -1070,11 +1089,11 @@ async function loadMxfileForMutation(): Promise<
             addHistory(sessionRef.id, sessionRef.xml, browserState?.svg || "")
             sessionRef.xml = newXml
             sessionRef.version++
-            // Page CRUD updates the structure that get_diagram would return,
-            // so refresh the workflow timestamp — subsequent edit_diagram
-            // calls don't need a redundant get_diagram round-trip.
-            sessionRef.lastGetDiagramTime = Date.now()
             setState(sessionRef.id, newXml)
+            // The model just wrote this exact state, so mark it as seen —
+            // subsequent edit_diagram calls don't need a redundant
+            // get_diagram round-trip.
+            sessionRef.lastSeenXml = newXml
             addHistory(sessionRef.id, newXml, "")
         },
     }

packages/mcp-server/tests/edit-gate.test.ts

@@ -0,0 +1,132 @@
+/**
+ * Unit tests for the edit_diagram workflow gate (edit-gate.ts).
+ *
+ * The gate replaced the old 30-second wall-clock rule (#885): an edit is
+ * allowed when the model has seen the current browser state, no matter how
+ * long ago — and rejected when the browser state moved since. "Seen" is
+ * judged structurally, so draw.io's re-serialisation of the same content
+ * (attribute order, whitespace, viewport attributes, wrapper shape) never
+ * reads as a user edit.
+ */
+
+import { DOMParser } from "linkedom"
+import { beforeAll, describe, expect, it } from "vitest"
+
+beforeAll(() => {
+    ;(globalThis as any).DOMParser = DOMParser
+})
+
+import { checkEditGate, contentFingerprint } from "../src/edit-gate.js"
+
+const XML_A = `<mxfile host="app.diagrams.net"><diagram id="p1" name="Page-1"><mxGraphModel><root><mxCell id="0"/><mxCell id="1" parent="0"/><mxCell id="box1" value="Hello" style="rounded=0;" vertex="1" parent="1"><mxGeometry x="40" y="40" width="120" height="60" as="geometry"/></mxCell></root></mxGraphModel></diagram></mxfile>`
+
+// The same document as draw.io re-serialises it on autosave: different host,
+// regenerated diagram id, viewport attributes on mxGraphModel, re-ordered
+// cell attributes, pretty-printed whitespace.
+const XML_A_RESERIALIZED = `<mxfile host="embed.diagrams.net">
+  <diagram id="regenerated-id" name="Page-1">
+    <mxGraphModel dx="1596" dy="743" grid="1" pageWidth="827" pageHeight="1169">
+      <root>
+        <mxCell id="0" />
+        <mxCell id="1" parent="0" />
+        <mxCell id="box1" parent="1" style="rounded=0;" value="Hello" vertex="1">
+          <mxGeometry height="60" width="120" x="40" y="40" as="geometry" />
+        </mxCell>
+      </root>
+    </mxGraphModel>
+  </diagram>
+</mxfile>`
+
+// A real user edit: box1 moved to a different position.
+const XML_B = XML_A.replace('x="40" y="40"', 'x="300" y="200"')
+
+// Bare mxGraphModel with identical page content to XML_A.
+const XML_A_BARE = `<mxGraphModel><root><mxCell id="0"/><mxCell id="1" parent="0"/><mxCell id="box1" value="Hello" style="rounded=0;" vertex="1" parent="1"><mxGeometry x="40" y="40" width="120" height="60" as="geometry"/></mxCell></root></mxGraphModel>`
+
+describe("checkEditGate", () => {
+    it("rejects when no diagram context was ever established", () => {
+        expect(checkEditGate("", XML_A)).toEqual({
+            ok: false,
+            reason: "no-context",
+        })
+    })
+
+    it("allows when the browser state is exactly what the model saw", () => {
+        expect(checkEditGate(XML_A, XML_A)).toEqual({ ok: true })
+    })
+
+    it("allows when the browser state is a re-serialisation of the same content", () => {
+        expect(checkEditGate(XML_A, XML_A_RESERIALIZED)).toEqual({ ok: true })
+    })
+
+    it("rejects when a cell actually changed", () => {
+        expect(checkEditGate(XML_A, XML_B)).toEqual({
+            ok: false,
+            reason: "stale",
+        })
+    })
+
+    it("rejects a real edit even when wrapped in re-serialisation noise", () => {
+        const movedAndReserialized = XML_A_RESERIALIZED.replace(
+            'x="40" y="40"',
+            'x="300" y="200"',
+        )
+        expect(checkEditGate(XML_A, movedAndReserialized)).toEqual({
+            ok: false,
+            reason: "stale",
+        })
+    })
+
+    it("allows when the store has no live entry to compare against", () => {
+        expect(checkEditGate(XML_A, "")).toEqual({ ok: true })
+    })
+
+    // A bare <mxGraphModel> push carries no page name, so the gate must not
+    // compare the invented "Page-1" wrapper name against the real one.
+    it("allows a bare mxGraphModel push when the page has a custom name", () => {
+        const seenRenamed = XML_A.replace('name="Page-1"', 'name="Arch"')
+        expect(checkEditGate(seenRenamed, XML_A_BARE)).toEqual({ ok: true })
+    })
+
+    it("still rejects a bare mxGraphModel push whose cells changed", () => {
+        const seenRenamed = XML_A.replace('name="Page-1"', 'name="Arch"')
+        const bareMoved = XML_A_BARE.replace('x="40" y="40"', 'x="300" y="200"')
+        expect(checkEditGate(seenRenamed, bareMoved)).toEqual({
+            ok: false,
+            reason: "stale",
+        })
+    })
+})
+
+describe("contentFingerprint", () => {
+    it("is invariant under draw.io re-serialisation", () => {
+        expect(contentFingerprint(XML_A)).toBe(
+            contentFingerprint(XML_A_RESERIALIZED),
+        )
+    })
+
+    it("treats a bare mxGraphModel like its one-page mxfile wrapping", () => {
+        expect(contentFingerprint(XML_A_BARE)).toBe(contentFingerprint(XML_A))
+    })
+
+    it("changes when a cell attribute changes", () => {
+        expect(contentFingerprint(XML_A)).not.toBe(contentFingerprint(XML_B))
+    })
+
+    it("changes when a page is renamed", () => {
+        const renamed = XML_A.replace('name="Page-1"', 'name="Renamed"')
+        expect(contentFingerprint(XML_A)).not.toBe(contentFingerprint(renamed))
+    })
+
+    it("changes when a page is added", () => {
+        const twoPages = XML_A.replace(
+            "</mxfile>",
+            `<diagram id="p2" name="Page-2"><mxGraphModel><root><mxCell id="0"/><mxCell id="1" parent="0"/></root></mxGraphModel></diagram></mxfile>`,
+        )
+        expect(contentFingerprint(XML_A)).not.toBe(contentFingerprint(twoPages))
+    })
+
+    it("falls back to the raw string for unparseable input", () => {
+        expect(contentFingerprint("not xml at all")).toBe("not xml at all")
+    })
+})