Keep native-context predictions stable without fallback tracking (#806)

Separate fallback tracking availability from native context caching and stop unchanged polls from resetting prediction batches. Cover scan timing, pass limits, selected choices and stale-context rejection.

Co-authored-by: Owen McGirr <owenmcgirr@Owens-Mac-Studio-2.local>

Author: enaboapps

src-tauri/src/prediction/mod.rs

@@ -152,6 +152,18 @@ struct Service {
     tracking: bool,
 }
 impl Service {
+    fn suggestions(
+        &mut self,
+        keyboard: &mut Keyboard,
+        batch: Option<worker::Batch>,
+        tracking: bool,
+    ) {
+        self.tracking = tracking;
+        if !tracking {
+            self.edit = None;
+        }
+        keyboard.predictions(batch, false);
+    }
     fn fail(&mut self, keyboard: &mut Keyboard) {
         self.client = None;
         self.failed = true;
@@ -270,12 +282,7 @@ pub fn poll(app: &AppHandle, keyboard: Option<&mut Keyboard>, enabled: bool, ign
                         batch,
                         tracking,
                     } if generation == s.generation => {
-                        s.tracking = tracking;
-                        if !tracking {
-                            s.edit = None;
-                            s.reset = true;
-                        }
-                        keyboard.predictions(batch, false);
+                        s.suggestions(keyboard, batch, tracking);
                     }
                     Response::Insert { generation, text } => {
                         s.accepting = false;

src-tauri/src/prediction/worker.rs

@@ -139,9 +139,16 @@ impl<A: Adapter> Engine<A> {
             .as_ref()
             .is_some_and(|old| self.adapter.same(old, &target).unwrap_or(false));
         let uninterrupted = epoch == self.activity;
-        if !same || !uninterrupted || reset || !healthy {
+        if !same || !uninterrupted || reset {
             self.clear();
         }
+        if !healthy {
+            self.fallback.clear();
+            self.boundary = false;
+            if self.snapshot.as_ref().is_some_and(|raw| raw.position == -1) {
+                self.clear();
+            }
+        }
         self.activity = epoch;
         if same && uninterrupted && !reset && self.tracked && healthy {
             if let Some(text) = edit {
@@ -369,6 +376,7 @@ mod tests {
         target: usize,
         protected: bool,
         unsupported: bool,
+        editable: bool,
         reads: usize,
         change_on_read: bool,
     }
@@ -385,6 +393,7 @@ mod tests {
                 target: 1,
                 protected: false,
                 unsupported: false,
+                editable: true,
                 reads: 0,
                 change_on_read: false,
             }
@@ -399,7 +408,7 @@ mod tests {
             Ok(self.protected)
         }
         fn editable(&mut self, _: &usize) -> bool {
-            true
+            self.editable
         }
         fn same(&mut self, a: &usize, b: &usize) -> Result<bool, Status> {
             Ok(a == b)
@@ -423,6 +432,100 @@ mod tests {
         e
     }
     #[test]
+    fn native_predictions_survive_polls_without_fallback_tracking() {
+        for unavailable in 0..3 {
+            let mut e = engine();
+            match unavailable {
+                0 => e.tracked = false,
+                1 => e.adapter.editable = false,
+                _ => e.observe = || (0, false),
+            }
+            let mut service = crate::prediction::Service::default();
+            let mut keyboard = crate::scan_keyboard::Keyboard::new(false);
+            keyboard.enable_predictions(true);
+            let mut original = None;
+            for _ in 0..40 {
+                let response = e.respond(Request::Query {
+                    generation: 0,
+                    edit: service.edit.take(),
+                    reset: std::mem::take(&mut service.reset),
+                    shift: false,
+                    caps: false,
+                });
+                let Response::Suggestions {
+                    batch, tracking, ..
+                } = response
+                else {
+                    panic!("expected suggestions");
+                };
+                assert!(!tracking);
+                let batch = batch.unwrap();
+                let expected = original.get_or_insert_with(|| batch.clone());
+                assert_eq!(batch.token, expected.token);
+                assert_eq!(batch.words, expected.words);
+                service.suggestions(&mut keyboard, Some(batch), tracking);
+                keyboard.advance(250, 1000);
+            }
+            let token = original.unwrap().token;
+            assert_eq!(e.accept(token, 0), Some("ter ".into()));
+            assert!(e.accept(token, 0).is_none());
+        }
+    }
+    #[test]
+    fn genuine_context_changes_invalidate_identical_native_candidates() {
+        for change in 0..4 {
+            let mut e = engine();
+            e.adapter.editable = false;
+            let original = e.query(None, false, false, false).unwrap();
+            match change {
+                0 => e.adapter.target += 1,
+                1 => e.adapter.raw.position += 1,
+                2 => e.activity = u64::MAX,
+                _ => {}
+            }
+            let replacement = e.query(None, change == 3, false, false).unwrap();
+            assert_eq!(replacement.words, original.words);
+            assert_ne!(replacement.token, original.token);
+            assert!(e.accept(original.token, 0).is_none());
+            assert_eq!(e.accept(replacement.token, 0), Some("ter ".into()));
+        }
+    }
+    #[test]
+    fn native_edits_selection_and_casing_refresh_without_tracking() {
+        let mut e = engine();
+        e.tracked = false;
+        let original = e.query(None, false, false, false).unwrap();
+        e.adapter.raw.before.push('t');
+        let typed = e.query(None, false, false, false).unwrap();
+        assert_ne!(typed.token, original.token);
+        assert!(e.accept(original.token, 0).is_none());
+        e.adapter.raw.before.pop();
+        let deleted = e.query(None, false, false, false).unwrap();
+        assert_ne!(deleted.token, typed.token);
+        let shifted = e.query(None, false, true, false).unwrap();
+        assert_ne!(shifted.token, deleted.token);
+        assert!(e.accept(deleted.token, 0).is_none());
+        e.adapter.raw.has_selection = true;
+        assert!(e.query(None, false, false, false).is_none());
+        assert!(e.accept(shifted.token, 0).is_none());
+        e.adapter.raw.has_selection = false;
+        let restored = e.query(None, false, false, false).unwrap();
+        e.adapter.protected = true;
+        assert!(e.query(None, false, false, false).is_none());
+        assert!(e.accept(restored.token, 0).is_none());
+    }
+    #[test]
+    fn losing_tracking_does_not_erase_a_pending_context_reset() {
+        let mut service = crate::prediction::Service {
+            reset: true,
+            edit: Some("queued".into()),
+            ..Default::default()
+        };
+        service.suggestions(&mut crate::scan_keyboard::Keyboard::new(false), None, false);
+        assert!(service.reset);
+        assert!(service.edit.is_none());
+    }
+    #[test]
     fn only_suffix_and_space_are_returned_once() {
         let mut e = engine();
         let b = e.query(None, false, false, false).unwrap();

src-tauri/src/scan_keyboard.rs

@@ -619,6 +619,77 @@ mod tests {
         assert_eq!(keyboard.scan.nav.index(), 1);
     }
 
+    #[test]
+    fn unchanged_predictions_preserve_navigation_timing_and_pass_limits() {
+        use crate::scan_preferences::{Direction, Pattern, Resolved};
+        for direction in [Direction::Forward, Direction::Reverse] {
+            for pattern in [Pattern::Grouped, Pattern::Linear] {
+                for automatic in [false, true] {
+                    let options = Resolved {
+                        direction,
+                        pattern,
+                        pass_limit: 2,
+                        ..Default::default()
+                    };
+                    let mut baseline = Keyboard::configured(false, options);
+                    let mut polled = Keyboard::configured(false, options);
+                    let batch = crate::prediction::worker::Batch {
+                        token: 7,
+                        words: vec!["water".into(), "walk".into()],
+                    };
+                    for k in [&mut baseline, &mut polled] {
+                        k.enable_predictions(true);
+                        k.predictions(Some(batch.clone()), false);
+                        k.restart();
+                    }
+                    let ticks = 12 * (polled.rows.iter().map(Vec::len).sum::<usize>() + 1);
+                    for tick in 0..ticks {
+                        polled.predictions(Some(batch.clone()), false);
+                        if automatic {
+                            baseline.advance(250, 1000);
+                            polled.advance(250, 1000);
+                        } else if tick % 4 == 3 {
+                            baseline.handle(Action::Next);
+                            polled.handle(Action::Next);
+                        }
+                        assert_eq!(
+                            polled.scan.position(&polled.rows),
+                            baseline.scan.position(&baseline.rows)
+                        );
+                        assert_eq!(polled.suspended(), baseline.suspended());
+                        assert_eq!(polled.disabled(), baseline.disabled());
+                        assert_eq!(polled.predictions.as_ref().unwrap().token, 7);
+                        assert!(polled.queued_predictions.is_none());
+                    }
+                    if automatic {
+                        assert!(polled.suspended());
+                    }
+                }
+            }
+        }
+    }
+    #[test]
+    fn unchanged_predictions_preserve_selected_key_and_execute_once() {
+        let mut k = Keyboard::new(false);
+        k.enable_predictions(true);
+        let batch = crate::prediction::worker::Batch {
+            token: 7,
+            words: vec!["water".into(), "walk".into()],
+        };
+        k.predictions(Some(batch.clone()), false);
+        k.handle(Action::Select);
+        k.advance(750, 1000);
+        k.predictions(Some(batch.clone()), false);
+        assert_eq!(k.scan.position(&k.rows), (0, Some(0)));
+        k.advance(250, 1000);
+        assert_eq!(k.scan.position(&k.rows), (0, Some(1)));
+        k.predictions(Some(batch), false);
+        assert_eq!(
+            k.handle(Action::Select),
+            Some(Output::Prediction { token: 7, index: 1 })
+        );
+        assert!(k.handle(Action::Select).is_none());
+    }
     #[test]
     fn predictions_skip_empty_slots_and_defer_acceptance() {
         let mut k = Keyboard::new(false);