engine/crates/server/src/lib.rs

//! dllm-server: the HTTP layer of span-02 (`POST /v1/systemone`, `GET /health`): request validation, prompt construction,
//! answer assembly (blend, choice / noul / score fields), and the GPU engine behind it (`Args` = the command line, `start_engine` loads the model,
//! the MTP draft and the vision tower and starts the scheduler).
pub mod assets;
pub mod wire;

use axum::{
    body::Bytes,
    extract::State,
    http::{header, HeaderMap, StatusCode},
    response::{IntoResponse, Response as HttpResponse},
    routing::{get, post},
    Json, Router,
};
use dllm_core::policy::{self, Answer, Policy, ThinkMode};
use dllm_core::tokenizer::{img_soft_max, ostr, state_text, Tok};
use serde_json::{json, Value};
use std::sync::Arc;
use wire::{AnswerOut, ApiError, OrderedMap, ReasoningOut, Request, Response, Usage};

pub use dllm_core::serving::{Item, ItemResult};

pub trait Backend: Send + Sync {
    /// All questions of ONE request share the state prefix (and images): the engine batches their branches in one forward.
    fn run(&self, items: &[Item], images: &[Vec<u8>], pol: &Policy) -> anyhow::Result<Vec<ItemResult>>;
}

pub struct AppState {
    pub tok: Tok,
    pub backend: Box<dyn Backend>,
    pub pol: Policy,
    pub api_key: Option<String>,
    pub max_len: usize,
    pub max_body: usize,
    pub served_name: String,
}

pub fn app(st: Arc<AppState>) -> Router {
    Router::new().route("/health", get(health)).route("/v1/systemone", post(systemone)).fallback(not_found).with_state(st)
}

/// any other path: 404 in the error format, after reading the body (a proxy in front then sees a complete exchange instead of a reset)
async fn not_found(State(st): State<Arc<AppState>>, body: axum::body::Body) -> HttpResponse {
    let _ = axum::body::to_bytes(body, st.max_body).await;
    err(ApiError::new(404, "not found: the API is POST /v1/systemone and GET /health").kind("not_found"))
}

async fn health(State(st): State<Arc<AppState>>) -> Json<Value> {
    Json(json!({"ok": true, "model": st.served_name}))
}

fn err(e: ApiError) -> HttpResponse {
    (StatusCode::from_u16(e.code).unwrap(), Json(json!({"error": {"type": e.kind, "message": e.msg}}))).into_response()
}

pub fn decode_images(v: &Option<Value>) -> Result<Vec<Vec<u8>>, ApiError> {
    let bad = || ApiError::new(422, "images must be a list of at most 4 data URIs");
    let Some(v) = v else { return Ok(vec![]) };
    if v.is_null() {
        return Ok(vec![]);
    }
    let a = v.as_array().filter(|a| a.len() <= 4).ok_or_else(bad)?;
    a.iter()
        .map(|s| {
            let s = s.as_str().ok_or_else(bad)?;
            let b64 = s.split_once(',').map(|x| x.1).filter(|_| s.starts_with("data:")).unwrap_or(s);
            base64_simd::STANDARD.decode_to_vec(b64.trim()).map_err(|_| ApiError::new(422, "cannot decode image as a base64 data URI"))
        })
        .collect()
}

/// thinking tokens reserved in the context check of one question with `rows` prompt rows (images expanded at the maximum): the hard cap, or with
/// `--think-ms` the cap the deterministic cost model gives it as a lone thinking question (the engine also clamps every thought to the room left
/// in the context, so a request whose real cap comes out larger still fits)
pub fn think_reserve(pol: &Policy, rows: usize, n_images: usize) -> usize {
    if pol.think_ms == 0 {
        return pol.cap;
    }
    let s = dllm_core::policy::ThinkShape { prefill_tokens: rows, images: n_images, image_bytes: 0, questions: 1, thinking: 1 };
    pol.think_cost.cap(pol.think_ms as f64, &s, pol.cap)
}

/// per-request policy: the server policy with `reasoning` applied
pub fn request_policy(base: &Policy, req: &Request) -> Result<Policy, ApiError> {
    let r = req.reasoning.as_ref();
    let (mode, budget, trace) = r.map_or((None, None, None), |r| (r.mode.as_deref(), r.budget_ms, r.trace));
    let m = match mode {
        None | Some("auto") => ThinkMode::Auto,
        Some("on") => ThinkMode::On,
        Some("off") => ThinkMode::Off,
        Some(o) => return Err(ApiError::new(422, format!("reasoning.mode must be \"auto\", \"on\" or \"off\", got {o:?}"))),
    };
    // wall-clock budget: the request's `budget_ms` if sent; mode "on" without one thinks to completion (up to the cap); otherwise the server's.
    // `--think-ms` (the deterministic token cap) applies in every case; `budget_ms` only adds or removes the wall-clock deadline.
    let budget_ms = match budget {
        Some(b) if !(b.is_finite() && (0.0..=MAX_BUDGET_MS).contains(&b)) => return Err(ApiError::new(422, format!("reasoning.budget_ms must be a number of milliseconds in [0, {MAX_BUDGET_MS}], got {b}"))),
        Some(b) => b.round() as u32,
        None if m == ThinkMode::On => 0,
        None => base.budget_ms,
    };
    Ok(Policy { return_trace: trace.unwrap_or(false), budget_ms, ..base.with_think(m) })
}

/// largest accepted per-request `budget_ms` (10 minutes)
pub const MAX_BUDGET_MS: f64 = 600_000.0;

/// TypeSafe score confidence: max(0, 1 - E|i - m| / MAD_unif), m = most likely level (first on ties), MAD_unif = mean |i - (n-1)/2|
pub fn score_confidence(p: &[f32]) -> f32 {
    let n = p.len();
    if n < 2 {
        return 1.0;
    }
    let m = policy::argmax_first(p) as f64;
    let mad: f64 = p.iter().enumerate().map(|(i, &q)| q as f64 * (i as f64 - m).abs()).sum();
    let c = (n as f64 - 1.0) / 2.0;
    let unif = (0..n).map(|i| (i as f64 - c).abs()).sum::<f64>() / n as f64;
    (1.0 - mad / unif).max(0.0) as f32
}

/// noul option texts: A = true side, B = false side (our trained order). No criteria (absent, null or an empty object `{}`, which DI
/// rows send, e.g. calibration-v1 phishing signals) renders exactly as before ("True"/"False"); explicit criteria must be an object with
/// both "true" and "false" (a null side keeps its default).
fn noul_options(key: &str, c: Option<&Value>) -> Result<Vec<String>, ApiError> {
    let Some(c) = c.filter(|c| !c.is_null() && c.as_object().map_or(true, |o| !o.is_empty())) else { return Ok(vec!["True".into(), "False".into()]) };
    let bad = || ApiError::new(422, format!("question {key}: noul criteria must be an object with exactly the keys \"true\" and \"false\""));
    let o = c.as_object().filter(|o| o.len() == 2).ok_or_else(bad)?;
    let side = |k: &str, d: &str| o.get(k).map(|v| if v.is_null() { d.to_string() } else { ostr(v) }).ok_or_else(bad);
    Ok(vec![side("true", "True")?, side("false", "False")?])
}

/// TypeSafe choice confidence: the returned choice's probability rescaled from uniform (0) to certain (1)
pub fn choice_confidence(p_max: f32, n: usize) -> f32 {
    if n < 2 {
        return 1.0;
    }
    let u = 1.0 / n as f64;
    ((p_max as f64 - u) / (1.0 - u)).clamp(0.0, 1.0) as f32
}

pub const SCORE_MAX_LEVELS: usize = 10;

pub fn build_items(st: &AppState, req: &Request, n_images: usize) -> Result<Vec<Item>, ApiError> {
    build_questions(st, req, n_images).map(|x| x.0)
}

/// `build_prompts` over the request's own state and questions
pub fn build_questions(st: &AppState, req: &Request, n_images: usize) -> Result<(Vec<Item>, Vec<Option<Vec<Value>>>), ApiError> {
    build_prompts(st, &req.state, &req.questions, n_images)
}

