diff --git a/packages/stage-ui/src/composables/use-inference-preload.ts b/packages/stage-ui/src/composables/use-inference-preload.ts index 54df5dd95..a1c279401 100644 --- a/packages/stage-ui/src/composables/use-inference-preload.ts +++ b/packages/stage-ui/src/composables/use-inference-preload.ts @@ -37,7 +37,7 @@ export function useInferencePreload(options: UseInferencePreloadOptions = {}) { await detectWebGPU() const providersStore = useProvidersStore() - const tasks: { modelId: string, loader: () => Promise }[] = [] + const tasks: { modelId: string, loader: (signal: AbortSignal) => Promise }[] = [] // Check if Kokoro TTS is configured if (providersStore.configuredProviders['kokoro-local']) { @@ -53,9 +53,9 @@ export function useInferencePreload(options: UseInferencePreloadOptions = {}) { if (modelDef) { tasks.push({ modelId: `kokoro-${modelDef.id}`, - loader: async () => { + loader: async (signal) => { const adapter = await getKokoroAdapter() - await adapter.loadModel(modelDef.quantization, modelDef.platform) + await adapter.loadModel(modelDef.quantization, modelDef.platform, { signal }) }, }) } diff --git a/packages/stage-ui/src/composables/use-model-preload.ts b/packages/stage-ui/src/composables/use-model-preload.ts index 1b7dd01ee..5f3191180 100644 --- a/packages/stage-ui/src/composables/use-model-preload.ts +++ b/packages/stage-ui/src/composables/use-model-preload.ts @@ -5,6 +5,10 @@ * so they're ready before the user first needs them. * Uses `setTimeout` to defer loading and avoid blocking the main * thread during app startup. + * + * Cancellation uses an AbortController so the signal can be forwarded + * to adapter methods that accept `options.signal` — a cancelled preload + * aborts any in-flight model load instead of just ignoring the result. */ import { onUnmounted, ref } from 'vue' @@ -12,8 +16,12 @@ import { onUnmounted, ref } from 'vue' export interface PreloadTask { /** Human-readable model name for logging */ modelId: string - /** The async function that loads the model */ - loader: () => Promise + /** + * The async function that loads the model. Receives an `AbortSignal` + * that will fire if the preload is cancelled; loaders should forward + * it to adapter methods (e.g. `adapter.loadModel(q, d, { signal })`). + */ + loader: (signal: AbortSignal) => Promise } export interface UseModelPreloadOptions { @@ -28,7 +36,7 @@ export function useModelPreload(options: UseModelPreloadOptions = {}) { const preloadedModels = ref([]) const failedModels = ref([]) - let cancelled = false + let abortController: AbortController | null = null let timeoutId: ReturnType | undefined /** @@ -39,27 +47,38 @@ export function useModelPreload(options: UseModelPreloadOptions = {}) { if (tasks.length === 0) return - cancelled = false + // Fresh controller per scheduling — abort any prior in-flight preload first + if (abortController && !abortController.signal.aborted) + abortController.abort(new Error('Preload superseded by new schedule')) + abortController = new AbortController() + const signal = abortController.signal timeoutId = setTimeout(async () => { - if (cancelled) + if (signal.aborted) return preloading.value = true for (const task of tasks) { - if (cancelled) + if (signal.aborted) break try { // eslint-disable-next-line no-console console.debug(`[Preload] Loading ${task.modelId}...`) - await task.loader() + await task.loader(signal) preloadedModels.value.push(task.modelId) // eslint-disable-next-line no-console console.debug(`[Preload] ${task.modelId} ready`) } catch (error) { + // AbortError is expected when the preload is cancelled — don't + // treat it as a failure. + if ((error as Error)?.name === 'AbortError') { + // eslint-disable-next-line no-console + console.debug(`[Preload] ${task.modelId} aborted`) + break + } // Preload failures are non-fatal — model will load on first use console.warn(`[Preload] ${task.modelId} failed:`, error) failedModels.value.push(task.modelId) @@ -71,7 +90,8 @@ export function useModelPreload(options: UseModelPreloadOptions = {}) { } function cancelPreload(): void { - cancelled = true + if (abortController && !abortController.signal.aborted) + abortController.abort(new Error('Preload cancelled')) if (timeoutId !== undefined) { clearTimeout(timeoutId) timeoutId = undefined diff --git a/packages/stage-ui/src/libs/inference/adapters/background-removal.ts b/packages/stage-ui/src/libs/inference/adapters/background-removal.ts index 75ac5cba2..4fa4b47a4 100644 --- a/packages/stage-ui/src/libs/inference/adapters/background-removal.ts +++ b/packages/stage-ui/src/libs/inference/adapters/background-removal.ts @@ -16,7 +16,7 @@ import { removeInferenceStatus, updateInferenceStatus } from '../../../composabl import { MODEL_IDS, MODEL_NAMES, TIMEOUTS } from '../constants' import { getGPUCoordinator, getLoadQueue, MODEL_VRAM_ESTIMATES } from '../coordinator' import { LOAD_PRIORITY } from '../load-queue' -import { createRequestId } from '../protocol' +import { createRequestId, InferenceAbortError, throwIfAborted } from '../protocol' // --------------------------------------------------------------------------- // Types @@ -26,14 +26,22 @@ export interface BackgroundRemovalAdapter { /** * Load the background removal model in the worker. * Must be called before `processImage()`. + * Pass `options.signal` to cancel; rejects with `InferenceAbortError`. */ - load: (onProgress?: (p: ProgressPayload) => void) => Promise + load: ( + onProgress?: (p: ProgressPayload) => void, + options?: { signal?: AbortSignal }, + ) => Promise /** * Remove the background from an image. * Returns a new ImageData with the background alpha set to 0. + * Pass `options.signal` to cancel; rejects with `InferenceAbortError`. */ - processImage: (imageData: ImageData) => Promise + processImage: ( + imageData: ImageData, + options?: { signal?: AbortSignal }, + ) => Promise /** Terminate the worker */ terminate: () => void @@ -88,7 +96,8 @@ export function createBackgroundRemovalAdapter(): BackgroundRemovalAdapter { /** * Wait for a specific message type from the worker, filtered by requestId. - * Uses the unified protocol message types. + * Uses the unified protocol message types. Honors `signal` to cancel the + * wait (and notify the worker to discard the result). */ function waitForMessage( w: Worker, @@ -96,25 +105,35 @@ export function createBackgroundRemovalAdapter(): BackgroundRemovalAdapter { targetType: string, timeout: number, onOther?: (data: any) => void, + signal?: AbortSignal, ): Promise { return new Promise((resolve, reject) => { let timeoutId: ReturnType | undefined + let abortListener: (() => void) | null = null - const handler = (event: MessageEvent) => { + const cleanup = (): void => { + if (timeoutId !== undefined) + clearTimeout(timeoutId) + w.removeEventListener('message', handler) + if (abortListener && signal) + signal.removeEventListener('abort', abortListener) + } + + const handler = (event: MessageEvent): void => { if (event.data.requestId !== requestId) return if (event.data.type === targetType) { - if (timeoutId !== undefined) - clearTimeout(timeoutId) - w.removeEventListener('message', handler) + cleanup() resolve(event.data as T) } else if (event.data.type === 'error') { - if (timeoutId !== undefined) - clearTimeout(timeoutId) - w.removeEventListener('message', handler) - reject(new Error(event.data.payload?.message ?? 'Worker error')) + cleanup() + const code = event.data.payload?.code + if (code === 'CANCELLED') + reject(new InferenceAbortError(event.data.payload?.message)) + else + reject(new Error(event.data.payload?.message ?? 'Worker error')) } else { onOther?.(event.data) @@ -124,18 +143,40 @@ export function createBackgroundRemovalAdapter(): BackgroundRemovalAdapter { w.addEventListener('message', handler) timeoutId = setTimeout(() => { - w.removeEventListener('message', handler) + cleanup() reject(new Error(`Background removal: timeout after ${timeout}ms`)) }, timeout) + + if (signal) { + if (signal.aborted) { + cleanup() + w.postMessage({ type: 'cancel', requestId: createRequestId(), targetRequestId: requestId }) + reject(new InferenceAbortError(typeof signal.reason === 'string' ? signal.reason : undefined)) + return + } + abortListener = () => { + cleanup() + w.postMessage({ type: 'cancel', requestId: createRequestId(), targetRequestId: requestId }) + const reason = signal.reason + reject(reason instanceof Error ? reason : new InferenceAbortError(typeof reason === 'string' ? reason : undefined)) + } + signal.addEventListener('abort', abortListener) + } }) } - async function load(onProgress?: (p: ProgressPayload) => void): Promise { + async function load( + onProgress?: (p: ProgressPayload) => void, + options?: { signal?: AbortSignal }, + ): Promise { + throwIfAborted(options?.signal) return operationMutex.runExclusive(async () => { + throwIfAborted(options?.signal) state = 'loading' updateInferenceStatus(MODEL_NAMES.BG_REMOVAL, { state: 'downloading', device: 'webgpu' }) return getLoadQueue().enqueue(MODEL_NAMES.BG_REMOVAL, LOAD_PRIORITY.BACKGROUND_REMOVAL, async () => { + throwIfAborted(options?.signal) const w = ensureWorker() const requestId = createRequestId() @@ -151,7 +192,7 @@ export function createBackgroundRemovalAdapter(): BackgroundRemovalAdapter { total: payload.total, }) } - }) + }, options?.signal) w.postMessage({ type: 'load-model', requestId, modelId: MODEL_IDS.BG_REMOVAL, device: 'webgpu' }) @@ -179,19 +220,31 @@ export function createBackgroundRemovalAdapter(): BackgroundRemovalAdapter { state = 'ready' updateInferenceStatus(MODEL_NAMES.BG_REMOVAL, { state: 'ready', device: actualDevice }) - }) + }, { signal: options?.signal }) }) } - async function processImage(imageData: ImageData): Promise { + async function processImage( + imageData: ImageData, + options?: { signal?: AbortSignal }, + ): Promise { + throwIfAborted(options?.signal) return defaultPerfTracer.withMeasure('inference', 'bg-removal-process', () => operationMutex.runExclusive(async () => { + throwIfAborted(options?.signal) if (!worker || (state !== 'ready' && state !== 'processing')) throw new Error('Model not loaded. Call load() first.') state = 'processing' const requestId = createRequestId() - const resultPromise = waitForMessage(worker, requestId, 'inference-result', PROCESS_TIMEOUT) + const resultPromise = waitForMessage( + worker, + requestId, + 'inference-result', + PROCESS_TIMEOUT, + undefined, + options?.signal, + ) // Send raw pixel data (transferable copy) const pixelsCopy = new Uint8ClampedArray(imageData.data) diff --git a/packages/stage-ui/src/libs/inference/adapters/kokoro.ts b/packages/stage-ui/src/libs/inference/adapters/kokoro.ts index c0629a2af..21705b87c 100644 --- a/packages/stage-ui/src/libs/inference/adapters/kokoro.ts +++ b/packages/stage-ui/src/libs/inference/adapters/kokoro.ts @@ -16,22 +16,36 @@ import { removeInferenceStatus, updateInferenceStatus } from '../../../composabl import { DEVICE_LOSS_WASM_THRESHOLD, MAX_RESTARTS, MODEL_NAMES, RESTART_DELAY_MS, TIMEOUTS } from '../constants' import { getGPUCoordinator, getLoadQueue, MODEL_VRAM_ESTIMATES } from '../coordinator' import { LOAD_PRIORITY } from '../load-queue' -import { classifyDeviceLossReason, classifyError, createRequestId } from '../protocol' +import { classifyDeviceLossReason, classifyError, createRequestId, InferenceAbortError, throwIfAborted } from '../protocol' // --------------------------------------------------------------------------- // Types // --------------------------------------------------------------------------- export interface KokoroAdapter { - /** Load a TTS model with the given quantization and device */ + /** + * Load a TTS model with the given quantization and device. + * Pass `options.signal` to cancel the load; the returned promise will + * reject with `InferenceAbortError` (name: `'AbortError'`). + */ loadModel: ( quantization: string, device: string, - options?: { onProgress?: (p: ProgressPayload) => void }, + options?: { + onProgress?: (p: ProgressPayload) => void + signal?: AbortSignal + }, ) => Promise - /** Generate speech audio from text */ - generate: (text: string, voice: VoiceKey) => Promise + /** + * Generate speech audio from text. + * Pass `options.signal` to cancel; rejects with `InferenceAbortError`. + */ + generate: ( + text: string, + voice: VoiceKey, + options?: { signal?: AbortSignal }, + ) => Promise /** Get the voices from the last loaded model */ getVoices: () => Voices @@ -118,6 +132,10 @@ function writeString(view: DataView, offset: number, str: string): void { /** * Wait for a specific message type from the worker, filtered by requestId. * Calls `callback` for interleaved messages (e.g. progress). + * + * If `signal` is provided and aborts, the returned Promise rejects with + * `InferenceAbortError` and a `cancel` message is sent to the worker so + * it can discard the result when it eventually arrives. */ function waitForWorkerMessage( worker: Worker, @@ -125,25 +143,35 @@ function waitForWorkerMessage( targetType: string, timeout: number, callback?: (data: any) => void, + signal?: AbortSignal, ): Promise { return new Promise((resolve, reject) => { let timeoutId: ReturnType | undefined + let abortListener: (() => void) | null = null - const handler = (event: MessageEvent) => { + const cleanup = (): void => { + if (timeoutId !== undefined) + clearTimeout(timeoutId) + worker.removeEventListener('message', handler) + if (abortListener && signal) + signal.removeEventListener('abort', abortListener) + } + + const handler = (event: MessageEvent): void => { if (event.data.requestId !== requestId) return if (event.data.type === targetType) { - if (timeoutId !== undefined) - clearTimeout(timeoutId) - worker.removeEventListener('message', handler) + cleanup() resolve(event.data as T) } else if (event.data.type === 'error') { - if (timeoutId !== undefined) - clearTimeout(timeoutId) - worker.removeEventListener('message', handler) - reject(new Error(event.data.payload?.message ?? 'Worker error')) + cleanup() + const code = event.data.payload?.code + if (code === 'CANCELLED') + reject(new InferenceAbortError(event.data.payload?.message)) + else + reject(new Error(event.data.payload?.message ?? 'Worker error')) } else { callback?.(event.data) @@ -153,9 +181,26 @@ function waitForWorkerMessage( worker.addEventListener('message', handler) timeoutId = setTimeout(() => { - worker.removeEventListener('message', handler) + cleanup() reject(new Error(`Kokoro: timeout after ${timeout}ms waiting for '${targetType}'`)) }, timeout) + + if (signal) { + if (signal.aborted) { + cleanup() + // Tell the worker to discard the result when it arrives + worker.postMessage({ type: 'cancel', requestId: createRequestId(), targetRequestId: requestId }) + reject(new InferenceAbortError(typeof signal.reason === 'string' ? signal.reason : undefined)) + return + } + abortListener = () => { + cleanup() + worker.postMessage({ type: 'cancel', requestId: createRequestId(), targetRequestId: requestId }) + const reason = signal.reason + reject(reason instanceof Error ? reason : new InferenceAbortError(typeof reason === 'string' ? reason : undefined)) + } + signal.addEventListener('abort', abortListener) + } }) } @@ -270,7 +315,10 @@ export function createKokoroAdapter(): KokoroAdapter { async function loadModel( quantization: string, device: string, - options?: { onProgress?: (p: ProgressPayload) => void }, + options?: { + onProgress?: (p: ProgressPayload) => void + signal?: AbortSignal + }, ): Promise { // NOTICE: Proactive WASM promotion. If this adapter has suffered repeated // WebGPU device-loss events, webgpu is unreliable on this device and we @@ -287,10 +335,11 @@ export function createKokoroAdapter(): KokoroAdapter { ) effectiveDevice = 'wasm' } - + throwIfAborted(options?.signal) await ensureStarted() return defaultPerfTracer.withMeasure('inference', 'kokoro-load-model', () => operationMutex.runExclusive(async () => { + throwIfAborted(options?.signal) state = 'loading' const modelStatusId = `kokoro-${quantization}` @@ -303,7 +352,9 @@ export function createKokoroAdapter(): KokoroAdapter { // Use the global load queue to serialize model loads across all adapters return getLoadQueue().enqueue(modelStatusId, LOAD_PRIORITY.TTS, async () => { + throwIfAborted(options?.signal) const requestId = createRequestId() + // Signal is also passed to the queue below for pending-entry removal const readyPromise = waitForWorkerMessage(worker!, requestId, 'model-ready', LOAD_MODEL_TIMEOUT, (data) => { if (data.type === 'progress') { @@ -320,7 +371,7 @@ export function createKokoroAdapter(): KokoroAdapter { updateInferenceStatus(modelStatusId, { progress }) options?.onProgress?.(progress) } - }) + }, options?.signal) worker!.postMessage({ type: 'load-model', @@ -351,15 +402,25 @@ export function createKokoroAdapter(): KokoroAdapter { if (!voices) throw new Error('Kokoro worker did not return voice metadata') return voices - }) + }, { signal: options?.signal }) }), { quantization, device: effectiveDevice }).catch((error) => { + // Don't route AbortError through handleWorkerError — cancellation is + // not a worker failure and shouldn't trigger restart logic. + if ((error as Error)?.name === 'AbortError') + throw error handleWorkerError(error instanceof Error ? error : new Error(String(error))) throw error }) } - async function generate(text: string, voice: VoiceKey): Promise { + async function generate( + text: string, + voice: VoiceKey, + options?: { signal?: AbortSignal }, + ): Promise { + throwIfAborted(options?.signal) return defaultPerfTracer.withMeasure('inference', 'kokoro-generate', () => operationMutex.runExclusive(async () => { + throwIfAborted(options?.signal) if (!worker) throw new Error('Worker not initialized. Call loadModel() first.') @@ -370,7 +431,14 @@ export function createKokoroAdapter(): KokoroAdapter { state = 'running' const requestId = createRequestId() - const resultPromise = waitForWorkerMessage(worker, requestId, 'inference-result', GENERATE_TIMEOUT) + const resultPromise = waitForWorkerMessage( + worker, + requestId, + 'inference-result', + GENERATE_TIMEOUT, + undefined, + options?.signal, + ) worker.postMessage({ type: 'run-inference', diff --git a/packages/stage-ui/src/libs/inference/adapters/whisper.ts b/packages/stage-ui/src/libs/inference/adapters/whisper.ts index c685cb300..e81c35f03 100644 --- a/packages/stage-ui/src/libs/inference/adapters/whisper.ts +++ b/packages/stage-ui/src/libs/inference/adapters/whisper.ts @@ -16,7 +16,7 @@ import { removeInferenceStatus, updateInferenceStatus } from '../../../composabl import { DEVICE_LOSS_WASM_THRESHOLD, MAX_RESTARTS, MODEL_NAMES, RESTART_DELAY_MS, TIMEOUTS } from '../constants' import { getGPUCoordinator, getLoadQueue, MODEL_VRAM_ESTIMATES } from '../coordinator' import { LOAD_PRIORITY } from '../load-queue' -import { classifyDeviceLossReason, classifyError, createRequestId } from '../protocol' +import { classifyDeviceLossReason, classifyError, createRequestId, InferenceAbortError, throwIfAborted } from '../protocol' // --------------------------------------------------------------------------- // Types @@ -47,11 +47,23 @@ export type WhisperEvent | { type: 'error', payload: { code: string, message: string } } export interface WhisperAdapter { - /** Load the Whisper model */ - load: (onProgress?: (p: ProgressPayload) => void) => Promise + /** + * Load the Whisper model. + * Pass `options.signal` to cancel the load; rejects with `InferenceAbortError`. + */ + load: ( + onProgress?: (p: ProgressPayload) => void, + options?: { signal?: AbortSignal }, + ) => Promise - /** Transcribe audio, returning the text result */ - transcribe: (input: WhisperTranscribeInput) => Promise + /** + * Transcribe audio, returning the text result. + * Pass `options.signal` to cancel; rejects with `InferenceAbortError`. + */ + transcribe: ( + input: WhisperTranscribeInput, + options?: { signal?: AbortSignal }, + ) => Promise /** Terminate the worker */ terminate: () => void @@ -188,6 +200,8 @@ export function createWhisperAdapter(workerUrl: string | URL): WhisperAdapter { /** * Wait for a specific unified protocol message type, filtered by requestId. + * If `signal` is provided and aborts, sends a `cancel` message to the + * worker and rejects with `InferenceAbortError`. */ function waitForMessage( w: Worker, @@ -195,25 +209,35 @@ export function createWhisperAdapter(workerUrl: string | URL): WhisperAdapter { targetType: string, timeout: number, onOther?: (data: any) => void, + signal?: AbortSignal, ): Promise { return new Promise((resolve, reject) => { let timeoutId: ReturnType | undefined + let abortListener: (() => void) | null = null - const handler = (event: MessageEvent) => { + const cleanup = (): void => { + if (timeoutId !== undefined) + clearTimeout(timeoutId) + w.removeEventListener('message', handler) + if (abortListener && signal) + signal.removeEventListener('abort', abortListener) + } + + const handler = (event: MessageEvent): void => { if (event.data.requestId !== requestId) return if (event.data.type === targetType) { - if (timeoutId !== undefined) - clearTimeout(timeoutId) - w.removeEventListener('message', handler) + cleanup() resolve(event.data as T) } else if (event.data.type === 'error') { - if (timeoutId !== undefined) - clearTimeout(timeoutId) - w.removeEventListener('message', handler) - reject(new Error(event.data.payload?.message ?? 'Worker error')) + cleanup() + const code = event.data.payload?.code + if (code === 'CANCELLED') + reject(new InferenceAbortError(event.data.payload?.message)) + else + reject(new Error(event.data.payload?.message ?? 'Worker error')) } else { onOther?.(event.data) @@ -223,14 +247,31 @@ export function createWhisperAdapter(workerUrl: string | URL): WhisperAdapter { w.addEventListener('message', handler) timeoutId = setTimeout(() => { - w.removeEventListener('message', handler) + cleanup() reject(new Error(`Whisper: timeout after ${timeout}ms waiting for '${targetType}'`)) }, timeout) + + if (signal) { + if (signal.aborted) { + cleanup() + w.postMessage({ type: 'cancel', requestId: createRequestId(), targetRequestId: requestId }) + reject(new InferenceAbortError(typeof signal.reason === 'string' ? signal.reason : undefined)) + return + } + abortListener = () => { + cleanup() + w.postMessage({ type: 'cancel', requestId: createRequestId(), targetRequestId: requestId }) + const reason = signal.reason + reject(reason instanceof Error ? reason : new InferenceAbortError(typeof reason === 'string' ? reason : undefined)) + } + signal.addEventListener('abort', abortListener) + } }) } async function load( onProgress?: (p: ProgressPayload) => void, + options?: { signal?: AbortSignal }, ): Promise { // NOTICE: Proactive WASM promotion after repeated device-loss events. // See kokoro.ts for rationale. Whisper always requests 'webgpu' from the @@ -242,12 +283,14 @@ export function createWhisperAdapter(workerUrl: string | URL): WhisperAdapter { + `promoting load from webgpu to wasm.`, ) } - + throwIfAborted(options?.signal) return operationMutex.runExclusive(async () => { + throwIfAborted(options?.signal) state = 'loading' updateInferenceStatus(MODEL_NAMES.WHISPER, { state: 'downloading', device: requestedDevice as any }) return getLoadQueue().enqueue(MODEL_NAMES.WHISPER, LOAD_PRIORITY.ASR, async () => { + throwIfAborted(options?.signal) const w = ensureWorker() const requestId = createRequestId() @@ -263,7 +306,7 @@ export function createWhisperAdapter(workerUrl: string | URL): WhisperAdapter { total: payload.total, }) } - }) + }, options?.signal) w.postMessage({ type: 'load-model', requestId, modelId: MODEL_NAMES.WHISPER, device: requestedDevice }) @@ -293,19 +336,31 @@ export function createWhisperAdapter(workerUrl: string | URL): WhisperAdapter { state = 'ready' updateInferenceStatus(MODEL_NAMES.WHISPER, { state: 'ready', device: actualDevice }) onSuccess() - }) + }, { signal: options?.signal }) }) } - async function transcribe(input: WhisperTranscribeInput): Promise { + async function transcribe( + input: WhisperTranscribeInput, + options?: { signal?: AbortSignal }, + ): Promise { + throwIfAborted(options?.signal) return defaultPerfTracer.withMeasure('inference', 'whisper-transcribe', () => operationMutex.runExclusive(async () => { + throwIfAborted(options?.signal) if (!worker || state !== 'ready') throw new Error('Model not loaded. Call load() first.') state = 'transcribing' const requestId = createRequestId() - const resultPromise = waitForMessage(worker, requestId, 'inference-result', TRANSCRIBE_TIMEOUT) + const resultPromise = waitForMessage( + worker, + requestId, + 'inference-result', + TRANSCRIBE_TIMEOUT, + undefined, + options?.signal, + ) worker.postMessage({ type: 'run-inference', diff --git a/packages/stage-ui/src/libs/inference/load-queue.test.ts b/packages/stage-ui/src/libs/inference/load-queue.test.ts index 8120792b1..a322afedd 100644 --- a/packages/stage-ui/src/libs/inference/load-queue.test.ts +++ b/packages/stage-ui/src/libs/inference/load-queue.test.ts @@ -91,4 +91,107 @@ describe('loadQueue', () => { await p expect(queue.active).toBeNull() }) + + describe('cancellation', () => { + it('should reject immediately if signal is already aborted', async () => { + const queue = createLoadQueue() + const controller = new AbortController() + controller.abort() + + let loaderCalled = false + const promise = queue.enqueue( + 'already-aborted', + 1, + async () => { + loaderCalled = true + return 'x' + }, + { signal: controller.signal }, + ) + + await expect(promise).rejects.toMatchObject({ name: 'AbortError' }) + expect(loaderCalled).toBe(false) + }) + + it('should remove a pending entry from the queue when its signal aborts', async () => { + const queue = createLoadQueue() + + // Hold the queue with a slow loader + let releaseHold!: () => void + const hold = queue.enqueue('hold', 10, () => new Promise(r => releaseHold = r)) + + const controller = new AbortController() + let loaderCalled = false + const pending = queue.enqueue( + 'pending', + 1, + async () => { + loaderCalled = true + }, + { signal: controller.signal }, + ) + + expect(queue.pending).toContain('pending') + + controller.abort(new Error('cancelled by test')) + await expect(pending).rejects.toThrow('cancelled by test') + expect(queue.pending).not.toContain('pending') + expect(loaderCalled).toBe(false) + + releaseHold() + await hold + }) + + it('should not interrupt an active loader — that is the loader\'s responsibility', async () => { + const queue = createLoadQueue() + const controller = new AbortController() + + // Loader that honors its own signal + const activePromise = queue.enqueue( + 'active', + 1, + async () => { + await new Promise((resolve, reject) => { + controller.signal.addEventListener('abort', () => { + reject(new Error('loader aborted')) + }) + }) + }, + { signal: controller.signal }, + ) + + // Give the loader a tick to start + await new Promise(r => setTimeout(r, 5)) + expect(queue.active).toBe('active') + + controller.abort() + await expect(activePromise).rejects.toThrow('loader aborted') + }) + + it('should recover after a cancelled entry and continue processing subsequent items', async () => { + const queue = createLoadQueue() + + // Hold the queue + let releaseHold!: () => void + const hold = queue.enqueue('hold', 10, () => new Promise(r => releaseHold = r)) + + const controller = new AbortController() + const cancelled = queue.enqueue( + 'cancelled', + 5, + async () => 'should-not-run', + { signal: controller.signal }, + ) + + const later = queue.enqueue('later', 5, async () => 'later-result') + + controller.abort() + await expect(cancelled).rejects.toMatchObject({ name: 'AbortError' }) + + releaseHold() + await hold + + expect(await later).toBe('later-result') + }) + }) }) diff --git a/packages/stage-ui/src/libs/inference/load-queue.ts b/packages/stage-ui/src/libs/inference/load-queue.ts index c17a1435e..d5c98d6c1 100644 --- a/packages/stage-ui/src/libs/inference/load-queue.ts +++ b/packages/stage-ui/src/libs/inference/load-queue.ts @@ -6,8 +6,17 @@ * are dequeued first. * * Default priorities: TTS = 10, ASR = 5, BackgroundRemoval = 1. + * + * Cancellation: pass an `AbortSignal` in `enqueueOptions` to `enqueue()`. + * When aborted, the entry is removed from the pending queue (if not yet + * active) and its promise is rejected with `InferenceAbortError`. If the + * entry is already running, the loader itself is responsible for honoring + * the same signal and rejecting accordingly — the queue cannot interrupt + * an in-flight async loader. */ +import { InferenceAbortError } from './protocol' + // --------------------------------------------------------------------------- // Types // --------------------------------------------------------------------------- @@ -18,6 +27,13 @@ interface QueueEntry { loader: () => Promise resolve: (value: T) => void reject: (error: unknown) => void + signal?: AbortSignal + abortHandler?: () => void +} + +export interface EnqueueOptions { + /** Abort the enqueued load. Rejects the returned promise with `InferenceAbortError`. */ + signal?: AbortSignal } export interface LoadQueue { @@ -26,7 +42,12 @@ export interface LoadQueue { * the loader completes. If another load is in progress, this * one waits in a priority queue. */ - enqueue: (modelId: string, priority: number, loader: () => Promise) => Promise + enqueue: ( + modelId: string, + priority: number, + loader: () => Promise, + options?: EnqueueOptions, + ) => Promise /** Model IDs waiting in the queue */ readonly pending: string[] @@ -44,6 +65,13 @@ export function createLoadQueue(): LoadQueue { let active: string | null = null let running = false + function detachAbortHandler(entry: QueueEntry): void { + if (entry.signal && entry.abortHandler) { + entry.signal.removeEventListener('abort', entry.abortHandler) + entry.abortHandler = undefined + } + } + async function processQueue(): Promise { if (running) return @@ -54,12 +82,24 @@ export function createLoadQueue(): LoadQueue { queue.sort((a, b) => b.priority - a.priority) const entry = queue.shift()! + // Skip already-aborted entries (the abort handler may have fired + // before this dequeue; it removes the entry from the array but we + // also guard here in case of races) + if (entry.signal?.aborted) { + detachAbortHandler(entry) + const reason = entry.signal.reason + entry.reject(reason instanceof Error ? reason : new InferenceAbortError()) + continue + } + active = entry.modelId try { const result = await entry.loader() + detachAbortHandler(entry) entry.resolve(result) } catch (error) { + detachAbortHandler(entry) entry.reject(error) } } @@ -72,9 +112,39 @@ export function createLoadQueue(): LoadQueue { modelId: string, priority: number, loader: () => Promise, + options?: EnqueueOptions, ): Promise { return new Promise((resolve, reject) => { - queue.push({ modelId, priority, loader, resolve, reject }) + const entry: QueueEntry = { + modelId, + priority, + loader, + resolve, + reject, + signal: options?.signal, + } + + if (options?.signal) { + if (options.signal.aborted) { + const reason = options.signal.reason + reject(reason instanceof Error ? reason : new InferenceAbortError()) + return + } + entry.abortHandler = () => { + // Remove from pending queue if still there. If the entry has + // already been dequeued (active load), the loader's own abort + // propagation will handle rejection. + const idx = queue.indexOf(entry) + if (idx >= 0) { + queue.splice(idx, 1) + const reason = options.signal!.reason + reject(reason instanceof Error ? reason : new InferenceAbortError()) + } + } + options.signal.addEventListener('abort', entry.abortHandler) + } + + queue.push(entry) processQueue() }) } diff --git a/packages/stage-ui/src/libs/inference/protocol.test.ts b/packages/stage-ui/src/libs/inference/protocol.test.ts index f359afb75..44d353a77 100644 --- a/packages/stage-ui/src/libs/inference/protocol.test.ts +++ b/packages/stage-ui/src/libs/inference/protocol.test.ts @@ -1,6 +1,6 @@ import { describe, expect, it } from 'vitest' -import { classifyDeviceLossReason, classifyError, isRecoverable } from './protocol' +import { classifyDeviceLossReason, classifyError, InferenceAbortError, isRecoverable, throwIfAborted } from './protocol' describe('classifyError', () => { it('should classify OOM errors', () => { @@ -79,6 +79,74 @@ describe('isRecoverable', () => { it('should mark UNKNOWN as not recoverable', () => { expect(isRecoverable('UNKNOWN')).toBe(false) }) + + it('should mark CANCELLED as not recoverable', () => { + expect(isRecoverable('CANCELLED')).toBe(false) + }) +}) + +describe('inferenceAbortError', () => { + it('should have name "AbortError" for DOM compatibility', () => { + const err = new InferenceAbortError() + expect(err.name).toBe('AbortError') + }) + + it('should carry code "CANCELLED"', () => { + const err = new InferenceAbortError() + expect(err.code).toBe('CANCELLED') + }) + + it('should accept a custom message', () => { + const err = new InferenceAbortError('user cancelled') + expect(err.message).toBe('user cancelled') + }) + + it('should use a default message when none provided', () => { + const err = new InferenceAbortError() + expect(err.message).toBe('The operation was aborted') + }) + + it('should be instanceof Error', () => { + expect(new InferenceAbortError()).toBeInstanceOf(Error) + }) +}) + +describe('throwIfAborted', () => { + it('should be a no-op when signal is undefined', () => { + expect(() => throwIfAborted(undefined)).not.toThrow() + }) + + it('should be a no-op when signal is not aborted', () => { + const controller = new AbortController() + expect(() => throwIfAborted(controller.signal)).not.toThrow() + }) + + it('should throw when signal is already aborted', () => { + const controller = new AbortController() + controller.abort() + expect(() => throwIfAborted(controller.signal)).toThrow() + }) + + it('should throw the signal\'s reason if it is an Error', () => { + const controller = new AbortController() + const reason = new Error('custom reason') + controller.abort(reason) + expect(() => throwIfAborted(controller.signal)).toThrow(reason) + }) + + it('should throw InferenceAbortError when reason is a string', () => { + const controller = new AbortController() + controller.abort('cancelled by string') + try { + throwIfAborted(controller.signal) + } + catch (err) { + expect((err as Error).name).toBe('AbortError') + expect((err as Error).message).toBe('cancelled by string') + return + } + throw new Error('should have thrown') + }) }) describe('classifyDeviceLossReason', () => { diff --git a/packages/stage-ui/src/libs/inference/protocol.ts b/packages/stage-ui/src/libs/inference/protocol.ts index f3aa01676..75795868c 100644 --- a/packages/stage-ui/src/libs/inference/protocol.ts +++ b/packages/stage-ui/src/libs/inference/protocol.ts @@ -53,6 +53,7 @@ export type InferenceErrorCode | 'DEVICE_LOST' | 'LOAD_FAILED' | 'INFERENCE_FAILED' + | 'CANCELLED' | 'UNKNOWN' export interface ErrorPayload { @@ -87,10 +88,31 @@ export interface UnloadModelRequest { requestId: string } +/** + * Cancel an in-flight or queued request. The worker should stop any + * ongoing work tied to `targetRequestId` and must NOT send a normal + * `model-ready` / `inference-result` response for that request; instead + * it should send an `ErrorResponse` with code `'CANCELLED'` so the + * adapter can reject the caller's promise deterministically. + * + * NOTE: Cancellation is best-effort. We cannot interrupt a synchronous + * transformers.js / ONNX Runtime call that is already executing on the + * worker thread. What the cancel signal does guarantee is that the + * adapter stops waiting and the worker discards the result when it + * eventually arrives. + */ +export interface CancelRequest { + type: 'cancel' + requestId: string + /** The requestId of the operation to cancel */ + targetRequestId: string +} + export type WorkerInboundMessage = | LoadModelRequest | RunInferenceRequest | UnloadModelRequest + | CancelRequest // --------------------------------------------------------------------------- // Worker → Main responses @@ -228,3 +250,27 @@ export function classifyDeviceLossReason(error: unknown): DeviceLossReason { export function isRecoverable(code: InferenceErrorCode): boolean { return code === 'TIMEOUT' || code === 'DEVICE_LOST' } + +/** + * Canonical error thrown by inference adapters when an operation is + * cancelled via AbortSignal. Matches the DOM convention of `name === 'AbortError'` + * so existing `if (err.name === 'AbortError')` checks work unchanged. + */ +export class InferenceAbortError extends Error { + override readonly name = 'AbortError' + readonly code = 'CANCELLED' as const + + constructor(message = 'The operation was aborted') { + super(message) + } +} + +/** Throw `InferenceAbortError` if the signal is already aborted. */ +export function throwIfAborted(signal: AbortSignal | undefined): void { + if (signal?.aborted) { + const reason = signal.reason + if (reason instanceof Error) + throw reason + throw new InferenceAbortError(typeof reason === 'string' ? reason : undefined) + } +} diff --git a/packages/stage-ui/src/libs/workers/worker.ts b/packages/stage-ui/src/libs/workers/worker.ts index 0ee1720c5..bc6594a97 100644 --- a/packages/stage-ui/src/libs/workers/worker.ts +++ b/packages/stage-ui/src/libs/workers/worker.ts @@ -174,6 +174,42 @@ async function base64ToFeatures(base64Audio: string): Promise { // Helpers // --------------------------------------------------------------------------- +/** + * RequestIds the main thread has asked us to cancel. When an in-flight + * operation resolves, we check this set before posting the result; if + * the id is present, we send a `CANCELLED` error instead so the adapter + * rejects the caller's promise deterministically. + * + * We cannot synchronously interrupt a transformers.js call already running + * on this thread (no abort primitive is exposed) — cancellation here is + * about not leaking the stale result, not about stopping GPU work. + */ +const cancelledRequestIds = new Set() + +function markCancelled(targetRequestId: string): void { + cancelledRequestIds.add(targetRequestId) + // Emit the error now so the adapter can resolve immediately even if + // the inference keeps running in the background. + const msg: ErrorResponse = { + type: 'error', + requestId: targetRequestId, + payload: { + code: 'CANCELLED', + message: 'Operation cancelled by caller', + recoverable: false, + }, + } + globalThis.postMessage(msg) +} + +function isCancelled(requestId: string): boolean { + return cancelledRequestIds.has(requestId) +} + +function clearCancelled(requestId: string): void { + cancelledRequestIds.delete(requestId) +} + function sendProgress(requestId: string, phase: 'download' | 'compile' | 'warmup' | 'inference', percent: number, message?: string, extra?: Record): void { const msg: ProgressResponse = { type: 'progress', @@ -243,16 +279,25 @@ async function loadModel(request: LoadModelRequest): Promise { max_new_tokens: 1, } as Record) - const ready: ModelReadyResponse = { - type: 'model-ready', - requestId, - modelId: MODEL_NAMES.WHISPER, - device: resolvedDevice, + if (isCancelled(requestId)) { + // Adapter already received a CANCELLED error; drop the stale result. + clearCancelled(requestId) + } + else { + const ready: ModelReadyResponse = { + type: 'model-ready', + requestId, + modelId: MODEL_NAMES.WHISPER, + device: resolvedDevice, + } + globalThis.postMessage(ready) } - globalThis.postMessage(ready) } catch (error) { - sendError(requestId, error, 'load') + if (isCancelled(requestId)) + clearCancelled(requestId) + else + sendError(requestId, error, 'load') } finally { currentLoadRequestId = null @@ -311,15 +356,23 @@ async function runInference(request: RunInferenceRequest): Promise const outputText = tokenizer.batch_decode(outputs as Tensor, { skip_special_tokens: true }) - const result: InferenceResultResponse = { - type: 'inference-result', - requestId, - output: { text: outputText }, + if (isCancelled(requestId)) { + clearCancelled(requestId) + } + else { + const result: InferenceResultResponse = { + type: 'inference-result', + requestId, + output: { text: outputText }, + } + globalThis.postMessage(result) } - globalThis.postMessage(result) } catch (error) { - sendError(requestId, error, 'inference') + if (isCancelled(requestId)) + clearCancelled(requestId) + else + sendError(requestId, error, 'inference') } finally { processing = false @@ -344,5 +397,8 @@ globalThis.addEventListener('message', async (event: MessageEvent() + +function markCancelled(targetRequestId: string): void { + cancelledRequestIds.add(targetRequestId) + const msg: ErrorResponse = { + type: 'error', + requestId: targetRequestId, + payload: { + code: 'CANCELLED', + message: 'Operation cancelled by caller', + recoverable: false, + }, + } + globalThis.postMessage(msg) +} + +function isCancelled(requestId: string): boolean { + return cancelledRequestIds.has(requestId) +} + +function clearCancelled(requestId: string): void { + cancelledRequestIds.delete(requestId) +} + /** * Detect whether WebGPU is available inside the worker. */ @@ -97,6 +122,10 @@ async function loadModel(request: LoadModelRequest): Promise { try { if (model && processor) { + if (isCancelled(requestId)) { + clearCancelled(requestId) + return + } const ready: ModelReadyResponse = { type: 'model-ready', requestId, @@ -129,6 +158,11 @@ async function loadModel(request: LoadModelRequest): Promise { processor = await AutoProcessor.from_pretrained(MODEL_ID, {}) + if (isCancelled(requestId)) { + clearCancelled(requestId) + return + } + const ready: ModelReadyResponse = { type: 'model-ready', requestId, @@ -138,7 +172,10 @@ async function loadModel(request: LoadModelRequest): Promise { globalThis.postMessage(ready) } catch (error) { - sendError(requestId, error, 'load') + if (isCancelled(requestId)) + clearCancelled(requestId) + else + sendError(requestId, error, 'load') } } @@ -169,6 +206,11 @@ async function runInference(request: RunInferenceRequest output[0].mul(255).to('uint8'), ).resize(width, height) + if (isCancelled(requestId)) { + clearCancelled(requestId) + return + } + const maskData = new Uint8Array(mask.data.buffer) const result: InferenceResultResponse = { @@ -180,7 +222,10 @@ async function runInference(request: RunInferenceRequest ;(globalThis as any).postMessage(result, [maskData.buffer]) } catch (error) { - sendError(requestId, error, 'inference') + if (isCancelled(requestId)) + clearCancelled(requestId) + else + sendError(requestId, error, 'inference') } } @@ -203,5 +248,8 @@ globalThis.addEventListener('message', async (event: MessageEvent = { cpu: [], } +// NOTICE: Cancellation tracking — see Whisper worker for the full rationale. +// We cannot interrupt a transformers.js call synchronously; this set lets us +// drop stale results when they arrive. +const cancelledRequestIds = new Set() + +function markCancelled(targetRequestId: string): void { + cancelledRequestIds.add(targetRequestId) + const msg: ErrorResponse = { + type: 'error', + requestId: targetRequestId, + payload: { + code: 'CANCELLED', + message: 'Operation cancelled by caller', + recoverable: false, + }, + } + globalThis.postMessage(msg) +} + +function isCancelled(requestId: string): boolean { + return cancelledRequestIds.has(requestId) +} + +function clearCancelled(requestId: string): void { + cancelledRequestIds.delete(requestId) +} + function sendError(requestId: string, error: unknown, phase?: 'load' | 'inference'): void { const message = error instanceof Error ? error.message : String(error) const code = classifyError(error, phase) @@ -96,6 +123,10 @@ async function loadModel(request: LoadModelRequest): Promise { try { // Check if we already have the correct model loaded if (ttsModel && currentQuantization === quantization && currentDevice === device) { + if (isCancelled(requestId)) { + clearCancelled(requestId) + return + } const ready: ModelReadyResponse = { type: 'model-ready', requestId, @@ -151,6 +182,10 @@ async function loadModel(request: LoadModelRequest): Promise { currentQuantization = quantization currentDevice = attempt.device + if (isCancelled(requestId)) { + clearCancelled(requestId) + return + } const ready: ModelReadyResponse = { type: 'model-ready', requestId, @@ -175,10 +210,16 @@ async function loadModel(request: LoadModelRequest): Promise { } // All attempts exhausted - sendError(requestId, lastError ?? new Error('All dtype/device combinations failed'), 'load') + if (isCancelled(requestId)) + clearCancelled(requestId) + else + sendError(requestId, lastError ?? new Error('All dtype/device combinations failed'), 'load') } catch (error) { - sendError(requestId, error, 'load') + if (isCancelled(requestId)) + clearCancelled(requestId) + else + sendError(requestId, error, 'load') } } @@ -190,6 +231,11 @@ async function runInference(request: RunInferenceRequest): if (!ttsModel) throw new Error('Model not loaded. Send load-model first.') + if (isCancelled(requestId)) { + clearCancelled(requestId) + return + } + const result: InferenceResultResponse = { type: 'inference-result', requestId, @@ -206,6 +252,11 @@ async function runInference(request: RunInferenceRequest): const { text, voice } = input const audioResult = await ttsModel.generate(text, { voice }) + if (isCancelled(requestId)) { + clearCancelled(requestId) + return + } + // Transfer raw PCM Float32Array directly — avoids WAV blob encode/decode overhead. const samples = audioResult.audio const result: InferenceResultResponse = { @@ -216,7 +267,10 @@ async function runInference(request: RunInferenceRequest): ;(globalThis as any).postMessage(result, [samples.buffer]) } catch (error) { - sendError(requestId, error, 'inference') + if (isCancelled(requestId)) + clearCancelled(requestId) + else + sendError(requestId, error, 'inference') } } @@ -240,6 +294,9 @@ globalThis.addEventListener('message', async (event: MessageEvent