From fc8ffc5655cefcb49e77edc6e5b649eeef64d0e9 Mon Sep 17 00:00:00 2001 From: "dayuan.jiang" Date: Mon, 12 Oct 2026 08:25:47 +0900 Subject: [PATCH] refactor(mcp-server): remove dead code and fold repeated tool plumbing - index.ts: registerTool() does the session check, try/catch and the write queue for all 15 tools; fail()/ok(), syncFromBrowser(), commitDocument() and describePages() replace the copied blocks; drop pickPageSelector, currentSession.version, the input ?? {} defaults; add_page reuses prepareNewModel() split out of prepareNewDiagram - xml-validation.ts: drop the hand-written syntax checks that saxes already reports, MAX_XML_SIZE and the stale header - http-server.ts: one pendingExport object, no unread response fields, no CORS/OPTIONS branch for same-origin requests, fewer exports - diagram-operations parseNewCell(), addHistory returns void, normalizeToMxfile without host, foldCells without layer, expandStyles removed, unused preview.css rule removed, no .d.ts output - xml-examples.ts styleRules() shares the Styles paragraphs with the web prompt - tests: setupFiles installs the DOM polyfill once, server-wiring uses the SDK client, mcp-core.test.ts drops the cases MCP tests already cover --- packages/mcp-server/src/compact-cells.ts | 4 +- packages/mcp-server/src/diagram-operations.ts | 133 +- packages/mcp-server/src/dom.ts | 2 +- packages/mcp-server/src/drawing-guide.ts | 16 +- packages/mcp-server/src/drawio-themes.ts | 4 +- packages/mcp-server/src/history.ts | 5 +- packages/mcp-server/src/http-server.ts | 73 +- packages/mcp-server/src/index.ts | 2228 ++++++----------- packages/mcp-server/src/new-diagram.ts | 40 +- packages/mcp-server/src/pages.ts | 5 +- packages/mcp-server/src/preview/preview.css | 1 - packages/mcp-server/src/preview/preview.js | 14 +- packages/mcp-server/src/shape-library.ts | 2 +- packages/mcp-server/src/style-classes.ts | 6 - packages/mcp-server/src/xml-examples.ts | 18 + packages/mcp-server/src/xml-validation.ts | 145 +- .../tests/diagram-operations.test.ts | 8 +- .../mcp-server/tests/edit-diagram.test.ts | 8 +- packages/mcp-server/tests/edit-gate.test.ts | 8 +- packages/mcp-server/tests/history.test.ts | 8 +- packages/mcp-server/tests/http-server.test.ts | 17 +- packages/mcp-server/tests/multi-page.test.ts | 10 +- packages/mcp-server/tests/new-diagram.test.ts | 8 +- packages/mcp-server/tests/persistence.test.ts | 7 +- .../mcp-server/tests/server-wiring.test.ts | 242 +- packages/mcp-server/tests/setup.ts | 5 + packages/mcp-server/tests/wrap-cells.test.ts | 24 +- .../mcp-server/tests/xml-reference.test.ts | 8 +- .../mcp-server/tests/xml-validation.test.ts | 10 +- packages/mcp-server/tsconfig.json | 2 - packages/mcp-server/vitest.config.ts | 1 + tests/unit/compact-cells.test.ts | 5 +- tests/unit/mcp-core.test.ts | 177 +- tests/unit/mcp-preview-recovery.test.ts | 2 +- tests/unit/style-classes.test.ts | 7 +- 35 files changed, 1068 insertions(+), 2185 deletions(-) create mode 100644 packages/mcp-server/tests/setup.ts diff --git a/packages/mcp-server/src/compact-cells.ts b/packages/mcp-server/src/compact-cells.ts index aaaea56d..00b84460 100644 --- a/packages/mcp-server/src/compact-cells.ts +++ b/packages/mcp-server/src/compact-cells.ts @@ -158,7 +158,7 @@ function plainGeometry(body: string): Map | null { * missing one. Everything else, including wrapped cells' extra data and * edge labels, is left as written. */ -export function foldCells(xml: string, layer = "1"): string { +export function foldCells(xml: string): string { return xml.replace(CELL_BLOCK, (block, attrText: string, body?: string) => { if (body === undefined) return block const attrs = readAttributes(attrText) @@ -174,7 +174,7 @@ export function foldCells(xml: string, layer = "1"): string { (a) => a.name !== "vertex" && a.name !== "edge" && - !(a.name === "parent" && a.value === layer), + !(a.name === "parent" && a.value === "1"), ) if (byName.get("vertex") === "1" && byName.get("edge") !== "1") { if (geometryKeys.join() !== "height,width,x,y") return block diff --git a/packages/mcp-server/src/diagram-operations.ts b/packages/mcp-server/src/diagram-operations.ts index 64dd810a..e68173c7 100644 --- a/packages/mcp-server/src/diagram-operations.ts +++ b/packages/mcp-server/src/diagram-operations.ts @@ -79,6 +79,32 @@ function cellAttr(cell: Element, name: string): string | null { return inner?.getAttribute(name) ?? null } +/** + * The cell an add or update operation sends in new_xml, imported into doc, + * or the error message for the model. + */ +function parseNewCell( + op: DiagramOperation, + parser: DOMParser, + doc: Document, +): Element | string { + if (!op.new_xml) return `new_xml is required for ${op.operation} operation` + // A cut-off cell (its XML still streaming in a preview) would parse in + // Chrome to a cell without its geometry + const wrapped = `${op.new_xml}` + const syntaxError = getXmlSyntaxError(wrapped) + if (syntaxError) return `new_xml is not well-formed XML: ${syntaxError}` + const newCell = parser + .parseFromString(wrapped, "text/xml") + .querySelector(CELL_SELECTOR) + if (!newCell) return "new_xml must contain an mxCell element" + const newCellId = newCell.getAttribute("id") + if (newCellId !== op.cell_id) { + return `ID mismatch: cell_id is "${op.cell_id}" but new_xml has id="${newCellId}"` + } + return doc.importNode(newCell, true) +} + /** * Apply diagram operations (update/add/delete) using ID-based lookup. * @@ -207,63 +233,19 @@ export function applyDiagramOperations( continue } - if (!op.new_xml) { + const newCell = parseNewCell(op, parser, doc) + if (typeof newCell === "string") { errors.push({ type: "update", cellId: op.cell_id, - message: "new_xml is required for update operation", + message: newCell, }) continue } - // A cut-off cell (its XML still streaming in a preview) would - // parse in Chrome to a cell without its geometry - const syntaxError = getXmlSyntaxError( - `${op.new_xml}`, - ) - if (syntaxError) { - errors.push({ - type: "update", - cellId: op.cell_id, - message: `new_xml is not well-formed XML: ${syntaxError}`, - }) - continue - } - - // Parse the new XML - const newDoc = parser.parseFromString( - `${op.new_xml}`, - "text/xml", - ) - const newCell = newDoc.querySelector(CELL_SELECTOR) - if (!newCell) { - errors.push({ - type: "update", - cellId: op.cell_id, - message: "new_xml must contain an mxCell element", - }) - continue - } - - // Validate ID matches - const newCellId = newCell.getAttribute("id") - if (newCellId !== op.cell_id) { - errors.push({ - type: "update", - cellId: op.cell_id, - message: `ID mismatch: cell_id is "${op.cell_id}" but new_xml has id="${newCellId}"`, - }) - continue - } - - // Import and replace the node - const importedNode = doc.importNode(newCell, true) - existingCell.parentNode?.replaceChild(importedNode, existingCell) - - // Update the map with the new element - cellMap.set(op.cell_id, importedNode) + existingCell.parentNode?.replaceChild(newCell, existingCell) + cellMap.set(op.cell_id, newCell) } else if (op.operation === "add") { - // Check if ID already exists if (cellMap.has(op.cell_id)) { errors.push({ type: "add", @@ -273,61 +255,18 @@ export function applyDiagramOperations( continue } - if (!op.new_xml) { + const newCell = parseNewCell(op, parser, doc) + if (typeof newCell === "string") { errors.push({ type: "add", cellId: op.cell_id, - message: "new_xml is required for add operation", + message: newCell, }) continue } - // A cut-off cell (its XML still streaming in a preview) would - // parse in Chrome to a cell without its geometry - const syntaxError = getXmlSyntaxError( - `${op.new_xml}`, - ) - if (syntaxError) { - errors.push({ - type: "add", - cellId: op.cell_id, - message: `new_xml is not well-formed XML: ${syntaxError}`, - }) - continue - } - - // Parse the new XML - const newDoc = parser.parseFromString( - `${op.new_xml}`, - "text/xml", - ) - const newCell = newDoc.querySelector(CELL_SELECTOR) - if (!newCell) { - errors.push({ - type: "add", - cellId: op.cell_id, - message: "new_xml must contain an mxCell element", - }) - continue - } - - // Validate ID matches - const newCellId = newCell.getAttribute("id") - if (newCellId !== op.cell_id) { - errors.push({ - type: "add", - cellId: op.cell_id, - message: `ID mismatch: cell_id is "${op.cell_id}" but new_xml has id="${newCellId}"`, - }) - continue - } - - // Import and append the node - const importedNode = doc.importNode(newCell, true) - root.appendChild(importedNode) - - // Add to map - cellMap.set(op.cell_id, importedNode) + root.appendChild(newCell) + cellMap.set(op.cell_id, newCell) } else if (op.operation === "delete") { // Protect root cells from deletion if (op.cell_id === "0" || op.cell_id === "1") { diff --git a/packages/mcp-server/src/dom.ts b/packages/mcp-server/src/dom.ts index 13972b3a..34e618ca 100644 --- a/packages/mcp-server/src/dom.ts +++ b/packages/mcp-server/src/dom.ts @@ -25,7 +25,7 @@ const escapeChars = (text: string, chars: RegExp) => * Serialize a linkedom node as XML. Attribute values escape tabs and line * breaks too, so multi-line labels (value="a b") survive a round trip. */ -export function serializeXml(node: Node): string { +function serializeXml(node: Node): string { switch (node.nodeType) { case 9: { // Document diff --git a/packages/mcp-server/src/drawing-guide.ts b/packages/mcp-server/src/drawing-guide.ts index 25d4c1f3..e63d8e04 100644 --- a/packages/mcp-server/src/drawing-guide.ts +++ b/packages/mcp-server/src/drawing-guide.ts @@ -2,19 +2,23 @@ * Drawing guide for the model, returned by start_session, get_drawing_guide * and the diagram-workflow prompt. * - * Adapted from the web app's system prompt (lib/system-prompts.ts) and its - * tool descriptions (app/api/chat/route.ts). When drawing rules change there, - * update this file too. + * The XML examples and the style rules are shared with the web app's system + * prompt (xml-examples.ts). The other rules are adapted from that prompt + * (lib/system-prompts.ts) and the web app's tool descriptions + * (app/api/chat/route.ts); when they change there, update this file too. */ import { indent, STYLE_CLASS_EXAMPLE, SWIMLANE_EXAMPLE, + styleRules, TWO_EDGES_EXAMPLE, WAYPOINT_EXAMPLE, } from "./xml-examples.ts" +const styles = styleRules("create_new_diagram or add_page", "the server") + export const DRAWING_GUIDE = `# Draw.io drawing guide ## Workflow @@ -84,12 +88,12 @@ Waypoints go inside in the edge geometry. Example: Hotfix (r ${indent(WAYPOINT_EXAMPLE)} ## Styles -Define each style used by several cells ONCE, as a named style before the cells, and use the name in the cells like a CSS class. A cell's style can combine a shape token, a name and overrides; later pairs win. Name only styles that two or more cells share; a style used by one cell stays inline. Names must not be draw.io's own style names: shapes such as text, ellipse, rhombus, swimlane, label, image, and colors such as blue, green, red, gray, yellow, orange, purple, pink. A definition applies to the call it is in: each create_new_diagram or add_page call defines the names it uses. The server expands the names, so the saved file is standard draw.io XML. +${styles.namedStyles} ${indent(STYLE_CLASS_EXAMPLE)} -- NEVER write html=1 or whiteSpace=wrap: the server adds html=1 to every cell and whiteSpace=wrap to shapes. Labels are HTML: use <br> for a line break and <b> for bold, never \\n; a literal < or > in a label is written &lt; or &gt;. -- Do NOT repeat what draw.io already uses. For a plain shape: rounded=0, align=center, verticalAlign=middle, fontSize=12, strokeWidth=1, fillColor=#ffffff, strokeColor=#000000, fontColor=#000000. For an edge: endArrow=classic, strokeColor=#000000. Writing one of them is right only when it overrides what a name or the shape sets: an edge is rounded by default, so rounded=0 on an edge is a real setting, and a text cell is left/top aligned by default, so there align=center or verticalAlign=middle are real settings. +- ${styles.noHtml} +- ${styles.noDefaults} - Shapes: rounded=1, fillColor=#hex, strokeColor=#hex - Edges: endArrow=block, open or none; startArrow=classic; curved=1; dashed=1; edgeStyle=orthogonalEdgeStyle - Text: fontSize=14, fontStyle=1 (bold), align=center or right diff --git a/packages/mcp-server/src/drawio-themes.ts b/packages/mcp-server/src/drawio-themes.ts index a43b0615..9c032ded 100644 --- a/packages/mcp-server/src/drawio-themes.ts +++ b/packages/mcp-server/src/drawio-themes.ts @@ -1,4 +1,4 @@ -export const DRAWIO_THEMES = [ +const DRAWIO_THEMES = [ "kennedy", "atlas", "dark", @@ -7,7 +7,7 @@ export const DRAWIO_THEMES = [ "simple", ] as const -export type DrawioTheme = (typeof DRAWIO_THEMES)[number] +type DrawioTheme = (typeof DRAWIO_THEMES)[number] export function isDrawioTheme(value: unknown): value is DrawioTheme { return ( diff --git a/packages/mcp-server/src/history.ts b/packages/mcp-server/src/history.ts index 197a13e7..56f9169b 100644 --- a/packages/mcp-server/src/history.ts +++ b/packages/mcp-server/src/history.ts @@ -18,7 +18,7 @@ export interface HistoryEntry { let nextEntryId = 0 const historyStore = new Map() -export function addHistory(sessionId: string, xml: string, svg = ""): number { +export function addHistory(sessionId: string, xml: string, svg = ""): void { let history = historyStore.get(sessionId) if (!history) { history = [] @@ -31,7 +31,7 @@ export function addHistory(sessionId: string, xml: string, svg = ""): number { const last = history[history.length - 1] if (last && last.xml === xml) { if (svg && !last.svg) last.svg = svg - return history.length - 1 + return } history.push({ id: nextEntryId++, xml, svg }) @@ -42,7 +42,6 @@ export function addHistory(sessionId: string, xml: string, svg = ""): number { } log.debug(`History: session=${sessionId}, entries=${history.length}`) - return history.length - 1 } export function getHistory(sessionId: string): HistoryEntry[] { diff --git a/packages/mcp-server/src/http-server.ts b/packages/mcp-server/src/http-server.ts index ced602ba..dad81a91 100644 --- a/packages/mcp-server/src/http-server.ts +++ b/packages/mcp-server/src/http-server.ts @@ -146,13 +146,18 @@ interface SessionState { lastPolled?: number // Last browser poll; an open tab keeps the session alive svg?: string // Cached SVG from last browser save syncRequested?: number // Timestamp when sync requested, cleared when browser responds - 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 + pendingExport?: PendingExport // Set by MCP tool to request browser export exportData?: string // Base64/SVG data returned by browser after export } +/** An export the browser is asked to make, sent to the preview page as is */ +interface PendingExport { + id: number // Number of the export, echoed with its result + format: ExportFormat + xml?: string // Single-page projection to load before a page-targeted export + options?: ExportOptions // Extra draw.io export parameters (PNG only) +} + /** draw.io export formats; xmlsvg is an SVG with the diagram embedded */ export type ExportFormat = "png" | "svg" | "xmlsvg" @@ -165,7 +170,7 @@ export interface ExportOptions { pageId?: string } -export const stateStore = new Map() +const stateStore = new Map() let server: http.Server | null = null let serverPort = 6002 @@ -220,10 +225,7 @@ export function setState( svg || (existing && existing.xml === xml ? existing.svg : undefined), syncRequested: undefined, // Clear sync request when browser pushes state - exportFormat: existing?.exportFormat, // Preserve pending export request - exportXml: existing?.exportXml, // Preserve pending projection - exportOptions: existing?.exportOptions, - exportId: existing?.exportId, + pendingExport: existing?.pendingExport, // Preserve pending export request exportData: existing?.exportData, // Preserve export result }) log.debug(`State updated: session=${sessionId}, version=${newVersion}`) @@ -266,12 +268,14 @@ export function requestExport( const state = stateStore.get(sessionId) if (!state) return false state.exportData = undefined - state.exportXml = projectionXml - state.exportOptions = options - state.exportFormat = format - // The browser sends this back with the result, so a late result of an + // The browser sends the id back with the result, so a late result of an // export that timed out is not taken for this one - state.exportId = ++lastExportId + state.pendingExport = { + id: ++lastExportId, + format, + xml: projectionXml, + options, + } return true } @@ -340,7 +344,7 @@ export function startHttpServer(port = 6002): Promise { }) } -export function stopHttpServer(): void { +function stopHttpServer(): void { if (server) { server.close() server = null @@ -421,19 +425,6 @@ function routeRequest( return } - const requestOrigin = req.headers.origin - if (requestOrigin === `http://localhost:${serverPort}`) { - res.setHeader("Access-Control-Allow-Origin", requestOrigin) - res.setHeader("Access-Control-Allow-Methods", "GET, POST, OPTIONS") - res.setHeader("Access-Control-Allow-Headers", "Content-Type") - } - - if (req.method === "OPTIONS") { - res.writeHead(204) - res.end() - return - } - if (url.pathname === "/" || url.pathname === "/index.html") { const sessionId = url.searchParams.get("mcp") || "" if (sessionId && !isValidSessionId(sessionId)) { @@ -497,10 +488,7 @@ function handleStateApi( stateId: state?.stateId ?? null, blank: !!state?.blank, syncRequested: !!state?.syncRequested, - exportFormat: state?.exportFormat || null, - exportXml: state?.exportXml || null, - exportOptions: state?.exportOptions || null, - exportId: state?.exportId ?? null, + export: state?.pendingExport ?? null, }), ) } else if (req.method === "POST") { @@ -519,12 +507,9 @@ function handleStateApi( // Browser is returning export data (png/svg) if (data.exportData !== undefined) { const state = stateStore.get(sessionId) - if (state && data.exportId === state.exportId) { + if (state && data.exportId === state.pendingExport?.id) { state.exportData = data.exportData - state.exportFormat = undefined - state.exportXml = undefined - state.exportOptions = undefined - state.exportId = undefined + state.pendingExport = undefined log.debug( `Export data received for session=${sessionId}`, ) @@ -554,11 +539,7 @@ function handleStateApi( "Content-Type": "application/json", }) res.end( - JSON.stringify({ - error: "Session was recreated", - stateChanged: true, - version: current.version, - }), + JSON.stringify({ error: "Session was recreated" }), ) return } @@ -579,7 +560,6 @@ function handleStateApi( res.end( JSON.stringify({ error: "Diagram changed on the server", - version: current.version, savedToHistory: saved, }), ) @@ -614,7 +594,6 @@ function handleStateApi( res.end( JSON.stringify({ error: "Diagram changed on the server", - version: current?.version, savedToHistory, }), ) @@ -669,7 +648,6 @@ function handleHistoryApi( id: entry.id, svg: entry.svg, })), - count: history.length, }), ) } @@ -737,15 +715,14 @@ function handleRestoreApi( return } - const newVersion = restoreHistoryEntry(sessionId, id) - if (newVersion === null) { + if (restoreHistoryEntry(sessionId, id) === null) { res.writeHead(404, { "Content-Type": "application/json" }) res.end(JSON.stringify({ error: "Entry not found" })) return } res.writeHead(200, { "Content-Type": "application/json" }) - res.end(JSON.stringify({ success: true, newVersion })) + res.end(JSON.stringify({ success: true })) } catch { res.writeHead(400, { "Content-Type": "application/json" }) res.end(JSON.stringify({ error: "Invalid JSON" })) diff --git a/packages/mcp-server/src/index.ts b/packages/mcp-server/src/index.ts index 905687ca..eab70e8f 100644 --- a/packages/mcp-server/src/index.ts +++ b/packages/mcp-server/src/index.ts @@ -18,13 +18,22 @@ * surface. */ +import { readFile, stat, writeFile } from "node:fs/promises" import { createRequire } from "node:module" -import { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js" +import { resolve } from "node:path" +import { + McpServer, + type ToolCallback, +} from "@modelcontextprotocol/sdk/server/mcp.js" import { StdioServerTransport } from "@modelcontextprotocol/sdk/server/stdio.js" -import type { CallToolResult } from "@modelcontextprotocol/sdk/types.js" +import type { ZodRawShapeCompat } from "@modelcontextprotocol/sdk/server/zod-compat.js" +import type { + CallToolResult, + ToolAnnotations, +} from "@modelcontextprotocol/sdk/types.js" import open from "open" import { z } from "zod" -import { expandCompactCells, foldCells } from "./compact-cells.ts" +import { foldCells } from "./compact-cells.ts" import { customInstructionsPath, guideWithCustomInstructions, @@ -66,12 +75,7 @@ import { } from "./http-server.ts" import { parseDrawioFileContent } from "./load-diagram.ts" import { log } from "./logger.ts" -import { - prepareNewDiagram, - reservedIdError, - takeStyleDefinitions, - truncatedCellError, -} from "./new-diagram.ts" +import { prepareNewDiagram, prepareNewModel } from "./new-diagram.ts" import { addPageToDoc, BLANK_MXFILE, @@ -86,13 +90,10 @@ import { projectPage, renamePageInDoc, serializeMxfile, - wrapCellsInModel, } from "./pages.ts" import { Autosaver, defaultDataDir, expandHome } from "./persistence.ts" import { getShapeLibrary, SHAPE_LIBRARY_LIST } from "./shape-library.ts" -import { addDefaultStyles, applyStyleClasses } from "./style-classes.ts" import { XML_REFERENCE } from "./xml-reference.ts" -import { validateAndFixXml } from "./xml-validation.ts" // DOMParser/XMLSerializer globals for the XML helpers (Node has neither) installDomPolyfill() @@ -139,17 +140,34 @@ function sessionState(sessionId: string) { } // Session state (single session for simplicity) -let currentSession: { +type Session = { id: string xml: string - version: number // 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 +} +let currentSession: Session | null = null + +const NO_SESSION = "No active session. Please call start_session first." + +/** The session of a tool registered with needsSession, checked before it ran */ +function activeSession(): Session { + if (!currentSession) throw new Error(NO_SESSION) + return currentSession +} + +/** A tool result of one text */ +const ok = (text: string) => ({ content: [{ type: "text" as const, text }] }) + +/** A tool error; every tool reports its errors with this prefix */ +const fail = (text: string) => ({ + content: [{ type: "text" as const, text: `Error: ${text}` }], + isError: true as const, +}) // Create MCP server. The version reported in the MCP handshake is read from // package.json so it can never drift from the published npm version again @@ -195,12 +213,42 @@ const server = new McpServer( // would drop the first one's change. start_session in the queue keeps a // session switch from landing in the middle of a write. const exclusive = createExclusive() -const registerWriteTool = ((name: string, config: any, handler: any) => + +/** + * Register a tool. The handler returns its result or throws: an exception + * comes back as a tool error (the SDK would report it too, without the + * "Error: " prefix and the log line). With needsSession the handler runs + * only when a session is active, so it can take activeSession() for + * granted; write puts it in the queue above. + */ +function registerTool( + name: string, + config: { + title?: string + description?: string + inputSchema?: Args + annotations?: ToolAnnotations + }, + handler: ToolCallback, + { write = false, needsSession = false } = {}, +) { + const run = async (...args: unknown[]) => { + if (needsSession && !currentSession) return fail(NO_SESSION) + try { + return await (handler as (...a: unknown[]) => unknown)(...args) + } catch (error) { + const message = + error instanceof Error ? error.message : String(error) + log.error(`${name} failed:`, message) + return fail(message) + } + } server.registerTool( name, config, - exclusive(handler), - )) as typeof server.registerTool + (write ? exclusive(run) : run) as ToolCallback, + ) +} // Shared Zod schema fragment for page-targeting parameters. // Every multi-page-aware tool reuses these three optional fields so the LLM @@ -238,24 +286,6 @@ const screenshotSchema = z "Also return a PNG of the result so you can check it for overlaps and edges crossing shapes. Needs the preview tab open and in front; adds 2 to 10 s. Default: the DRAWIO_AUTO_SCREENSHOT environment variable.", ) -/** - * Pull a clean PageSelector out of a tool's parsed input. - * Returns an empty object when none of the page_* fields are set, so callers - * can simply pass it through to the lower layers (they treat empty as "first - * page" by convention). - */ -function pickPageSelector(input: { - page_id?: string - page_name?: string - page_index?: number -}): PageSelector { - const selector: PageSelector = {} - if (input.page_id) selector.page_id = input.page_id - if (input.page_name) selector.page_name = input.page_name - if (input.page_index !== undefined) selector.page_index = input.page_index - return selector -} - /** Format a selector for human-readable error messages. */ function describeSelector(s: PageSelector): string { if (s.page_id) return `id="${s.page_id}"` @@ -264,6 +294,50 @@ function describeSelector(s: PageSelector): string { return "first page" } +/** One line of pages for the tool results: index, id, name and cell count */ +function describePages(xml: string, empty = "no pages parsed"): string { + const doc = parseMxfile(xml) + const pages = doc ? listPagesFromDoc(doc) : [] + return pages.length > 0 + ? `Pages (${pages.length}): ${pages.map((p) => `[${p.index}] id=${p.id} name="${p.name}" cells=${p.cellCount}`).join(" | ")}` + : empty +} + +/** + * Take the browser's state of the session (the user's manual edits) into + * the session as an : the embed/sync path can hand back a bare + * , and adopting it verbatim would strip a multi-page document + * down to one page on the next write. Returns the state; the edit gate + * compares against its raw xml. + */ +function syncFromBrowser(session: Session) { + const browserState = sessionState(session.id) + if (browserState?.xml) { + session.xml = normalizeToMxfile(browserState.xml) ?? browserState.xml + } + return browserState +} + +/** + * Make xml the session's document: the previous one goes to History (with + * the browser's image of it), the new one is pushed to the preview and + * recorded in History (its image comes later, from the browser). seen: the + * model wrote this exact XML knowing the state it built on, so edit_diagram + * may follow without a get_diagram round-trip. + */ +function commitDocument( + session: Session, + xml: string, + seen: boolean, + browserSvg?: string, +) { + if (session.xml) keepInHistory(session.id, session.xml, browserSvg) + session.xml = xml + setState(session.id, xml) + session.lastSeenXml = seen ? xml : "" + addHistory(session.id, xml, "") +} + // The drawing guide plus the user's own rules from instructions.md, read // on every call so edits apply without restarting the host const guideText = () => guideWithCustomInstructions(DRAWING_GUIDE) @@ -285,7 +359,7 @@ server.registerPrompt( ) // Tool: get_drawing_guide -server.registerTool( +registerTool( "get_drawing_guide", { title: "Get drawing guide", @@ -303,18 +377,11 @@ server.registerTool( }, annotations: { readOnlyHint: true, openWorldHint: false }, }, - async ({ topic }) => ({ - content: [ - { - type: "text", - text: topic ? XML_REFERENCE[topic] : guideText(), - }, - ], - }), + async ({ topic }) => ok(topic ? XML_REFERENCE[topic] : guideText()), ) // Tool: get_shape_library -server.registerTool( +registerTool( "get_shape_library", { title: "Get shape library", @@ -331,17 +398,12 @@ server.registerTool( }, async ({ library }) => { const found = await getShapeLibrary(library) - return found.ok - ? { content: [{ type: "text", text: found.text }] } - : { - content: [{ type: "text", text: `Error: ${found.error}` }], - isError: true, - } + return found.ok ? ok(found.text) : fail(found.error) }, ) // Tool: start_session -registerWriteTool( +registerTool( "start_session", { title: "Start session", @@ -361,73 +423,46 @@ registerWriteTool( }, annotations: { destructiveHint: false, openWorldHint: false }, }, - async (input) => { - // The only field is optional, so a client may send no arguments - const { session_id } = input ?? {} - try { - // Start embedded HTTP server - const port = await startHttpServer(config.port) + async ({ session_id }) => { + // Start embedded HTTP server + const port = await startHttpServer(config.port) - // Create session, or continue a saved one under its old id - const sessionId = - session_id ?? - `mcp-${Date.now().toString(36)}-${Math.random().toString(36).substring(2, 8)}` - currentSession = { - id: sessionId, - xml: "", - version: 0, - lastSeenXml: "", - } + // Create session, or continue a saved one under its old id + const sessionId = + session_id ?? + `mcp-${Date.now().toString(36)}-${Math.random().toString(36).substring(2, 8)}` + currentSession = { id: sessionId, xml: "", lastSeenXml: "" } - // Open browser - const browserUrl = `http://localhost:${port}?mcp=${sessionId}` - await open(browserUrl) + // Open browser + const browserUrl = `http://localhost:${port}?mcp=${sessionId}` + await open(browserUrl) - // A saved diagram comes back from its file. lastSeenXml stays - // empty, so edit_diagram asks for get_diagram first. - const restored = sessionState(sessionId) - let intro = "Session started successfully!" - if (session_id) { - // blank: the browser asked first and nothing was saved - if (restored && !restored.blank) { - const doc = parseMxfile(restored.xml) - const pages = doc ? listPagesFromDoc(doc) : [] - const pageSummary = - pages.length > 0 - ? `Pages (${pages.length}): ${pages.map((p) => `[${p.index}] id=${p.id} name="${p.name}" cells=${p.cellCount}`).join(" | ")}` - : "no pages parsed" - intro = `Resumed session ${sessionId}.\n\n${pageSummary}\n\nCall get_diagram before edit_diagram.` - } else { - intro = `No saved diagram for ${sessionId}; starting blank.` - } - } - - const savePath = autosaver.pathFor(sessionId) - const saveNote = savePath - ? `\n\nAuto-save: after every change the diagram is saved to ${savePath}. To continue this diagram in a later conversation, call start_session with session_id=${sessionId}. list_saved_diagrams lists older diagrams.` - : "" - const rulesNote = `\n\nYour own drawing rules: write them in ${customInstructionsPath()} (Markdown, up to 5000 characters, read on every call).` - - log.info(`Started session ${sessionId}, browser at ${browserUrl}`) - - return { - content: [ - { - type: "text", - text: `${intro}\n\nSession ID: ${sessionId}\nBrowser URL: ${browserUrl}\n\nThe browser will now show real-time diagram updates.${saveNote}${rulesNote}\n\n${guideText()}`, - }, - ], - } - } catch (error) { - const message = - error instanceof Error ? error.message : String(error) - log.error("start_session failed:", message) - return { - content: [{ type: "text", text: `Error: ${message}` }], - isError: true, + // A saved diagram comes back from its file. lastSeenXml stays + // empty, so edit_diagram asks for get_diagram first. + const restored = sessionState(sessionId) + let intro = "Session started successfully!" + if (session_id) { + // blank: the browser asked first and nothing was saved + if (restored && !restored.blank) { + intro = `Resumed session ${sessionId}.\n\n${describePages(restored.xml)}\n\nCall get_diagram before edit_diagram.` + } else { + intro = `No saved diagram for ${sessionId}; starting blank.` } } + + const savePath = autosaver.pathFor(sessionId) + const saveNote = savePath + ? `\n\nAuto-save: after every change the diagram is saved to ${savePath}. To continue this diagram in a later conversation, call start_session with session_id=${sessionId}. list_saved_diagrams lists older diagrams.` + : "" + const rulesNote = `\n\nYour own drawing rules: write them in ${customInstructionsPath()} (Markdown, up to 5000 characters, read on every call).` + + log.info(`Started session ${sessionId}, browser at ${browserUrl}`) + + return ok( + `${intro}\n\nSession ID: ${sessionId}\nBrowser URL: ${browserUrl}\n\nThe browser will now show real-time diagram updates.${saveNote}${rulesNote}\n\n${guideText()}`, + ) }, + { write: true }, ) /** "2026-10-04 09:30" in the server's local time */ @@ -437,7 +472,7 @@ function formatLocalTime(date: Date): string { } // Tool: list_saved_diagrams -server.registerTool( +registerTool( "list_saved_diagrams", { title: "List saved diagrams", @@ -448,67 +483,43 @@ server.registerTool( annotations: { readOnlyHint: true, openWorldHint: false }, }, async () => { - try { - const dir = autosaver.dataDir() - if (!dir) { - return { - content: [ - { - type: "text", - text: "Auto-save is off (DRAWIO_DATA_DIR=off), so there are no saved diagrams.", - }, - ], - } - } - const saved = autosaver.list() - if (saved.length === 0) { - return { - content: [ - { type: "text", text: `No saved diagrams in ${dir}.` }, - ], - } - } - const fs = await import("node:fs/promises") - const lines: string[] = [] - for (const { sessionId, path, savedAt } of saved) { - // A file may hold compressed pages (saved by draw.io itself) - let pages = "unreadable" - try { - const loaded = parseDrawioFileContent( - await fs.readFile(path, "utf-8"), - ) - const doc = loaded.ok ? parseMxfile(loaded.xml) : null - if (doc) { - pages = listPagesFromDoc(doc) - .map((p) => `${p.name} (${p.cellCount} cells)`) - .join(", ") - } - } catch (error) { - log.warn( - `Could not read the saved diagram ${path}: ${error}`, - ) - } - const current = - sessionId === currentSession?.id ? " (current)" : "" - lines.push( - `${sessionId} saved ${formatLocalTime(savedAt)} pages: ${pages} file: ${path}${current}`, - ) - } - return { content: [{ type: "text", text: lines.join("\n") }] } - } catch (error) { - const message = - error instanceof Error ? error.message : String(error) - log.error("list_saved_diagrams failed:", message) - return { - content: [{ type: "text", text: `Error: ${message}` }], - isError: true, - } + const dir = autosaver.dataDir() + if (!dir) { + return ok( + "Auto-save is off (DRAWIO_DATA_DIR=off), so there are no saved diagrams.", + ) } + const saved = autosaver.list() + if (saved.length === 0) return ok(`No saved diagrams in ${dir}.`) + const lines: string[] = [] + for (const { sessionId, path, savedAt } of saved) { + // A file may hold compressed pages (saved by draw.io itself) + let pages = "unreadable" + try { + const loaded = parseDrawioFileContent( + await readFile(path, "utf-8"), + ) + const doc = loaded.ok ? parseMxfile(loaded.xml) : null + if (doc) { + pages = listPagesFromDoc(doc) + .map((p) => `${p.name} (${p.cellCount} cells)`) + .join(", ") + } + } catch (error) { + log.warn(`Could not read the saved diagram ${path}: ${error}`) + } + const current = + sessionId === currentSession?.id ? " (current)" : "" + lines.push( + `${sessionId} saved ${formatLocalTime(savedAt)} pages: ${pages} file: ${path}${current}`, + ) + } + return ok(lines.join("\n")) }, ) // Tool: create_new_diagram -registerWriteTool( +registerTool( "create_new_diagram", { title: "Create new diagram", @@ -536,123 +547,74 @@ To clear the canvas to one blank page, send only the two root cells 0) { - log.info(`XML auto-fixed: ${prepared.fixes.join(", ")}`) - } - // Every later tool can assume session.xml is an mxfile - const xml = prepared.xml - - log.info(`Setting diagram content, ${xml.length} chars`) - - // Sync from browser state first - const browserState = sessionState(currentSession.id) - if (browserState?.xml) { - currentSession.xml = browserState.xml - } - - // Save user's state before AI overwrites (with cached SVG) - if (currentSession.xml) { - keepInHistory( - currentSession.id, - currentSession.xml, - browserState?.svg || "", - ) - } - - // Update session state - currentSession.xml = xml - currentSession.version++ - - // 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, "") - - // Report page count back to the caller so the LLM learns whether - // multi-page worked or fell back to single. - const doc = parseMxfile(xml) - const pages = doc ? listPagesFromDoc(doc) : [] - const pageSummary = - pages.length > 1 - ? `${pages.length} pages: ${pages.map((p) => `${p.index}:${p.name}`).join(", ")}` - : pages.length === 1 - ? `1 page: ${pages[0].name}` - : "no pages parsed" - - log.info(`Diagram content set successfully (${pageSummary})`) - - // Only the root cells: the model asked for an empty canvas - // (several empty pages get the page summary) - const cleared = !hasCells(xml) && pages.length <= 1 - const content: CallToolResult["content"] = [ - { - type: "text", - text: cleared - ? "Canvas cleared: one blank page. The previous diagram is in History." - : `Diagram content set successfully!\n\nThe diagram is now visible in your browser.\n\nXML length: ${xml.length} characters\n${pageSummary}`, - }, - ] - // The page on screen is the new document's first page - if (screenshot ?? config.autoScreenshot) { - const shot = await captureScreenshot( - currentSession.id, - xml, - {}, - AFTER_WRITE_NEXT, - ) - content.push( - ...(shot.ok - ? shot.content - : [ - { - type: "text" as const, - text: `Screenshot skipped: ${shot.note}`, - }, - ]), - ) - } - return { content } - } catch (error) { - const message = - error instanceof Error ? error.message : String(error) - log.error("create_new_diagram failed:", message) - return { - content: [{ type: "text", text: `Error: ${message}` }], - isError: true, - } + const session = activeSession() + const prepared = prepareNewDiagram(inputXml) + if (!prepared.ok) { + log.error(prepared.error) + return fail(prepared.error) } + if (prepared.fixes.length > 0) { + log.info(`XML auto-fixed: ${prepared.fixes.join(", ")}`) + } + // Every later tool can assume session.xml is an mxfile + const xml = prepared.xml + + log.info(`Setting diagram content, ${xml.length} chars`) + + // The user's state goes to History before the AI overwrites it + const browserState = syncFromBrowser(session) + commitDocument(session, xml, true, browserState?.svg) + + // Report page count back to the caller so the LLM learns whether + // multi-page worked or fell back to single. + const doc = parseMxfile(xml) + const pages = doc ? listPagesFromDoc(doc) : [] + const pageSummary = + pages.length > 1 + ? `${pages.length} pages: ${pages.map((p) => `${p.index}:${p.name}`).join(", ")}` + : pages.length === 1 + ? `1 page: ${pages[0].name}` + : "no pages parsed" + + log.info(`Diagram content set successfully (${pageSummary})`) + + // Only the root cells: the model asked for an empty canvas + // (several empty pages get the page summary) + const cleared = !hasCells(xml) && pages.length <= 1 + const content: CallToolResult["content"] = [ + { + type: "text", + text: cleared + ? "Canvas cleared: one blank page. The previous diagram is in History." + : `Diagram content set successfully!\n\nThe diagram is now visible in your browser.\n\nXML length: ${xml.length} characters\n${pageSummary}`, + }, + ] + // The page on screen is the new document's first page + if (screenshot ?? config.autoScreenshot) { + const shot = await captureScreenshot( + session.id, + xml, + {}, + AFTER_WRITE_NEXT, + ) + content.push( + ...(shot.ok + ? shot.content + : [ + { + type: "text" as const, + text: `Screenshot skipped: ${shot.note}`, + }, + ]), + ) + } + return { content } }, + { write: true, needsSession: true }, ) // Tool: load_diagram -registerWriteTool( +registerTool( "load_diagram", { title: "Load diagram file", @@ -681,154 +643,75 @@ registerWriteTool( annotations: { openWorldHint: false }, }, async ({ path, xml: inlineXml }) => { - try { - // Argument validation comes before the session check: a bad - // argument is a caller error and should be reported as such. - if (path !== undefined && inlineXml !== undefined) { - return { - content: [ - { - type: "text", - text: "Error: Provide either 'path' or 'xml', not both.", - }, - ], - isError: true, - } - } - if (path === undefined && inlineXml === undefined) { - return { - content: [ - { - type: "text", - text: "Error: Provide either 'path' (a .drawio file to read) or 'xml' (the file's content).", - }, - ], - isError: true, - } - } - if (!currentSession) { - return { - content: [ - { - type: "text", - text: "Error: No active session. Please call start_session first.", - }, - ], - isError: true, - } - } - - // Exactly one of path/xml is present (validated above). - let content = "" - let sourceLabel = "" - if (inlineXml !== undefined) { - content = inlineXml - sourceLabel = "inline XML" - } else if (path !== undefined) { - const fs = await import("node:fs/promises") - const nodePath = await import("node:path") - const absolutePath = nodePath.resolve(expandHome(path)) - - try { - // A pipe or device could be read forever, and the other - // write tools wait for this one - if (!(await fs.stat(absolutePath)).isFile()) { - throw new Error("not a regular file") - } - content = await fs.readFile(absolutePath, "utf-8") - } catch (e) { - const msg = e instanceof Error ? e.message : String(e) - return { - content: [ - { - type: "text", - text: `Error: Cannot read file ${absolutePath}: ${msg}`, - }, - ], - isError: true, - } - } - sourceLabel = absolutePath - } - - const loaded = parseDrawioFileContent(content) - if (!loaded.ok) { - return { - content: [{ type: "text", text: `Error: ${loaded.error}` }], - isError: true, - } - } - const xml = loaded.xml - - log.info( - `Loading diagram from ${sourceLabel} (${xml.length} chars)`, - ) - - // Save the user's current state before replacing (same flow as - // create_new_diagram). - const browserState = sessionState(currentSession.id) - if (browserState?.xml) { - currentSession.xml = browserState.xml - } - if (currentSession.xml) { - keepInHistory( - currentSession.id, - currentSession.xml, - browserState?.svg || "", - ) - } - - currentSession.xml = xml - currentSession.version++ - setState(currentSession.id, xml) - // Edit-gate semantics by source: - // - 'path': the model only supplied a path, so it doesn't know - // the file's cell IDs — keep the gate (one get_diagram first). - // - plain 'xml': the model supplied the exact content, same - // rationale as create_new_diagram — record it as seen. - // - compressed 'xml': the session now holds the decompressed - // form, which the model cannot derive from the compressed - // input — keep the gate. - const markSeen = - inlineXml !== undefined && !loaded.hadCompressedPages - currentSession.lastSeenXml = markSeen ? xml : "" - - addHistory(currentSession.id, xml, "") - - const doc = parseMxfile(xml) - const pages = doc ? listPagesFromDoc(doc) : [] - const pageSummary = - pages.length > 0 - ? `Pages (${pages.length}): ${pages.map((p) => `[${p.index}] id=${p.id} name="${p.name}" cells=${p.cellCount}`).join(" | ")}` - : "no pages parsed" - - log.info(`Diagram loaded (${pageSummary})`) - - const gateHint = markSeen - ? "" - : "\n\nCall get_diagram before edit_diagram — you haven't seen this file's cell IDs yet." - - return { - content: [ - { - type: "text", - text: `Diagram loaded from ${sourceLabel}!\n\nThe diagram is now visible in your browser.\n\n${pageSummary}${gateHint}`, - }, - ], - } - } catch (error) { - const message = - error instanceof Error ? error.message : String(error) - log.error("load_diagram failed:", message) - return { - content: [{ type: "text", text: `Error: ${message}` }], - isError: true, - } + // Argument validation comes before the session check: a bad + // argument is a caller error and should be reported as such. + if (path !== undefined && inlineXml !== undefined) { + return fail("Provide either 'path' or 'xml', not both.") } + if (path === undefined && inlineXml === undefined) { + return fail( + "Provide either 'path' (a .drawio file to read) or 'xml' (the file's content).", + ) + } + if (!currentSession) return fail(NO_SESSION) + const session = currentSession + + // Exactly one of path/xml is present (validated above). + let content = "" + let sourceLabel = "" + if (inlineXml !== undefined) { + content = inlineXml + sourceLabel = "inline XML" + } else if (path !== undefined) { + const absolutePath = resolve(expandHome(path)) + try { + // A pipe or device could be read forever, and the other + // write tools wait for this one + if (!(await stat(absolutePath)).isFile()) { + throw new Error("not a regular file") + } + content = await readFile(absolutePath, "utf-8") + } catch (e) { + const msg = e instanceof Error ? e.message : String(e) + return fail(`Cannot read file ${absolutePath}: ${msg}`) + } + sourceLabel = absolutePath + } + + const loaded = parseDrawioFileContent(content) + if (!loaded.ok) return fail(loaded.error) + const xml = loaded.xml + + log.info(`Loading diagram from ${sourceLabel} (${xml.length} chars)`) + + // The user's state goes to History before the file replaces it. + // Edit-gate semantics by source: + // - 'path': the model only supplied a path, so it doesn't know + // the file's cell IDs — keep the gate (one get_diagram first). + // - plain 'xml': the model supplied the exact content, same + // rationale as create_new_diagram — record it as seen. + // - compressed 'xml': the session now holds the decompressed + // form, which the model cannot derive from the compressed + // input — keep the gate. + const browserState = syncFromBrowser(session) + const markSeen = inlineXml !== undefined && !loaded.hadCompressedPages + commitDocument(session, xml, markSeen, browserState?.svg) + + const pageSummary = describePages(xml) + log.info(`Diagram loaded (${pageSummary})`) + + const gateHint = markSeen + ? "" + : "\n\nCall get_diagram before edit_diagram — you haven't seen this file's cell IDs yet." + + return ok( + `Diagram loaded from ${sourceLabel}!\n\nThe diagram is now visible in your browser.\n\n${pageSummary}${gateHint}`, + ) }, + { write: true }, ) // Tool: edit_diagram -registerWriteTool( +registerTool( "edit_diagram", { title: "Edit diagram", @@ -884,188 +767,125 @@ registerWriteTool( annotations: { openWorldHint: false }, }, async ({ operations, page_id, page_name, page_index, screenshot }) => { - try { - if (!currentSession) { - return { - content: [ - { - type: "text", - text: "Error: No active session. Please call start_session first.", - }, - ], - isError: true, - } - } + const session = activeSession() - // Fetch latest state from browser. Re-normalise to mxfile: the - // embed/sync path can hand back a bare , and adopting - // it verbatim would silently strip a multi-page document down to - // one page on the next write. - const browserState = sessionState(currentSession.id) - if (browserState?.xml) { - currentSession.xml = - normalizeToMxfile(browserState.xml) ?? browserState.xml - log.info("Fetched latest diagram state from browser") - } + // Fetch latest state from browser + const browserState = syncFromBrowser(session) + if (browserState?.xml) { + log.info("Fetched latest diagram state from browser") + } - if (!currentSession.xml) { - return { - content: [ - { - type: "text", - text: "Error: No diagram to edit. Please create a diagram first with create_new_diagram.", - }, - ], - isError: true, - } - } - - const pageSelector = pickPageSelector({ - page_id, - page_name, - page_index, - }) - - // 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 (!session.xml) { + return fail( + "No diagram to edit. Please create a diagram first with create_new_diagram.", ) - if (!gate.ok) { - log.warn( - gate.reason === "stale" - ? "edit_diagram rejected: the browser has changes the model has not seen" - : "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, - // unless other pages changed too. - const liveXml = browserState?.xml || currentSession.xml - // What the model saw, before markPageSeen records the live state - const before = currentSession.lastSeenXml - 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." - // Per-cell summary of the user's changes, when they can be compared - const summary = - gate.reason === "stale" - ? describeChanges(before, liveXml) - : "" - const next = - currentSession.lastSeenXml === liveXml - ? "Build your operations on this XML and retry." - : `${OTHER_PAGES_UNSEEN} Call get_diagram without a page selector, then retry.` - return { - content: [ - { - type: "text", - text: `Error: ${reason}\n\n${summary ? `${summary}\n\n` : ""}Current XML of ${describeSelector(pageSelector)}:\n\n${foldCells(targetPageXml(currentSession.xml, pageSelector))}\n\n${next}`, - }, - ], - isError: true, - } - } + } - log.info( - `Editing diagram with ${operations.length} operation(s) on ${describeSelector(pageSelector)}`, + const pageSelector: PageSelector = { page_id, page_name, page_index } + + // 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(session.lastSeenXml, browserState?.xml ?? "") + if (!gate.ok) { + log.warn( + gate.reason === "stale" + ? "edit_diagram rejected: the browser has changes the model has not seen" + : "edit_diagram rejected: the model has not seen the diagram yet", ) - - const outcome = editDiagram( - currentSession.xml, - operations as DiagramOperation[], + // The error carries the current page, so the model has now + // seen it and can retry without a get_diagram round-trip, + // unless other pages changed too. + const liveXml = browserState?.xml || session.xml + // What the model saw, before markPageSeen records the live state + const before = session.lastSeenXml + session.lastSeenXml = markPageSeen( + session.lastSeenXml, + liveXml, pageSelector, ) - if (!outcome.ok) { - log.warn(`Edit rejected: ${outcome.errors.join("; ")}`) - const text = outcome.pageError - ? `Error: ${outcome.errors[0]}` - : `Error: No changes were made because ${outcome.errors.length} operation(s) failed:\n${outcome.errors.map((e) => `- ${e}`).join("\n")}\n\nCurrent XML of ${describeSelector(pageSelector)}:\n\n${foldCells(targetPageXml(currentSession.xml, pageSelector))}\n\nFix the operations against this XML and retry.` - return { - content: [{ type: "text", text }], - isError: true, - } - } - if (outcome.fixes.length > 0) { - log.info(`new_xml auto-fixed: ${outcome.fixes.join("; ")}`) - } - const result = outcome.xml - - // Save the pre-edit state for undo (with cached SVG from browser). - // Done only once the edit applied: a rejected edit returns above - // without leaving a phantom history entry. - keepInHistory( - currentSession.id, - currentSession.xml, - browserState?.svg || "", + 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." + // Per-cell summary of the user's changes, when they can be compared + const summary = + gate.reason === "stale" ? describeChanges(before, liveXml) : "" + const next = + session.lastSeenXml === liveXml + ? "Build your operations on this XML and retry." + : `${OTHER_PAGES_UNSEEN} Call get_diagram without a page selector, then retry.` + return fail( + `${reason}\n\n${summary ? `${summary}\n\n` : ""}Current XML of ${describeSelector(pageSelector)}:\n\n${foldCells(targetPageXml(session.xml, pageSelector))}\n\n${next}`, ) - - // Update state - currentSession.xml = result - currentSession.version++ - - // 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, "") - - log.info(`Diagram edited successfully`) - - const content: CallToolResult["content"] = [ - { - type: "text", - text: `Diagram edited successfully!\n\nApplied ${outcome.applied} operation(s) on ${describeSelector(pageSelector)}.`, - }, - ] - // Without a selector the edit went to the first page, while the - // page on screen may be another tab: name the page, so the PNG - // shows the edited one (a pageId export leaves the view alone) - if (screenshot ?? config.autoScreenshot) { - const shot = await captureScreenshot( - currentSession.id, - result, - hasPageSelector(pageSelector) - ? pageSelector - : { page_index: 0 }, - AFTER_WRITE_NEXT, - ) - content.push( - ...(shot.ok - ? shot.content - : [ - { - type: "text" as const, - text: `Screenshot skipped: ${shot.note}`, - }, - ]), - ) - } - return { content } - } catch (error) { - const message = - error instanceof Error ? error.message : String(error) - log.error("edit_diagram failed:", message) - return { - content: [{ type: "text", text: `Error: ${message}` }], - isError: true, - } } + + log.info( + `Editing diagram with ${operations.length} operation(s) on ${describeSelector(pageSelector)}`, + ) + + const outcome = editDiagram( + session.xml, + operations as DiagramOperation[], + pageSelector, + ) + if (!outcome.ok) { + log.warn(`Edit rejected: ${outcome.errors.join("; ")}`) + return fail( + outcome.pageError + ? outcome.errors[0] + : `No changes were made because ${outcome.errors.length} operation(s) failed:\n${outcome.errors.map((e) => `- ${e}`).join("\n")}\n\nCurrent XML of ${describeSelector(pageSelector)}:\n\n${foldCells(targetPageXml(session.xml, pageSelector))}\n\nFix the operations against this XML and retry.`, + ) + } + if (outcome.fixes.length > 0) { + log.info(`new_xml auto-fixed: ${outcome.fixes.join("; ")}`) + } + const result = outcome.xml + + // The pre-edit state goes to History for undo, only once the edit + // applied: a rejected edit returns above without leaving a phantom + // history entry. + commitDocument(session, result, true, browserState?.svg) + + log.info(`Diagram edited successfully`) + + const content: CallToolResult["content"] = [ + { + type: "text", + text: `Diagram edited successfully!\n\nApplied ${outcome.applied} operation(s) on ${describeSelector(pageSelector)}.`, + }, + ] + // Without a selector the edit went to the first page, while the + // page on screen may be another tab: name the page, so the PNG + // shows the edited one (a pageId export leaves the view alone) + if (screenshot ?? config.autoScreenshot) { + const shot = await captureScreenshot( + session.id, + result, + hasPageSelector(pageSelector) + ? pageSelector + : { page_index: 0 }, + AFTER_WRITE_NEXT, + ) + content.push( + ...(shot.ok + ? shot.content + : [ + { + type: "text" as const, + text: `Screenshot skipped: ${shot.note}`, + }, + ]), + ) + } + return { content } }, + { write: true, needsSession: true }, ) // Tool: get_diagram -server.registerTool( +registerTool( "get_diagram", { title: "Get diagram", @@ -1080,142 +900,82 @@ server.registerTool( }, annotations: { readOnlyHint: true, openWorldHint: false }, }, - async (input) => { - // Defensive: when every field is optional an MCP client could in - // principle invoke us with no `arguments` field. The SDK's zod parse - // normally produces `{}` in that case, but we coalesce explicitly so - // a destructure of `undefined` can never throw before we reach the - // session-existence check. - const { page_id, page_name, page_index } = input ?? {} - try { - if (!currentSession) { - return { - content: [ - { - type: "text", - text: "Error: No active session. Please call start_session first.", - }, - ], - isError: true, - } - } - - // start_session may replace currentSession while this waits - const session = currentSession - // Request browser to push fresh state and wait for it (an - // expired session first gets its saved file back to sync) - let staleNote = "" - restoreSavedSession(session.id) - const syncRequested = requestSync(session.id) - if (syncRequested) { - const synced = await waitForSync(session.id) - if (!synced) { - log.warn("get_diagram: sync timeout - state may be stale") - staleNote = - "\n\nNote: the browser did not respond, so this XML may not include the user's latest manual edits (is the preview tab open?)." - } - } - - // Fetch latest state from browser, re-normalising to mxfile so a - // bare pushed back by the embed/sync path doesn't - // strip page structure (see edit_diagram for the same guard). - const browserState = sessionState(session.id) - if (browserState?.xml) { - session.xml = - normalizeToMxfile(browserState.xml) ?? browserState.xml - } - - if (!session.xml) { - return { - content: [ - { - type: "text", - text: "No diagram exists yet. Use create_new_diagram to create one.", - }, - ], - } - } - - 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. - const liveXml = browserState?.xml || session.xml - // What the user changed since the model last looked, computed - // before lastSeenXml is overwritten below - const changes = session.lastSeenXml - ? describeChanges(session.lastSeenXml, liveXml) - : "" - const changesNote = changes ? `\n\nNote: ${changes}` : "" - const doc = parseMxfile(session.xml) - const pages = doc ? listPagesFromDoc(doc) : [] - const pageList = pages.length - ? `Pages (${pages.length}): ${pages.map((p) => `[${p.index}] id=${p.id} name="${p.name}" cells=${p.cellCount}`).join(" | ")}` - : "No wrapper detected (legacy single-page session)." - - // No selector → return full mxfile - if (!hasPageSelector(pageSelector)) { - session.lastSeenXml = liveXml - return { - content: [ - { - type: "text", - text: `Current diagram XML:\n\n${foldCells(session.xml)}\n\n${pageList}${staleNote}${changesNote}`, - }, - ], - } - } - - // Selector → return a single-page projection - const projection = projectPage(session.xml, pageSelector) - if (!projection.ok) { - return { - content: [ - { - type: "text", - text: - projection.reason === "parse" - ? `Error: a page selector was given but the current session XML could not be parsed as a multi-page (it may be a legacy single-page document or malformed), so it has no addressable pages.\n\n${pageList}` - : `Error: Page ${describeSelector(pageSelector)} not found.\n\n${pageList}`, - }, - ], - isError: true, - } - } - // One page shown counts for all only if the others are as the - // model saw them last - session.lastSeenXml = markPageSeen( - session.lastSeenXml, - liveXml, - pageSelector, - ) - const otherPagesNote = - session.lastSeenXml === liveXml - ? "" - : `\n\nNote: ${OTHER_PAGES_UNSEEN} Call get_diagram without a page selector before editing.` - return { - content: [ - { - type: "text", - text: `Page ${projection.index} ("${projection.name}"):\n\n${foldCells(projection.xml)}\n\n${pageList}${staleNote}${changesNote}${otherPagesNote}`, - }, - ], - } - } catch (error) { - const message = - error instanceof Error ? error.message : String(error) - log.error("get_diagram failed:", message) - return { - content: [{ type: "text", text: `Error: ${message}` }], - isError: true, + async ({ page_id, page_name, page_index }) => { + // start_session may replace currentSession while this waits + const session = activeSession() + // Request browser to push fresh state and wait for it (an + // expired session first gets its saved file back to sync) + let staleNote = "" + restoreSavedSession(session.id) + const syncRequested = requestSync(session.id) + if (syncRequested) { + const synced = await waitForSync(session.id) + if (!synced) { + log.warn("get_diagram: sync timeout - state may be stale") + staleNote = + "\n\nNote: the browser did not respond, so this XML may not include the user's latest manual edits (is the preview tab open?)." } } + + // Fetch latest state from browser + const browserState = syncFromBrowser(session) + + if (!session.xml) { + return ok( + "No diagram exists yet. Use create_new_diagram to create one.", + ) + } + + const pageSelector: PageSelector = { 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. + const liveXml = browserState?.xml || session.xml + // What the user changed since the model last looked, computed + // before lastSeenXml is overwritten below + const changes = session.lastSeenXml + ? describeChanges(session.lastSeenXml, liveXml) + : "" + const changesNote = changes ? `\n\nNote: ${changes}` : "" + const pageList = describePages( + session.xml, + "No wrapper detected (legacy single-page session).", + ) + + // No selector → return full mxfile + if (!hasPageSelector(pageSelector)) { + session.lastSeenXml = liveXml + return ok( + `Current diagram XML:\n\n${foldCells(session.xml)}\n\n${pageList}${staleNote}${changesNote}`, + ) + } + + // Selector → return a single-page projection + const projection = projectPage(session.xml, pageSelector) + if (!projection.ok) { + return fail( + projection.reason === "parse" + ? `a page selector was given but the current session XML could not be parsed as a multi-page (it may be a legacy single-page document or malformed), so it has no addressable pages.\n\n${pageList}` + : `Page ${describeSelector(pageSelector)} not found.\n\n${pageList}`, + ) + } + // One page shown counts for all only if the others are as the + // model saw them last + session.lastSeenXml = markPageSeen( + session.lastSeenXml, + liveXml, + pageSelector, + ) + const otherPagesNote = + session.lastSeenXml === liveXml + ? "" + : `\n\nNote: ${OTHER_PAGES_UNSEEN} Call get_diagram without a page selector before editing.` + return ok( + `Page ${projection.index} ("${projection.name}"):\n\n${foldCells(projection.xml)}\n\n${pageList}${staleNote}${changesNote}${otherPagesNote}`, + ) }, + { needsSession: true }, ) // The browser bridge has one export slot per session, so export requests @@ -1251,10 +1011,7 @@ function exportViaBrowser( const live = getState(sessionId) if (live) { live.exportData = undefined - live.exportFormat = undefined - live.exportXml = undefined - live.exportOptions = undefined - live.exportId = undefined + live.pendingExport = undefined } return exportData }) @@ -1281,18 +1038,6 @@ function previewStalledNote(sessionId: string, next = "then retry") { return `The preview tab is not responding (browsers pause background tabs). Ask the user to bring the preview tab to the front (http://localhost:${getServerPort()}?mcp=${sessionId}), ${next}.` } -function previewStalledError(sessionId: string) { - return { - content: [ - { - type: "text" as const, - text: `Error: ${previewStalledNote(sessionId)}`, - }, - ], - isError: true, - } -} - /** The of the page a selector targets, for draw.io's pageId. */ function pageIdFor(xml: string, selector: PageSelector): string | undefined { const doc = parseMxfile(xml) @@ -1404,7 +1149,7 @@ async function captureScreenshot( } // Tool: screenshot_diagram -server.registerTool( +registerTool( "screenshot_diagram", { title: "Screenshot diagram", @@ -1415,48 +1160,22 @@ server.registerTool( inputSchema: { ...pageSelectorSchema }, annotations: { readOnlyHint: true, openWorldHint: false }, }, - async (input) => { - const { page_id, page_name, page_index } = input ?? {} - try { - if (!currentSession) { - return { - content: [ - { - type: "text", - text: "Error: No active session. Please call start_session first.", - }, - ], - isError: true, - } - } - const xml = - sessionState(currentSession.id)?.xml || currentSession.xml - const shot = await captureScreenshot( - currentSession.id, - xml, - pickPageSelector({ page_id, page_name, page_index }), - ) - if (!shot.ok) { - return { - content: [{ type: "text", text: `Error: ${shot.note}` }], - isError: true, - } - } - return { content: shot.content } - } catch (error) { - const message = - error instanceof Error ? error.message : String(error) - log.error("screenshot_diagram failed:", message) - return { - content: [{ type: "text", text: `Error: ${message}` }], - isError: true, - } - } + async ({ page_id, page_name, page_index }) => { + const session = activeSession() + const xml = sessionState(session.id)?.xml || session.xml + const shot = await captureScreenshot(session.id, xml, { + page_id, + page_name, + page_index, + }) + if (!shot.ok) return fail(shot.note) + return { content: shot.content } }, + { needsSession: true }, ) // Tool: export_diagram -server.registerTool( +registerTool( "export_diagram", { title: "Export diagram", @@ -1487,253 +1206,169 @@ server.registerTool( }, async ({ path: rawPath, format, page_id, page_name, page_index }) => { const path = expandHome(rawPath) - try { - if (!currentSession) { - return { - content: [ - { - type: "text", - text: "Error: No active session. Please call start_session first.", - }, - ], - isError: true, - } - } + // start_session may replace currentSession while this waits + const session = activeSession() - // start_session may replace currentSession while this waits - const session = currentSession + // Detect format from extension if not specified + const lowerPath = path.toLowerCase() + const detectedFormat = + format || + (lowerPath.endsWith(".drawio.svg") + ? "drawio.svg" + : lowerPath.endsWith(".png") + ? "png" + : lowerPath.endsWith(".svg") + ? "svg" + : "drawio") - // Detect format from extension if not specified - const lowerPath = path.toLowerCase() - const detectedFormat = - format || - (lowerPath.endsWith(".drawio.svg") - ? "drawio.svg" - : lowerPath.endsWith(".png") - ? "png" - : lowerPath.endsWith(".svg") - ? "svg" - : "drawio") - - // The .drawio file is written from the state, so get the - // user's latest edits into it first, as get_diagram does (the - // images are made by the browser from its canvas) - let syncNote = "" - if (detectedFormat === "drawio") { - restoreSavedSession(session.id) - if (!requestSync(session.id)) { - syncNote = - "\n\nNote: the preview was not reachable, so the file may not include the user's latest manual edits." - } else if (!(await waitForSync(session.id))) { - log.warn( - "export_diagram: sync timeout - state may be stale", - ) - syncNote = - "\n\nNote: the browser did not respond, so the file may not include the user's latest manual edits (is the preview tab open?)." - } - } - - // Fetch latest state, re-normalised to mxfile so a page - // selector works on a bare pushed by the browser - const browserState = sessionState(session.id) - if (browserState?.xml) { - session.xml = - normalizeToMxfile(browserState.xml) ?? browserState.xml - } - - if (!session.xml) { - return { - content: [ - { - type: "text", - text: "Error: No diagram to export. Please create a diagram first.", - }, - ], - isError: true, - } - } - - const pageSelector = pickPageSelector({ - page_id, - page_name, - page_index, - }) - - const fs = await import("node:fs/promises") - const nodePath = await import("node:path") - - // .drawio path - write XML directly (no browser round-trip). - if (detectedFormat === "drawio") { - let filePath = path - if (!filePath.toLowerCase().endsWith(".drawio")) { - filePath = `${filePath}.drawio` - } - const absolutePath = nodePath.resolve(filePath) - - let outXml = session.xml - if (hasPageSelector(pageSelector)) { - const projection = projectPage(session.xml, pageSelector) - if (!projection.ok) { - return { - content: [ - { - type: "text", - text: - projection.reason === "parse" - ? "Error: Cannot parse current session XML as ; cannot project a single page." - : `Error: Page ${describeSelector(pageSelector)} not found for export.`, - }, - ], - isError: true, - } - } - outXml = projection.xml - } - - await fs.writeFile(absolutePath, outXml, "utf-8") - log.info(`Diagram exported to ${absolutePath}`) - return { - content: [ - { - type: "text", - text: `Diagram exported successfully!\n\nFile: ${absolutePath}\nSize: ${outXml.length} characters${syncNote}`, - }, - ], - } - } - - // PNG or SVG: request browser to export via iframe. Replace a - // known extension that does not match the format. - let filePath = path - const suffix = `.${detectedFormat}` - if (!lowerPath.endsWith(suffix)) { - const known = [".drawio.svg", ".drawio", ".png", ".svg"].find( - (e) => lowerPath.endsWith(e), - ) - if (known) filePath = filePath.slice(0, -known.length) - filePath = `${filePath}${suffix}` - } - const absolutePath = nodePath.resolve(filePath) - // draw.io's name for an SVG with the diagram embedded - const browserFormat = - detectedFormat === "drawio.svg" ? "xmlsvg" : detectedFormat - - const state = sessionState(session.id) - if (!state) { - return { - content: [ - { - type: "text", - text: "Error: Session state not found. Is the browser open?", - }, - ], - isError: true, - } - } - if (previewStalled(session.id)) { - return previewStalledError(session.id) - } - - // ----------------------------------------------------------------- - // Page-targeted PNG/SVG export. - // - // drawio's JSON embed protocol has no working `selectPage` action, - // so to export a specific page we build a single-page - // projection and hand it to the browser bridge alongside the export - // request. The bridge loads the projection, waits for draw.io's own - // render, exports, then reloads the user's real document — entirely - // browser-side. The canonical session state is never mutated here, - // so there is no restore race and no concurrent-edit clobbering. - // ----------------------------------------------------------------- - // PNG: draw.io renders any page by id, without touching the - // page on screen. SVG export has no page option, so it still - // needs the projection below. - let projectionXml: string | undefined - let pngPageId: string | undefined - if (hasPageSelector(pageSelector) && detectedFormat === "png") { - pngPageId = pageIdFor(session.xml, pageSelector) - } - if (hasPageSelector(pageSelector) && !pngPageId) { - const projection = projectPage(session.xml, pageSelector) - if (!projection.ok) { - return { - content: [ - { - type: "text", - text: - projection.reason === "parse" - ? "Error: Cannot parse current session XML as ; cannot target page for export." - : `Error: Page ${describeSelector(pageSelector)} not found for export.`, - }, - ], - isError: true, - } - } - projectionXml = projection.xml - } - - const exportData = await exportViaBrowser( - session.id, - browserFormat, - projectionXml, - pngPageId ? { pageId: pngPageId } : undefined, - ) - - if (!exportData) { - return { - content: [ - { - type: "text", - text: projectionXml - ? "Error: Export timed out after loading the single-page projection. The browser may be closed or unresponsive." - : "Error: Export timed out. Make sure the browser tab is open and the diagram is loaded.", - }, - ], - isError: true, - } - } - - // Decode and write - if (detectedFormat === "png") { - const base64 = exportData.replace( - /^data:image\/png;base64,/, - "", - ) - await fs.writeFile(absolutePath, Buffer.from(base64, "base64")) - } else { - let svgContent = exportData - if (svgContent.startsWith("data:image/svg+xml;base64,")) { - const base64 = svgContent.replace( - /^data:image\/svg\+xml;base64,/, - "", - ) - svgContent = Buffer.from(base64, "base64").toString("utf-8") - } - await fs.writeFile(absolutePath, svgContent, "utf-8") - } - - const stat = await fs.stat(absolutePath) - log.info( - `Diagram exported to ${absolutePath} (${detectedFormat}, ${stat.size} bytes)`, - ) - return { - content: [ - { - type: "text", - text: `Diagram exported successfully!\n\nFile: ${absolutePath}\nFormat: ${detectedFormat}\nSize: ${stat.size} bytes`, - }, - ], - } - } catch (error) { - const message = - error instanceof Error ? error.message : String(error) - log.error("export_diagram failed:", message) - return { - content: [{ type: "text", text: `Error: ${message}` }], - isError: true, + // The .drawio file is written from the state, so get the + // user's latest edits into it first, as get_diagram does (the + // images are made by the browser from its canvas) + let syncNote = "" + if (detectedFormat === "drawio") { + restoreSavedSession(session.id) + if (!requestSync(session.id)) { + syncNote = + "\n\nNote: the preview was not reachable, so the file may not include the user's latest manual edits." + } else if (!(await waitForSync(session.id))) { + log.warn("export_diagram: sync timeout - state may be stale") + syncNote = + "\n\nNote: the browser did not respond, so the file may not include the user's latest manual edits (is the preview tab open?)." } } + + // Fetch latest state, so a page selector works on a bare + // pushed by the browser + const browserState = syncFromBrowser(session) + + if (!session.xml) { + return fail("No diagram to export. Please create a diagram first.") + } + + const pageSelector: PageSelector = { page_id, page_name, page_index } + + // .drawio path - write XML directly (no browser round-trip). + if (detectedFormat === "drawio") { + let filePath = path + if (!filePath.toLowerCase().endsWith(".drawio")) { + filePath = `${filePath}.drawio` + } + const absolutePath = resolve(filePath) + + let outXml = session.xml + if (hasPageSelector(pageSelector)) { + const projection = projectPage(session.xml, pageSelector) + if (!projection.ok) { + return fail( + projection.reason === "parse" + ? "Cannot parse current session XML as ; cannot project a single page." + : `Page ${describeSelector(pageSelector)} not found for export.`, + ) + } + outXml = projection.xml + } + + await writeFile(absolutePath, outXml, "utf-8") + log.info(`Diagram exported to ${absolutePath}`) + return ok( + `Diagram exported successfully!\n\nFile: ${absolutePath}\nSize: ${outXml.length} characters${syncNote}`, + ) + } + + // PNG or SVG: request browser to export via iframe. Replace a + // known extension that does not match the format. + let filePath = path + const suffix = `.${detectedFormat}` + if (!lowerPath.endsWith(suffix)) { + const known = [".drawio.svg", ".drawio", ".png", ".svg"].find((e) => + lowerPath.endsWith(e), + ) + if (known) filePath = filePath.slice(0, -known.length) + filePath = `${filePath}${suffix}` + } + const absolutePath = resolve(filePath) + // draw.io's name for an SVG with the diagram embedded + const browserFormat = + detectedFormat === "drawio.svg" ? "xmlsvg" : detectedFormat + + if (!browserState) { + return fail("Session state not found. Is the browser open?") + } + if (previewStalled(session.id)) { + return fail(previewStalledNote(session.id)) + } + + // ----------------------------------------------------------------- + // Page-targeted PNG/SVG export. + // + // drawio's JSON embed protocol has no working `selectPage` action, + // so to export a specific page we build a single-page + // projection and hand it to the browser bridge alongside the export + // request. The bridge loads the projection, waits for draw.io's own + // render, exports, then reloads the user's real document — entirely + // browser-side. The canonical session state is never mutated here, + // so there is no restore race and no concurrent-edit clobbering. + // ----------------------------------------------------------------- + // PNG: draw.io renders any page by id, without touching the + // page on screen. SVG export has no page option, so it still + // needs the projection below. + let projectionXml: string | undefined + let pngPageId: string | undefined + if (hasPageSelector(pageSelector) && detectedFormat === "png") { + pngPageId = pageIdFor(session.xml, pageSelector) + } + if (hasPageSelector(pageSelector) && !pngPageId) { + const projection = projectPage(session.xml, pageSelector) + if (!projection.ok) { + return fail( + projection.reason === "parse" + ? "Cannot parse current session XML as ; cannot target page for export." + : `Page ${describeSelector(pageSelector)} not found for export.`, + ) + } + projectionXml = projection.xml + } + + const exportData = await exportViaBrowser( + session.id, + browserFormat, + projectionXml, + pngPageId ? { pageId: pngPageId } : undefined, + ) + + if (!exportData) { + return fail( + projectionXml + ? "Export timed out after loading the single-page projection. The browser may be closed or unresponsive." + : "Export timed out. Make sure the browser tab is open and the diagram is loaded.", + ) + } + + // Decode and write + if (detectedFormat === "png") { + const base64 = exportData.replace(/^data:image\/png;base64,/, "") + await writeFile(absolutePath, Buffer.from(base64, "base64")) + } else { + let svgContent = exportData + if (svgContent.startsWith("data:image/svg+xml;base64,")) { + const base64 = svgContent.replace( + /^data:image\/svg\+xml;base64,/, + "", + ) + svgContent = Buffer.from(base64, "base64").toString("utf-8") + } + await writeFile(absolutePath, svgContent, "utf-8") + } + + const written = await stat(absolutePath) + log.info( + `Diagram exported to ${absolutePath} (${detectedFormat}, ${written.size} bytes)`, + ) + return ok( + `Diagram exported successfully!\n\nFile: ${absolutePath}\nFormat: ${detectedFormat}\nSize: ${written.size} bytes`, + ) }, + { needsSession: true }, ) /** @@ -1742,29 +1377,23 @@ server.registerTool( * parsed Document the caller can mutate, plus a writer that persists. */ async function loadMxfileForMutation(): Promise< - | { ok: true; doc: Document; writeBack: (newDoc: Document) => void } + | { ok: true; doc: Document; writeBack: () => void } | { ok: false; message: string } > { - if (!currentSession) { - return { - ok: false, - message: "No active session. Please call start_session first.", - } - } + if (!currentSession) return { ok: false, message: NO_SESSION } + const session = currentSession // Pull latest from browser so we don't clobber autosaved changes. - const browserState = sessionState(currentSession.id) - if (browserState?.xml) { - currentSession.xml = browserState.xml - } - if (!currentSession.xml) { + const browserState = syncFromBrowser(session) + if (!session.xml) { return { ok: false, message: "No diagram exists yet. Use create_new_diagram first, then page tools.", } } - // Make sure the in-memory shape is canonical mxfile before any CRUD. - const normalized = normalizeToMxfile(currentSession.xml) + // Make sure the in-memory shape is canonical mxfile before any CRUD + // (syncFromBrowser keeps XML it cannot normalise as it is). + const normalized = normalizeToMxfile(session.xml) if (!normalized) { return { ok: false, @@ -1772,46 +1401,37 @@ async function loadMxfileForMutation(): Promise< "Current session XML is neither nor ; cannot perform page operations.", } } - currentSession.xml = normalized + session.xml = normalized - const doc = parseMxfile(currentSession.xml) + const doc = parseMxfile(session.xml) if (!doc) { return { ok: false, message: "Failed to parse current session XML as .", } } - const sessionRef = currentSession return { ok: true, doc, - writeBack: (newDoc: Document) => { - const newXml = serializeMxfile(newDoc) - // The store may hold user edits the model has not seen yet. + writeBack: () => { + // The store may hold user edits the model has not seen yet; the + // result counts as seen only when the model built on the latest const sawLatest = checkEditGate( - sessionRef.lastSeenXml, + session.lastSeenXml, browserState?.xml ?? "", ).ok - // Save history before overwriting so the user can undo. - keepInHistory( - sessionRef.id, - sessionRef.xml, - browserState?.svg || "", + commitDocument( + session, + serializeMxfile(doc), + sawLatest, + browserState?.svg, ) - sessionRef.xml = newXml - sessionRef.version++ - setState(sessionRef.id, newXml) - // The model just wrote this exact state. If it had seen the state - // it built on, mark the result as seen so edit_diagram needs no - // extra get_diagram; otherwise edit_diagram must ask for one. - sessionRef.lastSeenXml = sawLatest ? newXml : "" - addHistory(sessionRef.id, newXml, "") }, } } // Tool: list_pages -server.registerTool( +registerTool( "list_pages", { title: "List pages", @@ -1821,53 +1441,20 @@ server.registerTool( annotations: { readOnlyHint: true, openWorldHint: false }, }, async () => { - try { - const loaded = await loadMxfileForMutation() - if (!loaded.ok) { - return { - content: [ - { type: "text", text: `Error: ${loaded.message}` }, - ], - isError: true, - } - } - const pages = listPagesFromDoc(loaded.doc) - if (pages.length === 0) { - return { - content: [ - { - type: "text", - text: "No pages in document.", - }, - ], - } - } - const lines = pages.map( - (p) => - ` [${p.index}] id="${p.id}" name="${p.name}" cells=${p.cellCount}`, - ) - return { - content: [ - { - type: "text", - text: `Pages (${pages.length}):\n${lines.join("\n")}`, - }, - ], - } - } catch (error) { - const message = - error instanceof Error ? error.message : String(error) - log.error("list_pages failed:", message) - return { - content: [{ type: "text", text: `Error: ${message}` }], - isError: true, - } - } + const loaded = await loadMxfileForMutation() + if (!loaded.ok) return fail(loaded.message) + const pages = listPagesFromDoc(loaded.doc) + if (pages.length === 0) return ok("No pages in document.") + const lines = pages.map( + (p) => + ` [${p.index}] id="${p.id}" name="${p.name}" cells=${p.cellCount}`, + ) + return ok(`Pages (${pages.length}):\n${lines.join("\n")}`) }, ) // Tool: add_page -registerWriteTool( +registerTool( "add_page", { title: "Add page", @@ -1901,114 +1488,40 @@ registerWriteTool( }, annotations: { destructiveHint: false, openWorldHint: false }, }, - async (input) => { - // All three fields optional — coalesce so a no-args call doesn't - // crash on destructure before we surface a proper MCP error. - const { name, id, xml } = input ?? {} - try { - const loaded = await loadMxfileForMutation() - if (!loaded.ok) { - return { - content: [ - { type: "text", text: `Error: ${loaded.message}` }, - ], - isError: true, - } - } + async ({ name, id, xml }) => { + const loaded = await loadMxfileForMutation() + if (!loaded.ok) return fail(loaded.message) - // If caller provided XML, validate it before splicing it in so we - // never get a half-broken mxfile written to the session. - // Named style definitions come out first and are expanded on the - // validated XML, like prepareNewDiagram - const { - classes, - xml: startXml, - error: styleError, - } = takeStyleDefinitions(xml ?? "") - if (styleError) { - return { - content: [{ type: "text", text: `Error: ${styleError}` }], - isError: true, - } - } - const reserved = startXml && reservedIdError(startXml) - if (reserved) { - return { - content: [{ type: "text", text: `Error: ${reserved}` }], - isError: true, - } - } - const truncated = startXml && truncatedCellError(startXml) - if (truncated) { - return { - content: [{ type: "text", text: `Error: ${truncated}` }], - isError: true, - } - } - let cleanXml: string | undefined = - startXml && wrapCellsInModel(startXml) - if (cleanXml) { - const { valid, error, fixed, fixes } = - validateAndFixXml(cleanXml) - if (fixed) { - cleanXml = fixed - log.info( - `add_page: starting XML auto-fixed: ${fixes.join(", ")}`, - ) - } - if (!valid && error) { - return { - content: [ - { - type: "text", - text: `Error: starting xml validation failed - ${error}`, - }, - ], - isError: true, - } - } - cleanXml = addDefaultStyles( - applyStyleClasses(expandCompactCells(cleanXml), classes), + // The starting content is prepared like a new diagram's page, so a + // half-broken model never gets into the session + let cleanXml: string | undefined + if (xml?.trim()) { + const prepared = prepareNewModel(xml) + if (!prepared.ok) return fail(prepared.error) + if (prepared.fixes.length > 0) { + log.info( + `add_page: starting XML auto-fixed: ${prepared.fixes.join(", ")}`, ) } - - let info - try { - info = addPageToDoc(loaded.doc, { id, name, xml: cleanXml }) - } catch (e) { - const msg = e instanceof Error ? e.message : String(e) - return { - content: [{ type: "text", text: `Error: ${msg}` }], - isError: true, - } - } - - loaded.writeBack(loaded.doc) - log.info( - `Added page id=${info.id} name="${info.name}" index=${info.index}`, - ) - return { - content: [ - { - type: "text", - text: `Page added.\n\nid=${info.id}\nname=${info.name}\nindex=${info.index}\ncells=${info.cellCount}`, - }, - ], - } - } catch (error) { - const message = - error instanceof Error ? error.message : String(error) - log.error("add_page failed:", message) - return { - content: [{ type: "text", text: `Error: ${message}` }], - isError: true, - } + cleanXml = prepared.xml } + + // (throws on a page id in use or a whole as content) + const info = addPageToDoc(loaded.doc, { id, name, xml: cleanXml }) + + loaded.writeBack() + log.info( + `Added page id=${info.id} name="${info.name}" index=${info.index}`, + ) + return ok( + `Page added.\n\nid=${info.id}\nname=${info.name}\nindex=${info.index}\ncells=${info.cellCount}`, + ) }, + { write: true }, ) // Tool: rename_page -registerWriteTool( +registerTool( "rename_page", { title: "Rename page", @@ -2024,73 +1537,33 @@ registerWriteTool( annotations: { destructiveHint: false, openWorldHint: false }, }, async ({ new_name, page_id, page_name, page_index }) => { - try { - const loaded = await loadMxfileForMutation() - if (!loaded.ok) { - return { - content: [ - { type: "text", text: `Error: ${loaded.message}` }, - ], - isError: true, - } - } + const loaded = await loadMxfileForMutation() + if (!loaded.ok) return fail(loaded.message) - const pageSelector = pickPageSelector({ - page_id, - page_name, - page_index, - }) - if (!hasPageSelector(pageSelector)) { - return { - content: [ - { - type: "text", - text: "Error: rename_page requires one of page_id, page_name, or page_index to identify the page.", - }, - ], - isError: true, - } - } - - const ok = renamePageInDoc(loaded.doc, pageSelector, new_name) - if (!ok) { - return { - content: [ - { - type: "text", - text: `Error: Page ${describeSelector(pageSelector)} not found.`, - }, - ], - isError: true, - } - } - - loaded.writeBack(loaded.doc) - log.info( - `Renamed page ${describeSelector(pageSelector)} → "${new_name}"`, + const pageSelector: PageSelector = { page_id, page_name, page_index } + if (!hasPageSelector(pageSelector)) { + return fail( + "rename_page requires one of page_id, page_name, or page_index to identify the page.", ) - return { - content: [ - { - type: "text", - text: `Page ${describeSelector(pageSelector)} renamed to "${new_name}".`, - }, - ], - } - } catch (error) { - const message = - error instanceof Error ? error.message : String(error) - log.error("rename_page failed:", message) - return { - content: [{ type: "text", text: `Error: ${message}` }], - isError: true, - } } + + if (!renamePageInDoc(loaded.doc, pageSelector, new_name)) { + return fail(`Page ${describeSelector(pageSelector)} not found.`) + } + + loaded.writeBack() + log.info( + `Renamed page ${describeSelector(pageSelector)} → "${new_name}"`, + ) + return ok( + `Page ${describeSelector(pageSelector)} renamed to "${new_name}".`, + ) }, + { write: true }, ) // Tool: delete_page -registerWriteTool( +registerTool( "delete_page", { title: "Delete page", @@ -2101,83 +1574,30 @@ registerWriteTool( }, annotations: { openWorldHint: false }, }, - async (input) => { - // All three fields are optional — coalesce so a no-args call returns - // a clean error message instead of crashing on destructure. - const { page_id, page_name, page_index } = input ?? {} - try { - const loaded = await loadMxfileForMutation() - if (!loaded.ok) { - return { - content: [ - { type: "text", text: `Error: ${loaded.message}` }, - ], - isError: true, - } - } + async ({ page_id, page_name, page_index }) => { + const loaded = await loadMxfileForMutation() + if (!loaded.ok) return fail(loaded.message) - const pageSelector = pickPageSelector({ - page_id, - page_name, - page_index, - }) - if (!hasPageSelector(pageSelector)) { - return { - content: [ - { - type: "text", - text: "Error: delete_page requires one of page_id, page_name, or page_index to identify the page.", - }, - ], - isError: true, - } - } - - const outcome = deletePageFromDoc(loaded.doc, pageSelector) - if (!outcome.ok) { - return { - content: [ - { - type: "text", - text: `Error: ${outcome.reason}.`, - }, - ], - isError: true, - } - } - - loaded.writeBack(loaded.doc) - log.info( - `Deleted page id=${outcome.deletedId} index=${outcome.deletedIndex}`, + const pageSelector: PageSelector = { page_id, page_name, page_index } + if (!hasPageSelector(pageSelector)) { + return fail( + "delete_page requires one of page_id, page_name, or page_index to identify the page.", ) - return { - content: [ - { - type: "text", - text: `Page deleted (id=${outcome.deletedId}, was at index ${outcome.deletedIndex}).`, - }, - ], - } - } catch (error) { - const message = - error instanceof Error ? error.message : String(error) - log.error("delete_page failed:", message) - return { - content: [{ type: "text", text: `Error: ${message}` }], - isError: true, - } } - }, -) -/** One line of pages for a history entry, for restore_version's texts. */ -function describePages(xml: string): string { - const doc = parseMxfile(xml) - const pages = doc ? listPagesFromDoc(doc) : [] - return pages.length > 0 - ? `Pages (${pages.length}): ${pages.map((p) => `[${p.index}] id=${p.id} name="${p.name}" cells=${p.cellCount}`).join(" | ")}` - : "no pages parsed" -} + const outcome = deletePageFromDoc(loaded.doc, pageSelector) + if (!outcome.ok) return fail(`${outcome.reason}.`) + + loaded.writeBack() + log.info( + `Deleted page id=${outcome.deletedId} index=${outcome.deletedIndex}`, + ) + return ok( + `Page deleted (id=${outcome.deletedId}, was at index ${outcome.deletedIndex}).`, + ) + }, + { write: true }, +) /** The versions restore_version can go to, numbered by steps_back. */ function describeVersions(entries: HistoryEntry[]): string { @@ -2198,7 +1618,7 @@ function describeVersions(entries: HistoryEntry[]): string { } // Tool: restore_version -registerWriteTool( +registerTool( "restore_version", { title: "Restore version", @@ -2216,100 +1636,50 @@ registerWriteTool( }, annotations: { destructiveHint: false, openWorldHint: false }, }, - async (input) => { - const stepsBack = input?.steps_back ?? 1 - try { - if (!currentSession) { - return { - content: [ - { - type: "text", - text: "Error: No active session. Please call start_session first.", - }, - ], - isError: true, - } - } - const session = currentSession + async ({ steps_back: stepsBack = 1 }) => { + const session = activeSession() - // The canvas, with the user's manual edits, is the state store - const current = sessionState(session.id)?.xml || session.xml - const candidates = otherVersions(session.id, current) - if (candidates.length === 0) { - return { - content: [ - { - type: "text", - text: "Error: No other version in History.", - }, - ], - isError: true, - } - } - if (candidates.length < stepsBack) { - return { - content: [ - { - type: "text", - text: `Error: Only ${candidates.length} other version(s) in History: ${describeVersions(candidates)}`, - }, - ], - isError: true, - } - } - const entry = candidates[stepsBack - 1] - - // Manual edits since the last entry are not in History yet; - // restoreHistoryEntry keeps them (a blank page excepted) - const hadManual = - contentFingerprint(current) !== - contentFingerprint(BLANK_MXFILE) && - !isInHistory(session.id, current) - - if (restoreHistoryEntry(session.id, entry.id) === null) { - return { - content: [ - { - type: "text", - text: "Error: That version is no longer in History.", - }, - ], - isError: true, - } - } - session.xml = entry.xml - session.version++ - // As load_diagram: the model may never have seen this XML, so - // the edit gate requires one get_diagram before edits - session.lastSeenXml = "" - - log.info( - `Restored the version ${stepsBack} step(s) back (entry ${entry.id})`, - ) - - const manualNote = hadManual - ? " Your manual changes were kept in History (steps_back=1 brings them back)." - : "" - const next = - describeVersions(otherVersions(session.id, entry.xml)) || "none" - return { - content: [ - { - type: "text", - text: `Restored the version ${stepsBack} step(s) back.${manualNote}\n\n${describePages(entry.xml)}\n\nVersions you can go to now: ${next}\n\nCall get_diagram before edit_diagram.`, - }, - ], - } - } catch (error) { - const message = - error instanceof Error ? error.message : String(error) - log.error("restore_version failed:", message) - return { - content: [{ type: "text", text: `Error: ${message}` }], - isError: true, - } + // The canvas, with the user's manual edits, is the state store + const current = sessionState(session.id)?.xml || session.xml + const candidates = otherVersions(session.id, current) + if (candidates.length === 0) { + return fail("No other version in History.") } + if (candidates.length < stepsBack) { + return fail( + `Only ${candidates.length} other version(s) in History: ${describeVersions(candidates)}`, + ) + } + const entry = candidates[stepsBack - 1] + + // Manual edits since the last entry are not in History yet; + // restoreHistoryEntry keeps them (a blank page excepted) + const hadManual = + contentFingerprint(current) !== contentFingerprint(BLANK_MXFILE) && + !isInHistory(session.id, current) + + if (restoreHistoryEntry(session.id, entry.id) === null) { + return fail("That version is no longer in History.") + } + session.xml = entry.xml + // As load_diagram: the model may never have seen this XML, so + // the edit gate requires one get_diagram before edits + session.lastSeenXml = "" + + log.info( + `Restored the version ${stepsBack} step(s) back (entry ${entry.id})`, + ) + + const manualNote = hadManual + ? " Your manual changes were kept in History (steps_back=1 brings them back)." + : "" + const next = + describeVersions(otherVersions(session.id, entry.xml)) || "none" + return ok( + `Restored the version ${stepsBack} step(s) back.${manualNote}\n\n${describePages(entry.xml)}\n\nVersions you can go to now: ${next}\n\nCall get_diagram before edit_diagram.`, + ) }, + { write: true, needsSession: true }, ) // Graceful shutdown handler diff --git a/packages/mcp-server/src/new-diagram.ts b/packages/mcp-server/src/new-diagram.ts index 68617ed7..8eef583c 100644 --- a/packages/mcp-server/src/new-diagram.ts +++ b/packages/mcp-server/src/new-diagram.ts @@ -121,21 +121,12 @@ export function truncatedCellError(input: string): string | null { return `The XML ends inside an unfinished mxCell, so the output was probably cut off. Nothing was drawn. It ended with:\n${input.slice(-300)}\nResend the complete diagram, or send the first cells now and add the remaining cells with edit_diagram add operations.` } -/** - * The named style definitions are taken out first (style-classes.ts). Bare - * cells then get the wrapper and root cells, since the strict parser rejects - * several top-level elements. Then the XML is validated and auto-fixed while - * it is still a bare model, where duplicate ids are renamed. Compact cells - * (compact-cells.ts) are expanded, the names expanded and the default styles - * added on the fixed XML, so repaired cells get them too, and finally it is - * turned into an . - */ /** * Take the named style definitions out of the model's XML (style-classes.ts). * Returns the error for the model when a definition never closed or when * nothing but definitions was sent. */ -export function takeStyleDefinitions(input: string): { +function takeStyleDefinitions(input: string): { classes: StyleClasses xml: string error: string | null @@ -152,10 +143,17 @@ export function takeStyleDefinitions(input: string): { return { classes, xml, error } } -export function prepareNewDiagram( - input: string, - page: { pageId?: string; pageName?: string } = {}, -): NewDiagram { +/** + * The model of one page from the XML the model wrote (also the starting + * content of add_page). The named style definitions are taken out first + * (style-classes.ts). Bare cells then get the wrapper and root cells, since + * the strict parser rejects several top-level elements. Then the XML is + * validated and auto-fixed while it is still a bare model, where duplicate + * ids are renamed. Compact cells (compact-cells.ts) are expanded, the names + * expanded and the default styles added on the fixed XML, so repaired cells + * get them too. + */ +export function prepareNewModel(input: string): NewDiagram { const { classes, xml: cells, @@ -182,12 +180,22 @@ export function prepareNewDiagram( error: `XML validation failed after expanding the cells - ${rewriteError}`, } } - const normalized = normalizeToMxfile(xml, page) + return { ok: true, xml, fixes } +} + +/** prepareNewModel, and the model turned into an */ +export function prepareNewDiagram( + input: string, + page: { pageId?: string; pageName?: string } = {}, +): NewDiagram { + const model = prepareNewModel(input) + if (!model.ok) return model + const normalized = normalizeToMxfile(model.xml, page) if (!normalized) { return { ok: false, error: "XML must be the mxCell elements of one page, a , or an with one or more children.", } } - return { ok: true, xml: normalized, fixes } + return { ok: true, xml: normalized, fixes: model.fixes } } diff --git a/packages/mcp-server/src/pages.ts b/packages/mcp-server/src/pages.ts index 1b83e55f..b7f5d70c 100644 --- a/packages/mcp-server/src/pages.ts +++ b/packages/mcp-server/src/pages.ts @@ -147,7 +147,7 @@ export function wrapCellsInModel(xml: string): string { */ export function normalizeToMxfile( xml: string, - opts: { pageId?: string; pageName?: string; host?: string } = {}, + opts: { pageId?: string; pageName?: string } = {}, ): string | null { const trimmed = xml.trim() if (!trimmed) return null @@ -156,9 +156,8 @@ export function normalizeToMxfile( const pageId = opts.pageId || generatePageId() const pageName = opts.pageName || "Page-1" - const host = opts.host || "app.diagrams.net" const inner = stripXmlDeclaration(trimmed) - return `${inner}` + return `${inner}` } /** diff --git a/packages/mcp-server/src/preview/preview.css b/packages/mcp-server/src/preview/preview.css index 15e85066..43af06b1 100644 --- a/packages/mcp-server/src/preview/preview.css +++ b/packages/mcp-server/src/preview/preview.css @@ -19,7 +19,6 @@ html, body { width: 100%; height: 100%; overflow: hidden; } display: flex; align-items: center; justify-content: center; overflow: hidden; } -#header .logo img { width: 20px; height: 20px; filter: brightness(0) invert(1); } #header .title { font-size: 15px; font-weight: 600; color: #1a1a2e; letter-spacing: -0.3px; diff --git a/packages/mcp-server/src/preview/preview.js b/packages/mcp-server/src/preview/preview.js index 6d79e132..2779b63d 100644 --- a/packages/mcp-server/src/preview/preview.js +++ b/packages/mcp-server/src/preview/preview.js @@ -328,17 +328,17 @@ async function poll() { // Plain export: capture whatever tab is currently displayed. // // Page-targeted export: the server sends a single-page - // projection in s.exportXml. We load it into the iframe, let + // projection in s.export.xml. 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; + if (s.export && !pendingMcpExport && isReady) { + pendingMcpExport = s.export.format; const seq = ++mcpExportSeq; - mcpExportId = s.exportId; - const extra = s.exportOptions || {}; + mcpExportId = s.export.id; + const extra = s.export.options || {}; const fireExport = () => { // mcpExport carries this export's number and is echoed // back in msg.message (see the handler). PNG: width @@ -349,10 +349,10 @@ async function poll() { : { action: 'export', format: pendingMcpExport, mcpExport: seq }; iframe.contentWindow.postMessage(JSON.stringify(exportOpts), '*'); }; - if (s.exportXml) { + if (s.export.xml) { projectionExportActive = true; // Load the projection without touching lastXml/server state. - iframe.contentWindow.postMessage(JSON.stringify({ action: 'load', xml: s.exportXml, autosave: 0 }), '*'); + iframe.contentWindow.postMessage(JSON.stringify({ action: 'load', xml: s.export.xml, autosave: 0 }), '*'); // Let draw.io render the loaded page before exporting // (same proven settle delay as the AI-preview path). setTimeout(fireExport, 600); diff --git a/packages/mcp-server/src/shape-library.ts b/packages/mcp-server/src/shape-library.ts index dc1d935c..a428dcfb 100644 --- a/packages/mcp-server/src/shape-library.ts +++ b/packages/mcp-server/src/shape-library.ts @@ -7,7 +7,7 @@ import { readFile } from "node:fs/promises" import { dirname, join, resolve } from "node:path" import { fileURLToPath } from "node:url" -export const SHAPE_LIBRARY_GROUPS: Record = { +const SHAPE_LIBRARY_GROUPS: Record = { Cloud: [ "aws4", "azure2", diff --git a/packages/mcp-server/src/style-classes.ts b/packages/mcp-server/src/style-classes.ts index 33681d43..6287bc2c 100644 --- a/packages/mcp-server/src/style-classes.ts +++ b/packages/mcp-server/src/style-classes.ts @@ -224,9 +224,3 @@ export function addDefaultStyles( return `${current}${separator}${added.join(";")};` }) } - -/** Expand the named styles of a model's XML and add the default styles */ -export function expandStyles(xml: string): string { - const { classes, xml: cells } = readStyleClasses(xml) - return addDefaultStyles(applyStyleClasses(cells, classes)) -} diff --git a/packages/mcp-server/src/xml-examples.ts b/packages/mcp-server/src/xml-examples.ts index e9af6ca4..0def58a2 100644 --- a/packages/mcp-server/src/xml-examples.ts +++ b/packages/mcp-server/src/xml-examples.ts @@ -30,6 +30,24 @@ export const STYLE_CLASS_EXAMPLE = ` ` +/** + * The style rules shared by the drawing guide and the web app's system + * prompt. They differ only in the tool whose call defines named styles + * ("create_new_diagram or add_page", "display_diagram") and in who expands + * them ("the server", "the app"). + */ +export function styleRules(tool: string, actor: string) { + const actorSentence = actor.charAt(0).toUpperCase() + actor.slice(1) + return { + /** Named styles, defined once and used like CSS classes */ + namedStyles: `Define each style used by several cells ONCE, as a named style before the cells, and use the name in the cells like a CSS class. A cell's style can combine a shape token, a name and overrides; later pairs win. Name only styles that two or more cells share; a style used by one cell stays inline. Names must not be draw.io's own style names: shapes such as text, ellipse, rhombus, swimlane, label, image, and colors such as blue, green, red, gray, yellow, orange, purple, pink. A definition applies to the call it is in: each ${tool} call defines the names it uses. ${actorSentence} expands the names, so the saved file is standard draw.io XML.`, + /** html=1 and whiteSpace=wrap are added automatically */ + noHtml: `NEVER write html=1 or whiteSpace=wrap: ${actor} adds html=1 to every cell and whiteSpace=wrap to shapes. Labels are HTML: use <br> for a line break and <b> for bold, never \\n; a literal < or > in a label is written &lt; or &gt;.`, + /** draw.io's own defaults are not written */ + noDefaults: `Do NOT repeat what draw.io already uses. For a plain shape: rounded=0, align=center, verticalAlign=middle, fontSize=12, strokeWidth=1, fillColor=#ffffff, strokeColor=#000000, fontColor=#000000. For an edge: endArrow=classic, strokeColor=#000000. Writing one of them is right only when it overrides what a name or the shape sets: an edge is rounded by default, so rounded=0 on an edge is a real setting, and a text cell is left/top aligned by default, so there align=center or verticalAlign=middle are real settings.`, + } +} + /** Indent every line, for an indented code block in Markdown */ export const indent = (text: string, prefix = " ") => text diff --git a/packages/mcp-server/src/xml-validation.ts b/packages/mcp-server/src/xml-validation.ts index cce08353..115d9db8 100644 --- a/packages/mcp-server/src/xml-validation.ts +++ b/packages/mcp-server/src/xml-validation.ts @@ -1,6 +1,6 @@ /** - * XML Validation and Auto-Fix for draw.io diagrams - * Copied from lib/utils.ts to avoid cross-package imports + * XML validation and auto-fix for draw.io diagrams, shared by the MCP server + * and the web app (contexts/diagram-context.tsx) */ import { readAttributes } from "./xml-attributes.ts" @@ -10,9 +10,6 @@ import { getXmlSyntaxError } from "./xml-syntax.ts" // Constants // ============================================================================ -/** Maximum XML size to process (1MB) - larger XMLs may cause performance issues */ -const MAX_XML_SIZE = 1_000_000 - /** Structural attributes that should not be duplicated in draw.io */ const STRUCTURAL_ATTRS = [ "edge", @@ -23,9 +20,6 @@ const STRUCTURAL_ATTRS = [ "connectable", ] -/** Valid XML entity names */ -const VALID_ENTITIES = new Set(["lt", "gt", "amp", "quot", "apos"]) - /** Element names draw.io understands (case-sensitive) */ const VALID_DRAWIO_TAGS = new Set([ "mxfile", @@ -261,77 +255,6 @@ function checkDuplicateIds(xml: string): string | null { return null } -/** Check for tag mismatches using parsed tags */ -function checkTagMismatches(xml: string): string | null { - const xmlWithoutComments = xml.replace(//g, "") - const tags = parseXmlTags(xmlWithoutComments) - const tagStack: string[] = [] - - for (const { tagName, isClosing, isSelfClosing } of tags) { - if (isClosing) { - if (tagStack.length === 0) { - return `Invalid XML: Closing tag without matching opening tag` - } - const expected = tagStack.pop() - if (expected?.toLowerCase() !== tagName.toLowerCase()) { - return `Invalid XML: Expected closing tag but found ` - } - } else if (!isSelfClosing) { - tagStack.push(tagName) - } - } - if (tagStack.length > 0) { - return `Invalid XML: Document has ${tagStack.length} unclosed tag(s): ${tagStack.join(", ")}` - } - return null -} - -/** Check for invalid character references */ -function checkCharacterReferences(xml: string): string | null { - const charRefPattern = /&#x?[^;]+;?/g - let charMatch - while ((charMatch = charRefPattern.exec(xml)) !== null) { - const ref = charMatch[0] - if (ref.startsWith("&#x")) { - if (!ref.endsWith(";")) { - return `Invalid XML: Missing semicolon after hex reference: ${ref}` - } - const hexDigits = ref.substring(3, ref.length - 1) - if (hexDigits.length === 0 || !/^[0-9a-fA-F]+$/.test(hexDigits)) { - return `Invalid XML: Invalid hex character reference: ${ref}` - } - } else if (ref.startsWith("&#")) { - if (!ref.endsWith(";")) { - return `Invalid XML: Missing semicolon after decimal reference: ${ref}` - } - const decDigits = ref.substring(2, ref.length - 1) - if (decDigits.length === 0 || !/^[0-9]+$/.test(decDigits)) { - return `Invalid XML: Invalid decimal character reference: ${ref}` - } - } - } - return null -} - -/** Check for invalid entity references */ -function checkEntityReferences(xml: string): string | null { - const xmlWithoutComments = xml.replace(//g, "") - const bareAmpPattern = /&(?!(?:lt|gt|amp|quot|apos|#))/g - if (bareAmpPattern.test(xmlWithoutComments)) { - return "Invalid XML: Found unescaped & character(s). Replace & with &" - } - const invalidEntityPattern = /&([a-zA-Z][a-zA-Z0-9]*);/g - let entityMatch - while ( - (entityMatch = invalidEntityPattern.exec(xmlWithoutComments)) !== null - ) { - if (!VALID_ENTITIES.has(entityMatch[1])) { - return `Invalid XML: Invalid entity reference: &${entityMatch[1]}; - use only valid XML entities (lt, gt, amp, quot, apos)` - } - } - return null -} - /** Check for nested mxCell tags using regex */ function checkNestedMxCells(xml: string): string | null { // Quoted values may hold ">", so the tag ends at the first ">" outside them @@ -427,13 +350,6 @@ export function validateMxCellStructure( xml: string, opts: { strict?: boolean } = {}, ): string | null { - // Size check for performance - if (xml.length > MAX_XML_SIZE) { - console.warn( - `[validateMxCellStructure] XML size (${xml.length}) exceeds ${MAX_XML_SIZE} bytes, may cause performance issues`, - ) - } - // 0. DOM-based checks. Syntax errors are caught by the strict check at // the end: linkedom's DOMParser never reports them. try { @@ -464,66 +380,24 @@ export function validateMxCellStructure( ) } - // 1. Check for CDATA wrapper (invalid at document root) - if (/^\s* from end" - } - - // 2. Check for duplicate structural attributes + // 1. Check for duplicate structural attributes const dupAttrError = checkDuplicateAttributes(xml) if (dupAttrError) { return dupAttrError } - // 3. Check for unescaped < in attribute values - const attrValuePattern = /=\s*"([^"]*)"/g - let attrValMatch - while ((attrValMatch = attrValuePattern.exec(xml)) !== null) { - const value = attrValMatch[1] - if (//g - let commentMatch - while ((commentMatch = commentPattern.exec(xml)) !== null) { - if (/--/.test(commentMatch[1])) { - return "Invalid XML: Comment contains -- (double hyphen) which is not allowed" - } - } - - // 8. Check for unescaped entity references and invalid entity names - const entityError = checkEntityReferences(xml) - if (entityError) { - return entityError - } - - // 9. Check for empty id attributes on mxCell + // 3. Check for empty id attributes on mxCell if (/]*\sid\s*=\s*["']\s*["'][^>]*>/g.test(xml)) { return "Invalid XML: Found mxCell element(s) with empty id attribute" } - // 10. Check for nested mxCell tags + // 4. Check for nested mxCell tags const nestedCellError = checkNestedMxCells(xml) if (nestedCellError) { return nestedCellError @@ -539,9 +413,10 @@ export function validateMxCellStructure( } } - // 11. Strict XML syntax check, run last so the checks above can give - // more specific messages. Catches what they miss, e.g. duplicate or - // unquoted attributes, which make draw.io refuse to load the diagram. + // 5. Strict XML syntax check (saxes), which reports every + // well-formedness error: unbalanced tags, bad character or entity + // references, an unescaped < or &, duplicate or unquoted attributes. + // Any of them makes draw.io refuse to load the diagram. const syntaxError = getXmlSyntaxError(xml) if (syntaxError) { return `Invalid XML: syntax error at ${syntaxError} Escape special characters in attribute values (< for <, & for &, " for "), quote every attribute value, and do not repeat an attribute.` diff --git a/packages/mcp-server/tests/diagram-operations.test.ts b/packages/mcp-server/tests/diagram-operations.test.ts index 97af0dcd..e43718ec 100644 --- a/packages/mcp-server/tests/diagram-operations.test.ts +++ b/packages/mcp-server/tests/diagram-operations.test.ts @@ -5,13 +5,7 @@ */ import { deflateRawSync } from "node:zlib" -import { beforeAll, describe, expect, it, vi } from "vitest" -import { installDomPolyfill } from "../src/dom.ts" - -beforeAll(() => { - installDomPolyfill() -}) - +import { describe, expect, it, vi } from "vitest" import { applyDiagramOperations } from "../src/diagram-operations.ts" const DOC = `` diff --git a/packages/mcp-server/tests/edit-diagram.test.ts b/packages/mcp-server/tests/edit-diagram.test.ts index 34368ef0..69dbf0a6 100644 --- a/packages/mcp-server/tests/edit-diagram.test.ts +++ b/packages/mcp-server/tests/edit-diagram.test.ts @@ -2,13 +2,7 @@ * Tests for the all-or-nothing edit_diagram core (src/edit-diagram.ts). */ -import { beforeAll, describe, expect, it } from "vitest" -import { installDomPolyfill } from "../src/dom.ts" - -beforeAll(() => { - installDomPolyfill() -}) - +import { describe, expect, it } from "vitest" import { editDiagram, targetPageXml } from "../src/edit-diagram.ts" import { validateMxCellStructure } from "../src/xml-validation.ts" diff --git a/packages/mcp-server/tests/edit-gate.test.ts b/packages/mcp-server/tests/edit-gate.test.ts index e5cf6a78..ca3e7e05 100644 --- a/packages/mcp-server/tests/edit-gate.test.ts +++ b/packages/mcp-server/tests/edit-gate.test.ts @@ -9,13 +9,7 @@ * reads as a user edit. */ -import { beforeAll, describe, expect, it } from "vitest" -import { installDomPolyfill } from "../src/dom.ts" - -beforeAll(() => { - installDomPolyfill() -}) - +import { describe, expect, it } from "vitest" import { checkEditGate, contentFingerprint, diff --git a/packages/mcp-server/tests/history.test.ts b/packages/mcp-server/tests/history.test.ts index e8be5391..36308c3e 100644 --- a/packages/mcp-server/tests/history.test.ts +++ b/packages/mcp-server/tests/history.test.ts @@ -2,8 +2,7 @@ * Tests for the diagram history: the versions restore_version can go to. */ -import { beforeAll, describe, expect, it } from "vitest" -import { installDomPolyfill } from "../src/dom.ts" +import { describe, expect, it } from "vitest" import { addHistory, clearHistory, @@ -12,11 +11,6 @@ import { otherVersions, } from "../src/history.ts" -beforeAll(() => { - // XML parsing for the content fingerprints - installDomPolyfill() -}) - const page = (cellId: string) => `` diff --git a/packages/mcp-server/tests/http-server.test.ts b/packages/mcp-server/tests/http-server.test.ts index d4cee359..7997d04b 100644 --- a/packages/mcp-server/tests/http-server.test.ts +++ b/packages/mcp-server/tests/http-server.test.ts @@ -8,7 +8,6 @@ import http from "node:http" import { afterAll, beforeAll, describe, expect, it } from "vitest" -import { installDomPolyfill } from "../src/dom.ts" import { addHistory, getHistory } from "../src/history.ts" import { drawioEmbedParams, @@ -28,8 +27,6 @@ import { let port = 0 beforeAll(async () => { - // XML parsing, as the server installs it at startup - installDomPolyfill() port = await startHttpServer(40000 + Math.floor(Math.random() * 10000)) }) @@ -330,15 +327,15 @@ describe("export requests", () => { const poll = JSON.parse( (await request(`/api/state?sessionId=${id}`)).body, ) - expect(poll.exportFormat).toBe("png") - expect(poll.exportOptions).toEqual({ width: 1000, pageId: "p2" }) + expect(poll.export.format).toBe("png") + expect(poll.export.options).toEqual({ width: 1000, pageId: "p2" }) await postJson("/api/state", { sessionId: id, exportData: "data:image/png;base64,AAAA", - exportId: poll.exportId, + exportId: poll.export.id, }) - expect(getState(id)?.exportOptions).toBeUndefined() + expect(getState(id)?.pendingExport).toBeUndefined() }) it("ignores a late result of an export that already timed out", async () => { @@ -353,10 +350,10 @@ describe("export requests", () => { await postJson("/api/state", { sessionId: id, exportData: "data:image/png;base64,LATE", - exportId: first.exportId, + exportId: first.export.id, }) expect(getState(id)?.exportData).toBeUndefined() - expect(getState(id)?.exportFormat).toBe("svg") + expect(getState(id)?.pendingExport?.format).toBe("svg") }) }) @@ -388,7 +385,7 @@ describe("a session state recreated after it was lost", () => { stateId, }) expect(res.status).toBe(409) - expect(JSON.parse(res.body).stateChanged).toBe(true) + expect(JSON.parse(res.body).error).toBe("Session was recreated") // The saved file was recovered first and is kept expect(getState("mcp-sid-restart")?.xml).toBe(SAVED) } diff --git a/packages/mcp-server/tests/multi-page.test.ts b/packages/mcp-server/tests/multi-page.test.ts index 5e9011ed..0d6baf65 100644 --- a/packages/mcp-server/tests/multi-page.test.ts +++ b/packages/mcp-server/tests/multi-page.test.ts @@ -15,15 +15,7 @@ * (diagram-operations.ts) — i.e. the layers underneath the MCP tool surface. */ -import { beforeAll, describe, expect, it } from "vitest" -import { installDomPolyfill } from "../src/dom.ts" - -// Install the DOM polyfill exactly as index.ts does at runtime — the -// helpers under test rely on it. -beforeAll(() => { - installDomPolyfill() -}) - +import { describe, expect, it } from "vitest" import { applyDiagramOperations } from "../src/diagram-operations.ts" import { addPageToDoc, diff --git a/packages/mcp-server/tests/new-diagram.test.ts b/packages/mcp-server/tests/new-diagram.test.ts index 9a70b44a..b2204460 100644 --- a/packages/mcp-server/tests/new-diagram.test.ts +++ b/packages/mcp-server/tests/new-diagram.test.ts @@ -4,13 +4,7 @@ * never sent. */ -import { beforeAll, describe, expect, it } from "vitest" -import { installDomPolyfill } from "../src/dom.ts" - -beforeAll(() => { - installDomPolyfill() -}) - +import { describe, expect, it } from "vitest" import { isMxCellXmlComplete, prepareNewDiagram, diff --git a/packages/mcp-server/tests/persistence.test.ts b/packages/mcp-server/tests/persistence.test.ts index 7c2882d4..6a86664e 100644 --- a/packages/mcp-server/tests/persistence.test.ts +++ b/packages/mcp-server/tests/persistence.test.ts @@ -15,14 +15,9 @@ import { } from "node:fs" import { homedir, tmpdir } from "node:os" import { join } from "node:path" -import { afterEach, beforeAll, describe, expect, it } from "vitest" -import { installDomPolyfill } from "../src/dom.ts" +import { afterEach, describe, expect, it } from "vitest" import { Autosaver, defaultDataDir } from "../src/persistence.ts" -beforeAll(() => { - installDomPolyfill() -}) - const DIAGRAM = `` const BLANK = `` const sleep = (ms: number) => new Promise((r) => setTimeout(r, ms)) diff --git a/packages/mcp-server/tests/server-wiring.test.ts b/packages/mcp-server/tests/server-wiring.test.ts index 18c139fe..0de5737c 100644 --- a/packages/mcp-server/tests/server-wiring.test.ts +++ b/packages/mcp-server/tests/server-wiring.test.ts @@ -1,7 +1,7 @@ /** * Server-wiring test: boot the actual MCP stdio server (from source via tsx) - * and drive it the way a real MCP client does — initialize handshake, - * tools/list — to catch registration/schema regressions that the unit tests + * and drive it with the SDK's own client — initialize handshake, tools/list, + * tools/call — to catch registration/schema regressions that the unit tests * (which import helpers directly) can't see. * * This replaces the old standalone tests/smoke.mjs, which spawned the BUILT @@ -13,11 +13,13 @@ * window via open(). The browser bridge is covered by the Playwright e2e suite. */ -import { type ChildProcessWithoutNullStreams, spawn } from "node:child_process" import { mkdtempSync, rmSync, writeFileSync } from "node:fs" import { tmpdir } from "node:os" import path from "node:path" import { fileURLToPath } from "node:url" +import { Client } from "@modelcontextprotocol/sdk/client/index.js" +import { StdioClientTransport } from "@modelcontextprotocol/sdk/client/stdio.js" +import type { CallToolResult } from "@modelcontextprotocol/sdk/types.js" import { afterAll, beforeAll, describe, expect, it } from "vitest" const __dirname = path.dirname(fileURLToPath(import.meta.url)) @@ -60,115 +62,69 @@ const MAX_DESCRIPTION = 2048 // developer's real ~/.next-ai-drawio/instructions.md must not leak in const CUSTOM_RULE = "Always use blue fill #dae8fc." -let proc: ChildProcessWithoutNullStreams -let stdoutBuf = "" -const pending = new Map< - number, - { resolve: (m: any) => void; reject: (e: Error) => void; timeout: any } ->() -let nextId = 1 -let initResp: any +let client: Client -function send(method: string, params: unknown, isNotification = false) { - const msg: Record = { jsonrpc: "2.0", method, params } - if (!isNotification) msg.id = nextId++ - proc.stdin.write(`${JSON.stringify(msg)}\n`) - if (isNotification) return Promise.resolve(undefined) - return new Promise((resolve, reject) => { - const id = msg.id as number - const timeout = setTimeout(() => { - pending.delete(id) - reject(new Error(`Timed out waiting for response to ${method}`)) - }, 15000) - pending.set(id, { resolve, reject, timeout }) - }) -} +const listTools = async () => (await client.listTools()).tools +const findTool = async (name: string) => + (await listTools()).find((t) => t.name === name) +const properties = async (name: string) => + (await findTool(name))?.inputSchema?.properties ?? {} + +/** A tool call's result; the SDK throws on a JSON-RPC error */ +const callTool = async (name: string, args: Record) => + (await client.callTool({ name, arguments: args })) as CallToolResult +const textOf = (result: CallToolResult) => + (result.content[0] as { text: string }).text beforeAll(async () => { writeFileSync(path.join(dataDir, "instructions.md"), CUSTOM_RULE) writeFileSync(path.join(dataDir, `${SAVED_ID}.drawio`), SAVED_DIAGRAM) - proc = spawn(tsxBin, [entry], { - stdio: ["pipe", "pipe", "pipe"], - env: { ...process.env, DRAWIO_DATA_DIR: dataDir }, - }) as ChildProcessWithoutNullStreams - - proc.stdout.on("data", (chunk: Buffer) => { - stdoutBuf += chunk.toString() - const lines = stdoutBuf.split("\n") - stdoutBuf = lines.pop() || "" - for (const line of lines) { - const trimmed = line.trim() - if (!trimmed) continue - let msg: any - try { - msg = JSON.parse(trimmed) - } catch { - // Non-JSON-RPC log line — ignore. - continue - } - const p = msg.id !== undefined ? pending.get(msg.id) : undefined - if (p) { - clearTimeout(p.timeout) - pending.delete(msg.id) - p.resolve(msg) - } - } - }) - - initResp = await send("initialize", { - protocolVersion: "2024-11-05", - capabilities: {}, - clientInfo: { name: "wiring-test", version: "0.0.0" }, - }) - expect(initResp.error, JSON.stringify(initResp.error)).toBeUndefined() - expect(initResp.result?.serverInfo?.name).toBeTruthy() - await send("notifications/initialized", {}, true) + client = new Client({ name: "wiring-test", version: "0.0.0" }) + await client.connect( + new StdioClientTransport({ + command: tsxBin, + args: [entry], + env: { ...process.env, DRAWIO_DATA_DIR: dataDir } as Record< + string, + string + >, + // The server's log lines + stderr: "ignore", + }), + ) + expect(client.getServerVersion()?.name).toBeTruthy() }, 30000) -afterAll(() => { - proc?.kill("SIGTERM") +afterAll(async () => { + await client?.close() rmSync(dataDir, { recursive: true, force: true }) }) describe("MCP server wiring", () => { it("registers all tools", async () => { - const resp = await send("tools/list", {}) - expect(resp.error, JSON.stringify(resp.error)).toBeUndefined() - const names: string[] = (resp.result?.tools ?? []).map( - (t: { name: string }) => t.name, - ) + const names = (await listTools()).map((t) => t.name) for (const expected of EXPECTED_TOOLS) { expect(names, `missing tool: ${expected}`).toContain(expected) } }) it("advertises page-selector params on edit_diagram", async () => { - const resp = await send("tools/list", {}) - const edit = resp.result.tools.find( - (t: { name: string }) => t.name === "edit_diagram", - ) - const props = edit?.inputSchema?.properties ?? {} + const props = await properties("edit_diagram") expect(props.page_id).toBeTruthy() expect(props.page_name).toBeTruthy() expect(props.page_index).toBeTruthy() }) it("advertises session_id on start_session", async () => { - const resp = await send("tools/list", {}) - const start = resp.result.tools.find( - (t: { name: string }) => t.name === "start_session", - ) + const start = await findTool("start_session") expect(start?.inputSchema?.properties?.session_id).toBeTruthy() expect(start?.inputSchema?.required ?? []).not.toContain("session_id") }) it("lists the saved diagrams with their pages, without a session", async () => { - const resp = await send("tools/call", { - name: "list_saved_diagrams", - arguments: {}, - }) - expect(resp.result.isError).toBeFalsy() - const text: string = resp.result.content[0].text + const result = await callTool("list_saved_diagrams", {}) + expect(result.isError).toBeFalsy() + const text = textOf(result) const line = text.split("\n").find((l) => l.startsWith(SAVED_ID)) expect(line, text).toBeTruthy() expect(line).toContain("pages: Flow (3 cells)") @@ -177,81 +133,59 @@ describe("MCP server wiring", () => { }) it("advertises screenshot on create_new_diagram and edit_diagram", async () => { - const resp = await send("tools/list", {}) for (const name of ["create_new_diagram", "edit_diagram"]) { - const tool = resp.result.tools.find( - (t: { name: string }) => t.name === name, - ) - const props = tool?.inputSchema?.properties ?? {} - expect(props.screenshot?.type, `${name} screenshot`).toBe("boolean") + const props = await properties(name) + const screenshot = props.screenshot as { type?: string } | undefined + expect(screenshot?.type, `${name} screenshot`).toBe("boolean") } }) it("advertises name/id/xml on add_page", async () => { - const resp = await send("tools/list", {}) - const addPage = resp.result.tools.find( - (t: { name: string }) => t.name === "add_page", - ) - const props = addPage?.inputSchema?.properties ?? {} + const props = await properties("add_page") expect(props.name).toBeTruthy() expect(props.id).toBeTruthy() expect(props.xml).toBeTruthy() }) it("advertises steps_back on restore_version", async () => { - const resp = await send("tools/list", {}) - const restore = resp.result.tools.find( - (t: { name: string }) => t.name === "restore_version", - ) - expect(restore?.inputSchema?.properties?.steps_back).toBeTruthy() + const props = await properties("restore_version") + expect(props.steps_back).toBeTruthy() }) it("refuses restore_version without a session", async () => { - const resp = await send("tools/call", { - name: "restore_version", - arguments: {}, - }) - expect(resp.result.isError).toBe(true) - expect(resp.result.content[0].text).toContain("start_session") + const result = await callTool("restore_version", {}) + expect(result.isError).toBe(true) + expect(textOf(result)).toContain("start_session") }) it("keeps every description and the instructions within the host limit", async () => { - const resp = await send("tools/list", {}) - for (const tool of resp.result.tools) { + for (const tool of await listTools()) { expect( - tool.description.length, + tool.description?.length, `${tool.name} description length`, ).toBeLessThanOrEqual(MAX_DESCRIPTION) } - const instructions: string = initResp.result.instructions + const instructions = client.getInstructions() ?? "" expect(instructions).toContain("start_session") expect(instructions.length).toBeLessThanOrEqual(MAX_DESCRIPTION) }) it("serves a shape library without a session", async () => { - const resp = await send("tools/call", { - name: "get_shape_library", - arguments: { library: "AWS4" }, - }) - expect(resp.result.isError).toBeFalsy() - expect(resp.result.content[0].text).toContain("mxgraph.aws4") + const result = await callTool("get_shape_library", { library: "AWS4" }) + expect(result.isError).toBeFalsy() + expect(textOf(result)).toContain("mxgraph.aws4") }) it("reports an unknown shape library with the available names", async () => { - const resp = await send("tools/call", { - name: "get_shape_library", - arguments: { library: "../secrets" }, + const result = await callTool("get_shape_library", { + library: "../secrets", }) - expect(resp.result.isError).toBe(true) - expect(resp.result.content[0].text).toContain("kubernetes") + expect(result.isError).toBe(true) + expect(textOf(result)).toContain("kubernetes") }) it("serves the drawing guide without a session", async () => { - const resp = await send("tools/call", { - name: "get_drawing_guide", - arguments: {}, - }) - const text: string = resp.result.content[0].text + const text = textOf(await callTool("get_drawing_guide", {})) expect(text).toContain("Edge routing rules") expect(text).toContain("do not describe the diagram") expect(text).toContain("clear the canvas") @@ -259,31 +193,21 @@ describe("MCP server wiring", () => { }) it("appends the user's instructions.md to the drawing guide", async () => { - const resp = await send("tools/call", { - name: "get_drawing_guide", - arguments: {}, - }) - const text: string = resp.result.content[0].text + const text = textOf(await callTool("get_drawing_guide", {})) expect( text.endsWith(`\n\n## Custom Instructions\n${CUSTOM_RULE}`), ).toBe(true) }) it("advertises topic on get_drawing_guide", async () => { - const resp = await send("tools/list", {}) - const guide = resp.result.tools.find( - (t: { name: string }) => t.name === "get_drawing_guide", - ) - expect(guide?.inputSchema?.properties?.topic).toBeTruthy() + const props = await properties("get_drawing_guide") + expect(props.topic).toBeTruthy() }) it("serves the table XML reference with topic=tables", async () => { - const resp = await send("tools/call", { - name: "get_drawing_guide", - arguments: { topic: "tables" }, - }) - expect(resp.result.isError).toBeFalsy() - const text: string = resp.result.content[0].text + const result = await callTool("get_drawing_guide", { topic: "tables" }) + expect(result.isError).toBeFalsy() + const text = textOf(result) expect(text).toContain("shape=table") expect(text).not.toContain("Edge routing rules") }) @@ -291,10 +215,7 @@ describe("MCP server wiring", () => { describe("load_diagram dual-source arguments", () => { it("advertises both optional 'path' and 'xml' sources", async () => { - const resp = await send("tools/list", {}) - const load = resp.result.tools.find( - (t: { name: string }) => t.name === "load_diagram", - ) + const load = await findTool("load_diagram") const props = load?.inputSchema?.properties ?? {} expect(props.path).toBeTruthy() expect(props.xml).toBeTruthy() @@ -307,32 +228,25 @@ describe("load_diagram dual-source arguments", () => { // Argument validation fires before the session check: no session // exists in this harness, so a both-args call must report the // mutual-exclusion error, not "No active session". - const resp = await send("tools/call", { - name: "load_diagram", - arguments: { path: "/tmp/x.drawio", xml: "" }, + const result = await callTool("load_diagram", { + path: "/tmp/x.drawio", + xml: "", }) - expect(resp.error, JSON.stringify(resp.error)).toBeUndefined() - expect(resp.result?.isError).toBe(true) - expect(resp.result?.content?.[0]?.text).toContain("not both") + expect(result.isError).toBe(true) + expect(textOf(result)).toContain("not both") }) it("rejects passing neither 'path' nor 'xml'", async () => { - const resp = await send("tools/call", { - name: "load_diagram", - arguments: {}, - }) - expect(resp.result?.isError).toBe(true) - expect(resp.result?.content?.[0]?.text).toContain("either 'path'") + const result = await callTool("load_diagram", {}) + expect(result.isError).toBe(true) + expect(textOf(result)).toContain("either 'path'") }) it("accepts 'xml' alone as a source (fails only on the missing session)", async () => { - const resp = await send("tools/call", { - name: "load_diagram", - arguments: { - xml: '', - }, + const result = await callTool("load_diagram", { + xml: '', }) - expect(resp.result?.isError).toBe(true) - expect(resp.result?.content?.[0]?.text).toContain("No active session") + expect(result.isError).toBe(true) + expect(textOf(result)).toContain("No active session") }) }) diff --git a/packages/mcp-server/tests/setup.ts b/packages/mcp-server/tests/setup.ts new file mode 100644 index 00000000..962b8bc1 --- /dev/null +++ b/packages/mcp-server/tests/setup.ts @@ -0,0 +1,5 @@ +// The DOMParser and XMLSerializer globals the XML helpers need, installed +// before every test file as index.ts does at startup +import { installDomPolyfill } from "../src/dom.ts" + +installDomPolyfill() diff --git a/packages/mcp-server/tests/wrap-cells.test.ts b/packages/mcp-server/tests/wrap-cells.test.ts index 176dc2a7..83cdb22f 100644 --- a/packages/mcp-server/tests/wrap-cells.test.ts +++ b/packages/mcp-server/tests/wrap-cells.test.ts @@ -3,13 +3,7 @@ * a page, like in the web app, and the server adds the wrapper and root cells. */ -import { beforeAll, describe, expect, it } from "vitest" -import { installDomPolyfill } from "../src/dom.ts" - -beforeAll(() => { - installDomPolyfill() -}) - +import { describe, expect, it } from "vitest" import { prepareNewDiagram, reservedIdError } from "../src/new-diagram.ts" import { hasCells, wrapCellsInModel } from "../src/pages.ts" import { validateAndFixXml } from "../src/xml-validation.ts" @@ -100,6 +94,22 @@ describe("hasCells", () => { hasCells(``), ).toBe(false) }) + + it("does not mistake ids that start with 0 or 1 for root cells", () => { + expect( + hasCells( + '', + ), + ).toBe(true) + }) + + it("counts a cell wrapped in a UserObject", () => { + expect( + hasCells( + '', + ), + ).toBe(true) + }) }) describe("reservedIdError", () => { diff --git a/packages/mcp-server/tests/xml-reference.test.ts b/packages/mcp-server/tests/xml-reference.test.ts index 6d867d79..f707aba4 100644 --- a/packages/mcp-server/tests/xml-reference.test.ts +++ b/packages/mcp-server/tests/xml-reference.test.ts @@ -4,13 +4,7 @@ * whether draw.io renders it as intended is checked by hand. */ -import { beforeAll, describe, expect, it } from "vitest" -import { installDomPolyfill } from "../src/dom.ts" - -beforeAll(() => { - installDomPolyfill() -}) - +import { describe, expect, it } from "vitest" import { prepareNewDiagram } from "../src/new-diagram.ts" import { hasCells } from "../src/pages.ts" import { XML_REFERENCE } from "../src/xml-reference.ts" diff --git a/packages/mcp-server/tests/xml-validation.test.ts b/packages/mcp-server/tests/xml-validation.test.ts index 4f246f67..87b32376 100644 --- a/packages/mcp-server/tests/xml-validation.test.ts +++ b/packages/mcp-server/tests/xml-validation.test.ts @@ -7,15 +7,9 @@ * leave valid parts of the document untouched. */ -import { beforeAll, describe, expect, it } from "vitest" -import { installDomPolyfill } from "../src/dom.ts" -import { getXmlSyntaxError } from "../src/xml-syntax.ts" - -beforeAll(() => { - installDomPolyfill() -}) - +import { describe, expect, it } from "vitest" import { addPageToDoc, parseMxfile, serializeMxfile } from "../src/pages.ts" +import { getXmlSyntaxError } from "../src/xml-syntax.ts" import { validateAndFixXml } from "../src/xml-validation.ts" /** Bare model with the root cells plus the given cells. */ diff --git a/packages/mcp-server/tsconfig.json b/packages/mcp-server/tsconfig.json index da1ec74c..ce2b0a4e 100644 --- a/packages/mcp-server/tsconfig.json +++ b/packages/mcp-server/tsconfig.json @@ -9,8 +9,6 @@ "esModuleInterop": true, "skipLibCheck": true, "forceConsistentCasingInFileNames": true, - "declaration": true, - "declarationMap": true, "sourceMap": true, "resolveJsonModule": true, "rewriteRelativeImportExtensions": true, diff --git a/packages/mcp-server/vitest.config.ts b/packages/mcp-server/vitest.config.ts index 3d3c44ba..a5a49d3f 100644 --- a/packages/mcp-server/vitest.config.ts +++ b/packages/mcp-server/vitest.config.ts @@ -4,6 +4,7 @@ export default defineConfig({ test: { include: ["tests/**/*.test.ts"], environment: "node", + setupFiles: ["tests/setup.ts"], // The package source uses Node16 module resolution with explicit .js // extensions in imports. Vitest+esbuild handles the .ts→.js mapping // transparently, so no extra alias config is needed. diff --git a/tests/unit/compact-cells.test.ts b/tests/unit/compact-cells.test.ts index 3aeb761d..95652c60 100644 --- a/tests/unit/compact-cells.test.ts +++ b/tests/unit/compact-cells.test.ts @@ -299,8 +299,9 @@ describe("review round: cells the model may write", () => { expect(outcome.xml).toContain( '', + // Folding keeps a parent other than the default layer "1" + expect(foldCells(file)).toContain( + '', ) }) diff --git a/tests/unit/mcp-core.test.ts b/tests/unit/mcp-core.test.ts index 3636ac4e..a5fc907b 100644 --- a/tests/unit/mcp-core.test.ts +++ b/tests/unit/mcp-core.test.ts @@ -14,11 +14,7 @@ import { wrapCellsInModel, } from "@/packages/mcp-server/src/pages.ts" import { getXmlSyntaxError } from "@/packages/mcp-server/src/xml-syntax.ts" -import { - autoFixXml, - validateAndFixXml, - validateMxCellStructure, -} from "@/packages/mcp-server/src/xml-validation.ts" +import { validateAndFixXml } from "@/packages/mcp-server/src/xml-validation.ts" const box = (id: string, parent = "1") => `` @@ -93,177 +89,6 @@ describe("MCP diagram modules with a browser DOM", () => { }) }) -// Repair cases fixed in the web app's own copy before it moved here -const page = (id: string, cells: string) => - `${cells}` - -describe("duplicate ids in multi-page documents", () => { - const shape = (id: string, value = "Box") => - `` - - it("accepts the same ids on different pages", () => { - const xml = `${page("p1", shape("2"))}${page("p2", shape("2"))}` - expect(validateMxCellStructure(xml)).toBeNull() - }) - - it("still reports duplicate ids within one page", () => { - const xml = `${page("p1", shape("2") + shape("2"))}${page("p2", "")}` - expect(validateMxCellStructure(xml)).toMatch(/duplicate cell ID/i) - }) - - it("does not rename the root cells of other pages when fixing", () => { - const xml = `${page("p1", shape("2", "R&D"))}${page("p2", shape("3"))}` - const result = validateAndFixXml(xml) - expect(result.valid).toBe(true) - expect(result.fixed).not.toContain("_dup") - expect(result.fixed).toContain("R&D") - }) - - it("renames a duplicate id in a bare model, as display_diagram has", () => { - // In an the duplicate is reported instead (above) - const xml = `${shape("d") + shape("d")}` - const { fixed } = autoFixXml(xml) - expect(fixed).toContain(' { - it("does not insert a space at the start of style values", () => { - const xml = `` - const { fixed } = autoFixXml(xml) - expect(fixed).toContain('style="rounded=1;whiteSpace=wrap;"') - }) - - it("adds a missing space between attributes", () => { - const xml = `` - expect(autoFixXml(xml).fixed).toContain('vertex="1" parent="1"') - }) - - it("keeps " inside rich text labels", () => { - const label = "<font color="#ff0000">Hello</font>" - const xml = `` - const result = validateAndFixXml(xml) - expect(result.valid).toBe(true) - expect(result.fixed).toContain(`value="${label}"`) - }) - - it("fixes an attribute delimited by "", () => { - const xml = `` - expect(autoFixXml(xml).fixed).toContain('dashPattern="1 1;"') - }) - - it("keeps cells written on one line next to multi-line cells", () => { - const xml = ` - - - - - - - - - -` - const result = validateAndFixXml(xml) - expect(result.valid).toBe(true) - for (const id of ["2", "e1", "3"]) { - expect(result.fixed).toContain(` { - const xml = `` - const result = validateAndFixXml(xml) - expect(result.valid).toBe(true) - expect(result.fixed).toContain(' { - it("returns true for empty diagram", () => { - const xml = '' - expect(hasCells(xml)).toBe(false) - }) - - it("returns false for diagram with content", () => { - const xml = - '' - expect(hasCells(xml)).toBe(true) - }) - - it("handles whitespace correctly", () => { - const xml = ' ' - expect(hasCells(xml)).toBe(false) - }) - - it("returns false for a shape drawn in draw.io with a random id", () => { - const xml = - '' - expect(hasCells(xml)).toBe(true) - }) - - it("does not mistake ids that start with 0 or 1 for root cells", () => { - const xml = - '' - expect(hasCells(xml)).toBe(true) - }) - - it("counts a cell wrapped in a UserObject", () => { - const xml = - '' - expect(hasCells(xml)).toBe(true) - }) -}) - -describe("applyDiagramOperations with wrapped cells", () => { - const xml = `` - - it("deletes a wrapped cell and its edges", () => { - const { result, errors } = applyDiagramOperations(xml, [ - { operation: "delete", cell_id: "5" }, - { operation: "delete", cell_id: "e1" }, - ]) - expect(errors).toEqual([]) - expect(result).not.toContain("UserObject") - expect(result).not.toContain('id="e1"') - expect(result).toContain('id="6"') - }) - - it("rejects adding a cell with the id of a wrapped cell", () => { - const { errors } = applyDiagramOperations(xml, [ - { - operation: "add", - cell_id: "5", - new_xml: '', - }, - ]) - expect(errors[0]?.message).toContain("already exists") - }) - - it("updates a wrapped cell", () => { - const { result, errors } = applyDiagramOperations(xml, [ - { - operation: "update", - cell_id: "5", - new_xml: - '', - }, - ]) - expect(errors).toEqual([]) - expect(result).toContain('label="New"') - expect(result).not.toContain('label="Docs"') - }) - - it("reports deleting a cell that does not exist", () => { - const { errors } = applyDiagramOperations(xml, [ - { operation: "delete", cell_id: "missing" }, - ]) - expect(errors[0]?.message).toContain("not found") - }) -}) - describe("edit_diagram on one page of a multi-page file", () => { // Page A has an edge "g"; on page B, "g" is a shape on layer "L" const twoPages = `${box("s")}${box("t")}${box("g", "L")}` diff --git a/tests/unit/mcp-preview-recovery.test.ts b/tests/unit/mcp-preview-recovery.test.ts index dbb28f91..2b11f0ce 100644 --- a/tests/unit/mcp-preview-recovery.test.ts +++ b/tests/unit/mcp-preview-recovery.test.ts @@ -445,7 +445,7 @@ describe("MCP preview History and a recreated session", () => { } const historyList = { status: 200, - body: { entries: [{ id: 5, index: 0, svg: "" }], count: 1 }, + body: { entries: [{ id: 5, index: 0, svg: "" }] }, } const modalOpen = () => document.getElementById("history-modal")?.classList.contains("open") diff --git a/tests/unit/style-classes.test.ts b/tests/unit/style-classes.test.ts index d71d244f..b0773f23 100644 --- a/tests/unit/style-classes.test.ts +++ b/tests/unit/style-classes.test.ts @@ -9,10 +9,15 @@ import { import { addDefaultStyles, applyStyleClasses, - expandStyles, readStyleClasses, } from "@/packages/mcp-server/src/style-classes.ts" +/** The named styles of a model's XML expanded, with the default styles */ +const expandStyles = (xml: string) => { + const { classes, xml: cells } = readStyleClasses(xml) + return addDefaultStyles(applyStyleClasses(cells, classes)) +} + const geometry = '' const cell = (id: string, style: string, extra = "") =>