Harden dictation cancellation and streamline audio-to-text processing

Sources/Dictation/DictationTranscriptPersistence.swift

@@ -0,0 +1,33 @@
+import Foundation
+
+/// Measures the writer itself, independently of Auto Enter and MainActor publication.
+struct DictationTranscriptPersistenceResult: Sendable {
+    let saved: SavedDictationTranscript?
+    let failureError: Error?
+    let startedAt: CFAbsoluteTime
+    let finishedAt: CFAbsoluteTime
+
+    var failureMessage: String? {
+        guard failureError != nil else { return nil }
+        return "Transcripted couldn't save a local copy of this dictation. Check your save location and available disk space."
+    }
+
+    static func measure(
+        now: () -> CFAbsoluteTime = CFAbsoluteTimeGetCurrent,
+        save: () throws -> SavedDictationTranscript
+    ) -> Self {
+        let startedAt = now()
+        do {
+            let saved = try save()
+            return Self(saved: saved, failureError: nil, startedAt: startedAt, finishedAt: now())
+        } catch {
+            return Self(saved: nil, failureError: error, startedAt: startedAt, finishedAt: now())
+        }
+    }
+}
+
+enum DictationSessionCompletionPolicy {
+    static func canPublish(sessionID: UUID, currentSessionID: UUID, isDictating: Bool, cancelled: Bool) -> Bool {
+        sessionID == currentSessionID && isDictating && !cancelled
+    }
+}

Sources/Speech/ASRInferenceWaiterQueue.swift

@@ -0,0 +1,40 @@
+import Foundation
+
+/// Cancellation removes only a pending waiter. Once handed off, the caller owns
+/// the reserved inference slot and must release it even if cancellation wins.
+@MainActor
+final class ASRInferenceWaiterQueue {
+    private struct Waiter {
+        let id: UUID
+        let continuation: CheckedContinuation<Void, Error>
+    }
+    private var waiters: [Waiter] = []
+    var count: Int { waiters.count }
+    var isEmpty: Bool { waiters.isEmpty }
+
+    func wait() async throws {
+        let id = UUID()
+        try await withTaskCancellationHandler {
+            try await withCheckedThrowingContinuation { (continuation: CheckedContinuation<Void, Error>) in
+                guard !Task.isCancelled else {
+                    continuation.resume(throwing: CancellationError())
+                    return
+                }
+                waiters.append(Waiter(id: id, continuation: continuation))
+            }
+        } onCancel: {
+            Task { @MainActor in self.cancel(id: id) }
+        }
+    }
+
+    /// The caller reserves the handoff before calling this method.
+    func resumeFirst() {
+        guard !waiters.isEmpty else { return }
+        waiters.removeFirst().continuation.resume()
+    }
+
+    private func cancel(id: UUID) {
+        guard let index = waiters.firstIndex(where: { $0.id == id }) else { return }
+        waiters.remove(at: index).continuation.resume(throwing: CancellationError())
+    }
+}

Sources/Speech/DictationSession.swift

@@ -131,14 +131,14 @@ extension DictationSession {
         isDictating: @escaping () -> Bool,
         onModelStateUpdate: @escaping (ParakeetModelState) -> Void
     ) async -> ModelWarmupOutcome {
+        let deadline = ProcessInfo.processInfo.systemUptime
+            + TranscriptedConstants.modelLoadWaitBudget
         if case .failed = appState.sttRouter.recordingModelDownloadState {
-            // A previous attempt failed; retry once before the wait loop
-            // treats .failed as terminal.
-            await appState.sttRouter.initializeRecordingModel()
+            // Retry once, observing the new attempt without awaiting native load.
+            appState.sttRouter.requestRecordingModelInitialization()
+            await appState.sttRouter.waitForRecordingModelLoadProgress(until: deadline)
         }
 
-        let deadline = ProcessInfo.processInfo.systemUptime
-            + TranscriptedConstants.modelLoadWaitBudget
         while ProcessInfo.processInfo.systemUptime < deadline {
             guard !Task.isCancelled, isDictating() else { return .aborted }
 
@@ -151,18 +151,12 @@ extension DictationSession {
             case .failed(let message):
                 return .failed(message)
             case .notLoaded, .cached:
-                let stateBefore = appState.sttRouter.recordingModelDownloadState.diagnosticName
-                await appState.sttRouter.initializeRecordingModel()
-                // If initialization bailed without progressing (e.g.
-                // mid-shutdown), sleep so this loop can't spin hot.
-                if appState.sttRouter.recordingModelDownloadState.diagnosticName == stateBefore {
-                    try? await Task.sleep(nanoseconds: TranscriptedConstants.modelLoadPollInterval)
-                }
+                appState.sttRouter.requestRecordingModelInitialization()
+                await appState.sttRouter.waitForRecordingModelLoadProgress(until: deadline)
             case .downloading, .loading:
-                // Downloads publish progress the overlay refreshes on a
-                // short poll; an in-flight load is joined directly so
-                // recording starts the moment it settles.
-                await appState.sttRouter.waitForRecordingModelLoadProgress()
+                // Observe progress without letting a shared native load outlive
+                // this caller's deadline or cancellation.
+                await appState.sttRouter.waitForRecordingModelLoadProgress(until: deadline)
             }
         }
 

Sources/Speech/ModelLoadProgressWaiter.swift

@@ -0,0 +1,45 @@
+import Combine
+import Foundation
+
+/// A caller-owned wait. Finishing releases its observation and timer, but never
+/// cancels the shared model initialization that other sessions may still need.
+@MainActor
+final class ModelLoadProgressWaiter {
+    private var continuation: CheckedContinuation<Void, Never>?
+    private var observation: AnyCancellable?
+    private var timer: Task<Void, Never>?
+
+    static func wait(for changes: AnyPublisher<Void, Never>, until deadline: TimeInterval) async {
+        let waiter = ModelLoadProgressWaiter()
+        await withTaskCancellationHandler {
+            await withCheckedContinuation { continuation in
+                waiter.continuation = continuation
+                let remaining = deadline - ProcessInfo.processInfo.systemUptime
+                guard !Task.isCancelled, remaining > 0 else {
+                    waiter.finish()
+                    return
+                }
+                waiter.observation = changes.sink { [weak waiter] in
+                    Task { @MainActor in waiter?.finish() }
+                }
+                waiter.timer = Task { @MainActor [weak waiter] in
+                    do { try await Task.sleep(for: .seconds(remaining)) }
+                    catch { return }
+                    waiter?.finish()
+                }
+            }
+        } onCancel: {
+            Task { @MainActor in waiter.finish() }
+        }
+    }
+
+    private func finish() {
+        guard let continuation else { return }
+        self.continuation = nil
+        observation?.cancel()
+        observation = nil
+        timer?.cancel()
+        timer = nil
+        continuation.resume()
+    }
+}

Sources/Speech/ParakeetAudioGraphOwnership.swift

@@ -399,3 +399,14 @@ enum ParakeetZombieRecoveryOwnershipPolicy {
             && expectedOwner.matches(generation: currentGraphGeneration, engine: currentEngine)
     }
 }
+
+/// A stopped recording can change without replacing its native graph (for
+/// example explicit discard). Both identities must survive async conversion.
+struct ParakeetRecordedSamplesClaim: Equatable {
+    let graphOwner: ParakeetAudioGraphOwnerToken
+    let revision: UInt64
+
+    func isCurrent(owner: ParakeetAudioGraphOwnerToken, revision: UInt64, cancelled: Bool) -> Bool {
+        !cancelled && owner == graphOwner && revision == self.revision
+    }
+}