Fix rapid Tab leaking into the host after a final-word accept

Cotabby/App/Coordinators/SuggestionCoordinator+Acceptance.swift

@@ -44,6 +44,23 @@ extension SuggestionCoordinator {
         // accept. `validateSessionForAcceptance` still rejects the accept if the session no longer
         // reconciles with the live AX state.
         guard interactionState.activeSession != nil else {
+            // A final-chunk accept tears the session down and regenerates the continuation
+            // asynchronously (see the `.exhausted` branch below). While that regen is in flight we
+            // keep owning Tab instead of leaking it into the host app as a real Tab. Swallow the
+            // press and remember it: the continuation that lands next accepts its first word, so
+            // rapid Tabbing keeps inserting words across the exhaustion boundary instead of jumping
+            // focus out of the field.
+            if isPostExhaustionAcceptanceArmed {
+                hasQueuedPostExhaustionAccept = true
+                logStage(
+                    "\(keyName)-held-for-regen",
+                    workID: currentWorkID,
+                    generation: latestGenerationNumber,
+                    message: "Held a rapid \(keyName) during post-acceptance regeneration; "
+                        + "the next continuation will accept its first word."
+                )
+                return true
+            }
             return passTabThrough(
                 reason: "Key passed through because no valid suggestion was ready."
             )
@@ -134,6 +151,11 @@ extension SuggestionCoordinator {
                 message: "Inserted the final suggestion chunk and queued a refresh.",
                 normalizedOutput: acceptedChunk
             )
+            // Keep owning Tab while the continuation regenerates so a fast follow-up press is
+            // swallowed and queued instead of leaking into the host app as a real Tab. Must run
+            // *after* the `hideOverlay` above, which routes through `onStateChange(.hidden)` and
+            // turns interception off; arming re-asserts it. See `armPostExhaustionAcceptance`.
+            armPostExhaustionAcceptance()
             // Wait for the host to actually publish the inserted text before regenerating. A bare
             // `schedulePrediction()` here reads pre-insertion AX in Chromium editors (the publish lags
             // the synthetic keystroke), so the model re-proposes the word just accepted and the next
@@ -194,6 +216,91 @@ extension SuggestionCoordinator {
         return false
     }
 
+    // MARK: - Post-Exhaustion Acceptance Window
+
+    /// How long Cotabby keeps owning Tab after a final-chunk accept while it waits for the
+    /// continuation to regenerate. This is only a backstop: the window normally ends much sooner —
+    /// when the next suggestion shows (overlay visible) or any teardown hides the overlay. It exists
+    /// so a regeneration that silently stalls can never trap Tab in the focused field. Sized to
+    /// comfortably outlast the host-publish poll ceiling plus a debounce and a typical on-device
+    /// generation.
+    static let postExhaustionAcceptanceWindowSeconds: TimeInterval = 0.8
+
+    /// Keeps the accept tap owning Tab for a brief window after a final-chunk accept, while the
+    /// continuation regenerates asynchronously.
+    ///
+    /// Accepting the last buffered word hides the overlay synchronously (which routes through
+    /// `onStateChange(.hidden)` and turns interception off) and then reschedules generation through
+    /// the host-publish poll + debounce + engine round-trip. That leaves a gap with no active session
+    /// and no visible overlay. A fast follow-up Tab in that gap used to hit the fail-open preflight
+    /// (`shouldConsumeAcceptKeyProvider` keys on overlay visibility) and the `activeSession != nil`
+    /// 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.
+    func armPostExhaustionAcceptance() {
+        isPostExhaustionAcceptanceArmed = true
+        hasQueuedPostExhaustionAccept = false
+        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 }
+            self.releasePostExhaustionAcceptanceWindow()
+        }
+    }
+
+    /// Clears the window flags and invalidates the backstop token, so a timer still pending from
+    /// `armPostExhaustionAcceptance` fires as a no-op once the window has ended. Interception is left
+    /// 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
+    }
+
+    /// 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 }
+        if !overlayState.isVisible {
+            inputMonitor.setAcceptInterceptionActive(false)
+        }
+        clearPostExhaustionAcceptanceWindow()
+    }
+
+    /// Once a regenerated continuation is on screen, accepts its first word if the user pressed Tab
+    /// while it was still loading. Keeps rapid Tabbing inserting words across the exhaustion boundary
+    /// 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()
+        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
+        // the rare miss for diagnosis instead of letting the swallowed Tab vanish without a trace.
+        if !acceptCurrentSuggestion() {
+            logStage(
+                "flush-queued-accept-failed",
+                workID: currentWorkID,
+                generation: latestGenerationNumber,
+                message: "Flushed a queued post-exhaustion Tab, but the follow-up acceptance returned false."
+            )
+        }
+    }
+
     /// Advances the active session from the user's directly typed characters when they match the
     /// next expected tail exactly. This avoids a wasteful regeneration for text the user already
     /// committed to the field themselves.

