From 3632b16c58d74d256f189c70fb514a5d62f65d3d Mon Sep 17 00:00:00 2001
From: Niko Maroulis <nik.maroulis@gmail.com>
Date: Sat, 4 Jul 2026 11:19:52 -0400
Subject: [PATCH 1/2] Add credo and fix all --strict findings

- Add credo ~> 1.7 (dev/test) and a .credo.exs that excludes the NIF
  entry points from FunctionArity (arities mirror the fixed C ABI)
- Replace length/1 emptiness checks and fix number formatting in tests
- Alphabetize alias groups
- Swap negated if-else branches in Budget.distribute/3 and MTP.init/2
- Flatten deep nesting: guard clauses in Strategy.Batch reduces,
  pattern-matched embed helpers in Schema, split_weights/check_gpu in
  Budget, and run_forward_pass/emit_tick_telemetry/continue_if_active
  extracted from Server.run_tick/1
- Dedupe HuggingFace request handling in Hub behind hf_get/3 and
  hf_api_error/2 with parse_search_results/parse_gguf_tree helpers
---
 .credo.exs                                |  20 ++++
 lib/llama_cpp_ex.ex                       |  14 +--
 lib/llama_cpp_ex/embedding.ex             |   2 +-
 lib/llama_cpp_ex/hub.ex                   | 127 ++++++++++------------
 lib/llama_cpp_ex/model_manager/budget.ex  |  53 ++++-----
 lib/llama_cpp_ex/mtp.ex                   |   6 +-
 lib/llama_cpp_ex/schema.ex                |  33 +++---
 lib/llama_cpp_ex/server.ex                | 101 +++++++++--------
 lib/llama_cpp_ex/server/strategy/batch.ex | 108 +++++++++---------
 mix.exs                                   |   1 +
 mix.lock                                  |   3 +
 test/batch_strategy_test.exs              |   2 +-
 test/llama_cpp_ex_test.exs                |  14 +--
 13 files changed, 249 insertions(+), 235 deletions(-)
 create mode 100644 .credo.exs

diff --git a/.credo.exs b/.credo.exs
new file mode 100644
index 0000000..49c0786
--- /dev/null
+++ b/.credo.exs
@@ -0,0 +1,20 @@
+# Credo configuration. Runs on top of Credo's default check set.
+%{
+  configs: [
+    %{
+      name: "default",
+      files: %{
+        included: ["lib/", "test/"],
+        excluded: []
+      },
+      checks: %{
+        extra: [
+          # NIF entry points mirror the positional C ABI in c_src/llama_cpp_ex —
+          # their arity is fixed by the native signatures, not a style choice.
+          # The keyword-based wrappers around them keep the public API ergonomic.
+          {Credo.Check.Refactor.FunctionArity, files: %{excluded: ["lib/llama_cpp_ex/nif.ex"]}}
+        ]
+      }
+    }
+  ]
+}
diff --git a/lib/llama_cpp_ex.ex b/lib/llama_cpp_ex.ex
index 7535496..86542a8 100644
--- a/lib/llama_cpp_ex.ex
+++ b/lib/llama_cpp_ex.ex
@@ -28,16 +28,16 @@ defmodule LlamaCppEx do
   """
 
   alias LlamaCppEx.{
-    Model,
-    Context,
-    Sampler,
-    Tokenizer,
     Chat,
-    Embedding,
-    Grammar,
     ChatCompletion,
     ChatCompletionChunk,
-    Thinking
+    Context,
+    Embedding,
+    Grammar,
+    Model,
+    Sampler,
+    Thinking,
+    Tokenizer
   }
 
   @context_opt_keys [
diff --git a/lib/llama_cpp_ex/embedding.ex b/lib/llama_cpp_ex/embedding.ex
index 2fbc827..59b2d14 100644
--- a/lib/llama_cpp_ex/embedding.ex
+++ b/lib/llama_cpp_ex/embedding.ex
@@ -1,7 +1,7 @@
 defmodule LlamaCppEx.Embedding do
   @moduledoc "Generate embeddings from text using an embedding model."
 
-  alias LlamaCppEx.{Model, Context, Tokenizer}
+  alias LlamaCppEx.{Context, Model, Tokenizer}
 
   @type t :: [float()]
 
diff --git a/lib/llama_cpp_ex/hub.ex b/lib/llama_cpp_ex/hub.ex
index c600d9b..48b7a20 100644
--- a/lib/llama_cpp_ex/hub.ex
+++ b/lib/llama_cpp_ex/hub.ex
@@ -95,37 +95,17 @@ defmodule LlamaCppEx.Hub do
   @spec search(String.t(), keyword()) :: {:ok, [map()]} | {:error, String.t()}
   def search(query, opts \\ []) do
     with :ok <- ensure_req() do
-      limit = Keyword.get(opts, :limit, 10)
-      sort = Keyword.get(opts, :sort, "downloads")
-      direction = Keyword.get(opts, :direction, -1)
-      headers = auth_headers(opts)
-
       params = [
         search: query,
         filter: "gguf",
-        sort: sort,
-        direction: direction,
-        limit: limit
+        sort: Keyword.get(opts, :sort, "downloads"),
+        direction: Keyword.get(opts, :direction, -1),
+        limit: Keyword.get(opts, :limit, 10)
       ]
 
-      req_opts = [headers: headers, params: params] ++ proxy_request_options(@hf_api_url, opts)
-
-      case Req.get(@hf_api_url, req_opts) do
+      case hf_get(@hf_api_url, opts, params: params) do
         {:ok, %{status: 200, body: body}} when is_list(body) ->
-          models =
-            Enum.map(body, fn m ->
-              %{
-                id: m["id"] || m["modelId"],
-                downloads: m["downloads"] || 0,
-                likes: m["likes"] || 0,
-                last_modified: m["lastModified"],
-                tags: m["tags"] || [],
-                private: m["private"] || false,
-                gated: m["gated"] || false
-              }
-            end)
-
-          {:ok, models}
+          {:ok, parse_search_results(body)}
 
         {:ok, %{status: status}} ->
           {:error, "HuggingFace search returned status #{status}"}
@@ -136,6 +116,20 @@ defmodule LlamaCppEx.Hub do
     end
   end
 
+  defp parse_search_results(body) do
+    Enum.map(body, fn m ->
+      %{
+        id: m["id"] || m["modelId"],
+        downloads: m["downloads"] || 0,
+        likes: m["likes"] || 0,
+        last_modified: m["lastModified"],
+        tags: m["tags"] || [],
+        private: m["private"] || false,
+        gated: m["gated"] || false
+      }
+    end)
+  end
+
   # --- Download ---
 
   @doc """
@@ -206,39 +200,21 @@ defmodule LlamaCppEx.Hub do
     with :ok <- ensure_req() do
       revision = Keyword.get(opts, :revision, "main")
       url = "#{@hf_api_url}/#{repo_id}/tree/#{revision}"
-      headers = auth_headers(opts)
-
-      case Req.get(url, [headers: headers] ++ proxy_request_options(url, opts)) do
-        {:ok, %{status: 200, body: body}} when is_list(body) ->
-          files =
-            body
-            |> Enum.filter(
-              &(&1["type"] == "file" and String.ends_with?(&1["path"] || "", ".gguf"))
-            )
-            |> Enum.map(fn f -> %{filename: f["path"], size: f["size"] || 0} end)
-            |> Enum.sort_by(& &1.size)
-
-          {:ok, files}
 
-        {:ok, %{status: 401}} ->
-          {:error, "authentication required — set HF_TOKEN or pass :token option"}
-
-        {:ok, %{status: 403}} ->
-          {:error,
-           "access denied — this may be a gated model requiring access approval at #{@hf_base_url}/#{repo_id}"}
-
-        {:ok, %{status: 404}} ->
-          {:error, "repository not found: #{repo_id}"}
-
-        {:ok, %{status: status}} ->
-          {:error, "HuggingFace API returned status #{status}"}
-
-        {:error, exception} ->
-          {:error, "network error: #{Exception.message(exception)}"}
+      case hf_get(url, opts) do
+        {:ok, %{status: 200, body: body}} when is_list(body) -> {:ok, parse_gguf_tree(body)}
+        other -> hf_api_error(other, repo_id)
       end
     end
   end
 
