inferencelayer 0.2.3

//! `lfm2-serve` — OpenAI-compatible multi-user serving on the wgpu engine.
//!
//! One engine thread owns the GPU and the continuous-batching [`Scheduler`]; axum handlers talk to
//! it over channels. Decoding is greedy by default, with optional `temperature`/`top_p`/`seed`
//! nucleus sampling; both paths are **bitwise batch-invariant**: a request's tokens are identical
//! no matter what traffic shares its batches (the property vLLM lacks and SGLang charges ~34%
//! for), and a seeded sampled request replays exactly under any load. Out-of-range or unimplemented
//! parameters are rejected with an OpenAI error envelope (400/404), never silently clamped.
//!
//! The wire surface (`chat.completion`/`text_completion` objects with faithful `finish_reason`,
//! unique response ids, `usage` incl. radix `cached_tokens`, `n` choices, `stop` strings and
//! `stop_token_ids`, client-disconnect cancellation, `/v1/models`, `/health`) lives in the
//! feature-gated `inferencelayer::serve` library so its pure logic is unit-tested there; this binary is
//! the async plumbing.
//!
//! Embeddings run on a SEPARATE engine thread (embedding models are different checkpoints and
//! touch none of the KV/scheduler machinery): handlers queue jobs over a channel, the embed
//! thread drains pending jobs into one ragged micro-batch per GPU dispatch, and `EmbedEngine`
//! auto-falls back to the native CPU encoder on GPU-less hosts.
//!
//! ```text
//! lfm2-serve [--model <dir>] [--embed-model <dir>] [--port 8210] [--batch 8] [--embed-batch 32]
//! POST /v1/completions        {"prompt": "...", "max_tokens": 64, "n": 1, "stop": ["\n"], ...}
//!
//! `--served-model-name <id>` advertises `<id>` instead of the checkpoint's directory name, so a
//! client that sends `"model": "org/name"` is served rather than 404'd (vLLM's flag of the same
//! name). Required to stand in for a vLLM server: an HF id has a slash and cannot be a dir name.
//! POST /v1/chat/completions   {"messages": [{"role": "user", "content": "..."}], ...}
//! POST /v1/embeddings         {"input": "..."|["...", ...], "dimensions": 256?}
//! GET  /v1/models             the loaded generation/embedding checkpoints
//! GET  /health                liveness
//! GET  /metrics               scheduler counters + radix cache size
//! ```
use std::collections::{HashMap, VecDeque};
use std::net::SocketAddr;
use std::path::PathBuf;
use std::sync::{Arc, Mutex};

use anyhow::{Context, Result};
use axum::Json;
use axum::extract::State;
use axum::response::IntoResponse;
use axum::response::sse::{Event, Sse};
use axum::routing::{get, post};
use inferencelayer::encoder_weights::{EncBatch, PosKind};
use inferencelayer::serve::chat::{self, ChatTurn, ToolCall};
use inferencelayer::serve::error::ApiError;
use inferencelayer::serve::stop;
use inferencelayer::serve::tools;
use inferencelayer::serve::types::{
    ChatReq, CommonParams, CompletionReq, Usage, created_epoch, response_id,
};
use inferencelayer::weights::Arch;
use inferencelayer::{
    EmbedEngine, EmbedOut, FinishReason, GpuCtx, Lfm2Gpu, RequestParams, SamplingParams, Scheduler,
    ServeStats, Weights, pooling,
};
use serde::{Deserialize, Serialize};
use tokio::sync::mpsc as amp;

/// Scratch capacity of the embed engine: total tokens per GPU dispatch (ragged batches are
/// chunked at sequence boundaries to fit; a single sequence is bounded by the model's own
/// max positions, which never exceeds this).
const EMBED_MAX_TOKENS: usize = 8192;

/// Which endpoint a streamed/whole response is shaped for.
#[derive(Clone, Copy, PartialEq, Eq)]
enum RespKind {
    Completion,
    Chat,
}

impl RespKind {
    /// The `object` field of a whole (non-streamed) response.
    fn object(self) -> &'static str {
        match self {
            RespKind::Completion => "text_completion",
            RespKind::Chat => "chat.completion",
        }
    }
    /// The `object` field of a streamed chunk.
    fn chunk_object(self) -> &'static str {
        match self {
            RespKind::Completion => "text_completion.chunk",
            RespKind::Chat => "chat.completion.chunk",
        }
    }
    /// Response-id prefix (`cmpl` / `chatcmpl`).
    fn id_prefix(self) -> &'static str {
        match self {
            RespKind::Completion => "cmpl",
            RespKind::Chat => "chatcmpl",
        }
    }
}

/// How a choice terminated, in wire terms: the OpenAI `finish_reason` plus the vendor `stop_reason`
/// detail (the specific stop string or token id, else null).
#[derive(Clone)]
struct Finish {
    reason: &'static str,
    stop_reason: Option<serde_json::Value>,
}

impl Finish {
    fn from_scheduler(fr: FinishReason) -> Self {
        match fr {
            FinishReason::Eos => Finish {
                reason: "stop",
                stop_reason: None,
            },
            FinishReason::StopTokenId(id) => Finish {
                reason: "stop",
                stop_reason: Some(serde_json::json!(id)),
            },
            FinishReason::Length => Finish {
                reason: "length",
                stop_reason: None,
            },
        }
    }
    /// Terminated by a stop STRING (scanned over the detokenized text in the engine thread).
    fn from_stop_string(which: &str) -> Self {
        Finish {
            reason: "stop",
            stop_reason: Some(serde_json::json!(which)),
        }
    }
}

/// One generation request handed to the engine thread. `n` choices share this prompt (the radix
/// cache shares their prefix KV) and differ only by a per-choice seed offset.
/// Reject a request that cannot be admitted, telling the caller WHY rather than dying quietly.
#[allow(dead_code)]
fn fail(r: &EngineRequest, msg: &str) {
    let _ = r.reply.send(EngineEvent::Rejected(msg.to_string()));
}

/// Cumulative tower-encode busy time / count — the under-load attribution counters
/// (`/metrics` vision_us): tower work runs handler-side, so scheduler stats cannot see it.
static VISION_US: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
static VISION_CALLS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
/// CPU preprocessing (PNG decode / resize / patchify) inside `prepare` — the part a pod's CPU
/// limit throttles; the remainder of vision_us is tower GPU forward + upload.
static VISION_PRE_US: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
/// Engine-thread time OUTSIDE step() while sequences are live (emission processing, detok,
/// channel sends, stats publish) — GPU-idle if nothing else is queued.
static EMIT_US: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);

/// The vision half of the serving loop: the tower, plus the prompt surgery around it.
///
/// The chat template emits ONE `<|image_pad|>` per image — expansion to one placeholder per merged
/// patch normally happens in HuggingFace's image processor, which needs torch. It happens here
/// instead, which is what lets the calling service render its prompt with a tokenizer and a jinja
/// renderer and nothing else.
enum Tower {
    /// Qwen3.5-VL family — on wgpu.
    Qwen(inferencelayer::vision_gpu::VisionGpu),
    /// GLM-OCR family. The CPU tower is the loader + fallback; the GPU tower (iso-gated against it)
    /// is used when present — MEASURED as the OCR bottleneck (~26 ms/vision-token ⇒ ~10 s/page on
    /// CPU while the GPU idles). `OSFKB_GLM_VISION_GPU=0` opts back into CPU.
    Glm(
        inferencelayer::vision_glm::GlmVisionTower,
        Option<inferencelayer::vision_glm_gpu::GlmVisionGpu>,
    ),
}

struct VisionCtx {
    tower: Tower,
    /// CUDA tower arm (`OSFKB_CUDA_TOWER=1` + successful device init; GLM only). Tried FIRST;
    /// the portable wgpu tower below stays the default and the fallback — see cuda_tower.rs
    /// for the portability contract.
    #[cfg(feature = "cudarc")]
    cuda: Option<inferencelayer::cuda_tower::CudaGlmTower>,
    image_token_id: u32,
    /// The tower's OWN device/queue — NOT the engine's. Tower forwards are chunky (a 760px page
    /// is ~1 s of GPU work in a handful of submissions), and on a shared queue every ENGINE
    /// step's poll waited on the queue timeline behind them: with 16 pages encoding, decode
    /// steps stalled for tower chunks they had nothing to do with. A separate logical device
    /// gives the tower its own timeline; the driver schedules both streams, and the tower's
    /// only output crosses as host floats (`VisionPrompt.embeds`), so nothing is shared.
    /// `None` = a second device could not be created; the tower shares the engine's queue as
    /// before (callers pass it to [`Self::prepare`]).
    vctx: Option<GpuCtx>,
}

