fix: convo component re-render every single token yield

web/screens/Thread/ThreadCenterPanel/ChatBody/index.tsx

@@ -1,6 +1,6 @@
-import { memo } from 'react'
+import { memo, useEffect, useState } from 'react'
 
-import { MessageStatus } from '@janhq/core'
+import { MessageStatus, ThreadMessage } from '@janhq/core'
 
 import { useAtomValue } from 'jotai'
 
@@ -17,33 +17,67 @@ import EmptyThread from './EmptyThread'
 
 import { getCurrentChatMessagesAtom } from '@/helpers/atoms/ChatMessage.atom'
 
-const ChatBody = () => {
+const ChatConfigurator = memo(() => {
   const messages = useAtomValue(getCurrentChatMessagesAtom)
 
+  const [current, setCurrent] = useState<ThreadMessage[]>([])
+
+  const isMessagesIdentificial = (
+    arr1: ThreadMessage[],
+    arr2: ThreadMessage[]
+  ): boolean => {
+    if (arr1.length !== arr2.length) return false
+    return arr1.every((item, index) => item.id === arr2[index].id)
+  }
+
+  useEffect(() => {
+    if (
+      messages.length !== current.length ||
+      !isMessagesIdentificial(messages, current)
+    ) {
+      setCurrent(messages)
+    }
+  }, [messages, current])
+
   const loadModelError = useAtomValue(loadModelErrorAtom)
 
   if (!messages.length) return <EmptyThread />
-
   return (
-    <ListContainer>
-      {messages.map((message, index) => (
-        <div key={message.id}>
-          {message.status !== MessageStatus.Error &&
-            message.content.length > 0 && (
+    <div className="flex h-full w-full flex-col">
+      <ChatBody loadModelError={loadModelError} messages={current} />
+    </div>
+  )
+})
+
+const ChatBody = memo(
+  ({
+    messages,
+    loadModelError,
+  }: {
+    messages: ThreadMessage[]
+    loadModelError?: string
+  }) => {
+    return (
+      <ListContainer>
+        {messages.map((message, index) => (
+          <div key={message.id}>
+            {message.status !== MessageStatus.Error && (
               <ChatItem {...message} key={message.id} />
             )}
 
-          {!loadModelError &&
-            index === messages.length - 1 &&
-            message.status !== MessageStatus.Pending &&
-            message.status !== MessageStatus.Ready && (
-              <ErrorMessage message={message} />
-            )}
-        </div>
-      ))}
-      {loadModelError && <LoadModelError />}
-    </ListContainer>
-  )
-}
+            {!loadModelError &&
+              index === messages.length - 1 &&
+              message.status !== MessageStatus.Pending &&
+              message.status !== MessageStatus.Ready && (
+                <ErrorMessage message={message} />
+              )}
+          </div>
+        ))}
+
+        {loadModelError && <LoadModelError />}
+      </ListContainer>
+    )
+  }
+)
 
-export default memo(ChatBody)
+export default memo(ChatConfigurator)

web/screens/Thread/ThreadCenterPanel/ChatItem/index.tsx

@@ -1,15 +1,42 @@
-import React, { forwardRef } from 'react'
+import React, { forwardRef, useEffect, useState } from 'react'
 
-import { ThreadMessage } from '@janhq/core'
+import {
+  events,
+  MessageEvent,
+  MessageStatus,
+  ThreadContent,
+  ThreadMessage,
+} from '@janhq/core'
 
 import SimpleTextMessage from '../SimpleTextMessage'
 
 type Ref = HTMLDivElement
 
-const ChatItem = forwardRef<Ref, ThreadMessage>((message, ref) => (
-  <div ref={ref} className="relative">
-    <SimpleTextMessage {...message} />
-  </div>
-))
+const ChatItem = forwardRef<Ref, ThreadMessage>((message, ref) => {
+  const [content, setContent] = useState<ThreadContent[]>(message.content)
+  const [status, setStatus] = useState<MessageStatus>(message.status)
+
+  function onMessageUpdate(data: ThreadMessage) {
+    if (data.id === message.id) {
+      setContent(data.content)
+      if (data.status !== status) setStatus(data.status)
+    }
+  }
+
+  useEffect(() => {
+    if (message.status === MessageStatus.Pending)
+      events.on(MessageEvent.OnMessageUpdate, onMessageUpdate)
+    return () => {
+      events.off(MessageEvent.OnMessageUpdate, onMessageUpdate)
+    }
+    // eslint-disable-next-line react-hooks/exhaustive-deps
+  }, [])
+
+  return (
+    <div ref={ref} className="relative">
+      <SimpleTextMessage {...message} content={content} status={status} />
+    </div>
+  )
+})
 
 export default ChatItem