fix(android): support String/ByteArray onToken callbacks in JNI streaming (#376)

* fix(android): support String and ByteArray JNI stream callbacks

Handle both onToken signatures in JNI streaming callback path and avoid stale callback/context usage by waiting for completion before releasing JNI global refs.

Also add native lib version marker checks in downloadJniLibs to prevent stale JNI binary reuse across version changes.

Refs #353

* fix(android): guard JNI stream token accumulation with mutex

* fix(android): ignore onToken return when JNI exception occurs

* fix(android): add timeout for JNI stream completion wait

Replace unbounded condition variable waits with wait_for(10m) and propagate timeout as stream error to avoid indefinite JNI thread blocking.

Author: ManthanNimodiya

sdk/runanywhere-commons/src/jni/runanywhere_commons_jni.cpp

@@ -17,6 +17,7 @@
 #include <jni.h>
 
 #include <condition_variable>
+#include <chrono>
 #include <cstring>
 #include <mutex>
 #include <string>
@@ -679,6 +680,9 @@ struct LLMStreamCallbackContext {
     JavaVM* jvm = nullptr;
     jobject callback = nullptr;
     jmethodID onTokenMethod = nullptr;
+    bool onTokenExpectsBytes = true;
+    std::mutex mtx;
+    std::condition_variable cv;
     std::string accumulated_text;
     int token_count = 0;
     bool is_complete = false;
@@ -694,9 +698,12 @@ static rac_bool_t llm_stream_callback_token(const char* token, void* user_data)
 
     auto* ctx = static_cast<LLMStreamCallbackContext*>(user_data);
 
-    // Accumulate token
-    ctx->accumulated_text += token;
-    ctx->token_count++;
+    // Accumulate token (thread-safe)
+    {
+        std::lock_guard<std::mutex> lock(ctx->mtx);
+        ctx->accumulated_text += token;
+        ctx->token_count++;
+    }
 
     // Call back to Kotlin
     if (ctx->jvm && ctx->callback && ctx->onTokenMethod) {
@@ -714,21 +721,27 @@ static rac_bool_t llm_stream_callback_token(const char* token, void* user_data)
         }
 
         if (env) {
-            jsize len = static_cast<jsize>(strlen(token));
-
-            jbyteArray jToken = env->NewByteArray(len);
-            env->SetByteArrayRegion(
-                jToken,
-                0,
-                len,
-                reinterpret_cast<const jbyte*>(token)
-            );
-
-            jboolean continueGen =
-                env->CallBooleanMethod(ctx->callback, ctx->onTokenMethod, jToken);
-            env->DeleteLocalRef(jToken);
+            jboolean continueGen = JNI_TRUE;
+
+            if (ctx->onTokenExpectsBytes) {
+                jsize len = static_cast<jsize>(strlen(token));
+                jbyteArray jToken = env->NewByteArray(len);
+                env->SetByteArrayRegion(
+                    jToken,
+                    0,
+                    len,
+                    reinterpret_cast<const jbyte*>(token)
+                );
+                continueGen = env->CallBooleanMethod(ctx->callback, ctx->onTokenMethod, jToken);
+                env->DeleteLocalRef(jToken);
+            } else {
+                jstring jToken = env->NewStringUTF(token);
+                continueGen = env->CallBooleanMethod(ctx->callback, ctx->onTokenMethod, jToken);
+                env->DeleteLocalRef(jToken);
+            }
 
-            if (env->ExceptionCheck()) {
+            const bool hadException = env->ExceptionCheck();
+            if (hadException) {
                 env->ExceptionDescribe();
                 env->ExceptionClear();
             }
@@ -737,6 +750,11 @@ static rac_bool_t llm_stream_callback_token(const char* token, void* user_data)
                 ctx->jvm->DetachCurrentThread();
             }
 
+            if (hadException) {
+                // Ignore callback return value when JNI exception was thrown.
+                return RAC_TRUE;
+            }
+
             if (!continueGen) {
                 LOGi("Streaming cancelled by callback");
                 return RAC_FALSE;  // Stop streaming
@@ -752,6 +770,7 @@ static void llm_stream_callback_complete(const rac_llm_result_t* result, void* u
         return;
 
     auto* ctx = static_cast<LLMStreamCallbackContext*>(user_data);
+    std::lock_guard<std::mutex> lock(ctx->mtx);
 
     LOGi("Streaming with callback complete: %d tokens", ctx->token_count);
 
@@ -767,6 +786,7 @@ static void llm_stream_callback_complete(const rac_llm_result_t* result, void* u
     }
 
     ctx->is_complete = true;
+    ctx->cv.notify_one();
 }
 
 static void llm_stream_callback_error(rac_result_t error_code, const char* error_message,
@@ -775,6 +795,7 @@ static void llm_stream_callback_error(rac_result_t error_code, const char* error
         return;
 
     auto* ctx = static_cast<LLMStreamCallbackContext*>(user_data);
+    std::lock_guard<std::mutex> lock(ctx->mtx);
 
     LOGe("Streaming with callback error: %d - %s", error_code,
          error_message ? error_message : "Unknown");
@@ -783,6 +804,7 @@ static void llm_stream_callback_error(rac_result_t error_code, const char* error
     ctx->error_code = error_code;
     ctx->error_message = error_message ? error_message : "Unknown error";
     ctx->is_complete = true;
+    ctx->cv.notify_one();
 }
 
 JNIEXPORT jstring JNICALL
@@ -847,7 +869,12 @@ Java_com_runanywhere_sdk_native_bridge_RunAnywhereBridge_racLlmComponentGenerate
     // Wait for streaming to complete
     {
         std::unique_lock<std::mutex> lock(ctx.mtx);
-        ctx.cv.wait(lock, [&ctx] { return ctx.is_complete; });
+        constexpr auto kStreamWaitTimeout = std::chrono::minutes(10);
+        if (!ctx.cv.wait_for(lock, kStreamWaitTimeout, [&ctx] { return ctx.is_complete; })) {
+            ctx.has_error = true;
+            ctx.error_message = "Streaming timed out waiting for completion callback";
+            ctx.is_complete = true;
+        }
     }
 
     if (ctx.has_error) {
@@ -904,7 +931,13 @@ Java_com_runanywhere_sdk_native_bridge_RunAnywhereBridge_racLlmComponentGenerate
     env->GetJavaVM(&jvm);
 
     jclass callbackClass = env->GetObjectClass(tokenCallback);
+    bool onTokenExpectsBytes = true;
     jmethodID onTokenMethod = env->GetMethodID(callbackClass, "onToken", "([B)Z");
+    if (!onTokenMethod) {
+        env->ExceptionClear();
+        onTokenMethod = env->GetMethodID(callbackClass, "onToken", "(Ljava/lang/String;)Z");
+        onTokenExpectsBytes = false;
+    }
 
     if (!onTokenMethod) {
         LOGe("racLlmComponentGenerateStreamWithCallback: could not find onToken method");
@@ -948,21 +981,34 @@ Java_com_runanywhere_sdk_native_bridge_RunAnywhereBridge_racLlmComponentGenerate
     ctx.jvm = jvm;
     ctx.callback = globalCallback;
     ctx.onTokenMethod = onTokenMethod;
+    ctx.onTokenExpectsBytes = onTokenExpectsBytes;
 
     LOGi("racLlmComponentGenerateStreamWithCallback calling rac_llm_component_generate_stream...");
 
     rac_result_t status = rac_llm_component_generate_stream(
         reinterpret_cast<rac_handle_t>(handle), promptStr.c_str(), &options,
         llm_stream_callback_token, llm_stream_callback_complete, llm_stream_callback_error, &ctx);
 
-    // Clean up global ref
-    env->DeleteGlobalRef(globalCallback);
-
     if (status != RAC_SUCCESS) {
+        env->DeleteGlobalRef(globalCallback);
         LOGe("rac_llm_component_generate_stream failed with status=%d", status);
         return nullptr;
     }
 
+    // Wait until completion/error before releasing callback/context.
+    {
+        std::unique_lock<std::mutex> lock(ctx.mtx);
+        constexpr auto kStreamWaitTimeout = std::chrono::minutes(10);
+        if (!ctx.cv.wait_for(lock, kStreamWaitTimeout, [&ctx] { return ctx.is_complete; })) {
+            ctx.has_error = true;
+            ctx.error_message = "Streaming timed out waiting for completion callback";
+            ctx.is_complete = true;
+        }
+    }
+
+    // Clean up global ref after callbacks have finished.
+    env->DeleteGlobalRef(globalCallback);
+
     if (ctx.has_error) {
         LOGe("Streaming failed: %s", ctx.error_message.c_str());
         return nullptr;

sdk/runanywhere-kotlin/build.gradle.kts

@@ -399,6 +399,7 @@ tasks.register("downloadJniLibs") {
     onlyIf { !testLocal }
 
     val outputDir = file("src/androidMain/jniLibs")
+    val nativeLibVersionMarker = file("$outputDir/.native_lib_version")
     val tempDir = file("${layout.buildDirectory.get()}/jni-temp")
 
     val releaseBaseUrl = "https://github.com/RunanywhereAI/runanywhere-sdks/releases/download/v$nativeLibVersion"
@@ -424,12 +425,25 @@ tasks.register("downloadJniLibs") {
             return@doLast
         }
 
+        // Check if libs already exist (CI pre-populates build/jniLibs/).
+        // Guard against stale libs from a different native version.
         val existingLibs = outputDir.walkTopDown().filter { it.extension == "so" }.count()
-        if (existingLibs > 0) {
-            logger.lifecycle("Skipping JNI download: $existingLibs .so files already in $outputDir")
+        val existingVersion = nativeLibVersionMarker.takeIf { it.exists() }?.readText()?.trim()
+        if (existingLibs > 0 && existingVersion == nativeLibVersion) {
+            logger.lifecycle(
+                "Skipping JNI download: $existingLibs .so files already in $outputDir " +
+                    "(native version v$nativeLibVersion)",
+            )
             return@doLast
         }
+        if (existingLibs > 0 && existingVersion != nativeLibVersion) {
+            logger.lifecycle(
+                "Refreshing JNI libs: found $existingLibs existing .so files " +
+                    "with version '${existingVersion ?: "unknown"}', expected '$nativeLibVersion'",
+            )
+        }
 
+        // Clean output directories for a fresh download
         outputDir.deleteRecursively()
         tempDir.deleteRecursively()
         outputDir.mkdirs()
@@ -497,6 +511,11 @@ tasks.register("downloadJniLibs") {
         logger.lifecycle("  Output: $outputDir")
         logger.lifecycle("═══════════════════════════════════════════════════════════════")
 
+        // Record native lib version to avoid reusing stale JNI binaries.
+        nativeLibVersionMarker.parentFile.mkdirs()
+        nativeLibVersionMarker.writeText(nativeLibVersion)
+
+        // List libraries per ABI
         abiDirs.forEach { abi ->
             val libs = file("$outputDir/$abi").listFiles()?.filter { it.extension == "so" }?.map { it.name } ?: emptyList()
             logger.lifecycle("$abi (${libs.size} libs):")
@@ -662,4 +681,3 @@ tasks.withType<PublishToMavenRepository>().configureEach {
         !dominated
     }
 }
-