Recover word prediction from one slow call without needing Retry

docs/word-prediction.md

@@ -12,7 +12,7 @@ A passive observer records only an activity counter and timestamp, never externa
 
 One Word prediction setting controls an offline SmolLM2-135M int8 ONNX model. The saved `enhancedWordPrediction` field is retained for compatibility but does not choose an engine. The model spells candidates from subword pieces and reads only the last 256 characters of the temporary buffer, beginning at a word boundary. It never learns from typing.
 
-For a first typed prefix, a fixed local context keeps the model from favoring website names at the start of a document; it adds no user text. The model loads in the worker while keyboard input remains available, which takes about a second. Suggestions are blank until loading finishes. Later keyboard opens normally skip this wait by using the spare worker described under Worker process. A passive badge beside the scan prompt distinguishes loading, a ready keyboard awaiting typed context, no matching suggestions, available suggestions, paused activity tracking, and prediction failure. It never shows typed text and is not a scan target. If loading or inference fails, or a call exceeds 1.5 seconds, the keyboard continues accepting input and offers **Retry predictions** in its toolbar. Retry restarts only the prediction worker, clears its private text context, and leaves the keyboard open; type a new prefix afterward. A 400 ms search budget bounds candidate exploration. A clipped buffer with no complete earlier word yields no suggestions.
+For a first typed prefix, a fixed local context keeps the model from favoring website names at the start of a document; it adds no user text. The model loads in the worker while keyboard input remains available, which takes about a second. Suggestions are blank until loading finishes. Later keyboard opens normally skip this wait by using the spare worker described under Worker process. A passive badge beside the scan prompt distinguishes loading, a ready keyboard awaiting typed context, no matching suggestions, available suggestions, paused activity tracking, and prediction failure. It never shows typed text and is not a scan target. If loading or inference fails, the keyboard continues accepting input and offers **Retry predictions** in its toolbar. A call that exceeds 1.5 seconds is tolerated: its suggestions are still shown, and only three such calls in a row count as a failure, since one stall on a busy machine does not mean the model cannot keep up. A call at normal speed clears the count. A call so slow that its reply misses the two-second deadline is handled by replacing the worker, described under Worker process. Retry restarts only the prediction worker, clears its private text context, and leaves the keyboard open; type a new prefix afterward. A 400 ms search budget bounds candidate exploration. A clipped buffer with no complete earlier word yields no suggestions.
 
 Candidates contain ASCII letters and apostrophes and must have sufficient model probability. There is no vocabulary filter: any word the model finds likely can be suggested, including swearing, because the person typing chose it. Single-letter candidates are limited to “a” and “I”. Up to five suggestions are shown: the first, third and fifth are the most likely single words, and the second and fourth are the two most likely two-word phrases that begin with one of the top three words, ranked by the probability of the pair. When fewer phrases are found within an extra 200 ms, single words fill the remaining slots, and the other way round. Accepting a phrase inserts the rest of its first word, a space, the second word and a trailing space.
 
@@ -62,7 +62,7 @@ The model files are too large to commit. `npm run prediction-model` downloads th
 
 ## Worker process
 
-A separate process owns the buffer, activity observer and model. Private bounded inherited pipes carry successful edits and results to native rendering. Text and suggestions are not sent to the React UI, diagnostic history or telemetry. One request is outstanding at a time with a two-second deadline. Timeout stops predictions until the keyboard is reopened; ordinary keyboard operation remains available. Closing the keyboard, ending scanning, or exiting kills and reaps the worker.
+A separate process owns the buffer, activity observer and model. Private bounded inherited pipes carry successful edits and results to native rendering. Text and suggestions are not sent to the React UI, diagnostic history or telemetry. One request is outstanding at a time with a two-second deadline. When a reply misses it, the worker is killed and reaped with its text context and a new worker starts by itself, so suggestions return after the model loads. As after any reset of the context, the new worker knows only what is typed from then on, so a word that was half typed is completed from its remaining letters alone; suggestions are reliable again from the next word. A second missed deadline within 60 seconds of that replacement stops prediction and offers Retry predictions, whatever the replacement answered in between. So does a third missed deadline while the same keyboard is open, however far apart they are. Reopening the keyboard or selecting Retry predictions starts the count again. An acceptance that was waiting on the missed reply fails and types nothing; the keyboard shows its error state for that acceptance while prediction itself recovers. Ordinary keyboard operation remains available throughout. Closing the keyboard, ending scanning, or exiting kills and reaps the worker.
 
 ### Spare worker
 

src-tauri/src/prediction/database.rs

@@ -14,6 +14,9 @@ const PHRASE_BUDGET: Duration = Duration::from_millis(200);
 /// Well under the parent's two-second reply deadline, with room for the
 /// budgets above plus the beam step that may overrun each of them.
 const SLOW_CALL: Duration = Duration::from_millis(1500);
