react-native-litert-lm 0.2.1 → 0.2.2

package/lib/memoryTracker.js

@@ -0,0 +1,155 @@
+"use strict";
+/**
+ * Memory tracking utilities for LiteRT-LM using real native memory metrics.
+ *
+ * Records real memory usage from OS-level APIs via `getMemoryUsage()`,
+ * and stores snapshots in a native-backed ArrayBuffer allocated via
+ * `NitroModules.createNativeArrayBuffer()` (v0.34+) for zero-copy interop.
+ *
+ * @example
+ * ```typescript
+ * import { createMemoryTracker } from 'react-native-litert-lm';
+ *
+ * const tracker = createMemoryTracker(100);
+ *
+ * // Record a real snapshot (typically called internally after inference)
+ * tracker.record({
+ *   timestamp: Date.now(),
+ *   nativeHeapBytes: usage.nativeHeapBytes,
+ *   residentBytes: usage.residentBytes,
+ *   availableMemoryBytes: usage.availableMemoryBytes,
+ * });
+ *
+ * console.log(`Peak RSS: ${tracker.getPeakMemory()} bytes`);
+ * ```
+ */
+Object.defineProperty(exports, "__esModule", { value: true });
+exports.createMemoryTracker = createMemoryTracker;
+exports.createNativeBuffer = createNativeBuffer;
+const react_native_nitro_modules_1 = require("react-native-nitro-modules");
+/** Number of Float64 fields per snapshot */
+const FIELDS_PER_SNAPSHOT = 4;
+/** Bytes per Float64 value */
+const BYTES_PER_FIELD = Float64Array.BYTES_PER_ELEMENT; // 8
+/**
+ * Create a new memory tracker backed by a native ArrayBuffer.
+ *
+ * @param maxSnapshots Maximum number of snapshots to store (default: 256)
+ * @returns A MemoryTracker instance
+ */
+function createMemoryTracker(maxSnapshots = 256) {
+    const bufferSize = maxSnapshots * FIELDS_PER_SNAPSHOT * BYTES_PER_FIELD;
+    // Use NitroModules.createNativeArrayBuffer for native-backed allocation.
+    const nativeBuffer = react_native_nitro_modules_1.NitroModules.createNativeArrayBuffer(bufferSize);
+    const view = new Float64Array(nativeBuffer);
+    let currentIndex = 0;
+    return {
+        record(snapshot) {
+            if (currentIndex >= maxSnapshots) {
+                return false;
+            }
+            const offset = currentIndex * FIELDS_PER_SNAPSHOT;
+            view[offset] = snapshot.timestamp;
+            view[offset + 1] = snapshot.nativeHeapBytes;
+            view[offset + 2] = snapshot.residentBytes;
+            view[offset + 3] = snapshot.availableMemoryBytes;
+            currentIndex++;
+            return true;
+        },
+        getSnapshots() {
+            const snapshots = [];
+            for (let i = 0; i < currentIndex; i++) {
+                const offset = i * FIELDS_PER_SNAPSHOT;
+                snapshots.push({
+                    timestamp: view[offset],
+                    nativeHeapBytes: view[offset + 1],
+                    residentBytes: view[offset + 2],
+                    availableMemoryBytes: view[offset + 3],
+                });
+            }
+            return snapshots;
+        },
+        getSnapshotCount() {
+            return currentIndex;
+        },
+        getCapacity() {
+            return maxSnapshots;
+        },
+        getPeakMemory() {
+            let peak = 0;
+            for (let i = 0; i < currentIndex; i++) {
+                const rss = view[i * FIELDS_PER_SNAPSHOT + 2];
+                if (rss > peak) {
+                    peak = rss;
+                }
+            }
+            return peak;
+        },
+        getLatestSnapshot() {
+            if (currentIndex === 0)
+                return undefined;
+            const offset = (currentIndex - 1) * FIELDS_PER_SNAPSHOT;
+            return {
+                timestamp: view[offset],
+                nativeHeapBytes: view[offset + 1],
+                residentBytes: view[offset + 2],
+                availableMemoryBytes: view[offset + 3],
+            };
+        },
+        getNativeBuffer() {
+            return nativeBuffer;
+        },
+        getView() {
+            return view;
+        },
+        reset() {
+            view.fill(0);
+            currentIndex = 0;
+        },
+        getSummary() {
+            let peakRss = 0;
+            let peakHeap = 0;
+            let sumRss = 0;
+            let firstRss = 0;
+            let lastRss = 0;
+            let lastHeap = 0;
+            for (let i = 0; i < currentIndex; i++) {
+                const offset = i * FIELDS_PER_SNAPSHOT;
+                const heap = view[offset + 1];
+                const rss = view[offset + 2];
+                if (rss > peakRss)
+                    peakRss = rss;
+                if (heap > peakHeap)
+                    peakHeap = heap;
+                sumRss += rss;
+                if (i === 0)
+                    firstRss = rss;
+                if (i === currentIndex - 1) {
+                    lastRss = rss;
+                    lastHeap = heap;
+                }
+            }
+            return {
+                snapshotCount: currentIndex,
+                peakResidentBytes: peakRss,
+                averageResidentBytes: currentIndex > 0 ? sumRss / currentIndex : 0,
+                currentResidentBytes: lastRss,
+                peakNativeHeapBytes: peakHeap,
+                currentNativeHeapBytes: lastHeap,
+                residentDeltaBytes: lastRss - firstRss,
+                trackerBufferSizeBytes: bufferSize,
+            };
+        },
+    };
+}
+/**
+ * Create a native ArrayBuffer for efficient data transfer.
+ *
+ * A convenience wrapper around `NitroModules.createNativeArrayBuffer()`.
+ *
+ * @param size Size in bytes
+ * @returns A native-backed ArrayBuffer
+ */
+function createNativeBuffer(size) {
+    return react_native_nitro_modules_1.NitroModules.createNativeArrayBuffer(size);
+}

