Revert last dictation to the literal transcript (AI edit undo)

Makefile

@@ -104,14 +104,14 @@ endif
 test: $(TEST_RUNNER)
 	@$(TEST_RUNNER)
 
-$(TEST_RUNNER): Sources/AppContextService.swift Sources/AppleFoundationModelsPostProcessor.swift Sources/DictionaryStore.swift Sources/TranscriptTidier.swift Sources/TransformStore.swift Sources/WakePhraseMatcher.swift Tests/AppContextServiceTests.swift Tests/DictionaryStoreTests.swift Tests/TranscriptTidierTests.swift Tests/TransformStoreTests.swift Tests/WakePhraseMatcherTests.swift Sources/ScratchCommandMatcher.swift Tests/ScratchCommandMatcherTests.swift
+$(TEST_RUNNER): Sources/AppContextService.swift Sources/AppleFoundationModelsPostProcessor.swift Sources/DictionaryStore.swift Sources/TranscriptTidier.swift Sources/TransformStore.swift Sources/WakePhraseMatcher.swift Tests/AppContextServiceTests.swift Tests/DictionaryStoreTests.swift Tests/TranscriptTidierTests.swift Tests/TransformStoreTests.swift Tests/WakePhraseMatcherTests.swift Sources/ScratchCommandMatcher.swift Tests/ScratchCommandMatcherTests.swift Sources/RawRevertEligibility.swift Tests/RawRevertEligibilityTests.swift
 	@mkdir -p "$(BUILD_DIR)"
 	swiftc \
 		-parse-as-library \
 		-o "$(TEST_RUNNER)" \
 		-sdk $(shell xcrun --show-sdk-path) \
 		-target $(ARCH)-apple-macosx26.0 \
-		Sources/AppContextService.swift Sources/AppleFoundationModelsPostProcessor.swift Sources/DictionaryStore.swift Sources/TranscriptTidier.swift Sources/TransformStore.swift Sources/WakePhraseMatcher.swift Tests/AppContextServiceTests.swift Tests/DictionaryStoreTests.swift Tests/TranscriptTidierTests.swift Tests/TransformStoreTests.swift Tests/WakePhraseMatcherTests.swift Sources/ScratchCommandMatcher.swift Tests/ScratchCommandMatcherTests.swift
+		Sources/AppContextService.swift Sources/AppleFoundationModelsPostProcessor.swift Sources/DictionaryStore.swift Sources/TranscriptTidier.swift Sources/TransformStore.swift Sources/WakePhraseMatcher.swift Tests/AppContextServiceTests.swift Tests/DictionaryStoreTests.swift Tests/TranscriptTidierTests.swift Tests/TransformStoreTests.swift Tests/WakePhraseMatcherTests.swift Sources/ScratchCommandMatcher.swift Tests/ScratchCommandMatcherTests.swift Sources/RawRevertEligibility.swift Tests/RawRevertEligibilityTests.swift
 
 icon: $(ICON_ICNS)
 

Sources/AppState.swift

@@ -1684,6 +1684,94 @@ final class AppState: ObservableObject, @unchecked Sendable {
         }
     }
 