/// items + per-item score legend (Some(levels as sent) for `score` questions; they run exactly like choice with keys "0".."n-1")
pub fn build_prompts(st: &AppState, state: &Value, questions: &serde_json::Map<String, Value>, n_images: usize) -> Result<(Vec<Item>, Vec<Option<Vec<Value>>>), ApiError> {
    if questions.is_empty() {
        return Err(ApiError::new(422, "questions must be a non-empty object"));
    }
    let stxt = state_text(state);
    let mut items = vec![];
    let mut legends = vec![];
    for (key, q) in questions {
        let q = q.as_object().ok_or_else(|| ApiError::new(422, format!("question {key}: unsupported type")))?;
        let ty = q.get("type").and_then(|x| x.as_str());
        let mut legend = None;
        if ty == Some("number") {
            // free-form number (our extension): no criteria; the answer is generated after the end-of-thought token
            let instr = q.get("instructions").map(ostr).unwrap_or_default();
            let b = st.tok.build_number(&stxt, &instr, n_images).map_err(|e| ApiError::new(422, format!("bad request: {e}")))?;
            let reserve = think_reserve(&st.pol, b.ids.len() + (img_soft_max() + 2) * n_images, n_images);
            let need = b.ids.len() + (img_soft_max() + 2) * n_images + reserve + 3 + dllm_runtime::engine::ANSWER_MAX;
            if need > st.max_len {
                return Err(ApiError::context(format!("This model's maximum context length is {} tokens. However, the request needs about {} tokens (prompt + thinking budget {} + answer).", st.max_len, need, reserve)));
            }
            items.push(Item { key: key.clone(), number: true, ids: b.ids, ..Default::default() });
            legends.push(None);
            continue;
        }
        let (noul, keys, opts): (bool, Vec<String>, Vec<String>) = match ty {
            Some("noul") => (true, vec!["true".into(), "false".into()], noul_options(key, q.get("criteria"))?),
            Some("choice") => {
                let c = q.get("criteria").and_then(|x| x.as_object()).filter(|c| !c.is_empty()).ok_or_else(|| ApiError::new(422, format!("question {key}: criteria must be a non-empty object")))?;
                // a null description defaults to its option key (TypeSafe / surogate)
                (false, c.keys().cloned().collect(), c.iter().map(|(k, v)| if v.is_null() { k.clone() } else { ostr(v) }).collect())
            }
            Some("score") => {
                let c = q.get("criteria").and_then(|x| x.as_array()).ok_or_else(|| ApiError::new(422, format!("question {key}: criteria must be an array of level descriptions")))?;
                if c.len() < 2 || c.len() > SCORE_MAX_LEVELS {
                    return Err(ApiError::new(422, format!("question {key}: score criteria must have between 2 and {SCORE_MAX_LEVELS} levels, got {}", c.len())));
                }
                legend = Some(c.clone());
                (false, (0..c.len()).map(|i| i.to_string()).collect(), c.iter().map(ostr).collect())
            }
            _ => return Err(ApiError::new(422, format!("question {key}: unsupported type"))),
        };
        let instr = q.get("instructions").map(ostr).unwrap_or_default();
        let b = st
            .tok
            .build(&stxt, &instr, &opts, n_images)
            .map_err(|e| ApiError::new(422, format!("bad request: {e}")))?
            .ok_or_else(|| ApiError::new(422, format!("question {key}: too many options per choice (max {})", st.tok.labels.len())))?;
        let reserve = think_reserve(&st.pol, b.ids.len() + (img_soft_max() + 2) * n_images, n_images);
        let need = b.ids.len() + (img_soft_max() + 2) * n_images + reserve + 3;
        if need > st.max_len {
            return Err(ApiError::context(format!("This model's maximum context length is {} tokens. However, the request needs about {} tokens (prompt + thinking budget {}).", st.max_len, need, reserve)));
        }
        items.push(Item { key: key.clone(), noul, keys, ids: b.ids, label_ids: b.label_ids, ..Default::default() });
        legends.push(legend);
    }
    Ok((items, legends))
}

async fn systemone(State(st): State<Arc<AppState>>, headers: HeaderMap, body: axum::body::Body) -> HttpResponse {
    let t0 = std::time::Instant::now();
    if let Some(k) = &st.api_key {
        let ok = headers.get(header::AUTHORIZATION).and_then(|h| h.to_str().ok()).map_or(false, |h| constant_eq(h.as_bytes(), format!("Bearer {k}").as_bytes()));
        if !ok {
            return err(ApiError::new(401, "invalid or missing bearer token"));
        }
    }
    let too_big = || err(ApiError::new(413, "request body too large (maximum context length / body size exceeded)"));
    if headers.get(header::CONTENT_LENGTH).and_then(|h| h.to_str().ok()).and_then(|h| h.parse::<usize>().ok()).map_or(false, |n| n > st.max_body) {
        return too_big();
    }
    let body: Bytes = match axum::body::to_bytes(body, st.max_body).await {
        Ok(b) => b,
        Err(_) => return too_big(),
    };
    match handle(&st, &headers, body, t0).await {
        Ok(r) => ok_json(&r),
        Err(e) => err(e),
    }
}

async fn handle(st: &Arc<AppState>, headers: &HeaderMap, body: Bytes, t0: std::time::Instant) -> Result<Response, ApiError> {
    let body = assets::decode_encoding(headers, body.to_vec(), st.max_body)?;
    let mut body = assets::split_body(headers, body).await?;
    let req: Request = parse_request(&body.json)?;
    // images in the order the model sees them: the DI kit `images` (as before), then the assets in the order sent, labelled "Image k" in the text
    let mut images = decode_images(&req.images)?;
    let assets = assets::resolve_assets(&req.assets, &mut body.parts)?;
    if images.len() + assets.len() > assets::MAX_IMAGES {
        return Err(if assets.is_empty() { ApiError::new(422, "images must be a list of at most 4 data URIs") } else { ApiError::asset(format!("at most {} images per request (images + assets), got {}", assets::MAX_IMAGES, images.len() + assets.len())) });
    }
    let rewritten = if assets.is_empty() {
        None
    } else {
        let map: Vec<(String, String)> = assets.iter().enumerate().map(|(i, (ph, _))| (ph.clone(), assets::image_label(images.len() + i))).collect();
        Some(assets::substitute(&req.state, &req.questions, &map))
    };
    let has_assets = !assets.is_empty();
    images.extend(assets.into_iter().map(|(_, b)| b));
    let soft: Vec<usize> = images.iter().map(|b| image_soft_tokens(b)).collect();
    let rpol = request_policy(&st.pol, &req)?;
    let (items, mut legends) = match &rewritten {
        None => build_questions(st, &req, images.len())?,
        Some((s, q)) => build_prompts(st, s, q, images.len())?,
    };
    // a score legend echoes the levels as sent (placeholders, not their labels)
    for ((_, q), l) in req.questions.iter().zip(legends.iter_mut()) {
        if let (Some(l), Some(c)) = (l.as_mut(), q.get("criteria").and_then(|c| c.as_array())) {
            *l = c.clone();
        }
    }
    let st2 = st.clone();
    let items2 = items.clone();
    let res = tokio::task::spawn_blocking(move || {
        dllm_core::serving::set_arrival(t0);
        st2.backend.run(&items2, &images, &rpol)
    })
    .await;
    let res = match res {
        Ok(Ok(r)) if r.len() == items.len() => r,
        Ok(Ok(_)) => return Err(ApiError::new(500, "inference failed: result count mismatch")),
        Ok(Err(e)) => {
            let e = backend_error(&format!("{e:#}"));
            return Err(if has_assets && e.code == 422 { ApiError::asset(e.msg) } else { e });
        }
        Err(e) => return Err(ApiError::new(500, format!("inference failed: {e}"))),
    };
    let mut answers = OrderedMap::new();
    let mut usage = Usage::default();
    let img_rows: usize = soft.iter().map(|n| n + 2).sum();
    for ((it, r), legend) in items.iter().zip(&res).zip(&legends) {
        let a = if it.number {
            number_answer(&st.tok, r)
        } else {
            let p0 = policy::softmax(&r.p0);
            let pt = r.think.as_ref().map(|z| policy::softmax(z));
            match policy::finish(&it.keys, it.noul, &p0, pt.as_deref()) {
                Answer::Noul(x) => AnswerOut::Noul { noul: x },
                Answer::Choice { probs, order, .. } if legend.is_some() => score_answer(&order, &probs, legend.as_ref().unwrap()),
                Answer::Choice { choice, probs, order } => AnswerOut::Choice { confidence: choice_confidence(probs[&choice], order.len()), choice, probabilities: OrderedMap(order.iter().map(|k| (k.clone(), probs[k])).collect()) },
            }
        };
        let reasoning = r.think.as_ref().map(|_| ReasoningOut { tokens: r.trace_tokens, trace: rpol.return_trace.then(|| st.tok.decode_text(&r.trace_ids)) });
        answers.insert(it.key.clone(), wire::Answer { answer: a, reasoning });
        usage.input_tokens += it.ids.len() + img_rows;
        usage.reasoning_tokens += r.trace_tokens;
        usage.output_tokens += r.answer_ids.len();
    }
    usage.latency_ms = t0.elapsed().as_secs_f64() * 1e3;
    let image_tokens = if soft.is_empty() { None } else { Some(soft.iter().sum()) };
    // like TypeSafe (reports the versioned id even for an alias): the served model id, whatever the request named
    Ok(Response { model: st.served_name.clone(), answers, usage, image_tokens })
}

/// soft tokens this image becomes under the current budget (Gemma-4 aspect-preserving resize; header-only probe); the budget if the header is unreadable
fn image_soft_tokens(bytes: &[u8]) -> usize {
    let cfg = dllm_vision::PreprocConfig { max_soft_tokens: img_soft_max(), ..Default::default() };
    dllm_vision::probe(bytes).ok().and_then(|(w, h)| dllm_vision::n_soft_for_size(w, h, &cfg).ok()).unwrap_or_else(img_soft_max)
}

/// number answer: the generated text, its value (None if it does not parse as a number) and confidence = exp(mean log-probability of the answer tokens)
fn number_answer(tok: &Tok, r: &ItemResult) -> AnswerOut {
    let text = tok.decode_text(&r.answer_ids).trim().to_string();
    let number = parse_number(&text);
    let confidence = if r.answer_ids.is_empty() { 0.0 } else { (r.answer_logprob / r.answer_ids.len() as f64).exp() as f32 };
    AnswerOut::Number { number, text, confidence }
}

/// a decimal number as the model writes it (training used `str()` of the value); thousands separators and surrounding spaces are tolerated
pub fn parse_number(text: &str) -> Option<f64> {
    let t: String = text.trim().chars().filter(|&c| c != ',' && c != ' ').collect();
    t.parse::<f64>().ok().filter(|v| v.is_finite())
}

