i/next-ai-draw-io
1
0
Fork 0
mirror of https://github.com/DayuanJiang/next-ai-draw-io.git synced 2026-10-08 02:37:46 +08:00
Code Issues Packages Projects Releases Wiki Activity

fix(mcp-server): keep both pages when get_diagram meets a page export, and more review fixes

Browse Source 
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 <parsererror>
  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 <diagram> 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.
This commit is contained in:
dayuan.jiang
2026-10-04 23:04:21 +09:00
parent 2bb69b58af
commit 504d2fa812
15 changed files with 189 additions and 26 deletions
Download Patch File Download Diff File Expand all files Collapse all files
packages/mcp-server/src/diagram-operations.ts
+5 -1
View File
@@ -185,7 +185,11 @@ export function applyDiagramOperations(
const cellMap = new Map<string, Element>()
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<string>()
packages/mcp-server/src/edit-diagram.ts
+7 -1
View File
@@ -24,7 +24,13 @@ function countTopLevelElements(fragment: string): number {
`<wrapper>${fragment}</wrapper>`,
"text/xml",
)
return doc.documentElement?.children.length ?? 0
// On a syntax error the browser adds a <parsererror> 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 <mxfile>, or the whole document. */
packages/mcp-server/src/http-server.ts
+14 -1
View File
@@ -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 }))
packages/mcp-server/src/index.ts
+2 -9
View File
@@ -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
})
packages/mcp-server/src/new-diagram.ts
+13
View File
@@ -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) &&
/<mxCell\b(?=[^>]*\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
packages/mcp-server/src/pages.ts
+1 -1
View File
@@ -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 <mxfile> root? */
export function isMxFile(xml: string): boolean {
packages/mcp-server/src/persistence.ts
+8 -1
View File
@@ -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 {
packages/mcp-server/src/preview/preview.js
+14 -9
View File
@@ -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
packages/mcp-server/src/xml-validation.ts
+3 -1
View File
@@ -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 <diagram> 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)
}
packages/mcp-server/tests/diagram-operations.test.ts
+25
View File
@@ -124,3 +124,28 @@ describe("pages without a <root>", () => {
expect(errors[0]?.message).toContain("could not be decompressed")
})
})
describe("a wrapped mxCell with its wrapper's id", () => {
const doc = `<mxfile><diagram id="p" name="P"><mxGraphModel><root><mxCell id="0"/><mxCell id="1" parent="0"/><UserObject id="u" label="A" link="https://example.com"><mxCell id="u" vertex="1" parent="1"><mxGeometry as="geometry"/></mxCell></UserObject></root></mxGraphModel></diagram></mxfile>`
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: `<UserObject id="u" label="B" link="https://example.com"><mxCell vertex="1" parent="1"><mxGeometry as="geometry"/></mxCell></UserObject>`,
},
])
expect(errors).toEqual([])
expect(result.match(/<UserObject/g)).toHaveLength(1)
expect(result).toContain('label="B"')
})
})
packages/mcp-server/tests/http-server.test.ts
+19
View File
@@ -258,9 +258,28 @@ describe("export requests", () => {
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, "<mxfile>x</mxfile>")
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", () => {
packages/mcp-server/tests/persistence.test.ts
+6 -1
View File
@@ -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"))
})
})
packages/mcp-server/tests/wrap-cells.test.ts
+32 -1
View File
@@ -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 = `<mxCell id="2" value="A" vertex="1" parent="1"><mxGeometry x="0" y="0" width="80" height="40" as="geometry"/></mxCell>`
@@ -66,3 +67,33 @@ describe("wrapCellsInModel", () => {
)
})
})
describe("hasCells", () => {
it("counts cells written with single quotes", () => {
expect(hasCells(`<mxCell id='2' vertex='1' parent='1'/>`)).toBe(true)
expect(hasCells(`<mxCell id='0'/><mxCell id='1' parent='0'/>`)).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 =
`<mxCell id="1" value="Start" vertex="1" parent="1"><mxGeometry x="0" y="0" width="80" height="40" as="geometry"/></mxCell>` +
`<mxCell id="3" edge="1" parent="1" source="1" target="2"><mxGeometry relative="1" as="geometry"/></mxCell>`
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(
`<mxCell id="0"/><mxCell id="1" parent="0"/>${A}`,
)
expect(out.ok).toBe(true)
})
})
packages/mcp-server/tests/xml-validation.test.ts
+10
View File
@@ -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(
`<mxfile><diagram id="p" name="P">\\n${model(`<mxCell id="2" vertex="1" parent="1"/>`)}</diagram></mxfile>`,
)
expect(r.valid).toBe(true)
expect(r.fixed).not.toContain("\\n")
})
})
tests/e2e/diagram-content.spec.ts
+30
View File
@@ -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 <parsererror> 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(
"</mxCell>",
"</mxcell>",
),
},
],
}),
])
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", {