merge: tail-drain and chunked ASR hardening into main

Resolve FlowContinuousCapture by combining tail-drain gate/drain logic
with flow-abcd route-change and interruption recovery.
This commit is contained in:
Rocky
2026-07-06 18:21:34 +08:00
14 changed files with 696 additions and 53 deletions
@@ -14,6 +14,13 @@ import os
private enum FlowCaptureConstants {
static let levelBarCount = 24
static let targetSampleRate: Double = 16_000
static let drainPollIntervalNs: UInt64 = 20_000_000
}
private enum UtteranceGatePhase: Equatable {
case idle
case recording
case draining
}
/// Thread-safe relay for utterance-scoped ASR snapshots.
@@ -158,7 +165,14 @@ public final class FlowContinuousCapture {
private let streamRelay = FlowCaptureStreamRelay()
private let prerollStore = FlowPrerollStore()
private let levelStore = FlowLevelStore(barCount: FlowCaptureConstants.levelBarCount)
private let isUtteranceActive = OSAllocatedUnfairLock(initialState: false)
private let gate = OSAllocatedUnfairLock(initialState: UtteranceGatePhase.idle)
private let drainTracker = FlowCaptureDrainTracker()
private let tailSampleCounter = OSAllocatedUnfairLock(initialState: 0)
private var audioConverter: AVAudioConverter?
private var targetFormat: AVAudioFormat?
private var hwFormat: AVAudioFormat?
private var drainPolicy = FlowCaptureTailDrainPolicy.flowDefault
private var didInstallTap = false
private var isRunning = false
@@ -197,15 +211,15 @@ public final class FlowContinuousCapture {
}
let inputNode = audioEngine.inputNode
let hwFormat = inputNode.outputFormat(forBus: 0)
guard hwFormat.sampleRate > 0, hwFormat.channelCount > 0 else {
let hardwareFormat = inputNode.outputFormat(forBus: 0)
guard hardwareFormat.sampleRate > 0, hardwareFormat.channelCount > 0 else {
throw StartError.invalidHardwareFormat(
sampleRate: hwFormat.sampleRate,
channels: Int(hwFormat.channelCount)
sampleRate: hardwareFormat.sampleRate,
channels: Int(hardwareFormat.channelCount)
)
}
guard let targetFormat = AVAudioFormat(
guard let resolvedTargetFormat = AVAudioFormat(
commonFormat: .pcmFormatFloat32,
sampleRate: FlowCaptureConstants.targetSampleRate,
channels: 1,
@@ -213,25 +227,40 @@ public final class FlowContinuousCapture {
) else {
throw StartError.formatCreateFailed
}
guard let converter = AVAudioConverter(from: hwFormat, to: targetFormat) else {
guard let converter = AVAudioConverter(from: hardwareFormat, to: resolvedTargetFormat) else {
throw StartError.converterCreateFailed
}
audioConverter = converter
targetFormat = resolvedTargetFormat
hwFormat = hardwareFormat
// Rebuild the tap so its bound hardware format matches the new route.
if didInstallTap {
inputNode.removeTap(onBus: 0)
didInstallTap = false
}
let gateLock = gate
let relay = streamRelay
let preroll = prerollStore
let levels = levelStore
let tracker = drainTracker
let tailCounter = tailSampleCounter
let policy = drainPolicy
let tap = Self.makeAudioTapBlock(
converter: converter,
targetFormat: targetFormat,
hwFormat: hwFormat,
utteranceFlag: isUtteranceActive,
levelStore: levelStore,
prerollStore: prerollStore,
streamRelay: streamRelay
targetFormat: resolvedTargetFormat,
hwFormat: hardwareFormat,
gate: gateLock,
levelStore: levels,
prerollStore: preroll,
streamRelay: relay,
drainTracker: tracker,
tailSampleCounter: tailCounter,
drainPolicy: policy
)
inputNode.installTap(onBus: 0, bufferSize: 4096, format: hwFormat, block: tap)
inputNode.installTap(onBus: 0, bufferSize: 4096, format: hardwareFormat, block: tap)
didInstallTap = true
audioEngine.prepare()
@@ -245,7 +274,9 @@ public final class FlowContinuousCapture {
/// Tear down the engine and release the audio session.
public func stop() {
removeSessionObservers()
isUtteranceActive.withLock { $0 = false }
gate.withLock { $0 = .idle }
drainTracker.reset()
tailSampleCounter.withLock { $0 = 0 }
streamRelay.finish()
if didInstallTap {
@@ -256,6 +287,9 @@ public final class FlowContinuousCapture {
audioEngine.stop()
}
isRunning = false
audioConverter = nil
targetFormat = nil
hwFormat = nil
try? AVAudioSession.sharedInstance().setActive(
false,
options: .notifyOthersOnDeactivation
@@ -288,8 +322,6 @@ public final class FlowContinuousCapture {
object: nil,
queue: .main
) { [weak self] note in
// Extract Sendable primitives here (Notification isn't Sendable)
// before hopping onto the main actor.
let reasonRaw = note.userInfo?[AVAudioSessionRouteChangeReasonKey] as? UInt
MainActor.assumeIsolated { self?.handleRouteChange(reasonRaw: reasonRaw) }
}
@@ -321,8 +353,6 @@ public final class FlowContinuousCapture {
guard isRunning else { return }
guard let reasonRaw,
let reason = AVAudioSession.RouteChangeReason(rawValue: reasonRaw) else { return }
// Only rebuild for real device swaps (plugging / unplugging a headset
// or AirPods). Ignore `.categoryChange`, which we trigger ourselves.
switch reason {
case .oldDeviceUnavailable, .newDeviceAvailable:
log.info("Audio route changed (\(reasonRaw, privacy: .public)) — rebuilding engine")
@@ -337,7 +367,6 @@ public final class FlowContinuousCapture {
let type = AVAudioSession.InterruptionType(rawValue: typeRaw) else { return }
switch type {
case .began:
// The system already paused our engine; wait for `.ended`.
log.info("Audio interruption began")
case .ended:
guard isRunning else { return }
@@ -373,38 +402,129 @@ public final class FlowContinuousCapture {
/// Begin forwarding downsampled buffers to ASR for one utterance.
public func beginUtterance() -> AsyncStream<AudioBufferSnapshot> {
let (stream, continuation) = AsyncStream<AudioBufferSnapshot>.makeStream()
drainTracker.reset()
tailSampleCounter.withLock { $0 = 0 }
// Bind the consumer before opening the gate so early tap frames
// are not dropped on the floor.
streamRelay.bind(continuation)
streamRelay.replay(prerollStore.drain())
isUtteranceActive.withLock { $0 = true }
gate.withLock { $0 = .recording }
return stream
}
/// Stop forwarding buffers; finishes the ASR stream.
public func endUtterance() {
isUtteranceActive.withLock { $0 = false }
/// Drain trailing PCM after the user stops, then finish the ASR stream.
public func endUtteranceAndDrain(
policy: FlowCaptureTailDrainPolicy = .flowDefault
) async -> FlowCaptureDrainReport {
let currentPhase = gate.withLock { $0 }
guard currentPhase == .recording else {
return .skipped
}
drainPolicy = policy
gate.withLock { $0 = .draining }
drainTracker.beginDrain()
var endedBySilence = false
while true {
let decision = drainTracker.shouldFinish(policy: policy)
if decision.finished {
endedBySilence = decision.endedBySilence
break
}
if Task.isCancelled { break }
try? await Task.sleep(nanoseconds: FlowCaptureConstants.drainPollIntervalNs)
}
let flushSamples = flushConverterTailToStream()
tailSampleCounter.withLock { $0 += flushSamples }
streamRelay.finish()
gate.withLock { $0 = .idle }
let tailSamples = tailSampleCounter.withLock { $0 }
let report = FlowCaptureDrainReport(
drainDurationSeconds: drainTracker.elapsedSeconds(),
endedBySilence: endedBySilence,
tailSampleCount: tailSamples
)
drainTracker.reset()
tailSampleCounter.withLock { $0 = 0 }
FlowPipelineDiagnostics.logDrain(report)
return report
}
/// Immediate stop without tail drain (abort / session teardown).
public func cancelUtterance() {
endUtterance()
gate.withLock { $0 = .idle }
drainTracker.reset()
tailSampleCounter.withLock { $0 = 0 }
streamRelay.finish()
}
public func currentAudioLevels() -> [Float] {
levelStore.snapshot()
}
// MARK: - Converter flush
@discardableResult
private func flushConverterTailToStream() -> Int {
guard let converter = audioConverter,
let targetFormat,
let hwFormat else {
return 0
}
var flushedSamples = 0
let capacity = AVAudioFrameCount(max(512, hwFormat.sampleRate / 20))
guard let outBuffer = AVAudioPCMBuffer(pcmFormat: targetFormat, frameCapacity: capacity) else {
return 0
}
var endOfStreamSignaled = false
while true {
outBuffer.frameLength = 0
var error: NSError?
let status = converter.convert(to: outBuffer, error: &error) { _, outStatus in
if endOfStreamSignaled {
outStatus.pointee = .noDataNow
return nil
}
endOfStreamSignaled = true
outStatus.pointee = .endOfStream
return nil
}
if status == .error || error != nil {
break
}
guard status == .haveData, outBuffer.frameLength > 0 else {
break
}
let snapshot = AudioBufferSnapshot(buffer: outBuffer)
guard !snapshot.samples.isEmpty else { break }
streamRelay.yield(snapshot)
flushedSamples += snapshot.samples.count
}
return flushedSamples
}
// MARK: - Audio tap (nonisolated runs on realtime thread)
private nonisolated static func makeAudioTapBlock(
converter: AVAudioConverter,
targetFormat: AVAudioFormat,
hwFormat: AVAudioFormat,
utteranceFlag: OSAllocatedUnfairLock<Bool>,
gate: OSAllocatedUnfairLock<UtteranceGatePhase>,
levelStore: FlowLevelStore,
prerollStore: FlowPrerollStore,
streamRelay: FlowCaptureStreamRelay
streamRelay: FlowCaptureStreamRelay,
drainTracker: FlowCaptureDrainTracker,
tailSampleCounter: OSAllocatedUnfairLock<Int>,
drainPolicy: FlowCaptureTailDrainPolicy
) -> @Sendable (AVAudioPCMBuffer, AVAudioTime) -> Void {
return { buffer, _ in
levelStore.update(from: buffer, barCount: FlowCaptureConstants.levelBarCount)
@@ -426,9 +546,15 @@ public final class FlowContinuousCapture {
let snapshot = AudioBufferSnapshot(buffer: outBuffer)
guard !snapshot.samples.isEmpty else { return }
if utteranceFlag.withLock({ $0 }) {
let phase = gate.withLock { $0 }
switch phase {
case .recording, .draining:
streamRelay.yield(snapshot)
} else {
if phase == .draining {
drainTracker.noteAudio(samples: snapshot.samples, policy: drainPolicy)
tailSampleCounter.withLock { $0 += snapshot.samples.count }
}
case .idle:
prerollStore.append(snapshot)
}
}