fn score_answer(order: &[String], probs: &std::collections::BTreeMap<String, f32>, legend: &[Value]) -> AnswerOut {
    let p: Vec<f32> = order.iter().map(|k| probs[k]).collect();
    let score = p.iter().enumerate().map(|(i, &q)| i as f32 * q).sum();
    AnswerOut::Score {
        score,
        legend: OrderedMap(legend.iter().enumerate().map(|(i, v)| (i.to_string(), v.clone())).collect()),
        probabilities: OrderedMap(order.iter().cloned().zip(p.iter().cloned()).collect()),
        confidence: score_confidence(&p),
    }
}

fn ok_json<T: serde::Serialize>(v: &T) -> HttpResponse {
    match serde_json::to_vec(v) {
        Ok(b) => (StatusCode::OK, [(header::CONTENT_TYPE, "application/json")], b).into_response(),
        Err(e) => err(ApiError::new(500, format!("serialisation failed: {e}"))),
    }
}

/// simd-json fast path (bodies are dominated by base64 image strings); serde_json re-parse only on failure, to classify the error (400 invalid JSON / 422 bad schema)
fn parse_request(body: &[u8]) -> Result<Request, ApiError> {
    let mut buf = body.to_vec();
    if let Ok(r) = simd_json::serde::from_slice::<Request>(&mut buf) {
        return Ok(r);
    }
    match serde_json::from_slice::<Request>(body) {
        Ok(r) => Ok(r),
        Err(e) if e.is_syntax() || e.is_eof() => Err(ApiError::new(400, "invalid JSON body")),
        Err(e) => Err(ApiError::new(422, format!("bad request: {e}"))),
    }
}

fn constant_eq(a: &[u8], b: &[u8]) -> bool {
    a.len() == b.len() && a.iter().zip(b).fold(0u8, |d, (x, y)| d | (x ^ y)) == 0
}

// ================================================================================================================ command line + engine

/// span-02 server: Gemma-4 26B-A4B decision model on the custom Rust/CUDA engine (sm_120a)
#[derive(clap::Parser, Debug, Clone)]
#[command(about = "span-02 decision server: POST /v1/systemone, GET /health")]
pub struct Args {
    #[arg(long, default_value = "127.0.0.1:8000")]
    pub bind: String,
    /// model directory: config.json, model.safetensors (engine-native quantised checkpoint), tokenizer.json, draft_vocab_16384.json
    #[arg(long)]
    pub model: String,
    /// tokenizer.json (default: <model>/tokenizer.json)
    #[arg(long)]
    pub tokenizer: Option<String>,
    /// MTP draft for speculative decoding of thoughts: directory of google/gemma-4-26B-A4B-it-assistant
    #[arg(long)]
    pub spec_draft: String,
    /// pruned draft vocabulary (default: <model>/draft_vocab_16384.json)
    #[arg(long)]
    pub spec_vocab: Option<String>,
    /// speculative step for up to N simultaneously thinking branches (1..=16; more branches use the plain batched greedy step)
    #[arg(long, default_value_t = 4)]
    pub spec_max_branches: usize,
    /// bearer token every request must carry (none = open)
    #[arg(long, env = "AR_API_KEY")]
    pub api_key: Option<String>,
    /// model id reported in every response's `model` field (the request's `model` is accepted and ignored)
    #[arg(long, default_value = "span-02")]
    pub served_name: String,
    /// think when the no-think answer's top probability is below tau
    #[arg(long, default_value_t = 0.8)]
    pub tau: f32,
    /// hard maximum of a thought in tokens
    #[arg(long, default_value_t = 4096)]
    pub cap: usize,
    /// per-request wall-clock budget (ms, 0 = off): thoughts are cut and read out before arrival + budget
    #[arg(long, default_value_t = 700)]
    pub budget_ms: u32,
    /// deterministic thought budget (ms, 0 = off): each request's per-branch thought cap = floor((T - fixed) / per_token) from the fixed cost model of
    /// --think-cost (prompt rows incl. image soft tokens, images, questions, thinking questions), clipped to --cap. No clock is read.
    #[arg(long, default_value_t = 0)]
    pub think_ms: u32,
    /// JSON cost model for --think-ms (think_cost.json, calibrated on an RTX PRO 6000)
    #[arg(long)]
    pub think_cost: Option<String>,
    /// context limit per question: prompt + image tokens + reserved thought tokens
    #[arg(long, default_value_t = 32768)]
    pub max_len: usize,
    #[arg(long, default_value_t = 64)]
    pub max_body_mb: usize,
    /// hard row limit of one forward (workspace, ~0.45 MB per row); no-think prefill packs fill it
    #[arg(long, default_value_t = 8192)]
    pub max_rows: usize,
    /// forward row cap while any branch is thinking or waiting for its readout (0 = --max-rows). Thinking deadlines are protected by the
    /// cut-before-forward check (which uses the expected duration of the forward about to run, incl. vision-tower encodes); a larger cap
    /// trades decode-step frequency for prefill throughput (round 3: 4096 vs 1024 = +17 % E3 policy req/s at c=64, text p99 unchanged)
    #[arg(long, default_value_t = 4096)]
    pub think_max_rows: usize,
    /// KV pages (16 tokens) of the full-attention layers; default max(2500, (max_len + cap) / 16 + 1024)
    #[arg(long)]
    pub pages: Option<usize>,
    /// KV memory budget (GB): the full-attention pool holds one max-length request plus 16K tokens, the sliding-window pool gets the rest
    #[arg(long, default_value_t = 9.0)]
    pub kv_gb: f64,
    /// Gemma-4 image budget in soft tokens per image (max_soft_tokens: 70/140/280/560/1120)
    #[arg(long, default_value_t = 280)]
    pub image_soft_tokens: usize,
    /// one request in the engine at a time (no cross-request batching): an answer then does not depend on concurrent load
    #[arg(long)]
    pub serial: bool,
    /// disable the cross-request radix prefix cache (no KV reuse across requests)
    #[arg(long)]
    pub no_prefix_cache: bool,
}

impl Args {
    /// the server's answer policy
    pub fn policy(&self) -> anyhow::Result<Policy> {
        let think_cost = match &self.think_cost {
            Some(p) => policy::ThinkCost::from_json(&std::fs::read_to_string(p)?).map_err(|e| anyhow::anyhow!("--think-cost {p}: {e:#}"))?,
            None => {
                anyhow::ensure!(self.think_ms == 0, "--think-ms needs --think-cost FILE (the calibrated cost model)");
                policy::ThinkCost::default()
            }
        };
        Ok(Policy { tau: self.tau, cap: self.cap, budget_ms: self.budget_ms, think_ms: self.think_ms, think_cost, ..Policy::default() })
    }
    fn model_file(&self, f: &str) -> String {
        std::path::Path::new(&self.model).join(f).to_string_lossy().into_owned()
    }
    pub fn tokenizer_path(&self) -> String {
        self.tokenizer.clone().unwrap_or_else(|| self.model_file("tokenizer.json"))
    }
}

/// the GPU engine as the server backend (`--serial`: one request in the engine at a time)
pub struct EngineBackend(dllm_runtime::Engine, Option<std::sync::Mutex<()>>);
impl Backend for EngineBackend {
    fn run(&self, items: &[Item], images: &[Vec<u8>], pol: &Policy) -> anyhow::Result<Vec<ItemResult>> {
        let _g = self.1.as_ref().map(|m| m.lock().unwrap_or_else(|e| e.into_inner()));
        self.0.run(items, images, pol)
    }
}

