fix: merge split assistant text blocks in desktop chat

This commit is contained in:
Relakkes Yang 2026-04-22 14:23:14 +08:00
parent 4b513271ee
commit d390eb6463
2 changed files with 182 additions and 37 deletions

View File

@ -130,6 +130,30 @@ describe('chatStore history mapping', () => {
expect(mapped[3]).toMatchObject({ parentToolUseId: 'agent-1' }) 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', () => { it('surfaces teammate prompt content when mapping member transcript history', () => {
const messages: MessageEntry[] = [ const messages: MessageEntry[] = [
{ {
@ -389,6 +413,68 @@ describe('chatStore history mapping', () => {
).toBe('permission_pending') ).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', () => { it('sends Computer Use approval payloads back over websocket', () => {
useChatStore.setState({ useChatStore.setState({
sessions: { sessions: {

View File

@ -103,6 +103,46 @@ const nextId = () => `msg-${++msgCounter}-${Date.now()}`
let pendingDelta = '' let pendingDelta = ''
let flushTimer: ReturnType<typeof setTimeout> | null = null let flushTimer: ReturnType<typeof setTimeout> | 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 */ /** Helper: immutably update a specific session within the sessions record */
function updateSessionIn( function updateSessionIn(
sessions: Record<string, PerSessionState>, sessions: Record<string, PerSessionState>,
@ -163,8 +203,7 @@ export const useChatStore = create<ChatStore>((set, get) => ({
if (session?.elapsedTimer) clearInterval(session.elapsedTimer) if (session?.elapsedTimer) clearInterval(session.elapsedTimer)
if (flushTimer) { clearTimeout(flushTimer); flushTimer = null } if (flushTimer) { clearTimeout(flushTimer); flushTimer = null }
if (pendingDelta) { if (pendingDelta) {
const text = pendingDelta const text = consumePendingDelta()
pendingDelta = ''
set((s) => ({ sessions: updateSessionIn(s.sessions, sessionId, (sess) => ({ streamingText: sess.streamingText + text })) })) set((s) => ({ sessions: updateSessionIn(s.sessions, sessionId, (sess) => ({ streamingText: sess.streamingText + text })) }))
} }
wsManager.disconnect(sessionId) wsManager.disconnect(sessionId)
@ -196,20 +235,11 @@ export const useChatStore = create<ChatStore>((set, get) => ({
clearTimeout(flushTimer) clearTimeout(flushTimer)
flushTimer = null flushTimer = null
} }
const bufferedDelta = pendingDelta const bufferedDelta = consumePendingDelta()
pendingDelta = '' const pendingAssistantText = `${session.streamingText}${bufferedDelta}`
const pendingAssistantText = `${session.streamingText}${bufferedDelta}`.trim()
const newMessages = pendingAssistantText const newMessages = pendingAssistantText.trim()
? [ ? appendAssistantTextMessage(session.messages, pendingAssistantText, Date.now())
...session.messages,
{
id: nextId(),
type: 'assistant_text' as const,
content: pendingAssistantText,
timestamp: Date.now(),
},
]
: [...session.messages] : [...session.messages]
if (!isMemberSession && allTasksDone) { if (!isMemberSession && allTasksDone) {
newMessages.push({ newMessages.push({
@ -307,8 +337,7 @@ export const useChatStore = create<ChatStore>((set, get) => ({
wsManager.send(sessionId, { type: 'stop_generation' }) wsManager.send(sessionId, { type: 'stop_generation' })
if (flushTimer) { clearTimeout(flushTimer); flushTimer = null } if (flushTimer) { clearTimeout(flushTimer); flushTimer = null }
if (pendingDelta) { if (pendingDelta) {
const text = pendingDelta const text = consumePendingDelta()
pendingDelta = ''
set((s) => ({ sessions: updateSessionIn(s.sessions, sessionId, (sess) => ({ streamingText: sess.streamingText + text })) })) set((s) => ({ sessions: updateSessionIn(s.sessions, sessionId, (sess) => ({ streamingText: sess.streamingText + text })) }))
} }
set((s) => { set((s) => {
@ -371,17 +400,17 @@ export const useChatStore = create<ChatStore>((set, get) => ({
case 'status': case 'status':
update((session) => { update((session) => {
const pendingText = session.streamingText.trim() const pendingText = `${session.streamingText}${consumePendingDelta()}`
const shouldFlush = pendingText && session.chatState === 'streaming' && msg.state !== 'streaming' const shouldFlush = pendingText.trim() && session.chatState === 'streaming' && msg.state !== 'streaming'
return { return {
chatState: msg.state, chatState: msg.state,
...(msg.verb && msg.verb !== 'Thinking' ? { statusVerb: msg.verb } : {}), ...(msg.verb && msg.verb !== 'Thinking' ? { statusVerb: msg.verb } : {}),
...(msg.tokens ? { tokenUsage: { ...session.tokenUsage, output_tokens: msg.tokens } } : {}), ...(msg.tokens ? { tokenUsage: { ...session.tokenUsage, output_tokens: msg.tokens } } : {}),
...(msg.state === 'idle' ? { activeThinkingId: null, statusVerb: '' } : {}), ...(msg.state === 'idle' ? { activeThinkingId: null, statusVerb: '' } : {}),
...(shouldFlush ? { ...(shouldFlush ? {
messages: [...session.messages, { id: nextId(), type: 'assistant_text' as const, content: pendingText, timestamp: Date.now() }], messages: appendAssistantTextMessage(session.messages, pendingText, Date.now()),
streamingText: '', streamingText: '',
} : {}), } : pendingText !== session.streamingText ? { streamingText: pendingText } : {}),
} }
}) })
if (msg.state === 'idle') { if (msg.state === 'idle') {
@ -398,15 +427,19 @@ export const useChatStore = create<ChatStore>((set, get) => ({
case 'content_start': { case 'content_start': {
const session = get().sessions[sessionId] const session = get().sessions[sessionId]
if (!session) break if (!session) break
const pendingText = session.streamingText.trim() const pendingText = `${session.streamingText}${consumePendingDelta()}`
if (pendingText) { if (msg.blockType !== 'text' && pendingText.trim()) {
update((s) => ({ 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: '', streamingText: '',
})) }))
} }
if (msg.blockType === 'text') { 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') { } else if (msg.blockType === 'tool_use') {
update(() => ({ update(() => ({
activeToolUseId: msg.toolUseId ?? null, activeToolUseId: msg.toolUseId ?? null,
@ -436,9 +469,9 @@ export const useChatStore = create<ChatStore>((set, get) => ({
case 'thinking': case 'thinking':
update((s) => { update((s) => {
const pendingText = s.streamingText.trim() const pendingText = `${s.streamingText}${consumePendingDelta()}`
const base = pendingText const base = pendingText.trim()
? [...s.messages, { id: nextId(), type: 'assistant_text' as const, content: pendingText, timestamp: Date.now() }] ? appendAssistantTextMessage(s.messages, pendingText, Date.now())
: s.messages : s.messages
const last = base[base.length - 1] const last = base[base.length - 1]
if (last && last.type === 'thinking') { if (last && last.type === 'thinking') {
@ -518,12 +551,14 @@ export const useChatStore = create<ChatStore>((set, get) => ({
case 'message_complete': { case 'message_complete': {
const session = get().sessions[sessionId] const session = get().sessions[sessionId]
if (!session) break if (!session) break
const text = session.streamingText const text = `${session.streamingText}${consumePendingDelta()}`
if (text) { if (text.trim()) {
update((s) => ({ update((s) => ({
messages: [...s.messages, { id: nextId(), type: 'assistant_text', content: text, timestamp: Date.now() }], messages: appendAssistantTextMessage(s.messages, text, Date.now()),
streamingText: '', streamingText: '',
})) }))
} else if (text !== session.streamingText) {
update(() => ({ streamingText: text }))
} }
if (session.elapsedTimer) clearInterval(session.elapsedTimer) if (session.elapsedTimer) clearInterval(session.elapsedTimer)
update(() => ({ update(() => ({
@ -539,12 +574,12 @@ export const useChatStore = create<ChatStore>((set, get) => ({
case 'error': case 'error':
update((s) => { update((s) => {
const pendingText = s.streamingText.trim() const pendingText = `${s.streamingText}${consumePendingDelta()}`
const newMessages = [...s.messages] let newMessages = s.messages
if (pendingText) { if (pendingText.trim()) {
newMessages.push({ id: nextId(), type: 'assistant_text' as const, content: pendingText, timestamp: Date.now() }) 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 { return {
messages: newMessages, messages: newMessages,
chatState: 'idle', chatState: 'idle',
@ -666,6 +701,30 @@ function extractVisibleTeammateMessageContents(text: string): string[] {
return contents 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 = { type HistoryMappingOptions = {
includeTeammateMessages?: boolean includeTeammateMessages?: boolean
} }
@ -775,7 +834,7 @@ export function mapHistoryMessagesToUiMessages(
if ((msg.type === 'assistant' || msg.type === 'tool_use') && Array.isArray(msg.content)) { if ((msg.type === 'assistant' || msg.type === 'tool_use') && Array.isArray(msg.content)) {
for (const block of msg.content as AssistantHistoryBlock[]) { for (const block of msg.content as AssistantHistoryBlock[]) {
if (block.type === 'thinking' && block.thinking) uiMessages.push({ id: nextId(), type: 'thinking', content: block.thinking, timestamp }) 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 }) 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 continue