From 504b8f72618e8d64a5497e83562662a6b6862ba8 Mon Sep 17 00:00:00 2001
From: Beingpax <prakashjoshipaxtek@gmail.com>
Date: Tue, 8 Jul 2025 12:53:21 +0545
Subject: [PATCH] added cancellation checks during transcription

---
 VoiceInk/Whisper/WhisperState.swift | 26 +++++++++++++++++++++-----
 1 file changed, 21 insertions(+), 5 deletions(-)

diff --git a/VoiceInk/Whisper/WhisperState.swift b/VoiceInk/Whisper/WhisperState.swift
index a4889a2..749af5b 100644
--- a/VoiceInk/Whisper/WhisperState.swift
+++ b/VoiceInk/Whisper/WhisperState.swift
@@ -305,6 +305,8 @@ class WhisperState: NSObject, ObservableObject, AVAudioRecorderDelegate {
             var text = try await transcriptionService.transcribe(audioURL: url, model: model)
             let transcriptionDuration = Date().timeIntervalSince(transcriptionStart)
             
+            if await checkCancellationAndCleanup() { return }
+            
             text = text.trimmingCharacters(in: .whitespacesAndNewlines)
             
             if UserDefaults.standard.bool(forKey: "IsWordReplacementEnabled") {
@@ -326,7 +328,8 @@ class WhisperState: NSObject, ObservableObject, AVAudioRecorderDelegate {
                enhancementService.isEnhancementEnabled,
                enhancementService.isConfigured {
                 do {
-                    if shouldCancelRecording { return }
+                    if await checkCancellationAndCleanup() { return }
+
                     let textForAI = promptDetectionResult?.processedText ?? text
                     let (enhancedText, enhancementDuration) = try await enhancementService.enhance(textForAI)
                     let newTranscription = Transcription(
@@ -382,6 +385,8 @@ class WhisperState: NSObject, ObservableObject, AVAudioRecorderDelegate {
 
             text += " "
 
+            if await checkCancellationAndCleanup() { return }
+
             SoundManager.shared.playStopSound()
             DispatchQueue.main.asyncAfter(deadline: .now() + 0.05) {
                 
@@ -399,8 +404,7 @@ class WhisperState: NSObject, ObservableObject, AVAudioRecorderDelegate {
                 await promptDetectionService.restoreOriginalSettings(result, to: enhancementService)
             }
             
-            await dismissMiniRecorder()
-            await cleanupModelResources()
+            await cleanupAndDismiss()
             
         } catch {
             if let permanentURL = permanentURL {
@@ -443,8 +447,7 @@ class WhisperState: NSObject, ObservableObject, AVAudioRecorderDelegate {
                 }
             }
             
-            await cleanupModelResources()
-            await dismissMiniRecorder()
+            await cleanupAndDismiss()
         }
     }
 
@@ -548,6 +551,19 @@ class WhisperState: NSObject, ObservableObject, AVAudioRecorderDelegate {
             setDefaultTranscriptionModel(updatedModel)
         }
     }
+    
+    private func checkCancellationAndCleanup() async -> Bool {
+        if shouldCancelRecording {
+            await cleanupAndDismiss()
+            return true
+        }
+        return false
+    }
+    
+    private func cleanupAndDismiss() async {
+        await cleanupModelResources()
+        await dismissMiniRecorder()
+    }
 }
 
 struct WhisperModel: Identifiable {