From 1a5e9759d892daaffd6e5af5ac15d33a9c4b029d Mon Sep 17 00:00:00 2001
From: Saoud Rizwan <7799382+saoudrizwan@users.noreply.github.com>
Date: Wed, 2 Oct 2024 21:19:10 -0400
Subject: [PATCH] Use key prop to re-render virtuoso when task changes

---
 webview-ui/src/components/chat/ChatView.tsx | 15 ++-------------
 1 file changed, 2 insertions(+), 13 deletions(-)

diff --git a/webview-ui/src/components/chat/ChatView.tsx b/webview-ui/src/components/chat/ChatView.tsx
index c9c208a..22568e6 100644
--- a/webview-ui/src/components/chat/ChatView.tsx
+++ b/webview-ui/src/components/chat/ChatView.tsx
@@ -51,7 +51,6 @@ const ChatView = ({ isHidden, showAnnouncement, hideAnnouncement, showHistoryVie
 
 	const [showScrollToBottom, setShowScrollToBottom] = useState(false)
 	const isAtBottomRef = useRef(false)
-	const taskMsgTsRef = useRef<number | undefined>(undefined)
 	const [didScrollUp, setDidScrollUp] = useState(false)
 	const lastScrollTopRef = useRef(0)
 	const [didClickScrollToBottom, setDidClickScrollToBottom] = useState(false)
@@ -484,18 +483,6 @@ const ChatView = ({ isHidden, showAnnouncement, hideAnnouncement, showHistoryVie
 		[visibleMessages, expandedRows, scrollToBottomAuto]
 	)
 
-	// scroll to bottom if task changes
-	// (this gets called when messages changes, so we use ref to ts to detect new task)
-	useEffect(() => {
-		if (task && task.ts !== taskMsgTsRef.current) {
-			taskMsgTsRef.current = task.ts
-			const timer = setTimeout(() => {
-				scrollToBottomSmooth()
-			}, 50)
-			return () => clearTimeout(timer)
-		}
-	}, [task, scrollToBottomSmooth, visibleMessages.length])
-
 	const handleRowHeightChange = useCallback(() => {
 		if (isAtBottomRef.current || didClickScrollToBottom || !didScrollUp) {
 			scrollToBottomSmooth()
@@ -590,6 +577,7 @@ const ChatView = ({ isHidden, showAnnouncement, hideAnnouncement, showHistoryVie
 				<>
 					<Virtuoso
 						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"
 						style={{
 							flexGrow: 1,
@@ -637,6 +625,7 @@ const ChatView = ({ isHidden, showAnnouncement, hideAnnouncement, showHistoryVie
 						}}
 						atBottomThreshold={10} // anything lower causes issues with followOutput
 						onScroll={handleScroll}
+						initialTopMostItemIndex={visibleMessages.length - 1}
 					/>
 					{showScrollToBottom ? (
 						<div