Skip to content

Commit 5aed1af

Browse files
InneiInnei
committed
fix(ai): fetch remote messages if not exist
Signed-off-by: Innei <[email protected]>
1 parent a939842 commit 5aed1af

File tree

19 files changed

+1147
-34
lines changed

19 files changed

+1147
-34
lines changed

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

Lines changed: 13 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -49,6 +49,7 @@ class AIChatSessionServiceStatic {
4949
createdAt: new Date(session.createdAt),
5050
// Use createdAt for updatedAt as we are syncing session instead of messages
5151
updatedAt: new Date(session.createdAt),
52+
isLocal: false,
5253
})
5354
})
5455
await this.loadSessionsFromDb()
@@ -84,11 +85,15 @@ class AIChatSessionServiceStatic {
8485
force?: boolean
8586
},
8687
): Promise<void> {
87-
const dbSession = await AIPersistService.getChatSession(session.chatId)
88+
const [dbSession, hasPersistedMessages] = await Promise.all([
89+
AIPersistService.getChatSession(session.chatId),
90+
AIPersistService.hasPersistedMessages(session.chatId),
91+
])
92+
8893
const lastUpdatedAt = dbSession ? dbSession.updatedAt : new Date(0)
8994
const hasUpToDateSession = lastUpdatedAt >= new Date(session.updatedAt)
9095

91-
if (!options?.force && hasUpToDateSession) {
96+
if (!options?.force && hasUpToDateSession && hasPersistedMessages) {
9297
// If local session is already up-to-date, skip fetching messages
9398
return
9499
}
@@ -104,6 +109,7 @@ class AIChatSessionServiceStatic {
104109
// Use createdAt for updatedAt
105110
// Because we are fetching session data instead of messages
106111
updatedAt: new Date(session.createdAt),
112+
isLocal: false,
107113
})
108114
await AIPersistService.upsertMessages(session.chatId, normalized)
109115

@@ -120,6 +126,11 @@ class AIChatSessionServiceStatic {
120126

121127
async syncSessionMessages(chatId: string) {
122128
try {
129+
const sessionRecord = await AIPersistService.getChatSession(chatId)
130+
if (sessionRecord?.isLocal) {
131+
return AIPersistService.loadUIMessages(chatId)
132+
}
133+
123134
const sessionResponse = await followApi.aiChatSessions.get({ chatId })
124135
const session = sessionResponse.data
125136

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

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -82,6 +82,8 @@ export const aiChatSessionStoreActions = {
8282
title: session.title || t("ai:common.new_chat"),
8383
createdAt: new Date(session.createdAt),
8484
updatedAt: new Date(session.updatedAt),
85+
isLocal: false,
86+
syncStatus: "synced" as const,
8587
})),
8688
)
8789
} catch (error) {

‎apps/desktop/layer/renderer/src/modules/ai-chat/components/message/AIMessageParts.tsx‎

Lines changed: 11 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -25,20 +25,21 @@ interface AIMessagePartsProps {
2525

2626
export const AIMessageParts: React.FC<AIMessagePartsProps> = React.memo(
2727
({ message, isLastMessage }) => {
28-
const [shouldStreamingAnimation, setShouldStreamingAnimation] = React.useState(false)
28+
// const [shouldStreamingAnimation, setShouldStreamingAnimation] = React.useState(false)
2929
const chatStatus = useChatStatus()
3030

31-
React.useEffect(() => {
32-
// Delay 2s to set shouldStreamingAnimation
33-
const timerId = setTimeout(() => {
34-
setShouldStreamingAnimation(true)
35-
}, 2000)
36-
return () => clearTimeout(timerId)
37-
}, [])
31+
// React.useEffect(() => {
32+
// I forgot why do this
33+
// // Delay 2s to set shouldStreamingAnimation
34+
// const timerId = setTimeout(() => {
35+
// setShouldStreamingAnimation(true)
36+
// }, 2000)
37+
// return () => clearTimeout(timerId)
38+
// }, [])
3839

3940
const shouldMessageAnimation = React.useMemo(() => {
40-
return chatStatus === "streaming" && shouldStreamingAnimation && isLastMessage
41-
}, [chatStatus, isLastMessage, shouldStreamingAnimation])
41+
return chatStatus === "streaming" && isLastMessage // && shouldStreamingAnimation
42+
}, [chatStatus, isLastMessage])
4243

4344
const chainThoughtParts = React.useMemo(() => {
4445
const parts = [] as (ChainReasoningPart[] | TextUIPart | ToolUIPart<BizUITools>)[]

‎apps/desktop/layer/renderer/src/modules/ai-chat/components/message/UserChatMessage.tsx‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -31,6 +31,7 @@ export const UserChatMessage: React.FC<UserChatMessageProps> = React.memo(({ mes
3131
const setEditingMessageId = useSetEditingMessageId()
3232

3333
const chatStatus = useChatStatus()
34+
3435
const isStreaming = chatStatus === "submitted" || chatStatus === "streaming"
3536
const isEditing = editingMessageId === messageId
3637

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

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@ import { useEventCallback } from "usehooks-ts"
33

44
import { AIChatSessionService } from "~/modules/ai-chat-session/service"
55

6+
import { AIPersistService } from "../services"
67
import { useChatActions } from "../store/hooks"
78
import type { BizUIMessage } from "../store/types"
89

@@ -20,6 +21,12 @@ export const useLoadMessages = (
2021
})
2122

2223
useEffect(() => {
24+
if (chatActions.get().isLocal) {
25+
AIPersistService.loadUIMessages(chatId)
26+
setIsLoading(false)
27+
28+
return
29+
}
2330
let mounted = true
2431
setIsLoading(true)
2532
setIsSyncingRemote(false)

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

Lines changed: 47 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -38,6 +38,17 @@ class AIPersistServiceStatic {
3838
})
3939
}
4040

41+
async hasPersistedMessages(chatId: string): Promise<boolean> {
42+
const existingMessage = await db.query.aiChatMessagesTable.findFirst({
43+
where: eq(aiChatMessagesTable.chatId, chatId),
44+
columns: {
45+
id: true,
46+
},
47+
})
48+
49+
return Boolean(existingMessage?.id === chatId)
50+
}
51+
4152
/**
4253
* Convert enhanced database message to BizUIMessage format for compatibility
4354
*/
@@ -74,6 +85,8 @@ class AIPersistServiceStatic {
7485
title?: string
7586
createdAt: Date
7687
updatedAt: Date
88+
isLocal: boolean
89+
syncStatus: "local" | "synced"
7790
} | null
7891
messages: BizUIMessage[]
7992
}> {
@@ -97,9 +110,14 @@ class AIPersistServiceStatic {
97110
await this.updateSessionTitle(sessionRaw.chatId, resolvedTitle)
98111
}
99112

113+
const isLocal = Boolean(sessionRaw.isLocal)
114+
const syncStatus: "local" | "synced" = isLocal ? "local" : "synced"
115+
100116
const session = {
101117
...sessionRaw,
102118
title: resolvedTitle ?? undefined,
119+
isLocal,
120+
syncStatus,
103121
}
104122

105123
return { session, messages }
@@ -282,10 +300,11 @@ class AIPersistServiceStatic {
282300
*/
283301
async ensureSession(
284302
chatId: string,
285-
options: { title?: string; createdAt?: Date; updatedAt?: Date } = {},
303+
options: { title?: string; createdAt?: Date; updatedAt?: Date; isLocal?: boolean } = {},
286304
): Promise<void> {
287305
const cachedExists = this.getSessionExistsFromCache(chatId)
288-
const shouldCheckDb = cachedExists !== true || options.title
306+
const shouldCheckDb =
307+
cachedExists !== true || options.title !== undefined || typeof options.isLocal === "boolean"
289308

290309
if (!shouldCheckDb) {
291310
return
@@ -311,6 +330,11 @@ class AIPersistServiceStatic {
311330
}
312331
}
313332

333+
if (typeof options.isLocal === "boolean" && existing.isLocal !== options.isLocal) {
334+
updates.isLocal = options.isLocal
335+
shouldUpdate = true
336+
}
337+
314338
if (shouldUpdate) {
315339
updates.updatedAt = new Date()
316340
await db.update(aiChatTable).set(updates).where(eq(aiChatTable.chatId, chatId))
@@ -325,14 +349,15 @@ class AIPersistServiceStatic {
325349

326350
async createSession(
327351
chatId: string,
328-
options: { title?: string; createdAt?: Date; updatedAt?: Date } = {},
352+
options: { title?: string; createdAt?: Date; updatedAt?: Date; isLocal?: boolean } = {},
329353
) {
330354
const now = new Date()
331355
await db.insert(aiChatTable).values({
332356
chatId,
333357
title: this.resolveSessionTitle(chatId, options.title, { createdAt: now, updatedAt: now }),
334358
createdAt: options.createdAt ?? now,
335359
updatedAt: options.updatedAt ?? now,
360+
isLocal: options.isLocal ?? true,
336361
})
337362
// Mark session as existing in cache
338363
this.markSessionExists(chatId, true)
@@ -369,6 +394,7 @@ class AIPersistServiceStatic {
369394
title: true,
370395
createdAt: true,
371396
updatedAt: true,
397+
isLocal: true,
372398
},
373399
})
374400
return result?.chatId ? result : null
@@ -381,6 +407,7 @@ class AIPersistServiceStatic {
381407
title: true,
382408
createdAt: true,
383409
updatedAt: true,
410+
isLocal: true,
384411
},
385412
orderBy: (t, { desc }) => desc(t.updatedAt),
386413
limit,
@@ -408,12 +435,19 @@ class AIPersistServiceStatic {
408435
}),
409436
)
410437

411-
return normalizedChats.map((chat) => ({
412-
chatId: chat.chatId,
413-
title: chat.title,
414-
createdAt: chat.createdAt,
415-
updatedAt: chat.updatedAt,
416-
}))
438+
return normalizedChats.map((chat) => {
439+
const isLocal = Boolean(chat.isLocal)
440+
const syncStatus: "local" | "synced" = isLocal ? "local" : "synced"
441+
442+
return {
443+
chatId: chat.chatId,
444+
title: chat.title,
445+
createdAt: chat.createdAt,
446+
updatedAt: chat.updatedAt,
447+
isLocal,
448+
syncStatus,
449+
}
450+
})
417451
}
418452

419453
async deleteSession(chatId: string) {
@@ -445,6 +479,10 @@ class AIPersistServiceStatic {
445479
.where(eq(aiChatTable.chatId, chatId))
446480
}
447481

482+
async markSessionSynced(chatId: string) {
483+
await this.ensureSession(chatId, { isLocal: false })
484+
}
485+
448486
async cleanupEmptySessions() {
449487
const emptySessions = await db.values<[string]>(
450488
sql`

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

Lines changed: 43 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -48,6 +48,42 @@ export class ChatSliceActions {
4848
return this.params[1]
4949
}
5050

51+
private computeSyncStatus(isLocal: boolean): "local" | "synced" {
52+
return isLocal ? "local" : "synced"
53+
}
54+
55+
private setSyncState(isLocal: boolean) {
56+
this.set((state) => {
57+
const nextStatus = this.computeSyncStatus(isLocal)
58+
if (state.isLocal === isLocal && state.syncStatus === nextStatus) {
59+
return state
60+
}
61+
return {
62+
isLocal,
63+
syncStatus: nextStatus,
64+
}
65+
})
66+
}
67+
68+
async markSessionSynced() {
69+
const currentChatId = this.get().chatId
70+
if (!currentChatId) {
71+
return
72+
}
73+
74+
if (!this.get().isLocal) {
75+
return
76+
}
77+
78+
this.setSyncState(false)
79+
80+
try {
81+
await AIPersistService.markSessionSynced(currentChatId)
82+
} catch (error) {
83+
console.error("Failed to mark chat session as synced:", error)
84+
}
85+
}
86+
5187
// Direct message management methods (delegating to chat instance state)
5288
setMessages = (
5389
messagesParam: BizUIMessage[] | ((messages: BizUIMessage[]) => BizUIMessage[]),
@@ -176,8 +212,8 @@ export class ChatSliceActions {
176212
},
177213
options,
178214
)
179-
const response = await this.chatInstance.sendMessage(messageObj, finalOptions)
180-
return response
215+
216+
return await this.chatInstance.sendMessage(messageObj, finalOptions)
181217
} catch (error) {
182218
this.setError(error as Error)
183219
throw error
@@ -193,8 +229,7 @@ export class ChatSliceActions {
193229
},
194230
options,
195231
)
196-
const response = await this.chatInstance.regenerate({ messageId, ...finalOptions })
197-
return response
232+
return await this.chatInstance.regenerate({ messageId, ...finalOptions })
198233
} catch (error) {
199234
this.setError(error as Error)
200235
throw error
@@ -252,6 +287,8 @@ export class ChatSliceActions {
252287
isStreaming: false,
253288
currentTitle: undefined,
254289
chatInstance: newChatInstance,
290+
isLocal: true,
291+
syncStatus: "local",
255292
}))
256293

257294
// Update the reference
@@ -293,6 +330,8 @@ export class ChatSliceActions {
293330
isStreaming: false,
294331
currentTitle: chatSession?.title || undefined,
295332
chatInstance: newChatInstance,
333+
isLocal: chatSession ? chatSession.isLocal : true,
334+
syncStatus: chatSession ? chatSession.syncStatus : "local",
296335
}))
297336

298337
newChatInstance.resumeStream()

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

Lines changed: 12 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -61,11 +61,18 @@ export class ZustandChatState implements ChatState<BizUIMessage> {
6161
})
6262

6363
this.#eventEmitter.on("status", ({ status }) => {
64-
this.updateZustandState((state) => ({
65-
...state,
66-
status,
67-
isStreaming: status === "streaming",
68-
}))
64+
this.updateZustandState((state) => {
65+
const isStreaming = status === "streaming"
66+
if (isStreaming) {
67+
void state.chatActions.markSessionSynced()
68+
}
69+
70+
return {
71+
...state,
72+
status,
73+
isStreaming,
74+
}
75+
})
6976
})
7077

7178
this.#eventEmitter.on("error", ({ error }) => {

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

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,8 @@ export interface ChatSlice {
1111
status: ChatStatus
1212
error: Error | undefined
1313
isStreaming: boolean
14+
isLocal: boolean
15+
syncStatus: "local" | "synced"
1416

1517
// UI state
1618
currentTitle: string | undefined

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

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -65,6 +65,21 @@ export const useHasMessages = () => {
6565
return store((state) => state.messages.length > 0)
6666
}
6767

68+
export const useIsLocalChat = () => {
69+
const store = useAIChatStore()
70+
return store((state) => state.isLocal)
71+
}
72+
73+
export const useSyncStatus = () => {
74+
const store = useAIChatStore()
75+
return store((state) => state.syncStatus)
76+
}
77+
78+
export const useSyncStateActions = () => {
79+
const store = useAIChatStore()
80+
return store((state) => state.chatActions)
81+
}
82+
6883
export const useChatBlockActions = () => useAIChatStore()((state) => state.blockActions)
6984
/**
7085
* Hook to get the chat status

0 commit comments

Comments
 (0)