Improve AI correction: loading indicator, field safety, memory, dismiss

- Show "…" placeholder immediately after debounce so users know the AI
  is working (KIND_AI_LOADING_FLAG = 0x20000)
- Skip AI correction in password, URL, and e-mail fields
  (mIsGeneralTextInput + mIsPasswordField guards in InputLogic)
- Protect suggestion strip from being cleared while AI suggestion is
  visible (mAiSuggestionVisible flag + clearAiSuggestionAndSetNeutral)
- Release Gemma 4 engine on TRIM_MEMORY_RUNNING_CRITICAL
- Long-press on AI suggestion dismisses it (dismissAiSuggestion)
- Warm-up model at IME startup for instant first correction

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

README.md

@@ -3,22 +3,25 @@
 HeliBoard fork that adds on-device spell and grammar correction powered by
 **Gemma 4 E2B-it** running locally via Google's **LiteRT-LM** SDK.
 
-After you finish typing a sentence (`.`, `!`, `?`), the model silently checks
+After you finish typing a sentence (`.`, `!`, `?`), the model checks it and —
-it and — if it finds an error — shows the corrected sentence in the suggestion
+if it finds an error — shows `…` immediately, then the corrected sentence in
-strip. Tap it once to replace the original text. Everything runs on-device,
+the suggestion strip. Tap it to replace the original text, or long-press to
-no network, no cloud API.
+dismiss. Everything runs on-device, no network, no cloud API.
 
 ---
 
 ## How it works
 
-```
+```text
 User types "guten rag."
           ↓
 InputLogic detects sentence-ending punctuation
+(skipped in password / URL / e-mail fields)
           ↓
 AiTriggerHook (500 ms debounce) extracts last sentence
           ↓
+"…" placeholder appears in suggestion strip immediately
+          ↓
 AiCorrectionEngine.correctSentence()
   → LiteRT-LM Engine (GPU first, CPU fallback)
   → Gemma 4 E2B-it .litertlm model
@@ -27,18 +30,19 @@ AiSuggestionManager.postSuggestion("Guten Tag.", "guten rag.")
           ↓
 Suggestion strip shows "Guten Tag." in italic/accent color
           ↓
