From 2b0a4c433c1687cdbca2e47ed41b3502d49c5014 Mon Sep 17 00:00:00 2001
From: Louis <louis@jan.ai>
Date: Tue, 17 Dec 2024 15:44:13 +0700
Subject: [PATCH] fix: performance issue with atom storage persistence

---
 web/helpers/atoms/ChatMessage.atom.ts | 7 +++++--
 1 file changed, 5 insertions(+), 2 deletions(-)

diff --git a/web/helpers/atoms/ChatMessage.atom.ts b/web/helpers/atoms/ChatMessage.atom.ts
index b0ec6c493..e4259d604 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 } from 'jotai/utils'
+import { atomWithStorage, createJSONStorage } from 'jotai/utils'
 
 import {
   getActiveThreadIdAtom,
@@ -16,13 +16,16 @@ import {
 import { TokenSpeed } from '@/types/token'
 
 const CHAT_MESSAGE_NAME = 'chatMessages'
+const storage = createJSONStorage<Record<string, ThreadMessage[]>>(
+  () => sessionStorage
+)
 /**
  * Stores all chat messages for all threads
  */
 export const chatMessages = atomWithStorage<Record<string, ThreadMessage[]>>(
   CHAT_MESSAGE_NAME,
   {},
-  undefined,
+  storage,
   { getOnInit: true }
 )