fix: Address multiple memory leaks in CodeBlock component

webview-ui/src/components/common/CodeBlock.tsx

@@ -232,6 +232,10 @@ const CodeBlock = memo(
 		const copyButtonWrapperRef = useRef<HTMLDivElement>(null)
 		const { showCopyFeedback, copyWithFeedback } = useCopyToClipboard()
 		const { t } = useAppTranslation()
+		const isMountedRef = useRef(true)
+		const buttonPositionTimeoutRef = useRef<NodeJS.Timeout | null>(null)
+		const collapseTimeout1Ref = useRef<NodeJS.Timeout | null>(null)
+		const collapseTimeout2Ref = useRef<NodeJS.Timeout | null>(null)
 
 		// Update current language when prop changes, but only if user hasn't
 		// made a selection.
@@ -243,17 +247,23 @@ const CodeBlock = memo(
 			}
 		}, [language, currentLanguage])
 
-		// Syntax highlighting with cached Shiki instance.
+		// Syntax highlighting with cached Shiki instance and mounted state management
 		useEffect(() => {
+			// Set mounted state at the beginning of this effect
+			isMountedRef.current = true
+
 			const fallback = `<pre style="padding: 0; margin: 0;"><code class="hljs language-${currentLanguage || "txt"}">${source || ""}</code></pre>`
 
 			const highlight = async () => {
 				// Show plain text if language needs to be loaded.
 				if (currentLanguage && !isLanguageLoaded(currentLanguage)) {
-					setHighlightedCode(fallback)
+					if (isMountedRef.current) {
+						setHighlightedCode(fallback)
+					}
 				}
 
 				const highlighter = await getHighlighter(currentLanguage)
+				if (!isMountedRef.current) return
 
 				const html = await highlighter.codeToHtml(source || "", {
 					lang: currentLanguage || "txt",
@@ -277,14 +287,36 @@ const CodeBlock = memo(
 						},
 					] as ShikiTransformer[],
 				})
+				if (!isMountedRef.current) return
 
-				setHighlightedCode(html)
+				if (isMountedRef.current) {
+					setHighlightedCode(html)
+				}
 			}
 
 			highlight().catch((e) => {
 				console.error("[CodeBlock] Syntax highlighting error:", e, "\nStack trace:", e.stack)
-				setHighlightedCode(fallback)
+				if (isMountedRef.current) {
+					setHighlightedCode(fallback)
+				}
 			})
+
+			// Cleanup function - manage mounted state and clear all timeouts
+			return () => {
+				isMountedRef.current = false
+				if (buttonPositionTimeoutRef.current) {
+					clearTimeout(buttonPositionTimeoutRef.current)
+					buttonPositionTimeoutRef.current = null
+				}
+				if (collapseTimeout1Ref.current) {
+					clearTimeout(collapseTimeout1Ref.current)
+					collapseTimeout1Ref.current = null
+				}
+				if (collapseTimeout2Ref.current) {
+					clearTimeout(collapseTimeout2Ref.current)
+					collapseTimeout2Ref.current = null
+				}
+			}
 		}, [source, currentLanguage, collapsedHeight])
 
 		// Check if content height exceeds collapsed height whenever content changes
@@ -455,8 +487,15 @@ const CodeBlock = memo(
 		// Update button position and scroll when highlightedCode changes
 		useEffect(() => {
 			if (highlightedCode) {
+				// Clear any existing timeout before setting a new one
+				if (buttonPositionTimeoutRef.current) {
+					clearTimeout(buttonPositionTimeoutRef.current)
+				}
 				// Update button position
-				setTimeout(updateCodeBlockButtonPosition, 0)
+				buttonPositionTimeoutRef.current = setTimeout(() => {
+					updateCodeBlockButtonPosition()
+					buttonPositionTimeoutRef.current = null // Optional: Clear ref after execution
+				}, 0)
 
 				// Scroll to bottom if needed (immediately after Shiki updates)
 				if (shouldScrollAfterHighlightRef.current) {
@@ -479,6 +518,12 @@ const CodeBlock = memo(
 					shouldScrollAfterHighlightRef.current = false
 				}
 			}
+			// Cleanup function for this effect
+			return () => {
+				if (buttonPositionTimeoutRef.current) {
+					clearTimeout(buttonPositionTimeoutRef.current)
+				}
+			}
 		}, [highlightedCode, updateCodeBlockButtonPosition])
 
 		// Advanced inertial scroll chaining
@@ -682,23 +727,30 @@ const CodeBlock = memo(
 						{showCollapseButton && (
 							<CodeBlockButton
 								onClick={() => {
-									// Get the current code block element and scrollable container
-									const codeBlock = codeBlockRef.current
-									const scrollContainer = document.querySelector('[data-virtuoso-scroller="true"]')
-									if (!codeBlock || !scrollContainer) return
-
+									// Get the current code block element
+									const codeBlock = codeBlockRef.current // Capture ref early
 									// Toggle window shade state
 									setWindowShade(!windowShade)
 
+									// Clear any previous timeouts
+									if (collapseTimeout1Ref.current) clearTimeout(collapseTimeout1Ref.current)
+									if (collapseTimeout2Ref.current) clearTimeout(collapseTimeout2Ref.current)
+
 									// After UI updates, ensure code block is visible and update button position
-									setTimeout(
+									collapseTimeout1Ref.current = setTimeout(
 										() => {
-											codeBlock.scrollIntoView({ behavior: "smooth", block: "nearest" })
-
-											// Wait for scroll to complete before updating button position
-											setTimeout(() => {
-												updateCodeBlockButtonPosition()
-											}, 50)
+											if (codeBlock) {
+												// Check if codeBlock element still exists
+												codeBlock.scrollIntoView({ behavior: "smooth", block: "nearest" })
+
+												// Wait for scroll to complete before updating button position
+												collapseTimeout2Ref.current = setTimeout(() => {
+													// updateCodeBlockButtonPosition itself should also check for refs if needed
+													updateCodeBlockButtonPosition()
+													collapseTimeout2Ref.current = null
+												}, 50)
+											}
+											collapseTimeout1Ref.current = null
 										},
 										WINDOW_SHADE_SETTINGS.transitionDelayS * 1000 + 50,
 									)