Merge pull request #4051 from vitorpamplona/claude/mlkit-genai-nullpointer-nge2pm

Fix ML Kit GenAI crash on future cancellation
2026-09-14 00:55:08 +00:00 · 2026-09-03 15:54:08 -04:00
parent 9585e66e60 8b47588d37
commit e6e05c5915
5 changed files with 308 additions and 70 deletions
@@ -445,6 +445,12 @@ open class ShortNotePostViewModel :
    private var lastComputedText: String = ""
    private var lastStatusCheckAt: Long = 0

+    /** True once a batch is past its debounce and its inferences have reached the model. */
+    private var aiInferenceRunning = false
+
+    /** Newest draft text seen while a batch was running, picked up when that batch ends. */
+    private var aiPendingText: String? = null
+
    /**
     * Whether there is anything to show. The user's preference is watched by the screen so
     * that turning the setting off hides the panel right away.
@@ -505,7 +511,7 @@ open class ShortNotePostViewModel :
        }
        if (text == lastComputedText) return

-        val assistant = writingAssistant ?: return
+        if (writingAssistant == null) return
        if (!isAiEnabledInSettings()) return

        if (aiStatus !is WritingAssistantStatus.Available) {
@@ -513,6 +519,17 @@ open class ShortNotePostViewModel :
            return
        }

+        // An inference that has reached the model cannot be recalled — cancelling its future
+        // is what used to crash the app, so the bridge detaches instead (see GenAiFutures).
+        // Abandoning a running batch would therefore leave seven rewrites burning on-device
+        // compute for text the user has already moved past, and every later keystroke would
+        // stack seven more on top. So only the debounce window is cancellable: once a batch
+        // is under way it is left to finish, and the newest text is picked up when it ends.
+        if (aiInferenceRunning) {
+            aiPendingText = text
+            return
+        }
+
        aiComputeJob?.cancel()
        aiResults = persistentMapOf()
        aiSelectedResult = null
@@ -521,35 +538,51 @@ open class ShortNotePostViewModel :
            viewModelScope.launch {
                delay(AI_DEBOUNCE_MS)

-                val results =
-                    coroutineScope {
-                        WritingTone.entries
-                            .map { tone ->
-                                async {
-                                    try {
-                                        assistant.transform(text, tone)
-                                    } catch (e: CancellationException) {
-                                        throw e
-                                    } catch (e: Exception) {
-                                        Log.w("ShortNotePostViewModel", "Could not run the $tone rewrite", e)
-                                        null
-                                    }
-                                }
-                            }.awaitAll()
-                            .filterNotNull()
-                            // A model that hands back the input unchanged has nothing to offer.
-                            .filter { it.transformedText.isNotBlank() && it.transformedText != it.originalText }
-                            .associateBy { it.tone }
-                            .toImmutableMap()
-                    }
+                aiInferenceRunning = true
+                try {
+                    runAiBatch(text)
+                } finally {
+                    aiInferenceRunning = false
+                }

-                aiResults = results
-                // Only now: a run that was cancelled must not mark this text as done, or
-                // coming back to it would show nothing.
-                lastComputedText = text
+                // The draft moved on while this batch ran: start the next one now.
+                val pending = aiPendingText
+                aiPendingText = null
+                if (pending != null && pending != text) precomputeAiResults()
            }
    }

+    private suspend fun runAiBatch(text: String) {
+        val assistant = writingAssistant ?: return
+
+        val results =
+            coroutineScope {
+                WritingTone.entries
+                    .map { tone ->
+                        async {
+                            try {
+                                assistant.transform(text, tone)
+                            } catch (e: CancellationException) {
+                                throw e
+                            } catch (e: Exception) {
+                                Log.w("ShortNotePostViewModel", "Could not run the $tone rewrite", e)
+                                null
+                            }
+                        }
+                    }.awaitAll()
+                    .filterNotNull()
+                    // A model that hands back the input unchanged has nothing to offer.
+                    .filter { it.transformedText.isNotBlank() && it.transformedText != it.originalText }
+                    .associateBy { it.tone }
+                    .toImmutableMap()
+            }
+
+        aiResults = results
+        // Only now: a run that was cancelled must not mark this text as done, or
+        // coming back to it would show nothing.
+        lastComputedText = text
+    }
+
    fun selectAiResult(tone: WritingTone) {
        aiSelectedResult = aiResults[tone]
    }
@@ -574,6 +607,7 @@ open class ShortNotePostViewModel :
    private fun resetAiState() {
        aiComputeJob?.cancel()
        aiComputeJob = null
+        aiPendingText = null
        aiResults = persistentMapOf()
        aiSelectedResult = null
        lastComputedText = ""
@@ -0,0 +1,85 @@
+/*
+ * Copyright (c) 2025 Vitor Pamplona
+ *
+ * Permission is hereby granted, free of charge, to any person obtaining a copy of
+ * this software and associated documentation files (the "Software"), to deal in
+ * the Software without restriction, including without limitation the rights to use,
+ * copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the
+ * Software, and to permit persons to whom the Software is furnished to do so,
+ * subject to the following conditions:
+ *
+ * The above copyright notice and this permission notice shall be included in all
+ * copies or substantial portions of the Software.
+ *
+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+ * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS
+ * FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR
+ * COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN
+ * AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION
+ * WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
+ */
+package com.vitorpamplona.amethyst.service.ai
+
+import com.google.common.util.concurrent.ListenableFuture
+import kotlinx.coroutines.CancellationException
+import kotlinx.coroutines.suspendCancellableCoroutine
+import java.util.concurrent.ExecutionException
+import java.util.concurrent.Executor
+import kotlin.coroutines.resume
+import kotlin.coroutines.resumeWithException
+
+/** Completes the continuation on whichever thread finished the future. */
+private val DIRECT_EXECUTOR = Executor { it.run() }
+
+/**
+ * Awaits an ML Kit GenAI [ListenableFuture], deliberately **detaching** when the caller is
+ * cancelled instead of cancelling the future.
+ *
+ * Cancelling one of these futures crashes the process from a thread we don't own. When
+ * `genai-rewriting` 1.0.0-beta1 starts an inference it asks AiCore for an
+ * `ICancellationCallback` and registers it as the future's cancellation listener with no null
+ * check (`zzbh.attachCompleter` builds `new zzbt(handle)`). AiCore is free to answer with a
+ * null binder — `Parcel.readStrongBinder()` then returns null and the handle is null — so the
+ * listener holds nothing, and the moment the future is cancelled `zzby.zzk(null)` dereferences
+ * it:
+ *
+ * ```
+ * Thread: AiCoreClientWorker-thread-5
+ * java.lang.NullPointerException: Attempt to invoke interface method
+ *   'void com.google.android.gms.internal.mlkit_genai_rewriting.zzp.zzd()' on a null object reference
+ *     at com.google.android.gms.internal.mlkit_genai_rewriting.zzby.zzk
+ *     at com.google.android.gms.internal.mlkit_genai_rewriting.zzbt.run
+ *     at java.util.concurrent.ThreadPoolExecutor.runWorker
+ * ```
+ *
+ * The NPE is thrown on ML Kit's own worker pool, so no `try`/`catch` on our side can see it:
+ * it goes straight to the default uncaught handler and takes the app down. The composer
+ * cancels these futures as a matter of course — a keystroke replaces the in-flight batch of
+ * tones, and leaving the composer cancels `viewModelScope` — which is what made a beta-library
+ * race a routine crash.
+ *
+ * Detaching leaves the inference to finish with nobody listening for its result. That wastes a
+ * little on-device compute; cancelling wastes the process. There is nothing to fix on the ML
+ * Kit side either — `genai-rewriting` has shipped no version past `1.0.0-beta1`.
+ */
+internal suspend fun <T> ListenableFuture<T>.awaitDetached(): T =
+    suspendCancellableCoroutine { continuation ->
+        addListener(
+            {
+                // A detached caller is already gone: don't even read the result.
+                if (continuation.isActive) {
+                    try {
+                        continuation.resume(get())
+                    } catch (e: CancellationException) {
+                        continuation.cancel(e)
+                    } catch (e: ExecutionException) {
+                        continuation.resumeWithException(e.cause ?: e)
+                    } catch (e: Exception) {
+                        continuation.resumeWithException(e)
+                    }
+                }
+            },
+            DIRECT_EXECUTOR,
+        )
+        // Deliberately no invokeOnCancellation { cancel(true) }: cancelling is the crash.
+    }
@@ -29,14 +29,15 @@ import com.google.mlkit.genai.imagedescription.ImageDescriber
 import com.google.mlkit.genai.imagedescription.ImageDescriberOptions
 import com.google.mlkit.genai.imagedescription.ImageDescription
 import com.google.mlkit.genai.imagedescription.ImageDescriptionRequest
