feat: implement virtualization for ChatView to handle long conversations

webview-ui/src/components/chat/ChatView.tsx

@@ -1,7 +1,7 @@
 import { forwardRef, useCallback, useEffect, useImperativeHandle, useMemo, useRef, useState } from "react"
 import { useDeepCompareEffect, useEvent, useMount } from "react-use"
 import debounce from "debounce"
-import { Virtuoso, type VirtuosoHandle } from "react-virtuoso"
+import { Virtuoso } from "react-virtuoso"
 import removeMd from "remove-markdown"
 import { VSCodeButton } from "@vscode/webview-ui-toolkit/react"
 import useSound from "use-sound"
@@ -58,6 +58,9 @@ import QueuedMessages from "./QueuedMessages"
 import { getLatestTodo } from "@roo/todo"
 import { QueuedMessage } from "@roo-code/types"
 
+// Import the new virtualization utilities
+import { useOptimizedVirtualization } from "./virtualization"
+
 export interface ChatViewProps {
 	isHidden: boolean
 	showAnnouncement: boolean
@@ -167,13 +170,7 @@ const ChatViewComponent: React.ForwardRefRenderFunction<ChatViewRef, ChatViewPro
 	const [primaryButtonText, setPrimaryButtonText] = useState<string | undefined>(undefined)
 	const [secondaryButtonText, setSecondaryButtonText] = useState<string | undefined>(undefined)
 	const [didClickCancel, setDidClickCancel] = useState(false)
-	const virtuosoRef = useRef<VirtuosoHandle>(null)
-	const [expandedRows, setExpandedRows] = useState<Record<number, boolean>>({})
-	const prevExpandedRowsRef = useRef<Record<number, boolean>>()
 	const scrollContainerRef = useRef<HTMLDivElement>(null)
-	const disableAutoScrollRef = useRef(false)
-	const [showScrollToBottom, setShowScrollToBottom] = useState(false)
-	const [isAtBottom, setIsAtBottom] = useState(false)
 	const lastTtsRef = useRef<string>("")
 	const [wasStreaming, setWasStreaming] = useState<boolean>(false)
 	const [showCheckpointWarning, setShowCheckpointWarning] = useState<boolean>(false)
@@ -189,6 +186,12 @@ const ChatViewComponent: React.ForwardRefRenderFunction<ChatViewRef, ChatViewPro
 	const userRespondedRef = useRef<boolean>(false)
 	const [currentFollowUpTs, setCurrentFollowUpTs] = useState<number | null>(null)
 
+	// Map the optimized state to the existing state variables for backward compatibility
+	const [expandedRows, setExpandedRows] = useState<Record<number, boolean>>({})
+	const prevExpandedRowsRef = useRef<Record<number, boolean>>()
+	const disableAutoScrollRef = useRef(false)
+	const [showScrollToBottom, setShowScrollToBottom] = useState(false)
+
 	const clineAskRef = useRef(clineAsk)
 	useEffect(() => {
 		clineAskRef.current = clineAsk
@@ -434,6 +437,7 @@ const ChatViewComponent: React.ForwardRefRenderFunction<ChatViewRef, ChatViewPro
 	}, [messages.length])
 
 	useEffect(() => {
+		// Clear expanded rows for new task
 		setExpandedRows({})
 		everVisibleMessagesTsRef.current.clear() // Clear for new task
 		setCurrentFollowUpTs(null) // Clear follow-up answered state for new task
@@ -480,6 +484,7 @@ const ChatViewComponent: React.ForwardRefRenderFunction<ChatViewRef, ChatViewPro
 		prevExpandedRowsRef.current = expandedRows // Store current state for next comparison
 	}, [expandedRows])
 
+	// Define isStreaming before using it in the virtualization hook
 	const isStreaming = useMemo(() => {
 		// Checking clineAsk isn't enough since messages effect may be called
 		// again for a tool for example, set clineAsk to its value, and if the
@@ -521,6 +526,88 @@ const ChatViewComponent: React.ForwardRefRenderFunction<ChatViewRef, ChatViewPro
 		return false
 	}, [modifiedMessages, clineAsk, enableButtons, primaryButtonText])
 
+	// Use the optimized virtualization hook (after isStreaming is defined)
+	const {
+		virtuosoRef,
+		viewportConfig,
+		stateManager,
+		scrollManager,
+		performanceMonitor,
+		handleScroll: handleVirtuosoScroll,
+		handleRangeChange,
+		handleScrollStateChange,
+		scrollToBottom: optimizedScrollToBottom,
+		isAtBottom,
+		showScrollToBottom: shouldShowScrollButton,
+		visibleRange,
+	} = useOptimizedVirtualization({
+		messages: modifiedMessages,
+		isStreaming,
+		isHidden,
+		onPerformanceIssue: (metric, value) => {
+			console.warn(`ChatView performance issue: ${metric} = ${value}`)
+		},
+	})
+
+	// Sync expanded rows with state manager
+	useEffect(() => {
+		if (stateManager) {
+			const newExpandedRows: Record<number, boolean> = {}
+			modifiedMessages.forEach((msg) => {
+				if (stateManager.isExpanded(msg.ts)) {
+					newExpandedRows[msg.ts] = true
+				}
+			})
+			setExpandedRows(newExpandedRows)
+		}
+	}, [modifiedMessages, stateManager, visibleRange])
+
+	// Sync scroll button visibility
+	useEffect(() => {
+		setShowScrollToBottom(shouldShowScrollButton || (disableAutoScrollRef.current && !isAtBottom))
+	}, [shouldShowScrollButton, isAtBottom])
+
+	// Clear state manager for new task
+	useEffect(() => {
+		if (stateManager) {
+			stateManager.clear()
+		}
+		if (scrollManager) {
+			scrollManager.reset()
+		}
+	}, [task?.ts, stateManager, scrollManager])
+
+	// Handle performance monitoring
+	useEffect(() => {
+		if (performanceMonitor) {
+			if (isHidden) {
+				performanceMonitor.stopMonitoring()
+			} else {
+				performanceMonitor.startMonitoring()
+			}
+		}
+	}, [isHidden, performanceMonitor])
+
+	// Update scroll manager when user expands rows
+	useEffect(() => {
+		const prev = prevExpandedRowsRef.current
+		let wasAnyRowExpandedByUser = false
+		if (prev) {
+			// Check if any row transitioned from false/undefined to true
+			for (const [tsKey, isExpanded] of Object.entries(expandedRows)) {
+				const ts = Number(tsKey)
+				if (isExpanded && !(prev[ts] ?? false)) {
+					wasAnyRowExpandedByUser = true
+					break
+				}
+			}
+		}
+
+		if (wasAnyRowExpandedByUser && scrollManager) {
+			scrollManager.forceUserScrolling()
+		}
+	}, [expandedRows, scrollManager])
+
 	const markFollowUpAsAnswered = useCallback(() => {
 		const lastFollowUpMessage = messagesRef.current.findLast((msg) => msg.ask === "followup")
 		if (lastFollowUpMessage) {
@@ -549,6 +636,14 @@ const ChatViewComponent: React.ForwardRefRenderFunction<ChatViewRef, ChatViewPro
 		disableAutoScrollRef.current = false
 	}, [])
 
+	// Reset scroll manager state on chat reset
+	const handleChatResetWithManagers = useCallback(() => {
+		handleChatReset()
+		if (scrollManager) {
+			scrollManager.resetUserScrolling()
+		}
+	}, [handleChatReset, scrollManager])
+
 	/**
 	 * Handles sending messages to the extension
 	 * @param text - The message text to send
@@ -607,7 +702,7 @@ const ChatViewComponent: React.ForwardRefRenderFunction<ChatViewRef, ChatViewPro
 						vscode.postMessage({ type: "askResponse", askResponse: "messageResponse", text, images })
 					}
 
-					handleChatReset()
+					handleChatResetWithManagers()
 				}
 			} catch (error) {
 				console.error("Error in handleSendMessage:", error)
@@ -619,7 +714,7 @@ const ChatViewComponent: React.ForwardRefRenderFunction<ChatViewRef, ChatViewPro
 				// but for now we'll just log it
 			}
 		},
-		[handleChatReset, markFollowUpAsAnswered, sendingDisabled], // messagesRef and clineAskRef are stable
+		[handleChatResetWithManagers, markFollowUpAsAnswered, sendingDisabled], // messagesRef and clineAskRef are stable
 	)
 
 	useEffect(() => {
@@ -842,7 +937,7 @@ const ChatViewComponent: React.ForwardRefRenderFunction<ChatViewRef, ChatViewPro
 				case "invoke":
 					switch (message.invoke!) {
 						case "newChat":
-							handleChatReset()
+							handleChatResetWithManagers()
 							break
 						case "sendMessage":
 							handleSendMessage(message.text ?? "", message.images ?? [])
@@ -877,7 +972,7 @@ const ChatViewComponent: React.ForwardRefRenderFunction<ChatViewRef, ChatViewPro
 			sendingDisabled,
 			enableButtons,
 			currentTaskItem,
-			handleChatReset,
+			handleChatResetWithManagers,
 			handleSendMessage,
 			handleSetChatBoxMessage,
 			handlePrimaryButtonClick,
@@ -1306,14 +1401,13 @@ const ChatViewComponent: React.ForwardRefRenderFunction<ChatViewRef, ChatViewPro
 		return result
 	}, [isCondensing, visibleMessages])
 
-	// scrolling
-
+	// scrolling - use the optimized versions
 	const scrollToBottomSmooth = useMemo(
 		() =>
-			debounce(() => virtuosoRef.current?.scrollTo({ top: Number.MAX_SAFE_INTEGER, behavior: "smooth" }), 10, {
+			debounce(() => optimizedScrollToBottom("smooth"), 10, {
 				immediate: true,
 			}),
-		[],
+		[optimizedScrollToBottom],
 	)
 
 	useEffect(() => {
@@ -1325,17 +1419,21 @@ const ChatViewComponent: React.ForwardRefRenderFunction<ChatViewRef, ChatViewPro
 	}, [scrollToBottomSmooth])
 
 	const scrollToBottomAuto = useCallback(() => {
-		virtuosoRef.current?.scrollTo({
-			top: Number.MAX_SAFE_INTEGER,
-			behavior: "auto", // Instant causes crash.
-		})
-	}, [])
+		optimizedScrollToBottom("auto")
+	}, [optimizedScrollToBottom])
 
 	const handleSetExpandedRow = useCallback(
 		(ts: number, expand?: boolean) => {
-			setExpandedRows((prev) => ({ ...prev, [ts]: expand === undefined ? !prev[ts] : expand }))
+			if (stateManager) {
+				const newExpanded = expand === undefined ? !stateManager.isExpanded(ts) : expand
+				stateManager.setState(ts, { isExpanded: newExpanded })
+				setExpandedRows((prev) => ({ ...prev, [ts]: newExpanded }))
+			} else {
+				// Fallback to local state if stateManager not ready
+				setExpandedRows((prev) => ({ ...prev, [ts]: expand === undefined ? !prev[ts] : expand }))
+			}
 		},
-		[setExpandedRows], // setExpandedRows is stable
+		[stateManager],
 	)
 
 	// Scroll when user toggles certain rows.
@@ -1363,26 +1461,37 @@ const ChatViewComponent: React.ForwardRefRenderFunction<ChatViewRef, ChatViewPro
 
 	useEffect(() => {
 		let timer: NodeJS.Timeout | undefined
-		if (!disableAutoScrollRef.current) {
+		if (
+			!disableAutoScrollRef.current &&
+			scrollManager &&
+			stateManager &&
+			scrollManager.shouldAutoScroll(stateManager.hasExpandedMessages())
+		) {
 			timer = setTimeout(() => scrollToBottomSmooth(), 50)
 		}
 		return () => {
 			if (timer) {
 				clearTimeout(timer)
 			}
 		}
-	}, [groupedMessages.length, scrollToBottomSmooth])
-
-	const handleWheel = useCallback((event: Event) => {
-		const wheelEvent = event as WheelEvent
-
-		if (wheelEvent.deltaY && wheelEvent.deltaY < 0) {
-			if (scrollContainerRef.current?.contains(wheelEvent.target as Node)) {
-				// User scrolled up
-				disableAutoScrollRef.current = true
+	}, [groupedMessages.length, scrollToBottomSmooth, scrollManager, stateManager])
+
+	const handleWheel = useCallback(
+		(event: Event) => {
+			const wheelEvent = event as WheelEvent
+
+			if (wheelEvent.deltaY && wheelEvent.deltaY < 0) {
+				if (scrollContainerRef.current?.contains(wheelEvent.target as Node)) {
+					// User scrolled up
+					disableAutoScrollRef.current = true
+					if (scrollManager) {
+						scrollManager.forceUserScrolling()
+					}
+				}
 			}
-		}
-	}, [])
+		},
+		[scrollManager],
+	)
 
 	useEvent("wheel", handleWheel, window, { passive: true }) // passive improves scrolling performance
 
@@ -1844,16 +1953,27 @@ const ChatViewComponent: React.ForwardRefRenderFunction<ChatViewRef, ChatViewPro
 							ref={virtuosoRef}
 							key={task.ts} // trick to make sure virtuoso re-renders when task changes, and we use initialTopMostItemIndex to start at the bottom
 							className="scrollable grow overflow-y-scroll mb-1"
-							// increasing top by 3_000 to prevent jumping around when user collapses a row
-							increaseViewportBy={{ top: 3_000, bottom: Number.MAX_SAFE_INTEGER }} // hack to make sure the last message is always rendered to get truly perfect scroll to bottom animation when new messages are added (Number.MAX_SAFE_INTEGER is safe for arithmetic operations, which is all virtuoso uses this value for in src/sizeRangeSystem.ts)
+							// Use dynamic viewport configuration based on device and state
+							increaseViewportBy={viewportConfig}
 							data={groupedMessages} // messages is the raw format returned by extension, modifiedMessages is the manipulated structure that combines certain messages of related type, and visibleMessages is the filtered structure that removes messages that should not be rendered
 							itemContent={itemContent}
-							atBottomStateChange={(isAtBottom) => {
-								setIsAtBottom(isAtBottom)
-								if (isAtBottom) {
+							onScroll={(e) => {
+								const target = e.currentTarget as HTMLElement
+								handleVirtuosoScroll(target.scrollTop)
+								handleScrollStateChange({
+									scrollTop: target.scrollTop,
+									scrollHeight: target.scrollHeight,
+									viewportHeight: target.clientHeight,
+								})
+							}}
+							rangeChanged={handleRangeChange}
+							atBottomStateChange={(atBottom) => {
+								if (atBottom) {
 									disableAutoScrollRef.current = false
+									if (scrollManager) {
+										scrollManager.resetUserScrolling()
+									}
 								}
-								setShowScrollToBottom(disableAutoScrollRef.current && !isAtBottom)
 							}}
 							atBottomThreshold={10} // anything lower causes issues with followOutput
 							initialTopMostItemIndex={groupedMessages.length - 1}