Cut keystroke-to-ghost latency on the llama engine to under 60 ms

Cotabby/App/Coordinators/Suggestion/SuggestionCoordinator+Input.swift

@@ -1,3 +1,4 @@
+import CoreGraphics
 import Foundation
 import Logging
 
@@ -220,10 +221,19 @@ extension SuggestionCoordinator {
             // null the freshly primed session out from under us.
             await self.awaitCachedGenerationContextResetIfNeeded()
             let prewarmContext = FocusedInputContext(snapshot: rawContext, generation: 0)
+            // Build the prompt from the same context sources the first real request will read
+            // (`SuggestionCoordinator+Prediction`): the pinned clipboard and any ready screen
+            // excerpt. The prefill is only worth its decode if the next request shares its bytes;
+            // leaving these out made the first request diverge right after the surface line and
+            // re-decode everything the prewarm had already paid for.
             let request = SuggestionRequestFactory.buildRequest(
                 context: prewarmContext,
                 settings: settings,
                 configuration: configuration,
+                clipboardContext: self.pinnedClipboardContext(rawContext: rawContext),
+                visualContextSummary: self.permissionManager.screenRecordingGranted
+                    ? self.visualContextCoordinator.excerpt(for: prewarmContext)
+                    : nil,
                 historyExamples: self.historyExamples(for: prewarmContext)
             ).request
             await suggestionEngine.prewarm(for: request)
@@ -334,11 +344,40 @@ extension SuggestionCoordinator {
             // *before* the host app processes the keystroke, so a synchronous `refreshNow()` here
             // reads pre-keystroke text and feeds it into generation. The result is a suggestion
             // that looks like the typed character was ignored — see
-            // `schedulePredictionAfterHostPublishDelay` for the full rationale.
+            // `schedulePredictionAfterHostPublishDelay` for the full rationale. A plain typed
+            // character can still start generating now, against the text the host is about to
+            // publish; the poll then validates that bet instead of gating the request.
+            speculateOnTypedCharacter(event)
             schedulePredictionAfterHostPublishDelay()
         }
     }
 
+    /// Starts the keystroke-time speculative generation for a plain typed character. The base is
+    /// the last captured snapshot, extended by the letters a visible suggestion was typed through
+    /// (`session`), which the snapshot may not include yet. When it still misses an unpublished key,
+    /// the speculated text simply misses and the publish poll falls back to the ordinary request.
+    private func speculateOnTypedCharacter(_ event: CapturedInputEvent, after session: ActiveSuggestionSession? = nil) {
+        guard event.kind == .textMutation, let typed = Self.speculatedCharacters(for: event),
+              let raw = focusModel.snapshot.context else { return }
+        let typedThrough = session.flatMap { session -> String? in
+            guard session.baseContext.elementIdentifier == raw.elementIdentifier else { return nil }
+            return session.baseContext.precedingText + String(session.fullText.prefix(session.consumedCharacterCount))
+        }
+        dispatchSpeculativeTypingGeneration(rawContext: raw, typed: typed, knownPrecedingText: typedThrough)
+    }
+
+    /// The character the host will most likely insert for `event`, or nil when it cannot be told
+    /// from the event alone. A key event's Unicode string does not always carry Shift (synthetic
+    /// events report "s" for a shifted S), so a letter follows Shift XOR Caps Lock; uppercasing an
+    /// already-capital letter changes nothing. A shifted symbol depends on the keyboard layout
+    /// ("/" becomes "?" on US layouts), so it is left to the ordinary publish-gated request.
+    static func speculatedCharacters(for event: CapturedInputEvent) -> String? {
+        guard event.characters.count == 1, let character = event.characters.first else { return nil }
+        let shift = event.flags.contains(.maskShift)
+        guard character.isLetter else { return shift ? nil : event.characters }
+        return shift != event.flags.contains(.maskAlphaShift) ? event.characters.uppercased() : event.characters
+    }
+
     /// Maximum wall time we'll wait for the host app to publish post-keystroke AX before giving
     /// up and generating against whatever's there. Chosen empirically: long enough to cover
     /// Chrome's slower contenteditable publish on a busy page, short enough that the user can