package/lib/modelFactory.d.ts

@@ -1,5 +1,18 @@
 import { LiteRTLM } from "./specs/LiteRTLM.nitro";
+import { MemoryTracker } from "./memoryTracker";
 /**
  * Creates a new LiteRT-LM inference engine instance.
+ *
+ * Optionally creates a native-backed memory tracker using
+ * `NitroModules.createNativeArrayBuffer()` (v0.34+) for efficient
+ * zero-copy memory usage tracking.
+ *
+ * @param options.enableMemoryTracking Enable automatic memory tracking (default: false)
+ * @param options.maxMemorySnapshots Maximum number of memory snapshots to store (default: 256)
  */
-export declare function createLLM(): LiteRTLM;
+export declare function createLLM(options?: {
+    enableMemoryTracking?: boolean;
+    maxMemorySnapshots?: number;
+}): LiteRTLM & {
+    memoryTracker?: MemoryTracker;
+};

package/lib/modelFactory.js

@@ -2,13 +2,45 @@
 Object.defineProperty(exports, "__esModule", { value: true });
 exports.createLLM = createLLM;
 const react_native_nitro_modules_1 = require("react-native-nitro-modules");
+const memoryTracker_1 = require("./memoryTracker");
 /**
  * Creates a new LiteRT-LM inference engine instance.
+ *
+ * Optionally creates a native-backed memory tracker using
+ * `NitroModules.createNativeArrayBuffer()` (v0.34+) for efficient
+ * zero-copy memory usage tracking.
+ *
+ * @param options.enableMemoryTracking Enable automatic memory tracking (default: false)
+ * @param options.maxMemorySnapshots Maximum number of memory snapshots to store (default: 256)
  */
