Fix clipboard restore timing for dictation insertion

Sources/OpenWisprLib/TextInserter.swift

@@ -3,32 +3,87 @@ import Foundation
 import Cocoa
 import Carbon.HIToolbox
 
+protocol TextInsertionPasteboard: AnyObject {
+    var pasteboardItems: [NSPasteboardItem]? { get }
+    var changeCount: Int { get }
+
+    @discardableResult
+    func clearContents() -> Int
+
+    @discardableResult
+    func setString(_ string: String, forType dataType: NSPasteboard.PasteboardType) -> Bool
+
+    @discardableResult
+    func writeItems(_ items: [NSPasteboardItem]) -> Bool
+}
+
+extension NSPasteboard: TextInsertionPasteboard {
+    func writeItems(_ items: [NSPasteboardItem]) -> Bool {
+        writeObjects(items)
+    }
+}
+
 class TextInserter {
+    typealias PasteboardProvider = () -> any TextInsertionPasteboard
+    typealias PasteAction = (CGKeyCode) -> Void
+    typealias RestoreScheduler = (_ delay: TimeInterval, _ action: @escaping () -> Void) -> Void
+
+    static let defaultRestoreDelay: TimeInterval = 1.0
+
     let pasteKeyCode: CGKeyCode
 
-    init() {
-        self.pasteKeyCode = TextInserter.resolveKeyCode(for: "v") ?? 9
+    private let pasteboardProvider: PasteboardProvider
+    private let pasteAction: PasteAction
+    private let restoreDelay: TimeInterval
+    private let scheduleRestore: RestoreScheduler
+
+    convenience init() {
+        let pasteKeyCode = TextInserter.resolveKeyCode(for: "v") ?? 9
+        self.init(
+            pasteKeyCode: pasteKeyCode,
+            pasteboardProvider: { NSPasteboard.general },
+            pasteAction: TextInserter.simulatePaste,
+            scheduleRestore: { delay, action in
+                DispatchQueue.main.asyncAfter(deadline: .now() + delay) {
+                    action()
+                }
+            }
+        )
+    }
+
+    init(
+        pasteKeyCode: CGKeyCode,
+        pasteboardProvider: @escaping PasteboardProvider,
+        pasteAction: @escaping PasteAction,
+        restoreDelay: TimeInterval = TextInserter.defaultRestoreDelay,
+        scheduleRestore: @escaping RestoreScheduler
+    ) {
+        self.pasteKeyCode = pasteKeyCode
+        self.pasteboardProvider = pasteboardProvider
+        self.pasteAction = pasteAction
+        self.restoreDelay = restoreDelay
+        self.scheduleRestore = scheduleRestore
     }
 
     func insert(text: String) {
-        let pasteboard = NSPasteboard.general
+        let pasteboard = pasteboardProvider()
         let savedItems = savePasteboard(pasteboard)
 
         pasteboard.clearContents()
         pasteboard.setString(text, forType: .string)
         let writeChangeCount = pasteboard.changeCount
 
-        simulatePaste()
+        pasteAction(pasteKeyCode)
 
-        DispatchQueue.main.asyncAfter(deadline: .now() + 0.4) {
+        scheduleRestore(restoreDelay) {
             // If something else has written to the pasteboard since our write
             // (user copied something, another tool wrote), do not clobber it.
             guard pasteboard.changeCount == writeChangeCount else { return }
             self.restorePasteboard(pasteboard, items: savedItems)
         }
     }
 
-    private func savePasteboard(_ pasteboard: NSPasteboard) -> [[(NSPasteboard.PasteboardType, Data)]] {
+    private func savePasteboard(_ pasteboard: any TextInsertionPasteboard) -> [[(NSPasteboard.PasteboardType, Data)]] {
         guard let items = pasteboard.pasteboardItems else { return [] }
         return items.map { item in
             item.types.compactMap { type in
@@ -38,7 +93,7 @@ class TextInserter {
         }
     }
 
-    private func restorePasteboard(_ pasteboard: NSPasteboard, items: [[(NSPasteboard.PasteboardType, Data)]]) {
+    private func restorePasteboard(_ pasteboard: any TextInsertionPasteboard, items: [[(NSPasteboard.PasteboardType, Data)]]) {
         pasteboard.clearContents()
         guard !items.isEmpty else { return }
         let pasteboardItems = items.map { entries -> NSPasteboardItem in
@@ -48,7 +103,7 @@ class TextInserter {
             }
             return item
         }
-        pasteboard.writeObjects(pasteboardItems)
+        pasteboard.writeItems(pasteboardItems)
     }
 
     private static func resolveKeyCode(for target: Character) -> CGKeyCode? {
@@ -95,9 +150,7 @@ class TextInserter {
         return nil
     }
 
-    private func simulatePaste() {
-        let keyCode = pasteKeyCode
-
+    private static func simulatePaste(keyCode: CGKeyCode) {
         guard let source = CGEventSource(stateID: .hidSystemState),
             let keyDown = CGEvent(keyboardEventSource: source, virtualKey: keyCode, keyDown: true),
             let keyUp = CGEvent(keyboardEventSource: source, virtualKey: keyCode, keyDown: false) else {

Tests/OpenWisprTests/TextInserterTests.swift

@@ -1,3 +1,4 @@
+import AppKit
 import XCTest
 @testable import OpenWisprLib
 
@@ -7,4 +8,158 @@ final class TextInserterTests: XCTestCase {
         let inserter = TextInserter()
         XCTAssertTrue(inserter.pasteKeyCode < 128, "Paste key code should be a valid virtual key code")
     }
+
+    func testInsertWritesTranscriptionBeforePastingAndSchedulesRestore() {
+        var events: [String] = []
+        let pasteboard = FakePasteboard(string: "old clipboard") { events.append($0) }
+        let scheduler = CapturingScheduler { events.append("schedule") }
+        let inserter = makeInserter(pasteboard: pasteboard, scheduler: scheduler) { _ in
+            events.append("paste:\(pasteboard.string(forType: .string) ?? "")")
+        }
+
+        inserter.insert(text: "new transcription")
+
+        XCTAssertEqual(events, [
+            "clear",
+            "set:new transcription",
+            "paste:new transcription",
+            "schedule",
+        ])
+        XCTAssertEqual(pasteboard.string(forType: .string), "new transcription")
+        guard let scheduledAction = scheduler.onlyScheduledAction() else { return }
+        XCTAssertEqual(scheduledAction.delay, TextInserter.defaultRestoreDelay, accuracy: 0.001)
+    }
+
+    func testInsertUsesSaferDefaultRestoreDelay() {
+        let pasteboard = FakePasteboard(string: "old clipboard")
+        let scheduler = CapturingScheduler()
+        let inserter = makeInserter(pasteboard: pasteboard, scheduler: scheduler)
+
+        inserter.insert(text: "new transcription")
+
+        XCTAssertEqual(TextInserter.defaultRestoreDelay, 1.0, accuracy: 0.001)
+        guard let scheduledAction = scheduler.onlyScheduledAction() else { return }
+        XCTAssertEqual(scheduledAction.delay, 1.0, accuracy: 0.001)
+        XCTAssertEqual(pasteboard.string(forType: .string), "new transcription")
+    }
+
+    func testScheduledRestoreRestoresOriginalClipboardWhenUnchanged() {
+        let pasteboard = FakePasteboard(string: "old clipboard")
+        let scheduler = CapturingScheduler()
+        let inserter = makeInserter(pasteboard: pasteboard, scheduler: scheduler)
+
+        inserter.insert(text: "new transcription")
+        scheduler.runOnlyScheduledAction()
+
+        XCTAssertEqual(pasteboard.string(forType: .string), "old clipboard")
+    }
+
+    func testScheduledRestoreSkipsRestoreWhenPasteboardChanged() {
+        let pasteboard = FakePasteboard(string: "old clipboard")
+        let scheduler = CapturingScheduler()
+        let inserter = makeInserter(pasteboard: pasteboard, scheduler: scheduler)
+
+        inserter.insert(text: "new transcription")
+        pasteboard.clearContents()
+        pasteboard.setString("newer clipboard", forType: .string)
+        scheduler.runOnlyScheduledAction()
+
+        XCTAssertEqual(pasteboard.string(forType: .string), "newer clipboard")
+        XCTAssertFalse(pasteboard.events.contains("writeItems"))
+    }
+
+    private func makeInserter(
+        pasteboard: FakePasteboard,
+        scheduler: CapturingScheduler,
+        pasteAction: @escaping (CGKeyCode) -> Void = { _ in }
+    ) -> TextInserter {
+        TextInserter(
+            pasteKeyCode: 9,
+            pasteboardProvider: { pasteboard },
+            pasteAction: pasteAction,
+            scheduleRestore: { delay, action in
+                scheduler.schedule(delay: delay, action: action)
+            }
+        )
+    }
+}
+
+private final class FakePasteboard: TextInsertionPasteboard {
+    private(set) var pasteboardItems: [NSPasteboardItem]?
+    private(set) var changeCount = 0
+    private(set) var events: [String] = []
+
+    private let onEvent: (String) -> Void
+
+    init(string: String? = nil, onEvent: @escaping (String) -> Void = { _ in }) {
+        self.onEvent = onEvent
+
+        if let string {
+            let item = NSPasteboardItem()
+            item.setString(string, forType: .string)
+            pasteboardItems = [item]
+        }
+    }
+
+    @discardableResult
+    func clearContents() -> Int {
+        pasteboardItems = []
+        return recordChange("clear")
+    }
+
+    @discardableResult
+    func setString(_ string: String, forType dataType: NSPasteboard.PasteboardType) -> Bool {
+        let item = NSPasteboardItem()
+        item.setString(string, forType: dataType)
+        pasteboardItems = [item]
+        _ = recordChange("set:\(string)")
+        return true
+    }
+
+    @discardableResult
+    func writeItems(_ items: [NSPasteboardItem]) -> Bool {
+        pasteboardItems = items
+        _ = recordChange("writeItems")
+        return true
+    }
+
+    func string(forType dataType: NSPasteboard.PasteboardType) -> String? {
+        pasteboardItems?.first?.string(forType: dataType)
+    }
+
+    private func recordChange(_ event: String) -> Int {
+        changeCount += 1
+        events.append(event)
+        onEvent(event)
+        return changeCount
+    }
+}
+
+private final class CapturingScheduler {
+    struct ScheduledAction {
+        let delay: TimeInterval
+        let action: () -> Void
+    }
+
+    private(set) var scheduledActions: [ScheduledAction] = []
+    private let onSchedule: () -> Void
+
+    init(onSchedule: @escaping () -> Void = {}) {
+        self.onSchedule = onSchedule
+    }
+
+    func schedule(delay: TimeInterval, action: @escaping () -> Void) {
+        scheduledActions.append(ScheduledAction(delay: delay, action: action))
+        onSchedule()
+    }
+
+    func onlyScheduledAction(file: StaticString = #filePath, line: UInt = #line) -> ScheduledAction? {
+        XCTAssertEqual(scheduledActions.count, 1, file: file, line: line)
+        guard scheduledActions.count == 1 else { return nil }
+        return scheduledActions[0]
+    }
+
+    func runOnlyScheduledAction(file: StaticString = #filePath, line: UInt = #line) {
+        onlyScheduledAction(file: file, line: line)?.action()
+    }
 }