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

feat(mcp): add multi-page (mxfile) support to MCP server (#862)

Browse Source 
* feat(mcp): add multi-page (mxfile) support

The MCP server's write path could only address a single drawio page even
though the underlying .drawio file format and the embedded editor both
natively support multi-page documents. A user asking for "a second page
with a CNN diagram" would hit the validator with the error
"Expected closing tag </root> but found </mxCell>" because the validator
assumed input was a bare <mxGraphModel> and could not walk past the
<mxfile><diagram>...</diagram></mxfile> wrapper.

This patch closes the gap end to end:

* New helper module `pages.ts` centralises page CRUD (normalize, parse,
  list, find, add, rename, delete) so every layer agrees that the
  canonical in-memory shape is always <mxfile>. normalizeToMxfile and
  addPageToDoc both strip any leading <?xml ?> declaration before
  embedding a fragment inside <diagram> (the declaration is only valid
  at document start). addPageToDoc explicitly rejects full <mxfile>
  inputs so a caller cannot accidentally nest a document inside a page.
* `xml-validation.ts` now detects an <mxfile> root and scopes the
  duplicate-id check per <diagram>. The legacy regex check would
  otherwise reject every multi-page doc, because cells "0" and "1"
  repeat in each page's <root> by design. The DOM-parse path is gated
  by a cheap regex pre-check so legacy bare <mxGraphModel> callers
  don't pay any extra cost. The autoFix duplicate-id rename step is
  also guarded against mxfile inputs — renaming those sentinel cells
  would silently break drawio's parent references.
* `diagram-operations.ts` accepts an optional PageSelector. For
  <mxfile> input it resolves the page first and scopes all
  querySelectorAll calls to that page's <root>, so a delete on page 2's
  cell "2" no longer touches page 1's cell "2".
* `create_new_diagram` accepts either a bare <mxGraphModel> (legacy,
  auto-wrapped into a single-page mxfile) or a full <mxfile> with N
  diagrams. All existing single-page callers keep working unchanged.
* `edit_diagram`, `get_diagram`, and `export_diagram` gain optional
  `page_id` / `page_name` / `page_index` parameters. When omitted they
  target the first page — the "active by convention" default. Tool
  handlers with all-optional input schemas coalesce missing arguments
  via `input ?? {}` so a no-args MCP invocation can't crash on
  destructure before reaching the session-existence check.
* New tools: `list_pages`, `add_page`, `rename_page`, `delete_page`.
* Page-targeted PNG/SVG export uses a "load + export + restore" dance:
  the server projects the target page into a single-page <mxfile>,
  pushes it into the transient state so the browser reloads the iframe
  with just that page, waits for drawio to render (~3s), triggers the
  export, captures the data, and then restores the original multi-page
  document. The dance is wrapped in `try/finally` so the restore runs
  unconditionally — even if an exception is thrown mid-dance, the
  user's multi-tab view is recovered before the function returns.
  The earlier attempt to use drawio's `selectPage` postMessage was a
  no-op because drawio's JSON embed protocol does not expose that
  action — silently exporting whatever tab happened to be active. The
  load-export-restore approach trades a brief visible tab-flicker for
  correctness: the exported image is guaranteed to match the requested
  page.
* Tool description strings reflect the multi-page semantics so the LLM
  client learns the new contract.
* Package version bumped 0.2.0 → 0.3.0 (additive surface — four new
  tools, three extended input schemas, canonical XML shape change).
* CI: `.github/workflows/test.yml` gains an explicit install + vitest
  run for the mcp-server package so the new multi-page invariants are
  covered by automation, not just local runs.

Backward compatibility: every existing single-page caller continues to
work without modification. The session.xml shape is normalised on every
write, removing the wrapper-injection hack from the .drawio download
path.

Tests: 43 unit tests under `packages/mcp-server/tests/multi-page.test.ts`
pin the validator's mxfile path, the page-scoped operations, the XML
declaration-prefix handling for both normalizeToMxfile and addPageToDoc,
addPageToDoc's rejection of full <mxfile> inputs, the single-page
projection used by export_diagram (a direct regression test for the
selectPage bug — two distinct page selectors must produce visually
different projections), and the Transformer + CNN motivating scenario.
A `tests/smoke.mjs` smoke test drives the built `dist/index.js` over
JSON-RPC and asserts all 9 tools register with the right input schemas.
Root vitest suite (107 tests) still green.

* fix(mcp): rewrite page-targeted export browser-side; harden edit/get

The page-targeted PNG/SVG export never worked: export_diagram swapped the
live session to a single-page projection, slept 3s, then wrote the export
flag onto a state object that setState() had already replaced in the store
Map — so the browser never saw the request and every such export timed out.
The swap+restore also clobbered concurrent edits.

Move the projection entirely browser-side: requestExport() hands a single
-page <mxfile> to the bridge via state.exportXml; the bridge loads it,
lets draw.io render, exports, then reloads the user's real document. The
canonical session state is never mutated, so there is no restore race and
no fixed-delay guessing. The export poll now re-reads the live store entry
each tick instead of a captured reference. autosave is suppressed and the
version-bump reload is skipped while a projection is on screen; if no real
document was captured, restore forces a server reload rather than leaving
the iframe stuck on the projection.

Also:
- edit_diagram now returns isError on a page-level failure (selector matched
  no page / page has no <root>) instead of reporting success-with-warnings
  and persisting a no-op; the pre-edit history snapshot is taken only after
  that gate so a failed edit leaves no phantom undo entry.
- edit_diagram/get_diagram re-normalise browser-pushed xml to mxfile so a
  bare <mxGraphModel> can't silently strip a multi-page document.
- get_diagram now errors (instead of silently returning the full doc) when a
  selector is given but the session isn't a parseable mxfile.
- page_id / page_name / add_page.id get .min(1) so empty strings can't
  silently target the first page.
- Extract pages.ts:projectPage(), collapsing three copies of the
  parse→find→serialise projection logic in index.ts.
- Replace the never-in-CI tests/smoke.mjs with tests/server-wiring.test.ts,
  which boots the server from source via tsx and runs under the existing
  vitest CI step.

* chore(mcp): set version to 0.2.1 for release

---------

Co-authored-by: dayuan.jiang <jdy.toh@gmail.com>
This commit is contained in:
Siddhant Shekhar
2026-06-16 05:45:50 +05:30
committed by GitHub
parent 8e42dd9da8
commit 5c884766a8
11 changed files with 3298 additions and 121 deletions
Download Patch File Download Diff File Expand all files Collapse all files

111
packages/mcp-server/src/http-server.ts
View File

@@ -93,6 +93,7 @@ interface SessionState {
svg?: string // Cached SVG from last browser save
syncRequested?: number // Timestamp when sync requested, cleared when browser responds
exportFormat?: "png" | "svg" // Set by MCP tool to request browser export
exportXml?: string // Single-page projection to load before a page-targeted export
exportData?: string // Base64/SVG data returned by browser after export
}
@@ -117,12 +118,37 @@ export function setState(sessionId: string, xml: string, svg?: string): number {
svg: svg || existing?.svg, // Preserve cached SVG if not provided
syncRequested: undefined, // Clear sync request when browser pushes state
exportFormat: existing?.exportFormat, // Preserve pending export request
exportXml: existing?.exportXml, // Preserve pending projection
exportData: existing?.exportData, // Preserve export result
})
log.debug(`State updated: session=${sessionId}, version=${newVersion}`)
return newVersion
}
/**
* Ask the browser bridge to export the current diagram as png/svg.
*
* When `projectionXml` is given (a single-page <mxfile>), the bridge loads it
* first, waits for draw.io's own load event, exports, then reloads the
* session's real document — so a page-targeted export never mutates the
* canonical session state and needs no fixed-delay guessing on the server.
*
* Returns false when the session is unknown. Callers should then poll
* `getState(sessionId)?.exportData` for the result.
*/
export function requestExport(
sessionId: string,
format: "png" | "svg",
projectionXml?: string,
): boolean {
const state = stateStore.get(sessionId)
if (!state) return false
state.exportData = undefined
state.exportXml = projectionXml
state.exportFormat = format
return true
}
export function requestSync(sessionId: string): boolean {
const state = stateStore.get(sessionId)
if (state) {
@@ -286,6 +312,7 @@ function handleStateApi(
version: state?.version || 0,
syncRequested: !!state?.syncRequested,
exportFormat: state?.exportFormat || null,
exportXml: state?.exportXml || null,
}),
)
} else if (req.method === "POST") {
@@ -305,6 +332,7 @@ function handleStateApi(
if (state) {
state.exportData = data.exportData
state.exportFormat = undefined
state.exportXml = undefined
log.debug(
`Export data received for session=${sessionId}`,
)
@@ -675,6 +703,8 @@ function getHtmlPage(sessionId: string): string {
let pendingSvgExport = null;
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
window.addEventListener('message', (e) => {
if (e.origin !== '${DRAWIO_ORIGIN}') return;
@@ -684,6 +714,10 @@ function getHtmlPage(sessionId: string): string {
isReady = true;
if (pendingXml) { loadDiagram(pendingXml); pendingXml = null; }
} else if ((msg.event === 'save' || msg.event === 'autosave') && msg.xml && msg.xml !== lastXml) {
// Ignore autosave while a single-page projection is on screen
// for a page-targeted export — otherwise we'd push the
// transient projection back as the canonical session state.
if (projectionExportActive) return;
// Request SVG export, then push state with SVG
pendingSvgExport = msg.xml;
iframe.contentWindow.postMessage(JSON.stringify({ action: 'export', format: 'svg' }), '*');
@@ -704,6 +738,9 @@ function getHtmlPage(sessionId: string): string {
headers: { 'Content-Type': 'application/json' },
body: JSON.stringify({ sessionId, exportData: d })
}).catch(() => {});
// Page-targeted export: restore the user's real
// multi-page document now that we have the image.
restoreFromProjection();
return;
}
}
@@ -761,6 +798,22 @@ 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.
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
}
}
async function pushState(xml, svg = '') {
if (!sessionId) return;
try {
@@ -786,20 +839,54 @@ function getHtmlPage(sessionId: string): string {
pendingSyncExport = true;
iframe.contentWindow.postMessage(JSON.stringify({ action: 'export', format: 'xml' }), '*');
}
// Load new diagram from server (before export, so we export latest)
if (s.version > currentVersion && s.xml) {
// Load new diagram from server (before export, so we export latest).
// While a page-targeted projection is on screen, skip the reload
// 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) {
currentVersion = s.version;
loadDiagram(s.xml, true);
}
// Handle export request from MCP server (png/svg) - after version update
// Handle export request from MCP server (png/svg).
//
// Plain export: capture whatever tab is currently displayed.
//
// Page-targeted export: the server sends a single-page <mxfile>
// projection in s.exportXml. We load it into the iframe, let
// draw.io render it, export, then reload the user's real
// document — all browser-side. The canonical session state is
// never mutated, so there is no server-side restore race and no
// dependence on poll timing. autosave is suppressed while the
// projection is showing (see projectionExportActive guard).
if (s.exportFormat && !pendingMcpExport && isReady) {
pendingMcpExport = s.exportFormat;
const exportOpts = s.exportFormat === 'png'
? { action: 'export', format: 'png', scale: 2 }
: { action: 'export', format: 'svg' };
iframe.contentWindow.postMessage(JSON.stringify(exportOpts), '*');
// Timeout: reset if draw.io never responds
setTimeout(() => { if (pendingMcpExport) { pendingMcpExport = null; } }, 8000);
const fireExport = () => {
const exportOpts = pendingMcpExport === 'png'
? { action: 'export', format: 'png', scale: 2 }
: { action: 'export', format: 'svg' };
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 }), '*');
// Let draw.io render the loaded page before exporting
// (same proven settle delay as the AI-preview path).
setTimeout(fireExport, 600);
} else {
fireExport();
}
// Timeout: reset if draw.io never responds, and restore the
// real document if a projection was left showing.
setTimeout(() => {
if (pendingMcpExport) {
pendingMcpExport = null;
restoreFromProjection();
}
}, 10000);
}
} catch {}
}
@@ -839,7 +926,11 @@ function getHtmlPage(sessionId: string): string {
saveConfirmBtn.textContent = 'Exporting...';
if (format === 'drawio') {
// Use lastXml directly instead of requesting export (avoids race with SVG exports)
// Use lastXml directly instead of requesting export (avoids race with SVG exports).
// session.xml is canonically <mxfile> after the multi-page refactor,
// so no wrapper injection is needed. The legacy fallback below
// remains only for documents that somehow slipped past
// normalisation (e.g. an older session loaded from external state).
let xmlData = lastXml || '';
if (xmlData && !xmlData.includes('<mxfile')) {
xmlData = '<mxfile host="mcp"><diagram name="Page-1">' + xmlData + '</diagram></mxfile>';