@@ -354,7 +393,16 @@ extension SuggestionCoordinator {
     /// own tap callback before the host has even processed the key and therefore always misses. A
     /// short first retry catches fast-publishing native apps in ~10ms instead of making them wait a
     /// full steady interval; slow Chromium hosts fall through to `hostPublishPollIntervalMs` below.
-    private static let hostPublishFirstPollIntervalMs = 10
+    private static let hostPublishFirstPollIntervalMs = 6
+
+    /// Native hosts publish a keystroke to AX within roughly 10-40 ms, so the wait polls densely
+    /// during that window and only then backs off to `hostPublishPollIntervalMs` for slow Chromium
+    /// hosts. With a 30 ms step from the start, a publish that landed just after the first retry
+    /// waited a further ~30 ms (live TextEdit trace: key-to-request 50 ms instead of ~20 ms), the
+    /// largest remaining wait before generation. The dense phase costs a handful of ~2-3 ms AX
+    /// reads per keystroke and ends after `hostPublishFastPhaseMs`.
+    private static let hostPublishFastPollIntervalMs = 8
+    private static let hostPublishFastPhaseMs = 60
 
     /// Schedules a fresh prediction once the host app has actually published the new
     /// contenteditable text to AX. The previous fix waited a fixed 150ms — see PR #376 — but the
@@ -402,6 +450,27 @@ extension SuggestionCoordinator {
         }
     }
 
+    /// Records why a publish did not match the speculated text, without logging either text: the
+    /// lengths, selections, and the first differing UTF-16 offset are enough to tell a lagging base
+    /// snapshot from a host transformation (smart punctuation, auto-pairing, a shifted key).
+    private func logSpeculationMiss(expected: FocusedInputContext, published: FocusedInputSnapshot) {
+        let lhs = Array(expected.precedingText.utf16), rhs = Array(published.precedingText.utf16)
+        let shared = zip(lhs, rhs).prefix { $0 == $1 }.count
+        CotabbyLogger.suggestion.debug(
+            "Speculation missed the published text",
+            metadata: [
+                "stage": .string("speculation-missed"),
+                "same_session": .stringConvertible(expected.sessionIdentity == published.sessionIdentity),
+                "expected_utf16": .stringConvertible(lhs.count),
+                "published_utf16": .stringConvertible(rhs.count),
+                "first_difference": .stringConvertible(shared),
+                "expected_selection": .stringConvertible(expected.selection.location),
+                "published_selection": .stringConvertible(published.selection.location),
+                "trailing_matches": .stringConvertible(expected.trailingText == published.trailingText)
+            ]
+        )
+    }
+
     /// The AX state captured at keystroke time, against which the host-publish poll detects "the
     /// host has processed the key", plus the uptime anchor for folding the wait into the debounce.
     private struct HostPublishBaseline {
@@ -471,6 +540,9 @@ extension SuggestionCoordinator {
                 )
                 return
             }