/// load the decoder, the vision tower and the MTP draft on GPU 0 and start the scheduler
pub fn start_engine(a: &Args) -> anyhow::Result<EngineBackend> {
    use dllm_core::{config::{AssistantConfig, LayerKind, ModelConfig}, loader::WeightStore};
    use dllm_kernels::Gpu;
    use dllm_model::draft::{Draft, DraftVocab};
    use dllm_model::gpu_model::{GpuModel, GpuOpts};
    use dllm_runtime::{Engine, EngineConfig, SpecDecoder, TowerEncoder, VisionPrep};
    anyhow::ensure!(a.image_soft_tokens > 0, "--image-soft-tokens must be > 0");
    anyhow::ensure!((1..=16).contains(&a.spec_max_branches), "--spec-max-branches must be 1..=16");
    let dir = std::path::PathBuf::from(&a.model);
    let cfg = ModelConfig::from_file(dir.join("config.json"))?;
    let st = WeightStore::open(&dir)?;
    // KV pools: bytes per 16-token page of the full-attention / sliding-window layers
    let pb = |k: LayerKind| -> usize { cfg.text.layers.iter().filter(|s| s.kind == k).map(|s| 2 * s.n_kv_heads * 16 * s.head_dim * 2).sum() };
    let (full_pb, sl_pb) = (pb(LayerKind::Full), pb(LayerKind::Sliding));
    let num_pages = a.pages.unwrap_or_else(|| 2500usize.max((a.max_len + a.cap + 15) / 16 + 1024));
    let sliding_pages = ((a.kv_gb * 1e9) as usize).saturating_sub(num_pages * full_pb) / sl_pb.max(1);
    anyhow::ensure!(sliding_pages >= 256, "sliding-window pool too small ({sliding_pages} pages): raise --kv-gb");
    eprintln!("KV pools: {} full-attention pages ({:.2} GB, {} tokens), {} sliding-window pages ({:.2} GB)", num_pages, (num_pages * full_pb) as f64 / 1e9, num_pages * 16, sliding_pages, ((sliding_pages + 1) * sl_pb) as f64 / 1e9);
    let opts = GpuOpts { max_rows: a.max_rows, num_pages, max_pos: a.max_len + 64, sliding_pages };
    let gpu = Gpu::new(dllm_kernels::gpu::CudaBackend::new(0)? as Arc<dyn dllm_kernels::Backend>);
    let prefix = if st.contains("model.language_model.norm.weight") { "model.language_model." } else { "model." };
    let m = GpuModel::load(gpu, &st, &cfg.text, prefix, opts)?;
    let tok = Arc::new(Tok::from_file(a.tokenizer_path())?);
    let vc = cfg.vision.clone().ok_or_else(|| anyhow::anyhow!("config.json has no vision_config"))?;
    let vis: Box<dyn dllm_runtime::VisionEncoder> = Box::new(TowerEncoder(dllm_model::vision::VisionTower::load(&m.gpu, &st, &vc, cfg.text.hidden_size)?));
    let prep: Arc<dyn dllm_runtime::ImagePrep> = Arc::new(VisionPrep(dllm_vision::PreprocConfig { max_soft_tokens: a.image_soft_tokens, ..Default::default() }));
    let ecfg = EngineConfig { max_len: a.max_len, prep: Some(prep), max_prefill_rows: (2 * a.max_rows).max(2048), think_max_rows: a.think_max_rows.min(a.max_rows), prefix_cache: !a.no_prefix_cache, ..Default::default() };
    let draft_dir = std::path::Path::new(&a.spec_draft);
    let acfg = AssistantConfig::from_file(draft_dir.join("config.json"))?;
    let vocab = DraftVocab::from_json_file(a.spec_vocab.clone().unwrap_or_else(|| a.model_file("draft_vocab_16384.json")))?;
    let draft = Draft::load(m.gpu.clone(), &WeightStore::open(draft_dir)?, &acfg, &cfg.text, &vocab, a.max_len + 64)?;
    let spec = dllm_model::spec::Spec::new(&m, draft, a.max_len + 64, 256)?;
    eprintln!("speculative decoding: draft {} (vocab {}), adaptive k, spec step for up to {} thinking branches", a.spec_draft, vocab.len(), a.spec_max_branches);
    let eng = Engine::start_with(ecfg, Box::new(m), Some(vis), tok, Box::new(SpecDecoder::new(spec, a.spec_max_branches)))?;
    eprintln!("image budget {} soft tokens/image, serial {}, prefix cache {}", a.image_soft_tokens, a.serial, !a.no_prefix_cache);
    Ok(EngineBackend(eng, a.serial.then(|| std::sync::Mutex::new(()))))
}

/// the server state with the GPU engine (also used by the GPU tests, in process)
pub fn engine_state(a: &Args) -> anyhow::Result<Arc<AppState>> {
    dllm_core::tokenizer::set_img_soft_max(a.image_soft_tokens);
    let pol = a.policy()?;
    if a.think_ms > 0 {
        eprintln!("think budget: deterministic {} ms, hard cap {}, cost model {}", a.think_ms, a.cap, serde_json::to_string(&pol.think_cost)?);
    }
    let be = start_engine(a)?;
    Ok(Arc::new(AppState { tok: Tok::from_file(a.tokenizer_path())?, backend: Box::new(be), pol, api_key: a.api_key.clone(), max_len: a.max_len, max_body: a.max_body_mb << 20, served_name: a.served_name.clone() }))
}

/// start the engine, warm it up and serve until the process ends
pub async fn serve(a: Args) -> anyhow::Result<()> {
    let st = engine_state(&a)?;
    warmup(&st);
    let l = tokio::net::TcpListener::bind(&a.bind).await?;
    eprintln!("listening on {}", a.bind);
    axum::serve(l, app(st)).await?;
    Ok(())
}

/// compile / allocate / capture every shape family once (text, many questions, noul, images, forced short thought) before accepting traffic
pub fn warmup(st: &Arc<AppState>) {
    let t0 = std::time::Instant::now();
    let mk = |n: usize, k: usize| -> Value { json!((0..n).map(|i| (format!("q{i}"), json!({"type": "choice", "instructions": format!("Which item number {i}?"), "criteria": (0..k).map(|j| (format!("o{j}"), json!(format!("item {j}")))).collect::<serde_json::Map<_, _>>()}))).collect::<serde_json::Map<_, _>>()) };
    let mut reqs = vec![
        json!({"state": "The color is red.", "questions": mk(1, 2)}),
        json!({"state": {"a": (0..300).collect::<Vec<_>>()}, "questions": mk(12, 5)}),
        json!({"state": "y", "questions": {"n": {"type": "noul", "instructions": "Is it red?"}}}),
        json!({"state": {"a": 3, "b": 4}, "questions": {"s": {"type": "number", "instructions": "What is a + b?"}}}),
    ];
    // image path: a small synthetic PNG (upscaled by the preprocessor to the full patch budget, i.e. the biggest tower shape), one and two images
    const PNG_B64: &str = "iVBORw0KGgoAAAANSUhEUgAAAGAAAABACAIAAABqVuVZAAABl0lEQVR4nN3TwQrCMBCE4RVTelDf/3U9CCJa2zSTZGYWcg38fMNeIuK2xN32PZax/eUaa4lofkvEDfiOvyv2/bi/YEDD++i+INDwProvAmSwP96PABnsj/cTT8zDtxnIY3+8vxnIY3+8n3ViNr5tQDb74/1tQDb74/2UE3PybQBy2h/vbwBy2h/vn39iZr5ngcz2x/vPApntj/dPPjE/31NAfvvj/aeA/PbH+2eemKVvPZDl/nh/PZDl/nj/tBNz9a0Ect0f768Ect0f759zYsa+NUDG++P9NUDG++P9E07M2/cQyHt/vP8QyHt/uH/4ibn7HgDR+/i++0D8Prbv2BPL4LsDJNFH990Bkuij+447sSS+/4BU+ui+/4BU+ui+g04sj+8mkFAf3XcTSKiP7jvixFL5/gJp9dF9f4G0+ui+3U8sm+8XkFwf3fcLSK6P7tv3xBL6fgIp9tF9P4EU++i+HU8sp+8bSLSP7vsGEu2j+/Y6sbS+LyDdPrrvC0i3j+7b5cQy+5ZYpfvoviVW6T66L35iuX2f8mlqD6Wvw2UAAAAASUVORK5CYII=";
    let img = format!("data:image/png;base64,{PNG_B64}");
    reqs.push(json!({"state": "x", "images": [img.clone()], "questions": mk(2, 4)}));
    reqs.push(json!({"state": "x", "images": [img.clone(), img], "questions": mk(1, 3)}));
    for tau in [0.8f32, 1.01] {
        for r in &reqs {
            let Ok(req) = serde_json::from_value::<wire::Request>(r.clone()) else { continue };
            let Ok(imgs) = decode_images(&req.images) else { continue };
            let Ok(items) = build_items(st, &req, imgs.len()) else { continue };
            let pol = Policy { tau, cap: if tau > 1.0 { 24 } else { st.pol.cap }, ..st.pol };
            if let Err(e) = st.backend.run(&items, &imgs, &pol) {
                eprintln!("warmup request failed: {e:#}");
            }
        }
    }
    eprintln!("warmup done in {:.1}s", t0.elapsed().as_secs_f32());
}

#[cfg(test)]
mod tests {
    use super::*;
    use axum::body::Body;
    use tower::ServiceExt;

    #[test]
    fn kv_capacity_is_400_max_tokens_exceeded_with_di_marker_and_transient_shortage_is_529() {
        let cap = backend_error("KV pool exhausted: KV cache capacity exceeded (3652 pages needed, 2500 pages total)");
        assert_eq!((cap.code, cap.kind), (400, "max_tokens_exceeded"));
        assert!(cap.msg.contains("maximum context length"), "{}", cap.msg);
        let sw = backend_error("sliding-window KV pool exhausted: KV cache capacity exceeded (3000 pages needed, 2496 pages total)");
        assert_eq!(sw.code, 400);
        let tr = backend_error("KV pool exhausted (2500 pages total, 300 needed, 100 free)");
        assert_eq!((tr.code, tr.kind), (529, "overloaded"));
        assert!(!tr.msg.contains("maximum context length"));
        assert_eq!(backend_error("boom").code, 500);
    }

    /// uniform distributions, never thinks
    struct UniformBackend;
    impl Backend for UniformBackend {
        fn run(&self, items: &[Item], _images: &[Vec<u8>], _pol: &Policy) -> anyhow::Result<Vec<ItemResult>> {
            Ok(items.iter().map(|it| ItemResult { p0: vec![1.0 / it.keys.len() as f32; it.keys.len()], think: None, trace_tokens: 0, out_tokens: 1, ..Default::default() }).collect())
        }
    }

