From 0c8fbb11d888fb6ec80fd816c6d0884dd3f0a3cb Mon Sep 17 00:00:00 2001
From: Saoud Rizwan <7799382+saoudrizwan@users.noreply.github.com>
Date: Fri, 13 Sep 2024 03:23:06 -0400
Subject: [PATCH] Fix scrolling behavior in chatview

---
 webview-ui/src/components/ChatView.tsx | 49 +++++++++++++++++++++-----
 1 file changed, 41 insertions(+), 8 deletions(-)

diff --git a/webview-ui/src/components/ChatView.tsx b/webview-ui/src/components/ChatView.tsx
index 710c065..eedbc06 100644
--- a/webview-ui/src/components/ChatView.tsx
+++ b/webview-ui/src/components/ChatView.tsx
@@ -50,13 +50,7 @@ const ChatView = ({ isHidden, showAnnouncement, hideAnnouncement, showHistoryVie
 	const [secondaryButtonText, setSecondaryButtonText] = useState<string | undefined>(undefined)
 	const virtuosoRef = useRef<VirtuosoHandle>(null)
 	const [expandedRows, setExpandedRows] = useState<Record<number, boolean>>({})
-
-	const toggleRowExpansion = (ts: number) => {
-		setExpandedRows((prev) => ({
-			...prev,
-			[ts]: !prev[ts],
-		}))
-	}
+	const [isAtBottom, setIsAtBottom] = useState(false)
 
 	useEffect(() => {
 		// if last message is an ask, show user ask UI
@@ -426,6 +420,43 @@ const ChatView = ({ isHidden, showAnnouncement, hideAnnouncement, showHistoryVie
 		})
 	}, [modifiedMessages])
 
+	const toggleRowExpansion = useCallback(
+		(ts: number) => {
+			const isCollapsing = expandedRows[ts] ?? false
+			const isLastMessage = visibleMessages.at(-1)?.ts === ts
+			setExpandedRows((prev) => ({
+				...prev,
+				[ts]: !prev[ts],
+			}))
+
+			if (isCollapsing && isAtBottom) {
+				const timer = setTimeout(() => {
+					if (virtuosoRef.current) {
+						virtuosoRef.current.scrollToIndex({
+							index: visibleMessages.length - 1,
+							align: "end",
+						})
+					}
+				}, 0)
+				return () => clearTimeout(timer)
+			}
+
+			if (!isCollapsing && isLastMessage) {
+				const timer = setTimeout(() => {
+					if (virtuosoRef.current) {
+						virtuosoRef.current.scrollToIndex({
+							index: visibleMessages.length - 1,
+							align: "start",
+							behavior: "smooth",
+						})
+					}
+				}, 0)
+				return () => clearTimeout(timer)
+			}
+		},
+		[isAtBottom, visibleMessages, expandedRows]
+	)
+
 	useEffect(() => {
 		// We use a setTimeout to ensure new content is rendered before scrolling to the bottom. virtuoso's followOutput would scroll to the bottom before the new content could render.
 		const timer = setTimeout(() => {
@@ -453,7 +484,7 @@ const ChatView = ({ isHidden, showAnnouncement, hideAnnouncement, showHistoryVie
 				isLast={index === visibleMessages.length - 1}
 			/>
 		),
-		[expandedRows, modifiedMessages, visibleMessages.length]
+		[expandedRows, modifiedMessages, visibleMessages.length, toggleRowExpansion]
 	)
 
 	return (
@@ -525,6 +556,8 @@ const ChatView = ({ isHidden, showAnnouncement, hideAnnouncement, showHistoryVie
 						increaseViewportBy={{ top: 1_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)
 						data={visibleMessages} // 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={setIsAtBottom}
+						atBottomThreshold={100}
 					/>
 					<div
 						style={{