From 1395aa40eb2997774cebb525c0aaf93d3e24119b Mon Sep 17 00:00:00 2001
From: Louis <louis@jan.ai>
Date: Tue, 17 Dec 2024 16:44:34 +0700
Subject: [PATCH] fix: use composing atom to switch between cache and local
 storage

---
 web/helpers/atoms/ChatMessage.atom.ts | 24 +++++++++++++++---------
 1 file changed, 15 insertions(+), 9 deletions(-)

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<Record<string, ThreadMessage[]>>(
-  () => sessionStorage
-)
 /**
  * Stores all chat messages for all threads
  */
-export const chatMessages = atomWithStorage<Record<string, ThreadMessage[]>>(
-  CHAT_MESSAGE_NAME,
-  {},
-  storage,
-  { getOnInit: true }
+export const chatMessagesStorage = atomWithStorage<
+  Record<string, ThreadMessage[]>
+>(CHAT_MESSAGE_NAME, {}, undefined, { getOnInit: true })
+
+export const cachedMessages = atom<Record<string, ThreadMessage[]>>()
+/**
+ * Retrieve chat messages for all threads
+ */
+export const chatMessages = atom(
+  (get) => get(cachedMessages) ?? get(chatMessagesStorage),
+  (_get, set, newValue: Record<string, ThreadMessage[]>) => {
+    set(cachedMessages, newValue)
+    ;(() => set(chatMessagesStorage, newValue))()
+  }
 )
 
 /**