diff --git a/scripts/test-ai-stream-ipc-batching.mjs b/scripts/test-ai-stream-ipc-batching.mjs new file mode 100644 index 00000000..58f5608a --- /dev/null +++ b/scripts/test-ai-stream-ipc-batching.mjs @@ -0,0 +1,234 @@ +#!/usr/bin/env node + +import test from 'node:test'; +import assert from 'node:assert/strict'; +import path from 'node:path'; +import { importTs } from './lib/ts-import.mjs'; + +const CHUNK_COUNT = 600; + +const { + AI_STREAM_IPC_FLUSH_MS, + createAIStreamIpcCoalescer, + forwardAIStreamChunksToIpc, +} = await importTs(path.resolve('src/main/ai-stream-ipc.ts')); + +function makeChunks(count = CHUNK_COUNT) { + return Array.from({ length: count }, (_, index) => `chunk-${index};`); +} + +async function* streamChunks(chunks) { + for (const chunk of chunks) { + yield chunk; + } +} + +function createManualScheduler() { + let nextHandle = 1; + const callbacks = new Map(); + + return { + 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; + }, + pendingCount() { + return callbacks.size; + }, + }; +} + +function createSender() { + const events = []; + return { + events, + send(channel, payload) { + events.push({ channel, payload }); + }, + }; +} + +function aiStreamChunkEvents(events) { + return events.filter((event) => event.channel === 'ai-stream-chunk'); +} + +function joinedChunkText(events) { + return aiStreamChunkEvents(events).map((event) => event.payload.chunk).join(''); +} + +async function simulateLegacyIpcSend(chunks) { + const sender = createSender(); + for await (const chunk of streamChunks(chunks)) { + sender.send('ai-stream-chunk', { requestId: 'request-1', chunk }); + } + return sender.events; +} + +async function simulateBatchedIpcSend(chunks) { + const sender = createSender(); + const scheduler = createManualScheduler(); + const controller = new AbortController(); + const coalescer = createAIStreamIpcCoalescer({ + requestId: 'request-1', + sender, + flushIntervalMs: AI_STREAM_IPC_FLUSH_MS, + scheduleFlush: scheduler.schedule, + cancelFlush: scheduler.cancel, + }); + + await forwardAIStreamChunksToIpc(streamChunks(chunks), coalescer, controller.signal); + + assert.equal(sender.events.length, 0, 'burst should wait for the scheduled IPC flush'); + assert.equal(scheduler.pendingCount(), 1, 'burst should schedule one coalesced IPC flush'); + coalescer.flush(); + assert.equal(scheduler.pendingCount(), 0); + + return sender.events; +} + +test('main-process AI stream IPC batching coalesces provider chunk bursts', async (t) => { + const chunks = makeChunks(); + const expectedText = chunks.join(''); + + const legacyEvents = await simulateLegacyIpcSend(chunks); + const batchedEvents = await simulateBatchedIpcSend(chunks); + + const legacyChunkSends = aiStreamChunkEvents(legacyEvents).length; + const batchedChunkSends = aiStreamChunkEvents(batchedEvents).length; + + assert.equal(legacyChunkSends, CHUNK_COUNT); + assert.equal(joinedChunkText(legacyEvents), expectedText); + assert.equal(joinedChunkText(batchedEvents), expectedText); + assert.ok(batchedChunkSends < legacyChunkSends / 10); + + t.diagnostic(`baseline ai-stream-chunk IPC sends: ${legacyChunkSends} for ${CHUNK_COUNT} provider chunks`); + t.diagnostic(`batched ai-stream-chunk IPC sends: ${batchedChunkSends} for ${CHUNK_COUNT} provider chunks`); +}); + +test('main-process AI stream IPC flushes pending text before done', () => { + const sender = createSender(); + const scheduler = createManualScheduler(); + const coalescer = createAIStreamIpcCoalescer({ + requestId: 'done-request', + sender, + scheduleFlush: scheduler.schedule, + cancelFlush: scheduler.cancel, + }); + + coalescer.appendChunk('final '); + coalescer.appendChunk('answer'); + coalescer.flush(); + sender.send('ai-stream-done', { requestId: 'done-request' }); + + assert.deepEqual(sender.events.map((event) => event.channel), [ + 'ai-stream-chunk', + 'ai-stream-done', + ]); + assert.equal(joinedChunkText(sender.events), 'final answer'); + assert.equal(scheduler.pendingCount(), 0); +}); + +test('main-process AI stream IPC flushes pending text before error', async () => { + const sender = createSender(); + const scheduler = createManualScheduler(); + const controller = new AbortController(); + const coalescer = createAIStreamIpcCoalescer({ + requestId: 'error-request', + sender, + scheduleFlush: scheduler.schedule, + cancelFlush: scheduler.cancel, + }); + + async function* failingStream() { + yield 'partial '; + yield 'answer'; + throw new Error('provider failed'); + } + + try { + await forwardAIStreamChunksToIpc(failingStream(), coalescer, controller.signal); + assert.fail('expected provider stream to fail'); + } catch (error) { + coalescer.flush(); + sender.send('ai-stream-error', { + requestId: 'error-request', + error: error instanceof Error ? error.message : 'AI request failed', + }); + } + + assert.deepEqual(sender.events.map((event) => event.channel), [ + 'ai-stream-chunk', + 'ai-stream-error', + ]); + assert.equal(joinedChunkText(sender.events), 'partial answer'); + assert.equal(sender.events[1].payload.error, 'provider failed'); + assert.equal(scheduler.pendingCount(), 0); +}); + +test('main-process AI stream IPC flushes pending text during cancellation cleanup', () => { + const sender = createSender(); + const scheduler = createManualScheduler(); + const controller = new AbortController(); + const coalescer = createAIStreamIpcCoalescer({ + requestId: 'cancel-request', + sender, + scheduleFlush: scheduler.schedule, + cancelFlush: scheduler.cancel, + }); + + coalescer.appendChunk('visible before cancel'); + assert.equal(coalescer.hasPendingChunk(), true); + assert.equal(coalescer.hasPendingFlush(), true); + + coalescer.flush(); + controller.abort(); + + assert.equal(controller.signal.aborted, true); + assert.equal(joinedChunkText(sender.events), 'visible before cancel'); + assert.equal(aiStreamChunkEvents(sender.events).length, 1); + assert.equal(coalescer.hasPendingChunk(), false); + assert.equal(coalescer.hasPendingFlush(), false); + assert.equal(scheduler.pendingCount(), 0); +}); + +test('main-process AI stream IPC flushes the previous buffer before request replacement', () => { + const sender = createSender(); + const scheduler = createManualScheduler(); + const previous = createAIStreamIpcCoalescer({ + requestId: 'same-request', + sender, + scheduleFlush: scheduler.schedule, + cancelFlush: scheduler.cancel, + }); + + previous.appendChunk('old buffered text'); + previous.flush(); + + const replacement = createAIStreamIpcCoalescer({ + requestId: 'same-request', + sender, + scheduleFlush: scheduler.schedule, + cancelFlush: scheduler.cancel, + }); + replacement.appendChunk('new buffered text'); + replacement.flush(); + + assert.deepEqual( + aiStreamChunkEvents(sender.events).map((event) => event.payload.chunk), + ['old buffered text', 'new buffered text'], + ); + assert.equal(scheduler.pendingCount(), 0); +}); diff --git a/src/main/ai-stream-ipc.ts b/src/main/ai-stream-ipc.ts new file mode 100644 index 00000000..662c56de --- /dev/null +++ b/src/main/ai-stream-ipc.ts @@ -0,0 +1,82 @@ +export const AI_STREAM_IPC_FLUSH_MS = 16; + +export interface AIStreamIpcSender { + send: (channel: 'ai-stream-chunk', payload: { requestId: string; chunk: string }) => void; +} + +export interface AIStreamIpcCoalescerOptions { + requestId: string; + sender: AIStreamIpcSender; + flushIntervalMs?: number; + scheduleFlush?: (callback: () => void, delayMs: number) => unknown; + cancelFlush?: (handle: unknown) => void; +} + +export interface AIStreamIpcCoalescer { + appendChunk: (chunk: string) => void; + cancelPendingFlush: () => void; + flush: () => boolean; + hasPendingChunk: () => boolean; + hasPendingFlush: () => boolean; +} + +export function createAIStreamIpcCoalescer({ + requestId, + sender, + flushIntervalMs = AI_STREAM_IPC_FLUSH_MS, + scheduleFlush = (callback, delayMs) => setTimeout(callback, delayMs), + cancelFlush = (handle) => clearTimeout(handle as ReturnType), +}: AIStreamIpcCoalescerOptions): AIStreamIpcCoalescer { + let pendingChunk = ''; + let flushHandle: unknown = null; + + const cancelPendingFlush = () => { + if (flushHandle === null) return; + cancelFlush(flushHandle); + flushHandle = null; + }; + + const flush = () => { + cancelPendingFlush(); + if (!pendingChunk) return false; + const chunk = pendingChunk; + pendingChunk = ''; + sender.send('ai-stream-chunk', { requestId, chunk }); + return true; + }; + + const scheduleNextFlush = () => { + if (flushHandle !== null) return; + flushHandle = scheduleFlush(() => { + flushHandle = null; + flush(); + }, flushIntervalMs); + }; + + return { + appendChunk(chunk) { + if (!chunk) return; + pendingChunk += chunk; + scheduleNextFlush(); + }, + cancelPendingFlush, + flush, + hasPendingChunk() { + return pendingChunk.length > 0; + }, + hasPendingFlush() { + return flushHandle !== null; + }, + }; +} + +export async function forwardAIStreamChunksToIpc( + chunks: AsyncIterable, + coalescer: AIStreamIpcCoalescer, + signal: AbortSignal +): Promise { + for await (const chunk of chunks) { + if (signal.aborted) break; + coalescer.appendChunk(chunk); + } +} diff --git a/src/main/main.ts b/src/main/main.ts index cf510437..39f98772 100644 --- a/src/main/main.ts +++ b/src/main/main.ts @@ -44,6 +44,11 @@ import { import type { AppSettings, BrowserProfileSetting, BrowserProfileFilters, BrowserProfileFilterKind, RelocateMode } from './settings-store'; import { recordRootSearchLaunchInState, type RootSearchRankingState } from '../shared/root-search-ranking-state'; import { streamAI, streamAIChat, isAIAvailable, transcribeAudio } from './ai-provider'; +import { + createAIStreamIpcCoalescer, + forwardAIStreamChunksToIpc, + type AIStreamIpcCoalescer, +} from './ai-stream-ipc'; import { scanAppRemnants } from './app-uninstaller'; import * as soulverCalculator from './soulver-calculator'; import { addMemory, buildMemoryContextSystemPrompt } from './memory'; @@ -3074,7 +3079,12 @@ function resolveAppIconDataUrl(appPath: string, size = 32): string | null { } let launcherEntryFrontmostApp: FrontmostAppContext | null = null; const registeredHotkeys = new Map(); // shortcut → commandId -const activeAIRequests = new Map(); // requestId → controller +type ActiveAIRequest = { + controller: AbortController; + stream: AIStreamIpcCoalescer; +}; +const activeAIRequests = new Map(); // requestId → active request +const activeOllamaPullRequests = new Map(); // requestId → controller const pendingOAuthCallbackUrls: string[] = []; let snippetExpanderProcess: any = null; let snippetExpanderStdoutBuffer = ''; @@ -3097,6 +3107,24 @@ let emojiPickerCurrentQuery = ''; let emojiPickerCurrentPrefixLen = 1; let emojiPickerSelectedIdx = 0; let nativeSpeechProcess: any = null; + +function startActiveAIRequest(requestId: string, sender: { send: (channel: 'ai-stream-chunk', payload: { requestId: string; chunk: string }) => void }): ActiveAIRequest { + activeAIRequests.get(requestId)?.stream.flush(); + const controller = new AbortController(); + const request = { + controller, + stream: createAIStreamIpcCoalescer({ requestId, sender }), + }; + activeAIRequests.set(requestId, request); + return request; +} + +function finishActiveAIRequest(requestId: string, request: ActiveAIRequest): void { + request.stream.flush(); + if (activeAIRequests.get(requestId) === request) { + activeAIRequests.delete(requestId); + } +} let nativeSpeechStdoutBuffer = ''; let nativeColorPickerPromise: Promise | null = null; let keyboardLockProcess: any = null; @@ -17644,8 +17672,8 @@ if let tiff = image?.tiffRepresentation { return; } - const controller = new AbortController(); - activeAIRequests.set(requestId, controller); + const activeRequest = startActiveAIRequest(requestId, event.sender); + const { controller, stream } = activeRequest; try { const memoryContextSystemPrompt = await buildMemoryContextSystemPrompt( @@ -17665,30 +17693,35 @@ if let tiff = image?.tiffRepresentation { signal: controller.signal, }); - for await (const chunk of gen) { - if (controller.signal.aborted) break; - event.sender.send('ai-stream-chunk', { requestId, chunk }); - } + await forwardAIStreamChunksToIpc(gen, stream, controller.signal); if (!controller.signal.aborted) { + stream.flush(); event.sender.send('ai-stream-done', { requestId }); } } catch (e: any) { if (!controller.signal.aborted) { + stream.flush(); event.sender.send('ai-stream-error', { requestId, error: e?.message || 'AI request failed' }); } } finally { - activeAIRequests.delete(requestId); + finishActiveAIRequest(requestId, activeRequest); } } ); ipcMain.handle('ai-cancel', (_event: any, requestId: string) => { - const controller = activeAIRequests.get(requestId); - if (controller) { - controller.abort(); + const activeRequest = activeAIRequests.get(requestId); + if (activeRequest) { + activeRequest.stream.flush(); + activeRequest.controller.abort(); activeAIRequests.delete(requestId); } + const ollamaPullController = activeOllamaPullRequests.get(requestId); + if (ollamaPullController) { + ollamaPullController.abort(); + activeOllamaPullRequests.delete(requestId); + } }); ipcMain.handle( @@ -17709,8 +17742,8 @@ if let tiff = image?.tiffRepresentation { return; } - const controller = new AbortController(); - activeAIRequests.set(requestId, controller); + const activeRequest = startActiveAIRequest(requestId, event.sender); + const { controller, stream } = activeRequest; try { const latestUser = [...(messages || [])].reverse().find((m) => m.role === 'user'); @@ -17731,20 +17764,19 @@ if let tiff = image?.tiffRepresentation { signal: controller.signal, }); - for await (const chunk of gen) { - if (controller.signal.aborted) break; - event.sender.send('ai-stream-chunk', { requestId, chunk }); - } + await forwardAIStreamChunksToIpc(gen, stream, controller.signal); if (!controller.signal.aborted) { + stream.flush(); event.sender.send('ai-stream-done', { requestId }); } } catch (e: any) { if (!controller.signal.aborted) { + stream.flush(); event.sender.send('ai-stream-error', { requestId, error: e?.message || 'AI request failed' }); } } finally { - activeAIRequests.delete(requestId); + finishActiveAIRequest(requestId, activeRequest); } } ); @@ -18270,7 +18302,7 @@ if let tiff = image?.tiffRepresentation { const mod = url.protocol === 'https:' ? require('https') : require('http'); const controller = new AbortController(); - activeAIRequests.set(requestId, controller); + activeOllamaPullRequests.set(requestId, controller); const body = JSON.stringify({ name: modelName, stream: true }); @@ -18291,7 +18323,7 @@ if let tiff = image?.tiffRepresentation { requestId, error: `HTTP ${res.statusCode}: ${errBody.slice(0, 200)}`, }); - activeAIRequests.delete(requestId); + activeOllamaPullRequests.delete(requestId); }); return; } @@ -18335,7 +18367,7 @@ if let tiff = image?.tiffRepresentation { if (!controller.signal.aborted) { event.sender.send('ollama-pull-done', { requestId }); } - activeAIRequests.delete(requestId); + activeOllamaPullRequests.delete(requestId); }); } ); @@ -18347,7 +18379,7 @@ if let tiff = image?.tiffRepresentation { error: err.message || 'Failed to pull model', }); } - activeAIRequests.delete(requestId); + activeOllamaPullRequests.delete(requestId); }); if (controller.signal.aborted) {