fix(mac,asr): harden menu-bar delivery, chunk retry, and polish validator

Retain the external target app for menu-bar paste and polish context, retry failed middle ASR chunks once, and stop false-positive path/number violations.
2026-07-29 18:30:36 +08:00
parent 34be2e8dd1
commit 4c929d8b8b
17 changed files with 646 additions and 39 deletions
@@ -68,7 +68,14 @@ enum MacAppContextService {
    }

    static func detectContext() -> AppContext {
-        guard let bundleId = frontmostBundleIdentifier() else { return .unknown }
+        detectContext(bundleIdentifier: frontmostBundleIdentifier())
+    }
+
+    /// Resolve polish context from the application captured for this dictation
+    /// session. This avoids reading OSGKeyboard itself after a popover steals
+    /// focus.
+    static func detectContext(bundleIdentifier bundleId: String?) -> AppContext {
+        guard let bundleId else { return .unknown }
        if let mapped = contextByBundleId[bundleId] { return mapped }
        if chatBundleIdsFromRegistry.contains(bundleId) { return .chat }
        if bundleId.hasPrefix("com.apple.Safari") || bundleId.contains("chrome") {
@@ -83,4 +90,12 @@ enum MacAppContextService {
        let context = detectContext()
        store.setDetectedAppContext(context)
    }
+
+    static func captureAndPersist(
+        application: NSRunningApplication?,
+        to store: AppGroupStore
+    ) {
+        let context = detectContext(bundleIdentifier: application?.bundleIdentifier)
+        store.setDetectedAppContext(context)
+    }
 }
@@ -8,7 +8,14 @@

@preconcurrency import AVFoundation

-final class MacAudioRecorder: @unchecked Sendable {
+protocol MacAudioRecording: Sendable {
+    func level() -> Float
+    func start() async throws
+    func makeSnapshotStream() -> AsyncStream<AudioBufferSnapshot>
+    func stop() -> [Float]
+}
+
+final class MacAudioRecorder: MacAudioRecording, @unchecked Sendable {
    enum RecorderError: Error, LocalizedError {
        case converterUnavailable
        case microphoneAccessDenied
@@ -60,6 +60,7 @@ enum MacDictationPipeline {
    static func run(
        samples: [Float],
        store: AppGroupStore,
+        targetAppBundleIdentifier: String? = nil,
        onPartial: (@Sendable (String) -> Void)? = nil
    ) async throws -> MacDictationResult {
        guard !samples.isEmpty else { throw MacDictationError.noAudio }
@@ -69,7 +70,11 @@ enum MacDictationPipeline {
        var localBias: LocalASRBiasPayload?

        if store.engineMode == "local" {
-            localBias = resolveLocalBias(store: store, locale: locale)
+            localBias = resolveLocalBias(
+                store: store,
+                locale: locale,
+                targetAppBundleIdentifier: targetAppBundleIdentifier
+            )
            raw = try await MacLocalASRService.transcribe(
                samples: samples,
                locale: locale,
@@ -103,6 +108,7 @@ enum MacDictationPipeline {
        stream: AsyncStream<AudioBufferSnapshot>,
        finishSignal: AsyncStream<Void>,
        store: AppGroupStore,
+        targetAppBundleIdentifier: String?,
        onPartial: @escaping @Sendable (String) -> Void
    ) async -> MacLiveASRCaptureResult {
        if store.engineMode == "local", MacLocalASRService.usesMLXLiveStreaming() {
@@ -110,6 +116,7 @@ enum MacDictationPipeline {
                audioStream: stream,
                finishSignal: finishSignal,
                store: store,
+                targetAppBundleIdentifier: targetAppBundleIdentifier,
                onPartial: onPartial
            )
        }
@@ -117,7 +124,11 @@ enum MacDictationPipeline {
        let locale = resolvedLocale(store: store)
        let localBias: LocalASRBiasPayload?
        if store.engineMode == "local" {
-            localBias = resolveLocalBias(store: store, locale: locale)
+            localBias = resolveLocalBias(
+                store: store,
+                locale: locale,
+                targetAppBundleIdentifier: targetAppBundleIdentifier
+            )
        } else {
            localBias = nil
        }
@@ -280,15 +291,15 @@ enum MacDictationPipeline {

    private static func resolveLocalBias(
        store: AppGroupStore,
-        locale: Locale
+        locale: Locale,
+        targetAppBundleIdentifier: String?
    ) -> LocalASRBiasPayload? {
-        MacAppContextService.captureAndPersist(to: store)
        let capabilities = MacLocalASRService.currentCapabilities()
        let bias = LocalASRBiasAdapter.adapt(
            LocalASRBiasRequest(
                dictionary: store.personalDictionary,
                locale: locale,
-                frontAppBundleId: MacAppContextService.frontmostBundleIdentifier(),
+                frontAppBundleId: targetAppBundleIdentifier,
                capabilities: capabilities
            )
        )
@@ -75,14 +75,21 @@ final class MacDictationViewModel: ObservableObject {
    @Published var config: ProviderConfig

    let defaults: UserDefaults
-    private let recorder = MacAudioRecorder()
-    private let hotkeyService = MacHotkeyService()
+    private let recorder: any MacAudioRecording
+    private let hotkeyService: MacHotkeyService
    private var levelTimer: Timer?
    private var sessionTimer: Timer?
    private var cancellables = Set<AnyCancellable>()
    /// In-flight `beginRecording` started by the hotkey — cancelled if the
    /// key is released before the engine is ready (avoids a stuck session).
    private var hotkeyBeginTask: Task<Void, Never>?
+    /// Button-triggered preparation needs the same cancellation semantics as
+    /// the hotkey path when the user clicks Stop before the engine is ready.
+    private var buttonBeginTask: Task<Void, Never>?
+    /// Captured before the menu-bar popover activates OSGKeyboard.
+    private var preparedPopoverTargetApplication: NSRunningApplication?
+    /// Frozen for one take so app switches during ASR cannot redirect delivery.
+    private var sessionTargetApplication: NSRunningApplication?
    /// Live chunked / streaming ASR while recording (cloud or MLX local).
    /// Finished in `finishRecording` so partials can become the final draft.
    private var liveCaptureTask: Task<MacLiveASRCaptureResult, Never>?
@@ -97,8 +104,15 @@ final class MacDictationViewModel: ObservableObject {
        static let hotkeyTrigger = MacHotkeyTrigger.storageKey
    }

-    init(defaults: UserDefaults = .standard) {
+    init(
+        defaults: UserDefaults = .standard,
+        recorder: any MacAudioRecording = MacAudioRecorder(),
+        hotkeyService: MacHotkeyService = MacHotkeyService(),
+        startHotkeyService: Bool = true
+    ) {
        self.defaults = defaults
+        self.recorder = recorder
+        self.hotkeyService = hotkeyService
        self.config = ProviderConfig(defaults: defaults)
        self.usageStatistics = UsageStatisticsStore(defaults: defaults)
        self.autoPasteEnabled = defaults.object(forKey: StoredKeys.autoPaste) as? Bool ?? true
@@ -111,7 +125,9 @@ final class MacDictationViewModel: ObservableObject {

        MacICloudSyncBootstrap.configure(defaults: defaults)
        statusMessage = MacL10n.string("mac.status.ready", language: config.uiLanguage)
-        wireHotkeyService()
+        if startHotkeyService {
+            wireHotkeyService()
+        }
        forwardNestedObjectChanges()
    }

@@ -132,6 +148,17 @@ final class MacDictationViewModel: ObservableObject {
        refreshForegroundAppName()
    }

+    /// Called immediately before the menu-bar popover activates the app.
+    func prepareForPopoverPresentation() {
+        let target = MacTextInsertionService.captureTargetApplication()
+        preparedPopoverTargetApplication = target
+        foregroundAppName = target?.localizedName
+    }
+
+    func clearPreparedPopoverTarget() {
+        preparedPopoverTargetApplication = nil
+    }
+
    func reloadConfigFromCloud() {
        config.reloadFromPersistedStorage()
        statusMessage = MacL10n.string("mac.status.ready", language: config.uiLanguage)
@@ -256,7 +283,10 @@ final class MacDictationViewModel: ObservableObject {
        if isRecording || isPreparingToRecord {
            cancelOrFinishRecording()
        } else {
-            Task { await beginRecording() }
+            buttonBeginTask?.cancel()
+            buttonBeginTask = Task { [weak self] in
+                await self?.beginRecording()
+            }
        }
    }

@@ -264,8 +294,12 @@ final class MacDictationViewModel: ObservableObject {
        guard !isProcessing, !isRecording, !isPreparingToRecord else { return }
        isPreparingToRecord = true
        let store = AppGroupStore(defaults: defaults)
-        MacAppContextService.captureAndPersist(to: store)
-        refreshForegroundAppName()
+        let targetApplication = preparedPopoverTargetApplication
+            ?? MacTextInsertionService.captureTargetApplication()
+        preparedPopoverTargetApplication = nil
+        sessionTargetApplication = targetApplication
+        MacAppContextService.captureAndPersist(application: targetApplication, to: store)
+        foregroundAppName = targetApplication?.localizedName

        do {
            try await recorder.start()
@@ -274,14 +308,20 @@ final class MacDictationViewModel: ObservableObject {
            isPreparingToRecord = false
            if Task.isCancelled {
                _ = recorder.stop()
+                sessionTargetApplication = nil
+                buttonBeginTask = nil
                return
            }
+            buttonBeginTask = nil
            isRecording = true
            transcript = ""
            isStreamingPartial = false
            statusMessage = MacL10n.string("mac.status.listening", language: config.uiLanguage)
            startTimers()
-            startLiveCaptureIfSupported(store: store)
+            startLiveCaptureIfSupported(
+                store: store,
+                targetAppBundleIdentifier: targetApplication?.bundleIdentifier
+            )
            // Tiny race: Option released between the cancel check and
            // `isRecording = true`. Treat it as end-of-hold and finish.
            if Task.isCancelled {
@@ -289,6 +329,8 @@ final class MacDictationViewModel: ObservableObject {
            }
        } catch {
            isPreparingToRecord = false
+            sessionTargetApplication = nil
+            buttonBeginTask = nil
            if !Task.isCancelled {
                statusMessage = error.localizedDescription
            }
@@ -307,6 +349,9 @@ final class MacDictationViewModel: ObservableObject {
        stopTimers()
        audioLevel = 0
        let store = AppGroupStore(defaults: defaults)
+        let targetApplication = sessionTargetApplication
+        let targetAppBundleIdentifier = targetApplication?.bundleIdentifier
+        sessionTargetApplication = nil
        let usesDeferredStop = MacDictationPipeline.supportsLivePartials(store: store)
            && store.engineMode == "local"
            && MacLocalASRService.usesMLXLiveStreaming()
@@ -350,6 +395,7 @@ final class MacDictationViewModel: ObservableObject {
                        result = try await MacDictationPipeline.run(
                            samples: capturedSamples,
                            store: store,
+                            targetAppBundleIdentifier: targetAppBundleIdentifier,
                            onPartial: { [weak self] partial in
                                Task { @MainActor in
                                    self?.transcript = partial
@@ -362,6 +408,7 @@ final class MacDictationViewModel: ObservableObject {
                    result = try await MacDictationPipeline.run(
                        samples: capturedSamples,
                        store: store,
+                        targetAppBundleIdentifier: targetAppBundleIdentifier,
                        onPartial: { [weak self] partial in
                            Task { @MainActor in
                                self?.transcript = partial
@@ -370,7 +417,10 @@ final class MacDictationViewModel: ObservableObject {
                    )
                }
                self.transcript = result.text
-                let pasted = try await self.deliver(result.text)
+                let pasted = try await self.deliver(
+                    result.text,
+                    targetApplication: targetApplication
+                )
                self.recordUsage(for: result.text)
                self.speechHistory.append(text: result.text)
                self.appendToOverview(result.text)
@@ -393,7 +443,10 @@ final class MacDictationViewModel: ObservableObject {
        }
    }

-    private func startLiveCaptureIfSupported(store: AppGroupStore) {
+    private func startLiveCaptureIfSupported(
+        store: AppGroupStore,
+        targetAppBundleIdentifier: String?
+    ) {
        guard MacDictationPipeline.supportsLivePartials(store: store) else { return }
        let stream = recorder.makeSnapshotStream()
        let (finishStream, finishContinuation) = AsyncStream<Void>.makeStream(
@@ -405,6 +458,7 @@ final class MacDictationViewModel: ObservableObject {
                stream: stream,
                finishSignal: finishStream,
                store: store,
+                targetAppBundleIdentifier: targetAppBundleIdentifier,
                onPartial: { [weak self] partial in
                    Task { @MainActor in
                        guard let self else { return }
@@ -449,22 +503,32 @@ final class MacDictationViewModel: ObservableObject {
    /// Stops an in-flight prepare, or finishes an active recording.
    private func cancelOrFinishRecording() {
        if isRecording {
+            buttonBeginTask = nil
            finishRecording()
            return
        }
        if isPreparingToRecord {
            hotkeyBeginTask?.cancel()
            hotkeyBeginTask = nil
-            // If the button-triggered prepare wasn't tracked by hotkeyBeginTask,
-            // still clear the preparing flag and stop any engine that raced in.
-            isPreparingToRecord = false
+            buttonBeginTask?.cancel()
+            buttonBeginTask = nil
+            // Keep the preparation gate closed until the cancelled start call
+            // actually returns; otherwise a rapid third click can start a
+            // second recorder task while the first one is still unwinding.
            cancelLiveCapture()
            _ = recorder.stop()
        }
    }

-    private func deliver(_ text: String) async throws -> Bool {
-        try await MacTextInsertionService.insert(text, autoPaste: autoPasteEnabled)
+    private func deliver(
+        _ text: String,
+        targetApplication: NSRunningApplication?
+    ) async throws -> Bool {
+        try await MacTextInsertionService.insert(
+            text,
+            autoPaste: autoPasteEnabled,
+            targetApp: targetApplication
+        )
    }

    private func statusAfterDelivery(
@@ -14,10 +14,15 @@ enum MacMLXLiveCapture {
        audioStream: AsyncStream<AudioBufferSnapshot>,
        finishSignal: AsyncStream<Void>,
        store: AppGroupStore,
+        targetAppBundleIdentifier: String?,
        onPartial: @escaping @Sendable (String) -> Void
    ) async -> MacLiveASRCaptureResult {
        let locale = Locale(identifier: store.localeId.isEmpty ? "zh-CN" : store.localeId)
-        let bias = resolveBias(store: store, locale: locale)
+        let bias = resolveBias(
+            store: store,
+            locale: locale,
+            targetAppBundleIdentifier: targetAppBundleIdentifier
+        )

        guard let model = MacLocalASRService.selectedModelDefinition(),
              model.backend == .mlx,
@@ -133,14 +138,17 @@ enum MacMLXLiveCapture {
        }
    }

-    private static func resolveBias(store: AppGroupStore, locale: Locale) -> LocalASRBiasPayload? {
-        MacAppContextService.captureAndPersist(to: store)
+    private static func resolveBias(
+        store: AppGroupStore,
+        locale: Locale,
+        targetAppBundleIdentifier: String?
+    ) -> LocalASRBiasPayload? {
        let capabilities = MacLocalASRService.currentCapabilities()
        let bias = LocalASRBiasAdapter.adapt(
            LocalASRBiasRequest(
                dictionary: store.personalDictionary,
                locale: locale,
-                frontAppBundleId: MacAppContextService.frontmostBundleIdentifier(),
+                frontAppBundleId: targetAppBundleIdentifier,
                capabilities: capabilities
            )
        )
@@ -12,6 +12,10 @@ import Carbon
 import Foundation

 enum MacTextInsertionService {
+    /// Paste has no completion callback. Keep the transcript available long
+    /// enough for slower apps to consume the event before restoring clipboard.
+    static let pasteboardRestoreDelayNanoseconds: UInt64 = 500_000_000
+
    enum InsertionError: Error, LocalizedError {
        case accessibilityNotGranted

@@ -72,6 +76,7 @@ enum MacTextInsertionService {
        let snapshot = snapshotItems(of: pasteboard)
        pasteboard.clearContents()
        pasteboard.setString(text, forType: .string)
+        let transcriptChangeCount = pasteboard.changeCount

        guard autoPaste else { return false }
        guard AXIsProcessTrusted() else { throw InsertionError.accessibilityNotGranted }
@@ -84,11 +89,25 @@ enum MacTextInsertionService {

        // Give the target app time to read the transcript off the
        // pasteboard, then restore whatever the user had on it.
-        try? await Task.sleep(nanoseconds: 300_000_000)
-        restoreItems(snapshot, to: pasteboard)
+        try? await Task.sleep(nanoseconds: pasteboardRestoreDelayNanoseconds)
+        if shouldRestorePasteboard(
+            transcriptChangeCount: transcriptChangeCount,
+            currentChangeCount: pasteboard.changeCount
+        ) {
+            restoreItems(snapshot, to: pasteboard)
+        }
        return true
    }

+    /// Do not overwrite clipboard content written by the user, target app, or
+    /// a clipboard manager while the synthesized paste was in flight.
+    static func shouldRestorePasteboard(
+        transcriptChangeCount: Int,
+        currentChangeCount: Int
+    ) -> Bool {
+        transcriptChangeCount == currentChangeCount
+    }
+
    /// Brings `app` forward and waits (up to ~1 s) until it is frontmost so
    /// the synthesized keystroke isn't swallowed mid-switch.
    @MainActor
@@ -119,12 +138,12 @@ enum MacTextInsertionService {
        }
    }

-    private static func restoreItems(
+    static func restoreItems(
        _ items: [[NSPasteboard.PasteboardType: Data]],
        to pasteboard: NSPasteboard
    ) {
-        guard !items.isEmpty else { return }
        pasteboard.clearContents()
+        guard !items.isEmpty else { return }
        pasteboard.writeObjects(items.map { flavours in
            let item = NSPasteboardItem()
            for (type, data) in flavours { item.setData(data, forType: type) }
@@ -104,7 +104,7 @@ enum MacMainWindow {
 /// AppKit rather than SwiftUI's `MenuBarExtra` because the latter is flaky
 /// when combined with a primary `Window` scene (the icon can silently vanish).
@MainActor
-final class MacAppDelegate: NSObject, NSApplicationDelegate {
+final class MacAppDelegate: NSObject, NSApplicationDelegate, NSPopoverDelegate {
    private var statusItem: NSStatusItem?
    private let popover = NSPopover()

@@ -183,6 +183,7 @@ final class MacAppDelegate: NSObject, NSApplicationDelegate {
    }

    private func configurePopover() {
+        popover.delegate = self
        popover.behavior = .transient
        popover.animates = true
        popover.contentSize = NSSize(width: 340, height: 420)
@@ -194,11 +195,18 @@ final class MacAppDelegate: NSObject, NSApplicationDelegate {
        if popover.isShown {
            popover.performClose(sender)
        } else {
+            // Capture before activation: once the popover becomes key,
+            // NSWorkspace reports OSGKeyboard instead of the user's target.
+            MacDictationViewModel.shared.prepareForPopoverPresentation()
            NSApp.activate(ignoringOtherApps: true)
            popover.show(relativeTo: button.bounds, of: button, preferredEdge: .minY)
            popover.contentViewController?.view.window?.makeKey()
        }
    }
+
+    func popoverDidClose(_ notification: Notification) {
+        MacDictationViewModel.shared.clearPreparedPopoverTarget()
+    }
 }

 /// SwiftUI content hosted inside the status-bar popover. Shares the single