From 9c85a36b2c1910af092a2003c21429f93ed9b165 Mon Sep 17 00:00:00 2001
From: Saoud Rizwan <7799382+saoudrizwan@users.noreply.github.com>
Date: Wed, 24 Jul 2024 12:19:14 -0400
Subject: [PATCH] Implement virtualized rendering of items in the chat view

---
 webview-ui/package-lock.json           | 13 ++++++
 webview-ui/package.json                |  1 +
 webview-ui/src/components/ChatRow.tsx  | 24 +++++++---
 webview-ui/src/components/ChatView.tsx | 61 ++++++++++----------------
 4 files changed, 55 insertions(+), 44 deletions(-)

diff --git a/webview-ui/package-lock.json b/webview-ui/package-lock.json
index e6f0dce..0b4af76 100644
--- a/webview-ui/package-lock.json
+++ b/webview-ui/package-lock.json
@@ -23,6 +23,7 @@
         "react-text-truncate": "^0.19.0",
         "react-textarea-autosize": "^8.5.3",
         "react-use": "^17.5.1",
+        "react-virtuoso": "^4.7.13",
         "rewire": "^7.0.0",
         "typescript": "^4.9.5",
         "web-vitals": "^2.1.4"
@@ -16567,6 +16568,18 @@
         "react-dom": "*"
       }
     },
+    "node_modules/react-virtuoso": {
+      "version": "4.7.13",
+      "resolved": "https://registry.npmjs.org/react-virtuoso/-/react-virtuoso-4.7.13.tgz",
+      "integrity": "sha512-rabPhipwJ8rdA6TDk1vdVqVoU6eOkWukqoC1pNQVBCsvjBvIeJMi9nO079s0L7EsRzAxFFQNahX+8vuuY4F1Qg==",
+      "engines": {
+        "node": ">=10"
+      },
+      "peerDependencies": {
+        "react": ">=16 || >=17 || >= 18",
+        "react-dom": ">=16 || >=17 || >= 18"
+      }
+    },
     "node_modules/read-cache": {
       "version": "1.0.0",
       "resolved": "https://registry.npmjs.org/read-cache/-/read-cache-1.0.0.tgz",
diff --git a/webview-ui/package.json b/webview-ui/package.json
index a82e0fb..b05a8bd 100644
--- a/webview-ui/package.json
+++ b/webview-ui/package.json
@@ -18,6 +18,7 @@
     "react-text-truncate": "^0.19.0",
     "react-textarea-autosize": "^8.5.3",
     "react-use": "^17.5.1",
+    "react-virtuoso": "^4.7.13",
     "rewire": "^7.0.0",
     "typescript": "^4.9.5",
     "web-vitals": "^2.1.4"
diff --git a/webview-ui/src/components/ChatRow.tsx b/webview-ui/src/components/ChatRow.tsx
index c2872b0..69e245c 100644
--- a/webview-ui/src/components/ChatRow.tsx
+++ b/webview-ui/src/components/ChatRow.tsx
@@ -114,7 +114,7 @@ const ChatRow: React.FC<ChatRowProps> = ({ message, syntaxHighlighterStyle }) =>
 							</div>
 						)
 					case "api_req_finished":
-						return null // Hide this message type
+						return null // we should never see this message type
 					case "text":
 						return <p style={contentStyle}>{message.text}</p>
 					case "user_feedback":
@@ -332,18 +332,14 @@ const ChatRow: React.FC<ChatRowProps> = ({ message, syntaxHighlighterStyle }) =>
 	}
 
 	// we need to return null here instead of in getContent since that way would result in padding being applied
