diff --git a/adapters/common/__tests__/message-buffer.test.ts b/adapters/common/__tests__/message-buffer.test.ts index e531bc67..9e8a037b 100644 --- a/adapters/common/__tests__/message-buffer.test.ts +++ b/adapters/common/__tests__/message-buffer.test.ts @@ -65,6 +65,39 @@ describe('MessageBuffer', () => { expect(flushed.length).toBe(0) }) + it('waits for an in-flight flush before complete resolves', async () => { + let releaseFlush!: () => void + let flushStarted = false + const flushed: Array<{ text: string; isComplete: boolean }> = [] + const buf = new MessageBuffer( + async (text, isComplete) => { + flushStarted = true + flushed.push({ text, isComplete }) + await new Promise((resolve) => { + releaseFlush = resolve + }) + }, + 10000, + 3, + ) + + buf.append('abcd') + await new Promise((resolve) => setTimeout(resolve, 0)) + expect(flushStarted).toBe(true) + + let completeResolved = false + const completing = buf.complete().then(() => { + completeResolved = true + }) + await new Promise((resolve) => setTimeout(resolve, 0)) + expect(completeResolved).toBe(false) + + releaseFlush() + await completing + expect(completeResolved).toBe(true) + expect(flushed).toEqual([{ text: 'abcd', isComplete: false }]) + }) + it('resets properly between messages', async () => { const flushed: string[] = [] const buf = new MessageBuffer( diff --git a/adapters/common/message-buffer.ts b/adapters/common/message-buffer.ts index 5698d7c6..b5a5006a 100644 --- a/adapters/common/message-buffer.ts +++ b/adapters/common/message-buffer.ts @@ -15,6 +15,7 @@ export class MessageBuffer { private timer: ReturnType | null = null private flushing = false private pendingComplete = false + private activeFlush: Promise | null = null constructor( private onFlush: FlushCallback, @@ -41,6 +42,7 @@ export class MessageBuffer { if (this.flushing) { // A flush is in-flight; mark pending so it fires after current flush finishes this.pendingComplete = true + await this.activeFlush return } await this.flush(true) @@ -69,23 +71,30 @@ export class MessageBuffer { clearTimeout(this.timer) this.timer = null } - if (this.flushing) return + if (this.flushing) { + await this.activeFlush + return + } if (this.buffer.length === 0) return this.flushing = true const text = this.buffer this.buffer = '' - try { - await this.onFlush(text, isComplete) - } catch (err) { - console.error('[MessageBuffer] Flush error:', err) - } finally { - this.flushing = false - // If complete() was called while we were flushing, do the final flush now - if (this.pendingComplete) { - this.pendingComplete = false - await this.flush(true) + this.activeFlush = (async () => { + try { + await this.onFlush(text, isComplete) + } catch (err) { + console.error('[MessageBuffer] Flush error:', err) + } finally { + this.flushing = false + this.activeFlush = null + // If complete() was called while we were flushing, do the final flush now. + if (this.pendingComplete) { + this.pendingComplete = false + await this.flush(true) + } } - } + })() + await this.activeFlush } } diff --git a/adapters/dingtalk/__tests__/stream-state.test.ts b/adapters/dingtalk/__tests__/stream-state.test.ts index 50a4e932..920096ee 100644 --- a/adapters/dingtalk/__tests__/stream-state.test.ts +++ b/adapters/dingtalk/__tests__/stream-state.test.ts @@ -28,6 +28,7 @@ describe('DingTalk streaming state', () => { it('completes the existing stream before dropping it for a permission request', async () => { const flushed: Array<{ text: string; complete: boolean }> = [] + let finalized = false const buffer = new MessageBuffer( async (text, complete) => { flushed.push({ text, complete }) @@ -41,11 +42,35 @@ describe('DingTalk streaming state', () => { aiCardBuffers: new Map([['chat-1', buffer]]), streamingCards: new Map([['chat-1', Promise.resolve(null)]]), streamingCardText: new Map([['chat-1', 'already streamed']]), + finalize: async () => { + finalized = true + }, } await finishAndResetDingTalkStreamingState(state, 'chat-1') expect(flushed).toEqual([{ text: 'pre-permission text', complete: true }]) + expect(finalized).toBe(true) + expect(state.aiCardBuffers.has('chat-1')).toBe(false) + expect(state.streamingCards.has('chat-1')).toBe(false) + expect(state.streamingCardText.has('chat-1')).toBe(false) + }) + + it('finalizes an already-flushed card even when the message buffer is empty', async () => { + let finalized = false + const buffer = new MessageBuffer(async () => {}, 100, 1000) + const state = { + aiCardBuffers: new Map([['chat-1', buffer]]), + streamingCards: new Map([['chat-1', Promise.resolve(null)]]), + streamingCardText: new Map([['chat-1', 'already streamed']]), + finalize: async () => { + finalized = true + }, + } + + await finishAndResetDingTalkStreamingState(state, 'chat-1') + + expect(finalized).toBe(true) expect(state.aiCardBuffers.has('chat-1')).toBe(false) expect(state.streamingCards.has('chat-1')).toBe(false) expect(state.streamingCardText.has('chat-1')).toBe(false) diff --git a/adapters/dingtalk/index.ts b/adapters/dingtalk/index.ts index 7682da8c..4b2bce32 100644 --- a/adapters/dingtalk/index.ts +++ b/adapters/dingtalk/index.ts @@ -402,7 +402,7 @@ async function handleServerMessage(chatId: string, msg: ServerMessage): Promise< case 'message_complete': runtime.state = 'idle' runtime.verb = undefined - await aiCardBuffers.get(chatId)?.complete() + await finishAndResetDingTalkStreamingState({ aiCardBuffers, streamingCards, streamingCardText, finalize: () => flushToAiCard(chatId, '', true) }, chatId) break case 'error': runtime.state = 'idle' @@ -425,7 +425,7 @@ async function sendPermissionRequest(chatId: string, msg: ServerMessage): Promis const runtime = getRuntimeState(chatId) runtime.pendingPermissionCount += 1 runtime.state = 'permission_pending' - await finishAndResetDingTalkStreamingState({ aiCardBuffers, streamingCards, streamingCardText }, chatId) + await finishAndResetDingTalkStreamingState({ aiCardBuffers, streamingCards, streamingCardText, finalize: () => flushToAiCard(chatId, '', true) }, chatId) const set = pendingPermissions.get(chatId) ?? new Set() set.add(msg.requestId) diff --git a/adapters/dingtalk/stream-state.ts b/adapters/dingtalk/stream-state.ts index 7d144c9b..7ea4e527 100644 --- a/adapters/dingtalk/stream-state.ts +++ b/adapters/dingtalk/stream-state.ts @@ -5,6 +5,7 @@ export type DingTalkStreamingState = { aiCardBuffers: Map streamingCards: Map> streamingCardText: Map + finalize?: () => Promise } export function resetDingTalkStreamingState( @@ -22,5 +23,8 @@ export async function finishAndResetDingTalkStreamingState( chatId: string, ): Promise { await state.aiCardBuffers.get(chatId)?.complete() + if (state.finalize && (state.streamingCards.has(chatId) || state.streamingCardText.has(chatId))) { + await state.finalize() + } resetDingTalkStreamingState(state, chatId) }