    /// deterministic pseudo-model: logits are a hash of the prompt ids; questions whose max prob < tau "think" (post-thought logits = another hash)
    struct HashBackend;
    impl Backend for HashBackend {
        fn run(&self, items: &[Item], _images: &[Vec<u8>], pol: &Policy) -> anyhow::Result<Vec<ItemResult>> {
            fn logits(ids: &[u32], n: usize, salt: u64) -> Vec<f32> {
                let mut h = 0xcbf29ce484222325u64 ^ salt;
                for &t in ids {
                    h = (h ^ t as u64).wrapping_mul(0x100000001b3);
                }
                (0..n)
                    .map(|i| {
                        let mut x = h ^ (i as u64).wrapping_mul(0x9e3779b97f4a7c15);
                        x ^= x >> 33;
                        x = x.wrapping_mul(0xff51afd7ed558ccd);
                        x ^= x >> 29;
                        ((x % 6000) as f32 / 1000.0) - 1.0
                    })
                    .collect()
            }
            Ok(items
                .iter()
                .map(|it| {
                    let n = it.keys.len();
                    let p0 = logits(&it.ids, n, 1);
                    let think = (!it.number && pol.gate != policy::GateMode::Never && policy::needs_think(&policy::softmax(&p0), pol.tau)).then(|| logits(&it.ids, n, 2));
                    let tt = if think.is_some() { 7.min(it.ids.len()) } else { 0 };
                    // fake trace = the last `tt` prompt ids (decodable text) when asked for
                    let trace_ids = if pol.return_trace { it.ids[it.ids.len() - tt..].to_vec() } else { vec![] };
                    ItemResult { p0, think, trace_tokens: tt, out_tokens: 1 + tt, trace_ids, ..Default::default() }
                })
                .collect())
        }
    }

    fn state(key: Option<&str>) -> Option<Arc<AppState>> {
        state_with(key, Box::new(UniformBackend))
    }
    fn state_with(key: Option<&str>, backend: Box<dyn Backend>) -> Option<Arc<AppState>> {
        let p = std::path::Path::new(env!("CARGO_MANIFEST_DIR")).join("../../../tokenizer.json");
        Some(Arc::new(AppState { tok: Tok::from_file(p).unwrap(), backend, pol: Policy::default(), api_key: key.map(String::from), max_len: 8192, max_body: 1 << 20, served_name: "span-02".into() }))
    }
    async fn post_json(app: Router, body: Value, auth: Option<&str>) -> (u16, Value) {
        let mut r = axum::http::Request::post("/v1/systemone").header("content-type", "application/json");
        if let Some(a) = auth {
            r = r.header("authorization", a);
        }
        let resp = app.oneshot(r.body(Body::from(body.to_string())).unwrap()).await.unwrap();
        let code = resp.status().as_u16();
        let b = axum::body::to_bytes(resp.into_body(), 1 << 20).await.unwrap();
        (code, serde_json::from_slice(&b).unwrap())
    }

    #[tokio::test]
    async fn wire_format_and_errors() {
        let Some(st) = state(Some("k")) else { return };
        let app = app(st);
        let body = json!({"model": "x", "state": {"hp": 7}, "questions": {
            "act": {"type": "choice", "instructions": "next?", "criteria": {"up": "move up", "down": "move down", "wait": "wait"}},
            "ok": {"type": "noul", "instructions": "healthy?"}}});
        let (c, v) = post_json(app.clone(), body.clone(), Some("Bearer k")).await;
        assert_eq!(c, 200, "{v}");
        // exactly what decision_index.engines.base.validate checks
        let a = &v["answers"]["act"];
        assert_eq!(a["type"], "choice");
        let p = a["probabilities"].as_object().unwrap();
        assert_eq!(p.len(), 3);
        assert!((p.values().map(|x| x.as_f64().unwrap()).sum::<f64>() - 1.0).abs() < 0.01);
        assert!(p.contains_key(a["choice"].as_str().unwrap()));
        assert_eq!(v["answers"]["ok"]["type"], "noul");
        assert!((v["answers"]["ok"]["noul"].as_f64().unwrap() - 0.5).abs() < 1e-6);
        assert!(a.get("reasoning").is_none(), "no reasoning field on an answer that did not think");
        assert_eq!(v["usage"].as_object().unwrap().keys().collect::<Vec<_>>(), ["input_tokens", "reasoning_tokens", "output_tokens", "latency_ms"]);
        assert_eq!((v["usage"]["reasoning_tokens"].as_u64(), v["usage"]["output_tokens"].as_u64()), (Some(0), Some(0)));
        let (c, e) = post_json(app.clone(), body.clone(), None).await;
        assert_eq!((c, &e), (401, &json!({"error": {"type": "authentication_error", "message": "invalid or missing bearer token"}})));
        let (c, e) = post_json(app.clone(), json!({"state": 1, "questions": {}}), Some("Bearer k")).await;
        assert_eq!((c, e["error"]["type"].as_str()), (422, Some("validation_error")));
        assert_eq!(post_json(app.clone(), json!({"state": 1, "questions": {"q": {"type": "weird"}}}), Some("Bearer k")).await.0, 422);
        assert_eq!(post_json(app.clone(), json!({"state": 1, "images": ["data:image/png;base64,@@"], "questions": {"q": {"type": "noul", "instructions": "x"}}}), Some("Bearer k")).await.0, 422);
        // the old alias route is gone: 404 in the error format
        let r = axum::http::Request::post("/v1/decisions").header("content-type", "application/json").body(Body::from(body.to_string())).unwrap();
        let resp = app.clone().oneshot(r).await.unwrap();
        assert_eq!(resp.status().as_u16(), 404);
        let e: Value = serde_json::from_slice(&axum::body::to_bytes(resp.into_body(), 1 << 20).await.unwrap()).unwrap();
        assert_eq!(e["error"]["type"], "not_found");
        let big = json!({"state": "x".repeat(200_000), "questions": {"q": {"type": "noul", "instructions": "x"}}});
        let (c, v) = post_json(app, big, Some("Bearer k")).await;
        assert!(c == 400 || c == 413, "{c} {v}");
    }

    #[tokio::test]
    async fn think_override_and_traces() {
        let Some(st) = state_with(None, Box::new(HashBackend)) else { return };
        let app = app(st);
        let qs = json!((0..16).map(|i| (format!("q{i}"), json!({"type": "choice", "instructions": format!("pick {i}"), "criteria": {"a": "alpha", "b": "beta", "c": "gamma"}}))).collect::<serde_json::Map<_, _>>());
        let body = |extra: Value| {
            let mut b = json!({"state": {"hp": 7}, "questions": qs.clone()});
            b.as_object_mut().unwrap().extend(extra.as_object().unwrap().clone());
            b
        };
        let thought = |v: &Value| v["answers"].as_object().unwrap().values().filter(|a| a.get("reasoning").is_some()).count();
        let traces = |v: &Value| v["answers"].as_object().unwrap().values().filter(|a| a["reasoning"].get("trace").is_some()).count();
        let strip = |v: &Value| {
            let mut a = v["answers"].clone();
            a.as_object_mut().unwrap().values_mut().for_each(|x| drop(x.as_object_mut().unwrap().remove("reasoning")));
            a
        };
        let (c, auto) = post_json(app.clone(), body(json!({})), None).await;
        assert_eq!(c, 200);
        let na = thought(&auto);
        assert!(na > 0 && na < 16, "hash backend should split at tau 0.8: {na}");
        assert_eq!(traces(&auto), 0, "no trace unless asked");
        let rt: u64 = auto["answers"].as_object().unwrap().values().filter_map(|a| a["reasoning"]["tokens"].as_u64()).sum();
        assert!(rt > 0 && auto["usage"]["reasoning_tokens"].as_u64() == Some(rt));
        // mode "auto" == no field; unknown top-level fields (incl. the removed aliases think / budget_ms / return_trace) are ignored
        for x in [json!({"reasoning": {"mode": "auto"}}), json!({"reasoning": {}}), json!({"reasoning": null}), json!({"think": "on", "return_trace": true, "budget_ms": 5})] {
            let (c, v) = post_json(app.clone(), body(x.clone()), None).await;
            assert_eq!((c, &v["answers"]), (200, &auto["answers"]), "{x}");
        }
        for x in [json!({"reasoning": {"mode": "on", "trace": true}}), json!({"reasoning": {"mode": "on", "trace": true, "budget_ms": 0}})] {
            let (_, on) = post_json(app.clone(), body(x.clone()), None).await;
            assert_eq!((thought(&on), traces(&on)), (16, 16), "{x}");
            assert!(on["answers"].as_object().unwrap().values().all(|a| !a["reasoning"]["trace"].as_str().unwrap().is_empty()));
        }
        let (_, off) = post_json(app.clone(), body(json!({"reasoning": {"mode": "off", "trace": true}})), None).await;
        assert_eq!((thought(&off), off["usage"]["reasoning_tokens"].as_u64()), (0, Some(0)));
        // traces only for questions that thought (auto); answers unchanged by asking for them
        let (_, at) = post_json(app.clone(), body(json!({"reasoning": {"trace": true}})), None).await;
        assert_eq!((thought(&at), traces(&at)), (na, na));
        assert_eq!(strip(&at), strip(&auto));
        for bad in [json!({"reasoning": {"mode": "maybe"}}), json!({"reasoning": {"mode": true}}), json!({"reasoning": {"trace": "yes"}}), json!({"reasoning": {"budget_ms": -1}}),
                    json!({"reasoning": {"modes": "on"}}), json!({"reasoning": "on"})] {
            let (c, v) = post_json(app.clone(), body(bad.clone()), None).await;
            assert_eq!((c, v["error"]["type"].as_str()), (422, Some("validation_error")), "{bad} {v}");
        }
    }

