diff --git a/app/api/chat/route.ts b/app/api/chat/route.ts index 97bf486e..ce75d158 100644 --- a/app/api/chat/route.ts +++ b/app/api/chat/route.ts @@ -23,7 +23,6 @@ import { findCachedResponse } from "@/lib/cached-responses" import { dropInvalidToolCalls, fixToolInputJson, - isMinimalDiagram, replaceHistoricalToolInputs, validateFileParts, } from "@/lib/chat-helpers" @@ -50,6 +49,7 @@ import { import { allowPrivateUrls, isPrivateUrl } from "@/lib/ssrf-protection" import { getSystemPrompt } from "@/lib/system-prompts" import { getUserIdFromRequest } from "@/lib/user-id" +import { hasCells } from "@/packages/mcp-server/src/pages.ts" // No explicit cap: a reasoning model can spend minutes planning before it emits // the tool call, so take whatever the host allows. Vercel's own default is 300s, @@ -164,7 +164,7 @@ async function handleChatRequest(req: Request): Promise { // === CACHE CHECK START === const isFirstMessage = messages.length === 1 - const isEmptyDiagram = !xml || xml.trim() === "" || isMinimalDiagram(xml) + const isEmptyDiagram = !xml || !hasCells(xml) if (isFirstMessage && isEmptyDiagram) { const lastMessage = messages[0] diff --git a/components/chat-message-display.tsx b/components/chat-message-display.tsx index e765b39f..62b5ee62 100644 --- a/components/chat-message-display.tsx +++ b/components/chat-message-display.tsx @@ -37,11 +37,12 @@ import { ScrollArea } from "@/components/ui/scroll-area" import { useDictionary } from "@/hooks/use-dictionary" import { getApiEndpoint } from "@/lib/base-path" import { - applyDiagramOperations, convertToLegalXml, extractCompleteMxCells, replaceNodes, } from "@/lib/utils" +import { applyDiagramOperations } from "@/packages/mcp-server/src/diagram-operations.ts" +import { BLANK_MXFILE } from "@/packages/mcp-server/src/pages.ts" // Helper to extract complete operations from streaming input function getCompleteOperations( @@ -362,9 +363,7 @@ export function ChatMessageDisplay({ try { // An empty canvas gets a default mxfile to put the cells in - const baseXML = - chartXML || - `` + const baseXML = chartXML || BLANK_MXFILE const replacedXML = replaceNodes(baseXML, convertedXml) previousXML.current = convertedXml onDisplayChart(replacedXML, true) diff --git a/components/chat-panel.tsx b/components/chat-panel.tsx index 6f3a1588..c16d366c 100644 --- a/components/chat-panel.tsx +++ b/components/chat-panel.tsx @@ -32,7 +32,6 @@ import { useSessionManager } from "@/hooks/use-session-manager" import { useValidateDiagram } from "@/hooks/use-validate-diagram" import { getApiEndpoint } from "@/lib/base-path" import { findCachedResponse } from "@/lib/cached-responses" -import { isMinimalDiagram } from "@/lib/chat-helpers" import type { DrawioTheme } from "@/lib/drawio-themes" import { formatMessage } from "@/lib/i18n/utils" import { isPdfFile, isTextFile } from "@/lib/pdf-utils" @@ -41,7 +40,8 @@ import { STORAGE_KEYS } from "@/lib/storage" import type { UrlData } from "@/lib/url-utils" import { type FileData, useFileProcessor } from "@/lib/use-file-processor" import { useQuotaManager } from "@/lib/use-quota-manager" -import { cn, formatXML, isRealDiagram, wrapWithMxFile } from "@/lib/utils" +import { cn, formatXML, isRealDiagram } from "@/lib/utils" +import { BLANK_MXFILE, hasCells } from "@/packages/mcp-server/src/pages.ts" import type { ValidationState } from "./chat/ValidationCard" import { APPENDED_FILE_SECTIONS_PATTERN, @@ -812,10 +812,7 @@ export default function ChatPanel({ if (input.trim() && !isProcessing && !isExtracting) { // Check if input matches a cached example (only when no messages // yet and the canvas is empty, same rule as the server) - if ( - messages.length === 0 && - isMinimalDiagram(chartXMLRef.current || "") - ) { + if (messages.length === 0 && !hasCells(chartXMLRef.current || "")) { // Pass the file name so a user's own file never matches an example const cached = findCachedResponse( input.trim(), @@ -859,7 +856,7 @@ export default function ChatPanel({ // Snapshot the canvas before the example so editing this message works xmlSnapshotsRef.current.set( 0, - chartXMLRef.current || wrapWithMxFile(""), + chartXMLRef.current || BLANK_MXFILE, ) setInput("") sessionStorage.removeItem(SESSION_STORAGE_INPUT_KEY) diff --git a/components/chat/types.ts b/components/chat/types.ts index 45505a96..b62df1bf 100644 --- a/components/chat/types.ts +++ b/components/chat/types.ts @@ -1,8 +1,6 @@ -export interface DiagramOperation { - operation: "update" | "add" | "delete" - cell_id: string - new_xml?: string -} +import type { DiagramOperation } from "@/packages/mcp-server/src/diagram-operations.ts" + +export type { DiagramOperation } export interface ToolPartLike { type: string diff --git a/components/dev-xml-simulator.tsx b/components/dev-xml-simulator.tsx index 530baab7..a636dde5 100644 --- a/components/dev-xml-simulator.tsx +++ b/components/dev-xml-simulator.tsx @@ -2,7 +2,7 @@ import { useEffect, useRef, useState } from "react" import { useDictionary } from "@/hooks/use-dictionary" -import { wrapWithMxFile } from "@/lib/utils" +import { prepareNewDiagram } from "@/packages/mcp-server/src/new-diagram.ts" // Dev XML presets for streaming simulator const DEV_XML_PRESETS: Record = { @@ -237,8 +237,8 @@ export function DevXmlSimulator({ }) // Display the final diagram - const fullXml = wrapWithMxFile(xml) - onDisplayChart(fullXml) + const prepared = prepareNewDiagram(xml) + if (prepared.ok) onDisplayChart(prepared.xml) setIsSimulating(false) } diff --git a/contexts/diagram-context.tsx b/contexts/diagram-context.tsx index af62ee16..451e175e 100644 --- a/contexts/diagram-context.tsx +++ b/contexts/diagram-context.tsx @@ -6,6 +6,10 @@ import type { DrawIoEmbedRef, EventExport } from "react-drawio" import { toast } from "sonner" import type { ExportFormat } from "@/components/save-dialog" import { getApiEndpoint } from "@/lib/base-path" +import { + BLANK_MXFILE, + normalizeToMxfile, +} from "@/packages/mcp-server/src/pages.ts" import { validateAndFixXml } from "@/packages/mcp-server/src/xml-validation.ts" import { extractDiagramXML, isRealDiagram } from "../lib/utils" @@ -262,7 +266,7 @@ export function DiagramProvider({ children }: { children: React.ReactNode }) { } const clearDiagram = () => { - const emptyDiagram = `` + const emptyDiagram = BLANK_MXFILE // Skip validation for trusted internal template (loadDiagram also sets chartXML) loadDiagram(emptyDiagram, true) setLatestSvg("") @@ -296,11 +300,11 @@ export function DiagramProvider({ children }: { children: React.ReactNode }) { const xml = fullDiagramXML?.trim() ? fullDiagramXML : extractDiagramXML(exportData) - let xmlContent = xml - if (!xml.includes("${xml}` - } - fileContent = xmlContent + fileContent = + normalizeToMxfile(xml, { + pageId: "page-1", + pageName: "Page-1", + }) ?? xml mimeType = "application/xml" extension = ".drawio" } else if (format === "png") { diff --git a/hooks/use-diagram-tool-handlers.ts b/hooks/use-diagram-tool-handlers.ts index 332ba8fa..2a316bf3 100644 --- a/hooks/use-diagram-tool-handlers.ts +++ b/hooks/use-diagram-tool-handlers.ts @@ -7,6 +7,7 @@ import type { import type { ValidationResult } from "@/lib/diagram-validator" import { formatValidationFeedback } from "@/lib/diagram-validator" import { isMxCellXmlComplete } from "@/lib/utils" +import { editDiagram } from "@/packages/mcp-server/src/edit-diagram.ts" import { prepareNewDiagram } from "@/packages/mcp-server/src/new-diagram.ts" const DEBUG = process.env.NODE_ENV === "development" @@ -401,27 +402,19 @@ ${finalXml} } } - const { applyDiagramOperations } = await import("@/lib/utils") - const { result: editedXml, errors } = applyDiagramOperations( - currentXml, - operations, - ) - - // Check for operation errors - if (errors.length > 0) { - const errorMessages = errors - .map( - (e) => - `- ${e.type} on cell_id="${e.cellId}": ${e.message}`, - ) - .join("\n") - + // All or nothing, checked like the MCP server's edit_diagram. + // The model sees the first page, so edits target it. + const outcome = editDiagram(currentXml, operations, {}) + if (!outcome.ok) { + const reason = outcome.pageError + ? outcome.errors[0] + : `No changes were made because ${outcome.errors.length} operation(s) failed:\n${outcome.errors.map((e) => `- ${e}`).join("\n")}` restoreOriginal() addToolOutput({ tool: "edit_diagram", toolCallId: toolCall.toolCallId, state: "output-error", - errorText: `Some operations failed:\n${errorMessages} + errorText: `${reason} Current diagram XML: \`\`\`xml @@ -435,36 +428,12 @@ Please check the cell IDs and retry.`, return } - // loadDiagram validates and returns error if invalid - const validationError = onDisplayChart(editedXml) - if (validationError) { - console.warn( - "[edit_diagram] Validation error:", - validationError, - ) - restoreOriginal() - addToolOutput({ - tool: "edit_diagram", - toolCallId: toolCall.toolCallId, - state: "output-error", - errorText: `Edit produced invalid XML: ${validationError} - -Current diagram XML: -\`\`\`xml -${currentXml} -\`\`\` - -Please fix the operations to avoid structural issues.`, - }) - // Clean up the shared original XML ref - editDiagramOriginalXmlRef.current.delete(toolCall.toolCallId) - return - } + onDisplayChart(outcome.xml, true) onExport() addToolOutput({ tool: "edit_diagram", toolCallId: toolCall.toolCallId, - output: `Successfully applied ${operations.length} operation(s) to the diagram.`, + output: `Successfully applied ${outcome.applied} operation(s) to the diagram.`, }) // Clean up the shared original XML ref editDiagramOriginalXmlRef.current.delete(toolCall.toolCallId) diff --git a/lib/chat-helpers.ts b/lib/chat-helpers.ts index 572a0794..ddd8a730 100644 --- a/lib/chat-helpers.ts +++ b/lib/chat-helpers.ts @@ -53,13 +53,6 @@ export function validateFileParts(messages: any[]): { return { valid: true } } -// Helper function to check if diagram is minimal/empty -// Empty means no mxCell besides the root cells "0" and "1". Cells drawn in -// draw.io get random ids, so checking for id="2" is not enough. -export function isMinimalDiagram(xml: string): boolean { - return !/]*\bid="(?![01]")/.test(xml) -} - // A tool-call input providers accept: a non-empty JSON object function isValidToolInput(input: unknown): boolean { return !!input && typeof input === "object" && Object.keys(input).length > 0 diff --git a/lib/utils.ts b/lib/utils.ts index 1b75d668..ad72509f 100644 --- a/lib/utils.ts +++ b/lib/utils.ts @@ -1,9 +1,6 @@ import { type ClassValue, clsx } from "clsx" import * as pako from "pako" import { twMerge } from "tailwind-merge" -import type { DiagramOperation } from "@/components/chat/types" - -export type { DiagramOperation } export function cn(...inputs: ClassValue[]) { return twMerge(clsx(inputs)) @@ -213,61 +210,6 @@ export function convertToLegalXml(xmlString: string): string { return result } -/** - * Wrap XML content with the full mxfile structure required by draw.io. - * Always adds root cells (id="0" and id="1") automatically. - * If input already contains root cells, they are removed to avoid duplication. - * LLM should only generate mxCell elements starting from id="2". - * @param xml - The XML string (bare mxCells, , , or full ) - * @returns Full mxfile-wrapped XML string with root cells included - */ -export function wrapWithMxFile(xml: string): string { - const ROOT_CELLS = '' - - if (!xml || !xml.trim()) { - return `${ROOT_CELLS}` - } - - // Already has full structure - if (xml.includes("${xml}` - } - - // Has wrapper - extract inner content - let content = xml - if (xml.includes("")) { - content = xml.replace(/<\/?root>/g, "").trim() - } - - // Strip trailing LLM wrapper tags (from any provider: Anthropic, DeepSeek, etc.) - // Find the last valid mxCell ending and remove everything after it - const lastSelfClose = content.lastIndexOf("/>") - const lastMxCellClose = content.lastIndexOf("") - const lastValidEnd = Math.max(lastSelfClose, lastMxCellClose) - if (lastValidEnd !== -1) { - const endOffset = lastMxCellClose > lastSelfClose ? 9 : 2 - const suffix = content.slice(lastValidEnd + endOffset) - // If suffix is only closing tags (wrapper tags), strip it - if (/^(\s*<\/[^>]+>)*\s*$/.test(suffix)) { - content = content.slice(0, lastValidEnd + endOffset) - } - } - - // Remove any existing root cells from content (LLM shouldn't include them, but handle it gracefully) - // Use flexible patterns that match both self-closing (/>) and non-self-closing (>) formats - content = content - .replace(/]*\bid=["']0["'][^>]*(?:\/>|><\/mxCell>)/g, "") - .replace(/]*\bid=["']1["'][^>]*(?:\/>|><\/mxCell>)/g, "") - .trim() - - return `${ROOT_CELLS}${content}` -} - /** * Replace nodes in a Draw.io XML diagram * @param currentXML - The original Draw.io XML string @@ -370,299 +312,6 @@ export function replaceNodes(currentXML: string, nodes: string): string { } } -// ============================================================================ -// ID-based Diagram Operations -// ============================================================================ - -export interface OperationError { - type: "update" | "add" | "delete" - cellId: string - message: string -} - -export interface ApplyOperationsResult { - result: string - errors: OperationError[] -} - -/** - * draw.io wraps cells that have links, tooltips or custom data in - * /, and the wrapper carries the id instead of the mxCell. - */ -function getCellWrapper(cell: Element): Element | null { - const parent = cell.parentElement - return parent?.tagName === "object" || parent?.tagName === "UserObject" - ? parent - : null -} - -/** Id of a cell, read from its wrapper when the mxCell has none */ -function getCellId(cell: Element): string | null { - return ( - cell.getAttribute("id") || - getCellWrapper(cell)?.getAttribute("id") || - null - ) -} - -/** Element to replace or remove for a cell (the wrapper if there is one) */ -function getCellNode(cell: Element): Element { - return getCellWrapper(cell) || cell -} - -/** - * Apply diagram operations (update/add/delete) using ID-based lookup. - * This replaces the text-matching approach with direct DOM manipulation. - * - * @param xmlContent - The full mxfile XML content - * @param operations - Array of operations to apply - * @returns Object with result XML and any errors - */ -export function applyDiagramOperations( - xmlContent: string, - operations: DiagramOperation[], -): ApplyOperationsResult { - const errors: OperationError[] = [] - - // Parse the XML - const parser = new DOMParser() - const doc = parser.parseFromString(xmlContent, "text/xml") - - // Check for parse errors - const parseError = doc.querySelector("parsererror") - if (parseError) { - return { - result: xmlContent, - errors: [ - { - type: "update", - cellId: "", - message: `XML parse error: ${parseError.textContent}`, - }, - ], - } - } - - // Find the root element (inside mxGraphModel) - const root = doc.querySelector("root") - if (!root) { - return { - result: xmlContent, - errors: [ - { - type: "update", - cellId: "", - message: "Could not find element in XML", - }, - ], - } - } - - // Build a map of cell IDs to elements (wrapper elements for wrapped cells) - const cellMap = new Map() - root.querySelectorAll("mxCell").forEach((cell) => { - const id = getCellId(cell) - if (id) cellMap.set(id, getCellNode(cell)) - }) - // Cells removed by delete operations in this batch - const deletedIds = new Set() - - // Process each operation - for (const op of operations) { - if (op.operation === "update") { - const existingCell = cellMap.get(op.cell_id) - if (!existingCell) { - errors.push({ - type: "update", - cellId: op.cell_id, - message: `Cell with id="${op.cell_id}" not found`, - }) - continue - } - - if (!op.new_xml) { - errors.push({ - type: "update", - cellId: op.cell_id, - message: "new_xml is required for update operation", - }) - continue - } - - // Parse the new XML - const newDoc = parser.parseFromString( - `${op.new_xml}`, - "text/xml", - ) - const newCell = newDoc.querySelector("mxCell") - if (!newCell) { - errors.push({ - type: "update", - cellId: op.cell_id, - message: "new_xml must contain an mxCell element", - }) - continue - } - - // Validate ID matches - const newCellId = getCellId(newCell) - 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 (with its wrapper, if any) - const importedNode = doc.importNode(getCellNode(newCell), true) - existingCell.parentNode?.replaceChild(importedNode, existingCell) - - // Update the map with the new element - cellMap.set(op.cell_id, importedNode) - } else if (op.operation === "add") { - // Check if ID already exists - if (cellMap.has(op.cell_id)) { - errors.push({ - type: "add", - cellId: op.cell_id, - message: `Cell with id="${op.cell_id}" already exists`, - }) - continue - } - - if (!op.new_xml) { - errors.push({ - type: "add", - cellId: op.cell_id, - message: "new_xml is required for add operation", - }) - continue - } - - // Parse the new XML - const newDoc = parser.parseFromString( - `${op.new_xml}`, - "text/xml", - ) - const newCell = newDoc.querySelector("mxCell") - if (!newCell) { - errors.push({ - type: "add", - cellId: op.cell_id, - message: "new_xml must contain an mxCell element", - }) - continue - } - - // Validate ID matches - const newCellId = getCellId(newCell) - 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 (with its wrapper, if any) - const importedNode = doc.importNode(getCellNode(newCell), true) - root.appendChild(importedNode) - - // Add to map - cellMap.set(op.cell_id, importedNode) - } else if (op.operation === "delete") { - // Protect root cells from deletion - if (op.cell_id === "0" || op.cell_id === "1") { - errors.push({ - type: "delete", - cellId: op.cell_id, - message: `Cannot delete root cell "${op.cell_id}"`, - }) - continue - } - - const existingCell = cellMap.get(op.cell_id) - if (!existingCell) { - // Cells cascade-deleted earlier in this batch are skipped silently - // (AI may redundantly list children/edges) - if (!deletedIds.has(op.cell_id)) { - errors.push({ - type: "delete", - cellId: op.cell_id, - message: `Cell with id="${op.cell_id}" not found`, - }) - } - continue - } - - // Cascade delete: collect all cells to delete (children + edges + self) - const cellsToDelete = new Set() - - // Recursive function to find all descendants - const collectDescendants = (cellId: string) => { - if (cellsToDelete.has(cellId)) return - cellsToDelete.add(cellId) - - // Find children (cells where parent === cellId) - const children = root.querySelectorAll( - `mxCell[parent="${cellId}"]`, - ) - children.forEach((child) => { - const childId = getCellId(child) - if (childId && childId !== "0" && childId !== "1") { - collectDescendants(childId) - } - }) - } - - // Collect the target cell and all its descendants - collectDescendants(op.cell_id) - - // Find edges referencing any of the cells to be deleted - // Also recursively collect children of those edges (e.g., edge labels) - for (const cellId of cellsToDelete) { - const referencingEdges = root.querySelectorAll( - `mxCell[source="${cellId}"], mxCell[target="${cellId}"]`, - ) - referencingEdges.forEach((edge) => { - const edgeId = getCellId(edge) - // Protect root cells from being added via edge references - if (edgeId && edgeId !== "0" && edgeId !== "1") { - // Recurse to collect edge's children (like labels) - collectDescendants(edgeId) - } - }) - } - - // Log what will be deleted - if (cellsToDelete.size > 1) { - console.log( - `[applyDiagramOperations] Cascade delete "${op.cell_id}" → deleting ${cellsToDelete.size} cells: ${Array.from(cellsToDelete).join(", ")}`, - ) - } - - // Delete all collected cells - for (const cellId of cellsToDelete) { - const cell = cellMap.get(cellId) - if (cell) { - cell.parentNode?.removeChild(cell) - cellMap.delete(cellId) - deletedIds.add(cellId) - } - } - } - } - - // Serialize back to string - const serializer = new XMLSerializer() - const result = serializer.serializeToString(doc) - - return { result, errors } -} - /** * Decode an xmlsvg export (SVG data URL) into uncompressed diagram XML. * Only the first page is returned; for the full multi-page document use the diff --git a/packages/mcp-server/src/http-server.ts b/packages/mcp-server/src/http-server.ts index 6acd200e..5b058fe1 100644 --- a/packages/mcp-server/src/http-server.ts +++ b/packages/mcp-server/src/http-server.ts @@ -40,6 +40,7 @@ import { updateLastHistorySvg, } from "./history.ts" import { log } from "./logger.ts" +import { BLANK_MXFILE } from "./pages.ts" // Configurable draw.io embed URL for private deployments const DRAWIO_BASE_URL = @@ -57,10 +58,6 @@ function getOrigin(url: string): string { const DRAWIO_ORIGIN = getOrigin(DRAWIO_BASE_URL) -// Minimal blank diagram used to bootstrap new sessions. -// This avoids the draw.io embed spinner (spin=1) getting stuck when no `load(xml)` is ever sent. -const DEFAULT_DIAGRAM_XML = `` - // Normalize URL for iframe src - ensure no double slashes function normalizeUrl(url: string): string { // Remove trailing slash to avoid double slashes @@ -91,7 +88,9 @@ function ensureSessionStateInitialized(sessionId: string): void { if (stateStore.has(sessionId)) return // Not a change worth saving: the browser fills it on its next push - setState(sessionId, DEFAULT_DIAGRAM_XML, undefined, false, false) + // A blank diagram keeps the draw.io spinner (spin=1) from waiting + // forever when no load(xml) is ever sent + setState(sessionId, BLANK_MXFILE, undefined, false, false) } interface SessionState { diff --git a/packages/mcp-server/src/pages.ts b/packages/mcp-server/src/pages.ts index 0a70d977..3e572419 100644 --- a/packages/mcp-server/src/pages.ts +++ b/packages/mcp-server/src/pages.ts @@ -86,6 +86,9 @@ function stripXmlDeclaration(xml: string): string { const ROOT_CELLS = '' +/** A one-page document with only the root cells */ +export const BLANK_MXFILE = `${ROOT_CELLS}` + /** * Turn a list of bare cells (optionally inside ) into a one-page * , adding the "0" and "1" root cells. The model then only diff --git a/scripts/test-diagram-operations.mjs b/scripts/test-diagram-operations.mjs deleted file mode 100644 index b58ad8dd..00000000 --- a/scripts/test-diagram-operations.mjs +++ /dev/null @@ -1,375 +0,0 @@ -/** - * Simple test script for applyDiagramOperations function - * Run with: node scripts/test-diagram-operations.mjs - */ - -import { JSDOM } from "jsdom" - -// Set up DOMParser for Node.js environment -const dom = new JSDOM() -globalThis.DOMParser = dom.window.DOMParser -globalThis.XMLSerializer = dom.window.XMLSerializer - -// Import the function (we'll inline it since it's not ESM exported) -function applyDiagramOperations(xmlContent, operations) { - const errors = [] - const parser = new DOMParser() - const doc = parser.parseFromString(xmlContent, "text/xml") - - const parseError = doc.querySelector("parsererror") - if (parseError) { - return { - result: xmlContent, - errors: [ - { - operation: "update", - cellId: "", - message: `XML parse error: ${parseError.textContent}`, - }, - ], - } - } - - const root = doc.querySelector("root") - if (!root) { - return { - result: xmlContent, - errors: [ - { - operation: "update", - cellId: "", - message: "Could not find element in XML", - }, - ], - } - } - - const cellMap = new Map() - root.querySelectorAll("mxCell").forEach((cell) => { - const id = cell.getAttribute("id") - if (id) cellMap.set(id, cell) - }) - - for (const op of operations) { - if (op.operation === "update") { - const existingCell = cellMap.get(op.cell_id) - if (!existingCell) { - errors.push({ - operation: "update", - cellId: op.cell_id, - message: `Cell with id="${op.cell_id}" not found`, - }) - continue - } - if (!op.new_xml) { - errors.push({ - operation: "update", - cellId: op.cell_id, - message: "new_xml is required for update operation", - }) - continue - } - const newDoc = parser.parseFromString( - `${op.new_xml}`, - "text/xml", - ) - const newCell = newDoc.querySelector("mxCell") - if (!newCell) { - errors.push({ - operation: "update", - cellId: op.cell_id, - message: "new_xml must contain an mxCell element", - }) - continue - } - const newCellId = newCell.getAttribute("id") - if (newCellId !== op.cell_id) { - errors.push({ - operation: "update", - cellId: op.cell_id, - message: `ID mismatch: cell_id is "${op.cell_id}" but new_xml has id="${newCellId}"`, - }) - continue - } - const importedNode = doc.importNode(newCell, true) - existingCell.parentNode?.replaceChild(importedNode, existingCell) - cellMap.set(op.cell_id, importedNode) - } else if (op.operation === "add") { - if (cellMap.has(op.cell_id)) { - errors.push({ - operation: "add", - cellId: op.cell_id, - message: `Cell with id="${op.cell_id}" already exists`, - }) - continue - } - if (!op.new_xml) { - errors.push({ - operation: "add", - cellId: op.cell_id, - message: "new_xml is required for add operation", - }) - continue - } - const newDoc = parser.parseFromString( - `${op.new_xml}`, - "text/xml", - ) - const newCell = newDoc.querySelector("mxCell") - if (!newCell) { - errors.push({ - operation: "add", - cellId: op.cell_id, - message: "new_xml must contain an mxCell element", - }) - continue - } - const newCellId = newCell.getAttribute("id") - if (newCellId !== op.cell_id) { - errors.push({ - operation: "add", - cellId: op.cell_id, - message: `ID mismatch: cell_id is "${op.cell_id}" but new_xml has id="${newCellId}"`, - }) - continue - } - const importedNode = doc.importNode(newCell, true) - root.appendChild(importedNode) - cellMap.set(op.cell_id, importedNode) - } else if (op.operation === "delete") { - const existingCell = cellMap.get(op.cell_id) - if (!existingCell) { - errors.push({ - operation: "delete", - cellId: op.cell_id, - message: `Cell with id="${op.cell_id}" not found`, - }) - continue - } - existingCell.parentNode?.removeChild(existingCell) - cellMap.delete(op.cell_id) - } - } - - const serializer = new XMLSerializer() - const result = serializer.serializeToString(doc) - return { result, errors } -} - -// Test data -const sampleXml = ` - - - - - - - - - - - - - - - - - - -` - -let passed = 0 -let failed = 0 - -function test(name, fn) { - try { - fn() - console.log(`✓ ${name}`) - passed++ - } catch (e) { - console.log(`✗ ${name}`) - console.log(` Error: ${e.message}`) - failed++ - } -} - -function assert(condition, message) { - if (!condition) throw new Error(message || "Assertion failed") -} - -// Tests -test("Update operation changes cell value", () => { - const { result, errors } = applyDiagramOperations(sampleXml, [ - { - operation: "update", - cell_id: "2", - new_xml: - '', - }, - ]) - assert( - errors.length === 0, - `Expected no errors, got: ${JSON.stringify(errors)}`, - ) - assert( - result.includes('value="Updated Box A"'), - "Updated value should be in result", - ) - assert( - !result.includes('value="Box A"'), - "Old value should not be in result", - ) -}) - -test("Update operation fails for non-existent cell", () => { - const { errors } = applyDiagramOperations(sampleXml, [ - { - operation: "update", - cell_id: "999", - new_xml: '', - }, - ]) - assert(errors.length === 1, "Should have one error") - assert( - errors[0].message.includes("not found"), - "Error should mention not found", - ) -}) - -test("Update operation fails on ID mismatch", () => { - const { errors } = applyDiagramOperations(sampleXml, [ - { - operation: "update", - cell_id: "2", - new_xml: '', - }, - ]) - assert(errors.length === 1, "Should have one error") - assert( - errors[0].message.includes("ID mismatch"), - "Error should mention ID mismatch", - ) -}) - -test("Add operation creates new cell", () => { - const { result, errors } = applyDiagramOperations(sampleXml, [ - { - operation: "add", - cell_id: "new1", - new_xml: - '', - }, - ]) - assert( - errors.length === 0, - `Expected no errors, got: ${JSON.stringify(errors)}`, - ) - assert(result.includes('id="new1"'), "New cell should be in result") - assert( - result.includes('value="New Box"'), - "New cell value should be in result", - ) -}) - -test("Add operation fails for duplicate ID", () => { - const { errors } = applyDiagramOperations(sampleXml, [ - { - operation: "add", - cell_id: "2", - new_xml: '', - }, - ]) - assert(errors.length === 1, "Should have one error") - assert( - errors[0].message.includes("already exists"), - "Error should mention already exists", - ) -}) - -test("Add operation fails on ID mismatch", () => { - const { errors } = applyDiagramOperations(sampleXml, [ - { - operation: "add", - cell_id: "new1", - new_xml: '', - }, - ]) - assert(errors.length === 1, "Should have one error") - assert( - errors[0].message.includes("ID mismatch"), - "Error should mention ID mismatch", - ) -}) - -test("Delete operation removes cell", () => { - const { result, errors } = applyDiagramOperations(sampleXml, [ - { operation: "delete", cell_id: "3" }, - ]) - assert( - errors.length === 0, - `Expected no errors, got: ${JSON.stringify(errors)}`, - ) - assert(!result.includes('id="3"'), "Deleted cell should not be in result") - assert(result.includes('id="2"'), "Other cells should remain") -}) - -test("Delete operation fails for non-existent cell", () => { - const { errors } = applyDiagramOperations(sampleXml, [ - { operation: "delete", cell_id: "999" }, - ]) - assert(errors.length === 1, "Should have one error") - assert( - errors[0].message.includes("not found"), - "Error should mention not found", - ) -}) - -test("Multiple operations in sequence", () => { - const { result, errors } = applyDiagramOperations(sampleXml, [ - { - operation: "update", - cell_id: "2", - new_xml: - '', - }, - { - operation: "add", - cell_id: "new1", - new_xml: - '', - }, - { operation: "delete", cell_id: "3" }, - ]) - assert( - errors.length === 0, - `Expected no errors, got: ${JSON.stringify(errors)}`, - ) - assert( - result.includes('value="Updated"'), - "Updated value should be present", - ) - assert(result.includes('id="new1"'), "Added cell should be present") - assert(!result.includes('id="3"'), "Deleted cell should not be present") -}) - -test("Invalid XML returns parse error", () => { - const { errors } = applyDiagramOperations(" { - const { errors } = applyDiagramOperations("", [ - { operation: "delete", cell_id: "1" }, - ]) - assert(errors.length === 1, "Should have one error") - assert( - errors[0].message.includes("root"), - "Error should mention root element", - ) -}) - -// Summary -console.log(`\n${passed} passed, ${failed} failed`) -process.exit(failed > 0 ? 1 : 0) diff --git a/tests/e2e/diagram-content.spec.ts b/tests/e2e/diagram-content.spec.ts index 6aabf697..6e1cb457 100644 --- a/tests/e2e/diagram-content.spec.ts +++ b/tests/e2e/diagram-content.spec.ts @@ -1,34 +1,45 @@ -import { expect, test } from "@playwright/test" +import { expect, type Page, test } from "@playwright/test" import { getIframe, sendMessage, waitForCompleteCount } from "./lib/fixtures" /** - * Checks what draw.io actually shows after display_diagram, not only the + * Checks what draw.io actually shows after the diagram tools, not only the * tool card. The tool input is streamed in chunks like a real model, and * the browser tool handler (not the server) completes the tool call. */ -function streamedToolCall(xml: string) { +function streamedToolCall(toolName: string, input: unknown) { const toolCallId = `call_${Math.random().toString(36).slice(2)}` - const input = JSON.stringify({ xml }) - const chunks = input.match(/[\s\S]{1,40}/g) ?? [] + const chunks = JSON.stringify(input).match(/[\s\S]{1,40}/g) ?? [] const events = [ { type: "start", messageId: `msg_${toolCallId}` }, - { type: "tool-input-start", toolCallId, toolName: "display_diagram" }, + { type: "tool-input-start", toolCallId, toolName }, ...chunks.map((inputTextDelta) => ({ type: "tool-input-delta", toolCallId, inputTextDelta, })), - { - type: "tool-input-available", - toolCallId, - toolName: "display_diagram", - input: { xml }, - }, + { type: "tool-input-available", toolCallId, toolName, input }, { type: "finish" }, ] return `${events.map((e) => `data: ${JSON.stringify(e)}\n\n`).join("")}data: [DONE]\n\n` } +const END_TURN = + 'data: {"type":"start"}\n\ndata: {"type":"finish"}\n\ndata: [DONE]\n\n' + +/** Answer each chat request with the next reply, then end the turn */ +async function mockReplies(p: Page, replies: string[]) { + await p.route("**/api/chat", async (route) => { + await route.fulfill({ + status: 200, + contentType: "text/event-stream", + body: replies.shift() ?? END_TURN, + }) + }) + await p.goto("/", { waitUntil: "networkidle" }) + await getIframe(p).waitFor({ state: "visible", timeout: 30000 }) + return p.frameLocator("iframe") +} + const cell = (id: string, label: string, x: number) => `` const page = (id: string, cells: string) => @@ -46,20 +57,10 @@ const NEW_CELLS = test("display_diagram replaces the document with the fixed diagram", async ({ page: p, }) => { - const replies = [TWO_PAGES, NEW_CELLS] - await p.route("**/api/chat", async (route) => { - const xml = replies.shift() - await route.fulfill({ - status: 200, - contentType: "text/event-stream", - body: xml - ? streamedToolCall(xml) - : 'data: {"type":"start"}\n\ndata: {"type":"finish"}\n\ndata: [DONE]\n\n', - }) - }) - await p.goto("/", { waitUntil: "networkidle" }) - await getIframe(p).waitFor({ state: "visible", timeout: 30000 }) - const canvas = p.frameLocator("iframe") + const canvas = await mockReplies(p, [ + streamedToolCall("display_diagram", { xml: TWO_PAGES }), + streamedToolCall("display_diagram", { xml: NEW_CELLS }), + ]) await sendMessage(p, "Draw two pages") await waitForCompleteCount(p, 1) @@ -79,3 +80,59 @@ test("display_diagram replaces the document with the fixed diagram", async ({ await expect(canvas.getByText("Old A")).toHaveCount(0) await expect(canvas.getByText("Second", { exact: true })).toHaveCount(0) }) + +test("edit_diagram applies all operations or none", async ({ page: p }) => { + const canvas = await mockReplies(p, [ + streamedToolCall("display_diagram", { + xml: cell("a", "Alpha", 40) + cell("b", "Beta", 220), + }), + streamedToolCall("edit_diagram", { + operations: [ + { operation: "delete", cell_id: "a" }, + { + operation: "update", + cell_id: "b", + new_xml: cell("b", "Beta two", 220), + }, + { + operation: "add", + cell_id: "c", + new_xml: cell("c", "Gamma", 400), + }, + ], + }), + // The first operation is fine, the second fails: nothing is kept + streamedToolCall("edit_diagram", { + operations: [ + { + operation: "update", + cell_id: "c", + new_xml: cell("c", "Broken", 400), + }, + { operation: "delete", cell_id: "missing" }, + ], + }), + ]) + + await sendMessage(p, "Draw two boxes") + await waitForCompleteCount(p, 1) + await expect(canvas.getByText("Alpha", { exact: true })).toBeVisible({ + timeout: 15000, + }) + + await sendMessage(p, "Change them") + await waitForCompleteCount(p, 2) + await expect(canvas.getByText("Gamma", { exact: true })).toBeVisible({ + timeout: 15000, + }) + await expect(canvas.getByText("Beta two", { exact: true })).toBeVisible() + await expect(canvas.getByText("Alpha", { exact: true })).toHaveCount(0) + + await sendMessage(p, "Change again") + await expect(p.getByText(/No changes were made/).first()).toBeAttached({ + timeout: 15000, + }) + await p.waitForTimeout(1000) + await expect(canvas.getByText("Gamma", { exact: true })).toBeVisible() + await expect(canvas.getByText("Broken", { exact: true })).toHaveCount(0) +}) diff --git a/tests/unit/chat-helpers.test.ts b/tests/unit/chat-helpers.test.ts index c3f672a9..f7b16bf2 100644 --- a/tests/unit/chat-helpers.test.ts +++ b/tests/unit/chat-helpers.test.ts @@ -6,7 +6,6 @@ import { describe, expect, it } from "vitest" import { dropInvalidToolCalls, fixToolInputJson, - isMinimalDiagram, replaceHistoricalToolInputs, validateFileParts, } from "@/lib/chat-helpers" @@ -95,36 +94,6 @@ describe("validateFileParts", () => { }) }) -describe("isMinimalDiagram", () => { - it("returns true for empty diagram", () => { - const xml = '' - expect(isMinimalDiagram(xml)).toBe(true) - }) - - it("returns false for diagram with content", () => { - const xml = - '' - expect(isMinimalDiagram(xml)).toBe(false) - }) - - it("handles whitespace correctly", () => { - const xml = ' ' - expect(isMinimalDiagram(xml)).toBe(true) - }) - - it("returns false for a shape drawn in draw.io with a random id", () => { - const xml = - '' - expect(isMinimalDiagram(xml)).toBe(false) - }) - - it("does not mistake ids that start with 0 or 1 for root cells", () => { - const xml = - '' - expect(isMinimalDiagram(xml)).toBe(false) - }) -}) - describe("replaceHistoricalToolInputs", () => { it("replaces display_diagram tool inputs with placeholder", () => { const messages = [ diff --git a/tests/unit/mcp-core.test.ts b/tests/unit/mcp-core.test.ts index 78a5a820..9a3de097 100644 --- a/tests/unit/mcp-core.test.ts +++ b/tests/unit/mcp-core.test.ts @@ -180,3 +180,86 @@ describe("autoFixXml", () => { 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") + }) +}) diff --git a/tests/unit/utils.test.ts b/tests/unit/utils.test.ts index f2ed8aff..15330f56 100644 --- a/tests/unit/utils.test.ts +++ b/tests/unit/utils.test.ts @@ -1,11 +1,5 @@ import { describe, expect, it } from "vitest" -import { - applyDiagramOperations, - cn, - extractCompleteMxCells, - isMxCellXmlComplete, - wrapWithMxFile, -} from "@/lib/utils" +import { cn, extractCompleteMxCells, isMxCellXmlComplete } from "@/lib/utils" describe("isMxCellXmlComplete", () => { it("returns false for empty/null input", () => { @@ -71,36 +65,6 @@ describe("isMxCellXmlComplete", () => { }) }) -describe("wrapWithMxFile", () => { - it("wraps empty string with default structure", () => { - const result = wrapWithMxFile("") - expect(result).toContain("") - expect(result).toContain("") - expect(result).toContain('') - expect(result).toContain('') - }) - - it("wraps raw mxCell content", () => { - const xml = '' - const result = wrapWithMxFile(xml) - expect(result).toContain("") - expect(result).toContain(xml) - expect(result).toContain("") - }) - - it("returns full mxfile unchanged", () => { - const fullXml = - '' - const result = wrapWithMxFile(fullXml) - expect(result).toBe(fullXml) - }) - - it("handles whitespace in input", () => { - const result = wrapWithMxFile(" ") - expect(result).toContain("") - }) -}) - describe("cn (class name utility)", () => { it("merges class names", () => { expect(cn("foo", "bar")).toBe("foo bar") @@ -132,50 +96,3 @@ describe("extractCompleteMxCells", () => { ) }) }) - -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") - }) -})