impl VisionCtx {
    /// Load the tower if this checkpoint has one. A text-only checkpoint returns `None` and the
    /// server behaves exactly as it always did.
    fn load(ctx: &GpuCtx, dir: &std::path::Path) -> Option<Self> {
        let cfg: serde_json::Value =
            serde_json::from_slice(&std::fs::read(dir.join("config.json")).ok()?).ok()?;
        let image_token_id = cfg.get("image_token_id")?.as_u64()? as u32;
        let arch = cfg
            .get("architectures")
            .and_then(|a| a.get(0))
            .and_then(|x| x.as_str())
            .unwrap_or("");
        // Same adapter, second logical device (see the field doc) — OPT-IN via
        // OSFKB_VISION_DEVICE=2: two wgpu devices mean two allocators that cannot share freed
        // chunks, and on a card as full as the prod V100 (NuExtract3 co-resident, ~27.5/32 GB)
        // the tower's scratch OOM-panicked at the first forward. Enable it only where VRAM has
        // real headroom.
        let vctx = if std::env::var("OSFKB_VISION_DEVICE").as_deref() == Ok("2") {
            match GpuCtx::new() {
                Ok(c) => {
                    eprintln!("vision tower: own device/queue (OSFKB_VISION_DEVICE=2)");
                    Some(c)
                }
                Err(e) => {
                    eprintln!("vision tower: no second device ({e}); sharing the engine queue");
                    None
                }
            }
        } else {
            None
        };
        let tctx = vctx.as_ref().unwrap_or(ctx);
        if arch == "GlmOcrForConditionalGeneration" {
            let tower = inferencelayer::vision_glm::GlmVisionTower::load(dir).ok()?;
            // CUDA tensor-core arm: opt-in AND probe-gated — any failure (no NVIDIA driver,
            // no libcublas, init error) logs and falls through to the portable wgpu tower.
            #[cfg(feature = "cudarc")]
            let cuda = if std::env::var("OSFKB_CUDA_TOWER").as_deref() == Ok("1") {
                match inferencelayer::cuda_tower::CudaGlmTower::new(&tower) {
                    Ok(c) => {
                        eprintln!("GLM vision tower: CUDA arm active (tensor cores)");
                        Some(c)
                    }
                    Err(e) => {
                        eprintln!("GLM vision tower: CUDA arm unavailable ({e:#}); using wgpu");
                        None
                    }
                }
            } else {
                None
            };
            let gpu = if std::env::var("OSFKB_GLM_VISION_GPU").as_deref() == Ok("0") {
                None
            } else {
                match inferencelayer::vision_glm_gpu::GlmVisionGpu::new(tctx, &tower) {
                    Ok(g) => Some(g),
                    Err(e) => {
                        eprintln!("GLM vision tower: GPU build failed ({e}); using the CPU tower");
                        None
                    }
                }
            };
            eprintln!(
                "GLM vision tower loaded (image_token_id={image_token_id}, {}) — images accepted",
                if gpu.is_some() { "GPU" } else { "CPU" }
            );
            return Some(Self {
                tower: Tower::Glm(tower, gpu),
                #[cfg(feature = "cudarc")]
                cuda,
                image_token_id,
                vctx,
            });
        }
        let cpu = inferencelayer::vision::VisionTower::load(dir).ok()?;
        let tower = inferencelayer::vision_gpu::VisionGpu::new(tctx, cpu).ok()?;
        eprintln!("vision tower loaded (image_token_id={image_token_id}) — images accepted");
        Some(Self {
            #[cfg(feature = "cudarc")]
            cuda: None,
            tower: Tower::Qwen(tower),
            image_token_id,
            vctx,
        })
    }

    /// `prompt_ids` (with one `image_token_id` per image) + raw image bytes → the EXPANDED prompt and
    /// the vision payload the scheduler substitutes into the residual stream.
    fn prepare(
        &self,
        engine_ctx: &GpuCtx,
        prompt_ids: &[u32],
        images: &[Vec<u8>],
    ) -> anyhow::Result<(Vec<u32>, inferencelayer::server::VisionPrompt)> {
        // Tower forwards run on the tower's OWN device (see the `vctx` field doc); the engine's
        // queue only as the no-second-device fallback.
        let ctx = self.vctx.as_ref().unwrap_or(engine_ctx);
        let cfg = match &self.tower {
            Tower::Qwen(t) => t.config().clone(),
            Tower::Glm(t, _) => t.cfg.clone(),
        };
        let cfg = &cfg;
        let pads = prompt_ids
            .iter()
            .filter(|&&t| t == self.image_token_id)
            .count();
        anyhow::ensure!(
            pads == images.len(),
            "prompt has {pads} image placeholder(s) but {} image(s) were sent",
            images.len()
        );

        let _t_pre = std::time::Instant::now();
        let patches: Vec<_> = images
            .iter()
            .map(|b| match &self.tower {
                Tower::Qwen(_) => inferencelayer::vision::preprocess_bytes(b, cfg),
                Tower::Glm(..) => inferencelayer::vision_glm::glm_preprocess_bytes(b, cfg),
            })
            .collect::<anyhow::Result<_>>()?;
        VISION_PRE_US.fetch_add(
            _t_pre.elapsed().as_micros() as u64,
            std::sync::atomic::Ordering::Relaxed,
        );

        // Expand each single placeholder into one per merged patch. The count is a property of the
        // IMAGE (its grid), which is why it cannot be known by a caller that has not preprocessed it.
        let mut expanded = Vec::with_capacity(
            prompt_ids.len() + patches.iter().map(|p| p.num_tokens(cfg)).sum::<usize>(),
        );
        let mut next = 0usize;
        for &t in prompt_ids {
            if t == self.image_token_id {
                let n = patches[next].num_tokens(cfg);
                expanded.extend(std::iter::repeat_n(self.image_token_id, n));
                next += 1;
            } else {
                expanded.push(t);
            }
        }

        let vp = match &self.tower {
            Tower::Qwen(t) => inferencelayer::vision::prepare_prompt_gpu(
                ctx,
                t,
                &expanded,
                self.image_token_id,
                &patches,
            )?,
            Tower::Glm(t, gpu) => {
                #[cfg(feature = "cudarc")]
                if let Some(c) = &self.cuda {
                    let vp = inferencelayer::vision_glm::glm_prepare_prompt_cuda(
                        c,
                        &t.cfg,
                        &expanded,
                        self.image_token_id,
                        &patches,
                    )?;
                    return Ok((expanded, vp));
                }
                match gpu {
                Some(g) => inferencelayer::vision_glm::glm_prepare_prompt_gpu(
                    ctx,
                    g,
                    &expanded,
                    self.image_token_id,
                    &patches,
                )?,
                None => inferencelayer::vision_glm::glm_prepare_prompt(
                    t,
                    &expanded,
                    self.image_token_id,
                    &patches,
                )?,
            }
            }
        };
        Ok((expanded, vp))
    }
}

struct EngineRequest {
    /// For an image request these are the EXPANDED ids (one placeholder per merged patch) — the
    /// vision tower now runs in the HTTP handler (spawn_blocking), NOT on the engine thread:
    /// under load, an on-thread tower serialized every request behind every other's encode
    /// (measured: 8-way concurrency was 0.38 pages/s vs 0.35 sequential — zero overlap).
    prompt_ids: Vec<u32>,
    /// The tower's output for an image request: embeddings + M-RoPE positions, ready for the
    /// scheduler. None for text.
    vprompt: Option<inferencelayer::server::VisionPrompt>,
    max_tokens: usize,
    n: usize,
    params: RequestParams,
    stop_strings: Vec<String>,
    reply: amp::UnboundedSender<EngineEvent>,
}

/// An event from the engine thread to a handler, tagged with the choice index it belongs to.
enum EngineEvent {
    Delta {
        choice: usize,
        text: String,
    },
    Done {
        choice: usize,
        text: String,
        completion_tokens: usize,
        prompt_tokens: usize,
        cached_tokens: usize,
        finish: Finish,
        logprobs: Vec<TokenLp>,
    },
    Error(String),
    /// The scheduler REFUSED the request — it does not fit the per-sequence cap or the KV pool, and
    /// no amount of waiting will change that. A distinct variant because it is a CLIENT error (the
    /// request is too big) and must not be reported as a 500: an OpenAI client seeing
    /// `internal_error` retries, and every retry is refused the same way.
    Rejected(String),
}

/// A unit of work for the embed thread: either an OpenAI `/v1/embeddings` request or a Cohere/Jina
/// `/v1/rerank` request. Both run on the one encoder; embeds coalesce into a shared micro-batch,
/// reranks run one at a time (query + documents in one batch).
enum EmbedJob {
    Embed(EmbedReq),
    Rerank(RerankReq),
}

/// One embedding request in flight: raw texts in, pooled vectors out.
struct EmbedReq {
    texts: Vec<String>,
    /// Matryoshka truncation (`dimensions` in the OpenAI API): truncate + re-L2-normalize.
    dimensions: Option<usize>,
    reply: tokio::sync::oneshot::Sender<Result<EmbedDone, String>>,
}

struct EmbedDone {
    vectors: Vec<Vec<f32>>,
    prompt_tokens: usize,
}

/// One rerank request: score every document against the query and return them score-descending.
struct RerankReq {
    query: String,
    documents: Vec<String>,
    top_n: Option<usize>,
    reply: tokio::sync::oneshot::Sender<Result<RerankDone, String>>,
}

struct RerankDone {
    /// `(original document index, relevance score)`, sorted score-descending (stable), `top_n` applied.
    ranked: Vec<(usize, f32)>,
    prompt_tokens: usize,
}