+  defp parse_gguf_tree(body) do
+    body
+    |> Enum.filter(&(&1["type"] == "file" and String.ends_with?(&1["path"] || "", ".gguf")))
+    |> Enum.map(fn f -> %{filename: f["path"], size: f["size"] || 0} end)
+    |> Enum.sort_by(& &1.size)
+  end
+
   # --- Model Info ---
 
   @doc """
@@ -253,30 +229,39 @@ defmodule LlamaCppEx.Hub do
   def get_model_info(repo_id, opts \\ []) do
     with :ok <- ensure_req() do
       url = "#{@hf_api_url}/#{repo_id}"
-      headers = auth_headers(opts)
 
-      case Req.get(url, [headers: headers] ++ proxy_request_options(url, opts)) do
-        {:ok, %{status: 200, body: body}} ->
-          {:ok, body}
+      case hf_get(url, opts) do
+        {:ok, %{status: 200, body: body}} -> {:ok, body}
+        other -> hf_api_error(other, repo_id)
+      end
+    end
+  end
 
-        {:ok, %{status: 401}} ->
-          {:error, "authentication required — set HF_TOKEN or pass :token option"}
+  # --- Shared HF API request helpers ---
 
-        {:ok, %{status: 403}} ->
-          {:error,
-           "access denied — this may be a gated model requiring access approval at #{@hf_base_url}/#{repo_id}"}
+  # Issues a GET with auth headers and proxy options applied.
+  defp hf_get(url, opts, extra_req_opts \\ []) do
+    req_opts = [headers: auth_headers(opts)] ++ extra_req_opts ++ proxy_request_options(url, opts)
+    Req.get(url, req_opts)
+  end
 
-        {:ok, %{status: 404}} ->
-          {:error, "repository not found: #{repo_id}"}
+  # Maps a non-200 HF API response or transport error to an error tuple.
+  defp hf_api_error({:ok, %{status: 401}}, _repo_id),
+    do: {:error, "authentication required — set HF_TOKEN or pass :token option"}
 
-        {:ok, %{status: status}} ->
-          {:error, "HuggingFace API returned status #{status}"}
+  defp hf_api_error({:ok, %{status: 403}}, repo_id),
+    do:
+      {:error,
+       "access denied — this may be a gated model requiring access approval at #{@hf_base_url}/#{repo_id}"}
 
-        {:error, exception} ->
-          {:error, "network error: #{Exception.message(exception)}"}
-      end
-    end
-  end
+  defp hf_api_error({:ok, %{status: 404}}, repo_id),
+    do: {:error, "repository not found: #{repo_id}"}
+
+  defp hf_api_error({:ok, %{status: status}}, _repo_id),
+    do: {:error, "HuggingFace API returned status #{status}"}
+
+  defp hf_api_error({:error, exception}, _repo_id),
+    do: {:error, "network error: #{Exception.message(exception)}"}
 
   # --- Public Helpers ---
 
diff --git a/lib/llama_cpp_ex/model_manager/budget.ex b/lib/llama_cpp_ex/model_manager/budget.ex
index 65ecc23..cb018f3 100644
--- a/lib/llama_cpp_ex/model_manager/budget.ex
+++ b/lib/llama_cpp_ex/model_manager/budget.ex
@@ -123,41 +123,40 @@ defmodule LlamaCppEx.ModelManager.Budget do
     n_gpu_layers = Keyword.get(opts, :n_gpu_layers, 99)
     offloaded? = n_gpu_layers != 0 and n_gpus > 0
 
-    if not offloaded? do
-      %{ram: file_bytes + kv, vram: %{}}
-    else
+    if offloaded? do
       kqv_on_gpu? = Keyword.get(opts, :offload_kqv, true)
       vram_total = file_bytes + if(kqv_on_gpu?, do: kv, else: 0)
       ram_total = if kqv_on_gpu?, do: 0, else: kv
       weights = device_weights(opts, n_gpus)
       vram = Map.new(weights, fn {i, w} -> {i, round(vram_total * w)} end)
       %{ram: ram_total, vram: vram}
+    else
+      %{ram: file_bytes + kv, vram: %{}}
     end
   end
 
   # Returns %{gpu_index => fraction}, summing to 1.0.
   defp device_weights(opts, n_gpus) do
     case Keyword.get(opts, :split_mode, :none) do
-      :none ->
-        %{Keyword.get(opts, :main_gpu, 0) => 1.0}
-
-      _layer_or_row ->
-        case Keyword.get(opts, :tensor_split, []) do
-          [] ->
-            frac = 1.0 / n_gpus
-            Map.new(0..(n_gpus - 1), fn i -> {i, frac} end)
-
-          weights ->
-            sum = Enum.sum(weights)
-
-            weights
-            |> Enum.with_index()
-            |> Enum.reject(fn {w, _i} -> w == 0 end)
-            |> Map.new(fn {w, i} -> {i, w / sum} end)
-        end
+      :none -> %{Keyword.get(opts, :main_gpu, 0) => 1.0}
+      _layer_or_row -> split_weights(Keyword.get(opts, :tensor_split, []), n_gpus)
     end
   end
 
+  defp split_weights([], n_gpus) do
+    frac = 1.0 / n_gpus
+    Map.new(0..(n_gpus - 1), fn i -> {i, frac} end)
+  end
+
+  defp split_weights(weights, _n_gpus) do
+    sum = Enum.sum(weights)
+
+    weights
+    |> Enum.with_index()
+    |> Enum.reject(fn {w, _i} -> w == 0 end)
+    |> Map.new(fn {w, i} -> {i, w / sum} end)
+  end
+
   # --- Fit check ---
 
   @doc "An empty usage accumulator."
@@ -193,16 +192,18 @@ defmodule LlamaCppEx.ModelManager.Budget do
   def check(%{mode: :per_device, ram: ram_limit, vram: vram}, placement, used) do
     with :ok <- fits(ram_limit, placement.ram, used.ram, :ram) do
       Enum.reduce_while(placement.vram, :ok, fn {gi, required}, :ok ->
-        limit = vram_limit(vram, gi)
-
-        case fits(limit, required, Map.get(used.vram, gi, 0), {:gpu, gi}) do
-          :ok -> {:cont, :ok}
-          error -> {:halt, error}
-        end
+        check_gpu(vram, gi, required, used)
       end)
     end
   end
 
+  defp check_gpu(vram, gi, required, used) do
+    case fits(vram_limit(vram, gi), required, Map.get(used.vram, gi, 0), {:gpu, gi}) do
+      :ok -> {:cont, :ok}
+      error -> {:halt, error}
+    end
+  end
+
   defp vram_limit(:infinity, _gi), do: :infinity
   defp vram_limit(map, gi), do: Map.get(map, gi, :infinity)
 
diff --git a/lib/llama_cpp_ex/mtp.ex b/lib/llama_cpp_ex/mtp.ex
index b02f215..2c473f1 100644
--- a/lib/llama_cpp_ex/mtp.ex
+++ b/lib/llama_cpp_ex/mtp.ex
@@ -89,9 +89,7 @@ defmodule LlamaCppEx.MTP do
   def init(%Model{} = model, opts \\ []) do
     n_draft = Keyword.get(opts, :n_draft, 3)
 
-    if not (is_integer(n_draft) and n_draft > 0) do
-      {:error, ":n_draft must be a positive integer"}
-    else
+    if is_integer(n_draft) and n_draft > 0 do
       base_ctx_opts = Keyword.take(opts, @context_opt_keys)
       main_opts = Keyword.merge(base_ctx_opts, ctx_type: :default)
       # Match upstream server: MTP draft context is created with n_rs_seq=0.
@@ -111,6 +109,8 @@ defmodule LlamaCppEx.MTP do
            n_draft: n_draft
          }}
       end
+    else
+      {:error, ":n_draft must be a positive integer"}
     end
   end
 
diff --git a/lib/llama_cpp_ex/schema.ex b/lib/llama_cpp_ex/schema.ex
index 415ecde..080e545 100644
--- a/lib/llama_cpp_ex/schema.ex
+++ b/lib/llama_cpp_ex/schema.ex
@@ -79,24 +79,8 @@ defmodule LlamaCppEx.Schema do
 
     {properties, required} =
       Enum.reduce(fields, {%{}, []}, fn field, {props, req} ->
-        if field in embeds do
-          embed_schema = module.__schema__(:embed, field)
-          embed_mod = embed_schema.related
-
-          case embed_schema.cardinality do
-            :one ->
-              nested = to_json_schema(embed_mod)
-              {Map.put(props, to_string(field), nested), [to_string(field) | req]}
-
-            :many ->
-              nested = to_json_schema(embed_mod)
-              item = %{"type" => "array", "items" => nested}
-              {Map.put(props, to_string(field), item), [to_string(field) | req]}
-          end
-        else
-          json_type = ecto_type_to_json(types[field])
-          {Map.put(props, to_string(field), json_type), [to_string(field) | req]}
-        end
+        property = field_property(module, field, types, embeds)
+        {Map.put(props, to_string(field), property), [to_string(field) | req]}
       end)
 
     %{
@@ -106,6 +90,19 @@ defmodule LlamaCppEx.Schema do
     }
   end
 
+  defp field_property(module, field, types, embeds) do
+    if field in embeds do
+      embed_property(module.__schema__(:embed, field))
+    else
+      ecto_type_to_json(types[field])
+    end
+  end
+
+  defp embed_property(%{cardinality: :one, related: embed_mod}), do: to_json_schema(embed_mod)
+
+  defp embed_property(%{cardinality: :many, related: embed_mod}),
+    do: %{"type" => "array", "items" => to_json_schema(embed_mod)}
+
   defp ecto_type_to_json(:string), do: %{"type" => "string"}
   defp ecto_type_to_json(:binary), do: %{"type" => "string"}
   defp ecto_type_to_json(:integer), do: %{"type" => "integer"}
diff --git a/lib/llama_cpp_ex/server.ex b/lib/llama_cpp_ex/server.ex
index 60d7636..829cc1c 100644
--- a/lib/llama_cpp_ex/server.ex
+++ b/lib/llama_cpp_ex/server.ex
@@ -103,7 +103,7 @@ defmodule LlamaCppEx.Server do
 
   require Logger
 
-  alias LlamaCppEx.{Model, Context, Sampler, Tokenizer}
+  alias LlamaCppEx.{Context, Model, Sampler, Tokenizer}
 
   defstruct [
     :model,
@@ -666,52 +666,65 @@ defmodule LlamaCppEx.Server do
     if entries == [] do
       state
     else
-      # Count decode tokens before sampling (slots may transition after)
-      n_decode =
-        Enum.count(state.slots, fn {_id, s} ->
-          s.state == :generating and s.batch_idx >= 0
-        end)
-
-      # Phase 3: Forward pass
-      tick_start = System.monotonic_time()
-
-      case LlamaCppEx.NIF.batch_eval(state.ctx.ref, entries) do
-        :ok ->
-          tick_end = System.monotonic_time()
-
-          # Phase 4: Sample
-          state = sample_generating_slots(state)
-          state = sample_completed_prefills(state)
-          state = advance_incomplete_prefills(state)
-
-          # Emit tick telemetry
-          :telemetry.execute(
-            [:llama_cpp_ex, :server, :tick],
-            %{
-              batch_size: length(entries),
-              decode_tokens: n_decode,
-              prefill_tokens: length(entries) - n_decode,
-              active_slots: Enum.count(state.slots, fn {_id, s} -> s.state != :idle end),
-              queue_depth: :queue.len(state.queue),
-              eval_ms: (tick_end - tick_start) / 1_000_000
-            },
-            %{server: self()}
-          )
-
-          # Phase 5: Continue
-          if Enum.any?(state.slots, fn {_id, slot} -> slot.state != :idle end) do
-            maybe_schedule_tick(state)
-          else
-            state
-          end
-
-        {:error, reason} ->
-          Logger.error("batch forward pass failed: #{reason}")
-          fail_all_active_slots(state, reason)
-      end
+      run_forward_pass(state, entries)
+    end
+  end
+
+  # Phases 3-5: forward pass, sampling, telemetry, and continuation
+  defp run_forward_pass(state, entries) do
+    # Count decode tokens before sampling (slots may transition after)
+    n_decode =
+      Enum.count(state.slots, fn {_id, s} ->
+        s.state == :generating and s.batch_idx >= 0
+      end)
+
+    # Phase 3: Forward pass
+    tick_start = System.monotonic_time()
+
+    case LlamaCppEx.NIF.batch_eval(state.ctx.ref, entries) do
+      :ok ->
+        tick_end = System.monotonic_time()
+
+        # Phase 4: Sample
+        state = sample_generating_slots(state)
+        state = sample_completed_prefills(state)
+        state = advance_incomplete_prefills(state)
+
+        emit_tick_telemetry(state, entries, n_decode, tick_end - tick_start)
+
+        # Phase 5: Continue
+        continue_if_active(state)
+
+      {:error, reason} ->
+        Logger.error("batch forward pass failed: #{reason}")
+        fail_all_active_slots(state, reason)
     end
   end
 
+  # Schedules the next tick while any slot is still active.
+  defp continue_if_active(state) do
+    if Enum.any?(state.slots, fn {_id, slot} -> slot.state != :idle end) do
+      maybe_schedule_tick(state)
+    else
+      state
+    end
+  end
+
+  defp emit_tick_telemetry(state, entries, n_decode, eval_native) do
+    :telemetry.execute(
+      [:llama_cpp_ex, :server, :tick],
+      %{
+        batch_size: length(entries),
+        decode_tokens: n_decode,
+        prefill_tokens: length(entries) - n_decode,
+        active_slots: Enum.count(state.slots, fn {_id, s} -> s.state != :idle end),
+        queue_depth: :queue.len(state.queue),
+        eval_ms: eval_native / 1_000_000
+      },
+      %{server: self()}
+    )
+  end
+
   # Phase 1: Check generating slots for completion
   defp finish_completed_slots(state) do
     generating_slots =
diff --git a/lib/llama_cpp_ex/server/strategy/batch.ex b/lib/llama_cpp_ex/server/strategy/batch.ex
index 8cd4ccf..4d28b2c 100644
--- a/lib/llama_cpp_ex/server/strategy/batch.ex
+++ b/lib/llama_cpp_ex/server/strategy/batch.ex
@@ -37,32 +37,31 @@ defmodule LlamaCppEx.Server.Strategy.Batch do
       |> Enum.sort_by(&elem(&1, 0))
 
     Enum.reduce(generating_slots, {entries, n_entries, slots, budget}, fn
+      {_seq_id, _slot}, {entries, n_entries, slots, budget} when budget <= 0 ->
+        {entries, n_entries, slots, budget}
+
       {seq_id, _slot}, {entries, n_entries, slots, budget} ->
-        if budget <= 0 do
-          {entries, n_entries, slots, budget}
-        else
-          slot = slots[seq_id]
-          token = slot.pending_token
-
-          piece = LlamaCppEx.NIF.token_to_piece(model_ref, token)
-
-          if slot.stream_pid && slot.stream_ref do
-            send(slot.stream_pid, {slot.stream_ref, {:token, piece}})
-          end
-
-          slot = %{
-            slot
-            | accumulated_pieces: [piece | slot.accumulated_pieces],
-              batch_idx: n_entries,
-              tokens_generated: slot.tokens_generated + 1,
-              generated_token_ids: [token | slot.generated_token_ids]
-          }
-
-          entry = {token, slot.pos, seq_id, true}
-          slots = Map.put(slots, seq_id, slot)
-
-          {[entry | entries], n_entries + 1, slots, budget - 1}
+        slot = slots[seq_id]
+        token = slot.pending_token
+
+        piece = LlamaCppEx.NIF.token_to_piece(model_ref, token)
+
+        if slot.stream_pid && slot.stream_ref do
+          send(slot.stream_pid, {slot.stream_ref, {:token, piece}})
         end
+
+        slot = %{
+          slot
+          | accumulated_pieces: [piece | slot.accumulated_pieces],
+            batch_idx: n_entries,
+            tokens_generated: slot.tokens_generated + 1,
+            generated_token_ids: [token | slot.generated_token_ids]
+        }
+
+        entry = {token, slot.pos, seq_id, true}
+        slots = Map.put(slots, seq_id, slot)
+
+        {[entry | entries], n_entries + 1, slots, budget - 1}
     end)
   end
 
@@ -78,40 +77,35 @@ defmodule LlamaCppEx.Server.Strategy.Batch do
       |> Enum.sort_by(&elem(&1, 0))
 
     Enum.reduce(prefilling_slots, {entries, n_entries, slots, budget}, fn
+      {_seq_id, _slot}, {entries, n_entries, slots, budget} when budget <= 0 ->
+        {entries, n_entries, slots, budget}
+
       {seq_id, _slot}, {entries, n_entries, slots, budget} ->
-        if budget <= 0 do
-          {entries, n_entries, slots, budget}
-        else
-          slot = slots[seq_id]
-          remaining = slot.n_prompt_tokens - slot.prefill_pos
-          chunk_len = min(budget, min(chunk_size, remaining))
-          is_last_chunk = slot.prefill_pos + chunk_len >= slot.n_prompt_tokens
-
-          tuple = slot.prompt_tokens_tuple
-
-          # Build this chunk's entries in O(chunk_len): elem/2 is O(1) on a tuple,
-          # and n_entries gives each entry's final batch index without length/1.
-          {entries, n_entries, last_batch_idx} =
-            Enum.reduce(0..(chunk_len - 1)//1, {entries, n_entries, -1}, fn i,
-                                                                            {entries, n_entries,
-                                                                             _last} ->
-              pos = slot.prefill_pos + i
-              token = elem(tuple, pos)
-              logits = is_last_chunk and i == chunk_len - 1
-              entry = {token, pos, seq_id, logits}
-              {[entry | entries], n_entries + 1, n_entries}
-            end)
-
-          slot =
-            if is_last_chunk do
-              %{slot | batch_idx: last_batch_idx, prefill_pos: slot.prefill_pos + chunk_len}
-            else
-              %{slot | batch_idx: -1, prefill_pos: slot.prefill_pos + chunk_len}
-            end
-
-          slots = Map.put(slots, seq_id, slot)
-          {entries, n_entries, slots, budget - chunk_len}
-        end
+        slot = slots[seq_id]
+        remaining = slot.n_prompt_tokens - slot.prefill_pos
+        chunk_len = min(budget, min(chunk_size, remaining))
+        is_last_chunk = slot.prefill_pos + chunk_len >= slot.n_prompt_tokens
+
+        tuple = slot.prompt_tokens_tuple
+
+        # Build this chunk's entries in O(chunk_len): elem/2 is O(1) on a tuple,
+        # and n_entries gives each entry's final batch index without length/1.
+        {entries, n_entries, last_batch_idx} =
+          Enum.reduce(0..(chunk_len - 1)//1, {entries, n_entries, -1}, fn i,
+                                                                          {entries, n_entries,
+                                                                           _last} ->
+            pos = slot.prefill_pos + i
+            token = elem(tuple, pos)
+            logits = is_last_chunk and i == chunk_len - 1
+            entry = {token, pos, seq_id, logits}
+            {[entry | entries], n_entries + 1, n_entries}
+          end)
+
+        batch_idx = if is_last_chunk, do: last_batch_idx, else: -1
+        slot = %{slot | batch_idx: batch_idx, prefill_pos: slot.prefill_pos + chunk_len}
+
+        slots = Map.put(slots, seq_id, slot)
+        {entries, n_entries, slots, budget - chunk_len}
     end)
   end
 end
diff --git a/mix.exs b/mix.exs
index f61ca53..89c4f2e 100644
--- a/mix.exs
+++ b/mix.exs
@@ -78,6 +78,7 @@ defmodule LlamaCppEx.MixProject do
       {:ecto, "~> 3.0", optional: true},
       {:req, "~> 0.5", optional: true},
       {:ex_doc, "~> 0.34", only: :dev, runtime: false},
+      {:credo, "~> 1.7", only: [:dev, :test], runtime: false},
       {:dialyxir, "~> 1.4", only: [:dev, :test], runtime: false},
       {:benchee, "~> 1.0", only: :bench, runtime: false},
       {:benchee_html, "~> 1.0", only: :bench, runtime: false}
diff --git a/mix.lock b/mix.lock
index bb4dc04..54364b3 100644
--- a/mix.lock
+++ b/mix.lock
@@ -2,6 +2,8 @@
   "benchee": {:hex, :benchee, "1.5.0", "4d812c31d54b0ec0167e91278e7de3f596324a78a096fd3d0bea68bb0c513b10", [:mix], [{:deep_merge, "~> 1.0", [hex: :deep_merge, repo: "hexpm", optional: false]}, {:statistex, "~> 1.1", [hex: :statistex, repo: "hexpm", optional: false]}, {:table, "~> 0.1.0", [hex: :table, repo: "hexpm", optional: true]}], "hexpm", "5b075393aea81b8ae74eadd1c28b1d87e8a63696c649d8293db7c4df3eb67535"},
   "benchee_html": {:hex, :benchee_html, "1.0.1", "1e247c0886c3fdb0d3f4b184b653a8d6fb96e4ad0d0389267fe4f36968772e24", [:mix], [{:benchee, ">= 0.99.0 and < 2.0.0", [hex: :benchee, repo: "hexpm", optional: false]}, {:benchee_json, "~> 1.0", [hex: :benchee_json, repo: "hexpm", optional: false]}], "hexpm", "b00a181af7152431901e08f3fc9f7197ed43ff50421a8347b0c80bf45d5b3fef"},
   "benchee_json": {:hex, :benchee_json, "1.0.0", "cc661f4454d5995c08fe10dd1f2f72f229c8f0fb1c96f6b327a8c8fc96a91fe5", [:mix], [{:benchee, ">= 0.99.0 and < 2.0.0", [hex: :benchee, repo: "hexpm", optional: false]}, {:jason, "~> 1.0", [hex: :jason, repo: "hexpm", optional: false]}], "hexpm", "da05d813f9123505f870344d68fb7c86a4f0f9074df7d7b7e2bb011a63ec231c"},
+  "bunt": {:hex, :bunt, "1.0.0", "081c2c665f086849e6d57900292b3a161727ab40431219529f13c4ddcf3e7a44", [:mix], [], "hexpm", "dc5f86aa08a5f6fa6b8096f0735c4e76d54ae5c9fa2c143e5a1fc7c1cd9bb6b5"},
+  "credo": {:hex, :credo, "1.7.19", "cc52129665fc7c15143d47838fda0f9cd6dac9ceced7bf4da6f85fcbfe64b12a", [:mix], [{:bunt, "~> 0.2.1 or ~> 1.0", [hex: :bunt, repo: "hexpm", optional: false]}, {:file_system, "~> 0.2 or ~> 1.0", [hex: :file_system, repo: "hexpm", optional: false]}, {:jason, "~> 1.0", [hex: :jason, repo: "hexpm", optional: false]}], "hexpm", "2d8bc95d5a7bb99dd2613621d4f08c6a3575c3fd4b62e6a2b48a100352a557b8"},
   "decimal": {:hex, :decimal, "2.3.0", "3ad6255aa77b4a3c4f818171b12d237500e63525c2fd056699967a3e7ea20f62", [:mix], [], "hexpm", "a4d66355cb29cb47c3cf30e71329e58361cfcb37c34235ef3bf1d7bf3773aeac"},
   "deep_merge": {:hex, :deep_merge, "1.0.0", "b4aa1a0d1acac393bdf38b2291af38cb1d4a52806cf7a4906f718e1feb5ee961", [:mix], [], "hexpm", "ce708e5f094b9cd4e8f2be4f00d2f4250c4095be93f8cd6d018c753894885430"},
   "dialyxir": {:hex, :dialyxir, "1.4.7", "dda948fcee52962e4b6c5b4b16b2d8fa7d50d8645bbae8b8685c3f9ecb7f5f4d", [:mix], [{:erlex, ">= 0.2.8", [hex: :erlex, repo: "hexpm", optional: false]}], "hexpm", "b34527202e6eb8cee198efec110996c25c5898f43a4094df157f8d28f27d9efe"},
@@ -10,6 +12,7 @@
   "elixir_make": {:hex, :elixir_make, "0.9.0", "6484b3cd8c0cee58f09f05ecaf1a140a8c97670671a6a0e7ab4dc326c3109726", [:mix], [], "hexpm", "db23d4fd8b757462ad02f8aa73431a426fe6671c80b200d9710caf3d1dd0ffdb"},
   "erlex": {:hex, :erlex, "0.2.8", "cd8116f20f3c0afe376d1e8d1f0ae2452337729f68be016ea544a72f767d9c12", [:mix], [], "hexpm", "9d66ff9fedf69e49dc3fd12831e12a8a37b76f8651dd21cd45fcf5561a8a7590"},
   "ex_doc": {:hex, :ex_doc, "0.40.1", "67542e4b6dde74811cfd580e2c0149b78010fd13001fda7cfeb2b2c2ffb1344d", [:mix], [{:earmark_parser, "~> 1.4.44", [hex: :earmark_parser, repo: "hexpm", optional: false]}, {:makeup_c, ">= 0.1.0", [hex: :makeup_c, repo: "hexpm", optional: true]}, {:makeup_elixir, "~> 0.14 or ~> 1.0", [hex: :makeup_elixir, repo: "hexpm", optional: false]}, {:makeup_erlang, "~> 0.1 or ~> 1.0", [hex: :makeup_erlang, repo: "hexpm", optional: false]}, {:makeup_html, ">= 0.1.0", [hex: :makeup_html, repo: "hexpm", optional: true]}], "hexpm", "bcef0e2d360d93ac19f01a85d58f91752d930c0a30e2681145feea6bd3516e00"},
+  "file_system": {:hex, :file_system, "1.1.1", "31864f4685b0148f25bd3fbef2b1228457c0c89024ad67f7a81a3ffbc0bbad3a", [:mix], [], "hexpm", "7a15ff97dfe526aeefb090a7a9d3d03aa907e100e262a0f8f7746b78f8f87a5d"},
   "finch": {:hex, :finch, "0.21.0", "b1c3b2d48af02d0c66d2a9ebfb5622be5c5ecd62937cf79a88a7f98d48a8290c", [:mix], [{:mime, "~> 1.0 or ~> 2.0", [hex: :mime, repo: "hexpm", optional: false]}, {:mint, "~> 1.6.2 or ~> 1.7", [hex: :mint, repo: "hexpm", optional: false]}, {:nimble_options, "~> 0.4 or ~> 1.0", [hex: :nimble_options, repo: "hexpm", optional: false]}, {:nimble_pool, "~> 1.1", [hex: :nimble_pool, repo: "hexpm", optional: false]}, {:telemetry, "~> 0.4 or ~> 1.0", [hex: :telemetry, repo: "hexpm", optional: false]}], "hexpm", "87dc6e169794cb2570f75841a19da99cfde834249568f2a5b121b809588a4377"},
   "fine": {:hex, :fine, "0.1.4", "b19a89c1476c7c57afb5f9314aed5960b5bc95d5277de4cb5ee8e1d1616ce379", [:mix], [], "hexpm", "be3324cc454a42d80951cf6023b9954e9ff27c6daa255483b3e8d608670303f5"},
   "hpax": {:hex, :hpax, "1.0.3", "ed67ef51ad4df91e75cc6a1494f851850c0bd98ebc0be6e81b026e765ee535aa", [:mix], [], "hexpm", "8eab6e1cfa8d5918c2ce4ba43588e894af35dbd8e91e6e55c817bca5847df34a"},
diff --git a/test/batch_strategy_test.exs b/test/batch_strategy_test.exs
index 4f10e17..5142089 100644
--- a/test/batch_strategy_test.exs
+++ b/test/batch_strategy_test.exs
@@ -1,7 +1,7 @@
 defmodule LlamaCppEx.Server.BatchStrategyTest do
   use ExUnit.Case, async: true
 
-  alias LlamaCppEx.Server.Strategy.{DecodeMaximal, PrefillPriority, Balanced}
+  alias LlamaCppEx.Server.Strategy.{Balanced, DecodeMaximal, PrefillPriority}
 
   # We need a mock model_ref for token_to_piece calls.
   # Since strategies call NIF.token_to_piece, we need a real model for integration.
diff --git a/test/llama_cpp_ex_test.exs b/test/llama_cpp_ex_test.exs
index 9a5fa76..7523956 100644
--- a/test/llama_cpp_ex_test.exs
+++ b/test/llama_cpp_ex_test.exs
@@ -209,7 +209,7 @@ defmodule LlamaCppExTest do
         text = "Hello, world!"
         {:ok, tokens} = LlamaCppEx.Tokenizer.encode(model, text, add_special: false)
         assert is_list(tokens)
-        assert length(tokens) > 0
+        assert tokens != []
 
         {:ok, decoded} = LlamaCppEx.Tokenizer.decode(model, tokens)
         assert decoded == text
@@ -259,7 +259,7 @@ defmodule LlamaCppExTest do
       end
 
       test "generate is deterministic with same seed", %{model: model} do
-        opts = [max_tokens: 16, seed: 12345, temp: 0.0]
+        opts = [max_tokens: 16, seed: 12_345, temp: 0.0]
         {:ok, text1} = LlamaCppEx.generate(model, "The answer is", opts)
         {:ok, text2} = LlamaCppEx.generate(model, "The answer is", opts)
         assert text1 == text2
@@ -279,7 +279,7 @@ defmodule LlamaCppExTest do
           |> LlamaCppEx.stream("Once upon a time", max_tokens: 16, seed: 42)
           |> Enum.to_list()
 
-        assert length(chunks) > 0
+        assert chunks != []
         assert Enum.all?(chunks, &is_binary/1)
 
         text = Enum.join(chunks)
@@ -389,7 +389,7 @@ defmodule LlamaCppExTest do
             )
             |> Enum.to_list()
 
-          assert length(chunks) > 0
+          assert chunks != []
         end
       end
 
@@ -710,7 +710,7 @@ defmodule LlamaCppExTest do
           LlamaCppEx.Server.stream(server, "Hello", max_tokens: 8)
           |> Enum.to_list()
 
-        assert length(chunks) > 0
+        assert chunks != []
         assert Enum.all?(chunks, &is_binary/1)
       end
 
@@ -896,7 +896,7 @@ defmodule LlamaCppExTest do
       test "embed single text", %{model: model} do
         {:ok, embedding} = LlamaCppEx.embed(model, "Hello world")
         assert is_list(embedding)
-        assert length(embedding) > 0
+        assert embedding != []
         assert Enum.all?(embedding, &is_float/1)
       end
 
@@ -958,7 +958,7 @@ defmodule LlamaCppExTest do
       test "embed with pooling_type option", %{model: model} do
         {:ok, embedding} = LlamaCppEx.embed(model, "Test", pooling_type: :last)
         assert is_list(embedding)
-        assert length(embedding) > 0
+        assert embedding != []
       end
     end
   else

From ccc5e04e4eb0d851696d6ab3b86c5642747b40d0 Mon Sep 17 00:00:00 2001
From: Niko Maroulis <nik.maroulis@gmail.com>
Date: Sat, 4 Jul 2026 11:36:28 -0400
Subject: [PATCH 2/2] Refactor tech-debt findings: dedup generation paths,
 narrow rescues
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

- llama_cpp_ex.ex: extract shared generation plumbing (gen_config,
  create_gen_resources, spawn_generator/stop_generator, chunk builder,
  finish_reason mapping) — the four entry points were ~90% copy-paste.
  generate/3 now honors all @context_opt_keys like the other entry
  points (it previously forwarded only 4 of them), and chat_completion/3
  now kills/drains its generator after collection so a timeout can't
  leave a runaway process and stray mailbox messages.
- server.ex: single idle_slot_fields/2 source of truth for slot resets
  (init/reset_slot/fail_all_active_slots each spelled out all 18 fields);
  shared request_measurements/2 behind the :done/:exception telemetry.
- model_manager.ex: with_route/3 dispatch skeleton for generate/stream/
  chat; merged the two build_ready_entry clauses; narrowed
  safe_backend_init's rescue-all to ErlangError/UndefinedFunctionError
  with a debug log.
- hub.ex: do_download_to rescues only File.Error/ErlangError so
  programming errors propagate.
- tokenizer.ex/chat.ex: NIF ErlangError wrappers now re-raise
  :not_loaded (packaging bug, not bad input) via NIF.error_tuple/3.
- embedding.ex: map_while_ok/2 replaces three hand-rolled short-circuit
  folds (also fixes O(n²) accumulation in decode_groups); NIF wrappers
  are thin delegates instead of no-op case passthroughs.
- tests: fix MTP n_rs_seq assertion — the draft context is intentionally
  created with n_rs_seq=0 (rollback uses cached hidden states); the old
  assertion contradicted the implementation and always failed.
---
 lib/llama_cpp_ex.ex               | 394 ++++++++++--------------------
 lib/llama_cpp_ex/chat.ex          |   2 +-
 lib/llama_cpp_ex/embedding.ex     |  87 +++----
 lib/llama_cpp_ex/hub.ex           |   4 +-
 lib/llama_cpp_ex/model_manager.ex | 148 +++++------
 lib/llama_cpp_ex/nif.ex           |  10 +
 lib/llama_cpp_ex/server.ex        | 205 +++++-----------
 lib/llama_cpp_ex/tokenizer.ex     |   4 +-
 test/llama_cpp_ex_test.exs        |   9 +-
 9 files changed, 311 insertions(+), 552 deletions(-)

diff --git a/lib/llama_cpp_ex.ex b/lib/llama_cpp_ex.ex
index 86542a8..6990631 100644
--- a/lib/llama_cpp_ex.ex
+++ b/lib/llama_cpp_ex.ex
@@ -63,6 +63,24 @@ defmodule LlamaCppEx do
     :swa_full
   ]
 
+  # Sampling options forwarded to Sampler.create/2 by the generation entry
+  # points. Keep in sync with the options documented on generate/3.
+  @sampler_opt_keys [
+    :seed,
+    :temp,
+    :top_k,
+    :top_p,
+    :min_p,
+    :penalty_repeat,
+    :penalty_freq,
+    :penalty_present,
+    :grammar,
+    :grammar_root
+  ]
+
+  # Chat-templating options split off before the rest flows to generation.
+  @chat_opt_keys [:add_assistant, :enable_thinking, :chat_template_kwargs]
+
   @doc """
   Initializes the llama.cpp backend. Call once at application start.
   """
@@ -170,38 +188,11 @@ defmodule LlamaCppEx do
   """
   @spec generate(Model.t(), String.t(), keyword()) :: {:ok, String.t()} | {:error, String.t()}
   def generate(%Model{} = model, prompt, opts \\ []) when is_binary(prompt) do
