fix(gpt-runner-web): optimize sidebar performance and fix chat history lost issue

packages/gpt-runner-web/client/src/store/zustand/global/chat.slice.ts

@@ -18,6 +18,7 @@ export interface ChatSlice {
   chatInstances: SingleChat[]
   updateActiveChatId: (activeChatId: string) => void
   getChatInstance: (chatId: string) => SingleChat | undefined
+  getChatInstancesBySingleFilePath: (singleFilePath: string) => SingleChat[]
   addChatInstance: (gptFileId: string, instance: Omit<SingleChat, 'id'>) => {
     chatSidebarTreeItem: SidebarTreeItem
     chatInstance: SingleChat
@@ -26,6 +27,7 @@ export interface ChatSlice {
     (chatId: string, chat: Partial<SingleChat>, replace: false): void
     (chatId: string, chat: SingleChat, replace: true): void
   }
+  updateChatInstances: (chatInstances: SingleChat[] | ((oldChatInstances: SingleChat[]) => SingleChat[])) => void
   removeChatInstance: (chatId: string) => void
   generateChatAnswer: (chatId: string, type?: GenerateAnswerType) => Promise<void>
   regenerateLastChatAnswer: (chatId: string) => Promise<void>
@@ -42,6 +44,8 @@ function getInitialState() {
 }
 
 const chatIdAbortCtrlMap = new Map<string, AbortController>()
+const chatIdChatInstanceMap = new Map<string, SingleChat>()
+const singleFilePathChatInstancesMap = new Map<string, SingleChat[]>()
 
 export const createChatSlice: StateCreator<
   ChatSlice & SidebarTreeSlice & FileTreeSlice,
@@ -54,9 +58,11 @@ export const createChatSlice: StateCreator<
     set({ activeChatId })
   },
   getChatInstance(chatId) {
-    return get().chatInstances.find(chatInstance => chatInstance.id === chatId)
+    return chatIdChatInstanceMap.get(chatId)
+  },
+  getChatInstancesBySingleFilePath(singleFilePath) {
+    return singleFilePathChatInstancesMap.get(singleFilePath) || []
   },
-
   addChatInstance(gptFileId, instance) {
     const state = get()
     const chatId = uuidv4()
@@ -65,7 +71,7 @@ export const createChatSlice: StateCreator<
     const finalInstance: SingleChat = {
       ...instance,
       id: chatId,
-      singleFilePath: gptFileIdTreeItem?.otherInfo?.path || '',
+      singleFilePath: gptFileIdTreeItem?.path || '',
     }
 
     const chatInfo = state.getChatInfo(finalInstance)
@@ -81,9 +87,7 @@ export const createChatSlice: StateCreator<
       return oldItem
     })
 