/// Generation-side serving state (present when `--model` was given).
#[derive(Clone)]
struct GenState {
    submit: std::sync::mpsc::Sender<EngineRequest>,
    tokenizer: Arc<tokenizers::Tokenizer>,
    arch: Arch,
    model_name: String,
    /// For handler-side vision encode (concurrent towers; the engine thread never blocks on one).
    gctx: Arc<GpuCtx>,
    vision: Option<Arc<VisionCtx>>,
    /// Caps CONCURRENT tower forwards (`OSFKB_VISION_CONCURRENCY`, default 2). More buys
    /// nothing — the GPU queue serializes them anyway — and each in-flight forward holds
    /// ~0.5 GB of scratch, so 8-way bursts spike VRAM for zero throughput (that spike is what
    /// OOM'd the second-device tower). Encode-vs-decode overlap only needs ~2 in flight.
    vision_gate: Arc<tokio::sync::Semaphore>,
}

/// Embedding-side serving state (present when `--embed-model` was given).
#[derive(Clone)]
struct EmbedState {
    submit: std::sync::mpsc::Sender<EmbedJob>,
    model_name: String,
    dimension: usize,
    /// The checkpoint produces per-token (ColBERT) embeddings — `/v1/embeddings` rejects it and
    /// points at `/v1/rerank` (which scores via MaxSim).
    is_per_token: bool,
}

#[derive(Clone)]
struct AppState {
    generation: Option<GenState>,
    embed: Option<EmbedState>,
    stats: Arc<Mutex<(ServeStats, usize, usize)>>,
}

fn main() -> Result<()> {
    // ABORT on any thread panic. A panicked ENGINE thread previously left the pod Running and
    // Ready (the HTTP layer lives on) while every completion hung to its client timeout — a
    // zombie that kube-proxy keeps feeding. Crashing the process turns that into a restart.
    let default_hook = std::panic::take_hook();
    std::panic::set_hook(Box::new(move |info| {
        default_hook(info);
        std::process::abort();
    }));
    // Serving default: SPIN-poll the GPU. The sleep-wakeup readback path pays OS scheduling latency
    // on EVERY per-token decode step — measured ~800 ms/token on Metal (1.2 tok/s for a 0.6B whose
    // GPU step is 2.7 ms), because `GpuCtx::read` blocks on a fence instead of spinning. Burning one
    // core to spin collapses it to GPU-bound (~240 tok/s here). Every lfm2-serve model is a
    // latency-sensitive decoder; this is the same trade moshi-serve already makes. `OSFKB_SPIN_POLL=0`
    // opts back out (e.g. a shared box that can't spare the core).
    if std::env::var_os("OSFKB_SPIN_POLL").is_none() {
        unsafe { std::env::set_var("OSFKB_SPIN_POLL", "1") };
    }
    let mut model = None;
    let mut embed_model = None;
    let mut port = 8210u16;
    let mut batch = 8usize;
    let mut embed_batch = 32usize;
    let mut vram_gb: Option<f64> = None;
    let mut kv_fraction = 0.35f64;
    // 0 = "let the engine decide": a wide split-fuse prefill budget when the adapter has
    // subgroups, off otherwise. Defaulting this to a hard 0 meant the wide plan was BUILT and
    // never engaged — prompts were chunked through the narrow k-plan for no reason.
    let mut max_batched: Option<usize> = None;
    // The name this server ADVERTISES, independent of where the weights sit on disk. OpenAI clients
    // send the model id they know and expect a 404 if it is not served — and an id like
    // `principled-intelligence/claim-extractor-4B-q-2605` can never be a directory name (it has a
    // slash in it), so without this the engine cannot be a drop-in for a vLLM server serving that
    // model. vLLM spells the same flag `--served-model-name`.
    let mut served_name: Option<String> = None;
    let mut args = std::env::args().skip(1);
    while let Some(a) = args.next() {
        match a.as_str() {
            "--model" => model = args.next().map(PathBuf::from),
            "--embed-model" => embed_model = args.next().map(PathBuf::from),
            "--port" => port = args.next().context("--port value")?.parse()?,
            "--batch" => batch = args.next().context("--batch value")?.parse()?,
            "--embed-batch" => embed_batch = args.next().context("--embed-batch value")?.parse()?,
            // KV-pool sizing: --vram-gb sets the device VRAM budget, --kv-fraction the share of it
            // (default 0.35) given to the paged KV cache. Precedence: these flags > env
            // OSFKB_KV_POOL_TOKENS > the MAX_T default.
            "--vram-gb" => vram_gb = Some(args.next().context("--vram-gb value")?.parse()?),
            "--kv-fraction" => kv_fraction = args.next().context("--kv-fraction value")?.parse()?,
            // Split-fuse: zero-decode steps pack up to W prompt columns through the wide KC=16 plan
            // (0 = off; ignored on adapters without subgroups).
            "--max-batched-tokens" => {
                max_batched = Some(args.next().context("--max-batched-tokens value")?.parse()?)
            }
            "--served-model-name" => {
                served_name = Some(args.next().context("--served-model-name value")?)
            }
            other => anyhow::bail!("unknown argument {other}"),
        }
    }
    anyhow::ensure!(
        kv_fraction > 0.0 && kv_fraction <= 1.0,
        "--kv-fraction must be in (0, 1]"
    );
    anyhow::ensure!(
        model.is_some() || embed_model.is_some(),
        "at least one of --model <dir> / --embed-model <dir> is required"
    );
    let stats: Arc<Mutex<(ServeStats, usize, usize)>> = Arc::new(Mutex::new((ServeStats::default(), 0, 0)));

    let embed = match embed_model {
        Some(dir) => Some(spawn_embed_engine(dir, embed_batch)?),
        None => None,
    };
    let generation = match model {
        Some(model) => Some(spawn_gen_engine(
            model,
            served_name,
            batch,
            max_batched,
            vram_gb,
            kv_fraction,
            stats.clone(),
        )?),
        None => None,
    };

    let app_state = AppState {
        generation,
        embed,
        stats,
    };
    let rt = tokio::runtime::Builder::new_multi_thread()
        .enable_all()
        .build()?;
    rt.block_on(async move {
        let app = axum::Router::new()
            .route("/v1/completions", post(completions))
            .route("/v1/chat/completions", post(chat_completions))
            .route("/v1/embeddings", post(embeddings))
            .route("/v1/rerank", post(rerank))
            .route("/v1/models", get(models))
            .route("/health", get(health))
            .route("/metrics", get(metrics))
            .with_state(app_state);
        let addr = SocketAddr::from(([0, 0, 0, 0], port));
        let listener = tokio::net::TcpListener::bind(addr).await?;
        eprintln!("lfm2-serve listening on {addr}");
        axum::serve(listener, app).await?;
        Ok(())
    })
}

/// Load the embedding checkpoint on the best device (GPU → CPU fallback) and start its engine
/// thread: drain everything pending into one ragged micro-batch per dispatch, fan results back.
fn spawn_embed_engine(dir: PathBuf, embed_batch: usize) -> Result<EmbedState> {
    let mut engine = EmbedEngine::auto(&dir, EMBED_MAX_TOKENS)?;
    eprintln!("embed backend: {}", engine.device());
    let cfg = engine.config();
    let dimension = cfg.hidden;
    let is_per_token = matches!(
        cfg.pooling,
        inferencelayer::pooling::Pooling::PerToken { .. }
    );
    let pos_offset = match cfg.pos_kind {
        PosKind::Learned { offset } => offset,
        PosKind::Rope { .. } => 0,
    };
    let max_seq = cfg.max_pos - pos_offset;
    // Sanitized load: strips baked-in padding (the ragged encoder must never see [PAD] rows).
    let tok = inferencelayer::load_encoder_tokenizer(&dir, max_seq)
        .map_err(|e| anyhow::anyhow!("embed tokenizer: {e}"))?;
    let model_name = dir
        .file_name()
        .map(|s| s.to_string_lossy().into_owned())
        .unwrap_or_else(|| "embed-model".into());
    let (tx, rx) = std::sync::mpsc::channel::<EmbedJob>();
    std::thread::spawn(move || {
        loop {
            let first = match rx.recv() {
                Ok(j) => j,
                Err(_) => return,
            };
            let mut jobs = vec![first];
            while jobs.len() < embed_batch {
                match rx.try_recv() {
                    Ok(j) => jobs.push(j),
                    Err(_) => break,
                }
            }
            // Coalesce embeds into one micro-batch; run each rerank on its own (query + docs batch).
            let mut embeds = Vec::new();
            for j in jobs {
                match j {
                    EmbedJob::Embed(req) => embeds.push(req),
                    EmbedJob::Rerank(req) => run_rerank(&mut engine, &tok, max_seq, req),
                }
            }
            if !embeds.is_empty() {
                run_embed_batch(&mut engine, &tok, max_seq, embeds);
            }
        }
    });
    Ok(EmbedState {
        submit: tx,
        model_name,
        dimension,
        is_per_token,
    })
}