-    opts = resolve_grammar_opts(opts)
-    max_tokens = Keyword.get(opts, :max_tokens, 256)
-    n_ctx = Keyword.get(opts, :n_ctx, 2048)
-
-    sampler_opts =
-      Keyword.take(opts, [
-        :seed,
-        :temp,
-        :top_k,
-        :top_p,
-        :min_p,
-        :penalty_repeat,
-        :penalty_freq,
-        :penalty_present,
-        :grammar,
-        :grammar_root
-      ])
-
-    # Tokenize prompt
+    cfg = opts |> resolve_grammar_opts() |> gen_config()
     {:ok, tokens} = Tokenizer.encode(model, prompt)
 
-    # Ensure context is large enough for prompt + generation
-    ctx_size = max(n_ctx, length(tokens) + max_tokens)
-
-    ctx_opts =
-      opts
-      |> Keyword.take([:n_threads, :n_threads_batch, :n_batch, :n_ubatch])
-      |> Keyword.put(:n_ctx, ctx_size)
-
-    with {:ok, ctx} <- Context.create(model, ctx_opts),
-         {:ok, sampler} <- Sampler.create(model, sampler_opts) do
-      Context.generate(ctx, sampler, tokens, max_tokens: max_tokens)
+    with {:ok, ctx, sampler} <- create_gen_resources(model, tokens, cfg) do
+      Context.generate(ctx, sampler, tokens, max_tokens: cfg.max_tokens)
     end
   end
 
