diff --git a/components/chat-message-display.tsx b/components/chat-message-display.tsx index 8f8fd993..e765b39f 100644 --- a/components/chat-message-display.tsx +++ b/components/chat-message-display.tsx @@ -41,7 +41,6 @@ import { convertToLegalXml, extractCompleteMxCells, replaceNodes, - validateAndFixXml, } from "@/lib/utils" // Helper to extract complete operations from streaming input @@ -345,73 +344,32 @@ export function ChatMessageDisplay({ } } + // Streaming preview of display_diagram: draw the complete cells written + // so far. The tool handler validates and loads the final diagram. const handleDisplayChart = useCallback( - (xml: string, showToast = false) => { - let currentXml = xml || "" + (xml: string) => { + const completeCells = extractCompleteMxCells(xml || "") + if (!completeCells) return + const convertedXml = convertToLegalXml(completeCells) + if (convertedXml === previousXML.current) return - // During streaming (showToast=false), extract only complete mxCell elements - // This allows progressive rendering even with partial/incomplete trailing XML - if (!showToast) { - const completeCells = extractCompleteMxCells(currentXml) - if (!completeCells) { - return - } - currentXml = completeCells - } + // Skip this update while the cells written so far don't parse + const testDoc = new DOMParser().parseFromString( + `${convertedXml}`, + "text/xml", + ) + if (testDoc.querySelector("parsererror")) return - const convertedXml = convertToLegalXml(currentXml) - if (convertedXml !== previousXML.current) { - // Parse and validate XML BEFORE calling replaceNodes - const parser = new DOMParser() - // Wrap in root element for parsing multiple mxCell elements - const testDoc = parser.parseFromString( - `${convertedXml}`, - "text/xml", - ) - const parseError = testDoc.querySelector("parsererror") - - if (parseError) { - // Only show toast if this is the final XML (not during streaming) - if (showToast) { - toast.error(dict.errors.malformedXml) - } - return // Skip this update - } - - try { - // If chartXML is empty, create a default mxfile structure to use with replaceNodes - // This ensures the XML is properly wrapped in mxfile/diagram/mxGraphModel format - const baseXML = - chartXML || - `` - const replacedXML = replaceNodes(baseXML, convertedXml) - - // During streaming (showToast=false), skip heavy validation for lower latency - // The quick DOM parse check above catches malformed XML - // Full validation runs on final output (showToast=true) - if (!showToast) { - previousXML.current = convertedXml - onDisplayChart(replacedXML, true) - return - } - - // Final output: run full validation and auto-fix - const validation = validateAndFixXml(replacedXML) - if (validation.valid) { - previousXML.current = convertedXml - // Use fixed XML if available, otherwise use original - const xmlToLoad = validation.fixed || replacedXML - onDisplayChart(xmlToLoad, true) - } else { - toast.error(dict.errors.validationFailed) - } - } catch (error) { - console.error("Error processing XML:", error) - // Only show toast if this is the final XML (not during streaming) - if (showToast) { - toast.error(dict.errors.failedToProcess) - } - } + try { + // An empty canvas gets a default mxfile to put the cells in + const baseXML = + chartXML || + `` + const replacedXML = replaceNodes(baseXML, convertedXml) + previousXML.current = convertedXml + onDisplayChart(replacedXML, true) + } catch (error) { + console.error("Error processing XML:", error) } }, [chartXML, onDisplayChart], @@ -497,10 +455,7 @@ export function ChatMessageDisplay({ return // Skip redundant processing } - if ( - state === "input-streaming" || - state === "input-available" - ) { + if (state === "input-streaming") { // Debounce streaming updates - queue the XML and process after delay pendingXmlRef.current = xml @@ -513,10 +468,7 @@ export function ChatMessageDisplay({ debounceTimeoutRef.current = null pendingXmlRef.current = null if (pendingXml) { - handleDisplayChart( - pendingXml, - false, - ) + handleDisplayChart(pendingXml) lastProcessedXmlRef.current.set( toolCallId, pendingXml, @@ -527,17 +479,16 @@ export function ChatMessageDisplay({ ) } } else if ( - state === "output-available" && !processedToolCalls.current.has(toolCallId) ) { - // Final output - process immediately (clear any pending debounce) + // Input complete: the tool handler loads the + // validated diagram, so drop a queued preview + // that would draw the raw cells over it if (debounceTimeoutRef.current) { clearTimeout(debounceTimeoutRef.current) debounceTimeoutRef.current = null pendingXmlRef.current = null } - // Show toast only if final XML is malformed - handleDisplayChart(xml, true) processedToolCalls.current.add(toolCallId) // Clean up the ref entry - tool is complete, no longer needed lastProcessedXmlRef.current.delete(toolCallId) diff --git a/contexts/diagram-context.tsx b/contexts/diagram-context.tsx index e4071012..af62ee16 100644 --- a/contexts/diagram-context.tsx +++ b/contexts/diagram-context.tsx @@ -6,11 +6,8 @@ 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 { - extractDiagramXML, - isRealDiagram, - validateAndFixXml, -} from "../lib/utils" +import { validateAndFixXml } from "@/packages/mcp-server/src/xml-validation.ts" +import { extractDiagramXML, isRealDiagram } from "../lib/utils" interface DiagramContextType { chartXML: string @@ -170,9 +167,11 @@ export function DiagramProvider({ children }: { children: React.ReactNode }) { ): string | null => { let xmlToLoad = chart - // Validate XML structure before loading (unless skipped for internal use) + // Validate XML structure before loading (unless skipped for internal + // use). Not strict: the XML may hold the user's own diagram, and the + // tool handlers check model XML strictly before it gets here. if (!skipValidation) { - const validation = validateAndFixXml(chart) + const validation = validateAndFixXml(chart, { strict: false }) if (!validation.valid) { console.warn( "[loadDiagram] Validation error:", diff --git a/hooks/use-diagram-tool-handlers.ts b/hooks/use-diagram-tool-handlers.ts index a4350789..332ba8fa 100644 --- a/hooks/use-diagram-tool-handlers.ts +++ b/hooks/use-diagram-tool-handlers.ts @@ -6,10 +6,14 @@ import type { } from "@/components/chat/ValidationCard" import type { ValidationResult } from "@/lib/diagram-validator" import { formatValidationFeedback } from "@/lib/diagram-validator" -import { isMxCellXmlComplete, wrapWithMxFile } from "@/lib/utils" +import { isMxCellXmlComplete } from "@/lib/utils" +import { prepareNewDiagram } from "@/packages/mcp-server/src/new-diagram.ts" const DEBUG = process.env.NODE_ENV === "development" +// display_diagram replaces the document with this one page +const NEW_PAGE = { pageId: "page-1", pageName: "Page-1" } + interface ToolCall { toolCallId: string toolName: string @@ -173,11 +177,12 @@ NEXT STEP: Call append_diagram with the continuation XML. const finalXml = xml partialXmlRef.current = "" // Reset any partial from previous truncation - // Wrap raw XML with full mxfile structure for draw.io - const fullXml = wrapWithMxFile(finalXml) - - // loadDiagram validates and returns error if invalid - const validationError = onDisplayChart(fullXml) + // Wrap, validate and auto-fix the model's XML like the MCP server's + // create_new_diagram, then load it + const prepared = prepareNewDiagram(finalXml, NEW_PAGE) + const validationError = prepared.ok + ? onDisplayChart(prepared.xml, true) + : prepared.error if (validationError) { console.warn("[display_diagram] Validation error:", validationError) @@ -545,8 +550,10 @@ Start your continuation with the NEXT character after where it stopped.`, const finalXml = partialXmlRef.current partialXmlRef.current = "" // Reset - const fullXml = wrapWithMxFile(finalXml) - const validationError = onDisplayChart(fullXml) + const prepared = prepareNewDiagram(finalXml, NEW_PAGE) + const validationError = prepared.ok + ? onDisplayChart(prepared.xml, true) + : prepared.error if (validationError) { addToolOutput({ diff --git a/lib/i18n/dictionaries/en.json b/lib/i18n/dictionaries/en.json index 77ec3e0a..ace8cda0 100644 --- a/lib/i18n/dictionaries/en.json +++ b/lib/i18n/dictionaries/en.json @@ -178,9 +178,6 @@ "networkError": "Network error. Please check your connection.", "retryLimit": "Auto-retry limit reached ({max}). Please try again manually.", "continuationRetryLimit": "Continuation retry limit reached ({max}). The diagram may be too complex.", - "validationFailed": "Diagram validation failed. Please try regenerating.", - "malformedXml": "AI generated invalid diagram XML. Please try regenerating.", - "failedToProcess": "Failed to process diagram. Please try regenerating.", "sessionCorrupted": "Session data was corrupted. Starting fresh.", "failedToSave": "Failed to save messages to localStorage", "failedToRestore": "Failed to restore from localStorage", diff --git a/lib/i18n/dictionaries/ja.json b/lib/i18n/dictionaries/ja.json index 71018e95..214892cf 100644 --- a/lib/i18n/dictionaries/ja.json +++ b/lib/i18n/dictionaries/ja.json @@ -178,9 +178,6 @@ "networkError": "ネットワークエラー。接続を確認してください。", "retryLimit": "自動再試行制限に達しました（{max}）。手動で再試行してください。", "continuationRetryLimit": "継続再試行制限に達しました（{max}）。ダイアグラムが複雑すぎる可能性があります。", - "validationFailed": "ダイアグラムの検証に失敗しました。再生成してみてください。", - "malformedXml": "AI が無効なダイアグラム XML を生成しました。再生成してみてください。", - "failedToProcess": "ダイアグラムの処理に失敗しました。再生成してみてください。", "sessionCorrupted": "セッションデータが破損しました。最初からやり直します。", "failedToSave": "localStorage へのメッセージの保存に失敗しました", "failedToRestore": "localStorage からの復元に失敗しました", diff --git a/lib/i18n/dictionaries/zh-Hant.json b/lib/i18n/dictionaries/zh-Hant.json index c6e78986..e4e12602 100644 --- a/lib/i18n/dictionaries/zh-Hant.json +++ b/lib/i18n/dictionaries/zh-Hant.json @@ -178,9 +178,6 @@ "networkError": "網路錯誤。請檢查您的連線。", "retryLimit": "已達自動重試限制（{max}）。請手動重試。", "continuationRetryLimit": "已達繼續重試限制（{max}）。圖表可能過於複雜。", - "validationFailed": "圖表驗證失敗。請嘗試重新產生。", - "malformedXml": "AI 產生的圖表 XML 無效。請嘗試重新產生。", - "failedToProcess": "無法處理圖表。請嘗試重新產生。", "sessionCorrupted": "工作階段資料已損壞。重新開始。", "failedToSave": "無法儲存訊息到 localStorage", "failedToRestore": "無法從 localStorage 還原", diff --git a/lib/i18n/dictionaries/zh.json b/lib/i18n/dictionaries/zh.json index ae549068..56d322f2 100644 --- a/lib/i18n/dictionaries/zh.json +++ b/lib/i18n/dictionaries/zh.json @@ -178,9 +178,6 @@ "networkError": "网络错误。请检查您的连接。", "retryLimit": "已达到自动重试限制（{max}）。请手动重试。", "continuationRetryLimit": "已达到继续重试限制（{max}）。图表可能过于复杂。", - "validationFailed": "图表验证失败。请尝试重新生成。", - "malformedXml": "AI 生成的图表 XML 无效。请尝试重新生成。", - "failedToProcess": "无法处理图表。请尝试重新生成。", "sessionCorrupted": "会话数据已损坏。重新开始。", "failedToSave": "无法保存消息到 localStorage", "failedToRestore": "无法从 localStorage 恢复", diff --git a/lib/utils.ts b/lib/utils.ts index b51ad5dd..1b75d668 100644 --- a/lib/utils.ts +++ b/lib/utils.ts @@ -28,29 +28,6 @@ export function isRealDiagram(xml: string | undefined | null): boolean { return !!xml && xml.length > MIN_REAL_DIAGRAM_LENGTH } -// ============================================================================ -// XML Validation/Fix Constants -// ============================================================================ - -/** Maximum XML size to process (1MB) - larger XMLs may cause performance issues */ -const MAX_XML_SIZE = 1_000_000 - -/** Maximum iterations for aggressive cell dropping to prevent infinite loops */ -const MAX_DROP_ITERATIONS = 10 - -/** Structural attributes that should not be duplicated in draw.io */ -const STRUCTURAL_ATTRS = [ - "edge", - "parent", - "source", - "target", - "vertex", - "connectable", -] - -/** Valid XML entity names */ -const VALID_ENTITIES = new Set(["lt", "gt", "amp", "quot", "apos"]) - // ============================================================================ // mxCell XML Helpers // ============================================================================ @@ -114,72 +91,6 @@ export function extractCompleteMxCells(xml: string | undefined | null): string { return (xml.match(cellPattern) || []).join("\n") } -// ============================================================================ -// XML Parsing Helpers -// ============================================================================ - -interface ParsedTag { - tag: string - tagName: string - isClosing: boolean - isSelfClosing: boolean - startIndex: number - endIndex: number -} - -/** - * Parse XML tags while properly handling quoted strings - * This is a shared utility used by both validation and fixing logic - */ -function parseXmlTags(xml: string): ParsedTag[] { - const tags: ParsedTag[] = [] - let i = 0 - - while (i < xml.length) { - const tagStart = xml.indexOf("<", i) - if (tagStart === -1) break - - // Find matching > by tracking quotes - let tagEnd = tagStart + 1 - let inQuote = false - let quoteChar = "" - - while (tagEnd < xml.length) { - const c = xml[tagEnd] - if (inQuote) { - if (c === quoteChar) inQuote = false - } else { - if (c === '"' || c === "'") { - inQuote = true - quoteChar = c - } else if (c === ">") { - break - } - } - tagEnd++ - } - - if (tagEnd >= xml.length) break - - const tag = xml.substring(tagStart, tagEnd + 1) - i = tagEnd + 1 - - const tagMatch = /^<(\/?)([a-zA-Z][a-zA-Z0-9:_-]*)/.exec(tag) - if (!tagMatch) continue - - tags.push({ - tag, - tagName: tagMatch[2], - isClosing: tagMatch[1] === "/", - isSelfClosing: tag.endsWith("/>"), - startIndex: tagStart, - endIndex: tagEnd, - }) - } - - return tags -} - /** * Format XML string with proper indentation and line breaks * @param xml - The XML string to format @@ -752,991 +663,6 @@ export function applyDiagramOperations( return { result, errors } } -// ============================================================================ -// Validation Helper Functions -// ============================================================================ - -/** Check for duplicate structural attributes in a tag */ -function checkDuplicateAttributes(xml: string): string | null { - const structuralSet = new Set(STRUCTURAL_ATTRS) - const tagPattern = /<[^>]+>/g - let tagMatch - while ((tagMatch = tagPattern.exec(xml)) !== null) { - const tag = tagMatch[0] - const attrPattern = /\s([a-zA-Z_:][a-zA-Z0-9_:.-]*)\s*=/g - const attributes = new Map() - let attrMatch - while ((attrMatch = attrPattern.exec(tag)) !== null) { - const attrName = attrMatch[1] - attributes.set(attrName, (attributes.get(attrName) || 0) + 1) - } - const duplicates = Array.from(attributes.entries()) - .filter(([name, count]) => count > 1 && structuralSet.has(name)) - .map(([name]) => name) - if (duplicates.length > 0) { - return `Invalid XML: Duplicate structural attribute(s): ${duplicates.join(", ")}. Remove duplicate attributes.` - } - } - return null -} - -/** Matches one page of a document (the last one may be unclosed) */ -const PAGE_PATTERN = /|$)/g - -const ID_ATTR_PATTERN = /\bid\s*=\s*["']([^"']+)["']/gi - -/** - * Split XML into pages. Ids only need to be unique within a page: every - * page of a multi-page document has its own root cells "0" and "1". - */ -function splitPages(xml: string): string[] { - return xml.match(PAGE_PATTERN) || [xml] -} - -/** Ids that appear more than once, with their counts */ -function findDuplicateIds(xml: string): Map { - const ids = new Map() - for (const match of xml.matchAll(ID_ATTR_PATTERN)) { - ids.set(match[1], (ids.get(match[1]) || 0) + 1) - } - return new Map(Array.from(ids).filter(([, count]) => count > 1)) -} - -/** Check for duplicate IDs in XML (per page) */ -function checkDuplicateIds(xml: string): string | null { - for (const page of splitPages(xml)) { - const duplicateIds = Array.from(findDuplicateIds(page)).map( - ([id, count]) => `'${id}' (${count}x)`, - ) - if (duplicateIds.length > 0) { - return `Invalid XML: Found duplicate ID(s): ${duplicateIds.slice(0, 3).join(", ")}. All id attributes must be unique.` - } - } - return null -} - -/** Rename repeated ids in one page (keeps the first occurrence) */ -function renameDuplicateIds(xml: string): { xml: string; renamed: number } { - const duplicateIds = findDuplicateIds(xml) - if (duplicateIds.size === 0) return { xml, renamed: 0 } - - const idCounters = new Map() - const renamedXml = xml.replace(ID_ATTR_PATTERN, (match, id) => { - if (!duplicateIds.has(id)) return match - - const count = idCounters.get(id) || 0 - idCounters.set(id, count + 1) - - if (count === 0) return match // Keep first occurrence - - // Rename subsequent occurrences (the id sits just before the closing quote) - return `${match.slice(0, -id.length - 1)}${id}_dup${count}${match.slice(-1)}` - }) - return { xml: renamedXml, renamed: duplicateIds.size } -} - -/** - * Returns a function telling whether a position is inside a quoted attribute - * value. Positions must be queried in increasing order: the scan resumes where - * it stopped instead of starting over, which keeps large documents fast. - */ -function createQuoteTracker(str: string): (pos: number) => boolean { - let i = 0 - let inQuote = false - let quoteChar = "" - return (pos: number) => { - for (; i < pos && i < str.length; i++) { - const c = str[i] - if (inQuote) { - if (c === quoteChar) inQuote = false - } else if (c === '"' || c === "'") { - // Only quotes that follow "=" open an attribute value - let j = i - 1 - while (j >= 0 && /\s/.test(str[j])) j-- - if (j >= 0 && str[j] === "=") { - inQuote = true - quoteChar = c - } - } - } - return inQuote - } -} - -/** 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 { - const cellTagPattern = /<\/?mxCell[^>]*>/g - const cellStack: number[] = [] - let cellMatch - while ((cellMatch = cellTagPattern.exec(xml)) !== null) { - const tag = cellMatch[0] - if (tag.startsWith("")) { - if (cellStack.length > 0) cellStack.pop() - } else if (!tag.endsWith("/>")) { - const isLabelOrGeometry = - /\sas\s*=\s*["'](valueLabel|geometry)["']/.test(tag) - if (!isLabelOrGeometry) { - cellStack.push(cellMatch.index) - if (cellStack.length > 1) { - return "Invalid XML: Found nested mxCell tags. Cells should be siblings, not nested inside other mxCell elements." - } - } - } - } - return null -} - -/** - * Validates draw.io XML structure for common issues - * Uses DOM parsing + additional regex checks for high accuracy - * @param xml - The XML string to validate - * @returns null if valid, error message string if invalid - */ -export function validateMxCellStructure(xml: string): 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. First use DOM parser to catch syntax errors (most accurate) - try { - const parser = new DOMParser() - const doc = parser.parseFromString(xml, "text/xml") - const parseError = doc.querySelector("parsererror") - if (parseError) { - return `Invalid XML: The XML contains syntax errors (likely unescaped special characters like <, >, & in attribute values). Please escape special characters: use < for <, > for >, & for &, " for ". Regenerate the diagram with properly escaped values.` - } - - // DOM-based checks for nested mxCell - const allCells = doc.querySelectorAll("mxCell") - for (const cell of allCells) { - if (cell.parentElement?.tagName === "mxCell") { - const id = cell.getAttribute("id") || "unknown" - return `Invalid XML: Found nested mxCell (id="${id}"). Cells should be siblings, not nested inside other mxCell elements.` - } - } - } catch (error) { - // Log unexpected DOMParser errors before falling back to regex checks - console.warn( - "[validateMxCellStructure] DOMParser threw unexpected error, falling back to regex validation:", - error, - ) - } - - // 1. Check for CDATA wrapper (invalid at document root) - if (/^\s* from end" - } - - // 2. 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 - 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 - const nestedCellError = checkNestedMxCells(xml) - if (nestedCellError) { - return nestedCellError - } - - return null -} - -/** - * Attempts to auto-fix common XML issues in draw.io diagrams - * @param xml - The XML string to fix - * @returns Object with fixed XML and list of fixes applied - */ -export function autoFixXml(xml: string): { fixed: string; fixes: string[] } { - let fixed = xml - const fixes: string[] = [] - - // 0. Fix JSON-escaped XML (common when XML is stored in JSON without unescaping) - // Only apply when we see JSON-escaped attribute patterns like =\"value\" - // Don't apply to legitimate \n in value attributes (draw.io uses these for line breaks) - if (/=\\"/.test(fixed)) { - // Replace literal \" with actual quotes - fixed = fixed.replace(/\\"/g, '"') - // Replace literal \n with actual newlines (only after confirming JSON-escaped) - fixed = fixed.replace(/\\n/g, "\n") - fixes.push("Fixed JSON-escaped XML") - } - - // 1. Remove CDATA wrapper (MUST be before text-before-root check) - if (/^\s*\s*$/, "") - fixes.push("Removed CDATA wrapper") - } - - // 1b. Strip trailing LLM wrapper tags (DeepSeek, Anthropic, etc.) - // These are closing tags after the last valid mxCell that break XML parsing - const lastSelfClose = fixed.lastIndexOf("/>") - const lastMxCellClose = fixed.lastIndexOf("") - const lastValidEnd = Math.max(lastSelfClose, lastMxCellClose) - if (lastValidEnd !== -1) { - const endOffset = lastMxCellClose > lastSelfClose ? 9 : 2 - const suffix = fixed.slice(lastValidEnd + endOffset) - // If suffix contains only closing tags (wrapper tags) or whitespace, strip it - if (/^(\s*<\/[^>]+>)+\s*$/.test(suffix)) { - fixed = fixed.slice(0, lastValidEnd + endOffset) - fixes.push("Stripped trailing LLM wrapper tags") - } - } - - // 2. Remove text before XML declaration or root element (only if it's garbage text, not valid XML) - const xmlStart = fixed.search(/<(\?xml|mxGraphModel|mxfile)/i) - if (xmlStart > 0 && !/^<[a-zA-Z]/.test(fixed.trim())) { - fixed = fixed.substring(xmlStart) - fixes.push("Removed text before XML root") - } - - // 2. Fix duplicate attributes (keep first occurrence, remove duplicates) - let dupAttrFixed = false - fixed = fixed.replace(/<[^>]+>/g, (tag) => { - let newTag = tag - - for (const attr of STRUCTURAL_ATTRS) { - // Find all occurrences of this attribute - const attrRegex = new RegExp( - `\\s${attr}\\s*=\\s*["'][^"']*["']`, - "gi", - ) - const matches = tag.match(attrRegex) - - if (matches && matches.length > 1) { - // Keep first, remove others - let firstKept = false - newTag = newTag.replace(attrRegex, (m) => { - if (!firstKept) { - firstKept = true - return m - } - dupAttrFixed = true - return "" - }) - } - } - return newTag - }) - if (dupAttrFixed) { - fixes.push("Removed duplicate structural attributes") - } - - // 3. Fix unescaped & characters (but not valid entities) - // Match & not followed by valid entity pattern - const ampersandPattern = - /&(?!(?:lt|gt|amp|quot|apos|#[0-9]+|#x[0-9a-fA-F]+);)/g - if (ampersandPattern.test(fixed)) { - fixed = fixed.replace( - /&(?!(?:lt|gt|amp|quot|apos|#[0-9]+|#x[0-9a-fA-F]+);)/g, - "&", - ) - fixes.push("Escaped unescaped & characters") - } - - // 3. Fix invalid entity names like &quot; -> " - // Common mistake: double-escaping - const invalidEntities = [ - { pattern: /&quot;/g, replacement: """, name: "&quot;" }, - { pattern: /&lt;/g, replacement: "<", name: "&lt;" }, - { pattern: /&gt;/g, replacement: ">", name: "&gt;" }, - { pattern: /&apos;/g, replacement: "'", name: "&apos;" }, - { pattern: /&amp;/g, replacement: "&", name: "&amp;" }, - ] - for (const { pattern, replacement, name } of invalidEntities) { - if (pattern.test(fixed)) { - fixed = fixed.replace(pattern, replacement) - fixes.push(`Fixed double-escaped entity ${name}`) - } - } - - // 3b. Fix malformed attribute values where " is used as delimiter instead of actual quotes - // Pattern: attr="value" should become attr="value" (the " was meant to be the quote delimiter) - // This commonly happens with dashPattern="1 1;" - // Matches inside another attribute value are kept: rich text labels like - // value="<font color="#ff0000">..." are valid. - const isInsideQuotesFor3b = createQuoteTracker(fixed) - let malformedQuotesFixed = false - fixed = fixed.replace( - /(\s[a-zA-Z][a-zA-Z0-9_:-]*)="([^&]*?)"/g, - (match: string, attr: string, value: string, offset: number) => { - if (isInsideQuotesFor3b(offset)) return match - malformedQuotesFixed = true - return `${attr}="${value}"` - }, - ) - if (malformedQuotesFixed) { - fixes.push( - 'Fixed malformed attribute quotes (="..." to ="...")', - ) - } - - // 3c. Fix malformed closing tags like -> - const malformedClosingTag = /<\/([a-zA-Z][a-zA-Z0-9]*)\s*\/>/g - if (malformedClosingTag.test(fixed)) { - fixed = fixed.replace(/<\/([a-zA-Z][a-zA-Z0-9]*)\s*\/>/g, "") - fixes.push("Fixed malformed closing tags ( to )") - } - - // 3d. Fix missing space between attributes like vertex="1"parent="1" - // Requires name=" right after the quote, so the opening quote of a value - // such as style="rounded=1;..." is not mistaken for a closing one. - const missingSpacePattern = /"([a-zA-Z_:][\w:.-]*=")/g - if (missingSpacePattern.test(fixed)) { - fixed = fixed.replace(missingSpacePattern, '" $1') - fixes.push("Added missing space between attributes") - } - - // 3e. Fix unescaped quotes in style color values like fillColor="#fff2e6" - // The " after Color= prematurely ends the style attribute. Remove it. - // Pattern: ;fillColor="#fff → ;fillColor=#fff (remove first ", keep second as style closer) - const quotedColorPattern = /;([a-zA-Z]*[Cc]olor)="#/ - if (quotedColorPattern.test(fixed)) { - fixed = fixed.replace(/;([a-zA-Z]*[Cc]olor)="#/g, ";$1=#") - fixes.push("Removed quotes around color values in style") - } - - // 4. Fix unescaped < and > in attribute values - // < is required to be escaped, > is not strictly required but we escape for consistency - const attrPattern = /(=\s*")([^"]*?)(<)([^"]*?)(")/g - let attrMatch - let hasUnescapedLt = false - while ((attrMatch = attrPattern.exec(fixed)) !== null) { - if (!attrMatch[3].startsWith("<")) { - hasUnescapedLt = true - break - } - } - if (hasUnescapedLt) { - // Replace < and > with < and > inside attribute values - fixed = fixed.replace(/=\s*"([^"]*)"/g, (_match, value) => { - const escaped = value.replace(//g, ">") - return `="${escaped}"` - }) - fixes.push("Escaped <> characters in attribute values") - } - - // 5. Fix invalid character references (remove malformed ones) - // Pattern: &#x followed by non-hex chars before ; - const invalidHexRefs: string[] = [] - fixed = fixed.replace(/&#x([^;]*);/g, (match, hex) => { - if (/^[0-9a-fA-F]+$/.test(hex) && hex.length > 0) { - return match // Valid hex ref, keep it - } - invalidHexRefs.push(match) - return "" // Remove invalid ref - }) - if (invalidHexRefs.length > 0) { - fixes.push( - `Removed ${invalidHexRefs.length} invalid hex character reference(s)`, - ) - } - - // 6. Fix invalid decimal character references - const invalidDecRefs: string[] = [] - fixed = fixed.replace(/&#([^x][^;]*);/g, (match, dec) => { - if (/^[0-9]+$/.test(dec) && dec.length > 0) { - return match // Valid decimal ref, keep it - } - invalidDecRefs.push(match) - return "" // Remove invalid ref - }) - if (invalidDecRefs.length > 0) { - fixes.push( - `Removed ${invalidDecRefs.length} invalid decimal character reference(s)`, - ) - } - - // 7. Fix invalid comment syntax (replace -- with - repeatedly until none left) - fixed = fixed.replace(//g, (match, content) => { - if (/--/.test(content)) { - // Keep replacing until no double hyphens remain - let fixedContent = content - while (/--/.test(fixedContent)) { - fixedContent = fixedContent.replace(/--/g, "-") - } - fixes.push("Fixed invalid comment syntax (removed double hyphens)") - return `` - } - return match - }) - - // 8. Fix tags that should be (common LLM mistake) - // This handles both opening and closing tags - const hasCellTags = /<\/?Cell[\s>]/i.test(fixed) - if (hasCellTags) { - console.log("[autoFixXml] Step 8: Found tags to fix") - const beforeFix = fixed - fixed = fixed.replace(//gi, "") - fixed = fixed.replace(/<\/Cell>/gi, "") - if (beforeFix !== fixed) { - console.log("[autoFixXml] Step 8: Fixed tags") - } - fixes.push("Fixed tags to ") - } - - // 8b. Fix common closing tag typos (MUST run before foreign tag removal) - const tagTypos = [ - { wrong: /<\/mxElement>/gi, right: "", name: "" }, - { wrong: /<\/mxcell>/g, right: "", name: "" }, // case sensitivity - { - wrong: /<\/mxgeometry>/g, - right: "", - name: "", - }, - { wrong: /<\/mxpoint>/g, right: "", name: "" }, - { - wrong: /<\/mxgraphmodel>/gi, - right: "", - name: "", - }, - ] - for (const { wrong, right, name } of tagTypos) { - const before = fixed - fixed = fixed.replace(wrong, right) - if (fixed !== before) { - fixes.push(`Fixed typo ${name} to ${right}`) - } - } - - // 8c. Remove non-draw.io tags (after typo fixes so lowercase variants are fixed first) - // IMPORTANT: Only remove tags at the element level, NOT inside quoted attribute values - // Tags like ,
inside value="text" should be preserved (they're HTML content) - const validDrawioTags = new Set([ - "mxfile", - "diagram", - "mxGraphModel", - "root", - "mxCell", - "mxGeometry", - "mxPoint", - "Array", - "Object", - // Wrappers of cells with links, tooltips or custom data - "object", - "UserObject", - "mxRectangle", - ]) - - const isInsideQuotesFor8c = createQuoteTracker(fixed) - const foreignTagPattern = /<\/?([a-zA-Z][a-zA-Z0-9_]*)[^>]*>/g - let foreignMatch - const foreignTags = new Set() - const foreignTagPositions: Array<{ - tag: string - start: number - end: number - }> = [] - - while ((foreignMatch = foreignTagPattern.exec(fixed)) !== null) { - const tagName = foreignMatch[1] - // Skip if this is a valid draw.io tag - if (validDrawioTags.has(tagName)) continue - // Skip if this tag is inside a quoted attribute value - if (isInsideQuotesFor8c(foreignMatch.index)) continue - - foreignTags.add(tagName) - foreignTagPositions.push({ - tag: tagName, - start: foreignMatch.index, - end: foreignMatch.index + foreignMatch[0].length, - }) - } - - if (foreignTagPositions.length > 0) { - // Remove tags from end to start to preserve indices - foreignTagPositions.sort((a, b) => b.start - a.start) - for (const { start, end } of foreignTagPositions) { - fixed = fixed.slice(0, start) + fixed.slice(end) - } - fixes.push( - `Removed foreign tags: ${Array.from(foreignTags).join(", ")}`, - ) - } - - // 10. Fix unclosed tags by appending missing closing tags - // Use parseXmlTags helper to track open tags - const tagStack: string[] = [] - const parsedTags = parseXmlTags(fixed) - - for (const { tagName, isClosing, isSelfClosing } of parsedTags) { - if (isClosing) { - // Find matching opening tag (may not be the last one if there's mismatch) - const lastIdx = tagStack.lastIndexOf(tagName) - if (lastIdx !== -1) { - tagStack.splice(lastIdx, 1) - } - } else if (!isSelfClosing) { - tagStack.push(tagName) - } - } - - // If there are unclosed tags, append closing tags in reverse order - // But first verify with simple count that they're actually unclosed - if (tagStack.length > 0) { - const tagsToClose: string[] = [] - for (const tagName of tagStack.reverse()) { - // Simple count check: only close if opens > closes - const openCount = ( - fixed.match(new RegExp(`<${tagName}[\\s>]`, "gi")) || [] - ).length - const closeCount = ( - fixed.match(new RegExp(``, "gi")) || [] - ).length - if (openCount > closeCount) { - tagsToClose.push(tagName) - } - } - if (tagsToClose.length > 0) { - const closingTags = tagsToClose.map((t) => ``).join("\n") - fixed = fixed.trimEnd() + "\n" + closingTags - fixes.push( - `Closed ${tagsToClose.length} unclosed tag(s): ${tagsToClose.join(", ")}`, - ) - } - } - - // 10b. Remove extra closing tags (more closes than opens) - // Need to properly count self-closing tags (they don't need closing tags) - // IMPORTANT: Only count tags at element level, NOT inside quoted attribute values - const tagCounts = new Map< - string, - { opens: number; closes: number; selfClosing: number } - >() - // Match full tags to detect self-closing by checking if ends with /> - const fullTagPattern = /<(\/?[a-zA-Z][a-zA-Z0-9]*)[^>]*>/g - const isInsideQuotesFor10b = createQuoteTracker(fixed) - let tagCountMatch - while ((tagCountMatch = fullTagPattern.exec(fixed)) !== null) { - // Skip tags inside quoted attribute values (e.g., value="Title") - if (isInsideQuotesFor10b(tagCountMatch.index)) continue - - const fullMatch = tagCountMatch[0] // e.g., "" or "" - const tagPart = tagCountMatch[1] // e.g., "mxCell" or "/mxCell" - const isClosing = tagPart.startsWith("/") - const isSelfClosing = fullMatch.endsWith("/>") - const tagName = isClosing ? tagPart.slice(1) : tagPart - - // Only count valid draw.io tags - skip partial/invalid tags like "mx" from streaming - if (!validDrawioTags.has(tagName)) continue - - let counts = tagCounts.get(tagName) - if (!counts) { - counts = { opens: 0, closes: 0, selfClosing: 0 } - tagCounts.set(tagName, counts) - } - if (isClosing) { - counts.closes++ - } else if (isSelfClosing) { - counts.selfClosing++ - } else { - counts.opens++ - } - } - - // Log tag counts for debugging - for (const [tagName, counts] of tagCounts) { - if ( - tagName === "mxCell" || - tagName === "mxGeometry" || - counts.opens !== counts.closes - ) { - console.log( - `[autoFixXml] Step 10b: ${tagName} - opens: ${counts.opens}, closes: ${counts.closes}, selfClosing: ${counts.selfClosing}`, - ) - } - } - - // Find tags with extra closing tags (self-closing tags are balanced, don't need closing) - for (const [tagName, counts] of tagCounts) { - const extraCloses = counts.closes - counts.opens // Only compare opens vs closes (self-closing are balanced) - if (extraCloses > 0) { - console.log( - `[autoFixXml] Step 10b: ${tagName} has ${counts.opens} opens, ${counts.closes} closes, removing ${extraCloses} extra`, - ) - // Remove extra closing tags from the end - let removed = 0 - const closeTagPattern = new RegExp(``, "g") - const matches = [...fixed.matchAll(closeTagPattern)] - // Remove from the end (last occurrences are likely the extras) - for ( - let i = matches.length - 1; - i >= 0 && removed < extraCloses; - i-- - ) { - const match = matches[i] - const idx = match.index ?? 0 - fixed = fixed.slice(0, idx) + fixed.slice(idx + match[0].length) - removed++ - } - if (removed > 0) { - console.log( - `[autoFixXml] Step 10b: Removed ${removed} extra `, - ) - fixes.push( - `Removed ${removed} extra closing tag(s)`, - ) - } - } - } - - // 10c. Remove trailing garbage after last XML tag (e.g., stray backslashes, text) - // Find the last valid closing tag or self-closing tag - const closingTagPattern = /<\/[a-zA-Z][a-zA-Z0-9]*>|\/>/g - let lastValidTagEnd = -1 - let closingMatch - while ((closingMatch = closingTagPattern.exec(fixed)) !== null) { - lastValidTagEnd = closingMatch.index + closingMatch[0].length - } - if (lastValidTagEnd > 0 && lastValidTagEnd < fixed.length) { - const trailing = fixed.slice(lastValidTagEnd).trim() - if (trailing) { - fixed = fixed.slice(0, lastValidTagEnd) - fixes.push("Removed trailing garbage after last XML tag") - } - } - - // 11. Fix nested mxCell by flattening - // Pattern A: ...... (duplicate ID) - // Pattern B: ...... (different ID - true nesting) - // These passes work line by line and would break valid cells written on a - // single line, so each one runs only when cells are really nested. - if (checkNestedMxCells(fixed)) { - const lines = fixed.split("\n") - const newLines: string[] = [] - let nestedFixed = 0 - let extraClosingToRemove = 0 - - // First pass: fix duplicate ID nesting (same as before) - for (let i = 0; i < lines.length; i++) { - const line = lines[i] - const nextLine = lines[i + 1] - - // Check if current line and next line are both mxCell opening tags with same ID - if ( - nextLine && - /") && - !nextLine.includes("/>") - ) { - const id1 = line.match(/\bid\s*=\s*["']([^"']+)["']/)?.[1] - const id2 = nextLine.match(/\bid\s*=\s*["']([^"']+)["']/)?.[1] - - if (id1 && id1 === id2) { - nestedFixed++ - extraClosingToRemove++ // Need to remove one later - continue // Skip this duplicate opening line - } - } - - // Remove extra if we have pending removals - if (extraClosingToRemove > 0 && /^\s*<\/mxCell>\s*$/.test(line)) { - extraClosingToRemove-- - continue // Skip this closing tag - } - - newLines.push(line) - } - - if (nestedFixed > 0) { - fixed = newLines.join("\n") - fixes.push(`Flattened ${nestedFixed} duplicate-ID nested mxCell(s)`) - } - } - - if (checkNestedMxCells(fixed)) { - // Second pass: fix true nesting (different IDs) - // Insert before nested child to close parent - const lines2 = fixed.split("\n") - const newLines: string[] = [] - let trueNestedFixed = 0 - let cellDepth = 0 - let pendingCloseRemoval = 0 - - for (let i = 0; i < lines2.length; i++) { - const line = lines2[i] - const trimmed = line.trim() - - // Track mxCell depth - const isOpenCell = - /") - const isCloseCell = trimmed === "" - - if (isOpenCell) { - if (cellDepth > 0) { - // Found nested cell - insert closing tag for parent before this line - const indent = line.match(/^(\s*)/)?.[1] || "" - newLines.push(indent + "") - trueNestedFixed++ - pendingCloseRemoval++ // Need to remove one later - } - cellDepth = 1 // Reset to 1 since we just opened a new cell - newLines.push(line) - } else if (isCloseCell) { - if (pendingCloseRemoval > 0) { - pendingCloseRemoval-- - // Skip this extra closing tag - } else { - cellDepth = Math.max(0, cellDepth - 1) - newLines.push(line) - } - } else { - newLines.push(line) - } - } - - if (trueNestedFixed > 0) { - fixed = newLines.join("\n") - fixes.push(`Fixed ${trueNestedFixed} true nested mxCell(s)`) - } - } - - // 12. Fix duplicate IDs by appending suffix, page by page (ids such as the - // root cells "0" and "1" legitimately repeat across pages) - let renamedIds = 0 - const renamePage = (page: string) => { - const { xml: renamed, renamed: count } = renameDuplicateIds(page) - renamedIds += count - return renamed - } - fixed = / 0) { - fixes.push(`Renamed ${renamedIds} duplicate ID(s)`) - } - - // 9. Fix empty id attributes by generating unique IDs - let emptyIdCount = 0 - fixed = fixed.replace( - /]*)\sid\s*=\s*["']\s*["']([^>]*)>/g, - (_match, before, after) => { - emptyIdCount++ - const newId = `cell_${Date.now()}_${emptyIdCount}` - return `` - }, - ) - if (emptyIdCount > 0) { - fixes.push(`Generated ${emptyIdCount} missing ID(s)`) - } - - // 13. Aggressive: drop broken mxCell elements that can't be fixed - // Only do this if DOM parser still finds errors after all other fixes - if (typeof DOMParser !== "undefined") { - let droppedCells = 0 - let maxIterations = MAX_DROP_ITERATIONS - while (maxIterations-- > 0) { - const parser = new DOMParser() - const doc = parser.parseFromString(fixed, "text/xml") - const parseError = doc.querySelector("parsererror") - if (!parseError) break // Valid now! - - const errText = parseError.textContent || "" - const match = errText.match(/(\d+):\d+:/) - if (!match) break - - const errLine = parseInt(match[1], 10) - 1 - const lines = fixed.split("\n") - - // Find the mxCell containing this error line - let cellStart = errLine - let cellEnd = errLine - - // Go back to find 0 && !lines[cellStart].includes(" or /> - while (cellEnd < lines.length - 1) { - if ( - lines[cellEnd].includes("") || - lines[cellEnd].trim().endsWith("/>") - ) { - break - } - cellEnd++ - } - - // Remove these lines - lines.splice(cellStart, cellEnd - cellStart + 1) - fixed = lines.join("\n") - droppedCells++ - } - if (droppedCells > 0) { - fixes.push(`Dropped ${droppedCells} unfixable mxCell element(s)`) - } - } - - return { fixed, fixes } -} - -/** - * Validates XML and attempts to fix if invalid - * @param xml - The XML string to validate and potentially fix - * @returns Object with validation result, fixed XML if applicable, and fixes applied - */ -export function validateAndFixXml(xml: string): { - valid: boolean - error: string | null - fixed: string | null - fixes: string[] -} { - // First validation attempt - let error = validateMxCellStructure(xml) - - if (!error) { - return { valid: true, error: null, fixed: null, fixes: [] } - } - - // Try to fix - const { fixed, fixes } = autoFixXml(xml) - console.log("[validateAndFixXml] Fixes applied:", fixes) - - // Validate the fixed version - error = validateMxCellStructure(fixed) - if (error) { - console.log("[validateAndFixXml] Still invalid after fix:", error) - } - - if (!error) { - return { valid: true, error: null, fixed, fixes } - } - - // Still invalid after fixes - but return the partially fixed XML - // so we can see what was fixed and what error remains - return { - valid: false, - error, - fixed: fixes.length > 0 ? fixed : null, - fixes, - } -} - /** * 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/index.ts b/packages/mcp-server/src/index.ts index 9fd7253a..aa58e170 100644 --- a/packages/mcp-server/src/index.ts +++ b/packages/mcp-server/src/index.ts @@ -45,6 +45,7 @@ import { } from "./http-server.ts" import { parseDrawioFileContent } from "./load-diagram.ts" import { log } from "./logger.ts" +import { prepareNewDiagram } from "./new-diagram.ts" import { addPageToDoc, deletePageFromDoc, @@ -341,44 +342,21 @@ Rules: cells are siblings (never nested), ids are unique per page and start from } } - // Bare cells get the wrapper and root cells first: the strict - // parser rejects several top-level elements. Then validate and - // auto-fix (works for both mxfile and mxGraphModel inputs). - let xml = wrapCellsInModel(inputXml) - const { valid, error, fixed, fixes } = validateAndFixXml(xml) - if (fixed) { - xml = fixed - log.info(`XML auto-fixed: ${fixes.join(", ")}`) - } - if (!valid && error) { - log.error(`XML validation failed: ${error}`) + const prepared = prepareNewDiagram(inputXml) + if (!prepared.ok) { + log.error(prepared.error) return { content: [ - { - type: "text", - text: `Error: XML validation failed - ${error}`, - }, + { type: "text", text: `Error: ${prepared.error}` }, ], isError: true, } } - - // Normalise to the canonical mxfile shape so every later tool can - // assume "session.xml is always an mxfile". Bare - // inputs are wrapped into a single-page mxfile here. - const normalized = normalizeToMxfile(xml) - if (!normalized) { - return { - content: [ - { - type: "text", - text: "Error: XML must be the mxCell elements of one page, a , or an with one or more children.", - }, - ], - isError: true, - } + if (prepared.fixes.length > 0) { + log.info(`XML auto-fixed: ${prepared.fixes.join(", ")}`) } - xml = normalized + // Every later tool can assume session.xml is an mxfile + const xml = prepared.xml log.info(`Setting diagram content, ${xml.length} chars`) diff --git a/packages/mcp-server/src/new-diagram.ts b/packages/mcp-server/src/new-diagram.ts new file mode 100644 index 00000000..11bd09e6 --- /dev/null +++ b/packages/mcp-server/src/new-diagram.ts @@ -0,0 +1,36 @@ +/** + * A whole new diagram written by the model, for the create_new_diagram tool + * and the web app's display_diagram tool. + */ +import { normalizeToMxfile, wrapCellsInModel } from "./pages.ts" +import { validateAndFixXml } from "./xml-validation.ts" + +export type NewDiagram = + | { ok: true; xml: string; fixes: string[] } + | { ok: false; error: string } + +/** + * Bare cells get the wrapper and root cells first, 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, and finally turned into an . + */ +export function prepareNewDiagram( + input: string, + page: { pageId?: string; pageName?: string } = {}, +): NewDiagram { + let xml = wrapCellsInModel(input) + const { valid, error, fixed, fixes } = validateAndFixXml(xml) + if (fixed) xml = fixed + if (!valid) { + return { ok: false, error: `XML validation failed - ${error}` } + } + const normalized = normalizeToMxfile(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 } +} diff --git a/tests/e2e/diagram-content.spec.ts b/tests/e2e/diagram-content.spec.ts new file mode 100644 index 00000000..6aabf697 --- /dev/null +++ b/tests/e2e/diagram-content.spec.ts @@ -0,0 +1,81 @@ +import { expect, test } from "@playwright/test" +import { getIframe, sendMessage, waitForCompleteCount } from "./lib/fixtures" + +/** + * Checks what draw.io actually shows after display_diagram, 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) { + const toolCallId = `call_${Math.random().toString(36).slice(2)}` + const input = JSON.stringify({ xml }) + const chunks = input.match(/[\s\S]{1,40}/g) ?? [] + const events = [ + { type: "start", messageId: `msg_${toolCallId}` }, + { type: "tool-input-start", toolCallId, toolName: "display_diagram" }, + ...chunks.map((inputTextDelta) => ({ + type: "tool-input-delta", + toolCallId, + inputTextDelta, + })), + { + type: "tool-input-available", + toolCallId, + toolName: "display_diagram", + input: { xml }, + }, + { type: "finish" }, + ] + return `${events.map((e) => `data: ${JSON.stringify(e)}\n\n`).join("")}data: [DONE]\n\n` +} + +const cell = (id: string, label: string, x: number) => + `` +const page = (id: string, cells: string) => + `${cells}` + +const TWO_PAGES = `${page("First", cell("a", "Old A", 40))}${page("Second", cell("b", "Old B", 40))}` +// Bare cells with a duplicate id and an unescaped &, which get fixed, and +// a linked cell whose label lives on its UserObject wrapper +const NEW_CELLS = + cell("2", "Alpha", 40) + + cell("2", "Beta", 220) + + cell("3", "R&D", 400) + + `` + +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") + + await sendMessage(p, "Draw two pages") + await waitForCompleteCount(p, 1) + await expect(canvas.getByText("Old A")).toBeVisible({ timeout: 15000 }) + await expect(canvas.getByText("Second", { exact: true })).toBeVisible() + + await sendMessage(p, "Start over with three boxes") + await waitForCompleteCount(p, 2) + // Give a late preview time to redraw the raw cells, as it used to + await p.waitForTimeout(1000) + for (const label of ["Alpha", "Beta", "R&D", "Docs"]) { + await expect(canvas.getByText(label, { exact: true })).toBeVisible({ + timeout: 15000, + }) + } + // The old pages are gone + await expect(canvas.getByText("Old A")).toHaveCount(0) + await expect(canvas.getByText("Second", { exact: true })).toHaveCount(0) +}) diff --git a/tests/e2e/file-upload.spec.ts b/tests/e2e/file-upload.spec.ts index 6348a752..b7e254ad 100644 --- a/tests/e2e/file-upload.spec.ts +++ b/tests/e2e/file-upload.spec.ts @@ -24,7 +24,8 @@ test.describe("File Upload", () => { await page.goto("/", { waitUntil: "networkidle" }) await getIframe(page).waitFor({ state: "visible", timeout: 30000 }) - const fileInput = page.locator('input[type="file"]') + // The chat attachment input; the template panel has its own file input + const fileInput = page.locator('input[type="file"][multiple]') await fileInput.setInputFiles({ name: "test-image.png", @@ -44,7 +45,8 @@ test.describe("File Upload", () => { await page.goto("/", { waitUntil: "networkidle" }) await getIframe(page).waitFor({ state: "visible", timeout: 30000 }) - const fileInput = page.locator('input[type="file"]') + // The chat attachment input; the template panel has its own file input + const fileInput = page.locator('input[type="file"][multiple]') await fileInput.setInputFiles({ name: "test-image.png", @@ -91,7 +93,8 @@ test.describe("File Upload", () => { await page.goto("/", { waitUntil: "networkidle" }) await getIframe(page).waitFor({ state: "visible", timeout: 30000 }) - const fileInput = page.locator('input[type="file"]') + // The chat attachment input; the template panel has its own file input + const fileInput = page.locator('input[type="file"][multiple]') await fileInput.setInputFiles({ name: "architecture.png", @@ -115,7 +118,8 @@ test.describe("File Upload", () => { await page.goto("/", { waitUntil: "networkidle" }) await getIframe(page).waitFor({ state: "visible", timeout: 30000 }) - const fileInput = page.locator('input[type="file"]') + // The chat attachment input; the template panel has its own file input + const fileInput = page.locator('input[type="file"][multiple]') const largeBuffer = Buffer.alloc(3 * 1024 * 1024, "x") await fileInput.setInputFiles({ diff --git a/tests/unit/mcp-core.test.ts b/tests/unit/mcp-core.test.ts index 72f3a2f2..78a5a820 100644 --- a/tests/unit/mcp-core.test.ts +++ b/tests/unit/mcp-core.test.ts @@ -14,7 +14,11 @@ import { wrapCellsInModel, } from "@/packages/mcp-server/src/pages.ts" import { getXmlSyntaxError } from "@/packages/mcp-server/src/xml-syntax.ts" -import { validateAndFixXml } from "@/packages/mcp-server/src/xml-validation.ts" +import { + autoFixXml, + validateAndFixXml, + validateMxCellStructure, +} from "@/packages/mcp-server/src/xml-validation.ts" const box = (id: string, parent = "1") => `` @@ -88,3 +92,91 @@ describe("MCP diagram modules with a browser DOM", () => { expect(decompressPageContent("not compressed")).toBeNull() }) }) + +// 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(' { }) }) -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)).toContain("duplicate ID") - }) - - 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 within a page", () => { - const xml = `${page("p1", 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(' { const xml = ``