feat(ai-chat): integrate session hydration and enhance chat history loading logic

- Added `hydrateSessionsFromLocalDb` to initialize app for session management.
- Updated `ChatHistoryDropdown` to conditionally render loading state based on session length.
- Refactored `useChatHistory` to utilize the AI chat session store for state management.
- Enhanced `AIChatSessionService` to include session synchronization and error handling.
- Exported store actions from `ai-chat-session` module for better state management.

Signed-off-by: Innei <tukon479@gmail.com>

apps/desktop/layer/renderer/src/initialize/index.ts

@@ -7,6 +7,7 @@ import { repository } from "@pkg"
 import { enableMapSet } from "immer"
 
 import { initI18n } from "~/i18n"
+import { hydrateSessionsFromLocalDb } from "~/modules/ai-chat-session"
 import { settingSyncQueue } from "~/modules/settings/helper/sync-queue"
 import { ElectronCloseEvent, ElectronShowEvent } from "~/providers/invalidate-query-provider"
 
@@ -55,6 +56,7 @@ export const initializeApp = async () => {
   initializeDayjs()
   registerHistoryStack()
 
+  hydrateSessionsFromLocalDb()
   // Set Environment
   document.documentElement.dataset.buildType = ELECTRON_BUILD ? "electron" : "web"
 

apps/desktop/layer/renderer/src/modules/ai-chat-session/index.ts

@@ -1,2 +1,3 @@
 export * from "./query"
 export * from "./service"
+export * from "./store"

apps/desktop/layer/renderer/src/modules/ai-chat-session/service.ts

@@ -1,10 +1,13 @@
 import type { AIChatMessage, AIChatSession, ListSessionsQuery } from "@follow-app/client-sdk"
 
+import { getI18n } from "~/i18n"
+
 import { followApi } from "../../lib/api-client"
 import { queryClient } from "../../lib/query-client"
 import { AIPersistService } from "../ai-chat/services"
 import type { BizUIMessage, BizUIMetadata } from "../ai-chat/store/types"
 import { aiChatSessionKeys } from "./query"
+import { aiChatSessionStoreActions } from "./store"
 
 // Hard cap on pagination to prevent excessive API calls while keeping initial sync fast.
 const MAX_PAGES = 10
@@ -29,42 +32,99 @@ class AIChatSessionServiceStatic {
     if (this.sync) {
       return { sessions: 0, messages: 0, failures: 0 }
     }
+
     this.sync = true
-    const res = await followApi.aiChatSessions.list(filters)
-    const sessions = res.data
 
-    let totalMessages = 0
-    let failures = 0
+    aiChatSessionStoreActions.setSyncing(true)
+    aiChatSessionStoreActions.clearError()
 
-    for (const session of sessions) {
-      const dbSession = await AIPersistService.getChatSession(session.chatId)
-      const lastUpdatedAt = dbSession ? dbSession.updatedAt : new Date(0)
-      if (lastUpdatedAt >= new Date(session.updatedAt)) {
-        // If local session is already up-to-date, skip fetching messages
-        continue
-      }
-      try {
-        // Ensure session exists locally first
-        await AIPersistService.ensureSession(session.chatId, {
-          title: session.title,
-          createdAt: new Date(session.createdAt),
-          updatedAt: new Date(session.updatedAt),
-        })
+    await this.loadSessionsFromDb().catch((error) => {
+      console.error("syncSessionsAndMessagesFromServer: failed to load local sessions", error)
+    })
 
-        // Fetch unseen messages (newer than lastSeenAt) without changing read state remotely
-        const unseenMessages = await this.fetchUnseenRemoteMessages(session.chatId, lastUpdatedAt)
-        if (unseenMessages.length === 0) continue
-
-        const normalized = unseenMessages.map(this.normalizeRemoteMessage)
-        await AIPersistService.upsertMessages(session.chatId, normalized)
-        totalMessages += normalized.length
-      } catch (err) {
-        // Keep going even if a single session fails
-        console.error("syncSessionsAndMessagesFromServer: failed for session", session.chatId, err)
-        failures += 1
-      }
+    let summary: { sessions: number; messages: number; failures: number } = {
+      sessions: 0,
+      messages: 0,
+      failures: 0,
     }
-    return { sessions: sessions.length, messages: totalMessages, failures }
+
+    try {
+      const res = await followApi.aiChatSessions.list(filters)
+      const sessions = res.data
+
+      let totalMessages = 0
+      let failures = 0
+
+      await Promise.allSettled(
+        sessions.map((session) =>
+          this.processSessionSync(session)
+            .then((count) => {
+              totalMessages += count
+            })
+            .catch((err) => {
+              console.error(
+                "syncSessionsAndMessagesFromServer: failed for session",
+                session.chatId,
+                err,
+              )
+              failures += 1
+            }),
+        ),
+      )
+
+      summary = { sessions: sessions.length, messages: totalMessages, failures }
+
+      await this.loadSessionsFromDb()
+      aiChatSessionStoreActions.setStats(summary)
+      aiChatSessionStoreActions.setLastSyncedAt(new Date())
+      return summary
+    } catch (error) {
+      aiChatSessionStoreActions.setError(
+        error instanceof Error ? error.message : "Failed to sync chat sessions",
+      )
+      throw error
+    } finally {
+      aiChatSessionStoreActions.setSyncing(false)
+      this.sync = false
+    }
+  }
+
+  private async processSessionSync(session: AIChatSession): Promise<number> {
+    const dbSession = await AIPersistService.getChatSession(session.chatId)
+    const lastUpdatedAt = dbSession ? dbSession.updatedAt : new Date(0)
+    if (lastUpdatedAt >= new Date(session.updatedAt)) {
+      return 0
+    }
+
+    await AIPersistService.ensureSession(session.chatId, {
+      title: session.title,
+      createdAt: new Date(session.createdAt),
+      updatedAt: new Date(session.updatedAt),
+    })
+
+    const unseenMessages = await this.fetchUnseenRemoteMessages(session.chatId, lastUpdatedAt)
+    if (unseenMessages.length === 0) {
+      return 0
+    }
+
+    const normalized = unseenMessages.map(this.normalizeRemoteMessage)
+    await AIPersistService.upsertMessages(session.chatId, normalized)
+
+    return normalized.length
+  }
+
+  async loadSessionsFromDb() {
+    const { t } = getI18n()
+    const rows = await AIPersistService.getChatSessions()
+    const normalized = rows.map((row) => ({
+      chatId: row.chatId,
+      title: row.title || t("ai:common.new_chat"),
+      createdAt: new Date(row.createdAt),
+      updatedAt: new Date(row.updatedAt),
+      messageCount: row.messageCount,
+    }))
+    aiChatSessionStoreActions.setSessions(normalized)
+    return normalized
   }
   /**
    * Fetch messages for a chat session from the remote API and persist (upsert) them locally.
@@ -97,6 +157,9 @@ class AIChatSessionServiceStatic {
       lastSeenAt: new Date().toISOString(),
     })
 
+    await this.loadSessionsFromDb()
+    aiChatSessionStoreActions.setLastSyncedAt(new Date())
+
     // Invalidate related queries so UI updates outside of hook-based mutation flows
     Promise.all([
       queryClient.invalidateQueries({ queryKey: aiChatSessionKeys.detail(session.chatId) }),

apps/desktop/layer/renderer/src/modules/ai-chat-session/store.ts

@@ -0,0 +1,75 @@
+import { createWithEqualityFn } from "zustand/traditional"
+
+import type { ChatSession } from "../ai-chat/types/ChatSession"
+import { AIChatSessionService } from "./service"
+
+export interface AIChatSessionSyncStats {
+  sessions: number
+  messages: number
+  failures: number
+}
+
+export interface AIChatSessionViewModelState {
+  sessions: ChatSession[]
+  isLoading: boolean
+  isSyncing: boolean
+  stats: AIChatSessionSyncStats
+  lastSyncedAt?: Date
+  error?: string
+}
+
+const createInitialStats = (): AIChatSessionSyncStats => ({
+  sessions: 0,
+  messages: 0,
+  failures: 0,
+})
+
+export const useAIChatSessionStore = createWithEqualityFn<AIChatSessionViewModelState>(() => ({
+  sessions: [],
+  isLoading: false,
+  isSyncing: false,
+  stats: createInitialStats(),
+  lastSyncedAt: undefined,
+  error: undefined,
+}))
+
+const { setState } = useAIChatSessionStore
+
+export const aiChatSessionStoreActions = {
+  setSessions: (sessions: ChatSession[]) =>
+    setState({
+      sessions,
+    }),
+  setLoading: (isLoading: boolean) =>
+    setState({
+      isLoading,
+    }),
+  setSyncing: (isSyncing: boolean) =>
+    setState({
+      isSyncing,
+    }),
+  setStats: (stats: AIChatSessionSyncStats) =>
+    setState({
+      stats: { ...stats },
+    }),
+  resetStats: () =>
+    setState({
+      stats: createInitialStats(),
+    }),
+  setLastSyncedAt: (lastSyncedAt?: Date) =>
+    setState({
+      lastSyncedAt,
+    }),
+  setError: (error?: string) =>
+    setState({
+      error,
+    }),
+  clearError: () =>
+    setState({
+      error: undefined,
+    }),
+}
+
+export const hydrateSessionsFromLocalDb = async () => {
+  return AIChatSessionService.loadSessionsFromDb()
+}

apps/desktop/layer/renderer/src/modules/ai-chat/components/layouts/ChatHistoryDropdown.tsx

@@ -95,7 +95,7 @@ export const ChatHistoryDropdown = ({
         {triggerElement || defaultTrigger}
       </DropdownMenuTrigger>
       <DropdownMenuContent align="start" className="max-h-96 w-72 overflow-y-auto">
-        {loading ? (
+        {loading && sessions.length === 0 ? (
           <div className="flex items-center justify-center py-8">
             <i className="i-mgc-loading-3-cute-re size-5 animate-spin text-text-secondary" />
           </div>

apps/desktop/layer/renderer/src/modules/ai-chat/hooks/useChatHistory.ts

@@ -1,41 +1,42 @@
-import { useCallback, useState } from "react"
+import { useCallback, useMemo } from "react"
 
-import { getI18n } from "~/i18n"
-import { AIPersistService } from "~/modules/ai-chat/services/index"
 import { AIChatSessionService } from "~/modules/ai-chat-session/service"
-
-import type { ChatSession } from "../types/ChatSession"
+import { aiChatSessionStoreActions, useAIChatSessionStore } from "~/modules/ai-chat-session/store"
 
 export const useChatHistory = () => {
-  const [sessions, setSessions] = useState<ChatSession[]>([])
-  const [loading, setLoading] = useState(false)
+  const state = useAIChatSessionStore((s) => s)
+
+  const { sessions } = state
+  const loading = state.isLoading || state.isSyncing
 
   const loadHistory = useCallback(async () => {
-    setLoading(true)
+    if (loading) return
+
+    aiChatSessionStoreActions.setLoading(true)
+    aiChatSessionStoreActions.clearError()
 
-    const { t } = getI18n()
     try {
-      await AIChatSessionService.syncSessionsAndMessagesFromServer()
-      const result = await AIPersistService.getChatSessions()
-      const sessions: ChatSession[] = result.map((row) => ({
-        chatId: row.chatId,
-        title: row.title || t("ai:common.new_chat"),
-        createdAt: new Date(row.createdAt),
-        updatedAt: new Date(row.updatedAt),
-        messageCount: row.messageCount,
-      }))
+      aiChatSessionStoreActions.setSyncing(true)
 
-      setSessions(sessions)
+      await AIChatSessionService.syncSessionsAndMessagesFromServer()
     } catch (error) {
       console.error("Failed to load chat history:", error)
+      aiChatSessionStoreActions.setError(error instanceof Error ? error.message : "Unknown error")
     } finally {
-      setLoading(false)
+      aiChatSessionStoreActions.setSyncing(false)
+      aiChatSessionStoreActions.setLoading(false)
     }
-  }, [])
+  }, [loading])
 
-  return {
-    sessions,
-    loading,
-    loadHistory,
-  }
+  return useMemo(
+    () => ({
+      sessions,
+      loading,
+      loadHistory,
+      stats: state.stats,
+      lastSyncedAt: state.lastSyncedAt,
+      error: state.error,
+    }),
+    [sessions, loading, loadHistory, state.stats, state.lastSyncedAt, state.error],
+  )
 }