@@ -222,72 +213,75 @@ defmodule LlamaCppEx do
   """
   @spec stream(Model.t(), String.t(), keyword()) :: Enumerable.t()
   def stream(%Model{} = model, prompt, opts \\ []) when is_binary(prompt) do
-    opts = resolve_grammar_opts(opts)
-    max_tokens = Keyword.get(opts, :max_tokens, 256)
-    n_ctx = Keyword.get(opts, :n_ctx, 2048)
-    timeout = Keyword.get(opts, :timeout, 60_000)
-
-    sampler_opts =
-      Keyword.take(opts, [
-        :seed,
-        :temp,
-        :top_k,
-        :top_p,
-        :min_p,
-        :penalty_repeat,
-        :penalty_freq,
-        :penalty_present,
-        :grammar,
-        :grammar_root
-      ])
-
-    ctx_opts =
-      Keyword.take(opts, @context_opt_keys)
+    cfg = opts |> resolve_grammar_opts() |> gen_config()
 
     Stream.resource(
       fn ->
         # Start: tokenize, create context+sampler, spawn generator
         {:ok, tokens} = Tokenizer.encode(model, prompt)
-        ctx_size = max(n_ctx, length(tokens) + max_tokens)
-        {:ok, ctx} = Context.create(model, Keyword.put(ctx_opts, :n_ctx, ctx_size))
-        {:ok, sampler} = Sampler.create(model, sampler_opts)
-
-        ref = make_ref()
-        parent = self()
-
-        gen_pid =
-          spawn_link(fn ->
-            LlamaCppEx.NIF.generate_tokens(
-              ctx.ref,
-              sampler.ref,
-              tokens,
-              max_tokens,
-              parent,
-              ref
-            )
-          end)
-
-        {ref, gen_pid, timeout}
+        {:ok, ctx, sampler} = create_gen_resources(model, tokens, cfg)
+        {ref, gen_pid} = spawn_generator(ctx, sampler, tokens, cfg.max_tokens)
+        {ref, gen_pid, cfg.timeout}
       end,
       fn {ref, _gen_pid, timeout} = state ->
         receive do
           {^ref, {:token, _id, text}} -> {[text], state}
-          {^ref, :eog} -> {:halt, state}
-          {^ref, :done} -> {:halt, state}
+          {^ref, outcome} when outcome in [:eog, :done] -> {:halt, state}
           {^ref, {:error, _reason}} -> {:halt, state}
         after
           timeout -> {:halt, state}
         end
       end,
-      fn {ref, gen_pid, _timeout} ->
-        # Kill generator if still running, flush remaining messages
-        Process.unlink(gen_pid)
-        Process.exit(gen_pid, :kill)
-        flush_stream_messages(ref)
-      end
+      fn {ref, gen_pid, _timeout} -> stop_generator(ref, gen_pid) end
     )
   end
 
+  # --- Shared generation plumbing ---
+
+  # Option handling shared by the generation entry points.
+  defp gen_config(opts) do
+    %{
+      max_tokens: Keyword.get(opts, :max_tokens, 256),
+      n_ctx: Keyword.get(opts, :n_ctx, 2048),
+      timeout: Keyword.get(opts, :timeout, 60_000),
+      sampler_opts: Keyword.take(opts, @sampler_opt_keys),
+      ctx_opts: Keyword.take(opts, @context_opt_keys)
+    }
+  end
+
+  defp split_chat_opts(opts), do: Keyword.split(opts, @chat_opt_keys)
+
+  # Creates a context sized to fit prompt + generation, plus its sampler.
+  defp create_gen_resources(model, tokens, cfg) do
+    ctx_size = max(cfg.n_ctx, length(tokens) + cfg.max_tokens)
+
+    with {:ok, ctx} <- Context.create(model, Keyword.put(cfg.ctx_opts, :n_ctx, ctx_size)),
+         {:ok, sampler} <- Sampler.create(model, cfg.sampler_opts) do
+      {:ok, ctx, sampler}
+    end
+  end
+
+  # Runs the NIF token generator in a linked process that sends
+  # {ref, {:token, id, text} | :eog | :done | {:error, reason}} to the caller.
+  defp spawn_generator(ctx, sampler, tokens, max_tokens) do
+    ref = make_ref()
+    parent = self()
+
+    gen_pid =
+      spawn_link(fn ->
+        LlamaCppEx.NIF.generate_tokens(ctx.ref, sampler.ref, tokens, max_tokens, parent, ref)
+      end)
+
+    {ref, gen_pid}
+  end
+
+  # Kills a generator (if still running) and drains its remaining messages.
+  defp stop_generator(ref, gen_pid) do
+    Process.unlink(gen_pid)
+    Process.exit(gen_pid, :kill)
+    flush_stream_messages(ref)
+  end
+
   defp flush_stream_messages(ref) do
     receive do
       {^ref, _} -> flush_stream_messages(ref)
@@ -296,6 +290,10 @@ defmodule LlamaCppEx do
     end
   end
 
+  # OpenAI-style finish_reason for a generator outcome message.
+  defp finish_reason(:eog), do: "stop"
+  defp finish_reason(:done), do: "length"
+
   @doc """
   Applies the chat template and generates a response.
 
