Best practice for complex persistent design

[gengjiawen]

Local Storage has a limit of 5Mb design. I want to store it in Indexed DB. The structure is consistent of Session and Message.
I only want to store messageId in session when persistent. But when I use it. I want to Message[] with actual content.
This is currently I made. wonder is this the good way todo it.

import { create } from 'zustand'
import { subscribeWithSelector } from 'zustand/middleware'
import * as idb from 'idb-keyval'

interface Message {
  id: string
  sessionId: string
  content: string
  role: 'user' | 'assistant'
  timestamp: number
}

interface ChatSession {
  id: string
  name: string
  createdAt: number
  lastUpdated: number
  messages: Message[] // In-memory messages
}

interface ChatStore {
  sessions: ChatSession[]
  currentSessionId: string | null
  isLoading: boolean
  error: string | null
  // Actions
  addSession: (name: string) => Promise<void>
  removeSession: (sessionId: string) => Promise<void>
  setCurrentSession: (sessionId: string) => void
  setError: (error: string | null) => void
}

const getMessageKey = (messageId: string) => `message:${messageId}`
const getSessionListKey = () => 'sessionList'

const serializeSession = (
  session: ChatSession
): Omit<ChatSession, 'messages'> & { messageIds: string[] } => {
  const { messages, ...rest } = session // Exclude messages
  const messageIds = messages.map((msg) => msg.id)
  return {
    ...rest, // Spread the rest of the session properties without messages
    messageIds,
  }
}

const deserializeSession = async (
  sessionData: Omit<ChatSession, 'messages'> & { messageIds: string[] }
): Promise<ChatSession> => {
  // Replace messageIds array with the actual messages
  const messages = await Promise.all(
    sessionData.messageIds.map((id) => idb.get<Message>(getMessageKey(id)))
  )
  const validMessages = messages.filter((msg): msg is Message => msg != null)
  return {
    ...sessionData,
    messages: validMessages,
  }
}

const useChatStore = create<ChatStore>()(
  subscribeWithSelector((set, get) => ({
    sessions: [],
    currentSessionId: null,
    isLoading: false,
    error: null,

    setError: (error) => set({ error }),

    addSession: async (name) => {
      try {
        set({ isLoading: true, error: null })

        const newSession: ChatSession = {
          id: crypto.randomUUID(),
          name,
          createdAt: Date.now(),
          lastUpdated: Date.now(),
          messages: [], // Initialize empty messages array
        }

        // Update sessions list
        const currentSessions = [...get().sessions, newSession]

        set({
          sessions: currentSessions,
          currentSessionId: newSession.id,
          isLoading: false,
        })
        // No need to manually save to IndexedDB here, the subscription will handle it
      } catch (error) {
        set({
          error: error instanceof Error ? error.message : 'An error occurred',
          isLoading: false,
        })
      }
    },

    removeSession: async (sessionId) => {
      try {
        set({ isLoading: true, error: null })

        // Get session's messages
        const session = get().sessions.find((s) => s.id === sessionId)
        if (session) {
          // Delete all messages from IndexedDB
          await Promise.all(
            session.messages.map((msg) => idb.del(getMessageKey(msg.id)))
          )
        }

        // Update sessions list
        const updatedSessions = get().sessions.filter(
          (session) => session.id !== sessionId
        )

        set((state) => ({
          sessions: updatedSessions,
          currentSessionId:
            state.currentSessionId === sessionId
              ? updatedSessions[0]?.id ?? null
              : state.currentSessionId,
          isLoading: false,
        }))
        // The subscription will handle saving the updated sessions
      } catch (error) {
        set({
          error: error instanceof Error ? error.message : 'An error occurred',
          isLoading: false,
        })
      }
    },

    setCurrentSession: (sessionId) => {
      set({ currentSessionId: sessionId })
    },
  }))
)

// Subscribe to sessions changes and persist to IndexedDB
useChatStore.subscribe(
  (state) => state.sessions,
  async (sessions) => {
    // Whenever sessions change, save to IndexedDB
    const serializedSessions = sessions.map(serializeSession)
    await idb.set(getSessionListKey(), serializedSessions)

    // Save or delete messages as needed
    const allMessageIds = new Set(
      sessions.flatMap((session) => session.messages.map((msg) => msg.id))
    )

    // Get all message keys in IndexedDB
    const allKeys = await idb.keys()
    const messageKeys = allKeys.filter(
      (key) => typeof key === 'string' && key.startsWith('message:')
    )

    // Messages that are no longer in sessions need to be deleted
    const keysToDelete = messageKeys.filter(
      (key) => !allMessageIds.has(key.replace('message:', ''))
    )
    await Promise.all(keysToDelete.map((key) => idb.del(key)))

    // Save current messages
    const messagesToSave = sessions.flatMap((session) => session.messages)
    await Promise.all(
      messagesToSave.map((msg) => idb.set(getMessageKey(msg.id), msg))
    )
  }
)

// Initialize store by loading sessions list and messages
const initializeStore = async () => {
  try {
    // Load serialized sessions from IndexedDB
    const serializedSessions =
      (await idb.get<(Omit<ChatSession, 'messages'> & { messageIds: string[] })[]>(
        getSessionListKey()
      )) || []

    // Deserialize sessions
    const sessions = await Promise.all(
      serializedSessions.map((sessionData) => deserializeSession(sessionData))
    )

    useChatStore.setState({ sessions })
  } catch (error) {
    console.error('Failed to initialize store:', error)
  }
}

initializeStore()

export default useChatStore