From 14c1aa8e1c6c5117cfdaa5192c11b7c950feecc1 Mon Sep 17 00:00:00 2001 From: Dayuan Jiang <34411969+DayuanJiang@users.noreply.github.com> Date: Sun, 21 Dec 2025 19:38:35 +0900 Subject: [PATCH] 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 | 43 ++++++++++++++++++++ packages/mcp-server/src/index.ts | 56 +++++++++++++++++++++++--- 2 files changed, 93 insertions(+), 6 deletions(-) diff --git a/packages/mcp-server/src/http-server.ts b/packages/mcp-server/src/http-server.ts index 27a2e7a..95ff6b2 100644 --- a/packages/mcp-server/src/http-server.ts +++ b/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() @@ -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 { + 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 { 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(' currentVersion && s.xml) { currentVersion = s.version; loadDiagram(s.xml, true); diff --git a/packages/mcp-server/src/index.ts b/packages/mcp-server/src/index.ts index ffb3912..e2af673 100644 --- a/packages/mcp-server/src/index.ts +++ b/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) {