@@ -315,11 +313,7 @@ defmodule LlamaCppEx do
   """
   @spec chat(Model.t(), [Chat.message()], keyword()) :: {:ok, String.t()} | {:error, String.t()}
   def chat(%Model{} = model, messages, opts \\ []) when is_list(messages) do
-    opts = resolve_grammar_opts(opts)
-
-    {chat_opts, gen_opts} =
-      Keyword.split(opts, [:add_assistant, :enable_thinking, :chat_template_kwargs])
-
+    {chat_opts, gen_opts} = opts |> resolve_grammar_opts() |> split_chat_opts()
     {:ok, prompt} = Chat.apply_template(model, messages, chat_opts)
     generate(model, prompt, gen_opts)
   end
@@ -332,11 +326,7 @@ defmodule LlamaCppEx do
   """
   @spec stream_chat(Model.t(), [Chat.message()], keyword()) :: Enumerable.t()
   def stream_chat(%Model{} = model, messages, opts \\ []) when is_list(messages) do
-    opts = resolve_grammar_opts(opts)
-
-    {chat_opts, gen_opts} =
-      Keyword.split(opts, [:add_assistant, :enable_thinking, :chat_template_kwargs])
-
+    {chat_opts, gen_opts} = opts |> resolve_grammar_opts() |> split_chat_opts()
     {:ok, prompt} = Chat.apply_template(model, messages, chat_opts)
     stream(model, prompt, gen_opts)
   end
@@ -365,53 +355,19 @@ defmodule LlamaCppEx do
   @spec chat_completion(Model.t(), [Chat.message()], keyword()) ::
           {:ok, ChatCompletion.t()} | {:error, term()}
   def chat_completion(%Model{} = model, messages, opts \\ []) when is_list(messages) do
-    opts = resolve_grammar_opts(opts)
-
-    {chat_opts, gen_opts} =
-      Keyword.split(opts, [:add_assistant, :enable_thinking, :chat_template_kwargs])
-
-    max_tokens = Keyword.get(gen_opts, :max_tokens, 256)
-    n_ctx = Keyword.get(gen_opts, :n_ctx, 2048)
-    timeout = Keyword.get(gen_opts, :timeout, 60_000)
-
-    sampler_opts =
-      Keyword.take(gen_opts, [
-        :seed,
-        :temp,
-        :top_k,
-        :top_p,
-        :min_p,
-        :penalty_repeat,
-        :penalty_freq,
-        :penalty_present,
-        :grammar,
-        :grammar_root
-      ])
-
-    ctx_opts =
-      Keyword.take(gen_opts, @context_opt_keys)
+    {chat_opts, gen_opts} = opts |> resolve_grammar_opts() |> split_chat_opts()
+    cfg = gen_config(gen_opts)
 
     with {:ok, prompt} <- Chat.apply_template(model, messages, chat_opts),
          {:ok, prompt_tokens} <- Tokenizer.encode(model, prompt) do
-      ctx_size = max(n_ctx, length(prompt_tokens) + max_tokens)
-      {:ok, ctx} = Context.create(model, Keyword.put(ctx_opts, :n_ctx, ctx_size))
-      {:ok, sampler} = Sampler.create(model, sampler_opts)
+      {:ok, ctx, sampler} = create_gen_resources(model, prompt_tokens, cfg)
+      {ref, gen_pid} = spawn_generator(ctx, sampler, prompt_tokens, cfg.max_tokens)
 
-      ref = make_ref()
-      parent = self()
+      {texts, finish_reason, completion_tokens} = collect_completion_tokens(ref, cfg.timeout)
 
-      spawn_link(fn ->
-        LlamaCppEx.NIF.generate_tokens(
-          ctx.ref,
-          sampler.ref,
-          prompt_tokens,
-          max_tokens,
-          parent,
-          ref
-        )
-      end)
-
-      {texts, finish_reason, completion_tokens} = collect_completion_tokens(ref, timeout)
+      # On timeout the generator may still be running — kill it and drain any
+      # stragglers so they don't pollute the caller's mailbox.
+      stop_generator(ref, gen_pid)
 
       raw_text = Enum.join(texts)
       enable_thinking = Keyword.get(chat_opts, :enable_thinking, false)
