sdk: update JNI and session generation to return structured result

*   Introduced `TokenGenerationResult` (Kotlin) and `Generation` (C++) to return both the generated token and a completion flag from the native layer.
*   Updated `LlamaSession::generate` to return the new structured result instead of a nullable string.
*   Implemented JNI reference caching in `llama_session_jni.cpp` for the new result class to improve performance.
*   Simplified `LlamaEngine` JNI callback logic by removing unnecessary global references during synchronous model loading.
*   Updated `LlamaChatSessionImpl` and tests to handle the new generation result format.
*   Refactored JNI parameter naming for consistency (e.g., `kSessionPtr` to `jSessionPtr`).

Author: whyisitworking

sdk/src/main/cpp/jni/llama_engine_jni.cpp

@@ -7,8 +7,14 @@
 
 // ── Shared helper ─────────────────────────────────────────────────────────────
 
-static NativeEngineParams readEngineParams(JNIEnv *env, jobject kConfig) {
-    auto configReader = JniConfigReader(env, kConfig);
+namespace jni_refs {
+    constexpr auto progress_listener_method = "onProgress";
+    constexpr auto progress_listener_method_sig = "(F)Z";
+}
+
+static NativeEngineParams readEngineParams(JNIEnv *env,
+                                           jobject jConfig) {
+    auto configReader = JniConfigReader(env, jConfig);
     return NativeEngineParams{
             .model_path = configReader.getString("modelPath"),
             .threads    = configReader.getInt("threads"),
@@ -22,9 +28,10 @@ static NativeEngineParams readEngineParams(JNIEnv *env, jobject kConfig) {
 extern "C"
 JNIEXPORT jlong JNICALL
 Java_com_suhel_llamabro_sdk_internal_LlamaEngineImpl_00024Jni_create(JNIEnv *env, jclass,
-                                                                     jobject kConfig) {
+                                                                     jobject jConfig) {
     try {
-        return reinterpret_cast<jlong>(new LlamaEngine(readEngineParams(env, kConfig)));
+        auto instance = new LlamaEngine(readEngineParams(env, jConfig));
+        return reinterpret_cast<jlong>(instance);
     } catch (const LlamaException &ex) {
         throwLlamaError(env, ex);
         return 0L;
@@ -37,39 +44,27 @@ extern "C"
 JNIEXPORT jlong JNICALL
 Java_com_suhel_llamabro_sdk_internal_LlamaEngineImpl_00024Jni_createWithProgress(JNIEnv *env,
                                                                                  jclass,
-                                                                                 jobject kConfig,
-                                                                                 jobject kListener) {
-    auto config = readEngineParams(env, kConfig);
+                                                                                 jobject jConfig,
+                                                                                 jobject jListener) {
+    auto config = readEngineParams(env, jConfig);
 
     // Resolve the callback method ID once before entering the blocking load
-    jclass listenerClass = env->GetObjectClass(kListener);
-    jmethodID onProgress = env->GetMethodID(listenerClass, "onProgress", "(F)Z");
-    env->DeleteLocalRef(listenerClass);
-
-    // Create a GlobalRef to guarantee the object survives the JNI frame securely
-    auto globalListener = env->NewGlobalRef(kListener);
-
-    // Get the JavaVM so the lambda can attach/detach thread or use the current env
-    JavaVM* jvm;
-    env->GetJavaVM(&jvm);
+    auto jListenerClass = env->GetObjectClass(jListener);
+    auto jOnProgress = env->GetMethodID(jListenerClass,
+                                        jni_refs::progress_listener_method,
+                                        jni_refs::progress_listener_method_sig);
+    env->DeleteLocalRef(jListenerClass);
 
-    config.progress_callback = [jvm, globalListener, onProgress](float progress) -> bool {
-        JNIEnv *currentEnv;
-        // In this specific codebase, llama_model_load_from_file is synchronous,
-        // but getting the env from JVM is the correct modern JNI pattern.
-        auto res = jvm->GetEnv(reinterpret_cast<void**>(&currentEnv), JNI_VERSION_1_6);
-        if (res == JNI_OK) {
-            return currentEnv->CallBooleanMethod(globalListener, onProgress, static_cast<jfloat>(progress)) == JNI_TRUE;
-        }
-        return false;
+    // It is safe to pass these refs to the callback because this method is synchronous
+    config.progress_callback = [env, jListener, jOnProgress](float progress) -> bool {
+        return env->CallBooleanMethod(jListener, jOnProgress,
+                                      static_cast<jfloat>(progress)) == JNI_TRUE;
     };
 
     try {
-        auto *engine = new LlamaEngine(config);
-        env->DeleteGlobalRef(globalListener); // Safe to delete after synchronous load
-        return reinterpret_cast<jlong>(engine);
+        auto instance = new LlamaEngine(config);
+        return reinterpret_cast<jlong>(instance);
     } catch (const LlamaException &ex) {
-        env->DeleteGlobalRef(globalListener); // Clean up on error
         throwLlamaError(env, ex);
         return 0L;
     }
@@ -78,6 +73,6 @@ Java_com_suhel_llamabro_sdk_internal_LlamaEngineImpl_00024Jni_createWithProgress
 extern "C"
 JNIEXPORT void JNICALL
 Java_com_suhel_llamabro_sdk_internal_LlamaEngineImpl_00024Jni_destroy(JNIEnv *, jclass,
-                                                                      jlong ptr) {
-    delete reinterpret_cast<LlamaEngine *>(ptr);
+                                                                      jlong jEnginePtr) {
+    delete reinterpret_cast<LlamaEngine *>(jEnginePtr);
 }

sdk/src/main/cpp/jni/llama_session_jni.cpp

@@ -4,16 +4,44 @@
 #include "utils/llama_exception.h"
 #include "engine.h"
 
+namespace jni_refs {
+    constexpr auto token_generation_result_class = "com/suhel/llamabro/sdk/internal/LlamaSessionImpl$NativeTokenGenerationResult";
+    constexpr auto token_generation_result_constructor_sig = "(Ljava/lang/String;Z)V";
+}
+
+static jclass jTokenGenerationResultClass = nullptr;
+static jmethodID jTokenGenerationResultConstructor = nullptr;
+
+static void cache_refs(JNIEnv *env) {
+    auto local = env->FindClass(jni_refs::token_generation_result_class);
+
+    jTokenGenerationResultClass = reinterpret_cast<jclass>(env->NewGlobalRef(local));
+    jTokenGenerationResultConstructor = env->GetMethodID(jTokenGenerationResultClass,
+                                                         "<init>",
+                                                         jni_refs::token_generation_result_constructor_sig);
+    env->DeleteLocalRef(local);
+}
+
+JNIEXPORT jint JNICALL JNI_OnLoad(JavaVM *vm, void *) {
+    JNIEnv *env;
+    if (vm->GetEnv(reinterpret_cast<void **>(&env), JNI_VERSION_1_6) != JNI_OK) {
+        return JNI_ERR;
+    }
+
+    cache_refs(env);
+    return JNI_VERSION_1_6;
+}
+
 // ── create ────────────────────────────────────────────────────────────────────
 
 extern "C"
 JNIEXPORT jlong JNICALL
 Java_com_suhel_llamabro_sdk_internal_LlamaSessionImpl_00024Jni_create(JNIEnv *env,
                                                                       jclass,
-                                                                      jlong kEnginePtr,
-                                                                      jobject kParams) {
-    auto engine = reinterpret_cast<LlamaEngine *>(kEnginePtr);
-    auto configReader = JniConfigReader(env, kParams);
+                                                                      jlong jEnginePtr,
+                                                                      jobject jParams) {
+    auto engine = reinterpret_cast<LlamaEngine *>(jEnginePtr);
+    auto configReader = JniConfigReader(env, jParams);
 
     auto config = NativeSessionParams{
             .context_size          = configReader.getInt("contextSize"),
@@ -46,16 +74,16 @@ Java_com_suhel_llamabro_sdk_internal_LlamaSessionImpl_00024Jni_create(JNIEnv *en
 extern "C"
 JNIEXPORT void JNICALL
 Java_com_suhel_llamabro_sdk_internal_LlamaSessionImpl_00024Jni_setSystemPrompt(JNIEnv *env, jclass,
-                                                                               jlong kSessionPtr,
-                                                                               jstring kText,
-                                                                               jboolean kAddSpecial) {
-    auto session = reinterpret_cast<LlamaSession *>(kSessionPtr);
-    auto text = env->GetStringUTFChars(kText, nullptr);
-    std::string textStr(text);
-    env->ReleaseStringUTFChars(kText, text);
+                                                                               jlong jSessionPtr,
+                                                                               jstring jText,
+                                                                               jboolean jAddSpecial) {
+    auto session = reinterpret_cast<LlamaSession *>(jSessionPtr);
+    auto text = env->GetStringUTFChars(jText, nullptr);
+    auto textStr = std::string(text);
+    env->ReleaseStringUTFChars(jText, text);
 
     try {
-        session->setSystemPrompt(textStr, kAddSpecial);
+        session->setSystemPrompt(textStr, jAddSpecial);
     } catch (const LlamaException &ex) {
         throwLlamaError(env, ex);
     }
@@ -64,16 +92,16 @@ Java_com_suhel_llamabro_sdk_internal_LlamaSessionImpl_00024Jni_setSystemPrompt(J
 extern "C"
 JNIEXPORT void JNICALL
 Java_com_suhel_llamabro_sdk_internal_LlamaSessionImpl_00024Jni_injectPrompt(JNIEnv *env, jclass,
-                                                                            jlong kSessionPtr,
-                                                                            jstring kText,
-                                                                            jboolean kAddSpecial) {
-    auto session = reinterpret_cast<LlamaSession *>(kSessionPtr);
-    auto text = env->GetStringUTFChars(kText, nullptr);
-    std::string textStr(text);
-    env->ReleaseStringUTFChars(kText, text);
+                                                                            jlong jSessionPtr,
+                                                                            jstring jText,
+                                                                            jboolean jAddSpecial) {
+    auto session = reinterpret_cast<LlamaSession *>(jSessionPtr);
+    auto text = env->GetStringUTFChars(jText, nullptr);
+    auto textStr = std::string(text);
+    env->ReleaseStringUTFChars(jText, text);
 
     try {
-        session->injectPrompt(textStr, kAddSpecial);
+        session->injectPrompt(textStr, jAddSpecial);
     } catch (const LlamaException &ex) {
         throwLlamaError(env, ex);
     }
@@ -84,9 +112,11 @@ Java_com_suhel_llamabro_sdk_internal_LlamaSessionImpl_00024Jni_injectPrompt(JNIE
 extern "C"
 JNIEXPORT void JNICALL
 Java_com_suhel_llamabro_sdk_internal_LlamaSessionImpl_00024Jni_clear(JNIEnv *env, jclass,
-                                                                     jlong kSessionPtr) {
+                                                                     jlong jSessionPtr) {
+    auto session = reinterpret_cast<LlamaSession *>(jSessionPtr);
+
     try {
-        reinterpret_cast<LlamaSession *>(kSessionPtr)->clear();
+        session->clear();
     } catch (const LlamaException &ex) {
         throwLlamaError(env, ex);
     }
@@ -97,25 +127,31 @@ Java_com_suhel_llamabro_sdk_internal_LlamaSessionImpl_00024Jni_clear(JNIEnv *env
 extern "C"
 JNIEXPORT void JNICALL
 Java_com_suhel_llamabro_sdk_internal_LlamaSessionImpl_00024Jni_abort(JNIEnv *, jclass,
-                                                                     jlong kSessionPtr) {
-    reinterpret_cast<LlamaSession *>(kSessionPtr)->abort();
+                                                                     jlong jSessionPtr) {
+    auto session = reinterpret_cast<LlamaSession *>(jSessionPtr);
+    session->abort();
 }
 
 // ── generate ─────────────────────────────────────────────────────────────────
 
 extern "C"
-JNIEXPORT jstring JNICALL
+JNIEXPORT jobject JNICALL
 Java_com_suhel_llamabro_sdk_internal_LlamaSessionImpl_00024Jni_generate(JNIEnv *env, jclass,
-                                                                        jlong kSessionPtr) {
+                                                                        jlong jSessionPtr) {
+    auto session = reinterpret_cast<LlamaSession *>(jSessionPtr);
+
     try {
-        auto result = reinterpret_cast<LlamaSession *>(kSessionPtr)->generate();
+        auto gen = session->generate();
+        auto token = gen.token;
 
-        if (result.has_value()) {
-            const auto &utf16 = result.value();
-            return env->NewString(reinterpret_cast<const jchar *>(utf16.data()),
-                                  static_cast<jsize>(utf16.size()));
-        }
-        return nullptr;
+        auto jToken = token.has_value()
+                      ? env->NewString(reinterpret_cast<const jchar *>(token.value().data()),
+                                       static_cast<jsize>(token.value().size()))
+                      : nullptr;
+        auto jIsComplete = static_cast<jboolean>(gen.is_complete);
+
+        return env->NewObject(jTokenGenerationResultClass, jTokenGenerationResultConstructor,
+                              jToken, jIsComplete);
     } catch (const LlamaException &ex) {
         throwLlamaError(env, ex);
         return nullptr;
@@ -127,6 +163,7 @@ Java_com_suhel_llamabro_sdk_internal_LlamaSessionImpl_00024Jni_generate(JNIEnv *
 extern "C"
 JNIEXPORT void JNICALL
 Java_com_suhel_llamabro_sdk_internal_LlamaSessionImpl_00024Jni_destroy(JNIEnv *, jclass,
-                                                                       jlong kSessionPtr) {
-    delete reinterpret_cast<LlamaSession *>(kSessionPtr);
+                                                                       jlong jSessionPtr) {
+    auto session = reinterpret_cast<LlamaSession *>(jSessionPtr);
+    delete session;
 }

sdk/src/main/cpp/session.cpp

@@ -201,7 +201,7 @@ bool LlamaSession::is_token_buffer_valid() {
     return !token_buffer.empty() && utils::llm_is_valid_utf8(token_buffer);
 }
 
-std::u16string LlamaSession::get_token_buffer_as_u16string() {
+std::u16string LlamaSession::get_and_clear_token_buffer() {
     auto result = utils::llm_utf8_to_utf16_sanitized(token_buffer);
     token_buffer.clear();
     return result;
@@ -221,7 +221,7 @@ void LlamaSession::injectPrompt(const std::string &user_message, bool add_specia
     ingest_prompt(user_message, false, add_special);
 }
 
-std::optional<std::u16string> LlamaSession::generate() {
+Generation LlamaSession::generate() {
     auto ctx = llama_context.get();
     auto model = llama_get_model(ctx);
     auto vocab = llama_model_get_vocab(model);
@@ -247,20 +247,24 @@ std::optional<std::u16string> LlamaSession::generate() {
                 }
             }
 
-            if (is_token_buffer_valid()) {
-                return get_token_buffer_as_u16string();
-            }
-            return std::nullopt;
+            return Generation{
+                    .token = is_token_buffer_valid()
+                             ? std::make_optional(get_and_clear_token_buffer())
+                             : std::nullopt,
+                    .is_complete = true,
+            };
         }
 
         auto piece = utils::token_to_piece(vocab, new_token, true);
         token_buffer.append(piece);
 
         if (!roll_kv_cache_if_needed(1)) {
-            if (is_token_buffer_valid()) {
-                return get_token_buffer_as_u16string();
-            }
-            return std::nullopt;
+            return Generation{
+                    .token = is_token_buffer_valid()
+                             ? std::make_optional(get_and_clear_token_buffer())
+                             : std::nullopt,
+                    .is_complete = true,
+            };
         }
 
         utils::batch_clear(llama_batch);
@@ -275,7 +279,12 @@ std::optional<std::u16string> LlamaSession::generate() {
         n_past += 1;
 
         if (is_token_buffer_valid()) {
-            return get_token_buffer_as_u16string();
+            return Generation{
+                    .token = is_token_buffer_valid()
+                             ? std::make_optional(get_and_clear_token_buffer())
+                             : std::nullopt,
+                    .is_complete = false,
+            };
         }
     }
 }

sdk/src/main/cpp/session.h

@@ -33,6 +33,11 @@ struct NativeSessionParams {
     int micro_batch_size;
 };
 
+struct Generation {
+    std::optional<std::u16string> token;
+    bool is_complete;
+};
+
 #include <atomic>
 
 class LlamaSession {
@@ -62,7 +67,7 @@ class LlamaSession {
 
     bool is_token_buffer_valid();
 
-    std::u16string get_token_buffer_as_u16string();
+    std::u16string get_and_clear_token_buffer();
 
 public:
     LlamaSession(llama_model *model, int threads, const NativeSessionParams &config);
@@ -81,7 +86,7 @@ class LlamaSession {
 
     void injectPrompt(const std::string &prompt, bool add_special);
 
-    std::optional<std::u16string> generate();
+    Generation generate();
 
     void clear();
 

sdk/src/main/java/com/suhel/llamabro/sdk/LlamaSession.kt

@@ -2,6 +2,7 @@ package com.suhel.llamabro.sdk
 
 import com.suhel.llamabro.sdk.model.ResourceState
 import com.suhel.llamabro.sdk.model.ModelConfig
+import com.suhel.llamabro.sdk.model.TokenGenerationResult
 import kotlinx.coroutines.flow.Flow
 
 /**
@@ -60,7 +61,7 @@ interface LlamaSession : AutoCloseable {
      *         End-of-Generation (EOG) token.
      * @throws LlamaError.DecodeFailed if the native sampling loop fails.
      */
-    suspend fun generate(): String?
+    suspend fun generate(): TokenGenerationResult
 
     /**
      * Clears the conversation history from the KV cache.