Cotabby/App/Coordinators/SuggestionCoordinator+Prediction.swift

@@ -266,6 +266,11 @@ extension SuggestionCoordinator {
             rawOutput: result.rawText,
             normalizedOutput: result.text
         )
+
+        // If the user pressed Tab while this continuation was still regenerating, accept its first
+        // word now so rapid Tabbing keeps inserting words across the exhaustion boundary instead of
+        // stalling once the previous suggestion ran out. No-op when nothing was queued.
+        flushQueuedPostExhaustionAcceptIfNeeded()
     }
 
     /// Converts a runtime or engine failure into visible coordinator state and clears stale UI.

Cotabby/App/Coordinators/SuggestionCoordinator.swift

@@ -86,6 +86,22 @@ final class SuggestionCoordinator: ObservableObject {
     /// accept on the last word. See `SuggestionSessionReconciler.isStaleAcceptanceEcho`.
     var lastAcceptedTail: AcceptedSuggestionTail?
 
+    /// 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
+
     init(
         permissionManager: any SuggestionPermissionProviding,
         focusModel: any SuggestionFocusProviding,
@@ -166,7 +182,11 @@ final class SuggestionCoordinator: ObservableObject {
         // before the tap passes the original key through.
         inputMonitor.shouldConsumeAcceptKeyProvider = { [weak self] in
             guard let self else { return false }
-            guard self.overlayState.isVisible else { return false }
+            // Keep owning the accept key through the brief post-acceptance regeneration window too,
+            // 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 }
             return true
         }
 
@@ -181,6 +201,12 @@ final class SuggestionCoordinator: ObservableObject {
                 self.inputMonitor.setAcceptInterceptionActive(true)
             case .hidden:
                 self.inputMonitor.setAcceptInterceptionActive(false)
+                // A hidden overlay ends any post-exhaustion Tab-ownership window. Every teardown and
+                // abort path hides the overlay, so ending the window here is the single catch-all
+                // that returns the accept key to the host (and cancels the backstop timer) once the
+                // window is genuinely over. The `.exhausted` accept re-arms *after* its own
+                // `hideOverlay` call, so this never cancels a window that was just opened.
+                self.clearPostExhaustionAcceptanceWindow()
             }
         }
 

CotabbyTests/SuggestionCoordinatorAcceptanceTests.swift

@@ -132,6 +132,144 @@ final class SuggestionCoordinatorAcceptanceTests: XCTestCase {
         }
     }
 
+    func test_rapidSecondAcceptDuringRegenerationIsConsumedNotPassedThrough() {
+        runOnMainActor {
+            let snapshot = CotabbyTestFixtures.focusedInputSnapshot(precedingText: "what's on your mind")
+            let context = FocusedInputContext(snapshot: snapshot, generation: 7)
+            let interactionState = SuggestionInteractionState()
+            _ = interactionState.startSession(
+                fullText: " today",
+                liveContext: context,
+                latency: 0.1
+            )
+            let overlayState = OverlayState.visible(
+                text: " today",
+                geometry: CotabbyTestFixtures.overlayGeometry(caretRect: context.caretRect),
+                mode: .inline
+            )
+            let inputMonitor = StubSuggestionInputMonitor()
+            let inserter = StubSuggestionInserter()
+            let coordinator = makeCoordinator(
+                snapshot: snapshot,
+                overlayState: overlayState,
+                inputMonitor: inputMonitor,
+                inserter: inserter,
+                interactionState: interactionState
+            )
+
+            // First Tab accepts the only remaining chunk, exhausts the session, and arms the window.
+            XCTAssertTrue(coordinator.acceptCurrentSuggestion())
+            XCTAssertEqual(inserter.insertedChunks, [" today"])
+            XCTAssertTrue(coordinator.isPostExhaustionAcceptanceArmed)
+            XCTAssertFalse(coordinator.overlayState.isVisible)
+            // Ownership of Tab was re-asserted even though the overlay is now hidden.
+            XCTAssertEqual(inputMonitor.acceptInterceptionRequests.last, true)
+            XCTAssertTrue(
+                inputMonitor.shouldConsumeAcceptKeyProvider(),
+                "The accept tap must keep owning Tab while the continuation regenerates."
+            )
+
+            // The rapid second Tab lands before the continuation regenerates. It must be swallowed
+            // (consumed) and queued — never forwarded to the host as a real Tab that moves focus.
+            XCTAssertTrue(
+                coordinator.acceptCurrentSuggestion(),
+                "A fast follow-up Tab during regeneration must be consumed, not passed through to the host."
+            )
+            XCTAssertEqual(
+                inserter.insertedChunks,
+                [" today"],
+                "The second Tab has nothing to insert yet; it is queued, not inserted."
+            )
+            XCTAssertTrue(coordinator.hasQueuedPostExhaustionAccept)
+        }
+    }
+
+    func test_postExhaustionWindowReleasesAcceptKeyWhenOverlayHides() {
+        runOnMainActor {
+            let snapshot = CotabbyTestFixtures.focusedInputSnapshot(precedingText: "what's on your mind")
+            let context = FocusedInputContext(snapshot: snapshot, generation: 7)
+            let interactionState = SuggestionInteractionState()
+            _ = interactionState.startSession(
+                fullText: " today",
+                liveContext: context,
+                latency: 0.1
+            )
+            let overlayState = OverlayState.visible(
+                text: " today",
+                geometry: CotabbyTestFixtures.overlayGeometry(caretRect: context.caretRect),
+                mode: .inline
+            )
+            let inputMonitor = StubSuggestionInputMonitor()
+            let inserter = StubSuggestionInserter()
+            let coordinator = makeCoordinator(
+                snapshot: snapshot,
+                overlayState: overlayState,
+                inputMonitor: inputMonitor,
+                inserter: inserter,
+                interactionState: interactionState
+            )
+
+            XCTAssertTrue(coordinator.acceptCurrentSuggestion())
+            XCTAssertTrue(coordinator.isPostExhaustionAcceptanceArmed)
+
+            // Any teardown that hides the overlay (focus change, typing, dismissal, an empty
+            // regeneration) must end the window so the user can Tab out of the field normally again.
+            coordinator.invalidateActiveSuggestion(reason: "Focus moved to another field.")
+
+            XCTAssertFalse(coordinator.isPostExhaustionAcceptanceArmed)
+            XCTAssertFalse(coordinator.hasQueuedPostExhaustionAccept)
+            XCTAssertFalse(
+                inputMonitor.shouldConsumeAcceptKeyProvider(),
+                "Once the window is released the accept tap should stop owning Tab."
+            )
+            XCTAssertFalse(
+                coordinator.acceptCurrentSuggestion(),
+                "With the window released and no suggestion, Tab must pass through to the host."
+            )
+        }
+    }
+
+    func test_queuedPostExhaustionAcceptInsertsNextWordWhenContinuationArrives() {
+        runOnMainActor {
+            let snapshot = CotabbyTestFixtures.focusedInputSnapshot(precedingText: "Hello")
+            let context = FocusedInputContext(snapshot: snapshot, generation: 7)
+            let interactionState = SuggestionInteractionState()
+            let session = interactionState.startSession(
+                fullText: " world again",
+                liveContext: context,
+                latency: 0.1
+            )
+            let overlayState = OverlayState.visible(
+                text: session.remainingText,
+                geometry: CotabbyTestFixtures.overlayGeometry(caretRect: context.caretRect),
+                mode: .inline
+            )
+            let inputMonitor = StubSuggestionInputMonitor()
+            let inserter = StubSuggestionInserter()
+            let coordinator = makeCoordinator(
+                snapshot: snapshot,
+                overlayState: overlayState,
+                inputMonitor: inputMonitor,
+                inserter: inserter,
+                interactionState: interactionState
+            )
+            // Simulate a Tab that was swallowed and queued while this continuation was still loading;
+            // `apply` calls `flushQueuedPostExhaustionAcceptIfNeeded` once the suggestion is on screen.
+            coordinator.isPostExhaustionAcceptanceArmed = true
+            coordinator.hasQueuedPostExhaustionAccept = true
+
+            coordinator.flushQueuedPostExhaustionAcceptIfNeeded()
+
+            XCTAssertEqual(
+                inserter.insertedChunks,
+                [" world"],
+                "The queued Tab should accept the continuation's first word."
+            )
+            XCTAssertFalse(coordinator.isPostExhaustionAcceptanceArmed)
+            XCTAssertFalse(coordinator.hasQueuedPostExhaustionAccept)
+        }
+    }
+
     @MainActor
     private func makeCoordinator(
         snapshot: FocusedInputSnapshot,