From 67721a45eed68db5a55a01c32b0d26a650fa41ec Mon Sep 17 00:00:00 2001
From: Louis <louis@jan.ai>
Date: Sat, 30 Nov 2024 16:31:34 +0700
Subject: [PATCH] fix: high performance convo rendering - virtualizing large
 convo

---
 web/package.json                              |   1 +
 .../ThreadCenterPanel/ChatBody/index.tsx      | 115 ++++++++++++++----
 .../ThreadCenterPanel/ChatItem/index.tsx      |   2 +-
 .../LoadModelError/index.tsx                  |   2 +-
 4 files changed, 95 insertions(+), 25 deletions(-)

diff --git a/web/package.json b/web/package.json
index 3fdc6889b..5cd3e2b08 100644
--- a/web/package.json
+++ b/web/package.json
@@ -16,6 +16,7 @@
   "dependencies": {
     "@janhq/core": "link:./core",
     "@janhq/joi": "link:./joi",
+    "@tanstack/react-virtual": "^3.10.9",
     "autoprefixer": "10.4.16",
     "class-variance-authority": "^0.7.0",
     "framer-motion": "^10.16.4",
diff --git a/web/screens/Thread/ThreadCenterPanel/ChatBody/index.tsx b/web/screens/Thread/ThreadCenterPanel/ChatBody/index.tsx
index ee001bcb1..41f0c1597 100644
--- a/web/screens/Thread/ThreadCenterPanel/ChatBody/index.tsx
+++ b/web/screens/Thread/ThreadCenterPanel/ChatBody/index.tsx
@@ -1,12 +1,10 @@
-import { memo, useEffect, useState } from 'react'
+import { memo, useEffect, useMemo, useRef, useState } from 'react'
 
-import { MessageStatus, ThreadMessage } from '@janhq/core'
+import { ThreadMessage } from '@janhq/core'
+import { useVirtualizer } from '@tanstack/react-virtual'
 
 import { useAtomValue } from 'jotai'
 
-import ErrorMessage from '@/containers/ErrorMessage'
-import ListContainer from '@/containers/ListContainer'
-
 import { loadModelErrorAtom } from '@/hooks/useActiveModel'
 
 import ChatItem from '../ChatItem'
@@ -21,6 +19,7 @@ const ChatConfigurator = memo(() => {
   const messages = useAtomValue(getCurrentChatMessagesAtom)
 
   const [current, setCurrent] = useState<ThreadMessage[]>([])
+  const loadModelError = useAtomValue(loadModelErrorAtom)
 
   const isMessagesIdentificial = (
     arr1: ThreadMessage[],
@@ -32,14 +31,12 @@ const ChatConfigurator = memo(() => {
 
   useEffect(() => {
     if (
-      messages.length !== current.length ||
+      messages?.length !== current?.length ||
       !isMessagesIdentificial(messages, current)
     ) {
       setCurrent(messages)
     }
-  }, [messages, current])
-
-  const loadModelError = useAtomValue(loadModelErrorAtom)
+  }, [messages, current, loadModelError])
 
   if (!messages.length) return <EmptyThread />
   return (
@@ -57,21 +54,93 @@ const ChatBody = memo(
     messages: ThreadMessage[]
     loadModelError?: string
   }) => {
-    return (
-      <ListContainer>
-        {messages.map((message, index) => (
-          <div key={message.id}>
-            <ChatItem
-              {...message}
-              key={message.id}
-              loadModelError={loadModelError}
-              isCurrentMessage={index === messages.length - 1}
-            />
-          </div>
-        ))}
+    // The scrollable element for your list
+    const parentRef = useRef(null)
 
-        {loadModelError && <LoadModelError />}
-      </ListContainer>
+    const count = useMemo(
+      () => (messages?.length ?? 0) + (loadModelError ? 1 : 0),
+      [messages, loadModelError]
+    )
+
+    // The virtualizer
+    const virtualizer = useVirtualizer({
+      count,
+      getScrollElement: () => parentRef.current,
+      estimateSize: () => 35,
+      overscan: 5,
+    })
+    useEffect(() => {
+      if (count > 0 && messages && virtualizer) {
+        virtualizer.scrollToIndex(count - 1)
+      }
+    }, [count, virtualizer, messages, loadModelError])
+
+    const items = virtualizer.getVirtualItems()
+    virtualizer.shouldAdjustScrollPositionOnItemSizeChange = (
+      item,
+      _,
+      instance
+    ) => {
+      return (
+        // item.start < (instance.scrollOffset ?? 0) &&
+        instance.scrollDirection !== 'backward'
+      )
+    }
+
+    return (
+      <div className="flex h-full w-full flex-col overflow-x-hidden">
+        <div
+          ref={parentRef}
+          className="List"
+          style={{
+            flex: 1,
+            height: '100%',
+            width: '100%',
+            overflowY: 'auto',
+            overflowX: 'hidden',
+            contain: 'strict',
+          }}
+        >
+          <div
+            style={{
+              height: virtualizer.getTotalSize(),
+              width: '100%',
+              position: 'relative',
+            }}
+          >
+            <div
+              style={{
+                position: 'absolute',
+                top: 0,
+                left: 0,
+                width: '100%',
+                transform: `translateY(${items[0]?.start ?? 0}px)`,
+              }}
+            >
+              {items.map((virtualRow) => (
+                <div
+                  key={virtualRow.key}
+                  data-index={virtualRow.index}
+                  ref={virtualizer.measureElement}
+                >
+                  {loadModelError && virtualRow.index === count - 1 ? (
+                    <LoadModelError />
+                  ) : (
+                    <ChatItem
+                      {...messages[virtualRow.index]}
+                      // key={messages[virtualRow.index]?.id}
+                      loadModelError={loadModelError}
+                      isCurrentMessage={
+                        virtualRow.index === messages?.length - 1
+                      }
+                    />
+                  )}
+                </div>
+              ))}
+            </div>
+          </div>
+        </div>
+      </div>
     )
   }
 )
diff --git a/web/screens/Thread/ThreadCenterPanel/ChatItem/index.tsx b/web/screens/Thread/ThreadCenterPanel/ChatItem/index.tsx
index 192c6f82a..92fd6394d 100644
--- a/web/screens/Thread/ThreadCenterPanel/ChatItem/index.tsx
+++ b/web/screens/Thread/ThreadCenterPanel/ChatItem/index.tsx
@@ -52,7 +52,7 @@ const ChatItem = forwardRef<Ref, Props>((message, ref) => {
 
   return (
     <>
-      {status !== MessageStatus.Error && content.length > 0 && (
+      {status !== MessageStatus.Error && content?.length > 0 && (
         <div ref={ref} className="relative">
           <SimpleTextMessage {...message} content={content} status={status} />
         </div>
diff --git a/web/screens/Thread/ThreadCenterPanel/LoadModelError/index.tsx b/web/screens/Thread/ThreadCenterPanel/LoadModelError/index.tsx
index 5829a6923..f17bf43c4 100644
--- a/web/screens/Thread/ThreadCenterPanel/LoadModelError/index.tsx
+++ b/web/screens/Thread/ThreadCenterPanel/LoadModelError/index.tsx
@@ -66,7 +66,7 @@ const LoadModelError = () => {
   }
 
   return (
-    <div className="mt-10">
+    <div className="flex flex-1">
       <div className="flex w-full flex-col items-center text-center font-medium">
         <p className="w-[90%]">
           <ErrorMessage />