-	if (message.say === "api_req_finished") {
-		return null // Don't render anything for this message type
-	}
-
-	if (message.type === "ask" && message.ask === "completion_result" && message.text === "") {
+	if (!shouldShowChatRow(message)) {
 		return null // Don't render anything for this message type
 	}
 
 	return (
 		<div
 			style={{
-				padding: "10px 0px 10px 0px",
+				padding: "10px 6px 10px 15px",
 			}}>
 			{renderContent()}
 			{isExpanded && message.say === "api_req_started" && (
@@ -359,4 +355,18 @@ const ChatRow: React.FC<ChatRowProps> = ({ message, syntaxHighlighterStyle }) =>
 	)
 }
 
+export const shouldShowChatRow = (message: ClaudeMessage) => {
+	// combineApiRequests removes this from modifiedMessages anyways
+	if (message.say === "api_req_finished") {
+		return false
+	}
+
+	// don't show a chat row for a completion_result ask without text. This specific type of message only occurs if Claude wants to execute a command as part of its completion result, in which case we interject the completion_result tool with the execute_command tool.
+	if (message.type === "ask" && message.ask === "completion_result" && message.text === "") {
+		return false
+	}
+
+	return true
+}
+
 export default ChatRow
diff --git a/webview-ui/src/components/ChatView.tsx b/webview-ui/src/components/ChatView.tsx
index f1944ab..b5172a4 100644
--- a/webview-ui/src/components/ChatView.tsx
+++ b/webview-ui/src/components/ChatView.tsx
@@ -9,8 +9,9 @@ import { combineCommandSequences } from "../utilities/combineCommandSequences"
 import { getApiMetrics } from "../utilities/getApiMetrics"
 import { getSyntaxHighlighterStyleFromTheme } from "../utilities/getSyntaxHighlighterStyleFromTheme"
 import { vscode } from "../utilities/vscode"
-import ChatRow from "./ChatRow"
+import ChatRow, { shouldShowChatRow } from "./ChatRow"
 import TaskHeader from "./TaskHeader"
+import { Virtuoso, type VirtuosoHandle } from "react-virtuoso"
 
 interface ChatViewProps {
 	messages: ClaudeMessage[]
@@ -39,7 +40,7 @@ const ChatView = ({ messages, isHidden, vscodeThemeName }: ChatViewProps) => {
 
 	const [syntaxHighlighterStyle, setSyntaxHighlighterStyle] = useState(vsDarkPlus)
 
-	const chatContainerRef = useRef<HTMLDivElement>(null)
+	const virtuosoRef = useRef<VirtuosoHandle>(null)
 
 	useEffect(() => {
 		if (!vscodeThemeName) return
@@ -49,31 +50,6 @@ const ChatView = ({ messages, isHidden, vscodeThemeName }: ChatViewProps) => {
 		}
 	}, [vscodeThemeName])
 
-	const scrollToBottom = (instant: boolean = false) => {
-		if (chatContainerRef.current) {
-			const scrollOptions: ScrollToOptions = {
-				top: chatContainerRef.current.scrollHeight,
-				behavior: instant ? "auto" : "smooth",
-			}
-			chatContainerRef.current.scrollTo(scrollOptions)
-		}
-	}
-
-	// scroll to bottom when new message is added
-	const visibleMessages = useMemo(
-		() =>
-			modifiedMessages.filter(
-				(message) => !(message.type === "ask" && message.ask === "completion_result" && message.text === "")
-			),
-		[modifiedMessages]
-	)
-	useEffect(() => {
-		const timer = setTimeout(() => {
-			scrollToBottom()
-		}, 0)
-		return () => clearTimeout(timer)
-	}, [visibleMessages])
-
 	useEffect(() => {
 		// if last message is an ask, show user ask UI
 
@@ -323,18 +299,27 @@ const ChatView = ({ messages, isHidden, vscodeThemeName }: ChatViewProps) => {
 					</p>
 				</div>
 			)}
-			<div
-				ref={chatContainerRef}
+			<Virtuoso
+				ref={virtuosoRef}
 				className="scrollable"
 				style={{
 					flexGrow: 1,
-					overflowY: "scroll",
-					padding: "0 6px 0 15px",
-				}}>
-				{modifiedMessages.map((message, index) => (
-					<ChatRow key={index} message={message} syntaxHighlighterStyle={syntaxHighlighterStyle} />
-				))}
-			</div>
+					overflowY: "scroll", // always show scrollbar
+				}}
+				followOutput={(isAtBottom) => {
+					// TODO: we can use isAtBottom to prevent scrolling if user is scrolled up, and show a 'scroll to bottom' button for better UX
+					const lastMessage = modifiedMessages.at(-1)
+					if (lastMessage && shouldShowChatRow(lastMessage)) {
+						return "smooth"
+					}
+					return false
+				}}
+				increaseViewportBy={{ top: 0, bottom: Infinity }} // hack to make sure the last message is always rendered to get truly perfect scroll to bottom animation when new messages are added
+				data={modifiedMessages}
+				itemContent={(index, message) => (
+					<ChatRow key={message.ts} message={message} syntaxHighlighterStyle={syntaxHighlighterStyle} />
+				)}
+			/>
 			<div
 				style={{
 					opacity: primaryButtonText || secondaryButtonText ? (enableButtons ? 1 : 0.5) : 0,
@@ -370,7 +355,9 @@ const ChatView = ({ messages, isHidden, vscodeThemeName }: ChatViewProps) => {
 					disabled={textAreaDisabled}
 					onChange={(e) => setInputValue(e.target.value)}
 					onKeyDown={handleKeyDown}
-					onHeightChange={() => scrollToBottom(true)}
+					onHeightChange={() =>
+						virtuosoRef.current?.scrollToIndex({ index: "LAST", align: "end", behavior: "auto" })
+					}
 					placeholder={task ? "Type a message..." : "Type your task here..."}
 					maxRows={10}
 					autoFocus={true}