From 88fdd6d8cdb07181c4a2d3ea7adad05b0648fccc Mon Sep 17 00:00:00 2001 From: vaiju1981 Date: Fri, 10 Jul 2026 14:17:49 -0700 Subject: [PATCH] Fix concurrency, exception-safety, and correctness bugs in JNI layer and Java API JNI / native fixes: - Guard rerank / embedding / streaming to_json() calls against uncaught C++ exceptions crossing the JNI boundary, which would otherwise abort the JVM. - Hold a jllama_context user reference for the attached NativeServer worker so closing the model cannot free the context out from under the running HTTP frontend (use-after-free). - Fix close()-during-inference hang: add a closing flag that unblocks readers parked in next()/wait_for_all(), and decrement the user count under the context mutex so the teardown wait cannot free the mutex/condvar before the releaser notifies (latent use-after-free in the ref-count itself). - Fix NativeServer jstring_to_utf8 calling getBytes with a Charset argument while resolving the String overload (deterministic break on every server start); cache the JNI IDs thread-safely. - Initialize parse_string_array slots to nullptr so an OOM path cannot free uninitialized memory. - Add a pending-exception check in parse_jstring. Java fixes: - Widen token counts to long in the protocol translators and ChatResponseParser.extractUsageField to avoid truncation above ~2.1B tokens. - Honor a pre-cancelled CancellationToken in complete() (stop resetting it at entry); fix completeAsJson javadoc (the params object is not mutated). Tests and docs: - Extract the OuteTTS codec filter into a testable helper; add rerank bounds and codec-filter unit tests, and a model-gated close()-during-inference test. - Remove the dead BUILD_SHARED_LIBS ON line in CMakeLists.txt. - Update plan and TODO docs to reflect the real teardown design. --- TODO.md | 12 +- llama/CMakeLists.txt | 1 - llama/src/main/cpp/jllama.cpp | 238 ++++++++++++++---- llama/src/main/cpp/jni_helpers.hpp | 68 +++++ llama/src/main/cpp/json_helpers.hpp | 31 ++- llama/src/main/cpp/native_server.cpp | 86 ++++++- llama/src/main/cpp/tts_engine.cpp | 16 +- llama/src/main/cpp/tts_engine.h | 21 ++ .../java/net/ladenthin/llama/LlamaModel.java | 9 +- .../llama/json/ChatResponseParser.java | 6 +- .../llama/server/AnthropicApiSupport.java | 12 +- .../server/AnthropicStreamTranslator.java | 14 +- .../llama/server/OllamaApiSupport.java | 4 +- .../llama/server/ResponsesApiSupport.java | 6 +- .../server/ResponsesStreamTranslator.java | 12 +- llama/src/test/cpp/test_json_helpers.cpp | 38 +++ llama/src/test/cpp/test_tts_wav.cpp | 41 +++ .../net/ladenthin/llama/LlamaModelTest.java | 54 +++- 18 files changed, 552 insertions(+), 117 deletions(-) diff --git a/TODO.md b/TODO.md index 4e23bde7..b45414ba 100644 --- a/TODO.md +++ b/TODO.md @@ -149,10 +149,14 @@ Feel free to contribute fixes — PRs welcome. `synchronized` guard (or `compareAndSet`) on the Java side so `close()` is idempotent and the native pointer is nulled before the second caller can reach it. -- **`ServerMetrics.getCumulativeTimings()` truncates cumulative token totals to `int`.** The - cumulative token counters are stored as `long` in the JSON but cast to `int` when - constructing `ServerMetrics`, silently truncating values above `Integer.MAX_VALUE` - (~2.1 billion tokens). Fix: widen the field and constructor parameter to `long`. +- **`ServerMetrics.getCumulativeTimings()` truncation — STALE / REFUTED.** This claim was + verified false: `getCumulativeTimings` reads `n_prompt_tokens_processed_total` / + `n_tokens_predicted_total` via `asLong(0L)` into `long` fields and constructs a `Timings` + with `long promptN` / `predictedN` (`ServerMetrics.java` / `Timings.java`) — no int cast + occurs. The real token-count truncation (now fixed) lived in the protocol shims + (`ResponsesStreamTranslator`, `AnthropicStreamTranslator`, `ResponsesApiSupport`, + `AnthropicApiSupport`, `OllamaApiSupport`, `ChatResponseParser.extractUsageField`), which used + `.asInt(0)`; those now use `.asLong(0L)` so sessions above ~2.1B tokens no longer overflow. - **Unbounded request-body read → OOM DoS.** The HTTP handler reads the entire request body into a `String`/`byte[]` before parsing it, with no size cap. A client that streams a diff --git a/llama/CMakeLists.txt b/llama/CMakeLists.txt index b351fefa..d6734299 100644 --- a/llama/CMakeLists.txt +++ b/llama/CMakeLists.txt @@ -21,7 +21,6 @@ endif() include(FetchContent) -set(BUILD_SHARED_LIBS ON) set(CMAKE_POSITION_INDEPENDENT_CODE ON) set(BUILD_SHARED_LIBS OFF) # Android NDK only declares posix_spawn_file_actions_t as a type alias but diff --git a/llama/src/main/cpp/jllama.cpp b/llama/src/main/cpp/jllama.cpp index ffd3c846..997e6a8a 100644 --- a/llama/src/main/cpp/jllama.cpp +++ b/llama/src/main/cpp/jllama.cpp @@ -29,6 +29,11 @@ #include #include +// Guards the Java -> native jllama_context handle so acquire_jllama_context_impl (every entry +// point) and delete() (close) cannot race on the same field. Declared extern in jni_helpers.hpp. +// Defined here at file scope (external linkage) so the header's extern reference resolves. +std::mutex g_ctx_mutex; + // We store some references to Java classes and their fields/methods here to speed up things for later and to fail // early on if anything can't be found. This happens when the JVM loads the shared library (see `JNI_OnLoad`). // The references remain valid throughout the whole life of the shared library, on `JNI_OnUnload` they are released. @@ -203,6 +208,12 @@ static void throw_invalid_request(JNIEnv *env, const std::exception &e) { env->ThrowNew(c_llama_error, get_result_error_message(br.error).c_str()); return false; } + if (br.is_terminated) { + // wait_for_all was cut short because close() set jctx->closing: the results vector is + // only partially filled (null entries), so don't dereference it - unwind cleanly instead. + env->ThrowNew(c_llama_error, "inference interrupted by close()"); + return false; + } return true; } @@ -286,7 +297,7 @@ static void populate_completion_task(server_task &task, jllama_context *jctx, in } task.params.res_type = res_type; rd.post_task(std::move(task)); - auto br = rd.wait_for_all([] { return false; }); + auto br = rd.wait_for_all([jctx] { return jctx->closing.load(); }); if (!batch_ok_or_throw(env, br)) return nullptr; return results_to_jstring(env, br.results); @@ -296,8 +307,23 @@ static void populate_completion_task(server_task &task, jllama_context *jctx, in * Convert a Java string to a std::string */ std::string parse_jstring(JNIEnv *env, jstring java_string) { + // A null receiver would make CallObjectMethod raise a pending JNI exception; bail out + // cleanly so callers (e.g. parse_json_params) can surface their own error instead of + // invoking ThrowNew under an already-pending exception (L6). + if (java_string == nullptr) { + return std::string(); + } + auto *const string_bytes = (jbyteArray)env->CallObjectMethod(java_string, m_get_bytes, o_utf_8); + // CallObjectMethod may have raised a pending exception (e.g. OOM) leaving string_bytes + // null; dereferencing it via GetArrayLength would be undefined behaviour. Bail out so the + // caller (parse_json_params) surfaces a parse error as a LlamaException instead of crashing + // (MED #5). + if (env->ExceptionCheck() || string_bytes == nullptr) { + return std::string(); + } + auto length = (size_t)env->GetArrayLength(string_bytes); jbyte *byte_elements = env->GetByteArrayElements(string_bytes, nullptr); @@ -385,21 +411,45 @@ char **parse_string_array(JNIEnv *env, const jobjectArray string_array, const js } for (jsize i = 0; i < length; i++) { + // Default to a null slot so an OOM/early-exit path below (e.g. GetByteArrayElements + // returning nullptr) cannot leave result[i] as uninitialized malloc garbage that + // free_string_array would later free() (MED #3). + result[i] = nullptr; auto *const javaString = static_cast(env->GetObjectArrayElement(string_array, i)); // A null array element (legal from a Java String[]) must not reach GetStringUTFChars. if (javaString == nullptr) { result[i] = strdup(""); continue; } - // GetStringUTFChars returns nullptr on allocation failure; never strdup(nullptr). - const char *cString = env->GetStringUTFChars(javaString, nullptr); - result[i] = strdup(cString != nullptr ? cString : ""); - if (cString != nullptr) { - env->ReleaseStringUTFChars(javaString, cString); + // Use the standard-UTF-8 byte path (String.getBytes("UTF-8")), NOT GetStringUTFChars + // which yields Modified-UTF8 (CESU-8) and would mangle non-BMP chars (emoji, many CJK + // extensions) in model paths and rerank documents (H1). + jbyteArray bytes = (jbyteArray)env->CallObjectMethod(javaString, m_get_bytes, o_utf_8); + if (bytes == nullptr) { + result[i] = strdup(""); + } else { + const jsize blen = env->GetArrayLength(bytes); + jbyte *bytesElems = env->GetByteArrayElements(bytes, nullptr); + if (bytesElems != nullptr && blen >= 0) { + result[i] = static_cast(malloc(static_cast(blen) + 1)); + if (result[i] != nullptr) { + memcpy(result[i], bytesElems, static_cast(blen)); + result[i][blen] = '\0'; + } + } + if (result[i] == nullptr) { + result[i] = strdup(""); + } + if (bytesElems != nullptr) { + env->ReleaseByteArrayElements(bytes, bytesElems, JNI_ABORT); + } } // Each GetObjectArrayElement returns a fresh local ref; release it so a large // array (e.g. a rerank with many documents) cannot overflow the local-ref table. env->DeleteLocalRef(javaString); + if (bytes != nullptr) { + env->DeleteLocalRef(bytes); + } } return result; @@ -470,18 +520,30 @@ JNIEnv *get_jni_env_or_null() noexcept { bool log_json; std::function log_callback; +// Guards the logger globals below so concurrent setLogger calls (and the +// trampoline reading them from arbitrary worker threads) cannot race (L4). +static std::mutex g_log_mutex; /** * Invoke the log callback if there is any. When JSON mode is enabled, * the message is formatted as a JSON object before forwarding. */ void log_callback_trampoline(ggml_log_level level, const char *text, void *user_data) { - if (log_callback != nullptr) { - if (log_json) { + // Copy the (small) callback state under the lock so we don't hold it across the + // user-provided callback, which may itself do arbitrary work. + std::function cb; + bool json_mode; + { + std::lock_guard lk(g_log_mutex); + cb = log_callback; + json_mode = log_json; + } + if (cb != nullptr) { + if (json_mode) { std::string json_text = format_log_as_json(level, text, std::time(nullptr)); - log_callback(level, json_text.c_str(), user_data); + cb(level, json_text.c_str(), user_data); } else { - log_callback(level, text, user_data); + cb(level, text, user_data); } } } @@ -491,11 +553,12 @@ void log_callback_trampoline(ggml_log_level level, const char *text, void *user_ // `jctx` and `ctx_server` in the caller's scope; returns the given sentinel // (omit for void functions) if the model is not loaded. #define REQUIRE_SERVER_CONTEXT(...) \ - auto *jctx = get_jllama_context(env, obj); \ - if (!jctx) { \ + jllama_context_guard _jctx_guard{acquire_jllama_context_impl(env, obj, f_model_pointer)}; \ + if (!_jctx_guard.ptr) { \ env->ThrowNew(c_llama_error, "Model is not loaded"); \ return __VA_ARGS__; \ } \ + auto *jctx = _jctx_guard.ptr; \ server_context *ctx_server = &jctx->server /** @@ -816,11 +879,17 @@ JNIEXPORT jstring JNICALL Java_net_ladenthin_llama_LlamaModel_getModelMetaJson(J return json_to_jstring(env, meta); } auto m = ctx_server->get_meta(); - // Read general.architecture from GGUF metadata via the llama C API. - char arch_buf[128] = {}; + // Read general.architecture from GGUF metadata via the llama C API. Size the buffer + // dynamically: llama_model_meta_val_str returns the required length when given a null/0 + // buffer, so a long architecture name is never silently truncated (L3). + std::string arch; const llama_model *mdl = llama_get_model(ctx_server->get_llama_context()); if (mdl) { - llama_model_meta_val_str(mdl, "general.architecture", arch_buf, sizeof(arch_buf)); + const int need = llama_model_meta_val_str(mdl, "general.architecture", nullptr, 0); + if (need > 0) { + arch.resize(static_cast(need)); + llama_model_meta_val_str(mdl, "general.architecture", arch.data(), arch.size() + 1); + } } json j = { {"vocab_type", m.model_vocab_type}, @@ -831,7 +900,7 @@ JNIEXPORT jstring JNICALL Java_net_ladenthin_llama_LlamaModel_getModelMetaJson(J {"size", m.model_size}, {"modalities", {{"vision", m.has_inp_image}, {"audio", m.has_inp_audio}}}, {"name", m.model_name}, - {"architecture", std::string(arch_buf)}, + {"architecture", arch}, {"ftype", m.model_ftype}, }; // Resolved default chat template (Jinja); empty when the model ships none. @@ -897,24 +966,30 @@ JNIEXPORT jstring JNICALL Java_net_ladenthin_llama_LlamaModel_receiveCompletionJ // a nullptr sentinel meaning "HTTP-headers-only, no body". Loop past // those so the Java iterator only ever sees real events. json response; - while (true) { - server_task_result_ptr result = rd->next([] { return false; }); + try { + while (true) { + server_task_result_ptr result = rd->next([jctx] { return jctx->closing.load(); }); - if (!result_ok_or_throw(env, result)) { - erase_reader(jctx, id_task); - return nullptr; - } + if (!result_ok_or_throw(env, result)) { + erase_reader(jctx, id_task); + return nullptr; + } - response = result->to_json(); - if (response.is_null()) { - continue; - } - response["stop"] = result->is_stop(); + response = result->to_json(); + if (response.is_null()) { + continue; + } + response["stop"] = result->is_stop(); - if (result->is_stop()) { - erase_reader(jctx, id_task); + if (result->is_stop()) { + erase_reader(jctx, id_task); + } + break; } - break; + } catch (const std::exception &e) { + // A throwing to_json() must surface as a LlamaException, not abort the JVM. + env->ThrowNew(c_llama_error, e.what()); + return nullptr; } return json_to_jstring(env, response); @@ -945,24 +1020,29 @@ JNIEXPORT jstring JNICALL Java_net_ladenthin_llama_LlamaModel_receiveChatComplet json payload; bool stop = false; - while (true) { - server_task_result_ptr result = rd->next([] { return false; }); + try { + while (true) { + server_task_result_ptr result = rd->next([jctx] { return jctx->closing.load(); }); - if (!result_ok_or_throw(env, result)) { - erase_reader(jctx, id_task); - return nullptr; - } + if (!result_ok_or_throw(env, result)) { + erase_reader(jctx, id_task); + return nullptr; + } - json chunk = result->to_json(); - if (chunk.is_null()) { - continue; - } - payload = std::move(chunk); - stop = result->is_stop(); - if (stop) { - erase_reader(jctx, id_task); + json chunk = result->to_json(); + if (chunk.is_null()) { + continue; + } + payload = std::move(chunk); + stop = result->is_stop(); + if (stop) { + erase_reader(jctx, id_task); + } + break; } - break; + } catch (const std::exception &e) { + env->ThrowNew(c_llama_error, e.what()); + return nullptr; } return json_to_jstring(env, wrap_stream_chunk(std::move(payload), stop)); @@ -992,10 +1072,14 @@ JNIEXPORT jfloatArray JNICALL Java_net_ladenthin_llama_LlamaModel_embed(JNIEnv * task.index = 0; rd.post_task(std::move(task)); - auto br = rd.wait_for_all([] { return false; }); + auto br = rd.wait_for_all([jctx] { return jctx->closing.load(); }); if (!batch_ok_or_throw(env, br)) return nullptr; + if (br.results.empty()) { + env->ThrowNew(c_llama_error, "embedding result is empty"); + return nullptr; + } auto *embd_result = dynamic_cast(br.results[0].get()); if (!embd_result || embd_result->embedding.empty() || embd_result->embedding[0].empty()) { env->ThrowNew(c_llama_error, "embedding result is empty"); @@ -1039,10 +1123,18 @@ JNIEXPORT jstring JNICALL Java_net_ladenthin_llama_LlamaModel_handleRerank(JNIEn } rd.post_tasks(std::move(tasks)); - auto br = rd.wait_for_all([] { return false; }); + auto br = rd.wait_for_all([jctx] { return jctx->closing.load(); }); if (!batch_ok_or_throw(env, br)) return nullptr; - return json_to_jstring(env, rerank_results_to_json(br.results, document_vector)); + // rerank_results_to_json throws std::invalid_argument on a malformed/out-of-range + // result index (M2); unwrap it into a LlamaException instead of letting it cross + // the JNI boundary (undefined behaviour / JVM abort). + try { + return json_to_jstring(env, rerank_results_to_json(br.results, document_vector)); + } catch (const std::exception &e) { + env->ThrowNew(c_llama_error, e.what()); + return nullptr; + } } JNIEXPORT jstring JNICALL Java_net_ladenthin_llama_LlamaModel_applyTemplate(JNIEnv *env, jobject obj, jstring jparams) { @@ -1157,9 +1249,21 @@ JNIEXPORT void JNICALL Java_net_ladenthin_llama_LlamaModel_delete(JNIEnv *env, j if (!jctx) return; - env->SetLongField(obj, f_model_pointer, 0); + // Null the Java handle under g_ctx_mutex so no NEW entry point can acquire this context + // (acquire_jllama_context_impl will see 0 and return null). In-flight calls that already + // hold a user reference keep running until they release on their own scope exit. + { + std::lock_guard lk(g_ctx_mutex); + env->SetLongField(obj, f_model_pointer, 0); + } if (!jctx->vocab_only) { + // Signal teardown to any in-flight reader BEFORE stopping the worker. The streaming / + // blocking should_stop lambdas observe this and make next()/wait_for_all() return + // within one poll (~1s), so the in-flight JNI call unwinds and releases its user + // reference — otherwise the reader would poll forever (worker only stops the task + // queue, not the results queue) and close() would hang (M4). + jctx->closing.store(true, std::memory_order_release); // Cancel any pending streaming readers before stopping the server. { std::lock_guard lk(jctx->readers_mutex); @@ -1182,6 +1286,16 @@ JNIEXPORT void JNICALL Java_net_ladenthin_llama_LlamaModel_delete(JNIEnv *env, j if (jctx->vocab_only_model) { llama_model_free(jctx->vocab_only_model); } + + // Wait for any in-flight JNI call (an entry point still inside its jllama_context_guard) + // to release its user reference before we free the context. The closing flag set above makes + // those calls unwind (rd->next()/wait_for_all() return early), so this wait is bounded to + // ~1s. Without it a concurrent complete()/receiveCompletionJson() could still be + // dereferencing jctx when delete runs (use-after-free, H2 / M4). + { + std::unique_lock lk(jctx->m); + jctx->cv.wait(lk, [&] { return jctx->users.load() == 0; }); + } delete jctx; } @@ -1192,8 +1306,14 @@ JNIEXPORT void JNICALL Java_net_ladenthin_llama_LlamaModel_cancelCompletion(JNIE JNIEXPORT void JNICALL Java_net_ladenthin_llama_LlamaModel_setLogger(JNIEnv *env, jclass clazz, jobject log_format, jobject jcallback) { + // Serialize the whole swap under the logger mutex (L4): first clear the live callback so + // no in-flight trampoline can copy a lambda that still references the global ref we are + // about to delete, then delete the old ref, then install the new one. + std::lock_guard lk(g_log_mutex); + log_callback = nullptr; if (o_log_callback != nullptr) { env->DeleteGlobalRef(o_log_callback); + o_log_callback = nullptr; } log_json = env->IsSameObject(log_format, o_log_format_json); @@ -1203,7 +1323,10 @@ JNIEXPORT void JNICALL Java_net_ladenthin_llama_LlamaModel_setLogger(JNIEnv *env llama_log_set(nullptr, nullptr); } else { o_log_callback = env->NewGlobalRef(jcallback); - log_callback = [](enum ggml_log_level level, const char *text, void *user_data) noexcept { + // Capture copies of the global ref and method id so the callback never dereferences the + // logger globals at call time (those may be swapped by a concurrent setLogger). + jobject cb_ref = o_log_callback; + log_callback = [cb_ref](enum ggml_log_level level, const char *text, void *user_data) noexcept { // Logging can fire from internal native threads with no JNIEnv; skip rather than // throw (an exception here would unwind through llama.cpp's C frames). JNIEnv *env = get_jni_env_or_null(); @@ -1218,7 +1341,7 @@ JNIEXPORT void JNICALL Java_net_ladenthin_llama_LlamaModel_setLogger(JNIEnv *env return; } jobject log_level = log_level_to_jobject(level); - env->CallVoidMethod(o_log_callback, m_biconsumer_accept, log_level, message); + env->CallVoidMethod(cb_ref, m_biconsumer_accept, log_level, message); env->DeleteLocalRef(message); }; // Always set the trampoline — it handles JSON formatting internally @@ -1373,7 +1496,7 @@ JNIEXPORT jstring JNICALL Java_net_ladenthin_llama_LlamaModel_handleEmbeddings(J } rd.post_tasks(std::move(tasks)); - auto br = rd.wait_for_all([] { return false; }); + auto br = rd.wait_for_all([jctx] { return jctx->closing.load(); }); if (!batch_ok_or_throw(env, br)) return nullptr; @@ -1385,7 +1508,12 @@ JNIEXPORT jstring JNICALL Java_net_ladenthin_llama_LlamaModel_handleEmbeddings(J ? format_embeddings_response_oaicompat( body, json_value(body, "model", std::string(DEFAULT_OAICOMPAT_MODEL)), responses, use_base64) : responses; - return json_to_jstring(env, out); + try { + return json_to_jstring(env, out); + } catch (const std::exception &e) { + env->ThrowNew(c_llama_error, e.what()); + return nullptr; + } } JNIEXPORT jstring JNICALL Java_net_ladenthin_llama_LlamaModel_handleTokenize(JNIEnv *env, jobject obj, jstring jcontent, @@ -1525,7 +1653,7 @@ JNIEXPORT jboolean JNICALL Java_net_ladenthin_llama_LlamaModel_configureParallel slot_sim_opt = parse_slot_prompt_similarity(config); n_threads_opt = parse_positive_int_config(config, "n_threads"); n_threads_batch_opt = parse_positive_int_config(config, "n_threads_batch"); - } catch (const std::invalid_argument &e) { + } catch (const std::exception &e) { env->ThrowNew(c_llama_error, e.what()); return JNI_FALSE; } @@ -1555,7 +1683,7 @@ JNIEXPORT jboolean JNICALL Java_net_ladenthin_llama_LlamaModel_configureParallel // contract). Live mutation uses server_context::set_slot_prompt_similarity(), // added upstream by https://github.com/ggml-org/llama.cpp/pull/22393 and carried // in this repo as patches/0003-pr22393-... until it merges upstream (the pinned - // llama.cpp is now b9829, which the patch applies against). not thread-safe per + // llama.cpp is now b9941, which the patch applies against). not thread-safe per // the upstream contract — main-thread only, which this JNI call is. if (slot_sim_opt.has_value()) { ctx_server->set_slot_prompt_similarity(*slot_sim_opt); diff --git a/llama/src/main/cpp/jni_helpers.hpp b/llama/src/main/cpp/jni_helpers.hpp index 5118d69f..c4f5c9bf 100644 --- a/llama/src/main/cpp/jni_helpers.hpp +++ b/llama/src/main/cpp/jni_helpers.hpp @@ -29,6 +29,7 @@ #include "nlohmann/json.hpp" #include +#include #include #include #include @@ -73,6 +74,24 @@ struct jllama_context { // that removes the map entry must not free a reader still in use — the receiver's copy // keeps it alive until next() returns. std::map> readers; + + // In-flight JNI call reference count, for safe teardown. close() waits for this to reach 0 + // before deleting the context, so a concurrent inference/receive call can never dereference + // a freed jllama_context (fixes the close()-vs-inference use-after-free, H2, and the + // close()-during-streaming hang, M4). acquire_jllama_context_impl increments it; the + // jllama_context_guard (set up by REQUIRE_SERVER_CONTEXT) decrements it on scope exit. + std::atomic users{0}; + // Set by delete() (teardown) before terminating the worker, so any in-flight + // reader parked in next()/wait_for_all() observes a stop signal and unwinds + // instead of polling forever (which would otherwise keep `users` > 0 and hang + // close()). The streaming/blocking should_stop lambdas read this. + std::atomic closing{false}; + // Serializes the "last user released" notification (paired with cv). The users + // decrement in release_jllama_context_impl happens UNDER this lock so delete()'s + // cv.wait cannot observe users==0 and free m/cv before the releaser finishes + // touching them (otherwise a use-after-free on the condvar/mutex themselves). + std::mutex m; + std::condition_variable cv; }; // Removes the streaming reader entry for `id_task` under the readers_mutex. @@ -112,6 +131,55 @@ inline void erase_reader(jllama_context *jctx, int id_task) { return reinterpret_cast(handle); // NOLINT(*-no-int-to-ptr) } +// --------------------------------------------------------------------------- +// acquire / release_jllama_context_impl + jllama_context_guard +// +// Reference-count in-flight JNI calls against the jllama_context so that +// close() can wait for all of them to drain before deleting the context. Every +// entry point that dereferences jctx goes through REQUIRE_SERVER_CONTEXT, which +// constructs a jllama_context_guard; the guard's destructor calls +// release_jllama_context_impl (even on early return), decrementing users. +// +// close() (delete()) first nulls the Java handle under g_ctx_mutex so no NEW +// acquire can succeed, terminates+joins the worker, then waits on jctx->cv for +// users to reach 0 — guaranteeing no call is mid-use when the context is freed +// (fixes H2 use-after-free and M4 close()-during-streaming hang). g_ctx_mutex is +// defined in jllama.cpp and declared extern here. +// --------------------------------------------------------------------------- +extern std::mutex g_ctx_mutex; + +[[nodiscard]] inline jllama_context *acquire_jllama_context_impl(JNIEnv *env, jobject obj, jfieldID field_id) { + std::lock_guard lk(g_ctx_mutex); + const jlong handle = env->GetLongField(obj, field_id); + if (handle == 0) { + return nullptr; + } + auto *jctx = reinterpret_cast(handle); // NOLINT(*-no-int-to-ptr) + jctx->users.fetch_add(1, std::memory_order_relaxed); + return jctx; +} + +inline void release_jllama_context_impl(jllama_context *jctx) { + if (jctx == nullptr) { + return; + } + // Decrement under jctx->m so delete()'s cv.wait (which holds jctx->m while checking the + // predicate) cannot observe users==0 and free m/cv before this releaser has finished + // notifying — that would be a use-after-free on the condvar/mutex themselves (H2/H3). + std::lock_guard lk(jctx->m); + // fetch_sub returns the previous value; when it was 1 we just dropped the last reference. + if (jctx->users.fetch_sub(1, std::memory_order_acq_rel) == 1) { + jctx->cv.notify_all(); + } +} + +// RAII guard: releases the jllama_context user reference at scope exit, including +// on early return. REQUIRE_SERVER_CONTEXT uses this so every entry point is covered. +struct jllama_context_guard { + jllama_context *ptr; + ~jllama_context_guard() { release_jllama_context_impl(ptr); } +}; + // --------------------------------------------------------------------------- // require_json_field_impl // diff --git a/llama/src/main/cpp/json_helpers.hpp b/llama/src/main/cpp/json_helpers.hpp index 6d460710..45e05927 100644 --- a/llama/src/main/cpp/json_helpers.hpp +++ b/llama/src/main/cpp/json_helpers.hpp @@ -36,6 +36,7 @@ #include "nlohmann/json.hpp" +#include #include #include #include @@ -90,7 +91,15 @@ json arr = json::array(); for (const auto &result : results) { const auto out = result->to_json(); - int index = out["index"].get(); + // Defensive: a malformed/absent "index" or an out-of-range value would otherwise + // throw json::type_error or index documents[] out of bounds (M2). + if (!out.contains("index")) { + throw std::invalid_argument("rerank result is missing the 'index' field"); + } + const int index = out["index"].get(); + if (index < 0 || static_cast(index) >= documents.size()) { + throw std::invalid_argument("rerank result index " + std::to_string(index) + " out of range"); + } float score = out["score"].get(); arr.push_back({{"document", documents[index]}, {"index", index}, {"score", score}}); } @@ -168,11 +177,13 @@ if (!config.contains("slot_prompt_similarity")) { return std::nullopt; } - const float v = config["slot_prompt_similarity"].get(); - if (v < 0.0f || v > 1.0f) { + // Coerce via double so an integer-valued config (e.g. 0 or 1) is accepted rather than + // throwing json::type_error; the range check below preserves the [0.0, 1.0] contract (L5). + const double v = config["slot_prompt_similarity"].get(); + if (v < 0.0 || v > 1.0) { throw std::invalid_argument("slot_prompt_similarity must be between 0.0 and 1.0"); } - return v; + return static_cast(v); } // --------------------------------------------------------------------------- @@ -188,11 +199,15 @@ if (!config.contains(key)) { return std::nullopt; } - const int v = config[key].get(); - if (v <= 0) { - throw std::invalid_argument(std::string(key) + " must be greater than 0"); + // Coerce via double so an integer-valued config is accepted rather than throwing + // json::type_error; require a positive integer per the documented contract (L5). + const double raw = config[key].get(); + // Reject non-positive, non-integer, or > INT_MAX values: a whole-number JSON value above + // INT_MAX would overflow static_cast (UB). + if (raw <= 0.0 || raw != std::floor(raw) || raw > 2147483647.0) { + throw std::invalid_argument(std::string(key) + " must be a positive integer"); } - return v; + return static_cast(raw); } // --------------------------------------------------------------------------- diff --git a/llama/src/main/cpp/native_server.cpp b/llama/src/main/cpp/native_server.cpp index 1234a3ff..7f79166f 100644 --- a/llama/src/main/cpp/native_server.cpp +++ b/llama/src/main/cpp/native_server.cpp @@ -46,6 +46,9 @@ struct native_server { int exit_code = -1; }; +// Standard-UTF8 extraction used for argv/commands (declared here; defined below). +static std::string jstring_to_utf8(JNIEnv *env, jstring js); + // Copies the forwarded Java argv into srv->args/argv with a synthetic argv[0]. The argv pointers // reference the std::string storage in `args`, which is filled once (with reserve) and never // mutated afterwards, so the pointers stay valid for the worker's lifetime. @@ -56,11 +59,7 @@ void fill_native_server_args(JNIEnv *env, jobjectArray jargs, native_server *srv for (jsize i = 0; i < n; ++i) { auto js = static_cast(env->GetObjectArrayElement(jargs, i)); if (js != nullptr) { - const char *chars = env->GetStringUTFChars(js, nullptr); - srv->args.emplace_back(chars != nullptr ? chars : ""); - if (chars != nullptr) { - env->ReleaseStringUTFChars(js, chars); - } + srv->args.emplace_back(jstring_to_utf8(env, js)); env->DeleteLocalRef(js); } else { srv->args.emplace_back(""); @@ -82,6 +81,62 @@ void throw_llama_exception(JNIEnv *env, const char *message) { } } +// Standard-UTF8 extraction (mirrors jllama.cpp parse_jstring): GetStringUTFChars yields +// Modified-UTF8 (CESU-8), which mangles non-BMP chars (emoji, CJK extensions) in argv such as +// model paths. We go through String.getBytes("UTF-8") instead (H1). +std::string jstring_to_utf8(JNIEnv *env, jstring js) { + if (js == nullptr) { + return ""; + } + // Resolve and cache the JNI class/method/field IDs once (server-startup is the only caller, + // but it can run many times — e.g. once per forwarded argv element). FindClass returns a + // frame-local ref, so promote the reused handles to global refs for caching (MED #6). + static jclass str_cls = nullptr; + static jmethodID get_bytes = nullptr; + static jobject utf8 = nullptr; + static std::once_flag init_flag; + std::call_once(init_flag, [&]() { + jclass local = env->FindClass("java/lang/String"); + if (local == nullptr) { + return; + } + jmethodID gb = env->GetMethodID(local, "getBytes", "(Ljava/nio/charset/Charset;)[B"); + jclass charset_cls = env->FindClass("java/nio/charset/StandardCharsets"); + jfieldID utf8_f = + charset_cls ? env->GetStaticFieldID(charset_cls, "UTF_8", "Ljava/nio/charset/Charset;") : nullptr; + if (gb == nullptr || charset_cls == nullptr || utf8_f == nullptr) { + if (charset_cls != nullptr) { + env->DeleteLocalRef(charset_cls); + } + env->DeleteLocalRef(local); + return; + } + jobject u = env->GetStaticObjectField(charset_cls, utf8_f); + env->DeleteLocalRef(charset_cls); + str_cls = static_cast(env->NewGlobalRef(local)); + get_bytes = gb; + utf8 = env->NewGlobalRef(u); + env->DeleteLocalRef(local); + env->DeleteLocalRef(u); + }); + if (str_cls == nullptr || get_bytes == nullptr || utf8 == nullptr) { + return ""; + } + + jbyteArray bytes = (jbyteArray)env->CallObjectMethod(js, get_bytes, utf8); + std::string out; + if (bytes != nullptr) { + const jsize blen = env->GetArrayLength(bytes); + jbyte *elems = env->GetByteArrayElements(bytes, nullptr); + if (elems != nullptr) { + out.assign(reinterpret_cast(elems), static_cast(blen)); + env->ReleaseByteArrayElements(bytes, elems, JNI_ABORT); + } + env->DeleteLocalRef(bytes); + } + return out; +} + } // namespace extern "C" { @@ -151,6 +206,13 @@ JNIEXPORT jlong JNICALL Java_net_ladenthin_llama_server_NativeServer_startAttach } auto *jctx = reinterpret_cast(model_handle); // NOLINT(*-no-int-to-ptr) + // Contract (M6): llama_server_attach drives the HTTP frontend on a NEW worker thread but + // must NOT start a second task-processing loop on the shared server_context — the model's + // own worker (spawned in load_model_impl) already runs start_loop(). Likewise, + // stopNativeServer's llama_server_request_shutdown() must terminate only the HTTP frontend, + // not the attached model's worker, so the LlamaModel remains usable after the server closes + // (documented order: close server first, then model). If a future upstream change breaks + // either assumption, this attach path needs revisiting. auto *srv = new native_server(); fill_native_server_args(env, jargs, srv); @@ -159,9 +221,15 @@ JNIEXPORT jlong JNICALL Java_net_ladenthin_llama_server_NativeServer_startAttach llama_server_set_embedded(true); server_context *ctx_server = &jctx->server; - srv->worker = std::thread([srv, ctx_server]() { + // Take a jllama_context user reference for the lifetime of the attach worker so that + // LlamaModel.close() cannot free jctx (and thus ctx_server) out from under the running + // HTTP frontend (use-after-free, MED #2). delete() in jllama.cpp waits on the user-count + // condition variable, so close() blocks until this thread exits. + jctx->users.fetch_add(1, std::memory_order_relaxed); + srv->worker = std::thread([srv, ctx_server, jctx]() { srv->exit_code = llama_server_attach(static_cast(srv->argv.size()), srv->argv.data(), *ctx_server); srv->finished.store(true); + release_jllama_context_impl(jctx); }); return reinterpret_cast(srv); @@ -173,11 +241,7 @@ JNIEXPORT void JNICALL Java_net_ladenthin_llama_server_NativeServer_setWorkerCom // model manager (server-models.cpp, patches/0008) reads when spawning worker instances. std::string value; if (jcommand != nullptr) { - const char *chars = env->GetStringUTFChars(jcommand, nullptr); - if (chars != nullptr) { - value = chars; - env->ReleaseStringUTFChars(jcommand, chars); - } + value = jstring_to_utf8(env, jcommand); } #if defined(_WIN32) _putenv_s("LLAMA_SERVER_WORKER_CMD", value.c_str()); // empty value removes the variable diff --git a/llama/src/main/cpp/tts_engine.cpp b/llama/src/main/cpp/tts_engine.cpp index 586b0571..5bbf1c69 100644 --- a/llama/src/main/cpp/tts_engine.cpp +++ b/llama/src/main/cpp/tts_engine.cpp @@ -28,6 +28,7 @@ #include #include +#include #include #include #include @@ -44,6 +45,9 @@ struct tts_engine { const llama_vocab *vocab = nullptr; outetts_version tts_version = OUTETTS_V0_2; int n_threads = 4; + // Serializes engine_synthesize: it drives llama_decode/llama_encode on the shared + // ctx_ttc/ctx_cts contexts, so two threads on one engine would race (M5). + std::mutex synthesize_mutex; }; tts_engine *engine_init(const std::string &ttc_model_path, const std::string &cts_model_path, int n_gpu_layers, @@ -98,6 +102,8 @@ bool engine_synthesize(tts_engine *engine, const std::string &text, int n_predic err = "engine is null"; return false; } + // Serialize against concurrent calls on the same engine (M5). + std::lock_guard engine_lock(engine->synthesize_mutex); const llama_vocab *vocab = engine->vocab; common_params_sampling sparams; @@ -185,12 +191,10 @@ bool engine_synthesize(tts_engine *engine, const std::string &text, int n_predic llama_batch_free(batch); common_sampler_free(smpl); - // Keep only audio codec tokens (151672..155772) and rebase to codec ids. - codes.erase(std::remove_if(codes.begin(), codes.end(), [](llama_token t) { return t < 151672 || t > 155772; }), - codes.end()); - for (auto &token : codes) { - token -= 151672; - } + // Keep only audio codec tokens and rebase to 0-based codec ids. The window is identical + // for OuteTTS V0_2 and V0_3 (verified against upstream tools/tts/tts.cpp, see M3 in + // docs/plan-bugs-and-issues.md) — delegated to the pure filter_ouetts_codec_tokens helper. + filter_ouetts_codec_tokens(codes); if (codes.empty()) { err = "no audio codes were generated"; return false; diff --git a/llama/src/main/cpp/tts_engine.h b/llama/src/main/cpp/tts_engine.h index ede04527..da07a8b1 100644 --- a/llama/src/main/cpp/tts_engine.h +++ b/llama/src/main/cpp/tts_engine.h @@ -9,6 +9,7 @@ #ifndef JLLAMA_TTS_ENGINE_H #define JLLAMA_TTS_ENGINE_H +#include #include #include #include @@ -31,6 +32,26 @@ bool engine_synthesize(tts_engine *engine, const std::string &text, int n_predic // Release both models / contexts. Safe on nullptr. void engine_free(tts_engine *engine); +// --------------------------------------------------------------------------- +// OuteTTS audio-codec token filtering (pure transform, testable without a model). +// +// The TTC model emits codec tokens in a high id window; the CTS vocoder only +// understands the 0-based codec ids inside [k_ouetts_codec_lo, k_ouetts_codec_hi]. +// This window is identical for OuteTTS V0_2 and V0_3 (verified against upstream +// tools/tts/tts.cpp — see M3 in docs/plan-bugs-and-issues.md), so it is NOT gated +// on tts_version. Keep only tokens in the window and rebase them to 0-based ids. +constexpr int32_t k_ouetts_codec_lo = 151672; +constexpr int32_t k_ouetts_codec_hi = 155772; + +inline void filter_ouetts_codec_tokens(std::vector &tokens) { + tokens.erase(std::remove_if(tokens.begin(), tokens.end(), + [](int32_t t) { return t < k_ouetts_codec_lo || t > k_ouetts_codec_hi; }), + tokens.end()); + for (auto &t : tokens) { + t -= k_ouetts_codec_lo; + } +} + } // namespace jllama_tts #endif // JLLAMA_TTS_ENGINE_H diff --git a/llama/src/main/java/net/ladenthin/llama/LlamaModel.java b/llama/src/main/java/net/ladenthin/llama/LlamaModel.java index 8d13b5d6..c4268786 100644 --- a/llama/src/main/java/net/ladenthin/llama/LlamaModel.java +++ b/llama/src/main/java/net/ladenthin/llama/LlamaModel.java @@ -324,7 +324,10 @@ public CompletableFuture chatCompleteTextAsync(InferenceParameters param * @return the text generated up to the point of stop or cancellation */ public String complete(InferenceParameters parameters, CancellationToken token) { - token.reset(); + // NOTE: do not reset the token here. A caller may pre-cancel it (or reuse one that + // carries a cancel from a previous aborted run that never reached finally); wiping it + // at entry would defeat cooperative cancellation. The finally block below resets the + // token after the call returns, so it stays reusable for the next call. InferenceParameters streaming = parameters.withStream(true); int taskId = requestCompletion(streaming.toString()); StringBuilder sb = new StringBuilder(); @@ -789,8 +792,8 @@ public String getMetrics() { /** * Run {@link #complete(InferenceParameters)} constrained to the supplied JSON Schema * and deserialize the result into an instance of {@code type}. The schema is applied - * via {@link net.ladenthin.llama.parameters.InferenceParameters#withJsonSchema(String)} for the duration of this call; - * the supplied {@code parameters} object is mutated. + * via {@link net.ladenthin.llama.parameters.InferenceParameters#withJsonSchema(String)} for the duration of this call. + * The supplied {@code parameters} object is NOT mutated — a new derivation carrying the schema is used internally. *