-    set(state => ({
-      chatInstances: [...state.chatInstances, finalInstance],
-    }))
+    state.updateChatInstances(preState => [...preState, finalInstance])
 
     return {
       chatSidebarTreeItem,
@@ -94,20 +98,52 @@ export const createChatSlice: StateCreator<
   updateChatInstance(chatId, chat, replace = false) {
     set(state => ({
       chatInstances: state.chatInstances.map((chatInstance) => {
-        if (chatInstance.id === chatId)
-          return replace ? chat as SingleChat : Object.assign(chatInstance, chat)
+        if (chatInstance.id === chatId) {
+          const oldChatInstance = chatInstance
+          const newChatInstance = replace ? chat as SingleChat : Object.assign(chatInstance, chat)
+
+          // remove old map
+          chatIdChatInstanceMap.delete(oldChatInstance.id)
+          singleFilePathChatInstancesMap.set(oldChatInstance.singleFilePath,
+            state.getChatInstancesBySingleFilePath(oldChatInstance.singleFilePath)
+              .filter(item => item.id !== oldChatInstance.id))
+
+          // add new map
+          chatIdChatInstanceMap.set(newChatInstance.id, newChatInstance)
+          singleFilePathChatInstancesMap.set(newChatInstance.singleFilePath,
+            [...state.getChatInstancesBySingleFilePath(newChatInstance.singleFilePath), newChatInstance])
+
+          return newChatInstance
+        }
 
         return chatInstance
       }),
     }))
   },
+  updateChatInstances(chatInstances) {
+    const state = get()
+    const finalChatInstances = typeof chatInstances === 'function' ? chatInstances(get().chatInstances) : chatInstances
 
+    // clear old map
+    chatIdChatInstanceMap.clear()
+    singleFilePathChatInstancesMap.clear()
+
+    finalChatInstances.forEach((chatInstance) => {
+      const { id, singleFilePath } = chatInstance
+
+      // add new map
+      singleFilePathChatInstancesMap.set(singleFilePath,
+        [...state.getChatInstancesBySingleFilePath(singleFilePath), chatInstance])
+
+      chatIdChatInstanceMap.set(id, chatInstance)
+    })
+
+    set({ chatInstances: finalChatInstances })
+  },
   removeChatInstance(chatId) {
     const state = get()
 
-    set(state => ({
-      chatInstances: state.chatInstances.filter(chatInstance => chatInstance.id !== chatId),
-    }))
+    state.updateChatInstances(chatInstances => chatInstances.filter(chatInstance => chatInstance.id !== chatId))
 
     const nextSidebarTree = travelTree(state.sidebarTree, (item) => {
       if (item.id === chatId)

packages/gpt-runner-web/client/src/store/zustand/global/sidebar-tree.slice.ts

@@ -106,27 +106,15 @@ export const createSidebarTreeSlice: StateCreator<
 
     const state = get()
 
-    // save current gptFileIdChatIdsMap for later use
-    const currentGptFileIdChatIdsMap: Map<string, string[]> = new Map()
-    travelTree(state.sidebarTree, (treeItem) => {
-      if (treeItem.otherInfo?.type === GptFileTreeItemType.File) {
-        const chatIds: string[] = []
-        treeItem.children?.forEach((child) => {
-          if (child.otherInfo?.type === GptFileTreeItemType.Chat)
-            chatIds.push(child.id)
-        })
-        const prevChatIds = currentGptFileIdChatIdsMap.get(treeItem.id) || []
-        const nextChatIds = [...new Set([...prevChatIds, ...chatIds])]
-        currentGptFileIdChatIdsMap.set(treeItem.id, nextChatIds)
-      }
-    })
+    // refresh map
+    state.updateChatInstances(state.chatInstances)
 
     const fetchGptFilesTreeRes = await fetchGptFilesTree({
       rootPath,
     })
 
     const filesInfoTree = fetchGptFilesTreeRes.data?.filesInfoTree || []
-    const gptFileIds: string[] = []
+    const gptFilePaths: string[] = []
     const treeItems = travelTree(filesInfoTree, (item) => {
       const oldIsExpanded = idTreeItemMap.get(item.id)?.isExpanded
 
@@ -140,12 +128,13 @@ export const createSidebarTreeSlice: StateCreator<
       }
 
       if (item.type === GptFileTreeItemType.File) {
-        gptFileIds.push(item.id)
+        gptFilePaths.push(item.path)
 
-        const chatIds = currentGptFileIdChatIdsMap.get(item.id) || []
+        // const chatIds = currentGptFileIdChatIdsMap.get(item.id) || []
+        const chatInstances = state.getChatInstancesBySingleFilePath(item.path)
 
-        result.children = chatIds.map((chatId) => {
-          const chatInfo = state.getChatInfo(chatId)
+        result.children = chatInstances.map((chatInstance) => {
+          const chatInfo = state.getChatInfo(chatInstance)
           chatInfo.parentId = item.id
 
           return state.chatInfo2SidebarTreeItem(chatInfo)
@@ -155,12 +144,14 @@ export const createSidebarTreeSlice: StateCreator<
       return result
     }) satisfies SidebarTreeItem[]
 
-    // remove gptFileIds that are not in the tree
-    Array.from(currentGptFileIdChatIdsMap.entries()).forEach(([gptFileId, chatIds]) => {
-      if (!gptFileIds.includes(gptFileId)) {
+    // remove gptFilePaths that are not in the tree
+    const oldGptFilePaths = state.chatInstances.map(chatInstance => chatInstance.singleFilePath)
+    oldGptFilePaths.forEach((gptFilePath) => {
+      if (!gptFilePaths.includes(gptFilePath)) {
         // remove chat instances that are not in the tree
-        chatIds.forEach((chatId) => {
-          state.removeChatInstance(chatId)
+        const chatInstances = state.getChatInstancesBySingleFilePath(gptFilePath)
+        chatInstances.forEach((chatInstance) => {
+          state.removeChatInstance(chatInstance.id)
         })
       }
     })
@@ -175,7 +166,7 @@ export const createSidebarTreeSlice: StateCreator<
       inputtingPrompt: gptFileTreeItem?.otherInfo?.singleFileConfig.userPrompt || '',
       systemPrompt: gptFileTreeItem?.otherInfo?.singleFileConfig.systemPrompt || '',
       messages: [],
-      singleFilePath: gptFileTreeItem?.otherInfo?.path || '',
+      singleFilePath: gptFileTreeItem?.path || '',
       status: ChatMessageStatus.Success,
       createAt: Date.now(),
     })