diff --git a/components/chat-message-display.tsx b/components/chat-message-display.tsx
index 68d72ae..b094f8e 100644
--- a/components/chat-message-display.tsx
+++ b/components/chat-message-display.tsx
@@ -31,6 +31,7 @@ import { getApiEndpoint } from "@/lib/base-path"
import {
applyDiagramOperations,
convertToLegalXml,
+ extractCompleteMxCells,
isMxCellXmlComplete,
replaceNodes,
validateAndFixXml,
@@ -315,12 +316,28 @@ export function ChatMessageDisplay({
const handleDisplayChart = useCallback(
(xml: string, showToast = false) => {
- const currentXml = xml || ""
+ let currentXml = xml || ""
+ const startTime = performance.now()
+
+ // 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
+ }
+
const convertedXml = convertToLegalXml(currentXml)
if (convertedXml !== previousXML.current) {
// Parse and validate XML BEFORE calling replaceNodes
const parser = new DOMParser()
- const testDoc = parser.parseFromString(convertedXml, "text/xml")
+ // Wrap in root element for parsing multiple mxCell elements
+ const testDoc = parser.parseFromString(
+ `${convertedXml}`,
+ "text/xml",
+ )
const parseError = testDoc.querySelector("parsererror")
if (parseError) {
@@ -347,7 +364,22 @@ export function ChatMessageDisplay({
``
const replacedXML = replaceNodes(baseXML, convertedXml)
- // Validate and auto-fix the XML
+ const xmlProcessTime = performance.now() - startTime
+
+ // 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
+ const loadStartTime = performance.now()
+ onDisplayChart(replacedXML, true)
+ console.log(
+ `[Streaming] XML processing: ${xmlProcessTime.toFixed(1)}ms, drawio load: ${(performance.now() - loadStartTime).toFixed(1)}ms`,
+ )
+ return
+ }
+
+ // Final output: run full validation and auto-fix
const validation = validateAndFixXml(replacedXML)
if (validation.valid) {
previousXML.current = convertedXml
@@ -360,18 +392,19 @@ export function ChatMessageDisplay({
)
}
// Skip validation in loadDiagram since we already validated above
+ const loadStartTime = performance.now()
onDisplayChart(xmlToLoad, true)
+ console.log(
+ `[Final] XML processing: ${xmlProcessTime.toFixed(1)}ms, validation+load: ${(performance.now() - loadStartTime).toFixed(1)}ms`,
+ )
} else {
console.error(
"[ChatMessageDisplay] XML validation failed:",
validation.error,
)
- // Only show toast if this is the final XML (not during streaming)
- if (showToast) {
- toast.error(
- "Diagram validation failed. Please try regenerating.",
- )
- }
+ toast.error(
+ "Diagram validation failed. Please try regenerating.",
+ )
}
} catch (error) {
console.error(
@@ -603,17 +636,10 @@ export function ChatMessageDisplay({
}
})
- // Cleanup: clear any pending debounce timeout on unmount
- return () => {
- if (debounceTimeoutRef.current) {
- clearTimeout(debounceTimeoutRef.current)
- debounceTimeoutRef.current = null
- }
- if (editDebounceTimeoutRef.current) {
- clearTimeout(editDebounceTimeoutRef.current)
- editDebounceTimeoutRef.current = null
- }
- }
+ // NOTE: Don't cleanup debounce timeouts here!
+ // The cleanup runs on every re-render (when messages changes),
+ // which would cancel the timeout before it fires.
+ // Let the timeouts complete naturally - they're harmless if component unmounts.
}, [messages, handleDisplayChart, chartXML])
const renderToolPart = (part: ToolPartLike) => {
diff --git a/lib/utils.ts b/lib/utils.ts
index 1266d10..6ee7ac2 100644
--- a/lib/utils.ts
+++ b/lib/utils.ts
@@ -61,6 +61,47 @@ export function isMxCellXmlComplete(xml: string | undefined | null): boolean {
return trimmed.endsWith("/>") || trimmed.endsWith("")
}
+/**
+ * Extract only complete mxCell elements from partial/streaming XML.
+ * This allows progressive rendering during streaming by ignoring incomplete trailing elements.
+ * @param xml - The partial XML string (may contain incomplete trailing mxCell)
+ * @returns XML string containing only complete mxCell elements
+ */
+export function extractCompleteMxCells(xml: string | undefined | null): string {
+ if (!xml) return ""
+
+ const completeCells: Array<{ index: number; text: string }> = []
+
+ // Match self-closing mxCell tags:
+ // Also match mxCell with nested mxGeometry: ...
+ const selfClosingPattern = /]*\/>/g
+ const nestedPattern = /]*>[\s\S]*?<\/mxCell>/g
+
+ // Find all self-closing mxCell elements
+ let match: RegExpExecArray | null
+ while ((match = selfClosingPattern.exec(xml)) !== null) {
+ completeCells.push({ index: match.index, text: match[0] })
+ }
+
+ // Find all mxCell elements with nested content (like mxGeometry)
+ while ((match = nestedPattern.exec(xml)) !== null) {
+ completeCells.push({ index: match.index, text: match[0] })
+ }
+
+ // Sort by position to maintain order
+ completeCells.sort((a, b) => a.index - b.index)
+
+ // Remove duplicates (a self-closing match might overlap with nested match)
+ const seen = new Set()
+ const uniqueCells = completeCells.filter((cell) => {
+ if (seen.has(cell.index)) return false
+ seen.add(cell.index)
+ return true
+ })
+
+ return uniqueCells.map((c) => c.text).join("\n")
+}
+
// ============================================================================
// XML Parsing Helpers
// ============================================================================