fix(mcp): capture SVG for AI-generated diagrams

- Sync browser state before saving history in display_diagram
- Save AI result to history (in addition to state before)
- Add SVG capture after browser loads AI diagrams
- Add /api/history-svg endpoint to update last entry's SVG
- Add updateLastHistorySvg() function to history module

.mcp.json

@@ -1,10 +0,0 @@
-{
-    "mcpServers": {
-        "drawio": {
-            "command": "node",
-            "args": [
-                "/Users/jiangdy/Documents/programming/next-ai-draw-io/packages/mcp-server/dist/index.js"
-            ]
-        }
-    }
-}

packages/mcp-server/src/history.ts

@@ -49,3 +49,14 @@ export function getHistoryEntry(
 export function clearHistory(sessionId: string): void {
     historyStore.delete(sessionId)
 }
+
+export function updateLastHistorySvg(sessionId: string, svg: string): boolean {
+    const history = historyStore.get(sessionId)
+    if (!history || history.length === 0) return false
+    const last = history[history.length - 1]
+    if (!last.svg) {
+        last.svg = svg
+        return true
+    }
+    return false
+}

packages/mcp-server/src/http-server.ts

@@ -9,6 +9,7 @@ import {
     clearHistory,
     getHistory,
     getHistoryEntry,
+    updateLastHistorySvg,
 } from "./history.js"
 import { log } from "./logger.js"
 
@@ -16,6 +17,7 @@ interface SessionState {
     xml: string
     version: number
     lastUpdated: Date
+    svg?: string // Cached SVG from last browser save
 }
 
 export const stateStore = new Map<string, SessionState>()
@@ -29,13 +31,14 @@ export function getState(sessionId: string): SessionState | undefined {
     return stateStore.get(sessionId)
 }
 
-export function setState(sessionId: string, xml: string): number {
+export function setState(sessionId: string, xml: string, svg?: string): number {
     const existing = stateStore.get(sessionId)
     const newVersion = (existing?.version || 0) + 1
     stateStore.set(sessionId, {
         xml,
         version: newVersion,
         lastUpdated: new Date(),
+        svg: svg || existing?.svg, // Preserve cached SVG if not provided
     })
     log.debug(`State updated: session=${sessionId}, version=${newVersion}`)
     return newVersion
@@ -128,6 +131,8 @@ function handleRequest(
         handleHistoryApi(req, res, url)
     } else if (url.pathname === "/api/restore") {
         handleRestoreApi(req, res)
+    } else if (url.pathname === "/api/history-svg") {
+        handleHistorySvgApi(req, res)
     } else {
         res.writeHead(404)
         res.end("Not Found")
@@ -161,13 +166,13 @@ function handleStateApi(
         })
         req.on("end", () => {
             try {
-                const { sessionId, xml } = JSON.parse(body)
+                const { sessionId, xml, svg } = JSON.parse(body)
                 if (!sessionId) {
                     res.writeHead(400, { "Content-Type": "application/json" })
                     res.end(JSON.stringify({ error: "sessionId required" }))
                     return
                 }
-                const version = setState(sessionId, xml)
+                const version = setState(sessionId, xml, svg)
                 res.writeHead(200, { "Content-Type": "application/json" })
                 res.end(JSON.stringify({ success: true, version }))
             } catch {
@@ -255,6 +260,39 @@ function handleRestoreApi(
     })
 }
 
+function handleHistorySvgApi(
+    req: http.IncomingMessage,
+    res: http.ServerResponse,
+): void {
+    if (req.method !== "POST") {
+        res.writeHead(405)
+        res.end("Method Not Allowed")
+        return
+    }
+
+    let body = ""
+    req.on("data", (chunk) => {
+        body += chunk
+    })
+    req.on("end", () => {
+        try {
+            const { sessionId, svg } = JSON.parse(body)
+            if (!sessionId || !svg) {
+                res.writeHead(400, { "Content-Type": "application/json" })
+                res.end(JSON.stringify({ error: "sessionId and svg required" }))
+                return
+            }
+
+            updateLastHistorySvg(sessionId, svg)
+            res.writeHead(200, { "Content-Type": "application/json" })
+            res.end(JSON.stringify({ success: true }))
+        } catch {
+            res.writeHead(400, { "Content-Type": "application/json" })
+            res.end(JSON.stringify({ error: "Invalid JSON" }))
+        }
+    })
+}
+
 function getHtmlPage(sessionId: string): string {
     return `<!DOCTYPE html>
 <html lang="en">
@@ -358,6 +396,8 @@ function getHtmlPage(sessionId: string): string {
         const iframe = document.getElementById('drawio');
         const statusEl = document.getElementById('status');
         let currentVersion = 0, isReady = false, pendingXml = null, lastXml = null;
+        let pendingSvgExport = null;
+        let pendingAiSvg = false;
 
         window.addEventListener('message', (e) => {
             if (e.origin !== 'https://embed.diagrams.net') return;
@@ -369,24 +409,49 @@ function getHtmlPage(sessionId: string): string {
                     statusEl.className = 'status connected';
                     if (pendingXml) { loadDiagram(pendingXml); pendingXml = null; }
                 } else if ((msg.event === 'save' || msg.event === 'autosave') && msg.xml && msg.xml !== lastXml) {
-                    pushState(msg.xml);
+                    // Request SVG export, then push state with SVG
+                    pendingSvgExport = msg.xml;
+                    iframe.contentWindow.postMessage(JSON.stringify({ action: 'export', format: 'svg' }), '*');
+                    // Fallback if export doesn't respond
+                    setTimeout(() => { if (pendingSvgExport === msg.xml) { pushState(msg.xml, ''); pendingSvgExport = null; } }, 2000);
+                } else if (msg.event === 'export' && msg.data) {
+                    let svg = msg.data;
+                    if (!svg.startsWith('data:')) svg = 'data:image/svg+xml;base64,' + btoa(unescape(encodeURIComponent(svg)));
+                    if (pendingSvgExport) {
+                        const xml = pendingSvgExport;
+                        pendingSvgExport = null;
+                        pushState(xml, svg);
+                    } else if (pendingAiSvg) {
+                        pendingAiSvg = false;
+                        fetch('/api/history-svg', {
+                            method: 'POST',
+                            headers: { 'Content-Type': 'application/json' },
+                            body: JSON.stringify({ sessionId, svg })
+                        }).catch(() => {});
+                    }
                 }
             } catch {}
         });
 
-        function loadDiagram(xml) {
+        function loadDiagram(xml, capturePreview = false) {
             if (!isReady) { pendingXml = xml; return; }
             lastXml = xml;
             iframe.contentWindow.postMessage(JSON.stringify({ action: 'load', xml, autosave: 1 }), '*');
+            if (capturePreview) {
+                setTimeout(() => {
+                    pendingAiSvg = true;
+                    iframe.contentWindow.postMessage(JSON.stringify({ action: 'export', format: 'svg' }), '*');
+                }, 500);
+            }
         }
 
-        async function pushState(xml) {
+        async function pushState(xml, svg = '') {
             if (!sessionId) return;
             try {
                 const r = await fetch('/api/state', {
                     method: 'POST',
                     headers: { 'Content-Type': 'application/json' },
-                    body: JSON.stringify({ sessionId, xml })
+                    body: JSON.stringify({ sessionId, xml, svg })
                 });
                 if (r.ok) { const d = await r.json(); currentVersion = d.version; lastXml = xml; }
             } catch (e) { console.error('Push failed:', e); }
@@ -400,7 +465,7 @@ function getHtmlPage(sessionId: string): string {
                 const s = await r.json();
                 if (s.version > currentVersion && s.xml) {
                     currentVersion = s.version;
-                    loadDiagram(s.xml);
+                    loadDiagram(s.xml, true);
                 }
             } catch {}
         }

packages/mcp-server/src/index.ts

@@ -48,22 +48,26 @@ server.prompt(
                 role: "user",
                 content: {
                     type: "text",
-                    text: `# Draw.io Diagram Workflow
+                    text: `# Draw.io Diagram Workflow Guidelines
 
 ## Creating a New Diagram
-1. Call start_session to open browser preview
+1. Call start_session to open the browser preview
-2. Use display_diagram with mxGraphModel XML
+2. Use display_diagram with complete mxGraphModel XML to create a new diagram
 
-## Adding Elements
+## Adding Elements to Existing Diagram
-Use edit_diagram with "add" operation - provide cell_id and new_xml
+1. Use edit_diagram with "add" operation
+2. Provide a unique cell_id and complete mxCell XML
+3. No need to call get_diagram first - the server fetches latest state automatically
 
-## Modifying/Deleting Elements
+## Modifying or Deleting Existing Elements
-1. Call get_diagram to see current cell IDs
+1. FIRST call get_diagram to see current cell IDs and structure
-2. Use edit_diagram with "update" or "delete" operations
+2. THEN call edit_diagram with "update" or "delete" operations
+3. For update, provide the cell_id and complete new mxCell XML
 
-## Notes
+## Important Notes
-- display_diagram REPLACES entire diagram
+- display_diagram REPLACES the entire diagram - only use for new diagrams
-- edit_diagram PRESERVES user's manual changes`,
+- edit_diagram PRESERVES user's manual changes (fetches browser state first)
+- Always use unique cell_ids when adding elements (e.g., "shape-1", "arrow-2")`,
                 },
             },
         ],
@@ -75,7 +79,9 @@ server.registerTool(
     "start_session",
     {
         description:
-            "Start a new diagram session and open browser for real-time preview.",
+            "Start a new diagram session and open the browser for real-time preview. " +
+            "Starts an embedded server and opens a browser window with draw.io. " +
+            "The browser will show diagram updates as they happen.",
         inputSchema: {},
     },
     async () => {
@@ -93,7 +99,7 @@ server.registerTool(
                 content: [
                     {
                         type: "text",
-                        text: `Session started!\n\nSession ID: ${sessionId}\nBrowser: ${browserUrl}`,
+                        text: `Session started successfully!\n\nSession ID: ${sessionId}\nBrowser URL: ${browserUrl}\n\nThe browser will now show real-time diagram updates.`,
                     },
                 ],
             }
@@ -115,7 +121,9 @@ server.registerTool(
     {
         description:
             "Display a NEW draw.io diagram from XML. REPLACES the entire diagram. " +
-            "Use edit_diagram to add elements to existing diagram.",
+            "Use this for creating new diagrams from scratch. " +
+            "To ADD elements to an existing diagram, use edit_diagram with 'add' operation instead. " +
+            "You should generate valid draw.io/mxGraph XML format.",
         inputSchema: {
             xml: z
                 .string()
@@ -129,7 +137,7 @@ server.registerTool(
                     content: [
                         {
                             type: "text",
-                            text: "Error: No active session. Call start_session first.",
+                            text: "Error: No active session. Please call start_session first.",
                         },
                     ],
                     isError: true,
@@ -154,16 +162,26 @@ server.registerTool(
                 }
             }
 
-            // Save current state before replacing
+            // Sync from browser state first
+            const browserState = getState(currentSession.id)
+            if (browserState?.xml) {
+                currentSession.xml = browserState.xml
+            }
+
+            // Save user's state before AI overwrites (with cached SVG)
             if (currentSession.xml) {
-                addHistory(currentSession.id, currentSession.xml)
+                addHistory(
+                    currentSession.id,
+                    currentSession.xml,
+                    browserState?.svg || "",
+                )
             }
 
             currentSession.xml = xml
             setState(currentSession.id, xml)
 
-            // Save new state
+            // Save AI result (no SVG yet)
-            addHistory(currentSession.id, xml)
+            addHistory(currentSession.id, xml, "")
 
             log.info(`Displayed diagram, ${xml.length} chars`)
 
@@ -217,7 +235,7 @@ server.registerTool(
                     content: [
                         {
                             type: "text",
-                            text: "Error: No active session. Call start_session first.",
+                            text: "Error: No active session. Please call start_session first.",
                         },
                     ],
                     isError: true,
@@ -242,8 +260,12 @@ server.registerTool(
                 }
             }
 
-            // Save before editing
+            // Save before editing (with cached SVG from browser)
-            addHistory(currentSession.id, currentSession.xml)
+            addHistory(
+                currentSession.id,
+                currentSession.xml,
+                browserState?.svg || "",
+            )
 
             // Validate operations
             const validatedOps = operations.map((op) => {
@@ -267,8 +289,8 @@ server.registerTool(
             currentSession.xml = result
             setState(currentSession.id, result)
 
-            // Save after editing
+            // Save AI result (no SVG yet)
-            addHistory(currentSession.id, result)
+            addHistory(currentSession.id, result, "")
 
             log.info(`Edited diagram: ${operations.length} operation(s)`)
 
@@ -306,7 +328,7 @@ server.registerTool(
                     content: [
                         {
                             type: "text",
-                            text: "Error: No active session. Call start_session first.",
+                            text: "Error: No active session. Please call start_session first.",
                         },
                     ],
                     isError: true,
@@ -364,7 +386,7 @@ server.registerTool(
                     content: [
                         {
                             type: "text",
-                            text: "Error: No active session. Call start_session first.",
+                            text: "Error: No active session. Please call start_session first.",
                         },
                     ],
                     isError: true,