From 19c6a26ef59b5f1348fcb2bedbd91ffd15ec1ce7 Mon Sep 17 00:00:00 2001
From: Louis <louis@jan.ai>
Date: Thu, 28 Nov 2024 13:23:03 +0700
Subject: [PATCH] fix: proper error message display

---
 .../ThreadCenterPanel/ChatBody/index.tsx      | 16 ++++-----
 .../ThreadCenterPanel/ChatItem/index.tsx      | 33 ++++++++++++++++---
 2 files changed, 35 insertions(+), 14 deletions(-)

diff --git a/web/screens/Thread/ThreadCenterPanel/ChatBody/index.tsx b/web/screens/Thread/ThreadCenterPanel/ChatBody/index.tsx
index d226c7e6f..ee001bcb1 100644
--- a/web/screens/Thread/ThreadCenterPanel/ChatBody/index.tsx
+++ b/web/screens/Thread/ThreadCenterPanel/ChatBody/index.tsx
@@ -61,16 +61,12 @@ const ChatBody = memo(
       <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} />
-              )}
+            <ChatItem
+              {...message}
+              key={message.id}
+              loadModelError={loadModelError}
+              isCurrentMessage={index === messages.length - 1}
+            />
           </div>
         ))}
 
diff --git a/web/screens/Thread/ThreadCenterPanel/ChatItem/index.tsx b/web/screens/Thread/ThreadCenterPanel/ChatItem/index.tsx
index 913333c6b..192c6f82a 100644
--- a/web/screens/Thread/ThreadCenterPanel/ChatItem/index.tsx
+++ b/web/screens/Thread/ThreadCenterPanel/ChatItem/index.tsx
@@ -8,21 +8,39 @@ import {
   ThreadMessage,
 } from '@janhq/core'
 
+import ErrorMessage from '@/containers/ErrorMessage'
+
 import SimpleTextMessage from '../SimpleTextMessage'
 
 type Ref = HTMLDivElement
 
-const ChatItem = forwardRef<Ref, ThreadMessage>((message, ref) => {
+type Props = {
+  loadModelError?: string
+  isCurrentMessage?: boolean
+} & ThreadMessage
+
+const ChatItem = forwardRef<Ref, Props>((message, ref) => {
   const [content, setContent] = useState<ThreadContent[]>(message.content)
   const [status, setStatus] = useState<MessageStatus>(message.status)
+  const [errorMessage, setErrorMessage] = useState<ThreadMessage | undefined>(
+    message.isCurrentMessage && message.status === MessageStatus.Error
+      ? message
+      : undefined
+  )
 
   function onMessageUpdate(data: ThreadMessage) {
     if (data.id === message.id) {
       setContent(data.content)
       if (data.status !== status) setStatus(data.status)
+      if (data.status === MessageStatus.Error && message.isCurrentMessage)
+        setErrorMessage(data)
     }
   }
 
+  useEffect(() => {
+    if (!message.isCurrentMessage && errorMessage) setErrorMessage(undefined)
+  }, [message, errorMessage])
+
   useEffect(() => {
     if (message.status === MessageStatus.Pending)
       events.on(MessageEvent.OnMessageUpdate, onMessageUpdate)
@@ -33,9 +51,16 @@ const ChatItem = forwardRef<Ref, ThreadMessage>((message, ref) => {
   }, [])
 
   return (
-    <div ref={ref} className="relative">
-      <SimpleTextMessage {...message} content={content} status={status} />
-    </div>
+    <>
+      {status !== MessageStatus.Error && content.length > 0 && (
+        <div ref={ref} className="relative">
+          <SimpleTextMessage {...message} content={content} status={status} />
+        </div>
+      )}
+      {errorMessage && !message.loadModelError && (
+        <ErrorMessage message={errorMessage} />
+      )}
+    </>
   )
 })