@@ -475,68 +431,25 @@ defmodule LlamaCppEx do
   """
   @spec stream_chat_completion(Model.t(), [Chat.message()], keyword()) :: Enumerable.t()
   def stream_chat_completion(%Model{} = model, messages, opts \\ []) when is_list(messages) do
-    opts = resolve_grammar_opts(opts)
-
-    {chat_opts, gen_opts} =
-      Keyword.split(opts, [:add_assistant, :enable_thinking, :chat_template_kwargs])
-
-    max_tokens = Keyword.get(gen_opts, :max_tokens, 256)
-    n_ctx = Keyword.get(gen_opts, :n_ctx, 2048)
-    timeout = Keyword.get(gen_opts, :timeout, 60_000)
-
-    sampler_opts =
-      Keyword.take(gen_opts, [
-        :seed,
-        :temp,
-        :top_k,
-        :top_p,
-        :min_p,
-        :penalty_repeat,
-        :penalty_freq,
-        :penalty_present,
-        :grammar,
-        :grammar_root
-      ])
-
-    ctx_opts =
-      Keyword.take(gen_opts, @context_opt_keys)
+    {chat_opts, gen_opts} = opts |> resolve_grammar_opts() |> split_chat_opts()
+    cfg = gen_config(gen_opts)
 
     Stream.resource(
       fn ->
         {:ok, prompt} = Chat.apply_template(model, messages, chat_opts)
         {:ok, tokens} = Tokenizer.encode(model, prompt)
-        ctx_size = max(n_ctx, length(tokens) + max_tokens)
-        {:ok, ctx} = Context.create(model, Keyword.put(ctx_opts, :n_ctx, ctx_size))
-        {:ok, sampler} = Sampler.create(model, sampler_opts)
-
-        id = "chatcmpl-" <> random_hex(12)
-        created = System.os_time(:second)
-        model_name = Model.desc(model)
-
-        ref = make_ref()
-        parent = self()
-
-        gen_pid =
-          spawn_link(fn ->
-            LlamaCppEx.NIF.generate_tokens(
-              ctx.ref,
-              sampler.ref,
-              tokens,
-              max_tokens,
-              parent,
-              ref
-            )
-          end)
+        {:ok, ctx, sampler} = create_gen_resources(model, tokens, cfg)
+        {ref, gen_pid} = spawn_generator(ctx, sampler, tokens, cfg.max_tokens)
 
         enable_thinking = Keyword.get(chat_opts, :enable_thinking, false)
 
         %{
           ref: ref,
           gen_pid: gen_pid,
-          timeout: timeout,
-          id: id,
-          created: created,
-          model: model_name,
+          timeout: cfg.timeout,
+          id: "chatcmpl-" <> random_hex(12),
+          created: System.os_time(:second),
+          model: Model.desc(model),
           phase: :first,
           enable_thinking: enable_thinking,
           thinking_parser:
@@ -545,80 +458,15 @@ defmodule LlamaCppEx do
       end,
       fn
         %{phase: :first} = state ->
-          chunk = %ChatCompletionChunk{
-            id: state.id,
-            object: "chat.completion.chunk",
-            created: state.created,
-            model: state.model,
-            choices: [%{index: 0, delta: %{role: "assistant", content: ""}, finish_reason: nil}]
-          }
-
-          {[chunk], %{state | phase: :streaming}}
+          {[chunk(state, %{role: "assistant", content: ""}, nil)], %{state | phase: :streaming}}
 
         %{phase: :streaming, ref: ref, timeout: timeout} = state ->
           receive do
             {^ref, {:token, _id, text}} ->
-              if state.enable_thinking do
-                {events, new_parser} = Thinking.feed(state.thinking_parser, text)
-                state = %{state | thinking_parser: new_parser}
-
-                chunks =
-                  Enum.map(events, fn
-                    {:thinking, t} ->
-                      %ChatCompletionChunk{
-                        id: state.id,
-                        object: "chat.completion.chunk",
-                        created: state.created,
-                        model: state.model,
-                        choices: [
-                          %{index: 0, delta: %{reasoning_content: t}, finish_reason: nil}
-                        ]
-                      }
-
-                    {:content, t} ->
-                      %ChatCompletionChunk{
-                        id: state.id,
-                        object: "chat.completion.chunk",
-                        created: state.created,
-                        model: state.model,
-                        choices: [%{index: 0, delta: %{content: t}, finish_reason: nil}]
-                      }
-                  end)
-
-                {chunks, state}
-              else
-                chunk = %ChatCompletionChunk{
-                  id: state.id,
-                  object: "chat.completion.chunk",
-                  created: state.created,
-                  model: state.model,
-                  choices: [%{index: 0, delta: %{content: text}, finish_reason: nil}]
-                }
-
-                {[chunk], state}
-              end
-
-            {^ref, :eog} ->
-              final_chunk = %ChatCompletionChunk{
-                id: state.id,
-                object: "chat.completion.chunk",
-                created: state.created,
-                model: state.model,
-                choices: [%{index: 0, delta: %{}, finish_reason: "stop"}]
-              }
-
-              {[final_chunk], %{state | phase: :done}}
-
-            {^ref, :done} ->
-              final_chunk = %ChatCompletionChunk{
-                id: state.id,
-                object: "chat.completion.chunk",
-                created: state.created,
-                model: state.model,
-                choices: [%{index: 0, delta: %{}, finish_reason: "length"}]
-              }
-
-              {[final_chunk], %{state | phase: :done}}
+              token_chunks(state, text)
+
+            {^ref, outcome} when outcome in [:eog, :done] ->
+              {[chunk(state, %{}, finish_reason(outcome))], %{state | phase: :done}}
 
             {^ref, {:error, _reason}} ->
               {:halt, state}
@@ -629,14 +477,39 @@ defmodule LlamaCppEx do
         %{phase: :done} = state ->
           {:halt, state}
       end,
-      fn %{ref: ref, gen_pid: gen_pid} ->
-        Process.unlink(gen_pid)
-        Process.exit(gen_pid, :kill)
-        flush_stream_messages(ref)
-      end
+      fn %{ref: ref, gen_pid: gen_pid} -> stop_generator(ref, gen_pid) end
     )
   end
 
+  # Emits the chunk(s) for one generated token, routing through the thinking
+  # parser when enabled (one token can yield reasoning and content chunks).
+  defp token_chunks(%{enable_thinking: true} = state, text) do
+    {events, new_parser} = Thinking.feed(state.thinking_parser, text)
+    state = %{state | thinking_parser: new_parser}
+
+    chunks =
+      Enum.map(events, fn
+        {:thinking, t} -> chunk(state, %{reasoning_content: t}, nil)
+        {:content, t} -> chunk(state, %{content: t}, nil)
+      end)
+
+    {chunks, state}
+  end
+
+  defp token_chunks(state, text) do
+    {[chunk(state, %{content: text}, nil)], state}
+  end
+
+  defp chunk(state, delta, finish_reason) do
+    %ChatCompletionChunk{
+      id: state.id,
+      object: "chat.completion.chunk",
+      created: state.created,
+      model: state.model,
+      choices: [%{index: 0, delta: delta, finish_reason: finish_reason}]
+    }
+  end
+
   defp collect_completion_tokens(ref, timeout) do
     collect_completion_tokens(ref, timeout, [], 0)
   end
@@ -646,11 +519,8 @@ defmodule LlamaCppEx do
       {^ref, {:token, _id, text}} ->
         collect_completion_tokens(ref, timeout, [text | texts], count + 1)
 
-      {^ref, :eog} ->
-        {Enum.reverse(texts), "stop", count}
-
-      {^ref, :done} ->
-        {Enum.reverse(texts), "length", count}
+      {^ref, outcome} when outcome in [:eog, :done] ->
+        {Enum.reverse(texts), finish_reason(outcome), count}
 
       {^ref, {:error, _reason}} ->
         {Enum.reverse(texts), "stop", count}
diff --git a/lib/llama_cpp_ex/chat.ex b/lib/llama_cpp_ex/chat.ex
index 6ef4c2b..c55da24 100644
--- a/lib/llama_cpp_ex/chat.ex
+++ b/lib/llama_cpp_ex/chat.ex
@@ -63,7 +63,7 @@ defmodule LlamaCppEx.Chat do
 
       {:ok, result}
     rescue
-      e in ErlangError -> {:error, "chat template failed: #{inspect(e.original)}"}
+      e in ErlangError -> LlamaCppEx.NIF.error_tuple(e, "chat template", __STACKTRACE__)
     end
   end
 end
diff --git a/lib/llama_cpp_ex/embedding.ex b/lib/llama_cpp_ex/embedding.ex
index 59b2d14..bfd0a43 100644
--- a/lib/llama_cpp_ex/embedding.ex
+++ b/lib/llama_cpp_ex/embedding.ex
@@ -73,7 +73,7 @@ defmodule LlamaCppEx.Embedding do
 
     # texts is non-empty here (embed_batch handles the [] case), so tokenized
     # and groups are non-empty and Enum.max/1 is safe.
-    with {:ok, tokenized} <- tokenize_all(model, texts) do
+    with {:ok, tokenized} <- map_while_ok(texts, &Tokenizer.encode(model, &1)) do
       longest = tokenized |> Enum.map(&length/1) |> Enum.max()
       budget = max(n_ctx, longest + 8)
       groups = group_by_budget(tokenized, budget, max_seqs)
@@ -91,20 +91,6 @@ defmodule LlamaCppEx.Embedding do
     end
   end
 
-  # Tokenizes every text up front, short-circuiting on the first error.
-  defp tokenize_all(model, texts) do
-    Enum.reduce_while(texts, {:ok, []}, fn text, {:ok, acc} ->
-      case Tokenizer.encode(model, text) do
-        {:ok, tokens} -> {:cont, {:ok, [tokens | acc]}}
-        {:error, _} = err -> {:halt, err}
-      end
-    end)
-    |> case do
-      {:ok, acc} -> {:ok, Enum.reverse(acc)}
-      {:error, _} = err -> err
-    end
-  end
-
   # Greedy bin-packing: each group's total tokens stays within `budget` and its
   # size within `max_seqs`. A single text longer than the budget gets its own
   # group (the context is sized to fit the longest text).
@@ -133,66 +119,49 @@ defmodule LlamaCppEx.Embedding do
   # Decodes each group in one batch and extracts per-sequence embeddings,
   # preserving the original text order.
   defp decode_groups(ctx, groups, normalize) do
-    Enum.reduce_while(groups, {:ok, []}, fn group, {:ok, acc} ->
-      sequences = group |> Enum.with_index() |> Enum.map(fn {tokens, i} -> {i, tokens} end)
-
-      with :ok <- embed_batch_decode(ctx, sequences),
-           {:ok, embs} <- collect_group_embeddings(ctx, length(group), normalize) do
-        {:cont, {:ok, acc ++ embs}}
-      else
-        {:error, _} = err -> {:halt, err}
-      end
-    end)
+    with {:ok, nested} <- map_while_ok(groups, &decode_group(ctx, &1, normalize)) do
+      {:ok, Enum.concat(nested)}
+    end
   end
 
-  defp collect_group_embeddings(ctx, count, normalize) do
-    Enum.reduce_while(0..(count - 1)//1, {:ok, []}, fn seq_id, {:ok, acc} ->
-      case get_embeddings(ctx, seq_id, normalize) do
-        {:ok, emb} -> {:cont, {:ok, [emb | acc]}}
-        {:error, _} = err -> {:halt, err}
-      end
-    end)
-    |> case do
-      {:ok, acc} -> {:ok, Enum.reverse(acc)}
-      {:error, _} = err -> err
+  defp decode_group(ctx, group, normalize) do
+    sequences = Enum.with_index(group, fn tokens, i -> {i, tokens} end)
+
+    with :ok <- embed_batch_decode(ctx, sequences) do
+      map_while_ok(0..(length(group) - 1)//1, &get_embeddings(ctx, &1, normalize))
     end
   end
 
   # --- Sequential fallback (one context per text) ---
 
   defp embed_batch_sequential(model, texts, opts) do
-    Enum.reduce_while(texts, {:ok, []}, fn text, {:ok, acc} ->
-      case embed(model, text, opts) do
-        {:ok, emb} -> {:cont, {:ok, [emb | acc]}}
+    map_while_ok(texts, &embed(model, &1, opts))
+  end
+
+  # Maps `fun` (returning {:ok, value} | {:error, _}) over `enum`, preserving
+  # order and short-circuiting on the first error.
+  defp map_while_ok(enum, fun) do
+    enum
+    |> Enum.reduce_while({:ok, []}, fn item, {:ok, acc} ->
+      case fun.(item) do
+        {:ok, value} -> {:cont, {:ok, [value | acc]}}
         {:error, _} = err -> {:halt, err}
       end
     end)
-    |> case do
-      {:ok, embeddings} -> {:ok, Enum.reverse(embeddings)}
+    |> then(fn
+      {:ok, acc} -> {:ok, Enum.reverse(acc)}
       {:error, _} = err -> err
-    end
+    end)
   end
 
   # --- NIF wrappers ---
 
-  defp embed_decode(%Context{ref: ref}, tokens, seq_id) do
-    case LlamaCppEx.NIF.embed_decode(ref, tokens, seq_id) do
-      :ok -> :ok
-      {:error, _} = err -> err
-    end
-  end
+  defp embed_decode(%Context{ref: ref}, tokens, seq_id),
+    do: LlamaCppEx.NIF.embed_decode(ref, tokens, seq_id)
 
-  defp embed_batch_decode(%Context{ref: ref}, sequences) do
-    case LlamaCppEx.NIF.embed_batch_decode(ref, sequences) do
-      :ok -> :ok
-      {:error, _} = err -> err
-    end
-  end
+  defp embed_batch_decode(%Context{ref: ref}, sequences),
+    do: LlamaCppEx.NIF.embed_batch_decode(ref, sequences)
 
-  defp get_embeddings(%Context{ref: ref}, seq_id, normalize) do
-    case LlamaCppEx.NIF.get_embeddings(ref, seq_id, normalize) do
-      {:ok, _} = result -> result
-      {:error, _} = err -> err
-    end
-  end
+  defp get_embeddings(%Context{ref: ref}, seq_id, normalize),
+    do: LlamaCppEx.NIF.get_embeddings(ref, seq_id, normalize)
 end
diff --git a/lib/llama_cpp_ex/hub.ex b/lib/llama_cpp_ex/hub.ex
index 48b7a20..72cc530 100644
--- a/lib/llama_cpp_ex/hub.ex
+++ b/lib/llama_cpp_ex/hub.ex
@@ -496,7 +496,9 @@ defmodule LlamaCppEx.Hub do
           {:error, reason}
       end
     rescue
-      e ->
+      # File.Error from rename!/write!/the File.stream! sink, ErlangError from
+      # lower-level IO — programming errors (KeyError, MatchError, …) propagate.
+      e in [File.Error, ErlangError] ->
         File.rm(tmp_dest)
         {:error, "download failed: #{Exception.message(e)}"}
     end
diff --git a/lib/llama_cpp_ex/model_manager.ex b/lib/llama_cpp_ex/model_manager.ex
index b7dada6..4627eef 100644
--- a/lib/llama_cpp_ex/model_manager.ex
+++ b/lib/llama_cpp_ex/model_manager.ex
@@ -196,18 +196,11 @@ defmodule LlamaCppEx.ModelManager do
   """
   @spec generate(id(), String.t(), keyword()) :: {:ok, String.t()} | {:error, term()}
   def generate(id, prompt, opts \\ []) do
