fix: resilient multi-stage audio retry, mid-stream disconnect recovery, and intelligent token stitching

src/AppDelegate.swift

@@ -137,6 +137,7 @@ public final class AppDelegate: NSObject, NSApplicationDelegate {
         
         // 2. Acquire pre-warmed / fast-connected WebSocket session
         let session = SonioxSessionPool.shared.acquireSession()
+        session.setRecordingActive(true)
         self.activeSession = session
         
         // Handle network degradation warning
@@ -160,6 +161,13 @@ public final class AppDelegate: NSObject, NSApplicationDelegate {
         session.onFinalResult = { [weak self] result in
             DispatchQueue.main.async {
                 guard let self = self else { return }
+                // If user is STILL speaking (Option held / recording active), do NOT finish early!
+                if self.audioRecorder.isRecording {
+                    print("AppDelegate: Session finished/failed while user is still speaking. Marking degraded; recovery will run upon release.")
+                    self.isNetworkDegraded = true
+                    HUDOverlayController.shared.show(state: .recording(level: self.currentAudioLevel, isLagging: true))
+                    return
+                }
                 self.handleFinalResult(result)
             }
         }
@@ -201,6 +209,7 @@ public final class AppDelegate: NSObject, NSApplicationDelegate {
         HotkeyManager.shared.resetKeyState()
         guard audioRecorder.isRecording else { return }
         
+        activeSession?.setRecordingActive(false)
         let pcmData = audioRecorder.stopRecording()
         self.lastRecordedAudio = pcmData
         
@@ -220,24 +229,31 @@ public final class AppDelegate: NSObject, NSApplicationDelegate {
         playSystemSound(name: "Pop")
         statusBarController.updateIcon(state: .transcribing)
         statusBarController.buildMenu(isRecording: false)
-        HUDOverlayController.shared.show(state: .transcribing)
         
         isBusyFinalizing = true
         
+        // Check if the live session suffered a network failure during recording
+        let sessionDropped = (self.activeSession == nil) || (self.activeSession?.hasNetworkFailure == true) || self.isNetworkDegraded
+        if sessionDropped {
+            print("AppDelegate: Live session dropped/degraded during recording. Triggering Smart Audio Recovery immediately.")
+            startSmartAudioRecovery(pcmData: pcmData, reason: "Live session dropped mid-recording")
+            return
+        }
+        
+        HUDOverlayController.shared.show(state: .transcribing)
+        
         // Safety timeout watchdog: finalize must complete within 2.5s
         DispatchQueue.main.asyncAfter(deadline: .now() + 2.5) { [weak self] in
             guard let self = self, self.isBusyFinalizing else { return }
-            print("AppDelegate: Finalize safety watchdog triggered — recovering transcript.")
-            self.isBusyFinalizing = false
-            self.statusBarController.updateIcon(state: .idle)
-            HUDOverlayController.shared.hide()
-            self.activeSession?.completeWithCurrentText()
+            print("AppDelegate: Finalize safety watchdog triggered — starting Smart Audio Recovery.")
+            self.startSmartAudioRecovery(pcmData: self.lastRecordedAudio, reason: "Finalize timeout watchdog (2.5s)")
         }
         
         activeSession?.finalizeStream()
     }
     
     private func handleFinalResult(_ result: Result<String, Error>) {
+        guard isBusyFinalizing else { return }
         self.activeSession = nil
         self.statusBarController.updateIcon(state: .idle)
         
@@ -252,33 +268,12 @@ public final class AppDelegate: NSObject, NSApplicationDelegate {
         }
         
         let hasSpeech = hasAudibleSpeech(pcmData: self.lastRecordedAudio)
-        
-        // Check if smart audio recovery is needed:
-        // STRICT: Only retry if there was an ACTUAL network failure/drop AND audible speech was recorded.
-        // NEVER retry when Soniox cleanly returned empty text for silence!
-        let needsRetry = (isNetworkError || self.isNetworkDegraded) && hasSpeech && self.lastRecordedAudio.count > 16000
+        let isTruncated = hasSpeech && self.lastRecordedAudio.count > 64000 && rawText.count < 10
+        let needsRetry = (isNetworkError || self.isNetworkDegraded || isTruncated || rawText.isEmpty) && hasSpeech && self.lastRecordedAudio.count > 16000
         
         if needsRetry {
-            print("AppDelegate: Network drop detected with audible voice. Triggering Smart Audio Recovery with \(self.lastRecordedAudio.count) bytes...")
-            HUDOverlayController.shared.show(state: .retrying)
-            
-            let audioCopy = self.lastRecordedAudio
-            let hints = SonioxSettings.shared.language.hints
-            let target = SonioxSettings.shared.targetLanguage.isEnabled ? SonioxSettings.shared.targetLanguage.rawValue : nil
-            
-            SonioxLiveSession.retryFullBufferedAudio(pcmData: audioCopy, languageHints: hints, targetLanguage: target, onProgress: nil) { [weak self] retryResult in
-                guard let self = self else { return }
-                self.isBusyFinalizing = false
-                
-                switch retryResult {
-                case .success(let recoveredText):
-                    let trimmed = recoveredText.trimmingCharacters(in: .whitespacesAndNewlines)
-                    let finalText = trimmed.isEmpty ? rawText : trimmed
-                    self.deliverOutputText(finalText)
-                case .failure:
-                    self.deliverOutputText(rawText)
-                }
-            }
+            let reason = isNetworkError ? "WebSocket network error" : (isTruncated ? "Truncated transcript (\(rawText.count) chars for \(self.lastRecordedAudio.count) bytes)" : "Empty transcript for audible speech")
+            startSmartAudioRecovery(pcmData: self.lastRecordedAudio, reason: reason)
             return
         }
         
@@ -286,6 +281,66 @@ public final class AppDelegate: NSObject, NSApplicationDelegate {
         deliverOutputText(rawText)
     }
     
+    private func startSmartAudioRecovery(pcmData: Data, reason: String) {
+        guard !pcmData.isEmpty else {
+            self.isBusyFinalizing = false
+            self.statusBarController.updateIcon(state: .idle)
+            HUDOverlayController.shared.hide()
+            return
+        }
+        
+        let durationSec = Double(pcmData.count) / 32000.0
+        print("AppDelegate: Starting Smart Audio Recovery [\(reason)] with \(pcmData.count) bytes (\(String(format: "%.1f", durationSec))s)...")
+        
+        // Show glowing amber recovery state on HUD
+        HUDOverlayController.shared.show(state: .retrying)
+        statusBarController.updateIcon(state: .transcribing)
+        
+        // Cancel active session to prevent resource leaks
+        self.activeSession?.cancel()
+        self.activeSession = nil
+        
+        // Concurrently prime a fresh session pool in background for the NEXT speech
+        SonioxSessionPool.shared.reconnectAll()
+        
+        let hints = SonioxSettings.shared.language.hints
+        let target = SonioxSettings.shared.targetLanguage.isEnabled ? SonioxSettings.shared.targetLanguage.rawValue : nil
+        
+        SonioxLiveSession.retryFullBufferedAudio(pcmData: pcmData, languageHints: hints, targetLanguage: target, onProgress: nil) { [weak self] retryResult in
+            guard let self = self else { return }
+            self.isBusyFinalizing = false
+            self.statusBarController.updateIcon(state: .idle)
+            
+            switch retryResult {
+            case .success(let recoveredText):
+                let trimmed = recoveredText.trimmingCharacters(in: .whitespacesAndNewlines)
+                print("AppDelegate: Smart Audio Recovery SUCCEEDED! Recovered \(trimmed.count) chars.")
+                VoiceCommandLogger.shared.logRecovery(
+                    audioBytes: pcmData.count,
+                    durationSec: durationSec,
+                    reason: reason,
+                    outcome: "SUCCESS",
+                    recoveredText: trimmed
+                )
+                self.deliverOutputText(trimmed)
+                
+            case .failure(let error):
+                print("AppDelegate: Smart Audio Recovery failed completely: \(error)")
+                VoiceCommandLogger.shared.logRecovery(
+                    audioBytes: pcmData.count,
+                    durationSec: durationSec,
+                    reason: reason,
+                    outcome: "FAILED: \(error.localizedDescription)",
+                    recoveredText: ""
+                )
+                HUDOverlayController.shared.show(state: .error)
+                DispatchQueue.main.asyncAfter(deadline: .now() + 0.8) {
+                    HUDOverlayController.shared.hide()
+                }
+            }
+        }
+    }
+    
     private func deliverOutputText(_ text: String) {
         guard !text.isEmpty else {
             HUDOverlayController.shared.hide()

src/InputBoxManager.swift

@@ -236,7 +236,7 @@ public final class InputBoxManager {
                 return
             }
             
-            // Case 3: Streamed prefix is prefix of trimmed final text
+            // Case 3: Streamed prefix is exact prefix of trimmed final text
             if trimmed.count > cleanStreamed.count && trimmed.hasPrefix(cleanStreamed) {
                 let suffix = String(trimmed.dropFirst(cleanStreamed.count))
                 let suffixToPaste = addTrailingSpace ? (suffix + " ") : suffix
@@ -244,7 +244,53 @@ public final class InputBoxManager {
                 pb.clearContents()
                 pb.setString(suffixToPaste, forType: .string)
                 self.postCleanCmdV()
+                return
+            }
+            
+            // Case 4: Word-level token overlap matching (robust to punctuation, ZWNJ, and acoustic model variations)
+            let streamedWords = cleanStreamed.components(separatedBy: .whitespacesAndNewlines).filter { !$0.isEmpty }
+            let finalWords = trimmed.components(separatedBy: .whitespacesAndNewlines).filter { !$0.isEmpty }
+            
+            let cleanWord: (String) -> String = { w in
+                w.replacingOccurrences(of: "[.,!?:;،؛؟\\\"'()\\[\\]{}«»–—\\-]", with: "", options: .regularExpression).lowercased()
+            }
+            let normStreamed = streamedWords.map(cleanWord)
+            let normFinal = finalWords.map(cleanWord)
+            
+            var matchedWordCount = 0
+            for i in 0..<min(normStreamed.count, normFinal.count) {
+                if normStreamed[i] == normFinal[i] {
+                    matchedWordCount = i + 1
+                } else {
+                    break
+                }
             }
+            
+            if matchedWordCount > 0 && finalWords.count > matchedWordCount {
+                let remainingWords = Array(finalWords.dropFirst(matchedWordCount))
+                let suffix = remainingWords.joined(separator: " ")
+                let suffixToPaste = addTrailingSpace ? (" " + suffix + " ") : (" " + suffix)
+                let pb = NSPasteboard.general
+                pb.clearContents()
+                pb.setString(suffixToPaste, forType: .string)
+                self.postCleanCmdV()
+                return
+            }
+            
+            // Case 5: Approximate fallback when streamed words were few and final words are many
+            if finalWords.count > streamedWords.count {
+                let remainingWords = Array(finalWords.dropFirst(streamedWords.count))
+                let suffix = remainingWords.joined(separator: " ")
+                let suffixToPaste = addTrailingSpace ? (" " + suffix + " ") : (" " + suffix)
+                let pb = NSPasteboard.general
+                pb.clearContents()
+                pb.setString(suffixToPaste, forType: .string)
+                self.postCleanCmdV()
+                return
+            }
+            
+            // Case 6: Fallback - paste the full text so speech is NEVER discarded
+            self.postCleanCmdV()
         }
     }
     

src/SonioxClient.swift

@@ -32,9 +32,16 @@ public final class SonioxLiveSession {
     public var onNetworkDegraded: ((Bool) -> Void)?
     
     public private(set) var hasNetworkFailure: Bool = false
+    public private(set) var isRecordingActive: Bool = false
     private var lastDataTime: Date = Date()
     private var heartbeatTimer: Timer?
     
+    public func setRecordingActive(_ active: Bool) {
+        lock.lock()
+        isRecordingActive = active
+        lock.unlock()
+    }
+    
     public var isReady: Bool {
         lock.lock()
         defer { lock.unlock() }
@@ -148,6 +155,7 @@ public final class SonioxLiveSession {
                 print("SonioxLiveSession: Error streaming audio chunk:", error)
                 self?.lock.lock()
                 self?.hasNetworkFailure = true
+                self?.isConnected = false
                 self?.lock.unlock()
                 self?.onNetworkDegraded?(true)
             }
@@ -228,9 +236,18 @@ public final class SonioxLiveSession {
                 print("SonioxLiveSession: WebSocket receive error:", error)
                 self.lock.lock()
                 self.hasNetworkFailure = true
+                self.isConnected = false
+                let activeRecording = self.isRecordingActive
                 self.lock.unlock()
                 self.onNetworkDegraded?(true)
-                self.completeWithCurrentText()
+                
+                // If recording is still active (user speaking), keep buffering in AudioRecorder;
+                // do NOT complete or deliver empty text prematurely!
+                if !activeRecording {
+                    self.completeWithCurrentText()
+                } else {
+                    print("SonioxLiveSession: Network dropped mid-recording. Audio is buffered for Smart Recovery upon key release.")
+                }
             }
         }
     }
@@ -321,6 +338,7 @@ public final class SonioxLiveSession {
         let fullNonFinal = currentNonFinalTokens.joined()
         let rawCombined = fullCommitted.isEmpty ? fullNonFinal : (fullCommitted + fullNonFinal)
         let cleaned = SonioxLiveSession.sanitizeText(rawCombined)
+        let failureOccurred = hasNetworkFailure
         
         let cb = onFinalResult
         webSocketTask?.cancel(with: .normalClosure, reason: nil)
@@ -329,7 +347,11 @@ public final class SonioxLiveSession {
         lock.unlock()
         
         DispatchQueue.main.async {
-            cb?(.success(cleaned))
+            if cleaned.isEmpty && failureOccurred {
+                cb?(.failure(NSError(domain: "SonioxLiveSession", code: -101, userInfo: [NSLocalizedDescriptionKey: "WebSocket network failure with empty transcript"])))
+            } else {
+                cb?(.success(cleaned))
+            }
         }
     }
     
@@ -353,23 +375,47 @@ public final class SonioxLiveSession {
         completion: @escaping (Result<String, Error>) -> Void
     ) {
         guard !pcmData.isEmpty else {
-            completion(.success(""))
+            completion(.failure(NSError(domain: "SonioxLiveSession", code: -102, userInfo: [NSLocalizedDescriptionKey: "Empty audio buffer for recovery"])))
             return
         }
         
         let effectiveTarget = targetLanguage ?? (SonioxSettings.shared.targetLanguage.isEnabled ? SonioxSettings.shared.targetLanguage.rawValue : nil)
-        print("SonioxLiveSession: Starting Smart Burst Retry with \(pcmData.count) bytes (\(Double(pcmData.count)/32000.0)s audio, target: \(effectiveTarget ?? "orig"))...")
-        
-        // Try Primary endpoint, failover to Fallback if needed
-        executeBurstSession(endpoint: primaryWsBase, pcmData: pcmData, languageHints: languageHints, targetLanguage: effectiveTarget, onProgress: onProgress) { result in
-            switch result {
-            case .success(let text):
-                completion(.success(text))
-            case .failure(let error):
-                print("SonioxLiveSession: Primary retry failed (\(error)), attempting fallback endpoint...")
-                executeBurstSession(endpoint: fallbackWsBase, pcmData: pcmData, languageHints: languageHints, targetLanguage: effectiveTarget, onProgress: onProgress, completion: completion)
+        let durationSec = Double(pcmData.count) / 32000.0
+        print("SonioxLiveSession: Initiating Smart Burst Recovery: \(pcmData.count) bytes (\(String(format: "%.1f", durationSec))s audio, target: \(effectiveTarget ?? "orig"))...")
+        
+        // Multi-stage resilient retry pipeline across dual endpoints with automatic backoff:
+        // Attempt 1: Primary WS
+        // Attempt 2: Fallback WS
+        // Attempt 3: (pause 0.8s for network interface/VPN to stabilize) Primary WS
+        // Attempt 4: Fallback WS
+        
+        func executeAttempt(_ step: Int) {
+            let isPrimary = (step == 1 || step == 3)
+            let endpoint = isPrimary ? primaryWsBase : fallbackWsBase
+            let name = isPrimary ? "Primary" : "Fallback"
+            print("SonioxLiveSession: Smart Recovery Attempt \(step)/4 on \(name) endpoint...")
+            
+            executeBurstSession(endpoint: endpoint, pcmData: pcmData, languageHints: languageHints, targetLanguage: effectiveTarget, onProgress: onProgress) { result in
+                switch result {
+                case .success(let text):
+                    print("SonioxLiveSession: Recovery Attempt \(step) SUCCEEDED (\(text.count) chars).")
+                    completion(.success(text))
+                case .failure(let error):
+                    print("SonioxLiveSession: Recovery Attempt \(step) failed: \(error)")
+                    if step < 4 {
+                        let delay = (step == 2) ? 0.8 : 0.2 // 0.8s pause before attempt 3 for network interface stabilization
+                        DispatchQueue.global().asyncAfter(deadline: .now() + delay) {
+                            executeAttempt(step + 1)
+                        }
+                    } else {
+                        print("SonioxLiveSession: All 4 recovery attempts exhausted: \(error)")
+                        completion(.failure(error))
+                    }
+                }
             }
         }
+        
+        executeAttempt(1)
     }
     
     private static func executeBurstSession(
@@ -402,9 +448,14 @@ public final class SonioxLiveSession {
         var request = URLRequest(url: url)
         request.setValue(userAgent, forHTTPHeaderField: "User-Agent")
         request.setValue(origin, forHTTPHeaderField: "Origin")
-        request.timeoutInterval = 25.0
+        
+        let durationSec = Double(pcmData.count) / 32000.0
+        // Dynamic watchdog: allow generous time for network reconnect + fast upload + STT inference
+        let watchdogTimeout = max(8.0, min(30.0, durationSec * 0.75 + 6.0))
+        request.timeoutInterval = watchdogTimeout + 5.0
         
         let config = URLSessionConfiguration.default
+        config.waitsForConnectivity = true
         let session = URLSession(configuration: config)
         let task = session.webSocketTask(with: request)
         task.resume()
@@ -414,16 +465,16 @@ public final class SonioxLiveSession {
         var isDone = false
         let lock = NSLock()
         
-        func finishOnce(with text: String) {
+        func finishOnce(with result: Result<String, Error>) {
             lock.lock()
             if isDone { lock.unlock(); return }
             isDone = true
             task.cancel(with: .normalClosure, reason: nil)
+            session.finishTasksAndInvalidate()
             lock.unlock()
             
-            let sanitized = sanitizeText(text)
             DispatchQueue.main.async {
-                completion(.success(sanitized))
+                completion(result)
             }
         }
         
@@ -478,16 +529,26 @@ public final class SonioxLiveSession {
                         
                         let ended = (json["session_ended"] as? Bool ?? false) || ((json["type"] as? String) == "session_done")
                         if gotFin || ended {
-                            finishOnce(with: combined)
+                            let sanitized = sanitizeText(combined)
+                            if !sanitized.isEmpty {
+                                finishOnce(with: .success(sanitized))
+                            } else {
+                                finishOnce(with: .failure(NSError(domain: "Soniox", code: -4, userInfo: [NSLocalizedDescriptionKey: "Server returned empty transcript"])))
+                            }
                             return
                         }
                     }
                     receiveNext()
                     
                 case .failure(let err):
-                    print("SonioxLiveSession: Burst receive error:", err)
+                    print("SonioxLiveSession: Burst receive error on \(endpoint):", err)
                     let combined = fullCommitted.joined() + fullNonFinal.joined()
-                    finishOnce(with: combined)
+                    let sanitized = sanitizeText(combined)
+                    if !sanitized.isEmpty {
+                        finishOnce(with: .success(sanitized))
+                    } else {
+                        finishOnce(with: .failure(err))
+                    }
                 }
             }
         }
@@ -499,6 +560,11 @@ public final class SonioxLiveSession {
             let chunkSize = 16000 // 0.5s audio
             var offset = 0
             while offset < pcmData.count {
+                lock.lock()
+                let done = isDone
+                lock.unlock()
+                if done { break }
+                
                 let length = min(chunkSize, pcmData.count - offset)
                 let chunk = pcmData.subdata(in: offset..<(offset + length))
                 task.send(URLSessionWebSocketTask.Message.data(chunk)) { _ in }
@@ -506,17 +572,29 @@ public final class SonioxLiveSession {
                 usleep(8000) // 8ms pace
             }
             
-            // Send trailing silence + finalize
-            let silence = Data(count: 3840)
-            task.send(URLSessionWebSocketTask.Message.data(silence)) { _ in }
-            let finalizeMsg = URLSessionWebSocketTask.Message.string("{\"type\": \"finalize\"}")
-            task.send(finalizeMsg) { _ in }
+            lock.lock()
+            let done = isDone
+            lock.unlock()
+            if !done {
+                // Send trailing silence + finalize
+                let silence = Data(count: 3840)
+                task.send(URLSessionWebSocketTask.Message.data(silence)) { _ in }
+                let finalizeMsg = URLSessionWebSocketTask.Message.string("{\"type\": \"finalize\"}")
+                task.send(finalizeMsg) { _ in }
+            }
         }
         
-        // Watchdog timeout for burst retry: max 2.5 seconds (never hangs UI)
-        DispatchQueue.global().asyncAfter(deadline: .now() + 2.5) {
+        // Dynamic Watchdog timeout
+        DispatchQueue.global().asyncAfter(deadline: .now() + watchdogTimeout) {
             let combined = fullCommitted.joined() + fullNonFinal.joined()
-            finishOnce(with: combined)
+            let sanitized = sanitizeText(combined)
+            if !sanitized.isEmpty {
+                print("SonioxLiveSession: Burst watchdog fired, salvaging partial text: \(sanitized.prefix(30))...")
+                finishOnce(with: .success(sanitized))
+            } else {
+                print("SonioxLiveSession: Burst watchdog timed out after \(watchdogTimeout)s with no text.")
+                finishOnce(with: .failure(NSError(domain: "Soniox", code: -5, userInfo: [NSLocalizedDescriptionKey: "Burst watchdog timed out after \(watchdogTimeout)s"])))
+            }
         }
     }
     

src/VoiceCommandLogger.swift

@@ -82,6 +82,30 @@ public final class VoiceCommandLogger {
         }
     }
     
+    public func logRecovery(
+        audioBytes: Int,
+        durationSec: Double,
+        reason: String,
+        outcome: String,
+        recoveredText: String
+    ) {
+        queue.async {
+            let timestamp = self.dateFormatter.string(from: Date())
+            let statusIcon = outcome.contains("SUCCESS") ? "✅" : "⚠️"
+            let entry = """
+            --------------------------------------------------------------------------------
+            [\(timestamp)] 🔄 SMART AUDIO RECOVERY
+              ⏱ Audio Length:       \(String(format: "%.1f", durationSec))s (\(audioBytes) bytes)
+              ⚡️ Trigger Reason:     \(reason)
+              \(statusIcon) Recovery Outcome:  \(outcome)
+              📝 Recovered Text:    "\(recoveredText)"
+            --------------------------------------------------------------------------------
+            
+            """
+            self.appendToFile(text: entry)
+        }
+    }
+    
     private func appendToFile(text: String) {
         guard let data = text.data(using: .utf8) else { return }
         if let fileHandle = try? FileHandle(forWritingTo: logFileURL) {