* Callers are responsible for producing a JSON Schema that matches the target type; * this project intentionally does not pull in a schema-from-POJO generator. Use the diff --git a/llama/src/main/java/net/ladenthin/llama/json/ChatResponseParser.java b/llama/src/main/java/net/ladenthin/llama/json/ChatResponseParser.java index 8b66e8fa..b27fc062 100644 --- a/llama/src/main/java/net/ladenthin/llama/json/ChatResponseParser.java +++ b/llama/src/main/java/net/ladenthin/llama/json/ChatResponseParser.java @@ -119,10 +119,10 @@ public String extractChoiceContent(JsonNode node) { * * @param node the parsed chat completion response * @param field the field name within {@code "usage"} - * @return the integer value, or {@code 0} if the field or the {@code "usage"} object is absent + * @return the long value, or {@code 0} if the field or the {@code "usage"} object is absent */ - public int extractUsageField(JsonNode node, String field) { - return node.path("usage").path(field).asInt(0); + public long extractUsageField(JsonNode node, String field) { + return node.path("usage").path(field).asLong(0); } /** diff --git a/llama/src/main/java/net/ladenthin/llama/server/AnthropicApiSupport.java b/llama/src/main/java/net/ladenthin/llama/server/AnthropicApiSupport.java index ea314181..6a8dbca5 100644 --- a/llama/src/main/java/net/ladenthin/llama/server/AnthropicApiSupport.java +++ b/llama/src/main/java/net/ladenthin/llama/server/AnthropicApiSupport.java @@ -254,14 +254,14 @@ static String toAnthropicResponse(String openAiCompletionJson, String model) { stopReason = anthropicStopReason(choice.path("finish_reason").asText("stop")); JsonNode openAiUsage = completion.path("usage"); if (openAiUsage.isObject()) { - int promptTokens = openAiUsage.path("prompt_tokens").asInt(0); - int cachedTokens = openAiUsage + long promptTokens = openAiUsage.path("prompt_tokens").asLong(0); + long cachedTokens = openAiUsage .path("prompt_tokens_details") .path("cached_tokens") - .asInt(0); - usage.put("input_tokens", Math.max(0, promptTokens - cachedTokens)); + .asLong(0); + usage.put("input_tokens", Math.max(0L, promptTokens - cachedTokens)); usage.put("cache_read_input_tokens", cachedTokens); - usage.put("output_tokens", openAiUsage.path("completion_tokens").asInt(0)); + usage.put("output_tokens", openAiUsage.path("completion_tokens").asLong(0)); } } catch (IOException e) { stopReason = "end_turn"; @@ -401,7 +401,7 @@ static String messageDeltaEvent(String stopReason) { } /** Final message delta carrying token usage collected from the trailing OpenAI usage chunk. */ - static String messageDeltaEvent(String stopReason, int inputTokens, int outputTokens, int cachedTokens) { + static String messageDeltaEvent(String stopReason, long inputTokens, long outputTokens, long cachedTokens) { ObjectNode data = OBJECT_MAPPER.createObjectNode(); data.put("type", "message_delta"); ObjectNode delta = data.putObject("delta"); diff --git a/llama/src/main/java/net/ladenthin/llama/server/AnthropicStreamTranslator.java b/llama/src/main/java/net/ladenthin/llama/server/AnthropicStreamTranslator.java index 188365fd..7f002e9d 100644 --- a/llama/src/main/java/net/ladenthin/llama/server/AnthropicStreamTranslator.java +++ b/llama/src/main/java/net/ladenthin/llama/server/AnthropicStreamTranslator.java @@ -32,9 +32,9 @@ final class AnthropicStreamTranslator { private int textBlockIndex = -1; private int nextIndex; private String finishReason = "stop"; - private int inputTokens; - private int outputTokens; - private int cachedTokens; + private long inputTokens; + private long outputTokens; + private long cachedTokens; AnthropicStreamTranslator(String id, String model) { this.id = id; @@ -65,12 +65,12 @@ String onChunk(String openAiChunkJson) { accumulator.accept(chunk); JsonNode usage = chunk.path("usage"); if (usage.isObject()) { - int promptTokens = usage.path("prompt_tokens").asInt(0); + long promptTokens = usage.path("prompt_tokens").asLong(0); cachedTokens = usage.path("prompt_tokens_details") .path("cached_tokens") - .asInt(0); - inputTokens = Math.max(0, promptTokens - cachedTokens); - outputTokens = usage.path("completion_tokens").asInt(0); + .asLong(0); + inputTokens = Math.max(0L, promptTokens - cachedTokens); + outputTokens = usage.path("completion_tokens").asLong(0); } JsonNode choice = chunk.path("choices").path(0); JsonNode content = choice.path("delta").path("content"); diff --git a/llama/src/main/java/net/ladenthin/llama/server/OllamaApiSupport.java b/llama/src/main/java/net/ladenthin/llama/server/OllamaApiSupport.java index 4ec6748d..21f566b7 100644 --- a/llama/src/main/java/net/ladenthin/llama/server/OllamaApiSupport.java +++ b/llama/src/main/java/net/ladenthin/llama/server/OllamaApiSupport.java @@ -211,8 +211,8 @@ static String toOllamaChatResponse(String openAiCompletionJson, String model) { } JsonNode usage = completion.path("usage"); if (usage.isObject()) { - root.put("prompt_eval_count", usage.path("prompt_tokens").asInt(0)); - root.put("eval_count", usage.path("completion_tokens").asInt(0)); + root.put("prompt_eval_count", usage.path("prompt_tokens").asLong(0)); + root.put("eval_count", usage.path("completion_tokens").asLong(0)); } } catch (IOException e) { // Defensive: an unexpected body still yields a valid, empty Ollama "done" response. diff --git a/llama/src/main/java/net/ladenthin/llama/server/ResponsesApiSupport.java b/llama/src/main/java/net/ladenthin/llama/server/ResponsesApiSupport.java index 5dd6a4e0..23d1ce03 100644 --- a/llama/src/main/java/net/ladenthin/llama/server/ResponsesApiSupport.java +++ b/llama/src/main/java/net/ladenthin/llama/server/ResponsesApiSupport.java @@ -202,8 +202,8 @@ static String toResponsesResponse(String openAiCompletionJson, String model, Str } JsonNode openAiUsage = completion.path("usage"); if (openAiUsage.isObject()) { - int in = openAiUsage.path("prompt_tokens").asInt(0); - int out = openAiUsage.path("completion_tokens").asInt(0); + long in = openAiUsage.path("prompt_tokens").asLong(0); + long out = openAiUsage.path("completion_tokens").asLong(0); usage.put("input_tokens", in); usage.put("output_tokens", out); usage.put("total_tokens", in + out); @@ -212,7 +212,7 @@ static String toResponsesResponse(String openAiCompletionJson, String model, Str openAiUsage .path("prompt_tokens_details") .path("cached_tokens") - .asInt(0)); + .asLong(0)); } } catch (IOException e) { // Defensive: an unexpected body still yields a valid, empty completed response. diff --git a/llama/src/main/java/net/ladenthin/llama/server/ResponsesStreamTranslator.java b/llama/src/main/java/net/ladenthin/llama/server/ResponsesStreamTranslator.java index eb50e8f0..51665a1b 100644 --- a/llama/src/main/java/net/ladenthin/llama/server/ResponsesStreamTranslator.java +++ b/llama/src/main/java/net/ladenthin/llama/server/ResponsesStreamTranslator.java @@ -39,9 +39,9 @@ final class ResponsesStreamTranslator { private boolean messageOpen; private int nextOutputIndex; private int messageOutputIndex = -1; - private int inputTokens; - private int outputTokens; - private int cachedTokens; + private long inputTokens; + private long outputTokens; + private long cachedTokens; ResponsesStreamTranslator(String model, String responseId) { this.model = model; @@ -160,9 +160,9 @@ private void captureUsage(JsonNode usage) { if (!usage.isObject()) { return; } - inputTokens = usage.path("prompt_tokens").asInt(0); - outputTokens = usage.path("completion_tokens").asInt(0); - cachedTokens = usage.path("prompt_tokens_details").path("cached_tokens").asInt(0); + inputTokens = usage.path("prompt_tokens").asLong(0); + outputTokens = usage.path("completion_tokens").asLong(0); + cachedTokens = usage.path("prompt_tokens_details").path("cached_tokens").asLong(0); } private ObjectNode messageItemShell() { diff --git a/llama/src/test/cpp/test_json_helpers.cpp b/llama/src/test/cpp/test_json_helpers.cpp index a5afb95e..e332bc79 100644 --- a/llama/src/test/cpp/test_json_helpers.cpp +++ b/llama/src/test/cpp/test_json_helpers.cpp @@ -241,6 +241,44 @@ TEST(RerankResultsToJson, PreservesInputOrder) { EXPECT_FLOAT_EQ(out[2].value("score", 0.0f), 0.6f); } +// A rerank result whose to_json() omits the "index" field entirely. The real +// upstream formatter always emits it, but a malformed/future result may not, and +// rerank_results_to_json must reject it rather than throw json::type_error (M2). +struct fake_rerank_result_no_index : server_task_result { + float score; + explicit fake_rerank_result_no_index(int id_, float sc) : score(sc) { id = id_; } + json to_json() override { return {{"score", score}}; } +}; + +static server_task_result_ptr make_rerank_no_index(int id_, float sc) { + return std::make_unique(id_, sc); +} + +TEST(RerankResultsToJson, MissingIndex_ThrowsInvalidArgument) { + std::vector results; + results.push_back(make_rerank_no_index(1, 0.5f)); + std::vector docs = {"doc A"}; + + EXPECT_THROW((void)rerank_results_to_json(results, docs), std::invalid_argument); +} + +TEST(RerankResultsToJson, IndexAboveRange_ThrowsInvalidArgument) { + // index 5 with only 2 documents is out of bounds — must throw, not OOB-index. + std::vector results; + results.push_back(make_rerank(1, 5, 0.5f)); + std::vector docs = {"doc zero", "doc one"}; + + EXPECT_THROW((void)rerank_results_to_json(results, docs), std::invalid_argument); +} + +TEST(RerankResultsToJson, NegativeIndex_ThrowsInvalidArgument) { + std::vector results; + results.push_back(make_rerank(1, -1, 0.5f)); + std::vector docs = {"only doc"}; + + EXPECT_THROW((void)rerank_results_to_json(results, docs), std::invalid_argument); +} + // ============================================================ // parse_encoding_format // ============================================================ diff --git a/llama/src/test/cpp/test_tts_wav.cpp b/llama/src/test/cpp/test_tts_wav.cpp index f7048061..a04b00fb 100644 --- a/llama/src/test/cpp/test_tts_wav.cpp +++ b/llama/src/test/cpp/test_tts_wav.cpp @@ -12,6 +12,8 @@ #include #include +#include "tts_engine.h" + using namespace jllama_tts; namespace { @@ -51,3 +53,42 @@ TEST(TtsWav, ClampsAndEncodesSamplesLittleEndian) { EXPECT_EQ(sample(2), -32767); EXPECT_EQ(sample(3), -32768); } + +// ============================================================ +// filter_ouetts_codec_tokens (M3 — V0_2 / V0_3 codec window) +// ============================================================ + +// The OuteTTS codec window is identical for V0_2 and V0_3 (verified against upstream +// tools/tts/tts.cpp). In-range tokens are kept and rebased to 0-based codec ids. +TEST(OuteTtsCodecFilter, InRangeTokensAreRebased) { + std::vector codes = {151672, 151673, 155772}; + filter_ouetts_codec_tokens(codes); + ASSERT_EQ(codes.size(), 3u); + EXPECT_EQ(codes[0], 0); + EXPECT_EQ(codes[1], 1); + EXPECT_EQ(codes[2], 155772 - 151672); +} + +// Tokens below / above the window (e.g. text, control, or speaker tokens) are dropped. +TEST(OuteTtsCodecFilter, OutOfRangeTokensAreDropped) { + std::vector codes = {151671, 198 /*newline*/, 155773, 151672}; + filter_ouetts_codec_tokens(codes); + ASSERT_EQ(codes.size(), 1u); + EXPECT_EQ(codes[0], 0); +} + +TEST(OuteTtsCodecFilter, EmptyInput_StaysEmpty) { + std::vector codes; + filter_ouetts_codec_tokens(codes); + EXPECT_TRUE(codes.empty()); +} + +TEST(OuteTtsCodecFilter, V02AndV03ShareWindow) { + // Whatever the engine version, the same window applies (M3). Exercise both boundaries. + std::vector v02 = {151672, 153000, 155772}; + std::vector v03 = {151672, 153000, 155772}; + filter_ouetts_codec_tokens(v02); + filter_ouetts_codec_tokens(v03); + EXPECT_EQ(v02, v03); + EXPECT_EQ(v02, (std::vector{0, 153000 - 151672, 155772 - 151672})); +} diff --git a/llama/src/test/java/net/ladenthin/llama/LlamaModelTest.java b/llama/src/test/java/net/ladenthin/llama/LlamaModelTest.java index d78c1eb8..39b1b64d 100644 --- a/llama/src/test/java/net/ladenthin/llama/LlamaModelTest.java +++ b/llama/src/test/java/net/ladenthin/llama/LlamaModelTest.java @@ -10,6 +10,7 @@ import java.io.*; import java.nio.charset.StandardCharsets; import java.util.*; +import java.util.concurrent.atomic.AtomicReference; import java.util.regex.Pattern; import net.ladenthin.llama.args.LogFormat; import net.ladenthin.llama.callback.CancellationToken; @@ -338,8 +339,57 @@ public void testCompleteWithCancellationToken() throws Exception { } /** - * Regression: {@link LlamaModel#completeAsync(InferenceParameters)} must - * complete with the same text {@link LlamaModel#complete(InferenceParameters)} + * Regression (H2 / M4 in docs/plan-bugs-and-issues.md): a {@link LlamaModel#close()} + * racing an in-flight {@link LlamaModel#complete(InferenceParameters)} on another thread + * must not crash the JVM (use-after-free) or hang. The native layer ref-counts in-flight + * JNI entry points and waits for them to drain before freeing the context, so close() + * either lets the in-flight call finish or surfaces a clean {@link LlamaException} — never + * a use-after-free. The in-flight thread must always terminate. + *

+ * Uses its own model instance (not the shared {@code model}) so closing it here does not + * perturb the other tests in this class. + */ + @Test + public void testCloseDuringInference() throws Exception { + Assumptions.assumeTrue( + new java.io.File("models/codellama-7b.Q2_K.gguf").exists(), "Model file not found, skipping"); + int gpuLayers = Integer.getInteger(TestConstants.PROP_TEST_NGL, TestConstants.DEFAULT_TEST_NGL); + try (LlamaModel localModel = new LlamaModel(new ModelParameters() + .setCtxSize(128) + .setModel(TestConstants.MODEL_PATH) + .setGpuLayers(gpuLayers) + .setFit(false))) { + + InferenceParameters params = new InferenceParameters(prefix).withNPredict(256); + AtomicReference inferenceError = new AtomicReference<>(); + Thread inferer = new Thread(() -> { + try { + localModel.complete(params); + } catch (Throwable t) { + // A clean LlamaException (context torn down mid-loop) is an acceptable + // degraded outcome; a hard Error (e.g. a JVM-level crash) is not. + inferenceError.set(t); + } + }); + inferer.start(); + // Let generation get underway, then close() from this (different) thread. + Thread.sleep(150); + localModel.close(); + + inferer.join(30000); + assertFalse(inferer.isAlive(), "in-flight inference thread must not hang after close()"); + Throwable t = inferenceError.get(); + if (t != null) { + assertFalse(t instanceof Error, "in-flight inference crashed: " + t); + } + // close() must remain idempotent and safe to call again. + localModel.close(); + } + } + + /** + * Regression: {@link LlamaModel#completeAsync(InferenceParameters)} + * must complete with the same text {@link LlamaModel#complete(InferenceParameters)} * would have produced, on a background thread. */ @Test