i/next-ai-draw-io
1
0
Fork 0
mirror of https://github.com/DayuanJiang/next-ai-draw-io.git synced 2026-10-12 12:39:52 +08:00
Code Issues Packages Projects Releases Wiki Activity

refactor(lib): reuse MCP helpers for streaming preview and SVG extraction

Browse Source 
- lib/utils.ts: remove convertToLegalXml and replaceNodes; the streaming
  preview now goes through wrapCellsInModel, autoFixXml and
  replacePageModel like the final commit path does
- lib/utils.ts: extractDiagramXML reuses extractSvgEmbeddedXml,
  pageElement and decompressPageContent instead of hand-written SVG
  parsing and pako inflation
- lib/use-file-processor.tsx: messages parameter is required, drop the
  English default texts only tests used
- hooks/use-dictionary.ts, lib/i18n/utils.ts: drop duplicate default
  exports nobody imported
- tests updated accordingly
This commit is contained in:
dayuan.jiang
2026-10-12 08:05:13 +09:00
parent 07f3e3c2a3
commit a3664367ce
8 changed files with 73 additions and 316 deletions
Download Patch File Download Diff File Expand all files Collapse all files
components/chat/use-diagram-streaming.ts
+21 -22
View File
@@ -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(
`<root>${convertedXml}</root>`,
"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)
hooks/use-dictionary.ts
-2
View File
@@ -25,5 +25,3 @@ export function useDictionary() {
}
return dict
}
export default useDictionary
lib/i18n/utils.ts
-2
View File
@@ -10,5 +10,3 @@ export function formatMessage(
return val === undefined ? match : String(val)
})
}
export default formatMessage
lib/use-file-processor.tsx
+1 -8
View File
@@ -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<File[]>([])
const [pdfData, setPdfData] = useState<Map<File, FileData>>(new Map())
lib/utils.ts
+17 -245
View File
@@ -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 <mxCell> tag does not have an mxGeometry child (e.g. <mxCell id="3">),
* it removes that tag from the output.
* Also removes orphaned <mxPoint> elements that aren't inside <Array> 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 <mxCell .../> or a block element
// <mxCell ...> ... </mxCell>. Unfinished ones are left out because they don't match.
const regex =
/<mxCell\b(?:[^<>"']|"[^"]*"|'[^']*')*?(?:\/>|>([\s\S]*?)<\/mxCell>)/g
let match: RegExpExecArray | null
let result = "<root>\n"
while ((match = regex.exec(xmlString)) !== null) {
// match[0] contains the entire matched mxCell block
let cellContent = match[0]
// Remove orphaned <mxPoint> elements that are directly inside <mxGeometry>
// without an 'as' attribute (like as="sourcePoint", as="targetPoint")
// and not inside <Array as="points">
// These cause "Could not add object mxPoint" errors in draw.io
// First check if there's an <Array as="points"> - if so, keep all mxPoints inside it
const hasArrayPoints = /<Array\s+as="points">/.test(cellContent)
if (!hasArrayPoints) {
// Remove mxPoint elements without 'as' attribute
cellContent = cellContent.replace(
/<mxPoint\b[^>]*\/>/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,
"&amp;",
)
// Fix unescaped < and > in attribute values for XML parsing
// HTML content in value attributes (e.g., <b>Title</b>) 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, "&lt;")
.replace(/>/g, "&gt;")
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 += "</root>"
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 <root>, wrap it
let nodesString = nodes
if (!nodes.includes("<root>")) {
nodesString = `<root>${nodes}</root>`
}
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 <root> 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 &amp; 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
tests/unit/use-diagram-streaming.test.tsx
+29
View File
@@ -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 = `<mxCell id="a" value="R & D <b>team</b>" vertex="1" parent="1"><mxGeometry x="0" y="0" width="80" height="40" as="geometry"/></mxCell>`
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 &amp; D &lt;b&gt;team&lt;/b&gt;"')
})
})
describe("the streaming preview of edit_diagram", () => {
tests/unit/use-file-processor.test.tsx
+5 -3
View File
@@ -22,6 +22,8 @@ function deferred<T>() {
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<void>
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<void>
act(() => {
@@ -82,7 +84,7 @@ describe("useFileProcessor", () => {
const pdf = deferred<string>()
vi.mocked(extractPdfText).mockReturnValue(pdf.promise)
const a = pdfFile()
const { result } = renderHook(() => useFileProcessor())
const { result } = renderHook(() => useFileProcessor(messages))
let first!: Promise<void>
act(() => {
tests/unit/utils.test.ts
-34
View File
@@ -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) => `<mxCell id="${id}" vertex="1" parent="1"/>`
const pageModel = (cells: string) =>
`<mxGraphModel><root><mxCell id="0"/><mxCell id="1" parent="0"/>${cells}</root></mxGraphModel>`
const pageOf = (id: string, cells: string) =>
`<diagram id="${id}" name="${id}">${pageModel(cells)}</diagram>`
/** 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 = `<mxfile>${pageOf("a", pageBox("A"))}${pageOf("b", pageBox("B"))}</mxfile>`
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) =>