fix(streaming): stop fallback no longer clobbers a slow engine finalize (tail-word loss)

OpenWhisp/Models/AppState.swift

@@ -3471,20 +3471,20 @@ class AppState: ObservableObject {
         // at the true end (insert / finishSessionUI).
         activeStreamingEngine.stop(cancel: false)
 
-        // Session-scoped fallback: a rapid follow-up session could otherwise be
-        // finalized early by THIS session's leftover timer.
-        //
-        // 2.0s (not 0.9s): AppleSpeechEngine's own 0.8s fallback guarantees a
-        // final for the Apple path, so this is a stuck-session guard (WhisperKit
-        // streaming teardown, engine deallocated) — with a wide margin so a busy
-        // main thread can't let this fire first and clobber the engine's genuine
-        // final (which arrives via an extra Task hop) with a stale partial.
+        // Session-scoped fallback (a rapid follow-up session must not be
+        // finalized by THIS session's leftover timer). Re-arming poll, not a
+        // one-shot: while the engine reports `isFinalizing` (Parakeet draining
+        // its decode backlog, WhisperKit's finalizeTail) keep waiting for the
+        // genuine final — a one-shot fired first on long dictations and its stale
+        // partial dropped the real final's tail words; see StreamingRoutePolicy.
         let sessionID = activeSessionID
+        let engine = activeStreamingEngine
         Task { @MainActor in
-            try? await Task.sleep(nanoseconds: 2_000_000_000)
-            if sessionID == self.activeSessionID && self.isAppleSpeechSession && self.isTranscribing && !self.appleDidCompleteFinal {
-                self.handleAppleSpeechFinal(self.streamingText, sessionID: sessionID)
-            }
+            await StreamingRoutePolicy.runStopFallback(
+                isSessionStillWaiting: { sessionID == self.activeSessionID && self.isAppleSpeechSession
+                    && self.isTranscribing && !self.appleDidCompleteFinal },
+                isEngineFinalizing: { engine.isFinalizing },
+                completeFallback: { self.handleAppleSpeechFinal(self.streamingText, sessionID: sessionID) })
         }
     }
 

OpenWhisp/Services/ParakeetStreamingEngine.swift

