fix: use token_to_piece_bytes for proper token decoding

Dockerfile

@@ -1,5 +1,6 @@
-# ── Build stage: cargo-chef for dependency caching ──
+# ── Build stage: cargo-chef ──
-FROM lukemathwalker/cargo-chef:latest-rust-1.97.0 AS chef
+FROM lukemathwalker/cargo-chef:latest-rust-1.89.0 AS chef
+RUN apt-get update && apt-get install -y --no-install-recommends libclang-dev cmake && rm -rf /var/lib/apt/lists/*
 WORKDIR /app
 
 FROM chef AS planner
@@ -24,6 +25,7 @@ RUN apt-get update && apt-get install -y --no-install-recommends \
     ca-certificates \
     curl \
     libssl3 \
+    libgomp1 \
     && rm -rf /var/lib/apt/lists/*
 
 RUN groupadd -g 1001 appgroup && \

src/main.rs

@@ -1,6 +1,6 @@
 use axum::{
     extract::State,
-    http::StatusCode,
+    http::{HeaderMap, StatusCode},
     response::Json,
     routing::{get, post},
     Router,
@@ -9,9 +9,10 @@ use llama_cpp_2::{
     context::params::LlamaContextParams,
     llama_backend::LlamaBackend,
     llama_batch::LlamaBatch,
-    model::{params::LlamaModelParams, AddBos, LlamaModel, Special},
+    model::{params::LlamaModelParams, AddBos, LlamaModel},
     sampling::LlamaSampler,
     token::LlamaToken,
+    TokenToStringError,
 };
 use serde::{Deserialize, Serialize};
 use std::num::NonZeroU32;
@@ -132,9 +133,28 @@ struct HealthResponse {
     model: String,
 }
 
-const DEFAULT_MODEL_PATH: &str = "/root/models/gguf/MiniCPM-V-4.6-Q4_K_M.gguf";
+const DEFAULT_MODEL_PATH: &str = "/models/MiniCPM-V-4.6-Q4_K_M.gguf";
 const EOS_TOKEN: i32 = 248044;
 
+fn check_auth(headers: &HeaderMap) -> Result<(), (StatusCode, String)> {
+    let api_key = std::env::var("API_KEY").unwrap_or_default();
+    if api_key.is_empty() {
+        return Ok(()); // no key configured = open
+    }
+    let header = headers
+        .get("authorization")
+        .and_then(|v| v.to_str().ok())
+        .unwrap_or("");
+    let expected = format!("Bearer {api_key}");
+    if header == expected || header == api_key {
+        return Ok(());
+    }
+    Err((
+        StatusCode::UNAUTHORIZED,
+        "{\"error\":\"unauthorized\",\"message\":\"Invalid API key\"}".into(),
+    ))
+}
+
 #[tokio::main]
 async fn main() {
     tracing_subscriber::fmt()
@@ -216,8 +236,11 @@ async fn list_models() -> Json<ModelsResponse> {
 
 async fn chat_completions(
     State(state): State<Arc<AppState>>,
+    headers: HeaderMap,
     Json(req): Json<ChatRequest>,
 ) -> Result<Json<ChatResponse>, (StatusCode, String)> {
+    info!("Chat: {} chars, max_tokens={:?}", req.messages.len(), req.max_tokens);
+    check_auth(&headers)?;
     let prompt = build_prompt(&req.messages);
     let max_tokens = req.max_tokens.unwrap_or(256).min(1024);
 
@@ -261,10 +284,7 @@ async fn chat_completions(
         current = inner.sample_token();
     }
 
-    let output_text = state
+    let output_text = decode_tokens(&state.model, &output_tokens);
-        .model
-        .tokens_to_str(&output_tokens, Special::Tokenize)
-        .unwrap_or_else(|_| "<decode error>".to_string());
 
     let completion_tokens = output_tokens.len() as u32;
 
@@ -309,3 +329,23 @@ fn build_prompt(messages: &[ChatMessage]) -> String {
     prompt.push_str("Assistant: ");
     prompt
 }
+
+fn decode_tokens(model: &LlamaModel, tokens: &[LlamaToken]) -> String {
+    let mut out = String::with_capacity(tokens.len() * 4);
+    for &token in tokens {
+        // Try with a reasonable initial buffer (32 bytes)
+        let bytes = match model.token_to_piece_bytes(token, 32, true, None) {
+            Ok(b) => b,
+            Err(TokenToStringError::InsufficientBufferSpace(neg)) => {
+                // Retry with the suggested buffer size
+                let size = (-neg).max(0).try_into().unwrap_or(256);
+                model.token_to_piece_bytes(token, size, true, None).unwrap_or_default()
+            }
+            _ => continue,
+        };
+        if let Ok(s) = String::from_utf8(bytes) {
+            out.push_str(&s);
+        }
+    }
+    out
+}