+            if let expected = pendingSpeculativeContext, let published = currentContext {
+                logSpeculationMiss(expected: expected, published: published)
+            }
             pendingSpeculativeContext = nil
             schedulePrediction(
                 consumedDelayMilliseconds: Self.elapsedMilliseconds(since: baseline.keystrokeUptimeNanoseconds)
@@ -482,8 +554,11 @@ extension SuggestionCoordinator {
         // produces no AX change (rare but possible — e.g. dead-key composition) would never get
         // its next prediction.
         // We already waited the short first interval before entering this method; remaining polls
-        // use the steady cadence until the ceiling.
-        let interval = Self.hostPublishPollIntervalMs
+        // stay dense while native hosts usually publish, then use the steady cadence until the
+        // ceiling.
+        let interval = elapsedMs < Self.hostPublishFastPhaseMs
+            ? Self.hostPublishFastPollIntervalMs
+            : Self.hostPublishPollIntervalMs
         let nextElapsed = elapsedMs + interval
         guard nextElapsed < Self.hostPublishWaitCeilingMs else {
             if baseline.requiresTextChange {
@@ -546,6 +621,7 @@ extension SuggestionCoordinator {
                 clearDiagnostics: false
             )
             if event.shouldSchedulePrediction {
+                speculateOnTypedCharacter(event, after: session)
                 schedulePredictionAfterHostPublishDelay()
             }
             return false

Cotabby/App/Coordinators/Suggestion/SuggestionCoordinator+Prediction.swift

@@ -26,12 +26,14 @@ extension SuggestionCoordinator {
             return
         }
 
-        // The debounce window adapts to the last generation latency: snappier when the model is
-        // fast, calmer when it is slow (fewer doomed generations to cancel). The configured value
-        // is the fallback until a first latency exists.
+        // The debounce window adapts to how quickly the engine produces text: none when a warm
+        // request shows its first text fast, calmer when the model is slow (fewer doomed
+        // generations to cancel). The configured value is the fallback until a latency exists.
         let debounceMilliseconds = DebouncePolicy.milliseconds(
             lastGenerationLatencyMilliseconds:
                 lastLatencyByEngine[settingsSnapshot.selectedEngine],
+            lastFirstTextLatencyMilliseconds:
+                typicalFirstTextLatency(engine: settingsSnapshot.selectedEngine),
             fallback: settingsSnapshot.debounceMilliseconds,
             engine: settingsSnapshot.selectedEngine
         )
@@ -252,11 +254,46 @@ extension SuggestionCoordinator {
         guard !userDefaults.bool(forKey: Self.speculativePrefetchDisabledDefaultsKey) else { return }
         guard !insertionChunk.isEmpty else { return }
 
-        let optimistic = SpeculativeAcceptanceContext.optimisticSnapshot(
-            after: rawContext,
-            inserting: insertionChunk
+        dispatchSpeculativeGeneration(
+            optimistic: SpeculativeAcceptanceContext.optimisticSnapshot(after: rawContext, inserting: insertionChunk),
+            message: "Started the post-acceptance generation against the expected post-insert text."
         )
+    }
 
+    /// Starts generating for a plain typed character at keystroke time, against the text the host
+    /// is about to publish, instead of after the host-publish poll sees it (10-50 ms in native
+    /// hosts, longer in Chromium). The poll that follows validates exactly as it does for
+    /// post-acceptance speculation: a matching publish keeps this request, anything else (smart
+    /// quotes, dead keys, IME, auto-pairing, a key the host had not yet applied to the snapshot)
+    /// schedules the ordinary request and retires this one. Results and streamed partials are only
+    /// shown once the live text matches, so a wrong bet costs one discarded generation, never a
+    /// ghost at the wrong caret.
+    ///
+    /// Endpoint requests keep their debounced path: speculative traffic must not multiply requests
+    /// to a configured server.
+    ///
+    /// `knownPrecedingText` is the text before the caret as Cotabby already knows it when the
+    /// captured snapshot lags: typing through a visible suggestion advances the session without a
+    /// new AX capture, so the snapshot can still miss those letters. It is used only when it
+    /// extends the snapshot's own text.
+    func dispatchSpeculativeTypingGeneration(
+        rawContext: FocusedInputSnapshot, typed characters: String, knownPrecedingText: String? = nil
+    ) {
+        guard settingsSnapshot.selectedEngine != .openAICompatible,
+              !userDefaults.bool(forKey: Self.speculativePrefetchDisabledDefaultsKey),
+              characters.count == 1, let typed = characters.first, !typed.isNewline,
+              rawContext.selection.length == 0, !rawContext.isSecure else { return }
+        let base = knownPrecedingText.flatMap { $0.hasPrefix(rawContext.precedingText) ? $0 : nil }
+            ?? rawContext.precedingText
+        dispatchSpeculativeGeneration(
+            optimistic: SpeculativeAcceptanceContext.optimisticSnapshot(after: rawContext, precedingText: base + characters),
+            message: "Started generating against the typed character before the host published it."
+        )
+    }
+
+    /// Shared body of the speculative paths: gate, build the request for `optimistic`, and run it
+    /// with no debounce while `pendingSpeculativeContext` lets the publish poll validate it.
+    private func dispatchSpeculativeGeneration(optimistic: FocusedInputSnapshot, message: String) {
         // Same pre-generation gates the ordinary cycle applies, minus their UI side effects: a
         // speculative request must not spend a decode on text the normal path would refuse (too
         // little text) or suppress (typo gate). The post-publish regeneration still runs the full
@@ -306,7 +343,7 @@ extension SuggestionCoordinator {
             "speculative-generating",
             workID: workID,
             generation: context.generation,
-            message: "Started the post-acceptance generation against the expected post-insert text.",
+            message: message,
             prompt: requestBuildResult.promptPreview
         )
     }
@@ -576,7 +613,25 @@ extension SuggestionCoordinator {
             }
             return rebased
         }
-        return liveContext.generation == partial.generation ? partial : nil
+        if liveContext.generation == partial.generation { return partial }
+        // A speculative request was built against text the host had not published yet, so any
+        // read of the pre-publish snapshot moves the live generation past the partial's. Like
+        // `apply`, accept it by content instead: once the live field reads exactly what was
+        // speculated, the partial is current; before that it stays hidden.
+        if let expected = pendingSpeculativeContext,
+           expected.sessionIdentity == liveContext.sessionIdentity,
+           expected.contentSignature == liveContext.contentSignature,
+           expected.generation == partial.generation {
+            return SuggestionResult(
+                generation: liveContext.generation,
+                rawText: partial.rawText,
+                text: partial.text,
+                latency: partial.latency,
+                suppressionReason: partial.suppressionReason,
+                spacingIsExact: partial.spacingIsExact
+            )
+        }
+        return nil
     }
 
     /// Runs the typo gate for the current word. Returns `true` when it handled the cycle by suppressing,

Cotabby/App/Coordinators/Suggestion/SuggestionCoordinator+TypingPrediction.swift

@@ -67,6 +67,10 @@ extension SuggestionCoordinator {
     /// presentation; the hidden buffer still lets a complete answer survive intervening typing.
     func receiveTypingPartial(_ result: SuggestionResult, workID: UInt64, showPartials: Bool) {
         guard workController.isCurrent(workID) else { return }
+        if firstTextRecordedWorkID != workID {
+            firstTextRecordedWorkID = workID
+            recordFirstTextLatency(Int((result.latency * 1_000).rounded()), engine: settingsSnapshot.selectedEngine)
+        }
         if var candidate = typingPrediction {
             guard !candidate.isFinal, result.generation == candidate.context.generation else { return }
             guard candidate.receive(result, final: false) else {

Cotabby/App/Coordinators/Suggestion/SuggestionCoordinator.swift

@@ -77,6 +77,24 @@ final class SuggestionCoordinator: ObservableObject {
     /// a slow HTTP backend retain an appropriate debounce without leaking that latency into the
     /// in-process engines.
     var lastLatencyByEngine: [SuggestionEngineKind: Int] = [:]
+    /// Latency to the first streamed text of the most recent requests per backend: what starting a
+    /// request costs once the KV prefix is warm. `DebouncePolicy` prefers it, because a superseded
+    /// request stops at its next token while a debounce delays every keystroke. Recorded once per
+    /// work item (`firstTextRecordedWorkID`), only for the request that is still current, and read
+    /// as the median of the last few so one cache miss (a context refresh re-decoding a whole
+    /// section) cannot put the next keystroke behind the slow tier.
+    var recentFirstTextLatenciesByEngine: [SuggestionEngineKind: [Int]] = [:]
+    var firstTextRecordedWorkID: UInt64?
+
+    func recordFirstTextLatency(_ milliseconds: Int, engine: SuggestionEngineKind) {
+        recentFirstTextLatenciesByEngine[engine, default: []].append(milliseconds)
+        recentFirstTextLatenciesByEngine[engine] = Array(recentFirstTextLatenciesByEngine[engine, default: []].suffix(3))
+    }
+
+    func typicalFirstTextLatency(engine: SuggestionEngineKind) -> Int? {
+        let recent = recentFirstTextLatenciesByEngine[engine, default: []].sorted()
+        return recent.isEmpty ? nil : recent[recent.count / 2]
+    }
     // Synchronous input/focus callbacks cannot directly `await`, so resets are represented as a
     // barrier task that the next generation must cross before it can ask the runtime for output.
     var cacheResetSequence: UInt64 = 0

Cotabby/Models/Onboarding/OnboardingTemplate.swift

@@ -86,16 +86,19 @@ enum OnboardingTemplate: String, CaseIterable, Identifiable, Equatable, Sendable
         }
     }
 
+    /// Short suggestions are both faster (fewer generated tokens) and more often exactly right, and
+    /// accepting one extends it with a continuation, so every tier but Powerful starts short.
+    /// Powerful keeps the long preset for users who explicitly choose longer suggestions.
     var wordCountPreset: SuggestionWordCountPreset {
         switch self {
         case .quick:
-            return .fourToSeven
+            return .twoToFour
         case .everyday:
-            return .sevenToTwelve
+            return .fourToSeven
         case .powerful:
             return .twelveToTwenty
         case .custom:
-            return .sevenToTwelve
+            return .fourToSeven
         }
     }