fix(mcp): sync browser state before get_diagram to prevent data loss (#342)

* fix(mcp): sync browser state before get_diagram to prevent data loss

- Add syncRequested flag to SessionState for browser sync coordination
- Add requestSync() and waitForSync() functions to http-server
- Browser polls for syncRequested flag and immediately pushes current state
- get_diagram now syncs fresh state from browser before returning
- edit_diagram requires get_diagram to be called within 30s to prevent stale edits
- Updated edit_diagram description to enforce workflow

* fix(mcp): make lastGetDiagramTime session-scoped and handle missing session in requestSync

- Move lastGetDiagramTime into currentSession object to prevent cross-session issues
- requestSync now returns boolean indicating if request was made
- Only wait for sync if session exists (avoids false-positive from undefined state)

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

@@ -18,6 +18,7 @@ interface SessionState {
     version: number
     lastUpdated: Date
     svg?: string // Cached SVG from last browser save
+    syncRequested?: number // Timestamp when sync requested, cleared when browser responds
 }
 
 export const stateStore = new Map<string, SessionState>()
@@ -39,11 +40,37 @@ export function setState(sessionId: string, xml: string, svg?: string): number {
         version: newVersion,
         lastUpdated: new Date(),
         svg: svg || existing?.svg, // Preserve cached SVG if not provided
+        syncRequested: undefined, // Clear sync request when browser pushes state
     })
     log.debug(`State updated: session=${sessionId}, version=${newVersion}`)
     return newVersion
 }
 
