@dr33m/react-native-litert-lm 0.6.0 → 0.6.1

package/android/build.gradle

@@ -29,7 +29,28 @@ android {
         }
         
         ndk {
-            abiFilters 'arm64-v8a'
+            // Follow the app's `reactNativeArchitectures` rather than pinning
+            // arm64. Pinning it broke x86_64 emulator builds of consuming apps:
+            // this module still configured an arm64-v8a build while the app —
+            // and so the NitroModules prefab it links against — was built
+            // x86_64 only, and LiteRTLM failed to link with undefined
+            // `margelo::nitro::JHybridObject::*` vtable symbols.
+            //
+            // The pin bought nothing: com.google.ai.edge.litertlm ships both
+            // arm64-v8a and x86_64.
+            def supportedAbis = ['arm64-v8a', 'x86_64']
+            def requestedAbis = (project.findProperty('reactNativeArchitectures')
+                    ?: rootProject.findProperty('reactNativeArchitectures')
+                    ?: supportedAbis.join(','))
+                .toString()
+                .split(',')
+                .collect { it.trim() }
+            def selectedAbis = requestedAbis.findAll { supportedAbis.contains(it) }
+            if (selectedAbis.isEmpty()) {
+                selectedAbis = supportedAbis
+            }
+            abiFilters.clear()
+            abiFilters.addAll(selectedAbis)
         }
     }
     

package/android/src/main/java/com/margelo/nitro/dev/litert/litertlm/HybridLiteRTLM.kt

@@ -11,6 +11,7 @@ import android.app.ActivityManager
 import android.content.Context
 import java.util.Collections
 import java.util.concurrent.CountDownLatch
+import java.util.concurrent.TimeUnit
 import java.util.concurrent.atomic.AtomicReference
 import androidx.annotation.Keep
 import com.facebook.proguard.annotations.DoNotStrip