-    case route(id) do
-      {:ok, {:server, pid, e}} ->
-        touch(e.id)
-        LlamaCppEx.Server.generate(pid, prompt, opts)
-
-      {:ok, {:direct, model, e}} ->
-        touch(e.id)
-        LlamaCppEx.generate(model, prompt, opts)
-
-      {:error, _} = err ->
-        err
-    end
+    with_route(
+      id,
+      &LlamaCppEx.Server.generate(&1, prompt, opts),
+      &LlamaCppEx.generate(&1, prompt, opts)
+    )
   end
 
   @doc """
@@ -218,17 +211,16 @@ defmodule LlamaCppEx.ModelManager do
   """
   @spec stream(id(), String.t(), keyword()) :: Enumerable.t()
   def stream(id, prompt, opts \\ []) do
-    case route(id) do
-      {:ok, {:server, pid, e}} ->
-        touch(e.id)
-        LlamaCppEx.Server.stream(pid, prompt, opts)
-
-      {:ok, {:direct, model, e}} ->
-        touch(e.id)
-        LlamaCppEx.stream(model, prompt, opts)
-
+    case with_route(
+           id,
+           &LlamaCppEx.Server.stream(&1, prompt, opts),
+           &LlamaCppEx.stream(&1, prompt, opts)
+         ) do
       {:error, reason} ->
         raise ArgumentError, "cannot stream from model #{inspect(id)}: #{inspect(reason)}"
+
+      stream ->
+        stream
     end
   end
 
@@ -236,25 +228,23 @@ defmodule LlamaCppEx.ModelManager do
   @spec chat(id(), [LlamaCppEx.Chat.message()], keyword()) ::
           {:ok, String.t()} | {:error, term()}
   def chat(id, messages, opts \\ []) do
-    case route(id) do
-      {:ok, {:server, pid, e}} ->
-        touch(e.id)
+    with_route(
+      id,
+      &server_chat(&1, messages, opts),
+      &LlamaCppEx.chat(&1, messages, opts)
+    )
+  end
 
-        {chat_opts, gen_opts} =
-          Keyword.split(opts, [:add_assistant, :enable_thinking, :chat_template_kwargs])
+  # A server-backed model exposes generate/stream only, so chat templating
+  # happens here before handing the rendered prompt to the server.
+  defp server_chat(pid, messages, opts) do
+    {chat_opts, gen_opts} =
+      Keyword.split(opts, [:add_assistant, :enable_thinking, :chat_template_kwargs])
 
-        model = LlamaCppEx.Server.get_model(pid)
+    model = LlamaCppEx.Server.get_model(pid)
 
-        with {:ok, prompt} <- LlamaCppEx.Chat.apply_template(model, messages, chat_opts) do
-          LlamaCppEx.Server.generate(pid, prompt, gen_opts)
-        end
-
-      {:ok, {:direct, model, e}} ->
-        touch(e.id)
-        LlamaCppEx.chat(model, messages, opts)
-
-      {:error, _} = err ->
-        err
+    with {:ok, prompt} <- LlamaCppEx.Chat.apply_template(model, messages, chat_opts) do
+      LlamaCppEx.Server.generate(pid, prompt, gen_opts)
     end
   end
 
@@ -304,6 +294,23 @@ defmodule LlamaCppEx.ModelManager do
     end
   end
 
+  # Shared dispatch skeleton: resolve the route, mark the model used, and hand
+  # the server pid or direct model to the matching callback. Errors pass through.
+  defp with_route(id, server_fun, direct_fun) do
+    case route(id) do
+      {:ok, {:server, pid, e}} ->
+        touch(e.id)
+        server_fun.(pid)
+
+      {:ok, {:direct, model, e}} ->
+        touch(e.id)
+        direct_fun.(model)
+
+      {:error, _} = err ->
+        err
+    end
+  end
+
   # --- Server callbacks ---
 
   @impl true
@@ -595,24 +602,10 @@ defmodule LlamaCppEx.ModelManager do
     end
   end
 
-  defp build_ready_entry(
-         state,
-         id,
-         source,
-         {:server, pid},
-         file_bytes,
-         placement,
-         capabilities,
-         opts
-       ) do
-    ref = Process.monitor(pid)
-
-    entry = %Entry{
+  defp build_ready_entry(state, id, source, backing, file_bytes, placement, capabilities, opts) do
+    base = [
       id: id,
       status: :ready,
-      mode: :server,
-      server_pid: pid,
-      monitor_ref: ref,
       source: source,
       capabilities: capabilities,
       byte_size: file_bytes,
@@ -620,36 +613,17 @@ defmodule LlamaCppEx.ModelManager do
       placement: placement,
       n_gpu_layers: Keyword.get(opts, :n_gpu_layers, 99),
       loaded_at: System.system_time(:second)
-    }
-
-    {entry, %{state | monitors: Map.put(state.monitors, ref, id)}}
-  end
-
-  defp build_ready_entry(
-         state,
-         id,
-         source,
-         {:direct, model},
-         file_bytes,
-         placement,
-         capabilities,
-         opts
-       ) do
-    entry = %Entry{
-      id: id,
-      status: :ready,
-      mode: :direct,
-      model: model,
-      source: source,
-      capabilities: capabilities,
-      byte_size: file_bytes,
-      est_bytes: placement_total(placement),
-      placement: placement,
-      n_gpu_layers: Keyword.get(opts, :n_gpu_layers, 99),
-      loaded_at: System.system_time(:second)
-    }
+    ]
+
+    case backing do
+      {:server, pid} ->
+        ref = Process.monitor(pid)
+        entry = struct!(Entry, [mode: :server, server_pid: pid, monitor_ref: ref] ++ base)
+        {entry, %{state | monitors: Map.put(state.monitors, ref, id)}}
 
-    {entry, state}
+      {:direct, model} ->
+        {struct!(Entry, [mode: :direct, model: model] ++ base), state}
+    end
   end
 
   defp placement_total(%{ram: ram, vram: vram}), do: ram + Enum.sum(Map.values(vram))
@@ -770,11 +744,15 @@ defmodule LlamaCppEx.ModelManager do
     ArgumentError -> []
   end
 
+  # Best-effort: the NIF may be missing entirely (test env / unsupported
+  # platform) — degrade to :ok like gpu_devices/0, but log what was skipped.
+  # init/0 is a thin NIF call, so only these two exceptions can occur;
+  # anything else propagates.
   defp safe_backend_init do
     LlamaCppEx.init()
   rescue
-    _ -> :ok
-  catch
-    _, _ -> :ok
+    e in [ErlangError, UndefinedFunctionError] ->
+      Logger.debug("backend init skipped: #{Exception.message(e)}")
+      :ok
   end
 end
diff --git a/lib/llama_cpp_ex/nif.ex b/lib/llama_cpp_ex/nif.ex
index 3ea5d5e..815f15b 100644
--- a/lib/llama_cpp_ex/nif.ex
+++ b/lib/llama_cpp_ex/nif.ex
@@ -7,6 +7,16 @@ defmodule LlamaCppEx.NIF do
     :erlang.load_nif(path, 0)
   end
 
+  # Converts an ErlangError raised by a NIF into an error tuple. The
+  # NIF-not-loaded error is re-raised: it signals a build/packaging problem,
+  # not bad input, and must not be flattened into a caller-visible string.
+  @doc false
+  def error_tuple(%ErlangError{original: :not_loaded} = e, _label, stacktrace),
+    do: reraise(e, stacktrace)
+
+  def error_tuple(%ErlangError{original: original}, label, _stacktrace),
+    do: {:error, "#{label} failed: #{inspect(original)}"}
+
   # Backend
   def backend_init, do: :erlang.nif_error(:not_loaded)
   def backend_free, do: :erlang.nif_error(:not_loaded)
diff --git a/lib/llama_cpp_ex/server.ex b/lib/llama_cpp_ex/server.ex
index 829cc1c..9f63092 100644
--- a/lib/llama_cpp_ex/server.ex
+++ b/lib/llama_cpp_ex/server.ex
@@ -371,33 +371,7 @@ defmodule LlamaCppEx.Server do
     slots =
       for seq_id <- 0..(n_parallel - 1), into: %{} do
         {:ok, sampler} = Sampler.create(model, sampler_opts)
