From 2ad4169c54f254beebaa1a43d4493005aa62a394 Mon Sep 17 00:00:00 2001 From: Whitewater Date: Fri, 31 Oct 2025 20:38:26 +0800 Subject: [PATCH] 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 --- .../src/modules/ai-chat-session/service.ts | 14 ++++---- .../modules/ai-chat/hooks/useLoadMessages.ts | 35 ++++--------------- .../src/modules/ai-chat/services/index.ts | 24 +++++++++++-- .../ai-chat/store/chat-core/chat-state.ts | 4 +-- 4 files changed, 35 insertions(+), 42 deletions(-) diff --git a/apps/desktop/layer/renderer/src/modules/ai-chat-session/service.ts b/apps/desktop/layer/renderer/src/modules/ai-chat-session/service.ts index 9d1d69111..3ac09a875 100644 --- a/apps/desktop/layer/renderer/src/modules/ai-chat-session/service.ts +++ b/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) diff --git a/apps/desktop/layer/renderer/src/modules/ai-chat/hooks/useLoadMessages.ts b/apps/desktop/layer/renderer/src/modules/ai-chat/hooks/useLoadMessages.ts index 8b30c5c06..9d6af2bca 100644 --- a/apps/desktop/layer/renderer/src/modules/ai-chat/hooks/useLoadMessages.ts +++ b/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) => { diff --git a/apps/desktop/layer/renderer/src/modules/ai-chat/services/index.ts b/apps/desktop/layer/renderer/src/modules/ai-chat/services/index.ts index ed909bd0c..12d2c385b 100644 --- a/apps/desktop/layer/renderer/src/modules/ai-chat/services/index.ts +++ b/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((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)) } diff --git a/apps/desktop/layer/renderer/src/modules/ai-chat/store/chat-core/chat-state.ts b/apps/desktop/layer/renderer/src/modules/ai-chat/store/chat-core/chat-state.ts index 205340a90..6c05ff329 100644 --- a/apps/desktop/layer/renderer/src/modules/ai-chat/store/chat-core/chat-state.ts +++ b/apps/desktop/layer/renderer/src/modules/ai-chat/store/chat-core/chat-state.ts @@ -147,14 +147,12 @@ export class ZustandChatState implements ChatState { 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 {