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.
2026-10-06 08:02:28 +00:00 · 2026-07-09 22:13:11 -07:00
parent 693ac1fb25
commit 5c6cbfc84e
4 changed files with 239 additions and 160 deletions
@@ -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) => {
@@ -35,6 +35,10 @@ export class SttService {
  private canceledOwners = new Set<string>()
  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<void>((resolve) => {
-      let settled = false
-      let timeout: ReturnType<typeof setTimeout> | 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<void>((resolve) => {
+        let settled = false
+        let timeout: ReturnType<typeof setTimeout> | 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
    }
  }

@@ -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
+  }
+}
@@ -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<WorkerMessage, { type: 'init' }>): {
-  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<WorkerMessage, { type: 'init' }>): void {
  try {
    sherpa = loadSherpa()
@@ -204,15 +137,33 @@ function handleInit(msg: Extract<WorkerMessage, { type: 'init' }>): 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<WorkerMessage, { type: 'feed' }>): void {
@@ -253,12 +204,23 @@ function handleFeed(msg: Extract<WorkerMessage, { type: 'feed' }>): 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 {