+    /// Whether "Revert Last Cleanup" can do anything: a dictation landed
+    /// recently enough to still be replaceable and its literal transcript
+    /// actually differs from what the cleanup pasted.
+    @MainActor
+    var canRevertLastDictationToRaw: Bool {
+        guard !isRecording, !isTranscribing else { return false }
+        guard let item = recentDictationHistoryItem(in: nil, now: Date()) else { return false }
+        return RawRevertEligibility.revertTarget(
+            rawTranscript: item.rawTranscript,
+            cleanedTranscript: item.postProcessedTranscript
+        ) != nil
+    }
+
+    /// Wispr Flow-style "AI edit undo": when the smart cleanup over-edited,
+    /// recover exactly what was said. If the last cleaned dictation is still
+    /// sitting immediately before the caret, select it with the same
+    /// accessibility machinery the wake-command REPLACE_PREVIOUS flow uses
+    /// and paste the literal (pre-cleanup) transcript over the selection.
+    @MainActor
+    func revertLastDictationToRaw() {
+        guard !isRecording, !isTranscribing else { return }
+        let context = fallbackContextAtStop()
+        guard let item = recentDictationHistoryItem(in: context, now: Date()),
+              let rawTranscript = RawRevertEligibility.revertTarget(
+                  rawTranscript: item.rawTranscript,
+                  cleanedTranscript: item.postProcessedTranscript
+              ) else {
+            overlayManager.showError("Nothing to revert: no recent cleaned dictation in this app.")
+            return
+        }
+
+        let cleanedTarget = item.postProcessedTranscript.trimmingCharacters(in: .whitespacesAndNewlines)
+        let pendingClipboardRestore = writeTranscriptToPasteboard(rawTranscript)
+        pasteAtCursorWhenShortcutReleased(performPaste: false) { [weak self] in
+            guard let self else { return }
+            guard self.contextService.selectTextImmediatelyBeforeCaret(matching: cleanedTarget) else {
+                // Nothing was selected, so pasting would duplicate text
+                // instead of replacing it. Leave the document alone.
+                self.restoreClipboardIfNeeded(pendingClipboardRestore)
+                self.overlayManager.showError("Couldn't revert: the dictated text is no longer at the cursor.")
+                return
+            }
+            self.pasteAtCursor()
+            self.restoreClipboardIfNeeded(pendingClipboardRestore)
+            self.recordRawRevert(of: item, rawTranscript: rawTranscript)
+        }
+    }
+
+    /// After a successful revert the literal transcript is what sits before
+    /// the caret, so it becomes the tracked previous text: follow-up wake
+    /// commands ("megaphone, make that shorter") must edit the reverted
+    /// words, and Paste Again must repeat them.
+    @MainActor
+    private func recordRawRevert(of item: PipelineHistoryItem, rawTranscript: String) {
+        let updatedItem = PipelineHistoryItem(
+            intent: item.intent,
+            selectedText: item.selectedText,
+            capturedSelection: item.capturedSelection,
+            id: item.id,
+            timestamp: item.timestamp,
+            rawTranscript: item.rawTranscript,
+            postProcessedTranscript: rawTranscript,
+            postProcessingPrompt: item.postProcessingPrompt,
+            systemPrompt: item.systemPrompt,
+            contextSummary: item.contextSummary,
+            contextSystemPrompt: item.contextSystemPrompt,
+            contextPrompt: item.contextPrompt,
+            contextScreenshotDataURL: item.contextScreenshotDataURL,
+            contextScreenshotStatus: item.contextScreenshotStatus,
+            postProcessingStatus: item.postProcessingStatus + " — reverted to literal transcript",
+            debugStatus: item.debugStatus,
+            customVocabulary: item.customVocabulary,
+            audioFileName: item.audioFileName,
+            contextAppName: item.contextAppName,
+            contextBundleIdentifier: item.contextBundleIdentifier,
+            contextWindowTitle: item.contextWindowTitle
+        )
+        do {
+            try pipelineHistoryStore.update(updatedItem)
+            pipelineHistory = pipelineHistoryStore.loadAllHistory()
+        } catch {
+            errorMessage = "Unable to save revert in run history: \(error.localizedDescription)"
+        }
+        lastTranscript = rawTranscript
+        statusText = "Reverted to literal transcript"
+        scheduleReadyStatusReset(after: 3, matching: ["Reverted to literal transcript"])
+    }
+
     func toggleRecording() {
         os_log(.info, log: recordingLog, "toggleRecording() called, isRecording=%{public}d", isRecording)
         cancelPendingShortcutStart()
@@ -2403,8 +2491,13 @@ final class AppState: ObservableObject, @unchecked Sendable {
         }
     }
 
+    /// The newest pipeline-history entry whose inserted text is recent
+    /// enough (`wakeCommandPreviousTextWindow`) to still count as "the
+    /// previous dictation". Pass a context to additionally require that the
+    /// entry was dictated into the same app and window; pass nil when the
+    /// destination cannot be known yet (e.g. menu item validation).
     @MainActor