    #[tokio::test]
    async fn budget_ms_field() {
        let base = Policy { budget_ms: 700, think_ms: 0, ..Policy::default() };
        let req = |extra: Value| -> Request {
            let mut b = json!({"state": 1, "questions": {"q": {"type": "noul", "instructions": "x"}}});
            b.as_object_mut().unwrap().extend(extra.as_object().unwrap().clone());
            serde_json::from_value(b).unwrap()
        };
        let bud = |extra: Value| request_policy(&base, &req(extra)).map(|p| p.budget_ms).map_err(|e| e.code);
        assert_eq!(bud(json!({})), Ok(700), "server budget by default");
        assert_eq!(bud(json!({"reasoning": {"mode": "auto"}})), Ok(700));
        assert_eq!(bud(json!({"budget_ms": 0, "think": "on"})), Ok(700), "top-level fields are not reasoning controls");
        assert_eq!(bud(json!({"reasoning": {"mode": "on"}})), Ok(0), "mode on without budget_ms: full thought");
        assert_eq!(bud(json!({"reasoning": {"budget_ms": 1200.4}})), Ok(1200));
        assert_eq!(bud(json!({"reasoning": {"budget_ms": -1}})), Err(422));
        assert_eq!(bud(json!({"reasoning": {"mode": "on", "budget_ms": 2500}})), Ok(2500));
        assert_eq!(bud(json!({"reasoning": {"budget_ms": 0}})), Ok(0));
        assert_eq!(bud(json!({"reasoning": {"budget_ms": 600001}})), Err(422));
        // think-ms (deterministic cap) is untouched by budget_ms
        let det = Policy { budget_ms: 0, think_ms: 750, ..Policy::default() };
        let p = request_policy(&det, &req(json!({"reasoning": {"mode": "on", "budget_ms": 300}}))).ok().unwrap();
        assert_eq!((p.think_ms, p.budget_ms), (750, 300));
        // over HTTP: a bad value is a 422, a string is a 422 schema error
        let Some(st) = state_with(None, Box::new(HashBackend)) else { return };
        let app = app(st);
        let body = |b: Value| json!({"state": 1, "questions": {"q": {"type": "noul", "instructions": "x"}}, "reasoning": {"budget_ms": b}});
        assert_eq!(post_json(app.clone(), body(json!(-5)), None).await.0, 422);
        assert_eq!(post_json(app.clone(), body(json!("100")), None).await.0, 422);
        assert_eq!(post_json(app, body(json!(0)), None).await.0, 200);
    }

    #[tokio::test]
    async fn number_questions() {
        assert_eq!(parse_number("391"), Some(391.0));
        assert_eq!(parse_number(" -12.5 "), Some(-12.5));
        assert_eq!(parse_number("1,250"), Some(1250.0));
        assert_eq!(parse_number("1e-05"), Some(1e-5));
        assert_eq!(parse_number("about 3"), None);
        assert_eq!(parse_number(""), None);
        let Some(st) = state_with(None, Box::new(HashBackend)) else { return };
        // the prompt is the trained number prompt: no OPTIONS block, "Answer with the number only."
        let req: Request = serde_json::from_value(json!({"state": {"a": 3}, "questions": {"n": {"type": "number", "instructions": "a + 4?"}, "c": {"type": "choice", "instructions": "x", "criteria": {"a": "1", "b": "2"}}}})).unwrap();
        let (items, _) = build_questions(&st, &req, 0).ok().unwrap();
        assert!(items[0].number && items[0].keys.is_empty() && items[0].label_ids.is_empty() && !items[1].number);
        let text = st.tok.decode(&items[0].ids);
        assert!(text.contains("exactly one number") && text.ends_with("Answer with the number only.<turn|>\n<|turn>model\n<|channel>thought\n") && !text.contains("OPTIONS"), "{text}");
        let (c, v) = post_json(app(st), json!({"state": 1, "questions": {"n": {"type": "number", "instructions": "how many?"}}}), None).await;
        assert_eq!(c, 200, "{v}");
        assert_eq!(v["answers"]["n"].as_object().unwrap().keys().collect::<Vec<_>>(), ["type", "number", "text", "confidence"]);
        assert_eq!(v["answers"]["n"]["type"], "number");
    }

    #[test]
    fn score_confidence_typesafe_examples() {
        assert!((score_confidence(&[0.0, 0.5, 0.5]) - 0.25).abs() < 1e-6);
        assert_eq!(score_confidence(&[0.5, 0.0, 0.5]), 0.0);
        assert!((score_confidence(&[0.0, 0.57, 0.43]) - 0.355).abs() < 1e-5);
        assert!((score_confidence(&[0.0, 0.0, 1.0]) - 1.0).abs() < 1e-6);
        assert!((score_confidence(&[0.25; 4]) - 0.0).abs() < 1e-6);
        assert!((score_confidence(&[0.0, 0.05, 0.9, 0.05, 0.0]) - (1.0 - 0.1 / 1.2)).abs() < 1e-5); // docs' explorer default ~0.92
    }

    #[tokio::test]
    async fn score_parsing_and_shape() {
        let Some(st) = state_with(None, Box::new(HashBackend)) else { return };
        let app = app(st);
        let levels = json!(["calm", {"desc": "annoyed", "w": 1}, ["angry", 2]]);
        let body = json!({"state": {"msg": "hi"}, "questions": {
            "c": {"type": "choice", "instructions": "pick", "criteria": {"a": "alpha", "b": "beta"}},
            "s": {"type": "score", "instructions": "how frustrated?", "criteria": levels.clone()},
            "n": {"type": "noul", "instructions": "ok?"}}});
        let (c, v) = post_json(app.clone(), body, None).await;
        assert_eq!(c, 200, "{v}");
        assert_eq!(v["answers"].as_object().unwrap().keys().collect::<Vec<_>>(), ["c", "s", "n"]);
        let s = &v["answers"]["s"];
        assert_eq!(s.as_object().unwrap().keys().collect::<Vec<_>>(), ["type", "score", "legend", "probabilities", "confidence"]);
        assert_eq!(s["type"], "score");
        assert_eq!(s["legend"], json!({"0": "calm", "1": {"desc": "annoyed", "w": 1}, "2": ["angry", 2]}));
        let p = s["probabilities"].as_object().unwrap();
        assert_eq!(p.keys().collect::<Vec<_>>(), ["0", "1", "2"]);
        let pv: Vec<f32> = p.values().map(|x| x.as_f64().unwrap() as f32).collect();
        assert!((pv.iter().sum::<f32>() - 1.0).abs() < 1e-3);
        let sc = s["score"].as_f64().unwrap() as f32;
        assert!((sc - (pv[1] + 2.0 * pv[2])).abs() < 1e-4 && (0.0..=2.0).contains(&sc));
        assert!((s["confidence"].as_f64().unwrap() as f32 - score_confidence(&pv)).abs() < 1e-6);
        assert!(v["answers"]["c"].get("confidence").is_some());
        // score runs exactly as a choice question over the same options (same prompt -> same probabilities)
        let as_choice = json!({"state": {"msg": "hi"}, "questions": {"s": {"type": "choice", "instructions": "how frustrated?", "criteria": {"0": "calm", "1": {"desc": "annoyed", "w": 1}, "2": ["angry", 2]}}}});
        let (_, vc) = post_json(app.clone(), as_choice, None).await;
        assert_eq!(vc["answers"]["s"]["probabilities"], s["probabilities"]);
        for bad in [json!(["one"]), json!((0..11).map(|i| format!("l{i}")).collect::<Vec<_>>()), json!({"0": "a", "1": "b"}), json!("a"), json!([])] {
            let (c, v) = post_json(app.clone(), json!({"state": 1, "questions": {"s": {"type": "score", "instructions": "x", "criteria": bad}}}), None).await;
            assert_eq!(c, 422, "{v}");
        }
        let (c, _) = post_json(app.clone(), json!({"state": 1, "questions": {"s": {"type": "score", "instructions": "x"}}}), None).await;
        assert_eq!(c, 422);
        let ten = json!((0..10).map(|i| format!("l{i}")).collect::<Vec<_>>());
        let (c, v) = post_json(app, json!({"state": 1, "questions": {"s": {"type": "score", "instructions": "x", "criteria": ten}}}), None).await;
        assert_eq!(c, 200, "{v}");
        assert_eq!(v["answers"]["s"]["probabilities"].as_object().unwrap().len(), 10);
    }

    #[test]
    fn choice_confidence_typesafe_example() {
        assert!((choice_confidence(0.6, 3) - 0.4).abs() < 1e-6);
        assert_eq!(choice_confidence(1.0, 1), 1.0);
        assert!((choice_confidence(1.0 / 3.0, 3)).abs() < 1e-6);
        assert!((choice_confidence(1.0, 4) - 1.0).abs() < 1e-6);
        assert!((choice_confidence(0.9, 3) - 0.85).abs() < 1e-6);
    }