/// Tokenize all texts of the drained jobs, chunk into ragged batches within the engine's token
/// budget, encode, and distribute per-job results (with per-job Matryoshka truncation).
fn run_embed_batch(
    engine: &mut EmbedEngine,
    tok: &tokenizers::Tokenizer,
    max_seq: usize,
    jobs: Vec<EmbedReq>,
) {
    let texts: Vec<String> = jobs.iter().flat_map(|j| j.texts.iter().cloned()).collect();
    let encodings = match tok.encode_batch(texts, true) {
        Ok(e) => e,
        Err(e) => {
            for j in jobs {
                let _ = j.reply.send(Err(format!("tokenize: {e}")));
            }
            return;
        }
    };
    let mut id_lists: Vec<Vec<u32>> = Vec::with_capacity(encodings.len());
    for enc in &encodings {
        let mut ids = enc.get_ids().to_vec();
        if ids.len() > max_seq {
            eprintln!(
                "embed: truncating a {}-token input to the model's {max_seq}-token window",
                ids.len()
            );
            ids.truncate(max_seq);
        }
        id_lists.push(ids);
    }
    // Empty tokenizations cannot be embedded (a mean over zero rows has no value) — fail the
    // owning job, keep the rest of the drain alive.
    let mut vectors: Vec<Option<Vec<f32>>> = vec![None; id_lists.len()];
    let mut chunk: Vec<usize> = Vec::new();
    let mut chunk_tokens = 0usize;
    let mut failed = false;
    let flush = |chunk: &mut Vec<usize>,
                 chunk_tokens: &mut usize,
                 vectors: &mut Vec<Option<Vec<f32>>>,
                 engine: &mut EmbedEngine,
                 id_lists: &[Vec<u32>]|
     -> bool {
        if chunk.is_empty() {
            return true;
        }
        let batch = EncBatch::from_seqs(chunk.iter().map(|&i| id_lists[i].clone()));
        let ok = match engine.encode(&batch) {
            Ok(EmbedOut::Pooled(vs)) => {
                for (&i, v) in chunk.iter().zip(vs) {
                    vectors[i] = Some(v);
                }
                true
            }
            Ok(EmbedOut::PerToken(_)) => false,
            Err(e) => {
                eprintln!("embed: encode failed: {e:#}");
                false
            }
        };
        chunk.clear();
        *chunk_tokens = 0;
        ok
    };
    for (i, ids) in id_lists.iter().enumerate() {
        if ids.is_empty() {
            continue;
        }
        if chunk_tokens + ids.len() > EMBED_MAX_TOKENS
            && !flush(
                &mut chunk,
                &mut chunk_tokens,
                &mut vectors,
                engine,
                &id_lists,
            )
        {
            failed = true;
            break;
        }
        chunk.push(i);
        chunk_tokens += ids.len();
    }
    if !failed {
        failed = !flush(
            &mut chunk,
            &mut chunk_tokens,
            &mut vectors,
            engine,
            &id_lists,
        );
    }

    // Fan results back out per job, in input order.
    let mut cursor = 0usize;
    for j in jobs {
        let n = j.texts.len();
        let span = cursor..cursor + n;
        cursor += n;
        if failed {
            let _ = j.reply.send(Err("embedding engine failed".to_string()));
            continue;
        }
        let mut out = Vec::with_capacity(n);
        let mut prompt_tokens = 0usize;
        let mut err: Option<String> = None;
        for i in span {
            prompt_tokens += id_lists[i].len();
            match &vectors[i] {
                Some(v) => out.push(match j.dimensions {
                    Some(d) => pooling::matryoshka_truncate(v, d),
                    None => v.clone(),
                }),
                None => {
                    err = Some("input tokenized to zero tokens".to_string());
                    break;
                }
            }
        }
        let _ = match err {
            Some(e) => j.reply.send(Err(e)),
            None => j.reply.send(Ok(EmbedDone {
                vectors: out,
                prompt_tokens,
            })),
        };
    }
}

/// Score every document against the query and reply with them ranked score-descending. The query
/// is sequence 0, documents follow; late-interaction (PerToken/ColBERT) checkpoints score with
/// MaxSim, pooled checkpoints with cosine (a dot of the already-L2-normalized vectors).
fn run_rerank(
    engine: &mut EmbedEngine,
    tok: &tokenizers::Tokenizer,
    max_seq: usize,
    req: RerankReq,
) {
    let mut texts = Vec::with_capacity(1 + req.documents.len());
    texts.push(req.query.clone());
    texts.extend(req.documents.iter().cloned());
    let encodings = match tok.encode_batch(texts, true) {
        Ok(e) => e,
        Err(e) => {
            let _ = req.reply.send(Err(format!("tokenize: {e}")));
            return;
        }
    };
    let id_lists: Vec<Vec<u32>> = encodings
        .iter()
        .map(|e| {
            let mut ids = e.get_ids().to_vec();
            ids.truncate(max_seq);
            ids
        })
        .collect();
    let prompt_tokens: usize = id_lists.iter().map(Vec::len).sum();
    // Encode all sequences, chunked to the engine's token budget (the checkpoint is uniformly
    // pooled OR per-token, so exactly one accumulator fills).
    let n = id_lists.len();
    let mut pooled: Vec<Vec<f32>> = Vec::with_capacity(n);
    let mut per_token: Vec<Vec<Vec<f32>>> = Vec::with_capacity(n);
    let mut i = 0;
    while i < n {
        let mut j = i;
        let mut toks = 0usize;
        while j < n && (j == i || toks + id_lists[j].len() <= EMBED_MAX_TOKENS) {
            toks += id_lists[j].len();
            j += 1;
        }
        let batch = EncBatch::from_seqs(id_lists[i..j].iter().cloned());
        match engine.encode(&batch) {
            Ok(EmbedOut::Pooled(vs)) => pooled.extend(vs),
            Ok(EmbedOut::PerToken(ss)) => per_token.extend(ss),
            Err(e) => {
                let _ = req.reply.send(Err(format!("encode: {e:#}")));
                return;
            }
        }
        i = j;
    }
    let ndocs = req.documents.len();
    let mut ranked: Vec<(usize, f32)> = if !per_token.is_empty() {
        let q = &per_token[0];
        (0..ndocs)
            .map(|d| (d, pooling::maxsim(q, &per_token[d + 1])))
            .collect()
    } else {
        let q = &pooled[0];
        (0..ndocs)
            .map(|d| {
                let doc = &pooled[d + 1];
                let s = q.iter().zip(doc).map(|(a, b)| a * b).sum::<f32>();
                (d, s)
            })
            .collect()
    };
    // Stable score-descending sort (ties keep input order), then apply top_n.
    ranked.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
    if let Some(k) = req.top_n {
        ranked.truncate(k);
    }
    let _ = req.reply.send(Ok(RerankDone {
        ranked,
        prompt_tokens,
    }));
}