@@ -53,6 +54,9 @@ class HybridLiteRTLM : HybridLiteRTLMSpec() {
 
     companion object {
         private const val TAG = "HybridLiteRTLM"
+
+        /** How long a teardown waits for a cancelled decode loop to stop. */
+        private const val CANCEL_SETTLE_TIMEOUT_SECONDS = 2L
         private val initLock = Any()
 
         /** Cached result of OpenCL availability probe (null = not yet checked). */
@@ -88,6 +92,17 @@ class HybridLiteRTLM : HybridLiteRTLMSpec() {
     @Volatile
     private var isClosed = false
 
+    /**
+     * Latch of the streaming worker currently decoding, if any.
+     *
+     * Published so a teardown can wait for the decode loop to actually stop.
+     * Closing a conversation out from under a running worker deletes native
+     * state the worker is still reading, which is a crash the Kotlin side
+     * cannot catch.
+     */
+    @Volatile
+    private var activeStreamingLatch: CountDownLatch? = null
+
     private val modelStore = HybridModelStore()
     private var loadedModelPath: String? = null
 
@@ -240,8 +255,13 @@ class HybridLiteRTLM : HybridLiteRTLMSpec() {
                     // Only Gemma 3n bundles vision/audio executors; Gemma 4 E2B is text-only.
                     // Passing vision/audio backends to a text-only model causes
                     // vision_litert_compiled_model_executor init failures.
+                    //
+                    // "gemma3" used to be on this list and was wrong: Gemma 3 1B
+                    // is text-only, and matching it here attached vision and
+                    // audio executors that the model has no weights for. Only
+                    // Gemma 3n is multimodal, and "3n" already catches it.
                     val modelFileName = modelPath.substringAfterLast("/").lowercase()
-                    val isMultimodal = config?.multimodal ?: (modelFileName.contains("3n") || modelFileName.contains("gemma3"))
+                    val isMultimodal = config?.multimodal ?: modelFileName.contains("3n")
     
                     // Get cache directory from application context
                     val cacheDirectory = LiteRTLMInitProvider.applicationContext?.cacheDir?.absolutePath
@@ -252,6 +272,10 @@ class HybridLiteRTLM : HybridLiteRTLMSpec() {
                         ExperimentalFlags.enableSpeculativeDecoding = true
                     }
 
+                    // The most recent engine failure, kept so the thrown error can say
+                    // what went wrong instead of only that everything was tried.
+                    var lastEngineError: String? = null
+
                     // Helper: attempt engine creation with given backends, return null on failure
                     fun tryCreateEngine(
                         mainBackend: com.google.ai.edge.litertlm.Backend,
@@ -278,6 +302,11 @@ class HybridLiteRTLM : HybridLiteRTLMSpec() {
                             }
                             Engine(cfg).also { it.initialize() }
                         } catch (e: Exception) {
+                            // Held onto because it is the only description of what
+                            // actually went wrong. Reporting only "tried everything"
+                            // turned a precise engine error ("Invalid magic number")
+                            // into an unactionable one by the time it reached JS.
+                            lastEngineError = e.message
                             Log.w(TAG, "Engine creation failed with backend $mainBackend: ${e.message}")
                             null
                         }
@@ -322,47 +351,44 @@ class HybridLiteRTLM : HybridLiteRTLMSpec() {
     
                     if (isClosed) return@synchronized
 
-                    // Attempt primary backend
-                    var eng = tryCreateEngine(lmBackend, lmVisionBackend, lmAudioBackend)
-
-                    // Fallback sequence if GPU/NPU fails to initialize (mirrors iOS behavior)
-                    if (eng == null && backend != Backend.CPU) {
-                        val requestedName = if (backend == Backend.GPU) "GPU" else "NPU"
-                        Log.w(TAG, "$requestedName backend failed — trying fallback chain...")
-
-                        // Fallback 1: CPU main + GPU vision + CPU audio
-                        eng = tryCreateEngine(
-                            com.google.ai.edge.litertlm.Backend.CPU(),
-                            if (isMultimodal) com.google.ai.edge.litertlm.Backend.GPU() else null,
-                            if (isMultimodal) com.google.ai.edge.litertlm.Backend.CPU() else null
-                        )
-
-                        // Fallback 2: Full CPU for all modalities
-                        if (eng == null) {
-                            eng = tryCreateEngine(
-                                com.google.ai.edge.litertlm.Backend.CPU(),
-                                if (isMultimodal) com.google.ai.edge.litertlm.Backend.CPU() else null,
-                                if (isMultimodal) com.google.ai.edge.litertlm.Backend.CPU() else null
-                            )
-                        }
-
-                        // Fallback 3: Text-only CPU (no vision/audio executors)
-                        if (eng == null) {
-                            eng = tryCreateEngine(
-                                com.google.ai.edge.litertlm.Backend.CPU(),
-                                null,
-                                null
-                            )
-                        }
-
-                        if (eng != null) {
-                            Log.w(TAG, "$requestedName backend unavailable — fell back to CPU successfully")
+                    // Engine configurations to try, each dropping one more
+                    // capability than the last and ending at plain text-only CPU.
+                    //
+                    // The text-only row matters even when CPU was the backend
+                    // asked for. A text-only model that gets sniffed as
+                    // multimodal fails with vision and audio executors attached
+                    // and loads the moment they are dropped, and gating this
+                    // whole chain on "the caller wanted GPU or NPU" left the
+                    // default CPU path with no recovery at all: one failed
+                    // attempt and straight to an error.
+                    fun cpu() = com.google.ai.edge.litertlm.Backend.CPU()
+                    val attempts = mutableListOf(
+                        Triple(lmBackend, lmVisionBackend, lmAudioBackend)
+                    )
+                    if (isMultimodal) {
+                        attempts += Triple(cpu(), com.google.ai.edge.litertlm.Backend.GPU(), cpu())
+                        attempts += Triple(cpu(), cpu(), cpu())
+                    }
+                    attempts += Triple(cpu(), null, null)
+
+                    var eng: Engine? = null
+                    for ((index, attempt) in attempts.withIndex()) {
+                        if (isClosed) return@synchronized
+                        val (main, vision, audio) = attempt
+                        eng = tryCreateEngine(main, vision, audio)
+                        if (eng == null) continue
+                        if (index > 0) {
+                            Log.w(TAG, "Primary backend failed; loaded on fallback #$index (main=$main, vision=$vision, audio=$audio)")
+                            // Every fallback runs the main graph on CPU, so the
+                            // reported active backend has to follow.
                             backend = Backend.CPU
                         }
+                        break
                     }
 
                     engine = eng ?: throw RuntimeException(
-                        "Failed to create LiteRT-LM engine. Tried primary backend and all CPU fallbacks."
+                        "Failed to create LiteRT-LM engine after ${attempts.size} attempts: " +
+                            (lastEngineError ?: "no error reported")
                     )
                     Log.i(TAG, "Engine created and initialized successfully")
     
@@ -557,6 +583,9 @@ class HybridLiteRTLM : HybridLiteRTLMSpec() {
     }
 
     override fun resetConversation() {
+        // A reset that recreates the conversation under a running decode loop
+        // deletes the native state that loop is still reading.
+        cancelInFlightGeneration()
         synchronized(history) {
             history.clear()
         }
@@ -667,6 +696,42 @@ class HybridLiteRTLM : HybridLiteRTLMSpec() {
         }
     }
 
+    /**
+     * Signal the native decode loop to stop, then wait for the streaming worker
+     * to settle.
+     *
+     * `cancelProcess()` only asks; the worker keeps running until it notices.
+     * Anything that then closes or recreates the conversation would be deleting
+     * native state while that worker still reads it, so every such path waits
+     * here first. The timeout is a backstop: a worker that never settles is
+     * worth a warning and a teardown anyway, because the alternative is hanging
+     * the caller forever.
+     */
+    private fun cancelInFlightGeneration() {
+        val latch = activeStreamingLatch ?: return
+        try {
+            conversation?.cancelProcess()
+            if (!latch.await(CANCEL_SETTLE_TIMEOUT_SECONDS, TimeUnit.SECONDS)) {
+                Log.w(TAG, "cancelInFlightGeneration: worker did not settle in time")
+            }
+        } catch (e: Exception) {
+            Log.w(TAG, "cancelInFlightGeneration: ${e.message}")
+        }
+    }
+
+    /**
+     * Free the engine and its KV cache while leaving this instance usable.
+     *
+     * The distinction from `close()` is the whole point: `close()` sets
+     * `isClosed`, after which even `loadModel()` is refused, so using it to
+     * answer memory pressure permanently retired an instance that JS still
+     * held and believed in.
+     */
+    fun releaseUnderMemoryPressure() {
+        Log.w(TAG, "Releasing engine under memory pressure; instance stays reloadable")
+        cleanupInternal()
+    }
+
     override fun close() {
         Log.d(TAG, "Closing resources")
         isClosed = true
@@ -674,6 +739,9 @@ class HybridLiteRTLM : HybridLiteRTLMSpec() {
     }
 
     private fun cleanupInternal() {
+        // Outside the lock: the decode loop does not need initLock, and waiting
+        // for it while holding the lock would block loadModel behind it.
+        cancelInFlightGeneration()
         synchronized(initLock) {
             try {
                 conversation?.close()
@@ -737,11 +805,12 @@ class HybridLiteRTLM : HybridLiteRTLMSpec() {
             ),
             systemInstruction = systemPrompt?.let { Contents.of(Content.Text(it)) },
             initialMessages = initialMessages,
-            tools = lmTools ?: emptyList()
+            tools = lmTools ?: emptyList(),
+            // Available since LiteRT-LM 0.15.0; before that the Kotlin SDK had no
+            // way to express it and every Android reply ran to the engine default.
+            maxOutputToken = maxOutputTokens,
         )
-        // TODO: maxOutputTokens is not configurable on Android — the Kotlin SDK's
-        // ConversationConfig does not expose this parameter. Only EngineConfig.maxNumTokens
-        // (context budget) is supported. maxOutputTokens is effective on iOS only.
+
         //
         // Upstream is actively adding max_output_tokens across API surfaces:
         //   - C API:    PR #2470 (merged 2026-06-04)
@@ -896,6 +965,7 @@ class HybridLiteRTLM : HybridLiteRTLMSpec() {
                 if (onToken != null) {
                     // ── Streaming path ────────────────────────────────────────────────
                     val latch = CountDownLatch(1)
+                    activeStreamingLatch = latch
                     val errorRef = AtomicReference<Throwable?>(null)
                     val fullResponseBuilder = StringBuilder()
                     val thinkingBuilder = StringBuilder()
@@ -922,7 +992,11 @@ class HybridLiteRTLM : HybridLiteRTLMSpec() {
                         latch.countDown()
                     }
 
-                    latch.await()
+                    try {
+                        latch.await()
+                    } finally {
+                        activeStreamingLatch = null
+                    }
                     errorRef.get()?.let { throw RuntimeException("execute streaming failed: ${it.message}", it) }
                     val capturedToolCalls = synchronized(pendingToolCalls) {
                         pendingToolCalls.toTypedArray().also { pendingToolCalls.clear() }

package/android/src/main/java/com/margelo/nitro/dev/litert/litertlm/LiteRTLMRegistry.kt

@@ -1,5 +1,6 @@
 package com.margelo.nitro.dev.litert.litertlm
 
+import android.content.ComponentCallbacks2
 import java.util.Collections
 import java.util.WeakHashMap
 import android.util.Log
@@ -10,7 +11,7 @@ import android.util.Log
  */
 object LiteRTLMRegistry {
     private const val TAG = "LiteRTLMRegistry"
-    
+
     // Use WeakSet-like structure to prevent leaks
     private val instances = Collections.newSetFromMap(WeakHashMap<HybridLiteRTLM, Boolean>())
 
@@ -20,13 +21,31 @@ object LiteRTLMRegistry {
         }
     }
 
+    /**
+     * Whether a trim level is a real emergency worth dropping engines for.
+     *
+     * Trim levels are event codes, not a severity scale, so a `>=` comparison
+     * is the wrong shape entirely: TRIM_MEMORY_UI_HIDDEN is 20 and fires every
+     * time the screen locks or the reader presses home, which says nothing
+     * about memory. Testing `level >= TRIM_MEMORY_RUNNING_LOW` (10) therefore
+     * matched a screen lock and tore the engine down on it.
+     *
+     * Only two codes mean the process dies without intervention:
+     * RUNNING_CRITICAL (foreground, memory critical) and COMPLETE (cached and
+     * next in line to be killed).
+     */
+    fun isMemoryEmergency(level: Int): Boolean =
+        level == ComponentCallbacks2.TRIM_MEMORY_RUNNING_CRITICAL ||
+            level >= ComponentCallbacks2.TRIM_MEMORY_COMPLETE
+
     fun onTrimMemory(level: Int) {
-        Log.w(TAG, "Received memory warning (level=$level). Releasing resources...")
+        Log.w(TAG, "Memory emergency (level=$level). Releasing engines...")
         synchronized(instances) {
-            instances.forEach { it.close() }
-            // Note: We don't clear the set here, as close() should be idempotent
-            // and the instance might still be ref-counted by JS. 
-            // We just ensure the HEAVY native resources are gone.
+            // Release the heavy native resources but keep each instance
+            // reloadable. close() would set isClosed, and an instance that
+            // refuses loadModel() afterwards is worse than one holding memory:
+            // JS still points at it and has no way to tell it has been retired.
+            instances.forEach { it.releaseUnderMemoryPressure() }
         }
     }
 }

package/android/src/main/java/dev/litert/litertlm/LiteRTLMInitProvider.kt

@@ -20,7 +20,10 @@ class LiteRTLMInitProvider : ContentProvider() {
         
         applicationContext?.registerComponentCallbacks(object : android.content.ComponentCallbacks2 {
             override fun onTrimMemory(level: Int) {
-                if (level >= android.content.ComponentCallbacks2.TRIM_MEMORY_RUNNING_LOW) {
+                // Trim levels are event codes, not severities: a `>=` gate here
+                // matched TRIM_MEMORY_UI_HIDDEN (20), which fires on every screen
+                // lock, and dropped the loaded engine along with it.
+                if (com.margelo.nitro.dev.litert.litertlm.LiteRTLMRegistry.isMemoryEmergency(level)) {
                     com.margelo.nitro.dev.litert.litertlm.LiteRTLMRegistry.onTrimMemory(level)
                 }
             }

package/cpp/include/README.md

@@ -1,32 +1,28 @@
 # LiteRT-LM Headers Fallback
 
-This directory contains the LiteRT-LM C API header (`litert_lm_engine.h`) used by the iOS C++ implementation.
+This directory contains the LiteRT-LM C API headers (`litert_lm_engine.h`,
+`capabilities_c.h`) used by the iOS C++ implementation. They are vendored
+copies pinned to the LiteRT-LM version in `package.json` (`litertLm.version`).
 
 ## If Headers Are Missing
 
-If you get compilation errors like `litert_lm_engine.h: No such file or directory`, you need to manually copy the LiteRT-LM C API header here:
+If you get compilation errors like `litert_lm_engine.h: No such file or directory`, re-download the headers (replace `v0.15.0` with the pinned version):
 
-1. Clone LiteRT-LM repository:
-
-   ```bash
-   git clone https://github.com/google-ai-edge/LiteRT-LM.git /tmp/LiteRT-LM
-   cd /tmp/LiteRT-LM && git checkout v0.10.2
-   ```
-
-2. Copy the header:
-   ```bash
-   cp /tmp/LiteRT-LM/c/litert_lm_engine.h ./
-   ```
+```bash
+curl -sL "https://raw.githubusercontent.com/google-ai-edge/LiteRT-LM/v0.15.0/c/engine.h" -o litert_lm_engine.h
+curl -sL "https://raw.githubusercontent.com/google-ai-edge/LiteRT-LM/v0.15.0/schema/capabilities/capabilities_c.h" -o capabilities_c.h
+```
 
-The expected directory structure after copying:
+The expected directory structure:
 
 ```
 cpp/include/
-├── litert_lm_engine.h   # LiteRT-LM C API header
+├── litert_lm_engine.h   # LiteRT-LM C API header (upstream c/engine.h)
+├── capabilities_c.h     # Model capability probing (upstream schema/capabilities/capabilities_c.h)
 ├── stb_image.h          # Image loading for multimodal
 └── README.md
 ```
 
 ## Note
 
-On **Android**, headers are provided by the `litertlm-android` AAR via Prefab — this directory is only needed for the **iOS** build which uses the raw C API via the prebuilt XCFramework.
+On **Android**, headers are provided by the `litertlm-android` AAR via Prefab — this directory is only needed for the **iOS** build which uses the raw C API via the prebuilt XCFramework (which bundles its own copies of these headers).