fix(stage-ui): use stable chat history keys in the stage UI (#1312)

---------

Co-authored-by: Stable Genius <259448942+stablegenius49@users.noreply.github.com>
Co-authored-by: Garfield Lee <Garfield550@users.noreply.github.com>
Co-authored-by-agent: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>

packages/stage-ui/src/components/scenarios/chat/history.vue

@@ -8,6 +8,8 @@ import ChatAssistantItem from './assistant-item.vue'
 import ChatErrorItem from './error-item.vue'
 import ChatUserItem from './user-item.vue'
 
+import { getChatHistoryItemKey } from './message-key'
+
 const props = withDefaults(defineProps<{
   messages: ChatHistoryItem[]
   streamingMessage?: ChatAssistantMessage & { createdAt?: number }
@@ -73,7 +75,7 @@ const renderMessages = computed<ChatHistoryItem[]>(() => {
 
 <template>
   <div ref="chatHistoryRef" v-auto-animate flex="~ col" relative h-full w-full overflow-y-auto rounded-xl px="<sm:2" py="<sm:2" :class="variant === 'mobile' ? 'gap-1' : 'gap-2'">
-    <template v-for="(message, index) in renderMessages" :key="message?.createdAt ?? index">
+    <template v-for="(message, index) in renderMessages" :key="getChatHistoryItemKey(message, index)">
       <div v-if="message.role === 'error'">
         <ChatErrorItem
           :message="message"

packages/stage-ui/src/components/scenarios/chat/message-key.test.ts

@@ -0,0 +1,42 @@
+import type { ChatHistoryItem } from '../../../types/chat'
+
+import assert from 'node:assert/strict'
+
+import { describe, it } from 'vitest'
+
+import { getChatHistoryItemKey } from './message-key'
+
+describe('getChatHistoryItemKey', () => {
+  it('prefers stable message ids when available', () => {
+    const createdAt = 1700000000000
+
+    const userMessage: ChatHistoryItem = { role: 'user', content: 'hi', createdAt, id: 'user-1' }
+    const assistantMessage: ChatHistoryItem = { role: 'assistant', content: 'hello', createdAt, id: 'assistant-1', slices: [], tool_results: [] }
+
+    assert.equal(getChatHistoryItemKey(userMessage, 0), 'user-1')
+    assert.equal(getChatHistoryItemKey(assistantMessage, 1), 'assistant-1')
+  })
+
+  it('falls back to a role + timestamp + index composite when ids are missing', () => {
+    const createdAt = 1700000000000
+
+    const userMessage: ChatHistoryItem = { role: 'user', content: 'hi', createdAt }
+    const assistantMessage: ChatHistoryItem = { role: 'assistant', content: 'hello', createdAt, slices: [], tool_results: [] }
+
+    assert.equal(getChatHistoryItemKey(userMessage, 0), 'user:1700000000000:0')
+    assert.equal(getChatHistoryItemKey(assistantMessage, 1), 'assistant:1700000000000:1')
+  })
+
+  it('falls back to index when message is missing', () => {
+    assert.equal(getChatHistoryItemKey(undefined, 0), 0)
+    assert.equal(getChatHistoryItemKey(undefined, 1), 1)
+  })
+
+  it('falls back to a role + index composite when ids and timestamps are missing', () => {
+    const userMessage: ChatHistoryItem = { role: 'user', content: 'hi' }
+    const assistantMessage: ChatHistoryItem = { role: 'assistant', content: 'hello', slices: [], tool_results: [] }
+
+    assert.equal(getChatHistoryItemKey(userMessage, 0), 'user:0')
+    assert.equal(getChatHistoryItemKey(assistantMessage, 1), 'assistant:1')
+  })
+})

packages/stage-ui/src/components/scenarios/chat/message-key.ts

@@ -0,0 +1,14 @@
+import type { ChatHistoryItem } from '../../../types/chat'
+
+export function getChatHistoryItemKey(message: ChatHistoryItem | undefined, index: number): string | number {
+  if (!message)
+    return index
+
+  if (message.id)
+    return message.id
+
+  if (message.createdAt != null)
+    return `${message.role}:${message.createdAt}:${index}`
+
+  return `${message.role}:${index}`
+}

packages/stage-ui/src/stores/chat/session-message-merge.test.ts

@@ -0,0 +1,50 @@
+import type { ChatHistoryItem } from '../../types/chat'
+
+import assert from 'node:assert/strict'
+
+import { describe, it } from 'vitest'
+
+import { mergeLoadedSessionMessages } from './session-message-merge'
+
+describe('mergeLoadedSessionMessages', () => {
+  it('keeps stored history when the in-memory session only has the placeholder system message', () => {
+    const storedMessages: ChatHistoryItem[] = [
+      { role: 'system', content: 'system', createdAt: 1, id: 'system-stored' },
+      { role: 'assistant', content: 'saved reply', createdAt: 2, id: 'assistant-1', slices: [], tool_results: [] },
+    ]
+    const currentMessages: ChatHistoryItem[] = [
+      { role: 'system', content: 'system', createdAt: 3, id: 'system-current' },
+    ]
+
+    assert.equal(mergeLoadedSessionMessages(storedMessages, currentMessages), storedMessages)
+  })
+
+  it('appends in-flight messages when IndexedDB finishes loading after a new send starts', () => {
+    const storedMessages: ChatHistoryItem[] = [
+      { role: 'system', content: 'system', createdAt: 1, id: 'system-stored' },
+      { role: 'assistant', content: 'older reply', createdAt: 2, id: 'assistant-1', slices: [], tool_results: [] },
+    ]
+    const currentMessages: ChatHistoryItem[] = [
+      { role: 'system', content: 'system', createdAt: 3, id: 'system-current' },
+      { role: 'user', content: 'latest prompt', createdAt: 4, id: 'user-2' },
+    ]
+
+    assert.deepEqual(mergeLoadedSessionMessages(storedMessages, currentMessages), [
+      ...storedMessages,
+      currentMessages[1],
+    ])
+  })
+
+  it('does not duplicate messages that are already present in storage', () => {
+    const storedMessages: ChatHistoryItem[] = [
+      { role: 'system', content: 'system', createdAt: 1, id: 'system-stored' },
+      { role: 'user', content: 'latest prompt', createdAt: 4 },
+    ]
+    const currentMessages: ChatHistoryItem[] = [
+      { role: 'system', content: 'system', createdAt: 3, id: 'system-current' },
+      { role: 'user', content: 'latest prompt', createdAt: 4 },
+    ]
+
+    assert.equal(mergeLoadedSessionMessages(storedMessages, currentMessages), storedMessages)
+  })
+})

packages/stage-ui/src/stores/chat/session-message-merge.ts

@@ -0,0 +1,59 @@
+import type { ChatHistoryItem } from '../../types/chat'
+
+function extractMessageContent(message: ChatHistoryItem) {
+  if (typeof message.content === 'string')
+    return message.content
+
+  if (Array.isArray(message.content)) {
+    return message.content.map((part) => {
+      if (typeof part === 'string')
+        return part
+      if (part && typeof part === 'object' && 'text' in part)
+        return String(part.text ?? '')
+      return ''
+    }).join('')
+  }
+
+  return ''
+}
+
+function getMessageFingerprint(message: ChatHistoryItem) {
+  return [
+    message.id ?? '',
+    message.role,
+    message.createdAt ?? '',
+    extractMessageContent(message),
+  ].join('\u001F')
+}
+
+export function mergeLoadedSessionMessages(storedMessages: ChatHistoryItem[], currentMessages: ChatHistoryItem[]) {
+  if (currentMessages.length === 0)
+    return storedMessages
+
+  const currentNonSystemMessages = currentMessages.filter((message, index) => index !== 0 || message.role !== 'system')
+  if (currentNonSystemMessages.length === 0)
+    return storedMessages
+
+  const seen = new Set(storedMessages.map(getMessageFingerprint))
+  const extraMessages = currentNonSystemMessages.filter((message) => {
+    const fingerprint = getMessageFingerprint(message)
+    if (seen.has(fingerprint))
+      return false
+    seen.add(fingerprint)
+    return true
+  })
+
+  if (extraMessages.length === 0)
+    return storedMessages
+
+  const systemMessage = storedMessages[0]?.role === 'system'
+    ? storedMessages[0]
+    : currentMessages[0]?.role === 'system'
+      ? currentMessages[0]
+      : undefined
+
+  if (storedMessages.length === 0 && systemMessage)
+    return [systemMessage, ...extraMessages]
+
+  return [...storedMessages, ...extraMessages]
+}

packages/stage-ui/src/stores/chat/session-store.ts

@@ -10,6 +10,7 @@ import { useLocalFirstRequest } from '../../composables/use-local-first'
 import { chatSessionsRepo } from '../../database/repos/chat-sessions.repo'
 import { useAuthStore } from '../auth'
 import { useAiriCardStore } from '../modules/airi-card'
+import { mergeLoadedSessionMessages } from './session-message-merge'
 
 export const useChatSessionStore = defineStore('chat-session', () => {
   const { userId, isAuthenticated } = storeToRefs(useAuthStore())
@@ -259,9 +260,15 @@ export const useChatSessionStore = defineStore('chat-session', () => {
     const loadPromise = (async () => {
       const stored = await chatSessionsRepo.getSession(sessionId)
       if (stored) {
+        const currentMessages = sessionMessages.value[sessionId] ?? []
+        const mergedMessages = mergeLoadedSessionMessages(stored.messages, currentMessages)
+
         sessionMetas.value[sessionId] = stored.meta
-        sessionMessages.value[sessionId] = stored.messages
+        sessionMessages.value[sessionId] = mergedMessages
         ensureGeneration(sessionId)
+
+        if (mergedMessages !== stored.messages)
+          await persistSession(sessionId)
       }
       loadedSessions.add(sessionId)
     })()