diff --git a/web/helpers/atoms/ChatMessage.atom.ts b/web/helpers/atoms/ChatMessage.atom.ts index e4259d604..7cdeb6946 100644 --- a/web/helpers/atoms/ChatMessage.atom.ts +++ b/web/helpers/atoms/ChatMessage.atom.ts @@ -6,7 +6,7 @@ import { } from '@janhq/core' import { atom } from 'jotai' -import { atomWithStorage, createJSONStorage } from 'jotai/utils' +import { atomWithStorage } from 'jotai/utils' import { getActiveThreadIdAtom, @@ -16,17 +16,23 @@ import { import { TokenSpeed } from '@/types/token' const CHAT_MESSAGE_NAME = 'chatMessages' -const storage = createJSONStorage>( - () => sessionStorage -) /** * Stores all chat messages for all threads */ -export const chatMessages = atomWithStorage>( - CHAT_MESSAGE_NAME, - {}, - storage, - { getOnInit: true } +export const chatMessagesStorage = atomWithStorage< + Record +>(CHAT_MESSAGE_NAME, {}, undefined, { getOnInit: true }) + +export const cachedMessages = atom>() +/** + * Retrieve chat messages for all threads + */ +export const chatMessages = atom( + (get) => get(cachedMessages) ?? get(chatMessagesStorage), + (_get, set, newValue: Record) => { + set(cachedMessages, newValue) + ;(() => set(chatMessagesStorage, newValue))() + } ) /**