-function createLLM() {
+function createLLM(options) {
     const native = react_native_nitro_modules_1.NitroModules.createHybridObject("LiteRTLM");
+    const enableTracking = options?.enableMemoryTracking ?? false;
+    const tracker = enableTracking
+        ? (0, memoryTracker_1.createMemoryTracker)(options?.maxMemorySnapshots ?? 256)
+        : undefined;
+    /**
+     * Record a real memory snapshot using OS-level APIs via getMemoryUsage().
+     */
+    const recordMemorySnapshot = () => {
+        if (!tracker)
+            return;
+        try {
+            const usage = native.getMemoryUsage();
+            tracker.record({
+                timestamp: Date.now(),
+                nativeHeapBytes: usage.nativeHeapBytes,
+                residentBytes: usage.residentBytes,
+                availableMemoryBytes: usage.availableMemoryBytes,
+            });
+        }
+        catch {
+            // Ignore errors during memory tracking - it's non-critical
+        }
+    };
     return {
         ...native,
+        memoryTracker: tracker,
         loadModel: async (pathOrUrl, config) => {
             let modelPath = pathOrUrl;
             // Check if it's a URL
@@ -24,17 +56,47 @@ function createLLM() {
                 });
                 console.log(`Model downloaded to: ${modelPath}`);
             }
-            return native.loadModel(modelPath, config);
+            const result = await native.loadModel(modelPath, config);
+            // Record initial memory snapshot after model load
+            if (tracker) {
+                tracker.reset();
+                recordMemorySnapshot();
+            }
+            return result;
+        },
+        sendMessage: async (...args) => {
+            const result = await native.sendMessage(...args);
+            recordMemorySnapshot();
+            return result;
+        },
+        sendMessageAsync: (...args) => {
+            const [message, onToken] = args;
+            native.sendMessageAsync(message, (token, done) => {
+                onToken(token, done);
+                if (done) {
+                    recordMemorySnapshot();
+                }
+            });
+        },
+        sendMessageWithImage: async (...args) => {
+            const result = await native.sendMessageWithImage(...args);
+            recordMemorySnapshot();
+            return result;
+        },
+        sendMessageWithAudio: async (...args) => {
+            const result = await native.sendMessageWithAudio(...args);
+            recordMemorySnapshot();
+            return result;
         },
-        // Bind valid methods to native instance
-        sendMessage: native.sendMessage.bind(native),
-        sendMessageAsync: native.sendMessageAsync.bind(native),
-        sendMessageWithImage: native.sendMessageWithImage.bind(native),
-        sendMessageWithAudio: native.sendMessageWithAudio.bind(native),
         getHistory: native.getHistory.bind(native),
-        resetConversation: native.resetConversation.bind(native),
+        resetConversation: () => {
+            native.resetConversation();
+            // KV cache is cleared on reset, record the drop
+            recordMemorySnapshot();
+        },
         isReady: native.isReady.bind(native),
         getStats: native.getStats.bind(native),
+        getMemoryUsage: native.getMemoryUsage.bind(native),
         close: native.close.bind(native),
         downloadModel: native.downloadModel.bind(native),
         deleteModel: native.deleteModel.bind(native),

package/lib/specs/LiteRTLM.nitro.d.ts

@@ -88,6 +88,20 @@ export interface GenerationStats {
     /** Tokens per second */
     tokensPerSecond: number;
 }
+/**
+ * Real memory usage statistics from the native runtime.
+ * Measured from OS-level APIs, not estimated.
+ */
+export interface MemoryUsage {
+    /** Native heap allocated bytes (Debug.getNativeHeapAllocatedSize on Android, malloc_size on iOS) */
+    nativeHeapBytes: number;
+    /** Total process resident set size (RSS) in bytes */
+    residentBytes: number;
+    /** Available system memory in bytes */
+    availableMemoryBytes: number;
+    /** Whether the system considers memory low */
+    isLowMemory: boolean;
+}
 /**
  * LiteRT-LM: High-performance LLM inference engine.
  * Supports Gemma 3n, Phi-4, Qwen, and other .litertlm models.
@@ -175,6 +189,11 @@ export interface LiteRTLM extends HybridObject<{
      * Get the last generation statistics.
      */
     getStats(): GenerationStats;
+    /**
+     * Get real memory usage from the native runtime.
+     * Uses OS-level APIs to report actual memory consumption.
+     */
+    getMemoryUsage(): MemoryUsage;
     /**
      * Release all native resources.
      * Call this when done with the LLM instance.

package/nitrogen/generated/android/LiteRTLMOnLoad.cpp

@@ -23,26 +23,29 @@
 namespace margelo::nitro::litertlm {
 
 int initialize(JavaVM* vm) {
+  return facebook::jni::initialize(vm, []() {
+    ::margelo::nitro::litertlm::registerAllNatives();
+  });
+}
+
+void registerAllNatives() {
   using namespace margelo::nitro;
   using namespace margelo::nitro::litertlm;
-  using namespace facebook;
-
-  return facebook::jni::initialize(vm, [] {
-    // Register native JNI methods
-    margelo::nitro::litertlm::JHybridLiteRTLMSpec::registerNatives();
-    margelo::nitro::litertlm::JFunc_void_double_cxx::registerNatives();
-    margelo::nitro::litertlm::JFunc_void_std__string_bool_cxx::registerNatives();
-
-    // Register Nitro Hybrid Objects
-    HybridObjectRegistry::registerHybridObjectConstructor(
-      "LiteRTLM",
-      []() -> std::shared_ptr<HybridObject> {
-        static DefaultConstructableObject<JHybridLiteRTLMSpec::javaobject> object("com/margelo/nitro/dev/litert/litertlm/HybridLiteRTLM");
-        auto instance = object.create();
-        return instance->cthis()->shared();
-      }
-    );
-  });
+
+  // Register native JNI methods
+  margelo::nitro::litertlm::JHybridLiteRTLMSpec::registerNatives();
+  margelo::nitro::litertlm::JFunc_void_double_cxx::registerNatives();
+  margelo::nitro::litertlm::JFunc_void_std__string_bool_cxx::registerNatives();
+
+  // Register Nitro Hybrid Objects
+  HybridObjectRegistry::registerHybridObjectConstructor(
+    "LiteRTLM",
+    []() -> std::shared_ptr<HybridObject> {
+      static DefaultConstructableObject<JHybridLiteRTLMSpec::javaobject> object("com/margelo/nitro/dev/litert/litertlm/HybridLiteRTLM");
+      auto instance = object.create();
+      return instance->cthis()->shared();
+    }
+  );
 }
 
 } // namespace margelo::nitro::litertlm

package/nitrogen/generated/android/LiteRTLMOnLoad.hpp

@@ -6,20 +6,29 @@
 ///
 
 #include <jni.h>
+#include <functional>
 #include <NitroModules/NitroDefines.hpp>
 
 namespace margelo::nitro::litertlm {
 
+  [[deprecated("Use registerNatives() instead.")]]
+  int initialize(JavaVM* vm);
+
   /**
-   * Initializes the native (C++) part of LiteRTLM, and autolinks all Hybrid Objects.
-   * Call this in your `JNI_OnLoad` function (probably inside `cpp-adapter.cpp`).
+   * Register the native (C++) part of LiteRTLM, and autolinks all Hybrid Objects.
+   * Call this in your `JNI_OnLoad` function (probably inside `cpp-adapter.cpp`),
+   * inside a `facebook::jni::initialize(vm, ...)` call.
    * Example:
    * ```cpp (cpp-adapter.cpp)
    * JNIEXPORT jint JNICALL JNI_OnLoad(JavaVM* vm, void*) {
-   *   return margelo::nitro::litertlm::initialize(vm);
+   *   return facebook::jni::initialize(vm, []() {
+   *     // register all LiteRTLM HybridObjects
+   *     margelo::nitro::litertlm::registerNatives();
+   *     // any other custom registrations go here.
+   *   });
    * }
    * ```
    */
-  int initialize(JavaVM* vm);
+  void registerAllNatives();
 
 } // namespace margelo::nitro::litertlm

