Extract coordinator state mechanisms

Cotabby/App/Coordinators/SuggestionCoordinator+Acceptance.swift

@@ -335,8 +335,7 @@ extension SuggestionCoordinator {
         // A final-chunk accept tears the session down and regenerates the continuation
         // asynchronously (see the `.exhausted` branch). While that regen is in flight we keep owning
         // Tab instead of leaking it into the host app as a real Tab.
-        if isPostExhaustionAcceptanceArmed {
-            hasQueuedPostExhaustionAccept = true
+        if postExhaustionAcceptanceState.queueAcceptIfArmed() {
             logStage(
                 "\(keyName)-held-for-regen",
                 workID: currentWorkID,
@@ -467,20 +466,18 @@ extension SuggestionCoordinator {
     /// guard, so the accept tap forwarded the original Tab to the host and focus jumped out of the
     /// field. Re-asserting interception here keeps the tail tap installed and owning Tab across the
     /// regen window (its mach port otherwise lingers only ~50ms), and `shouldConsumeAcceptKeyProvider`
-    /// also consults `isPostExhaustionAcceptanceArmed` so the key is still routed in while the overlay
-    /// is hidden. A token-keyed backstop guarantees the window can never trap Tab.
+    /// also consults the extracted state so the key is still routed in while the overlay is hidden.
+    /// A token-keyed backstop guarantees the window can never trap Tab.
     func armPostExhaustionAcceptance() {
-        isPostExhaustionAcceptanceArmed = true
-        hasQueuedPostExhaustionAccept = false
+        let generation = postExhaustionAcceptanceState.arm()
         inputMonitor.setAcceptInterceptionActive(true)
-        postExhaustionAcceptanceGeneration &+= 1
-        let generation = postExhaustionAcceptanceGeneration
         DispatchQueue.main.asyncAfter(
             deadline: .now() + Self.postExhaustionAcceptanceWindowSeconds
         ) { [weak self] in
             // Only the generation that scheduled this timer may act on it; a newer accept (or an
             // already-released window) bumped the token, so this fires as a no-op.
-            guard let self, self.postExhaustionAcceptanceGeneration == generation else { return }
+            guard let self,
+                  self.postExhaustionAcceptanceState.ownsBackstop(generation: generation) else { return }
             self.releasePostExhaustionAcceptanceWindow()
         }
     }
@@ -490,18 +487,15 @@ extension SuggestionCoordinator {
     /// to the caller: whether Tab ownership should drop depends on why the window ended (a fresh
     /// suggestion keeps owning it; a teardown or the backstop drops it).
     func clearPostExhaustionAcceptanceWindow() {
-        isPostExhaustionAcceptanceArmed = false
-        hasQueuedPostExhaustionAccept = false
-        // Cancel any pending backstop, which is keyed to the generation captured at arm time.
-        postExhaustionAcceptanceGeneration &+= 1
+        postExhaustionAcceptanceState.clear()
     }
 
     /// Ends the post-exhaustion window and returns the accept key to the host unless a suggestion is
     /// now visible (in which case the normal overlay path keeps owning it). Idempotent. This is the
     /// backstop release; the common, prompt release is `onStateChange(.hidden)` ending the window as
     /// soon as any teardown hides the overlay.
     func releasePostExhaustionAcceptanceWindow() {
-        guard isPostExhaustionAcceptanceArmed || hasQueuedPostExhaustionAccept else { return }
+        guard postExhaustionAcceptanceState.needsRelease else { return }
         if !overlayState.isVisible {
             inputMonitor.setAcceptInterceptionActive(false)
         }
@@ -513,10 +507,9 @@ extension SuggestionCoordinator {
     /// instead of stalling. Bounded to one queued accept so mashing Tab cannot run away. Called at the
     /// end of `apply`'s success path, after the new session and overlay exist.
     func flushQueuedPostExhaustionAcceptIfNeeded() {
-        let shouldAccept = isPostExhaustionAcceptanceArmed && hasQueuedPostExhaustionAccept
-        // Normal acceptance has resumed now that a fresh suggestion is visible, so end the window
-        // regardless of whether a press was queued (this also cancels the now-redundant backstop).
-        clearPostExhaustionAcceptanceWindow()
+        // Consuming also ends the window when no press was queued, invalidating the backstop now that
+        // normal acceptance has resumed with a fresh visible suggestion.
+        let shouldAccept = postExhaustionAcceptanceState.consumeQueuedAccept()
         guard shouldAccept else { return }
         // A queued accept can still legitimately fail (the new continuation no longer reconciles with
         // live AX, or insertion fails). `acceptSuggestion` cleans up its own state on failure, so log

Cotabby/App/Coordinators/SuggestionCoordinator+Prediction.swift

@@ -265,14 +265,9 @@ extension SuggestionCoordinator {
     /// result (or failure) only while it is still the current work. Extracted from
     /// `generateFromCurrentFocus` so that function stays within the project's complexity budget.
     private func dispatchGeneration(request: SuggestionRequest, workID: UInt64) {
-        // A new generation starts a new stream; the previous request's rendered-partial state
-        // must not gate the new partials' monotonic checks. `isStreamDrainScheduled` is left
-        // alone on purpose: an already-enqueued drain block cannot be unscheduled, and it
-        // self-heals either way — it finds nil and clears the flag, or it finds a partial the
-        // new generation queued in the meantime and renders it under the same work-id guards.
-        // Resetting the flag here would instead double-schedule a drain for one partial.
-        streamRenderedText = nil
-        pendingStreamPartial = nil
+        // A new generation starts a new stream. The state value deliberately preserves an already
+        // scheduled drain callback while dropping the old request's partial and rendered text.
+        suggestionStreamingState.beginGeneration()
         // Streaming the ghost text token-by-token is opt-in. Read the flag here on the main actor so
         // the work closure captures a plain Bool. When off, the closure passes no `onPartial`, so the
         // engine skips its per-token main-actor hops entirely and the suggestion appears once, fully
@@ -361,22 +356,18 @@ extension SuggestionCoordinator {
         guard workController.isCurrent(workID) else {
             return
         }
-        pendingStreamPartial = PendingStreamPartial(result: partial, workID: workID)
-        guard !isStreamDrainScheduled else {
+        guard suggestionStreamingState.enqueue(partial, workID: workID) else {
             return
         }
-        isStreamDrainScheduled = true
         DispatchQueue.main.async { [weak self] in
             self?.drainStreamedPartial()
         }
     }
 
     private func drainStreamedPartial() {
-        isStreamDrainScheduled = false
-        guard let pending = pendingStreamPartial else {
+        guard let pending = suggestionStreamingState.drain() else {
             return
         }
-        pendingStreamPartial = nil
         applyStreamedPartial(pending.result, workID: pending.workID)
     }
 
@@ -393,10 +384,7 @@ extension SuggestionCoordinator {
         guard workController.isCurrent(workID) else {
             return
         }
-        guard StreamedGhostTextPolicy.isRenderableExtension(
-            candidate: partial.text,
-            currentlyRendered: streamRenderedText
-        ) else {
+        guard suggestionStreamingState.canRender(partial.text) else {
             return
         }
         guard let rawContext = focusModel.snapshot.context else {
@@ -423,7 +411,7 @@ extension SuggestionCoordinator {
             liveContext: liveContext,
             latency: partial.latency
         )
-        streamRenderedText = partial.text
+        suggestionStreamingState.recordRendered(partial.text)
         presentOverlay(
             text: partial.text,
             at: liveContext.caretRect,
@@ -1085,8 +1073,7 @@ extension SuggestionCoordinator {
         // typed, focus changed, predictions disabled). The final-chunk accept re-sets it afterward.
         lastAcceptedTail = nil
         // Stream bookkeeping follows the session it was rendering for.
-        streamRenderedText = nil
-        pendingStreamPartial = nil
+        suggestionStreamingState.clearSession()
         latestSuggestionPreview = nil
         latestFullSuggestionPreview = nil
         latestRemainingSuggestionPreview = nil

Cotabby/App/Coordinators/SuggestionCoordinator.swift

@@ -101,19 +101,9 @@ final class SuggestionCoordinator: ObservableObject {
     }
 
     var clipboardPrefaceMemo: ClipboardPrefaceMemo?
-    /// Streamed-render bookkeeping. Partial results hop in from the engine while a decode is
-    /// still running; they are coalesced (latest wins, drained once per runloop turn) so
-    /// token-rate deliveries cannot stack session and overlay layout work on the main actor, and
-    /// `streamRenderedText` carries the monotonic-extension state for `StreamedGhostTextPolicy`.
-    /// All of it is scoped to the current work id and reset when a new generation dispatches.
-    struct PendingStreamPartial {
-        let result: SuggestionResult
-        let workID: UInt64
-    }
-
-    var pendingStreamPartial: PendingStreamPartial?
-    var isStreamDrainScheduled = false
-    var streamRenderedText: String?
+    /// Coalescing and monotonic-render state for engine partials. The coordinator owns scheduling
+    /// and presentation; the value owns the stream's pure state transitions.
+    var suggestionStreamingState = SuggestionStreamingState()
 
     /// Monotonic cancellation token for the "wait until the host publishes typed text to AX" loop.
     ///
@@ -155,21 +145,9 @@ final class SuggestionCoordinator: ObservableObject {
     /// stands down instead of scheduling a duplicate regeneration.
     var pendingSpeculativeSignature: String?
 
-    /// Monotonic token for the post-exhaustion "keep owning Tab" window. Bumped on every arm so a
-    /// stale backstop timer (or a window superseded by a newer accept) no-ops instead of releasing a
-    /// window it no longer owns. See `armPostExhaustionAcceptance`.
-    var postExhaustionAcceptanceGeneration: UInt64 = 0
-    /// True while Cotabby keeps the accept tap owning Tab in the gap between a final-chunk accept and
-    /// the regenerated continuation appearing. Accepting the last buffered word hides the overlay and
-    /// reschedules generation asynchronously; without this the fail-open accept-tap preflight (which
-    /// keys on overlay visibility) would forward a fast follow-up Tab to the host as a real Tab and
-    /// focus would jump out of the field. The window self-releases when the next suggestion shows,
-    /// when any teardown hides the overlay, or via a backstop timer. See `armPostExhaustionAcceptance`.
-    var isPostExhaustionAcceptanceArmed = false
-    /// Set when a Tab is swallowed during that window. The next continuation that lands accepts its
-    /// first word, so rapid Tabbing keeps inserting words across the exhaustion boundary instead of
-    /// stalling. Bounded to a single queued accept so mashing Tab cannot run away.
-    var hasQueuedPostExhaustionAccept = false
+    /// Pure state for the bounded "keep owning Tab" window after a final-chunk acceptance. The
+    /// coordinator continues to own the timer and input-monitor effects around these transitions.
+    var postExhaustionAcceptanceState = PostExhaustionAcceptanceState()
 
     init(
         permissionManager: any SuggestionPermissionProviding,
@@ -263,7 +241,7 @@ final class SuggestionCoordinator: ObservableObject {
             // even though the overlay is hidden then. Otherwise a fast follow-up Tab in that gap
             // falls through to the host app as a real Tab and focus jumps out of the field — the
             // "rapid Tab breaks, slow Tab is fine" report. See `armPostExhaustionAcceptance`.
-            guard self.overlayState.isVisible || self.isPostExhaustionAcceptanceArmed else { return false }
+            guard self.overlayState.isVisible || self.postExhaustionAcceptanceState.isArmed else { return false }
             return true
         }
 

Cotabby/Support/Suggestion/PostExhaustionAcceptanceState.swift

@@ -0,0 +1,59 @@
+/// Pure state machine for the short Tab-ownership window after a suggestion is exhausted.
+///
+/// The coordinator owns one instance for its lifetime and supplies the side effects: installing the
+/// input interception, scheduling the timeout, and accepting a regenerated continuation. Keeping
+/// those effects outside this value makes the hard rules explicit: only one accept may be queued,
+/// every arm invalidates older timeout callbacks, and consuming or clearing the window is atomic.
+struct PostExhaustionAcceptanceState: Equatable {
+    private(set) var backstopGeneration: UInt64 = 0
+    private(set) var isArmed = false
+    private(set) var hasQueuedAccept = false
+
+    /// Whether a release has any state to clear. Kept separate from `isArmed` so an inconsistent
+    /// queued flag still fails safe by returning Tab ownership to the host.
+    var needsRelease: Bool {
+        isArmed || hasQueuedAccept
+    }
+
+    /// Opens a fresh window and returns the identity its timeout callback must capture.
+    @discardableResult
+    mutating func arm() -> UInt64 {
+        isArmed = true
+        hasQueuedAccept = false
+        backstopGeneration &+= 1
+        return backstopGeneration
+    }
+
+    /// Queues one rapid accept while the continuation is unavailable.
+    ///
+    /// Repeated presses deliberately collapse into one Boolean so a user mashing Tab cannot cause a
+    /// regenerated suggestion to accept multiple unseen chunks.
+    @discardableResult
+    mutating func queueAcceptIfArmed() -> Bool {
+        guard isArmed else {
+            return false
+        }
+
+        hasQueuedAccept = true
+        return true
+    }
+
+    /// Returns true only for the timeout belonging to the current arm operation.
+    func ownsBackstop(generation: UInt64) -> Bool {
+        backstopGeneration == generation
+    }
+
+    /// Closes the window and invalidates every timeout callback already in flight.
+    mutating func clear() {
+        isArmed = false
+        hasQueuedAccept = false
+        backstopGeneration &+= 1
+    }
+
+    /// Atomically closes the window and reports whether the regenerated continuation owes one accept.
+    mutating func consumeQueuedAccept() -> Bool {
+        let shouldAccept = isArmed && hasQueuedAccept
+        clear()
+        return shouldAccept
+    }
+}

Cotabby/Support/Suggestion/SuggestionStreamingState.swift

@@ -0,0 +1,69 @@
+/// Main-actor bookkeeping for one streamed suggestion delivery.
+///
+/// `SuggestionCoordinator` owns this value for its lifetime. The value keeps token-rate engine
+/// callbacks from leaking three related invariants into the coordinator: pending partials are
+/// latest-wins, at most one drain is scheduled per runloop turn, and visible text may only grow
+/// monotonically. It does not schedule work or render UI; the coordinator remains responsible for
+/// those side effects.
+struct SuggestionStreamingState {
+    /// One partial paired with the replaceable-work identity that produced it.
+    struct PendingPartial {
+        let result: SuggestionResult
+        let workID: UInt64
+    }
+
+    private(set) var pendingPartial: PendingPartial?
+    private(set) var isDrainScheduled = false
+    private(set) var renderedText: String?
+
+    /// Starts a new stream without clearing an already-enqueued drain callback.
+    ///
+    /// A scheduled callback cannot be cancelled. Preserving `isDrainScheduled` lets the existing
+    /// callback drain a replacement partial from the new generation; clearing it here could enqueue
+    /// two drains for the same runloop turn.
+    mutating func beginGeneration() {
+        renderedText = nil
+        pendingPartial = nil
+    }
+
+    /// Stores the newest partial and returns whether the coordinator must schedule a drain.
+    @discardableResult
+    mutating func enqueue(_ result: SuggestionResult, workID: UInt64) -> Bool {
+        pendingPartial = PendingPartial(result: result, workID: workID)
+        guard !isDrainScheduled else {
+            return false
+        }
+
+        isDrainScheduled = true
+        return true
+    }
+
+    /// Ends the scheduled-drain window and returns the newest partial, if one survived teardown.
+    mutating func drain() -> PendingPartial? {
+        isDrainScheduled = false
+        defer { pendingPartial = nil }
+        return pendingPartial
+    }
+
+    /// Applies the backend-independent monotonic rendering rule to the current stream.
+    func canRender(_ candidate: String) -> Bool {
+        StreamedGhostTextPolicy.isRenderableExtension(
+            candidate: candidate,
+            currentlyRendered: renderedText
+        )
+    }
+
+    /// Records text only after all coordinator freshness and seam guards have accepted it.
+    mutating func recordRendered(_ text: String) {
+        renderedText = text
+    }
+
+    /// Drops state associated with a torn-down suggestion session.
+    ///
+    /// As with `beginGeneration`, a scheduled callback remains responsible for clearing the drain
+    /// flag when it eventually runs.
+    mutating func clearSession() {
+        renderedText = nil
+        pendingPartial = nil
+    }
+}