+export function requestSync(sessionId: string): boolean {
+    const state = stateStore.get(sessionId)
+    if (state) {
+        state.syncRequested = Date.now()
+        log.debug(`Sync requested for session=${sessionId}`)
+        return true
+    }
+    log.debug(`Sync requested for non-existent session=${sessionId}`)
+    return false
+}
+
+export async function waitForSync(
+    sessionId: string,
+    timeoutMs = 3000,
+): Promise<boolean> {
+    const start = Date.now()
+    while (Date.now() - start < timeoutMs) {
+        const state = stateStore.get(sessionId)
+        if (!state?.syncRequested) return true // Sync completed
+        await new Promise((r) => setTimeout(r, 100))
+    }
+    log.warn(`Sync timeout for session=${sessionId}`)
+    return false // Timeout
+}
+
 export function startHttpServer(port = 6002): Promise<number> {
     return new Promise((resolve, reject) => {
         if (server) {
@@ -157,6 +184,7 @@ function handleStateApi(
             JSON.stringify({
                 xml: state?.xml || null,
                 version: state?.version || 0,
+                syncRequested: !!state?.syncRequested,
             }),
         )
     } else if (req.method === "POST") {
@@ -415,6 +443,13 @@ function getHtmlPage(sessionId: string): string {
                     // Fallback if export doesn't respond
                     setTimeout(() => { if (pendingSvgExport === msg.xml) { pushState(msg.xml, ''); pendingSvgExport = null; } }, 2000);
                 } else if (msg.event === 'export' && msg.data) {
+                    // Handle sync export (XML format) - server requested fresh state
+                    if (pendingSyncExport && !msg.data.startsWith('data:') && !msg.data.startsWith('<svg')) {
+                        pendingSyncExport = false;
+                        pushState(msg.data, '');
+                        return;
+                    }
+                    // Handle SVG export
                     let svg = msg.data;
                     if (!svg.startsWith('data:')) svg = 'data:image/svg+xml;base64,' + btoa(unescape(encodeURIComponent(svg)));
                     if (pendingSvgExport) {
@@ -457,12 +492,20 @@ function getHtmlPage(sessionId: string): string {
             } catch (e) { console.error('Push failed:', e); }
         }
 
+        let pendingSyncExport = false;
+
         async function poll() {
             if (!sessionId) return;
             try {
                 const r = await fetch('/api/state?sessionId=' + encodeURIComponent(sessionId));
                 if (!r.ok) return;
                 const s = await r.json();
+                // Handle sync request - server needs fresh state
+                if (s.syncRequested && !pendingSyncExport) {
+                    pendingSyncExport = true;
+                    iframe.contentWindow.postMessage(JSON.stringify({ action: 'export', format: 'xml' }), '*');
+                }
+                // Load new diagram from server
                 if (s.version > currentVersion && s.xml) {
                     currentVersion = s.version;
                     loadDiagram(s.xml, true);

packages/mcp-server/src/index.ts

@@ -35,7 +35,13 @@ import {
     type DiagramOperation,
 } from "./diagram-operations.js"
 import { addHistory } from "./history.js"
-import { getState, setState, startHttpServer } from "./http-server.js"
+import {
+    getState,
+    requestSync,
+    setState,
+    startHttpServer,
+    waitForSync,
+} from "./http-server.js"
 import { log } from "./logger.js"
 import { validateAndFixXml } from "./xml-validation.js"
 
@@ -49,6 +55,7 @@ let currentSession: {
     id: string
     xml: string
     version: number
+    lastGetDiagramTime: number // Track when get_diagram was last called (for enforcing workflow)
 } | null = null
 
 // Create MCP server
@@ -114,6 +121,7 @@ server.registerTool(
                 id: sessionId,
                 xml: "",
                 version: 0,
+                lastGetDiagramTime: 0,
             }
 
             // Open browser
@@ -245,11 +253,14 @@ server.registerTool(
     "edit_diagram",
     {
         description:
-            "Edit the current diagram by ID-based operations (update/add/delete cells). " +
-            "ALWAYS fetches the latest state from browser first, so user's manual changes are preserved.\n\n" +
-            "IMPORTANT workflow:\n" +
-            "- For ADD operations: Can use directly - just provide new unique cell_id and new_xml.\n" +
-            "- For UPDATE/DELETE: Call get_diagram FIRST to see current cell IDs, then edit.\n\n" +
+            "Edit the current diagram by ID-based operations (update/add/delete cells).\n\n" +
+            "⚠️ REQUIRED: You MUST call get_diagram BEFORE this tool!\n" +
+            "This fetches the latest state from the browser including any manual user edits.\n" +
+            "Skipping get_diagram WILL cause user's changes to be LOST.\n\n" +
+            "Workflow:\n" +
+            "1. Call get_diagram to see current cell IDs and structure\n" +
+            "2. Use the returned XML to construct your edit operations\n" +
+            "3. Call edit_diagram with your operations\n\n" +
             "Operations:\n" +
             "- add: Add a new cell. Provide cell_id (new unique id) and new_xml.\n" +
             "- update: Replace an existing cell by its id. Provide cell_id and complete new_xml.\n" +
@@ -288,6 +299,27 @@ server.registerTool(
                 }
             }
 
+            // Enforce workflow: require get_diagram to be called first
+            const timeSinceGet = Date.now() - currentSession.lastGetDiagramTime
+            if (timeSinceGet > 30000) {
+                // 30 seconds
+                log.warn(
+                    "edit_diagram called without recent get_diagram - rejecting to prevent data loss",
+                )
+                return {
+                    content: [
+                        {
+                            type: "text",
+                            text:
+                                "Error: You must call get_diagram first before edit_diagram.\n\n" +
+                                "This ensures you have the latest diagram state including any manual edits the user made in the browser. " +
+                                "Please call get_diagram, then use that XML to construct your edit operations.",
+                        },
+                    ],
+                    isError: true,
+                }
+            }
+
             // Fetch latest state from browser
             const browserState = getState(currentSession.id)
             if (browserState?.xml) {
@@ -411,6 +443,18 @@ server.registerTool(
                 }
             }
 
+            // Request browser to push fresh state and wait for it
+            const syncRequested = requestSync(currentSession.id)
+            if (syncRequested) {
+                const synced = await waitForSync(currentSession.id)
+                if (!synced) {
+                    log.warn("get_diagram: sync timeout - state may be stale")
+                }
+            }
+
+            // Mark that get_diagram was called (for edit_diagram workflow check)
+            currentSession.lastGetDiagramTime = Date.now()
+
             // Fetch latest state from browser
             const browserState = getState(currentSession.id)
             if (browserState?.xml) {