package/nitrogen/generated/android/c++/JHybridLiteRTLMSpec.cpp

@@ -13,6 +13,8 @@ namespace margelo::nitro::litertlm { struct Message; }
 namespace margelo::nitro::litertlm { enum class Role; }
 // Forward declaration of `GenerationStats` to properly resolve imports.
 namespace margelo::nitro::litertlm { struct GenerationStats; }
+// Forward declaration of `MemoryUsage` to properly resolve imports.
+namespace margelo::nitro::litertlm { struct MemoryUsage; }
 // Forward declaration of `LLMConfig` to properly resolve imports.
 namespace margelo::nitro::litertlm { struct LLMConfig; }
 // Forward declaration of `Backend` to properly resolve imports.
@@ -29,6 +31,8 @@ namespace margelo::nitro::litertlm { enum class Backend; }
 #include "JRole.hpp"
 #include "GenerationStats.hpp"
 #include "JGenerationStats.hpp"
+#include "MemoryUsage.hpp"
+#include "JMemoryUsage.hpp"
 #include "LLMConfig.hpp"
 #include <optional>
 #include "JLLMConfig.hpp"
@@ -204,6 +208,11 @@ namespace margelo::nitro::litertlm {
     auto __result = method(_javaPart);
     return __result->toCpp();
   }