/// Load the generation checkpoint and start the continuous-batching engine thread (the only GPU
/// owner). It admits requests between steps, fans each `n`-choice request into `n` scheduler
/// submits, scans stop strings over the detokenized text, and cancels a sequence the moment its
/// client disconnects (a reply-channel send failure).
fn spawn_gen_engine(
    model: PathBuf,
    served_name: Option<String>,
    batch: usize,
    max_batched: Option<usize>,
    vram_gb: Option<f64>,
    kv_fraction: f64,
    stats: Arc<Mutex<(ServeStats, usize, usize)>>,
) -> Result<GenState> {
    let tokenizer = tokenizers::Tokenizer::from_file(model.join("tokenizer.json"))
        .map_err(|e| anyhow::anyhow!("tokenizer: {e}"))?;
    let eos = read_eos_ids(&model, &tokenizer);
    // Q1/GGUF (Bonsai-class) serving: a `*.gguf` in the model dir switches the loader. The env
    // flags must be set BEFORE any pipeline is built: Q1 selects the sign-weight kernel plan,
    // and the DeltaNet state slabs default to ~194 slots x ~2 MB x DN-layers — far past a
    // 32 GB card on a 48-DN-layer 27B. Serving needs exactly 2 + MAX_SLOTS (solo + trash +
    // the scheduler's 16 sequence slots); an explicit env still wins.
    let gguf_model = std::fs::read_dir(&model).ok().and_then(|rd| {
        rd.filter_map(|e| e.ok().map(|e| e.path()))
            .find(|p| p.extension().is_some_and(|x| x == "gguf"))
    });
    if gguf_model.is_some() {
        unsafe { std::env::set_var("OSFKB_Q1_WEIGHTS", "1") };
        if std::env::var("OSFKB_DN_SLOTS").is_err() {
            unsafe { std::env::set_var("OSFKB_DN_SLOTS", "18") };
        }
    }
    eprintln!("eos ids: {eos:?}");

    let ctx = Arc::new(GpuCtx::new()?);
    eprintln!("backend: {}", ctx.backend);
    let w = match &gguf_model {
        Some(g) => {
            eprintln!("loading Q1 GGUF: {}", g.display());
            inferencelayer::gguf::load_bonsai(&ctx, g)?
        }
        None => Weights::load(&ctx, &model)?,
    };
    let arch = w.cfg.arch;
    // KV-pool sizing precedence: --vram-gb flag > env (OSFKB_KV_POOL_TOKENS, honored by `new`) >
    // MAX_T default. When a VRAM budget is given, derive pool_tokens from the checkpoint's
    // per-token KV footprint.
    let gpu = match vram_gb {
        Some(gb) => {
            let opts = inferencelayer::EngineOpts::from_vram(
                &w.cfg,
                (gb * 1024.0 * 1024.0 * 1024.0) as u64,
                kv_fraction,
            );
            eprintln!(
                "kv pool: {} tokens ({} blocks) from {gb} GB × {kv_fraction} ({} B/token)",
                opts.pool_tokens,
                opts.pool_tokens / 16,
                inferencelayer::EngineOpts::kv_bytes_per_token(&w.cfg),
            );
            Lfm2Gpu::new_with_opts(&ctx, w, opts)
        }
        None => Lfm2Gpu::new(&ctx, w),
    };
    let model_name = served_name.unwrap_or_else(|| {
        model
            .file_name()
            .map(|s| s.to_string_lossy().into_owned())
            .unwrap_or_else(|| "model".into())
    });

    let (submit_tx, submit_rx) = std::sync::mpsc::channel::<EngineRequest>();
    let stats_engine = stats;
    let tok_engine = tokenizer.clone();

    // The vision tower loads BEFORE the engine thread spawns and is shared with the HTTP
    // handlers: encode runs there (spawn_blocking, concurrently across requests), never on the
    // engine thread.
    let vision = VisionCtx::load(&ctx, &model).map(Arc::new);
    let vision_handlers = vision.clone();
    let ctx_handlers = ctx.clone();

    std::thread::spawn(move || {
        // Unset ⇒ the engine's default wide budget (the KC=16 split-fuse prefill plan when the
        // adapter has subgroups). Prompts then ride the batched prefill instead of being fed
        // through the narrow k-plan, which is what the wide plan was built for.
        let wide = max_batched.unwrap_or_else(|| inferencelayer::server::default_wide_budget(&ctx));
        let mut sched = Scheduler::new_with_wide(&gpu, &ctx, batch, wide, eos);
        let mut live: HashMap<u64, Live> = HashMap::new();
        loop {
            // Blocking wait when idle; opportunistic drain when busy.
            if sched.pending() == 0 && live.is_empty() {
                match submit_rx.recv() {
                    Ok(r) => admit(&mut sched, &mut live, r),
                    Err(_) => return,
                }
            }
            while let Ok(r) = submit_rx.try_recv() {
                admit(&mut sched, &mut live, r);
            }
            let ems = match sched.step(&gpu, &ctx) {
                Ok(e) => e,
                Err(e) => {
                    for (_, l) in live.drain() {
                        let _ = l.reply.send(EngineEvent::Error(format!("engine: {e}")));
                    }
                    continue;
                }
            };
            let _t_emit = std::time::Instant::now();
            // Attach each freshly-admitted sequence's radix prefix-hit count (usage cached_tokens).
            for (id, cached) in sched.drain_admissions() {
                if let Some(l) = live.get_mut(&id) {
                    l.cached_tokens = cached;
                }
            }
            for e in ems {
                let Some(l) = live.get_mut(&e.id) else {
                    continue;
                };
                l.ids.push(e.token);
                if let Some(lp) = &e.logprobs {
                    let token = decode_piece(&tok_engine, lp.chosen);
                    let top = lp
                        .top
                        .iter()
                        .map(|(id, v)| (decode_piece(&tok_engine, *id), *v))
                        .collect();
                    l.logprobs.push(TokenLp {
                        token,
                        logprob: lp.chosen_logprob,
                        top,
                    });
                }
                let full = tok_engine.decode(&l.ids, true).unwrap_or_default();
                // Terminal? A stop STRING (excluded from the output) wins over the scheduler's own
                // EOS/length/stop-token finish; otherwise stream the delta minus a held-back tail
                // that could still grow into a stop string.
                let (visible_end, finish) =
                    if let Some((pos, which)) = stop::first_stop(&full, &l.stop_strings) {
                        (pos, Some(Finish::from_stop_string(&which)))
                    } else if let Some(fr) = e.finish {
                        (full.len(), Some(Finish::from_scheduler(fr)))
                    } else {
                        (stop::safe_stream_end(&full, &l.stop_strings), None)
                    };
                let visible = &full[..visible_end];
                if visible.len() > l.sent.len() && visible.starts_with(&l.sent) {
                    let delta = visible[l.sent.len()..].to_string();
                    if l.reply
                        .send(EngineEvent::Delta {
                            choice: l.choice,
                            text: delta,
                        })
                        .is_err()
                    {
                        // Client gone: stop spending the GPU on a dead socket.
                        sched.cancel(e.id);
                        live.remove(&e.id);
                        continue;
                    }
                    l.sent = visible.to_string();
                }
                if let Some(finish) = finish {
                    let stopped_early = finish.reason == "stop"
                        && matches!(&finish.stop_reason, Some(v) if v.is_string());
                    let l = live.remove(&e.id).expect("live entry");
                    let _ = l.reply.send(EngineEvent::Done {
                        choice: l.choice,
                        text: visible.to_string(),
                        completion_tokens: l.ids.len(),
                        prompt_tokens: l.prompt_tokens,
                        cached_tokens: l.cached_tokens,
                        finish,
                        logprobs: l.logprobs,
                    });
                    // A stop STRING finishes the choice before the scheduler would: free the slot.
                    if stopped_early {
                        sched.cancel(e.id);
                    }
                }
            }
            let mut st = stats_engine.lock().expect("stats lock");
            *st = (sched.stats.clone(), sched.cache_entries(), sched.kv_free_blocks());
            drop(st);
            EMIT_US.fetch_add(
                _t_emit.elapsed().as_micros() as u64,
                std::sync::atomic::Ordering::Relaxed,
            );
        }

        fn admit(sched: &mut Scheduler, live: &mut HashMap<u64, Live>, r: EngineRequest) {
            // Vision already ran in the handler; the scheduler sees an ordinary (longer) prompt
            // plus a block of embeddings to substitute.
            let (prompt_ids, vprompt) = (r.prompt_ids.clone(), r.vprompt.clone());
            let prompt_tokens = prompt_ids.len();
            for i in 0..r.n {
                let mut params = r.params.clone();
                // Per-choice seed so sampled `n` explores distinct trajectories; greedy `n` (temp
                // 0) is seed-independent, hence identical, which is the correct greedy semantics.
                params.sampling.seed = params.sampling.seed.wrapping_add(i as u64);
                let submitted = match &vprompt {
                    Some(v) => sched.submit_with_vision(
                        prompt_ids.clone(),
                        r.max_tokens.max(1),
                        params,
                        v.clone(),
                    ),
                    None => {
                        sched.submit_with_params(prompt_ids.clone(), r.max_tokens.max(1), params)
                    }
                };
                match submitted {
                    Ok(id) => {
                        live.insert(
                            id,
                            Live {
                                reply: r.reply.clone(),
                                choice: i,
                                ids: Vec::new(),
                                sent: String::new(),
                                prompt_tokens,
                                cached_tokens: 0,
                                stop_strings: r.stop_strings.clone(),
                                logprobs: Vec::new(),
                            },
                        );
                    }
                    Err(e) => {
                        let _ = r.reply.send(EngineEvent::Rejected(e.to_string()));
                    }
                }
            }
        }
    });

    Ok(GenState {
        submit: submit_tx,
        tokenizer: Arc::new(tokenizer),
        arch,
        model_name,
        gctx: ctx_handlers,
        vision: vision_handlers,
        vision_gate: Arc::new(tokio::sync::Semaphore::new(
            std::env::var("OSFKB_VISION_CONCURRENCY")
                .ok()
                .and_then(|v| v.parse().ok())
                .filter(|&n: &usize| n > 0)
                .unwrap_or(2),
        )),
    })
}

/// One in-flight choice's engine-thread bookkeeping.
struct Live {
    reply: amp::UnboundedSender<EngineEvent>,
    choice: usize,
    ids: Vec<u32>,
    /// Visible text already streamed (detok is not strictly monotonic, so we diff against it).
    sent: String,
    prompt_tokens: usize,
    cached_tokens: usize,
    stop_strings: Vec<String>,
    /// Accumulated per-token logprobs (empty unless the request asked for them).
    logprobs: Vec<TokenLp>,
}

/// eos ids: config.json `eos_token_id` (int or array) ∪ the chat end-of-turn token.
fn read_eos_ids(model: &std::path::Path, tok: &tokenizers::Tokenizer) -> Vec<u32> {
    let mut eos = Vec::new();
    if let Ok(cfg) = std::fs::read_to_string(model.join("config.json"))
        && let Ok(v) = serde_json::from_str::<serde_json::Value>(&cfg)
    {
        // Multimodal wrappers (GLM-OCR) nest the language model's ids under `text_config`.
        for spot in [&v["eos_token_id"], &v["text_config"]["eos_token_id"]] {
            match spot {
                serde_json::Value::Number(n) => eos.extend(n.as_u64().map(|x| x as u32)),
                serde_json::Value::Array(a) => {
                    eos.extend(a.iter().filter_map(|x| x.as_u64()).map(|x| x as u32));
                }
                _ => {}
            }
        }
        eos.dedup();
    }
    for t in ["<|im_end|>", "<end_of_turn>"] {
        if let Some(id) = tok.token_to_id(t)
            && !eos.contains(&id)
        {
            eos.push(id);
        }
    }
    eos
}