-
-        slot = %{
-          state: :idle,
-          sampler: sampler,
-          from: nil,
-          stream_pid: nil,
-          stream_ref: nil,
-          prompt_tokens: [],
-          prompt_tokens_tuple: {},
-          prefill_pos: 0,
-          pos: 0,
-          pending_token: nil,
-          batch_idx: -1,
-          tokens_generated: 0,
-          max_tokens: 0,
-          accumulated_pieces: [],
-          t_start: nil,
-          t_first_token: nil,
-          n_prompt_tokens: 0,
-          # Prefix cache fields
-          cached_tokens: [],
-          cached_pos: 0,
-          generated_token_ids: [],
-          n_prefix_cache_tokens: 0
-        }
-
-        {seq_id, slot}
+        {seq_id, Map.put(idle_slot_fields([], 0), :sampler, sampler)}
       end
 
     state = %__MODULE__{
@@ -869,59 +843,40 @@ defmodule LlamaCppEx.Server do
   end
 
   defp emit_request_done(slot, seq_id, t_end, stop_reason) do
-    duration_ns = t_end - slot.t_start
-    duration_ms = duration_ns / 1_000_000
-
-    ttft_ms =
-      if slot.t_first_token do
-        (slot.t_first_token - slot.t_start) / 1_000_000
-      else
-        duration_ms
-      end
-
-    gen_duration_s = (t_end - (slot.t_first_token || slot.t_start)) / 1_000_000_000
-    prompt_duration_s = ttft_ms / 1000
-
-    prompt_eval_rate =
-      if prompt_duration_s > 0, do: slot.n_prompt_tokens / prompt_duration_s, else: 0.0
-
-    generation_rate =
-      if gen_duration_s > 0, do: slot.tokens_generated / gen_duration_s, else: 0.0
-
-    mode = slot_mode(slot)
+    m = request_measurements(slot, t_end)
 
     Logger.debug(
-      "slot #{seq_id} done (#{stop_reason}): #{slot.n_prompt_tokens} prompt tokens (#{Float.round(prompt_eval_rate, 1)} t/s), " <>
-        "#{slot.tokens_generated} generated (#{Float.round(generation_rate, 1)} t/s), " <>
-        "ttft #{Float.round(ttft_ms, 1)}ms, total #{Float.round(duration_ms, 1)}ms"
+      "slot #{seq_id} done (#{stop_reason}): #{slot.n_prompt_tokens} prompt tokens (#{Float.round(m.prompt_eval_rate, 1)} t/s), " <>
+        "#{slot.tokens_generated} generated (#{Float.round(m.generation_rate, 1)} t/s), " <>
+        "ttft #{Float.round(m.ttft_ms, 1)}ms, total #{Float.round(m.duration_ms, 1)}ms"
     )
 
-    prefix_cache_ratio =
-      if slot.n_prompt_tokens > 0,
-        do: slot.n_prefix_cache_tokens / slot.n_prompt_tokens,
-        else: 0.0
-
     :telemetry.execute(
       [:llama_cpp_ex, :server, :request, :done],
-      %{
-        prompt_tokens: slot.n_prompt_tokens,
-        generated_tokens: slot.tokens_generated,
-        duration_ms: duration_ms,
-        ttft_ms: ttft_ms,
-        prompt_eval_rate: prompt_eval_rate,
-        generation_rate: generation_rate,
-        prefix_cache_tokens: slot.n_prefix_cache_tokens,
-        prefix_cache_ratio: prefix_cache_ratio
-      },
-      %{server: self(), seq_id: seq_id, mode: mode, stop_reason: stop_reason}
+      m,
+      %{server: self(), seq_id: seq_id, mode: slot_mode(slot), stop_reason: stop_reason}
     )
   end
 
   defp emit_request_exception(slot, seq_id, reason) do
-    # Mirror `:done`'s measurement shape so dashboards can aggregate them.
-    # Timings are best-effort: if generation never started, ttft/duration
-    # fall back to the wall time since slot acquisition.
-    t_end = System.monotonic_time()
+    # Mirrors `:done`'s measurement shape so dashboards can aggregate them.
+    :telemetry.execute(
+      [:llama_cpp_ex, :server, :request, :exception],
+      request_measurements(slot, System.monotonic_time()),
+      %{
+        server: self(),
+        seq_id: seq_id,
+        mode: slot_mode(slot),
+        stop_reason: :error,
+        reason: reason
+      }
+    )
+  end
+
+  # Shared measurements for the :done / :exception request events. Timings are
+  # best-effort: if generation never started, ttft/duration fall back to the
+  # wall time since slot acquisition.
+  defp request_measurements(slot, t_end) do
     t_start = slot.t_start || t_end
     duration_ms = (t_end - t_start) / 1_000_000
 
@@ -942,26 +897,16 @@ defmodule LlamaCppEx.Server do
         do: slot.n_prefix_cache_tokens / slot.n_prompt_tokens,
         else: 0.0
 
-    :telemetry.execute(
-      [:llama_cpp_ex, :server, :request, :exception],
-      %{
-        prompt_tokens: slot.n_prompt_tokens,
-        generated_tokens: slot.tokens_generated,
-        duration_ms: duration_ms,
-        ttft_ms: ttft_ms,
-        prompt_eval_rate: prompt_eval_rate,
-        generation_rate: generation_rate,
-        prefix_cache_tokens: slot.n_prefix_cache_tokens,
-        prefix_cache_ratio: prefix_cache_ratio
-      },
-      %{
-        server: self(),
-        seq_id: seq_id,
-        mode: slot_mode(slot),
-        stop_reason: :error,
-        reason: reason
-      }
-    )
+    %{
+      prompt_tokens: slot.n_prompt_tokens,
+      generated_tokens: slot.tokens_generated,
+      duration_ms: duration_ms,
+      ttft_ms: ttft_ms,
+      prompt_eval_rate: prompt_eval_rate,
+      generation_rate: generation_rate,
+      prefix_cache_tokens: slot.n_prefix_cache_tokens,
+      prefix_cache_ratio: prefix_cache_ratio
+    }
   end
 
   defp reset_slot(state, seq_id) do
@@ -978,33 +923,41 @@ defmodule LlamaCppEx.Server do
         {[], 0}
       end
 
-    slot = %{
-      slot
-      | state: :idle,
-        from: nil,
-        stream_pid: nil,
-        stream_ref: nil,
-        prompt_tokens: [],
-        prompt_tokens_tuple: {},
-        prefill_pos: 0,
-        pos: 0,
-        pending_token: nil,
-        batch_idx: -1,
-        tokens_generated: 0,
-        max_tokens: 0,
-        accumulated_pieces: [],
-        t_start: nil,
-        t_first_token: nil,
-        n_prompt_tokens: 0,
-        cached_tokens: cached_tokens,
-        cached_pos: cached_pos,
-        generated_token_ids: [],
-        n_prefix_cache_tokens: 0
-    }
+    slot = Map.merge(slot, idle_slot_fields(cached_tokens, cached_pos))
 
     put_in(state.slots[seq_id], slot)
   end
 
+  # The single source of truth for a slot's per-request fields. init/1,
+  # reset_slot/2, and fail_all_active_slots/2 all build from this map, so a
+  # new slot field cannot silently carry stale data across requests. The
+  # prefix-cache carry-over is the only caller-controlled part; :sampler is
+  # the only field that lives outside it.
+  defp idle_slot_fields(cached_tokens, cached_pos) do
+    %{
+      state: :idle,
+      from: nil,
+      stream_pid: nil,
+      stream_ref: nil,
+      prompt_tokens: [],
+      prompt_tokens_tuple: {},
+      prefill_pos: 0,
+      pos: 0,
+      pending_token: nil,
+      batch_idx: -1,
+      tokens_generated: 0,
+      max_tokens: 0,
+      accumulated_pieces: [],
+      t_start: nil,
+      t_first_token: nil,
+      n_prompt_tokens: 0,
+      cached_tokens: cached_tokens,
+      cached_pos: cached_pos,
+      generated_token_ids: [],
+      n_prefix_cache_tokens: 0
+    }
+  end
+
   # Builds the final completion string from the reverse-ordered piece list.
   defp accumulated_text(slot) do
     slot.accumulated_pieces |> Enum.reverse() |> IO.iodata_to_binary()
@@ -1029,31 +982,7 @@ defmodule LlamaCppEx.Server do
       LlamaCppEx.NIF.memory_seq_rm(state.ctx.ref, seq_id, 0, -1)
       Sampler.reset(slot.sampler)
 
-      slot = %{
-        slot
-        | state: :idle,
-          from: nil,
-          stream_pid: nil,
-          stream_ref: nil,
-          prompt_tokens: [],
-          prompt_tokens_tuple: {},
-          prefill_pos: 0,
-          pos: 0,
-          pending_token: nil,
-          batch_idx: -1,
-          tokens_generated: 0,
-          max_tokens: 0,
-          accumulated_pieces: [],
-          t_start: nil,
-          t_first_token: nil,
-          n_prompt_tokens: 0,
-          cached_tokens: [],
-          cached_pos: 0,
-          generated_token_ids: [],
-          n_prefix_cache_tokens: 0
-      }
-
-      put_in(state.slots[seq_id], slot)
+      put_in(state.slots[seq_id], Map.merge(slot, idle_slot_fields([], 0)))
     end)
   end
 
diff --git a/lib/llama_cpp_ex/tokenizer.ex b/lib/llama_cpp_ex/tokenizer.ex
index 77612ff..0f621dc 100644
--- a/lib/llama_cpp_ex/tokenizer.ex
+++ b/lib/llama_cpp_ex/tokenizer.ex
@@ -19,7 +19,7 @@ defmodule LlamaCppEx.Tokenizer do
     parse_special = Keyword.get(opts, :parse_special, true)
     {:ok, LlamaCppEx.NIF.tokenize(ref, text, add_special, parse_special)}
   rescue
-    e in ErlangError -> {:error, "tokenize failed: #{inspect(e.original)}"}
+    e in ErlangError -> LlamaCppEx.NIF.error_tuple(e, "tokenize", __STACKTRACE__)
   end
 
   @doc """
@@ -29,7 +29,7 @@ defmodule LlamaCppEx.Tokenizer do
   def decode(%LlamaCppEx.Model{ref: ref}, tokens) when is_list(tokens) do
     {:ok, LlamaCppEx.NIF.detokenize(ref, tokens)}
   rescue
-    e in ErlangError -> {:error, "detokenize failed: #{inspect(e.original)}"}
+    e in ErlangError -> LlamaCppEx.NIF.error_tuple(e, "detokenize", __STACKTRACE__)
   end
 
   @doc """
diff --git a/test/llama_cpp_ex_test.exs b/test/llama_cpp_ex_test.exs
index 7523956..0e6401e 100644
--- a/test/llama_cpp_ex_test.exs
+++ b/test/llama_cpp_ex_test.exs
@@ -981,10 +981,11 @@ defmodule LlamaCppExTest do
         %{model: model, mtp: mtp}
       end
 
-      test "draft context exposes rollback capacity", %{mtp: mtp} do
-        # n_rs_seq on the MTP draft ctx should be at least the configured n_draft;
-        # the default target ctx should report 0 (no rollback).
-        assert LlamaCppEx.Context.n_rs_seq(mtp.mtp_ctx) >= mtp.n_draft
+      test "contexts use no recurrent-state rollback slots", %{mtp: mtp} do
+        # Matching upstream server, the draft ctx is created with n_rs_seq=0 —
+        # MTP rolls back via cached hidden states (pending_h / verify_h), not
+        # recurrent-state snapshots. Both contexts report 0.
+        assert LlamaCppEx.Context.n_rs_seq(mtp.mtp_ctx) == 0
         assert LlamaCppEx.Context.n_rs_seq(mtp.main_ctx) == 0
       end