-User taps → original text replaced with correction
+Tap  → original text replaced with correction
+Long-press → suggestion dismissed
 ```
 
 ### Components
 
 | File | Role |
-|------|------|
+| ---- | ---- |
-| `app/…/ai/AiCorrectionEngine.kt` | LiteRT-LM wrapper; GPU→CPU fallback |
+| `app/…/ai/AiCorrectionEngine.kt` | LiteRT-LM wrapper; GPU→CPU fallback; warm-up on start |
-| `app/…/ai/AiTriggerHook.kt` | Sentence detection + debounce |
+| `app/…/ai/AiTriggerHook.kt` | Sentence detection, debounce, loading placeholder |
-| `app/…/ai/AiSuggestionManager.kt` | Posts corrected text to suggestion strip |
+| `app/…/ai/AiSuggestionManager.kt` | Posts corrections and loading state to suggestion strip |
 | `patches/0001-InputLogic-ai-hook.patch` | Hook in `InputLogic.java` after `commitCodePoint` |
-| `patches/0002-LatinIME-ai-lifecycle.patch` | AI object lifecycle in `LatinIME.java` |
+| `patches/0002-LatinIME-ai-lifecycle.patch` | AI object lifecycle, pick intercept, memory trim |
 | `patches/0003-SuggestionStrip-ai-style.patch` | Italic + accent color for AI suggestions |
 
 ---
@@ -84,26 +88,71 @@ Languages & input → On-screen keyboard.
 
 ---
 
+## Features
+
+### AI correction flow
+
+- Triggers after `.` `!` `?` `…` at the end of a sentence
+- 500 ms debounce prevents firing on rapid punctuation (e.g. `...`)
+- `…` placeholder appears immediately so the user knows the AI is working
+- Correction shown in **italic** with accent color; normal suggestions are unaffected
+- Tap to accept, long-press to dismiss
+
+### Field-type safety
+
+AI correction is automatically skipped in:
+
+- Password fields
+- URL / URI fields
+- E-mail address fields
+- Web password / phonetic input fields
+
+### Model warm-up
+
+The Gemma 4 engine is loaded in the background as soon as the keyboard
+service starts (`onCreate`), so the first correction after a sentence
+appears without the initial loading delay.
+
+### Memory management
+
+Under critical memory pressure (`TRIM_MEMORY_RUNNING_CRITICAL`), the engine
+is released automatically. It reloads on the next sentence.
+
+### Backend selection
+
+`AiCorrectionEngine` tries **GPU** (WebGPU/Vulkan via LiteRT-LM) first.
+If GPU initialization fails, it falls back to **CPU** immediately. A runtime
+GPU inference error triggers a permanent CPU switch for the session.
+
+Tested on:
+
+- Google Pixel 7 (Tensor G2) — CPU backend (OpenCL not available)
+- Devices with Mali-G710 — GPU via WebGPU/Vulkan
+
+---
+
 ## Patch details
 
 The three patches modify HeliBoard source files in `heliboard/`:
 
 **0001 — InputLogic hook**  
 After every committed code point, checks `AiTriggerHook.isSentenceEnder()`.
-On a sentence-ender, reads up to 500 chars before the cursor and calls
+Only fires for general text input (not password/URL/email fields). Reads up
-`onSentenceEndDetected()`.
+to 500 chars before the cursor and calls `onSentenceEndDetected()`.
 
 **0002 — LatinIME lifecycle**  
 Instantiates `AiCorrectionEngine`, `AiSuggestionManager`, and `AiTriggerHook`
-in `onCreate()`, cleans them up in `onDestroy()`. Overrides
+in `onCreate()` and triggers warm-up. Cleans up in `onDestroy()` and
-`showAiSuggestion()` to bypass the normal suggestions-enabled gate and
+`onTrimMemory()`. Overrides `showAiSuggestion()` with a `mAiSuggestionVisible`
-intercepts `pickSuggestionManually()` to do a clean sentence replacement
+guard that prevents `setNeutralSuggestionStrip()` from clearing an active AI
-(finish composing → delete original → commit corrected text).
+suggestion. Intercepts `pickSuggestionManually()` to handle loading placeholders
+(ignore tap) and accepted corrections (delete original → commit corrected text).
+Implements `dismissAiSuggestion()` for long-press dismiss.
 
 **0003 — SuggestionStrip styling**  
-Detects the `KIND_AI_FLAG` bit (`0x10000`) on a `SuggestedWordInfo` and
+Detects `KIND_AI_FLAG` (`0x10000`) for accepted corrections (italic + accent
-applies italic style + auto-correct accent color to visually distinguish AI
+color) and `KIND_AI_LOADING_FLAG` (`0x20000`) for the loading placeholder
-suggestions from normal word predictions.
+(normal color, no italic, not tappable).
 
 ---
 
@@ -117,19 +166,6 @@ in `.gitignore`).
 
 ---
 
-## Backend selection
-
-`AiCorrectionEngine` tries **GPU** (WebGPU/Vulkan via LiteRT-LM) first.
-If GPU initialization fails, it falls back to **CPU** immediately. If a GPU
-inference error occurs at runtime, it permanently switches to CPU for the
-session and recreates the engine.
-
-Tested on:
-- Google Pixel 7 (Tensor G2) — CPU backend (OpenCL not available)
-- Devices with Mali-G710 — GPU via WebGPU/Vulkan
-
----
-
 ## License
 
 The AI integration layer (`app/src/main/java/helium314/keyboard/latin/ai/`)

app/src/main/java/helium314/keyboard/latin/ai/AiCorrectionEngine.kt

@@ -9,7 +9,9 @@ import com.google.ai.edge.litertlm.ConversationConfig
 import com.google.ai.edge.litertlm.Engine
 import com.google.ai.edge.litertlm.EngineConfig
 import com.google.ai.edge.litertlm.SamplerConfig
+import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.Dispatchers
+import kotlinx.coroutines.launch
 import kotlinx.coroutines.withContext
 import java.io.File
 
@@ -114,6 +116,21 @@ class AiCorrectionEngine(private val context: Context) {
         }
     }
 
+    /**
+     * Pre-loads the model in the background so the first correction request
+     * doesn't have to wait. Call once after IME startup.
+     */
+    fun warmUp(scope: CoroutineScope) {
+        scope.launch(Dispatchers.IO) {
+            try {
+                getOrCreate()
+                Log.i(TAG, "Warm-up complete")
+            } catch (ex: Exception) {
+                Log.w(TAG, "Warm-up failed (model may not be present yet): ${ex.message}")
+            }
+        }
+    }
+
     fun close() {
         engine?.close()
         engine = null

app/src/main/java/helium314/keyboard/latin/ai/AiSuggestionManager.kt

@@ -45,6 +45,10 @@ class AiSuggestionManager {
         const val KIND_AI_FLAG = 0x10000
         const val KIND_AI_CORRECTION = SuggestedWordInfo.KIND_CORRECTION or KIND_AI_FLAG
 
+        /** Marks a transient "loading" placeholder — not tappable. */
+        const val KIND_AI_LOADING_FLAG = 0x20000
+        const val KIND_AI_LOADING = SuggestedWordInfo.KIND_CORRECTION or KIND_AI_LOADING_FLAG
+
         /**
          * Factory method accessible from Java (LatinIME.java patch).
          * Returns a [CoroutineScope] tied to a SupervisorJob so individual
@@ -126,6 +130,20 @@ class AiSuggestionManager {
     /** Returns the original (pre-correction) sentence, or null if none is pending. */
     fun getOriginalSentence(): String? = originalSentence
 
+    /** Shows a non-interactive "…" placeholder while inference is running. */
+    fun showLoadingPlaceholder() {
+        val wordInfo = SuggestedWordInfo(
+            "…", "", Int.MAX_VALUE, KIND_AI_LOADING,
+            null, SuggestedWordInfo.NOT_AN_INDEX, SuggestedWordInfo.NOT_A_CONFIDENCE
+        )
+        val words = SuggestedWords(
+            arrayListOf(wordInfo), null, wordInfo,
+            false, false, false,
+            SuggestedWords.INPUT_STYLE_PREDICTION, SuggestedWords.NOT_A_SEQUENCE_NUMBER
+        )
+        accessor?.showAiSuggestion(words)
+    }
+
     /**
      * Clears the AI suggestion from the strip.
      * Called when the correction matches the original (no change needed)
@@ -134,6 +152,6 @@ class AiSuggestionManager {
     fun clearSuggestion() {
         _currentSuggestion.value = null
         originalSentence = null
-        accessor?.setNeutralSuggestionStrip()
+        accessor?.clearAiSuggestionAndSetNeutral()
     }
 }

app/src/main/java/helium314/keyboard/latin/ai/AiTriggerHook.kt

@@ -63,6 +63,7 @@ class AiTriggerHook(
         debounceJob?.cancel()
         debounceJob = scope.launch {
             delay(DEBOUNCE_MS)
+            manager.showLoadingPlaceholder()
             Log.d(TAG, "Triggering AI correction for: \"$sentence\"")
             val corrected = engine.correctSentence(sentence)
             if (corrected != sentence) {