+import kotlinx.coroutines.CancellationException
 import kotlinx.coroutines.Dispatchers
 import kotlinx.coroutines.withContext

 /**
- * Unified alt-text suggestion service.
+ * Alt-text suggestion service backed by Gemini Nano through AICore.
 *
- * Prefers Gemini-Nano-backed `genai-image-description` for full descriptive sentences when
- * the device supports AICore; falls back to the legacy keyword `image-labeling` model otherwise.
+ * Produces a full descriptive sentence via `genai-image-description`, or null when the device
+ * cannot run the model.
 */
 class MLKitImageLabelService(
    private val context: Context,
@@ -63,17 +64,26 @@ class MLKitImageLabelService(
        withContext(Dispatchers.IO) {
            val client = ensureDescriber() ?: return@withContext null
            try {
+                // awaitDetached rather than ListenableFuture.get(): describing an image takes
+                // seconds, and get() would hold an IO thread for all of it — uninterruptibly,
+                // so backing out of the composer would leave the thread pinned until AICore
+                // answered. Detaching also keeps us off the cancellation path that crashes
+                // these clients; see GenAiFutures.
                val status =
-                    cachedGenAiStatus ?: client.checkFeatureStatus().get().also { cachedGenAiStatus = it }
+                    cachedGenAiStatus ?: client.checkFeatureStatus().awaitDetached().also { cachedGenAiStatus = it }
                if (status != FeatureStatus.AVAILABLE) return@withContext null
                val bitmap = loadDownscaledBitmap(uri) ?: return@withContext null
                val request = ImageDescriptionRequest.builder(bitmap).build()
                client
                    .runInference(request)
-                    .get()
+                    .awaitDetached()
                    .description
                    .trim()
                    .takeIf { it.isNotEmpty() }
+            } catch (e: CancellationException) {
+                // Now that the awaits suspend, cancelling the caller lands here: it must not
+                // be swallowed as "no suggestion", or the coroutine would keep running.
+                throw e
            } catch (_: Exception) {
                null
            }
@@ -113,8 +123,6 @@ class MLKitImageLabelService(
    }

    companion object {
-        private const val MIN_CONFIDENCE = 0.6f
-        private const val MAX_LABELS = 5
        private const val TARGET_DIM_PX = 1024
    }
 }
@@ -22,7 +22,6 @@ package com.vitorpamplona.amethyst.service.ai

 import android.content.Context
 import com.google.android.gms.tasks.Tasks
-import com.google.common.util.concurrent.ListenableFuture
 import com.google.mlkit.genai.common.DownloadCallback
 import com.google.mlkit.genai.common.FeatureStatus
 import com.google.mlkit.genai.common.GenAiException
@@ -42,15 +41,10 @@ import com.vitorpamplona.amethyst.service.lang.LanguageTranslatorService
 import com.vitorpamplona.quartz.utils.Log
 import kotlinx.coroutines.CancellationException
 import kotlinx.coroutines.Dispatchers
-import kotlinx.coroutines.suspendCancellableCoroutine
 import kotlinx.coroutines.sync.Mutex
 import kotlinx.coroutines.sync.withLock
 import kotlinx.coroutines.withContext
 import java.util.concurrent.ConcurrentHashMap
-import java.util.concurrent.ExecutionException
-import java.util.concurrent.Executor
-import kotlin.coroutines.resume
-import kotlin.coroutines.resumeWithException

 /**
 * On-device writing assistance backed by ML Kit GenAI (Gemini Nano through AICore).
@@ -146,12 +140,12 @@ class MLKitWritingAssistant(
                    val rewriter = getRewriter(RewriterOptions.OutputType.REPHRASE, WritingLanguage.ENGLISH)
                    if (!downloadRequested) {
                        downloadRequested = true
-                        rewriter.downloadFeature(SilentDownloadCallback).await()
+                        rewriter.downloadFeature(SilentDownloadCallback).awaitDetached()
                        // The proofreader ships as its own feature: fetch it too, or the
                        // CORRECT tone would stay missing forever. A failure here is not
                        // fatal — the rewriting tones still work.
                        try {
-                            getProofreader(WritingLanguage.ENGLISH).downloadFeature(SilentDownloadCallback).await()
+                            getProofreader(WritingLanguage.ENGLISH).downloadFeature(SilentDownloadCallback).awaitDetached()
                        } catch (e: CancellationException) {
                            throw e
                        } catch (e: Exception) {
@@ -169,7 +163,7 @@ class MLKitWritingAssistant(
        }

    private suspend fun statusOf(rewriter: Rewriter): WritingAssistantStatus =
-        when (rewriter.checkFeatureStatus().await()) {
+        when (rewriter.checkFeatureStatus().awaitDetached()) {
            FeatureStatus.AVAILABLE -> WritingAssistantStatus.Available
            FeatureStatus.DOWNLOADING -> WritingAssistantStatus.Downloading
            FeatureStatus.DOWNLOADABLE -> WritingAssistantStatus.Downloadable
@@ -231,7 +225,7 @@ class MLKitWritingAssistant(
    ): String {
        // Building a client touches disk and another process; awaiting the inference does not.
        val rewriter = withContext(Dispatchers.IO) { getRewriter(outputType, language) }
-        val result = rewriter.runInference(RewritingRequest.builder(text).build()).await()
+        val result = rewriter.runInference(RewritingRequest.builder(text).build()).awaitDetached()
        return result.results.firstOrNull()?.text ?: text
    }

@@ -240,7 +234,7 @@ class MLKitWritingAssistant(
        language: WritingLanguage,
    ): String {
        val proofreader = withContext(Dispatchers.IO) { getProofreader(language) }
-        val result = proofreader.runInference(ProofreadingRequest.builder(text).build()).await()
+        val result = proofreader.runInference(ProofreadingRequest.builder(text).build()).awaitDetached()
        return result.results.firstOrNull()?.text ?: text
    }

@@ -250,29 +244,6 @@ class MLKitWritingAssistant(
        proofreaders.keys.toList().forEach { proofreaders.remove(it)?.close() }
    }

-    /**
-     * Bridges a [ListenableFuture] into a cancellable suspend call: cancelling the caller
-     * cancels the inference instead of leaving it running on a thread nobody waits for.
-     */
-    private suspend fun <T> ListenableFuture<T>.await(): T =
-        suspendCancellableCoroutine { continuation ->
-            addListener(
-                {
-                    try {
-                        continuation.resume(get())
-                    } catch (e: CancellationException) {
-                        continuation.cancel(e)
-                    } catch (e: ExecutionException) {
-                        continuation.resumeWithException(e.cause ?: e)
-                    } catch (e: Exception) {
-                        continuation.resumeWithException(e)
-                    }
-                },
-                DIRECT_EXECUTOR,
-            )
-            continuation.invokeOnCancellation { cancel(true) }
-        }
-
    /**
     * The two ML Kit APIs declare their own language constants. They happen to share the
     * same numbering today; this enum keeps our call sites from depending on that.
@@ -313,8 +284,5 @@ class MLKitWritingAssistant(

    companion object {
        private const val TAG = "MLKitWritingAssistant"
-
-        /** Completes the continuation on whichever thread finished the future. */
-        private val DIRECT_EXECUTOR = Executor { it.run() }
    }
 }
