diff --git a/components/chat/use-diagram-streaming.ts b/components/chat/use-diagram-streaming.ts index 8ebe5df3..e511d207 100644 --- a/components/chat/use-diagram-streaming.ts +++ b/components/chat/use-diagram-streaming.ts @@ -5,25 +5,25 @@ import type { RefObject } from "react" import { useCallback, useEffect, useRef } from "react" import type { DiagramOperation, ToolPartLike } from "@/components/chat/types" import { useDiagram } from "@/contexts/diagram-context" -import { pageSelectorFor } from "@/lib/diagram-pages" -import { - convertToLegalXml, - extractCompleteMxCells, - replaceNodes, -} from "@/lib/utils" +import { pageSelectorFor, replacePageModel } from "@/lib/diagram-pages" +import { extractCompleteMxCells } from "@/lib/utils" import { defaultLayerOf, expandCompactCells, } from "@/packages/mcp-server/src/compact-cells.ts" import { applyDiagramOperations } from "@/packages/mcp-server/src/diagram-operations.ts" import { targetPageXml } from "@/packages/mcp-server/src/edit-diagram.ts" -import { BLANK_MXFILE } from "@/packages/mcp-server/src/pages.ts" +import { + BLANK_MXFILE, + wrapCellsInModel, +} from "@/packages/mcp-server/src/pages.ts" import { addDefaultStyles, applyStyleClasses, edgeIdsOf, readStyleClasses, } from "@/packages/mcp-server/src/style-classes.ts" +import { autoFixXml } from "@/packages/mcp-server/src/xml-validation.ts" // Helper to extract complete operations from streaming input function getCompleteOperations( @@ -91,32 +91,31 @@ export function useDiagramStreaming({ const { classes, xml: cellsXml } = readStyleClasses(xml || "") const completeCells = extractCompleteMxCells(cellsXml) if (!completeCells) return - const convertedXml = addDefaultStyles( + // The same wrapping, fixes and expansion as the final diagram + // (prepareNewDiagram), so the preview matches what is committed + const model = addDefaultStyles( applyStyleClasses( - expandCompactCells(convertToLegalXml(completeCells)), + expandCompactCells( + autoFixXml(wrapCellsInModel(completeCells)).fixed, + ), classes, ), ) - if (convertedXml === previousXML.current) return + if (model === previousXML.current) return // Skip this update while the cells written so far don't parse - const testDoc = new DOMParser().parseFromString( - `${convertedXml}`, - "text/xml", - ) + const testDoc = new DOMParser().parseFromString(model, "text/xml") if (testDoc.querySelector("parsererror")) return try { - // Replace the model's page's cells so other pages stay - // intact. An empty canvas gets a default mxfile to put the - // cells in. - const baseXML = chartXML || BLANK_MXFILE - const replacedXML = replaceNodes( - baseXML, - convertedXml, + // Replace the model's page so other pages stay intact; an + // empty canvas gets the blank mxfile to put the page in + const replacedXML = replacePageModel( + chartXML, turnPageIdRef?.current, + model, ) - previousXML.current = convertedXml + previousXML.current = model loadDiagram(replacedXML, true, "preview") } catch (error) { console.error("Error processing XML:", error) diff --git a/hooks/use-dictionary.ts b/hooks/use-dictionary.ts index 42a3db8b..dccfc486 100644 --- a/hooks/use-dictionary.ts +++ b/hooks/use-dictionary.ts @@ -25,5 +25,3 @@ export function useDictionary() { } return dict } - -export default useDictionary diff --git a/lib/i18n/utils.ts b/lib/i18n/utils.ts index 122dd721..05a13236 100644 --- a/lib/i18n/utils.ts +++ b/lib/i18n/utils.ts @@ -10,5 +10,3 @@ export function formatMessage( return val === undefined ? match : String(val) }) } - -export default formatMessage diff --git a/lib/use-file-processor.tsx b/lib/use-file-processor.tsx index e4637577..6c3909c9 100644 --- a/lib/use-file-processor.tsx +++ b/lib/use-file-processor.tsx @@ -23,18 +23,11 @@ export interface FileProcessorMessages { readFailed: string } -const DEFAULT_MESSAGES: FileProcessorMessages = { - tooLong: "{name}: Content exceeds {limit}k character limit ({size}k chars)", - readFailed: "Failed to read file: {name}", -} - /** * Hook for processing file uploads, especially PDFs and text files. * Handles text extraction, character limit validation, and cleanup. */ -export function useFileProcessor( - messages: FileProcessorMessages = DEFAULT_MESSAGES, -) { +export function useFileProcessor(messages: FileProcessorMessages) { const [files, setFiles] = useState([]) const [pdfData, setPdfData] = useState>(new Map()) diff --git a/lib/utils.ts b/lib/utils.ts index 178404db..c2162baf 100644 --- a/lib/utils.ts +++ b/lib/utils.ts @@ -1,6 +1,10 @@ import { type ClassValue, clsx } from "clsx" -import * as pako from "pako" import { twMerge } from "tailwind-merge" +import { pageElement } from "@/lib/diagram-pages" +import { + decompressPageContent, + extractSvgEmbeddedXml, +} from "@/packages/mcp-server/src/load-diagram.ts" import { hasCells } from "@/packages/mcp-server/src/pages.ts" export function cn(...inputs: ClassValue[]) { @@ -100,197 +104,6 @@ export function formatXML(xml: string, indent: string = " "): string { return formatted.trim() } -/** - * Efficiently converts a potentially incomplete XML string to a legal XML string by closing any open tags properly. - * Additionally, if an tag does not have an mxGeometry child (e.g. ), - * it removes that tag from the output. - * Also removes orphaned elements that aren't inside or don't have proper 'as' attribute. - * @param xmlString The potentially incomplete XML string - * @returns A legal XML string with properly closed tags and removed incomplete mxCell elements. - */ -export function convertToLegalXml(xmlString: string): string { - // This regex will match either self-closing or a block element - // ... . Unfinished ones are left out because they don't match. - const regex = - /"']|"[^"]*"|'[^']*')*?(?:\/>|>([\s\S]*?)<\/mxCell>)/g - let match: RegExpExecArray | null - let result = "\n" - - while ((match = regex.exec(xmlString)) !== null) { - // match[0] contains the entire matched mxCell block - let cellContent = match[0] - - // Remove orphaned elements that are directly inside - // without an 'as' attribute (like as="sourcePoint", as="targetPoint") - // and not inside - // These cause "Could not add object mxPoint" errors in draw.io - // First check if there's an - if so, keep all mxPoints inside it - const hasArrayPoints = //.test(cellContent) - if (!hasArrayPoints) { - // Remove mxPoint elements without 'as' attribute - cellContent = cellContent.replace( - /]*\/>/g, - (pointMatch) => { - // Keep if it has an 'as' attribute - if (/\sas=/.test(pointMatch)) { - return pointMatch - } - // Remove orphaned mxPoint - return "" - }, - ) - } - - // Fix unescaped & characters in attribute values (but not valid entities) - // This prevents DOMParser from failing on content like "semantic & missing-step" - cellContent = cellContent.replace( - /&(?!(?:lt|gt|amp|quot|apos|#[0-9]+|#x[0-9a-fA-F]+);)/g, - "&", - ) - - // Fix unescaped < and > in attribute values for XML parsing - // HTML content in value attributes (e.g., Title) needs to be escaped - // This is critical because DOMParser will fail on unescaped < > in attributes - if (/=\s*"[^"]*<[^"]*"/.test(cellContent)) { - cellContent = cellContent.replace( - /=\s*"([^"]*)"/g, - (_match, value) => { - const escaped = value - .replace(//g, ">") - return `="${escaped}"` - }, - ) - } - - // Indent each line of the matched block for readability. - const formatted = cellContent - .split("\n") - .map((line) => " " + line.trim()) - .filter((line) => line.trim()) // Remove empty lines from removed mxPoints - .join("\n") - result += formatted + "\n" - } - result += "" - - return result -} - -/** - * Replace the cells of one page of a Draw.io XML document - * @param currentXML - The original Draw.io XML string - * @param nodes - The XML string containing new nodes to replace in the diagram - * @param pageId - The page whose cells are replaced; the first page when - * missing or not found - * @returns The updated XML string with replaced nodes - */ -export function replaceNodes( - currentXML: string, - nodes: string, - pageId?: string | null, -): string { - // Check for valid inputs - if (!currentXML || !nodes) { - throw new Error("Both currentXML and nodes must be provided") - } - - try { - // Parse the XML strings to create DOM objects - const parser = new DOMParser() - const currentDoc = parser.parseFromString(currentXML, "text/xml") - - // Handle nodes input - if it doesn't contain , wrap it - let nodesString = nodes - if (!nodes.includes("")) { - nodesString = `${nodes}` - } - - const nodesDoc = parser.parseFromString(nodesString, "text/xml") - - // The page to draw on, when the document has it - const page = pageId - ? Array.from(currentDoc.getElementsByTagName("diagram")).find( - (diagram) => diagram.getAttribute("id") === pageId, - ) - : undefined - const scope: ParentNode = page ?? currentDoc - // Find the root element in the current document - let currentRoot = scope.querySelector("mxGraphModel > root") - if (!currentRoot) { - // If no root element is found, create the proper structure - const mxGraphModel = - scope.querySelector("mxGraphModel") || - currentDoc.createElement("mxGraphModel") - - if (!currentDoc.contains(mxGraphModel)) { - ;(page ?? currentDoc).appendChild(mxGraphModel) - } - - currentRoot = currentDoc.createElement("root") - mxGraphModel.appendChild(currentRoot) - } - - // Find the root element in the nodes document - const nodesRoot = nodesDoc.querySelector("root") - if (!nodesRoot) { - throw new Error( - "Invalid nodes: Could not find or create element", - ) - } - - // Clear all existing child elements from the current root - while (currentRoot.firstChild) { - currentRoot.removeChild(currentRoot.firstChild) - } - - // Ensure the base cells exist - const hasCell0 = Array.from(nodesRoot.childNodes).some( - (node) => - node.nodeName === "mxCell" && - (node as Element).getAttribute("id") === "0", - ) - - const hasCell1 = Array.from(nodesRoot.childNodes).some( - (node) => - node.nodeName === "mxCell" && - (node as Element).getAttribute("id") === "1", - ) - - // Copy all child nodes from the nodes root to the current root - Array.from(nodesRoot.childNodes).forEach((node) => { - const importedNode = currentDoc.importNode(node, true) - currentRoot.appendChild(importedNode) - }) - - // Add default cells if they don't exist - if (!hasCell0) { - const cell0 = currentDoc.createElement("mxCell") - cell0.setAttribute("id", "0") - currentRoot.insertBefore(cell0, currentRoot.firstChild) - } - - if (!hasCell1) { - const cell1 = currentDoc.createElement("mxCell") - cell1.setAttribute("id", "1") - cell1.setAttribute("parent", "0") - - // Insert after cell0 if possible - const cell0 = currentRoot.querySelector('mxCell[id="0"]') - if (cell0?.nextSibling) { - currentRoot.insertBefore(cell1, cell0.nextSibling) - } else { - currentRoot.appendChild(cell1) - } - } - - // Convert the modified DOM back to a string - const serializer = new XMLSerializer() - return serializer.serializeToString(currentDoc) - } catch (error) { - throw new Error(`Error replacing nodes: ${error}`) - } -} - /** * Decode an xmlsvg export (SVG data URL) into uncompressed diagram XML. * One page is returned: the one with the given id, else the first; for the @@ -301,67 +114,26 @@ export function extractDiagramXML( pageId?: string | null, ): string { try { - // 1. Parse the SVG string (using built-in DOMParser in a browser-like environment) - const svgString = atob(xml_svg_string.slice(26)) - const parser = new DOMParser() - const svgDoc = parser.parseFromString(svgString, "image/svg+xml") - const svgElement = svgDoc.querySelector("svg") - - if (!svgElement) { + // The data URL holds the SVG; its root element's content attribute + // holds the mxfile, whose pages are compressed + const svg = atob(xml_svg_string.slice(26)) + const content = extractSvgEmbeddedXml(svg) + if (content === undefined) { throw new Error("No SVG element found in the input string.") } - // 2. Extract the 'content' attribute - const encodedContent = svgElement.getAttribute("content") - - if (!encodedContent) { + if (!content) { throw new Error("SVG element does not have a 'content' attribute.") } - - // 3. getAttribute already decoded the entities: decoding again - // would break values that hold them (a page named "R & D") - const xmlContent = encodedContent - - // 4. Parse the XML content - const xmlDoc = parser.parseFromString(xmlContent, "text/xml") - const diagrams = Array.from(xmlDoc.getElementsByTagName("diagram")) - const diagramElement = - (pageId && - diagrams.find( - (diagram) => diagram.getAttribute("id") === pageId, - )) || - diagrams[0] - - if (!diagramElement) { + const doc = new DOMParser().parseFromString(content, "text/xml") + const diagram = pageElement(doc, pageId) + if (!diagram) { throw new Error("No diagram element found") } - // 5. Extract base64 encoded data - const base64EncodedData = diagramElement.textContent - - if (!base64EncodedData) { + const xml = decompressPageContent(diagram.textContent ?? "") + if (!xml) { throw new Error("No encoded data found in the diagram element") } - - // 6. Decode base64 data - const binaryString = atob(base64EncodedData) - - // 7. Convert binary string to Uint8Array - const len = binaryString.length - const bytes = new Uint8Array(len) - for (let i = 0; i < len; i++) { - bytes[i] = binaryString.charCodeAt(i) - } - - // 8. Decompress data using pako (equivalent to zlib.decompress with wbits=-15) - const decompressedData = pako.inflate(bytes, { windowBits: -15 }) - - // 9. Convert the decompressed data to a string - const decoder = new TextDecoder("utf-8") - const decodedString = decoder.decode(decompressedData) - - // Decode URL-encoded content (equivalent to Python's urllib.parse.unquote) - const urlDecodedString = decodeURIComponent(decodedString) - - return urlDecodedString + return xml } catch (error) { console.error("Error extracting diagram XML:", error) throw error // Re-throw for caller handling diff --git a/tests/unit/use-diagram-streaming.test.tsx b/tests/unit/use-diagram-streaming.test.tsx index dd0e992f..c68c68a7 100644 --- a/tests/unit/use-diagram-streaming.test.tsx +++ b/tests/unit/use-diagram-streaming.test.tsx @@ -103,6 +103,35 @@ describe("the streaming preview of display_diagram", () => { expect(loaded).not.toContain("mxStyle") expect(loaded).not.toContain('id="b"') }) + + it("fixes an unescaped & and HTML in a label, as the final diagram does", () => { + mocks.loadDiagram.mockClear() + const xml = `` + renderHook(() => + useDiagramStreaming({ + messages: [ + { + id: "m1", + role: "assistant", + parts: [ + { + type: "tool-display_diagram", + toolCallId: "draw-2", + state: "input-streaming", + input: { xml }, + }, + ], + }, + ] as any, + processedToolCallsRef: { current: new Set() }, + editDiagramOriginalXmlRef: { current: new Map() }, + loadedMessageIdsRef: { current: new Set() }, + }), + ) + expect(mocks.loadDiagram).toHaveBeenCalledTimes(1) + const loaded = mocks.loadDiagram.mock.calls[0][0] + expect(loaded).toContain('value="R & D <b>team</b>"') + }) }) describe("the streaming preview of edit_diagram", () => { diff --git a/tests/unit/use-file-processor.test.tsx b/tests/unit/use-file-processor.test.tsx index 626c309c..6d937833 100644 --- a/tests/unit/use-file-processor.test.tsx +++ b/tests/unit/use-file-processor.test.tsx @@ -22,6 +22,8 @@ function deferred() { const pdfFile = () => new File(["%PDF"], "slow.pdf", { type: "application/pdf" }) const textFile = () => new File(["notes"], "notes.txt", { type: "text/plain" }) +// The toast texts the hook formats; toast is mocked, so any texts do +const messages = { tooLong: "{name} too long", readFailed: "failed {name}" } describe("useFileProcessor", () => { beforeEach(() => { @@ -35,7 +37,7 @@ describe("useFileProcessor", () => { vi.mocked(extractTextFileContent).mockResolvedValue("notes") const a = pdfFile() const b = textFile() - const { result } = renderHook(() => useFileProcessor()) + const { result } = renderHook(() => useFileProcessor(messages)) let done!: Promise act(() => { @@ -59,7 +61,7 @@ describe("useFileProcessor", () => { vi.mocked(extractTextFileContent).mockResolvedValue("notes") const a = pdfFile() const b = textFile() - const { result } = renderHook(() => useFileProcessor()) + const { result } = renderHook(() => useFileProcessor(messages)) let first!: Promise act(() => { @@ -82,7 +84,7 @@ describe("useFileProcessor", () => { const pdf = deferred() vi.mocked(extractPdfText).mockReturnValue(pdf.promise) const a = pdfFile() - const { result } = renderHook(() => useFileProcessor()) + const { result } = renderHook(() => useFileProcessor(messages)) let first!: Promise act(() => { diff --git a/tests/unit/utils.test.ts b/tests/unit/utils.test.ts index e0e89fdc..0426b20c 100644 --- a/tests/unit/utils.test.ts +++ b/tests/unit/utils.test.ts @@ -5,7 +5,6 @@ import { extractCompleteMxCells, extractDiagramXML, isRealDiagram, - replaceNodes, } from "@/lib/utils" import { BLANK_MXFILE } from "@/packages/mcp-server/src/pages.ts" @@ -70,39 +69,6 @@ describe("extractCompleteMxCells", () => { const pageBox = (id: string) => `` const pageModel = (cells: string) => `${cells}` -const pageOf = (id: string, cells: string) => - `${pageModel(cells)}` -/** Ids of the shapes on a page (not the root cells) */ -const shapesOn = (xml: string, pageId: string) => - Array.from( - new DOMParser() - .parseFromString(xml, "text/xml") - .querySelectorAll(`diagram[id="${pageId}"] mxCell`), - ) - .map((cell) => cell.getAttribute("id")) - .filter((id) => id !== "0" && id !== "1") - -describe("replaceNodes", () => { - const two = `${pageOf("a", pageBox("A"))}${pageOf("b", pageBox("B"))}` - - it("replaces the first page's cells by default", () => { - const out = replaceNodes(two, pageBox("N")) - expect(shapesOn(out, "a")).toEqual(["N"]) - expect(shapesOn(out, "b")).toEqual(["B"]) - }) - - it("replaces the given page's cells and keeps the others", () => { - const out = replaceNodes(two, pageBox("N"), "b") - expect(shapesOn(out, "a")).toEqual(["A"]) - expect(shapesOn(out, "b")).toEqual(["N"]) - }) - - it("uses the first page when the given page is not there", () => { - const out = replaceNodes(two, pageBox("N"), "zzz") - expect(shapesOn(out, "a")).toEqual(["N"]) - expect(shapesOn(out, "b")).toEqual(["B"]) - }) -}) describe("extractDiagramXML", () => { const packed = (model: string) =>