/// The generation state, or a 400 telling the caller this server has no `--model`.
fn require_gen(st: &AppState) -> Result<GenState, ApiError> {
    st.generation.clone().ok_or_else(|| {
        ApiError::invalid_request("no generation model loaded (start with --model <dir>)")
    })
}

/// Reject a request naming a model this server did not load (OpenAI returns 404 for that).
fn check_model(model: &Option<String>, generation: &GenState) -> Result<(), ApiError> {
    match model {
        Some(m) if m != &generation.model_name => Err(ApiError::not_found(format!(
            "model `{m}` not found (this server serves `{}`)",
            generation.model_name
        ))
        .with_param("model")
        .with_code("model_not_found")),
        _ => Ok(()),
    }
}

/// Map the validated wire params onto the engine's per-request bundle.
fn build_request_params(c: &CommonParams) -> RequestParams {
    RequestParams {
        sampling: SamplingParams {
            temperature: c.temperature,
            top_p: c.top_p,
            seed: c.seed,
        },
        stop_token_ids: c.stop_token_ids.clone(),
        presence_penalty: c.presence_penalty,
        frequency_penalty: c.frequency_penalty,
        repetition_penalty: c.repetition_penalty,
        logit_bias: c.logit_bias.clone(),
        top_k: c.top_k,
        min_p: c.min_p,
        logprobs: c.logprobs,
    }
}

/// One emitted token's logprob info in wire form (token ids decoded to their text pieces).
#[derive(Clone)]
struct TokenLp {
    token: String,
    logprob: f32,
    top: Vec<(String, f32)>,
}

/// Decode a single token id to its text piece (for logprobs display).
fn decode_piece(tok: &tokenizers::Tokenizer, id: u32) -> String {
    tok.decode(&[id], false).unwrap_or_default()
}

/// Build the OpenAI `logprobs` object for a choice, or `null` when none were collected. The chat
/// spelling nests under `content`; the completion spelling uses the parallel-array form. In both,
/// the chosen token's `logprob` is exactly its entry inside `top_logprobs`.
fn logprobs_json(kind: RespKind, lps: &[TokenLp]) -> serde_json::Value {
    if lps.is_empty() {
        return serde_json::Value::Null;
    }
    let top_obj = |lp: &TokenLp| -> Vec<serde_json::Value> {
        lp.top
            .iter()
            .map(|(t, v)| serde_json::json!({"token": t, "logprob": v}))
            .collect()
    };
    match kind {
        RespKind::Chat => {
            let content: Vec<serde_json::Value> = lps
                .iter()
                .map(|lp| {
                    serde_json::json!({
                        "token": lp.token,
                        "logprob": lp.logprob,
                        "top_logprobs": top_obj(lp),
                    })
                })
                .collect();
            serde_json::json!({ "content": content })
        }
        RespKind::Completion => serde_json::json!({
            "tokens": lps.iter().map(|lp| lp.token.clone()).collect::<Vec<_>>(),
            "token_logprobs": lps.iter().map(|lp| lp.logprob).collect::<Vec<_>>(),
            "top_logprobs": lps
                .iter()
                .map(|lp| lp.top.iter().map(|(t, v)| (t.clone(), *v)).collect::<std::collections::BTreeMap<_, _>>())
                .collect::<Vec<_>>(),
        }),
    }
}

async fn completions(
    State(st): State<AppState>,
    Json(req): Json<CompletionReq>,
) -> axum::response::Response {
    let generation = match require_gen(&st) {
        Ok(g) => g,
        Err(e) => return e.into_response(),
    };
    if let Err(e) = check_model(&req.model, &generation) {
        return e.into_response();
    }
    let common = match req.common() {
        Ok(c) => c,
        Err(e) => return e.into_response(),
    };
    let prompt = match req.prompt_text() {
        Ok(p) => p,
        Err(e) => return e.into_response(),
    };
    let ids = match generation.tokenizer.encode(prompt, true) {
        Ok(e) => e.get_ids().to_vec(),
        Err(e) => return ApiError::invalid_request(format!("tokenize: {e}")).into_response(),
    };
    let images = match decode_images(&req.images) {
        Ok(v) => v,
        Err(e) => return e.into_response(),
    };
    run_with_images(generation, ids, images, common, RespKind::Completion, false).await
}

/// Base64 (raw, or a `data:...;base64,` URI) → bytes. The image is NOT decoded here — that is the
/// engine thread's job, on the thread that owns the GPU.
fn decode_images(b64s: &[String]) -> Result<Vec<Vec<u8>>, ApiError> {
    b64s.iter()
        .map(|s| {
            let payload = match s.split_once(";base64,") {
                Some((_, tail)) => tail,
                None => s.as_str(),
            };
            inferencelayer::serve::b64::decode(payload.trim()).map_err(|e| {
                ApiError::invalid_request(format!("image base64: {e}")).with_param("images")
            })
        })
        .collect()
}

async fn chat_completions(
    State(st): State<AppState>,
    Json(req): Json<ChatReq>,
) -> axum::response::Response {
    let generation = match require_gen(&st) {
        Ok(g) => g,
        Err(e) => return e.into_response(),
    };
    if let Err(e) = check_model(&req.model, &generation) {
        return e.into_response();
    }
    let common = match req.common() {
        Ok(c) => c,
        Err(e) => return e.into_response(),
    };
    // Tools to advertise this request (empty under `tool_choice: "none"` or when none were sent).
    let tool_defs = match req.tools_to_render() {
        Ok(t) => t,
        Err(e) => return e.into_response(),
    };
    // Streaming tool-call deltas (the incremental `tool_calls[].function.arguments` fragment shape)
    // are a distinct unimplemented feature — reject loudly rather than silently buffer.
    if common.stream && !tool_defs.is_empty() {
        return ApiError::invalid_request(
            "streaming is not supported together with tools — send stream:false for tool calls",
        )
        .with_param("stream")
        .into_response();
    }
    let tools_block = (!tool_defs.is_empty()).then(|| tools::tools_system_block(&tool_defs));
    let turns: Vec<ChatTurn> = req
        .messages
        .iter()
        .map(|m| ChatTurn {
            role: &m.role,
            content: &m.content,
            tool_calls: m
                .tool_calls
                .iter()
                .map(|tc| ToolCall {
                    name: &tc.function.name,
                    arguments: &tc.function.arguments,
                })
                .collect(),
        })
        .collect();
    let prompt = chat::render_chat_prompt(generation.arch, &turns, tools_block.as_deref());
    let ids = match generation.tokenizer.encode(prompt.as_str(), true) {
        Ok(e) => e.get_ids().to_vec(),
        Err(e) => return ApiError::invalid_request(format!("tokenize: {e}")).into_response(),
    };
    // Parse `<tool_call>` blocks out of the reply exactly when we advertised tools.
    run(
        generation,
        ids,
        common,
        RespKind::Chat,
        !tool_defs.is_empty(),
    )
    .await
}

// ============================================================================
// /v1/embeddings
// ============================================================================

#[derive(Deserialize)]
#[serde(untagged)]
enum EmbedInput {
    One(String),
    Many(Vec<String>),
}

#[derive(Deserialize)]
struct EmbeddingsReq {
    input: EmbedInput,
    #[serde(default)]
    #[allow(dead_code)]
    // accepted for OpenAI compatibility; the loaded checkpoint is the model
    model: Option<String>,
    #[serde(default)]
    encoding_format: Option<String>,
    /// Matryoshka truncation: keep the first N dims, re-L2-normalize.
    #[serde(default)]
    dimensions: Option<usize>,
}

#[derive(Serialize)]
struct EmbeddingItem {
    object: &'static str,
    index: usize,
    /// A float array (`encoding_format: "float"`) or a base64 string of the f32 LE bytes
    /// (`encoding_format: "base64"`).
    embedding: serde_json::Value,
}

#[derive(Serialize)]
struct EmbedUsage {
    prompt_tokens: usize,
    total_tokens: usize,
}

#[derive(Serialize)]
struct EmbeddingsResp {
    object: &'static str,
    data: Vec<EmbeddingItem>,
    model: String,
    usage: EmbedUsage,
}

