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 // ============================================================================