@@ -0,0 +1,143 @@
+/*
+ * Copyright (c) 2025 Vitor Pamplona
+ *
+ * Permission is hereby granted, free of charge, to any person obtaining a copy of
+ * this software and associated documentation files (the "Software"), to deal in
+ * the Software without restriction, including without limitation the rights to use,
+ * copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the
+ * Software, and to permit persons to whom the Software is furnished to do so,
+ * subject to the following conditions:
+ *
+ * The above copyright notice and this permission notice shall be included in all
+ * copies or substantial portions of the Software.
+ *
+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+ * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS
+ * FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR
+ * COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN
+ * AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION
+ * WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
+ */
+package com.vitorpamplona.amethyst.service.ai
+
+import com.google.common.util.concurrent.ListenableFuture
+import kotlinx.coroutines.ExperimentalCoroutinesApi
+import kotlinx.coroutines.async
+import kotlinx.coroutines.launch
+import kotlinx.coroutines.test.advanceUntilIdle
+import kotlinx.coroutines.test.runTest
+import org.junit.Assert.assertEquals
+import org.junit.Assert.assertFalse
+import org.junit.Assert.assertTrue
+import org.junit.Assert.fail
+import org.junit.Test
+import java.util.concurrent.CancellationException
+import java.util.concurrent.ExecutionException
+import java.util.concurrent.Executor
+import java.util.concurrent.TimeUnit
+
+/**
+ * The ML Kit GenAI futures must never be cancelled: `genai-rewriting` 1.0.0-beta1 registers a
+ * possibly-null `ICancellationCallback` as the cancellation listener, so a cancel dereferences
+ * null on its own worker thread and kills the process. See [awaitDetached].
+ */
+class GenAiFuturesTest {
+    /** Minimal hand-rolled future: records cancellation and completes on demand. */
+    private class FakeFuture<T> : ListenableFuture<T> {
+        private val listeners = mutableListOf<Pair<Runnable, Executor>>()
+        private var value: T? = null
+        private var failure: Throwable? = null
+        private var done = false
+
+        var cancelCalls = 0
+            private set
+
+        fun complete(result: T) = finish { value = result }
+
+        fun fail(error: Throwable) = finish { failure = error }
+
+        private fun finish(set: () -> Unit) {
+            if (done) return
+            set()
+            done = true
+            listeners.forEach { (runnable, executor) -> executor.execute(runnable) }
+            listeners.clear()
+        }
+
+        override fun addListener(
+            listener: Runnable,
+            executor: Executor,
+        ) {
+            if (done) executor.execute(listener) else listeners.add(listener to executor)
+        }
+
+        override fun cancel(mayInterruptIfRunning: Boolean): Boolean {
+            cancelCalls++
+            finish { failure = CancellationException("cancelled") }
+            return true
+        }
+
+        override fun isCancelled() = failure is CancellationException
+
+        override fun isDone() = done
+
+        override fun get(): T {
+            check(done) { "not done" }
+            failure?.let { if (it is CancellationException) throw it else throw ExecutionException(it) }
+            @Suppress("UNCHECKED_CAST")
+            return value as T
+        }
+
+        override fun get(
+            timeout: Long,
+            unit: TimeUnit,
+        ): T = get()
+    }
+
+    @OptIn(ExperimentalCoroutinesApi::class)
+    @Test
+    fun `returns the value the future completes with`() =
+        runTest {
+            val future = FakeFuture<String>()
+            val awaited = async { future.awaitDetached() }
+            advanceUntilIdle()
+
+            future.complete("rewritten")
+
+            assertEquals("rewritten", awaited.await())
+        }
+
+    @Test
+    fun `unwraps the cause out of an ExecutionException`() =
+        runTest {
+            val future = FakeFuture<String>()
+            future.fail(IllegalStateException("AICore is not available"))
+
+            try {
+                future.awaitDetached()
+                fail("Expected the cause to surface")
+            } catch (e: IllegalStateException) {
+                assertEquals("AICore is not available", e.message)
+            }
+        }
+
+    /** The regression this file exists for. */
+    @OptIn(ExperimentalCoroutinesApi::class)
+    @Test
+    fun `cancelling the caller does not cancel the future`() =
+        runTest {
+            val future = FakeFuture<String>()
+            val job = launch { future.awaitDetached() }
+            advanceUntilIdle()
+
+            job.cancel()
+            advanceUntilIdle()
+
+            assertEquals(0, future.cancelCalls)
+            assertFalse(future.isDone)
+
+            // And the inference still finishing afterwards must not blow up on the detached caller.
+            future.complete("rewritten")
+            assertTrue(job.isCancelled)
+        }
+}