diff --git a/desktop/src/stores/chatStore.test.ts b/desktop/src/stores/chatStore.test.ts index 75058a32..2474343a 100644 --- a/desktop/src/stores/chatStore.test.ts +++ b/desktop/src/stores/chatStore.test.ts @@ -695,6 +695,117 @@ describe('chatStore history mapping', () => { }) }) + it('does not duplicate a hydrated assistant reply when live output replays after reconnect', () => { + useChatStore.setState({ + sessions: { + [TEST_SESSION_ID]: makeSession({ + messages: [ + { + id: 'live-user', + type: 'user_text', + content: 'live prompt', + transcriptMessageId: 'transcript-user-1', + timestamp: 1, + }, + { + id: 'live-assistant', + type: 'assistant_text', + content: 'live answer', + transcriptMessageId: 'transcript-assistant-1', + timestamp: 2, + }, + ], + streamingText: 'live answer', + chatState: 'streaming', + }), + }, + }) + + useChatStore.getState().handleServerMessage(TEST_SESSION_ID, { + type: 'message_complete', + usage: { input_tokens: 1, output_tokens: 2 }, + }) + + expect(useChatStore.getState().sessions[TEST_SESSION_ID]?.messages).toMatchObject([ + { + id: 'live-user', + type: 'user_text', + transcriptMessageId: 'transcript-user-1', + }, + { + id: 'live-assistant', + type: 'assistant_text', + content: 'live answer', + transcriptMessageId: 'transcript-assistant-1', + }, + ]) + expect(notifyDesktopMock).not.toHaveBeenCalled() + }) + + it('collapses duplicate assistant replies after transcript id hydration', async () => { + vi.mocked(sessionsApi.getMessages).mockResolvedValueOnce({ + messages: [ + { + id: 'transcript-user-1', + type: 'user', + timestamp: '2026-04-06T00:00:00.000Z', + content: 'live prompt', + }, + { + id: 'transcript-assistant-1', + type: 'assistant', + timestamp: '2026-04-06T00:00:01.000Z', + content: 'live answer', + }, + ], + }) + + useChatStore.setState({ + sessions: { + [TEST_SESSION_ID]: makeSession({ + messages: [ + { + id: 'live-user', + type: 'user_text', + content: 'live prompt', + transcriptMessageId: 'transcript-user-1', + timestamp: 1, + }, + { + id: 'live-assistant', + type: 'assistant_text', + content: 'live answer', + transcriptMessageId: 'transcript-assistant-1', + timestamp: 2, + }, + { + id: 'replayed-assistant', + type: 'assistant_text', + content: 'live answer', + timestamp: 3, + }, + ], + }), + }, + }) + + await useChatStore.getState().loadHistory(TEST_SESSION_ID) + + expect(useChatStore.getState().sessions[TEST_SESSION_ID]?.messages).toMatchObject([ + { + id: 'live-user', + type: 'user_text', + transcriptMessageId: 'transcript-user-1', + }, + { + id: 'live-assistant', + type: 'assistant_text', + content: 'live answer', + transcriptMessageId: 'transcript-assistant-1', + }, + ]) + }) + it('retries transcript id hydration after the assistant message is persisted', async () => { vi.useFakeTimers() vi.mocked(sessionsApi.getMessages) diff --git a/desktop/src/stores/chatStore.ts b/desktop/src/stores/chatStore.ts index 4ac7bcb7..7a384c82 100644 --- a/desktop/src/stores/chatStore.ts +++ b/desktop/src/stores/chatStore.ts @@ -352,9 +352,21 @@ function appendAssistantTextMessage( model?: string, transcriptMessageId?: string, ): UIMessage[] { - if (!content.trim()) return messages + const trimmedContent = content.trim() + if (!trimmedContent) return messages const last = messages[messages.length - 1] + // Wake/reconnect replay can resend persisted assistant text without a + // transcript id. Ignore chunks that are already present in the hydrated tail. + if ( + last?.type === 'assistant_text' && + last.transcriptMessageId && + !transcriptMessageId && + last.content.trim().includes(trimmedContent) + ) { + return messages + } + const canMergeIntoLast = last?.type === 'assistant_text' && ( @@ -581,12 +593,38 @@ function mergeRestoredTranscriptMessageIds( return changed ? merged : messages } +function dropDuplicateTranscriptTextMessages(messages: UIMessage[]): UIMessage[] { + const seen = new Set() + const deduped: UIMessage[] = [] + let changed = false + + for (const message of messages) { + if ( + (message.type === 'user_text' || message.type === 'assistant_text') && + message.transcriptMessageId + ) { + const key = `${message.type}:${message.transcriptMessageId}:${message.content.trim()}` + if (seen.has(key)) { + changed = true + continue + } + seen.add(key) + } + + deduped.push(message) + } + + return changed ? deduped : messages +} + function mergeRestoredHistoryIntoLiveMessages( messages: UIMessage[], restoredMessages: UIMessage[], ): UIMessage[] { return mergeRestoredTerminalGoalEvents( - mergeRestoredTranscriptMessageIds(messages, restoredMessages), + dropDuplicateTranscriptTextMessages( + mergeRestoredTranscriptMessageIds(messages, restoredMessages), + ), restoredMessages, ) } @@ -1493,7 +1531,8 @@ export const useChatStore = create((set, get) => ({ apiRetry: null, })) useTabStore.getState().updateTabStatus(sessionId, 'idle') - const notification = wasAgentRunning + const appendedCompletionMessage = completionMessages !== session.messages + const notification = wasAgentRunning && appendedCompletionMessage ? buildAgentCompletionNotification(sessionId, completionMessages, text) : null if (notification) {