From d7ad7c5acdaaf4183e6d2ea57a76f59ae3dc2366 Mon Sep 17 00:00:00 2001
From: Beingpax <prakashjoshipaxtek@gmail.com>
Date: Sun, 2 Mar 2025 19:07:43 +0545
Subject: [PATCH] refactor: decouple screen capture from recording start and
 tie it to AI enhancement toggle with race condition protection

---
 VoiceInk/Services/AIEnhancementService.swift | 11 +++++++++++
 VoiceInk/WhisperState.swift                  | 12 +-----------
 2 files changed, 12 insertions(+), 11 deletions(-)

diff --git a/VoiceInk/Services/AIEnhancementService.swift b/VoiceInk/Services/AIEnhancementService.swift
index fa8896d..015a323 100644
--- a/VoiceInk/Services/AIEnhancementService.swift
+++ b/VoiceInk/Services/AIEnhancementService.swift
@@ -22,6 +22,16 @@ class AIEnhancementService: ObservableObject {
                 // Select the first prompt (default) if none is selected
                 selectedPromptId = customPrompts.first?.id
             }
+            
+            // Cancel any existing capture task
+            currentCaptureTask?.cancel()
+            
+            // Trigger screen capture when enhancement is enabled and screen capture is on
+            if isEnhancementEnabled && useScreenCaptureContext {
+                currentCaptureTask = Task {
+                    await captureScreenContext()
+                }
+            }
         }
     }        
     @Published var useClipboardContext: Bool {
@@ -67,6 +77,7 @@ class AIEnhancementService: ObservableObject {
     
     private let aiService: AIService
     private let screenCaptureService: ScreenCaptureService
+    private var currentCaptureTask: Task<Void, Never>?
     private let maxRetries = 3
     private let baseTimeout: TimeInterval = 4
     private let rateLimitInterval: TimeInterval = 1.0 // 1 request per second
diff --git a/VoiceInk/WhisperState.swift b/VoiceInk/WhisperState.swift
index cf16c4e..efe8ebd 100644
--- a/VoiceInk/WhisperState.swift
+++ b/VoiceInk/WhisperState.swift
@@ -292,17 +292,7 @@ class WhisperState: NSObject, ObservableObject, AVAudioRecorderDelegate {
                             // Step 1: Apply power mode configuration
                             await ActiveWindowService.shared.applyConfigurationForCurrentApp()
                             
-                            // Step 2: Handle screen capture if enabled by the configuration
-                            if let enhancementService = self.enhancementService,
-                               enhancementService.isEnhancementEnabled &&
-                               enhancementService.useScreenCaptureContext {
-                                await MainActor.run {
-                                    self.messageLog += "Capturing screen context...\n"
-                                }
-                                await enhancementService.captureScreenContext()
-                            }
-                            
-                            // Step 3: Load model if needed
+                            // Step 2: Load model if needed
                             if let currentModel = self.currentModel, self.whisperContext == nil {
                                 do {
                                     try await self.loadModel(currentModel)