From 4eb522f3f640a8ac4e01bd209eded7f7eda25eed Mon Sep 17 00:00:00 2001 From: digua Date: Wed, 19 Aug 2026 18:28:13 +0800 Subject: [PATCH] fix(ai): drop queued stream text after stop Stop generation now cancels the text batcher instead of flushing it, and queued deltas are ignored once the assistant message is no longer streaming. --- src/stores/aiChat.ts | 45 +++++++++----------------- src/stores/aiChatStreamBatcher.test.ts | 32 +++++++++++++++++- src/stores/aiChatStreamBatcher.ts | 41 +++++++++++++++++++++++ 3 files changed, 88 insertions(+), 30 deletions(-) diff --git a/src/stores/aiChat.ts b/src/stores/aiChat.ts index c8370987..b239acb6 100644 --- a/src/stores/aiChat.ts +++ b/src/stores/aiChat.ts @@ -44,7 +44,7 @@ import { toSerializableContentBlocks, } from './aiChatContentBlocks' import { createToolLifecycleTracker, type ToolStatus } from './aiToolLifecycle' -import { createAIStreamTextBatcher } from './aiChatStreamBatcher' +import { applyQueuedStreamTextDeltas, createAIStreamTextBatcher } from './aiChatStreamBatcher' export type { ToolStatus } from './aiToolLifecycle' @@ -666,6 +666,7 @@ export const useAIChatStore = defineStore('aiChatRuntime', () => { completeTurn: (hadToolCalls: boolean) => void settleProcessDuration: () => void flushPendingText: () => void + discardPendingText: () => void } function createStreamBlockHelpers(targetBuffer: AIChatBuffer, getAiMessageIndex: () => number): StreamBlockHelpers { @@ -705,35 +706,16 @@ export const useAIChatStore = defineStore('aiChatRuntime', () => { const streamTextBatcher = createAIStreamTextBatcher((deltas) => { const idx = getAiMessageIndex() const message = targetBuffer.messages[idx] - const blocks = message.contentBlocks || [] - let content = message.content - let changed = false - - for (const delta of deltas) { - const lastBlock = blocks[blocks.length - 1] - if (delta.type === 'content') { - if (delta.content.trim().length === 0 && (!lastBlock || lastBlock.type !== 'text')) continue - if (lastBlock?.type === 'text') lastBlock.text += delta.content - else blocks.push({ type: 'text', text: delta.content }) - content += delta.content - changed = true - continue - } - - const thinkTag = delta.thinkTag || 'think' - if (lastBlock?.type === 'think' && lastBlock.tag === thinkTag) { - lastBlock.text += delta.content - changed = true - } else if (delta.content.trim().length > 0) { - blocks.push({ type: 'think', tag: thinkTag, text: delta.content }) - changed = true - } - } - - if (changed) applyAIMessageUpdates({ contentBlocks: [...blocks], content }) + const applied = applyQueuedStreamTextDeltas(message, deltas) + if (!applied) return + applyAIMessageUpdates({ + content: applied.content, + contentBlocks: applied.contentBlocks as ContentBlock[], + }) }) const flushPendingText = () => streamTextBatcher.flush() + const discardPendingText = () => streamTextBatcher.cancel() const updateAIMessage = (updates: Partial) => { flushPendingText() @@ -939,6 +921,7 @@ export const useAIChatStore = defineStore('aiChatRuntime', () => { completeTurn, settleProcessDuration, flushPendingText, + discardPendingText, } } @@ -1067,6 +1050,7 @@ export const useAIChatStore = defineStore('aiChatRuntime', () => { completeTurn, settleProcessDuration, flushPendingText, + discardPendingText, } = createStreamBlockHelpers(targetBuffer, () => aiMessageIndex) const currentAssistantId = targetBuffer.assistantId ?? getDefaultGeneralAssistantId(state.locale) @@ -1314,11 +1298,12 @@ export const useAIChatStore = defineStore('aiChatRuntime', () => { setActiveTaskMeta(chatKey, content, agentReqId, resolvedAIChatId) const result = await agentPromise - flushPendingText() if (state.isAborted) { + discardPendingText() clearActiveTask(chatKey, agentReqId) return { success: false, reason: 'aborted' } } + flushPendingText() if (thisRequestId !== state.currentRequestId) { clearActiveTask(chatKey, agentReqId) @@ -1694,6 +1679,7 @@ export const useAIChatStore = defineStore('aiChatRuntime', () => { completeTurn, settleProcessDuration, flushPendingText, + discardPendingText, } = createStreamBlockHelpers(targetBuffer, () => aiMessageIndex) try { @@ -1884,12 +1870,13 @@ export const useAIChatStore = defineStore('aiChatRuntime', () => { setActiveTaskMeta(chatKey, content, agentReqId, state.currentAIChatId) const result = await agentPromise - flushPendingText() if (state.isAborted) { + discardPendingText() restoreOriginal() clearActiveTask(chatKey, agentReqId) return { success: false, reason: 'aborted' } } + flushPendingText() if (thisRequestId !== state.currentRequestId) { restoreOriginal() clearActiveTask(chatKey, agentReqId) diff --git a/src/stores/aiChatStreamBatcher.test.ts b/src/stores/aiChatStreamBatcher.test.ts index 0e681908..6baab20f 100644 --- a/src/stores/aiChatStreamBatcher.test.ts +++ b/src/stores/aiChatStreamBatcher.test.ts @@ -1,6 +1,6 @@ import assert from 'node:assert/strict' import test from 'node:test' -import { createAIStreamTextBatcher, type AIStreamTextDelta } from './aiChatStreamBatcher' +import { applyQueuedStreamTextDeltas, createAIStreamTextBatcher, type AIStreamTextDelta } from './aiChatStreamBatcher' function createManualScheduler() { let nextHandle = 1 @@ -69,6 +69,36 @@ test('cancels the pending frame without applying stale deltas', () => { batcher.push({ type: 'content', content: 'stale' }) batcher.cancel() frame.runFrame() + batcher.flush() assert.deepEqual(applied, []) }) + +test('does not append queued text after the assistant message has stopped', () => { + const message = { + isStreaming: false, + content: 'hello\n\n_（已停止生成）_', + contentBlocks: [{ type: 'text', text: 'hello' }], + } + + const applied = applyQueuedStreamTextDeltas(message, [{ type: 'content', content: ' world' }]) + + assert.equal(applied, null) + assert.equal(message.content, 'hello\n\n_（已停止生成）_') + assert.deepEqual(message.contentBlocks, [{ type: 'text', text: 'hello' }]) +}) + +test('appends queued text while the assistant message is still streaming', () => { + const message = { + isStreaming: true, + content: 'hel', + contentBlocks: [{ type: 'text', text: 'hel' }], + } + + const applied = applyQueuedStreamTextDeltas(message, [{ type: 'content', content: 'lo' }]) + + assert.deepEqual(applied, { + content: 'hello', + contentBlocks: [{ type: 'text', text: 'hello' }], + }) +}) diff --git a/src/stores/aiChatStreamBatcher.ts b/src/stores/aiChatStreamBatcher.ts index 02b4f2be..40f2e54c 100644 --- a/src/stores/aiChatStreamBatcher.ts +++ b/src/stores/aiChatStreamBatcher.ts @@ -2,6 +2,47 @@ export type AIStreamTextDelta = | { type: 'content'; content: string } | { type: 'think'; content: string; thinkTag?: string } +export interface StreamTextTarget { + isStreaming?: boolean + content: string + contentBlocks?: Array<{ type: string; text?: string; tag?: string }> +} + +export function applyQueuedStreamTextDeltas( + message: StreamTextTarget, + deltas: AIStreamTextDelta[] +): { content: string; contentBlocks: NonNullable } | null { + if (!message.isStreaming) return null + + const blocks = message.contentBlocks ? [...message.contentBlocks] : [] + let content = message.content + let changed = false + + for (const delta of deltas) { + const lastBlock = blocks[blocks.length - 1] + if (delta.type === 'content') { + if (delta.content.trim().length === 0 && lastBlock?.type !== 'text') continue + if (lastBlock?.type === 'text') lastBlock.text = (lastBlock.text || '') + delta.content + else blocks.push({ type: 'text', text: delta.content }) + content += delta.content + changed = true + continue + } + + const thinkTag = delta.thinkTag || 'think' + if (lastBlock?.type === 'think' && lastBlock.tag === thinkTag) { + lastBlock.text = (lastBlock.text || '') + delta.content + changed = true + } else if (delta.content.trim().length > 0) { + blocks.push({ type: 'think', tag: thinkTag, text: delta.content }) + changed = true + } + } + + if (!changed) return { content: message.content, contentBlocks: message.contentBlocks || [] } + return { content, contentBlocks: blocks } +} + interface FrameScheduler { request: (callback: () => void) => number cancel: (handle: number) => void