From e7e932059a9f22640acd01d0fdb339384541949b Mon Sep 17 00:00:00 2001
From: Louis <louis@jan.ai>
Date: Fri, 29 Sep 2023 15:01:28 +0700
Subject: [PATCH] fix: state management - get current conv messages with
 selectAtom

---
 web/app/_components/ChatBody/index.tsx | 33 +++++++++++++++++++-------
 web/app/_hooks/useSendChatMessage.ts   | 19 +++++++++++++--
 2 files changed, 41 insertions(+), 11 deletions(-)

diff --git a/web/app/_components/ChatBody/index.tsx b/web/app/_components/ChatBody/index.tsx
index dbf8eb5da..4ec776e78 100644
--- a/web/app/_components/ChatBody/index.tsx
+++ b/web/app/_components/ChatBody/index.tsx
@@ -5,14 +5,24 @@ import ChatItem from "../ChatItem";
 import { ChatMessage } from "@/_models/ChatMessage";
 import useChatMessages from "@/_hooks/useChatMessages";
 import {
-  currentChatMessagesAtom,
+  chatMessages,
+  getActiveConvoIdAtom,
   showingTyping,
 } from "@/_helpers/JotaiWrapper";
 import { useAtomValue } from "jotai";
+import { selectAtom } from "jotai/utils";
 import LoadingIndicator from "../LoadingIndicator";
 
 const ChatBody: React.FC = () => {
-  const messages = useAtomValue(currentChatMessagesAtom);
+  const activeConversationId = useAtomValue(getActiveConvoIdAtom) ?? "";
+  const messageList = useAtomValue(
+    selectAtom(
+      chatMessages,
+      useCallback((v) => v[activeConversationId], [activeConversationId])
+    )
+  );
+  const [content, setContent] = useState<React.JSX.Element[]>([]);
+
   const isTyping = useAtomValue(showingTyping);
   const [offset, setOffset] = useState(0);
   const { loading, hasMore } = useChatMessages(offset);
@@ -35,13 +45,18 @@ const ChatBody: React.FC = () => {
     [loading, hasMore]
   );
 
-  const content = messages.map((message, index) => {
-    if (messages.length === index + 1) {
-      // @ts-ignore
-      return <ChatItem ref={lastPostRef} message={message} key={message.id} />;
-    }
-    return <ChatItem message={message} key={message.id} />;
-  });
+  React.useEffect(() => {
+    const list = messageList?.map((message, index) => {
+      if (messageList?.length === index + 1) {
+        return (
+          // @ts-ignore
+          <ChatItem ref={lastPostRef} message={message} key={message.id} />
+        );
+      }
+      return <ChatItem message={message} key={message.id} />;
+    });
+    setContent(list);
+  }, [messageList, lastPostRef]);
 
   return (
     <div className="flex flex-col-reverse flex-1 py-4 overflow-y-auto scroll">
diff --git a/web/app/_hooks/useSendChatMessage.ts b/web/app/_hooks/useSendChatMessage.ts
index 7bde91cb1..a6312e570 100644
--- a/web/app/_hooks/useSendChatMessage.ts
+++ b/web/app/_hooks/useSendChatMessage.ts
@@ -1,14 +1,16 @@
 import {
   addNewMessageAtom,
-  currentChatMessagesAtom,
+  chatMessages,
   currentConversationAtom,
   currentPromptAtom,
   currentStreamingMessageAtom,
+  getActiveConvoIdAtom,
   showingTyping,
   updateMessageAtom,
 } from "@/_helpers/JotaiWrapper";
 
 import { useAtom, useAtomValue, useSetAtom } from "jotai";
+import { selectAtom } from "jotai/utils";
 import { DataService } from "../../shared/coreService";
 import {
   MessageSenderType,
@@ -16,13 +18,21 @@ import {
   toChatMessage,
 } from "@/_models/ChatMessage";
 import { executeSerial } from "@/_services/pluginService";
+import { useCallback } from "react";
 
 export default function useSendChatMessage() {
   const currentConvo = useAtomValue(currentConversationAtom);
   const updateStreamMessage = useSetAtom(currentStreamingMessageAtom);
   const addNewMessage = useSetAtom(addNewMessageAtom);
   const updateMessage = useSetAtom(updateMessageAtom);
-  const chatMessagesHistory = useAtomValue(currentChatMessagesAtom);
+  const activeConversationId = useAtomValue(getActiveConvoIdAtom) ?? "";
+
+  const chatMessagesHistory = useAtomValue(
+    selectAtom(
+      chatMessages,
+      useCallback((v) => v[activeConversationId], [activeConversationId])
+    )
+  );
   const [currentPrompt, setCurrentPrompt] = useAtom(currentPromptAtom);
   const [, setIsTyping] = useAtom(showingTyping);
   const sendChatMessage = async () => {
@@ -120,6 +130,11 @@ export default function useSendChatMessage() {
         }
       }
     }
+    updateMessage(
+      responseChatMessage.id,
+      responseChatMessage.conversationId,
+      answer
+    );
     await executeSerial(DataService.UPDATE_MESSAGE, {
       ...newResponse,
       message: answer,