diff --git a/desktop/src/stores/chatStore.test.ts b/desktop/src/stores/chatStore.test.ts index e8d6ea68..aab57fdc 100644 --- a/desktop/src/stores/chatStore.test.ts +++ b/desktop/src/stores/chatStore.test.ts @@ -157,6 +157,30 @@ describe('chatStore history mapping', () => { expect(mapped[3]).toMatchObject({ parentToolUseId: 'agent-1' }) }) + it('merges consecutive assistant text blocks when restoring transcript history', () => { + const messages: MessageEntry[] = [ + { + id: 'assistant-merge-1', + type: 'assistant', + timestamp: '2026-04-06T00:00:00.000Z', + model: 'opus', + content: [ + { type: 'text', text: '第一段：Windows 下的桌面端输出。' }, + { type: 'text', text: '\r\n第二段：刷新后也不应该被拆开。' }, + ], + }, + ] + + const mapped = mapHistoryMessagesToUiMessages(messages) + + expect(mapped).toMatchObject([ + { + type: 'assistant_text', + content: '第一段：Windows 下的桌面端输出。\r\n第二段：刷新后也不应该被拆开。', + }, + ]) + }) + it('surfaces teammate prompt content when mapping member transcript history', () => { const messages: MessageEntry[] = [ { @@ -532,6 +556,68 @@ describe('chatStore history mapping', () => { ).toBe('permission_pending') }) + it('keeps delayed text blocks from one streamed assistant turn in a single message', () => { + vi.useFakeTimers() + + useChatStore.setState({ + sessions: { + [TEST_SESSION_ID]: { + messages: [], + chatState: 'idle', + connectionState: 'connected', + streamingText: '', + streamingToolInput: '', + activeToolUseId: null, + activeToolName: null, + activeThinkingId: null, + pendingPermission: null, + pendingComputerUsePermission: null, + tokenUsage: { input_tokens: 0, output_tokens: 0 }, + elapsedSeconds: 0, + statusVerb: '', + slashCommands: [], + agentTaskNotifications: {}, + elapsedTimer: null, + }, + }, + }) + + useChatStore.getState().handleServerMessage(TEST_SESSION_ID, { + type: 'content_start', + blockType: 'text', + }) + useChatStore.getState().handleServerMessage(TEST_SESSION_ID, { + type: 'content_delta', + text: '第一段：先到达。', + }) + vi.advanceTimersByTime(60) + + useChatStore.getState().handleServerMessage(TEST_SESSION_ID, { + type: 'content_start', + blockType: 'text', + }) + useChatStore.getState().handleServerMessage(TEST_SESSION_ID, { + type: 'content_delta', + text: '\r\n第二段：稍后到达，但仍属于同一轮回复。', + }) + vi.advanceTimersByTime(60) + + 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([ + { + type: 'assistant_text', + content: '第一段：先到达。\r\n第二段：稍后到达，但仍属于同一轮回复。', + }, + ]) + + vi.runOnlyPendingTimers() + vi.useRealTimers() + }) + it('sends Computer Use approval payloads back over websocket', () => { useChatStore.setState({ sessions: { diff --git a/desktop/src/stores/chatStore.ts b/desktop/src/stores/chatStore.ts index 94d87f38..5c78f7d7 100644 --- a/desktop/src/stores/chatStore.ts +++ b/desktop/src/stores/chatStore.ts @@ -128,6 +128,46 @@ const nextId = () => `msg-${++msgCounter}-${Date.now()}` let pendingDelta = '' let flushTimer: ReturnType | null = null +function consumePendingDelta(): string { + if (flushTimer) { + clearTimeout(flushTimer) + flushTimer = null + } + const text = pendingDelta + pendingDelta = '' + return text +} + +function appendAssistantTextMessage( + messages: UIMessage[], + content: string, + timestamp: number, + model?: string, +): UIMessage[] { + if (!content.trim()) return messages + + const last = messages[messages.length - 1] + if (last?.type === 'assistant_text') { + const merged: UIMessage = { + ...last, + content: last.content + content, + ...(model ?? last.model ? { model: model ?? last.model } : {}), + } + return [...messages.slice(0, -1), merged] + } + + return [ + ...messages, + { + id: nextId(), + type: 'assistant_text', + content, + timestamp, + ...(model ? { model } : {}), + }, + ] +} + /** Helper: immutably update a specific session within the sessions record */ function updateSessionIn( sessions: Record, @@ -200,8 +240,7 @@ export const useChatStore = create((set, get) => ({ if (session?.elapsedTimer) clearInterval(session.elapsedTimer) if (flushTimer) { clearTimeout(flushTimer); flushTimer = null } if (pendingDelta) { - const text = pendingDelta - pendingDelta = '' + const text = consumePendingDelta() set((s) => ({ sessions: updateSessionIn(s.sessions, sessionId, (sess) => ({ streamingText: sess.streamingText + text })) })) } wsManager.disconnect(sessionId) @@ -241,20 +280,11 @@ export const useChatStore = create((set, get) => ({ clearTimeout(flushTimer) flushTimer = null } - const bufferedDelta = pendingDelta - pendingDelta = '' - const pendingAssistantText = `${session.streamingText}${bufferedDelta}`.trim() + const bufferedDelta = consumePendingDelta() + const pendingAssistantText = `${session.streamingText}${bufferedDelta}` - const newMessages = pendingAssistantText - ? [ - ...session.messages, - { - id: nextId(), - type: 'assistant_text' as const, - content: pendingAssistantText, - timestamp: Date.now(), - }, - ] + const newMessages = pendingAssistantText.trim() + ? appendAssistantTextMessage(session.messages, pendingAssistantText, Date.now()) : [...session.messages] if (!isMemberSession && allTasksDone) { newMessages.push({ @@ -357,8 +387,7 @@ export const useChatStore = create((set, get) => ({ wsManager.send(sessionId, { type: 'stop_generation' }) if (flushTimer) { clearTimeout(flushTimer); flushTimer = null } if (pendingDelta) { - const text = pendingDelta - pendingDelta = '' + const text = consumePendingDelta() set((s) => ({ sessions: updateSessionIn(s.sessions, sessionId, (sess) => ({ streamingText: sess.streamingText + text })) })) } set((s) => { @@ -481,17 +510,17 @@ export const useChatStore = create((set, get) => ({ case 'status': update((session) => { - const pendingText = session.streamingText.trim() - const shouldFlush = pendingText && session.chatState === 'streaming' && msg.state !== 'streaming' + const pendingText = `${session.streamingText}${consumePendingDelta()}` + const shouldFlush = pendingText.trim() && session.chatState === 'streaming' && msg.state !== 'streaming' return { chatState: msg.state, ...(msg.verb && msg.verb !== 'Thinking' ? { statusVerb: msg.verb } : {}), ...(msg.tokens ? { tokenUsage: { ...session.tokenUsage, output_tokens: msg.tokens } } : {}), ...(msg.state === 'idle' ? { activeThinkingId: null, statusVerb: '' } : {}), ...(shouldFlush ? { - messages: [...session.messages, { id: nextId(), type: 'assistant_text' as const, content: pendingText, timestamp: Date.now() }], + messages: appendAssistantTextMessage(session.messages, pendingText, Date.now()), streamingText: '', - } : {}), + } : pendingText !== session.streamingText ? { streamingText: pendingText } : {}), } }) if (msg.state === 'idle') { @@ -508,15 +537,19 @@ export const useChatStore = create((set, get) => ({ case 'content_start': { const session = get().sessions[sessionId] if (!session) break - const pendingText = session.streamingText.trim() - if (pendingText) { + const pendingText = `${session.streamingText}${consumePendingDelta()}` + if (msg.blockType !== 'text' && pendingText.trim()) { update((s) => ({ - messages: [...s.messages, { id: nextId(), type: 'assistant_text' as const, content: pendingText, timestamp: Date.now() }], + messages: appendAssistantTextMessage(s.messages, pendingText, Date.now()), streamingText: '', })) } if (msg.blockType === 'text') { - update(() => ({ streamingText: '', chatState: 'streaming', activeThinkingId: null })) + update((s) => ({ + ...(pendingText !== s.streamingText ? { streamingText: pendingText } : {}), + chatState: 'streaming', + activeThinkingId: null, + })) } else if (msg.blockType === 'tool_use') { update(() => ({ activeToolUseId: msg.toolUseId ?? null, @@ -546,9 +579,9 @@ export const useChatStore = create((set, get) => ({ case 'thinking': update((s) => { - const pendingText = s.streamingText.trim() - const base = pendingText - ? [...s.messages, { id: nextId(), type: 'assistant_text' as const, content: pendingText, timestamp: Date.now() }] + const pendingText = `${s.streamingText}${consumePendingDelta()}` + const base = pendingText.trim() + ? appendAssistantTextMessage(s.messages, pendingText, Date.now()) : s.messages const last = base[base.length - 1] if (last && last.type === 'thinking') { @@ -643,12 +676,14 @@ export const useChatStore = create((set, get) => ({ case 'message_complete': { const session = get().sessions[sessionId] if (!session) break - const text = session.streamingText - if (text) { + const text = `${session.streamingText}${consumePendingDelta()}` + if (text.trim()) { update((s) => ({ - messages: [...s.messages, { id: nextId(), type: 'assistant_text', content: text, timestamp: Date.now() }], + messages: appendAssistantTextMessage(s.messages, text, Date.now()), streamingText: '', })) + } else if (text !== session.streamingText) { + update(() => ({ streamingText: text })) } if (session.elapsedTimer) clearInterval(session.elapsedTimer) update(() => ({ @@ -664,12 +699,12 @@ export const useChatStore = create((set, get) => ({ case 'error': update((s) => { - const pendingText = s.streamingText.trim() - const newMessages = [...s.messages] - if (pendingText) { - newMessages.push({ id: nextId(), type: 'assistant_text' as const, content: pendingText, timestamp: Date.now() }) + const pendingText = `${s.streamingText}${consumePendingDelta()}` + let newMessages = s.messages + if (pendingText.trim()) { + newMessages = appendAssistantTextMessage(newMessages, pendingText, Date.now()) } - newMessages.push({ id: nextId(), type: 'error', message: msg.message, code: msg.code, timestamp: Date.now() }) + newMessages = [...newMessages, { id: nextId(), type: 'error', message: msg.message, code: msg.code, timestamp: Date.now() }] return { messages: newMessages, chatState: 'idle', @@ -791,6 +826,30 @@ function extractVisibleTeammateMessageContents(text: string): string[] { return contents } +function pushAssistantHistoryText( + messages: UIMessage[], + content: string, + timestamp: number, + model?: string, +): void { + if (!content.trim()) return + + const last = messages[messages.length - 1] + if (last?.type === 'assistant_text') { + last.content += content + if (model && !last.model) last.model = model + return + } + + messages.push({ + id: nextId(), + type: 'assistant_text', + content, + timestamp, + ...(model ? { model } : {}), + }) +} + type HistoryMappingOptions = { includeTeammateMessages?: boolean } @@ -900,7 +959,7 @@ export function mapHistoryMessagesToUiMessages( if ((msg.type === 'assistant' || msg.type === 'tool_use') && Array.isArray(msg.content)) { for (const block of msg.content as AssistantHistoryBlock[]) { if (block.type === 'thinking' && block.thinking) uiMessages.push({ id: nextId(), type: 'thinking', content: block.thinking, timestamp }) - else if (block.type === 'text' && block.text) uiMessages.push({ id: nextId(), type: 'assistant_text', content: block.text, timestamp, model: msg.model }) + else if (block.type === 'text' && block.text) pushAssistantHistoryText(uiMessages, block.text, timestamp, msg.model) else if (block.type === 'tool_use') uiMessages.push({ id: nextId(), type: 'tool_use', toolName: block.name ?? 'unknown', toolUseId: block.id ?? '', input: block.input, timestamp, parentToolUseId: msg.parentToolUseId }) } continue