    #[tokio::test]
    async fn official_fields_noul_criteria_model() {
        let Some(st) = state_with(None, Box::new(HashBackend)) else { return };
        let app = app(st.clone());
        let q = |c: Option<Value>| {
            let mut n = json!({"type": "noul", "instructions": "refund requested?"});
            if let Some(c) = c {
                n["criteria"] = c;
            }
            json!({"model": "jev-latest", "state": "I want my money back", "questions": {"n": n, "c": {"type": "choice", "instructions": "tone", "criteria": {"calm": null, "angry": "upset"}}}})
        };
        let (c, plain) = post_json(app.clone(), q(None), None).await;
        assert_eq!(c, 200, "{plain}");
        assert_eq!(plain["model"], "span-02");
        assert_eq!(plain["answers"]["n"].as_object().unwrap().keys().collect::<Vec<_>>(), ["type", "noul"]);
        let ch = &plain["answers"]["c"];
        assert_eq!(ch.as_object().unwrap().keys().collect::<Vec<_>>(), ["type", "choice", "confidence", "probabilities"]);
        let pm = ch["probabilities"][ch["choice"].as_str().unwrap()].as_f64().unwrap() as f32;
        assert!((ch["confidence"].as_f64().unwrap() as f32 - choice_confidence(pm, 2)).abs() < 1e-6);
        // the default descriptions are exactly today's rendering; null criteria == absent
        let (_, dflt) = post_json(app.clone(), q(Some(json!({"false": "False", "true": "True"}))), None).await;
        assert_eq!(dflt["answers"], plain["answers"]);
        let (_, nul) = post_json(app.clone(), q(Some(Value::Null)), None).await;
        assert_eq!(nul["answers"], plain["answers"]);
        let mk = |c: Value| build_questions(&st, &serde_json::from_value(q(Some(c))).unwrap(), 0).map(|x| x.0[0].ids.clone());
        let base = build_questions(&st, &serde_json::from_value(q(None)).unwrap(), 0).ok().unwrap().0[0].ids.clone();
        assert_eq!(mk(json!({"true": null, "false": null})).ok().unwrap(), base);
        // DI rows send `"criteria": {}` on noul questions (calibration-v1): same as absent, not a 422
        assert_eq!(mk(json!({})).ok().unwrap(), base);
        assert_ne!(mk(json!({"true": "A refund is requested", "false": {"no": ["refund"]}})).ok().unwrap(), base);
        let (c, _) = post_json(app.clone(), q(Some(json!({"true": "A refund is requested", "false": {"no": ["refund"]}}))), None).await;
        assert_eq!(c, 200);
        for bad in [json!({"true": "x"}), json!({"true": "x", "false": "y", "maybe": "z"}), json!(["x", "y"]), json!("x"), json!({"yes": "x", "no": "y"})] {
            let (c, v) = post_json(app.clone(), q(Some(bad)), None).await;
            assert_eq!(c, 422, "{v}");
        }
        // choice null description renders as its key
        let ids = |cr: Value| build_questions(&st, &serde_json::from_value(json!({"state": 1, "questions": {"c": {"type": "choice", "instructions": "t", "criteria": cr}}})).unwrap(), 0).ok().unwrap().0[0].ids.clone();
        assert_eq!(ids(json!({"calm": null, "angry": "upset"})), ids(json!({"calm": "calm", "angry": "upset"})));
        let mut b = q(None);
        b.as_object_mut().unwrap().remove("model");
        assert_eq!(post_json(app, b, None).await.1["model"], "span-02");
    }

    const PNG_B64: &str = "iVBORw0KGgoAAAANSUhEUgAAAGAAAABACAIAAABqVuVZAAABl0lEQVR4nN3TwQrCMBCE4RVTelDf/3U9CCJa2zSTZGYWcg38fMNeIuK2xN32PZax/eUaa4lofkvEDfiOvyv2/bi/YEDD++i+INDwProvAmSwP96PABnsj/cTT8zDtxnIY3+8vxnIY3+8n3ViNr5tQDb74/1tQDb74/2UE3PybQBy2h/vbwBy2h/vn39iZr5ngcz2x/vPApntj/dPPjE/31NAfvvj/aeA/PbH+2eemKVvPZDl/nh/PZDl/nj/tBNz9a0Ect0f768Ect0f759zYsa+NUDG++P9NUDG++P9E07M2/cQyHt/vP8QyHt/uH/4ibn7HgDR+/i++0D8Prbv2BPL4LsDJNFH990Bkuij+447sSS+/4BU+ui+/4BU+ui+g04sj+8mkFAf3XcTSKiP7jvixFL5/gJp9dF9f4G0+ui+3U8sm+8XkFwf3fcLSK6P7tv3xBL6fgIp9tF9P4EU++i+HU8sp+8bSLSP7vsGEu2j+/Y6sbS+LyDdPrrvC0i3j+7b5cQy+5ZYpfvoviVW6T66L35iuX2f8mlqD6Wvw2UAAAAASUVORK5CYII=";
    const JPG_B64: &str = "/9j/4AAQSkZJRgABAQAAAQABAAD/2wBDABALDA4MChAODQ4SERATGCgaGBYWGDEjJR0oOjM9PDkzODdASFxOQERXRTc4UG1RV19iZ2hnPk1xeXBkeFxlZ2P/2wBDARESEhgVGC8aGi9jQjhCY2NjY2NjY2NjY2NjY2NjY2NjY2NjY2NjY2NjY2NjY2NjY2NjY2NjY2NjY2NjY2NjY2P/wAARCAAQABADASIAAhEBAxEB/8QAHwAAAQUBAQEBAQEAAAAAAAAAAAECAwQFBgcICQoL/8QAtRAAAgEDAwIEAwUFBAQAAAF9AQIDAAQRBRIhMUEGE1FhByJxFDKBkaEII0KxwRVS0fAkM2JyggkKFhcYGRolJicoKSo0NTY3ODk6Q0RFRkdISUpTVFVWV1hZWmNkZWZnaGlqc3R1dnd4eXqDhIWGh4iJipKTlJWWl5iZmqKjpKWmp6ipqrKztLW2t7i5usLDxMXGx8jJytLT1NXW19jZ2uHi4+Tl5ufo6erx8vP09fb3+Pn6/8QAHwEAAwEBAQEBAQEBAQAAAAAAAAECAwQFBgcICQoL/8QAtREAAgECBAQDBAcFBAQAAQJ3AAECAxEEBSExBhJBUQdhcRMiMoEIFEKRobHBCSMzUvAVYnLRChYkNOEl8RcYGRomJygpKjU2Nzg5OkNERUZHSElKU1RVVldYWVpjZGVmZ2hpanN0dXZ3eHl6goOEhYaHiImKkpOUlZaXmJmaoqOkpaanqKmqsrO0tba3uLm6wsPExcbHyMnK0tPU1dbX2Nna4uPk5ebn6Onq8vP09fb3+Pn6/9oADAMBAAIRAxEAPwDHooorhPqD/9k=";

    /// records the prompt text and image count of the last run; uniform answers
    struct RecBackend(Arc<std::sync::Mutex<(Vec<String>, usize)>>, Arc<Tok>);
    impl Backend for RecBackend {
        fn run(&self, items: &[Item], images: &[Vec<u8>], _pol: &Policy) -> anyhow::Result<Vec<ItemResult>> {
            *self.0.lock().unwrap() = (items.iter().map(|it| self.1.decode(&it.ids)).collect(), images.len());
            UniformBackend.run(items, images, _pol)
        }
    }
    fn rec_app() -> (Router, Arc<std::sync::Mutex<(Vec<String>, usize)>>) {
        let p = std::path::Path::new(env!("CARGO_MANIFEST_DIR")).join("../../../tokenizer.json");
        let rec = Arc::new(std::sync::Mutex::new((vec![], 0)));
        let st = state_with(None, Box::new(RecBackend(rec.clone(), Arc::new(Tok::from_file(p).unwrap())))).unwrap();
        (app(st), rec)
    }
    async fn post_raw(app: Router, body: Vec<u8>, headers: &[(&str, &str)]) -> (u16, Value) {
        let mut r = axum::http::Request::post("/v1/systemone");
        for (k, v) in headers {
            r = r.header(*k, *v);
        }
        let resp = app.oneshot(r.body(Body::from(body)).unwrap()).await.unwrap();
        let code = resp.status().as_u16();
        let b = axum::body::to_bytes(resp.into_body(), 1 << 20).await.unwrap();
        (code, serde_json::from_slice(&b).unwrap_or(Value::Null))
    }