@@ -81,6 +81,19 @@ final class ParakeetStreamingEngine: NSObject, StreamingTranscriptionEngine {
 
     @MainActor private var lastPartial = ""
     @MainActor private var didStop = false
+    /// Non-cancel stops enqueued whose `runStop` hasn't completed yet — the
+    /// engine is still working toward a genuine final (feed drain + `finish()`
+    /// flush). A counter, not a bool: on a fast stop→start→stop, session N+1's
+    /// stop is already counted while session N's runStop is still queued, and
+    /// N's completion must not clear N+1's pending finalize.
+    @MainActor private var pendingFinalizeStops = 0
+
+    /// True while a non-cancel stop is still draining/flushing toward its final.
+    /// Read by AppState's stuck-session fallback; main-actor callers only (same
+    /// contract as `selectDevice`).
+    var isFinalizing: Bool {
+        MainActor.assumeIsolated { pendingFinalizeStops > 0 }
+    }
     /// True while a settle-check for a configuration-change notification is
     /// already scheduled — further notifications in that window coalesce into
     /// the pending check instead of stacking checks (a storm posts several).
@@ -462,6 +475,13 @@ final class ParakeetStreamingEngine: NSObject, StreamingTranscriptionEngine {
             // session's final with the successor's sessionID, straight through
             // AppState's session fence.
             let final = self.onFinal
+            // Report "finalizing" from ENQUEUE until runStop completes, so the
+            // stuck-session fallback (StreamingRoutePolicy.runStopFallback) keeps
+            // waiting for the genuine final instead of completing the session
+            // with the stale last partial while the feed drain / finish() flush
+            // runs — enqueue time matters because runStop can sit behind a slow
+            // predecessor in the chain.
+            if !cancel { self.pendingFinalizeStops += 1 }
             self.lifecycle.enqueue {
                 await self.runStop(cancel: cancel, final: final)
             }
@@ -481,6 +501,11 @@ final class ParakeetStreamingEngine: NSObject, StreamingTranscriptionEngine {
 
     @MainActor
     private func runStop(cancel: Bool, final deliverFinal: ((String) -> Void)?) async {
+        // Balance stop(cancel:false)'s increment on EVERY exit path (no-session
+        // early return, finish() error): the fallback poll must see finalizing
+        // end even when no final gets delivered — its grace poll then fires the
+        // stale-partial completion instead of wedging the session.
+        defer { if !cancel { pendingFinalizeStops = max(0, pendingFinalizeStops - 1) } }
         didStop = true
         // Supersede the session so late partial hops are dropped.
         generation += 1

OpenWhisp/Services/StreamingRoutePolicy.swift

@@ -105,6 +105,77 @@ public enum StreamingRoutePolicy {
         return pendingStop ? .beginListeningThenStop : .beginListening
     }
 
+    /// Stuck-session stop-fallback poll interval (seconds). 2.0 (not 0.9):
+    /// AppleSpeechEngine's own 0.8s fallback guarantees a final for the Apple
+    /// path, so the fallback is a stuck-session guard — with a wide margin so a
+    /// busy main thread can't let it fire first and clobber the engine's genuine
+    /// final (which arrives via an extra Task hop) with a stale partial.
+    public static let stopFallbackInterval: TimeInterval = 2
+    /// Hard cap (seconds) on how long the stop fallback keeps waiting for an
+    /// engine that reports `isFinalizing`. A finalize that outlasts this is
+    /// treated as hung and the session completes with the last partial — the
+    /// stuck-session guarantee the fallback exists for. Generous on purpose: a
+    /// real decode-backlog drain is far faster than realtime (Parakeet RTFx ~5),
+    /// so only a genuine hang ever reaches it.
+    public static let stopFallbackMaxWait: TimeInterval = 30
+
+    /// The stuck-session fallback armed by a streaming stop (AppState's
+    /// `stopAppleSpeech`): if the engine's genuine final never arrives, complete
+    /// the session with the last streamed partial rather than wedging at
+    /// "Finalizing...".
+    ///
+    /// Re-arming poll loop, NOT a one-shot deadline — the one-shot was the
+    /// tail-word-loss bug: a long dictation leaves Parakeet's `runStop` draining
+    /// the queued audio feed through the decoder (then `finish()`), which can
+    /// outlast any fixed window; the timer fired first with the STALE partial,
+    /// set the did-complete flag, and the genuine final (holding the tail words)
+    /// was dropped by the completion guard. While the engine reports
+    /// `isFinalizing`, the fallback re-arms instead of firing.
+    ///
+    /// One extra GRACE poll after finalizing ends: the engine clears its
+    /// finalizing state when its stop chain returns, but the final it delivered
+    /// still rides a main-actor Task hop into the completion handler — completing
+    /// in that gap would clobber it. By the next poll the hop has long landed and
+    /// `isSessionStillWaiting` reports the session done.
+    ///
+    /// - Parameters:
+    ///   - sleep: injectable wait (tests script it; production sleeps for real).
+    ///   - isSessionStillWaiting: the caller's session fence — still the same
+    ///     session, still finalizing, no final handled yet. `false` ends the loop
+    ///     silently (the timer's job is done or moot).
+    ///   - isEngineFinalizing: the active engine's `isFinalizing`.
+    ///   - completeFallback: complete the session with the stale partial. Called
+    ///     at most once, as the loop's last act.
+    @MainActor
+    public static func runStopFallback(
+        interval: TimeInterval = stopFallbackInterval,
+        maxWait: TimeInterval = stopFallbackMaxWait,
+        sleep: (TimeInterval) async -> Void = {
+            try? await Task.sleep(nanoseconds: UInt64($0 * 1_000_000_000))
+        },
+        isSessionStillWaiting: () -> Bool,
+        isEngineFinalizing: () -> Bool,
+        completeFallback: () -> Void
+    ) async {
+        var elapsed: TimeInterval = 0
+        var wasFinalizing = false
+        while true {
+            await sleep(interval)
+            elapsed += interval
+            guard isSessionStillWaiting() else { return }
+            let finalizing = isEngineFinalizing()
+            if (finalizing || wasFinalizing) && elapsed < maxWait {
+                // `wasFinalizing` is what grants the single grace poll: it holds
+                // fire for exactly one interval after finalizing ends, then (if
+                // the final still never landed) the fallback fires after all.
+                wasFinalizing = finalizing
+                continue
+            }
+            completeFallback()
+            return
+        }
+    }
+
     /// Arming-timeout fallback (seconds): if `onStarted` hasn't fired this long
     /// after `start()` was issued, flip to Listening anyway so a signal-wiring
     /// bug can't wedge the session at "Starting…" forever. Deliberately generous

OpenWhisp/Services/TranscriptionEngine.swift

@@ -128,6 +128,18 @@ public protocol StreamingTranscriptionEngine: AnyObject {
     /// non-default mic still recorded the built-in one.
     func selectDevice(_ deviceID: String)
 
+    /// True while a non-cancel `stop` is still actively producing the session's
+    /// genuine final — e.g. Parakeet draining its queued audio feed through the
+    /// decoder before `finish()`, or WhisperKit's `finalizeTail` re-decode. The
+    /// stuck-session fallback (`StreamingRoutePolicy.runStopFallback`) polls this
+    /// instead of firing at a fixed deadline: a long dictation's decode backlog
+    /// can outlast any fixed window, and completing the session with the stale
+    /// last partial would silently drop the tail words the engine is about to
+    /// deliver. Engines whose final is synchronous or internally guaranteed
+    /// (Apple Speech's 0.8s synthesized final, SpeechAnalyzer's stop-time
+    /// delivery) keep the default `false`.
+    var isFinalizing: Bool { get }
+
     /// Begin streaming recognition for `language` ("auto" = current locale).
     ///
     /// `prompt` carries the session's vocabulary/bias context (the same
@@ -149,6 +161,11 @@ public protocol StreamingTranscriptionEngine: AnyObject {
 }
 
 extension StreamingTranscriptionEngine {
+    /// Default: the engine's final needs no extended wait (it arrives promptly or
+    /// is internally guaranteed), so the stuck-session fallback keeps its original
+    /// single-deadline behavior.
+    public var isFinalizing: Bool { false }
+
     /// Convenience for call sites (and the test fake) that carry no bias context.
     /// Not a way to smuggle a prompt past the capability gate — it forwards the
     /// empty string, which every engine honors trivially (it's a no-op).

OpenWhisp/Services/WhisperKitStreamingEngine.swift

@@ -58,6 +58,19 @@ final class WhisperKitStreamingEngine: NSObject, StreamingTranscriptionEngine {
     // unconfirmed/hypothesis tail.
     @MainActor private var lastConfirmedText: String = ""
     @MainActor private var didFinish: Bool = false
+    /// Non-cancel stops enqueued whose `runStop` hasn't completed yet — the
+    /// engine is still working toward a genuine final (the finalizeTail
+    /// re-decode). A counter, not a bool: on a fast stop→start→stop, session
+    /// N+1's stop is already counted while session N's runStop is still queued,
+    /// and N's completion must not clear N+1's pending finalize. Mirrors
+    /// ParakeetStreamingEngine.
+    @MainActor private var pendingFinalizeStops = 0
+
+    /// True while a non-cancel stop is still flushing toward its final. Read by
+    /// AppState's stuck-session fallback; main-actor callers only.
+    var isFinalizing: Bool {
+        MainActor.assumeIsolated { pendingFinalizeStops > 0 }
+    }
     /// Whether this stream's `onStarted` already fired. AudioStreamTranscriber
     /// installs the tap inside `handle.start()`, which doesn't return until the
     /// stream ENDS — so the first state diff (they arrive per buffer, energy
@@ -213,6 +226,11 @@ final class WhisperKitStreamingEngine: NSObject, StreamingTranscriptionEngine {
             // final with the successor's sessionID, straight through AppState's
             // session fence.
             let final = self.onFinal
+            // Report "finalizing" from ENQUEUE until runStop completes, so the
+            // stuck-session fallback (StreamingRoutePolicy.runStopFallback) keeps
+            // waiting for the genuine final instead of completing the session
+            // with the stale last partial while finalizeTail decodes the tail.
+            if !cancel { self.pendingFinalizeStops += 1 }
             // Strong capture (see start): teardown must run even if the caller drops
             // its reference to this engine immediately after enqueueing the stop.
             self.lifecycle.enqueue {
@@ -226,6 +244,9 @@ final class WhisperKitStreamingEngine: NSObject, StreamingTranscriptionEngine {
     /// before any queued start proceeds.
     @MainActor
     private func runStop(cancel: Bool, final deliverFinal: ((String) -> Void)?) async {
+        // Balance stop(cancel:false)'s increment on every exit path — the
+        // fallback poll must see finalizing end even when no final is delivered.
+        defer { if !cancel { pendingFinalizeStops = max(0, pendingFinalizeStops - 1) } }
         removeConfigChangeObserver()
         let handle = transcriber
         transcriber = nil