From 079f2060441ffd9e2ba689b4a680bc3969c9ffec Mon Sep 17 00:00:00 2001 From: Louis Date: Thu, 12 Jun 2025 09:10:00 +0700 Subject: [PATCH] fix: avoid persisting threads and messages on local storage (#5249) --- web-app/src/hooks/useMessages.ts | 100 +++--- web-app/src/hooks/useThreads.ts | 534 +++++++++++++++---------------- 2 files changed, 306 insertions(+), 328 deletions(-) diff --git a/web-app/src/hooks/useMessages.ts b/web-app/src/hooks/useMessages.ts index 4c38e73b7..3a83b5a48 100644 --- a/web-app/src/hooks/useMessages.ts +++ b/web-app/src/hooks/useMessages.ts @@ -1,7 +1,5 @@ import { create } from 'zustand' import { ThreadMessage } from '@janhq/core' -import { createJSONStorage, persist } from 'zustand/middleware' -import { localStorageKey } from '@/constants/localStorage' import { createMessage, deleteMessage as deleteMessageExt, @@ -16,59 +14,51 @@ type MessageState = { deleteMessage: (threadId: string, messageId: string) => void } -export const useMessages = create()( - persist( - (set, get) => ({ - messages: {}, - getMessages: (threadId) => { - return get().messages[threadId] || [] +export const useMessages = create()((set, get) => ({ + messages: {}, + getMessages: (threadId) => { + return get().messages[threadId] || [] + }, + setMessages: (threadId, messages) => { + set((state) => ({ + messages: { + ...state.messages, + [threadId]: messages, }, - setMessages: (threadId, messages) => { - set((state) => ({ - messages: { - ...state.messages, - [threadId]: messages, - }, - })) + })) + }, + addMessage: (message) => { + const currentAssistant = useAssistant.getState().currentAssistant + const newMessage = { + ...message, + created_at: message.created_at || Date.now(), + metadata: { + ...message.metadata, + assistant: currentAssistant, }, - addMessage: (message) => { - const currentAssistant = useAssistant.getState().currentAssistant - const newMessage = { - ...message, - created_at: message.created_at || Date.now(), - metadata: { - ...message.metadata, - assistant: currentAssistant, - }, - } - createMessage(newMessage).then((createdMessage) => { - set((state) => ({ - messages: { - ...state.messages, - [message.thread_id]: [ - ...(state.messages[message.thread_id] || []), - createdMessage, - ], - }, - })) - }) - }, - deleteMessage: (threadId, messageId) => { - deleteMessageExt(threadId, messageId) - set((state) => ({ - messages: { - ...state.messages, - [threadId]: - state.messages[threadId]?.filter( - (message) => message.id !== messageId - ) || [], - }, - })) - }, - }), - { - name: localStorageKey.messages, - storage: createJSONStorage(() => localStorage), } - ) -) + createMessage(newMessage).then((createdMessage) => { + set((state) => ({ + messages: { + ...state.messages, + [message.thread_id]: [ + ...(state.messages[message.thread_id] || []), + createdMessage, + ], + }, + })) + }) + }, + deleteMessage: (threadId, messageId) => { + deleteMessageExt(threadId, messageId) + set((state) => ({ + messages: { + ...state.messages, + [threadId]: + state.messages[threadId]?.filter( + (message) => message.id !== messageId + ) || [], + }, + })) + }, +})) diff --git a/web-app/src/hooks/useThreads.ts b/web-app/src/hooks/useThreads.ts index 5118974c6..806749b73 100644 --- a/web-app/src/hooks/useThreads.ts +++ b/web-app/src/hooks/useThreads.ts @@ -1,6 +1,4 @@ import { create } from 'zustand' -import { persist, createJSONStorage } from 'zustand/middleware' -import { localStorageKey } from '@/constants/localStorage' import { ulid } from 'ulidx' import { createThread, deleteThread, updateThread } from '@/services/threads' import { Fzf } from 'fzf' @@ -30,281 +28,271 @@ type ThreadState = { searchIndex: Fzf | null } -export const useThreads = create()( - persist( - (set, get) => ({ - threads: {}, - searchIndex: null, - setThreads: (threads) => { - threads.forEach((thread, index) => { - thread.order = index + 1 - updateThread({ - ...thread, - order: index + 1, - }) - }) - const threadMap = threads.reduce( - (acc: Record, thread) => { - acc[thread.id] = thread - return acc - }, - {} as Record - ) - set({ - threads: threadMap, - searchIndex: new Fzf(Object.values(threadMap), { - selector: (item: Thread) => item.title, - }), - }) +export const useThreads = create()((set, get) => ({ + threads: {}, + searchIndex: null, + setThreads: (threads) => { + threads.forEach((thread, index) => { + thread.order = index + 1 + updateThread({ + ...thread, + order: index + 1, + }) + }) + const threadMap = threads.reduce( + (acc: Record, thread) => { + acc[thread.id] = thread + return acc }, - getFilteredThreads: (searchTerm: string) => { - const { threads, searchIndex } = get() + {} as Record + ) + set({ + threads: threadMap, + searchIndex: new Fzf(Object.values(threadMap), { + selector: (item: Thread) => item.title, + }), + }) + }, + getFilteredThreads: (searchTerm: string) => { + const { threads, searchIndex } = get() - // If no search term, return all threads - if (!searchTerm) { - // return all threads - return Object.values(threads) + // If no search term, return all threads + if (!searchTerm) { + // return all threads + return Object.values(threads) + } + + let currentIndex = searchIndex + if (!currentIndex?.find) { + currentIndex = new Fzf(Object.values(threads), { + selector: (item: Thread) => item.title, + }) + set({ searchIndex: currentIndex }) + } + + // Use the index to search and return matching threads + const fzfResults = currentIndex.find(searchTerm) + return fzfResults.map( + (result: { item: Thread; positions: Set }) => { + const thread = result.item // Fzf stores the original item here + // Ensure result.positions is an array, default to empty if undefined + const positions = Array.from(result.positions) || [] + const highlightedTitle = highlightFzfMatch(thread.title, positions) + return { + ...thread, + title: highlightedTitle, // Override title with highlighted version } - - let currentIndex = searchIndex - if (!currentIndex?.find) { - currentIndex = new Fzf(Object.values(threads), { - selector: (item: Thread) => item.title, - }) - set({ searchIndex: currentIndex }) - } - - // Use the index to search and return matching threads - const fzfResults = currentIndex.find(searchTerm) - return fzfResults.map( - (result: { item: Thread; positions: Set }) => { - const thread = result.item // Fzf stores the original item here - // Ensure result.positions is an array, default to empty if undefined - const positions = Array.from(result.positions) || [] - const highlightedTitle = highlightFzfMatch(thread.title, positions) - return { - ...thread, - title: highlightedTitle, // Override title with highlighted version - } - } - ) - }, - toggleFavorite: (threadId) => { - set((state) => { - updateThread({ + } + ) + }, + toggleFavorite: (threadId) => { + set((state) => { + updateThread({ + ...state.threads[threadId], + isFavorite: !state.threads[threadId].isFavorite, + }) + return { + threads: { + ...state.threads, + [threadId]: { ...state.threads[threadId], isFavorite: !state.threads[threadId].isFavorite, - }) - return { - threads: { - ...state.threads, - [threadId]: { - ...state.threads[threadId], - isFavorite: !state.threads[threadId].isFavorite, - }, - }, + }, + }, + } + }) + }, + deleteThread: (threadId) => { + set((state) => { + // eslint-disable-next-line @typescript-eslint/no-unused-vars + const { [threadId]: _, ...remainingThreads } = state.threads + deleteThread(threadId) + return { + threads: remainingThreads, + searchIndex: new Fzf(Object.values(remainingThreads), { + selector: (item: Thread) => item.title, + }), + } + }) + }, + deleteAllThreads: () => { + set((state) => { + const allThreadIds = Object.keys(state.threads) + allThreadIds.forEach((threadId) => { + deleteThread(threadId) + }) + return { + threads: {}, + searchIndex: null, // Or new Fzf([], {selector...}) + } + }) + }, + unstarAllThreads: () => { + set((state) => { + const updatedThreads = Object.keys(state.threads).reduce( + (acc, threadId) => { + acc[threadId] = { + ...state.threads[threadId], + isFavorite: false, } - }) - }, - deleteThread: (threadId) => { - set((state) => { - // eslint-disable-next-line @typescript-eslint/no-unused-vars - const { [threadId]: _, ...remainingThreads } = state.threads - deleteThread(threadId) - return { - threads: remainingThreads, - searchIndex: new Fzf(Object.values(remainingThreads), { - selector: (item: Thread) => item.title, - }), - } - }) - }, - deleteAllThreads: () => { - set((state) => { - const allThreadIds = Object.keys(state.threads) - allThreadIds.forEach((threadId) => { - deleteThread(threadId) - }) - return { - threads: {}, - searchIndex: null, // Or new Fzf([], {selector...}) - } - }) - }, - unstarAllThreads: () => { - set((state) => { - const updatedThreads = Object.keys(state.threads).reduce( - (acc, threadId) => { - acc[threadId] = { - ...state.threads[threadId], - isFavorite: false, - } - return acc - }, - {} as Record - ) - Object.values(updatedThreads).forEach((thread) => { - updateThread({ ...thread, isFavorite: false }) - }) - return { threads: updatedThreads } - }) - }, - getFavoriteThreads: () => { - return Object.values(get().threads).filter( - (thread) => thread.isFavorite - ) - }, - getThreadById: (threadId: string) => { - return get().threads[threadId] - }, - setCurrentThreadId: (threadId) => { - set({ currentThreadId: threadId }) - }, - createThread: async (model, title, assistant) => { - const newThread: Thread = { - id: ulid(), - title: title ?? 'New Thread', - model, - // order: 1, // Will be set properly by setThreads - updated: Date.now() / 1000, - assistants: assistant ? [assistant] : [], - } - return await createThread(newThread).then((createdThread) => { - set((state) => { - // Get all existing threads as an array - const existingThreads = Object.values(state.threads) - - // Create new array with the new thread at the beginning - const reorderedThreads = [createdThread, ...existingThreads] - - // Use setThreads to handle proper ordering (this will assign order 1, 2, 3...) - get().setThreads(reorderedThreads) - - return { - currentThreadId: createdThread.id, - } - }) - return createdThread - }) - }, - updateCurrentThreadAssistant: (assistant) => { - set((state) => { - if (!state.currentThreadId) return { ...state } - const currentThread = state.getCurrentThread() - if (currentThread) - updateThread({ - ...currentThread, - assistants: [{ ...assistant, model: currentThread.model }], - }) - return { - threads: { - ...state.threads, - [state.currentThreadId as string]: { - ...state.threads[state.currentThreadId as string], - assistants: [assistant], - }, - }, - } - }) - }, - updateCurrentThreadModel: (model) => { - set((state) => { - if (!state.currentThreadId) return { ...state } - const currentThread = state.getCurrentThread() - if (currentThread) updateThread({ ...currentThread, model }) - return { - threads: { - ...state.threads, - [state.currentThreadId as string]: { - ...state.threads[state.currentThreadId as string], - model, - }, - }, - } - }) - }, - renameThread: (threadId, newTitle) => { - set((state) => { - const thread = state.threads[threadId] - if (!thread) return state - const updatedThread = { - ...thread, - title: newTitle, - } - updateThread(updatedThread) // External call, order is fine - const newThreads = { ...state.threads, [threadId]: updatedThread } - return { - threads: newThreads, - searchIndex: new Fzf(Object.values(newThreads), { - selector: (item: Thread) => item.title, - }), - } - }) - }, - getCurrentThread: () => { - const { currentThreadId, threads } = get() - return currentThreadId ? threads[currentThreadId] : undefined - }, - updateThreadTimestamp: (threadId) => { - set((state) => { - const thread = state.threads[threadId] - if (!thread) return state - - // If the thread is already at order 1, just update the timestamp - if (thread.order === 1) { - const updatedThread = { - ...thread, - updated: Date.now() / 1000, - } - updateThread(updatedThread) - - return { - threads: { - ...state.threads, - [threadId]: updatedThread, - }, - } - } - - // Update the thread with new timestamp and set it to order 1 (top) - const updatedThread = { - ...thread, - updated: Date.now() / 1000, - order: 1, - } - - // Update all other threads to increment their order by 1 - const updatedThreads = { ...state.threads } - Object.keys(updatedThreads).forEach(id => { - if (id !== threadId) { - const otherThread = updatedThreads[id] - updatedThreads[id] = { - ...otherThread, - order: (otherThread.order || 1) + 1, - } - // Update the backend for other threads - updateThread(updatedThreads[id]) - } - }) - - // Set the updated thread - updatedThreads[threadId] = updatedThread - - // Update the backend for the main thread - updateThread(updatedThread) - - return { - threads: updatedThreads, - searchIndex: new Fzf(Object.values(updatedThreads), { - selector: (item: Thread) => item.title, - }), - } - }) - }, - }), - { - name: localStorageKey.threads, - storage: createJSONStorage(() => localStorage), + return acc + }, + {} as Record + ) + Object.values(updatedThreads).forEach((thread) => { + updateThread({ ...thread, isFavorite: false }) + }) + return { threads: updatedThreads } + }) + }, + getFavoriteThreads: () => { + return Object.values(get().threads).filter((thread) => thread.isFavorite) + }, + getThreadById: (threadId: string) => { + return get().threads[threadId] + }, + setCurrentThreadId: (threadId) => { + set({ currentThreadId: threadId }) + }, + createThread: async (model, title, assistant) => { + const newThread: Thread = { + id: ulid(), + title: title ?? 'New Thread', + model, + // order: 1, // Will be set properly by setThreads + updated: Date.now() / 1000, + assistants: assistant ? [assistant] : [], } - ) -) + return await createThread(newThread).then((createdThread) => { + set((state) => { + // Get all existing threads as an array + const existingThreads = Object.values(state.threads) + + // Create new array with the new thread at the beginning + const reorderedThreads = [createdThread, ...existingThreads] + + // Use setThreads to handle proper ordering (this will assign order 1, 2, 3...) + get().setThreads(reorderedThreads) + + return { + currentThreadId: createdThread.id, + } + }) + return createdThread + }) + }, + updateCurrentThreadAssistant: (assistant) => { + set((state) => { + if (!state.currentThreadId) return { ...state } + const currentThread = state.getCurrentThread() + if (currentThread) + updateThread({ + ...currentThread, + assistants: [{ ...assistant, model: currentThread.model }], + }) + return { + threads: { + ...state.threads, + [state.currentThreadId as string]: { + ...state.threads[state.currentThreadId as string], + assistants: [assistant], + }, + }, + } + }) + }, + updateCurrentThreadModel: (model) => { + set((state) => { + if (!state.currentThreadId) return { ...state } + const currentThread = state.getCurrentThread() + if (currentThread) updateThread({ ...currentThread, model }) + return { + threads: { + ...state.threads, + [state.currentThreadId as string]: { + ...state.threads[state.currentThreadId as string], + model, + }, + }, + } + }) + }, + renameThread: (threadId, newTitle) => { + set((state) => { + const thread = state.threads[threadId] + if (!thread) return state + const updatedThread = { + ...thread, + title: newTitle, + } + updateThread(updatedThread) // External call, order is fine + const newThreads = { ...state.threads, [threadId]: updatedThread } + return { + threads: newThreads, + searchIndex: new Fzf(Object.values(newThreads), { + selector: (item: Thread) => item.title, + }), + } + }) + }, + getCurrentThread: () => { + const { currentThreadId, threads } = get() + return currentThreadId ? threads[currentThreadId] : undefined + }, + updateThreadTimestamp: (threadId) => { + set((state) => { + const thread = state.threads[threadId] + if (!thread) return state + + // If the thread is already at order 1, just update the timestamp + if (thread.order === 1) { + const updatedThread = { + ...thread, + updated: Date.now() / 1000, + } + updateThread(updatedThread) + + return { + threads: { + ...state.threads, + [threadId]: updatedThread, + }, + } + } + + // Update the thread with new timestamp and set it to order 1 (top) + const updatedThread = { + ...thread, + updated: Date.now() / 1000, + order: 1, + } + + // Update all other threads to increment their order by 1 + const updatedThreads = { ...state.threads } + Object.keys(updatedThreads).forEach((id) => { + if (id !== threadId) { + const otherThread = updatedThreads[id] + updatedThreads[id] = { + ...otherThread, + order: (otherThread.order || 1) + 1, + } + // Update the backend for other threads + updateThread(updatedThreads[id]) + } + }) + + // Set the updated thread + updatedThreads[threadId] = updatedThread + + // Update the backend for the main thread + updateThread(updatedThread) + + return { + threads: updatedThreads, + searchIndex: new Fzf(Object.values(updatedThreads), { + selector: (item: Thread) => item.title, + }), + } + }) + }, +}))