-    private func recentTextForWakeCommand(in context: AppContext, now: Date) -> String? {
+    private func recentDictationHistoryItem(in context: AppContext?, now: Date) -> PipelineHistoryItem? {
         pipelineHistory.first { item in
             let text = item.postProcessedTranscript.trimmingCharacters(in: .whitespacesAndNewlines)
             guard !text.isEmpty else { return false }
@@ -2413,6 +2506,7 @@ final class AppState: ObservableObject, @unchecked Sendable {
             }
             let age = now.timeIntervalSince(item.timestamp)
             guard age >= 0, age <= wakeCommandPreviousTextWindow else { return false }
+            guard let context else { return true }
 
             let previousBundleID = item.contextBundleIdentifier?.trimmingCharacters(in: .whitespacesAndNewlines)
             let currentBundleID = context.bundleIdentifier?.trimmingCharacters(in: .whitespacesAndNewlines)
@@ -2430,7 +2524,12 @@ final class AppState: ObservableObject, @unchecked Sendable {
                 return false
             }
             return true
-        }?.postProcessedTranscript
+        }
+    }
+
+    @MainActor
+    private func recentTextForWakeCommand(in context: AppContext, now: Date) -> String? {
+        recentDictationHistoryItem(in: context, now: now)?.postProcessedTranscript
     }
 
     private func processTranscript(

Sources/MenuBarView.swift

@@ -156,6 +156,15 @@ struct MenuBarView: View {
                     .frame(maxWidth: 280, alignment: .leading)
             }
 
+            // Wispr Flow-style "AI edit undo": swap the last cleaned
+            // dictation still at the caret for the literal transcript.
+            // Disabled when no recent dictation is revertible or the cleanup
+            // made no edits (raw == cleaned).
+            Button("Revert Last Cleanup") {
+                appState.revertLastDictationToRaw()
+            }
+            .disabled(!appState.canRevertLastDictationToRaw)
+
             Menu("History") {
                 if recentHistoryItems.isEmpty {
                     Text("No transcripts yet")

Sources/RawRevertEligibility.swift

@@ -0,0 +1,17 @@
+import Foundation
+
+/// Decides whether a dictation can be reverted to its literal (pre-cleanup)
+/// transcript — Wispr Flow calls this "AI edit undo". Pure string logic so it
+/// stays unit-testable without AppKit; the accessibility half of the feature
+/// (is the cleaned text still at the caret?) lives in AppContextService.
+enum RawRevertEligibility {
+    /// Returns the raw transcript that should replace the cleaned text, or
+    /// nil when reverting would be pointless: nothing was dictated, nothing
+    /// was pasted, or the cleanup made no edits at all.
+    static func revertTarget(rawTranscript: String, cleanedTranscript: String) -> String? {
+        let raw = rawTranscript.trimmingCharacters(in: .whitespacesAndNewlines)
+        let cleaned = cleanedTranscript.trimmingCharacters(in: .whitespacesAndNewlines)
+        guard !raw.isEmpty, !cleaned.isEmpty, raw != cleaned else { return nil }
+        return raw
+    }
+}

Tests/AppContextServiceTests.swift

@@ -25,6 +25,7 @@ struct AppContextServiceTests {
         WakePhraseMatcherTests.run()
         TransformStoreTests.run()
         ScratchCommandMatcherTests.run()
+        RawRevertEligibilityTests.run()
         print("MegaphoneTests passed")
     }
 

Tests/RawRevertEligibilityTests.swift

@@ -0,0 +1,76 @@
+import Foundation
+
+struct RawRevertEligibilityTests {
+    static func run() {
+        testRevertTargetWhenCleanupEditedTheText()
+        testNoRevertWhenCleanupMadeNoEdits()
+        testWhitespaceOnlyDifferencesAreNotEdits()
+        testEmptySidesAreNeverRevertible()
+        testRevertTargetIsTrimmed()
+    }
+
+    private static func testRevertTargetWhenCleanupEditedTheText() {
+        expectEqual(
+            RawRevertEligibility.revertTarget(
+                rawTranscript: "um so basically ship it on friday",
+                cleanedTranscript: "Ship it on Friday."
+            ),
+            "um so basically ship it on friday"
+        )
+    }
+
+    private static func testNoRevertWhenCleanupMadeNoEdits() {
+        expect(
+            RawRevertEligibility.revertTarget(
+                rawTranscript: "Ship it on Friday.",
+                cleanedTranscript: "Ship it on Friday."
+            ) == nil,
+            "Identical raw and cleaned transcripts must not be revertible"
+        )
+    }
+
+    private static func testWhitespaceOnlyDifferencesAreNotEdits() {
+        expect(
+            RawRevertEligibility.revertTarget(
+                rawTranscript: "  Ship it on Friday.\n",
+                cleanedTranscript: "Ship it on Friday."
+            ) == nil,
+            "Surrounding whitespace alone is not a cleanup edit"
+        )
+    }
+
+    private static func testEmptySidesAreNeverRevertible() {
+        expect(
+            RawRevertEligibility.revertTarget(rawTranscript: "", cleanedTranscript: "Hello.") == nil,
+            "An empty raw transcript has nothing to revert to"
+        )
+        expect(
+            RawRevertEligibility.revertTarget(rawTranscript: "hello", cleanedTranscript: "") == nil,
+            "Nothing was pasted, so there is nothing to replace"
+        )
+        expect(
+            RawRevertEligibility.revertTarget(rawTranscript: "   \n", cleanedTranscript: "Hello.") == nil,
+            "A whitespace-only raw transcript has nothing to revert to"
+        )
+    }
+
+    private static func testRevertTargetIsTrimmed() {
+        expectEqual(
+            RawRevertEligibility.revertTarget(
+                rawTranscript: "  hey there\n",
+                cleanedTranscript: "Hey there."
+            ),
+            "hey there"
+        )
+    }
+
+    private static func expectEqual(_ actual: String?, _ expected: String, file: StaticString = #file, line: UInt = #line) {
+        expect(actual == expected, "Expected \(expected.debugDescription), got \((actual ?? "nil").debugDescription)", file: file, line: line)
+    }
+
+    private static func expect(_ condition: Bool, _ message: String, file: StaticString = #file, line: UInt = #line) {
+        if !condition {
+            fatalError("\(file):\(line): \(message)")
+        }
+    }
+}