diff --git a/Cotabby.xcodeproj/project.pbxproj b/Cotabby.xcodeproj/project.pbxproj index dd0e859..830814c 100644 --- a/Cotabby.xcodeproj/project.pbxproj +++ b/Cotabby.xcodeproj/project.pbxproj @@ -12,6 +12,7 @@ 19E177EA2FC1B7890067E267 /* CotabbyInference in Frameworks */ = {isa = PBXBuildFile; productRef = 19E177E92FC1B7890067E267 /* CotabbyInference */; }; 7A1E3B0C2FC8A00100CARET1 /* AXTextGeometryResolverTests.swift in Sources */ = {isa = PBXBuildFile; fileRef = 7A1E3B0D2FC8A00100CARET2 /* AXTextGeometryResolverTests.swift */; }; 8B6282F0C1CCA0746D96B914 /* DownloadOutcomeClassifierTests.swift in Sources */ = {isa = PBXBuildFile; fileRef = 562F89255AF340C15A0554BE /* DownloadOutcomeClassifierTests.swift */; }; + H10000012FC0000100FFF001 /* PrefixCorrectionFilterTests.swift in Sources */ = {isa = PBXBuildFile; fileRef = H10000112FC0000100FFF011 /* PrefixCorrectionFilterTests.swift */; }; A1C3E0112F90000100AAA001 /* Sparkle in Frameworks */ = {isa = PBXBuildFile; productRef = A1C3E0102F90000100AAA001 /* Sparkle */; }; A404828463CADB2ECDAE7AF3 /* LlamaPromptRendererTests.swift in Sources */ = {isa = PBXBuildFile; fileRef = F29623C5C0A67B992D383A3C /* LlamaPromptRendererTests.swift */; }; ADEFEE12C197DB6C990E3812 /* SuggestionRequestFactoryTests.swift in Sources */ = {isa = PBXBuildFile; fileRef = 8B1121FFDAD30C7F62E15289 /* SuggestionRequestFactoryTests.swift */; }; @@ -53,6 +54,7 @@ 193741492F81DE7000BEC04F /* Cotabby.app */ = {isa = PBXFileReference; explicitFileType = wrapper.application; includeInIndex = 0; path = Cotabby.app; sourceTree = BUILT_PRODUCTS_DIR; }; 3FBFA92FA44AA317135426FB /* CotabbyTests.xctest */ = {isa = PBXFileReference; explicitFileType = wrapper.cfbundle; includeInIndex = 0; path = CotabbyTests.xctest; sourceTree = BUILT_PRODUCTS_DIR; }; 562F89255AF340C15A0554BE /* DownloadOutcomeClassifierTests.swift */ = {isa = PBXFileReference; includeInIndex = 1; lastKnownFileType = sourcecode.swift; path = DownloadOutcomeClassifierTests.swift; sourceTree = ""; }; + H10000112FC0000100FFF011 /* PrefixCorrectionFilterTests.swift */ = {isa = PBXFileReference; includeInIndex = 1; lastKnownFileType = sourcecode.swift; path = PrefixCorrectionFilterTests.swift; sourceTree = ""; }; 5A39EC1A44E9160E13E2AE77 /* SuggestionAvailabilityEvaluatorTests.swift */ = {isa = PBXFileReference; includeInIndex = 1; lastKnownFileType = sourcecode.swift; path = SuggestionAvailabilityEvaluatorTests.swift; sourceTree = ""; }; 7A1E3B0D2FC8A00100CARET2 /* AXTextGeometryResolverTests.swift */ = {isa = PBXFileReference; includeInIndex = 1; lastKnownFileType = sourcecode.swift; path = AXTextGeometryResolverTests.swift; sourceTree = ""; }; 8B1121FFDAD30C7F62E15289 /* SuggestionRequestFactoryTests.swift */ = {isa = PBXFileReference; includeInIndex = 1; lastKnownFileType = sourcecode.swift; path = SuggestionRequestFactoryTests.swift; sourceTree = ""; }; @@ -157,6 +159,7 @@ G20000112FB0000100FFF011 /* ClipboardRelevanceFilterTests.swift */, G10000122FB0000100FFF012 /* ClipboardContentDistillerTests.swift */, G10000112FB0000100FFF011 /* WordCountFormatterTests.swift */, + H10000112FC0000100FFF011 /* PrefixCorrectionFilterTests.swift */, ); path = CotabbyTests; sourceTree = ""; @@ -303,6 +306,7 @@ G20000012FB0000100FFF001 /* ClipboardRelevanceFilterTests.swift in Sources */, G10000022FB0000100FFF002 /* ClipboardContentDistillerTests.swift in Sources */, G10000012FB0000100FFF001 /* WordCountFormatterTests.swift in Sources */, + H10000012FC0000100FFF001 /* PrefixCorrectionFilterTests.swift in Sources */, ); runOnlyForDeploymentPostprocessing = 0; }; diff --git a/Cotabby/App/Coordinators/PrefixCorrectionCoordinator.swift b/Cotabby/App/Coordinators/PrefixCorrectionCoordinator.swift new file mode 100644 index 0000000..e98f287 --- /dev/null +++ b/Cotabby/App/Coordinators/PrefixCorrectionCoordinator.swift @@ -0,0 +1,213 @@ +import Combine +import Foundation +import Logging + +/// File overview: +/// Owns the settled-pause-driven typo-correction loop. Subscribes to focus snapshots, +/// detects when the user has stopped typing for `settledDuration`, asks the correction +/// engine for a proposed fix, validates it through `PrefixCorrectionFilter`, and writes +/// the accepted result back to the focused field via `PrefixCorrectionWriter`. +/// +/// All collaborators are injected by capability — engine, writer, and settings/state +/// queries are closures or narrow protocols. The coordinator owns the state machine, +/// cancellation discipline, and gate logic; everything else can be stubbed. +/// +/// Cancellation discipline (mirrors `SuggestionWorkController`): +/// - Each settled event gets a monotonically increasing work ID. +/// - The async correction call may complete after the user has typed more characters. +/// Before writing, the coordinator re-reads the live focus snapshot and drops the +/// result if the prefix changed under it. +/// - `lastSubmittedPrefix` is tracked per-bundle-identifier so a correction that simply +/// re-surfaces (because writing the fix triggered another settled event) is not +/// re-sent to the model — saves a roundtrip and keeps the loop naturally idempotent. +@MainActor +final class PrefixCorrectionCoordinator { + /// 800ms is a balance between "the user is mid-thought" and "the user has moved on." + /// Short enough that the fix lands while the user can still see their original typo; + /// long enough that pause-to-think doesn't trigger a correction mid-sentence. + static let defaultSettledDuration: TimeInterval = 0.8 + /// Below this character count the model has almost no signal to work with and short + /// fragments are usually still being typed. + static let defaultMinimumPrefixCharacterCount = 12 + /// Above this the backspace burst becomes user-visible flicker. The feature is + /// already opt-in and per-app, so a hard cap is safer than trying to optimize. + static let defaultMaximumPrefixCharacterCount = 500 + + private let focusModel: any SuggestionFocusProviding + private let correctionEngine: any PrefixCorrecting + private let writer: PrefixCorrectionWriter + private let isCorrectionEnabled: @MainActor () -> Bool + private let isAutocompleteBusy: @MainActor () -> Bool + private let settledDuration: TimeInterval + private let minimumPrefixCharacterCount: Int + private let maximumPrefixCharacterCount: Int + + private var cancellables = Set() + private var latestWorkID: UInt64 = 0 + private var inflightTask: Task? + /// Most recent prefix that was submitted to the engine, keyed by the bundle it came + /// from. Lets us short-circuit "we just corrected this, the publisher fired again + /// with the corrected text" without burning another LLM call. + private var lastSubmittedPrefix: [String: String] = [:] + + /// Debug-only hook fired right after a correction is written, with the original prefix and the + /// accepted replacement. Set by app composition only when the debug overlay is active; nil + /// otherwise so production does no extra work. + var onCorrectionApplied: (@MainActor (_ original: String, _ corrected: String) -> Void)? + + init( + focusModel: any SuggestionFocusProviding, + correctionEngine: any PrefixCorrecting, + writer: PrefixCorrectionWriter, + isCorrectionEnabled: @escaping @MainActor () -> Bool, + isAutocompleteBusy: @escaping @MainActor () -> Bool, + settledDuration: TimeInterval = defaultSettledDuration, + minimumPrefixCharacterCount: Int = defaultMinimumPrefixCharacterCount, + maximumPrefixCharacterCount: Int = defaultMaximumPrefixCharacterCount + ) { + self.focusModel = focusModel + self.correctionEngine = correctionEngine + self.writer = writer + self.isCorrectionEnabled = isCorrectionEnabled + self.isAutocompleteBusy = isAutocompleteBusy + self.settledDuration = settledDuration + self.minimumPrefixCharacterCount = minimumPrefixCharacterCount + self.maximumPrefixCharacterCount = maximumPrefixCharacterCount + } + + func start() { + guard cancellables.isEmpty else { return } + + focusModel.snapshotPublisher + .compactMap { snapshot -> SettledKey? in + guard let context = snapshot.context else { return nil } + return SettledKey( + bundleIdentifier: snapshot.bundleIdentifier, + precedingText: context.precedingText, + selectionLength: context.selection.length, + isSecure: context.isSecure + ) + } + .removeDuplicates() + .debounce(for: .seconds(settledDuration), scheduler: DispatchQueue.main) + .sink { [weak self] _ in + self?.handleSettled() + } + .store(in: &cancellables) + } + + func stop() { + cancellables.removeAll() + inflightTask?.cancel() + inflightTask = nil + } + + // MARK: - Settled-event handling + + /// One settled event = one attempt. Bumps the work ID immediately so any in-flight + /// attempt becomes stale, then spawns a fresh Task to drive the engine + filter + + /// writer pipeline. + private func handleSettled() { + latestWorkID &+= 1 + let workID = latestWorkID + + inflightTask?.cancel() + + // Re-read fresh from the focus model rather than trusting whatever the publisher + // delivered — the snapshot may have moved between debounce-firing and this sink. + focusModel.refreshNow() + let snapshot = focusModel.snapshot + guard let context = snapshot.context else { return } + + guard passesGate(snapshot: snapshot, context: context) else { return } + + let bundleKey = snapshot.bundleIdentifier ?? "" + if lastSubmittedPrefix[bundleKey] == context.precedingText { + // Already asked about this exact prefix in this app — nothing to do. + return + } + lastSubmittedPrefix[bundleKey] = context.precedingText + + let originalPrefix = context.precedingText + let originalLength = originalPrefix.count + + inflightTask = Task { [weak self] in + guard let self else { return } + await self.runCorrection( + workID: workID, + originalPrefix: originalPrefix, + originalLength: originalLength + ) + } + } + + private func runCorrection(workID: UInt64, originalPrefix: String, originalLength: Int) async { + let proposal: String? + do { + proposal = try await correctionEngine.proposeCorrection(for: originalPrefix) + } catch is CancellationError { + return + } catch SuggestionClientError.cancelled { + return + } catch { + CotabbyLogger.suggestion.debug("Prefix-correction engine error: \(error.localizedDescription)") + return + } + + guard !Task.isCancelled, workID == latestWorkID else { return } + guard let proposal else { return } + + // The user may have typed more characters during the LLM round-trip. If the live + // prefix no longer matches what we submitted, drop the result — it would clobber + // characters that didn't exist when the model formed its answer. + focusModel.refreshNow() + guard let liveContext = focusModel.snapshot.context, + liveContext.precedingText == originalPrefix, + liveContext.selection.length == 0 + else { + return + } + + guard let accepted = PrefixCorrectionFilter.acceptedCorrection( + original: originalPrefix, + proposed: proposal + ) else { + CotabbyLogger.suggestion.debug("Prefix-correction filter rejected proposal") + return + } + + // Record the accepted output so the re-trigger from our own write doesn't ask + // the engine to "correct" already-corrected text. + let bundleKey = focusModel.snapshot.bundleIdentifier ?? "" + lastSubmittedPrefix[bundleKey] = accepted + + _ = writer.replacePrefix(originalLength: originalLength, with: accepted) + onCorrectionApplied?(originalPrefix, accepted) + } + + // MARK: - Gating + + private func passesGate(snapshot: FocusSnapshot, context: FocusedInputSnapshot) -> Bool { + guard isCorrectionEnabled() else { return false } + guard !context.isSecure else { return false } + guard context.selection.length == 0 else { return false } + guard TerminalAppDetector.isTerminal(bundleIdentifier: snapshot.bundleIdentifier) == false else { + return false + } + guard !isAutocompleteBusy() else { return false } + guard context.precedingText.count >= minimumPrefixCharacterCount else { return false } + guard context.precedingText.count <= maximumPrefixCharacterCount else { return false } + return correctionEngine.isAvailable + } + + // MARK: - De-dup key + + /// Compound key used to suppress duplicate settled events from a single sink. Any of + /// these changing means "the user has done something interesting since last time." + private struct SettledKey: Equatable { + let bundleIdentifier: String? + let precedingText: String + let selectionLength: Int + let isSecure: Bool + } +} diff --git a/Cotabby/App/Core/CotabbyAppEnvironment.swift b/Cotabby/App/Core/CotabbyAppEnvironment.swift index 9e531d0..a5432cd 100644 --- a/Cotabby/App/Core/CotabbyAppEnvironment.swift +++ b/Cotabby/App/Core/CotabbyAppEnvironment.swift @@ -23,6 +23,7 @@ final class CotabbyAppEnvironment { let foundationModelAvailabilityService: FoundationModelAvailabilityService let clipboardContextProvider: ClipboardContextProvider let suggestionCoordinator: SuggestionCoordinator + let prefixCorrectionCoordinator: PrefixCorrectionCoordinator let welcomeCoordinator: WelcomeCoordinator let settingsCoordinator: SettingsCoordinator let activationIndicatorController: ActivationIndicatorController @@ -144,6 +145,34 @@ final class CotabbyAppEnvironment { configuration: configuration ) + let prefixCorrectionEngine: any PrefixCorrecting = { + #if canImport(FoundationModels) + if #available(macOS 26.0, *) { + return FoundationModelPrefixCorrectionEngine( + availabilityService: foundationModelAvailabilityService + ) + } + #endif + return UnavailablePrefixCorrectionEngine() + }() + let prefixCorrectionWriter = PrefixCorrectionWriter(suppressionController: suppressionController) + let prefixCorrectionCoordinator = PrefixCorrectionCoordinator( + focusModel: focusModel, + correctionEngine: prefixCorrectionEngine, + writer: prefixCorrectionWriter, + isCorrectionEnabled: { suggestionSettings.isPrefixAutoCorrectEnabled }, + isAutocompleteBusy: { [weak suggestionCoordinator] in + guard let state = suggestionCoordinator?.state else { return false } + switch state { + case .debouncing, .generating: + return true + default: + return false + } + } + ) + prefixCorrectionCoordinator.start() + self.permissionManager = permissionManager self.runtimeModel = runtimeModel self.modelDownloadManager = modelDownloadManager @@ -156,6 +185,7 @@ final class CotabbyAppEnvironment { self.foundationModelAvailabilityService = foundationModelAvailabilityService self.clipboardContextProvider = clipboardContextProvider self.suggestionCoordinator = suggestionCoordinator + self.prefixCorrectionCoordinator = prefixCorrectionCoordinator self.welcomeCoordinator = welcomeCoordinator self.settingsCoordinator = settingsCoordinator self.activationIndicatorController = activationIndicatorController @@ -163,6 +193,14 @@ final class CotabbyAppEnvironment { ? FocusDebugOverlayController() : nil + // Surface prefix auto-corrections in the debug overlay when it's active. Production builds + // leave the hook nil so no extra work happens per correction. + if let focusDebugOverlayController = self.focusDebugOverlayController { + prefixCorrectionCoordinator.onCorrectionApplied = { [weak focusDebugOverlayController] original, corrected in + focusDebugOverlayController?.recordAutoCorrect(original: original, corrected: corrected) + } + } + // Update the AX polling timer whenever the user changes the poll interval setting. suggestionSettings.$focusPollIntervalMilliseconds .removeDuplicates() diff --git a/Cotabby/Models/SuggestionSettingsModel.swift b/Cotabby/Models/SuggestionSettingsModel.swift index a9d4385..dc07c82 100644 --- a/Cotabby/Models/SuggestionSettingsModel.swift +++ b/Cotabby/Models/SuggestionSettingsModel.swift @@ -29,6 +29,7 @@ final class SuggestionSettingsModel: ObservableObject { @Published private(set) var acceptanceKeyLabel: String @Published private(set) var fullAcceptanceKeyCode: CGKeyCode @Published private(set) var fullAcceptanceKeyLabel: String + @Published private(set) var isPrefixAutoCorrectEnabled: Bool private let userDefaults: UserDefaults private static let isGloballyEnabledDefaultsKey = "cotabbyGloballyEnabled" @@ -49,6 +50,7 @@ final class SuggestionSettingsModel: ObservableObject { private static let acceptanceKeyLabelDefaultsKey = "cotabbyAcceptanceKeyLabel" private static let fullAcceptanceKeyCodeDefaultsKey = "cotabbyFullAcceptanceKeyCode" private static let fullAcceptanceKeyLabelDefaultsKey = "cotabbyFullAcceptanceKeyLabel" + private static let prefixAutoCorrectEnabledDefaultsKey = "cotabbyPrefixAutoCorrectEnabled" static let defaultAcceptanceKeyCode: CGKeyCode = 48 static let defaultAcceptanceKeyLabel = "Tab" @@ -138,6 +140,9 @@ final class SuggestionSettingsModel: ObservableObject { let resolvedFullAcceptanceKeyLabel = userDefaults.string(forKey: Self.fullAcceptanceKeyLabelDefaultsKey) ?? Self.defaultFullAcceptanceKeyLabel + let resolvedPrefixAutoCorrectEnabled = + userDefaults.object(forKey: Self.prefixAutoCorrectEnabledDefaultsKey) as? Bool ?? false + isGloballyEnabled = resolvedGloballyEnabled disabledAppRules = resolvedDisabledAppRules showIndicator = resolvedShowIndicator @@ -155,6 +160,7 @@ final class SuggestionSettingsModel: ObservableObject { acceptanceKeyLabel = resolvedAcceptanceKeyLabel fullAcceptanceKeyCode = resolvedFullAcceptanceKeyCode fullAcceptanceKeyLabel = resolvedFullAcceptanceKeyLabel + isPrefixAutoCorrectEnabled = resolvedPrefixAutoCorrectEnabled userDefaults.set(resolvedGloballyEnabled, forKey: Self.isGloballyEnabledDefaultsKey) persistDisabledAppRules(resolvedDisabledAppRules) @@ -173,6 +179,7 @@ final class SuggestionSettingsModel: ObservableObject { userDefaults.set(resolvedAcceptanceKeyLabel, forKey: Self.acceptanceKeyLabelDefaultsKey) userDefaults.set(Int(resolvedFullAcceptanceKeyCode), forKey: Self.fullAcceptanceKeyCodeDefaultsKey) userDefaults.set(resolvedFullAcceptanceKeyLabel, forKey: Self.fullAcceptanceKeyLabelDefaultsKey) + userDefaults.set(resolvedPrefixAutoCorrectEnabled, forKey: Self.prefixAutoCorrectEnabledDefaultsKey) } /// Legacy compatibility shim. Reads through to `showIndicator`. @@ -338,6 +345,14 @@ final class SuggestionSettingsModel: ObservableObject { } } + // MARK: - Prefix auto-correct + + func setPrefixAutoCorrectEnabled(_ enabled: Bool) { + guard isPrefixAutoCorrectEnabled != enabled else { return } + isPrefixAutoCorrectEnabled = enabled + userDefaults.set(enabled, forKey: Self.prefixAutoCorrectEnabledDefaultsKey) + } + func setShowIndicator(_ show: Bool) { guard showIndicator != show else { return diff --git a/Cotabby/Models/SuggestionSubsystemContracts.swift b/Cotabby/Models/SuggestionSubsystemContracts.swift index 763be75..47c37ad 100644 --- a/Cotabby/Models/SuggestionSubsystemContracts.swift +++ b/Cotabby/Models/SuggestionSubsystemContracts.swift @@ -85,6 +85,20 @@ protocol SuggestionOverlayControlling: AnyObject { func hide(reason: String) } +@MainActor +protocol PrefixCorrecting: AnyObject { + /// Whether the engine can currently service a correction request (e.g., Apple Intelligence + /// is downloaded and available). Used by the coordinator to skip cycles when the backend is + /// down rather than queueing requests that will throw. + var isAvailable: Bool { get } + + /// Returns a proposed correction of `prefix`. The implementation must not perform any user- + /// visible side effects — the coordinator owns the apply-or-drop decision after the safety + /// filter runs. Throws on backend failures; returns `nil` if the model returned an empty or + /// unusable response. + func proposeCorrection(for prefix: String) async throws -> String? +} + @MainActor protocol VisualContextCoordinating: AnyObject { var status: VisualContextStatus { get } diff --git a/Cotabby/Services/Runtime/FoundationModelPrefixCorrectionEngine.swift b/Cotabby/Services/Runtime/FoundationModelPrefixCorrectionEngine.swift new file mode 100644 index 0000000..fbd50c6 --- /dev/null +++ b/Cotabby/Services/Runtime/FoundationModelPrefixCorrectionEngine.swift @@ -0,0 +1,131 @@ +import Foundation +import Logging + +#if canImport(FoundationModels) +import FoundationModels +#endif + +/// File overview: +/// Adapts Apple's on-device Foundation Models framework to Cotabby's `PrefixCorrecting` +/// capability. The coordinator uses this to ask Apple Intelligence for a typo-fixed version +/// of the user's currently-typed prefix; the safety filter downstream decides whether the +/// returned text is conservative enough to apply. +/// +/// Why Apple Intelligence only: prefix auto-correct demands tight instruction-following +/// (no rephrasing, no capitalization changes, no extra tokens). The bundled local llama +/// model isn't reliable at this task and would silently rewrite the user's prose. Routing +/// is deliberately one-engine for v1 — additional backends can adopt the protocol later. +#if canImport(FoundationModels) +@available(macOS 26.0, *) +@MainActor +final class FoundationModelPrefixCorrectionEngine { + private let availabilityService: FoundationModelAvailabilityService + + init(availabilityService: FoundationModelAvailabilityService) { + self.availabilityService = availabilityService + } + + var isAvailable: Bool { + availabilityService.refresh() + return availabilityService.isAvailable + } + + func proposeCorrection(for prefix: String) async throws -> String? { + availabilityService.refresh() + guard availabilityService.isAvailable else { + let message = availabilityService.userVisibleMessage + CotabbyLogger.suggestion.debug("Prefix-correction unavailable: \(message)") + return nil + } + guard let model = availabilityService.systemLanguageModel else { + return nil + } + + let trimmedPrefix = prefix.trimmingCharacters(in: .whitespacesAndNewlines) + guard !trimmedPrefix.isEmpty else { return nil } + + do { + let startTime = Date() + let session = LanguageModelSession( + model: model, + instructions: Self.correctionInstructions + ) + // Deterministic decoding so the same prefix yields the same correction and the safety + // filter can reason about the output shape predictably. + let options = GenerationOptions( + sampling: .greedy, + temperature: 0.0, + maximumResponseTokens: tokenBudget(for: prefix) + ) + let response = try await session.respond(to: prefix, options: options) + try Task.checkCancellation() + + let raw = response.content + let cleaned = strippedResponse(raw) + let latencyMs = Int(Date().timeIntervalSince(startTime) * 1000) + CotabbyLogger.suggestion.debug( + "Prefix-correction: in=\(prefix.count) chars, out=\(cleaned.count) chars, latency=\(latencyMs)ms" + ) + return cleaned.isEmpty ? nil : cleaned + } catch is CancellationError { + throw SuggestionClientError.cancelled + } catch let error as LanguageModelSession.GenerationError { + CotabbyLogger.suggestion.debug("Prefix-correction generation error: \(error.localizedDescription)") + // Swallow into nil rather than throwing — a failed correction should be invisible to + // the user, not surfaced as an autocomplete error. + return nil + } catch { + CotabbyLogger.suggestion.debug("Prefix-correction unexpected error: \(error.localizedDescription)") + return nil + } + } + + // MARK: - Prompting + + private static let correctionInstructions: String = """ + You correct spelling typos in text from an inline autocomplete tool. + + Rules — apply without exception: + - Only fix obvious misspellings of individual words. + - Never add, remove, reorder, or rephrase words. + - Never change capitalization, punctuation, spacing, or line breaks. + - Never add quotes, prefixes, suffixes, explanations, or commentary. + - If there are no typos, return the input unchanged. + + Output only the corrected text. + """ + + /// Token budget sized to "input length plus a little slack" because typo-fixes do not grow + /// the text appreciably. Anything longer is already suspicious and the safety filter will + /// reject it, but a tighter budget also lets us cut off runaway generation. + private func tokenBudget(for prefix: String) -> Int { + // ~3 chars per token, generous upward rounding plus 16 tokens of slack. + max(32, prefix.count / 3 + 16) + } + + /// Models occasionally bracket their output in quotes or prepend "Corrected: ". The safety + /// filter would reject those, but stripping the most common wrappers here makes the filter's + /// real-world hit rate noticeably higher. + private func strippedResponse(_ raw: String) -> String { + var text = raw.trimmingCharacters(in: .whitespacesAndNewlines) + if let openQuote = text.first, openQuote == "\"" || openQuote == "“", + let closeQuote = text.last, closeQuote == "\"" || closeQuote == "”", + text.count >= 2 { + text = String(text.dropFirst().dropLast()).trimmingCharacters(in: .whitespacesAndNewlines) + } + return text + } +} + +@available(macOS 26.0, *) +extension FoundationModelPrefixCorrectionEngine: PrefixCorrecting {} +#endif + +/// Always-unavailable fallback used when the FoundationModels SDK is missing or the +/// host macOS is older than the supported Apple Intelligence release. The coordinator +/// gates on `isAvailable` before calling, so this drops every correction silently. +@MainActor +final class UnavailablePrefixCorrectionEngine: PrefixCorrecting { + var isAvailable: Bool { false } + func proposeCorrection(for prefix: String) async throws -> String? { nil } +} diff --git a/Cotabby/Services/Suggestion/PrefixCorrectionWriter.swift b/Cotabby/Services/Suggestion/PrefixCorrectionWriter.swift new file mode 100644 index 0000000..71dd444 --- /dev/null +++ b/Cotabby/Services/Suggestion/PrefixCorrectionWriter.swift @@ -0,0 +1,84 @@ +import ApplicationServices +import Foundation +import Logging + +/// File overview: +/// Performs the synthetic-event write that replaces a user's typed prefix with a corrected +/// version. Mirrors `SuggestionInserter`'s pattern: post raw CGEvents directly to the HID +/// event tap and register the expected key-down count with `InputSuppressionController` +/// so the global input monitor ignores its own writes. +/// +/// Write strategy: backspace × original-length, then a single Unicode keystroke for the +/// corrected text. This is the "fake delete all, retype" approach — simple, app-agnostic, +/// and avoids the diff/arrow-key bookkeeping that an in-place edit would require. The +/// coordinator gates against very long prefixes so the backspace burst stays bounded. +@MainActor +final class PrefixCorrectionWriter { + private static let backspaceVirtualKey: CGKeyCode = 51 // kVK_Delete + + private let suppressionController: InputSuppressionController + + init(suppressionController: InputSuppressionController) { + self.suppressionController = suppressionController + } + + /// Deletes the last `originalLength` graphemes from the focused field and types + /// `correctedPrefix` in their place. Returns false if any synthetic event could not be + /// created or if the inputs are degenerate. + func replacePrefix(originalLength: Int, with correctedPrefix: String) -> Bool { + let normalized = correctedPrefix.replacingOccurrences(of: "\r", with: "") + guard originalLength > 0, !normalized.isEmpty else { + CotabbyLogger.suggestion.warning("Prefix-correction write skipped: empty input") + return false + } + + // The unicode keystroke event counts as one key-down. Total suppression budget is the + // backspace burst plus that one event. + let expectedKeyDowns = originalLength + 1 + suppressionController.registerSyntheticInsertion(expectedKeyDownCount: expectedKeyDowns) + + for _ in 0.. Bool { + guard let keyDown = CGEvent(keyboardEventSource: nil, virtualKey: Self.backspaceVirtualKey, keyDown: true), + let keyUp = CGEvent(keyboardEventSource: nil, virtualKey: Self.backspaceVirtualKey, keyDown: false) + else { + return false + } + keyDown.post(tap: .cghidEventTap) + keyUp.post(tap: .cghidEventTap) + return true + } + + private func postUnicodeString(_ text: String) -> Bool { + guard let keyDown = CGEvent(keyboardEventSource: nil, virtualKey: 0, keyDown: true), + let keyUp = CGEvent(keyboardEventSource: nil, virtualKey: 0, keyDown: false) + else { + return false + } + let utf16 = Array(text.utf16) + keyDown.keyboardSetUnicodeString(stringLength: utf16.count, unicodeString: utf16) + keyUp.keyboardSetUnicodeString(stringLength: utf16.count, unicodeString: utf16) + keyDown.post(tap: .cghidEventTap) + keyUp.post(tap: .cghidEventTap) + return true + } +} diff --git a/Cotabby/Services/UI/FocusDebugOverlayController.swift b/Cotabby/Services/UI/FocusDebugOverlayController.swift index 08b8c7e..13ee1da 100644 --- a/Cotabby/Services/UI/FocusDebugOverlayController.swift +++ b/Cotabby/Services/UI/FocusDebugOverlayController.swift @@ -31,6 +31,8 @@ final class FocusDebugOverlayController { private var latestVisualContextStatus: VisualContextStatus = .idle private var latestVisualContextExcerptCharacterCount: Int? private var latestPollEvent: FocusPollingEvent? + private var latestAutoCorrect: AutoCorrectDebugEvent? + private var autoCorrectExpiryTask: Task? func update(for snapshot: FocusSnapshot) { guard let context = snapshot.context else { @@ -62,11 +64,28 @@ final class FocusDebugOverlayController { renderBottomStatusPanel() } + /// Flashes the most recent prefix auto-correct in the bottom panel, then clears it after a few + /// seconds so the panel reflects "just corrected" rather than a stale event. + func recordAutoCorrect(original: String, corrected: String) { + latestAutoCorrect = AutoCorrectDebugEvent(original: original, corrected: corrected, timestamp: Date()) + renderBottomStatusPanel() + + autoCorrectExpiryTask?.cancel() + autoCorrectExpiryTask = Task { [weak self] in + try? await Task.sleep(for: .seconds(6)) + guard !Task.isCancelled else { return } + self?.latestAutoCorrect = nil + self?.renderBottomStatusPanel() + } + } + func hide() { hideFocusGeometry() latestPollEvent = nil latestVisualContextStatus = .idle latestVisualContextExcerptCharacterCount = nil + autoCorrectExpiryTask?.cancel() + latestAutoCorrect = nil bottomStatusPanel.orderOut(nil) } @@ -139,6 +158,7 @@ final class FocusDebugOverlayController { visualContextStatus: latestVisualContextStatus, excerptCharacterCount: latestVisualContextExcerptCharacterCount, pollEvent: latestPollEvent, + autoCorrect: latestAutoCorrect, maxWidth: maxWidth )) contentView.layoutSubtreeIfNeeded() @@ -162,7 +182,7 @@ final class FocusDebugOverlayController { } private var shouldShowBottomStatusPanel: Bool { - latestVisualContextStatus != .idle || latestPollEvent != nil + latestVisualContextStatus != .idle || latestPollEvent != nil || latestAutoCorrect != nil } // MARK: - Helpers @@ -240,6 +260,7 @@ private struct BottomDebugStatusView: View { let visualContextStatus: VisualContextStatus let excerptCharacterCount: Int? let pollEvent: FocusPollingEvent? + let autoCorrect: AutoCorrectDebugEvent? let maxWidth: CGFloat private var stages: [VisualContextDebugStage] { @@ -313,6 +334,38 @@ private struct BottomDebugStatusView: View { .lineLimit(1) } } + + if let autoCorrect { + Divider() + .overlay(Color.white.opacity(0.16)) + + HStack(spacing: 7) { + Image(systemName: "checkmark.circle.fill") + .font(.system(size: 9)) + .foregroundStyle(.orange) + + Text("Auto-correct") + .font(.system(size: 10, weight: .bold, design: .monospaced)) + .foregroundStyle(.orange) + + Text(Self.tail(autoCorrect.original)) + .font(.system(size: 10, weight: .medium, design: .monospaced)) + .foregroundStyle(.white.opacity(0.6)) + .lineLimit(1) + .strikethrough() + + Image(systemName: "arrow.right") + .font(.system(size: 8)) + .foregroundStyle(.white.opacity(0.5)) + + Text(Self.tail(autoCorrect.corrected)) + .font(.system(size: 10, weight: .semibold, design: .monospaced)) + .foregroundStyle(.green) + .lineLimit(1) + + Spacer(minLength: 0) + } + } } .padding(.horizontal, 12) .padding(.vertical, 10) @@ -394,6 +447,20 @@ private struct BottomDebugStatusView: View { return .blocked } } + + /// Shows the trailing slice of a prefix (where edits usually land), with newlines collapsed, + /// so the row stays single-line in the panel. + private static func tail(_ text: String, max: Int = 32) -> String { + let collapsed = text.replacingOccurrences(of: "\n", with: "⏎") + return collapsed.count <= max ? collapsed : "…" + String(collapsed.suffix(max)) + } +} + +/// One prefix auto-correct event, surfaced in the debug overlay's bottom panel. +private struct AutoCorrectDebugEvent { + let original: String + let corrected: String + let timestamp: Date } private struct VisualContextStagePill: View { diff --git a/Cotabby/Support/PrefixCorrectionFilter.swift b/Cotabby/Support/PrefixCorrectionFilter.swift new file mode 100644 index 0000000..56c5709 --- /dev/null +++ b/Cotabby/Support/PrefixCorrectionFilter.swift @@ -0,0 +1,146 @@ +import Foundation + +/// File overview: +/// Decides whether an LLM-proposed correction to the user's prefix is "typo-shaped" +/// enough to apply, or whether it has drifted into rewording, repunctuation, or +/// recapitalization that the user did not ask for. +/// +/// The filter is the safety net for prefix auto-correct. Even with a tight prompt the +/// model will sometimes rephrase, change capitalization, or "improve" punctuation. +/// Those changes are silently destructive because the write replaces the user's typed +/// text without a diff UI. Only changes that look like single-word spelling fixes are +/// allowed through. +enum PrefixCorrectionFilter { + /// Returns `proposed` when it is a safe typo-fix of `original`, or `nil` to drop it. + /// + /// Rules — all must hold: + /// - Same number of word/separator tokens, in the same order. + /// - Inter-word separators (whitespace, punctuation) are byte-identical. + /// - For each word pair that differs: + /// - Both words are at least `minimumWordLength` characters. + /// - Case shape matches (all-lower, all-upper, capitalized, or mixed). + /// - Edit distance ≤ `max(2, length / 3)` using the longer of the two words. + static func acceptedCorrection(original: String, proposed: String) -> String? { + guard original != proposed else { return nil } + + let originalTokens = tokenize(original) + let proposedTokens = tokenize(proposed) + guard originalTokens.count == proposedTokens.count else { return nil } + + for (originalToken, proposedToken) in zip(originalTokens, proposedTokens) { + switch (originalToken, proposedToken) { + case let (.separator(originalRun), .separator(proposedRun)): + guard originalRun == proposedRun else { return nil } + case let (.word(originalWord), .word(proposedWord)): + guard isTypoShapedChange(original: originalWord, proposed: proposedWord) else { return nil } + default: + // Boundary mismatch: a word in one stream lines up with a separator in the other. + return nil + } + } + + return proposed + } + + // MARK: - Tokenization + + private static let minimumWordLength = 3 + + private enum Token: Equatable { + case word(String) + case separator(String) + } + + /// Splits `text` into alternating runs of Unicode letters and everything else. + private static func tokenize(_ text: String) -> [Token] { + var tokens: [Token] = [] + var current = "" + var currentIsWord = false + + for scalar in text.unicodeScalars { + let scalarIsLetter = CharacterSet.letters.contains(scalar) + if current.isEmpty { + current.unicodeScalars.append(scalar) + currentIsWord = scalarIsLetter + continue + } + + if scalarIsLetter == currentIsWord { + current.unicodeScalars.append(scalar) + } else { + tokens.append(currentIsWord ? .word(current) : .separator(current)) + current = String(scalar) + currentIsWord = scalarIsLetter + } + } + + if !current.isEmpty { + tokens.append(currentIsWord ? .word(current) : .separator(current)) + } + + return tokens + } + + // MARK: - Per-word shape check + + private static func isTypoShapedChange(original: String, proposed: String) -> Bool { + if original == proposed { return true } + guard original.count >= minimumWordLength, proposed.count >= minimumWordLength else { + return false + } + guard caseShape(of: original) == caseShape(of: proposed) else { return false } + let distance = levenshteinDistance(original.lowercased(), proposed.lowercased()) + let allowed = Swift.max(2, Swift.max(original.count, proposed.count) / 3) + return distance <= allowed + } + + private enum CaseShape: Equatable { + case allLower + case allUpper + case capitalized + case mixed + } + + /// Categorizes a word by its capitalization pattern so the filter can reject changes + /// that swap between shapes (the model "fixing" capitalization the user didn't ask for). + private static func caseShape(of word: String) -> CaseShape { + let letters = word.filter { $0.isLetter } + guard let first = letters.first else { return .mixed } + + let rest = letters.dropFirst() + let allLower = letters.allSatisfy(\.isLowercase) + if allLower { return .allLower } + let allUpper = letters.allSatisfy(\.isUppercase) + if allUpper { return .allUpper } + if first.isUppercase, rest.allSatisfy(\.isLowercase) { return .capitalized } + return .mixed + } + + /// Standard two-row Levenshtein. Words are short, so the simple implementation is fine. + private static func levenshteinDistance(_ lhs: String, _ rhs: String) -> Int { + let lhsChars = Array(lhs) + let rhsChars = Array(rhs) + let lhsLength = lhsChars.count + let rhsLength = rhsChars.count + if lhsLength == 0 { return rhsLength } + if rhsLength == 0 { return lhsLength } + + var previous = Array(0...rhsLength) + var current = Array(repeating: 0, count: rhsLength + 1) + + for row in 1...lhsLength { + current[0] = row + for col in 1...rhsLength { + let cost = lhsChars[row - 1] == rhsChars[col - 1] ? 0 : 1 + current[col] = Swift.min( + previous[col] + 1, + current[col - 1] + 1, + previous[col - 1] + cost + ) + } + swap(&previous, ¤t) + } + + return previous[rhsLength] + } +} diff --git a/Cotabby/UI/MenuBarView.swift b/Cotabby/UI/MenuBarView.swift index 73fd2c9..4dcc340 100644 --- a/Cotabby/UI/MenuBarView.swift +++ b/Cotabby/UI/MenuBarView.swift @@ -75,6 +75,14 @@ struct MenuBarView: View { .toggleStyle(.switch) .controlSize(.small) + Toggle("Auto-Correct Typos", isOn: prefixAutoCorrectEnabledBinding) + .toggleStyle(.switch) + .controlSize(.small) + .disabled(!foundationModelAvailabilityService.isAvailable) + .help(foundationModelAvailabilityService.isAvailable + ? "Apple Intelligence rewrites obvious typos in the text you just typed." + : "Requires Apple Intelligence: \(foundationModelAvailabilityService.userVisibleMessage)") + if let application = focusModel.latestExternalApplication, !TerminalAppDetector.isTerminal(bundleIdentifier: application.bundleIdentifier) { Toggle("Enable in \(application.applicationName)", isOn: appEnabledBinding(for: application)) @@ -264,6 +272,13 @@ struct MenuBarView: View { ) } + private var prefixAutoCorrectEnabledBinding: Binding { + Binding( + get: { suggestionSettings.isPrefixAutoCorrectEnabled }, + set: { suggestionSettings.setPrefixAutoCorrectEnabled($0) } + ) + } + private func appEnabledBinding(for application: FocusedApplicationIdentity) -> Binding { Binding( get: { diff --git a/Cotabby/UI/SettingsView.swift b/Cotabby/UI/SettingsView.swift index 4f21b9d..4a0266e 100644 --- a/Cotabby/UI/SettingsView.swift +++ b/Cotabby/UI/SettingsView.swift @@ -144,6 +144,12 @@ struct SettingsView: View { Section("General") { Toggle("Enable Globally", isOn: globallyEnabledBinding) + Toggle("Auto-Correct Typos", isOn: prefixAutoCorrectEnabledBinding) + .disabled(!foundationModelAvailabilityService.isAvailable) + .help(foundationModelAvailabilityService.isAvailable + ? "Apple Intelligence rewrites obvious typos in the text you just typed." + : "Requires Apple Intelligence: \(foundationModelAvailabilityService.userVisibleMessage)") + Toggle("Show Indicator", isOn: showIndicatorBinding) Toggle("Allow Multi-line Suggestions", isOn: multiLineEnabledBinding) @@ -626,6 +632,13 @@ struct SettingsView: View { ) } + private var prefixAutoCorrectEnabledBinding: Binding { + Binding( + get: { suggestionSettings.isPrefixAutoCorrectEnabled }, + set: { suggestionSettings.setPrefixAutoCorrectEnabled($0) } + ) + } + private var debounceMillisecondsBinding: Binding { Binding( get: { suggestionSettings.debounceMilliseconds }, diff --git a/CotabbyTests/PrefixCorrectionFilterTests.swift b/CotabbyTests/PrefixCorrectionFilterTests.swift new file mode 100644 index 0000000..fcf36cb --- /dev/null +++ b/CotabbyTests/PrefixCorrectionFilterTests.swift @@ -0,0 +1,187 @@ +import XCTest +@testable import Cotabby + +final class PrefixCorrectionFilterTests: XCTestCase { + + // MARK: - Identity short-circuit + + func test_identicalText_returnsNil() { + XCTAssertNil(PrefixCorrectionFilter.acceptedCorrection( + original: "the quick brown fox", + proposed: "the quick brown fox" + )) + } + + // MARK: - Accepted typo fixes + + func test_singleWordTypo_isAccepted() { + XCTAssertEqual( + PrefixCorrectionFilter.acceptedCorrection(original: "teh", proposed: "the"), + "the" + ) + } + + func test_multiWordPrefix_singleTypoFix_isAccepted() { + XCTAssertEqual( + PrefixCorrectionFilter.acceptedCorrection( + original: "teh quick brown fox", + proposed: "the quick brown fox" + ), + "the quick brown fox" + ) + } + + func test_multipleTypoFixes_isAccepted() { + XCTAssertEqual( + PrefixCorrectionFilter.acceptedCorrection( + original: "teh quick brwn fox", + proposed: "the quick brown fox" + ), + "the quick brown fox" + ) + } + + func test_longerWordTypo_isAccepted() { + XCTAssertEqual( + PrefixCorrectionFilter.acceptedCorrection( + original: "definately yes", + proposed: "definitely yes" + ), + "definitely yes" + ) + } + + func test_capitalizedTypoPreservingShape_isAccepted() { + XCTAssertEqual( + PrefixCorrectionFilter.acceptedCorrection(original: "Helo", proposed: "Hello"), + "Hello" + ) + } + + func test_allUppercaseTypoPreservingShape_isAccepted() { + XCTAssertEqual( + PrefixCorrectionFilter.acceptedCorrection(original: "TEH", proposed: "THE"), + "THE" + ) + } + + // MARK: - Rejected: structural changes + + func test_addedWord_isRejected() { + XCTAssertNil(PrefixCorrectionFilter.acceptedCorrection( + original: "the quick fox", + proposed: "the quick brown fox" + )) + } + + func test_removedWord_isRejected() { + XCTAssertNil(PrefixCorrectionFilter.acceptedCorrection( + original: "the the quick fox", + proposed: "the quick fox" + )) + } + + func test_addedTrailingPunctuation_isRejected() { + XCTAssertNil(PrefixCorrectionFilter.acceptedCorrection( + original: "the quick fox", + proposed: "the quick fox." + )) + } + + func test_changedSeparator_isRejected() { + // Model "fixed" comma to comma+space. + XCTAssertNil(PrefixCorrectionFilter.acceptedCorrection( + original: "hello,world", + proposed: "hello, world" + )) + } + + func test_collapsedWhitespace_isRejected() { + XCTAssertNil(PrefixCorrectionFilter.acceptedCorrection( + original: "hello world", + proposed: "hello world" + )) + } + + // MARK: - Rejected: case changes + + func test_capitalizationAdded_isRejected() { + // Model promoted lowercase start to capital — rewriting voice, not fixing a typo. + XCTAssertNil(PrefixCorrectionFilter.acceptedCorrection( + original: "the quick fox", + proposed: "The quick fox" + )) + } + + func test_caseShapeMismatch_isRejected() { + XCTAssertNil(PrefixCorrectionFilter.acceptedCorrection( + original: "teh", + proposed: "The" + )) + } + + // MARK: - Rejected: too-short words + + func test_twoCharWordChange_isRejected() { + // "im" → "I'm" would also fail on separators, but a pure two-char change is itself rejected. + XCTAssertNil(PrefixCorrectionFilter.acceptedCorrection( + original: "im here", + proposed: "is here" + )) + } + + func test_singleCharWordChange_isRejected() { + XCTAssertNil(PrefixCorrectionFilter.acceptedCorrection( + original: "i am", + proposed: "a am" + )) + } + + // MARK: - Rejected: edit distance too large + + func test_wordReplacedWithUnrelatedWord_isRejected() { + XCTAssertNil(PrefixCorrectionFilter.acceptedCorrection( + original: "the cat sat", + proposed: "the dog sat" + )) + } + + func test_donutToDoughnut_isRejected() { + // Real word change, not a typo. Distance 3, threshold = max(2, max(5,8)/3) = 2. + XCTAssertNil(PrefixCorrectionFilter.acceptedCorrection( + original: "donut shop", + proposed: "doughnut shop" + )) + } + + // MARK: - Mixed scenarios + + func test_typoFixSurroundedByPunctuation_isAccepted() { + XCTAssertEqual( + PrefixCorrectionFilter.acceptedCorrection( + original: "Hello, teh world!", + proposed: "Hello, the world!" + ), + "Hello, the world!" + ) + } + + func test_typoFixWithNewlines_isAccepted() { + XCTAssertEqual( + PrefixCorrectionFilter.acceptedCorrection( + original: "first line\nteh second", + proposed: "first line\nthe second" + ), + "first line\nthe second" + ) + } + + func test_emptyStrings_returnNil() { + XCTAssertNil(PrefixCorrectionFilter.acceptedCorrection(original: "", proposed: "")) + } + + func test_emptyToNonEmpty_isRejected() { + // Token counts differ (0 vs 1). + XCTAssertNil(PrefixCorrectionFilter.acceptedCorrection(original: "", proposed: "hello")) + } +}