fix: skip streaming messages in message persist (#4637)

* fix: skip streaming messages in message processing

* fix(ai-chat): streamline message synchronization logic

Removed redundant message loading check in syncSessionMessages method to improve efficiency. Now directly fetches and persists messages when necessary.

* fix(ai-chat): update session time handling after message persistence

Ensure session time is updated based on the latest message creation date after successfully saving messages.

* fix: simplify message sync logic

* fix: update session updatedAt to use createdAt during sync

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

@@ -47,7 +47,8 @@ class AIChatSessionServiceStatic {
         await AIPersistService.ensureSession(session.chatId, {
           title: session.title,
           createdAt: new Date(session.createdAt),
-          updatedAt: new Date(session.updatedAt),
+          // Use createdAt for updatedAt as we are syncing session instead of messages
+          updatedAt: new Date(session.createdAt),
         })
       })
       await this.loadSessionsFromDb()
@@ -100,7 +101,9 @@ class AIChatSessionServiceStatic {
     await AIPersistService.ensureSession(session.chatId, {
       title: session.title,
       createdAt: new Date(session.createdAt),
-      updatedAt: new Date(session.updatedAt),
+      // Use createdAt for updatedAt
+      // Because we are fetching session data instead of messages
+      updatedAt: new Date(session.createdAt),
     })
     await AIPersistService.upsertMessages(session.chatId, normalized)
 
@@ -117,11 +120,6 @@ class AIChatSessionServiceStatic {
 
   async syncSessionMessages(chatId: string) {
     try {
-      const existingMessages = await AIPersistService.loadUIMessages(chatId)
-      if (existingMessages.length > 0) {
-        return existingMessages
-      }
-
       const sessionResponse = await followApi.aiChatSessions.get({ chatId })
       const session = sessionResponse.data
 
@@ -129,7 +127,7 @@ class AIChatSessionServiceStatic {
         return AIPersistService.loadUIMessages(chatId)
       }
 
-      await this.fetchAndPersistMessages(session, { force: true })
+      await this.fetchAndPersistMessages(session)
       return AIPersistService.loadUIMessages(chatId)
     } catch (error) {
       console.error("syncSessionMessages: failed", error)

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

@@ -3,9 +3,8 @@ import { useEventCallback } from "usehooks-ts"
 
 import { AIChatSessionService } from "~/modules/ai-chat-session/service"
 
-import { AIPersistService } from "../services"
 import { useChatActions } from "../store/hooks"
-import type { BizUIMessage, BizUIMetadata } from "../store/types"
+import type { BizUIMessage } from "../store/types"
 
 export const useLoadMessages = (
   chatId: string,
@@ -24,35 +23,13 @@ export const useLoadMessages = (
     let mounted = true
     setIsLoading(true)
     setIsSyncingRemote(false)
-    AIPersistService.loadMessages(chatId)
+    AIChatSessionService.syncSessionMessages(chatId)
       .then(async (messages) => {
-        if (mounted) {
-          const messagesToSet: BizUIMessage[] = messages.map((message) => ({
-            id: message.id,
-            parts: message.messageParts as any[],
-            role: message.role,
-            metadata: message.metadata as BizUIMetadata,
-            createdAt: message.createdAt,
-          }))
-          const existingMessages = chatActions.getMessages()
-
-          if (messagesToSet.length === 0) {
-            if (existingMessages.length > 0) {
-              onLoadEventCallback(existingMessages)
-              return existingMessages
-            }
-
-            setIsSyncingRemote(true)
-            const syncedMessages = await AIChatSessionService.syncSessionMessages(chatId)
-            chatActions.setMessages(syncedMessages)
-            onLoadEventCallback(syncedMessages)
-            return syncedMessages
-          }
-
-          chatActions.setMessages(messagesToSet)
-          onLoadEventCallback(messagesToSet)
-          return messagesToSet
+        if (!mounted) {
+          return []
         }
+        chatActions.setMessages(messages)
+        onLoadEventCallback(messages)
         return messages
       })
       .catch((error) => {

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

@@ -129,6 +129,10 @@ class AIPersistServiceStatic {
       const cleanParts = [] as typeof message.parts
 
       for (const part of message.parts) {
+        // Skip streaming messages
+        if ("state" in part && part.state === "streaming") {
+          return acc
+        }
         if (isDataBlockPart(part)) {
           const nextPart = structuredClone(part)
           for (const block of nextPart.data) {
@@ -168,8 +172,24 @@ class AIPersistServiceStatic {
           metadata: sql`excluded.metadata`,
           finishedAt: sql`excluded.finished_at`,
           status: sql`excluded.status`,
+          createdAt: sql`excluded.created_at`,
         },
       })
+
+    const date = results.reduce<Date | null>((latest, msg) => {
+      const date = msg.createdAt ? new Date(msg.createdAt) : null
+      if (date === null) {
+        return latest
+      }
+      if (!latest || date > latest) {
+        return date
+      }
+      return latest
+    }, null)
+    if (date) {
+      // Update session time after successfully saving messages
+      await AIPersistService.updateSessionTime(chatId, date)
+    }
   }
 
   /**
@@ -416,11 +436,11 @@ class AIPersistServiceStatic {
       .where(eq(aiChatTable.chatId, chatId))
   }
 
-  async updateSessionTime(chatId: string) {
+  async updateSessionTime(chatId: string, date: Date = new Date()) {
     await db
       .update(aiChatTable)
       .set({
-        updatedAt: new Date(Date.now()),
+        updatedAt: date,
       })
       .where(eq(aiChatTable.chatId, chatId))
   }

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

@@ -147,14 +147,12 @@ export class ZustandChatState implements ChatState<BizUIMessage> {
         await AIPersistService.ensureSession(this.chatId)
         // Save messages using incremental updates
         await AIPersistService.replaceAllMessages(this.chatId, this.#messages)
-        // Update session time after successfully saving messages
-        await AIPersistService.updateSessionTime(this.chatId)
       } catch (error) {
         console.error("Failed to persist messages:", error)
       }
     },
     100,
-    { leading: true, trailing: true },
+    { leading: false, trailing: true },
   )
 
   #fillMessageCreatedAt(message: SendingUIMessage | BizUIMessage): BizUIMessage {