diff --git a/scripts/test-raycast-ai-abort-listeners.mjs b/scripts/test-raycast-ai-abort-listeners.mjs new file mode 100644 index 00000000..7e4e3dd7 --- /dev/null +++ b/scripts/test-raycast-ai-abort-listeners.mjs @@ -0,0 +1,203 @@ +#!/usr/bin/env node + +import test from 'node:test'; +import assert from 'node:assert/strict'; +import fs from 'node:fs'; +import path from 'node:path'; +import { transform } from 'esbuild'; + +const RAYCAST_API_PATH = path.resolve('src/renderer/src/raycast-api/index.tsx'); + +test('AI.ask detaches caller abort listeners on every terminal path', async (t) => { + await t.test('stream success', async () => { + const { AI, electron } = await loadRaycastAIShim(); + const tracked = createTrackedAbortController(); + + const stream = AI.ask('success prompt', { signal: tracked.signal }); + const requestId = electron.aiAskCalls[0].requestId; + + electron.handlers.chunk({ requestId, chunk: 'hello ' }); + electron.handlers.chunk({ requestId, chunk: 'world' }); + electron.handlers.done({ requestId }); + + assert.equal(await stream, 'hello world'); + assertDetached('success', tracked); + }); + + await t.test('stream error', async () => { + const { AI, electron } = await loadRaycastAIShim(); + const tracked = createTrackedAbortController(); + + const stream = AI.ask('error prompt', { signal: tracked.signal }); + const requestId = electron.aiAskCalls[0].requestId; + + electron.handlers.error({ requestId, error: 'stream failed' }); + + await assert.rejects(stream, /stream failed/); + assertDetached('stream error', tracked); + }); + + await t.test('aiAsk rejection', async () => { + const { AI } = await loadRaycastAIShim({ rejectAiAsk: true }); + const tracked = createTrackedAbortController(); + + const stream = AI.ask('ipc rejection prompt', { signal: tracked.signal }); + + await assert.rejects(stream, /aiAsk rejected/); + assertDetached('aiAsk rejection', tracked); + }); + + await t.test('explicit abort', async () => { + const { AI, electron } = await loadRaycastAIShim(); + const tracked = createTrackedAbortController(); + + const stream = AI.ask('abort prompt', { signal: tracked.signal }); + const requestId = electron.aiAskCalls[0].requestId; + + tracked.controller.abort(); + + await assert.rejects(stream, /Request aborted/); + assert.deepEqual(electron.cancellations, [requestId]); + assertDetached('explicit abort', tracked); + }); +}); + +function assertDetached(label, tracked) { + const stats = tracked.stats(); + console.log( + `AI.ask abort listeners after ${label}: active=${stats.active}, added=${stats.added}, removed=${stats.removed}` + ); + assert.equal(stats.active, 0, `${label} should not retain caller abort listeners`); + assert.equal(stats.added, 1, `${label} should attach one caller abort listener`); + assert.equal(stats.removed, 1, `${label} should detach the caller abort listener`); +} + +async function loadRaycastAIShim(options = {}) { + const electron = createFakeElectron(options); + const importNonce = `${Date.now()}-${Math.random()}`; + const testWindow = { + electron, + addEventListener() {}, + removeEventListener() {}, + }; + const testDocument = { + visibilityState: 'visible', + addEventListener() {}, + removeEventListener() {}, + }; + + globalThis.window = testWindow; + globalThis.document = testDocument; + + const source = fs.readFileSync(RAYCAST_API_PATH, 'utf8'); + const start = source.indexOf('type AICreativity ='); + const aiExport = source.indexOf('export const AI =', start); + const end = source.indexOf('// =====================================================================', aiExport); + + assert.notEqual(start, -1, 'expected to find AI section start'); + assert.notEqual(aiExport, -1, 'expected to find AI export'); + assert.notEqual(end, -1, 'expected to find AI section end'); + + const moduleSource = ` +async function refreshAIAvailabilityCache() {} +${source.slice(start, end)} +export const __testImportNonce = ${JSON.stringify(importNonce)}; +`; + + const { code } = await transform(moduleSource, { + loader: 'tsx', + format: 'esm', + target: 'es2020', + }); + const dataUrl = 'data:text/javascript;base64,' + Buffer.from(code).toString('base64'); + const module = await import(dataUrl); + return { AI: module.AI, electron }; +} + +function createFakeElectron(options = {}) { + const handlers = { + chunk: undefined, + done: undefined, + error: undefined, + }; + const aiAskCalls = []; + const cancellations = []; + + return { + handlers, + aiAskCalls, + cancellations, + onAIStreamChunk(handler) { + handlers.chunk = handler; + return () => { + if (handlers.chunk === handler) handlers.chunk = undefined; + }; + }, + onAIStreamDone(handler) { + handlers.done = handler; + return () => { + if (handlers.done === handler) handlers.done = undefined; + }; + }, + onAIStreamError(handler) { + handlers.error = handler; + return () => { + if (handlers.error === handler) handlers.error = undefined; + }; + }, + aiAsk(requestId, prompt, askOptions) { + aiAskCalls.push({ requestId, prompt, askOptions }); + if (options.rejectAiAsk) { + return Promise.reject(new Error('aiAsk rejected')); + } + return Promise.resolve(); + }, + aiCancel(requestId) { + cancellations.push(requestId); + return Promise.resolve(); + }, + }; +} + +function createTrackedAbortController() { + const controller = new AbortController(); + const signal = controller.signal; + const originalAddEventListener = signal.addEventListener.bind(signal); + const originalRemoveEventListener = signal.removeEventListener.bind(signal); + const activeAbortListeners = new Set(); + let added = 0; + let removed = 0; + + Object.defineProperty(signal, 'addEventListener', { + configurable: true, + value(type, listener, options) { + if (type === 'abort') { + activeAbortListeners.add(listener); + added += 1; + } + return originalAddEventListener(type, listener, options); + }, + }); + + Object.defineProperty(signal, 'removeEventListener', { + configurable: true, + value(type, listener, options) { + if (type === 'abort' && activeAbortListeners.delete(listener)) { + removed += 1; + } + return originalRemoveEventListener(type, listener, options); + }, + }); + + return { + controller, + signal, + stats() { + return { + active: activeAbortListeners.size, + added, + removed, + }; + }, + }; +} diff --git a/src/renderer/src/raycast-api/index.tsx b/src/renderer/src/raycast-api/index.tsx index 95940389..e58cc08c 100644 --- a/src/renderer/src/raycast-api/index.tsx +++ b/src/renderer/src/raycast-api/index.tsx @@ -1646,9 +1646,30 @@ class StreamingPromise implements PromiseLike { } // Global IPC listener registry — routes chunks to the right StreamingPromise -const _activeStreams = new Map(); +type ActiveAIStream = { + sp: StreamingPromise; + fullText: string; + abortSignal?: AbortSignal; + abortHandler?: () => void; +}; + +const _activeStreams = new Map(); let _aiListenersRegistered = false; +function finalizeAIStream(requestId: string): ActiveAIStream | undefined { + const entry = _activeStreams.get(requestId); + if (!entry) return undefined; + + if (entry.abortSignal && entry.abortHandler) { + entry.abortSignal.removeEventListener('abort', entry.abortHandler); + entry.abortSignal = undefined; + entry.abortHandler = undefined; + } + + _activeStreams.delete(requestId); + return entry; +} + function ensureAIListeners(): void { if (_aiListenersRegistered) return; _aiListenersRegistered = true; @@ -1665,18 +1686,16 @@ function ensureAIListeners(): void { }); electron.onAIStreamDone?.((data: { requestId: string }) => { - const entry = _activeStreams.get(data.requestId); + const entry = finalizeAIStream(data.requestId); if (entry) { entry.sp._complete(entry.fullText); - _activeStreams.delete(data.requestId); } }); electron.onAIStreamError?.((data: { requestId: string; error: string }) => { - const entry = _activeStreams.get(data.requestId); + const entry = finalizeAIStream(data.requestId); if (entry) { entry.sp._error(new Error(data.error)); - _activeStreams.delete(data.requestId); } }); } @@ -1729,10 +1748,9 @@ export const AI = { model: options?.model, creativity, }).catch((err: any) => { - const entry = _activeStreams.get(requestId); + const entry = finalizeAIStream(requestId); if (entry) { entry.sp._error(err); - _activeStreams.delete(requestId); } }); @@ -1740,17 +1758,22 @@ export const AI = { if (options?.signal) { if (options.signal.aborted) { electron.aiCancel?.(requestId); - setTimeout(() => sp._error(new Error('Request aborted')), 0); - _activeStreams.delete(requestId); + const entry = finalizeAIStream(requestId); + setTimeout(() => entry?.sp._error(new Error('Request aborted')), 0); } else { - options.signal.addEventListener('abort', () => { + const abortHandler = () => { electron.aiCancel?.(requestId); - const entry = _activeStreams.get(requestId); + const entry = finalizeAIStream(requestId); if (entry) { entry.sp._error(new Error('Request aborted')); - _activeStreams.delete(requestId); } - }, { once: true }); + }; + const entry = _activeStreams.get(requestId); + if (entry) { + entry.abortSignal = options.signal; + entry.abortHandler = abortHandler; + options.signal.addEventListener('abort', abortHandler, { once: true }); + } } }