async fn embeddings(
    State(st): State<AppState>,
    Json(req): Json<EmbeddingsReq>,
) -> axum::response::Response {
    let Some(es) = st.embed.clone() else {
        return (
            axum::http::StatusCode::BAD_REQUEST,
            "no embedding model loaded (start with --embed-model <dir>)",
        )
            .into_response();
    };
    let base64 = match req.encoding_format.as_deref() {
        None | Some("float") => false,
        Some("base64") => true,
        Some(fmt) => {
            return (
                axum::http::StatusCode::BAD_REQUEST,
                format!("encoding_format `{fmt}` not supported (float | base64)"),
            )
                .into_response();
        }
    };
    // Late-interaction (ColBERT) checkpoints have no single per-input vector — point at /v1/rerank.
    if es.is_per_token {
        return (
            axum::http::StatusCode::BAD_REQUEST,
            "this checkpoint produces per-token (ColBERT) embeddings; use POST /v1/rerank",
        )
            .into_response();
    }
    let texts = match req.input {
        EmbedInput::One(s) => vec![s],
        EmbedInput::Many(v) => v,
    };
    if texts.is_empty() {
        return (
            axum::http::StatusCode::BAD_REQUEST,
            "input must not be empty",
        )
            .into_response();
    }
    if let Some(d) = req.dimensions
        && (d == 0 || d > es.dimension)
    {
        return (
            axum::http::StatusCode::BAD_REQUEST,
            format!("dimensions must be in 1..={} for this model", es.dimension),
        )
            .into_response();
    }
    let (tx, rx) = tokio::sync::oneshot::channel();
    if es
        .submit
        .send(EmbedJob::Embed(EmbedReq {
            texts,
            dimensions: req.dimensions,
            reply: tx,
        }))
        .is_err()
    {
        return (
            axum::http::StatusCode::SERVICE_UNAVAILABLE,
            "embedding engine thread gone",
        )
            .into_response();
    }
    match rx.await {
        Ok(Ok(done)) => {
            let data: Vec<EmbeddingItem> = done
                .vectors
                .into_iter()
                .enumerate()
                .map(|(index, embedding)| EmbeddingItem {
                    object: "embedding",
                    index,
                    embedding: if base64 {
                        serde_json::Value::String(inferencelayer::serve::b64::encode_f32_le(
                            &embedding,
                        ))
                    } else {
                        serde_json::json!(embedding)
                    },
                })
                .collect();
            Json(EmbeddingsResp {
                object: "list",
                data,
                model: es.model_name.clone(),
                usage: EmbedUsage {
                    prompt_tokens: done.prompt_tokens,
                    total_tokens: done.prompt_tokens,
                },
            })
            .into_response()
        }
        Ok(Err(e)) => (axum::http::StatusCode::BAD_REQUEST, e).into_response(),
        Err(_) => (
            axum::http::StatusCode::INTERNAL_SERVER_ERROR,
            "embedding engine dropped the request",
        )
            .into_response(),
    }
}

// ============================================================================
// /v1/rerank (Cohere/Jina shape)
// ============================================================================

#[derive(Deserialize)]
struct RerankApiReq {
    #[serde(default)]
    #[allow(dead_code)] // accepted for compatibility; the loaded checkpoint is the model
    model: Option<String>,
    query: String,
    documents: Vec<String>,
    #[serde(default)]
    top_n: Option<usize>,
    #[serde(default)]
    return_documents: bool,
}

async fn rerank(
    State(st): State<AppState>,
    Json(req): Json<RerankApiReq>,
) -> axum::response::Response {
    let Some(es) = st.embed.clone() else {
        return (
            axum::http::StatusCode::BAD_REQUEST,
            "no embedding model loaded (start with --embed-model <dir>)",
        )
            .into_response();
    };
    if req.documents.is_empty() {
        return (
            axum::http::StatusCode::BAD_REQUEST,
            "documents must not be empty",
        )
            .into_response();
    }
    let (tx, rx) = tokio::sync::oneshot::channel();
    if es
        .submit
        .send(EmbedJob::Rerank(RerankReq {
            query: req.query,
            documents: req.documents.clone(),
            top_n: req.top_n,
            reply: tx,
        }))
        .is_err()
    {
        return (
            axum::http::StatusCode::SERVICE_UNAVAILABLE,
            "embedding engine thread gone",
        )
            .into_response();
    }
    match rx.await {
        Ok(Ok(done)) => {
            let results: Vec<serde_json::Value> = done
                .ranked
                .iter()
                .map(|(idx, score)| {
                    let mut obj = serde_json::json!({"index": idx, "relevance_score": score});
                    if req.return_documents {
                        obj["document"] = serde_json::json!({"text": req.documents[*idx]});
                    }
                    obj
                })
                .collect();
            Json(serde_json::json!({
                "model": es.model_name,
                "results": results,
                "usage": {"total_tokens": done.prompt_tokens},
            }))
            .into_response()
        }
        Ok(Err(e)) => (axum::http::StatusCode::BAD_REQUEST, e).into_response(),
        Err(_) => (
            axum::http::StatusCode::INTERNAL_SERVER_ERROR,
            "embedding engine dropped the request",
        )
            .into_response(),
    }
}

/// One completed choice (non-streaming assembly).
struct ChoiceOut {
    text: String,
    finish: Finish,
    logprobs: Vec<TokenLp>,
}

/// Build one chat choice object. When `parse_tools`, `<tool_call>` blocks are pulled out of the
/// text into OpenAI `tool_calls` and `finish_reason` is overridden to `"tool_calls"` (the leftover
/// prose becomes `content`, or JSON `null` when the reply was calls-only). Otherwise it is a plain
/// assistant message. The stop-driven `finish_reason` is kept when no tool call was actually parsed,
/// so a tools-enabled request that answers in prose still reports `"stop"`.
fn chat_choice_json(
    index: usize,
    c: &ChoiceOut,
    parse_tools: bool,
    logprobs: serde_json::Value,
) -> serde_json::Value {
    if parse_tools {
        let (content, calls) = tools::parse_tool_calls(&c.text);
        if !calls.is_empty() {
            let tool_calls: Vec<serde_json::Value> = calls
                .iter()
                .map(|tc| {
                    serde_json::json!({
                        "id": response_id("call"),
                        "type": "function",
                        "function": {"name": tc.name, "arguments": tc.arguments},
                    })
                })
                .collect();
            let content = if content.is_empty() {
                serde_json::Value::Null
            } else {
                serde_json::Value::String(content)
            };
            return serde_json::json!({
                "index": index,
                "message": {"role": "assistant", "content": content, "tool_calls": tool_calls},
                "finish_reason": "tool_calls",
                "stop_reason": c.finish.stop_reason,
                "logprobs": logprobs,
            });
        }
    }
    serde_json::json!({
        "index": index,
        "message": {"role": "assistant", "content": c.text},
        "finish_reason": c.finish.reason,
        "stop_reason": c.finish.stop_reason,
        "logprobs": logprobs,
    })
}

async fn run(
    generation: GenState,
    ids: Vec<u32>,
    common: CommonParams,
    kind: RespKind,
    parse_tools: bool,
) -> axum::response::Response {
    run_with_images(generation, ids, Vec::new(), common, kind, parse_tools).await
}

async fn run_with_images(
    generation: GenState,
    ids: Vec<u32>,
    images: Vec<Vec<u8>>,
    common: CommonParams,
    kind: RespKind,
    // Chat + tools only: extract `<tool_call>` blocks from each choice's text into structured
    // `tool_calls`. Always false for completions and for tool-less chat.
    parse_tools: bool,
) -> axum::response::Response {
    let (tx, mut rx) = amp::unbounded_channel();
    let n = common.n;
    // Vision encode HERE (concurrently across requests, off the engine thread). The tower is
    // GPU-resident and &self-shareable; spawn_blocking keeps the async runtime unblocked.
    let (ids, vprompt) = if images.is_empty() {
        (ids, None)
    } else {
        let Some(vx) = generation.vision.clone() else {
            return ApiError::invalid_request(
                "this checkpoint has no vision tower; images are not accepted",
            )
            .into_response();
        };
        let gctx = generation.gctx.clone();
        let _permit = generation
            .vision_gate
            .clone()
            .acquire_owned()
            .await
            .expect("vision gate never closes");
        let _t_vis = std::time::Instant::now();
        match tokio::task::spawn_blocking(move || vx.prepare(&gctx, &ids, &images)).await {
            Ok(Ok((expanded, vp))) => {
                VISION_US.fetch_add(
                    _t_vis.elapsed().as_micros() as u64,
                    std::sync::atomic::Ordering::Relaxed,
                );
                VISION_CALLS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
                (expanded, Some(vp))
            }
            Ok(Err(e)) => {
                return ApiError::invalid_request(format!("image: {e:#}")).into_response()
            }
            Err(e) => return ApiError::unavailable(format!("vision task: {e}")).into_response(),
        }
    };
    let request = EngineRequest {
        prompt_ids: ids,
        vprompt,
        max_tokens: common.max_tokens,
        n,
        params: build_request_params(&common),
        stop_strings: common.stop.clone(),
        reply: tx,
    };
    if generation.submit.send(request).is_err() {
        return ApiError::unavailable("engine thread gone").into_response();
    }
    if common.stream {
        return Sse::new(SseStream {
            rx,
            model: generation.model_name.clone(),
            kind,
            id: response_id(kind.id_prefix()),
            created: created_epoch(),
            remaining: n,
            role_sent: vec![false; n],
            include_usage: common.include_usage,
            prompt_tokens: 0,
            completion_tokens: 0,
            cached_tokens: 0,
            queued: VecDeque::new(),
            ended: false,
            finalized: false,
        })
        .into_response();
    }
    // Non-streaming: collect one Done per choice, then assemble the grouped response.
    let mut choices: Vec<Option<ChoiceOut>> = (0..n).map(|_| None).collect();
    let mut remaining = n;
    let mut prompt_tokens = 0;
    let mut completion_tokens = 0;
    let mut cached_tokens = 0;
    while remaining > 0 {
        match rx.recv().await {
            Some(EngineEvent::Delta { .. }) => {}
            Some(EngineEvent::Done {
                choice,
                text,
                completion_tokens: ct,
                prompt_tokens: pt,
                cached_tokens: cached,
                finish,
                logprobs,
            }) => {
                prompt_tokens = pt;
                cached_tokens = cached;
                completion_tokens += ct;
                if let Some(slot) = choices.get_mut(choice) {
                    *slot = Some(ChoiceOut {
                        text,
                        finish,
                        logprobs,
                    });
                }
                remaining -= 1;
            }
            Some(EngineEvent::Error(e)) => return ApiError::internal(e).into_response(),
            Some(EngineEvent::Rejected(e)) => {
                return ApiError::invalid_request(e)
                    .with_code("context_length_exceeded")
                    .with_param("max_tokens")
                    .into_response();
            }
            None => break, // engine thread gone before every choice finished
        }
    }
    let usage = Usage::new(prompt_tokens, completion_tokens, cached_tokens);
    let choice_json: Vec<serde_json::Value> = choices
        .into_iter()
        .enumerate()
        .map(|(i, c)| {
            let c = c.unwrap_or(ChoiceOut {
                text: String::new(),
                finish: Finish {
                    reason: "stop",
                    stop_reason: None,
                },
                logprobs: Vec::new(),
            });
            let logprobs = logprobs_json(kind, &c.logprobs);
            match kind {
                RespKind::Completion => serde_json::json!({
                    "index": i,
                    "text": c.text,
                    "finish_reason": c.finish.reason,
                    "stop_reason": c.finish.stop_reason,
                    "logprobs": logprobs,
                }),
                RespKind::Chat => chat_choice_json(i, &c, parse_tools, logprobs),
            }
        })
        .collect();
    Json(serde_json::json!({
        "id": response_id(kind.id_prefix()),
        "object": kind.object(),
        "created": created_epoch(),
        "model": generation.model_name,
        "choices": choice_json,
        "usage": usage,
    }))
    .into_response()
}

