diff --git a/scripts/test-ai-chat-stream-batching.mjs b/scripts/test-ai-chat-stream-batching.mjs new file mode 100644 index 00000000..066a2c4b --- /dev/null +++ b/scripts/test-ai-chat-stream-batching.mjs @@ -0,0 +1,263 @@ +#!/usr/bin/env node + +import assert from 'node:assert/strict'; +import fs from 'node:fs'; +import path from 'node:path'; +import { performance } from 'node:perf_hooks'; +import vm from 'node:vm'; +import { createRequire } from 'node:module'; + +const require = createRequire(import.meta.url); +const ts = require('typescript'); + +const CHUNK_COUNT = 240; + +function loadTsModule(filePath) { + const resolvedPath = path.resolve(filePath); + const source = fs.readFileSync(resolvedPath, 'utf8'); + const transpiled = ts.transpileModule(source, { + compilerOptions: { + module: ts.ModuleKind.CommonJS, + target: ts.ScriptTarget.ES2022, + esModuleInterop: true, + importsNotUsedAsValues: ts.ImportsNotUsedAsValues.Remove, + }, + fileName: resolvedPath, + }); + + const module = { exports: {} }; + vm.runInNewContext(transpiled.outputText, { + module, + exports: module.exports, + require, + console, + setTimeout, + clearTimeout, + }, { filename: resolvedPath }); + return module.exports; +} + +const { + AI_CHAT_STREAM_FLUSH_MS, + createAiChatStreamBuffer, +} = loadTsModule('src/renderer/src/utils/ai-chat-stream-buffer.ts'); + +function makeChunks(count = CHUNK_COUNT) { + return Array.from({ length: count }, (_, index) => `chunk-${index}\n`); +} + +function simulateUnbatchedStreaming(chunks) { + let content = ''; + let visibleUpdates = 0; + let markdownReparseChars = 0; + let layoutMeasurements = 0; + + const startedAt = performance.now(); + for (const chunk of chunks) { + content += chunk; + visibleUpdates += 1; + markdownReparseChars += content.length; + layoutMeasurements += 1; + } + + return { + content, + elapsedMs: performance.now() - startedAt, + layoutMeasurements, + markdownReparseChars, + visibleUpdates, + }; +} + +function createManualScheduler() { + let nextHandle = 1; + const callbacks = new Map(); + + return { + get size() { + return callbacks.size; + }, + schedule(callback) { + const handle = nextHandle; + nextHandle += 1; + callbacks.set(handle, callback); + return handle; + }, + cancel(handle) { + callbacks.delete(handle); + }, + runNext() { + const next = callbacks.entries().next(); + if (next.done) return false; + const [handle, callback] = next.value; + callbacks.delete(handle); + callback(); + return true; + }, + }; +} + +function createRenderCounters() { + return { + content: '', + layoutMeasurements: 0, + markdownReparseChars: 0, + visibleUpdates: 0, + onFlush(content) { + this.content = content; + this.visibleUpdates += 1; + this.layoutMeasurements += 1; + this.markdownReparseChars += content.length; + }, + }; +} + +function simulateBatchedBurstStreaming(chunks) { + const scheduler = createManualScheduler(); + const counters = createRenderCounters(); + const startedAt = performance.now(); + const buffer = createAiChatStreamBuffer({ + flushIntervalMs: AI_CHAT_STREAM_FLUSH_MS, + onFlush: (content) => counters.onFlush(content), + scheduleFlush: scheduler.schedule, + cancelFlush: scheduler.cancel, + }); + + for (const chunk of chunks) { + buffer.append(chunk); + } + assert.equal(scheduler.size, 1); + buffer.flushNow(); + assert.equal(scheduler.size, 0); + + return { + content: counters.content, + elapsedMs: performance.now() - startedAt, + layoutMeasurements: counters.layoutMeasurements, + markdownReparseChars: counters.markdownReparseChars, + visibleUpdates: counters.visibleUpdates, + }; +} + +function createConversationHarness() { + const scheduler = createManualScheduler(); + let messages = [ + { id: 'user-1', role: 'user', content: 'Question', createdAt: 1 }, + { id: 'assistant-1', role: 'assistant', content: '', createdAt: 2 }, + ]; + const persisted = []; + const counters = createRenderCounters(); + const buffer = createAiChatStreamBuffer({ + flushIntervalMs: AI_CHAT_STREAM_FLUSH_MS, + onFlush: (content) => { + counters.onFlush(content); + messages = messages.map((message) => ( + message.id === 'assistant-1' + ? { ...message, content } + : message + )); + }, + scheduleFlush: scheduler.schedule, + cancelFlush: scheduler.cancel, + }); + + return { + append(chunk) { + buffer.append(chunk); + }, + complete() { + buffer.flushNow(); + persisted.push(messages.map((message) => ({ ...message }))); + buffer.reset(); + }, + fail(error) { + const currentContent = buffer.getContent(); + buffer.append(`${currentContent ? '\n\n' : ''}Error: ${error}`); + buffer.flushNow(); + persisted.push(messages.map((message) => ({ ...message }))); + buffer.reset(); + }, + get counters() { + return counters; + }, + get messages() { + return messages; + }, + get persisted() { + return persisted; + }, + get schedulerSize() { + return scheduler.size; + }, + flushScheduled() { + return scheduler.runNext(); + }, + }; +} + +function test(name, fn) { + fn(); + console.log(`PASS ${name}`); +} + +const chunks = makeChunks(); +const baseline = simulateUnbatchedStreaming(chunks); +const batched = simulateBatchedBurstStreaming(chunks); + +assert.equal(baseline.content, chunks.join('')); +assert.equal(baseline.visibleUpdates, CHUNK_COUNT); +assert.equal(baseline.layoutMeasurements, CHUNK_COUNT); +assert.equal(batched.content, baseline.content); +assert.ok(batched.visibleUpdates < baseline.visibleUpdates); +assert.ok(batched.markdownReparseChars < baseline.markdownReparseChars); + +test('scheduled flush exposes partial streaming content', () => { + const harness = createConversationHarness(); + harness.append('Hello'); + assert.equal(harness.schedulerSize, 1); + assert.equal(harness.flushScheduled(), true); + assert.equal(harness.messages[1].content, 'Hello'); + harness.append(' world'); + assert.equal(harness.flushScheduled(), true); + assert.equal(harness.messages[1].content, 'Hello world'); + assert.equal(harness.counters.visibleUpdates, 2); +}); + +test('completion forces final flush before persistence', () => { + const harness = createConversationHarness(); + chunks.forEach((chunk) => harness.append(chunk)); + assert.equal(harness.messages[1].content, ''); + harness.complete(); + assert.equal(harness.persisted.length, 1); + assert.equal(harness.persisted[0][1].content, chunks.join('')); + assert.equal(harness.counters.visibleUpdates, 1); +}); + +test('error forces authoritative content plus error before persistence', () => { + const harness = createConversationHarness(); + harness.append('partial'); + harness.append(' answer'); + harness.fail('network failed'); + assert.equal(harness.persisted.length, 1); + assert.equal(harness.persisted[0][1].content, 'partial answer\n\nError: network failed'); + assert.equal(harness.counters.visibleUpdates, 1); +}); + +console.log(JSON.stringify({ + mode: 'unbatched-baseline', + chunks: CHUNK_COUNT, + visibleUpdates: baseline.visibleUpdates, + layoutMeasurements: baseline.layoutMeasurements, + markdownReparseChars: baseline.markdownReparseChars, + elapsedMs: Number(baseline.elapsedMs.toFixed(3)), +}, null, 2)); + +console.log(JSON.stringify({ + mode: 'batched-burst', + chunks: CHUNK_COUNT, + flushWindowMs: AI_CHAT_STREAM_FLUSH_MS, + visibleUpdates: batched.visibleUpdates, + layoutMeasurements: batched.layoutMeasurements, + markdownReparseChars: batched.markdownReparseChars, + elapsedMs: Number(batched.elapsedMs.toFixed(3)), +}, null, 2)); diff --git a/src/renderer/src/hooks/useAiChat.ts b/src/renderer/src/hooks/useAiChat.ts index fea32c94..072cfb4d 100644 --- a/src/renderer/src/hooks/useAiChat.ts +++ b/src/renderer/src/hooks/useAiChat.ts @@ -20,6 +20,7 @@ import type { AiChatMessage as AiMessage, AiChatSnapshot, } from '../../types/electron'; +import { createAiChatStreamBuffer } from '../utils/ai-chat-stream-buffer'; export type { AiConversation, AiMessage }; @@ -80,9 +81,65 @@ export function useAiChat({ onExitAiMode, setAiMode }: UseAiChatOptions): UseAiC const streamingMessageIdRef = useRef(null); const activeConversationIdRef = useRef(null); const messagesRef = useRef([]); + const streamBufferRef = useRef | null>(null); const aiInputRef = useRef(null); const aiResponseRef = useRef(null); + const setMessagesSnapshot = useCallback((nextMessages: AiMessage[]) => { + messagesRef.current = nextMessages; + setMessages(nextMessages); + }, []); + + const updateMessagesSnapshot = useCallback((updater: (current: AiMessage[]) => AiMessage[]) => { + const current = messagesRef.current; + const next = updater(current); + if (next === current) return current; + messagesRef.current = next; + setMessages(next); + return next; + }, []); + + const applyStreamingContent = useCallback( + (messageId: string, content: string) => { + updateMessagesSnapshot((current) => { + let changed = false; + const next = current.map((message) => { + if (message.id !== messageId) return message; + if (message.content === content) return message; + changed = true; + return { ...message, content }; + }); + return changed ? next : current; + }); + }, + [updateMessagesSnapshot] + ); + + if (streamBufferRef.current === null) { + streamBufferRef.current = createAiChatStreamBuffer({ + onFlush: (content) => { + const messageId = streamingMessageIdRef.current; + if (messageId) { + applyStreamingContent(messageId, content); + } + }, + }); + } + + const flushStreamingBuffer = useCallback(() => { + streamBufferRef.current?.flushNow(); + }, []); + + const resetStreamingBuffer = useCallback((content = '') => { + streamBufferRef.current?.reset(content); + }, []); + + useEffect(() => { + return () => { + streamBufferRef.current?.cancel(); + }; + }, []); + useEffect(() => { activeConversationIdRef.current = activeConversationId; }, [activeConversationId]); @@ -104,7 +161,7 @@ export function useAiChat({ onExitAiMode, setAiMode }: UseAiChatOptions): UseAiC const nextActive = nextConversations.find((conversation) => conversation.id === activeId); if (nextActive) { - setMessages(nextActive.messages); + setMessagesSnapshot(nextActive.messages); return; } @@ -112,7 +169,7 @@ export function useAiChat({ onExitAiMode, setAiMode }: UseAiChatOptions): UseAiC activeConversationIdRef.current = null; setActiveConversationId(null); } - }, []); + }, [setMessagesSnapshot]); const refreshSnapshot = useCallback(() => { void window.electron.getAiChatSnapshot() @@ -146,37 +203,29 @@ export function useAiChat({ onExitAiMode, setAiMode }: UseAiChatOptions): UseAiC const appendToStreamingMessage = (chunk: string) => { const msgId = streamingMessageIdRef.current; if (!msgId) return; - setMessages((prev) => - prev.map((message) => ( - message.id === msgId - ? { ...message, content: message.content + chunk } - : message - )) - ); + streamBufferRef.current?.append(chunk); }; const finalizeConversation = () => { const conversationId = activeConversationIdRef.current; if (!conversationId) return; - setMessages((current) => { - const existing = conversations.find((conversation) => conversation.id === conversationId); - const updatedConversation: AiConversation = { - id: conversationId, - title: - existing?.title && existing.title !== 'New Chat' - ? existing.title - : makeTitle(current.find((message) => message.role === 'user')?.content || 'New Chat'), - messages: current, - createdAt: existing?.createdAt ?? Date.now(), - updatedAt: Date.now(), - source: existing?.source || 'local', - ...(existing?.sourceConversationId ? { sourceConversationId: existing.sourceConversationId } : {}), - ...(existing?.metadata ? { metadata: existing.metadata } : {}), - }; - persistConversation(updatedConversation); - return current; - }); + const current = messagesRef.current; + const existing = conversations.find((conversation) => conversation.id === conversationId); + const updatedConversation: AiConversation = { + id: conversationId, + title: + existing?.title && existing.title !== 'New Chat' + ? existing.title + : makeTitle(current.find((message) => message.role === 'user')?.content || 'New Chat'), + messages: current, + createdAt: existing?.createdAt ?? Date.now(), + updatedAt: Date.now(), + source: existing?.source || 'local', + ...(existing?.sourceConversationId ? { sourceConversationId: existing.sourceConversationId } : {}), + ...(existing?.metadata ? { metadata: existing.metadata } : {}), + }; + persistConversation(updatedConversation); }; const handleChunk = (data: { requestId: string; chunk: string }) => { @@ -187,10 +236,12 @@ export function useAiChat({ onExitAiMode, setAiMode }: UseAiChatOptions): UseAiC const handleDone = (data: { requestId: string }) => { if (data.requestId === aiRequestIdRef.current) { + flushStreamingBuffer(); aiStreamingRef.current = false; setAiStreaming(false); - streamingMessageIdRef.current = null; finalizeConversation(); + streamingMessageIdRef.current = null; + resetStreamingBuffer(); } }; @@ -199,20 +250,14 @@ export function useAiChat({ onExitAiMode, setAiMode }: UseAiChatOptions): UseAiC aiStreamingRef.current = false; const msgId = streamingMessageIdRef.current; if (msgId) { - setMessages((prev) => - prev.map((message) => - message.id === msgId - ? { - ...message, - content: message.content + (message.content ? '\n\n' : '') + `Error: ${data.error}`, - } - : message - ) - ); + const currentContent = streamBufferRef.current?.getContent() ?? ''; + streamBufferRef.current?.append(`${currentContent ? '\n\n' : ''}Error: ${data.error}`); + flushStreamingBuffer(); } setAiStreaming(false); - streamingMessageIdRef.current = null; finalizeConversation(); + streamingMessageIdRef.current = null; + resetStreamingBuffer(); } }; @@ -227,7 +272,7 @@ export function useAiChat({ onExitAiMode, setAiMode }: UseAiChatOptions): UseAiC removeDone?.(); removeError?.(); }; - }, [hasBeenActivated, conversations, persistConversation]); + }, [hasBeenActivated, conversations, flushStreamingBuffer, persistConversation, resetStreamingBuffer]); useEffect(() => { if (aiResponseRef.current) { @@ -283,16 +328,16 @@ export function useAiChat({ onExitAiMode, setAiMode }: UseAiChatOptions): UseAiC content: '', createdAt: Date.now(), }; + resetStreamingBuffer(); streamingMessageIdRef.current = assistantMessage.id; - setMessages((prev) => { - const next = [...prev, userMessage, assistantMessage]; - sendChatTurn([...prev, userMessage]); - return next; - }); + const currentMessages = messagesRef.current; + const nextMessages = [...currentMessages, userMessage, assistantMessage]; + setMessagesSnapshot(nextMessages); + sendChatTurn([...currentMessages, userMessage]); setAiQuery(''); }, - [aiAvailable, persistConversation, sendChatTurn] + [aiAvailable, persistConversation, resetStreamingBuffer, sendChatTurn, setMessagesSnapshot] ); const startAiChat = useCallback( @@ -301,7 +346,8 @@ export function useAiChat({ onExitAiMode, setAiMode }: UseAiChatOptions): UseAiC setHasBeenActivated(true); activeConversationIdRef.current = null; setActiveConversationId(null); - setMessages([]); + resetStreamingBuffer(); + setMessagesSnapshot([]); setAiMode(true); const trimmed = searchQuery.trim(); if (trimmed) { @@ -310,10 +356,11 @@ export function useAiChat({ onExitAiMode, setAiMode }: UseAiChatOptions): UseAiC setAiQuery(''); } }, - [aiAvailable, setAiMode, sendMessage] + [aiAvailable, resetStreamingBuffer, setAiMode, sendMessage, setMessagesSnapshot] ); const stopStreaming = useCallback(() => { + flushStreamingBuffer(); if (aiRequestIdRef.current && aiStreamingRef.current) { window.electron.aiCancel(aiRequestIdRef.current); } @@ -321,13 +368,21 @@ export function useAiChat({ onExitAiMode, setAiMode }: UseAiChatOptions): UseAiC setAiStreaming(false); const messageId = streamingMessageIdRef.current; if (messageId) { - setMessages((prev) => - prev.map((message) => (message.id === messageId ? { ...message, cancelled: true } : message)) - ); + updateMessagesSnapshot((current) => { + let changed = false; + const next = current.map((message) => { + if (message.id !== messageId) return message; + if (message.cancelled) return message; + changed = true; + return { ...message, cancelled: true }; + }); + return changed ? next : current; + }); } streamingMessageIdRef.current = null; aiRequestIdRef.current = null; - }, []); + resetStreamingBuffer(); + }, [flushStreamingBuffer, resetStreamingBuffer, updateMessagesSnapshot]); const newChat = useCallback(() => { if (aiRequestIdRef.current && aiStreamingRef.current) { @@ -338,11 +393,12 @@ export function useAiChat({ onExitAiMode, setAiMode }: UseAiChatOptions): UseAiC streamingMessageIdRef.current = null; activeConversationIdRef.current = null; setActiveConversationId(null); - setMessages([]); + resetStreamingBuffer(); + setMessagesSnapshot([]); setAiStreaming(false); setAiQuery(''); setTimeout(() => aiInputRef.current?.focus(), 0); - }, []); + }, [resetStreamingBuffer, setMessagesSnapshot]); const selectConversation = useCallback((id: string) => { if (aiRequestIdRef.current && aiStreamingRef.current) { @@ -351,6 +407,7 @@ export function useAiChat({ onExitAiMode, setAiMode }: UseAiChatOptions): UseAiC aiRequestIdRef.current = null; aiStreamingRef.current = false; streamingMessageIdRef.current = null; + resetStreamingBuffer(); setAiStreaming(false); setConversations((current) => { @@ -358,13 +415,13 @@ export function useAiChat({ onExitAiMode, setAiMode }: UseAiChatOptions): UseAiC if (conversation) { activeConversationIdRef.current = id; setActiveConversationId(id); - setMessages(conversation.messages); + setMessagesSnapshot(conversation.messages); } return current; }); setAiQuery(''); setTimeout(() => aiInputRef.current?.focus(), 0); - }, []); + }, [resetStreamingBuffer, setMessagesSnapshot]); const deleteConversation = useCallback((id: string) => { setConversations((prev) => prev.filter((conversation) => conversation.id !== id)); @@ -372,7 +429,8 @@ export function useAiChat({ onExitAiMode, setAiMode }: UseAiChatOptions): UseAiC if (activeConversationIdRef.current === id) { activeConversationIdRef.current = null; setActiveConversationId(null); - setMessages([]); + resetStreamingBuffer(); + setMessagesSnapshot([]); if (aiRequestIdRef.current && aiStreamingRef.current) { window.electron.aiCancel(aiRequestIdRef.current); } @@ -381,9 +439,10 @@ export function useAiChat({ onExitAiMode, setAiMode }: UseAiChatOptions): UseAiC streamingMessageIdRef.current = null; setAiStreaming(false); } - }, []); + }, [resetStreamingBuffer, setMessagesSnapshot]); const exitAiMode = useCallback(() => { + flushStreamingBuffer(); if (aiRequestIdRef.current && aiStreamingRef.current) { window.electron.aiCancel(aiRequestIdRef.current); } @@ -393,8 +452,9 @@ export function useAiChat({ onExitAiMode, setAiMode }: UseAiChatOptions): UseAiC setAiMode(false); setAiStreaming(false); setAiQuery(''); + resetStreamingBuffer(); onExitAiMode?.(); - }, [setAiMode, onExitAiMode]); + }, [flushStreamingBuffer, resetStreamingBuffer, setAiMode, onExitAiMode]); useEffect(() => { if (messages.length === 0 && !aiQuery && !aiStreaming) return; diff --git a/src/renderer/src/utils/ai-chat-stream-buffer.ts b/src/renderer/src/utils/ai-chat-stream-buffer.ts new file mode 100644 index 00000000..ec66da81 --- /dev/null +++ b/src/renderer/src/utils/ai-chat-stream-buffer.ts @@ -0,0 +1,73 @@ +export const AI_CHAT_STREAM_FLUSH_MS = 40; + +export interface AiChatStreamBufferOptions { + flushIntervalMs?: number; + onFlush: (content: string) => void; + scheduleFlush?: (callback: () => void, delayMs: number) => unknown; + cancelFlush?: (handle: unknown) => void; +} + +export interface AiChatStreamBuffer { + append: (chunk: string) => void; + cancel: () => void; + flushNow: () => boolean; + getContent: () => string; + hasPendingFlush: () => boolean; + reset: (content?: string) => void; +} + +export function createAiChatStreamBuffer({ + flushIntervalMs = AI_CHAT_STREAM_FLUSH_MS, + onFlush, + scheduleFlush = (callback, delayMs) => setTimeout(callback, delayMs), + cancelFlush = (handle) => clearTimeout(handle as ReturnType), +}: AiChatStreamBufferOptions): AiChatStreamBuffer { + let content = ''; + let visibleContent = ''; + let flushHandle: unknown = null; + + const clearScheduledFlush = () => { + if (flushHandle === null) return; + cancelFlush(flushHandle); + flushHandle = null; + }; + + const flushNow = () => { + clearScheduledFlush(); + if (visibleContent === content) return false; + visibleContent = content; + onFlush(content); + return true; + }; + + const scheduleNextFlush = () => { + if (flushHandle !== null) return; + flushHandle = scheduleFlush(() => { + flushHandle = null; + flushNow(); + }, flushIntervalMs); + }; + + return { + append(chunk) { + if (!chunk) return; + content += chunk; + scheduleNextFlush(); + }, + cancel() { + clearScheduledFlush(); + }, + flushNow, + getContent() { + return content; + }, + hasPendingFlush() { + return flushHandle !== null; + }, + reset(nextContent = '') { + clearScheduledFlush(); + content = nextContent; + visibleContent = nextContent; + }, + }; +}