From 4e42cc2a3f8992c8aa9a09d5091dc00ceac3660a Mon Sep 17 00:00:00 2001
From: Kelvin Amoaba <97001695+AmoabaKelvin@users.noreply.github.com>
Date: Mon, 9 Nov 2026 20:49:37 +0000
Subject: [PATCH] fix(speech): keep a slow GPU model load on the GPU while it
 still makes progress (#40472)

The GPU stall watch gave the whole model load one 25 s budget, so a cold disk
or busy machine restarted a healthy load on the CPU and could miss the 60 s
start limit. The speech process now reports each load step, and the watch
restarts on it; a step that hangs for 25 s still moves to the CPU.
---
 src/main/speech/stt-recognizer.ts             |  2 ++
 .../speech/stt-transcribe-cpp-recognizer.ts   |  5 ++++
 .../speech/stt-worker-cpu-fallback.test.ts    | 27 ++++++++++++++++++-
 src/main/speech/stt-worker-cpu-fallback.ts    |  9 +++++--
 4 files changed, 40 insertions(+), 3 deletions(-)

diff --git a/src/main/speech/stt-recognizer.ts b/src/main/speech/stt-recognizer.ts
index 2b91dd8d43f..1998151130c 100644
--- a/src/main/speech/stt-recognizer.ts
+++ b/src/main/speech/stt-recognizer.ts
@@ -15,6 +15,8 @@ export type SttRecognizerEvent =
   | { type: 'chunk'; streamId: string; text: string; samples: number }
   /** Pause-cut streams only: speech was heard that no chunk has committed yet. */
   | { type: 'speech-pending'; streamId: string; pending: boolean }
+  /** A load step finished, so a slow load is told apart from a hung GPU. */
+  | { type: 'loading' }
 
 export type SttRecognizerPost = (event: SttRecognizerEvent) => void
 
diff --git a/src/main/speech/stt-transcribe-cpp-recognizer.ts b/src/main/speech/stt-transcribe-cpp-recognizer.ts
index e8c62efda14..9d5b51f0852 100644
--- a/src/main/speech/stt-transcribe-cpp-recognizer.ts
+++ b/src/main/speech/stt-transcribe-cpp-recognizer.ts
@@ -91,7 +91,12 @@ export async function createTranscribeCppRecognizer(
   if (!gguf) {
     throw new Error(`No .gguf model file in: ${init.files.join(', ')}`)
   }
+  // Device setup (Metal shader library, Vulkan drivers) and reading the weights each report done,
+  // so a cold disk or a busy machine isn't taken for a GPU that hung.
+  await transcribe.initialize()
+  post({ type: 'loading' })
   const model = await transcribe.TranscribeModel.load(join(init.modelDir, gguf))
+  post({ type: 'loading' })
   // The child runs one call at a time, so every stream can share one session.
   const session = model.createSession({ nThreads: performanceCoreThreads() })
   const { languages } = model.capabilities
diff --git a/src/main/speech/stt-worker-cpu-fallback.test.ts b/src/main/speech/stt-worker-cpu-fallback.test.ts
index 582acc144d1..ffc7a6cb42c 100644
--- a/src/main/speech/stt-worker-cpu-fallback.test.ts
+++ b/src/main/speech/stt-worker-cpu-fallback.test.ts
@@ -1,5 +1,5 @@
 import { EventEmitter } from 'node:events'
-import { describe, expect, it } from 'vitest'
+import { describe, expect, it, vi } from 'vitest'
 import { spawnSttWorkerWithCpuFallback } from './stt-worker-cpu-fallback'
 
 class FakeChild extends EventEmitter {
@@ -115,6 +115,31 @@ describe('GPU speech process falling back to the CPU', () => {
     expect(marks()).toBe(0)
   })
 
+  it('keeps a slow GPU load that still finishes its steps, and moves a silent one to the CPU', () => {
+    vi.useFakeTimers()
+    try {
+      const { children, events } = setup()
+      const [gpu] = children
+      vi.advanceTimersByTime(20_000)
+      gpu.emit('message', { type: 'loading' })
+      vi.advanceTimersByTime(20_000)
+      gpu.emit('message', { type: 'loading' })
+      vi.advanceTimersByTime(20_000)
+      gpu.emit('message', { type: 'ready' })
+      expect(children).toHaveLength(1)
+      expect(events).toEqual([{ type: 'ready' }])
+
+      const silent = setup()
+      vi.advanceTimersByTime(20_000)
+      silent.children[0].emit('message', { type: 'loading' })
+      vi.advanceTimersByTime(25_000)
+      expect(silent.children).toHaveLength(2)
+      expect(silent.children[1].cpuOnly).toBe(true)
+    } finally {
+      vi.useRealTimers()
+    }
+  })
+
   it('fails the worker when the CPU replacement cannot load the model mid-dictation', () => {
     const { worker, children } = setup()
     const exits: number[] = []
diff --git a/src/main/speech/stt-worker-cpu-fallback.ts b/src/main/speech/stt-worker-cpu-fallback.ts
index 63bdeba05a4..d86c3d1b15b 100644
--- a/src/main/speech/stt-worker-cpu-fallback.ts
+++ b/src/main/speech/stt-worker-cpu-fallback.ts
@@ -17,7 +17,7 @@ type ChildMessage =
   | { type: 'chunk'; streamId: string; text: string; samples: number }
   | { type: 'speech-pending'; streamId: string; pending: boolean }
   | {
-      type: 'ready' | 'stopped' | 'error' | 'partial' | 'final'
+      type: 'ready' | 'loading' | 'stopped' | 'error' | 'partial' | 'final'
       streamId?: string
       text?: string
       error?: string
@@ -33,7 +33,8 @@ type StreamBacklog = {
   pending: 'stop' | 'discard' | null
 }
 
-// Well under the session's 60 s start and stop timeouts, so a CPU replay still fits; covers first-run shader compiles.
+// Per load step or decode. Well under the session's 60 s start and stop timeouts, so a CPU replay still fits;
+// covers first-run shader compiles.
 const GPU_STALL_TIMEOUT_MS = 25_000
 
 /**
@@ -132,6 +133,10 @@ class SttCpuFallbackWorker extends EventEmitter implements SttWorkerProcess {
       this.emit('exit', 1)
       return
     }
+    if (message.type === 'loading') {
+      this.watchForStall(true)
+      return
+    }
     if (message.type === 'chunk') {
       const { streamId, text, samples } = message
       const stream = this.streams.get(streamId)