Skip to content
4 changes: 4 additions & 0 deletions Cotabby.xcodeproj/project.pbxproj
Original file line number Diff line number Diff line change
Expand Up @@ -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 */; };
Expand Down Expand Up @@ -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 = "<group>"; };
H10000112FC0000100FFF011 /* PrefixCorrectionFilterTests.swift */ = {isa = PBXFileReference; includeInIndex = 1; lastKnownFileType = sourcecode.swift; path = PrefixCorrectionFilterTests.swift; sourceTree = "<group>"; };
5A39EC1A44E9160E13E2AE77 /* SuggestionAvailabilityEvaluatorTests.swift */ = {isa = PBXFileReference; includeInIndex = 1; lastKnownFileType = sourcecode.swift; path = SuggestionAvailabilityEvaluatorTests.swift; sourceTree = "<group>"; };
7A1E3B0D2FC8A00100CARET2 /* AXTextGeometryResolverTests.swift */ = {isa = PBXFileReference; includeInIndex = 1; lastKnownFileType = sourcecode.swift; path = AXTextGeometryResolverTests.swift; sourceTree = "<group>"; };
8B1121FFDAD30C7F62E15289 /* SuggestionRequestFactoryTests.swift */ = {isa = PBXFileReference; includeInIndex = 1; lastKnownFileType = sourcecode.swift; path = SuggestionRequestFactoryTests.swift; sourceTree = "<group>"; };
Expand Down Expand Up @@ -157,6 +159,7 @@
G20000112FB0000100FFF011 /* ClipboardRelevanceFilterTests.swift */,
G10000122FB0000100FFF012 /* ClipboardContentDistillerTests.swift */,
G10000112FB0000100FFF011 /* WordCountFormatterTests.swift */,
H10000112FC0000100FFF011 /* PrefixCorrectionFilterTests.swift */,
);
path = CotabbyTests;
sourceTree = "<group>";
Expand Down Expand Up @@ -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;
};
Expand Down
213 changes: 213 additions & 0 deletions Cotabby/App/Coordinators/PrefixCorrectionCoordinator.swift
Original file line number Diff line number Diff line change
@@ -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<AnyCancellable>()
private var latestWorkID: UInt64 = 0
private var inflightTask: Task<Void, Never>?
/// 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
Comment on lines +124 to +129
Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P1 Dedup entry written before the async call permanently suppresses retry on failure

lastSubmittedPrefix[bundleKey] = context.precedingText is written at line 127, before the Task is spawned. If runCorrection later exits through any error path (the catch block in the engine, or Task.isCancelled), the entry is never cleared. The next time the same prefix settles — e.g., after focus moves away and back while the user left the typo unchanged — handleSettled finds lastSubmittedPrefix[bundleKey] == context.precedingText and returns early, so the correction is permanently skipped for that session.

The intended guard is "we just wrote a correction and the publisher re-fired with the corrected text," but setting the entry before the call conflates that case with transient failures. Fix: move the write into runCorrection, setting it only after proposeCorrection returns (nil or a proposal), and skipping the write entirely on the catch path so errors are retryable.

Fix in Codex Fix in Claude Code


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
}
}
38 changes: 38 additions & 0 deletions Cotabby/App/Core/CotabbyAppEnvironment.swift
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand All @@ -156,13 +185,22 @@ final class CotabbyAppEnvironment {
self.foundationModelAvailabilityService = foundationModelAvailabilityService
self.clipboardContextProvider = clipboardContextProvider
self.suggestionCoordinator = suggestionCoordinator
self.prefixCorrectionCoordinator = prefixCorrectionCoordinator
self.welcomeCoordinator = welcomeCoordinator
self.settingsCoordinator = settingsCoordinator
self.activationIndicatorController = activationIndicatorController
self.focusDebugOverlayController = FocusDebugOverlayController.isEnabled
? 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()
Expand Down
15 changes: 15 additions & 0 deletions Cotabby/Models/SuggestionSettingsModel.swift
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand All @@ -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"
Expand Down Expand Up @@ -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
Expand All @@ -155,6 +160,7 @@ final class SuggestionSettingsModel: ObservableObject {
acceptanceKeyLabel = resolvedAcceptanceKeyLabel
fullAcceptanceKeyCode = resolvedFullAcceptanceKeyCode
fullAcceptanceKeyLabel = resolvedFullAcceptanceKeyLabel
isPrefixAutoCorrectEnabled = resolvedPrefixAutoCorrectEnabled

userDefaults.set(resolvedGloballyEnabled, forKey: Self.isGloballyEnabledDefaultsKey)
persistDisabledAppRules(resolvedDisabledAppRules)
Expand All @@ -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`.
Expand Down Expand Up @@ -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
Expand Down
Loading
Loading