From 5c6cbfc84e1eb1db6f739b4cc31c819deb2c4c2b Mon Sep 17 00:00:00 2001 From: Jinjing <6427696+AmethystLiang@users.noreply.github.com> Date: Thu, 9 Jul 2026 22:13:11 -0700 Subject: [PATCH] fix(speech): harden offline STT chunk decode warm-reuse lifecycle (#8007) Keep offline multi-chunk decoding safe for warm workers: refresh streams after every decode attempt, continue sibling chunks after a failure, reset session state on stop without dropping the stopped signal, and drop in-flight feeds while stop is in progress so residual audio cannot contaminate the next dictation. Extract model path/hotwords config to stay under max-lines. --- src/main/speech/stt-service.test.ts | 20 +++ src/main/speech/stt-service.ts | 170 +++++++++++---------- src/main/speech/stt-worker-model-config.ts | 80 ++++++++++ src/main/speech/stt-worker.ts | 129 ++++++---------- 4 files changed, 239 insertions(+), 160 deletions(-) create mode 100644 src/main/speech/stt-worker-model-config.ts diff --git a/src/main/speech/stt-service.test.ts b/src/main/speech/stt-service.test.ts index 5138831d8ae..31736082ab8 100644 --- a/src/main/speech/stt-service.test.ts +++ b/src/main/speech/stt-service.test.ts @@ -239,6 +239,26 @@ describe('SttService', () => { expect(worker!.messages.filter((message) => message.type === 'feed')).toHaveLength(0) }) + it('drops in-flight audio while stop is waiting for the worker', async () => { + const service = new SttService({ + getModelState: vi.fn().mockResolvedValue({ id: 'model-a', status: 'ready' }), + getModelDir: vi.fn().mockReturnValue('/tmp/model-a') + } as never) + + await service.startDictation('model-a', vi.fn(), undefined, 'desktop:1') + const worker = getLastWorker() + expect(worker).toBeDefined() + worker!.emitStoppedOnStop = false + + const stopPromise = service.stopDictation('desktop:1') + await Promise.resolve() + service.feedAudio(new Float32Array([1]), 16000, 'desktop:1') + expect(worker!.messages.filter((message) => message.type === 'feed')).toHaveLength(0) + + worker!.emit('message', { type: 'stopped' }) + await stopPromise + }) + it('rejects deletion prep while the target model is starting', async () => { let resolveModelState: (state: { id: string; status: string }) => void = () => {} const modelStatePromise = new Promise<{ id: string; status: string }>((resolve) => { diff --git a/src/main/speech/stt-service.ts b/src/main/speech/stt-service.ts index d9d2129dd71..6974dd5945e 100644 --- a/src/main/speech/stt-service.ts +++ b/src/main/speech/stt-service.ts @@ -35,6 +35,10 @@ export class SttService { private canceledOwners = new Set() private eventSink: SttEventSink | null = null private idleTeardownTimer: NodeJS.Timeout | null = null + // Why: stop resolves only after the worker flushes; in-flight feedAudio IPC + // must not enqueue samples after that flush or they stick on the warm worker + // and contaminate the next dictation session. + private stopping = false // Why: warm workers intentionally keep lifecycle listeners while reusable; // stale workers must not retain this service after error, exit, or teardown. private cleanupWorkerLifecycleListeners: (() => void) | null = null @@ -259,6 +263,9 @@ export class SttService { } feedAudio(samples: Float32Array, sampleRate: number, owner = 'desktop'): void { + if (this.stopping) { + return + } const currentOwner = this.activeOwner ?? this.startingOwner if (!currentOwner) { return @@ -288,92 +295,97 @@ export class SttService { throw new Error('dictation_owner_mismatch') } - if (this.cloudSession) { - const session = this.cloudSession - this.cloudSession = null - try { - const text = await session.finish() - if (text) { - this.eventSink?.({ type: 'final', text }) + this.stopping = true + try { + if (this.cloudSession) { + const session = this.cloudSession + this.cloudSession = null + try { + const text = await session.finish() + if (text) { + this.eventSink?.({ type: 'final', text }) + } + } catch (error) { + this.eventSink?.({ + type: 'error', + error: error instanceof Error ? error.message : String(error) + }) + } finally { + this.eventSink?.({ type: 'stopped' }) + this.activeModelId = null + this.activeHotwordsFilePath = undefined + this.activeOwner = null + this.eventSink = null } - } catch (error) { - this.eventSink?.({ - type: 'error', - error: error instanceof Error ? error.message : String(error) - }) - } finally { - this.eventSink?.({ type: 'stopped' }) - this.activeModelId = null - this.activeHotwordsFilePath = undefined - this.activeOwner = null - this.eventSink = null - } - return - } - - const worker = this.worker - if (!worker) { - return - } - worker.postMessage({ type: 'stop' }) - - let forcedTeardown = false - await new Promise((resolve) => { - let settled = false - let timeout: ReturnType | null = null - - const cleanup = (): void => { - if (timeout) { - clearTimeout(timeout) - timeout = null - } - worker.off('message', onStopped) + return } - const finish = (): void => { - if (settled) { - return - } - settled = true - cleanup() - resolve() + const worker = this.worker + if (!worker) { + return } + worker.postMessage({ type: 'stop' }) - const onStopped = (msg: { type: string; text?: string; error?: string }) => { - if (msg.type === 'stopped') { - finish() + let forcedTeardown = false + await new Promise((resolve) => { + let settled = false + let timeout: ReturnType | null = null + + const cleanup = (): void => { + if (timeout) { + clearTimeout(timeout) + timeout = null + } + worker.off('message', onStopped) + } + + const finish = (): void => { + if (settled) { + return + } + settled = true + cleanup() + resolve() + } + + const onStopped = (msg: { type: string; text?: string; error?: string }) => { + if (msg.type === 'stopped') { + finish() + } + } + + timeout = setTimeout(() => { + if (settled) { + return + } + settled = true + forcedTeardown = true + cleanup() + // Why: a worker that cannot finish dictation is no longer reusable; do + // not keep it in the warm-worker slot or retain its message listeners. + this.cleanupActiveWorkerLifecycleListeners() + worker.removeAllListeners() + void worker.terminate().finally(resolve) + }, STOP_DICTATION_TIMEOUT_MS) + + worker.on('message', onStopped) + }) + + if (this.worker === worker) { + if (forcedTeardown) { + this.worker = null + this.activeModelId = null + this.activeHotwordsFilePath = undefined + this.activeOwner = null + this.eventSink = null + } else { + this.activeOwner = null + this.eventSink = null + this.scheduleIdleTeardown() } } - - timeout = setTimeout(() => { - if (settled) { - return - } - settled = true - forcedTeardown = true - cleanup() - // Why: a worker that cannot finish dictation is no longer reusable; do - // not keep it in the warm-worker slot or retain its message listeners. - this.cleanupActiveWorkerLifecycleListeners() - worker.removeAllListeners() - void worker.terminate().finally(resolve) - }, STOP_DICTATION_TIMEOUT_MS) - - worker.on('message', onStopped) - }) - - if (this.worker === worker) { - if (forcedTeardown) { - this.worker = null - this.activeModelId = null - this.activeHotwordsFilePath = undefined - this.activeOwner = null - this.eventSink = null - } else { - this.activeOwner = null - this.eventSink = null - this.scheduleIdleTeardown() - } + } finally { + this.stopping = false } } diff --git a/src/main/speech/stt-worker-model-config.ts b/src/main/speech/stt-worker-model-config.ts new file mode 100644 index 00000000000..6a2a0e89dd9 --- /dev/null +++ b/src/main/speech/stt-worker-model-config.ts @@ -0,0 +1,80 @@ +import { readdirSync } from 'node:fs' + +// Why: different models name their ONNX files differently (e.g. +// encoder.int8.onnx vs tiny-encoder.onnx vs encoder-epoch-99-avg-1.onnx). +// We resolve the actual path from the manifest's files list by searching +// for the role name anywhere in the filename. +export function resolveFile( + files: string[], + role: string, + modelDir: string, + ext = '.onnx' +): string { + const match = files.find((f) => f.includes(role) && f.endsWith(ext)) + if (!match) { + throw new Error(`No *${role}*${ext} found in model files: ${files.join(', ')}`) + } + return `${modelDir}/${match}` +} + +export function resolveTokens(files: string[], modelDir: string): string { + const match = files.find((f) => f.endsWith('tokens.txt')) + if (!match) { + throw new Error(`No *tokens.txt found in model files: ${files.join(', ')}`) + } + return `${modelDir}/${match}` +} + +// Why: BPE models need a vocab file for hotwords token matching. The file +// ships in the model archive but isn't listed in the manifest. We discover +// it at runtime to avoid breaking existing downloads. +function discoverBpeVocab(modelDir: string): string | undefined { + try { + const entries = readdirSync(modelDir) + const vocabFile = entries.find((f) => f.endsWith('.vocab')) + return vocabFile ? `${modelDir}/${vocabFile}` : undefined + } catch { + return undefined + } +} + +export type HotwordsConfig = { + decodingMethod: string + hotwordsFile?: string + hotwordsScore?: number + modelingUnit?: string + bpeVocab?: string +} + +export function buildHotwordsConfig(opts: { + modelDir: string + modelType: string + hotwordsFilePath?: string + modelingUnit?: string +}): HotwordsConfig { + if (opts.modelType !== 'transducer' || !opts.hotwordsFilePath) { + return { decodingMethod: 'greedy_search' } + } + + const unit = opts.modelingUnit + if (unit?.includes('bpe')) { + const bpeVocab = discoverBpeVocab(opts.modelDir) + if (!bpeVocab) { + return { decodingMethod: 'greedy_search' } + } + return { + decodingMethod: 'modified_beam_search', + hotwordsFile: opts.hotwordsFilePath, + hotwordsScore: 1.5, + modelingUnit: unit, + bpeVocab + } + } + + return { + decodingMethod: 'modified_beam_search', + hotwordsFile: opts.hotwordsFilePath, + hotwordsScore: 1.5, + modelingUnit: unit + } +} diff --git a/src/main/speech/stt-worker.ts b/src/main/speech/stt-worker.ts index 2f4f9b025d4..1e6cd8bb75c 100644 --- a/src/main/speech/stt-worker.ts +++ b/src/main/speech/stt-worker.ts @@ -1,8 +1,8 @@ /* oxlint-disable typescript-eslint/no-explicit-any -- sherpa-onnx native addon has no type definitions */ import { parentPort, workerData } from 'node:worker_threads' -import { readdirSync } from 'node:fs' import { resampleToRate } from './stt-audio-resample' import { OfflineAudioChunker } from './stt-offline-audio-chunker' +import { buildHotwordsConfig, resolveFile, resolveTokens } from './stt-worker-model-config' type WorkerMessage = | { @@ -39,73 +39,6 @@ function loadSherpa(): any { return require(modulePath) } -// Why: different models name their ONNX files differently (e.g. -// encoder.int8.onnx vs tiny-encoder.onnx vs encoder-epoch-99-avg-1.onnx). -// We resolve the actual path from the manifest's files list by searching -// for the role name anywhere in the filename. -function resolveFile(files: string[], role: string, modelDir: string, ext = '.onnx'): string { - const match = files.find((f) => f.includes(role) && f.endsWith(ext)) - if (!match) { - throw new Error(`No *${role}*${ext} found in model files: ${files.join(', ')}`) - } - return `${modelDir}/${match}` -} - -function resolveTokens(files: string[], modelDir: string): string { - const match = files.find((f) => f.endsWith('tokens.txt')) - if (!match) { - throw new Error(`No *tokens.txt found in model files: ${files.join(', ')}`) - } - return `${modelDir}/${match}` -} - -// Why: BPE models need a vocab file for hotwords token matching. The file -// ships in the model archive but isn't listed in the manifest. We discover -// it at runtime to avoid breaking existing downloads. -function discoverBpeVocab(modelDir: string): string | undefined { - try { - const entries = readdirSync(modelDir) - const vocabFile = entries.find((f) => f.endsWith('.vocab')) - return vocabFile ? `${modelDir}/${vocabFile}` : undefined - } catch { - return undefined - } -} - -function buildHotwordsConfig(msg: Extract): { - decodingMethod: string - hotwordsFile?: string - hotwordsScore?: number - modelingUnit?: string - bpeVocab?: string -} { - if (msg.modelType !== 'transducer' || !msg.hotwordsFilePath) { - return { decodingMethod: 'greedy_search' } - } - - const unit = msg.modelingUnit - if (unit?.includes('bpe')) { - const bpeVocab = discoverBpeVocab(msg.modelDir) - if (!bpeVocab) { - return { decodingMethod: 'greedy_search' } - } - return { - decodingMethod: 'modified_beam_search', - hotwordsFile: msg.hotwordsFilePath, - hotwordsScore: 1.5, - modelingUnit: unit, - bpeVocab - } - } - - return { - decodingMethod: 'modified_beam_search', - hotwordsFile: msg.hotwordsFilePath, - hotwordsScore: 1.5, - modelingUnit: unit - } -} - function handleInit(msg: Extract): void { try { sherpa = loadSherpa() @@ -204,15 +137,33 @@ function handleInit(msg: Extract): void { } } -// Why: an offline stream is single-use — decode one bounded chunk, then -// recreate the stream so the recognizer is ready for the next chunk. +// Why: an offline stream is single-use — always mint a fresh stream after an +// attempt so a failed decode cannot leave a spent stream on a warm worker. function decodeOfflineChunk(samples: Float32Array): string { - sherpa.acceptWaveformOffline(stream, { sampleRate: offlineSampleRate, samples }) - sherpa.decodeOfflineStream(recognizer, stream) - const resultJson = sherpa.getOfflineStreamResultAsJson(stream) - stream = sherpa.createOfflineStream(recognizer) - const result = JSON.parse(resultJson) - return result?.text?.trim() ?? '' + try { + sherpa.acceptWaveformOffline(stream, { sampleRate: offlineSampleRate, samples }) + sherpa.decodeOfflineStream(recognizer, stream) + const resultJson = sherpa.getOfflineStreamResultAsJson(stream) + const result = JSON.parse(resultJson) + return result?.text?.trim() ?? '' + } finally { + if (sherpa && recognizer) { + stream = sherpa.createOfflineStream(recognizer) + } + } +} + +// Why: warm-worker reuse must not see residual audio or a spent offline stream. +// Recovery failures are swallowed so they cannot skip the stopped lifecycle signal. +function resetOfflineSessionState(): void { + try { + offlineChunker = new OfflineAudioChunker(offlineSampleRate) + if (sherpa && recognizer) { + stream = sherpa.createOfflineStream(recognizer) + } + } catch { + // Non-fatal: prefer posting stopped over stranding stopDictation for 60s. + } } function handleFeed(msg: Extract): void { @@ -253,12 +204,23 @@ function handleFeed(msg: Extract): void { // the whole app (#7925). Decode bounded chunks as they fill instead; // each consumer already appends multiple 'final' segments per session. const readyChunks = offlineChunker?.push(new Float32Array(samples)) ?? [] + // Why: keep decoding later ready windows after one failure — push() has + // already removed them from the chunker, and decodeOfflineChunk refreshes + // the stream in finally so a spent stream cannot poison the next attempt. + let firstError: unknown = null for (const chunk of readyChunks) { - const text = decodeOfflineChunk(chunk) - if (text) { - parentPort?.postMessage({ type: 'final', text }) + try { + const text = decodeOfflineChunk(chunk) + if (text) { + parentPort?.postMessage({ type: 'final', text }) + } + } catch (err) { + firstError ??= err } } + if (firstError) { + throw firstError + } } } catch (err) { parentPort?.postMessage({ type: 'error', error: String(err) }) @@ -294,12 +256,17 @@ function handleStop(): void { parentPort?.postMessage({ type: 'final', text }) } } + resetOfflineSessionState() } } catch (err) { + if (!isStreaming) { + resetOfflineSessionState() + } parentPort?.postMessage({ type: 'error', error: String(err) }) + } finally { + // Why: stopDictation waits on this signal; recovery must never prevent it. + parentPort?.postMessage({ type: 'stopped' }) } - - parentPort?.postMessage({ type: 'stopped' }) } function handleTeardown(): void {