i/next-ai-draw-io
1
0
Fork 0
mirror of https://github.com/DayuanJiang/next-ai-draw-io.git synced 2026-01-02 14:22:28 +08:00
Code Issues Packages Projects Releases Wiki Activity

fix: prevent browser crash during long streaming sessions (#262)

Browse Source 
- Debounce streaming diagram updates (150ms) to reduce handleDisplayChart calls by 93%
- Debounce localStorage writes (1s) to prevent blocking main thread
- Limit diagramHistory to 20 entries to prevent unbounded memory growth
- Clean up debounce timeout on component unmount to prevent memory leaks
- Add console timing markers for performance profiling

Fixes #78
This commit is contained in:
Dayuan Jiang
2025-12-14 21:23:14 +09:00
committed by GitHub
parent 55821301dd
commit 78a77e102d
4 changed files with 205 additions and 23 deletions
Download Patch File Download Diff File Expand all files Collapse all files

92
components/chat-message-display.tsx
View File

@@ -193,6 +193,14 @@ export function ChatMessageDisplay({
const messagesEndRef = useRef<HTMLDivElement>(null)
const previousXML = useRef<string>("")
const processedToolCalls = processedToolCallsRef
// Track the last processed XML per toolCallId to skip redundant processing during streaming
const lastProcessedXmlRef = useRef<Map<string, string>>(new Map())
// Debounce streaming diagram updates - store pending XML and timeout
const pendingXmlRef = useRef<string | null>(null)
const debounceTimeoutRef = useRef<ReturnType<typeof setTimeout> | null>(
null,
)
const STREAMING_DEBOUNCE_MS = 150 // Only update diagram every 150ms during streaming
const [expandedTools, setExpandedTools] = useState<Record<string, boolean>>(
{},
)
@@ -284,12 +292,15 @@ export function ChatMessageDisplay({
const handleDisplayChart = useCallback(
(xml: string, showToast = false) => {
console.time("perf:handleDisplayChart")
const currentXml = xml || ""
const convertedXml = convertToLegalXml(currentXml)
if (convertedXml !== previousXML.current) {
// Parse and validate XML BEFORE calling replaceNodes
console.time("perf:DOMParser")
const parser = new DOMParser()
const testDoc = parser.parseFromString(convertedXml, "text/xml")
console.timeEnd("perf:DOMParser")
const parseError = testDoc.querySelector("parsererror")
if (parseError) {
@@ -305,6 +316,7 @@ export function ChatMessageDisplay({
"AI generated invalid diagram XML. Please try regenerating.",
)
}
console.timeEnd("perf:handleDisplayChart")
return // Skip this update
}
@@ -314,10 +326,14 @@ export function ChatMessageDisplay({
const baseXML =
chartXML ||
`<mxfile><diagram name="Page-1" id="page-1"><mxGraphModel><root><mxCell id="0"/><mxCell id="1" parent="0"/></root></mxGraphModel></diagram></mxfile>`
console.time("perf:replaceNodes")
const replacedXML = replaceNodes(baseXML, convertedXml)
console.timeEnd("perf:replaceNodes")
// Validate and auto-fix the XML
console.time("perf:validateAndFixXml")
const validation = validateAndFixXml(replacedXML)
console.timeEnd("perf:validateAndFixXml")
if (validation.valid) {
previousXML.current = convertedXml
// Use fixed XML if available, otherwise use original
@@ -354,6 +370,9 @@ export function ChatMessageDisplay({
)
}
}
console.timeEnd("perf:handleDisplayChart")
} else {
console.timeEnd("perf:handleDisplayChart")
}
},
[chartXML, onDisplayChart],
@@ -372,7 +391,17 @@ export function ChatMessageDisplay({
}, [editingMessageId])
useEffect(() => {
messages.forEach((message) => {
console.time("perf:message-display-useEffect")
let processedCount = 0
let skippedCount = 0
let debouncedCount = 0
// Only process the last message for streaming performance
// Previous messages are already processed and won't change
const messagesToProcess =
messages.length > 0 ? [messages[messages.length - 1]] : []
messagesToProcess.forEach((message) => {
if (message.parts) {
message.parts.forEach((part) => {
if (part.type?.startsWith("tool-")) {
@@ -391,25 +420,82 @@ export function ChatMessageDisplay({
input?.xml
) {
const xml = input.xml as string
// Skip if XML hasn't changed since last processing
const lastXml =
lastProcessedXmlRef.current.get(toolCallId)
if (lastXml === xml) {
skippedCount++
return // Skip redundant processing
}
if (
state === "input-streaming" ||
state === "input-available"
) {
// During streaming, don't show toast (XML may be incomplete)
handleDisplayChart(xml, false)
// Debounce streaming updates - queue the XML and process after delay
pendingXmlRef.current = xml
if (!debounceTimeoutRef.current) {
// No pending timeout - set one up
debounceTimeoutRef.current = setTimeout(
() => {
const pendingXml =
pendingXmlRef.current
debounceTimeoutRef.current = null
pendingXmlRef.current = null
if (pendingXml) {
console.log(
"perf:debounced-handleDisplayChart executing",
)
handleDisplayChart(
pendingXml,
false,
)
lastProcessedXmlRef.current.set(
toolCallId,
pendingXml,
)
}
},
STREAMING_DEBOUNCE_MS,
)
}
debouncedCount++
} else if (
state === "output-available" &&
!processedToolCalls.current.has(toolCallId)
) {
// Final output - process immediately (clear any pending debounce)
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)
processedCount++
}
}
}
})
}
})
console.log(
`perf:message-display-useEffect processed=${processedCount} skipped=${skippedCount} debounced=${debouncedCount}`,
)
console.timeEnd("perf:message-display-useEffect")
// Cleanup: clear any pending debounce timeout on unmount
return () => {
if (debounceTimeoutRef.current) {
clearTimeout(debounceTimeoutRef.current)
debounceTimeoutRef.current = null
}
}
}, [messages, handleDisplayChart])
const renderToolPart = (part: ToolPartLike) => {

56
components/chat-panel.tsx
View File

@@ -233,6 +233,15 @@ export default function ChatPanel({
// Persist processed tool call IDs so collapsing the chat doesn't replay old tool outputs
const processedToolCallsRef = useRef<Set<string>>(new Set())
// Debounce timeout for localStorage writes (prevents blocking during streaming)
const localStorageDebounceRef = useRef<ReturnType<
typeof setTimeout
> | null>(null)
const xmlStorageDebounceRef = useRef<ReturnType<typeof setTimeout> | null>(
null,
)
const LOCAL_STORAGE_DEBOUNCE_MS = 1000 // Save at most once per second
const {
messages,
sendMessage,
@@ -728,32 +737,71 @@ Continue from EXACTLY where you stopped.`,
}, 500)
}, [isDrawioReady, onDisplayChart])
// Save messages to localStorage whenever they change
// Save messages to localStorage whenever they change (debounced to prevent blocking during streaming)
useEffect(() => {
if (!hasRestoredRef.current) return
// Clear any pending save
if (localStorageDebounceRef.current) {
clearTimeout(localStorageDebounceRef.current)
}
// Debounce: save after 1 second of no changes
localStorageDebounceRef.current = setTimeout(() => {
try {
localStorage.setItem(STORAGE_MESSAGES_KEY, JSON.stringify(messages))
console.time("perf:localStorage-messages")
localStorage.setItem(
STORAGE_MESSAGES_KEY,
JSON.stringify(messages),
)
console.timeEnd("perf:localStorage-messages")
} catch (error) {
console.error("Failed to save messages to localStorage:", error)
}
}, LOCAL_STORAGE_DEBOUNCE_MS)
// Cleanup on unmount
return () => {
if (localStorageDebounceRef.current) {
clearTimeout(localStorageDebounceRef.current)
}
}
}, [messages])
// Save diagram XML to localStorage whenever it changes
// Save diagram XML to localStorage whenever it changes (debounced)
useEffect(() => {
if (!canSaveDiagram) return
if (chartXML && chartXML.length > 300) {
if (!chartXML || chartXML.length <= 300) return
// Clear any pending save
if (xmlStorageDebounceRef.current) {
clearTimeout(xmlStorageDebounceRef.current)
}
// Debounce: save after 1 second of no changes
xmlStorageDebounceRef.current = setTimeout(() => {
console.time("perf:localStorage-xml")
localStorage.setItem(STORAGE_DIAGRAM_XML_KEY, chartXML)
console.timeEnd("perf:localStorage-xml")
}, LOCAL_STORAGE_DEBOUNCE_MS)
return () => {
if (xmlStorageDebounceRef.current) {
clearTimeout(xmlStorageDebounceRef.current)
}
}
}, [chartXML, canSaveDiagram])
// Save XML snapshots to localStorage whenever they change
const saveXmlSnapshots = useCallback(() => {
try {
console.time("perf:localStorage-snapshots")
const snapshotsArray = Array.from(xmlSnapshotsRef.current.entries())
localStorage.setItem(
STORAGE_XML_SNAPSHOTS_KEY,
JSON.stringify(snapshotsArray),
)
console.timeEnd("perf:localStorage-snapshots")
} catch (error) {
console.error(
"Failed to save XML snapshots to localStorage:",

17
contexts/diagram-context.tsx
View File

@@ -86,16 +86,20 @@ export function DiagramProvider({ children }: { children: React.ReactNode }) {
chart: string,
skipValidation?: boolean,
): string | null => {
console.time("perf:loadDiagram")
let xmlToLoad = chart
// Validate XML structure before loading (unless skipped for internal use)
if (!skipValidation) {
console.time("perf:loadDiagram-validation")
const validation = validateAndFixXml(chart)
console.timeEnd("perf:loadDiagram-validation")
if (!validation.valid) {
console.warn(
"[loadDiagram] Validation error:",
validation.error,
)
console.timeEnd("perf:loadDiagram")
return validation.error
}
// Use fixed XML if auto-fix was applied
@@ -112,11 +116,14 @@ export function DiagramProvider({ children }: { children: React.ReactNode }) {
setChartXML(xmlToLoad)
if (drawioRef.current) {
console.time("perf:drawio-iframe-load")
drawioRef.current.load({
xml: xmlToLoad,
})
console.timeEnd("perf:drawio-iframe-load")
}
console.timeEnd("perf:loadDiagram")
return null
}
@@ -138,14 +145,20 @@ export function DiagramProvider({ children }: { children: React.ReactNode }) {
setLatestSvg(data.data)
// Only add to history if this was a user-initiated export
// Limit to 20 entries to prevent memory leaks during long sessions
const MAX_HISTORY_SIZE = 20
if (expectHistoryExportRef.current) {
setDiagramHistory((prev) => [
setDiagramHistory((prev) => {
const newHistory = [
...prev,
{
svg: data.data,
xml: extractedXML,
},
])
]
// Keep only the last MAX_HISTORY_SIZE entries (circular buffer)
return newHistory.slice(-MAX_HISTORY_SIZE)
})
expectHistoryExportRef.current = false
}

47
lib/utils.ts
View File

@@ -823,6 +823,8 @@ function checkNestedMxCells(xml: string): string | null {
* @returns null if valid, error message string if invalid
*/
export function validateMxCellStructure(xml: string): string | null {
console.time("perf:validateMxCellStructure")
console.log(`perf:validateMxCellStructure XML size: ${xml.length} bytes`)
// Size check for performance
if (xml.length > MAX_XML_SIZE) {
console.warn(
@@ -832,8 +834,10 @@ export function validateMxCellStructure(xml: string): string | null {
// 0. First use DOM parser to catch syntax errors (most accurate)
try {
console.time("perf:validate-DOMParser")
const parser = new DOMParser()
const doc = parser.parseFromString(xml, "text/xml")
console.timeEnd("perf:validate-DOMParser")
const parseError = doc.querySelector("parsererror")
if (parseError) {
const actualError = parseError.textContent || "Unknown parse error"
@@ -841,6 +845,7 @@ export function validateMxCellStructure(xml: string): string | null {
"[validateMxCellStructure] DOMParser error:",
actualError,
)
console.timeEnd("perf:validateMxCellStructure")
return `Invalid XML: The XML contains syntax errors (likely unescaped special characters like <, >, & in attribute values). Please escape special characters: use &lt; for <, &gt; for >, &amp; for &, &quot; for ". Regenerate the diagram with properly escaped values.`
}
@@ -849,6 +854,7 @@ export function validateMxCellStructure(xml: string): string | null {
for (const cell of allCells) {
if (cell.parentElement?.tagName === "mxCell") {
const id = cell.getAttribute("id") || "unknown"
console.timeEnd("perf:validateMxCellStructure")
return `Invalid XML: Found nested mxCell (id="${id}"). Cells should be siblings, not nested inside other mxCell elements.`
}
}
@@ -862,12 +868,18 @@ export function validateMxCellStructure(xml: string): string | null {
// 1. Check for CDATA wrapper (invalid at document root)
if (/^\s*<!\[CDATA\[/.test(xml)) {
console.timeEnd("perf:validateMxCellStructure")
return "Invalid XML: XML is wrapped in CDATA section - remove <![CDATA[ from start and ]]> from end"
}
// 2. Check for duplicate structural attributes
console.time("perf:checkDuplicateAttributes")
const dupAttrError = checkDuplicateAttributes(xml)
if (dupAttrError) return dupAttrError
console.timeEnd("perf:checkDuplicateAttributes")
if (dupAttrError) {
console.timeEnd("perf:validateMxCellStructure")
return dupAttrError
}
// 3. Check for unescaped < in attribute values
const attrValuePattern = /=\s*"([^"]*)"/g
@@ -875,44 +887,67 @@ export function validateMxCellStructure(xml: string): string | null {
while ((attrValMatch = attrValuePattern.exec(xml)) !== null) {
const value = attrValMatch[1]
if (/</.test(value) && !/&lt;/.test(value)) {
console.timeEnd("perf:validateMxCellStructure")
return "Invalid XML: Unescaped < character in attribute values. Replace < with &lt;"
}
}
// 4. Check for duplicate IDs
console.time("perf:checkDuplicateIds")
const dupIdError = checkDuplicateIds(xml)
if (dupIdError) return dupIdError
console.timeEnd("perf:checkDuplicateIds")
if (dupIdError) {
console.timeEnd("perf:validateMxCellStructure")
return dupIdError
}
// 5. Check for tag mismatches
console.time("perf:checkTagMismatches")
const tagMismatchError = checkTagMismatches(xml)
if (tagMismatchError) return tagMismatchError
console.timeEnd("perf:checkTagMismatches")
if (tagMismatchError) {
console.timeEnd("perf:validateMxCellStructure")
return tagMismatchError
}
// 6. Check invalid character references
const charRefError = checkCharacterReferences(xml)
if (charRefError) return charRefError
if (charRefError) {
console.timeEnd("perf:validateMxCellStructure")
return charRefError
}
// 7. Check for invalid comment syntax (-- inside comments)
const commentPattern = /<!--([\s\S]*?)-->/g
let commentMatch
while ((commentMatch = commentPattern.exec(xml)) !== null) {
if (/--/.test(commentMatch[1])) {
console.timeEnd("perf:validateMxCellStructure")
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
if (entityError) {
console.timeEnd("perf:validateMxCellStructure")
return entityError
}
// 9. Check for empty id attributes on mxCell
if (/<mxCell[^>]*\sid\s*=\s*["']\s*["'][^>]*>/g.test(xml)) {
console.timeEnd("perf:validateMxCellStructure")
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
if (nestedCellError) {
console.timeEnd("perf:validateMxCellStructure")
return nestedCellError
}
console.timeEnd("perf:validateMxCellStructure")
return null
}