/// SSE stream over the engine's reply channel, multiplexing the `n` choices into OpenAI chunk
/// objects (hand-implemented `Stream`: one tiny trait dep instead of a combinator stack).
struct SseStream {
    rx: amp::UnboundedReceiver<EngineEvent>,
    model: String,
    kind: RespKind,
    id: String,
    created: u64,
    remaining: usize,
    /// Chat only: whether each choice's first `delta` (which carries `role`) has been sent.
    role_sent: Vec<bool>,
    include_usage: bool,
    prompt_tokens: usize,
    completion_tokens: usize,
    cached_tokens: usize,
    queued: VecDeque<Event>,
    ended: bool,
    finalized: bool,
}

impl SseStream {
    fn delta_chunk(&mut self, choice: usize, text: &str) -> serde_json::Value {
        let choices = match self.kind {
            RespKind::Completion => serde_json::json!([{"index": choice, "text": text}]),
            RespKind::Chat => {
                let delta = if !self.role_sent[choice] {
                    self.role_sent[choice] = true;
                    serde_json::json!({"role": "assistant", "content": text})
                } else {
                    serde_json::json!({"content": text})
                };
                serde_json::json!([{"index": choice, "delta": delta}])
            }
        };
        serde_json::json!({
            "id": self.id,
            "object": self.kind.chunk_object(),
            "created": self.created,
            "model": self.model,
            "choices": choices,
        })
    }

    fn finish_chunk(&mut self, choice: usize, finish: &Finish) -> serde_json::Value {
        let choices = match self.kind {
            RespKind::Completion => serde_json::json!([{
                "index": choice,
                "text": "",
                "finish_reason": finish.reason,
                "stop_reason": finish.stop_reason,
            }]),
            RespKind::Chat => {
                let delta = if !self.role_sent[choice] {
                    self.role_sent[choice] = true;
                    serde_json::json!({"role": "assistant", "content": ""})
                } else {
                    serde_json::json!({})
                };
                serde_json::json!([{
                    "index": choice,
                    "delta": delta,
                    "finish_reason": finish.reason,
                    "stop_reason": finish.stop_reason,
                }])
            }
        };
        serde_json::json!({
            "id": self.id,
            "object": self.kind.chunk_object(),
            "created": self.created,
            "model": self.model,
            "choices": choices,
        })
    }

    /// Push the optional usage chunk (when `stream_options.include_usage`) and the `[DONE]`
    /// sentinel exactly once, at the end of the stream.
    fn finalize(&mut self) {
        if self.finalized {
            return;
        }
        self.finalized = true;
        if self.include_usage {
            let usage = Usage::new(
                self.prompt_tokens,
                self.completion_tokens,
                self.cached_tokens,
            );
            let chunk = serde_json::json!({
                "id": self.id,
                "object": self.kind.chunk_object(),
                "created": self.created,
                "model": self.model,
                "choices": [],
                "usage": usage,
            });
            self.queued
                .push_back(Event::default().data(chunk.to_string()));
        }
        self.queued.push_back(Event::default().data("[DONE]"));
        self.ended = true;
    }
}

impl futures_core::Stream for SseStream {
    type Item = Result<Event, std::convert::Infallible>;
    fn poll_next(
        mut self: std::pin::Pin<&mut Self>,
        cx: &mut std::task::Context<'_>,
    ) -> std::task::Poll<Option<Self::Item>> {
        use std::task::Poll;
        loop {
            if let Some(ev) = self.queued.pop_front() {
                return Poll::Ready(Some(Ok(ev)));
            }
            if self.ended {
                return Poll::Ready(None);
            }
            match self.rx.poll_recv(cx) {
                Poll::Pending => return Poll::Pending,
                Poll::Ready(None) => {
                    // Engine thread gone: finalize whatever we have.
                    self.finalize();
                }
                Poll::Ready(Some(EngineEvent::Delta { choice, text })) => {
                    let chunk = self.delta_chunk(choice, &text);
                    self.queued
                        .push_back(Event::default().data(chunk.to_string()));
                }
                Poll::Ready(Some(EngineEvent::Done {
                    choice,
                    completion_tokens,
                    prompt_tokens,
                    cached_tokens,
                    finish,
                    ..
                })) => {
                    self.prompt_tokens = prompt_tokens;
                    self.cached_tokens = cached_tokens;
                    self.completion_tokens += completion_tokens;
                    let chunk = self.finish_chunk(choice, &finish);
                    self.queued
                        .push_back(Event::default().data(chunk.to_string()));
                    self.remaining = self.remaining.saturating_sub(1);
                    if self.remaining == 0 {
                        self.finalize();
                    }
                }
                Poll::Ready(Some(EngineEvent::Error(e)))
                | Poll::Ready(Some(EngineEvent::Rejected(e))) => {
                    self.queued
                        .push_back(Event::default().data(format!("{{\"error\":{e:?}}}")));
                    self.finalize();
                }
            }
        }
    }
}

async fn models(State(st): State<AppState>) -> impl IntoResponse {
    let created = created_epoch();
    let mut data = Vec::new();
    if let Some(g) = &st.generation {
        data.push(serde_json::json!({
            "id": g.model_name, "object": "model", "created": created, "owned_by": "inferencelayer",
        }));
    }
    if let Some(e) = &st.embed {
        data.push(serde_json::json!({
            "id": e.model_name, "object": "model", "created": created, "owned_by": "inferencelayer",
        }));
    }
    Json(serde_json::json!({"object": "list", "data": data}))
}

async fn health() -> impl IntoResponse {
    Json(serde_json::json!({"status": "ok"}))
}

async fn metrics(State(st): State<AppState>) -> impl IntoResponse {
    let (s, cache, free_blocks) = st.stats.lock().expect("stats lock").clone();
    Json(serde_json::json!({
        "steps": s.steps,
        "kv_free_blocks": free_blocks,
        "tokens_out": s.tokens_out,
        "prefill_tokens": s.prefill_tokens,
        "cache_hit_tokens": s.cache_hit_tokens,
        "admitted": s.admitted,
        "finished": s.finished,
        "rejected_full": s.rejected_full,
        "cancelled": s.cancelled,
        "preempted": s.preempted,
        "wide_steps": s.wide_steps,
        "spec_rounds": s.spec_rounds,
        "spec_drafted": s.spec_drafted,
        "spec_accepted": s.spec_accepted,
        "pld_hits": s.pld_hits,
        "radix_entries": cache,
        "wide_us": s.wide_us,
        "decode_us": s.decode_us,
        "vision_us": VISION_US.load(std::sync::atomic::Ordering::Relaxed),
        "vision_calls": VISION_CALLS.load(std::sync::atomic::Ordering::Relaxed),
        "vision_pre_us": VISION_PRE_US.load(std::sync::atomic::Ordering::Relaxed),
        "emit_us": EMIT_US.load(std::sync::atomic::Ordering::Relaxed),
        "step_us": s.step_us,
        "build_us": s.build_us,
        "post_us": s.post_us,
        "chain_us": s.chain_us,
        "chained_steps": s.chained_steps,
    }))
}