From c91f92f4d09ee3b2345fd5720aa380e8de2a9806 Mon Sep 17 00:00:00 2001
From: Serwan Asaad <serwan.asaad@quantum-machines.co>
Date: Mon, 22 Sep 2025 13:49:07 +0200
Subject: [PATCH] Fix: Enable Clipboard context with custom AI enhancers

---
 VoiceInk/Services/AIEnhancementService.swift |  5 ++-
 VoiceInk/Services/SelectedTextService.swift  | 44 +++++++++++---------
 2 files changed, 27 insertions(+), 22 deletions(-)

diff --git a/VoiceInk/Services/AIEnhancementService.swift b/VoiceInk/Services/AIEnhancementService.swift
index a827b5f..82051b8 100644
--- a/VoiceInk/Services/AIEnhancementService.swift
+++ b/VoiceInk/Services/AIEnhancementService.swift
@@ -132,8 +132,9 @@ class AIEnhancementService: ObservableObject {
     }
 
     private func getSystemMessage(for mode: EnhancementPrompt) -> String {
+        let clipboardSnapshot = NSPasteboard.general.string(forType: .string)
         let selectedText = SelectedTextService.fetchSelectedText()
-
+        
         if let activePrompt = activePrompt,
            activePrompt.id == PredefinedPrompts.assistantPromptId,
            let selectedText = selectedText, !selectedText.isEmpty {
@@ -149,7 +150,7 @@ class AIEnhancementService: ObservableObject {
         }
 
         let clipboardContext = if useClipboardContext,
-                              let clipboardText = NSPasteboard.general.string(forType: .string),
+                              let clipboardText = clipboardSnapshot ?? NSPasteboard.general.string(forType: .string),
                               !clipboardText.isEmpty {
             "\n\n<CLIPBOARD_CONTEXT>\n\(clipboardText)\n</CLIPBOARD_CONTEXT>"
         } else {
diff --git a/VoiceInk/Services/SelectedTextService.swift b/VoiceInk/Services/SelectedTextService.swift
index f026107..c58dc31 100644
--- a/VoiceInk/Services/SelectedTextService.swift
+++ b/VoiceInk/Services/SelectedTextService.swift
@@ -1,11 +1,7 @@
 import Foundation
 import AppKit
-
 class SelectedTextService {
     
-    // Private pasteboard type to avoid clipboard history pollution
-    private static let privatePasteboardType = NSPasteboard.PasteboardType("com.prakashjoshipax.VoiceInk.transient")
-
     static func fetchSelectedText() -> String? {
         // Don't check for selected text within VoiceInk itself
         guard let frontmostApp = NSWorkspace.shared.frontmostApplication,
@@ -14,10 +10,15 @@ class SelectedTextService {
         }
 
         let pasteboard = NSPasteboard.general
-        
-        // Save original clipboard content
+        let originalClipboardText = pasteboard.string(forType: .string)
+
+        // Save original clipboard content (all UTIs with their data)
         let originalPasteboardItems = pasteboard.pasteboardItems?.map { item in
-            (item.types, item.data(forType: item.types.first ?? .string))
+            item.types.reduce(into: [NSPasteboard.PasteboardType: Data]()) { acc, type in
+                if let data = item.data(forType: type) {
+                    acc[type] = data
+                }
+            }
         }
 
         // Clear clipboard to prepare for selection detection
@@ -43,24 +44,27 @@ class SelectedTextService {
 
         // Read the copied text
         let selectedText = pasteboard.string(forType: .string)
-        
+
         // Restore original clipboard content
         pasteboard.clearContents()
-        if let originalItems = originalPasteboardItems {
-            for (types, data) in originalItems {
-                if let data = data {
-                    let pasteboardItem = NSPasteboardItem()
-                    pasteboardItem.setData(data, forType: types.first ?? .string)
-                    pasteboard.writeObjects([pasteboardItem])
+        if let originalItems = originalPasteboardItems, !originalItems.isEmpty {
+            let restoredItems: [NSPasteboardItem] = originalItems.compactMap { dataMap in
+                guard !dataMap.isEmpty else { return nil }
+                let item = NSPasteboardItem()
+                for (type, data) in dataMap {
+                    item.setData(data, forType: type)
                 }
+                return item
             }
+            if !restoredItems.isEmpty {
+                pasteboard.writeObjects(restoredItems)
+            } else if let originalClipboardText {
+                _ = pasteboard.setString(originalClipboardText, forType: .string)
+            }
+        } else if let originalClipboardText {
+            _ = pasteboard.setString(originalClipboardText, forType: .string)
         }
-        
-        // Clear clipboard history by writing transient data
-        let transientItem = NSPasteboardItem()
-        transientItem.setString("", forType: privatePasteboardType)
-        pasteboard.writeObjects([transientItem])
 
         return selectedText
     }
-}
\ No newline at end of file
+}