    #[tokio::test]
    async fn assets_render_as_numbered_images_in_order() {
        let (app, rec) = rec_app();
        let body = json!({"model": "jev-latest", "state": {"photo": "a3f1c2e8-found"},
            "assets": {"a3f1c2e8-found": {"type": "image", "format": "jpg", "data_base64": JPG_B64}, "77b0d4a1-poster": {"format": "png", "data_base64": PNG_B64}},
            "questions": {"same": {"type": "noul", "instructions": "Is the dog in `photo` the same dog as in this picture: 77b0d4a1-poster?"},
                          "which": {"type": "choice", "instructions": "Which shows a dog?", "criteria": {"found": "a3f1c2e8-found", "poster": {"image": "77b0d4a1-poster"}}},
                          "s": {"type": "score", "instructions": "x", "criteria": ["a3f1c2e8-found", "b"]}}});
        let (c, v) = post_json(app.clone(), body.clone(), None).await;
        assert_eq!(c, 200, "{v}");
        let (texts, nimg) = rec.lock().unwrap().clone();
        assert_eq!(nimg, 2);
        assert!(texts[0].starts_with("<bos><|turn>system\n") && texts[0].contains("<|turn>user\n<|image|><|image|>SHARED STATE (JSON string):\n{\"photo\": \"Image 1\"}"), "{}", texts[0]);
        assert!(texts[0].contains("the same dog as in this picture: Image 2?"), "{}", texts[0]);
        assert!(texts[1].contains("A: Image 1\nB: {\"image\": \"Image 2\"}"), "{}", texts[1]);
        assert!(!texts.iter().any(|t| t.contains("a3f1c2e8") || t.contains("77b0d4a1")));
        assert_eq!(v["answers"]["s"]["legend"], json!({"0": "a3f1c2e8-found", "1": "b"}), "legend echoes the placeholder");
        assert!(v["image_tokens"].as_u64().unwrap() > 0);
        // `images` come first: Image 1 = the data URI, the asset becomes Image 2
        let b2 = json!({"state": "is asset-1 red?", "images": [format!("data:image/png;base64,{PNG_B64}")], "assets": {"asset-1": {"format": "jpeg", "data_base64": JPG_B64}}, "questions": {"q": {"type": "noul", "instructions": "x"}}});
        assert_eq!(post_json(app.clone(), b2, None).await.0, 200);
        let (texts, nimg) = rec.lock().unwrap().clone();
        assert!(nimg == 2 && texts[0].contains("\"is Image 2 red?\""), "{}", texts[0]);
        // a single `images` entry renders exactly as before (no labels)
        let b3 = json!({"state": "s", "images": [PNG_B64], "questions": {"q": {"type": "noul", "instructions": "x"}}});
        assert_eq!(post_json(app.clone(), b3, None).await.0, 200);
        let want = Tok::from_file(std::path::Path::new(env!("CARGO_MANIFEST_DIR")).join("../../../tokenizer.json")).unwrap().build(&state_text(&json!("s")), "x", &["True".into(), "False".into()], 1).unwrap().unwrap().text;
        assert_eq!(rec.lock().unwrap().0[0], want);
        // referencing is optional: one unreferenced asset renders byte-identically to the same image sent via `images` (single image, no label)
        let b4 = json!({"state": "s", "assets": {"photo-1": {"type": "image", "format": "png", "data_base64": PNG_B64}}, "questions": {"q": {"type": "noul", "instructions": "x"}}});
        let (c, v) = post_json(app.clone(), b4, None).await;
        assert_eq!(c, 200, "{v}");
        let (texts, nimg) = rec.lock().unwrap().clone();
        assert_eq!((texts[0].as_str(), nimg), (want.as_str(), 1));
        // two assets, one referenced: both images shown in send order, only the referenced one labelled in the text
        let b5 = json!({"state": "is ph red?", "assets": {"other": {"format": "png", "data_base64": PNG_B64}, "ph": {"format": "jpg", "data_base64": JPG_B64}}, "questions": {"q": {"type": "noul", "instructions": "x"}}});
        assert_eq!(post_json(app.clone(), b5, None).await.0, 200);
        let (texts, nimg) = rec.lock().unwrap().clone();
        assert!(nimg == 2 && texts[0].contains("<|image|><|image|>") && texts[0].contains("\"is Image 2 red?\""), "{}", texts[0]);
        // rejections (TypeSafe: 400 per rejected asset)
        let q = json!({"q": {"type": "noul", "instructions": "about ph?"}});
        for (name, assets) in [
            ("type", json!({"ph": {"type": "audio", "format": "png", "data_base64": PNG_B64}})),
            ("format", json!({"ph": {"format": "gif", "data_base64": PNG_B64}})),
            ("format_missing", json!({"ph": {"data_base64": PNG_B64}})),
            ("format_mismatch", json!({"ph": {"format": "jpg", "data_base64": PNG_B64}})),
            ("bad_base64", json!({"ph": {"format": "png", "data_base64": "@@@"}})),
            ("not_image", json!({"ph": {"format": "png", "data_base64": "aGVsbG8="}})),
            ("no_data", json!({"ph": {"format": "png"}})),
            ("missing_part", json!({"ph": {"part": "p1"}})),
            ("empty_placeholder", json!({" ": {"format": "png", "data_base64": PNG_B64}})),
            ("not_object", json!(["x"])),
            ("five", json!((0..5).map(|i| (format!("ph{i}"), json!({"format": "png", "data_base64": PNG_B64}))).collect::<serde_json::Map<_, _>>())),
        ] {
            let st = if name == "five" { json!("ph0 ph1 ph2 ph3 ph4") } else { json!("s") };
            let (c, v) = post_json(app.clone(), json!({"state": st, "assets": assets, "questions": q}), None).await;
            assert_eq!((c, v["error"]["type"].as_str()), (400, Some("invalid_asset")), "{name}: {v}");
        }
    }

    #[tokio::test]
    async fn multipart_and_compressed_bodies() {
        let (app, rec) = rec_app();
        let req = json!({"state": {"photo": "<p>"}, "assets": {"<p>": {"type": "image", "part": "door-photo"}}, "questions": {"q": {"type": "noul", "instructions": "Is the door in `photo` open?"}}});
        let jpg = base64_simd::STANDARD.decode_to_vec(JPG_B64).unwrap();
        let mp = |parts: &[(&str, Option<&str>, &[u8])]| {
            let mut b = vec![];
            for (name, ct, data) in parts {
                b.extend(format!("--XyZ\r\nContent-Disposition: form-data; name=\"{name}\"; filename=\"{name}\"\r\n").as_bytes());
                if let Some(ct) = ct {
                    b.extend(format!("Content-Type: {ct}\r\n").as_bytes());
                }
                b.extend(b"\r\n");
                b.extend(*data);
                b.extend(b"\r\n");
            }
            b.extend(b"--XyZ--\r\n");
            b
        };
        let rj = req.to_string();
        let ct = [("content-type", "multipart/form-data; boundary=XyZ")];
        let (c, v) = post_raw(app.clone(), mp(&[("request", Some("application/json"), rj.as_bytes()), ("door-photo", Some("image/jpeg"), &jpg)]), &ct).await;
        assert_eq!(c, 200, "{v}");
        let (texts, nimg) = rec.lock().unwrap().clone();
        assert!(nimg == 1 && texts[0].contains("{\"photo\": \"Image 1\"}"), "{}", texts[0]);
        for (name, parts) in [
            ("unused_part", vec![("request", Some("application/json"), rj.as_bytes()), ("door-photo", None, &jpg[..]), ("extra", None, &jpg[..])]),
            ("bad_ctype", vec![("request", Some("application/json"), rj.as_bytes()), ("door-photo", Some("text/plain"), &jpg[..])]),
            ("no_request", vec![("door-photo", None, &jpg[..])]),
            ("missing_part", vec![("request", None, rj.as_bytes())]),
        ] {
            let (c, v) = post_raw(app.clone(), mp(&parts), &ct).await;
            assert_eq!(c, 400, "{name}: {v}");
        }
        // Content-Encoding: zstd (the TypeSafe Python SDK default) and gzip give the same answers as the plain body
        let plain = json!({"state": "s", "questions": {"q": {"type": "choice", "instructions": "x", "criteria": {"a": "1", "b": "2"}}}}).to_string();
        let j = [("content-type", "application/json")];
        let (c0, v0) = post_raw(app.clone(), plain.clone().into_bytes(), &j).await;
        assert_eq!(c0, 200);
        let z = zstd::encode_all(plain.as_bytes(), 3).unwrap();
        let (c, v) = post_raw(app.clone(), z, &[j[0], ("content-encoding", "zstd")]).await;
        assert_eq!((c, &v["answers"]), (200, &v0["answers"]), "{v}");
        let mut g = flate2::write::GzEncoder::new(vec![], flate2::Compression::default());
        std::io::Write::write_all(&mut g, plain.as_bytes()).unwrap();
        let (c, v) = post_raw(app.clone(), g.finish().unwrap(), &[j[0], ("content-encoding", "gzip")]).await;
        assert_eq!((c, &v["answers"]), (200, &v0["answers"]), "{v}");
        let (c, v) = post_raw(app.clone(), plain.clone().into_bytes(), &[j[0], ("content-encoding", "br")]).await;
        assert_eq!((c, v["error"]["type"].as_str()), (415, Some("unsupported_encoding")));
        let (c, _) = post_raw(app.clone(), plain.clone().into_bytes(), &[j[0], ("content-encoding", "gzip, zstd")]).await;
        assert_eq!(c, 415);
        let (c, v) = post_raw(app.clone(), b"not zstd".to_vec(), &[j[0], ("content-encoding", "zstd")]).await;
        assert_eq!((c, v["error"]["type"].as_str()), (400, Some("invalid_request")));
        let (c, v) = post_raw(app.clone(), b"{not json".to_vec(), &j).await;
        assert_eq!((c, &v), (400, &json!({"error": {"type": "invalid_request", "message": "invalid JSON body"}})));
    }

    #[tokio::test]
    async fn health_is_minimal_and_context_overflow_is_400() {
        let Some(st) = state(None) else { return };
        let resp = app(st.clone()).oneshot(axum::http::Request::get("/health").body(Body::empty()).unwrap()).await.unwrap();
        let b = axum::body::to_bytes(resp.into_body(), 1 << 20).await.unwrap();
        assert_eq!(serde_json::from_slice::<Value>(&b).unwrap(), json!({"ok": true, "model": "span-02"}));
        let (c, v) = post_json(app(st), json!({"state": "x ".repeat(9000), "questions": {"q": {"type": "noul", "instructions": "x"}}}), None).await;
        assert_eq!((c, v["error"]["type"].as_str()), (400, Some("max_tokens_exceeded")));
        assert!(v["error"]["message"].as_str().unwrap().contains("maximum context length"));
    }
}

/// HTTP error for a failed backend run. A request larger than the whole KV pool can never fit: 422 with the DI capacity marker
/// ("maximum context length"), so clients do not retry it. A transient shortage (concurrent requests hold the KV cache) stays 529.
pub fn backend_error(m: &str) -> ApiError {
    if m.contains("cannot decode image") {
        ApiError::new(422, "cannot decode image as a base64 data URI")
    } else if m.contains("KV cache capacity exceeded") {
        ApiError::context(format!("This model's maximum context length is exceeded: the request does not fit in the server's KV cache ({m})"))
    } else if m.contains("KV pool exhausted") {
        // transient overload (concurrent requests hold the KV cache): TypeSafe's 529, clients retry after a short delay
        ApiError::new(529, format!("overloaded, retry after a short delay ({m})"))
    } else {
        ApiError::new(500, format!("inference failed: {m}"))
    }
}