+  MemoryUsage JHybridLiteRTLMSpec::getMemoryUsage() {
+    static const auto method = javaClassStatic()->getMethod<jni::local_ref<JMemoryUsage>()>("getMemoryUsage");
+    auto __result = method(_javaPart);
+    return __result->toCpp();
+  }
   void JHybridLiteRTLMSpec::close() {
     static const auto method = javaClassStatic()->getMethod<void()>("close");
     method(_javaPart);

package/nitrogen/generated/android/c++/JHybridLiteRTLMSpec.hpp

@@ -66,6 +66,7 @@ namespace margelo::nitro::litertlm {
     void resetConversation() override;
     bool isReady() override;
     GenerationStats getStats() override;
+    MemoryUsage getMemoryUsage() override;
     void close() override;
 
   private:

package/nitrogen/generated/android/c++/JMemoryUsage.hpp

@@ -0,0 +1,69 @@
+///
+/// JMemoryUsage.hpp
+/// This file was generated by nitrogen. DO NOT MODIFY THIS FILE.
+/// https://github.com/mrousavy/nitro
+/// Copyright © Marc Rousavy @ Margelo
+///
+
+#pragma once
+
+#include <fbjni/fbjni.h>
+#include "MemoryUsage.hpp"
+
+
+
+namespace margelo::nitro::litertlm {
+
+  using namespace facebook;
+
+  /**
+   * The C++ JNI bridge between the C++ struct "MemoryUsage" and the the Kotlin data class "MemoryUsage".
+   */
+  struct JMemoryUsage final: public jni::JavaClass<JMemoryUsage> {
+  public:
+    static auto constexpr kJavaDescriptor = "Lcom/margelo/nitro/dev/litert/litertlm/MemoryUsage;";
+
+  public:
+    /**
+     * Convert this Java/Kotlin-based struct to the C++ struct MemoryUsage by copying all values to C++.
+     */
+    [[maybe_unused]]
+    [[nodiscard]]
+    MemoryUsage toCpp() const {
+      static const auto clazz = javaClassStatic();
+      static const auto fieldNativeHeapBytes = clazz->getField<double>("nativeHeapBytes");
+      double nativeHeapBytes = this->getFieldValue(fieldNativeHeapBytes);
+      static const auto fieldResidentBytes = clazz->getField<double>("residentBytes");
+      double residentBytes = this->getFieldValue(fieldResidentBytes);
+      static const auto fieldAvailableMemoryBytes = clazz->getField<double>("availableMemoryBytes");
+      double availableMemoryBytes = this->getFieldValue(fieldAvailableMemoryBytes);
+      static const auto fieldIsLowMemory = clazz->getField<jboolean>("isLowMemory");
+      jboolean isLowMemory = this->getFieldValue(fieldIsLowMemory);
+      return MemoryUsage(
+        nativeHeapBytes,
+        residentBytes,
+        availableMemoryBytes,
+        static_cast<bool>(isLowMemory)
+      );
+    }
+
+  public:
+    /**
+     * Create a Java/Kotlin-based struct by copying all values from the given C++ struct to Java.
+     */
+    [[maybe_unused]]
+    static jni::local_ref<JMemoryUsage::javaobject> fromCpp(const MemoryUsage& value) {
+      using JSignature = JMemoryUsage(double, double, double, jboolean);
+      static const auto clazz = javaClassStatic();
+      static const auto create = clazz->getStaticMethod<JSignature>("fromCpp");
+      return create(
+        clazz,
+        value.nativeHeapBytes,
+        value.residentBytes,
+        value.availableMemoryBytes,
+        value.isLowMemory
+      );
+    }
+  };
+
+} // namespace margelo::nitro::litertlm

package/nitrogen/generated/android/kotlin/com/margelo/nitro/dev/litert/litertlm/HybridLiteRTLMSpec.kt

@@ -100,6 +100,10 @@ abstract class HybridLiteRTLMSpec: HybridObject() {
   @Keep
   abstract fun getStats(): GenerationStats
   
+  @DoNotStrip
+  @Keep
+  abstract fun getMemoryUsage(): MemoryUsage
+  
   @DoNotStrip
   @Keep
   abstract fun close(): Unit

package/nitrogen/generated/android/kotlin/com/margelo/nitro/dev/litert/litertlm/MemoryUsage.kt

@@ -0,0 +1,47 @@
+///
+/// MemoryUsage.kt
+/// This file was generated by nitrogen. DO NOT MODIFY THIS FILE.
+/// https://github.com/mrousavy/nitro
+/// Copyright © Marc Rousavy @ Margelo
+///
+
+package com.margelo.nitro.dev.litert.litertlm
+
+import androidx.annotation.Keep
+import com.facebook.proguard.annotations.DoNotStrip
+
+
+/**
+ * Represents the JavaScript object/struct "MemoryUsage".
+ */
+@DoNotStrip
+@Keep
+data class MemoryUsage(
+  @DoNotStrip
+  @Keep
+  val nativeHeapBytes: Double,
+  @DoNotStrip
+  @Keep
+  val residentBytes: Double,
+  @DoNotStrip
+  @Keep
+  val availableMemoryBytes: Double,
+  @DoNotStrip
+  @Keep
+  val isLowMemory: Boolean
+) {
+  /* primary constructor */
+
+  companion object {
+    /**
+     * Constructor called from C++
+     */
+    @DoNotStrip
+    @Keep
+    @Suppress("unused")
+    @JvmStatic
+    private fun fromCpp(nativeHeapBytes: Double, residentBytes: Double, availableMemoryBytes: Double, isLowMemory: Boolean): MemoryUsage {
+      return MemoryUsage(nativeHeapBytes, residentBytes, availableMemoryBytes, isLowMemory)
+    }
+  }
+}

package/nitrogen/generated/shared/c++/HybridLiteRTLMSpec.cpp

@@ -25,6 +25,7 @@ namespace margelo::nitro::litertlm {
       prototype.registerHybridMethod("resetConversation", &HybridLiteRTLMSpec::resetConversation);
       prototype.registerHybridMethod("isReady", &HybridLiteRTLMSpec::isReady);
       prototype.registerHybridMethod("getStats", &HybridLiteRTLMSpec::getStats);
+      prototype.registerHybridMethod("getMemoryUsage", &HybridLiteRTLMSpec::getMemoryUsage);
       prototype.registerHybridMethod("close", &HybridLiteRTLMSpec::close);
     });
   }

package/nitrogen/generated/shared/c++/HybridLiteRTLMSpec.hpp

@@ -19,6 +19,8 @@ namespace margelo::nitro::litertlm { struct LLMConfig; }
 namespace margelo::nitro::litertlm { struct Message; }
 // Forward declaration of `GenerationStats` to properly resolve imports.
 namespace margelo::nitro::litertlm { struct GenerationStats; }
+// Forward declaration of `MemoryUsage` to properly resolve imports.
+namespace margelo::nitro::litertlm { struct MemoryUsage; }
 
 #include <NitroModules/Promise.hpp>
 #include <string>
@@ -28,6 +30,7 @@ namespace margelo::nitro::litertlm { struct GenerationStats; }
 #include "Message.hpp"
 #include <vector>
 #include "GenerationStats.hpp"
+#include "MemoryUsage.hpp"
 
 namespace margelo::nitro::litertlm {
 
@@ -71,6 +74,7 @@ namespace margelo::nitro::litertlm {
       virtual void resetConversation() = 0;
       virtual bool isReady() = 0;
       virtual GenerationStats getStats() = 0;
+      virtual MemoryUsage getMemoryUsage() = 0;
       virtual void close() = 0;
 
     protected: