From f0240c3088cdf3b4c0c63e453fa2e817c0ccc5f0 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Sun, 26 Jul 2026 10:15:31 +0000 Subject: [PATCH] feat(flow): add whole-utterance batch ASR fallback (P1) - Accumulate utterance PCM in FlowContinuousCapture for batch retry - Run full-utterance transcribeChunk when stitched final lags partial - Refactor Mac MLX tail drain to shared FlowUtteranceEndCoordinator - Add FlowUtterancePCMStore, UtteranceBatchFallbackPolicy, and tests Co-authored-by: Rocky --- CHANGELOG.md | 4 ++ OSGKeyboard/Services/FlowSessionManager.swift | 60 +++++++++++++++++++ OSGKeyboardMac/MacMLXLiveCapture.swift | 14 ++++- .../Services/FlowContinuousCapture.swift | 14 +++++ .../Utilities/FlowPipelineDiagnostics.swift | 12 ++++ .../Utilities/FlowUtterancePCMStore.swift | 46 ++++++++++++++ .../UtteranceBatchFallbackPolicy.swift | 41 +++++++++++++ .../FlowUtterancePCMStoreTests.swift | 23 +++++++ .../UtteranceBatchFallbackPolicyTests.swift | 45 ++++++++++++++ 9 files changed, 256 insertions(+), 3 deletions(-) create mode 100644 OSGKeyboardShared/Utilities/FlowUtterancePCMStore.swift create mode 100644 OSGKeyboardShared/Utilities/UtteranceBatchFallbackPolicy.swift create mode 100644 OSGKeyboardTests/FlowUtterancePCMStoreTests.swift create mode 100644 OSGKeyboardTests/UtteranceBatchFallbackPolicyTests.swift diff --git a/CHANGELOG.md b/CHANGELOG.md index 1cdc9e9..9fe9d9d 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -9,6 +9,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ### Fixed - **Flow tail ASR drop**: after mic stop, iOS Flow now uses a longer silence drain (350 ms), a fixed 150 ms post-roll, expanded final-chunk ASR recovery, and a partial transcript guard so weak trailing syllables are less likely to disappear from the result. / **Flow 尾音识别丢失**：松手后 iOS Flow 采用更长的静音排空（350 ms）、固定 150 ms 尾音保留、增强末块 ASR 恢复与 partial 兜底，降低弱尾音从结果中消失的概率。 +- **Flow batch ASR fallback**: when pipelined chunk output is clearly shorter than the live partial, the host re-transcribes the full utterance PCM captured during recording (Mac-style safety net). / **Flow 整句 ASR 兜底**：流水线拼接结果明显短于实时 partial 时，主 App 对录音期间累积的整段 PCM 重新识别（对齐 Mac 双保险）。 + +### Changed +- **Mac MLX tail drain**: streaming capture now uses the shared `FlowUtteranceEndCoordinator` (silence drain + post-roll) instead of an inline poll loop. / **Mac MLX 尾音排空**：流式采集改用 Shared 层 `FlowUtteranceEndCoordinator`（静音排空 + post-roll），替代内联轮询循环。 ### Added - **Mac MLX streaming ASR**: local dictation uses Qwen3-ASR via mlx-audio-swift with overlay partial preview, tail drain, vocabulary prompt, and polish-before-insert. / **Mac MLX 流式 ASR**：本地听写改用 mlx-audio-swift 的 Qwen3-ASR，支持浮层 partial 预览、尾部截断、词库 prompt 与润色后再插入。 diff --git a/OSGKeyboard/Services/FlowSessionManager.swift b/OSGKeyboard/Services/FlowSessionManager.swift index 91933db..64d3b3a 100644 --- a/OSGKeyboard/Services/FlowSessionManager.swift +++ b/OSGKeyboard/Services/FlowSessionManager.swift @@ -66,6 +66,8 @@ final class FlowSessionManager: ObservableObject { private var lastFinal = "" /// Partial stitched text captured when the user stops recording. private var bestPartialSnapshot = "" + /// Full utterance PCM for batch ASR fallback after pipelined chunking. + private var utterancePCMSamples: [Float] = [] private var chunkWarnings: [String] = [] private var lastReadyTraceSignature = "" private var lastCommandFingerprint = "" @@ -359,6 +361,7 @@ final class FlowSessionManager: ObservableObject { currentPartial = "" lastFinal = "" bestPartialSnapshot = "" + utterancePCMSamples = [] chunkWarnings = [] FlowSessionBridge.setHostReady(false) } @@ -1105,6 +1108,7 @@ final class FlowSessionManager: ObservableObject { currentPartial = "" lastFinal = "" bestPartialSnapshot = "" + utterancePCMSamples = [] chunkWarnings = [] let localeId = store.localeId @@ -1214,6 +1218,7 @@ final class FlowSessionManager: ObservableObject { guard let self else { return } let drainReport = await self.capture.endUtteranceAndDrain() FlowDiagnostics.logDrain(drainReport) + self.utterancePCMSamples = self.capture.consumeUtteranceSamples() await self.finalizeUtterance( sessionId: drainingSessionId, utteranceId: drainingUtteranceId, @@ -1239,6 +1244,7 @@ final class FlowSessionManager: ObservableObject { currentPartial = "" lastFinal = "" bestPartialSnapshot = "" + utterancePCMSamples = [] chunkWarnings = [] currentUtteranceId = nil currentCommandSeq = 0 @@ -1266,6 +1272,7 @@ final class FlowSessionManager: ObservableObject { currentPartial = "" lastFinal = "" bestPartialSnapshot = "" + utterancePCMSamples = [] chunkWarnings = [] storeCurrentError(message, kind: kind) currentUtteranceId = nil @@ -1289,6 +1296,7 @@ final class FlowSessionManager: ObservableObject { currentPartial = "" lastFinal = "" bestPartialSnapshot = "" + utterancePCMSamples = [] chunkWarnings = [] storeCurrentError(message, kind: kind) currentUtteranceId = nil @@ -1348,6 +1356,14 @@ final class FlowSessionManager: ObservableObject { if text.isEmpty { text = currentPartial.trimmingCharacters(in: .whitespacesAndNewlines) } + + if UtteranceBatchFallbackPolicy.shouldRunBatchFallback( + stitchedFinal: lastFinal, + partialSnapshot: bestPartialSnapshot + ), !utterancePCMSamples.isEmpty { + text = await runBatchASRFallback(currentText: text) + } + utterancePCMSamples = [] guard !text.isEmpty else { let key = (asrTask?.isCancelled == true || Task.isCancelled) ? "flow.error.recognitionInterrupted" @@ -1441,6 +1457,7 @@ final class FlowSessionManager: ObservableObject { currentPartial = "" lastFinal = "" bestPartialSnapshot = "" + utterancePCMSamples = [] chunkWarnings = [] chunkedPipeline = nil debug("utterance finalized length=\(text.count)") @@ -1569,6 +1586,49 @@ final class FlowSessionManager: ObservableObject { ) } + /// Re-transcribe the full utterance PCM when pipelined chunking likely dropped tail text. + private func runBatchASRFallback(currentText: String) async -> String { + let samples = utterancePCMSamples + guard !samples.isEmpty else { return currentText } + + let locale = SpeechLocaleResolver.resolve(store.localeId) + let stitched = lastFinal.trimmingCharacters(in: .whitespacesAndNewlines) + let partial = bestPartialSnapshot.trimmingCharacters(in: .whitespacesAndNewlines) + + FlowDiagnostics.log( + "batch fallback start samples=\(samples.count) stitchedLen=\(stitched.count) partialLen=\(partial.count)" + ) + + let asrService = asr + let result = await Task.detached(priority: .userInitiated) { [asrService] in + await asrService.transcribeChunk(samples: samples, locale: locale) + }.value + + switch result { + case .success(let batchText): + let trimmedBatch = batchText.trimmingCharacters(in: .whitespacesAndNewlines) + guard !trimmedBatch.isEmpty else { return currentText } + let resolved = UtteranceBatchFallbackPolicy.preferredTranscript( + batch: trimmedBatch, + stitchedFinal: stitched, + partialSnapshot: partial, + current: currentText + ) + FlowPipelineDiagnostics.logBatchFallback( + sampleCount: samples.count, + stitchedLength: stitched.count, + partialLength: partial.count, + batchLength: trimmedBatch.count + ) + return resolved + case .failure(let message): + FlowDiagnostics.log("batch fallback failed: \(message)") + return currentText + case .cancelled: + return currentText + } + } + private func asrWaitTimeout() -> TimeInterval { // v0.2.0: local engine is iOS `SpeechAnalyzer` only, so the // previous Qwen3-specific timeout collapses into the shared diff --git a/OSGKeyboardMac/MacMLXLiveCapture.swift b/OSGKeyboardMac/MacMLXLiveCapture.swift index 46f070f..c26eb81 100644 --- a/OSGKeyboardMac/MacMLXLiveCapture.swift +++ b/OSGKeyboardMac/MacMLXLiveCapture.swift @@ -43,6 +43,7 @@ enum MacMLXLiveCapture { let drainTracker = FlowCaptureDrainTracker() let draining = OSAllocatedUnfairLock(initialState: false) + let drainComplete = OSAllocatedUnfairLock(initialState: false) let pendingFeed = OSAllocatedUnfairLock(initialState: [Float]()) let feedIntervalSamples = 1_600 // 100 ms @ 16 kHz @@ -56,6 +57,11 @@ enum MacMLXLiveCapture { for await _ in finishSignal { draining.withLock { $0 = true } drainTracker.beginDrain() + _ = await FlowUtteranceEndCoordinator.awaitTailCapture( + tracker: drainTracker, + policy: tailDrainPolicy + ) + drainComplete.withLock { $0 = true } break } } @@ -63,10 +69,12 @@ enum MacMLXLiveCapture { group.addTask { for await snapshot in audioStream { if Task.isCancelled { break } + if drainComplete.withLock({ $0 }) { break } if draining.withLock({ $0 }) { - drainTracker.noteAudio(samples: snapshot.samples, policy: tailDrainPolicy) - let decision = drainTracker.shouldFinish(policy: tailDrainPolicy) - if decision.finished { break } + drainTracker.noteAudio( + samples: snapshot.samples, + policy: tailDrainPolicy + ) } pendingFeed.withLock { buffer in buffer.append(contentsOf: snapshot.samples) diff --git a/OSGKeyboardShared/Services/FlowContinuousCapture.swift b/OSGKeyboardShared/Services/FlowContinuousCapture.swift index 7c82350..5cf02fc 100644 --- a/OSGKeyboardShared/Services/FlowContinuousCapture.swift +++ b/OSGKeyboardShared/Services/FlowContinuousCapture.swift @@ -281,6 +281,9 @@ public final class FlowContinuousCapture { private let gate = OSAllocatedUnfairLock(initialState: UtteranceGatePhase.idle) private let drainTracker = FlowCaptureDrainTracker() private let tailSampleCounter = OSAllocatedUnfairLock(initialState: 0) + private let utterancePCMStore = FlowUtterancePCMStore( + maxSampleCount: Int(FlowSessionKeys.maxUtteranceDuration) * 16_000 + ) private var downsampler: AdaptiveDownsampler? private var targetFormat: AVAudioFormat? @@ -412,6 +415,7 @@ public final class FlowContinuousCapture { let proof = audioProofStore let tracker = drainTracker let tailCounter = tailSampleCounter + let pcmStore = utterancePCMStore let policy = drainPolicy let tap = Self.makeAudioTapBlock( downsampler: downsampler, @@ -422,6 +426,7 @@ public final class FlowContinuousCapture { streamRelay: relay, drainTracker: tracker, tailSampleCounter: tailCounter, + utterancePCMStore: pcmStore, drainPolicy: policy ) // `format: nil` binds the tap to the input node's *live* format. Passing @@ -645,6 +650,7 @@ public final class FlowContinuousCapture { let (stream, continuation) = AsyncStream.makeStream() drainTracker.reset() tailSampleCounter.withLock { $0 = 0 } + utterancePCMStore.reset() // Bind the consumer before opening the gate so early tap frames // are not dropped on the floor. streamRelay.bind(continuation) @@ -697,11 +703,17 @@ public final class FlowContinuousCapture { return report } + /// Returns the utterance PCM accumulated during the last recording cycle. + public func consumeUtteranceSamples() -> [Float] { + utterancePCMStore.consume() + } + /// Immediate stop without tail drain (abort / session teardown). public func cancelUtterance() { gate.withLock { $0 = .idle } drainTracker.reset() tailSampleCounter.withLock { $0 = 0 } + utterancePCMStore.reset() streamRelay.finish() } @@ -720,6 +732,7 @@ public final class FlowContinuousCapture { streamRelay: FlowCaptureStreamRelay, drainTracker: FlowCaptureDrainTracker, tailSampleCounter: OSAllocatedUnfairLock, + utterancePCMStore: FlowUtterancePCMStore, drainPolicy: FlowCaptureTailDrainPolicy ) -> @Sendable (AVAudioPCMBuffer, AVAudioTime) -> Void { return { buffer, _ in @@ -739,6 +752,7 @@ public final class FlowContinuousCapture { let phase = gate.withLock { $0 } switch phase { case .recording, .draining: + utterancePCMStore.append(snapshot.samples) streamRelay.yield(snapshot) if phase == .draining { drainTracker.noteAudio(samples: snapshot.samples, policy: drainPolicy) diff --git a/OSGKeyboardShared/Utilities/FlowPipelineDiagnostics.swift b/OSGKeyboardShared/Utilities/FlowPipelineDiagnostics.swift index c1e3474..7a3a02d 100644 --- a/OSGKeyboardShared/Utilities/FlowPipelineDiagnostics.swift +++ b/OSGKeyboardShared/Utilities/FlowPipelineDiagnostics.swift @@ -25,6 +25,18 @@ public enum FlowPipelineDiagnostics { OSGLog.asr.info("finalChunkRecovery \(action) chunk=\(chunkIndex)") } + public static func logBatchFallback( + sampleCount: Int, + stitchedLength: Int, + partialLength: Int, + batchLength: Int + ) { + OSGLog.flow.info( + "batchFallback samples=\(sampleCount) stitchedLen=\(stitchedLength) " + + "partialLen=\(partialLength) batchLen=\(batchLength)" + ) + } + public static func logChunkFinalize( chunkCount: Int, lastChunkSamples: Int, diff --git a/OSGKeyboardShared/Utilities/FlowUtterancePCMStore.swift b/OSGKeyboardShared/Utilities/FlowUtterancePCMStore.swift new file mode 100644 index 0000000..8adc313 --- /dev/null +++ b/OSGKeyboardShared/Utilities/FlowUtterancePCMStore.swift @@ -0,0 +1,46 @@ +// FlowUtterancePCMStore.swift +// OSGKeyboard · Shared +// +// Thread-safe rolling buffer of 16 kHz mono utterance PCM for whole-utterance +// batch ASR fallback when pipelined chunking drops weak tail segments. + +import Foundation + +public final class FlowUtterancePCMStore: @unchecked Sendable { + private let lock = OSAllocatedUnfairLock() + private var samples: [Float] = [] + private let maxSampleCount: Int + + public init(maxSampleCount: Int) { + self.maxSampleCount = max(1, maxSampleCount) + } + + public func reset() { + lock.withLock { + samples.removeAll(keepingCapacity: false) + } + } + + public func append(_ chunk: [Float]) { + guard !chunk.isEmpty else { return } + lock.withLock { + samples.append(contentsOf: chunk) + if samples.count > maxSampleCount { + samples.removeFirst(samples.count - maxSampleCount) + } + } + } + + public var sampleCount: Int { + lock.withLock { samples.count } + } + + /// Returns accumulated samples and clears the store. + public func consume() -> [Float] { + lock.withLock { + let out = samples + samples.removeAll(keepingCapacity: false) + return out + } + } +} diff --git a/OSGKeyboardShared/Utilities/UtteranceBatchFallbackPolicy.swift b/OSGKeyboardShared/Utilities/UtteranceBatchFallbackPolicy.swift new file mode 100644 index 0000000..a29a632 --- /dev/null +++ b/OSGKeyboardShared/Utilities/UtteranceBatchFallbackPolicy.swift @@ -0,0 +1,41 @@ +// UtteranceBatchFallbackPolicy.swift +// OSGKeyboard · Shared +// +// Decides when to re-run ASR on the full utterance PCM after pipelined +// chunking, and how to pick the best transcript among candidates. + +import Foundation + +public enum UtteranceBatchFallbackPolicy { + public static let defaultCharacterAdvantage = UtteranceTranscriptGuard.defaultPartialAdvantage + + /// True when chunked output likely lost trailing content vs the live partial. + public static func shouldRunBatchFallback( + stitchedFinal: String, + partialSnapshot: String, + minimumCharacterAdvantage: Int = defaultCharacterAdvantage + ) -> Bool { + let final = stitchedFinal.trimmingCharacters(in: .whitespacesAndNewlines) + let partial = partialSnapshot.trimmingCharacters(in: .whitespacesAndNewlines) + + if final.isEmpty, !partial.isEmpty { return true } + if partial.isEmpty { return false } + return partial.count >= final.count + minimumCharacterAdvantage + } + + /// Prefer the longest non-empty transcript after batch ASR completes. + public static func preferredTranscript( + batch: String, + stitchedFinal: String, + partialSnapshot: String, + current: String + ) -> String { + let candidates = [batch, current, stitchedFinal, partialSnapshot] + .map { $0.trimmingCharacters(in: .whitespacesAndNewlines) } + .filter { !$0.isEmpty } + guard let best = candidates.max(by: { $0.count < $1.count }) else { + return current.trimmingCharacters(in: .whitespacesAndNewlines) + } + return best + } +} diff --git a/OSGKeyboardTests/FlowUtterancePCMStoreTests.swift b/OSGKeyboardTests/FlowUtterancePCMStoreTests.swift new file mode 100644 index 0000000..e4f4af6 --- /dev/null +++ b/OSGKeyboardTests/FlowUtterancePCMStoreTests.swift @@ -0,0 +1,23 @@ +// FlowUtterancePCMStoreTests.swift +// OSGKeyboardTests + +import XCTest +@testable import OSGKeyboardShared + +final class FlowUtterancePCMStoreTests: XCTestCase { + + func testAppendAndConsume() { + let store = FlowUtterancePCMStore(maxSampleCount: 100) + store.append([1, 2, 3]) + store.append([4, 5]) + XCTAssertEqual(store.sampleCount, 5) + XCTAssertEqual(store.consume(), [1, 2, 3, 4, 5]) + XCTAssertEqual(store.sampleCount, 0) + } + + func testTrimsOldestWhenOverCap() { + let store = FlowUtterancePCMStore(maxSampleCount: 4) + store.append([1, 2, 3, 4, 5]) + XCTAssertEqual(store.consume(), [2, 3, 4, 5]) + } +} diff --git a/OSGKeyboardTests/UtteranceBatchFallbackPolicyTests.swift b/OSGKeyboardTests/UtteranceBatchFallbackPolicyTests.swift new file mode 100644 index 0000000..0a77ad7 --- /dev/null +++ b/OSGKeyboardTests/UtteranceBatchFallbackPolicyTests.swift @@ -0,0 +1,45 @@ +// UtteranceBatchFallbackPolicyTests.swift +// OSGKeyboardTests + +import XCTest +@testable import OSGKeyboardShared + +final class UtteranceBatchFallbackPolicyTests: XCTestCase { + + func testShouldRunWhenPartialClearlyLonger() { + XCTAssertTrue( + UtteranceBatchFallbackPolicy.shouldRunBatchFallback( + stitchedFinal: "今天很好", + partialSnapshot: "今天很好，我们一起去公园吧" + ) + ) + } + + func testShouldRunWhenFinalEmptyButPartialPresent() { + XCTAssertTrue( + UtteranceBatchFallbackPolicy.shouldRunBatchFallback( + stitchedFinal: "", + partialSnapshot: "最后一段" + ) + ) + } + + func testShouldNotRunWhenPartialNotLonger() { + XCTAssertFalse( + UtteranceBatchFallbackPolicy.shouldRunBatchFallback( + stitchedFinal: "今天很好，我们一起去公园吧", + partialSnapshot: "今天很好" + ) + ) + } + + func testPreferredTranscriptPicksLongestCandidate() { + let resolved = UtteranceBatchFallbackPolicy.preferredTranscript( + batch: "今天很好，我们一起去公园吧", + stitchedFinal: "今天很好", + partialSnapshot: "今天很好，我们", + current: "今天很好，我们" + ) + XCTAssertEqual(resolved, "今天很好，我们一起去公园吧") + } +}