From 88fdd6d8cdb07181c4a2d3ea7adad05b0648fccc Mon Sep 17 00:00:00 2001
From: vaiju1981 <vaiju1981@gmail.com>
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 <stdexcept>
 #include <thread>
 
+// 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<jstring>(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<char *>(malloc(static_cast<size_t>(blen) + 1));
+                if (result[i] != nullptr) {
+                    memcpy(result[i], bytesElems, static_cast<size_t>(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<void(ggml_log_level, const char *, void *)> 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<void(ggml_log_level, const char *, void *)> cb;
+    bool json_mode;
+    {
+        std::lock_guard<std::mutex> 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<size_t>(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<server_task_result_embd *>(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<std::mutex> 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<std::mutex> 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<std::mutex> 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<std::mutex> 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 <atomic>
+#include <condition_variable>
 #include <map>
 #include <memory>
 #include <mutex>
@@ -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<int, std::shared_ptr<server_response_reader>> 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<int> 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<bool> 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<jllama_context *>(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<std::mutex> lk(g_ctx_mutex);
+    const jlong handle = env->GetLongField(obj, field_id);
+    if (handle == 0) {
+        return nullptr;
+    }
+    auto *jctx = reinterpret_cast<jllama_context *>(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<std::mutex> 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 <cmath>
 #include <optional>
 #include <stdexcept>
 #include <string>
@@ -90,7 +91,15 @@
     json arr = json::array();
     for (const auto &result : results) {
         const auto out = result->to_json();
-        int index = out["index"].get<int>();
+        // 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<int>();
+        if (index < 0 || static_cast<size_t>(index) >= documents.size()) {
+            throw std::invalid_argument("rerank result index " + std::to_string(index) + " out of range");
+        }
         float score = out["score"].get<float>();
         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<float>();
-    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<double>();
+    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<float>(v);
 }
 
 // ---------------------------------------------------------------------------
@@ -188,11 +199,15 @@
     if (!config.contains(key)) {
         return std::nullopt;
     }
-    const int v = config[key].get<int>();
-    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<double>();
+    // Reject non-positive, non-integer, or > INT_MAX values: a whole-number JSON value above
+    // INT_MAX would overflow static_cast<int> (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<int>(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<jstring>(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<jclass>(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<char *>(elems), static_cast<size_t>(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<jllama_context *>(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<int>(srv->argv.size()), srv->argv.data(), *ctx_server);
         srv->finished.store(true);
+        release_jllama_context_impl(jctx);
     });
 
     return reinterpret_cast<jlong>(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 <algorithm>
 #include <cstdint>
+#include <mutex>
 #include <regex>
 #include <string>
 #include <vector>
@@ -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<std::mutex> 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 <algorithm>
 #include <cstdint>
 #include <string>
 #include <vector>
@@ -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<int32_t> &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<String> 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.
      * <p>
      * 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<fake_rerank_result_no_index>(id_, sc);
+}
+
+TEST(RerankResultsToJson, MissingIndex_ThrowsInvalidArgument) {
+    std::vector<server_task_result_ptr> results;
+    results.push_back(make_rerank_no_index(1, 0.5f));
+    std::vector<std::string> 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<server_task_result_ptr> results;
+    results.push_back(make_rerank(1, 5, 0.5f));
+    std::vector<std::string> docs = {"doc zero", "doc one"};
+
+    EXPECT_THROW((void)rerank_results_to_json(results, docs), std::invalid_argument);
+}
+
+TEST(RerankResultsToJson, NegativeIndex_ThrowsInvalidArgument) {
+    std::vector<server_task_result_ptr> results;
+    results.push_back(make_rerank(1, -1, 0.5f));
+    std::vector<std::string> 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 <gtest/gtest.h>
 #include <vector>
 
+#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<int32_t> 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<int32_t> 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<int32_t> 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<int32_t> v02 = {151672, 153000, 155772};
+    std::vector<int32_t> v03 = {151672, 153000, 155772};
+    filter_ouetts_codec_tokens(v02);
+    filter_ouetts_codec_tokens(v03);
+    EXPECT_EQ(v02, v03);
+    EXPECT_EQ(v02, (std::vector<int32_t>{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.
+     * <p>
+     * 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<Throwable> 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
