feat: migrate on-device Qwen3 ASR to CoreML for background Flow dictation

Replace MLX GPU inference with CoreML bundles so transcription continues while the host app is backgrounded. Adds model download and warm-up, vendored Qwen3Speech, and updates onboarding, settings, and copy for the ~1.6 GB CoreML package (iOS 18+).
2026-06-23 00:46:58 +08:00
parent 5e5122f172
commit df1c5ff32c
160 changed files with 22080 additions and 492 deletions
@@ -3,12 +3,14 @@
 //
 // Session Owner for TypeWhisper-style Flow dictation: continuous
 // `.playAndRecord` capture for the whole session, utterance gating for
-// ASR, optional LLM polish, and App Group result delivery.
+// ASR and cloud LLM polish, with App Group result delivery.

 import Foundation
 import AVFoundation
 import Speech
 import OSGKeyboardShared
+import UIKit
+import SwiftUI

@MainActor
 final class FlowSessionManager: ObservableObject {
@@ -19,9 +21,22 @@ final class FlowSessionManager: ObservableObject {
    @Published private(set) var sessionWarning: String?

    private let capture = FlowContinuousCapture()
-    private let asr: ASRService = ASRServiceFactory.make()
-    private let polisher = PolishingService()
    private let store = AppGroupStore()
+    /// Cloud-engine polish only; local engine delivers raw ASR text.
+    private var polisher: PolishingService {
+        PolishingService()
+    }
+    /// Cached ASR instance shared with `OnDeviceModelWarmup`.
+    private var sessionASR: ASRService?
+    private var asr: ASRService {
+        if let sessionASR { return sessionASR }
+        let service = OnDeviceModelWarmup.shared.asrService(
+            engineMode: store.engineMode,
+            localBackend: store.localASRBackend
+        )
+        sessionASR = service
+        return service
+    }

    private var pollingTask: Task<Void, Never>?
    private var heartbeatTask: Task<Void, Never>?
@@ -33,10 +48,13 @@ final class FlowSessionManager: ObservableObject {
    private var isUtteranceProcessing = false
    private var finalizeTask: Task<Void, Never>?
    private var asrTask: Task<Void, Never>?
+    private var chunkedPipeline: ChunkedUtterancePipeline?
    private var currentPartial = ""
    private var lastFinal = ""
+    private var chunkWarnings: [String] = []
    /// True while the host app scene is `.active` — drives foreground renewal.
    private var isAppForeground = false
+    private var backgroundTaskID: UIBackgroundTaskIdentifier = .invalid

    init() {
        Task { @MainActor [weak self] in
@@ -126,6 +144,9 @@ final class FlowSessionManager: ObservableObject {
        startLevelPublishing()
        scheduleExpiry(after: remaining)

+        OnDeviceModelWarmup.shared.warmUpIfNeeded()
+        bindSessionASR()
+
        debug("Flow session restored (\(Int(remaining))s remaining)")
    }

@@ -149,13 +170,17 @@ final class FlowSessionManager: ObservableObject {
        if isUtteranceRecording || isUtteranceProcessing {
            capture.cancelUtterance()
            asrTask?.cancel()
+            Task { await chunkedPipeline?.cancel() }
            asr.cancel()
        }
        asrTask = nil
+        chunkedPipeline = nil
        isUtteranceRecording = false
        isUtteranceProcessing = false

        capture.stop()
+        endBackgroundKeepAlive()
+        sessionASR = nil
        FlowSessionBridge.markSessionInactive()
        FlowSessionDarwin.postSessionChanged()
        isActive = false
@@ -179,6 +204,81 @@ final class FlowSessionManager: ObservableObject {
        }
    }

+    /// Full scene lifecycle — keeps Flow + ASR alive across app switches.
+    func handleScenePhase(_ phase: ScenePhase) {
+        switch phase {
+        case .active:
+            FlowAppLifecycle.shared.setForeground(true)
+            setAppForeground(true)
+            resumeAfterForeground()
+        case .inactive:
+            writeHeartbeatIfActive()
+        case .background:
+            FlowAppLifecycle.shared.setForeground(false)
+            setAppForeground(false)
+            beginBackgroundKeepAlive()
+        @unknown default:
+            break
+        }
+    }
+
+    private func writeHeartbeatIfActive() {
+        guard isActive else { return }
+        FlowSessionBridge.writeHeartbeat()
+    }
+
+    private func beginBackgroundKeepAlive() {
+        guard isActive else { return }
+        FlowSessionBridge.writeHeartbeat()
+
+        guard backgroundTaskID == .invalid else { return }
+        backgroundTaskID = UIApplication.shared.beginBackgroundTask { [weak self] in
+            self?.endBackgroundKeepAlive()
+        }
+        debug("background keep-alive started")
+    }
+
+    private func endBackgroundKeepAlive() {
+        guard backgroundTaskID != .invalid else { return }
+        UIApplication.shared.endBackgroundTask(backgroundTaskID)
+        backgroundTaskID = .invalid
+        debug("background keep-alive ended")
+    }
+
+    private func resumeAfterForeground() {
+        guard isActive else {
+            endBackgroundKeepAlive()
+            return
+        }
+
+        FlowSessionBridge.writeHeartbeat()
+        endBackgroundKeepAlive()
+
+        Task { @MainActor [weak self] in
+            await self?.reactivateCaptureIfNeeded()
+            OnDeviceModelWarmup.shared.ensureReadyAfterBackground()
+            self?.bindSessionASR()
+        }
+    }
+
+    private func reactivateCaptureIfNeeded() async {
+        guard isActive else { return }
+
+        if capture.running {
+            capture.reassertIfRunning()
+            return
+        }
+
+        do {
+            try capture.start()
+            debug("capture restarted after foreground")
+        } catch {
+            let message = (error as? LocalizedError)?.errorDescription ?? error.localizedDescription
+            sessionWarning = message
+            debug("capture restart failed: \(message)")
+        }
+    }
+
    /// Extend the session before it expires while the host app stays in foreground.
    private func renewSessionIfNeededWhileForeground() {
        guard isActive, isAppForeground else { return }
@@ -220,14 +320,24 @@ final class FlowSessionManager: ObservableObject {
        startLevelPublishing()
        scheduleExpiry(after: duration)

+        OnDeviceModelWarmup.shared.warmUpIfNeeded()
+        bindSessionASR()
+
        debug("Flow session started (\(Int(duration))s), continuous capture running")
    }

+    private func bindSessionASR() {
+        sessionASR = OnDeviceModelWarmup.shared.asrService(
+            engineMode: store.engineMode,
+            localBackend: store.localASRBackend
+        )
+    }
+
    private func permissionWarningMessage() -> String {
        if AppPermissions.micStatus != .granted {
-            return NSLocalizedString("flow.error.micRequired", comment: "")
+            return AppL10n.string("flow.error.micRequired")
        }
-        return NSLocalizedString("flow.error.speechRequired", comment: "")
+        return AppL10n.string("flow.error.speechRequired")
    }

    // MARK: - Polling
@@ -259,19 +369,20 @@ final class FlowSessionManager: ObservableObject {

    private func beginUtterance() {
        guard capture.running else {
-            failUtterance(message: NSLocalizedString("flow.error.audioUnavailable", comment: ""))
+            failUtterance(message: AppL10n.string("flow.error.audioUnavailable"))
            return
        }
-        // Mirror `LiveDictationController.start`: only begin when the previous
-        // utterance fully finished. Never cancel an in-flight analyzer here —
-        // that was the source of intermittent CancellationError / noSpeech.
        guard !isUtteranceProcessing else {
            debug("beginUtterance ignored — previous utterance still processing")
            return
        }

+        // Honor engine / ASR backend changes without restarting the session.
+        bindSessionASR()
+
        currentPartial = ""
        lastFinal = ""
+        chunkWarnings = []

        let localeId = store.localeId
        FlowSessionBridge.setTranscriptionLanguage(localeId)
@@ -279,28 +390,42 @@ final class FlowSessionManager: ObservableObject {

        let locale = SpeechLocaleResolver.resolve(localeId)
        let stream = capture.beginUtterance()
-        let events = asr.transcribe(stream: stream, locale: locale)
+        let pipeline = ChunkedUtterancePipeline(asr: asr, locale: locale)
+        chunkedPipeline = pipeline

        isUtteranceRecording = true
+        FlowDiagnostics.log(
+            "beginUtterance engine=\(store.engineMode) asr=\(store.localASRBackend.rawValue) " +
+            "modelsInMemory=\(OnDeviceModelStatus.modelsLoadedInMemory()) " +
+            "asrType=\(type(of: asr)) pipelined=true max=\(Int(FlowSessionKeys.maxUtteranceDuration))s"
+        )

-        asrTask = Task { @MainActor [weak self] in
-            guard let self else { return }
-            for await event in events {
-                switch event {
-                case .capability:
-                    break
-                case .partial(let text):
-                    self.currentPartial = text
-                case .final(let text):
-                    self.lastFinal = text.trimmingCharacters(in: .whitespacesAndNewlines)
-                    self.currentPartial = ""
-                case .error(let message):
-                    self.debug("asr error: \(message)")
-                    if self.isUtteranceRecording {
-                        self.failUtterance(message: message)
-                    } else if self.isUtteranceProcessing {
-                        self.finishProcessing(withError: message)
+        asrTask = Task.detached(priority: .userInitiated) { [weak manager = self] in
+            let outcome = await pipeline.transcribe(stream: stream) { partial in
+                Task { @MainActor in
+                    manager?.currentPartial = partial
+                }
+            }
+            await MainActor.run {
+                guard let manager else { return }
+                FlowDiagnostics.log(
+                    "chunkedASR finished partialLen=\(manager.currentPartial.count) " +
+                    "finalPending=\(manager.lastFinal.isEmpty)"
+                )
+                switch outcome {
+                case .success(let success):
+                    manager.lastFinal = success.text
+                    manager.chunkWarnings = success.chunkWarnings
+                    manager.currentPartial = ""
+                case .failure(let message):
+                    manager.debug("asr error: \(message)")
+                    if manager.isUtteranceRecording {
+                        manager.failUtterance(message: message)
+                    } else if manager.isUtteranceProcessing {
+                        manager.finishProcessing(withError: message)
                    }
+                case .cancelled:
+                    break
                }
            }
        }
@@ -335,10 +460,13 @@ final class FlowSessionManager: ObservableObject {
        finalizeTask?.cancel()
        finalizeTask = nil
        asrTask?.cancel()
+        Task { await chunkedPipeline?.cancel() }
+        chunkedPipeline = nil
        asr.cancel()
        capture.cancelUtterance()
        currentPartial = ""
        lastFinal = ""
+        chunkWarnings = []
        FlowSessionBridge.setRecordingState(.idle)
        debug("utterance aborted")
    }
@@ -349,10 +477,13 @@ final class FlowSessionManager: ObservableObject {
        finalizeTask?.cancel()
        finalizeTask = nil
        asrTask?.cancel()
+        Task { await chunkedPipeline?.cancel() }
+        chunkedPipeline = nil
        asr.cancel()
        capture.cancelUtterance()
        currentPartial = ""
        lastFinal = ""
+        chunkWarnings = []
        FlowSessionBridge.storeTranscriptionError(message)
        FlowSessionBridge.setRecordingState(.idle)
        debug("utterance failed: \(message)")
@@ -362,26 +493,43 @@ final class FlowSessionManager: ObservableObject {
        isUtteranceProcessing = false
        finalizeTask?.cancel()
        finalizeTask = nil
+        chunkedPipeline = nil
        currentPartial = ""
        lastFinal = ""
+        chunkWarnings = []
        FlowSessionBridge.storeTranscriptionError(message)
        FlowSessionBridge.setRecordingState(.idle)
        debug("utterance processing failed: \(message)")
    }

    private func finalizeUtterance() async {
+        let pipelineStarted = Date()
        defer {
            isUtteranceProcessing = false
            FlowSessionBridge.setRecordingState(.idle)
        }

-        let deadline = Date().addingTimeInterval(30)
-        while Date() < deadline {
+        let asrWait = asrWaitTimeout()
+        FlowDiagnostics.log(
+            "finalize start asrWait=\(Int(asrWait))s engine=\(store.engineMode) " +
+            "backend=\(store.localASRBackend.rawValue)"
+        )
+
+        let asrDeadline = Date().addingTimeInterval(asrWait)
+        while Date() < asrDeadline {
            if !lastFinal.isEmpty { break }
            if asrTask?.isCancelled == true { break }
            try? await Task.sleep(nanoseconds: 100_000_000)
        }

+        if lastFinal.isEmpty, let asrTask {
+            FlowDiagnostics.log("ASR wait elapsed — awaiting asrTask completion")
+            _ = await asrTask.value
+        }
+
+        let asrElapsed = Date().timeIntervalSince(pipelineStarted)
+        FlowDiagnostics.log("ASR phase done in \(String(format: "%.1f", asrElapsed))s finalLen=\(lastFinal.count)")
+
        var text = lastFinal.trimmingCharacters(in: .whitespacesAndNewlines)
        if text.isEmpty {
            text = currentPartial.trimmingCharacters(in: .whitespacesAndNewlines)
@@ -390,36 +538,72 @@ final class FlowSessionManager: ObservableObject {
            let key = (asrTask?.isCancelled == true)
                ? "flow.error.recognitionInterrupted"
                : "flow.error.noSpeech"
+            FlowDiagnostics.log("finalize failed: empty transcript after \(String(format: "%.1f", asrElapsed))s")
            FlowSessionBridge.storeTranscriptionError(
-                NSLocalizedString(key, comment: "")
+                AppL10n.string(key)
            )
            return
        }

        let engineMode = store.engineMode
-        let modeId = store.modeId
-        let shouldPolish = engineMode != "local" && modeId == "polish"
+
+        if engineMode == "local" {
+            let warning = Self.chunkWarningMessage(chunkWarnings)
+            FlowSessionBridge.storeTranscriptionResult(text, polishWarning: warning)
+            FlowDiagnostics.log(
+                "finalize ASR-only total=\(String(format: "%.1f", Date().timeIntervalSince(pipelineStarted)))s " +
+                "len=\(text.count)"
+            )
+            SpeechHistoryStore.shared.append(text: text, engineMode: engineMode)
+            currentPartial = ""
+            lastFinal = ""
+            chunkWarnings = []
+            debug("utterance finalized length=\(text.count)")
+            return
+        }

        var delivered = text
-        if shouldPolish {
-            do {
-                let polished = try await polisher.polish(text)
-                delivered = polished
-                FlowSessionBridge.storeTranscriptionResult(polished)
-            } catch {
-                FlowSessionBridge.storeTranscriptionResult(text)
-            }
-        } else {
-            FlowSessionBridge.storeTranscriptionResult(text)
+        let chunkNote = Self.chunkWarningMessage(chunkWarnings)
+        let polishStarted = Date()
+        do {
+            let polished = try await polisher.polish(text)
+            delivered = polished
+            FlowSessionBridge.storeTranscriptionResult(polished, polishWarning: chunkNote)
+            FlowDiagnostics.log(
+                "polish done in \(String(format: "%.1f", Date().timeIntervalSince(polishStarted)))s " +
+                "total=\(String(format: "%.1f", Date().timeIntervalSince(pipelineStarted)))s"
+            )
+        } catch {
+            FlowDiagnostics.log(
+                "polish failed after \(String(format: "%.1f", Date().timeIntervalSince(polishStarted)))s: " +
+                "\(error.localizedDescription)"
+            )
+            FlowSessionBridge.storeTranscriptionResult(text, polishWarning: chunkNote)
        }

        SpeechHistoryStore.shared.append(text: delivered, engineMode: engineMode)

        currentPartial = ""
        lastFinal = ""
+        chunkWarnings = []
+        chunkedPipeline = nil
        debug("utterance finalized length=\(text.count)")
    }

+    private static func chunkWarningMessage(_ warnings: [String]) -> String? {
+        guard !warnings.isEmpty else { return nil }
+        return warnings.joined(separator: "\n")
+    }
+
+    private func asrWaitTimeout() -> TimeInterval {
+        if store.engineMode == "local" {
+            return store.localASRBackend == .qwen3ASR
+                ? FlowSessionKeys.localQwen3ASRWaitTimeout
+                : FlowSessionKeys.localASRWaitTimeout
+        }
+        return FlowSessionKeys.cloudASRWaitTimeout
+    }
+
    // MARK: - Level publishing (main thread only)

    private func startLevelPublishing() {
@@ -461,8 +645,6 @@ final class FlowSessionManager: ObservableObject {
    }

    private func debug(_ message: String) {
-        #if DEBUG
-        print("🌊[FlowSession] \(message)")
-        #endif
+        FlowDiagnostics.log(message)
    }
 }