+/// Slow calls in a row that mark the model unavailable. One stall, such as
+/// a busy disk or a waking laptop, is not a model that cannot keep up.
+const SLOW_CALLS: u8 = 3;
 
 #[derive(Clone, Copy, Debug, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
 pub enum Status {
@@ -33,6 +36,7 @@ enum State {
 pub struct Database {
     state: State,
     slow_call: Duration,
+    slow_calls: u8,
 }
 impl Database {
     pub fn open(model: &Path) -> Self {
@@ -44,6 +48,7 @@ impl Database {
         Self {
             state: State::Loading(rx),
             slow_call: SLOW_CALL,
+            slow_calls: 0,
         }
     }
     pub fn status(&mut self) -> Status {
@@ -75,11 +80,15 @@ impl Database {
             start + SEARCH_BUDGET,
             start + SEARCH_BUDGET + PHRASE_BUDGET,
         );
-        if start.elapsed() > self.slow_call || result.is_err() {
+        self.slow_calls = if start.elapsed() > self.slow_call {
+            self.slow_calls + 1
+        } else {
+            0
+        };
+        let Some(mut prediction) = result.ok().filter(|_| self.slow_calls < SLOW_CALLS) else {
             self.state = State::Unavailable;
             return (Status::Unavailable, Prediction::default());
-        }
-        let mut prediction = result.unwrap();
+        };
         prediction.words.truncate(5);
         prediction.phrases.truncate(2);
         (Status::Ready, prediction)
@@ -89,6 +98,7 @@ impl Database {
         Self {
             state: State::Ready(Box::new(FakeModel)),
             slow_call: SLOW_CALL,
+            slow_calls: 0,
         }
     }
 }
@@ -142,6 +152,7 @@ mod tests {
         let mut db = Database {
             state: State::Loading(rx),
             slow_call: SLOW_CALL,
+            slow_calls: 0,
         };
         assert_eq!(
             db.predict(&context()),
@@ -178,8 +189,31 @@ mod tests {
             db.predict(&context()),
             (Status::Unavailable, Prediction::default())
         );
-        db.state = State::Ready(Box::new(Fake(Ok(vec!["late"]), Duration::from_millis(20))));
-        db.slow_call = Duration::from_millis(10);
+    }
+    #[test]
+    fn only_repeated_slow_calls_make_the_model_unavailable() {
+        // Wide margins, so a busy machine cannot make a prompt call look slow.
+        let late = || State::Ready(Box::new(Fake(Ok(vec!["late"]), Duration::from_millis(250))));
+        let prompt = || State::Ready(Box::new(Fake(Ok(vec!["prompt"]), Duration::ZERO)));
+        let mut db = Database::fixture();
+        db.slow_call = Duration::from_millis(200);
+        // A slow call still answers, and a prompt one forgives it.
+        {
+            db.state = late();
+            for _ in 1..SLOW_CALLS {
+                let (status, prediction) = db.predict(&context());
+                assert_eq!(
+                    (status, prediction.words),
+                    (Status::Ready, vec!["late".to_owned()])
+                );
+            }
+            db.state = prompt();
+            assert_eq!(db.predict(&context()).0, Status::Ready);
+        }
+        db.state = late();
+        for _ in 1..SLOW_CALLS {
+            assert_eq!(db.predict(&context()).0, Status::Ready);
+        }
         assert_eq!(
             db.predict(&context()),
             (Status::Unavailable, Prediction::default())
@@ -188,5 +222,9 @@ mod tests {
             db.predict(&context()),
             (Status::Unavailable, Prediction::default())
         );
+        // A model error is not forgiven, however fast.
+        let mut db = Database::fixture();
+        db.state = State::Ready(Box::new(Fake(Err(()), Duration::ZERO)));
+        assert_eq!(db.predict(&context()).0, Status::Unavailable);
     }
 }

src-tauri/src/prediction/mod.rs

@@ -135,6 +135,12 @@ impl Drop for Client {
         self.terminate();
     }
 }
+/// A second missed reply deadline this soon after a worker was replaced
+/// means the model cannot keep up, rather than one stall.
+const RECOVERY: Duration = Duration::from_secs(60);
+/// Replacements one keyboard gets, however far apart. A slow scanner on a
+/// machine that stalls on every call would otherwise never be offered Retry.
+const RECOVERIES: u8 = 2;
 /// How long a spare worker waits for the keyboard to be opened again.
 const SPARE_IDLE: Duration = Duration::from_secs(120);
 
@@ -170,6 +176,10 @@ fn expire(spare: &mut Option<Spare>, now: Instant) {
 #[derive(Default)]
 struct Service {
     client: Option<Client>,
+    /// When the worker was last replaced after a missed deadline, and how
+    /// often since the keyboard opened.
+    recovered: Option<Instant>,
+    recoveries: u8,
     /// What the worker was started with, to start its replacement.
     model: Option<PathBuf>,
     ignored: Vec<u32>,
@@ -295,13 +305,40 @@ impl Service {
         keyboard.prediction_loading = status == database::Status::Loading;
         keyboard.prediction_tracking(tracking);
     }
+    /// A reply missed its deadline. One stall should not cost a switch user a
+    /// trip to Retry, so the worker and its private text context are killed
+    /// and a new worker starts on the next poll. A second missed deadline
+    /// within `RECOVERY` of that, or one after `RECOVERIES` replacements,
+    /// fails for good whatever the replacement answered in between, so a
+    /// model that cannot keep up never loops.
+    fn timed_out(&mut self, keyboard: &mut Keyboard, now: Instant) {
+        if self.recoveries >= RECOVERIES
+            || self
+                .recovered
+                .is_some_and(|at| now.duration_since(at) < RECOVERY)
+        {
+            self.fail(keyboard);
+            return;
+        }
+        self.discard(keyboard, false);
+        self.recovered = Some(now);
+        self.recoveries += 1;
+        keyboard.prediction_loading = true;
+        self.reset = true;
+        self.generation = self.generation.wrapping_add(1);
+        self.last = None;
+    }
     fn fail(&mut self, keyboard: &mut Keyboard) {
-        self.client = None;
         self.failed = true;
+        self.discard(keyboard, true);
+    }
+    /// Kills and reaps the worker and fails an acceptance waiting on it.
+    fn discard(&mut self, keyboard: &mut Keyboard, failed: bool) {
+        self.client = None;
         self.tracking = false;
         self.edit.clear();
         self.outstanding = None;
-        keyboard.predictions(None, true);
+        keyboard.predictions(None, failed);
         if self.accepting || self.accept.is_some() {
             keyboard.failed();
         }
@@ -531,7 +568,7 @@ pub fn poll(app: &AppHandle, keyboard: Option<&mut Keyboard>, enabled: bool, ign
             keyboard.prediction_loading = true;
         }
         if s.response_timed_out(Instant::now()) {
-            s.fail(keyboard);
+            s.timed_out(keyboard, Instant::now());
             return;
         }
         let reply = s.client.as_ref().unwrap().replies.try_recv();
@@ -781,6 +818,75 @@ mod tests {
         stop();
     }
     #[test]
+    fn one_missed_deadline_replaces_the_worker_and_a_second_fails() {
+        let mut keyboard = Keyboard::new(false);
+        let mut service = Service {
+            client: Some(sleeper()),
+            outstanding: Some(Instant::now()),
+            edit_revision: 1,
+            edit: vec![InputScope::capture().record(Edit::Append("test".into()))],
+            ..Default::default()
+        };
+        let now = Instant::now();
+        let held = service.client.as_ref().unwrap().child.id();
+        assert!(alive(held));
+        service.timed_out(&mut keyboard, now);
+        assert!(service.client.is_none() && !alive(held));
+        assert!(!service.failed && service.reset);
+        assert!(service.edit.is_empty() && service.outstanding.is_none());
+        assert!(!keyboard.error);
+        // Whatever the replacement answers, a second miss soon after fails.
+        let batch = worker::Batch {
+            token: 1,
+            words: vec!["water".into()],
+        };
+        service.received_suggestions(
+            &mut keyboard,
+            service.generation,
+            1,
+            Some(batch),
+            true,
+            database::Status::Ready,
+        );
+        assert!(service.tracking && !keyboard.prediction_loading);
+        service.client = Some(sleeper());
+        let held = service.client.as_ref().unwrap().child.id();
+        service.timed_out(&mut keyboard, now + RECOVERY - Duration::from_secs(1));
+        assert!(service.client.is_none() && !alive(held));
+        assert!(service.failed);
+
+        // A stall long after the last one is recovered from again, but one
+        // keyboard only gets so many replacements.
+        let mut service = Service {
+            recovered: Some(now),
+            recoveries: 1,
+            ..Default::default()
+        };
+        service.timed_out(&mut keyboard, now + RECOVERY);
+        assert!(!service.failed && keyboard.prediction_loading);
+        assert_eq!(service.recovered, Some(now + RECOVERY));
+        assert_eq!(service.recoveries, RECOVERIES);
+        service.timed_out(&mut keyboard, now + RECOVERY * 10);
+        assert!(service.failed);
+    }
+    #[test]
+    fn a_missed_deadline_fails_the_acceptance_waiting_on_it() {
+        for (accepting, accept) in [(true, None), (false, Some((1, 0)))] {
+            let mut keyboard = Keyboard::new(false);
+            let mut service = Service {
+                client: Some(sleeper()),
+                outstanding: Some(Instant::now()),
+                accepting,
+                accept,
+                ..Default::default()
+            };
+            service.timed_out(&mut keyboard, Instant::now());
+            assert!(keyboard.error);
+            assert!(!service.failed && !service.accepting && service.accept.is_none());
+            assert!(service.client.is_none());
+        }
+    }
+    #[test]
     fn invalid_edit_acknowledgements_fail_closed() {
         let mut service = Service {
             edit_revision: 3,