@mentra/bluetooth-sdk 3.3.0-dev.516 → 3.3.0-dev.523

package/ios/Source/stt/SherpaOnnxTranscriber.swift

@@ -1,5 +1,4 @@
 import Foundation
-import UIKit
 
 /**
  * SherpaOnnxTranscriber handles real-time audio transcription using Sherpa-ONNX.
@@ -11,24 +10,30 @@ class SherpaOnnxTranscriber {
     private static let TAG = "SherpaOnnxTranscriber"
 
     private static let SAMPLE_RATE = 16000 // Sherpa-ONNX model's required sample rate
-    private static let QUEUE_CAPACITY = 100 // Max number of audio buffers to keep in queue
+    private final class Session {
+        let recognizer: SherpaOnnxRecognizer
+        var lastPartialResult = ""
 
-    private let pcmQueue = DispatchQueue(label: "com.augmentos.sherpaonnx.pcmQueue", qos: .userInteractive)
-    private var pcmBuffers = [Data]()
-    private var isRunning = false
-    private var processingQueue: DispatchQueue?
-    private var processingTask: DispatchWorkItem?
-
-    /// The underlying Sherpa-ONNX objects
-    private var recognizer: SherpaOnnxRecognizer?
+        init(_ recognizer: SherpaOnnxRecognizer) {
+            self.recognizer = recognizer
+        }
+    }
 
-    private var lastPartialResult = ""
+    private lazy var worker = DemandDrivenAudioWorker<Session>(
+        makeProcessor: { Self.makeRecognizer().map(Session.init) },
+        process: { [weak self] session, data, generation in
+            self?.process(session, data: data, generation: generation)
+        }
+    )
 
-    /// Parent context
-    private weak var context: UIViewController?
+    /// Construct the worker on the manager's actor before audio can arrive.
+    init() {
+        _ = worker
+    }
 
-    /// Session start time for relative timestamps
-    private var transcriptionSessionStart: Date
+    func setActive(_ active: Bool) {
+        worker.setActive(active)
+    }
 
     /// Dynamic model path support
     private static var customModelPath: String? {
@@ -65,23 +70,9 @@ class SherpaOnnxTranscriber {
         return nil
     }
 
-    /**
-     * Constructor that accepts a UIViewController to load model assets.
-     */
-    init(context: UIViewController) {
-        self.context = context
-        transcriptionSessionStart = Date()
-    }
-
-    deinit {
-        shutdown()
-    }
-
-    /**
-     * Initialize the Sherpa-ONNX recognizer.
-     * Loads models and configuration, sets up processing thread.
-     */
-    func initialize() {
+    /// Runs exclusively on the worker queue, never on the main/audio delivery thread.
+    private static func makeRecognizer() -> SherpaOnnxRecognizer? {
+        var recognizer: SherpaOnnxRecognizer?
         do {
             var tokensPath: String
             var modelType = "unknown"
@@ -201,30 +192,26 @@ class SherpaOnnxTranscriber {
                 Bridge.log("No Sherpa ONNX model available. Transcription will be disabled.")
                 Bridge.log("Please download a model using the model downloader in settings.")
                 recognizer = nil
-                isRunning = false
-                return
+                return nil
             }
 
             if recognizer == nil {
                 throw NSError(domain: "SherpaOnnxTranscriber", code: 2, userInfo: [NSLocalizedDescriptionKey: "Failed to create recognizer"])
             }
 
-            startProcessingTask()
-            isRunning = true
-
             Bridge.log("Sherpa-ONNX ASR initialized successfully with \(modelType) model")
+            return recognizer
 
         } catch {
             Bridge.log("Failed to initialize Sherpa-ONNX: \(error.localizedDescription)")
+            return nil
         }
     }
 
-    /**
-     * Handle transcription results - send only to delegate
-     */
-    private func handleTranscriptionResult(text: String, isFinal: Bool) {
-        // Forward to delegate if set
+    private func handleTranscriptionResult(text: String, isFinal: Bool, generation: UInt64) {
         DispatchQueue.main.async { [weak self] in
+            // A decode or main-queue delivery may finish after stop/model replacement.
+            guard let self, self.worker.isCurrent(generation) else { return }
             if isFinal {
                 STTTools.didReceiveFinalTranscription(text)
             } else {
@@ -233,123 +220,30 @@ class SherpaOnnxTranscriber {
         }
     }
 
-    /**
-     * Feed PCM audio data (16-bit little endian) into the transcriber.
-     * This method should be called continuously with short chunks (e.g., 100-300ms).
-     *
-     * Audio is queued directly; microphone VAD gating is not applied in the SDK.
-     */
     func acceptAudio(pcm16le: Data) {
-        guard isRunning else {
-            return
-        }
-
-        queueAudioData(pcm16le)
-    }
-
-    private func queueAudioData(_ pcm16le: Data) {
-        pcmQueue.async { [weak self] in
-            guard let self = self else { return }
-
-            let queueSizeBefore = self.pcmBuffers.count
-            self.pcmBuffers.append(pcm16le)
-
-            // Keep queue size manageable
-            if self.pcmBuffers.count > Self.QUEUE_CAPACITY {
-                let removedBuffer = self.pcmBuffers.removeFirst()
-                Bridge.log("⚠️ Audio queue overflow - dropped buffer of \(removedBuffer.count) bytes")
-            }
-        }
+        worker.accept(pcm16le)
     }
 
-    /**
-     * Start a background task to continuously consume audio and decode using Sherpa.
-     */
-    private func startProcessingTask() {
-        Bridge.log("🚀 Starting Sherpa-ONNX processing task...")
-
-        processingQueue = DispatchQueue(label: "com.augmentos.sherpaonnx.processor", qos: .userInitiated)
-
-        let workItem = DispatchWorkItem { [weak self] in
-            self?.runLoop()
+    private func process(_ session: Session, data: Data, generation: UInt64) {
+        guard worker.isCurrent(generation) else { return }
+        let recognizer = session.recognizer
+        recognizer.acceptWaveform(samples: toFloatArray(from: data), sampleRate: Self.SAMPLE_RATE)
+        while recognizer.isReady() {
+            guard worker.isCurrent(generation) else { return }
+            recognizer.decode()
         }
-
-        processingTask = workItem
-        processingQueue?.async(execute: workItem)
-    }
-
-    /**
-     * Main processing loop that handles transcription in real-time.
-     * Pulls audio from queue, feeds into Sherpa, emits partial/final results.
-     */
-    private func runLoop() {
-        Bridge.log("🔄 Sherpa-ONNX processing loop started")
-
-        while isRunning {
-            // Pull data from queue
-            var audioData: Data?
-
-            pcmQueue.sync {
-                if !self.pcmBuffers.isEmpty {
-                    audioData = self.pcmBuffers.removeFirst()
-                }
-            }
-
-            if let data = audioData {
-                // Synchronize access to recognizer to prevent race conditions
-                objc_sync_enter(self)
-                defer { objc_sync_exit(self) }
-
-                guard let recognizer = recognizer else {
-                    Bridge.log("⚠️ Recognizer not available, skipping audio chunk")
-                    continue
-                }
-
-                do {
-                    // Convert PCM to float [-1.0, 1.0]
-                    let floatBuf = toFloatArray(from: data)
-
-                    // Pass audio data to the Sherpa-ONNX stream
-                    recognizer.acceptWaveform(samples: floatBuf, sampleRate: Self.SAMPLE_RATE)
-
-                    // Decode continuously while model is ready
-                    var decodeCount = 0
-                    while recognizer.isReady() {
-                        recognizer.decode()
-                        decodeCount += 1
-                    }
-
-                    // If utterance endpoint detected
-                    if recognizer.isEndpoint() {
-                        let result = recognizer.getResult()
-                        let finalText = result.text.trimmingCharacters(in: .whitespacesAndNewlines)
-
-                        if !finalText.isEmpty {
-                            handleTranscriptionResult(text: finalText, isFinal: true)
-                        }
-
-                        recognizer.reset() // Start new utterance
-                        lastPartialResult = ""
-                    } else {
-                        // Emit partial results if changed
-                        let result = recognizer.getResult()
-                        let partial = result.text.trimmingCharacters(in: .whitespacesAndNewlines)
-
-                        if partial != lastPartialResult, !partial.isEmpty {
-                            handleTranscriptionResult(text: partial, isFinal: false)
-                            lastPartialResult = partial
-                        }
-                    }
-                } catch {
-                    Bridge.log("❌ Error processing audio: \(error.localizedDescription)")
-                }
-            } else {
-                // Sleep briefly to avoid tight CPU loop if no audio is available
-                Thread.sleep(forTimeInterval: 0.01)
+        guard worker.isCurrent(generation) else { return }
+        let text = recognizer.getResult().text.trimmingCharacters(in: .whitespacesAndNewlines)
+        if recognizer.isEndpoint() {
+            if !text.isEmpty {
+                handleTranscriptionResult(text: text, isFinal: true, generation: generation)
             }
+            recognizer.reset()
+            session.lastPartialResult = ""
+        } else if !text.isEmpty, text != session.lastPartialResult {
+            handleTranscriptionResult(text: text, isFinal: false, generation: generation)
+            session.lastPartialResult = text
         }
-
-        Bridge.log("ASR processing thread stopped")
     }
 
     /**
@@ -377,45 +271,11 @@ class SherpaOnnxTranscriber {
         return samples
     }
 
-    /**
-     * Stop transcription processing.
-     * This shuts down the processing thread and releases Sherpa-ONNX resources.
-     */
     func shutdown() {
-        Bridge.log("🛑 Shutting down SherpaOnnxTranscriber...")
-
-        isRunning = false
-        processingTask?.cancel()
-
-        // Synchronize access to recognizer during shutdown
-        objc_sync_enter(self)
-        defer { objc_sync_exit(self) }
-
-        // The recognizer will be automatically cleaned up by ARC when set to nil
-        if recognizer != nil {
-            Bridge.log("🧹 Cleaning up Sherpa-ONNX recognizer")
-            recognizer = nil
-        }
-
-        // Clear any remaining audio buffers
-        pcmQueue.sync {
-            let remainingBuffers = self.pcmBuffers.count
-            if remainingBuffers > 0 {
-                Bridge.log("🗑️ Clearing \(remainingBuffers) remaining audio buffers")
-            }
-            self.pcmBuffers.removeAll()
-        }
-
-        Bridge.log("✅ SherpaOnnxTranscriber shutdown complete")
+        worker.setActive(false)
     }
 
-    /**
-     * Restarts the transcriber after a model change.
-     * Shuts down existing resources, clears buffers, and reinitializes the recognizer.
-     */
     func restart() {
-        Bridge.log("♻️ Restarting SherpaOnnxTranscriber...")
-        shutdown()
-        initialize()
+        worker.restart()
     }
 }

package/ios/Source/utils/DemandDrivenAudioWorker.swift

@@ -0,0 +1,150 @@
+import Foundation
+
+/// Owns a processor on one serial queue. Empty queues schedule no recurring work.
+/// Demand changes invalidate queued audio/results immediately; creation, processing,
+/// and destruction finish in order without blocking the caller (including the UI).
+final class DemandDrivenAudioWorker<Processor> {
+    struct Snapshot {
+        let active: Bool
+        let initialized: Bool
+        let processing: Bool
+        let queuedBuffers: Int
+        let processedBuffers: UInt64
+        let droppedBuffers: UInt64
+    }
+
+    private let lock = NSLock()
+    private let queue: DispatchQueue
+    private let capacity: Int
+    private let makeProcessor: () -> Processor?
+    private let process: (Processor, Data, UInt64) -> Void
+    private var active = false
+    private var generation: UInt64 = 0
+    private var buffers = [Data]()
+    private var scheduled = false
+    private var initialized = false
+    private var processing = false
+    private var processedBuffers: UInt64 = 0
+    private var droppedBuffers: UInt64 = 0
+
+    // Accessed only on queue, including release of the previous processor.
+    private var processor: Processor?
+    private var processorGeneration: UInt64?
+    private var attemptedInitialization = false
+
+    init(capacity: Int = 100,
+         queue: DispatchQueue = DispatchQueue(label: "com.mentra.speech.processor", qos: .userInitiated),
+         makeProcessor: @escaping () -> Processor?,
+         process: @escaping (Processor, Data, UInt64) -> Void)
+    {
+        precondition(capacity > 0)
+        self.capacity = capacity
+        self.queue = queue
+        self.makeProcessor = makeProcessor
+        self.process = process
+    }
+
+    func setActive(_ value: Bool) {
+        lock.lock()
+        defer { lock.unlock() }
+        guard active != value else { return }
+        active = value
+        invalidateLocked()
+    }
+
+    /// Reload a changed model only if there is still local-transcription demand.
+    func restart() {
+        lock.lock()
+        defer { lock.unlock() }
+        invalidateLocked()
+    }
+
+    private func invalidateLocked() {
+        generation &+= 1
+        buffers.removeAll(keepingCapacity: true)
+        initialized = false
+        scheduleLocked()
+    }
+
+    func accept(_ data: Data) {
+        guard !data.isEmpty else { return }
+        lock.lock()
+        defer { lock.unlock() }
+        guard active else { return }
+        // Bound even audio arriving while model loading or decoding is blocked.
+        if buffers.count == capacity {
+            buffers.removeFirst()
+            droppedBuffers &+= 1
+        }
+        buffers.append(data)
+        scheduleLocked()
+    }
+
+    func isCurrent(_ token: UInt64) -> Bool {
+        lock.lock()
+        defer { lock.unlock() }
+        return active && generation == token
+    }
+
+    func snapshot() -> Snapshot {
+        lock.lock()
+        defer { lock.unlock() }
+        return Snapshot(active: active, initialized: initialized, processing: processing,
+                        queuedBuffers: buffers.count, processedBuffers: processedBuffers,
+                        droppedBuffers: droppedBuffers)
+    }
+
+    private func scheduleLocked() {
+        guard !scheduled else { return }
+        scheduled = true
+        queue.async { self.drain() }
+    }
+
+    private func drain() {
+        while true {
+            lock.lock()
+            let token = generation
+            if processorGeneration != token {
+                lock.unlock()
+                // ARC destruction stays serialized with inference and model loading.
+                processor = nil
+                processorGeneration = token
+                attemptedInitialization = false
+                continue
+            }
+            guard active else {
+                scheduled = false
+                lock.unlock()
+                return
+            }
+            if !attemptedInitialization {
+                attemptedInitialization = true
+                lock.unlock()
+                processor = makeProcessor()
+                lock.lock()
+                if generation == token { initialized = processor != nil }
+                lock.unlock()
+                continue
+            }
+            guard let processor else {
+                buffers.removeAll(keepingCapacity: true)
+                scheduled = false
+                lock.unlock()
+                return // Missing model: retry on a demand transition or model restart.
+            }
+            guard !buffers.isEmpty else {
+                scheduled = false
+                lock.unlock()
+                return
+            }
+            let data = buffers.removeFirst()
+            processing = true
+            lock.unlock()
+            process(processor, data, token)
+            lock.lock()
+            processedBuffers &+= 1
+            processing = false
+            lock.unlock()
+        }
+    }
+}

package/ios/Tests/DemandDrivenAudioWorkerTests.swift

@@ -0,0 +1,157 @@
+import Foundation
+@testable import MentraBluetoothSDK
+import XCTest
+
+final class DemandDrivenAudioWorkerTests: XCTestCase {
+    private final class Processor {
+        let onRelease: () -> Void
+        init(onRelease: @escaping () -> Void = {}) {
+            self.onRelease = onRelease
+        }
+
+        deinit { onRelease() }
+    }
+
+    func testNoDemandAndEmptyActiveQueueScheduleNoRecurringWork() {
+        let queue = DispatchQueue(label: "test.speech")
+        var creations = 0
+        var processed = [Data]()
+        let worker = DemandDrivenAudioWorker(queue: queue, makeProcessor: {
+            creations += 1
+            return Processor()
+        }, process: { _, data, _ in processed.append(data) })
+        worker.accept(Data([1]))
+        queue.sync {}
+        XCTAssertEqual(creations, 0)
+        worker.setActive(true)
+        worker.setActive(true)
+        queue.sync {}
+        XCTAssertEqual(creations, 1)
+        XCTAssertTrue(worker.snapshot().initialized)
+        XCTAssertFalse(worker.snapshot().processing)
+        // A serial barrier completes with an empty queue; no sleeping consumer holds it.
+        queue.sync {}
+        XCTAssertTrue(processed.isEmpty)
+        worker.setActive(false)
+        queue.sync {}
+        XCTAssertFalse(worker.snapshot().initialized)
+    }
+
+    func testFirstAudioIsRetainedDuringAsynchronousInitializationAndBounded() {
+        let queue = DispatchQueue(label: "test.speech")
+        let loading = DispatchSemaphore(value: 0)
+        let release = DispatchSemaphore(value: 0)
+        var processed = [Data]()
+        let worker = DemandDrivenAudioWorker(capacity: 3, queue: queue, makeProcessor: {
+            XCTAssertFalse(Thread.isMainThread)
+            loading.signal()
+            release.wait()
+            return Processor()
+        }, process: { _, data, _ in processed.append(data) })
+        worker.setActive(true)
+        XCTAssertEqual(loading.wait(timeout: .now() + 2), .success)
+        worker.accept(Data([1]))
+        worker.accept(Data([2]))
+        XCTAssertEqual(worker.snapshot().queuedBuffers, 2)
+        release.signal()
+        queue.sync {}
+        XCTAssertEqual(processed, [Data([1]), Data([2])])
+        worker.setActive(false)
+        queue.sync {}
+    }
+
+    func testStopDuringLoadingDropsOldAudioAndCannotResurrectProcessor() {
+        let queue = DispatchQueue(label: "test.speech")
+        let loading = DispatchSemaphore(value: 0)
+        let release = DispatchSemaphore(value: 0)
+        var released = 0
+        var processed = 0
+        let worker = DemandDrivenAudioWorker(capacity: 2, queue: queue, makeProcessor: {
+            loading.signal()
+            release.wait()
+            return Processor { released += 1 }
+        }, process: { _, _, _ in processed += 1 })
+        worker.setActive(true)
+        XCTAssertEqual(loading.wait(timeout: .now() + 2), .success)
+        for byte: UInt8 in 1 ... 10 {
+            worker.accept(Data([byte]))
+        }
+        XCTAssertEqual(worker.snapshot().queuedBuffers, 2)
+        XCTAssertEqual(worker.snapshot().droppedBuffers, 8)
+        worker.setActive(false)
+        XCTAssertEqual(worker.snapshot().queuedBuffers, 0)
+        release.signal()
+        queue.sync {}
+        XCTAssertEqual(processed, 0)
+        XCTAssertEqual(released, 1)
+        XCTAssertFalse(worker.snapshot().initialized)
+    }
+
+    func testRestartAndRapidDemandChangesSerializeOldAndNewProcessors() {
+        let queue = DispatchQueue(label: "test.speech")
+        let processing = DispatchSemaphore(value: 0)
+        let release = DispatchSemaphore(value: 0)
+        var creations = 0
+        var destructions = 0
+        var tokens = [UInt64]()
+        var dataSeen = [Data]()
+        let worker = DemandDrivenAudioWorker(queue: queue, makeProcessor: {
+            XCTAssertEqual(creations, destructions)
+            creations += 1
+            return Processor { destructions += 1 }
+        }, process: { _, data, token in
+            dataSeen.append(data)
+            tokens.append(token)
+            if data == Data([1]) {
+                processing.signal()
+                release.wait()
+            }
+        })
+        worker.setActive(true)
+        worker.accept(Data([1]))
+        XCTAssertEqual(processing.wait(timeout: .now() + 2), .success)
+        worker.accept(Data([2]))
+        worker.restart()
+        worker.setActive(false)
+        worker.setActive(true)
+        worker.accept(Data([3]))
+        release.signal()
+        queue.sync {}
+        XCTAssertEqual(dataSeen, [Data([1]), Data([3])])
+        XCTAssertFalse(worker.isCurrent(tokens[0]))
+        XCTAssertTrue(worker.isCurrent(tokens[1]))
+        XCTAssertEqual(creations, 2)
+        worker.setActive(false)
+        queue.sync {}
+        XCTAssertEqual(destructions, 2)
+        XCTAssertFalse(worker.isCurrent(tokens[1]))
+    }
+
+    func testMissingModelDoesNotRetryEveryFrameAndRestartWithoutDemandDoesNotLoad() {
+        let queue = DispatchQueue(label: "test.speech")
+        var attempts = 0
+        let worker = DemandDrivenAudioWorker<Processor>(queue: queue, makeProcessor: {
+            attempts += 1
+            return nil
+        }, process: { _, _, _ in XCTFail("Missing processor") })
+        worker.restart()
+        queue.sync {}
+        XCTAssertEqual(attempts, 0)
+        worker.setActive(true)
+        for _ in 0 ..< 5 {
+            worker.accept(Data([1]))
+            queue.sync {}
+        }
+        XCTAssertEqual(attempts, 1)
+        XCTAssertEqual(worker.snapshot().queuedBuffers, 0)
+        worker.restart()
+        queue.sync {}
+        XCTAssertEqual(attempts, 2)
+        worker.setActive(false)
+        worker.setActive(true)
+        queue.sync {}
+        XCTAssertEqual(attempts, 3)
+        worker.setActive(false)
+        queue.sync {}
+    }
+}